net/core/sock_map.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2017 - 2018 Covalent IO, Inc. http://covalent.io */
   3
   4 #include <linux/bpf.h>
   5 #include <linux/filter.h>
   6 #include <linux/errno.h>
   7 #include <linux/file.h>
   8 #include <linux/net.h>
   9 #include <linux/workqueue.h>
  10 #include <linux/skmsg.h>
  11 #include <linux/list.h>
  12 #include <linux/jhash.h>
  13
  14 struct bpf_stab {
  15         struct bpf_map map;
  16         struct sock **sks;
  17         struct sk_psock_progs progs;
  18         raw_spinlock_t lock;
  19 };
  20
  21 #define SOCK_CREATE_FLAG_MASK                           \
  22         (BPF_F_NUMA_NODE | BPF_F_RDONLY | BPF_F_WRONLY)
  23
  24 static struct bpf_map *sock_map_alloc(union bpf_attr *attr)
  25 {
  26         struct bpf_stab *stab;
  27         u64 cost;
  28         int err;
  29
  30         if (!capable(CAP_NET_ADMIN))
  31                 return ERR_PTR(-EPERM);
  32         if (attr->max_entries == 0 ||
  33             attr->key_size    != 4 ||
  34             attr->value_size  != 4 ||
  35             attr->map_flags & ~SOCK_CREATE_FLAG_MASK)
  36                 return ERR_PTR(-EINVAL);
  37
  38         stab = kzalloc(sizeof(*stab), GFP_USER);
  39         if (!stab)
  40                 return ERR_PTR(-ENOMEM);
  41
  42         bpf_map_init_from_attr(&stab->map, attr);
  43         raw_spin_lock_init(&stab->lock);
  44
  45         /* Make sure page count doesn't overflow. */
  46         cost = (u64) stab->map.max_entries * sizeof(struct sock *);
  47         if (cost >= U32_MAX - PAGE_SIZE) {
  48                 err = -EINVAL;
  49                 goto free_stab;
  50         }
  51
  52         err = bpf_map_charge_init(&stab->map.memory,
  53                                   round_up(cost, PAGE_SIZE) >> PAGE_SHIFT);
  54         if (err)
  55                 goto free_stab;
  56
  57         stab->sks = bpf_map_area_alloc(stab->map.max_entries *
  58                                        sizeof(struct sock *),
  59                                        stab->map.numa_node);
  60         if (stab->sks)
  61                 return &stab->map;
  62         err = -ENOMEM;
  63         bpf_map_charge_finish(&stab->map.memory);
  64 free_stab:
  65         kfree(stab);
  66         return ERR_PTR(err);
  67 }
  68
  69 int sock_map_get_from_fd(const union bpf_attr *attr, struct bpf_prog *prog)
  70 {
  71         u32 ufd = attr->target_fd;
  72         struct bpf_map *map;
  73         struct fd f;
  74         int ret;
  75
  76         f = fdget(ufd);
  77         map = __bpf_map_get(f);
  78         if (IS_ERR(map))
  79                 return PTR_ERR(map);
  80         ret = sock_map_prog_update(map, prog, attr->attach_type);
  81         fdput(f);
  82         return ret;
  83 }
  84
  85 static void sock_map_sk_acquire(struct sock *sk)
  86         __acquires(&sk->sk_lock.slock)
  87 {
  88         lock_sock(sk);
  89         preempt_disable();
  90         rcu_read_lock();
  91 }
  92
  93 static void sock_map_sk_release(struct sock *sk)
  94         __releases(&sk->sk_lock.slock)
  95 {
  96         rcu_read_unlock();
  97         preempt_enable();
  98         release_sock(sk);
  99 }
 100
 101 static void sock_map_add_link(struct sk_psock *psock,
 102                               struct sk_psock_link *link,
 103                               struct bpf_map *map, void *link_raw)
 104 {
 105         link->link_raw = link_raw;
 106         link->map = map;
 107         spin_lock_bh(&psock->link_lock);
 108         list_add_tail(&link->list, &psock->link);
 109         spin_unlock_bh(&psock->link_lock);
 110 }
 111
 112 static void sock_map_del_link(struct sock *sk,
 113                               struct sk_psock *psock, void *link_raw)
 114 {
 115         struct sk_psock_link *link, *tmp;
 116         bool strp_stop = false;
 117
 118         spin_lock_bh(&psock->link_lock);
 119         list_for_each_entry_safe(link, tmp, &psock->link, list) {
 120                 if (link->link_raw == link_raw) {
 121                         struct bpf_map *map = link->map;
 122                         struct bpf_stab *stab = container_of(map, struct bpf_stab,
 123                                                              map);
 124                         if (psock->parser.enabled && stab->progs.skb_parser)
 125                                 strp_stop = true;
 126                         list_del(&link->list);
 127                         sk_psock_free_link(link);
 128                 }
 129         }
 130         spin_unlock_bh(&psock->link_lock);
 131         if (strp_stop) {
 132                 write_lock_bh(&sk->sk_callback_lock);
 133                 sk_psock_stop_strp(sk, psock);
 134                 write_unlock_bh(&sk->sk_callback_lock);
 135         }
 136 }
 137
 138 static void sock_map_unref(struct sock *sk, void *link_raw)
 139 {
 140         struct sk_psock *psock = sk_psock(sk);
 141
 142         if (likely(psock)) {
 143                 sock_map_del_link(sk, psock, link_raw);
 144                 sk_psock_put(sk, psock);
 145         }
 146 }
 147
 148 static int sock_map_link(struct bpf_map *map, struct sk_psock_progs *progs,
 149                          struct sock *sk)
 150 {
 151         struct bpf_prog *msg_parser, *skb_parser, *skb_verdict;
 152         bool skb_progs, sk_psock_is_new = false;
 153         struct sk_psock *psock;
 154         int ret;
 155
 156         skb_verdict = READ_ONCE(progs->skb_verdict);
 157         skb_parser = READ_ONCE(progs->skb_parser);
 158         skb_progs = skb_parser && skb_verdict;
 159         if (skb_progs) {
 160                 skb_verdict = bpf_prog_inc_not_zero(skb_verdict);
 161                 if (IS_ERR(skb_verdict))
 162                         return PTR_ERR(skb_verdict);
 163                 skb_parser = bpf_prog_inc_not_zero(skb_parser);
 164                 if (IS_ERR(skb_parser)) {
 165                         bpf_prog_put(skb_verdict);
 166                         return PTR_ERR(skb_parser);
 167                 }
 168         }
 169
 170         msg_parser = READ_ONCE(progs->msg_parser);
 171         if (msg_parser) {
 172                 msg_parser = bpf_prog_inc_not_zero(msg_parser);
 173                 if (IS_ERR(msg_parser)) {
 174                         ret = PTR_ERR(msg_parser);
 175                         goto out;
 176                 }
 177         }
 178
 179         psock = sk_psock_get_checked(sk);
 180         if (IS_ERR(psock)) {
 181                 ret = PTR_ERR(psock);
 182                 goto out_progs;
 183         }
 184
 185         if (psock) {
 186                 if ((msg_parser && READ_ONCE(psock->progs.msg_parser)) ||
 187                     (skb_progs  && READ_ONCE(psock->progs.skb_parser))) {
 188                         sk_psock_put(sk, psock);
 189                         ret = -EBUSY;
 190                         goto out_progs;
 191                 }
 192         } else {
 193                 psock = sk_psock_init(sk, map->numa_node);
 194                 if (!psock) {
 195                         ret = -ENOMEM;
 196                         goto out_progs;
 197                 }
 198                 sk_psock_is_new = true;
 199         }
 200
 201         if (msg_parser)
 202                 psock_set_prog(&psock->progs.msg_parser, msg_parser);
 203         if (sk_psock_is_new) {
 204                 ret = tcp_bpf_init(sk);
 205                 if (ret < 0)
 206                         goto out_drop;
 207         } else {
 208                 tcp_bpf_reinit(sk);
 209         }
 210
 211         write_lock_bh(&sk->sk_callback_lock);
 212         if (skb_progs && !psock->parser.enabled) {
 213                 ret = sk_psock_init_strp(sk, psock);
 214                 if (ret) {
 215                         write_unlock_bh(&sk->sk_callback_lock);
 216                         goto out_drop;
 217                 }
 218                 psock_set_prog(&psock->progs.skb_verdict, skb_verdict);
 219                 psock_set_prog(&psock->progs.skb_parser, skb_parser);
 220                 sk_psock_start_strp(sk, psock);
 221         }
 222         write_unlock_bh(&sk->sk_callback_lock);
 223         return 0;
 224 out_drop:
 225         sk_psock_put(sk, psock);
 226 out_progs:
 227         if (msg_parser)
 228                 bpf_prog_put(msg_parser);
 229 out:
 230         if (skb_progs) {
 231                 bpf_prog_put(skb_verdict);
 232                 bpf_prog_put(skb_parser);
 233         }
 234         return ret;
 235 }
 236
 237 static void sock_map_free(struct bpf_map *map)
 238 {
 239         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 240         int i;
 241
 242         synchronize_rcu();
 243         rcu_read_lock();
 244         raw_spin_lock_bh(&stab->lock);
 245         for (i = 0; i < stab->map.max_entries; i++) {
 246                 struct sock **psk = &stab->sks[i];
 247                 struct sock *sk;
 248
 249                 sk = xchg(psk, NULL);
 250                 if (sk)
 251                         sock_map_unref(sk, psk);
 252         }
 253         raw_spin_unlock_bh(&stab->lock);
 254         rcu_read_unlock();
 255
 256         bpf_map_area_free(stab->sks);
 257         kfree(stab);
 258 }
 259
 260 static void sock_map_release_progs(struct bpf_map *map)
 261 {
 262         psock_progs_drop(&container_of(map, struct bpf_stab, map)->progs);
 263 }
 264
 265 static struct sock *__sock_map_lookup_elem(struct bpf_map *map, u32 key)
 266 {
 267         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 268
 269         WARN_ON_ONCE(!rcu_read_lock_held());
 270
 271         if (unlikely(key >= map->max_entries))
 272                 return NULL;
 273         return READ_ONCE(stab->sks[key]);
 274 }
 275
 276 static void *sock_map_lookup(struct bpf_map *map, void *key)
 277 {
 278         return ERR_PTR(-EOPNOTSUPP);
 279 }
 280
 281 static int __sock_map_delete(struct bpf_stab *stab, struct sock *sk_test,
 282                              struct sock **psk)
 283 {
 284         struct sock *sk;
 285
 286         raw_spin_lock_bh(&stab->lock);
 287         sk = *psk;
 288         if (!sk_test || sk_test == sk)
 289                 *psk = NULL;
 290         raw_spin_unlock_bh(&stab->lock);
 291         if (unlikely(!sk))
 292                 return -EINVAL;
 293         sock_map_unref(sk, psk);
 294         return 0;
 295 }
 296
 297 static void sock_map_delete_from_link(struct bpf_map *map, struct sock *sk,
 298                                       void *link_raw)
 299 {
 300         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 301
 302         __sock_map_delete(stab, sk, link_raw);
 303 }
 304
 305 static int sock_map_delete_elem(struct bpf_map *map, void *key)
 306 {
 307         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 308         u32 i = *(u32 *)key;
 309         struct sock **psk;
 310
 311         if (unlikely(i >= map->max_entries))
 312                 return -EINVAL;
 313
 314         psk = &stab->sks[i];
 315         return __sock_map_delete(stab, NULL, psk);
 316 }
 317
 318 static int sock_map_get_next_key(struct bpf_map *map, void *key, void *next)
 319 {
 320         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 321         u32 i = key ? *(u32 *)key : U32_MAX;
 322         u32 *key_next = next;
 323
 324         if (i == stab->map.max_entries - 1)
 325                 return -ENOENT;
 326         if (i >= stab->map.max_entries)
 327                 *key_next = 0;
 328         else
 329                 *key_next = i + 1;
 330         return 0;
 331 }
 332
 333 static int sock_map_update_common(struct bpf_map *map, u32 idx,
 334                                   struct sock *sk, u64 flags)
 335 {
 336         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 337         struct sk_psock_link *link;
 338         struct sk_psock *psock;
 339         struct sock *osk;
 340         int ret;
 341
 342         WARN_ON_ONCE(!rcu_read_lock_held());
 343         if (unlikely(flags > BPF_EXIST))
 344                 return -EINVAL;
 345         if (unlikely(idx >= map->max_entries))
 346                 return -E2BIG;
 347
 348         link = sk_psock_init_link();
 349         if (!link)
 350                 return -ENOMEM;
 351
 352         ret = sock_map_link(map, &stab->progs, sk);
 353         if (ret < 0)
 354                 goto out_free;
 355
 356         psock = sk_psock(sk);
 357         WARN_ON_ONCE(!psock);
 358
 359         raw_spin_lock_bh(&stab->lock);
 360         osk = stab->sks[idx];
 361         if (osk && flags == BPF_NOEXIST) {
 362                 ret = -EEXIST;
 363                 goto out_unlock;
 364         } else if (!osk && flags == BPF_EXIST) {
 365                 ret = -ENOENT;
 366                 goto out_unlock;
 367         }
 368
 369         sock_map_add_link(psock, link, map, &stab->sks[idx]);
 370         stab->sks[idx] = sk;
 371         if (osk)
 372                 sock_map_unref(osk, &stab->sks[idx]);
 373         raw_spin_unlock_bh(&stab->lock);
 374         return 0;
 375 out_unlock:
 376         raw_spin_unlock_bh(&stab->lock);
 377         if (psock)
 378                 sk_psock_put(sk, psock);
 379 out_free:
 380         sk_psock_free_link(link);
 381         return ret;
 382 }
 383
 384 static bool sock_map_op_okay(const struct bpf_sock_ops_kern *ops)
 385 {
 386         return ops->op == BPF_SOCK_OPS_PASSIVE_ESTABLISHED_CB ||
 387                ops->op == BPF_SOCK_OPS_ACTIVE_ESTABLISHED_CB;
 388 }
 389
 390 static bool sock_map_sk_is_suitable(const struct sock *sk)
 391 {
 392         return sk->sk_type == SOCK_STREAM &&
 393                sk->sk_protocol == IPPROTO_TCP;
 394 }
 395
 396 static int sock_map_update_elem(struct bpf_map *map, void *key,
 397                                 void *value, u64 flags)
 398 {
 399         u32 ufd = *(u32 *)value;
 400         u32 idx = *(u32 *)key;
 401         struct socket *sock;
 402         struct sock *sk;
 403         int ret;
 404
 405         sock = sockfd_lookup(ufd, &ret);
 406         if (!sock)
 407                 return ret;
 408         sk = sock->sk;
 409         if (!sk) {
 410                 ret = -EINVAL;
 411                 goto out;
 412         }
 413         if (!sock_map_sk_is_suitable(sk) ||
 414             sk->sk_state != TCP_ESTABLISHED) {
 415                 ret = -EOPNOTSUPP;
 416                 goto out;
 417         }
 418
 419         sock_map_sk_acquire(sk);
 420         ret = sock_map_update_common(map, idx, sk, flags);
 421         sock_map_sk_release(sk);
 422 out:
 423         fput(sock->file);
 424         return ret;
 425 }
 426
 427 BPF_CALL_4(bpf_sock_map_update, struct bpf_sock_ops_kern *, sops,
 428            struct bpf_map *, map, void *, key, u64, flags)
 429 {
 430         WARN_ON_ONCE(!rcu_read_lock_held());
 431
 432         if (likely(sock_map_sk_is_suitable(sops->sk) &&
 433                    sock_map_op_okay(sops)))
 434                 return sock_map_update_common(map, *(u32 *)key, sops->sk,
 435                                               flags);
 436         return -EOPNOTSUPP;
 437 }
 438
 439 const struct bpf_func_proto bpf_sock_map_update_proto = {
 440         .func           = bpf_sock_map_update,
 441         .gpl_only       = false,
 442         .pkt_access     = true,
 443         .ret_type       = RET_INTEGER,
 444         .arg1_type      = ARG_PTR_TO_CTX,
 445         .arg2_type      = ARG_CONST_MAP_PTR,
 446         .arg3_type      = ARG_PTR_TO_MAP_KEY,
 447         .arg4_type      = ARG_ANYTHING,
 448 };
 449
 450 BPF_CALL_4(bpf_sk_redirect_map, struct sk_buff *, skb,
 451            struct bpf_map *, map, u32, key, u64, flags)
 452 {
 453         struct tcp_skb_cb *tcb = TCP_SKB_CB(skb);
 454
 455         if (unlikely(flags & ~(BPF_F_INGRESS)))
 456                 return SK_DROP;
 457         tcb->bpf.flags = flags;
 458         tcb->bpf.sk_redir = __sock_map_lookup_elem(map, key);
 459         if (!tcb->bpf.sk_redir)
 460                 return SK_DROP;
 461         return SK_PASS;
 462 }
 463
 464 const struct bpf_func_proto bpf_sk_redirect_map_proto = {
 465         .func           = bpf_sk_redirect_map,
 466         .gpl_only       = false,
 467         .ret_type       = RET_INTEGER,
 468         .arg1_type      = ARG_PTR_TO_CTX,
 469         .arg2_type      = ARG_CONST_MAP_PTR,
 470         .arg3_type      = ARG_ANYTHING,
 471         .arg4_type      = ARG_ANYTHING,
 472 };
 473
 474 BPF_CALL_4(bpf_msg_redirect_map, struct sk_msg *, msg,
 475            struct bpf_map *, map, u32, key, u64, flags)
 476 {
 477         if (unlikely(flags & ~(BPF_F_INGRESS)))
 478                 return SK_DROP;
 479         msg->flags = flags;
 480         msg->sk_redir = __sock_map_lookup_elem(map, key);
 481         if (!msg->sk_redir)
 482                 return SK_DROP;
 483         return SK_PASS;
 484 }
 485
 486 const struct bpf_func_proto bpf_msg_redirect_map_proto = {
 487         .func           = bpf_msg_redirect_map,
 488         .gpl_only       = false,
 489         .ret_type       = RET_INTEGER,
 490         .arg1_type      = ARG_PTR_TO_CTX,
 491         .arg2_type      = ARG_CONST_MAP_PTR,
 492         .arg3_type      = ARG_ANYTHING,
 493         .arg4_type      = ARG_ANYTHING,
 494 };
 495
 496 const struct bpf_map_ops sock_map_ops = {
 497         .map_alloc              = sock_map_alloc,
 498         .map_free               = sock_map_free,
 499         .map_get_next_key       = sock_map_get_next_key,
 500         .map_update_elem        = sock_map_update_elem,
 501         .map_delete_elem        = sock_map_delete_elem,
 502         .map_lookup_elem        = sock_map_lookup,
 503         .map_release_uref       = sock_map_release_progs,
 504         .map_check_btf          = map_check_no_btf,
 505 };
 506
 507 struct bpf_htab_elem {
 508         struct rcu_head rcu;
 509         u32 hash;
 510         struct sock *sk;
 511         struct hlist_node node;
 512         u8 key[0];
 513 };
 514
 515 struct bpf_htab_bucket {
 516         struct hlist_head head;
 517         raw_spinlock_t lock;
 518 };
 519
 520 struct bpf_htab {
 521         struct bpf_map map;
 522         struct bpf_htab_bucket *buckets;
 523         u32 buckets_num;
 524         u32 elem_size;
 525         struct sk_psock_progs progs;
 526         atomic_t count;
 527 };
 528
 529 static inline u32 sock_hash_bucket_hash(const void *key, u32 len)
 530 {
 531         return jhash(key, len, 0);
 532 }
 533
 534 static struct bpf_htab_bucket *sock_hash_select_bucket(struct bpf_htab *htab,
 535                                                        u32 hash)
 536 {
 537         return &htab->buckets[hash & (htab->buckets_num - 1)];
 538 }
 539
 540 static struct bpf_htab_elem *
 541 sock_hash_lookup_elem_raw(struct hlist_head *head, u32 hash, void *key,
 542                           u32 key_size)
 543 {
 544         struct bpf_htab_elem *elem;
 545
 546         hlist_for_each_entry_rcu(elem, head, node) {
 547                 if (elem->hash == hash &&
 548                     !memcmp(&elem->key, key, key_size))
 549                         return elem;
 550         }
 551
 552         return NULL;
 553 }
 554
 555 static struct sock *__sock_hash_lookup_elem(struct bpf_map *map, void *key)
 556 {
 557         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 558         u32 key_size = map->key_size, hash;
 559         struct bpf_htab_bucket *bucket;
 560         struct bpf_htab_elem *elem;
 561
 562         WARN_ON_ONCE(!rcu_read_lock_held());
 563
 564         hash = sock_hash_bucket_hash(key, key_size);
 565         bucket = sock_hash_select_bucket(htab, hash);
 566         elem = sock_hash_lookup_elem_raw(&bucket->head, hash, key, key_size);
 567
 568         return elem ? elem->sk : NULL;
 569 }
 570
 571 static void sock_hash_free_elem(struct bpf_htab *htab,
 572                                 struct bpf_htab_elem *elem)
 573 {
 574         atomic_dec(&htab->count);
 575         kfree_rcu(elem, rcu);
 576 }
 577
 578 static void sock_hash_delete_from_link(struct bpf_map *map, struct sock *sk,
 579                                        void *link_raw)
 580 {
 581         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 582         struct bpf_htab_elem *elem_probe, *elem = link_raw;
 583         struct bpf_htab_bucket *bucket;
 584
 585         WARN_ON_ONCE(!rcu_read_lock_held());
 586         bucket = sock_hash_select_bucket(htab, elem->hash);
 587
 588         /* elem may be deleted in parallel from the map, but access here
 589          * is okay since it's going away only after RCU grace period.
 590          * However, we need to check whether it's still present.
 591          */
 592         raw_spin_lock_bh(&bucket->lock);
 593         elem_probe = sock_hash_lookup_elem_raw(&bucket->head, elem->hash,
 594                                                elem->key, map->key_size);
 595         if (elem_probe && elem_probe == elem) {
 596                 hlist_del_rcu(&elem->node);
 597                 sock_map_unref(elem->sk, elem);
 598                 sock_hash_free_elem(htab, elem);
 599         }
 600         raw_spin_unlock_bh(&bucket->lock);
 601 }
 602
 603 static int sock_hash_delete_elem(struct bpf_map *map, void *key)
 604 {
 605         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 606         u32 hash, key_size = map->key_size;
 607         struct bpf_htab_bucket *bucket;
 608         struct bpf_htab_elem *elem;
 609         int ret = -ENOENT;
 610
 611         hash = sock_hash_bucket_hash(key, key_size);
 612         bucket = sock_hash_select_bucket(htab, hash);
 613
 614         raw_spin_lock_bh(&bucket->lock);
 615         elem = sock_hash_lookup_elem_raw(&bucket->head, hash, key, key_size);
 616         if (elem) {
 617                 hlist_del_rcu(&elem->node);
 618                 sock_map_unref(elem->sk, elem);
 619                 sock_hash_free_elem(htab, elem);
 620                 ret = 0;
 621         }
 622         raw_spin_unlock_bh(&bucket->lock);
 623         return ret;
 624 }
 625
 626 static struct bpf_htab_elem *sock_hash_alloc_elem(struct bpf_htab *htab,
 627                                                   void *key, u32 key_size,
 628                                                   u32 hash, struct sock *sk,
 629                                                   struct bpf_htab_elem *old)
 630 {
 631         struct bpf_htab_elem *new;
 632
 633         if (atomic_inc_return(&htab->count) > htab->map.max_entries) {
 634                 if (!old) {
 635                         atomic_dec(&htab->count);
 636                         return ERR_PTR(-E2BIG);
 637                 }
 638         }
 639
 640         new = kmalloc_node(htab->elem_size, GFP_ATOMIC | __GFP_NOWARN,
 641                            htab->map.numa_node);
 642         if (!new) {
 643                 atomic_dec(&htab->count);
 644                 return ERR_PTR(-ENOMEM);
 645         }
 646         memcpy(new->key, key, key_size);
 647         new->sk = sk;
 648         new->hash = hash;
 649         return new;
 650 }
 651
 652 static int sock_hash_update_common(struct bpf_map *map, void *key,
 653                                    struct sock *sk, u64 flags)
 654 {
 655         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 656         u32 key_size = map->key_size, hash;
 657         struct bpf_htab_elem *elem, *elem_new;
 658         struct bpf_htab_bucket *bucket;
 659         struct sk_psock_link *link;
 660         struct sk_psock *psock;
 661         int ret;
 662
 663         WARN_ON_ONCE(!rcu_read_lock_held());
 664         if (unlikely(flags > BPF_EXIST))
 665                 return -EINVAL;
 666
 667         link = sk_psock_init_link();
 668         if (!link)
 669                 return -ENOMEM;
 670
 671         ret = sock_map_link(map, &htab->progs, sk);
 672         if (ret < 0)
 673                 goto out_free;
 674
 675         psock = sk_psock(sk);
 676         WARN_ON_ONCE(!psock);
 677
 678         hash = sock_hash_bucket_hash(key, key_size);
 679         bucket = sock_hash_select_bucket(htab, hash);
 680
 681         raw_spin_lock_bh(&bucket->lock);
 682         elem = sock_hash_lookup_elem_raw(&bucket->head, hash, key, key_size);
 683         if (elem && flags == BPF_NOEXIST) {
 684                 ret = -EEXIST;
 685                 goto out_unlock;
 686         } else if (!elem && flags == BPF_EXIST) {
 687                 ret = -ENOENT;
 688                 goto out_unlock;
 689         }
 690
 691         elem_new = sock_hash_alloc_elem(htab, key, key_size, hash, sk, elem);
 692         if (IS_ERR(elem_new)) {
 693                 ret = PTR_ERR(elem_new);
 694                 goto out_unlock;
 695         }
 696
 697         sock_map_add_link(psock, link, map, elem_new);
 698         /* Add new element to the head of the list, so that
 699          * concurrent search will find it before old elem.
 700          */
 701         hlist_add_head_rcu(&elem_new->node, &bucket->head);
 702         if (elem) {
 703                 hlist_del_rcu(&elem->node);
 704                 sock_map_unref(elem->sk, elem);
 705                 sock_hash_free_elem(htab, elem);
 706         }
 707         raw_spin_unlock_bh(&bucket->lock);
 708         return 0;
 709 out_unlock:
 710         raw_spin_unlock_bh(&bucket->lock);
 711         sk_psock_put(sk, psock);
 712 out_free:
 713         sk_psock_free_link(link);
 714         return ret;
 715 }
 716
 717 static int sock_hash_update_elem(struct bpf_map *map, void *key,
 718                                  void *value, u64 flags)
 719 {
 720         u32 ufd = *(u32 *)value;
 721         struct socket *sock;
 722         struct sock *sk;
 723         int ret;
 724
 725         sock = sockfd_lookup(ufd, &ret);
 726         if (!sock)
 727                 return ret;
 728         sk = sock->sk;
 729         if (!sk) {
 730                 ret = -EINVAL;
 731                 goto out;
 732         }
 733         if (!sock_map_sk_is_suitable(sk) ||
 734             sk->sk_state != TCP_ESTABLISHED) {
 735                 ret = -EOPNOTSUPP;
 736                 goto out;
 737         }
 738
 739         sock_map_sk_acquire(sk);
 740         ret = sock_hash_update_common(map, key, sk, flags);
 741         sock_map_sk_release(sk);
 742 out:
 743         fput(sock->file);
 744         return ret;
 745 }
 746
 747 static int sock_hash_get_next_key(struct bpf_map *map, void *key,
 748                                   void *key_next)
 749 {
 750         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 751         struct bpf_htab_elem *elem, *elem_next;
 752         u32 hash, key_size = map->key_size;
 753         struct hlist_head *head;
 754         int i = 0;
 755
 756         if (!key)
 757                 goto find_first_elem;
 758         hash = sock_hash_bucket_hash(key, key_size);
 759         head = &sock_hash_select_bucket(htab, hash)->head;
 760         elem = sock_hash_lookup_elem_raw(head, hash, key, key_size);
 761         if (!elem)
 762                 goto find_first_elem;
 763
 764         elem_next = hlist_entry_safe(rcu_dereference_raw(hlist_next_rcu(&elem->node)),
 765                                      struct bpf_htab_elem, node);
 766         if (elem_next) {
 767                 memcpy(key_next, elem_next->key, key_size);
 768                 return 0;
 769         }
 770
 771         i = hash & (htab->buckets_num - 1);
 772         i++;
 773 find_first_elem:
 774         for (; i < htab->buckets_num; i++) {
 775                 head = &sock_hash_select_bucket(htab, i)->head;
 776                 elem_next = hlist_entry_safe(rcu_dereference_raw(hlist_first_rcu(head)),
 777                                              struct bpf_htab_elem, node);
 778                 if (elem_next) {
 779                         memcpy(key_next, elem_next->key, key_size);
 780                         return 0;
 781                 }
 782         }
 783
 784         return -ENOENT;
 785 }
 786
 787 static struct bpf_map *sock_hash_alloc(union bpf_attr *attr)
 788 {
 789         struct bpf_htab *htab;
 790         int i, err;
 791         u64 cost;
 792
 793         if (!capable(CAP_NET_ADMIN))
 794                 return ERR_PTR(-EPERM);
 795         if (attr->max_entries == 0 ||
 796             attr->key_size    == 0 ||
 797             attr->value_size  != 4 ||
 798             attr->map_flags & ~SOCK_CREATE_FLAG_MASK)
 799                 return ERR_PTR(-EINVAL);
 800         if (attr->key_size > MAX_BPF_STACK)
 801                 return ERR_PTR(-E2BIG);
 802
 803         htab = kzalloc(sizeof(*htab), GFP_USER);
 804         if (!htab)
 805                 return ERR_PTR(-ENOMEM);
 806
 807         bpf_map_init_from_attr(&htab->map, attr);
 808
 809         htab->buckets_num = roundup_pow_of_two(htab->map.max_entries);
 810         htab->elem_size = sizeof(struct bpf_htab_elem) +
 811                           round_up(htab->map.key_size, 8);
 812         if (htab->buckets_num == 0 ||
 813             htab->buckets_num > U32_MAX / sizeof(struct bpf_htab_bucket)) {
 814                 err = -EINVAL;
 815                 goto free_htab;
 816         }
 817
 818         cost = (u64) htab->buckets_num * sizeof(struct bpf_htab_bucket) +
 819                (u64) htab->elem_size * htab->map.max_entries;
 820         if (cost >= U32_MAX - PAGE_SIZE) {
 821                 err = -EINVAL;
 822                 goto free_htab;
 823         }
 824
 825         htab->buckets = bpf_map_area_alloc(htab->buckets_num *
 826                                            sizeof(struct bpf_htab_bucket),
 827                                            htab->map.numa_node);
 828         if (!htab->buckets) {
 829                 err = -ENOMEM;
 830                 goto free_htab;
 831         }
 832
 833         for (i = 0; i < htab->buckets_num; i++) {
 834                 INIT_HLIST_HEAD(&htab->buckets[i].head);
 835                 raw_spin_lock_init(&htab->buckets[i].lock);
 836         }
 837
 838         return &htab->map;
 839 free_htab:
 840         kfree(htab);
 841         return ERR_PTR(err);
 842 }
 843
 844 static void sock_hash_free(struct bpf_map *map)
 845 {
 846         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
 847         struct bpf_htab_bucket *bucket;
 848         struct bpf_htab_elem *elem;
 849         struct hlist_node *node;
 850         int i;
 851
 852         synchronize_rcu();
 853         rcu_read_lock();
 854         for (i = 0; i < htab->buckets_num; i++) {
 855                 bucket = sock_hash_select_bucket(htab, i);
 856                 raw_spin_lock_bh(&bucket->lock);
 857                 hlist_for_each_entry_safe(elem, node, &bucket->head, node) {
 858                         hlist_del_rcu(&elem->node);
 859                         sock_map_unref(elem->sk, elem);
 860                 }
 861                 raw_spin_unlock_bh(&bucket->lock);
 862         }
 863         rcu_read_unlock();
 864
 865         bpf_map_area_free(htab->buckets);
 866         kfree(htab);
 867 }
 868
 869 static void sock_hash_release_progs(struct bpf_map *map)
 870 {
 871         psock_progs_drop(&container_of(map, struct bpf_htab, map)->progs);
 872 }
 873
 874 BPF_CALL_4(bpf_sock_hash_update, struct bpf_sock_ops_kern *, sops,
 875            struct bpf_map *, map, void *, key, u64, flags)
 876 {
 877         WARN_ON_ONCE(!rcu_read_lock_held());
 878
 879         if (likely(sock_map_sk_is_suitable(sops->sk) &&
 880                    sock_map_op_okay(sops)))
 881                 return sock_hash_update_common(map, key, sops->sk, flags);
 882         return -EOPNOTSUPP;
 883 }
 884
 885 const struct bpf_func_proto bpf_sock_hash_update_proto = {
 886         .func           = bpf_sock_hash_update,
 887         .gpl_only       = false,
 888         .pkt_access     = true,
 889         .ret_type       = RET_INTEGER,
 890         .arg1_type      = ARG_PTR_TO_CTX,
 891         .arg2_type      = ARG_CONST_MAP_PTR,
 892         .arg3_type      = ARG_PTR_TO_MAP_KEY,
 893         .arg4_type      = ARG_ANYTHING,
 894 };
 895
 896 BPF_CALL_4(bpf_sk_redirect_hash, struct sk_buff *, skb,
 897            struct bpf_map *, map, void *, key, u64, flags)
 898 {
 899         struct tcp_skb_cb *tcb = TCP_SKB_CB(skb);
 900
 901         if (unlikely(flags & ~(BPF_F_INGRESS)))
 902                 return SK_DROP;
 903         tcb->bpf.flags = flags;
 904         tcb->bpf.sk_redir = __sock_hash_lookup_elem(map, key);
 905         if (!tcb->bpf.sk_redir)
 906                 return SK_DROP;
 907         return SK_PASS;
 908 }
 909
 910 const struct bpf_func_proto bpf_sk_redirect_hash_proto = {
 911         .func           = bpf_sk_redirect_hash,
 912         .gpl_only       = false,
 913         .ret_type       = RET_INTEGER,
 914         .arg1_type      = ARG_PTR_TO_CTX,
 915         .arg2_type      = ARG_CONST_MAP_PTR,
 916         .arg3_type      = ARG_PTR_TO_MAP_KEY,
 917         .arg4_type      = ARG_ANYTHING,
 918 };
 919
 920 BPF_CALL_4(bpf_msg_redirect_hash, struct sk_msg *, msg,
 921            struct bpf_map *, map, void *, key, u64, flags)
 922 {
 923         if (unlikely(flags & ~(BPF_F_INGRESS)))
 924                 return SK_DROP;
 925         msg->flags = flags;
 926         msg->sk_redir = __sock_hash_lookup_elem(map, key);
 927         if (!msg->sk_redir)
 928                 return SK_DROP;
 929         return SK_PASS;
 930 }
 931
 932 const struct bpf_func_proto bpf_msg_redirect_hash_proto = {
 933         .func           = bpf_msg_redirect_hash,
 934         .gpl_only       = false,
 935         .ret_type       = RET_INTEGER,
 936         .arg1_type      = ARG_PTR_TO_CTX,
 937         .arg2_type      = ARG_CONST_MAP_PTR,
 938         .arg3_type      = ARG_PTR_TO_MAP_KEY,
 939         .arg4_type      = ARG_ANYTHING,
 940 };
 941
 942 const struct bpf_map_ops sock_hash_ops = {
 943         .map_alloc              = sock_hash_alloc,
 944         .map_free               = sock_hash_free,
 945         .map_get_next_key       = sock_hash_get_next_key,
 946         .map_update_elem        = sock_hash_update_elem,
 947         .map_delete_elem        = sock_hash_delete_elem,
 948         .map_lookup_elem        = sock_map_lookup,
 949         .map_release_uref       = sock_hash_release_progs,
 950         .map_check_btf          = map_check_no_btf,
 951 };
 952
 953 static struct sk_psock_progs *sock_map_progs(struct bpf_map *map)
 954 {
 955         switch (map->map_type) {
 956         case BPF_MAP_TYPE_SOCKMAP:
 957                 return &container_of(map, struct bpf_stab, map)->progs;
 958         case BPF_MAP_TYPE_SOCKHASH:
 959                 return &container_of(map, struct bpf_htab, map)->progs;
 960         default:
 961                 break;
 962         }
 963
 964         return NULL;
 965 }
 966
 967 int sock_map_prog_update(struct bpf_map *map, struct bpf_prog *prog,
 968                          u32 which)
 969 {
 970         struct sk_psock_progs *progs = sock_map_progs(map);
 971
 972         if (!progs)
 973                 return -EOPNOTSUPP;
 974
 975         switch (which) {
 976         case BPF_SK_MSG_VERDICT:
 977                 psock_set_prog(&progs->msg_parser, prog);
 978                 break;
 979         case BPF_SK_SKB_STREAM_PARSER:
 980                 psock_set_prog(&progs->skb_parser, prog);
 981                 break;
 982         case BPF_SK_SKB_STREAM_VERDICT:
 983                 psock_set_prog(&progs->skb_verdict, prog);
 984                 break;
 985         default:
 986                 return -EOPNOTSUPP;
 987         }
 988
 989         return 0;
 990 }
 991
 992 void sk_psock_unlink(struct sock *sk, struct sk_psock_link *link)
 993 {
 994         switch (link->map->map_type) {
 995         case BPF_MAP_TYPE_SOCKMAP:
 996                 return sock_map_delete_from_link(link->map, sk,
 997                                                  link->link_raw);
 998         case BPF_MAP_TYPE_SOCKHASH:
 999                 return sock_hash_delete_from_link(link->map, sk,
1000                                                   link->link_raw);
1001         default:
1002                 break;
1003         }
1004 }