Repository: stripe/stripe-cli
Pull request #759 (2021-09-29 22:29:48)

Adding telemetry to gPRC calls
 ### Reviewers
r? @vcheung-stripe 
cc @stripe/developer-products

### Summary
This change adds additional logic in the middleware interceptors to add the `telemetryClient` and `eventMetadata` to the context passed into the RPC functions. This allows us to get telemetry events for each separate RPC call instead of sharing one metadata for the whole of the daemon process. 

An instance of an `eventMetadata` is created on each method call to gRPC.

We use the same instance of `telemetryClient` for the duration the RPC service is up. The instance of the `telemeryClient` is passed to the RPC service when it is created by the daemon command. 

### Testing
An example of data sent to the analytics service through manual testing
```
 p_invocation_id=1beaa848-7919-407d-81e5-96aeb9d6c24c p_user_agent="Stripe/v1 stripe-cli/master" p_os=darwin p_cli_version=master p_livemode=false p_generated_resource=false p_event_value= p_client_id=stripe-cli p_event_id=01f20da5-258b-4fe4-8a96-52813f1b9bc3 p_request_id=req_t1MWVF7KlGFz54 p_command_path=/rpc.StripeCLI/LogsTail p_created=1632951950 p_merchant=acct_1HzAHfIMIxZxvOV4 p_event_name="API Request"

p_invocation_id=eb1fb4ac-dff7-4de7-9e11-561b14e1072b p_user_agent="Stripe/v1 stripe-cli/master" p_os=darwin p_cli_version=master p_livemode=false p_generated_resource=false p_event_value= p_client_id=stripe-cli p_event_id=006011c9-c4cf-4a9c-aa6e-aa0ade98a679 p_request_id=req_Sseb1LmF1NhqVj p_command_path=/rpc.StripeCLI/Listen p_created=1632951941 p_merchant=acct_1HzAHfIMIxZxvOV4 p_event_name="API Request"
```

Regarding unit tests: 
I wasn't sure what the best method of adding unit tests to test the new interceptor logic was. Calling `DialContext` allows us to hit the interceptor logic in the middleware.go class but I'm not sure how to assert that the context updated in the way we expected. In the other tests, we pass in a `ctx` directly into the gRPC calls which would bypass the logic we are trying to test. Let me know if anyone has ideas.


vcheung-stripe: Regarding tests: I don't know of a way to test that the context is being set correctly in the gRPC calls. It might be worth unit testing `updateContextWithTelemetry` by itself.
gracegoo-stripe: > I'm very happy with how decoupled this is from the gRPC logic! Where is the user_agent being set? Or will that be in a future PR?

Oops that's a great catch, I want to set it with this PR! I'll update shortly
gracegoo-stripe: > Regarding tests: I don't know of a way to test that the context is being set correctly in the gRPC calls. It might be worth unit testing `updateContextWithTelemetry` by itself.

Sounds good to me
gracegoo-stripe: Added logic to extract user_agent from the incoming metadata. Example of new params getting sent to r.stripe.com:
```

p_invocation_id=449f13f4-2fd3-4853-bd3e-643b24a06d74 p_user_agent=grpc-node-js/1.3.7 p_os=darwin p_cli_version=master p_livemode=false p_generated_resource=false p_event_value= p_client_id=stripe-cli p_event_id=530821f2-b5f8-4f37-b9fe-4fe80661386d p_request_id=req_hC5y5zlQbs1MhR p_command_path=/rpc.StripeCLI/LogsTail p_created=1633025624 p_merchant=acct_1HzAHfIMIxZxvOV4 p_event_name="API Request"
```

Patch:
diff --git a/cmd/stripe/main.go b/cmd/stripe/main.go
index adffde4..e359ea9 100644
--- a/cmd/stripe/main.go
+++ b/cmd/stripe/main.go
@@ -22,7 +22,7 @@ func main() {
 			Timeout: time.Second * 3,
 		}
 		telemetryClient := &stripe.AnalyticsTelemetryClient{HTTPClient: httpClient}
-		contextWithTelemetry := stripe.WithAnalyticsTelemetryClient(ctx, telemetryClient)
+		contextWithTelemetry := stripe.WithTelemetryClient(ctx, telemetryClient)
 
 		cmd.Execute(contextWithTelemetry)
 
diff --git a/pkg/cmd/daemon.go b/pkg/cmd/daemon.go
index 4ca4718..b803161 100644
--- a/pkg/cmd/daemon.go
+++ b/pkg/cmd/daemon.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/stripe/stripe-cli/pkg/config"
 	"github.com/stripe/stripe-cli/pkg/rpcservice"
+	"github.com/stripe/stripe-cli/pkg/stripe"
 	"github.com/stripe/stripe-cli/pkg/validators"
 )
 
@@ -37,11 +38,12 @@ Currently, stripe daemon only supports a subset of CLI commands. Documentation i
 }
 
 func (dc *daemonCmd) runDaemonCmd(cmd *cobra.Command, args []string) {
+	telemetryClient := stripe.GetTelemetryClient(cmd.Context())
 	srv := rpcservice.New(&rpcservice.Config{
 		Port:    dc.port,
 		Log:     log.StandardLogger(),
 		UserCfg: dc.cfg,
-	})
+	}, telemetryClient)
 
 	ctx := withSIGTERMCancel(cmd.Context(), func() {
 		log.WithFields(log.Fields{
diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go
index 075e3be..482ef90 100644
--- a/pkg/cmd/root.go
+++ b/pkg/cmd/root.go
@@ -18,6 +18,7 @@ import (
 	"github.com/stripe/stripe-cli/pkg/cmd/resource"
 	"github.com/stripe/stripe-cli/pkg/config"
 	"github.com/stripe/stripe-cli/pkg/stripe"
+	"github.com/stripe/stripe-cli/pkg/useragent"
 	"github.com/stripe/stripe-cli/pkg/validators"
 	"github.com/stripe/stripe-cli/pkg/version"
 )
@@ -54,6 +55,8 @@ var rootCmd = &cobra.Command{
 		telemetryMetadata := stripe.GetEventMetadata(cmd.Context())
 		telemetryMetadata.SetCobraCommandContext(cmd)
 		telemetryMetadata.SetMerchant(merchant)
+		telemetryMetadata.SetUserAgent(useragent.GetEncodedUserAgent())
+
 	},
 }
 
diff --git a/pkg/rpcservice/middleware.go b/pkg/rpcservice/middleware.go
index b3e8dd2..a49f011 100644
--- a/pkg/rpcservice/middleware.go
+++ b/pkg/rpcservice/middleware.go
@@ -3,9 +3,12 @@ package rpcservice
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/stripe/stripe-cli/pkg/stripe"
+
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -14,6 +17,23 @@ import (
 
 const requiredHeader = "sec-x-stripe-cli"
 
+// WrappedServerStream wraps a ServerSteam so that we can pass values through context.
+// https://pkg.go.dev/github.com/grpc-ecosystem/go-grpc-middleware#hdr-Writing_Your_Own
+type WrappedServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+// Context returns the context for this stream.
+func (w WrappedServerStream) Context() context.Context {
+	return w.ctx
+}
+
+func newWrappedStream(stream grpc.ServerStream, methodName string, server *RPCService) grpc.ServerStream {
+	newCtx := updateContextWithTelemetry(stream.Context(), methodName, server)
+	return &WrappedServerStream{stream, newCtx}
+}
+
 // Only allow requests from clients that have the required header. This helps prevent malicious
 // websites from making requests. See https://fetch.spec.whatwg.org/#forbidden-header-name
 func authorize(ctx context.Context) error {
@@ -29,6 +49,36 @@ func authorize(ctx context.Context) error {
 	return nil
 }
 
+// Populate the context with:
+// 1. The telemetry client from the RPC Service
+// 2. The event metadata
+func updateContextWithTelemetry(ctx context.Context, methodName string, server *RPCService) context.Context {
+	// If the context is nil for whatever reason, create an empty one
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	// if getting the config errors, don't fail running the command
+	merchant, _ := server.cfg.UserCfg.Profile.GetAccountID()
+	useragent := getUserAgentFromGrpcMetadata(ctx)
+
+	telemetryMetadata := stripe.NewEventMetadata()
+	telemetryMetadata.SetMerchant(merchant)
+	telemetryMetadata.SetCommandPath(methodName)
+	telemetryMetadata.SetUserAgent(useragent)
+
+	newCtx := stripe.WithEventMetadata(stripe.WithTelemetryClient(ctx, server.TelemetryClient), telemetryMetadata)
+	return newCtx
+}
+
+func getUserAgentFromGrpcMetadata(ctx context.Context) string {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return ""
+	}
+	return strings.Join(md["user-agent"], ",")
+}
+
 // Middleware for stream requests
 func serverStreamInterceptor(
 	srv interface{},
@@ -39,10 +89,11 @@ func serverStreamInterceptor(
 	log.WithFields(log.Fields{
 		"prefix": "gRPC",
 	}).Debugf("Streaming method invoked: %v", info.FullMethod)
-	if err := authorize(stream.Context()); err != nil {
+	wrappedStream := newWrappedStream(stream, info.FullMethod, srv.(*RPCService))
+	if err := authorize(wrappedStream.Context()); err != nil {
 		return err
 	}
-	return handler(srv, stream)
+	return handler(srv, wrappedStream)
 }
 
 // Middleware for unary requests
@@ -55,8 +106,9 @@ func serverUnaryInterceptor(
 	log.WithFields(log.Fields{
 		"prefix": "gRPC",
 	}).Debugf("Unary method invoked: %v, req: %v", info.FullMethod, req)
-	if err := authorize(ctx); err != nil {
+	newCtx := updateContextWithTelemetry(ctx, info.FullMethod, info.Server.(*RPCService))
+	if err := authorize(newCtx); err != nil {
 		return nil, err
 	}
-	return handler(ctx, req)
+	return handler(newCtx, req)
 }
diff --git a/pkg/rpcservice/rpc_service.go b/pkg/rpcservice/rpc_service.go
index 454341e..0260766 100644
--- a/pkg/rpcservice/rpc_service.go
+++ b/pkg/rpcservice/rpc_service.go
@@ -14,6 +14,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/stripe/stripe-cli/pkg/config"
+	"github.com/stripe/stripe-cli/pkg/stripe"
 	"github.com/stripe/stripe-cli/rpc"
 )
 
@@ -34,6 +35,9 @@ type RPCService struct {
 	cfg *Config
 
 	grpcServer *grpc.Server
+
+	// TelemetryClient to use for sending telemetry events
+	TelemetryClient stripe.TelemetryClient
 }
 
 // ConfigOutput is the config that clients will need to connect to the gRPC server. This is printed
@@ -47,7 +51,10 @@ type ConfigOutput struct {
 }
 
 // New creates a new RPC service
-func New(cfg *Config) *RPCService {
+func New(
+	cfg *Config,
+	telemetryClient stripe.TelemetryClient,
+) *RPCService {
 	if cfg.Log == nil {
 		cfg.Log = &log.Logger{Out: ioutil.Discard}
 	}
@@ -58,8 +65,9 @@ func New(cfg *Config) *RPCService {
 	)
 
 	return &RPCService{
-		cfg:        cfg,
-		grpcServer: grpcServer,
+		cfg:             cfg,
+		grpcServer:      grpcServer,
+		TelemetryClient: telemetryClient,
 	}
 }
 
diff --git a/pkg/stripe/analytics_telemetry.go b/pkg/stripe/analytics_telemetry.go
index 4e7977a..476f674 100644
--- a/pkg/stripe/analytics_telemetry.go
+++ b/pkg/stripe/analytics_telemetry.go
@@ -16,7 +16,6 @@ import (
 
 	"github.com/spf13/cobra"
 
-	"github.com/stripe/stripe-cli/pkg/useragent"
 	"github.com/stripe/stripe-cli/pkg/version"
 )
 
@@ -57,6 +56,10 @@ type AnalyticsTelemetryClient struct {
 	HTTPClient *http.Client
 }
 
+// NoOpTelemetryClient does not call any endpoint and returns an empty response
+type NoOpTelemetryClient struct {
+}
+
 //
 // Public functions
 //
@@ -65,7 +68,6 @@ type AnalyticsTelemetryClient struct {
 func NewEventMetadata() *CLIAnalyticsEventMetadata {
 	return &CLIAnalyticsEventMetadata{
 		InvocationID: uuid.NewString(),
-		UserAgent:    useragent.GetEncodedUserAgent(),
 		CLIVersion:   version.Version,
 		OS:           runtime.GOOS,
 	}
@@ -85,8 +87,8 @@ func GetEventMetadata(ctx context.Context) *CLIAnalyticsEventMetadata {
 	return nil
 }
 
-// WithAnalyticsTelemetryClient returns a new copy of context.Context with the provided AnalyticsTelemetryClient
-func WithAnalyticsTelemetryClient(ctx context.Context, client *AnalyticsTelemetryClient) context.Context {
+// WithTelemetryClient returns a new copy of context.Context with the provided telemetryClient
+func WithTelemetryClient(ctx context.Context, client TelemetryClient) context.Context {
 	return context.WithValue(ctx, telemetryClientKey{}, client)
 }
 
@@ -118,6 +120,16 @@ func (e *CLIAnalyticsEventMetadata) SetMerchant(merchant string) {
 	e.Merchant = merchant
 }
 
+// SetUserAgent sets the userAgent on the CLIAnalyticsEventContext object
+func (e *CLIAnalyticsEventMetadata) SetUserAgent(userAgent string) {
+	e.UserAgent = userAgent
+}
+
+// SetCommandPath sets the commandPath on the CLIAnalyticsEventContext object
+func (e *CLIAnalyticsEventMetadata) SetCommandPath(commandPath string) {
+	e.CommandPath = commandPath
+}
+
 // SendAPIRequestEvent is a special function for API requests
 func (a *AnalyticsTelemetryClient) SendAPIRequestEvent(ctx context.Context, requestID string, livemode bool) (*http.Response, error) {
 	a.wg.Add(1)
@@ -194,6 +206,16 @@ func (a *AnalyticsTelemetryClient) Wait() {
 	a.wg.Wait()
 }
 
+// SendAPIRequestEvent does nothing
+func (a *NoOpTelemetryClient) SendAPIRequestEvent(ctx context.Context, requestID string, livemode bool) (*http.Response, error) {
+	return nil, nil
+}
+
+// SendEvent does nothing
+func (a *NoOpTelemetryClient) SendEvent(ctx context.Context, eventName string, eventValue string) (*http.Response, error) {
+	return nil, nil
+}
+
 // TelemetryOptedOut returns true if the user has opted out of telemetry,
 // false otherwise.
 func TelemetryOptedOut(optoutVar string) bool {


Test patch:
diff --git a/pkg/rpcservice/middleware_test.go b/pkg/rpcservice/middleware_test.go
index 9539230..9a76326 100644
--- a/pkg/rpcservice/middleware_test.go
+++ b/pkg/rpcservice/middleware_test.go
@@ -12,6 +12,8 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/stripe/stripe-cli/pkg/config"
+	"github.com/stripe/stripe-cli/pkg/stripe"
 	"github.com/stripe/stripe-cli/rpc"
 )
 
@@ -60,3 +62,50 @@ func TestRejectRequestIfMetadataEmpty(t *testing.T) {
 
 	assert.Equal(t, expected.Error(), err.Error())
 }
+
+func TestUpdateContextWithTelemetry(t *testing.T) {
+	telemetryClient := &stripe.NoOpTelemetryClient{}
+	config := &Config{
+		UserCfg: &config.Config{
+			Profile: config.Profile{
+				AccountID: "acct_xxx",
+			},
+		}}
+	rpcService := New(config, telemetryClient)
+
+	// Add grpc Metadata to context
+	md := metadata.Pairs("user-agent", "unit_test")
+
+	ctx := metadata.NewIncomingContext(context.Background(), md)
+
+	newCtx := updateContextWithTelemetry(ctx, "method", rpcService)
+
+	eventMetadata := stripe.GetEventMetadata(newCtx)
+	assert.NotNil(t, eventMetadata)
+	assert.Equal(t, eventMetadata.Merchant, "acct_xxx")
+	assert.Equal(t, eventMetadata.CommandPath, "method")
+	assert.Equal(t, eventMetadata.UserAgent, "unit_test")
+	assert.Equal(t, stripe.GetTelemetryClient(newCtx), telemetryClient)
+}
+
+func TestGetUserAgentFromGRPCMetadata(t *testing.T) {
+	// No grpc metadata
+	assert.Equal(t, getUserAgentFromGrpcMetadata(context.Background()), "")
+}
+
+func TestGetUserAgentFromGRPCMetadataWithNoUserAgent(t *testing.T) {
+	// no user-agent key
+	md := metadata.Pairs("hello", "world")
+
+	ctx := metadata.NewIncomingContext(context.Background(), md)
+	assert.Equal(t, getUserAgentFromGrpcMetadata(ctx), "")
+}
+
+func TestGetUserAgentFromGRPCMetadataWitMultipleUserAgents(t *testing.T) {
+	// no user-agent key
+	md := metadata.Pairs("user-agent", "unit_test")
+	md.Append("user-agent", "hello")
+
+	ctx := metadata.NewIncomingContext(context.Background(), md)
+	assert.Equal(t, getUserAgentFromGrpcMetadata(ctx), "unit_test,hello")
+}
diff --git a/pkg/rpcservice/rpc_service_test.go b/pkg/rpcservice/rpc_service_test.go
index dd59d05..68b6a05 100644
--- a/pkg/rpcservice/rpc_service_test.go
+++ b/pkg/rpcservice/rpc_service_test.go
@@ -25,7 +25,7 @@ func init() {
 				DeviceName: "rpc_test_device_name",
 			},
 		},
-	})
+	}, nil)
 
 	rpc.RegisterStripeCLIServer(srv.grpcServer, srv)
 
diff --git a/pkg/stripe/analytics_telemetry_test.go b/pkg/stripe/analytics_telemetry_test.go
index de4a7d5..fb0f414 100644
--- a/pkg/stripe/analytics_telemetry_test.go
+++ b/pkg/stripe/analytics_telemetry_test.go
@@ -13,6 +13,41 @@ import (
 )
 
 // Context Tests
+func TestEventMetadataWithGet(t *testing.T) {
+	ctx := context.Background()
+	event := &CLIAnalyticsEventMetadata{
+		InvocationID: "hello",
+		UserAgent:    "uesr",
+		CLIVersion:   "1.0",
+		OS:           "os",
+	}
+	newCtx := WithEventMetadata(ctx, event)
+
+	require.Equal(t, GetEventMetadata(newCtx), event)
+}
+
+func TestGetEventMetadata_DoesNotExistInCtx(t *testing.T) {
+	ctx := context.Background()
+	require.Nil(t, GetEventMetadata(ctx))
+}
+
+func TestTelemetryClientWithGet(t *testing.T) {
+	ctx := context.Background()
+	url, _ := url.Parse("http://hello.com")
+	telemetryClient := &AnalyticsTelemetryClient{
+		BaseURL:    url,
+		HTTPClient: &http.Client{},
+	}
+	newCtx := WithTelemetryClient(ctx, telemetryClient)
+
+	require.Equal(t, GetTelemetryClient(newCtx), telemetryClient)
+}
+
+func TestGetTelemetryClient_DoesNotExistInCtx(t *testing.T) {
+	ctx := context.Background()
+	require.Nil(t, GetTelemetryClient(ctx))
+}
+
 func TestSetCobraCommandContext(t *testing.T) {
 	tel := NewEventMetadata()
 	cmd := &cobra.Command{