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 }; 27 28 static void nft_ng_inc_eval(const struct nft_expr *expr, 29 struct nft_regs *regs, 30 const struct nft_pktinfo *pkt) 31 { 32 struct nft_ng_inc *priv = nft_expr_priv(expr); 33 u32 nval, oval; 34 35 do { 36 oval = atomic_read(&priv->counter); 37 nval = (oval + 1 < priv->modulus) ? oval + 1 : 0; 38 } while (atomic_cmpxchg(&priv->counter, oval, nval) != oval); 39 40 memcpy(®s->data[priv->dreg], &priv->counter, sizeof(u32)); 41 } 42 43 static const struct nla_policy nft_ng_policy[NFTA_NG_MAX + 1] = { 44 [NFTA_NG_DREG] = { .type = NLA_U32 }, 45 [NFTA_NG_MODULUS] = { .type = NLA_U32 }, 46 [NFTA_NG_TYPE] = { .type = NLA_U32 }, 47 }; 48 49 static int nft_ng_inc_init(const struct nft_ctx *ctx, 50 const struct nft_expr *expr, 51 const struct nlattr * const tb[]) 52 { 53 struct nft_ng_inc *priv = nft_expr_priv(expr); 54 55 priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS])); 56 if (priv->modulus == 0) 57 return -ERANGE; 58 59 priv->dreg = nft_parse_register(tb[NFTA_NG_DREG]); 60 atomic_set(&priv->counter, 0); 61 62 return nft_validate_register_store(ctx, priv->dreg, NULL, 63 NFT_DATA_VALUE, sizeof(u32)); 64 } 65 66 static int nft_ng_dump(struct sk_buff *skb, enum nft_registers dreg, 67 u32 modulus, enum nft_ng_types type) 68 { 69 if (nft_dump_register(skb, NFTA_NG_DREG, dreg)) 70 goto nla_put_failure; 71 if (nla_put_be32(skb, NFTA_NG_MODULUS, htonl(modulus))) 72 goto nla_put_failure; 73 if (nla_put_be32(skb, NFTA_NG_TYPE, htonl(type))) 74 goto nla_put_failure; 75 76 return 0; 77 78 nla_put_failure: 79 return -1; 80 } 81 82 static int nft_ng_inc_dump(struct sk_buff *skb, const struct nft_expr *expr) 83 { 84 const struct nft_ng_inc *priv = nft_expr_priv(expr); 85 86 return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_INCREMENTAL); 87 } 88 89 struct nft_ng_random { 90 enum nft_registers dreg:8; 91 u32 modulus; 92 }; 93 94 static void nft_ng_random_eval(const struct nft_expr *expr, 95 struct nft_regs *regs, 96 const struct nft_pktinfo *pkt) 97 { 98 struct nft_ng_random *priv = nft_expr_priv(expr); 99 struct rnd_state *state = this_cpu_ptr(&nft_numgen_prandom_state); 100 101 regs->data[priv->dreg] = reciprocal_scale(prandom_u32_state(state), 102 priv->modulus); 103 } 104 105 static int nft_ng_random_init(const struct nft_ctx *ctx, 106 const struct nft_expr *expr, 107 const struct nlattr * const tb[]) 108 { 109 struct nft_ng_random *priv = nft_expr_priv(expr); 110 111 priv->modulus = ntohl(nla_get_be32(tb[NFTA_NG_MODULUS])); 112 if (priv->modulus == 0) 113 return -ERANGE; 114 115 prandom_init_once(&nft_numgen_prandom_state); 116 117 priv->dreg = nft_parse_register(tb[NFTA_NG_DREG]); 118 119 return nft_validate_register_store(ctx, priv->dreg, NULL, 120 NFT_DATA_VALUE, sizeof(u32)); 121 } 122 123 static int nft_ng_random_dump(struct sk_buff *skb, const struct nft_expr *expr) 124 { 125 const struct nft_ng_random *priv = nft_expr_priv(expr); 126 127 return nft_ng_dump(skb, priv->dreg, priv->modulus, NFT_NG_RANDOM); 128 } 129 130 static struct nft_expr_type nft_ng_type; 131 static const struct nft_expr_ops nft_ng_inc_ops = { 132 .type = &nft_ng_type, 133 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_inc)), 134 .eval = nft_ng_inc_eval, 135 .init = nft_ng_inc_init, 136 .dump = nft_ng_inc_dump, 137 }; 138 139 static const struct nft_expr_ops nft_ng_random_ops = { 140 .type = &nft_ng_type, 141 .size = NFT_EXPR_SIZE(sizeof(struct nft_ng_random)), 142 .eval = nft_ng_random_eval, 143 .init = nft_ng_random_init, 144 .dump = nft_ng_random_dump, 145 }; 146 147 static const struct nft_expr_ops * 148 nft_ng_select_ops(const struct nft_ctx *ctx, const struct nlattr * const tb[]) 149 { 150 u32 type; 151 152 if (!tb[NFTA_NG_DREG] || 153 !tb[NFTA_NG_MODULUS] || 154 !tb[NFTA_NG_TYPE]) 155 return ERR_PTR(-EINVAL); 156 157 type = ntohl(nla_get_be32(tb[NFTA_NG_TYPE])); 158 159 switch (type) { 160 case NFT_NG_INCREMENTAL: 161 return &nft_ng_inc_ops; 162 case NFT_NG_RANDOM: 163 return &nft_ng_random_ops; 164 } 165 166 return ERR_PTR(-EINVAL); 167 } 168 169 static struct nft_expr_type nft_ng_type __read_mostly = { 170 .name = "numgen", 171 .select_ops = &nft_ng_select_ops, 172 .policy = nft_ng_policy, 173 .maxattr = NFTA_NG_MAX, 174 .owner = THIS_MODULE, 175 }; 176 177 static int __init nft_ng_module_init(void) 178 { 179 return nft_register_expr(&nft_ng_type); 180 } 181 182 static void __exit nft_ng_module_exit(void) 183 { 184 nft_unregister_expr(&nft_ng_type); 185 } 186 187 module_init(nft_ng_module_init); 188 module_exit(nft_ng_module_exit); 189 190 MODULE_LICENSE("GPL"); 191 MODULE_AUTHOR("Laura Garcia <nevola@gmail.com>"); 192 MODULE_ALIAS_NFT_EXPR("numgen"); 193