Merge "Add IPv6 Handling for DSCP Policies and Support Interfaces with MAC Addresses" am: ba18a9ee43

Original change: https://android-review.googlesource.com/c/platform/packages/modules/Connectivity/+/1972139

Change-Id: I9abc9a43fb3eca26aa84ce7bdac256ab5715f7d6
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/bpf_progs/dscp_policy.c b/bpf_progs/dscp_policy.c
index 9989e6b..d5df7ef 100644
--- a/bpf_progs/dscp_policy.c
+++ b/bpf_progs/dscp_policy.c
@@ -16,6 +16,7 @@
 
 #include <linux/types.h>
 #include <linux/bpf.h>
+#include <linux/if_packet.h>
 #include <linux/ip.h>
 #include <linux/ipv6.h>
 #include <linux/if_ether.h>
@@ -27,249 +28,294 @@
 #include <string.h>
 
 #include "bpf_helpers.h"
+#include "dscp_policy.h"
 
-#define MAX_POLICIES 16
-#define MAP_A 1
-#define MAP_B 2
-
-#define STRUCT_SIZE(name, size) _Static_assert(sizeof(name) == (size), "Incorrect struct size.")
-
-// TODO: these are already defined in /system/netd/bpf_progs/bpf_net_helpers.h
-// should they be moved to common location?
-static uint64_t (*bpf_get_socket_cookie)(struct __sk_buff* skb) =
-        (void*)BPF_FUNC_get_socket_cookie;
-static int (*bpf_skb_store_bytes)(struct __sk_buff* skb, __u32 offset, const void* from, __u32 len,
-                                  __u64 flags) = (void*)BPF_FUNC_skb_store_bytes;
-static int (*bpf_l3_csum_replace)(struct __sk_buff* skb, __u32 offset, __u64 from, __u64 to,
-                                  __u64 flags) = (void*)BPF_FUNC_l3_csum_replace;
-
-typedef struct {
-    // Add family here to match __sk_buff ?
-    struct in_addr srcIp;
-    struct in_addr dstIp;
-    __be16 srcPort;
-    __be16 dstPort;
-    uint8_t proto;
-    uint8_t dscpVal;
-    uint8_t pad[2];
-} Ipv4RuleEntry;
-STRUCT_SIZE(Ipv4RuleEntry, 2 * 4 + 2 * 2 + 2 * 1 + 2);  // 16, 4 for in_addr
-
-#define SRC_IP_MASK     1
-#define DST_IP_MASK     2
-#define SRC_PORT_MASK   4
-#define DST_PORT_MASK   8
-#define PROTO_MASK      16
-
-typedef struct {
-    struct in6_addr srcIp;
-    struct in6_addr dstIp;
-    __be16 srcPort;
-    __be16 dstPortStart;
-    __be16 dstPortEnd;
-    uint8_t proto;
-    uint8_t dscpVal;
-    uint8_t mask;
-    uint8_t pad[3];
-} Ipv4Policy;
-STRUCT_SIZE(Ipv4Policy, 2 * 16 + 3 * 2 + 3 * 1 + 3);  // 44
-
-typedef struct {
-    struct in6_addr srcIp;
-    struct in6_addr dstIp;
-    __be16 srcPort;
-    __be16 dstPortStart;
-    __be16 dstPortEnd;
-    uint8_t proto;
-    uint8_t dscpVal;
-    uint8_t mask;
-    // should we override this struct to include the param bitmask for linear search?
-    // For mapping socket to policies, all the params should match exactly since we can
-    // pull any missing from the sock itself.
-} Ipv6RuleEntry;
-STRUCT_SIZE(Ipv6RuleEntry, 2 * 16 + 3 * 2 + 3 * 1 + 3);  // 44
-
-// TODO: move to using 1 map. Map v4 address to 0xffff::v4
-DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_A, HASH, uint64_t, Ipv4RuleEntry, MAX_POLICIES,
-        AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_B, HASH, uint64_t, Ipv4RuleEntry, MAX_POLICIES,
-        AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_A, HASH, uint64_t, Ipv6RuleEntry, MAX_POLICIES,
-        AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_B, HASH, uint64_t, Ipv6RuleEntry, MAX_POLICIES,
-        AID_SYSTEM)
 DEFINE_BPF_MAP_GRW(switch_comp_map, ARRAY, int, uint64_t, 1, AID_SYSTEM)
 
-DEFINE_BPF_MAP_GRW(ipv4_dscp_policies_map, ARRAY, uint32_t, Ipv4Policy, MAX_POLICIES,
+DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_A, HASH, uint64_t, RuleEntry, MAX_POLICIES,
         AID_SYSTEM)
-DEFINE_BPF_MAP_GRW(ipv6_dscp_policies_map, ARRAY, uint32_t, Ipv6RuleEntry, MAX_POLICIES,
+DEFINE_BPF_MAP_GRW(ipv4_socket_to_policies_map_B, HASH, uint64_t, RuleEntry, MAX_POLICIES,
+        AID_SYSTEM)
+DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_A, HASH, uint64_t, RuleEntry, MAX_POLICIES,
+        AID_SYSTEM)
+DEFINE_BPF_MAP_GRW(ipv6_socket_to_policies_map_B, HASH, uint64_t, RuleEntry, MAX_POLICIES,
         AID_SYSTEM)
 
-DEFINE_BPF_PROG_KVER("schedcls/set_dscp", AID_ROOT, AID_SYSTEM,
-                     schedcls_set_dscp, KVER(5, 4, 0))
-(struct __sk_buff* skb) {
-    int one = 0;
-    uint64_t* selectedMap = bpf_switch_comp_map_lookup_elem(&one);
+DEFINE_BPF_MAP_GRW(ipv4_dscp_policies_map, ARRAY, uint32_t, DscpPolicy, MAX_POLICIES,
+        AID_SYSTEM)
+DEFINE_BPF_MAP_GRW(ipv6_dscp_policies_map, ARRAY, uint32_t, DscpPolicy, MAX_POLICIES,
+        AID_SYSTEM)
+
+static inline __always_inline void match_policy(struct __sk_buff* skb, bool ipv4, bool is_eth) {
+    void* data = (void*)(long)skb->data;
+    const void* data_end = (void*)(long)skb->data_end;
+
+    const int l2_header_size = is_eth ? sizeof(struct ethhdr) : 0;
+    struct ethhdr* eth = is_eth ? data : NULL;
+
+    if (data + l2_header_size > data_end) return;
+
+    int zero = 0;
+    int hdr_size = 0;
+    uint64_t* selectedMap = bpf_switch_comp_map_lookup_elem(&zero);
 
     // use this with HASH map so map lookup only happens once policies have been added?
     if (!selectedMap) {
-        return TC_ACT_PIPE;
+        return;
     }
 
     // used for map lookup
     uint64_t cookie = bpf_get_socket_cookie(skb);
+    if (!cookie)
+        return;
 
-    // Do we need separate maps for ipv4/ipv6
-    if (skb->protocol == htons(ETH_P_IP)) { //maybe bpf_htons()
-        Ipv4RuleEntry* v4Policy;
-        if (*selectedMap == MAP_A) {
-            v4Policy = bpf_ipv4_socket_to_policies_map_A_lookup_elem(&cookie);
-        } else {
-            v4Policy = bpf_ipv4_socket_to_policies_map_B_lookup_elem(&cookie);
-        }
-
-        // How to use bitmask here to compare params efficiently?
-        // TODO: add BPF_PROG_TYPE_SK_SKB prog type to Loader?
-
-        void* data = (void*)(long)skb->data;
-        const void* data_end = (void*)(long)skb->data_end;
-        const struct iphdr* const iph = data;
-
+    uint16_t sport = 0;
+    uint16_t dport = 0;
+    uint8_t protocol = 0; // TODO: Use are reserved value? Or int (-1) and cast to uint below?
+    struct in6_addr srcIp = {};
+    struct in6_addr dstIp = {};
+    uint8_t tos = 0; // Only used for IPv4
+    uint8_t priority = 0; // Only used for IPv6
+    uint8_t flow_lbl = 0; // Only used for IPv6
+    if (ipv4) {
+        const struct iphdr* const iph = is_eth ? (void*)(eth + 1) : data;
         // Must have ipv4 header
-        if (data + sizeof(*iph) > data_end) return TC_ACT_PIPE;
+        if (data + l2_header_size + sizeof(*iph) > data_end) return;
 
         // IP version must be 4
-        if (iph->version != 4) return TC_ACT_PIPE;
+        if (iph->version != 4) return;
 
         // We cannot handle IP options, just standard 20 byte == 5 dword minimal IPv4 header
-        if (iph->ihl != 5) return TC_ACT_PIPE;
+        if (iph->ihl != 5) return;
 
-        if (iph->protocol != IPPROTO_UDP) return TC_ACT_PIPE;
+        // V4 mapped address in in6_addr sets 10/11 position to 0xff.
+        srcIp.s6_addr32[2] = htonl(0x0000ffff);
+        dstIp.s6_addr32[2] = htonl(0x0000ffff);
 
-        struct udphdr *udp;
-        udp = data + sizeof(struct iphdr); //sizeof(struct ethhdr)
+        // Copy IPv4 address into in6_addr for easy comparison below.
+        srcIp.s6_addr32[3] = iph->saddr;
+        dstIp.s6_addr32[3] = iph->daddr;
+        protocol = iph->protocol;
+        tos = iph->tos;
+        hdr_size = sizeof(struct iphdr);
+    } else {
+        struct ipv6hdr* ip6h = is_eth ? (void*)(eth + 1) : data;
+        // Must have ipv6 header
+        if (data + l2_header_size + sizeof(*ip6h) > data_end) return;
 
-        if ((void*)(udp + 1) > data_end) return TC_ACT_PIPE;
+        if (ip6h->version != 6) return;
 
-        // Source/destination port in udphdr are stored in be16, need to convert to le16.
-        // This can be done via ntohs or htons. Is there a more preferred way?
-        // Cached policy was found.
-        if (v4Policy && iph->saddr == v4Policy->srcIp.s_addr &&
-                    iph->daddr == v4Policy->dstIp.s_addr &&
-                    ntohs(udp->source) == v4Policy->srcPort &&
-                    ntohs(udp->dest) == v4Policy->dstPort &&
-                    iph->protocol == v4Policy->proto) {
-            // set dscpVal in packet. Least sig 2 bits of TOS
-            // reference ipv4_change_dsfield()
+        srcIp = ip6h->saddr;
+        dstIp = ip6h->daddr;
+        protocol = ip6h->nexthdr;
+        priority = ip6h->priority;
+        flow_lbl = ip6h->flow_lbl[0];
+        hdr_size = sizeof(struct ipv6hdr);
+    }
 
-            // TODO: fix checksum...
-            int ecn = iph->tos & 3;
-            uint8_t newDscpVal = (v4Policy->dscpVal << 2) + ecn;
-            int oldDscpVal = iph->tos >> 2;
+    switch (protocol) {
+        case IPPROTO_UDP:
+        case IPPROTO_UDPLITE:
+        {
+            struct udphdr *udp;
+            udp = data + hdr_size;
+            if ((void*)(udp + 1) > data_end) return;
+            sport = udp->source;
+            dport = udp->dest;
+        }
+        break;
+        case IPPROTO_TCP:
+        {
+            struct tcphdr *tcp;
+            tcp = data + hdr_size;
+            if ((void*)(tcp + 1) > data_end) return;
+            sport = tcp->source;
+            dport = tcp->dest;
+        }
+        break;
+        default:
+            return;
+    }
+
+    RuleEntry* existingRule;
+    if (ipv4) {
+        if (*selectedMap == MAP_A) {
+            existingRule = bpf_ipv4_socket_to_policies_map_A_lookup_elem(&cookie);
+        } else {
+            existingRule = bpf_ipv4_socket_to_policies_map_B_lookup_elem(&cookie);
+        }
+    } else {
+        if (*selectedMap == MAP_A) {
+            existingRule = bpf_ipv6_socket_to_policies_map_A_lookup_elem(&cookie);
+        } else {
+            existingRule = bpf_ipv6_socket_to_policies_map_B_lookup_elem(&cookie);
+        }
+    }
+
+    if (existingRule && v6_equal(srcIp, existingRule->srcIp) &&
+                v6_equal(dstIp, existingRule->dstIp) &&
+                skb->ifindex == existingRule->ifindex &&
+                ntohs(sport) == htons(existingRule->srcPort) &&
+                ntohs(dport) == htons(existingRule->dstPort) &&
+                protocol == existingRule->proto) {
+        if (ipv4) {
+            int ecn = tos & 3;
+            uint8_t newDscpVal = (existingRule->dscpVal << 2) + ecn;
+            int oldDscpVal = tos >> 2;
             bpf_l3_csum_replace(skb, 1, oldDscpVal, newDscpVal, sizeof(uint8_t));
             bpf_skb_store_bytes(skb, 1, &newDscpVal, sizeof(uint8_t), 0);
-            return TC_ACT_PIPE;
+        } else {
+            uint8_t new_priority = (existingRule->dscpVal >> 2) + 0x60;
+            uint8_t new_flow_label = ((existingRule->dscpVal & 0xf) << 6) + (priority >> 6);
+            bpf_skb_store_bytes(skb, 0, &new_priority, sizeof(uint8_t), 0);
+            bpf_skb_store_bytes(skb, 1, &new_flow_label, sizeof(uint8_t), 0);
+        }
+        return;
+    }
+
+    // Linear scan ipv4_dscp_policies_map since no stored params match skb.
+    int bestScore = -1;
+    uint32_t bestMatch = 0;
+
+    for (register uint64_t i = 0; i < MAX_POLICIES; i++) {
+        int score = 0;
+        uint8_t tempMask = 0;
+        // Using a uint64 in for loop prevents infinite loop during BPF load,
+        // but the key is uint32, so convert back.
+        uint32_t key = i;
+
+        DscpPolicy* policy;
+        if (ipv4) {
+            policy = bpf_ipv4_dscp_policies_map_lookup_elem(&key);
+        } else {
+            policy = bpf_ipv6_dscp_policies_map_lookup_elem(&key);
         }
 
-        // linear scan ipv4_dscp_policies_map, stored socket params do not match actual
-        int bestScore = -1;
-        uint32_t bestMatch = 0;
+        // If the policy lookup failed, presentFields is 0, or iface index does not match
+        // index on skb buff, then we can continue to next policy.
+        if (!policy || policy->presentFields == 0 || policy->ifindex != skb->ifindex)
+            continue;
 
-        for (register uint64_t i = 0; i < MAX_POLICIES; i++) {
-            int score = 0;
-            uint8_t tempMask = 0;
-            // Using a uint62 in for loop prevents infinite loop during BPF load,
-            // but the key is uint32, so convert back.
-            uint32_t key = i;
-            Ipv4Policy* policy = bpf_ipv4_dscp_policies_map_lookup_elem(&key);
+        if ((policy->presentFields & SRC_IP_MASK_FLAG) == SRC_IP_MASK_FLAG &&
+                v6_equal(srcIp, policy->srcIp)) {
+            score++;
+            tempMask |= SRC_IP_MASK_FLAG;
+        }
+        if ((policy->presentFields & DST_IP_MASK_FLAG) == DST_IP_MASK_FLAG &&
+                v6_equal(dstIp, policy->dstIp)) {
+            score++;
+            tempMask |= DST_IP_MASK_FLAG;
+        }
+        if ((policy->presentFields & SRC_PORT_MASK_FLAG) == SRC_PORT_MASK_FLAG &&
+                ntohs(sport) == htons(policy->srcPort)) {
+            score++;
+            tempMask |= SRC_PORT_MASK_FLAG;
+        }
+        if ((policy->presentFields & DST_PORT_MASK_FLAG) == DST_PORT_MASK_FLAG &&
+                ntohs(dport) >= htons(policy->dstPortStart) &&
+                ntohs(dport) <= htons(policy->dstPortEnd)) {
+            score++;
+            tempMask |= DST_PORT_MASK_FLAG;
+        }
+        if ((policy->presentFields & PROTO_MASK_FLAG) == PROTO_MASK_FLAG &&
+                protocol == policy->proto) {
+            score++;
+            tempMask |= PROTO_MASK_FLAG;
+        }
 
-            // if mask is 0 continue, key does not have corresponding policy value
-            if (policy && policy->mask != 0) {
-                if ((policy->mask & SRC_IP_MASK) == SRC_IP_MASK &&
-                        iph->saddr == policy->srcIp.s6_addr32[3]) {
-                    score++;
-                    tempMask |= SRC_IP_MASK;
-                }
-                if ((policy->mask & DST_IP_MASK) == DST_IP_MASK &&
-                        iph->daddr == policy->dstIp.s6_addr32[3]) {
-                    score++;
-                    tempMask |= DST_IP_MASK;
-                }
-                if ((policy->mask & SRC_PORT_MASK) == SRC_PORT_MASK &&
-                        ntohs(udp->source) == htons(policy->srcPort)) {
-                    score++;
-                    tempMask |= SRC_PORT_MASK;
-                }
-                if ((policy->mask & DST_PORT_MASK) == DST_PORT_MASK &&
-                        ntohs(udp->dest) >= htons(policy->dstPortStart) &&
-                        ntohs(udp->dest) <= htons(policy->dstPortEnd)) {
-                    score++;
-                    tempMask |= DST_PORT_MASK;
-                }
-                if ((policy->mask & PROTO_MASK) == PROTO_MASK &&
-                        iph->protocol == policy->proto) {
-                    score++;
-                    tempMask |= PROTO_MASK;
-                }
+        if (score > bestScore && tempMask == policy->presentFields) {
+            bestMatch = i;
+            bestScore = score;
+        }
+    }
 
-                if (score > bestScore && tempMask == policy->mask) {
-                    bestMatch = i;
-                    bestScore = score;
-                }
+    uint8_t new_tos= 0; // Can 0 be used as default forwarding value?
+    uint8_t new_priority = 0;
+    uint8_t new_flow_lbl = 0;
+    if (bestScore > 0) {
+        DscpPolicy* policy;
+        if (ipv4) {
+            policy = bpf_ipv4_dscp_policies_map_lookup_elem(&bestMatch);
+        } else {
+            policy = bpf_ipv6_dscp_policies_map_lookup_elem(&bestMatch);
+        }
+
+        if (policy) {
+            // TODO: if DSCP value is already set ignore?
+            if (ipv4) {
+                int ecn = tos & 3;
+                new_tos = (policy->dscpVal << 2) + ecn;
+            } else {
+                new_priority = (policy->dscpVal >> 2) + 0x60;
+                new_flow_lbl = ((policy->dscpVal & 0xf) << 6) + (flow_lbl >> 6);
+
+                // Set IPv6 curDscp value to stored value and recalulate priority
+                // and flow label during next use.
+                new_tos = policy->dscpVal;
             }
         }
+    } else return;
 
-        uint8_t newDscpVal = 0; // Can 0 be used as default forwarding value?
-        uint8_t curDscp = iph->tos & 252;
-        if (bestScore > 0) {
-            Ipv4Policy* policy = bpf_ipv4_dscp_policies_map_lookup_elem(&bestMatch);
-            if (policy) {
-                // TODO: if DSCP value is already set ignore?
-                // TODO: update checksum, for testing increment counter...
-                int ecn = iph->tos & 3;
-                newDscpVal = (policy->dscpVal << 2) + ecn;
-            }
-        }
+    RuleEntry value = {
+        .srcIp = srcIp,
+        .dstIp = dstIp,
+        .ifindex = skb->ifindex,
+        .srcPort = sport,
+        .dstPort = dport,
+        .proto = protocol,
+        .dscpVal = new_tos,
+    };
 
-        Ipv4RuleEntry value = {
-            .srcIp.s_addr = iph->saddr,
-            .dstIp.s_addr = iph->daddr,
-            .srcPort = udp->source,
-            .dstPort = udp->dest,
-            .proto = iph->protocol,
-            .dscpVal = newDscpVal,
-        };
-
-        if (!cookie)
-            return TC_ACT_PIPE;
-
-        // Update map
+    //Update map with new policy.
+    if (ipv4) {
         if (*selectedMap == MAP_A) {
             bpf_ipv4_socket_to_policies_map_A_update_elem(&cookie, &value, BPF_ANY);
         } else {
             bpf_ipv4_socket_to_policies_map_B_update_elem(&cookie, &value, BPF_ANY);
         }
-
-        // Need to store bytes after updating map or program will not load.
-        if (newDscpVal != curDscp) {
-            // 1 is the offset (Version/Header length)
-            int oldDscpVal = iph->tos >> 2;
-            bpf_l3_csum_replace(skb, 1, oldDscpVal, newDscpVal, sizeof(uint8_t));
-            bpf_skb_store_bytes(skb, 1, &newDscpVal, sizeof(uint8_t), 0);
-        }
-
-    } else if (skb->protocol == htons(ETH_P_IPV6)) { //maybe bpf_htons()
-        Ipv6RuleEntry* v6Policy;
+    } else {
         if (*selectedMap == MAP_A) {
-            v6Policy = bpf_ipv6_socket_to_policies_map_A_lookup_elem(&cookie);
+            bpf_ipv6_socket_to_policies_map_A_update_elem(&cookie, &value, BPF_ANY);
         } else {
-            v6Policy = bpf_ipv6_socket_to_policies_map_B_lookup_elem(&cookie);
+            bpf_ipv6_socket_to_policies_map_B_update_elem(&cookie, &value, BPF_ANY);
         }
+    }
 
-        if (!v6Policy)
-            return TC_ACT_PIPE;
+    // Need to store bytes after updating map or program will not load.
+    if (ipv4 && new_tos != (tos & 252)) {
+        int oldDscpVal = tos >> 2;
+        bpf_l3_csum_replace(skb, 1, oldDscpVal, new_tos, sizeof(uint8_t));
+        bpf_skb_store_bytes(skb, 1, &new_tos, sizeof(uint8_t), 0);
+    } else if (!ipv4 && (new_priority != priority || new_flow_lbl != flow_lbl)) {
+        bpf_skb_store_bytes(skb, 0, &new_priority, sizeof(uint8_t), 0);
+        bpf_skb_store_bytes(skb, 1, &new_flow_lbl, sizeof(uint8_t), 0);
+    }
+    return;
+}
 
-        // TODO: Add code to process IPv6 packet.
+DEFINE_BPF_PROG_KVER("schedcls/set_dscp_ether", AID_ROOT, AID_SYSTEM,
+                     schedcls_set_dscp_ether, KVER(5, 4, 0))
+(struct __sk_buff* skb) {
+
+    if (skb->pkt_type != PACKET_HOST) return TC_ACT_PIPE;
+
+    if (skb->protocol == htons(ETH_P_IP)) {
+        match_policy(skb, true, true);
+    } else if (skb->protocol == htons(ETH_P_IPV6)) {
+        match_policy(skb, false, true);
+    }
+
+    // Always return TC_ACT_PIPE
+    return TC_ACT_PIPE;
+}
+
+DEFINE_BPF_PROG_KVER("schedcls/set_dscp_raw_ip", AID_ROOT, AID_SYSTEM,
+                     schedcls_set_dscp_raw_ip, KVER(5, 4, 0))
+(struct __sk_buff* skb) {
+    if (skb->protocol == htons(ETH_P_IP)) {
+        match_policy(skb, true, false);
+    } else if (skb->protocol == htons(ETH_P_IPV6)) {
+        match_policy(skb, false, false);
     }
 
     // Always return TC_ACT_PIPE
diff --git a/bpf_progs/dscp_policy.h b/bpf_progs/dscp_policy.h
new file mode 100644
index 0000000..777c4ff
--- /dev/null
+++ b/bpf_progs/dscp_policy.h
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#define MAX_POLICIES 16
+#define MAP_A 1
+#define MAP_B 2
+
+#define SRC_IP_MASK_FLAG     1
+#define DST_IP_MASK_FLAG     2
+#define SRC_PORT_MASK_FLAG   4
+#define DST_PORT_MASK_FLAG   8
+#define PROTO_MASK_FLAG      16
+
+#define STRUCT_SIZE(name, size) _Static_assert(sizeof(name) == (size), "Incorrect struct size.")
+
+#ifndef v6_equal
+#define v6_equal(a, b)    (a.s6_addr32[0] == b.s6_addr32[0] && \
+                 a.s6_addr32[1] == b.s6_addr32[1] && \
+                 a.s6_addr32[2] == b.s6_addr32[2] && \
+                 a.s6_addr32[3] == b.s6_addr32[3])
+#endif
+
+// TODO: these are already defined in packages/modules/Connectivity/bpf_progs/bpf_net_helpers.h.
+// smove to common location in future.
+static uint64_t (*bpf_get_socket_cookie)(struct __sk_buff* skb) =
+        (void*)BPF_FUNC_get_socket_cookie;
+static int (*bpf_skb_store_bytes)(struct __sk_buff* skb, __u32 offset, const void* from, __u32 len,
+                                  __u64 flags) = (void*)BPF_FUNC_skb_store_bytes;
+static int (*bpf_l3_csum_replace)(struct __sk_buff* skb, __u32 offset, __u64 from, __u64 to,
+                                  __u64 flags) = (void*)BPF_FUNC_l3_csum_replace;
+static long (*bpf_skb_ecn_set_ce)(struct __sk_buff* skb) =
+        (void*)BPF_FUNC_skb_ecn_set_ce;
+
+typedef struct {
+    struct in6_addr srcIp;
+    struct in6_addr dstIp;
+    uint32_t ifindex;
+    __be16 srcPort;
+    __be16 dstPortStart;
+    __be16 dstPortEnd;
+    uint8_t proto;
+    uint8_t dscpVal;
+    uint8_t presentFields;
+    uint8_t pad[3];
+} DscpPolicy;
+STRUCT_SIZE(DscpPolicy, 2 * 16 + 4 + 3 * 2 + 3 * 1 + 3);  // 48
+
+typedef struct {
+    struct in6_addr srcIp;
+    struct in6_addr dstIp;
+    __u32 ifindex;
+    __be16 srcPort;
+    __be16 dstPort;
+    __u8 proto;
+    __u8 dscpVal;
+    __u8 pad[2];
+} RuleEntry;
+STRUCT_SIZE(RuleEntry, 2 * 16 + 1 * 4 + 2 * 2 + 2 * 1 + 2);  // 44
\ No newline at end of file
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index d79bdb8..a0e75ec 100644
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -3669,7 +3669,7 @@
                 }
                 case NetworkAgent.EVENT_REMOVE_ALL_DSCP_POLICIES: {
                     if (mDscpPolicyTracker != null) {
-                        mDscpPolicyTracker.removeAllDscpPolicies(nai);
+                        mDscpPolicyTracker.removeAllDscpPolicies(nai, true);
                     }
                     break;
                 }
@@ -4410,6 +4410,9 @@
     }
 
     private void destroyNativeNetwork(@NonNull NetworkAgentInfo nai) {
+        if (mDscpPolicyTracker != null) {
+            mDscpPolicyTracker.removeAllDscpPolicies(nai, false);
+        }
         try {
             mNetd.networkDestroy(nai.network.getNetId());
         } catch (RemoteException | ServiceSpecificException e) {
diff --git a/service/src/com/android/server/connectivity/DscpPolicyTracker.java b/service/src/com/android/server/connectivity/DscpPolicyTracker.java
index 53b276e..de9dfe3 100644
--- a/service/src/com/android/server/connectivity/DscpPolicyTracker.java
+++ b/service/src/com/android/server/connectivity/DscpPolicyTracker.java
@@ -19,6 +19,7 @@
 import static android.net.NetworkAgent.DSCP_POLICY_STATUS_DELETED;
 import static android.net.NetworkAgent.DSCP_POLICY_STATUS_INSUFFICIENT_PROCESSING_RESOURCES;
 import static android.net.NetworkAgent.DSCP_POLICY_STATUS_POLICY_NOT_FOUND;
+import static android.net.NetworkAgent.DSCP_POLICY_STATUS_REQUEST_DECLINED;
 import static android.net.NetworkAgent.DSCP_POLICY_STATUS_SUCCESS;
 import static android.system.OsConstants.ETH_P_ALL;
 
@@ -37,6 +38,7 @@
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.NetworkInterface;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -66,30 +68,68 @@
 
     private final BpfMap<Struct.U32, DscpPolicyValue> mBpfDscpIpv4Policies;
     private final BpfMap<Struct.U32, DscpPolicyValue> mBpfDscpIpv6Policies;
-    private final SparseIntArray mPolicyIdToBpfMapIndex;
+
+    // The actual policy rules used by the BPF code to process packets
+    // are in mBpfDscpIpv4Policies and mBpfDscpIpv4Policies. Both of
+    // these can contain up to MAX_POLICIES rules.
+    //
+    // A given policy always consumes one entry in both the IPv4 and
+    // IPv6 maps even if if's an IPv4-only or IPv6-only policy.
+    //
+    // Each interface index has a SparseIntArray of rules which maps a
+    // policy ID to the index of the corresponding rule in the maps.
+    // mIfaceIndexToPolicyIdBpfMapIndex maps the interface index to
+    // the per-interface SparseIntArray.
+    private final HashMap<Integer, SparseIntArray> mIfaceIndexToPolicyIdBpfMapIndex;
 
     public DscpPolicyTracker() throws ErrnoException {
         mAttachedIfaces = new HashSet<String>();
-
-        mPolicyIdToBpfMapIndex = new SparseIntArray(MAX_POLICIES);
+        mIfaceIndexToPolicyIdBpfMapIndex = new HashMap<Integer, SparseIntArray>();
         mBpfDscpIpv4Policies = new BpfMap<Struct.U32, DscpPolicyValue>(IPV4_POLICY_MAP_PATH,
                 BpfMap.BPF_F_RDWR, Struct.U32.class, DscpPolicyValue.class);
         mBpfDscpIpv6Policies = new BpfMap<Struct.U32, DscpPolicyValue>(IPV6_POLICY_MAP_PATH,
                 BpfMap.BPF_F_RDWR, Struct.U32.class, DscpPolicyValue.class);
     }
 
+    private boolean isUnusedIndex(int index) {
+        for (SparseIntArray ifacePolicies : mIfaceIndexToPolicyIdBpfMapIndex.values()) {
+            if (ifacePolicies.indexOfValue(index) >= 0) return false;
+        }
+        return true;
+    }
+
     private int getFirstFreeIndex() {
+        if (mIfaceIndexToPolicyIdBpfMapIndex.size() == 0) return 0;
         for (int i = 0; i < MAX_POLICIES; i++) {
-            if (mPolicyIdToBpfMapIndex.indexOfValue(i) < 0) return i;
+            if (isUnusedIndex(i)) {
+                return i;
+            }
         }
         return MAX_POLICIES;
     }
 
+    private int findIndex(int policyId, int ifIndex) {
+        SparseIntArray ifacePolicies = mIfaceIndexToPolicyIdBpfMapIndex.get(ifIndex);
+        if (ifacePolicies != null) {
+            final int existingIndex = ifacePolicies.get(policyId, -1);
+            if (existingIndex != -1) {
+                return existingIndex;
+            }
+        }
+
+        final int firstIndex = getFirstFreeIndex();
+        if (firstIndex >= MAX_POLICIES) {
+            // New policy is being added, but max policies has already been reached.
+            return -1;
+        }
+        return firstIndex;
+    }
+
     private void sendStatus(NetworkAgentInfo nai, int policyId, int status) {
         try {
             nai.networkAgent.onDscpPolicyStatusUpdated(policyId, status);
         } catch (RemoteException e) {
-            Log.d(TAG, "Failed update policy status: ", e);
+            Log.e(TAG, "Failed update policy status: ", e);
         }
     }
 
@@ -107,37 +147,43 @@
                         || policy.getSourceAddress() instanceof Inet6Address));
     }
 
-    private int addDscpPolicyInternal(DscpPolicy policy) {
+    private int getIfaceIndex(NetworkAgentInfo nai) {
+        String iface = nai.linkProperties.getInterfaceName();
+        NetworkInterface netIface;
+        try {
+            netIface = NetworkInterface.getByName(iface);
+        } catch (IOException e) {
+            Log.e(TAG, "Unable to get iface index for " + iface + ": " + e);
+            netIface = null;
+        }
+        return (netIface != null) ? netIface.getIndex() : 0;
+    }
+
+    private int addDscpPolicyInternal(DscpPolicy policy, int ifIndex) {
         // If there is no existing policy with a matching ID, and we are already at
         // the maximum number of policies then return INSUFFICIENT_PROCESSING_RESOURCES.
-        final int existingIndex = mPolicyIdToBpfMapIndex.get(policy.getPolicyId(), -1);
-        if (existingIndex == -1 && mPolicyIdToBpfMapIndex.size() >= MAX_POLICIES) {
-            return DSCP_POLICY_STATUS_INSUFFICIENT_PROCESSING_RESOURCES;
+        SparseIntArray ifacePolicies = mIfaceIndexToPolicyIdBpfMapIndex.get(ifIndex);
+        if (ifacePolicies == null) {
+            ifacePolicies = new SparseIntArray(MAX_POLICIES);
         }
 
         // Currently all classifiers are supported, if any are removed return
         // DSCP_POLICY_STATUS_REQUESTED_CLASSIFIER_NOT_SUPPORTED,
         // and for any other generic error DSCP_POLICY_STATUS_REQUEST_DECLINED
 
-        int addIndex = 0;
-        // If a policy with a matching ID exists, replace it, otherwise use the next free
-        // index for the policy.
-        if (existingIndex != -1) {
-            addIndex = mPolicyIdToBpfMapIndex.get(policy.getPolicyId());
-        } else {
-            addIndex = getFirstFreeIndex();
+        final int addIndex = findIndex(policy.getPolicyId(), ifIndex);
+        if (addIndex == -1) {
+            return DSCP_POLICY_STATUS_INSUFFICIENT_PROCESSING_RESOURCES;
         }
 
         try {
-            mPolicyIdToBpfMapIndex.put(policy.getPolicyId(), addIndex);
-
             // Add v4 policy to mBpfDscpIpv4Policies if source and destination address
-            // are both null or if they are both instances of Inet6Address.
+            // are both null or if they are both instances of Inet4Address.
             if (matchesIpv4(policy)) {
                 mBpfDscpIpv4Policies.insertOrReplaceEntry(
                         new Struct.U32(addIndex),
                         new DscpPolicyValue(policy.getSourceAddress(),
-                            policy.getDestinationAddress(),
+                            policy.getDestinationAddress(), ifIndex,
                             policy.getSourcePort(), policy.getDestinationPortRange(),
                             (short) policy.getProtocol(), (short) policy.getDscpValue()));
             }
@@ -148,10 +194,16 @@
                 mBpfDscpIpv6Policies.insertOrReplaceEntry(
                         new Struct.U32(addIndex),
                         new DscpPolicyValue(policy.getSourceAddress(),
-                                policy.getDestinationAddress(),
+                                policy.getDestinationAddress(), ifIndex,
                                 policy.getSourcePort(), policy.getDestinationPortRange(),
                                 (short) policy.getProtocol(), (short) policy.getDscpValue()));
             }
+
+            ifacePolicies.put(policy.getPolicyId(), addIndex);
+            // Only add the policy to the per interface map if the policy was successfully
+            // added to both bpf maps above. It is safe to assume that if insert fails for
+            // one map then it fails for both.
+            mIfaceIndexToPolicyIdBpfMapIndex.put(ifIndex, ifacePolicies);
         } catch (ErrnoException e) {
             Log.e(TAG, "Failed to insert policy into map: ", e);
             return DSCP_POLICY_STATUS_INSUFFICIENT_PROCESSING_RESOURCES;
@@ -166,6 +218,7 @@
      *
      * DSCP_POLICY_STATUS_SUCCESS - if policy was added successfully
      * DSCP_POLICY_STATUS_INSUFFICIENT_PROCESSING_RESOURCES - if max policies were already set
+     * DSCP_POLICY_STATUS_REQUEST_DECLINED - Interface index was invalid
      */
     public void addDscpPolicy(NetworkAgentInfo nai, DscpPolicy policy) {
         if (!mAttachedIfaces.contains(nai.linkProperties.getInterfaceName())) {
@@ -177,11 +230,19 @@
             }
         }
 
-        int status = addDscpPolicyInternal(policy);
+        final int ifIndex = getIfaceIndex(nai);
+        if (ifIndex == 0) {
+            Log.e(TAG, "Iface index is invalid");
+            sendStatus(nai, policy.getPolicyId(), DSCP_POLICY_STATUS_REQUEST_DECLINED);
+            return;
+        }
+
+        int status = addDscpPolicyInternal(policy, ifIndex);
         sendStatus(nai, policy.getPolicyId(), status);
     }
 
-    private void removePolicyFromMap(NetworkAgentInfo nai, int policyId, int index) {
+    private void removePolicyFromMap(NetworkAgentInfo nai, int policyId, int index,
+            boolean sendCallback) {
         int status = DSCP_POLICY_STATUS_POLICY_NOT_FOUND;
         try {
             mBpfDscpIpv4Policies.replaceEntry(new Struct.U32(index), DscpPolicyValue.NONE);
@@ -191,7 +252,9 @@
             Log.e(TAG, "Failed to delete policy from map: ", e);
         }
 
-        sendStatus(nai, policyId, status);
+        if (sendCallback) {
+            sendStatus(nai, policyId, status);
+        }
     }
 
     /**
@@ -204,36 +267,44 @@
             return;
         }
 
-        if (mPolicyIdToBpfMapIndex.get(policyId, -1) != -1) {
-            removePolicyFromMap(nai, policyId, mPolicyIdToBpfMapIndex.get(policyId));
-            mPolicyIdToBpfMapIndex.delete(policyId);
+        SparseIntArray ifacePolicies = mIfaceIndexToPolicyIdBpfMapIndex.get(getIfaceIndex(nai));
+        if (ifacePolicies == null) return;
+
+        final int existingIndex = ifacePolicies.get(policyId, -1);
+        if (existingIndex == -1) {
+            Log.e(TAG, "Policy " + policyId + " does not exist in map.");
+            sendStatus(nai, policyId, DSCP_POLICY_STATUS_POLICY_NOT_FOUND);
+            return;
         }
 
-        // TODO: detach should only occur if no more policies are present on the nai's iface.
-        if (mPolicyIdToBpfMapIndex.size() == 0) {
+        removePolicyFromMap(nai, policyId, existingIndex, true);
+        ifacePolicies.delete(policyId);
+
+        if (ifacePolicies.size() == 0) {
             detachProgram(nai.linkProperties.getInterfaceName());
         }
     }
 
     /**
-     * Remove all DSCP policies and detach program.
+     * Remove all DSCP policies and detach program. Send callback if requested.
      */
-    // TODO: Remove all should only remove policies from corresponding nai iface.
-    public void removeAllDscpPolicies(NetworkAgentInfo nai) {
+    public void removeAllDscpPolicies(NetworkAgentInfo nai, boolean sendCallback) {
         if (!mAttachedIfaces.contains(nai.linkProperties.getInterfaceName())) {
             // Nothing to remove since program is not attached. Send update for policy
             // id 0. The status update must contain a policy ID, and 0 is an invalid id.
-            sendStatus(nai, 0, DSCP_POLICY_STATUS_SUCCESS);
+            if (sendCallback) {
+                sendStatus(nai, 0, DSCP_POLICY_STATUS_SUCCESS);
+            }
             return;
         }
 
-        for (int i = 0; i < mPolicyIdToBpfMapIndex.size(); i++) {
-            removePolicyFromMap(nai, mPolicyIdToBpfMapIndex.keyAt(i),
-                    mPolicyIdToBpfMapIndex.valueAt(i));
+        SparseIntArray ifacePolicies = mIfaceIndexToPolicyIdBpfMapIndex.get(getIfaceIndex(nai));
+        if (ifacePolicies == null) return;
+        for (int i = 0; i < ifacePolicies.size(); i++) {
+            removePolicyFromMap(nai, ifacePolicies.keyAt(i), ifacePolicies.valueAt(i),
+                    sendCallback);
         }
-        mPolicyIdToBpfMapIndex.clear();
-
-        // Can detach program since no policies are active.
+        ifacePolicies.clear();
         detachProgram(nai.linkProperties.getInterfaceName());
     }
 
@@ -241,12 +312,12 @@
      * Attach BPF program
      */
     private boolean attachProgram(@NonNull String iface) {
-        // TODO: attach needs to be per iface not program.
-
         try {
             NetworkInterface netIface = NetworkInterface.getByName(iface);
+            boolean isEth = TcUtils.isEthernet(iface);
+            String path = PROG_PATH + (isEth ? "_ether" : "_raw_ip");
             TcUtils.tcFilterAddDevBpf(netIface.getIndex(), false, PRIO_DSCP, (short) ETH_P_ALL,
-                    PROG_PATH);
+                    path);
         } catch (IOException e) {
             Log.e(TAG, "Unable to attach to TC on " + iface + ": " + e);
             return false;
@@ -264,9 +335,9 @@
             if (netIface != null) {
                 TcUtils.tcFilterDelDev(netIface.getIndex(), false, PRIO_DSCP, (short) ETH_P_ALL);
             }
+            mAttachedIfaces.remove(iface);
         } catch (IOException e) {
             Log.e(TAG, "Unable to detach to TC on " + iface + ": " + e);
         }
-        mAttachedIfaces.remove(iface);
     }
 }
diff --git a/service/src/com/android/server/connectivity/DscpPolicyValue.java b/service/src/com/android/server/connectivity/DscpPolicyValue.java
index cb40306..6e4e7eb 100644
--- a/service/src/com/android/server/connectivity/DscpPolicyValue.java
+++ b/service/src/com/android/server/connectivity/DscpPolicyValue.java
@@ -31,29 +31,31 @@
 public class DscpPolicyValue extends Struct {
     private static final String TAG = DscpPolicyValue.class.getSimpleName();
 
-    // TODO: add the interface index.
     @Field(order = 0, type = Type.ByteArray, arraysize = 16)
     public final byte[] src46;
 
     @Field(order = 1, type = Type.ByteArray, arraysize = 16)
     public final byte[] dst46;
 
-    @Field(order = 2, type = Type.UBE16)
-    public final int srcPort;
+    @Field(order = 2, type = Type.U32)
+    public final long ifIndex;
 
     @Field(order = 3, type = Type.UBE16)
-    public final int dstPortStart;
+    public final int srcPort;
 
     @Field(order = 4, type = Type.UBE16)
+    public final int dstPortStart;
+
+    @Field(order = 5, type = Type.UBE16)
     public final int dstPortEnd;
 
-    @Field(order = 5, type = Type.U8)
+    @Field(order = 6, type = Type.U8)
     public final short proto;
 
-    @Field(order = 6, type = Type.U8)
+    @Field(order = 7, type = Type.U8)
     public final short dscp;
 
-    @Field(order = 7, type = Type.U8, padding = 3)
+    @Field(order = 8, type = Type.U8, padding = 3)
     public final short mask;
 
     private static final int SRC_IP_MASK = 0x1;
@@ -69,6 +71,7 @@
         return true;
     }
 
+    // TODO:  move to frameworks/libs/net and have this and BpfCoordinator import it.
     private byte[] toIpv4MappedAddressBytes(InetAddress ia) {
         final byte[] addr6 = new byte[16];
         if (ia != null) {
@@ -117,13 +120,12 @@
         return mask;
     }
 
-    // This constructor is necessary for BpfMap#getValue since all values must be
-    // in the constructor.
-    public DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final int srcPort,
-            final int dstPortStart, final int dstPortEnd, final short proto,
+    private DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final long ifIndex,
+            final int srcPort, final int dstPortStart, final int dstPortEnd, final short proto,
             final short dscp) {
         this.src46 = toAddressField(src46);
         this.dst46 = toAddressField(dst46);
+        this.ifIndex = ifIndex;
 
         // These params need to be stored as 0 because uints are used in BpfMap.
         // If they are -1 BpfMap write will throw errors.
@@ -138,15 +140,15 @@
         this.mask = makeMask(this.src46, this.dst46, srcPort, dstPortStart, proto, dscp);
     }
 
-    public DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final int srcPort,
-            final Range<Integer> dstPort, final short proto,
+    public DscpPolicyValue(final InetAddress src46, final InetAddress dst46, final long ifIndex,
+            final int srcPort, final Range<Integer> dstPort, final short proto,
             final short dscp) {
-        this(src46, dst46, srcPort, dstPort != null ? dstPort.getLower() : -1,
+        this(src46, dst46, ifIndex, srcPort, dstPort != null ? dstPort.getLower() : -1,
                 dstPort != null ? dstPort.getUpper() : -1, proto, dscp);
     }
 
     public static final DscpPolicyValue NONE = new DscpPolicyValue(
-            null /* src46 */, null /* dst46 */, -1 /* srcPort */,
+            null /* src46 */, null /* dst46 */, 0 /* ifIndex */, -1 /* srcPort */,
             -1 /* dstPortStart */, -1 /* dstPortEnd */, (short) -1 /* proto */,
             (short) 0 /* dscp */);
 
@@ -170,9 +172,9 @@
 
         try {
             return String.format(
-                    "src46: %s, dst46: %s, srcPort: %d, dstPortStart: %d, dstPortEnd: %d,"
-                    + " protocol: %d, dscp %s", srcIpString, dstIpString, srcPort, dstPortStart,
-                    dstPortEnd, proto, dscp);
+                    "src46: %s, dst46: %s, ifIndex: %d, srcPort: %d, dstPortStart: %d,"
+                    + " dstPortEnd: %d, protocol: %d, dscp %s", srcIpString, dstIpString,
+                    ifIndex, srcPort, dstPortStart, dstPortEnd, proto, dscp);
         } catch (IllegalArgumentException e) {
             return String.format("String format error: " + e);
         }
diff --git a/tests/cts/net/src/android/net/cts/DscpPolicyTest.kt b/tests/cts/net/src/android/net/cts/DscpPolicyTest.kt
index e8add6b..1e42fe6 100644
--- a/tests/cts/net/src/android/net/cts/DscpPolicyTest.kt
+++ b/tests/cts/net/src/android/net/cts/DscpPolicyTest.kt
@@ -48,9 +48,11 @@
 import android.platform.test.annotations.AppModeFull
 import android.system.Os
 import android.system.OsConstants.AF_INET
+import android.system.OsConstants.AF_INET6
 import android.system.OsConstants.IPPROTO_UDP
 import android.system.OsConstants.SOCK_DGRAM
 import android.system.OsConstants.SOCK_NONBLOCK
+import android.util.Log
 import android.util.Range
 import androidx.test.InstrumentationRegistry
 import androidx.test.runner.AndroidJUnit4
@@ -71,6 +73,8 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import java.net.Inet4Address
+import java.net.Inet6Address
+import java.net.InetAddress
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.util.regex.Pattern
@@ -81,6 +85,9 @@
 
 private const val MAX_PACKET_LENGTH = 1500
 
+private const val IP4_PREFIX_LEN = 32
+private const val IP6_PREFIX_LEN = 128
+
 private val instrumentation: Instrumentation
     get() = InstrumentationRegistry.getInstrumentation()
 
@@ -97,6 +104,9 @@
     private val LOCAL_IPV4_ADDRESS = InetAddresses.parseNumericAddress("192.0.2.1")
     private val TEST_TARGET_IPV4_ADDR =
             InetAddresses.parseNumericAddress("8.8.8.8") as Inet4Address
+    private val LOCAL_IPV6_ADDRESS = InetAddresses.parseNumericAddress("2001:db8::1")
+    private val TEST_TARGET_IPV6_ADDR =
+            InetAddresses.parseNumericAddress("2001:4860:4860::8888") as Inet6Address
 
     private val realContext = InstrumentationRegistry.getContext()
     private val cm = realContext.getSystemService(ConnectivityManager::class.java)
@@ -132,7 +142,9 @@
         runAsShell(MANAGE_TEST_NETWORKS) {
             val tnm = realContext.getSystemService(TestNetworkManager::class.java)
 
-            iface = tnm.createTunInterface(Array(1) { LinkAddress(LOCAL_IPV4_ADDRESS, 32) })
+            iface = tnm.createTunInterface(arrayOf(
+                    LinkAddress(LOCAL_IPV4_ADDRESS, IP4_PREFIX_LEN),
+                    LinkAddress(LOCAL_IPV6_ADDRESS, IP6_PREFIX_LEN)))
             assertNotNull(iface)
         }
 
@@ -154,6 +166,8 @@
 
         // reader.stop() cleans up tun fd
         reader.handler.post { reader.stop() }
+        if (iface.fileDescriptor.fileDescriptor != null)
+            Os.close(iface.fileDescriptor.fileDescriptor)
         handlerThread.quitSafely()
     }
 
@@ -196,9 +210,11 @@
             }
         }
         val lp = LinkProperties().apply {
-            addLinkAddress(LinkAddress(LOCAL_IPV4_ADDRESS, 32))
+            addLinkAddress(LinkAddress(LOCAL_IPV4_ADDRESS, IP4_PREFIX_LEN))
+            addLinkAddress(LinkAddress(LOCAL_IPV6_ADDRESS, IP6_PREFIX_LEN))
             addRoute(RouteInfo(IpPrefix("0.0.0.0/0"), null, null))
-            setInterfaceName(iface.getInterfaceName())
+            addRoute(RouteInfo(InetAddress.getByName("fe80::1234")))
+            setInterfaceName(specifier)
         }
         val config = NetworkAgentConfig.Builder().build()
         val agent = TestableNetworkAgent(context, handlerThread.looper, nc, lp, config)
@@ -218,47 +234,114 @@
         eachByte -> "%02x".format(eachByte)
     }
 
-    fun checkDscpValue(
+    fun sendPacket(
         agent: TestableNetworkAgent,
-        callback: TestableNetworkCallback,
-        dscpValue: Int = 0,
-        dstPort: Int = 0
+        sendV6: Boolean,
+        dstPort: Int = 0,
     ) {
         val testString = "test string"
         val testPacket = ByteBuffer.wrap(testString.toByteArray(Charsets.UTF_8))
         var packetFound = false
 
-        val socket = Os.socket(AF_INET, SOCK_DGRAM or SOCK_NONBLOCK, IPPROTO_UDP)
+        val socket = Os.socket(if (sendV6) AF_INET6 else AF_INET, SOCK_DGRAM or SOCK_NONBLOCK,
+                IPPROTO_UDP)
         agent.network.bindSocket(socket)
 
         val originalPacket = testPacket.readAsArray()
-        Os.sendto(socket, originalPacket, 0 /* bytesOffset */, originalPacket.size,
-                0 /* flags */, TEST_TARGET_IPV4_ADDR, dstPort)
-
+        Os.sendto(socket, originalPacket, 0 /* bytesOffset */, originalPacket.size, 0 /* flags */,
+                if(sendV6) TEST_TARGET_IPV6_ADDR else TEST_TARGET_IPV4_ADDR, dstPort)
         Os.close(socket)
+    }
+
+    fun parseV4PacketDscp(buffer : ByteBuffer) : Int {
+        val ip_ver = buffer.get()
+        val tos = buffer.get()
+        val length = buffer.getShort()
+        val id = buffer.getShort()
+        val offset = buffer.getShort()
+        val ttl = buffer.get()
+        val ipType = buffer.get()
+        val checksum = buffer.getShort()
+        return tos.toInt().shr(2)
+    }
+
+    fun parseV6PacketDscp(buffer : ByteBuffer) : Int {
+        val ip_ver = buffer.get()
+        val tc = buffer.get()
+        val fl = buffer.getShort()
+        val length = buffer.getShort()
+        val proto = buffer.get()
+        val hop = buffer.get()
+        // DSCP is bottom 4 bits of ip_ver and top 2 of tc.
+        val ip_ver_bottom = ip_ver.toInt().and(0xf)
+        val tc_dscp = tc.toInt().shr(6)
+        return ip_ver_bottom.toInt().shl(2) + tc_dscp
+    }
+
+    fun parsePacketIp(
+        buffer : ByteBuffer,
+        sendV6 : Boolean,
+    ) : Boolean {
+        val ipAddr = if (sendV6) ByteArray(16) else ByteArray(4)
+        buffer.get(ipAddr)
+        val srcIp = if (sendV6) Inet6Address.getByAddress(ipAddr)
+                else Inet4Address.getByAddress(ipAddr)
+        buffer.get(ipAddr)
+        val dstIp = if (sendV6) Inet6Address.getByAddress(ipAddr)
+                else Inet4Address.getByAddress(ipAddr)
+
+        Log.e(TAG, "IP Src:" + srcIp + " dst: " + dstIp)
+
+        if ((sendV6 && srcIp == LOCAL_IPV6_ADDRESS && dstIp == TEST_TARGET_IPV6_ADDR) ||
+                (!sendV6 && srcIp == LOCAL_IPV4_ADDRESS && dstIp == TEST_TARGET_IPV4_ADDR)) {
+            Log.e(TAG, "IP return true");
+            return true
+        }
+        Log.e(TAG, "IP return false");
+        return false
+    }
+
+    fun parsePacketPort(
+        buffer : ByteBuffer,
+        srcPort : Int,
+        dstPort : Int
+    ) : Boolean {
+        if (srcPort == 0 && dstPort == 0) return true
+
+        val packetSrcPort = buffer.getShort().toInt()
+        val packetDstPort = buffer.getShort().toInt()
+
+        Log.e(TAG, "Port Src:" + packetSrcPort + " dst: " + packetDstPort)
+
+        if ((srcPort == 0 || (srcPort != 0 && srcPort == packetSrcPort)) &&
+                (dstPort == 0 || (dstPort != 0 && dstPort == packetDstPort))) {
+            Log.e(TAG, "Port return true");
+            return true
+        }
+        Log.e(TAG, "Port return false");
+        return false
+    }
+
+    fun validatePacket(
+        agent : TestableNetworkAgent,
+        sendV6 : Boolean = false,
+        dscpValue : Int = 0,
+        dstPort : Int = 0,
+    ) {
+        var packetFound = false;
+        sendPacket(agent, sendV6, dstPort)
+        // TODO: grab source port from socket in sendPacket
+
+        Log.e(TAG, "find DSCP value:" + dscpValue)
         generateSequence { reader.poll(PACKET_TIMEOUT_MS) }.forEach { packet ->
             val buffer = ByteBuffer.wrap(packet, 0, packet.size).order(ByteOrder.BIG_ENDIAN)
-            val ip_ver = buffer.get()
-            val tos = buffer.get()
-            val length = buffer.getShort()
-            val id = buffer.getShort()
-            val offset = buffer.getShort()
-            val ttl = buffer.get()
-            val ipType = buffer.get()
-            val checksum = buffer.getShort()
+            val dscp = if (sendV6) parseV6PacketDscp(buffer) else parseV4PacketDscp(buffer)
+            Log.e(TAG, "DSCP value:" + dscp)
 
-            val ipAddr = ByteArray(4)
-            buffer.get(ipAddr)
-            val srcIp = Inet4Address.getByAddress(ipAddr)
-            buffer.get(ipAddr)
-            val dstIp = Inet4Address.getByAddress(ipAddr)
-            val packetSrcPort = buffer.getShort().toInt()
-            val packetDstPort = buffer.getShort().toInt()
-
-            // TODO: Add source port comparison.
-            if (srcIp == LOCAL_IPV4_ADDRESS && dstIp == TEST_TARGET_IPV4_ADDR &&
-                    packetDstPort == dstPort) {
-                assertEquals(dscpValue, (tos.toInt().shr(2)))
+            // TODO: Add source port comparison. Use 0 for now.
+            if (parsePacketIp(buffer, sendV6) && parsePacketPort(buffer, 0, dstPort)) {
+                Log.e(TAG, "DSCP value found")
+                assertEquals(dscpValue, dscp)
                 packetFound = true
             }
         }
@@ -275,12 +358,12 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(policyId, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
-            checkDscpValue(agent, callback, dstPort = portNumber)
         }
     }
 
     @Test
-    fun testDscpPolicyAddPolicies(): Unit = createConnectedNetworkAgent().let { (agent, callback) ->
+    fun testDscpPolicyAddPolicies(): Unit = createConnectedNetworkAgent().let {
+                (agent, callback) ->
         val policy = DscpPolicy.Builder(1, 1)
                 .setDestinationPortRange(Range(4444, 4444)).build()
         agent.sendAddDscpPolicy(policy)
@@ -288,8 +371,7 @@
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
         }
-
-        checkDscpValue(agent, callback, dscpValue = 1, dstPort = 4444)
+        validatePacket(agent, dscpValue = 1, dstPort = 4444)
 
         agent.sendRemoveDscpPolicy(1)
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
@@ -298,15 +380,54 @@
         }
 
         val policy2 = DscpPolicy.Builder(1, 4)
-                .setDestinationPortRange(Range(5555, 5555)).setSourceAddress(LOCAL_IPV4_ADDRESS)
-                .setDestinationAddress(TEST_TARGET_IPV4_ADDR).setProtocol(IPPROTO_UDP).build()
+                .setDestinationPortRange(Range(5555, 5555))
+                .setDestinationAddress(TEST_TARGET_IPV4_ADDR)
+                .setSourceAddress(LOCAL_IPV4_ADDRESS)
+                .setProtocol(IPPROTO_UDP).build()
         agent.sendAddDscpPolicy(policy2)
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
         }
 
-        checkDscpValue(agent, callback, dscpValue = 4, dstPort = 5555)
+        validatePacket(agent, dscpValue = 4, dstPort = 5555)
+
+        agent.sendRemoveDscpPolicy(1)
+        agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
+            assertEquals(1, it.policyId)
+            assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
+        }
+    }
+
+    @Test
+    fun testDscpPolicyAddV6Policies(): Unit = createConnectedNetworkAgent().let {
+                (agent, callback) ->
+        val policy = DscpPolicy.Builder(1, 1)
+                .setDestinationPortRange(Range(4444, 4444)).build()
+        agent.sendAddDscpPolicy(policy)
+        agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
+            assertEquals(1, it.policyId)
+            assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
+        }
+        validatePacket(agent, true, dscpValue = 1, dstPort = 4444)
+
+        agent.sendRemoveDscpPolicy(1)
+        agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
+            assertEquals(1, it.policyId)
+            assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
+        }
+
+        val policy2 = DscpPolicy.Builder(1, 4)
+                .setDestinationPortRange(Range(5555, 5555))
+                .setDestinationAddress(TEST_TARGET_IPV6_ADDR)
+                .setSourceAddress(LOCAL_IPV6_ADDRESS)
+                .setProtocol(IPPROTO_UDP).build()
+        agent.sendAddDscpPolicy(policy2)
+        agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
+            assertEquals(1, it.policyId)
+            assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
+        }
+        validatePacket(agent, true, dscpValue = 4, dstPort = 5555)
 
         agent.sendRemoveDscpPolicy(1)
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
@@ -324,7 +445,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 1111)
+            validatePacket(agent, dscpValue = 1, dstPort = 1111)
         }
 
         val policy2 = DscpPolicy.Builder(2, 1).setDestinationPortRange(Range(2222, 2222)).build()
@@ -332,7 +453,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(2, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 2222)
+            validatePacket(agent, dscpValue = 1, dstPort = 2222)
         }
 
         val policy3 = DscpPolicy.Builder(3, 1).setDestinationPortRange(Range(3333, 3333)).build()
@@ -340,13 +461,16 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(3, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 3333)
+            validatePacket(agent, dscpValue = 1, dstPort = 3333)
         }
 
         /* Remove Policies and check CE is no longer set */
         doRemovePolicyTest(agent, callback, 1)
+        validatePacket(agent, dscpValue = 0, dstPort = 1111)
         doRemovePolicyTest(agent, callback, 2)
+        validatePacket(agent, dscpValue = 0, dstPort = 2222)
         doRemovePolicyTest(agent, callback, 3)
+        validatePacket(agent, dscpValue = 0, dstPort = 3333)
     }
 
     @Test
@@ -357,7 +481,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 1111)
+            validatePacket(agent, dscpValue = 1, dstPort = 1111)
         }
         doRemovePolicyTest(agent, callback, 1)
 
@@ -366,7 +490,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(2, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 2222)
+            validatePacket(agent, dscpValue = 1, dstPort = 2222)
         }
         doRemovePolicyTest(agent, callback, 2)
 
@@ -375,7 +499,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(3, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 3333)
+            validatePacket(agent, dscpValue = 1, dstPort = 3333)
         }
         doRemovePolicyTest(agent, callback, 3)
     }
@@ -389,7 +513,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 1111)
+            validatePacket(agent, dscpValue = 1, dstPort = 1111)
         }
 
         val policy2 = DscpPolicy.Builder(2, 1).setDestinationPortRange(Range(2222, 2222)).build()
@@ -397,7 +521,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(2, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 2222)
+            validatePacket(agent, dscpValue = 1, dstPort = 2222)
         }
 
         val policy3 = DscpPolicy.Builder(3, 1).setDestinationPortRange(Range(3333, 3333)).build()
@@ -405,7 +529,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(3, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 3333)
+            validatePacket(agent, dscpValue = 1, dstPort = 3333)
         }
 
         /* Remove Policies and check CE is no longer set */
@@ -423,14 +547,15 @@
     }
 
     @Test
-    fun testRemoveAllDscpPolicies(): Unit = createConnectedNetworkAgent().let { (agent, callback) ->
+    fun testRemoveAllDscpPolicies(): Unit = createConnectedNetworkAgent().let {
+                (agent, callback) ->
         val policy = DscpPolicy.Builder(1, 1)
                 .setDestinationPortRange(Range(1111, 1111)).build()
         agent.sendAddDscpPolicy(policy)
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 1111)
+            validatePacket(agent, dscpValue = 1, dstPort = 1111)
         }
 
         val policy2 = DscpPolicy.Builder(2, 1)
@@ -439,7 +564,7 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(2, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 2222)
+            validatePacket(agent, dscpValue = 1, dstPort = 2222)
         }
 
         val policy3 = DscpPolicy.Builder(3, 1)
@@ -448,24 +573,24 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(3, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 3333)
+            validatePacket(agent, dscpValue = 1, dstPort = 3333)
         }
 
         agent.sendRemoveAllDscpPolicies()
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
-            checkDscpValue(agent, callback, dstPort = 1111)
+            validatePacket(agent, false, dstPort = 1111)
         }
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(2, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
-            checkDscpValue(agent, callback, dstPort = 2222)
+            validatePacket(agent, false, dstPort = 2222)
         }
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(3, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_DELETED, it.status)
-            checkDscpValue(agent, callback, dstPort = 3333)
+            validatePacket(agent, false, dstPort = 3333)
         }
     }
 
@@ -477,12 +602,9 @@
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
             assertEquals(1, it.policyId)
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 4444)
+            validatePacket(agent, dscpValue = 1, dstPort = 4444)
         }
 
-        // TODO: send packet on socket and confirm that changing the DSCP policy
-        // updates the mark to the new value.
-
         val policy2 = DscpPolicy.Builder(1, 1).setDestinationPortRange(Range(5555, 5555)).build()
         agent.sendAddDscpPolicy(policy2)
         agent.expectCallback<OnDscpPolicyStatusUpdated>().let {
@@ -490,8 +612,8 @@
             assertEquals(DSCP_POLICY_STATUS_SUCCESS, it.status)
 
             // Sending packet with old policy should fail
-            checkDscpValue(agent, callback, dstPort = 4444)
-            checkDscpValue(agent, callback, dscpValue = 1, dstPort = 5555)
+            validatePacket(agent, dscpValue = 0, dstPort = 4444)
+            validatePacket(agent, dscpValue = 1, dstPort = 5555)
         }
 
         agent.sendRemoveDscpPolicy(1)