Home | History | Annotate | Download | only in split_link 
      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 }
    258