Files
kcp-go/sess.go
T

1335 lines
36 KiB
Go

// The MIT License (MIT)
//
// # Copyright (c) 2015 xtaci
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
// [THE GENERALIZED DATA PIPELINE FOR KCP-GO]
//
// Outgoing Data Pipeline: Incoming Data Pipeline:
// Stream (Input Data) Packet Network (Network Interface Card)
// | |
// v v
// KCP Output (Reliable Transport Layer) Reader/Listener (Reception Queue)
// | |
// v v
// FEC Encoding (Forward Error Correction) Decryption (Data Security)
// | |
// v v
// CRC32 Checksum (Error Detection) CRC32 Checksum (Error Detection)
// | |
// v v
// Encryption (Data Security) FEC Decoding (Forward Error Correction)
// | |
// v v
// TxQueue (Transmission Queue) KCP Input (Reliable Transport Layer)
// | |
// v v
// Packet Network (Network Transmission) Stream (Input Data)
package kcp
import (
"context"
"crypto/rand"
"encoding/binary"
"hash/crc32"
"io"
"net"
"sync"
"sync/atomic"
"time"
"github.com/pkg/errors"
"golang.org/x/net/ipv4"
"golang.org/x/net/ipv6"
"golang.org/x/time/rate"
)
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
// accept backlog
acceptBacklog = 128
// dev backlog
devBacklog = 2048
// max latency for consecutive FEC encoding, in millisecond
maxFECEncodeLatency = 500
// max batch size
maxBatchSize = 64
)
var (
errInvalidOperation = errors.New("invalid operation")
errTimeout = timeoutError{}
errNotOwner = errors.New("not the owner of this connection")
)
// timeoutError implements net.Error
type timeoutError struct{}
func (timeoutError) Error() string { return "timeout" }
func (timeoutError) Timeout() bool { return true }
func (timeoutError) Temporary() bool { return true }
type (
// UDPSession defines a KCP session implemented by UDP
UDPSession struct {
conn net.PacketConn // the underlying packet connection
ownConn bool // true if we created conn internally, false if provided by caller
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
// FEC codec
fecDecoder *fecDecoder
fecEncoder *fecEncoder
// settings
remote net.Addr // remote peer address
rd atomic.Value // read deadline
wd atomic.Value // 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
dieOnce sync.Once
chReadEvent chan struct{} // notify Read() can be called without blocking
chWriteEvent chan struct{} // notify Write() can be called without blocking
// socket error handling
socketReadError atomic.Value
socketWriteError atomic.Value
chSocketReadError chan struct{}
chSocketWriteError chan struct{}
socketReadErrorOnce sync.Once
socketWriteErrorOnce sync.Once
// packets waiting to be sent on wire
chPostProcessing chan []byte
// platform-dependent optimizations
platform platform
// rate limiter (bytes per second)
rateLimiter atomic.Value
mu sync.Mutex
}
setReadBuffer interface {
SetReadBuffer(bytes int) error
}
setWriteBuffer interface {
SetWriteBuffer(bytes int) error
}
setDSCP interface {
SetDSCP(int) error
}
)
// newUDPSession create a new udp session for client or server
func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn net.PacketConn, ownConn bool, remote net.Addr, block BlockCrypt) *UDPSession {
sess := new(UDPSession)
sess.die = make(chan struct{})
sess.chReadEvent = make(chan struct{}, 1)
sess.chWriteEvent = make(chan struct{}, 1)
sess.chSocketReadError = make(chan struct{})
sess.chSocketWriteError = make(chan struct{})
sess.chPostProcessing = make(chan []byte, devBacklog)
sess.remote = remote
sess.conn = conn
sess.ownConn = ownConn
sess.l = l
sess.block = block
sess.recvbuf = make([]byte, mtuLimit)
sess.initPlatform()
// calculate additional header size introduced by encryption
switch block := sess.block.(type) {
case nil:
sess.headerSize = 0
case *aeadCrypt:
sess.headerSize = block.NonceSize()
default:
sess.headerSize = cryptHeaderSize
}
// FEC codec initialization
sess.fecDecoder = newFECDecoder(dataShards, parityShards)
sess.fecEncoder = newFECEncoder(dataShards, parityShards, sess.headerSize)
// calculate additional header size introduced by FEC
if sess.fecEncoder != nil {
sess.headerSize += fecHeaderSizePlus2
}
sess.kcp = NewKCP(conv, func(buf []byte, size int) {
// A basic check for the minimum packet size
if size >= IKCP_OVERHEAD {
// make a copy
bts := defaultBufferPool.Get()[:size+sess.headerSize]
// copy the data to a new buffer, and reserve header space
copy(bts[sess.headerSize:], buf)
// delivery to post processing (non-blocking to avoid deadlock under lock)
select {
case sess.chPostProcessing <- bts:
case <-sess.die:
return
default:
// drop and recycle to avoid blocking; KCP will retransmit if needed
defaultBufferPool.Put(bts)
}
}
})
// Set Default MTU
if !sess.SetMtu(IKCP_MTU_DEF) {
panic("Overhead too large")
}
// create post-processing goroutine
go sess.postProcess()
if sess.l == nil { // it's a client connection
go sess.readLoop()
atomic.AddUint64(&DefaultSnmp.ActiveOpens, 1)
} else {
atomic.AddUint64(&DefaultSnmp.PassiveOpens, 1)
}
// start per-session updater
SystemTimedSched.Put(sess.update, time.Now())
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) {
RESET_TIMER:
var timeout *time.Timer
// deadline for current reading operation
var c <-chan time.Time
if trd, ok := s.rd.Load().(time.Time); ok && !trd.IsZero() {
timeout = time.NewTimer(time.Until(trd))
c = timeout.C
defer timeout.Stop()
}
for {
s.mu.Lock()
// bufptr points to the current position of recvbuf,
// if previous 'b' is insufficient to accommodate the data, the
// remaining data will be stored in bufptr for next read.
if len(s.bufptr) > 0 {
n = copy(b, s.bufptr)
s.bufptr = s.bufptr[n:]
s.mu.Unlock()
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(n))
return n, nil
}
if size := s.kcp.PeekSize(); size > 0 { // peek data size from kcp
// if 'b' is large enough to accommodate the data, read directly
// from kcp.recv() to 'b', like 'DMA'.
if len(b) >= size {
s.kcp.Recv(b)
s.mu.Unlock()
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(size))
return size, nil
}
// otherwise, read to recvbuf first, then copy to 'b'.
// dynamically adjust the buffer size to the maximum of 'packet size' when necessary.
if cap(s.recvbuf) < size {
// usually recvbuf has a size of maximum packet size
s.recvbuf = make([]byte, size)
}
// resize the length of recvbuf to match the data size
s.recvbuf = s.recvbuf[:size]
s.kcp.Recv(s.recvbuf) // read data to recvbuf first
n = copy(b, s.recvbuf) // then copy bytes to 'b' as many as possible
s.bufptr = s.recvbuf[n:] // pointer update
s.mu.Unlock()
atomic.AddUint64(&DefaultSnmp.BytesReceived, uint64(n))
return n, nil
}
s.mu.Unlock()
// if it runs here, that means we have to block the call, and wait until the
// next data packet arrives.
select {
case <-s.chReadEvent:
if timeout != nil {
timeout.Stop()
goto RESET_TIMER
}
case <-c:
return 0, errors.WithStack(errTimeout)
case <-s.chSocketReadError:
return 0, s.socketReadError.Load().(error)
case <-s.die:
return 0, errors.WithStack(io.ErrClosedPipe)
}
}
}
// Write implements net.Conn
func (s *UDPSession) Write(b []byte) (n int, err error) { return s.WriteBuffers([][]byte{b}) }
// WriteBuffers write a vector of byte slices to the underlying connection
func (s *UDPSession) WriteBuffers(v [][]byte) (n int, err error) {
RESET_TIMER:
var timeout *time.Timer
var c <-chan time.Time
if twd, ok := s.wd.Load().(time.Time); ok && !twd.IsZero() {
timeout = time.NewTimer(time.Until(twd))
c = timeout.C
defer timeout.Stop()
}
for {
// check for connection close and socket error
select {
case <-s.chSocketWriteError:
return 0, s.socketWriteError.Load().(error)
case <-s.die:
return 0, errors.WithStack(io.ErrClosedPipe)
default:
}
s.mu.Lock()
// make sure write do not overflow the max sliding window on both side
waitsnd := s.kcp.WaitSnd()
if waitsnd < int(s.kcp.snd_wnd) {
// transmit all data sequentially, make sure every packet size is within 'mss'
for _, b := range v {
n += len(b)
// handle each slice for packet splitting
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:]
}
}
}
waitsnd = s.kcp.WaitSnd()
if waitsnd >= int(s.kcp.snd_wnd) || !s.writeDelay {
// put the packets on wire immediately if the inflight window is full
// or if we've specified write no delay(NO merging of outgoing bytes)
// we don't have to wait until the periodical update() procedure uncorks.
s.kcp.flush(IKCP_FLUSH_FULL)
}
s.mu.Unlock()
atomic.AddUint64(&DefaultSnmp.BytesSent, uint64(n))
return n, nil
}
s.mu.Unlock()
// if it runs here, that means we have to block the call, and wait until the
// transmit buffer to become available again.
select {
case <-s.chWriteEvent:
if timeout != nil {
timeout.Stop()
goto RESET_TIMER
}
case <-c:
return 0, errors.WithStack(errTimeout)
case <-s.chSocketWriteError:
return 0, s.socketWriteError.Load().(error)
case <-s.die:
return 0, errors.WithStack(io.ErrClosedPipe)
}
}
}
func (s *UDPSession) isClosed() bool {
select {
case <-s.die:
return true
default:
return false
}
}
// Close closes the connection.
func (s *UDPSession) Close() error {
var once bool
s.dieOnce.Do(func() {
close(s.die)
once = true
})
if !once {
return errors.WithStack(io.ErrClosedPipe)
}
atomic.AddUint64(&DefaultSnmp.CurrEstab, ^uint64(0))
// try best to send all queued messages especially the data in txqueue
s.mu.Lock()
s.kcp.flush((IKCP_FLUSH_FULL))
s.mu.Unlock()
if s.l != nil { // belongs to listener
s.l.closeSession(s.remote)
return nil
}
if s.ownConn { // 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.rd.Store(t)
s.wd.Store(t)
s.notifyReadEvent()
s.notifyWriteEvent()
return nil
}
// SetReadDeadline implements the Conn SetReadDeadline method.
func (s *UDPSession) SetReadDeadline(t time.Time) error {
s.rd.Store(t)
s.notifyReadEvent()
return nil
}
// SetWriteDeadline implements the Conn SetWriteDeadline method.
func (s *UDPSession) SetWriteDeadline(t time.Time) error {
s.wd.Store(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 {
mtu = min(mtuLimit, mtu)
mtu -= s.headerSize
if aead, ok := s.block.(*aeadCrypt); ok {
mtu -= aead.Overhead()
}
s.mu.Lock()
defer s.mu.Unlock()
ret := s.kcp.SetMtu(mtu) // kcp mtu is not including udp header
return ret == 0
}
// Deprecated: 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
}
// (deprecated)
//
// SetDUP duplicates udp packets for kcp output.
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 in IPv4 header, or 8bit Traffic Class in IPv6 header.
//
// if the underlying connection has implemented `func SetDSCP(int) error`, SetDSCP() will invoke
// this function instead.
//
// It has 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 {
return errInvalidOperation
}
// interface enabled
if ts, ok := s.conn.(setDSCP); ok {
return ts.SetDSCP(dscp)
}
if nc, ok := s.conn.(net.Conn); ok {
var succeed bool
if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err == nil {
succeed = true
}
if err := ipv6.NewConn(nc).SetTrafficClass(dscp); err == nil {
succeed = true
}
if succeed {
return nil
}
}
return 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 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 errInvalidOperation
}
// SetRateLimit sets the rate limit for this session in bytes per second,
// by setting to 0 will disable rate limiting.
func (s *UDPSession) SetRateLimit(bytesPerSecond uint32) {
var limiter *rate.Limiter
if bytesPerSecond == 0 {
limiter = rate.NewLimiter(rate.Inf, maxBatchSize*mtuLimit)
} else {
limiter = rate.NewLimiter(rate.Limit(bytesPerSecond), maxBatchSize*mtuLimit)
}
s.rateLimiter.Store(limiter)
}
// SetLogger configures the kcp trace logger
func (s *UDPSession) SetLogger(mask KCPLogType, logger logoutput_callback) {
s.kcp.SetLogger(mask, logger)
}
// Control applys a procedure to the underly socket fd.
// CAUTION: BE VERY CAREFUL TO USE THIS FUNCTION, YOU MAY BREAK THE PROTOCOL.
func (s *UDPSession) Control(f func(conn net.PacketConn) error) error {
if !s.ownConn {
return errNotOwner
}
s.mu.Lock()
defer s.mu.Unlock()
return f(s.conn)
}
// a goroutine to handle post processing of kcp and make the critical section smaller
// pipeline for outgoing packets (from ARQ to network)
//
// KCP output -> FEC encoding -> CRC32 integrity -> Encryption -> TxQueue
func (s *UDPSession) postProcess() {
txqueue := make([]ipv4.Message, 0, devBacklog)
chDie := s.die
ctx := context.Background()
bytesToSend := 0
for {
select {
case buf := <-s.chPostProcessing: // dequeue from post processing
var ecc [][]byte
// 1. FEC encoding
if s.fecEncoder != nil {
ecc = s.fecEncoder.encode(buf, maxFECEncodeLatency)
}
// 2. Encryption
switch block := s.block.(type) {
case nil:
case *aeadCrypt: // AEAD mode
nonceSize := block.NonceSize()
dst := buf[:nonceSize]
nonce := buf[:nonceSize]
plaintext := buf[nonceSize:]
fillRand(nonce)
buf = block.Seal(dst, nonce, plaintext, nil)
for k := range ecc {
dst := ecc[k][:nonceSize]
nonce := ecc[k][:nonceSize]
plaintext := ecc[k][nonceSize:]
fillRand(nonce)
ecc[k] = block.Seal(dst, nonce, plaintext, nil)
}
default: // Cipher Feedback (CFB) mode
fillRand(buf[:nonceSize])
checksum := crc32.ChecksumIEEE(buf[cryptHeaderSize:])
binary.LittleEndian.PutUint32(buf[nonceSize:], checksum)
block.Encrypt(buf, buf)
for k := range ecc {
fillRand(ecc[k][:nonceSize])
checksum := crc32.ChecksumIEEE(ecc[k][cryptHeaderSize:])
binary.LittleEndian.PutUint32(ecc[k][nonceSize:], checksum)
block.Encrypt(ecc[k], ecc[k])
}
}
// 3. TxQueue
var msg ipv4.Message
msg.Addr = s.remote
// original copy, move buf to txqueue directly
msg.Buffers = [][]byte{buf}
bytesToSend += len(buf)
txqueue = append(txqueue, msg)
// dup copies for testing if set
for i := 0; i < s.dup; i++ {
bts := defaultBufferPool.Get()[:len(buf)]
copy(bts, buf)
msg.Buffers = [][]byte{bts}
bytesToSend += len(bts)
txqueue = append(txqueue, msg)
}
// parity
for k := range ecc {
bts := defaultBufferPool.Get()[:len(ecc[k])]
copy(bts, ecc[k])
msg.Buffers = [][]byte{bts}
bytesToSend += len(bts)
txqueue = append(txqueue, msg)
}
// transmit when chPostProcessing is empty or we've reached max batch size
if len(s.chPostProcessing) == 0 || len(txqueue) >= maxBatchSize {
if limiter, ok := s.rateLimiter.Load().(*rate.Limiter); ok {
err := limiter.WaitN(ctx, bytesToSend)
if err != nil {
panic(err)
}
}
s.tx(txqueue)
s.kcp.debugLog(IKCP_LOG_OUTPUT, "conv", s.kcp.conv, "datalen", bytesToSend)
// recycle
for k := range txqueue {
defaultBufferPool.Put(txqueue[k].Buffers[0])
txqueue[k].Buffers = nil
}
txqueue = txqueue[:0]
bytesToSend = 0
}
// re-enable die channel
chDie = s.die
case <-chDie:
// remaining packets in txqueue should be sent out
if len(s.chPostProcessing) > 0 {
chDie = nil // block chDie temporarily
continue
}
return
}
}
}
// sess update to trigger protocol
func (s *UDPSession) update() {
select {
case <-s.die:
default:
s.mu.Lock()
interval := s.kcp.flush(IKCP_FLUSH_FULL)
waitsnd := s.kcp.WaitSnd()
if waitsnd < int(s.kcp.snd_wnd) {
s.notifyWriteEvent()
}
s.mu.Unlock()
// self-synchronized timed scheduling
SystemTimedSched.Put(s.update, time.Now().Add(time.Duration(interval)*time.Millisecond))
}
}
// GetConv gets conversation id of a session
func (s *UDPSession) GetConv() uint32 { return s.kcp.conv }
// GetRTO gets current rto of the session
func (s *UDPSession) GetRTO() uint32 {
s.mu.Lock()
defer s.mu.Unlock()
return s.kcp.rx_rto
}
// GetSRTT gets current srtt of the session
func (s *UDPSession) GetSRTT() int32 {
s.mu.Lock()
defer s.mu.Unlock()
return s.kcp.rx_srtt
}
// GetRTTVar gets current rtt variance of the session
func (s *UDPSession) GetSRTTVar() int32 {
s.mu.Lock()
defer s.mu.Unlock()
return s.kcp.rx_rttvar
}
func (s *UDPSession) notifyReadEvent() {
select {
case s.chReadEvent <- struct{}{}:
default:
}
}
func (s *UDPSession) notifyWriteEvent() {
select {
case s.chWriteEvent <- struct{}{}:
default:
}
}
func (s *UDPSession) notifyReadError(err error) {
s.socketReadErrorOnce.Do(func() {
s.socketReadError.Store(err)
close(s.chSocketReadError)
})
}
func (s *UDPSession) notifyWriteError(err error) {
s.socketWriteErrorOnce.Do(func() {
s.socketWriteError.Store(err)
close(s.chSocketWriteError)
})
}
// packet input pipeline:
// network -> [decryption ->] [crc32 ->] [FEC ->] [KCP input ->] stream -> application
func (s *UDPSession) packetInput(data []byte) {
switch block := s.block.(type) {
case nil:
case *aeadCrypt:
nonceSize := block.NonceSize()
if len(data) < nonceSize+block.Overhead() {
return
}
nonce := data[:nonceSize]
ciphertext := data[nonceSize:]
plaintext, err := block.Open(ciphertext[:0], nonce, ciphertext, nil)
if err != nil {
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
return
}
data = plaintext
default:
// decryption and crc32 check
if len(data) < cryptHeaderSize {
return
}
block.Decrypt(data, data)
data = data[nonceSize:]
checksum := crc32.ChecksumIEEE(data[crcSize:])
if checksum != binary.LittleEndian.Uint32(data) {
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
return
}
data = data[crcSize:]
}
// basic check for minimum packet size
if len(data) < IKCP_OVERHEAD {
atomic.AddUint64(&DefaultSnmp.KCPInErrors, 1)
return
}
s.kcpInput(data)
}
// kcpInput inputs a decrypted and crc32-checked packet into kcp with FEC handling
func (s *UDPSession) kcpInput(data []byte) {
atomic.AddUint64(&DefaultSnmp.InPkts, 1)
atomic.AddUint64(&DefaultSnmp.InBytes, uint64(len(data)))
// 16bit kcp cmd [81-84] and frg [0-255] will not overlap with FEC type 0x00f1 0x00f2
fecFlag := binary.LittleEndian.Uint16(data[4:])
switch fecFlag {
case typeData, typeParity: // packet with FEC
if len(data) < fecHeaderSizePlus2 {
atomic.AddUint64(&DefaultSnmp.InErrs, 1)
return
}
var kcpInErrors uint64
f := fecPacket(data)
// lock
s.mu.Lock()
defer s.mu.Unlock()
// if fecDecoder is not initialized, create one with default parameter
// lazy initialization
if s.fecDecoder == nil {
s.fecDecoder = newFECDecoder(1, 1)
}
// FEC decoding
recovers := s.fecDecoder.decode(f)
if f.flag() == typeData {
if ret := s.kcp.Input(data[fecHeaderSizePlus2:], IKCP_PACKET_REGULAR, s.ackNoDelay); ret != 0 {
kcpInErrors++
}
}
// If there're some packets recovered from FEC, feed them into kcp
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], IKCP_PACKET_FEC, s.ackNoDelay); ret != 0 {
kcpInErrors++
}
}
}
// recycle the buffer
defaultBufferPool.Put(r)
}
// to notify the readers to receive the data if there's any
if n := s.kcp.PeekSize(); n > 0 {
s.notifyReadEvent()
}
// to notify the writers if the window size allows to send more packets
// and the remote window size is not full.
waitsnd := s.kcp.WaitSnd()
if waitsnd < int(s.kcp.snd_wnd) {
s.notifyWriteEvent()
}
if kcpInErrors > 0 {
atomic.AddUint64(&DefaultSnmp.KCPInErrors, kcpInErrors)
}
default: // packet without FEC
s.mu.Lock()
defer s.mu.Unlock()
if ret := s.kcp.Input(data, IKCP_PACKET_REGULAR, s.ackNoDelay); ret != 0 {
atomic.AddUint64(&DefaultSnmp.KCPInErrors, 1)
}
if n := s.kcp.PeekSize(); n > 0 {
s.notifyReadEvent()
}
waitsnd := s.kcp.WaitSnd()
if waitsnd < int(s.kcp.snd_wnd) {
s.notifyWriteEvent()
}
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
conn net.PacketConn // the underlying packet connection
ownConn bool // true if we created conn internally, false if provided by caller
sessions map[string]*UDPSession // all sessions accepted by this Listener
sessionLock sync.RWMutex
chAccepts chan *UDPSession // Listen() backlog
chSessionClosed chan net.Addr // session close queue
die chan struct{} // notify the listener has closed
dieOnce sync.Once
// socket error handling
socketReadError atomic.Value
chSocketReadError chan struct{}
socketReadErrorOnce sync.Once
rd atomic.Value // read deadline for Accept()
}
)
// packet input stage
func (l *Listener) packetInput(data []byte, addr net.Addr) {
switch block := l.block.(type) {
case nil:
case *aeadCrypt:
nonceSize := block.NonceSize()
if len(data) < nonceSize+block.Overhead() {
return
}
nonce := data[:nonceSize]
ciphertext := data[nonceSize:]
plaintext, err := block.Open(ciphertext[:0], nonce, ciphertext, nil)
if err != nil {
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
return
}
data = plaintext
default:
// decryption and crc32 check
if len(data) < cryptHeaderSize {
return
}
block.Decrypt(data, data)
data = data[nonceSize:]
checksum := crc32.ChecksumIEEE(data[crcSize:])
if checksum != binary.LittleEndian.Uint32(data) {
atomic.AddUint64(&DefaultSnmp.InCsumErrors, 1)
return
}
data = data[crcSize:]
}
// basic check for minimum packet size
if len(data) < IKCP_OVERHEAD {
return
}
// look for existing session
l.sessionLock.RLock()
s, exist := l.sessions[addr.String()]
l.sessionLock.RUnlock()
var conv, sn uint32
hasConv := false
// try to get conversation id from the packet
// 16bit kcp cmd [81-84] and frg [0-255] will not overlap with FEC type 0x00f1 0x00f2
fecFlag := binary.LittleEndian.Uint16(data[4:])
switch fecFlag {
case typeData:
// data packet of FEC, conversation id inside
if len(data) < fecHeaderSizePlus2+IKCP_OVERHEAD {
break
}
hasConv = true
conv = binary.LittleEndian.Uint32(data[fecHeaderSizePlus2:])
sn = binary.LittleEndian.Uint32(data[fecHeaderSizePlus2+IKCP_SN_OFFSET:])
case typeParity:
// parity packet of FEC, conversation id inside
default:
// packet without FEC
hasConv = true
conv = binary.LittleEndian.Uint32(data)
sn = binary.LittleEndian.Uint32(data[IKCP_SN_OFFSET:])
}
// on an existing connection
if exist {
// If we have a valid conversation id or we cannot get conversation id from the packet,
// just feed the data into the existing session.
if !hasConv || conv == s.kcp.conv {
s.kcpInput(data)
return
}
// conversation id mismatched, only accept reset packet with sn == 0
if sn != 0 {
return
}
// Close will remove the session from listener's session map,
// So we can create a new session with the same addr below.
s.Close()
}
// The connection does not exist, try to create a new one.
// But if we don't have a valid conversation id, nothing we can do here except dropping the packet.
if !hasConv {
return
}
// Now we have a valid conversation id here without a session object, create a new session.
// do not let the new sessions overwhelm accept queue
if len(l.chAccepts) >= cap(l.chAccepts) {
return
}
// new session
s = newUDPSession(conv, l.dataShards, l.parityShards, l, l.conn, false, addr, l.block)
s.kcpInput(data)
l.sessionLock.Lock()
l.sessions[addr.String()] = s
l.sessionLock.Unlock()
l.chAccepts <- s
}
func (l *Listener) notifyReadError(err error) {
l.socketReadErrorOnce.Do(func() {
l.socketReadError.Store(err)
close(l.chSocketReadError)
// propagate read error to all sessions
l.sessionLock.RLock()
for _, s := range l.sessions {
s.notifyReadError(err)
}
l.sessionLock.RUnlock()
})
}
// SetReadBuffer sets the socket read buffer for the Listener
func (l *Listener) SetReadBuffer(bytes int) error {
if conn, ok := l.conn.(setReadBuffer); ok {
return conn.SetReadBuffer(bytes)
}
return errInvalidOperation
}
// SetWriteBuffer sets the socket write buffer for the Listener
func (l *Listener) SetWriteBuffer(bytes int) error {
if conn, ok := l.conn.(setWriteBuffer); ok {
return conn.SetWriteBuffer(bytes)
}
return errInvalidOperation
}
// SetDSCP sets the 6bit DSCP field in IPv4 header, or 8bit Traffic Class in IPv6 header.
//
// if the underlying connection has implemented `func SetDSCP(int) error`, SetDSCP() will invoke
// this function instead.
func (l *Listener) SetDSCP(dscp int) error {
// interface enabled
if conn, ok := l.conn.(setDSCP); ok {
return conn.SetDSCP(dscp)
}
conn, ok := l.conn.(net.Conn)
if !ok {
return errInvalidOperation
}
var succeed bool
if err := ipv4.NewConn(conn).SetTOS(dscp << 2); err == nil {
succeed = true
}
if err := ipv6.NewConn(conn).SetTrafficClass(dscp); err == nil {
succeed = true
}
if succeed {
return nil
}
return 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() {
timer := time.NewTimer(time.Until(tdeadline))
defer timer.Stop()
timeout = timer.C
}
select {
case <-timeout:
return nil, errors.WithStack(errTimeout)
case c := <-l.chAccepts:
return c, nil
case <-l.chSocketReadError:
return nil, l.socketReadError.Load().(error)
case <-l.die:
return nil, errors.WithStack(io.ErrClosedPipe)
}
}
// 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 {
return errInvalidOperation
}
// Close stops listening on the UDP address, and closes the socket
func (l *Listener) Close() error {
var once bool
l.dieOnce.Do(func() {
close(l.die)
once = true
})
if !once {
return errors.WithStack(io.ErrClosedPipe)
}
if l.ownConn {
return l.conn.Close()
}
return nil
}
// Control applys a procedure to the underly socket fd.
// CAUTION: BE VERY CAREFUL TO USE THIS FUNCTION, YOU MAY BREAK THE PROTOCOL.
func (l *Listener) Control(f func(conn net.PacketConn) error) error {
l.sessionLock.Lock()
defer l.sessionLock.Unlock()
return f(l.conn)
}
// 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.
//
// 'block' is the block encryption algorithm to encrypt packets.
//
// 'dataShards', 'parityShards' specify how many parity packets will be generated following the data packets.
//
// Check https://github.com/klauspost/reedsolomon for details
func ListenWithOptions(laddr string, block BlockCrypt, dataShards, parityShards int) (*Listener, error) {
udpaddr, err := net.ResolveUDPAddr("udp", laddr)
if err != nil {
return nil, errors.WithStack(err)
}
conn, err := net.ListenUDP("udp", udpaddr)
if err != nil {
return nil, errors.WithStack(err)
}
return serveConn(block, dataShards, parityShards, conn, true)
}
// ServeConn serves KCP protocol for a single packet connection.
func ServeConn(block BlockCrypt, dataShards, parityShards int, conn net.PacketConn) (*Listener, error) {
return serveConn(block, dataShards, parityShards, conn, false)
}
func serveConn(block BlockCrypt, dataShards, parityShards int, conn net.PacketConn, ownConn bool) (*Listener, error) {
l := new(Listener)
l.conn = conn
l.ownConn = ownConn
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.chSocketReadError = make(chan struct{})
go l.monitor()
return l, nil
}
// Dial connects to the remote address "raddr" on the network "udp" without encryption and FEC
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
//
// 'block' is the block encryption algorithm to encrypt packets.
//
// 'dataShards', 'parityShards' specify how many parity packets will be generated following the data packets.
//
// Check https://github.com/klauspost/reedsolomon for details
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.WithStack(err)
}
network := "udp4"
if udpaddr.IP.To4() == nil {
network = "udp"
}
conn, err := net.ListenUDP(network, nil)
if err != nil {
return nil, errors.WithStack(err)
}
var convid uint32
binary.Read(rand.Reader, binary.LittleEndian, &convid)
return newUDPSession(convid, dataShards, parityShards, nil, conn, true, udpaddr, block), nil
}
// NewConn4 establishes a session and talks KCP protocol over a packet connection.
func NewConn4(convid uint32, raddr net.Addr, block BlockCrypt, dataShards, parityShards int, ownConn bool, conn net.PacketConn) (*UDPSession, error) {
return newUDPSession(convid, dataShards, parityShards, nil, conn, ownConn, raddr, block), nil
}
// NewConn3 establishes a session and talks KCP protocol over a packet connection.
func NewConn3(convid uint32, raddr net.Addr, block BlockCrypt, dataShards, parityShards int, conn net.PacketConn) (*UDPSession, error) {
return newUDPSession(convid, dataShards, parityShards, nil, conn, false, raddr, block), nil
}
// NewConn2 establishes a session and talks KCP protocol over a packet connection.
func NewConn2(raddr net.Addr, block BlockCrypt, dataShards, parityShards int, conn net.PacketConn) (*UDPSession, error) {
var convid uint32
binary.Read(rand.Reader, binary.LittleEndian, &convid)
return NewConn3(convid, 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.WithStack(err)
}
return NewConn2(udpaddr, block, dataShards, parityShards, conn)
}