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                     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
157static 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
179static void Fallback() {
180  Fallback(NULL);
181}
182
183static 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
199static 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.
218int 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