Index: invite.module
===================================================================
RCS file: /cvs/drupal-contrib/contributions/modules/invite/invite.module,v
retrieving revision 1.25.2.9
diff -u -p -r1.25.2.9 invite.module
--- invite.module	4 Oct 2009 14:35:51 -0000	1.25.2.9
+++ invite.module	5 Oct 2009 12:15:24 -0000
@@ -468,7 +468,20 @@ function invite_user($op, &$edit, &$acco
         }
       }
       if ($invite) {
-        _invite_accept($invite, $account);
+        // Process the invitation and assign target roles to the user. These
+        // will be saved by user_save().
+        $roles = invite_process($invite, $account);
+
+        if ($roles) {
+          if (!isset($edit['roles']) || !is_array($edit['roles'])) {
+            $edit['roles'] = array();
+          }
+          $edit['roles'] += $roles;
+
+          // We need to notify other modules of the role change, otherwise they
+          // will be unaware of it.
+          user_module_invoke('update', $edit, $account);
+        }
 
         // Flag the inviting user, this triggers status notifications and
         // saves us some queries otherwise.
@@ -487,45 +500,59 @@ function invite_user($op, &$edit, &$acco
 }
 
 /**
- * Set an invitation's status to accepted.
+ * Process a user that accepted an invitation.
  *
  * @param $invite
  *   An invite object.
  * @param $account
- *   The user object of the invitee.
+ *   The user object that accepted the invitation.
+ * @return
+ *   Array of target roles for the invited user.
  */
-function _invite_accept($invite, $account) {
+function invite_process($invite, $account) {
   // Update the invitation record.
   db_query("UPDATE {invite} SET email = '%s', invitee = %d, joined = %d WHERE reg_code = '%s'", $account->mail, $account->uid, time(), $invite->reg_code);
+
   // Delete all invites to these e-mail addresses, except this one.
   db_query("DELETE FROM {invite} WHERE (email = '%s' OR email = '%s') AND reg_code <> '%s'", $invite->email, $account->mail, $invite->reg_code);
+
   // Add all users who invited this particular e-mail address to the
   // notification queue.
   db_query("INSERT INTO {invite_notifications} (uid, invitee) SELECT uid, %d from {invite} WHERE (email = '%s' OR email = '%s') AND canceled = 0", $account->uid, $invite->email, $account->mail);
-  // Escalate the invitee's role.
-  _invite_escalate_role($account);
-  // Unblock user account.
+
+  // Unblock the user account.
   db_query("UPDATE {users} SET status = 1 WHERE uid = %d", $account->uid);
+
+  // Determine target roles for invited user.
+  $roles = invite_target_roles($invite);
+
+  // Allow other modules to act on the role escalation.
+  $args = array('invitee' => $account, 'inviter' => $invite->inviter, 'roles' => &$roles);
+  module_invoke_all('invite', 'escalate', $args);
+
+  return $roles;
 }
 
 /**
- * Escalates an invited user's role, based on the role(s) of the inviter.
+ * Determine target roles based on the roles of an inviter.
  *
- * @param $account
- *   The user object of the invitee.
+ * @param $invite
+ *   An invite object.
+ * @return
+ *   Array of target roles for an invited user.
  */
-function _invite_escalate_role($account) {
+function invite_target_roles($invite) {
+  $targets = array();
+
   // Add a dummy entry to retrieve the default target role setting.
   $roles = array('default' => 'default');
 
   // Add roles of inviter.
-  $inviter_uid = db_result(db_query("SELECT uid FROM {invite} WHERE invitee = %d", $account->uid));
-  if ($inviter_uid && $inviter = user_load(array('uid' => $inviter_uid))) {
-    $roles = array_merge($roles, array_intersect($inviter->roles, user_roles(FALSE, 'send invitations')));
+  if ($invite->inviter) {
+    $roles = array_merge($roles, array_intersect($invite->inviter->roles, user_roles(FALSE, 'send invitations')));
   }
 
   // Map to configured target roles.
-  $targets = array();
   foreach ($roles as $rid => $role) {
     $target = variable_get('invite_target_role_'. $rid, DRUPAL_AUTHENTICATED_RID);
     if ($target != DRUPAL_AUTHENTICATED_RID) {
@@ -533,19 +560,7 @@ function _invite_escalate_role($account)
     }
   }
 
-  // Notify other modules of changed user.
-  $edit = array('roles' => $targets);
-  user_module_invoke('update', $edit, $account);
-
-  // Save new user role(s).
-  foreach ($targets as $target) {
-    db_query("DELETE FROM {users_roles} WHERE uid = %d AND rid = %d", $account->uid, $target);
-    db_query("INSERT INTO {users_roles} (uid, rid) VALUES (%d, %d)", $account->uid, $target);
-  }
-
-  // Notify other modules of role escalation.
-  $args = array('invitee' => $account, 'inviter' => $inviter, 'roles' => $targets);
-  module_invoke_all('invite', 'escalate', $args);
+  return $targets;
 }
 
 /**