add IP distance extraction example

example_sock_extract_dist_test.go

@@ -0,0 +1,246 @@
+package ebpf_test
+
+// This code is derived from https://github.com/cloudflare/cloudflare-blog/tree/master/2018-03-ebpf
+//
+// Copyright (c) 2015-2017 Cloudflare, Inc. All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//    * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//    * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//    * Neither the name of the Cloudflare, Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import (
+	"fmt"
+	"net"
+	"syscall"
+	"time"
+
+	"github.com/newtools/ebpf"
+	"github.com/newtools/ebpf/asm"
+)
+
+// ExampleExtractDistance shows how to extract the network distance of
+// an IP host.
+func Example_extractDistance() {
+	var addr = &net.TCPAddr{
+		IP:   net.ParseIP("1.1.1.1"),
+		Port: 53,
+	}
+
+	// Call our own socket functions, so that we can attach
+	// the eBPF before calling connect. On Go 1.11 you can use
+	// Dialer.Control to achieve the same thing.
+	fd, err := socket(addr)
+	if err != nil {
+		panic(err)
+	}
+	defer syscall.Close(fd)
+
+	ttls, err := attachBPF(fd)
+	if err != nil {
+		panic(err)
+	}
+	defer ttls.Close()
+
+	if err := connect(fd, addr, time.Second); err != nil {
+		panic(err)
+	}
+
+	minDist, err := minDistance(ttls)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(addr, "is", minDist, "hops away")
+
+	if err := detachBPF(fd); err != nil {
+		panic(err)
+	}
+}
+
+func attachBPF(fd int) (*ebpf.Map, error) {
+	const ETH_P_IPV6 uint16 = 0x86DD
+	const SO_ATTACH_BPF = 50
+
+	ttls, err := ebpf.NewMap(&ebpf.MapSpec{
+		Type:       ebpf.Hash,
+		KeySize:    4,
+		ValueSize:  8,
+		MaxEntries: 4,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	insns := asm.Instructions{
+		// r1 has ctx
+		// r0 = ctx[16] (aka protocol)
+		asm.LoadMem(asm.R0, asm.R1, 16, asm.Word),
+
+		// Perhaps ipv6
+		asm.LoadImm(asm.R2, int64(ETH_P_IPV6), asm.Half),
+		asm.HostTo(asm.BE, asm.R2, asm.Half),
+		asm.JEq.Reg(asm.R0, asm.R2, "ipv6"),
+
+		// otherwise assume ipv4
+		// 8th byte in IPv4 is TTL
+		// LDABS requires ctx in R6
+		asm.Mov.Reg(asm.R6, asm.R1),
+		asm.LoadAbs(-0x100000+8, asm.Byte),
+		asm.Ja.Label("store-ttl"),
+
+		// 7th byte in IPv6 is Hop count
+		// LDABS requires ctx in R6
+		asm.Mov.Reg(asm.R6, asm.R1).Sym("ipv6"),
+		asm.LoadAbs(-0x100000+7, asm.Byte),
+
+		// stash the load result into FP[-4]
+		asm.StoreMem(asm.RFP, -4, asm.R0, asm.Word).Sym("store-ttl"),
+		// stash the &FP[-4] into r2
+		asm.Mov.Reg(asm.R2, asm.RFP),
+		asm.Add.Imm(asm.R2, -4),
+
+		// r1 must point to map
+		asm.LoadImm(asm.R1, int64(ttls.FD()), asm.DWord),
+		asm.MapLookupElement.Call(),
+
+		// load ok? inc. Otherwise? jmp to mapupdate
+		asm.JEq.Imm(asm.R0, 0, "update-map"),
+		asm.Mov.Imm(asm.R1, 1),
+		asm.XAdd(asm.R0, asm.R1),
+		asm.Ja.Label("exit"),
+
+		// MapUpdate
+		// r1 has map ptr
+		asm.LoadImm(asm.R1, int64(ttls.FD()), asm.DWord),
+		// r2 has key -> &FP[-4]
+		asm.Mov.Reg(asm.R2, asm.RFP),
+		asm.Add.Imm(asm.R2, -4),
+		// r3 has value -> &FP[-16] , aka 1
+		asm.StoreImm(asm.RFP, -16, 1, asm.DWord),
+		asm.Mov.Reg(asm.R3, asm.RFP),
+		asm.Add.Imm(asm.R3, -16),
+		// r4 has flags, 0
+		asm.Mov.Imm(asm.R4, 0),
+		asm.MapUpdateElement.Call(),
+
+		// set exit code to -1, don't trunc packet
+		asm.Mov.Imm(asm.R0, -1).Sym("exit"),
+		asm.Return(),
+	}
+
+	prog, err := ebpf.NewProgram(&ebpf.ProgramSpec{
+		Type:         ebpf.SocketFilter,
+		License:      "GPL",
+		Instructions: insns,
+	})
+	if err != nil {
+		ttls.Close()
+		return nil, err
+	}
+	defer prog.Close()
+
+	err = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, SO_ATTACH_BPF, prog.FD())
+	if err != nil {
+		ttls.Close()
+		return nil, err
+	}
+
+	return ttls, nil
+}
+
+func detachBPF(fd int) error {
+	const SO_DETACH_BPF = 27
+	return syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, SO_DETACH_BPF, 0)
+}
+
+func minDistance(ttls *ebpf.Map) (int, error) {
+	var (
+		entries      = ttls.Iterate()
+		ttl, minDist uint32
+		count        uint64
+	)
+	for entries.Next(&ttl, &count) {
+		var dist uint32
+		switch {
+		case ttl > 128:
+			dist = 255 - ttl
+		case ttl > 64:
+			dist = 128 - ttl
+		case ttl > 32:
+			dist = 64 - ttl
+		default:
+			dist = 32 - ttl
+		}
+		if minDist > dist {
+			minDist = dist
+		}
+	}
+	return int(minDist), entries.Err()
+}
+
+func socket(dst *net.TCPAddr) (int, error) {
+	var domain int
+	if dst.IP.To4() != nil {
+		domain = syscall.AF_INET
+	} else {
+		domain = syscall.AF_INET6
+	}
+
+	return syscall.Socket(domain, syscall.SOCK_STREAM, 0)
+}
+
+func connect(fd int, dst *net.TCPAddr, timeout time.Duration) error {
+	var domain int
+	if dst.IP.To4() != nil {
+		domain = syscall.AF_INET
+	} else {
+		domain = syscall.AF_INET6
+	}
+
+	var sa syscall.Sockaddr
+	if domain == syscall.AF_INET {
+		var x [4]byte
+		copy(x[:], dst.IP.To4())
+		sa = &syscall.SockaddrInet4{Port: dst.Port, Addr: x}
+	} else {
+		var x [16]byte
+		copy(x[:], dst.IP.To16())
+		sa = &syscall.SockaddrInet6{Port: dst.Port, Addr: x}
+	}
+
+	if ns := timeout.Nanoseconds(); ns > 0 {
+		// Set SO_SNDTIMEO
+		var tv syscall.Timeval
+		tv.Sec = ns / 1000000000
+		tv.Usec = (ns - tv.Sec*1000000000) / 1000
+
+		if err := syscall.SetsockoptTimeval(fd, syscall.SOL_SOCKET, syscall.SO_SNDTIMEO, &tv); err != nil {
+			return err
+		}
+	}
+
+	// This is blocking.
+	return syscall.Connect(fd, sa)
+}