From 6abbb21e086a79827618272972db0d95969fe558 Mon Sep 17 00:00:00 2001
From: vruttantmankad <vruttant.mankad@gojek.com>
Date: Mon, 29 Apr 2024 15:37:30 +0530
Subject: [PATCH] login module as config

---
 src/ziggurat/config.clj       | 23 +++++++++--------------
 test/ziggurat/config_test.clj | 10 ++++++----
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/src/ziggurat/config.clj b/src/ziggurat/config.clj
index 2a03f8be..8ca3484f 100644
--- a/src/ziggurat/config.clj
+++ b/src/ziggurat/config.clj
@@ -201,23 +201,17 @@
         (.setProperty p sk nv))))
   p)
 
-(def jaas-template
-  {"PLAIN"         "org.apache.kafka.common.security.plain.PlainLoginModule"
-   "SCRAM-SHA-512" "org.apache.kafka.common.security.scram.ScramLoginModule"
-   "SCRAM-SHA-256" "org.apache.kafka.common.security.scram.ScramLoginModule"})
-
 (defn create-jaas-properties
-  [user-name password mechanism]
-  (let [jaas-template (get jaas-template mechanism)]
-    (format "%s required username=\"%s\" password=\"%s\";" jaas-template user-name password)))
+  [user-name password login-module]
+  (format "%s required username=\"%s\" password=\"%s\";" login-module user-name password))
 
 (defn- add-jaas-properties
   [properties jaas-config]
   (if (some? jaas-config)
     (let [username  (get jaas-config :username)
           password  (get jaas-config :password)
-          mechanism (get jaas-config :mechanism)
-          jaas_props (create-jaas-properties username password mechanism)]
+          login-module (get jaas-config :login-module)
+          jaas_props (create-jaas-properties username password login-module)]
       (doto properties
         (.put SaslConfigs/SASL_JAAS_CONFIG jaas_props)))
     properties))
@@ -265,20 +259,21 @@
 
     SASL properties are only set if [:ziggurat :sasl :enabled] returns true.
 
-    Creates JAAS template if values are provided in the map provided agains this key sequence
-      [:ziggurat :ssl :jaas].
+    Creates JAAS template if values are provided in the map provided against this key sequence
+      [:ziggurat :sasl :jaas].
 
     Example of sasl-config-map
     {:enabled true
      :protocol <>
+     :mechanism <>
       {:jaas
         {:username <>
          :password <>
-         :mechanism}}}
+         :login-module <>}}}
   "
   (let [sasl-configs-enabled (:enabled sasl-config-map)
         jaas-config         (get sasl-config-map :jaas)
-        mechanism           (get jaas-config :mechanism)
+        mechanism           (get sasl-config-map :mechanism)
         protocol            (get sasl-config-map :protocol)]
     (if (true? sasl-configs-enabled)
       (as-> properties pr
diff --git a/test/ziggurat/config_test.clj b/test/ziggurat/config_test.clj
index dfbc708a..48b76b93 100644
--- a/test/ziggurat/config_test.clj
+++ b/test/ziggurat/config_test.clj
@@ -322,9 +322,10 @@
       (with-redefs [ssl-config (constantly {:enabled true
                                             :ssl-keystore-location "/some/location"
                                             :ssl-keystore-password "some-password"
+                                            :mechanism "SCRAM-SHA-512"
                                             :jaas {:username "myuser"
                                                    :password "mypassword"
-                                                   :mechanism "SCRAM-SHA-512"}})]
+                                                   :login-module "org.apache.kafka.common.security.scram.ScramLoginModule"}})]
         (let [streams-config-map {:auto-offset-reset  :latest}
               props              (build-streams-config-properties streams-config-map)
               auto-offset-reset  (.getProperty props "auto.offset.reset")
@@ -334,7 +335,7 @@
           (is (= auto-offset-reset "latest"))
           (is (= ssl-ks-location  "/some/location"))
           (is (= ssl-ks-password  "some-password"))
-          (is (= sasl-jaas-config (create-jaas-properties "myuser" "mypassword" "SCRAM-SHA-512"))))))
+          (is (= sasl-jaas-config (create-jaas-properties "myuser" "mypassword" "org.apache.kafka.common.security.scram.ScramLoginModule"))))))
     (testing "ssl properties DO NOT create jaas template if no value is provided for key sequence [:ziggurat :ssl :jaas]"
       (with-redefs [ssl-config (constantly {:enabled true
                                             :ssl-keystore-location "/some/location"
@@ -352,9 +353,10 @@
     (testing "sasl properties create jaas template from the map provided in [:ziggurat :sasl :jaas]"
       (with-redefs [sasl-config (constantly {:enabled true
                                              :protocol "SASL_PLAINTEXT"
+                                             :mechanism "SCRAM-SHA-256"
                                             :jaas {:username "myuser"
                                                    :password "mypassword"
-                                                   :mechanism "SCRAM-SHA-256"}})]
+                                                   :login-module "org.apache.kafka.common.security.scram.ScramLoginModule"}})]
         (let [streams-config-map {:auto-offset-reset  :latest}
               props              (build-streams-config-properties streams-config-map)
               auto-offset-reset  (.getProperty props "auto.offset.reset")
@@ -363,7 +365,7 @@
               sasl-mechanism     (.getProperty props "sasl.mechanism")]
           (is (= auto-offset-reset "latest"))
           (is (= sasl-protocol "SASL_PLAINTEXT"))
-          (is (= sasl-jaas-config (create-jaas-properties "myuser" "mypassword" "SCRAM-SHA-256"))))))))
+          (is (= sasl-jaas-config (create-jaas-properties "myuser" "mypassword" "org.apache.kafka.common.security.scram.ScramLoginModule"))))))))
 
 (deftest test-set-property
   (testing "set-property with empty (with spaces) value"