--- rolesignup.module	2007-01-20 16:49:48.000000000 -0500
+++ rolesignup.module	2007-08-18 15:34:33.000000000 -0400
@@ -31,8 +31,15 @@
 function rolesignup_form_alter($form_id, &$form) {
   if ($form_id == "user_register") {
     global $user;
+    //Invite support
+    if(module_exists('invite')) {
+      if($invite_code=arg(2)) {
+        $_SESSION['invite_code'] = $invite_code;
+      }
+    }
     if (!user_access('administer users') && !$_SESSION['role']) {
-      drupal_goto('user/register/role', drupal_get_destination());
+//      drupal_goto('user/register/role', drupal_get_destination());
+      drupal_goto('user/register/role');
     } else {
       $form['role'] = array(
         '#type' => 'value',
@@ -41,16 +48,17 @@
     }
     $form['rolelink'] = array(
       '#value' => theme('select_role'),
-      '#weight' => -11,      
+      '#weight' => -11,
       );
   }
-} 
+}
 
 function role_register_page() {
   $roles = user_roles(1, 'register for role');
   if (isset($roles[arg(3)])) {
     $_SESSION['role'] = arg(3);
-    drupal_goto('user/register');
+    $url = 'user/register/'.$_SESSION['invite_code'];
+    drupal_goto($url);
   }
   return theme('rolelist', $roles);
 }
@@ -63,6 +71,7 @@
     case 'insert':
     $roles = user_roles(1,'register for role');
     if ($roles[$_SESSION['role']]) {
+      db_query("DELETE FROM {users_roles} WHERE uid=%d AND rid=%d", $account->uid, $_SESSION['role']);
       db_query("INSERT INTO {users_roles} (uid, rid) VALUES (%d, %d)", $account->uid, $_SESSION['role']);
     }
   }