From 6f688ac2c2e37305e9c1dbe9c83fed102ea253d8 Mon Sep 17 00:00:00 2001
From: michel-laterman <michel.laterman@elastic.co>
Date: Mon, 23 Sep 2024 11:53:03 -0700
Subject: [PATCH] smaller fixes

---
 internal/pkg/api/router.go | 15 ++++++++++-----
 internal/pkg/api/server.go | 23 +++++++++++++++++++----
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/internal/pkg/api/router.go b/internal/pkg/api/router.go
index d71ff274d..4eb50a7f0 100644
--- a/internal/pkg/api/router.go
+++ b/internal/pkg/api/router.go
@@ -9,18 +9,21 @@ import (
 	"regexp"
 	"strings"
 
-	"github.com/elastic/fleet-server/v7/internal/pkg/config"
-	"github.com/elastic/fleet-server/v7/internal/pkg/limit"
-	"github.com/elastic/fleet-server/v7/internal/pkg/logger"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
 	"github.com/rs/zerolog"
 	"go.elastic.co/apm/module/apmchiv5/v2"
 	"go.elastic.co/apm/v2"
 
+	"github.com/elastic/fleet-server/v7/internal/pkg/config"
+	"github.com/elastic/fleet-server/v7/internal/pkg/limit"
+	"github.com/elastic/fleet-server/v7/internal/pkg/logger"
+
 	opamp "github.com/open-telemetry/opamp-go/server"
 )
 
+const opampPath = "/v1/opamp"
+
 func newRouter(cfg *config.ServerLimits, si ServerInterface, tracer *apm.Tracer, handlerFn opamp.HTTPHandlerFunc) http.Handler {
 	r := chi.NewRouter()
 	if tracer != nil {
@@ -30,7 +33,7 @@ func newRouter(cfg *config.ServerLimits, si ServerInterface, tracer *apm.Tracer,
 	r.Use(logger.Middleware) // Attach middlewares to router directly so the occur before any request parsing/validation
 	r.Use(middleware.Recoverer)
 	r.Use(Limiter(cfg).middleware)
-	r.HandleFunc("/opamp", http.HandlerFunc(
+	r.HandleFunc(opampPath, http.HandlerFunc(
 		func(w http.ResponseWriter, r *http.Request) {
 			handlerFn(w, r)
 		},
@@ -87,7 +90,7 @@ func pathToOperation(path string) string {
 	if path == "/api/status" {
 		return "status"
 	}
-	if path == "/opamp" {
+	if path == opampPath {
 		return "opamp"
 	}
 	if path == "/api/fleet/uploads" {
@@ -147,6 +150,8 @@ func (l *limiter) middleware(next http.Handler) http.Handler {
 			l.getPGPKey.Wrap("getPGPKey", &cntGetPGP, zerolog.DebugLevel)(next).ServeHTTP(w, r)
 		case "status":
 			l.status.Wrap("status", &cntStatus, zerolog.DebugLevel)(next).ServeHTTP(w, r)
+		case "audit-unenroll":
+			l.auditUnenroll.Wrap("audit-unenroll", &cntAuditUnenroll, zerolog.DebugLevel)(next).ServeHTTP(w, r)
 		default:
 			// no tracking or limits
 			next.ServeHTTP(w, r)
diff --git a/internal/pkg/api/server.go b/internal/pkg/api/server.go
index fe62ab22d..d8c3a311a 100644
--- a/internal/pkg/api/server.go
+++ b/internal/pkg/api/server.go
@@ -13,6 +13,8 @@ import (
 	"net"
 	"net/http"
 
+	"go.elastic.co/apm/v2"
+
 	"github.com/elastic/elastic-agent-libs/transport/tlscommon"
 	"github.com/elastic/fleet-server/v7/internal/pkg/build"
 	"github.com/elastic/fleet-server/v7/internal/pkg/bulk"
@@ -20,8 +22,8 @@ import (
 	"github.com/elastic/fleet-server/v7/internal/pkg/limit"
 	"github.com/elastic/fleet-server/v7/internal/pkg/logger"
 	"github.com/elastic/fleet-server/v7/internal/pkg/policy"
-	"go.elastic.co/apm/v2"
 
+	"github.com/open-telemetry/opamp-go/protobufs"
 	opamp "github.com/open-telemetry/opamp-go/server"
 	"github.com/open-telemetry/opamp-go/server/types"
 	"github.com/rs/zerolog"
@@ -58,10 +60,23 @@ func NewServer(addr string, cfg *config.Server, ct *CheckinT, et *EnrollerT, at
 	handlerFn, contextWithConn, _ := ompampServer.Attach(opamp.Settings{
 		Callbacks: opamp.CallbacksStruct{
 			OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-				fmt.Printf("TEST ")
 				return types.ConnectionResponse{
-					Accept:              true,
-					ConnectionCallbacks: opamp.ConnectionCallbacksStruct{},
+					Accept: true,
+					ConnectionCallbacks: opamp.ConnectionCallbacksStruct{
+						OnConnectedFunc: func(ctx context.Context, _ types.Connection) {
+							zerolog.Ctx(ctx).Info().Msg("Opamp connection started.")
+						},
+						OnMessageFunc: func(ctx context.Context, _ types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+							// TODO: opamp logic goes here
+							zerolog.Ctx(ctx).Info().Msg("Opamp message recieved.")
+							return &protobufs.ServerToAgent{
+								InstanceUid: message.InstanceUid,
+							}
+						},
+						OnConnectionClose: func(ctx context.Context, _ types.Connection) {
+							zerolog.Ctx(ctx).Info().Msg("Opamp connection ended.")
+						},
+					},
 				}
 
 			},