From 994c8501e86b1eaa1481d256a5e449e7975ca3be Mon Sep 17 00:00:00 2001
From: Matt Mundell <matt@mundell.me>
Date: Tue, 15 Aug 2023 19:25:06 +0200
Subject: [PATCH] Add close attribute to all GMP commands

---
 src/gmp.c      | 48 +++++++++++++++++++++++++++++++++++++++++++++---
 src/gmp_base.h |  1 +
 src/gmpd.c     |  3 +++
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/gmp.c b/src/gmp.c
index c4be583e7..69b9d4305 100644
--- a/src/gmp.c
+++ b/src/gmp.c
@@ -293,6 +293,11 @@ check_public_key (const char *key_str)
 
 /* GMP parser. */
 
+/**
+ * @brief User data for XML parser.
+ */
+gmp_parser_t *global_gmp_parser;
+
 static int
 process_gmp (gmp_parser_t *, const gchar *, gchar **);
 
@@ -314,6 +319,7 @@ gmp_parser_new (int (*write_to_client) (const char*, void*), void* write_to_clie
   gmp_parser->client_writer = write_to_client;
   gmp_parser->client_writer_data = write_to_client_data;
   gmp_parser->read_over = 0;
+  gmp_parser->close_connection = 0;
   gmp_parser->disabled_commands = g_strdupv (disable);
   return gmp_parser;
 }
@@ -330,6 +336,33 @@ gmp_parser_free (gmp_parser_t *gmp_parser)
   g_free (gmp_parser);
 }
 
+/**
+ * @brief Create global GMP parser.
+ *
+ * @param[in]  write_to_client       Function to write to client.
+ * @param[in]  write_to_client_data  Argument to \p write_to_client.
+ * @param[in]  disable               Commands to disable.  Copied, and freed by
+ *                                   gmp_parser_free.
+ */
+static gmp_parser_t *
+global_gmp_parser_new (int (*write_to_client) (const char*, void*),
+                       void* write_to_client_data,
+                       gchar **disable)
+{
+  global_gmp_parser = gmp_parser_new (write_to_client, write_to_client_data, disable);
+  return global_gmp_parser;
+}
+
+/**
+ * @brief Free global GMP parser.
+ */
+static void
+global_gmp_parser_free ()
+{
+  gmp_parser_free (global_gmp_parser);
+  global_gmp_parser = NULL;
+}
+
 /**
  * @brief Check if command has been disabled.
  *
@@ -4629,6 +4662,7 @@ gmp_xml_handle_start_element (/* unused */ GMarkupParseContext* context,
                               gpointer user_data,
                               GError **error)
 {
+  const gchar* attr;
   gmp_parser_t *gmp_parser = (gmp_parser_t*) user_data;
   int (*write_to_client) (const char *, void*)
     = (int (*) (const char *, void*)) gmp_parser->client_writer;
@@ -4636,6 +4670,10 @@ gmp_xml_handle_start_element (/* unused */ GMarkupParseContext* context,
 
   g_debug ("   XML  start: %s (%i)", element_name, client_state);
 
+  if (find_attribute (attribute_names, attribute_values,
+                      "close", &attr))
+    gmp_parser->close_connection = strcmp (attr, "0");
+
   if (gmp_parser->read_over)
     gmp_parser->read_over++;
   else switch (client_state)
@@ -26719,9 +26757,8 @@ init_gmp_process (const db_conn_info_t *database,
   xml_context = g_markup_parse_context_new
                  (&xml_parser,
                   0,
-                  gmp_parser_new (write_to_client, write_to_client_data,
-                                  disable),
-                  (GDestroyNotify) gmp_parser_free);
+                  global_gmp_parser_new (write_to_client, write_to_client_data, disable),
+                  (GDestroyNotify) global_gmp_parser_free);
 }
 
 /**
@@ -26739,6 +26776,7 @@ init_gmp_process (const db_conn_info_t *database,
  * \endif
  *
  * @return 0 success,
+ *         1 success and client requested connection close,
  *         -1 error,
  *         -4 XML syntax error.
  */
@@ -26788,6 +26826,10 @@ process_gmp_client_input ()
       return err;
     }
   from_client_end = from_client_start = 0;
+
+  if (global_gmp_parser && global_gmp_parser->close_connection)
+    return 1;
+
   return 0;
 }
 
diff --git a/src/gmp_base.h b/src/gmp_base.h
index f9d25c965..fb3fceaec 100644
--- a/src/gmp_base.h
+++ b/src/gmp_base.h
@@ -28,6 +28,7 @@ typedef struct
 {
   int (*client_writer) (const char *, void *); ///< Writes to the client.
   void *client_writer_data;                    ///< Argument to client_writer.
+  int close_connection;      ///< Client requested connection by closed after.
   int importing;             ///< Whether the current op is importing.
   int read_over;             ///< Read over any child elements.
   int parent_state;          ///< Parent state when reading over.
diff --git a/src/gmpd.c b/src/gmpd.c
index 056ee9bb4..39954e23e 100644
--- a/src/gmpd.c
+++ b/src/gmpd.c
@@ -597,6 +597,9 @@ serve_gmp (gvm_connection_t *client_connection, const db_conn_info_t *database,
           if (ret == 0)
             /* Processed all input. */
             ;
+          else if (ret == 1)
+            /* Processed all input.  Client want connection closed after output. */
+            close_connection = 1;
           else if (ret == -1 || ret == -4)
             {
               /* Error.  Write rest of to_client to client, so that the