net_sched: cls_flow: use proto_ports_offset() to support AH message

Signed-off-by: Changli Gao <xiaosuo@gmail.com>
Signed-off-by: David S. Miller <davem@davemloft.net>

diff --git a/net/sched/cls_flow.c b/net/sched/cls_flow.c
index e17096e3913ccb983cf8abd533da20d30cf37b62..cd709f1294df1bd536943555885ef6a147efcac4 100644 (file)
--- a/net/sched/cls_flow.c
+++ b/net/sched/cls_flow.c
@@ -111,44 +111,41 @@ static u32 flow_get_proto(struct sk_buff *skb)
        }
 }
 
-static int has_ports(u8 protocol)
-{
-       switch (protocol) {
-       case IPPROTO_TCP:
-       case IPPROTO_UDP:
-       case IPPROTO_UDPLITE:
-       case IPPROTO_SCTP:
-       case IPPROTO_DCCP:
-       case IPPROTO_ESP:
-               return 1;
-       default:
-               return 0;
-       }
-}
-
 static u32 flow_get_proto_src(struct sk_buff *skb)
 {
        switch (skb->protocol) {
        case htons(ETH_P_IP): {
                struct iphdr *iph;
+               int poff;
 
                if (!pskb_network_may_pull(skb, sizeof(*iph)))
                        break;
                iph = ip_hdr(skb);
-               if (!(iph->frag_off&htons(IP_MF|IP_OFFSET)) &&
-                   has_ports(iph->protocol) &&
-                   pskb_network_may_pull(skb, iph->ihl * 4 + 2))
-                       return ntohs(*(__be16 *)((void *)iph + iph->ihl * 4));
+               if (iph->frag_off & htons(IP_MF|IP_OFFSET))
+                       break;
+               poff = proto_ports_offset(iph->protocol);
+               if (poff >= 0 &&
+                   pskb_network_may_pull(skb, iph->ihl * 4 + 2 + poff)) {
+                       iph = ip_hdr(skb);
+                       return ntohs(*(__be16 *)((void *)iph + iph->ihl * 4 +
+                                                poff));
+               }
                break;
        }
        case htons(ETH_P_IPV6): {
                struct ipv6hdr *iph;
+               int poff;
 
-               if (!pskb_network_may_pull(skb, sizeof(*iph) + 2))
+               if (!pskb_network_may_pull(skb, sizeof(*iph)))
                        break;
                iph = ipv6_hdr(skb);
-               if (has_ports(iph->nexthdr))
-                       return ntohs(*(__be16 *)&iph[1]);
+               poff = proto_ports_offset(iph->nexthdr);
+               if (poff >= 0 &&
+                   pskb_network_may_pull(skb, sizeof(*iph) + poff + 2)) {
+                       iph = ipv6_hdr(skb);
+                       return ntohs(*(__be16 *)((void *)iph + sizeof(*iph) +
+                                                poff));
+               }
                break;
        }
        }
@@ -161,24 +158,36 @@ static u32 flow_get_proto_dst(struct sk_buff *skb)
        switch (skb->protocol) {
        case htons(ETH_P_IP): {
                struct iphdr *iph;
+               int poff;
 
                if (!pskb_network_may_pull(skb, sizeof(*iph)))
                        break;
                iph = ip_hdr(skb);
-               if (!(iph->frag_off&htons(IP_MF|IP_OFFSET)) &&
-                   has_ports(iph->protocol) &&
-                   pskb_network_may_pull(skb, iph->ihl * 4 + 4))
-                       return ntohs(*(__be16 *)((void *)iph + iph->ihl * 4 + 2));
+               if (iph->frag_off & htons(IP_MF|IP_OFFSET))
+                       break;
+               poff = proto_ports_offset(iph->protocol);
+               if (poff >= 0 &&
+                   pskb_network_may_pull(skb, iph->ihl * 4 + 4 + poff)) {
+                       iph = ip_hdr(skb);
+                       return ntohs(*(__be16 *)((void *)iph + iph->ihl * 4 +
+                                                2 + poff));
+               }
                break;
        }
        case htons(ETH_P_IPV6): {
                struct ipv6hdr *iph;
+               int poff;
 
-               if (!pskb_network_may_pull(skb, sizeof(*iph) + 4))
+               if (!pskb_network_may_pull(skb, sizeof(*iph)))
                        break;
                iph = ipv6_hdr(skb);
-               if (has_ports(iph->nexthdr))
-                       return ntohs(*(__be16 *)((void *)&iph[1] + 2));
+               poff = proto_ports_offset(iph->nexthdr);
+               if (poff >= 0 &&
+                   pskb_network_may_pull(skb, sizeof(*iph) + poff + 4)) {
+                       iph = ipv6_hdr(skb);
+                       return ntohs(*(__be16 *)((void *)iph + sizeof(*iph) +
+                                                poff + 2));
+               }
                break;
        }
        }