From be05717a88c7986db0e94b0cd46c43cdb50e6935 Mon Sep 17 00:00:00 2001 From: Fraser Isbester Date: Wed, 8 Jul 2026 11:30:26 -0700 Subject: [PATCH 1/2] feat: first-run setup screen + VPC-aware connection management MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make tusk easy to start against real databases. Onboarding: when a connection can't be established — no config, an unknown profile, or an unreachable DB — tusk now opens an interactive setup screen instead of printing a cobra error and exiting. The user enters a connection (or a tunnel), tusk validates it with a real Ping, optionally saves it as a profile (comment-preserving), and drops straight into the TUI. Subsequent launches connect through. Connectivity: new TUI-free internal/connect package turns a profile into a reachable DSN, establishing a tunnel first when needed and tearing it down on exit. Methods: - direct (default): use the URL/fields as-is. - kube-port-forward: spawn `kubectl port-forward` into a cluster (GKE), wait for the local port, connect, kill the process group on exit. - exec: run an arbitrary tunnel command (SSH bastion, cloud-sql-proxy, …) with {local_port} substitution. Both tusk and the tuskd daemon use it, so headless enforcement works against VPC databases too. Config gains a `connect:` block (ConnectConfig) and omitempty on Profile fields; SaveProfile persists a profile via the same comment-preserving YAML node edit as SaveProfileRules, carrying existing rules over. Tests: connect method dispatch, kube argv, {local_port} substitution, freePort/waitForPort (incl. process-death), and spawn/teardown; SaveProfile round-trips. Verified end-to-end against local Postgres: direct onboarding saves a working profile and re-runs connect through, and an exec tunnel (via a stand-in forwarder) drives tuskd against a real DB and cleans up. Closes #11 --- README.md | 56 +++++++ cmd/tusk/main.go | 128 ++++++++++----- cmd/tuskd/main.go | 11 +- internal/config/config.go | 175 ++++++++++++++------ internal/config/config_profile_test.go | 97 +++++++++++ internal/connect/connect.go | 94 +++++++++++ internal/connect/connect_suite_test.go | 13 ++ internal/connect/connect_test.go | 133 +++++++++++++++ internal/connect/methods.go | 132 +++++++++++++++ internal/tui/setup.go | 215 +++++++++++++++++++++++++ 10 files changed, 965 insertions(+), 89 deletions(-) create mode 100644 internal/config/config_profile_test.go create mode 100644 internal/connect/connect.go create mode 100644 internal/connect/connect_suite_test.go create mode 100644 internal/connect/connect_test.go create mode 100644 internal/connect/methods.go create mode 100644 internal/tui/setup.go diff --git a/README.md b/README.md index 80920e7..394a972 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,12 @@ tusk 'postgres://user:pass@localhost:5432/mydb' tusk -P production ``` +**First run:** if no connection can be established — no config yet, an unknown +profile (`tusk -P nope`), or an unreachable database — tusk drops into an +interactive setup screen instead of erroring out. Enter a connection (or a +tunnel, see below), and tusk validates it and optionally saves it as a profile +so the next launch connects straight through. + ## Daemon `tuskd` runs a profile's rules engine headlessly — the same rules, evaluation, and @@ -142,6 +148,56 @@ profiles: **Lock**: `blocked_pid`, `blocking_pid`, `blocked_user`, `blocking_user`, `blocked_app`, `blocking_app`, `lock_type`, `mode`, `wait_duration` +## Connecting to remote / VPC databases + +Production and staging databases usually aren't reachable directly — they live +inside a VPC. A profile's `connect:` block tells tusk how to reach them, and tusk +establishes the tunnel itself (and tears it down on exit). This works for both +`tusk` and the `tuskd` daemon. + +**Kubernetes port-forward** (e.g. Postgres in a GKE cluster) — tusk runs +`kubectl port-forward` for you: + +```yaml +profiles: + staging: + connect: + via: kube-port-forward + context: gke_myproject_us-central1_staging # optional; omit for current context + namespace: databases + target: svc/postgres # svc/pod/deployment/statefulset + remote_port: 5432 + # local_port: 0 # 0 = auto-pick a free port + user: readonly + database: appdb + sslmode: disable +``` + +**Arbitrary tunnel command** (`exec`) — the escape hatch for SSH bastions, the +Cloud SQL Auth Proxy, or anything else. `{local_port}` is substituted with an +auto-picked free port that tusk then connects to: + +```yaml +profiles: + bastion: + connect: + via: exec + command: ["ssh", "-N", "-L", "{local_port}:db.internal:5432", "bastion.example.com"] + user: readonly + database: appdb + cloudsql: + connect: + via: exec + command: ["cloud-sql-proxy", "--port", "{local_port}", "myproj:us-central1:pg"] + user: readonly + database: appdb +``` + +For tunnel methods, credentials and the database name come from the profile +fields (`user`/`password`/`database`/`sslmode`); host and port are supplied by +the tunnel. The default method is `direct`, which uses `url`/fields as-is. +`kubectl` / `cloud-sql-proxy` / `ssh` must be on your `PATH`. + ## Views | Command | Description | diff --git a/cmd/tusk/main.go b/cmd/tusk/main.go index 0f34579..c5e8d09 100644 --- a/cmd/tusk/main.go +++ b/cmd/tusk/main.go @@ -9,11 +9,24 @@ import ( "github.com/spf13/cobra" "github.com/fraser-isbester/tusk/internal/config" + "github.com/fraser-isbester/tusk/internal/connect" "github.com/fraser-isbester/tusk/internal/db" "github.com/fraser-isbester/tusk/internal/rules" "github.com/fraser-isbester/tusk/internal/tui" ) +// session is a validated, live connection plus everything the TUI needs to +// render it. The closer tears down any tunnel and must run after db.Close. +type session struct { + db *db.DB + closeTunnel func() error + profileName string + color string + connUser string + readonly bool + ruleConfigs []rules.RuleConfig +} + func main() { var profileFlag string @@ -22,62 +35,52 @@ func main() { Short: "Tusk — a terminal UI for PostgreSQL administration", Long: "A k9s-style terminal UI for real-time PostgreSQL management.", Args: cobra.MaximumNArgs(1), - RunE: func(cmd *cobra.Command, args []string) error { + RunE: func(_ *cobra.Command, args []string) error { cfg, err := config.Load() if err != nil { return fmt.Errorf("loading config: %w", err) } - var connStr string - var profileName string - var profileColor string - var connUser string - var readonly bool - var ruleConfigs []rules.RuleConfig - - if len(args) > 0 { - connStr = args[0] - profileName = "cli" - } else { - profile, err := cfg.ResolveProfile(profileFlag) - if err != nil { - return fmt.Errorf("resolving profile: %w", err) + // Try to bring up a connection from args/profile. If we can't + // (unresolved profile, unreachable DB), drop into the setup screen + // instead of exiting — get people into the TUI ASAP. + sess, reason := establish(cfg, profileFlag, args) + if sess == nil { + res, ok := tui.RunSetup(cfg, reason) + if !ok { + return nil // user quit setup } - connStr = profile.ConnectionString() - profileName = profileFlag - if profileName == "" { - profileName = cfg.DefaultProfile + sess = &session{ + db: res.DB, + closeTunnel: res.Closer, + profileName: res.ProfileName, + color: res.Profile.Color, + connUser: res.Profile.User, + readonly: res.Profile.Readonly, + ruleConfigs: res.Profile.Rules, } - profileColor = profile.Color - connUser = profile.User - readonly = profile.Readonly - ruleConfigs = profile.Rules } + defer sess.db.Close() + defer func() { _ = sess.closeTunnel() }() - database, err := db.New(connStr) - if err != nil { - return fmt.Errorf("connecting to database: %w", err) - } - defer database.Close() - - // Query actual connection user if not set from profile - if connUser == "" { + // Query the actual connection user if the profile didn't set one. + if sess.connUser == "" { var user string - if err := database.Pool().QueryRow(context.Background(), "SELECT current_user").Scan(&user); err == nil { - connUser = user + if err := sess.db.Pool().QueryRow(context.Background(), "SELECT current_user").Scan(&user); err == nil { + sess.connUser = user } } var engine *rules.Engine - if len(ruleConfigs) > 0 { - compiled, err := rules.BuildRules(ruleConfigs, readonly) + if len(sess.ruleConfigs) > 0 { + compiled, err := rules.BuildRules(sess.ruleConfigs, sess.readonly) if err != nil { return fmt.Errorf("compiling rules: %w", err) } - engine = rules.NewEngine(compiled, database, 5*time.Minute, 1000) + engine = rules.NewEngine(compiled, sess.db, 5*time.Minute, 1000) } - app := tui.NewApp(database, cfg, profileName, profileColor, connUser, readonly, engine) + app := tui.NewApp(sess.db, cfg, sess.profileName, sess.color, sess.connUser, sess.readonly, engine) if err := app.Run(); err != nil { return fmt.Errorf("running tusk: %w", err) } @@ -91,3 +94,54 @@ func main() { os.Exit(1) } } + +// establish attempts to open and validate a connection. It returns a live +// session on success, or (nil, reason) when the caller should fall back to the +// interactive setup screen — reason explains why (shown to the user). +func establish(cfg *config.Config, profileFlag string, args []string) (*session, string) { + if len(args) > 0 { + return connectProfile(config.Profile{URL: args[0]}, "cli") + } + + profile, err := cfg.ResolveProfile(profileFlag) + if err != nil { + return nil, err.Error() + } + name := profileFlag + if name == "" { + name = cfg.DefaultProfile + } + return connectProfile(profile, name) +} + +// connectProfile opens any tunnel, builds the pool, and validates it with a +// ping. On any failure it tears everything down and returns a reason string so +// the caller can offer setup. +func connectProfile(profile config.Profile, name string) (*session, string) { + dsn, closer, err := connect.Open(context.Background(), profile) + if err != nil { + return nil, err.Error() + } + database, err := db.New(dsn) + if err != nil { + _ = closer() + return nil, err.Error() + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := database.Pool().Ping(ctx); err != nil { + database.Close() + _ = closer() + return nil, fmt.Sprintf("could not connect to %q: %s", name, err) + } + + return &session{ + db: database, + closeTunnel: closer, + profileName: name, + color: profile.Color, + connUser: profile.User, + readonly: profile.Readonly, + ruleConfigs: profile.Rules, + }, "" +} diff --git a/cmd/tuskd/main.go b/cmd/tuskd/main.go index 1840fb1..f2c3833 100644 --- a/cmd/tuskd/main.go +++ b/cmd/tuskd/main.go @@ -16,6 +16,7 @@ import ( "github.com/spf13/cobra" "github.com/fraser-isbester/tusk/internal/config" + "github.com/fraser-isbester/tusk/internal/connect" "github.com/fraser-isbester/tusk/internal/db" "github.com/fraser-isbester/tusk/internal/rules" ) @@ -52,7 +53,15 @@ func main() { return fmt.Errorf("no rules configured in profile %q", profileName) } - database, err := db.New(profile.ConnectionString()) + // Establish any tunnel (e.g. kube-port-forward into a VPC) before + // connecting; the daemon needs the same reachability as the TUI. + dsn, closeTunnel, err := connect.Open(context.Background(), profile) + if err != nil { + return fmt.Errorf("establishing connection: %w", err) + } + defer func() { _ = closeTunnel() }() + + database, err := db.New(dsn) if err != nil { return fmt.Errorf("connecting to database: %w", err) } diff --git a/internal/config/config.go b/internal/config/config.go index c00d588..f53d288 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -20,29 +20,46 @@ type Config struct { DefaultProfile string `yaml:"default_profile"` } -// Profile represents a single PostgreSQL connection profile. +// Profile represents a single PostgreSQL connection profile. Optional fields +// use omitempty so profiles written back to disk (by the setup screen) stay +// clean and don't bake in defaults. type Profile struct { - Host string `yaml:"host"` - Port int `yaml:"port"` - User string `yaml:"user"` - Password string `yaml:"password"` - Database string `yaml:"database"` - SSLMode string `yaml:"sslmode"` - URL string `yaml:"url"` - Readonly bool `yaml:"readonly"` - Color string `yaml:"color"` - RefreshInterval time.Duration `yaml:"refresh_interval"` - Rules []rules.RuleConfig `yaml:"rules"` + Host string `yaml:"host,omitempty"` + Port int `yaml:"port,omitempty"` + User string `yaml:"user,omitempty"` + Password string `yaml:"password,omitempty"` + Database string `yaml:"database,omitempty"` + SSLMode string `yaml:"sslmode,omitempty"` + URL string `yaml:"url,omitempty"` + Readonly bool `yaml:"readonly,omitempty"` + Color string `yaml:"color,omitempty"` + RefreshInterval time.Duration `yaml:"refresh_interval,omitempty"` + Connect *ConnectConfig `yaml:"connect,omitempty"` + Rules []rules.RuleConfig `yaml:"rules,omitempty"` +} + +// ConnectConfig describes how tusk should reach a database that isn't directly +// accessible — e.g. a kubectl port-forward into a VPC. When Via is empty or +// "direct", the profile's URL/fields are used as-is. For tunnel methods +// (kube-port-forward, exec) the profile supplies credentials and database name +// while the tunnel supplies the local host:port. +type ConnectConfig struct { + Via string `yaml:"via,omitempty"` // "" | direct | kube-port-forward | exec + Context string `yaml:"context,omitempty"` // kube: kubeconfig context + Namespace string `yaml:"namespace,omitempty"` // kube: namespace + Target string `yaml:"target,omitempty"` // kube: svc/pod/deploy/statefulset + RemotePort int `yaml:"remote_port,omitempty"` // kube: remote port (default 5432) + LocalPort int `yaml:"local_port,omitempty"` // 0 = auto-pick a free port + Command []string `yaml:"command,omitempty"` // exec: argv, {local_port} substituted } // ConnectionString returns a PostgreSQL connection string for this profile. // If URL is set directly, it is returned as-is. Otherwise, the string is -// assembled from the individual fields. +// assembled from the individual fields using localhost/5432 defaults. func (p Profile) ConnectionString() string { if p.URL != "" { return p.URL } - host := p.Host if host == "" { host = "localhost" @@ -51,6 +68,13 @@ func (p Profile) ConnectionString() string { if port == 0 { port = 5432 } + return p.DSN(host, port) +} + +// DSN builds a connection string against the given host and port using the +// profile's credentials, database, and sslmode (ignoring Host/Port/URL). It is +// used when a tunnel supplies the local endpoint. +func (p Profile) DSN(host string, port int) string { user := p.User if user == "" { user = "postgres" @@ -63,7 +87,6 @@ func (p Profile) ConnectionString() string { if sslmode == "" { sslmode = "disable" } - if p.Password != "" { return fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=%s", url.PathEscape(user), url.PathEscape(p.Password), host, port, db, sslmode) @@ -199,52 +222,84 @@ func SaveProfileRules(profileName string, ruleConfigs []rules.RuleConfig) error if err != nil { return err } - - // Load the existing document (preserving comments/formatting), or start a - // fresh one if the file does not exist yet. - root := &yaml.Node{Kind: yaml.MappingNode} - if data, readErr := os.ReadFile(path); readErr == nil { //nolint:gosec // path from configPath(), not user input - var doc yaml.Node - if parseErr := yaml.Unmarshal(data, &doc); parseErr != nil { - return fmt.Errorf("parsing config %s: %w", path, parseErr) - } - if len(doc.Content) > 0 && doc.Content[0].Kind == yaml.MappingNode { - root = doc.Content[0] - } - } else if !os.IsNotExist(readErr) { - return fmt.Errorf("reading config %s: %w", path, readErr) - } - - // profiles: - profiles := mapValue(root, "profiles") - if profiles == nil || profiles.Kind != yaml.MappingNode { - profiles = &yaml.Node{Kind: yaml.MappingNode} - setMapValue(root, "profiles", profiles) + root, err := loadConfigRoot(path) + if err != nil { + return err } - // profiles.: + profiles := ensureProfiles(root) profile := mapValue(profiles, profileName) if profile == nil || profile.Kind != yaml.MappingNode { profile = &yaml.Node{Kind: yaml.MappingNode} setMapValue(profiles, profileName, profile) } - // profiles..rules: (marshal the configs to a fresh sequence node) rulesNode, err := toNode(ruleConfigs) if err != nil { return fmt.Errorf("encoding rules: %w", err) } setMapValue(profile, "rules", rulesNode) + ensureDefaultProfile(root, profileName) - // default_profile: fill in for a freshly created file so it round-trips. - if mapValue(root, "default_profile") == nil { - setMapValue(root, "default_profile", &yaml.Node{ - Kind: yaml.ScalarNode, Tag: "!!str", Value: profileName, - }) + return writeConfigRoot(path, root) +} + +// SaveProfile writes p to profiles. in ~/.config/tusk/config.yaml, +// preserving comments and every other profile (same node-edit approach as +// SaveProfileRules). An existing profile's rules: node is carried over when p +// has no rules, so saving connection details from the setup screen never drops +// rules a user added via the editor. The write is atomic. +func SaveProfile(name string, p Profile) error { + path, err := configPath() + if err != nil { + return err + } + root, err := loadConfigRoot(path) + if err != nil { + return err } - // Encode with 2-space indent to match the documented config style and keep - // diffs minimal (yaml.Marshal defaults to 4). + profiles := ensureProfiles(root) + profNode, err := toNode(p) + if err != nil { + return fmt.Errorf("encoding profile: %w", err) + } + // Preserve existing rules if we're only updating connection details. + if existing := mapValue(profiles, name); existing != nil && existing.Kind == yaml.MappingNode && len(p.Rules) == 0 { + if r := mapValue(existing, "rules"); r != nil { + setMapValue(profNode, "rules", r) + } + } + setMapValue(profiles, name, profNode) + ensureDefaultProfile(root, name) + + return writeConfigRoot(path, root) +} + +// loadConfigRoot reads the config file into its root mapping node (preserving +// comments), or returns a fresh mapping node if the file does not exist. +func loadConfigRoot(path string) (*yaml.Node, error) { + root := &yaml.Node{Kind: yaml.MappingNode} + data, err := os.ReadFile(path) //nolint:gosec // path from configPath(), not user input + if err != nil { + if os.IsNotExist(err) { + return root, nil + } + return nil, fmt.Errorf("reading config %s: %w", path, err) + } + var doc yaml.Node + if err := yaml.Unmarshal(data, &doc); err != nil { + return nil, fmt.Errorf("parsing config %s: %w", path, err) + } + if len(doc.Content) > 0 && doc.Content[0].Kind == yaml.MappingNode { + root = doc.Content[0] + } + return root, nil +} + +// writeConfigRoot atomically writes root to path with 2-space indent (matching +// the documented config style and minimizing diffs). +func writeConfigRoot(path string, root *yaml.Node) error { var buf bytes.Buffer enc := yaml.NewEncoder(&buf) enc.SetIndent(2) @@ -252,19 +307,37 @@ func SaveProfileRules(profileName string, ruleConfigs []rules.RuleConfig) error return fmt.Errorf("marshaling config: %w", err) } _ = enc.Close() - out := buf.Bytes() - dir := filepath.Dir(path) - if mkErr := os.MkdirAll(dir, 0o750); mkErr != nil { - return fmt.Errorf("creating config dir: %w", mkErr) + if err := os.MkdirAll(filepath.Dir(path), 0o750); err != nil { + return fmt.Errorf("creating config dir: %w", err) } tmp := path + ".tmp" - if writeErr := os.WriteFile(tmp, out, 0o600); writeErr != nil { - return fmt.Errorf("writing config: %w", writeErr) + if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil { + return fmt.Errorf("writing config: %w", err) } return os.Rename(tmp, path) } +// ensureProfiles returns the profiles mapping node, creating it if absent. +func ensureProfiles(root *yaml.Node) *yaml.Node { + profiles := mapValue(root, "profiles") + if profiles == nil || profiles.Kind != yaml.MappingNode { + profiles = &yaml.Node{Kind: yaml.MappingNode} + setMapValue(root, "profiles", profiles) + } + return profiles +} + +// ensureDefaultProfile sets default_profile to name when it is not already set, +// so a freshly created file round-trips. +func ensureDefaultProfile(root *yaml.Node, name string) { + if mapValue(root, "default_profile") == nil { + setMapValue(root, "default_profile", &yaml.Node{ + Kind: yaml.ScalarNode, Tag: "!!str", Value: name, + }) + } +} + // mapValue returns the value node for key in a mapping node, or nil if absent. // Mapping content alternates key, value, key, value, ... func mapValue(m *yaml.Node, key string) *yaml.Node { diff --git a/internal/config/config_profile_test.go b/internal/config/config_profile_test.go new file mode 100644 index 0000000..b37a0f9 --- /dev/null +++ b/internal/config/config_profile_test.go @@ -0,0 +1,97 @@ +package config_test + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/fraser-isbester/tusk/internal/config" + "github.com/fraser-isbester/tusk/internal/rules" +) + +var _ = Describe("SaveProfile", func() { + var cfgPath string + + BeforeEach(func() { + home := GinkgoT().TempDir() + GinkgoT().Setenv("HOME", home) + cfgPath = filepath.Join(home, ".config", "tusk", "config.yaml") + }) + + read := func() string { + data, err := os.ReadFile(cfgPath) //nolint:gosec // test-controlled temp path + Expect(err).NotTo(HaveOccurred()) + return string(data) + } + + It("creates a clean profile and round-trips through Load", func() { + err := config.SaveProfile("staging", config.Profile{ + User: "readonly", + Database: "appdb", + Connect: &config.ConnectConfig{ + Via: "kube-port-forward", + Context: "gke_proj_staging", + Namespace: "databases", + Target: "svc/postgres", + RemotePort: 5432, + }, + }) + Expect(err).NotTo(HaveOccurred()) + + out := read() + Expect(out).To(ContainSubstring("default_profile: staging")) + Expect(out).To(ContainSubstring("via: kube-port-forward")) + Expect(out).To(ContainSubstring("target: svc/postgres")) + // omitempty keeps unset fields out of the file. + Expect(out).NotTo(ContainSubstring("host:")) + Expect(out).NotTo(ContainSubstring("password")) + Expect(out).NotTo(ContainSubstring("readonly:")) + Expect(out).NotTo(ContainSubstring("local_port")) + + cfg, err := config.Load() + Expect(err).NotTo(HaveOccurred()) + p, err := cfg.ResolveProfile("staging") + Expect(err).NotTo(HaveOccurred()) + Expect(p.User).To(Equal("readonly")) + Expect(p.Connect).NotTo(BeNil()) + Expect(p.Connect.Via).To(Equal("kube-port-forward")) + Expect(p.Connect.Target).To(Equal("svc/postgres")) + }) + + It("preserves comments and other profiles", func() { + Expect(os.MkdirAll(filepath.Dir(cfgPath), 0o750)).To(Succeed()) + original := "# my config\ndefault_profile: dev\nprofiles:\n dev:\n url: \"postgres://localhost/dev\"\n" + Expect(os.WriteFile(cfgPath, []byte(original), 0o600)).To(Succeed()) + + Expect(config.SaveProfile("prod", config.Profile{URL: "postgres://prod/db"})).To(Succeed()) + + out := read() + Expect(out).To(ContainSubstring("# my config")) + Expect(out).To(ContainSubstring("default_profile: dev")) // not overwritten + Expect(out).To(ContainSubstring("postgres://localhost/dev")) + Expect(out).To(ContainSubstring("postgres://prod/db")) + }) + + It("does not drop rules when re-saving connection details", func() { + enabled := true + Expect(config.SaveProfileRules("dev", []rules.RuleConfig{ + {Name: "r1", Resource: "query", When: "true", Action: "log", Enabled: &enabled}, + })).To(Succeed()) + + // Re-save the profile with new connection info but no rules. + Expect(config.SaveProfile("dev", config.Profile{URL: "postgres://new/db"})).To(Succeed()) + + out := read() + Expect(out).To(ContainSubstring("postgres://new/db")) + Expect(out).To(ContainSubstring("name: r1")) // rules preserved + + cfg, err := config.Load() + Expect(err).NotTo(HaveOccurred()) + p, err := cfg.ResolveProfile("dev") + Expect(err).NotTo(HaveOccurred()) + Expect(p.Rules).To(HaveLen(1)) + Expect(p.URL).To(Equal("postgres://new/db")) + }) +}) diff --git a/internal/connect/connect.go b/internal/connect/connect.go new file mode 100644 index 0000000..916db59 --- /dev/null +++ b/internal/connect/connect.go @@ -0,0 +1,94 @@ +// Package connect turns a connection profile into a reachable PostgreSQL DSN. +// +// Tusk is a dev tool and the databases worth watching usually live inside a VPC +// with no public endpoint. A profile's connect block declares how to reach it — +// directly, or by having tusk establish a tunnel first (a kubectl port-forward +// into a cluster, or an arbitrary command such as an SSH tunnel or the Cloud SQL +// Auth Proxy). Open establishes the tunnel, returns a DSN pointing at the local +// endpoint, and hands back a closer that tears the tunnel down on exit. +// +// The package has no TUI dependency so both the tusk TUI and the tuskd daemon +// use it. +package connect + +import ( + "context" + "fmt" + "net" + "time" + + "github.com/fraser-isbester/tusk/internal/config" +) + +// establishTimeout bounds how long we wait for a tunnel's local port to start +// accepting connections before giving up. +const establishTimeout = 15 * time.Second + +// Open resolves profile into a DSN that is reachable from localhost, starting a +// tunnel first when the profile's connect method requires one. The returned +// closer must be called when the connection is no longer needed; for the direct +// method it is a no-op. +func Open(ctx context.Context, profile config.Profile) (dsn string, closer func() error, err error) { + method := "direct" + if profile.Connect != nil && profile.Connect.Via != "" { + method = profile.Connect.Via + } + + switch method { + case "direct": + return profile.ConnectionString(), noopCloser, nil + case "kube-port-forward": + return openKubePortForward(ctx, profile) + case "exec": + return openExec(ctx, profile) + default: + return "", nil, fmt.Errorf("unknown connect method %q (want direct, kube-port-forward, or exec)", method) + } +} + +func noopCloser() error { return nil } + +// freePort asks the OS for an unused TCP port by binding to :0 and releasing it. +// There is an inherent race between release and reuse, but it is the standard +// approach and the window is tiny for a short-lived local forward. +func freePort() (int, error) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return 0, err + } + port := l.Addr().(*net.TCPAddr).Port + if cerr := l.Close(); cerr != nil { + return 0, cerr + } + return port, nil +} + +// waitForPort blocks until a TCP connection to addr succeeds or the deadline +// passes. It also aborts if ctx is canceled or dead reports the tunnel process +// exited early (returning that as the error, so callers surface why). +func waitForPort(ctx context.Context, addr string, timeout time.Duration, dead <-chan error) error { + deadline := time.Now().Add(timeout) + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + + for { + conn, err := net.DialTimeout("tcp", addr, 500*time.Millisecond) + if err == nil { + _ = conn.Close() + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case derr := <-dead: + if derr != nil { + return fmt.Errorf("tunnel process exited: %w", derr) + } + return fmt.Errorf("tunnel process exited before %s became reachable", addr) + case <-ticker.C: + } + if time.Now().After(deadline) { + return fmt.Errorf("timed out after %s waiting for %s", timeout, addr) + } + } +} diff --git a/internal/connect/connect_suite_test.go b/internal/connect/connect_suite_test.go new file mode 100644 index 0000000..ae93ab0 --- /dev/null +++ b/internal/connect/connect_suite_test.go @@ -0,0 +1,13 @@ +package connect + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestConnect(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Connect Suite") +} diff --git a/internal/connect/connect_test.go b/internal/connect/connect_test.go new file mode 100644 index 0000000..963d75d --- /dev/null +++ b/internal/connect/connect_test.go @@ -0,0 +1,133 @@ +package connect + +import ( + "context" + "net" + "strconv" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/fraser-isbester/tusk/internal/config" +) + +var _ = Describe("Open", func() { + It("returns the profile DSN unchanged for the direct method", func() { + p := config.Profile{URL: "postgres://u:p@localhost:5432/db"} //nolint:gosec // test fixture + dsn, closer, err := Open(context.Background(), p) + Expect(err).NotTo(HaveOccurred()) + Expect(dsn).To(Equal("postgres://u:p@localhost:5432/db")) + Expect(closer()).To(Succeed()) + }) + + It("treats an empty connect block as direct", func() { + p := config.Profile{URL: "postgres://localhost/db", Connect: &config.ConnectConfig{}} + dsn, _, err := Open(context.Background(), p) + Expect(err).NotTo(HaveOccurred()) + Expect(dsn).To(Equal("postgres://localhost/db")) + }) + + It("rejects an unknown method", func() { + p := config.Profile{Connect: &config.ConnectConfig{Via: "carrier-pigeon"}} + _, _, err := Open(context.Background(), p) + Expect(err).To(MatchError(ContainSubstring("unknown connect method"))) + }) + + It("errors when kube-port-forward has no target", func() { + p := config.Profile{Connect: &config.ConnectConfig{Via: "kube-port-forward"}} + _, _, err := Open(context.Background(), p) + Expect(err).To(MatchError(ContainSubstring("requires a target"))) + }) + + It("errors when exec has no command", func() { + p := config.Profile{Connect: &config.ConnectConfig{Via: "exec"}} + _, _, err := Open(context.Background(), p) + Expect(err).To(MatchError(ContainSubstring("requires a command"))) + }) + + Describe("exec method", func() { + It("connects once the local port is reachable and tears the process down", func() { + // Pre-open a listener to stand in for the far end of a tunnel, and + // use a long-lived `sleep` as the tunnel process. spawnTunnel should + // see the port is reachable and return a DSN + working closer. + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer func() { _ = ln.Close() }() + port := ln.Addr().(*net.TCPAddr).Port + + p := config.Profile{ + User: "readonly", Database: "appdb", + Connect: &config.ConnectConfig{Via: "exec", LocalPort: port, Command: []string{"sleep", "30"}}, + } + dsn, closer, err := Open(context.Background(), p) + Expect(err).NotTo(HaveOccurred()) + Expect(dsn).To(ContainSubstring("readonly@127.0.0.1:")) + Expect(dsn).To(ContainSubstring("/appdb?sslmode=disable")) + Expect(closer()).To(Succeed()) + }) + + It("fails fast with the tunnel's exit when the command dies early", func() { + p := config.Profile{ + Connect: &config.ConnectConfig{Via: "exec", Command: []string{"sh", "-c", "exit 3"}}, + } + _, _, err := Open(context.Background(), p) + Expect(err).To(MatchError(ContainSubstring("tunnel process exited"))) + }) + }) +}) + +var _ = Describe("helpers", func() { + Describe("kubectlArgs", func() { + It("includes context and namespace when set", func() { + c := &config.ConnectConfig{Context: "ctx", Namespace: "ns", Target: "svc/pg"} + Expect(kubectlArgs(c, 5000, 5432)).To(Equal( + []string{"port-forward", "--context", "ctx", "-n", "ns", "svc/pg", "5000:5432"})) + }) + It("omits context and namespace when empty", func() { + c := &config.ConnectConfig{Target: "pod/pg"} + Expect(kubectlArgs(c, 6000, 5432)).To(Equal( + []string{"port-forward", "pod/pg", "6000:5432"})) + }) + }) + + Describe("substituteLocalPort", func() { + It("replaces the token everywhere", func() { + out := substituteLocalPort([]string{"ssh", "-L", "{local_port}:db:5432", "x{local_port}"}, 42) + Expect(out).To(Equal([]string{"ssh", "-L", "42:db:5432", "x42"})) + }) + }) + + Describe("freePort", func() { + It("returns a bindable port", func() { + port, err := freePort() + Expect(err).NotTo(HaveOccurred()) + Expect(port).To(BeNumerically(">", 0)) + ln, err := net.Listen("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port))) + Expect(err).NotTo(HaveOccurred()) + _ = ln.Close() + }) + }) + + Describe("waitForPort", func() { + It("returns nil once the address is reachable", func() { + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer func() { _ = ln.Close() }() + err = waitForPort(context.Background(), ln.Addr().String(), 2*time.Second, make(chan error)) + Expect(err).NotTo(HaveOccurred()) + }) + + It("times out when nothing is listening", func() { + err := waitForPort(context.Background(), "127.0.0.1:1", 300*time.Millisecond, make(chan error)) + Expect(err).To(MatchError(ContainSubstring("timed out"))) + }) + + It("fails fast when the tunnel process dies", func() { + dead := make(chan error, 1) + dead <- nil + err := waitForPort(context.Background(), "127.0.0.1:1", 5*time.Second, dead) + Expect(err).To(MatchError(ContainSubstring("exited"))) + }) + }) +}) diff --git a/internal/connect/methods.go b/internal/connect/methods.go new file mode 100644 index 0000000..5b2dfe0 --- /dev/null +++ b/internal/connect/methods.go @@ -0,0 +1,132 @@ +package connect + +import ( + "context" + "fmt" + "os/exec" + "strconv" + "strings" + "syscall" + + "github.com/fraser-isbester/tusk/internal/config" +) + +// openKubePortForward starts a `kubectl port-forward` to the profile's target +// and returns a DSN pointing at the local end of the forward. +func openKubePortForward(ctx context.Context, profile config.Profile) (string, func() error, error) { + c := profile.Connect + if c.Target == "" { + return "", nil, fmt.Errorf("kube-port-forward requires a target (e.g. svc/postgres)") + } + remotePort := c.RemotePort + if remotePort == 0 { + remotePort = 5432 + } + localPort := c.LocalPort + if localPort == 0 { + p, err := freePort() + if err != nil { + return "", nil, fmt.Errorf("allocating local port: %w", err) + } + localPort = p + } + + args := kubectlArgs(c, localPort, remotePort) + dsn := profile.DSN("127.0.0.1", localPort) + target := fmt.Sprintf("%s:%d", c.Target, remotePort) + return spawnTunnel(ctx, "kubectl", args, localPort, dsn, target) +} + +// kubectlArgs builds the kubectl port-forward argument vector. +func kubectlArgs(c *config.ConnectConfig, localPort, remotePort int) []string { + args := []string{"port-forward"} + if c.Context != "" { + args = append(args, "--context", c.Context) + } + if c.Namespace != "" { + args = append(args, "-n", c.Namespace) + } + args = append(args, c.Target, fmt.Sprintf("%d:%d", localPort, remotePort)) + return args +} + +// openExec runs an arbitrary tunnel command (SSH, cloud-sql-proxy, …) and +// returns a DSN pointing at the local port. The token {local_port} in the +// command is replaced with the chosen port. +func openExec(ctx context.Context, profile config.Profile) (string, func() error, error) { + c := profile.Connect + if len(c.Command) == 0 { + return "", nil, fmt.Errorf("exec connect method requires a command") + } + localPort := c.LocalPort + if localPort == 0 { + p, err := freePort() + if err != nil { + return "", nil, fmt.Errorf("allocating local port: %w", err) + } + localPort = p + } + + argv := substituteLocalPort(c.Command, localPort) + dsn := profile.DSN("127.0.0.1", localPort) + return spawnTunnel(ctx, argv[0], argv[1:], localPort, dsn, argv[0]) +} + +// substituteLocalPort replaces every {local_port} token in argv with port. +func substituteLocalPort(argv []string, port int) []string { + out := make([]string, len(argv)) + for i, a := range argv { + out[i] = strings.ReplaceAll(a, "{local_port}", strconv.Itoa(port)) + } + return out +} + +// spawnTunnel starts cmd, waits for the local port to accept connections, and +// returns the DSN plus a closer that kills the process group. If the process +// exits or the port never opens within establishTimeout, it kills the process +// and returns an error that names the target. +func spawnTunnel(ctx context.Context, name string, args []string, localPort int, dsn, target string) (string, func() error, error) { + cmd := exec.Command(name, args...) //nolint:gosec // command/args come from the user's own config + // Run in its own process group so we can kill the whole tree on teardown + // (kubectl and proxies spawn children). + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + var stderr strings.Builder + cmd.Stderr = &stderr + + if err := cmd.Start(); err != nil { + return "", nil, fmt.Errorf("starting %s: %w", name, err) + } + + kill := func() error { + if cmd.Process == nil { + return nil + } + // Negative PID signals the whole process group. + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM) + _, _ = cmd.Process.Wait() + return nil + } + + // Watch for early exit so waitForPort can fail fast with the reason. + dead := make(chan error, 1) + go func() { dead <- cmd.Wait() }() + + addr := fmt.Sprintf("127.0.0.1:%d", localPort) + if err := waitForPort(ctx, addr, establishTimeout, dead); err != nil { + _ = kill() + msg := strings.TrimSpace(stderr.String()) + if msg != "" { + return "", nil, fmt.Errorf("%s tunnel to %s failed: %w: %s", name, target, err, lastLine(msg)) + } + return "", nil, fmt.Errorf("%s tunnel to %s failed: %w", name, target, err) + } + + return dsn, kill, nil +} + +// lastLine returns the last non-empty line of s, the most useful part of a +// tool's stderr for an error message. +func lastLine(s string) string { + lines := strings.Split(strings.TrimRight(s, "\n"), "\n") + return lines[len(lines)-1] +} diff --git a/internal/tui/setup.go b/internal/tui/setup.go new file mode 100644 index 0000000..dc62e66 --- /dev/null +++ b/internal/tui/setup.go @@ -0,0 +1,215 @@ +package tui + +import ( + "context" + "fmt" + "strconv" + "strings" + "time" + + "github.com/gdamore/tcell/v2" + "github.com/rivo/tview" + + "github.com/fraser-isbester/tusk/internal/config" + "github.com/fraser-isbester/tusk/internal/connect" + "github.com/fraser-isbester/tusk/internal/db" + "github.com/fraser-isbester/tusk/internal/tui/theme" +) + +// SetupResult is a working connection produced by the first-run setup screen. +type SetupResult struct { + DB *db.DB + Closer func() error // tears down any tunnel; call after DB.Close + ProfileName string + Profile config.Profile +} + +// RunSetup shows a first-run connection screen and returns a validated, live +// connection once the user connects successfully. It returns ok=false if the +// user quits without connecting. reason explains why setup was launched (e.g. +// "profile \"dev\" not found") and is shown at the top. +func RunSetup(cfg *config.Config, reason string) (*SetupResult, bool) { + app := tview.NewApplication() + form := tview.NewForm() + form.SetBackgroundColor(tcell.ColorDefault) + form.SetFieldBackgroundColor(tcell.NewRGBColor(0x1a, 0x1a, 0x1a)) + form.SetFieldTextColor(tcell.ColorWhite) + form.SetLabelColor(theme.ColorLabel) + form.SetButtonBackgroundColor(theme.ColorBorder) + form.SetButtonTextColor(tcell.ColorWhite) + form.SetBorder(true).SetBorderColor(theme.ColorBorder) + form.SetTitle(" Connect to PostgreSQL ").SetTitleColor(theme.ColorLogo) + + status := tview.NewTextView().SetDynamicColors(true) + status.SetBackgroundColor(tcell.ColorDefault) + if reason != "" { + status.SetText(fmt.Sprintf(" [#D78700]%s — set up a connection to continue[-]", reason)) + } else { + status.SetText(" [#808080]Enter connection details, then Connect[-]") + } + + methods := []string{"direct", "kube-port-forward", "exec"} + + // Field values, kept in closure vars so the form can be rebuilt per method + // without losing input. Default the profile name to the one that was + // expected but unusable, so saving fixes the dangling default_profile. + name := cfg.DefaultProfile + if name == "" { + name = "default" + } + method := "direct" + url := "" + kctx, kns, ktarget, kremote := "", "", "", "5432" + command := "" + user, pass, dbname, sslmode := "", "", "", "disable" + save := true + + var result *SetupResult + var rebuild func() + + setStatus := func(color, msg string) { status.SetText(fmt.Sprintf(" [%s]%s[-]", color, msg)) } + + connectFn := func() { + profile := config.Profile{} + switch method { + case "direct": + if strings.TrimSpace(url) == "" { + setStatus("#FF5F5F", "Connection URL is required") + return + } + profile.URL = url + case "kube-port-forward": + if strings.TrimSpace(ktarget) == "" { + setStatus("#FF5F5F", "Target is required (e.g. svc/postgres)") + return + } + profile.User, profile.Password, profile.Database, profile.SSLMode = user, pass, dbname, sslmode + remote, _ := strconv.Atoi(strings.TrimSpace(kremote)) + profile.Connect = &config.ConnectConfig{ + Via: "kube-port-forward", Context: kctx, Namespace: kns, Target: ktarget, RemotePort: remote, + } + case "exec": + if strings.TrimSpace(command) == "" { + setStatus("#FF5F5F", "Command is required") + return + } + profile.User, profile.Password, profile.Database, profile.SSLMode = user, pass, dbname, sslmode + profile.Connect = &config.ConnectConfig{Via: "exec", Command: strings.Fields(command)} + } + + setStatus("#808080", "Connecting…") + app.ForceDraw() + + dsn, closer, err := connect.Open(context.Background(), profile) + if err != nil { + setStatus("#FF5F5F", err.Error()) + return + } + database, err := db.New(dsn) + if err != nil { + _ = closer() + setStatus("#FF5F5F", err.Error()) + return + } + pingCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := database.Pool().Ping(pingCtx); err != nil { + database.Close() + _ = closer() + setStatus("#FF5F5F", fmt.Sprintf("connection failed: %s", err.Error())) + return + } + + profileName := strings.TrimSpace(name) + if profileName == "" { + profileName = "default" + } + if save { + if err := config.SaveProfile(profileName, profile); err != nil { + setStatus("#FF5F5F", fmt.Sprintf("connected, but saving profile failed: %s", err.Error())) + // Keep the connection — persistence is best-effort. + } else { + if cfg.Profiles == nil { + cfg.Profiles = map[string]config.Profile{} + } + cfg.Profiles[profileName] = profile + if cfg.DefaultProfile == "" { + cfg.DefaultProfile = profileName + } + } + } + + result = &SetupResult{DB: database, Closer: closer, ProfileName: profileName, Profile: profile} + app.Stop() + } + + rebuild = func() { + form.Clear(true) + form.AddInputField("Profile name", name, 30, nil, func(t string) { name = t }) + form.AddDropDown("Connect via", methods, indexOf(methods, method), func(opt string, _ int) { + if opt != method { + method = opt + rebuild() + } + }) + switch method { + case "direct": + form.AddInputField("Connection URL", url, 0, nil, func(t string) { url = t }) + case "kube-port-forward": + form.AddInputField("Context", kctx, 0, nil, func(t string) { kctx = t }) + form.AddInputField("Namespace", kns, 0, nil, func(t string) { kns = t }) + form.AddInputField("Target (svc/pod/…)", ktarget, 0, nil, func(t string) { ktarget = t }) + form.AddInputField("Remote port", kremote, 10, nil, func(t string) { kremote = t }) + addCredFields(form, &user, &pass, &dbname, &sslmode) + case "exec": + form.AddInputField("Command ({local_port})", command, 0, nil, func(t string) { command = t }) + addCredFields(form, &user, &pass, &dbname, &sslmode) + } + form.AddCheckbox("Save profile", save, func(c bool) { save = c }) + form.AddButton("Connect", connectFn) + form.AddButton("Quit", func() { app.Stop() }) + app.SetFocus(form) + } + rebuild() + + layout := tview.NewFlex().SetDirection(tview.FlexRow). + AddItem(nil, 0, 1, false). + AddItem(tview.NewFlex(). + AddItem(nil, 0, 1, false). + AddItem(tview.NewFlex().SetDirection(tview.FlexRow). + AddItem(form, 0, 1, true). + AddItem(status, 1, 0, false), 74, 0, true). + AddItem(nil, 0, 1, false), 24, 0, true). + AddItem(nil, 0, 1, false) + layout.SetBackgroundColor(tcell.ColorDefault) + + app.SetInputCapture(func(evt *tcell.EventKey) *tcell.EventKey { + if evt.Key() == tcell.KeyEscape { + app.Stop() + return nil + } + return evt + }) + + if err := app.SetRoot(layout, true).EnableMouse(true).Run(); err != nil { + return nil, false + } + return result, result != nil +} + +// addCredFields adds the shared credential/database fields used by tunnel methods. +func addCredFields(form *tview.Form, user, pass, dbname, sslmode *string) { + form.AddInputField("DB user", *user, 30, nil, func(t string) { *user = t }) + form.AddPasswordField("DB password", *pass, 30, '*', func(t string) { *pass = t }) + form.AddInputField("Database", *dbname, 30, nil, func(t string) { *dbname = t }) + form.AddInputField("SSL mode", *sslmode, 20, nil, func(t string) { *sslmode = t }) +} + +func indexOf(ss []string, s string) int { + for i, v := range ss { + if v == s { + return i + } + } + return 0 +} From f5e81527a9044b2219efa4eda909f859e70ce240 Mon Sep 17 00:00:00 2001 From: Fraser Isbester Date: Wed, 8 Jul 2026 11:32:27 -0700 Subject: [PATCH 2/2] fix: bump go directive to 1.26.5 for crypto/tls advisory (GO-2026-5856) --- go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go.mod b/go.mod index 29adc6b..f2cffaa 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/fraser-isbester/tusk -go 1.26.4 +go 1.26.5 require ( github.com/gdamore/tcell/v2 v2.13.9