golang.org/x/net@v0.25.1-0.20240516223405-c87a5b62e243/internal/socket/mmsghdr_unix.go

// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build aix || linux || netbsd

package socket

import (
	"net"
	"os"
	"sync"
	"syscall"
)

type mmsghdrs []mmsghdr

func (hs mmsghdrs) unpack(ms []Message, parseFn func([]byte, string) (net.Addr, error), hint string) error {
	for i := range hs {
		ms[i].N = int(hs[i].Len)
		ms[i].NN = hs[i].Hdr.controllen()
		ms[i].Flags = hs[i].Hdr.flags()
		if parseFn != nil {
			var err error
			ms[i].Addr, err = parseFn(hs[i].Hdr.name(), hint)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// mmsghdrsPacker packs Message-slices into mmsghdrs (re-)using pre-allocated buffers.
type mmsghdrsPacker struct {
	// hs are the pre-allocated mmsghdrs.
	hs mmsghdrs
	// sockaddrs is the pre-allocated buffer for the Hdr.Name buffers.
	// We use one large buffer for all messages and slice it up.
	sockaddrs []byte
	// vs are the pre-allocated iovecs.
	// We allocate one large buffer for all messages and slice it up. This allows to reuse the buffer
	// if the number of buffers per message is distributed differently between calls.
	vs []iovec
}

func (p *mmsghdrsPacker) prepare(ms []Message) {
	n := len(ms)
	if n <= cap(p.hs) {
		p.hs = p.hs[:n]
	} else {
		p.hs = make(mmsghdrs, n)
	}
	if n*sizeofSockaddrInet6 <= cap(p.sockaddrs) {
		p.sockaddrs = p.sockaddrs[:n*sizeofSockaddrInet6]
	} else {
		p.sockaddrs = make([]byte, n*sizeofSockaddrInet6)
	}

	nb := 0
	for _, m := range ms {
		nb += len(m.Buffers)
	}
	if nb <= cap(p.vs) {
		p.vs = p.vs[:nb]
	} else {
		p.vs = make([]iovec, nb)
	}
}

func (p *mmsghdrsPacker) pack(ms []Message, parseFn func([]byte, string) (net.Addr, error), marshalFn func(net.Addr, []byte) int) mmsghdrs {
	p.prepare(ms)
	hs := p.hs
	vsRest := p.vs
	saRest := p.sockaddrs
	for i := range hs {
		nvs := len(ms[i].Buffers)
		vs := vsRest[:nvs]
		vsRest = vsRest[nvs:]

		var sa []byte
		if parseFn != nil {
			sa = saRest[:sizeofSockaddrInet6]
			saRest = saRest[sizeofSockaddrInet6:]
		} else if marshalFn != nil {
			n := marshalFn(ms[i].Addr, saRest)
			if n > 0 {
				sa = saRest[:n]
				saRest = saRest[n:]
			}
		}
		hs[i].Hdr.pack(vs, ms[i].Buffers, ms[i].OOB, sa)
	}
	return hs
}

// syscaller is a helper to invoke recvmmsg and sendmmsg via the RawConn.Read/Write interface.
// It is reusable, to amortize the overhead of allocating a closure for the function passed to
// RawConn.Read/Write.
type syscaller struct {
	n     int
	operr error
	hs    mmsghdrs
	flags int

	boundRecvmmsgF func(uintptr) bool
	boundSendmmsgF func(uintptr) bool
}

func (r *syscaller) init() {
	r.boundRecvmmsgF = r.recvmmsgF
	r.boundSendmmsgF = r.sendmmsgF
}

func (r *syscaller) recvmmsg(c syscall.RawConn, hs mmsghdrs, flags int) (int, error) {
	r.n = 0
	r.operr = nil
	r.hs = hs
	r.flags = flags
	if err := c.Read(r.boundRecvmmsgF); err != nil {
		return r.n, err
	}
	if r.operr != nil {
		return r.n, os.NewSyscallError("recvmmsg", r.operr)
	}
	return r.n, nil
}

func (r *syscaller) recvmmsgF(s uintptr) bool {
	r.n, r.operr = recvmmsg(s, r.hs, r.flags)
	return ioComplete(r.flags, r.operr)
}

func (r *syscaller) sendmmsg(c syscall.RawConn, hs mmsghdrs, flags int) (int, error) {
	r.n = 0
	r.operr = nil
	r.hs = hs
	r.flags = flags
	if err := c.Write(r.boundSendmmsgF); err != nil {
		return r.n, err
	}
	if r.operr != nil {
		return r.n, os.NewSyscallError("sendmmsg", r.operr)
	}
	return r.n, nil
}

func (r *syscaller) sendmmsgF(s uintptr) bool {
	r.n, r.operr = sendmmsg(s, r.hs, r.flags)
	return ioComplete(r.flags, r.operr)
}

// mmsgTmps holds reusable temporary helpers for recvmmsg and sendmmsg.
type mmsgTmps struct {
	packer    mmsghdrsPacker
	syscaller syscaller
}

var defaultMmsgTmpsPool = mmsgTmpsPool{
	p: sync.Pool{
		New: func() interface{} {
			tmps := new(mmsgTmps)
			tmps.syscaller.init()
			return tmps
		},
	},
}

type mmsgTmpsPool struct {
	p sync.Pool
}

func (p *mmsgTmpsPool) Get() *mmsgTmps {
	m := p.p.Get().(*mmsgTmps)
	// Clear fields up to the len (not the cap) of the slice,
	// assuming that the previous caller only used that many elements.
	for i := range m.packer.sockaddrs {
		m.packer.sockaddrs[i] = 0
	}
	m.packer.sockaddrs = m.packer.sockaddrs[:0]
	for i := range m.packer.vs {
		m.packer.vs[i] = iovec{}
	}
	m.packer.vs = m.packer.vs[:0]
	for i := range m.packer.hs {
		m.packer.hs[i].Len = 0
		m.packer.hs[i].Hdr = msghdr{}
	}
	m.packer.hs = m.packer.hs[:0]
	return m
}

func (p *mmsgTmpsPool) Put(tmps *mmsgTmps) {
	p.p.Put(tmps)
}