Leverage HTTP header Host to generate a more usable config (#299)

center/src/main/java/com/microsoft/hydralab/center/controller/AgentManageController.java

@@ -14,6 +14,7 @@
 import com.microsoft.hydralab.common.util.AttachmentService;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Value;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.security.access.prepost.PreAuthorize;
@@ -227,40 +228,30 @@ public Result<AgentUser> getAgentInfo(@CurrentSecurityContext SysUser requestor,
     @GetMapping("/api/agent/downloadAgentConfigFile/{agentId}")
     public Result downloadAgentConfigFile(@CurrentSecurityContext SysUser requestor,
                                           @PathVariable(value = "agentId") String agentId,
+                                          @RequestHeader(HttpHeaders.HOST) String host,
                                           HttpServletResponse response) throws IOException {
         if (!agentManageService.checkAgentAuthorization(requestor, agentId)) {
             return Result.error(HttpStatus.UNAUTHORIZED.value(), "Authentication failed");
         }
-        File agentConfigFile = agentManageService.generateAgentConfigFile(agentId);
+        File agentConfigFile = agentManageService.generateAgentConfigFile(agentId, host);
         if (agentConfigFile == null) {
             return Result.error(HttpStatus.BAD_REQUEST.value(), "The file was not downloaded");
         }
 
-        ServletOutputStream out = null;
-        FileInputStream in = null;
-        try {
-            in = new FileInputStream(agentConfigFile);
+        try (FileInputStream in = new FileInputStream(agentConfigFile);
+             ServletOutputStream out = response.getOutputStream()) {
+
             response.setContentType("application/octet-stream;charset=UTF-8");
-            response.setHeader("Content-Disposition","attachment;filename=" + agentConfigFile.getName());
-            out = response.getOutputStream();
-            int len = 0;
+            response.setHeader("Content-Disposition", "attachment;filename=" + agentConfigFile.getName());
+            int len;
             byte[] buffer = new byte[1024 * 10];
             while ((len = in.read(buffer)) != -1) {
                 out.write(buffer, 0, len);
             }
             out.flush();
-        } catch (Exception e) {
-            e.printStackTrace();
-            return Result.error(HttpStatus.INTERNAL_SERVER_ERROR.value(), "Internal server error");
         } finally {
             response.flushBuffer();
-            try {
-                out.close();
-                in.close();
-                agentConfigFile.delete();
-            } catch (Exception e) {
-                e.printStackTrace();
-            }
+            agentConfigFile.delete();
         }
 
         return Result.ok();

center/src/main/java/com/microsoft/hydralab/center/service/AgentManageService.java

@@ -132,12 +132,12 @@ public void updateAgentTeam(String teamId, String teamName){
          agentUserRepository.saveAll(agents);
     }
 
-    public File generateAgentConfigFile(String agentId) {
+    public File generateAgentConfigFile(String agentId, String host) {
         AgentUser agentUser = getAgent(agentId);
         if (agentUser != null) {
             try {
                 File tempFolder = new File(CenterConstant.CENTER_TEMP_FILE_DIR);
-                if (!tempFolder.exists()){
+                if (!tempFolder.exists()) {
                     if (!tempFolder.mkdirs()) {
                         throw new RuntimeException("mkdirs fail for: " + tempFolder);
                     }
@@ -152,7 +152,7 @@ public File generateAgentConfigFile(String agentId) {
                         "  # register to Hydra Lab Center\n" +
                         "  registry:\n" +
                         "    # The server hostname:port of Hydra Lab Center. If nginx enabled, switch to port of nginx\n" +
-                        "    server: 'localhost:9886'\n" +
+                        "    server: '" + host + "'\n" +
                         "    # The Agent info registered in Hydra Lab Center, for instance if it's running on localhost, the URL would be: http://localhost:9886/portal/index.html#/auth\n" +
                         "    name: " + agentUser.getName() + "\n" +
                         "    id: " + agentUser.getId() + "\n" +

center/src/test/java/com/microsoft/hydralab/center/service/AgentManageServiceTest.java

@@ -1,28 +1,35 @@
 package com.microsoft.hydralab.center.service;
 
 import com.microsoft.hydralab.common.entity.center.AgentUser;
+import org.apache.commons.io.IOUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 
 import java.io.File;
+import java.io.FileReader;
+import java.io.IOException;
 
 import static org.mockito.Mockito.doReturn;
 
 public class AgentManageServiceTest {
 
     @Test
-    void generateAgentConfig() {
+    void generateAgentConfig() throws IOException {
         String agentId = "test_agent_id";
+        String host = "localhost:9886";
+
         AgentManageService agentManageService = Mockito.spy(AgentManageService.class);
         AgentUser agent = new AgentUser();
         agent.setName("Agent Name");
         agent.setId(agentId);
         agent.setSecret("Agent Secret");
         doReturn(agent).when(agentManageService).getAgent(agentId);
 
-        File file = agentManageService.generateAgentConfigFile(agentId);
+        File file = agentManageService.generateAgentConfigFile(agentId, host);
         Assertions.assertNotNull(file, "Get agent user error");
         Assertions.assertTrue(file.length() > 0, "Write agent config file error");
+
+        Assertions.assertTrue(IOUtils.toString(new FileReader(file)).contains(host));
     }
 }