2 files changed, 79 insertions, 33 deletions

diff --git a/app/src/main/java/com/wireguard/config/InetEndpoint.java b/app/src/main/java/com/wireguard/config/InetEndpoint.java
new file mode 100644
index 00000000..4896ee97
--- /dev/null
+++ b/app/src/main/java/com/wireguard/config/InetEndpoint.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright © 2017-2018 WireGuard LLC. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.wireguard.config;
+
+import android.annotation.SuppressLint;
+
+import com.wireguard.android.Application;
+import com.wireguard.android.R;
+
+import java.net.Inet4Address;
+import java.net.InetAddress;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.net.UnknownHostException;
+
+import javax.annotation.Nullable;
+
+public class InetEndpoint {
+    private final String host;
+    private final int port;
+    @Nullable private InetAddress resolvedHost;
+
+    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));
+        final URI uri;
+        try {
+            uri = new URI("wg://" + endpoint);
+        } catch (final URISyntaxException e) {
+            throw new IllegalArgumentException(e);
+        }
+        host = uri.getHost();
+        port = uri.getPort();
+    }
+
+    public String getHost() {
+        return host;
+    }
+
+    public int getPort() {
+        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;
+                }
+            }
+            if (resolvedHost == null)
+                resolvedHost = candidates[0];
+        }
+        return String.format(resolvedHost instanceof Inet4Address ?
+                "[%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);
+    }
+}
diff --git a/app/src/main/java/com/wireguard/config/Peer.java b/app/src/main/java/com/wireguard/config/Peer.java
index 803bf9a7..914516ba 100644
--- a/app/src/main/java/com/wireguard/config/Peer.java
+++ b/app/src/main/java/com/wireguard/config/Peer.java
@@ -37,7 +37,7 @@ import java9.lang.Iterables;
 
 public class Peer {
     private final List<InetNetwork> allowedIPsList;
-    @Nullable private InetSocketAddress endpoint;
+    @Nullable private InetEndpoint endpoint;
     private int persistentKeepalive;
     @Nullable private String preSharedKey;
     @Nullable private String publicKey;
@@ -67,19 +67,15 @@ public class Peer {
     }
 
     @Nullable
-    public InetSocketAddress getEndpoint() {
+    public InetEndpoint getEndpoint() {
         return endpoint;
     }
 
-    @SuppressLint("DefaultLocale")
     @Nullable
     private String getEndpointString() {
         if (endpoint == null)
             return null;
-        if (endpoint.getHostString().contains(":") && !endpoint.getHostString().contains("["))
-            return String.format("[%s]:%d", endpoint.getHostString(), endpoint.getPort());
-        else
-            return String.format("%s:%d", endpoint.getHostString(), endpoint.getPort());
+        return endpoint.getEndpoint();
     }
 
     public int getPersistentKeepalive() {
@@ -103,21 +99,10 @@ public class Peer {
         return publicKey;
     }
 
-    @SuppressLint("DefaultLocale")
     public String getResolvedEndpointString() throws UnknownHostException {
         if (endpoint == null)
             throw new UnknownHostException("{empty}");
-        if (endpoint.isUnresolved())
-            endpoint = new InetSocketAddress(endpoint.getHostString(), endpoint.getPort());
-        if (endpoint.isUnresolved())
-            throw new UnknownHostException(endpoint.getHostString());
-        if (endpoint.getAddress() instanceof Inet6Address)
-            return String.format("[%s]:%d",
-                    endpoint.getAddress().getHostAddress(),
-                    endpoint.getPort());
-        return String.format("%s:%d",
-                endpoint.getAddress().getHostAddress(),
-                endpoint.getPort());
+        return endpoint.getResolvedEndpoint();
     }
 
     public void parse(final String line) {
@@ -150,24 +135,14 @@ public class Peer {
         addAllowedIPs(Attribute.stringToList(allowedIPsString));
     }
 
-    private void setEndpoint(@Nullable final InetSocketAddress endpoint) {
+    private void setEndpoint(@Nullable final InetEndpoint endpoint) {
         this.endpoint = endpoint;
     }
 
     private void setEndpointString(@Nullable final String endpoint) {
-        if (endpoint != null && !endpoint.isEmpty()) {
-            final InetSocketAddress constructedEndpoint;
-            if (endpoint.indexOf('/') != -1 || endpoint.indexOf('?') != -1 || endpoint.indexOf('#') != -1)
-                throw new IllegalArgumentException(context.getString(R.string.tunnel_error_forbidden_endpoint_chars));
-            final URI uri;
-            try {
-                uri = new URI("wg://" + endpoint);
-            } catch (final URISyntaxException e) {
-                throw new IllegalArgumentException(e);
-            }
-            constructedEndpoint = InetSocketAddress.createUnresolved(uri.getHost(), uri.getPort());
-            setEndpoint(constructedEndpoint);
-        } else
+        if (endpoint != null && !endpoint.isEmpty())
+            setEndpoint(new InetEndpoint(endpoint));
+        else
             setEndpoint(null);
     }