Fix dual-stack IP allocation

4 files changed, 226 insertions, 55 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
index c714d1d2231..074a20fc82d 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
@@ -13,9 +13,11 @@ import java.net.Inet6Address;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Represents IP addresses owned by a node.
@@ -61,23 +63,26 @@ public class IP {
         }
 
         /**
-         * Find a free allocation in this pool
+         * Find a free allocation in this pool. Note that the allocation is not final until it is assigned to a node
          *
          * @param nodes All nodes in the repository
          * @return An allocation from the pool, if any can be made
          */
-        public Optional<Allocation> findAllocation(NodeList nodes) {
+        public Optional<Allocation> findAllocation(NodeList nodes, NameResolver resolver) {
             Set<String> unusedAddresses = findUnused(nodes);
-            Optional<String> ipv6Address = unusedAddresses.stream()
-                                                          .filter(addr -> InetAddresses.forString(addr) instanceof Inet6Address)
-                                                          .findFirst();
-            Optional<String> ipv4Address = unusedAddresses.stream()
-                                                          .filter(addr -> InetAddresses.forString(addr) instanceof Inet4Address)
-                                                          .findFirst();
-
-            // An allocation must contain one IPv6 address, but IPv4 is optional. All hosts have IPv6 addresses that
-            // can be used by containers, while the availability of IPv4 addresses depends on the cloud provider.
-            return ipv6Address.map(address -> new Allocation(address, ipv4Address));
+            Optional<Allocation> allocation = unusedAddresses.stream()
+                                                             .filter(IP::isV6)
+                                                             .findFirst()
+                                                             .map(addr -> Allocation.resolveFrom(addr, resolver));
+            allocation.flatMap(Allocation::ipv4Address).ifPresent(ipv4Address -> {
+               if (!unusedAddresses.contains(ipv4Address)) {
+                   throw new IllegalArgumentException("Allocation resolved " + ipv4Address + " from hostname " +
+                                                      allocation.get().hostname +
+                                                      ", but that address is not available in the address pool of " +
+                                                      owner.hostname());
+               }
+            });
+            return allocation;
         }
 
         /** Find all unused addresses in this pool
@@ -110,16 +115,12 @@ public class IP {
         private static Set<String> requireAddresses(Set<String> addresses) {
             Objects.requireNonNull(addresses, "ipAddressPool must be non-null");
 
-            long ipv6AddrCount = addresses.stream()
-                                          .filter(addr -> InetAddresses.forString(addr) instanceof Inet6Address)
-                                          .count();
+            long ipv6AddrCount = addresses.stream().filter(IP::isV6).count();
             if (ipv6AddrCount == addresses.size()) {
                 return addresses; // IPv6-only pool is valid
             }
 
-            long ipv4AddrCount = addresses.stream()
-                                          .filter(addr -> InetAddresses.forString(addr) instanceof Inet4Address)
-                                          .count();
+            long ipv4AddrCount = addresses.stream().filter(IP::isV4).count();
             if (ipv4AddrCount == ipv6AddrCount) {
                 return addresses;
             }
@@ -135,32 +136,66 @@ public class IP {
     /** An IP address allocation from a pool */
     public static class Allocation {
 
+        private final String hostname;
         private final String ipv6Address;
         private final Optional<String> ipv4Address;
 
-        private Allocation(String ipv6Address, Optional<String> ipv4Address) {
-            this.ipv6Address = Objects.requireNonNull(ipv6Address, "ipv6Address must be non-null");
-            this.ipv4Address = Objects.requireNonNull(ipv4Address, "ipv4Address must be non-null");
+        private Allocation(String hostname, String ipv6Address, Optional<String> ipv4Address) {
+            Objects.requireNonNull(ipv6Address, "ipv6Address must be non-null");
+            if (!isV6(ipv6Address)) {
+                throw new IllegalArgumentException("Invalid IPv6 address '" + ipv6Address + "'");
+            }
+
+            Objects.requireNonNull(ipv4Address, "ipv4Address must be non-null");
+            if (ipv4Address.isPresent() && !isV4(ipv4Address.get())) {
+                throw new IllegalArgumentException("Invalid IPv4 address '" + ipv4Address + "'");
+            }
+            this.hostname = Objects.requireNonNull(hostname, "hostname must be non-null");
+            this.ipv6Address = ipv6Address;
+            this.ipv4Address = ipv4Address;
         }
 
         /**
-         * Resolve and return the hostname of this allocation
+         * Resolve the IP addresses and hostname of this allocation
          *
-         * @param resolver DNS name resolver to use when resolving the hostname
+         * @param ipv6Address Unassigned IPv6 address
+         * @param resolver DNS name resolver to use
          * @throws IllegalArgumentException if DNS is misconfigured
-         * @return The hostname
+         * @return The allocation containing 1 IPv6 address and 1 IPv4 address (if hostname is dual-stack)
          */
-        public String resolveHostname(NameResolver resolver) {
+        public static Allocation resolveFrom(String ipv6Address, NameResolver resolver) {
             String hostname6 = resolver.getHostname(ipv6Address).orElseThrow(() -> new IllegalArgumentException("Could not resolve IP address: " + ipv6Address));
-            if (ipv4Address.isPresent()) {
-                String hostname4 = resolver.getHostname(ipv4Address.get()).orElseThrow(() -> new IllegalArgumentException("Could not resolve IP address: " + ipv4Address));
+            List<String> ipv4Addresses = resolver.getAllByNameOrThrow(hostname6).stream()
+                                                 .filter(IP::isV4)
+                                                 .collect(Collectors.toList());
+            if (ipv4Addresses.size() > 1) {
+                throw new IllegalArgumentException("Hostname " + hostname6 + " resolved to more than 1 IPv4 address: " + ipv4Addresses);
+            }
+            Optional<String> ipv4Address = ipv4Addresses.stream().findFirst();
+            ipv4Address.ifPresent(addr -> {
+                String hostname4 = resolver.getHostname(addr).orElseThrow(() -> new IllegalArgumentException("Could not resolve IP address: " + ipv4Address));
                 if (!hostname6.equals(hostname4)) {
                     throw new IllegalArgumentException(String.format("Hostnames resolved from each IP address do not " +
                                                                      "point to the same hostname [%s -> %s, %s -> %s]",
-                                                                     ipv6Address, hostname6, ipv4Address.get(), hostname4));
+                                                                     ipv6Address, hostname6, addr, hostname4));
                 }
-            }
-            return hostname6;
+            });
+            return new Allocation(hostname6, ipv6Address, ipv4Address);
+        }
+
+        /** Hostname pointing to the IP addresses in this */
+        public String hostname() {
+            return hostname;
+        }
+
+        /** IPv6 address in this allocation */
+        public String ipv6Address() {
+            return ipv6Address;
+        }
+
+        /** IPv4 address in this allocation */
+        public Optional<String> ipv4Address() {
+            return ipv4Address;
         }
 
         /** All IP addresses in this */
@@ -179,4 +214,14 @@ public class IP {
 
     }
 
+    /** Returns whether given string is an IPv4 address */
+    public static boolean isV4(String ipAddress) {
+        return InetAddresses.forString(ipAddress) instanceof Inet4Address;
+    }
+
+    /** Returns whether given string is an IPv6 address */
+    public static boolean isV6(String ipAddress) {
+        return InetAddresses.forString(ipAddress) instanceof Inet6Address;
+    }
+
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
index 892dcdcb36c..6624a3c8b59 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
@@ -134,21 +134,19 @@ public class NodePrioritizer {
 
             log.log(LogLevel.DEBUG, "Trying to add new Docker node on " + node);
 
-            Optional<IP.Allocation> allocation = node.ipAddressPool().findAllocation(list);
-            if (!allocation.isPresent()) continue; // No free addresses in this pool
-
-            String hostname;
+            Optional<IP.Allocation> allocation;
             try {
-                hostname = allocation.get().resolveHostname(nameResolver);
-            } catch (IllegalArgumentException e) {
-                log.log(LogLevel.WARNING, "Failed to resolve hostname for allocation: " + allocation.get() + ", skipping", e);
+                allocation = node.ipAddressPool().findAllocation(list, nameResolver);
+                if (!allocation.isPresent()) continue; // No free addresses in this pool
+            } catch (Exception e) {
+                log.log(LogLevel.WARNING, "Failed to resolve hostname for allocation, skipping", e);
                 continue;
             }
 
-            Node newNode = Node.createDockerNode("fake-" + hostname,
+            Node newNode = Node.createDockerNode("fake-" + allocation.get().hostname(),
                                                  allocation.get().addresses(),
                                                  Collections.emptySet(),
-                                                 hostname,
+                                                 allocation.get().hostname(),
                                                  Optional.of(node.hostname()),
                                                  getFlavor(requestedNodes),
                                                  NodeType.tenant);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNameResolver.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNameResolver.java
index e1b1fc85f0a..325dabe08f1 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNameResolver.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNameResolver.java
@@ -22,10 +22,16 @@ import java.util.concurrent.ThreadLocalRandom;
  */
 public class MockNameResolver implements NameResolver {
 
-    private boolean allowInvocation = true;
-    private boolean mockAnyLookup = false;
     private final Map<String, Set<String>> records = new HashMap<>();
 
+    private boolean mockAnyLookup = false;
+    private boolean explicitReverseRecords = false;
+
+    public MockNameResolver addReverseRecord(String ipAddress, String hostname) {
+        addRecord(ipAddress, hostname);
+        return this;
+    }
+
     public MockNameResolver addRecord(String hostname, String... ipAddress) {
         Objects.requireNonNull(hostname, "hostname must be non-null");
         Arrays.stream(ipAddress).forEach(ip -> Objects.requireNonNull(ip, "ipAddress must be non-null"));
@@ -39,28 +45,23 @@ public class MockNameResolver implements NameResolver {
         return this;
     }
 
-    public MockNameResolver reset() {
-        this.allowInvocation = true;
-        this.mockAnyLookup = false;
-        this.records.clear();
+    public MockNameResolver removeRecord(String name) {
+        records.remove(name);
         return this;
     }
 
-    public MockNameResolver failIfInvoked() {
-        this.allowInvocation = false;
+    public MockNameResolver mockAnyLookup() {
+        this.mockAnyLookup = true;
         return this;
     }
 
-    public MockNameResolver mockAnyLookup() {
-        this.mockAnyLookup = true;
+    public MockNameResolver explicitReverseRecords() {
+        this.explicitReverseRecords = true;
         return this;
     }
 
     @Override
     public Set<String> getAllByNameOrThrow(String hostname) {
-        if (!allowInvocation) {
-            throw new IllegalStateException("Expected getAllByName to not be invoked for hostname: " + hostname);
-        }
         if (records.containsKey(hostname)) {
             return records.get(hostname);
         }
@@ -74,14 +75,18 @@ public class MockNameResolver implements NameResolver {
 
     @Override
     public Optional<String> getHostname(String ipAddress) {
-        for (String host : records.keySet()) {
-            if (records.get(host).contains(ipAddress)) return Optional.of(host);
+        if (!explicitReverseRecords) {
+            return records.entrySet().stream()
+                          .filter(kv -> kv.getValue().contains(ipAddress))
+                          .map(Map.Entry::getKey)
+                          .findFirst();
         }
         if (mockAnyLookup) {
             String hostname = UUID.randomUUID().toString();
             records.put(hostname, Collections.singleton(ipAddress));
         }
-        return Optional.empty();
+        return Optional.ofNullable(records.get(ipAddress))
+                       .flatMap(values -> values.stream().findFirst());
     }
 
     private static String randomIpAddress() {
@@ -92,4 +97,5 @@ public class MockNameResolver implements NameResolver {
                 random.nextInt(1, 255 + 1),
                 random.nextInt(1, 255 + 1));
     }
+
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/node/IPTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/node/IPTest.java
index fc28e1fdc8f..0b2a866cc99 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/node/IPTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/node/IPTest.java
@@ -3,19 +3,40 @@ package com.yahoo.vespa.hosted.provision.node;
 
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
+import com.yahoo.config.provision.NodeFlavors;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.provisioning.FlavorConfigBuilder;
+import com.yahoo.vespa.hosted.provision.testutils.MockNameResolver;
+import org.junit.Assert;
+import org.junit.Before;
 import org.junit.Test;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.Optional;
 import java.util.Set;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 /**
  * @author mpolden
  */
 public class IPTest {
 
+    private static final NodeFlavors nodeFlavors = FlavorConfigBuilder.createDummies("default");
+    private static final NodeList emptyList = new NodeList(Collections.emptyList());
+
+    private MockNameResolver resolver;
+
+    @Before
+    public void before() {
+        resolver = new MockNameResolver().explicitReverseRecords();
+    }
+
     @Test
     public void test_natural_order() {
         Set<String> ipAddresses = ImmutableSet.of(
@@ -50,4 +71,105 @@ public class IPTest {
         );
     }
 
+    @Test
+    public void test_find_allocation_single_stack() {
+        IP.AddressPool pool = createNode(ImmutableSet.of(
+                "::1",
+                "::2",
+                "::3"
+        )).ipAddressPool();
+
+        resolver.addRecord("host1", "::2");
+        resolver.addRecord("host2", "::3");
+        resolver.addRecord("host3", "::1");
+        resolver.addReverseRecord("::3", "host2");
+        resolver.addReverseRecord("::1", "host3");
+        resolver.addReverseRecord("::2", "host1");
+
+        Optional<IP.Allocation> allocation = pool.findAllocation(emptyList, resolver);
+        assertEquals("::1", allocation.get().ipv6Address());
+        Assert.assertFalse(allocation.get().ipv4Address().isPresent());
+        assertEquals("host3", allocation.get().hostname());
+
+        // Allocation fails if DNS record is missing
+        resolver.removeRecord("host3");
+        try {
+            pool.findAllocation(emptyList, resolver);
+            fail("Expected exception");
+        } catch (Exception e) {
+            assertEquals("java.net.UnknownHostException: Could not resolve: host3", e.getMessage());
+        }
+    }
+
+    @Test
+    public void test_find_allocation_dual_stack() {
+        IP.AddressPool pool = dualStackPool();
+        Optional<IP.Allocation> allocation = pool.findAllocation(emptyList, resolver);
+        assertEquals("::1", allocation.get().ipv6Address());
+        assertEquals("127.0.0.2", allocation.get().ipv4Address().get());
+        assertEquals("host3", allocation.get().hostname());
+    }
+
+    @Test
+    public void test_find_allocation_multiple_ipv4_addresses() {
+        IP.AddressPool pool = dualStackPool();
+        resolver.addRecord("host3", "127.0.0.127");
+        try {
+            pool.findAllocation(emptyList, resolver);
+            fail("Expected exception");
+        } catch (IllegalArgumentException e) {
+            assertEquals("Hostname host3 resolved to more than 1 IPv4 address: [127.0.0.2, 127.0.0.127]",
+                         e.getMessage());
+        }
+    }
+
+    @Test
+    public void test_find_allocation_invalid_ipv4_reverse_record() {
+        IP.AddressPool pool = dualStackPool();
+        resolver.removeRecord("127.0.0.2")
+                .addReverseRecord("127.0.0.2", "host5");
+        try {
+            pool.findAllocation(emptyList, resolver);
+            //fail("Expected exception");
+        } catch (IllegalArgumentException e) {
+            assertEquals("Hostnames resolved from each IP address do not point to the same hostname " +
+                         "[::1 -> host3, 127.0.0.2 -> host5]", e.getMessage());
+        }
+    }
+
+    private IP.AddressPool dualStackPool() {
+        Node node = createNode(ImmutableSet.of(
+                "127.0.0.1",
+                "127.0.0.2",
+                "127.0.0.3",
+                "::1",
+                "::2",
+                "::3"
+        ));
+
+        // IPv4 addresses
+        resolver.addRecord("host1", "127.0.0.3")
+                .addRecord("host2", "127.0.0.1")
+                .addRecord("host3", "127.0.0.2")
+                .addReverseRecord("127.0.0.1", "host2")
+                .addReverseRecord("127.0.0.2", "host3")
+                .addReverseRecord("127.0.0.3", "host1");
+
+        // IPv6 addresses
+        resolver.addRecord("host1", "::2")
+                .addRecord("host2", "::3")
+                .addRecord("host3", "::1")
+                .addReverseRecord("::3", "host2")
+                .addReverseRecord("::1", "host3")
+                .addReverseRecord("::2", "host1");
+
+        return node.ipAddressPool();
+    }
+
+    private static Node createNode(Set<String> ipAddresses) {
+        return Node.create("id1", Collections.singleton("127.0.0.1"), ipAddresses,
+                           "host1", Optional.empty(), nodeFlavors.getFlavorOrThrow("default"),
+                           NodeType.host);
+    }
+
 }