From 568a77b1bf7f01333399f8a6bc21780ae741f21f Mon Sep 17 00:00:00 2001
From: Jonathan Wood <jonathan.wood@nominum.com>
Date: Sat, 4 Jul 2015 17:50:52 -0700
Subject: [PATCH] Refuse to write TOML fragments.

---
 .../com/moandjiezana/toml/TomlWriter.java     | 15 +++++--
 .../com/moandjiezana/toml/TomlWriterTest.java | 45 ++++++++++++++++---
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/moandjiezana/toml/TomlWriter.java b/src/main/java/com/moandjiezana/toml/TomlWriter.java
index d4b30e2..e698e3c 100644
--- a/src/main/java/com/moandjiezana/toml/TomlWriter.java
+++ b/src/main/java/com/moandjiezana/toml/TomlWriter.java
@@ -1,7 +1,5 @@
 package com.moandjiezana.toml;
 
-import static com.moandjiezana.toml.ValueWriters.WRITERS;
-
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
@@ -13,6 +11,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
 
+import static com.moandjiezana.toml.MapValueWriter.*;
+import static com.moandjiezana.toml.ObjectValueWriter.*;
+import static com.moandjiezana.toml.ValueWriters.WRITERS;
+
 /**
  * <p>Converts Objects to TOML</p>
  *
@@ -144,9 +146,16 @@ public class TomlWriter {
    * @param from the object to be written
    * @param target the Writer to which TOML will be written. The Writer is not closed.
    * @throws IOException if target.write() fails
+   * @throws IllegalStateException if
    */
   public void write(Object from, Writer target) throws IOException {
     WriterContext context = new WriterContext(indentationPolicy, datePolicy, target);
-    WRITERS.findWriterFor(from).write(from, context);
+
+    ValueWriter writer = WRITERS.findWriterFor(from);
+    if (writer == MAP_VALUE_WRITER || writer == OBJECT_VALUE_WRITER) {
+      WRITERS.findWriterFor(from).write(from, context);
+    } else {
+      throw new IllegalStateException("Top-level value must not be a primitive or array.");
+    }
   }
 }
diff --git a/src/test/java/com/moandjiezana/toml/TomlWriterTest.java b/src/test/java/com/moandjiezana/toml/TomlWriterTest.java
index 4ec2642..708db04 100644
--- a/src/test/java/com/moandjiezana/toml/TomlWriterTest.java
+++ b/src/test/java/com/moandjiezana/toml/TomlWriterTest.java
@@ -281,12 +281,17 @@ public class TomlWriterTest {
 
   @Test
   public void should_write_strings_to_toml_utf8() throws UnsupportedEncodingException {
-    String input = " é foo € \b \t \n \f \r \" \\ ";
-    assertEquals("\" \\u00E9 foo \\u20AC \\b \\t \\n \\f \\r \\\" \\\\ \"", new TomlWriter().write(input));
+    class Utf8Test {
+      String input;
+    }
+
+    Utf8Test utf8Test = new Utf8Test();
+    utf8Test.input = " é foo € \b \t \n \f \r \" \\ ";
+    assertEquals("input = \" \\u00E9 foo \\u20AC \\b \\t \\n \\f \\r \\\" \\\\ \"\n", new TomlWriter().write(utf8Test));
 
     // Check unicode code points greater than 0XFFFF
-    input = " \uD801\uDC28 \uD840\uDC0B ";
-    assertEquals("\" \\U00010428 \\U0002000B \"", new TomlWriter().write(input));
+    utf8Test.input = " \uD801\uDC28 \uD840\uDC0B ";
+    assertEquals("input = \" \\U00010428 \\U0002000B \"\n", new TomlWriter().write(utf8Test));
   }
 
   @Test
@@ -438,11 +443,41 @@ public class TomlWriterTest {
     
     assertEquals(expected, writer.write(o));
   }
-  
+
   private static class SimpleTestClass {
     int a = 1;
   }
 
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_string_fragment() {
+    new TomlWriter().write("fragment");
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_boolean_fragment() {
+    new TomlWriter().write(true);
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_number_fragment() {
+    new TomlWriter().write(42);
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_date_fragment() {
+    new TomlWriter().write(new Date());
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_array_fragment() {
+    new TomlWriter().write(new int[2]);
+  }
+
+  @Test(expected = IllegalStateException.class)
+  public void should_refuse_to_write_table_array_fragment() {
+    new TomlWriter().write(new SimpleTestClass[2]);
+  }
+
   @Test
   public void should_write_to_writer() throws IOException {
     StringWriter output = new StringWriter();