safer conversion

This commit is contained in:
xtaci
2019-05-23 23:56:01 +08:00
parent c749e12021
commit 5e990bd411
7 changed files with 115 additions and 72 deletions
+48
View File
@@ -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
}
}
}
+2 -41
View File
@@ -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()
}
+23 -7
View File
@@ -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)
+9 -7
View File
@@ -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
+25
View File
@@ -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))
}
+1 -17
View File
@@ -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)
}
+7
View File
@@ -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 {