refactor: use MEMORY type cache (#36)

* refactor: use memory cache * doc: add function doc for kinit * refactor: set env KRB5CCNAME for default ccname when load nif but also allow kinit/3 to override when the name is not empty string * test: cover custom ccname in tests
kafka4beam · Aug 29, 2024 · 0010a80 · 0010a80
1 parent 2e91b77
commit 0010a80
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 33 deletions.
diff --git a/c_src/sasl_auth.c b/c_src/sasl_auth.c
@@ -1,3 +1,4 @@
+#define _DEFAULT_SOURCE 200112L
 #include <erl_nif.h>
 #include <krb5.h>
 #include <sasl/sasl.h>
@@ -31,6 +32,8 @@ static ERL_NIF_TERM ATOM_NOT_CONTROLLING_PROCESS;
     ERROR_TUPLE(env, enif_make_tuple2(env, enif_make_int(env, code), sasl_error(env, state)));
 
 #define KT_NAME_LEN 1024
+#define DEFAULT_CCNAME "MEMORY:krb5cc_sasl_auth"
+
 
 typedef struct {
     sasl_conn_t* conn;
@@ -142,6 +145,7 @@ static int load(ErlNifEnv* env, void** UNUSED(priv), ERL_NIF_TERM UNUSED(info))
     int cli_result = sasl_client_init(NULL);
     sasl_server_connection_nif_resource_type = init_resource_type(env, "sasl_auth_srv_state");
     int srv_result = sasl_server_init(NULL, "sasl_auth");
+    setenv("KRB5CCNAME", DEFAULT_CCNAME, 1);
     return !sasl_client_connection_nif_resource_type && !(cli_result == SASL_OK)
         && !sasl_server_connection_nif_resource_type && !(srv_result == SASL_OK);
 }
@@ -669,13 +673,21 @@ static ERL_NIF_TERM sasl_krb5_kt_default_name(ErlNifEnv* env, int UNUSED(argc),
     }
 }
 
+static void enif_free_non_null(void* ptr) {
+    if (ptr != NULL) {
+        enif_free(ptr);
+    }
+}
+
 static ERL_NIF_TERM sasl_kinit(ErlNifEnv* env, int UNUSED(argc), const ERL_NIF_TERM argv[])
 {
 
-    ErlNifBinary keytab;
-    ErlNifBinary principal_in;
-    unsigned char* principal_mut = NULL;
-    unsigned char* keytab_in = NULL;
+    ErlNifBinary keytab_bin;
+    ErlNifBinary principal_bin;
+    ErlNifBinary ccname_bin;
+    unsigned char* principal_char = NULL;
+    unsigned char* keytab_char = NULL;
+    unsigned char* ccname_char = NULL;
     ERL_NIF_TERM ret;
     ERL_NIF_TERM error_tag;
     ERL_NIF_TERM error_code;
@@ -687,45 +699,67 @@ static ERL_NIF_TERM sasl_kinit(ErlNifEnv* env, int UNUSED(argc), const ERL_NIF_T
     krb5_context context = NULL;
     krb5_creds creds = { .magic = 0 };
     krb5_keytab kt_handle = NULL;
-    krb5_ccache defcache = NULL;
+    krb5_ccache ccache = NULL;
     krb5_get_init_creds_opt* options = NULL;
 
     const char* krb_error_msg;
     const char* tag;
     int handle_alive = 0;
     int cache_alive = 0;
 
-    if ((!enif_inspect_binary(env, argv[0], &keytab)
-            || !enif_inspect_binary(env, argv[1], &principal_in)))
+    if ( !enif_inspect_binary(env, argv[0], &keytab_bin) ) {
         return enif_make_badarg(env);
+    }
+    if ( !enif_inspect_binary(env, argv[1], &principal_bin) ) {
+        return enif_make_badarg(env);
+    }
+    if ( !enif_inspect_binary(env, argv[2], &ccname_bin) ) {
+        return enif_make_badarg(env);
+    }
 
-    principal_mut = copy_bin(principal_in);
-    if (principal_mut == NULL) {
-        return ERROR_TUPLE(env, ATOM_OOM);
+    keytab_char = copy_bin(keytab_bin);
+    if (keytab_char == NULL) {
+        ret = ERROR_TUPLE(env, ATOM_OOM);
+        goto kinit_free_chars;
     }
-    keytab_in = copy_bin(keytab);
-    if (keytab_in == NULL) {
-        return ERROR_TUPLE(env, ATOM_OOM);
+
+    principal_char = copy_bin(principal_bin);
+    if (principal_char == NULL) {
+        ret = ERROR_TUPLE(env, ATOM_OOM);
+        goto kinit_free_chars;
+    }
+
+    ccname_char = copy_bin(ccname_bin);
+    if (ccname_char == NULL) {
+        ret = ERROR_TUPLE(env, ATOM_OOM);
+        goto kinit_free_chars;
     }
 
     if ((error = krb5_init_context(&context)) != 0) {
         tag = "krb5_parse_context";
         goto kinit_finish;
     }
 
-    if ((error = krb5_parse_name(context, (const char*)principal_mut, &principal)) != 0) {
+    if ((error = krb5_parse_name(context, (const char*)principal_char, &principal)) != 0) {
         tag = "krb5_parse_name";
         goto kinit_finish;
     }
 
-    if ((error = krb5_kt_resolve(context, (const char*)keytab_in, &kt_handle)) != 0) {
+    if ((error = krb5_kt_resolve(context, (const char*)keytab_char, &kt_handle)) != 0) {
         tag = "krb5_kt_resolve";
         goto kinit_finish;
     }
 
     handle_alive = 1;
 
-    if ((error = krb5_cc_default(context, &defcache)) != 0) {
+    /* NOTWORKING: (error = krb5_cc_resolve(context, (const char*)ccname_char, &ccache))
+     *
+     * krb5 doc says krb5_cc_default is essentially krb5_cc_resolve with default ccname, but it does not work.
+     * So we set environment variable KRB5CCNAME and call krb5_cc_default instead */
+    if (ccname_char[0] != 0) {
+        setenv("KRB5CCNAME", (const char*)ccname_char, 1);
+    }
+    if ((error = krb5_cc_default(context, &ccache)) != 0) {
         tag = "krb5_cc_default";
         goto kinit_finish;
     }
@@ -740,15 +774,14 @@ static ERL_NIF_TERM sasl_kinit(ErlNifEnv* env, int UNUSED(argc), const ERL_NIF_T
 /* It's not clear why this call fails on mac. For the time being initially keytab init must
  * be done on the command line */
 #if (!defined __APPLE__ || !defined __MACH__)
-    if ((error = krb5_get_init_creds_opt_set_out_ccache(context, options, defcache)) != 0) {
+    if ((error = krb5_get_init_creds_opt_set_out_ccache(context, options, ccache)) != 0) {
         tag = "krb5_get_init_creds_opt_set_out_ccache";
         goto kinit_finish;
     }
 #endif
 
-    if ((error
-            = krb5_get_init_creds_keytab(context, &creds, principal, kt_handle, 0, NULL, options))
-        != 0) {
+
+    if ((error = krb5_get_init_creds_keytab(context, &creds, principal, kt_handle, 0, NULL, options)) != 0) {
         tag = "krb5_get_init_creds_keytab";
         goto kinit_finish;
     }
@@ -768,7 +801,7 @@ static ERL_NIF_TERM sasl_kinit(ErlNifEnv* env, int UNUSED(argc), const ERL_NIF_T
     }
 
     if (1 == cache_alive) {
-        krb5_cc_close(context, defcache);
+        krb5_cc_close(context, ccache);
     }
 
     if (1 == handle_alive) {
@@ -782,8 +815,11 @@ static ERL_NIF_TERM sasl_kinit(ErlNifEnv* env, int UNUSED(argc), const ERL_NIF_T
     }
     krb5_free_context(context);
 
-    enif_free(principal_mut);
-    enif_free(keytab_in);
+kinit_free_chars:
+
+    enif_free_non_null(principal_char);
+    enif_free_non_null(keytab_char);
+    enif_free_non_null(ccname_char);
     return ret;
 }
 
@@ -793,7 +829,7 @@ static ErlNifFunc nif_funcs[]
           { "sasl_client_start", 1, sasl_cli_start, ERL_NIF_DIRTY_JOB_CPU_BOUND },
           { "sasl_client_step", 2, sasl_cli_step, ERL_NIF_DIRTY_JOB_CPU_BOUND },
           { "sasl_client_done", 1, sasl_cli_done, ERL_NIF_DIRTY_JOB_CPU_BOUND },
-          { "sasl_kinit", 2, sasl_kinit, ERL_NIF_DIRTY_JOB_CPU_BOUND },
+          { "sasl_kinit", 3, sasl_kinit, ERL_NIF_DIRTY_JOB_CPU_BOUND },
           { "sasl_server_new", 3, sasl_srv_new, ERL_NIF_DIRTY_JOB_CPU_BOUND },
           { "sasl_server_start", 2, sasl_srv_start, ERL_NIF_DIRTY_JOB_CPU_BOUND },
           { "sasl_server_step", 2, sasl_srv_step, ERL_NIF_DIRTY_JOB_CPU_BOUND },

diff --git a/scripts/int-test.sh b/scripts/int-test.sh
@@ -6,7 +6,7 @@
 set -euo pipefail
 
 KRB5_IMAGE='sasl_auth_dockerfile.ubuntu22.04'
-ERLANGE_IMAGE='ghcr.io/emqx/emqx-builder/5.3-9:1.15.7-26.2.5-3-ubuntu20.04'
+ERLANGE_IMAGE='ghcr.io/emqx/emqx-builder/5.3-9:1.15.7-26.2.5-3-ubuntu22.04'
 #ERLANGE_IMAGE="$KRB5_IMAGE"
 NET='example.com'
 REALM='EXAMPLE.COM'

diff --git a/scripts/int_test.erl b/scripts/int_test.erl
@@ -19,7 +19,8 @@ do_start(I) ->
 run_cli(I) ->
     Service = env("SERVICE"),
     maybe
-        ok = sasl_auth:kinit("cli.keytab", env("CLI_PRINC")),
+        ok = sasl_auth:kinit("cli.keytab", env("CLI_PRINC"), "MEMORY:test"),
+        io:format("done kinit~n", []),
         {ok, C} ?= sasl_auth:client_new(Service, env("SRV"), env("CLI_PRINC"), env("CLI_NAME")),
         Pid = wait_for_server(I),
         {ok, AvaialbeMecs} ?= sasl_auth:client_listmech(C),

diff --git a/src/sasl_auth.erl b/src/sasl_auth.erl
@@ -8,6 +8,7 @@
 -export([
     init/0,
     kinit/2,
+    kinit/3,
     client_new/3,
     client_new/4,
     client_listmech/1,
@@ -64,6 +65,7 @@
 -type state() :: reference().
 -type keytab_path() :: file:filename_all().
 -type principal() :: string() | binary().
+-type ccname() :: string() | binary().
 -type service_name() :: string() | binary().
 -type host() :: string() | binary().
 -type user() :: string() | binary().
@@ -158,10 +160,31 @@ init() ->
             ok
     end.
 
--spec kinit(KeyTabPath :: keytab_path(), Principal :: principal()) ->
+%% @doc Initialize credentials from a keytab file and principal.
+%% It makes use of the per Erlang node static MEMORY type ccache.
+-spec kinit(keytab_path(), principal()) ->
     ok | {error, {binary(), integer(), binary()}}.
 kinit(KeyTabPath, Principal) ->
-    sasl_kinit(null_terminate(KeyTabPath), null_terminate(Principal)).
+    kinit(KeyTabPath, Principal, <<>>).
+
+%% @hidden Initialize credentials from a keytab file and principal.
+%% The argument CCname is provided for application's flexibility to decide
+%% which credentials cache type or name to use.
+%% When set to empty string, the default cache name `MEMORY:krb5cc_sasl_auth'
+%% is used.
+%% e.g. `FILE:/tmp/krb5cc_mycache' or `MEMORY:krbcc5_mycache'
+%%
+%% CAUTION: Changing credentials cache name at runtime is not tested!
+%% CAUTION: There is currently a lack of call to krb5_cc_destroy, creating
+%%          many caches at runtime may lead to memory leak.
+-spec kinit(keytab_path(), principal(), ccname()) ->
+    ok | {error, {binary(), integer(), binary()}}.
+kinit(KeyTabPath, Principal, Ccname) ->
+    sasl_kinit(
+        null_terminate(KeyTabPath),
+        null_terminate(Principal),
+        null_terminate(Ccname)
+    ).
 
 %% @doc Initialize a client context. User client's principal as client's username.
 %% This is the default behaviour before version 2.1.1, however may not work when
@@ -303,7 +326,7 @@ krb5_kt_default_name() -> sasl_krb5_kt_default_name().
 
 sasl_krb5_kt_default_name() -> not_loaded(?LINE).
 
-sasl_kinit(_, _) -> not_loaded(?LINE).
+sasl_kinit(_KeyTab, _Principal, _CacheName) -> not_loaded(?LINE).
 
 sasl_client_new(_Service, _Host, _Principal, _User) -> not_loaded(?LINE).
 

diff --git a/test/sasl_auth_SUITE.erl b/test/sasl_auth_SUITE.erl
@@ -9,6 +9,7 @@ all() ->
         default_keytab_test,
         kinit_test,
         simple_test,
+        custom_ccname,
         kinit_keytab_fail_test,
         kinit_invalid_principal_test,
         concurrency_test,
@@ -28,10 +29,6 @@ init_per_suite(Config) ->
     ServiceKeyTab = get_env(service_keytab, "SASL_AUTH_KAFKA_KEY_TAB"),
     ServiceName = get_env(service_principal, "SASL_AUTH_KAFKA_PRINCIPAL"),
     Service = {service, <<"kafka">>},
-
-    ok = sasl_auth:kinit(element(2, UserKeyTab), element(2, UserPrincipal)),
-    %% Unable to kinit with service keytab here, only one keytab can be kinit at a time.
-
     [
         UserKeyTab,
         UserPrincipal,
@@ -46,6 +43,17 @@ init_per_suite(Config) ->
 end_per_suite(_Config) ->
     ok.
 
+init_per_testcase(Case, Config) ->
+    ClientKeyTab = ?config(user_keytab, Config),
+    ClientPrincipal = ?config(user_principal, Config),
+    case Case of
+        custom_ccname ->
+            ok = sasl_auth:kinit(ClientKeyTab, ClientPrincipal, "MEMORY:custom_ccname");
+        _ ->
+            ok = sasl_auth:kinit(ClientKeyTab, ClientPrincipal)
+    end,
+    Config.
+
 default_keytab_test(_Config) ->
     ?assertMatch(<<"FILE:", _/binary>>, sasl_auth:krb5_kt_default_name()).
 
@@ -55,6 +63,12 @@ kinit_test(Config) ->
     ok = sasl_auth:kinit(KeyTab, Principal).
 
 simple_test(Config) ->
+    client_server_interaction(Config).
+
+custom_ccname(Config) ->
+    client_server_interaction(Config).
+
+client_server_interaction(Config) ->
     {ok, CliConn} = setup_default_client(Config),
     {ok, [_ | _]} = sasl_auth:client_listmech(CliConn),
     {ok, {sasl_continue, ClientToken}} = sasl_auth:client_start(CliConn),