probe/internal/pcap/packet.go

package pcap

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

type PacketType int

const (
	PacketTypeUnknown PacketType = iota
	PacketTypeARP
	PacketTypeICMPv4
	PacketTypeICMPv6
	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.LayerTypeICMPv4)
	if layer != nil {
		return PacketTypeICMPv4
	}

	layer = packet.Layer(layers.LayerTypeICMPv6)
	if layer != nil {
		return PacketTypeICMPv6
	}

	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 PacketTypeICMPv4:
		handlePacketICMP4(ps, packet)
	case PacketTypeICMPv6:
		handlePacketICMP6(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 handlePacketICMP4(ps *pCapScan, packet gopacket.Packet) {
	ps.icmp4ListenerChanMtx.RLock()
	defer ps.icmp4ListenerChanMtx.RUnlock()

	// icmp4Layer := packet.Layer(layers.LayerTypeICMPv4)
	// icmp4 := icmp4Layer.(*layers.ICMPv4)

	for _, ch := range ps.icmp4ListenerChans {
		ch <- packet
	}
}

func handlePacketICMP6(ps *pCapScan, packet gopacket.Packet) {
	ps.icmp6ListenerChanMtx.RLock()
	defer ps.icmp6ListenerChanMtx.RUnlock()

	// icmp6Layer := packet.Layer(layers.LayerTypeICMPv6)
	// icmp6 := icmp6Layer.(*layers.ICMPv6)

	for _, ch := range ps.icmp6ListenerChans {
		ch <- packet
	}
}

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

	ip := ipLayer.(*layers.IPv4).SrcIP.String()
	layer := packet.Layer(layers.LayerTypeTCP)
	tcp, _ := layer.(*layers.TCP)

	ps.tcpListenerChanMtx.RLock()
	defer func() {
		ps.tcpListenerChanMtx.RUnlock()
	}()

	chs, ok := ps.tcpListenerChans[ip]
	if ok {
		// for _, ch := range chs {
		// 	ch <- tcp
		// }
		for _, ch := range chs {
			select {
			case ch <- tcp:
			default:
			}
		}

	}
}

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
		}
	}
}