1 /* 2 * Copyright (c) 2016 Laura Garcia <nevola@gmail.com> 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 */ 9 10 #include <linux/kernel.h> 11 #include <linux/init.h> 12 #include <linux/module.h> 13 #include <linux/netlink.h> 14 #include <linux/netfilter.h> 15 #include <linux/netfilter/nf_tables.h> 16 #include <linux/static_key.h> 17 #include <net/netfilter/nf_tables.h> 18 #include <net/netfilter/nf_tables_core.h> 19 20 static DEFINE_PER_CPU(struct rnd_state, nft_numgen_prandom_state); 21 22 struct nft_ng_inc { 23 enum nft_registers dreg:8; 24 u32 modulus; 25 atomic_t counter; 26 u32 offset; 27 struct nft_set *map; 28 }; 29 30 static u32 nft_ng_inc_gen(struct nft_ng_inc *priv) 31 { 32 u32 nval, oval; 33 34 do { 35 oval = atomic_read(&priv->counter); 36 nval = (oval + 1 < priv->modulus) ? oval + 1 : 0; 37 } while (atomic_cmpxchg(&priv->counter, oval, nval) != oval); 38 39 return nval + priv->offset; 40 } 41 42 static void nft_ng_inc_eval(const struct nft_expr *expr, 43 struct nft_regs *regs, 44 const struct nft_pktinfo *pkt) 45 { 46 struct nft_ng_inc *priv = nft_expr_priv(expr); 47 48 regs->data[priv->dreg] = nft_ng_inc_gen(priv); 49 } 50 51 static void nft_ng_inc_map_eval(const struct nft_expr *expr, 52 struct nft_regs *regs, 53 const struct nft_pktinfo *pkt) 54 { 55 struct nft_ng_inc *priv = nft_expr_priv(expr); 56 const struct nft_set *map = priv->map; 57 const struct nft_set_ext *ext; 58 u32 result; 59 bool found; 60 61 result = nft_ng_inc_gen(priv); 62 found = map->ops->lookup(nft_net(pkt), map, &result, &ext); 63 64 if (!found) 65 return; 66 67 nft_data_copy(®s->data[priv->dreg], 68 nft_set_ext_data(ext), map->dlen); 69 } 70 71 static const struct nla_policy nft_ng_policy[NFTA_NG_MAX + 1] = { 72 [NFTA_NG_DREG] = { .type = NLA_U32 }, 73 [NFTA_NG_MODULUS] = { .type = NLA_U32 }, 74 [NFTA_NG_TYPE] = { .type = NLA_U32 }, 75 [NFTA_NG_OFFSET] = { .type = NLA_U32 }, 76 [NFTA_NG_SET_NAME] = { .type = NLA_STRING, 77 .len = NFT_SET_MAXNAMELEN - 1 }, 78 [NFTA_NG_SET_ID] = { .type = NLA_U32 }, 79 }; 80 81 static int nft_ng_inc_init(const struct nft_ctx *ctx, 82 const struct nft_expr *expr, 83 const struct nlattr * const tb[]) 84 { 85 struct nft_ng_inc *priv = nft_expr_priv(expr); 86 87 if (tb[NFTA_NG_OFFSET]) 88 priv->offset = ntohl(nla_get_be32(tb[NFTA_NG_OFFSET])); 89 90 priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS])); 91 if (priv->modulus == 0) 92 return -ERANGE; 93 94 if (priv->offset + priv->modulus - 1 < priv->offset) 95 return -EOVERFLOW; 96 97 priv->dreg = nft_parse_register(tb[NFTA_NG_DREG]); 98 atomic_set(&priv->counter, priv->modulus - 1); 99 100 return nft_validate_register_store(ctx, priv->dreg, NULL, 101 NFT_DATA_VALUE, sizeof(u32)); 102 } 103 104 static int nft_ng_inc_map_init(const struct nft_ctx *ctx, 105 const struct nft_expr *expr, 106 const struct nlattr * const tb[]) 107 { 108 struct nft_ng_inc *priv = nft_expr_priv(expr); 109 u8 genmask = nft_genmask_next(ctx->net); 110 111 nft_ng_inc_init(ctx, expr, tb); 112 113 priv->map = nft_set_lookup_global(ctx->net, ctx->table, 114 tb[NFTA_NG_SET_NAME], 115 tb[NFTA_NG_SET_ID], genmask); 116 117 return PTR_ERR_OR_ZERO(priv->map); 118 } 119 120 static int nft_ng_dump(struct sk_buff *skb, enum nft_registers dreg, 121 u32 modulus, enum nft_ng_types type, u32 offset) 122 { 123 if (nft_dump_register(skb, NFTA_NG_DREG, dreg)) 124 goto nla_put_failure; 125 if (nla_put_be32(skb, NFTA_NG_MODULUS, htonl(modulus))) 126 goto nla_put_failure; 127 if (nla_put_be32(skb, NFTA_NG_TYPE, htonl(type))) 128 goto nla_put_failure; 129 if (nla_put_be32(skb, NFTA_NG_OFFSET, htonl(offset))) 130 goto nla_put_failure; 131 132 return 0; 133 134 nla_put_failure: 135 return -1; 136 } 137 138 static int nft_ng_inc_dump(struct sk_buff *skb, const struct nft_expr *expr) 139 { 140 const struct nft_ng_inc *priv = nft_expr_priv(expr); 141 142 return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_INCREMENTAL, 143 priv->offset); 144 } 145 146 static int nft_ng_inc_map_dump(struct sk_buff *skb, 147 const struct nft_expr *expr) 148 { 149 const struct nft_ng_inc *priv = nft_expr_priv(expr); 150 151 if (nft_ng_dump(skb, priv->dreg, priv->modulus, 152 NFT_NG_INCREMENTAL, priv->offset) || 153 nla_put_string(skb, NFTA_NG_SET_NAME, priv->map->name)) 154 goto nla_put_failure; 155 156 return 0; 157 158 nla_put_failure: 159 return -1; 160 } 161 162 struct nft_ng_random { 163 enum nft_registers dreg:8; 164 u32 modulus; 165 u32 offset; 166 struct nft_set *map; 167 }; 168 169 static u32 nft_ng_random_gen(struct nft_ng_random *priv) 170 { 171 struct rnd_state *state = this_cpu_ptr(&nft_numgen_prandom_state); 172 173 return reciprocal_scale(prandom_u32_state(state), priv->modulus) + 174 priv->offset; 175 } 176 177 static void nft_ng_random_eval(const struct nft_expr *expr, 178 struct nft_regs *regs, 179 const struct nft_pktinfo *pkt) 180 { 181 struct nft_ng_random *priv = nft_expr_priv(expr); 182 183 regs->data[priv->dreg] = nft_ng_random_gen(priv); 184 } 185 186 static void nft_ng_random_map_eval(const struct nft_expr *expr, 187 struct nft_regs *regs, 188 const struct nft_pktinfo *pkt) 189 { 190 struct nft_ng_random *priv = nft_expr_priv(expr); 191 const struct nft_set *map = priv->map; 192 const struct nft_set_ext *ext; 193 u32 result; 194 bool found; 195 196 result = nft_ng_random_gen(priv); 197 found = map->ops->lookup(nft_net(pkt), map, &result, &ext); 198 if (!found) 199 return; 200 201 nft_data_copy(®s->data[priv->dreg], 202 nft_set_ext_data(ext), map->dlen); 203 } 204 205 static int nft_ng_random_init(const struct nft_ctx *ctx, 206 const struct nft_expr *expr, 207 const struct nlattr * const tb[]) 208 { 209 struct nft_ng_random *priv = nft_expr_priv(expr); 210 211 if (tb[NFTA_NG_OFFSET]) 212 priv->offset = ntohl(nla_get_be32(tb[NFTA_NG_OFFSET])); 213 214 priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS])); 215 if (priv->modulus == 0) 216 return -ERANGE; 217 218 if (priv->offset + priv->modulus - 1 < priv->offset) 219 return -EOVERFLOW; 220 221 prandom_init_once(&nft_numgen_prandom_state); 222 223 priv->dreg = nft_parse_register(tb[NFTA_NG_DREG]); 224 225 return nft_validate_register_store(ctx, priv->dreg, NULL, 226 NFT_DATA_VALUE, sizeof(u32)); 227 } 228 229 static int nft_ng_random_map_init(const struct nft_ctx *ctx, 230 const struct nft_expr *expr, 231 const struct nlattr * const tb[]) 232 { 233 struct nft_ng_random *priv = nft_expr_priv(expr); 234 u8 genmask = nft_genmask_next(ctx->net); 235 236 nft_ng_random_init(ctx, expr, tb); 237 priv->map = nft_set_lookup_global(ctx->net, ctx->table, 238 tb[NFTA_NG_SET_NAME], 239 tb[NFTA_NG_SET_ID], genmask); 240 if (IS_ERR(priv->map)) 241 return PTR_ERR(priv->map); 242 243 return 0; 244 } 245 246 static int nft_ng_random_dump(struct sk_buff *skb, const struct nft_expr *expr) 247 { 248 const struct nft_ng_random *priv = nft_expr_priv(expr); 249 250 return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_RANDOM, 251 priv->offset); 252 } 253 254 static int nft_ng_random_map_dump(struct sk_buff *skb, 255 const struct nft_expr *expr) 256 { 257 const struct nft_ng_random *priv = nft_expr_priv(expr); 258 259 if (nft_ng_dump(skb, priv->dreg, priv->modulus, 260 NFT_NG_RANDOM, priv->offset) || 261 nla_put_string(skb, NFTA_NG_SET_NAME, priv->map->name)) 262 goto nla_put_failure; 263 264 return 0; 265 266 nla_put_failure: 267 return -1; 268 } 269 270 static struct nft_expr_type nft_ng_type; 271 static const struct nft_expr_ops nft_ng_inc_ops = { 272 .type = &nft_ng_type, 273 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_inc)), 274 .eval = nft_ng_inc_eval, 275 .init = nft_ng_inc_init, 276 .dump = nft_ng_inc_dump, 277 }; 278 279 static const struct nft_expr_ops nft_ng_inc_map_ops = { 280 .type = &nft_ng_type, 281 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_inc)), 282 .eval = nft_ng_inc_map_eval, 283 .init = nft_ng_inc_map_init, 284 .dump = nft_ng_inc_map_dump, 285 }; 286 287 static const struct nft_expr_ops nft_ng_random_ops = { 288 .type = &nft_ng_type, 289 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_random)), 290 .eval = nft_ng_random_eval, 291 .init = nft_ng_random_init, 292 .dump = nft_ng_random_dump, 293 }; 294 295 static const struct nft_expr_ops nft_ng_random_map_ops = { 296 .type = &nft_ng_type, 297 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_random)), 298 .eval = nft_ng_random_map_eval, 299 .init = nft_ng_random_map_init, 300 .dump = nft_ng_random_map_dump, 301 }; 302 303 static const struct nft_expr_ops * 304 nft_ng_select_ops(const struct nft_ctx *ctx, const struct nlattr * const tb[]) 305 { 306 u32 type; 307 308 if (!tb[NFTA_NG_DREG] || 309 !tb[NFTA_NG_MODULUS] || 310 !tb[NFTA_NG_TYPE]) 311 return ERR_PTR(-EINVAL); 312 313 type = ntohl(nla_get_be32(tb[NFTA_NG_TYPE])); 314 315 switch (type) { 316 case NFT_NG_INCREMENTAL: 317 if (tb[NFTA_NG_SET_NAME]) 318 return &nft_ng_inc_map_ops; 319 return &nft_ng_inc_ops; 320 case NFT_NG_RANDOM: 321 if (tb[NFTA_NG_SET_NAME]) 322 return &nft_ng_random_map_ops; 323 return &nft_ng_random_ops; 324 } 325 326 return ERR_PTR(-EINVAL); 327 } 328 329 static struct nft_expr_type nft_ng_type __read_mostly = { 330 .name = "numgen", 331 .select_ops = nft_ng_select_ops, 332 .policy = nft_ng_policy, 333 .maxattr = NFTA_NG_MAX, 334 .owner = THIS_MODULE, 335 }; 336 337 static int __init nft_ng_module_init(void) 338 { 339 return nft_register_expr(&nft_ng_type); 340 } 341 342 static void __exit nft_ng_module_exit(void) 343 { 344 nft_unregister_expr(&nft_ng_type); 345 } 346 347 module_init(nft_ng_module_init); 348 module_exit(nft_ng_module_exit); 349 350 MODULE_LICENSE("GPL"); 351 MODULE_AUTHOR("Laura Garcia <nevola@gmail.com>"); 352 MODULE_ALIAS_NFT_EXPR("numgen"); 353