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