Fix: stop FileHandler after error responses

chaoscenter/graphql/server/pkg/handlers/file_handler.go

@@ -23,46 +23,51 @@ func FileHandler(mongodbOperator mongodb.MongoOperator) gin.HandlerFunc {
 		infraId, err := chaos_infrastructure.InfraValidateJWT(token)
 		if err != nil {
 			logrus.Error(err)
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte(err.Error()))
+			return
 		}
 
 		infra, err := dbChaosInfra.NewInfrastructureOperator(mongodbOperator).GetInfra(infraId)
 		if err != nil {
 			logrus.Error(err)
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte(err.Error()))
+			return
 		}
 
 		reqHeader, ok := c.Value("request-header").(http.Header)
 		if !ok {
 			logrus.Error("unable to parse referer header")
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte("unable to parse referer header"))
+			return
 		}
-
 		referrer := reqHeader.Get("Referer")
 		if referrer == "" {
 			logrus.Error("unable to parse referer header")
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte("unable to parse referer header"))
+			return
 		}
 
 		referrerURL, err := url.Parse(referrer)
 		if err != nil {
 			logrus.Error(err)
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte(err.Error()))
+			return
 		}
 
 		response, err := chaos_infrastructure.GetK8sInfraYaml(fmt.Sprintf("%s://%s", referrerURL.Scheme, referrerURL.Host), infra)
 		if err != nil {
 			logrus.Error(err)
-			utils.WriteHeaders(&c.Writer, 500)
+			utils.WriteHeaders(&c.Writer, http.StatusInternalServerError)
 			c.Writer.Write([]byte(err.Error()))
+			return
 		}
 
-		utils.WriteHeaders(&c.Writer, 200)
+		utils.WriteHeaders(&c.Writer, http.StatusOK)
 		c.Writer.Write(response)
 	}
 }

chaoscenter/graphql/server/pkg/handlers/handlers_test.go

@@ -66,3 +66,20 @@ func TestReadinessHandler_DatabaseUpWithoutLitmus(t *testing.T) {
 	assert.Contains(t, w.Body.String(), `"database":"down"`)
 	mockOp.AssertExpectations(t)
 }
+
+func TestFileHandlerInvalidTokenReturnsAfterError(t *testing.T) {
+	mockOp := new(mocks.MongoOperator)
+
+	w := httptest.NewRecorder()
+	ctx := GetTestGinContext(w)
+	ctx.Params = []gin.Param{{Key: "key", Value: "not-a-jwt.yaml"}}
+
+	handlers.FileHandler(mockOp)(ctx)
+
+	assert.Equal(t, http.StatusInternalServerError, w.Code)
+	assert.Contains(t, w.Body.String(), "token contains an invalid number of segments")
+	assert.NotContains(t, w.Body.String(), "mongo: no documents in result")
+	assert.NotContains(t, w.Body.String(), "unable to parse referer header")
+	mockOp.AssertNotCalled(t, "Get", mock.Anything, mock.Anything, mock.Anything)
+	mockOp.AssertExpectations(t)
+}