fix ticket 903782: initialise document dict for HTML parsing in iterparse to prevent segfaults on cleanup

scoder · scoder · commit 6744a095ce80 · 2011-12-14T22:11:27.000+01:00
--HG--
extra : transplant_source : G%BE%D7Pe%05%03%B8C%24%10%25I%7D%834%5B%E2%C9%02
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -32,15 +32,18 @@ Bugs fixed
   a descendant combinator. For example, "div> .foo" was parsed the same as
   "div>* .foo" instead of "div>.foo". Patch by Simon Sapin.
 
-* lxml.html.diff no longer raises an exception when hitting
-  'img' tags without 'src' attribute.
+* Fixed a crash when using ``iterparse()`` for HTML parsing and
+  requesting start events.
 
 * Fixed parsing of more selectors in cssselect.  Whitespace before
   pseudo-elements and pseudo-classes is significant as it is a
   descendant combinator.
   "E :pseudo" should parse the same as "E \*:pseudo", not "E:pseudo".
   Patch by Simon Sapin.
 
+* lxml.html.diff no longer raises an exception when hitting
+  'img' tags without 'src' attribute.
+
 Other changes
 --------------
 
diff --git a/src/lxml/iterparse.pxi b/src/lxml/iterparse.pxi
@@ -64,6 +64,7 @@ cdef class _IterparseContext(_ParserContext):
     cdef xmlparser.endElementNsSAX2Func   _origSaxEnd
     cdef xmlparser.startElementSAXFunc    _origSaxStartNoNs
     cdef xmlparser.endElementSAXFunc      _origSaxEndNoNs
+    cdef xmlparser.startDocumentSAXFunc   _origSaxStartDocument
     cdef xmlparser.commentSAXFunc         _origSaxComment
     cdef xmlparser.processingInstructionSAXFunc _origSaxPI
     cdef _Element  _root
@@ -88,6 +89,8 @@ cdef class _IterparseContext(_ParserContext):
         cdef xmlparser.xmlSAXHandler* sax
         _ParserContext._initParserContext(self, c_ctxt)
         sax = c_ctxt.sax
+        self._origSaxStartDocument = sax.startDocument
+        sax.startDocument = _iterparseSaxStartDocument
         self._origSaxStart = sax.startElementNs
         self._origSaxStartNoNs = sax.startElement
         # only override start event handler if needed
@@ -231,6 +234,16 @@ cdef inline void _pushSaxEvent(_IterparseContext context,
         context._c_ctxt.disableSAX = 1
         context._store_raised()
 
+cdef void _iterparseSaxStartDocument(void* ctxt):
+    cdef xmlparser.xmlParserCtxt* c_ctxt
+    c_ctxt = <xmlparser.xmlParserCtxt*>ctxt
+    context = <_IterparseContext>c_ctxt._private
+    context._origSaxStartDocument(ctxt)
+    if c_ctxt.myDoc and c_ctxt.dict and not c_ctxt.myDoc.dict:
+        # I have no idea why libxml2 disables this - we need it
+        c_ctxt.dictNames = 1
+        c_ctxt.myDoc.dict = c_ctxt.dict
+
 cdef void _iterparseSaxStart(void* ctxt, char* localname, char* prefix,
                              char* URI, int nb_namespaces, char** namespaces,
                              int nb_attributes, int nb_defaulted,
diff --git a/src/lxml/tests/test_htmlparser.py b/src/lxml/tests/test_htmlparser.py
@@ -288,6 +288,22 @@ def test_html_iterparse_file(self):
             [],
             [ event for (event, element) in events if event != 'end' ])
 
+    def test_html_iterparse_start(self):
+        iterparse = self.etree.iterparse
+        f = BytesIO(
+            '<html><head><title>TITLE</title><body><p>P</p></body></html>')
+
+        iterator = iterparse(f, html=True, events=('start',))
+        self.assertEquals(None, iterator.root)
+
+        events = list(iterator)
+        root = iterator.root
+        self.assert_(root is not None)
+        self.assertEquals(
+            [('start', root), ('start', root[0]), ('start', root[0][0]),
+                ('start', root[1]), ('start', root[1][0])],
+            events)
+
 def test_suite():
     suite = unittest.TestSuite()
     suite.addTests([unittest.makeSuite(HtmlParserTestCase)])
diff --git a/src/lxml/xmlparser.pxd b/src/lxml/xmlparser.pxd
@@ -40,6 +40,8 @@ cdef extern from "libxml/parser.h":
 
     ctypedef void (*endDocumentSAXFunc)(void* ctx)
 
+    ctypedef void (*startDocumentSAXFunc)(void* ctx)
+
     ctypedef void (*referenceSAXFunc)(void * ctx, char* name)
 
     cdef int XML_SAX2_MAGIC
@@ -68,6 +70,7 @@ cdef extern from "libxml/tree.h":
         referenceSAXFunc                reference
         commentSAXFunc                  comment
         processingInstructionSAXFunc	processingInstruction
+        startDocumentSAXFunc            startDocument
         endDocumentSAXFunc              endDocument
         int                             initialized
 
@@ -86,6 +89,7 @@ cdef extern from "libxml/parser.h":
     ctypedef struct xmlParserCtxt:
         xmlDoc* myDoc
         xmlDict* dict
+        int dictNames
         void* _private
         bint wellFormed
         bint recovery
