FakeDNS.java

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.kudu.client;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import javax.annotation.concurrent.GuardedBy;

import com.google.common.base.Throwables;
import com.google.common.net.InetAddresses;

/**
 * Fake DNS resolver which allows our tests to work well even though we use
 * strange loopback IP addresses (127.x.y.z) with no corresponding reverse
 * DNS.
 *
 * This overrides the reverse lookups for such IPs to return the same address
 * in String form.
 *
 * Without this class, reverse DNS lookups for such addresses often take
 * 5 seconds to return, causing timeouts and overall test slowness.
 *
 * In the future this class might also be extended to test more interesting
 * DNS-related scenarios.
 */
public class FakeDNS {
 static FakeDNS instance = new FakeDNS();

 @GuardedBy("this")
 private Map<String, InetAddress> forwardResolutions = new HashMap<>();

 @GuardedBy("this")
 private Map<InetAddress, String> reverseResolutions = new HashMap<>();

 /** whether the fake resolver has been installed */
 @GuardedBy("this")
 private boolean installed = false;

 private FakeDNS() {}
 public static FakeDNS getInstance() {
 return instance;
 }

 public synchronized void addForwardResolution(String hostname, InetAddress ip) {
 forwardResolutions.put(hostname, ip);
 }

 public synchronized void addReverseResolution(InetAddress ip, String hostname) {
 reverseResolutions.put(ip, hostname);
 }

 /**
 * Install the fake DNS resolver into the Java runtime.
 */
 public synchronized void install() {
 if (installed) return;
 try {
 try {
 // Override the NameService in Java 9 or later.
 final Class<?> nameServiceInterface = Class.forName("java.net.InetAddress$NameService");
 Field field = InetAddress.class.getDeclaredField("nameService");
 // Get the default NameService to fallback to.
 Method method = InetAddress.class.getDeclaredMethod("createNameService");
 method.setAccessible(true);
 Object fallbackNameService = method.invoke(null);
 // Create a proxy instance to set on the InetAddress field which will handle
 // all NameService calls.
 Object proxy = Proxy.newProxyInstance(nameServiceInterface.getClassLoader(),
 new Class<?>[]{nameServiceInterface}, new NameServiceListener(fallbackNameService));
 field.setAccessible(true);
 field.set(InetAddress.class, proxy);
 } catch (final ClassNotFoundException | NoSuchFieldException e) {
 // Override the NameService in Java 8 or earlier.
 final Class<?> nameServiceInterface = Class.forName("sun.net.spi.nameservice.NameService");
 Field field = InetAddress.class.getDeclaredField("nameServices");
 // Get the default NameService to fallback to.
 Method method = InetAddress.class.getDeclaredMethod("createNSProvider", String.class);
 method.setAccessible(true);
 Object fallbackNameService = method.invoke(null, "default");
 // Create a proxy instance to set on the InetAddress field which will handle
 // all NameService calls.
 Object proxy = Proxy.newProxyInstance(nameServiceInterface.getClassLoader(),
 new Class<?>[]{nameServiceInterface}, new NameServiceListener(fallbackNameService));
 field.setAccessible(true);
 // Java 8 or earlier takes a list of NameServices
 field.set(InetAddress.class, Arrays.asList(proxy));
 }
 } catch (Exception e) {
 throw new RuntimeException(e);
 }
 installed = true;
 }

 /**
 * The NameService in all versions of Java has the same interface, so we
 * can use the same InvocationHandler as our proxy instance for both
 * java.net.InetAddress$NameService and sun.net.spi.nameservice.NameService.
 */
 private class NameServiceListener implements InvocationHandler {

 private final Object fallbackNameService;

 // Creates a NameServiceListener with a NameService implementation to
 // fallback to. The parameter is untyped so we can handle the NameService
 // type in all versions of Java with reflection.
 NameServiceListener(Object fallbackNameService) {
 this.fallbackNameService = fallbackNameService;
 }

 private InetAddress[] lookupAllHostAddr(String host) throws UnknownHostException {
 InetAddress inetAddress;
 synchronized(FakeDNS.this) {
 inetAddress = forwardResolutions.get(host);
 }
 if (inetAddress != null) {
 return new InetAddress[]{inetAddress};
 }

 try {
 Method method = fallbackNameService.getClass()
 .getDeclaredMethod("lookupAllHostAddr", String.class);
 method.setAccessible(true);
 return (InetAddress[]) method.invoke(fallbackNameService, host);
 } catch (ReflectiveOperationException e) {
 Throwables.propagateIfPossible(e.getCause(), UnknownHostException.class);
 throw new AssertionError("unexpected reflection issue", e);
 }
 }

 private String getHostByAddr(byte[] addr) throws UnknownHostException {
 if (addr[0] == 127) {
 return InetAddresses.toAddrString(InetAddress.getByAddress(addr));
 }

 String hostname;
 synchronized (FakeDNS.this) {
 hostname = reverseResolutions.get(InetAddress.getByAddress(addr));
 }
 if (hostname != null) {
 return hostname;
 }

 try {
 Method method = fallbackNameService.getClass()
 .getDeclaredMethod("getHostByAddr", byte[].class);
 method.setAccessible(true);
 return (String) method.invoke(fallbackNameService, (Object) addr);
 } catch (ReflectiveOperationException e) {
 Throwables.propagateIfPossible(e.getCause(), UnknownHostException.class);
 throw new AssertionError("unexpected reflection issue", e);
 }
 }

 @Override
 public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
 switch (method.getName()) {
 case "lookupAllHostAddr":
 return lookupAllHostAddr((String) args[0]);
 case "getHostByAddr":
 return getHostByAddr((byte[]) args[0]);
 default:
 throw new UnsupportedOperationException();
 }
 }
 }
}