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