1 // SPDX-License-Identifier: GPL-2.0-only
2 /*
3  * Copyright (c) 2016 Pablo Neira Ayuso <pablo@netfilter.org>
4  */
5 
6 #include <linux/kernel.h>
7 #include <linux/init.h>
8 #include <linux/module.h>
9 #include <linux/atomic.h>
10 #include <linux/netlink.h>
11 #include <linux/netfilter.h>
12 #include <linux/netfilter/nf_tables.h>
13 #include <net/netfilter/nf_tables.h>
14 
15 struct nft_quota {
16 	atomic64_t	quota;
17 	unsigned long	flags;
18 	atomic64_t	*consumed;
19 };
20 
21 static inline bool nft_overquota(struct nft_quota *priv,
22 				 const struct sk_buff *skb,
23 				 bool *report)
24 {
25 	u64 consumed = atomic64_add_return(skb->len, priv->consumed);
26 	u64 quota = atomic64_read(&priv->quota);
27 
28 	if (report)
29 		*report = consumed >= quota;
30 
31 	return consumed > quota;
32 }
33 
34 static inline bool nft_quota_invert(struct nft_quota *priv)
35 {
36 	return priv->flags & NFT_QUOTA_F_INV;
37 }
38 
39 static inline void nft_quota_do_eval(struct nft_quota *priv,
40 				     struct nft_regs *regs,
41 				     const struct nft_pktinfo *pkt)
42 {
43 	if (nft_overquota(priv, pkt->skb, NULL) ^ nft_quota_invert(priv))
44 		regs->verdict.code = NFT_BREAK;
45 }
46 
47 static const struct nla_policy nft_quota_policy[NFTA_QUOTA_MAX + 1] = {
48 	[NFTA_QUOTA_BYTES]	= { .type = NLA_U64 },
49 	[NFTA_QUOTA_FLAGS]	= { .type = NLA_U32 },
50 	[NFTA_QUOTA_CONSUMED]	= { .type = NLA_U64 },
51 };
52 
53 #define NFT_QUOTA_DEPLETED_BIT	1	/* From NFT_QUOTA_F_DEPLETED. */
54 
55 static void nft_quota_obj_eval(struct nft_object *obj,
56 			       struct nft_regs *regs,
57 			       const struct nft_pktinfo *pkt)
58 {
59 	struct nft_quota *priv = nft_obj_data(obj);
60 	bool overquota, report;
61 
62 	overquota = nft_overquota(priv, pkt->skb, &report);
63 	if (overquota ^ nft_quota_invert(priv))
64 		regs->verdict.code = NFT_BREAK;
65 
66 	if (report &&
67 	    !test_and_set_bit(NFT_QUOTA_DEPLETED_BIT, &priv->flags))
68 		nft_obj_notify(nft_net(pkt), obj->key.table, obj, 0, 0,
69 			       NFT_MSG_NEWOBJ, 0, nft_pf(pkt), 0, GFP_ATOMIC);
70 }
71 
72 static int nft_quota_do_init(const struct nlattr * const tb[],
73 			     struct nft_quota *priv)
74 {
75 	unsigned long flags = 0;
76 	u64 quota, consumed = 0;
77 
78 	if (!tb[NFTA_QUOTA_BYTES])
79 		return -EINVAL;
80 
81 	quota = be64_to_cpu(nla_get_be64(tb[NFTA_QUOTA_BYTES]));
82 	if (quota > S64_MAX)
83 		return -EOVERFLOW;
84 
85 	if (tb[NFTA_QUOTA_CONSUMED]) {
86 		consumed = be64_to_cpu(nla_get_be64(tb[NFTA_QUOTA_CONSUMED]));
87 		if (consumed > quota)
88 			return -EINVAL;
89 	}
90 
91 	if (tb[NFTA_QUOTA_FLAGS]) {
92 		flags = ntohl(nla_get_be32(tb[NFTA_QUOTA_FLAGS]));
93 		if (flags & ~NFT_QUOTA_F_INV)
94 			return -EINVAL;
95 		if (flags & NFT_QUOTA_F_DEPLETED)
96 			return -EOPNOTSUPP;
97 	}
98 
99 	priv->consumed = kmalloc(sizeof(*priv->consumed), GFP_KERNEL_ACCOUNT);
100 	if (!priv->consumed)
101 		return -ENOMEM;
102 
103 	atomic64_set(&priv->quota, quota);
104 	priv->flags = flags;
105 	atomic64_set(priv->consumed, consumed);
106 
107 	return 0;
108 }
109 
110 static void nft_quota_do_destroy(const struct nft_ctx *ctx,
111 				 struct nft_quota *priv)
112 {
113 	kfree(priv->consumed);
114 }
115 
116 static int nft_quota_obj_init(const struct nft_ctx *ctx,
117 			      const struct nlattr * const tb[],
118 			      struct nft_object *obj)
119 {
120 	struct nft_quota *priv = nft_obj_data(obj);
121 
122 	return nft_quota_do_init(tb, priv);
123 }
124 
125 static void nft_quota_obj_update(struct nft_object *obj,
126 				 struct nft_object *newobj)
127 {
128 	struct nft_quota *newpriv = nft_obj_data(newobj);
129 	struct nft_quota *priv = nft_obj_data(obj);
130 	u64 newquota;
131 
132 	newquota = atomic64_read(&newpriv->quota);
133 	atomic64_set(&priv->quota, newquota);
134 	priv->flags = newpriv->flags;
135 }
136 
137 static int nft_quota_do_dump(struct sk_buff *skb, struct nft_quota *priv,
138 			     bool reset)
139 {
140 	u64 consumed, consumed_cap, quota;
141 	u32 flags = priv->flags;
142 
143 	/* Since we inconditionally increment consumed quota for each packet
144 	 * that we see, don't go over the quota boundary in what we send to
145 	 * userspace.
146 	 */
147 	consumed = atomic64_read(priv->consumed);
148 	quota = atomic64_read(&priv->quota);
149 	if (consumed >= quota) {
150 		consumed_cap = quota;
151 		flags |= NFT_QUOTA_F_DEPLETED;
152 	} else {
153 		consumed_cap = consumed;
154 	}
155 
156 	if (nla_put_be64(skb, NFTA_QUOTA_BYTES, cpu_to_be64(quota),
157 			 NFTA_QUOTA_PAD) ||
158 	    nla_put_be64(skb, NFTA_QUOTA_CONSUMED, cpu_to_be64(consumed_cap),
159 			 NFTA_QUOTA_PAD) ||
160 	    nla_put_be32(skb, NFTA_QUOTA_FLAGS, htonl(flags)))
161 		goto nla_put_failure;
162 
163 	if (reset) {
164 		atomic64_sub(consumed, priv->consumed);
165 		clear_bit(NFT_QUOTA_DEPLETED_BIT, &priv->flags);
166 	}
167 	return 0;
168 
169 nla_put_failure:
170 	return -1;
171 }
172 
173 static int nft_quota_obj_dump(struct sk_buff *skb, struct nft_object *obj,
174 			      bool reset)
175 {
176 	struct nft_quota *priv = nft_obj_data(obj);
177 
178 	return nft_quota_do_dump(skb, priv, reset);
179 }
180 
181 static void nft_quota_obj_destroy(const struct nft_ctx *ctx,
182 				  struct nft_object *obj)
183 {
184 	struct nft_quota *priv = nft_obj_data(obj);
185 
186 	return nft_quota_do_destroy(ctx, priv);
187 }
188 
189 static struct nft_object_type nft_quota_obj_type;
190 static const struct nft_object_ops nft_quota_obj_ops = {
191 	.type		= &nft_quota_obj_type,
192 	.size		= sizeof(struct nft_quota),
193 	.init		= nft_quota_obj_init,
194 	.destroy	= nft_quota_obj_destroy,
195 	.eval		= nft_quota_obj_eval,
196 	.dump		= nft_quota_obj_dump,
197 	.update		= nft_quota_obj_update,
198 };
199 
200 static struct nft_object_type nft_quota_obj_type __read_mostly = {
201 	.type		= NFT_OBJECT_QUOTA,
202 	.ops		= &nft_quota_obj_ops,
203 	.maxattr	= NFTA_QUOTA_MAX,
204 	.policy		= nft_quota_policy,
205 	.owner		= THIS_MODULE,
206 };
207 
208 static void nft_quota_eval(const struct nft_expr *expr,
209 			   struct nft_regs *regs,
210 			   const struct nft_pktinfo *pkt)
211 {
212 	struct nft_quota *priv = nft_expr_priv(expr);
213 
214 	nft_quota_do_eval(priv, regs, pkt);
215 }
216 
217 static int nft_quota_init(const struct nft_ctx *ctx,
218 			  const struct nft_expr *expr,
219 			  const struct nlattr * const tb[])
220 {
221 	struct nft_quota *priv = nft_expr_priv(expr);
222 
223 	return nft_quota_do_init(tb, priv);
224 }
225 
226 static int nft_quota_dump(struct sk_buff *skb,
227 			  const struct nft_expr *expr, bool reset)
228 {
229 	struct nft_quota *priv = nft_expr_priv(expr);
230 
231 	return nft_quota_do_dump(skb, priv, reset);
232 }
233 
234 static void nft_quota_destroy(const struct nft_ctx *ctx,
235 			      const struct nft_expr *expr)
236 {
237 	struct nft_quota *priv = nft_expr_priv(expr);
238 
239 	return nft_quota_do_destroy(ctx, priv);
240 }
241 
242 static int nft_quota_clone(struct nft_expr *dst, const struct nft_expr *src, gfp_t gfp)
243 {
244 	struct nft_quota *priv_dst = nft_expr_priv(dst);
245 	struct nft_quota *priv_src = nft_expr_priv(src);
246 
247 	priv_dst->quota = priv_src->quota;
248 	priv_dst->flags = priv_src->flags;
249 
250 	priv_dst->consumed = kmalloc(sizeof(*priv_dst->consumed), gfp);
251 	if (!priv_dst->consumed)
252 		return -ENOMEM;
253 
254 	*priv_dst->consumed = *priv_src->consumed;
255 
256 	return 0;
257 }
258 
259 static struct nft_expr_type nft_quota_type;
260 static const struct nft_expr_ops nft_quota_ops = {
261 	.type		= &nft_quota_type,
262 	.size		= NFT_EXPR_SIZE(sizeof(struct nft_quota)),
263 	.eval		= nft_quota_eval,
264 	.init		= nft_quota_init,
265 	.destroy	= nft_quota_destroy,
266 	.clone		= nft_quota_clone,
267 	.dump		= nft_quota_dump,
268 	.reduce		= NFT_REDUCE_READONLY,
269 };
270 
271 static struct nft_expr_type nft_quota_type __read_mostly = {
272 	.name		= "quota",
273 	.ops		= &nft_quota_ops,
274 	.policy		= nft_quota_policy,
275 	.maxattr	= NFTA_QUOTA_MAX,
276 	.flags		= NFT_EXPR_STATEFUL,
277 	.owner		= THIS_MODULE,
278 };
279 
280 static int __init nft_quota_module_init(void)
281 {
282 	int err;
283 
284 	err = nft_register_obj(&nft_quota_obj_type);
285 	if (err < 0)
286 		return err;
287 
288 	err = nft_register_expr(&nft_quota_type);
289 	if (err < 0)
290 		goto err1;
291 
292 	return 0;
293 err1:
294 	nft_unregister_obj(&nft_quota_obj_type);
295 	return err;
296 }
297 
298 static void __exit nft_quota_module_exit(void)
299 {
300 	nft_unregister_expr(&nft_quota_type);
301 	nft_unregister_obj(&nft_quota_obj_type);
302 }
303 
304 module_init(nft_quota_module_init);
305 module_exit(nft_quota_module_exit);
306 
307 MODULE_LICENSE("GPL");
308 MODULE_AUTHOR("Pablo Neira Ayuso <pablo@netfilter.org>");
309 MODULE_ALIAS_NFT_EXPR("quota");
310 MODULE_ALIAS_NFT_OBJ(NFT_OBJECT_QUOTA);
311 MODULE_DESCRIPTION("Netfilter nftables quota module");
312