From 7f9120ac0cb3a5ee87011412eda54c7db6b5f003 Mon Sep 17 00:00:00 2001
From: Davor Gajic <davor@nine.ch>
Date: Wed, 16 Oct 2024 14:10:37 +0200
Subject: [PATCH] make the token and control urls configurable

---
 cmd/k8s-operator/operator.go | 13 ++++++++++---
 cmd/k8s-operator/sts.go      | 11 +++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/cmd/k8s-operator/operator.go b/cmd/k8s-operator/operator.go
index d8dd403cc6097..e667163aa6e46 100644
--- a/cmd/k8s-operator/operator.go
+++ b/cmd/k8s-operator/operator.go
@@ -121,6 +121,8 @@ func initTSNet(zlog *zap.SugaredLogger) (*tsnet.Server, *tailscale.Client) {
 	var (
 		clientIDPath     = defaultEnv("CLIENT_ID_FILE", "")
 		clientSecretPath = defaultEnv("CLIENT_SECRET_FILE", "")
+		tokenURL         = defaultEnv("TOKEN_URL", "https://login.tailscale.com/api/v2/oauth/token")
+		controlURL       = defaultEnv("CONTROL_URL", "")
 		hostname         = defaultEnv("OPERATOR_HOSTNAME", "tailscale-operator")
 		kubeSecret       = defaultEnv("OPERATOR_SECRET", "")
 		operatorTags     = defaultEnv("OPERATOR_INITIAL_TAGS", "tag:k8s-operator")
@@ -140,15 +142,19 @@ func initTSNet(zlog *zap.SugaredLogger) (*tsnet.Server, *tailscale.Client) {
 	credentials := clientcredentials.Config{
 		ClientID:     string(clientID),
 		ClientSecret: string(clientSecret),
-		TokenURL:     "https://login.tailscale.com/api/v2/oauth/token",
+		TokenURL:     tokenURL,
 	}
 	tsClient := tailscale.NewClient("-", nil)
+	if controlURL != "" {
+		tsClient.BaseURL = controlURL
+	}
 	tsClient.UserAgent = "tailscale-k8s-operator"
 	tsClient.HTTPClient = credentials.Client(context.Background())
 
 	s := &tsnet.Server{
-		Hostname: hostname,
-		Logf:     zlog.Named("tailscaled").Debugf,
+		ControlURL: controlURL,
+		Hostname:   hostname,
+		Logf:       zlog.Named("tailscaled").Debugf,
 	}
 	if kubeSecret != "" {
 		st, err := kubestore.New(logger.Discard, kubeSecret)
@@ -271,6 +277,7 @@ func runReconcilers(opts reconcilerOpts) {
 		proxyImage:             opts.proxyImage,
 		proxyPriorityClassName: opts.proxyPriorityClassName,
 		tsFirewallMode:         opts.proxyFirewallMode,
+		controlUrl:             opts.tsServer.ControlURL,
 	}
 	err = builder.
 		ControllerManagedBy(mgr).
diff --git a/cmd/k8s-operator/sts.go b/cmd/k8s-operator/sts.go
index 6378a82636939..e17bac446ea05 100644
--- a/cmd/k8s-operator/sts.go
+++ b/cmd/k8s-operator/sts.go
@@ -122,6 +122,7 @@ type tailscaleSTSConfig struct {
 	Hostname string
 	Tags     []string // if empty, use defaultTags
 
+	ControlURL string
 	// Connector specifies a configuration of a Connector instance if that's
 	// what this StatefulSet should be created for.
 	Connector *connector
@@ -150,6 +151,7 @@ type tailscaleSTSReconciler struct {
 	proxyImage             string
 	proxyPriorityClassName string
 	tsFirewallMode         string
+	controlUrl             string
 }
 
 func (sts tailscaleSTSReconciler) validate() error {
@@ -186,6 +188,10 @@ func (a *tailscaleSTSReconciler) Provision(ctx context.Context, logger *zap.Suga
 	}
 	sts.ProxyClass = proxyClass
 
+	if a.controlUrl != "" {
+		sts.ControlURL = a.controlUrl
+	}
+
 	secretName, tsConfigHash, configs, err := a.createOrGetSecret(ctx, logger, sts, hsvc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create or get API key secret: %w", err)
@@ -830,6 +836,11 @@ func tailscaledConfig(stsC *tailscaleSTSConfig, newAuthkey string, oldSecret *co
 		}
 		conf.AuthKey = key
 	}
+
+	if stsC.ControlURL != "" {
+		conf.ServerURL = &stsC.ControlURL
+	}
+
 	capVerConfigs := make(map[tailcfg.CapabilityVersion]ipn.ConfigVAlpha)
 	capVerConfigs[95] = *conf
 	// legacy config should not contain NoStatefulFiltering field.