Sebi/ssh-netbox-wrapper
1
0 Fork 0
Code Issues 2 Pull Requests 10 Actions Releases 22 Activity

Fix host resolution and concurrent cache updates

Browse Source
This commit is contained in:
Sebastian Unterschütz committed 2026-10-27 22:23:59 +01:00
1 parent 92b1c40cb8
commit fa7446e5f7
11 files changed
+648 -113
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

cmd/netssh/connection_test.go
+197
View File
@@ -0,0 +1,197 @@
package main
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/cache"
"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/netbox"
)
type connectionTestHost struct {
ID int `json:"id"`
Name string `json:"name"`
PrimaryIP4 *struct {
Address string `json:"address"`
} `json:"primary_ip4,omitempty"`
}
func connectionTestServer(t *testing.T, devices, vms []connectionTestHost) *httptest.Server {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
var hosts []connectionTestHost
switch r.URL.Path {
case "/api/dcim/devices/":
hosts = devices
case "/api/virtualization/virtual-machines/":
hosts = vms
case "/api/ipam/ip-addresses/":
fmt.Fprint(w, `{"count":0,"results":[]}`)
return
default:
http.NotFound(w, r)
return
}
query := strings.ToLower(r.URL.Query().Get("name__ic"))
matched := make([]connectionTestHost, 0)
for _, host := range hosts {
if strings.Contains(strings.ToLower(host.Name), query) {
matched = append(matched, host)
}
}
offset, _ := strconv.Atoi(r.URL.Query().Get("offset"))
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
if limit <= 0 {
limit = 50
}
start, end := min(offset, len(matched)), min(offset+limit, len(matched))
json.NewEncoder(w).Encode(map[string]any{"count": len(matched), "results": matched[start:end]})
}))
t.Cleanup(srv.Close)
return srv
}
func TestFindSSHHostBeyondFirstPage(t *testing.T) {
var devices []connectionTestHost
for i := range 50 {
devices = append(devices, connectionTestHost{ID: i + 1, Name: fmt.Sprintf("web-extra-%02d", i)})
}
devices = append(devices, connectionTestHost{ID: 51, Name: "web"})
srv := connectionTestServer(t, devices, nil)
host, err := findSSHHost(context.Background(), netbox.NewClient(srv.URL, "token", 0), "web", nil)
if err != nil || host.ID != 51 {
t.Fatalf("host beyond first page not resolved: %+v, %v", host, err)
}
}
func TestFindSSHHostDomainExpansion(t *testing.T) {
for _, domain := range []string{"example.com", ".example.com", " .example.com "} {
t.Run(domain, func(t *testing.T) {
srv := connectionTestServer(t, []connectionTestHost{{ID: 1, Name: "web.example.com"}}, nil)
host, err := findSSHHost(context.Background(), netbox.NewClient(srv.URL, "token", 0), "web", []string{domain})
if err != nil || host.Name != "web.example.com" {
t.Fatalf("domain expansion failed: %+v, %v", host, err)
}
})
}
}
func TestFindSSHHostRejectsAmbiguity(t *testing.T) {
srv := connectionTestServer(t, []connectionTestHost{{ID: 1, Name: "web"}, {ID: 2, Name: "web.example.com"}}, []connectionTestHost{{ID: 3, Name: "web"}})
_, err := findSSHHost(context.Background(), netbox.NewClient(srv.URL, "token", 0), "web", []string{"example.com"})
if err == nil || errors.Is(err, netbox.ErrHostNotFound) || !strings.Contains(err.Error(), "2 exact matches") {
t.Fatalf("ambiguous name accepted or expanded: %v", err)
}
}
func TestFindSSHHostDoesNotHideAPIErrors(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, "unavailable", http.StatusServiceUnavailable)
}))
defer srv.Close()
_, err := findSSHHost(context.Background(), netbox.NewClient(srv.URL, "token", 0), "web", []string{"example.com"})
if err == nil || errors.Is(err, netbox.ErrHostNotFound) {
t.Fatalf("API failure hidden as missing host: %v", err)
}
}
func TestFindSSHHostNotFound(t *testing.T) {
srv := connectionTestServer(t, nil, nil)
_, err := findSSHHost(context.Background(), netbox.NewClient(srv.URL, "token", 0), "missing", []string{"example.com"})
if !errors.Is(err, netbox.ErrHostNotFound) {
t.Fatalf("want missing-host error, got %v", err)
}
}
func TestCacheRefreshHonorsResolverAndRemovesUnsafeEntries(t *testing.T) {
for _, tc := range []struct {
name, strategies string
wantCached bool
}{
{"management only", "[management_subnet]", false},
{"explicit primary fallback", "[management_subnet, primary_ip]", true},
} {
t.Run(tc.name, func(t *testing.T) {
host := connectionTestHost{ID: 1, Name: "web", PrimaryIP4: &struct {
Address string `json:"address"`
}{Address: "192.0.2.1/24"}}
srv := connectionTestServer(t, []connectionTestHost{host}, nil)
path := prepareRefreshTest(t, srv.URL, tc.strategies)
c := cache.New(path, 3600)
c.Upsert(cache.Entry{Name: "web", IP: "192.0.2.99", Kind: "device"})
if err := c.Save(); err != nil {
t.Fatal(err)
}
cmd := rootCmd()
cmd.SetArgs([]string{"cache", "refresh"})
if err := cmd.Execute(); err != nil {
t.Fatal(err)
}
if err := c.Load(); err != nil {
t.Fatal(err)
}
got, fresh := c.Get("web")
if tc.wantCached {
if !fresh || got.IP != "192.0.2.1" {
t.Fatalf("explicit fallback failed: %+v", got)
}
} else if got.Name != "" {
t.Fatalf("unsafe primary IP retained: %+v", got)
}
})
}
}
func TestCacheRefreshSkipsAmbiguousNames(t *testing.T) {
host := connectionTestHost{ID: 1, Name: "web", PrimaryIP4: &struct {
Address string `json:"address"`
}{Address: "192.0.2.1/24"}}
vm := host
vm.ID = 2
srv := connectionTestServer(t, []connectionTestHost{host}, []connectionTestHost{vm})
path := prepareRefreshTest(t, srv.URL, "[primary_ip]")
c := cache.New(path, 3600)
c.Upsert(cache.Entry{Name: "web", IP: "192.0.2.99", Kind: "device"})
if err := c.Save(); err != nil {
t.Fatal(err)
}
cmd := rootCmd()
cmd.SetArgs([]string{"cache", "refresh"})
if err := cmd.Execute(); err != nil {
t.Fatal(err)
}
if err := c.Load(); err != nil {
t.Fatal(err)
}
if len(c.All()) != 0 {
t.Fatalf("ambiguous host cached: %+v", c.All())
}
}
func prepareRefreshTest(t *testing.T, baseURL, strategies string) string {
t.Helper()
configDir := t.TempDir()
t.Setenv("HOME", t.TempDir())
t.Setenv("XDG_CONFIG_HOME", configDir)
t.Setenv("XDG_CACHE_HOME", t.TempDir())
path := filepath.Join(t.TempDir(), "hosts.json")
t.Setenv("NETSSH_CACHE_PATH", path)
t.Setenv("NETSSH_NETBOX_URL", baseURL)
t.Setenv("NETSSH_NETBOX_TOKEN", "token")
t.Setenv("NETSSH_SHELL_ALIAS", "")
data := fmt.Sprintf("resolver:\n strategies: %s\n management_subnets: [10.0.0.0/8]\n", strategies)
if err := os.WriteFile(filepath.Join(configDir, "netssh.yaml"), []byte(data), 0o600); err != nil {
t.Fatal(err)
}
return path
}
cmd/netssh/main.go
+51 -54
View File
@@ -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 {