add bcc socket filter support and example

Author: zhangbo1882

bcc/module.go

@@ -197,6 +197,11 @@ func (bpf *Module) LoadUprobe(name string) (int, error) {
 	return bpf.Load(name, C.BPF_PROG_TYPE_KPROBE, 0, 0)
 }
 
+// LoadSocketFilter loads a program of type BPF_PROG_TYPE_SOCKET_FILTER.
+func (bpf *Module) LoadSocketFilter(name string) (int, error) {
+	return bpf.Load(name, C.BPF_PROG_TYPE_SOCKET_FILTER, 0, 0)
+}
+
 // Load a program.
 func (bpf *Module) Load(name string, progType int, logLevel, logSize uint) (int, error) {
 	fd, ok := bpf.funcs[name]
@@ -269,6 +274,15 @@ func (bpf *Module) attachUProbe(evName string, attachType uint32, path string, a
 	return nil
 }
 
+// AttachSocketFilter attach a socket filter to a function
+func (bpf *Module) AttachSocketFilter(sockFd, socketFilterFd int) error {
+	ret, err := C.bpf_attach_socket(C.int(sockFd), C.int(socketFilterFd))
+	if ret != 0 {
+		return fmt.Errorf("error attaching BPF socket filter: %v", err)
+	}
+	return nil
+}
+
 // AttachKprobe attaches a kprobe fd to a function.
 func (bpf *Module) AttachKprobe(fnName string, fd int, maxActive int) error {
 	evName := "p_" + kprobeRegexp.ReplaceAllString(fnName, "_")

examples/bcc/protocol_count/README.md

@@ -0,0 +1,32 @@
+# usage
+
+1 In current directory (example/bcc/protocol_count)
+
+```bash
+go build net_protocol.go
+sudo ./net_protocol
+```
+
+2  Open another terminal
+
+```bash
+ping 127.0.0.1 -c 10
+```
+
+3  Result
+
+```
+TCP: 0, UDP: 0, ICMP: 0
+TCP: 0, UDP: 0, ICMP: 4
+TCP: 0, UDP: 0, ICMP: 24
+TCP: 0, UDP: 0, ICMP: 40
+TCP: 0, UDP: 0, ICMP: 40
+TCP: 0, UDP: 0, ICMP: 40
+TCP: 4, UDP: 0, ICMP: 40
+```
+
+# Misc
+
+Since we  run ping for the loop interface, there are  4 packets for one ping including( egress  send, ingress receive, egress reply and ingress reply)
+
+

examples/bcc/protocol_count/net_protocol.c

@@ -0,0 +1,32 @@
+#include <linux/if_packet.h>
+#include <linux/in.h>
+#include <linux/ip.h>
+#include <linux/string.h>
+#include <linux/tcp.h>
+#include <linux/types.h>
+#include <linux/udp.h>
+
+#ifndef offsetof
+#define offsetof(TYPE, MEMBER) ((size_t) & ((TYPE *)0)->MEMBER)
+#endif
+
+#define SEC(NAME) __attribute__((section(NAME), used))
+
+unsigned long long load_byte(void *skb,unsigned long long off) asm("llvm.bpf.load.byte");
+BPF_HASH(countmap,u32,u32,32);
+int protocol_count(struct __sk_buff *skb) {
+    
+  int proto = load_byte(skb, ETH_HLEN + offsetof(struct iphdr, protocol));
+  int one = 1;
+  
+  int *el = countmap.lookup(&proto);
+  
+  if (el) {
+    (*el)++;
+  } else {
+    el = &one;
+  }
+  countmap.update(&proto,el);
+  
+  return 0;
+}

examples/bcc/protocol_count/net_protocol.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io/ioutil"
+	"syscall"
+	"time"
+	"unsafe"
+
+	"github.com/iovisor/gobpf/bcc"
+)
+
+const (
+	PROTOCOL_FUNC  = "protocol_count"
+	PROTOCOL_COUNT = "./net_protocol.c"
+	LOOP           = 1
+)
+
+func ReadBPFFile(file string) (string, error) {
+	content, err := ioutil.ReadFile(file)
+	if err != nil {
+		return "", err
+	}
+	return string(content), nil
+}
+
+func Htons(i uint16) uint16 {
+	b := make([]byte, 2)
+	binary.BigEndian.PutUint16(b, i)
+	return *(*uint16)(unsafe.Pointer(&b[0]))
+}
+
+func OpenRawSock(index int) (int, error) {
+	sock, err := syscall.Socket(syscall.AF_PACKET, syscall.SOCK_RAW, int(Htons(syscall.ETH_P_ALL)))
+	if err != nil {
+		return 0, err
+	}
+	sll := syscall.SockaddrLinklayer{
+		Ifindex:  index,
+		Protocol: Htons(syscall.ETH_P_ALL),
+	}
+	if err := syscall.Bind(sock, &sll); err != nil {
+		return 0, err
+	}
+	return sock, nil
+}
+func ProtocolCount(cSource string) {
+	source, err := ReadBPFFile(cSource)
+	if err != nil {
+		fmt.Errorf("read BPF file error: %v", err)
+		return
+	}
+	m := bcc.NewModule(source, []string{})
+	defer m.Close()
+
+	socketFilter, err := m.LoadSocketFilter(PROTOCOL_FUNC)
+	if err != nil {
+		fmt.Errorf("socket filter %s not found, err: %v", PROTOCOL_FUNC, err)
+		return
+	}
+
+	fd, err := OpenRawSock(LOOP)
+	if err != nil {
+		fmt.Errorf("unable to open a raw socket: %s", err)
+		return
+	}
+	defer syscall.Close(fd)
+
+	if err := m.AttachSocketFilter(fd, socketFilter); err != nil {
+		fmt.Errorf("failed trying to attach socket filter: %s", err)
+		return
+	}
+
+	table := bcc.NewTable(m.TableId("countmap"), m)
+	var tcp, udp, icmp, leafInt, keyInt uint32
+	hostEndian := bcc.GetHostByteOrder()
+	for {
+		iter := table.Iter()
+		for iter.Next() {
+			key, leaf := iter.Key(), iter.Leaf()
+			if err := binary.Read(bytes.NewBuffer(key), hostEndian, &keyInt); err != nil {
+				continue
+			}
+			if err := binary.Read(bytes.NewBuffer(leaf), hostEndian, &leafInt); err != nil {
+				continue
+			}
+			switch keyInt {
+			case syscall.IPPROTO_TCP:
+				tcp = leafInt
+			case syscall.IPPROTO_UDP:
+				udp = leafInt
+			case syscall.IPPROTO_ICMP:
+				icmp = leafInt
+			}
+		}
+		fmt.Printf("TCP: %v, UDP: %v, ICMP: %v\n", tcp, udp, icmp)
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func main() {
+	ProtocolCount(PROTOCOL_COUNT)
+}