Add filterSensitiveLog for Union types

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/StructuredMemberWriter.java

@@ -25,9 +25,11 @@
 import software.amazon.smithy.model.shapes.MapShape;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.SimpleShape;
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.traits.IdempotencyTokenTrait;
 import software.amazon.smithy.model.traits.SensitiveTrait;
+import software.amazon.smithy.model.traits.StreamingTrait;
 
 /**
  * Generates objects, interfaces, enums, etc.
@@ -74,27 +76,37 @@ void writeFilterSensitiveLog(TypeScriptWriter writer, String objectParam) {
         writer.write("...$L,", objectParam);
         for (MemberShape member : members) {
             if (isMemberOverwriteRequired(member, new HashSet<String>())) {
-                Shape memberTarget = model.expectShape(member.getTarget());
                 String memberName = getSanitizedMemberName(member);
-                String memberParam = String.format("%s.%s", objectParam, memberName);
                 writer.openBlock("...($1L.$2L && { $2L: ", "}),", objectParam, memberName, () -> {
-                    if (member.getMemberTrait(model, SensitiveTrait.class).isPresent()) {
-                        // member is Sensitive, hide the value.
-                        writer.write("SENSITIVE_STRING");
-                    } else if (memberTarget instanceof StructureShape) {
-                        writeStructureFilterSensitiveLog(writer, memberTarget, memberParam);
-                    } else if (memberTarget instanceof CollectionShape) {
-                        MemberShape collectionMember = ((CollectionShape) memberTarget).getMember();
-                        writeCollectionFilterSensitiveLog(writer, collectionMember, memberParam);
-                    } else if (memberTarget instanceof MapShape) {
-                        MemberShape mapMember = ((MapShape) memberTarget).getValue();
-                        writeMapFilterSensitiveLog(writer, mapMember, memberParam);
-                    }
+                    String memberParam = String.format("%s.%s", objectParam, memberName);
+                    writeMemberFilterSensitiveLog(writer, member, memberParam);
                 });
             }
         }
     }
 
+    void writeMemberFilterSensitiveLog(TypeScriptWriter writer, MemberShape member, String memberParam) {
+        Shape memberTarget = model.expectShape(member.getTarget());
+        if (member.getMemberTrait(model, SensitiveTrait.class).isPresent()) {
+            // member is Sensitive, hide the value.
+            writer.write("SENSITIVE_STRING");
+        } else if (memberTarget instanceof SimpleShape) {
+            writer.write(memberParam);
+        } else if (memberTarget.isStructureShape() || memberTarget.isUnionShape()) {
+            writeStructureFilterSensitiveLog(writer, memberTarget, memberParam);
+        } else if (memberTarget instanceof CollectionShape) {
+            MemberShape collectionMember = ((CollectionShape) memberTarget).getMember();
+            writeCollectionFilterSensitiveLog(writer, collectionMember, memberParam);
+        } else if (memberTarget instanceof MapShape) {
+            MemberShape mapMember = ((MapShape) memberTarget).getValue();
+            writeMapFilterSensitiveLog(writer, mapMember, memberParam);
+        } else {
+            throw new CodegenException(String.format(
+                "MemberFilterSensitiveLog attempted for %s", memberTarget.getType()
+            ));
+        }
+    }
+
     /**
      * Recursively writes filterSensitiveLog for StructureShape.
      */
@@ -106,10 +118,13 @@ private void writeStructureFilterSensitiveLog(
         if (structureTarget.hasTrait(SensitiveTrait.class)) {
             // member is Sensitive, hide the value.
             writer.write("SENSITIVE_STRING");
-            return;
+        } else if (structureTarget.hasTrait(StreamingTrait.class) && structureTarget.isUnionShape()) {
+            // disable logging for StreamingTrait
+            writer.write("'STREAMING_CONTENT'");
+        } else {
+            // Call filterSensitiveLog on Structure.
+            writer.write("$T.filterSensitiveLog($L)", symbolProvider.toSymbol(structureTarget), structureParam);
         }
-        // Call filterSensitiveLog on Structure.
-        writer.write("$T.filterSensitiveLog($L)", symbolProvider.toSymbol(structureTarget), structureParam);
     }
 
     /**
@@ -123,32 +138,15 @@ private void writeCollectionFilterSensitiveLog(
         if (collectionMember.getMemberTrait(model, SensitiveTrait.class).isPresent()) {
             // member is Sensitive, hide the value.
             writer.write("SENSITIVE_STRING");
-            return;
+        } else if (model.expectShape(collectionMember.getTarget()) instanceof SimpleShape) {
+            writer.write(collectionParam);
+        } else {
+            writer.openBlock("$L.map(", ")", collectionParam, () -> {
+                String itemParam = "item";
+                writer.write("$L => ", itemParam);
+                writeMemberFilterSensitiveLog(writer, collectionMember, itemParam);
+            });
         }
-
-        writer.openBlock("$L.map(", ")", collectionParam, () -> {
-            String itemParam = "item";
-            Shape collectionMemberTarget = model.expectShape(collectionMember.getTarget());
-            writer.write("$L => ", itemParam);
-            if (collectionMemberTarget instanceof StructureShape) {
-                writeStructureFilterSensitiveLog(writer, collectionMemberTarget, itemParam);
-            } else if (collectionMemberTarget instanceof CollectionShape) {
-                MemberShape nestedCollectionMember = ((CollectionShape) collectionMemberTarget).getMember();
-                writeCollectionFilterSensitiveLog(writer, nestedCollectionMember, itemParam);
-            } else if (collectionMemberTarget instanceof MapShape) {
-                MemberShape mapMember = ((MapShape) collectionMemberTarget).getValue();
-                writeMapFilterSensitiveLog(writer, mapMember, itemParam);
-            } else {
-                // This path should not reach because of recursive isMemberOverwriteRequired.
-                throw new CodegenException(String.format(
-                    "CollectionFilterSensitiveLog attempted for %s while it was not required",
-                    collectionMemberTarget.getType()
-                ));
-                // For quick-fix in case of high severity issue:
-                // comment out the exception above and uncomment the line below.
-                // writer.write("$1L => $1L", itemParam);
-            }
-        });
     }
 
     /**
@@ -158,41 +156,23 @@ private void writeMapFilterSensitiveLog(TypeScriptWriter writer, MemberShape map
         if (mapMember.getMemberTrait(model, SensitiveTrait.class).isPresent()) {
             // member is Sensitive, hide the value.
             writer.write("SENSITIVE_STRING");
-            return;
+        } else if (model.expectShape(mapMember.getTarget()) instanceof SimpleShape) {
+            writer.write(mapParam);
+        } else {
+            String accParam = "acc"; // accumulator for the reducer
+            String keyParam = "key"; // key of the Object.entries() key-value pair
+            String valueParam = "value"; // value of the Object.entries() key-value pair
+
+            // Reducer is common to all shapes.
+            writer.openBlock("Object.entries($L).reduce(($L: any, [$L, $L]: [string, $T]) => ({", "}), {})",
+                mapParam, accParam, keyParam, valueParam, symbolProvider.toSymbol(mapMember), () -> {
+                    writer.write("...$L,", accParam);
+                    writer.openBlock("[$L]: ", ",", keyParam, () -> {
+                        writeMemberFilterSensitiveLog(writer, mapMember, valueParam);
+                    });
+                }
+            );
         }
-
-        String accParam = "acc"; // accumulator for the reducer
-        String keyParam = "key"; // key of the Object.entries() key-value pair
-        String valueParam = "value"; // value of the Object.entries() key-value pair
-
-        // Reducer is common to all shapes.
-        writer.openBlock("Object.entries($L).reduce(($L: any, [$L, $L]: [string, $T]) => ({", "}), {})",
-            mapParam, accParam, keyParam, valueParam, symbolProvider.toSymbol(mapMember), () -> {
-                writer.write("...$L,", accParam);
-                Shape mapMemberTarget = model.expectShape(mapMember.getTarget());
-                writer.openBlock("[$L]: ", ",", keyParam, () -> {
-                    if (mapMemberTarget instanceof StructureShape) {
-                        writeStructureFilterSensitiveLog(writer, mapMemberTarget, valueParam);
-                    } else if (mapMemberTarget instanceof CollectionShape) {
-                        MemberShape collectionMember = ((CollectionShape) mapMemberTarget).getMember();
-                        writeCollectionFilterSensitiveLog(writer, collectionMember, valueParam);
-                    } else if (mapMemberTarget instanceof MapShape) {
-                        MemberShape nestedMapMember = ((MapShape) mapMemberTarget).getValue();
-                        writeMapFilterSensitiveLog(writer, nestedMapMember, valueParam);
-                    } else {
-                        // This path should not reach because of recursive isMemberOverwriteRequired.
-                        throw new CodegenException(String.format(
-                            "MapFilterSensitiveLog attempted for %s while it was not required",
-                            mapMemberTarget.getType()
-                        ));
-                        // For quick-fix in case of high severity issue:
-                        // comment out the exception above and uncomment the line below.
-                        // writer.write("$L", valueParam);
-                    }
-
-                });
-            }
-        );
     }
 
     /**
@@ -209,7 +189,10 @@ private boolean isMemberOverwriteRequired(MemberShape member, Set<String> parent
 
         Shape memberTarget = model.expectShape(member.getTarget());
         parents.add(symbolProvider.toMemberName(member));
-        if (memberTarget instanceof StructureShape) {
+        if (memberTarget.isUnionShape()) {
+            // always call filterSensitiveLog for UnionShape
+            return true;
+        } else if (memberTarget.isStructureShape()) {
             Collection<MemberShape> structureMemberList = ((StructureShape) memberTarget).getAllMembers().values();
             for (MemberShape structureMember: structureMemberList) {
                 if (!parents.contains(symbolProvider.toMemberName(structureMember))

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/UnionGenerator.java

@@ -101,6 +101,16 @@
  *     if (value.bear !== undefined) return visitor.bear(value.bear);
  *     return visitor._(value.$unknown[0], value.$unknown[1]);
  *   }
+ *   export const filterSensitiveLog = (obj: Attacker) => {
+ *     if (obj.lion !== undefined)
+ *       return { lion: Lion.filterSensitiveLog(obj.lion) };
+ *     if (obj.tiger !== undefined)
+ *       return { tiger: Tiger.filterSensitiveLog(obj.tiger) };
+ *     if (obj.bear !== undefined)
+ *       return { bear: Bear.filterSensitiveLog(obj.bear) };
+ *     if (obj.$unknown !== undefined)
+ *       return { [obj.$unknown[0]]: 'UNKNOWN' };
+ *   }
  * }
  * }</pre>
  *
@@ -148,13 +158,15 @@ public void run() {
             writeUnionMemberInterfaces();
             writeVisitorType();
             writeVisitorFunction();
+            writeFilterSensitiveLog();
         });
     }
 
     private void writeUnionMemberInterfaces() {
         writer.openBlock("interface $$Base {", "}", () -> {
             writer.write("__type?: $S;", shape.getId().getName());
         });
+        writer.write("");
 
         for (MemberShape member : shape.getAllMembers().values()) {
             String name = variantMap.get(member.getMemberName());
@@ -170,6 +182,7 @@ private void writeUnionMemberInterfaces() {
                 }
                 writer.write("$$unknown?: never;");
             });
+            writer.write("");
         }
 
         // Write out the unknown variant.
@@ -179,6 +192,7 @@ private void writeUnionMemberInterfaces() {
             }
             writer.write("$$unknown: [string, any];");
         });
+        writer.write("");
     }
 
     private void writeVisitorType() {
@@ -189,6 +203,7 @@ private void writeVisitorType() {
             }
             writer.write("_: (name: string, value: any) => T;");
         });
+        writer.write("");
     }
 
     private void writeVisitorFunction() {
@@ -203,5 +218,28 @@ private void writeVisitorFunction() {
         }
         writer.write("return visitor._(value.$$unknown[0], value.$$unknown[1]);");
         writer.dedent().write("}");
+        writer.write("");
+    }
+
+    private void writeFilterSensitiveLog() {
+        String objectParam = "obj";
+        writer.openBlock("export const filterSensitiveLog = ($L: $L): any => {", "}",
+            objectParam, symbol.getName(),
+            () -> {
+                for (MemberShape member : shape.getAllMembers().values()) {
+                    String memberName = symbolProvider.toMemberName(member);
+                    StructuredMemberWriter structuredMemberWriter = new StructuredMemberWriter(
+                            model, symbolProvider, shape.getAllMembers().values());
+                    writer.openBlock("if (${1L}.${2L} !== undefined) return {${2L}: ", "};",
+                        objectParam, memberName, () -> {
+                            String memberParam = String.format("%s.%s", objectParam, memberName);
+                            structuredMemberWriter.writeMemberFilterSensitiveLog(writer, member, memberParam);
+                        }
+                    );
+                }
+                writer.write("if (${1L}.$$unknown !== undefined) return {[${1L}.$$unknown[0]]: 'UNKNOWN'};",
+                    objectParam);
+            }
+        );
     }
 }