Ratchet: Prep for next key

router/java/src/net/i2p/router/crypto/ratchet/ECIESAEADEngine.java

@@ -760,6 +760,7 @@ public final class ECIESAEADEngine {
     private class PLCallback implements RatchetPayload.PayloadCallback {
         public final List<GarlicClove> cloveSet = new ArrayList<GarlicClove>(3);
         public long datetime;
+        public NextSessionKey nextKey;
 
         public void gotDateTime(long time) {
             if (_log.shouldDebug())
@@ -780,6 +781,12 @@ public final class ECIESAEADEngine {
             cloveSet.add(clove);
         }
 
+        public void gotNextKey(NextSessionKey next) {
+            if (_log.shouldDebug())
+                _log.debug("Got NEXTKEY block: " + next);
+            nextKey = next;
+        }
+
         public void gotTermination(int reason, long count) {
             if (_log.shouldDebug())
                 _log.debug("Got TERMINATION block, reason: " + reason + " count: " + count);

router/java/src/net/i2p/router/crypto/ratchet/NextSessionKey.java

@@ -0,0 +1,31 @@
+package net.i2p.router.crypto.ratchet;
+
+import net.i2p.data.SessionKey;
+
+/**
+ * A session key and key ID.
+ *
+ * @since 0.9.44
+ */
+class NextSessionKey extends SessionKey {
+    private final int _id;
+
+    public NextSessionKey(byte[] data, int id) {
+        super(data);
+        _id = id;
+    }
+
+    public int getID() {
+        return _id;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder buf = new StringBuilder(64);
+        buf.append("[NextSessionKey: ");
+        buf.append(toBase64());
+        buf.append(" ID: ").append(_id);
+        buf.append(']');
+        return buf.toString();
+    }
+}

router/java/src/net/i2p/router/crypto/ratchet/RatchetEntry.java

@@ -1,23 +1,37 @@
 package net.i2p.router.crypto.ratchet;
 
-import net.i2p.data.SessionTag;
+import net.i2p.data.SessionKey;
 
 /**
+ *  Simple object with outbound tag, key, and nonce,
+ *  and an optional next key.
+ *  The object returned from SKM.consumeNextAvailableTag() to the engine encrypt.
  *
  *  @since 0.9.44
  */
 class RatchetEntry {
     public final RatchetSessionTag tag;
     public final SessionKeyAndNonce key;
+    public final int keyID;
+    public final int pn;
+    public final NextSessionKey nextKey;
 
     /** outbound - calculated key */
-    public RatchetEntry(RatchetSessionTag tag, SessionKeyAndNonce key) {
+    public RatchetEntry(RatchetSessionTag tag, SessionKeyAndNonce key, int keyID, int pn) {
+        this(tag, key, keyID, pn, null);
+    }
+
+    public RatchetEntry(RatchetSessionTag tag, SessionKeyAndNonce key, int keyID, int pn,
+                        NextSessionKey nextKey) {
         this.tag = tag;
         this.key = key;
+        this.keyID = keyID;
+        this.pn = pn;
+        this.nextKey = nextKey;
     }
 
     @Override
     public String toString() {
-        return "RatchetEntry[" + tag + ' ' + key + ']';
+        return "RatchetEntry[" + tag.toBase64() + ' ' + key + ']';
     }
 }

router/java/src/net/i2p/router/crypto/ratchet/RatchetPayload.java

@@ -56,6 +56,11 @@ class RatchetPayload {
          */
         public void gotTermination(int reason, long lastReceived);
 
+        /**
+         *  @param nextKey the next one
+         */
+        public void gotNextKey(NextSessionKey nextKey);
+
         /**
          *  For stats.
          *  @param paddingLength the number of padding bytes, not including the 3-byte block header
@@ -117,6 +122,16 @@ class RatchetPayload {
                     cb.gotGarlic(clove);
                     break;
 
+                case BLOCK_NEXTKEY:
+                    if (len != 34)
+                        throw new IOException("Bad length for NEXTKEY: " + len);
+                    int id = (int) DataHelper.fromLong(payload, i, 2);
+                    byte[] data = new byte[32];
+                    System.arraycopy(payload, i + 2, data, 0, 32);
+                    NextSessionKey nsk = new NextSessionKey(data, id);
+                    cb.gotNextKey(nsk);
+                    break;
+
                 case BLOCK_TERMINATION:
                     if (isHandshake)
                         throw new IOException("Illegal block in handshake: " + type);
@@ -284,6 +299,25 @@ class RatchetPayload {
         }
     }
 
+    public static class NextKeyBlock extends Block {
+        private final NextSessionKey next;
+
+        public NextKeyBlock(NextSessionKey nextKey) {
+            super(BLOCK_NEXTKEY);
+            next = nextKey;
+        }
+
+        public int getDataLength() {
+            return 34;
+        }
+
+        public int writeData(byte[] tgt, int off) {
+            DataHelper.toLong(tgt, off, 2, next.getID());
+            System.arraycopy(next.getData(), 0, tgt, off + 2, 32);
+            return off + 34;
+        }
+    }
+
     public static class TerminationBlock extends Block {
         private final byte rsn;
         private final long rcvd;

router/java/src/net/i2p/router/crypto/ratchet/RatchetSKM.java

@@ -1060,7 +1060,8 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
                             if (tag != null) {
                                 set.setDate(now);
                                 SessionKeyAndNonce skn = set.consumeNextKey();
-                                return new RatchetEntry(tag, skn);
+                                // TODO key ID and PN
+                                return new RatchetEntry(tag, skn, 0, 0);
                             } else if (_log.shouldInfo()) {
                                 _log.info("Removing empty " + set);
                             }

router/java/src/net/i2p/router/crypto/ratchet/RatchetTagSet.java

@@ -354,6 +354,8 @@ class RatchetTagSet implements TagSetHandle {
 
     private void storeNextTag() {
         RatchetSessionTag tag = consumeNext();
+        if (tag == null)
+            return;
         _sessionTags.put(_lastTag, tag);
         if (_lsnr != null)
             _lsnr.addTag(tag, this);
@@ -363,9 +365,11 @@ class RatchetTagSet implements TagSetHandle {
      *  For outbound only.
      *  Call before consumeNextKey();
      *
-     *  @return a tag or null
+     *  @return a tag or null if we ran out
      */
     public RatchetSessionTag consumeNext() {
+        if (_lastTag >= MAX)
+            return null;
         byte[] tmp = new byte[32];
         hkdf.calculate(_sesstag_ck, _sesstag_constant, INFO_4, _sesstag_ck, tmp, 0);
         byte[] tag = new byte[TAGLEN];