pkg/hostagent: Use HostAgent.DialContextToGuestIP() if the IP is accessible directly.

Signed-off-by: Norio Nomura <[email protected]>

Author: norio-nomura

pkg/hostagent/hostagent.go

@@ -514,23 +514,28 @@ func (a *HostAgent) startRoutinesAndWait(ctx context.Context, errCh <-chan error
 	return a.driver.Stop(ctx)
 }
 
-func (a *HostAgent) Info(_ context.Context) (*hostagentapi.Info, error) {
+// GuestIP returns the guest IP address, or nil if not accessible by direct IP.
+func (a *HostAgent) GuestIP() net.IP {
 	a.guestIPMu.RLock()
 	defer a.guestIPMu.RUnlock()
+	return a.guestIP
+}
+
+func (a *HostAgent) Info(_ context.Context) (*hostagentapi.Info, error) {
+	guestIP := a.GuestIP()
 	info := &hostagentapi.Info{
-		GuestIP:      a.guestIP,
+		GuestIP:      guestIP,
 		SSHLocalPort: a.sshLocalPort,
 	}
 	return info, nil
 }
 
 func (a *HostAgent) sshAddressPort() (sshAddress string, sshPort int) {
-	a.guestIPMu.RLock()
-	defer a.guestIPMu.RUnlock()
 	sshAddress = a.instSSHAddress
 	sshPort = a.sshLocalPort
-	if a.guestIP != nil {
-		sshAddress = a.guestIP.String()
+	guestIP := a.GuestIP()
+	if guestIP != nil {
+		sshAddress = guestIP.String()
 		sshPort = 22
 		logrus.Debugf("Using the guest IP address %q directly", sshAddress)
 	}
@@ -891,9 +896,17 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 		if useSSHFwd {
 			a.portForwarder.OnEvent(ctx, ev)
 		} else {
-			dialContext := a.DialContextToGuestIP()
-			if dialContext == nil {
-				dialContext = portfwd.DialContextToGRPCTunnel(client)
+			dialContext := func(ctx context.Context, network, address string) (net.Conn, error) {
+				dialContext := a.DialContextToGuestIP()
+				if dialContext == nil {
+					return portfwd.DialContextToGRPCTunnel(client)(ctx, network, address)
+				}
+				if host, _, err := net.SplitHostPort(address); err != nil {
+					return nil, err
+				} else if ip := net.ParseIP(host); ip.IsUnspecified() || ip.Equal(a.GuestIP()) {
+					return a.DialContextToGuestIP()(ctx, network, address)
+				}
+				return portfwd.DialContextToGRPCTunnel(client)(ctx, network, address)
 			}
 			a.grpcPortForwarder.OnEvent(ctx, dialContext, ev)
 		}
@@ -911,9 +924,8 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 // DialContextToGuestIP returns a DialContext function that connects to the guest IP directly.
 // If the guest IP is not known, it returns nil.
 func (a *HostAgent) DialContextToGuestIP() func(ctx context.Context, network, address string) (net.Conn, error) {
-	a.guestIPMu.RLock()
-	defer a.guestIPMu.RUnlock()
-	if a.guestIP == nil {
+	guestIP := a.GuestIP()
+	if guestIP == nil {
 		return nil
 	}
 	return func(ctx context.Context, network, address string) (net.Conn, error) {
@@ -924,7 +936,7 @@ func (a *HostAgent) DialContextToGuestIP() func(ctx context.Context, network, ad
 		}
 		// Host part of address is ignored, because it already has been checked by forwarding rules
 		// and we want to connect to the guest IP directly.
-		return d.DialContext(ctx, network, net.JoinHostPort(a.guestIP.String(), port))
+		return d.DialContext(ctx, network, net.JoinHostPort(guestIP.String(), port))
 	}
 }