From 01d6cea0175a620edcc2744f0909549ff1081af2 Mon Sep 17 00:00:00 2001
From: zzz <zzz@mail.i2p>
Date: Tue, 5 Nov 2019 17:15:37 +0000
Subject: [PATCH] Ratchet: Simplify lookup of OB session from IB NSR

---
 .../i2p/router/crypto/ratchet/RatchetSKM.java | 48 +++++++++++--------
 1 file changed, 27 insertions(+), 21 deletions(-)

diff --git a/router/java/src/net/i2p/router/crypto/ratchet/RatchetSKM.java b/router/java/src/net/i2p/router/crypto/ratchet/RatchetSKM.java
index 95df6f58ed..1f42b1d6c7 100644
--- a/router/java/src/net/i2p/router/crypto/ratchet/RatchetSKM.java
+++ b/router/java/src/net/i2p/router/crypto/ratchet/RatchetSKM.java
@@ -244,28 +244,26 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
                 }
                 boolean found = false;
                 for (OutboundSession sess : pending) {
-                    for (RatchetTagSet ts : sess.getTagSets()) {
-                        if (ts.getHandshakeState().equals(oldState)) {
-                            if (!found) {
-                                found = true;
-                                sess.updateSession(state);
-                                boolean ok = addSession(sess);
-                                if (_log.shouldDebug()) {
-                                    if (ok)
-                                        _log.debug("Update Alice session from NSR to ES for "  + target);
-                                    else
-                                        _log.debug("Session already updated from NSR to ES for "  + target);
-                                }
-                            } else {
-                                if (_log.shouldDebug())
-                                    _log.debug("Dup tagset " + ts + " for "  + target);
+                    if (oldState.equals(sess.getHandshakeState())) {
+                        if (!found) {
+                            found = true;
+                            sess.updateSession(state);
+                            boolean ok = addSession(sess);
+                            if (_log.shouldDebug()) {
+                                if (ok)
+                                    _log.debug("Update Alice session from NSR to ES for "  + target);
+                                else
+                                    _log.debug("Session already updated from NSR to ES for "  + target);
                             }
                         } else {
-                            // TODO
-                            // remove old tags
                             if (_log.shouldDebug())
-                                _log.debug("Remove tagset " + ts + " for "  + target);
+                                _log.debug("Other pending session " + sess + " for "  + target);
                         }
+                    } else {
+                        // TODO
+                        // remove old tags
+                        if (_log.shouldDebug())
+                            _log.debug("Other pending session " + sess + " for "  + target);
                     }
                 }
                 if (found) {
@@ -786,6 +784,7 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
      */
     private class OutboundSession {
         private final PublicKey _target;
+        private final HandshakeState _state;
         private SessionKey _currentKey;
         private final long _established;
         private long _lastUsed;
@@ -839,6 +838,7 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
                                                          rk, tk,
                                                          _established, _sentTagSetID.getAndIncrement());
                 _tagSets.add(tagset);
+                _state = null;
                 if (_log.shouldDebug())
                     _log.debug("New OB Session, rk = " + rk + " tk = " + tk + " 1st tagset: " + tagset);
             } else {
@@ -847,9 +847,8 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
                 RatchetTagSet tagset = new RatchetTagSet(_hkdf, RatchetSKM.this, state,
                                                          rk, tk,
                                                          _established, _rcvTagSetID.getAndIncrement(), 5, 5);
-                // store the IB tagset as OB so we can lookup the state
-                // TODO just store the state
-                _unackedTagSets.add(tagset);
+-               // store the state so we can find the right session when we receive the NSR
+                _state = state;
                 if (_log.shouldDebug())
                     _log.debug("New IB Session, rk = " + rk + " tk = " + tk + " 1st tagset: " + tagset);
             }
@@ -986,6 +985,13 @@ public class RatchetSKM extends SessionKeyManager implements SessionTagListener
             return _target;
         }
 
+        /**
+         *  Original outbound state, null for inbound.
+         */
+        public HandshakeState getHandshakeState() {
+            return _state;
+        }
+
         public SessionKey getCurrentKey() {
             return _currentKey;
         }