Properly handle single statement IAM policies

AWS will not wrap the Statements in single statement policies in an array,
they will just have the single statement object.

This adds a custom UnmarshalJSON methods which deal with this in a way similar
to Principal (another "adaptable" field.)

iam_policy.go

@@ -8,6 +8,16 @@ type IAMPolicyDocument struct {
 	Statement []IAMPolicyStatement
 }
 
+// Avoid infinite loops when we just want to unmarshal the struct normally.
+type iamPolicyDocumentCopy IAMPolicyDocument
+
+// iamPolicyDocumentSingleStatement is used for parsing policy documents with a
+// single statement.
+type iamPolicyDocumentSingleStatement struct {
+	Version   string `json:",omitempty"`
+	Statement IAMPolicyStatement
+}
+
 // ToJSON returns the JSON representation of the policy document or
 // panics if the object cannot be marshaled.
 func (i IAMPolicyDocument) ToJSON() string {
@@ -18,6 +28,31 @@ func (i IAMPolicyDocument) ToJSON() string {
 	return string(buf)
 }
 
+// UnmarshalJSON sets the object from the provided JSON representation. This has
+// been added to handle the special case of a single statement versus an array.
+func (i *IAMPolicyDocument) UnmarshalJSON(data []byte) error {
+	// Handle single statement policy documents
+	var v iamPolicyDocumentSingleStatement
+	err := json.Unmarshal(data, &v)
+	if err == nil {
+		i.Version = v.Version
+		i.Statement = []IAMPolicyStatement{v.Statement}
+		return nil
+	}
+
+	// Handle multiple statements
+	var v2 iamPolicyDocumentCopy
+	err = json.Unmarshal(data, &v2)
+	if err != nil {
+		return err
+	}
+
+	i.Version = v2.Version
+	i.Statement = v2.Statement
+
+	return nil
+}
+
 // IAMPrincipal represents a principal in an IAM policy
 type IAMPrincipal struct {
 	AWS           *StringListExpr `json:",omitempty"`
@@ -38,7 +73,7 @@ type IAMPolicyStatement struct {
 	Condition    interface{}     `json:",omitempty"`
 }
 
-// Avoid infinite loops when we just want to marshal the struct normally
+// Avoid infinite loops when we just want to marshal the struct normally.
 type iamPrincipalCopy IAMPrincipal
 
 // MarshalJSON returns a JSON representation of the object. This has been added

iam_policy_test.go

@@ -45,23 +45,22 @@ func (testSuite *IAMPolicyTest) TestIAMPolicyStatement(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(string(b), Equals, `{"Version":"2012-10-17","Statement":[{"Sid":"AddPerm","Effect":"Allow","Principal":"*","Action":["s3:GetObject"],"Resource":["arn:aws:s3:::examplebucket/*"]}]}`)
 
-	// Try other variations of Principal
+	// Try other variations of Principal, as well as a single statement rather than an array
 	p = `{
    "Version":"2012-10-17",
    "Id":"PolicyForCloudFrontPrivateContent",
-   "Statement":[
-     {
+   "Statement": {
        "Sid":" Grant a CloudFront Origin Identity access to support private content",
        "Effect":"Allow",
        "Principal":{"CanonicalUser":"79a59df900b949e55d96a1e698fbacedfd6e09d98eacf8f8d5218e7cd47ef2be"},
        "Action":"s3:GetObject",
        "Resource":"arn:aws:s3:::example-bucket/*"
-     }
-   ]
+    }
 }`
 
 	err = json.Unmarshal([]byte(p), &v)
 	c.Assert(err, IsNil)
+	s = v.Statement[0]
 	c.Assert(s.Principal.AWS, IsNil)
 	c.Assert(s.Principal.CanonicalUser, DeepEquals, StringList(String("79a59df900b949e55d96a1e698fbacedfd6e09d98eacf8f8d5218e7cd47ef2be")))
 	c.Assert(s.Principal.Federated, IsNil)
@@ -93,6 +92,7 @@ func (testSuite *IAMPolicyTest) TestIAMPolicyStatement(c *C) {
 
 	err = json.Unmarshal([]byte(p), &v)
 	c.Assert(err, IsNil)
+	s = v.Statement[0]
 	c.Assert(s.Principal.AWS, IsNil)
 	c.Assert(s.Principal.CanonicalUser, IsNil)
 	c.Assert(s.Principal.Federated, IsNil)