mirror of
https://github.com/ehang-io/nps.git
synced 2025-09-06 07:06:53 +00:00
Ip限制 npc代理连接
This commit is contained in:
刘河
785
vender/github.com/xtaci/kcp/crypt.go
Normal file
785
vender/github.com/xtaci/kcp/crypt.go
Normal file
@@ -0,0 +1,785 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/des"
|
||||
"crypto/sha1"
|
||||
|
||||
"github.com/templexxx/xor"
|
||||
"github.com/tjfoc/gmsm/sm4"
|
||||
|
||||
"golang.org/x/crypto/blowfish"
|
||||
"golang.org/x/crypto/cast5"
|
||||
"golang.org/x/crypto/pbkdf2"
|
||||
"golang.org/x/crypto/salsa20"
|
||||
"golang.org/x/crypto/tea"
|
||||
"golang.org/x/crypto/twofish"
|
||||
"golang.org/x/crypto/xtea"
|
||||
)
|
||||
|
||||
var (
|
||||
initialVector = []byte{167, 115, 79, 156, 18, 172, 27, 1, 164, 21, 242, 193, 252, 120, 230, 107}
|
||||
saltxor = `sH3CIVoF#rWLtJo6`
|
||||
)
|
||||
|
||||
// BlockCrypt defines encryption/decryption methods for a given byte slice.
|
||||
// Notes on implementing: the data to be encrypted contains a builtin
|
||||
// nonce at the first 16 bytes
|
||||
type BlockCrypt interface {
|
||||
// Encrypt encrypts the whole block in src into dst.
|
||||
// Dst and src may point at the same memory.
|
||||
Encrypt(dst, src []byte)
|
||||
|
||||
// Decrypt decrypts the whole block in src into dst.
|
||||
// Dst and src may point at the same memory.
|
||||
Decrypt(dst, src []byte)
|
||||
}
|
||||
|
||||
type salsa20BlockCrypt struct {
|
||||
key [32]byte
|
||||
}
|
||||
|
||||
// NewSalsa20BlockCrypt https://en.wikipedia.org/wiki/Salsa20
|
||||
func NewSalsa20BlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(salsa20BlockCrypt)
|
||||
copy(c.key[:], key)
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *salsa20BlockCrypt) Encrypt(dst, src []byte) {
|
||||
salsa20.XORKeyStream(dst[8:], src[8:], src[:8], &c.key)
|
||||
copy(dst[:8], src[:8])
|
||||
}
|
||||
func (c *salsa20BlockCrypt) Decrypt(dst, src []byte) {
|
||||
salsa20.XORKeyStream(dst[8:], src[8:], src[:8], &c.key)
|
||||
copy(dst[:8], src[:8])
|
||||
}
|
||||
|
||||
type sm4BlockCrypt struct {
|
||||
encbuf [sm4.BlockSize]byte
|
||||
decbuf [2 * sm4.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewSM4BlockCrypt https://github.com/tjfoc/gmsm/tree/master/sm4
|
||||
func NewSM4BlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(sm4BlockCrypt)
|
||||
block, err := sm4.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *sm4BlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *sm4BlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type twofishBlockCrypt struct {
|
||||
encbuf [twofish.BlockSize]byte
|
||||
decbuf [2 * twofish.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewTwofishBlockCrypt https://en.wikipedia.org/wiki/Twofish
|
||||
func NewTwofishBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(twofishBlockCrypt)
|
||||
block, err := twofish.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *twofishBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *twofishBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type tripleDESBlockCrypt struct {
|
||||
encbuf [des.BlockSize]byte
|
||||
decbuf [2 * des.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewTripleDESBlockCrypt https://en.wikipedia.org/wiki/Triple_DES
|
||||
func NewTripleDESBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(tripleDESBlockCrypt)
|
||||
block, err := des.NewTripleDESCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *tripleDESBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *tripleDESBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type cast5BlockCrypt struct {
|
||||
encbuf [cast5.BlockSize]byte
|
||||
decbuf [2 * cast5.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewCast5BlockCrypt https://en.wikipedia.org/wiki/CAST-128
|
||||
func NewCast5BlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(cast5BlockCrypt)
|
||||
block, err := cast5.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *cast5BlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *cast5BlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type blowfishBlockCrypt struct {
|
||||
encbuf [blowfish.BlockSize]byte
|
||||
decbuf [2 * blowfish.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewBlowfishBlockCrypt https://en.wikipedia.org/wiki/Blowfish_(cipher)
|
||||
func NewBlowfishBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(blowfishBlockCrypt)
|
||||
block, err := blowfish.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *blowfishBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *blowfishBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type aesBlockCrypt struct {
|
||||
encbuf [aes.BlockSize]byte
|
||||
decbuf [2 * aes.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewAESBlockCrypt https://en.wikipedia.org/wiki/Advanced_Encryption_Standard
|
||||
func NewAESBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(aesBlockCrypt)
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *aesBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *aesBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type teaBlockCrypt struct {
|
||||
encbuf [tea.BlockSize]byte
|
||||
decbuf [2 * tea.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewTEABlockCrypt https://en.wikipedia.org/wiki/Tiny_Encryption_Algorithm
|
||||
func NewTEABlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(teaBlockCrypt)
|
||||
block, err := tea.NewCipherWithRounds(key, 16)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *teaBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *teaBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type xteaBlockCrypt struct {
|
||||
encbuf [xtea.BlockSize]byte
|
||||
decbuf [2 * xtea.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
// NewXTEABlockCrypt https://en.wikipedia.org/wiki/XTEA
|
||||
func NewXTEABlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(xteaBlockCrypt)
|
||||
block, err := xtea.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.block = block
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *xteaBlockCrypt) Encrypt(dst, src []byte) { encrypt(c.block, dst, src, c.encbuf[:]) }
|
||||
func (c *xteaBlockCrypt) Decrypt(dst, src []byte) { decrypt(c.block, dst, src, c.decbuf[:]) }
|
||||
|
||||
type simpleXORBlockCrypt struct {
|
||||
xortbl []byte
|
||||
}
|
||||
|
||||
// NewSimpleXORBlockCrypt simple xor with key expanding
|
||||
func NewSimpleXORBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
c := new(simpleXORBlockCrypt)
|
||||
c.xortbl = pbkdf2.Key(key, []byte(saltxor), 32, mtuLimit, sha1.New)
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *simpleXORBlockCrypt) Encrypt(dst, src []byte) { xor.Bytes(dst, src, c.xortbl) }
|
||||
func (c *simpleXORBlockCrypt) Decrypt(dst, src []byte) { xor.Bytes(dst, src, c.xortbl) }
|
||||
|
||||
type noneBlockCrypt struct{}
|
||||
|
||||
// NewNoneBlockCrypt does nothing but copying
|
||||
func NewNoneBlockCrypt(key []byte) (BlockCrypt, error) {
|
||||
return new(noneBlockCrypt), nil
|
||||
}
|
||||
|
||||
func (c *noneBlockCrypt) Encrypt(dst, src []byte) { copy(dst, src) }
|
||||
func (c *noneBlockCrypt) Decrypt(dst, src []byte) { copy(dst, src) }
|
||||
|
||||
// packet encryption with local CFB mode
|
||||
func encrypt(block cipher.Block, dst, src, buf []byte) {
|
||||
switch block.BlockSize() {
|
||||
case 8:
|
||||
encrypt8(block, dst, src, buf)
|
||||
case 16:
|
||||
encrypt16(block, dst, src, buf)
|
||||
default:
|
||||
encryptVariant(block, dst, src, buf)
|
||||
}
|
||||
}
|
||||
|
||||
// optimized encryption for the ciphers which works in 8-bytes
|
||||
func encrypt8(block cipher.Block, dst, src, buf []byte) {
|
||||
tbl := buf[:8]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / 8
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
s := src[base:][0:64]
|
||||
d := dst[base:][0:64]
|
||||
// 1
|
||||
xor.BytesSrc1(d[0:8], s[0:8], tbl)
|
||||
block.Encrypt(tbl, d[0:8])
|
||||
// 2
|
||||
xor.BytesSrc1(d[8:16], s[8:16], tbl)
|
||||
block.Encrypt(tbl, d[8:16])
|
||||
// 3
|
||||
xor.BytesSrc1(d[16:24], s[16:24], tbl)
|
||||
block.Encrypt(tbl, d[16:24])
|
||||
// 4
|
||||
xor.BytesSrc1(d[24:32], s[24:32], tbl)
|
||||
block.Encrypt(tbl, d[24:32])
|
||||
// 5
|
||||
xor.BytesSrc1(d[32:40], s[32:40], tbl)
|
||||
block.Encrypt(tbl, d[32:40])
|
||||
// 6
|
||||
xor.BytesSrc1(d[40:48], s[40:48], tbl)
|
||||
block.Encrypt(tbl, d[40:48])
|
||||
// 7
|
||||
xor.BytesSrc1(d[48:56], s[48:56], tbl)
|
||||
block.Encrypt(tbl, d[48:56])
|
||||
// 8
|
||||
xor.BytesSrc1(d[56:64], s[56:64], tbl)
|
||||
block.Encrypt(tbl, d[56:64])
|
||||
base += 64
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 6:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 5:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 4:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 3:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 2:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 1:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 8
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
|
||||
// optimized encryption for the ciphers which works in 16-bytes
|
||||
func encrypt16(block cipher.Block, dst, src, buf []byte) {
|
||||
tbl := buf[:16]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / 16
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
s := src[base:][0:128]
|
||||
d := dst[base:][0:128]
|
||||
// 1
|
||||
xor.BytesSrc1(d[0:16], s[0:16], tbl)
|
||||
block.Encrypt(tbl, d[0:16])
|
||||
// 2
|
||||
xor.BytesSrc1(d[16:32], s[16:32], tbl)
|
||||
block.Encrypt(tbl, d[16:32])
|
||||
// 3
|
||||
xor.BytesSrc1(d[32:48], s[32:48], tbl)
|
||||
block.Encrypt(tbl, d[32:48])
|
||||
// 4
|
||||
xor.BytesSrc1(d[48:64], s[48:64], tbl)
|
||||
block.Encrypt(tbl, d[48:64])
|
||||
// 5
|
||||
xor.BytesSrc1(d[64:80], s[64:80], tbl)
|
||||
block.Encrypt(tbl, d[64:80])
|
||||
// 6
|
||||
xor.BytesSrc1(d[80:96], s[80:96], tbl)
|
||||
block.Encrypt(tbl, d[80:96])
|
||||
// 7
|
||||
xor.BytesSrc1(d[96:112], s[96:112], tbl)
|
||||
block.Encrypt(tbl, d[96:112])
|
||||
// 8
|
||||
xor.BytesSrc1(d[112:128], s[112:128], tbl)
|
||||
block.Encrypt(tbl, d[112:128])
|
||||
base += 128
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 6:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 5:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 4:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 3:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 2:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 1:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += 16
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
|
||||
func encryptVariant(block cipher.Block, dst, src, buf []byte) {
|
||||
blocksize := block.BlockSize()
|
||||
tbl := buf[:blocksize]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / blocksize
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
// 1
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 2
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 3
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 4
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 5
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 6
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 7
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
|
||||
// 8
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 6:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 5:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 4:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 3:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 2:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 1:
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
block.Encrypt(tbl, dst[base:])
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
|
||||
// decryption
|
||||
func decrypt(block cipher.Block, dst, src, buf []byte) {
|
||||
switch block.BlockSize() {
|
||||
case 8:
|
||||
decrypt8(block, dst, src, buf)
|
||||
case 16:
|
||||
decrypt16(block, dst, src, buf)
|
||||
default:
|
||||
decryptVariant(block, dst, src, buf)
|
||||
}
|
||||
}
|
||||
|
||||
func decrypt8(block cipher.Block, dst, src, buf []byte) {
|
||||
tbl := buf[0:8]
|
||||
next := buf[8:16]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / 8
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
s := src[base:][0:64]
|
||||
d := dst[base:][0:64]
|
||||
// 1
|
||||
block.Encrypt(next, s[0:8])
|
||||
xor.BytesSrc1(d[0:8], s[0:8], tbl)
|
||||
// 2
|
||||
block.Encrypt(tbl, s[8:16])
|
||||
xor.BytesSrc1(d[8:16], s[8:16], next)
|
||||
// 3
|
||||
block.Encrypt(next, s[16:24])
|
||||
xor.BytesSrc1(d[16:24], s[16:24], tbl)
|
||||
// 4
|
||||
block.Encrypt(tbl, s[24:32])
|
||||
xor.BytesSrc1(d[24:32], s[24:32], next)
|
||||
// 5
|
||||
block.Encrypt(next, s[32:40])
|
||||
xor.BytesSrc1(d[32:40], s[32:40], tbl)
|
||||
// 6
|
||||
block.Encrypt(tbl, s[40:48])
|
||||
xor.BytesSrc1(d[40:48], s[40:48], next)
|
||||
// 7
|
||||
block.Encrypt(next, s[48:56])
|
||||
xor.BytesSrc1(d[48:56], s[48:56], tbl)
|
||||
// 8
|
||||
block.Encrypt(tbl, s[56:64])
|
||||
xor.BytesSrc1(d[56:64], s[56:64], next)
|
||||
base += 64
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 6:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 5:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 4:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 3:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 2:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 1:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 8
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
|
||||
func decrypt16(block cipher.Block, dst, src, buf []byte) {
|
||||
tbl := buf[0:16]
|
||||
next := buf[16:32]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / 16
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
s := src[base:][0:128]
|
||||
d := dst[base:][0:128]
|
||||
// 1
|
||||
block.Encrypt(next, s[0:16])
|
||||
xor.BytesSrc1(d[0:16], s[0:16], tbl)
|
||||
// 2
|
||||
block.Encrypt(tbl, s[16:32])
|
||||
xor.BytesSrc1(d[16:32], s[16:32], next)
|
||||
// 3
|
||||
block.Encrypt(next, s[32:48])
|
||||
xor.BytesSrc1(d[32:48], s[32:48], tbl)
|
||||
// 4
|
||||
block.Encrypt(tbl, s[48:64])
|
||||
xor.BytesSrc1(d[48:64], s[48:64], next)
|
||||
// 5
|
||||
block.Encrypt(next, s[64:80])
|
||||
xor.BytesSrc1(d[64:80], s[64:80], tbl)
|
||||
// 6
|
||||
block.Encrypt(tbl, s[80:96])
|
||||
xor.BytesSrc1(d[80:96], s[80:96], next)
|
||||
// 7
|
||||
block.Encrypt(next, s[96:112])
|
||||
xor.BytesSrc1(d[96:112], s[96:112], tbl)
|
||||
// 8
|
||||
block.Encrypt(tbl, s[112:128])
|
||||
xor.BytesSrc1(d[112:128], s[112:128], next)
|
||||
base += 128
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 6:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 5:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 4:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 3:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 2:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 1:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += 16
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
|
||||
func decryptVariant(block cipher.Block, dst, src, buf []byte) {
|
||||
blocksize := block.BlockSize()
|
||||
tbl := buf[:blocksize]
|
||||
next := buf[blocksize:]
|
||||
block.Encrypt(tbl, initialVector)
|
||||
n := len(src) / blocksize
|
||||
base := 0
|
||||
repeat := n / 8
|
||||
left := n % 8
|
||||
for i := 0; i < repeat; i++ {
|
||||
// 1
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
base += blocksize
|
||||
|
||||
// 2
|
||||
block.Encrypt(tbl, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], next)
|
||||
base += blocksize
|
||||
|
||||
// 3
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
base += blocksize
|
||||
|
||||
// 4
|
||||
block.Encrypt(tbl, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], next)
|
||||
base += blocksize
|
||||
|
||||
// 5
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
base += blocksize
|
||||
|
||||
// 6
|
||||
block.Encrypt(tbl, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], next)
|
||||
base += blocksize
|
||||
|
||||
// 7
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
base += blocksize
|
||||
|
||||
// 8
|
||||
block.Encrypt(tbl, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], next)
|
||||
base += blocksize
|
||||
}
|
||||
|
||||
switch left {
|
||||
case 7:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 6:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 5:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 4:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 3:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 2:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 1:
|
||||
block.Encrypt(next, src[base:])
|
||||
xor.BytesSrc1(dst[base:], src[base:], tbl)
|
||||
tbl, next = next, tbl
|
||||
base += blocksize
|
||||
fallthrough
|
||||
case 0:
|
||||
xor.BytesSrc0(dst[base:], src[base:], tbl)
|
||||
}
|
||||
}
|
||||
289
vender/github.com/xtaci/kcp/crypt_test.go
Normal file
289
vender/github.com/xtaci/kcp/crypt_test.go
Normal file
@@ -0,0 +1,289 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/aes"
|
||||
"crypto/md5"
|
||||
"crypto/rand"
|
||||
"crypto/sha1"
|
||||
"hash/crc32"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSM4(t *testing.T) {
|
||||
bc, err := NewSM4BlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestAES(t *testing.T) {
|
||||
bc, err := NewAESBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestTEA(t *testing.T) {
|
||||
bc, err := NewTEABlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestXOR(t *testing.T) {
|
||||
bc, err := NewSimpleXORBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestBlowfish(t *testing.T) {
|
||||
bc, err := NewBlowfishBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestNone(t *testing.T) {
|
||||
bc, err := NewNoneBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestCast5(t *testing.T) {
|
||||
bc, err := NewCast5BlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func Test3DES(t *testing.T) {
|
||||
bc, err := NewTripleDESBlockCrypt(pass[:24])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestTwofish(t *testing.T) {
|
||||
bc, err := NewTwofishBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestXTEA(t *testing.T) {
|
||||
bc, err := NewXTEABlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func TestSalsa20(t *testing.T) {
|
||||
bc, err := NewSalsa20BlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cryptTest(t, bc)
|
||||
}
|
||||
|
||||
func cryptTest(t *testing.T, bc BlockCrypt) {
|
||||
data := make([]byte, mtuLimit)
|
||||
io.ReadFull(rand.Reader, data)
|
||||
dec := make([]byte, mtuLimit)
|
||||
enc := make([]byte, mtuLimit)
|
||||
bc.Encrypt(enc, data)
|
||||
bc.Decrypt(dec, enc)
|
||||
if !bytes.Equal(data, dec) {
|
||||
t.Fail()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkSM4(b *testing.B) {
|
||||
bc, err := NewSM4BlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkAES128(b *testing.B) {
|
||||
bc, err := NewAESBlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkAES192(b *testing.B) {
|
||||
bc, err := NewAESBlockCrypt(pass[:24])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkAES256(b *testing.B) {
|
||||
bc, err := NewAESBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkTEA(b *testing.B) {
|
||||
bc, err := NewTEABlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkXOR(b *testing.B) {
|
||||
bc, err := NewSimpleXORBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkBlowfish(b *testing.B) {
|
||||
bc, err := NewBlowfishBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkNone(b *testing.B) {
|
||||
bc, err := NewNoneBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkCast5(b *testing.B) {
|
||||
bc, err := NewCast5BlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func Benchmark3DES(b *testing.B) {
|
||||
bc, err := NewTripleDESBlockCrypt(pass[:24])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkTwofish(b *testing.B) {
|
||||
bc, err := NewTwofishBlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkXTEA(b *testing.B) {
|
||||
bc, err := NewXTEABlockCrypt(pass[:16])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func BenchmarkSalsa20(b *testing.B) {
|
||||
bc, err := NewSalsa20BlockCrypt(pass[:32])
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
benchCrypt(b, bc)
|
||||
}
|
||||
|
||||
func benchCrypt(b *testing.B, bc BlockCrypt) {
|
||||
data := make([]byte, mtuLimit)
|
||||
io.ReadFull(rand.Reader, data)
|
||||
dec := make([]byte, mtuLimit)
|
||||
enc := make([]byte, mtuLimit)
|
||||
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(enc) * 2))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
bc.Encrypt(enc, data)
|
||||
bc.Decrypt(dec, enc)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCRC32(b *testing.B) {
|
||||
content := make([]byte, 1024)
|
||||
b.SetBytes(int64(len(content)))
|
||||
for i := 0; i < b.N; i++ {
|
||||
crc32.ChecksumIEEE(content)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCsprngSystem(b *testing.B) {
|
||||
data := make([]byte, md5.Size)
|
||||
b.SetBytes(int64(len(data)))
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
io.ReadFull(rand.Reader, data)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCsprngMD5(b *testing.B) {
|
||||
var data [md5.Size]byte
|
||||
b.SetBytes(md5.Size)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
data = md5.Sum(data[:])
|
||||
}
|
||||
}
|
||||
func BenchmarkCsprngSHA1(b *testing.B) {
|
||||
var data [sha1.Size]byte
|
||||
b.SetBytes(sha1.Size)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
data = sha1.Sum(data[:])
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCsprngNonceMD5(b *testing.B) {
|
||||
var ng nonceMD5
|
||||
ng.Init()
|
||||
b.SetBytes(md5.Size)
|
||||
data := make([]byte, md5.Size)
|
||||
for i := 0; i < b.N; i++ {
|
||||
ng.Fill(data)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCsprngNonceAES128(b *testing.B) {
|
||||
var ng nonceAES128
|
||||
ng.Init()
|
||||
|
||||
b.SetBytes(aes.BlockSize)
|
||||
data := make([]byte, aes.BlockSize)
|
||||
for i := 0; i < b.N; i++ {
|
||||
ng.Fill(data)
|
||||
}
|
||||
}
|
||||
52
vender/github.com/xtaci/kcp/entropy.go
Normal file
52
vender/github.com/xtaci/kcp/entropy.go
Normal file
@@ -0,0 +1,52 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/md5"
|
||||
"crypto/rand"
|
||||
"io"
|
||||
)
|
||||
|
||||
// Entropy defines a entropy source
|
||||
type Entropy interface {
|
||||
Init()
|
||||
Fill(nonce []byte)
|
||||
}
|
||||
|
||||
// nonceMD5 nonce generator for packet header
|
||||
type nonceMD5 struct {
|
||||
seed [md5.Size]byte
|
||||
}
|
||||
|
||||
func (n *nonceMD5) Init() { /*nothing required*/ }
|
||||
|
||||
func (n *nonceMD5) Fill(nonce []byte) {
|
||||
if n.seed[0] == 0 { // entropy update
|
||||
io.ReadFull(rand.Reader, n.seed[:])
|
||||
}
|
||||
n.seed = md5.Sum(n.seed[:])
|
||||
copy(nonce, n.seed[:])
|
||||
}
|
||||
|
||||
// nonceAES128 nonce generator for packet headers
|
||||
type nonceAES128 struct {
|
||||
seed [aes.BlockSize]byte
|
||||
block cipher.Block
|
||||
}
|
||||
|
||||
func (n *nonceAES128) Init() {
|
||||
var key [16]byte //aes-128
|
||||
io.ReadFull(rand.Reader, key[:])
|
||||
io.ReadFull(rand.Reader, n.seed[:])
|
||||
block, _ := aes.NewCipher(key[:])
|
||||
n.block = block
|
||||
}
|
||||
|
||||
func (n *nonceAES128) Fill(nonce []byte) {
|
||||
if n.seed[0] == 0 { // entropy update
|
||||
io.ReadFull(rand.Reader, n.seed[:])
|
||||
}
|
||||
n.block.Encrypt(n.seed[:], n.seed[:])
|
||||
copy(nonce, n.seed[:])
|
||||
}
|
||||
311
vender/github.com/xtaci/kcp/fec.go
Normal file
311
vender/github.com/xtaci/kcp/fec.go
Normal file
@@ -0,0 +1,311 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/klauspost/reedsolomon"
|
||||
)
|
||||
|
||||
const (
|
||||
fecHeaderSize = 6
|
||||
fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size
|
||||
typeData = 0xf1
|
||||
typeFEC = 0xf2
|
||||
)
|
||||
|
||||
type (
|
||||
// fecPacket is a decoded FEC packet
|
||||
fecPacket struct {
|
||||
seqid uint32
|
||||
flag uint16
|
||||
data []byte
|
||||
}
|
||||
|
||||
// fecDecoder for decoding incoming packets
|
||||
fecDecoder struct {
|
||||
rxlimit int // queue size limit
|
||||
dataShards int
|
||||
parityShards int
|
||||
shardSize int
|
||||
rx []fecPacket // ordered receive queue
|
||||
|
||||
// caches
|
||||
decodeCache [][]byte
|
||||
flagCache []bool
|
||||
|
||||
// zeros
|
||||
zeros []byte
|
||||
|
||||
// RS decoder
|
||||
codec reedsolomon.Encoder
|
||||
}
|
||||
)
|
||||
|
||||
func newFECDecoder(rxlimit, dataShards, parityShards int) *fecDecoder {
|
||||
if dataShards <= 0 || parityShards <= 0 {
|
||||
return nil
|
||||
}
|
||||
if rxlimit < dataShards+parityShards {
|
||||
return nil
|
||||
}
|
||||
|
||||
dec := new(fecDecoder)
|
||||
dec.rxlimit = rxlimit
|
||||
dec.dataShards = dataShards
|
||||
dec.parityShards = parityShards
|
||||
dec.shardSize = dataShards + parityShards
|
||||
codec, err := reedsolomon.New(dataShards, parityShards)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
dec.codec = codec
|
||||
dec.decodeCache = make([][]byte, dec.shardSize)
|
||||
dec.flagCache = make([]bool, dec.shardSize)
|
||||
dec.zeros = make([]byte, mtuLimit)
|
||||
return dec
|
||||
}
|
||||
|
||||
// decodeBytes a fec packet
|
||||
func (dec *fecDecoder) decodeBytes(data []byte) fecPacket {
|
||||
var pkt fecPacket
|
||||
pkt.seqid = binary.LittleEndian.Uint32(data)
|
||||
pkt.flag = binary.LittleEndian.Uint16(data[4:])
|
||||
// allocate memory & copy
|
||||
buf := xmitBuf.Get().([]byte)[:len(data)-6]
|
||||
copy(buf, data[6:])
|
||||
pkt.data = buf
|
||||
return pkt
|
||||
}
|
||||
|
||||
// decode a fec packet
|
||||
func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) {
|
||||
// insertion
|
||||
n := len(dec.rx) - 1
|
||||
insertIdx := 0
|
||||
for i := n; i >= 0; i-- {
|
||||
if pkt.seqid == dec.rx[i].seqid { // de-duplicate
|
||||
xmitBuf.Put(pkt.data)
|
||||
return nil
|
||||
} else if _itimediff(pkt.seqid, dec.rx[i].seqid) > 0 { // insertion
|
||||
insertIdx = i + 1
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// insert into ordered rx queue
|
||||
if insertIdx == n+1 {
|
||||
dec.rx = append(dec.rx, pkt)
|
||||
} else {
|
||||
dec.rx = append(dec.rx, fecPacket{})
|
||||
copy(dec.rx[insertIdx+1:], dec.rx[insertIdx:]) // shift right
|
||||
dec.rx[insertIdx] = pkt
|
||||
}
|
||||
|
||||
// shard range for current packet
|
||||
shardBegin := pkt.seqid - pkt.seqid%uint32(dec.shardSize)
|
||||
shardEnd := shardBegin + uint32(dec.shardSize) - 1
|
||||
|
||||
// max search range in ordered queue for current shard
|
||||
searchBegin := insertIdx - int(pkt.seqid%uint32(dec.shardSize))
|
||||
if searchBegin < 0 {
|
||||
searchBegin = 0
|
||||
}
|
||||
searchEnd := searchBegin + dec.shardSize - 1
|
||||
if searchEnd >= len(dec.rx) {
|
||||
searchEnd = len(dec.rx) - 1
|
||||
}
|
||||
|
||||
// re-construct datashards
|
||||
if searchEnd-searchBegin+1 >= dec.dataShards {
|
||||
var numshard, numDataShard, first, maxlen int
|
||||
|
||||
// zero caches
|
||||
shards := dec.decodeCache
|
||||
shardsflag := dec.flagCache
|
||||
for k := range dec.decodeCache {
|
||||
shards[k] = nil
|
||||
shardsflag[k] = false
|
||||
}
|
||||
|
||||
// shard assembly
|
||||
for i := searchBegin; i <= searchEnd; i++ {
|
||||
seqid := dec.rx[i].seqid
|
||||
if _itimediff(seqid, shardEnd) > 0 {
|
||||
break
|
||||
} else if _itimediff(seqid, shardBegin) >= 0 {
|
||||
shards[seqid%uint32(dec.shardSize)] = dec.rx[i].data
|
||||
shardsflag[seqid%uint32(dec.shardSize)] = true
|
||||
numshard++
|
||||
if dec.rx[i].flag == typeData {
|
||||
numDataShard++
|
||||
}
|
||||
if numshard == 1 {
|
||||
first = i
|
||||
}
|
||||
if len(dec.rx[i].data) > maxlen {
|
||||
maxlen = len(dec.rx[i].data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if numDataShard == dec.dataShards {
|
||||
// case 1: no loss on data shards
|
||||
dec.rx = dec.freeRange(first, numshard, dec.rx)
|
||||
} else if numshard >= dec.dataShards {
|
||||
// case 2: loss on data shards, but it's recoverable from parity shards
|
||||
for k := range shards {
|
||||
if shards[k] != nil {
|
||||
dlen := len(shards[k])
|
||||
shards[k] = shards[k][:maxlen]
|
||||
copy(shards[k][dlen:], dec.zeros)
|
||||
}
|
||||
}
|
||||
if err := dec.codec.ReconstructData(shards); err == nil {
|
||||
for k := range shards[:dec.dataShards] {
|
||||
if !shardsflag[k] {
|
||||
recovered = append(recovered, shards[k])
|
||||
}
|
||||
}
|
||||
}
|
||||
dec.rx = dec.freeRange(first, numshard, dec.rx)
|
||||
}
|
||||
}
|
||||
|
||||
// keep rxlimit
|
||||
if len(dec.rx) > dec.rxlimit {
|
||||
if dec.rx[0].flag == typeData { // track the unrecoverable data
|
||||
atomic.AddUint64(&DefaultSnmp.FECShortShards, 1)
|
||||
}
|
||||
dec.rx = dec.freeRange(0, 1, dec.rx)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// free a range of fecPacket, and zero for GC recycling
|
||||
func (dec *fecDecoder) freeRange(first, n int, q []fecPacket) []fecPacket {
|
||||
for i := first; i < first+n; i++ { // recycle buffer
|
||||
xmitBuf.Put(q[i].data)
|
||||
}
|
||||
copy(q[first:], q[first+n:])
|
||||
for i := 0; i < n; i++ { // dereference data
|
||||
q[len(q)-1-i].data = nil
|
||||
}
|
||||
return q[:len(q)-n]
|
||||
}
|
||||
|
||||
type (
|
||||
// fecEncoder for encoding outgoing packets
|
||||
fecEncoder struct {
|
||||
dataShards int
|
||||
parityShards int
|
||||
shardSize int
|
||||
paws uint32 // Protect Against Wrapped Sequence numbers
|
||||
next uint32 // next seqid
|
||||
|
||||
shardCount int // count the number of datashards collected
|
||||
maxSize int // track maximum data length in datashard
|
||||
|
||||
headerOffset int // FEC header offset
|
||||
payloadOffset int // FEC payload offset
|
||||
|
||||
// caches
|
||||
shardCache [][]byte
|
||||
encodeCache [][]byte
|
||||
|
||||
// zeros
|
||||
zeros []byte
|
||||
|
||||
// RS encoder
|
||||
codec reedsolomon.Encoder
|
||||
}
|
||||
)
|
||||
|
||||
func newFECEncoder(dataShards, parityShards, offset int) *fecEncoder {
|
||||
if dataShards <= 0 || parityShards <= 0 {
|
||||
return nil
|
||||
}
|
||||
enc := new(fecEncoder)
|
||||
enc.dataShards = dataShards
|
||||
enc.parityShards = parityShards
|
||||
enc.shardSize = dataShards + parityShards
|
||||
enc.paws = (0xffffffff/uint32(enc.shardSize) - 1) * uint32(enc.shardSize)
|
||||
enc.headerOffset = offset
|
||||
enc.payloadOffset = enc.headerOffset + fecHeaderSize
|
||||
|
||||
codec, err := reedsolomon.New(dataShards, parityShards)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
enc.codec = codec
|
||||
|
||||
// caches
|
||||
enc.encodeCache = make([][]byte, enc.shardSize)
|
||||
enc.shardCache = make([][]byte, enc.shardSize)
|
||||
for k := range enc.shardCache {
|
||||
enc.shardCache[k] = make([]byte, mtuLimit)
|
||||
}
|
||||
enc.zeros = make([]byte, mtuLimit)
|
||||
return enc
|
||||
}
|
||||
|
||||
// encodes the packet, outputs parity shards if we have collected quorum datashards
|
||||
// notice: the contents of 'ps' will be re-written in successive calling
|
||||
func (enc *fecEncoder) encode(b []byte) (ps [][]byte) {
|
||||
enc.markData(b[enc.headerOffset:])
|
||||
binary.LittleEndian.PutUint16(b[enc.payloadOffset:], uint16(len(b[enc.payloadOffset:])))
|
||||
|
||||
// copy data to fec datashards
|
||||
sz := len(b)
|
||||
enc.shardCache[enc.shardCount] = enc.shardCache[enc.shardCount][:sz]
|
||||
copy(enc.shardCache[enc.shardCount], b)
|
||||
enc.shardCount++
|
||||
|
||||
// track max datashard length
|
||||
if sz > enc.maxSize {
|
||||
enc.maxSize = sz
|
||||
}
|
||||
|
||||
// Generation of Reed-Solomon Erasure Code
|
||||
if enc.shardCount == enc.dataShards {
|
||||
// fill '0' into the tail of each datashard
|
||||
for i := 0; i < enc.dataShards; i++ {
|
||||
shard := enc.shardCache[i]
|
||||
slen := len(shard)
|
||||
copy(shard[slen:enc.maxSize], enc.zeros)
|
||||
}
|
||||
|
||||
// construct equal-sized slice with stripped header
|
||||
cache := enc.encodeCache
|
||||
for k := range cache {
|
||||
cache[k] = enc.shardCache[k][enc.payloadOffset:enc.maxSize]
|
||||
}
|
||||
|
||||
// encoding
|
||||
if err := enc.codec.Encode(cache); err == nil {
|
||||
ps = enc.shardCache[enc.dataShards:]
|
||||
for k := range ps {
|
||||
enc.markFEC(ps[k][enc.headerOffset:])
|
||||
ps[k] = ps[k][:enc.maxSize]
|
||||
}
|
||||
}
|
||||
|
||||
// counters resetting
|
||||
enc.shardCount = 0
|
||||
enc.maxSize = 0
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (enc *fecEncoder) markData(data []byte) {
|
||||
binary.LittleEndian.PutUint32(data, enc.next)
|
||||
binary.LittleEndian.PutUint16(data[4:], typeData)
|
||||
enc.next++
|
||||
}
|
||||
|
||||
func (enc *fecEncoder) markFEC(data []byte) {
|
||||
binary.LittleEndian.PutUint32(data, enc.next)
|
||||
binary.LittleEndian.PutUint16(data[4:], typeFEC)
|
||||
enc.next = (enc.next + 1) % enc.paws
|
||||
}
|
||||
43
vender/github.com/xtaci/kcp/fec_test.go
Normal file
43
vender/github.com/xtaci/kcp/fec_test.go
Normal file
@@ -0,0 +1,43 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func BenchmarkFECDecode(b *testing.B) {
|
||||
const dataSize = 10
|
||||
const paritySize = 3
|
||||
const payLoad = 1500
|
||||
decoder := newFECDecoder(1024, dataSize, paritySize)
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(payLoad)
|
||||
for i := 0; i < b.N; i++ {
|
||||
if rand.Int()%(dataSize+paritySize) == 0 { // random loss
|
||||
continue
|
||||
}
|
||||
var pkt fecPacket
|
||||
pkt.seqid = uint32(i)
|
||||
if i%(dataSize+paritySize) >= dataSize {
|
||||
pkt.flag = typeFEC
|
||||
} else {
|
||||
pkt.flag = typeData
|
||||
}
|
||||
pkt.data = make([]byte, payLoad)
|
||||
decoder.decode(pkt)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkFECEncode(b *testing.B) {
|
||||
const dataSize = 10
|
||||
const paritySize = 3
|
||||
const payLoad = 1500
|
||||
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(payLoad)
|
||||
encoder := newFECEncoder(dataSize, paritySize, 0)
|
||||
for i := 0; i < b.N; i++ {
|
||||
data := make([]byte, payLoad)
|
||||
encoder.encode(data)
|
||||
}
|
||||
}
|
||||
1012
vender/github.com/xtaci/kcp/kcp.go
Normal file
1012
vender/github.com/xtaci/kcp/kcp.go
Normal file
File diff suppressed because it is too large
Load Diff
302
vender/github.com/xtaci/kcp/kcp_test.go
Normal file
302
vender/github.com/xtaci/kcp/kcp_test.go
Normal file
@@ -0,0 +1,302 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"container/list"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func iclock() int32 {
|
||||
return int32(currentMs())
|
||||
}
|
||||
|
||||
type DelayPacket struct {
|
||||
_ptr []byte
|
||||
_size int
|
||||
_ts int32
|
||||
}
|
||||
|
||||
func (p *DelayPacket) Init(size int, src []byte) {
|
||||
p._ptr = make([]byte, size)
|
||||
p._size = size
|
||||
copy(p._ptr, src[:size])
|
||||
}
|
||||
|
||||
func (p *DelayPacket) ptr() []byte { return p._ptr }
|
||||
func (p *DelayPacket) size() int { return p._size }
|
||||
func (p *DelayPacket) ts() int32 { return p._ts }
|
||||
func (p *DelayPacket) setts(ts int32) { p._ts = ts }
|
||||
|
||||
type DelayTunnel struct{ *list.List }
|
||||
type LatencySimulator struct {
|
||||
current int32
|
||||
lostrate, rttmin, rttmax, nmax int
|
||||
p12 DelayTunnel
|
||||
p21 DelayTunnel
|
||||
r12 *rand.Rand
|
||||
r21 *rand.Rand
|
||||
}
|
||||
|
||||
// lostrate: 往返一周丢包率的百分比，默认 10%
|
||||
// rttmin：rtt最小值，默认 60
|
||||
// rttmax：rtt最大值，默认 125
|
||||
//func (p *LatencySimulator)Init(int lostrate = 10, int rttmin = 60, int rttmax = 125, int nmax = 1000):
|
||||
func (p *LatencySimulator) Init(lostrate, rttmin, rttmax, nmax int) {
|
||||
p.r12 = rand.New(rand.NewSource(9))
|
||||
p.r21 = rand.New(rand.NewSource(99))
|
||||
p.p12 = DelayTunnel{list.New()}
|
||||
p.p21 = DelayTunnel{list.New()}
|
||||
p.current = iclock()
|
||||
p.lostrate = lostrate / 2 // 上面数据是往返丢包率，单程除以2
|
||||
p.rttmin = rttmin / 2
|
||||
p.rttmax = rttmax / 2
|
||||
p.nmax = nmax
|
||||
}
|
||||
|
||||
// 发送数据
|
||||
// peer - 端点0/1，从0发送，从1接收；从1发送从0接收
|
||||
func (p *LatencySimulator) send(peer int, data []byte, size int) int {
|
||||
rnd := 0
|
||||
if peer == 0 {
|
||||
rnd = p.r12.Intn(100)
|
||||
} else {
|
||||
rnd = p.r21.Intn(100)
|
||||
}
|
||||
//println("!!!!!!!!!!!!!!!!!!!!", rnd, p.lostrate, peer)
|
||||
if rnd < p.lostrate {
|
||||
return 0
|
||||
}
|
||||
pkt := &DelayPacket{}
|
||||
pkt.Init(size, data)
|
||||
p.current = iclock()
|
||||
delay := p.rttmin
|
||||
if p.rttmax > p.rttmin {
|
||||
delay += rand.Int() % (p.rttmax - p.rttmin)
|
||||
}
|
||||
pkt.setts(p.current + int32(delay))
|
||||
if peer == 0 {
|
||||
p.p12.PushBack(pkt)
|
||||
} else {
|
||||
p.p21.PushBack(pkt)
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
// 接收数据
|
||||
func (p *LatencySimulator) recv(peer int, data []byte, maxsize int) int32 {
|
||||
var it *list.Element
|
||||
if peer == 0 {
|
||||
it = p.p21.Front()
|
||||
if p.p21.Len() == 0 {
|
||||
return -1
|
||||
}
|
||||
} else {
|
||||
it = p.p12.Front()
|
||||
if p.p12.Len() == 0 {
|
||||
return -1
|
||||
}
|
||||
}
|
||||
pkt := it.Value.(*DelayPacket)
|
||||
p.current = iclock()
|
||||
if p.current < pkt.ts() {
|
||||
return -2
|
||||
}
|
||||
if maxsize < pkt.size() {
|
||||
return -3
|
||||
}
|
||||
if peer == 0 {
|
||||
p.p21.Remove(it)
|
||||
} else {
|
||||
p.p12.Remove(it)
|
||||
}
|
||||
maxsize = pkt.size()
|
||||
copy(data, pkt.ptr()[:maxsize])
|
||||
return int32(maxsize)
|
||||
}
|
||||
|
||||
//=====================================================================
|
||||
//=====================================================================
|
||||
|
||||
// 模拟网络
|
||||
var vnet *LatencySimulator
|
||||
|
||||
// 测试用例
|
||||
func test(mode int) {
|
||||
// 创建模拟网络：丢包率10%，Rtt 60ms~125ms
|
||||
vnet = &LatencySimulator{}
|
||||
vnet.Init(10, 60, 125, 1000)
|
||||
|
||||
// 创建两个端点的 kcp对象，第一个参数 conv是会话编号，同一个会话需要相同
|
||||
// 最后一个是 user参数，用来传递标识
|
||||
output1 := func(buf []byte, size int) {
|
||||
if vnet.send(0, buf, size) != 1 {
|
||||
}
|
||||
}
|
||||
output2 := func(buf []byte, size int) {
|
||||
if vnet.send(1, buf, size) != 1 {
|
||||
}
|
||||
}
|
||||
kcp1 := NewKCP(0x11223344, output1)
|
||||
kcp2 := NewKCP(0x11223344, output2)
|
||||
|
||||
current := uint32(iclock())
|
||||
slap := current + 20
|
||||
index := 0
|
||||
next := 0
|
||||
var sumrtt uint32
|
||||
count := 0
|
||||
maxrtt := 0
|
||||
|
||||
// 配置窗口大小：平均延迟200ms，每20ms发送一个包，
|
||||
// 而考虑到丢包重发，设置最大收发窗口为128
|
||||
kcp1.WndSize(128, 128)
|
||||
kcp2.WndSize(128, 128)
|
||||
|
||||
// 判断测试用例的模式
|
||||
if mode == 0 {
|
||||
// 默认模式
|
||||
kcp1.NoDelay(0, 10, 0, 0)
|
||||
kcp2.NoDelay(0, 10, 0, 0)
|
||||
} else if mode == 1 {
|
||||
// 普通模式，关闭流控等
|
||||
kcp1.NoDelay(0, 10, 0, 1)
|
||||
kcp2.NoDelay(0, 10, 0, 1)
|
||||
} else {
|
||||
// 启动快速模式
|
||||
// 第二个参数 nodelay-启用以后若干常规加速将启动
|
||||
// 第三个参数 interval为内部处理时钟，默认设置为 10ms
|
||||
// 第四个参数 resend为快速重传指标，设置为2
|
||||
// 第五个参数 为是否禁用常规流控，这里禁止
|
||||
kcp1.NoDelay(1, 10, 2, 1)
|
||||
kcp2.NoDelay(1, 10, 2, 1)
|
||||
}
|
||||
|
||||
buffer := make([]byte, 2000)
|
||||
var hr int32
|
||||
|
||||
ts1 := iclock()
|
||||
|
||||
for {
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
current = uint32(iclock())
|
||||
kcp1.Update()
|
||||
kcp2.Update()
|
||||
|
||||
// 每隔 20ms，kcp1发送数据
|
||||
for ; current >= slap; slap += 20 {
|
||||
buf := new(bytes.Buffer)
|
||||
binary.Write(buf, binary.LittleEndian, uint32(index))
|
||||
index++
|
||||
binary.Write(buf, binary.LittleEndian, uint32(current))
|
||||
// 发送上层协议包
|
||||
kcp1.Send(buf.Bytes())
|
||||
//println("now", iclock())
|
||||
}
|
||||
|
||||
// 处理虚拟网络：检测是否有udp包从p1->p2
|
||||
for {
|
||||
hr = vnet.recv(1, buffer, 2000)
|
||||
if hr < 0 {
|
||||
break
|
||||
}
|
||||
// 如果 p2收到udp，则作为下层协议输入到kcp2
|
||||
kcp2.Input(buffer[:hr], true, false)
|
||||
}
|
||||
|
||||
// 处理虚拟网络：检测是否有udp包从p2->p1
|
||||
for {
|
||||
hr = vnet.recv(0, buffer, 2000)
|
||||
if hr < 0 {
|
||||
break
|
||||
}
|
||||
// 如果 p1收到udp，则作为下层协议输入到kcp1
|
||||
kcp1.Input(buffer[:hr], true, false)
|
||||
//println("@@@@", hr, r)
|
||||
}
|
||||
|
||||
// kcp2接收到任何包都返回回去
|
||||
for {
|
||||
hr = int32(kcp2.Recv(buffer[:10]))
|
||||
// 没有收到包就退出
|
||||
if hr < 0 {
|
||||
break
|
||||
}
|
||||
// 如果收到包就回射
|
||||
buf := bytes.NewReader(buffer)
|
||||
var sn uint32
|
||||
binary.Read(buf, binary.LittleEndian, &sn)
|
||||
kcp2.Send(buffer[:hr])
|
||||
}
|
||||
|
||||
// kcp1收到kcp2的回射数据
|
||||
for {
|
||||
hr = int32(kcp1.Recv(buffer[:10]))
|
||||
buf := bytes.NewReader(buffer)
|
||||
// 没有收到包就退出
|
||||
if hr < 0 {
|
||||
break
|
||||
}
|
||||
var sn uint32
|
||||
var ts, rtt uint32
|
||||
binary.Read(buf, binary.LittleEndian, &sn)
|
||||
binary.Read(buf, binary.LittleEndian, &ts)
|
||||
rtt = uint32(current) - ts
|
||||
|
||||
if sn != uint32(next) {
|
||||
// 如果收到的包不连续
|
||||
//for i:=0;i<8 ;i++ {
|
||||
//println("---", i, buffer[i])
|
||||
//}
|
||||
println("ERROR sn ", count, "<->", next, sn)
|
||||
return
|
||||
}
|
||||
|
||||
next++
|
||||
sumrtt += rtt
|
||||
count++
|
||||
if rtt > uint32(maxrtt) {
|
||||
maxrtt = int(rtt)
|
||||
}
|
||||
|
||||
//println("[RECV] mode=", mode, " sn=", sn, " rtt=", rtt)
|
||||
}
|
||||
|
||||
if next > 100 {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
ts1 = iclock() - ts1
|
||||
|
||||
names := []string{"default", "normal", "fast"}
|
||||
fmt.Printf("%s mode result (%dms):\n", names[mode], ts1)
|
||||
fmt.Printf("avgrtt=%d maxrtt=%d\n", int(sumrtt/uint32(count)), maxrtt)
|
||||
}
|
||||
|
||||
func TestNetwork(t *testing.T) {
|
||||
test(0) // 默认模式，类似 TCP：正常模式，无快速重传，常规流控
|
||||
test(1) // 普通模式，关闭流控等
|
||||
test(2) // 快速模式，所有开关都打开，且关闭流控
|
||||
}
|
||||
|
||||
func BenchmarkFlush(b *testing.B) {
|
||||
kcp := NewKCP(1, func(buf []byte, size int) {})
|
||||
kcp.snd_buf = make([]segment, 1024)
|
||||
for k := range kcp.snd_buf {
|
||||
kcp.snd_buf[k].xmit = 1
|
||||
kcp.snd_buf[k].resendts = currentMs() + 10000
|
||||
}
|
||||
b.ResetTimer()
|
||||
b.ReportAllocs()
|
||||
var mu sync.Mutex
|
||||
for i := 0; i < b.N; i++ {
|
||||
mu.Lock()
|
||||
kcp.flush(false)
|
||||
mu.Unlock()
|
||||
}
|
||||
}
|
||||
963
vender/github.com/xtaci/kcp/sess.go
Normal file
963
vender/github.com/xtaci/kcp/sess.go
Normal file
@@ -0,0 +1,963 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/binary"
|
||||
"hash/crc32"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
"golang.org/x/net/ipv4"
|
||||
"golang.org/x/net/ipv6"
|
||||
)
|
||||
|
||||
type errTimeout struct {
|
||||
error
|
||||
}
|
||||
|
||||
func (errTimeout) Timeout() bool { return true }
|
||||
func (errTimeout) Temporary() bool { return true }
|
||||
func (errTimeout) Error() string { return "i/o timeout" }
|
||||
|
||||
const (
|
||||
// 16-bytes nonce for each packet
|
||||
nonceSize = 16
|
||||
|
||||
// 4-bytes packet checksum
|
||||
crcSize = 4
|
||||
|
||||
// overall crypto header size
|
||||
cryptHeaderSize = nonceSize + crcSize
|
||||
|
||||
// maximum packet size
|
||||
mtuLimit = 1500
|
||||
|
||||
// FEC keeps rxFECMulti* (dataShard+parityShard) ordered packets in memory
|
||||
rxFECMulti = 3
|
||||
|
||||
// accept backlog
|
||||
acceptBacklog = 128
|
||||
)
|
||||
|
||||
const (
|
||||
errBrokenPipe = "broken pipe"
|
||||
errInvalidOperation = "invalid operation"
|
||||
)
|
||||
|
||||
var (
|
||||
// a system-wide packet buffer shared among sending, receiving and FEC
|
||||
// to mitigate high-frequency memory allocation for packets
|
||||
xmitBuf sync.Pool
|
||||
)
|
||||
|
||||
func init() {
|
||||
xmitBuf.New = func() interface{} {
|
||||
return make([]byte, mtuLimit)
|
||||
}
|
||||
}
|
||||
|
||||
type (
|
||||
// UDPSession defines a KCP session implemented by UDP
|
||||
UDPSession struct {
|
||||
updaterIdx int // record slice index in updater
|
||||
conn net.PacketConn // the underlying packet connection
|
||||
kcp *KCP // KCP ARQ protocol
|
||||
l *Listener // pointing to the Listener object if it's been accepted by a Listener
|
||||
block BlockCrypt // block encryption object
|
||||
|
||||
// kcp receiving is based on packets
|
||||
// recvbuf turns packets into stream
|
||||
recvbuf []byte
|
||||
bufptr []byte
|
||||
// header extended output buffer, if has header
|
||||
ext []byte
|
||||
|
||||
// FEC codec
|
||||
fecDecoder *fecDecoder
|
||||
fecEncoder *fecEncoder
|
||||
|
||||
// settings
|
||||
remote net.Addr // remote peer address
|
||||
rd time.Time // read deadline
|
||||
wd time.Time // write deadline
|
||||
headerSize int // the header size additional to a KCP frame
|
||||
ackNoDelay bool // send ack immediately for each incoming packet(testing purpose)
|
||||
writeDelay bool // delay kcp.flush() for Write() for bulk transfer
|
||||
dup int // duplicate udp packets(testing purpose)
|
||||
|
||||
// notifications
|
||||
die chan struct{} // notify current session has Closed
|
||||
chReadEvent chan struct{} // notify Read() can be called without blocking
|
||||
chWriteEvent chan struct{} // notify Write() can be called without blocking
|
||||
chReadError chan error // notify PacketConn.Read() have an error
|
||||
chWriteError chan error // notify PacketConn.Write() have an error
|
||||
|
||||
// nonce generator
|
||||
nonce Entropy
|
||||
|
||||
isClosed bool // flag the session has Closed
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
setReadBuffer interface {
|
||||
SetReadBuffer(bytes int) error
|
||||
}
|
||||
|
||||
setWriteBuffer interface {
|
||||
SetWriteBuffer(bytes int) error
|
||||
}
|
||||
)
|
||||
|
||||
// newUDPSession create a new udp session for client or server
|
||||
func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn net.PacketConn, remote net.Addr, block BlockCrypt) *UDPSession {
|
||||
sess := new(UDPSession)
|
||||
sess.die = make(chan struct{})
|
||||
sess.nonce = new(nonceAES128)
|
||||
sess.nonce.Init()
|
||||
sess.chReadEvent = make(chan struct{}, 1)
|
||||
sess.chWriteEvent = make(chan struct{}, 1)
|
||||
sess.chReadError = make(chan error, 1)
|
||||
sess.chWriteError = make(chan error, 1)
|
||||
sess.remote = remote
|
||||
sess.conn = conn
|
||||
sess.l = l
|
||||
sess.block = block
|
||||
sess.recvbuf = make([]byte, mtuLimit)
|
||||
|
||||
// FEC codec initialization
|
||||
sess.fecDecoder = newFECDecoder(rxFECMulti*(dataShards+parityShards), dataShards, parityShards)
|
||||
if sess.block != nil {
|
||||
sess.fecEncoder = newFECEncoder(dataShards, parityShards, cryptHeaderSize)
|
||||
} else {
|
||||
sess.fecEncoder = newFECEncoder(dataShards, parityShards, 0)
|
||||
}
|
||||
|
||||
// calculate additional header size introduced by FEC and encryption
|
||||
if sess.block != nil {
|
||||
sess.headerSize += cryptHeaderSize
|
||||
}
|
||||
if sess.fecEncoder != nil {
|
||||
sess.headerSize += fecHeaderSizePlus2
|
||||
}
|
||||
|
||||
// we only need to allocate extended packet buffer if we have the additional header
|
||||
if sess.headerSize > 0 {
|
||||
sess.ext = make([]byte, mtuLimit)
|
||||
}
|
||||
|
||||
sess.kcp = NewKCP(conv, func(buf []byte, size int) {
|
||||
if size >= IKCP_OVERHEAD {
|
||||
sess.output(buf[:size])
|
||||
}
|
||||
})
|
||||
sess.kcp.SetMtu(IKCP_MTU_DEF - sess.headerSize)
|
||||
|
||||
// register current session to the global updater,
|
||||
// which call sess.update() periodically.
|
||||
updater.addSession(sess)
|
||||
|
||||
if sess.l == nil { // it's a client connection
|
||||
go sess.readLoop()
|
||||
atomic.AddUint64(&DefaultSnmp.ActiveOpens, 1)
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.PassiveOpens, 1)
|
||||
}
|
||||
currestab := atomic.AddUint64(&DefaultSnmp.CurrEstab, 1)
|
||||
maxconn := atomic.LoadUint64(&DefaultSnmp.MaxConn)
|
||||
if currestab > maxconn {
|
||||
atomic.CompareAndSwapUint64(&DefaultSnmp.MaxConn, maxconn, currestab)
|
||||
}
|
||||
|
||||
return sess
|
||||
}
|
||||
|
||||
// Read implements net.Conn
|
||||
func (s *UDPSession) Read(b []byte) (n int, err error) {
|
||||
for {
|
||||
s.mu.Lock()
|
||||
if len(s.bufptr) > 0 { // copy from buffer into b
|
||||
n = copy(b, s.bufptr)
|
||||
s.bufptr = s.bufptr[n:]
|
||||
s.mu.Unlock()
|
||||
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(n))
|
||||
return n, nil
|
||||
}
|
||||
|
||||
if s.isClosed {
|
||||
s.mu.Unlock()
|
||||
return 0, errors.New(errBrokenPipe)
|
||||
}
|
||||
|
||||
if size := s.kcp.PeekSize(); size > 0 { // peek data size from kcp
|
||||
if len(b) >= size { // receive data into 'b' directly
|
||||
s.kcp.Recv(b)
|
||||
s.mu.Unlock()
|
||||
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(size))
|
||||
return size, nil
|
||||
}
|
||||
|
||||
// if necessary resize the stream buffer to guarantee a sufficent buffer space
|
||||
if cap(s.recvbuf) < size {
|
||||
s.recvbuf = make([]byte, size)
|
||||
}
|
||||
|
||||
// resize the length of recvbuf to correspond to data size
|
||||
s.recvbuf = s.recvbuf[:size]
|
||||
s.kcp.Recv(s.recvbuf)
|
||||
n = copy(b, s.recvbuf) // copy to 'b'
|
||||
s.bufptr = s.recvbuf[n:] // pointer update
|
||||
s.mu.Unlock()
|
||||
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(n))
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// deadline for current reading operation
|
||||
var timeout *time.Timer
|
||||
var c <-chan time.Time
|
||||
if !s.rd.IsZero() {
|
||||
if time.Now().After(s.rd) {
|
||||
s.mu.Unlock()
|
||||
return 0, errTimeout{}
|
||||
}
|
||||
|
||||
delay := s.rd.Sub(time.Now())
|
||||
timeout = time.NewTimer(delay)
|
||||
c = timeout.C
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// wait for read event or timeout
|
||||
select {
|
||||
case <-s.chReadEvent:
|
||||
case <-c:
|
||||
case <-s.die:
|
||||
case err = <-s.chReadError:
|
||||
if timeout != nil {
|
||||
timeout.Stop()
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
if timeout != nil {
|
||||
timeout.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Write implements net.Conn
|
||||
func (s *UDPSession) Write(b []byte) (n int, err error) {
|
||||
for {
|
||||
s.mu.Lock()
|
||||
if s.isClosed {
|
||||
s.mu.Unlock()
|
||||
return 0, errors.New(errBrokenPipe)
|
||||
}
|
||||
|
||||
// controls how much data will be sent to kcp core
|
||||
// to prevent the memory from exhuasting
|
||||
if s.kcp.WaitSnd() < int(s.kcp.snd_wnd) {
|
||||
n = len(b)
|
||||
for {
|
||||
if len(b) <= int(s.kcp.mss) {
|
||||
s.kcp.Send(b)
|
||||
break
|
||||
} else {
|
||||
s.kcp.Send(b[:s.kcp.mss])
|
||||
b = b[s.kcp.mss:]
|
||||
}
|
||||
}
|
||||
|
||||
// flush immediately if the queue is full
|
||||
if s.kcp.WaitSnd() >= int(s.kcp.snd_wnd) || !s.writeDelay {
|
||||
s.kcp.flush(false)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
atomic.AddUint64(&DefaultSnmp.BytesSent, uint64(n))
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// deadline for current writing operation
|
||||
var timeout *time.Timer
|
||||
var c <-chan time.Time
|
||||
if !s.wd.IsZero() {
|
||||
if time.Now().After(s.wd) {
|
||||
s.mu.Unlock()
|
||||
return 0, errTimeout{}
|
||||
}
|
||||
delay := s.wd.Sub(time.Now())
|
||||
timeout = time.NewTimer(delay)
|
||||
c = timeout.C
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// wait for write event or timeout
|
||||
select {
|
||||
case <-s.chWriteEvent:
|
||||
case <-c:
|
||||
case <-s.die:
|
||||
case err = <-s.chWriteError:
|
||||
if timeout != nil {
|
||||
timeout.Stop()
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
if timeout != nil {
|
||||
timeout.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close closes the connection.
|
||||
func (s *UDPSession) Close() error {
|
||||
// remove current session from updater & listener(if necessary)
|
||||
updater.removeSession(s)
|
||||
if s.l != nil { // notify listener
|
||||
s.l.closeSession(s.remote)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.isClosed {
|
||||
return errors.New(errBrokenPipe)
|
||||
}
|
||||
close(s.die)
|
||||
s.isClosed = true
|
||||
atomic.AddUint64(&DefaultSnmp.CurrEstab, ^uint64(0))
|
||||
if s.l == nil { // client socket close
|
||||
return s.conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// LocalAddr returns the local network address. The Addr returned is shared by all invocations of LocalAddr, so do not modify it.
|
||||
func (s *UDPSession) LocalAddr() net.Addr { return s.conn.LocalAddr() }
|
||||
|
||||
// RemoteAddr returns the remote network address. The Addr returned is shared by all invocations of RemoteAddr, so do not modify it.
|
||||
func (s *UDPSession) RemoteAddr() net.Addr { return s.remote }
|
||||
|
||||
// SetDeadline sets the deadline associated with the listener. A zero time value disables the deadline.
|
||||
func (s *UDPSession) SetDeadline(t time.Time) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.rd = t
|
||||
s.wd = t
|
||||
s.notifyReadEvent()
|
||||
s.notifyWriteEvent()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetReadDeadline implements the Conn SetReadDeadline method.
|
||||
func (s *UDPSession) SetReadDeadline(t time.Time) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.rd = t
|
||||
s.notifyReadEvent()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetWriteDeadline implements the Conn SetWriteDeadline method.
|
||||
func (s *UDPSession) SetWriteDeadline(t time.Time) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.wd = t
|
||||
s.notifyWriteEvent()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetWriteDelay delays write for bulk transfer until the next update interval
|
||||
func (s *UDPSession) SetWriteDelay(delay bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.writeDelay = delay
|
||||
}
|
||||
|
||||
// SetWindowSize set maximum window size
|
||||
func (s *UDPSession) SetWindowSize(sndwnd, rcvwnd int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.kcp.WndSize(sndwnd, rcvwnd)
|
||||
}
|
||||
|
||||
// SetMtu sets the maximum transmission unit(not including UDP header)
|
||||
func (s *UDPSession) SetMtu(mtu int) bool {
|
||||
if mtu > mtuLimit {
|
||||
return false
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.kcp.SetMtu(mtu - s.headerSize)
|
||||
return true
|
||||
}
|
||||
|
||||
// SetStreamMode toggles the stream mode on/off
|
||||
func (s *UDPSession) SetStreamMode(enable bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if enable {
|
||||
s.kcp.stream = 1
|
||||
} else {
|
||||
s.kcp.stream = 0
|
||||
}
|
||||
}
|
||||
|
||||
// SetACKNoDelay changes ack flush option, set true to flush ack immediately,
|
||||
func (s *UDPSession) SetACKNoDelay(nodelay bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.ackNoDelay = nodelay
|
||||
}
|
||||
|
||||
// SetDUP duplicates udp packets for kcp output, for testing purpose only
|
||||
func (s *UDPSession) SetDUP(dup int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.dup = dup
|
||||
}
|
||||
|
||||
// SetNoDelay calls nodelay() of kcp
|
||||
// https://github.com/skywind3000/kcp/blob/master/README.en.md#protocol-configuration
|
||||
func (s *UDPSession) SetNoDelay(nodelay, interval, resend, nc int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.kcp.NoDelay(nodelay, interval, resend, nc)
|
||||
}
|
||||
|
||||
// SetDSCP sets the 6bit DSCP field of IP header, no effect if it's accepted from Listener
|
||||
func (s *UDPSession) SetDSCP(dscp int) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.l == nil {
|
||||
if nc, ok := s.conn.(net.Conn); ok {
|
||||
if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err != nil {
|
||||
return ipv6.NewConn(nc).SetTrafficClass(dscp)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// SetReadBuffer sets the socket read buffer, no effect if it's accepted from Listener
|
||||
func (s *UDPSession) SetReadBuffer(bytes int) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.l == nil {
|
||||
if nc, ok := s.conn.(setReadBuffer); ok {
|
||||
return nc.SetReadBuffer(bytes)
|
||||
}
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// SetWriteBuffer sets the socket write buffer, no effect if it's accepted from Listener
|
||||
func (s *UDPSession) SetWriteBuffer(bytes int) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.l == nil {
|
||||
if nc, ok := s.conn.(setWriteBuffer); ok {
|
||||
return nc.SetWriteBuffer(bytes)
|
||||
}
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// post-processing for sending a packet from kcp core
|
||||
// steps:
|
||||
// 0. Header extending
|
||||
// 1. FEC packet generation
|
||||
// 2. CRC32 integrity
|
||||
// 3. Encryption
|
||||
// 4. WriteTo kernel
|
||||
func (s *UDPSession) output(buf []byte) {
|
||||
var ecc [][]byte
|
||||
|
||||
// 0. extend buf's header space(if necessary)
|
||||
ext := buf
|
||||
if s.headerSize > 0 {
|
||||
ext = s.ext[:s.headerSize+len(buf)]
|
||||
copy(ext[s.headerSize:], buf)
|
||||
}
|
||||
|
||||
// 1. FEC encoding
|
||||
if s.fecEncoder != nil {
|
||||
ecc = s.fecEncoder.encode(ext)
|
||||
}
|
||||
|
||||
// 2&3. crc32 & encryption
|
||||
if s.block != nil {
|
||||
s.nonce.Fill(ext[:nonceSize])
|
||||
checksum := crc32.ChecksumIEEE(ext[cryptHeaderSize:])
|
||||
binary.LittleEndian.PutUint32(ext[nonceSize:], checksum)
|
||||
s.block.Encrypt(ext, ext)
|
||||
|
||||
for k := range ecc {
|
||||
s.nonce.Fill(ecc[k][:nonceSize])
|
||||
checksum := crc32.ChecksumIEEE(ecc[k][cryptHeaderSize:])
|
||||
binary.LittleEndian.PutUint32(ecc[k][nonceSize:], checksum)
|
||||
s.block.Encrypt(ecc[k], ecc[k])
|
||||
}
|
||||
}
|
||||
|
||||
// 4. WriteTo kernel
|
||||
nbytes := 0
|
||||
npkts := 0
|
||||
for i := 0; i < s.dup+1; i++ {
|
||||
if n, err := s.conn.WriteTo(ext, s.remote); err == nil {
|
||||
nbytes += n
|
||||
npkts++
|
||||
} else {
|
||||
s.notifyWriteError(err)
|
||||
}
|
||||
}
|
||||
|
||||
for k := range ecc {
|
||||
if n, err := s.conn.WriteTo(ecc[k], s.remote); err == nil {
|
||||
nbytes += n
|
||||
npkts++
|
||||
} else {
|
||||
s.notifyWriteError(err)
|
||||
}
|
||||
}
|
||||
atomic.AddUint64(&DefaultSnmp.OutPkts, uint64(npkts))
|
||||
atomic.AddUint64(&DefaultSnmp.OutBytes, uint64(nbytes))
|
||||
}
|
||||
|
||||
// kcp update, returns interval for next calling
|
||||
func (s *UDPSession) update() (interval time.Duration) {
|
||||
s.mu.Lock()
|
||||
waitsnd := s.kcp.WaitSnd()
|
||||
interval = time.Duration(s.kcp.flush(false)) * time.Millisecond
|
||||
if s.kcp.WaitSnd() < waitsnd {
|
||||
s.notifyWriteEvent()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
// GetConv gets conversation id of a session
|
||||
func (s *UDPSession) GetConv() uint32 { return s.kcp.conv }
|
||||
|
||||
func (s *UDPSession) notifyReadEvent() {
|
||||
select {
|
||||
case s.chReadEvent <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UDPSession) notifyWriteEvent() {
|
||||
select {
|
||||
case s.chWriteEvent <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UDPSession) notifyWriteError(err error) {
|
||||
select {
|
||||
case s.chWriteError <- err:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UDPSession) kcpInput(data []byte) {
|
||||
var kcpInErrors, fecErrs, fecRecovered, fecParityShards uint64
|
||||
|
||||
if s.fecDecoder != nil {
|
||||
if len(data) > fecHeaderSize { // must be larger than fec header size
|
||||
f := s.fecDecoder.decodeBytes(data)
|
||||
if f.flag == typeData || f.flag == typeFEC { // header check
|
||||
if f.flag == typeFEC {
|
||||
fecParityShards++
|
||||
}
|
||||
recovers := s.fecDecoder.decode(f)
|
||||
|
||||
s.mu.Lock()
|
||||
waitsnd := s.kcp.WaitSnd()
|
||||
if f.flag == typeData {
|
||||
if ret := s.kcp.Input(data[fecHeaderSizePlus2:], true, s.ackNoDelay); ret != 0 {
|
||||
kcpInErrors++
|
||||
}
|
||||
}
|
||||
|
||||
for _, r := range recovers {
|
||||
if len(r) >= 2 { // must be larger than 2bytes
|
||||
sz := binary.LittleEndian.Uint16(r)
|
||||
if int(sz) <= len(r) && sz >= 2 {
|
||||
if ret := s.kcp.Input(r[2:sz], false, s.ackNoDelay); ret == 0 {
|
||||
fecRecovered++
|
||||
} else {
|
||||
kcpInErrors++
|
||||
}
|
||||
} else {
|
||||
fecErrs++
|
||||
}
|
||||
} else {
|
||||
fecErrs++
|
||||
}
|
||||
}
|
||||
|
||||
// to notify the readers to receive the data
|
||||
if n := s.kcp.PeekSize(); n > 0 {
|
||||
s.notifyReadEvent()
|
||||
}
|
||||
// to notify the writers when queue is shorter(e.g. ACKed)
|
||||
if s.kcp.WaitSnd() < waitsnd {
|
||||
s.notifyWriteEvent()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
|
||||
}
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
|
||||
}
|
||||
} else {
|
||||
s.mu.Lock()
|
||||
waitsnd := s.kcp.WaitSnd()
|
||||
if ret := s.kcp.Input(data, true, s.ackNoDelay); ret != 0 {
|
||||
kcpInErrors++
|
||||
}
|
||||
if n := s.kcp.PeekSize(); n > 0 {
|
||||
s.notifyReadEvent()
|
||||
}
|
||||
if s.kcp.WaitSnd() < waitsnd {
|
||||
s.notifyWriteEvent()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
atomic.AddUint64(&DefaultSnmp.InPkts, 1)
|
||||
atomic.AddUint64(&DefaultSnmp.InBytes, uint64(len(data)))
|
||||
if fecParityShards > 0 {
|
||||
atomic.AddUint64(&DefaultSnmp.FECParityShards, fecParityShards)
|
||||
}
|
||||
if kcpInErrors > 0 {
|
||||
atomic.AddUint64(&DefaultSnmp.KCPInErrors, kcpInErrors)
|
||||
}
|
||||
if fecErrs > 0 {
|
||||
atomic.AddUint64(&DefaultSnmp.FECErrs, fecErrs)
|
||||
}
|
||||
if fecRecovered > 0 {
|
||||
atomic.AddUint64(&DefaultSnmp.FECRecovered, fecRecovered)
|
||||
}
|
||||
}
|
||||
|
||||
// the read loop for a client session
|
||||
func (s *UDPSession) readLoop() {
|
||||
buf := make([]byte, mtuLimit)
|
||||
var src string
|
||||
for {
|
||||
if n, addr, err := s.conn.ReadFrom(buf); err == nil {
|
||||
// make sure the packet is from the same source
|
||||
if src == "" { // set source address
|
||||
src = addr.String()
|
||||
} else if addr.String() != src {
|
||||
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
|
||||
continue
|
||||
}
|
||||
|
||||
if n >= s.headerSize+IKCP_OVERHEAD {
|
||||
data := buf[:n]
|
||||
dataValid := false
|
||||
if s.block != nil {
|
||||
s.block.Decrypt(data, data)
|
||||
data = data[nonceSize:]
|
||||
checksum := crc32.ChecksumIEEE(data[crcSize:])
|
||||
if checksum == binary.LittleEndian.Uint32(data) {
|
||||
data = data[crcSize:]
|
||||
dataValid = true
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
|
||||
}
|
||||
} else if s.block == nil {
|
||||
dataValid = true
|
||||
}
|
||||
|
||||
if dataValid {
|
||||
s.kcpInput(data)
|
||||
}
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
|
||||
}
|
||||
} else {
|
||||
s.chReadError <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type (
|
||||
// Listener defines a server which will be waiting to accept incoming connections
|
||||
Listener struct {
|
||||
block BlockCrypt // block encryption
|
||||
dataShards int // FEC data shard
|
||||
parityShards int // FEC parity shard
|
||||
fecDecoder *fecDecoder // FEC mock initialization
|
||||
conn net.PacketConn // the underlying packet connection
|
||||
|
||||
sessions map[string]*UDPSession // all sessions accepted by this Listener
|
||||
sessionLock sync.Mutex
|
||||
chAccepts chan *UDPSession // Listen() backlog
|
||||
chSessionClosed chan net.Addr // session close queue
|
||||
headerSize int // the additional header to a KCP frame
|
||||
die chan struct{} // notify the listener has closed
|
||||
rd atomic.Value // read deadline for Accept()
|
||||
wd atomic.Value
|
||||
}
|
||||
)
|
||||
|
||||
// monitor incoming data for all connections of server
|
||||
func (l *Listener) monitor() {
|
||||
// a cache for session object last used
|
||||
var lastAddr string
|
||||
var lastSession *UDPSession
|
||||
buf := make([]byte, mtuLimit)
|
||||
for {
|
||||
if n, from, err := l.conn.ReadFrom(buf); err == nil {
|
||||
if n >= l.headerSize+IKCP_OVERHEAD {
|
||||
data := buf[:n]
|
||||
dataValid := false
|
||||
if l.block != nil {
|
||||
l.block.Decrypt(data, data)
|
||||
data = data[nonceSize:]
|
||||
checksum := crc32.ChecksumIEEE(data[crcSize:])
|
||||
if checksum == binary.LittleEndian.Uint32(data) {
|
||||
data = data[crcSize:]
|
||||
dataValid = true
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
|
||||
}
|
||||
} else if l.block == nil {
|
||||
dataValid = true
|
||||
}
|
||||
|
||||
if dataValid {
|
||||
addr := from.String()
|
||||
var s *UDPSession
|
||||
var ok bool
|
||||
|
||||
// the packets received from an address always come in batch,
|
||||
// cache the session for next packet, without querying map.
|
||||
if addr == lastAddr {
|
||||
s, ok = lastSession, true
|
||||
} else {
|
||||
l.sessionLock.Lock()
|
||||
if s, ok = l.sessions[addr]; ok {
|
||||
lastSession = s
|
||||
lastAddr = addr
|
||||
}
|
||||
l.sessionLock.Unlock()
|
||||
}
|
||||
|
||||
if !ok { // new session
|
||||
if len(l.chAccepts) < cap(l.chAccepts) { // do not let the new sessions overwhelm accept queue
|
||||
var conv uint32
|
||||
convValid := false
|
||||
if l.fecDecoder != nil {
|
||||
isfec := binary.LittleEndian.Uint16(data[4:])
|
||||
if isfec == typeData {
|
||||
conv = binary.LittleEndian.Uint32(data[fecHeaderSizePlus2:])
|
||||
convValid = true
|
||||
}
|
||||
} else {
|
||||
conv = binary.LittleEndian.Uint32(data)
|
||||
convValid = true
|
||||
}
|
||||
|
||||
if convValid { // creates a new session only if the 'conv' field in kcp is accessible
|
||||
s := newUDPSession(conv, l.dataShards, l.parityShards, l, l.conn, from, l.block)
|
||||
s.kcpInput(data)
|
||||
l.sessionLock.Lock()
|
||||
l.sessions[addr] = s
|
||||
l.sessionLock.Unlock()
|
||||
l.chAccepts <- s
|
||||
}
|
||||
}
|
||||
} else {
|
||||
s.kcpInput(data)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
|
||||
}
|
||||
} else {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SetReadBuffer sets the socket read buffer for the Listener
|
||||
func (l *Listener) SetReadBuffer(bytes int) error {
|
||||
if nc, ok := l.conn.(setReadBuffer); ok {
|
||||
return nc.SetReadBuffer(bytes)
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// SetWriteBuffer sets the socket write buffer for the Listener
|
||||
func (l *Listener) SetWriteBuffer(bytes int) error {
|
||||
if nc, ok := l.conn.(setWriteBuffer); ok {
|
||||
return nc.SetWriteBuffer(bytes)
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// SetDSCP sets the 6bit DSCP field of IP header
|
||||
func (l *Listener) SetDSCP(dscp int) error {
|
||||
if nc, ok := l.conn.(net.Conn); ok {
|
||||
if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err != nil {
|
||||
return ipv6.NewConn(nc).SetTrafficClass(dscp)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return errors.New(errInvalidOperation)
|
||||
}
|
||||
|
||||
// Accept implements the Accept method in the Listener interface; it waits for the next call and returns a generic Conn.
|
||||
func (l *Listener) Accept() (net.Conn, error) {
|
||||
return l.AcceptKCP()
|
||||
}
|
||||
|
||||
// AcceptKCP accepts a KCP connection
|
||||
func (l *Listener) AcceptKCP() (*UDPSession, error) {
|
||||
var timeout <-chan time.Time
|
||||
if tdeadline, ok := l.rd.Load().(time.Time); ok && !tdeadline.IsZero() {
|
||||
timeout = time.After(tdeadline.Sub(time.Now()))
|
||||
}
|
||||
|
||||
select {
|
||||
case <-timeout:
|
||||
return nil, &errTimeout{}
|
||||
case c := <-l.chAccepts:
|
||||
return c, nil
|
||||
case <-l.die:
|
||||
return nil, errors.New(errBrokenPipe)
|
||||
}
|
||||
}
|
||||
|
||||
// SetDeadline sets the deadline associated with the listener. A zero time value disables the deadline.
|
||||
func (l *Listener) SetDeadline(t time.Time) error {
|
||||
l.SetReadDeadline(t)
|
||||
l.SetWriteDeadline(t)
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetReadDeadline implements the Conn SetReadDeadline method.
|
||||
func (l *Listener) SetReadDeadline(t time.Time) error {
|
||||
l.rd.Store(t)
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetWriteDeadline implements the Conn SetWriteDeadline method.
|
||||
func (l *Listener) SetWriteDeadline(t time.Time) error {
|
||||
l.wd.Store(t)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close stops listening on the UDP address. Already Accepted connections are not closed.
|
||||
func (l *Listener) Close() error {
|
||||
close(l.die)
|
||||
return l.conn.Close()
|
||||
}
|
||||
|
||||
// closeSession notify the listener that a session has closed
|
||||
func (l *Listener) closeSession(remote net.Addr) (ret bool) {
|
||||
l.sessionLock.Lock()
|
||||
defer l.sessionLock.Unlock()
|
||||
if _, ok := l.sessions[remote.String()]; ok {
|
||||
delete(l.sessions, remote.String())
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Addr returns the listener's network address, The Addr returned is shared by all invocations of Addr, so do not modify it.
|
||||
func (l *Listener) Addr() net.Addr { return l.conn.LocalAddr() }
|
||||
|
||||
// Listen listens for incoming KCP packets addressed to the local address laddr on the network "udp",
|
||||
func Listen(laddr string) (net.Listener, error) { return ListenWithOptions(laddr, nil, 0, 0) }
|
||||
|
||||
// ListenWithOptions listens for incoming KCP packets addressed to the local address laddr on the network "udp" with packet encryption,
|
||||
// rdataShards, parityShards defines Reed-Solomon Erasure Coding parametes
|
||||
func ListenWithOptions(laddr string, block BlockCrypt, dataShards, parityShards int) (*Listener, error) {
|
||||
udpaddr, err := net.ResolveUDPAddr("udp", laddr)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net.ResolveUDPAddr")
|
||||
}
|
||||
conn, err := net.ListenUDP("udp", udpaddr)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net.ListenUDP")
|
||||
}
|
||||
|
||||
return ServeConn(block, dataShards, parityShards, conn)
|
||||
}
|
||||
|
||||
// ServeConn serves KCP protocol for a single packet connection.
|
||||
func ServeConn(block BlockCrypt, dataShards, parityShards int, conn net.PacketConn) (*Listener, error) {
|
||||
l := new(Listener)
|
||||
l.conn = conn
|
||||
l.sessions = make(map[string]*UDPSession)
|
||||
l.chAccepts = make(chan *UDPSession, acceptBacklog)
|
||||
l.chSessionClosed = make(chan net.Addr)
|
||||
l.die = make(chan struct{})
|
||||
l.dataShards = dataShards
|
||||
l.parityShards = parityShards
|
||||
l.block = block
|
||||
l.fecDecoder = newFECDecoder(rxFECMulti*(dataShards+parityShards), dataShards, parityShards)
|
||||
|
||||
// calculate header size
|
||||
if l.block != nil {
|
||||
l.headerSize += cryptHeaderSize
|
||||
}
|
||||
if l.fecDecoder != nil {
|
||||
l.headerSize += fecHeaderSizePlus2
|
||||
}
|
||||
|
||||
go l.monitor()
|
||||
return l, nil
|
||||
}
|
||||
|
||||
// Dial connects to the remote address "raddr" on the network "udp"
|
||||
func Dial(raddr string) (net.Conn, error) { return DialWithOptions(raddr, nil, 0, 0) }
|
||||
|
||||
// DialWithOptions connects to the remote address "raddr" on the network "udp" with packet encryption
|
||||
func DialWithOptions(raddr string, block BlockCrypt, dataShards, parityShards int) (*UDPSession, error) {
|
||||
// network type detection
|
||||
udpaddr, err := net.ResolveUDPAddr("udp", raddr)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net.ResolveUDPAddr")
|
||||
}
|
||||
network := "udp4"
|
||||
if udpaddr.IP.To4() == nil {
|
||||
network = "udp"
|
||||
}
|
||||
|
||||
conn, err := net.ListenUDP(network, nil)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net.DialUDP")
|
||||
}
|
||||
|
||||
return NewConn(raddr, block, dataShards, parityShards, conn)
|
||||
}
|
||||
|
||||
// NewConn establishes a session and talks KCP protocol over a packet connection.
|
||||
func NewConn(raddr string, block BlockCrypt, dataShards, parityShards int, conn net.PacketConn) (*UDPSession, error) {
|
||||
udpaddr, err := net.ResolveUDPAddr("udp", raddr)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "net.ResolveUDPAddr")
|
||||
}
|
||||
|
||||
var convid uint32
|
||||
binary.Read(rand.Reader, binary.LittleEndian, &convid)
|
||||
return newUDPSession(convid, dataShards, parityShards, nil, conn, udpaddr, block), nil
|
||||
}
|
||||
|
||||
// monotonic reference time point
|
||||
var refTime time.Time = time.Now()
|
||||
|
||||
// currentMs returns current elasped monotonic milliseconds since program startup
|
||||
func currentMs() uint32 { return uint32(time.Now().Sub(refTime) / time.Millisecond) }
|
||||
475
vender/github.com/xtaci/kcp/sess_test.go
Normal file
475
vender/github.com/xtaci/kcp/sess_test.go
Normal file
@@ -0,0 +1,475 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
_ "net/http/pprof"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"golang.org/x/crypto/pbkdf2"
|
||||
)
|
||||
|
||||
const portEcho = "127.0.0.1:9999"
|
||||
const portSink = "127.0.0.1:19999"
|
||||
const portTinyBufferEcho = "127.0.0.1:29999"
|
||||
const portListerner = "127.0.0.1:9998"
|
||||
|
||||
var key = []byte("testkey")
|
||||
var pass = pbkdf2.Key(key, []byte(portSink), 4096, 32, sha1.New)
|
||||
|
||||
func init() {
|
||||
go func() {
|
||||
log.Println(http.ListenAndServe("localhost:6060", nil))
|
||||
}()
|
||||
|
||||
go echoServer()
|
||||
go sinkServer()
|
||||
go tinyBufferEchoServer()
|
||||
println("beginning tests, encryption:salsa20, fec:10/3")
|
||||
}
|
||||
|
||||
func dialEcho() (*UDPSession, error) {
|
||||
//block, _ := NewNoneBlockCrypt(pass)
|
||||
//block, _ := NewSimpleXORBlockCrypt(pass)
|
||||
//block, _ := NewTEABlockCrypt(pass[:16])
|
||||
//block, _ := NewAESBlockCrypt(pass)
|
||||
block, _ := NewSalsa20BlockCrypt(pass)
|
||||
sess, err := DialWithOptions(portEcho, block, 10, 3)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sess.SetStreamMode(true)
|
||||
sess.SetStreamMode(false)
|
||||
sess.SetStreamMode(true)
|
||||
sess.SetWindowSize(1024, 1024)
|
||||
sess.SetReadBuffer(16 * 1024 * 1024)
|
||||
sess.SetWriteBuffer(16 * 1024 * 1024)
|
||||
sess.SetStreamMode(true)
|
||||
sess.SetNoDelay(1, 10, 2, 1)
|
||||
sess.SetMtu(1400)
|
||||
sess.SetMtu(1600)
|
||||
sess.SetMtu(1400)
|
||||
sess.SetACKNoDelay(true)
|
||||
sess.SetACKNoDelay(false)
|
||||
sess.SetDeadline(time.Now().Add(time.Minute))
|
||||
return sess, err
|
||||
}
|
||||
|
||||
func dialSink() (*UDPSession, error) {
|
||||
sess, err := DialWithOptions(portSink, nil, 0, 0)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sess.SetStreamMode(true)
|
||||
sess.SetWindowSize(1024, 1024)
|
||||
sess.SetReadBuffer(16 * 1024 * 1024)
|
||||
sess.SetWriteBuffer(16 * 1024 * 1024)
|
||||
sess.SetStreamMode(true)
|
||||
sess.SetNoDelay(1, 10, 2, 1)
|
||||
sess.SetMtu(1400)
|
||||
sess.SetACKNoDelay(false)
|
||||
sess.SetDeadline(time.Now().Add(time.Minute))
|
||||
return sess, err
|
||||
}
|
||||
|
||||
func dialTinyBufferEcho() (*UDPSession, error) {
|
||||
//block, _ := NewNoneBlockCrypt(pass)
|
||||
//block, _ := NewSimpleXORBlockCrypt(pass)
|
||||
//block, _ := NewTEABlockCrypt(pass[:16])
|
||||
//block, _ := NewAESBlockCrypt(pass)
|
||||
block, _ := NewSalsa20BlockCrypt(pass)
|
||||
sess, err := DialWithOptions(portTinyBufferEcho, block, 10, 3)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return sess, err
|
||||
}
|
||||
|
||||
//////////////////////////
|
||||
func listenEcho() (net.Listener, error) {
|
||||
//block, _ := NewNoneBlockCrypt(pass)
|
||||
//block, _ := NewSimpleXORBlockCrypt(pass)
|
||||
//block, _ := NewTEABlockCrypt(pass[:16])
|
||||
//block, _ := NewAESBlockCrypt(pass)
|
||||
block, _ := NewSalsa20BlockCrypt(pass)
|
||||
return ListenWithOptions(portEcho, block, 10, 3)
|
||||
}
|
||||
func listenTinyBufferEcho() (net.Listener, error) {
|
||||
//block, _ := NewNoneBlockCrypt(pass)
|
||||
//block, _ := NewSimpleXORBlockCrypt(pass)
|
||||
//block, _ := NewTEABlockCrypt(pass[:16])
|
||||
//block, _ := NewAESBlockCrypt(pass)
|
||||
block, _ := NewSalsa20BlockCrypt(pass)
|
||||
return ListenWithOptions(portTinyBufferEcho, block, 10, 3)
|
||||
}
|
||||
|
||||
func listenSink() (net.Listener, error) {
|
||||
return ListenWithOptions(portSink, nil, 0, 0)
|
||||
}
|
||||
|
||||
func echoServer() {
|
||||
l, err := listenEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
kcplistener := l.(*Listener)
|
||||
kcplistener.SetReadBuffer(4 * 1024 * 1024)
|
||||
kcplistener.SetWriteBuffer(4 * 1024 * 1024)
|
||||
kcplistener.SetDSCP(46)
|
||||
for {
|
||||
s, err := l.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// coverage test
|
||||
s.(*UDPSession).SetReadBuffer(4 * 1024 * 1024)
|
||||
s.(*UDPSession).SetWriteBuffer(4 * 1024 * 1024)
|
||||
go handleEcho(s.(*UDPSession))
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func sinkServer() {
|
||||
l, err := listenSink()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
kcplistener := l.(*Listener)
|
||||
kcplistener.SetReadBuffer(4 * 1024 * 1024)
|
||||
kcplistener.SetWriteBuffer(4 * 1024 * 1024)
|
||||
kcplistener.SetDSCP(46)
|
||||
for {
|
||||
s, err := l.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
go handleSink(s.(*UDPSession))
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func tinyBufferEchoServer() {
|
||||
l, err := listenTinyBufferEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
for {
|
||||
s, err := l.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
go handleTinyBufferEcho(s.(*UDPSession))
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
///////////////////////////
|
||||
|
||||
func handleEcho(conn *UDPSession) {
|
||||
conn.SetStreamMode(true)
|
||||
conn.SetWindowSize(4096, 4096)
|
||||
conn.SetNoDelay(1, 10, 2, 1)
|
||||
conn.SetDSCP(46)
|
||||
conn.SetMtu(1400)
|
||||
conn.SetACKNoDelay(false)
|
||||
conn.SetReadDeadline(time.Now().Add(time.Hour))
|
||||
conn.SetWriteDeadline(time.Now().Add(time.Hour))
|
||||
buf := make([]byte, 65536)
|
||||
for {
|
||||
n, err := conn.Read(buf)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
conn.Write(buf[:n])
|
||||
}
|
||||
}
|
||||
|
||||
func handleSink(conn *UDPSession) {
|
||||
conn.SetStreamMode(true)
|
||||
conn.SetWindowSize(4096, 4096)
|
||||
conn.SetNoDelay(1, 10, 2, 1)
|
||||
conn.SetDSCP(46)
|
||||
conn.SetMtu(1400)
|
||||
conn.SetACKNoDelay(false)
|
||||
conn.SetReadDeadline(time.Now().Add(time.Hour))
|
||||
conn.SetWriteDeadline(time.Now().Add(time.Hour))
|
||||
buf := make([]byte, 65536)
|
||||
for {
|
||||
_, err := conn.Read(buf)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func handleTinyBufferEcho(conn *UDPSession) {
|
||||
conn.SetStreamMode(true)
|
||||
buf := make([]byte, 2)
|
||||
for {
|
||||
n, err := conn.Read(buf)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
conn.Write(buf[:n])
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////
|
||||
|
||||
func TestTimeout(t *testing.T) {
|
||||
cli, err := dialEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
buf := make([]byte, 10)
|
||||
|
||||
//timeout
|
||||
cli.SetDeadline(time.Now().Add(time.Second))
|
||||
<-time.After(2 * time.Second)
|
||||
n, err := cli.Read(buf)
|
||||
if n != 0 || err == nil {
|
||||
t.Fail()
|
||||
}
|
||||
cli.Close()
|
||||
}
|
||||
|
||||
func TestSendRecv(t *testing.T) {
|
||||
cli, err := dialEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
cli.SetWriteDelay(true)
|
||||
cli.SetDUP(1)
|
||||
const N = 100
|
||||
buf := make([]byte, 10)
|
||||
for i := 0; i < N; i++ {
|
||||
msg := fmt.Sprintf("hello%v", i)
|
||||
cli.Write([]byte(msg))
|
||||
if n, err := cli.Read(buf); err == nil {
|
||||
if string(buf[:n]) != msg {
|
||||
t.Fail()
|
||||
}
|
||||
} else {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
cli.Close()
|
||||
}
|
||||
|
||||
func TestTinyBufferReceiver(t *testing.T) {
|
||||
cli, err := dialTinyBufferEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
const N = 100
|
||||
snd := byte(0)
|
||||
fillBuffer := func(buf []byte) {
|
||||
for i := 0; i < len(buf); i++ {
|
||||
buf[i] = snd
|
||||
snd++
|
||||
}
|
||||
}
|
||||
|
||||
rcv := byte(0)
|
||||
check := func(buf []byte) bool {
|
||||
for i := 0; i < len(buf); i++ {
|
||||
if buf[i] != rcv {
|
||||
return false
|
||||
}
|
||||
rcv++
|
||||
}
|
||||
return true
|
||||
}
|
||||
sndbuf := make([]byte, 7)
|
||||
rcvbuf := make([]byte, 7)
|
||||
for i := 0; i < N; i++ {
|
||||
fillBuffer(sndbuf)
|
||||
cli.Write(sndbuf)
|
||||
if n, err := io.ReadFull(cli, rcvbuf); err == nil {
|
||||
if !check(rcvbuf[:n]) {
|
||||
t.Fail()
|
||||
}
|
||||
} else {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
cli.Close()
|
||||
}
|
||||
|
||||
func TestClose(t *testing.T) {
|
||||
cli, err := dialEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
buf := make([]byte, 10)
|
||||
|
||||
cli.Close()
|
||||
if cli.Close() == nil {
|
||||
t.Fail()
|
||||
}
|
||||
n, err := cli.Write(buf)
|
||||
if n != 0 || err == nil {
|
||||
t.Fail()
|
||||
}
|
||||
n, err = cli.Read(buf)
|
||||
if n != 0 || err == nil {
|
||||
t.Fail()
|
||||
}
|
||||
cli.Close()
|
||||
}
|
||||
|
||||
func TestParallel1024CLIENT_64BMSG_64CNT(t *testing.T) {
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1024)
|
||||
for i := 0; i < 1024; i++ {
|
||||
go parallel_client(&wg)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func parallel_client(wg *sync.WaitGroup) (err error) {
|
||||
cli, err := dialEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
err = echo_tester(cli, 64, 64)
|
||||
wg.Done()
|
||||
return
|
||||
}
|
||||
|
||||
func BenchmarkEchoSpeed4K(b *testing.B) {
|
||||
speedclient(b, 4096)
|
||||
}
|
||||
|
||||
func BenchmarkEchoSpeed64K(b *testing.B) {
|
||||
speedclient(b, 65536)
|
||||
}
|
||||
|
||||
func BenchmarkEchoSpeed512K(b *testing.B) {
|
||||
speedclient(b, 524288)
|
||||
}
|
||||
|
||||
func BenchmarkEchoSpeed1M(b *testing.B) {
|
||||
speedclient(b, 1048576)
|
||||
}
|
||||
|
||||
func speedclient(b *testing.B, nbytes int) {
|
||||
b.ReportAllocs()
|
||||
cli, err := dialEcho()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if err := echo_tester(cli, nbytes, b.N); err != nil {
|
||||
b.Fail()
|
||||
}
|
||||
b.SetBytes(int64(nbytes))
|
||||
}
|
||||
|
||||
func BenchmarkSinkSpeed4K(b *testing.B) {
|
||||
sinkclient(b, 4096)
|
||||
}
|
||||
|
||||
func BenchmarkSinkSpeed64K(b *testing.B) {
|
||||
sinkclient(b, 65536)
|
||||
}
|
||||
|
||||
func BenchmarkSinkSpeed256K(b *testing.B) {
|
||||
sinkclient(b, 524288)
|
||||
}
|
||||
|
||||
func BenchmarkSinkSpeed1M(b *testing.B) {
|
||||
sinkclient(b, 1048576)
|
||||
}
|
||||
|
||||
func sinkclient(b *testing.B, nbytes int) {
|
||||
b.ReportAllocs()
|
||||
cli, err := dialSink()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
sink_tester(cli, nbytes, b.N)
|
||||
b.SetBytes(int64(nbytes))
|
||||
}
|
||||
|
||||
func echo_tester(cli net.Conn, msglen, msgcount int) error {
|
||||
buf := make([]byte, msglen)
|
||||
for i := 0; i < msgcount; i++ {
|
||||
// send packet
|
||||
if _, err := cli.Write(buf); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// receive packet
|
||||
nrecv := 0
|
||||
for {
|
||||
n, err := cli.Read(buf)
|
||||
if err != nil {
|
||||
return err
|
||||
} else {
|
||||
nrecv += n
|
||||
if nrecv == msglen {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func sink_tester(cli *UDPSession, msglen, msgcount int) error {
|
||||
// sender
|
||||
buf := make([]byte, msglen)
|
||||
for i := 0; i < msgcount; i++ {
|
||||
if _, err := cli.Write(buf); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestSNMP(t *testing.T) {
|
||||
t.Log(DefaultSnmp.Copy())
|
||||
t.Log(DefaultSnmp.Header())
|
||||
t.Log(DefaultSnmp.ToSlice())
|
||||
DefaultSnmp.Reset()
|
||||
t.Log(DefaultSnmp.ToSlice())
|
||||
}
|
||||
|
||||
func TestListenerClose(t *testing.T) {
|
||||
l, err := ListenWithOptions(portListerner, nil, 10, 3)
|
||||
if err != nil {
|
||||
t.Fail()
|
||||
}
|
||||
l.SetReadDeadline(time.Now().Add(time.Second))
|
||||
l.SetWriteDeadline(time.Now().Add(time.Second))
|
||||
l.SetDeadline(time.Now().Add(time.Second))
|
||||
time.Sleep(2 * time.Second)
|
||||
if _, err := l.Accept(); err == nil {
|
||||
t.Fail()
|
||||
}
|
||||
|
||||
l.Close()
|
||||
fakeaddr, _ := net.ResolveUDPAddr("udp6", "127.0.0.1:1111")
|
||||
if l.closeSession(fakeaddr) {
|
||||
t.Fail()
|
||||
}
|
||||
}
|
||||
164
vender/github.com/xtaci/kcp/snmp.go
Normal file
164
vender/github.com/xtaci/kcp/snmp.go
Normal file
@@ -0,0 +1,164 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Snmp defines network statistics indicator
|
||||
type Snmp struct {
|
||||
BytesSent uint64 // bytes sent from upper level
|
||||
BytesReceived uint64 // bytes received to upper level
|
||||
MaxConn uint64 // max number of connections ever reached
|
||||
ActiveOpens uint64 // accumulated active open connections
|
||||
PassiveOpens uint64 // accumulated passive open connections
|
||||
CurrEstab uint64 // current number of established connections
|
||||
InErrs uint64 // UDP read errors reported from net.PacketConn
|
||||
InCsumErrors uint64 // checksum errors from CRC32
|
||||
KCPInErrors uint64 // packet iput errors reported from KCP
|
||||
InPkts uint64 // incoming packets count
|
||||
OutPkts uint64 // outgoing packets count
|
||||
InSegs uint64 // incoming KCP segments
|
||||
OutSegs uint64 // outgoing KCP segments
|
||||
InBytes uint64 // UDP bytes received
|
||||
OutBytes uint64 // UDP bytes sent
|
||||
RetransSegs uint64 // accmulated retransmited segments
|
||||
FastRetransSegs uint64 // accmulated fast retransmitted segments
|
||||
EarlyRetransSegs uint64 // accmulated early retransmitted segments
|
||||
LostSegs uint64 // number of segs infered as lost
|
||||
RepeatSegs uint64 // number of segs duplicated
|
||||
FECRecovered uint64 // correct packets recovered from FEC
|
||||
FECErrs uint64 // incorrect packets recovered from FEC
|
||||
FECParityShards uint64 // FEC segments received
|
||||
FECShortShards uint64 // number of data shards that's not enough for recovery
|
||||
}
|
||||
|
||||
func newSnmp() *Snmp {
|
||||
return new(Snmp)
|
||||
}
|
||||
|
||||
// Header returns all field names
|
||||
func (s *Snmp) Header() []string {
|
||||
return []string{
|
||||
"BytesSent",
|
||||
"BytesReceived",
|
||||
"MaxConn",
|
||||
"ActiveOpens",
|
||||
"PassiveOpens",
|
||||
"CurrEstab",
|
||||
"InErrs",
|
||||
"InCsumErrors",
|
||||
"KCPInErrors",
|
||||
"InPkts",
|
||||
"OutPkts",
|
||||
"InSegs",
|
||||
"OutSegs",
|
||||
"InBytes",
|
||||
"OutBytes",
|
||||
"RetransSegs",
|
||||
"FastRetransSegs",
|
||||
"EarlyRetransSegs",
|
||||
"LostSegs",
|
||||
"RepeatSegs",
|
||||
"FECParityShards",
|
||||
"FECErrs",
|
||||
"FECRecovered",
|
||||
"FECShortShards",
|
||||
}
|
||||
}
|
||||
|
||||
// ToSlice returns current snmp info as slice
|
||||
func (s *Snmp) ToSlice() []string {
|
||||
snmp := s.Copy()
|
||||
return []string{
|
||||
fmt.Sprint(snmp.BytesSent),
|
||||
fmt.Sprint(snmp.BytesReceived),
|
||||
fmt.Sprint(snmp.MaxConn),
|
||||
fmt.Sprint(snmp.ActiveOpens),
|
||||
fmt.Sprint(snmp.PassiveOpens),
|
||||
fmt.Sprint(snmp.CurrEstab),
|
||||
fmt.Sprint(snmp.InErrs),
|
||||
fmt.Sprint(snmp.InCsumErrors),
|
||||
fmt.Sprint(snmp.KCPInErrors),
|
||||
fmt.Sprint(snmp.InPkts),
|
||||
fmt.Sprint(snmp.OutPkts),
|
||||
fmt.Sprint(snmp.InSegs),
|
||||
fmt.Sprint(snmp.OutSegs),
|
||||
fmt.Sprint(snmp.InBytes),
|
||||
fmt.Sprint(snmp.OutBytes),
|
||||
fmt.Sprint(snmp.RetransSegs),
|
||||
fmt.Sprint(snmp.FastRetransSegs),
|
||||
fmt.Sprint(snmp.EarlyRetransSegs),
|
||||
fmt.Sprint(snmp.LostSegs),
|
||||
fmt.Sprint(snmp.RepeatSegs),
|
||||
fmt.Sprint(snmp.FECParityShards),
|
||||
fmt.Sprint(snmp.FECErrs),
|
||||
fmt.Sprint(snmp.FECRecovered),
|
||||
fmt.Sprint(snmp.FECShortShards),
|
||||
}
|
||||
}
|
||||
|
||||
// Copy make a copy of current snmp snapshot
|
||||
func (s *Snmp) Copy() *Snmp {
|
||||
d := newSnmp()
|
||||
d.BytesSent = atomic.LoadUint64(&s.BytesSent)
|
||||
d.BytesReceived = atomic.LoadUint64(&s.BytesReceived)
|
||||
d.MaxConn = atomic.LoadUint64(&s.MaxConn)
|
||||
d.ActiveOpens = atomic.LoadUint64(&s.ActiveOpens)
|
||||
d.PassiveOpens = atomic.LoadUint64(&s.PassiveOpens)
|
||||
d.CurrEstab = atomic.LoadUint64(&s.CurrEstab)
|
||||
d.InErrs = atomic.LoadUint64(&s.InErrs)
|
||||
d.InCsumErrors = atomic.LoadUint64(&s.InCsumErrors)
|
||||
d.KCPInErrors = atomic.LoadUint64(&s.KCPInErrors)
|
||||
d.InPkts = atomic.LoadUint64(&s.InPkts)
|
||||
d.OutPkts = atomic.LoadUint64(&s.OutPkts)
|
||||
d.InSegs = atomic.LoadUint64(&s.InSegs)
|
||||
d.OutSegs = atomic.LoadUint64(&s.OutSegs)
|
||||
d.InBytes = atomic.LoadUint64(&s.InBytes)
|
||||
d.OutBytes = atomic.LoadUint64(&s.OutBytes)
|
||||
d.RetransSegs = atomic.LoadUint64(&s.RetransSegs)
|
||||
d.FastRetransSegs = atomic.LoadUint64(&s.FastRetransSegs)
|
||||
d.EarlyRetransSegs = atomic.LoadUint64(&s.EarlyRetransSegs)
|
||||
d.LostSegs = atomic.LoadUint64(&s.LostSegs)
|
||||
d.RepeatSegs = atomic.LoadUint64(&s.RepeatSegs)
|
||||
d.FECParityShards = atomic.LoadUint64(&s.FECParityShards)
|
||||
d.FECErrs = atomic.LoadUint64(&s.FECErrs)
|
||||
d.FECRecovered = atomic.LoadUint64(&s.FECRecovered)
|
||||
d.FECShortShards = atomic.LoadUint64(&s.FECShortShards)
|
||||
return d
|
||||
}
|
||||
|
||||
// Reset values to zero
|
||||
func (s *Snmp) Reset() {
|
||||
atomic.StoreUint64(&s.BytesSent, 0)
|
||||
atomic.StoreUint64(&s.BytesReceived, 0)
|
||||
atomic.StoreUint64(&s.MaxConn, 0)
|
||||
atomic.StoreUint64(&s.ActiveOpens, 0)
|
||||
atomic.StoreUint64(&s.PassiveOpens, 0)
|
||||
atomic.StoreUint64(&s.CurrEstab, 0)
|
||||
atomic.StoreUint64(&s.InErrs, 0)
|
||||
atomic.StoreUint64(&s.InCsumErrors, 0)
|
||||
atomic.StoreUint64(&s.KCPInErrors, 0)
|
||||
atomic.StoreUint64(&s.InPkts, 0)
|
||||
atomic.StoreUint64(&s.OutPkts, 0)
|
||||
atomic.StoreUint64(&s.InSegs, 0)
|
||||
atomic.StoreUint64(&s.OutSegs, 0)
|
||||
atomic.StoreUint64(&s.InBytes, 0)
|
||||
atomic.StoreUint64(&s.OutBytes, 0)
|
||||
atomic.StoreUint64(&s.RetransSegs, 0)
|
||||
atomic.StoreUint64(&s.FastRetransSegs, 0)
|
||||
atomic.StoreUint64(&s.EarlyRetransSegs, 0)
|
||||
atomic.StoreUint64(&s.LostSegs, 0)
|
||||
atomic.StoreUint64(&s.RepeatSegs, 0)
|
||||
atomic.StoreUint64(&s.FECParityShards, 0)
|
||||
atomic.StoreUint64(&s.FECErrs, 0)
|
||||
atomic.StoreUint64(&s.FECRecovered, 0)
|
||||
atomic.StoreUint64(&s.FECShortShards, 0)
|
||||
}
|
||||
|
||||
// DefaultSnmp is the global KCP connection statistics collector
|
||||
var DefaultSnmp *Snmp
|
||||
|
||||
func init() {
|
||||
DefaultSnmp = newSnmp()
|
||||
}
|
||||
104
vender/github.com/xtaci/kcp/updater.go
Normal file
104
vender/github.com/xtaci/kcp/updater.go
Normal file
@@ -0,0 +1,104 @@
|
||||
package kcp
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var updater updateHeap
|
||||
|
||||
func init() {
|
||||
updater.init()
|
||||
go updater.updateTask()
|
||||
}
|
||||
|
||||
// entry contains a session update info
|
||||
type entry struct {
|
||||
ts time.Time
|
||||
s *UDPSession
|
||||
}
|
||||
|
||||
// a global heap managed kcp.flush() caller
|
||||
type updateHeap struct {
|
||||
entries []entry
|
||||
mu sync.Mutex
|
||||
chWakeUp chan struct{}
|
||||
}
|
||||
|
||||
func (h *updateHeap) Len() int { return len(h.entries) }
|
||||
func (h *updateHeap) Less(i, j int) bool { return h.entries[i].ts.Before(h.entries[j].ts) }
|
||||
func (h *updateHeap) Swap(i, j int) {
|
||||
h.entries[i], h.entries[j] = h.entries[j], h.entries[i]
|
||||
h.entries[i].s.updaterIdx = i
|
||||
h.entries[j].s.updaterIdx = j
|
||||
}
|
||||
|
||||
func (h *updateHeap) Push(x interface{}) {
|
||||
h.entries = append(h.entries, x.(entry))
|
||||
n := len(h.entries)
|
||||
h.entries[n-1].s.updaterIdx = n - 1
|
||||
}
|
||||
|
||||
func (h *updateHeap) Pop() interface{} {
|
||||
n := len(h.entries)
|
||||
x := h.entries[n-1]
|
||||
h.entries[n-1].s.updaterIdx = -1
|
||||
h.entries[n-1] = entry{} // manual set nil for GC
|
||||
h.entries = h.entries[0 : n-1]
|
||||
return x
|
||||
}
|
||||
|
||||
func (h *updateHeap) init() {
|
||||
h.chWakeUp = make(chan struct{}, 1)
|
||||
}
|
||||
|
||||
func (h *updateHeap) addSession(s *UDPSession) {
|
||||
h.mu.Lock()
|
||||
heap.Push(h, entry{time.Now(), s})
|
||||
h.mu.Unlock()
|
||||
h.wakeup()
|
||||
}
|
||||
|
||||
func (h *updateHeap) removeSession(s *UDPSession) {
|
||||
h.mu.Lock()
|
||||
if s.updaterIdx != -1 {
|
||||
heap.Remove(h, s.updaterIdx)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func (h *updateHeap) wakeup() {
|
||||
select {
|
||||
case h.chWakeUp <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (h *updateHeap) updateTask() {
|
||||
var timer <-chan time.Time
|
||||
for {
|
||||
select {
|
||||
case <-timer:
|
||||
case <-h.chWakeUp:
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
hlen := h.Len()
|
||||
for i := 0; i < hlen; i++ {
|
||||
entry := &h.entries[0]
|
||||
if time.Now().After(entry.ts) {
|
||||
interval := entry.s.update()
|
||||
entry.ts = time.Now().Add(interval)
|
||||
heap.Fix(h, 0)
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if hlen > 0 {
|
||||
timer = time.After(h.entries[0].ts.Sub(time.Now()))
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user