Skip to content

Commit 9ef9eb3

Browse files
Fanglidingmaoxikun
Fangliding
authored and
maoxikun
committed
Refine must2 and apply NewAesGcm() to all usage (XTLS#5011)
* Refine must2 and apply NewAesGcm() to all usage * Remove unused package * Fix test (cherry picked from commit b1107b9)
1 parent 3408f7b commit 9ef9eb3

File tree

12 files changed

+37
-96
lines changed

12 files changed

+37
-96
lines changed

‎app/dns/dnscommon_test.go‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,31 +18,31 @@ func Test_parseResponse(t *testing.T) {
1818

1919
ans := new(dns.Msg)
2020
ans.Id = 0
21-
p = append(p, common.Must2(ans.Pack()).([]byte))
21+
p = append(p, common.Must2(ans.Pack()))
2222

2323
p = append(p, []byte{})
2424

2525
ans = new(dns.Msg)
2626
ans.Id = 1
2727
ans.Answer = append(ans.Answer,
28-
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")).(dns.RR),
29-
common.Must2(dns.NewRR("google.com. IN CNAME fake.google.com")).(dns.RR),
30-
common.Must2(dns.NewRR("google.com. IN A 8.8.8.8")).(dns.RR),
31-
common.Must2(dns.NewRR("google.com. IN A 8.8.4.4")).(dns.RR),
28+
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")),
29+
common.Must2(dns.NewRR("google.com. IN CNAME fake.google.com")),
30+
common.Must2(dns.NewRR("google.com. IN A 8.8.8.8")),
31+
common.Must2(dns.NewRR("google.com. IN A 8.8.4.4")),
3232
)
33-
p = append(p, common.Must2(ans.Pack()).([]byte))
33+
p = append(p, common.Must2(ans.Pack()))
3434

3535
ans = new(dns.Msg)
3636
ans.Id = 2
3737
ans.Answer = append(ans.Answer,
38-
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")).(dns.RR),
39-
common.Must2(dns.NewRR("google.com. IN CNAME fake.google.com")).(dns.RR),
40-
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")).(dns.RR),
41-
common.Must2(dns.NewRR("google.com. IN CNAME test.google.com")).(dns.RR),
42-
common.Must2(dns.NewRR("google.com. IN AAAA 2001::123:8888")).(dns.RR),
43-
common.Must2(dns.NewRR("google.com. IN AAAA 2001::123:8844")).(dns.RR),
38+
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")),
39+
common.Must2(dns.NewRR("google.com. IN CNAME fake.google.com")),
40+
common.Must2(dns.NewRR("google.com. IN CNAME m.test.google.com")),
41+
common.Must2(dns.NewRR("google.com. IN CNAME test.google.com")),
42+
common.Must2(dns.NewRR("google.com. IN AAAA 2001::123:8888")),
43+
common.Must2(dns.NewRR("google.com. IN AAAA 2001::123:8844")),
4444
)
45-
p = append(p, common.Must2(ans.Pack()).([]byte))
45+
p = append(p, common.Must2(ans.Pack()))
4646

4747
tests := []struct {
4848
name string

‎common/common.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@ func Must(err error) {
2323
}
2424

2525
// Must2 panics if the second parameter is not nil, otherwise returns the first parameter.
26-
func Must2(v interface{}, err error) interface{} {
26+
// This is useful when function returned "sth, err" and avoid many "if err != nil"
27+
// Internal usage only, if user input can cause err, it must be handled
28+
func Must2[T any](v T, err error) T {
2729
Must(err)
2830
return v
2931
}

‎common/crypto/aes.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,7 @@ func NewAesCTRStream(key []byte, iv []byte) cipher.Stream {
3232

3333
// NewAesGcm creates a AEAD cipher based on AES-GCM.
3434
func NewAesGcm(key []byte) cipher.AEAD {
35-
block, err := aes.NewCipher(key)
36-
common.Must(err)
37-
aead, err := cipher.NewGCM(block)
38-
common.Must(err)
35+
block := common.Must2(aes.NewCipher(key))
36+
aead := common.Must2(cipher.NewGCM(block))
3937
return aead
4038
}

‎common/crypto/auth_test.go‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@ package crypto_test
22

33
import (
44
"bytes"
5-
"crypto/aes"
6-
"crypto/cipher"
75
"crypto/rand"
86
"io"
97
"testing"
@@ -18,11 +16,8 @@ import (
1816
func TestAuthenticationReaderWriter(t *testing.T) {
1917
key := make([]byte, 16)
2018
rand.Read(key)
21-
block, err := aes.NewCipher(key)
22-
common.Must(err)
2319

24-
aead, err := cipher.NewGCM(block)
25-
common.Must(err)
20+
aead := NewAesGcm(key)
2621

2722
const payloadSize = 1024 * 80
2823
rawPayload := make([]byte, payloadSize)
@@ -71,7 +66,7 @@ func TestAuthenticationReaderWriter(t *testing.T) {
7166
t.Error(r)
7267
}
7368

74-
_, err = reader.ReadMultiBuffer()
69+
_, err := reader.ReadMultiBuffer()
7570
if err != io.EOF {
7671
t.Error("error: ", err)
7772
}
@@ -80,11 +75,8 @@ func TestAuthenticationReaderWriter(t *testing.T) {
8075
func TestAuthenticationReaderWriterPacket(t *testing.T) {
8176
key := make([]byte, 16)
8277
common.Must2(rand.Read(key))
83-
block, err := aes.NewCipher(key)
84-
common.Must(err)
8578

86-
aead, err := cipher.NewGCM(block)
87-
common.Must(err)
79+
aead := NewAesGcm(key)
8880

8981
cache := buf.New()
9082
iv := make([]byte, 12)

‎proxy/dokodemo/dokodemo.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,15 +91,15 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn st
9191
}
9292
}
9393
if dest.Port == 0 {
94-
dest.Port = net.Port(common.Must2(strconv.Atoi(port)).(int))
94+
dest.Port = net.Port(common.Must2(strconv.Atoi(port)))
9595
}
9696
if d.portMap != nil && d.portMap[port] != "" {
9797
h, p, _ := net.SplitHostPort(d.portMap[port])
9898
if len(h) > 0 {
9999
dest.Address = net.ParseAddress(h)
100100
}
101101
if len(p) > 0 {
102-
dest.Port = net.Port(common.Must2(strconv.Atoi(p)).(int))
102+
dest.Port = net.Port(common.Must2(strconv.Atoi(p)))
103103
}
104104
}
105105
}

‎proxy/shadowsocks/config.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package shadowsocks
22

33
import (
44
"bytes"
5-
"crypto/aes"
65
"crypto/cipher"
76
"crypto/md5"
87
"crypto/sha1"
@@ -58,11 +57,7 @@ func (a *MemoryAccount) CheckIV(iv []byte) error {
5857
}
5958

6059
func createAesGcm(key []byte) cipher.AEAD {
61-
block, err := aes.NewCipher(key)
62-
common.Must(err)
63-
gcm, err := cipher.NewGCM(block)
64-
common.Must(err)
65-
return gcm
60+
return crypto.NewAesGcm(key)
6661
}
6762

6863
func createChaCha20Poly1305(key []byte) cipher.AEAD {

‎proxy/vmess/aead/encrypt.go‎

Lines changed: 5 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,13 @@ package aead
22

33
import (
44
"bytes"
5-
"crypto/aes"
6-
"crypto/cipher"
75
"crypto/rand"
86
"encoding/binary"
97
"io"
108
"time"
119

1210
"github.com/xtls/xray-core/common"
11+
"github.com/xtls/xray-core/common/crypto"
1312
)
1413

1514
func SealVMessAEADHeader(key [16]byte, data []byte) []byte {
@@ -34,15 +33,7 @@ func SealVMessAEADHeader(key [16]byte, data []byte) []byte {
3433

3534
payloadHeaderLengthAEADNonce := KDF(key[:], KDFSaltConstVMessHeaderPayloadLengthAEADIV, string(generatedAuthID[:]), string(connectionNonce))[:12]
3635

37-
payloadHeaderLengthAEADAESBlock, err := aes.NewCipher(payloadHeaderLengthAEADKey)
38-
if err != nil {
39-
panic(err.Error())
40-
}
41-
42-
payloadHeaderAEAD, err := cipher.NewGCM(payloadHeaderLengthAEADAESBlock)
43-
if err != nil {
44-
panic(err.Error())
45-
}
36+
payloadHeaderAEAD := crypto.NewAesGcm(payloadHeaderLengthAEADKey)
4637

4738
payloadHeaderLengthAEADEncrypted = payloadHeaderAEAD.Seal(nil, payloadHeaderLengthAEADNonce, aeadPayloadLengthSerializedByte, generatedAuthID[:])
4839
}
@@ -54,15 +45,7 @@ func SealVMessAEADHeader(key [16]byte, data []byte) []byte {
5445

5546
payloadHeaderAEADNonce := KDF(key[:], KDFSaltConstVMessHeaderPayloadAEADIV, string(generatedAuthID[:]), string(connectionNonce))[:12]
5647

57-
payloadHeaderAEADAESBlock, err := aes.NewCipher(payloadHeaderAEADKey)
58-
if err != nil {
59-
panic(err.Error())
60-
}
61-
62-
payloadHeaderAEAD, err := cipher.NewGCM(payloadHeaderAEADAESBlock)
63-
if err != nil {
64-
panic(err.Error())
65-
}
48+
payloadHeaderAEAD := crypto.NewAesGcm(payloadHeaderAEADKey)
6649

6750
payloadHeaderAEADEncrypted = payloadHeaderAEAD.Seal(nil, payloadHeaderAEADNonce, data, generatedAuthID[:])
6851
}
@@ -104,15 +87,7 @@ func OpenVMessAEADHeader(key [16]byte, authid [16]byte, data io.Reader) ([]byte,
10487

10588
payloadHeaderLengthAEADNonce := KDF(key[:], KDFSaltConstVMessHeaderPayloadLengthAEADIV, string(authid[:]), string(nonce[:]))[:12]
10689

107-
payloadHeaderAEADAESBlock, err := aes.NewCipher(payloadHeaderLengthAEADKey)
108-
if err != nil {
109-
panic(err.Error())
110-
}
111-
112-
payloadHeaderLengthAEAD, err := cipher.NewGCM(payloadHeaderAEADAESBlock)
113-
if err != nil {
114-
panic(err.Error())
115-
}
90+
payloadHeaderLengthAEAD := crypto.NewAesGcm(payloadHeaderLengthAEADKey)
11691

11792
decryptedAEADHeaderLengthPayload, erropenAEAD := payloadHeaderLengthAEAD.Open(nil, payloadHeaderLengthAEADNonce, payloadHeaderLengthAEADEncrypted[:], authid[:])
11893

@@ -145,15 +120,7 @@ func OpenVMessAEADHeader(key [16]byte, authid [16]byte, data io.Reader) ([]byte,
145120
return nil, false, bytesRead, err
146121
}
147122

148-
payloadHeaderAEADAESBlock, err := aes.NewCipher(payloadHeaderAEADKey)
149-
if err != nil {
150-
panic(err.Error())
151-
}
152-
153-
payloadHeaderAEAD, err := cipher.NewGCM(payloadHeaderAEADAESBlock)
154-
if err != nil {
155-
panic(err.Error())
156-
}
123+
payloadHeaderAEAD := crypto.NewAesGcm(payloadHeaderAEADKey)
157124

158125
decryptedAEADHeaderPayload, erropenAEAD := payloadHeaderAEAD.Open(nil, payloadHeaderAEADNonce, payloadHeaderAEADEncrypted, authid[:])
159126

‎proxy/vmess/encoding/client.go‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@ package encoding
33
import (
44
"bytes"
55
"context"
6-
"crypto/aes"
7-
"crypto/cipher"
86
"crypto/rand"
97
"crypto/sha256"
108
"encoding/binary"
@@ -182,8 +180,7 @@ func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.Respon
182180
aeadResponseHeaderLengthEncryptionKey := vmessaead.KDF16(c.responseBodyKey[:], vmessaead.KDFSaltConstAEADRespHeaderLenKey)
183181
aeadResponseHeaderLengthEncryptionIV := vmessaead.KDF(c.responseBodyIV[:], vmessaead.KDFSaltConstAEADRespHeaderLenIV)[:12]
184182

185-
aeadResponseHeaderLengthEncryptionKeyAESBlock := common.Must2(aes.NewCipher(aeadResponseHeaderLengthEncryptionKey)).(cipher.Block)
186-
aeadResponseHeaderLengthEncryptionAEAD := common.Must2(cipher.NewGCM(aeadResponseHeaderLengthEncryptionKeyAESBlock)).(cipher.AEAD)
183+
aeadResponseHeaderLengthEncryptionAEAD := crypto.NewAesGcm(aeadResponseHeaderLengthEncryptionKey)
187184

188185
var aeadEncryptedResponseHeaderLength [18]byte
189186
var decryptedResponseHeaderLength int
@@ -205,8 +202,7 @@ func (c *ClientSession) DecodeResponseHeader(reader io.Reader) (*protocol.Respon
205202
aeadResponseHeaderPayloadEncryptionKey := vmessaead.KDF16(c.responseBodyKey[:], vmessaead.KDFSaltConstAEADRespHeaderPayloadKey)
206203
aeadResponseHeaderPayloadEncryptionIV := vmessaead.KDF(c.responseBodyIV[:], vmessaead.KDFSaltConstAEADRespHeaderPayloadIV)[:12]
207204

208-
aeadResponseHeaderPayloadEncryptionKeyAESBlock := common.Must2(aes.NewCipher(aeadResponseHeaderPayloadEncryptionKey)).(cipher.Block)
209-
aeadResponseHeaderPayloadEncryptionAEAD := common.Must2(cipher.NewGCM(aeadResponseHeaderPayloadEncryptionKeyAESBlock)).(cipher.AEAD)
205+
aeadResponseHeaderPayloadEncryptionAEAD := crypto.NewAesGcm(aeadResponseHeaderPayloadEncryptionKey)
210206

211207
encryptedResponseHeaderBuffer := make([]byte, decryptedResponseHeaderLength+16)
212208

‎proxy/vmess/encoding/server.go‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@ package encoding
22

33
import (
44
"bytes"
5-
"crypto/aes"
6-
"crypto/cipher"
75
"crypto/sha256"
86
"encoding/binary"
97
"hash/fnv"
@@ -350,8 +348,7 @@ func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, wr
350348
aeadResponseHeaderLengthEncryptionKey := vmessaead.KDF16(s.responseBodyKey[:], vmessaead.KDFSaltConstAEADRespHeaderLenKey)
351349
aeadResponseHeaderLengthEncryptionIV := vmessaead.KDF(s.responseBodyIV[:], vmessaead.KDFSaltConstAEADRespHeaderLenIV)[:12]
352350

353-
aeadResponseHeaderLengthEncryptionKeyAESBlock := common.Must2(aes.NewCipher(aeadResponseHeaderLengthEncryptionKey)).(cipher.Block)
354-
aeadResponseHeaderLengthEncryptionAEAD := common.Must2(cipher.NewGCM(aeadResponseHeaderLengthEncryptionKeyAESBlock)).(cipher.AEAD)
351+
aeadResponseHeaderLengthEncryptionAEAD := crypto.NewAesGcm(aeadResponseHeaderLengthEncryptionKey)
355352

356353
aeadResponseHeaderLengthEncryptionBuffer := bytes.NewBuffer(nil)
357354

@@ -365,8 +362,7 @@ func (s *ServerSession) EncodeResponseHeader(header *protocol.ResponseHeader, wr
365362
aeadResponseHeaderPayloadEncryptionKey := vmessaead.KDF16(s.responseBodyKey[:], vmessaead.KDFSaltConstAEADRespHeaderPayloadKey)
366363
aeadResponseHeaderPayloadEncryptionIV := vmessaead.KDF(s.responseBodyIV[:], vmessaead.KDFSaltConstAEADRespHeaderPayloadIV)[:12]
367364

368-
aeadResponseHeaderPayloadEncryptionKeyAESBlock := common.Must2(aes.NewCipher(aeadResponseHeaderPayloadEncryptionKey)).(cipher.Block)
369-
aeadResponseHeaderPayloadEncryptionAEAD := common.Must2(cipher.NewGCM(aeadResponseHeaderPayloadEncryptionKeyAESBlock)).(cipher.AEAD)
365+
aeadResponseHeaderPayloadEncryptionAEAD := crypto.NewAesGcm(aeadResponseHeaderPayloadEncryptionKey)
370366

371367
aeadEncryptedHeaderPayload := aeadResponseHeaderPayloadEncryptionAEAD.Seal(nil, aeadResponseHeaderPayloadEncryptionIV, aeadEncryptedHeaderBuffer.Bytes(), nil)
372368
common.Must2(io.Copy(writer, bytes.NewReader(aeadEncryptedHeaderPayload)))

‎transport/internet/kcp/cryptreal.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
11
package kcp
22

33
import (
4-
"crypto/aes"
54
"crypto/cipher"
65
"crypto/sha256"
76

8-
"github.com/xtls/xray-core/common"
7+
"github.com/xtls/xray-core/common/crypto"
98
)
109

1110
func NewAEADAESGCMBasedOnSeed(seed string) cipher.AEAD {
1211
hashedSeed := sha256.Sum256([]byte(seed))
13-
aesBlock := common.Must2(aes.NewCipher(hashedSeed[:16])).(cipher.Block)
14-
return common.Must2(cipher.NewGCM(aesBlock)).(cipher.AEAD)
12+
return crypto.NewAesGcm(hashedSeed[:])
1513
}

0 commit comments

Comments
 (0)