Add check for mutability of init payer

coral-xyz · Feb 7, 2022 · 2103e2b · 2103e2b
1 parent a662577
commit 2103e2b
Showing 1 changed file with 62 additions and 32 deletions.
diff --git a/lang/syn/src/parser/accounts/mod.rs b/lang/syn/src/parser/accounts/mod.rs
@@ -39,55 +39,85 @@ pub fn parse(strct: &syn::ItemStruct) -> ParseResult<AccountsStruct> {
 
 fn constraints_cross_checks(fields: &[AccountField]) -> ParseResult<()> {
     // INIT
-    let init_field = fields.iter().find(|f| {
-        if let AccountField::Field(field) = f {
-            field.constraints.init.is_some()
-        } else {
-            false
-        }
-    });
-    if let Some(init_field) = init_field {
+    let init_fields: Vec<&Field> = fields
+        .iter()
+        .filter_map(|f| match f {
+            AccountField::Field(field) if field.constraints.init.is_some() => Some(field),
+            _ => None,
+        })
+        .collect();
+
+    if !init_fields.is_empty() {
         // init needs system program.
         if fields.iter().all(|f| f.ident() != "system_program") {
             return Err(ParseError::new(
-                init_field.ident().span(),
+                init_fields[0].ident.span(),
                 "the init constraint requires \
                 the system_program field to exist in the account \
                 validation struct. Use the program type to add \
                 the system_program field to your validation struct.",
             ));
         }
-        if let AccountField::Field(field) = init_field {
-            let kind = &field.constraints.init.as_ref().unwrap().kind;
-            // init token/a_token/mint needs token program.
-            match kind {
-                InitKind::Program { .. } => (),
-                InitKind::Token { .. }
-                | InitKind::AssociatedToken { .. }
-                | InitKind::Mint { .. } => {
-                    if fields.iter().all(|f| f.ident() != "token_program") {
-                        return Err(ParseError::new(
-                            init_field.ident().span(),
-                            "the init constraint requires \
+
+        let kind = &init_fields[0].constraints.init.as_ref().unwrap().kind;
+        // init token/a_token/mint needs token program.
+        match kind {
+            InitKind::Program { .. } => (),
+            InitKind::Token { .. } | InitKind::AssociatedToken { .. } | InitKind::Mint { .. } => {
+                if fields.iter().all(|f| f.ident() != "token_program") {
+                    return Err(ParseError::new(
+                        init_fields[0].ident.span(),
+                        "the init constraint requires \
                             the token_program field to exist in the account \
                             validation struct. Use the program type to add \
                             the token_program field to your validation struct.",
-                        ));
-                    }
+                    ));
                 }
             }
-            // a_token needs associated token program.
-            if let InitKind::AssociatedToken { .. } = kind {
-                if fields
-                    .iter()
-                    .all(|f| f.ident() != "associated_token_program")
-                {
-                    return Err(ParseError::new(
-                        init_field.ident().span(),
-                        "the init constraint requires \
+        }
+        // a_token needs associated token program.
+        if let InitKind::AssociatedToken { .. } = kind {
+            if fields
+                .iter()
+                .all(|f| f.ident() != "associated_token_program")
+            {
+                return Err(ParseError::new(
+                    init_fields[0].ident.span(),
+                    "the init constraint requires \
                     the associated_token_program field to exist in the account \
                     validation struct. Use the program type to add \
                     the associated_token_program field to your validation struct.",
+                ));
+            }
+        }
+
+        for field in init_fields {
+            // Get payer for init-ed account
+            let associated_payer_name = match field.constraints.init.clone().unwrap().payer.unwrap()
+            {
+                // composite payer, check not supported
+                Expr::Field(_) => continue,
+                field_name => field_name.to_token_stream().to_string(),
+            };
+
+            // Check payer is mutable
+            let associated_payer_field = fields.iter().find_map(|f| match f {
+                AccountField::Field(field) if *f.ident() == associated_payer_name => Some(field),
+                _ => None,
+            });
+            match associated_payer_field {
+                Some(associated_payer_field) => {
+                    if !associated_payer_field.constraints.is_mutable() {
+                        return Err(ParseError::new(
+                            field.ident.span(),
+                            "the payer specified for an init constraint must be mutable.",
+                        ));
+                    }
+                }
+                _ => {
+                    return Err(ParseError::new(
+                        field.ident.span(),
+                        "the payer specified does not exist.",
                     ));
                 }
             }