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 <net/strparser.h>
  42 #include <net/tcp.h>
  43
  44 struct bpf_stab {
  45         struct bpf_map map;
  46         struct sock **sock_map;
  47         struct bpf_prog *bpf_parse;
  48         struct bpf_prog *bpf_verdict;
  49 };
  50
  51 enum smap_psock_state {
  52         SMAP_TX_RUNNING,
  53 };
  54
  55 struct smap_psock_map_entry {
  56         struct list_head list;
  57         struct sock **entry;
  58 };
  59
  60 struct smap_psock {
  61         struct rcu_head rcu;
  62         /* refcnt is used inside sk_callback_lock */
  63         u32 refcnt;
  64
  65         /* datapath variables */
  66         struct sk_buff_head rxqueue;
  67         bool strp_enabled;
  68
  69         /* datapath error path cache across tx work invocations */
  70         int save_rem;
  71         int save_off;
  72         struct sk_buff *save_skb;
  73
  74         struct strparser strp;
  75         struct bpf_prog *bpf_parse;
  76         struct bpf_prog *bpf_verdict;
  77         struct list_head maps;
  78
  79         /* Back reference used when sock callback trigger sockmap operations */
  80         struct sock *sock;
  81         unsigned long state;
  82
  83         struct work_struct tx_work;
  84         struct work_struct gc_work;
  85
  86         void (*save_data_ready)(struct sock *sk);
  87         void (*save_write_space)(struct sock *sk);
  88         void (*save_state_change)(struct sock *sk);
  89 };
  90
  91 static inline struct smap_psock *smap_psock_sk(const struct sock *sk)
  92 {
  93         return rcu_dereference_sk_user_data(sk);
  94 }
  95
  96 /* compute the linear packet data range [data, data_end) for skb when
  97  * sk_skb type programs are in use.
  98  */
  99 static inline void bpf_compute_data_end_sk_skb(struct sk_buff *skb)
 100 {
 101         TCP_SKB_CB(skb)->bpf.data_end = skb->data + skb_headlen(skb);
 102 }
 103
 104 enum __sk_action {
 105         __SK_DROP = 0,
 106         __SK_PASS,
 107         __SK_REDIRECT,
 108 };
 109
 110 static int smap_verdict_func(struct smap_psock *psock, struct sk_buff *skb)
 111 {
 112         struct bpf_prog *prog = READ_ONCE(psock->bpf_verdict);
 113         int rc;
 114
 115         if (unlikely(!prog))
 116                 return __SK_DROP;
 117
 118         skb_orphan(skb);
 119         /* We need to ensure that BPF metadata for maps is also cleared
 120          * when we orphan the skb so that we don't have the possibility
 121          * to reference a stale map.
 122          */
 123         TCP_SKB_CB(skb)->bpf.map = NULL;
 124         skb->sk = psock->sock;
 125         bpf_compute_data_end_sk_skb(skb);
 126         preempt_disable();
 127         rc = (*prog->bpf_func)(skb, prog->insnsi);
 128         preempt_enable();
 129         skb->sk = NULL;
 130
 131         /* Moving return codes from UAPI namespace into internal namespace */
 132         return rc == SK_PASS ?
 133                 (TCP_SKB_CB(skb)->bpf.map ? __SK_REDIRECT : __SK_PASS) :
 134                 __SK_DROP;
 135 }
 136
 137 static void smap_do_verdict(struct smap_psock *psock, struct sk_buff *skb)
 138 {
 139         struct sock *sk;
 140         int rc;
 141
 142         rc = smap_verdict_func(psock, skb);
 143         switch (rc) {
 144         case __SK_REDIRECT:
 145                 sk = do_sk_redirect_map(skb);
 146                 if (likely(sk)) {
 147                         struct smap_psock *peer = smap_psock_sk(sk);
 148
 149                         if (likely(peer &&
 150                                    test_bit(SMAP_TX_RUNNING, &peer->state) &&
 151                                    !sock_flag(sk, SOCK_DEAD) &&
 152                                    sock_writeable(sk))) {
 153                                 skb_set_owner_w(skb, sk);
 154                                 skb_queue_tail(&peer->rxqueue, skb);
 155                                 schedule_work(&peer->tx_work);
 156                                 break;
 157                         }
 158                 }
 159         /* Fall through and free skb otherwise */
 160         case __SK_DROP:
 161         default:
 162                 kfree_skb(skb);
 163         }
 164 }
 165
 166 static void smap_report_sk_error(struct smap_psock *psock, int err)
 167 {
 168         struct sock *sk = psock->sock;
 169
 170         sk->sk_err = err;
 171         sk->sk_error_report(sk);
 172 }
 173
 174 static void smap_release_sock(struct smap_psock *psock, struct sock *sock);
 175
 176 /* Called with lock_sock(sk) held */
 177 static void smap_state_change(struct sock *sk)
 178 {
 179         struct smap_psock_map_entry *e, *tmp;
 180         struct smap_psock *psock;
 181         struct socket_wq *wq;
 182         struct sock *osk;
 183
 184         rcu_read_lock();
 185
 186         /* Allowing transitions into an established syn_recv states allows
 187          * for early binding sockets to a smap object before the connection
 188          * is established.
 189          */
 190         switch (sk->sk_state) {
 191         case TCP_SYN_SENT:
 192         case TCP_SYN_RECV:
 193         case TCP_ESTABLISHED:
 194                 break;
 195         case TCP_CLOSE_WAIT:
 196         case TCP_CLOSING:
 197         case TCP_LAST_ACK:
 198         case TCP_FIN_WAIT1:
 199         case TCP_FIN_WAIT2:
 200         case TCP_LISTEN:
 201                 break;
 202         case TCP_CLOSE:
 203                 /* Only release if the map entry is in fact the sock in
 204                  * question. There is a case where the operator deletes
 205                  * the sock from the map, but the TCP sock is closed before
 206                  * the psock is detached. Use cmpxchg to verify correct
 207                  * sock is removed.
 208                  */
 209                 psock = smap_psock_sk(sk);
 210                 if (unlikely(!psock))
 211                         break;
 212                 write_lock_bh(&sk->sk_callback_lock);
 213                 list_for_each_entry_safe(e, tmp, &psock->maps, list) {
 214                         osk = cmpxchg(e->entry, sk, NULL);
 215                         if (osk == sk) {
 216                                 list_del(&e->list);
 217                                 smap_release_sock(psock, sk);
 218                         }
 219                 }
 220                 write_unlock_bh(&sk->sk_callback_lock);
 221                 break;
 222         default:
 223                 psock = smap_psock_sk(sk);
 224                 if (unlikely(!psock))
 225                         break;
 226                 smap_report_sk_error(psock, EPIPE);
 227                 break;
 228         }
 229
 230         wq = rcu_dereference(sk->sk_wq);
 231         if (skwq_has_sleeper(wq))
 232                 wake_up_interruptible_all(&wq->wait);
 233         rcu_read_unlock();
 234 }
 235
 236 static void smap_read_sock_strparser(struct strparser *strp,
 237                                      struct sk_buff *skb)
 238 {
 239         struct smap_psock *psock;
 240
 241         rcu_read_lock();
 242         psock = container_of(strp, struct smap_psock, strp);
 243         smap_do_verdict(psock, skb);
 244         rcu_read_unlock();
 245 }
 246
 247 /* Called with lock held on socket */
 248 static void smap_data_ready(struct sock *sk)
 249 {
 250         struct smap_psock *psock;
 251
 252         rcu_read_lock();
 253         psock = smap_psock_sk(sk);
 254         if (likely(psock)) {
 255                 write_lock_bh(&sk->sk_callback_lock);
 256                 strp_data_ready(&psock->strp);
 257                 write_unlock_bh(&sk->sk_callback_lock);
 258         }
 259         rcu_read_unlock();
 260 }
 261
 262 static void smap_tx_work(struct work_struct *w)
 263 {
 264         struct smap_psock *psock;
 265         struct sk_buff *skb;
 266         int rem, off, n;
 267
 268         psock = container_of(w, struct smap_psock, tx_work);
 269
 270         /* lock sock to avoid losing sk_socket at some point during loop */
 271         lock_sock(psock->sock);
 272         if (psock->save_skb) {
 273                 skb = psock->save_skb;
 274                 rem = psock->save_rem;
 275                 off = psock->save_off;
 276                 psock->save_skb = NULL;
 277                 goto start;
 278         }
 279
 280         while ((skb = skb_dequeue(&psock->rxqueue))) {
 281                 rem = skb->len;
 282                 off = 0;
 283 start:
 284                 do {
 285                         if (likely(psock->sock->sk_socket))
 286                                 n = skb_send_sock_locked(psock->sock,
 287                                                          skb, off, rem);
 288                         else
 289                                 n = -EINVAL;
 290                         if (n <= 0) {
 291                                 if (n == -EAGAIN) {
 292                                         /* Retry when space is available */
 293                                         psock->save_skb = skb;
 294                                         psock->save_rem = rem;
 295                                         psock->save_off = off;
 296                                         goto out;
 297                                 }
 298                                 /* Hard errors break pipe and stop xmit */
 299                                 smap_report_sk_error(psock, n ? -n : EPIPE);
 300                                 clear_bit(SMAP_TX_RUNNING, &psock->state);
 301                                 kfree_skb(skb);
 302                                 goto out;
 303                         }
 304                         rem -= n;
 305                         off += n;
 306                 } while (rem);
 307                 kfree_skb(skb);
 308         }
 309 out:
 310         release_sock(psock->sock);
 311 }
 312
 313 static void smap_write_space(struct sock *sk)
 314 {
 315         struct smap_psock *psock;
 316         void (*write_space)(struct sock *sk);
 317
 318         rcu_read_lock();
 319         psock = smap_psock_sk(sk);
 320         if (likely(psock && test_bit(SMAP_TX_RUNNING, &psock->state)))
 321                 schedule_work(&psock->tx_work);
 322         write_space = psock->save_write_space;
 323         rcu_read_unlock();
 324         write_space(sk);
 325 }
 326
 327 static void smap_stop_sock(struct smap_psock *psock, struct sock *sk)
 328 {
 329         if (!psock->strp_enabled)
 330                 return;
 331         sk->sk_data_ready = psock->save_data_ready;
 332         sk->sk_write_space = psock->save_write_space;
 333         sk->sk_state_change = psock->save_state_change;
 334         psock->save_data_ready = NULL;
 335         psock->save_write_space = NULL;
 336         psock->save_state_change = NULL;
 337         strp_stop(&psock->strp);
 338         psock->strp_enabled = false;
 339 }
 340
 341 static void smap_destroy_psock(struct rcu_head *rcu)
 342 {
 343         struct smap_psock *psock = container_of(rcu,
 344                                                   struct smap_psock, rcu);
 345
 346         /* Now that a grace period has passed there is no longer
 347          * any reference to this sock in the sockmap so we can
 348          * destroy the psock, strparser, and bpf programs. But,
 349          * because we use workqueue sync operations we can not
 350          * do it in rcu context
 351          */
 352         schedule_work(&psock->gc_work);
 353 }
 354
 355 static void smap_release_sock(struct smap_psock *psock, struct sock *sock)
 356 {
 357         psock->refcnt--;
 358         if (psock->refcnt)
 359                 return;
 360
 361         smap_stop_sock(psock, sock);
 362         clear_bit(SMAP_TX_RUNNING, &psock->state);
 363         rcu_assign_sk_user_data(sock, NULL);
 364         call_rcu_sched(&psock->rcu, smap_destroy_psock);
 365 }
 366
 367 static int smap_parse_func_strparser(struct strparser *strp,
 368                                        struct sk_buff *skb)
 369 {
 370         struct smap_psock *psock;
 371         struct bpf_prog *prog;
 372         int rc;
 373
 374         rcu_read_lock();
 375         psock = container_of(strp, struct smap_psock, strp);
 376         prog = READ_ONCE(psock->bpf_parse);
 377
 378         if (unlikely(!prog)) {
 379                 rcu_read_unlock();
 380                 return skb->len;
 381         }
 382
 383         /* Attach socket for bpf program to use if needed we can do this
 384          * because strparser clones the skb before handing it to a upper
 385          * layer, meaning skb_orphan has been called. We NULL sk on the
 386          * way out to ensure we don't trigger a BUG_ON in skb/sk operations
 387          * later and because we are not charging the memory of this skb to
 388          * any socket yet.
 389          */
 390         skb->sk = psock->sock;
 391         bpf_compute_data_end_sk_skb(skb);
 392         rc = (*prog->bpf_func)(skb, prog->insnsi);
 393         skb->sk = NULL;
 394         rcu_read_unlock();
 395         return rc;
 396 }
 397
 398
 399 static int smap_read_sock_done(struct strparser *strp, int err)
 400 {
 401         return err;
 402 }
 403
 404 static int smap_init_sock(struct smap_psock *psock,
 405                           struct sock *sk)
 406 {
 407         static const struct strp_callbacks cb = {
 408                 .rcv_msg = smap_read_sock_strparser,
 409                 .parse_msg = smap_parse_func_strparser,
 410                 .read_sock_done = smap_read_sock_done,
 411         };
 412
 413         return strp_init(&psock->strp, sk, &cb);
 414 }
 415
 416 static void smap_init_progs(struct smap_psock *psock,
 417                             struct bpf_stab *stab,
 418                             struct bpf_prog *verdict,
 419                             struct bpf_prog *parse)
 420 {
 421         struct bpf_prog *orig_parse, *orig_verdict;
 422
 423         orig_parse = xchg(&psock->bpf_parse, parse);
 424         orig_verdict = xchg(&psock->bpf_verdict, verdict);
 425
 426         if (orig_verdict)
 427                 bpf_prog_put(orig_verdict);
 428         if (orig_parse)
 429                 bpf_prog_put(orig_parse);
 430 }
 431
 432 static void smap_start_sock(struct smap_psock *psock, struct sock *sk)
 433 {
 434         if (sk->sk_data_ready == smap_data_ready)
 435                 return;
 436         psock->save_data_ready = sk->sk_data_ready;
 437         psock->save_write_space = sk->sk_write_space;
 438         psock->save_state_change = sk->sk_state_change;
 439         sk->sk_data_ready = smap_data_ready;
 440         sk->sk_write_space = smap_write_space;
 441         sk->sk_state_change = smap_state_change;
 442         psock->strp_enabled = true;
 443 }
 444
 445 static void sock_map_remove_complete(struct bpf_stab *stab)
 446 {
 447         bpf_map_area_free(stab->sock_map);
 448         kfree(stab);
 449 }
 450
 451 static void smap_gc_work(struct work_struct *w)
 452 {
 453         struct smap_psock_map_entry *e, *tmp;
 454         struct smap_psock *psock;
 455
 456         psock = container_of(w, struct smap_psock, gc_work);
 457
 458         /* no callback lock needed because we already detached sockmap ops */
 459         if (psock->strp_enabled)
 460                 strp_done(&psock->strp);
 461
 462         cancel_work_sync(&psock->tx_work);
 463         __skb_queue_purge(&psock->rxqueue);
 464
 465         /* At this point all strparser and xmit work must be complete */
 466         if (psock->bpf_parse)
 467                 bpf_prog_put(psock->bpf_parse);
 468         if (psock->bpf_verdict)
 469                 bpf_prog_put(psock->bpf_verdict);
 470
 471         list_for_each_entry_safe(e, tmp, &psock->maps, list) {
 472                 list_del(&e->list);
 473                 kfree(e);
 474         }
 475
 476         sock_put(psock->sock);
 477         kfree(psock);
 478 }
 479
 480 static struct smap_psock *smap_init_psock(struct sock *sock,
 481                                           struct bpf_stab *stab)
 482 {
 483         struct smap_psock *psock;
 484
 485         psock = kzalloc_node(sizeof(struct smap_psock),
 486                              GFP_ATOMIC | __GFP_NOWARN,
 487                              stab->map.numa_node);
 488         if (!psock)
 489                 return ERR_PTR(-ENOMEM);
 490
 491         psock->sock = sock;
 492         skb_queue_head_init(&psock->rxqueue);
 493         INIT_WORK(&psock->tx_work, smap_tx_work);
 494         INIT_WORK(&psock->gc_work, smap_gc_work);
 495         INIT_LIST_HEAD(&psock->maps);
 496         psock->refcnt = 1;
 497
 498         rcu_assign_sk_user_data(sock, psock);
 499         sock_hold(sock);
 500         return psock;
 501 }
 502
 503 static struct bpf_map *sock_map_alloc(union bpf_attr *attr)
 504 {
 505         struct bpf_stab *stab;
 506         int err = -EINVAL;
 507         u64 cost;
 508
 509         if (!capable(CAP_NET_ADMIN))
 510                 return ERR_PTR(-EPERM);
 511
 512         /* check sanity of attributes */
 513         if (attr->max_entries == 0 || attr->key_size != 4 ||
 514             attr->value_size != 4 || attr->map_flags & ~BPF_F_NUMA_NODE)
 515                 return ERR_PTR(-EINVAL);
 516
 517         if (attr->value_size > KMALLOC_MAX_SIZE)
 518                 return ERR_PTR(-E2BIG);
 519
 520         stab = kzalloc(sizeof(*stab), GFP_USER);
 521         if (!stab)
 522                 return ERR_PTR(-ENOMEM);
 523
 524         /* mandatory map attributes */
 525         stab->map.map_type = attr->map_type;
 526         stab->map.key_size = attr->key_size;
 527         stab->map.value_size = attr->value_size;
 528         stab->map.max_entries = attr->max_entries;
 529         stab->map.map_flags = attr->map_flags;
 530         stab->map.numa_node = bpf_map_attr_numa_node(attr);
 531
 532         /* make sure page count doesn't overflow */
 533         cost = (u64) stab->map.max_entries * sizeof(struct sock *);
 534         if (cost >= U32_MAX - PAGE_SIZE)
 535                 goto free_stab;
 536
 537         stab->map.pages = round_up(cost, PAGE_SIZE) >> PAGE_SHIFT;
 538
 539         /* if map size is larger than memlock limit, reject it early */
 540         err = bpf_map_precharge_memlock(stab->map.pages);
 541         if (err)
 542                 goto free_stab;
 543
 544         err = -ENOMEM;
 545         stab->sock_map = bpf_map_area_alloc(stab->map.max_entries *
 546                                             sizeof(struct sock *),
 547                                             stab->map.numa_node);
 548         if (!stab->sock_map)
 549                 goto free_stab;
 550
 551         return &stab->map;
 552 free_stab:
 553         kfree(stab);
 554         return ERR_PTR(err);
 555 }
 556
 557 static void smap_list_remove(struct smap_psock *psock, struct sock **entry)
 558 {
 559         struct smap_psock_map_entry *e, *tmp;
 560
 561         list_for_each_entry_safe(e, tmp, &psock->maps, list) {
 562                 if (e->entry == entry) {
 563                         list_del(&e->list);
 564                         break;
 565                 }
 566         }
 567 }
 568
 569 static void sock_map_free(struct bpf_map *map)
 570 {
 571         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 572         int i;
 573
 574         synchronize_rcu();
 575
 576         /* At this point no update, lookup or delete operations can happen.
 577          * However, be aware we can still get a socket state event updates,
 578          * and data ready callabacks that reference the psock from sk_user_data
 579          * Also psock worker threads are still in-flight. So smap_release_sock
 580          * will only free the psock after cancel_sync on the worker threads
 581          * and a grace period expire to ensure psock is really safe to remove.
 582          */
 583         rcu_read_lock();
 584         for (i = 0; i < stab->map.max_entries; i++) {
 585                 struct smap_psock *psock;
 586                 struct sock *sock;
 587
 588                 sock = xchg(&stab->sock_map[i], NULL);
 589                 if (!sock)
 590                         continue;
 591
 592                 write_lock_bh(&sock->sk_callback_lock);
 593                 psock = smap_psock_sk(sock);
 594                 /* This check handles a racing sock event that can get the
 595                  * sk_callback_lock before this case but after xchg happens
 596                  * causing the refcnt to hit zero and sock user data (psock)
 597                  * to be null and queued for garbage collection.
 598                  */
 599                 if (likely(psock)) {
 600                         smap_list_remove(psock, &stab->sock_map[i]);
 601                         smap_release_sock(psock, sock);
 602                 }
 603                 write_unlock_bh(&sock->sk_callback_lock);
 604         }
 605         rcu_read_unlock();
 606
 607         sock_map_remove_complete(stab);
 608 }
 609
 610 static int sock_map_get_next_key(struct bpf_map *map, void *key, void *next_key)
 611 {
 612         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 613         u32 i = key ? *(u32 *)key : U32_MAX;
 614         u32 *next = (u32 *)next_key;
 615
 616         if (i >= stab->map.max_entries) {
 617                 *next = 0;
 618                 return 0;
 619         }
 620
 621         if (i == stab->map.max_entries - 1)
 622                 return -ENOENT;
 623
 624         *next = i + 1;
 625         return 0;
 626 }
 627
 628 struct sock  *__sock_map_lookup_elem(struct bpf_map *map, u32 key)
 629 {
 630         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 631
 632         if (key >= map->max_entries)
 633                 return NULL;
 634
 635         return READ_ONCE(stab->sock_map[key]);
 636 }
 637
 638 static int sock_map_delete_elem(struct bpf_map *map, void *key)
 639 {
 640         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 641         struct smap_psock *psock;
 642         int k = *(u32 *)key;
 643         struct sock *sock;
 644
 645         if (k >= map->max_entries)
 646                 return -EINVAL;
 647
 648         sock = xchg(&stab->sock_map[k], NULL);
 649         if (!sock)
 650                 return -EINVAL;
 651
 652         write_lock_bh(&sock->sk_callback_lock);
 653         psock = smap_psock_sk(sock);
 654         if (!psock)
 655                 goto out;
 656
 657         if (psock->bpf_parse)
 658                 smap_stop_sock(psock, sock);
 659         smap_list_remove(psock, &stab->sock_map[k]);
 660         smap_release_sock(psock, sock);
 661 out:
 662         write_unlock_bh(&sock->sk_callback_lock);
 663         return 0;
 664 }
 665
 666 /* Locking notes: Concurrent updates, deletes, and lookups are allowed and are
 667  * done inside rcu critical sections. This ensures on updates that the psock
 668  * will not be released via smap_release_sock() until concurrent updates/deletes
 669  * complete. All operations operate on sock_map using cmpxchg and xchg
 670  * operations to ensure we do not get stale references. Any reads into the
 671  * map must be done with READ_ONCE() because of this.
 672  *
 673  * A psock is destroyed via call_rcu and after any worker threads are cancelled
 674  * and syncd so we are certain all references from the update/lookup/delete
 675  * operations as well as references in the data path are no longer in use.
 676  *
 677  * Psocks may exist in multiple maps, but only a single set of parse/verdict
 678  * programs may be inherited from the maps it belongs to. A reference count
 679  * is kept with the total number of references to the psock from all maps. The
 680  * psock will not be released until this reaches zero. The psock and sock
 681  * user data data use the sk_callback_lock to protect critical data structures
 682  * from concurrent access. This allows us to avoid two updates from modifying
 683  * the user data in sock and the lock is required anyways for modifying
 684  * callbacks, we simply increase its scope slightly.
 685  *
 686  * Rules to follow,
 687  *  - psock must always be read inside RCU critical section
 688  *  - sk_user_data must only be modified inside sk_callback_lock and read
 689  *    inside RCU critical section.
 690  *  - psock->maps list must only be read & modified inside sk_callback_lock
 691  *  - sock_map must use READ_ONCE and (cmp)xchg operations
 692  *  - BPF verdict/parse programs must use READ_ONCE and xchg operations
 693  */
 694 static int sock_map_ctx_update_elem(struct bpf_sock_ops_kern *skops,
 695                                     struct bpf_map *map,
 696                                     void *key, u64 flags)
 697 {
 698         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 699         struct smap_psock_map_entry *e = NULL;
 700         struct bpf_prog *verdict, *parse;
 701         struct sock *osock, *sock;
 702         struct smap_psock *psock;
 703         u32 i = *(u32 *)key;
 704         int err;
 705
 706         if (unlikely(flags > BPF_EXIST))
 707                 return -EINVAL;
 708
 709         if (unlikely(i >= stab->map.max_entries))
 710                 return -E2BIG;
 711
 712         sock = READ_ONCE(stab->sock_map[i]);
 713         if (flags == BPF_EXIST && !sock)
 714                 return -ENOENT;
 715         else if (flags == BPF_NOEXIST && sock)
 716                 return -EEXIST;
 717
 718         sock = skops->sk;
 719
 720         /* 1. If sock map has BPF programs those will be inherited by the
 721          * sock being added. If the sock is already attached to BPF programs
 722          * this results in an error.
 723          */
 724         verdict = READ_ONCE(stab->bpf_verdict);
 725         parse = READ_ONCE(stab->bpf_parse);
 726
 727         if (parse && verdict) {
 728                 /* bpf prog refcnt may be zero if a concurrent attach operation
 729                  * removes the program after the above READ_ONCE() but before
 730                  * we increment the refcnt. If this is the case abort with an
 731                  * error.
 732                  */
 733                 verdict = bpf_prog_inc_not_zero(stab->bpf_verdict);
 734                 if (IS_ERR(verdict))
 735                         return PTR_ERR(verdict);
 736
 737                 parse = bpf_prog_inc_not_zero(stab->bpf_parse);
 738                 if (IS_ERR(parse)) {
 739                         bpf_prog_put(verdict);
 740                         return PTR_ERR(parse);
 741                 }
 742         }
 743
 744         write_lock_bh(&sock->sk_callback_lock);
 745         psock = smap_psock_sk(sock);
 746
 747         /* 2. Do not allow inheriting programs if psock exists and has
 748          * already inherited programs. This would create confusion on
 749          * which parser/verdict program is running. If no psock exists
 750          * create one. Inside sk_callback_lock to ensure concurrent create
 751          * doesn't update user data.
 752          */
 753         if (psock) {
 754                 if (READ_ONCE(psock->bpf_parse) && parse) {
 755                         err = -EBUSY;
 756                         goto out_progs;
 757                 }
 758                 psock->refcnt++;
 759         } else {
 760                 psock = smap_init_psock(sock, stab);
 761                 if (IS_ERR(psock)) {
 762                         err = PTR_ERR(psock);
 763                         goto out_progs;
 764                 }
 765
 766                 set_bit(SMAP_TX_RUNNING, &psock->state);
 767         }
 768
 769         e = kzalloc(sizeof(*e), GFP_ATOMIC | __GFP_NOWARN);
 770         if (!e) {
 771                 err = -ENOMEM;
 772                 goto out_progs;
 773         }
 774         e->entry = &stab->sock_map[i];
 775
 776         /* 3. At this point we have a reference to a valid psock that is
 777          * running. Attach any BPF programs needed.
 778          */
 779         if (parse && verdict && !psock->strp_enabled) {
 780                 err = smap_init_sock(psock, sock);
 781                 if (err)
 782                         goto out_free;
 783                 smap_init_progs(psock, stab, verdict, parse);
 784                 smap_start_sock(psock, sock);
 785         }
 786
 787         /* 4. Place psock in sockmap for use and stop any programs on
 788          * the old sock assuming its not the same sock we are replacing
 789          * it with. Because we can only have a single set of programs if
 790          * old_sock has a strp we can stop it.
 791          */
 792         list_add_tail(&e->list, &psock->maps);
 793         write_unlock_bh(&sock->sk_callback_lock);
 794
 795         osock = xchg(&stab->sock_map[i], sock);
 796         if (osock) {
 797                 struct smap_psock *opsock = smap_psock_sk(osock);
 798
 799                 write_lock_bh(&osock->sk_callback_lock);
 800                 if (osock != sock && parse)
 801                         smap_stop_sock(opsock, osock);
 802                 smap_list_remove(opsock, &stab->sock_map[i]);
 803                 smap_release_sock(opsock, osock);
 804                 write_unlock_bh(&osock->sk_callback_lock);
 805         }
 806         return 0;
 807 out_free:
 808         smap_release_sock(psock, sock);
 809 out_progs:
 810         if (verdict)
 811                 bpf_prog_put(verdict);
 812         if (parse)
 813                 bpf_prog_put(parse);
 814         write_unlock_bh(&sock->sk_callback_lock);
 815         kfree(e);
 816         return err;
 817 }
 818
 819 int sock_map_prog(struct bpf_map *map, struct bpf_prog *prog, u32 type)
 820 {
 821         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 822         struct bpf_prog *orig;
 823
 824         if (unlikely(map->map_type != BPF_MAP_TYPE_SOCKMAP))
 825                 return -EINVAL;
 826
 827         switch (type) {
 828         case BPF_SK_SKB_STREAM_PARSER:
 829                 orig = xchg(&stab->bpf_parse, prog);
 830                 break;
 831         case BPF_SK_SKB_STREAM_VERDICT:
 832                 orig = xchg(&stab->bpf_verdict, prog);
 833                 break;
 834         default:
 835                 return -EOPNOTSUPP;
 836         }
 837
 838         if (orig)
 839                 bpf_prog_put(orig);
 840
 841         return 0;
 842 }
 843
 844 static void *sock_map_lookup(struct bpf_map *map, void *key)
 845 {
 846         return NULL;
 847 }
 848
 849 static int sock_map_update_elem(struct bpf_map *map,
 850                                 void *key, void *value, u64 flags)
 851 {
 852         struct bpf_sock_ops_kern skops;
 853         u32 fd = *(u32 *)value;
 854         struct socket *socket;
 855         int err;
 856
 857         socket = sockfd_lookup(fd, &err);
 858         if (!socket)
 859                 return err;
 860
 861         skops.sk = socket->sk;
 862         if (!skops.sk) {
 863                 fput(socket->file);
 864                 return -EINVAL;
 865         }
 866
 867         if (skops.sk->sk_type != SOCK_STREAM ||
 868             skops.sk->sk_protocol != IPPROTO_TCP) {
 869                 fput(socket->file);
 870                 return -EOPNOTSUPP;
 871         }
 872
 873         err = sock_map_ctx_update_elem(&skops, map, key, flags);
 874         fput(socket->file);
 875         return err;
 876 }
 877
 878 static void sock_map_release(struct bpf_map *map)
 879 {
 880         struct bpf_stab *stab = container_of(map, struct bpf_stab, map);
 881         struct bpf_prog *orig;
 882
 883         orig = xchg(&stab->bpf_parse, NULL);
 884         if (orig)
 885                 bpf_prog_put(orig);
 886         orig = xchg(&stab->bpf_verdict, NULL);
 887         if (orig)
 888                 bpf_prog_put(orig);
 889 }
 890
 891 const struct bpf_map_ops sock_map_ops = {
 892         .map_alloc = sock_map_alloc,
 893         .map_free = sock_map_free,
 894         .map_lookup_elem = sock_map_lookup,
 895         .map_get_next_key = sock_map_get_next_key,
 896         .map_update_elem = sock_map_update_elem,
 897         .map_delete_elem = sock_map_delete_elem,
 898         .map_release_uref = sock_map_release,
 899 };
 900
 901 BPF_CALL_4(bpf_sock_map_update, struct bpf_sock_ops_kern *, bpf_sock,
 902            struct bpf_map *, map, void *, key, u64, flags)
 903 {
 904         WARN_ON_ONCE(!rcu_read_lock_held());
 905         return sock_map_ctx_update_elem(bpf_sock, map, key, flags);
 906 }
 907
 908 const struct bpf_func_proto bpf_sock_map_update_proto = {
 909         .func           = bpf_sock_map_update,
 910         .gpl_only       = false,
 911         .pkt_access     = true,
 912         .ret_type       = RET_INTEGER,
 913         .arg1_type      = ARG_PTR_TO_CTX,
 914         .arg2_type      = ARG_CONST_MAP_PTR,
 915         .arg3_type      = ARG_PTR_TO_MAP_KEY,
 916         .arg4_type      = ARG_ANYTHING,
 917 };