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                      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 }
    259