From 6615676d825e02d271fe7a9ca78a77ac3773ab93 Mon Sep 17 00:00:00 2001
From: Pablo Neira Ayuso <pablo@netfilter.org>
Date: Mon, 3 Aug 2015 15:50:03 +0200
Subject: src: add per-bytes limit

This example show how to accept packets below the ratelimit:

... limit rate 1024 mbytes/second counter accept

You need a Linux kernel >= 4.3-rc1.

Signed-off-by: Pablo Neira Ayuso <pablo@netfilter.org>
---
 src/datatype.c            | 55 +++++++++++++++++++++++++++++++++++++++++++++++
 src/netlink_delinearize.c |  1 +
 src/netlink_linearize.c   |  1 +
 src/parser_bison.y        | 17 +++++++++++++++
 src/statement.c           | 43 ++++++++++++++++++++++++++++++++++--
 5 files changed, 115 insertions(+), 2 deletions(-)

(limited to 'src')

diff --git a/src/datatype.c b/src/datatype.c
index f79f5d2d..e5a486fb 100644
--- a/src/datatype.c
+++ b/src/datatype.c
@@ -976,3 +976,58 @@ void concat_type_destroy(const struct datatype *dtype)
 		xfree(dtype);
 	}
 }
+
+static struct error_record *time_unit_parse(const struct location *loc,
+					    const char *str, uint64_t *unit)
+{
+	if (strcmp(str, "second") == 0)
+		*unit = 1ULL;
+	else if (strcmp(str, "minute") == 0)
+		*unit = 1ULL * 60;
+	else if (strcmp(str, "hour") == 0)
+		*unit = 1ULL * 60 * 60;
+	else if (strcmp(str, "day") == 0)
+		*unit = 1ULL * 60 * 60 * 24;
+	else if (strcmp(str, "week") == 0)
+		*unit = 1ULL * 60 * 60 * 24 * 7;
+	else
+		return error(loc, "Wrong rate format");
+
+	return NULL;
+}
+
+static struct error_record *data_unit_parse(const struct location *loc,
+					    const char *str, uint64_t *rate)
+{
+	if (strncmp(str, "bytes", strlen("bytes")) == 0)
+		*rate = 1ULL;
+	else if (strncmp(str, "kbytes", strlen("kbytes")) == 0)
+		*rate = 1024;
+	else if (strncmp(str, "mbytes", strlen("mbytes")) == 0)
+		*rate = 1024 * 1024;
+	else
+		return error(loc, "Wrong rate format");
+
+	return NULL;
+}
+
+struct error_record *rate_parse(const struct location *loc, const char *str,
+				uint64_t *rate, uint64_t *unit)
+{
+	struct error_record *erec;
+	const char *slash;
+
+	slash = strchr(str, '/');
+	if (!slash)
+		return error(loc, "wrong rate format");
+
+	erec = data_unit_parse(loc, str, rate);
+	if (erec != NULL)
+		return erec;
+
+	erec = time_unit_parse(loc, slash + 1, unit);
+	if (erec != NULL)
+		return erec;
+
+	return NULL;
+}
diff --git a/src/netlink_delinearize.c b/src/netlink_delinearize.c
index dc6338c2..4c639a16 100644
--- a/src/netlink_delinearize.c
+++ b/src/netlink_delinearize.c
@@ -583,6 +583,7 @@ static void netlink_parse_limit(struct netlink_parse_ctx *ctx,
 	stmt = limit_stmt_alloc(loc);
 	stmt->limit.rate = nftnl_expr_get_u64(nle, NFTNL_EXPR_LIMIT_RATE);
 	stmt->limit.unit = nftnl_expr_get_u64(nle, NFTNL_EXPR_LIMIT_UNIT);
+	stmt->limit.type = nftnl_expr_get_u32(nle, NFTNL_EXPR_LIMIT_TYPE);
 	list_add_tail(&stmt->list, &ctx->rule->stmts);
 }
 
diff --git a/src/netlink_linearize.c b/src/netlink_linearize.c
index b2cb98dd..47092d33 100644
--- a/src/netlink_linearize.c
+++ b/src/netlink_linearize.c
@@ -708,6 +708,7 @@ static void netlink_gen_limit_stmt(struct netlink_linearize_ctx *ctx,
 	nle = alloc_nft_expr("limit");
 	nftnl_expr_set_u64(nle, NFTNL_EXPR_LIMIT_RATE, stmt->limit.rate);
 	nftnl_expr_set_u64(nle, NFTNL_EXPR_LIMIT_UNIT, stmt->limit.unit);
+	nftnl_expr_set_u32(nle, NFTNL_EXPR_LIMIT_TYPE, stmt->limit.type);
 	nftnl_rule_add_expr(ctx->nlr, nle);
 }
 
diff --git a/src/parser_bison.y b/src/parser_bison.y
index cfb6b707..ec44a2cd 100644
--- a/src/parser_bison.y
+++ b/src/parser_bison.y
@@ -1446,6 +1446,23 @@ limit_stmt		:	LIMIT	RATE	NUM	SLASH	time_unit
 				$$ = limit_stmt_alloc(&@$);
 				$$->limit.rate	= $3;
 				$$->limit.unit	= $5;
+				$$->limit.type	= NFT_LIMIT_PKTS;
+			}
+			|	LIMIT RATE	NUM	STRING
+			{
+				struct error_record *erec;
+				uint64_t rate, unit;
+
+				erec = rate_parse(&@$, $4, &rate, &unit);
+				if (erec != NULL) {
+					erec_queue(erec, state->msgs);
+					YYERROR;
+				}
+
+				$$ = limit_stmt_alloc(&@$);
+				$$->limit.rate	= rate * $3;
+				$$->limit.unit	= unit;
+				$$->limit.type	= NFT_LIMIT_PKT_BYTES;
 			}
 			;
 
diff --git a/src/statement.c b/src/statement.c
index 9ebc5938..ba7b8be2 100644
--- a/src/statement.c
+++ b/src/statement.c
@@ -185,10 +185,49 @@ static const char *get_unit(uint64_t u)
 	return "error";
 }
 
+static const char *data_unit[] = {
+	"bytes",
+	"kbytes",
+	"mbytes",
+	NULL
+};
+
+static const char *get_rate(uint64_t byte_rate, uint64_t *rate)
+{
+	uint64_t res, prev, rest;
+	int i;
+
+	res = prev = byte_rate;
+	for (i = 0;; i++) {
+		rest = res % 1024;
+		res /= 1024;
+		if (res <= 1 && rest != 0)
+			break;
+		if (data_unit[i + 1] == NULL)
+			break;
+		prev = res;
+	}
+	*rate = prev;
+	return data_unit[i];
+}
+
 static void limit_stmt_print(const struct stmt *stmt)
 {
-	printf("limit rate %" PRIu64 "/%s",
-	       stmt->limit.rate, get_unit(stmt->limit.unit));
+	const char *data_unit;
+	uint64_t rate;
+
+	switch (stmt->limit.type) {
+	case NFT_LIMIT_PKTS:
+		printf("limit rate %" PRIu64 "/%s",
+		       stmt->limit.rate, get_unit(stmt->limit.unit));
+		break;
+	case NFT_LIMIT_PKT_BYTES:
+		data_unit = get_rate(stmt->limit.rate, &rate);
+
+		printf("limit rate %" PRIu64 " %s/%s",
+		       rate, data_unit, get_unit(stmt->limit.unit));
+		break;
+	}
 }
 
 static const struct stmt_ops limit_stmt_ops = {
-- 
cgit v1.2.3