From 45652d8430a523f1603c19a294f33c150ca2bd1d Mon Sep 17 00:00:00 2001
From: XieHan <XieHan@sogou-inc.com>
Date: Thu, 14 Oct 2021 17:13:10 +0800
Subject: [PATCH] add WFServerBase::get_listen_addr() for starting server on a
 random port

---
 src/server/WFServer.cc | 54 ++++++++++--------------------------------
 src/server/WFServer.h  | 16 +++++++++++--
 2 files changed, 26 insertions(+), 44 deletions(-)

diff --git a/src/server/WFServer.cc b/src/server/WFServer.cc
index 2dc79379..f314f2db 100644
--- a/src/server/WFServer.cc
+++ b/src/server/WFServer.cc
@@ -19,6 +19,7 @@
 
 #include <sys/types.h>
 #include <sys/socket.h>
+#include <errno.h>
 #include <unistd.h>
 #include <stdio.h>
 #include <atomic>
@@ -114,24 +115,24 @@ int WFServerBase::init(const struct sockaddr *bind_addr, socklen_t addrlen,
 
 int WFServerBase::create_listen_fd()
 {
-	int listen_fd  = this->listen_fd;
-
-	if (listen_fd < 0)
+	if (this->listen_fd < 0)
 	{
 		const struct sockaddr *bind_addr;
 		socklen_t addrlen;
 		int reuse = 1;
 
 		this->get_addr(&bind_addr, &addrlen);
-		listen_fd = socket(bind_addr->sa_family, SOCK_STREAM, 0);
-		if (listen_fd >= 0)
+		this->listen_fd = socket(bind_addr->sa_family, SOCK_STREAM, 0);
+		if (this->listen_fd >= 0)
 		{
-			setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR,
+			setsockopt(this->listen_fd, SOL_SOCKET, SO_REUSEADDR,
 					   &reuse, sizeof (int));
 		}
 	}
+	else
+		this->listen_fd = dup(this->listen_fd);
 
-	return listen_fd;
+	return this->listen_fd;
 }
 
 WFConnection *WFServerBase::new_connection(int accept_fd)
@@ -179,6 +180,7 @@ int WFServerBase::start(const struct sockaddr *bind_addr, socklen_t addrlen,
 			SSL_CTX_free(ssl_ctx);
 	}
 
+	this->listen_fd = -1;
 	return -1;
 }
 
@@ -212,54 +214,22 @@ int WFServerBase::start(int family, const char *host, unsigned short port,
 	return ret;
 }
 
-static int __get_addr_bound(int sockfd, struct sockaddr *addr, socklen_t *len)
-{
-	int family;
-	socklen_t i;
-
-	if (getsockname(sockfd, addr, len) < 0)
-		return -1;
-
-	family = addr->sa_family;
-	addr->sa_family = 0;
-	for (i = 0; i < *len; i++)
-	{
-		if (((char *)addr)[i])
-			break;
-	}
-
-	if (i == *len)
-	{
-		errno = EINVAL;
-		return -1;
-	}
-
-	addr->sa_family = family;
-	return 0;
-}
-
 int WFServerBase::serve(int listen_fd,
 						const char *cert_file, const char *key_file)
 {
 	struct sockaddr_storage ss;
 	socklen_t len = sizeof ss;
-	int ret;
 
-	if (__get_addr_bound(listen_fd, (struct sockaddr *)&ss, &len) < 0)
-		return -1;
-
-	listen_fd = dup(listen_fd);
-	if (listen_fd < 0)
+	if (getsockname(listen_fd, (struct sockaddr *)&ss, &len) < 0)
 		return -1;
 
 	this->listen_fd = listen_fd;
-	ret = start((struct sockaddr *)&ss, len, cert_file, key_file);
-	this->listen_fd = -1;
-	return ret;
+	return start((struct sockaddr *)&ss, len, cert_file, key_file);
 }
 
 void WFServerBase::shutdown()
 {
+	this->listen_fd = -1;
 	this->scheduler->unbind(this);
 }
 
diff --git a/src/server/WFServer.h b/src/server/WFServer.h
index 6cebc7a3..7439f8b3 100644
--- a/src/server/WFServer.h
+++ b/src/server/WFServer.h
@@ -22,11 +22,12 @@
 
 #include <sys/types.h>
 #include <sys/socket.h>
-#include <openssl/ssl.h>
+#include <errno.h>
 #include <functional>
 #include <atomic>
 #include <mutex>
 #include <condition_variable>
+#include <openssl/ssl.h>
 #include "WFTaskFactory.h"
 
 struct WFServerParams
@@ -56,8 +57,8 @@ public:
 		conn_count(0)
 	{
 		this->params = *params;
-		this->listen_fd = -1;
 		this->unbind_finish = false;
+		this->listen_fd = -1;
 	}
 
 public:
@@ -144,6 +145,17 @@ public:
 public:
 	size_t get_conn_count() const { return this->conn_count; }
 
+	/* Get the listening address. This is often used after starting
+	 * server on a random port (start() with port == 0). */
+	int get_listen_addr(struct sockaddr *addr, socklen_t *addrlen) const
+	{
+		if (this->listen_fd >= 0)
+			return getsockname(this->listen_fd, addr, addrlen);
+
+		errno = ENOTCONN;
+		return -1;
+	}
+
 protected:
 	/* Override this function to implement server that supports TLS SNI.
 	 * "servername" will be NULL if client does not set a host name.