kernel/bpf/sockmap.c

   1 /* Copyright (c) 2017 Covalent IO, Inc. http://covalent.io
   2  *
   3  * This program is free software; you can redistribute it and/or
   4  * modify it under the terms of version 2 of the GNU General Public
   5  * License as published by the Free Software Foundation.
   6  *
   7  * This program is distributed in the hope that it will be useful, but
   8  * WITHOUT ANY WARRANTY; without even the implied warranty of
   9  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  10  * General Public License for more details.
  11  */
  12
  13 /* A BPF sock_map is used to store sock objects. This is primarly used
  14  * for doing socket redirect with BPF helper routines.
  15  *
  16  * A sock map may have BPF programs attached to it, currently a program
  17  * used to parse packets and a program to provide a verdict and redirect
  18  * decision on the packet are supported. Any programs attached to a sock
  19  * map are inherited by sock objects when they are added to the map. If
  20  * no BPF programs are attached the sock object may only be used for sock
  21  * redirect.
  22  *
  23  * A sock object may be in multiple maps, but can only inherit a single
  24  * parse or verdict program. If adding a sock object to a map would result
  25  * in having multiple parsing programs the update will return an EBUSY error.
  26  *
  27  * For reference this program is similar to devmap used in XDP context
  28  * reviewing these together may be useful. For an example please review
  29  * ./samples/bpf/sockmap/.
  30  */
  31 #include <linux/bpf.h>
  32 #include <net/sock.h>
  33 #include <linux/filter.h>
  34 #include <linux/errno.h>
  35 #include <linux/file.h>
  36 #include <linux/kernel.h>
  37 #include <linux/net.h>
  38 #include <linux/skbuff.h>
  39 #include <linux/workqueue.h>
  40 #include <linux/list.h>
  41 #include <linux/mm.h>
  42 #include <net/strparser.h>
  43 #include <net/tcp.h>
  44 #include <linux/ptr_ring.h>
  45 #include <net/inet_common.h>
  46 #include <linux/sched/signal.h>
  47
  48 #define SOCK_CREATE_FLAG_MASK \
  49         (BPF_F_NUMA_NODE | BPF_F_RDONLY | BPF_F_WRONLY)
  50
  51 struct bpf_sock_progs {
  52         struct bpf_prog *bpf_tx_msg;
  53         struct bpf_prog *bpf_parse;
  54         struct bpf_prog *bpf_verdict;
  55 };
  56
  57 struct bpf_stab {
  58         struct bpf_map map;
  59         struct sock **sock_map;
  60         struct bpf_sock_progs progs;
  61 };
  62
  63 struct bucket {
  64         struct hlist_head head;
  65         raw_spinlock_t lock;
  66 };
  67
  68 struct bpf_htab {
  69         struct bpf_map map;
  70         struct bucket *buckets;
  71         atomic_t count;
  72         u32 n_buckets;
  73         u32 elem_size;
  74         struct bpf_sock_progs progs;
  75 };
  76
  77 struct htab_elem {
  78         struct rcu_head rcu;
  79         struct hlist_node hash_node;
  80         u32 hash;
  81         struct sock *sk;
  82         char key[0];
  83 };
  84
  85 enum smap_psock_state {
  86         SMAP_TX_RUNNING,
  87 };
  88
  89 struct smap_psock_map_entry {
  90         struct list_head list;
  91         struct sock **entry;
  92         struct htab_elem *hash_link;
  93         struct bpf_htab *htab;
  94 };
  95
  96 struct smap_psock {
  97         struct rcu_head rcu;
  98         refcount_t refcnt;
  99
 100         /* datapath variables */
 101         struct sk_buff_head rxqueue;
 102         bool strp_enabled;
 103
 104         /* datapath error path cache across tx work invocations */
 105         int save_rem;
 106         int save_off;
 107         struct sk_buff *save_skb;
 108
 109         /* datapath variables for tx_msg ULP */
 110         struct sock *sk_redir;
 111         int apply_bytes;
 112         int cork_bytes;
 113         int sg_size;
 114         int eval;
 115         struct sk_msg_buff *cork;
 116         struct list_head ingress;
 117
 118         struct strparser strp;
 119         struct bpf_prog *bpf_tx_msg;
 120         struct bpf_prog *bpf_parse;
 121         struct bpf_prog *bpf_verdict;
 122         struct list_head maps;
 123
 124         /* Back reference used when sock callback trigger sockmap operations */
 125         struct sock *sock;
 126         unsigned long state;
 127
 128         struct work_struct tx_work;
 129         struct work_struct gc_work;
 130
 131         struct proto *sk_proto;
 132         void (*save_close)(struct sock *sk, long timeout);
 133         void (*save_data_ready)(struct sock *sk);
 134         void (*save_write_space)(struct sock *sk);
 135 };
 136
 137 static void smap_release_sock(struct smap_psock *psock, struct sock *sock);
 138 static int bpf_tcp_recvmsg(struct sock *sk, struct msghdr *msg, size_t len,
 139                            int nonblock, int flags, int *addr_len);
 140 static int bpf_tcp_sendmsg(struct sock *sk, struct msghdr *msg, size_t size);
 141 static int bpf_tcp_sendpage(struct sock *sk, struct page *page,
 142                             int offset, size_t size, int flags);
 143
 144 static inline struct smap_psock *smap_psock_sk(const struct sock *sk)
 145 {
 146         return rcu_dereference_sk_user_data(sk);
 147 }
 148
 149 static bool bpf_tcp_stream_read(const struct sock *sk)
 150 {
 151         struct smap_psock *psock;
 152         bool empty = true;
 153
 154         rcu_read_lock();
 155         psock = smap_psock_sk(sk);
 156         if (unlikely(!psock))
 157                 goto out;
 158         empty = list_empty(&psock->ingress);
 159 out:
 160         rcu_read_unlock();
 161         return !empty;
 162 }
 163
 164 static struct proto tcp_bpf_proto;
 165 static int bpf_tcp_init(struct sock *sk)
 166 {
 167         struct smap_psock *psock;
 168
 169         rcu_read_lock();
 170         psock = smap_psock_sk(sk);
 171         if (unlikely(!psock)) {
 172                 rcu_read_unlock();
 173                 return -EINVAL;
 174         }
 175
 176         if (unlikely(psock->sk_proto)) {
 177                 rcu_read_unlock();
 178                 return -EBUSY;
 179         }
 180
 181         psock->save_close = sk->sk_prot->close;
 182         psock->sk_proto = sk->sk_prot;
 183
 184         if (psock->bpf_tx_msg) {
 185                 tcp_bpf_proto.sendmsg = bpf_tcp_sendmsg;
 186                 tcp_bpf_proto.sendpage = bpf_tcp_sendpage;
 187                 tcp_bpf_proto.recvmsg = bpf_tcp_recvmsg;
 188                 tcp_bpf_proto.stream_memory_read = bpf_tcp_stream_read;
 189         }
 190
 191         sk->sk_prot = &tcp_bpf_proto;
 192         rcu_read_unlock();
 193         return 0;
 194 }
 195
 196 static void smap_release_sock(struct smap_psock *psock, struct sock *sock);
 197 static int free_start_sg(struct sock *sk, struct sk_msg_buff *md);
 198
 199 static void bpf_tcp_release(struct sock *sk)
 200 {
 201         struct smap_psock *psock;
 202
 203         rcu_read_lock();
 204         psock = smap_psock_sk(sk);
 205         if (unlikely(!psock))
 206                 goto out;
 207
 208         if (psock->cork) {
 209                 free_start_sg(psock->sock, psock->cork);
 210                 kfree(psock->cork);
 211                 psock->cork = NULL;
 212         }
 213
 214         if (psock->sk_proto) {
 215                 sk->sk_prot = psock->sk_proto;
 216                 psock->sk_proto = NULL;
 217         }
 218 out:
 219         rcu_read_unlock();
 220 }
 221
 222 static void free_htab_elem(struct bpf_htab *htab, struct htab_elem *l)
 223 {
 224         atomic_dec(&htab->count);
 225         kfree_rcu(l, rcu);
 226 }
 227
 228 static void bpf_tcp_close(struct sock *sk, long timeout)
 229 {
 230         void (*close_fun)(struct sock *sk, long timeout);
 231         struct smap_psock_map_entry *e, *tmp;
 232         struct sk_msg_buff *md, *mtmp;
 233         struct smap_psock *psock;
 234         struct sock *osk;
 235
 236         rcu_read_lock();
 237         psock = smap_psock_sk(sk);
 238         if (unlikely(!psock)) {
 239                 rcu_read_unlock();
 240                 return sk->sk_prot->close(sk, timeout);
 241         }
 242
 243         /* The psock may be destroyed anytime after exiting the RCU critial
 244          * section so by the time we use close_fun the psock may no longer
 245          * be valid. However, bpf_tcp_close is called with the sock lock
 246          * held so the close hook and sk are still valid.
 247          */
 248         close_fun = psock->save_close;
 249
 250         write_lock_bh(&sk->sk_callback_lock);
 251         if (psock->cork) {
 252                 free_start_sg(psock->sock, psock->cork);
 253                 kfree(psock->cork);
 254                 psock->cork = NULL;
 255         }
 256
 257         list_for_each_entry_safe(md, mtmp, &psock->ingress, list) {
 258                 list_del(&md->list);
 259                 free_start_sg(psock->sock, md);
 260                 kfree(md);
 261         }
 262
 263         list_for_each_entry_safe(e, tmp, &psock->maps, list) {
 264                 if (e->entry) {
 265                         osk = cmpxchg(e->entry, sk, NULL);
 266                         if (osk == sk) {
 267                                 list_del(&e->list);
 268                                 smap_release_sock(psock, sk);
 269                         }
 270                 } else {
 271                         hlist_del_rcu(&e->hash_link->hash_node);
 272                         smap_release_sock(psock, e->hash_link->sk);
 273                         free_htab_elem(e->htab, e->hash_link);
 274                 }
 275         }
 276         write_unlock_bh(&sk->sk_callback_lock);
 277         rcu_read_unlock();
 278         close_fun(sk, timeout);
 279 }
 280
 281 enum __sk_action {
 282         __SK_DROP = 0,
 283         __SK_PASS,
 284         __SK_REDIRECT,
 285         __SK_NONE,
 286 };
 287
 288 static struct tcp_ulp_ops bpf_tcp_ulp_ops __read_mostly = {
 289         .name           = "bpf_tcp",
 290         .uid            = TCP_ULP_BPF,
 291         .user_visible   = false,
 292         .owner          = NULL,
 293         .init           = bpf_tcp_init,
 294         .release        = bpf_tcp_release,
 295 };
 296
 297 static int memcopy_from_iter(struct sock *sk,
 298                              struct sk_msg_buff *md,
 299                              struct iov_iter *from, int bytes)
 300 {
 301         struct scatterlist *sg = md->sg_data;
 302         int i = md->sg_curr, rc = -ENOSPC;
 303
 304         do {
 305                 int copy;
 306                 char *to;
 307
 308                 if (md->sg_copybreak >= sg[i].length) {
 309                         md->sg_copybreak = 0;
 310
 311                         if (++i == MAX_SKB_FRAGS)
 312                                 i = 0;
 313
 314                         if (i == md->sg_end)
 315                                 break;
 316                 }
 317
 318                 copy = sg[i].length - md->sg_copybreak;
 319                 to = sg_virt(&sg[i]) + md->sg_copybreak;
 320                 md->sg_copybreak += copy;
 321
 322                 if (sk->sk_route_caps & NETIF_F_NOCACHE_COPY)
 323                         rc = copy_from_iter_nocache(to, copy, from);
 324                 else
 325                         rc = copy_from_iter(to, copy, from);
 326
 327                 if (rc != copy) {
 328                         rc = -EFAULT;
 329                         goto out;
 330                 }
 331
 332                 bytes -= copy;
 333                 if (!bytes)
 334                         break;
 335
 336                 md->sg_copybreak = 0;
 337                 if (++i == MAX_SKB_FRAGS)
 338                         i = 0;
 339         } while (i != md->sg_end);
 340 out:
 341         md->sg_curr = i;
 342         return rc;
 343 }
 344
 345 static int bpf_tcp_push(struct sock *sk, int apply_bytes,
 346                         struct sk_msg_buff *md,
 347                         int flags, bool uncharge)
 348 {
 349         bool apply = apply_bytes;
 350         struct scatterlist *sg;
 351         int offset, ret = 0;
 352         struct page *p;
 353         size_t size;
 354
 355         while (1) {
 356                 sg = md->sg_data + md->sg_start;
 357                 size = (apply && apply_bytes < sg->length) ?
 358                         apply_bytes : sg->length;
 359                 offset = sg->offset;
 360
 361                 tcp_rate_check_app_limited(sk);
 362                 p = sg_page(sg);
 363 retry:
 364                 ret = do_tcp_sendpages(sk, p, offset, size, flags);
 365                 if (ret != size) {
 366                         if (ret > 0) {
 367                                 if (apply)
 368                                         apply_bytes -= ret;
 369
 370                                 sg->offset += ret;
 371                                 sg->length -= ret;
 372                                 size -= ret;
 373                                 offset += ret;
 374                                 if (uncharge)
 375                                         sk_mem_uncharge(sk, ret);
 376                                 goto retry;
 377                         }
 378
 379                         return ret;
 380                 }
 381
 382                 if (apply)
 383                         apply_bytes -= ret;
 384                 sg->offset += ret;
 385                 sg->length -= ret;
 386                 if (uncharge)
 387                         sk_mem_uncharge(sk, ret);
 388
 389                 if (!sg->length) {
 390                         put_page(p);
 391                         md->sg_start++;
 392                         if (md->sg_start == MAX_SKB_FRAGS)
 393                                 md->sg_start = 0;
 394                         sg_init_table(sg, 1);
 395
 396                         if (md->sg_start == md->sg_end)
 397                                 break;
 398                 }
 399
 400                 if (apply && !apply_bytes)
 401                         break;
 402         }
 403         return 0;
 404 }
 405
 406 static inline void bpf_compute_data_pointers_sg(struct sk_msg_buff *md)
 407 {
 408         struct scatterlist *sg = md->sg_data + md->sg_start;
 409
 410         if (md->sg_copy[md->sg_start]) {
 411                 md->data = md->data_end = 0;
 412         } else {
 413                 md->data = sg_virt(sg);
 414                 md->data_end = md->data + sg->length;
 415         }
 416 }
 417
 418 static void return_mem_sg(struct sock *sk, int bytes, struct sk_msg_buff *md)
 419 {
 420         struct scatterlist *sg = md->sg_data;
 421         int i = md->sg_start;
 422
 423         do {
 424                 int uncharge = (bytes < sg[i].length) ? bytes : sg[i].length;
 425
 426                 sk_mem_uncharge(sk, uncharge);
 427                 bytes -= uncharge;
 428                 if (!bytes)
 429                         break;
 430                 i++;
 431                 if (i == MAX_SKB_FRAGS)
 432                         i = 0;
 433         } while (i != md->sg_end);
 434 }
 435
 436 static void free_bytes_sg(struct sock *sk, int bytes,
 437                           struct sk_msg_buff *md, bool charge)
 438 {
 439         struct scatterlist *sg = md->sg_data;
 440         int i = md->sg_start, free;
 441
 442         while (bytes && sg[i].length) {
 443                 free = sg[i].length;
 444                 if (bytes < free) {
 445                         sg[i].length -= bytes;
 446                         sg[i].offset += bytes;
 447                         if (charge)
 448                                 sk_mem_uncharge(sk, bytes);
 449                         break;
 450                 }
 451
 452                 if (charge)
 453                         sk_mem_uncharge(sk, sg[i].length);
 454                 put_page(sg_page(&sg[i]));
 455                 bytes -= sg[i].length;
 456                 sg[i].length = 0;
 457                 sg[i].page_link = 0;
 458                 sg[i].offset = 0;
 459                 i++;
 460
 461                 if (i == MAX_SKB_FRAGS)
 462                         i = 0;
 463         }
 464         md->sg_start = i;
 465 }
 466
 467 static int free_sg(struct sock *sk, int start, struct sk_msg_buff *md)
 468 {
 469         struct scatterlist *sg = md->sg_data;
 470         int i = start, free = 0;
 471
 472         while (sg[i].length) {
 473                 free += sg[i].length;
 474                 sk_mem_uncharge(sk, sg[i].length);
 475                 put_page(sg_page(&sg[i]));
 476                 sg[i].length = 0;
 477                 sg[i].page_link = 0;
 478                 sg[i].offset = 0;
 479                 i++;
 480
 481                 if (i == MAX_SKB_FRAGS)
 482                         i = 0;
 483         }
 484
 485         return free;
 486 }
 487
 488 static int free_start_sg(struct sock *sk, struct sk_msg_buff *md)
 489 {
 490         int free = free_sg(sk, md->sg_start, md);
 491
 492         md->sg_start = md->sg_end;
 493         return free;
 494 }
 495
 496 static int free_curr_sg(struct sock *sk, struct sk_msg_buff *md)
 497 {
 498         return free_sg(sk, md->sg_curr, md);
 499 }
 500
 501 static int bpf_map_msg_verdict(int _rc, struct sk_msg_buff *md)
 502 {
 503         return ((_rc == SK_PASS) ?
 504                (md->sk_redir ? __SK_REDIRECT : __SK_PASS) :
 505                __SK_DROP);
 506 }
 507
 508 static unsigned int smap_do_tx_msg(struct sock *sk,
 509                                    struct smap_psock *psock,
 510                                    struct sk_msg_buff *md)
 511 {
 512         struct bpf_prog *prog;
 513         unsigned int rc, _rc;
 514
 515         preempt_disable();
 516         rcu_read_lock();
 517
 518         /* If the policy was removed mid-send then default to 'accept' */
 519         prog = READ_ONCE(psock->bpf_tx_msg);
 520         if (unlikely(!prog)) {
 521                 _rc = SK_PASS;
 522                 goto verdict;
 523         }
 524
 525         bpf_compute_data_pointers_sg(md);
 526         rc = (*prog->bpf_func)(md, prog->insnsi);
 527         psock->apply_bytes = md->apply_bytes;
 528
 529         /* Moving return codes from UAPI namespace into internal namespace */
 530         _rc = bpf_map_msg_verdict(rc, md);
 531
 532         /* The psock has a refcount on the sock but not on the map and because
 533          * we need to drop rcu read lock here its possible the map could be
 534          * removed between here and when we need it to execute the sock
 535          * redirect. So do the map lookup now for future use.
 536          */
 537         if (_rc == __SK_REDIRECT) {
 538                 if (psock->sk_redir)
 539                         sock_put(psock->sk_redir);
 540                 psock->sk_redir = do_msg_redirect_map(md);
 541                 if (!psock->sk_redir) {
 542                         _rc = __SK_DROP;
 543                         goto verdict;
 544                 }
 545                 sock_hold(psock->sk_redir);
 546         }
 547 verdict:
 548         rcu_read_unlock();
 549         preempt_enable();
 550
 551         return _rc;
 552 }
 553
 554 static int bpf_tcp_ingress(struct sock *sk, int apply_bytes,
 555                            struct smap_psock *psock,
 556                            struct sk_msg_buff *md, int flags)
 557 {
 558         bool apply = apply_bytes;
 559         size_t size, copied = 0;
 560         struct sk_msg_buff *r;
 561         int err = 0, i;
 562
 563         r = kzalloc(sizeof(struct sk_msg_buff), __GFP_NOWARN | GFP_KERNEL);
 564         if (unlikely(!r))
 565                 return -ENOMEM;
 566
 567         lock_sock(sk);
 568         r->sg_start = md->sg_start;
 569         i = md->sg_start;
 570
 571         do {
 572                 size = (apply && apply_bytes < md->sg_data[i].length) ?
 573                         apply_bytes : md->sg_data[i].length;
 574
 575                 if (!sk_wmem_schedule(sk, size)) {
 576                         if (!copied)
 577                                 err = -ENOMEM;
 578                         break;
 579                 }
 580
 581                 sk_mem_charge(sk, size);
 582                 r->sg_data[i] = md->sg_data[i];
 583                 r->sg_data[i].length = size;
 584                 md->sg_data[i].length -= size;
 585                 md->sg_data[i].offset += size;
 586                 copied += size;
 587
 588                 if (md->sg_data[i].length) {
 589                         get_page(sg_page(&r->sg_data[i]));
 590                         r->sg_end = (i + 1) == MAX_SKB_FRAGS ? 0 : i + 1;
 591                 } else {
 592                         i++;
 593                         if (i == MAX_SKB_FRAGS)
 594                                 i = 0;
 595                         r->sg_end = i;
 596                 }
 597
 598                 if (apply) {
 599                         apply_bytes -= size;
 600                         if (!apply_bytes)
 601                                 break;
 602                 }
 603         } while (i != md->sg_end);
 604
 605         md->sg_start = i;
 606
 607         if (!err) {
 608                 list_add_tail(&r->list, &psock->ingress);
 609                 sk->sk_data_ready(sk);
 610         } else {
 611                 free_start_sg(sk, r);
 612                 kfree(r);
 613         }
 614
 615         release_sock(sk);
 616         return err;
 617 }
 618
 619 static int bpf_tcp_sendmsg_do_redirect(struct sock *sk, int send,
 620                                        struct sk_msg_buff *md,
 621                                        int flags)
 622 {
 623         bool ingress = !!(md->flags & BPF_F_INGRESS);
 624         struct smap_psock *psock;
 625         struct scatterlist *sg;
 626         int err = 0;
 627
 628         sg = md->sg_data;
 629
 630         rcu_read_lock();
 631         psock = smap_psock_sk(sk);
 632         if (unlikely(!psock))
 633                 goto out_rcu;
 634
 635         if (!refcount_inc_not_zero(&psock->refcnt))
 636                 goto out_rcu;
 637
 638         rcu_read_unlock();
 639
 640         if (ingress) {
 641                 err = bpf_tcp_ingress(sk, send, psock, md, flags);
 642         } else {
 643                 lock_sock(sk);
 644                 err = bpf_tcp_push(sk, send, md, flags, false);
 645                 release_sock(sk);
 646         }
 647         smap_release_sock(psock, sk);
 648         if (unlikely(err))
 649                 goto out;
 650         return 0;
 651 out_rcu:
 652         rcu_read_unlock();
 653 out:
 654         free_bytes_sg(NULL, send, md, false);
 655         return err;
 656 }
 657
 658 static inline void bpf_md_init(struct smap_psock *psock)
 659 {
 660         if (!psock->apply_bytes) {
 661                 psock->eval =  __SK_NONE;
 662                 if (psock->sk_redir) {
 663                         sock_put(psock->sk_redir);
 664                         psock->sk_redir = NULL;
 665                 }
 666         }
 667 }
 668
 669 static void apply_bytes_dec(struct smap_psock *psock, int i)
 670 {
 671         if (psock->apply_bytes) {
 672                 if (psock->apply_bytes < i)
 673                         psock->apply_bytes = 0;
 674                 else
 675                         psock->apply_bytes -= i;
 676         }
 677 }
 678
 679 static int bpf_exec_tx_verdict(struct smap_psock *psock,
 680                                struct sk_msg_buff *m,
 681                                struct sock *sk,
 682                                int *copied, int flags)
 683 {
 684         bool cork = false, enospc = (m->sg_start == m->sg_end);
 685         struct sock *redir;
 686         int err = 0;
 687         int send;
 688
 689 more_data:
 690         if (psock->eval == __SK_NONE)
 691                 psock->eval = smap_do_tx_msg(sk, psock, m);
 692
 693         if (m->cork_bytes &&
 694             m->cork_bytes > psock->sg_size && !enospc) {
 695                 psock->cork_bytes = m->cork_bytes - psock->sg_size;
 696                 if (!psock->cork) {
 697                         psock->cork = kcalloc(1,
 698                                         sizeof(struct sk_msg_buff),
 699                                         GFP_ATOMIC | __GFP_NOWARN);
 700
 701                         if (!psock->cork) {
 702                                 err = -ENOMEM;
 703                                 goto out_err;
 704                         }
 705                 }
 706                 memcpy(psock->cork, m, sizeof(*m));
 707                 goto out_err;
 708         }
 709
 710         send = psock->sg_size;
 711         if (psock->apply_bytes && psock->apply_bytes < send)
 712                 send = psock->apply_bytes;
 713
 714         switch (psock->eval) {
 715         case __SK_PASS:
 716                 err = bpf_tcp_push(sk, send, m, flags, true);
 717                 if (unlikely(err)) {
 718                         *copied -= free_start_sg(sk, m);
 719                         break;
 720                 }
 721
 722                 apply_bytes_dec(psock, send);
 723                 psock->sg_size -= send;
 724                 break;
 725         case __SK_REDIRECT:
 726                 redir = psock->sk_redir;
 727                 apply_bytes_dec(psock, send);
 728
 729                 if (psock->cork) {
 730                         cork = true;
 731                         psock->cork = NULL;
 732                 }
 733
 734                 return_mem_sg(sk, send, m);
 735                 release_sock(sk);
 736
 737                 err = bpf_tcp_sendmsg_do_redirect(redir, send, m, flags);
 738                 lock_sock(sk);
 739
 740                 if (unlikely(err < 0)) {
 741                         free_start_sg(sk, m);
 742                         psock->sg_size = 0;
 743                         if (!cork)
 744                                 *copied -= send;
 745                 } else {
 746                         psock->sg_size -= send;
 747                 }
 748
 749                 if (cork) {
 750                         free_start_sg(sk, m);
 751                         psock->sg_size = 0;
 752                         kfree(m);
 753                         m = NULL;
 754                         err = 0;
 755                 }
 756                 break;
 757         case __SK_DROP:
 758         default:
 759                 free_bytes_sg(sk, send, m, true);
 760                 apply_bytes_dec(psock, send);
 761                 *copied -= send;
 762                 psock->sg_size -= send;
 763                 err = -EACCES;
 764                 break;
 765         }
 766
 767         if (likely(!err)) {
 768                 bpf_md_init(psock);
 769                 if (m &&
 770                     m->sg_data[m->sg_start].page_link &&
 771                     m->sg_data[m->sg_start].length)
 772                         goto more_data;
 773         }
 774
 775 out_err:
 776         return err;
 777 }
 778
 779 static int bpf_wait_data(struct sock *sk,
 780                          struct smap_psock *psk, int flags,
 781                          long timeo, int *err)
 782 {
 783         int rc;
 784
 785         DEFINE_WAIT_FUNC(wait, woken_wake_function);
 786
 787         add_wait_queue(sk_sleep(sk), &wait);
 788         sk_set_bit(SOCKWQ_ASYNC_WAITDATA, sk);
 789         rc = sk_wait_event(sk, &timeo,
 790                            !list_empty(&psk->ingress) ||
 791                            !skb_queue_empty(&sk->sk_receive_queue),
 792                            &wait);
 793         sk_clear_bit(SOCKWQ_ASYNC_WAITDATA, sk);
 794         remove_wait_queue(sk_sleep(sk), &wait);
 795
 796         return rc;
 797 }
 798
 799 static int bpf_tcp_recvmsg(struct sock *sk, struct msghdr *msg, size_t len,
 800                            int nonblock, int flags, int *addr_len)
 801 {
 802         struct iov_iter *iter = &msg->msg_iter;
 803         struct smap_psock *psock;
 804         int copied = 0;
 805
 806         if (unlikely(flags & MSG_ERRQUEUE))
 807                 return inet_recv_error(sk, msg, len, addr_len);
 808
 809         rcu_read_lock();
 810         psock = smap_psock_sk(sk);
 811         if (unlikely(!psock))
 812                 goto out;
 813
 814         if (unlikely(!refcount_inc_not_zero(&psock->refcnt)))
 815                 goto out;
 816         rcu_read_unlock();
 817
 818         if (!skb_queue_empty(&sk->sk_receive_queue))
 819                 return tcp_recvmsg(sk, msg, len, nonblock, flags, addr_len);
 820
 821         lock_sock(sk);
 822 bytes_ready:
 823         while (copied != len) {
 824                 struct scatterlist *sg;
 825                 struct sk_msg_buff *md;
 826                 int i;
 827
 828                 md = list_first_entry_or_null(&psock->ingress,
 829                                               struct sk_msg_buff, list);
 830                 if (unlikely(!md))
 831                         break;
 832                 i = md->sg_start;
 833                 do {
 834                         struct page *page;
 835                         int n, copy;
 836
 837                         sg = &md->sg_data[i];
 838                         copy = sg->length;
 839                         page = sg_page(sg);
 840
 841                         if (copied + copy > len)
 842                                 copy = len - copied;
 843
 844                         n = copy_page_to_iter(page, sg->offset, copy, iter);
 845                         if (n != copy) {
 846                                 md->sg_start = i;
 847                                 release_sock(sk);
 848                                 smap_release_sock(psock, sk);
 849                                 return -EFAULT;
 850                         }
 851
 852                         copied += copy;
 853                         sg->offset += copy;
 854                         sg->length -= copy;
 855                         sk_mem_uncharge(sk, copy);
 856
 857                         if (!sg->length) {
 858                                 i++;
 859                                 if (i == MAX_SKB_FRAGS)
 860                                         i = 0;
 861                                 if (!md->skb)
 862                                         put_page(page);
 863                         }
 864                         if (copied == len)
 865                                 break;
 866                 } while (i != md->sg_end);
 867                 md->sg_start = i;
 868
 869                 if (!sg->length && md->sg_start == md->sg_end) {
 870                         list_del(&md->list);
 871                         if (md->skb)
 872                                 consume_skb(md->skb);
 873                         kfree(md);
 874                 }
 875         }
 876
 877         if (!copied) {
 878                 long timeo;
 879                 int data;
 880                 int err = 0;
 881
 882                 timeo = sock_rcvtimeo(sk, nonblock);
 883                 data = bpf_wait_data(sk, psock, flags, timeo, &err);
 884
 885                 if (data) {
 886                         if (!skb_queue_empty(&sk->sk_receive_queue)) {
 887                                 release_sock(sk);
 888                                 smap_release_sock(psock, sk);
 889                                 copied = tcp_recvmsg(sk, msg, len, nonblock, flags, addr_len);
 890                                 return copied;
 891                         }
 892                         goto bytes_ready;
 893                 }
 894
 895                 if (err)
 896                         copied = err;
 897         }
 898
 899         release_sock(sk);
 900         smap_release_sock(psock, sk);
 901         return copied;
 902 out:
 903         rcu_read_unlock();
 904         return tcp_recvmsg(sk, msg, len, nonblock, flags, addr_len);
 905 }
 906
 907
 908 static int bpf_tcp_sendmsg(struct sock *sk, struct msghdr *msg, size_t size)
 909 {
 910         int flags = msg->msg_flags | MSG_NO_SHARED_FRAGS;
 911         struct sk_msg_buff md = {0};
 912         unsigned int sg_copy = 0;
 913         struct smap_psock *psock;
 914         int copied = 0, err = 0;
 915         struct scatterlist *sg;
 916         long timeo;
 917
 918         /* Its possible a sock event or user removed the psock _but_ the ops
 919          * have not been reprogrammed yet so we get here. In this case fallback
 920          * to tcp_sendmsg. Note this only works because we _only_ ever allow
 921          * a single ULP there is no hierarchy here.
 922          */
 923         rcu_read_lock();
 924         psock = smap_psock_sk(sk);
 925         if (unlikely(!psock)) {
 926                 rcu_read_unlock();
 927                 return tcp_sendmsg(sk, msg, size);
 928         }
 929
 930         /* Increment the psock refcnt to ensure its not released while sending a
 931          * message. Required because sk lookup and bpf programs are used in
 932          * separate rcu critical sections. Its OK if we lose the map entry
 933          * but we can't lose the sock reference.
 934          */
 935         if (!refcount_inc_not_zero(&psock->refcnt)) {
 936                 rcu_read_unlock();
 937                 return tcp_sendmsg(sk, msg, size);
 938         }
 939
 940         sg = md.sg_data;
 941         sg_init_marker(sg, MAX_SKB_FRAGS);
 942         rcu_read_unlock();
 943
 944         lock_sock(sk);
 945         timeo = sock_sndtimeo(sk, msg->msg_flags & MSG_DONTWAIT);
 946
 947         while (msg_data_left(msg)) {
 948                 struct sk_msg_buff *m;
 949                 bool enospc = false;
 950                 int copy;
 951
 952                 if (sk->sk_err) {
 953                         err = sk->sk_err;
 954                         goto out_err;
 955                 }
 956
 957                 copy = msg_data_left(msg);
 958                 if (!sk_stream_memory_free(sk))
 959                         goto wait_for_sndbuf;
 960
 961                 m = psock->cork_bytes ? psock->cork : &md;
 962                 m->sg_curr = m->sg_copybreak ? m->sg_curr : m->sg_end;
 963                 err = sk_alloc_sg(sk, copy, m->sg_data,
 964                                   m->sg_start, &m->sg_end, &sg_copy,
 965                                   m->sg_end - 1);
 966                 if (err) {
 967                         if (err != -ENOSPC)
 968                                 goto wait_for_memory;
 969                         enospc = true;
 970                         copy = sg_copy;
 971                 }
 972
 973                 err = memcopy_from_iter(sk, m, &msg->msg_iter, copy);
 974                 if (err < 0) {
 975                         free_curr_sg(sk, m);
 976                         goto out_err;
 977                 }
 978
 979                 psock->sg_size += copy;
 980                 copied += copy;
 981                 sg_copy = 0;
 982
 983                 /* When bytes are being corked skip running BPF program and
 984                  * applying verdict unless there is no more buffer space. In
 985                  * the ENOSPC case simply run BPF prorgram with currently
 986                  * accumulated data. We don't have much choice at this point
 987                  * we could try extending the page frags or chaining complex
 988                  * frags but even in these cases _eventually_ we will hit an
 989                  * OOM scenario. More complex recovery schemes may be
 990                  * implemented in the future, but BPF programs must handle
 991                  * the case where apply_cork requests are not honored. The
 992                  * canonical method to verify this is to check data length.
 993                  */
 994                 if (psock->cork_bytes) {
 995                         if (copy > psock->cork_bytes)
 996                                 psock->cork_bytes = 0;
 997                         else
 998                                 psock->cork_bytes -= copy;
 999
1000                         if (psock->cork_bytes && !enospc)
1001                                 goto out_cork;
1002
1003                         /* All cork bytes accounted for re-run filter */
1004                         psock->eval = __SK_NONE;
1005                         psock->cork_bytes = 0;
1006                 }
1007
1008                 err = bpf_exec_tx_verdict(psock, m, sk, &copied, flags);
1009                 if (unlikely(err < 0))
1010                         goto out_err;
1011                 continue;
1012 wait_for_sndbuf:
1013                 set_bit(SOCK_NOSPACE, &sk->sk_socket->flags);
1014 wait_for_memory:
1015                 err = sk_stream_wait_memory(sk, &timeo);
1016                 if (err)
1017                         goto out_err;
1018         }
1019 out_err:
1020         if (err < 0)
1021                 err = sk_stream_error(sk, msg->msg_flags, err);
1022 out_cork:
1023         release_sock(sk);
1024         smap_release_sock(psock, sk);
1025         return copied ? copied : err;
1026 }
1027
1028 static int bpf_tcp_sendpage(struct sock *sk, struct page *page,
1029                             int offset, size_t size, int flags)
1030 {
1031         struct sk_msg_buff md = {0}, *m = NULL;
1032         int err = 0, copied = 0;
1033         struct smap_psock *psock;
1034         struct scatterlist *sg;
1035         bool enospc = false;
1036
1037         rcu_read_lock();
1038         psock = smap_psock_sk(sk);
1039         if (unlikely(!psock))
1040                 goto accept;
1041
1042         if (!refcount_inc_not_zero(&psock->refcnt))
1043                 goto accept;
1044         rcu_read_unlock();
1045
1046         lock_sock(sk);
1047
1048         if (psock->cork_bytes) {
1049                 m = psock->cork;
1050                 sg = &m->sg_data[m->sg_end];
1051         } else {
1052                 m = &md;
1053                 sg = m->sg_data;
1054                 sg_init_marker(sg, MAX_SKB_FRAGS);
1055         }
1056
1057         /* Catch case where ring is full and sendpage is stalled. */
1058         if (unlikely(m->sg_end == m->sg_start &&
1059             m->sg_data[m->sg_end].length))
1060                 goto out_err;
1061
1062         psock->sg_size += size;
1063         sg_set_page(sg, page, size, offset);
1064         get_page(page);
1065         m->sg_copy[m->sg_end] = true;
1066         sk_mem_charge(sk, size);
1067         m->sg_end++;
1068         copied = size;
1069
1070         if (m->sg_end == MAX_SKB_FRAGS)
1071                 m->sg_end = 0;
1072
1073         if (m->sg_end == m->sg_start)
1074                 enospc = true;
1075
1076         if (psock->cork_bytes) {
1077                 if (size > psock->cork_bytes)
1078                         psock->cork_bytes = 0;
1079                 else
1080                         psock->cork_bytes -= size;
1081
1082                 if (psock->cork_bytes && !enospc)
1083                         goto out_err;
1084
1085                 /* All cork bytes accounted for re-run filter */
1086                 psock->eval = __SK_NONE;
1087                 psock->cork_bytes = 0;
1088         }
1089
1090         err = bpf_exec_tx_verdict(psock, m, sk, &copied, flags);
1091 out_err:
1092         release_sock(sk);
1093         smap_release_sock(psock, sk);
1094         return copied ? copied : err;
1095 accept:
1096         rcu_read_unlock();
1097         return tcp_sendpage(sk, page, offset, size, flags);
1098 }
1099
1100 static void bpf_tcp_msg_add(struct smap_psock *psock,
1101                             struct sock *sk,
1102                             struct bpf_prog *tx_msg)
1103 {
1104         struct bpf_prog *orig_tx_msg;
1105
1106         orig_tx_msg = xchg(&psock->bpf_tx_msg, tx_msg);
1107         if (orig_tx_msg)
1108                 bpf_prog_put(orig_tx_msg);
1109 }
1110
1111 static int bpf_tcp_ulp_register(void)
1112 {
1113         tcp_bpf_proto = tcp_prot;
1114         tcp_bpf_proto.close = bpf_tcp_close;
1115         /* Once BPF TX ULP is registered it is never unregistered. It
1116          * will be in the ULP list for the lifetime of the system. Doing
1117          * duplicate registers is not a problem.
1118          */
1119         return tcp_register_ulp(&bpf_tcp_ulp_ops);
1120 }
1121
1122 static int smap_verdict_func(struct smap_psock *psock, struct sk_buff *skb)
1123 {
1124         struct bpf_prog *prog = READ_ONCE(psock->bpf_verdict);
1125         int rc;
1126
1127         if (unlikely(!prog))
1128                 return __SK_DROP;
1129
1130         skb_orphan(skb);
1131         /* We need to ensure that BPF metadata for maps is also cleared
1132          * when we orphan the skb so that we don't have the possibility
1133          * to reference a stale map.
1134          */
1135         TCP_SKB_CB(skb)->bpf.sk_redir = NULL;
1136         skb->sk = psock->sock;
1137         bpf_compute_data_pointers(skb);
1138         preempt_disable();
1139         rc = (*prog->bpf_func)(skb, prog->insnsi);
1140         preempt_enable();
1141         skb->sk = NULL;
1142
1143         /* Moving return codes from UAPI namespace into internal namespace */
1144         return rc == SK_PASS ?
1145                 (TCP_SKB_CB(skb)->bpf.sk_redir ? __SK_REDIRECT : __SK_PASS) :
1146                 __SK_DROP;
1147 }
1148
1149 static int smap_do_ingress(struct smap_psock *psock, struct sk_buff *skb)
1150 {
1151         struct sock *sk = psock->sock;
1152         int copied = 0, num_sg;
1153         struct sk_msg_buff *r;
1154
1155         r = kzalloc(sizeof(struct sk_msg_buff), __GFP_NOWARN | GFP_ATOMIC);
1156         if (unlikely(!r))
1157                 return -EAGAIN;
1158
1159         if (!sk_rmem_schedule(sk, skb, skb->len)) {
1160                 kfree(r);
1161                 return -EAGAIN;
1162         }
1163
1164         sg_init_table(r->sg_data, MAX_SKB_FRAGS);
1165         num_sg = skb_to_sgvec(skb, r->sg_data, 0, skb->len);
1166         if (unlikely(num_sg < 0)) {
1167                 kfree(r);
1168                 return num_sg;
1169         }
1170         sk_mem_charge(sk, skb->len);
1171         copied = skb->len;
1172         r->sg_start = 0;
1173         r->sg_end = num_sg == MAX_SKB_FRAGS ? 0 : num_sg;
1174         r->skb = skb;
1175         list_add_tail(&r->list, &psock->ingress);
1176         sk->sk_data_ready(sk);
1177         return copied;
1178 }
1179
1180 static void smap_do_verdict(struct smap_psock *psock, struct sk_buff *skb)
1181 {
1182         struct smap_psock *peer;
1183         struct sock *sk;
1184         __u32 in;
1185         int rc;
1186
1187         rc = smap_verdict_func(psock, skb);
1188         switch (rc) {
1189         case __SK_REDIRECT:
1190                 sk = do_sk_redirect_map(skb);
1191                 if (!sk) {
1192                         kfree_skb(skb);
1193                         break;
1194                 }
1195
1196                 peer = smap_psock_sk(sk);
1197                 in = (TCP_SKB_CB(skb)->bpf.flags) & BPF_F_INGRESS;
1198
1199                 if (unlikely(!peer || sock_flag(sk, SOCK_DEAD) ||
1200                              !test_bit(SMAP_TX_RUNNING, &peer->state))) {
1201                         kfree_skb(skb);
1202                         break;
1203                 }
1204
1205                 if (!in && sock_writeable(sk)) {
1206                         skb_set_owner_w(skb, sk);
1207                         skb_queue_tail(&peer->rxqueue, skb);
1208                         schedule_work(&peer->tx_work);
1209                         break;
1210                 } else if (in &&
1211                            atomic_read(&sk->sk_rmem_alloc) <= sk->sk_rcvbuf) {
1212                         skb_queue_tail(&peer->rxqueue, skb);
1213                         schedule_work(&peer->tx_work);
1214                         break;
1215                 }
1216         /* Fall through and free skb otherwise */
1217         case __SK_DROP:
1218         default:
1219                 kfree_skb(skb);
1220         }
1221 }
1222
1223 static void smap_report_sk_error(struct smap_psock *psock, int err)
1224 {
1225         struct sock *sk = psock->sock;
1226
1227         sk->sk_err = err;
1228         sk->sk_error_report(sk);
1229 }
1230
1231 static void smap_read_sock_strparser(struct strparser *strp,
1232                                      struct sk_buff *skb)
1233 {
1234         struct smap_psock *psock;
1235
1236         rcu_read_lock();
1237         psock = container_of(strp, struct smap_psock, strp);
1238         smap_do_verdict(psock, skb);
1239         rcu_read_unlock();
1240 }
1241
1242 /* Called with lock held on socket */
1243 static void smap_data_ready(struct sock *sk)
1244 {
1245         struct smap_psock *psock;
1246
1247         rcu_read_lock();
1248         psock = smap_psock_sk(sk);
1249         if (likely(psock)) {
1250                 write_lock_bh(&sk->sk_callback_lock);
1251                 strp_data_ready(&psock->strp);
1252                 write_unlock_bh(&sk->sk_callback_lock);
1253         }
1254         rcu_read_unlock();
1255 }
1256
1257 static void smap_tx_work(struct work_struct *w)
1258 {
1259         struct smap_psock *psock;
1260         struct sk_buff *skb;
1261         int rem, off, n;
1262
1263         psock = container_of(w, struct smap_psock, tx_work);
1264
1265         /* lock sock to avoid losing sk_socket at some point during loop */
1266         lock_sock(psock->sock);
1267         if (psock->save_skb) {
1268                 skb = psock->save_skb;
1269                 rem = psock->save_rem;
1270                 off = psock->save_off;
1271                 psock->save_skb = NULL;
1272                 goto start;
1273         }
1274
1275         while ((skb = skb_dequeue(&psock->rxqueue))) {
1276                 __u32 flags;
1277
1278                 rem = skb->len;
1279                 off = 0;
1280 start:
1281                 flags = (TCP_SKB_CB(skb)->bpf.flags) & BPF_F_INGRESS;
1282                 do {
1283                         if (likely(psock->sock->sk_socket)) {
1284                                 if (flags)
1285                                         n = smap_do_ingress(psock, skb);
1286                                 else
1287                                         n = skb_send_sock_locked(psock->sock,
1288                                                                  skb, off, rem);
1289                         } else {
1290                                 n = -EINVAL;
1291                         }
1292
1293                         if (n <= 0) {
1294                                 if (n == -EAGAIN) {
1295                                         /* Retry when space is available */
1296                                         psock->save_skb = skb;
1297                                         psock->save_rem = rem;
1298                                         psock->save_off = off;
1299                                         goto out;
1300                                 }
1301                                 /* Hard errors break pipe and stop xmit */
1302                                 smap_report_sk_error(psock, n ? -n : EPIPE);
1303                                 clear_bit(SMAP_TX_RUNNING, &psock->state);
1304                                 kfree_skb(skb);
1305                                 goto out;
1306                         }
1307                         rem -= n;
1308                         off += n;
1309                 } while (rem);
1310
1311                 if (!flags)
1312                         kfree_skb(skb);
1313         }
1314 out:
1315         release_sock(psock->sock);
1316 }
1317
1318 static void smap_write_space(struct sock *sk)
1319 {
1320         struct smap_psock *psock;
1321
1322         rcu_read_lock();
1323         psock = smap_psock_sk(sk);
1324         if (likely(psock && test_bit(SMAP_TX_RUNNING, &psock->state)))
1325                 schedule_work(&psock->tx_work);
1326         rcu_read_unlock();
1327 }
1328
1329 static void smap_stop_sock(struct smap_psock *psock, struct sock *sk)
1330 {
1331         if (!psock->strp_enabled)
1332                 return;
1333         sk->sk_data_ready = psock->save_data_ready;
1334         sk->sk_write_space = psock->save_write_space;
1335         psock->save_data_ready = NULL;
1336         psock->save_write_space = NULL;
1337         strp_stop(&psock->strp);
1338         psock->strp_enabled = false;
1339 }
1340
1341 static void smap_destroy_psock(struct rcu_head *rcu)
1342 {
1343         struct smap_psock *psock = container_of(rcu,
1344                                                   struct smap_psock, rcu);
1345
1346         /* Now that a grace period has passed there is no longer
1347          * any reference to this sock in the sockmap so we can
1348          * destroy the psock, strparser, and bpf programs. But,
1349          * because we use workqueue sync operations we can not
1350          * do it in rcu context
1351          */
1352         schedule_work(&psock->gc_work);
1353 }
1354
1355 static void smap_release_sock(struct smap_psock *psock, struct sock *sock)
1356 {
1357         if (refcount_dec_and_test(&psock->refcnt)) {
1358                 tcp_cleanup_ulp(sock);
1359                 smap_stop_sock(psock, sock);
1360                 clear_bit(SMAP_TX_RUNNING, &psock->state);
1361                 rcu_assign_sk_user_data(sock, NULL);
1362                 call_rcu_sched(&psock->rcu, smap_destroy_psock);
1363         }
1364 }
1365
1366 static int smap_parse_func_strparser(struct strparser *strp,
1367                                        struct sk_buff *skb)
1368 {
1369         struct smap_psock *psock;
1370         struct bpf_prog *prog;
1371         int rc;
1372
1373         rcu_read_lock();
1374         psock = container_of(strp, struct smap_psock, strp);
1375         prog = READ_ONCE(psock->bpf_parse);
1376
1377         if (unlikely(!prog)) {
1378                 rcu_read_unlock();
1379                 return skb->len;
1380         }
1381
1382         /* Attach socket for bpf program to use if needed we can do this
1383          * because strparser clones the skb before handing it to a upper
1384          * layer, meaning skb_orphan has been called. We NULL sk on the
1385          * way out to ensure we don't trigger a BUG_ON in skb/sk operations
1386          * later and because we are not charging the memory of this skb to
1387          * any socket yet.
1388          */
1389         skb->sk = psock->sock;
1390         bpf_compute_data_pointers(skb);
1391         rc = (*prog->bpf_func)(skb, prog->insnsi);
1392         skb->sk = NULL;
1393         rcu_read_unlock();
1394         return rc;
1395 }
1396
1397 static int smap_read_sock_done(struct strparser *strp, int err)
1398 {
1399         return err;
1400 }
1401
1402 static int smap_init_sock(struct smap_psock *psock,
1403                           struct sock *sk)
1404 {
1405         static const struct strp_callbacks cb = {
1406                 .rcv_msg = smap_read_sock_strparser,
1407                 .parse_msg = smap_parse_func_strparser,
1408                 .read_sock_done = smap_read_sock_done,
1409         };
1410
1411         return strp_init(&psock->strp, sk, &cb);
1412 }
1413
1414 static void smap_init_progs(struct smap_psock *psock,
1415                             struct bpf_prog *verdict,
1416                             struct bpf_prog *parse)
1417 {
1418         struct bpf_prog *orig_parse, *orig_verdict;
1419
1420         orig_parse = xchg(&psock->bpf_parse, parse);
1421         orig_verdict = xchg(&psock->bpf_verdict, verdict);
1422
1423         if (orig_verdict)
1424                 bpf_prog_put(orig_verdict);
1425         if (orig_parse)
1426                 bpf_prog_put(orig_parse);
1427 }
1428
1429 static void smap_start_sock(struct smap_psock *psock, struct sock *sk)
1430 {
1431         if (sk->sk_data_ready == smap_data_ready)
1432                 return;
1433         psock->save_data_ready = sk->sk_data_ready;
1434         psock->save_write_space = sk->sk_write_space;
1435         sk->sk_data_ready = smap_data_ready;
1436         sk->sk_write_space = smap_write_space;
1437         psock->strp_enabled = true;
1438 }
1439
1440 static void sock_map_remove_complete(struct bpf_stab *stab)
1441 {
1442         bpf_map_area_free(stab->sock_map);
1443         kfree(stab);
1444 }
1445
1446 static void smap_gc_work(struct work_struct *w)
1447 {
1448         struct smap_psock_map_entry *e, *tmp;
1449         struct sk_msg_buff *md, *mtmp;
1450         struct smap_psock *psock;
1451
1452         psock = container_of(w, struct smap_psock, gc_work);
1453
1454         /* no callback lock needed because we already detached sockmap ops */
1455         if (psock->strp_enabled)
1456                 strp_done(&psock->strp);
1457
1458         cancel_work_sync(&psock->tx_work);
1459         __skb_queue_purge(&psock->rxqueue);
1460
1461         /* At this point all strparser and xmit work must be complete */
1462         if (psock->bpf_parse)
1463                 bpf_prog_put(psock->bpf_parse);
1464         if (psock->bpf_verdict)
1465                 bpf_prog_put(psock->bpf_verdict);
1466         if (psock->bpf_tx_msg)
1467                 bpf_prog_put(psock->bpf_tx_msg);
1468
1469         if (psock->cork) {
1470                 free_start_sg(psock->sock, psock->cork);
1471                 kfree(psock->cork);
1472         }
1473
1474         list_for_each_entry_safe(md, mtmp, &psock->ingress, list) {
1475                 list_del(&md->list);
1476                 free_start_sg(psock->sock, md);
1477                 kfree(md);
1478         }
1479
1480         list_for_each_entry_safe(e, tmp, &psock->maps, list) {
1481                 list_del(&e->list);
1482                 kfree(e);
1483         }
1484
1485         if (psock->sk_redir)
1486                 sock_put(psock->sk_redir);
1487
1488         sock_put(psock->sock);
1489         kfree(psock);
1490 }
1491
1492 static struct smap_psock *smap_init_psock(struct sock *sock, int node)
1493 {
1494         struct smap_psock *psock;
1495
1496         psock = kzalloc_node(sizeof(struct smap_psock),
1497                              GFP_ATOMIC | __GFP_NOWARN,
1498                              node);
1499         if (!psock)
1500                 return ERR_PTR(-ENOMEM);
1501
1502         psock->eval =  __SK_NONE;
1503         psock->sock = sock;
1504         skb_queue_head_init(&psock->rxqueue);
1505         INIT_WORK(&psock->tx_work, smap_tx_work);
1506         INIT_WORK(&psock->gc_work, smap_gc_work);
1507         INIT_LIST_HEAD(&psock->maps);
1508         INIT_LIST_HEAD(&psock->ingress);
1509         refcount_set(&psock->refcnt, 1);
1510
1511         rcu_assign_sk_user_data(sock, psock);
1512         sock_hold(sock);
1513         return psock;
1514 }
1515
1516 static struct bpf_map *sock_map_alloc(union bpf_attr *attr)
1517 {
1518         struct bpf_stab *stab;
1519         u64 cost;
1520         int err;
1521
1522         if (!capable(CAP_NET_ADMIN))
1523                 return ERR_PTR(-EPERM);
1524
1525         /* check sanity of attributes */
1526         if (attr->max_entries == 0 || attr->key_size != 4 ||
1527             attr->value_size != 4 || attr->map_flags & ~SOCK_CREATE_FLAG_MASK)
1528                 return ERR_PTR(-EINVAL);
1529
1530         err = bpf_tcp_ulp_register();
1531         if (err && err != -EEXIST)
1532                 return ERR_PTR(err);
1533
1534         stab = kzalloc(sizeof(*stab), GFP_USER);
1535         if (!stab)
1536                 return ERR_PTR(-ENOMEM);
1537
1538         bpf_map_init_from_attr(&stab->map, attr);
1539
1540         /* make sure page count doesn't overflow */
1541         cost = (u64) stab->map.max_entries * sizeof(struct sock *);
1542         err = -EINVAL;
1543         if (cost >= U32_MAX - PAGE_SIZE)
1544                 goto free_stab;
1545
1546         stab->map.pages = round_up(cost, PAGE_SIZE) >> PAGE_SHIFT;
1547
1548         /* if map size is larger than memlock limit, reject it early */
1549         err = bpf_map_precharge_memlock(stab->map.pages);
1550         if (err)
1551                 goto free_stab;
1552
1553         err = -ENOMEM;
1554         stab->sock_map = bpf_map_area_alloc(stab->map.max_entries *
1555                                             sizeof(struct sock *),
1556                                             stab->map.numa_node);
1557         if (!stab->sock_map)
1558                 goto free_stab;
1559
1560         return &stab->map;
1561 free_stab:
1562         kfree(stab);
1563         return ERR_PTR(err);
1564 }
1565
1566 static void smap_list_remove(struct smap_psock *psock,
1567                              struct sock **entry,
1568                              struct htab_elem *hash_link)
1569 {
1570         struct smap_psock_map_entry *e, *tmp;
1571
1572         list_for_each_entry_safe(e, tmp, &psock->maps, list) {
1573                 if (e->entry == entry || e->hash_link == hash_link) {
1574                         list_del(&e->list);
1575                         break;
1576                 }
1577         }
1578 }
1579
1580 static void sock_map_free(struct bpf_map *map)
1581 {
1582         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1583         int i;
1584
1585         synchronize_rcu();
1586
1587         /* At this point no update, lookup or delete operations can happen.
1588          * However, be aware we can still get a socket state event updates,
1589          * and data ready callabacks that reference the psock from sk_user_data
1590          * Also psock worker threads are still in-flight. So smap_release_sock
1591          * will only free the psock after cancel_sync on the worker threads
1592          * and a grace period expire to ensure psock is really safe to remove.
1593          */
1594         rcu_read_lock();
1595         for (i = 0; i < stab->map.max_entries; i++) {
1596                 struct smap_psock *psock;
1597                 struct sock *sock;
1598
1599                 sock = xchg(&stab->sock_map[i], NULL);
1600                 if (!sock)
1601                         continue;
1602
1603                 write_lock_bh(&sock->sk_callback_lock);
1604                 psock = smap_psock_sk(sock);
1605                 /* This check handles a racing sock event that can get the
1606                  * sk_callback_lock before this case but after xchg happens
1607                  * causing the refcnt to hit zero and sock user data (psock)
1608                  * to be null and queued for garbage collection.
1609                  */
1610                 if (likely(psock)) {
1611                         smap_list_remove(psock, &stab->sock_map[i], NULL);
1612                         smap_release_sock(psock, sock);
1613                 }
1614                 write_unlock_bh(&sock->sk_callback_lock);
1615         }
1616         rcu_read_unlock();
1617
1618         sock_map_remove_complete(stab);
1619 }
1620
1621 static int sock_map_get_next_key(struct bpf_map *map, void *key, void *next_key)
1622 {
1623         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1624         u32 i = key ? *(u32 *)key : U32_MAX;
1625         u32 *next = (u32 *)next_key;
1626
1627         if (i >= stab->map.max_entries) {
1628                 *next = 0;
1629                 return 0;
1630         }
1631
1632         if (i == stab->map.max_entries - 1)
1633                 return -ENOENT;
1634
1635         *next = i + 1;
1636         return 0;
1637 }
1638
1639 struct sock  *__sock_map_lookup_elem(struct bpf_map *map, u32 key)
1640 {
1641         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1642
1643         if (key >= map->max_entries)
1644                 return NULL;
1645
1646         return READ_ONCE(stab->sock_map[key]);
1647 }
1648
1649 static int sock_map_delete_elem(struct bpf_map *map, void *key)
1650 {
1651         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1652         struct smap_psock *psock;
1653         int k = *(u32 *)key;
1654         struct sock *sock;
1655
1656         if (k >= map->max_entries)
1657                 return -EINVAL;
1658
1659         sock = xchg(&stab->sock_map[k], NULL);
1660         if (!sock)
1661                 return -EINVAL;
1662
1663         write_lock_bh(&sock->sk_callback_lock);
1664         psock = smap_psock_sk(sock);
1665         if (!psock)
1666                 goto out;
1667
1668         if (psock->bpf_parse)
1669                 smap_stop_sock(psock, sock);
1670         smap_list_remove(psock, &stab->sock_map[k], NULL);
1671         smap_release_sock(psock, sock);
1672 out:
1673         write_unlock_bh(&sock->sk_callback_lock);
1674         return 0;
1675 }
1676
1677 /* Locking notes: Concurrent updates, deletes, and lookups are allowed and are
1678  * done inside rcu critical sections. This ensures on updates that the psock
1679  * will not be released via smap_release_sock() until concurrent updates/deletes
1680  * complete. All operations operate on sock_map using cmpxchg and xchg
1681  * operations to ensure we do not get stale references. Any reads into the
1682  * map must be done with READ_ONCE() because of this.
1683  *
1684  * A psock is destroyed via call_rcu and after any worker threads are cancelled
1685  * and syncd so we are certain all references from the update/lookup/delete
1686  * operations as well as references in the data path are no longer in use.
1687  *
1688  * Psocks may exist in multiple maps, but only a single set of parse/verdict
1689  * programs may be inherited from the maps it belongs to. A reference count
1690  * is kept with the total number of references to the psock from all maps. The
1691  * psock will not be released until this reaches zero. The psock and sock
1692  * user data data use the sk_callback_lock to protect critical data structures
1693  * from concurrent access. This allows us to avoid two updates from modifying
1694  * the user data in sock and the lock is required anyways for modifying
1695  * callbacks, we simply increase its scope slightly.
1696  *
1697  * Rules to follow,
1698  *  - psock must always be read inside RCU critical section
1699  *  - sk_user_data must only be modified inside sk_callback_lock and read
1700  *    inside RCU critical section.
1701  *  - psock->maps list must only be read & modified inside sk_callback_lock
1702  *  - sock_map must use READ_ONCE and (cmp)xchg operations
1703  *  - BPF verdict/parse programs must use READ_ONCE and xchg operations
1704  */
1705
1706 static int __sock_map_ctx_update_elem(struct bpf_map *map,
1707                                       struct bpf_sock_progs *progs,
1708                                       struct sock *sock,
1709                                       struct sock **map_link,
1710                                       void *key)
1711 {
1712         struct bpf_prog *verdict, *parse, *tx_msg;
1713         struct smap_psock_map_entry *e = NULL;
1714         struct smap_psock *psock;
1715         bool new = false;
1716         int err;
1717
1718         /* 1. If sock map has BPF programs those will be inherited by the
1719          * sock being added. If the sock is already attached to BPF programs
1720          * this results in an error.
1721          */
1722         verdict = READ_ONCE(progs->bpf_verdict);
1723         parse = READ_ONCE(progs->bpf_parse);
1724         tx_msg = READ_ONCE(progs->bpf_tx_msg);
1725
1726         if (parse && verdict) {
1727                 /* bpf prog refcnt may be zero if a concurrent attach operation
1728                  * removes the program after the above READ_ONCE() but before
1729                  * we increment the refcnt. If this is the case abort with an
1730                  * error.
1731                  */
1732                 verdict = bpf_prog_inc_not_zero(verdict);
1733                 if (IS_ERR(verdict))
1734                         return PTR_ERR(verdict);
1735
1736                 parse = bpf_prog_inc_not_zero(parse);
1737                 if (IS_ERR(parse)) {
1738                         bpf_prog_put(verdict);
1739                         return PTR_ERR(parse);
1740                 }
1741         }
1742
1743         if (tx_msg) {
1744                 tx_msg = bpf_prog_inc_not_zero(tx_msg);
1745                 if (IS_ERR(tx_msg)) {
1746                         if (parse && verdict) {
1747                                 bpf_prog_put(parse);
1748                                 bpf_prog_put(verdict);
1749                         }
1750                         return PTR_ERR(tx_msg);
1751                 }
1752         }
1753
1754         write_lock_bh(&sock->sk_callback_lock);
1755         psock = smap_psock_sk(sock);
1756
1757         /* 2. Do not allow inheriting programs if psock exists and has
1758          * already inherited programs. This would create confusion on
1759          * which parser/verdict program is running. If no psock exists
1760          * create one. Inside sk_callback_lock to ensure concurrent create
1761          * doesn't update user data.
1762          */
1763         if (psock) {
1764                 if (READ_ONCE(psock->bpf_parse) && parse) {
1765                         err = -EBUSY;
1766                         goto out_progs;
1767                 }
1768                 if (READ_ONCE(psock->bpf_tx_msg) && tx_msg) {
1769                         err = -EBUSY;
1770                         goto out_progs;
1771                 }
1772                 if (!refcount_inc_not_zero(&psock->refcnt)) {
1773                         err = -EAGAIN;
1774                         goto out_progs;
1775                 }
1776         } else {
1777                 psock = smap_init_psock(sock, map->numa_node);
1778                 if (IS_ERR(psock)) {
1779                         err = PTR_ERR(psock);
1780                         goto out_progs;
1781                 }
1782
1783                 set_bit(SMAP_TX_RUNNING, &psock->state);
1784                 new = true;
1785         }
1786
1787         if (map_link) {
1788                 e = kzalloc(sizeof(*e), GFP_ATOMIC | __GFP_NOWARN);
1789                 if (!e) {
1790                         err = -ENOMEM;
1791                         goto out_progs;
1792                 }
1793         }
1794
1795         /* 3. At this point we have a reference to a valid psock that is
1796          * running. Attach any BPF programs needed.
1797          */
1798         if (tx_msg)
1799                 bpf_tcp_msg_add(psock, sock, tx_msg);
1800         if (new) {
1801                 err = tcp_set_ulp_id(sock, TCP_ULP_BPF);
1802                 if (err)
1803                         goto out_free;
1804         }
1805
1806         if (parse && verdict && !psock->strp_enabled) {
1807                 err = smap_init_sock(psock, sock);
1808                 if (err)
1809                         goto out_free;
1810                 smap_init_progs(psock, verdict, parse);
1811                 smap_start_sock(psock, sock);
1812         }
1813
1814         /* 4. Place psock in sockmap for use and stop any programs on
1815          * the old sock assuming its not the same sock we are replacing
1816          * it with. Because we can only have a single set of programs if
1817          * old_sock has a strp we can stop it.
1818          */
1819         if (map_link) {
1820                 e->entry = map_link;
1821                 list_add_tail(&e->list, &psock->maps);
1822         }
1823         write_unlock_bh(&sock->sk_callback_lock);
1824         return err;
1825 out_free:
1826         kfree(e);
1827         smap_release_sock(psock, sock);
1828 out_progs:
1829         if (parse && verdict) {
1830                 bpf_prog_put(parse);
1831                 bpf_prog_put(verdict);
1832         }
1833         if (tx_msg)
1834                 bpf_prog_put(tx_msg);
1835         write_unlock_bh(&sock->sk_callback_lock);
1836         kfree(e);
1837         return err;
1838 }
1839
1840 static int sock_map_ctx_update_elem(struct bpf_sock_ops_kern *skops,
1841                                     struct bpf_map *map,
1842                                     void *key, u64 flags)
1843 {
1844         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1845         struct bpf_sock_progs *progs = &stab->progs;
1846         struct sock *osock, *sock;
1847         u32 i = *(u32 *)key;
1848         int err;
1849
1850         if (unlikely(flags > BPF_EXIST))
1851                 return -EINVAL;
1852
1853         if (unlikely(i >= stab->map.max_entries))
1854                 return -E2BIG;
1855
1856         sock = READ_ONCE(stab->sock_map[i]);
1857         if (flags == BPF_EXIST && !sock)
1858                 return -ENOENT;
1859         else if (flags == BPF_NOEXIST && sock)
1860                 return -EEXIST;
1861
1862         sock = skops->sk;
1863         err = __sock_map_ctx_update_elem(map, progs, sock, &stab->sock_map[i],
1864                                          key);
1865         if (err)
1866                 goto out;
1867
1868         osock = xchg(&stab->sock_map[i], sock);
1869         if (osock) {
1870                 struct smap_psock *opsock = smap_psock_sk(osock);
1871
1872                 write_lock_bh(&osock->sk_callback_lock);
1873                 smap_list_remove(opsock, &stab->sock_map[i], NULL);
1874                 smap_release_sock(opsock, osock);
1875                 write_unlock_bh(&osock->sk_callback_lock);
1876         }
1877 out:
1878         return err;
1879 }
1880
1881 int sock_map_prog(struct bpf_map *map, struct bpf_prog *prog, u32 type)
1882 {
1883         struct bpf_sock_progs *progs;
1884         struct bpf_prog *orig;
1885
1886         if (map->map_type == BPF_MAP_TYPE_SOCKMAP) {
1887                 struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1888
1889                 progs = &stab->progs;
1890         } else if (map->map_type == BPF_MAP_TYPE_SOCKHASH) {
1891                 struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
1892
1893                 progs = &htab->progs;
1894         } else {
1895                 return -EINVAL;
1896         }
1897
1898         switch (type) {
1899         case BPF_SK_MSG_VERDICT:
1900                 orig = xchg(&progs->bpf_tx_msg, prog);
1901                 break;
1902         case BPF_SK_SKB_STREAM_PARSER:
1903                 orig = xchg(&progs->bpf_parse, prog);
1904                 break;
1905         case BPF_SK_SKB_STREAM_VERDICT:
1906                 orig = xchg(&progs->bpf_verdict, prog);
1907                 break;
1908         default:
1909                 return -EOPNOTSUPP;
1910         }
1911
1912         if (orig)
1913                 bpf_prog_put(orig);
1914
1915         return 0;
1916 }
1917
1918 static void *sock_map_lookup(struct bpf_map *map, void *key)
1919 {
1920         return NULL;
1921 }
1922
1923 static int sock_map_update_elem(struct bpf_map *map,
1924                                 void *key, void *value, u64 flags)
1925 {
1926         struct bpf_sock_ops_kern skops;
1927         u32 fd = *(u32 *)value;
1928         struct socket *socket;
1929         int err;
1930
1931         socket = sockfd_lookup(fd, &err);
1932         if (!socket)
1933                 return err;
1934
1935         skops.sk = socket->sk;
1936         if (!skops.sk) {
1937                 fput(socket->file);
1938                 return -EINVAL;
1939         }
1940
1941         if (skops.sk->sk_type != SOCK_STREAM ||
1942             skops.sk->sk_protocol != IPPROTO_TCP) {
1943                 fput(socket->file);
1944                 return -EOPNOTSUPP;
1945         }
1946
1947         err = sock_map_ctx_update_elem(&skops, map, key, flags);
1948         fput(socket->file);
1949         return err;
1950 }
1951
1952 static void sock_map_release(struct bpf_map *map)
1953 {
1954         struct bpf_sock_progs *progs;
1955         struct bpf_prog *orig;
1956
1957         if (map->map_type == BPF_MAP_TYPE_SOCKMAP) {
1958                 struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
1959
1960                 progs = &stab->progs;
1961         } else {
1962                 struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
1963
1964                 progs = &htab->progs;
1965         }
1966
1967         orig = xchg(&progs->bpf_parse, NULL);
1968         if (orig)
1969                 bpf_prog_put(orig);
1970         orig = xchg(&progs->bpf_verdict, NULL);
1971         if (orig)
1972                 bpf_prog_put(orig);
1973
1974         orig = xchg(&progs->bpf_tx_msg, NULL);
1975         if (orig)
1976                 bpf_prog_put(orig);
1977 }
1978
1979 static struct bpf_map *sock_hash_alloc(union bpf_attr *attr)
1980 {
1981         struct bpf_htab *htab;
1982         int i, err;
1983         u64 cost;
1984
1985         if (!capable(CAP_NET_ADMIN))
1986                 return ERR_PTR(-EPERM);
1987
1988         /* check sanity of attributes */
1989         if (attr->max_entries == 0 || attr->value_size != 4 ||
1990             attr->map_flags & ~SOCK_CREATE_FLAG_MASK)
1991                 return ERR_PTR(-EINVAL);
1992
1993         if (attr->key_size > MAX_BPF_STACK)
1994                 /* eBPF programs initialize keys on stack, so they cannot be
1995                  * larger than max stack size
1996                  */
1997                 return ERR_PTR(-E2BIG);
1998
1999         err = bpf_tcp_ulp_register();
2000         if (err && err != -EEXIST)
2001                 return ERR_PTR(err);
2002
2003         htab = kzalloc(sizeof(*htab), GFP_USER);
2004         if (!htab)
2005                 return ERR_PTR(-ENOMEM);
2006
2007         bpf_map_init_from_attr(&htab->map, attr);
2008
2009         htab->n_buckets = roundup_pow_of_two(htab->map.max_entries);
2010         htab->elem_size = sizeof(struct htab_elem) +
2011                           round_up(htab->map.key_size, 8);
2012         err = -EINVAL;
2013         if (htab->n_buckets == 0 ||
2014             htab->n_buckets > U32_MAX / sizeof(struct bucket))
2015                 goto free_htab;
2016
2017         cost = (u64) htab->n_buckets * sizeof(struct bucket) +
2018                (u64) htab->elem_size * htab->map.max_entries;
2019
2020         if (cost >= U32_MAX - PAGE_SIZE)
2021                 goto free_htab;
2022
2023         htab->map.pages = round_up(cost, PAGE_SIZE) >> PAGE_SHIFT;
2024         err = bpf_map_precharge_memlock(htab->map.pages);
2025         if (err)
2026                 goto free_htab;
2027
2028         err = -ENOMEM;
2029         htab->buckets = bpf_map_area_alloc(
2030                                 htab->n_buckets * sizeof(struct bucket),
2031                                 htab->map.numa_node);
2032         if (!htab->buckets)
2033                 goto free_htab;
2034
2035         for (i = 0; i < htab->n_buckets; i++) {
2036                 INIT_HLIST_HEAD(&htab->buckets[i].head);
2037                 raw_spin_lock_init(&htab->buckets[i].lock);
2038         }
2039
2040         return &htab->map;
2041 free_htab:
2042         kfree(htab);
2043         return ERR_PTR(err);
2044 }
2045
2046 static inline struct bucket *__select_bucket(struct bpf_htab *htab, u32 hash)
2047 {
2048         return &htab->buckets[hash & (htab->n_buckets - 1)];
2049 }
2050
2051 static inline struct hlist_head *select_bucket(struct bpf_htab *htab, u32 hash)
2052 {
2053         return &__select_bucket(htab, hash)->head;
2054 }
2055
2056 static void sock_hash_free(struct bpf_map *map)
2057 {
2058         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
2059         int i;
2060
2061         synchronize_rcu();
2062
2063         /* At this point no update, lookup or delete operations can happen.
2064          * However, be aware we can still get a socket state event updates,
2065          * and data ready callabacks that reference the psock from sk_user_data
2066          * Also psock worker threads are still in-flight. So smap_release_sock
2067          * will only free the psock after cancel_sync on the worker threads
2068          * and a grace period expire to ensure psock is really safe to remove.
2069          */
2070         rcu_read_lock();
2071         for (i = 0; i < htab->n_buckets; i++) {
2072                 struct hlist_head *head = select_bucket(htab, i);
2073                 struct hlist_node *n;
2074                 struct htab_elem *l;
2075
2076                 hlist_for_each_entry_safe(l, n, head, hash_node) {
2077                         struct sock *sock = l->sk;
2078                         struct smap_psock *psock;
2079
2080                         hlist_del_rcu(&l->hash_node);
2081                         write_lock_bh(&sock->sk_callback_lock);
2082                         psock = smap_psock_sk(sock);
2083                         /* This check handles a racing sock event that can get
2084                          * the sk_callback_lock before this case but after xchg
2085                          * causing the refcnt to hit zero and sock user data
2086                          * (psock) to be null and queued for garbage collection.
2087                          */
2088                         if (likely(psock)) {
2089                                 smap_list_remove(psock, NULL, l);
2090                                 smap_release_sock(psock, sock);
2091                         }
2092                         write_unlock_bh(&sock->sk_callback_lock);
2093                         kfree(l);
2094                 }
2095         }
2096         rcu_read_unlock();
2097         bpf_map_area_free(htab->buckets);
2098         kfree(htab);
2099 }
2100
2101 static struct htab_elem *alloc_sock_hash_elem(struct bpf_htab *htab,
2102                                               void *key, u32 key_size, u32 hash,
2103                                               struct sock *sk,
2104                                               struct htab_elem *old_elem)
2105 {
2106         struct htab_elem *l_new;
2107
2108         if (atomic_inc_return(&htab->count) > htab->map.max_entries) {
2109                 if (!old_elem) {
2110                         atomic_dec(&htab->count);
2111                         return ERR_PTR(-E2BIG);
2112                 }
2113         }
2114         l_new = kmalloc_node(htab->elem_size, GFP_ATOMIC | __GFP_NOWARN,
2115                              htab->map.numa_node);
2116         if (!l_new)
2117                 return ERR_PTR(-ENOMEM);
2118
2119         memcpy(l_new->key, key, key_size);
2120         l_new->sk = sk;
2121         l_new->hash = hash;
2122         return l_new;
2123 }
2124
2125 static struct htab_elem *lookup_elem_raw(struct hlist_head *head,
2126                                          u32 hash, void *key, u32 key_size)
2127 {
2128         struct htab_elem *l;
2129
2130         hlist_for_each_entry_rcu(l, head, hash_node) {
2131                 if (l->hash == hash && !memcmp(&l->key, key, key_size))
2132                         return l;
2133         }
2134
2135         return NULL;
2136 }
2137
2138 static inline u32 htab_map_hash(const void *key, u32 key_len)
2139 {
2140         return jhash(key, key_len, 0);
2141 }
2142
2143 static int sock_hash_get_next_key(struct bpf_map *map,
2144                                   void *key, void *next_key)
2145 {
2146         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
2147         struct htab_elem *l, *next_l;
2148         struct hlist_head *h;
2149         u32 hash, key_size;
2150         int i = 0;
2151
2152         WARN_ON_ONCE(!rcu_read_lock_held());
2153
2154         key_size = map->key_size;
2155         if (!key)
2156                 goto find_first_elem;
2157         hash = htab_map_hash(key, key_size);
2158         h = select_bucket(htab, hash);
2159
2160         l = lookup_elem_raw(h, hash, key, key_size);
2161         if (!l)
2162                 goto find_first_elem;
2163         next_l = hlist_entry_safe(
2164                      rcu_dereference_raw(hlist_next_rcu(&l->hash_node)),
2165                      struct htab_elem, hash_node);
2166         if (next_l) {
2167                 memcpy(next_key, next_l->key, key_size);
2168                 return 0;
2169         }
2170
2171         /* no more elements in this hash list, go to the next bucket */
2172         i = hash & (htab->n_buckets - 1);
2173         i++;
2174
2175 find_first_elem:
2176         /* iterate over buckets */
2177         for (; i < htab->n_buckets; i++) {
2178                 h = select_bucket(htab, i);
2179
2180                 /* pick first element in the bucket */
2181                 next_l = hlist_entry_safe(
2182                                 rcu_dereference_raw(hlist_first_rcu(h)),
2183                                 struct htab_elem, hash_node);
2184                 if (next_l) {
2185                         /* if it's not empty, just return it */
2186                         memcpy(next_key, next_l->key, key_size);
2187                         return 0;
2188                 }
2189         }
2190
2191         /* iterated over all buckets and all elements */
2192         return -ENOENT;
2193 }
2194
2195 static int sock_hash_ctx_update_elem(struct bpf_sock_ops_kern *skops,
2196                                      struct bpf_map *map,
2197                                      void *key, u64 map_flags)
2198 {
2199         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
2200         struct bpf_sock_progs *progs = &htab->progs;
2201         struct htab_elem *l_new = NULL, *l_old;
2202         struct smap_psock_map_entry *e = NULL;
2203         struct hlist_head *head;
2204         struct smap_psock *psock;
2205         u32 key_size, hash;
2206         struct sock *sock;
2207         struct bucket *b;
2208         int err;
2209
2210         sock = skops->sk;
2211
2212         if (sock->sk_type != SOCK_STREAM ||
2213             sock->sk_protocol != IPPROTO_TCP)
2214                 return -EOPNOTSUPP;
2215
2216         if (unlikely(map_flags > BPF_EXIST))
2217                 return -EINVAL;
2218
2219         e = kzalloc(sizeof(*e), GFP_ATOMIC | __GFP_NOWARN);
2220         if (!e)
2221                 return -ENOMEM;
2222
2223         WARN_ON_ONCE(!rcu_read_lock_held());
2224         key_size = map->key_size;
2225         hash = htab_map_hash(key, key_size);
2226         b = __select_bucket(htab, hash);
2227         head = &b->head;
2228
2229         err = __sock_map_ctx_update_elem(map, progs, sock, NULL, key);
2230         if (err)
2231                 goto err;
2232
2233         /* bpf_map_update_elem() can be called in_irq() */
2234         raw_spin_lock_bh(&b->lock);
2235         l_old = lookup_elem_raw(head, hash, key, key_size);
2236         if (l_old && map_flags == BPF_NOEXIST) {
2237                 err = -EEXIST;
2238                 goto bucket_err;
2239         }
2240         if (!l_old && map_flags == BPF_EXIST) {
2241                 err = -ENOENT;
2242                 goto bucket_err;
2243         }
2244
2245         l_new = alloc_sock_hash_elem(htab, key, key_size, hash, sock, l_old);
2246         if (IS_ERR(l_new)) {
2247                 err = PTR_ERR(l_new);
2248                 goto bucket_err;
2249         }
2250
2251         psock = smap_psock_sk(sock);
2252         if (unlikely(!psock)) {
2253                 err = -EINVAL;
2254                 goto bucket_err;
2255         }
2256
2257         e->hash_link = l_new;
2258         e->htab = container_of(map, struct bpf_htab, map);
2259         list_add_tail(&e->list, &psock->maps);
2260
2261         /* add new element to the head of the list, so that
2262          * concurrent search will find it before old elem
2263          */
2264         hlist_add_head_rcu(&l_new->hash_node, head);
2265         if (l_old) {
2266                 psock = smap_psock_sk(l_old->sk);
2267
2268                 hlist_del_rcu(&l_old->hash_node);
2269                 smap_list_remove(psock, NULL, l_old);
2270                 smap_release_sock(psock, l_old->sk);
2271                 free_htab_elem(htab, l_old);
2272         }
2273         raw_spin_unlock_bh(&b->lock);
2274         return 0;
2275 bucket_err:
2276         raw_spin_unlock_bh(&b->lock);
2277 err:
2278         kfree(e);
2279         psock = smap_psock_sk(sock);
2280         if (psock)
2281                 smap_release_sock(psock, sock);
2282         return err;
2283 }
2284
2285 static int sock_hash_update_elem(struct bpf_map *map,
2286                                 void *key, void *value, u64 flags)
2287 {
2288         struct bpf_sock_ops_kern skops;
2289         u32 fd = *(u32 *)value;
2290         struct socket *socket;
2291         int err;
2292
2293         socket = sockfd_lookup(fd, &err);
2294         if (!socket)
2295                 return err;
2296
2297         skops.sk = socket->sk;
2298         if (!skops.sk) {
2299                 fput(socket->file);
2300                 return -EINVAL;
2301         }
2302
2303         err = sock_hash_ctx_update_elem(&skops, map, key, flags);
2304         fput(socket->file);
2305         return err;
2306 }
2307
2308 static int sock_hash_delete_elem(struct bpf_map *map, void *key)
2309 {
2310         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
2311         struct hlist_head *head;
2312         struct bucket *b;
2313         struct htab_elem *l;
2314         u32 hash, key_size;
2315         int ret = -ENOENT;
2316
2317         key_size = map->key_size;
2318         hash = htab_map_hash(key, key_size);
2319         b = __select_bucket(htab, hash);
2320         head = &b->head;
2321
2322         raw_spin_lock_bh(&b->lock);
2323         l = lookup_elem_raw(head, hash, key, key_size);
2324         if (l) {
2325                 struct sock *sock = l->sk;
2326                 struct smap_psock *psock;
2327
2328                 hlist_del_rcu(&l->hash_node);
2329                 write_lock_bh(&sock->sk_callback_lock);
2330                 psock = smap_psock_sk(sock);
2331                 /* This check handles a racing sock event that can get the
2332                  * sk_callback_lock before this case but after xchg happens
2333                  * causing the refcnt to hit zero and sock user data (psock)
2334                  * to be null and queued for garbage collection.
2335                  */
2336                 if (likely(psock)) {
2337                         smap_list_remove(psock, NULL, l);
2338                         smap_release_sock(psock, sock);
2339                 }
2340                 write_unlock_bh(&sock->sk_callback_lock);
2341                 free_htab_elem(htab, l);
2342                 ret = 0;
2343         }
2344         raw_spin_unlock_bh(&b->lock);
2345         return ret;
2346 }
2347
2348 struct sock  *__sock_hash_lookup_elem(struct bpf_map *map, void *key)
2349 {
2350         struct bpf_htab *htab = container_of(map, struct bpf_htab, map);
2351         struct hlist_head *head;
2352         struct htab_elem *l;
2353         u32 key_size, hash;
2354         struct bucket *b;
2355         struct sock *sk;
2356
2357         key_size = map->key_size;
2358         hash = htab_map_hash(key, key_size);
2359         b = __select_bucket(htab, hash);
2360         head = &b->head;
2361
2362         raw_spin_lock_bh(&b->lock);
2363         l = lookup_elem_raw(head, hash, key, key_size);
2364         sk = l ? l->sk : NULL;
2365         raw_spin_unlock_bh(&b->lock);
2366         return sk;
2367 }
2368
2369 const struct bpf_map_ops sock_map_ops = {
2370         .map_alloc = sock_map_alloc,
2371         .map_free = sock_map_free,
2372         .map_lookup_elem = sock_map_lookup,
2373         .map_get_next_key = sock_map_get_next_key,
2374         .map_update_elem = sock_map_update_elem,
2375         .map_delete_elem = sock_map_delete_elem,
2376         .map_release_uref = sock_map_release,
2377 };
2378
2379 const struct bpf_map_ops sock_hash_ops = {
2380         .map_alloc = sock_hash_alloc,
2381         .map_free = sock_hash_free,
2382         .map_lookup_elem = sock_map_lookup,
2383         .map_get_next_key = sock_hash_get_next_key,
2384         .map_update_elem = sock_hash_update_elem,
2385         .map_delete_elem = sock_hash_delete_elem,
2386 };
2387
2388 BPF_CALL_4(bpf_sock_map_update, struct bpf_sock_ops_kern *, bpf_sock,
2389            struct bpf_map *, map, void *, key, u64, flags)
2390 {
2391         WARN_ON_ONCE(!rcu_read_lock_held());
2392         return sock_map_ctx_update_elem(bpf_sock, map, key, flags);
2393 }
2394
2395 const struct bpf_func_proto bpf_sock_map_update_proto = {
2396         .func           = bpf_sock_map_update,
2397         .gpl_only       = false,
2398         .pkt_access     = true,
2399         .ret_type       = RET_INTEGER,
2400         .arg1_type      = ARG_PTR_TO_CTX,
2401         .arg2_type      = ARG_CONST_MAP_PTR,
2402         .arg3_type      = ARG_PTR_TO_MAP_KEY,
2403         .arg4_type      = ARG_ANYTHING,
2404 };
2405
2406 BPF_CALL_4(bpf_sock_hash_update, struct bpf_sock_ops_kern *, bpf_sock,
2407            struct bpf_map *, map, void *, key, u64, flags)
2408 {
2409         WARN_ON_ONCE(!rcu_read_lock_held());
2410         return sock_hash_ctx_update_elem(bpf_sock, map, key, flags);
2411 }
2412
2413 const struct bpf_func_proto bpf_sock_hash_update_proto = {
2414         .func           = bpf_sock_hash_update,
2415         .gpl_only       = false,
2416         .pkt_access     = true,
2417         .ret_type       = RET_INTEGER,
2418         .arg1_type      = ARG_PTR_TO_CTX,
2419         .arg2_type      = ARG_CONST_MAP_PTR,
2420         .arg3_type      = ARG_PTR_TO_MAP_KEY,
2421         .arg4_type      = ARG_ANYTHING,
2422 };