1  // SPDX-License-Identifier: GPL-2.0-only
2  /*
3   * Copyright (c) 2016 Laura Garcia <nevola@gmail.com>
4   */
5  
6  #include <linux/kernel.h>
7  #include <linux/init.h>
8  #include <linux/module.h>
9  #include <linux/netlink.h>
10  #include <linux/netfilter.h>
11  #include <linux/netfilter/nf_tables.h>
12  #include <linux/random.h>
13  #include <linux/static_key.h>
14  #include <net/netfilter/nf_tables.h>
15  #include <net/netfilter/nf_tables_core.h>
16  
17  struct nft_ng_inc {
18  	u8			dreg;
19  	u32			modulus;
20  	atomic_t		*counter;
21  	u32			offset;
22  };
23  
nft_ng_inc_gen(struct nft_ng_inc * priv)24  static u32 nft_ng_inc_gen(struct nft_ng_inc *priv)
25  {
26  	u32 nval, oval;
27  
28  	do {
29  		oval = atomic_read(priv->counter);
30  		nval = (oval + 1 < priv->modulus) ? oval + 1 : 0;
31  	} while (atomic_cmpxchg(priv->counter, oval, nval) != oval);
32  
33  	return nval + priv->offset;
34  }
35  
nft_ng_inc_eval(const struct nft_expr * expr,struct nft_regs * regs,const struct nft_pktinfo * pkt)36  static void nft_ng_inc_eval(const struct nft_expr *expr,
37  			    struct nft_regs *regs,
38  			    const struct nft_pktinfo *pkt)
39  {
40  	struct nft_ng_inc *priv = nft_expr_priv(expr);
41  
42  	regs->data[priv->dreg] = nft_ng_inc_gen(priv);
43  }
44  
45  static const struct nla_policy nft_ng_policy[NFTA_NG_MAX + 1] = {
46  	[NFTA_NG_DREG]		= { .type = NLA_U32 },
47  	[NFTA_NG_MODULUS]	= { .type = NLA_U32 },
48  	[NFTA_NG_TYPE]		= { .type = NLA_U32 },
49  	[NFTA_NG_OFFSET]	= { .type = NLA_U32 },
50  };
51  
nft_ng_inc_init(const struct nft_ctx * ctx,const struct nft_expr * expr,const struct nlattr * const tb[])52  static int nft_ng_inc_init(const struct nft_ctx *ctx,
53  			   const struct nft_expr *expr,
54  			   const struct nlattr * const tb[])
55  {
56  	struct nft_ng_inc *priv = nft_expr_priv(expr);
57  	int err;
58  
59  	if (tb[NFTA_NG_OFFSET])
60  		priv->offset = ntohl(nla_get_be32(tb[NFTA_NG_OFFSET]));
61  
62  	priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS]));
63  	if (priv->modulus == 0)
64  		return -ERANGE;
65  
66  	if (priv->offset + priv->modulus - 1 < priv->offset)
67  		return -EOVERFLOW;
68  
69  	priv->counter = kmalloc(sizeof(*priv->counter), GFP_KERNEL_ACCOUNT);
70  	if (!priv->counter)
71  		return -ENOMEM;
72  
73  	atomic_set(priv->counter, priv->modulus - 1);
74  
75  	err = nft_parse_register_store(ctx, tb[NFTA_NG_DREG], &priv->dreg,
76  				       NULL, NFT_DATA_VALUE, sizeof(u32));
77  	if (err < 0)
78  		goto err;
79  
80  	return 0;
81  err:
82  	kfree(priv->counter);
83  
84  	return err;
85  }
86  
nft_ng_inc_reduce(struct nft_regs_track * track,const struct nft_expr * expr)87  static bool nft_ng_inc_reduce(struct nft_regs_track *track,
88  				 const struct nft_expr *expr)
89  {
90  	const struct nft_ng_inc *priv = nft_expr_priv(expr);
91  
92  	nft_reg_track_cancel(track, priv->dreg, NFT_REG32_SIZE);
93  
94  	return false;
95  }
96  
nft_ng_dump(struct sk_buff * skb,enum nft_registers dreg,u32 modulus,enum nft_ng_types type,u32 offset)97  static int nft_ng_dump(struct sk_buff *skb, enum nft_registers dreg,
98  		       u32 modulus, enum nft_ng_types type, u32 offset)
99  {
100  	if (nft_dump_register(skb, NFTA_NG_DREG, dreg))
101  		goto nla_put_failure;
102  	if (nla_put_be32(skb, NFTA_NG_MODULUS, htonl(modulus)))
103  		goto nla_put_failure;
104  	if (nla_put_be32(skb, NFTA_NG_TYPE, htonl(type)))
105  		goto nla_put_failure;
106  	if (nla_put_be32(skb, NFTA_NG_OFFSET, htonl(offset)))
107  		goto nla_put_failure;
108  
109  	return 0;
110  
111  nla_put_failure:
112  	return -1;
113  }
114  
nft_ng_inc_dump(struct sk_buff * skb,const struct nft_expr * expr,bool reset)115  static int nft_ng_inc_dump(struct sk_buff *skb,
116  			   const struct nft_expr *expr, bool reset)
117  {
118  	const struct nft_ng_inc *priv = nft_expr_priv(expr);
119  
120  	return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_INCREMENTAL,
121  			   priv->offset);
122  }
123  
nft_ng_inc_destroy(const struct nft_ctx * ctx,const struct nft_expr * expr)124  static void nft_ng_inc_destroy(const struct nft_ctx *ctx,
125  			       const struct nft_expr *expr)
126  {
127  	const struct nft_ng_inc *priv = nft_expr_priv(expr);
128  
129  	kfree(priv->counter);
130  }
131  
132  struct nft_ng_random {
133  	u8			dreg;
134  	u32			modulus;
135  	u32			offset;
136  };
137  
nft_ng_random_gen(const struct nft_ng_random * priv)138  static u32 nft_ng_random_gen(const struct nft_ng_random *priv)
139  {
140  	return reciprocal_scale(get_random_u32(), priv->modulus) + priv->offset;
141  }
142  
nft_ng_random_eval(const struct nft_expr * expr,struct nft_regs * regs,const struct nft_pktinfo * pkt)143  static void nft_ng_random_eval(const struct nft_expr *expr,
144  			       struct nft_regs *regs,
145  			       const struct nft_pktinfo *pkt)
146  {
147  	struct nft_ng_random *priv = nft_expr_priv(expr);
148  
149  	regs->data[priv->dreg] = nft_ng_random_gen(priv);
150  }
151  
nft_ng_random_init(const struct nft_ctx * ctx,const struct nft_expr * expr,const struct nlattr * const tb[])152  static int nft_ng_random_init(const struct nft_ctx *ctx,
153  			      const struct nft_expr *expr,
154  			      const struct nlattr * const tb[])
155  {
156  	struct nft_ng_random *priv = nft_expr_priv(expr);
157  
158  	if (tb[NFTA_NG_OFFSET])
159  		priv->offset = ntohl(nla_get_be32(tb[NFTA_NG_OFFSET]));
160  
161  	priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS]));
162  	if (priv->modulus == 0)
163  		return -ERANGE;
164  
165  	if (priv->offset + priv->modulus - 1 < priv->offset)
166  		return -EOVERFLOW;
167  
168  	return nft_parse_register_store(ctx, tb[NFTA_NG_DREG], &priv->dreg,
169  					NULL, NFT_DATA_VALUE, sizeof(u32));
170  }
171  
nft_ng_random_dump(struct sk_buff * skb,const struct nft_expr * expr,bool reset)172  static int nft_ng_random_dump(struct sk_buff *skb,
173  			      const struct nft_expr *expr, bool reset)
174  {
175  	const struct nft_ng_random *priv = nft_expr_priv(expr);
176  
177  	return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_RANDOM,
178  			   priv->offset);
179  }
180  
nft_ng_random_reduce(struct nft_regs_track * track,const struct nft_expr * expr)181  static bool nft_ng_random_reduce(struct nft_regs_track *track,
182  				 const struct nft_expr *expr)
183  {
184  	const struct nft_ng_random *priv = nft_expr_priv(expr);
185  
186  	nft_reg_track_cancel(track, priv->dreg, NFT_REG32_SIZE);
187  
188  	return false;
189  }
190  
191  static struct nft_expr_type nft_ng_type;
192  static const struct nft_expr_ops nft_ng_inc_ops = {
193  	.type		= &nft_ng_type,
194  	.size		= NFT_EXPR_SIZE(sizeof(struct nft_ng_inc)),
195  	.eval		= nft_ng_inc_eval,
196  	.init		= nft_ng_inc_init,
197  	.destroy	= nft_ng_inc_destroy,
198  	.dump		= nft_ng_inc_dump,
199  	.reduce		= nft_ng_inc_reduce,
200  };
201  
202  static const struct nft_expr_ops nft_ng_random_ops = {
203  	.type		= &nft_ng_type,
204  	.size		= NFT_EXPR_SIZE(sizeof(struct nft_ng_random)),
205  	.eval		= nft_ng_random_eval,
206  	.init		= nft_ng_random_init,
207  	.dump		= nft_ng_random_dump,
208  	.reduce		= nft_ng_random_reduce,
209  };
210  
211  static const struct nft_expr_ops *
nft_ng_select_ops(const struct nft_ctx * ctx,const struct nlattr * const tb[])212  nft_ng_select_ops(const struct nft_ctx *ctx, const struct nlattr * const tb[])
213  {
214  	u32 type;
215  
216  	if (!tb[NFTA_NG_DREG]	 ||
217  	    !tb[NFTA_NG_MODULUS] ||
218  	    !tb[NFTA_NG_TYPE])
219  		return ERR_PTR(-EINVAL);
220  
221  	type = ntohl(nla_get_be32(tb[NFTA_NG_TYPE]));
222  
223  	switch (type) {
224  	case NFT_NG_INCREMENTAL:
225  		return &nft_ng_inc_ops;
226  	case NFT_NG_RANDOM:
227  		return &nft_ng_random_ops;
228  	}
229  
230  	return ERR_PTR(-EINVAL);
231  }
232  
233  static struct nft_expr_type nft_ng_type __read_mostly = {
234  	.name		= "numgen",
235  	.select_ops	= nft_ng_select_ops,
236  	.policy		= nft_ng_policy,
237  	.maxattr	= NFTA_NG_MAX,
238  	.owner		= THIS_MODULE,
239  };
240  
nft_ng_module_init(void)241  static int __init nft_ng_module_init(void)
242  {
243  	return nft_register_expr(&nft_ng_type);
244  }
245  
nft_ng_module_exit(void)246  static void __exit nft_ng_module_exit(void)
247  {
248  	nft_unregister_expr(&nft_ng_type);
249  }
250  
251  module_init(nft_ng_module_init);
252  module_exit(nft_ng_module_exit);
253  
254  MODULE_LICENSE("GPL");
255  MODULE_AUTHOR("Laura Garcia <nevola@gmail.com>");
256  MODULE_ALIAS_NFT_EXPR("numgen");
257  MODULE_DESCRIPTION("nftables number generator module");
258