Add context to tunnel connect

pkg/agent/run.go

@@ -47,7 +47,7 @@ func run(ctx context.Context, cfg cmds.Agent) error {
 		return err
 	}
 
-	if err := tunnel.Setup(nodeConfig); err != nil {
+	if err := tunnel.Setup(ctx, nodeConfig); err != nil {
 		return err
 	}
 

pkg/agent/tunnel/tunnel.go

@@ -52,7 +52,7 @@ func getAddresses(endpoint *v1.Endpoints) []string {
 	return serverAddresses
 }
 
-func Setup(config *config.Node) error {
+func Setup(ctx context.Context, config *config.Node) error {
 	restConfig, err := clientcmd.BuildConfigFromFlags("", config.AgentConfig.KubeConfigNode)
 	if err != nil {
 		return err
@@ -133,7 +133,18 @@ func Setup(config *config.Node) error {
 		}
 	}()
 
-	wg.Wait()
+	wait := make(chan int, 1)
+	go func() {
+		wg.Wait()
+		wait <- 0
+	}()
+
+	select {
+	case <-ctx.Done():
+		logrus.Error("tunnel context canceled while waiting for connection")
+		return ctx.Err()
+	case <-wait:
+	}
 
 	return nil
 }