8000 Improved buffer handling · python/cpython@2d807c9 · GitHub
[go: up one dir, main page]
Skip to content

Commit 2d807c9

Browse files
zoobazooba
committed
Improved buffer handling
1 parent cafd311 commit 2d807c9

File tree

1 file changed

+63
-56
lines changed

1 file changed

+63
-56
lines changed

‎Modules/posixmodule.c

Lines changed: 63 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -4230,30 +4230,43 @@ static PyObject *
42304230
os_listdrives_impl(PyObject *module)
42314231
/*[clinic end generated code: output=aaece9dacdf682b5 input=1af9ccc9e583798e]*/
42324232
{
4233-
DWORD buflen = 64;
42344233
wchar_t defaultBuffer[64];
4234+
int success = 0;
4235+
DWORD buflen = Py_ARRAY_LENGTH(defaultBuffer);
42354236
LPWSTR buffer = defaultBuffer;
42364237
if (PySys_Audit("os.listdrives", NULL) < 0) {
42374238
return NULL;
42384239
}
4239-
Py_BEGIN_ALLOW_THREADS;
4240-
buflen = GetLogicalDriveStringsW(buflen, buffer);
4241-
if (buflen >= Py_ARRAY_LENGTH(defaultBuffer)) {
4242-
Py_BLOCK_THREADS;
4243-
buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * (buflen + 1));
4244-
Py_UNBLOCK_THREADS;
4245-
if (buffer) {
4246-
buflen = GetLogicalDriveStringsW(buflen, buffer);
4240+
while (1) {
4241+
DWORD requiredlen;
4242+
Py_BEGIN_ALLOW_THREADS;
4243+
requiredlen = GetLogicalDriveStringsW(buflen, buffer);
4244+
Py_END_ALLOW_THREADS;
4245+
if (!requiredlen) {
4246+
PyErr_SetFromWindowsErr(0);
4247+
break;
4248+
} else if (requiredlen >= buflen) {
4249+
buflen = requiredlen;
4250+
if (buffer != defaultBuffer) {
4251+
PyMem_Free((void *)buffer);
4252+
}
4253+
buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * (buflen + 1));
4254+
if (!buffer) {
4255+
PyErr_NoMemory();
4256+
break;
4257+
}
4258+
} else {
4259+
success = 1;
4260+
buflen = requiredlen;
4261+
break;
42474262
}
42484263
}
4249-
Py_END_ALLOW_THREADS;
4250-
42514264

42524265
PyObject *str = NULL;
42534266
PyObject *nullchar = NULL;
42544267
PyObject *r = NULL;
42554268

4256-
if (!buffer) {
4269+
if (!success || !buffer) {
42574270
goto exit;
42584271
}
42594272

@@ -4292,7 +4305,7 @@ os_listvolumes_impl(PyObject *module)
42924305
{
42934306
PyObject *result = PyList_New(0);
42944307
HANDLE find = INVALID_HANDLE_VALUE;
4295-
wchar_t buffer[256];
4308+
wchar_t buffer[MAX_PATH + 1];
42964309
if (!result) {
42974310
return NULL;
42984311
}
@@ -4303,7 +4316,7 @@ os_listvolumes_impl(PyObject *module)
43034316

43044317
int err = 0;
43054318
Py_BEGIN_ALLOW_THREADS;
4306-
find = FindFirstVolumeW(buffer, 256);
4319+
find = FindFirstVolumeW(buffer, Py_ARRAY_LENGTH(buffer));
43074320
if (find == INVALID_HANDLE_VALUE) {
43084321
err = GetLastError();
43094322
}
@@ -4319,7 +4332,7 @@ os_listvolumes_impl(PyObject *module)
43194332
Py_DECREF(s);
43204333

43214334
Py_BEGIN_ALLOW_THREADS;
4322-
if (!FindNextVolumeW(find, buffer, 256)) {
4335+
if (!FindNextVolumeW(find, buffer, Py_ARRAY_LENGTH(buffer))) {
43234336
err = GetLastError();
43244337
}
43254338
Py_END_ALLOW_THREADS;
@@ -4332,9 +4345,8 @@ os_listvolumes_impl(PyObject *module)
43324345
}
43334346
if (err && err != ERROR_NO_MORE_FILES) {
43344347
PyErr_SetFromWindowsErr(err);
4335-
if (result) {
4336-
Py_CLEAR(result);
4337-
}
4348+
Py_XDECREF(result);
4349+
result = NULL;
43384350
}
43394351
return result;
43404352
}
@@ -4355,60 +4367,55 @@ static PyObject *
43554367
os_listmounts_impl(PyObject *module, path_t *volume)
43564368
/*[clinic end generated code: output=06da49679de4512e input=a8a27178e3f67845]*/
43574369
{
4358-
DWORD buflen = 64;
4359-
wchar_t defaultBuffer[64];
4360-
LPWSTR buffer = defaultBuffer;
4361-
int err = 0;
4370+
wchar_t default_buffer[MAX_PATH + 1];
4371+
DWORD buflen = Py_ARRAY_LENGTH(default_buffer);
4372+
LPWSTR buffer = default_buffer;
4373+
PyObject *str = NULL;
4374+
PyObject *nullchar = NULL;
4375+
PyObject *result = NULL;
43624376
if (PySys_Audit("os.listmounts", "O", volume->object) < 0) {
43634377
return NULL;
43644378
}
4365-
Py_BEGIN_ALLOW_THREADS;
4366-
if (!GetVolumePathNamesForVolumeNameW(volume->wide, buffer, buflen, &buflen)) {
4367-
err = GetLastError();
4368-
if (err == ERROR_MORE_DATA) {
4369-
Py_BLOCK_THREADS;
4370-
buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * buflen);
4371-
Py_UNBLOCK_THREADS;
4372-
if (buffer) {
4373-
if (!GetVolumePathNamesForVolumeNameW(volume->wide, buffer,
4374-
buflen, &buflen)) {
4375-
err = GetLastError();
4376-
}
4377-
}
4379+
while (1) {
4380+
BOOL success;
4381+
Py_BEGIN_ALLOW_THREADS
4382+
success = GetVolumePathNamesForVolumeNameW(volume->wide, buffer,
4383+
buflen, &buflen);
4384+
Py_END_ALLOW_THREADS
4385+
if (success) {
4386+
break;
4387+
}
4388+
if (GetLastError() != ERROR_MORE_DATA) {
4389+
PyErr_SetFromWindowsErr(0);
4390+
goto exit;
4391+
}
4392+
if (buffer != default_buffer) {
4393+
PyMem_Free((void *)buffer);
4394+
}
4395+
buffer = (wchar_t*)PyMem_Malloc(sizeof(wchar_t) * buflen);
4396+
if (!buffer) {
4397+
PyErr_NoMemory();
4398+
goto exit;
43784399
}
4379-
}
4380-
Py_END_ALLOW_THREADS;
4381-
4382-
PyObject *str = NULL;
4383-
PyObject *nullchar = NULL;
4384-
PyObject *r = NULL;
4385-
4386-
if (err) {
4387-
PyErr_SetFromWindowsErr(err);
4388-
goto exit;
4389-
}
4390-
if (!buffer) {
4391-
goto exit;
43924400
}
43934401
if (buflen < 2) {
4394-
r = PyList_New(0);
4402+
result = PyList_New(0);
43954403
goto exit;
43964404
}
4397-
4398-
/* buflen includes two null terminators (one for the last string
4399-
and one for the array of strings */
4405+
// buflen includes two null terminators, one for the last string
4406+
// and one for the array of strings.
44004407
str = PyUnicode_FromWideChar(buffer, buflen - 2);
44014408
nullchar = PyUnicode_FromStringAndSize("\0", 1);
44024409
if (str && nullchar) {
4403-
r = PyUnicode_Split(str, nullchar, buflen);
4410+
result = PyUnicode_Split(str, nullchar, -1);
44044411
}
44054412
exit:
4406-
if (buffer && buffer != defaultBuffer) {
4407-
PyMem_Free((void *)buffer);
4413+
if (buffer != default_buffer) {
4414+
PyMem_Free(buffer);
44084415
}
44094416
Py_XDECREF(nullchar);
44104417
Py_XDECREF(str);
4411-
return r;
4418+
return result;
44124419
}
44134420

44144421

0 commit comments

Comments
 (0)
0