add support for first-n-connections

pkg/client/cli/intercept/command.go

@@ -18,13 +18,14 @@ import (
 )
 
 type Command struct {
-	Name           string // Command[0] || `${Command[0]}-${--namespace}` // which depends on a combinationof --workload and --namespace
-	AgentName      string // --workload || Command[0] // only valid if !localOnly
-	Port           string // --port // only valid if !localOnly
-	ServiceName    string // --service // only valid if !localOnly
-	Address        string // --address // only valid if !localOnly
-	LocalOnly      bool   // --local-only
-	LocalMountPort uint16 // --local-mount-port
+	Name              string // Command[0] || `${Command[0]}-${--namespace}` // which depends on a combinationof --workload and --namespace
+	AgentName         string // --workload || Command[0] // only valid if !localOnly
+	Port              string // --port // only valid if !localOnly
+	ServiceName       string // --service // only valid if !localOnly
+	Address           string // --address // only valid if !localOnly
+	LocalOnly         bool   // --local-only
+	LocalMountPort    uint16 // --local-mount-port
+	FirstNConnections uint32 // --first-n-connections
 
 	Replace bool // whether --replace was passed
 
@@ -117,7 +118,8 @@ func (a *Command) AddFlags(cmd *cobra.Command) {
 	flagSet.BoolVarP(&a.Replace, "replace", "", false,
 		`Indicates if the traffic-agent should replace application containers in workload pods. `+
 			`The default behavior is for the agent sidecar to be installed alongside existing containers.`)
-
+	// the default value of 0 is used to indicate that all connections should be intercepted
+	flagSet.Uint32Var(&a.FirstNConnections, "first-n-connections", 0, `Limit the number of connections to intercept`)
 	// Hide these flags. They are still functional but deprecated. Using them will yield a deprecation message.
 	flagSet.Lookup("local-only").Hidden = true
 	flagSet.Lookup("namespace").Hidden = true

pkg/client/userd/trafficmgr/intercept.go

@@ -68,6 +68,9 @@ type intercept struct {
 
 	// Use bridged ftp/sftp mount through this local port
 	localMountPort int32
+
+	// nConnections is the number of connections that have been intercepted
+	nConnections uint32
 }
 
 // interceptResult is what gets written to the awaitIntercept's waitCh channel when the
@@ -140,7 +143,17 @@ func (ic *intercept) localPorts() []string {
 }
 
 func (ic *intercept) shouldForward() bool {
-	return len(ic.localPorts()) > 0
+	accumulator := true
+	b := []bool{}
+	b = append(b, len(ic.localPorts()) > 0)
+	if ic.Spec.FirstNConnections > 0 {
+		b = append(b, ic.nConnections < ic.Spec.FirstNConnections)
+
+	}
+	for _, v := range b {
+		accumulator = accumulator && v
+	}
+	return accumulator
 }
 
 // startForwards starts port forwards and mounts for the given podInterceptKey.
@@ -241,6 +254,7 @@ func (lpf *podIntercepts) start(ctx context.Context, ic *intercept, rd daemon.Da
 	}
 	if ic.shouldForward() {
 		ic.startForwards(ctx, &lp.wg)
+		ic.nConnections++
 	}
 	dlog.Debugf(ctx, "Started mounts and port-forwards for %+v", fk)
 	lpf.alivePods[fk] = lp

rpc/manager/manager.proto

@@ -142,6 +142,9 @@ message InterceptSpec {
 
   // Whether to replace the running container.
   bool replace = 22;
+
+  // How many connections should be used for the intercept
+  uint32 first_n_connections = 23;
 }
 
 enum InterceptDispositionType {