#include <linux/kernel.h>
 #include <linux/skbuff.h>
 #include <linux/rtnetlink.h>
+#include <net/inet_ecn.h>
 #include <net/netlink.h>
 #include <net/pkt_sched.h>
 #include <net/pkt_cls.h>
 static unsigned int skbmod_net_id;
 static struct tc_action_ops act_skbmod_ops;
 
-#define MAX_EDIT_LEN ETH_HLEN
 static int tcf_skbmod_act(struct sk_buff *skb, const struct tc_action *a,
                          struct tcf_result *res)
 {
        struct tcf_skbmod *d = to_skbmod(a);
-       int action;
+       int action, max_edit_len, err;
        struct tcf_skbmod_params *p;
        u64 flags;
-       int err;
 
        tcf_lastuse_update(&d->tcf_tm);
        bstats_cpu_update(this_cpu_ptr(d->common.cpu_bstats), skb);
        if (unlikely(action == TC_ACT_SHOT))
                goto drop;
 
-       if (!skb->dev || skb->dev->type != ARPHRD_ETHER)
-               return action;
+       max_edit_len = skb_mac_header_len(skb);
+       p = rcu_dereference_bh(d->skbmod_p);
+       flags = p->flags;
+
+       /* tcf_skbmod_init() guarantees "flags" to be one of the following:
+        *      1. a combination of SKBMOD_F_{DMAC,SMAC,ETYPE}
+        *      2. SKBMOD_F_SWAPMAC
+        *      3. SKBMOD_F_ECN
+        * SKBMOD_F_ECN only works with IP packets; all other flags only work with Ethernet
+        * packets.
+        */
+       if (flags == SKBMOD_F_ECN) {
+               switch (skb_protocol(skb, true)) {
+               case cpu_to_be16(ETH_P_IP):
+               case cpu_to_be16(ETH_P_IPV6):
+                       max_edit_len += skb_network_header_len(skb);
+                       break;
+               default:
+                       goto out;
+               }
+       } else if (!skb->dev || skb->dev->type != ARPHRD_ETHER) {
+               goto out;
+       }
 
-       /* XXX: if you are going to edit more fields beyond ethernet header
-        * (example when you add IP header replacement or vlan swap)
-        * then MAX_EDIT_LEN needs to change appropriately
-       */
-       err = skb_ensure_writable(skb, MAX_EDIT_LEN);
+       err = skb_ensure_writable(skb, max_edit_len);
        if (unlikely(err)) /* best policy is to drop on the floor */
                goto drop;
 
-       p = rcu_dereference_bh(d->skbmod_p);
-       flags = p->flags;
        if (flags & SKBMOD_F_DMAC)
                ether_addr_copy(eth_hdr(skb)->h_dest, p->eth_dst);
        if (flags & SKBMOD_F_SMAC)
                ether_addr_copy(eth_hdr(skb)->h_source, (u8 *)tmpaddr);
        }
 
+       if (flags & SKBMOD_F_ECN)
+               INET_ECN_set_ce(skb);
+
+out:
        return action;
 
 drop:
        index = parm->index;
        if (parm->flags & SKBMOD_F_SWAPMAC)
                lflags = SKBMOD_F_SWAPMAC;
+       if (parm->flags & SKBMOD_F_ECN)
+               lflags = SKBMOD_F_ECN;
 
        err = tcf_idr_check_alloc(tn, &index, a, bind);
        if (err < 0)