#include <crypto/aead.h>
 #include <crypto/aes.h>
 #include "crypto.h"
+#include "msg.h"
+#include "bcast.h"
 
 #define TIPC_TX_GRACE_PERIOD   msecs_to_jiffies(5000) /* 5s */
 #define TIPC_TX_LASTING_TIME   msecs_to_jiffies(10000) /* 10s */
 
 /* Max TFMs number per key */
 int sysctl_tipc_max_tfms __read_mostly = TIPC_MAX_TFMS_DEF;
+/* Key exchange switch, default: on */
+int sysctl_tipc_key_exchange_enabled __read_mostly = 1;
 
 /**
  * struct tipc_key - TIPC keys' status indicator
  * @mode: crypto mode is applied to the key
  * @hint[]: a hint for user key
  * @rcu: struct rcu_head
+ * @key: the aead key
+ * @gen: the key's generation
  * @seqno: the key seqno (cluster scope)
  * @refcnt: the key reference counter
  */
        u8 mode;
        char hint[2 * TIPC_AEAD_HINT_LEN + 1];
        struct rcu_head rcu;
+       struct tipc_aead_key *key;
+       u16 gen;
 
        atomic64_t seqno ____cacheline_aligned;
        refcount_t refcnt ____cacheline_aligned;
  * @node: TIPC node (RX)
  * @aead: array of pointers to AEAD keys for encryption/decryption
  * @peer_rx_active: replicated peer RX active key index
+ * @key_gen: TX/RX key generation
  * @key: the key states
+ * @skey_mode: session key's mode
+ * @skey: received session key
+ * @wq: common workqueue on TX crypto
+ * @work: delayed work sched for TX/RX
+ * @key_distr: key distributing state
  * @stats: the crypto statistics
  * @name: the crypto name
  * @sndnxt: the per-peer sndnxt (TX)
  * @working: the crypto is working or not
  * @key_master: flag indicates if master key exists
  * @legacy_user: flag indicates if a peer joins w/o master key (for bwd comp.)
+ * @nokey: no key indication
  * @lock: tipc_key lock
  */
 struct tipc_crypto {
        struct tipc_node *node;
        struct tipc_aead __rcu *aead[KEY_MAX + 1];
        atomic_t peer_rx_active;
+       u16 key_gen;
        struct tipc_key key;
+       u8 skey_mode;
+       struct tipc_aead_key *skey;
+       struct workqueue_struct *wq;
+       struct delayed_work work;
+#define KEY_DISTR_SCHED                1
+#define KEY_DISTR_COMPL                2
+       atomic_t key_distr;
+
        struct tipc_crypto_stats __percpu *stats;
        char name[48];
 
                        u8 working:1;
                        u8 key_master:1;
                        u8 legacy_user:1;
+                       u8 nokey: 1;
                };
                u8 flags;
        };
 static char *tipc_crypto_key_dump(struct tipc_crypto *c, char *buf);
 static char *tipc_key_change_dump(struct tipc_key old, struct tipc_key new,
                                  char *buf);
+static int tipc_crypto_key_xmit(struct net *net, struct tipc_aead_key *skey,
+                               u16 gen, u8 mode, u32 dnode);
+static bool tipc_crypto_key_rcv(struct tipc_crypto *rx, struct tipc_msg *hdr);
+static void tipc_crypto_work_rx(struct work_struct *work);
+
 #define is_tx(crypto) (!(crypto)->node)
 #define is_rx(crypto) (!is_tx(crypto))
 
                kfree(head);
        }
        free_percpu(aead->tfm_entry);
+       kzfree(aead->key);
        kfree(aead);
 }
 
        tmp->mode = mode;
        tmp->cloned = NULL;
        tmp->authsize = TIPC_AES_GCM_TAG_SIZE;
+       tmp->key = kmemdup(ukey, tipc_aead_key_size(ukey), GFP_KERNEL);
        memcpy(&tmp->salt, ukey->key + keylen, TIPC_AES_GCM_SALT_SIZE);
        atomic_set(&tmp->users, 0);
        atomic64_set(&tmp->seqno, 0);
        ehdr->tx_key = tx_key;
        ehdr->destined = (__rx) ? 1 : 0;
        ehdr->rx_key_active = (__rx) ? __rx->key.active : 0;
-       ehdr->rx_nokey = (__rx) ? !__rx->key.keys : 0;
+       ehdr->rx_nokey = (__rx) ? __rx->nokey : 0;
        ehdr->master_key = aead->crypto->key_master;
        ehdr->reserved_1 = 0;
        ehdr->reserved_2 = 0;
 
 attach:
        aead->crypto = c;
+       aead->gen = (is_tx(c)) ? ++c->key_gen : c->key_gen;
        tipc_aead_rcu_replace(c->aead[new_key], aead, &c->lock);
        if (likely(c->key.keys != key.keys))
                tipc_crypto_key_set_state(c, key.passive, key.active,
                                          key.pending);
        c->working = 1;
+       c->nokey = 0;
        c->key_master |= master_key;
        rc = new_key;
 
 
 void tipc_crypto_key_flush(struct tipc_crypto *c)
 {
+       struct tipc_crypto *tx, *rx;
        int k;
 
        spin_lock_bh(&c->lock);
+       if (is_rx(c)) {
+               /* Try to cancel pending work */
+               rx = c;
+               tx = tipc_net(rx->net)->crypto_tx;
+               if (cancel_delayed_work(&rx->work)) {
+                       kfree(rx->skey);
+                       rx->skey = NULL;
+                       atomic_xchg(&rx->key_distr, 0);
+                       tipc_node_put(rx->node);
+               }
+               /* RX stopping => decrease TX key users if any */
+               k = atomic_xchg(&rx->peer_rx_active, 0);
+               if (k) {
+                       tipc_aead_users_dec(tx->aead[k], 0);
+                       /* Mark the point TX key users changed */
+                       tx->timer1 = jiffies;
+               }
+       }
+
        c->flags = 0;
        tipc_crypto_key_set_state(c, 0, 0, 0);
        for (k = KEY_MIN; k <= KEY_MAX; k++)
                tipc_crypto_key_detach(c->aead[k], &c->lock);
-       atomic_set(&c->peer_rx_active, 0);
        atomic64_set(&c->sndnxt, 0);
        spin_unlock_bh(&c->lock);
 }
  * decreased correspondingly.
  *
  * It also considers if peer has no key, then we need to make own master key
- * (if any) taking over i.e. starting grace period.
+ * (if any) taking over i.e. starting grace period and also trigger key
+ * distributing process.
  *
  * The "per-peer" sndnxt is also reset when the peer key has switched.
  */
        struct tipc_msg *hdr = buf_msg(skb);
        u32 self = tipc_own_addr(rx->net);
        u8 cur, new;
+       unsigned long delay;
 
        /* Update RX 'key_master' flag according to peer, also mark "legacy" if
         * a peer has no master key.
                return;
 
        /* Case 1: Peer has no keys, let's make master key take over */
-       if (ehdr->rx_nokey)
+       if (ehdr->rx_nokey) {
                /* Set or extend grace period */
                tx->timer2 = jiffies;
+               /* Schedule key distributing for the peer if not yet */
+               if (tx->key.keys &&
+                   !atomic_cmpxchg(&rx->key_distr, 0, KEY_DISTR_SCHED)) {
+                       get_random_bytes(&delay, 2);
+                       delay %= 5;
+                       delay = msecs_to_jiffies(500 * ++delay);
+                       if (queue_delayed_work(tx->wq, &rx->work, delay))
+                               tipc_node_get(rx->node);
+               }
+       } else {
+               /* Cancel a pending key distributing if any */
+               atomic_xchg(&rx->key_distr, 0);
+       }
 
        /* Case 2: Peer RX active key has changed, let's update own TX users */
        cur = atomic_read(&rx->peer_rx_active);
        if (!c)
                return -ENOMEM;
 
+       /* Allocate workqueue on TX */
+       if (!node) {
+               c->wq = alloc_ordered_workqueue("tipc_crypto", 0);
+               if (!c->wq) {
+                       kfree(c);
+                       return -ENOMEM;
+               }
+       }
+
        /* Allocate statistic structure */
        c->stats = alloc_percpu_gfp(struct tipc_crypto_stats, GFP_ATOMIC);
        if (!c->stats) {
        c->flags = 0;
        c->net = net;
        c->node = node;
+       get_random_bytes(&c->key_gen, 2);
        tipc_crypto_key_set_state(c, 0, 0, 0);
+       atomic_set(&c->key_distr, 0);
        atomic_set(&c->peer_rx_active, 0);
        atomic64_set(&c->sndnxt, 0);
        c->timer1 = jiffies;
                  (is_rx(c)) ? tipc_node_get_id_str(c->node) :
                               tipc_own_id_string(c->net));
 
+       if (is_rx(c))
+               INIT_DELAYED_WORK(&c->work, tipc_crypto_work_rx);
+
        *crypto = c;
        return 0;
 }
 
 void tipc_crypto_stop(struct tipc_crypto **crypto)
 {
-       struct tipc_crypto *c = *crypto, *tx, *rx;
+       struct tipc_crypto *c = *crypto;
        u8 k;
 
        if (!c)
                return;
 
-       rcu_read_lock();
-       /* RX stopping? => decrease TX key users if any */
-       if (is_rx(c)) {
-               rx = c;
-               tx = tipc_net(rx->net)->crypto_tx;
-               k = atomic_read(&rx->peer_rx_active);
-               if (k) {
-                       tipc_aead_users_dec(tx->aead[k], 0);
-                       /* Mark the point TX key users changed */
-                       tx->timer1 = jiffies;
-               }
-       }
+       /* Flush any queued works & destroy wq */
+       if (is_tx(c))
+               destroy_workqueue(c->wq);
 
        /* Release AEAD keys */
+       rcu_read_lock();
        for (k = KEY_MIN; k <= KEY_MAX; k++)
                tipc_aead_put(rcu_dereference(c->aead[k]));
        rcu_read_unlock();
                }
                if (user == LINK_CONFIG ||
                    (user == LINK_PROTOCOL && type == RESET_MSG) ||
+                   (user == MSG_CRYPTO && type == KEY_DISTR_MSG) ||
                    time_before(jiffies, tx->timer2 + TIPC_TX_GRACE_PERIOD)) {
                        if (__rx && __rx->key_master &&
                            !atomic_read(&__rx->peer_rx_active))
        struct tipc_aead *aead = NULL;
        struct tipc_key key;
        int rc = -ENOKEY;
-       u8 tx_key;
+       u8 tx_key, n;
 
        tx_key = ((struct tipc_ehdr *)(*skb)->data)->tx_key;
 
                if (rc == -ENOKEY) {
                        kfree_skb(*skb);
                        *skb = NULL;
-                       if (rx)
+                       if (rx) {
+                               /* Mark rx->nokey only if we dont have a
+                                * pending received session key, nor a newer
+                                * one i.e. in the next slot.
+                                */
+                               n = key_next(tx_key);
+                               rx->nokey = !(rx->skey ||
+                                             rcu_access_pointer(rx->aead[n]));
+                               pr_debug_ratelimited("%s: nokey %d, key %d/%x\n",
+                                                    rx->name, rx->nokey,
+                                                    tx_key, rx->key.keys);
                                tipc_node_put(rx->node);
+                       }
                        this_cpu_inc(stats->stat[STAT_NOKEYS]);
                        return rc;
                } else if (rc == -EBADMSG) {
        i += scnprintf(buf + i, 32 - i, "]");
        return buf;
 }
+
+/**
+ * tipc_crypto_msg_rcv - Common 'MSG_CRYPTO' processing point
+ * @net: the struct net
+ * @skb: the receiving message buffer
+ */
+void tipc_crypto_msg_rcv(struct net *net, struct sk_buff *skb)
+{
+       struct tipc_crypto *rx;
+       struct tipc_msg *hdr;
+
+       if (unlikely(skb_linearize(skb)))
+               goto exit;
+
+       hdr = buf_msg(skb);
+       rx = tipc_node_crypto_rx_by_addr(net, msg_prevnode(hdr));
+       if (unlikely(!rx))
+               goto exit;
+
+       switch (msg_type(hdr)) {
+       case KEY_DISTR_MSG:
+               if (tipc_crypto_key_rcv(rx, hdr))
+                       goto exit;
+               break;
+       default:
+               break;
+       }
+
+       tipc_node_put(rx->node);
+
+exit:
+       kfree_skb(skb);
+}
+
+/**
+ * tipc_crypto_key_distr - Distribute a TX key
+ * @tx: the TX crypto
+ * @key: the key's index
+ * @dest: the destination tipc node, = NULL if distributing to all nodes
+ *
+ * Return: 0 in case of success, otherwise < 0
+ */
+int tipc_crypto_key_distr(struct tipc_crypto *tx, u8 key,
+                         struct tipc_node *dest)
+{
+       struct tipc_aead *aead;
+       u32 dnode = tipc_node_get_addr(dest);
+       int rc = -ENOKEY;
+
+       if (!sysctl_tipc_key_exchange_enabled)
+               return 0;
+
+       if (key) {
+               rcu_read_lock();
+               aead = tipc_aead_get(tx->aead[key]);
+               if (likely(aead)) {
+                       rc = tipc_crypto_key_xmit(tx->net, aead->key,
+                                                 aead->gen, aead->mode,
+                                                 dnode);
+                       tipc_aead_put(aead);
+               }
+               rcu_read_unlock();
+       }
+
+       return rc;
+}
+
+/**
+ * tipc_crypto_key_xmit - Send a session key
+ * @net: the struct net
+ * @skey: the session key to be sent
+ * @gen: the key's generation
+ * @mode: the key's mode
+ * @dnode: the destination node address, = 0 if broadcasting to all nodes
+ *
+ * The session key 'skey' is packed in a TIPC v2 'MSG_CRYPTO/KEY_DISTR_MSG'
+ * as its data section, then xmit-ed through the uc/bc link.
+ *
+ * Return: 0 in case of success, otherwise < 0
+ */
+static int tipc_crypto_key_xmit(struct net *net, struct tipc_aead_key *skey,
+                               u16 gen, u8 mode, u32 dnode)
+{
+       struct sk_buff_head pkts;
+       struct tipc_msg *hdr;
+       struct sk_buff *skb;
+       u16 size, cong_link_cnt;
+       u8 *data;
+       int rc;
+
+       size = tipc_aead_key_size(skey);
+       skb = tipc_buf_acquire(INT_H_SIZE + size, GFP_ATOMIC);
+       if (!skb)
+               return -ENOMEM;
+
+       hdr = buf_msg(skb);
+       tipc_msg_init(tipc_own_addr(net), hdr, MSG_CRYPTO, KEY_DISTR_MSG,
+                     INT_H_SIZE, dnode);
+       msg_set_size(hdr, INT_H_SIZE + size);
+       msg_set_key_gen(hdr, gen);
+       msg_set_key_mode(hdr, mode);
+
+       data = msg_data(hdr);
+       *((__be32 *)(data + TIPC_AEAD_ALG_NAME)) = htonl(skey->keylen);
+       memcpy(data, skey->alg_name, TIPC_AEAD_ALG_NAME);
+       memcpy(data + TIPC_AEAD_ALG_NAME + sizeof(__be32), skey->key,
+              skey->keylen);
+
+       __skb_queue_head_init(&pkts);
+       __skb_queue_tail(&pkts, skb);
+       if (dnode)
+               rc = tipc_node_xmit(net, &pkts, dnode, 0);
+       else
+               rc = tipc_bcast_xmit(net, &pkts, &cong_link_cnt);
+
+       return rc;
+}
+
+/**
+ * tipc_crypto_key_rcv - Receive a session key
+ * @rx: the RX crypto
+ * @hdr: the TIPC v2 message incl. the receiving session key in its data
+ *
+ * This function retrieves the session key in the message from peer, then
+ * schedules a RX work to attach the key to the corresponding RX crypto.
+ *
+ * Return: "true" if the key has been scheduled for attaching, otherwise
+ * "false".
+ */
+static bool tipc_crypto_key_rcv(struct tipc_crypto *rx, struct tipc_msg *hdr)
+{
+       struct tipc_crypto *tx = tipc_net(rx->net)->crypto_tx;
+       struct tipc_aead_key *skey = NULL;
+       u16 key_gen = msg_key_gen(hdr);
+       u16 size = msg_data_sz(hdr);
+       u8 *data = msg_data(hdr);
+
+       spin_lock(&rx->lock);
+       if (unlikely(rx->skey || (key_gen == rx->key_gen && rx->key.keys))) {
+               pr_err("%s: key existed <%p>, gen %d vs %d\n", rx->name,
+                      rx->skey, key_gen, rx->key_gen);
+               goto exit;
+       }
+
+       /* Allocate memory for the key */
+       skey = kmalloc(size, GFP_ATOMIC);
+       if (unlikely(!skey)) {
+               pr_err("%s: unable to allocate memory for skey\n", rx->name);
+               goto exit;
+       }
+
+       /* Copy key from msg data */
+       skey->keylen = ntohl(*((__be32 *)(data + TIPC_AEAD_ALG_NAME)));
+       memcpy(skey->alg_name, data, TIPC_AEAD_ALG_NAME);
+       memcpy(skey->key, data + TIPC_AEAD_ALG_NAME + sizeof(__be32),
+              skey->keylen);
+
+       /* Sanity check */
+       if (unlikely(size != tipc_aead_key_size(skey))) {
+               kfree(skey);
+               skey = NULL;
+               goto exit;
+       }
+
+       rx->key_gen = key_gen;
+       rx->skey_mode = msg_key_mode(hdr);
+       rx->skey = skey;
+       rx->nokey = 0;
+       mb(); /* for nokey flag */
+
+exit:
+       spin_unlock(&rx->lock);
+
+       /* Schedule the key attaching on this crypto */
+       if (likely(skey && queue_delayed_work(tx->wq, &rx->work, 0)))
+               return true;
+
+       return false;
+}
+
+/**
+ * tipc_crypto_work_rx - Scheduled RX works handler
+ * @work: the struct RX work
+ *
+ * The function processes the previous scheduled works i.e. distributing TX key
+ * or attaching a received session key on RX crypto.
+ */
+static void tipc_crypto_work_rx(struct work_struct *work)
+{
+       struct delayed_work *dwork = to_delayed_work(work);
+       struct tipc_crypto *rx = container_of(dwork, struct tipc_crypto, work);
+       struct tipc_crypto *tx = tipc_net(rx->net)->crypto_tx;
+       unsigned long delay = msecs_to_jiffies(5000);
+       bool resched = false;
+       u8 key;
+       int rc;
+
+       /* Case 1: Distribute TX key to peer if scheduled */
+       if (atomic_cmpxchg(&rx->key_distr,
+                          KEY_DISTR_SCHED,
+                          KEY_DISTR_COMPL) == KEY_DISTR_SCHED) {
+               /* Always pick the newest one for distributing */
+               key = tx->key.pending ?: tx->key.active;
+               rc = tipc_crypto_key_distr(tx, key, rx->node);
+               if (unlikely(rc))
+                       pr_warn("%s: unable to distr key[%d] to %s, err %d\n",
+                               tx->name, key, tipc_node_get_id_str(rx->node),
+                               rc);
+
+               /* Sched for key_distr releasing */
+               resched = true;
+       } else {
+               atomic_cmpxchg(&rx->key_distr, KEY_DISTR_COMPL, 0);
+       }
+
+       /* Case 2: Attach a pending received session key from peer if any */
+       if (rx->skey) {
+               rc = tipc_crypto_key_init(rx, rx->skey, rx->skey_mode, false);
+               if (unlikely(rc < 0))
+                       pr_warn("%s: unable to attach received skey, err %d\n",
+                               rx->name, rc);
+               switch (rc) {
+               case -EBUSY:
+               case -ENOMEM:
+                       /* Resched the key attaching */
+                       resched = true;
+                       break;
+               default:
+                       synchronize_rcu();
+                       kfree(rx->skey);
+                       rx->skey = NULL;
+                       break;
+               }
+       }
+
+       if (resched && queue_delayed_work(tx->wq, &rx->work, delay))
+               return;
+
+       tipc_node_put(rx->node);
+}