[FSTORE-1015] Add option to add Keytab to Kafka connector

java/beam/src/main/java/com/logicalclocks/hsfs/beam/engine/BeamEngine.java

@@ -32,6 +32,8 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 public class BeamEngine extends EngineBase {
   private static BeamEngine INSTANCE = null;
@@ -91,6 +93,18 @@ public Map<String, String> getKafkaConfig(FeatureGroupBase featureGroup, Map<Str
     storageConnector.setSslKeystoreLocation(addFile(storageConnector.getSslKeystoreLocation()));
 
     Map<String, String> config = storageConnector.kafkaOptions();
+    // add keytab file
+    String saslJaasConfig = config.get("sasl.jaas.config");
+    if (saslJaasConfig != null) {
+      Pattern pattern = Pattern.compile("keyTab=[\"'](.+?)[\"']");
+      Matcher matcher = pattern.matcher(saslJaasConfig);
+      while (matcher.find()) {
+        String originalKeytabLocation = matcher.group(1);
+        String newKeytabLocation = addFile(originalKeytabLocation);
+        saslJaasConfig = saslJaasConfig.replace(originalKeytabLocation, newKeytabLocation);
+      }
+      config.put("sasl.jaas.config", saslJaasConfig);
+    }
 
     if (writeOptions != null) {
       config.putAll(writeOptions);

java/flink/src/main/java/com/logicalclocks/hsfs/flink/engine/FlinkEngine.java

@@ -44,6 +44,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 import static org.apache.flink.configuration.ConfigOptions.key;
 
@@ -141,6 +143,18 @@ public Map<String, String> getKafkaConfig(FeatureGroupBase featureGroup, Map<Str
     storageConnector.setSslKeystoreLocation(addFile(storageConnector.getSslKeystoreLocation()));
 
     Map<String, String> config = storageConnector.kafkaOptions();
+    // add keytab file
+    String saslJaasConfig = config.get("sasl.jaas.config");
+    if (saslJaasConfig != null) {
+      Pattern pattern = Pattern.compile("keyTab=[\"'](.+?)[\"']");
+      Matcher matcher = pattern.matcher(saslJaasConfig);
+      while (matcher.find()) {
+        String originalKeytabLocation = matcher.group(1);
+        String newKeytabLocation = addFile(originalKeytabLocation);
+        saslJaasConfig = saslJaasConfig.replace(originalKeytabLocation, newKeytabLocation);
+      }
+      config.put("sasl.jaas.config", saslJaasConfig);
+    }
 
     if (writeOptions != null) {
       config.putAll(writeOptions);

java/spark/src/main/java/com/logicalclocks/hsfs/spark/engine/SparkEngine.java

@@ -144,7 +144,7 @@ public static void setInstance(SparkEngine sparkEngine) {
 
   private HudiEngine hudiEngine = new HudiEngine();
 
-  private SparkEngine() {
+  protected SparkEngine() {
     sparkSession = SparkSession.builder()
         .enableHiveSupport()
         .getOrCreate();
@@ -955,6 +955,18 @@ public Map<String, String> getKafkaConfig(FeatureGroupBase featureGroup, Map<Str
     storageConnector.setSslKeystoreLocation(addFile(storageConnector.getSslKeystoreLocation()));
 
     Map<String, String> config = storageConnector.sparkOptions();
+    // add keytab file
+    String saslJaasConfig = config.get("kafka.sasl.jaas.config");
+    if (saslJaasConfig != null) {
+      Pattern pattern = Pattern.compile("keyTab=[\"'](.+?)[\"']");
+      Matcher matcher = pattern.matcher(saslJaasConfig);
+      while (matcher.find()) {
+        String originalKeytabLocation = matcher.group(1);
+        String newKeytabLocation = addFile(originalKeytabLocation);
+        saslJaasConfig = saslJaasConfig.replace(originalKeytabLocation, newKeytabLocation);
+      }
+      config.put("kafka.sasl.jaas.config", saslJaasConfig);
+    }
 
     if (writeOptions != null) {
       config.putAll(writeOptions);

java/spark/src/test/java/com/logicalclocks/hsfs/spark/engine/TestSparkEngine.java

@@ -226,6 +226,35 @@ public void testGetKafkaConfigExternalClientInternalKafka() throws FeatureStoreE
         Assertions.assertEquals(Boolean.FALSE, externalArg.getValue());
     }
 
+    @Test
+    public void testGetKafkaKerberos() throws FeatureStoreException, IOException {
+        // Arrange
+        HopsworksClient hopsworksClient = Mockito.mock(HopsworksClient.class);
+        hopsworksClient.setInstance(new HopsworksClient(Mockito.mock(HopsworksHttpClient.class), "host"));
+        System.setProperty(HopsworksInternalClient.REST_ENDPOINT_SYS, "");
+
+        SparkEngine sparkEngine = Mockito.spy(SparkEngine.class);
+        StorageConnectorApi storageConnectorApi = Mockito.mock(StorageConnectorApi.class);
+        sparkEngine.storageConnectorApi = storageConnectorApi;
+        StorageConnector.KafkaConnector kafkaConnector = new StorageConnector.KafkaConnector();
+        kafkaConnector.setBootstrapServers("testServer:123");
+        kafkaConnector.setExternalKafka(Boolean.TRUE);
+        kafkaConnector.setSecurityProtocol(SecurityProtocol.SSL);
+        kafkaConnector.setOptions(Collections.singletonList(new Option("sasl.jaas.config", "com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab=\"/home/laurent/my.keytab\" storeKey=true useTicketCache=false serviceName=\"kafka\" principal=\"laurent@kafka.com\";")));
+        kafkaConnector.setSslEndpointIdentificationAlgorithm(SslEndpointIdentificationAlgorithm.EMPTY);
+
+        Mockito.when(sparkEngine.addFile(Mockito.anyString())).thenReturn("result_from_add_file");
+
+        Mockito.when(storageConnectorApi.getKafkaStorageConnector(Mockito.any(), Mockito.anyBoolean()))
+            .thenReturn(kafkaConnector);
+
+        // Act
+        Map<String, String> result = sparkEngine.getKafkaConfig(new FeatureGroup(), new HashMap<>());
+
+        // Assert
+        Assertions.assertEquals("com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab=\"result_from_add_file\" storeKey=true useTicketCache=false serviceName=\"kafka\" principal=\"laurent@kafka.com\";", result.get("kafka.sasl.jaas.config"));
+    }
+
     @Test
     public void testMakeQueryName() {
         SparkEngine sparkEngine = SparkEngine.getInstance();

python/hsfs/storage_connector.py

@@ -944,6 +944,19 @@ def __init__(
         self._external_kafka = external_kafka
         self._pem_files_created = False
 
+        # add keytab file
+        sasl_jaas_config = self._options.get("sasl.jaas.config")
+        if sasl_jaas_config:
+            for option in re.findall("keyTab=[\"'](.+?)[\"']", sasl_jaas_config):
+                original_keytab_location = option
+                new_keytab_location = engine.get_instance().add_file(
+                    original_keytab_location
+                )
+                sasl_jaas_config = sasl_jaas_config.replace(
+                    original_keytab_location, new_keytab_location
+                )
+            self._options["sasl.jaas.config"] = sasl_jaas_config
+
     @property
     def bootstrap_servers(self):
         """Bootstrap servers string."""
@@ -1064,12 +1077,17 @@ def confluent_options(self):
                 config["ssl.key.location"] = client_key_path
             elif key == "sasl.jaas.config":
                 groups = re.search(
-                    "(.+?) .*username=[\"'](.+?)[\"'] .*password=[\"'](.+?)[\"']",
-                    value,
+                    "(.+) (required|requisite|sufficient|optional)(.*)", value
                 )
+                mechanism = groups.group(1)
+                # flag = groups.group(2)
+                options = groups.group(3)
+
+                option_dict = {}
+                for option in re.findall(r"\s(\w+)=[\"'](.+?)[\"']", options):
+                    option_dict[option[0]] = option[1]
+
                 if "sasl.mechanisms" not in config:
-                    mechanism = groups.group(1)
-                    mechanism_value = None
                     if (
                         mechanism
                         == "org.apache.kafka.common.security.plain.PlainLoginModule"
@@ -1085,9 +1103,30 @@ def confluent_options(self):
                         == "org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule"
                     ):
                         mechanism_value = "OAUTHBEARER"
-                config["sasl.mechanisms"] = mechanism_value
-                config["sasl.username"] = groups.group(2)
-                config["sasl.password"] = groups.group(3)
+                    else:
+                        mechanism_value = "GSSAPI"
+                    config["sasl.mechanisms"] = mechanism_value
+
+                    if mechanism_value == "GSSAPI":
+                        service_name = option_dict.get("serviceName")
+                        if service_name:
+                            config["sasl.kerberos.service.name"] = service_name
+
+                        principal = option_dict.get("principal")
+                        if principal:
+                            config["sasl.kerberos.principal"] = principal
+
+                        key_tab = option_dict.get("keyTab")
+                        if key_tab:
+                            config["sasl.kerberos.keytab"] = key_tab
+                    else:
+                        username = option_dict.get("username")
+                        if username:
+                            config["sasl.username"] = username
+
+                        password = option_dict.get("password")
+                        if password:
+                            config["sasl.password"] = password
             elif key == "ssl.endpoint.identification.algorithm":
                 config[key] = "none" if value == "" else value
             elif key == "queued.max.requests":

python/tests/test_storage_connector.py

@@ -740,6 +740,108 @@ def test_confluent_options_jaas_double_quotes(self, mocker):
             "sasl.username": "222",
         }
 
+    def test_kafka_options_kerberos(self, mocker, backend_fixtures):
+        # Arrange
+        mock_engine_get_instance = mocker.patch("hsfs.engine.get_instance")
+        mock_engine_get_instance.return_value.add_file.side_effect = [
+            None,
+            None,
+            "result_from_add_file",
+        ]
+        mocker.patch("hsfs.client.get_instance")
+        sc = storage_connector.KafkaConnector(
+            1,
+            "kafka_connector",
+            0,
+            external_kafka=True,
+            options=[
+                {
+                    "name": "sasl.jaas.config",
+                    "value": 'com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab="/home/laurent/my.keytab" storeKey=true useTicketCache=false serviceName="kafka" principal="laurent@kafka.com";',
+                }
+            ],
+        )
+
+        # Act
+        config = sc.kafka_options()
+
+        # Assert
+        assert config == {
+            "bootstrap.servers": None,
+            "security.protocol": None,
+            "ssl.endpoint.identification.algorithm": None,
+            "sasl.jaas.config": 'com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab="result_from_add_file" storeKey=true useTicketCache=false serviceName="kafka" principal="laurent@kafka.com";',
+        }
+
+    def test_spark_options_kerberos(self, mocker, backend_fixtures):
+        # Arrange
+        mock_engine_get_instance = mocker.patch("hsfs.engine.get_instance")
+        mock_engine_get_instance.return_value.add_file.side_effect = [
+            None,
+            None,
+            "result_from_add_file",
+        ]
+        mocker.patch("hsfs.client.get_instance")
+        sc = storage_connector.KafkaConnector(
+            1,
+            "kafka_connector",
+            0,
+            external_kafka=True,
+            options=[
+                {
+                    "name": "sasl.jaas.config",
+                    "value": 'com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab="/home/laurent/my.keytab" storeKey=true useTicketCache=false serviceName="kafka" principal="laurent@kafka.com";',
+                }
+            ],
+        )
+
+        # Act
+        config = sc.spark_options()
+
+        # Assert
+        assert config == {
+            "kafka.bootstrap.servers": None,
+            "kafka.security.protocol": None,
+            "kafka.ssl.endpoint.identification.algorithm": None,
+            "kafka.sasl.jaas.config": 'com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab="result_from_add_file" storeKey=true useTicketCache=false serviceName="kafka" principal="laurent@kafka.com";',
+        }
+
+    def test_confluent_options_kerberos(self, mocker):
+        # Arrange
+        mock_engine_get_instance = mocker.patch("hsfs.engine.get_instance")
+        mock_engine_get_instance.return_value.add_file.side_effect = [
+            None,
+            None,
+            "result_from_add_file",
+        ]
+        mocker.patch("hsfs.client.get_instance")
+        sc = storage_connector.KafkaConnector(
+            1,
+            "kafka_connector",
+            0,
+            external_kafka=True,
+            options=[
+                {
+                    "name": "sasl.jaas.config",
+                    "value": 'com.sun.security.auth.module.Krb5LoginModule required useKeyTab=true keyTab="/home/laurent/my.keytab" storeKey=true useTicketCache=false serviceName="kafka" principal="laurent@kafka.com";',
+                }
+            ],
+        )
+
+        # Act
+        config = sc.confluent_options()
+
+        # Assert
+        assert config == {
+            "bootstrap.servers": None,
+            "security.protocol": None,
+            "ssl.endpoint.identification.algorithm": None,
+            "sasl.mechanisms": "GSSAPI",
+            "sasl.kerberos.service.name": "kafka",
+            "sasl.kerberos.principal": "laurent@kafka.com",
+            "sasl.kerberos.keytab": "result_from_add_file",
+        }
+
 
 class TestGcsConnector:
     def test_from_response_json(self, backend_fixtures):