package vpnruntime

import "encoding/binary"

func normalizeIPv4PacketChecksums(packet []byte) bool {
	if len(packet) < 20 || packet[0]>>4 != 4 {
		return false
	}
	ihl := int(packet[0]&0x0f) * 4
	if ihl < 20 || len(packet) < ihl {
		return false
	}
	totalLen := int(binary.BigEndian.Uint16(packet[2:4]))
	if totalLen <= 0 || totalLen > len(packet) {
		totalLen = len(packet)
	}
	if totalLen < ihl {
		return false
	}

	packet[10], packet[11] = 0, 0
	binary.BigEndian.PutUint16(packet[10:12], checksum(packet[:ihl]))

	proto := packet[9]
	payload := packet[ihl:totalLen]
	switch proto {
	case 6:
		if len(payload) < 20 {
			return true
		}
		payload[16], payload[17] = 0, 0
		binary.BigEndian.PutUint16(payload[16:18], transportChecksum(packet, payload, proto))
	case 17:
		if len(payload) < 8 {
			return true
		}
		payload[6], payload[7] = 0, 0
		sum := transportChecksum(packet, payload, proto)
		if sum == 0 {
			sum = 0xffff
		}
		binary.BigEndian.PutUint16(payload[6:8], sum)
	case 1:
		if len(payload) < 4 {
			return true
		}
		payload[2], payload[3] = 0, 0
		binary.BigEndian.PutUint16(payload[2:4], checksum(payload))
	}
	return true
}

func transportChecksum(ipHeader []byte, payload []byte, proto byte) uint16 {
	pseudo := make([]byte, 12+len(payload))
	copy(pseudo[0:4], ipHeader[12:16])
	copy(pseudo[4:8], ipHeader[16:20])
	pseudo[8] = 0
	pseudo[9] = proto
	binary.BigEndian.PutUint16(pseudo[10:12], uint16(len(payload)))
	copy(pseudo[12:], payload)
	return checksum(pseudo)
}

func checksum(data []byte) uint16 {
	var sum uint32
	for len(data) >= 2 {
		sum += uint32(binary.BigEndian.Uint16(data[:2]))
		data = data[2:]
	}
	if len(data) == 1 {
		sum += uint32(data[0]) << 8
	}
	for (sum >> 16) != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}