misc: small fix or general refactoring i did not bother commenting

packets/queue.go

@@ -71,17 +71,7 @@ func NewQueue(iface *bnet.Endpoint) (*Queue, error) {
 	return q, nil
 }
 
-func (q *Queue) worker() {
-	for pkt := range q.source.Packets() {
-		if q.active == false {
-			return
-		}
-
-		pktSize := uint64(len(pkt.Data()))
-
-		atomic.AddUint64(&q.PktReceived, 1)
-		atomic.AddUint64(&q.Received, pktSize)
-
+func (q *Queue) trackProtocols(pkt gopacket.Packet) {
 	// gather protocols stats
 	pktLayers := pkt.Layers()
 	for _, layer := range pktLayers {
@@ -98,51 +88,65 @@ func (q *Queue) worker() {
 		}
 		q.Unlock()
 	}
+}
 
-		// check for new ipv4 endpoints
+func (q *Queue) trackActivity(eth *layers.Ethernet, ip4 *layers.IPv4, address net.IP, pktSize uint64) {
+	q.Lock()
+	defer q.Unlock()
+
+	// detrmine direction
+	isSent := bytes.Compare(address, ip4.SrcIP) == 0
+	// push to activity channel
+	q.Activities <- Activity{
+		IP:     address,
+		MAC:    eth.SrcMAC,
+		Source: isSent,
+	}
+
+	// initialize or update stats
+	addr := address.String()
+	if _, found := q.Traffic[addr]; found == false {
+		if isSent {
+			q.Traffic[addr] = &Traffic{Sent: pktSize}
+		} else {
+			q.Traffic[addr] = &Traffic{Received: pktSize}
+		}
+	} else {
+		if isSent {
+			q.Traffic[addr].Sent += pktSize
+		} else {
+			q.Traffic[addr].Received += pktSize
+		}
+	}
+}
+
+func (q *Queue) worker() {
+	for pkt := range q.source.Packets() {
+		if q.active == false {
+			return
+		}
+
+		q.trackProtocols(pkt)
+
+		pktSize := uint64(len(pkt.Data()))
+
+		atomic.AddUint64(&q.PktReceived, 1)
+		atomic.AddUint64(&q.Received, pktSize)
+
+		// decode eth and ipv4 layers
 		leth := pkt.Layer(layers.LayerTypeEthernet)
 		lip4 := pkt.Layer(layers.LayerTypeIPv4)
-
 		if leth != nil && lip4 != nil {
 			eth := leth.(*layers.Ethernet)
 			ip4 := lip4.(*layers.IPv4)
 
+			// coming from our network
 			if bytes.Compare(q.iface.IP, ip4.SrcIP) != 0 && q.iface.Net.Contains(ip4.SrcIP) {
-				q.Lock()
-				q.Activities <- Activity{
-					IP:     ip4.SrcIP,
-					MAC:    eth.SrcMAC,
-					Source: true,
+				q.trackActivity(eth, ip4, ip4.SrcIP, pktSize)
 			}
-
-				addr := ip4.SrcIP.String()
-				if _, found := q.Traffic[addr]; found == false {
-					q.Traffic[addr] = &Traffic{
-						Sent: pktSize,
-					}
-				} else {
-					q.Traffic[addr].Sent += pktSize
-				}
-				q.Unlock()
-			}
-
+			// coming to our network
 			if bytes.Compare(q.iface.IP, ip4.DstIP) != 0 && q.iface.Net.Contains(ip4.DstIP) {
-				q.Lock()
-				q.Activities <- Activity{
-					IP:     ip4.DstIP,
-					MAC:    eth.SrcMAC,
-					Source: false,
-				}
-
-				addr := ip4.DstIP.String()
-				if _, found := q.Traffic[addr]; found == false {
-					q.Traffic[addr] = &Traffic{
-						Received: pktSize,
-					}
-				} else {
-					q.Traffic[addr].Received += pktSize
-				}
-				q.Unlock()
+				q.trackActivity(eth, ip4, ip4.DstIP, pktSize)
 			}
 		}
 	}
@@ -152,17 +156,18 @@ func (q *Queue) Send(raw []byte) error {
 	q.Lock()
 	defer q.Unlock()
 
-	if q.active {
-		err := q.handle.WritePacketData(raw)
-		if err == nil {
-			q.Sent += uint64(len(raw))
-		} else {
-			q.Errors += 1
-		}
-		return err
-	} else {
+	if q.active == false {
 		return fmt.Errorf("Packet queue is not active.")
 	}
+
+	if err := q.handle.WritePacketData(raw); err != nil {
+		q.Errors += 1
+		return err
+	} else {
+		q.Sent += uint64(len(raw))
+	}
+
+	return nil
 }
 
 func (q *Queue) Stop() {