tee.c

   1 /*
   2  * CertViewer - tee for Windows
   3  * Copyright (c) 2023 "dEajL3kA" <Cumpoing79@web.de>
   4  *
   5  * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
   6  * associated documentation files (the "Software"), to deal in the Software without restriction,
   7  * including without limitation the rights to use, copy, modify, merge, publish, distribute,
   8  * sub license, and/or sell copies of the Software, and to permit persons to whom the Software is
   9  * furnished to do so, subject to the following conditions: The above copyright notice and this
  10  * permission notice shall be included in all copies or substantial portions of the Software.
  11  *
  12  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
  13  * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  14  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
  15  * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
  16  * OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  17  */
  18 #define WIN32_LEAN_AND_MEAN 1
  19 #include <Windows.h>
  20 #include <ShellAPI.h>
  21
  22 #define BUFFSIZE 4096U
  23 #pragma warning(disable: 4706)
  24
  25  // --------------------------------------------------------------------------
  26  // Utilities
  27  // --------------------------------------------------------------------------
  28
  29 static BOOL is_terminal(const HANDLE handle)
  30 {
  31     DWORD mode;
  32     return GetConsoleMode(handle, &mode);
  33 }
  34
  35 static wchar_t *concat_3(const wchar_t *const strA, const wchar_t *const strB, const wchar_t *const strC)
  36 {
  37     const size_t lenA = lstrlenW(strA), lenB = lstrlenW(strB), lenC = lstrlenW(strC);
  38     wchar_t *const buffer = (wchar_t*) LocalAlloc(LPTR, sizeof(wchar_t) * (lenA + lenB + lenC + 1U));
  39     if (buffer)
  40     {
  41         lstrcpyW(buffer, strA);
  42         lstrcpyW(buffer + lenA, strB);
  43         lstrcpyW(buffer + lenA + lenB, strC);
  44     }
  45     return buffer;
  46 }
  47
  48 #define SAFE_CLOSE_HANDLE(HANDLE) do \
  49 { \
  50     if (HANDLE) \
  51     { \
  52         CloseHandle((HANDLE)); \
  53     } \
  54 } \
  55 while (0)
  56
  57 // --------------------------------------------------------------------------
  58 // Console CTRL+C handler
  59 // --------------------------------------------------------------------------
  60
  61 static volatile BOOL g_stop = FALSE;
  62
  63 static BOOL WINAPI console_handler(const DWORD ctrlType)
  64 {
  65     switch (ctrlType)
  66     {
  67     case CTRL_C_EVENT:
  68     case CTRL_BREAK_EVENT:
  69     case CTRL_CLOSE_EVENT:
  70         g_stop = TRUE;
  71         return TRUE;
  72     default:
  73         return FALSE;
  74     }
  75 }
  76
  77 // --------------------------------------------------------------------------
  78 // Text output
  79 // --------------------------------------------------------------------------
  80
  81 static char *utf16_to_utf8(const wchar_t *const input)
  82 {
  83     const int buff_size = WideCharToMultiByte(CP_UTF8, 0, input, -1, NULL, 0, NULL, NULL);
  84     if (buff_size > 0)
  85     {
  86         char *const buffer = (char*)LocalAlloc(LPTR, buff_size);
  87         if (buffer)
  88         {
  89             const int result = WideCharToMultiByte(CP_UTF8, 0, input, -1, buffer, buff_size, NULL, NULL);
  90             if ((result > 0) && (result <= buff_size))
  91             {
  92                 return buffer;
  93             }
  94             LocalFree(buffer);
  95         }
  96     }
  97     return NULL;
  98 }
  99
 100 static BOOL write_text(const HANDLE handle, const wchar_t *const text)
 101 {
 102     BOOL result = FALSE;
 103     DWORD written;
 104     if (GetConsoleMode(handle, &written))
 105     {
 106         result = WriteConsoleW(handle, text, lstrlenW(text), &written, NULL);
 107     }
 108     else
 109     {
 110         char *const utf8_text = utf16_to_utf8(text);
 111         if (utf8_text)
 112         {
 113             result = WriteFile(handle, utf8_text, lstrlenA(utf8_text), &written, NULL);
 114             LocalFree(utf8_text);
 115         }
 116     }
 117     return result;
 118 }
 119
 120 // --------------------------------------------------------------------------
 121 // Writer thread
 122 // --------------------------------------------------------------------------
 123
 124 static BYTE buffer[2U][BUFFSIZE];
 125 static DWORD bytesTotal[2U] = { 0U, 0U };
 126 static volatile ULONG_PTR index = 0U;
 127
 128 typedef struct
 129 {
 130     HANDLE hOutput,hError;
 131     BOOL flush;
 132     HANDLE hEventReady[2U], hEventCompleted;
 133 }
 134 thread_t;
 135
 136 static DWORD WINAPI writer_thread_start_routine(const LPVOID lpThreadParameter)
 137 {
 138     DWORD bytesWritten;
 139     const thread_t *const param = (thread_t*) lpThreadParameter;
 140
 141     for (;;)
 142     {
 143         switch (WaitForMultipleObjects(2U, param->hEventReady, FALSE, INFINITE))
 144         {
 145         case WAIT_OBJECT_0:
 146             break;
 147         case WAIT_OBJECT_0 + 1U:
 148             SetEvent(param->hEventCompleted);
 149             return 0U;
 150         default:
 151             write_text(param->hError, L"[tee] System error: Failed to wait for event!\n");
 152             return 1U;
 153         }
 154
 155         const ULONG_PTR myIndex = index;
 156
 157         for (DWORD offset = 0U; offset < bytesTotal[myIndex]; offset += bytesWritten)
 158         {
 159             const BOOL result = WriteFile(param->hOutput, buffer[myIndex] + offset, bytesTotal[myIndex] - offset, &bytesWritten, NULL);
 160             if ((!result) || (!bytesWritten))
 161             {
 162                 write_text(param->hError, L"[tee] Error: Not all data could be written!\n");
 163                 break;
 164             }
 165         }
 166
 167         SetEvent(param->hEventCompleted);
 168
 169         if (param->flush)
 170         {
 171             FlushFileBuffers(param->hOutput);
 172         }
 173     }
 174 }
 175
 176 // --------------------------------------------------------------------------
 177 // MAIN
 178 // --------------------------------------------------------------------------
 179
 180 int wmain(const int argc, const wchar_t *const argv[])
 181 {
 182     HANDLE hThreads[2U] = { NULL, NULL };
 183     HANDLE hEventStop = NULL, hEventThrdReady[2U] = { NULL, NULL }, hEventCompleted[2U] = { NULL, NULL };
 184     HANDLE hMyFile = INVALID_HANDLE_VALUE;
 185     int exitCode = 1, argOff = 1;
 186     BOOL append = FALSE, flush = FALSE, ignore = FALSE;
 187     thread_t threadData[2U];
 188
 189     /* Initialize standard streams */
 190     const HANDLE hStdIn = GetStdHandle(STD_INPUT_HANDLE), hStdOut = GetStdHandle(STD_OUTPUT_HANDLE), hStdErr = GetStdHandle(STD_ERROR_HANDLE);
 191     if ((hStdIn == INVALID_HANDLE_VALUE) || (hStdOut == INVALID_HANDLE_VALUE) || (hStdErr == INVALID_HANDLE_VALUE))
 192     {
 193         FatalExit(-1);
 194     }
 195
 196     /* Set up CRTL+C handler */
 197     SetConsoleCtrlHandler(console_handler, TRUE);
 198
 199     /* Print manpage */
 200     if ((argc < 2) || (lstrcmpiW(argv[1], L"/?") == 0) || (lstrcmpiW(argv[1], L"--help") == 0))
 201     {
 202         write_text(hStdErr, L"tee for Windows [" TEXT(__DATE__) L"]\n\n");
 203         write_text(hStdErr, L"Usage:\n");
 204         write_text(hStdErr, L"  your_program.exe [...] | tee.exe [options] <output_file>\n\n");
 205         write_text(hStdErr, L"Options:\n");
 206         write_text(hStdErr, L"  --append   Append to the existing file, instead of truncating\n");
 207         write_text(hStdErr, L"  --flush    Flush output file after each write operation\n");
 208         write_text(hStdErr, L"  --ignore   Ignore the interrupt signal (SIGINT), e.g. CTRL+C\n\n");
 209         return 1;
 210     }
 211
 212     /* Parse command-line options */
 213     while (argOff < argc)
 214     {
 215         if ((argv[argOff][0U] == L'-') && (argv[argOff][1U] == L'-'))
 216         {
 217             const wchar_t *const option = argv[argOff++] + 2U;
 218             if (*option == L'\0')
 219             {
 220                 break;
 221             }
 222             else if (lstrcmpiW(option, L"append") == 0)
 223             {
 224                 append = TRUE;
 225             }
 226             else if (lstrcmpiW(option, L"flush") == 0)
 227             {
 228                 flush = TRUE;
 229             }
 230             else if (lstrcmpiW(option, L"ignore") == 0)
 231             {
 232                 ignore = TRUE;
 233             }
 234             else
 235             {
 236                 wchar_t *const message = concat_3(L"[tee] Error: Invalid option \"--", option, L"\" encountered!\n");
 237                 if (message)
 238                 {
 239                     write_text(hStdErr, message);
 240                     LocalFree(message);
 241                 }
 242                 return 1;
 243             }
 244         }
 245         else
 246         {
 247             break; /* stop option processing */
 248         }
 249     }
 250
 251     /* Check output file name */
 252     if (argOff >= argc)
 253     {
 254         write_text(hStdErr, L"[tee] Error: Output file name is missing!\n");
 255         return 1;
 256     }
 257
 258     /* Create events */
 259     if (!(hEventStop = CreateEventW(NULL, TRUE, FALSE, NULL)))
 260     {
 261         write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 262         goto cleanup;
 263     }
 264     for (size_t i = 0U; i < 2U; ++i)
 265     {
 266         if (!(hEventThrdReady[i] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 267         {
 268             write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 269             goto cleanup;
 270         }
 271         if (!(hEventCompleted[i] = CreateEventW(NULL, FALSE, FALSE, NULL)))
 272         {
 273             write_text(hStdErr, L"[tee] System error: Failed to create event!\n\n");
 274             goto cleanup;
 275         }
 276     }
 277
 278     /* Open output file */
 279     if ((hMyFile = CreateFileW(argv[argOff], GENERIC_WRITE, FILE_SHARE_READ, NULL, append ? OPEN_ALWAYS : CREATE_ALWAYS, 0U, NULL)) == INVALID_HANDLE_VALUE)
 280     {
 281         write_text(hStdErr, L"[tee] Error: Failed to open the output file for writing!\n");
 282         goto cleanup;
 283     }
 284
 285     /* Seek to the end of the file */
 286     if (append)
 287     {
 288         LARGE_INTEGER offset = { .QuadPart = 0LL };
 289         if (!SetFilePointerEx(hMyFile, offset, NULL, FILE_END))
 290         {
 291             write_text(hStdErr, L"[tee] Error: Failed to move the file pointer to the end of the file!\n");
 292             goto cleanup;
 293         }
 294     }
 295
 296     /* Set up thread data */
 297     for (size_t i = 0; i < 2U; ++i)
 298     {
 299         threadData[i].hOutput = (i > 0U) ? hMyFile : hStdOut;
 300         threadData[i].hError = hStdErr;
 301         threadData[i].flush = flush && (!is_terminal(threadData[i].hOutput));
 302         threadData[i].hEventReady[0U] = hEventThrdReady[i];
 303         threadData[i].hEventReady[1U] = hEventStop;
 304         threadData[i].hEventCompleted = hEventCompleted[i];
 305     }
 306
 307     /* Start threads */
 308     for (size_t i = 0; i < 2U; ++i)
 309     {
 310         if (!(hThreads[i] = CreateThread(NULL, 0U, writer_thread_start_routine, &threadData[i], 0U, NULL)))
 311         {
 312             write_text(hStdErr, L"[tee] System error: Failed to create thread!\n");
 313             goto cleanup;
 314         }
 315     }
 316
 317     /* Are we reading from a pipe? */
 318     const BOOL isPipeInput = (GetFileType(hStdIn) == FILE_TYPE_PIPE);
 319
 320     /* Initialize index */
 321     ULONG_PTR myIndex = 1U - index;
 322
 323     /* Process all input from STDIN stream */
 324     do
 325     {
 326         for (size_t i = 0U; i < 2U; ++i)
 327         {
 328             if (!SetEvent(hEventThrdReady[i]))
 329             {
 330                 write_text(hStdErr, L"[tee] System error: Failed to signal event!\n");
 331                 goto cleanup;
 332             }
 333         }
 334
 335         if (!ReadFile(hStdIn, buffer[myIndex], BUFFSIZE, &bytesTotal[myIndex], NULL))
 336         {
 337             if (GetLastError() != ERROR_BROKEN_PIPE)
 338             {
 339                 write_text(hStdErr, L"[tee] Error: Failed to read input data!\n");
 340                 goto cleanup;
 341             }
 342             break;
 343         }
 344
 345         if ((!bytesTotal[myIndex]) && (!isPipeInput)) /*pipes may return zero bytes, even when more data can become available later!*/
 346         {
 347             break;
 348         }
 349
 350         const DWORD waitResult = WaitForMultipleObjects(2U, hEventCompleted, TRUE, INFINITE);
 351         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 352         {
 353             write_text(hStdErr, L"[tee] System error: Failed to wait for events!\n");
 354             goto cleanup;
 355         }
 356
 357         myIndex = (ULONG_PTR) InterlockedExchangePointer((PVOID*)&index, (PVOID)myIndex);
 358     }
 359     while ((!g_stop) || ignore);
 360
 361     exitCode = 0;
 362
 363 cleanup:
 364
 365     /* Stop the worker threads */
 366     if (hEventStop)
 367     {
 368         SetEvent(hEventStop);
 369     }
 370
 371     /* Wait for worker threads to exit */
 372     if (hThreads[0U])
 373     {
 374         const DWORD waitResult = WaitForMultipleObjects(hThreads[1U] ? 2U : 1U, hThreads, TRUE, 12000U);
 375         if ((waitResult != WAIT_OBJECT_0) && (waitResult != WAIT_OBJECT_0 + 1U))
 376         {
 377             for (DWORD i = 0U; i < (hThreads[1U] ? 2U : 1U); ++i)
 378             {
 379                 if (WaitForSingleObject(hThreads[i], 125U) != WAIT_OBJECT_0)
 380                 {
 381                     write_text(hStdErr, L"[tee] Error: Worker thread did not exit cleanly!\n");
 382                     TerminateThread(hThreads[i], 1U);
 383                 }
 384             }
 385         }
 386     }
 387
 388     /* Close worker threads */
 389     for (DWORD i = 0U; i < 2U; ++i)
 390     {
 391         SAFE_CLOSE_HANDLE(hThreads[i]);
 392     }
 393
 394     /* Close output file */
 395     if (hMyFile != INVALID_HANDLE_VALUE)
 396     {
 397         if (flush)
 398         {
 399             FlushFileBuffers(hMyFile);
 400         }
 401         CloseHandle(hMyFile);
 402     }
 403
 404     /* Close events */
 405     for (size_t i = 0U; i < 2U; ++i)
 406     {
 407         SAFE_CLOSE_HANDLE(hEventThrdReady[i]);
 408         SAFE_CLOSE_HANDLE(hEventCompleted[i]);
 409     }
 410     SAFE_CLOSE_HANDLE(hEventStop);
 411
 412     /* Exit */
 413     return exitCode;
 414 }
 415
 416 // --------------------------------------------------------------------------
 417 // CRT Startup
 418 // --------------------------------------------------------------------------
 419
 420 #pragma warning(disable: 4702)
 421
 422 int wmainCRTStartup(void)
 423 {
 424     SetErrorMode(SEM_FAILCRITICALERRORS);
 425
 426     int nArgs;
 427     LPWSTR *const szArglist = CommandLineToArgvW(GetCommandLineW(), &nArgs);
 428     if (!szArglist)
 429     {
 430         FatalExit(-1);
 431     }
 432
 433     const int retval = wmain(nArgs, szArglist);
 434     LocalFree(szArglist);
 435     ExitProcess((UINT)retval);
 436
 437     return 0;
 438 }