core: handle long dns txt records properly, parse service config, and add tests

core/src/main/java/io/grpc/internal/DnsNameResolver.java

@@ -21,6 +21,11 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Verify;
+import com.google.gson.Gson;
+import com.google.gson.JsonArray;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonSyntaxException;
 import io.grpc.Attributes;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.NameResolver;
@@ -34,7 +39,12 @@ import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Locale;
+import java.util.Map.Entry;
+import java.util.Random;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -62,14 +72,21 @@ final class DnsNameResolver extends NameResolver {
   private static final Logger logger = Logger.getLogger(DnsNameResolver.class.getName());
 
   private static final boolean JNDI_AVAILABLE = jndiAvailable();
+  private static final String LOCAL_HOST_NAME = getHostName();
 
   // From https://github.com/grpc/proposal/blob/master/A2-service-configs-in-dns.md
   private static final String SERVICE_CONFIG_NAME_PREFIX = "_grpc_config.";
   // From https://github.com/grpc/proposal/blob/master/A5-grpclb-in-dns.md
   private static final String GRPCLB_NAME_PREFIX = "_grpclb._tcp.";
+  // From https://github.com/grpc/proposal/blob/master/A2-service-configs-in-dns.md
+  private static final String SERVICE_CONFIG_PREFIX = "_grpc_config=";
+  private static final Set<String> SERVICE_CONFIG_CHOICE_KEYS =
+      Collections.unmodifiableSet(
+          new HashSet<String>(
+              Arrays.asList("clientLanguage", "percentage", "clientHostname", "serviceConfig")));
 
   private static final String JNDI_PROPERTY =
-      System.getProperty("io.grpc.internal.DnsNameResolverProvider.enable_jndi", "false");
+      System.getProperty("io.grpc.internal.DnsNameResolverProvider.enable_jndi", "true");
 
   @VisibleForTesting
   static boolean enableJndi = Boolean.parseBoolean(JNDI_PROPERTY);
@@ -95,6 +112,8 @@ final class DnsNameResolver extends NameResolver {
   @GuardedBy("this")
   private Listener listener;
 
+  private final Random random = new Random();
+
   DnsNameResolver(@Nullable String nsAuthority, String name, Attributes params,
       Resource<ScheduledExecutorService> timerServiceResource,
       Resource<ExecutorService> executorResource,
@@ -194,13 +213,29 @@ final class DnsNameResolver extends NameResolver {
 
           Attributes.Builder attrs = Attributes.newBuilder();
           if (!resolvedInetAddrs.txtRecords.isEmpty()) {
-            // TODO(carl-mastrangelo): re enable this
+            JsonObject serviceConfig = null;
-            /*
+            try {
-            attrs.set(
+              JsonArray allServiceConfigChoices = parseTxtResults(resolvedInetAddrs.txtRecords);
-                GrpcAttributes.NAME_RESOLVER_ATTR_SERVICE_CONFIG,
+              for (JsonElement serviceConfigChoice : allServiceConfigChoices) {
-                Collections.unmodifiableList(new ArrayList<String>(resolvedInetAddrs.txtRecords)));
+                try {
-            */
+                  serviceConfig = maybeChooseServiceConfig(
+                      serviceConfigChoice.getAsJsonObject(), random, LOCAL_HOST_NAME);
+                  if (serviceConfig != null) {
+                    break;
+                  }
+                } catch (RuntimeException e) {
+                  logger.log(Level.WARNING, "Bad service config choice " + serviceConfigChoice, e);
+                }
+              }
+              if (serviceConfig != null) {
+                attrs.set(GrpcAttributes.NAME_RESOLVER_ATTR_SERVICE_CONFIG, serviceConfig);
+              }
 
+            } catch (RuntimeException e) {
+              logger.log(Level.WARNING, "Can't parse service Configs", e);
+            }
+          } else {
+            logger.log(Level.FINE, "No TXT records found for {0}", new Object[]{host});
           }
           savedListener.onAddresses(servers, attrs.build());
         } finally {
@@ -387,7 +422,6 @@ final class DnsNameResolver extends NameResolver {
       try {
         serviceConfigTxtRecords = getAllRecords("TXT", "dns:///" + serviceConfigHostname);
       } catch (NamingException e) {
-
         if (logger.isLoggable(Level.FINE)) {
           logger.log(Level.FINE, "Unable to look up " + serviceConfigHostname, e);
         }
@@ -451,7 +485,7 @@ final class DnsNameResolver extends NameResolver {
           NamingEnumeration<?> rrValues = rrEntry.getAll();
           try {
             while (rrValues.hasMore()) {
-              records.add(String.valueOf(rrValues.next()));
+              records.add(unquote(String.valueOf(rrValues.next())));
             }
           } finally {
             rrValues.close();
@@ -463,4 +497,124 @@ final class DnsNameResolver extends NameResolver {
       return records;
     }
   }
+
+  @VisibleForTesting
+  static JsonArray parseTxtResults(List<String> txtRecords) {
+    JsonArray serviceConfigs = new JsonArray();
+    Gson gson = new Gson();
+
+    for (String txtRecord : txtRecords) {
+      if (txtRecord.startsWith(SERVICE_CONFIG_PREFIX)) {
+        JsonArray choices;
+        try {
+          choices =
+              gson.fromJson(txtRecord.substring(SERVICE_CONFIG_PREFIX.length()), JsonArray.class);
+        } catch (JsonSyntaxException e) {
+          logger.log(Level.WARNING, "Bad service config: " + txtRecord, e);
+          continue;
+        }
+        serviceConfigs.addAll(choices);
+      } else {
+        logger.log(Level.FINE, "Ignoring non service config {0}", new Object[]{txtRecord});
+      }
+    }
+
+    return serviceConfigs;
+  }
+
+  /**
+   * Determines if a given Service Config choice applies, and if so, returns it.
+   *
+   * @see <a href="https://github.com/grpc/proposal/blob/master/A2-service-configs-in-dns.md">
+   *   Service Config in DNS</a>
+   * @param choice The service config choice.
+   * @return The service config object or {@code null} if this choice does not apply.
+   */
+  @VisibleForTesting
+  @Nullable
+  @SuppressWarnings("BetaApi") // Verify isn't all that beta
+  static JsonObject maybeChooseServiceConfig(JsonObject choice, Random random, String hostname) {
+    for (Entry<String, ?> entry : choice.entrySet()) {
+      Verify.verify(SERVICE_CONFIG_CHOICE_KEYS.contains(entry.getKey()), "Bad key: %s", entry);
+    }
+    if (choice.has("clientLanguage")) {
+      JsonArray clientLanguages = choice.get("clientLanguage").getAsJsonArray();
+      if (clientLanguages.size() != 0) {
+        boolean javaPresent = false;
+        for (JsonElement clientLanguage : clientLanguages) {
+          String lang = clientLanguage.getAsString().toLowerCase(Locale.ROOT);
+          if ("java".equals(lang)) {
+            javaPresent = true;
+            break;
+          }
+        }
+        if (!javaPresent) {
+          return null;
+        }
+      }
+    }
+    if (choice.has("percentage")) {
+      int pct = choice.get("percentage").getAsInt();
+      Verify.verify(pct >= 0 && pct <= 100, "Bad percentage", choice.get("percentage"));
+      if (pct == 0) {
+        return null;
+      } else if (pct != 100 && pct < random.nextInt(100)) {
+        return null;
+      }
+    }
+    if (choice.has("clientHostname")) {
+      JsonArray clientHostnames = choice.get("clientHostname").getAsJsonArray();
+      if (clientHostnames.size() != 0) {
+        boolean hostnamePresent = false;
+        for (JsonElement clientHostname : clientHostnames) {
+          if (clientHostname.getAsString().equals(hostname)) {
+            hostnamePresent = true;
+            break;
+          }
+        }
+        if (!hostnamePresent) {
+          return null;
+        }
+      }
+    }
+    return choice.getAsJsonObject("serviceConfig");
+  }
+
+  /**
+   * Undo the quoting done in {@link com.sun.jndi.dns.ResourceRecord#decodeTxt}.
+   */
+  @VisibleForTesting
+  static String unquote(String txtRecord) {
+    StringBuilder sb = new StringBuilder(txtRecord.length());
+    boolean inquote = false;
+    for (int i = 0; i < txtRecord.length(); i++) {
+      char c = txtRecord.charAt(i);
+      if (!inquote) {
+        if (c == ' ') {
+          continue;
+        } else if (c == '"') {
+          inquote = true;
+          continue;
+        }
+      } else {
+        if (c == '"') {
+          inquote = false;
+          continue;
+        } else if (c == '\\') {
+          c = txtRecord.charAt(++i);
+          assert c == '"' || c == '\\';
+        }
+      }
+      sb.append(c);
+    }
+    return sb.toString();
+  }
+
+  private static final String getHostName() {
+    try {
+      return InetAddress.getLocalHost().getHostName();
+    } catch (UnknownHostException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

core/src/test/java/io/grpc/internal/DnsNameResolverTest.java

@@ -19,6 +19,7 @@ package io.grpc.internal;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -32,6 +33,9 @@ import static org.mockito.Mockito.when;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.Iterables;
+import com.google.gson.JsonArray;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonPrimitive;
 import io.grpc.Attributes;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.NameResolver;
@@ -50,6 +54,7 @@ import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
+import java.util.Random;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -58,6 +63,7 @@ import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -70,7 +76,8 @@ import org.mockito.MockitoAnnotations;
 @RunWith(JUnit4.class)
 public class DnsNameResolverTest {
 
-  @Rule public final Timeout globalTimeout = Timeout.seconds(10);
+  @Rule public final Timeout globalTimeout = Timeout.seconds(100000);
+  @Rule public final ExpectedException thrown = ExpectedException.none();
 
   private static final int DEFAULT_PORT = 887;
   private static final Attributes NAME_RESOLVER_PARAMS =
@@ -385,6 +392,202 @@ public class DnsNameResolverTest {
     assertTrue(((InetSocketAddress) socketAddress).isUnresolved());
   }
 
+  @Test
+  public void unquoteRemovesJndiFormatting() {
+    assertEquals("blah", DnsNameResolver.unquote("blah"));
+    assertEquals("", DnsNameResolver.unquote("\"\""));
+    assertEquals("blahblah", DnsNameResolver.unquote("blah blah"));
+    assertEquals("blahfoo blah", DnsNameResolver.unquote("blah \"foo blah\""));
+    assertEquals("blah blah", DnsNameResolver.unquote("\"blah blah\""));
+    assertEquals("blah\"blah", DnsNameResolver.unquote("\"blah\\\"blah\""));
+    assertEquals("blah\\blah", DnsNameResolver.unquote("\"blah\\\\blah\""));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_failsOnMisspelling() {
+    JsonObject bad = new JsonObject();
+    bad.add("parcentage", new JsonPrimitive(1));
+    thrown.expectMessage("Bad key");
+
+    DnsNameResolver.maybeChooseServiceConfig(bad, new Random(), "host");
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageMatchesJava() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    langs.add("java");
+    choice.add("clientLanguage", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageDoesntMatchGo() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    langs.add("go");
+    choice.add("clientLanguage", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageCaseInsensitive() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    langs.add("JAVA");
+    choice.add("clientLanguage", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageMatchesEmtpy() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    choice.add("clientLanguage", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageMatchesMulti() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    langs.add("go");
+    langs.add("java");
+    choice.add("clientLanguage", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_percentageZeroAlwaysFails() {
+    JsonObject choice = new JsonObject();
+    choice.add("percentage", new JsonPrimitive(0));
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_percentageHundredAlwaysSucceeds() {
+    JsonObject choice = new JsonObject();
+    choice.add("percentage", new JsonPrimitive(100));
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_percentageAboveMatches() {
+    JsonObject choice = new JsonObject();
+    choice.add("percentage", new JsonPrimitive(50));
+    choice.add("serviceConfig", new JsonObject());
+
+    Random r = new Random() {
+      @Override
+      public int nextInt(int bound) {
+        return 49;
+      }
+    };
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, r, "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_percentageAtMatches() {
+    JsonObject choice = new JsonObject();
+    choice.add("percentage", new JsonPrimitive(50));
+    choice.add("serviceConfig", new JsonObject());
+
+    Random r = new Random() {
+      @Override
+      public int nextInt(int bound) {
+        return 50;
+      }
+    };
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, r, "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_percentageBelowFails() {
+    JsonObject choice = new JsonObject();
+    choice.add("percentage", new JsonPrimitive(50));
+    choice.add("serviceConfig", new JsonObject());
+
+    Random r = new Random() {
+      @Override
+      public int nextInt(int bound) {
+        return 51;
+      }
+    };
+
+    assertNull(DnsNameResolver.maybeChooseServiceConfig(choice, r, "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_hostnameMatches() {
+    JsonObject choice = new JsonObject();
+    JsonArray hosts = new JsonArray();
+    hosts.add("localhost");
+    choice.add("clientHostname", hosts);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "localhost"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_hostnameDoesntMatch() {
+    JsonObject choice = new JsonObject();
+    JsonArray hosts = new JsonArray();
+    hosts.add("localhorse");
+    choice.add("clientHostname", hosts);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "localhost"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_clientLanguageCaseSensitive() {
+    JsonObject choice = new JsonObject();
+    JsonArray hosts = new JsonArray();
+    hosts.add("LOCALHOST");
+    choice.add("clientHostname", hosts);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "localhost"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_hostnameMatchesEmtpy() {
+    JsonObject choice = new JsonObject();
+    JsonArray hosts = new JsonArray();
+    choice.add("clientHostname", hosts);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "host"));
+  }
+
+  @Test
+  public void maybeChooseServiceConfig_hostnameMatchesMulti() {
+    JsonObject choice = new JsonObject();
+    JsonArray langs = new JsonArray();
+    langs.add("localhorse");
+    langs.add("localhost");
+    choice.add("clientHostname", langs);
+    choice.add("serviceConfig", new JsonObject());
+
+    assertNotNull(DnsNameResolver.maybeChooseServiceConfig(choice, new Random(), "localhost"));
+  }
+
   private void testInvalidUri(URI uri) {
     try {
       provider.newNameResolver(uri, NAME_RESOLVER_PARAMS);