set http response code on error, some testing

Author: squito

core/src/main/scala/org/apache/spark/status/StatusJsonHandler.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.status
 
-import javax.servlet.http.HttpServletRequest
+import javax.servlet.http.{HttpServletResponse, HttpServlet, HttpServletRequest}
 
 import com.fasterxml.jackson.annotation.JsonInclude
 import com.fasterxml.jackson.databind.{SerializationFeature, ObjectMapper}
@@ -31,8 +31,7 @@ import org.eclipse.jetty.servlet.{ServletHolder, ServletContextHandler}
 import scala.util.matching.Regex
 
 import org.apache.spark.{Logging, SecurityManager}
-import org.apache.spark.deploy.history.{OneApplicationJsonRoute, HistoryServer, AllApplicationsJsonRoute}
-import org.apache.spark.ui.JettyUtils._
+import org.apache.spark.deploy.history.{OneApplicationJsonRoute, AllApplicationsJsonRoute}
 
 
 /**
@@ -60,9 +59,9 @@ private[spark] class JsonRequestHandler(uiRoot: UIRoot, securityManager: Securit
     s"/applications/$noSlash+/jobs/?".r -> new AllJobsJsonRoute(this),
     s"/applications/$noSlash+/executors/?".r -> new ExecutorsJsonRoute(this),
     s"/applications/$noSlash+/stages/?".r -> new AllStagesJsonRoute(this),
-    s"/applications/$noSlash+/stages/\\d+/?".r -> new OneStageJsonRoute(this),
+    s"/applications/$noSlash+/stages/$noSlash+/?".r -> new OneStageJsonRoute(this),
     s"/applications/$noSlash+/storage/rdd/?".r -> new AllRDDJsonRoute(this),
-    s"/applications/$noSlash+/storage/rdd/\\d+/?".r -> new RDDJsonRoute(this)
+    s"/applications/$noSlash+/storage/rdd/$noSlash+/?".r -> new RDDJsonRoute(this)
   )
 
   private val jsonMapper = {
@@ -76,17 +75,36 @@ private[spark] class JsonRequestHandler(uiRoot: UIRoot, securityManager: Securit
   val jsonContextHandler = {
 
     //TODO throw out all the JettyUtils stuff, so I can set the response status code, etc.
-    val params = new ServletParams(
-    {
-      (request: HttpServletRequest) =>
-        route(request).map{jsonRoute =>
-          logInfo("handling route: " + request.getPathInfo)
-          val responseObj = jsonRoute.renderJson(request)
-          jsonMapper.writeValueAsString(responseObj)
-        }.getOrElse(throw new IllegalArgumentException("unmatched route"))
-    }, "text/json")
+    val servlet = new HttpServlet {
+      override def doGet(request: HttpServletRequest, response: HttpServletResponse) {
+        if (securityManager.checkUIViewPermissions(request.getRemoteUser)) {
+          response.setContentType("text/json;charset=utf-8")
+          route(request) match {
+            case Some(jsonRoute) =>
+              response.setHeader("Cache-Control", "no-cache, no-store, must-revalidate")
+              try {
+                val responseObj = jsonRoute.renderJson(request)
+                val result = jsonMapper.writeValueAsString(responseObj)
+                response.setStatus(HttpServletResponse.SC_OK)
+                response.getWriter.println(result)
+              } catch {
+                case iae: IllegalArgumentException =>
+                  response.setStatus(HttpServletResponse.SC_BAD_REQUEST)
+                  response.getOutputStream.print(iae.getMessage())
+              }
+            case None =>
+              println("no match for path: " + request.getPathInfo)
+              response.setStatus(HttpServletResponse.SC_NOT_FOUND)
+          }
+        } else {
+          response.setStatus(HttpServletResponse.SC_UNAUTHORIZED)
+          response.setHeader("Cache-Control", "no-cache, no-store, must-revalidate")
+          response.sendError(HttpServletResponse.SC_UNAUTHORIZED,
+            "User is not authorized to access this page.")
+        }
+      }
+    }
     val path = "/json/v1"
-    val servlet = createServlet(params, securityManager)
     val contextHandler = new ServletContextHandler
     val holder = new ServletHolder(servlet)
     contextHandler.setContextPath(path)
@@ -148,4 +166,3 @@ object RouteUtils {
     }
   }
 }
-

core/src/main/scala/org/apache/spark/ui/jobs/OneStageJsonRoute.scala

@@ -45,7 +45,7 @@ class OneStageJsonRoute(parent: JsonRequestHandler) extends StatusJsonRoute[Stag
               }
               AllStagesJsonRoute.stageUiToStageData(status, stageInfo, stageUiData, includeDetails = true)
             case None =>
-              throw new IllegalArgumentException("no stage found with id:" + stageId)
+              throw new IllegalArgumentException("unknown stage: " + stageId)
           }
 
         case None =>

core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala

@@ -16,8 +16,9 @@
  */
 package org.apache.spark.deploy.history
 
-import java.io.{FileInputStream, FileWriter, File}
-import java.net.URL
+import java.io.{IOException, FileInputStream, FileWriter, File}
+import java.net.{HttpURLConnection, URL}
+import javax.servlet.http.HttpServletResponse
 
 import org.apache.commons.io.{FileUtils, IOUtils}
 
@@ -72,7 +73,11 @@ class HistoryServerSuite extends FunSuite with BeforeAndAfter with Matchers {
   // resource folder
   cases.foreach{case(name, path) =>
       test(name){
-        val json = getUrl(path)
+        val (code, jsonOpt, errOpt) = getContentAndCode(path)
+        code should be (HttpServletResponse.SC_OK)
+        jsonOpt should be ('defined)
+        errOpt should be (None)
+        val json = jsonOpt.get
         val exp = IOUtils.toString(new FileInputStream(new File(expRoot, path + "/json_expectation")))
         json should be (exp)
       }
@@ -83,7 +88,42 @@ class HistoryServerSuite extends FunSuite with BeforeAndAfter with Matchers {
   }
 
   test("response codes on bad paths") {
-    pending
+    val badAppId = getContentAndCode("applications/foobar")
+    badAppId._1 should be (HttpServletResponse.SC_BAD_REQUEST)
+    badAppId._3 should be (Some("unknown app: foobar"))
+
+    val badStageId = getContentAndCode("applications/local-1422981780767/stages/12345")
+    badStageId._1 should be (HttpServletResponse.SC_BAD_REQUEST)
+    badStageId._3 should be (Some("unknown stage: 12345"))
+
+    val badStageId2 = getContentAndCode("applications/local-1422981780767/stages/flimflam")
+    badStageId2._1 should be (HttpServletResponse.SC_BAD_REQUEST)
+    badStageId2._3 should be (Some("no valid stageId in path"))
+
+
+    getContentAndCode("foobar")._1 should be (HttpServletResponse.SC_NOT_FOUND)
+
+  }
+
+  def getContentAndCode(path: String): (Int, Option[String], Option[String]) = {
+    val url = new URL(s"http://localhost:$port/json/v1/$path")
+    val connection = url.openConnection().asInstanceOf[HttpURLConnection]
+    connection.setRequestMethod("GET")
+    connection.connect()
+    val code = connection.getResponseCode()
+    val inString = try {
+      val in = Option(connection.getInputStream())
+      in.map{IOUtils.toString}
+    } catch {
+      case io: IOException => None
+    }
+    val errString = try {
+      val err = Option(connection.getErrorStream())
+      err.map{IOUtils.toString}
+    } catch {
+      case io: IOException => None
+    }
+    (code, inString, errString)
   }