Changeset 3547 for trunk/src/lib/maybe_con_fwrite.c

Timestamp:

Jan 29, 2022, 3:39:47 AM (3 years ago)

Author:

bird

Message:

lib: Changed the console-optimization wrappers to use the get_crt_codepage() function as that's more appropriate than GetConsoleCP() for use with _cputws(). Also made them avoid the heap for smaller writes that could fit into a reasonable (2KiB) stack buffer.

File:

• trunk/src/lib/maybe_con_fwrite.c (modified) (2 diffs)

trunk/src/lib/maybe_con_fwrite.c

@@ -57 +57 @@
 #ifdef KBUILD_OS_WINDOWS
     /*
-     * If it's a TTY, do our own conversion to wide char and
-     * call WriteConsoleW directly.
+     * If it's a TTY, do our own conversion to wide char and call _cputws.
      */
     if (   cbUnit > 0
         && cUnits > 0
+        && cbUnit < (unsigned)INT_MAX / 4
+        && cUnits < (unsigned)INT_MAX / 4
         && (pFile == stdout || pFile == stderr))
     {
@@ -73 +74 @@
                 if (is_console_handle((intptr_t)hCon))
                 {
-                    size_t   cbToWrite = cbUnit * cUnits;
-                    size_t   cwcTmp    = cbToWrite * 2 + 16;
-                    wchar_t *pawcTmp   = (wchar_t *)malloc(cwcTmp * sizeof(wchar_t));
-                    if (pawcTmp)
+                    /* Use a stack buffer if we can, falling back on the heap for larger writes: */
+                    wchar_t  awcBuf[1024];
+                    wchar_t *pawcBuf;
+                    wchar_t *pawcBufFree = NULL;
+                    size_t   cbToWrite   = cbUnit * cUnits;
+                    size_t   cwcBuf      = cbToWrite * 2 + 16;
+                    if (cwcBuf < sizeof(awcBuf) / sizeof(awcBuf[0]))
                     {
-                        int           cwcToWrite;
-                        static UINT s_uConsoleCp = 0;
-                        if (s_uConsoleCp == 0)
-                            s_uConsoleCp = GetConsoleCP();
-
-                        cwcToWrite = MultiByteToWideChar(s_uConsoleCp, 0 /*dwFlags*/, pvBuf, (int)cbToWrite,
-                                                         pawcTmp, (int)(cwcTmp - 1));
+                        pawcBuf = awcBuf;
+                        cwcBuf  = sizeof(awcBuf) / sizeof(awcBuf[0]);
+                    }
+                    else
+                        pawcBufFree = pawcBuf = (wchar_t *)malloc(cwcBuf * sizeof(wchar_t));
+                    if (pawcBuf)
+                    {
+                        int cwcToWrite = MultiByteToWideChar(get_crt_codepage(), 0 /*dwFlags*/,
+                                                             pvBuf, (int)cbToWrite,
+                                                             pawcBuf, (int)(cwcBuf - 1));
                         if (cwcToWrite > 0)
                         {
                             int rc;
-                            pawcTmp[cwcToWrite] = '\0';
+                            pawcBuf[cwcToWrite] = '\0';
 
                             /* Let the CRT do the rest.  At least the Visual C++ 2010 CRT
                                sources indicates _cputws will do the right thing.  */
                             fflush(pFile);
-                            rc = _cputws(pawcTmp);
-                            free(pawcTmp);
+                            rc = _cputws(pawcBuf);
+                            if (pawcBufFree)
+                                free(pawcBufFree);
                             if (rc >= 0)
                                 return cUnits;
                             return 0;
                         }
-                        free(pawcTmp);
+                        free(pawcBufFree);
                     }
                 }