#include "Python.h"

#include "pycore_fileutils.h" // _Py_BEGIN_SUPPRESS_IPH

#include "pycore_object.h" // _PyObject_GC_UNTRACK()

#include "pycore_pyerrors.h" // _PyErr_ChainExceptions1()

#ifdef HAVE_WINDOWS_CONSOLE_IO

#ifdef HAVE_SYS_TYPES_H

#include <sys/types.h>

#ifdef HAVE_SYS_STAT_H

#include <sys/stat.h>

#include <stddef.h> /* For offsetof */

#ifndef WIN32_LEAN_AND_MEAN

#define WIN32_LEAN_AND_MEAN

#include <windows.h>

#include <fcntl.h>

#include "_iomodule.h"

#if BUFSIZ < (16*1024)

#define SMALLCHUNK (2*1024)

#elif (BUFSIZ >= (2 << 25))

#error "unreasonable BUFSIZ > 64 MiB defined"

#define SMALLCHUNK BUFSIZ

#define BUFMAX (32*1024*1024)

#define SMALLBUF 4

char _get_console_type(HANDLE handle) {

DWORD mode, peek_count;

if (handle == INVALID_HANDLE_VALUE)

return '\0';

if (!GetConsoleMode(handle, &mode))

return '\0';

/* Peek at the handle to see whether it is an input or output handle */

if (GetNumberOfConsoleInputEvents(handle, &peek_count))

return 'r';

return 'w';

}

char _PyIO_get_console_type(PyObject *path_or_fd) {

int fd = PyLong_AsLong(path_or_fd);

PyErr_Clear();

if (fd >= 0) {

HANDLE handle = _Py_get_osfhandle_noraise(fd);

if (handle == INVALID_HANDLE_VALUE)

return '\0';

return _get_console_type(handle);

}

PyObject *decoded;

wchar_t *decoded_wstr;

if (!PyUnicode_FSDecoder(path_or_fd, &decoded)) {

PyErr_Clear();

return '\0';

}

decoded_wstr = PyUnicode_AsWideCharString(decoded, NULL);

Py_CLEAR(decoded);

if (!decoded_wstr) {

PyErr_Clear();

return '\0';

}

char m = '\0';

if (!_wcsicmp(decoded_wstr, L"CONIN$")) {

m = 'r';

} else if (!_wcsicmp(decoded_wstr, L"CONOUT$")) {

m = 'w';

} else if (!_wcsicmp(decoded_wstr, L"CON")) {

m = 'x';

}

if (m) {

PyMem_Free(decoded_wstr);

return m;

}

DWORD length;

wchar_t name_buf[MAX_PATH], *pname_buf = name_buf;

length = GetFullPathNameW(decoded_wstr, MAX_PATH, pname_buf, NULL);

if (length > MAX_PATH) {

pname_buf = PyMem_New(wchar_t, length);

if (pname_buf)

length = GetFullPathNameW(decoded_wstr, length, pname_buf, NULL);

else

length = 0;

}

PyMem_Free(decoded_wstr);

if (length) {

wchar_t *name = pname_buf;

if (length >= 4 && name[3] == L'\\' &&

(name[2] == L'.' || name[2] == L'?') &&

name[1] == L'\\' && name[0] == L'\\') {

name += 4;

}

if (!_wcsicmp(name, L"CONIN$")) {

m = 'r';

} else if (!_wcsicmp(name, L"CONOUT$")) {

m = 'w';

} else if (!_wcsicmp(name, L"CON")) {

m = 'x';

}

}

if (pname_buf != name_buf)

PyMem_Free(pname_buf);

return m;

}

static DWORD

_find_last_utf8_boundary(const char *buf, DWORD len)

{

/* This function never returns 0, returns the original len instead */

DWORD count = 1;

if (len == 0 || (buf[len - 1] & 0x80) == 0) {

return len;

}

for (;; count++) {

if (count > 3 || count >= len) {

return len;

}

if ((buf[len - count] & 0xc0) != 0x80) {

return len - count;

}

}

}

typedef struct {

PyObject_HEAD

int fd;

unsigned int created : 1;

unsigned int readable : 1;

unsigned int writable : 1;

unsigned int closefd : 1;

char finalizing;

unsigned int blksize;

PyObject *weakreflist;

PyObject *dict;

char buf[SMALLBUF];

wchar_t wbuf;

} winconsoleio;

_PyWindowsConsoleIO_closed(PyObject *self)

{

return ((winconsoleio *)self)->fd == -1;

}

static int

internal_close(winconsoleio *self)

{

if (self->fd != -1) {

if (self->closefd) {

_Py_BEGIN_SUPPRESS_IPH

close(self->fd);

_Py_END_SUPPRESS_IPH

}

self->fd = -1;

}

return 0;

}

static PyObject *

_io__WindowsConsoleIO_close_impl(winconsoleio *self, PyTypeObject *cls)

{

PyObject *res;

PyObject *exc;

int rc;

_PyIO_State *state = get_io_state_by_cls(cls);

res = PyObject_CallMethodOneArg((PyObject*)state->PyRawIOBase_Type,

&_Py_ID(close), (PyObject*)self);

if (!self->closefd) {

self->fd = -1;

return res;

}

if (res == NULL) {

exc = PyErr_GetRaisedException();

}

rc = internal_close(self);

if (res == NULL) {

_PyErr_ChainExceptions1(exc);

}

if (rc < 0) {

Py_CLEAR(res);

}

return res;

}

static PyObject *

winconsoleio_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

winconsoleio *self;

assert(type != NULL && type->tp_alloc != NULL);

self = (winconsoleio *) type->tp_alloc(type, 0);

if (self != NULL) {

self->fd = -1;

self->created = 0;

self->readable = 0;

self->writable = 0;

self->closefd = 0;

self->blksize = 0;

self->weakreflist = NULL;

}

return (PyObject *) self;

}

static int

_io__WindowsConsoleIO___init___impl(winconsoleio *self, PyObject *nameobj,

const char *mode, int closefd,

PyObject *opener)

{

const char *s;

wchar_t *name = NULL;

char console_type = '\0';

int ret = 0;

int rwa = 0;

int fd = -1;

int fd_is_own = 0;

HANDLE handle = NULL;

#ifndef NDEBUG

_PyIO_State *state = find_io_state_by_def(Py_TYPE(self));

assert(PyObject_TypeCheck(self, state->PyWindowsConsoleIO_Type));

if (self->fd >= 0) {

if (self->closefd) {

/* Have to close the existing file first. */

if (internal_close(self) < 0)

return -1;

}

else

self->fd = -1;

}

fd = PyLong_AsInt(nameobj);

if (fd < 0) {

if (!PyErr_Occurred()) {

PyErr_SetString(PyExc_ValueError,

"negative file descriptor");

return -1;

}

PyErr_Clear();

}

self->fd = fd;

if (fd < 0) {

PyObject *decodedname;

int d = PyUnicode_FSDecoder(nameobj, (void*)&decodedname);

if (!d)

return -1;

name = PyUnicode_AsWideCharString(decodedname, NULL);

console_type = _PyIO_get_console_type(decodedname);

Py_CLEAR(decodedname);

if (name == NULL)

return -1;

}

s = mode;

while (*s) {

switch (*s++) {

case '+':

case 'a':

case 'b':

case 'x':

break;

case 'r':

if (rwa)

goto bad_mode;

rwa = 1;

self->readable = 1;

if (console_type == 'x')

console_type = 'r';

break;

case 'w':

if (rwa)

goto bad_mode;

rwa = 1;

self->writable = 1;

if (console_type == 'x')

console_type = 'w';

break;

default:

PyErr_Format(PyExc_ValueError,

"invalid mode: %.200s", mode);

goto error;

}

}

if (!rwa)

goto bad_mode;

if (fd >= 0) {

handle = _Py_get_osfhandle_noraise(fd);

self->closefd = 0;

} else {

DWORD access = GENERIC_READ;

self->closefd = 1;

if (!closefd) {

PyErr_SetString(PyExc_ValueError,

"Cannot use closefd=False with file name");

goto error;

}

if (self->writable)

access = GENERIC_WRITE;

Py_BEGIN_ALLOW_THREADS

/* Attempt to open for read/write initially, then fall back

handle = CreateFileW(name, GENERIC_READ | GENERIC_WRITE,

FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);

if (handle == INVALID_HANDLE_VALUE)

handle = CreateFileW(name, access,

FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);

Py_END_ALLOW_THREADS

if (handle == INVALID_HANDLE_VALUE) {

PyErr_SetExcFromWindowsErrWithFilenameObject(PyExc_OSError, GetLastError(), nameobj);

goto error;

}

if (self->writable)

self->fd = _Py_open_osfhandle_noraise(handle, _O_WRONLY | _O_BINARY);

else

self->fd = _Py_open_osfhandle_noraise(handle, _O_RDONLY | _O_BINARY);

if (self->fd < 0) {

PyErr_SetFromErrnoWithFilenameObject(PyExc_OSError, nameobj);

CloseHandle(handle);

goto error;

}

}

if (console_type == '\0')

console_type = _get_console_type(handle);

if (self->writable && console_type != 'w') {

PyErr_SetString(PyExc_ValueError,

"Cannot open console input buffer for writing");

goto error;

}

if (self->readable && console_type != 'r') {

PyErr_SetString(PyExc_ValueError,

"Cannot open console output buffer for reading");

goto error;

}

self->blksize = DEFAULT_BUFFER_SIZE;

memset(self->buf, 0, 4);

if (PyObject_SetAttr((PyObject *)self, &_Py_ID(name), nameobj) < 0)

goto error;

goto done;

bad_mode:

PyErr_SetString(PyExc_ValueError,

"Must have exactly one of read or write mode");

error:

ret = -1;

internal_close(self);

done:

if (name)

PyMem_Free(name);

return ret;

}

static int

winconsoleio_traverse(winconsoleio *self, visitproc visit, void *arg)

{

Py_VISIT(Py_TYPE(self));

Py_VISIT(self->dict);

return 0;

}

static int

winconsoleio_clear(winconsoleio *self)

{

Py_CLEAR(self->dict);

return 0;

}

static void

winconsoleio_dealloc(winconsoleio *self)

{

PyTypeObject *tp = Py_TYPE(self);

self->finalizing = 1;

if (_PyIOBase_finalize((PyObject *) self) < 0)

return;

_PyObject_GC_UNTRACK(self);

if (self->weakreflist != NULL)

PyObject_ClearWeakRefs((PyObject *) self);

Py_CLEAR(self->dict);

tp->tp_free((PyObject *)self);

Py_DECREF(tp);

}

static PyObject *

err_closed(void)

{

PyErr_SetString(PyExc_ValueError, "I/O operation on closed file");

return NULL;

}

static PyObject *

err_mode(_PyIO_State *state, const char *action)

{

return PyErr_Format(state->unsupported_operation,

"Console buffer does not support %s", action);

}

static PyObject *

_io__WindowsConsoleIO_fileno_impl(winconsoleio *self)

{

if (self->fd < 0)

return err_closed();

return PyLong_FromLong(self->fd);

}

static PyObject *

_io__WindowsConsoleIO_readable_impl(winconsoleio *self)

{

if (self->fd == -1)

return err_closed();

return PyBool_FromLong((long) self->readable);

}

static PyObject *

_io__WindowsConsoleIO_writable_impl(winconsoleio *self)

{

if (self->fd == -1)

return err_closed();

return PyBool_FromLong((long) self->writable);

}

static DWORD

_buflen(winconsoleio *self)

{

for (DWORD i = 0; i < SMALLBUF; ++i) {

if (!self->buf[i])

return i;

}

return SMALLBUF;

}

static DWORD

_copyfrombuf(winconsoleio *self, char *buf, DWORD len)

{

DWORD n = 0;

while (self->buf[0] && len--) {

buf[n++] = self->buf[0];

for (int i = 1; i < SMALLBUF; ++i)

self->buf[i - 1] = self->buf[i];

self->buf[SMALLBUF - 1] = 0;

}

return n;

}

static wchar_t *

read_console_w(HANDLE handle, DWORD maxlen, DWORD *readlen) {

int err = 0, sig = 0;

wchar_t *buf = (wchar_t*)PyMem_Malloc(maxlen * sizeof(wchar_t));

if (!buf)

goto error;

*readlen = 0;

//DebugBreak();

Py_BEGIN_ALLOW_THREADS

DWORD off = 0;

while (off < maxlen) {

DWORD n = (DWORD)-1;

DWORD len = min(maxlen - off, BUFSIZ);

SetLastError(0);

BOOL res = ReadConsoleW(handle, &buf[off], len, &n, NULL);

if (!res) {

err = GetLastError();

break;

}

if (n == (DWORD)-1 && (err = GetLastError()) == ERROR_OPERATION_ABORTED) {

break;

}

if (n == 0) {

err = GetLastError();

if (err != ERROR_OPERATION_ABORTED)

break;

err = 0;

HANDLE hInterruptEvent = _PyOS_SigintEvent();

if (WaitForSingleObjectEx(hInterruptEvent, 100, FALSE)

== WAIT_OBJECT_0) {

ResetEvent(hInterruptEvent);

Py_BLOCK_THREADS

sig = PyErr_CheckSignals();

Py_UNBLOCK_THREADS

if (sig < 0)

break;

}

}

*readlen += n;

/* If we didn't read a full buffer that time, don't try

if (n < len)

break;

/* If the buffer ended with a newline, break out */

if (buf[*readlen - 1] == '\n')

break;

/* If the buffer ends with a high surrogate, expand the

WORD char_type;

if (off + BUFSIZ >= maxlen &&

GetStringTypeW(CT_CTYPE3, &buf[*readlen - 1], 1, &char_type) &&

char_type == C3_HIGHSURROGATE) {

wchar_t *newbuf;

maxlen += 1;

Py_BLOCK_THREADS

newbuf = (wchar_t*)PyMem_Realloc(buf, maxlen * sizeof(wchar_t));

Py_UNBLOCK_THREADS

if (!newbuf) {

sig = -1;

break;

}

buf = newbuf;

/* Only advance by n and not BUFSIZ in this case */

off += n;

continue;

}

off += BUFSIZ;

}

Py_END_ALLOW_THREADS

if (sig)

goto error;

if (err) {

PyErr_SetFromWindowsErr(err);

goto error;

}

if (*readlen > 0 && buf[0] == L'\x1a') {

PyMem_Free(buf);

buf = (wchar_t *)PyMem_Malloc(sizeof(wchar_t));

if (!buf)

goto error;

buf[0] = L'\0';

*readlen = 0;

}

return buf;

error:

if (buf)

PyMem_Free(buf);

return NULL;

}

static Py_ssize_t

readinto(_PyIO_State *state, winconsoleio *self, char *buf, Py_ssize_t len)

{

if (self->fd == -1) {

err_closed();

return -1;

}

if (!self->readable) {

err_mode(state, "reading");

return -1;

}

if (len == 0)

return 0;

if (len > BUFMAX) {

PyErr_Format(PyExc_ValueError, "cannot read more than %d bytes", BUFMAX);

return -1;

}

HANDLE handle = _Py_get_osfhandle(self->fd);

if (handle == INVALID_HANDLE_VALUE)

return -1;

/* Each character may take up to 4 bytes in the final buffer.

DWORD wlen = (DWORD)(len / 4);

if (wlen == 0) {

wlen = 1;

}

DWORD read_len = _copyfrombuf(self, buf, (DWORD)len);

if (read_len) {

buf = &buf[read_len];

len -= read_len;

wlen -= 1;

}

if (len == read_len || wlen == 0)

return read_len;

DWORD n;

wchar_t *wbuf = read_console_w(handle, wlen, &n);

if (wbuf == NULL)

return -1;

if (n == 0) {

PyMem_Free(wbuf);

return read_len;

}

int err = 0;

DWORD u8n = 0;

Py_BEGIN_ALLOW_THREADS

if (len < 4) {

if (WideCharToMultiByte(CP_UTF8, 0, wbuf, n,

self->buf, sizeof(self->buf) / sizeof(self->buf[0]),

NULL, NULL))

u8n = _copyfrombuf(self, buf, (DWORD)len);

} else {

u8n = WideCharToMultiByte(CP_UTF8, 0, wbuf, n,

buf, (DWORD)len, NULL, NULL);

}

if (u8n) {

read_len += u8n;

u8n = 0;

} else {

err = GetLastError();

if (err == ERROR_INSUFFICIENT_BUFFER) {

/* Calculate the needed buffer for a more useful error, as this

u8n = WideCharToMultiByte(CP_UTF8, 0, wbuf, n,

NULL, 0, NULL, NULL);

}

}

Py_END_ALLOW_THREADS

PyMem_Free(wbuf);

if (u8n) {

PyErr_Format(PyExc_SystemError,

"Buffer had room for %zd bytes but %u bytes required",

len, u8n);

return -1;

}

if (err) {

PyErr_SetFromWindowsErr(err);

return -1;

}

return read_len;

}

static PyObject *

_io__WindowsConsoleIO_readinto_impl(winconsoleio *self, PyTypeObject *cls,

Py_buffer *buffer)

{

_PyIO_State *state = get_io_state_by_cls(cls);

Py_ssize_t len = readinto(state, self, buffer->buf, buffer->len);

if (len < 0)

return NULL;

return PyLong_FromSsize_t(len);

}

static DWORD

new_buffersize(winconsoleio *self, DWORD currentsize)

{

DWORD addend;

/* Expand the buffer by an amount proportional to the current size,

if (currentsize > 65536)

addend = currentsize >> 3;

else

addend = 256 + currentsize;

if (addend < SMALLCHUNK)

/* Avoid tiny read() calls. */

addend = SMALLCHUNK;

return addend + currentsize;

}

static PyObject *

_io__WindowsConsoleIO_readall_impl(winconsoleio *self)

{

wchar_t *buf;

DWORD bufsize, n, len = 0;

PyObject *bytes;

DWORD bytes_size, rn;

HANDLE handle;

if (self->fd == -1)

return err_closed();

handle = _Py_get_osfhandle(self->fd);

if (handle == INVALID_HANDLE_VALUE)

return NULL;

bufsize = BUFSIZ;

buf = (wchar_t*)PyMem_Malloc((bufsize + 1) * sizeof(wchar_t));

if (buf == NULL)

return NULL;

while (1) {

wchar_t *subbuf;

if (len >= (Py_ssize_t)bufsize) {

DWORD newsize = new_buffersize(self, len);

if (newsize > BUFMAX)

break;

if (newsize < bufsize) {

PyErr_SetString(PyExc_OverflowError,

"unbounded read returned more bytes "

"than a Python bytes object can hold");

PyMem_Free(buf);

return NULL;

}

bufsize = newsize;

wchar_t *tmp = PyMem_Realloc(buf,

(bufsize + 1) * sizeof(wchar_t));

if (tmp == NULL) {

PyMem_Free(buf);

return NULL;

}

buf = tmp;

}

subbuf = read_console_w(handle, bufsize - len, &n);

if (subbuf == NULL) {

PyMem_Free(buf);

return NULL;

}

if (n > 0)

wcsncpy_s(&buf[len], bufsize - len + 1, subbuf, n);

PyMem_Free(subbuf);

/* when the read is empty we break */

if (n == 0)

break;

len += n;

}

if (len == 0 && _buflen(self) == 0) {

/* when the result starts with ^Z we return an empty buffer */

PyMem_Free(buf);

return PyBytes_FromStringAndSize(NULL, 0);

}

if (len) {

Py_BEGIN_ALLOW_THREADS

bytes_size = WideCharToMultiByte(CP_UTF8, 0, buf, len,

NULL, 0, NULL, NULL);

Py_END_ALLOW_THREADS

if (!bytes_size) {

DWORD err = GetLastError();

PyMem_Free(buf);

return PyErr_SetFromWindowsErr(err);

}

} else {

bytes_size = 0;

}

bytes_size += _buflen(self);

bytes = PyBytes_FromStringAndSize(NULL, bytes_size);

rn = _copyfrombuf(self, PyBytes_AS_STRING(bytes), bytes_size);

if (len) {

Py_BEGIN_ALLOW_THREADS

bytes_size = WideCharToMultiByte(CP_UTF8, 0, buf, len,

&PyBytes_AS_STRING(bytes)[rn], bytes_size - rn, NULL, NULL);

Py_END_ALLOW_THREADS

if (!bytes_size) {

DWORD err = GetLastError();

PyMem_Free(buf);

Py_CLEAR(bytes);

return PyErr_SetFromWindowsErr(err);

}

/* add back the number of preserved bytes */

bytes_size += rn;

}

PyMem_Free(buf);

if (bytes_size < (size_t)PyBytes_GET_SIZE(bytes)) {

if (_PyBytes_Resize(&bytes, n * sizeof(wchar_t)) < 0) {

Py_CLEAR(bytes);

return NULL;

}

}

return bytes;

}

static PyObject *

_io__WindowsConsoleIO_read_impl(winconsoleio *self, PyTypeObject *cls,

Py_ssize_t size)

{

PyObject *bytes;

Py_ssize_t bytes_size;

if (self->fd == -1)

return err_closed();

if (!self->readable) {

_PyIO_State *state = get_io_state_by_cls(cls);

return err_mode(state, "reading");

}

if (size < 0)

return _io__WindowsConsoleIO_readall_impl(self);

if (size > BUFMAX) {

PyErr_Format(PyExc_ValueError, "cannot read more than %d bytes", BUFMAX);

return NULL;

}

bytes = PyBytes_FromStringAndSize(NULL, size);

if (bytes == NULL)

return NULL;

_PyIO_State *state = get_io_state_by_cls(cls);

bytes_size = readinto(state, self, PyBytes_AS_STRING(bytes),

PyBytes_GET_SIZE(bytes));

if (bytes_size < 0) {

Py_CLEAR(bytes);

return NULL;

}

if (bytes_size < PyBytes_GET_SIZE(bytes)) {

if (_PyBytes_Resize(&bytes, bytes_size) < 0) {

Py_CLEAR(bytes);

return NULL;

}

}

return bytes;

}

static PyObject *

_io__WindowsConsoleIO_write_impl(winconsoleio *self, PyTypeObject *cls,

Py_buffer *b)

{

BOOL res = TRUE;

wchar_t *wbuf;

DWORD len, wlen, n = 0;

HANDLE handle;

if (self->fd == -1)

return err_closed();

if (!self->writable) {