Fix host resolution and concurrent cache updates

cmd/netssh/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
@@ -118,9 +119,13 @@ func runSSHWrapper(args []string) {
 	// If a stale match is found, use the canonical name for the NetBox re-fetch below.
 	lookupHost := parsed.Host
 	shortcutsEnabled := len(cfg.Shortcuts.Domains) > 0 || cfg.Shortcuts.StripHyphens
-	if shortcutsEnabled {
+	if cached, _ := c.Get(lookupHost); shortcutsEnabled && cached.Name == "" {
 		normalize := shortcuts.MakeNormalizer(cfg.Shortcuts)
-		if entry, found, fresh := c.GetByShortcut(lookupHost, normalize); found {
+		entry, found, fresh, err := c.GetByShortcut(lookupHost, normalize)
+		if err != nil {
+			fatalf("%v", err)
+		}
+		if found {
 			if fresh {
 				c.MarkUsed(entry.Name)
 				_ = c.Save()
@@ -140,42 +145,9 @@ func runSSHWrapper(args []string) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	entries, err := nbClient.Search(ctx, lookupHost, netbox.SearchOptions{})
+	target, err := findSSHHost(ctx, nbClient, lookupHost, cfg.Shortcuts.Domains)
 	if err != nil {
-		fatalf("NetBox search failed: %v", err)
-	}
-
-	var target *netbox.HostEntry
-	for i, e := range entries {
-		if strings.EqualFold(e.Name, lookupHost) {
-			target = &entries[i]
-			break
-		}
-	}
-
-	// Domain expansion: if still no match and the input has no dots, try appending
-	// each configured domain suffix and re-querying NetBox.
-	if target == nil && !strings.Contains(parsed.Host, ".") && len(cfg.Shortcuts.Domains) > 0 {
-		for _, domain := range cfg.Shortcuts.Domains {
-			expanded := parsed.Host + domain
-			expandedEntries, searchErr := nbClient.Search(ctx, expanded, netbox.SearchOptions{})
-			if searchErr != nil {
-				continue
-			}
-			for i, e := range expandedEntries {
-				if strings.EqualFold(e.Name, expanded) {
-					target = &expandedEntries[i]
-					break
-				}
-			}
-			if target != nil {
-				break
-			}
-		}
-	}
-
-	if target == nil {
-		fatalf("host %q not found in NetBox", parsed.Host)
+		fatalf("NetBox: %v", err)
 	}
 
 	chain, err := resolver.New(cfg.Resolver)
@@ -183,7 +155,7 @@ func runSSHWrapper(args []string) {
 		fatalf("resolver: %v", err)
 	}
 
-	ip, err := chain.Resolve(ctx, target, nbClient)
+	ip, err := chain.Resolve(ctx, &target, nbClient)
 	if err != nil {
 		fatalf("IP resolution for %q: %v", parsed.Host, err)
 	}
@@ -195,6 +167,26 @@ func runSSHWrapper(args []string) {
 	connect(ip, parsed, args)
 }
 
+// findSSHHost uses the paginated, ambiguity-checking lookup for connection
+// targets. Only a genuine missing host permits trying configured suffixes.
+func findSSHHost(ctx context.Context, client *netbox.Client, name string, domains []string) (netbox.HostEntry, error) {
+	target, err := client.FindHost(ctx, name, "")
+	if !errors.Is(err, netbox.ErrHostNotFound) || strings.Contains(name, ".") {
+		return target, err
+	}
+	for _, domain := range domains {
+		domain = strings.TrimPrefix(strings.TrimSpace(domain), ".")
+		if domain == "" {
+			continue
+		}
+		target, err = client.FindHost(ctx, name+"."+domain, "")
+		if !errors.Is(err, netbox.ErrHostNotFound) {
+			return target, err
+		}
+	}
+	return netbox.HostEntry{}, fmt.Errorf("%w: %q", netbox.ErrHostNotFound, name)
+}
+
 func connect(ip string, parsed *internalssh.ParsedArgs, originalArgs []string) {
 	newArgs := internalssh.ReplaceHost(originalArgs, parsed.DestIdx, ip)
 	if err := internalssh.Exec(newArgs); err != nil {
@@ -233,7 +225,7 @@ func runTUI() {
 	}
 
 	// Revalidate stale or live TUI results through the configured resolver chain.
-	if cached, fresh := c.Get(host.Name); fresh {
+	if cached, fresh := c.Get(host.Name); fresh && cached.Kind == host.Kind {
 		host.IP = cached.IP
 	} else {
 		if nbClient == nil {
@@ -533,26 +525,31 @@ func cacheRefreshCmd() *cobra.Command {
 				return fmt.Errorf("resolver: %w", err)
 			}
 			written := 0
+			nameCounts := make(map[string]int)
+			for _, e := range entries {
+				nameCounts[strings.ToLower(e.Name)]++
+			}
 			for i := range entries {
 				e := &entries[i]
+				if nameCounts[strings.ToLower(e.Name)] > 1 {
+					c.Delete(e.Name)
+					fmt.Fprintf(os.Stderr, "warning: skipping ambiguous hostname %q\n", e.Name)
+					continue
+				}
 
-				var ip string
-				if chain != nil {
-					if resolved, err := chain.Resolve(ctx, e, nbClient); err == nil {
-						ip = resolved
+				ip, err := chain.Resolve(ctx, e, nbClient)
+				if err != nil {
+					// Do not turn a failed configured strategy into a fresh
+					// primary-IP cache hit that bypasses the resolver.
+					if ctx.Err() != nil {
+						return ctx.Err()
 					}
+					c.Delete(e.Name)
+					fmt.Fprintf(os.Stderr, "warning: skipping %s: %v\n", e.Name, err)
+					continue
 				}
-				if ip == "" {
-					ip = e.PrimaryIP4
-				}
-				if ip == "" {
-					ip = e.PrimaryIP6
-				}
-
-				if ip != "" {
-					c.Upsert(cache.Entry{Name: e.Name, IP: ip, Kind: e.Kind, Tags: e.Tags})
-					written++
-				}
+				c.Upsert(cache.Entry{Name: e.Name, IP: ip, Kind: e.Kind, Tags: e.Tags})
+				written++
 			}
 
 			if err := c.Save(); err != nil {