split_link.cc revision 03b57e008b61dfcb1fbad3aea950ae0e001748b0 
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
29using namespace std;
30
31// Don't use stderr for errors because VS has large buffers on them, leading
32// to confusing error output.
33static void Fatal(const wchar_t* msg) {
34  wprintf(L"split_link fatal error: %s\n", msg);
35  exit(1);
36}
37
38static 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
55static 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
91static 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
102static 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
156static 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
178static void Fallback() {
179  Fallback(NULL);
180}
181
182static 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
198static 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.
217int 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