2 files changed, 160 insertions, 0 deletions
diff --git a/pkg/systemd/notifyproxy/notifyproxy.go b/pkg/systemd/notifyproxy/notifyproxy.go
new file mode 100644
index 000000000..9e6eb4cf0
--- /dev/null
+++ b/pkg/systemd/notifyproxy/notifyproxy.go
@@ -0,0 +1,102 @@
+package notifyproxy
+
+import (
+	"io/ioutil"
+	"net"
+	"os"
+	"strings"
+	"syscall"
+
+	"github.com/coreos/go-systemd/v22/daemon"
+	"github.com/sirupsen/logrus"
+)
+
+// SendMessage sends the specified message to the specified socket.
+// No message is sent if no socketPath is provided and the NOTIFY_SOCKET
+// variable is not set either.
+func SendMessage(socketPath string, message string) error {
+	if socketPath == "" {
+		socketPath, _ = os.LookupEnv("NOTIFY_SOCKET")
+		if socketPath == "" {
+			return nil
+		}
+	}
+	socketAddr := &net.UnixAddr{
+		Name: socketPath,
+		Net:  "unixgram",
+	}
+	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	_, err = conn.Write([]byte(message))
+	return err
+}
+
+// NotifyProxy can be used to proxy notify messages.
+type NotifyProxy struct {
+	connection *net.UnixConn
+	socketPath string
+}
+
+// New creates a NotifyProxy.  The specified temp directory can be left empty.
+func New(tmpDir string) (*NotifyProxy, error) {
+	tempFile, err := ioutil.TempFile(tmpDir, "-podman-notify-proxy.sock")
+	if err != nil {
+		return nil, err
+	}
+	defer tempFile.Close()
+
+	socketPath := tempFile.Name()
+	if err := syscall.Unlink(socketPath); err != nil { // Unlink the socket so we can bind it
+		return nil, err
+	}
+
+	socketAddr := &net.UnixAddr{
+		Name: socketPath,
+		Net:  "unixgram",
+	}
+	conn, err := net.ListenUnixgram(socketAddr.Net, socketAddr)
+	if err != nil {
+		return nil, err
+	}
+
+	return &NotifyProxy{connection: conn, socketPath: socketPath}, nil
+}
+
+// SocketPath returns the path of the socket the proxy is listening on.
+func (p *NotifyProxy) SocketPath() string {
+	return p.socketPath
+}
+
+// close closes the listener and removes the socket.
+func (p *NotifyProxy) close() error {
+	defer os.Remove(p.socketPath)
+	return p.connection.Close()
+}
+
+// WaitAndClose waits until receiving the `READY` notify message and close the
+// listener. Note that the this function must only be executed inside a systemd
+// service which will kill the process after a given timeout.
+func (p *NotifyProxy) WaitAndClose() error {
+	defer func() {
+		if err := p.close(); err != nil {
+			logrus.Errorf("Closing notify proxy: %v", err)
+		}
+	}()
+
+	for {
+		buf := make([]byte, 1024)
+		num, err := p.connection.Read(buf)
+		if err != nil {
+			return err
+		}
+		for _, s := range strings.Split(string(buf[:num]), "\n") {
+			if s == daemon.SdNotifyReady {
+				return nil
+			}
+		}
+	}
+}
diff --git a/pkg/systemd/notifyproxy/notifyproxy_test.go b/pkg/systemd/notifyproxy/notifyproxy_test.go
new file mode 100644
index 000000000..edad95659
--- /dev/null
+++ b/pkg/systemd/notifyproxy/notifyproxy_test.go
@@ -0,0 +1,58 @@
+package notifyproxy
+
+import (
+	"testing"
+	"time"
+
+	"github.com/coreos/go-systemd/v22/daemon"
+	"github.com/stretchr/testify/require"
+)
+
+// Helper function to send the specified message over the socket of the proxy.
+func sendMessage(t *testing.T, proxy *NotifyProxy, message string) {
+	err := SendMessage(proxy.SocketPath(), message)
+	require.NoError(t, err)
+}
+
+func TestNotifyProxy(t *testing.T) {
+	proxy, err := New("")
+	require.NoError(t, err)
+	require.FileExists(t, proxy.SocketPath())
+	require.NoError(t, proxy.close())
+	require.NoFileExists(t, proxy.SocketPath())
+}
+
+func TestWaitAndClose(t *testing.T) {
+	proxy, err := New("")
+	require.NoError(t, err)
+	require.FileExists(t, proxy.SocketPath())
+
+	ch := make(chan error)
+
+	go func() {
+		ch <- proxy.WaitAndClose()
+	}()
+
+	sendMessage(t, proxy, "foo\n")
+	time.Sleep(250 * time.Millisecond)
+	select {
+	case err := <-ch:
+		t.Fatalf("Should stil be waiting but received %v", err)
+	default:
+	}
+
+	sendMessage(t, proxy, daemon.SdNotifyReady+"\nsomething else")
+	done := func() bool {
+		for i := 0; i < 10; i++ {
+			select {
+			case err := <-ch:
+				require.NoError(t, err, "Waiting should succeed")
+				return true
+			default:
+				time.Sleep(time.Duration(i*250) * time.Millisecond)
+			}
+		}
+		return false
+	}()
+	require.True(t, done, "READY MESSAGE SHOULD HAVE ARRIVED")
+}