Feedback

app/invocation/invocation_action_card.tsx

@@ -34,7 +34,7 @@ import { BuildBuddyError, HTTPStatusError } from "../util/errors";
 import { Profile, readProfile } from "../trace/trace_events";
 import TraceViewer from "../trace/trace_viewer";
 import Spinner from "../components/spinner/spinner";
-import { timestampToDate } from "../util/proto";
+import { MessageClass, timestampToDate } from "../util/proto";
 
 type Timestamp = google_timestamp.protobuf.Timestamp;
 type ITimestamp = google_timestamp.protobuf.ITimestamp;
@@ -564,36 +564,21 @@ export default class InvocationActionCardComponent extends React.Component<Props
       .catch((e) => console.error(e));
   }
 
-  private getExecutionAuxiliaryMetadata(): execution_stats.ExecutionAuxiliaryMetadata | null | undefined {
-    const auxiliaryMetadata = this.state.actionResult?.executionMetadata?.auxiliaryMetadata;
-    if (!auxiliaryMetadata || auxiliaryMetadata.length == 0) {
-      return null;
-    }
-    for (const metadata of auxiliaryMetadata) {
-      if (metadata.typeUrl === execution_stats.ExecutionAuxiliaryMetadata.getTypeUrl()) {
-        return execution_stats.ExecutionAuxiliaryMetadata.decode(metadata.value);
-      }
-    }
-    return null;
-  }
-
-  // For firecracker actions, VM metadata is stored in the auxiliary metadata field
-  // of the execution metadata. Try to decode it into an object if it exists.
-  private getFirecrackerVMMetadata(): firecracker.VMMetadata | null | undefined {
-    const auxiliaryMetadata = this.state.actionResult?.executionMetadata?.auxiliaryMetadata;
-    if (!auxiliaryMetadata || auxiliaryMetadata.length == 0) {
-      return null;
-    }
-    for (const metadata of auxiliaryMetadata) {
-      if (metadata.typeUrl === firecracker.VMMetadata.getTypeUrl()) {
-        return firecracker.VMMetadata.decode(metadata.value);
+  /**
+   * Looks for the given message type in auxiliary metadata and returns the
+   * decoded message if found.
+   */
+  private getAuxiliaryMetadata<T>(messageClass: MessageClass<T>): T | null | undefined {
+    for (const metadata of this.state.actionResult?.executionMetadata?.auxiliaryMetadata ?? []) {
+      if (metadata.typeUrl === messageClass.getTypeUrl()) {
+        return messageClass.decode(metadata.value);
       }
     }
     return null;
   }
 
   private getVMPreviousTaskHref(): string {
-    const vmMetadata = this.getFirecrackerVMMetadata();
+    const vmMetadata = this.getAuxiliaryMetadata(firecracker.VMMetadata);
     const task = vmMetadata?.lastExecutedTask;
     if (!task?.executeResponseDigest || !task?.invocationId || !task?.actionDigest) return "";
     return `/invocation/${task.invocationId}?actionDigest=${digestToString(
@@ -890,7 +875,8 @@ export default class InvocationActionCardComponent extends React.Component<Props
 
   private getPlatformOverrides(): Map<string, string> {
     const overrides = new Map<string, string>();
-    for (const prop of this.getExecutionAuxiliaryMetadata()?.platformOverrides?.properties ?? []) {
+    const executionAuxiliaryMetadata = this.getAuxiliaryMetadata(execution_stats.ExecutionAuxiliaryMetadata);
+    for (const prop of executionAuxiliaryMetadata?.platformOverrides?.properties ?? []) {
       let value = prop.value ?? "";
       // TODO: this redaction is also done on the server and can be removed
       // after some time.
@@ -906,7 +892,7 @@ export default class InvocationActionCardComponent extends React.Component<Props
   render() {
     const digest = parseActionDigest(this.props.search.get("actionDigest") ?? "");
     if (!digest) return <></>;
-    const vmMetadata = this.getFirecrackerVMMetadata();
+    const vmMetadata = this.getAuxiliaryMetadata(firecracker.VMMetadata);
     const executionId = this.getExecutionId();
     const platformOverrides = this.getPlatformOverrides();
 

app/util/BUILD

@@ -98,6 +98,7 @@ ts_library(
         "//proto:timestamp_ts_proto",
         "@npm//@types/long",
         "@npm//long",
+        "@npm//protobufjs",
     ],
 )
 

app/util/proto.ts

@@ -1,7 +1,27 @@
 import Long from "long";
+import { Reader } from "protobufjs";
 import { google as google_timestamp } from "../../proto/timestamp_ts_proto";
 import { google as google_duration } from "../../proto/duration_ts_proto";
 
+/**
+ * Generic interface exposed by message classes.
+ *
+ * This exposes the static methods from the generated message classes. Note:
+ * class objects are not the same as _instances_ of the class:
+ *
+ * ```
+ * // Create an instance of `message FooMsg`:
+ * const fooInstance = new FooMsg({ someField: 'someValue' });
+ * // Note the types of `fooInstance` and `FooMsg`:
+ * let value1: FooMsg = fooInstance;
+ * let value2: MessageClass<FooMsg> = FooMsg;
+ * ```
+ */
+export type MessageClass<T> = {
+  decode(source: Reader | Uint8Array, length?: number): T;
+  getTypeUrl(): string;
+};
+
 export function dateToTimestamp(date: Date): google_timestamp.protobuf.Timestamp {
   const timestampMillis = date.getTime();
   return new google_timestamp.protobuf.Timestamp({

enterprise/server/remote_execution/execution_server/execution_server.go

@@ -1213,15 +1213,15 @@ func (s *ExecutionServer) fetchActionAndCommand(ctx context.Context, actionResou
 func redactCachedExecuteResponse(ctx context.Context, rsp *repb.ExecuteResponse) {
 	md := rsp.GetResult().GetExecutionMetadata()
 	for _, auxAny := range md.GetAuxiliaryMetadata() {
-		if auxAny.GetTypeUrl() == "type.googleapis.com/"+string((&espb.ExecutionAuxiliaryMetadata{}).ProtoReflect().Descriptor().FullName()) {
+		if auxAny.MessageIs(&espb.ExecutionAuxiliaryMetadata{}) {
 			redactExecutionAuxiliaryMetadata(ctx, auxAny)
 		}
 	}
 }
 
 func redactExecutionAuxiliaryMetadata(ctx context.Context, auxAny *anypb.Any) {
 	md := &espb.ExecutionAuxiliaryMetadata{}
-	if err := proto.Unmarshal(auxAny.GetValue(), md); err != nil {
+	if err := auxAny.UnmarshalTo(md); err != nil {
 		log.CtxErrorf(ctx, "Failed to unmarshal ExecutionAuxiliaryMetadata: %s", err)
 		return
 	}
@@ -1235,12 +1235,9 @@ func redactExecutionAuxiliaryMetadata(ctx context.Context, auxAny *anypb.Any) {
 		}
 	}
 
-	b, err := proto.Marshal(md)
-	if err != nil {
+	if err := auxAny.MarshalFrom(md); err != nil {
 		log.CtxErrorf(ctx, "Failed to marshal ExecutionAuxiliaryMetadata: %s", err)
-		return
 	}
-	auxAny.Value = b
 }
 
 func executionDuration(md *repb.ExecutedActionMetadata) (time.Duration, error) {