[LIVY-552][WIP] Add JWTFilter to validate incoming JWT tokens

pom.xml

@@ -158,6 +158,8 @@
     <maven.compiler.source>${java.version}</maven.compiler.source>
     <maven.compiler.target>${java.version}</maven.compiler.target>
 
+    <nimbus-jwt.version>4.41.2</nimbus-jwt.version>
+
   </properties>
 
   <repositories>

server/pom.xml

@@ -79,6 +79,12 @@
       <artifactId>guava</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>com.nimbusds</groupId>
+      <artifactId>nimbus-jose-jwt</artifactId>
+      <version>${nimbus-jwt.version}</version>
+    </dependency>
+
     <dependency>
       <groupId>io.dropwizard.metrics</groupId>
       <artifactId>metrics-core</artifactId>

server/src/main/scala/org/apache/livy/LivyConf.scala

@@ -168,6 +168,10 @@ object LivyConf {
   val THRIFT_DELEGATION_TOKEN_RENEW_INTERVAL =
     Entry("livy.server.thrift.delegation.token.renew-interval", "1d")
 
+  val JWT_FILTER_ENABLED = Entry("livy.jwt.filter.enabled", false)
+  val JWT_HEADER_NAME = Entry("livy.jwt.header.name", "Authentication")
+  val JWT_SIGNATURE_PUBLIC_KEY_PATH = Entry("livy.jwt.public-key.path", null)
+
   /**
    * Recovery mode of Livy. Possible values:
    * off: Default. Turn off recovery. Every time Livy shuts down, it stops and forgets all sessions.

server/src/main/scala/org/apache/livy/server/LivyServer.scala

@@ -37,6 +37,7 @@ import org.scalatra.servlet.{MultipartConfig, ServletApiImplicits}
 import org.apache.livy._
 import org.apache.livy.server.batch.BatchSessionServlet
 import org.apache.livy.server.interactive.InteractiveSessionServlet
+import org.apache.livy.server.jwt.JWTFilter
 import org.apache.livy.server.recovery.{SessionStore, StateStore}
 import org.apache.livy.server.ui.UIServlet
 import org.apache.livy.sessions.{BatchSessionManager, InteractiveSessionManager}
@@ -268,14 +269,17 @@ class LivyServer extends Logging {
 
     if (livyConf.getBoolean(CSRF_PROTECTION)) {
       info("CSRF protection is enabled.")
-      val csrfHolder = new FilterHolder(new CsrfFilter())
-      server.context.addFilter(csrfHolder, "/*", EnumSet.allOf(classOf[DispatcherType]))
+      addFilter(new CsrfFilter())
     }
 
     if (accessManager.isAccessControlOn) {
       info("Access control is enabled")
-      val accessHolder = new FilterHolder(new AccessFilter(accessManager))
-      server.context.addFilter(accessHolder, "/*", EnumSet.allOf(classOf[DispatcherType]))
+      addFilter(new AccessFilter(accessManager))
+    }
+
+    if (livyConf.getBoolean(JWT_FILTER_ENABLED)) {
+      info("JWT Authentication is enabled")
+      addFilter(JWTFilter(livyConf))
     }
 
     server.start()
@@ -377,6 +381,11 @@ class LivyServer extends Logging {
         "Session recovery requires YARN.")
     }
   }
+
+  private def addFilter(filter: Filter): Unit = {
+    val filterHolder = new FilterHolder(filter)
+    server.context.addFilter(filterHolder, "/*", EnumSet.allOf(classOf[DispatcherType]))
+  }
 }
 
 object LivyServer {

server/src/main/scala/org/apache/livy/server/jwt/JWSVerifierProvider.scala

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.livy.server.jwt
+
+import java.io.{ByteArrayInputStream, File, InputStream}
+import java.security.cert.{CertificateFactory, X509Certificate}
+import java.security.interfaces.RSAPublicKey
+
+import com.nimbusds.jose.JWSVerifier
+import com.nimbusds.jose.crypto.RSASSAVerifier
+import org.apache.commons.io.FileUtils
+
+import org.apache.livy.LivyConf
+import org.apache.livy.LivyConf.JWT_SIGNATURE_PUBLIC_KEY_PATH
+
+/**
+ * A provider of JWSVerifier instances based on the given LivyConf.
+ * @param livyConf Used to determine the public key path to use for the verifier
+ */
+class JWSVerifierProvider(livyConf: LivyConf) {
+
+  private[this] lazy val verifier = {
+    val publicKeyFile = livyConf.get(JWT_SIGNATURE_PUBLIC_KEY_PATH)
+    require(publicKeyFile != null,
+      s"${JWT_SIGNATURE_PUBLIC_KEY_PATH.key} must be set to verify JWT signatures.")
+    val publicKey = parseRSAPublicKey(publicKeyFile)
+    new RSASSAVerifier(publicKey)
+  }
+
+  def get(): JWSVerifier = verifier
+
+  private[this] def parseRSAPublicKey(publicKeyFile: String): RSAPublicKey = {
+    val is = new ByteArrayInputStream(FileUtils.readFileToByteArray(
+      new File(publicKeyFile)))
+    try {
+      parseRSAPublicKeyFromInputStream(is)
+    } finally {
+      is.close()
+    }
+  }
+
+  private[this] def parseRSAPublicKeyFromInputStream(inputStream: InputStream): RSAPublicKey = {
+    val factory = CertificateFactory.getInstance("X.509")
+    val certificate = factory.generateCertificate(inputStream).asInstanceOf[X509Certificate]
+    certificate.getPublicKey.asInstanceOf[RSAPublicKey]
+  }
+}

server/src/main/scala/org/apache/livy/server/jwt/JWTFilter.scala

@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.livy.server.jwt
+
+import javax.servlet.{Filter, FilterChain, FilterConfig, ServletRequest, ServletResponse}
+import javax.servlet.http.{HttpServletRequest, HttpServletResponse}
+
+import scala.util.control.NonFatal
+
+import org.apache.livy.{LivyConf, Logging}
+import org.apache.livy.LivyConf.JWT_HEADER_NAME
+
+/**
+ * A JWTFilter that utilizes a JWTValidator to determine if a JWT token included in a request
+ * is valid and accept/reject the request.
+ * @param jwtValidator validator used to determine the validity of the jwt token
+ */
+class JWTFilter(jwtValidator: JWTValidator, livyConf: LivyConf) extends Filter with Logging {
+
+  private val headerName = livyConf.get(JWT_HEADER_NAME)
+
+  override def init(filterConfig: FilterConfig): Unit = {}
+
+  override def destroy(): Unit = {}
+
+  override def doFilter(servletRequest: ServletRequest,
+                        servletResponse: ServletResponse,
+                        filterChain: FilterChain): Unit = {
+    val httpRequest = servletRequest.asInstanceOf[HttpServletRequest]
+    val httpServletResponse = servletResponse.asInstanceOf[HttpServletResponse]
+
+    val jwtToken = httpRequest.getHeader(headerName)
+    if (jwtToken == null) {
+      httpServletResponse.sendError(HttpServletResponse.SC_BAD_REQUEST,
+        s"Header: ${headerName} is missing in the request.")
+    } else {
+      if (validateToken(jwtToken)) {
+        filterChain.doFilter(httpRequest, servletResponse)
+      } else {
+        httpServletResponse.sendError(HttpServletResponse.SC_UNAUTHORIZED,
+          "JWT token included in request failed validation.")
+      }
+    }
+  }
+
+  private def validateToken(jwtToken: String): Boolean = {
+    try {
+      jwtValidator.isValid(jwtToken)
+    } catch {
+      case NonFatal(e) =>
+        warn("Exception while validating JWTToken", e)
+        false
+    }
+  }
+}
+
+object JWTFilter {
+
+  def apply(livyConf: LivyConf): JWTFilter = {
+    val jwsVerifier = new JWSVerifierProvider(livyConf).get()
+    val jwtValidator = new JWTValidator(jwsVerifier)
+    new JWTFilter(jwtValidator, livyConf)
+  }
+}

server/src/main/scala/org/apache/livy/server/jwt/JWTValidator.scala

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.livy.server.jwt
+
+import java.util.Date
+
+import com.nimbusds.jose.JWSVerifier
+import com.nimbusds.jwt.SignedJWT
+
+import org.apache.livy.Logging
+
+/**
+ * A validator for JSON Web Tokens (JWT).
+ * @param jwsVerifier JWSVerifier used to verify the JWT token's signature
+ */
+class JWTValidator(jwsVerifier: JWSVerifier) extends Logging {
+
+  def isValid(token: String): Boolean = {
+    val signedJwt = SignedJWT.parse(token)
+
+    if (!isSignatureValid(signedJwt)) {
+      warn("Signature of JWT token could not be verified.")
+      false
+    } else if (!isExpired(signedJwt)) {
+      warn("Expiration time validation of JWT token failed.")
+      false
+    } else true
+  }
+
+  private def isSignatureValid(signedJWT: SignedJWT): Boolean = {
+    signedJWT.getSignature != null && signedJWT.verify(jwsVerifier)
+  }
+
+  private def isExpired(signedJWT: SignedJWT): Boolean = {
+    val expires = signedJWT.getJWTClaimsSet.getExpirationTime
+    expires == null || new Date().before(expires)
+  }
+}

server/src/test/scala/org/apache/livy/server/jwt/BaseJWTSpec.scala

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.livy.server.jwt
+
+import java.security.interfaces.{RSAPrivateKey, RSAPublicKey}
+import java.util.Date
+
+import com.nimbusds.jose.{JWSAlgorithm, JWSHeader}
+import com.nimbusds.jose.crypto.RSASSASigner
+import com.nimbusds.jwt.{JWTClaimsSet, SignedJWT}
+import org.joda.time.DateTime
+import sun.security.tools.keytool.CertAndKeyGen
+import sun.security.x509.X500Name
+
+trait BaseJWTSpec {
+
+  private final val BIT_STRENGTH = 512
+  private final val EXPIRATION = 60 * 60 * 24 // One Day
+  lazy val generator = createGenerator()
+  lazy val cert = generator.getSelfCertificate(
+    new X500Name("CN=LivyServer,O=Apache,L=ORG,C=DE"), EXPIRATION)
+  lazy val privateKey = generator.getPrivateKey.asInstanceOf[RSAPrivateKey]
+  lazy val publicKey = generator.getPublicKeyAnyway.asInstanceOf[RSAPublicKey]
+  lazy val signer = new RSASSASigner(privateKey)
+  lazy val unexpiredToken = createToken(new DateTime().plusDays(1), sign = true)
+  lazy val expiredToken = createToken(new DateTime().minusDays(1), sign = true)
+  lazy val unsignedToken = createToken(new DateTime().plusDays(1), sign = false)
+  lazy val unexpiredTokenWithNoExpiration = createToken(null, sign = true)
+
+  def createNewPublicKey(): RSAPublicKey = {
+    val newGenerator = createGenerator()
+    newGenerator.getPublicKeyAnyway.asInstanceOf[RSAPublicKey]
+  }
+
+  def createClaimsSet(issuer: String, expirationTime: Date): JWTClaimsSet = {
+    val builder = new JWTClaimsSet.Builder().issuer(issuer)
+    if (expirationTime != null) {
+      builder.expirationTime(expirationTime)
+    }
+    builder.build()
+  }
+
+  def createToken(expirationTime: DateTime, sign: Boolean): SignedJWT = {
+    val expiration = if (expirationTime == null) null else expirationTime.toDate
+    val signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.RS256),
+      createClaimsSet("jwt-test", expiration))
+    if (sign) {
+      signedJWT.sign(signer)
+    }
+    signedJWT
+  }
+
+  private def createGenerator(): CertAndKeyGen = {
+    val newGenerator = new CertAndKeyGen("RSA", "SHA256WithRSA", null)
+    newGenerator.generate(BIT_STRENGTH)
+    newGenerator
+  }
+}