From 5f622370d3f5986fa721df016271c064474bd51c Mon Sep 17 00:00:00 2001
From: Peter Lemenkov <lemenkov@gmail.com>
Date: Tue, 7 Jul 2020 00:03:24 +0200
Subject: Unset NOTIFY_SOCKET env if requested

Signed-off-by: Peter Lemenkov <lemenkov@gmail.com>
---
 src/sd_notify.erl       |  5 ++++-
 test/sd_notify_test.erl | 28 ++++++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/src/sd_notify.erl b/src/sd_notify.erl
index 9d0b2e2..72f1324 100644
--- a/src/sd_notify.erl
+++ b/src/sd_notify.erl
@@ -41,7 +41,7 @@ sd_notifyf(UnsetEnv, Format, Data) ->
 sd_pid_notifyf(Pid, UnsetEnv, Format, Data) ->
 	sd_pid_notify_with_fds(Pid, UnsetEnv, lists:flatten(io_lib:format(Format, Data)), []).
 
-sd_pid_notify_with_fds(_Pid, _UnsetEnv, Call, _Fds) ->
+sd_pid_notify_with_fds(_Pid, UnsetEnv, Call, _Fds) ->
 	error_logger:info_msg("systemd: ~p", [Call]),
 	case os:getenv("NOTIFY_SOCKET") of
 		false -> {error, not_configured};
@@ -52,6 +52,9 @@ sd_pid_notify_with_fds(_Pid, _UnsetEnv, Call, _Fds) ->
 				{ok, Socket} ->
 					Result = gen_udp:send(Socket, {local,Path}, 0, Call),
 					gen_udp:close(Socket),
+
+					UnsetEnv == true andalso os:unsetenv("NOTIFY_SOCKET"),
+
 					Result
 			end
 	end.
diff --git a/test/sd_notify_test.erl b/test/sd_notify_test.erl
index 86766bc..21a9f9d 100644
--- a/test/sd_notify_test.erl
+++ b/test/sd_notify_test.erl
@@ -2,8 +2,7 @@
 
 -include_lib("eunit/include/eunit.hrl").
 
-
-sd_notify_test_() ->
+sd_notify_basic_test_() ->
 	{ok, CWD} = file:get_cwd(),
 	FakeNotifyUnixSockName = CWD ++ "/fake-sock-" ++ integer_to_list(erlang:phash2(make_ref())),
 	{ok, FakeNotifyUnixSock} = gen_udp:open(0, [{ifaddr, {local, FakeNotifyUnixSockName}}, {active, false}, list]),
@@ -26,3 +25,28 @@ sd_notify_test_() ->
 		]
 
 	}.
+
+sd_notify_unsetenv_test_() ->
+	{ok, CWD} = file:get_cwd(),
+	FakeNotifyUnixSockName = CWD ++ "/fake-sock-" ++ integer_to_list(erlang:phash2(make_ref())),
+	{ok, FakeNotifyUnixSock} = gen_udp:open(0, [{ifaddr, {local, FakeNotifyUnixSockName}}, {active, false}, list]),
+	os:putenv("NOTIFY_SOCKET", FakeNotifyUnixSockName),
+
+	{setup,
+		fun() -> ok end,
+		fun(_) -> ok = gen_udp:close(FakeNotifyUnixSock), ok = file:delete(FakeNotifyUnixSockName)  end,
+		[
+			{
+				"Try sending message",
+				fun() ->
+					TestMessage = integer_to_list(erlang:phash2(make_ref())),
+					ok = sd_notify:sd_pid_notify_with_fds(0, true, TestMessage, []),
+					{ok, {_Address, _Port, _Packet}} = gen_udp:recv(FakeNotifyUnixSock, length(TestMessage), 1000),
+					Ret = os:getenv("NOTIFY_SOCKET"),
+					?assertEqual(Ret, false)
+				end
+			}
+
+		]
+
+	}.
-- 
cgit v1.2.1