优化抓包性能

vendor/github.com/google/gopacket/pcapgo/capture.go

@@ -0,0 +1,286 @@
+// Copyright 2012 Google, Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style license
+// that can be found in the LICENSE file in the root of the source
+// tree.
+// +build linux,go1.7
+
+package pcapgo
+
+import (
+	"fmt"
+	"net"
+	"runtime"
+	"sync"
+	"syscall"
+	"time"
+	"unsafe"
+
+	"golang.org/x/net/bpf"
+	"golang.org/x/sys/unix"
+
+	"github.com/google/gopacket"
+)
+
+var hdrLen = unix.CmsgSpace(0)
+var auxLen = unix.CmsgSpace(int(unsafe.Sizeof(unix.TpacketAuxdata{})))
+var timensLen = unix.CmsgSpace(int(unsafe.Sizeof(unix.Timespec{})))
+var timeLen = unix.CmsgSpace(int(unsafe.Sizeof(unix.Timeval{})))
+
+func htons(data uint16) uint16 { return data<<8 | data>>8 }
+
+// EthernetHandle holds shared buffers and file descriptor of af_packet socket
+type EthernetHandle struct {
+	fd     int
+	buffer []byte
+	oob    []byte
+	ancil  []interface{}
+	mu     sync.Mutex
+	intf   int
+	addr   net.HardwareAddr
+}
+
+// readOne reads a packet from the handle and returns a capture info + vlan info
+func (h *EthernetHandle) readOne() (ci gopacket.CaptureInfo, vlan int, haveVlan bool, err error) {
+	// we could use unix.Recvmsg, but that does a memory allocation (for the returned sockaddr) :(
+	var msg unix.Msghdr
+	var sa unix.RawSockaddrLinklayer
+
+	msg.Name = (*byte)(unsafe.Pointer(&sa))
+	msg.Namelen = uint32(unsafe.Sizeof(sa))
+
+	var iov unix.Iovec
+	if len(h.buffer) > 0 {
+		iov.Base = &h.buffer[0]
+		iov.SetLen(len(h.buffer))
+	}
+	msg.Iov = &iov
+	msg.Iovlen = 1
+
+	if len(h.oob) > 0 {
+		msg.Control = &h.oob[0]
+		msg.SetControllen(len(h.oob))
+	}
+
+	// use msg_trunc so we know packet size without auxdata, which might be missing
+	n, _, e := syscall.Syscall(unix.SYS_RECVMSG, uintptr(h.fd), uintptr(unsafe.Pointer(&msg)), uintptr(unix.MSG_TRUNC))
+
+	if e != 0 {
+		return gopacket.CaptureInfo{}, 0, false, fmt.Errorf("couldn't read packet: %s", e)
+	}
+
+	if sa.Family == unix.AF_PACKET {
+		ci.InterfaceIndex = int(sa.Ifindex)
+	} else {
+		ci.InterfaceIndex = h.intf
+	}
+
+	// custom aux parsing so we don't allocate stuff (unix.ParseSocketControlMessage allocates a slice)
+	// we're getting at most 2 cmsgs anyway and know which ones they are (auxdata + timestamp(ns))
+	oob := h.oob[:msg.Controllen]
+	gotAux := false
+
+	for len(oob) > hdrLen { // > hdrLen, because we also need something after the cmsg header
+		hdr := (*unix.Cmsghdr)(unsafe.Pointer(&oob[0]))
+		switch {
+		case hdr.Level == unix.SOL_PACKET && hdr.Type == unix.PACKET_AUXDATA && len(oob) >= auxLen:
+			aux := (*unix.TpacketAuxdata)(unsafe.Pointer(&oob[hdrLen]))
+			ci.CaptureLength = int(n)
+			ci.Length = int(aux.Len)
+			vlan = int(aux.Vlan_tci)
+			haveVlan = (aux.Status & unix.TP_STATUS_VLAN_VALID) != 0
+			gotAux = true
+		case hdr.Level == unix.SOL_SOCKET && hdr.Type == unix.SO_TIMESTAMPNS && len(oob) >= timensLen:
+			tstamp := (*unix.Timespec)(unsafe.Pointer(&oob[hdrLen]))
+			ci.Timestamp = time.Unix(int64(tstamp.Sec), int64(tstamp.Nsec))
+		case hdr.Level == unix.SOL_SOCKET && hdr.Type == unix.SO_TIMESTAMP && len(oob) >= timeLen:
+			tstamp := (*unix.Timeval)(unsafe.Pointer(&oob[hdrLen]))
+			ci.Timestamp = time.Unix(int64(tstamp.Sec), int64(tstamp.Usec)*1000)
+		}
+		oob = oob[unix.CmsgSpace(int(hdr.Len))-hdrLen:]
+	}
+
+	if !gotAux {
+		// fallback for no aux cmsg
+		ci.CaptureLength = int(n)
+		ci.Length = int(n)
+		haveVlan = false
+	}
+
+	// fix up capture length if we needed to truncate
+	if ci.CaptureLength > len(h.buffer) {
+		ci.CaptureLength = len(h.buffer)
+	}
+
+	if ci.Timestamp.IsZero() {
+		// we got no timestamp info -> emulate it
+		ci.Timestamp = time.Now()
+	}
+
+	return ci, vlan, haveVlan, nil
+}
+
+// ReadPacketData implements gopacket.PacketDataSource. If this was captured on a vlan, the vlan id will be in the AncillaryData[0]
+func (h *EthernetHandle) ReadPacketData() ([]byte, gopacket.CaptureInfo, error) {
+	h.mu.Lock()
+	ci, vlan, haveVlan, err := h.readOne()
+	if err != nil {
+		h.mu.Unlock()
+		return nil, gopacket.CaptureInfo{}, fmt.Errorf("couldn't read packet data: %s", err)
+	}
+
+	b := make([]byte, ci.CaptureLength)
+	copy(b, h.buffer)
+	h.mu.Unlock()
+
+	if haveVlan {
+		ci.AncillaryData = []interface{}{vlan}
+
+	}
+
+	return b, ci, nil
+}
+
+// ZeroCopyReadPacketData implements gopacket.ZeroCopyPacketDataSource. If this was captured on a vlan, the vlan id will be in the AncillaryData[0].
+// This function does not allocate memory. Beware that the next call to ZeroCopyReadPacketData will overwrite existing slices (returned data AND AncillaryData)!
+// Due to shared buffers this must not be called concurrently
+func (h *EthernetHandle) ZeroCopyReadPacketData() ([]byte, gopacket.CaptureInfo, error) {
+	ci, vlan, haveVlan, err := h.readOne()
+	if err != nil {
+		return nil, gopacket.CaptureInfo{}, fmt.Errorf("couldn't read packet data: %s", err)
+	}
+
+	if haveVlan {
+		h.ancil[0] = vlan
+		ci.AncillaryData = h.ancil
+	}
+
+	return h.buffer[:ci.CaptureLength], ci, nil
+}
+
+// Close closes the underlying socket
+func (h *EthernetHandle) Close() {
+	if h.fd != -1 {
+		unix.Close(h.fd)
+		h.fd = -1
+		runtime.SetFinalizer(h, nil)
+	}
+}
+
+// SetCaptureLength sets the maximum capture length to the given value
+func (h *EthernetHandle) SetCaptureLength(len int) error {
+	if len < 0 {
+		return fmt.Errorf("illegal capture length %d. Must be at least 0", len)
+	}
+	h.buffer = make([]byte, len)
+	return nil
+}
+
+// GetCaptureLength returns the maximum capture length
+func (h *EthernetHandle) GetCaptureLength() int {
+	return len(h.buffer)
+}
+
+// SetBPF attaches the given BPF filter to the socket. After this, only the packets for which the filter returns a value greater than zero are received.
+// If a filter was already attached, it will be overwritten. To remove the filter, provide an empty slice.
+func (h *EthernetHandle) SetBPF(filter []bpf.RawInstruction) error {
+	if len(filter) == 0 {
+		return unix.SetsockoptInt(h.fd, unix.SOL_SOCKET, unix.SO_DETACH_FILTER, 0)
+	}
+	f := make([]unix.SockFilter, len(filter))
+	for i := range filter {
+		f[i].Code = filter[i].Op
+		f[i].Jf = filter[i].Jf
+		f[i].Jt = filter[i].Jt
+		f[i].K = filter[i].K
+	}
+	fprog := &unix.SockFprog{
+		Len:    uint16(len(filter)),
+		Filter: &f[0],
+	}
+	return unix.SetsockoptSockFprog(h.fd, unix.SOL_SOCKET, unix.SO_ATTACH_FILTER, fprog)
+}
+
+// LocalAddr returns the local network address
+func (h *EthernetHandle) LocalAddr() net.HardwareAddr {
+	// Hardware Address might have changed. Fetch new one and fall back to the stored one if fetching interface fails
+	intf, err := net.InterfaceByIndex(h.intf)
+	if err == nil {
+		h.addr = intf.HardwareAddr
+	}
+	return h.addr
+}
+
+// SetPromiscuous sets promiscous mode to the required value. If it is enabled, traffic not destined for the interface will also be captured.
+func (h *EthernetHandle) SetPromiscuous(b bool) error {
+	mreq := unix.PacketMreq{
+		Ifindex: int32(h.intf),
+		Type:    unix.PACKET_MR_PROMISC,
+	}
+
+	opt := unix.PACKET_ADD_MEMBERSHIP
+	if !b {
+		opt = unix.PACKET_DROP_MEMBERSHIP
+	}
+
+	return unix.SetsockoptPacketMreq(h.fd, unix.SOL_PACKET, opt, &mreq)
+}
+
+// Stats returns number of packets and dropped packets. This will be the number of packets/dropped packets since the last call to stats (not the cummulative sum!).
+func (h *EthernetHandle) Stats() (*unix.TpacketStats, error) {
+	return unix.GetsockoptTpacketStats(h.fd, unix.SOL_PACKET, unix.PACKET_STATISTICS)
+}
+
+// NewEthernetHandle implements pcap.OpenLive for network devices.
+// If you want better performance have a look at github.com/google/gopacket/afpacket.
+// SetCaptureLength can be used to limit the maximum capture length.
+func NewEthernetHandle(ifname string) (*EthernetHandle, error) {
+	intf, err := net.InterfaceByName(ifname)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't query interface %s: %s", ifname, err)
+	}
+
+	fd, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, int(htons(unix.ETH_P_ALL)))
+	if err != nil {
+		return nil, fmt.Errorf("couldn't open packet socket: %s", err)
+	}
+
+	addr := unix.SockaddrLinklayer{
+		Protocol: htons(unix.ETH_P_ALL),
+		Ifindex:  intf.Index,
+	}
+
+	if err := unix.Bind(fd, &addr); err != nil {
+		return nil, fmt.Errorf("couldn't bind to interface %s: %s", ifname, err)
+	}
+
+	ooblen := 0
+
+	if err := unix.SetsockoptInt(fd, unix.SOL_PACKET, unix.PACKET_AUXDATA, 1); err != nil {
+		// we can't get auxdata -> no vlan info
+	} else {
+		ooblen += auxLen
+	}
+
+	if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_TIMESTAMPNS, 1); err != nil {
+		// no nanosecond resolution :( -> try ms
+		if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_TIMESTAMP, 1); err != nil {
+			// if this doesn't work we well use time.Now() -> ignore errors here
+		} else {
+			ooblen += timeLen
+		}
+	} else {
+		ooblen += timensLen
+	}
+
+	handle := &EthernetHandle{
+		fd:     fd,
+		buffer: make([]byte, intf.MTU),
+		oob:    make([]byte, ooblen),
+		ancil:  make([]interface{}, 1),
+		intf:   intf.Index,
+		addr:   intf.HardwareAddr,
+	}
+	runtime.SetFinalizer(handle, (*EthernetHandle).Close)
+	return handle, nil
+}