github.com/dylandreimerink/gobpfld@v0.6.1-0.20220205171531-e79c330ad608/cmd/examples/xsk_echo_reply/main.go

package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"

	"github.com/dylandreimerink/gobpfld"
	"github.com/dylandreimerink/gobpfld/bpftypes"
	"github.com/dylandreimerink/gobpfld/ebpf"
	"github.com/vishvananda/netlink"
	"golang.org/x/sys/unix"

	_ "net/http/pprof"
)

var (
	ifname = flag.String("ifname", "eth0", "name of the network interface to bind to")
	ip     = flag.String("ip", "", "the ipv4 ip we will use as ping target")
)

func main() {
	flag.Parse()

	if ifname == nil || *ifname == "" {
		fmt.Fprint(os.Stderr, "flag 'ifname' is required\n")
		os.Exit(1)
	}

	if ip == nil || *ip == "" {
		fmt.Fprint(os.Stderr, "flag 'ip' is required\n")
		os.Exit(1)
	}

	linkip := net.ParseIP(*ip).To4()
	if linkip.Equal(net.IPv4zero) {
		fmt.Fprint(os.Stderr, "flag 'ip' contains an invalid IPv4\n")
		os.Exit(1)
	}

	linkName := *ifname
	link, err := netlink.LinkByName(linkName)
	if err != nil {
		fmt.Fprintf(os.Stderr, "get link by name: %s\n", err.Error())
		os.Exit(1)
	}
	linkMAC := link.Attrs().HardwareAddr

	// Create an new XSK socket, bound to queue 0.
	// NOTE this example only works on non-multi queue NIC's
	xsksock, err := gobpfld.NewXSKSocket(gobpfld.XSKSettings{
		NetDevIfIndex: link.Attrs().Index,
		QueueID:       0,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "new socket: %s\n", err.Error())
		os.Exit(1)
	}

	// Block forever until we can read or write (default behavour is to never block)
	xsksock.SetReadTimeout(-1)
	xsksock.SetWriteTimeout(-1)

	// Generate a program which will bypass all traffic to userspace
	program := &gobpfld.ProgramXDP{
		AbstractBPFProgram: gobpfld.AbstractBPFProgram{
			Name:        gobpfld.MustNewObjName("xsk_bypass"),
			ProgramType: bpftypes.BPF_PROG_TYPE_XDP,
			License:     "GPL",
			Maps: map[string]gobpfld.BPFMap{
				"xskmap": &gobpfld.XSKMap{
					AbstractMap: gobpfld.AbstractMap{
						Name: gobpfld.MustNewObjName("xskmap"),
						Definition: gobpfld.BPFMapDef{
							Type:       bpftypes.BPF_MAP_TYPE_XSKMAP,
							KeySize:    4, // SizeOf(uint32)
							ValueSize:  4, // SizeOf(uint32)
							MaxEntries: 5,
						},
					},
				},
			},
			MapFDLocations: map[string][]uint64{
				"xskmap": {
					// LoadConstant64bit is the 2nd instruction in this program. So the first byte of the
					// 2nd instruction is the width of a instruction * 1 to skip the first 1 instruction
					uint64(ebpf.BPFInstSize) * 1,
				},
			},
			// Instructions for this program:
			// int xsk_bypass(struct xdp_md *ctx)
			// {
			// 	return bpf_redirect_map(&xsks_map, ctx->rx_queue_index, XDP_PASS);
			// }
			//
			// NOTE this program only works in linux kernel >= 5.3
			// https://elixir.bootlin.com/linux/v5.12.2/source/tools/lib/bpf/xsk.c#L416
			Instructions: ebpf.MustEncode([]ebpf.Instruction{
				// load ((xdp_md) ctx)->rx_queue_index into R2 (used as 2nd parameter)
				/* r2 = *(u32 *)(r1 + 16) */
				&ebpf.LoadMemory{
					Dest:   ebpf.BPF_REG_2,
					Src:    ebpf.BPF_REG_1,
					Size:   ebpf.BPF_W,
					Offset: 16,
				},
				// Set R1(first parameter) to the address of the xskmap.
				// Which will be set during loading
				/* r1 = xskmap[] */
				&ebpf.LoadConstant64bit{
					Dest: ebpf.BPF_REG_1,
				},
				// Move XDP_PASS into R3 (third argument)
				/* r3 = XDP_PASS */
				&ebpf.Mov64{
					Dest:  ebpf.BPF_REG_3,
					Value: ebpf.XDP_PASS,
				},
				/* call bpf_redirect_map */
				&ebpf.CallHelper{
					Function: 51,
				},
				&ebpf.Exit{}, // exit
			}),
		},
	}

	xskmap := program.Maps["xskmap"].(*gobpfld.XSKMap)
	log, err := program.Load(gobpfld.ProgXDPLoadOpts{
		VerifierLogLevel: bpftypes.BPFLogLevelBasic,
	})

	program.DecodeToReader(os.Stdout)
	fmt.Fprintln(os.Stderr, log)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while loading program: %s\n", err.Error())
		os.Exit(1)
	}

	// Set the xsksocket for queue ID 0
	err = xskmap.Set(0, xsksock)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while setting xsksock in map: %s\n", err.Error())
		os.Exit(1)
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, unix.SIGTERM, unix.SIGINT)

	err = program.Attach(gobpfld.ProgXDPAttachOpts{
		InterfaceName: linkName,
		Replace:       true,
	})

	if err != nil {
		fmt.Fprintf(os.Stderr, "error while attaching program to loopback device: %s\n", err.Error())
		os.Exit(1)
	}

	done := false
	for !done {
		select {
		case <-sigChan:
			done = true

		default:
			// Seperator to distinguish between frames
			fmt.Println("------")
			lease, err := xsksock.ReadLease()
			if err != nil {
				fmt.Fprintf(os.Stderr, "read lease: %s\n", err.Error())
				break
			}

			if lease == nil {
				break
			}

			err = HandleFrame(lease, linkMAC, linkip)
			if err != nil {
				fmt.Fprintf(os.Stderr, "echo reply: %s", err.Error())
			}
		}
	}

	fmt.Println("Detaching XPD program and stopping")

	err = program.XDPLinkDetach(gobpfld.BPFProgramXDPLinkDetachSettings{
		All: true,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "error while detaching program: %s\n", err.Error())
		os.Exit(1)
	}

	os.Exit(0)
}

func HandleFrame(lease *gobpfld.XSKLease, linkMac net.HardwareAddr, linkIP net.IP) error {
	var err error

	// Swap the MAC addresses
	fmt.Printf("Src MAC: %X, Dst MAC: %X, EthType: %X\n", lease.Data[0:6], lease.Data[6:12], lease.Data[12:14])
	swapMac := make([]byte, 6)
	copy(swapMac, lease.Data[0:6])
	copy(lease.Data[0:6], lease.Data[6:12])
	copy(lease.Data[6:12], swapMac)

	// If EtherType == 0x0806(ARP)
	if lease.Data[12] == 0x08 && lease.Data[13] == 0x06 {
		return HandleARP(lease, linkMac, linkIP)
	}

	// If EtherType == 0x0800(IPv4)
	if lease.Data[12] == 0x08 && lease.Data[13] == 0x00 {
		return EchoReply(lease, linkIP)
	}

	err = lease.Release()
	if err != nil {
		return fmt.Errorf("release lease: %w", err)
	}
	return nil
}

// Since we have no network stack we also need to respond to ARP messages
func HandleARP(lease *gobpfld.XSKLease, linkMac net.HardwareAddr, linkIP net.IP) error {
	// 14-20 left unchanged

	// Change opcode from request to reply
	lease.Data[21] = 2

	// Copy the MAC of our link into the ethernet SRC
	copy(lease.Data[6:12], linkMac)
	// Copy the MAC of our link into the ARP target (will be swapped to source)
	copy(lease.Data[32:38], linkMac)

	// Ignore ARP requests whic hare not for us, otherwise we are ARP spoofing
	if !net.IP(lease.Data[38:42]).Equal(linkIP) {
		err := lease.Release()
		if err != nil {
			return fmt.Errorf("release lease: %w", err)
		}
		return nil
	}

	fmt.Println("respond to arp")

	// Swap sender and target fields
	swap := make([]byte, 10)
	copy(swap, lease.Data[22:32])
	copy(lease.Data[22:32], lease.Data[32:42])
	copy(lease.Data[32:42], swap)

	err := lease.Write()
	if err != nil {
		return fmt.Errorf("release lease: %w", err)
	}
	return nil
}

func EchoReply(lease *gobpfld.XSKLease, linkIP net.IP) error {
	var err error

	// Ignore IP traffic that is not for us
	if !net.IP(lease.Data[30:34]).Equal(linkIP) {
		err = lease.Release()
		if err != nil {
			return fmt.Errorf("release lease: %w", err)
		}
		return nil
	}

	fmt.Printf("IPv4 Src: %X, Dst: %X, Proto: %X\n", lease.Data[26:30], lease.Data[30:34], lease.Data[23])
	swapIP := make([]byte, 4)
	copy(swapIP, lease.Data[26:30])
	copy(lease.Data[26:30], lease.Data[30:34])
	copy(lease.Data[30:34], swapIP)

	fmt.Printf("IPv4 Checksum in: %X\n", lease.Data[24:26])

	// Zero the checksum
	lease.Data[24] = 0x00
	lease.Data[25] = 0x00

	// Calculate header checksum
	// https://en.wikipedia.org/wiki/IPv4_header_checksum
	var sum uint32
	for i := 14; i < 34; i += 2 {
		sum += uint32(lease.Data[i]) << 8
		sum += uint32(lease.Data[i+1])
	}
	for {
		// Break when sum is less or equals to 0xFFFF
		if sum <= 65535 {
			break
		}
		// Add carry to the sum
		sum = (sum >> 16) + uint32(uint16(sum))
	}
	checkSum := ^uint16(sum)
	lease.Data[24] = byte(checkSum >> 8)
	lease.Data[25] = byte(checkSum & 0xFF)

	fmt.Printf("IPv4 Checksum out: %X\n", lease.Data[24:26])

	// If Protocol != 0x01(ICMPv4)
	if lease.Data[23] != 0x01 {
		err = lease.Release()
		if err != nil {
			return fmt.Errorf("release lease: %w", err)
		}
		return nil
	}

	fmt.Printf("ICMPv4 Type: %X, Code: %X\n", lease.Data[34], lease.Data[35])
	// Set type to 0 to get 0,0 = Echo Reply
	lease.Data[34] = 0

	fmt.Printf("ICMPv4 Checksum in: %X\n", lease.Data[36:38])

	// clear icmp checksum
	lease.Data[36] = 0x00
	lease.Data[37] = 0x00

	// Calculate ICMP checksum
	sum = 0
	for i := 34; i < len(lease.Data); i += 2 {
		sum += uint32(lease.Data[i]) << 8
		sum += uint32(lease.Data[i+1])
	}
	for {
		// Break when sum is less or equals to 0xFFFF
		if sum <= 65535 {
			break
		}
		// Add carry to the sum
		sum = (sum >> 16) + uint32(uint16(sum))
	}
	checkSum = ^uint16(sum)
	lease.Data[36] = byte(checkSum >> 8)
	lease.Data[37] = byte(checkSum & 0xFF)

	fmt.Printf("ICMPv4 Checksum out: %X\n", lease.Data[36:38])

	// Now that we have converted the request into a reply in the same memory buffer
	// we can just write this buffer back to the network interface
	err = lease.Write()
	if err != nil {
		return fmt.Errorf("write lease: %w", err)
	}
	return nil
}