Update WgcCaptureSession to handle portrait oriented screen capture.

WgcCaptureSession would crash when copying the frame data for an image
from a portrait oriented monitor. This is because we were using the
height of the image multiplied by the rowpitch of the buffer to
determine the size of the data to be copied. However, in portrait
mode the height measures the same dimension as the rowpitch, leading
to us overrunning the frame buffer.

The fix is to use the height and width of the image multiplied by
the number of bytes per pixel to determine how much data to copy
out of the buffer, and only use the rowpitch to advance the pointer
in the source data buffer. This has the added benefit of giving us
contiguous data, reducing the size of the DesktopFrame that we output.

Bug: webrtc:12490
Change-Id: I4c26f8864cb57ac566a742af70fea1da504b9706
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/209501
Reviewed-by: Joe Downing <joedow@chromium.org>
Commit-Queue: Austin Orion <auorion@microsoft.com>
Cr-Commit-Position: refs/heads/master@{#33532}
diff --git a/modules/desktop_capture/win/test_support/test_window.cc b/modules/desktop_capture/win/test_support/test_window.cc
index dc94ee0..d5fa9ed 100644
--- a/modules/desktop_capture/win/test_support/test_window.cc
+++ b/modules/desktop_capture/win/test_support/test_window.cc
@@ -17,6 +17,26 @@
 const int kWindowHeight = 200;
 const int kWindowWidth = 300;
 
+LRESULT CALLBACK WindowProc(HWND hwnd,
+                            UINT msg,
+                            WPARAM w_param,
+                            LPARAM l_param) {
+  switch (msg) {
+    case WM_PAINT:
+      PAINTSTRUCT paint_struct;
+      HDC hdc = BeginPaint(hwnd, &paint_struct);
+
+      // Paint the window so the color is consistent and we can inspect the
+      // pixels in tests and know what to expect.
+      FillRect(hdc, &paint_struct.rcPaint,
+               CreateSolidBrush(RGB(kTestWindowRValue, kTestWindowGValue,
+                                    kTestWindowBValue)));
+
+      EndPaint(hwnd, &paint_struct);
+  }
+  return DefWindowProc(hwnd, msg, w_param, l_param);
+}
+
 }  // namespace
 
 WindowInfo CreateTestWindow(const WCHAR* window_title,
@@ -25,7 +45,7 @@
   WindowInfo info;
   ::GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
                            GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
-                       reinterpret_cast<LPCWSTR>(&::DefWindowProc),
+                       reinterpret_cast<LPCWSTR>(&WindowProc),
                        &info.window_instance);
 
   WNDCLASSEXW wcex;
@@ -33,7 +53,7 @@
   wcex.cbSize = sizeof(wcex);
   wcex.style = CS_HREDRAW | CS_VREDRAW;
   wcex.hInstance = info.window_instance;
-  wcex.lpfnWndProc = &::DefWindowProc;
+  wcex.lpfnWndProc = &WindowProc;
   wcex.lpszClassName = kWindowClass;
   info.window_class = ::RegisterClassExW(&wcex);
 
diff --git a/modules/desktop_capture/win/test_support/test_window.h b/modules/desktop_capture/win/test_support/test_window.h
index a5962b5..7c7676c 100644
--- a/modules/desktop_capture/win/test_support/test_window.h
+++ b/modules/desktop_capture/win/test_support/test_window.h
@@ -17,6 +17,14 @@
 
 namespace webrtc {
 
+typedef unsigned char uint8_t;
+
+// Define an arbitrary color for the test window with unique R, G, and B values
+// so consumers can verify captured content in tests.
+const uint8_t kTestWindowRValue = 191;
+const uint8_t kTestWindowGValue = 99;
+const uint8_t kTestWindowBValue = 12;
+
 struct WindowInfo {
   HWND hwnd;
   HINSTANCE window_instance;
diff --git a/modules/desktop_capture/win/wgc_capture_session.cc b/modules/desktop_capture/win/wgc_capture_session.cc
index 0fdb6ec..7ff2f93 100644
--- a/modules/desktop_capture/win/wgc_capture_session.cc
+++ b/modules/desktop_capture/win/wgc_capture_session.cc
@@ -186,29 +186,36 @@
   if (FAILED(hr))

     return hr;

 

-  // If the size has changed since the last capture, we must be sure to choose

-  // the smaller of the two sizes. Otherwise we might overrun our buffer, or

+  // If the size has changed since the last capture, we must be sure to use

+  // the smaller dimensions. Otherwise we might overrun our buffer, or

   // read stale data from the last frame.

-  int previous_area = previous_size_.Width * previous_size_.Height;

-  int new_area = new_size.Width * new_size.Height;

-  auto smaller_size = previous_area < new_area ? previous_size_ : new_size;

+  int image_height = std::min(previous_size_.Height, new_size.Height);

+  int image_width = std::min(previous_size_.Width, new_size.Width);

+  int row_data_length = image_width * DesktopFrame::kBytesPerPixel;

 

   // Make a copy of the data pointed to by |map_info.pData| so we are free to

   // unmap our texture.

-  uint8_t* data = static_cast<uint8_t*>(map_info.pData);

-  int data_size = smaller_size.Height * map_info.RowPitch;

-  std::vector<uint8_t> image_data(data, data + data_size);

-  DesktopSize size(smaller_size.Width, smaller_size.Height);

+  uint8_t* src_data = static_cast<uint8_t*>(map_info.pData);

+  std::vector<uint8_t> image_data;

+  image_data.reserve(image_height * row_data_length);

+  uint8_t* image_data_ptr = image_data.data();

+  for (int i = 0; i < image_height; i++) {

+    memcpy(image_data_ptr, src_data, row_data_length);

+    image_data_ptr += row_data_length;

+    src_data += map_info.RowPitch;

+  }

 

   // Transfer ownership of |image_data| to the output_frame.

-  *output_frame = std::make_unique<WgcDesktopFrame>(

-      size, static_cast<int>(map_info.RowPitch), std::move(image_data));

+  DesktopSize size(image_width, image_height);

+  *output_frame = std::make_unique<WgcDesktopFrame>(size, row_data_length,

+                                                    std::move(image_data));

 

   d3d_context->Unmap(mapped_texture_.Get(), 0);

 

   // If the size changed, we must resize the texture and frame pool to fit the

   // new size.

-  if (previous_area != new_area) {

+  if (previous_size_.Height != new_size.Height ||

+      previous_size_.Width != new_size.Width) {

     hr = CreateMappedTexture(texture_2D, new_size.Width, new_size.Height);

     if (FAILED(hr))

       return hr;

diff --git a/modules/desktop_capture/win/wgc_capturer_win_unittest.cc b/modules/desktop_capture/win/wgc_capturer_win_unittest.cc
index 01af044..25866c2 100644
--- a/modules/desktop_capture/win/wgc_capturer_win_unittest.cc
+++ b/modules/desktop_capture/win/wgc_capturer_win_unittest.cc
@@ -34,10 +34,10 @@
 
 const int kSmallWindowWidth = 200;
 const int kSmallWindowHeight = 100;
-const int kWindowWidth = 300;
-const int kWindowHeight = 200;
+const int kMediumWindowWidth = 300;
+const int kMediumWindowHeight = 200;
 const int kLargeWindowWidth = 400;
-const int kLargeWindowHeight = 300;
+const int kLargeWindowHeight = 500;
 
 // The size of the image we capture is slightly smaller than the actual size of
 // the window.
@@ -69,10 +69,11 @@
     EXPECT_TRUE(com_initializer_->Succeeded());
   }
 
-  void SetUpForWindowCapture() {
+  void SetUpForWindowCapture(int window_width = kMediumWindowWidth,
+                             int window_height = kMediumWindowHeight) {
     capturer_ = WgcCapturerWin::CreateRawWindowCapturer(
         DesktopCaptureOptions::CreateDefault());
-    CreateWindowOnSeparateThread();
+    CreateWindowOnSeparateThread(window_width, window_height);
     StartWindowThreadMessageLoop();
     source_id_ = GetTestWindowIdFromSourceList();
   }
@@ -93,14 +94,15 @@
   // without blocking the test thread. This is necessary if we are interested in
   // having GraphicsCaptureItem events (i.e. the Closed event) fire, and it more
   // closely resembles how capture works in the wild.
-  void CreateWindowOnSeparateThread() {
+  void CreateWindowOnSeparateThread(int window_width, int window_height) {
     window_thread_ = rtc::Thread::Create();
     window_thread_->SetName(kWindowThreadName, nullptr);
     window_thread_->Start();
-    window_thread_->Invoke<void>(RTC_FROM_HERE, [this]() {
+    window_thread_->Invoke<void>(RTC_FROM_HERE, [this, window_width,
+                                                 window_height]() {
       window_thread_id_ = GetCurrentThreadId();
       window_info_ =
-          CreateTestWindow(kWindowTitle, kWindowHeight, kWindowWidth);
+          CreateTestWindow(kWindowTitle, window_height, window_width);
       window_open_ = true;
 
       while (!IsWindowResponding(window_info_.hwnd)) {
@@ -181,6 +183,39 @@
     EXPECT_TRUE(frame_);
   }
 
+  void ValidateFrame(int expected_width, int expected_height) {
+    EXPECT_EQ(frame_->size().width(), expected_width - kWindowWidthSubtrahend);
+    EXPECT_EQ(frame_->size().height(),
+              expected_height - kWindowHeightSubtrahend);
+
+    // Verify the buffer contains as much data as it should, and that the right
+    // colors are found.
+    int data_length = frame_->stride() * frame_->size().height();
+
+    // The first and last pixel should have the same color because they will be
+    // from the border of the window.
+    // Pixels have 4 bytes of data so the whole pixel needs a uint32_t to fit.
+    uint32_t first_pixel = static_cast<uint32_t>(*frame_->data());
+    uint32_t last_pixel = static_cast<uint32_t>(
+        *(frame_->data() + data_length - DesktopFrame::kBytesPerPixel));
+    EXPECT_EQ(first_pixel, last_pixel);
+
+    // Let's also check a pixel from the middle of the content area, which the
+    // TestWindow will paint a consistent color for us to verify.
+    uint8_t* middle_pixel = frame_->data() + (data_length / 2);
+
+    int sub_pixel_offset = DesktopFrame::kBytesPerPixel / 4;
+    EXPECT_EQ(*middle_pixel, kTestWindowBValue);
+    middle_pixel += sub_pixel_offset;
+    EXPECT_EQ(*middle_pixel, kTestWindowGValue);
+    middle_pixel += sub_pixel_offset;
+    EXPECT_EQ(*middle_pixel, kTestWindowRValue);
+    middle_pixel += sub_pixel_offset;
+
+    // The window is opaque so we expect 0xFF for the Alpha channel.
+    EXPECT_EQ(*middle_pixel, 0xFF);
+  }
+
   // DesktopCapturer::Callback interface
   // The capturer synchronously invokes this method before |CaptureFrame()|
   // returns.
@@ -276,30 +311,44 @@
   EXPECT_FALSE(capturer_->SelectSource(source_id_));
 }
 
-TEST_F(WgcCapturerWinTest, ResizeWindowMidCapture) {
-  SetUpForWindowCapture();
+TEST_F(WgcCapturerWinTest, IncreaseWindowSizeMidCapture) {
+  SetUpForWindowCapture(kSmallWindowWidth, kSmallWindowHeight);
   EXPECT_TRUE(capturer_->SelectSource(source_id_));
 
   capturer_->Start(this);
   DoCapture();
-  EXPECT_EQ(frame_->size().width(), kWindowWidth - kWindowWidthSubtrahend);
-  EXPECT_EQ(frame_->size().height(), kWindowHeight - kWindowHeightSubtrahend);
+  ValidateFrame(kSmallWindowWidth, kSmallWindowHeight);
 
-  ResizeTestWindow(window_info_.hwnd, kLargeWindowWidth, kLargeWindowHeight);
+  ResizeTestWindow(window_info_.hwnd, kSmallWindowWidth, kMediumWindowHeight);
   DoCapture();
   // We don't expect to see the new size until the next capture, as the frame
-  // pool hadn't had a chance to resize yet.
+  // pool hadn't had a chance to resize yet to fit the new, larger image.
   DoCapture();
-  EXPECT_EQ(frame_->size().width(), kLargeWindowWidth - kWindowWidthSubtrahend);
-  EXPECT_EQ(frame_->size().height(),
-            kLargeWindowHeight - kWindowHeightSubtrahend);
+  ValidateFrame(kSmallWindowWidth, kMediumWindowHeight);
 
-  ResizeTestWindow(window_info_.hwnd, kSmallWindowWidth, kSmallWindowHeight);
+  ResizeTestWindow(window_info_.hwnd, kLargeWindowWidth, kMediumWindowHeight);
   DoCapture();
   DoCapture();
-  EXPECT_EQ(frame_->size().width(), kSmallWindowWidth - kWindowWidthSubtrahend);
-  EXPECT_EQ(frame_->size().height(),
-            kSmallWindowHeight - kWindowHeightSubtrahend);
+  ValidateFrame(kLargeWindowWidth, kMediumWindowHeight);
+}
+
+TEST_F(WgcCapturerWinTest, ReduceWindowSizeMidCapture) {
+  SetUpForWindowCapture(kLargeWindowWidth, kLargeWindowHeight);
+  EXPECT_TRUE(capturer_->SelectSource(source_id_));
+
+  capturer_->Start(this);
+  DoCapture();
+  ValidateFrame(kLargeWindowWidth, kLargeWindowHeight);
+
+  ResizeTestWindow(window_info_.hwnd, kLargeWindowWidth, kMediumWindowHeight);
+  // We expect to see the new size immediately because the image data has shrunk
+  // and will fit in the existing buffer.
+  DoCapture();
+  ValidateFrame(kLargeWindowWidth, kMediumWindowHeight);
+
+  ResizeTestWindow(window_info_.hwnd, kSmallWindowWidth, kMediumWindowHeight);
+  DoCapture();
+  ValidateFrame(kSmallWindowWidth, kMediumWindowHeight);
 }
 
 TEST_F(WgcCapturerWinTest, MinimizeWindowMidCapture) {
@@ -329,8 +378,7 @@
 
   capturer_->Start(this);
   DoCapture();
-  EXPECT_EQ(frame_->size().width(), kWindowWidth - kWindowWidthSubtrahend);
-  EXPECT_EQ(frame_->size().height(), kWindowHeight - kWindowHeightSubtrahend);
+  ValidateFrame(kMediumWindowWidth, kMediumWindowHeight);
 
   CloseTestWindow();