Merge pull request #1817 from HubSpot/file_download_proxy

Proxy download from Mesos slave over Singularity.
HubSpot · Aug 1, 2018 · 4c119da · 4c119da
2 parents 9910a1e + a76151f
commit 4c119da
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 3 deletions.
diff --git a/SingularityService/src/main/java/com/hubspot/singularity/resources/TaskResource.java b/SingularityService/src/main/java/com/hubspot/singularity/resources/TaskResource.java
@@ -5,6 +5,9 @@
 import static com.hubspot.singularity.WebExceptions.checkNotFound;
 import static com.hubspot.singularity.WebExceptions.notFound;
 
+import java.io.InputStream;
+import java.net.ConnectException;
+import java.nio.file.Paths;
 import java.util.Collections;
 import java.util.List;
 
@@ -69,9 +72,11 @@
 import com.hubspot.singularity.api.SingularityTaskMetadataRequest;
 import com.hubspot.singularity.auth.SingularityAuthorizationHelper;
 import com.hubspot.singularity.config.ApiPaths;
+import com.hubspot.singularity.config.MesosConfiguration;
 import com.hubspot.singularity.config.SingularityTaskMetadataConfiguration;
 import com.hubspot.singularity.data.DisasterManager;
 import com.hubspot.singularity.data.RequestManager;
+import com.hubspot.singularity.data.SandboxManager.SlaveNotFoundException;
 import com.hubspot.singularity.data.SingularityValidator;
 import com.hubspot.singularity.data.SlaveManager;
 import com.hubspot.singularity.data.TaskManager;
@@ -95,6 +100,8 @@
 public class TaskResource extends AbstractLeaderAwareResource {
   private static final Logger LOG = LoggerFactory.getLogger(TaskResource.class);
 
+  private final AsyncHttpClient httpClient;
+  private final MesosConfiguration configuration;
   private final TaskManager taskManager;
   private final RequestManager requestManager;
   private final SlaveManager slaveManager;
@@ -109,7 +116,7 @@ public class TaskResource extends AbstractLeaderAwareResource {
   @Inject
   public TaskResource(TaskRequestManager taskRequestManager, TaskManager taskManager, SlaveManager slaveManager, MesosClient mesosClient, SingularityTaskMetadataConfiguration taskMetadataConfiguration,
                       SingularityAuthorizationHelper authorizationHelper, RequestManager requestManager, SingularityValidator validator, DisasterManager disasterManager,
-                      AsyncHttpClient httpClient, LeaderLatch leaderLatch, ObjectMapper objectMapper, RequestHelper requestHelper) {
+                      AsyncHttpClient httpClient, LeaderLatch leaderLatch, ObjectMapper objectMapper, RequestHelper requestHelper, MesosConfiguration configuration) {
     super(httpClient, leaderLatch, objectMapper);
     this.taskManager = taskManager;
     this.taskRequestManager = taskRequestManager;
@@ -121,6 +128,8 @@ public TaskResource(TaskRequestManager taskRequestManager, TaskManager taskManag
     this.validator = validator;
     this.disasterManager = disasterManager;
     this.requestHelper = requestHelper;
+    this.httpClient = httpClient;
+    this.configuration = configuration;
   }
 
   @GET
@@ -609,4 +618,42 @@ public List<SingularityTaskShellCommandUpdate> getShellCommandHisotryUpdates(
     SingularityTaskId taskIdObj = getTaskIdFromStr(taskId);
     return taskManager.getTaskShellCommandUpdates(new SingularityTaskShellCommandRequestId(taskIdObj, commandName, commandTimestamp));
   }
+
+  @GET
+  @Path("/download/")
+  @Produces(MediaType.APPLICATION_OCTET_STREAM)
+  @Operation(summary = "Proxy a file download from a Mesos Slave through Singularity")
+  public Response downloadFileOverProxy(
+      @Parameter(required = true, description = "Mesos slave hostname") @QueryParam("slaveHostname") String slaveHostname,
+      @Parameter(required = true, description = "Full file path to file on Mesos slave to be downloaded") @QueryParam("path") String fileFullPath
+  ) {
+    String httpPrefix = configuration.getSlaveHttpsPort().isPresent() ? "https" : "http";
+    int httpPort = configuration.getSlaveHttpsPort().isPresent() ? configuration.getSlaveHttpsPort().get() : configuration.getSlaveHttpPort();
+
+    String url = String.format("%s://%s:%s/files/download.json",
+            httpPrefix, slaveHostname, httpPort);
+
+    try {
+      final InputStream responseStream = httpClient.prepareGet(url)
+          .addQueryParameter("path", fileFullPath)
+          .execute()
+          .get()
+          .getResponseBodyAsStream();
+
+      // Strip file path down to just a file name if we can
+      java.nio.file.Path filePath = Paths.get(fileFullPath).getFileName();
+      String fileName = filePath != null ? filePath.toString() : fileFullPath;
+
+      final String headerValue = String.format("attachment; filename=\"%s\"", fileName);
+      return Response.ok(responseStream).header("Content-Disposition", headerValue).build();
+
+    } catch (Exception e) {
+      if (e.getCause().getClass() == ConnectException.class) {
+        throw new SlaveNotFoundException(e);
+      } else {
+        throw new RuntimeException(e);
+      }
+    }
+
+  }
 }
diff --git a/SingularityUI/app/components/taskDetail/TaskDetail.jsx b/SingularityUI/app/components/taskDetail/TaskDetail.jsx
@@ -179,8 +179,7 @@ class TaskDetail extends Component {
         }
 
         file.fullPath = `${files.fullPathToRoot}/${files.currentDirectory}/${file.name}`;
-        file.downloadLink = `${httpPrefix}://${files.slaveHostname}:${httpPort}/files/download.json?path=${file.fullPath}`;
-
+        file.downloadLink = `${config.apiRoot}/tasks/download?slaveHostname=${files.slaveHostname}&path=${file.fullPath}`;
         file.isRecentlyModified = Date.now() / 1000 - file.mtime <= RECENTLY_MODIFIED_SECONDS;
 
         if (!file.isDirectory) {