From 578bfff9b92adaae29c8fb09835b839697ba5e6c Mon Sep 17 00:00:00 2001
From: Dean Sheather <dean@deansheather.com>
Date: Mon, 27 Nov 2023 15:43:51 +1000
Subject: [PATCH] fix: avoid panic on re-register timer (#23)

---
 tunnelsdk/tunnel.go | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tunnelsdk/tunnel.go b/tunnelsdk/tunnel.go
index 2a27505..cd30cc9 100644
--- a/tunnelsdk/tunnel.go
+++ b/tunnelsdk/tunnel.go
@@ -2,10 +2,12 @@ package tunnelsdk
 
 import (
 	"context"
+	"crypto/rand"
 	"crypto/sha512"
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"math/big"
 	"net"
 	"net/netip"
 	"net/url"
@@ -194,6 +196,9 @@ func (c *Client) LaunchTunnel(ctx context.Context, cfg TunnelConfig) (*Tunnel, e
 	if len(res.TunnelURLs) == 0 {
 		return nil, xerrors.Errorf("no tunnel urls returned from server")
 	}
+	if res.ReregisterWait <= 0 {
+		return nil, xerrors.Errorf("invalid reregister wait time: %s", res.ReregisterWait)
+	}
 
 	primaryURL, err := url.Parse(res.TunnelURLs[0])
 	if err != nil {
@@ -247,6 +252,17 @@ func (c *Client) LaunchTunnel(ctx context.Context, cfg TunnelConfig) (*Tunnel, e
 				cfg.Log.Warn(ctx, "periodically re-register tunnel", slog.Error(err))
 			}
 
+			// If we failed to re-register, try again in 30 seconds plus a
+			// random amount of time between 0 and 30 seconds.
+			if res.ReregisterWait <= 0 {
+				res.ReregisterWait = 30 * time.Second
+				i, err := rand.Int(rand.Reader, big.NewInt(30))
+				if err != nil {
+					i = big.NewInt(30)
+				}
+				res.ReregisterWait += time.Duration(i.Int64()) * time.Second
+			}
+
 			ticker.Reset(res.ReregisterWait)
 			cancel()
 		}