packet.go

package main

import (
	"encoding/binary"
	"net/netip"
)

func isMulticastPacket(packet []byte) bool {
	if len(packet) == 0 {
		return false
	}

	switch packet[0] >> 4 {
	case 4:
		if len(packet) < 20 {
			return false
		}
		addr, ok := netip.AddrFromSlice(packet[16:20])
		return ok && addr.IsMulticast()
	case 6:
		if len(packet) < 40 {
			return false
		}
		addr, ok := netip.AddrFromSlice(packet[24:40])
		return ok && addr.IsMulticast()
	default:
		return false
	}
}

func cloneAsUnicast(packet []byte, target netip.Addr) ([]byte, bool) {
	if len(packet) == 0 {
		return nil, false
	}

	clone := append([]byte(nil), packet...)

	switch clone[0] >> 4 {
	case 4:
		if !target.Is4() || len(clone) < 20 {
			return nil, false
		}
		copy(clone[16:20], target.AsSlice())
		rewriteIPv4Checksums(clone)
		return clone, true
	case 6:
		if !target.Is6() || len(clone) < 40 {
			return nil, false
		}
		if _, _, ok := ipv6TransportHeader(clone); !ok {
			return nil, false
		}
		copy(clone[24:40], target.AsSlice())
		rewriteIPv6Checksums(clone)
		return clone, true
	default:
		return nil, false
	}
}

func rewriteIPv4Checksums(packet []byte) {
	ihl := int(packet[0]&0x0f) * 4
	if ihl < 20 || len(packet) < ihl {
		return
	}

	packet[10], packet[11] = 0, 0
	binary.BigEndian.PutUint16(packet[10:12], checksum(packet[:ihl]))
	if binary.BigEndian.Uint16(packet[6:8])&0x3fff != 0 {
		return
	}

	protocol := packet[9]
	transport := packet[ihl:]
	switch protocol {
	case 6:
		if len(transport) < 18 {
			return
		}
		transport[16], transport[17] = 0, 0
		binary.BigEndian.PutUint16(transport[16:18], checksumWithPseudoHeader(packet[12:16], packet[16:20], protocol, transport))
	case 17:
		if len(transport) < 8 {
			return
		}
		transport[6], transport[7] = 0, 0
		binary.BigEndian.PutUint16(transport[6:8], udpChecksum(checksumWithPseudoHeader(packet[12:16], packet[16:20], protocol, transport)))
	case 1:
		if len(transport) < 4 {
			return
		}
		transport[2], transport[3] = 0, 0
		binary.BigEndian.PutUint16(transport[2:4], checksum(transport))
	}
}

func rewriteIPv6Checksums(packet []byte) {
	if len(packet) < 40 {
		return
	}

	nextHeader, transportOffset, ok := ipv6TransportHeader(packet)
	if !ok {
		return
	}
	transport := packet[transportOffset:]
	switch nextHeader {
	case 6:
		if len(transport) < 18 {
			return
		}
		transport[16], transport[17] = 0, 0
		binary.BigEndian.PutUint16(transport[16:18], checksumWithPseudoHeader(packet[8:24], packet[24:40], nextHeader, transport))
	case 17:
		if len(transport) < 8 {
			return
		}
		transport[6], transport[7] = 0, 0
		binary.BigEndian.PutUint16(transport[6:8], udpChecksum(checksumWithPseudoHeader(packet[8:24], packet[24:40], nextHeader, transport)))
	case 58:
		if len(transport) < 4 {
			return
		}
		transport[2], transport[3] = 0, 0
		binary.BigEndian.PutUint16(transport[2:4], checksumWithPseudoHeader(packet[8:24], packet[24:40], nextHeader, transport))
	}
}

func ipv6TransportHeader(packet []byte) (byte, int, bool) {
	if len(packet) < 40 {
		return 0, 0, false
	}

	nextHeader := packet[6]
	offset := 40
	for {
		switch nextHeader {
		case 0, 43, 60:
			if len(packet) < offset+2 {
				return 0, 0, false
			}
			extLen := (int(packet[offset+1]) + 1) * 8
			if extLen < 8 || len(packet) < offset+extLen {
				return 0, 0, false
			}
			nextHeader = packet[offset]
			offset += extLen
		case 51:
			if len(packet) < offset+2 {
				return 0, 0, false
			}
			extLen := (int(packet[offset+1]) + 2) * 4
			if extLen < 8 || len(packet) < offset+extLen {
				return 0, 0, false
			}
			nextHeader = packet[offset]
			offset += extLen
		case 44:
			return 0, 0, false
		default:
			return nextHeader, offset, true
		}
	}
}

func checksumWithPseudoHeader(src, dst []byte, protocol byte, payload []byte) uint16 {
	sum := checksumSum(src) + checksumSum(dst)
	if len(src) == 16 {
		sum += uint32(len(payload) >> 16)
	}
	sum += uint32(len(payload) & 0xffff)
	sum += uint32(protocol)
	sum += checksumSum(payload)
	return finishChecksum(sum)
}

func checksum(packet []byte) uint16 {
	return finishChecksum(checksumSum(packet))
}

func udpChecksum(sum uint16) uint16 {
	if sum == 0 {
		return 0xffff
	}
	return sum
}

func checksumSum(packet []byte) uint32 {
	var sum uint32
	for len(packet) > 1 {
		sum += uint32(binary.BigEndian.Uint16(packet[:2]))
		packet = packet[2:]
	}
	if len(packet) == 1 {
		sum += uint32(packet[0]) << 8
	}
	return sum
}

func finishChecksum(sum uint32) uint16 {
	for sum > 0xffff {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}