fixes for test coverage

Author: zjx20

batchconn_generic.go

@@ -14,15 +14,17 @@ func toBatchConn(c net.PacketConn) batchConn {
 }
 
 func readBatchUnavailable(xconn batchConn, err error) bool {
+	ret := false
 	if detector, ok := xconn.(batchErrDetector); ok {
-		return detector.ReadBatchUnavailable(err)
+		ret = detector.ReadBatchUnavailable(err)
 	}
-	return false
+	return ret
 }
 
 func writeBatchUnavailable(xconn batchConn, err error) bool {
+	ret := false
 	if detector, ok := xconn.(batchErrDetector); ok {
-		return detector.WriteBatchUnavailable(err)
+		ret = detector.WriteBatchUnavailable(err)
 	}
-	return false
+	return ret
 }

batchconn_linux.go

@@ -14,8 +14,8 @@ func toBatchConn(c net.PacketConn) batchConn {
 	if xconn, ok := c.(batchConn); ok {
 		return xconn
 	}
+	var xconn batchConn
 	if _, ok := c.(*net.UDPConn); ok {
-		var xconn batchConn
 		addr, err := net.ResolveUDPAddr("udp", c.LocalAddr().String())
 		if err == nil {
 			if addr.IP.To4() != nil {
@@ -24,9 +24,8 @@ func toBatchConn(c net.PacketConn) batchConn {
 				xconn = ipv6.NewPacketConn(c)
 			}
 		}
-		return xconn
 	}
-	return nil
+	return xconn
 }
 
 func isPacketConn(xconn batchConn) bool {
@@ -53,10 +52,11 @@ func readBatchUnavailable(xconn batchConn, err error) bool {
 		}
 		return false
 	}
+	ret := false
 	if detector, ok := xconn.(batchErrDetector); ok {
-		return detector.ReadBatchUnavailable(err)
+		ret = detector.ReadBatchUnavailable(err)
 	}
-	return false
+	return ret
 }
 
 func writeBatchUnavailable(xconn batchConn, err error) bool {
@@ -73,8 +73,9 @@ func writeBatchUnavailable(xconn batchConn, err error) bool {
 		}
 		return false
 	}
+	ret := false
 	if detector, ok := xconn.(batchErrDetector); ok {
-		return detector.WriteBatchUnavailable(err)
+		ret = detector.WriteBatchUnavailable(err)
 	}
-	return false
+	return ret
 }

readloop.go

@@ -33,13 +33,11 @@ func (s *UDPSession) defaultReadLoop() {
 	for {
 		if n, addr, err := s.conn.ReadFrom(buf); err == nil {
 			// make sure the packet is from the same source
-			if addr.String() != src {
-				if len(src) == 0 { // set source address
-					src = addr.String()
-				} else {
-					atomic.AddUint64(&DefaultSnmp.InErrs, 1)
-					continue
-				}
+			if src == "" { // set source address
+				src = addr.String()
+			} else if addr.String() != src {
+				atomic.AddUint64(&DefaultSnmp.InErrs, 1)
+				continue
 			}
 			s.packetInput(buf[:n])
 		} else {

sess_test.go

@@ -90,6 +90,8 @@ func dialTinyBufferEcho(port int) (*UDPSession, error) {
 }
 
 //////////////////////////
+type listenFn func(port int) (net.Listener, error)
+
 func listenEcho(port int) (net.Listener, error) {
 	//block, _ := NewNoneBlockCrypt(pass)
 	//block, _ := NewSimpleXORBlockCrypt(pass)
@@ -113,7 +115,7 @@ func listenNoEncryption(port int) (net.Listener, error) {
 
 func server(
 	port int,
-	listen func(port int) (net.Listener, error),
+	listen listenFn,
 	handle func(*UDPSession),
 ) net.Listener {
 	l, err := listen(port)
@@ -748,8 +750,25 @@ func (c *customBatchConn) WriteBatchUnavailable(err error) bool {
 	return err.Error() == "unsupported"
 }
 
+func batchListenFn(opt func(pconn *customBatchConn)) listenFn {
+	return func(port int) (net.Listener, error) {
+		udpaddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("127.0.0.1:%v", port))
+		if err != nil {
+			return nil, err
+		}
+		conn, err := net.ListenUDP("udp", udpaddr)
+		if err != nil {
+			return nil, err
+		}
+		pconn := &customBatchConn{UDPConn: conn}
+		opt(pconn)
+		return serveConn(nil, 0, 0, pconn, true)
+	}
+}
+
 func TestCustomBatchConn(t *testing.T) {
-	l := server(0, listenNoEncryption, handleEcho)
+	listen := batchListenFn(func(pconn *customBatchConn) {})
+	l := server(0, listen, handleEcho)
 	defer l.Close()
 
 	// Create a net.PacketConn not owned by the UDPSession.
@@ -794,7 +813,12 @@ func TestCustomBatchConn(t *testing.T) {
 }
 
 func TestCustomBatchConnFallback(t *testing.T) {
-	l := server(0, listenNoEncryption, handleEcho)
+	// should fallback to defaultMonitor()
+	listen := batchListenFn(func(pconn *customBatchConn) {
+		pconn.disableReadBatch = true
+		pconn.disableWriteBatch = true
+	})
+	l := server(0, listen, handleEcho)
 	defer l.Close()
 
 	// Create a net.PacketConn not owned by the UDPSession.