DnsResolver.java

DnsResolver.java

import java.util.*;
import java.util.stream.*;

import javax.naming.Context;
import javax.naming.NameNotFoundException;
import javax.naming.NamingEnumeration;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;

public class DnsResolver {

  public static void main(String... args) {

    if (args.length == 0) {
      System.out.println("java Dig.java DOMAIN [RECORD_TYPE] [DNS_SERVER]");
    }

    var domain = args[0];
    var recordType = args.length == 1 ? "A" : args[1];
    var dnsServer = args.length == 3 ? new String[] { args[2] } : new String[0];

    var dnsResolver = new DnsResolver(dnsServer);
    var result = dnsResolver.resolve(domain, recordType);

    result.forEach(v -> System.out.println(v));
  }

  private final Set<String> dnsServers;
  private final String formattedDnsServers;
  private InitialDirContext context;

  public DnsResolver(String... dnsServers) {
    this.dnsServers = Set.of(dnsServers);
    if (!this.dnsServers.isEmpty()) {
      // NOTE:
      // This is done as specified in "DNS Service Provider for the Java 
      // Naming Directory Interface" [1].
      // 
      // If multiple servers are used, they should be in the following format:
      // "dns://host:port dns://host2:port2 .."
      //
      // [1]: https://docs.oracle.com/javase/8/docs/technotes/guides/jndi/jndi-dns.html 
      this.formattedDnsServers = this.dnsServers.stream()
        .map(server -> "dns://" + server).collect(Collectors.joining(" "));
    } else {
      this.formattedDnsServers = "dns:";
    }

    this.context = getNewContext();
  }

  public List<String> resolve(String domain, String recordType) {
    try {
      Attributes dnsResult = context.getAttributes(domain, new String[] { recordType });
      Attribute recordTypeResult = dnsResult.get(recordType);

      if (recordTypeResult == null) {
        return List.of();
      }

      List<String> result = new ArrayList<>();
      for (int index = 0; index < recordTypeResult.size(); index++) {
        result.add((String) recordTypeResult.get(index));
      }
      return List.copyOf(result);
    } catch (NamingException resolvingFail) {
      throw new DnsLookupException("DNS Lookup failed for" + domain, resolvingFail);
    }
  }

  private InitialDirContext getNewContext() {
    Properties env = new Properties();
    env.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.dns.DnsContextFactory");
    env.put(Context.PROVIDER_URL, formattedDnsServers);
    env.put("com.sun.jndi.dns.timeout.retries", "1");
    try {
      return new InitialDirContext(env);
    }
    catch (NamingException ex) {
      throw new DnsLookupException("Cannot create InitialDirContext for DNS lookup", ex);
    }
  }

  public static class DnsLookupException extends RuntimeException {
    public DnsLookupException(String message, Throwable cause) {
      super(message, cause);
    }
  }
}