nebula/diag/registry_test.go

package diag

import (
	"bytes"
	"flag"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type testFlags struct {
	Json bool
}

// testCommand builds a command carrying a flag set, recording what the callback was actually
// handed so a test can assert on it.
func testCommand(name string, seen *any, args *[]string) *Command {
	return &Command{
		Name:             name,
		ShortDescription: name + " short description",
		Flags: func() (*flag.FlagSet, any) {
			fl := flag.NewFlagSet("", flag.ContinueOnError)
			f := &testFlags{}
			fl.BoolVar(&f.Json, "json", false, "outputs json")
			return fl, f
		},
		Callback: func(fs any, a []string, w StringWriter) error {
			if seen != nil {
				*seen = fs
			}
			if args != nil {
				*args = a
			}
			return w.WriteLine("ran " + name)
		},
	}
}

func newTestRegistry(t *testing.T) (*Registry, *bytes.Buffer, StringWriter) {
	t.Helper()
	buf := &bytes.Buffer{}
	return NewRegistry(), buf, NewWriter(buf)
}

func TestRegistryDispatch(t *testing.T) {
	t.Run("a new registry knows help and nothing else", func(t *testing.T) {
		r, buf, w := newTestRegistry(t)

		require.NoError(t, r.DispatchArgs([]string{"help"}, w))
		assert.Contains(t, buf.String(), "help -")
	})

	t.Run("empty args dump the command list, matching an empty line on the console", func(t *testing.T) {
		r, buf, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", nil, nil))

		require.NoError(t, r.DispatchArgs(nil, w))
		assert.Contains(t, buf.String(), "Available commands:")
		assert.Contains(t, buf.String(), "do-thing - do-thing short description")
	})

	t.Run("an unknown command reports ErrUnknownCommand and still tells the user", func(t *testing.T) {
		r, buf, w := newTestRegistry(t)

		err := r.DispatchArgs([]string{"nope"}, w)
		require.ErrorIs(t, err, ErrUnknownCommand)
		assert.Contains(t, buf.String(), "Did not understand: nope")
		assert.Contains(t, buf.String(), "Available commands:")
	})

	// This is the hazard the ctl transport has to preserve: every callback in ssh.go begins by
	// type asserting fs to its own concrete flags struct. Reach a callback without going
	// through Command.Flags and every one of them fails.
	t.Run("a callback is handed the concrete struct its Flags callback returned", func(t *testing.T) {
		var seen any
		r, _, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", &seen, nil))

		require.NoError(t, r.DispatchArgs([]string{"do-thing", "-json"}, w))

		flags, ok := seen.(*testFlags)
		require.True(t, ok, "callback was handed %T, not *testFlags", seen)
		assert.True(t, flags.Json)
	})

	t.Run("positional arguments survive flag parsing", func(t *testing.T) {
		var args []string
		r, _, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", nil, &args))

		require.NoError(t, r.DispatchArgs([]string{"do-thing", "-json", "10.0.0.1"}, w))
		assert.Equal(t, []string{"10.0.0.1"}, args)
	})

	// Documents stdlib flag behaviour rather than endorsing it: parsing stops at the first
	// positional, so a flag written after one is silently a positional too.
	t.Run("a flag after a positional is not parsed as a flag", func(t *testing.T) {
		var seen any
		var args []string
		r, _, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", &seen, &args))

		require.NoError(t, r.DispatchArgs([]string{"do-thing", "10.0.0.1", "-json"}, w))
		assert.False(t, seen.(*testFlags).Json)
		assert.Equal(t, []string{"10.0.0.1", "-json"}, args)
	})

	t.Run("a bad flag reports ErrUsage and writes the usage text", func(t *testing.T) {
		r, buf, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", nil, nil))

		err := r.DispatchArgs([]string{"do-thing", "-nope"}, w)
		require.ErrorIs(t, err, ErrUsage)
		assert.Contains(t, buf.String(), "flag provided but not defined")
	})

	t.Run("-h anywhere routes to help instead of running the command", func(t *testing.T) {
		var seen any
		r, buf, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", &seen, nil))

		require.NoError(t, r.DispatchArgs([]string{"do-thing", "-h"}, w))
		assert.Nil(t, seen, "the callback should not have run")
		assert.Contains(t, buf.String(), "do-thing - do-thing short description")
		assert.Contains(t, buf.String(), "-json")
	})

	t.Run("Dispatch splits a line the way a shell would", func(t *testing.T) {
		var args []string
		r, _, w := newTestRegistry(t)
		r.RegisterCommand(testCommand("do-thing", nil, &args))

		require.NoError(t, r.Dispatch(`do-thing "/tmp/a path.pb.gz"`, w))
		assert.Equal(t, []string{"/tmp/a path.pb.gz"}, args)
	})

	t.Run("Match returns names by prefix for tab completion", func(t *testing.T) {
		r, _, _ := newTestRegistry(t)
		r.RegisterCommand(testCommand("print-cert", nil, nil))
		r.RegisterCommand(testCommand("print-tunnel", nil, nil))
		r.RegisterCommand(testCommand("version", nil, nil))

		assert.Equal(t, []string{"print-cert", "print-tunnel"}, r.Match("print-"))
	})
}

// A clone is what keeps the ssh session's `logout` command from being visible to every other
// session, and to nebula ctl.
func TestRegistryCloneIsolation(t *testing.T) {
	parent, _, w := newTestRegistry(t)
	parent.RegisterCommand(testCommand("shared", nil, nil))

	child := parent.Clone()
	child.RegisterCommand(testCommand("logout", nil, nil))

	require.NoError(t, child.DispatchArgs([]string{"logout"}, w))

	buf := &bytes.Buffer{}
	err := parent.DispatchArgs([]string{"logout"}, NewWriter(buf))
	assert.ErrorIs(t, err, ErrUnknownCommand)

	buf.Reset()
	require.NoError(t, child.DispatchArgs([]string{"shared"}, NewWriter(buf)))
	assert.True(t, strings.HasPrefix(buf.String(), "ran shared"))
}