From 57b79aa852973bbc9ac3ab6a952061b83c1d008b Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Wed, 18 Aug 2021 23:21:11 +0100
Subject: [PATCH] Set timeout, add lastupdate field

This commit makes two reasonably major changes:

Set a default timeout for the go HTTP server (which gin uses), which
allows us to actually have broken long poll sessions fail so we can have
the client re-establish them.
The current 10s number is chosen randomly and we need more testing to
ensure that the feature work as intended.

The second is adding a last updated field to keep track of the last time
we had an update that needs to be propagated to all of our
clients/nodes. This will be used to keep track of our machines and if
they are up to date or need us to push an update.
---
 app.go | 52 ++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 40 insertions(+), 12 deletions(-)

diff --git a/app.go b/app.go
index fcf287f7..255a7df1 100644
--- a/app.go
+++ b/app.go
@@ -58,7 +58,10 @@ type Headscale struct {
 	aclPolicy *ACLPolicy
 	aclRules  *[]tailcfg.FilterRule
 
-	clientsPolling sync.Map
+	clientsUpdateChannels sync.Map
+
+	lastStateChangeMutex sync.RWMutex
+	lastStateChange      time.Time
 }
 
 // NewHeadscale returns the Headscale app
@@ -85,12 +88,13 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
 	}
 
 	h := Headscale{
-		cfg:        cfg,
-		dbType:     cfg.DBtype,
-		dbString:   dbString,
-		privateKey: privKey,
-		publicKey:  &pubKey,
-		aclRules:   &tailcfg.FilterAllowAll, // default allowall
+		cfg:             cfg,
+		dbType:          cfg.DBtype,
+		dbString:        dbString,
+		privateKey:      privKey,
+		publicKey:       &pubKey,
+		aclRules:        &tailcfg.FilterAllowAll, // default allowall
+		lastStateChange: time.Now().UTC(),
 	}
 
 	err = h.initDB()
@@ -168,6 +172,13 @@ func (h *Headscale) Serve() error {
 	go h.watchForKVUpdates(5000)
 	go h.expireEphemeralNodes(5000)
 
+	s := &http.Server{
+		Addr:         h.cfg.Addr,
+		Handler:      r,
+		ReadTimeout:  10 * time.Second,
+		WriteTimeout: 10 * time.Second,
+	}
+
 	if h.cfg.TLSLetsEncryptHostname != "" {
 		if !strings.HasPrefix(h.cfg.ServerURL, "https://") {
 			log.Warn().Msg("Listening with TLS but ServerURL does not start with https://")
@@ -179,9 +190,11 @@ func (h *Headscale) Serve() error {
 			Cache:      autocert.DirCache(h.cfg.TLSLetsEncryptCacheDir),
 		}
 		s := &http.Server{
-			Addr:      h.cfg.Addr,
-			TLSConfig: m.TLSConfig(),
-			Handler:   r,
+			Addr:         h.cfg.Addr,
+			TLSConfig:    m.TLSConfig(),
+			Handler:      r,
+			ReadTimeout:  10 * time.Second,
+			WriteTimeout: 10 * time.Second,
 		}
 		if h.cfg.TLSLetsEncryptChallengeType == "TLS-ALPN-01" {
 			// Configuration via autocert with TLS-ALPN-01 (https://tools.ietf.org/html/rfc8737)
@@ -206,12 +219,27 @@ func (h *Headscale) Serve() error {
 		if !strings.HasPrefix(h.cfg.ServerURL, "http://") {
 			log.Warn().Msg("Listening without TLS but ServerURL does not start with http://")
 		}
-		err = r.Run(h.cfg.Addr)
+		err = s.ListenAndServe()
 	} else {
 		if !strings.HasPrefix(h.cfg.ServerURL, "https://") {
 			log.Warn().Msg("Listening with TLS but ServerURL does not start with https://")
 		}
-		err = r.RunTLS(h.cfg.Addr, h.cfg.TLSCertPath, h.cfg.TLSKeyPath)
+		err = s.ListenAndServeTLS(h.cfg.TLSCertPath, h.cfg.TLSKeyPath)
 	}
 	return err
 }
+
+func (h *Headscale) setLastStateChangeToNow() {
+	h.lastStateChangeMutex.Lock()
+
+	now := time.Now().UTC()
+	h.lastStateChange = now
+
+	h.lastStateChangeMutex.Unlock()
+}
+
+func (h *Headscale) getLastStateChange() time.Time {
+	h.lastStateChangeMutex.RLock()
+	defer h.lastStateChangeMutex.RUnlock()
+	return h.lastStateChange
+}