From 8cbe3bee31e605d535783431bfb74c3e2867ee37 Mon Sep 17 00:00:00 2001
From: Severin Gehwolf <sgehwolf@redhat.com>
Date: Wed, 15 Apr 2020 17:41:17 +0200
Subject: [PATCH] Fix CVE-2017-18640 and add a test

---
 collector/pom.xml                             |  2 +-
 .../java/io/prometheus/jmx/JmxCollector.java  | 12 ++++++++-
 .../io/prometheus/jmx/JmxCollectorTest.java   | 26 ++++++++++++++++---
 collector/src/test/resources/testyaml.config  |  9 +++++++
 4 files changed, 44 insertions(+), 5 deletions(-)
 create mode 100644 collector/src/test/resources/testyaml.config

diff --git a/collector/pom.xml b/collector/pom.xml
index 8e90d58..b58d94a 100644
--- a/collector/pom.xml
+++ b/collector/pom.xml
@@ -29,7 +29,7 @@
     <dependency>
       <groupId>org.yaml</groupId>
       <artifactId>snakeyaml</artifactId>
-      <version>1.16</version>
+      <version>1.26</version>
     </dependency>
   </dependencies>
 
diff --git a/collector/src/main/java/io/prometheus/jmx/JmxCollector.java b/collector/src/main/java/io/prometheus/jmx/JmxCollector.java
index f06d3ef..a5fc96f 100644
--- a/collector/src/main/java/io/prometheus/jmx/JmxCollector.java
+++ b/collector/src/main/java/io/prometheus/jmx/JmxCollector.java
@@ -3,6 +3,7 @@ package io.prometheus.jmx;
 import io.prometheus.client.Collector;
 import io.prometheus.client.Counter;
 import org.yaml.snakeyaml.Yaml;
+import org.yaml.snakeyaml.error.YAMLException;
 
 import javax.management.MalformedObjectNameException;
 import javax.management.ObjectName;
@@ -71,7 +72,16 @@ public class JmxCollector extends Collector implements Collector.Describable {
 
     public JmxCollector(File in) throws IOException, MalformedObjectNameException {
         configFile = in;
-        config = loadConfig((Map<String, Object>)new Yaml().load(new FileReader(in)));
+        // Be defensive about loading yaml from a user. In some cases YAMLException
+        // will be thrown for bad configs
+        Map<String, Object> yamlConfig = null;
+        try {
+            yamlConfig = (Map<String, Object>)new Yaml().load(new FileReader(in));
+        } catch (YAMLException e) {
+            System.err.println("YAML configuration error: " + e.getMessage());
+            throw new IllegalArgumentException(e);
+        }
+        config = loadConfig(yamlConfig);
         config.lastUpdate = configFile.lastModified();
     }
 
diff --git a/collector/src/test/java/io/prometheus/jmx/JmxCollectorTest.java b/collector/src/test/java/io/prometheus/jmx/JmxCollectorTest.java
index dc35734..3d4ecf0 100644
--- a/collector/src/test/java/io/prometheus/jmx/JmxCollectorTest.java
+++ b/collector/src/test/java/io/prometheus/jmx/JmxCollectorTest.java
@@ -3,15 +3,21 @@ package io.prometheus.jmx;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import io.prometheus.client.Collector;
-import io.prometheus.client.CollectorRegistry;
+import java.io.File;
 import java.lang.management.ManagementFactory;
+
 import javax.management.MBeanServer;
-import org.junit.Test;
+
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Test;
+import org.yaml.snakeyaml.error.YAMLException;
+
+import io.prometheus.client.Collector;
+import io.prometheus.client.CollectorRegistry;
 
 
 public class JmxCollectorTest {
@@ -252,4 +258,18 @@ public class JmxCollectorTest {
       Thread.sleep(2000);
       assertEquals(1.0, registry.getSampleValue("boolean_Test_True", new String[]{}, new String[]{}), .001);
     }
+
+    @Test
+    public void testBillionLaughs() throws Exception {
+      File configFile = new File(getClass().getResource("/testyaml.config").getPath());
+      assertTrue(configFile.exists());
+      try {
+        JmxCollector jc = new JmxCollector(configFile);
+        fail("Excected yaml exception due to billion laughs");
+      } catch (IllegalArgumentException e) {
+        Throwable ex = e.getCause();
+        String prefix = YAMLException.class.getName() + ": ";
+        assertEquals(prefix + "Number of aliases for non-scalar nodes exceeds the specified max=50", e.getMessage());
+      }
+    }
 }
diff --git a/collector/src/test/resources/testyaml.config b/collector/src/test/resources/testyaml.config
new file mode 100644
index 0000000..4a3ed69
--- /dev/null
+++ b/collector/src/test/resources/testyaml.config
@@ -0,0 +1,9 @@
+a: &a ["lol","lol","lol","lol","lol","lol","lol","lol","lol"]
+b: &b [*a,*a,*a,*a,*a,*a,*a,*a,*a]
+c: &c [*b,*b,*b,*b,*b,*b,*b,*b,*b]
+d: &d [*c,*c,*c,*c,*c,*c,*c,*c,*c]
+e: &e [*d,*d,*d,*d,*d,*d,*d,*d,*d]
+f: &f [*e,*e,*e,*e,*e,*e,*e,*e,*e]
+g: &g [*f,*f,*f,*f,*f,*f,*f,*f,*f]
+h: &h [*g,*g,*g,*g,*g,*g,*g,*g,*g]
+i: &i [*h,*h,*h,*h,*h,*h,*h,*h,*h]
\ No newline at end of file
-- 
2.21.1