From a89e5592f2ee6ac73304a74b17c6f09320fff572 Mon Sep 17 00:00:00 2001
From: Cheng Zhao <zcbenz@gmail.com>
Date: Wed, 8 Jul 2015 15:55:53 +0800
Subject: [PATCH] Use stock buttons on Windows when possible

---
 atom/browser/ui/message_box_win.cc | 67 ++++++++++++++++++++++++++----
 1 file changed, 58 insertions(+), 9 deletions(-)

diff --git a/atom/browser/ui/message_box_win.cc b/atom/browser/ui/message_box_win.cc
index 6fcd7f5f900..eaa002cd8d8 100644
--- a/atom/browser/ui/message_box_win.cc
+++ b/atom/browser/ui/message_box_win.cc
@@ -4,11 +4,15 @@
 
 #include "atom/browser/ui/message_box.h"
 
+#include <map>
+#include <vector>
+
 #include <windows.h>
 #include <commctrl.h>
 
 #include "atom/browser/native_window_views.h"
 #include "base/callback.h"
+#include "base/strings/string_util.h"
 #include "base/strings/utf_string_conversions.h"
 #include "base/threading/thread.h"
 #include "base/win/scoped_gdi_object.h"
@@ -23,6 +27,47 @@ namespace {
 // a large number to avoid conflicts with Windows.
 const int kIDStart = 100;
 
+// Get the common ID from button's name.
+struct CommonButtonID {
+  int button;
+  int id;
+};
+CommonButtonID GetCommonID(const base::string16& button) {
+  base::string16 lower = base::StringToLowerASCII(button);
+  if (lower == L"ok")
+    return { TDCBF_OK_BUTTON, IDOK };
+  else if (lower == L"yes")
+    return { TDCBF_YES_BUTTON, IDYES };
+  else if (lower == L"no")
+    return { TDCBF_NO_BUTTON, IDNO };
+  else if (lower == L"cancel")
+    return { TDCBF_CANCEL_BUTTON, IDCANCEL };
+  else if (lower == L"retry")
+    return { TDCBF_RETRY_BUTTON, IDRETRY };
+  else if (lower == L"close")
+    return { TDCBF_CLOSE_BUTTON, IDCLOSE };
+  return { -1, -1 };
+}
+
+// Determine whether the buttons are common buttons, if so map common ID
+// to button ID.
+void MapToCommonID(const std::vector<base::string16>& buttons,
+                   std::map<int, int>* id_map,
+                   TASKDIALOG_COMMON_BUTTON_FLAGS* button_flags,
+                   std::vector<TASKDIALOG_BUTTON>* dialog_buttons) {
+  for (size_t i = 0; i < buttons.size(); ++i) {
+    auto common = GetCommonID(buttons[i]);
+    if (common.button != -1) {
+      // It is a common button.
+      (*id_map)[common.id] = i;
+      (*button_flags) |= common.button;
+    } else {
+      // It is a custom button.
+      dialog_buttons->push_back({i + kIDStart, buttons[i].c_str()});
+    }
+  }
+}
+
 int ShowMessageBoxUTF16(HWND parent,
                         MessageBoxType type,
                         const std::vector<base::string16>& buttons,
@@ -31,10 +76,6 @@ int ShowMessageBoxUTF16(HWND parent,
                         const base::string16& message,
                         const base::string16& detail,
                         const gfx::ImageSkia& icon) {
-  std::vector<TASKDIALOG_BUTTON> dialog_buttons;
-  for (size_t i = 0; i < buttons.size(); ++i)
-    dialog_buttons.push_back({i + kIDStart, buttons[i].c_str()});
-
   TASKDIALOG_FLAGS flags = TDF_SIZE_TO_CONTENT;  // show all content.
   if (cancel_id != 0)
     flags |= TDF_ALLOW_DIALOG_CANCELLATION;  // allow dialog to be cancelled.
@@ -45,8 +86,6 @@ int ShowMessageBoxUTF16(HWND parent,
   config.hInstance      = GetModuleHandle(NULL);
   config.dwFlags        = flags;
   config.pszWindowTitle = title.c_str();
-  config.pButtons       = &dialog_buttons.front();
-  config.cButtons       = dialog_buttons.size();
 
   base::win::ScopedHICON hicon;
   if (!icon.isNull()) {
@@ -76,12 +115,22 @@ int ShowMessageBoxUTF16(HWND parent,
     config.pszContent = detail.c_str();
   }
 
+  // Iterate through the buttons, put common buttons in dwCommonButtons
+  // and custom buttons in pButtons.
+  std::map<int, int> id_map;
+  std::vector<TASKDIALOG_BUTTON> dialog_buttons;
+  MapToCommonID(buttons, &id_map, &config.dwCommonButtons, &dialog_buttons);
+  config.pButtons = &dialog_buttons.front();
+  config.cButtons = dialog_buttons.size();
+
   int id = 0;
   TaskDialogIndirect(&config, &id, NULL, NULL);
-  if (id == 0 || id == IDCANCEL)
-    return cancel_id;
-  else
+  if (id_map.find(id) != id_map.end())  // common button.
+    return id_map[id];
+  else if (id >= kIDStart)  // custom button.
     return id - kIDStart;
+  else
+    return cancel_id;
 }
 
 void RunMessageBoxInNewThread(base::Thread* thread,