bpf_progs/dscp_policy.c - platform/packages/modules/Connectivity - Git at Google

 /*
  * Copyright (C) 2021 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.
  */

 #include <linux/types.h>
 #include <linux/bpf.h>
 #include <linux/ip.h>
 #include <linux/ipv6.h>
 #include <linux/if_ether.h>
 #include <linux/pkt_cls.h>
 #include <linux/tcp.h>
 #include <stdint.h>
 #include <netinet/in.h>
 #include <netinet/udp.h>
 #include <string.h>

 #include "bpf_helpers.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,
         AID_SYSTEM)
 DEFINE_BPF_MAP_GRW(ipv6_dscp_policies_map, ARRAY, uint32_t, Ipv6RuleEntry, 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);

     // use this with HASH map so map lookup only happens once policies have been added?
     if (!selectedMap) {
         return TC_ACT_PIPE;
     }

     // used for map lookup
     uint64_t cookie = bpf_get_socket_cookie(skb);

     // 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;

         // Must have ipv4 header
         if (data + sizeof(*iph) > data_end) return TC_ACT_PIPE;

         // IP version must be 4
         if (iph->version != 4) return TC_ACT_PIPE;

         // We cannot handle IP options, just standard 20 byte == 5 dword minimal IPv4 header
         if (iph->ihl != 5) return TC_ACT_PIPE;

         if (iph->protocol != IPPROTO_UDP) return TC_ACT_PIPE;

         struct udphdr *udp;
         udp = data + sizeof(struct iphdr); //sizeof(struct ethhdr)

         if ((void*)(udp + 1) > data_end) return TC_ACT_PIPE;

         // 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()

             // TODO: fix checksum...
             int ecn = iph->tos & 3;
             uint8_t newDscpVal = (v4Policy->dscpVal << 2) + ecn;
             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);
             return TC_ACT_PIPE;
         }

         // linear scan ipv4_dscp_policies_map, stored socket params do not match actual
         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 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 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->mask) {
                     bestMatch = i;
                     bestScore = score;
                 }
             }
         }

         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;
             }
         }

         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
         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;
         if (*selectedMap == MAP_A) {
             v6Policy = bpf_ipv6_socket_to_policies_map_A_lookup_elem(&cookie);
         } else {
             v6Policy = bpf_ipv6_socket_to_policies_map_B_lookup_elem(&cookie);
         }

         if (!v6Policy)
             return TC_ACT_PIPE;

         // TODO: Add code to process IPv6 packet.
     }

     // Always return TC_ACT_PIPE
     return TC_ACT_PIPE;
 }

 LICENSE("Apache 2.0");
 CRITICAL("Connectivity");
	/*
	* Copyright (C) 2021 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.
	*/

	#include <linux/types.h>
	#include <linux/bpf.h>
	#include <linux/ip.h>
	#include <linux/ipv6.h>
	#include <linux/if_ether.h>
	#include <linux/pkt_cls.h>
	#include <linux/tcp.h>
	#include <stdint.h>
	#include <netinet/in.h>
	#include <netinet/udp.h>
	#include <string.h>

	#include "bpf_helpers.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,
	AID_SYSTEM)
	DEFINE_BPF_MAP_GRW(ipv6_dscp_policies_map, ARRAY, uint32_t, Ipv6RuleEntry, 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);

	// use this with HASH map so map lookup only happens once policies have been added?
	if (!selectedMap) {
	return TC_ACT_PIPE;
	}

	// used for map lookup
	uint64_t cookie = bpf_get_socket_cookie(skb);

	// 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;

	// Must have ipv4 header
	if (data + sizeof(*iph) > data_end) return TC_ACT_PIPE;

	// IP version must be 4
	if (iph->version != 4) return TC_ACT_PIPE;

	// We cannot handle IP options, just standard 20 byte == 5 dword minimal IPv4 header
	if (iph->ihl != 5) return TC_ACT_PIPE;

	if (iph->protocol != IPPROTO_UDP) return TC_ACT_PIPE;

	struct udphdr *udp;
	udp = data + sizeof(struct iphdr); //sizeof(struct ethhdr)

	if ((void*)(udp + 1) > data_end) return TC_ACT_PIPE;

	// 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()

	// TODO: fix checksum...
	int ecn = iph->tos & 3;
	uint8_t newDscpVal = (v4Policy->dscpVal << 2) + ecn;
	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);
	return TC_ACT_PIPE;
	}

	// linear scan ipv4_dscp_policies_map, stored socket params do not match actual
	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 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 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->mask) {
	bestMatch = i;
	bestScore = score;
	}
	}
	}

	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;
	}
	}

	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
	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;
	if (*selectedMap == MAP_A) {
	v6Policy = bpf_ipv6_socket_to_policies_map_A_lookup_elem(&cookie);
	} else {
	v6Policy = bpf_ipv6_socket_to_policies_map_B_lookup_elem(&cookie);
	}

	if (!v6Policy)
	return TC_ACT_PIPE;

	// TODO: Add code to process IPv6 packet.
	}

	// Always return TC_ACT_PIPE
	return TC_ACT_PIPE;
	}

	LICENSE("Apache 2.0");
	CRITICAL("Connectivity");