rehandshake command works in udp mode as well.

4 files changed, 41 insertions, 14 deletions

diff --git a/src/common.c b/src/common.c
index f35f16f9f7..65e9c721b7 100644
--- a/src/common.c
+++ b/src/common.c
@@ -1029,3 +1029,19 @@ print_list (const char *priorities, int verbose)
     }
 }
 
+int check_command(gnutls_session_t session, const char* str)
+{
+int len = strlen(str);
+
+  if (len > 2 && str[0] == str[1] && str[0] == '*')
+    {
+      if (strncmp(str, "**REHANDSHAKE**",
+          sizeof ("**REHANDSHAKE**") - 1) == 0)
+        {
+          fprintf (stderr, "*** Sending rehandshake request\n");
+                   gnutls_rehandshake (session);
+          return 1;
+        }
+    }
+  return 0;
+}
diff --git a/src/common.h b/src/common.h
index da42835556..0495bdbf57 100644
--- a/src/common.h
+++ b/src/common.h
@@ -58,3 +58,4 @@ int cert_verify (gnutls_session_t session, const char* hostname);
 
 const char *raw_to_string (const unsigned char *raw, size_t raw_size);
 void pkcs11_common (void);
+int check_command(gnutls_session_t session, const char* str);
diff --git a/src/serv.c b/src/serv.c
index e3a62e17cf..07ca20004e 100644
--- a/src/serv.c
+++ b/src/serv.c
@@ -839,15 +839,8 @@ get_response (gnutls_session_t session, char *request,
     {
       strip (request);
       fprintf (stderr, "received: %s\n", request);
-      if (request[0] == request[1] && request[0] == '*')
+      if (check_command(session, *response))
         {
-          if (strncmp
-              (request, "**REHANDSHAKE**",
-               sizeof ("**REHANDSHAKE**") - 1) == 0)
-            {
-              fprintf (stderr, "*** Sending rehandshake request\n");
-              gnutls_rehandshake (session);
-            }
           *response = NULL;
           *response_length = 0;
           return;
diff --git a/src/udp-serv.c b/src/udp-serv.c
index c4a376b499..478435fa4b 100644
--- a/src/udp-serv.c
+++ b/src/udp-serv.c
@@ -34,6 +34,7 @@
 #include <string.h>
 #include <unistd.h>
 #include <errno.h>
+#include <common.h>
 #include "udp-serv.h"
 #include "list.h"
 
@@ -147,6 +148,18 @@ void udp_server(const char* name, int port, int mtu)
               ret = gnutls_record_recv_seq(session, buffer, MAX_BUFFER, sequence);
             } while(ret == GNUTLS_E_AGAIN || ret == GNUTLS_E_INTERRUPTED);
 
+            if (ret == GNUTLS_E_REHANDSHAKE)
+              {
+                fprintf (stderr, "*** Received hello message\n");
+                do
+                  {
+                    ret = gnutls_handshake (session);
+                  }
+                while (ret == GNUTLS_E_INTERRUPTED ||
+                       ret == GNUTLS_E_AGAIN);
+                
+                if (ret == 0) continue;
+              }
             if (ret < 0)
               {
                 fprintf(stderr, "Error in recv(): %s\n", gnutls_strerror(ret));
@@ -157,16 +170,20 @@ void udp_server(const char* name, int port, int mtu)
                 printf("EOF\n\n");
                 break;
               }
+              
             buffer[ret] = 0;
             printf("received[%.2x%.2x%.2x%.2x%.2x%.2x%.2x%.2x]: %s\n", sequence[0], sequence[1], sequence[2],
-              sequence[3], sequence[4], sequence[5], sequence[6], sequence[7], buffer);
+                   sequence[3], sequence[4], sequence[5], sequence[6], sequence[7], buffer);
 
-            /* reply back */
-            ret = gnutls_record_send(session, buffer, ret);
-            if (ret < 0)
+            if (check_command(session, buffer) == 0)
               {
-                fprintf(stderr, "Error in send(): %s\n", gnutls_strerror(ret));
-                break;
+                /* reply back */
+                ret = gnutls_record_send(session, buffer, ret);
+                if (ret < 0)
+                  {
+                    fprintf(stderr, "Error in send(): %s\n", gnutls_strerror(ret));
+                    break;
+                  }
               }
           }
       }