1/* PickleBuffer object implementation */
2
3#define PY_SSIZE_T_CLEAN
4#include "Python.h"
5#include <stddef.h>
6
7typedef struct {
8    PyObject_HEAD
9    /* The view exported by the original object */
10    Py_buffer view;
11    PyObject *weakreflist;
12} PyPickleBufferObject;
13
14/* C API */
15
16PyObject *
17PyPickleBuffer_FromObject(PyObject *base)
18{
19    PyTypeObject *type = &PyPickleBuffer_Type;
20    PyPickleBufferObject *self;
21
22    self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
23    if (self == NULL) {
24        return NULL;
25    }
26    self->view.obj = NULL;
27    self->weakreflist = NULL;
28    if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
29        Py_DECREF(self);
30        return NULL;
31    }
32    return (PyObject *) self;
33}
34
35const Py_buffer *
36PyPickleBuffer_GetBuffer(PyObject *obj)
37{
38    PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
39
40    if (!PyPickleBuffer_Check(obj)) {
41        PyErr_Format(PyExc_TypeError,
42                     "expected PickleBuffer, %.200s found",
43                     Py_TYPE(obj)->tp_name);
44        return NULL;
45    }
46    if (self->view.obj == NULL) {
47        PyErr_SetString(PyExc_ValueError,
48                        "operation forbidden on released PickleBuffer object");
49        return NULL;
50    }
51    return &self->view;
52}
53
54int
55PyPickleBuffer_Release(PyObject *obj)
56{
57    PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
58
59    if (!PyPickleBuffer_Check(obj)) {
60        PyErr_Format(PyExc_TypeError,
61                     "expected PickleBuffer, %.200s found",
62                     Py_TYPE(obj)->tp_name);
63        return -1;
64    }
65    PyBuffer_Release(&self->view);
66    return 0;
67}
68
69static PyObject *
70picklebuf_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
71{
72    PyPickleBufferObject *self;
73    PyObject *base;
74    char *keywords[] = {"", NULL};
75
76    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:PickleBuffer",
77                                     keywords, &base)) {
78        return NULL;
79    }
80
81    self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
82    if (self == NULL) {
83        return NULL;
84    }
85    self->view.obj = NULL;
86    self->weakreflist = NULL;
87    if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
88        Py_DECREF(self);
89        return NULL;
90    }
91    return (PyObject *) self;
92}
93
94static int
95picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
96{
97    Py_VISIT(self->view.obj);
98    return 0;
99}
100
101static int
102picklebuf_clear(PyPickleBufferObject *self)
103{
104    PyBuffer_Release(&self->view);
105    return 0;
106}
107
108static void
109picklebuf_dealloc(PyPickleBufferObject *self)
110{
111    PyObject_GC_UnTrack(self);
112    if (self->weakreflist != NULL)
113        PyObject_ClearWeakRefs((PyObject *) self);
114    PyBuffer_Release(&self->view);
115    Py_TYPE(self)->tp_free((PyObject *) self);
116}
117
118/* Buffer API */
119
120static int
121picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
122{
123    if (self->view.obj == NULL) {
124        PyErr_SetString(PyExc_ValueError,
125                        "operation forbidden on released PickleBuffer object");
126        return -1;
127    }
128    return PyObject_GetBuffer(self->view.obj, view, flags);
129}
130
131static void
132picklebuf_releasebuf(PyPickleBufferObject *self, Py_buffer *view)
133{
134    /* Since our bf_getbuffer redirects to the original object, this
135     * implementation is never called.  It only exists to signal that
136     * buffers exported by PickleBuffer have non-trivial releasing
137     * behaviour (see check in Python/getargs.c).
138     */
139}
140
141static PyBufferProcs picklebuf_as_buffer = {
142    .bf_getbuffer = (getbufferproc) picklebuf_getbuf,
143    .bf_releasebuffer = (releasebufferproc) picklebuf_releasebuf,
144};
145
146/* Methods */
147
148static PyObject *
149picklebuf_raw(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
150{
151    if (self->view.obj == NULL) {
152        PyErr_SetString(PyExc_ValueError,
153                        "operation forbidden on released PickleBuffer object");
154        return NULL;
155    }
156    if (self->view.suboffsets != NULL
157        || !PyBuffer_IsContiguous(&self->view, 'A')) {
158        PyErr_SetString(PyExc_BufferError,
159                        "cannot extract raw buffer from non-contiguous buffer");
160        return NULL;
161    }
162    PyObject *m = PyMemoryView_FromObject((PyObject *) self);
163    if (m == NULL) {
164        return NULL;
165    }
166    PyMemoryViewObject *mv = (PyMemoryViewObject *) m;
167    assert(mv->view.suboffsets == NULL);
168    /* Mutate memoryview instance to make it a "raw" memoryview */
169    mv->view.format = "B";
170    mv->view.ndim = 1;
171    mv->view.itemsize = 1;
172    /* shape = (length,) */
173    mv->view.shape = &mv->view.len;
174    /* strides = (1,) */
175    mv->view.strides = &mv->view.itemsize;
176    /* Fix memoryview state flags */
177    /* XXX Expose memoryobject.c's init_flags() instead? */
178    mv->flags = _Py_MEMORYVIEW_C | _Py_MEMORYVIEW_FORTRAN;
179    return m;
180}
181
182PyDoc_STRVAR(picklebuf_raw_doc,
183"raw($self, /)\n--\n\
184\n\
185Return a memoryview of the raw memory underlying this buffer.\n\
186Will raise BufferError is the buffer isn't contiguous.");
187
188static PyObject *
189picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
190{
191    PyBuffer_Release(&self->view);
192    Py_RETURN_NONE;
193}
194
195PyDoc_STRVAR(picklebuf_release_doc,
196"release($self, /)\n--\n\
197\n\
198Release the underlying buffer exposed by the PickleBuffer object.");
199
200static PyMethodDef picklebuf_methods[] = {
201    {"raw",     (PyCFunction) picklebuf_raw,     METH_NOARGS, picklebuf_raw_doc},
202    {"release", (PyCFunction) picklebuf_release, METH_NOARGS, picklebuf_release_doc},
203    {NULL,      NULL}
204};
205
206PyTypeObject PyPickleBuffer_Type = {
207    PyVarObject_HEAD_INIT(NULL, 0)
208    .tp_name = "pickle.PickleBuffer",
209    .tp_doc = PyDoc_STR("Wrapper for potentially out-of-band buffers"),
210    .tp_basicsize = sizeof(PyPickleBufferObject),
211    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
212    .tp_new = picklebuf_new,
213    .tp_dealloc = (destructor) picklebuf_dealloc,
214    .tp_traverse = (traverseproc) picklebuf_traverse,
215    .tp_clear = (inquiry) picklebuf_clear,
216    .tp_weaklistoffset = offsetof(PyPickleBufferObject, weakreflist),
217    .tp_as_buffer = &picklebuf_as_buffer,
218    .tp_methods = picklebuf_methods,
219};
220