Skip to content
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions Lib/test/test_zoneinfo/test_zoneinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import struct
import tempfile
import unittest
import zoneinfo
from datetime import date, datetime, time, timedelta, timezone
from functools import cached_property

Expand Down Expand Up @@ -1551,6 +1552,46 @@ def __eq__(self, other):
except CustomError:
pass

def test_weak_cache_descriptor_use_after_free(self):
from zoneinfo import ZoneInfo

available_zones = sorted(zoneinfo.available_timezones())
if "UTC" not in available_zones:
raise unittest.SkipTest("No time zone data available")

class Cache:
def __init__(self):
self.data = {}

def get(self, key, default=None):
return self.data.get(key, default)

def setdefault(self, key, default):
return self.data.setdefault(key, default)

def clear(self, *args, **kwargs):
self.data.clear()

class BombDescriptor:
def __get__(self, obj, owner):
return Cache()

class EvilZoneInfo(ZoneInfo):
pass

EvilZoneInfo._weak_cache = BombDescriptor()

zone1 = EvilZoneInfo("UTC")

self.assertIsNotNone(zone1)
self.assertEqual(str(zone1), "UTC")

EvilZoneInfo.clear_cache()

zone2 = EvilZoneInfo("UTC")
self.assertIsNotNone(zone2)
self.assertEqual(str(zone2), "UTC")


class CZoneInfoCacheTest(ZoneInfoCacheTest):
module = c_zoneinfo
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fix zoneinfo use-after-free with descriptor _weak_cache. Return new
reference from get_weak_cache() instead of borrowed reference to prevent
premature object free.
16 changes: 9 additions & 7 deletions Modules/_zoneinfo.c
Original file line number Diff line number Diff line change
Expand Up @@ -292,16 +292,11 @@ static PyObject *
get_weak_cache(zoneinfo_state *state, PyTypeObject *type)
{
if (type == state->ZoneInfoType) {
Py_INCREF(state->ZONEINFO_WEAK_CACHE);
return state->ZONEINFO_WEAK_CACHE;
}
else {
PyObject *cache =
PyObject_GetAttrString((PyObject *)type, "_weak_cache");
// We are assuming that the type lives at least as long as the function
// that calls get_weak_cache, and that it holds a reference to the
// cache, so we'll return a "borrowed reference".
Py_XDECREF(cache);
return cache;
return PyObject_GetAttrString((PyObject *)type, "_weak_cache");
}
}

Expand All @@ -328,26 +323,30 @@ zoneinfo_ZoneInfo_impl(PyTypeObject *type, PyObject *key)
PyObject *weak_cache = get_weak_cache(state, type);
instance = PyObject_CallMethod(weak_cache, "get", "O", key, Py_None);
if (instance == NULL) {
Py_DECREF(weak_cache);
return NULL;
}

if (instance == Py_None) {
Py_DECREF(instance);
PyObject *tmp = zoneinfo_new_instance(state, type, key);
if (tmp == NULL) {
Py_DECREF(weak_cache);
return NULL;
}

instance =
PyObject_CallMethod(weak_cache, "setdefault", "OO", key, tmp);
Py_DECREF(tmp);
if (instance == NULL) {
Py_DECREF(weak_cache);
return NULL;
}
((PyZoneInfo_ZoneInfo *)instance)->source = SOURCE_CACHE;
}

update_strong_cache(state, type, key, instance);
Py_DECREF(weak_cache);
return instance;
}

Expand Down Expand Up @@ -510,12 +509,14 @@ zoneinfo_ZoneInfo_clear_cache_impl(PyTypeObject *type, PyTypeObject *cls,
PyObject *item = NULL;
PyObject *pop = PyUnicode_FromString("pop");
if (pop == NULL) {
Py_DECREF(weak_cache);
return NULL;
}

PyObject *iter = PyObject_GetIter(only_keys);
if (iter == NULL) {
Py_DECREF(pop);
Py_DECREF(weak_cache);
return NULL;
}

Expand All @@ -540,6 +541,7 @@ zoneinfo_ZoneInfo_clear_cache_impl(PyTypeObject *type, PyTypeObject *cls,
Py_DECREF(pop);
}

Py_DECREF(weak_cache);
if (PyErr_Occurred()) {
return NULL;
}
Expand Down
Loading