From c2a5b0b525b70af108b2b1135f0779fbb062554f Mon Sep 17 00:00:00 2001
From: AJ ONeal <coolaj86@gmail.com>
Date: Tue, 2 Nov 2021 01:30:02 -0600
Subject: [PATCH] feature: forcefully exit after given duration or at given
 wall clock time

---
 README.md              |  4 ++
 cmd/telebit/telebit.go | 89 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 40120d1..6e64fb7 100644
--- a/README.md
+++ b/README.md
@@ -175,6 +175,10 @@ Usage of telebit:
     	a list of resolvers in the format 8.8.8.8:53,8.8.4.4:53
   --env string
     	path to .env file
+  --exit-after string (such as "12h")
+      forcefully exit after the given duration
+  --exit-at string (such as "15:04")
+      forcefully exit at the given wall clock time
   --leeway duration
     	allow for time drift / skew (hard-coded to 15 minutes) (default 15m0s)
   LISTEN
diff --git a/cmd/telebit/telebit.go b/cmd/telebit/telebit.go
index a843a38..d83822b 100644
--- a/cmd/telebit/telebit.go
+++ b/cmd/telebit/telebit.go
@@ -78,6 +78,12 @@ var (
 
 	// defaultRelay should be set when compiled for the client
 	defaultRelay = "" //"https://telebit.app"
+
+	// exitAfter will cause the program to forcefully exit after the given duration
+	exitAfter time.Duration
+
+	// exitAt will cause the program to forcefully exit at the given time
+	exitAt string
 )
 
 var bindAddrs []string
@@ -136,6 +142,36 @@ func main() {
 		}
 	}()
 
+	if len(exitAt) > 0 || exitAfter > 0 {
+		go func() {
+			// ex: "Nov 1st, 2021 11:30pm"
+			start := time.Now()
+			end, d, err := getEnd(start, exitAt, exitAfter)
+			if nil != err {
+				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+				os.Exit(exitBadArguments)
+				return
+			}
+
+			for {
+				time.Sleep(1 * time.Minute)
+				now := time.Now()
+				// check monotonic time (duration)
+				if now.Sub(start) > d || now.Sub(end) > 0 {
+					break
+				}
+				// check wall clock time
+				if now.UTC().Format(time.RFC3339) > end.UTC().Format(time.RFC3339) {
+					break
+				}
+			}
+
+			// TODO does this restart automatically?
+			prg.Stop()
+			os.Exit(exitRetry)
+		}()
+	}
+
 	// call svc.Run to start your program/service
 	// svc.Run will call Init, Start, and Stop
 	if err := svc.Run(&prg); err != nil {
@@ -143,6 +179,51 @@ func main() {
 	}
 }
 
+func getEnd(start time.Time, exitAt string, exitAfter time.Duration) (time.Time, time.Duration, error) {
+	var end time.Time
+
+	if exitAfter > 0 {
+		if 0 == len(exitAt) {
+			end = start.Add(exitAfter)
+		}
+	}
+
+	if len(exitAt) > 0 {
+		// ex: "2:30am"
+		clock, err := time.Parse(exitAt, time.Kitchen)
+		if nil != err {
+			clock, err = time.Parse(exitAt, "15:04")
+			if nil != err {
+				return time.Time{}, 0, fmt.Errorf("could not parse --exit-at='%v'", exitAt)
+			}
+		}
+		// ex: "Nov 1st, 2021 2:30am"
+		end = time.Date(
+			start.Year(), start.Month(), start.Day(),
+			clock.Hour(), clock.Minute(), 0, 0,
+			start.Location(),
+		)
+		// ex: "-21h"
+		diff := end.Sub(start)
+		if diff < 0 {
+			// ex: "Nov 2nd, 2021 2:30am"
+			end = end.Add(24 * time.Hour)
+			end = time.Date(
+				end.Year(), end.Month(), end.Day(),
+				clock.Hour(), clock.Minute(), 0, 0,
+				end.Location(),
+			)
+			// ex: "3h"
+			diff = end.Sub(start)
+		}
+		if 0 == exitAfter {
+			exitAfter = diff
+		}
+	}
+
+	return end, exitAfter, nil
+}
+
 // implements svc.Service
 type program struct {
 	LogFile *os.File
@@ -188,6 +269,7 @@ func (p *program) Start() error {
 
 func (p *program) Stop() error {
 	log.Printf("Can't stop. Doing nothing instead.\n")
+	// TODO telebit.Shutdown()
 	return nil
 }
 
@@ -241,11 +323,14 @@ func parseFlagsAndENVs() {
 	flag.StringVar(&config.token, "token", "", "an auth token for the server (instead of generating --secret); use --token=false to ignore any $TOKEN in env")
 	flag.DurationVar(&config.leeway, "leeway", 15*time.Minute, "allow for time drift / skew (hard-coded to 15 minutes)")
 
-	bindAddrsStr := flag.String("listen", "", "list of bind addresses on which to listen, such as localhost:80, or :443")
+	bindAddrsStr := flag.String("listen", "", "list of bind addresses on which to listenfor relay clients, such as localhost:80, or :443")
 	tlsLocals := flag.String("tls-locals", "", "like --locals, but TLS will be used to connect to the local port")
-	locals := flag.String("locals", "", "a list of <from-domain>:<to-port>")
+	locals := flag.String("locals", "", "a list of <from-domain>:<to-port> to use as reverse proxy")
 	portToPorts := flag.String("port-forward", "", "a list of <from-port>:<to-port> for raw port-forwarding")
 
+	flag.DurationVar(&exitAfter, "exit-after", 0, "force quite after the specified duration")
+	flag.StringVar(&exitAt, "exit-at", "", "force quite after the specified time (ex: 15:04)")
+
 	flag.Parse()
 
 	if len(*envpath) > 0 {