Fix host resolution and concurrent cache updates

internal/cache/cache.go

@@ -2,13 +2,17 @@ package cache
 
 import (
 	"encoding/json"
-	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/storage"
+	"fmt"
 	"os"
 	"path/filepath"
 	"sort"
 	"strings"
 	"sync"
 	"time"
+
+	"golang.org/x/sys/unix"
+
+	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/storage"
 )
 
 type Entry struct {
@@ -20,8 +24,18 @@ type Entry struct {
 	LastUsed time.Time `json:"last_used,omitempty"`
 }
 
+type changeKind uint8
+
+const (
+	upsertChange changeKind = iota
+	deleteChange
+	usageChange
+)
+
 type Cache struct {
 	mu           sync.RWMutex
+	dirty        map[string]changeKind
+	cleared      bool
 	entries      map[string]Entry
 	path         string
 	ttl          time.Duration
@@ -39,6 +53,7 @@ func New(path string, ttlSeconds int) *Cache {
 	}
 	return &Cache{
 		entries:      make(map[string]Entry),
+		dirty:        make(map[string]changeKind),
 		path:         path,
 		ttl:          time.Duration(ttlSeconds) * time.Second,
 		refreshStamp: stamp,
@@ -94,6 +109,8 @@ func (c *Cache) Load() error {
 		return err
 	}
 
+	c.dirty = make(map[string]changeKind)
+	c.cleared = false
 	c.entries = make(map[string]Entry, len(df.Entries))
 	for _, e := range df.Entries {
 		c.entries[e.Name] = e
@@ -101,23 +118,80 @@ func (c *Cache) Load() error {
 	return nil
 }
 
+// Save merges only local changes under a process-wide file lock. Atomic rename
+// alone prevents torn JSON but cannot prevent lost updates from another process.
 func (c *Cache) Save() error {
-	c.mu.RLock()
-	df := diskFormat{Entries: make([]Entry, 0, len(c.entries))}
-	for _, e := range c.entries {
-		df.Entries = append(df.Entries, e)
-	}
-	c.mu.RUnlock()
-
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	if err := os.MkdirAll(filepath.Dir(c.path), 0o755); err != nil {
 		return err
 	}
+	lock, err := os.OpenFile(c.path+".lock", os.O_CREATE|os.O_RDWR, 0o600)
+	if err != nil {
+		return err
+	}
+	defer lock.Close()
+	if err := unix.Flock(int(lock.Fd()), unix.LOCK_EX); err != nil {
+		return err
+	}
+	defer unix.Flock(int(lock.Fd()), unix.LOCK_UN)
 
+	merged := make(map[string]Entry)
+	if !c.cleared {
+		data, err := os.ReadFile(c.path)
+		if err != nil && !os.IsNotExist(err) {
+			return err
+		}
+		if err == nil {
+			var df diskFormat
+			if err := json.Unmarshal(data, &df); err != nil {
+				return err
+			}
+			for _, e := range df.Entries {
+				merged[e.Name] = e
+			}
+		}
+	}
+	for name, change := range c.dirty {
+		if change == deleteChange {
+			delete(merged, name)
+			continue
+		}
+		e := c.entries[name]
+		old, exists := merged[name]
+		if change == usageChange && !exists {
+			// A connection finishing after a refresh or clear must not
+			// restore an entry that the other process removed.
+			continue
+		}
+		if exists {
+			lastUsed := e.LastUsed
+			if change == usageChange || old.CachedAt.After(e.CachedAt) {
+				e = old
+			}
+			if old.LastUsed.After(lastUsed) {
+				lastUsed = old.LastUsed
+			}
+			e.LastUsed = lastUsed
+		}
+		merged[name] = e
+	}
+	df := diskFormat{Entries: make([]Entry, 0, len(merged))}
+	for _, e := range merged {
+		df.Entries = append(df.Entries, e)
+	}
+	sort.Slice(df.Entries, func(i, j int) bool { return df.Entries[i].Name < df.Entries[j].Name })
 	data, err := json.MarshalIndent(df, "", "  ")
 	if err != nil {
 		return err
 	}
-	return storage.WriteFile(c.path, data, 0o600)
+	if err := storage.WriteFile(c.path, data, 0o600); err != nil {
+		return err
+	}
+	c.entries = merged
+	c.dirty = make(map[string]changeKind)
+	c.cleared = false
+	return nil
 }
 
 func (c *Cache) Upsert(e Entry) {
@@ -127,9 +201,18 @@ func (c *Cache) Upsert(e Entry) {
 		e.LastUsed = old.LastUsed
 	}
 	c.entries[e.Name] = e
+	c.dirty[e.Name] = upsertChange
 	c.mu.Unlock()
 }
 
+// Delete removes an entry when a refresh can no longer resolve it safely.
+func (c *Cache) Delete(name string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	delete(c.entries, name)
+	c.dirty[name] = deleteChange
+}
+
 // MarkUsed records the current time as LastUsed for the named entry.
 // It is a no-op if the entry does not exist.
 func (c *Cache) MarkUsed(name string) {
@@ -138,6 +221,9 @@ func (c *Cache) MarkUsed(name string) {
 	if e, ok := c.entries[name]; ok {
 		e.LastUsed = time.Now()
 		c.entries[name] = e
+		if _, changed := c.dirty[name]; !changed {
+			c.dirty[name] = usageChange
+		}
 	}
 }
 
@@ -179,19 +265,20 @@ func (c *Cache) Search(prefix string) []Entry {
 	return out
 }
 
-// GetByShortcut scans all entries and returns the first whose normalized name matches
-// the normalized shortcut. Returns (entry, found, fresh).
-func (c *Cache) GetByShortcut(shortcut string, normalize func(string) string) (entry Entry, found bool, fresh bool) {
+// GetByShortcut rejects collisions instead of choosing a random map entry.
+func (c *Cache) GetByShortcut(shortcut string, normalize func(string) string) (entry Entry, found bool, fresh bool, err error) {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 	norm := normalize(shortcut)
 	for _, e := range c.entries {
 		if normalize(e.Name) == norm {
-			isFresh := c.ttl > 0 && time.Since(e.CachedAt) < c.ttl
-			return e, true, isFresh
+			if found {
+				return Entry{}, false, false, fmt.Errorf("ambiguous shortcut %q; use a full hostname", shortcut)
+			}
+			entry, found = e, true
 		}
 	}
-	return Entry{}, false, false
+	return entry, found, found && c.ttl > 0 && time.Since(entry.CachedAt) < c.ttl, nil
 }
 
 // Get returns an entry and reports whether it is still within the TTL.
@@ -212,6 +299,8 @@ func (c *Cache) Get(name string) (entry Entry, fresh bool) {
 func (c *Cache) Clear() {
 	c.mu.Lock()
 	c.entries = make(map[string]Entry)
+	c.dirty = make(map[string]changeKind)
+	c.cleared = true
 	c.mu.Unlock()
 }
 

internal/cache/cache_test.go

@@ -314,7 +314,7 @@ func TestGetByShortcut_MatchFresh(t *testing.T) {
 		return s
 	}
 
-	e, found, fresh := c.GetByShortcut("web01", normalize)
+	e, found, fresh, _ := c.GetByShortcut("web01", normalize)
 	if !found {
 		t.Fatal("expected entry to be found")
 	}
@@ -340,7 +340,7 @@ func TestGetByShortcut_MatchStale(t *testing.T) {
 		return s
 	}
 
-	_, found, fresh := c.GetByShortcut("web01", normalize)
+	_, found, fresh, _ := c.GetByShortcut("web01", normalize)
 	if !found {
 		t.Fatal("expected entry to be found even when stale")
 	}
@@ -353,7 +353,7 @@ func TestGetByShortcut_NotFound(t *testing.T) {
 	c := New("", 3600)
 	c.Upsert(Entry{Name: "db01.example.com", IP: "10.0.0.2", Kind: "device"})
 
-	_, found, _ := c.GetByShortcut("web01", func(s string) string { return s })
+	_, found, _, _ := c.GetByShortcut("web01", func(s string) string { return s })
 	if found {
 		t.Error("should not find an entry that does not match the shortcut")
 	}
@@ -361,7 +361,7 @@ func TestGetByShortcut_NotFound(t *testing.T) {
 
 func TestGetByShortcut_EmptyCache(t *testing.T) {
 	c := New("", 3600)
-	_, found, _ := c.GetByShortcut("web01", func(s string) string { return s })
+	_, found, _, _ := c.GetByShortcut("web01", func(s string) string { return s })
 	if found {
 		t.Error("empty cache should return found=false")
 	}
@@ -381,7 +381,7 @@ func TestGetByShortcut_MultiDomain(t *testing.T) {
 		return s
 	}
 
-	e, found, _ := c.GetByShortcut("web01", normalize)
+	e, found, _, _ := c.GetByShortcut("web01", normalize)
 	if !found {
 		t.Fatal("should match entry with second configured domain")
 	}

internal/cache/concurrency_test.go

@@ -0,0 +1,182 @@
+package cache
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestShortcutRejectsAmbiguousTargets(t *testing.T) {
+	c := New("", 3600)
+	c.Upsert(Entry{Name: "web.prod.example", IP: "10.0.0.1"})
+	c.Upsert(Entry{Name: "web.dev.example", IP: "10.0.0.2"})
+	normalize := func(s string) string { return strings.SplitN(s, ".", 2)[0] }
+	for range 100 {
+		entry, found, fresh, err := c.GetByShortcut("web", normalize)
+		if err == nil || found || fresh || entry.Name != "" {
+			t.Fatalf("ambiguous shortcut returned target: %+v, %v, %v, %v", entry, found, fresh, err)
+		}
+	}
+}
+
+func TestSaveMergesIndependentlyLoadedCaches(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	a, b := New(path, 3600), New(path, 3600)
+	for _, c := range []*Cache{a, b} {
+		if err := c.Load(); err != nil {
+			t.Fatal(err)
+		}
+	}
+	a.Upsert(Entry{Name: "a", IP: "10.0.0.1"})
+	b.Upsert(Entry{Name: "b", IP: "10.0.0.2"})
+	for _, c := range []*Cache{a, b} {
+		if err := c.Save(); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := a.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(a.All()) != 2 {
+		t.Fatalf("lost an update: %+v", a.All())
+	}
+
+	// Saving a stale usage update must retain newer IPs and both usage times.
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	b.MarkUsed("a")
+	used, _ := b.Get("a")
+	a.Upsert(Entry{Name: "a", IP: "10.0.0.3"})
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	got, _ := b.Get("a")
+	if got.IP != "10.0.0.3" || !got.LastUsed.Equal(used.LastUsed) {
+		t.Fatalf("lost refreshed IP or usage: %+v", got)
+	}
+
+	b.Delete("a")
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := a.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if got, _ := a.Get("a"); got.Name != "" {
+		t.Fatalf("deleted entry survived: %+v", got)
+	}
+	a.Clear()
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(b.All()) != 0 {
+		t.Fatal("clear merged old entries back in")
+	}
+}
+
+func TestSaveAcrossProcesses(t *testing.T) {
+	if path := os.Getenv("NETSSH_TEST_CACHE_WRITER_PATH"); path != "" {
+		c := New(path, 3600)
+		if err := c.Load(); err != nil {
+			t.Fatal(err)
+		}
+		c.Upsert(Entry{Name: os.Getenv("NETSSH_TEST_CACHE_WRITER_NAME"), IP: "10.0.0.1"})
+		if err := c.Save(); err != nil {
+			t.Fatal(err)
+		}
+		return
+	}
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	const writers = 8
+	cmds := make([]*exec.Cmd, writers)
+	for i := range cmds {
+		cmd := exec.Command(os.Args[0], "-test.run=^TestSaveAcrossProcesses$")
+		cmd.Env = append(os.Environ(), "NETSSH_TEST_CACHE_WRITER_PATH="+path, fmt.Sprintf("NETSSH_TEST_CACHE_WRITER_NAME=host-%d", i))
+		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
+		if err := cmd.Start(); err != nil {
+			t.Fatal(err)
+		}
+		cmds[i] = cmd
+	}
+	for _, cmd := range cmds {
+		if err := cmd.Wait(); err != nil {
+			t.Error(err)
+		}
+	}
+	c := New(path, 3600)
+	if err := c.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(c.All()) != writers {
+		t.Fatalf("lost concurrent updates: got %d, want %d", len(c.All()), writers)
+	}
+}
+
+func TestSavePreservesNewerUsageFromDisk(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	a := New(path, 3600)
+	a.Upsert(Entry{Name: "web", IP: "10.0.0.1", LastUsed: time.Now()})
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	b := New(path, 3600)
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	a.MarkUsed("web")
+	used, _ := a.Get("web")
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	b.Upsert(Entry{Name: "web", IP: "10.0.0.2"})
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	got, _ := b.Get("web")
+	if got.IP != "10.0.0.2" || !got.LastUsed.Equal(used.LastUsed) {
+		t.Fatalf("lost data: %+v", got)
+	}
+}
+
+func TestUsageSaveDoesNotRestoreRemovedHosts(t *testing.T) {
+	for _, clear := range []bool{false, true} {
+		t.Run(fmt.Sprintf("clear=%v", clear), func(t *testing.T) {
+			path := filepath.Join(t.TempDir(), "hosts.json")
+			a := New(path, 3600)
+			a.Upsert(Entry{Name: "web", IP: "10.0.0.1"})
+			if err := a.Save(); err != nil {
+				t.Fatal(err)
+			}
+			b := New(path, 3600)
+			if err := b.Load(); err != nil {
+				t.Fatal(err)
+			}
+			b.MarkUsed("web")
+			if clear {
+				a.Clear()
+			} else {
+				a.Delete("web")
+			}
+			if err := a.Save(); err != nil {
+				t.Fatal(err)
+			}
+			if err := b.Save(); err != nil {
+				t.Fatal(err)
+			}
+			if len(b.All()) != 0 {
+				t.Fatalf("usage update restored removed entry: %+v", b.All())
+			}
+		})
+	}
+}

internal/netbox/client.go

@@ -351,6 +351,9 @@ func stripPrefix(cidr string) string {
 	return cidr
 }
 
+// ErrHostNotFound distinguishes a missing exact match from API failures or ambiguity.
+var ErrHostNotFound = errors.New("host not found")
+
 // FindHost refuses ambiguous names so writes cannot target an arbitrary server.
 func (c *Client) FindHost(ctx context.Context, name, kind string) (HostEntry, error) {
 	if kind != "" && kind != "device" && kind != "vm" {
@@ -366,6 +369,9 @@ func (c *Client) FindHost(ctx context.Context, name, kind string) (HostEntry, er
 			matches = append(matches, e)
 		}
 	}
+	if len(matches) == 0 {
+		return HostEntry{}, fmt.Errorf("%w: %q", ErrHostNotFound, name)
+	}
 	if len(matches) != 1 {
 		return HostEntry{}, fmt.Errorf("host %q: found %d exact matches; use --kind or a unique hostname", name, len(matches))
 	}

internal/ssh/args.go

@@ -4,7 +4,7 @@ import "strings"
 
 // flagsWithArg lists all SSH flags that consume the following argument.
 var flagsWithArg = map[byte]bool{
-	'b': true, 'c': true, 'D': true, 'E': true, 'e': true,
+	'B': true, 'P': true, 'b': true, 'c': true, 'D': true, 'E': true, 'e': true,
 	'F': true, 'I': true, 'i': true, 'J': true, 'L': true,
 	'l': true, 'm': true, 'o': true, 'O': true, 'p': true,
 	'Q': true, 'R': true, 'S': true, 'w': true, 'W': true,
@@ -32,18 +32,16 @@ func Parse(args []string) *ParsedArgs {
 		}
 
 		if strings.HasPrefix(arg, "-") && len(arg) > 1 {
-			flag := arg[1]
-			if flagsWithArg[flag] {
-				if len(arg) > 2 {
-					// argument is attached, e.g. -p2222
-					i++
-				} else {
-					// argument is the next element, e.g. -p 2222
-					i += 2
+			// OpenSSH accepts clusters such as -vp 2222 and -vladmin.
+			for j := 1; j < len(arg); j++ {
+				if flagsWithArg[arg[j]] {
+					if j == len(arg)-1 {
+						i++
+					}
+					break
 				}
-			} else {
-				i++
 			}
+			i++
 			continue
 		}
 
@@ -86,35 +84,41 @@ func ReplaceHost(args []string, destIdx int, newHost string) []string {
 	return result
 }
 
-// HasPortFlag reports whether a port was specified via -p in args.
-func HasPortFlag(args []string) bool {
-	if parsed := Parse(args); parsed != nil {
-		args = args[:parsed.DestIdx]
-	}
-	for i, a := range args {
-		if a == "-p" && i+1 < len(args) {
-			return true
-		}
-		if len(a) > 2 && a[0] == '-' && a[1] == 'p' {
-			return true
-		}
-	}
-	return false
-}
+// HasPortFlag reports whether a port was explicitly supplied.
+func HasPortFlag(args []string) bool { return hasOption(args, 'p', "Port") }
 
-// HasUserFlag reports whether a user was specified via -l in args.
-// Used to avoid overriding an explicit -l with the configured default user.
-func HasUserFlag(args []string) bool {
-	if parsed := Parse(args); parsed != nil {
-		args = args[:parsed.DestIdx]
-	}
-	for i, a := range args {
-		if a == "-l" && i+1 < len(args) {
-			return true
+// HasUserFlag reports whether a login user was explicitly supplied.
+func HasUserFlag(args []string) bool { return hasOption(args, 'l', "User") }
+
+func hasOption(args []string, wanted byte, keyword string) bool {
+	for i := 0; i < len(args); i++ {
+		arg := args[i]
+		if arg == "--" || len(arg) < 2 || arg[0] != '-' {
+			break
 		}
-		// handle attached form: -lroot
-		if len(a) > 2 && a[0] == '-' && a[1] == 'l' {
-			return true
+		for j := 1; j < len(arg); j++ {
+			flag := arg[j]
+			if !flagsWithArg[flag] {
+				continue
+			}
+			value := arg[j+1:]
+			if value == "" {
+				i++
+				if i >= len(args) {
+					return false
+				}
+				value = args[i]
+			}
+			if flag == wanted {
+				return true
+			}
+			if flag == 'o' {
+				fields := strings.FieldsFunc(value, func(r rune) bool { return r == '=' || r == ' ' || r == '\t' })
+				if len(fields) >= 2 && strings.EqualFold(fields[0], keyword) {
+					return true
+				}
+			}
+			break
 		}
 	}
 	return false

internal/ssh/args_regression_test.go

@@ -0,0 +1,46 @@
+package ssh
+
+import "testing"
+
+func TestParseOptionClusters(t *testing.T) {
+	for _, args := range [][]string{
+		{"-vp", "2222", "server"},
+		{"-vp2222", "server"},
+		{"-vAl", "admin", "server"},
+		{"-vladmin", "server"},
+		{"-B", "eth0", "server"},
+		{"-P", "production", "server"},
+	} {
+		t.Run(args[0], func(t *testing.T) {
+			assertParsed(t, Parse(args), "server", "", len(args)-1)
+		})
+	}
+}
+
+func TestExplicitSSHOptions(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		args       []string
+		user, port bool
+	}{
+		{"cluster port separated", []string{"-vp", "2222", "server"}, false, true},
+		{"cluster user attached", []string{"-vladmin", "server"}, true, false},
+		{"o separated", []string{"-o", "User=admin", "-o", "Port=2222", "server"}, true, true},
+		{"o attached and case insensitive", []string{"-ouSeR=admin", "-oport=2222", "server"}, true, true},
+		{"o whitespace", []string{"-o", " User = admin", "-o", "Port\t2222", "server"}, true, true},
+		{"cluster o", []string{"-voUser=admin", "server"}, true, false},
+		{"option value looks like flag", []string{"-i", "-lroot", "-F", "-p2222", "server"}, false, false},
+		{"remote flags", []string{"server", "-p2222", "-o", "User=admin"}, false, false},
+		{"after double dash", []string{"--", "-lroot", "-p2222"}, false, false},
+		{"missing value", []string{"-vp"}, false, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := HasUserFlag(tc.args); got != tc.user {
+				t.Errorf("HasUserFlag(%v) = %v, want %v", tc.args, got, tc.user)
+			}
+			if got := HasPortFlag(tc.args); got != tc.port {
+				t.Errorf("HasPortFlag(%v) = %v, want %v", tc.args, got, tc.port)
+			}
+		})
+	}
+}