gh-87868: Sort and remove duplicates in getenvironment() (GH-102731)

(cherry picked from commit c31be58da8)

Co-authored-by: AN Long <aisk@users.noreply.github.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
Co-authored-by: Pieter Eendebak <pieter.eendebak@gmail.com>
Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>
This commit is contained in:
Miss Islington (bot) 2024-01-11 23:27:12 +01:00 committed by GitHub
parent a956e510f6
commit d15e1ac828
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 210 additions and 5 deletions

View file

@ -39,11 +39,12 @@
#include "structmember.h" // PyMemberDef
#define WINDOWS_LEAN_AND_MEAN
#include "windows.h"
#include <crtdbg.h>
#include "winreparse.h"
#include "pycore_runtime.h" // _Py_ID
#if defined(MS_WIN32) && !defined(MS_WIN64)
#define HANDLE_TO_PYNUM(handle) \
PyLong_FromUnsignedLong((unsigned long) handle)
@ -781,12 +782,162 @@ gethandle(PyObject* obj, const char* name)
return ret;
}
static PyObject *
sortenvironmentkey(PyObject *module, PyObject *item)
{
PyObject *result = NULL;
PyObject *locale = PyUnicode_FromWideChar(LOCALE_NAME_INVARIANT, -1);
if (locale) {
result = _winapi_LCMapStringEx_impl(NULL, locale, LCMAP_UPPERCASE, item);
Py_DECREF(locale);
}
return result;
}
static PyMethodDef sortenvironmentkey_def = {
"sortenvironmentkey", _PyCFunction_CAST(sortenvironmentkey), METH_O, "",
};
static int
sort_environment_keys(PyObject *keys)
{
PyObject *keyfunc = PyCFunction_New(&sortenvironmentkey_def, NULL);
if (keyfunc == NULL) {
return -1;
}
PyObject *kwnames = Py_BuildValue("(s)", "key");
if (kwnames == NULL) {
Py_DECREF(keyfunc);
return -1;
}
PyObject *args[] = { keys, keyfunc };
PyObject *ret = PyObject_VectorcallMethod(&_Py_ID(sort), args, 1, kwnames);
Py_DECREF(keyfunc);
Py_DECREF(kwnames);
if (ret == NULL) {
return -1;
}
Py_DECREF(ret);
return 0;
}
static int
compare_string_ordinal(PyObject *str1, PyObject *str2, int *result)
{
wchar_t *s1 = PyUnicode_AsWideCharString(str1, NULL);
if (s1 == NULL) {
return -1;
}
wchar_t *s2 = PyUnicode_AsWideCharString(str2, NULL);
if (s2 == NULL) {
PyMem_Free(s1);
return -1;
}
*result = CompareStringOrdinal(s1, -1, s2, -1, TRUE);
PyMem_Free(s1);
PyMem_Free(s2);
return 0;
}
static PyObject *
dedup_environment_keys(PyObject *keys)
{
PyObject *result = PyList_New(0);
if (result == NULL) {
return NULL;
}
// Iterate over the pre-ordered keys, check whether the current key is equal
// to the next key (ignoring case), if different, insert the current value
// into the result list. If they are equal, do nothing because we always
// want to keep the last inserted one.
for (Py_ssize_t i = 0; i < PyList_GET_SIZE(keys); i++) {
PyObject *key = PyList_GET_ITEM(keys, i);
// The last key will always be kept.
if (i + 1 == PyList_GET_SIZE(keys)) {
if (PyList_Append(result, key) < 0) {
Py_DECREF(result);
return NULL;
}
continue;
}
PyObject *next_key = PyList_GET_ITEM(keys, i + 1);
int compare_result;
if (compare_string_ordinal(key, next_key, &compare_result) < 0) {
Py_DECREF(result);
return NULL;
}
if (compare_result == CSTR_EQUAL) {
continue;
}
if (PyList_Append(result, key) < 0) {
Py_DECREF(result);
return NULL;
}
}
return result;
}
static PyObject *
normalize_environment(PyObject *environment)
{
PyObject *keys = PyMapping_Keys(environment);
if (keys == NULL) {
return NULL;
}
if (sort_environment_keys(keys) < 0) {
Py_DECREF(keys);
return NULL;
}
PyObject *normalized_keys = dedup_environment_keys(keys);
Py_DECREF(keys);
if (normalized_keys == NULL) {
return NULL;
}
PyObject *result = PyDict_New();
if (result == NULL) {
Py_DECREF(normalized_keys);
return NULL;
}
for (int i = 0; i < PyList_GET_SIZE(normalized_keys); i++) {
PyObject *key = PyList_GET_ITEM(normalized_keys, i);
PyObject *value = PyObject_GetItem(environment, key);
if (value == NULL) {
Py_DECREF(normalized_keys);
Py_DECREF(result);
return NULL;
}
int ret = PyObject_SetItem(result, key, value);
Py_DECREF(value);
if (ret < 0) {
Py_DECREF(normalized_keys);
Py_DECREF(result);
return NULL;
}
}
Py_DECREF(normalized_keys);
return result;
}
static wchar_t *
getenvironment(PyObject* environment)
{
Py_ssize_t i, envsize, totalsize;
wchar_t *buffer = NULL, *p, *end;
PyObject *keys, *values;
PyObject *normalized_environment = NULL;
PyObject *keys = NULL;
PyObject *values = NULL;
/* convert environment dictionary to windows environment string */
if (! PyMapping_Check(environment)) {
@ -795,11 +946,16 @@ getenvironment(PyObject* environment)
return NULL;
}
keys = PyMapping_Keys(environment);
if (!keys) {
normalized_environment = normalize_environment(environment);
if (normalize_environment == NULL) {
return NULL;
}
values = PyMapping_Values(environment);
keys = PyMapping_Keys(normalized_environment);
if (!keys) {
goto error;
}
values = PyMapping_Values(normalized_environment);
if (!values) {
goto error;
}
@ -891,6 +1047,7 @@ getenvironment(PyObject* environment)
cleanup:
error:
Py_XDECREF(normalized_environment);
Py_XDECREF(keys);
Py_XDECREF(values);
return buffer;