From 3a887b3c65e093e52d38a785a864e9b12efaac2b Mon Sep 17 00:00:00 2001
From: networkException <git@nwex.de>
Date: Mon, 30 Sep 2024 21:26:22 +0200
Subject: [PATCH] feat(server): support listener passed by systemd for socket
 activation

This patch adds support for server to be used with systemd's
socket activation by expecting an open file descriptor to be
announced using environment variables.

To make use of this set either HTTPListenNetwork or
GRPCListenNetwork to "sd_listen_fd" and HTTPListenAddress
or GRPCListenAddress either to an ASCII string passed as
the "FileDescriptorName" from a systemd.socket unit or
"LISTEN_FD_$n", where $n is the file descriptor number.

By default, when using "sd_listen_fd", the address will
be "LISTEN_FD_3".

See https://www.freedesktop.org/software/systemd/man/latest/systemd.socket.html#FileDescriptorName=
---
 go.mod           |  2 +-
 server/server.go | 63 ++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/go.mod b/go.mod
index 771fc1e18..0089b8e79 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/armon/go-metrics v0.3.10
 	github.com/aws/aws-sdk-go v1.44.321
 	github.com/cespare/xxhash/v2 v2.3.0
+	github.com/coreos/go-systemd/v22 v22.5.0
 	github.com/cristalhq/hedgedhttp v0.9.1
 	github.com/davecgh/go-spew v1.1.1
 	github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb
@@ -61,7 +62,6 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd // indirect
 	github.com/coreos/go-semver v0.3.0 // indirect
-	github.com/coreos/go-systemd/v22 v22.5.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
diff --git a/server/server.go b/server/server.go
index 7b8e7593d..0e66ad8fd 100644
--- a/server/server.go
+++ b/server/server.go
@@ -7,12 +7,14 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"math"
 	"net"
 	"net/http"
 	_ "net/http/pprof" // anonymous import to get the pprof handler registered
+	"os"
 	"strconv"
 	"strings"
 	"time"
@@ -38,6 +40,8 @@ import (
 	"github.com/grafana/dskit/log"
 	"github.com/grafana/dskit/middleware"
 	"github.com/grafana/dskit/signals"
+
+	"github.com/coreos/go-systemd/v22/activation"
 )
 
 // Listen on the named network
@@ -47,6 +51,13 @@ const (
 	DefaultNetwork = "tcp"
 	// NetworkTCPV4 for IPV4 only
 	NetworkTCPV4 = "tcp4"
+	// NetworkSystemdListenFd for using a passed open file descriptor
+	NetworkSystemdListenFd = "sd_listen_fd"
+)
+
+const (
+	// See `SD_LISTEN_FDS(3)`
+	NetworkSystemdListenFdsStart = 3
 )
 
 // SignalHandler used by Server.
@@ -248,6 +259,43 @@ func NewWithMetrics(cfg Config, metrics *Metrics) (*Server, error) {
 	return newServer(cfg, metrics)
 }
 
+func listen(network string, address string, port int, systemdListenFiles []*os.File) (net.Listener, error) {
+	if network == "" {
+		network = DefaultNetwork
+	}
+
+	switch network {
+	case DefaultNetwork, NetworkTCPV4:
+		return net.Listen(network, net.JoinHostPort(address, strconv.Itoa(port)))
+	case NetworkSystemdListenFd:
+		if address == "" {
+			address = "LISTEN_FD_" + strconv.Itoa(NetworkSystemdListenFdsStart)
+		}
+
+		listeners := map[string]net.Listener{}
+		for index, file := range systemdListenFiles {
+			if listener, err := net.FileListener(file); err == nil {
+				listeners[file.Name()] = listener
+				listeners["LISTEN_FD_"+strconv.Itoa(index+NetworkSystemdListenFdsStart)] = listener
+				file.Close()
+			}
+		}
+
+		bs, _ := json.Marshal(listeners)
+		fmt.Println(string(bs))
+
+		listener, hasListener := listeners[address]
+
+		if !hasListener {
+			return nil, fmt.Errorf("could not listen on 'sd_listen_fd', no file descriptor given for name '%s'", address)
+		}
+
+		return listener, nil
+	default:
+		return nil, fmt.Errorf("cannot listen on unknown network '%s'", network)
+	}
+}
+
 func newServer(cfg Config, metrics *Metrics) (*Server, error) {
 	// If user doesn't supply a logging implementation, by default instantiate go-kit.
 	logger := cfg.Log
@@ -260,15 +308,14 @@ func newServer(cfg Config, metrics *Metrics) (*Server, error) {
 		gatherer = prometheus.DefaultGatherer
 	}
 
-	network := cfg.HTTPListenNetwork
-	if network == "" {
-		network = DefaultNetwork
-	}
+	systemdListenFiles := activation.Files(true)
+
 	// Setup listeners first, so we can fail early if the port is in use.
-	httpListener, err := net.Listen(network, net.JoinHostPort(cfg.HTTPListenAddress, strconv.Itoa(cfg.HTTPListenPort)))
+	httpListener, err := listen(cfg.HTTPListenNetwork, cfg.HTTPListenAddress, cfg.HTTPListenPort, systemdListenFiles)
 	if err != nil {
 		return nil, err
 	}
+
 	httpListener = middleware.CountingListener(httpListener, metrics.TCPConnections.WithLabelValues("http"))
 	if cfg.HTTPLogClosedConnectionsWithoutResponse {
 		httpListener = middleware.NewZeroResponseListener(httpListener, level.Warn(logger))
@@ -279,11 +326,7 @@ func newServer(cfg Config, metrics *Metrics) (*Server, error) {
 		httpListener = netutil.LimitListener(httpListener, cfg.HTTPConnLimit)
 	}
 
-	network = cfg.GRPCListenNetwork
-	if network == "" {
-		network = DefaultNetwork
-	}
-	grpcListener, err := net.Listen(network, net.JoinHostPort(cfg.GRPCListenAddress, strconv.Itoa(cfg.GRPCListenPort)))
+	grpcListener, err := listen(cfg.GRPCListenNetwork, cfg.GRPCListenAddress, cfg.GRPCListenPort, systemdListenFiles)
 	if err != nil {
 		return nil, err
 	}