Merge pull request #6 from AudriusButkevicius/batch

Add batch read support, update QUIC to RFC9000 (fixes #5)

Author: AudriusButkevicius

conn.go

@@ -4,7 +4,10 @@ import (
 	"io"
 	"net"
 	"sync/atomic"
+	"syscall"
 	"time"
+
+	"golang.org/x/net/ipv4"
 )
 
 type filteredConn struct {
@@ -14,7 +17,7 @@ type filteredConn struct {
 	source   *PacketFilter
 	priority int
 
-	recvBuffer chan packet
+	recvBuffer chan messageWithError
 
 	filter Filter
 
@@ -76,24 +79,113 @@ func (r *filteredConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
 	select {
 	case <-timeout:
 		return 0, nil, errTimeout
-	case pkt := <-r.recvBuffer:
-		n := pkt.n
-		err := pkt.err
-		if l := len(b); l < n {
-			n = l
-			if err == nil {
-				err = io.ErrShortBuffer
-			}
+	case msg := <-r.recvBuffer:
+		n, _, err := copyBuffers(msg, b, nil)
+
+		r.source.returnBuffers(msg.Message)
+
+		return n, msg.Addr, err
+	case <-r.closed:
+		return 0, nil, errClosed
+	}
+}
+
+func (r *filteredConn) ReadBatch(ms []ipv4.Message, flags int) (int, error) {
+	if flags != 0 {
+		return 0, errNotSupported
+	}
+
+	if len(ms) == 0 {
+		return 0, nil
+	}
+
+	var timeout <-chan time.Time
+
+	if deadline, ok := r.deadline.Load().(time.Time); ok && !deadline.IsZero() {
+		timer := time.NewTimer(deadline.Sub(time.Now()))
+		timeout = timer.C
+		defer timer.Stop()
+	}
+
+	msgs := make([]messageWithError, 0, len(ms))
+
+	defer func() {
+		for _, msg := range msgs {
+			r.source.returnBuffers(msg.Message)
 		}
-		copy(b, pkt.buf[:n])
-		r.source.bufPool.Put(pkt.buf[:r.source.packetSize])
-		if pkt.oobBuf != nil {
-			r.source.bufPool.Put(pkt.oobBuf[:r.source.packetSize])
+	}()
+
+	// We must read at least one message.
+	select {
+	//goland:noinspection GoNilness
+	case <-timeout:
+		return 0, errTimeout
+	case msg := <-r.recvBuffer:
+		msgs = append(msgs, msg)
+		if msg.Err != nil {
+			return 0, msg.Err
 		}
-		return n, pkt.addr, err
 	case <-r.closed:
-		return 0, nil, errClosed
+		return 0, errClosed
 	}
+
+	// After that, it's best effort. If there are messages, we read them.
+	// If not, we break out and return what we got.
+loop:
+	for len(msgs) != len(ms) {
+		select {
+		case msg := <-r.recvBuffer:
+			msgs = append(msgs, msg)
+			if msg.Err != nil {
+				return 0, msg.Err
+			}
+		case <-r.closed:
+			return 0, errClosed
+		default:
+			break loop
+		}
+	}
+
+	for i, msg := range msgs {
+		if len(ms[i].Buffers) != 1 {
+			return 0, errNotSupported
+		}
+
+		n, nn, err := copyBuffers(msg, ms[i].Buffers[0], ms[i].OOB)
+		if err != nil {
+			return 0, err
+		}
+
+		ms[i].N = n
+		ms[i].NN = nn
+		ms[i].Flags = msg.Flags
+		ms[i].Addr = msg.Addr
+	}
+
+	return len(msgs), nil
+}
+
+func copyBuffers(msg messageWithError, buf, oobBuf []byte) (n, nn int, err error) {
+	if msg.Err != nil {
+		return 0, 0, msg.Err
+	}
+
+	if len(buf) < msg.N {
+		return 0, 0, io.ErrShortBuffer
+	}
+
+	copy(buf, msg.Buffers[0][:msg.N])
+
+	// Truncate, probably?
+	oobn := msg.NN
+	if oobl := len(oobBuf); oobl < oobn {
+		oobn = oobl
+	}
+	if oobn > 0 {
+		copy(oobBuf, msg.OOB[:oobn])
+	}
+
+	return msg.N, oobn, nil
 }
 
 // Close closes the filtered connection, removing it's filters
@@ -107,3 +199,22 @@ func (r *filteredConn) Close() error {
 	r.source.removeConn(r)
 	return nil
 }
+
+func (r *filteredConn) SetReadBuffer(sz int) error {
+	if srb, ok := r.source.conn.(interface{ SetReadBuffer(int) error }); ok {
+		return srb.SetReadBuffer(sz)
+	}
+	return errNotSupported
+}
+
+func (r *filteredConn) SyscallConn() (syscall.RawConn, error) {
+	if r.source.oobConn != nil {
+		return r.source.oobConn.SyscallConn()
+	}
+	if scon, ok := r.source.conn.(interface {
+		SyscallConn() (syscall.RawConn, error)
+	}); ok {
+		return scon.SyscallConn()
+	}
+	return nil, errNotSupported
+}

conn_oob.go

@@ -1,17 +1,13 @@
 package pfilter
 
 import (
-	"io"
 	"net"
 	"time"
-)
 
-type oobPacketConn interface {
-	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
-	WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
-}
+	"github.com/lucas-clemente/quic-go"
+)
 
-var _ oobPacketConn = (*filteredConnObb)(nil)
+var _ quic.OOBCapablePacketConn = (*filteredConnObb)(nil)
 
 type filteredConnObb struct {
 	*filteredConn
@@ -39,30 +35,17 @@ func (r *filteredConnObb) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *n
 	select {
 	case <-timeout:
 		return 0, 0, 0, nil, errTimeout
-	case pkt := <-r.recvBuffer:
-		err := pkt.err
+	case msg := <-r.recvBuffer:
+		n, nn, err := copyBuffers(msg, b, oob)
 
-		n := pkt.n
-		if l := len(b); l < n {
-			n = l
-			if err == nil {
-				err = io.ErrShortBuffer
-			}
-		}
-		copy(b, pkt.buf[:n])
+		r.source.returnBuffers(msg.Message)
 
-		oobn := pkt.oobn
-		if oobl := len(oob); oobl < oobn {
-			oobn = oobl
-		}
-		if oobn > 0 {
-			copy(oob, pkt.oobBuf[:oobn])
+		udpAddr, ok := msg.Addr.(*net.UDPAddr)
+		if !ok && err == nil {
+			err = errNotSupported
 		}
 
-		r.source.bufPool.Put(pkt.buf[:r.source.packetSize])
-		r.source.bufPool.Put(pkt.oobBuf[:r.source.packetSize])
-
-		return n, oobn, pkt.flags, pkt.udpAddr, err
+		return n, nn, msg.Flags, udpAddr, err
 	case <-r.closed:
 		return 0, 0, 0, nil, errClosed
 	}

conn_test.go

@@ -50,7 +50,6 @@ func BenchmarkPacketConnPfilter(b *testing.B) {
 func benchmark(b *testing.B, client io.Writer, server io.Reader, sz int) {
 	data := make([]byte, sz)
 	if _, err := rand.Read(data); err != nil {
-
 		b.Fatal(err)
 	}
 
@@ -60,17 +59,17 @@ func benchmark(b *testing.B, client io.Writer, server io.Reader, sz int) {
 	for i := 0; i < b.N; i++ {
 		wg.Add(2)
 		go func() {
+			defer wg.Done()
 			if err := sendMsg(client, data); err != nil {
 				b.Fatal(err)
 			}
-			wg.Done()
 		}()
 		go func() {
+			defer wg.Done()
 			if err := recvMsg(server, data); err != nil {
 				b.Fatal(err)
 			}
 			total += sz
-			wg.Done()
 		}()
 		wg.Wait()
 	}
@@ -89,9 +88,12 @@ func (r *readerWrapper) Read(buf []byte) (int, error) {
 
 func sendMsg(c io.Writer, buf []byte) error {
 	n, err := c.Write(buf)
-	if n != len(buf) || err != nil {
+	if err != nil {
 		return err
 	}
+	if n != len(buf) {
+		return io.ErrShortWrite
+	}
 	return nil
 }