From 5e990bd411b41704e0a94a23bb55f9dd07edf68d Mon Sep 17 00:00:00 2001 From: xtaci Date: Thu, 23 May 2019 23:56:01 +0800 Subject: [PATCH] safer conversion --- readloop.go | 48 +++++++++++++++++++++++++++++++++++++++++++++ readloop_generic.go | 43 ++-------------------------------------- readloop_linux.go | 30 +++++++++++++++++++++------- sess.go | 16 ++++++++------- tx.go | 25 +++++++++++++++++++++++ tx_generic.go | 18 +---------------- tx_linux.go | 7 +++++++ 7 files changed, 115 insertions(+), 72 deletions(-) create mode 100644 readloop.go create mode 100644 tx.go diff --git a/readloop.go b/readloop.go new file mode 100644 index 0000000..bc48e84 --- /dev/null +++ b/readloop.go @@ -0,0 +1,48 @@ +package kcp + +import ( + "sync/atomic" + + "github.com/pkg/errors" +) + +func (s *UDPSession) defaultReadLoop() { + 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) defaultMonitor() { + 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/readloop_generic.go b/readloop_generic.go index 9f67d42..5dbe4f4 100644 --- a/readloop_generic.go +++ b/readloop_generic.go @@ -2,49 +2,10 @@ 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 - } - } + s.defaultReadLoop() } 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 - } - } + l.defaultMonitor() } diff --git a/readloop_linux.go b/readloop_linux.go index c984bd7..1ae6f74 100644 --- a/readloop_linux.go +++ b/readloop_linux.go @@ -13,6 +13,13 @@ import ( // the read loop for a client session func (s *UDPSession) readLoop() { + // default version + if s.xconn == nil { + s.defaultReadLoop() + return + } + + // x/net version var src string msgs := make([]ipv4.Message, batchSize) for k := range msgs { @@ -48,14 +55,23 @@ func (s *UDPSession) readLoop() { // monitor incoming data for all connections of server func (l *Listener) monitor() { - addr, err := net.ResolveUDPAddr("udp", l.conn.LocalAddr().String()) + // default version + if s.xconn == nil { + l.defaultMonitor() + return + } + + // x/net version var xconn batchConn - if err != nil { - xconn = ipv4.NewPacketConn(l.conn) - } else if addr.IP.To4() != nil { - xconn = ipv4.NewPacketConn(l.conn) - } else { - xconn = ipv6.NewPacketConn(l.conn) + if _, ok := l.conn.(*net.UDPConn); ok { + addr, err := net.ResolveUDPAddr("udp", l.conn.LocalAddr().String()) + if err == nil { + if addr.IP.To4() != nil { + xconn = ipv4.NewPacketConn(l.conn) + } else { + xconn = ipv6.NewPacketConn(l.conn) + } + } } msgs := make([]ipv4.Message, batchSize) diff --git a/sess.go b/sess.go index c76c910..48217ae 100644 --- a/sess.go +++ b/sess.go @@ -136,13 +136,15 @@ func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn sess.recvbuf = make([]byte, mtuLimit) // cast to writebatch conn - addr, err := net.ResolveUDPAddr("udp", conn.LocalAddr().String()) - if err != nil { - sess.xconn = ipv4.NewPacketConn(conn) - } else if addr.IP.To4() != nil { - sess.xconn = ipv4.NewPacketConn(conn) - } else { - sess.xconn = ipv6.NewPacketConn(conn) + if _, ok := conn.(*net.UDPConn); ok { + addr, err := net.ResolveUDPAddr("udp", conn.LocalAddr().String()) + if err == nil { + if addr.IP.To4() != nil { + sess.xconn = ipv4.NewPacketConn(conn) + } else { + sess.xconn = ipv6.NewPacketConn(conn) + } + } } // FEC codec initialization diff --git a/tx.go b/tx.go new file mode 100644 index 0000000..9fd0bf4 --- /dev/null +++ b/tx.go @@ -0,0 +1,25 @@ +package kcp + +import ( + "sync/atomic" + + "github.com/pkg/errors" + "golang.org/x/net/ipv4" +) + +func (s *UDPSession) defaultTx(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/tx_generic.go b/tx_generic.go index e5d09cc..0b4f349 100644 --- a/tx_generic.go +++ b/tx_generic.go @@ -3,25 +3,9 @@ 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)) + s.defaultTx(txqueue) } diff --git a/tx_linux.go b/tx_linux.go index 6e03971..b6f2e11 100644 --- a/tx_linux.go +++ b/tx_linux.go @@ -10,6 +10,13 @@ import ( ) func (s *UDPSession) tx(txqueue []ipv4.Message) { + // default version + if s.xconn == nil { + s.defaultTx(txqueue) + return + } + + // x/net version nbytes := 0 npkts := 0 for len(txqueue) > 0 {