cmd/tetra: tracing policy commands rework

Commands were implemented in a way that most of them would not return an exit code != 0 on error and no messages on success. This fixes the issues for the tracing policy subcommands and cleanup the implem. Signed-off-by: Mahe Tardy <[email protected]>
cilium · Oct 6, 2023 · 7a4ea3a · 7a4ea3a
1 parent 18e4632
commit 7a4ea3a
Show file tree

Hide file tree

Showing 3 changed files with 142 additions and 96 deletions.
diff --git a/cmd/tetra/common/client.go b/cmd/tetra/common/client.go
@@ -115,3 +115,48 @@ func CliRunErr(fn func(ctx context.Context, cli tetragon.FineGuidanceSensorsClie
 func CliRun(fn func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient)) {
 	CliRunErr(fn, func(_ error) {})
 }
+
+type ConnectedClient struct {
+	Client tetragon.FineGuidanceSensorsClient
+	Ctx    context.Context
+	conn   *grpc.ClientConn
+	cancel context.CancelFunc
+}
+
+// Close cleanup resources, it closes the connection and cancel the context
+func (c ConnectedClient) Close() {
+	c.conn.Close()
+	c.cancel()
+}
+
+// NewConnectedClient return a connected client to a tetragon server, caller
+// must call Close() on the client. On failure to connect, this function calls
+// Fatal() thus stopping execution.
+func NewConnectedClient() ConnectedClient {
+	c := ConnectedClient{}
+	c.Ctx, c.cancel = signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+
+	var serverAddr string
+	var err error
+
+	backoff := time.Second
+	attempts := 0
+	for {
+		c.conn, serverAddr, err = connect(c.Ctx)
+		if err != nil {
+			if attempts < viper.GetInt(KeyRetries) {
+				// Exponential backoff
+				attempts++
+				logger.GetLogger().WithField("server-address", serverAddr).WithField("attempts", attempts).WithError(err).Error("Connection attempt failed, retrying...")
+				time.Sleep(backoff)
+				backoff *= 2
+				continue
+			}
+			logger.GetLogger().WithField("server-address", serverAddr).WithField("attempts", attempts).WithError(err).Fatal("Failed to connect to server")
+		}
+		break
+	}
+
+	c.Client = tetragon.NewFineGuidanceSensorsClient(c.conn)
+	return c
+}
diff --git a/cmd/tetra/main.go b/cmd/tetra/main.go
@@ -25,8 +25,9 @@ func main() {
 
 func New() *cobra.Command {
 	rootCmd = &cobra.Command{
-		Use:   "tetra",
-		Short: "Tetragon CLI",
+		Use:          "tetra",
+		Short:        "Tetragon CLI",
+		SilenceUsage: true,
 		Run: func(cmd *cobra.Command, args []string) {
 			cmd.Help()
 		},

diff --git a/cmd/tetra/tracingpolicy/tracingpolicy.go b/cmd/tetra/tracingpolicy/tracingpolicy.go
@@ -4,15 +4,13 @@
 package tracingpolicy
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"strings"
 
 	"github.com/cilium/tetragon/api/v1/tetragon"
 	"github.com/cilium/tetragon/cmd/tetra/common"
 	"github.com/cilium/tetragon/cmd/tetra/tracingpolicy/generate"
-	"github.com/cilium/tetragon/pkg/logger"
 	"github.com/spf13/cobra"
 )
 
@@ -27,43 +25,88 @@ func New() *cobra.Command {
 		Use:   "add <yaml_file>",
 		Short: "add a new sensor based on a tracing policy",
 		Args:  cobra.ExactArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			common.CliRun(func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient) {
-				addTracingPolicy(ctx, cli, args[0])
+		RunE: func(cmd *cobra.Command, args []string) error {
+			c := common.NewConnectedClient()
+			defer c.Close()
+
+			yamlb, err := os.ReadFile(args[0])
+			if err != nil {
+				return fmt.Errorf("failed to read yaml file %s: %w", args[0], err)
+			}
+
+			_, err = c.Client.AddTracingPolicy(c.Ctx, &tetragon.AddTracingPolicyRequest{
+				Yaml: string(yamlb),
 			})
+			if err != nil {
+				return fmt.Errorf("failed to add tracing policy: %w", err)
+			}
+			cmd.Printf("tracing policy %q added\n", args[0])
+
+			return nil
 		},
 	}
 
 	tpDelCmd := &cobra.Command{
 		Use:   "delete <name>",
 		Short: "delete a tracing policy",
 		Args:  cobra.ExactArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			common.CliRun(func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient) {
-				deleteTracingPolicy(ctx, cli, args[0])
+		RunE: func(cmd *cobra.Command, args []string) error {
+			c := common.NewConnectedClient()
+			defer c.Close()
+
+			_, err := c.Client.DeleteTracingPolicy(c.Ctx, &tetragon.DeleteTracingPolicyRequest{
+				Name: args[0],
 			})
+			if err != nil {
+				return fmt.Errorf("failed to delete tracing policy: %w", err)
+			}
+			cmd.Printf("tracing policy %q deleted\n", args[0])
+
+			return nil
 		},
 	}
 
 	tpEnableCmd := &cobra.Command{
 		Use:   "enable <name>",
 		Short: "enable a tracing policy",
+		Long:  "Enable a disabled tracing policy. Use disable to re-disable the tracing policy.",
 		Args:  cobra.ExactArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			common.CliRun(func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient) {
-				enableTracingPolicy(ctx, cli, args[0])
+		RunE: func(cmd *cobra.Command, args []string) error {
+			c := common.NewConnectedClient()
+			defer c.Close()
+
+			_, err := c.Client.EnableTracingPolicy(c.Ctx, &tetragon.EnableTracingPolicyRequest{
+				Name: args[0],
 			})
+			if err != nil {
+				return fmt.Errorf("failed to enable tracing policy: %w", err)
+			}
+			cmd.Printf("tracing policy %q enabled\n", args[0])
+
+			return nil
 		},
 	}
 
 	tpDisableCmd := &cobra.Command{
 		Use:   "disable <name>",
 		Short: "disable a tracing policy",
+		Long:  "Disable an enabled tracing policy. Use enable to re-enable the tracing policy.",
 		Args:  cobra.ExactArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			common.CliRun(func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient) {
-				disableTracingPolicy(ctx, cli, args[0])
+		RunE: func(cmd *cobra.Command, args []string) error {
+			c := common.NewConnectedClient()
+			defer c.Close()
+
+			_, err := c.Client.DisableTracingPolicy(c.Ctx, &tetragon.DisableTracingPolicyRequest{
+				Name: args[0],
 			})
+
+			if err != nil {
+				return fmt.Errorf("failed to disable tracing policy: %w", err)
+			}
+
+			cmd.Printf("tracing policy %q disabled\n", args[0])
+
+			return nil
 		},
 	}
 
@@ -78,90 +121,47 @@ func New() *cobra.Command {
 			}
 			return nil
 		},
-		Run: func(cmd *cobra.Command, args []string) {
-			common.CliRun(func(ctx context.Context, cli tetragon.FineGuidanceSensorsClient) {
-				listTracingPolicies(ctx, cli, tpListOutputFlag)
-			})
-		},
-	}
-	tpListFlags := tpListCmd.Flags()
-	tpListFlags.StringVarP(&tpListOutputFlag, common.KeyOutput, "o", "text", "Output format. text or json")
-
-	tpCmd.AddCommand(tpAddCmd, tpDelCmd, tpEnableCmd, tpDisableCmd, tpListCmd, generate.New())
-	return tpCmd
-}
-
-func addTracingPolicy(ctx context.Context, client tetragon.FineGuidanceSensorsClient, yamlFname string) {
-	yamlb, err := os.ReadFile(yamlFname)
-	if err != nil {
-		fmt.Printf("failed to read yaml file %s: %s\n", yamlFname, err)
-		return
-	}
-
-	_, err = client.AddTracingPolicy(ctx, &tetragon.AddTracingPolicyRequest{
-		Yaml: string(yamlb),
-	})
-	if err != nil {
-		fmt.Printf("failed to add tracing policy: %s\n", err)
-	}
-}
-
-func deleteTracingPolicy(ctx context.Context, client tetragon.FineGuidanceSensorsClient, name string) {
-	_, err := client.DeleteTracingPolicy(ctx, &tetragon.DeleteTracingPolicyRequest{
-		Name: name,
-	})
-	if err != nil {
-		fmt.Printf("failed to delete tracing policy: %s\n", err)
-	}
-}
-
-func enableTracingPolicy(ctx context.Context, client tetragon.FineGuidanceSensorsClient, name string) {
-	_, err := client.EnableTracingPolicy(ctx, &tetragon.EnableTracingPolicyRequest{
-		Name: name,
-	})
-	if err != nil {
-		fmt.Printf("failed to enable tracing policy: %s\n", err)
-	}
-}
+		RunE: func(cmd *cobra.Command, args []string) error {
+			c := common.NewConnectedClient()
+			defer c.Close()
 
-func disableTracingPolicy(ctx context.Context, client tetragon.FineGuidanceSensorsClient, name string) {
-	_, err := client.DisableTracingPolicy(ctx, &tetragon.DisableTracingPolicyRequest{
-		Name: name,
-	})
-	if err != nil {
-		fmt.Printf("failed to disable tracing policy: %s\n", err)
-	}
-}
+			res, err := c.Client.ListTracingPolicies(c.Ctx, &tetragon.ListTracingPoliciesRequest{})
+			if err != nil || res == nil {
+				return fmt.Errorf("failed to list tracing policies: %w", err)
+			}
 
-func listTracingPoliciesText(res *tetragon.ListTracingPoliciesResponse) {
-	for _, pol := range res.Policies {
-		namespace := pol.Namespace
-		if namespace == "" {
-			namespace = "(global)"
-		}
+			if tpListOutputFlag == "json" {
+				b, err := res.MarshalJSON()
+				if err != nil {
+					return fmt.Errorf("failed to generate json: %w", err)
+				}
+				cmd.Println(string(b))
+			} else {
+				for _, pol := range res.Policies {
+					namespace := pol.Namespace
+					if namespace == "" {
+						namespace = "(global)"
+					}
+
+					sensors := strings.Join(pol.Sensors, ",")
+					cmd.Printf("%d %s (%s) %s %s\n", pol.Id, pol.Name, pol.Info, namespace, sensors)
+				}
+			}
 
-		sensors := strings.Join(pol.Sensors, ",")
-		fmt.Printf("%d %s (%s) %s %s\n", pol.Id, pol.Name, pol.Info, namespace, sensors)
+			return nil
+		},
 	}
-}
-
-func listTracingPolicies(
-	ctx context.Context,
-	client tetragon.FineGuidanceSensorsClient,
-	output string) {
+	tpListFlags := tpListCmd.Flags()
+	tpListFlags.StringVarP(&tpListOutputFlag, common.KeyOutput, "o", "text", "Output format. text or json")
 
-	res, err := client.ListTracingPolicies(ctx, &tetragon.ListTracingPoliciesRequest{})
-	if err != nil || res == nil {
-		logger.GetLogger().WithError(err).Fatal("failed to list tracing policies")
-	}
+	tpCmd.AddCommand(
+		tpAddCmd,
+		tpDelCmd,
+		tpEnableCmd,
+		tpDisableCmd,
+		tpListCmd,
+		generate.New(),
+	)
 
-	if output == "json" {
-		b, err := res.MarshalJSON()
-		if err != nil {
-			logger.GetLogger().WithError(err).Fatal("failed to generate json")
-		}
-		fmt.Println(string(b))
-	} else {
-		listTracingPoliciesText(res)
-	}
+	return tpCmd
 }