2 files changed, 24 insertions, 45 deletions

diff --git a/extensions/PhabBugz/lib/Feed.pm b/extensions/PhabBugz/lib/Feed.pm
index 4799bd0a3..4d2732f94 100644
--- a/extensions/PhabBugz/lib/Feed.pm
+++ b/extensions/PhabBugz/lib/Feed.pm
@@ -385,8 +385,8 @@ sub process_revision_change {
         $story_text);
     INFO($log_message);
 
-    # Pre setup before making changes
-    my $old_user = set_phab_user();
+    # change to the phabricator user, which returns a guard that restores the previous user.
+    my $restore_prev_user = set_phab_user();
     my $bug = $revision->bug;
 
     # Check to make sure bug id is valid and author can see it
@@ -619,8 +619,6 @@ sub process_revision_change {
         Bugzilla::BugMail::Send($bug_id, { changer => $rev_attachment->attacher });
     }
 
-    Bugzilla->set_user($old_user);
-
     INFO('SUCCESS: Revision D' . $revision->id . ' processed');
 }
 
@@ -639,7 +637,7 @@ sub process_new_user {
     my $bug_user  = $phab_user->bugzilla_user;
 
     # Pre setup before querying DB
-    my $old_user = set_phab_user();
+    my $restore_prev_user = set_phab_user();
 
     # CHECK AND WARN FOR POSSIBLE USERNAME SQUATTING
     INFO("Checking for username squatters");
@@ -758,8 +756,6 @@ sub process_new_user {
         }
     }
 
-    Bugzilla->set_user($old_user);
-
     INFO('SUCCESS: User ' . $phab_user->id . ' processed');
 }
 
@@ -868,11 +864,8 @@ sub add_flag_comment {
     my ( $bug, $attachment, $comment, $user, $old_flags, $new_flags, $timestamp )
         = @$params{qw(bug attachment comment user old_flags new_flags timestamp)};
 
-    my $old_user;
-    if ($user) {
-        $old_user = Bugzilla->user;
-        Bugzilla->set_user($user);
-    }
+    # when this function returns, Bugzilla->user will return to its previous value.
+    my $restore_prev_user = Bugzilla->set_user($user, scope_guard => 1);
 
     INFO("Flag comment: $comment");
     $bug->add_comment(
@@ -886,8 +879,6 @@ sub add_flag_comment {
 
     $attachment->set_flags( $old_flags, $new_flags );
     $attachment->update($timestamp);
-
-    Bugzilla->set_user($old_user) if $old_user;
 }
 
 1;
diff --git a/extensions/PhabBugz/lib/Util.pm b/extensions/PhabBugz/lib/Util.pm
index 4e846badc..9c79c1855 100644
--- a/extensions/PhabBugz/lib/Util.pm
+++ b/extensions/PhabBugz/lib/Util.pm
@@ -65,37 +65,26 @@ sub create_revision_attachment {
     }
 
     # If submitter, then switch to that user when creating attachment
-    my ($old_user, $attachment);
-    try {
-        if ($submitter) {
-            $old_user = Bugzilla->user;
-            $submitter->{groups} = [ Bugzilla::Group->get_all ]; # We need to always be able to add attachment
-            Bugzilla->set_user($submitter);
+    local $submitter->{groups} = [ Bugzilla::Group->get_all ]; # We need to always be able to add attachment
+    my $restore_prev_user = Bugzilla->set_user($submitter, scope_guard => 1);
+
+    my $attachment = Bugzilla::Attachment->create(
+        {
+            bug         => $bug,
+            creation_ts => $timestamp,
+            data        => $revision_uri,
+            description => $revision->title,
+            filename    => 'phabricator-D' . $revision->id . '-url.txt',
+            ispatch     => 0,
+            isprivate   => 0,
+            mimetype    => PHAB_CONTENT_TYPE,
         }
+    );
 
-        $attachment = Bugzilla::Attachment->create(
-            {
-                bug         => $bug,
-                creation_ts => $timestamp,
-                data        => $revision_uri,
-                description => $revision->title,
-                filename    => 'phabricator-D' . $revision->id . '-url.txt',
-                ispatch     => 0,
-                isprivate   => 0,
-                mimetype    => PHAB_CONTENT_TYPE,
-            }
-        );
+    # Insert a comment about the new attachment into the database.
+    $bug->add_comment($revision->summary, { type       => CMT_ATTACHMENT_CREATED,
+                                            extra_data => $attachment->id });
 
-        # Insert a comment about the new attachment into the database.
-        $bug->add_comment($revision->summary, { type       => CMT_ATTACHMENT_CREATED,
-                                                extra_data => $attachment->id });
-    }
-    catch {
-        die $_;
-    }
-    finally {
-        Bugzilla->set_user($old_user) if $old_user;
-    };
 
     return $attachment;
 }
@@ -210,11 +199,10 @@ sub request {
 }
 
 sub set_phab_user {
-    my $old_user = Bugzilla->user;
     my $user = Bugzilla::User->new( { name => PHAB_AUTOMATION_USER } );
     $user->{groups} = [ Bugzilla::Group->get_all ];
-    Bugzilla->set_user($user);
-    return $old_user;
+
+    return Bugzilla->set_user($user, scope_guard => 1);
 }
 
 sub get_needs_review {