From e11beb8f6f40921d89ba42d7b02299aeae3a0624 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Mon, 2 Jun 2025 14:16:19 -0400
Subject: [PATCH] include the agent_enrollment indirection in agentPingHandler
 (#967)

---
 controller/agentPing.go | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/controller/agentPing.go b/controller/agentPing.go
index 27b6994b..5291eb4c 100644
--- a/controller/agentPing.go
+++ b/controller/agentPing.go
@@ -17,7 +17,26 @@ func newAgentPingHandler() *agentPingHandler {
 }
 
 func (h *agentPingHandler) Handle(params agent.PingParams, principal *rest_model_zrok.Principal) middleware.Responder {
-	acli, aconn, err := agentController.NewAgentClient(params.Body.EnvZID, cfg.AgentController)
+	trx, err := str.Begin()
+	if err != nil {
+		logrus.Errorf("error starting transaction for '%v': %v", principal.Email, err)
+		return agent.NewPingInternalServerError()
+	}
+	defer trx.Rollback()
+
+	env, err := str.FindEnvironmentForAccount(params.Body.EnvZID, int(principal.ID), trx)
+	if err != nil {
+		logrus.Errorf("error finding environment '%v' for '%v': %v", params.Body.EnvZID, principal.Email, err)
+		return agent.NewPingUnauthorized()
+	}
+
+	ae, err := str.FindAgentEnrollmentForEnvironment(env.Id, trx)
+	if err != nil {
+		logrus.Errorf("error finding agent enrollment for '%v' (%v): %v", params.Body.EnvZID, principal.Email, err)
+		return agent.NewPingBadGateway()
+	}
+
+	acli, aconn, err := agentController.NewAgentClient(ae.Token, cfg.AgentController)
 	if err != nil {
 		logrus.Errorf("error creating agent client for '%v' (%v): %v", params.Body.EnvZID, principal.Email, err)
 		return agent.NewPingInternalServerError()