1 files changed, 81 insertions, 30 deletions

diff --git a/app/src/main/java/com/wireguard/config/InetEndpoint.java b/app/src/main/java/com/wireguard/config/InetEndpoint.java
index 3efe4203..06d0ca80 100644
--- a/app/src/main/java/com/wireguard/config/InetEndpoint.java
+++ b/app/src/main/java/com/wireguard/config/InetEndpoint.java
@@ -5,36 +5,68 @@
 
 package com.wireguard.config;
 
-import android.annotation.SuppressLint;
+import android.support.annotation.Nullable;
 
-import com.wireguard.android.Application;
-import com.wireguard.android.R;
+import org.threeten.bp.Duration;
+import org.threeten.bp.Instant;
 
 import java.net.Inet4Address;
-import java.net.Inet6Address;
 import java.net.InetAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.UnknownHostException;
+import java.util.regex.Pattern;
 
-import javax.annotation.Nullable;
+import java9.util.Optional;
+
+
+/**
+ * An external endpoint (host and port) used to connect to a WireGuard {@link Peer}.
+ * <p>
+ * Instances of this class are externally immutable.
+ */
+public final class InetEndpoint {
+    private static final Pattern BARE_IPV6 = Pattern.compile("^[^\\[]*:");
+    private static final Pattern FORBIDDEN_CHARACTERS = Pattern.compile("[/?#]");
 
-public class InetEndpoint {
     private final String host;
+    private final boolean isResolved;
+    private final Object lock = new Object();
     private final int port;
-    @Nullable private InetAddress resolvedHost;
+    private Instant lastResolution = Instant.EPOCH;
+    @Nullable private InetEndpoint resolved;
 
-    public InetEndpoint(@Nullable final String endpoint) {
-        if (endpoint.indexOf('/') != -1 || endpoint.indexOf('?') != -1 || endpoint.indexOf('#') != -1)
-            throw new IllegalArgumentException(Application.get().getString(R.string.tunnel_error_forbidden_endpoint_chars));
+    private InetEndpoint(final String host, final boolean isResolved, final int port) {
+        this.host = host;
+        this.isResolved = isResolved;
+        this.port = port;
+    }
+
+    public static InetEndpoint parse(final String endpoint) {
+        if (FORBIDDEN_CHARACTERS.matcher(endpoint).find())
+            throw new IllegalArgumentException("Forbidden characters in Endpoint");
         final URI uri;
         try {
             uri = new URI("wg://" + endpoint);
         } catch (final URISyntaxException e) {
             throw new IllegalArgumentException(e);
         }
-        host = uri.getHost();
-        port = uri.getPort();
+        try {
+            InetAddresses.parse(uri.getHost());
+            // Parsing ths host as a numeric address worked, so we don't need to do DNS lookups.
+            return new InetEndpoint(uri.getHost(), true, uri.getPort());
+        } catch (final IllegalArgumentException ignored) {
+            // Failed to parse the host as a numeric address, so it must be a DNS hostname/FQDN.
+            return new InetEndpoint(uri.getHost(), false, uri.getPort());
+        }
+    }
+
+    @Override
+    public boolean equals(final Object obj) {
+        if (!(obj instanceof InetEndpoint))
+            return false;
+        final InetEndpoint other = (InetEndpoint) obj;
+        return host.equals(other.host) && port == other.port;
     }
 
     public String getHost() {
@@ -45,28 +77,47 @@ public class InetEndpoint {
         return port;
     }
 
-    @SuppressLint("DefaultLocale")
-    public String getResolvedEndpoint() throws UnknownHostException {
-        if (resolvedHost == null) {
-            final InetAddress[] candidates = InetAddress.getAllByName(host);
-            if (candidates.length == 0)
-                throw new UnknownHostException(host);
-            for (final InetAddress addr : candidates) {
-                if (addr instanceof Inet4Address) {
-                    resolvedHost = addr;
-                    break;
+    /**
+     * Generate an {@code InetEndpoint} instance with the same port and the host resolved using DNS
+     * to a numeric address. If the host is already numeric, the existing instance may be returned.
+     * Because this function may perform network I/O, it must not be called from the main thread.
+     *
+     * @return the resolved endpoint, or {@link Optional#empty()}
+     */
+    public Optional<InetEndpoint> getResolved() {
+        if (isResolved)
+            return Optional.of(this);
+        synchronized (lock) {
+            //TODO(zx2c4): Implement a real timeout mechanism using DNS TTL
+            if (Duration.between(lastResolution, Instant.now()).toMinutes() > 1) {
+                try {
+                    // Prefer v4 endpoints over v6 to work around DNS64 and IPv6 NAT issues.
+                    final InetAddress[] candidates = InetAddress.getAllByName(host);
+                    InetAddress address = candidates[0];
+                    for (final InetAddress candidate : candidates) {
+                        if (candidate instanceof Inet4Address) {
+                            address = candidate;
+                            break;
+                        }
+                    }
+                    resolved = new InetEndpoint(address.getHostAddress(), true, port);
+                    lastResolution = Instant.now();
+                } catch (final UnknownHostException e) {
+                    resolved = null;
                 }
             }
-            if (resolvedHost == null)
-                resolvedHost = candidates[0];
+            return Optional.ofNullable(resolved);
         }
-        return String.format(resolvedHost instanceof Inet6Address ?
-                "[%s]:%d" : "%s:%d", resolvedHost.getHostAddress(), port);
     }
 
-    @SuppressLint("DefaultLocale")
-    public String getEndpoint() {
-        return String.format(host.contains(":") && !host.contains("[") ?
-                "[%s]:%d" : "%s:%d", host, port);
+    @Override
+    public int hashCode() {
+        return host.hashCode() ^ port;
+    }
+
+    @Override
+    public String toString() {
+        final boolean isBareIpv6 = isResolved && BARE_IPV6.matcher(host).matches();
+        return (isBareIpv6 ? '[' + host + ']' : host) + ':' + port;
     }
 }