1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.hbase.data
10 
11 import com.github.benmanes.caffeine.cache.{CacheLoader, Caffeine, LoadingCache}
12 import com.typesafe.scalalogging.LazyLogging
13 import org.apache.hadoop.conf.Configuration
14 import org.apache.hadoop.hbase.client.{Connection, ConnectionFactory}
15 import org.apache.hadoop.hbase.security.User
16 import org.apache.hadoop.hbase.security.token.AuthenticationTokenIdentifier
17 import org.apache.hadoop.hbase.{HBaseConfiguration, HConstants}
18 import org.apache.hadoop.security.authentication.util.KerberosUtil
19 import org.apache.hadoop.security.{SecurityUtil, UserGroupInformation}
20 import org.locationtech.geomesa.hbase.data.HBaseDataStoreFactory.{HBaseGeoMesaKeyTab, HBaseGeoMesaPrincipal}
21 import org.locationtech.geomesa.hbase.data.HBaseDataStoreParams.{ConfigPathsParam, ConfigsParam, ConnectionParam, ZookeeperParam}
22 import org.locationtech.geomesa.hbase.utils.HBaseVersions
23 import org.locationtech.geomesa.utils.hadoop.HadoopUtils
24 import org.locationtech.geomesa.utils.io.CloseWithLogging
25 
26 import java.io.{ByteArrayInputStream, Closeable}
27 import java.nio.charset.StandardCharsets
28 import java.security.PrivilegedExceptionAction
29 import scala.util.{Failure, Success, Try}
30 
31 object HBaseConnectionPool extends LazyLogging {
32 
33   import scala.collection.JavaConverters._
34 
35   private val configs: LoadingCache[ConfigKey, Configuration] = Caffeine.newBuilder().build(
36     new CacheLoader[ConfigKey, Configuration] {
37 
38       // add common resources from system property - lazy to allow object initialization if there's an error
39       private lazy val configuration = {
40         val base = HBaseConfiguration.create()
41         HBaseDataStoreFactory.ConfigPathProperty.option.foreach(addResources(base, _))
42         base
43       }
44 
45       override def load(key: ConfigKey): Configuration = {
46         val conf = new Configuration(configuration)
47         // Make sure that current user is always logged-in user
48         conf.set("hbase.client.userprovider.class", "org.locationtech.geomesa.hbase.data.LoginUserProvider")
49         // add the explicit props first, they may be needed for loading the path resources
50         key.xml.foreach(xml => conf.addResource(new ByteArrayInputStream(xml.getBytes(StandardCharsets.UTF_8))))
51         key.paths.foreach(addResources(conf, _))
52         key.zookeepers.foreach(zk => conf.set(HConstants.ZOOKEEPER_QUORUM, zk))
53         if (key.zookeepers.isEmpty && conf.get(HConstants.ZOOKEEPER_QUORUM) == "localhost") {
54           logger.warn("HBase connection is set to localhost - " +
55               "this may indicate that 'hbase-site.xml' is not on the classpath")
56         }
57         conf
58       }
59 
60       private def addResources(conf: Configuration, paths: String): Unit =
61         paths.split(',').map(_.trim).filterNot(_.isEmpty).foreach(HadoopUtils.addResource(conf, _))
62     }
63   )
64 
65   private val connections: LoadingCache[(Configuration, Boolean), CachedConnection] =  Caffeine.newBuilder().build(
66     new CacheLoader[(Configuration, Boolean), CachedConnection] {
67       override def load(key: (Configuration, Boolean)): CachedConnection = {
68         createConnection(key._1, key._2) match {
69           case SingletonConnection(connection, kerberos) => CachedConnection(connection, kerberos)
70           case c => throw new NotImplementedError(s"Expected SingletonConnection but got $c")
71         }
72       }
73     }
74   )
75 
76   Runtime.getRuntime.addShutdownHook(new Thread() {
77     override def run(): Unit =
78       CloseWithLogging(connections.asMap().values().asScala.flatMap { case CachedConnection(c, k) => Seq(c) ++ k })
79   })
80 
81   /**
82    * Get (or create) a cached configuration
83    *
84    * @param params data store params
85    * @return
86    */
87   def getConfiguration(params: java.util.Map[String, _]): Configuration = {
88     val zk = ZookeeperParam.lookupOpt(params)
89     val paths = ConfigPathsParam.lookupOpt(params)
90     val xml = ConfigsParam.lookupOpt(params)
91     configs.get(ConfigKey(zk, paths, xml))
92   }
93 
94   /**
95    * Get (or create) a cached connection
96    *
97    * @param params data store params
98    * @param validate validate the connection after creation, or not
99    * @return
100    */
101   def getConnection(params: java.util.Map[String, _], validate: Boolean): ConnectionWrapper = {
102     if (ConnectionParam.exists(params)) {
103       ProvidedConnection(ConnectionParam.lookup(params))
104     } else {
105       val conf = getConfiguration(params)
106       logger.debug(s"Connecting to HBase instance at ${conf.get(HConstants.ZOOKEEPER_QUORUM)}")
107       if (HBaseDataStoreParams.CacheConnectionsParam.lookup(params)) {
108         connections.get((conf, validate))
109       } else {
110         createConnection(conf, validate)
111       }
112     }
113   }
114 
115   /**
116    * Create a new connection (not pooled)
117    *
118    * @param conf hbase configuration
119    * @param validate validate the connection after creation, or not
120    * @return
121    */
122   def createConnection(conf: Configuration, validate: Boolean): ConnectionWrapper = {
123     if (User.isHBaseSecurityEnabled(conf)) {
124       configureSecurity(conf)
125       val action = new PrivilegedExceptionAction[ConnectionWrapper]() {
126         override def run(): ConnectionWrapper = doCreateConnection(conf, validate)
127       }
128       val user = UserGroupInformation.getLoginUser
129       logger.info(s"Creating Secured HBase connection with user $user")
130       user.doAs(action)
131     } else {
132       logger.info(s"Creating unsecured HBase connection")
133       doCreateConnection(conf, validate)
134     }
135   }
136 
137   private def doCreateConnection(conf: Configuration, validate: Boolean): ConnectionWrapper = {
138     if (validate) {
139       logger.debug("Checking configuration availability")
140       HBaseVersions.checkAvailable(conf)
141     }
142     val connection = ConnectionFactory.createConnection(conf)
143     val kerberos = if (User.isHBaseSecurityEnabled(conf)) { Some(HadoopUtils.kerberosTicketRenewer()) } else { None }
144     SingletonConnection(connection, kerberos)
145   }
146 
147   /**
148    * Configures hadoop security, based on the configuration.
149    *
150    * Note: hadoop security is configured globally - having different security settings in a single JVM
151    * will likely result in errors
152    *
153    * @param conf conf
154    */
155   def configureSecurity(conf: Configuration): Unit = synchronized {
156     import AuthenticationTokenIdentifier.AUTH_TOKEN_TYPE
157 
158     if (User.isHBaseSecurityEnabled(conf)) {
159       val currentUser = UserGroupInformation.getLoginUser
160       if (currentUser.getCredentials.getAllTokens.asScala.exists(_.getKind == AUTH_TOKEN_TYPE)) {
161         logger.debug("Using existing HBase authentication token")
162       } else {
163         val keytab = conf.get(HBaseGeoMesaKeyTab)
164         val rawPrincipal = conf.get(HBaseGeoMesaPrincipal)
165 
166         if (keytab == null || rawPrincipal == null) {
167           lazy val missing =
168             Seq(HBaseGeoMesaKeyTab -> keytab, HBaseGeoMesaPrincipal -> rawPrincipal).collect { case (k, null) => k }
169           logger.warn(s"Security is enabled but missing credentials under '${missing.mkString("' and '")}'")
170         } else {
171           val principal = fullPrincipal(rawPrincipal)
172 
173           lazy val principalMsg =
174             s"'$principal'${if (principal == rawPrincipal) { "" } else { s" (original '$rawPrincipal')"}}"
175           logger.debug(
176             s"Using Kerberos with principal $principalMsg, keytab '$keytab', " +
177                 s"and Hadoop authentication method '${SecurityUtil.getAuthenticationMethod(conf)}'")
178 
179           if (currentUser.hasKerberosCredentials && currentUser.getUserName == principal) {
180             logger.debug(s"User '$principal' is already authenticated")
181           } else {
182             if (currentUser.hasKerberosCredentials) {
183               logger.warn(
184                 s"Changing global authenticated Hadoop user from '${currentUser.getUserName}' to '$principal' -" +
185                     "this will affect any connections still using the old user")
186             }
187             UserGroupInformation.setConfiguration(conf)
188             UserGroupInformation.loginUserFromKeytab(principal, keytab)
189 
190             logger.debug(s"Logged into Hadoop with user '${UserGroupInformation.getLoginUser.getUserName}'")
191           }
192         }
193       }
194     }
195   }
196 
197   /**
198    * Replace _HOST with the current host and add the default realm if nothing is specified.
199    *
200    * `SecurityUtil.getServerPrincipal` will replace the _HOST but only if there is already a realm.
201    *
202    * @param principal kerberos principal
203    * @return
204    */
205   private def fullPrincipal(principal: String): String = {
206     if (principal.indexOf('@') != -1) {
207       // we have a realm so this should be work to replace _HOST if present
208       SecurityUtil.getServerPrincipal(principal, null: String)
209     } else {
210       // try to add the default realm and replace _HOST if present
211       Try(KerberosUtil.getDefaultRealm) match {
212         case Success(realm) => SecurityUtil.getServerPrincipal(s"$principal@$realm", null: String)
213         case Failure(e) =>
214           logger.warn(s"Unable to get default Kerberos realm: $e")
215           if (!principal.contains(SecurityUtil.HOSTNAME_PATTERN)) { principal } else {
216             // append a fake realm so that the _HOST replacement works and then remove it afterwards
217             SecurityUtil.getServerPrincipal(s"$principal@foo", null: String).dropRight(4)
218           }
219       }
220     }
221   }
222 
223   /**
224    * Managed connection. The connection itself should not be closed - instead close the wrapper to handle
225    * lifecycle events appropriately.
226    */
227   sealed trait ConnectionWrapper extends Closeable {
228     val connection: Connection
229   }
230 
231   /**
232    * An unshared connection
233    *
234    * @param connection connection
235    * @param kerberos kerberos ticket renewal thread
236    */
237   case class SingletonConnection(connection: Connection, kerberos: Option[Closeable]) extends ConnectionWrapper {
238     override def close(): Unit = CloseWithLogging(kerberos.toSeq ++ Seq(connection))
239   }
240 
241   /**
242    * A shared, cached connection
243    *
244    * @param connection connection
245    * @param kerberos kerberos ticket renewal thread
246    */
247   case class CachedConnection(connection: Connection, kerberos: Option[Closeable]) extends ConnectionWrapper {
248     override def close(): Unit = {}
249   }
250 
251   /**
252    * Provided connection - no lifecycle management is performed
253    *
254    * @param connection connection
255    */
256   case class ProvidedConnection(connection: Connection) extends ConnectionWrapper {
257     override def close(): Unit = {}
258   }
259 
260   private case class ConfigKey(zookeepers: Option[String], paths: Option[String], xml: Option[String])
261 }