From f985ed567f7cc7a8b1d8427d5c6d608b66269fbc Mon Sep 17 00:00:00 2001
From: Nuno Diegues <nuno@cloudflare.com>
Date: Tue, 21 Sep 2021 11:02:59 +0100
Subject: [PATCH] TUN-5128: Enforce maximum grace period

This maximum grace period will be honored by Cloudflare edge such that
either side will close the connection after unregistration at most
by this time (3min as of this commit):
 - If the connection is unused, it is already closed as soon as possible.
 - If the connection is still used, it is closed on the cloudflared configured grace-period.

Even if cloudflared does not close the connection by the grace-period time,
the edge will do so.
---
 cmd/cloudflared/tunnel/cmd.go           |  6 +++++-
 cmd/cloudflared/tunnel/configuration.go | 15 ++++++++++++++-
 connection/connection.go                |  1 +
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/cmd/cloudflared/tunnel/cmd.go b/cmd/cloudflared/tunnel/cmd.go
index 1ed2289a..5afd801f 100644
--- a/cmd/cloudflared/tunnel/cmd.go
+++ b/cmd/cloudflared/tunnel/cmd.go
@@ -384,7 +384,11 @@ func StartServer(
 		observer.RegisterSink(app)
 	}
 
-	return waitToShutdown(&wg, cancel, errC, graceShutdownC, c.Duration("grace-period"), log)
+	gracePeriod, err := gracePeriod(c)
+	if err != nil {
+		return err
+	}
+	return waitToShutdown(&wg, cancel, errC, graceShutdownC, gracePeriod, log)
 }
 
 func waitToShutdown(wg *sync.WaitGroup,
diff --git a/cmd/cloudflared/tunnel/configuration.go b/cmd/cloudflared/tunnel/configuration.go
index dc264c31..901d192c 100644
--- a/cmd/cloudflared/tunnel/configuration.go
+++ b/cmd/cloudflared/tunnel/configuration.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/google/uuid"
 	homedir "github.com/mitchellh/go-homedir"
@@ -260,9 +261,13 @@ func prepareTunnelConfig(
 	}
 
 	originProxy := origin.NewOriginProxy(ingressRules, warpRoutingService, tags, log)
+	gracePeriod, err := gracePeriod(c)
+	if err != nil {
+		return nil, ingress.Ingress{}, err
+	}
 	connectionConfig := &connection.Config{
 		OriginProxy:     originProxy,
-		GracePeriod:     c.Duration("grace-period"),
+		GracePeriod:     gracePeriod,
 		ReplaceExisting: c.Bool("force"),
 	}
 	muxerConfig := &connection.MuxerConfig{
@@ -300,6 +305,14 @@ func prepareTunnelConfig(
 	}, ingressRules, nil
 }
 
+func gracePeriod(c *cli.Context) (time.Duration, error) {
+	period := c.Duration("grace-period")
+	if period > connection.MaxGracePeriod {
+		return time.Duration(0), fmt.Errorf("grace-period must be equal or less than %v", connection.MaxGracePeriod)
+	}
+	return period, nil
+}
+
 func isWarpRoutingEnabled(warpConfig config.WarpRoutingConfig, isNamedTunnel bool) bool {
 	return warpConfig.Enabled && isNamedTunnel
 }
diff --git a/connection/connection.go b/connection/connection.go
index dbe5ef1e..f061672c 100644
--- a/connection/connection.go
+++ b/connection/connection.go
@@ -18,6 +18,7 @@ import (
 const (
 	lbProbeUserAgentPrefix = "Mozilla/5.0 (compatible; Cloudflare-Traffic-Manager/1.0; +https://www.cloudflare.com/traffic-manager/;"
 	LogFieldConnIndex      = "connIndex"
+	MaxGracePeriod         = time.Minute * 3
 )
 
 var switchingProtocolText = fmt.Sprintf("%d %s", http.StatusSwitchingProtocols, http.StatusText(http.StatusSwitchingProtocols))