/* SPDX-License-Identifier: GPL-2.0-only */

#ifndef __NET_PSP_HELPERS_H
#define __NET_PSP_HELPERS_H

#include <linux/skbuff.h>
#include <linux/rcupdate.h>
#include <linux/udp.h>
#include <net/sock.h>
#include <net/tcp.h>
#include <net/psp/types.h>

struct inet_timewait_sock;

/* Driver-facing API */
struct psp_dev *
psp_dev_create(struct net_device *netdev, struct psp_dev_ops *psd_ops,
               struct psp_dev_caps *psd_caps, void *priv_ptr);
void psp_dev_unregister(struct psp_dev *psd);
bool psp_dev_encapsulate(struct net *net, struct sk_buff *skb, __be32 spi,
                         u8 ver, __be16 sport);
int psp_dev_rcv(struct sk_buff *skb, u16 dev_id, u8 generation, bool strip_icv);

/* Kernel-facing API */
void psp_assoc_put(struct psp_assoc *pas);

static inline void *psp_assoc_drv_data(struct psp_assoc *pas)
{
        return pas->drv_data;
}

#if IS_ENABLED(CONFIG_INET_PSP)
unsigned int psp_key_size(u32 version);
void psp_sk_assoc_free(struct sock *sk);
void psp_twsk_init(struct inet_timewait_sock *tw, const struct sock *sk);
void psp_twsk_assoc_free(struct inet_timewait_sock *tw);
void psp_reply_set_decrypted(const struct sock *sk, struct sk_buff *skb);

static inline struct psp_assoc *psp_sk_assoc(const struct sock *sk)
{
        return rcu_dereference_check(sk->psp_assoc, lockdep_sock_is_held(sk));
}

static inline void
psp_enqueue_set_decrypted(struct sock *sk, struct sk_buff *skb)
{
        struct psp_assoc *pas;

        pas = psp_sk_assoc(sk);
        if (pas && pas->tx.spi)
                skb->decrypted = 1;
}

static inline unsigned long
__psp_skb_coalesce_diff(const struct sk_buff *one, const struct sk_buff *two,
                        unsigned long diffs)
{
        struct psp_skb_ext *a, *b;

        a = skb_ext_find(one, SKB_EXT_PSP);
        b = skb_ext_find(two, SKB_EXT_PSP);

        diffs |= (!!a) ^ (!!b);
        if (!diffs && unlikely(a))
                diffs |= memcmp(a, b, sizeof(*a));
        return diffs;
}

static inline bool
psp_is_allowed_nondata(struct sk_buff *skb, struct psp_assoc *pas)
{
        bool fin = !!(TCP_SKB_CB(skb)->tcp_flags & TCPHDR_FIN);
        u32 end_seq = TCP_SKB_CB(skb)->end_seq;
        u32 seq = TCP_SKB_CB(skb)->seq;
        bool pure_fin;

        pure_fin = fin && end_seq - seq == 1;

        return seq == end_seq || (pure_fin && seq == pas->upgrade_seq);
}

static inline bool
psp_pse_matches_pas(struct psp_skb_ext *pse, struct psp_assoc *pas)
{
        return pse && pas->rx.spi == pse->spi &&
               pas->generation == pse->generation &&
               pas->version == pse->version &&
               pas->dev_id == pse->dev_id;
}

static inline enum skb_drop_reason
__psp_sk_rx_policy_check(struct sk_buff *skb, struct psp_assoc *pas)
{
        struct psp_skb_ext *pse = skb_ext_find(skb, SKB_EXT_PSP);

        if (!pas)
                return pse ? SKB_DROP_REASON_PSP_INPUT : 0;

        if (likely(psp_pse_matches_pas(pse, pas))) {
                if (unlikely(!pas->peer_tx))
                        pas->peer_tx = 1;

                return 0;
        }

        if (!pse) {
                if (!pas->tx.spi ||
                    (!pas->peer_tx && psp_is_allowed_nondata(skb, pas)))
                        return 0;
        }

        return SKB_DROP_REASON_PSP_INPUT;
}

static inline enum skb_drop_reason
psp_sk_rx_policy_check(struct sock *sk, struct sk_buff *skb)
{
        return __psp_sk_rx_policy_check(skb, psp_sk_assoc(sk));
}

static inline enum skb_drop_reason
psp_twsk_rx_policy_check(struct inet_timewait_sock *tw, struct sk_buff *skb)
{
        return __psp_sk_rx_policy_check(skb, rcu_dereference(tw->psp_assoc));
}

static inline struct psp_assoc *psp_sk_get_assoc_rcu(const struct sock *sk)
{
        struct psp_assoc *pas;
        int state;

        state = READ_ONCE(sk->sk_state);
        if (!sk_is_inet(sk) || state == TCP_NEW_SYN_RECV)
                return NULL;

        pas = state == TCP_TIME_WAIT ?
                      rcu_dereference(inet_twsk(sk)->psp_assoc) :
                      rcu_dereference(sk->psp_assoc);
        return pas;
}

static inline struct psp_assoc *psp_skb_get_assoc_rcu(struct sk_buff *skb)
{
        if (!skb->decrypted || !skb->sk)
                return NULL;

        return psp_sk_get_assoc_rcu(skb->sk);
}

static inline unsigned int psp_sk_overhead(const struct sock *sk)
{
        int psp_encap = sizeof(struct udphdr) + PSP_HDR_SIZE + PSP_TRL_SIZE;
        bool has_psp = rcu_access_pointer(sk->psp_assoc);

        return has_psp ? psp_encap : 0;
}
#else
static inline void psp_sk_assoc_free(struct sock *sk) { }
static inline void
psp_twsk_init(struct inet_timewait_sock *tw, const struct sock *sk) { }
static inline void psp_twsk_assoc_free(struct inet_timewait_sock *tw) { }
static inline void
psp_reply_set_decrypted(const struct sock *sk, struct sk_buff *skb) { }

static inline struct psp_assoc *psp_sk_assoc(const struct sock *sk)
{
        return NULL;
}

static inline void
psp_enqueue_set_decrypted(struct sock *sk, struct sk_buff *skb) { }

static inline unsigned long
__psp_skb_coalesce_diff(const struct sk_buff *one, const struct sk_buff *two,
                        unsigned long diffs)
{
        return diffs;
}

static inline enum skb_drop_reason
psp_sk_rx_policy_check(struct sock *sk, struct sk_buff *skb)
{
        return 0;
}

static inline enum skb_drop_reason
psp_twsk_rx_policy_check(struct inet_timewait_sock *tw, struct sk_buff *skb)
{
        return 0;
}

static inline struct psp_assoc *psp_skb_get_assoc_rcu(struct sk_buff *skb)
{
        return NULL;
}

static inline unsigned int psp_sk_overhead(const struct sock *sk)
{
        return 0;
}
#endif

static inline unsigned long
psp_skb_coalesce_diff(const struct sk_buff *one, const struct sk_buff *two)
{
        return __psp_skb_coalesce_diff(one, two, 0);
}

#endif /* __NET_PSP_HELPERS_H */