From 866dae790bf2fd82cf1e0ee0e84f57e5c6b8e555 Mon Sep 17 00:00:00 2001
From: Kasper Kondzielski <kghost0@gmail.com>
Date: Thu, 26 Sep 2024 23:47:15 +0200
Subject: [PATCH] fix: Don't assume that all service bindings existed before
 (#2405)

---
 .../aws/traits/auth/diff/SigV4Migration.java  | 38 ++++++++++---------
 .../diffs/no-old-sigv4-new-service.a.smithy   | 21 ++++++++++
 .../diffs/no-old-sigv4-new-service.b.smithy   | 30 +++++++++++++++
 .../diffs/no-old-sigv4-new-service.events     |  1 +
 4 files changed, 73 insertions(+), 17 deletions(-)
 create mode 100644 smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.a.smithy
 create mode 100644 smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.b.smithy
 create mode 100644 smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.events

diff --git a/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/diff/SigV4Migration.java b/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/diff/SigV4Migration.java
index 890c9fda867..87e91abf114 100644
--- a/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/diff/SigV4Migration.java
+++ b/smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/diff/SigV4Migration.java
@@ -83,23 +83,27 @@ public List<ValidationEvent> evaluate(Differences differences) {
             }
             // Validate Operation effective auth schemes
             for (ServiceShape newServiceShape : newOperationServiceBindings) {
-                ServiceShape oldServiceShape = oldModel.expectShape(newServiceShape.getId(), ServiceShape.class);
-                OperationShape oldOperationShape = change.getOldShape();
-                List<ShapeId> oldAuthSchemes = oldServiceIndex
-                    .getEffectiveAuthSchemes(oldServiceShape, oldOperationShape)
-                    .keySet()
-                    .stream()
-                    .collect(Collectors.toList());
-                List<ShapeId> newAuthSchemes = newServiceIndex
-                    .getEffectiveAuthSchemes(newServiceShape, newOperationShape)
-                    .keySet()
-                    .stream()
-                    .collect(Collectors.toList());
-                validateMigration(
-                    newOperationShape,
-                    oldAuthSchemes,
-                    newAuthSchemes,
-                    events);
+                oldModel.getShape(newServiceShape.getId())
+                    .filter(Shape::isServiceShape)
+                    .map(s -> (ServiceShape) s)
+                    .ifPresent(oldServiceShape -> {
+                        OperationShape oldOperationShape = change.getOldShape();
+                        List<ShapeId> oldAuthSchemes = oldServiceIndex
+                            .getEffectiveAuthSchemes(oldServiceShape, oldOperationShape)
+                            .keySet()
+                            .stream()
+                            .collect(Collectors.toList());
+                        List<ShapeId> newAuthSchemes = newServiceIndex
+                            .getEffectiveAuthSchemes(newServiceShape, newOperationShape)
+                            .keySet()
+                            .stream()
+                            .collect(Collectors.toList());
+                        validateMigration(
+                            newOperationShape,
+                            oldAuthSchemes,
+                            newAuthSchemes,
+                            events);
+                    });
             }
         }
 
diff --git a/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.a.smithy b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.a.smithy
new file mode 100644
index 00000000000..e75d0177d05
--- /dev/null
+++ b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.a.smithy
@@ -0,0 +1,21 @@
+$version: "2.0"
+
+metadata suppressions = [
+    {
+        id: "ModifiedTrait.Update.smithy.api#auth"
+        namespace: "ns.foo"
+    }
+]
+
+namespace ns.foo
+
+@auth([])
+@httpBearerAuth
+service Service {
+    operations: [
+        Operation
+    ]
+}
+
+@auth([])
+operation Operation {}
diff --git a/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.b.smithy b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.b.smithy
new file mode 100644
index 00000000000..44870b7e793
--- /dev/null
+++ b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.b.smithy
@@ -0,0 +1,30 @@
+$version: "2.0"
+
+metadata suppressions = [
+    {
+        id: "ModifiedTrait.Update.smithy.api#auth"
+        namespace: "ns.foo"
+    }
+]
+
+namespace ns.foo
+
+@auth([httpBearerAuth])
+@httpBearerAuth
+service Service {
+    operations: [
+        Operation
+    ]
+}
+
+@auth([httpBearerAuth])
+@httpBearerAuth
+service NewService {
+    operations: [
+        Operation
+    ]
+}
+
+
+@auth([httpBearerAuth])
+operation Operation {}
diff --git a/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.events b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.events
new file mode 100644
index 00000000000..7ebdd3d2dc3
--- /dev/null
+++ b/smithy-aws-traits/src/test/resources/software/amazon/smithy/aws/traits/diffs/no-old-sigv4-new-service.events
@@ -0,0 +1 @@
+[NOTE] ns.foo#NewService: Added service `ns.foo#NewService` | AddedShape