Make sure archive and restore comamnds do not leak files

As reported by Roman Leventov in
percona/postgres#582 we had an issue where we
easily could leak tmeporary files and directories when on failure,
especially when the wrapped restore or archive command fails.

The solution in this commit is to move the cleanup of temporary files
and directories to and atexit() callback plus add a signal handler for
SIGINT and SIGTERM. The signal handler makes sure to kill itself after
running cleanup so the parent process gets the right error.

We have a tiny race condition between creating the temporary directory
and registering the signal handler but that does not seem worth fixing
since it would just leak a directory, not the WAL file, the race
conditions being tiny.

Author: jeltz

src/bin/pg_tde_archive_decrypt.c

@@ -8,6 +8,8 @@
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
 
+#include <signal.h>
+
 #define TMPFS_DIRECTORY "/dev/shm"
 
 static bool
@@ -134,6 +136,27 @@ usage(const char *progname)
 			 "\n"), progname, progname);
 }
 
+static char tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_archiveXXXXXX";
+static char tmppath[MAXPGPATH];
+
+static void
+cleanup_tmpdir(void)
+{
+	if (unlink(tmppath) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove file \"%s\": %m", tmppath);
+	if (rmdir(tmpdir) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
+}
+
+/* Clean up temporary files and dirs on SIGINT/SIGTERM */
+static void
+trapsig(int signum)
+{
+	cleanup_tmpdir();
+	signal(signum, SIG_DFL);
+	kill(getpid(), signum);
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -143,8 +166,6 @@ main(int argc, char *argv[])
 	char	   *command;
 	char	   *sep;
 	char	   *sourcename;
-	char		tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_archiveXXXXXX";
-	char		tmppath[MAXPGPATH];
 	bool		issegment;
 	int			rc;
 
@@ -199,6 +220,10 @@ main(int argc, char *argv[])
 		s = stpcpy(s, "/");
 		stpcpy(s, sourcename);
 
+		signal(SIGTERM, trapsig);
+		signal(SIGINT, trapsig);
+		atexit(cleanup_tmpdir);
+
 		command = replace_percent_placeholders(command,
 											   "ARCHIVE-COMMAND", "fp",
 											   targetname, tmppath);
@@ -228,13 +253,5 @@ main(int argc, char *argv[])
 
 	free(command);
 
-	if (issegment)
-	{
-		if (unlink(tmppath) < 0)
-			pg_log_warning("could not remove file \"%s\": %m", tmppath);
-		if (rmdir(tmpdir) < 0)
-			pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
-	}
-
 	return 0;
 }

src/bin/pg_tde_restore_encrypt.c

@@ -8,6 +8,8 @@
 #include "access/pg_tde_fe_init.h"
 #include "access/pg_tde_xlog_smgr.h"
 
+#include <signal.h>
+
 #define TMPFS_DIRECTORY "/dev/shm"
 
 /*
@@ -129,6 +131,27 @@ usage(const char *progname)
 			 "\n"), progname, progname);
 }
 
+static char	tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_restoreXXXXXX";
+static char	tmppath[MAXPGPATH];
+
+static void
+cleanup_tmpdir(void)
+{
+	if (unlink(tmppath) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove file \"%s\": %m", tmppath);
+	if (rmdir(tmpdir) < 0 && errno != ENOENT)
+		pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
+}
+
+/* Clean up temporary files and dirs on SIGINT/SIGTERM */
+static void
+trapsig(int signum)
+{
+	cleanup_tmpdir();
+	signal(signum, SIG_DFL);
+	kill(getpid(), signum);
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -138,8 +161,6 @@ main(int argc, char *argv[])
 	char	   *command;
 	char	   *sep;
 	char	   *targetname;
-	char		tmpdir[MAXPGPATH] = TMPFS_DIRECTORY "/pg_tde_restoreXXXXXX";
-	char		tmppath[MAXPGPATH];
 	bool		issegment;
 	int			rc;
 
@@ -194,6 +215,10 @@ main(int argc, char *argv[])
 		s = stpcpy(s, "/");
 		stpcpy(s, targetname);
 
+		signal(SIGTERM, trapsig);
+		signal(SIGINT, trapsig);
+		atexit(cleanup_tmpdir);
+
 		command = replace_percent_placeholders(command,
 											   "RESTORE-COMMAND", "fp",
 											   sourcename, tmppath);
@@ -222,14 +247,7 @@ main(int argc, char *argv[])
 	free(command);
 
 	if (issegment)
-	{
 		write_encrypted_segment(targetpath, sourcename, tmppath);
 
-		if (unlink(tmppath) < 0)
-			pg_log_warning("could not remove file \"%s\": %m", tmppath);
-		if (rmdir(tmpdir) < 0)
-			pg_log_warning("could not remove directory \"%s\": %m", tmpdir);
-	}
-
 	return 0;
 }