jax_fingerprint: Address user arg slot exhaustion

When multiple jax_fingerprint.so instances are loaded (e.g., one per
fingerprinting method), each instance was reserving its own user arg
slot. ATS has a limited number of slots (~4 per type), causing methods
like JA3 and JA4 to fail silently when slots were exhausted.

Solution: Share a single user arg slot per type (TS_USER_ARGS_VCONN or
TS_USER_ARGS_TXN) across all jax_fingerprint instances. A ContextMap
stores JAxContext instances keyed by method name.

Author: bneradt

plugins/experimental/jax_fingerprint/context_map.h

@@ -0,0 +1,109 @@
+/** @file
+ *
+ * Shared context map for jax_fingerprint plugin instances.
+ *
+ * When multiple jax_fingerprint.so instances are loaded (e.g., one per
+ * fingerprinting method), they share a single user arg slot. This map
+ * stores the JAxContext for each method, keyed by method name.
+ *
+ * @section license License
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "config.h"
+#include "context.h"
+
+#include <string>
+#include <string_view>
+#include <unordered_map>
+
+/**
+ * @brief Container holding JAxContext instances for multiple methods.
+ *
+ * ATS has a limited number of user arg slots (~4 per type). When loading
+ * many jax_fingerprint instances, we share a single slot and store all
+ * contexts in this map, keyed by method name.
+ */
+class ContextMap
+{
+public:
+  ~ContextMap()
+  {
+    for (auto &pair : m_contexts) {
+      delete pair.second;
+    }
+  }
+
+  /**
+   * @brief Store a context for a method.
+   * @param[in] method_name The method name (e.g., "JA3", "JA4").
+   * @param[in] ctx The context to store. Ownership is transferred.
+   */
+  void
+  set(std::string_view method_name, JAxContext *ctx)
+  {
+    auto it = m_contexts.find(method_name);
+    if (it != m_contexts.end()) {
+      delete it->second;
+      it->second = ctx;
+    } else {
+      m_contexts.emplace(std::string{method_name}, ctx);
+    }
+  }
+
+  /**
+   * @brief Retrieve a context for a method.
+   * @param[in] method_name The method name.
+   * @return The context, or nullptr if not found.
+   */
+  JAxContext *
+  get(std::string_view method_name)
+  {
+    auto it = m_contexts.find(method_name);
+    return it != m_contexts.end() ? it->second : nullptr;
+  }
+
+  /**
+   * @brief Remove a context for a method.
+   * @param[in] method_name The method name.
+   */
+  void
+  remove(std::string_view method_name)
+  {
+    auto it = m_contexts.find(method_name);
+    if (it != m_contexts.end()) {
+      delete it->second;
+      m_contexts.erase(it);
+    }
+  }
+
+  /**
+   * @brief Check if the map is empty.
+   * @return True if no contexts are stored.
+   */
+  bool
+  empty() const
+  {
+    return m_contexts.empty();
+  }
+
+private:
+  std::unordered_map<std::string, JAxContext *, StringHash, std::equal_to<>> m_contexts;
+};

plugins/experimental/jax_fingerprint/plugin.cc

@@ -282,8 +282,7 @@ handle_http_txn_close(void *edata, PluginConfig &config)
 {
   TSHttpTxn txnp = static_cast<TSHttpTxn>(edata);
 
-  delete get_user_arg(txnp, config);
-  set_user_arg(txnp, config, nullptr);
+  cleanup_user_arg(txnp, config);
 
   TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
   return TS_SUCCESS;
@@ -294,8 +293,7 @@ handle_vconn_close(void *edata, PluginConfig &config)
 {
   TSVConn vconn = static_cast<TSVConn>(edata);
 
-  delete get_user_arg(vconn, config);
-  set_user_arg(vconn, config, nullptr);
+  cleanup_user_arg(vconn, config);
 
   TSVConnReenable(vconn);
   return TS_SUCCESS;

plugins/experimental/jax_fingerprint/userarg.cc

@@ -22,32 +22,88 @@
 #include "plugin.h"
 #include "config.h"
 #include "userarg.h"
+#include "context_map.h"
+
+namespace
+{
+
+// Shared user arg indices for all jax_fingerprint instances. ATS has limited
+// slots (~4 per type), so we share one slot per type and use a ContextMap.
+int  vconn_user_arg_index = -1;
+int  txn_user_arg_index   = -1;
+bool vconn_slot_reserved  = false;
+bool txn_slot_reserved    = false;
+
+} // anonymous namespace
 
 int
 reserve_user_arg(PluginConfig &config)
 {
-  std::string name  = PLUGIN_NAME;
-  name             += config.method.name;
-
   TSUserArgType type;
+  int          *shared_index;
+  bool         *reserved_flag;
+
   if (config.method.type == Method::Type::CONNECTION_BASED) {
-    type = TS_USER_ARGS_VCONN;
+    type          = TS_USER_ARGS_VCONN;
+    shared_index  = &vconn_user_arg_index;
+    reserved_flag = &vconn_slot_reserved;
   } else {
-    type = TS_USER_ARGS_TXN;
+    type          = TS_USER_ARGS_TXN;
+    shared_index  = &txn_user_arg_index;
+    reserved_flag = &txn_slot_reserved;
+  }
+
+  // Only reserve the slot once per type; subsequent calls reuse it.
+  if (!*reserved_flag) {
+    int ret = TSUserArgIndexReserve(type, PLUGIN_NAME, "shared JAx context map", shared_index);
+    if (ret == TS_SUCCESS) {
+      *reserved_flag = true;
+      Dbg(dbg_ctl, "Reserved shared user_arg slot: type=%d, index=%d", static_cast<int>(type), *shared_index);
+    } else {
+      Dbg(dbg_ctl, "Failed to reserve shared user_arg slot: type=%d", static_cast<int>(type));
+      return ret;
+    }
   }
-  int ret = TSUserArgIndexReserve(type, name.c_str(), "used to pass JAx context between hooks", &config.user_arg_index);
-  Dbg(dbg_ctl, "user_arg_name: %s, user_arg_index: %d", name.c_str(), config.user_arg_index);
-  return ret;
+
+  config.user_arg_index = *shared_index;
+  Dbg(dbg_ctl, "Using shared user_arg: method=%.*s, index=%d", static_cast<int>(config.method.name.size()),
+      config.method.name.data(), config.user_arg_index);
+  return TS_SUCCESS;
 }
 
 void
 set_user_arg(void *container, PluginConfig &config, JAxContext *ctx)
 {
-  TSUserArgSet(container, config.user_arg_index, static_cast<void *>(ctx));
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map == nullptr) {
+    map = new ContextMap();
+    TSUserArgSet(container, config.user_arg_index, static_cast<void *>(map));
+  }
+  map->set(config.method.name, ctx);
 }
 
 JAxContext *
 get_user_arg(void *container, PluginConfig &config)
 {
-  return static_cast<JAxContext *>(TSUserArgGet(container, config.user_arg_index));
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map == nullptr) {
+    return nullptr;
+  }
+  return map->get(config.method.name);
+}
+
+void
+cleanup_user_arg(void *container, PluginConfig &config)
+{
+  ContextMap *map = static_cast<ContextMap *>(TSUserArgGet(container, config.user_arg_index));
+  if (map != nullptr) {
+    // Remove this plugin's context from the map.
+    map->remove(config.method.name);
+
+    // If the map is now empty, delete it and clear the user arg.
+    if (map->empty()) {
+      delete map;
+      TSUserArgSet(container, config.user_arg_index, nullptr);
+    }
+  }
 }

plugins/experimental/jax_fingerprint/userarg.h

@@ -30,3 +30,5 @@
 int         reserve_user_arg(PluginConfig &config);
 void        set_user_arg(void *container, PluginConfig &config, JAxContext *ctx);
 JAxContext *get_user_arg(void *container, PluginConfig &config);
+void        cleanup_user_arg(void *container, PluginConfig &config);
+void        cleanup_user_arg(void *container, PluginConfig &config);

tests/gold_tests/pluginTest/jax_fingerprint/jax_fingerprint.test.py

@@ -337,3 +337,100 @@ def _configure_client(self, tr: 'TestRun') -> None:
 # Remap plugin sets headers on both routes, but only the SNI-allowed
 # connection has a vconn context, so only that request gets headers.
 JaxFingerprintTest('Hybrid JA4 servernames', 'JA4', 'hybrid', servernames='jax.server.test')
+
+# ======================================================================
+# All Methods Test - Verify shared context map works with multiple methods
+# ======================================================================
+
+
+class AllMethodsTest:
+    '''Test multiple fingerprint methods loaded simultaneously.
+
+    When multiple jax_fingerprint instances are loaded, they share user arg
+    slots via a ContextMap. This test verifies that all methods produce
+    fingerprints when sharing the SSL_CLIENT_HELLO_HOOK.
+    '''
+
+    _dns_counter: int = 0
+    _server_counter: int = 0
+    _ts_counter: int = 0
+    _client_counter: int = 0
+
+    def __init__(self, name: str) -> None:
+        '''Configure test with multiple methods loaded.'''
+        self._name = name
+        self._replay_file = 'jax_fingerprint_all_methods.replay.yaml'
+
+        tr = Test.AddTestRun(name)
+        self._configure_dns(tr)
+        self._configure_server(tr)
+        self._configure_trafficserver()
+        self._configure_client(tr)
+
+    def _configure_dns(self, tr: 'TestRun') -> None:
+        '''Configure a nameserver for the test.'''
+        name = f'dns_all{AllMethodsTest._dns_counter}'
+        self._dns = tr.MakeDNServer(name, default='127.0.0.1')
+        AllMethodsTest._dns_counter += 1
+
+    def _configure_server(self, tr: 'TestRun') -> None:
+        '''Configure the origin (verifier) server.'''
+        name = f'server_all{AllMethodsTest._server_counter}'
+        self._server = tr.AddVerifierServerProcess(name, self._replay_file)
+        AllMethodsTest._server_counter += 1
+
+        # Verify all headers were forwarded to the origin.
+        for header in ['x-ja3', 'x-ja4', 'x-ja4h']:
+            self._server.Streams.All += Testers.ContainsExpression(
+                rf'{header}:', f'Verify {header} header was forwarded.', reflags=re.IGNORECASE)
+
+    def _configure_trafficserver(self) -> None:
+        '''Configure Traffic Server with multiple methods.'''
+        name = f'ts_all{AllMethodsTest._ts_counter}'
+        self._ts = Test.MakeATSProcess(name, enable_cache=False, enable_tls=True)
+        AllMethodsTest._ts_counter += 1
+
+        self._ts.addDefaultSSLFiles()
+        self._ts.Disk.ssl_multicert_config.AddLine('dest_ip=* ssl_cert_name=server.pem ssl_key_name=server.key')
+
+        server_port = self._server.Variables.https_port
+
+        self._ts.Disk.records_config.update(
+            {
+                'proxy.config.ssl.server.cert.path': self._ts.Variables.SSLDir,
+                'proxy.config.ssl.server.private_key.path': self._ts.Variables.SSLDir,
+                'proxy.config.ssl.client.verify.server.policy': 'PERMISSIVE',
+                'proxy.config.dns.nameservers': f"127.0.0.1:{self._dns.Variables.Port}",
+                'proxy.config.dns.resolv_conf': 'NULL',
+                'proxy.config.proxy_name': 'test.proxy.test',
+                'proxy.config.diags.debug.enabled': 1,
+                'proxy.config.diags.debug.tags': 'jax_fingerprint',
+            })
+
+        # Load multiple methods - all share the same user arg slot via ContextMap.
+        self._ts.Disk.plugin_config.AddLines(
+            [
+                'jax_fingerprint.so --method JA3 --header x-ja3 --standalone',
+                'jax_fingerprint.so --method JA4 --header x-ja4 --standalone',
+                'jax_fingerprint.so --method JA4H --header x-ja4h --standalone',
+            ])
+
+        self._ts.Disk.remap_config.AddLine(f'map https://jax.server.test https://jax.backend.test:{server_port}')
+
+    def _configure_client(self, tr: 'TestRun') -> None:
+        '''Configure the verifier client.'''
+        name = f'client_all{AllMethodsTest._client_counter}'
+        p = tr.AddVerifierClientProcess(
+            name, self._replay_file, http_ports=[self._ts.Variables.port], https_ports=[self._ts.Variables.ssl_port])
+        AllMethodsTest._client_counter += 1
+
+        p.StartBefore(self._dns)
+        p.StartBefore(self._server)
+        p.StartBefore(self._ts)
+        tr.StillRunningAfter = self._ts
+
+
+# --- All Methods Test --------------------------------------------------------
+
+# Multiple methods loaded simultaneously, verifying shared context map works.
+AllMethodsTest('Multiple methods loaded simultaneously')