Allow hooks to be called before or after generation/receiving.

4 files changed, 24 insertions, 14 deletions

diff --git a/lib/gnutls_handshake.c b/lib/gnutls_handshake.c
index a90d1a580d..6539458b69 100644
--- a/lib/gnutls_handshake.c
+++ b/lib/gnutls_handshake.c
@@ -1143,17 +1143,15 @@ _gnutls_send_empty_handshake (gnutls_session_t session,
 
 inline
 static int call_hook_func(gnutls_session_t session, gnutls_handshake_description_t type,
-                          unsigned post, unsigned incoming)
+                          int post, unsigned incoming)
 {
-  if (session->internals.h_hook == NULL)
-    return 0;
-  else 
+  if (session->internals.h_hook != NULL)
     {
-      if (session->internals.h_type == type || session->internals.h_type == GNUTLS_HANDSHAKE_ANY)
+      if ((session->internals.h_type == type || session->internals.h_type == GNUTLS_HANDSHAKE_ANY) && 
+      	(session->internals.h_post == post || session->internals.h_post == GNUTLS_HOOK_BOTH))
         return session->internals.h_hook(session, type, post, incoming);
-
-      return 0;
     }
+  return 0;
 }
 
 /* This function sends a handshake message of type 'type' containing the
@@ -1222,7 +1220,7 @@ _gnutls_send_handshake (gnutls_session_t session, mbuffer_st * bufel,
         return ret;
       }
 
-  ret = call_hook_func(session, type, 0, 0);
+  ret = call_hook_func(session, type, GNUTLS_HOOK_PRE, 0);
   if (ret < 0)
     {
       gnutls_assert ();
@@ -1264,7 +1262,7 @@ _gnutls_send_handshake (gnutls_session_t session, mbuffer_st * bufel,
       break;
     }
 
-  ret = call_hook_func(session, type, 1, 0);
+  ret = call_hook_func(session, type, GNUTLS_HOOK_POST, 0);
   if (ret < 0)
     {
       gnutls_assert ();
@@ -1397,7 +1395,7 @@ _gnutls_recv_handshake (gnutls_session_t session,
 
   session->internals.last_handshake_in = hsk.htype;
 
-  ret = call_hook_func(session, hsk.htype, 0, 1);
+  ret = call_hook_func(session, hsk.htype, GNUTLS_HOOK_PRE, 1);
   if (ret < 0)
     {
       gnutls_assert ();
@@ -1475,7 +1473,7 @@ _gnutls_recv_handshake (gnutls_session_t session,
       goto cleanup;
     }
 
-  ret2 = call_hook_func(session, hsk.htype, 1, 1);
+  ret2 = call_hook_func(session, hsk.htype, GNUTLS_HOOK_POST, 1);
   if (ret2 < 0)
     {
       ret = ret2;
diff --git a/lib/gnutls_int.h b/lib/gnutls_int.h
index 258ad5daab..e4901adcaa 100644
--- a/lib/gnutls_int.h
+++ b/lib/gnutls_int.h
@@ -870,6 +870,7 @@ typedef struct
   /* handshake hook function */
   gnutls_handshake_hook_func h_hook;
   unsigned int h_type; /* the hooked type */
+  int h_post; /* whether post-generation/receive */
 
   /* holds the selected certificate and key.
    * use _gnutls_selected_certs_deinit() and _gnutls_selected_certs_set()
diff --git a/lib/gnutls_state.c b/lib/gnutls_state.c
index a5b8ec8c99..a7e8a6b81c 100644
--- a/lib/gnutls_state.c
+++ b/lib/gnutls_state.c
@@ -1410,14 +1410,19 @@ gnutls_handshake_set_random (gnutls_session_t session, const gnutls_datum_t* ran
 
 /**
  * gnutls_handshake_set_hook_function:
- * @session: is a #gnutls_session_t structure.
- * @htype: the %gnutls_handshake_description_t of the message to hook at.
+ * @session: is a #gnutls_session_t structure
+ * @htype: the %gnutls_handshake_description_t of the message to hook at
+ * @post: %GNUTLS_HOOK_* depending on when the hook function should be called
  * @func: is the function to be called
  *
  * This function will set a callback to be called after or before the specified
  * handshake message has been received or generated. This is a
  * generalization of gnutls_handshake_set_post_client_hello_function().
  *
+ * To call the hook function prior to the message being sent/generation use 
+ * %GNUTLS_HOOK_PRE as @post parameter, %GNUTLS_HOOK_POST to call
+ * after, and %GNUTLS_HOOK_BOTH for both cases.
+ *
  * This callback must return 0 on success or a gnutls error code to
  * terminate the handshake.
  *
@@ -1433,9 +1438,11 @@ gnutls_handshake_set_random (gnutls_session_t session, const gnutls_datum_t* ran
 void
 gnutls_handshake_set_hook_function (gnutls_session_t session,
                                     unsigned int htype,
+                                    int post,
 				    gnutls_handshake_hook_func func)
 {
   session->internals.h_hook = func;
   session->internals.h_type = htype;
+  session->internals.h_post = post;
 }
 
diff --git a/lib/includes/gnutls/gnutls.h.in b/lib/includes/gnutls/gnutls.h.in
index 35d5ed3ecf..fba8b2a354 100644
--- a/lib/includes/gnutls/gnutls.h.in
+++ b/lib/includes/gnutls/gnutls.h.in
@@ -1187,10 +1187,14 @@ gnutls_ecc_curve_t gnutls_ecc_curve_get(gnutls_session_t session);
    *
    * Returns: Non zero on error.
    */
+#define GNUTLS_HOOK_POST (1)
+#define GNUTLS_HOOK_PRE (0)
+#define GNUTLS_HOOK_BOTH (-1)
 
-  typedef int (*gnutls_handshake_hook_func) (gnutls_session_t, unsigned int htype, unsigned int post, unsigned int incoming);
+  typedef int (*gnutls_handshake_hook_func) (gnutls_session_t, unsigned int htype, unsigned post, unsigned int incoming);
   void gnutls_handshake_set_hook_function (gnutls_session_t session,
 				           unsigned int htype,
+				           int post,
 				           gnutls_handshake_hook_func func);
 
   typedef int (*gnutls_handshake_post_client_hello_func) (gnutls_session_t);