Fix missing instance ticket token JNI binding

lib/api/AuthTokensController.cpp

@@ -22,10 +22,17 @@ namespace MAT_NS_BEGIN {
         LOG_TRACE("destructor");
     }
 
+    std::mutex& AuthTokensController::GetMutex()
+    {
+        static std::mutex mutex;
+        return mutex;
+    }
+
     status_t  AuthTokensController::SetTicketToken(TicketType type, char const* tokenValue)
     {
         if (nullptr != tokenValue)
         {
+            std::lock_guard<std::mutex> lock(GetMutex());
             if (type == TicketType::TicketType_MSA_Device ||
                 type == TicketType::TicketType_AAD ||
                 type == TicketType::TicketType_XAuth_Device ||
@@ -35,7 +42,10 @@ namespace MAT_NS_BEGIN {
             }
             else
             {
-                m_tickets.push_back(TICKETS_PREPEND_STRING + std::to_string(type));
+                if (m_userTokens.find(type) == m_userTokens.end())
+                {
+                    m_tickets.push_back(TICKETS_PREPEND_STRING + std::to_string(type));
+                }
                 m_userTokens[type] = std::string(tokenValue);
             }
             return STATUS_SUCCESS;
@@ -45,6 +55,7 @@ namespace MAT_NS_BEGIN {
 
     status_t  AuthTokensController::Clear()
     {
+        std::lock_guard<std::mutex> lock(GetMutex());
         m_deviceTokens.clear();
         m_userTokens.clear();
         m_tickets.clear();
@@ -53,12 +64,14 @@ namespace MAT_NS_BEGIN {
 
     status_t  AuthTokensController::SetStrictMode(bool value)
     {
+        std::lock_guard<std::mutex> lock(GetMutex());
         m_IsStrictModeEnabled = value;
         return STATUS_EFAIL;
     }
 
     bool  AuthTokensController::GetStrictMode()
     {
+        std::lock_guard<std::mutex> lock(GetMutex());
         return m_IsStrictModeEnabled;
     }
 

lib/api/AuthTokensController.hpp

@@ -11,6 +11,7 @@
 #include "IAuthTokensController.hpp"
 
 #include <map>
+#include <mutex>
 #include <vector>
 
 // *INDENT-OFF*
@@ -24,6 +25,11 @@ namespace MAT_NS_BEGIN
     public:
 
         AuthTokensController();
+
+        // Hold this lock throughout SDK reads of the reference-returning public getters.
+        // Sharing it across controllers preserves the public interface and object layout.
+        static std::mutex& GetMutex();
+
         /// <summary>
         /// Destroy the telemetry logging system instance. Calls `FlushAndTeardown()` implicitly.
         /// </summary>

lib/decorators/BaseDecorator.cpp

@@ -2,8 +2,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 //
-#include "mat/config.h"
 #include "BaseDecorator.hpp"
+#include "api/AuthTokensController.hpp"
+#include "mat/config.h"
 
 namespace MAT_NS_BEGIN {
 
@@ -54,6 +55,7 @@ namespace MAT_NS_BEGIN {
         }
 
         //set Tickets
+        std::lock_guard<std::mutex> lock(AuthTokensController::GetMutex());
         if ((m_owner.GetAuthTokensController()) &&
             (m_owner.GetAuthTokensController()->GetTickets().size() > 0))
         {

lib/http/HttpRequestEncoder.cpp

@@ -4,8 +4,9 @@
 //
 
 #include "HttpRequestEncoder.hpp"
-#include "utils/StringUtils.hpp"
+#include "api/AuthTokensController.hpp"
 #include "pal/PAL.hpp"
+#include "utils/StringUtils.hpp"
 
 #include <memory>
 #include <string>
@@ -55,7 +56,7 @@ namespace MAT_NS_BEGIN {
         ctx->httpRequest->GetHeaders().set("Content-Type", "application/bond-compact-binary");
         ctx->httpRequest->GetHeaders().set("Upload-Time", toString(PAL::getUtcSystemTimeMs()));
 
-
+        std::unique_lock<std::mutex> tokensLock(AuthTokensController::GetMutex());
         if (GetAuthTokensController() != nullptr && GetAuthTokensController()->GetDeviceTokens().size() > 0)
         {
             std::map<TicketType, std::string>& map = GetAuthTokensController()->GetDeviceTokens();
@@ -134,6 +135,8 @@ namespace MAT_NS_BEGIN {
                 ctx->httpRequest->GetHeaders().set("Tickets", ticketHeader);
             }
         }
+        tokensLock.unlock();
+
         //strict mode
         if (GetAuthTokensController() != nullptr && true == GetAuthTokensController()->GetStrictMode())
         {

lib/include/public/IAuthTokensController.hpp

@@ -15,6 +15,10 @@ namespace MAT_NS_BEGIN
     /// <summary>
     /// This class is used to manage the Events  logging system
     /// </summary>
+    /// <remarks>
+    /// The reference-returning getters expose live state. Application code must not
+    /// access these references concurrently with token updates or Clear().
+    /// </remarks>
     class MATSDK_LIBABI IAuthTokensController
     {
     public:

lib/jni/LogManager_jni.cpp

@@ -1242,6 +1242,36 @@ Java_com_microsoft_applications_events_LogManagerProvider_00024LogManagerImpl_na
     }
     return logManager->ResumeTransmission();
 }
+
+extern "C" JNIEXPORT jint JNICALL
+Java_com_microsoft_applications_events_LogManagerProvider_00024LogManagerImpl_nativeSetIntTicketToken(
+    JNIEnv* env,
+    jobject /* thiz */,
+    jlong nativeLogManager,
+    jint type,
+    jstring tokenValue)
+{
+    auto logManager = getLogManager(nativeLogManager);
+    if (!logManager)
+    {
+        return STATUS_EFAIL;
+    }
+
+    auto ticketValue = JStringToStdString(env, tokenValue);
+    if (env->ExceptionCheck())
+    {
+        return STATUS_EFAIL;
+    }
+
+    auto controller = logManager->GetAuthTokensController();
+    if (!controller)
+    {
+        return STATUS_EFAIL;
+    }
+    return static_cast<jint>(controller->SetTicketToken(
+        static_cast<TicketType>(type), ticketValue.c_str()));
+}
+
 extern "C" JNIEXPORT jint JNICALL
 Java_com_microsoft_applications_events_LogManagerProvider_00024LogManagerImpl_nativeSetTransmitProfileTP(
     JNIEnv* env,

tests/unittests/HttpRequestEncoderTests.cpp

@@ -1,9 +1,14 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 
+#include "api/AuthTokensController.hpp"
 #include "common/Common.hpp"
 #include "common/MockIHttpClient.hpp"
-#include "http/HttpRequestEncoder.hpp"
+#include "common/MockITelemetrySystem.hpp"
 #include "config/RuntimeConfig_Default.hpp"
+#include "decorators/BaseDecorator.hpp"
+#include "http/HttpRequestEncoder.hpp"
+
+#include <future>
 
 using namespace testing;
 using namespace MAT;
@@ -126,3 +131,173 @@ TEST_F(HttpRequestEncoderTests, DispatchDataViewerEventCorrectly)
 
     EXPECT_THAT(mockEncoder.dataPacket, Eq(std::vector<uint8_t>{1, 127, 255}));
 }
+
+namespace
+{
+    class TokenLogManager : public NullLogManager
+    {
+       public:
+        AuthTokensController tokens;
+
+        IAuthTokensController* GetAuthTokensController() override
+        {
+            return &tokens;
+        }
+    };
+
+    class TokenTelemetrySystem : public MockITelemetrySystem
+    {
+       public:
+        TokenLogManager manager;
+
+        ILogManager& getLogManager() override
+        {
+            return manager;
+        }
+    };
+}
+
+TEST(AuthTokensControllerTests, RefreshingUserTokensDoesNotDuplicateTicketKeys)
+{
+    TokenLogManager manager;
+    auto& tokens = manager.tokens;
+    const TicketType userTypes[] = {TicketType_MSA_User, TicketType_XAuth_User, TicketType_AAD_User};
+    for (int refresh = 0; refresh < 100; ++refresh)
+    {
+        const auto value = std::to_string(refresh);
+        for (auto type : userTypes)
+        {
+            ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(type, value.c_str()));
+        }
+        ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_AAD, value.c_str()));
+    }
+
+    const std::vector<std::string> expectedTickets{"10001", "10003", "10005"};
+    EXPECT_EQ(expectedTickets, tokens.GetTickets());
+    for (auto type : userTypes)
+    {
+        EXPECT_EQ("99", tokens.GetUserTokens().at(type));
+    }
+    EXPECT_EQ("99", tokens.GetDeviceTokens().at(TicketType_AAD));
+
+    BaseDecorator decorator(manager);
+    CsProtocol::Record record;
+    ASSERT_TRUE(decorator.decorate(record));
+    ASSERT_EQ(1u, record.extProtocol.size());
+    ASSERT_EQ(1u, record.extProtocol[0].ticketKeys.size());
+    EXPECT_EQ(expectedTickets, record.extProtocol[0].ticketKeys[0]);
+
+    ASSERT_EQ(STATUS_SUCCESS, tokens.Clear());
+    EXPECT_TRUE(tokens.GetTickets().empty());
+    EXPECT_TRUE(tokens.GetUserTokens().empty());
+    EXPECT_TRUE(tokens.GetDeviceTokens().empty());
+    ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_MSA_User, "after-clear"));
+    ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_MSA_User, "refreshed"));
+    EXPECT_THAT(tokens.GetTickets(), ElementsAre("10001"));
+    EXPECT_EQ("refreshed", tokens.GetUserTokens().at(TicketType_MSA_User));
+}
+
+TEST_F(HttpRequestEncoderTests, TokenRefreshIsSafeDuringEncodingAndDecoration)
+{
+    TokenTelemetrySystem tokenSystem;
+    auto& tokens = tokenSystem.manager.tokens;
+    const std::string firstValue(256, 'a');
+    const std::string secondValue(256, 'b');
+    ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_MSA_User, firstValue.c_str()));
+    ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_XAuth_User, firstValue.c_str()));
+    ASSERT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_AAD, firstValue.c_str()));
+
+    HttpRequestEncoder tokenEncoder(tokenSystem, mockHttpClient);
+    BaseDecorator decorator(tokenSystem.manager);
+    std::promise<void> start;
+    auto ready = start.get_future().share();
+    auto firstWriter = std::async(std::launch::async, [&]()
+                                  {
+        ready.wait();
+        for (int refresh = 0; refresh < 2000; ++refresh)
+        {
+            const auto& value = refresh % 2 == 0 ? firstValue : secondValue;
+            EXPECT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_MSA_User, value.c_str()));
+            EXPECT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_AAD, value.c_str()));
+        } });
+    auto secondWriter = std::async(std::launch::async, [&]()
+                                   {
+        ready.wait();
+        for (int refresh = 0; refresh < 2000; ++refresh)
+        {
+            const auto& value = refresh % 2 == 0 ? firstValue : secondValue;
+            EXPECT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_XAuth_User, value.c_str()));
+        } });
+
+    const auto ticketHeader = [](const std::string& msa, const std::string& xauth)
+    {
+        return "\"10001\"=\"p:" + msa + "\";\"10003\"=\"x:XBL3.0 x=" + xauth + "\"";
+    };
+    start.set_value();
+    for (int event = 0; event < 1000; ++event)
+    {
+        auto ctx = std::make_shared<EventsUploadContext>();
+        tokenEncoder.encode(ctx);
+        const auto* request = static_cast<SimpleHttpRequest*>(ctx->httpRequest);
+        EXPECT_THAT(request->m_headers, Contains(Pair("Aad-Token", AnyOf(Eq(firstValue), Eq(secondValue)))));
+        EXPECT_THAT(request->m_headers, Contains(Pair("Tickets",
+                                                      AnyOf(Eq(ticketHeader(firstValue, firstValue)), Eq(ticketHeader(firstValue, secondValue)),
+                                                            Eq(ticketHeader(secondValue, firstValue)), Eq(ticketHeader(secondValue, secondValue))))));
+
+        CsProtocol::Record record;
+        ASSERT_TRUE(decorator.decorate(record));
+        ASSERT_EQ(1u, record.extProtocol.size());
+        ASSERT_EQ(1u, record.extProtocol[0].ticketKeys.size());
+        EXPECT_THAT(record.extProtocol[0].ticketKeys[0], ElementsAre("10001", "10003"));
+    }
+    firstWriter.get();
+    secondWriter.get();
+}
+
+TEST_F(HttpRequestEncoderTests, TokenClearIsSafeDuringEncodingAndDecoration)
+{
+    TokenTelemetrySystem tokenSystem;
+    auto& tokens = tokenSystem.manager.tokens;
+    HttpRequestEncoder tokenEncoder(tokenSystem, mockHttpClient);
+    BaseDecorator decorator(tokenSystem.manager);
+    const std::string value(256, 'a');
+    std::promise<void> start;
+    auto ready = start.get_future().share();
+    auto writer = std::async(std::launch::async, [&]()
+                             {
+        ready.wait();
+        for (int refresh = 0; refresh < 2000; ++refresh)
+        {
+            EXPECT_EQ(STATUS_SUCCESS, tokens.Clear());
+            EXPECT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_MSA_User, value.c_str()));
+            EXPECT_EQ(STATUS_SUCCESS, tokens.SetTicketToken(TicketType_AAD, value.c_str()));
+            tokens.SetStrictMode(refresh % 2 == 0);
+        } });
+
+    start.set_value();
+    for (int event = 0; event < 1000; ++event)
+    {
+        auto ctx = std::make_shared<EventsUploadContext>();
+        tokenEncoder.encode(ctx);
+        const auto* request = static_cast<SimpleHttpRequest*>(ctx->httpRequest);
+        auto deviceHeader = request->m_headers.find("Aad-Token");
+        if (deviceHeader != request->m_headers.end())
+        {
+            EXPECT_EQ(value, deviceHeader->second);
+        }
+        auto userHeader = request->m_headers.find("Tickets");
+        if (userHeader != request->m_headers.end())
+        {
+            EXPECT_EQ("\"10001\"=\"p:" + value + "\"", userHeader->second);
+        }
+
+        CsProtocol::Record record;
+        ASSERT_TRUE(decorator.decorate(record));
+        if (!record.extProtocol.empty())
+        {
+            ASSERT_EQ(1u, record.extProtocol[0].ticketKeys.size());
+            EXPECT_THAT(record.extProtocol[0].ticketKeys[0], ElementsAre("10001"));
+        }
+    }
+    writer.get();
+}