// Copyright (c) 2013-2014 Sandstorm Development Group, Inc. and contributors // Licensed under the MIT License: // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal // in the Software without restriction, including without limitation the rights // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell // copies of the Software, and to permit persons to whom the Software is // furnished to do so, subject to the following conditions: // // The above copyright notice and this permission notice shall be included in // all copies or substantial portions of the Software. // // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN // THE SOFTWARE. #include "main.h" #include "debug.h" #include "arena.h" #include "miniposix.h" #include #include #include #include #include #if _WIN32 #define WIN32_LEAN_AND_MEAN #ifndef NOMINMAX #define NOMINMAX 1 #endif #include #include "windows-sanity.h" #else #include #endif namespace kj { // ======================================================================================= TopLevelProcessContext::TopLevelProcessContext(StringPtr programName) : programName(programName), cleanShutdown(getenv("KJ_CLEAN_SHUTDOWN") != nullptr) { printStackTraceOnCrash(); } StringPtr TopLevelProcessContext::getProgramName() { return programName; } void TopLevelProcessContext::exit() { int exitCode = hadErrors ? 1 : 0; if (cleanShutdown) { #if KJ_NO_EXCEPTIONS // This is the best we can do. warning("warning: KJ_CLEAN_SHUTDOWN may not work correctly when compiled " "with -fno-exceptions."); ::exit(exitCode); #else throw CleanShutdownException { exitCode }; #endif } _exit(exitCode); } #if _WIN32 void setStandardIoMode(int fd) { // Set mode to binary if the fd is not a console. HANDLE handle = reinterpret_cast(_get_osfhandle(fd)); DWORD consoleMode; if (GetConsoleMode(handle, &consoleMode)) { // It's a console. } else { KJ_SYSCALL(_setmode(fd, _O_BINARY)); } } #else void setStandardIoMode(int fd) {} #endif static void writeLineToFd(int fd, StringPtr message) { // Write the given message to the given file descriptor with a trailing newline iff the message // is non-empty and doesn't already have a trailing newline. We use writev() to do this in a // single system call without any copying (OS permitting). if (message.size() == 0) { return; } #if _WIN32 KJ_STACK_ARRAY(char, newlineExpansionBuffer, 2 * (message.size() + 1), 128, 512); char* p = newlineExpansionBuffer.begin(); for(char ch : message) { if(ch == '\n') { *(p++) = '\r'; } *(p++) = ch; } if(!message.endsWith("\n")) { *(p++) = '\r'; *(p++) = '\n'; } size_t newlineExpandedSize = p - newlineExpansionBuffer.begin(); KJ_ASSERT(newlineExpandedSize <= newlineExpansionBuffer.size()); HANDLE handle = reinterpret_cast(_get_osfhandle(fd)); DWORD consoleMode; bool redirectedToFile = !GetConsoleMode(handle, &consoleMode); DWORD writtenSize; if(redirectedToFile) { WriteFile(handle, newlineExpansionBuffer.begin(), newlineExpandedSize, &writtenSize, nullptr); } else { KJ_STACK_ARRAY(wchar_t, buffer, newlineExpandedSize, 128, 512); size_t finalSize = MultiByteToWideChar( CP_UTF8, 0, newlineExpansionBuffer.begin(), newlineExpandedSize, buffer.begin(), buffer.size()); KJ_ASSERT(finalSize <= buffer.size()); WriteConsoleW(handle, buffer.begin(), finalSize, &writtenSize, nullptr); } #else // Unfortunately the writev interface requires non-const pointers even though it won't modify // the data. struct iovec vec[2]; vec[0].iov_base = const_cast(message.begin()); vec[0].iov_len = message.size(); vec[1].iov_base = const_cast("\n"); vec[1].iov_len = 1; struct iovec* pos = vec; // Only use the second item in the vec if the message doesn't already end with \n. uint count = message.endsWith("\n") ? 1 : 2; for (;;) { ssize_t n = writev(fd, pos, count); if (n < 0) { if (errno == EINTR) { continue; } else { // This function is meant for writing to stdout and stderr. If writes fail on those FDs // there's not a whole lot we can reasonably do, so just ignore it. return; } } // Update chunks to discard what was successfully written. for (;;) { if (count == 0) { // Done writing. return; } else if (pos->iov_len <= implicitCast(n)) { // Wrote this entire chunk. n -= pos->iov_len; ++pos; --count; } else { // Wrote only part of this chunk. Adjust the pointer and then retry. pos->iov_base = reinterpret_cast(pos->iov_base) + n; pos->iov_len -= n; break; } } } #endif } void TopLevelProcessContext::warning(StringPtr message) { writeLineToFd(STDERR_FILENO, message); } void TopLevelProcessContext::error(StringPtr message) { hadErrors = true; writeLineToFd(STDERR_FILENO, message); } void TopLevelProcessContext::exitError(StringPtr message) { error(message); exit(); } void TopLevelProcessContext::exitInfo(StringPtr message) { writeLineToFd(STDOUT_FILENO, message); exit(); } void TopLevelProcessContext::increaseLoggingVerbosity() { // At the moment, there is only one log level that isn't enabled by default. _::Debug::setLogLevel(_::Debug::Severity::INFO); } // ======================================================================================= int runMainAndExit(ProcessContext& context, MainFunc&& func, int argc, char* argv[]) { setStandardIoMode(STDIN_FILENO); setStandardIoMode(STDOUT_FILENO); setStandardIoMode(STDERR_FILENO); #if !KJ_NO_EXCEPTIONS try { #endif KJ_ASSERT(argc > 0); KJ_STACK_ARRAY(StringPtr, params, argc - 1, 8, 32); for (int i = 1; i < argc; i++) { params[i - 1] = argv[i]; } KJ_IF_MAYBE(exception, runCatchingExceptions([&]() { func(argv[0], params); })) { context.error(str("*** Uncaught exception ***\n", *exception)); } context.exit(); #if !KJ_NO_EXCEPTIONS } catch (const TopLevelProcessContext::CleanShutdownException& e) { return e.exitCode; } #endif KJ_CLANG_KNOWS_THIS_IS_UNREACHABLE_BUT_GCC_DOESNT } // ======================================================================================= struct MainBuilder::Impl { inline Impl(ProcessContext& context, StringPtr version, StringPtr briefDescription, StringPtr extendedDescription) : context(context), version(version), briefDescription(briefDescription), extendedDescription(extendedDescription) {} ProcessContext& context; StringPtr version; StringPtr briefDescription; StringPtr extendedDescription; Arena arena; struct CharArrayCompare { inline bool operator()(const ArrayPtr& a, const ArrayPtr& b) const { int cmp = memcmp(a.begin(), b.begin(), min(a.size(), b.size())); if (cmp == 0) { return a.size() < b.size(); } else { return cmp < 0; } } }; struct Option { ArrayPtr names; bool hasArg; union { Function* func; Function* funcWithArg; }; StringPtr argTitle; StringPtr helpText; }; class OptionDisplayOrder; std::map shortOptions; std::map, Option*, CharArrayCompare> longOptions; struct SubCommand { Function func; StringPtr helpText; }; std::map subCommands; struct Arg { StringPtr title; Function callback; uint minCount; uint maxCount; }; Vector args; Maybe> finalCallback; Option& addOption(std::initializer_list names, bool hasArg, StringPtr helpText) { KJ_REQUIRE(names.size() > 0, "option must have at least one name"); Option& option = arena.allocate