client_data_usage/cdu_db.c

/* cdu_db.c
 *
 * This file defines client data usage database.
 *
 * Author: Cradlepoint Technology, Inc.  <source@cradlepoint.com>
 *		Adrian Sitterle <asitterle@cradlepoint.com>
 *
 * Copyright (C) 2019 Cradlepoint Technology, Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2
 * as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 */

#include <linux/version.h>
#include <linux/types.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/jhash.h>
#include <linux/list_sort.h>
#include <linux/etherdevice.h>
#include <linux/netfilter/xt_usage.h>
#include <linux/netfilter/x_tables.h>
#include <net/netfilter/nf_conntrack.h>
#include <net/netfilter/nf_conntrack_core.h>
#include <net/netfilter/nf_conntrack_acct.h>

#include "cdu_db.h"
#include "cdu_xt_target.h"

struct hlist_head cdu_db_usage_hash[USAGE_HASH_BITS] __read_mostly;
static int cdu_db_usage_count __read_mostly;
struct list_head cdu_db_usage_list; /* Sorted list */
DEFINE_SPINLOCK(cdu_db_hash_lock);

static unsigned int jhash_rnd __read_mostly;

/* Use tuple 0 (original dir) to deterine lan ip info */
static const union nf_inet_addr *get_lan_ip_from_ct(const struct nf_conn *ct)
{

	if (ct->orig_direction == CDU_DIR_DEST)
		return &ct->tuplehash[IP_CT_DIR_REPLY].tuple.src.u3;
	else if (ct->orig_direction == CDU_DIR_SRC)
		return &ct->tuplehash[IP_CT_DIR_ORIGINAL].tuple.src.u3;

	CDU_DEBUG("Direction is not set, but should be at this point. Don't know which ip to use.\n");
	return NULL;
}

/* Deterine lan mac info provided by hook in nf_conntrack_in() */
static const unsigned char *get_lan_mac_from_ct(const struct nf_conn *ct)
{

	if (ct->orig_direction != CDU_DIR_NOT_SET)
		return ct->mac;

	CDU_DEBUG("Direction is not set, but should be at this point. Don't know which mac to use.\n");
	return NULL;
}

/* Use tuple 0 (original dir) to deterine lan l3 proto info */
static unsigned char get_lan_family_from_ct(const struct nf_conn *ct)
{
	return ct->tuplehash[IP_CT_DIR_ORIGINAL].tuple.src.l3num;
}

/* Calculate htable bucket key */
static inline unsigned int calc_hash_key_from_ct(const struct nf_conn *ct)
{
	unsigned char family;
	const union nf_inet_addr *addr;

	family = get_lan_family_from_ct(ct);
	addr = get_lan_ip_from_ct(ct);

	if (!addr)
		return 0;

	if (family == NFPROTO_IPV4)
		return jhash((const char *)&addr->ip, IP_ALEN, jhash_rnd) & (USAGE_HASH_BITS - 1);

	if (family == NFPROTO_IPV6)
		return jhash((const char *)&addr->ip6, IP6_ALEN, jhash_rnd) & (USAGE_HASH_BITS - 1);

	return 0;
}

/* Try to find an entry in the cdu_db_usage_hash table based on information in ct */
static struct usage_entry *cdu_db_find_entry(const struct nf_conn *ct)
{
	unsigned int key;
	struct usage_entry *obj = NULL;
	unsigned char family;
	const union nf_inet_addr *paddr;
	const unsigned char *mac;

	family = get_lan_family_from_ct(ct);
	paddr = get_lan_ip_from_ct(ct);
	mac = get_lan_mac_from_ct(ct);

	if (!paddr || !mac)
		return NULL;

	key = calc_hash_key_from_ct(ct);
	if (hlist_empty(&cdu_db_usage_hash[key]))
		return NULL;

	hlist_for_each_entry(obj, &cdu_db_usage_hash[key], node) {
		CDU_DEBUG("searching... obj key=%u ip=%pI4 ip6=%pI6 mac=%pM\n", key, &obj->addr.ip, &obj->addr.ip6, obj->mac);
		/* find first ip, which will always be last updated entry due to hlist_add_head */
		if ((family == NFPROTO_IPV4 && (obj->addr.ip == paddr->ip)) || (family == NFPROTO_IPV6 && !memcmp(obj->addr.ip6, paddr->ip6, IP6_ALEN))) {
			if (!mac || ether_addr_equal(mac, obj->mac)) {
				CDU_DEBUG("Found entry in htable ^\n");
				return obj;
			}
		}
	}

	CDU_DEBUG("Did not find entry in htable.\n");
	return NULL;
}

static int cdu_last_activity_time_cmp(void * priv, struct list_head *a, struct list_head *b)
{
	struct usage_entry * obj_a = container_of(a, struct usage_entry, lnode);
	struct usage_entry * obj_b = container_of(b, struct usage_entry, lnode);

	if (obj_a->last_time < obj_b->last_time)
		return -1;

	if(obj_a->last_time > obj_b->last_time)
		return 1;

	return 0;

}

static struct list_head const* sort_cdu_db(void)
{
#ifdef DEBUG_ENABLE
	struct timespec start;
	struct timespec end;
	ktime_get_boottime_ts64(&start);
#endif

	if (list_empty(&cdu_db_usage_list))
		return NULL;

	list_sort(NULL, &cdu_db_usage_list, cdu_last_activity_time_cmp);
#ifdef DEBUG_ENABLE
	ktime_get_boottime_ts64(&end);
	CDU_DEBUG("%s took %lu ns\n", __func__, end.tv_nsec - start.tv_nsec);
#endif
	return &cdu_db_usage_list;
}

static void cdu_drop_clients(void)
{
	struct list_head const * sorted_list;
	struct list_head  * pos, *n;
	struct usage_entry *obj;
	unsigned int client_drop_count = CDU_USAGE_DROP_CHUNK;
	sorted_list = sort_cdu_db();

	if (!sorted_list)
		return;
	list_for_each_safe(pos, n, sorted_list) {
		if (--client_drop_count == 0)
			return;
		obj = container_of(pos, struct usage_entry, lnode);
		CDU_INFO("DROPPPING MAC = %pm, IP = %pI4, IP6 = %pI6, Last_Time = %lu, First_Time = %lu, Connect_Time = %lu\n",
					obj->mac, &obj->addr.ip, &obj->addr.ip6,
					obj->last_time, obj->first_time, obj->last_time - obj->first_time);
		cdu_db_remove(obj);
	}
}

/* Add new entry to the hash table, based on information in ct */
static struct usage_entry *cdu_db_add_entry(const struct nf_conn *ct)
{
	struct usage_entry *obj;
	struct timespec64 now;
	const union nf_inet_addr *paddr;
	const unsigned char *mac;

	paddr = get_lan_ip_from_ct(ct);
	mac = get_lan_mac_from_ct(ct);

	if (!paddr || !mac)
		return NULL;

	if (unlikely(cdu_db_usage_count >= CDU_USAGE_MAXENTRY)) {
		CDU_ERROR("Usage limit exceeded, dropping oldest clients\n");
		cdu_drop_clients();
	}

	obj = kzalloc(sizeof(struct usage_entry), GFP_ATOMIC);

	ktime_get_boottime_ts64(&now);
	obj->last_time = obj->first_time = (unsigned long) now.tv_sec;

	ether_addr_copy(obj->mac, mac);
	obj->family = get_lan_family_from_ct(ct);

	if (obj->family == NFPROTO_IPV4)
		obj->addr.ip = paddr->ip;
	else if (obj->family == NFPROTO_IPV6)
		memcpy(&obj->addr.ip6, &paddr->ip6, IP6_ALEN);

	obj->key = calc_hash_key_from_ct(ct);

	hlist_add_head(&obj->node, &cdu_db_usage_hash[obj->key]);
	list_add(&obj->lnode, &cdu_db_usage_list);
	cdu_db_usage_count++;

	if (obj->family == NFPROTO_IPV4)
		CDU_INFO("found new client ip=%pI4 mac=%pM (%d clients)\n", &obj->addr.ip, obj->mac, cdu_db_usage_count);
	else if (obj->family == NFPROTO_IPV6)
		CDU_INFO("found new client ip6=%pI6 mac=%pM (%d clients)\n", &obj->addr.ip6, obj->mac, cdu_db_usage_count);

	return obj;
}

/* external entry point for cdu_xt_target to add missing db entries
 * handles locking to avoid external needing to lock/find/add/unlock
 * necessary to utilize unconditional add_entry
 */
void cdu_db_add_entry_ifmissing(const struct nf_conn *ct)
{
	spin_lock_bh(&cdu_db_hash_lock);

	if (!cdu_db_find_entry(ct))
		cdu_db_add_entry(ct);

	spin_unlock_bh(&cdu_db_hash_lock);
}

/* Update information of a client in htable.
 * persist = CDU_DB_UPDATE_PERMA_SET (1), if called at the end of conntrack flow
 * persist = CDU_DB_UPDATE_EPHEM_SET (0), if called in the middle of a conntrack flow
 */
void cdu_db_update_entry(const struct nf_conn *ct, bool persist, bool clear)
{
	struct nf_conn_acct *acct;
	struct timespec64 now;
	u_int64_t bytes_up, bytes_dn, packets_up, packets_dn;

#ifdef ZERO_EPHEM_FLOWS
	struct nf_conn_counter *zeroed;
#endif
	struct nf_conn_counter *counter;
	struct usage_entry *obj = NULL;
	int up_dir = 0, dl_dir = 0;
	int dataset = STATSET(persist);

	if (ct->orig_direction == CDU_DIR_DEST)
		up_dir = 1;
	else if (ct->orig_direction == CDU_DIR_SRC)
		dl_dir = 1;
	else {
		CDU_DEBUG("Unknown direction\n");
		return;
	}

	acct = nf_conn_acct_find(ct);
	if (!acct) {
		CDU_DEBUG("Could not find conntrack acct entry for this ct.\n");
		return;
	}
	counter = acct->counter;

	spin_lock_bh(&cdu_db_hash_lock);

	obj = cdu_db_find_entry(ct);
	if (!obj) {
		obj = cdu_db_add_entry(ct);
		if (!obj)
			goto update_exit;
	}

	ktime_get_boottime_ts64(&now);
	obj->last_time = (unsigned long) now.tv_sec;

	bytes_up = atomic64_read(&counter[up_dir].bytes);
	packets_up = atomic64_read(&counter[up_dir].packets);
	bytes_dn = atomic64_read(&counter[dl_dir].bytes);
	packets_dn = atomic64_read(&counter[dl_dir].packets);

	obj->upload[dataset].bytes += bytes_up;
	obj->upload[dataset].packets += packets_up;
	obj->download[dataset].bytes += bytes_dn;
	obj->download[dataset].packets += packets_dn;

#ifdef ZERO_EPHEM_FLOWS
	zeroed  = acct->zeroed;

	/* remember what we need to subtract later */
	obj->upload[CLEARED_COUNTER].bytes += atomic64_read(&zeroed[up_dir].bytes);
	obj->upload[CLEARED_COUNTER].packets += atomic64_read(&zeroed[up_dir].packets);
	obj->download[CLEARED_COUNTER].bytes += atomic64_read(&zeroed[dl_dir].bytes);
	obj->download[CLEARED_COUNTER].packets += atomic64_read(&zeroed[dl_dir].packets);
	if (clear) {
		atomic64_set(&zeroed[up_dir].bytes, bytes_up);
		atomic64_set(&zeroed[up_dir].packets, packets_up);
		atomic64_set(&zeroed[dl_dir].bytes, bytes_dn);
		atomic64_set(&zeroed[dl_dir].packets, packets_dn);
	}
#endif

	CDU_DEBUG("Flow stats (set:%s) up: packets=%llu bytes %llu, dl: packets=%llu bytes %llu\n",
			(dataset)?"PERMA":"EPHEM",
			(uint64_t)atomic64_read(&counter[up_dir].packets), (uint64_t)atomic64_read(&counter[up_dir].bytes),
			(uint64_t)atomic64_read(&counter[dl_dir].packets), (uint64_t)atomic64_read(&counter[dl_dir].bytes));
	CDU_DEBUG("Updated %s client stats upload packets=%llu bytes %llu\n",
			(persist == CDU_DB_UPDATE_PERMA_SET)?"PERMA":"EPHEM",
			obj->upload[dataset].packets, obj->upload[dataset].bytes);
	CDU_DEBUG("Updated %s client stats download packets=%llu bytes %llu\n",
			(persist == CDU_DB_UPDATE_PERMA_SET)?"PERMA":"EPHEM",
			obj->download[dataset].packets, obj->download[dataset].bytes);

update_exit:
	spin_unlock_bh(&cdu_db_hash_lock);
}


void cdu_db_clear_stats(struct usage_entry *obj, bool persist)
{
	int dataset = STATSET(persist);

#ifdef ZERO_EPHEM_FLOWS
	if (dataset) {
		/* resetting ephemeral, also reset our 'cleared' total */
		obj->upload[CLEARED_COUNTER].bytes =
		obj->upload[CLEARED_COUNTER].packets =
		obj->download[CLEARED_COUNTER].bytes =
		obj->download[CLEARED_COUNTER].packets = 0;
	}
#endif
	obj->upload[dataset].bytes =
	obj->upload[dataset].packets =
	obj->download[dataset].bytes =
	obj->download[dataset].packets = 0;
}


void cdu_db_get_stats(struct usage_stats *upload, struct usage_stats *download, const struct usage_entry *obj)
{
	int direction; /* direction (upload or download) */
	struct usage_stats *target[] = {upload, download}; /* targets for the stats, in same order as data store */
	const struct usage_stats *store[] = {obj->upload, obj->download}; /* data stores, in same order as targets */
	int dataset; /* two datasets (ephemeral and persistent) to be summed */

	for (direction = 0; direction < 2; direction++) {
		target[direction]->bytes = target[direction]->packets = 0;
		for (dataset = 0; dataset < 2; dataset++) {
			target[direction]->bytes += store[direction][dataset].bytes;
			target[direction]->packets += store[direction][dataset].packets;
		}
#ifdef ZERO_EPHEM_FLOWS
#ifdef DEBUG_ENABLE
		if (store[direction][CLEARED_COUNTER].packets) {
			if (obj->family == NFPROTO_IPV4)
				CDU_DEBUG("Flow stats for %pI4, subtracting %llu from %llu bytes (%llu from %llu packets)\n",
					&obj->addr.ip,
					store[direction][CLEARED_COUNTER].bytes, target[direction]->bytes,
					store[direction][CLEARED_COUNTER].packets, target[direction]->packets
					);
			else
				CDU_DEBUG("Flow stats for %pI6, subtracting %llu from %llu bytes (%llu from %llu packets)\n",
					&obj->addr.ip6,
					store[direction][CLEARED_COUNTER].bytes, target[direction]->bytes,
					store[direction][CLEARED_COUNTER].packets, target[direction]->packets
					);
		}
#endif
		/* subtract counts from cleared partial sessions
		 * since the counts are from previous partial sessions,
		 * they will be less than or equal to the new counts
		 */
		target[direction]->bytes -= store[direction][CLEARED_COUNTER].bytes;
		target[direction]->packets -= store[direction][CLEARED_COUNTER].packets;
#endif
	}
}


void cdu_db_read_conn(bool clear)
{
	struct nf_conntrack_tuple_hash *h;
	const struct hlist_nulls_node *nn;
	struct nf_conn *ct_iter;
	int i;

	local_bh_disable();
	for (i = 0; i < nf_conntrack_htable_size; i++) {
		spin_lock(&nf_conntrack_locks[i % CONNTRACK_LOCKS]);
		if (i < nf_conntrack_htable_size) {
			hlist_nulls_for_each_entry(h, nn, &nf_conntrack_hash[i], hnnode) {
				if (NF_CT_DIRECTION(h) == IP_CT_DIR_ORIGINAL) {
					ct_iter = nf_ct_tuplehash_to_ctrack(h);
					if (ct_iter->orig_direction != CDU_DIR_NOT_SET) {
						CDU_DEBUG("ct iter %p (bucket: %d)\n", ct_iter, i);
						cdu_db_update_entry(ct_iter, CDU_DB_UPDATE_EPHEM_SET, clear);
					}
				}
			}
		}
		spin_unlock(&nf_conntrack_locks[i % CONNTRACK_LOCKS]);
	}
	local_bh_enable();
}


inline void cdu_db_remove(struct usage_entry *obj)
{
	cdu_db_usage_count--;
	hlist_del(&obj->node);
	list_del(&obj->lnode);
	kfree(obj);
}

void cdu_db_flush(void)
{
	struct usage_entry *obj;
	struct hlist_node *n;
	int i = 0;

	for (i = 0; i < USAGE_HASH_BITS; i++) {
		spin_lock_bh(&cdu_db_hash_lock);
		if (!hlist_empty(&cdu_db_usage_hash[i])) {
			hlist_for_each_entry_safe(obj, n, &cdu_db_usage_hash[i], node) {
				cdu_db_remove(obj);
			}
		}
		spin_unlock_bh(&cdu_db_hash_lock);
	}
}

void cdu_db_uninit(void)
{
	cdu_db_flush();
}

int cdu_db_init(void)
{
	int i;

	/* initialize empty hashtable buckets */
	for (i = 0; i < USAGE_HASH_BITS; i++)
		INIT_HLIST_HEAD(&cdu_db_usage_hash[i]);

	INIT_LIST_HEAD(&cdu_db_usage_list);

	get_random_bytes(&jhash_rnd, sizeof(jhash_rnd));

	CDU_DEBUG("db initialized\n");

	return 0;
}