winget-cli

commit 59d869ed1668b1595ede537cdf1048c66b614574
parent c0e74094c548d159c81ff98816010b2d5a010622
Author: yao-msft <50888816+yao-msft@users.noreply.github.com>
Date:   Tue, 21 Sep 2021 13:01:22 -0700

Remove motw before calling IAttachmentExecute::Save (#1491)


Diffstat:
M
src/AppInstallerCLICore/Workflows/InstallFlow.cpp
 | 
2
+-
M
src/AppInstallerCommonCore/Downloader.cpp
 | 
35
++++++++++++++++++++++++++++++++++-
M
src/AppInstallerCommonCore/Public/AppInstallerDownloader.h
 | 
6
+++++-

3 files changed, 40 insertions(+), 3 deletions(-)
diff --git a/src/AppInstallerCLICore/Workflows/InstallFlow.cpp b/src/AppInstallerCLICore/Workflows/InstallFlow.cpp
@@ -409,7 +409,7 @@ namespace AppInstaller::CLI::Workflow
             else if (WI_IsFlagSet(context.GetFlags(), Execution::ContextFlag::InstallerHashMatched))
             {
                 const auto& installer = context.Get<Execution::Data::Installer>();
-                HRESULT hr = Utility::ApplyMotwUsingIAttachmentExecuteIfApplicable(context.Get<Execution::Data::InstallerPath>(), installer.value().Url);
+                HRESULT hr = Utility::ApplyMotwUsingIAttachmentExecuteIfApplicable(context.Get<Execution::Data::InstallerPath>(), installer.value().Url, URLZONE_INTERNET);
 
                 // Not using SUCCEEDED(hr) to check since there are cases file is missing after a successful scan
                 if (hr != S_OK)
diff --git a/src/AppInstallerCommonCore/Downloader.cpp b/src/AppInstallerCommonCore/Downloader.cpp
@@ -263,7 +263,28 @@ namespace AppInstaller::Utility
         AICLI_LOG(Core, Info, << "Finished applying motw");
     }
 
-    HRESULT ApplyMotwUsingIAttachmentExecuteIfApplicable(const std::filesystem::path& filePath, const std::string& source)
+    void RemoveMotwIfApplicable(const std::filesystem::path& filePath)
+    {
+        AICLI_LOG(Core, Info, << "Started removing motw to " << filePath);
+
+        if (!IsNTFS(filePath))
+        {
+            AICLI_LOG(Core, Info, << "File system is not NTFS. Skipped removing motw");
+            return;
+        }
+
+        Microsoft::WRL::ComPtr<IZoneIdentifier> zoneIdentifier;
+        THROW_IF_FAILED(CoCreateInstance(CLSID_PersistentZoneIdentifier, nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&zoneIdentifier)));
+        THROW_IF_FAILED(zoneIdentifier->Remove());
+
+        Microsoft::WRL::ComPtr<IPersistFile> persistFile;
+        THROW_IF_FAILED(zoneIdentifier.As(&persistFile));
+        THROW_IF_FAILED(persistFile->Save(filePath.c_str(), TRUE));
+
+        AICLI_LOG(Core, Info, << "Finished removing motw");
+    }
+
+    HRESULT ApplyMotwUsingIAttachmentExecuteIfApplicable(const std::filesystem::path& filePath, const std::string& source, URLZONE zoneIfScanFailure)
     {
         AICLI_LOG(Core, Info, << "Started applying motw using IAttachmentExecute to " << filePath);
 
@@ -281,7 +302,19 @@ namespace AppInstaller::Utility
             RETURN_IF_FAILED(CoCreateInstance(CLSID_AttachmentServices, nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&attachmentExecute)));
             RETURN_IF_FAILED(attachmentExecute->SetLocalPath(filePath.c_str()));
             RETURN_IF_FAILED(attachmentExecute->SetSource(Utility::ConvertToUTF16(source).c_str()));
+
+            // IAttachmentExecute::Save() expects the local file to be clean(i.e. it won't clear existing motw if it thinks the source url is trusted)
+            RemoveMotwIfApplicable(filePath);
+
             aesSaveResult = attachmentExecute->Save();
+
+            // Reapply desired zone upon scan failure.
+            // Not using SUCCEEDED(hr) to check since there are cases file is missing after a successful scan
+            if (aesSaveResult != S_OK && std::filesystem::exists(filePath))
+            {
+                ApplyMotwIfApplicable(filePath, zoneIfScanFailure);
+            }
+
             RETURN_IF_FAILED(aesSaveResult);
             return S_OK;
         };
diff --git a/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h b/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h
@@ -66,8 +66,12 @@ namespace AppInstaller::Utility
     // Apply Mark of the web if the target file is on NTFS, otherwise does nothing.
     void ApplyMotwIfApplicable(const std::filesystem::path& filePath, URLZONE zone);
 
+    // Remove Mark of the web if the target file is on NTFS, otherwise does nothing.
+    void RemoveMotwIfApplicable(const std::filesystem::path& filePath);
+
     // Apply Mark of the web using IAttachmentExecute::Save if the target file is on NTFS, otherwise does nothing.
     // This method only does a best effort since Attachment Execution Service may be disabled.
     // If IAttachmentExecute::Save is successfully invoked and the scan failed, the failure HRESULT is returned.
-    HRESULT ApplyMotwUsingIAttachmentExecuteIfApplicable(const std::filesystem::path& filePath, const std::string& source);
+    // zoneIfScanFailure: URLZONE to apply if IAttachmentExecute::Save scan failed.
+    HRESULT ApplyMotwUsingIAttachmentExecuteIfApplicable(const std::filesystem::path& filePath, const std::string& source, URLZONE zoneIfScanFailure);
 }