Check for proper base types when receiving arguments of type Type.

3 files changed, 22 insertions, 20 deletions

diff --git a/src/conn.c b/src/conn.c
index aaed18f..03fb1e6 100644
--- a/src/conn.c
+++ b/src/conn.c
@@ -27,7 +27,6 @@ xpybConn_make_ext(xpybConn *self, PyObject *key)
     PyObject *result, *arglist;
     xpybExt *ext;
     const xcb_query_extension_reply_t *reply;
-    int rc;
 
     /* Look up the callable object in the global dictionary. */
     result = PyDict_GetItem(xpybModule_extdict, key);
@@ -49,18 +48,6 @@ xpybConn_make_ext(xpybConn *self, PyObject *key)
     if (ext == NULL)
 	return NULL;
 
-    /* Make sure what we got is actually an xcb.Extension */
-    rc = PyObject_IsInstance((PyObject *)ext, (PyObject *)&xpybExt_type);
-    switch (rc) {
-    case 1:
-	break;
-    case 0:
-	PyErr_SetString(xpybExcept_ext, "Invalid extension object returned.");
-    default:
-	Py_DECREF(ext);
-	return NULL;
-    }
-
     /* Get the opcode and base numbers for actual (non-core) extensions. */
     if (key != Py_None) {
 	reply = xcb_get_extension_data(self->conn, &((xpybExtkey *)key)->key);
diff --git a/src/ext.c b/src/ext.c
index d2a0412..90dca2d 100644
--- a/src/ext.c
+++ b/src/ext.c
@@ -102,13 +102,18 @@ xpybExt_send_request(xpybExt *self, PyObject *args, PyObject *kw)
     const void *data;
     Py_ssize_t size;
 
-    /* Parse arguments and determine number of data strings. */
+    /* Parse and check arguments */
     if (!PyArg_ParseTupleAndKeywords(args, kw, "O!O!|O!", kwlist,
 				     &xpybRequest_type, &request,
 				     &xpybCookie_type, &cookie,
 				     &PyType_Type, &reply))
 	return NULL;
 
+    if (!PyType_IsSubtype(reply, &xpybReply_type)) {
+	PyErr_SetString(xpybExcept_base, "Type not derived from xcb.Reply.");
+	return NULL;
+    }
+
     /* Set up request structure */
     xcb_req.count = 2;
     xcb_req.ext = (self->key != (xpybExtkey *)Py_None) ? &self->key->key : 0;
diff --git a/src/module.c b/src/module.c
index dabb63c..43aa722 100644
--- a/src/module.c
+++ b/src/module.c
@@ -64,12 +64,17 @@ err1:
 static PyObject *
 xpyb_add_core(PyObject *self, PyObject *args)
 {
-    PyObject *value;
+    PyTypeObject *value;
 
     if (!PyArg_ParseTuple(args, "O!", &PyType_Type, &value))
 	return NULL;
 
-    if (PyDict_SetItem(xpybModule_extdict, Py_None, value) < 0)
+    if (!PyType_IsSubtype(value, &xpybExt_type)) {
+	PyErr_SetString(xpybExcept_base, "Type not derived from xcb.Extension.");
+	return NULL;
+    }
+
+    if (PyDict_SetItem(xpybModule_extdict, Py_None, (PyObject *)value) < 0)
 	return NULL;
 
     Py_RETURN_NONE;
@@ -78,13 +83,18 @@ xpyb_add_core(PyObject *self, PyObject *args)
 static PyObject *
 xpyb_add_ext(PyObject *self, PyObject *args)
 {
-    PyObject *key = Py_None;
-    PyObject *value;
+    PyTypeObject *value;
+    PyObject *key;
 
-    if (!PyArg_ParseTuple(args, "O!|O!", &PyType_Type, &value, &xpybExtkey_type, &key))
+    if (!PyArg_ParseTuple(args, "O!O!", &PyType_Type, &value, &xpybExtkey_type, &key))
 	return NULL;
 
-    if (PyDict_SetItem(xpybModule_extdict, key, value) < 0)
+    if (!PyType_IsSubtype(value, &xpybExt_type)) {
+	PyErr_SetString(xpybExcept_base, "Type not derived from xcb.Extension.");
+	return NULL;
+    }
+
+    if (PyDict_SetItem(xpybModule_extdict, key, (PyObject *)value) < 0)
 	return NULL;
 
     Py_RETURN_NONE;