Inhibit system shutdown, sleep and idle when formatting and erasing devices

Formatting and ATA secure erase are both potentially long-running
tasks that can easily take hours to complete. Especially for ATA
secure erase, losing power while the command is pending can be very
very bad.

Here's udisks taking an inhibitor lock:

 # gdbus call --system --dest org.freedesktop.login1 --object-path /org/freedesktop/login1 --method org.freedesktop.login1.Manager.ListInhibitors
 ([('shutdown:sleep:idle', 'Disk Manager', 'Formatting Device', 'block', uint32 0, uint32 14172)],)

On non-systemd systems, this does nothing.

Signed-off-by: David Zeuthen <zeuthen@gmail.com>

6 files changed, 148 insertions, 0 deletions

diff --git a/doc/udisks2-sections.txt b/doc/udisks2-sections.txt
index 9871c1f..5c4f1b2 100644
--- a/doc/udisks2-sections.txt
+++ b/doc/udisks2-sections.txt
@@ -341,6 +341,9 @@ udisks_daemon_util_on_other_seat
 udisks_daemon_util_dup_object
 udisks_daemon_util_escape
 udisks_daemon_util_escape_and_quote
+UDisksInhibitCookie
+udisks_daemon_util_inhibit_system_sync
+udisks_daemon_util_uninhibit_system_sync
 </SECTION>
 
 <SECTION>
diff --git a/src/udisksdaemontypes.h b/src/udisksdaemontypes.h
index a00e343..58c09b8 100644
--- a/src/udisksdaemontypes.h
+++ b/src/udisksdaemontypes.h
@@ -103,6 +103,9 @@ typedef struct _UDisksLinuxPartition UDisksLinuxPartition;
 struct _UDisksLinuxPartitionTable;
 typedef struct _UDisksLinuxPartitionTable UDisksLinuxPartitionTable;
 
+struct UDisksInhibitCookie;
+typedef struct UDisksInhibitCookie UDisksInhibitCookie;
+
 /**
  * UDisksThreadedJobFunc:
  * @job: A #UDisksThreadedJob.
diff --git a/src/udisksdaemonutil.c b/src/udisksdaemonutil.c
index 06982d1..23bc3ed 100644
--- a/src/udisksdaemonutil.c
+++ b/src/udisksdaemonutil.c
@@ -21,6 +21,7 @@
 #include "config.h"
 #include <glib/gi18n-lib.h>
 #include <glib/gstdio.h>
+#include <gio/gunixfdlist.h>
 
 #include <stdio.h>
 
@@ -1238,3 +1239,133 @@ udisks_daemon_util_file_set_contents (const gchar  *filename,
   g_free (tmpl);
   return ret;
 }
+
+/* ---------------------------------------------------------------------------------------------------- */
+
+/**
+ * UDisksInhibitCookie:
+ *
+ * Opaque data structure used in udisks_daemon_util_inhibit_system_sync() and
+ * udisks_daemon_util_uninhibit_system_sync().
+ */
+struct UDisksInhibitCookie
+{
+  /*< private >*/
+  gint32 magic;
+#ifdef HAVE_LIBSYSTEMD_LOGIN
+  gint fd;
+#endif
+};
+
+/**
+ * udisks_daemon_util_inhibit_system_sync:
+ * @reason: A human readable explanation of why the system is being inhibited.
+ *
+ * Tries to inhibit the system.
+ *
+ * Right now only
+ * <ulink url="http://www.freedesktop.org/wiki/Software/systemd/inhibit">systemd</ulink>
+ * inhibitors are supported but other inhibitors can be added in the future.
+ *
+ * Returns: A cookie that can be used with udisks_daemon_util_uninhibit_system_sync().
+ */
+UDisksInhibitCookie *
+udisks_daemon_util_inhibit_system_sync (const gchar  *reason)
+{
+  g_return_val_if_fail (reason != NULL, NULL);
+
+#ifdef HAVE_LIBSYSTEMD_LOGIN
+  UDisksInhibitCookie *ret;
+  GDBusConnection *connection = NULL;
+  GVariant *value = NULL;
+  GUnixFDList *fd_list = NULL;
+  gint32 index = -1;
+  GError *error = NULL;
+
+  connection = g_bus_get_sync (G_BUS_TYPE_SYSTEM, NULL, &error);
+  if (connection == NULL)
+    {
+      udisks_error ("Error getting system bus: %s (%s, %d)",
+                    error->message, g_quark_to_string (error->domain), error->code);
+      g_clear_error (&error);
+      goto out;
+    }
+
+  value = g_dbus_connection_call_with_unix_fd_list_sync (connection,
+                                                         "org.freedesktop.login1",
+                                                         "/org/freedesktop/login1",
+                                                         "org.freedesktop.login1.Manager",
+                                                         "Inhibit",
+                                                         g_variant_new ("(ssss)",
+                                                                        "sleep:shutdown:idle", /* what */
+                                                                        "Disk Manager",        /* who */
+                                                                        reason,                /* why */
+                                                                        "block"),              /* mode */
+                                                         G_VARIANT_TYPE ("(h)"),
+                                                         G_DBUS_CALL_FLAGS_NONE,
+                                                         -1,       /* default timeout */
+                                                         NULL,     /* fd_list */
+                                                         &fd_list, /* out_fd_list */
+                                                         NULL, /* GCancellable */
+                                                         &error);
+  if (value == NULL)
+    {
+      udisks_error ("Error inhibiting: %s (%s, %d)",
+                    error->message, g_quark_to_string (error->domain), error->code);
+      g_clear_error (&error);
+      goto out;
+    }
+
+  g_variant_get (value, "(h)", &index);
+  g_assert (index >= 0 && index < g_unix_fd_list_get_length (fd_list));
+
+  ret = g_new0 (UDisksInhibitCookie, 1);
+  ret->magic = 0xdeadbeef;
+  ret->fd = g_unix_fd_list_get (fd_list, index, &error);
+  if (ret->fd == -1)
+    {
+      udisks_error ("Error getting fd: %s (%s, %d)",
+                    error->message, g_quark_to_string (error->domain), error->code);
+      g_clear_error (&error);
+      g_free (ret);
+      ret = NULL;
+      goto out;
+    }
+
+ out:
+  if (value != NULL)
+    g_variant_unref (value);
+  g_clear_object (&fd_list);
+  g_clear_object (&connection);
+  return ret;
+#else
+  /* non-systemd: just return a dummy pointer */
+  return (UDisksInhibitCookie* ) &udisks_daemon_util_inhibit_system_sync;
+#endif
+}
+
+/**
+ * udisks_daemon_util_uninhibit_system_sync:
+ * @cookie: %NULL or a cookie obtained from udisks_daemon_util_inhibit_system_sync().
+ *
+ * Does nothing if @cookie is %NULL, otherwise uninhibits.
+ */
+void
+udisks_daemon_util_uninhibit_system_sync (UDisksInhibitCookie *cookie)
+{
+#ifdef HAVE_LIBSYSTEMD_LOGIN
+  if (cookie != NULL)
+    {
+      g_assert (cookie->magic == 0xdeadbeef);
+      if (close (cookie->fd) != 0)
+        {
+          udisks_error ("Error closing inhbit-fd: %m");
+        }
+      g_free (cookie);
+    }
+#else
+  /* non-systemd: check dummy pointer */
+  g_warn_if_fail (cookie == (UDisksInhibitCookie* ) &udisks_daemon_util_inhibit_system_sync);
+#endif
+}
+
diff --git a/src/udisksdaemonutil.h b/src/udisksdaemonutil.h
index 0fc645c..7c07e36 100644
--- a/src/udisksdaemonutil.h
+++ b/src/udisksdaemonutil.h
@@ -84,6 +84,9 @@ gboolean udisks_daemon_util_file_set_contents (const gchar  *filename,
                                                gint          mode_for_new_file,
                                                GError      **error);
 
+UDisksInhibitCookie *udisks_daemon_util_inhibit_system_sync   (const gchar          *reason);
+void                 udisks_daemon_util_uninhibit_system_sync (UDisksInhibitCookie  *cookie);
+
 
 G_END_DECLS
 
diff --git a/src/udiskslinuxblock.c b/src/udiskslinuxblock.c
index f4b9a5c..159c253 100644
--- a/src/udiskslinuxblock.c
+++ b/src/udiskslinuxblock.c
@@ -1914,6 +1914,7 @@ handle_format (UDisksBlock           *block,
   const gchar *label = NULL;
   gchar *escaped_device = NULL;
   gboolean was_partitioned = FALSE;
+  UDisksInhibitCookie *inhibit_cookie = NULL;
 
   error = NULL;
   object = udisks_daemon_util_dup_object (block, &error);
@@ -2008,6 +2009,8 @@ handle_format (UDisksBlock           *block,
                                                     invocation))
     goto out;
 
+  inhibit_cookie = udisks_daemon_util_inhibit_system_sync (N_("Formatting Device"));
+
   escaped_device = udisks_daemon_util_escape_and_quote (udisks_block_get_device (block));
 
   was_partitioned = (udisks_object_peek_partition_table (object) != NULL);
@@ -2312,6 +2315,7 @@ handle_format (UDisksBlock           *block,
   udisks_block_complete_format (block, invocation);
 
  out:
+  udisks_daemon_util_uninhibit_system_sync (inhibit_cookie);
   g_free (escaped_device);
   g_free (mapped_name);
   g_free (command);
diff --git a/src/udiskslinuxdriveata.c b/src/udiskslinuxdriveata.c
index 52ecace..88f447b 100644
--- a/src/udiskslinuxdriveata.c
+++ b/src/udiskslinuxdriveata.c
@@ -2315,6 +2315,7 @@ handle_security_erase_unit (UDisksDriveAta        *_drive,
   uid_t caller_uid;
   gid_t caller_gid;
   gboolean enhanced = FALSE;
+  UDisksInhibitCookie *inhibit_cookie = NULL;
 
   object = udisks_daemon_util_dup_object (drive, &error);
   if (object == NULL)
@@ -2369,6 +2370,8 @@ handle_security_erase_unit (UDisksDriveAta        *_drive,
                                                     invocation))
     goto out;
 
+  inhibit_cookie = udisks_daemon_util_inhibit_system_sync (N_("Formatting Device"));
+
   if (!udisks_linux_drive_ata_secure_erase_sync (drive, caller_uid, enhanced, &error))
     {
       g_dbus_method_invocation_return_gerror (invocation, error);
@@ -2381,6 +2384,7 @@ handle_security_erase_unit (UDisksDriveAta        *_drive,
   udisks_drive_ata_complete_security_erase_unit (_drive, invocation);
 
  out:
+  udisks_daemon_util_uninhibit_system_sync (inhibit_cookie);
   g_clear_object (&block_object);
   g_clear_object (&object);
   return TRUE; /* returning TRUE means that we handled the method invocation */