NetworkTool/server/internal/discovery/packetinfo_linux.go

//go:build linux

package discovery

import (
	"encoding/binary"
	"net"
	"syscall"
	"unsafe"
)

func enableLocalAddrControl(conn *net.UDPConn) error {
	rawConn, err := conn.SyscallConn()
	if err != nil {
		return err
	}
	var controlErr error
	err = rawConn.Control(func(fd uintptr) {
		controlErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IP, syscall.IP_PKTINFO, 1)
		if controlErr == nil {
			controlErr = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_BROADCAST, 1)
		}
	})
	if err != nil {
		return err
	}
	return controlErr
}

func readFromUDP(conn *net.UDPConn, buf []byte) (int, *net.UDPAddr, udpPacketInfo, error) {
	oob := make([]byte, 128)
	n, oobn, _, remote, err := conn.ReadMsgUDP(buf, oob)
	if err != nil {
		return 0, nil, udpPacketInfo{}, err
	}
	packetInfo := parsePacketInfo(oob[:oobn])
	return n, remote, packetInfo, nil
}

func writeToUDP(conn *net.UDPConn, payload []byte, remote *net.UDPAddr, packetInfo udpPacketInfo, localIP net.IP) (int, error) {
	if packetInfo.ifIndex <= 0 {
		return conn.WriteToUDP(payload, remote)
	}

	oob := make([]byte, syscall.CmsgSpace(syscall.SizeofInet4Pktinfo))
	header := (*syscall.Cmsghdr)(unsafe.Pointer(&oob[0]))
	header.Level = syscall.IPPROTO_IP
	header.Type = syscall.IP_PKTINFO
	header.SetLen(syscall.CmsgLen(syscall.SizeofInet4Pktinfo))

	info := (*syscall.Inet4Pktinfo)(unsafe.Pointer(&oob[syscall.CmsgLen(0)]))
	info.Ifindex = int32(packetInfo.ifIndex)
	if ipv4 := localIP.To4(); ipv4 != nil {
		copy(info.Spec_dst[:], ipv4)
	}

	n, _, err := conn.WriteMsgUDP(payload, oob, remote)
	return n, err
}

func parsePacketInfo(oob []byte) udpPacketInfo {
	messages, err := syscall.ParseSocketControlMessage(oob)
	if err != nil {
		return udpPacketInfo{}
	}
	for _, message := range messages {
		if message.Header.Level != syscall.IPPROTO_IP || message.Header.Type != syscall.IP_PKTINFO || len(message.Data) < 12 {
			continue
		}
		specDst := [4]byte(message.Data[4:8])
		addr := [4]byte(message.Data[8:12])
		return udpPacketInfo{
			localIP: packetInfoIP(specDst, addr),
			ifIndex: int(int32(binary.LittleEndian.Uint32(message.Data[0:4]))),
		}
	}
	return udpPacketInfo{}
}

func packetInfoIP(specDst [4]byte, addr [4]byte) net.IP {
	if ip := ipv4FromBytes(specDst); ip != nil {
		return ip
	}
	return ipv4FromBytes(addr)
}

func ipv4FromBytes(value [4]byte) net.IP {
	if value == [4]byte{} || value == [4]byte{255, 255, 255, 255} {
		return nil
	}
	return net.IPv4(value[0], value[1], value[2], value[3]).To4()
}