#!/usr/bin/env escript

%%% You need an OTP with fixes for sendmsg types
%%% see https://github.com/erlang/otp/pull/2400

-module(sendsock).
-mode(compile).

-export([main/1]).


main(["proxy"|Opts]) when length(Opts) =< 2 ->
  {Port, Path} = case Opts of
    [] ->
      {13456, <<"test.sock">>};
    [[D|_] = PortStr] when $0 =< D, D =< $9 ->
      {list_to_integer(PortStr), <<"test.sock">>};
    [PathStr] ->
      {13456, list_to_binary(PathStr)};
    [PortStr, PathStr] ->
      {list_to_integer(PortStr), list_to_binary(PathStr)}
  end,
  proxy(Port, Path);

main(["worker"|Opts]) when length(Opts) =< 1 ->
  Path = case Opts of
    [] -> <<"test.sock">>;
    [PathStr] -> list_to_binary(PathStr)
  end,
  echo(Path);

main(_) ->
  Self = escript:script_name(),
  io:format([
    "Demo: pass an open TCP socket to another process over UNIX socket\n",
    "Usage:\n",
    "    * ", Self, " proxy <TCPPort> <UNIXSock>\n"
    "    * ", Self, " worker <UNIXSock>\n", "\n",
    "Demo:\n",
    "    1. Start a proxy, see the tcp listen socket owned by it\n",
    "    2. Start a worker\n",
    "    3. Connect to a specified port via telnet, see the accepted socket owned by the worker\n",
    "    4. Stop the proxy and see the worker is still acting as echo server\n"
            ]),
  ok.


proxy(Port, Path) ->
  {ok, UListen} = socket:open(local, stream, default),
  _ = file:delete(Path),
  {ok, _} = socket:bind(UListen, #{family => local, path => Path}),
  ok = socket:listen(UListen),
  io:format("Listening UNIX socket ~s for worker connection~n", [Path]),

  {ok, TListen} = gen_tcp:listen(Port, [{active, false}, {reuseaddr, true}]),
  io:format("Listening TCP port ~w for client connections~n", [Port]),

  proxy_loop(TListen, UListen, undefined).

proxy_loop(TListen, UListen, USock) ->
  {ok, S} = gen_tcp:accept(TListen),
  {ok, Peer} = inet:peername(S),
  {ok, FD} = inet:getfd(S),
  io:format("passing socket ~w [fd ~w] (peer ~w)\n", [S, FD, Peer]),
  {ok, USock1} = pass_fd(FD, UListen, USock),
  gen_tcp:close(S),
  proxy_loop(TListen, UListen, USock1).

pass_fd(FD, UListen, undefined) ->
  {ok, USock} = socket:accept(UListen),
  io:format("Worker connected~n"),
  pass_fd(FD, UListen, USock);
pass_fd(FD, UListen, USock) ->
  R = socket:sendmsg(USock, #{iov => [<<"hello">>], ctrl => [#{level => socket, type => rights, data => <<FD:32/native>>}]}),
  case R of
    ok ->
      {ok, USock};
    Err ->
      io:format("sendmsg error: ~p, waiting for a new worker~n", [Err]),
      pass_fd(FD, UListen, undefined)
  end.


echo(Path) ->
  {ok, U} = socket:open(local, stream, default),
  connect_loop(U, Path).

connect_loop(U, Path) ->
  case socket:connect(U, #{family => local, path => Path}) of
    ok ->
      io:format("Connected to ~s~n", [Path]),
      echo_acc_loop(U, Path);
    Err ->
      io:format("Failed to connect to socket ~s: ~p~n", [Path, Err]),
      timer:sleep(5000),
      connect_loop(U, Path)
  end.

echo_acc_loop(U, Path) ->
  case socket:recvmsg(U) of
    {ok, Msg} ->
      handle_msg(Msg),
      echo_acc_loop(U, Path);
    Err ->
      io:format("Failed to recvmsg: ~p~n", [Err]),
      socket:close(U),
      echo(Path)
  end.

handle_msg(#{ctrl := [#{level := socket, type := rights, data := <<FD:32/native, Rest/binary>>}]}) ->
  {ok, S} = gen_tcp:fdopen(FD, [binary, {active, false}]),
  {ok, Peer} = inet:peername(S),
  Extra = [ io_lib:format(" (~w extra data bytes: ~999P)", [byte_size(Rest), Rest, 8]) || byte_size(Rest) > 0],
  io:format("received socket ~w [fd ~w] (peer ~w)~s\n", [S, FD, Peer, Extra]),
  Pid = spawn(fun() -> echo_loop(S) end),
  gen_tcp:controlling_process(S, Pid);
handle_msg(_) ->
  ok.


echo_loop(S) ->
  case gen_tcp:recv(S, 0) of
    {ok, Data} ->
      ok = gen_tcp:send(S, ["> ", Data]),
      echo_loop(S);
    {error, closed} ->
      io:format("socket ~w closed\n", [S])
  end.