// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.vespa.hosted.provision.persistence;

import com.google.common.collect.ImmutableSet;

import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import java.util.Optional;
import java.util.Set;

/**
 * Implementation of a name resolver that always use a DNS server to resolve the given name. The intention is to avoid
 * possibly incorrect/incomplete records in /etc/hosts.
 *
 * @author mpolden
 */
public class DnsNameResolver implements NameResolver {

    private static final String TYPE_A = "A";
    private static final String TYPE_AAAA = "AAAA";

    /** Resolve IP addresses for given host name */
    @Override
    public Set<String> getAllByNameOrThrow(String hostname) {
        try {
            DirContext ctx = new InitialDirContext();
            Attributes attributes = ctx.getAttributes("dns:/" + hostname, new String[] {TYPE_A, TYPE_AAAA});
            Optional<String> inet4Address = getStringAttribute(attributes, TYPE_A);
            Optional<String> inet6Address = getStringAttribute(attributes, TYPE_AAAA);

            ImmutableSet.Builder<String> ipAddresses = ImmutableSet.builder();
            inet4Address.ifPresent(ipAddresses::add);
            inet6Address.ifPresent(ipAddresses::add);
            return ipAddresses.build();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static Optional<String> getStringAttribute(Attributes attributes, String key) throws NamingException {
        Optional<Attribute> attribute = Optional.ofNullable(attributes.get(key));
        if (attribute.isPresent()) {
            return Optional.ofNullable(attribute.get().get()).map(Object::toString);
        }
        return Optional.empty();
    }
}

@@ -527,16 +527,17 @@
 ways use
+s
  a DNS s
@@ -737,107 +737,8 @@
  %7B%0A%0A
-    private static final String TYPE_A = %22A%22;%0A    private static final String TYPE_AAAA = %22AAAA%22;%0A%0A
     
@@ -891,158 +891,157 @@
     
-DirContext ctx = new InitialDirContext();%0A            Attributes attributes = ctx.getAttributes(%22dns:/%22 + hostname, new String%5B%5D %7BTYPE_A, TYPE_AAAA%7D);
+Optional%3CString%3E cname = lookupName(hostname, Type.CNAME);%0A            if (cname.isPresent()) %7B%0A                hostname = cname.get();%0A            %7D
 %0A   
@@ -1085,44 +1085,34 @@
 s = 
-getStringAttribute(attributes, TYPE_
+lookupName(hostname, Type.
 A);%0A
@@ -1159,44 +1159,34 @@
 s = 
-getStringAttribute(attributes, TYPE_
+lookupName(hostname, Type.
 AAAA
@@ -1524,15 +1524,8 @@
 ate 
-static 
 Opti
@@ -1541,86 +1541,210 @@
 ng%3E 
-getStringAttribute(Attributes attributes, String key) throws NamingException %7B
+lookupName(String name, Type type) throws NamingException %7B%0A        DirContext ctx = new InitialDirContext();%0A        Attributes attributes = ctx.getAttributes(%22dns:/%22 + name, new String%5B%5D%7Btype.value%7D);
 %0A   
@@ -1819,11 +1819,18 @@
 get(
-key
+type.value
 ));%0A
@@ -1999,11 +1999,214 @@
 );%0A    %7D
+%0A%0A    private enum Type %7B%0A%0A        A(%22A%22),%0A        AAAA(%22AAAA%22),%0A        CNAME(%22CNAME%22);%0A%0A        private final String value;%0A%0A        Type(String value) %7B%0A            this.value = value;%0A        %7D%0A    %7D
 %0A%7D%0A

Support CNAME records