// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2024-2025 Meta Platforms, Inc. and affiliates. */
#include <vmlinux.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"

#define EDEADLK 35
#define ETIMEDOUT 110

struct arr_elem {
        struct bpf_res_spin_lock lock;
};

struct {
        __uint(type, BPF_MAP_TYPE_ARRAY);
        __uint(max_entries, 64);
        __type(key, int);
        __type(value, struct arr_elem);
} arrmap SEC(".maps");

struct bpf_res_spin_lock lockA __hidden SEC(".data.A");
struct bpf_res_spin_lock lockB __hidden SEC(".data.B");

SEC("tc")
int res_spin_lock_test(struct __sk_buff *ctx)
{
        struct arr_elem *elem1, *elem2;
        int r;

        elem1 = bpf_map_lookup_elem(&arrmap, &(int){0});
        if (!elem1)
                return -1;
        elem2 = bpf_map_lookup_elem(&arrmap, &(int){0});
        if (!elem2)
                return -1;

        r = bpf_res_spin_lock(&elem1->lock);
        if (r)
                return r;
        r = bpf_res_spin_lock(&elem2->lock);
        if (!r) {
                bpf_res_spin_unlock(&elem2->lock);
                bpf_res_spin_unlock(&elem1->lock);
                return -1;
        }
        bpf_res_spin_unlock(&elem1->lock);
        return r != -EDEADLK;
}

SEC("tc")
int res_spin_lock_test_AB(struct __sk_buff *ctx)
{
        int r;

        r = bpf_res_spin_lock(&lockA);
        if (r)
                return !r;
        /* Only unlock if we took the lock. */
        if (!bpf_res_spin_lock(&lockB))
                bpf_res_spin_unlock(&lockB);
        bpf_res_spin_unlock(&lockA);
        return 0;
}

int err;

SEC("tc")
int res_spin_lock_test_BA(struct __sk_buff *ctx)
{
        int r;

        r = bpf_res_spin_lock(&lockB);
        if (r)
                return !r;
        if (!bpf_res_spin_lock(&lockA))
                bpf_res_spin_unlock(&lockA);
        else
                err = -EDEADLK;
        bpf_res_spin_unlock(&lockB);
        return err ?: 0;
}

SEC("tc")
int res_spin_lock_test_held_lock_max(struct __sk_buff *ctx)
{
        struct bpf_res_spin_lock *locks[48] = {};
        struct arr_elem *e;
        u64 time_beg, time;
        int ret = 0, i;

        _Static_assert(ARRAY_SIZE(((struct rqspinlock_held){}).locks) == 31,
                       "RES_NR_HELD assumed to be 31");

        for (i = 0; i < 34; i++) {
                int key = i;

                /* We cannot pass in i as it will get spilled/filled by the compiler and
                 * loses bounds in verifier state.
                 */
                e = bpf_map_lookup_elem(&arrmap, &key);
                if (!e)
                        return 1;
                locks[i] = &e->lock;
        }

        for (; i < 48; i++) {
                int key = i - 2;

                /* We cannot pass in i as it will get spilled/filled by the compiler and
                 * loses bounds in verifier state.
                 */
                e = bpf_map_lookup_elem(&arrmap, &key);
                if (!e)
                        return 1;
                locks[i] = &e->lock;
        }

        time_beg = bpf_ktime_get_ns();
        for (i = 0; i < 34; i++) {
                if (bpf_res_spin_lock(locks[i]))
                        goto end;
        }

        /* Trigger AA, after exhausting entries in the held lock table. This
         * time, only the timeout can save us, as AA detection won't succeed.
         */
        ret = bpf_res_spin_lock(locks[34]);
        if (!ret) {
                bpf_res_spin_unlock(locks[34]);
                ret = 1;
                goto end;
        }

        ret = ret != -ETIMEDOUT ? 2 : 0;

end:
        for (i = i - 1; i >= 0; i--)
                bpf_res_spin_unlock(locks[i]);
        time = bpf_ktime_get_ns() - time_beg;
        /* Time spent should be easily above our limit (1/4 s), since AA
         * detection won't be expedited due to lack of held lock entry.
         */
        return ret ?: (time > 1000000000 / 4 ? 0 : 1);
}

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