// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

//go:build linux

package dataplane

import (
	"encoding/binary"
	"fmt"
	"net/netip"
	"sync/atomic"
	"syscall"
	"time"

	"echo-lot.app/server/internal/session"
)

// Crafted IPv4 fragmentation (spec §5 frag_send).
//
// Letting the kernel fragment an oversized datagram — which is what big_send with df=false does —
// answers one question: do fragments get through at all. It cannot answer the more interesting
// one, because the kernel always emits fragments in order, first one first.
//
// The classic middlebox fault is precisely about that ordering. Only the *first* fragment carries
// the UDP header, and therefore the ports; a stateful firewall or NAT that has not seen it has no
// flow to match later fragments against. Plenty of implementations drop them. Others hold them
// briefly and reassemble; others leak. The difference is invisible to any test that sends
// fragments in order, and it shows up in the real world as "large DNS answers fail on this
// network" or "the VPN works until the MTU drops".
//
// So this builds the fragments by hand and controls their order and timing. That needs a raw
// socket (CAP_NET_RAW); when we do not have one the capability is not advertised, rather than
// advertised and failing later.

// FragMode is how a fragmented datagram is put on the wire.
type FragMode string

const (
	// FragInOrder is the baseline: first fragment first, as the kernel would. A path that fails
	// this fails everything, and it tells the others apart from a path that drops all fragments.
	FragInOrder FragMode = "in_order"
	// FragReversed sends the last fragment first. This is the one that finds stateful devices
	// which need the first fragment to build state.
	FragReversed FragMode = "reversed"
	// FragFirstLast holds the first fragment back until the others have arrived, which tests
	// whether the path buffers non-first fragments at all and for how long.
	FragFirstLast FragMode = "first_last"
)

var fragIPID atomic.Uint32

// RawFragSupported reports whether crafted fragments can actually be sent here.
//
// Checked by opening the socket rather than by inspecting capabilities: the question is "will
// this work", and a permission model has more ways to say no than a capability bit has to say yes
// (user namespaces, seccomp, LSM). Advertising a capability we cannot deliver would turn a
// missing feature into a failed measurement.
func RawFragSupported() bool {
	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		return false
	}
	_ = syscall.Close(fd)
	return true
}

// FragResult is what happened to one crafted fragment burst.
type FragResult struct {
	Mode      FragMode `json:"mode"`
	SizeBytes int      `json:"size_bytes"`
	Fragments int      `json:"fragments"`
	Sent      bool     `json:"sent"`
	Err       string   `json:"err,omitempty"`
}

// FragSend emits one ELT1 packet of sizeBytes as hand-built IPv4 fragments, in the given order.
//
// The datagram is assembled whole and then cut up, so what the client reassembles — if it
// reassembles — is a normal, HMAC-valid packet indistinguishable from any other. That matters:
// the client must not be able to tell a crafted fragment burst from a kernel one, or it would be
// measuring our sender rather than the path.
func (s *Server) FragSend(
	sess *session.Session, g *session.Grant, sizeBytes int, mode FragMode, fragSize int,
) (FragResult, error) {
	res := FragResult{Mode: mode, SizeBytes: sizeBytes}

	target := sess.DataSource()
	if !target.IsValid() {
		return res, fmt.Errorf("no observed data-plane source")
	}
	if !target.Addr().Unmap().Is4() {
		// IPv6 has no in-network fragmentation: only the source may fragment, via an extension
		// header. Worth building, but it is a different mechanism and belongs in its own code
		// path rather than pretending this one covers it.
		return res, fmt.Errorf("crafted fragmentation is IPv4-only for now")
	}
	conn := s.connFor(target, sess.DataLocal())
	if conn == nil {
		return res, fmt.Errorf("no data-plane socket matches target family")
	}
	local := sess.DataLocal()
	if !local.IsValid() {
		return res, fmt.Errorf("session has no recorded local address")
	}

	if sizeBytes < HeaderSize+32 {
		sizeBytes = HeaderSize + 32
	}
	if sizeBytes > 8000 {
		sizeBytes = 8000
	}
	if !g.Allow(sizeBytes) {
		return res, fmt.Errorf("grant exhausted")
	}

	// The ELT1 packet, signed exactly as any other, then wrapped in UDP.
	payload := make([]byte, sizeBytes-HeaderSize)
	binary.BigEndian.PutUint32(payload[0:4], uint32(sizeBytes))
	// [8:16]: the action id, as on every granted packet (spec §5 correlation); the mode string
	// sits past the reserved slot.
	putActionID(payload, g.ActionID)
	copy(payload[16:], mode)
	elt := s.buildPacket(sess, TypeFragData, 0, payload)

	udp := buildUDP(local, target, elt)

	// Fragment offsets are in 8-byte units, so every fragment except the last must be a multiple
	// of 8. A payload that is not is not an error — it is a fragment that no host will reassemble.
	if fragSize <= 0 {
		fragSize = 576
	}
	fragSize = (fragSize / 8) * 8
	if fragSize < 8 {
		fragSize = 8
	}

	fragments := splitIPv4(local.Addr(), target.Addr(), udp, fragSize, uint16(fragIPID.Add(1)))
	res.Fragments = len(fragments)

	fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_RAW)
	if err != nil {
		res.Err = err.Error()
		return res, err
	}
	defer syscall.Close(fd)
	if err := syscall.SetsockoptInt(fd, syscall.IPPROTO_IP, syscall.IP_HDRINCL, 1); err != nil {
		res.Err = err.Error()
		return res, err
	}

	dst := syscall.SockaddrInet4{}
	copy(dst.Addr[:], target.Addr().Unmap().AsSlice())

	send := func(pkt []byte) error { return syscall.Sendto(fd, pkt, 0, &dst) }

	switch mode {
	case FragReversed:
		for i := len(fragments) - 1; i >= 0; i-- {
			if err := send(fragments[i]); err != nil {
				res.Err = err.Error()
				return res, err
			}
			time.Sleep(time.Millisecond)
		}
	case FragFirstLast:
		for i := 1; i < len(fragments); i++ {
			if err := send(fragments[i]); err != nil {
				res.Err = err.Error()
				return res, err
			}
			time.Sleep(time.Millisecond)
		}
		// Long enough to be a real test of whether anything holds fragments, short enough to stay
		// inside the usual 30-second reassembly timeout by a wide margin.
		time.Sleep(250 * time.Millisecond)
		if err := send(fragments[0]); err != nil {
			res.Err = err.Error()
			return res, err
		}
	default:
		for _, f := range fragments {
			if err := send(f); err != nil {
				res.Err = err.Error()
				return res, err
			}
			time.Sleep(time.Millisecond)
		}
	}
	res.Sent = true
	return res, nil
}

// buildUDP wraps a payload in a UDP header with a computed checksum.
//
// The checksum is optional in IPv4 and it would be less code to send zero, but a zero-checksum
// datagram is dropped by some middleboxes — and that drop would be recorded as a fragmentation
// failure, which is exactly the wrong conclusion.
func buildUDP(src, dst netip.AddrPort, payload []byte) []byte {
	out := make([]byte, 8+len(payload))
	binary.BigEndian.PutUint16(out[0:2], src.Port())
	binary.BigEndian.PutUint16(out[2:4], dst.Port())
	binary.BigEndian.PutUint16(out[4:6], uint16(8+len(payload)))
	copy(out[8:], payload)

	// Pseudo-header + UDP header + data, per RFC 768.
	var sum uint32
	s4, d4 := src.Addr().Unmap().As4(), dst.Addr().Unmap().As4()
	for _, b := range [][]byte{s4[:], d4[:]} {
		sum += uint32(binary.BigEndian.Uint16(b[0:2]))
		sum += uint32(binary.BigEndian.Uint16(b[2:4]))
	}
	sum += uint32(syscall.IPPROTO_UDP)
	sum += uint32(len(out))
	for i := 0; i+1 < len(out); i += 2 {
		sum += uint32(binary.BigEndian.Uint16(out[i : i+2]))
	}
	if len(out)%2 == 1 {
		sum += uint32(out[len(out)-1]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xFFFF) + (sum >> 16)
	}
	ck := ^uint16(sum)
	if ck == 0 {
		ck = 0xFFFF // 0 means "no checksum" in IPv4; the all-ones form is the same value
	}
	binary.BigEndian.PutUint16(out[6:8], ck)
	return out
}

// splitIPv4 cuts a UDP datagram into IPv4 fragments of at most fragSize payload bytes each.
//
// Every fragment carries the same IP ID — that is what marks them as one datagram — and every one
// but the last sets MF. The kernel fills in the header checksum and total length for us under
// IP_HDRINCL (raw(7)); the ID it only fills when zero, which is why it is set explicitly here.
func splitIPv4(src, dst netip.Addr, udp []byte, fragSize int, id uint16) [][]byte {
	s4, d4 := src.Unmap().As4(), dst.Unmap().As4()
	var out [][]byte
	for off := 0; off < len(udp); off += fragSize {
		end := off + fragSize
		if end > len(udp) {
			end = len(udp)
		}
		chunk := udp[off:end]
		more := end < len(udp)

		hdr := make([]byte, 20, 20+len(chunk))
		hdr[0] = 0x45 // IPv4, 5 words of header
		hdr[1] = 0    // DSCP/ECN
		binary.BigEndian.PutUint16(hdr[2:4], uint16(20+len(chunk)))
		binary.BigEndian.PutUint16(hdr[4:6], id)
		flagsOff := uint16(off / 8)
		if more {
			flagsOff |= 0x2000 // MF
		}
		binary.BigEndian.PutUint16(hdr[6:8], flagsOff)
		hdr[8] = 64 // TTL
		hdr[9] = syscall.IPPROTO_UDP
		// hdr[10:12] checksum left zero: the kernel computes it under IP_HDRINCL.
		copy(hdr[12:16], s4[:])
		copy(hdr[16:20], d4[:])
		out = append(out, append(hdr, chunk...))
	}
	return out
}