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 | |
24 | static u32 nft_ng_inc_gen(struct nft_ng_inc *priv) |
25 | { |
26 | u32 nval, oval; |
27 | |
28 | do { |
29 | oval = atomic_read(v: priv->counter); |
30 | nval = (oval + 1 < priv->modulus) ? oval + 1 : 0; |
31 | } while (atomic_cmpxchg(v: priv->counter, old: oval, new: nval) != oval); |
32 | |
33 | return nval + priv->offset; |
34 | } |
35 | |
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 | |
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(size: sizeof(*priv->counter), GFP_KERNEL); |
70 | if (!priv->counter) |
71 | return -ENOMEM; |
72 | |
73 | atomic_set(v: priv->counter, i: priv->modulus - 1); |
74 | |
75 | err = nft_parse_register_store(ctx, attr: tb[NFTA_NG_DREG], dreg: &priv->dreg, |
76 | NULL, type: NFT_DATA_VALUE, len: sizeof(u32)); |
77 | if (err < 0) |
78 | goto err; |
79 | |
80 | return 0; |
81 | err: |
82 | kfree(objp: priv->counter); |
83 | |
84 | return err; |
85 | } |
86 | |
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, dreg: priv->dreg, NFT_REG32_SIZE); |
93 | |
94 | return false; |
95 | } |
96 | |
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, attr: NFTA_NG_DREG, reg: dreg)) |
101 | goto nla_put_failure; |
102 | if (nla_put_be32(skb, attrtype: NFTA_NG_MODULUS, htonl(modulus))) |
103 | goto nla_put_failure; |
104 | if (nla_put_be32(skb, attrtype: NFTA_NG_TYPE, htonl(type))) |
105 | goto nla_put_failure; |
106 | if (nla_put_be32(skb, attrtype: NFTA_NG_OFFSET, htonl(offset))) |
107 | goto nla_put_failure; |
108 | |
109 | return 0; |
110 | |
111 | nla_put_failure: |
112 | return -1; |
113 | } |
114 | |
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, dreg: priv->dreg, modulus: priv->modulus, type: NFT_NG_INCREMENTAL, |
121 | offset: priv->offset); |
122 | } |
123 | |
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(objp: priv->counter); |
130 | } |
131 | |
132 | struct nft_ng_random { |
133 | u8 dreg; |
134 | u32 modulus; |
135 | u32 offset; |
136 | }; |
137 | |
138 | static u32 nft_ng_random_gen(const struct nft_ng_random *priv) |
139 | { |
140 | return reciprocal_scale(val: get_random_u32(), ep_ro: priv->modulus) + priv->offset; |
141 | } |
142 | |
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 | |
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, attr: tb[NFTA_NG_DREG], dreg: &priv->dreg, |
169 | NULL, type: NFT_DATA_VALUE, len: sizeof(u32)); |
170 | } |
171 | |
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, dreg: priv->dreg, modulus: priv->modulus, type: NFT_NG_RANDOM, |
178 | offset: priv->offset); |
179 | } |
180 | |
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, dreg: 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 * |
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(error: -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(error: -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 | |
241 | static int __init nft_ng_module_init(void) |
242 | { |
243 | return nft_register_expr(&nft_ng_type); |
244 | } |
245 | |
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 | |