xref: /third_party/python/Modules/_dbmmodule.c (revision 7db96d56) 
1
2/* DBM module using dictionary interface */
3
4
5#define PY_SSIZE_T_CLEAN
6#include "Python.h"
7
8#include <sys/types.h>
9#include <sys/stat.h>
10#include <fcntl.h>
11
12/* Some Linux systems install gdbm/ndbm.h, but not ndbm.h.  This supports
13 * whichever configure was able to locate.
14 */
15#if defined(USE_GDBM_COMPAT)
16  #ifdef HAVE_GDBM_NDBM_H
17    #include <gdbm/ndbm.h>
18  #elif HAVE_GDBM_DASH_NDBM_H
19    #include <gdbm-ndbm.h>
20  #else
21    #error "No gdbm/ndbm.h or gdbm-ndbm.h available"
22  #endif
23  static const char which_dbm[] = "GNU gdbm";
24#elif defined(USE_NDBM)
25  #include <ndbm.h>
26  static const char which_dbm[] = "GNU gdbm";
27#elif defined(USE_BERKDB)
28  #ifndef DB_DBM_HSEARCH
29    #define DB_DBM_HSEARCH 1
30  #endif
31  #include <db.h>
32  static const char which_dbm[] = "Berkeley DB";
33#else
34  #error "No ndbm.h available!"
35#endif
36
37typedef struct {
38    PyTypeObject *dbm_type;
39    PyObject *dbm_error;
40} _dbm_state;
41
42static inline _dbm_state*
43get_dbm_state(PyObject *module)
44{
45    void *state = PyModule_GetState(module);
46    assert(state != NULL);
47    return (_dbm_state *)state;
48}
49
50/*[clinic input]
51module _dbm
52class _dbm.dbm "dbmobject *" "&Dbmtype"
53[clinic start generated code]*/
54/*[clinic end generated code: output=da39a3ee5e6b4b0d input=9b1aa8756d16150e]*/
55
56typedef struct {
57    PyObject_HEAD
58    int flags;
59    int di_size;        /* -1 means recompute */
60    DBM *di_dbm;
61} dbmobject;
62
63#include "clinic/_dbmmodule.c.h"
64
65#define check_dbmobject_open(v, err)                                \
66    if ((v)->di_dbm == NULL) {                                      \
67        PyErr_SetString(err, "DBM object has already been closed"); \
68        return NULL;                                                \
69    }
70
71static PyObject *
72newdbmobject(_dbm_state *state, const char *file, int flags, int mode)
73{
74    dbmobject *dp = PyObject_GC_New(dbmobject, state->dbm_type);
75    if (dp == NULL) {
76        return NULL;
77    }
78    dp->di_size = -1;
79    dp->flags = flags;
80    PyObject_GC_Track(dp);
81
82    /* See issue #19296 */
83    if ( (dp->di_dbm = dbm_open((char *)file, flags, mode)) == 0 ) {
84        PyErr_SetFromErrnoWithFilename(state->dbm_error, file);
85        Py_DECREF(dp);
86        return NULL;
87    }
88    return (PyObject *)dp;
89}
90
91/* Methods */
92static int
93dbm_traverse(dbmobject *dp, visitproc visit, void *arg)
94{
95    Py_VISIT(Py_TYPE(dp));
96    return 0;
97}
98
99static void
100dbm_dealloc(dbmobject *dp)
101{
102    PyObject_GC_UnTrack(dp);
103    if (dp->di_dbm) {
104        dbm_close(dp->di_dbm);
105    }
106    PyTypeObject *tp = Py_TYPE(dp);
107    tp->tp_free(dp);
108    Py_DECREF(tp);
109}
110
111static Py_ssize_t
112dbm_length(dbmobject *dp)
113{
114    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
115    assert(state != NULL);
116    if (dp->di_dbm == NULL) {
117             PyErr_SetString(state->dbm_error, "DBM object has already been closed");
118             return -1;
119    }
120    if ( dp->di_size < 0 ) {
121        datum key;
122        int size;
123
124        size = 0;
125        for ( key=dbm_firstkey(dp->di_dbm); key.dptr;
126              key = dbm_nextkey(dp->di_dbm))
127            size++;
128        dp->di_size = size;
129    }
130    return dp->di_size;
131}
132
133static PyObject *
134dbm_subscript(dbmobject *dp, PyObject *key)
135{
136    datum drec, krec;
137    Py_ssize_t tmp_size;
138    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
139    assert(state != NULL);
140    if (!PyArg_Parse(key, "s#", &krec.dptr, &tmp_size)) {
141        return NULL;
142    }
143
144    krec.dsize = tmp_size;
145    check_dbmobject_open(dp, state->dbm_error);
146    drec = dbm_fetch(dp->di_dbm, krec);
147    if ( drec.dptr == 0 ) {
148        PyErr_SetObject(PyExc_KeyError, key);
149        return NULL;
150    }
151    if ( dbm_error(dp->di_dbm) ) {
152        dbm_clearerr(dp->di_dbm);
153        PyErr_SetString(state->dbm_error, "");
154        return NULL;
155    }
156    return PyBytes_FromStringAndSize(drec.dptr, drec.dsize);
157}
158
159static int
160dbm_ass_sub(dbmobject *dp, PyObject *v, PyObject *w)
161{
162    datum krec, drec;
163    Py_ssize_t tmp_size;
164
165    if ( !PyArg_Parse(v, "s#", &krec.dptr, &tmp_size) ) {
166        PyErr_SetString(PyExc_TypeError,
167                        "dbm mappings have bytes or string keys only");
168        return -1;
169    }
170    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
171    assert(state != NULL);
172    krec.dsize = tmp_size;
173    if (dp->di_dbm == NULL) {
174             PyErr_SetString(state->dbm_error, "DBM object has already been closed");
175             return -1;
176    }
177    dp->di_size = -1;
178    if (w == NULL) {
179        if ( dbm_delete(dp->di_dbm, krec) < 0 ) {
180            dbm_clearerr(dp->di_dbm);
181            /* we might get a failure for reasons like file corrupted,
182               but we are not able to distinguish it */
183            if (dp->flags & O_RDWR) {
184                PyErr_SetObject(PyExc_KeyError, v);
185            }
186            else {
187                PyErr_SetString(state->dbm_error, "cannot delete item from database");
188            }
189            return -1;
190        }
191    } else {
192        if ( !PyArg_Parse(w, "s#", &drec.dptr, &tmp_size) ) {
193            PyErr_SetString(PyExc_TypeError,
194                 "dbm mappings have bytes or string elements only");
195            return -1;
196        }
197        drec.dsize = tmp_size;
198        if ( dbm_store(dp->di_dbm, krec, drec, DBM_REPLACE) < 0 ) {
199            dbm_clearerr(dp->di_dbm);
200            PyErr_SetString(state->dbm_error,
201                            "cannot add item to database");
202            return -1;
203        }
204    }
205    if ( dbm_error(dp->di_dbm) ) {
206        dbm_clearerr(dp->di_dbm);
207        PyErr_SetString(state->dbm_error, "");
208        return -1;
209    }
210    return 0;
211}
212
213/*[clinic input]
214_dbm.dbm.close
215
216Close the database.
217[clinic start generated code]*/
218
219static PyObject *
220_dbm_dbm_close_impl(dbmobject *self)
221/*[clinic end generated code: output=c8dc5b6709600b86 input=046db72377d51be8]*/
222{
223    if (self->di_dbm) {
224        dbm_close(self->di_dbm);
225    }
226    self->di_dbm = NULL;
227    Py_RETURN_NONE;
228}
229
230/*[clinic input]
231_dbm.dbm.keys
232
233    cls: defining_class
234
235Return a list of all keys in the database.
236[clinic start generated code]*/
237
238static PyObject *
239_dbm_dbm_keys_impl(dbmobject *self, PyTypeObject *cls)
240/*[clinic end generated code: output=f2a593b3038e5996 input=d3706a28fc051097]*/
241{
242    PyObject *v, *item;
243    datum key;
244    int err;
245
246    _dbm_state *state = PyType_GetModuleState(cls);
247    assert(state != NULL);
248    check_dbmobject_open(self, state->dbm_error);
249    v = PyList_New(0);
250    if (v == NULL) {
251        return NULL;
252    }
253    for (key = dbm_firstkey(self->di_dbm); key.dptr;
254         key = dbm_nextkey(self->di_dbm)) {
255        item = PyBytes_FromStringAndSize(key.dptr, key.dsize);
256        if (item == NULL) {
257            Py_DECREF(v);
258            return NULL;
259        }
260        err = PyList_Append(v, item);
261        Py_DECREF(item);
262        if (err != 0) {
263            Py_DECREF(v);
264            return NULL;
265        }
266    }
267    return v;
268}
269
270static int
271dbm_contains(PyObject *self, PyObject *arg)
272{
273    dbmobject *dp = (dbmobject *)self;
274    datum key, val;
275    Py_ssize_t size;
276
277    _dbm_state *state = PyType_GetModuleState(Py_TYPE(dp));
278    assert(state != NULL);
279    if ((dp)->di_dbm == NULL) {
280        PyErr_SetString(state->dbm_error,
281                        "DBM object has already been closed");
282         return -1;
283    }
284    if (PyUnicode_Check(arg)) {
285        key.dptr = (char *)PyUnicode_AsUTF8AndSize(arg, &size);
286        key.dsize = size;
287        if (key.dptr == NULL)
288            return -1;
289    }
290    else if (!PyBytes_Check(arg)) {
291        PyErr_Format(PyExc_TypeError,
292                     "dbm key must be bytes or string, not %.100s",
293                     Py_TYPE(arg)->tp_name);
294        return -1;
295    }
296    else {
297        key.dptr = PyBytes_AS_STRING(arg);
298        key.dsize = PyBytes_GET_SIZE(arg);
299    }
300    val = dbm_fetch(dp->di_dbm, key);
301    return val.dptr != NULL;
302}
303
304/*[clinic input]
305_dbm.dbm.get
306    cls: defining_class
307    key: str(accept={str, robuffer}, zeroes=True)
308    default: object = None
309    /
310
311Return the value for key if present, otherwise default.
312[clinic start generated code]*/
313
314static PyObject *
315_dbm_dbm_get_impl(dbmobject *self, PyTypeObject *cls, const char *key,
316                  Py_ssize_t key_length, PyObject *default_value)
317/*[clinic end generated code: output=b4e55f8b6d482bc4 input=66b993b8349fa8c1]*/
318{
319    datum dbm_key, val;
320    _dbm_state *state = PyType_GetModuleState(cls);
321    assert(state != NULL);
322    dbm_key.dptr = (char *)key;
323    dbm_key.dsize = key_length;
324    check_dbmobject_open(self, state->dbm_error);
325    val = dbm_fetch(self->di_dbm, dbm_key);
326    if (val.dptr != NULL) {
327        return PyBytes_FromStringAndSize(val.dptr, val.dsize);
328    }
329
330    Py_INCREF(default_value);
331    return default_value;
332}
333
334/*[clinic input]
335_dbm.dbm.setdefault
336    cls: defining_class
337    key: str(accept={str, robuffer}, zeroes=True)
338    default: object(c_default="NULL") = b''
339    /
340
341Return the value for key if present, otherwise default.
342
343If key is not in the database, it is inserted with default as the value.
344[clinic start generated code]*/
345
346static PyObject *
347_dbm_dbm_setdefault_impl(dbmobject *self, PyTypeObject *cls, const char *key,
348                         Py_ssize_t key_length, PyObject *default_value)
349/*[clinic end generated code: output=9c2f6ea6d0fb576c input=126a3ff15c5f8232]*/
350{
351    datum dbm_key, val;
352    Py_ssize_t tmp_size;
353    _dbm_state *state = PyType_GetModuleState(cls);
354    assert(state != NULL);
355    dbm_key.dptr = (char *)key;
356    dbm_key.dsize = key_length;
357    check_dbmobject_open(self, state->dbm_error);
358    val = dbm_fetch(self->di_dbm, dbm_key);
359    if (val.dptr != NULL) {
360        return PyBytes_FromStringAndSize(val.dptr, val.dsize);
361    }
362    if (default_value == NULL) {
363        default_value = PyBytes_FromStringAndSize(NULL, 0);
364        if (default_value == NULL) {
365            return NULL;
366        }
367        val.dptr = NULL;
368        val.dsize = 0;
369    }
370    else {
371        if ( !PyArg_Parse(default_value, "s#", &val.dptr, &tmp_size) ) {
372            PyErr_SetString(PyExc_TypeError,
373                "dbm mappings have bytes or string elements only");
374            return NULL;
375        }
376        val.dsize = tmp_size;
377        Py_INCREF(default_value);
378    }
379    if (dbm_store(self->di_dbm, dbm_key, val, DBM_INSERT) < 0) {
380        dbm_clearerr(self->di_dbm);
381        PyErr_SetString(state->dbm_error, "cannot add item to database");
382        Py_DECREF(default_value);
383        return NULL;
384    }
385    return default_value;
386}
387
388static PyObject *
389dbm__enter__(PyObject *self, PyObject *args)
390{
391    Py_INCREF(self);
392    return self;
393}
394
395static PyObject *
396dbm__exit__(PyObject *self, PyObject *args)
397{
398    return _dbm_dbm_close_impl((dbmobject *)self);
399}
400
401static PyMethodDef dbm_methods[] = {
402    _DBM_DBM_CLOSE_METHODDEF
403    _DBM_DBM_KEYS_METHODDEF
404    _DBM_DBM_GET_METHODDEF
405    _DBM_DBM_SETDEFAULT_METHODDEF
406    {"__enter__", dbm__enter__, METH_NOARGS, NULL},
407    {"__exit__",  dbm__exit__, METH_VARARGS, NULL},
408    {NULL,  NULL}           /* sentinel */
409};
410
411static PyType_Slot dbmtype_spec_slots[] = {
412    {Py_tp_dealloc, dbm_dealloc},
413    {Py_tp_traverse, dbm_traverse},
414    {Py_tp_methods, dbm_methods},
415    {Py_sq_contains, dbm_contains},
416    {Py_mp_length, dbm_length},
417    {Py_mp_subscript, dbm_subscript},
418    {Py_mp_ass_subscript, dbm_ass_sub},
419    {0, 0}
420};
421
422
423static PyType_Spec dbmtype_spec = {
424    .name = "_dbm.dbm",
425    .basicsize = sizeof(dbmobject),
426    // Calling PyType_GetModuleState() on a subclass is not safe.
427    // dbmtype_spec does not have Py_TPFLAGS_BASETYPE flag
428    // which prevents to create a subclass.
429    // So calling PyType_GetModuleState() in this file is always safe.
430    .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION |
431              Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_IMMUTABLETYPE),
432    .slots = dbmtype_spec_slots,
433};
434
435/* ----------------------------------------------------------------- */
436
437/*[clinic input]
438
439_dbm.open as dbmopen
440
441    filename: object
442        The filename to open.
443
444    flags: str="r"
445        How to open the file.  "r" for reading, "w" for writing, etc.
446
447    mode: int(py_default="0o666") = 0o666
448        If creating a new file, the mode bits for the new file
449        (e.g. os.O_RDWR).
450
451    /
452
453Return a database object.
454
455[clinic start generated code]*/
456
457static PyObject *
458dbmopen_impl(PyObject *module, PyObject *filename, const char *flags,
459             int mode)
460/*[clinic end generated code: output=9527750f5df90764 input=d8cf50a9f81218c8]*/
461{
462    int iflags;
463    _dbm_state *state =  get_dbm_state(module);
464    assert(state != NULL);
465    if (strcmp(flags, "r") == 0) {
466        iflags = O_RDONLY;
467    }
468    else if (strcmp(flags, "w") == 0) {
469        iflags = O_RDWR;
470    }
471    else if (strcmp(flags, "rw") == 0) {
472        /* Backward compatibility */
473        iflags = O_RDWR|O_CREAT;
474    }
475    else if (strcmp(flags, "c") == 0) {
476        iflags = O_RDWR|O_CREAT;
477    }
478    else if (strcmp(flags, "n") == 0) {
479        iflags = O_RDWR|O_CREAT|O_TRUNC;
480    }
481    else {
482        PyErr_SetString(state->dbm_error,
483                        "arg 2 to open should be 'r', 'w', 'c', or 'n'");
484        return NULL;
485    }
486
487    PyObject *filenamebytes;
488    if (!PyUnicode_FSConverter(filename, &filenamebytes)) {
489        return NULL;
490    }
491
492    const char *name = PyBytes_AS_STRING(filenamebytes);
493    if (strlen(name) != (size_t)PyBytes_GET_SIZE(filenamebytes)) {
494        Py_DECREF(filenamebytes);
495        PyErr_SetString(PyExc_ValueError, "embedded null character");
496        return NULL;
497    }
498    PyObject *self = newdbmobject(state, name, iflags, mode);
499    Py_DECREF(filenamebytes);
500    return self;
501}
502
503static PyMethodDef dbmmodule_methods[] = {
504    DBMOPEN_METHODDEF
505    { 0, 0 },
506};
507
508static int
509_dbm_exec(PyObject *module)
510{
511    _dbm_state *state = get_dbm_state(module);
512    state->dbm_type = (PyTypeObject *)PyType_FromModuleAndSpec(module,
513                                                        &dbmtype_spec, NULL);
514    if (state->dbm_type == NULL) {
515        return -1;
516    }
517    state->dbm_error = PyErr_NewException("_dbm.error", PyExc_OSError, NULL);
518    if (state->dbm_error == NULL) {
519        return -1;
520    }
521    if (PyModule_AddStringConstant(module, "library", which_dbm) < 0) {
522        return -1;
523    }
524    if (PyModule_AddType(module, (PyTypeObject *)state->dbm_error) < 0) {
525        return -1;
526    }
527    return 0;
528}
529
530static int
531_dbm_module_traverse(PyObject *module, visitproc visit, void *arg)
532{
533    _dbm_state *state = get_dbm_state(module);
534    Py_VISIT(state->dbm_error);
535    Py_VISIT(state->dbm_type);
536    return 0;
537}
538
539static int
540_dbm_module_clear(PyObject *module)
541{
542    _dbm_state *state = get_dbm_state(module);
543    Py_CLEAR(state->dbm_error);
544    Py_CLEAR(state->dbm_type);
545    return 0;
546}
547
548static void
549_dbm_module_free(void *module)
550{
551    _dbm_module_clear((PyObject *)module);
552}
553
554static PyModuleDef_Slot _dbmmodule_slots[] = {
555    {Py_mod_exec, _dbm_exec},
556    {0, NULL}
557};
558
559static struct PyModuleDef _dbmmodule = {
560    PyModuleDef_HEAD_INIT,
561    .m_name = "_dbm",
562    .m_size = sizeof(_dbm_state),
563    .m_methods = dbmmodule_methods,
564    .m_slots = _dbmmodule_slots,
565    .m_traverse = _dbm_module_traverse,
566    .m_clear = _dbm_module_clear,
567    .m_free = _dbm_module_free,
568};
569
570PyMODINIT_FUNC
571PyInit__dbm(void)
572{
573    return PyModuleDef_Init(&_dbmmodule);
574}
575