From 0f1c37864d7bb1fb2cfced7ff61d0c2c4d5cf430 Mon Sep 17 00:00:00 2001
From: Michael Dowling <mtdowling@gmail.com>
Date: Thu, 24 Aug 2023 13:23:41 -0500
Subject: [PATCH] Catch exceptions when creating traits

Some traits like the http trait will throw exceptions that are not
SourceExceptions. This previously would cause the CLI to stop on
the first exception and not show pretty error output. This fixes that
by catching RuntimeExceptions to turn them into ValidationEvents.
---
 .../smithy/model/loader/LoaderTraitMap.java   | 13 ++++++++++
 .../model/loader/ModelAssemblerTest.java      | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java b/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
index 132d1661b89..1e15e276c92 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/loader/LoaderTraitMap.java
@@ -16,6 +16,8 @@
 package software.amazon.smithy.model.loader;
 
 import static java.lang.String.format;
+import static software.amazon.smithy.model.validation.Severity.ERROR;
+import static software.amazon.smithy.model.validation.Validator.MODEL_ERROR;
 
 import java.util.Collections;
 import java.util.HashMap;
@@ -107,6 +109,17 @@ private Trait createTrait(ShapeId target, ShapeId traitId, Node traitValue) {
             String message = format("Error creating trait `%s`: ", Trait.getIdiomaticTraitName(traitId));
             events.add(ValidationEvent.fromSourceException(e, message, target));
             return null;
+        } catch (RuntimeException e) {
+            events.add(ValidationEvent.builder()
+                    .id(MODEL_ERROR)
+                    .severity(ERROR)
+                    .shapeId(target)
+                    .sourceLocation(traitValue)
+                    .message(format("Error creating trait `%s`: %s",
+                                    Trait.getIdiomaticTraitName(traitId),
+                                    e.getMessage()))
+                    .build());
+            return null;
         }
     }
 
diff --git a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
index c50afdaaa93..6f3950fc00c 100644
--- a/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
+++ b/smithy-model/src/test/java/software/amazon/smithy/model/loader/ModelAssemblerTest.java
@@ -1317,4 +1317,28 @@ public void modelLoadingErrorsAreEmittedToListener() {
         assertThat(result.getValidationEvents(Severity.ERROR), hasSize(1));
         assertThat(events, equalTo(result.getValidationEvents()));
     }
+
+    @Test
+    public void exceptionsThrownWhenCreatingTraitsDontCrashSmithy() {
+        String document = "{\n"
+                          + "\"smithy\": \"" + Model.MODEL_VERSION + "\",\n"
+                          + "    \"shapes\": {\n"
+                          + "        \"ns.foo#Test\": {\n"
+                          + "            \"type\": \"string\",\n"
+                          + "            \"traits\": {\"smithy.foo#baz\": true}\n"
+                          + "        }\n"
+                          + "    }\n"
+                          + "}";
+        ValidatedResult<Model> result = new ModelAssembler()
+                .addUnparsedModel(SourceLocation.NONE.getFilename(), document)
+                .putProperty(ModelAssembler.ALLOW_UNKNOWN_TRAITS, true)
+                .traitFactory((traitId, target, value) -> {
+                    throw new RuntimeException("Oops!");
+                })
+                .assemble();
+
+        assertThat(result.getValidationEvents(Severity.ERROR), not(empty()));
+        assertThat(result.getValidationEvents(Severity.ERROR).get(0).getMessage(),
+                   equalTo("Error creating trait `smithy.foo#baz`: Oops!"));
+    }
 }