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->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 __read_mostly = { 155 .type = NFT_OBJECT_QUOTA, 156 .size = sizeof(struct nft_quota), 157 .maxattr = NFTA_QUOTA_MAX, 158 .policy = nft_quota_policy, 159 .init = nft_quota_obj_init, 160 .eval = nft_quota_obj_eval, 161 .dump = nft_quota_obj_dump, 162 .owner = THIS_MODULE, 163 }; 164 165 static void nft_quota_eval(const struct nft_expr *expr, 166 struct nft_regs *regs, 167 const struct nft_pktinfo *pkt) 168 { 169 struct nft_quota *priv = nft_expr_priv(expr); 170 171 nft_quota_do_eval(priv, regs, pkt); 172 } 173 174 static int nft_quota_init(const struct nft_ctx *ctx, 175 const struct nft_expr *expr, 176 const struct nlattr * const tb[]) 177 { 178 struct nft_quota *priv = nft_expr_priv(expr); 179 180 return nft_quota_do_init(tb, priv); 181 } 182 183 static int nft_quota_dump(struct sk_buff *skb, const struct nft_expr *expr) 184 { 185 struct nft_quota *priv = nft_expr_priv(expr); 186 187 return nft_quota_do_dump(skb, priv, false); 188 } 189 190 static struct nft_expr_type nft_quota_type; 191 static const struct nft_expr_ops nft_quota_ops = { 192 .type = &nft_quota_type, 193 .size = NFT_EXPR_SIZE(sizeof(struct nft_quota)), 194 .eval = nft_quota_eval, 195 .init = nft_quota_init, 196 .dump = nft_quota_dump, 197 }; 198 199 static struct nft_expr_type nft_quota_type __read_mostly = { 200 .name = "quota", 201 .ops = &nft_quota_ops, 202 .policy = nft_quota_policy, 203 .maxattr = NFTA_QUOTA_MAX, 204 .flags = NFT_EXPR_STATEFUL, 205 .owner = THIS_MODULE, 206 }; 207 208 static int __init nft_quota_module_init(void) 209 { 210 int err; 211 212 err = nft_register_obj(&nft_quota_obj); 213 if (err < 0) 214 return err; 215 216 err = nft_register_expr(&nft_quota_type); 217 if (err < 0) 218 goto err1; 219 220 return 0; 221 err1: 222 nft_unregister_obj(&nft_quota_obj); 223 return err; 224 } 225 226 static void __exit nft_quota_module_exit(void) 227 { 228 nft_unregister_expr(&nft_quota_type); 229 nft_unregister_obj(&nft_quota_obj); 230 } 231 232 module_init(nft_quota_module_init); 233 module_exit(nft_quota_module_exit); 234 235 MODULE_LICENSE("GPL"); 236 MODULE_AUTHOR("Pablo Neira Ayuso <pablo@netfilter.org>"); 237 MODULE_ALIAS_NFT_EXPR("quota"); 238 MODULE_ALIAS_NFT_OBJ(NFT_OBJECT_QUOTA); 239