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())
+			}
+		})
+	}
+}