Add canary service to GKE (#2594)

core/src/main/java/google/registry/tools/CurlCommand.java

@@ -21,6 +21,7 @@ import com.beust.jcommander.IStringConverter;
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
 import com.beust.jcommander.converters.IParameterSplitter;
+import com.google.common.base.Ascii;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
@@ -103,7 +104,10 @@ class CurlCommand implements CommandWithConnection {
       throw new IllegalArgumentException("You may not specify a body for a get method.");
     }
 
-    Service service = useGke ? GkeService.valueOf(serviceName) : GaeService.valueOf(serviceName);
+    Service service =
+        useGke
+            ? GkeService.valueOf(Ascii.toUpperCase(serviceName))
+            : GaeService.valueOf(Ascii.toUpperCase(serviceName));
 
     ServiceConnection connectionToService = connection.withService(service, canary);
     String response =

core/src/main/java/google/registry/tools/ServiceConnection.java

@@ -14,7 +14,6 @@
 
 package google.registry.tools;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.base.Verify.verify;
@@ -35,7 +34,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.io.CharStreams;
 import com.google.common.net.MediaType;
-import com.google.re2j.Matcher;
 import com.google.re2j.Pattern;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.request.Action.GaeService;
@@ -60,6 +58,8 @@ public class ServiceConnection {
   /** Pattern to heuristically extract title tag contents in HTML responses. */
   protected static final Pattern HTML_TITLE_TAG_PATTERN = Pattern.compile("<title>(.*?)</title>");
 
+  private static final String CANARY_HEADER = "canary";
+
   private final Service service;
   private final boolean useCanary;
   private final HttpRequestFactory requestFactory;
@@ -70,9 +70,6 @@ public class ServiceConnection {
   }
 
   private ServiceConnection(Service service, HttpRequestFactory requestFactory, boolean useCanary) {
-    // Currently, only GAE supports connecting to canary.
-    // TODO (jianglai): decide how to implement canary for GKE.
-    checkArgument(useCanary == false || service instanceof GaeService, "Canary is only for GAE");
     this.service = service;
     this.requestFactory = requestFactory;
     this.useCanary = useCanary;
@@ -80,15 +77,17 @@ public class ServiceConnection {
 
   /** Returns a copy of this connection that talks to a different service endpoint. */
   public ServiceConnection withService(Service service, boolean useCanary) {
+    Class<? extends Service> oldServiceClazz = this.service.getClass();
+    Class<? extends Service> newServiceClazz = service.getClass();
+    if (oldServiceClazz != newServiceClazz) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Cannot switch from %s to %s",
+              oldServiceClazz.getSimpleName(), newServiceClazz.getSimpleName()));
+    }
     return new ServiceConnection(service, requestFactory, useCanary);
   }
 
-  /** Returns the contents of the title tag in the given HTML, or null if not found. */
-  private static String extractHtmlTitle(String html) {
-    Matcher matcher = HTML_TITLE_TAG_PATTERN.matcher(html);
-    return (matcher.find() ? matcher.group(1) : null);
-  }
-
   /** Returns the HTML from the connection error stream, if any, otherwise the empty string. */
   private static String getErrorHtmlAsString(HttpResponse response) throws IOException {
     return CharStreams.toString(new InputStreamReader(response.getContent(), UTF_8));
@@ -107,19 +106,22 @@ public class ServiceConnection {
     HttpHeaders headers = request.getHeaders();
     headers.setCacheControl("no-cache");
     headers.put(X_REQUESTED_WITH, ImmutableList.of("RegistryTool"));
+    if (useCanary) {
+      headers.set(CANARY_HEADER, "true");
+    }
     request.setHeaders(headers);
     request.setFollowRedirects(false);
     request.setThrowExceptionOnExecuteError(false);
     request.setUnsuccessfulResponseHandler(
         (request1, response, supportsRetry) -> {
-          String errorTitle = extractHtmlTitle(getErrorHtmlAsString(response));
+          String error = getErrorHtmlAsString(response);
           throw new IOException(
               String.format(
                   "Error from %s: %d %s%s",
                   request1.getUrl().toString(),
                   response.getStatusCode(),
                   response.getStatusMessage(),
-                  (errorTitle == null ? "" : ": " + errorTitle)));
+                  error));
         });
     HttpResponse response = null;
     try {
@@ -135,8 +137,8 @@ public class ServiceConnection {
   @VisibleForTesting
   URL getServer() {
     URL url = service.getServiceUrl();
-    if (useCanary) {
-      verify(!isNullOrEmpty(url.getHost()), "Null host in url");
+    verify(!isNullOrEmpty(url.getHost()), "Null host in url");
+    if (useCanary && service instanceof GaeService) {
       url =
           makeUrl(
               String.format(

core/src/test/java/google/registry/tools/ServiceConnectionTest.java

@@ -16,23 +16,65 @@ package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.request.Action.GaeService.DEFAULT;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.HttpResponse;
+import com.google.common.collect.ImmutableMap;
+import google.registry.request.Action.GkeService;
+import java.io.ByteArrayInputStream;
 import org.junit.jupiter.api.Test;
 
 /** Unit tests for {@link google.registry.tools.ServiceConnection}. */
 public class ServiceConnectionTest {
 
   @Test
-  void testServerUrl_notCanary() {
+  void testSuccess_serverUrl_notCanary() {
     ServiceConnection connection = new ServiceConnection(false, null).withService(DEFAULT, false);
     String serverUrl = connection.getServer().toString();
     assertThat(serverUrl).isEqualTo("https://default.example.com"); // See default-config.yaml
   }
 
   @Test
-  void testServerUrl_canary() {
+  void testFailure_mixedService() throws Exception {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> {
+              new ServiceConnection(true, null).withService(DEFAULT, true);
+            });
+    assertThat(thrown).hasMessageThat().contains("Cannot switch from GkeService to GaeService");
+  }
+
+  @Test
+  void testSuccess_serverUrl_gae_canary() {
     ServiceConnection connection = new ServiceConnection(false, null).withService(DEFAULT, true);
     String serverUrl = connection.getServer().toString();
     assertThat(serverUrl).isEqualTo("https://nomulus-dot-default.example.com");
   }
+
+  @Test
+  void testSuccess_serverUrl_gke_canary() throws Exception {
+    HttpRequestFactory factory = mock(HttpRequestFactory.class);
+    HttpRequest request = mock(HttpRequest.class);
+    HttpHeaders headers = mock(HttpHeaders.class);
+    HttpResponse response = mock(HttpResponse.class);
+    when(request.getHeaders()).thenReturn(headers);
+    when(factory.buildGetRequest(any(GenericUrl.class))).thenReturn(request);
+    when(request.execute()).thenReturn(response);
+    when(response.getContent()).thenReturn(ByteArrayInputStream.nullInputStream());
+    ServiceConnection connection =
+        new ServiceConnection(true, factory).withService(GkeService.PUBAPI, true);
+    String serverUrl = connection.getServer().toString();
+    assertThat(serverUrl).isEqualTo("https://pubapi.registry.test");
+    connection.sendGetRequest("/path", ImmutableMap.of());
+    verify(headers).set("canary", "true");
+  }
 }