From 65b467e448d7450f832dbf7dbf2eac610d0c3a2e Mon Sep 17 00:00:00 2001
From: Hellojack <106379370+H1JK@users.noreply.github.com>
Date: Fri, 16 Jun 2023 00:22:53 +0800
Subject: [PATCH] REALITY protocol: Add ChaCha20-Poly1305 auth mode (#2212)

https://github.com/XTLS/REALITY/pull/4
---
 go.mod                                |  2 +-
 go.sum                                |  4 ++--
 transport/internet/reality/reality.go | 17 +++++++++++++----
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/go.mod b/go.mod
index c15f8890..8fd18217 100644
--- a/go.mod
+++ b/go.mod
@@ -19,7 +19,7 @@ require (
 	github.com/seiflotfy/cuckoofilter v0.0.0-20220411075957-e3b120b3f5fb
 	github.com/stretchr/testify v1.8.4
 	github.com/v2fly/ss-bloomring v0.0.0-20210312155135-28617310f63e
-	github.com/xtls/reality v0.0.0-20230331223127-176a94313eda
+	github.com/xtls/reality v0.0.0-20230613075828-e07c3b04b983
 	golang.org/x/crypto v0.10.0
 	golang.org/x/net v0.11.0
 	golang.org/x/sync v0.3.0
diff --git a/go.sum b/go.sum
index de4084a7..1f53d6d2 100644
--- a/go.sum
+++ b/go.sum
@@ -166,8 +166,8 @@ github.com/v2fly/ss-bloomring v0.0.0-20210312155135-28617310f63e h1:5QefA066A1tF
 github.com/v2fly/ss-bloomring v0.0.0-20210312155135-28617310f63e/go.mod h1:5t19P9LBIrNamL6AcMQOncg/r10y3Pc01AbHeMhwlpU=
 github.com/viant/assertly v0.4.8/go.mod h1:aGifi++jvCrUaklKEKT0BU95igDNaqkvz+49uaYMPRU=
 github.com/viant/toolbox v0.24.0/go.mod h1:OxMCG57V0PXuIP2HNQrtJf2CjqdmbrOx5EkMILuUhzM=
-github.com/xtls/reality v0.0.0-20230331223127-176a94313eda h1:psRJD2RrZbnI0OWyHvXfgYCPqlRM5q5SPDcjDoDBWhE=
-github.com/xtls/reality v0.0.0-20230331223127-176a94313eda/go.mod h1:rkuAY1S9F8eI8gDiPDYvACE8e2uwkyg8qoOTuwWov7Y=
+github.com/xtls/reality v0.0.0-20230613075828-e07c3b04b983 h1:AMyzgjkh54WocjQSlCnT1LhDc/BKiUqtNOv40AkpURs=
+github.com/xtls/reality v0.0.0-20230613075828-e07c3b04b983/go.mod h1:rkuAY1S9F8eI8gDiPDYvACE8e2uwkyg8qoOTuwWov7Y=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
 go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
diff --git a/transport/internet/reality/reality.go b/transport/internet/reality/reality.go
index 6355687b..ff2d3f38 100644
--- a/transport/internet/reality/reality.go
+++ b/transport/internet/reality/reality.go
@@ -30,12 +30,16 @@ import (
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/core"
 	"github.com/xtls/xray-core/transport/internet/tls"
+	"golang.org/x/crypto/chacha20poly1305"
 	"golang.org/x/crypto/hkdf"
 	"golang.org/x/net/http2"
 )
 
 //go:generate go run github.com/xtls/xray-core/common/errors/errorgen
 
+//go:linkname aesgcmPreferred github.com/refraction-networking/utls.aesgcmPreferred
+func aesgcmPreferred(ciphers []uint16) bool
+
 type Conn struct {
 	*reality.Conn
 }
@@ -136,11 +140,16 @@ func UClient(c net.Conn, config *Config, ctx context.Context, dest net.Destinati
 		if _, err := hkdf.New(sha256.New, uConn.AuthKey, hello.Random[:20], []byte("REALITY")).Read(uConn.AuthKey); err != nil {
 			return nil, err
 		}
-		if config.Show {
-			fmt.Printf("REALITY localAddr: %v\tuConn.AuthKey[:16]: %v\n", localAddr, uConn.AuthKey[:16])
+		var aead cipher.AEAD
+		if aesgcmPreferred(hello.CipherSuites) {
+			block, _ := aes.NewCipher(uConn.AuthKey)
+			aead, _ = cipher.NewGCM(block)
+		} else {
+			aead, _ = chacha20poly1305.New(uConn.AuthKey)
+		}
+		if config.Show {
+			fmt.Printf("REALITY localAddr: %v\tuConn.AuthKey[:16]: %v\tAEAD: %T\n", localAddr, uConn.AuthKey[:16], aead)
 		}
-		block, _ := aes.NewCipher(uConn.AuthKey)
-		aead, _ := cipher.NewGCM(block)
 		aead.Seal(hello.SessionId[:0], hello.Random[20:], hello.SessionId[:16], hello.Raw)
 		copy(hello.Raw[39:], hello.SessionId)
 	}