1 files changed, 84 insertions, 34 deletions

diff --git a/extensions/libct_proto_tcp.c b/extensions/libct_proto_tcp.c
index 45ee29b..5fa3652 100644
--- a/extensions/libct_proto_tcp.c
+++ b/extensions/libct_proto_tcp.c
@@ -10,18 +10,19 @@
 #include <stdio.h>
 #include <getopt.h>
 #include <stdlib.h>
-#include <string.h>
 #include <netinet/in.h> /* For htons */
-#include <linux/netfilter_ipv4/ip_conntrack_tuple.h>
-#include <linux/netfilter_ipv4/ip_conntrack.h>
+#include <linux/netfilter_ipv4/ip_conntrack_netlink.h>
 #include "libct_proto.h"
+#include "libctnetlink.h"
 
 static struct option opts[] = {
 	{"orig-port-src", 1, 0, '1'},
 	{"orig-port-dst", 1, 0, '2'},
 	{"reply-port-src", 1, 0, '3'},
 	{"reply-port-dst", 1, 0, '4'},
-	{"state", 1, 0, '5'},
+	{"mask-port-src", 1, 0, '5'},
+	{"mask-port-dst", 1, 0, '6'},
+	{"state", 1, 0, '7'},
 	{0, 0, 0, 0}
 };
 
@@ -38,7 +39,13 @@ enum tcp_param_flags {
 	REPL_DPORT_BIT = 3,
 	REPL_DPORT = (1 << REPL_DPORT_BIT),
 
-	STATE_BIT = 4,
+	MASK_SPORT_BIT = 4,
+	MASK_SPORT = (1 << MASK_SPORT_BIT),
+
+	MASK_DPORT_BIT = 5,
+	MASK_DPORT = (1 << MASK_DPORT_BIT),
+
+	STATE_BIT = 6,
 	STATE = (1 << STATE_BIT)
 };
 
@@ -55,48 +62,63 @@ static const char *states[] = {
 	"LISTEN"
 };
 
-static void help()
+void help()
 {
 	fprintf(stdout, "--orig-port-src        original source port\n");
 	fprintf(stdout, "--orig-port-dst        original destination port\n");
 	fprintf(stdout, "--reply-port-src       reply source port\n");
 	fprintf(stdout, "--reply-port-dst       reply destination port\n");
+	fprintf(stdout, "--mask-port-src	mask source port\n");
+	fprintf(stdout, "--mask-port-dst	mask destination port\n");
 	fprintf(stdout, "--state                TCP state, fe. ESTABLISHED\n");
 }
 
-static int parse(char c, char *argv[], 
-	   struct ip_conntrack_tuple *orig,
-	   struct ip_conntrack_tuple *reply,
-	   union ip_conntrack_proto *proto,
-	   unsigned int *flags)
+int parse_options(char c, char *argv[], 
+		  struct ctnl_tuple *orig,
+		  struct ctnl_tuple *reply,
+		  struct ctnl_tuple *mask,
+		  union ctnl_protoinfo *proto,
+		  unsigned int *flags)
 {
 	switch(c) {
 		case '1':
 			if (optarg) {
-				orig->src.u.tcp.port = htons(atoi(optarg));
+				orig->l4src.tcp.port = htons(atoi(optarg));
 				*flags |= ORIG_SPORT;
 			}
 			break;
 		case '2':
 			if (optarg) {
-				orig->dst.u.tcp.port = htons(atoi(optarg));
+				orig->l4dst.tcp.port = htons(atoi(optarg));
 				*flags |= ORIG_DPORT;
 			}
 			break;
 		case '3':
 			if (optarg) {
-				reply->src.u.tcp.port = htons(atoi(optarg));
+				reply->l4src.tcp.port = htons(atoi(optarg));
 				*flags |= REPL_SPORT;
 			}
 			break;
 		case '4':
 			if (optarg) {
-				reply->dst.u.tcp.port = htons(atoi(optarg));
+				reply->l4dst.tcp.port = htons(atoi(optarg));
 				*flags |= REPL_DPORT;
 			}
 			break;
 		case '5':
 			if (optarg) {
+				mask->l4src.tcp.port = htons(atoi(optarg));
+				*flags |= MASK_SPORT;
+			}
+			break;
+		case '6':
+			if (optarg) {
+				mask->l4src.tcp.port = htons(atoi(optarg));
+				*flags |= MASK_DPORT;
+			}
+			break;
+		case '7':
+			if (optarg) {
 				int i;
 				for (i=0; i<10; i++) {
 					if (strcmp(optarg, states[i]) == 0) {
@@ -114,39 +136,67 @@ static int parse(char c, char *argv[],
 	return 1;
 }
 
-static int final_check(unsigned int flags)
+int final_check(unsigned int flags,
+		struct ctnl_tuple *orig,
+		struct ctnl_tuple *reply)
 {
-	if ((flags & ORIG_SPORT) && (flags & ORIG_DPORT))
+	if ((flags & (ORIG_SPORT|ORIG_DPORT)) 
+	    && !(flags & (REPL_SPORT|REPL_DPORT))) {
+		reply->l4src.tcp.port = orig->l4dst.tcp.port;
+		reply->l4dst.tcp.port = orig->l4src.tcp.port;
 		return 1;
-	else if ((flags & REPL_SPORT) && (flags & REPL_DPORT))
+	} else if (!(flags & (ORIG_SPORT|ORIG_DPORT))
+	            && (flags & (REPL_SPORT|REPL_DPORT))) {
+		orig->l4src.tcp.port = reply->l4dst.tcp.port;
+		orig->l4dst.tcp.port = reply->l4src.tcp.port;
+		return 1;
+	}
+	if ((flags & (ORIG_SPORT|ORIG_DPORT)) 
+	    && ((flags & (REPL_SPORT|REPL_DPORT))))
 		return 1;
 
 	return 0;
 }
 
-static void print_tuple(struct ip_conntrack_tuple *t)
+void parse_proto(struct nfattr *cda[], struct ctnl_tuple *tuple)
+{
+	if (cda[CTA_PROTO_TCP_SRC-1])
+		tuple->l4src.tcp.port =
+			*(u_int16_t *)NFA_DATA(cda[CTA_PROTO_TCP_SRC-1]);
+	if (cda[CTA_PROTO_TCP_DST-1])
+		tuple->l4dst.tcp.port =
+			*(u_int16_t *)NFA_DATA(cda[CTA_PROTO_TCP_DST-1]);
+}
+
+void parse_protoinfo(struct nfattr *cda[], struct ctnl_conntrack *ct)
+{
+	if (cda[CTA_PROTOINFO_TCP_STATE-1])
+                ct->protoinfo.tcp.state =
+                        *(u_int8_t *)NFA_DATA(cda[CTA_PROTOINFO_TCP_STATE-1]);
+}
+
+void print_protoinfo(union ctnl_protoinfo *protoinfo)
 {
-	fprintf(stdout, "sport=%d dport=%d ", ntohs(t->src.u.tcp.port), 
-				             ntohs(t->dst.u.tcp.port));
+	fprintf(stdout, "%s ", states[protoinfo->tcp.state]);
 }
 
-static void print_proto(union ip_conntrack_proto *proto)
+void print_proto(struct ctnl_tuple *tuple)
 {
-	if (proto->tcp.state > sizeof(states)/sizeof(char *))
-		fprintf(stdout, "[%u] ", states[proto->tcp.state]);
-	else
-		fprintf(stdout, "[%s] ", states[proto->tcp.state]);
+	fprintf(stdout, "sport=%u dport=%u ", htons(tuple->l4src.tcp.port),
+					      htons(tuple->l4dst.tcp.port));
 }
 
 static struct ctproto_handler tcp = {
-	.name 		= "tcp",
-	.protonum	= 6,
-	.parse		= parse,
-	.print_tuple	= print_tuple,
-	.print_proto	= print_proto,
-	.final_check	= final_check,
-	.help		= help,
-	.opts		= opts,
+	.name 			= "tcp",
+	.protonum		= 6,
+	.parse_opts		= parse_options,
+	.parse_protoinfo	= parse_protoinfo,
+	.parse_proto		= parse_proto,
+	.print_proto		= print_proto,
+	.print_protoinfo	= print_protoinfo,
+	.final_check		= final_check,
+	.help			= help,
+	.opts			= opts
 };
 
 void __attribute__ ((constructor)) init(void);