• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2     An implementation of Windows console I/O
3 
4     Classes defined here: _WindowsConsoleIO
5 
6     Written by Steve Dower
7 */
8 
9 #define PY_SSIZE_T_CLEAN
10 #include "Python.h"
11 
12 #ifdef MS_WINDOWS
13 
14 #include "structmember.h"
15 #ifdef HAVE_SYS_TYPES_H
16 #include <sys/types.h>
17 #endif
18 #ifdef HAVE_SYS_STAT_H
19 #include <sys/stat.h>
20 #endif
21 #include <stddef.h> /* For offsetof */
22 
23 #define WIN32_LEAN_AND_MEAN
24 #include <windows.h>
25 #include <fcntl.h>
26 
27 #include "_iomodule.h"
28 
29 /* BUFSIZ determines how many characters can be typed at the console
30    before it starts blocking. */
31 #if BUFSIZ < (16*1024)
32 #define SMALLCHUNK (2*1024)
33 #elif (BUFSIZ >= (2 << 25))
34 #error "unreasonable BUFSIZ > 64 MiB defined"
35 #else
36 #define SMALLCHUNK BUFSIZ
37 #endif
38 
39 /* BUFMAX determines how many bytes can be read in one go. */
40 #define BUFMAX (32*1024*1024)
41 
42 /* SMALLBUF determines how many utf-8 characters will be
43    buffered within the stream, in order to support reads
44    of less than one character */
45 #define SMALLBUF 4
46 
_get_console_type(HANDLE handle)47 char _get_console_type(HANDLE handle) {
48     DWORD mode, peek_count;
49 
50     if (handle == INVALID_HANDLE_VALUE)
51         return '\0';
52 
53     if (!GetConsoleMode(handle, &mode))
54         return '\0';
55 
56     /* Peek at the handle to see whether it is an input or output handle */
57     if (GetNumberOfConsoleInputEvents(handle, &peek_count))
58         return 'r';
59     return 'w';
60 }
61 
_PyIO_get_console_type(PyObject * path_or_fd)62 char _PyIO_get_console_type(PyObject *path_or_fd) {
63     int fd = PyLong_AsLong(path_or_fd);
64     PyErr_Clear();
65     if (fd >= 0) {
66         HANDLE handle;
67         _Py_BEGIN_SUPPRESS_IPH
68         handle = (HANDLE)_get_osfhandle(fd);
69         _Py_END_SUPPRESS_IPH
70         if (handle == INVALID_HANDLE_VALUE)
71             return '\0';
72         return _get_console_type(handle);
73     }
74 
75     PyObject *decoded;
76     wchar_t *decoded_wstr;
77 
78     if (!PyUnicode_FSDecoder(path_or_fd, &decoded)) {
79         PyErr_Clear();
80         return '\0';
81     }
82     decoded_wstr = PyUnicode_AsWideCharString(decoded, NULL);
83     Py_CLEAR(decoded);
84     if (!decoded_wstr) {
85         PyErr_Clear();
86         return '\0';
87     }
88 
89     char m = '\0';
90     if (!_wcsicmp(decoded_wstr, L"CONIN$")) {
91         m = 'r';
92     } else if (!_wcsicmp(decoded_wstr, L"CONOUT$")) {
93         m = 'w';
94     } else if (!_wcsicmp(decoded_wstr, L"CON")) {
95         m = 'x';
96     }
97     if (m) {
98         PyMem_Free(decoded_wstr);
99         return m;
100     }
101 
102     DWORD length;
103     wchar_t name_buf[MAX_PATH], *pname_buf = name_buf;
104 
105     length = GetFullPathNameW(decoded_wstr, MAX_PATH, pname_buf, NULL);
106     if (length > MAX_PATH) {
107         pname_buf = PyMem_New(wchar_t, length);
108         if (pname_buf)
109             length = GetFullPathNameW(decoded_wstr, length, pname_buf, NULL);
110         else
111             length = 0;
112     }
113     PyMem_Free(decoded_wstr);
114 
115     if (length) {
116         wchar_t *name = pname_buf;
117         if (length >= 4 && name[3] == L'\\' &&
118             (name[2] == L'.' || name[2] == L'?') &&
119             name[1] == L'\\' && name[0] == L'\\') {
120             name += 4;
121         }
122         if (!_wcsicmp(name, L"CONIN$")) {
123             m = 'r';
124         } else if (!_wcsicmp(name, L"CONOUT$")) {
125             m = 'w';
126         } else if (!_wcsicmp(name, L"CON")) {
127             m = 'x';
128         }
129     }
130 
131     if (pname_buf != name_buf)
132         PyMem_Free(pname_buf);
133     return m;
134 }
135 
136 
137 /*[clinic input]
138 module _io
139 class _io._WindowsConsoleIO "winconsoleio *" "&PyWindowsConsoleIO_Type"
140 [clinic start generated code]*/
141 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=e897fdc1fba4e131]*/
142 
143 typedef struct {
144     PyObject_HEAD
145     HANDLE handle;
146     int fd;
147     unsigned int created : 1;
148     unsigned int readable : 1;
149     unsigned int writable : 1;
150     unsigned int closehandle : 1;
151     char finalizing;
152     unsigned int blksize;
153     PyObject *weakreflist;
154     PyObject *dict;
155     char buf[SMALLBUF];
156     wchar_t wbuf;
157 } winconsoleio;
158 
159 PyTypeObject PyWindowsConsoleIO_Type;
160 
161 _Py_IDENTIFIER(name);
162 
163 int
_PyWindowsConsoleIO_closed(PyObject * self)164 _PyWindowsConsoleIO_closed(PyObject *self)
165 {
166     return ((winconsoleio *)self)->handle == INVALID_HANDLE_VALUE;
167 }
168 
169 
170 /* Returns 0 on success, -1 with exception set on failure. */
171 static int
internal_close(winconsoleio * self)172 internal_close(winconsoleio *self)
173 {
174     if (self->handle != INVALID_HANDLE_VALUE) {
175         if (self->closehandle) {
176             if (self->fd >= 0) {
177                 _Py_BEGIN_SUPPRESS_IPH
178                 close(self->fd);
179                 _Py_END_SUPPRESS_IPH
180             }
181             CloseHandle(self->handle);
182         }
183         self->handle = INVALID_HANDLE_VALUE;
184         self->fd = -1;
185     }
186     return 0;
187 }
188 
189 /*[clinic input]
190 _io._WindowsConsoleIO.close
191 
192 Close the handle.
193 
194 A closed handle cannot be used for further I/O operations.  close() may be
195 called more than once without error.
196 [clinic start generated code]*/
197 
198 static PyObject *
_io__WindowsConsoleIO_close_impl(winconsoleio * self)199 _io__WindowsConsoleIO_close_impl(winconsoleio *self)
200 /*[clinic end generated code: output=27ef95b66c29057b input=185617e349ae4c7b]*/
201 {
202     PyObject *res;
203     PyObject *exc, *val, *tb;
204     int rc;
205     _Py_IDENTIFIER(close);
206     res = _PyObject_CallMethodIdObjArgs((PyObject*)&PyRawIOBase_Type,
207                                         &PyId_close, self, NULL);
208     if (!self->closehandle) {
209         self->handle = INVALID_HANDLE_VALUE;
210         return res;
211     }
212     if (res == NULL)
213         PyErr_Fetch(&exc, &val, &tb);
214     rc = internal_close(self);
215     if (res == NULL)
216         _PyErr_ChainExceptions(exc, val, tb);
217     if (rc < 0)
218         Py_CLEAR(res);
219     return res;
220 }
221 
222 static PyObject *
winconsoleio_new(PyTypeObject * type,PyObject * args,PyObject * kwds)223 winconsoleio_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
224 {
225     winconsoleio *self;
226 
227     assert(type != NULL && type->tp_alloc != NULL);
228 
229     self = (winconsoleio *) type->tp_alloc(type, 0);
230     if (self != NULL) {
231         self->handle = INVALID_HANDLE_VALUE;
232         self->fd = -1;
233         self->created = 0;
234         self->readable = 0;
235         self->writable = 0;
236         self->closehandle = 0;
237         self->blksize = 0;
238         self->weakreflist = NULL;
239     }
240 
241     return (PyObject *) self;
242 }
243 
244 /*[clinic input]
245 _io._WindowsConsoleIO.__init__
246     file as nameobj: object
247     mode: str = "r"
248     closefd: bool(accept={int}) = True
249     opener: object = None
250 
251 Open a console buffer by file descriptor.
252 
253 The mode can be 'rb' (default), or 'wb' for reading or writing bytes. All
254 other mode characters will be ignored. Mode 'b' will be assumed if it is
255 omitted. The *opener* parameter is always ignored.
256 [clinic start generated code]*/
257 
258 static int
_io__WindowsConsoleIO___init___impl(winconsoleio * self,PyObject * nameobj,const char * mode,int closefd,PyObject * opener)259 _io__WindowsConsoleIO___init___impl(winconsoleio *self, PyObject *nameobj,
260                                     const char *mode, int closefd,
261                                     PyObject *opener)
262 /*[clinic end generated code: output=3fd9cbcdd8d95429 input=06ae4b863c63244b]*/
263 {
264     const char *s;
265     wchar_t *name = NULL;
266     char console_type = '\0';
267     int ret = 0;
268     int rwa = 0;
269     int fd = -1;
270     int fd_is_own = 0;
271 
272     assert(PyWindowsConsoleIO_Check(self));
273     if (self->handle >= 0) {
274         if (self->closehandle) {
275             /* Have to close the existing file first. */
276             if (internal_close(self) < 0)
277                 return -1;
278         }
279         else
280             self->handle = INVALID_HANDLE_VALUE;
281     }
282 
283     if (PyFloat_Check(nameobj)) {
284         PyErr_SetString(PyExc_TypeError,
285                         "integer argument expected, got float");
286         return -1;
287     }
288 
289     fd = _PyLong_AsInt(nameobj);
290     if (fd < 0) {
291         if (!PyErr_Occurred()) {
292             PyErr_SetString(PyExc_ValueError,
293                             "negative file descriptor");
294             return -1;
295         }
296         PyErr_Clear();
297     }
298     self->fd = fd;
299 
300     if (fd < 0) {
301         PyObject *decodedname;
302 
303         int d = PyUnicode_FSDecoder(nameobj, (void*)&decodedname);
304         if (!d)
305             return -1;
306 
307         name = PyUnicode_AsWideCharString(decodedname, NULL);
308         console_type = _PyIO_get_console_type(decodedname);
309         Py_CLEAR(decodedname);
310         if (name == NULL)
311             return -1;
312     }
313 
314     s = mode;
315     while (*s) {
316         switch (*s++) {
317         case '+':
318         case 'a':
319         case 'b':
320         case 'x':
321             break;
322         case 'r':
323             if (rwa)
324                 goto bad_mode;
325             rwa = 1;
326             self->readable = 1;
327             if (console_type == 'x')
328                 console_type = 'r';
329             break;
330         case 'w':
331             if (rwa)
332                 goto bad_mode;
333             rwa = 1;
334             self->writable = 1;
335             if (console_type == 'x')
336                 console_type = 'w';
337             break;
338         default:
339             PyErr_Format(PyExc_ValueError,
340                          "invalid mode: %.200s", mode);
341             goto error;
342         }
343     }
344 
345     if (!rwa)
346         goto bad_mode;
347 
348     if (fd >= 0) {
349         _Py_BEGIN_SUPPRESS_IPH
350         self->handle = (HANDLE)_get_osfhandle(fd);
351         _Py_END_SUPPRESS_IPH
352         self->closehandle = 0;
353     } else {
354         DWORD access = GENERIC_READ;
355 
356         self->closehandle = 1;
357         if (!closefd) {
358             PyErr_SetString(PyExc_ValueError,
359                 "Cannot use closefd=False with file name");
360             goto error;
361         }
362 
363         if (self->writable)
364             access = GENERIC_WRITE;
365 
366         Py_BEGIN_ALLOW_THREADS
367         /* Attempt to open for read/write initially, then fall back
368            on the specific access. This is required for modern names
369            CONIN$ and CONOUT$, which allow reading/writing state as
370            well as reading/writing content. */
371         self->handle = CreateFileW(name, GENERIC_READ | GENERIC_WRITE,
372             FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
373         if (self->handle == INVALID_HANDLE_VALUE)
374             self->handle = CreateFileW(name, access,
375                 FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
376         Py_END_ALLOW_THREADS
377 
378         if (self->handle == INVALID_HANDLE_VALUE) {
379             PyErr_SetExcFromWindowsErrWithFilenameObject(PyExc_OSError, GetLastError(), nameobj);
380             goto error;
381         }
382     }
383 
384     if (console_type == '\0')
385         console_type = _get_console_type(self->handle);
386 
387     if (self->writable && console_type != 'w') {
388         PyErr_SetString(PyExc_ValueError,
389             "Cannot open console input buffer for writing");
390         goto error;
391     }
392     if (self->readable && console_type != 'r') {
393         PyErr_SetString(PyExc_ValueError,
394             "Cannot open console output buffer for reading");
395         goto error;
396     }
397 
398     self->blksize = DEFAULT_BUFFER_SIZE;
399     memset(self->buf, 0, 4);
400 
401     if (_PyObject_SetAttrId((PyObject *)self, &PyId_name, nameobj) < 0)
402         goto error;
403 
404     goto done;
405 
406 bad_mode:
407     PyErr_SetString(PyExc_ValueError,
408                     "Must have exactly one of read or write mode");
409 error:
410     ret = -1;
411     internal_close(self);
412 
413 done:
414     if (name)
415         PyMem_Free(name);
416     return ret;
417 }
418 
419 static int
winconsoleio_traverse(winconsoleio * self,visitproc visit,void * arg)420 winconsoleio_traverse(winconsoleio *self, visitproc visit, void *arg)
421 {
422     Py_VISIT(self->dict);
423     return 0;
424 }
425 
426 static int
winconsoleio_clear(winconsoleio * self)427 winconsoleio_clear(winconsoleio *self)
428 {
429     Py_CLEAR(self->dict);
430     return 0;
431 }
432 
433 static void
winconsoleio_dealloc(winconsoleio * self)434 winconsoleio_dealloc(winconsoleio *self)
435 {
436     self->finalizing = 1;
437     if (_PyIOBase_finalize((PyObject *) self) < 0)
438         return;
439     _PyObject_GC_UNTRACK(self);
440     if (self->weakreflist != NULL)
441         PyObject_ClearWeakRefs((PyObject *) self);
442     Py_CLEAR(self->dict);
443     Py_TYPE(self)->tp_free((PyObject *)self);
444 }
445 
446 static PyObject *
err_closed(void)447 err_closed(void)
448 {
449     PyErr_SetString(PyExc_ValueError, "I/O operation on closed file");
450     return NULL;
451 }
452 
453 static PyObject *
err_mode(const char * action)454 err_mode(const char *action)
455 {
456     _PyIO_State *state = IO_STATE();
457     if (state != NULL)
458         PyErr_Format(state->unsupported_operation,
459                      "Console buffer does not support %s", action);
460     return NULL;
461 }
462 
463 /*[clinic input]
464 _io._WindowsConsoleIO.fileno
465 
466 Return the underlying file descriptor (an integer).
467 
468 fileno is only set when a file descriptor is used to open
469 one of the standard streams.
470 
471 [clinic start generated code]*/
472 
473 static PyObject *
_io__WindowsConsoleIO_fileno_impl(winconsoleio * self)474 _io__WindowsConsoleIO_fileno_impl(winconsoleio *self)
475 /*[clinic end generated code: output=006fa74ce3b5cfbf input=079adc330ddaabe6]*/
476 {
477     if (self->fd < 0 && self->handle != INVALID_HANDLE_VALUE) {
478         _Py_BEGIN_SUPPRESS_IPH
479         if (self->writable)
480             self->fd = _open_osfhandle((intptr_t)self->handle, _O_WRONLY | _O_BINARY);
481         else
482             self->fd = _open_osfhandle((intptr_t)self->handle, _O_RDONLY | _O_BINARY);
483         _Py_END_SUPPRESS_IPH
484     }
485     if (self->fd < 0)
486         return err_mode("fileno");
487     return PyLong_FromLong(self->fd);
488 }
489 
490 /*[clinic input]
491 _io._WindowsConsoleIO.readable
492 
493 True if console is an input buffer.
494 [clinic start generated code]*/
495 
496 static PyObject *
_io__WindowsConsoleIO_readable_impl(winconsoleio * self)497 _io__WindowsConsoleIO_readable_impl(winconsoleio *self)
498 /*[clinic end generated code: output=daf9cef2743becf0 input=6be9defb5302daae]*/
499 {
500     if (self->handle == INVALID_HANDLE_VALUE)
501         return err_closed();
502     return PyBool_FromLong((long) self->readable);
503 }
504 
505 /*[clinic input]
506 _io._WindowsConsoleIO.writable
507 
508 True if console is an output buffer.
509 [clinic start generated code]*/
510 
511 static PyObject *
_io__WindowsConsoleIO_writable_impl(winconsoleio * self)512 _io__WindowsConsoleIO_writable_impl(winconsoleio *self)
513 /*[clinic end generated code: output=e0a2ad7eae5abf67 input=cefbd8abc24df6a0]*/
514 {
515     if (self->handle == INVALID_HANDLE_VALUE)
516         return err_closed();
517     return PyBool_FromLong((long) self->writable);
518 }
519 
520 static DWORD
_buflen(winconsoleio * self)521 _buflen(winconsoleio *self)
522 {
523     for (DWORD i = 0; i < SMALLBUF; ++i) {
524         if (!self->buf[i])
525             return i;
526     }
527     return SMALLBUF;
528 }
529 
530 static DWORD
_copyfrombuf(winconsoleio * self,char * buf,DWORD len)531 _copyfrombuf(winconsoleio *self, char *buf, DWORD len)
532 {
533     DWORD n = 0;
534 
535     while (self->buf[0] && len--) {
536         buf[n++] = self->buf[0];
537         for (int i = 1; i < SMALLBUF; ++i)
538             self->buf[i - 1] = self->buf[i];
539         self->buf[SMALLBUF - 1] = 0;
540     }
541 
542     return n;
543 }
544 
545 static wchar_t *
read_console_w(HANDLE handle,DWORD maxlen,DWORD * readlen)546 read_console_w(HANDLE handle, DWORD maxlen, DWORD *readlen) {
547     int err = 0, sig = 0;
548 
549     wchar_t *buf = (wchar_t*)PyMem_Malloc(maxlen * sizeof(wchar_t));
550     if (!buf)
551         goto error;
552 
553     *readlen = 0;
554 
555     //DebugBreak();
556     Py_BEGIN_ALLOW_THREADS
557     DWORD off = 0;
558     while (off < maxlen) {
559         DWORD n = (DWORD)-1;
560         DWORD len = min(maxlen - off, BUFSIZ);
561         SetLastError(0);
562         BOOL res = ReadConsoleW(handle, &buf[off], len, &n, NULL);
563 
564         if (!res) {
565             err = GetLastError();
566             break;
567         }
568         if (n == (DWORD)-1 && (err = GetLastError()) == ERROR_OPERATION_ABORTED) {
569             break;
570         }
571         if (n == 0) {
572             err = GetLastError();
573             if (err != ERROR_OPERATION_ABORTED)
574                 break;
575             err = 0;
576             HANDLE hInterruptEvent = _PyOS_SigintEvent();
577             if (WaitForSingleObjectEx(hInterruptEvent, 100, FALSE)
578                     == WAIT_OBJECT_0) {
579                 ResetEvent(hInterruptEvent);
580                 Py_BLOCK_THREADS
581                 sig = PyErr_CheckSignals();
582                 Py_UNBLOCK_THREADS
583                 if (sig < 0)
584                     break;
585             }
586         }
587         *readlen += n;
588 
589         /* If we didn't read a full buffer that time, don't try
590            again or we will block a second time. */
591         if (n < len)
592             break;
593         /* If the buffer ended with a newline, break out */
594         if (buf[*readlen - 1] == '\n')
595             break;
596         /* If the buffer ends with a high surrogate, expand the
597            buffer and read an extra character. */
598         WORD char_type;
599         if (off + BUFSIZ >= maxlen &&
600             GetStringTypeW(CT_CTYPE3, &buf[*readlen - 1], 1, &char_type) &&
601             char_type == C3_HIGHSURROGATE) {
602             wchar_t *newbuf;
603             maxlen += 1;
604             Py_BLOCK_THREADS
605             newbuf = (wchar_t*)PyMem_Realloc(buf, maxlen * sizeof(wchar_t));
606             Py_UNBLOCK_THREADS
607             if (!newbuf) {
608                 sig = -1;
609                 break;
610             }
611             buf = newbuf;
612             /* Only advance by n and not BUFSIZ in this case */
613             off += n;
614             continue;
615         }
616 
617         off += BUFSIZ;
618     }
619 
620     Py_END_ALLOW_THREADS
621 
622     if (sig)
623         goto error;
624     if (err) {
625         PyErr_SetFromWindowsErr(err);
626         goto error;
627     }
628 
629     if (*readlen > 0 && buf[0] == L'\x1a') {
630         PyMem_Free(buf);
631         buf = (wchar_t *)PyMem_Malloc(sizeof(wchar_t));
632         if (!buf)
633             goto error;
634         buf[0] = L'\0';
635         *readlen = 0;
636     }
637 
638     return buf;
639 
640 error:
641     if (buf)
642         PyMem_Free(buf);
643     return NULL;
644 }
645 
646 
647 static Py_ssize_t
readinto(winconsoleio * self,char * buf,Py_ssize_t len)648 readinto(winconsoleio *self, char *buf, Py_ssize_t len)
649 {
650     if (self->handle == INVALID_HANDLE_VALUE) {
651         err_closed();
652         return -1;
653     }
654     if (!self->readable) {
655         err_mode("reading");
656         return -1;
657     }
658     if (len == 0)
659         return 0;
660     if (len > BUFMAX) {
661         PyErr_Format(PyExc_ValueError, "cannot read more than %d bytes", BUFMAX);
662         return -1;
663     }
664 
665     /* Each character may take up to 4 bytes in the final buffer.
666        This is highly conservative, but necessary to avoid
667        failure for any given Unicode input (e.g. \U0010ffff).
668        If the caller requests fewer than 4 bytes, we buffer one
669        character.
670     */
671     DWORD wlen = (DWORD)(len / 4);
672     if (wlen == 0) {
673         wlen = 1;
674     }
675 
676     DWORD read_len = _copyfrombuf(self, buf, (DWORD)len);
677     if (read_len) {
678         buf = &buf[read_len];
679         len -= read_len;
680         wlen -= 1;
681     }
682     if (len == read_len || wlen == 0)
683         return read_len;
684 
685     DWORD n;
686     wchar_t *wbuf = read_console_w(self->handle, wlen, &n);
687     if (wbuf == NULL)
688         return -1;
689     if (n == 0) {
690         PyMem_Free(wbuf);
691         return read_len;
692     }
693 
694     int err = 0;
695     DWORD u8n = 0;
696 
697     Py_BEGIN_ALLOW_THREADS
698     if (len < 4) {
699         if (WideCharToMultiByte(CP_UTF8, 0, wbuf, n,
700                 self->buf, sizeof(self->buf) / sizeof(self->buf[0]),
701                 NULL, NULL))
702             u8n = _copyfrombuf(self, buf, (DWORD)len);
703     } else {
704         u8n = WideCharToMultiByte(CP_UTF8, 0, wbuf, n,
705             buf, (DWORD)len, NULL, NULL);
706     }
707 
708     if (u8n) {
709         read_len += u8n;
710         u8n = 0;
711     } else {
712         err = GetLastError();
713         if (err == ERROR_INSUFFICIENT_BUFFER) {
714             /* Calculate the needed buffer for a more useful error, as this
715                 means our "/ 4" logic above is insufficient for some input.
716             */
717             u8n = WideCharToMultiByte(CP_UTF8, 0, wbuf, n,
718                 NULL, 0, NULL, NULL);
719         }
720     }
721     Py_END_ALLOW_THREADS
722 
723     PyMem_Free(wbuf);
724 
725     if (u8n) {
726         PyErr_Format(PyExc_SystemError,
727             "Buffer had room for %d bytes but %d bytes required",
728             len, u8n);
729         return -1;
730     }
731     if (err) {
732         PyErr_SetFromWindowsErr(err);
733         return -1;
734     }
735 
736     return read_len;
737 }
738 
739 /*[clinic input]
740 _io._WindowsConsoleIO.readinto
741     buffer: Py_buffer(accept={rwbuffer})
742     /
743 
744 Same as RawIOBase.readinto().
745 [clinic start generated code]*/
746 
747 static PyObject *
_io__WindowsConsoleIO_readinto_impl(winconsoleio * self,Py_buffer * buffer)748 _io__WindowsConsoleIO_readinto_impl(winconsoleio *self, Py_buffer *buffer)
749 /*[clinic end generated code: output=66d1bdfa3f20af39 input=4ed68da48a6baffe]*/
750 {
751     Py_ssize_t len = readinto(self, buffer->buf, buffer->len);
752     if (len < 0)
753         return NULL;
754 
755     return PyLong_FromSsize_t(len);
756 }
757 
758 static DWORD
new_buffersize(winconsoleio * self,DWORD currentsize)759 new_buffersize(winconsoleio *self, DWORD currentsize)
760 {
761     DWORD addend;
762 
763     /* Expand the buffer by an amount proportional to the current size,
764        giving us amortized linear-time behavior.  For bigger sizes, use a
765        less-than-double growth factor to avoid excessive allocation. */
766     if (currentsize > 65536)
767         addend = currentsize >> 3;
768     else
769         addend = 256 + currentsize;
770     if (addend < SMALLCHUNK)
771         /* Avoid tiny read() calls. */
772         addend = SMALLCHUNK;
773     return addend + currentsize;
774 }
775 
776 /*[clinic input]
777 _io._WindowsConsoleIO.readall
778 
779 Read all data from the console, returned as bytes.
780 
781 Return an empty bytes object at EOF.
782 [clinic start generated code]*/
783 
784 static PyObject *
_io__WindowsConsoleIO_readall_impl(winconsoleio * self)785 _io__WindowsConsoleIO_readall_impl(winconsoleio *self)
786 /*[clinic end generated code: output=e6d312c684f6e23b input=4024d649a1006e69]*/
787 {
788     wchar_t *buf;
789     DWORD bufsize, n, len = 0;
790     PyObject *bytes;
791     DWORD bytes_size, rn;
792 
793     if (self->handle == INVALID_HANDLE_VALUE)
794         return err_closed();
795 
796     bufsize = BUFSIZ;
797 
798     buf = (wchar_t*)PyMem_Malloc((bufsize + 1) * sizeof(wchar_t));
799     if (buf == NULL)
800         return NULL;
801 
802     while (1) {
803         wchar_t *subbuf;
804 
805         if (len >= (Py_ssize_t)bufsize) {
806             DWORD newsize = new_buffersize(self, len);
807             if (newsize > BUFMAX)
808                 break;
809             if (newsize < bufsize) {
810                 PyErr_SetString(PyExc_OverflowError,
811                                 "unbounded read returned more bytes "
812                                 "than a Python bytes object can hold");
813                 PyMem_Free(buf);
814                 return NULL;
815             }
816             bufsize = newsize;
817 
818             wchar_t *tmp = PyMem_Realloc(buf,
819                                          (bufsize + 1) * sizeof(wchar_t));
820             if (tmp == NULL) {
821                 PyMem_Free(buf);
822                 return NULL;
823             }
824             buf = tmp;
825         }
826 
827         subbuf = read_console_w(self->handle, bufsize - len, &n);
828 
829         if (subbuf == NULL) {
830             PyMem_Free(buf);
831             return NULL;
832         }
833 
834         if (n > 0)
835             wcsncpy_s(&buf[len], bufsize - len + 1, subbuf, n);
836 
837         PyMem_Free(subbuf);
838 
839         /* when the read is empty we break */
840         if (n == 0)
841             break;
842 
843         len += n;
844     }
845 
846     if (len == 0 && _buflen(self) == 0) {
847         /* when the result starts with ^Z we return an empty buffer */
848         PyMem_Free(buf);
849         return PyBytes_FromStringAndSize(NULL, 0);
850     }
851 
852     if (len) {
853         Py_BEGIN_ALLOW_THREADS
854         bytes_size = WideCharToMultiByte(CP_UTF8, 0, buf, len,
855             NULL, 0, NULL, NULL);
856         Py_END_ALLOW_THREADS
857 
858         if (!bytes_size) {
859             DWORD err = GetLastError();
860             PyMem_Free(buf);
861             return PyErr_SetFromWindowsErr(err);
862         }
863     } else {
864         bytes_size = 0;
865     }
866 
867     bytes_size += _buflen(self);
868     bytes = PyBytes_FromStringAndSize(NULL, bytes_size);
869     rn = _copyfrombuf(self, PyBytes_AS_STRING(bytes), bytes_size);
870 
871     if (len) {
872         Py_BEGIN_ALLOW_THREADS
873         bytes_size = WideCharToMultiByte(CP_UTF8, 0, buf, len,
874             &PyBytes_AS_STRING(bytes)[rn], bytes_size - rn, NULL, NULL);
875         Py_END_ALLOW_THREADS
876 
877         if (!bytes_size) {
878             DWORD err = GetLastError();
879             PyMem_Free(buf);
880             Py_CLEAR(bytes);
881             return PyErr_SetFromWindowsErr(err);
882         }
883 
884         /* add back the number of preserved bytes */
885         bytes_size += rn;
886     }
887 
888     PyMem_Free(buf);
889     if (bytes_size < (size_t)PyBytes_GET_SIZE(bytes)) {
890         if (_PyBytes_Resize(&bytes, n * sizeof(wchar_t)) < 0) {
891             Py_CLEAR(bytes);
892             return NULL;
893         }
894     }
895     return bytes;
896 }
897 
898 /*[clinic input]
899 _io._WindowsConsoleIO.read
900     size: Py_ssize_t(accept={int, NoneType}) = -1
901     /
902 
903 Read at most size bytes, returned as bytes.
904 
905 Only makes one system call when size is a positive integer,
906 so less data may be returned than requested.
907 Return an empty bytes object at EOF.
908 [clinic start generated code]*/
909 
910 static PyObject *
_io__WindowsConsoleIO_read_impl(winconsoleio * self,Py_ssize_t size)911 _io__WindowsConsoleIO_read_impl(winconsoleio *self, Py_ssize_t size)
912 /*[clinic end generated code: output=57df68af9f4b22d0 input=8bc73bc15d0fa072]*/
913 {
914     PyObject *bytes;
915     Py_ssize_t bytes_size;
916 
917     if (self->handle == INVALID_HANDLE_VALUE)
918         return err_closed();
919     if (!self->readable)
920         return err_mode("reading");
921 
922     if (size < 0)
923         return _io__WindowsConsoleIO_readall_impl(self);
924     if (size > BUFMAX) {
925         PyErr_Format(PyExc_ValueError, "cannot read more than %d bytes", BUFMAX);
926         return NULL;
927     }
928 
929     bytes = PyBytes_FromStringAndSize(NULL, size);
930     if (bytes == NULL)
931         return NULL;
932 
933     bytes_size = readinto(self, PyBytes_AS_STRING(bytes), PyBytes_GET_SIZE(bytes));
934     if (bytes_size < 0) {
935         Py_CLEAR(bytes);
936         return NULL;
937     }
938 
939     if (bytes_size < PyBytes_GET_SIZE(bytes)) {
940         if (_PyBytes_Resize(&bytes, bytes_size) < 0) {
941             Py_CLEAR(bytes);
942             return NULL;
943         }
944     }
945 
946     return bytes;
947 }
948 
949 /*[clinic input]
950 _io._WindowsConsoleIO.write
951     b: Py_buffer
952     /
953 
954 Write buffer b to file, return number of bytes written.
955 
956 Only makes one system call, so not all of the data may be written.
957 The number of bytes actually written is returned.
958 [clinic start generated code]*/
959 
960 static PyObject *
_io__WindowsConsoleIO_write_impl(winconsoleio * self,Py_buffer * b)961 _io__WindowsConsoleIO_write_impl(winconsoleio *self, Py_buffer *b)
962 /*[clinic end generated code: output=775bdb16fbf9137b input=be35fb624f97c941]*/
963 {
964     BOOL res = TRUE;
965     wchar_t *wbuf;
966     DWORD len, wlen, n = 0;
967 
968     if (self->handle == INVALID_HANDLE_VALUE)
969         return err_closed();
970     if (!self->writable)
971         return err_mode("writing");
972 
973     if (!b->len) {
974         return PyLong_FromLong(0);
975     }
976     if (b->len > BUFMAX)
977         len = BUFMAX;
978     else
979         len = (DWORD)b->len;
980 
981     Py_BEGIN_ALLOW_THREADS
982     wlen = MultiByteToWideChar(CP_UTF8, 0, b->buf, len, NULL, 0);
983 
984     /* issue11395 there is an unspecified upper bound on how many bytes
985        can be written at once. We cap at 32k - the caller will have to
986        handle partial writes.
987        Since we don't know how many input bytes are being ignored, we
988        have to reduce and recalculate. */
989     while (wlen > 32766 / sizeof(wchar_t)) {
990         len /= 2;
991         wlen = MultiByteToWideChar(CP_UTF8, 0, b->buf, len, NULL, 0);
992     }
993     Py_END_ALLOW_THREADS
994 
995     if (!wlen)
996         return PyErr_SetFromWindowsErr(0);
997 
998     wbuf = (wchar_t*)PyMem_Malloc(wlen * sizeof(wchar_t));
999 
1000     Py_BEGIN_ALLOW_THREADS
1001     wlen = MultiByteToWideChar(CP_UTF8, 0, b->buf, len, wbuf, wlen);
1002     if (wlen) {
1003         res = WriteConsoleW(self->handle, wbuf, wlen, &n, NULL);
1004         if (res && n < wlen) {
1005             /* Wrote fewer characters than expected, which means our
1006              * len value may be wrong. So recalculate it from the
1007              * characters that were written. As this could potentially
1008              * result in a different value, we also validate that value.
1009              */
1010             len = WideCharToMultiByte(CP_UTF8, 0, wbuf, n,
1011                 NULL, 0, NULL, NULL);
1012             if (len) {
1013                 wlen = MultiByteToWideChar(CP_UTF8, 0, b->buf, len,
1014                     NULL, 0);
1015                 assert(wlen == len);
1016             }
1017         }
1018     } else
1019         res = 0;
1020     Py_END_ALLOW_THREADS
1021 
1022     if (!res) {
1023         DWORD err = GetLastError();
1024         PyMem_Free(wbuf);
1025         return PyErr_SetFromWindowsErr(err);
1026     }
1027 
1028     PyMem_Free(wbuf);
1029     return PyLong_FromSsize_t(len);
1030 }
1031 
1032 static PyObject *
winconsoleio_repr(winconsoleio * self)1033 winconsoleio_repr(winconsoleio *self)
1034 {
1035     if (self->handle == INVALID_HANDLE_VALUE)
1036         return PyUnicode_FromFormat("<_io._WindowsConsoleIO [closed]>");
1037 
1038     if (self->readable)
1039         return PyUnicode_FromFormat("<_io._WindowsConsoleIO mode='rb' closefd=%s>",
1040             self->closehandle ? "True" : "False");
1041     if (self->writable)
1042         return PyUnicode_FromFormat("<_io._WindowsConsoleIO mode='wb' closefd=%s>",
1043             self->closehandle ? "True" : "False");
1044 
1045     PyErr_SetString(PyExc_SystemError, "_WindowsConsoleIO has invalid mode");
1046     return NULL;
1047 }
1048 
1049 /*[clinic input]
1050 _io._WindowsConsoleIO.isatty
1051 
1052 Always True.
1053 [clinic start generated code]*/
1054 
1055 static PyObject *
_io__WindowsConsoleIO_isatty_impl(winconsoleio * self)1056 _io__WindowsConsoleIO_isatty_impl(winconsoleio *self)
1057 /*[clinic end generated code: output=9eac09d287c11bd7 input=9b91591dbe356f86]*/
1058 {
1059     if (self->handle == INVALID_HANDLE_VALUE)
1060         return err_closed();
1061 
1062     Py_RETURN_TRUE;
1063 }
1064 
1065 static PyObject *
winconsoleio_getstate(winconsoleio * self)1066 winconsoleio_getstate(winconsoleio *self)
1067 {
1068     PyErr_Format(PyExc_TypeError,
1069                  "cannot serialize '%s' object", Py_TYPE(self)->tp_name);
1070     return NULL;
1071 }
1072 
1073 #include "clinic/winconsoleio.c.h"
1074 
1075 static PyMethodDef winconsoleio_methods[] = {
1076     _IO__WINDOWSCONSOLEIO_READ_METHODDEF
1077     _IO__WINDOWSCONSOLEIO_READALL_METHODDEF
1078     _IO__WINDOWSCONSOLEIO_READINTO_METHODDEF
1079     _IO__WINDOWSCONSOLEIO_WRITE_METHODDEF
1080     _IO__WINDOWSCONSOLEIO_CLOSE_METHODDEF
1081     _IO__WINDOWSCONSOLEIO_READABLE_METHODDEF
1082     _IO__WINDOWSCONSOLEIO_WRITABLE_METHODDEF
1083     _IO__WINDOWSCONSOLEIO_FILENO_METHODDEF
1084     _IO__WINDOWSCONSOLEIO_ISATTY_METHODDEF
1085     {"__getstate__", (PyCFunction)winconsoleio_getstate, METH_NOARGS, NULL},
1086     {NULL,           NULL}             /* sentinel */
1087 };
1088 
1089 /* 'closed' and 'mode' are attributes for compatibility with FileIO. */
1090 
1091 static PyObject *
get_closed(winconsoleio * self,void * closure)1092 get_closed(winconsoleio *self, void *closure)
1093 {
1094     return PyBool_FromLong((long)(self->handle == INVALID_HANDLE_VALUE));
1095 }
1096 
1097 static PyObject *
get_closefd(winconsoleio * self,void * closure)1098 get_closefd(winconsoleio *self, void *closure)
1099 {
1100     return PyBool_FromLong((long)(self->closehandle));
1101 }
1102 
1103 static PyObject *
get_mode(winconsoleio * self,void * closure)1104 get_mode(winconsoleio *self, void *closure)
1105 {
1106     return PyUnicode_FromString(self->readable ? "rb" : "wb");
1107 }
1108 
1109 static PyGetSetDef winconsoleio_getsetlist[] = {
1110     {"closed", (getter)get_closed, NULL, "True if the file is closed"},
1111     {"closefd", (getter)get_closefd, NULL,
1112         "True if the file descriptor will be closed by close()."},
1113     {"mode", (getter)get_mode, NULL, "String giving the file mode"},
1114     {NULL},
1115 };
1116 
1117 static PyMemberDef winconsoleio_members[] = {
1118     {"_blksize", T_UINT, offsetof(winconsoleio, blksize), 0},
1119     {"_finalizing", T_BOOL, offsetof(winconsoleio, finalizing), 0},
1120     {NULL}
1121 };
1122 
1123 PyTypeObject PyWindowsConsoleIO_Type = {
1124     PyVarObject_HEAD_INIT(NULL, 0)
1125     "_io._WindowsConsoleIO",
1126     sizeof(winconsoleio),
1127     0,
1128     (destructor)winconsoleio_dealloc,           /* tp_dealloc */
1129     0,                                          /* tp_print */
1130     0,                                          /* tp_getattr */
1131     0,                                          /* tp_setattr */
1132     0,                                          /* tp_reserved */
1133     (reprfunc)winconsoleio_repr,                /* tp_repr */
1134     0,                                          /* tp_as_number */
1135     0,                                          /* tp_as_sequence */
1136     0,                                          /* tp_as_mapping */
1137     0,                                          /* tp_hash */
1138     0,                                          /* tp_call */
1139     0,                                          /* tp_str */
1140     PyObject_GenericGetAttr,                    /* tp_getattro */
1141     0,                                          /* tp_setattro */
1142     0,                                          /* tp_as_buffer */
1143     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE
1144         | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_FINALIZE,       /* tp_flags */
1145     _io__WindowsConsoleIO___init____doc__,      /* tp_doc */
1146     (traverseproc)winconsoleio_traverse,        /* tp_traverse */
1147     (inquiry)winconsoleio_clear,                /* tp_clear */
1148     0,                                          /* tp_richcompare */
1149     offsetof(winconsoleio, weakreflist),        /* tp_weaklistoffset */
1150     0,                                          /* tp_iter */
1151     0,                                          /* tp_iternext */
1152     winconsoleio_methods,                       /* tp_methods */
1153     winconsoleio_members,                       /* tp_members */
1154     winconsoleio_getsetlist,                    /* tp_getset */
1155     0,                                          /* tp_base */
1156     0,                                          /* tp_dict */
1157     0,                                          /* tp_descr_get */
1158     0,                                          /* tp_descr_set */
1159     offsetof(winconsoleio, dict),               /* tp_dictoffset */
1160     _io__WindowsConsoleIO___init__,             /* tp_init */
1161     PyType_GenericAlloc,                        /* tp_alloc */
1162     winconsoleio_new,                           /* tp_new */
1163     PyObject_GC_Del,                            /* tp_free */
1164     0,                                          /* tp_is_gc */
1165     0,                                          /* tp_bases */
1166     0,                                          /* tp_mro */
1167     0,                                          /* tp_cache */
1168     0,                                          /* tp_subclasses */
1169     0,                                          /* tp_weaklist */
1170     0,                                          /* tp_del */
1171     0,                                          /* tp_version_tag */
1172     0,                                          /* tp_finalize */
1173 };
1174 
1175 PyAPI_DATA(PyObject *) _PyWindowsConsoleIO_Type = (PyObject*)&PyWindowsConsoleIO_Type;
1176 
1177 #endif /* MS_WINDOWS */
1178