Add handlers for when the rewriting is finished. (#27)

* Allow handlers for the end of parsing. * Make the end handler a document content handler. * Allow appending at the end of the document. * Append chunks on the fly. * Remove unneeded example. * Propagate error to `end()`. * Propagate the error immediately. * Store a reference to the OutputSink in DocumentEnd. * Test if appending at the end works. * Extend the C API to support document end handlers. * Refactor condition. * Pass user data to the document end handler. * Document the C API function. * Add regression test.
cloudflare · Nov 26, 2019 · 1aedf36 · 1aedf36
1 parent b09b940
commit 1aedf36
Show file tree

Hide file tree

Showing 23 changed files with 420 additions and 28 deletions.
diff --git a/c-api/include/lol_html.h b/c-api/include/lol_html.h
@@ -22,6 +22,7 @@ extern "C" {
 typedef struct lol_html_HtmlRewriterBuilder lol_html_rewriter_builder_t;
 typedef struct lol_html_HtmlRewriter lol_html_rewriter_t;
 typedef struct lol_html_Doctype lol_html_doctype_t;
+typedef struct lol_html_DocumentEnd lol_html_doc_end_t;
 typedef struct lol_html_Comment lol_html_comment_t;
 typedef struct lol_html_TextChunk lol_html_text_chunk_t;
 typedef struct lol_html_Element lol_html_element_t;
@@ -75,7 +76,7 @@ lol_html_rewriter_builder_t *lol_html_rewriter_builder_new();
 // Content handlers
 //---------------------------------------------------------------------
 // Rewriter directive that should be returned from each content handler.
-// If LOL_HTML_STOP directive is returned then rewriting stops immidiately
+// If LOL_HTML_STOP directive is returned then rewriting stops immediately
 // and `write()` or `end()` methods of the rewriter return an error code.
 typedef enum {
     LOL_HTML_CONTINUE,
@@ -102,12 +103,17 @@ typedef lol_html_rewriter_directive_t (*lol_html_element_handler_t)(
     void *user_data
 );
 
+typedef lol_html_rewriter_directive_t (*lol_html_doc_end_handler_t)(
+    lol_html_doc_end_t *doc_end,
+    void *user_data
+);
+
 // Selector
 //---------------------------------------------------------------------
 
 // Parses given CSS selector string.
 //
-// Returns NULL if parsing error occures. The actual error message
+// Returns NULL if parsing error occurs. The actual error message
 // can be obtained using `lol_html_take_last_error` function.
 //
 // WARNING: Selector SHOULD NOT be deallocated if there are any active rewriter
@@ -137,8 +143,8 @@ void lol_html_selector_free(lol_html_selector_t *selector);
 // passed to the handler on each invocation along with the rewritable
 // unit argument.
 //
-// If any of handlers return LOL_HTML_STOP directive is then rewriting
-// stops immidiately and `write()` or `end()` of the rewriter methods
+// If any of handlers return LOL_HTML_STOP directive then rewriting
+// stops immediately and `write()` or `end()` of the rewriter methods
 // return an error code.
 //
 // WARNING: Pointers passed to handlers are valid only during the
@@ -150,7 +156,9 @@ void lol_html_rewriter_builder_add_document_content_handlers(
     lol_html_comment_handler_t comment_handler,
     void *comment_handler_user_data,
     lol_html_text_handler_handler_t text_handler,
-    void *text_handler_user_data
+    void *text_handler_user_data,
+    lol_html_doc_end_handler_t doc_end_handler,
+    void *doc_end_user_data
 );
 
 // Adds element content handlers to the builder for the
@@ -167,8 +175,8 @@ void lol_html_rewriter_builder_add_document_content_handlers(
 // passed to the handler on each invocation along with the rewritable
 // unit argument.
 //
-// If any of handlers return LOL_HTML_STOP directive is then rewriting
-// stops immidiately and `write()` or `end()` of the rewriter methods
+// If any of handlers return LOL_HTML_STOP directive then rewriting
+// stops immediately and `write()` or `end()` of the rewriter methods
 // return an error code.
 //
 // Returns 0 in case of success and -1 otherwise. The actual error message
@@ -651,6 +659,19 @@ void lol_html_element_user_data_set(
 // Returns user data attached to the text chunk.
 void *lol_html_element_user_data_get(const lol_html_element_t *element);
 
+// Inserts the content at the end of the document, either as raw text or as HTML.
+//
+// The content should be a valid UTF-8 string.
+//
+// Returns 0 if successful, and -1 otherwise. The actual error message
+// can be obtained using the `lol_html_take_last_error` function.
+int lol_html_doc_end_append(
+    lol_html_doc_end_t *doc_end,
+    const char *content,
+    size_t content_len,
+    bool is_html
+);
+
 #if defined(__cplusplus)
 }  // extern C
 #endif

diff --git a/c-api/src/document_end.rs b/c-api/src/document_end.rs
@@ -0,0 +1,11 @@
+use super::*;
+
+#[no_mangle]
+pub extern "C" fn lol_html_doc_end_append(
+    document_end: *mut DocumentEnd,
+    content: *const c_char,
+    content_len: size_t,
+    is_html: bool,
+) -> c_int {
+    content_insertion_fn_body! { document_end.append(content, content_len, is_html) }
+}
diff --git a/c-api/src/lib.rs b/c-api/src/lib.rs
@@ -117,6 +117,7 @@ macro_rules! get_user_data {
 
 mod comment;
 mod doctype;
+mod document_end;
 mod element;
 mod errors;
 mod rewriter;

diff --git a/c-api/src/rewriter_builder.rs b/c-api/src/rewriter_builder.rs
@@ -11,6 +11,7 @@ type ElementHandler = unsafe extern "C" fn(*mut Element, *mut c_void) -> Rewrite
 type DoctypeHandler = unsafe extern "C" fn(*mut Doctype, *mut c_void) -> RewriterDirective;
 type CommentsHandler = unsafe extern "C" fn(*mut Comment, *mut c_void) -> RewriterDirective;
 type TextHandler = unsafe extern "C" fn(*mut TextChunk, *mut c_void) -> RewriterDirective;
+type DocumentEndHandler = unsafe extern "C" fn(*mut DocumentEnd, *mut c_void) -> RewriterDirective;
 
 struct ExternHandler<F> {
     func: Option<F>,
@@ -52,6 +53,7 @@ pub struct ExternDocumentContentHandlers {
     doctype: ExternHandler<DoctypeHandler>,
     comments: ExternHandler<CommentsHandler>,
     text: ExternHandler<TextHandler>,
+    end: ExternHandler<DocumentEndHandler>,
 }
 
 impl ExternDocumentContentHandlers {
@@ -61,6 +63,7 @@ impl ExternDocumentContentHandlers {
         add_handler!(handlers, self.doctype);
         add_handler!(handlers, self.comments);
         add_handler!(handlers, self.text);
+        add_handler!(handlers, self.end);
 
         handlers
     }
@@ -126,13 +129,16 @@ pub extern "C" fn lol_html_rewriter_builder_add_document_content_handlers(
     comments_handler_user_data: *mut c_void,
     text_handler: Option<TextHandler>,
     text_handler_user_data: *mut c_void,
+    document_end_handler: Option<DocumentEndHandler>,
+    document_end_handler_user_data: *mut c_void,
 ) {
     let builder = to_ref_mut!(builder);
 
     let handlers = ExternDocumentContentHandlers {
         doctype: ExternHandler::new(doctype_handler, doctype_handler_user_data),
         comments: ExternHandler::new(comments_handler, comments_handler_user_data),
         text: ExternHandler::new(text_handler, text_handler_user_data),
+        end: ExternHandler::new(document_end_handler, document_end_handler_user_data),
     };
 
     builder.document_content_handlers.push(handlers);

diff --git a/c-api/tests/src/test.c b/c-api/tests/src/test.c
@@ -14,6 +14,7 @@ int run_tests() {
     subtest("Comment API", test_comment_api);
     subtest("Text chunk API", test_text_chunk_api);
     subtest("Element API", element_api_test);
+    subtest("Document end API", document_end_api_test);
     subtest("Memory limiting", test_memory_limiting);
     return done_testing();
 }
diff --git a/c-api/tests/src/test_comment_api.c b/c-api/tests/src/test_comment_api.c
@@ -44,6 +44,8 @@ static void test_get_set_comment_text(void *user_data) {
         &get_set_comment_text,
         user_data,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -84,6 +86,8 @@ static void test_insert_before_and_after_comment(void *user_data) {
         &insert_before_and_after_comment,
         user_data,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -131,6 +135,8 @@ static void test_get_set_user_data(void *user_data) {
         &get_set_user_data,
         user_data,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -213,6 +219,8 @@ static void test_insert_after_comment(lol_html_selector_t *selector, void *user_
         &insert_after_comment,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -258,6 +266,8 @@ static void test_remove_comment(void *user_data) {
         &remove_comment,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -287,6 +297,8 @@ static void test_stop(void *user_data) {
         &stop_rewriting,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 

diff --git a/c-api/tests/src/test_doctype_api.c b/c-api/tests/src/test_doctype_api.c
@@ -45,6 +45,8 @@ static void test_get_doctype_fields(void *user_data) {
         NULL,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -94,6 +96,8 @@ static void test_get_user_data(void *user_data) {
         NULL,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -128,6 +132,8 @@ static void test_stop(void *user_data) {
         NULL,
         NULL,
         NULL,
+        NULL,
+        NULL,
         NULL
     );
 

diff --git a/c-api/tests/src/test_document_end_api.c b/c-api/tests/src/test_document_end_api.c
@@ -0,0 +1,107 @@
+#include "../../include/lol_html.h"
+#include "deps/picotest/picotest.h"
+#include "tests.h"
+#include "test_util.h"
+
+static int EXPECTED_USER_DATA = 43;
+
+//-------------------------------------------------------------------------
+EXPECT_OUTPUT(
+    append_to_empty_doc_output_sink,
+    "<!--appended text-->hello &amp; world",
+    &EXPECTED_USER_DATA,
+    sizeof(EXPECTED_USER_DATA)
+);
+
+static lol_html_rewriter_directive_t append_to_empty_doc(
+    lol_html_doc_end_t *doc_end,
+    void *user_data
+) {
+    note("Append at at the end of an empty document");
+    ok(*(int*)user_data == EXPECTED_USER_DATA);
+
+    const char *append_html = "<!--appended text-->";
+    ok(!lol_html_doc_end_append(doc_end, append_html, strlen(append_html), true));
+
+    const char *append_text = "hello & world";
+    ok(!lol_html_doc_end_append(doc_end, append_text, strlen(append_text), false));
+
+    return LOL_HTML_CONTINUE;
+}
+
+static void test_append_to_empty_doc(void *user_data) {
+    lol_html_rewriter_builder_t *builder = lol_html_rewriter_builder_new();
+
+    lol_html_rewriter_builder_add_document_content_handlers(
+        builder,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        append_to_empty_doc,
+        user_data
+    );
+
+    run_rewriter(
+        builder,
+        "",
+        append_to_empty_doc_output_sink,
+        user_data
+    );
+}
+
+//-------------------------------------------------------------------------
+EXPECT_OUTPUT(
+    append_at_end_output_sink,
+    "<html><div>Hello</div></html><!--appended text-->hello &amp; world",
+    &EXPECTED_USER_DATA,
+    sizeof(EXPECTED_USER_DATA)
+);
+
+static lol_html_rewriter_directive_t append_at_end(
+    lol_html_doc_end_t *doc_end,
+    void *user_data
+) {
+    note("Append at at the end");
+    ok(*(int*)user_data == EXPECTED_USER_DATA);
+
+    const char *append_html = "<!--appended text-->";
+    ok(!lol_html_doc_end_append(doc_end, append_html, strlen(append_html), true));
+
+    const char *append_text = "hello & world";
+    ok(!lol_html_doc_end_append(doc_end, append_text, strlen(append_text), false));
+
+    return LOL_HTML_CONTINUE;
+}
+
+static void test_append_at_end(void *user_data) {
+    lol_html_rewriter_builder_t *builder = lol_html_rewriter_builder_new();
+
+    lol_html_rewriter_builder_add_document_content_handlers(
+        builder,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        NULL,
+        append_at_end,
+        user_data
+    );
+
+    run_rewriter(
+        builder,
+        "<html><div>Hello</div></html>",
+        append_at_end_output_sink,
+        user_data
+    );
+}
+
+void document_end_api_test() {
+    int user_data = 43;
+
+    test_append_to_empty_doc(&user_data);
+    test_append_at_end(&user_data);
+}
diff --git a/c-api/tests/src/test_text_chunk_api.c b/c-api/tests/src/test_text_chunk_api.c
@@ -52,6 +52,8 @@ static void test_insert_before_and_after_text_chunk(void *user_data) {
         NULL,
         NULL,
         &insert_before_and_after_text_chunk,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -95,7 +97,9 @@ static void test_modify_user_data(void *user_data) {
         NULL,
         NULL,
         &modify_user_data,
-        user_data
+        user_data,
+        NULL,
+        NULL
     );
 
     run_rewriter(builder, "Hey 42", modify_user_data_output_sink, user_data);
@@ -177,6 +181,8 @@ static void test_insert_after_chunk(void *user_data) {
         NULL,
         NULL,
         &insert_after_chunk,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -216,6 +222,8 @@ static void test_remove_chunk(void *user_data) {
         NULL,
         NULL,
         &remove_chunk,
+        NULL,
+        NULL,
         NULL
     );
 
@@ -263,6 +271,8 @@ static void test_stop(void *user_data) {
         NULL,
         NULL,
         &stop_rewriting,
+        NULL,
+        NULL,
         NULL
     );