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