WSLA: Add initial wsladiag tool with basic --list command and placeholder ListSessions

CMakeLists.txt

@@ -431,6 +431,7 @@ add_subdirectory(src/windows/wslhost)
 add_subdirectory(src/windows/wslrelay)
 add_subdirectory(src/windows/wslinstall)
 add_subdirectory(src/windows/wslaclient)
+add_subdirectory(src/windows/wsladiag)
 
 if (WSL_BUILD_WSL_SETTINGS)
     add_subdirectory(src/windows/libwsl)

msipackage/CMakeLists.txt

@@ -3,7 +3,7 @@ set(OUTPUT_PACKAGE ${BIN}/wsl.msi)
 set(PACKAGE_WIX_IN ${CMAKE_CURRENT_LIST_DIR}/package.wix.in)
 set(PACKAGE_WIX ${BIN}/package.wix)
 set(CAB_CACHE ${BIN}/cab)
-set(BINARIES wsl.exe;wslg.exe;wslhost.exe;wslrelay.exe;wslservice.exe;wslserviceproxystub.dll;init;initrd.img;wslinstall.dll;wslaserviceproxystub.dll;wslaservice.exe)
+set(BINARIES wsl.exe;wslg.exe;wslhost.exe;wslrelay.exe;wslservice.exe;wslserviceproxystub.dll;init;initrd.img;wslinstall.dll;wslaserviceproxystub.dll;wslaservice.exe;wsladiag.exe)
 
 if (WSL_BUILD_WSL_SETTINGS)
     list(APPEND BINARIES_DEPENDENCIES "wslsettings/wslsettings.dll;wslsettings/wslsettings.exe;libwsl.dll")
@@ -39,7 +39,7 @@ add_custom_command(
 
 add_custom_target(msipackage DEPENDS ${OUTPUT_PACKAGE})
 set_target_properties(msipackage PROPERTIES EXCLUDE_FROM_ALL FALSE SOURCES ${PACKAGE_WIX_IN})
-add_dependencies(msipackage wsl wslg wslservice wslhost wslrelay wslserviceproxystub init initramfs wslinstall msixgluepackage wslaservice wslaserviceproxystub)
+add_dependencies(msipackage wsl wslg wslservice wslhost wslrelay wslserviceproxystub init initramfs wslinstall msixgluepackage wslaservice wslaserviceproxystub wsladiag)
 
 if (WSL_BUILD_WSL_SETTINGS)
     add_dependencies(msipackage wslsettings libwsl)

msipackage/package.wix.in

@@ -27,6 +27,7 @@
                 </File>
 
                 <File Id="wslg.exe" Name="wslg.exe" Source="${BIN}/wslg.exe" />
+                <File Id="wsladiag.exe" Name="wsladiag.exe" Source="${BIN}/wsladiag.exe" />
                 <File Id="wslhost.exe" Name="wslhost.exe" Source="${BIN}/wslhost.exe" />
                 <File Id="wslrelay.exe" Name="wslrelay.exe" Source="${BIN}/wslrelay.exe" />
                 <File Id="wslserviceproxystub.dll" Name="wslserviceproxystub.dll" Source="${BIN}/wslserviceproxystub.dll" />
@@ -378,7 +379,6 @@
                 </Component>
             </DirectoryRef>
         <?endif?>
-
         <DirectoryRef Id="TOOLSFOLDER">
             <Component Id="tools" Guid="F0C8D6BA-1502-41E7-BF72-D93DFA134733" Bitness="always64" UninstallWhenSuperseded="yes">
                 <File Id="init" Source="${BIN}/init" Checksum="yes" />

src/windows/common/WslTelemetry.cpp

@@ -33,7 +33,7 @@ TRACELOGGING_DEFINE_PROVIDER(
     TraceLoggingOptionMicrosoftTelemetry());
 
 TRACELOGGING_DEFINE_PROVIDER(
-    WslaServiceTelemetryProvider,
+    WslaTelemetryProvider,
     "Microsoft.Windows.Wsla",
     // {0383CE62-8F86-4766-AFB2-9D66A7FB1E90}
     (0x383ce62, 0x8f86, 0x4766, 0xaf, 0xb2, 0x9d, 0x66, 0xa7, 0xfb, 0x1e, 0x90),

src/windows/common/WslTelemetry.h

@@ -28,7 +28,7 @@ extern "C" {
 #endif
 TRACELOGGING_DECLARE_PROVIDER(LxssTelemetryProvider);
 TRACELOGGING_DECLARE_PROVIDER(WslServiceTelemetryProvider);
-TRACELOGGING_DECLARE_PROVIDER(WslaServiceTelemetryProvider);
+TRACELOGGING_DECLARE_PROVIDER(WslaTelemetryProvider);
 #ifdef __cplusplus
 }
 #endif

src/windows/wsladiag/CMakeLists.txt

@@ -0,0 +1,16 @@
+
+
+set(SOURCES
+    main.cpp
+)
+
+add_executable(wsladiag ${SOURCES})
+
+target_link_libraries(wsladiag
+    ${COMMON_LINK_LIBRARIES}
+    common
+)
+
+target_precompile_headers(wsladiag REUSE_FROM common)
+
+set_target_properties(wsladiag PROPERTIES FOLDER windows)

src/windows/wsladiag/main.cpp

@@ -0,0 +1,139 @@
+/*++
+
+Copyright (c) Microsoft. All rights reserved.
+
+Module Name:
+
+    main.cpp
+
+Abstract:
+
+    Entry point for the wsladiag tool, performs WSL runtime initialization and parses --list/--help.
+
+--*/
+
+#include "precomp.h"
+#include "CommandLine.h"
+#include "wslutil.h"
+#include "wslaservice.h"
+#include "WslSecurity.h"
+
+using namespace wsl::shared;
+namespace wslutil = wsl::windows::common::wslutil;
+
+int wsladiag_main(std::wstring_view commandLine)
+{
+    wslutil::ConfigureCrt();
+    wslutil::InitializeWil();
+
+    WslTraceLoggingInitialize(WslaTelemetryProvider, !wsl::shared::OfficialBuild);
+    auto cleanupTelemetry = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, []() { WslTraceLoggingUninitialize(); });
+
+    wslutil::SetCrtEncoding(_O_U8TEXT);
+
+    auto coInit = wil::CoInitializeEx(COINIT_MULTITHREADED);
+    wslutil::CoInitializeSecurity();
+
+    WSADATA data{};
+    THROW_IF_WIN32_ERROR(WSAStartup(MAKEWORD(2, 2), &data));
+    auto wsaCleanup = wil::scope_exit_log(WI_DIAGNOSTICS_INFO, []() { WSACleanup(); });
+
+    // Command-line parsing using ArgumentParser.
+    ArgumentParser parser(std::wstring{commandLine}, L"wsladiag");
+
+    bool help = false;
+    bool list = false;
+
+    parser.AddArgument(list, L"--list");
+    parser.AddArgument(help, L"--help", L'h'); //  short option is a single wide char
+    parser.Parse();
+
+    auto printUsage = []() {
+        wslutil::PrintMessage(
+            L"wsladiag - WSLA diagnostics tool\n"
+            L"Usage:\n"
+            L"  wsladiag --list    List WSLA sessions\n"
+            L"  wsladiag --help    Show this help",
+            stderr);
+    };
+
+    // If '--help' was requested, print usage and exit.
+    if (help)
+    {
+        printUsage();
+        return 0;
+    }
+
+    if (!list)
+    {
+        // No recognized command → show usage
+        printUsage();
+        return 0;
+    }
+
+    // --list: Call WSLA service COM interface to retrieve and display sessions.
+
+    try
+    {
+        wil::com_ptr<IWSLAUserSession> userSession;
+        THROW_IF_FAILED(CoCreateInstance(__uuidof(WSLAUserSession), nullptr, CLSCTX_LOCAL_SERVER, IID_PPV_ARGS(&userSession)));
+
+        wsl::windows::common::security::ConfigureForCOMImpersonation(userSession.get());
+
+        wil::unique_cotaskmem_array_ptr<WSLA_SESSION_INFORMATION> sessions;
+
+        THROW_IF_FAILED(userSession->ListSessions(&sessions, sessions.size_address<ULONG>()));
+
+        if (sessions.size() == 0)
+        {
+            wslutil::PrintMessage(L"No WSLA sessions found.\n", stdout);
+        }
+        else
+        {
+            wslutil::PrintMessage(std::format(L"Found {} WSLA session{}:\n", sessions.size(), sessions.size() > 1 ? L"s" : L""), stdout);
+
+            wslutil::PrintMessage(L"ID\tCreator PID\tDisplay Name\n", stdout);
+            wslutil::PrintMessage(L"--\t-----------\t------------\n", stdout);
+
+            for (const auto& session : sessions)
+            {
+                const auto * displayName = session.DisplayName;
+                if (displayName[0] == L'\0')
+                {
+                    displayName = L"<unnamed>";
+                }
+
+                wslutil::PrintMessage(
+                    std::format(L"{}\t{}\t\t{}\n", session.SessionId, session.CreatorPid, displayName), stdout);
+            }
+        }
+
+        return 0;
+    }
+    catch (...)
+    {
+        const auto hr = wil::ResultFromCaughtException();
+        const std::wstring hrMessage = wslutil::ErrorCodeToString(hr);
+
+        if (!hrMessage.empty())
+        {
+            wslutil::PrintMessage(std::format(L"Error listing WSLA sessions: 0x{:08x} - {}\n", static_cast<unsigned int>(hr), hrMessage), stderr);
+        }
+        else
+        {
+            wslutil::PrintMessage(std::format(L"Error listing WSLA sessions: 0x{:08x}\n", static_cast<unsigned int>(hr)), stderr);
+        }
+
+        return 1;
+    }
+}
+
+int wmain(int /*argc*/, wchar_t** /*argv*/)
+{
+    try
+    {
+        // Use raw Unicode command line so ArgumentParser gets original input.
+        return wsladiag_main(GetCommandLineW());
+    }
+    CATCH_RETURN();
+}

src/windows/wslaservice/exe/ServiceMain.cpp

@@ -67,7 +67,7 @@ try
 
     // Initialize telemetry.
     // TODO-WSLA: Create a dedicated WSLA provider
-    WslTraceLoggingInitialize(WslaServiceTelemetryProvider, !wsl::shared::OfficialBuild);
+    WslTraceLoggingInitialize(WslaTelemetryProvider, !wsl::shared::OfficialBuild);
 
     WSL_LOG("Service starting", TraceLoggingLevel(WINEVENT_LEVEL_INFO));
 

src/windows/wslaservice/exe/WSLASession.cpp

@@ -20,7 +20,10 @@ Module Name:
 
 using wsl::windows::service::wsla::WSLASession;
 
-WSLASession::WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl, const VIRTUAL_MACHINE_SETTINGS& VmSettings) :
+WSLASession::WSLASession(ULONG id, const WSLA_SESSION_SETTINGS& Settings,
+                                   WSLAUserSessionImpl& userSessionImpl,
+                                   const VIRTUAL_MACHINE_SETTINGS& VmSettings) :
+    m_id(id),
     m_sessionSettings(Settings),
     m_userSession(&userSessionImpl),
     m_virtualMachine(wil::MakeOrThrow<WSLAVirtualMachine>(VmSettings, userSessionImpl.GetUserSid(), &userSessionImpl)),

src/windows/wslaservice/exe/WSLASession.h

@@ -23,8 +23,12 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     : public Microsoft::WRL::RuntimeClass<Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>, IWSLASession, IFastRundown>
 {
 public:
-    WSLASession(const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl, const VIRTUAL_MACHINE_SETTINGS& VmSettings);
-    ~WSLASession();
+    WSLASession(ULONG id, const WSLA_SESSION_SETTINGS& Settings, WSLAUserSessionImpl& userSessionImpl, const VIRTUAL_MACHINE_SETTINGS& VmSettings);
+     ~WSLASession();
+     ULONG GetId() const noexcept
+    {
+        return m_id;
+    }
 
     IFACEMETHOD(GetDisplayName)(LPWSTR* DisplayName) override;
 
@@ -51,6 +55,7 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     void OnUserSessionTerminating();
 
 private:
+    ULONG m_id = 0;
     WSLA_SESSION_SETTINGS m_sessionSettings; // TODO: Revisit to see if we should have session settings as a member or not
     WSLAUserSessionImpl* m_userSession = nullptr;
     Microsoft::WRL::ComPtr<WSLAVirtualMachine> m_virtualMachine;

src/windows/wslaservice/exe/WSLAUserSession.cpp

@@ -50,19 +50,51 @@ PSID WSLAUserSessionImpl::GetUserSid() const
 HRESULT wsl::windows::service::wsla::WSLAUserSessionImpl::CreateSession(
     const WSLA_SESSION_SETTINGS* Settings, const VIRTUAL_MACHINE_SETTINGS* VmSettings, IWSLASession** WslaSession)
 {
-    auto session = wil::MakeOrThrow<WSLASession>(*Settings, *this, *VmSettings);
-
-    std::lock_guard lock(m_wslaSessionsLock);
-    auto it = m_sessions.emplace(session.Get());
+    ULONG id = m_nextSessionId++;
+    auto session = wil::MakeOrThrow<WSLASession>(id, *Settings, *this, *VmSettings);
 
+    {
+        std::lock_guard lock(m_wslaSessionsLock);
+        auto it = m_sessions.emplace(session.Get());
+        m_wslaSessions.emplace_back(session);
     // Client now owns the session.
     // TODO: Add a flag for the client to specify that the session should outlive its process.
-
+    }
+
     THROW_IF_FAILED(session.CopyTo(__uuidof(IWSLASession), (void**)WslaSession));
 
     return S_OK;
 }
 
+HRESULT wsl::windows::service::wsla::WSLAUserSessionImpl::ListSessions(_Out_ WSLA_SESSION_INFORMATION** Sessions, _Out_ ULONG* SessionsCount)
+{
+    const auto count = m_wslaSessions.size();
+    auto output = wil::make_unique_cotaskmem<WSLA_SESSION_INFORMATION[]>(m_wslaSessions.size());
+    std::lock_guard lock(m_wslaSessionsLock);
+    for (size_t i = 0; i < m_wslaSessions.size(); ++i)
+    {
+        output[i].SessionId = m_wslaSessions[i]->GetId();
+        output[i].CreatorPid = 0; // placeholder until we populate this later
+        PWSTR tempName = nullptr;
+
+        RETURN_IF_FAILED(m_wslaSessions[i]->GetDisplayName(&tempName));
+
+        if (tempName)
+        {
+            wcscpy_s(output[i].DisplayName, tempName);
+            CoTaskMemFree(tempName);
+        }
+        else
+        {
+            output[i].DisplayName[0] = L'\0';
+        }
+
+    }
+    *Sessions = output.release();
+    *SessionsCount = static_cast<ULONG>(count);
+    return S_OK;
+}
+
 wsl::windows::service::wsla::WSLAUserSession::WSLAUserSession(std::weak_ptr<WSLAUserSessionImpl>&& Session) :
     m_session(std::move(Session))
 {
@@ -89,9 +121,21 @@ try
 CATCH_RETURN();
 
 HRESULT wsl::windows::service::wsla::WSLAUserSession::ListSessions(WSLA_SESSION_INFORMATION** Sessions, ULONG* SessionsCount)
+try
 {
-    return E_NOTIMPL;
+    if (!Sessions || !SessionsCount)
+    {
+        return E_INVALIDARG;
+    }
+
+    auto session = m_session.lock();
+    RETURN_HR_IF(RPC_E_DISCONNECTED, !session);
+
+    RETURN_IF_FAILED(session->ListSessions(Sessions, SessionsCount));
+    return S_OK;
 }
+CATCH_RETURN();
+
 HRESULT wsl::windows::service::wsla::WSLAUserSession::OpenSession(ULONG Id, IWSLASession** Session)
 {
     return E_NOTIMPL;

src/windows/wslaservice/exe/WSLAUserSession.h

@@ -15,6 +15,8 @@ Module Name:
 #pragma once
 #include "WSLAVirtualMachine.h"
 #include "WSLASession.h"
+#include <atomic>
+#include <vector>
 
 namespace wsl::windows::service::wsla {
 
@@ -30,14 +32,17 @@ class WSLAUserSessionImpl
     PSID GetUserSid() const;
 
     HRESULT CreateSession(const WSLA_SESSION_SETTINGS* Settings, const VIRTUAL_MACHINE_SETTINGS* VmSettings, IWSLASession** WslaSession);
-
+    HRESULT ListSessions(_Out_ WSLA_SESSION_INFORMATION** Sessions, _Out_ ULONG* SessionsCount);
+    void OnVmTerminated(WSLAVirtualMachine* machine);
     void OnSessionTerminated(WSLASession* Session);
 
 private:
     wil::unique_tokeninfo_ptr<TOKEN_USER> m_tokenInfo;
-
+    std::atomic<ULONG> m_nextSessionId{1};
     std::recursive_mutex m_wslaSessionsLock;
     std::recursive_mutex m_lock;
+    // Track active sessions for diagnostics / ListSessions.
+    std::vector<Microsoft::WRL::ComPtr<WSLASession>> m_wslaSessions;
 
     // TODO-WSLA: Consider using a weak_ptr to easily destroy when the last client reference is released.
     std::unordered_set<WSLASession*> m_sessions;

src/windows/wslaservice/exe/WSLAUserSessionFactory.cpp

@@ -50,11 +50,15 @@ HRESULT WSLAUserSessionFactory::CreateInstance(_In_ IUnknown* pUnkOuter, _In_ RE
         THROW_HR_IF(CO_E_SERVER_STOPPING, !g_sessions.has_value());
 
         auto session = std::find_if(g_sessions->begin(), g_sessions->end(), [&tokenInfo](auto it) {
-            return EqualSid(it->GetUserSid(), &tokenInfo->User.Sid);
+            return EqualSid(it->GetUserSid(), tokenInfo->User.Sid);
         });
 
         if (session == g_sessions->end())
         {
+            wil::unique_hlocal_string sid;
+            THROW_IF_WIN32_BOOL_FALSE(ConvertSidToStringSid(tokenInfo->User.Sid, &sid));
+            WSL_LOG("WSLAUserSession created", TraceLoggingValue(sid.get(), "sid"));
+
             session = g_sessions->insert(g_sessions->end(), std::make_shared<WSLAUserSessionImpl>(userToken.get(), std::move(tokenInfo)));
         }