Refactor SaveImageImpl and ExportContainerImpl to eliminate code duplication

_codeql_detected_source_root

@@ -0,0 +1 @@
+.

src/windows/wslaservice/exe/DockerHTTPClient.cpp

@@ -80,6 +80,11 @@ std::vector<docker_schema::DeletedImage> wsl::windows::service::wsla::DockerHTTP
         std::format("http://localhost/images/{}?force={}&noprune={}", Image, Force ? "true" : "false", NoPrune ? "true" : "false"));
 }
 
+std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> DockerHTTPClient::SaveImage(const std::string& NameOrId)
+{
+    return SendRequestWithContext(verb::get, std::format("http://localhost/images/{}/get", NameOrId), {}, nullptr, {});
+}
+
 std::vector<docker_schema::ContainerInfo> DockerHTTPClient::ListContainers(bool all)
 {
     auto url = std::format("http://localhost/containers/json?all={}", all ? "true" : "false");
@@ -152,6 +157,11 @@ wil::unique_socket DockerHTTPClient::AttachContainer(const std::string& Id)
     return std::move(socket);
 }
 
+std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> DockerHTTPClient::ExportContainer(const std::string& ContainerNameOrID)
+{
+    return SendRequestWithContext(verb::get, std::format("http://localhost/containers/{}/export", ContainerNameOrID), {}, nullptr, {});
+}
+
 wil::unique_socket DockerHTTPClient::ContainerLogs(const std::string& Id, WSLALogsFlags Flags, ULONGLONG Since, ULONGLONG Until, ULONGLONG Tail)
 {
     auto url = std::format(
@@ -409,12 +419,13 @@ std::unique_ptr<DockerHTTPClient::HTTPRequestContext> DockerHTTPClient::SendRequ
     return std::move(context);
 }
 
-std::pair<uint32_t, wil::unique_socket> DockerHTTPClient::SendRequest(
+std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> DockerHTTPClient::SendRequestWithContext(
     verb Method,
     const std::string& Url,
     const std::string& Body,
     const OnResponseBytes& OnResponse,
-    const std::map<boost::beast::http::field, std::string>& Headers)
+    const std::map<boost::beast::http::field, std::string>& Headers,
+    std::string* errorJson)
 {
     // Write the request
     auto context = SendRequestImpl(Method, Url, Body, Headers);
@@ -501,5 +512,16 @@ std::pair<uint32_t, wil::unique_socket> DockerHTTPClient::SendRequest(
         }
     }
 
-    return {parser.get().result_int(), wil::unique_socket{context->stream.release()}};
+    return {parser.get().result_int(), std::move(context)};
 }
+
+std::pair<uint32_t, wil::unique_socket> DockerHTTPClient::SendRequest(
+    verb Method,
+    const std::string& Url,
+    const std::string& Body,
+    const OnResponseBytes& OnResponse,
+    const std::map<boost::beast::http::field, std::string>& Headers)
+{
+    auto result = SendRequestWithContext(Method, Url, Body, OnResponse, Headers);
+    return {result.first, wil::unique_socket{result.second->stream.release()}};
+}

src/windows/wslaservice/exe/DockerHTTPClient.h

@@ -89,6 +89,7 @@ class DockerHTTPClient
     wil::unique_socket AttachContainer(const std::string& Id);
     void ResizeContainerTty(const std::string& Id, ULONG Rows, ULONG Columns);
     wil::unique_socket ContainerLogs(const std::string& Id, WSLALogsFlags Flags, ULONGLONG Since, ULONGLONG Until, ULONGLONG Tail);
+    std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> ExportContainer(const std::string& ContainerID);
 
     // Image management.
     std::unique_ptr<HTTPRequestContext> PullImage(const char* Name, const char* Tag);
@@ -97,6 +98,7 @@ class DockerHTTPClient
     void TagImage(const std::string& Id, const std::string& Repo, const std::string& Tag);
     std::vector<common::docker_schema::Image> ListImages();
     std::vector<common::docker_schema::DeletedImage> DeleteImage(const char* Image, bool Force, bool NoPrune); // Image can be ID or Repo:Tag.
+    std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> SaveImage(const std::string& NameOrId);
 
     // Exec.
     common::docker_schema::CreateExecResponse CreateExec(const std::string& Container, const common::docker_schema::CreateExec& Request);
@@ -148,6 +150,14 @@ class DockerHTTPClient
         const OnResponseBytes& OnResponse,
         const std::map<boost::beast::http::field, std::string>& Headers = {});
 
+    std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>> SendRequestWithContext(
+        boost::beast::http::verb Method,
+        const std::string& Url,
+        const std::string& Body,
+        const OnResponseBytes& OnResponse,
+        const std::map<boost::beast::http::field, std::string>& Headers = {},
+        std::string* errorJson = nullptr); // Out parameter to receive error response body. If null, error body is left unread and caller is expected to handle reading the error message body
+
     template <typename TRequest = common::docker_schema::EmptyRequest, typename TResponse = TRequest::TResponse>
     auto Transaction(boost::beast::http::verb Method, const std::string& Url, const TRequest& RequestObject = {})
     {

src/windows/wslaservice/exe/WSLAContainer.cpp

@@ -758,6 +758,11 @@ void WSLAContainerImpl::Inspect(LPSTR* Output)
     }
 }
 
+void WSLAContainerImpl::GetID(LPSTR* Output)
+{
+    *Output = wil::make_unique_ansistring<wil::unique_cotaskmem_ansistring>(m_id.data()).release();
+}
+
 void WSLAContainerImpl::Logs(WSLALogsFlags Flags, ULONG* Stdout, ULONG* Stderr, ULONGLONG Since, ULONGLONG Until, ULONGLONG Tail)
 {
     std::lock_guard lock(m_lock);
@@ -847,6 +852,12 @@ HRESULT WSLAContainer::Inspect(LPSTR* Output)
     return CallImpl(&WSLAContainerImpl::Inspect, Output);
 }
 
+HRESULT WSLAContainer::GetID(LPSTR* Id)
+{
+    *Id = nullptr;
+    return CallImpl(&WSLAContainerImpl::GetID, Id);
+}
+
 HRESULT WSLAContainer::Delete()
 try
 {

src/windows/wslaservice/exe/WSLAContainer.h

@@ -74,6 +74,7 @@ class WSLAContainerImpl
     void GetInitProcess(_Out_ IWSLAProcess** process);
     void Exec(_In_ const WSLA_PROCESS_OPTIONS* Options, _Out_ IWSLAProcess** Process, _Out_ int* Errno);
     void Inspect(LPSTR* Output);
+    void GetID(LPSTR* Id);
     void Logs(WSLALogsFlags Flags, ULONG* Stdout, ULONG* Stderr, ULONGLONG Since, ULONGLONG Until, ULONGLONG Tail);
 
     IWSLAContainer& ComWrapper();
@@ -145,6 +146,7 @@ class DECLSPEC_UUID("B1F1C4E3-C225-4CAE-AD8A-34C004DE1AE4") WSLAContainer
     IFACEMETHOD(Exec)(_In_ const WSLA_PROCESS_OPTIONS* Options, _Out_ IWSLAProcess** Process, _Out_ int* Errno) override;
     IFACEMETHOD(Start)() override;
     IFACEMETHOD(Inspect)(_Out_ LPSTR* Output) override;
+    IFACEMETHOD(GetID)(_Out_ LPSTR* Output) override;
     IFACEMETHOD(Logs)(_In_ WSLALogsFlags Flags, _Out_ ULONG* Stdout, _Out_ ULONG* Stderr, _In_ ULONGLONG Since, _In_ ULONGLONG Until, _In_ ULONGLONG Tail) override;
     IFACEMETHOD(GetId)(_Out_ WSLAContainerId Id) override;
     IFACEMETHOD(GetName)(_Out_ LPSTR* Name) override;

src/windows/wslaservice/exe/WSLASession.cpp

@@ -57,6 +57,68 @@ bool IsContainerNameValid(LPCSTR Name)
     return length > 0 && length <= WSLA_MAX_CONTAINER_NAME_LENGTH;
 }
 
+/// <summary>
+/// Handles image-like operations (save image, export container) by managing the HTTP response
+/// and relaying data to an output handle.
+/// </summary>
+/// <param name="RequestCodePair">The HTTP response code and request context from the Docker client</param>
+/// <param name="OutputHandle">Handle to the output file where data will be written</param>
+/// <param name="SessionTerminatingEvent">Event handle to monitor for session termination</param>
+/// <param name="ErrorInfo">Optional pointer to receive error information</param>
+/// <param name="FailureMessageFormat">Format string for error messages (must contain %hs for error text)</param>
+/// <param name="OnCompletedCallback">Optional callback invoked when data relay completes (not when entire operation succeeds)</param>
+/// <exception cref="HRESULT">Throws E_FAIL if the HTTP response code is not 200, or E_ABORT if session terminates</exception>
+void HandleImageLikeOperation(
+    std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>>& RequestCodePair,
+    ULONG OutputHandle,
+    HANDLE SessionTerminatingEvent,
+    WSLA_ERROR_INFO* ErrorInfo,
+    const char* FailureMessageFormat,
+    std::function<void()> OnCompletedCallback = nullptr)
+{
+    wil::unique_handle imageFileHandle{wsl::windows::common::wslutil::DuplicateHandleFromCallingProcess(ULongToHandle(OutputHandle))};
+
+    relay::MultiHandleWait io;
+    auto onCompleted = [&, OnCompletedCallback]() {
+        io.Cancel();
+        if (OnCompletedCallback)
+        {
+            OnCompletedCallback();
+        }
+    };
+    std::string errorJson;
+    auto accumulateError = [&](const gsl::span<char>& buffer) {
+        // If the operation failed, accumulate the error message.
+        errorJson.append(buffer.data(), buffer.size());
+    };
+
+    if (RequestCodePair.first != 200)
+    {
+        io.AddHandle(std::make_unique<relay::ReadHandle>(
+            common::relay::HandleWrapper{RequestCodePair.second->stream.native_handle()}, accumulateError));
+    }
+    else
+    {
+        io.AddHandle(std::make_unique<relay::RelayHandle<relay::HTTPChunkBasedReadHandle>>(
+            common::relay::HandleWrapper{RequestCodePair.second->stream.native_handle()},
+            common::relay::HandleWrapper{std::move(imageFileHandle), std::move(onCompleted)}));
+        io.AddHandle(std::make_unique<relay::EventHandle>(SessionTerminatingEvent, [&]() { THROW_HR(E_ABORT); }));
+    }
+
+    io.Run({});
+
+    if (RequestCodePair.first != 200)
+    {
+        // Operation failed, parse the error message.
+        auto error = wsl::shared::FromJson<docker_schema::ErrorResponse>(errorJson.c_str());
+        if (ErrorInfo != nullptr)
+        {
+            ErrorInfo->UserErrorMessage = wil::make_unique_ansistring<wil::unique_cotaskmem_ansistring>(error.message.c_str()).release();
+        }
+        THROW_HR_MSG(E_FAIL, FailureMessageFormat, error.message.c_str());
+    }
+}
+
 } // namespace
 
 WSLASession::WSLASession(ULONG id, const WSLA_SESSION_SETTINGS& Settings, wil::unique_tokeninfo_ptr<TOKEN_USER>&& TokenInfo, bool Elevated) :
@@ -515,6 +577,47 @@ void WSLASession::ImportImageImpl(DockerHTTPClient::HTTPRequestContext& Request,
     }
 }
 
+HRESULT WSLASession::ExportContainer(ULONG OutHandle, LPCSTR ContainerID, IProgressCallback* ProgressCallback, WSLA_ERROR_INFO* Error)
+try
+{
+    UNREFERENCED_PARAMETER(ProgressCallback);
+    RETURN_HR_IF_NULL(E_POINTER, ContainerID);
+    std::lock_guard lock{m_lock};
+    THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_INVALID_STATE), !m_dockerClient.has_value());
+    auto retVal = m_dockerClient->ExportContainer(ContainerID);
+    ExportContainerImpl(retVal, OutHandle, Error);
+    return S_OK;
+}
+CATCH_RETURN();
+
+void WSLASession::ExportContainerImpl(
+    std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>>& RequestCodePair, ULONG OutputHandle, WSLA_ERROR_INFO* Error)
+{
+    auto onCompletedCallback = []() {
+        WSL_LOG("OnCompletedCalledForExport", TraceLoggingValue("OnCompletedCalledForExport", "Content"));
+    };
+
+    HandleImageLikeOperation(RequestCodePair, OutputHandle, m_sessionTerminatingEvent.get(), Error, "Container export failed: %hs", onCompletedCallback);
+}
+
+HRESULT WSLASession::SaveImage(ULONG OutHandle, LPCSTR ImageNameOrID, IProgressCallback* ProgressCallback, WSLA_ERROR_INFO* Error)
+try
+{
+    UNREFERENCED_PARAMETER(ProgressCallback);
+    RETURN_HR_IF_NULL(E_POINTER, ImageNameOrID);
+    std::lock_guard lock{m_lock};
+    THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_INVALID_STATE), !m_dockerClient.has_value());
+    auto retVal = m_dockerClient->SaveImage(ImageNameOrID);
+    SaveImageImpl(retVal, OutHandle, Error);
+    return S_OK;
+}
+CATCH_RETURN();
+
+void WSLASession::SaveImageImpl(std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>>& RequestCodePair, ULONG OutputHandle, WSLA_ERROR_INFO* Error)
+{
+    HandleImageLikeOperation(RequestCodePair, OutputHandle, m_sessionTerminatingEvent.get(), Error, "Image save failed: %hs");
+}
+
 HRESULT WSLASession::ListImages(WSLA_IMAGE_INFORMATION** Images, ULONG* Count)
 try
 {

src/windows/wslaservice/exe/WSLASession.h

@@ -64,6 +64,7 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
         _Inout_opt_ WSLA_ERROR_INFO* ErrorInfo) override;
     IFACEMETHOD(LoadImage)(_In_ ULONG ImageHandle, _In_ IProgressCallback* ProgressCallback, _In_ ULONGLONG ContentLength) override;
     IFACEMETHOD(ImportImage)(_In_ ULONG ImageHandle, _In_ LPCSTR ImageName, _In_ IProgressCallback* ProgressCallback, _In_ ULONGLONG ContentLength) override;
+    IFACEMETHOD(SaveImage)(_In_ ULONG OutputHandle, _In_ LPCSTR ImageNameOrID, _In_ IProgressCallback* ProgressCallback, WSLA_ERROR_INFO* Error) override;
     IFACEMETHOD(ListImages)(_Out_ WSLA_IMAGE_INFORMATION** Images, _Out_ ULONG* Count) override;
     IFACEMETHOD(DeleteImage)(
         _In_ const WSLA_DELETE_IMAGE_OPTIONS* Options,
@@ -75,6 +76,7 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     IFACEMETHOD(CreateContainer)(_In_ const WSLA_CONTAINER_OPTIONS* Options, _Out_ IWSLAContainer** Container, _Inout_opt_ WSLA_ERROR_INFO* Error) override;
     IFACEMETHOD(OpenContainer)(_In_ LPCSTR Id, _In_ IWSLAContainer** Container) override;
     IFACEMETHOD(ListContainers)(_Out_ WSLA_CONTAINER** Images, _Out_ ULONG* Count) override;
+    IFACEMETHOD(ExportContainer)(_In_ ULONG OutputHandle, _In_ LPCSTR ContainerID, _In_ IProgressCallback* ProgressCallback, _Inout_opt_ WSLA_ERROR_INFO* Error) override;
 
     // VM management.
     IFACEMETHOD(CreateRootNamespaceProcess)(_In_ const WSLA_PROCESS_OPTIONS* Options, _Out_ IWSLAProcess** VirtualMachine, _Out_ int* Errno) override;
@@ -109,6 +111,9 @@ class DECLSPEC_UUID("4877FEFC-4977-4929-A958-9F36AA1892A4") WSLASession
     void ImportImageImpl(DockerHTTPClient::HTTPRequestContext& Request, ULONG InputHandle);
     void RecoverExistingContainers();
 
+    void SaveImageImpl(std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>>& RequestCodePair, ULONG OutputHandle, WSLA_ERROR_INFO* Error);
+    void ExportContainerImpl(std::pair<uint32_t, std::unique_ptr<DockerHTTPClient::HTTPRequestContext>>& RequestCodePair, ULONG OutputHandle, WSLA_ERROR_INFO* Error);
+
     std::optional<DockerHTTPClient> m_dockerClient;
     std::optional<WSLAVirtualMachine> m_virtualMachine;
     std::optional<ContainerEventTracker> m_eventTracker;

src/windows/wslaservice/inc/wslaservice.idl

@@ -302,6 +302,7 @@ interface IWSLAContainer : IUnknown
     HRESULT GetInitProcess([out] IWSLAProcess** Process);
     HRESULT Exec([in] const struct WSLA_PROCESS_OPTIONS* Options, [out] IWSLAProcess** Process, [out] int* Errno);
     HRESULT Inspect([out] LPSTR* Output);
+    HRESULT GetID([out] LPSTR * Id);
     HRESULT Logs([in] WSLALogsFlags Flags, [out] ULONG* Stdout, [out] ULONG* Stderr, [in] ULONGLONG Since, [in] ULONGLONG Until, [in] ULONGLONG Tail);
     HRESULT GetId([out, string] WSLAContainerId Id);
     HRESULT GetName([out, string] LPSTR* Name);
@@ -344,13 +345,15 @@ interface IWSLASession : IUnknown
     HRESULT PullImage([in] LPCSTR ImageUri, [in, unique] const struct WSLA_REGISTRY_AUTHENTICATION_INFORMATION* RegistryAuthenticationInformation, [in, unique] IProgressCallback* ProgressCallback, [in, out, unique] WSLA_ERROR_INFO* ErrorInfo);
     HRESULT LoadImage([in] ULONG ImageHandle, [in, unique] IProgressCallback* ProgressCallback, [in] ULONGLONG ContentLength);
     HRESULT ImportImage([in] ULONG ImageHandle, [in] LPCSTR ImageName, [in, unique] IProgressCallback* ProgressCallback, [in] ULONGLONG ContentLength);
+    HRESULT SaveImage([in] ULONG OutputHandle, [in] LPCSTR ImageNameOrID, [ in, unique ] IProgressCallback * ProgressCallback, [in, out, unique] WSLA_ERROR_INFO * ErrorInfo);
     HRESULT ListImages([out, size_is(, *Count)] struct WSLA_IMAGE_INFORMATION** Images, [out] ULONG* Count);
     HRESULT DeleteImage([in] const struct WSLA_DELETE_IMAGE_OPTIONS* Options, [out, size_is(, *Count)] struct WSLA_DELETED_IMAGE_INFORMATION** DeletedImages, [out] ULONG* Count, [in, out, unique] WSLA_ERROR_INFO* ErrorInfo);
 
     // Container management.
     HRESULT CreateContainer([in] const struct WSLA_CONTAINER_OPTIONS* Options, [out] IWSLAContainer** Container, [in, out, unique] WSLA_ERROR_INFO* ErrorInfo);
     HRESULT OpenContainer([in] LPCSTR Id, [out] IWSLAContainer** Container);
     HRESULT ListContainers([out, size_is(, *Count)] struct WSLA_CONTAINER** Images, [out] ULONG* Count);
+    HRESULT ExportContainer([in] ULONG OutputHandle, [in] LPCSTR ContainerID, [ in, unique ] IProgressCallback * ProgressCallback, [in,out,unique] WSLA_ERROR_INFO * ErrorInfo);
 
     // Create a process at the VM level. This is meant for debugging.
     HRESULT CreateRootNamespaceProcess([in] const struct WSLA_PROCESS_OPTIONS* Options, [out] IWSLAProcess** Process, [out] int* Errno);