added ParseNextLayer function

Author: shadowy-pycoder

layers/dns.go

@@ -8,6 +8,8 @@ import (
 	"strings"
 )
 
+// TODO (shadowy-pycoder): add MarshalJSON
+
 const headerSizeDNS = 12
 
 type DNSFlags struct {
@@ -835,6 +837,9 @@ func parseRData(payload, tail []byte, typ uint16, rdl int) (fmt.Stringer, []byte
 			DataLen:            uint16(rdl),
 		}
 	case 65:
+		if len(tail) < 3 {
+			return nil, nil, ErrSliceBounds
+		}
 		priority := binary.BigEndian.Uint16(tail[0:2])
 		nameLength := tail[2]
 		var target string

layers/layers.go

@@ -2,27 +2,30 @@
 package layers
 
 import (
+	"bytes"
+	"encoding/binary"
 	"fmt"
+	"net/netip"
 	"unsafe"
 )
 
 const maxLenSummary = 110
 
-var LayerMap = map[string]Layer{
-	"ETH":    &EthernetFrame{},
-	"IPv4":   &IPv4Packet{},
-	"IPv6":   &IPv6Packet{},
-	"ARP":    &ARPPacket{},
-	"TCP":    &TCPSegment{},
-	"UDP":    &UDPSegment{},
-	"ICMP":   &ICMPSegment{},
-	"ICMPv6": &ICMPv6Segment{},
-	"DNS":    &DNSMessage{},
-	"FTP":    &FTPMessage{},
-	"HTTP":   &HTTPMessage{},
-	"SNMP":   &SNMPMessage{},
-	"SSH":    &SSHMessage{},
-	"TLS":    &TLSMessage{},
+var Layers = []string{
+	"ETH",
+	"IPv4",
+	"IPv6",
+	"ARP",
+	"TCP",
+	"UDP",
+	"ICMP",
+	"ICMPv6",
+	"DNS",
+	"FTP",
+	"HTTP",
+	"SNMP",
+	"SSH",
+	"TLS",
 }
 
 var (
@@ -46,6 +49,129 @@ type Layer interface {
 	Summary() string
 }
 
+func parseNextLayerFallback(data []byte) Layer {
+	for _, layer := range Layers {
+		next := GetNextLayer(layer)
+		if err := next.Parse(data); err == nil {
+			return next
+		}
+	}
+	return nil
+}
+
+func parseNextLayerFromBytes(data []byte) Layer {
+	buf := make([]byte, 0, len(data))
+	buf = append(buf, data...)
+	var next Layer
+	firstByte := buf[0]
+	if firstByte >= 0x45 && firstByte <= 0x4F {
+		next = GetNextLayer("IPv4")
+		if err := next.Parse(buf); err == nil {
+			return next
+		}
+	}
+	if firstByte>>4 == 6 {
+		next = GetNextLayer("IPv6")
+		if err := next.Parse(buf); err == nil {
+			return next
+		}
+	}
+	if firstByte == HandshakeTLSVal {
+		next = GetNextLayer("TLS")
+		if err := next.Parse(buf); err == nil {
+			return next
+		}
+	}
+	if len(buf) > 3 {
+		b1 := binary.BigEndian.Uint16(buf[0:2])
+		b2 := binary.BigEndian.Uint16(buf[2:4])
+		if b1 == 1 && (b2 == 0x0800 || b2 == 0x86dd) {
+			next = GetNextLayer("ARP")
+			if err := next.Parse(buf); err == nil {
+				return next
+			}
+		}
+	}
+	if len(buf) > 15 {
+		b1 := binary.BigEndian.Uint16(buf[12:14])
+		if b1 == 0x0806 || b1 == 0x0800 || b1 == 0x86dd {
+			next = GetNextLayer("ETH")
+			if err := next.Parse(buf); err == nil {
+				return next
+			}
+		}
+	}
+	if bytes.Contains(buf, protohttp10) || bytes.Contains(buf, protohttp11) {
+		next = GetNextLayer("HTTP")
+		if err := next.Parse(buf); err == nil {
+			return next
+		}
+	}
+	if bytes.Contains(buf, protoSSH) {
+		next = GetNextLayer("SSH")
+		if err := next.Parse(buf); err == nil {
+			return next
+		}
+	}
+	return nil
+}
+
+func addrMatch(src, dst *netip.AddrPort, ports []uint16) bool {
+	var srcPort, dstPort uint16
+	if src != nil {
+		srcPort = src.Port()
+	}
+	if dst != nil {
+		dstPort = src.Port()
+	}
+	for _, port := range ports {
+		if srcPort == port || dstPort == port {
+			return true
+		}
+	}
+	return false
+}
+
+func parseNextLayerFromAddress(data []byte, src, dst *netip.AddrPort) Layer {
+	var next Layer
+	switch {
+	case addrMatch(src, dst, []uint16{53, 5353, 853, 5355}):
+		next = GetNextLayer("DNS")
+	case addrMatch(src, dst, []uint16{80, 8080, 8000, 8888, 81, 591, 5911}):
+		next = GetNextLayer("HTTP")
+	case addrMatch(src, dst, []uint16{161, 162, 10161, 10162, 1161, 2161}):
+		next = GetNextLayer("SNMP")
+	case addrMatch(src, dst, []uint16{21, 20, 2121, 8021}):
+		next = GetNextLayer("FTP")
+	case addrMatch(src, dst, []uint16{22, 2222, 2200, 222, 2022}):
+		next = GetNextLayer("SSH")
+	case addrMatch(src, dst, []uint16{443, 465, 993, 995, 8443, 9443, 10443, 8444}):
+		next = GetNextLayer("TLS")
+	default:
+		return nil
+	}
+	if err := next.Parse(data); err == nil {
+		return next
+	} else {
+		return nil
+	}
+}
+
+func ParseNextLayer(data []byte, src, dst *netip.AddrPort) Layer {
+	buf := make([]byte, 0, len(data))
+	buf = append(buf, data...)
+	var next Layer
+	if src != nil || dst != nil {
+		if next = parseNextLayerFromAddress(buf, src, dst); next != nil {
+			return next
+		}
+	}
+	if next = parseNextLayerFromBytes(buf); next != nil {
+		return next
+	}
+	return parseNextLayerFallback(buf)
+}
+
 func GetNextLayer(layer string) Layer {
 	// TODO (shadowy-pycoder): add this to NextLayer, choose by ports, parse, use fallback on error
 	switch layer {

layers/ssh.go

@@ -9,6 +9,8 @@ import (
 
 const messageSizeSSH = 6
 
+var protoSSH = []byte("SSH-")
+
 type Message struct {
 	PacketLength     uint32
 	PaddingLength    uint8
@@ -87,8 +89,6 @@ func (s *SSHMessage) Parse(data []byte) error {
 	}
 	buf := make([]byte, 0, len(data))
 	buf = append(buf, data...)
-	s.Protocol = ""
-	s.Messages = nil
 	if bytes.HasSuffix(buf, crlf) {
 		s.Protocol = bytesToStr(bytes.TrimSuffix(buf, crlf))
 		return nil