Add OpenTelemetry traces (#110)

Author: RebeccaMahany

client.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/osquery/osquery-go/gen/osquery"
+	"github.com/osquery/osquery-go/traces"
 	"github.com/osquery/osquery-go/transport"
 
 	"github.com/apache/thrift/lib/go/thrift"
@@ -107,6 +108,9 @@ func (c *ExtensionManagerClient) Call(registry, item string, request osquery.Ext
 
 // CallContext requests a call to an extension (or core) registry plugin.
 func (c *ExtensionManagerClient) CallContext(ctx context.Context, registry, item string, request osquery.ExtensionPluginRequest) (*osquery.ExtensionResponse, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.CallContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -121,6 +125,9 @@ func (c *ExtensionManagerClient) Extensions() (osquery.InternalExtensionList, er
 
 // ExtensionsContext requests the list of active registered extensions.
 func (c *ExtensionManagerClient) ExtensionsContext(ctx context.Context) (osquery.InternalExtensionList, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.ExtensionsContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -135,6 +142,9 @@ func (c *ExtensionManagerClient) RegisterExtension(info *osquery.InternalExtensi
 
 // RegisterExtensionContext registers the extension plugins with the osquery process.
 func (c *ExtensionManagerClient) RegisterExtensionContext(ctx context.Context, info *osquery.InternalExtensionInfo, registry osquery.ExtensionRegistry) (*osquery.ExtensionStatus, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.RegisterExtensionContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -149,6 +159,9 @@ func (c *ExtensionManagerClient) DeregisterExtension(uuid osquery.ExtensionRoute
 
 // DeregisterExtensionContext de-registers the extension plugins with the osquery process.
 func (c *ExtensionManagerClient) DeregisterExtensionContext(ctx context.Context, uuid osquery.ExtensionRouteUUID) (*osquery.ExtensionStatus, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.DeregisterExtensionContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -163,6 +176,9 @@ func (c *ExtensionManagerClient) Options() (osquery.InternalOptionList, error) {
 
 // OptionsContext requests the list of bootstrap or configuration options.
 func (c *ExtensionManagerClient) OptionsContext(ctx context.Context) (osquery.InternalOptionList, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.OptionsContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -181,6 +197,9 @@ func (c *ExtensionManagerClient) Query(sql string) (*osquery.ExtensionResponse,
 // Consider using the QueryRow or QueryRows helpers for a more friendly
 // interface.
 func (c *ExtensionManagerClient) QueryContext(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.QueryContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}
@@ -199,6 +218,9 @@ func (c *ExtensionManagerClient) QueryRows(sql string) ([]map[string]string, err
 // results. It handles checking both the transport level errors and the osquery
 // internal errors by returning a normal Go error type.
 func (c *ExtensionManagerClient) QueryRowsContext(ctx context.Context, sql string) ([]map[string]string, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.QueryRowsContext")
+	defer span.End()
+
 	res, err := c.QueryContext(ctx, sql)
 	if err != nil {
 		return nil, errors.Wrap(err, "transport error in query")
@@ -222,6 +244,9 @@ func (c *ExtensionManagerClient) QueryRow(sql string) (map[string]string, error)
 // QueryRowContext behaves similarly to QueryRows, but it returns an error if the
 // query does not return exactly one row.
 func (c *ExtensionManagerClient) QueryRowContext(ctx context.Context, sql string) (map[string]string, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.QueryRowContext")
+	defer span.End()
+
 	res, err := c.QueryRowsContext(ctx, sql)
 	if err != nil {
 		return nil, err
@@ -239,6 +264,9 @@ func (c *ExtensionManagerClient) GetQueryColumns(sql string) (*osquery.Extension
 
 // GetQueryColumnsContext requests the columns returned by the parsed query.
 func (c *ExtensionManagerClient) GetQueryColumnsContext(ctx context.Context, sql string) (*osquery.ExtensionResponse, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerClient.GetQueryColumnsContext")
+	defer span.End()
+
 	if err := c.lock.Lock(ctx); err != nil {
 		return nil, err
 	}

go.mod

@@ -5,11 +5,16 @@ require (
 	github.com/apache/thrift v0.16.0
 	github.com/pkg/errors v0.8.0
 	github.com/stretchr/testify v1.8.3
+	go.opentelemetry.io/otel v1.16.0
+	go.opentelemetry.io/otel/trace v1.16.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/go-logr/logr v1.2.4 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	go.opentelemetry.io/otel/metric v1.16.0 // indirect
 	golang.org/x/sys v0.1.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -4,13 +4,25 @@ github.com/apache/thrift v0.16.0 h1:qEy6UW60iVOlUy+b9ZR0d5WzUWYGOo4HfopoyBaNmoY=
 github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ=
+github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
 github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+go.opentelemetry.io/otel v1.16.0 h1:Z7GVAX/UkAXPKsy94IU+i6thsQS4nb7LviLpnaNeW8s=
+go.opentelemetry.io/otel v1.16.0/go.mod h1:vl0h9NUa1D5s1nv3A5vZOYWn8av4K8Ml6JDeHrT/bx4=
+go.opentelemetry.io/otel/metric v1.16.0 h1:RbrpwVG1Hfv85LgnZ7+txXioPDoh6EdbZHo26Q3hqOo=
+go.opentelemetry.io/otel/metric v1.16.0/go.mod h1:QE47cpOmkwipPiefDwo2wDzwJrlfxxNYodqc4xnGCo4=
+go.opentelemetry.io/otel/trace v1.16.0 h1:8JRpaObFoW0pxuVPapkgH8UhHQj+bJW8jJsCZEu5MQs=
+go.opentelemetry.io/otel/trace v1.16.0/go.mod h1:Yt9vYq1SdNz3xdjZZK7wcXv1qv2pwLkqr2QVwea0ef0=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=

plugin/config/config.go

@@ -7,6 +7,7 @@ import (
 	"context"
 
 	"github.com/osquery/osquery-go/gen/osquery"
+	"github.com/osquery/osquery-go/traces"
 )
 
 // GenerateConfigsFunc returns the configurations generated by this plugin.
@@ -55,6 +56,9 @@ const requestActionKey = "action"
 const genConfigAction = "genConfig"
 
 func (t *Plugin) Call(ctx context.Context, request osquery.ExtensionPluginRequest) osquery.ExtensionResponse {
+	ctx, span := traces.StartSpan(ctx, "Config.Call", "action", request[requestActionKey])
+	defer span.End()
+
 	switch request[requestActionKey] {
 	case genConfigAction:
 		configs, err := t.generate(ctx)

plugin/distributed/distributed.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"github.com/osquery/osquery-go/gen/osquery"
+	"github.com/osquery/osquery-go/traces"
 )
 
 // GetQueriesResult contains the information about which queries the
@@ -234,6 +235,9 @@ func convertRows(rows []interface{}) ([]map[string]string, error) {
 }
 
 func (t *Plugin) Call(ctx context.Context, request osquery.ExtensionPluginRequest) osquery.ExtensionResponse {
+	ctx, span := traces.StartSpan(ctx, "Distributed.Call", "action", request[requestActionKey])
+	defer span.End()
+
 	switch request[requestActionKey] {
 	case getQueriesAction:
 		queries, err := t.getQueries(ctx)

plugin/table/table.go

@@ -7,6 +7,7 @@ import (
 	"strconv"
 
 	"github.com/osquery/osquery-go/gen/osquery"
+	"github.com/osquery/osquery-go/traces"
 	"github.com/pkg/errors"
 )
 
@@ -52,6 +53,9 @@ func (t *Plugin) Routes() osquery.ExtensionPluginResponse {
 }
 
 func (t *Plugin) Call(ctx context.Context, request osquery.ExtensionPluginRequest) osquery.ExtensionResponse {
+	ctx, span := traces.StartSpan(ctx, "Table.Call", "action", request["action"])
+	defer span.End()
+
 	ok := osquery.ExtensionStatus{Code: 0, Message: "OK"}
 	switch request["action"] {
 	case "generate":

server.go

@@ -9,6 +9,7 @@ import (
 	"github.com/apache/thrift/lib/go/thrift"
 
 	"github.com/osquery/osquery-go/gen/osquery"
+	"github.com/osquery/osquery-go/traces"
 	"github.com/osquery/osquery-go/transport"
 	"github.com/pkg/errors"
 )
@@ -287,6 +288,12 @@ func (s *ExtensionManagerServer) Ping(ctx context.Context) (*osquery.ExtensionSt
 // Call routes a call from the osquery process to the appropriate registered
 // plugin.
 func (s *ExtensionManagerServer) Call(ctx context.Context, registry string, item string, request osquery.ExtensionPluginRequest) (*osquery.ExtensionResponse, error) {
+	ctx, span := traces.StartSpan(ctx, "ExtensionManagerServer.Call",
+		"registry", registry,
+		"item", item,
+	)
+	defer span.End()
+
 	subreg, ok := s.registry[registry]
 	if !ok {
 		return &osquery.ExtensionResponse{
@@ -307,7 +314,7 @@ func (s *ExtensionManagerServer) Call(ctx context.Context, registry string, item
 		}, nil
 	}
 
-	response := plugin.Call(context.Background(), request)
+	response := plugin.Call(ctx, request)
 	return &response, nil
 }
 

traces/traces.go

@@ -0,0 +1,79 @@
+// Package traces allows for instrumenting osquery-go with OpenTelemetry traces.
+// Unless the consuming application specifically configures a trace exporter, all tracing is a no-op.
+package traces
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
+)
+
+const instrumentationPkg = "github.com/osquery/osquery-go"
+
+var internalVersion string
+
+// osqueryGoVersion returns the version of osquery-go, to be used to set the instrumentation
+// version `otel.scope.version`. It looks through build info to determine the current version
+// of the osquery-go package. Once determined, it saves the version to avoid performing this
+// work again. If the version cannot be determined, the version defaults to the current version,
+// which is hardcoded.
+func osqueryGoVersion() string {
+	if internalVersion != "" {
+		return internalVersion
+	}
+	if info, ok := debug.ReadBuildInfo(); ok {
+		for _, dep := range info.Deps {
+			if dep == nil {
+				continue
+			}
+			if dep.Path == instrumentationPkg {
+				internalVersion = dep.Version
+				return internalVersion
+			}
+		}
+	}
+
+	// Couldn't get the version from runtime build info -- save and return 0.0.0,
+	// which is the current osquery-go version.
+	internalVersion = "0.0.0"
+	return internalVersion
+}
+
+// By default, use the global tracer provider, which is a no-op provider.
+var tracerProvider = otel.GetTracerProvider()
+
+// SetTracerProvider allows consuming libraries to set a custom/non-global tracer provider.
+func SetTracerProvider(tp trace.TracerProvider) {
+	tracerProvider = tp
+}
+
+// OsqueryGoTracer provides a tracer with a standardized name and version.
+// It should be used to start a span that requires `SpanStartOption`s that are
+// not supported by `StartSpan` below -- i.e., any `SpanStartOption` besides
+// `WithAttributes`.
+func OsqueryGoTracer() trace.Tracer {
+	return tracerProvider.Tracer(instrumentationPkg, trace.WithInstrumentationVersion(osqueryGoVersion()))
+}
+
+// StartSpan is a wrapper around trace.Tracer.Start that simplifies passing in span attributes.
+// `keyVals` should be a list of pairs, where the first in the pair is a string representing
+// the attribute key and the second in the pair is the attribute value.
+// The caller is always responsible for ending the returned span.
+// Any spans requiring more specific configuration can be created manually via OsqueryGoTracer().Start.
+func StartSpan(ctx context.Context, spanName string, keyVals ...string) (context.Context, trace.Span) {
+	attrs := make([]attribute.KeyValue, 0)
+
+	for i := 0; i < len(keyVals); i += 2 {
+		// Ensure all attributes are appropriately namespaced
+		key := fmt.Sprintf("osquery-go.%s", keyVals[i])
+		attrs = append(attrs, attribute.String(key, keyVals[i+1]))
+	}
+
+	opts := []trace.SpanStartOption{trace.WithAttributes(attrs...)}
+
+	return OsqueryGoTracer().Start(ctx, spanName, opts...)
+}