From 10ee72461254643bd152a7f3f6112edb6f517d4b Mon Sep 17 00:00:00 2001
From: Matt Caswell <matt@openssl.org>
Date: Tue, 12 Sep 2017 16:19:09 +0100
Subject: [PATCH] Enable the cookie callbacks to work even in TLS in the apps

Reviewed-by: Ben Kaduk <kaduk@mit.edu>
(Merged from https://github.com/openssl/openssl/pull/4435)
---
 apps/apps.h     |  2 ++
 apps/s_cb.c     | 29 ++++++++++++++++++-----------
 apps/s_socket.c | 14 +++++++++++++-
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/apps/apps.h b/apps/apps.h
index 6d0d7015f8..601797b1b6 100644
--- a/apps/apps.h
+++ b/apps/apps.h
@@ -46,6 +46,8 @@ extern BIO *bio_out;
 extern BIO *bio_err;
 extern const unsigned char tls13_aes128gcmsha256_id[];
 extern const unsigned char tls13_aes256gcmsha384_id[];
+extern BIO_ADDR *ourpeer;
+
 BIO *dup_bio_in(int format);
 BIO *dup_bio_out(int format);
 BIO *dup_bio_err(int format);
diff --git a/apps/s_cb.c b/apps/s_cb.c
index c7c9ecb170..575fb048c0 100644
--- a/apps/s_cb.c
+++ b/apps/s_cb.c
@@ -686,9 +686,9 @@ int generate_cookie_callback(SSL *ssl, unsigned char *cookie,
                              unsigned int *cookie_len)
 {
     unsigned char *buffer;
-    size_t length;
+    size_t length = 0;
     unsigned short port;
-    BIO_ADDR *peer = NULL;
+    BIO_ADDR *lpeer = NULL, *peer = NULL;
 
     /* Initialize a random secret */
     if (!cookie_initialized) {
@@ -699,17 +699,24 @@ int generate_cookie_callback(SSL *ssl, unsigned char *cookie,
         cookie_initialized = 1;
     }
 
-    peer = BIO_ADDR_new();
-    if (peer == NULL) {
-        BIO_printf(bio_err, "memory full\n");
-        return 0;
+    if (SSL_is_dtls(ssl)) {
+        lpeer = peer = BIO_ADDR_new();
+        if (peer == NULL) {
+            BIO_printf(bio_err, "memory full\n");
+            return 0;
+        }
+
+        /* Read peer information */
+        (void)BIO_dgram_get_peer(SSL_get_rbio(ssl), peer);
+    } else {
+        peer = ourpeer;
     }
 
-    /* Read peer information */
-    (void)BIO_dgram_get_peer(SSL_get_rbio(ssl), peer);
-
     /* Create buffer with peer's address and port */
-    BIO_ADDR_rawaddress(peer, NULL, &length);
+    if (!BIO_ADDR_rawaddress(peer, NULL, &length)) {
+        BIO_printf(bio_err, "Failed getting peer address\n");
+        return 0;
+    }
     OPENSSL_assert(length != 0);
     port = BIO_ADDR_rawport(peer);
     length += sizeof(port);
@@ -723,7 +730,7 @@ int generate_cookie_callback(SSL *ssl, unsigned char *cookie,
          buffer, length, cookie, cookie_len);
 
     OPENSSL_free(buffer);
-    BIO_ADDR_free(peer);
+    BIO_ADDR_free(lpeer);
 
     return 1;
 }
diff --git a/apps/s_socket.c b/apps/s_socket.c
index 74cf8d28e2..a9e46f9949 100644
--- a/apps/s_socket.c
+++ b/apps/s_socket.c
@@ -35,6 +35,9 @@ typedef unsigned int u_int;
 # include <openssl/bio.h>
 # include <openssl/err.h>
 
+/* Keep track of our peer's address for the cookie callback */
+BIO_ADDR *ourpeer = NULL;
+
 /*
  * init_client - helper routine to set up socket communication
  * @sock: pointer to storage of resulting socket.
@@ -212,8 +215,15 @@ int do_server(int *accept_sock, const char *host, const char *port,
         *accept_sock = asock;
     for (;;) {
         if (type == SOCK_STREAM) {
+            BIO_ADDR_free(ourpeer);
+            ourpeer = BIO_ADDR_new();
+            if (ourpeer == NULL) {
+                BIO_closesocket(asock);
+                ERR_print_errors(bio_err);
+                goto end;
+            }
             do {
-                sock = BIO_accept_ex(asock, NULL, 0);
+                sock = BIO_accept_ex(asock, ourpeer, 0);
             } while (sock < 0 && BIO_sock_should_retry(sock));
             if (sock < 0) {
                 ERR_print_errors(bio_err);
@@ -264,6 +274,8 @@ int do_server(int *accept_sock, const char *host, const char *port,
     if (family == AF_UNIX)
         unlink(host);
 # endif
+    BIO_ADDR_free(ourpeer);
+    ourpeer = NULL;
     return ret;
 }