Fetch authkey in two phases

Initially see if the provisioned server or global key exist in the app
usage pool. If not then fallback to the already existing environment
variables

Signed-off-by: Connor Kelly <connor.r.kelly@gmail.com>

Author: clly

caddyfile.go

@@ -10,6 +10,11 @@ func init() {
 	httpcaddyfile.RegisterGlobalOption("tailscale", parseApp)
 }
 
+const (
+	authUsageKey = "auth_key"
+	ephemeralKey = "ephemeral"
+)
+
 func parseApp(d *caddyfile.Dispenser, _ any) (any, error) {
 	app := &TSApp{
 		Servers: make(map[string]TSServer),

module.go

@@ -106,12 +106,8 @@ func getServer(_, addr string) (*tsnetServerDestructor, error) {
 		}
 
 		if host != "" {
-			// Set authkey to "TS_AUTHKEY_<HOST>".  If empty,
-			// fall back to "TS_AUTHKEY".
-			s.AuthKey = os.Getenv("TS_AUTHKEY_" + strings.ToUpper(host))
-			if s.AuthKey == "" {
-				s.AuthKey = os.Getenv("TS_AUTHKEY")
-			}
+			s.AuthKey = getAuthKey(host)
+			log.Println("auth_key", s.AuthKey)
 
 			// Set config directory for tsnet.  By default, tsnet will use the name of the
 			// running program, but we include the hostname as well so that a single
@@ -137,6 +133,21 @@ func getServer(_, addr string) (*tsnetServerDestructor, error) {
 	return s.(*tsnetServerDestructor), nil
 }
 
+func getAuthKey(host string) string {
+	storedAuthKey, loaded := app.LoadOrStore(authUsageKey, "")
+	if loaded {
+		return storedAuthKey.(string)
+	}
+
+	// Set authkey to "TS_AUTHKEY_<HOST>".  If empty,
+	// fall back to "TS_AUTHKEY".
+	authKey := os.Getenv("TS_AUTHKEY_" + strings.ToUpper(host))
+	if authKey == "" {
+		authKey = os.Getenv("TS_AUTHKEY")
+	}
+	return authKey
+}
+
 type TailscaleAuth struct {
 	localclient *tailscale.LocalClient
 }

module_test.go

@@ -0,0 +1,65 @@
+package tscaddy
+
+import (
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/caddyserver/caddy/v2"
+)
+
+func Test_GetAuthKey(t *testing.T) {
+
+	const testkey = "abcdefghijklmnopqrstuvwxyz"
+	const testHostKey = "1234567890"
+	const testenvkey = "zyxwvutsrqponmlkjihgfedca"
+	const testHost = "unittest"
+
+	tests := map[string]struct {
+		host    string
+		skipApp bool
+		want    string
+	}{
+		"default key from module": {
+			want: testkey,
+			host: "testhost",
+		},
+		"default key from environment": {
+			want:    testenvkey,
+			skipApp: true,
+			host:    "testhost",
+		},
+		"host key from module": {
+			want: testHostKey,
+			host: testHost,
+		},
+		"host key from environment": {
+			want: testHostKey,
+			host: testHost,
+		},
+	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			app = caddy.NewUsagePool()
+			if !tt.skipApp {
+				app.LoadOrStore(authUsageKey, testkey)
+				app.LoadOrStore(testHost, TSServer{
+					AuthKey: testHostKey,
+				})
+			}
+			os.Setenv("TS_AUTHKEY", testenvkey)
+			hostKey := fmt.Sprintf("TS_AUTHKEY_%s", strings.ToUpper(testHost))
+			os.Setenv(hostKey, testHostKey)
+			t.Cleanup(func() {
+				os.Unsetenv("TS_AUTHKEY")
+				os.Unsetenv(hostKey)
+			})
+
+			got := getAuthKey(tt.host)
+			if got != tt.want {
+				t.Errorf("GetAuthKey() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}