wip: spin egress container for each mcp server

Closes: #124

Author: yrobla

cmd/thv/app/rm.go

@@ -6,6 +6,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/stacklok/toolhive/pkg/lifecycle"
+	"github.com/stacklok/toolhive/pkg/logger"
 )
 
 var rmCmd = &cobra.Command{
@@ -37,9 +38,33 @@ func rmCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 
 	// Delete container.
-	if err := manager.DeleteContainer(ctx, containerName, rmForce); err != nil {
+	if err := manager.DeleteContainer(ctx, containerName, rmForce, true); err != nil {
 		return fmt.Errorf("failed to delete container: %v", err)
 	}
 
+	// Delete associated egress container.
+	egressContainerName := containerName + "-egress"
+	if err := manager.DeleteContainer(ctx, egressContainerName, rmForce, false); err != nil {
+		// just log the error and continue
+		logger.Warnf("failed to delete egress container %q: %v", egressContainerName, err)
+	}
+
+	// Delete networks if there are no containers using them.
+	toolHiveContainers, err := manager.ListContainers(ctx, listAll)
+	if err != nil {
+		return fmt.Errorf("failed to list containers: %v", err)
+	}
+	fmt.Println("ToolHive containers:", toolHiveContainers)
+
+	if len(toolHiveContainers) == 0 {
+		// remove networks
+		if err := manager.DeleteNetwork(ctx, "toolhive-internal"); err != nil {
+			return fmt.Errorf("failed to delete network: %v", err)
+		}
+		if err := manager.DeleteNetwork(ctx, "toolhive-external"); err != nil {
+			return fmt.Errorf("failed to delete network: %v", err)
+		}
+	}
+
 	return nil
 }

cmd/thv/app/run.go

@@ -245,6 +245,11 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to retrieve or pull image: %v", err)
 	}
 
+	// pull the egress image if it is not already pulled
+	if err := pullImage(ctx, config.EgressImage, rt); err != nil {
+		return fmt.Errorf("failed to retrieve or pull egress image: %v", err)
+	}
+
 	// Configure the RunConfig with transport, ports, permissions, etc.
 	if err := configureRunConfig(config, runTransport, runPort, runTargetPort, runEnv); err != nil {
 		return err

cmd/thv/app/stop.go

@@ -46,5 +46,16 @@ func stopCmdFunc(cmd *cobra.Command, args []string) error {
 		}
 	}
 
+	// Stop associated egress container
+	egressContainerName := containerName + "-egress"
+	err = manager.StopContainer(ctx, egressContainerName)
+	if err != nil {
+		if errors.Is(err, lifecycle.ErrContainerNotFound) {
+			logger.Infof("Egress container %s is not running", egressContainerName)
+		} else {
+			return fmt.Errorf("failed to stop egress container %q: %w", egressContainerName, err)
+		}
+	}
+
 	return nil
 }

pkg/api/v1/servers.go

@@ -145,7 +145,7 @@ func (s *ServerRoutes) deleteServer(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	name := chi.URLParam(r, "name")
 	forceDelete := r.URL.Query().Get("force") == "true"
-	err := s.manager.DeleteContainer(ctx, name, forceDelete)
+	err := s.manager.DeleteContainer(ctx, name, forceDelete, true)
 	if err != nil {
 		if errors.Is(err, lifecycle.ErrContainerNotFound) {
 			http.Error(w, "Server not found", http.StatusNotFound)

pkg/container/docker/client.go

@@ -80,8 +80,8 @@ func NewClient(ctx context.Context) (*Client, error) {
 
 		c, err := NewClientWithSocketPath(ctx, socketPath, runtimeType)
 		if err != nil {
-			logger.Debugf("Failed to create client for %s: %v", sp, err)
 			lastErr = err
+			logger.Debugf("Failed to create client for %s: %v", sp, err)
 			continue
 		}
 
@@ -989,6 +989,11 @@ func (c *Client) getPermissionConfigFromProfile(
 		return nil, fmt.Errorf("unsupported transport type: %s", transportType)
 	}
 
+	// Add necessary capabilities for egress containers
+	if profile.Name == permissions.ProfileEgress {
+		config.CapAdd = append(config.CapAdd, "CAP_SETUID", "CAP_SETGID")
+	}
+
 	return config, nil
 }
 
@@ -1334,3 +1339,54 @@ func (c *Client) handleExistingContainer(
 	// Container was removed and needs to be recreated
 	return false, nil
 }
+
+// CreateNetwork creates a network following configuration.
+func (c *Client) CreateNetwork(
+	ctx context.Context,
+	name string,
+	labels map[string]string,
+	internal bool,
+) (string, error) {
+	// Check if the network already exists
+	networks, err := c.client.NetworkList(ctx, network.ListOptions{
+		Filters: filters.NewArgs(filters.Arg("name", name)),
+	})
+	if err != nil {
+		return "", fmt.Errorf("failed to list networks: %w", err)
+	}
+	if len(networks) > 0 {
+		// Network already exists, return its ID
+		return networks[0].ID, nil
+	}
+
+	networkCreate := network.CreateOptions{
+		Driver:   "bridge",
+		Internal: internal,
+		Labels:   labels,
+	}
+
+	resp, err := c.client.NetworkCreate(ctx, name, networkCreate)
+	if err != nil {
+		return "", err
+	}
+	return resp.ID, nil
+}
+
+// DeleteNetwork deletes a network by name.
+func (c *Client) DeleteNetwork(ctx context.Context, name string) error {
+	// find the network by name
+	networks, err := c.client.NetworkList(ctx, network.ListOptions{
+		Filters: filters.NewArgs(filters.Arg("name", name)),
+	})
+	if err != nil {
+		return err
+	}
+	if len(networks) == 0 {
+		return fmt.Errorf("network %s not found", name)
+	}
+
+	if err := c.client.NetworkRemove(ctx, networks[0].ID); err != nil {
+		return fmt.Errorf("failed to remove network %s: %w", name, err)
+	}
+	return nil
+}

pkg/container/kubernetes/client.go

@@ -15,6 +15,8 @@ import (
 	backoff "github.com/cenkalti/backoff/v4"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	networkingv1 "k8s.io/api/networking/v1"
+
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -567,6 +569,85 @@ func (*Client) StopContainer(_ context.Context, _ string) error {
 	return nil
 }
 
+func (c *Client) CreateNetwork(ctx context.Context, name string, labels map[string]string, internal bool) (string, error) {
+	namespace := getCurrentNamespace()
+
+	// Check if the NetworkPolicy already exists
+	_, err := c.client.NetworkingV1().NetworkPolicies(namespace).Get(ctx, name, metav1.GetOptions{})
+	if err == nil {
+		return name, nil // NetworkPolicy already exists
+	}
+	if !errors.IsNotFound(err) {
+		return "", fmt.Errorf("failed to check if NetworkPolicy exists: %w", err)
+	}
+
+	// Define the NetworkPolicy spec based on the 'internal' flag
+	policyTypes := []networkingv1.PolicyType{networkingv1.PolicyTypeIngress}
+	var ingressRules []networkingv1.NetworkPolicyIngressRule
+
+	if internal {
+		// Restrict ingress to pods with the same labels
+		ingressRules = []networkingv1.NetworkPolicyIngressRule{
+			{
+				From: []networkingv1.NetworkPolicyPeer{
+					{
+						PodSelector: &metav1.LabelSelector{
+							MatchLabels: labels,
+						},
+					},
+				},
+			},
+		}
+	} else {
+		// Allow all ingress traffic
+		ingressRules = []networkingv1.NetworkPolicyIngressRule{
+			{
+				From: []networkingv1.NetworkPolicyPeer{
+					{
+						NamespaceSelector: &metav1.LabelSelector{},
+					},
+				},
+			},
+		}
+	}
+
+	// Create the NetworkPolicy object
+	policy := &networkingv1.NetworkPolicy{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+			Labels:    labels,
+		},
+		Spec: networkingv1.NetworkPolicySpec{
+			PodSelector: metav1.LabelSelector{
+				MatchLabels: labels,
+			},
+			PolicyTypes: policyTypes,
+			Ingress:     ingressRules,
+		},
+	}
+
+	// Create the NetworkPolicy in Kubernetes
+	_, err = c.client.NetworkingV1().NetworkPolicies(namespace).Create(ctx, policy, metav1.CreateOptions{})
+	if err != nil {
+		return "", fmt.Errorf("failed to create NetworkPolicy: %w", err)
+	}
+
+	return name, nil
+}
+
+func (c *Client) DeleteNetwork(ctx context.Context, name string) error {
+	namespace := getCurrentNamespace() // Implement this function to retrieve the current namespace
+
+	err := c.client.NetworkingV1().NetworkPolicies(namespace).Delete(ctx, name, metav1.DeleteOptions{})
+	if err != nil {
+		return fmt.Errorf("failed to delete network policy %q in namespace %q: %w", name, namespace, err)
+	}
+
+	fmt.Printf("Network policy %q in namespace %q has been deleted successfully.\n", name, namespace)
+	return nil
+}
+
 // waitForStatefulSetReady waits for a statefulset to be ready using the watch API
 func waitForStatefulSetReady(ctx context.Context, clientset kubernetes.Interface, namespace, name string) error {
 	// Create a field selector to watch only this specific statefulset

pkg/container/runtime/types.go

@@ -88,6 +88,12 @@ type Runtime interface {
 
 	// BuildImage builds a Docker image from a Dockerfile in the specified context directory
 	BuildImage(ctx context.Context, contextDir, imageName string) error
+
+	// CreateNetwork creates a network
+	CreateNetwork(ctx context.Context, networkName string, labels map[string]string, internal bool) (string, error)
+
+	// DeleteNetwork deletes a network
+	DeleteNetwork(ctx context.Context, networkName string) error
 }
 
 // Monitor defines the interface for container monitoring

pkg/labels/labels.go

@@ -43,6 +43,12 @@ func AddStandardLabels(labels map[string]string, containerName, containerBaseNam
 	labels[LabelToolType] = "mcp"
 }
 
+// AddNetworkLabels adds network-related labels to a network
+func AddNetworkLabels(labels map[string]string, networkName string) {
+	labels[LabelEnabled] = "true"
+	labels[LabelName] = networkName
+}
+
 // FormatToolHiveFilter formats a filter for ToolHive containers
 func FormatToolHiveFilter() string {
 	return fmt.Sprintf("%s=true", LabelEnabled)

pkg/lifecycle/manager.go

@@ -30,7 +30,9 @@ type Manager interface {
 	// ListContainers lists all ToolHive-managed containers.
 	ListContainers(ctx context.Context, listAll bool) ([]rt.ContainerInfo, error)
 	// DeleteContainer deletes a container and its associated proxy process.
-	DeleteContainer(ctx context.Context, name string, forceDelete bool) error
+	DeleteContainer(ctx context.Context, name string, forceDelete bool, removeConfig bool) error
+	// DeleteNetwork deletes a network.
+	DeleteNetwork(ctx context.Context, name string) error
 	// StopContainer stops a container and its associated proxy process.
 	StopContainer(ctx context.Context, name string) error
 	// RunContainer runs a container in the foreground.
@@ -86,7 +88,7 @@ func (d *defaultManager) ListContainers(ctx context.Context, listAll bool) ([]rt
 	return toolHiveContainers, nil
 }
 
-func (d *defaultManager) DeleteContainer(ctx context.Context, name string, forceDelete bool) error {
+func (d *defaultManager) DeleteContainer(ctx context.Context, name string, forceDelete bool, removeConfig bool) error {
 	// We need several fields from the container struct for deletion.
 	container, err := d.findContainerByName(ctx, name)
 	if err != nil {
@@ -114,24 +116,26 @@ func (d *defaultManager) DeleteContainer(ctx context.Context, name string, force
 		return fmt.Errorf("failed to remove container: %v", err)
 	}
 
-	// Get the base name from the container labels
-	baseName := labels.GetContainerBaseName(containerLabels)
-	if baseName != "" {
-		// Delete the saved state if it exists
-		if err := runner.DeleteSavedConfig(ctx, baseName); err != nil {
-			logger.Warnf("Warning: Failed to delete saved state: %v", err)
-		} else {
-			logger.Infof("Saved state for %s removed", baseName)
+	if removeConfig {
+		// Get the base name from the container labels
+		baseName := labels.GetContainerBaseName(containerLabels)
+		if baseName != "" {
+			// Delete the saved state if it exists
+			if err := runner.DeleteSavedConfig(ctx, baseName); err != nil {
+				logger.Warnf("Warning: Failed to delete saved state: %v", err)
+			} else {
+				logger.Infof("Saved state for %s removed", baseName)
+			}
 		}
-	}
 
-	logger.Infof("Container %s removed", name)
+		logger.Infof("Container %s removed", name)
 
-	if shouldRemoveClientConfig() {
-		if err := removeClientConfigurations(name); err != nil {
-			logger.Warnf("Warning: Failed to remove client configurations: %v", err)
-		} else {
-			logger.Infof("Client configurations for %s removed", name)
+		if shouldRemoveClientConfig() {
+			if err := removeClientConfigurations(name); err != nil {
+				logger.Warnf("Warning: Failed to remove client configurations: %v", err)
+			} else {
+				logger.Infof("Client configurations for %s removed", name)
+			}
 		}
 	}
 
@@ -445,6 +449,16 @@ func (d *defaultManager) stopContainer(ctx context.Context, containerID, contain
 	return nil
 }
 
+func (d *defaultManager) DeleteNetwork(ctx context.Context, name string) error {
+	// Remove the network
+	logger.Infof("Removing network %s...", name)
+	if err := d.runtime.DeleteNetwork(ctx, name); err != nil {
+		return fmt.Errorf("failed to remove network: %v", err)
+	}
+
+	return nil
+}
+
 func shouldRemoveClientConfig() bool {
 	c := config.GetConfig()
 	return len(c.Clients.RegisteredClients) > 0 || c.Clients.AutoDiscovery