// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2024 Benjamin Tissoires
 */

#include "bpf_experimental.h"
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"
#include "../test_kmods/bpf_testmod_kfunc.h"

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

struct hmap_elem {
 int counter;
 struct bpf_timer timer; /* unused */
 struct bpf_spin_lock lock; /* unused */
 struct bpf_wq work;
};

struct {
 __uint(type, BPF_MAP_TYPE_HASH);
 __uint(max_entries, 1000);
 __type(key, int);
 __type(value, struct hmap_elem);
} hmap SEC(".maps");

struct {
 __uint(type, BPF_MAP_TYPE_HASH);
 __uint(map_flags, BPF_F_NO_PREALLOC);
 __uint(max_entries, 1000);
 __type(key, int);
 __type(value, struct hmap_elem);
} hmap_malloc SEC(".maps");

struct elem {
 int ok_offset;
 struct bpf_wq w;
};

struct {
 __uint(type, BPF_MAP_TYPE_ARRAY);
 __uint(max_entries, 2);
 __type(key, int);
 __type(value, struct elem);
} array SEC(".maps");

struct {
 __uint(type, BPF_MAP_TYPE_LRU_HASH);
 __uint(max_entries, 4);
 __type(key, int);
 __type(value, struct elem);
} lru SEC(".maps");

__u32 ok;
__u32 ok_sleepable;

static int test_elem_callback(void *map, int *key,
  int (callback_fn)(void *map, int *key, void *value))
{
 struct elem init = {}, *val;
 struct bpf_wq *wq;

 if ((ok & (1 << *key) ||
     (ok_sleepable & (1 << *key))))
  return -22;

 if (map == &lru &&
     bpf_map_update_elem(map, key, &init, 0))
  return -1;

 val = bpf_map_lookup_elem(map, key);
 if (!val)
  return -2;

 val->ok_offset = *key;

 wq = &val->w;
 if (bpf_wq_init(wq, map, 0) != 0)
  return -3;

 if (bpf_wq_set_callback(wq, callback_fn, 0))
  return -4;

 if (bpf_wq_start(wq, 0))
  return -5;

 return 0;
}

static int test_hmap_elem_callback(void *map, int *key,
  int (callback_fn)(void *map, int *key, void *value))
{
 struct hmap_elem init = {}, *val;
 struct bpf_wq *wq;

 if ((ok & (1 << *key) ||
     (ok_sleepable & (1 << *key))))
  return -22;

 if (bpf_map_update_elem(map, key, &init, 0))
  return -1;

 val = bpf_map_lookup_elem(map, key);
 if (!val)
  return -2;

 wq = &val->work;
 if (bpf_wq_init(wq, map, 0) != 0)
  return -3;

 if (bpf_wq_set_callback(wq, callback_fn, 0))
  return -4;

 if (bpf_wq_start(wq, 0))
  return -5;

 return 0;
}

/* callback for non sleepable workqueue */
static int wq_callback(void *map, int *key, void *value)
{
 bpf_kfunc_common_test();
 ok |= (1 << *key);
 return 0;
}

/* callback for sleepable workqueue */
static int wq_cb_sleepable(void *map, int *key, void *value)
{
 struct elem *data = (struct elem *)value;
 int offset = data->ok_offset;

 if (*key != offset)
  return 0;

 bpf_kfunc_call_test_sleepable();
 ok_sleepable |= (1 << offset);
 return 0;
}

SEC("tc")
/* test that workqueues can be used from an array */
__retval(0)
long test_call_array_sleepable(void *ctx)
{
 int key = 0;

 return test_elem_callback(&array, &key, wq_cb_sleepable);
}

SEC("syscall")
/* Same test than above but from a sleepable context. */
__retval(0)
long test_syscall_array_sleepable(void *ctx)
{
 int key = 1;

 return test_elem_callback(&array, &key, wq_cb_sleepable);
}

SEC("tc")
/* test that workqueues can be used from a hashmap */
__retval(0)
long test_call_hash_sleepable(void *ctx)
{
 int key = 2;

 return test_hmap_elem_callback(&hmap, &key, wq_callback);
}

SEC("tc")
/* test that workqueues can be used from a hashmap with NO_PREALLOC. */
__retval(0)
long test_call_hash_malloc_sleepable(void *ctx)
{
 int key = 3;

 return test_hmap_elem_callback(&hmap_malloc, &key, wq_callback);
}

SEC("tc")
/* test that workqueues can be used from a LRU map */
__retval(0)
long test_call_lru_sleepable(void *ctx)
{
 int key = 4;

 return test_elem_callback(&lru, &key, wq_callback);
}