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