From 801dc967f83c405d928b0cf1abc55c64f10dac40 Mon Sep 17 00:00:00 2001
From: = <torjus@usit.uio.no>
Date: Sat, 6 Nov 2021 01:29:29 +0100
Subject: [PATCH] Add support for systemd sd_notify

---
 apiary.service |  1 +
 cmd/apiary.go  | 51 ++++++++++++++++++++++++++++++++++++++++++++++++++
 go.mod         |  1 +
 go.sum         |  2 ++
 version.go     |  2 +-
 5 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/apiary.service b/apiary.service
index a9e41b9..395437b 100644
--- a/apiary.service
+++ b/apiary.service
@@ -3,6 +3,7 @@ Description=Apiary SSH Honeypot
 After=postgresql.service
 
 [Service]
+Type=notify
 User=apiary
 ExecStart=/usr/bin/apiary serve
 Restart=on-failure
diff --git a/cmd/apiary.go b/cmd/apiary.go
index 2f8a2d1..baa610c 100644
--- a/cmd/apiary.go
+++ b/cmd/apiary.go
@@ -8,6 +8,7 @@ import (
 	"os/signal"
 	"time"
 
+	"github.com/coreos/go-systemd/daemon"
 	sshlib "github.com/gliderlabs/ssh"
 	"github.com/urfave/cli/v2"
 	"github.uio.no/torjus/apiary"
@@ -165,6 +166,56 @@ func ActionServe(c *cli.Context) error {
 		}
 	}()
 
+	// If run by systemd, enable watchdog and notify ready
+	go func() {
+		notifyCtx, cancel := context.WithCancel(rootCtx)
+		defer cancel()
+
+		_, ok := os.LookupEnv("NOTIFY_SOCKET")
+		if !ok {
+			return
+		}
+		loggers.rootLogger.Info("Systemd notify socket detected. Sending ready and enabling watchdog.")
+		ok, err := daemon.SdNotify(false, daemon.SdNotifyReady)
+		if !ok {
+			loggers.rootLogger.Info("Systemd notify not enabled.")
+			return
+		}
+		if err != nil {
+			loggers.rootLogger.Warnw("Unable to connect to NOTIFY_SOCKET.", "error", err)
+			return
+		}
+		loggers.rootLogger.Debug("Sent READY=1 to NOTIFY_SOCKET.")
+
+		if _, err := daemon.SdNotify(false, "WATCHDOG_USEC=10000000"); err != nil {
+			loggers.rootLogger.Warnw("Unable to connect to NOTIFY_SOCKET to set watchdog timeout.", "error", err)
+			return
+		}
+		loggers.rootLogger.Debug("Sent WATCHDOG_USEC=10000000 to NOTIFY_SOCKET.")
+		if _, err := daemon.SdNotify(false, "WATCHDOG_USEC=10000000"); err != nil {
+			loggers.rootLogger.Warnw("Unable to connect to NOTIFY_SOCKET to set watchdog timeout.", "error", err)
+			return
+		}
+		timeout, err := daemon.SdWatchdogEnabled(false)
+		if err != nil {
+			loggers.rootLogger.Warnw("Unable to connect to NOTIFY_SOCKET to get watchdog timeout.", "error", err)
+			return
+		}
+		ticker := time.NewTicker(timeout / 2)
+		for {
+			healthy := s.IsHealthy()
+			select {
+			case <-ticker.C:
+				if healthy == nil {
+					daemon.SdNotify(false, daemon.SdNotifyWatchdog)
+				}
+			case <-notifyCtx.Done():
+				loggers.rootLogger.Debugw("Notify context cancelled.")
+				return
+			}
+		}
+	}()
+
 	go func() {
 		<-serversCtx.Done()
 
diff --git a/go.mod b/go.mod
index 6c70d7a..5a88576 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.16
 
 require (
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
+	github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf
 	github.com/cpuguy83/go-md2man/v2 v2.0.1 // indirect
 	github.com/fujiwara/shapeio v1.0.0
 	github.com/gliderlabs/ssh v0.3.3
diff --git a/go.sum b/go.sum
index 0a9c139..204c14f 100644
--- a/go.sum
+++ b/go.sum
@@ -61,6 +61,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/coreos/go-systemd v0.0.0-20190321100706-95778dfbb74e/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
+github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf h1:iW4rZ826su+pqaw19uhpSCzhj44qo35pNgKFGqzDKkU=
+github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4=
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/cpuguy83/go-md2man/v2 v2.0.1 h1:r/myEWzV9lfsM1tFLgDyu0atFtJ1fXn261LKYj/3DxU=
 github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
diff --git a/version.go b/version.go
index 95f7900..ba27e0d 100644
--- a/version.go
+++ b/version.go
@@ -5,7 +5,7 @@ import (
 	"runtime"
 )
 
-var Version = "v0.1.21"
+var Version = "v0.1.22"
 var Build string
 
 func FullVersion() string {