Merge branch 'python_add_missing_callbacks' of https://github.com/fhriley/unbound into fhriley-python_add_missing_callbacks

gthess · gthess · commit 4776f749bec4 · 2023-07-21T15:19:12.000+02:00
diff --git a/pythonmod/examples/inplace_callbacks.py b/pythonmod/examples/inplace_callbacks.py
@@ -34,6 +34,9 @@
  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  POSSIBILITY OF SUCH DAMAGE.
 '''
+
+import os
+
 #Try:
 # - dig @localhost nlnetlabs.nl +ednsopt=65002:
 #       This query *could* be answered from cache. If so, unbound will reply
@@ -242,6 +245,36 @@ def inplace_query_callback(qinfo, flags, qstate, addr, zone, region, **kwargs):
     return True
 
 
+def inplace_query_response_callback(qstate, response, **kwargs):
+    """
+    Function that will be registered as an inplace callback function.
+    It will be called after receiving a reply from a backend server.
+
+    :param qstate: module qstate. opt_lists are available here;
+    :param response: struct dns_msg. The reply received from the backend server;
+    :param **kwargs: Dictionary that may contain parameters added in a future
+                     release.
+    """
+    log_dns_msg(
+        "python: incoming reply from {}{}".format(qstate.reply.addr, os.linesep),
+        response.qinfo, response.rep
+    )
+    return True
+
+
+def inplace_edns_back_parsed_call(qstate, **kwargs):
+    """
+    Function that will be registered as an inplace callback function.
+    It will be called after EDNS is parsed on a reply from a backend server..
+
+    :param qstate: module qstate. opt_lists are available here;
+    :param **kwargs: Dictionary that may contain parameters added in a future
+                     release.
+    """
+    log_info("python: edns parsed")
+    return True
+
+
 def init_standard(id, env):
     """
     New version of the init function.
@@ -281,6 +314,16 @@ def init_standard(id, env):
     if not register_inplace_cb_query(inplace_query_callback, env, id):
         return False
 
+    # Register the inplace_edns_back_parsed_call function as an inplace callback
+    # for when a reply is received from a backend server.
+    if not register_inplace_cb_query_response(inplace_query_response_callback, env, id):
+        return False
+
+    # Register the inplace_edns_back_parsed_call function as an inplace callback
+    # for when EDNS is parsed on a reply from a backend server.
+    if not register_inplace_cb_edns_back_parsed_call(inplace_edns_back_parsed_call, env, id):
+        return False
+
     return True
 
 
diff --git a/pythonmod/interface.i b/pythonmod/interface.i
@@ -1687,6 +1687,82 @@ int edns_opt_list_append(struct edns_option** list, uint16_t code, size_t len,
         if (ret) Py_INCREF(py_cb);
         return ret;
     }
+
+    int python_inplace_cb_query_response(struct module_qstate* qstate,
+        struct dns_msg* response, int id, void* python_callback)
+    {
+        int res = 0;
+        PyObject *func = python_callback;
+
+        PyGILState_STATE gstate = PyGILState_Ensure();
+
+        PyObject *py_qstate = SWIG_NewPointerObj((void*) qstate, SWIGTYPE_p_module_qstate, 0);
+        PyObject *py_response = SWIG_NewPointerObj((void*) response, SWIGTYPE_p_dns_msg, 0);
+
+        PyObject *py_args = Py_BuildValue("(OO)", py_qstate, py_response);
+        PyObject *py_kwargs = Py_BuildValue("{}");
+        PyObject *result = PyObject_Call(func, py_args, py_kwargs);
+        if (result) {
+            res = PyInt_AsLong(result);
+        }
+
+        Py_XDECREF(py_qstate);
+        Py_XDECREF(py_response);
+
+        Py_XDECREF(py_args);
+        Py_XDECREF(py_kwargs);
+        Py_XDECREF(result);
+
+        PyGILState_Release(gstate);
+
+        return res;
+    }
+
+    static int register_inplace_cb_query_response(PyObject* py_cb,
+        struct module_env* env, int id)
+    {
+        int ret = inplace_cb_register(python_inplace_cb_query_response,
+            inplace_cb_query_response, (void*) py_cb, env, id);
+        if (ret) Py_INCREF(py_cb);
+        return ret;
+    }
+
+    int python_inplace_cb_edns_back_parsed_call(struct module_qstate* qstate,
+        int id, void* python_callback)
+    {
+        int res = 0;
+        PyObject *func = python_callback;
+
+        PyGILState_STATE gstate = PyGILState_Ensure();
+
+        PyObject *py_qstate = SWIG_NewPointerObj((void*) qstate, SWIGTYPE_p_module_qstate, 0);
+
+        PyObject *py_args = Py_BuildValue("(O)", py_qstate);
+        PyObject *py_kwargs = Py_BuildValue("{}");
+        PyObject *result = PyObject_Call(func, py_args, py_kwargs);
+        if (result) {
+            res = PyInt_AsLong(result);
+        }
+
+        Py_XDECREF(py_qstate);
+
+        Py_XDECREF(py_args);
+        Py_XDECREF(py_kwargs);
+        Py_XDECREF(result);
+
+        PyGILState_Release(gstate);
+
+        return res;
+    }
+
+    static int register_inplace_cb_edns_back_parsed_call(PyObject* py_cb,
+        struct module_env* env, int id)
+    {
+        int ret = inplace_cb_register(python_inplace_cb_edns_back_parsed_call,
+            inplace_cb_edns_back_parsed, (void*) py_cb, env, id);
+        if (ret) Py_INCREF(py_cb);
+        return ret;
+    }
 %}
 /* C declarations */
 int inplace_cb_register(void* cb, enum inplace_cb_list_type type, void* cbarg,
@@ -1703,3 +1779,7 @@ static int register_inplace_cb_reply_servfail(PyObject* py_cb,
     struct module_env* env, int id);
 static int register_inplace_cb_query(PyObject *py_cb,
     struct module_env* env, int id);
+static int register_inplace_cb_query_response(PyObject *py_cb,
+    struct module_env* env, int id);
+static int register_inplace_cb_edns_back_parsed_call(PyObject *py_cb,
+    struct module_env* env, int id);
diff --git a/pythonmod/pythonmod.h b/pythonmod/pythonmod.h
@@ -82,4 +82,12 @@ int python_inplace_cb_query_generic(
         uint8_t* zone, size_t zonelen, struct regional* region, int id,
         void* python_callback);
 
+/** Declared here for fptr_wlist access. The definition is in interface.i. */
+int python_inplace_cb_query_response(struct module_qstate* qstate,
+    struct dns_msg* response, int id, void* python_callback);
+
+/** Declared here for fptr_wlist access. The definition is in interface.i. */
+int python_inplace_cb_edns_back_parsed_call(struct module_qstate* qstate,
+    int id, void* python_callback);
+
 #endif /* PYTHONMOD_H */
diff --git a/util/fptr_wlist.c b/util/fptr_wlist.c
@@ -659,6 +659,10 @@ int fptr_whitelist_inplace_cb_edns_back_parsed(
 #else
 	(void)fptr;
 #endif
+#ifdef WITH_PYTHONMODULE
+    if(fptr == &python_inplace_cb_edns_back_parsed_call)
+        return 1;
+#endif
 #ifdef WITH_DYNLIBMODULE
     if(fptr == &dynlib_inplace_cb_edns_back_parsed)
             return 1;
@@ -675,6 +679,10 @@ int fptr_whitelist_inplace_cb_query_response(
 #else
 	(void)fptr;
 #endif
+#ifdef WITH_PYTHONMODULE
+    if(fptr == &python_inplace_cb_query_response)
+        return 1;
+#endif
 #ifdef WITH_DYNLIBMODULE
     if(fptr == &dynlib_inplace_cb_query_response)
             return 1;
