BPF_CALL_4(bpf_sk_storage_get_tracing, struct bpf_map *, map, struct sock *, sk,
           void *, value, u64, flags)
 {
-       if (!in_serving_softirq() && !in_task())
+       if (in_irq() || in_nmi())
                return (unsigned long)NULL;
 
        return (unsigned long)____bpf_sk_storage_get(map, sk, value, flags);
 BPF_CALL_2(bpf_sk_storage_delete_tracing, struct bpf_map *, map,
           struct sock *, sk)
 {
-       if (!in_serving_softirq() && !in_task())
+       if (in_irq() || in_nmi())
                return -EPERM;
 
        return ____bpf_sk_storage_delete(map, sk);