Add post-listen callback

6 files changed, 112 insertions, 10 deletions

diff --git a/doc/src/guide/listeners.asciidoc b/doc/src/guide/listeners.asciidoc
index 0aee43c..779497c 100644
--- a/doc/src/guide/listeners.asciidoc
+++ b/doc/src/guide/listeners.asciidoc
@@ -177,6 +177,44 @@ file must not exist: Ranch must be able to create it.
     ]}, echo_protocol, []
 ).
 
+=== Performing additional setup steps on a listening socket
+
+If it is necessary to perform additional setup steps on the listening
+socket, it is possible to specify a function with the transport option
+`post_listen_callback`. This function will be called after the listening
+socket has been created but before accepting connections on it,
+with the socket as the single argument.
+
+The function must return either the atom `ok`, after which the listener
+will start accepting connections on the socket, or a tuple
+`{error, Reason}` which will cause the listener to fail starting with
+`Reason`.
+
+.Setting permissions on a UNIX Domain socket file
+
+[source,erlang]
+----
+PostListenCb = fun (Sock) ->
+    case ranch_tcp:sockname(Sock) of
+        {ok, {local, SockFile}} ->
+            file:change_mode(SockFile, 8#777);
+	{ok, _} ->
+	    ok;
+	Error = {error, _} ->
+            Error
+    end
+end,
+
+{ok, _} = ranch:start_listener(tcp_echo,
+    ranch_tcp, #{
+        socket_opts => [
+            {ip, {local, "/tmp/ranch_echo.sock"}},
+            {port, 0}],
+        post_listen_callback => PostListenCb},
+    echo_protocol, []
+).
+----
+
 === Accepting connections on an existing socket
 
 If you want to accept connections on an existing socket, you can write
diff --git a/doc/src/manual/ranch.asciidoc b/doc/src/manual/ranch.asciidoc
index ed57236..6dffb7f 100644
--- a/doc/src/manual/ranch.asciidoc
+++ b/doc/src/manual/ranch.asciidoc
@@ -93,14 +93,15 @@ Unique name used to refer to a listener.
 [source,erlang]
 ----
 transport_opts(SocketOpts) = #{
-    connection_type   => worker | supervisor,
-    handshake_timeout => timeout(),
-    max_connections   => max_conns(),
-    logger            => module(),
-    num_acceptors     => pos_integer(),
-    num_conns_sups    => pos_integer(),
-    shutdown          => timeout() | brutal_kill,
-    socket_opts       => SocketOpts
+    connection_type      => worker | supervisor,
+    handshake_timeout    => timeout(),
+    max_connections      => max_conns(),
+    logger               => module(),
+    num_acceptors        => pos_integer(),
+    num_conns_sups       => pos_integer(),
+    post_listen_callback => fun((term()) -> ok | {error, term()}),
+    shutdown             => timeout() | brutal_kill,
+    socket_opts          => SocketOpts
 }
 ----
 
@@ -137,6 +138,16 @@ num_conns_sups - see below::
 Number of processes that supervise connection processes.
 If not specified, defaults to be equal to `num_acceptors`.
 
+post_listen_callback (fun(_ListenSock) -> ok end)::
+
+A function which will be called after a listen socket has been successfully
+created, with the socket as argument. It can be used to perform any
+necessary setup steps on the socket.
++
+If the callback function returns `ok`, the listener will start accepting
+connections on the socket. If it returns `{error, Reason}`, the listener
+will fail to start.
+
 shutdown (5000)::
 
 Maximum allowed time for children to stop on listener shutdown.
diff --git a/doc/src/manual/ranch.set_transport_options.asciidoc b/doc/src/manual/ranch.set_transport_options.asciidoc
index 8c2eacb..125d037 100644
--- a/doc/src/manual/ranch.set_transport_options.asciidoc
+++ b/doc/src/manual/ranch.set_transport_options.asciidoc
@@ -29,6 +29,7 @@ Changes to the following options will take effect...
 * only after the listener has been suspended and resumed:
 ** `num_acceptors`
 ** `num_listen_sockets`
+** `post_listen_callback`
 ** `socket_opts`
 * only when the entire listener is restarted:
 ** `connection_type`
diff --git a/src/ranch.erl b/src/ranch.erl
index a83402d..f36c145 100644
--- a/src/ranch.erl
+++ b/src/ranch.erl
@@ -58,11 +58,12 @@
 -type transport_opts(SocketOpts) :: #{
 	connection_type => worker | supervisor,
 	handshake_timeout => timeout(),
-	max_connections => max_conns(),
 	logger => module(),
+	max_connections => max_conns(),
 	num_acceptors => pos_integer(),
 	num_conns_sups => pos_integer(),
 	num_listen_sockets => pos_integer(),
+	post_listen_callback => fun((term()) -> ok | {error, term()}),
 	shutdown => timeout() | brutal_kill,
 	socket_opts => SocketOpts
 }.
@@ -131,6 +132,8 @@ validate_transport_opt(num_conns_sups, Value, _) ->
 validate_transport_opt(num_listen_sockets, Value, Opts) ->
 	is_integer(Value) andalso Value > 0
 		andalso Value =< maps:get(num_acceptors, Opts, 10);
+validate_transport_opt(post_listen_callback, Value, _) ->
+	is_function(Value, 1);
 validate_transport_opt(shutdown, brutal_kill, _) ->
 	true;
 validate_transport_opt(shutdown, infinity, _) ->
diff --git a/src/ranch_acceptors_sup.erl b/src/ranch_acceptors_sup.erl
index 801fec5..76155b8 100644
--- a/src/ranch_acceptors_sup.erl
+++ b/src/ranch_acceptors_sup.erl
@@ -77,7 +77,13 @@ start_listen_sockets(Ref, NumListenSockets, Transport, TransOpts0, Logger) when
 start_listen_socket(Ref, Transport, TransOpts, Logger) ->
 	case Transport:listen(TransOpts) of
 		{ok, Socket} ->
-			Socket;
+			PostListenCb = maps:get(post_listen_callback, TransOpts, fun (_) -> ok end),
+			case PostListenCb(Socket) of
+				ok ->
+					Socket;
+				{error, Reason} ->
+					listen_error(Ref, Transport, TransOpts, Reason, Logger)
+			end;
 		{error, Reason} ->
 			listen_error(Ref, Transport, TransOpts, Reason, Logger)
 	end.
diff --git a/test/acceptor_SUITE.erl b/test/acceptor_SUITE.erl
index 6ca9bd3..1028614 100644
--- a/test/acceptor_SUITE.erl
+++ b/test/acceptor_SUITE.erl
@@ -77,6 +77,8 @@ groups() ->
 		misc_info_embedded,
 		misc_metrics,
 		misc_opts_logger,
+		misc_post_listen_callback,
+		misc_post_listen_callback_error,
 		misc_set_transport_options,
 		misc_wait_for_connections,
 		misc_multiple_ip_local_socket_opts
@@ -340,6 +342,47 @@ misc_opts_logger(_) ->
 warning(Format, Args) ->
 	misc_opts_logger ! {warning, Format, Args}.
 
+misc_post_listen_callback(_) ->
+	doc("Ensure that the post-listen callback works."),
+	Name = name(),
+	Self = self(),
+	Ref = make_ref(),
+	PostListenCb = fun (LSock) ->
+		ok = ranch_tcp:setopts(LSock, [{send_timeout, 1000}]),
+		Self ! {post_listen, Ref, LSock},
+		ok
+	end,
+	{ok, _} = ranch:start_listener(Name,
+		ranch_tcp, #{post_listen_callback => PostListenCb,
+			socket_opts => [{send_timeout, infinity}]},
+		echo_protocol, []),
+	receive
+		{post_listen, Ref, LSock} ->
+			{ok, [{send_timeout, 1000}]} = ranch_tcp:getopts(LSock, [send_timeout]),
+			ok
+	after 1000 ->
+		error(timeout)
+	end,
+	Port = ranch:get_port(Name),
+	{ok, S} = gen_tcp:connect("localhost", Port, [binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(S, <<"Test">>),
+	{ok, <<"Test">>} = gen_tcp:recv(S, 4, 1000),
+	ok = ranch:stop_listener(Name),
+	{error, closed} = gen_tcp:recv(S, 0, 1000),
+	%% Make sure the listener stopped.
+	{'EXIT', _} = begin catch ranch:get_port(Name) end,
+	ok.
+
+misc_post_listen_callback_error(_) ->
+	doc("Ensure that starting a listener fails when the post-listen callback returns an error."),
+	Name = name(),
+	PostListenCb = fun (_) -> {error, test} end,
+	{error, _} = ranch:start_listener(Name,
+		ranch_tcp, #{post_listen_callback => PostListenCb},
+		echo_protocol, []),
+	{'EXIT', _} = begin catch ranch:get_port(Name) end,
+	ok.
+
 misc_repeated_remove(_) ->
 	doc("Ensure repeated removal of connection does not crash the connection supervisor."),
 	Name = name(),