From c627b68d9a92b88b7b6beac18fd413f098ae668a Mon Sep 17 00:00:00 2001
From: JohnMcPMS <johnmcp@microsoft.com>
Date: Fri, 2 Jun 2023 14:39:06 -0700
Subject: [PATCH] Don't read past the end of the HTTP stream

---
 src/AppInstallerCLITests/Downloader.cpp       | 36 +++++++++++++++++++
 .../HttpStream/HttpClientWrapper.h            |  4 +--
 .../HttpStream/HttpLocalCache.cpp             |  5 ++-
 .../HttpStream/HttpLocalCache.h               |  4 +--
 .../Public/AppInstallerDownloader.h           |  1 +
 5 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/AppInstallerCLITests/Downloader.cpp b/src/AppInstallerCLITests/Downloader.cpp
index f53b2bc54a..95f8337818 100644
--- a/src/AppInstallerCLITests/Downloader.cpp
+++ b/src/AppInstallerCLITests/Downloader.cpp
@@ -4,6 +4,7 @@
 #include "TestCommon.h"
 #include "AppInstallerDownloader.h"
 #include "AppInstallerSHA256.h"
+#include "HttpStream/HttpLocalCache.h"
 
 using namespace AppInstaller;
 using namespace AppInstaller::Utility;
@@ -68,3 +69,38 @@ TEST_CASE("DownloadInvalidUrl", "[Downloader]")
 
     REQUIRE_THROWS(Download("blargle-flargle-fluff", tempFile.GetPath(), DownloadType::Installer, callback, true));
 }
+
+TEST_CASE("HttpStream_ReadLastFullPage", "[HttpStream]")
+{
+    Microsoft::WRL::ComPtr<IStream> stream;
+    STATSTG stat = { 0 };
+
+    for (size_t i = 0; i < 10; ++i)
+    {
+        stream = GetReadOnlyStreamFromURI("https://cdn.winget.microsoft.com/cache/source.msix");
+
+        stat = { 0 };
+        REQUIRE(stream->Stat(&stat, STATFLAG_NONAME) == S_OK);
+
+        if (stat.cbSize.QuadPart > 0)
+        {
+            break;
+        }
+
+        Sleep(500);
+    }
+
+    {
+        INFO("https://cdn.winget.microsoft.com/cache/source.msix gave back a 0 byte file");
+        REQUIRE(stream);
+    }
+
+    LARGE_INTEGER seek;
+    seek.QuadPart = (stat.cbSize.QuadPart / HttpStream::HttpLocalCache::PAGE_SIZE) * HttpStream::HttpLocalCache::PAGE_SIZE;
+    REQUIRE(stream->Seek(seek, STREAM_SEEK_SET, nullptr) == S_OK);
+
+    std::unique_ptr<BYTE[]> buffer = std::make_unique<BYTE[]>(HttpStream::HttpLocalCache::PAGE_SIZE);
+    ULONG read = 0;
+    REQUIRE(stream->Read(buffer.get(), static_cast<ULONG>(HttpStream::HttpLocalCache::PAGE_SIZE), &read) >= S_OK);
+    REQUIRE(read == (stat.cbSize.QuadPart % HttpStream::HttpLocalCache::PAGE_SIZE));
+}
diff --git a/src/AppInstallerCommonCore/HttpStream/HttpClientWrapper.h b/src/AppInstallerCommonCore/HttpStream/HttpClientWrapper.h
index c300839259..599c171224 100644
--- a/src/AppInstallerCommonCore/HttpStream/HttpClientWrapper.h
+++ b/src/AppInstallerCommonCore/HttpStream/HttpClientWrapper.h
@@ -1,8 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
-
 #pragma once
-
+#include <winrt/Windows.Foundation.h>
+#include <winrt/Windows.Web.Http.h>
 
 namespace AppInstaller::Utility::HttpStream
 {
diff --git a/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.cpp b/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.cpp
index 6a3150df35..55cc639792 100644
--- a/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.cpp
+++ b/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.cpp
@@ -217,6 +217,9 @@ namespace AppInstaller::Utility::HttpStream
         UINT32 trimStartIndex,
         UINT32 size)
     {
+        uint32_t bufferLength = originalBuffer.Length();
+        THROW_HR_IF(E_INVALIDARG, trimStartIndex > bufferLength);
+
         originalBuffer.as<::IInspectable>();
 
         // Get the byte array from the IBuffer object
@@ -228,7 +231,7 @@ namespace AppInstaller::Utility::HttpStream
 
         // Create the array of bytes holding the trimmed bytes
         IBuffer trimmedBuffer = CryptographicBuffer::CreateFromByteArray(
-            { byteBuffer + trimStartIndex, byteBuffer + trimStartIndex + size });
+            { byteBuffer + trimStartIndex, std::min(size, bufferLength - trimStartIndex) });
 
         return trimmedBuffer;
     }
diff --git a/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.h b/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.h
index 9484f36a52..e7fab8a318 100644
--- a/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.h
+++ b/src/AppInstallerCommonCore/HttpStream/HttpLocalCache.h
@@ -18,8 +18,8 @@ namespace AppInstaller::Utility::HttpStream
     class HttpLocalCache
     {
     public:
-        const UINT32 PAGE_SIZE = 2 << 16;   // each entry in the cache is 64 KB
-        const UINT32 MAX_PAGES = 200;       // cache size capped at 12.5 MB (200 * 64KB)
+        static constexpr UINT32 PAGE_SIZE = 2 << 16;   // each entry in the cache is 64 KB
+        static constexpr UINT32 MAX_PAGES = 200;       // cache size capped at 12.5 MB (200 * 64KB)
 
         // Returns a buffer matching the requested range by reading the parts of the range that are cached
         // and downloading the rest using the provided httpClientWrapper object
diff --git a/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h b/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h
index 65cf9d78b4..9b96d18504 100644
--- a/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h
+++ b/src/AppInstallerCommonCore/Public/AppInstallerDownloader.h
@@ -4,6 +4,7 @@
 #include <AppInstallerProgress.h>
 
 #include <urlmon.h>
+#include <wrl/client.h>
 
 #include <filesystem>
 #include <optional>