POC opamp

go.mod

@@ -21,6 +21,8 @@ require (
 	github.com/mailru/easyjson v0.7.7
 	github.com/miolini/datacounter v1.0.3
 	github.com/oapi-codegen/runtime v1.1.1
+	github.com/oklog/ulid/v2 v2.1.0
+	github.com/open-telemetry/opamp-go v0.15.0
 	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58
 	github.com/prometheus/client_golang v1.19.0
 	github.com/rs/xid v1.5.0
@@ -38,7 +40,7 @@ require (
 	golang.org/x/sync v0.8.0
 	golang.org/x/time v0.5.0
 	google.golang.org/grpc v1.63.2
-	google.golang.org/protobuf v1.33.0
+	google.golang.org/protobuf v1.34.1
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -64,6 +66,7 @@ require (
 	github.com/golang/glog v1.2.0 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/uuid v1.6.0 // indirect
+	github.com/gorilla/websocket v1.5.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect

go.sum

@@ -90,6 +90,8 @@ github.com/google/pprof v0.0.0-20230426061923-93006964c1fc h1:AGDHt781oIcL4EFk7c
 github.com/google/pprof v0.0.0-20230426061923-93006964c1fc/go.mod h1:79YE0hCXdHag9sBkw2o+N/YnZtTkXi0UT9Nnixa5eYk=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
+github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
 github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
 github.com/hashicorp/go-version v1.6.0 h1:feTTfFNnjP967rlCxM/I9g701jU+RN74YKx2mOkIeek=
@@ -125,12 +127,17 @@ github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6U
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/oapi-codegen/runtime v1.1.1 h1:EXLHh0DXIJnWhdRPN2w4MXAzFyE4CskzhNLUmtpMYro=
 github.com/oapi-codegen/runtime v1.1.1/go.mod h1:SK9X900oXmPWilYR5/WKPzt3Kqxn/uS/+lbpREv+eCg=
+github.com/oklog/ulid/v2 v2.1.0 h1:+9lhoxAP56we25tyYETBBY1YLA2SaoLvUFgrP2miPJU=
+github.com/oklog/ulid/v2 v2.1.0/go.mod h1:rcEKHmBBKfef9DhnvX7y1HZBYxjXb0cP5ExxNsTT1QQ=
+github.com/open-telemetry/opamp-go v0.15.0 h1:X2TWhEsGQ8GP7Uos3Ic9v/1aFUqoECZXKS7xAF5HqsA=
+github.com/open-telemetry/opamp-go v0.15.0/go.mod h1:QyPeN56JXlcZt5yG5RMdZ50Ju+zMFs1Ihy/hwHyF8Oo=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.0.2 h1:9yCKha/T5XdGtO0q9Q9a6T5NUCsTn/DrBg0D7ufOcFM=
 github.com/opencontainers/image-spec v1.0.2/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y=
+github.com/pborman/getopt v0.0.0-20170112200414-7148bc3a4c30/go.mod h1:85jBQOZwpVEaDAr341tbn15RS4fCAsIst0qp7i8ex1o=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -287,8 +294,8 @@ google.golang.org/genproto/googleapis/rpc v0.0.0-20240415180920-8c6c420018be h1:
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240415180920-8c6c420018be/go.mod h1:WtryC6hu0hhx87FDGxWCDptyssuo68sk10vYjF+T9fY=
 google.golang.org/grpc v1.63.2 h1:MUeiw1B2maTVZthpU5xvASfTh3LDbxHd6IJ6QQVU+xM=
 google.golang.org/grpc v1.63.2/go.mod h1:WAX/8DgncnokcFUldAxq7GeB5DXHDbMF+lLvDomNkRA=
-google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
-google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg=
+google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=

internal/pkg/api/router.go

@@ -9,24 +9,35 @@ 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"
 )
 
-func newRouter(cfg *config.ServerLimits, si ServerInterface, tracer *apm.Tracer) http.Handler {
+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 {
 		r.Use(apmchiv5.Middleware(apmchiv5.WithTracer(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(opampPath, http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			handlerFn(w, r)
+		},
+	))
 	return HandlerWithOptions(si, ChiServerOptions{
 		BaseRouter:       r,
 		ErrorHandlerFunc: ErrorResp,
@@ -79,6 +90,9 @@ func pathToOperation(path string) string {
 	if path == "/api/status" {
 		return "status"
 	}
+	if path == opampPath {
+		return "opamp"
+	}
 	if path == "/api/fleet/uploads" {
 		return "uploadBegin"
 	}

internal/pkg/api/server.go

@@ -13,29 +13,37 @@ 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"
+	"github.com/elastic/fleet-server/v7/internal/pkg/cache"
 	"github.com/elastic/fleet-server/v7/internal/pkg/config"
+	"github.com/elastic/fleet-server/v7/internal/pkg/dl"
 	"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/opamp"
 	"github.com/elastic/fleet-server/v7/internal/pkg/policy"
-	"go.elastic.co/apm/v2"
 
+	"github.com/open-telemetry/opamp-go/protobufs"
+	opampserver "github.com/open-telemetry/opamp-go/server"
+	"github.com/open-telemetry/opamp-go/server/types"
 	"github.com/rs/zerolog"
 )
 
 type server struct {
-	cfg     *config.Server
-	addr    string
-	handler http.Handler
+	cfg             *config.Server
+	addr            string
+	handler         http.Handler
+	contextWithConn opampserver.ConnContext
 }
 
 // NewServer creates a new HTTP api for the passed addr.
 //
 // The server has a listener specific conn limit and endpoint specific rate-limits.
 // The underlying API structs (such as *CheckinT) may be shared between servers.
-func NewServer(addr string, cfg *config.Server, ct *CheckinT, et *EnrollerT, at *ArtifactT, ack *AckT, st *StatusT, sm policy.SelfMonitor, bi build.Info, ut *UploadT, ft *FileDeliveryT, pt *PGPRetrieverT, audit *AuditT, bulker bulk.Bulk, tracer *apm.Tracer) *server {
+func NewServer(addr string, cfg *config.Server, ct *CheckinT, et *EnrollerT, at *ArtifactT, ack *AckT, st *StatusT, sm policy.SelfMonitor, bi build.Info, ut *UploadT, ft *FileDeliveryT, pt *PGPRetrieverT, audit *AuditT, bulker bulk.Bulk, cache cache.Cache, pm policy.Monitor, tracer *apm.Tracer) *server { // this is messy, we have an open issue to refactor
 	a := &apiServer{
 		ct:     ct,
 		et:     et,
@@ -50,10 +58,87 @@ func NewServer(addr string, cfg *config.Server, ct *CheckinT, et *EnrollerT, at
 		audit:  audit,
 		bulker: bulker,
 	}
+
+	ompampServer := opampserver.New(nil)
+	op := opamp.NewHandler(bulker, cache, pm)
+	handlerFn, contextWithConn, _ := ompampServer.Attach(opampserver.Settings{
+		Callbacks: opampserver.CallbacksStruct{
+			OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+				var policyID string
+				var namespaces []string
+				agent, err := authAgent(request, nil, bulker, cache)
+				if errors.Is(err, ErrAgentNotFound) { // No agent associated, get the enrollment token's associated policyID for a register on first use flow
+					enrollKey, err := authAPIKey(request, bulker, cache)
+					if err != nil {
+						zerolog.Ctx(request.Context()).Warn().Err(err).Msg("Opamp registration api key auth failed.")
+						return types.ConnectionResponse{
+							Accept:         false,
+							HTTPStatusCode: http.StatusUnauthorized,
+						}
+					}
+					// TODO handle static enrollment tokens
+					key, ok := cache.GetEnrollmentAPIKey(enrollKey.ID)
+					if !ok {
+						rec, err := dl.FindEnrollmentAPIKey(request.Context(), bulker, dl.QueryEnrollmentAPIKeyByID, dl.FieldAPIKeyID, enrollKey.ID)
+						if err != nil {
+							return types.ConnectionResponse{
+								Accept:         false,
+								HTTPStatusCode: http.StatusInternalServerError,
+							}
+						}
+						if !rec.Active {
+							return types.ConnectionResponse{
+								Accept:         false,
+								HTTPStatusCode: http.StatusUnauthorized,
+							}
+						}
+						cache.SetEnrollmentAPIKey(enrollKey.ID, rec, int64(len(rec.APIKey)))
+						key = rec
+					}
+					policyID = key.PolicyID
+					namespaces = key.Namespaces
+				} else if err != nil {
+					zerolog.Ctx(request.Context()).Warn().Err(err).Msg("Opamp request api key auth failed.")
+					return types.ConnectionResponse{
+						Accept:         false,
+						HTTPStatusCode: http.StatusUnauthorized,
+					}
+				}
+				return types.ConnectionResponse{
+					Accept: true,
+					ConnectionCallbacks: opampserver.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 {
+							zerolog.Ctx(ctx).Info().Msg("Opamp message received.")
+							response, err := op.Process(ctx, agent, policyID, namespaces, message)
+							if err != nil {
+								zerolog.Ctx(ctx).Error().Err(err).Msg("Error processing opamp request.")
+								return &protobufs.ServerToAgent{
+									InstanceUid: message.InstanceUid,
+									ErrorResponse: &protobufs.ServerErrorResponse{
+										ErrorMessage: err.Error(),
+									},
+								}
+							}
+							return response
+						},
+						OnConnectionCloseFunc: func(_ types.Connection) {
+							zerolog.Ctx(request.Context()).Info().Msg("Opamp connection ended.") // FIXME getting context from request might be messy
+						},
+					},
+				}
+
+			},
+		},
+	})
+
 	return &server{
-		addr:    addr,
-		cfg:     cfg,
-		handler: newRouter(&cfg.Limits, a, tracer),
+		addr:            addr,
+		cfg:             cfg,
+		handler:         newRouter(&cfg.Limits, a, tracer, handlerFn),
+		contextWithConn: contextWithConn,
 	}
 }
 
@@ -75,6 +160,7 @@ func (s *server) Run(ctx context.Context) error {
 		BaseContext:       func(net.Listener) context.Context { return ctx },
 		ErrorLog:          errLogger(ctx),
 		ConnState:         diagConn,
+		ConnContext:       s.contextWithConn,
 	}
 
 	var listenCfg net.ListenConfig