Fix invite code being consumed on failed registration in invite-only mode (#1576)

* Fix:Bug:invite-code-inproper-burn

* added docstring according to coderabbitai warning

* updated feedback for merge

app/controllers/registrations_controller.rb

@@ -6,7 +6,6 @@ class RegistrationsController < ApplicationController
   before_action :ensure_signup_open, if: :self_hosted?
   before_action :set_user, only: :create
   before_action :set_invitation
-  before_action :claim_invite_code, only: :create, if: :invite_code_required?
   before_action :validate_password_requirements, only: :create
 
   def new
@@ -29,10 +28,10 @@ class RegistrationsController < ApplicationController
       @user.role = User.role_for_new_family_creator
     end
 
-    if @user.save
-      @invitation&.update!(accepted_at: Time.current)
-      @session = create_session_for(@user)
+    if signup_with_invite_claim!
       redirect_to root_path, notice: t(".success")
+    elsif @invite_code_invalid
+      redirect_to new_registration_path, alert: t("registrations.create.invalid_invite_code")
     else
       render :new, status: :unprocessable_entity, alert: t(".failure")
     end
@@ -55,10 +54,30 @@ class RegistrationsController < ApplicationController
       specific_param ? params[specific_param] : params
     end
 
-    def claim_invite_code
-      unless InviteCode.claim! params[:user][:invite_code]
-        redirect_to new_registration_path, alert: t("registrations.create.invalid_invite_code")
+    # Keep save+claim atomic so failed signups never burn valid invite codes.
+    def signup_with_invite_claim!
+      invite_code = user_params[:invite_code]
+      @invite_code_invalid = invite_code_required? && invite_code.blank?
+      return false if @invite_code_invalid
+
+      success = false
+
+      ActiveRecord::Base.transaction do
+        unless @user.save
+          raise ActiveRecord::Rollback
+        end
+
+        if invite_code_required? && !InviteCode.claim!(invite_code)
+          @invite_code_invalid = true
+          raise ActiveRecord::Rollback
+        end
+
+        @invitation&.update!(accepted_at: Time.current)
+        @session = create_session_for(@user)
+        success = true
       end
+
+      success
     end
 
     def validate_password_requirements

test/controllers/registrations_controller_test.rb

@@ -59,15 +59,46 @@ class RegistrationsControllerTest < ActionDispatch::IntegrationTest
       end
 
       assert_difference "User.count", +1 do
+        invite_code = InviteCode.generate!
         post registration_url, params: { user: {
           email: "john@example.com",
           password: "Password1!",
-          invite_code: InviteCode.generate! } }
+          invite_code: invite_code } }
         assert_redirected_to root_url
+        assert_not InviteCode.exists?(token: invite_code)
       end
     end
   end
 
+  test "invite code is not consumed when signup fails validation" do
+    with_env_overrides REQUIRE_INVITE_CODE: "true" do
+      invite_code = InviteCode.generate!
+
+      assert_no_difference "User.count" do
+        post registration_url, params: { user: {
+          email: "validationfail@example.com",
+          password: "weak",
+          invite_code: invite_code } }
+      end
+
+      assert_response :unprocessable_entity
+      assert InviteCode.exists?(token: invite_code)
+    end
+  end
+
+  test "invalid invite code does not create a user" do
+    with_env_overrides REQUIRE_INVITE_CODE: "true" do
+      assert_no_difference "User.count" do
+        post registration_url, params: { user: {
+          email: "valid@example.com",
+          password: "Password1!",
+          invite_code: "invalid-token-that-does-not-exist" } }
+      end
+
+      assert_redirected_to new_registration_url
+    end
+  end
+
   test "creating account from guest invitation assigns guest role and intro layout" do
     invitation = invitations(:one)
     invitation.update!(role: "guest", email: "guest-signup@example.com")