Require valid annotation member names. (#620)

2 files changed, 26 insertions, 0 deletions

diff --git a/src/main/java/com/squareup/javapoet/AnnotationSpec.java b/src/main/java/com/squareup/javapoet/AnnotationSpec.java
index b5d6ab6..0c96756 100644
--- a/src/main/java/com/squareup/javapoet/AnnotationSpec.java
+++ b/src/main/java/com/squareup/javapoet/AnnotationSpec.java
@@ -27,6 +27,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import javax.lang.model.SourceVersion;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.ExecutableElement;
@@ -36,6 +37,7 @@ import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.SimpleAnnotationValueVisitor7;
 
 import static com.squareup.javapoet.Util.characterLiteralWithoutSingleQuotes;
+import static com.squareup.javapoet.Util.checkArgument;
 import static com.squareup.javapoet.Util.checkNotNull;
 
 /** A generated annotation on a declaration. */
@@ -206,6 +208,8 @@ public final class AnnotationSpec {
     }
 
     public Builder addMember(String name, CodeBlock codeBlock) {
+      checkNotNull(name, "name == null");
+      checkArgument(SourceVersion.isName(name), "not a valid name: %s", name);
       List<CodeBlock> values = members.get(name);
       if (values == null) {
         values = new ArrayList<>();
@@ -223,6 +227,7 @@ public final class AnnotationSpec {
     Builder addMemberForValue(String memberName, Object value) {
       checkNotNull(memberName, "memberName == null");
       checkNotNull(value, "value == null, constant non-null value expected for %s", memberName);
+      checkArgument(SourceVersion.isName(memberName), "not a valid name: %s", memberName);
       if (value instanceof Class<?>) {
         return addMember(memberName, "$T.class", value);
       }
diff --git a/src/test/java/com/squareup/javapoet/AnnotationSpecTest.java b/src/test/java/com/squareup/javapoet/AnnotationSpecTest.java
index a530111..49606c7 100644
--- a/src/test/java/com/squareup/javapoet/AnnotationSpecTest.java
+++ b/src/test/java/com/squareup/javapoet/AnnotationSpecTest.java
@@ -25,6 +25,7 @@ import org.junit.Rule;
 import org.junit.Test;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.fail;
 
 public final class AnnotationSpecTest {
 
@@ -350,6 +351,26 @@ public final class AnnotationSpecTest {
         + "}\n");
   }
 
+  @Test public void disallowsNullMemberName() {
+    AnnotationSpec.Builder builder = AnnotationSpec.builder(HasDefaultsAnnotation.class);
+    try {
+      AnnotationSpec.Builder $L = builder.addMember(null, "$L", "");
+      fail($L.build().toString());
+    } catch (NullPointerException e) {
+      assertThat(e).hasMessageThat().isEqualTo("name == null");
+    }
+  }
+
+  @Test public void requiresValidMemberName() {
+    AnnotationSpec.Builder builder = AnnotationSpec.builder(HasDefaultsAnnotation.class);
+    try {
+      AnnotationSpec.Builder $L = builder.addMember("@", "$L", "");
+      fail($L.build().toString());
+    } catch (IllegalArgumentException e) {
+      assertThat(e).hasMessageThat().isEqualTo("not a valid name: @");
+    }
+  }
+
   private String toString(TypeSpec typeSpec) {
     return JavaFile.builder("com.squareup.tacos", typeSpec).build().toString();
   }