1 files changed, 33 insertions, 13 deletions

diff --git a/tunnel/src/main/java/com/wireguard/android/backend/Statistics.java b/tunnel/src/main/java/com/wireguard/android/backend/Statistics.java
index 5d658019..322e766f 100644
--- a/tunnel/src/main/java/com/wireguard/android/backend/Statistics.java
+++ b/tunnel/src/main/java/com/wireguard/android/backend/Statistics.java
@@ -11,6 +11,7 @@ import android.util.Pair;
 import com.wireguard.crypto.Key;
 import com.wireguard.util.NonNullForAll;
 
+import java.time.LocalDateTime;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -19,7 +20,19 @@ import java.util.Map;
  */
 @NonNullForAll
 public class Statistics {
-    private final Map<Key, Pair<Long, Long>> peerBytes = new HashMap<>();
+    private static class Stat {
+        long rx;
+        long tx;
+        LocalDateTime lastHandshake;
+
+        Stat(long rx, long tx, LocalDateTime lastHandshake) {
+            this.rx = rx;
+            this.tx = tx;
+            this.lastHandshake = lastHandshake;
+        }
+    }
+
+    private final Map<Key, Stat> peerBytes = new HashMap<>();
     private long lastTouched = SystemClock.elapsedRealtime();
 
     Statistics() {
@@ -34,8 +47,8 @@ public class Statistics {
      * @param tx  The transmitted traffic for the {@link com.wireguard.config.Peer} referenced by
      *            the provided {@link Key}. This value is in bytes.
      */
-    void add(final Key key, final long rx, final long tx) {
-        peerBytes.put(key, Pair.create(rx, tx));
+    void add(final Key key, final long rx, final long tx, final LocalDateTime lastHandshake) {
+        peerBytes.put(key, new Stat(rx, tx, lastHandshake));
         lastTouched = SystemClock.elapsedRealtime();
     }
 
@@ -56,10 +69,10 @@ public class Statistics {
      * @return a long representing the number of bytes received by this peer.
      */
     public long peerRx(final Key peer) {
-        final Pair<Long, Long> rxTx = peerBytes.get(peer);
-        if (rxTx == null)
+        final Stat stat = peerBytes.get(peer);
+        if (stat == null)
             return 0;
-        return rxTx.first;
+        return stat.rx;
     }
 
     /**
@@ -70,10 +83,17 @@ public class Statistics {
      * @return a long representing the number of bytes transmitted by this peer.
      */
     public long peerTx(final Key peer) {
-        final Pair<Long, Long> rxTx = peerBytes.get(peer);
-        if (rxTx == null)
+        final Stat stat = peerBytes.get(peer);
+        if (stat == null)
             return 0;
-        return rxTx.second;
+        return stat.tx;
+    }
+
+    public LocalDateTime peerLastHandshake(final Key peer) {
+        final Stat stat = peerBytes.get(peer);
+        if (stat == null)
+            return null;
+        return stat.lastHandshake;
     }
 
     /**
@@ -93,8 +113,8 @@ public class Statistics {
      */
     public long totalRx() {
         long rx = 0;
-        for (final Pair<Long, Long> val : peerBytes.values()) {
-            rx += val.first;
+        for (final Stat val : peerBytes.values()) {
+            rx += val.rx;
         }
         return rx;
     }
@@ -106,8 +126,8 @@ public class Statistics {
      */
     public long totalTx() {
         long tx = 0;
-        for (final Pair<Long, Long> val : peerBytes.values()) {
-            tx += val.second;
+        for (final Stat val : peerBytes.values()) {
+            tx += val.tx;
         }
         return tx;
     }