WIP: conn_linux: add recvmmsg

conn_linux.go

@@ -4,6 +4,7 @@ package socket
 
 import (
 	"context"
+	"io"
 	"os"
 	"unsafe"
 
@@ -115,3 +116,21 @@ func (c *Conn) Waitid(idType int, info *unix.Siginfo, options int, rusage *unix.
 		return unix.Waitid(idType, fd, info, options, rusage)
 	})
 }
+
+// Recvmmsg wraps receivemmsg(2).
+func (c *Conn) Recvmmsg(ctx context.Context, p, oob [][]byte, flags int) (int, []int, []int, []int, error) {
+	type ret struct {
+		n                   int
+		ns, oobn, recvflags []int
+	}
+
+	r, err := readT(ctx, c, "recvmmsg", func(fd int) (ret, error) {
+		n, ns, oobn, recvflags, err := recvmmsg(fd, p, oob, flags)
+		return ret{n, ns, oobn, recvflags}, err
+	})
+	if r.n == 0 && err == nil && c.facts.zeroReadIsEOF {
+		return 0, []int{}, []int{}, []int{}, io.EOF
+	}
+
+	return r.n, r.ns, r.oobn, r.recvflags, err
+}

conn_linux_test.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"runtime"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/mdlayher/socket"
@@ -211,3 +212,79 @@ func TestLinuxBindToDevice(t *testing.T) {
 		t.Fatalf("unexpected interface index (-want +got):\n%s", diff)
 	}
 }
+
+func TestRecvmmsgMultipleMessages(t *testing.T) {
+	t.Parallel()
+
+	// Verify that recvmmsg receives multiple datagrams in a single call.
+	c, err := socket.Socket(unix.AF_INET, unix.SOCK_DGRAM, 0, "udp4", nil)
+	if err != nil {
+		t.Fatalf("failed to open socket: %v", err)
+	}
+	defer c.Close()
+
+	if err := c.Bind(&unix.SockaddrInet4{Port: 0}); err != nil {
+		t.Fatalf("failed to bind socket: %v", err)
+	}
+
+	sa, err := c.Getsockname()
+	if err != nil {
+		t.Fatalf("failed to getsockname: %v", err)
+	}
+
+	// Send multiple datagrams to the bound socket.
+	const numMsgs = 4
+	for i := range numMsgs {
+		if err := c.Sendto(context.Background(), []byte{byte(i)}, 0, sa); err != nil {
+			t.Fatalf("failed to sendto[%d]: %v", i, err)
+		}
+	}
+
+	p := make([][]byte, numMsgs)
+	for i := range p {
+		p[i] = make([]byte, 64)
+	}
+
+	ctx := context.Background()
+	n, ns, _, _, err := c.Recvmmsg(ctx, p, nil, 0)
+	if err != nil {
+		t.Fatalf("Recvmmsg failed: %v", err)
+	}
+
+	if n != numMsgs {
+		t.Fatalf("expected %d messages in one recvmmsg call, got %d", numMsgs, n)
+	}
+
+	for i := range numMsgs {
+		if ns[i] != 1 {
+			t.Errorf("message %d: expected length 1, got %d", i, ns[i])
+		}
+		if p[i][0] != byte(i) {
+			t.Errorf("message %d: expected payload %d, got %d", i, i, p[i][0])
+		}
+	}
+}
+
+func TestRecvmmsgContextCanceled(t *testing.T) {
+	t.Parallel()
+
+	c, err := socket.Socket(unix.AF_INET, unix.SOCK_DGRAM, 0, "udp4", nil)
+	if err != nil {
+		t.Fatalf("failed to open socket: %v", err)
+	}
+	defer c.Close()
+
+	if err := c.Bind(&unix.SockaddrInet4{Port: 0}); err != nil {
+		t.Fatalf("failed to bind socket: %v", err)
+	}
+
+	// Context will be canceled while blocked in recvmmsg.
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	p := [][]byte{make([]byte, 64)}
+	_, _, _, _, err = c.Recvmmsg(ctx, p, nil, 0)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("expected context.DeadlineExceeded, got: %v", err)
+	}
+}

syscall_linux.go

@@ -0,0 +1,96 @@
+//go:build linux
+
+package socket
+
+import (
+	"runtime"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+// mmsghdr matches the Linux kernel's struct mmsghdr layout.
+type mmsghdr struct {
+	Hdr    unix.Msghdr
+	Msglen uint32
+	_      [4]byte // padding to 8-byte alignment
+}
+
+func recvmmsg(fd int, ps, oobs [][]byte, flags int) (n int, ns, oobns, msgFlags []int, err error) {
+	vlen := len(ps)
+	if vlen == 0 {
+		return 0, nil, nil, nil, unix.EINVAL
+	}
+	if oobs != nil && len(oobs) != vlen {
+		return 0, nil, nil, nil, unix.EINVAL
+	}
+
+	iovs := make([]unix.Iovec, vlen)
+	msgs := make([]mmsghdr, vlen)
+	for i := range vlen {
+		if len(ps[i]) > 0 {
+			iovs[i].Base = &ps[i][0]
+			iovs[i].SetLen(len(ps[i]))
+			msgs[i].Hdr.Iov = &iovs[i]
+			msgs[i].Hdr.SetIovlen(1)
+		}
+		if oobs != nil && len(oobs[i]) > 0 {
+			msgs[i].Hdr.Control = &oobs[i][0]
+			msgs[i].Hdr.SetControllen(len(oobs[i]))
+		}
+	}
+
+	r, _, errno := unix.Syscall6(
+		unix.SYS_RECVMMSG,
+		uintptr(fd),
+		uintptr(unsafe.Pointer(&msgs[0])),
+		uintptr(vlen),
+		uintptr(flags),
+		0,
+		0,
+	)
+	runtime.KeepAlive(ps)
+	runtime.KeepAlive(oobs)
+	runtime.KeepAlive(iovs)
+	runtime.KeepAlive(msgs)
+
+	if errno != 0 {
+		return 0, nil, nil, nil, errnoErr(errno)
+	}
+	n = int(r)
+
+	ns = make([]int, n)
+	oobns = make([]int, n)
+	msgFlags = make([]int, n)
+	for i := range n {
+		ns[i] = int(msgs[i].Msglen)
+		oobns[i] = int(msgs[i].Hdr.Controllen)
+		msgFlags[i] = int(msgs[i].Hdr.Flags)
+	}
+	return n, ns, oobns, msgFlags, nil
+}
+
+// Do the interface allocations only once for common
+// Errno values.
+var (
+	errEAGAIN error = syscall.EAGAIN
+	errEINVAL error = syscall.EINVAL
+	errENOENT error = syscall.ENOENT
+)
+
+// errnoErr returns common boxed Errno values, to prevent
+// allocations at runtime.
+func errnoErr(e syscall.Errno) error {
+	switch e {
+	case 0:
+		return nil
+	case unix.EAGAIN:
+		return errEAGAIN
+	case unix.EINVAL:
+		return errEINVAL
+	case unix.ENOENT:
+		return errENOENT
+	}
+	return e
+}