mc-proxy/cmd/mcproxyctl/root.go

package main

import (
	"context"
	"os"

	"github.com/spf13/cobra"

	"git.wntrmute.dev/mc/mc-proxy/client/mcproxy"
)

const defaultSocketPath = "/srv/mc-proxy/mc-proxy.sock"

type contextKey string

const clientKey contextKey = "client"

func rootCmd() *cobra.Command {
	var socketPath string

	cmd := &cobra.Command{
		Use:   "mcproxyctl",
		Short: "Control a running mc-proxy instance",
		Long:  "mcproxyctl is a command-line tool for administrating a running mc-proxy instance via the gRPC admin API.",
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			// Skip client setup for help commands
			if cmd.Name() == "help" || cmd.Name() == "completion" {
				return nil
			}

			// Environment variable override
			if envSocket := os.Getenv("MCPROXY_SOCKET"); envSocket != "" && socketPath == defaultSocketPath {
				socketPath = envSocket
			}

			client, err := mcproxy.Dial(socketPath)
			if err != nil {
				return err
			}

			ctx := context.WithValue(cmd.Context(), clientKey, client)
			cmd.SetContext(ctx)
			return nil
		},
		PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
			client := clientFromContext(cmd.Context())
			if client != nil {
				return client.Close()
			}
			return nil
		},
		SilenceUsage: true,
	}

	cmd.PersistentFlags().StringVarP(&socketPath, "socket", "s", defaultSocketPath, "Unix socket path")

	cmd.AddCommand(statusCmd())
	cmd.AddCommand(healthCmd())
	cmd.AddCommand(routesCmd())
	cmd.AddCommand(firewallCmd())
	cmd.AddCommand(policiesCmd())

	return cmd
}

func clientFromContext(ctx context.Context) *mcproxy.Client {
	if ctx == nil {
		return nil
	}
	client, _ := ctx.Value(clientKey).(*mcproxy.Client)
	return client
}