Added custom header to create artifact end point.

varshavaradarajan · varshavaradarajan · commit 914974d96782 · 2016-08-03T09:44:16.000+05:30
diff --git a/server/src/com/thoughtworks/go/server/controller/ArtifactsController.java b/server/src/com/thoughtworks/go/server/controller/ArtifactsController.java
@@ -20,6 +20,7 @@
 import com.thoughtworks.go.domain.JobIdentifier;
 import com.thoughtworks.go.domain.exception.IllegalArtifactLocationException;
 import com.thoughtworks.go.server.cache.ZipArtifactCache;
+import com.thoughtworks.go.server.security.HeaderConstraint;
 import com.thoughtworks.go.server.service.ArtifactsService;
 import com.thoughtworks.go.server.service.ConsoleActivityMonitor;
 import com.thoughtworks.go.server.service.ConsoleService;
@@ -31,6 +32,7 @@
 import com.thoughtworks.go.server.web.FileModelAndView;
 import com.thoughtworks.go.server.web.ResponseCodeView;
 import com.thoughtworks.go.util.ArtifactLogUtil;
+import com.thoughtworks.go.util.SystemEnvironment;
 import org.apache.commons.io.IOUtils;
 import org.apache.log4j.Logger;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -67,10 +69,11 @@ public class ArtifactsController {
     private final ArtifactFolderViewFactory folderViewFactory;
     private final ArtifactFolderViewFactory jsonViewFactory;
     private final ArtifactFolderViewFactory zipViewFactory;
+    private HeaderConstraint headerConstraint;
 
     @Autowired
     ArtifactsController(ArtifactsService artifactsService, RestfulService restfulService, ZipArtifactCache zipArtifactCache,
-                        ConsoleActivityMonitor consoleActivityMonitor, ConsoleService consoleService) {
+                        ConsoleActivityMonitor consoleActivityMonitor, ConsoleService consoleService, SystemEnvironment systemEnvironment) {
         this.artifactsService = artifactsService;
         this.restfulService = restfulService;
         this.consoleActivityMonitor = consoleActivityMonitor;
@@ -79,6 +82,7 @@ public class ArtifactsController {
         this.folderViewFactory = FileModelAndView.htmlViewFactory();
         this.jsonViewFactory = FileModelAndView.jsonViewfactory();
         this.zipViewFactory = zipViewFactory(zipArtifactCache);
+        this.headerConstraint = new HeaderConstraint(systemEnvironment);
     }
 
 
@@ -135,6 +139,9 @@ public ModelAndView postArtifact(@RequestParam("pipelineName") String pipelineNa
                                      @RequestParam(value = "attempt", required = false) Integer attempt,
                                      MultipartHttpServletRequest request) throws Exception {
         JobIdentifier jobIdentifier;
+        if(!headerConstraint.isSatisfied(request)) {
+            return ResponseCodeView.create(HttpServletResponse.SC_BAD_REQUEST, "Missing required header 'Confirm'");
+        }
         try {
             jobIdentifier = restfulService.findJob(pipelineName, counterOrLabel, stageName, stageCounter,
                     buildName, buildId);
diff --git a/server/test/unit/com/thoughtworks/go/server/controller/ArtifactsControllerTest.java b/server/test/unit/com/thoughtworks/go/server/controller/ArtifactsControllerTest.java
@@ -24,11 +24,13 @@
 import com.thoughtworks.go.server.service.RestfulService;
 import com.thoughtworks.go.server.web.ArtifactFolderViewFactory;
 import com.thoughtworks.go.server.web.ResponseCodeView;
+import com.thoughtworks.go.util.SystemEnvironment;
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockMultipartFile;
 import org.springframework.mock.web.MockMultipartHttpServletRequest;
+import org.springframework.web.multipart.MultipartHttpServletRequest;
 import org.springframework.web.servlet.ModelAndView;
 
 import javax.servlet.http.HttpServletResponse;
@@ -52,6 +54,7 @@ public class ArtifactsControllerTest {
     private RestfulService restfulService;
     private ArtifactsService artifactService;
     private ConsoleService consoleService;
+    private SystemEnvironment systemEnvironment;
 
     @Before
     public void setUp() {
@@ -60,8 +63,8 @@ public void setUp() {
         restfulService = mock(RestfulService.class);
         artifactService = mock(ArtifactsService.class);
         consoleService = mock(ConsoleService.class);
-
-        artifactsController = new ArtifactsController(artifactService, restfulService, mock(ZipArtifactCache.class), consoleActivityMonitor, consoleService);
+        systemEnvironment = mock(SystemEnvironment.class);
+        artifactsController = new ArtifactsController(artifactService, restfulService, mock(ZipArtifactCache.class), consoleActivityMonitor, consoleService, systemEnvironment);
 
         request = new MockHttpServletRequest();
     }
@@ -102,7 +105,7 @@ public void shouldReturnHttpErrorCodeWhenChecksumFileSaveFails() throws Exceptio
     @Test
     public void shouldFunnelAll_GET_calls() throws Exception {
         final ModelAndView returnVal = new ModelAndView();
-        ArtifactsController controller = new ArtifactsController(artifactService, restfulService, mock(ZipArtifactCache.class), consoleActivityMonitor, consoleService) {
+        ArtifactsController controller = new ArtifactsController(artifactService, restfulService, mock(ZipArtifactCache.class), consoleActivityMonitor, consoleService, systemEnvironment) {
             @Override ModelAndView getArtifact(String filePath, ArtifactFolderViewFactory folderViewFactory, String pipelineName, String counterOrLabel, String stageName, String stageCounter,
                                                String buildName, String sha, String serverAlias) throws Exception {
                 return returnVal;
@@ -113,4 +116,17 @@ public void shouldFunnelAll_GET_calls() throws Exception {
         assertThat(controller.getArtifactAsZip("pipeline", "counter", "stage", "2", "job", "file_name", "sha1"), sameInstance(returnVal));
         assertThat(controller.getArtifactAsJson("pipeline", "counter", "stage", "2", "job", "file_name", "sha1"), sameInstance(returnVal));
     }
+
+    @Test
+    public void shouldReturnBadRequestIfRequiredHeadersAreMissingOnACreateArtifactRequest() throws Exception {
+        MultipartHttpServletRequest multipartHttpServletRequest = new MockMultipartHttpServletRequest();
+
+        when(systemEnvironment.isApiSafeModeEnabled()).thenReturn(true);
+        ModelAndView modelAndView = artifactsController.postArtifact("pipeline", "invalid-label", "stage", "stage-counter", "job-name", 3L, "file-path", 3, multipartHttpServletRequest);
+        ResponseCodeView codeView = (ResponseCodeView) modelAndView.getView();
+
+        assertThat(codeView.getStatusCode(), is(HttpServletResponse.SC_BAD_REQUEST));
+        assertThat(codeView.getContent(), is("Missing required header 'Confirm'"));
+
+    }
 }
