From 41b5946bac0ef3391e616e404b0a747d10e9de50 Mon Sep 17 00:00:00 2001
From: Piotr Findeisen <piotr.findeisen@gmail.com>
Date: Mon, 25 Feb 2019 13:12:30 +0100
Subject: [PATCH] Secure nextUri with a slug

---
 .../presto/server/protocol/Query.java         |  9 +++++++
 .../server/protocol/StatementResource.java    | 24 +++++++++++++++----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java
index 76fe95c0eddbe..5c1936cb472a9 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/protocol/Query.java
@@ -85,7 +85,9 @@
 import static io.airlift.concurrent.MoreFutures.addSuccessCallback;
 import static io.airlift.concurrent.MoreFutures.addTimeout;
 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
@@ -94,6 +96,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;
@@ -242,6 +245,11 @@ public QueryId getQueryId()
         return queryId;
     }
 
+    public boolean isSlugValid(String slug)
+    {
+        return this.slug.equals(slug);
+    }
+
     public synchronized Optional<String> getSetCatalog()
     {
         return setCatalog;
@@ -571,6 +579,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/com/facebook/presto/server/protocol/StatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java
index 3f10205c2e6db..c4c03b88aea0c 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/protocol/StatementResource.java
@@ -34,6 +34,7 @@
 import org.weakref.jmx.Managed;
 import org.weakref.jmx.Nested;
 
+import javax.annotation.Nullable;
 import javax.annotation.PreDestroy;
 import javax.inject.Inject;
 import javax.servlet.http.HttpServletRequest;
@@ -171,10 +172,11 @@ 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,
@@ -182,7 +184,7 @@ public void getQueryResults(
             @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;
@@ -194,6 +196,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,
@@ -262,12 +274,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();
         }