From e100ddcb2061c6c014e534de24f8c11727d7d3ee Mon Sep 17 00:00:00 2001
From: Kyle Erf <erf@mongodb.com>
Date: Tue, 14 Jan 2014 02:19:03 -0500
Subject: [PATCH] PYTHON-623: do not allow sending _ids with a key starting
 with $

---
 bson/__init__.py    | 6 ++++--
 bson/_cbsonmodule.c | 5 ++---
 test/test_bson.py   | 6 ++++++
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/bson/__init__.py b/bson/__init__.py
index 7c3ac4a7e..f0e13c7a4 100644
--- a/bson/__init__.py
+++ b/bson/__init__.py
@@ -477,10 +477,12 @@ def _dict_to_bson(dict, check_keys, uuid_subtype, top_level=True):
     try:
         elements = []
         if top_level and "_id" in dict:
-            elements.append(_element_to_bson("_id", dict["_id"], False, uuid_subtype))
+            elements.append(_element_to_bson("_id", dict["_id"],
+                                             check_keys, uuid_subtype))
         for (key, value) in dict.iteritems():
             if not top_level or key != "_id":
-                elements.append(_element_to_bson(key, value, check_keys, uuid_subtype))
+                elements.append(_element_to_bson(key, value,
+                                                 check_keys, uuid_subtype))
     except AttributeError:
         raise TypeError("encoder expected a mapping type but got: %r" % dict)
 
diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c
index c8901857c..30c8ad670 100644
--- a/bson/_cbsonmodule.c
+++ b/bson/_cbsonmodule.c
@@ -1341,9 +1341,8 @@ int write_dict(PyObject* self, buffer_t buffer,
     if (top_level) {
         PyObject* _id = PyDict_GetItemString(dict, "_id");
         if (_id) {
-            /* Don't bother checking keys, but do make sure we're allowed to
-             * write _id */
-            if (!write_pair(self, buffer, "_id", 3, _id, 0, uuid_subtype, 1)) {
+            if (!write_pair(self, buffer, "_id", 3,
+                            _id, check_keys, uuid_subtype, 1)) {
                 return 0;
             }
         }
diff --git a/test/test_bson.py b/test/test_bson.py
index 6f8a37029..d6f637c7d 100644
--- a/test/test_bson.py
+++ b/test/test_bson.py
@@ -676,6 +676,12 @@ class TestBSON(unittest.TestCase):
         self.assertTrue(Timestamp(1, 0) <= Timestamp(1, 0))
         self.assertFalse(Timestamp(1, 0) > Timestamp(1, 0))
 
+    def test_bad_id_keys(self):
+        self.assertRaises(InvalidDocument, BSON.encode,
+                          {"_id": {"$bad": 123}}, True)
+        self.assertRaises(InvalidDocument, BSON.encode,
+                          {"_id": {'$oid': "52d0b971b3ba219fdeb4170e"}}, True)
+        BSON.encode({"_id": {'$oid': "52d0b971b3ba219fdeb4170e"}})
 
 if __name__ == "__main__":
     unittest.main()