1/*
2 * Copyright (c) 2008-2009 Patrick McHardy <kaber@trash.net>
3 *
4 * This program is free software; you can redistribute it and/or modify
5 * it under the terms of the GNU General Public License version 2 as
6 * published by the Free Software Foundation.
7 *
8 * Development of this code funded by Astaro AG (http://www.astaro.com/)
9 */
10
11#include <linux/kernel.h>
12#include <linux/init.h>
13#include <linux/module.h>
14#include <linux/spinlock.h>
15#include <linux/netlink.h>
16#include <linux/netfilter.h>
17#include <linux/netfilter/nf_tables.h>
18#include <net/netfilter/nf_tables.h>
19
20static DEFINE_SPINLOCK(limit_lock);
21
22struct nft_limit {
23	u64		last;
24	u64		tokens;
25	u64		tokens_max;
26	u64		rate;
27	u64		nsecs;
28	u32		burst;
29};
30
31static inline bool nft_limit_eval(struct nft_limit *limit, u64 cost)
32{
33	u64 now, tokens;
34	s64 delta;
35
36	spin_lock_bh(&limit_lock);
37	now = ktime_get_ns();
38	tokens = limit->tokens + now - limit->last;
39	if (tokens > limit->tokens_max)
40		tokens = limit->tokens_max;
41
42	limit->last = now;
43	delta = tokens - cost;
44	if (delta >= 0) {
45		limit->tokens = delta;
46		spin_unlock_bh(&limit_lock);
47		return false;
48	}
49	limit->tokens = tokens;
50	spin_unlock_bh(&limit_lock);
51	return true;
52}
53
54static int nft_limit_init(struct nft_limit *limit,
55			  const struct nlattr * const tb[])
56{
57	u64 unit;
58
59	if (tb[NFTA_LIMIT_RATE] == NULL ||
60	    tb[NFTA_LIMIT_UNIT] == NULL)
61		return -EINVAL;
62
63	limit->rate = be64_to_cpu(nla_get_be64(tb[NFTA_LIMIT_RATE]));
64	unit = be64_to_cpu(nla_get_be64(tb[NFTA_LIMIT_UNIT]));
65	limit->nsecs = unit * NSEC_PER_SEC;
66	if (limit->rate == 0 || limit->nsecs < unit)
67		return -EOVERFLOW;
68	limit->tokens = limit->tokens_max = limit->nsecs;
69
70	if (tb[NFTA_LIMIT_BURST]) {
71		u64 rate;
72
73		limit->burst = ntohl(nla_get_be32(tb[NFTA_LIMIT_BURST]));
74
75		rate = limit->rate + limit->burst;
76		if (rate < limit->rate)
77			return -EOVERFLOW;
78
79		limit->rate = rate;
80	}
81	limit->last = ktime_get_ns();
82
83	return 0;
84}
85
86static int nft_limit_dump(struct sk_buff *skb, const struct nft_limit *limit,
87			  enum nft_limit_type type)
88{
89	u64 secs = div_u64(limit->nsecs, NSEC_PER_SEC);
90	u64 rate = limit->rate - limit->burst;
91
92	if (nla_put_be64(skb, NFTA_LIMIT_RATE, cpu_to_be64(rate)) ||
93	    nla_put_be64(skb, NFTA_LIMIT_UNIT, cpu_to_be64(secs)) ||
94	    nla_put_be32(skb, NFTA_LIMIT_BURST, htonl(limit->burst)) ||
95	    nla_put_be32(skb, NFTA_LIMIT_TYPE, htonl(type)))
96		goto nla_put_failure;
97	return 0;
98
99nla_put_failure:
100	return -1;
101}
102
103struct nft_limit_pkts {
104	struct nft_limit	limit;
105	u64			cost;
106};
107
108static void nft_limit_pkts_eval(const struct nft_expr *expr,
109				struct nft_regs *regs,
110				const struct nft_pktinfo *pkt)
111{
112	struct nft_limit_pkts *priv = nft_expr_priv(expr);
113
114	if (nft_limit_eval(&priv->limit, priv->cost))
115		regs->verdict.code = NFT_BREAK;
116}
117
118static const struct nla_policy nft_limit_policy[NFTA_LIMIT_MAX + 1] = {
119	[NFTA_LIMIT_RATE]	= { .type = NLA_U64 },
120	[NFTA_LIMIT_UNIT]	= { .type = NLA_U64 },
121	[NFTA_LIMIT_BURST]	= { .type = NLA_U32 },
122	[NFTA_LIMIT_TYPE]	= { .type = NLA_U32 },
123};
124
125static int nft_limit_pkts_init(const struct nft_ctx *ctx,
126			       const struct nft_expr *expr,
127			       const struct nlattr * const tb[])
128{
129	struct nft_limit_pkts *priv = nft_expr_priv(expr);
130	int err;
131
132	err = nft_limit_init(&priv->limit, tb);
133	if (err < 0)
134		return err;
135
136	priv->cost = div_u64(priv->limit.nsecs, priv->limit.rate);
137	return 0;
138}
139
140static int nft_limit_pkts_dump(struct sk_buff *skb, const struct nft_expr *expr)
141{
142	const struct nft_limit_pkts *priv = nft_expr_priv(expr);
143
144	return nft_limit_dump(skb, &priv->limit, NFT_LIMIT_PKTS);
145}
146
147static struct nft_expr_type nft_limit_type;
148static const struct nft_expr_ops nft_limit_pkts_ops = {
149	.type		= &nft_limit_type,
150	.size		= NFT_EXPR_SIZE(sizeof(struct nft_limit_pkts)),
151	.eval		= nft_limit_pkts_eval,
152	.init		= nft_limit_pkts_init,
153	.dump		= nft_limit_pkts_dump,
154};
155
156static void nft_limit_pkt_bytes_eval(const struct nft_expr *expr,
157				     struct nft_regs *regs,
158				     const struct nft_pktinfo *pkt)
159{
160	struct nft_limit *priv = nft_expr_priv(expr);
161	u64 cost = div_u64(priv->nsecs * pkt->skb->len, priv->rate);
162
163	if (nft_limit_eval(priv, cost))
164		regs->verdict.code = NFT_BREAK;
165}
166
167static int nft_limit_pkt_bytes_init(const struct nft_ctx *ctx,
168				    const struct nft_expr *expr,
169				    const struct nlattr * const tb[])
170{
171	struct nft_limit *priv = nft_expr_priv(expr);
172
173	return nft_limit_init(priv, tb);
174}
175
176static int nft_limit_pkt_bytes_dump(struct sk_buff *skb,
177				    const struct nft_expr *expr)
178{
179	const struct nft_limit *priv = nft_expr_priv(expr);
180
181	return nft_limit_dump(skb, priv, NFT_LIMIT_PKT_BYTES);
182}
183
184static const struct nft_expr_ops nft_limit_pkt_bytes_ops = {
185	.type		= &nft_limit_type,
186	.size		= NFT_EXPR_SIZE(sizeof(struct nft_limit)),
187	.eval		= nft_limit_pkt_bytes_eval,
188	.init		= nft_limit_pkt_bytes_init,
189	.dump		= nft_limit_pkt_bytes_dump,
190};
191
192static const struct nft_expr_ops *
193nft_limit_select_ops(const struct nft_ctx *ctx,
194		     const struct nlattr * const tb[])
195{
196	if (tb[NFTA_LIMIT_TYPE] == NULL)
197		return &nft_limit_pkts_ops;
198
199	switch (ntohl(nla_get_be32(tb[NFTA_LIMIT_TYPE]))) {
200	case NFT_LIMIT_PKTS:
201		return &nft_limit_pkts_ops;
202	case NFT_LIMIT_PKT_BYTES:
203		return &nft_limit_pkt_bytes_ops;
204	}
205	return ERR_PTR(-EOPNOTSUPP);
206}
207
208static struct nft_expr_type nft_limit_type __read_mostly = {
209	.name		= "limit",
210	.select_ops	= nft_limit_select_ops,
211	.policy		= nft_limit_policy,
212	.maxattr	= NFTA_LIMIT_MAX,
213	.flags		= NFT_EXPR_STATEFUL,
214	.owner		= THIS_MODULE,
215};
216
217static int __init nft_limit_module_init(void)
218{
219	return nft_register_expr(&nft_limit_type);
220}
221
222static void __exit nft_limit_module_exit(void)
223{
224	nft_unregister_expr(&nft_limit_type);
225}
226
227module_init(nft_limit_module_init);
228module_exit(nft_limit_module_exit);
229
230MODULE_LICENSE("GPL");
231MODULE_AUTHOR("Patrick McHardy <kaber@trash.net>");
232MODULE_ALIAS_NFT_EXPR("limit");
233