diff --git a/conn_linux.go b/conn_linux.go index b1cdffb..e42ae9c 100644 --- a/conn_linux.go +++ b/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 +} diff --git a/conn_linux_test.go b/conn_linux_test.go index e9124ca..c136725 100644 --- a/conn_linux_test.go +++ b/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) + } +} diff --git a/syscall_linux.go b/syscall_linux.go new file mode 100644 index 0000000..85854ff --- /dev/null +++ b/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 +}