From 7b88bc2b8bad95ad26455cfa586d5242b673cd61 Mon Sep 17 00:00:00 2001
From: kstich <kevin@kstich.com>
Date: Tue, 15 Jun 2021 14:58:16 -0700
Subject: [PATCH] Disallow loading DTDs, etc. in protocol tests

---
 .../traits/ProtocolTestCaseValidator.java     | 11 +++++
 .../errorfiles/invalid-xml-with-dtd.errors    |  1 +
 .../errorfiles/invalid-xml-with-dtd.smithy    | 46 +++++++++++++++++++
 3 files changed, 58 insertions(+)
 create mode 100644 smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.errors
 create mode 100644 smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.smithy

diff --git a/smithy-protocol-test-traits/src/main/java/software/amazon/smithy/protocoltests/traits/ProtocolTestCaseValidator.java b/smithy-protocol-test-traits/src/main/java/software/amazon/smithy/protocoltests/traits/ProtocolTestCaseValidator.java
index e5e295abe04..defa679872b 100644
--- a/smithy-protocol-test-traits/src/main/java/software/amazon/smithy/protocoltests/traits/ProtocolTestCaseValidator.java
+++ b/smithy-protocol-test-traits/src/main/java/software/amazon/smithy/protocoltests/traits/ProtocolTestCaseValidator.java
@@ -66,6 +66,17 @@ abstract class ProtocolTestCaseValidator<T extends Trait> extends AbstractValida
         this.traitClass = traitClass;
         this.descriptor = descriptor;
         documentBuilderFactory = DocumentBuilderFactory.newInstance();
+
+        // Disallow loading DTDs and more for protocol test contents.
+        try {
+            documentBuilderFactory.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
+            documentBuilderFactory.setXIncludeAware(false);
+            documentBuilderFactory.setExpandEntityReferences(false);
+            documentBuilderFactory.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
+            documentBuilderFactory.setFeature("http://xml.org/sax/features/external-general-entities", false);
+        } catch (ParserConfigurationException e) {
+            throw new RuntimeException(e);
+        }
     }
 
     @Override
diff --git a/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.errors b/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.errors
new file mode 100644
index 00000000000..3de1436bc59
--- /dev/null
+++ b/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.errors
@@ -0,0 +1 @@
+[DANGER] smithy.example#SayHello: Invalid application/xml content in `smithy.test#httpRequestTests` protocol test case `foo1` | HttpRequestTestsInput
diff --git a/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.smithy b/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.smithy
new file mode 100644
index 00000000000..1a74ab9d5f7
--- /dev/null
+++ b/smithy-protocol-test-traits/src/test/resources/software/amazon/smithy/protocoltests/traits/errorfiles/invalid-xml-with-dtd.smithy
@@ -0,0 +1,46 @@
+namespace smithy.example
+
+use smithy.test#httpRequestTests
+
+@trait
+@protocolDefinition
+structure testProtocol {}
+
+@http(method: "POST", uri: "/")
+@httpRequestTests([
+    {
+        id: "foo1",
+        protocol: testProtocol,
+        method: "POST",
+        uri: "/",
+        params: {
+            type: true
+        },
+        bodyMediaType: "application/xml",
+        body: """
+        <!DOCTYPE root [
+            <!ENTITY hifi "hifi">
+            <!ENTITY hifi1 "&hifi;&hifi;&hifi;">
+            <!ENTITY hifi2 "&hifi1;&hifi1;&hifi1;">
+            <!ENTITY hifi3 "&hifi2;&hifi2;&hifi2;">
+        ]>
+        <XmlNamespacesResponse xmlns="https://example.com/">
+            <nested>
+                <foo xmlns:baz="http://baz.com">Foo</foo>
+                <values xmlns="http://qux.com">
+                    <member xmlns="http://bux.com">Bar</member>
+                    <member xmlns="http://bux.com">Baz</member>
+                </values>
+            </nested>
+            <RequestId>requestid</RequestId>
+        </XmlNamespacesResponse>
+        """
+    }
+])
+operation SayHello {
+    input: SayHelloInput
+}
+
+structure SayHelloInput {
+    type: Boolean
+}