// SPDX-License-Identifier: GPL-2.0
/* Converted from tools/testing/selftests/bpf/verifier/jeq_infer_not_null.c */

#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"

struct {
        __uint(type, BPF_MAP_TYPE_XSKMAP);
        __uint(max_entries, 1);
        __type(key, int);
        __type(value, int);
} map_xskmap SEC(".maps");

/* This is equivalent to the following program:
 *
 *   r6 = skb->sk;
 *   r7 = sk_fullsock(r6);
 *   r0 = sk_fullsock(r6);
 *   if (r0 == 0) return 0;    (a)
 *   if (r0 != r7) return 0;   (b)
 *   *r7->type;                (c)
 *   return 0;
 *
 * It is safe to dereference r7 at point (c), because of (a) and (b).
 * The test verifies that relation r0 == r7 is propagated from (b) to (c).
 */
SEC("cgroup/skb")
__description("jne/jeq infer not null, PTR_TO_SOCKET_OR_NULL -> PTR_TO_SOCKET for JNE false branch")
__success __failure_unpriv __msg_unpriv("R7 pointer comparison")
__retval(0)
__naked void socket_for_jne_false_branch(void)
{
        asm volatile ("                                 \
        /* r6 = skb->sk; */                             \
        r6 = *(u64*)(r1 + %[__sk_buff_sk]);             \
        /* if (r6 == 0) return 0; */                    \
        if r6 == 0 goto l0_%=;                          \
        /* r7 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        r7 = r0;                                        \
        /* r0 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        /* if (r0 == null) return 0; */                 \
        if r0 == 0 goto l0_%=;                          \
        /* if (r0 == r7) r0 = *(r7->type); */           \
        if r0 != r7 goto l0_%=;         /* Use ! JNE ! */\
        r0 = *(u32*)(r7 + %[bpf_sock_type]);            \
l0_%=:  /* return 0 */                                  \
        r0 = 0;                                         \
        exit;                                           \
"       :
        : __imm(bpf_sk_fullsock),
          __imm_const(__sk_buff_sk, offsetof(struct __sk_buff, sk)),
          __imm_const(bpf_sock_type, offsetof(struct bpf_sock, type))
        : __clobber_all);
}

/* Same as above, but verify that another branch of JNE still
 * prohibits access to PTR_MAYBE_NULL.
 */
SEC("cgroup/skb")
__description("jne/jeq infer not null, PTR_TO_SOCKET_OR_NULL unchanged for JNE true branch")
__failure __msg("R7 invalid mem access 'sock_or_null'")
__failure_unpriv __msg_unpriv("R7 pointer comparison")
__naked void unchanged_for_jne_true_branch(void)
{
        asm volatile ("                                 \
        /* r6 = skb->sk */                              \
        r6 = *(u64*)(r1 + %[__sk_buff_sk]);             \
        /* if (r6 == 0) return 0; */                    \
        if r6 == 0 goto l0_%=;                          \
        /* r7 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        r7 = r0;                                        \
        /* r0 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        /* if (r0 == null) return 0; */                 \
        if r0 != 0 goto l0_%=;                          \
        /* if (r0 == r7) return 0; */                   \
        if r0 != r7 goto l1_%=;         /* Use ! JNE ! */\
        goto l0_%=;                                     \
l1_%=:  /* r0 = *(r7->type); */                         \
        r0 = *(u32*)(r7 + %[bpf_sock_type]);            \
l0_%=:  /* return 0 */                                  \
        r0 = 0;                                         \
        exit;                                           \
"       :
        : __imm(bpf_sk_fullsock),
          __imm_const(__sk_buff_sk, offsetof(struct __sk_buff, sk)),
          __imm_const(bpf_sock_type, offsetof(struct bpf_sock, type))
        : __clobber_all);
}

/* Same as a first test, but not null should be inferred for JEQ branch */
SEC("cgroup/skb")
__description("jne/jeq infer not null, PTR_TO_SOCKET_OR_NULL -> PTR_TO_SOCKET for JEQ true branch")
__success __failure_unpriv __msg_unpriv("R7 pointer comparison")
__retval(0)
__naked void socket_for_jeq_true_branch(void)
{
        asm volatile ("                                 \
        /* r6 = skb->sk; */                             \
        r6 = *(u64*)(r1 + %[__sk_buff_sk]);             \
        /* if (r6 == null) return 0; */                 \
        if r6 == 0 goto l0_%=;                          \
        /* r7 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        r7 = r0;                                        \
        /* r0 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        /* if (r0 == null) return 0; */                 \
        if r0 == 0 goto l0_%=;                          \
        /* if (r0 != r7) return 0; */                   \
        if r0 == r7 goto l1_%=;         /* Use ! JEQ ! */\
        goto l0_%=;                                     \
l1_%=:  /* r0 = *(r7->type); */                         \
        r0 = *(u32*)(r7 + %[bpf_sock_type]);            \
l0_%=:  /* return 0; */                                 \
        r0 = 0;                                         \
        exit;                                           \
"       :
        : __imm(bpf_sk_fullsock),
          __imm_const(__sk_buff_sk, offsetof(struct __sk_buff, sk)),
          __imm_const(bpf_sock_type, offsetof(struct bpf_sock, type))
        : __clobber_all);
}

/* Same as above, but verify that another branch of JNE still
 * prohibits access to PTR_MAYBE_NULL.
 */
SEC("cgroup/skb")
__description("jne/jeq infer not null, PTR_TO_SOCKET_OR_NULL unchanged for JEQ false branch")
__failure __msg("R7 invalid mem access 'sock_or_null'")
__failure_unpriv __msg_unpriv("R7 pointer comparison")
__naked void unchanged_for_jeq_false_branch(void)
{
        asm volatile ("                                 \
        /* r6 = skb->sk; */                             \
        r6 = *(u64*)(r1 + %[__sk_buff_sk]);             \
        /* if (r6 == null) return 0; */                 \
        if r6 == 0 goto l0_%=;                          \
        /* r7 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        r7 = r0;                                        \
        /* r0 = sk_fullsock(skb); */                    \
        r1 = r6;                                        \
        call %[bpf_sk_fullsock];                        \
        /* if (r0 == null) return 0; */                 \
        if r0 == 0 goto l0_%=;                          \
        /* if (r0 != r7) r0 = *(r7->type); */           \
        if r0 == r7 goto l0_%=;         /* Use ! JEQ ! */\
        r0 = *(u32*)(r7 + %[bpf_sock_type]);            \
l0_%=:  /* return 0; */                                 \
        r0 = 0;                                         \
        exit;                                           \
"       :
        : __imm(bpf_sk_fullsock),
          __imm_const(__sk_buff_sk, offsetof(struct __sk_buff, sk)),
          __imm_const(bpf_sock_type, offsetof(struct bpf_sock, type))
        : __clobber_all);
}

/* Maps are treated in a different branch of `mark_ptr_not_null_reg`,
 * so separate test for maps case.
 */
SEC("xdp")
__description("jne/jeq infer not null, PTR_TO_MAP_VALUE_OR_NULL -> PTR_TO_MAP_VALUE")
__success __retval(0)
__naked void null_ptr_to_map_value(void)
{
        asm volatile ("                                 \
        /* r9 = &some stack to use as key */            \
        r1 = 0;                                         \
        *(u32*)(r10 - 8) = r1;                          \
        r9 = r10;                                       \
        r9 += -8;                                       \
        /* r8 = process local map */                    \
        r8 = %[map_xskmap] ll;                          \
        /* r6 = map_lookup_elem(r8, r9); */             \
        r1 = r8;                                        \
        r2 = r9;                                        \
        call %[bpf_map_lookup_elem];                    \
        r6 = r0;                                        \
        /* r7 = map_lookup_elem(r8, r9); */             \
        r1 = r8;                                        \
        r2 = r9;                                        \
        call %[bpf_map_lookup_elem];                    \
        r7 = r0;                                        \
        /* if (r6 == 0) return 0; */                    \
        if r6 == 0 goto l0_%=;                          \
        /* if (r6 != r7) return 0; */                   \
        if r6 != r7 goto l0_%=;                         \
        /* read *r7; */                                 \
        r0 = *(u32*)(r7 + %[bpf_xdp_sock_queue_id]);    \
l0_%=:  /* return 0; */                                 \
        r0 = 0;                                         \
        exit;                                           \
"       :
        : __imm(bpf_map_lookup_elem),
          __imm_addr(map_xskmap),
          __imm_const(bpf_xdp_sock_queue_id, offsetof(struct bpf_xdp_sock, queue_id))
        : __clobber_all);
}

char _license[] SEC("license") = "GPL";