From 331cefb8832a535a0b40e5be0c47aa975110a405 Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Tue, 29 Aug 2023 19:14:12 -0300
Subject: [PATCH] feat(agent): make agent wait until connection end

This prevents the agent from exiting right after connecting to the
server
---
 agent/client/client.go        | 4 ++++
 agent/initialization/start.go | 2 ++
 go.work.sum                   | 3 +++
 3 files changed, 9 insertions(+)

diff --git a/agent/client/client.go b/agent/client/client.go
index 8d72b6fbd5..a68028681c 100644
--- a/agent/client/client.go
+++ b/agent/client/client.go
@@ -64,6 +64,10 @@ func (c *Client) Start(ctx context.Context) error {
 	return nil
 }
 
+func (c *Client) WaitUntilDisconnected() {
+	<-c.done
+}
+
 func (c *Client) SessionConfiguration() *SessionConfig {
 	if c.sessionConfig == nil {
 		return nil
diff --git a/agent/initialization/start.go b/agent/initialization/start.go
index ae91a7de91..65403d5ce9 100644
--- a/agent/initialization/start.go
+++ b/agent/initialization/start.go
@@ -37,4 +37,6 @@ func Start(config config.Config) {
 	if err != nil {
 		log.Fatal(err)
 	}
+
+	client.WaitUntilDisconnected()
 }
diff --git a/go.work.sum b/go.work.sum
index cb692fc292..3e735ec812 100644
--- a/go.work.sum
+++ b/go.work.sum
@@ -177,6 +177,7 @@ github.com/hjson/hjson-go/v4 v4.0.0/go.mod h1:KaYt3bTw3zhBjYqnXkYywcYctk0A2nxeEF
 github.com/jackc/puddle v1.1.3 h1:JnPg/5Q9xVJGfjsO5CPUOjnJps1JaRUm8I9FXVCFK94=
 github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU=
 github.com/knadh/koanf v1.4.3/go.mod h1:5FAkuykKXZvLqhAbP4peWgM5CTcZmn7L1d27k/a+kfg=
+github.com/labstack/gommon v0.3.0 h1:JEeO0bvc78PKdyHxloTKiF8BD5iGrH8T6MSeGvSgob0=
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
 github.com/lyft/protoc-gen-star v0.6.0/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=
 github.com/magiconair/properties v1.8.6/go.mod h1:y3VJvCyxH9uVvJTWEGAELF3aiYNyPKd5NZ3oSwXrF60=
@@ -200,6 +201,8 @@ github.com/spf13/cobra v1.5.0/go.mod h1:dWXEIy2H428czQCjInthrTRUg7yKbok+2Qi/yBIJ
 github.com/stvp/go-udp-testing v0.0.0-20201019212854-469649b16807/go.mod h1:7jxmlfBCDBXRzr0eAQJ48XC1hBu1np4CS5+cHEYfwpc=
 github.com/tklauser/go-sysconf v0.3.10/go.mod h1:C8XykCvCb+Gn0oNCWPIlcb0RuglQTYaQ2hGm7jmxEFk=
 github.com/tklauser/numcpus v0.4.0/go.mod h1:1+UI3pD8NW14VMwdgJNJ1ESk2UnwhAnz5hMwiKKqXCQ=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/fasttemplate v1.2.1 h1:TVEnxayobAdVkhQfrfes2IzOB6o+z4roRkPF52WA1u4=
 github.com/yusufpapurcu/wmi v1.2.2/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
 go.etcd.io/etcd/api/v3 v3.5.4/go.mod h1:5GB2vv4A4AOn3yk7MftYGHkUfGtDHnEraIjym4dYz5A=
 go.etcd.io/etcd/api/v3 v3.5.6/go.mod h1:KFtNaxGDw4Yx/BA4iPPwevUTAuqcsPxzyX8PHydchN8=