Secure nextUri with a slug

trinodb · Mar 28, 2019 · c57cb19 · c57cb19
1 parent 4b6a52d
commit c57cb19
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/Query.java b/presto-main/src/main/java/io/prestosql/server/protocol/Query.java
@@ -96,7 +96,9 @@
 import static io.prestosql.util.Failures.toFailure;
 import static io.prestosql.util.MoreLists.mappedCopy;
 import static java.lang.String.format;
+import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
+import static java.util.UUID.randomUUID;
 
 @ThreadSafe
 class Query
@@ -105,6 +107,7 @@ class Query
 
     private final QueryManager queryManager;
     private final QueryId queryId;
+    private final String slug = "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", "");
 
     @GuardedBy("this")
     private final ExchangeClient exchangeClient;
@@ -253,6 +256,11 @@ public QueryId getQueryId()
         return queryId;
     }
 
+    public boolean isSlugValid(String slug)
+    {
+        return this.slug.equals(slug);
+    }
+
     public synchronized Optional<String> getSetCatalog()
     {
         return setCatalog;
@@ -584,6 +592,7 @@ private synchronized URI createNextResultsUri(String scheme, UriInfo uriInfo)
                 .scheme(scheme)
                 .replacePath("/v1/statement")
                 .path(queryId.toString())
+                .path(slug)
                 .path(String.valueOf(resultId.incrementAndGet()))
                 .replaceQuery("")
                 .build();

diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/StatementResource.java b/presto-main/src/main/java/io/prestosql/server/protocol/StatementResource.java
@@ -31,6 +31,7 @@
 import io.prestosql.spi.QueryId;
 import io.prestosql.spi.block.BlockEncodingSerde;
 
+import javax.annotation.Nullable;
 import javax.annotation.PreDestroy;
 import javax.inject.Inject;
 import javax.servlet.http.HttpServletRequest;
@@ -164,18 +165,19 @@ public Response createQuery(
     }
 
     @GET
-    @Path("{queryId}/{token}")
+    @Path("{queryId}/{slug}/{token}")
     @Produces(MediaType.APPLICATION_JSON)
     public void getQueryResults(
             @PathParam("queryId") QueryId queryId,
+            @PathParam("slug") String slug,
             @PathParam("token") long token,
             @QueryParam("maxWait") Duration maxWait,
             @QueryParam("targetResultSize") DataSize targetResultSize,
             @HeaderParam(X_FORWARDED_PROTO) String proto,
             @Context UriInfo uriInfo,
             @Suspended AsyncResponse asyncResponse)
     {
-        Query query = queries.get(queryId);
+        Query query = getQuery(queryId, slug);
         if (query == null) {
             asyncResponse.resume(Response.status(Status.NOT_FOUND).build());
             return;
@@ -187,6 +189,16 @@ public void getQueryResults(
         asyncQueryResults(query, OptionalLong.of(token), maxWait, targetResultSize, uriInfo, proto, asyncResponse);
     }
 
+    @Nullable
+    private Query getQuery(QueryId queryId, String slug)
+    {
+        Query query = queries.get(queryId);
+        if (query != null && query.isSlugValid(slug)) {
+            return query;
+        }
+        return null;
+    }
+
     private void asyncQueryResults(
             Query query,
             OptionalLong token,
@@ -255,12 +267,14 @@ private static Response toResponse(Query query, QueryResults queryResults)
     }
 
     @DELETE
-    @Path("{queryId}/{token}")
+    @Path("{queryId}/{slug}/{token}")
     @Produces(MediaType.APPLICATION_JSON)
-    public Response cancelQuery(@PathParam("queryId") QueryId queryId,
+    public Response cancelQuery(
+            @PathParam("queryId") QueryId queryId,
+            @PathParam("slug") String slug,
             @PathParam("token") long token)
     {
-        Query query = queries.get(queryId);
+        Query query = getQuery(queryId, slug);
         if (query == null) {
             return Response.status(Status.NOT_FOUND).build();
         }