Repository: spotify/annoy (language: python)

Issue:
Missing error checks in Python extension code
The code in https://github.com/spotify/annoy/blob/master/src/annoymodule.cc that interacts with the CPython API mostly doesn't check the return value of called functions. This can lead to crashes:

```
% python -c 'import annoy
class X:
    def __len__(self): return 42
ann = annoy.AnnoyIndex(42)
ann.get_nns_by_vector(X(), 42)
'
<string>:4: FutureWarning: The default argument for metric will be removed in future version of Annoy. Please pass metric='angular' explicitly.
zsh: segmentation fault  python -c 
```
Similarly, calling `get_nns_by_vector()` with a vector of objects that aren't floats throws a `SystemError`.

There are also no checks on the return value of functions like `PyList_New`. In theory, this could lead to segfaults, because the function returns NULL on error. I believe that only happens if the system is out of memory though.


Discussion / hints:
is it easy to fix? sounds worth doing

Fix patch:
diff --git a/src/annoymodule.cc b/src/annoymodule.cc
index 286412f4..fb0f4c15 100644
--- a/src/annoymodule.cc
+++ b/src/annoymodule.cc
@@ -237,16 +237,25 @@ py_an_save(py_annoy *self, PyObject *args, PyObject *kwargs) {
 PyObject*
 get_nns_to_python(const vector<int32_t>& result, const vector<float>& distances, int include_distances) {
   PyObject* l = PyList_New(result.size());
+  if (l == NULL) {
+    return NULL;
+  }
   for (size_t i = 0; i < result.size(); i++)
     PyList_SetItem(l, i, PyInt_FromLong(result[i]));
   if (!include_distances)
     return l;
 
   PyObject* d = PyList_New(distances.size());
+  if (d == NULL) {
+    return NULL;
+  }
   for (size_t i = 0; i < distances.size(); i++)
     PyList_SetItem(d, i, PyFloat_FromDouble(distances[i]));
 
   PyObject* t = PyTuple_New(2);
+  if (t == NULL) {
+    return NULL;
+  }
   PyTuple_SetItem(t, 0, l);
   PyTuple_SetItem(t, 1, d);
 
@@ -293,24 +302,31 @@ py_an_get_nns_by_item(py_annoy *self, PyObject *args, PyObject *kwargs) {
 
 bool
 convert_list_to_vector(PyObject* v, int f, vector<float>* w) {
-  if (PyObject_Size(v) == -1) {
-    char buf[256];
-    snprintf(buf, 256, "Expected an iterable, got an object of type \"%s\"", v->ob_type->tp_name);
-    PyErr_SetString(PyExc_ValueError, buf);
+  Py_ssize_t length = PyObject_Size(v);
+  if (length == -1) {
     return false;
   }
-  if (PyObject_Size(v) != f) {
-    char buf[128];
-    snprintf(buf, 128, "Vector has wrong length (expected %d, got %ld)", f, PyObject_Size(v));
-    PyErr_SetString(PyExc_IndexError, buf);
+  if (length != f) {
+    PyErr_Format(PyExc_IndexError, "Vector has wrong length (expected %d, got %ld)", f, length);
     return false;
   }
+
   for (int z = 0; z < f; z++) {
     PyObject *key = PyInt_FromLong(z);
+    if (key == NULL) {
+      return false;
+    }
     PyObject *pf = PyObject_GetItem(v, key);
-    (*w)[z] = PyFloat_AsDouble(pf);
     Py_DECREF(key);
+    if (pf == NULL) {
+      return false;
+    }
+    double value = PyFloat_AsDouble(pf);
     Py_DECREF(pf);
+    if (value == -1.0 && PyErr_Occurred()) {
+      return false;
+    }
+    (*w)[z] = value;
   }
   return true;
 }
@@ -357,6 +373,9 @@ py_an_get_item_vector(py_annoy *self, PyObject *args) {
   vector<float> v(self->f);
   self->ptr->get_item(item, &v[0]);
   PyObject* l = PyList_New(self->f);
+  if (l == NULL) {
+    return NULL;
+  }
   for (int z = 0; z < self->f; z++) {
     PyList_SetItem(l, z, PyFloat_FromDouble(v[z]));
   }


Test patch:
diff --git a/test/types_test.py b/test/types_test.py
index f820d5df..b79a6d08 100644
--- a/test/types_test.py
+++ b/test/types_test.py
@@ -13,6 +13,7 @@
 # the License.
 
 import numpy
+import sys
 import random
 from common import TestCase
 from annoy import AnnoyIndex
@@ -56,3 +57,63 @@ def test_range_errors(self, n_points=1000, n_trees=10):
             self.assertRaises(IndexError, i.get_distance, 0, bad_index)
             self.assertRaises(IndexError, i.get_nns_by_item, bad_index, 1)
             self.assertRaises(IndexError, i.get_item_vector, bad_index)
+
+    def test_missing_len(self):
+        """
+        We should get a helpful error message if our vector doesn't have a
+        __len__ method.
+        """
+        class FakeCollection:
+            pass
+        
+        i = AnnoyIndex(10, 'euclidean')
+        # Python 2.7 raises an AttributeError instead of a TypeError like newer versions of Python.
+        if sys.version_info.major == 2:
+            with self.assertRaises(AttributeError, msg="FakeCollection instance has no attribute '__len__'"):
+                i.add_item(1, FakeCollection())
+        else:
+            with self.assertRaises(TypeError, msg="object of type 'FakeCollection' has no len()"):
+                i.add_item(1, FakeCollection())
+
+    def test_missing_getitem(self):
+        """
+        We should get a helpful error message if our vector doesn't have a
+        __getitem__ method.
+        """
+        class FakeCollection:
+            def __len__(self):
+                return 5
+        
+        i = AnnoyIndex(5, 'euclidean')
+        # Python 2.7 raises an AttributeError instead of a TypeError like newer versions of Python.
+        if sys.version_info.major == 2:
+            with self.assertRaises(AttributeError, msg="FakeCollection instance has no attribute '__getitem__'"):
+                i.add_item(1, FakeCollection())
+        else:
+            with self.assertRaises(TypeError, msg="'FakeCollection' object is not subscriptable"):
+                i.add_item(1, FakeCollection())
+
+    def test_short(self):
+        """
+        Ensure we handle our vector not being long enough.
+        """
+        class FakeCollection:
+            def __len__(self):
+                return 3
+            
+            def __getitem__(self, i):
+                raise IndexError 
+        
+        i = AnnoyIndex(3, 'euclidean')
+        with self.assertRaises(IndexError):
+            i.add_item(1, FakeCollection())
+    
+    def test_non_float(self):
+        """
+        We should error gracefully if non-floats are provided in our vector.
+        """
+        array_strings = ["1", "2", "3"]
+        
+        i = AnnoyIndex(3, 'euclidean')
+        with self.assertRaises(TypeError, msg="must be real number, not str"):
+            i.add_item(1, array_strings)