From: Mikko Rasa <tdb@tdb.fi>
Date: Sun, 15 Jun 2014 12:20:52 +0000 (+0300)
Subject: Restructure packet handling
X-Git-Url: http://git.tdb.fi/?p=netvis.git;a=commitdiff_plain;h=597a794cd648dd7b4c9d442fdc39f5961a38aba0

Restructure packet handling
---

diff --git a/source/host.h b/source/host.h
index ec2741d..c41610f 100644
--- a/source/host.h
+++ b/source/host.h
@@ -35,6 +35,7 @@ public:
 	void set_name(const std::string &);
 	const std::string &get_name() const { return name; }
 	void set_local(bool);
+	bool is_local() const { return local; }
 
 	void set_position(const Vector2 &);
 	const Vector2 &get_position() const { return pos; }
diff --git a/source/netvis.cpp b/source/netvis.cpp
index 52f5864..22231b1 100644
--- a/source/netvis.cpp
+++ b/source/netvis.cpp
@@ -9,10 +9,6 @@ Distributed unter the GPL
 #include <cstdlib>
 #include <cmath>
 #include <signal.h>
-#include <netinet/ether.h>
-#include <netinet/ip.h>
-#include <netinet/tcp.h>
-#include <netinet/udp.h>
 #include <msp/core/getopt.h>
 #include <msp/debug/profilingscope.h>
 #include <msp/gl/blend.h>
@@ -23,6 +19,7 @@ Distributed unter the GPL
 #include <msp/gl/projection.h>
 #include <msp/gl/texture2d.h>
 #include <msp/gl/transform.h>
+#include <msp/io/print.h>
 #include <msp/strings/format.h>
 #include <msp/time/units.h>
 #include <msp/time/utils.h>
@@ -278,79 +275,120 @@ Port &NetVis::get_port(unsigned number)
 	return *port;
 }
 
-void NetVis::capture_handler(unsigned char *user, const pcap_pkthdr *, const unsigned char *data)
+void NetVis::capture_handler(unsigned char *user, const pcap_pkthdr *cap, const unsigned char *data)
 {
 	NetVis *self = reinterpret_cast<NetVis *>(user);
 
+	CaptureContext ctx;
+	ctx.cap_hdr = cap;
 	const ethhdr *eth = reinterpret_cast<const ethhdr *>(data);
-	if(ntohs(eth->h_proto)==ETH_P_IP)
+	self->handle_ethernet(ctx, eth, cap->caplen);
+}
+
+void NetVis::handle_ethernet(CaptureContext &ctx, const ethhdr *eth, unsigned len)
+{
+	ctx.size = ctx.cap_hdr->len-sizeof(ethhdr);
+
+	int proto = ntohs(eth->h_proto);
+	if(proto==ETH_P_IP)
 	{
 		const iphdr *ip = reinterpret_cast<const iphdr *>(eth+1);
+		handle_ipv4(ctx, ip, len-sizeof(ethhdr));
+	}
+	else
+		IO::print("Unknown protocol in eth: %d\n", proto);
+}
 
-		unsigned size = ntohs(ip->tot_len);
+void NetVis::handle_ipv4(CaptureContext &ctx, const iphdr *ip, unsigned len)
+{
+	ctx.src_host = &get_host(ntohl(ip->saddr));
+	if((ntohl(ip->daddr)&0xFF)!=0xFF)
+		ctx.dst_host = &get_host(ntohl(ip->daddr));
 
-		Port *sport = 0;
-		Port *dport = 0;
-		if(ip->protocol==IPPROTO_TCP)
-		{
-			const tcphdr *tcp = reinterpret_cast<const tcphdr *>(ip+1);
-			sport = &self->get_port(ntohs(tcp->source));
-			dport = &self->get_port(ntohs(tcp->dest));
-		}
-		else if(ip->protocol==IPPROTO_UDP)
-		{
-			const udphdr *udp = reinterpret_cast<const udphdr *>(ip+1);
-			sport = &self->get_port(ntohs(udp->source));
-			dport = &self->get_port(ntohs(udp->dest));
-		}
+	if(ip->protocol==IPPROTO_TCP)
+	{
+		const tcphdr *tcp = reinterpret_cast<const tcphdr *>(ip+1);
+		handle_tcp(ctx, tcp, len-sizeof(iphdr));
+	}
+	else if(ip->protocol==IPPROTO_UDP)
+	{
+		const udphdr *udp = reinterpret_cast<const udphdr *>(ip+1);
+		handle_udp(ctx, udp, len-sizeof(iphdr));
+	}
+	else
+		IO::print("Unknown protocol in ip: %d\n", ip->protocol);
+}
 
-		Port *port = 0;
-		if(sport && dport)
-		{
-			if(sport->is_registered()!=dport->is_registered())
-			{
-				if(sport->is_registered())
-					port = sport;
-				else
-					port = dport;
-			}
-			else if(sport->get_number()<dport->get_number())
-				port = sport;
-			else
-				port = dport;
-		}
-		else
-			port = &self->get_port(0);
 
-		Host &shost = self->get_host(ntohl(ip->saddr));
-		Host *dhost = 0;
-		if((ntohl(ip->daddr)&0xFF)!=0xFF)
-			dhost = &self->get_host(ntohl(ip->daddr));
+void NetVis::handle_tcp(CaptureContext &ctx, const tcphdr *tcp, unsigned)
+{
+	ctx.src_port = &get_port(ntohs(tcp->source));
+	ctx.dst_port = &get_port(ntohs(tcp->dest));
+	handle_packet(ctx);
+}
+
+void NetVis::handle_udp(CaptureContext &ctx, const udphdr *udp, unsigned)
+{
+	ctx.src_port = &get_port(ntohs(udp->source));
+	ctx.dst_port = &get_port(ntohs(udp->dest));
+	handle_packet(ctx);
+}
 
-		float throttle = shost.send_packet();
-		if(throttle<1)
+void NetVis::handle_packet(CaptureContext &ctx)
+{
+	Port *port = 0;
+	if(ctx.src_port && ctx.dst_port)
+	{
+		if(ctx.src_port->is_registered()!=ctx.dst_port->is_registered())
 		{
-			self->packets.push_back(new Packet(shost, dhost, port->get_color(), size));
-			self->packets.back()->tick(-throttle*Msp::Time::sec);
+			if(ctx.src_port->is_registered())
+				port = ctx.src_port;
+			else
+				port = ctx.dst_port;
 		}
+		else if(ctx.src_port->get_number()<ctx.dst_port->get_number())
+			port = ctx.src_port;
+		else
+			port = ctx.dst_port;
+	}
+	else
+		port = &get_port(0);
 
-		shost.add_activity(size);
-		if(dhost)
-			dhost->add_activity(size);
-
-		if(sport)
-			sport->add_activity(size);
-		if(dport)
-			dport->add_activity(size);
-
-		if((ntohl(ip->saddr)&self->localnet_mask)==self->localnet)
-			self->history->activity(0, size);
-		else if((ntohl(ip->daddr)&self->localnet_mask)==self->localnet)
-			self->history->activity(size, 0);
+	float throttle = ctx.src_host->send_packet();
+	if(throttle<1)
+	{
+		packets.push_back(new Packet(*ctx.src_host, ctx.dst_host, port->get_color(), ctx.size));
+		packets.back()->tick(-throttle*Msp::Time::sec);
 	}
+
+	ctx.src_host->add_activity(ctx.size);
+	if(ctx.dst_host)
+		ctx.dst_host->add_activity(ctx.size);
+
+	if(ctx.src_port)
+		ctx.src_port->add_activity(ctx.size);
+	if(ctx.dst_port)
+		ctx.dst_port->add_activity(ctx.size);
+
+	bool local_src = ctx.src_host->is_local();
+	bool local_dst = (ctx.dst_host && ctx.dst_host->is_local());
+	if(local_src && !local_dst)
+		history->activity(0, ctx.size);
+	else if(local_dst && !local_src)
+		history->activity(ctx.size, 0);
 }
 
 void NetVis::sighandler(int)
 {
 	exit(0);
 }
+
+
+NetVis::CaptureContext::CaptureContext():
+	cap_hdr(0),
+	src_host(0),
+	src_port(0),
+	dst_host(0),
+	dst_port(0),
+	size(0)
+{ }
diff --git a/source/netvis.h b/source/netvis.h
index 46b96a5..68eed05 100644
--- a/source/netvis.h
+++ b/source/netvis.h
@@ -11,6 +11,10 @@ Distributed unter the GPL
 #include <list>
 #include <map>
 #include <pcap.h>
+#include <netinet/ether.h>
+#include <netinet/ip.h>
+#include <netinet/tcp.h>
+#include <netinet/udp.h>
 #include <msp/core/application.h>
 #include <msp/debug/profiler.h>
 #include <msp/graphics/display.h>
@@ -29,6 +33,18 @@ class Resolver;
 class NetVis: public Msp::RegisteredApplication<NetVis>
 {
 private:
+	struct CaptureContext
+	{
+		const pcap_pkthdr *cap_hdr;
+		Host *src_host;
+		Port *src_port;
+		Host *dst_host;
+		Port *dst_port;
+		unsigned size;
+
+		CaptureContext();
+	};
+
 	std::string iface;
 	pcap_t *pcap;
 	Resolver *resolver;
@@ -67,6 +83,11 @@ private:
 	void create_history_texture();
 
 	static void capture_handler(unsigned char *, const pcap_pkthdr *, const unsigned char *);
+	void handle_ethernet(CaptureContext &, const ethhdr *, unsigned);
+	void handle_ipv4(CaptureContext &, const iphdr *, unsigned);
+	void handle_tcp(CaptureContext &, const tcphdr *, unsigned);
+	void handle_udp(CaptureContext &, const udphdr *, unsigned);
+	void handle_packet(CaptureContext &);
 
 	void sighandler(int);
 };