1 files changed, 45 insertions, 17 deletions

diff --git a/kernel/include/linux/netfilter/ip_set_getport.h b/kernel/include/linux/netfilter/ip_set_getport.h
index cf150d3..e4d469d 100644
--- a/kernel/include/linux/netfilter/ip_set_getport.h
+++ b/kernel/include/linux/netfilter/ip_set_getport.h
@@ -2,13 +2,14 @@
 #define _IP_SET_GETPORT_H
 
 #ifdef __KERNEL__
+#include <linux/icmp.h>
+#include <linux/icmpv6.h>
 #include <linux/netfilter_ipv6/ip6_tables.h>
 #include <net/ip.h>
 
 #define IPSET_INVALID_PORT	65536
 
 /* We must handle non-linear skbs */
-
 static inline bool
 get_port(const struct sk_buff *skb, int protocol, unsigned int protooff,
 	 bool src, u16 *port, u8 *proto)
@@ -38,13 +39,32 @@ get_port(const struct sk_buff *skb, int protocol, unsigned int protooff,
 	     	*port = src ? uh->source : uh->dest;
 	     	break;
 	}
-	default:
-		if (*proto == IPSET_IPPROTO_TCPUDP)
+	case IPPROTO_ICMP: {
+		struct icmphdr _icmph;
+		const struct icmphdr *ic;
+		
+		ic = skb_header_pointer(skb, protooff, sizeof(_icmph), &_icmph);
+		if (ic == NULL)
+			return false;
+		
+		*port = (ic->type << 8) & ic->code;
+		break;
+	}
+	case IPPROTO_ICMPV6: {
+		struct icmp6hdr _icmph;
+		const struct icmp6hdr *ic;
+		
+		ic = skb_header_pointer(skb, protooff, sizeof(_icmph), &_icmph);
+		if (ic == NULL)
 			return false;
+		
+		*port = (ic->icmp6_type << 8) & ic->icmp6_code;
+		break;
+	}
+	default:
 		break;
 	}
-	if (*proto != IPSET_IPPROTO_TCPUDP)
-		*proto = protocol;
+	*proto = protocol;
 
 	return true;
 }
@@ -56,9 +76,6 @@ get_ip4_port(const struct sk_buff *skb, bool src, u16 *port, u8 *proto)
 	unsigned int protooff = ip_hdrlen(skb);
 	int protocol = iph->protocol;
 
-	if (!(*proto >= IPSET_IPPROTO_TCPUDP || *proto == protocol))
-		return false;
-
 	/* See comments at tcp_match in ip_tables.c */
 	if (ntohs(iph->frag_off) & IP_OFFSET)
 		return false;
@@ -77,21 +94,32 @@ get_ip6_port(const struct sk_buff *skb, bool src, u16 *port, u8 *proto)
 	if (protocol < 0 || fragoff)
 		return false;
 
-	if (!(*proto >= IPSET_IPPROTO_TCPUDP || *proto == protocol))
-		return false;
-
 	return get_port(skb, protocol, protooff, src, port, proto);
 }
 
 static inline bool
 get_ip_port(const struct sk_buff *skb, u8 pf, bool src, u16 *port)
 {
-	u8 proto = IPSET_IPPROTO_TCPUDP;
-
-	if (pf == AF_INET)
-		return get_ip4_port(skb, src, port, &proto);
-	else
-		return get_ip6_port(skb, src, port, &proto);
+	bool ret;
+	u8 proto;
+
+	switch (pf) {
+	case AF_INET:
+		ret = get_ip4_port(skb, src, port, &proto);
+	case AF_INET6:
+		ret = get_ip6_port(skb, src, port, &proto);
+	default:
+		return false;
+	}
+	if (!ret)
+		return ret;
+	switch (proto) {
+	case IPPROTO_TCP:
+	case IPPROTO_UDP:
+		return true;
+	default:
+		return false;
+	}
 }
 #endif /* __KERNEL__ */