Add user claims creation

.gitignore

@@ -15,4 +15,7 @@ Gemfile.lock
 .rspec_status
 
 # ruby version
-.ruby-version
+.ruby-version
+
+# asdf version manager
+.tool-versions

lib/firebase/admin/auth/client.rb

@@ -105,6 +105,18 @@ def create_session_cookie(id_token, valid_duration = 432000)
           @user_manager.create_session_cookie(id_token, valid_duration)
         end
 
+        # Sets custom claims for a user.
+        #
+        # @param [String] uid The id of the user.
+        # @param [Hash, nil] custom_claims The custom claims to set for the user. Pass nil to remove all custom claims.
+        #
+        # @raise [SetCustomUserClaimsError] if the operation fails.
+        #
+        # @return [UserRecord]
+        def set_custom_user_claims(uid, custom_claims)
+          @user_manager.set_custom_user_claims(uid, custom_claims)
+        end
+
         private
 
         # Checks if an ID token has been revoked.

lib/firebase/admin/auth/error.rb

@@ -18,6 +18,9 @@ class ExpiredTokenError < Error; end
 
       # Raised when a user cannot be created.
       class CreateUserError < Error; end
+
+      # Raised when setting custom user claims fails.
+      class SetCustomUserClaimsError < Error; end
     end
   end
 end

lib/firebase/admin/auth/user_manager.rb

@@ -87,6 +87,23 @@ def create_session_cookie(id_token, valid_duration = 432000)
           @client.post("projects/#{@project_id}:createSessionCookie", payload).body
         end
 
+        # Sets custom claims for a user.
+        #
+        # @param [String] uid The id of the user.
+        # @param [Hash, nil] custom_claims The custom claims to set for the user. Pass nil to remove all custom claims.
+        #
+        # @raise [ArgumentError] if uid is missing or invalid.
+        # @raise [SetCustomUserClaimsError] if the operation fails.
+        def set_custom_user_claims(uid, custom_claims)
+          payload = {
+            localId: validate_uid(uid, required: true),
+            customAttributes: custom_claims.nil? ? nil : custom_claims.to_json
+          }
+          res = @client.post(with_path("accounts:update"), payload).body
+          raise SetCustomUserClaimsError, "failed to set custom claims: #{res}" if res&.fetch("localId", nil).nil?
+          get_user_by(uid: uid)
+        end
+
         private
 
         def with_path(path)

spec/unit/firebase/admin/auth/client_spec.rb

@@ -38,4 +38,56 @@
       expect(session_cookie["validDuration"])
     end
   end
+
+  describe "#set_custom_user_claims" do
+    let(:uid) { "test-uid" }
+    let(:claims) { {admin: true} }
+
+    context "when setting custom claims" do
+      before do
+        stub_auth_request(:post, "/accounts:update")
+          .with(body: hash_including({localId: uid, customAttributes: claims.to_json}))
+          .to_return({body: {localId: uid}.to_json, headers: {content_type: "application/json; charset=utf-8"}})
+        stub_auth_request(:post, "/accounts:lookup")
+          .to_return({body: {users: [{localId: uid, customAttributes: {admin: true}.to_json}]}.to_json, headers: {content_type: "application/json; charset=utf-8"}})
+      end
+
+      it "sets custom claims for the user and returns the user record" do
+        user = @app.auth.set_custom_user_claims(uid, claims)
+        expect(user).to be_a(Firebase::Admin::Auth::UserRecord)
+        expect(user.uid).to eq(uid)
+        expect(user.custom_claims["admin"]).to eq(true)
+      end
+    end
+
+    context "when removing all custom claims" do
+      before do
+        stub_auth_request(:post, "/accounts:update")
+          .with(body: hash_including({localId: uid, customAttributes: nil}))
+          .to_return({body: {localId: uid}.to_json, headers: {content_type: "application/json; charset=utf-8"}})
+        stub_auth_request(:post, "/accounts:lookup")
+          .to_return({body: {users: [{localId: uid, customAttributes: nil}]}.to_json, headers: {content_type: "application/json; charset=utf-8"}})
+      end
+
+      it "removes all custom claims for the user" do
+        user = @app.auth.set_custom_user_claims(uid, nil)
+        expect(user).to be_a(Firebase::Admin::Auth::UserRecord)
+        expect(user.uid).to eq(uid)
+        expect(user.custom_claims).to be_nil.or eq({})
+      end
+    end
+
+    context "when the operation fails" do
+      before do
+        stub_auth_request(:post, "/accounts:update")
+          .to_return({body: {error: "something went wrong"}.to_json, headers: {content_type: "application/json; charset=utf-8"}})
+      end
+
+      it "raises SetCustomUserClaimsError" do
+        expect {
+          @app.auth.set_custom_user_claims(uid, claims)
+        }.to raise_error(Firebase::Admin::Auth::SetCustomUserClaimsError)
+      end
+    end
+  end
 end