probe/pcap/packet.go

package pcap

import (
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
)

type PacketType int

const (
	PacketTypeUnknown PacketType = iota
	PacketTypeARP
	PacketTypeTCP
	PacketTypeUDP
)

func getPacketType(packet gopacket.Packet) PacketType {
	if packet == nil {
		return PacketTypeUnknown
	}
	layer := packet.Layer(layers.LayerTypeARP)
	if layer != nil {
		return PacketTypeARP
	}

	layer = packet.Layer(layers.LayerTypeTCP)
	if layer != nil {
		if _, ok := layer.(*layers.TCP); ok {
			return PacketTypeTCP
		}
	}

	layer = packet.Layer(layers.LayerTypeUDP)
	if layer != nil {
		if _, ok := layer.(*layers.UDP); ok {
			return PacketTypeUDP
		}
	}
	return PacketTypeUnknown
}

func handlePacket(ps *pCapScan, packet gopacket.Packet) {
	switch getPacketType(packet) {
	case PacketTypeARP:
		handlePacketARP(ps, packet)
	case PacketTypeTCP:
		handlePacketTCP(ps, packet)
	case PacketTypeUDP:
		handlePacketUDP(ps, packet)
	default:
	}
}

func handlePacketARP(ps *pCapScan, packet gopacket.Packet) {
	ps.arpListenerChanMtx.RLock()
	defer ps.arpListenerChanMtx.RUnlock()

	arpLayer := packet.Layer(layers.LayerTypeARP)
	arp := arpLayer.(*layers.ARP)

	for _, ch := range ps.arpListenerChans {
		ch <- arp
	}
}

func handlePacketTCP(ps *pCapScan, packet gopacket.Packet) {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if nil == ipLayer {
		return
	}

	ip := ipLayer.(*layers.IPv4).SrcIP.String()

	ps.tcpListenerChanMtx.RLock()
	defer ps.tcpListenerChanMtx.RUnlock()

	chs, ok := ps.tcpListenerChans[ip]
	if ok {

		layer := packet.Layer(layers.LayerTypeTCP)
		tcp, _ := layer.(*layers.TCP)

		for _, ch := range chs {
			ch <- tcp
		}
	}
}

func handlePacketUDP(ps *pCapScan, packet gopacket.Packet) {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if nil == ipLayer {
		return
	}

	ip := ipLayer.(*layers.IPv4).SrcIP.String()

	ps.udpListenerChanMtx.RLock()
	defer ps.udpListenerChanMtx.RUnlock()

	chs, ok := ps.udpListenerChans[ip]
	if ok {
		for _, ch := range chs {
			ch <- packet
		}
	}
}