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                      &cmd[0],
 139                      NULL,
 140                      NULL,
 141                      TRUE,
 142                      0,
 143                      NULL,
 144                      NULL,
 145                      &startup_info, &process_info)) {
 146     wstring error = ErrorMessageToString(GetLastError());
 147     Fatal(error.c_str());
 148   }
 149   CloseHandle(process_info.hThread);
 150   WaitForSingleObject(process_info.hProcess, INFINITE);
 151   GetExitCodeProcess(process_info.hProcess, &exit_code);
 152   CloseHandle(process_info.hProcess);
 153   exit(exit_code);
 154 }
 155
 156 static void Fallback(const wchar_t* msg) {
 157   wchar_t original_link[1024];
 158   DWORD type;
 159   DWORD size = sizeof(original_link);
 160   if (SHGetValue(HKEY_CURRENT_USER,
 161                  L"Software\\Chromium\\split_link_installed",
 162                  NULL,
 163                  &type,
 164                  original_link,
 165                  &size) != ERROR_SUCCESS || type != REG_SZ) {
 166     Fatal(L"Couldn't retrieve linker location from "
 167           L"HKCU\\Software\\Chromium\\split_link_installed.");
 168   }
 169   if (getenv("SPLIT_LINK_DEBUG")) {
 170     wprintf(L"  got original linker '%s'\n", original_link);
 171     fflush(stdout);
 172   }
 173   vector<wstring> link_binary;
 174   link_binary.push_back(original_link);
 175   RunLinker(link_binary, msg);
 176 }
 177
 178 static void Fallback() {
 179   Fallback(NULL);
 180 }
 181
 182 static unsigned char* SlurpFile(const wchar_t* path, size_t* length) {
 183   HANDLE file = CreateFile(
 184       path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 185   if (file == INVALID_HANDLE_VALUE)
 186     Fallback(L"couldn't open file");
 187   LARGE_INTEGER file_size;
 188   if (!GetFileSizeEx(file, &file_size))
 189     Fallback(L"couldn't get file size");
 190   *length = static_cast<size_t>(file_size.QuadPart);
 191   unsigned char* buffer = static_cast<unsigned char*>(malloc(*length));
 192   DWORD bytes_read = 0;
 193   if (!ReadFile(file, buffer, *length, &bytes_read, NULL))
 194     Fallback(L"couldn't read file");
 195   return buffer;
 196 }
 197
 198 static bool SplitLinkRequested(const wchar_t* rsp_path) {
 199   size_t length;
 200   unsigned char* data = SlurpFile(rsp_path, &length);
 201   bool flag_found = false;
 202   if (data[0] == 0xff && data[1] == 0xfe) {
 203     // UTF-16LE
 204     wstring wide(reinterpret_cast<wchar_t*>(&data[2]),
 205                  length / sizeof(wchar_t) - 1);
 206     flag_found = wide.find(L"/splitlink") != wide.npos;
 207   } else {
 208     string narrow(reinterpret_cast<char*>(data), length);
 209     flag_found = narrow.find("/splitlink") != narrow.npos;
 210   }
 211   free(data);
 212   return flag_found;
 213 }
 214
 215 // If /splitlink is on the command line, delegate to split_link.py, otherwise
 216 // fallback to standard linker.
 217 int wmain(int argc, wchar_t** argv) {
 218   int rsp_file_index = -1;
 219
 220   if (argc < 2)
 221     Fallback();
 222
 223   for (int i = 1; i < argc; ++i) {
 224     if (argv[i][0] == L'@') {
 225       rsp_file_index = i;
 226       break;
 227     }
 228   }
 229
 230   if (rsp_file_index == -1)
 231     Fallback(L"couldn't find a response file in argv");
 232
 233   if (getenv("SPLIT_LINK_DEBUG")) {
 234     wstring backup_copy(&argv[rsp_file_index][1]);
 235     backup_copy += L".copy";
 236     wchar_t buf[1024];
 237     swprintf(buf,
 238              sizeof(buf),
 239              L"copy %s %s",
 240              &argv[rsp_file_index][1],
 241              backup_copy.c_str());
 242     if (_wsystem(buf) == 0)
 243       wprintf(L"Saved original rsp as %s\n", backup_copy.c_str());
 244     else
 245       wprintf(L"'%s' failed.", buf);
 246   }
 247
 248   if (SplitLinkRequested(&argv[rsp_file_index][1])) {
 249     vector<wstring> link_binary;
 250     link_binary.push_back(WPYTHON_PATH);
 251     link_binary.push_back(WSPLIT_LINK_SCRIPT_PATH);
 252     RunLinker(link_binary, NULL);
 253   }
 254
 255   // Otherwise, run regular linker silently.
 256   Fallback();
 257 }