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