From d5d31b6d5d6215db9cc23a09916673f1a93bbda1 Mon Sep 17 00:00:00 2001 From: hikari <13061980190@163.com> Date: Sat, 18 May 2019 00:14:29 -0700 Subject: [PATCH] merge kcp-go kcptun. update p2p protocol --- kcp-go/batchconn.go | 12 + kcp-go/fec.go | 113 ++++---- kcp-go/kcp.go | 136 +++++---- kcp-go/readloop_generic.go | 50 ++++ kcp-go/readloop_linux.go | 80 +++++ kcp-go/sess.go | 578 +++++++++++++++++++------------------ kcp-go/tx_generic.go | 27 ++ kcp-go/tx_linux.go | 31 ++ kcp-go/updater.go | 8 +- p2pclient/config.go | 1 + p2pclient/main.go | 175 ++++------- p2pserver/main.go | 227 +++++++-------- 12 files changed, 795 insertions(+), 643 deletions(-) create mode 100644 kcp-go/batchconn.go create mode 100755 kcp-go/readloop_generic.go create mode 100755 kcp-go/readloop_linux.go create mode 100644 kcp-go/tx_generic.go create mode 100644 kcp-go/tx_linux.go diff --git a/kcp-go/batchconn.go b/kcp-go/batchconn.go new file mode 100644 index 0000000..6c30701 --- /dev/null +++ b/kcp-go/batchconn.go @@ -0,0 +1,12 @@ +package kcp + +import "golang.org/x/net/ipv4" + +const ( + batchSize = 16 +) + +type batchConn interface { + WriteBatch(ms []ipv4.Message, flags int) (int, error) + ReadBatch(ms []ipv4.Message, flags int) (int, error) +} diff --git a/kcp-go/fec.go b/kcp-go/fec.go index 366637b..4501c67 100644 --- a/kcp-go/fec.go +++ b/kcp-go/fec.go @@ -11,36 +11,34 @@ const ( fecHeaderSize = 6 fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size typeData = 0xf1 - typeFEC = 0xf2 + typeParity = 0xf2 ) -type ( - // fecPacket is a decoded FEC packet - fecPacket struct { - seqid uint32 - flag uint16 - data []byte - } +// fecPacket is a decoded FEC packet +type fecPacket []byte - // fecDecoder for decoding incoming packets - fecDecoder struct { - rxlimit int // queue size limit - dataShards int - parityShards int - shardSize int - rx []fecPacket // ordered receive queue +func (bts fecPacket) seqid() uint32 { return binary.LittleEndian.Uint32(bts) } +func (bts fecPacket) flag() uint16 { return binary.LittleEndian.Uint16(bts[4:]) } +func (bts fecPacket) data() []byte { return bts[6:] } - // caches - decodeCache [][]byte - flagCache []bool +// fecDecoder for decoding incoming packets +type fecDecoder struct { + rxlimit int // queue size limit + dataShards int + parityShards int + shardSize int + rx []fecPacket // ordered receive queue - // zeros - zeros []byte + // caches + decodeCache [][]byte + flagCache []bool - // RS decoder - codec reedsolomon.Encoder - } -) + // zeros + zeros []byte + + // RS decoder + codec reedsolomon.Encoder +} func newFECDecoder(rxlimit, dataShards, parityShards int) *fecDecoder { if dataShards <= 0 || parityShards <= 0 { @@ -66,33 +64,24 @@ func newFECDecoder(rxlimit, dataShards, parityShards int) *fecDecoder { 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) { +func (dec *fecDecoder) decode(in 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) + if in.seqid() == dec.rx[i].seqid() { // de-duplicate return nil - } else if _itimediff(pkt.seqid, dec.rx[i].seqid) > 0 { // insertion + } else if _itimediff(in.seqid(), dec.rx[i].seqid()) > 0 { // insertion insertIdx = i + 1 break } } + // make a copy + pkt := fecPacket(xmitBuf.Get().([]byte)[:len(in)]) + copy(pkt, in) + // insert into ordered rx queue if insertIdx == n+1 { dec.rx = append(dec.rx, pkt) @@ -103,11 +92,11 @@ func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) { } // shard range for current packet - shardBegin := pkt.seqid - pkt.seqid%uint32(dec.shardSize) + 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)) + searchBegin := insertIdx - int(pkt.seqid()%uint32(dec.shardSize)) if searchBegin < 0 { searchBegin = 0 } @@ -130,21 +119,21 @@ func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) { // shard assembly for i := searchBegin; i <= searchEnd; i++ { - seqid := dec.rx[i].seqid + 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 + shards[seqid%uint32(dec.shardSize)] = dec.rx[i].data() shardsflag[seqid%uint32(dec.shardSize)] = true numshard++ - if dec.rx[i].flag == typeData { + 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 len(dec.rx[i].data()) > maxlen { + maxlen = len(dec.rx[i].data()) } } } @@ -159,11 +148,14 @@ func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) { dlen := len(shards[k]) shards[k] = shards[k][:maxlen] copy(shards[k][dlen:], dec.zeros) + } else { + shards[k] = xmitBuf.Get().([]byte)[:0] } } if err := dec.codec.ReconstructData(shards); err == nil { for k := range shards[:dec.dataShards] { if !shardsflag[k] { + // recovered data should be recycled recovered = append(recovered, shards[k]) } } @@ -174,7 +166,7 @@ func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) { // keep rxlimit if len(dec.rx) > dec.rxlimit { - if dec.rx[0].flag == typeData { // track the unrecoverable data + if dec.rx[0].flag() == typeData { // track the unrecoverable data atomic.AddUint64(&DefaultSnmp.FECShortShards, 1) } dec.rx = dec.freeRange(0, 1, dec.rx) @@ -182,15 +174,16 @@ func (dec *fecDecoder) decode(pkt fecPacket) (recovered [][]byte) { return } -// free a range of fecPacket, and zero for GC recycling +// free a range of fecPacket func (dec *fecDecoder) freeRange(first, n int, q []fecPacket) []fecPacket { for i := first; i < first+n; i++ { // recycle buffer - xmitBuf.Put(q[i].data) + xmitBuf.Put([]byte(q[i])) + } + + if first == 0 && n < cap(q)/2 { + return q[n:] } 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] } @@ -252,13 +245,16 @@ func newFECEncoder(dataShards, parityShards, offset int) *fecEncoder { // 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) { + // The header format: + // | FEC SEQID(4B) | FEC TYPE(2B) | SIZE (2B) | PAYLOAD(SIZE-2) | + // |<-headerOffset |<-payloadOffset enc.markData(b[enc.headerOffset:]) binary.LittleEndian.PutUint16(b[enc.payloadOffset:], uint16(len(b[enc.payloadOffset:]))) - // copy data to fec datashards + // copy data from payloadOffset to fec shard cache sz := len(b) enc.shardCache[enc.shardCount] = enc.shardCache[enc.shardCount][:sz] - copy(enc.shardCache[enc.shardCount], b) + copy(enc.shardCache[enc.shardCount][enc.payloadOffset:], b[enc.payloadOffset:]) enc.shardCount++ // track max datashard length @@ -285,7 +281,7 @@ func (enc *fecEncoder) encode(b []byte) (ps [][]byte) { if err := enc.codec.Encode(cache); err == nil { ps = enc.shardCache[enc.dataShards:] for k := range ps { - enc.markFEC(ps[k][enc.headerOffset:]) + enc.markParity(ps[k][enc.headerOffset:]) ps[k] = ps[k][:enc.maxSize] } } @@ -304,8 +300,9 @@ func (enc *fecEncoder) markData(data []byte) { enc.next++ } -func (enc *fecEncoder) markFEC(data []byte) { +func (enc *fecEncoder) markParity(data []byte) { binary.LittleEndian.PutUint32(data, enc.next) - binary.LittleEndian.PutUint16(data[4:], typeFEC) + binary.LittleEndian.PutUint16(data[4:], typeParity) + // sequence wrap will only happen at parity shard enc.next = (enc.next + 1) % enc.paws } diff --git a/kcp-go/kcp.go b/kcp-go/kcp.go index 6bfb04e..c5d6603 100644 --- a/kcp-go/kcp.go +++ b/kcp-go/kcp.go @@ -1,4 +1,3 @@ -// Package kcp - A Fast and Reliable ARQ Protocol package kcp import ( @@ -145,8 +144,9 @@ type KCP struct { acklist []ackItem - buffer []byte - output output_callback + buffer []byte + reserved int + output output_callback } type ackItem struct { @@ -154,8 +154,11 @@ type ackItem struct { ts uint32 } -// NewKCP create a new kcp control object, 'conv' must equal in two endpoint -// from the same connection. +// NewKCP create a new kcp state machine +// +// 'conv' must be equal in the connection peers, or else data will be silently rejected. +// +// 'output' function will be called whenever these is data to be sent on wire. func NewKCP(conv uint32, output output_callback) *KCP { kcp := new(KCP) kcp.conv = conv @@ -164,7 +167,7 @@ func NewKCP(conv uint32, output output_callback) *KCP { kcp.rmt_wnd = IKCP_WND_RCV kcp.mtu = IKCP_MTU_DEF kcp.mss = kcp.mtu - IKCP_OVERHEAD - kcp.buffer = make([]byte, (kcp.mtu+IKCP_OVERHEAD)*3) + kcp.buffer = make([]byte, kcp.mtu) kcp.rx_rto = IKCP_RTO_DEF kcp.rx_minrto = IKCP_RTO_MIN kcp.interval = IKCP_INTERVAL @@ -189,6 +192,19 @@ func (kcp *KCP) delSegment(seg *segment) { } } +// ReserveBytes keeps n bytes untouched from the beginning of the buffer, +// the output_callback function should be aware of this. +// +// Return false if n >= mss +func (kcp *KCP) ReserveBytes(n int) bool { + if n >= int(kcp.mtu-IKCP_OVERHEAD) || n < 0 { + return false + } + kcp.reserved = n + kcp.mss = kcp.mtu - IKCP_OVERHEAD - uint32(n) + return true +} + // PeekSize checks the size of next message in the recv queue func (kcp *KCP) PeekSize() (length int) { if len(kcp.rcv_queue) == 0 { @@ -214,19 +230,21 @@ func (kcp *KCP) PeekSize() (length int) { return } -// Recv is user/upper level recv: returns size, returns below zero for EAGAIN +// Receive data from kcp state machine +// +// Return number of bytes read. +// +// Return -1 when there is no readable data. +// +// Return -2 if len(buffer) is smaller than kcp.PeekSize(). func (kcp *KCP) Recv(buffer []byte) (n int) { - if len(kcp.rcv_queue) == 0 { + peeksize := kcp.PeekSize() + if peeksize < 0 { return -1 } - peeksize := kcp.PeekSize() - if peeksize < 0 { - return -2 - } - if peeksize > len(buffer) { - return -3 + return -2 } var fast_recover bool @@ -386,6 +404,10 @@ func (kcp *KCP) parse_ack(sn uint32) { for k := range kcp.snd_buf { seg := &kcp.snd_buf[k] if sn == seg.sn { + // mark and free space, but leave the segment here, + // and wait until `una` to delete this, then we don't + // have to shift the segments behind forward, + // which is an expensive operation for large window seg.acked = 1 kcp.delSegment(seg) break @@ -489,8 +511,12 @@ func (kcp *KCP) parse_data(newseg segment) bool { return repeat } -// Input when you received a low level packet (eg. UDP packet), call it -// regular indicates a regular packet has received(not from FEC) +// Input a packet into kcp state machine. +// +// 'regular' indicates it's a real data packet from remote, and it means it's not generated from ReedSolomon +// codecs. +// +// 'ackNoDelay' will trigger immediate ACK, but surely it will not be efficient in bandwidth func (kcp *KCP) Input(data []byte, regular, ackNoDelay bool) int { snd_una := kcp.snd_una if len(data) < IKCP_OVERHEAD { @@ -634,14 +660,28 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { seg.una = kcp.rcv_nxt buffer := kcp.buffer - // flush acknowledges - ptr := buffer - for i, ack := range kcp.acklist { + ptr := buffer[kcp.reserved:] // keep n bytes untouched + + // makeSpace makes room for writing + makeSpace := func(space int) { size := len(buffer) - len(ptr) - if size+IKCP_OVERHEAD > int(kcp.mtu) { + if size+space > int(kcp.mtu) { kcp.output(buffer, size) - ptr = buffer + ptr = buffer[kcp.reserved:] } + } + + // flush bytes in buffer if there is any + flushBuffer := func() { + size := len(buffer) - len(ptr) + if size > kcp.reserved { + kcp.output(buffer, size) + } + } + + // flush acknowledges + for i, ack := range kcp.acklist { + makeSpace(IKCP_OVERHEAD) // filter jitters caused by bufferbloat if ack.sn >= kcp.rcv_nxt || len(kcp.acklist)-1 == i { seg.sn, seg.ts = ack.sn, ack.ts @@ -651,10 +691,7 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { kcp.acklist = kcp.acklist[0:0] if ackOnly { // flash remain ack segments - size := len(buffer) - len(ptr) - if size > 0 { - kcp.output(buffer, size) - } + flushBuffer() return kcp.interval } @@ -685,22 +722,14 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { // flush window probing commands if (kcp.probe & IKCP_ASK_SEND) != 0 { seg.cmd = IKCP_CMD_WASK - size := len(buffer) - len(ptr) - if size+IKCP_OVERHEAD > int(kcp.mtu) { - kcp.output(buffer, size) - ptr = buffer - } + makeSpace(IKCP_OVERHEAD) ptr = seg.encode(ptr) } // flush window probing commands if (kcp.probe & IKCP_ASK_TELL) != 0 { seg.cmd = IKCP_CMD_WINS - size := len(buffer) - len(ptr) - if size+IKCP_OVERHEAD > int(kcp.mtu) { - kcp.output(buffer, size) - ptr = buffer - } + makeSpace(IKCP_OVERHEAD) ptr = seg.encode(ptr) } @@ -779,20 +808,14 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { } if needsend { - current = currentMs() // time update for a blocking call + current = currentMs() segment.xmit++ segment.ts = current segment.wnd = seg.wnd segment.una = seg.una - size := len(buffer) - len(ptr) need := IKCP_OVERHEAD + len(segment.data) - - if size+need > int(kcp.mtu) { - kcp.output(buffer, size) - ptr = buffer - } - + makeSpace(need) ptr = segment.encode(ptr) copy(ptr, segment.data) ptr = ptr[len(segment.data):] @@ -809,10 +832,7 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { } // flash remain segments - size := len(buffer) - len(ptr) - if size > 0 { - kcp.output(buffer, size) - } + flushBuffer() // counter updates sum := lostSegs @@ -864,6 +884,8 @@ func (kcp *KCP) flush(ackOnly bool) uint32 { return uint32(minrto) } +// (deprecated) +// // Update updates state (call it repeatedly, every 10ms-100ms), or you can ask // ikcp_check when to call it again (without ikcp_input/_send calling). // 'current' - current timestamp in millisec. @@ -892,6 +914,8 @@ func (kcp *KCP) Update() { } } +// (deprecated) +// // Check determines when should you invoke ikcp_update: // returns when you should invoke ikcp_update in millisec, if there // is no ikcp_input/_send calling. you can call ikcp_update in that @@ -947,12 +971,16 @@ func (kcp *KCP) SetMtu(mtu int) int { if mtu < 50 || mtu < IKCP_OVERHEAD { return -1 } - buffer := make([]byte, (mtu+IKCP_OVERHEAD)*3) + if kcp.reserved >= int(kcp.mtu-IKCP_OVERHEAD) || kcp.reserved < 0 { + return -1 + } + + buffer := make([]byte, mtu) if buffer == nil { return -2 } kcp.mtu = uint32(mtu) - kcp.mss = kcp.mtu - IKCP_OVERHEAD + kcp.mss = kcp.mtu - IKCP_OVERHEAD - uint32(kcp.reserved) kcp.buffer = buffer return 0 } @@ -1006,7 +1034,13 @@ func (kcp *KCP) WaitSnd() int { } // remove front n elements from queue +// if the number of elements to remove is more than half of the size. +// just shift the rear elements to front, otherwise just reslice q to q[n:] +// then the cost of runtime.growslice can always be less than n/2 func (kcp *KCP) remove_front(q []segment, n int) []segment { - newn := copy(q, q[n:]) - return q[:newn] + if n > cap(q)/2 { + newn := copy(q, q[n:]) + return q[:newn] + } + return q[n:] } diff --git a/kcp-go/readloop_generic.go b/kcp-go/readloop_generic.go new file mode 100755 index 0000000..9f67d42 --- /dev/null +++ b/kcp-go/readloop_generic.go @@ -0,0 +1,50 @@ +// +build !linux + +package kcp + +import ( + "sync/atomic" + + "github.com/pkg/errors" +) + +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 { + s.packetInput(buf[:n]) + } else { + atomic.AddUint64(&DefaultSnmp.InErrs, 1) + } + } else { + s.notifyReadError(errors.WithStack(err)) + return + } + } +} + +func (l *Listener) monitor() { + buf := make([]byte, mtuLimit) + for { + if n, from, err := l.conn.ReadFrom(buf); err == nil { + if n >= l.headerSize+IKCP_OVERHEAD { + l.packetInput(buf[:n], from) + } else { + atomic.AddUint64(&DefaultSnmp.InErrs, 1) + } + } else { + l.notifyReadError(errors.WithStack(err)) + return + } + } +} diff --git a/kcp-go/readloop_linux.go b/kcp-go/readloop_linux.go new file mode 100755 index 0000000..cc51ffe --- /dev/null +++ b/kcp-go/readloop_linux.go @@ -0,0 +1,80 @@ +// +build linux + +package kcp + +import ( + "net" + "sync/atomic" + + "github.com/pkg/errors" + "golang.org/x/net/ipv4" + "golang.org/x/net/ipv6" +) + +// the read loop for a client session +func (s *UDPSession) readLoop() { + var src string + msgs := make([]ipv4.Message, batchSize) + for k := range msgs { + msgs[k].Buffers = [][]byte{make([]byte, mtuLimit)} + } + + for { + if count, err := s.xconn.ReadBatch(msgs, 0); err == nil { + for i := 0; i < count; i++ { + msg := &msgs[i] + // make sure the packet is from the same source + if src == "" { // set source address if nil + src = msg.Addr.String() + s.remote = msg.Addr// added by hikari + } else if msg.Addr.String() != src { + atomic.AddUint64(&DefaultSnmp.InErrs, 1) + continue + } + + if msg.N < s.headerSize+IKCP_OVERHEAD { + atomic.AddUint64(&DefaultSnmp.InErrs, 1) + continue + } + + // source and size has validated + s.packetInput(msg.Buffers[0][:msg.N]) + } + } else { + s.notifyReadError(errors.WithStack(err)) + return + } + } +} + +// monitor incoming data for all connections of server +func (l *Listener) monitor() { + addr, _ := net.ResolveUDPAddr("udp", l.conn.LocalAddr().String()) + var xconn batchConn + if addr.IP.To4() != nil { + xconn = ipv4.NewPacketConn(l.conn) + } else { + xconn = ipv6.NewPacketConn(l.conn) + } + + msgs := make([]ipv4.Message, batchSize) + for k := range msgs { + msgs[k].Buffers = [][]byte{make([]byte, mtuLimit)} + } + + for { + if count, err := xconn.ReadBatch(msgs, 0); err == nil { + for i := 0; i < count; i++ { + msg := &msgs[i] + if msg.N >= l.headerSize+IKCP_OVERHEAD { + l.packetInput(msg.Buffers[0][:msg.N], msg.Addr) + } else { + atomic.AddUint64(&DefaultSnmp.InErrs, 1) + } + } + } else { + l.notifyReadError(errors.WithStack(err)) + return + } + } +} diff --git a/kcp-go/sess.go b/kcp-go/sess.go index 4bded00..91e9fd3 100755 --- a/kcp-go/sess.go +++ b/kcp-go/sess.go @@ -1,10 +1,17 @@ +// Package kcp-go is a Reliable-UDP library for golang. +// +// This library intents to provide a smooth, resilient, ordered, +// error-checked and anonymous delivery of streams over UDP packets. +// +// The interfaces of this package aims to be compatible with +// net.Conn in standard library, but offers powerful features for advanced users. package kcp import ( "crypto/rand" "encoding/binary" "hash/crc32" - "log" + "io" "net" "sync" "sync/atomic" @@ -15,14 +22,6 @@ import ( "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 @@ -43,9 +42,10 @@ const ( acceptBacklog = 128 ) -const ( - errBrokenPipe = "broken pipe" - errInvalidOperation = "invalid operation" +var ( + errInvalidOperation = errors.New("invalid operation") + errTimeout = errors.New("timeout") + firstPacket = true ) var ( @@ -73,8 +73,6 @@ type ( // recvbuf turns packets into stream recvbuf []byte bufptr []byte - // header extended output buffer, if has header - ext []byte // FEC codec fecDecoder *fecDecoder @@ -91,16 +89,26 @@ type ( // 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 - chReadError chan error // notify PacketConn.Read() have an error - chWriteError chan error // notify PacketConn.Write() have an error + + // socket error handling + socketReadError atomic.Value + socketWriteError atomic.Value + chSocketReadError chan struct{} + chSocketWriteError chan struct{} + socketReadErrorOnce sync.Once + socketWriteErrorOnce sync.Once // nonce generator nonce Entropy - isClosed bool // flag the session has Closed - mu sync.Mutex + // packets waiting to be sent on wire + txqueue []ipv4.Message + xconn batchConn // for x/net + + mu sync.Mutex } setReadBuffer interface { @@ -120,14 +128,22 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn 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.chSocketReadError = make(chan struct{}) + sess.chSocketWriteError = make(chan struct{}) sess.remote = remote sess.conn = conn sess.l = l sess.block = block sess.recvbuf = make([]byte, mtuLimit) + // cast to writebatch conn + addr, _ := net.ResolveUDPAddr("udp", conn.LocalAddr().String()) + if addr.IP.To4() != nil { + sess.xconn = ipv4.NewPacketConn(conn) + } else { + sess.xconn = ipv6.NewPacketConn(conn) + } + // FEC codec initialization sess.fecDecoder = newFECDecoder(rxFECMulti*(dataShards+parityShards), dataShards, parityShards) if sess.block != nil { @@ -144,17 +160,12 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn 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 { + if size >= IKCP_OVERHEAD+sess.headerSize { sess.output(buf[:size]) } }) - sess.kcp.SetMtu(IKCP_MTU_DEF - sess.headerSize) + sess.kcp.ReserveBytes(sess.headerSize) // register current session to the global updater, // which call sess.update() periodically. @@ -166,6 +177,7 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn } else { atomic.AddUint64(&DefaultSnmp.PassiveOpens, 1) } + currestab := atomic.AddUint64(&DefaultSnmp.CurrEstab, 1) maxconn := atomic.LoadUint64(&DefaultSnmp.MaxConn) if currestab > maxconn { @@ -187,11 +199,6 @@ func (s *UDPSession) Read(b []byte) (n int, err error) { 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) @@ -221,7 +228,7 @@ func (s *UDPSession) Read(b []byte) (n int, err error) { if !s.rd.IsZero() { if time.Now().After(s.rd) { s.mu.Unlock() - return 0, errTimeout{} + return 0, errors.WithStack(errTimeout) } delay := s.rd.Sub(time.Now()) @@ -230,63 +237,66 @@ func (s *UDPSession) Read(b []byte) (n int, err error) { } s.mu.Unlock() - // wait for read event or timeout + // wait for read event or timeout or error 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() + 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) { +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) { for { - s.mu.Lock() - if s.isClosed { - s.mu.Unlock() - return 0, errors.New(errBrokenPipe) + select { + case <-s.chSocketWriteError: + return 0, s.socketWriteError.Load().(error) + case <-s.die: + return 0, errors.WithStack(io.ErrClosedPipe) + default: } - // controls how much data will be sent to kcp core - // to prevent the memory from exhuasting + s.mu.Lock() 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:] + for _, b := range v { + 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.uncork() } 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{} + return 0, errors.WithStack(errTimeout) } delay := s.wd.Sub(time.Now()) timeout = time.NewTimer(delay) @@ -294,44 +304,52 @@ func (s *UDPSession) Write(b []byte) (n int, err error) { } 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() + 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) } } } +// uncork sends data in txqueue if there is any +func (s *UDPSession) uncork() { + if len(s.txqueue) > 0 { + s.tx(s.txqueue) + s.txqueue = s.txqueue[:0] + } + return +} + // 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) - } + var once bool + s.dieOnce.Do(func() { + close(s.die) + once = true + }) - s.mu.Lock() - defer s.mu.Unlock() - if s.isClosed { - return errors.New(errBrokenPipe) + if once { + // remove from updater + updater.removeSession(s) + atomic.AddUint64(&DefaultSnmp.CurrEstab, ^uint64(0)) + + if s.l != nil { // belongs to listener + s.l.closeSession(s.remote) + return nil + } else { // client socket close + return s.conn.Close() + } + } else { + return errors.WithStack(io.ErrClosedPipe) } - 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. @@ -391,7 +409,7 @@ func (s *UDPSession) SetMtu(mtu int) bool { s.mu.Lock() defer s.mu.Unlock() - s.kcp.SetMtu(mtu - s.headerSize) + s.kcp.SetMtu(mtu) return true } @@ -413,7 +431,9 @@ func (s *UDPSession) SetACKNoDelay(nodelay bool) { s.ackNoDelay = nodelay } -// SetDUP duplicates udp packets for kcp output, for testing purpose only +// (deprecated) +// +// SetDUP duplicates udp packets for kcp output. func (s *UDPSession) SetDUP(dup int) { s.mu.Lock() defer s.mu.Unlock() @@ -428,19 +448,23 @@ func (s *UDPSession) SetNoDelay(nodelay, interval, resend, nc int) { s.kcp.NoDelay(nodelay, interval, resend, nc) } -// SetDSCP sets the 6bit DSCP field of IP header, no effect if it's accepted from Listener +// SetDSCP sets the 6bit DSCP field in IPv4 header, or 8bit Traffic Class in IPv6 header. +// +// 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 { if nc, ok := s.conn.(net.Conn); ok { - if err := ipv4.NewConn(nc).SetTOS(dscp << 2); err != nil { + addr, _ := net.ResolveUDPAddr("udp", nc.LocalAddr().String()) + if addr.IP.To4() != nil { + return ipv4.NewConn(nc).SetTOS(dscp << 2) + } else { return ipv6.NewConn(nc).SetTrafficClass(dscp) } - return nil } } - return errors.New(errInvalidOperation) + return errInvalidOperation } // SetReadBuffer sets the socket read buffer, no effect if it's accepted from Listener @@ -452,7 +476,7 @@ func (s *UDPSession) SetReadBuffer(bytes int) error { return nc.SetReadBuffer(bytes) } } - return errors.New(errInvalidOperation) + return errInvalidOperation } // SetWriteBuffer sets the socket write buffer, no effect if it's accepted from Listener @@ -464,37 +488,29 @@ func (s *UDPSession) SetWriteBuffer(bytes int) error { return nc.SetWriteBuffer(bytes) } } - return errors.New(errInvalidOperation) + return 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 +// 4. TxQueue 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) + ecc = s.fecEncoder.encode(buf) } // 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) + s.nonce.Fill(buf[:nonceSize]) + checksum := crc32.ChecksumIEEE(buf[cryptHeaderSize:]) + binary.LittleEndian.PutUint32(buf[nonceSize:], checksum) + s.block.Encrypt(buf, buf) for k := range ecc { s.nonce.Fill(ecc[k][:nonceSize]) @@ -504,28 +520,23 @@ func (s *UDPSession) output(buf []byte) { } } - // 4. WriteTo kernel - nbytes := 0 - npkts := 0 + // 4. TxQueue + var msg ipv4.Message 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) - } + bts := xmitBuf.Get().([]byte)[:len(buf)] + copy(bts, buf) + msg.Buffers = [][]byte{bts} + msg.Addr = s.remote + s.txqueue = append(s.txqueue, msg) } for k := range ecc { - if n, err := s.conn.WriteTo(ecc[k], s.remote); err == nil { - nbytes += n - npkts++ - } else { - s.notifyWriteError(err) - } + bts := xmitBuf.Get().([]byte)[:len(ecc[k])] + copy(bts, ecc[k]) + msg.Buffers = [][]byte{bts} + msg.Addr = s.remote + s.txqueue = append(s.txqueue, msg) } - atomic.AddUint64(&DefaultSnmp.OutPkts, uint64(npkts)) - atomic.AddUint64(&DefaultSnmp.OutBytes, uint64(nbytes)) } // kcp update, returns interval for next calling @@ -536,6 +547,7 @@ func (s *UDPSession) update() (interval time.Duration) { if s.kcp.WaitSnd() < waitsnd { s.notifyWriteEvent() } + s.uncork() s.mu.Unlock() return } @@ -557,10 +569,39 @@ func (s *UDPSession) notifyWriteEvent() { } } +func (s *UDPSession) notifyReadError(err error) { + s.socketReadErrorOnce.Do(func() { + s.socketReadError.Store(err) + close(s.chSocketReadError) + }) +} + func (s *UDPSession) notifyWriteError(err error) { - select { - case s.chWriteError <- err: - default: + s.socketWriteErrorOnce.Do(func() { + s.socketWriteError.Store(err) + close(s.chSocketWriteError) + }) +} + +// packet input stage +func (s *UDPSession) packetInput(data []byte) { + 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) } } @@ -569,16 +610,16 @@ func (s *UDPSession) kcpInput(data []byte) { 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 { + f := fecPacket(data) + if f.flag() == typeData || f.flag() == typeParity { // header check + if f.flag() == typeParity { fecParityShards++ } recovers := s.fecDecoder.decode(f) s.mu.Lock() waitsnd := s.kcp.WaitSnd() - if f.flag == typeData { + if f.flag() == typeData { if ret := s.kcp.Input(data[fecHeaderSizePlus2:], true, s.ackNoDelay); ret != 0 { kcpInErrors++ } @@ -599,6 +640,8 @@ func (s *UDPSession) kcpInput(data []byte) { } else { fecErrs++ } + // recycle the recovers + xmitBuf.Put(r) } // to notify the readers to receive the data @@ -609,6 +652,7 @@ func (s *UDPSession) kcpInput(data []byte) { if s.kcp.WaitSnd() < waitsnd { s.notifyWriteEvent() } + s.uncork() s.mu.Unlock() } else { atomic.AddUint64(&DefaultSnmp.InErrs, 1) @@ -628,6 +672,7 @@ func (s *UDPSession) kcpInput(data []byte) { if s.kcp.WaitSnd() < waitsnd { s.notifyWriteEvent() } + s.uncork() s.mu.Unlock() } @@ -645,57 +690,7 @@ func (s *UDPSession) kcpInput(data []byte) { 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 - firstPacket := true - 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 { - if firstPacket{ - log.Println("firstPacket valided from", addr.String()) - firstPacket = false - //remote upd ip may change - s.remote = addr - } - s.kcpInput(data) - } - } else { - atomic.AddUint64(&DefaultSnmp.InErrs, 1) - } - } else { - s.chReadError <- err - return - } - } } type ( @@ -712,98 +707,91 @@ type ( 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 + + 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() } ) -// 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) +// packet input stage +func (l *Listener) packetInput(data []byte, addr net.Addr) { + 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 { + l.sessionLock.Lock() + s, ok := l.sessions[addr.String()] + l.sessionLock.Unlock() + + if !ok { // new address:port + 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 if l.block == nil { - dataValid = true + } else { + conv = binary.LittleEndian.Uint32(data) + convValid = 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) - } + 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, addr, l.block) + s.kcpInput(data) + l.sessionLock.Lock() + l.sessions[addr.String()] = s + l.sessionLock.Unlock() + l.chAccepts <- s } - } else { - atomic.AddUint64(&DefaultSnmp.InErrs, 1) } } else { - return + s.kcpInput(data) } } } +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.Lock() + for _, s := range l.sessions { + s.notifyReadError(err) + } + l.sessionLock.Unlock() + }) +} + // 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) + return errInvalidOperation } // SetWriteBuffer sets the socket write buffer for the Listener @@ -811,18 +799,20 @@ func (l *Listener) SetWriteBuffer(bytes int) error { if nc, ok := l.conn.(setWriteBuffer); ok { return nc.SetWriteBuffer(bytes) } - return errors.New(errInvalidOperation) + return errInvalidOperation } -// SetDSCP sets the 6bit DSCP field of IP header +// SetDSCP sets the 6bit DSCP field in IPv4 header, or 8bit Traffic Class in IPv6 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 { + addr, _ := net.ResolveUDPAddr("udp", nc.LocalAddr().String()) + if addr.IP.To4() != nil { + return ipv4.NewConn(nc).SetTOS(dscp << 2) + } else { return ipv6.NewConn(nc).SetTrafficClass(dscp) } - return nil } - return errors.New(errInvalidOperation) + return errInvalidOperation } // Accept implements the Accept method in the Listener interface; it waits for the next call and returns a generic Conn. @@ -839,11 +829,13 @@ func (l *Listener) AcceptKCP() (*UDPSession, error) { select { case <-timeout: - return nil, &errTimeout{} + 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.New(errBrokenPipe) + return nil, errors.WithStack(io.ErrClosedPipe) } } @@ -861,15 +853,21 @@ func (l *Listener) SetReadDeadline(t time.Time) error { } // SetWriteDeadline implements the Conn SetWriteDeadline method. -func (l *Listener) SetWriteDeadline(t time.Time) error { - l.wd.Store(t) - return nil -} +func (l *Listener) SetWriteDeadline(t time.Time) error { return errInvalidOperation } -// Close stops listening on the UDP address. Already Accepted connections are not closed. +// Close stops listening on the UDP address, and closes the socket func (l *Listener) Close() error { - close(l.die) - return l.conn.Close() + var once bool + l.dieOnce.Do(func() { + close(l.die) + once = true + }) + + if once { + return l.conn.Close() + } else { + return errors.WithStack(io.ErrClosedPipe) + } } // closeSession notify the listener that a session has closed @@ -889,16 +887,21 @@ 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, -// dataShards, parityShards defines Reed-Solomon Erasure Coding parameters +// 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' specifiy 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.Wrap(err, "net.ResolveUDPAddr") + return nil, errors.WithStack(err) } conn, err := net.ListenUDP("udp", udpaddr) if err != nil { - return nil, errors.Wrap(err, "net.ListenUDP") + return nil, errors.WithStack(err) } return ServeConn(block, dataShards, parityShards, conn) @@ -916,6 +919,7 @@ func ServeConn(block BlockCrypt, dataShards, parityShards int, conn net.PacketCo l.parityShards = parityShards l.block = block l.fecDecoder = newFECDecoder(rxFECMulti*(dataShards+parityShards), dataShards, parityShards) + l.chSocketReadError = make(chan struct{}) // calculate header size if l.block != nil { @@ -929,15 +933,21 @@ func ServeConn(block BlockCrypt, dataShards, parityShards int, conn net.PacketCo return l, nil } -// Dial connects to the remote address "raddr" on the network "udp" +// 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' specifiy 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.Wrap(err, "net.ResolveUDPAddr") + return nil, errors.WithStack(err) } network := "udp4" if udpaddr.IP.To4() == nil { @@ -946,7 +956,7 @@ func DialWithOptions(raddr string, block BlockCrypt, dataShards, parityShards in conn, err := net.ListenUDP(network, nil) if err != nil { - return nil, errors.Wrap(err, "net.DialUDP") + return nil, errors.WithStack(err) } return NewConn(raddr, block, dataShards, parityShards, conn) @@ -956,7 +966,7 @@ func DialWithOptions(raddr string, block BlockCrypt, dataShards, parityShards in 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") + return nil, errors.WithStack(err) } var convid uint32 @@ -964,6 +974,12 @@ func NewConn(raddr string, block BlockCrypt, dataShards, parityShards int, conn 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) } + func NewP2pConn(udpConn net.PacketConn, raddr string, block BlockCrypt, dataShards, parityShards int) (*UDPSession, error){ udpaddr, err := net.ResolveUDPAddr("udp", raddr) if err != nil { @@ -971,10 +987,4 @@ func NewP2pConn(udpConn net.PacketConn, raddr string, block BlockCrypt, dataShar } return newUDPSession(0x1, dataShards, parityShards, nil, udpConn, 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) } +} \ No newline at end of file diff --git a/kcp-go/tx_generic.go b/kcp-go/tx_generic.go new file mode 100644 index 0000000..e5d09cc --- /dev/null +++ b/kcp-go/tx_generic.go @@ -0,0 +1,27 @@ +// +build !linux + +package kcp + +import ( + "sync/atomic" + + "github.com/pkg/errors" + "golang.org/x/net/ipv4" +) + +func (s *UDPSession) tx(txqueue []ipv4.Message) { + nbytes := 0 + npkts := 0 + for k := range txqueue { + if n, err := s.conn.WriteTo(txqueue[k].Buffers[0], txqueue[k].Addr); err == nil { + nbytes += n + npkts++ + xmitBuf.Put(txqueue[k].Buffers[0]) + } else { + s.notifyWriteError(errors.WithStack(err)) + break + } + } + atomic.AddUint64(&DefaultSnmp.OutPkts, uint64(npkts)) + atomic.AddUint64(&DefaultSnmp.OutBytes, uint64(nbytes)) +} diff --git a/kcp-go/tx_linux.go b/kcp-go/tx_linux.go new file mode 100644 index 0000000..6e03971 --- /dev/null +++ b/kcp-go/tx_linux.go @@ -0,0 +1,31 @@ +// +build linux + +package kcp + +import ( + "sync/atomic" + + "github.com/pkg/errors" + "golang.org/x/net/ipv4" +) + +func (s *UDPSession) tx(txqueue []ipv4.Message) { + nbytes := 0 + npkts := 0 + for len(txqueue) > 0 { + if n, err := s.xconn.WriteBatch(txqueue, 0); err == nil { + for k := range txqueue[:n] { + nbytes += len(txqueue[k].Buffers[0]) + xmitBuf.Put(txqueue[k].Buffers[0]) + } + npkts += n + txqueue = txqueue[n:] + } else { + s.notifyWriteError(errors.WithStack(err)) + break + } + } + + atomic.AddUint64(&DefaultSnmp.OutPkts, uint64(npkts)) + atomic.AddUint64(&DefaultSnmp.OutBytes, uint64(nbytes)) +} diff --git a/kcp-go/updater.go b/kcp-go/updater.go index 9a90c82..b36d1b9 100644 --- a/kcp-go/updater.go +++ b/kcp-go/updater.go @@ -76,10 +76,10 @@ func (h *updateHeap) wakeup() { } func (h *updateHeap) updateTask() { - var timer <-chan time.Time + timer := time.NewTimer(0) for { select { - case <-timer: + case <-timer.C: case <-h.chWakeUp: } @@ -87,7 +87,7 @@ func (h *updateHeap) updateTask() { hlen := h.Len() for i := 0; i < hlen; i++ { entry := &h.entries[0] - if time.Now().After(entry.ts) { + if !time.Now().Before(entry.ts) { interval := entry.s.update() entry.ts = time.Now().Add(interval) heap.Fix(h, 0) @@ -97,7 +97,7 @@ func (h *updateHeap) updateTask() { } if hlen > 0 { - timer = time.After(h.entries[0].ts.Sub(time.Now())) + timer.Reset(h.entries[0].ts.Sub(time.Now())) } h.mu.Unlock() } diff --git a/p2pclient/config.go b/p2pclient/config.go index 32effb7..b266d12 100755 --- a/p2pclient/config.go +++ b/p2pclient/config.go @@ -45,3 +45,4 @@ func parseJSONConfig(config *Config, path string) error { return json.NewDecoder(file).Decode(config) } + diff --git a/p2pclient/main.go b/p2pclient/main.go index faa8ee7..4aba2ed 100755 --- a/p2pclient/main.go +++ b/p2pclient/main.go @@ -3,24 +3,20 @@ package main import ( "bufio" "crypto/sha1" - "encoding/csv" "encoding/json" - "fmt" + "github.com/hikaricai/p2p_tun/kcp-go" + "github.com/xtaci/kcptun/generic" "io" "log" "math/rand" "net" "os" + "sync" "sync/atomic" "time" - "golang.org/x/crypto/pbkdf2" - - "github.com/hikaricai/p2p_tun/kcp-go" "github.com/urfave/cli" "github.com/xtaci/smux" - - "path/filepath" ) var ( @@ -29,6 +25,8 @@ var ( // SALT is use for pbkdf2 key expansion SALT = "kcp-go" isServer = false + pingInterval = 10 + xmitBuf sync.Pool ) func handleLocalTcp(sess *smux.Session, p1 io.ReadWriteCloser, quiet bool) { @@ -36,7 +34,6 @@ func handleLocalTcp(sess *smux.Session, p1 io.ReadWriteCloser, quiet bool) { log.Println("stream opened") defer log.Println("stream closed") } - defer p1.Close() p2, err := sess.OpenStream() if err != nil { @@ -44,19 +41,20 @@ func handleLocalTcp(sess *smux.Session, p1 io.ReadWriteCloser, quiet bool) { } defer p2.Close() - // start tunnel - p1die := make(chan struct{}) - buf1 := make([]byte, 65535) - go func() { io.CopyBuffer(p1, p2, buf1); close(p1die) }() + streamCopy := func(dst io.Writer, src io.ReadCloser) chan struct{} { + die := make(chan struct{}) + go func() { + buf := xmitBuf.Get().([]byte) + generic.CopyBuffer(dst, src, buf) + xmitBuf.Put(buf) + close(die) + }() + return die + } - p2die := make(chan struct{}) - buf2 := make([]byte, 65535) - go func() { io.CopyBuffer(p2, p1, buf2); close(p2die) }() - - // wait for tunnel termination select { - case <-p1die: - case <-p2die: + case <-streamCopy(p1, p2): + case <-streamCopy(p2, p1): } } @@ -73,6 +71,9 @@ func main() { // add more log flags for debugging log.SetFlags(log.LstdFlags | log.Lshortfile) } + xmitBuf.New = func() interface{} { + return make([]byte, 32768) + } myApp := cli.NewApp() myApp.Name = "p2pclient" myApp.Usage = "client(with kcptun)" @@ -93,11 +94,6 @@ func main() { Value: "127.0.0.1:4000", Usage: "kcp server address", }, - cli.StringFlag{ - Name: "bindudp, b", - Value: ":29900", - Usage: "bind local udp", - }, cli.StringFlag{ Name: "key, k", Value: "1234", @@ -111,7 +107,7 @@ func main() { }, cli.StringFlag{ Name: "crypt", - Value: "aes", + Value: "none", Usage: "aes, aes-128, aes-192, salsa20, blowfish, twofish, cast5, 3des, tea, xtea, xor, sm4, none", }, cli.StringFlag{ @@ -146,12 +142,12 @@ func main() { }, cli.IntFlag{ Name: "datashard,ds", - Value: 10, + Value: 0, Usage: "set reed-solomon erasure coding - datashard", }, cli.IntFlag{ Name: "parityshard,ps", - Value: 3, + Value: 0, Usage: "set reed-solomon erasure coding - parityshard", }, cli.IntFlag{ @@ -228,12 +224,10 @@ func main() { config.ListenTcp = c.String("listentcp") config.RemoteUdp = c.String("remoteudp") config.TargetTcp = c.String("targettcp") - config.BindUdp = c.String("bindudp") config.Key = c.String("key") config.Passwd = c.String("passwd") config.Crypt = c.String("crypt") config.Mode = c.String("mode") - config.AutoExpire = c.Int("autoexpire") config.MTU = c.Int("mtu") config.SndWnd = c.Int("sndwnd") config.RcvWnd = c.Int("rcvwnd") @@ -248,9 +242,6 @@ func main() { config.NoCongestion = c.Int("nc") config.SockBuf = c.Int("sockbuf") config.KeepAlive = c.Int("keepalive") - config.Log = c.String("log") - config.SnmpLog = c.String("snmplog") - config.SnmpPeriod = c.Int("snmpperiod") config.Quiet = c.Bool("quiet") if c.String("c") != "" { @@ -258,14 +249,6 @@ func main() { checkError(err) } - // log redirect - if config.Log != "" { - f, err := os.OpenFile(config.Log, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) - checkError(err) - defer f.Close() - log.SetOutput(f) - } - switch config.Mode { case "normal": config.NoDelay, config.Interval, config.Resend, config.NoCongestion = 0, 40, 2, 1 @@ -279,23 +262,6 @@ func main() { log.Println("version:", VERSION) - log.Println("encryption:", config.Crypt) - log.Println("nodelay parameters:", config.NoDelay, config.Interval, config.Resend, config.NoCongestion) - log.Println("remote udp address:", config.RemoteUdp) - log.Println("sndwnd:", config.SndWnd, "rcvwnd:", config.RcvWnd) - log.Println("compression:", !config.NoComp) - log.Println("mtu:", config.MTU) - log.Println("datashard:", config.DataShard, "parityshard:", config.ParityShard) - log.Println("acknodelay:", config.AckNodelay) - log.Println("dscp:", config.DSCP) - log.Println("sockbuf:", config.SockBuf) - log.Println("keepalive:", config.KeepAlive) - log.Println("autoexpire:", config.AutoExpire) - log.Println("snmplog:", config.SnmpLog) - log.Println("snmpperiod:", config.SnmpPeriod) - log.Println("quiet:", config.Quiet) - go snmpLogger(config.SnmpLog, config.SnmpPeriod) - chTCPConn := make(chan *net.TCPConn, 16) go tcpListener(chTCPConn, &config) @@ -305,7 +271,7 @@ func main() { if err == nil{ p2pHandle(&config, peerAddr, chTCPConn) }else{ - time.Sleep(5*time.Second) + time.Sleep(time.Duration(pingInterval)*time.Second) } } } @@ -313,9 +279,9 @@ func main() { } func p2pHandle(config *Config, peerAddr string, chTCPConn chan *net.TCPConn){ - udpAddr, err := net.ResolveUDPAddr("udp", config.BindUdp) + udpAddr, err := net.ResolveUDPAddr("udp4", config.BindUdp) checkError(err) - udpconn, err := net.ListenUDP("udp", udpAddr) + udpconn, err := net.ListenUDP("udp4", udpAddr) checkError(err) defer udpconn.Close() @@ -339,9 +305,9 @@ func p2pHandle(config *Config, peerAddr string, chTCPConn chan *net.TCPConn){ } func tcpListener(chTCPConn chan *net.TCPConn, config *Config){ - listenTcpAddr, err := net.ResolveTCPAddr("tcp", config.ListenTcp) + listenTcpAddr, err := net.ResolveTCPAddr("tcp4", config.ListenTcp) checkError(err) - listener, err := net.ListenTCP("tcp", listenTcpAddr) + listener, err := net.ListenTCP("tcp4", listenTcpAddr) checkError(err) log.Println("listening on:", listener.Addr()) for{ @@ -355,11 +321,12 @@ func tcpListener(chTCPConn chan *net.TCPConn, config *Config){ } func getPeerAddr(config *Config)(string, error){ - - udpAddr, err := net.ResolveUDPAddr("udp", config.BindUdp) + udpAddr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0") checkError(err) - udpconn, err := net.ListenUDP("udp", udpAddr) + udpconn, err := net.ListenUDP("udp4", udpAddr) checkError(err) + config.BindUdp = udpconn.LocalAddr().String() + log.Println("config.BindUdp is ", config.BindUdp) defer udpconn.Close() kcpConn, err := newKcpConn(udpconn, config, config.RemoteUdp) @@ -375,14 +342,12 @@ func getPeerAddr(config *Config)(string, error){ finMess := phaseJsonMess("fin", "good bye") var peerAddr string log.Println("writing mess") - n, err := kcpConn.Write(pairMess) + _, err = kcpConn.Write(pairMess) if err != nil { log.Println("kcpConn.Write", err) return "", err } - log.Println("writen ", n) for { - pair_s := false log.Println("waiting for server") line, err := reader.ReadString('\n') if err != nil { @@ -397,10 +362,12 @@ func getPeerAddr(config *Config)(string, error){ atomic.StoreInt32(&dataReady, 1) log.Println("rcv ping") case "pair_s": - pair_s = true - fallthrough + isServer = true + peerAddr = mess.Data + log.Println("peer addr is ", peerAddr) + kcpConn.Write(finMess) case "pair_c": - isServer = pair_s + isServer = false peerAddr = mess.Data log.Println("peer addr is ", peerAddr) kcpConn.Write(finMess) @@ -411,21 +378,18 @@ func getPeerAddr(config *Config)(string, error){ } func pingCheck(conn *kcp.UDPSession, dataReady *int32, chPing chan struct {}){ - tickerDie := time.NewTicker(30*time.Second) - defer tickerDie.Stop() jsonPingMess := phaseJsonMess("ping", "hello") - tickerPing := time.NewTicker(10 * time.Second) + tickerPing := time.NewTicker(time.Duration(pingInterval)*time.Second) defer tickerPing.Stop() defer log.Println("pingCheck return") for { select { - case <-tickerDie.C: + case <-tickerPing.C: if !atomic.CompareAndSwapInt32(dataReady, 1, 0) { log.Println("ping timeout") conn.Close() return } - case <-tickerPing.C: conn.Write(jsonPingMess) case <- chPing: return @@ -448,7 +412,6 @@ func phaseJsonMess(cmd string, data string) []byte { } func newSmuxSession(udpconn net.PacketConn, config *Config, remoteAddr string) (*smux.Session, error) { - kcpconn, err := newKcpConn(udpconn, config, remoteAddr) if err != nil { return nil, err @@ -456,9 +419,9 @@ func newSmuxSession(udpconn net.PacketConn, config *Config, remoteAddr string) ( smuxConfig := smux.DefaultConfig() smuxConfig.MaxReceiveBuffer = config.SockBuf smuxConfig.KeepAliveInterval = time.Duration(config.KeepAlive) * time.Second + smuxConfig.KeepAliveTimeout = time.Duration(config.KeepAlive) * time.Second*3 // stream multiplex var smuxSession *smux.Session - if isServer { smuxSession, err = smux.Server(kcpconn, smuxConfig) } else { @@ -549,54 +512,22 @@ func handleTargetTcp(addr string, session *smux.Session, quiet bool) { defer p1.Close() defer p2.Close() - // start tunnel - p1die := make(chan struct{}) - buf1 := make([]byte, 65535) - go func() { io.CopyBuffer(p1, p2, buf1); close(p1die) }() + streamCopy := func(dst io.Writer, src io.ReadCloser) chan struct{} { + die := make(chan struct{}) + go func() { + buf := xmitBuf.Get().([]byte) + generic.CopyBuffer(dst, src, buf) + xmitBuf.Put(buf) + close(die) + }() + return die + } - p2die := make(chan struct{}) - buf2 := make([]byte, 65535) - go func() { io.CopyBuffer(p2, p1, buf2); close(p2die) }() - - // wait for tunnel termination select { - case <-p1die: - case <-p2die: + case <-streamCopy(p1, p2): + case <-streamCopy(p2, p1): } }() } } -func snmpLogger(path string, interval int) { - if path == "" || interval == 0 { - return - } - ticker := time.NewTicker(time.Duration(interval) * time.Second) - defer ticker.Stop() - for { - select { - case <-ticker.C: - // split path into dirname and filename - logdir, logfile := filepath.Split(path) - // only format logfile - f, err := os.OpenFile(logdir+time.Now().Format(logfile), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) - if err != nil { - log.Println(err) - return - } - w := csv.NewWriter(f) - // write header in empty file - if stat, err := f.Stat(); err == nil && stat.Size() == 0 { - if err := w.Write(append([]string{"Unix"}, kcp.DefaultSnmp.Header()...)); err != nil { - log.Println(err) - } - } - if err := w.Write(append([]string{fmt.Sprint(time.Now().Unix())}, kcp.DefaultSnmp.ToSlice()...)); err != nil { - log.Println(err) - } - kcp.DefaultSnmp.Reset() - w.Flush() - f.Close() - } - } -} diff --git a/p2pserver/main.go b/p2pserver/main.go index 8092364..4c41d0e 100755 --- a/p2pserver/main.go +++ b/p2pserver/main.go @@ -3,9 +3,7 @@ package main import ( "bufio" "crypto/sha1" - "encoding/csv" "encoding/json" - "fmt" "github.com/hikaricai/p2p_tun/kcp-go" "log" "math/rand" @@ -16,8 +14,6 @@ import ( "golang.org/x/crypto/pbkdf2" - "path/filepath" - "github.com/urfave/cli" ) @@ -26,6 +22,7 @@ var ( VERSION = "SELFBUILD" // SALT is use for pbkdf2 key expansion SALT = "kcp-go" + pingInterval = 10 ) func checkError(err error) { @@ -48,7 +45,7 @@ func main() { myApp.Flags = []cli.Flag{ cli.StringFlag{ Name: "listen,l", - Value: ":4000", + Value: "0.0.0.0:4000", Usage: "kcp server listen address", }, cli.StringFlag{ @@ -59,7 +56,7 @@ func main() { }, cli.StringFlag{ Name: "crypt", - Value: "aes", + Value: "none", Usage: "aes, aes-128, aes-192, salsa20, blowfish, twofish, cast5, 3des, tea, xtea, xor, sm4, none", }, cli.StringFlag{ @@ -84,12 +81,12 @@ func main() { }, cli.IntFlag{ Name: "datashard,ds", - Value: 10, + Value: 0, Usage: "set reed-solomon erasure coding - datashard", }, cli.IntFlag{ Name: "parityshard,ps", - Value: 3, + Value: 0, Usage: "set reed-solomon erasure coding - parityshard", }, cli.IntFlag{ @@ -185,7 +182,6 @@ func main() { config.NoCongestion = c.Int("nc") config.SockBuf = c.Int("sockbuf") config.KeepAlive = c.Int("keepalive") - config.Log = c.String("log") config.SnmpLog = c.String("snmplog") config.SnmpPeriod = c.Int("snmpperiod") config.Pprof = c.Bool("pprof") @@ -196,15 +192,6 @@ func main() { err := parseJSONConfig(&config, c.String("c")) checkError(err) } - - // log redirect - if config.Log != "" { - f, err := os.OpenFile(config.Log, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) - checkError(err) - defer f.Close() - log.SetOutput(f) - } - switch config.Mode { case "normal": config.NoDelay, config.Interval, config.Resend, config.NoCongestion = 0, 40, 2, 1 @@ -253,20 +240,6 @@ func main() { lis, err := kcp.ListenWithOptions(config.Listen, block, config.DataShard, config.ParityShard) checkError(err) log.Println("listening on:", lis.Addr()) - log.Println("encryption:", config.Crypt) - log.Println("nodelay parameters:", config.NoDelay, config.Interval, config.Resend, config.NoCongestion) - log.Println("sndwnd:", config.SndWnd, "rcvwnd:", config.RcvWnd) - log.Println("compression:", !config.NoComp) - log.Println("mtu:", config.MTU) - log.Println("datashard:", config.DataShard, "parityshard:", config.ParityShard) - log.Println("acknodelay:", config.AckNodelay) - log.Println("dscp:", config.DSCP) - log.Println("sockbuf:", config.SockBuf) - log.Println("keepalive:", config.KeepAlive) - log.Println("snmplog:", config.SnmpLog) - log.Println("snmpperiod:", config.SnmpPeriod) - log.Println("pprof:", config.Pprof) - log.Println("quiet:", config.Quiet) if err := lis.SetDSCP(config.DSCP); err != nil { log.Println("SetDSCP:", err) @@ -278,8 +251,6 @@ func main() { log.Println("SetWriteBuffer:", err) } - go snmpLogger(config.SnmpLog, config.SnmpPeriod) - for { log.Println("listening new kcp") if conn, err := lis.AcceptKCP(); err == nil { @@ -305,28 +276,54 @@ type DigHoleMess struct { Data string } +type UdpSess struct { + conn *kcp.UDPSession + pingFlag int32 + isFin bool +} type P2PSession struct { - addr string - chPair chan string - conn1 *kcp.UDPSession - conn2 *kcp.UDPSession + key string + sess1 *UdpSess + sess2 *UdpSess chFin chan struct{} + mu sync.Mutex } var keymap = make(map[string]*P2PSession) var keymu sync.Mutex - -var session *P2PSession; - +func registSession(session *P2PSession, sess *UdpSess) (ret bool){ + ret = true + session.mu.Lock() + if session.sess1 == nil{ + session.sess1 = sess + }else if session.sess2 == nil{ + session.sess2 = sess + }else{ + ret = false + } + if ret == true && session.sess1 != nil && session.sess2 != nil{ + conn1 := session.sess1.conn + conn2 := session.sess2.conn + addr1 := conn1.RemoteAddr().String() + addr2 := conn2.RemoteAddr().String() + jsonPair1Mess := phaseJsonMess("pair_s", addr2) + jsonPair2Mess := phaseJsonMess("pair_c", addr1) + log.Println("pairing ", addr1,addr2) + conn1.Write(jsonPair1Mess) + conn2.Write(jsonPair2Mess) + } + session.mu.Unlock() + return ret +} func handleClient(conn *kcp.UDPSession) { + var session *P2PSession; + key := "" reader := bufio.NewReader(conn) defer conn.Close() - var dataReady int32 - var chThreadDie = make(chan struct{}) - var ok bool; - defer close(chThreadDie) - go timeout(conn, &dataReady, chThreadDie) + var ok bool + var registed = false + sess := &UdpSess{conn,1,false} for { line, err := reader.ReadString('\n') if err != nil { @@ -342,74 +339,90 @@ func handleClient(conn *kcp.UDPSession) { case "login": remoteAddr := conn.RemoteAddr().String() log.Println("login from ", remoteAddr) - key := mess.Data + key = mess.Data log.Println("key is ", key) - keymu.Lock() session, ok = keymap[key] if ok { - peerAddr := session.addr - log.Println("find peer and addr is", peerAddr) - delete(keymap, key) - session.conn2 = conn - session.chPair <- remoteAddr - jsonPairMess := phaseJsonMess("pair_c", peerAddr) - conn.Write(jsonPairMess) + registed = registSession(session,sess) } else { - log.Println("no peer, registed") - session = &P2PSession{remoteAddr, make(chan string), conn, nil,make(chan struct{})} + log.Println("make new P2PSession") + session = new(P2PSession) + session.key = key + session.sess1 = sess + session.chFin = make(chan struct{}) keymap[key] = session - go p2pSessionHandler(session, chThreadDie) + registed = true + go p2pSessionHandler(session) } keymu.Unlock() - case "ping": - atomic.StoreInt32(&dataReady, 1) jsonPingMess := phaseJsonMess("ping", "hello") conn.Write(jsonPingMess) + case "ping": log.Println("rcv ping from ", conn.RemoteAddr().String()) + if registed == false{ + keymu.Lock() + session, ok = keymap[key] + if ok { + registed = registSession(session,sess) + } else { + log.Println("make new P2PSession") + session = new(P2PSession) + session.key = key + session.sess1 = sess + session.chFin = make(chan struct{}) + keymap[key] = session + registed = true + go p2pSessionHandler(session) + } + keymu.Unlock() + } + atomic.StoreInt32(&sess.pingFlag, 1) + jsonPingMess := phaseJsonMess("ping", "hello") + conn.Write(jsonPingMess) case "fin": log.Println("fin from", conn.RemoteAddr().String()) + sess.isFin = true session.chFin <- struct{}{} } } } -func timeout(conn *kcp.UDPSession, dataReady *int32, chThreadDie chan struct{}){ - tickerDie := time.NewTicker(30*time.Second) - defer tickerDie.Stop() +func p2pSessionHandler(session *P2PSession){ + tickerPing := time.NewTicker(time.Duration(pingInterval)*time.Second) + defer tickerPing.Stop() for { select { - case <-tickerDie.C: - if !atomic.CompareAndSwapInt32(dataReady, 1, 0) { - log.Println("ping timeout") - conn.Close() - return + case <-tickerPing.C: + if session.sess1 != nil && !atomic.CompareAndSwapInt32(&session.sess1.pingFlag, 1, 0) { + session.mu.Lock() + session.sess1.conn.Close() + session.sess1 = nil + log.Println("sess1 ping timeout") + session.mu.Unlock() + } + if session.sess2 != nil && !atomic.CompareAndSwapInt32(&session.sess2.pingFlag, 1, 0) { + session.mu.Lock() + session.sess2.conn.Close() + session.sess2 = nil + log.Println("sess2 ping timeout") + session.mu.Unlock() } - case <-chThreadDie: - return - } - } -} - -func p2pSessionHandler(session *P2PSession, chThreadDie chan struct{}){ - finCnt :=0 - for { - select { - case <-chThreadDie: - return - case peerAddr := <-session.chPair: - jsonPairMess := phaseJsonMess("pair_s", peerAddr) - session.conn1.Write(jsonPairMess) case <-session.chFin: - finCnt++ - if finCnt == 2{ - jsonPairMess := phaseJsonMess("fin", "bye") - session.conn1.Write(jsonPairMess) - session.conn2.Write(jsonPairMess) + if session.sess1 == nil || session.sess2 == nil{ + continue + } + if session.sess1.isFin == true && session.sess2.isFin == true{ + jsonFinMess := phaseJsonMess("fin", "bye") + session.sess1.conn.Write(jsonFinMess) + session.sess2.conn.Write(jsonFinMess) time.Sleep(time.Second) - session.conn1.Close() - session.conn2.Close() + session.sess1.conn.Close() + session.sess2.conn.Close() + keymu.Lock() + delete(keymap,session.key) + keymu.Unlock() return } } @@ -424,37 +437,3 @@ func phaseJsonMess(cmd string, data string) []byte { } return append(jsonMess, '\n') } - -func snmpLogger(path string, interval int) { - if path == "" || interval == 0 { - return - } - ticker := time.NewTicker(time.Duration(interval) * time.Second) - defer ticker.Stop() - for { - select { - case <-ticker.C: - // split path into dirname and filename - logdir, logfile := filepath.Split(path) - // only format logfile - f, err := os.OpenFile(logdir+time.Now().Format(logfile), os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666) - if err != nil { - log.Println(err) - return - } - w := csv.NewWriter(f) - // write header in empty file - if stat, err := f.Stat(); err == nil && stat.Size() == 0 { - if err := w.Write(append([]string{"Unix"}, kcp.DefaultSnmp.Header()...)); err != nil { - log.Println(err) - } - } - if err := w.Write(append([]string{fmt.Sprint(time.Now().Unix())}, kcp.DefaultSnmp.ToSlice()...)); err != nil { - log.Println(err) - } - kcp.DefaultSnmp.Reset() - w.Flush() - f.Close() - } - } -}