Implement statistics

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

5 files changed, 221 insertions, 6 deletions

diff --git a/app/src/main/java/com/wireguard/android/backend/GoBackend.java b/app/src/main/java/com/wireguard/android/backend/GoBackend.java
index e85f2b0d..3e8e1ec9 100644
--- a/app/src/main/java/com/wireguard/android/backend/GoBackend.java
+++ b/app/src/main/java/com/wireguard/android/backend/GoBackend.java
@@ -24,6 +24,8 @@ import com.wireguard.android.util.SharedLibraryLoader;
 import com.wireguard.config.Config;
 import com.wireguard.config.InetNetwork;
 import com.wireguard.config.Peer;
+import com.wireguard.crypto.Key;
+import com.wireguard.crypto.KeyFormatException;
 
 import java.net.InetAddress;
 import java.util.Collections;
@@ -47,6 +49,8 @@ public final class GoBackend implements Backend {
         this.context = context;
     }
 
+    private static native String wgGetConfig(int handle);
+
     private static native int wgGetSocketV4(int handle);
 
     private static native int wgGetSocketV6(int handle);
@@ -90,7 +94,45 @@ public final class GoBackend implements Backend {
 
     @Override
     public Statistics getStatistics(final Tunnel tunnel) {
-        return new Statistics();
+        final Statistics stats = new Statistics();
+        if (tunnel != currentTunnel) {
+            return stats;
+        }
+        final String config = wgGetConfig(currentTunnelHandle);
+        Key key = null;
+        long rx = 0, tx = 0;
+        for (final String line : config.split("\\n")) {
+            if (line.startsWith("public_key=")) {
+                if (key != null)
+                    stats.add(key, rx, tx);
+                rx = 0;
+                tx = 0;
+                try {
+                    key = Key.fromHex(line.substring(11));
+                } catch (final KeyFormatException ignored) {
+                    key = null;
+                }
+            } else if (line.startsWith("rx_bytes=")) {
+                if (key == null)
+                    continue;
+                try {
+                    rx = Long.parseLong(line.substring(9));
+                } catch (final NumberFormatException ignored) {
+                    rx = 0;
+                }
+            } else if (line.startsWith("tx_bytes=")) {
+                if (key == null)
+                    continue;
+                try {
+                    tx = Long.parseLong(line.substring(9));
+                } catch (final NumberFormatException ignored) {
+                    tx = 0;
+                }
+            }
+        }
+        if (key != null)
+            stats.add(key, rx, tx);
+        return stats;
     }
 
     @Override
diff --git a/app/src/main/java/com/wireguard/android/backend/WgQuickBackend.java b/app/src/main/java/com/wireguard/android/backend/WgQuickBackend.java
index 71427d8a..99b90af2 100644
--- a/app/src/main/java/com/wireguard/android/backend/WgQuickBackend.java
+++ b/app/src/main/java/com/wireguard/android/backend/WgQuickBackend.java
@@ -15,11 +15,13 @@ import com.wireguard.android.model.Tunnel;
 import com.wireguard.android.model.Tunnel.State;
 import com.wireguard.android.model.Tunnel.Statistics;
 import com.wireguard.config.Config;
+import com.wireguard.crypto.Key;
 
 import java.io.File;
 import java.io.FileOutputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
@@ -83,7 +85,24 @@ public final class WgQuickBackend implements Backend {
 
     @Override
     public Statistics getStatistics(final Tunnel tunnel) {
-        return new Statistics();
+        final Statistics stats = new Statistics();
+        final Collection<String> output = new ArrayList<>();
+        try {
+            if (Application.getRootShell().run(output, String.format("wg show '%s' transfer", tunnel.getName())) != 0)
+                return stats;
+        } catch (final Exception ignored) {
+            return stats;
+        }
+        for (final String line : output) {
+            final String[] parts = line.split("\\t");
+            if (parts.length != 3)
+                continue;
+            try {
+                stats.add(Key.fromBase64(parts[0]), Long.parseLong(parts[1]), Long.parseLong(parts[2]));
+            } catch (final Exception ignored) {
+            }
+        }
+        return stats;
     }
 
     @Override
diff --git a/app/src/main/java/com/wireguard/android/fragment/TunnelDetailFragment.java b/app/src/main/java/com/wireguard/android/fragment/TunnelDetailFragment.java
index 8d2be476..f28a7b67 100644
--- a/app/src/main/java/com/wireguard/android/fragment/TunnelDetailFragment.java
+++ b/app/src/main/java/com/wireguard/android/fragment/TunnelDetailFragment.java
@@ -7,6 +7,8 @@ package com.wireguard.android.fragment;
 
 import android.os.Bundle;
 import androidx.annotation.Nullable;
+import androidx.databinding.DataBindingUtil;
+
 import android.view.LayoutInflater;
 import android.view.Menu;
 import android.view.MenuInflater;
@@ -15,7 +17,13 @@ import android.view.ViewGroup;
 
 import com.wireguard.android.R;
 import com.wireguard.android.databinding.TunnelDetailFragmentBinding;
+import com.wireguard.android.databinding.TunnelDetailPeerBinding;
 import com.wireguard.android.model.Tunnel;
+import com.wireguard.android.model.Tunnel.State;
+import com.wireguard.crypto.Key;
+
+import java.util.Timer;
+import java.util.TimerTask;
 
 /**
  * Fragment that shows details about a specific tunnel.
@@ -23,6 +31,20 @@ import com.wireguard.android.model.Tunnel;
 
 public class TunnelDetailFragment extends BaseFragment {
     @Nullable private TunnelDetailFragmentBinding binding;
+    @Nullable private Timer timer;
+    @Nullable private State lastState = State.TOGGLE;
+
+    private static class StatsTimerTask extends TimerTask {
+        final TunnelDetailFragment tdf;
+        private StatsTimerTask(final TunnelDetailFragment tdf) {
+            this.tdf = tdf;
+        }
+
+        @Override
+        public void run() {
+            tdf.updateStats();
+        }
+    }
 
     @Override
     public void onCreate(@Nullable final Bundle savedInstanceState) {
@@ -36,6 +58,22 @@ public class TunnelDetailFragment extends BaseFragment {
     }
 
     @Override
+    public void onStop() {
+        super.onStop();
+        if (timer != null) {
+            timer.cancel();
+            timer = null;
+        }
+    }
+
+    @Override
+    public void onResume() {
+        super.onResume();
+        timer = new Timer();
+        timer.scheduleAtFixedRate(new StatsTimerTask(this), 0, 1000);
+    }
+
+    @Override
     public View onCreateView(final LayoutInflater inflater, @Nullable final ViewGroup container,
                              @Nullable final Bundle savedInstanceState) {
         super.onCreateView(inflater, container, savedInstanceState);
@@ -59,6 +97,8 @@ public class TunnelDetailFragment extends BaseFragment {
             binding.setConfig(null);
         else
             newTunnel.getConfigAsync().thenAccept(binding::setConfig);
+        lastState = State.TOGGLE;
+        updateStats();
     }
 
     @Override
@@ -72,4 +112,52 @@ public class TunnelDetailFragment extends BaseFragment {
         super.onViewStateRestored(savedInstanceState);
     }
 
+    private String formatBytes(final long bytes) {
+        if (bytes < 1024)
+            return getContext().getString(R.string.transfer_bytes, bytes);
+        else if (bytes < 1024*1024)
+            return getContext().getString(R.string.transfer_kibibytes, bytes/1024.0);
+        else if (bytes < 1024*1024*1024)
+            return getContext().getString(R.string.transfer_mibibytes, bytes/(1024.0*1024.0));
+        else if (bytes < 1024*1024*1024*1024)
+            return getContext().getString(R.string.transfer_gibibytes, bytes/(1024.0*1024.0*1024.0));
+        return getContext().getString(R.string.transfer_tibibytes, bytes/(1024.0*1024.0*1024.0)/1024.0);
+    }
+
+    private void updateStats() {
+        if (binding == null || !isResumed())
+            return;
+        final State state = binding.getTunnel().getState();
+        if (state != State.UP && lastState == state)
+            return;
+        lastState = state;
+        binding.getTunnel().getStatisticsAsync().whenComplete((statistics, throwable) -> {
+            if (throwable != null) {
+                for (int i = 0; i < binding.peersLayout.getChildCount(); ++i) {
+                    final TunnelDetailPeerBinding peer = DataBindingUtil.getBinding(binding.peersLayout.getChildAt(i));
+                    if (peer == null)
+                        continue;
+                    peer.transferLabel.setVisibility(View.GONE);
+                    peer.transferText.setVisibility(View.GONE);
+                }
+                return;
+            }
+            for (int i = 0; i < binding.peersLayout.getChildCount(); ++i) {
+                final TunnelDetailPeerBinding peer = DataBindingUtil.getBinding(binding.peersLayout.getChildAt(i));
+                if (peer == null)
+                    continue;
+                final Key publicKey = peer.getItem().getPublicKey();
+                final long rx = statistics.peerRx(publicKey);
+                final long tx = statistics.peerTx(publicKey);
+                if (rx == 0 && tx == 0) {
+                    peer.transferLabel.setVisibility(View.GONE);
+                    peer.transferText.setVisibility(View.GONE);
+                    continue;
+                }
+                peer.transferText.setText(getContext().getString(R.string.transfer_rx_tx, formatBytes(rx), formatBytes(tx)));
+                peer.transferLabel.setVisibility(View.VISIBLE);
+                peer.transferText.setVisibility(View.VISIBLE);
+            }
+        });
+    }
 }
diff --git a/app/src/main/java/com/wireguard/android/model/Tunnel.java b/app/src/main/java/com/wireguard/android/model/Tunnel.java
index 49e78a22..87b607d0 100644
--- a/app/src/main/java/com/wireguard/android/model/Tunnel.java
+++ b/app/src/main/java/com/wireguard/android/model/Tunnel.java
@@ -5,6 +5,9 @@
 
 package com.wireguard.android.model;
 
+import android.os.SystemClock;
+import android.util.Pair;
+
 import androidx.databinding.BaseObservable;
 import androidx.databinding.Bindable;
 import androidx.annotation.Nullable;
@@ -12,8 +15,11 @@ import androidx.annotation.Nullable;
 import com.wireguard.android.BR;
 import com.wireguard.android.util.ExceptionLoggers;
 import com.wireguard.config.Config;
+import com.wireguard.crypto.Key;
 import com.wireguard.util.Keyed;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.regex.Pattern;
 
 import java9.util.concurrent.CompletableFuture;
@@ -85,15 +91,13 @@ public class Tunnel extends BaseObservable implements Keyed<String> {
     @Bindable
     @Nullable
     public Statistics getStatistics() {
-        // FIXME: Check age of statistics.
-        if (statistics == null)
+        if (statistics == null || statistics.isStale())
             TunnelManager.getTunnelStatistics(this).whenComplete(ExceptionLoggers.E);
         return statistics;
     }
 
     public CompletionStage<Statistics> getStatisticsAsync() {
-        // FIXME: Check age of statistics.
-        if (statistics == null)
+        if (statistics == null || statistics.isStale())
             return TunnelManager.getTunnelStatistics(this);
         return CompletableFuture.completedFuture(statistics);
     }
@@ -154,5 +158,48 @@ public class Tunnel extends BaseObservable implements Keyed<String> {
     }
 
     public static class Statistics extends BaseObservable {
+        private long lastTouched = SystemClock.elapsedRealtime();
+        private final Map<Key, Pair<Long, Long>> peerBytes = new HashMap<>();
+
+        public void add(final Key key, final long rx, final long tx) {
+            peerBytes.put(key, Pair.create(rx, tx));
+            lastTouched = SystemClock.elapsedRealtime();
+        }
+
+        private boolean isStale() {
+            return SystemClock.elapsedRealtime() - lastTouched > 900;
+        }
+
+        public Key[] peers() {
+            return peerBytes.keySet().toArray(new Key[0]);
+        }
+
+        public long peerRx(final Key peer) {
+            if (!peerBytes.containsKey(peer))
+                return 0;
+            return peerBytes.get(peer).first;
+        }
+
+        public long peerTx(final Key peer) {
+            if (!peerBytes.containsKey(peer))
+                return 0;
+            return peerBytes.get(peer).second;
+        }
+
+        public long totalRx() {
+            long rx = 0;
+            for (final Pair<Long, Long> val : peerBytes.values()) {
+                rx += val.first;
+            }
+            return rx;
+        }
+
+        public long totalTx() {
+            long tx = 0;
+            for (final Pair<Long, Long> val : peerBytes.values()) {
+                tx += val.second;
+            }
+            return tx;
+        }
     }
 }
diff --git a/app/src/main/java/com/wireguard/crypto/Key.java b/app/src/main/java/com/wireguard/crypto/Key.java
index f743ddd2..6648a5f3 100644
--- a/app/src/main/java/com/wireguard/crypto/Key.java
+++ b/app/src/main/java/com/wireguard/crypto/Key.java
@@ -7,6 +7,7 @@ package com.wireguard.crypto;
 
 import com.wireguard.crypto.KeyFormatException.Type;
 
+import java.security.MessageDigest;
 import java.security.SecureRandom;
 import java.util.Arrays;
 
@@ -247,6 +248,24 @@ public final class Key {
         return new String(output);
     }
 
+    @Override
+    public int hashCode() {
+        int ret = 0;
+        for (int i = 0; i < key.length / 4; ++i)
+            ret ^= (key[i * 4 + 0] >> 0) + (key[i * 4 + 1] >> 8) + (key[i * 4 + 2] >> 16) + (key[i * 4 + 3] >> 24);
+        return ret;
+    }
+
+    @Override
+    public boolean equals(final Object obj) {
+        if (obj == this)
+            return true;
+        if (obj == null || obj.getClass() != getClass())
+            return false;
+        final Key other = (Key) obj;
+        return MessageDigest.isEqual(key, other.key);
+    }
+
     /**
      * The supported formats for encoding a WireGuard key.
      */