// Package inet recalculates length fields and checksums of IPv4/IPv6 packets
// and their transport headers (TCP, UDP, ICMP, ICMPv6, SCTP).
package inet

import (
	"encoding/binary"
	"errors"
	"fmt"
	"hash/crc32"
)

// IP protocol numbers.
const (
	protoHopByHop = 0
	protoICMP     = 1
	protoTCP      = 6
	protoUDP      = 17
	protoRouting  = 43
	protoFragment = 44
	protoAH       = 51
	protoICMPv6   = 58
	protoDstOpts  = 60
	protoSCTP     = 132
)

// Change records a field rewritten by Fix.
type Change struct {
	Field string
	Old   string
	New   string
}

// Result describes what Fix did.
type Result struct {
	Changes []Change // fields that were rewritten (including ones whose value did not change)
	Notes   []string // informational messages, e.g. why a checksum was skipped
}

func (r *Result) change16(field string, b []byte, v uint16) {
	old := binary.BigEndian.Uint16(b)
	binary.BigEndian.PutUint16(b, v)
	r.Changes = append(r.Changes, Change{field, fmt.Sprintf("0x%04x (%d)", old, old), fmt.Sprintf("0x%04x (%d)", v, v)})
}

func (r *Result) notef(format string, args ...any) {
	r.Notes = append(r.Notes, fmt.Sprintf(format, args...))
}

// Fix rewrites, in place, the length fields and checksums of the IPv4 or IPv6
// packet in pkt. pkt must start at the IP header and extend to the end of the
// packet; everything after the IP header is taken as the IP payload.
//
// hdrLen is the number of leading bytes of pkt that are headers, i.e. where
// the user's header input ends and the payload begins, or 0 if unknown. When
// known it is used to set the TCP data offset and the length of the SCTP
// chunk that the payload belongs to; when unknown those fields are left alone.
// The IPv4 IHL and IPv6 extension header lengths are always taken as given.
func Fix(pkt []byte, hdrLen int) (Result, error) {
	if len(pkt) == 0 {
		return Result{}, errors.New("no IP header data")
	}
	switch v := pkt[0] >> 4; v {
	case 4:
		return fixIPv4(pkt, hdrLen)
	case 6:
		return fixIPv6(pkt, hdrLen)
	default:
		return Result{}, fmt.Errorf("not an IPv4/IPv6 packet: IP version nibble is %d", v)
	}
}

func fixIPv4(pkt []byte, hdrLen int) (Result, error) {
	var r Result
	if len(pkt) < 20 {
		return r, fmt.Errorf("IPv4 header truncated: %d bytes, need at least 20", len(pkt))
	}
	ihl := int(pkt[0]&0x0f) * 4
	if ihl < 20 {
		return r, fmt.Errorf("invalid IPv4 header length (IHL) %d bytes", ihl)
	}
	if ihl > len(pkt) {
		return r, fmt.Errorf("IPv4 header length (IHL) %d exceeds packet length %d", ihl, len(pkt))
	}
	if len(pkt) > 0xffff {
		return r, fmt.Errorf("IPv4 packet length %d exceeds 65535", len(pkt))
	}
	if off := badIPv4Option(pkt[20:ihl]); off >= 0 {
		r.notef("IPv4 option at offset %d does not fit in the header length (IHL %d bytes); check IHL", 20+off, ihl)
	}

	r.change16("IPv4 total length", pkt[2:4], uint16(len(pkt)))
	binary.BigEndian.PutUint16(pkt[10:12], 0)
	r.change16("IPv4 header checksum", pkt[10:12], checksum(pkt[:ihl], 0))

	flagsFrag := binary.BigEndian.Uint16(pkt[6:8])
	if flagsFrag&0x2000 != 0 || flagsFrag&0x1fff != 0 {
		r.notef("IPv4 packet is a fragment; transport header not changed")
		return r, nil
	}

	proto := pkt[9]
	payload := pkt[ihl:]
	pseudo := sum(pkt[12:16], 0)
	pseudo = sum(pkt[16:20], pseudo)
	pseudo += uint32(proto) + uint32(len(payload))
	fixTransport(&r, proto, payload, transportHdrLen(&r, hdrLen, ihl), pseudo, false)
	return r, nil
}

// badIPv4Option walks IPv4 options and returns the offset of the first option
// that runs past the end of opts, or -1 if they are well formed.
func badIPv4Option(opts []byte) int {
	for i := 0; i < len(opts); {
		switch opts[i] {
		case 0: // end of option list, the rest is padding
			return -1
		case 1: // no-op
			i++
			continue
		}
		if i+1 >= len(opts) || opts[i+1] < 2 || i+int(opts[i+1]) > len(opts) {
			return i
		}
		i += int(opts[i+1])
	}
	return -1
}

func fixIPv6(pkt []byte, hdrLen int) (Result, error) {
	var r Result
	if len(pkt) < 40 {
		return r, fmt.Errorf("IPv6 header truncated: %d bytes, need 40", len(pkt))
	}
	if len(pkt)-40 > 0xffff {
		return r, fmt.Errorf("IPv6 payload length %d exceeds 65535 (jumbograms not supported)", len(pkt)-40)
	}
	r.change16("IPv6 payload length", pkt[4:6], uint16(len(pkt)-40))

	dst := pkt[24:40]
	next := pkt[6]
	off := 40
	for isIPv6ExtHeader(next) {
		if off+8 > len(pkt) {
			r.notef("IPv6 extension header truncated; transport header not changed")
			return r, nil
		}
		var hlen int
		switch next {
		case protoFragment:
			hlen = 8
			fragOffMore := binary.BigEndian.Uint16(pkt[off+2 : off+4])
			if fragOffMore&0xfff9 != 0 { // fragment offset != 0 or M flag set
				r.notef("IPv6 packet is a fragment; transport header not changed")
				return r, nil
			}
		case protoAH:
			hlen = (int(pkt[off+1]) + 2) * 4
		default:
			hlen = (int(pkt[off+1]) + 1) * 8
		}
		if off+hlen > len(pkt) {
			r.notef("IPv6 extension header truncated; transport header not changed")
			return r, nil
		}
		if next == protoRouting {
			if d := routingFinalDst(pkt[off : off+hlen]); d != nil {
				dst = d
			}
		}
		next = pkt[off]
		off += hlen
	}

	payload := pkt[off:]
	pseudo := sum(pkt[8:24], 0)
	pseudo = sum(dst, pseudo)
	pseudo += uint32(len(payload)) + uint32(next)
	fixTransport(&r, next, payload, transportHdrLen(&r, hdrLen, off), pseudo, true)
	return r, nil
}

// transportHdrLen converts the header length given to Fix into the number of
// transport-layer header bytes, or -1 if unknown.
func transportHdrLen(r *Result, hdrLen, ipHdrLen int) int {
	if hdrLen <= 0 {
		return -1
	}
	if hdrLen < ipHdrLen {
		r.notef("IP header (%d bytes including options/extension headers) extends past the layer 3 field (%d bytes)", ipHdrLen, hdrLen)
		return 0
	}
	return hdrLen - ipHdrLen
}

func isIPv6ExtHeader(proto byte) bool {
	switch proto {
	case protoHopByHop, protoRouting, protoFragment, protoAH, protoDstOpts:
		return true
	}
	return false
}

// routingFinalDst returns the final destination address carried in an IPv6
// routing header, which the upper-layer pseudo header must use (RFC 8200 8.1).
// It returns nil when the IPv6 destination address field should be used.
func routingFinalDst(rh []byte) []byte {
	routingType, segmentsLeft := rh[2], rh[3]
	if segmentsLeft == 0 {
		return nil
	}
	switch routingType {
	case 0: // deprecated source route: last address in the list
		n := int(rh[1]) / 2
		if n > 0 && 8+n*16 <= len(rh) {
			return rh[8+(n-1)*16 : 8+n*16]
		}
	case 2, 4: // Mobile IPv6 home address / Segment Routing: first address in the list
		if len(rh) >= 24 {
			return rh[8:24]
		}
	}
	return nil
}

// fixTransport fixes the transport header at the start of payload. hdrLen is
// the number of transport header bytes the user entered before the payload,
// or -1 if unknown.
func fixTransport(r *Result, proto byte, payload []byte, hdrLen int, pseudo uint32, ipv6 bool) {
	switch {
	case proto == protoTCP:
		if len(payload) < 20 {
			r.notef("TCP header truncated (%d bytes); TCP checksum not recalculated", len(payload))
			return
		}
		fixTCPDataOffset(r, payload, hdrLen)
		binary.BigEndian.PutUint16(payload[16:18], 0)
		r.change16("TCP checksum", payload[16:18], checksum(payload, pseudo))

	case proto == protoUDP:
		if len(payload) < 8 {
			r.notef("UDP header truncated (%d bytes); UDP length and checksum not recalculated", len(payload))
			return
		}
		r.change16("UDP length", payload[4:6], uint16(len(payload)))
		binary.BigEndian.PutUint16(payload[6:8], 0)
		cs := checksum(payload, pseudo)
		if cs == 0 {
			cs = 0xffff // zero means "no checksum" in UDP
		}
		r.change16("UDP checksum", payload[6:8], cs)

	case proto == protoICMP && !ipv6:
		if len(payload) < 4 {
			r.notef("ICMP header truncated (%d bytes); ICMP checksum not recalculated", len(payload))
			return
		}
		binary.BigEndian.PutUint16(payload[2:4], 0)
		r.change16("ICMP checksum", payload[2:4], checksum(payload, 0))

	case proto == protoICMPv6 && ipv6:
		if len(payload) < 4 {
			r.notef("ICMPv6 header truncated (%d bytes); ICMPv6 checksum not recalculated", len(payload))
			return
		}
		binary.BigEndian.PutUint16(payload[2:4], 0)
		r.change16("ICMPv6 checksum", payload[2:4], checksum(payload, pseudo))

	case proto == protoSCTP:
		if len(payload) < 12 {
			r.notef("SCTP common header truncated (%d bytes); SCTP checksum not recalculated", len(payload))
			return
		}
		fixSCTPChunkLength(r, payload, hdrLen)
		old := binary.BigEndian.Uint32(payload[8:12])
		binary.BigEndian.PutUint32(payload[8:12], 0)
		// CRC32c is transmitted least significant byte first (RFC 9260 appendix A).
		binary.LittleEndian.PutUint32(payload[8:12], crc32.Checksum(payload, castagnoli))
		v := binary.BigEndian.Uint32(payload[8:12])
		r.Changes = append(r.Changes, Change{"SCTP checksum (CRC32c)", fmt.Sprintf("0x%08x", old), fmt.Sprintf("0x%08x", v)})

	default:
		r.notef("IP protocol %d: no transport checksum to recalculate", proto)
	}
}

// fixTCPDataOffset sets the TCP data offset to the TCP header length (with
// options) entered in the layer 3 field.
func fixTCPDataOffset(r *Result, tcp []byte, hdrLen int) {
	switch {
	case hdrLen < 0:
		return
	case hdrLen == 0:
		r.notef("TCP header is in the payload field; TCP data offset not changed")
		return
	case hdrLen < 20 || hdrLen > 60 || hdrLen%4 != 0:
		r.notef("TCP header in the layer 3 field is %d bytes (must be 20-60 and a multiple of 4); TCP data offset not changed", hdrLen)
		return
	}
	old := int(tcp[12] >> 4)
	tcp[12] = byte(hdrLen/4)<<4 | tcp[12]&0x0f
	r.Changes = append(r.Changes, Change{"TCP data offset",
		fmt.Sprintf("%d (%d bytes)", old, old*4), fmt.Sprintf("%d (%d bytes)", hdrLen/4, hdrLen)})
}

// fixSCTPChunkLength sets the length of the SCTP chunk that the payload
// belongs to: the last chunk whose header starts in the layer 3 field. Chunks
// before it are skipped using their own length fields. The chunk is taken to
// extend to the end of the packet, so the payload must not include padding
// unless the chunk length is already correct.
func fixSCTPChunkLength(r *Result, sctp []byte, hdrLen int) {
	switch {
	case len(sctp) == 12: // no chunks
		return
	case hdrLen < 0:
		r.notef("SCTP chunk lengths are only recalculated when the layer 3 field is used; not changed")
		return
	case hdrLen < 16:
		r.notef("SCTP chunk header is not in the layer 3 field; chunk length not changed")
		return
	}

	off := 12
	for {
		clen := int(binary.BigEndian.Uint16(sctp[off+2:]))
		next := off + (clen+3)&^3
		if clen >= 4 && next+4 <= hdrLen {
			off = next
			continue
		}

		rem := len(sctp) - off
		newLen := rem
		if clen >= 4 && clen <= rem && (clen+3)&^3 == rem {
			newLen = clen // already correct; the rest is padding
		}
		if newLen > 0xffff {
			r.notef("SCTP chunk would be %d bytes, more than 65535; chunk length not changed", newLen)
			return
		}
		r.change16(fmt.Sprintf("SCTP %s chunk length", sctpChunkName(sctp[off])), sctp[off+2:off+4], uint16(newLen))
		if rem%4 != 0 {
			r.notef("SCTP %s chunk is %d bytes, not padded to a multiple of 4", sctpChunkName(sctp[off]), rem)
		}
		return
	}
}

func sctpChunkName(t byte) string {
	names := map[byte]string{
		0: "DATA", 1: "INIT", 2: "INIT ACK", 3: "SACK", 4: "HEARTBEAT", 5: "HEARTBEAT ACK",
		6: "ABORT", 7: "SHUTDOWN", 8: "SHUTDOWN ACK", 9: "ERROR", 10: "COOKIE ECHO",
		11: "COOKIE ACK", 14: "SHUTDOWN COMPLETE", 64: "I-DATA",
	}
	if n, ok := names[t]; ok {
		return n
	}
	return fmt.Sprintf("type %d", t)
}

var castagnoli = crc32.MakeTable(crc32.Castagnoli)

// sum adds data as big-endian 16-bit words to s (RFC 1071).
func sum(data []byte, s uint32) uint32 {
	n := len(data) &^ 1
	for i := 0; i < n; i += 2 {
		s += uint32(data[i])<<8 | uint32(data[i+1])
	}
	if len(data)%2 == 1 {
		s += uint32(data[len(data)-1]) << 8
	}
	return s
}

// checksum returns the Internet checksum of data, starting from the partial sum s.
func checksum(data []byte, s uint32) uint16 {
	s = sum(data, s)
	for s>>16 != 0 {
		s = s&0xffff + s>>16
	}
	return ^uint16(s)
}