src/tools/win/split_link/split_link.cc

   1 // Copyright (c) 2013 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include <windows.h>
   6 #include <shlwapi.h>
   7
   8 #include <stdio.h>
   9 #include <stdlib.h>
  10
  11 #include <algorithm>
  12 #include <iterator>
  13 #include <string>
  14 #include <vector>
  15
  16 #ifndef SPLIT_LINK_SCRIPT_PATH
  17 #error SPLIT_LINK_SCRIPT_PATH must be defined as the path to "split_link.py".
  18 #endif
  19
  20 #ifndef PYTHON_PATH
  21 #error PYTHON_PATH must be defined to be the path to the python binary.
  22 #endif
  23
  24 #define WIDEN2(x) L ## x
  25 #define WIDEN(x) WIDEN2(x)
  26 #define WPYTHON_PATH WIDEN(PYTHON_PATH)
  27 #define WSPLIT_LINK_SCRIPT_PATH WIDEN(SPLIT_LINK_SCRIPT_PATH)
  28
  29 using namespace std;
  30
  31 // Don't use stderr for errors because VS has large buffers on them, leading
  32 // to confusing error output.
  33 static void Fatal(const wchar_t* msg) {
  34   wprintf(L"split_link fatal error: %s\n", msg);
  35   exit(1);
  36 }
  37
  38 static wstring ErrorMessageToString(DWORD err) {
  39   wchar_t* msg_buf = NULL;
  40   DWORD rc = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
  41                                FORMAT_MESSAGE_FROM_SYSTEM,
  42                            NULL,
  43                            err,
  44                            MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
  45                            reinterpret_cast<LPTSTR>(&msg_buf),
  46                            0,
  47                            NULL);
  48   if (!rc)
  49     return L"unknown error";
  50   wstring ret(msg_buf);
  51   LocalFree(msg_buf);
  52   return ret;
  53 }
  54
  55 static void ArgvQuote(const std::wstring& argument,
  56                       std::wstring* command_line) {
  57   // Don't quote unless we actually need to.
  58   if (!argument.empty() &&
  59       argument.find_first_of(L" \t\n\v\"") == argument.npos) {
  60     command_line->append(argument);
  61   } else {
  62     command_line->push_back(L'"');
  63     for (std::wstring::const_iterator it = argument.begin();; ++it) {
  64       int num_backslashes = 0;
  65       while (it != argument.end() && *it == L'\\') {
  66         ++it;
  67         ++num_backslashes;
  68       }
  69       if (it == argument.end()) {
  70         // Escape all backslashes, but let the terminating double quotation
  71         // mark we add below be interpreted as a metacharacter.
  72         command_line->append(num_backslashes * 2, L'\\');
  73         break;
  74       } else if (*it == L'"') {
  75         // Escape all backslashes and the following double quotation mark.
  76         command_line->append(num_backslashes * 2 + 1, L'\\');
  77         command_line->push_back(*it);
  78       } else {
  79         // Backslashes aren't special here.
  80         command_line->append(num_backslashes, L'\\');
  81         command_line->push_back(*it);
  82       }
  83     }
  84     command_line->push_back(L'"');
  85   }
  86 }
  87
  88 // Does the opposite of CommandLineToArgvW. Suitable for CreateProcess, but
  89 // not for cmd.exe. |args| should include the program name as argv[0].
  90 // See http://blogs.msdn.com/b/twistylittlepassagesallalike/archive/2011/04/23/everyone-quotes-arguments-the-wrong-way.aspx
  91 static wstring BuildCommandLine(const vector<wstring>& args) {
  92   std::wstring result;
  93   for (size_t i = 0; i < args.size(); ++i) {
  94     ArgvQuote(args[i], &result);
  95     if (i < args.size() - 1) {
  96       result += L" ";
  97     }
  98   }
  99   return result;
 100 }
 101
 102 static void RunLinker(const vector<wstring>& prefix, const wchar_t* msg) {
 103   if (msg) {
 104     wprintf(L"split_link failed (%s), trying to fallback to standard link.\n",
 105             msg);
 106     wprintf(L"Original command line: %s\n", GetCommandLine());
 107     fflush(stdout);
 108   }
 109
 110   STARTUPINFO startup_info = { sizeof(STARTUPINFO) };
 111   PROCESS_INFORMATION process_info;
 112   DWORD exit_code;
 113
 114   GetStartupInfo(&startup_info);
 115
 116   if (getenv("SPLIT_LINK_DEBUG")) {
 117     wprintf(L"  original command line '%s'\n", GetCommandLine());
 118     fflush(stdout);
 119   }
 120
 121   int num_args;
 122   LPWSTR* args = CommandLineToArgvW(GetCommandLine(), &num_args);
 123   if (!args)
 124     Fatal(L"Couldn't parse command line.");
 125   vector<wstring> argv;
 126   argv.insert(argv.end(), prefix.begin(), prefix.end());
 127   for (int i = 1; i < num_args; ++i)  // Skip old argv[0].
 128     argv.push_back(args[i]);
 129   LocalFree(args);
 130
 131   wstring cmd = BuildCommandLine(argv);
 132
 133   if (getenv("SPLIT_LINK_DEBUG")) {
 134     wprintf(L"  running '%s'\n", cmd.c_str());
 135     fflush(stdout);
 136   }
 137   if (!CreateProcess(NULL,
 138                      reinterpret_cast<LPWSTR>(const_cast<wchar_t *>(
 139                              cmd.c_str())),
 140                      NULL,
 141                      NULL,
 142                      TRUE,
 143                      0,
 144                      NULL,
 145                      NULL,
 146                      &startup_info, &process_info)) {
 147     wstring error = ErrorMessageToString(GetLastError());
 148     Fatal(error.c_str());
 149   }
 150   CloseHandle(process_info.hThread);
 151   WaitForSingleObject(process_info.hProcess, INFINITE);
 152   GetExitCodeProcess(process_info.hProcess, &exit_code);
 153   CloseHandle(process_info.hProcess);
 154   exit(exit_code);
 155 }
 156
 157 static void Fallback(const wchar_t* msg) {
 158   wchar_t original_link[1024];
 159   DWORD type;
 160   DWORD size = sizeof(original_link);
 161   if (SHGetValue(HKEY_CURRENT_USER,
 162                  L"Software\\Chromium\\split_link_installed",
 163                  NULL,
 164                  &type,
 165                  original_link,
 166                  &size) != ERROR_SUCCESS || type != REG_SZ) {
 167     Fatal(L"Couldn't retrieve linker location from "
 168           L"HKCU\\Software\\Chromium\\split_link_installed.");
 169   }
 170   if (getenv("SPLIT_LINK_DEBUG")) {
 171     wprintf(L"  got original linker '%s'\n", original_link);
 172     fflush(stdout);
 173   }
 174   vector<wstring> link_binary;
 175   link_binary.push_back(original_link);
 176   RunLinker(link_binary, msg);
 177 }
 178
 179 static void Fallback() {
 180   Fallback(NULL);
 181 }
 182
 183 static unsigned char* SlurpFile(const wchar_t* path, size_t* length) {
 184   HANDLE file = CreateFile(
 185       path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 186   if (file == INVALID_HANDLE_VALUE)
 187     Fallback(L"couldn't open file");
 188   LARGE_INTEGER file_size;
 189   if (!GetFileSizeEx(file, &file_size))
 190     Fallback(L"couldn't get file size");
 191   *length = static_cast<size_t>(file_size.QuadPart);
 192   unsigned char* buffer = static_cast<unsigned char*>(malloc(*length));
 193   DWORD bytes_read = 0;
 194   if (!ReadFile(file, buffer, *length, &bytes_read, NULL))
 195     Fallback(L"couldn't read file");
 196   return buffer;
 197 }
 198
 199 static bool SplitLinkRequested(const wchar_t* rsp_path) {
 200   size_t length;
 201   unsigned char* data = SlurpFile(rsp_path, &length);
 202   bool flag_found = false;
 203   if (data[0] == 0xff && data[1] == 0xfe) {
 204     // UTF-16LE
 205     wstring wide(reinterpret_cast<wchar_t*>(&data[2]),
 206                  length / sizeof(wchar_t) - 1);
 207     flag_found = wide.find(L"/splitlink") != wide.npos;
 208   } else {
 209     string narrow(reinterpret_cast<char*>(data), length);
 210     flag_found = narrow.find("/splitlink") != narrow.npos;
 211   }
 212   free(data);
 213   return flag_found;
 214 }
 215
 216 // If /splitlink is on the command line, delegate to split_link.py, otherwise
 217 // fallback to standard linker.
 218 int wmain(int argc, wchar_t** argv) {
 219   int rsp_file_index = -1;
 220
 221   if (argc < 2)
 222     Fallback();
 223
 224   for (int i = 1; i < argc; ++i) {
 225     if (argv[i][0] == L'@') {
 226       rsp_file_index = i;
 227       break;
 228     }
 229   }
 230
 231   if (rsp_file_index == -1)
 232     Fallback(L"couldn't find a response file in argv");
 233
 234   if (getenv("SPLIT_LINK_DEBUG")) {
 235     wstring backup_copy(&argv[rsp_file_index][1]);
 236     backup_copy += L".copy";
 237     wchar_t buf[1024];
 238     swprintf(buf,
 239              sizeof(buf),
 240              L"copy %s %s",
 241              &argv[rsp_file_index][1],
 242              backup_copy.c_str());
 243     if (_wsystem(buf) == 0)
 244       wprintf(L"Saved original rsp as %s\n", backup_copy.c_str());
 245     else
 246       wprintf(L"'%s' failed.", buf);
 247   }
 248
 249   if (SplitLinkRequested(&argv[rsp_file_index][1])) {
 250     vector<wstring> link_binary;
 251     link_binary.push_back(WPYTHON_PATH);
 252     link_binary.push_back(WSPLIT_LINK_SCRIPT_PATH);
 253     RunLinker(link_binary, NULL);
 254   }
 255
 256   // Otherwise, run regular linker silently.
 257   Fallback();
 258 }