fix handler related and some nit

Signed-off-by: KevinLiang10 <[email protected]>

ipn/ipnlocal/local.go

@@ -3442,11 +3442,13 @@ func (b *LocalBackend) setVIPServicesTCPPortsInterceptedLocked(svcPorts map[stri
 	svcAddrPorts := make(map[netip.Addr]func(uint16) bool)
 	// Only set the intercept function if the service has been assigned a VIP.
 	for svcName, ports := range svcPorts {
-		if addrs, ok := vipServiceIPMap[svcName]; ok {
-			interceptFn := generateInterceptTCPPortFunc(ports)
-			for _, addr := range addrs {
-				svcAddrPorts[addr] = interceptFn
-			}
+		addrs, ok := vipServiceIPMap[svcName]
+		if !ok {
+			continue
+		}
+		interceptFn := generateInterceptTCPPortFunc(ports)
+		for _, addr := range addrs {
+			svcAddrPorts[addr] = interceptFn
 		}
 	}
 
@@ -4214,7 +4216,7 @@ func (b *LocalBackend) TCPHandlerForDst(src, dst netip.AddrPort) (handler func(c
 		}
 	}
 
-	// TODO(corp#26001): Get handler for VIP services and Local IPs using
+	// TODO(tailscale/corp#26001): Get handler for VIP services and Local IPs using
 	// the same function.
 	if handler := b.tcpHandlerForVIPService(dst, src); handler != nil {
 		return handler, opts
@@ -6023,7 +6025,7 @@ func (b *LocalBackend) reloadServeConfigLocked(prefs ipn.PrefsView) {
 // b.mu must be held.
 func (b *LocalBackend) setTCPPortsInterceptedFromNetmapAndPrefsLocked(prefs ipn.PrefsView) {
 	handlePorts := make([]uint16, 0, 4)
-	vipServicesPorts := make(map[string][]uint16)
+	var vipServicesPorts map[string][]uint16
 
 	if prefs.Valid() && prefs.RunSSH() && envknob.CanSSHD() {
 		handlePorts = append(handlePorts, 22)
@@ -6055,9 +6057,9 @@ func (b *LocalBackend) setTCPPortsInterceptedFromNetmapAndPrefsLocked(prefs ipn.
 				}
 			}
 			if _, ok := vipServicesPorts[svc]; !ok {
-				vipServicesPorts[svc] = servicePorts
+				mak.Set(&vipServicesPorts, svc, servicePorts)
 			} else {
-				vipServicesPorts[svc] = append(vipServicesPorts[svc], servicePorts...)
+				mak.Set(&vipServicesPorts, svc, append(vipServicesPorts[svc], servicePorts...))
 			}
 		}
 

ipn/ipnlocal/serve.go

@@ -55,7 +55,7 @@ var serveHTTPContextKey ctxkey.Key[*serveHTTPContext]
 
 type serveHTTPContext struct {
 	SrcAddr       netip.AddrPort
-	ForVIPService bool
+	ForVIPService string // VIP service name, empty string means local
 	DestPort      uint16
 
 	// provides funnel-specific context, nil if not funneled
@@ -471,7 +471,7 @@ func (b *LocalBackend) tcpHandlerForVIPService(dstAddr, srcAddr netip.AddrPort)
 			BaseContext: func(_ net.Listener) context.Context {
 				return serveHTTPContextKey.WithValue(context.Background(), &serveHTTPContext{
 					SrcAddr:       srcAddr,
-					ForVIPService: true,
+					ForVIPService: dstSvc,
 					DestPort:      dport,
 				})
 			},
@@ -481,7 +481,7 @@ func (b *LocalBackend) tcpHandlerForVIPService(dstAddr, srcAddr netip.AddrPort)
 			// hostnames, but for services this getTLSServeCetForPort will need a version that also take
 			// in the hostname. How to store the TLS cert is still being discussed.
 			hs.TLSConfig = &tls.Config{
-				GetCertificate: b.getTLSServeCertForPort(dport, true),
+				GetCertificate: b.getTLSServeCertForPort(dport, dstSvc),
 			}
 			return func(c net.Conn) error {
 				return hs.ServeTLS(netutil.NewOneConnListener(c, nil), "", "")
@@ -568,7 +568,7 @@ func (b *LocalBackend) tcpHandlerForServe(dport uint16, srcAddr netip.AddrPort,
 		}
 		if tcph.HTTPS() {
 			hs.TLSConfig = &tls.Config{
-				GetCertificate: b.getTLSServeCertForPort(dport, false),
+				GetCertificate: b.getTLSServeCertForPort(dport, ""),
 			}
 			return func(c net.Conn) error {
 				return hs.ServeTLS(netutil.NewOneConnListener(c, nil), "", "")
@@ -1006,7 +1006,7 @@ func allNumeric(s string) bool {
 	return s != ""
 }
 
-func (b *LocalBackend) webServerConfig(hostname string, forVIPService bool, port uint16) (c ipn.WebServerConfigView, ok bool) {
+func (b *LocalBackend) webServerConfig(hostname string, forVIPService string, port uint16) (c ipn.WebServerConfigView, ok bool) {
 	key := ipn.HostPort(fmt.Sprintf("%s:%v", hostname, port))
 
 	b.mu.Lock()
@@ -1015,13 +1015,13 @@ func (b *LocalBackend) webServerConfig(hostname string, forVIPService bool, port
 	if !b.serveConfig.Valid() {
 		return c, false
 	}
-	if forVIPService {
-		return b.serveConfig.FindServiceWeb(key)
+	if forVIPService != "" {
+		return b.serveConfig.FindServiceWeb(forVIPService, key)
 	}
 	return b.serveConfig.FindWeb(key)
 }
 
-func (b *LocalBackend) getTLSServeCertForPort(port uint16, forVIPService bool) func(hi *tls.ClientHelloInfo) (*tls.Certificate, error) {
+func (b *LocalBackend) getTLSServeCertForPort(port uint16, forVIPService string) func(hi *tls.ClientHelloInfo) (*tls.Certificate, error) {
 	return func(hi *tls.ClientHelloInfo) (*tls.Certificate, error) {
 		if hi == nil || hi.ServerName == "" {
 			return nil, errors.New("no SNI ServerName")

ipn/serve.go

@@ -626,9 +626,9 @@ func (v ServeConfigView) FindServiceTCP(svcName string, port uint16) (res TCPPor
 	return svcCfg.TCP().GetOk(port)
 }
 
-func (v ServeConfigView) FindServiceWeb(hp HostPort) (res WebServerConfigView, ok bool) {
-	for _, service := range v.Services().All() {
-		if res, ok := service.Web().GetOk(hp); ok {
+func (v ServeConfigView) FindServiceWeb(svcName string, hp HostPort) (res WebServerConfigView, ok bool) {
+	if svcCfg, ok := v.Services().GetOk(svcName); ok {
+		if res, ok := svcCfg.Web().GetOk(hp); ok {
 			return res, ok
 		}
 	}

types/netmap/IPServiceMappings.go

@@ -1,22 +0,0 @@
-// Copyright (c) Tailscale Inc & AUTHORS
-// SPDX-License-Identifier: BSD-3-Clause
-
-package netmap
-
-import "net/netip"
-
-// IPServiceMappings maps IP addresses to service names. This is the inverse of
-// [ServiceIPMappings], and is used to inform clients which services is an VIP
-// address associated with. This is set to b.ipVIPServiceMap every time the
-// netmap is updated. This is used to reduce the cost for looking up the service
-// name for the dst IP address in the netStack packet processing workflow.
-//
-// This is of the form:
-//
-//	{
-//	  "100.65.32.1": "svc:samba",
-//	  "fd7a:115c:a1e0::1234": "svc:samba",
-//	  "100.102.42.3": "svc:web",
-//	  "fd7a:115c:a1e0::abcd": "svc:web",
-//	}
-type IPServiceMappings map[netip.Addr]string

types/netmap/netmap.go

@@ -105,7 +105,7 @@ func (nm *NetworkMap) GetAddresses() views.Slice[netip.Prefix] {
 // VIP addresses that correspond to the service. The service names are
 // with the prefix "svc:".
 //
-// TODO(corp##25997): cache the result of decoding the capmap so that
+// TODO(tailscale/corp##25997): cache the result of decoding the capmap so that
 // we don't have to decode it multiple times after each netmap update.
 func (nm *NetworkMap) GetVIPServiceIPMap() tailcfg.ServiceIPMappings {
 	if nm == nil {
@@ -425,3 +425,19 @@ const (
 	_ WGConfigFlags = 1 << iota
 	AllowSubnetRoutes
 )
+
+// IPServiceMappings maps IP addresses to service names. This is the inverse of
+// [ServiceIPMappings], and is used to inform clients which services is an VIP
+// address associated with. This is set to b.ipVIPServiceMap every time the
+// netmap is updated. This is used to reduce the cost for looking up the service
+// name for the dst IP address in the netStack packet processing workflow.
+//
+// This is of the form:
+//
+//	{
+//	  "100.65.32.1": "svc:samba",
+//	  "fd7a:115c:a1e0::1234": "svc:samba",
+//	  "100.102.42.3": "svc:web",
+//	  "fd7a:115c:a1e0::abcd": "svc:web",
+//	}
+type IPServiceMappings map[netip.Addr]string