diff --git a/src/util/util.go b/src/util/util.go index a7a54562..c20421d5 100644 --- a/src/util/util.go +++ b/src/util/util.go @@ -103,11 +103,11 @@ func GetFlowKey(bs []byte) uint64 { // Get the IP protocol version from the packet switch bs[0] & 0xf0 { case 0x40: // IPv4 packet - // Check the packet meets minimum UDP packet length - if len(bs) >= 24 { - // Is the protocol TCP, UDP or SCTP? - if bs[9] == 0x06 || bs[9] == 0x11 || bs[9] == 0x84 { - ihl := bs[0] & 0x0f * 4 // Header length + ihl := (bs[0] & 0x0f) * 4 // whole IPv4 header length (min 20) + // 8 is minimum UDP packet length + if ihl >= 20 && len(bs)-int(ihl) >= 8 { + switch bs[9] /* protocol */ { + case 0x06 /* TCP */, 0x11 /* UDP */, 0x84 /* SCTP */ : flowkey = uint64(bs[9])<<32 /* proto */ | uint64(bs[ihl+0])<<24 | uint64(bs[ihl+1])<<16 /* sport */ | uint64(bs[ihl+2])<<8 | uint64(bs[ihl+3]) /* dport */ @@ -119,8 +119,8 @@ func GetFlowKey(bs []byte) uint64 { // If the flowlabel isn't present, make protokey from proto | sport | dport // if the packet meets minimum UDP packet length if flowkey == 0 && len(bs) >= 48 { - // Is the protocol TCP, UDP or SCTP? - if bs[6] == 0x06 || bs[6] == 0x11 || bs[6] == 0x84 { + switch bs[9] /* protocol */ { + case 0x06 /* TCP */, 0x11 /* UDP */, 0x84 /* SCTP */ : flowkey = uint64(bs[6])<<32 /* proto */ | uint64(bs[40])<<24 | uint64(bs[41])<<16 /* sport */ | uint64(bs[42])<<8 | uint64(bs[43]) /* dport */