1 /* 2 * Copyright (c) 2016 Pablo Neira Ayuso <pablo@netfilter.org> 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 #include <linux/kernel.h> 10 #include <linux/init.h> 11 #include <linux/module.h> 12 #include <linux/atomic.h> 13 #include <linux/netlink.h> 14 #include <linux/netfilter.h> 15 #include <linux/netfilter/nf_tables.h> 16 #include <net/netfilter/nf_tables.h> 17 18 struct nft_quota { 19 u64 quota; 20 unsigned long flags; 21 atomic64_t consumed; 22 }; 23 24 static inline bool nft_overquota(struct nft_quota *priv, 25 const struct sk_buff *skb) 26 { 27 return atomic64_add_return(skb->len, &priv->consumed) >= priv->quota; 28 } 29 30 static inline bool nft_quota_invert(struct nft_quota *priv) 31 { 32 return priv->flags & NFT_QUOTA_F_INV; 33 } 34 35 static inline void nft_quota_do_eval(struct nft_quota *priv, 36 struct nft_regs *regs, 37 const struct nft_pktinfo *pkt) 38 { 39 if (nft_overquota(priv, pkt->skb) ^ nft_quota_invert(priv)) 40 regs->verdict.code = NFT_BREAK; 41 } 42 43 static const struct nla_policy nft_quota_policy[NFTA_QUOTA_MAX + 1] = { 44 [NFTA_QUOTA_BYTES] = { .type = NLA_U64 }, 45 [NFTA_QUOTA_FLAGS] = { .type = NLA_U32 }, 46 [NFTA_QUOTA_CONSUMED] = { .type = NLA_U64 }, 47 }; 48 49 #define NFT_QUOTA_DEPLETED_BIT 1 /* From NFT_QUOTA_F_DEPLETED. */ 50 51 static void nft_quota_obj_eval(struct nft_object *obj, 52 struct nft_regs *regs, 53 const struct nft_pktinfo *pkt) 54 { 55 struct nft_quota *priv = nft_obj_data(obj); 56 bool overquota; 57 58 overquota = nft_overquota(priv, pkt->skb); 59 if (overquota ^ nft_quota_invert(priv)) 60 regs->verdict.code = NFT_BREAK; 61 62 if (overquota && 63 !test_and_set_bit(NFT_QUOTA_DEPLETED_BIT, &priv->flags)) 64 nft_obj_notify(nft_net(pkt), obj->key.table, obj, 0, 0, 65 NFT_MSG_NEWOBJ, nft_pf(pkt), 0, GFP_ATOMIC); 66 } 67 68 static int nft_quota_do_init(const struct nlattr * const tb[], 69 struct nft_quota *priv) 70 { 71 unsigned long flags = 0; 72 u64 quota, consumed = 0; 73 74 if (!tb[NFTA_QUOTA_BYTES]) 75 return -EINVAL; 76 77 quota = be64_to_cpu(nla_get_be64(tb[NFTA_QUOTA_BYTES])); 78 if (quota > S64_MAX) 79 return -EOVERFLOW; 80 81 if (tb[NFTA_QUOTA_CONSUMED]) { 82 consumed = be64_to_cpu(nla_get_be64(tb[NFTA_QUOTA_CONSUMED])); 83 if (consumed > quota) 84 return -EINVAL; 85 } 86 87 if (tb[NFTA_QUOTA_FLAGS]) { 88 flags = ntohl(nla_get_be32(tb[NFTA_QUOTA_FLAGS])); 89 if (flags & ~NFT_QUOTA_F_INV) 90 return -EINVAL; 91 if (flags & NFT_QUOTA_F_DEPLETED) 92 return -EOPNOTSUPP; 93 } 94 95 priv->quota = quota; 96 priv->flags = flags; 97 atomic64_set(&priv->consumed, consumed); 98 99 return 0; 100 } 101 102 static int nft_quota_obj_init(const struct nft_ctx *ctx, 103 const struct nlattr * const tb[], 104 struct nft_object *obj) 105 { 106 struct nft_quota *priv = nft_obj_data(obj); 107 108 return nft_quota_do_init(tb, priv); 109 } 110 111 static int nft_quota_do_dump(struct sk_buff *skb, struct nft_quota *priv, 112 bool reset) 113 { 114 u64 consumed, consumed_cap; 115 u32 flags = priv->flags; 116 117 /* Since we inconditionally increment consumed quota for each packet 118 * that we see, don't go over the quota boundary in what we send to 119 * userspace. 120 */ 121 consumed = atomic64_read(&priv->consumed); 122 if (consumed >= priv->quota) { 123 consumed_cap = priv->quota; 124 flags |= NFT_QUOTA_F_DEPLETED; 125 } else { 126 consumed_cap = consumed; 127 } 128 129 if (nla_put_be64(skb, NFTA_QUOTA_BYTES, cpu_to_be64(priv->quota), 130 NFTA_QUOTA_PAD) || 131 nla_put_be64(skb, NFTA_QUOTA_CONSUMED, cpu_to_be64(consumed_cap), 132 NFTA_QUOTA_PAD) || 133 nla_put_be32(skb, NFTA_QUOTA_FLAGS, htonl(flags))) 134 goto nla_put_failure; 135 136 if (reset) { 137 atomic64_sub(consumed, &priv->consumed); 138 clear_bit(NFT_QUOTA_DEPLETED_BIT, &priv->flags); 139 } 140 return 0; 141 142 nla_put_failure: 143 return -1; 144 } 145 146 static int nft_quota_obj_dump(struct sk_buff *skb, struct nft_object *obj, 147 bool reset) 148 { 149 struct nft_quota *priv = nft_obj_data(obj); 150 151 return nft_quota_do_dump(skb, priv, reset); 152 } 153 154 static struct nft_object_type nft_quota_obj_type; 155 static const struct nft_object_ops nft_quota_obj_ops = { 156 .type = &nft_quota_obj_type, 157 .size = sizeof(struct nft_quota), 158 .init = nft_quota_obj_init, 159 .eval = nft_quota_obj_eval, 160 .dump = nft_quota_obj_dump, 161 }; 162 163 static struct nft_object_type nft_quota_obj_type __read_mostly = { 164 .type = NFT_OBJECT_QUOTA, 165 .ops = &nft_quota_obj_ops, 166 .maxattr = NFTA_QUOTA_MAX, 167 .policy = nft_quota_policy, 168 .owner = THIS_MODULE, 169 }; 170 171 static void nft_quota_eval(const struct nft_expr *expr, 172 struct nft_regs *regs, 173 const struct nft_pktinfo *pkt) 174 { 175 struct nft_quota *priv = nft_expr_priv(expr); 176 177 nft_quota_do_eval(priv, regs, pkt); 178 } 179 180 static int nft_quota_init(const struct nft_ctx *ctx, 181 const struct nft_expr *expr, 182 const struct nlattr * const tb[]) 183 { 184 struct nft_quota *priv = nft_expr_priv(expr); 185 186 return nft_quota_do_init(tb, priv); 187 } 188 189 static int nft_quota_dump(struct sk_buff *skb, const struct nft_expr *expr) 190 { 191 struct nft_quota *priv = nft_expr_priv(expr); 192 193 return nft_quota_do_dump(skb, priv, false); 194 } 195 196 static struct nft_expr_type nft_quota_type; 197 static const struct nft_expr_ops nft_quota_ops = { 198 .type = &nft_quota_type, 199 .size = NFT_EXPR_SIZE(sizeof(struct nft_quota)), 200 .eval = nft_quota_eval, 201 .init = nft_quota_init, 202 .dump = nft_quota_dump, 203 }; 204 205 static struct nft_expr_type nft_quota_type __read_mostly = { 206 .name = "quota", 207 .ops = &nft_quota_ops, 208 .policy = nft_quota_policy, 209 .maxattr = NFTA_QUOTA_MAX, 210 .flags = NFT_EXPR_STATEFUL, 211 .owner = THIS_MODULE, 212 }; 213 214 static int __init nft_quota_module_init(void) 215 { 216 int err; 217 218 err = nft_register_obj(&nft_quota_obj_type); 219 if (err < 0) 220 return err; 221 222 err = nft_register_expr(&nft_quota_type); 223 if (err < 0) 224 goto err1; 225 226 return 0; 227 err1: 228 nft_unregister_obj(&nft_quota_obj_type); 229 return err; 230 } 231 232 static void __exit nft_quota_module_exit(void) 233 { 234 nft_unregister_expr(&nft_quota_type); 235 nft_unregister_obj(&nft_quota_obj_type); 236 } 237 238 module_init(nft_quota_module_init); 239 module_exit(nft_quota_module_exit); 240 241 MODULE_LICENSE("GPL"); 242 MODULE_AUTHOR("Pablo Neira Ayuso <pablo@netfilter.org>"); 243 MODULE_ALIAS_NFT_EXPR("quota"); 244 MODULE_ALIAS_NFT_OBJ(NFT_OBJECT_QUOTA); 245