adding server url to proxygroups when a custom tailnet has been specifiedchaosinthecrd/adding-server-url-multi-tailnet

Signed-off-by: chaosinthecrd <tom@tmlabs.co.uk>

5 files changed, 129 insertions, 59 deletions

diff --git a/cmd/k8s-operator/api-server-proxy-pg.go b/cmd/k8s-operator/api-server-proxy-pg.go
index d5fc63311..0900fd0aa 100644
--- a/cmd/k8s-operator/api-server-proxy-pg.go
+++ b/cmd/k8s-operator/api-server-proxy-pg.go
@@ -78,12 +78,9 @@ func (r *KubeAPIServerTSServiceReconciler) Reconcile(ctx context.Context, req re
 	serviceName := serviceNameForAPIServerProxy(pg)
 	logger = logger.With("Tailscale Service", serviceName)
 
-	tailscaleClient := r.tsClient
-	if pg.Spec.Tailnet != "" {
-		tailscaleClient, err = clientForTailnet(ctx, r.Client, r.tsNamespace, pg.Spec.Tailnet)
-		if err != nil {
-			return res, fmt.Errorf("failed to get tailscale client: %w", err)
-		}
+	tailscaleClient, err := r.getClient(ctx, pg.Spec.Tailnet)
+	if err != nil {
+		return res, fmt.Errorf("failed to get tailscale client: %w", err)
 	}
 
 	if markedForDeletion(pg) {
@@ -108,6 +105,22 @@ func (r *KubeAPIServerTSServiceReconciler) Reconcile(ctx context.Context, req re
 	return reconcile.Result{}, nil
 }
 
+// getClient returns the appropriate Tailscale client for the given tailnet.
+// If no tailnet is specified, returns the default client.
+func (r *KubeAPIServerTSServiceReconciler) getClient(ctx context.Context, tailnetName string) (tsClient,
+	error) {
+	if tailnetName == "" {
+		return r.tsClient, nil
+	}
+
+	tc, _, err := clientForTailnet(ctx, r.Client, r.tsNamespace, tailnetName)
+	if err != nil {
+		return nil, err
+	}
+
+	return tc, nil
+}
+
 // maybeProvision ensures that a Tailscale Service for this ProxyGroup exists
 // and is up to date.
 //
diff --git a/cmd/k8s-operator/proxygroup.go b/cmd/k8s-operator/proxygroup.go
index f4eab202e..44629b71b 100644
--- a/cmd/k8s-operator/proxygroup.go
+++ b/cmd/k8s-operator/proxygroup.go
@@ -119,20 +119,9 @@ func (r *ProxyGroupReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 		return reconcile.Result{}, fmt.Errorf("failed to get tailscale.com ProxyGroup: %w", err)
 	}
 
-	tailscaleClient := r.tsClient
-	if pg.Spec.Tailnet != "" {
-		tc, err := clientForTailnet(ctx, r.Client, r.tsNamespace, pg.Spec.Tailnet)
-		if err != nil {
-			oldPGStatus := pg.Status.DeepCopy()
-			nrr := &notReadyReason{
-				reason:  reasonProxyGroupTailnetUnavailable,
-				message: err.Error(),
-			}
-
-			return reconcile.Result{}, errors.Join(err, r.maybeUpdateStatus(ctx, logger, pg, oldPGStatus, nrr, make(map[string][]netip.AddrPort)))
-		}
-
-		tailscaleClient = tc
+	tailscaleClient, loginUrl, err := r.getClientAndLoginURL(ctx, pg.Spec.Tailnet)
+	if err != nil {
+		return reconcile.Result{}, fmt.Errorf("failed to get tailscale client and loginUrl: %w", err)
 	}
 
 	if markedForDeletion(pg) {
@@ -162,7 +151,7 @@ func (r *ProxyGroupReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 	}
 
 	oldPGStatus := pg.Status.DeepCopy()
-	staticEndpoints, nrr, err := r.reconcilePG(ctx, tailscaleClient, pg, logger)
+	staticEndpoints, nrr, err := r.reconcilePG(ctx, tailscaleClient, loginUrl, pg, logger)
 	return reconcile.Result{}, errors.Join(err, r.maybeUpdateStatus(ctx, logger, pg, oldPGStatus, nrr, staticEndpoints))
 }
 
@@ -170,7 +159,7 @@ func (r *ProxyGroupReconciler) Reconcile(ctx context.Context, req reconcile.Requ
 // for deletion. It is separated out from Reconcile to make a clear separation
 // between reconciling the ProxyGroup, and posting the status of its created
 // resources onto the ProxyGroup status field.
-func (r *ProxyGroupReconciler) reconcilePG(ctx context.Context, tailscaleClient tsClient, pg *tsapi.ProxyGroup, logger *zap.SugaredLogger) (map[string][]netip.AddrPort, *notReadyReason, error) {
+func (r *ProxyGroupReconciler) reconcilePG(ctx context.Context, tailscaleClient tsClient, loginUrl string, pg *tsapi.ProxyGroup, logger *zap.SugaredLogger) (map[string][]netip.AddrPort, *notReadyReason, error) {
 	if !slices.Contains(pg.Finalizers, FinalizerName) {
 		// This log line is printed exactly once during initial provisioning,
 		// because once the finalizer is in place this block gets skipped. So,
@@ -211,7 +200,7 @@ func (r *ProxyGroupReconciler) reconcilePG(ctx context.Context, tailscaleClient
 		return notReady(reasonProxyGroupInvalid, fmt.Sprintf("invalid ProxyGroup spec: %v", err))
 	}
 
-	staticEndpoints, nrr, err := r.maybeProvision(ctx, tailscaleClient, pg, proxyClass)
+	staticEndpoints, nrr, err := r.maybeProvision(ctx, tailscaleClient, loginUrl, pg, proxyClass)
 	if err != nil {
 		return nil, nrr, err
 	}
@@ -297,7 +286,7 @@ func (r *ProxyGroupReconciler) validate(ctx context.Context, pg *tsapi.ProxyGrou
 	return errors.Join(errs...)
 }
 
-func (r *ProxyGroupReconciler) maybeProvision(ctx context.Context, tailscaleClient tsClient, pg *tsapi.ProxyGroup, proxyClass *tsapi.ProxyClass) (map[string][]netip.AddrPort, *notReadyReason, error) {
+func (r *ProxyGroupReconciler) maybeProvision(ctx context.Context, tailscaleClient tsClient, loginUrl string, pg *tsapi.ProxyGroup, proxyClass *tsapi.ProxyClass) (map[string][]netip.AddrPort, *notReadyReason, error) {
 	logger := r.logger(pg.Name)
 	r.mu.Lock()
 	r.ensureAddedToGaugeForProxyGroup(pg)
@@ -320,7 +309,7 @@ func (r *ProxyGroupReconciler) maybeProvision(ctx context.Context, tailscaleClie
 		}
 	}
 
-	staticEndpoints, err := r.ensureConfigSecretsCreated(ctx, tailscaleClient, pg, proxyClass, svcToNodePorts)
+	staticEndpoints, err := r.ensureConfigSecretsCreated(ctx, tailscaleClient, loginUrl, pg, proxyClass, svcToNodePorts)
 	if err != nil {
 		var selectorErr *FindStaticEndpointErr
 		if errors.As(err, &selectorErr) {
@@ -735,6 +724,7 @@ func (r *ProxyGroupReconciler) deleteTailnetDevice(ctx context.Context, tailscal
 func (r *ProxyGroupReconciler) ensureConfigSecretsCreated(
 	ctx context.Context,
 	tailscaleClient tsClient,
+	loginUrl string,
 	pg *tsapi.ProxyGroup,
 	proxyClass *tsapi.ProxyClass,
 	svcToNodePorts map[string]uint16,
@@ -870,8 +860,8 @@ func (r *ProxyGroupReconciler) ensureConfigSecretsCreated(
 				}
 			}
 
-			if r.loginServer != "" {
-				cfg.ServerURL = &r.loginServer
+			if loginUrl != "" {
+				cfg.ServerURL = ptr.To(loginUrl)
 			}
 
 			if proxyClass != nil && proxyClass.Spec.TailscaleConfig != nil {
@@ -899,7 +889,7 @@ func (r *ProxyGroupReconciler) ensureConfigSecretsCreated(
 				return nil, err
 			}
 
-			configs, err := pgTailscaledConfig(pg, proxyClass, i, authKey, endpoints[nodePortSvcName], existingAdvertiseServices, r.loginServer)
+			configs, err := pgTailscaledConfig(pg, loginUrl, proxyClass, i, authKey, endpoints[nodePortSvcName], existingAdvertiseServices)
 			if err != nil {
 				return nil, fmt.Errorf("error creating tailscaled config: %w", err)
 			}
@@ -1056,7 +1046,7 @@ func (r *ProxyGroupReconciler) ensureRemovedFromGaugeForProxyGroup(pg *tsapi.Pro
 	gaugeAPIServerProxyGroupResources.Set(int64(r.apiServerProxyGroups.Len()))
 }
 
-func pgTailscaledConfig(pg *tsapi.ProxyGroup, pc *tsapi.ProxyClass, idx int32, authKey *string, staticEndpoints []netip.AddrPort, oldAdvertiseServices []string, loginServer string) (tailscaledConfigs, error) {
+func pgTailscaledConfig(pg *tsapi.ProxyGroup, loginServer string, pc *tsapi.ProxyClass, idx int32, authKey *string, staticEndpoints []netip.AddrPort, oldAdvertiseServices []string) (tailscaledConfigs, error) {
 	conf := &ipn.ConfigVAlpha{
 		Version:           "alpha0",
 		AcceptDNS:         "false",
@@ -1197,6 +1187,29 @@ func (r *ProxyGroupReconciler) getRunningProxies(ctx context.Context, pg *tsapi.
 	return devices, nil
 }
 
+// getClientAndLoginURL returns the appropriate Tailscale client and resolved login URL
+// for the given tailnet name. If no tailnet is specified, returns the default client
+// and login server. Applies fallback to the operator's login server if the tailnet
+// doesn't specify a custom login URL.
+func (r *ProxyGroupReconciler) getClientAndLoginURL(ctx context.Context, tailnetName string) (tsClient,
+	string, error) {
+	if tailnetName == "" {
+		return r.tsClient, r.loginServer, nil
+	}
+
+	tc, loginUrl, err := clientForTailnet(ctx, r.Client, r.tsNamespace, tailnetName)
+	if err != nil {
+		return nil, "", err
+	}
+
+	// Apply fallback if tailnet doesn't specify custom login URL
+	if loginUrl == "" {
+		loginUrl = r.loginServer
+	}
+
+	return tc, loginUrl, nil
+}
+
 type nodeMetadata struct {
 	ordinal     int
 	stateSecret *corev1.Secret
diff --git a/cmd/k8s-operator/sts.go b/cmd/k8s-operator/sts.go
index e81fe2d66..4345b1420 100644
--- a/cmd/k8s-operator/sts.go
+++ b/cmd/k8s-operator/sts.go
@@ -198,14 +198,9 @@ func IsHTTPSEnabledOnTailnet(tsnetServer tsnetServer) bool {
 // Provision ensures that the StatefulSet for the given service is running and
 // up to date.
 func (a *tailscaleSTSReconciler) Provision(ctx context.Context, logger *zap.SugaredLogger, sts *tailscaleSTSConfig) (*corev1.Service, error) {
-	tailscaleClient := a.tsClient
-	if sts.Tailnet != "" {
-		tc, err := clientForTailnet(ctx, a.Client, a.operatorNamespace, sts.Tailnet)
-		if err != nil {
-			return nil, err
-		}
-
-		tailscaleClient = tc
+	tailscaleClient, loginUrl, err := a.getClientAndLoginURL(ctx, sts.Tailnet)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get tailscale client and loginUrl: %w", err)
 	}
 
 	// Do full reconcile.
@@ -227,7 +222,7 @@ func (a *tailscaleSTSReconciler) Provision(ctx context.Context, logger *zap.Suga
 	}
 	sts.ProxyClass = proxyClass
 
-	secretNames, err := a.provisionSecrets(ctx, tailscaleClient, logger, sts, hsvc)
+	secretNames, err := a.provisionSecrets(ctx, tailscaleClient, loginUrl, sts, hsvc, logger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create or get API key secret: %w", err)
 	}
@@ -248,13 +243,36 @@ func (a *tailscaleSTSReconciler) Provision(ctx context.Context, logger *zap.Suga
 	return hsvc, nil
 }
 
+// getClientAndLoginURL returns the appropriate Tailscale client and resolved login URL
+// for the given tailnet name. If no tailnet is specified, returns the default client
+// and login server. Applies fallback to the operator's login server if the tailnet
+// doesn't specify a custom login URL.
+func (a *tailscaleSTSReconciler) getClientAndLoginURL(ctx context.Context, tailnetName string) (tsClient,
+	string, error) {
+	if tailnetName == "" {
+		return a.tsClient, a.loginServer, nil
+	}
+
+	tc, loginUrl, err := clientForTailnet(ctx, a.Client, a.operatorNamespace, tailnetName)
+	if err != nil {
+		return nil, "", err
+	}
+
+	// Apply fallback if tailnet doesn't specify custom login URL
+	if loginUrl == "" {
+		loginUrl = a.loginServer
+	}
+
+	return tc, loginUrl, nil
+}
+
 // Cleanup removes all resources associated that were created by Provision with
 // the given labels. It returns true when all resources have been removed,
 // otherwise it returns false and the caller should retry later.
 func (a *tailscaleSTSReconciler) Cleanup(ctx context.Context, tailnet string, logger *zap.SugaredLogger, labels map[string]string, typ string) (done bool, _ error) {
 	tailscaleClient := a.tsClient
 	if tailnet != "" {
-		tc, err := clientForTailnet(ctx, a.Client, a.operatorNamespace, tailnet)
+		tc, _, err := clientForTailnet(ctx, a.Client, a.operatorNamespace, tailnet)
 		if err != nil {
 			logger.Errorf("failed to get tailscale client: %v", err)
 			return false, nil
@@ -385,7 +403,7 @@ func (a *tailscaleSTSReconciler) reconcileHeadlessService(ctx context.Context, l
 	return createOrUpdate(ctx, a.Client, a.operatorNamespace, hsvc, func(svc *corev1.Service) { svc.Spec = hsvc.Spec })
 }
 
-func (a *tailscaleSTSReconciler) provisionSecrets(ctx context.Context, tailscaleClient tsClient, logger *zap.SugaredLogger, stsC *tailscaleSTSConfig, hsvc *corev1.Service) ([]string, error) {
+func (a *tailscaleSTSReconciler) provisionSecrets(ctx context.Context, tailscaleClient tsClient, loginUrl string, stsC *tailscaleSTSConfig, hsvc *corev1.Service, logger *zap.SugaredLogger) ([]string, error) {
 	secretNames := make([]string, stsC.Replicas)
 
 	// Start by ensuring we have Secrets for the desired number of replicas. This will handle both creating and scaling
@@ -434,7 +452,7 @@ func (a *tailscaleSTSReconciler) provisionSecrets(ctx context.Context, tailscale
 			}
 		}
 
-		configs, err := tailscaledConfig(stsC, authKey, orig, hostname)
+		configs, err := tailscaledConfig(stsC, loginUrl, authKey, orig, hostname)
 		if err != nil {
 			return nil, fmt.Errorf("error creating tailscaled config: %w", err)
 		}
@@ -1063,7 +1081,7 @@ func isMainContainer(c *corev1.Container) bool {
 
 // tailscaledConfig takes a proxy config, a newly generated auth key if generated and a Secret with the previous proxy
 // state and auth key and returns tailscaled config files for currently supported proxy versions.
-func tailscaledConfig(stsC *tailscaleSTSConfig, newAuthkey string, oldSecret *corev1.Secret, hostname string) (tailscaledConfigs, error) {
+func tailscaledConfig(stsC *tailscaleSTSConfig, loginUrl string, newAuthkey string, oldSecret *corev1.Secret, hostname string) (tailscaledConfigs, error) {
 	conf := &ipn.ConfigVAlpha{
 		Version:             "alpha0",
 		AcceptDNS:           "false",
@@ -1102,6 +1120,10 @@ func tailscaledConfig(stsC *tailscaleSTSConfig, newAuthkey string, oldSecret *co
 		conf.AuthKey = key
 	}
 
+	if loginUrl != "" {
+		conf.ServerURL = ptr.To(loginUrl)
+	}
+
 	capVerConfigs := make(map[tailcfg.CapabilityVersion]ipn.ConfigVAlpha)
 	capVerConfigs[107] = *conf
 
diff --git a/cmd/k8s-operator/tailnet.go b/cmd/k8s-operator/tailnet.go
index a9343272f..692cf0e10 100644
--- a/cmd/k8s-operator/tailnet.go
+++ b/cmd/k8s-operator/tailnet.go
@@ -21,19 +21,19 @@ import (
 	tsapi "tailscale.com/k8s-operator/apis/v1alpha1"
 )
 
-func clientForTailnet(ctx context.Context, cl client.Client, namespace, name string) (tsClient, error) {
+func clientForTailnet(ctx context.Context, cl client.Client, namespace, name string) (tsClient, string, error) {
 	var tn tsapi.Tailnet
 	if err := cl.Get(ctx, client.ObjectKey{Name: name}, &tn); err != nil {
-		return nil, fmt.Errorf("failed to get tailnet %q: %w", name, err)
+		return nil, "", fmt.Errorf("failed to get tailnet %q: %w", name, err)
 	}
 
 	if !operatorutils.TailnetIsReady(&tn) {
-		return nil, fmt.Errorf("tailnet %q is not ready", name)
+		return nil, "", fmt.Errorf("tailnet %q is not ready", name)
 	}
 
 	var secret corev1.Secret
 	if err := cl.Get(ctx, client.ObjectKey{Name: tn.Spec.Credentials.SecretName, Namespace: namespace}, &secret); err != nil {
-		return nil, fmt.Errorf("failed to get Secret %q in namespace %q: %w", tn.Spec.Credentials.SecretName, namespace, err)
+		return nil, "", fmt.Errorf("failed to get Secret %q in namespace %q: %w", tn.Spec.Credentials.SecretName, namespace, err)
 	}
 
 	baseURL := ipn.DefaultControlURL
@@ -55,7 +55,7 @@ func clientForTailnet(ctx context.Context, cl client.Client, namespace, name str
 	ts.HTTPClient = httpClient
 	ts.BaseURL = baseURL
 
-	return ts, nil
+	return ts, baseURL, nil
 }
 
 func clientFromProxyGroup(ctx context.Context, cl client.Client, obj client.Object, namespace string, def tsClient) (tsClient, error) {
@@ -73,7 +73,7 @@ func clientFromProxyGroup(ctx context.Context, cl client.Client, obj client.Obje
 		return def, nil
 	}
 
-	tailscaleClient, err := clientForTailnet(ctx, cl, namespace, pg.Spec.Tailnet)
+	tailscaleClient, _, err := clientForTailnet(ctx, cl, namespace, pg.Spec.Tailnet)
 	if err != nil {
 		return nil, err
 	}
diff --git a/cmd/k8s-operator/tsrecorder.go b/cmd/k8s-operator/tsrecorder.go
index 60ed24a70..777188b13 100644
--- a/cmd/k8s-operator/tsrecorder.go
+++ b/cmd/k8s-operator/tsrecorder.go
@@ -99,14 +99,9 @@ func (r *RecorderReconciler) Reconcile(ctx context.Context, req reconcile.Reques
 		return reconcile.Result{}, nil
 	}
 
-	tailscaleClient := r.tsClient
-	if tsr.Spec.Tailnet != "" {
-		tc, err := clientForTailnet(ctx, r.Client, r.tsNamespace, tsr.Spec.Tailnet)
-		if err != nil {
-			return setStatusReady(tsr, metav1.ConditionFalse, reasonRecorderTailnetUnavailable, err.Error())
-		}
-
-		tailscaleClient = tc
+	tailscaleClient, loginUrl, err := r.getClientAndLoginURL(ctx, tsr.Spec.Tailnet)
+	if err != nil {
+		return setStatusReady(tsr, metav1.ConditionFalse, reasonRecorderTailnetUnavailable, err.Error())
 	}
 
 	if markedForDeletion(tsr) {
@@ -149,7 +144,7 @@ func (r *RecorderReconciler) Reconcile(ctx context.Context, req reconcile.Reques
 		return setStatusReady(tsr, metav1.ConditionFalse, reasonRecorderInvalid, message)
 	}
 
-	if err = r.maybeProvision(ctx, tailscaleClient, tsr); err != nil {
+	if err = r.maybeProvision(ctx, tailscaleClient, loginUrl, tsr); err != nil {
 		reason := reasonRecorderCreationFailed
 		message := fmt.Sprintf("failed creating Recorder: %s", err)
 		if strings.Contains(err.Error(), optimisticLockErrorMsg) {
@@ -167,7 +162,30 @@ func (r *RecorderReconciler) Reconcile(ctx context.Context, req reconcile.Reques
 	return setStatusReady(tsr, metav1.ConditionTrue, reasonRecorderCreated, reasonRecorderCreated)
 }
 
-func (r *RecorderReconciler) maybeProvision(ctx context.Context, tailscaleClient tsClient, tsr *tsapi.Recorder) error {
+// getClientAndLoginURL returns the appropriate Tailscale client and resolved login URL
+// for the given tailnet name. If no tailnet is specified, returns the default client
+// and login server. Applies fallback to the operator's login server if the tailnet
+// doesn't specify a custom login URL.
+func (r *RecorderReconciler) getClientAndLoginURL(ctx context.Context, tailnetName string) (tsClient,
+	string, error) {
+	if tailnetName == "" {
+		return r.tsClient, r.loginServer, nil
+	}
+
+	tc, loginUrl, err := clientForTailnet(ctx, r.Client, r.tsNamespace, tailnetName)
+	if err != nil {
+		return nil, "", err
+	}
+
+	// Apply fallback if tailnet doesn't specify custom login URL
+	if loginUrl == "" {
+		loginUrl = r.loginServer
+	}
+
+	return tc, loginUrl, nil
+}
+
+func (r *RecorderReconciler) maybeProvision(ctx context.Context, tailscaleClient tsClient, loginUrl string, tsr *tsapi.Recorder) error {
 	logger := r.logger(tsr.Name)
 
 	r.mu.Lock()
@@ -234,7 +252,11 @@ func (r *RecorderReconciler) maybeProvision(ctx context.Context, tailscaleClient
 		return fmt.Errorf("error creating RoleBinding: %w", err)
 	}
 
-	ss := tsrStatefulSet(tsr, r.tsNamespace, r.loginServer)
+	if r.loginServer != "" && loginUrl == "" {
+		loginUrl = r.loginServer
+	}
+
+	ss := tsrStatefulSet(tsr, r.tsNamespace, loginUrl)
 	_, err = createOrUpdate(ctx, r.Client, r.tsNamespace, ss, func(s *appsv1.StatefulSet) {
 		s.ObjectMeta.Labels = ss.ObjectMeta.Labels
 		s.ObjectMeta.Annotations = ss.ObjectMeta.Annotations