tee-like utility demonstrating usage of WIN32 pipes

rooibos.cpp

// tee-like utility demonstrating usage of pipes.
//
// This program was originally intended to demonstrate overlapped i/o
// and WaitForMultipleObjects() but this proved to be quite complex and
// there is no need to wait concurrently on the pipe and child process
// as the child first emits its output (in the pipe) *then* terminates.
//
// The following command:
//      rooibos foo.out foo bar
// is equivalent to:
//      foo bar 2>&1 | tee foo.out
//
// Build with:
// cl /W4 /Zi /EHsc rooibos.cpp 
//
// TODO: i18n (uses narrow characters for the time being out of lazyness)

#define _CRT_SECURE_NO_WARNINGS

#include <cassert>
#include <cstring>
#include <iostream>
#include <sstream>
#include <iterator>

#include <windows.h>

// RAII wrapper over HANDLE.
class Handle {
public:
    Handle(HANDLE h) : h_(h) {}

    Handle(const Handle&) = delete;
    Handle(Handle&& other) {
        h_ = other.h_;
        other.h_ = INVALID_HANDLE_VALUE;
    }

    Handle& operator=(const Handle&) = delete;
    Handle& operator=(Handle&&) = delete;

    ~Handle() {
        if (h_ != INVALID_HANDLE_VALUE) {
            CloseHandle(h_);
        }
    }

    HANDLE get() const { return h_; }

    void close() {
        if (h_ != INVALID_HANDLE_VALUE) {
            CloseHandle(h_);
            h_ = INVALID_HANDLE_VALUE;
        }
    }

    static Handle fromRaw(HANDLE& raw) {
        Handle h(raw);
        raw = INVALID_HANDLE_VALUE;
        return h;
    }

private:
    HANDLE h_;
};

int main(int argc, char *argv[]) {
    if (argc < 3) {
        std::cerr << "usage: rooibos output_file command args...\n";
        return 1;
    }

    //
    // Create output file.
    //

    Handle hOutFile = CreateFileA(argv[1], GENERIC_WRITE, /*dwShareMode=*/0,
                                  /*lpSecurityAttributes=*/nullptr,
                                  CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL,
                                  /*hTemplateFile=*/nullptr);
    if (hOutFile.get() == INVALID_HANDLE_VALUE) {
        std::cerr << "failed to open " << argv[1] << ": " << GetLastError();
        return 1;
    }

    //
    // Format command line in a form that pleases CreateProcess().
    // Way too many copies! ;-)
    // TODO: Handle space-separated arguments.
    //

    std::ostringstream oss;
    std::copy(argv + 2, argv + argc, std::ostream_iterator<const char *>(oss, " "));
    auto cmdline = std::make_unique<char[]>(oss.str().size() + 1);
    std::strcpy(cmdline.get(), oss.str().c_str());

    //
    // Create a pipe to collect child process output.
    //

    constexpr int kPipeBufSize = 64 * 1024;
    SECURITY_ATTRIBUTES saPipe = { sizeof(saPipe) };
    saPipe.bInheritHandle = true;
    HANDLE hPipeParentRaw;
    HANDLE hPipeChildRaw;
    bool success = CreatePipe(&hPipeParentRaw, &hPipeChildRaw,
                              /*lpSecurityAttr=*/&saPipe, kPipeBufSize);
    if (!success) {
        std::cerr << "failed to create pipe: " << GetLastError();
        return 1;
    }
    auto hPipeParent = Handle::fromRaw(hPipeParentRaw);
    auto hPipeChild = Handle::fromRaw(hPipeChildRaw);

    //
    // Spawn child process with output redirected to pipe.
    //

    STARTUPINFO si = { sizeof(STARTUPINFO) };
    si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
    si.hStdOutput = hPipeChild.get();
    si.hStdError = hPipeChild.get();
    si.dwFlags = STARTF_USESTDHANDLES;

    PROCESS_INFORMATION procinfo;
    success = CreateProcessA(argv[2], cmdline.get(),
                             /*lpProcAttributes=*/nullptr,
                             /*lpThreadAttributes=*/nullptr,
                             /*bInheritHandles=*/true,
                             /*dwCreationFlags=*/0,
                             /*lpEnv=*/nullptr,
                             /*lpCwd=*/nullptr,
                             &si, &procinfo);
    if (!success) {
        std::cerr << "failed to spawn " << cmdline.get() << ": " << GetLastError() << '\n';
        return 1;
    }
    auto hChildProc = Handle::fromRaw(procinfo.hProcess);
    auto hChildThread = Handle::fromRaw(procinfo.hThread);

    // Without this, we will deadlock waiting for ERROR_BROKEN_PIPE.
    hPipeChild.close();

    //
    // Copy all data from pipe onto stdout and output file.
    //

    HANDLE hStdout = GetStdHandle(STD_OUTPUT_HANDLE);
    char pipeOutput[kPipeBufSize];
    DWORD nread;
    for (;;) {
        success = ReadFile(hPipeParent.get(), pipeOutput, sizeof(pipeOutput),
                           &nread, /*lpOverlapped=*/nullptr);
        if (!success) {
            if (GetLastError() == ERROR_BROKEN_PIPE) {
                break;
            } else {
                std::cerr << "error while reading pipe: " << GetLastError() << '\n';
                return 1;
            }
        }

        if (nread == 0) {
            break;
        }

        DWORD nwritten;
        success = WriteFile(hStdout, pipeOutput, nread, &nwritten,
                            /*lpOverlapped=*/nullptr);
        if (!success) {
            std::cerr << "error while writing to stdout: " << GetLastError() << '\n';
            return 1;
        }
        // TODO: handle partial writes
        assert(nread == nwritten);

        success = WriteFile(hOutFile.get(), pipeOutput, nread, &nwritten,
                            /*lpOverlapped=*/nullptr);
        if (!success) {
            std::cerr << "error while writing to stdout: " << GetLastError() << '\n';
            return 1;
        }
        // TODO: handle partial writes
        assert(nread == nwritten);
    }

    //
    // Wait for child process to complete.
    //

    if (WaitForSingleObject(hChildProc.get(), INFINITE) != WAIT_OBJECT_0) {
        std::cerr << "error in WaitForMultipleObjects(): " << GetLastError() << '\n';
        return 1;
    }

    return 0;
}