least-latency regex and function to method

ngrok · Feb 22, 2024 · e3926df · e3926df
1 parent 90c99c9
commit e3926df
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/internal/tunnel/client/reconnecting.go b/internal/tunnel/client/reconnecting.go
@@ -199,7 +199,7 @@ func (s *reconnectingSession) listenTunnel(listen func(*session) (Tunnel, error)
 		}
 		// connect this tunnel to the other legs
 		for _, session := range s.sessions[1:] {
-			if e := reconnectTunnelToSession(session.raw, tun.(*tunnel), make(map[string]*tunnel), tun.ID()); e != nil {
+			if e := s.reconnectTunnelToSession(session.raw, tun.(*tunnel), make(map[string]*tunnel), tun.ID()); e != nil {
 				return nil, e
 			}
 			// use locking method
@@ -338,7 +338,7 @@ func (s *reconnectingSession) connect(acceptErr error, connSession *session) err
 		// reconnected tunnels, which may have different IDs
 		newTunnels := make(map[string]*tunnel, len(session.tunnels))
 		for oldID, t := range session.tunnels {
-			if err := reconnectTunnelToSession(raw, t, newTunnels, oldID); err != nil {
+			if err := s.reconnectTunnelToSession(raw, t, newTunnels, oldID); err != nil {
 				return err
 			}
 		}
@@ -406,7 +406,7 @@ func (s *reconnectingSession) connect(acceptErr error, connSession *session) err
 	}
 }
 
-func reconnectTunnelToSession(raw RawSession, t *tunnel, newTunnels map[string]*tunnel, oldID string) error {
+func (s *reconnectingSession) reconnectTunnelToSession(raw RawSession, t *tunnel, newTunnels map[string]*tunnel, oldID string) error {
 	// set the returned token for reconnection
 	tCfg := t.RemoteBindConfig()
 	t.bindExtra.Token = tCfg.Token

diff --git a/session.go b/session.go
@@ -79,6 +79,8 @@ var defaultCACert []byte
 
 const defaultServer = "connect.ngrok-agent.com:443"
 
+var leastLatencyServer = regexp.MustCompile(`^connect\.([a-z]+?-)?ngrok-agent\.com(\.lan)?:443`)
+
 // Dialer is the interface a custom connection dialer must implement for use
 // with the [WithDialer] option.
 type Dialer interface {
@@ -716,13 +718,10 @@ func Connect(ctx context.Context, opts ...ConnectOption) (Session, error) {
 
 		// store any connect server addresses for use in subsequent legs
 		if cfg.EnableMultiLeg && legNumber == 0 && len(resp.Extra.ConnectAddresses) > 1 {
-			testServerAddr := strings.Replace(cfg.ServerAddr, ".lan", "", 1)
-			testServerAddr = strings.Replace(testServerAddr, ".dev-", ".", 1)
-			testServerAddr = strings.Replace(testServerAddr, ".stage-", ".", 1)
 			overrideAdditionalServers := len(cfg.AdditionalServerAddrs) == 0
 			for i, ca := range resp.Extra.ConnectAddresses {
 				if i == 0 {
-					if testServerAddr == defaultServer {
+					if leastLatencyServer.MatchString(cfg.ServerAddr) {
 						// lock in the leg 0 region
 						logger.Debug("first leg using region", "region", resp.Extra.Region, "server", ca.ServerAddr)
 						cfg.ServerAddr = ca.ServerAddr