save some basic code

sql/thriftserver/src/main/scala/org/apache/spark/sql/thriftserver/AbstractService.scala

@@ -0,0 +1,137 @@
+/*
+ * 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.spark.sql.thriftserver
+
+import java.util
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.hive.conf.HiveConf
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.thriftserver.Service._
+
+/**
+ * Construct the service.
+ *
+ * @param name
+ * service name
+ */
+abstract class AbstractService(val name: String) extends Service with Logging {
+
+  /**
+   * Service state: initially {@link STATE#NOTINITED}.
+   */
+  private var state: STATE = NOTINITED
+  /**
+   * Service start time. Will be zero until the service is started.
+   */
+  private var startTime = 0L
+  /**
+   * The configuration. Will be null until the service is initialized.
+   */
+  private var hiveConf: HiveConf = null
+  /**
+   * List of state change listeners; it is final to ensure
+   * that it will never be null.
+   */
+  final private val listeners = new util.ArrayList[ServiceStateChangeListener]
+
+  def getServiceState: Service.STATE = synchronized {
+    state
+  }
+
+  /**
+   * {@inheritDoc }
+   *
+   * @throws IllegalStateException
+   * if the current service state does not permit
+   * this action
+   */
+  def init(hiveConf: HiveConf): Unit = {
+    ensureCurrentState(NOTINITED)
+    this.hiveConf = hiveConf
+    changeState(INITED)
+    logInfo("Service:" + getName + " is inited.")
+  }
+
+  def start(): Unit = {
+    startTime = System.currentTimeMillis
+    ensureCurrentState(INITED)
+    changeState(STARTED)
+    logInfo("Service:" + getName + " is started.")
+  }
+
+  def stop(): Unit = {
+    if ((state eq STOPPED) ||
+      (state eq INITED) ||
+      (state eq NOTINITED)) {
+      // already stopped, or else it was never
+      // started (eg another service failing canceled startup)
+      return
+    }
+    ensureCurrentState(STARTED)
+    changeState(STOPPED)
+    logInfo("Service:" + getName + " is stopped.")
+  }
+
+  def register(l: ServiceStateChangeListener): Unit = {
+    listeners.add(l)
+  }
+
+  def unregister(l: ServiceStateChangeListener): Unit = {
+    listeners.remove(l)
+  }
+
+  def getName: String = name
+
+  def getHiveConf: HiveConf = hiveConf
+
+  def getStartTime: Long = startTime
+
+  /**
+   * Verify that a service is in a given state.
+   *
+   * @param currentState
+   * the desired state
+   * @throws IllegalStateException
+   * if the service state is different from
+   * the desired state
+   */
+  private def ensureCurrentState(currentState: STATE): Unit = {
+    ServiceOperations.ensureCurrentState(state, currentState)
+  }
+
+  /**
+   * Change to a new state and notify all listeners.
+   * This is a private method that is only invoked from synchronized methods,
+   * which avoid having to clone the listener list. It does imply that
+   * the state change listener methods should be short lived, as they
+   * will delay the state transition.
+   *
+   * @param newState
+   * new service state
+   */
+  private def changeState(newState: STATE): Unit = {
+    state = newState
+    // notify listeners
+    for (l <- listeners.asScala) {
+      l.stateChanged(this)
+    }
+  }
+}

sql/thriftserver/src/main/scala/org/apache/spark/sql/thriftserver/CompositeServices.scala

@@ -0,0 +1,113 @@
+/*
+ * 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.spark.sql.thriftserver
+
+import java.util
+import java.util.Collections
+
+import scala.collection.JavaConverters._
+
+import org.apache.hadoop.hive.conf.HiveConf
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.thriftserver.Service.STOPPED
+
+class CompositeService(name: String) extends AbstractService(name) with Logging {
+
+
+  private val serviceList: util.ArrayList[Service] = new util.ArrayList[Service]
+
+  def getServices: util.Collection[Service] = Collections.unmodifiableList(serviceList)
+
+  protected def addService(service: Service): Unit = {
+    serviceList.add(service)
+  }
+
+  protected def removeService(service: Service): Boolean = serviceList.remove(service)
+
+  override def init(hiveConf: HiveConf): Unit = {
+    for (service <- serviceList.asScala) {
+      service.init(hiveConf)
+    }
+    super.init(hiveConf)
+  }
+
+  override def start(): Unit = {
+    var i = 0
+    try {
+      val n = serviceList.size
+      while (i < n) {
+        val service = serviceList.get(i)
+        service.start
+        i += 1
+      }
+      super.start
+    } catch {
+      case e: Throwable =>
+        logError("Error starting services " + getName, e)
+        // Note that the state of the failed service is still INITED and not
+        // STARTED. Even though the last service is not started completely, still
+        // call stop() on all services including failed service to make sure cleanup
+        // happens.
+        stop(i)
+        throw new ServiceException("Failed to Start " + getName, e)
+    }
+  }
+
+  override def stop(): Unit = {
+    if (this.getServiceState eq STOPPED) {
+      // The base composite-service is already stopped, don't do anything again.
+      return
+    }
+    if (serviceList.size > 0) {
+      stop(serviceList.size - 1)
+    }
+    super.stop
+  }
+
+  private def stop(numOfServicesStarted: Int): Unit = {
+    // stop in reserve order of start
+    var i = numOfServicesStarted
+    while (i >= 0) {
+      val service = serviceList.get(i)
+      try {
+        service.stop
+      } catch {
+        case t: Throwable =>
+          logInfo("Error stopping " + service.getName, t)
+      }
+      i -= 1
+    }
+  }
+
+  /**
+   * JVM Shutdown hook for CompositeService which will stop the given
+   * CompositeService gracefully in case of JVM shutdown.
+   */
+  class CompositeServiceShutdownHook(val compositeService: CompositeService) extends Runnable {
+    override def run(): Unit = {
+      try // Stop the Composite Service
+      compositeService.stop
+      catch {
+        case t: Throwable =>
+          logInfo("Error stopping " + compositeService.getName, t)
+      }
+    }
+  }
+
+}

sql/thriftserver/src/main/scala/org/apache/spark/sql/thriftserver/CookieSigner.scala

@@ -0,0 +1,102 @@
+/*
+ * 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.spark.sql.thriftserver
+
+import java.security.{MessageDigest, NoSuchAlgorithmException}
+
+import org.apache.commons.codec.binary.Base64
+
+import org.apache.spark.internal.Logging
+
+/**
+ * The cookie signer generates a signature based on SHA digest
+ * and appends it to the cookie value generated at the
+ * server side. It uses SHA digest algorithm to sign and verify signatures.
+ */
+object CookieSigner {
+  private val SIGNATURE = "&s="
+  private val SHA_STRING = "SHA"
+}
+
+class CookieSigner extends Logging {
+  private var secretBytes: Array[Byte] = null
+
+  def this(secret: Array[Byte]) {
+    this()
+    if (secret == null) {
+      throw new IllegalArgumentException("NULL Secret Bytes")
+    }
+    this.secretBytes = secret.clone
+  }
+
+
+  /**
+   * Sign the cookie given the string token as input.
+   *
+   * @param str Input token
+   * @return Signed token that can be used to create a cookie
+   */
+  def signCookie(str: String): String = {
+    if (str == null || str.isEmpty) {
+      throw new IllegalArgumentException("NULL or empty string to sign")
+    }
+    val signature = getSignature(str)
+    logDebug("Signature generated for " + str + " is " + signature)
+    str + CookieSigner.SIGNATURE + signature
+  }
+
+  /**
+   * Verify a signed string and extracts the original string.
+   *
+   * @param signedStr The already signed string
+   * @return Raw Value of the string without the signature
+   */
+  def verifyAndExtract(signedStr: String): String = {
+    val index = signedStr.lastIndexOf(CookieSigner.SIGNATURE)
+    if (index == -1) {
+      throw new IllegalArgumentException("Invalid input sign: " + signedStr)
+    }
+    val originalSignature = signedStr.substring(index + CookieSigner.SIGNATURE.length)
+    val rawValue = signedStr.substring(0, index)
+    val currentSignature = getSignature(rawValue)
+    logDebug("Signature generated for " + rawValue + " inside verify is " + currentSignature)
+    if (!(originalSignature == currentSignature)) {
+      throw new IllegalArgumentException("Invalid sign, original = " +
+        originalSignature + " current = " + currentSignature)
+    }
+    rawValue
+  }
+
+  /**
+   * Get the signature of the input string based on SHA digest algorithm.
+   *
+   * @param str Input token
+   * @return Signed String
+   */
+  private def getSignature(str: String) = try {
+    val md = MessageDigest.getInstance(CookieSigner.SHA_STRING)
+    md.update(str.getBytes)
+    md.update(secretBytes)
+    val digest = md.digest
+    new Base64(0).encodeToString(digest)
+  } catch {
+    case ex: NoSuchAlgorithmException =>
+      throw new RuntimeException("Invalid SHA digest String: " +
+        CookieSigner.SHA_STRING + " " + ex.getMessage, ex)
+  }
+}