shaka-packager/tools/win/split_link/split_link.cc

259 lines
7.7 KiB
C++

// Copyright (c) 2013 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include <windows.h>
#include <shlwapi.h>
#include <stdio.h>
#include <stdlib.h>
#include <algorithm>
#include <iterator>
#include <string>
#include <vector>
#ifndef SPLIT_LINK_SCRIPT_PATH
#error SPLIT_LINK_SCRIPT_PATH must be defined as the path to "split_link.py".
#endif
#ifndef PYTHON_PATH
#error PYTHON_PATH must be defined to be the path to the python binary.
#endif
#define WIDEN2(x) L ## x
#define WIDEN(x) WIDEN2(x)
#define WPYTHON_PATH WIDEN(PYTHON_PATH)
#define WSPLIT_LINK_SCRIPT_PATH WIDEN(SPLIT_LINK_SCRIPT_PATH)
using namespace std;
// Don't use stderr for errors because VS has large buffers on them, leading
// to confusing error output.
static void Fatal(const wchar_t* msg) {
wprintf(L"split_link fatal error: %s\n", msg);
exit(1);
}
static wstring ErrorMessageToString(DWORD err) {
wchar_t* msg_buf = NULL;
DWORD rc = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
FORMAT_MESSAGE_FROM_SYSTEM,
NULL,
err,
MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
reinterpret_cast<LPTSTR>(&msg_buf),
0,
NULL);
if (!rc)
return L"unknown error";
wstring ret(msg_buf);
LocalFree(msg_buf);
return ret;
}
static void ArgvQuote(const std::wstring& argument,
std::wstring* command_line) {
// Don't quote unless we actually need to.
if (!argument.empty() &&
argument.find_first_of(L" \t\n\v\"") == argument.npos) {
command_line->append(argument);
} else {
command_line->push_back(L'"');
for (std::wstring::const_iterator it = argument.begin();; ++it) {
int num_backslashes = 0;
while (it != argument.end() && *it == L'\\') {
++it;
++num_backslashes;
}
if (it == argument.end()) {
// Escape all backslashes, but let the terminating double quotation
// mark we add below be interpreted as a metacharacter.
command_line->append(num_backslashes * 2, L'\\');
break;
} else if (*it == L'"') {
// Escape all backslashes and the following double quotation mark.
command_line->append(num_backslashes * 2 + 1, L'\\');
command_line->push_back(*it);
} else {
// Backslashes aren't special here.
command_line->append(num_backslashes, L'\\');
command_line->push_back(*it);
}
}
command_line->push_back(L'"');
}
}
// Does the opposite of CommandLineToArgvW. Suitable for CreateProcess, but
// not for cmd.exe. |args| should include the program name as argv[0].
// See http://blogs.msdn.com/b/twistylittlepassagesallalike/archive/2011/04/23/everyone-quotes-arguments-the-wrong-way.aspx
static wstring BuildCommandLine(const vector<wstring>& args) {
std::wstring result;
for (size_t i = 0; i < args.size(); ++i) {
ArgvQuote(args[i], &result);
if (i < args.size() - 1) {
result += L" ";
}
}
return result;
}
static void RunLinker(const vector<wstring>& prefix, const wchar_t* msg) {
if (msg) {
wprintf(L"split_link failed (%s), trying to fallback to standard link.\n",
msg);
wprintf(L"Original command line: %s\n", GetCommandLine());
fflush(stdout);
}
STARTUPINFO startup_info = { sizeof(STARTUPINFO) };
PROCESS_INFORMATION process_info;
DWORD exit_code;
GetStartupInfo(&startup_info);
if (getenv("SPLIT_LINK_DEBUG")) {
wprintf(L" original command line '%s'\n", GetCommandLine());
fflush(stdout);
}
int num_args;
LPWSTR* args = CommandLineToArgvW(GetCommandLine(), &num_args);
if (!args)
Fatal(L"Couldn't parse command line.");
vector<wstring> argv;
argv.insert(argv.end(), prefix.begin(), prefix.end());
for (int i = 1; i < num_args; ++i) // Skip old argv[0].
argv.push_back(args[i]);
LocalFree(args);
wstring cmd = BuildCommandLine(argv);
if (getenv("SPLIT_LINK_DEBUG")) {
wprintf(L" running '%s'\n", cmd.c_str());
fflush(stdout);
}
if (!CreateProcess(NULL,
reinterpret_cast<LPWSTR>(const_cast<wchar_t *>(
cmd.c_str())),
NULL,
NULL,
TRUE,
0,
NULL,
NULL,
&startup_info, &process_info)) {
wstring error = ErrorMessageToString(GetLastError());
Fatal(error.c_str());
}
CloseHandle(process_info.hThread);
WaitForSingleObject(process_info.hProcess, INFINITE);
GetExitCodeProcess(process_info.hProcess, &exit_code);
CloseHandle(process_info.hProcess);
exit(exit_code);
}
static void Fallback(const wchar_t* msg) {
wchar_t original_link[1024];
DWORD type;
DWORD size = sizeof(original_link);
if (SHGetValue(HKEY_CURRENT_USER,
L"Software\\Chromium\\split_link_installed",
NULL,
&type,
original_link,
&size) != ERROR_SUCCESS || type != REG_SZ) {
Fatal(L"Couldn't retrieve linker location from "
L"HKCU\\Software\\Chromium\\split_link_installed.");
}
if (getenv("SPLIT_LINK_DEBUG")) {
wprintf(L" got original linker '%s'\n", original_link);
fflush(stdout);
}
vector<wstring> link_binary;
link_binary.push_back(original_link);
RunLinker(link_binary, msg);
}
static void Fallback() {
Fallback(NULL);
}
static unsigned char* SlurpFile(const wchar_t* path, size_t* length) {
HANDLE file = CreateFile(
path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
if (file == INVALID_HANDLE_VALUE)
Fallback(L"couldn't open file");
LARGE_INTEGER file_size;
if (!GetFileSizeEx(file, &file_size))
Fallback(L"couldn't get file size");
*length = static_cast<size_t>(file_size.QuadPart);
unsigned char* buffer = static_cast<unsigned char*>(malloc(*length));
DWORD bytes_read = 0;
if (!ReadFile(file, buffer, *length, &bytes_read, NULL))
Fallback(L"couldn't read file");
return buffer;
}
static bool SplitLinkRequested(const wchar_t* rsp_path) {
size_t length;
unsigned char* data = SlurpFile(rsp_path, &length);
bool flag_found = false;
if (data[0] == 0xff && data[1] == 0xfe) {
// UTF-16LE
wstring wide(reinterpret_cast<wchar_t*>(&data[2]),
length / sizeof(wchar_t) - 1);
flag_found = wide.find(L"/splitlink") != wide.npos;
} else {
string narrow(reinterpret_cast<char*>(data), length);
flag_found = narrow.find("/splitlink") != narrow.npos;
}
free(data);
return flag_found;
}
// If /splitlink is on the command line, delegate to split_link.py, otherwise
// fallback to standard linker.
int wmain(int argc, wchar_t** argv) {
int rsp_file_index = -1;
if (argc < 2)
Fallback();
for (int i = 1; i < argc; ++i) {
if (argv[i][0] == L'@') {
rsp_file_index = i;
break;
}
}
if (rsp_file_index == -1)
Fallback(L"couldn't find a response file in argv");
if (getenv("SPLIT_LINK_DEBUG")) {
wstring backup_copy(&argv[rsp_file_index][1]);
backup_copy += L".copy";
wchar_t buf[1024];
swprintf(buf,
sizeof(buf),
L"copy %s %s",
&argv[rsp_file_index][1],
backup_copy.c_str());
if (_wsystem(buf) == 0)
wprintf(L"Saved original rsp as %s\n", backup_copy.c_str());
else
wprintf(L"'%s' failed.", buf);
}
if (SplitLinkRequested(&argv[rsp_file_index][1])) {
vector<wstring> link_binary;
link_binary.push_back(WPYTHON_PATH);
link_binary.push_back(WSPLIT_LINK_SCRIPT_PATH);
RunLinker(link_binary, NULL);
}
// Otherwise, run regular linker silently.
Fallback();
}