Make symbol table path handling more robust to handle D2 URIs

RB=1896287 R=kbalasub A=kbalasub
linkedin · Dec 6, 2019 · 3a80c39 · 3a80c39
1 parent 88036a7
commit 3a80c39
Show file tree

Hide file tree

Showing 10 changed files with 172 additions and 94 deletions.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,5 +1,11 @@
+28.0.12
+-------
+
 28.0.11
 -------
+(RB=1896287)
+Make symbol table path handling more robust to handle D2 URIs
+
 (RB=1903191)
 Allow incompatible rest model changes in restli-int-test-server
 

diff --git a/gradle.properties b/gradle.properties
@@ -1,4 +1,4 @@
-version=28.0.10
+version=28.0.11
 sonatypeUsername=please_set_in_home_dir_if_uploading_to_maven_central
 sonatypePassword=please_set_in_home_dir_if_uploading_to_maven_central
 org.gradle.configureondemand=true

diff --git a/restli-common/src/main/java/com/linkedin/restli/common/RestConstants.java b/restli-common/src/main/java/com/linkedin/restli/common/RestConstants.java
@@ -51,6 +51,7 @@ public interface RestConstants
   String HEADER_RESTLI_PROTOCOL_VERSION = "X-RestLi-Protocol-Version";
   String CONTENT_TYPE_PARAM_SYMBOL_TABLE = "symbol-table";
   String HEADER_CONTENT_ID = "Content-ID";
+  String HEADER_SERVICE_SCOPED_PATH = "x-restli-service-scoped-path";
 
   // Default supported mime types.
   Set<String> SUPPORTED_MIME_TYPES = new LinkedHashSet<>(

diff --git a/...rver/src/main/java/com/linkedin/restli/server/symbol/RestLiSymbolTableRequestHandler.java b/...rver/src/main/java/com/linkedin/restli/server/symbol/RestLiSymbolTableRequestHandler.java
@@ -49,7 +49,7 @@
  */
 public class RestLiSymbolTableRequestHandler implements NonResourceRequestHandler
 {
-  public static final String SYMBOL_TABLE_URI_PATH = "/symbolTable";
+  public static final String SYMBOL_TABLE_URI_PATH = "symbolTable";
   private static final Logger LOGGER = LoggerFactory.getLogger(RestLiSymbolTableRequestHandler.class);
   private static final int DEFAULT_CACHE_SIZE = 100;
 
@@ -76,10 +76,32 @@ public RestLiSymbolTableRequestHandler(int cacheSize)
   @Override
   public boolean shouldHandle(Request request)
   {
+    // we don't check the method here because we want to return 405 if it is anything but GET
     final String path = request.getURI().getRawPath();
+    if (path == null)
+    {
+      return false;
+    }
 
-    // we don't check the method here because we want to return 405 if it is anything but GET
-    return path != null && path.startsWith(SYMBOL_TABLE_URI_PATH);
+    final List<UriComponent.PathSegment> pathSegments = UriComponent.decodePath(path, true);
+    if (pathSegments.size() < 2)
+    {
+      return false;
+    }
+
+    //
+    // When path is service scoped, URI is in the form of /<SERVICE>/symbolTable, else it
+    // is in the form of /symbolTable or /symbolTable/<TABLENAME>
+    //
+    boolean isServiceScopedPath = request.getHeaders().containsKey(RestConstants.HEADER_SERVICE_SCOPED_PATH);
+    if (isServiceScopedPath)
+    {
+      return (pathSegments.size() == 3 && pathSegments.get(2).getPath().equals(SYMBOL_TABLE_URI_PATH));
+    }
+    else
+    {
+      return ((pathSegments.size() == 2 || pathSegments.size() == 3) && pathSegments.get(1).getPath().equals(SYMBOL_TABLE_URI_PATH));
+    }
   }
 
   @Override
@@ -117,22 +139,24 @@ public void handleRequest(RestRequest request, RequestContext requestContext, Ca
     final List<UriComponent.PathSegment> pathSegments = UriComponent.decodePath(path, true);
 
     final SymbolTableProvider provider = SymbolTableProviderHolder.INSTANCE.getSymbolTableProvider();
-    SymbolTable symbolTable;
+    SymbolTable symbolTable = null;
     try
     {
-      if (pathSegments.size() == 2)
+      boolean isServiceScopedPath = request.getHeaders().containsKey(RestConstants.HEADER_SERVICE_SCOPED_PATH);
+      if (isServiceScopedPath)
       {
         symbolTable = provider.getResponseSymbolTable(request.getURI(), request.getHeaders());
       }
-      else if (pathSegments.size() == 3)
-      {
-        symbolTable = provider.getSymbolTable(pathSegments.get(2).getPath());
-      }
       else
       {
-        LOGGER.error("Invalid path " + path);
-        callback.onError(RestException.forError(HttpStatus.S_400_BAD_REQUEST.getCode(), "Invalid path"));
-        return;
+        if (pathSegments.size() == 2)
+        {
+          symbolTable = provider.getResponseSymbolTable(request.getURI(), request.getHeaders());
+        }
+        else if (pathSegments.size() == 3)
+        {
+          symbolTable = provider.getSymbolTable(pathSegments.get(2).getPath());
+        }
       }
     }
     catch (IllegalStateException e)

diff --git a/...-server/src/test/java/com/linkedin/restli/server/TestRestLiSymbolTableRequestHandler.java b/...-server/src/test/java/com/linkedin/restli/server/TestRestLiSymbolTableRequestHandler.java
@@ -66,28 +66,43 @@ public void tearDown()
   @Test
   public void testShouldNotHandleEmptyPath()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service")).build();
+    RestRequest request = new RestRequestBuilder(URI.create("/")).build();
     Assert.assertFalse(_requestHandler.shouldHandle(request));
   }
 
   @Test
   public void testShouldNotHandleNonSymbolTablePath()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/someResource")).build();
+    RestRequest request = new RestRequestBuilder(URI.create("/someResource")).build();
     Assert.assertFalse(_requestHandler.shouldHandle(request));
   }
 
   @Test
   public void testShouldHandleSymbolTablePath()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/symbolTable")).build();
+    RestRequest request = new RestRequestBuilder(URI.create("/symbolTable")).build();
     Assert.assertTrue(_requestHandler.shouldHandle(request));
   }
 
+  @Test
+  public void testShouldHandleServiceScopedSymbolTablePathWhenHeaderSet()
+  {
+    RestRequest request = new RestRequestBuilder(URI.create("/service/symbolTable"))
+        .setHeaders(Collections.singletonMap(RestConstants.HEADER_SERVICE_SCOPED_PATH, "true")).build();
+    Assert.assertTrue(_requestHandler.shouldHandle(request));
+  }
+
+  @Test
+  public void testShouldNotHandleServiceScopedSymbolTablePathWhenNoHeaderSet()
+  {
+    RestRequest request = new RestRequestBuilder(URI.create("/service/symbolTable")).build();
+    Assert.assertFalse(_requestHandler.shouldHandle(request));
+  }
+
   @Test
   public void testNonGetRequest405()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/symbolTable"))
+    RestRequest request = new RestRequestBuilder(URI.create("/symbolTable"))
         .setMethod(HttpMethod.POST.name()).build();
 
     CompletableFuture<RestResponse> future = new CompletableFuture<>();
@@ -111,7 +126,7 @@ public void onSuccess(RestResponse result) {
   @Test
   public void testInvalidAcceptTypeRequest406()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/symbolTable"))
+    RestRequest request = new RestRequestBuilder(URI.create("/symbolTable"))
         .setHeader(RestConstants.HEADER_ACCEPT, "application/randomType").build();
 
     CompletableFuture<RestResponse> future = new CompletableFuture<>();
@@ -134,7 +149,7 @@ public void onSuccess(RestResponse result) {
   @Test
   public void testSelfSymbolTableNotFound404()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/symbolTable")).build();
+    RestRequest request = new RestRequestBuilder(URI.create("/symbolTable")).build();
 
     CompletableFuture<RestResponse> future = new CompletableFuture<>();
     _requestHandler.handleRequest(request, mock(RequestContext.class), new Callback<RestResponse>() {
@@ -156,7 +171,7 @@ public void onSuccess(RestResponse result) {
   @Test
   public void testOtherSymbolTableNotFound404()
   {
-    RestRequest request = new RestRequestBuilder(URI.create("d2://service/symbolTable/SomeName")).build();
+    RestRequest request = new RestRequestBuilder(URI.create("/symbolTable/SomeName")).build();
 
     CompletableFuture<RestResponse> future = new CompletableFuture<>();
     _requestHandler.handleRequest(request, mock(RequestContext.class), new Callback<RestResponse>() {
@@ -180,7 +195,7 @@ public void testReturnSelfSymbolTable() throws Exception
   {
     SymbolTable symbolTable =
         new InMemorySymbolTable("TestName", ImmutableList.of("Haha", "Hehe", "Hoho"));
-    URI uri = URI.create("d2://service/symbolTable");
+    URI uri = URI.create("/symbolTable");
     RestRequest request = new RestRequestBuilder(uri).build();
     when(_symbolTableProvider.getResponseSymbolTable(eq(uri), eq(Collections.emptyMap()))).thenReturn(symbolTable);
 
@@ -206,12 +221,43 @@ public void onSuccess(RestResponse result) {
     Assert.assertEquals(symbolTable, SymbolTableSerializer.fromByteString(response.getEntity(), ContentType.PROTOBUF.getCodec()));
   }
 
+  @Test
+  public void testReturnSelfSymbolTableWhenCalledWithServiceScope() throws Exception
+  {
+    SymbolTable symbolTable =
+        new InMemorySymbolTable("TestName", ImmutableList.of("Haha", "Hehe", "Hoho"));
+    URI uri = URI.create("/service/symbolTable");
+    RestRequest request = new RestRequestBuilder(uri).setHeader(RestConstants.HEADER_SERVICE_SCOPED_PATH, "true").build();
+    when(_symbolTableProvider.getResponseSymbolTable(eq(uri), eq(Collections.singletonMap(RestConstants.HEADER_SERVICE_SCOPED_PATH, "true")))).thenReturn(symbolTable);
+
+    CompletableFuture<RestResponse> future = new CompletableFuture<>();
+    _requestHandler.handleRequest(request, mock(RequestContext.class), new Callback<RestResponse>() {
+      @Override
+      public void onError(Throwable e) {
+        future.completeExceptionally(e);
+      }
+
+      @Override
+      public void onSuccess(RestResponse result) {
+        future.complete(result);
+      }
+    });
+
+    Assert.assertFalse(future.isCompletedExceptionally());
+    Assert.assertTrue(future.isDone());
+
+    RestResponse response = future.get();
+    Assert.assertEquals(response.getStatus(), HttpStatus.S_200_OK.getCode());
+    Assert.assertEquals(response.getHeader(RestConstants.HEADER_CONTENT_TYPE), ContentType.PROTOBUF.getHeaderKey());
+    Assert.assertEquals(symbolTable, SymbolTableSerializer.fromByteString(response.getEntity(), ContentType.PROTOBUF.getCodec()));
+  }
+
   @Test
   public void testReturnOtherSymbolTable() throws Exception
   {
     SymbolTable symbolTable =
         new InMemorySymbolTable("TestName", ImmutableList.of("Haha", "Hehe", "Hoho"));
-    URI uri = URI.create("d2://service/symbolTable/OtherTable");
+    URI uri = URI.create("/symbolTable/OtherTable");
     RestRequest request = new RestRequestBuilder(uri).build();
     when(_symbolTableProvider.getSymbolTable(eq("OtherTable"))).thenReturn(symbolTable);
 
@@ -242,7 +288,7 @@ public void testReturnOtherSymbolTableNonDefaultAcceptType() throws Exception
   {
     SymbolTable symbolTable =
         new InMemorySymbolTable("TestName", ImmutableList.of("Haha", "Hehe", "Hoho"));
-    URI uri = URI.create("d2://service/symbolTable/OtherTable");
+    URI uri = URI.create("/symbolTable/OtherTable");
     RestRequest request =
         new RestRequestBuilder(uri).setHeader(RestConstants.HEADER_ACCEPT, ContentType.JSON.getHeaderKey()).build();
     when(_symbolTableProvider.getSymbolTable(eq("OtherTable"))).thenReturn(symbolTable);

diff --git a/restli-tools/src/main/java/com/linkedin/restli/tools/symbol/RestLiSymbolTableProvider.java b/restli-tools/src/main/java/com/linkedin/restli/tools/symbol/RestLiSymbolTableProvider.java
@@ -37,6 +37,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -80,7 +81,6 @@ public class RestLiSymbolTableProvider implements SymbolTableProvider, ResourceD
    * Default timeout in milliseconds to use when fetching symbols from other services.
    */
   private static final long DEFAULT_TIMEOUT_MILLIS = 100;
-  private static final String HTTPS_PREFIX = "https://";
 
   private final Client _client;
   private final String _uriPrefix;
@@ -95,21 +95,20 @@ public class RestLiSymbolTableProvider implements SymbolTableProvider, ResourceD
    *
    * @param client             The {@link Client} to use to make requests to remote services to fetch their symbol tables.
    * @param uriPrefix          The URI prefix to use when invoking remote services by name (and not by hostname:port)
-   * @param symbolTablePrefix  The prefix to use for symbol tables vended by this instance.
    * @param cacheSize          The size of the caches used to store symbol tables.
-   * @param serverHostName     The hostname of the machine on which the current server is running.
-   * @param serverPort         The https port on which the current server is running.
+   * @param symbolTablePrefix  The prefix to use for symbol tables vended by this instance.
+   * @param serverNodeUri      The URI on which the current service is running. This should also include the context
+   *                           and servlet path (if applicable).
    */
   public RestLiSymbolTableProvider(Client client,
       String uriPrefix,
-      String symbolTablePrefix,
       int cacheSize,
-      String serverHostName,
-      int serverPort)
+      String symbolTablePrefix,
+      String serverNodeUri)
   {
     _client = client;
     _uriPrefix = uriPrefix;
-    _symbolTableNameHandler = new SymbolTableNameHandler(symbolTablePrefix, serverHostName, serverPort);
+    _symbolTableNameHandler = new SymbolTableNameHandler(symbolTablePrefix, serverNodeUri);
     _serviceNameToSymbolTableCache = Caffeine.newBuilder().maximumSize(cacheSize).build();
     _symbolTableNameToSymbolTableCache = Caffeine.newBuilder().maximumSize(cacheSize).build();
   }
@@ -120,7 +119,7 @@ public SymbolTable getSymbolTable(String symbolTableName)
     try
     {
       Tuple3<String, String, Boolean> tuple = _symbolTableNameHandler.extractTableInfo(symbolTableName);
-      String hostNameAndPort = tuple._1();
+      String serverNodeUri = tuple._1();
       String tableName = tuple._2();
       boolean isLocal = tuple._3();
 
@@ -143,10 +142,9 @@ public SymbolTable getSymbolTable(String symbolTableName)
         throw new IllegalStateException("Unable to fetch symbol table with name: " + symbolTableName);
       }
 
-      // Ok, we didn't find it in the cache, let's go query the host the table was served from.
-      URI symbolTableUri = new URI(
-          HTTPS_PREFIX + hostNameAndPort + RestLiSymbolTableRequestHandler.SYMBOL_TABLE_URI_PATH + "/" + tableName);
-      symbolTable = fetchRemoteSymbolTable(symbolTableUri);
+      // Ok, we didn't find it in the cache, let's go query the service the table was served from.
+      URI symbolTableUri = new URI(serverNodeUri + "/" + RestLiSymbolTableRequestHandler.SYMBOL_TABLE_URI_PATH + "/" + tableName);
+      symbolTable = fetchRemoteSymbolTable(symbolTableUri, Collections.emptyMap());
 
       if (symbolTable != null)
       {
@@ -175,11 +173,14 @@ public SymbolTable getRequestSymbolTable(URI requestUri)
       return symbolTable;
     }
 
-    // Ok, we didn't find it in the cache, let's go query the other service.
+    // Ok, we didn't find it in the cache, let's go query the other service using the URI prefix. In this case, we
+    // make sure to set the {@link RestConstants#HEADER_SERVICE_SCOPED_PATH} header to true to indicate that this
+    // path, post resolution must be interpreted as a service scoped path.
     try
     {
-      URI symbolTableUri = new URI(_uriPrefix + serviceName + RestLiSymbolTableRequestHandler.SYMBOL_TABLE_URI_PATH);
-      symbolTable = fetchRemoteSymbolTable(symbolTableUri);
+      URI symbolTableUri = new URI(_uriPrefix + serviceName + "/" + RestLiSymbolTableRequestHandler.SYMBOL_TABLE_URI_PATH);
+      symbolTable = fetchRemoteSymbolTable(symbolTableUri,
+          Collections.singletonMap(RestConstants.HEADER_SERVICE_SCOPED_PATH, Boolean.TRUE.toString()));
 
       if (symbolTable != null)
       {
@@ -214,11 +215,11 @@ public void onInitialized(Map<String, ResourceDefinition> resourceDefinitions)
     _defaultResponseSymbolTableName = _symbolTableNameHandler.extractTableInfo(_defaultResponseSymbolTable.getName())._2();
   }
 
-  SymbolTable fetchRemoteSymbolTable(URI symbolTableUri)
+  SymbolTable fetchRemoteSymbolTable(URI symbolTableUri, Map<String, String> requestHeaders)
   {
     try
     {
-      Future<RestResponse> future = _client.restRequest(new RestRequestBuilder(symbolTableUri).build());
+      Future<RestResponse> future = _client.restRequest(new RestRequestBuilder(symbolTableUri).setHeaders(requestHeaders).build());
       RestResponse restResponse = future.get(DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
       int status = restResponse.getStatus();
       if (status == HttpStatus.S_200_OK.getCode())
@@ -233,8 +234,8 @@ SymbolTable fetchRemoteSymbolTable(URI symbolTableUri)
             ContentType.getContentType(restResponse.getHeader(RestConstants.HEADER_CONTENT_TYPE))
                 .orElseThrow(() -> new IOException("Could not parse response content type"));
 
-        // Deserialize, and rename to replace hostname with current hostname.
-        return SymbolTableSerializer.fromByteString(byteString, contentType.getCodec(), _symbolTableNameHandler::replaceHostName);
+        // Deserialize, and rename to replace url prefix with current url prefix.
+        return SymbolTableSerializer.fromByteString(byteString, contentType.getCodec(), _symbolTableNameHandler::replaceServerNodeUri);
       }
 
       throw new IOException("Unexpected response status: " + status);