Add scrubbing to octtrpc

Allow enabling adding the request/response to the ttrpc span as
attributes, since there are other customers of our code.

Allow scrubbing the `proto.Message` request and response messages in the
ttrpc server interceptor by specifying an arbitrary function to update
a clone of the payloads.

Signed-off-by: Hamza El-Saawy <hamzaelsaawy@microsoft.com>

cmd/containerd-shim-runhcs-v1/exec_hcs.go

@@ -517,7 +517,7 @@ func (he *hcsExec) waitForContainerExit() {
 		trace.StringAttribute("tid", he.tid),
 		trace.StringAttribute("eid", he.id))
 
-	// wait for container or process to exit and ckean up resrources
+	// wait for container or process to exit and clean up resources
 	select {
 	case <-he.c.WaitChannel():
 		// Container exited first. We need to force the process into the exited

cmd/containerd-shim-runhcs-v1/serve.go

@@ -198,10 +198,12 @@ var serveCommand = cli.Command{
 
 		s, err := ttrpc.NewServer(
 			ttrpc.WithUnaryServerInterceptor(octtrpc.ServerInterceptor(
-				octtrpc.WithAttributes( // todo (helsaawy) set these in resource when we switch to OTel
+				octtrpc.WithAttributes( // TODO (helsaawy) set these in resource when we switch to OTel
 					trace.StringAttribute(logfields.ShimID, svc.tid),
 					trace.BoolAttribute(logfields.IsSandbox, svc.isSandbox),
 				),
+				octtrpc.WithAddMessage(),
+				octtrpc.WithAddMessageHook(hcslog.ScrubShimTTRPC),
 			)))
 		if err != nil {
 			return err

internal/gcs/bridge.go

@@ -244,7 +244,7 @@ func (brdg *bridge) RPC(ctx context.Context, proc rpcProc, req requestMessage, r
 		brdg.log.WithField("reason", ctx.Err()).Warn("ignoring response to bridge message")
 		return ctx.Err()
 	case <-t.C:
-		//todo: dont kill bridge on message timeout
+		// TODO: don't kill bridge on message timeout
 		brdg.kill(errors.New("message timeout"))
 		<-call.ch
 		return call.Err()
@@ -390,7 +390,7 @@ func (brdg *bridge) writeMessage(buf *bytes.Buffer, enc *json.Encoder, typ msgTy
 	// Update the message header with the size.
 	binary.LittleEndian.PutUint32(buf.Bytes()[hdrOffSize:], uint32(buf.Len()))
 
-	if brdg.log.Logger.GetLevel() >= logrus.DebugLevel {
+	if brdg.log.Logger.IsLevelEnabled(logrus.DebugLevel) {
 		b := buf.Bytes()[hdrSize:]
 		switch typ {
 		// container environment vars are in rpCreate for linux; rpcExecuteProcess for windows

internal/guest/bridge/bridge.go

@@ -301,7 +301,7 @@ func (b *Bridge) ListenAndServe(bridgeIn io.ReadCloser, bridgeOut io.WriteCloser
 					trace.StringAttribute(logfields.ContainerID, base.ContainerID))
 
 				entry := log.G(ctx)
-				if entry.Logger.GetLevel() >= logrus.DebugLevel {
+				if entry.Logger.IsLevelEnabled(logrus.DebugLevel) {
 					s := string(message)
 					switch header.Type {
 					case prot.ComputeSystemCreateV1:

internal/guest/network/netns.go

@@ -168,7 +168,7 @@ func NetNSConfig(ctx context.Context, ifStr string, nsPid int, adapter *prot.Net
 	}
 
 	// Add some debug logging
-	if entry.Logger.GetLevel() >= logrus.DebugLevel {
+	if entry.Logger.IsLevelEnabled(logrus.DebugLevel) {
 		curNS, _ := netns.Get()
 		// Refresh link attributes/state
 		link, _ = netlink.LinkByIndex(link.Attrs().Index)

internal/log/format.go

@@ -62,7 +62,7 @@ func formatAddr(a net.Addr) string {
 func Format(ctx context.Context, v interface{}) string {
 	b, err := encode(v)
 	if err != nil {
-		// logging errors aren't really warning worthy, and can potentially spam a lot of logs out
+		// logging-related errors aren't really warning worthy, and can potentially spam a lot of logs out
 		G(ctx).WithFields(logrus.Fields{
 			logrus.ErrorKey: err,
 			"type":          fmt.Sprintf("%T", v),

internal/log/scrub.go

@@ -6,6 +6,11 @@ import (
 	"errors"
 	"sync/atomic"
 
+	task "github.com/containerd/containerd/api/runtime/task/v2"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
+
 	hcsschema "github.com/Microsoft/hcsshim/internal/hcs/schema2"
 )
 
@@ -22,6 +27,7 @@ var (
 	// case sensitive keywords, so "env" is not a substring on "Environment"
 	_scrubKeywords = [][]byte{[]byte("env"), []byte("Environment")}
 
+	// TODO (go1.19) atomic.Bool
 	_scrub int32
 )
 
@@ -40,8 +46,27 @@ func IsScrubbingEnabled() bool {
 	return v != 0
 }
 
+func ScrubShimTTRPC(m proto.Message) {
+	if !IsScrubbingEnabled() {
+		return
+	}
+
+	switch t := m.(type) {
+	case *task.ExecProcessRequest:
+		// the spec will be logged elsewhere, so scrub the entire spec wholesale to avoid
+		// needing to unmarshall from then re-marshal back to an anypb.Any
+		//
+		// ignore errors with creating an anypb.Any and use nil response since nil is a
+		// "valid" proto.Message that it signifies an invlaid message
+		// see comment here: https://pkg.go.dev/google.golang.org/protobuf@v1.31.0/proto#Equal
+		a, _ := anypb.New(wrapperspb.String(_scrubbedReplacement))
+		t.Spec = a
+	default:
+	}
+}
+
 // ScrubProcessParameters scrubs HCS Create Process requests with config parameters of
-// type internal/hcs/schema2.ScrubProcessParameters (aka hcsshema.ScrubProcessParameters)
+// type internal/hcs/schema2.ScrubProcessParameters (aka hcsschema.ScrubProcessParameters)
 func ScrubProcessParameters(s string) (string, error) {
 	// todo: deal with v1 ProcessConfig
 	b := []byte(s)

pkg/octtrpc/interceptor.go

@@ -10,6 +10,7 @@ import (
 	"go.opencensus.io/trace/propagation"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/Microsoft/hcsshim/internal/log"
 	"github.com/Microsoft/hcsshim/internal/oc"
@@ -18,6 +19,21 @@ import (
 type options struct {
 	sampler trace.Sampler
 	attrs   []trace.Attribute
+	// add the request/response messages as span attributes
+	addMsg bool
+	// hook to update/modify a copy of the request/response messages before adding them as span attributes
+	msgAttrHook func(proto.Message)
+}
+
+func (o *options) msgHook(v any) any {
+	m, ok := v.(proto.Message)
+	if !ok || o.msgAttrHook == nil {
+		return v
+	}
+
+	m = proto.Clone(m)
+	o.msgAttrHook(m)
+	return m
 }
 
 // Option represents an option function that can be used with the OC TTRPC
@@ -39,6 +55,37 @@ func WithAttributes(attr ...trace.Attribute) Option {
 	}
 }
 
+// these are (currently) ServerInterceptor-specific options, but we cannot create a new [ServerOption] type
+// since that would break our API
+
+// WithAddMessage adds the request and response messages as attributes to the ttrpc method span.
+//
+// [ServerInterceptor] only.
+func WithAddMessage() Option {
+	return func(opts *options) {
+		opts.addMsg = true
+	}
+}
+
+// WithAddMessageHook specifies a hook to modify the ttrpc request or response messages
+// before adding them to the ttrpc method span.
+// This is intended to allow scrubbing sensitive fields from a the message.
+//
+// The function will be called with a clone of the original message (via [proto.Clone]) only if:
+//   - the interceptor is created with the [WithAddMessage] option
+//   - the function is non-nil
+//   - the ttrpc request or response are of type [proto.Message]
+//
+// Since ttrpc is a gRPC replacement, we are guaranteed that the messages will
+// implement [proto.Message].
+//
+// [ServerInterceptor] only.
+func WithAddMessageHook(f func(proto.Message)) Option {
+	return func(opts *options) {
+		opts.msgAttrHook = f
+	}
+}
+
 const metadataTraceContextKey = "octtrpc.tracecontext"
 
 func convertMethodName(name string) string {
@@ -124,18 +171,19 @@ func ServerInterceptor(opts ...Option) ttrpc.UnaryServerInterceptor {
 		}
 		defer span.End()
 		defer func() {
-			if err == nil {
-				span.AddAttributes(trace.StringAttribute("response", log.Format(ctx, resp)))
+			if o.addMsg && err == nil {
+				span.AddAttributes(trace.StringAttribute("response", log.Format(ctx, o.msgHook(resp))))
 			}
 			setSpanStatus(span, err)
 		}()
 		if len(o.attrs) > 0 {
 			span.AddAttributes(o.attrs...)
 		}
 
-		return method(ctx, func(req interface{}) (err error) {
-			if err = unmarshal(req); err == nil {
-				span.AddAttributes(trace.StringAttribute("request", log.Format(ctx, req)))
+		return method(ctx, func(req interface{}) error {
+			err := unmarshal(req)
+			if o.addMsg {
+				span.AddAttributes(trace.StringAttribute("request", log.Format(ctx, o.msgHook(req))))
 			}
 			return err
 		})