org.apache.thrift.transport.TSocket Scala Examples

package yaooqinn.kyuubi.auth

import org.apache.hive.service.cli.thrift.TCLIService.{Iface, Processor}
import org.apache.kyuubi.Logging
import org.apache.thrift.TException
import org.apache.thrift.protocol.TProtocol
import org.apache.thrift.transport.{TSaslClientTransport, TSaslServerTransport, TSocket, TTransport}


class TSetIpAddressProcessor[I <: Iface](iface: Iface)
  extends Processor[Iface](iface) with Logging {

  import TSetIpAddressProcessor._

  @throws[TException]
  override def process(in: TProtocol, out: TProtocol): Boolean = {
    setIpAddress(in)
    setUserName(in)
    try {
      super.process(in, out)
    } finally {
      THREAD_LOCAL_USER_NAME.remove()
      THREAD_LOCAL_IP_ADDRESS.remove()
    }
  }

  private def setUserName(in: TProtocol): Unit = {
    val transport = in.getTransport
    transport match {
      case transport1: TSaslServerTransport =>
        val userName = transport1.getSaslServer.getAuthorizationID
        THREAD_LOCAL_USER_NAME.set(userName)
      case _ =>
    }
  }

  private def setIpAddress(in: TProtocol): Unit = {
    val transport = in.getTransport
    val tSocket = getUnderlyingSocketFromTransport(transport)
    if (tSocket == null) {
      warn("Unknown Transport, cannot determine ipAddress")
    } else {
      THREAD_LOCAL_IP_ADDRESS.set(tSocket.getSocket.getInetAddress.getHostAddress)
    }
  }

  private def getUnderlyingSocketFromTransport(transport: TTransport): TSocket = transport match {
    case t: TSaslServerTransport => getUnderlyingSocketFromTransport(t.getUnderlyingTransport)
    case t: TSaslClientTransport => getUnderlyingSocketFromTransport(t.getUnderlyingTransport)
    case t: TSocket => t
    case _ => null
  }
}

object TSetIpAddressProcessor {

  private val THREAD_LOCAL_IP_ADDRESS = new ThreadLocal[String]() {
    override protected def initialValue: String = null
  }
  private val THREAD_LOCAL_USER_NAME = new ThreadLocal[String]() {
    override protected def initialValue: String = null
  }

  def getUserIpAddress: String = THREAD_LOCAL_IP_ADDRESS.get

  def getUserName: String = THREAD_LOCAL_USER_NAME.get
} 

package yaooqinn.kyuubi.auth

import java.security.Security
import javax.security.auth.login.LoginException

import org.apache.spark.{KyuubiSparkUtil, SparkConf, SparkFunSuite}
import org.apache.thrift.transport.{TSaslServerTransport, TSocket}

import yaooqinn.kyuubi.auth.PlainSaslServer.SaslPlainProvider
import yaooqinn.kyuubi.server.KyuubiServer

class PlainSaslHelperSuite extends SparkFunSuite {

  test("Plain Sasl Helper") {
    val conf = new SparkConf(loadDefaults = true)
    KyuubiSparkUtil.setupCommonConfig(conf)
    val server = new KyuubiServer()
    val fe = server.feService
    val tProcessorFactory = PlainSaslHelper.getProcessFactory(fe)
    assert(!tProcessorFactory.isAsyncProcessor)
    val tSocket = new TSocket("0.0.0.0", 0)
    val tProcessor = tProcessorFactory.getProcessor(tSocket)
    assert(tProcessor.isInstanceOf[TSetIpAddressProcessor[_]])
    intercept[LoginException](PlainSaslHelper.getTransportFactory("KERBEROS", conf))
    intercept[LoginException](PlainSaslHelper.getTransportFactory("NOSASL", conf))
    intercept[LoginException](PlainSaslHelper.getTransportFactory("ELSE", conf))
    val tTransportFactory = PlainSaslHelper.getTransportFactory("NONE", conf)
    assert(tTransportFactory.isInstanceOf[TSaslServerTransport.Factory])
    Security.getProviders.exists(_.isInstanceOf[SaslPlainProvider])
  }

  test("Sasl Plain Provider") {
    val saslPlainProvider = new SaslPlainProvider()
    assert(saslPlainProvider.containsKey("SaslServerFactory.PLAIN"))
    assert(saslPlainProvider.getName === "KyuubiSaslPlain")

  }


} 

package unicorn.bigtable.cassandra

import java.util.Properties
import java.net.{InetAddress, UnknownHostException}
import scala.collection.JavaConversions._
import org.apache.cassandra.locator.SimpleSnitch
import org.apache.cassandra.thrift.Cassandra.Client
import org.apache.cassandra.thrift.{ConsistencyLevel, KsDef, CfDef}
import org.apache.thrift.transport.TFramedTransport
import org.apache.thrift.transport.TSocket
import org.apache.thrift.protocol.TBinaryProtocol
import unicorn.bigtable._
import unicorn.util.Logging


  override def compactTable(name: String): Unit = {
    // fail silently
    log.warn("Cassandra client API doesn't support compaction")
  }
}

object Cassandra {
  def apply(host: String, port: Int): Cassandra = {
    // For ultra-wide row, we set the maxLength to 16MB.
    // Note that we also need to set the server side configuration
    // thrift_framed_transport_size_in_mb in cassandra.yaml
    // In case of ultra-wide row, it is better to use intra row scan.
    val transport = new TFramedTransport(new TSocket(host, port), 16 * 1024 * 1024)
    transport.open

    new Cassandra(transport)
  }
}