package net

import (
	"bytes"
	"encoding/binary"
	"net"
	"os"
	"runtime"
	"testing"
	"time"

	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netns"
	"golang.org/x/sys/unix"
)

// TestPortLockInKernel builds a bridge with two pods in throwaway network namespaces, locks pod A's port and sends
// frames from A: what A may send reaches B, what claims anything but A's own addresses does not. It needs root.
func TestPortLockInKernel(t *testing.T) {
	if os.Geteuid() != 0 {
		t.Skip("needs root (creates network namespaces)")
	}
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	orig, err := netns.Get()
	if err != nil {
		t.Fatal(err)
	}
	defer orig.Close()
	defer netns.Set(orig)

	hub, err := netns.New()
	if err != nil {
		t.Fatal(err)
	}
	defer hub.Close()
	podA, _ := netns.New()
	defer podA.Close()
	podB, _ := netns.New()
	defer podB.Close()
	must(t, netns.Set(hub))

	br := &netlink.Bridge{LinkAttrs: netlink.LinkAttrs{Name: "br0"}}
	must(t, netlink.LinkAdd(br))
	must(t, netlink.LinkSetUp(br))
	ipA, ipB := net.IPv4(10, 99, 0, 1).To4(), net.IPv4(10, 99, 0, 2).To4()
	hostA, macA := attach(t, hub, podA, br, "vethwepltsta", ipA)
	_, macB := attach(t, hub, podB, br, "vethwepltstb", ipB)

	// as the CNI plugin does it: from the pod's end of the veth, in the pod's namespace
	must(t, LockContainerPort(podA, "eth0", ipA))
	if ok, err := PortLocked(hostA, macA, ipA); err != nil || !ok {
		dumpIngress(t, hostA)
		t.Fatalf("port not locked after LockContainerPort: %v %v", ok, err)
	}
	// locking again changes nothing (the filters are compared, not rebuilt)
	must(t, LockPort(hostA, macA, ipA))
	if ok, _ := PortLocked(hostA, macA, ipA); !ok {
		t.Fatal("port lost its lock on a repeated LockPort")
	}

	send := socketIn(t, podA, "eth0")
	recv := socketIn(t, podB, "eth0")
	defer unix.Close(send)
	defer unix.Close(recv)
	other, _ := net.ParseMAC("02:00:00:00:00:99")
	spoofIP := net.IPv4(10, 99, 0, 200).To4()

	cases := []struct {
		name   string
		frame  []byte
		passes bool
	}{
		{"ipv4 from its own addresses", ipv4Frame(macB, macA, ipA, ipB, "own-ip"), true},
		{"ipv4 from another MAC", ipv4Frame(macB, other, ipA, ipB, "mac-spoof"), false},
		{"ipv4 from another address", ipv4Frame(macB, macA, spoofIP, ipB, "ip-spoof"), false},
		{"arp for itself", arpFrame(macA, macA, ipA, ipB, "own-arp"), true},
		{"arp probe", arpFrame(macA, macA, net.IPv4zero.To4(), ipB, "probe"), true},
		{"arp claiming another address", arpFrame(macA, macA, spoofIP, ipB, "arp-ip-spoof"), false},
		{"arp claiming another MAC", arpFrame(macA, other, ipA, ipB, "arp-sha-spoof"), false},
		{"arp from another MAC", arpFrame(other, other, ipA, ipB, "arp-eth-spoof"), false},
		{"ipv6", etherFrame(macB, macA, 0x86dd, []byte("ipv6-frame")), false},
		{"other ethertype", etherFrame(macB, macA, 0x88b5, []byte("experimental")), false},
	}
	for _, c := range cases {
		drain(recv)
		if _, err := unix.Write(send, c.frame); err != nil {
			t.Fatalf("%s: send: %v", c.name, err)
		}
		got := sawPayload(recv, c.frame[len(c.frame)-12:], 300*time.Millisecond)
		if got != c.passes {
			t.Errorf("%s: reached B = %v, want %v", c.name, got, c.passes)
		}
	}
}

func dumpIngress(t *testing.T, link netlink.Link) {
	have, _ := netlink.FilterList(link, netlink.MakeHandle(0xffff, 0))
	for _, f := range have {
		t.Logf("have %s %+v", f.Type(), f.Attrs())
		if u, ok := f.(*netlink.U32); ok && u.Sel != nil {
			t.Logf("  sel flags %x keys %+v actions %+v", u.Sel.Flags, u.Sel.Keys, u.Actions)
		}
		if m, ok := f.(*netlink.MatchAll); ok {
			t.Logf("  actions %+v", m.Actions)
		}
	}
}

func must(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}

// attach plugs a pod namespace into the bridge through a veth named host on the hub side, eth0 on the pod side.
func attach(t *testing.T, hub, pod netns.NsHandle, br netlink.Link, host string, ip net.IP) (netlink.Link, net.HardwareAddr) {
	t.Helper()
	must(t, netlink.LinkAdd(&netlink.Veth{LinkAttrs: netlink.LinkAttrs{Name: host}, PeerName: "tmp" + host[len(host)-2:]}))
	h, err := netlink.LinkByName(host)
	must(t, err)
	must(t, netlink.LinkSetMaster(h, br))
	must(t, netlink.LinkSetUp(h))
	p, err := netlink.LinkByName("tmp" + host[len(host)-2:])
	must(t, err)
	must(t, netlink.LinkSetNsFd(p, int(pod)))
	must(t, netns.Set(pod))
	p, err = netlink.LinkByName("tmp" + host[len(host)-2:])
	must(t, err)
	must(t, netlink.LinkSetName(p, "eth0"))
	must(t, netlink.AddrAdd(p, &netlink.Addr{IPNet: &net.IPNet{IP: ip, Mask: net.CIDRMask(24, 32)}}))
	must(t, netlink.LinkSetUp(p))
	mac := p.Attrs().HardwareAddr
	must(t, netns.Set(hub))
	return h, mac
}

// socketIn opens a raw packet socket on an interface of a namespace (the socket stays in that namespace).
func socketIn(t *testing.T, ns netns.NsHandle, ifName string) int {
	t.Helper()
	cur, _ := netns.Get()
	defer cur.Close()
	must(t, netns.Set(ns))
	defer netns.Set(cur)
	link, err := netlink.LinkByName(ifName)
	must(t, err)
	fd, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, int(htons(unix.ETH_P_ALL)))
	must(t, err)
	must(t, unix.Bind(fd, &unix.SockaddrLinklayer{Protocol: htons(unix.ETH_P_ALL), Ifindex: link.Attrs().Index}))
	return fd
}

func htons(v uint16) uint16 { return v<<8 | v>>8 }

func drain(fd int) {
	buf := make([]byte, 2048)
	for {
		if _, _, err := unix.Recvfrom(fd, buf, unix.MSG_DONTWAIT); err != nil {
			return
		}
	}
}

func sawPayload(fd int, marker []byte, within time.Duration) bool {
	buf := make([]byte, 2048)
	deadline := time.Now().Add(within)
	for time.Now().Before(deadline) {
		n, _, err := unix.Recvfrom(fd, buf, unix.MSG_DONTWAIT)
		if err != nil {
			time.Sleep(5 * time.Millisecond)
			continue
		}
		if bytes.Contains(buf[:n], marker) {
			return true
		}
	}
	return false
}

func etherFrame(dst, src net.HardwareAddr, ethertype uint16, payload []byte) []byte {
	f := make([]byte, 14, 14+len(payload)+12)
	copy(f[0:6], dst)
	copy(f[6:12], src)
	binary.BigEndian.PutUint16(f[12:14], ethertype)
	f = append(f, payload...)
	// a fixed-size tail the receiver looks for, so each case is told apart
	tail := make([]byte, 12)
	copy(tail, payload)
	return append(f, tail...)
}

func ipv4Frame(dst, src net.HardwareAddr, from, to net.IP, tag string) []byte {
	body := []byte(tag)
	h := make([]byte, 20)
	h[0] = 0x45
	binary.BigEndian.PutUint16(h[2:4], uint16(20+len(body)+12))
	h[8] = 64
	h[9] = 253 // experimental protocol: nothing answers it
	copy(h[12:16], from)
	copy(h[16:20], to)
	binary.BigEndian.PutUint16(h[10:12], checksum(h))
	return etherFrame(dst, src, unix.ETH_P_IP, append(h, body...))
}

func arpFrame(ethSrc, sha net.HardwareAddr, spa, tpa net.IP, tag string) []byte {
	a := make([]byte, 28)
	binary.BigEndian.PutUint16(a[0:2], 1)
	binary.BigEndian.PutUint16(a[2:4], unix.ETH_P_IP)
	a[4], a[5] = 6, 4
	binary.BigEndian.PutUint16(a[6:8], 1) // request
	copy(a[8:14], sha)
	copy(a[14:18], spa)
	copy(a[24:28], tpa)
	bcast, _ := net.ParseMAC("ff:ff:ff:ff:ff:ff")
	return etherFrame(bcast, ethSrc, unix.ETH_P_ARP, append(a, []byte(tag)...))
}

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