add support for systemd notify

cmd/root.go

@@ -7,10 +7,12 @@ import (
 	log2 "log"
 	"net/http"
 	"os"
+	"os/signal"
 	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/NYTimes/logrotate"
@@ -26,6 +28,7 @@ import (
 
 	"github.com/pterodactyl/wings/config"
 	"github.com/pterodactyl/wings/environment"
+	"github.com/pterodactyl/wings/internal/notify"
 	"github.com/pterodactyl/wings/loggers/cli"
 	"github.com/pterodactyl/wings/remote"
 	"github.com/pterodactyl/wings/router"
@@ -308,43 +311,57 @@ func rootCmdRun(cmd *cobra.Command, _ []string) {
 	}
 
 	// Check if the server should run with TLS but using autocert.
-	if autotls {
-		m := autocert.Manager{
-			Prompt:     autocert.AcceptTOS,
-			Cache:      autocert.DirCache(path.Join(sys.RootDirectory, "/.tls-cache")),
-			HostPolicy: autocert.HostWhitelist(tlshostname),
-		}
-
-		log.WithField("hostname", tlshostname).Info("webserver is now listening with auto-TLS enabled; certificates will be automatically generated by Let's Encrypt")
-
-		// Hook autocert into the main http server.
-		s.TLSConfig.GetCertificate = m.GetCertificate
-		s.TLSConfig.NextProtos = append(s.TLSConfig.NextProtos, acme.ALPNProto) // enable tls-alpn ACME challenges
-
-		// Start the autocert server.
-		go func() {
-			if err := http.ListenAndServe(":http", m.HTTPHandler(nil)); err != nil {
-				log.WithError(err).Error("failed to serve autocert http server")
+	go func(s *http.Server, api config.ApiConfiguration, sys config.SystemConfiguration, autotls bool, tlshostname string) {
+		if autotls {
+			m := autocert.Manager{
+				Prompt:     autocert.AcceptTOS,
+				Cache:      autocert.DirCache(path.Join(sys.RootDirectory, "/.tls-cache")),
+				HostPolicy: autocert.HostWhitelist(tlshostname),
 			}
-		}()
-		// Start the main http server with TLS using autocert.
-		if err := s.ListenAndServeTLS("", ""); err != nil {
-			log.WithFields(log.Fields{"auto_tls": true, "tls_hostname": tlshostname, "error": err}).Fatal("failed to configure HTTP server using auto-tls")
+
+			log.WithField("hostname", tlshostname).Info("webserver is now listening with auto-TLS enabled; certificates will be automatically generated by Let's Encrypt")
+
+			// Hook autocert into the main http server.
+			s.TLSConfig.GetCertificate = m.GetCertificate
+			s.TLSConfig.NextProtos = append(s.TLSConfig.NextProtos, acme.ALPNProto) // enable tls-alpn ACME challenges
+
+			// Start the autocert server.
+			go func() {
+				if err := http.ListenAndServe(":http", m.HTTPHandler(nil)); err != nil {
+					log.WithError(err).Error("failed to serve autocert http server")
+				}
+			}()
+			// Start the main http server with TLS using autocert.
+			if err := s.ListenAndServeTLS("", ""); err != nil {
+				log.WithFields(log.Fields{"auto_tls": true, "tls_hostname": tlshostname, "error": err}).Fatal("failed to configure HTTP server using auto-tls")
+			}
+			return
 		}
-		return
+
+		// Check if main http server should run with TLS. Otherwise reset the TLS
+		// config on the server and then serve it over normal HTTP.
+		if api.Ssl.Enabled {
+			if err := s.ListenAndServeTLS(strings.ToLower(api.Ssl.CertificateFile), strings.ToLower(api.Ssl.KeyFile)); err != nil {
+				log.WithFields(log.Fields{"auto_tls": false, "error": err}).Fatal("failed to configure HTTPS server")
+			}
+			return
+		}
+		s.TLSConfig = nil
+		if err := s.ListenAndServe(); err != nil {
+			log.WithField("error", err).Fatal("failed to configure HTTP server")
+		}
+	}(s, api, sys, autotls, tlshostname)
+
+	if err := notify.Readiness(); err != nil {
+		log.WithField("error", err).Error("failed to notify systemd of readiness state")
 	}
 
-	// Check if main http server should run with TLS. Otherwise reset the TLS
-	// config on the server and then serve it over normal HTTP.
-	if api.Ssl.Enabled {
-		if err := s.ListenAndServeTLS(strings.ToLower(api.Ssl.CertificateFile), strings.ToLower(api.Ssl.KeyFile)); err != nil {
-			log.WithFields(log.Fields{"auto_tls": false, "error": err}).Fatal("failed to configure HTTPS server")
-		}
-		return
-	}
-	s.TLSConfig = nil
-	if err := s.ListenAndServe(); err != nil {
-		log.WithField("error", err).Fatal("failed to configure HTTP server")
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
+	<-c
+
+	if err := notify.Stopping(); err != nil {
+		log.WithField("error", err).Error("failed to notify systemd of stopping state")
 	}
 }
 

internal/notify/notify.go

@@ -0,0 +1,19 @@
+// Package notify handles notifying the operating system of the program's state.
+//
+// For linux based operating systems, this is done through the systemd socket
+// set by "NOTIFY_SOCKET" environment variable.
+//
+// Currently, no other operating systems are supported.
+package notify
+
+func Readiness() error {
+	return readiness()
+}
+
+func Reloading() error {
+	return reloading()
+}
+
+func Stopping() error {
+	return stopping()
+}

internal/notify/notify_linux.go

@@ -0,0 +1,48 @@
+package notify
+
+import (
+	"io"
+	"net"
+	"os"
+	"strings"
+)
+
+func notify(path string, r io.Reader) error {
+	s := &net.UnixAddr{
+		Name: path,
+		Net:  "unixgram",
+	}
+	c, err := net.DialUnix(s.Net, nil, s)
+	if err != nil {
+		return err
+	}
+	defer c.Close()
+
+	if _, err := io.Copy(c, r); err != nil {
+		return err
+	}
+	return nil
+}
+
+func socketNotify(payload string) error {
+	v, ok := os.LookupEnv("NOTIFY_SOCKET")
+	if !ok || v == "" {
+		return nil
+	}
+	if err := notify(v, strings.NewReader(payload)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func readiness() error {
+	return socketNotify("READY=1")
+}
+
+func reloading() error {
+	return socketNotify("RELOADING=1")
+}
+
+func stopping() error {
+	return socketNotify("STOPPING=1")
+}

internal/notify/notify_other.go

@@ -0,0 +1,15 @@
+// +build !linux
+
+package notify
+
+func readiness() error {
+	return nil
+}
+
+func reloading() error {
+	return nil
+}
+
+func stopping() error {
+	return nil
+}