Add geo node enable api

- Add helper method to geo node for use in upcoming MR
- Add api helper to allow looking up disabled nodes
- Add method to jwt_decoder to include disabled geo nodes

Relates to https://gitlab.com/gitlab-org/gitlab/issues/35913

ee/app/models/geo_node.rb

@@ -196,6 +196,10 @@ class GeoNode < ApplicationRecord
     geo_api_url('status')
   end
 
+  def node_api_url(node)
+    api_url("geo_nodes/#{node.id}")
+  end
+
   def snapshot_url(repository)
     url = api_url("projects/#{repository.project.id}/snapshot")
     url += "?wiki=1" if repository.repo_type.wiki?

ee/lib/api/geo.rb

@@ -12,15 +12,11 @@ module API
           params.slice(*valid_attributes)
         end
 
-        def jwt_decoder
-          ::Gitlab::Geo::JwtRequestDecoder.new(headers['Authorization'])
-        end
-
         # Check if a Geo request is legit or fail the flow
         #
         # @param [Hash] attributes to be matched against JWT
         def authorize_geo_transfer!(**attributes)
-          unauthorized! unless jwt_decoder.valid_attributes?(**attributes)
+          unauthorized! unless geo_jwt_decoder.valid_attributes?(**attributes)
         end
       end
 
@@ -32,7 +28,7 @@ module API
         check_gitlab_geo_request_ip!
         authorize_geo_transfer!(replicable_name: params[:replicable_name], id: params[:id])
 
-        decoded_params = jwt_decoder.decode
+        decoded_params = geo_jwt_decoder.decode
         service = ::Geo::BlobUploadService.new(replicable_name: params[:replicable_name],
                                                blob_id: params[:id],
                                                decoded_params: decoded_params)
@@ -62,7 +58,7 @@ module API
         check_gitlab_geo_request_ip!
         authorize_geo_transfer!(file_type: params[:type], file_id: params[:id])
 
-        decoded_params = jwt_decoder.decode
+        decoded_params = geo_jwt_decoder.decode
         service = ::Geo::FileUploadService.new(params, decoded_params)
         response = service.execute
 

ee/lib/api/geo_nodes.rb

@@ -6,7 +6,24 @@ module API
     include APIGuard
     include ::Gitlab::Utils::StrongMemoize
 
-    before { authenticated_as_admin! }
+    before { authenticate_admin_or_geo_node! }
+
+    helpers do
+      def authenticate_admin_or_geo_node!
+        if gitlab_geo_node_token?
+          bad_request! unless update_geo_nodes_endpoint?
+          check_gitlab_geo_request_ip!
+          allow_paused_nodes!
+          authenticate_by_gitlab_geo_node_token!
+        else
+          authenticated_as_admin!
+        end
+      end
+
+      def update_geo_nodes_endpoint?
+        request.put? && request.path.match?(/\/geo_nodes\/\d+/)
+      end
+    end
 
     resource :geo_nodes do
       # Add a new Geo node

ee/lib/ee/api/helpers.rb

@@ -20,6 +20,19 @@ module EE
         render_api_error!(e.to_s, 401)
       end
 
+      def geo_jwt_decoder
+        return unless gitlab_geo_node_token?
+
+        strong_memoize(:geo_jwt_decoder) do
+          ::Gitlab::Geo::JwtRequestDecoder.new(headers['Authorization'])
+        end
+      end
+
+      # Update the jwt_decoder to allow authorization of disabled (paused) nodes
+      def allow_paused_nodes!
+        geo_jwt_decoder.include_disabled!
+      end
+
       def check_gitlab_geo_request_ip!
         unauthorized! unless ::Gitlab::Geo.allowed_ip?(request.ip)
       end
@@ -39,10 +52,9 @@ module EE
       end
 
       def authorization_header_valid?
-        auth_header = headers['Authorization']
-        return unless auth_header
+        return unless gitlab_geo_node_token?
 
-        scope = ::Gitlab::Geo::JwtRequestDecoder.new(auth_header).decode.try { |x| x[:scope] }
+        scope = geo_jwt_decoder.decode.try { |x| x[:scope] }
         scope == ::Gitlab::Geo::API_SCOPE
       end
 

ee/lib/gitlab/geo/jwt_request_decoder.rb

@@ -16,6 +16,7 @@ module Gitlab
       attr_reader :auth_header
 
       def initialize(auth_header)
+        @include_disabled = false
         @auth_header = auth_header
       end
 
@@ -28,6 +29,10 @@ module Gitlab
         end
       end
 
+      def include_disabled!
+        @include_disabled = true
+      end
+
       # Check if set of attributes match against attributes decoded from JWT
       #
       # @param [Hash] attributes to be matched against JWT
@@ -85,12 +90,21 @@ module Gitlab
 
       # rubocop: disable CodeReuse/ActiveRecord
       def hmac_secret(access_key)
-        @hmac_secret ||= begin
-                           geo_node = GeoNode.find_by(access_key: access_key, enabled: true)
-                           geo_node&.secret_access_key
-                         end
-      end
+        node_params = { access_key: access_key }
 
+        # By default, we fail authorization for requests from disabled nodes because
+        # it is a convenient place to block nearly all requests from disabled
+        # secondaries. The `include_disabled` option can safely override this
+        # check for `enabled`.
+        #
+        # A request is authorized if the access key in the Authorization header
+        # matches the access key of the requesting node, **and** the decoded data
+        # matches the requested resource.
+        node_params[:enabled] = true unless @include_disabled
+
+        @geo_node ||= GeoNode.find_by(node_params)
+        @geo_node&.secret_access_key
+      end
       # rubocop: enable CodeReuse/ActiveRecord
 
       def decode_auth_header

ee/spec/lib/ee/api/helpers_spec.rb

@@ -65,10 +65,12 @@ describe EE::API::Helpers do
   end
 
   describe '#authenticate_by_gitlab_geo_node_token!' do
+    let(:invalid_geo_auth_header) { "#{::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE}...Test" }
+
     it 'rescues from ::Gitlab::Geo::InvalidDecryptionKeyError' do
       expect_any_instance_of(::Gitlab::Geo::JwtRequestDecoder).to receive(:decode) { raise ::Gitlab::Geo::InvalidDecryptionKeyError }
 
-      header 'Authorization', 'test'
+      header 'Authorization', invalid_geo_auth_header
       get 'protected', params: { current_user: 'test' }
 
       expect(Gitlab::Json.parse(last_response.body)).to eq({ 'message' => 'Gitlab::Geo::InvalidDecryptionKeyError' })
@@ -77,7 +79,7 @@ describe EE::API::Helpers do
     it 'rescues from ::Gitlab::Geo::InvalidSignatureTimeError' do
       allow_any_instance_of(::Gitlab::Geo::JwtRequestDecoder).to receive(:decode) { raise ::Gitlab::Geo::InvalidSignatureTimeError }
 
-      header 'Authorization', 'test'
+      header 'Authorization', invalid_geo_auth_header
       get 'protected', params: { current_user: 'test' }
 
       expect(Gitlab::Json.parse(last_response.body)).to eq({ 'message' => 'Gitlab::Geo::InvalidSignatureTimeError' })

ee/spec/lib/gitlab/geo/jwt_request_decoder_spec.rb

@@ -26,6 +26,14 @@ describe Gitlab::Geo::JwtRequestDecoder do
       expect(subject.decode).to be_nil
     end
 
+    it 'decodes when node is disabled if `include_disabled!` is called first' do
+      primary_node.update_attribute(:enabled, false)
+
+      subject.include_disabled!
+
+      expect(subject.decode).to eq(data)
+    end
+
     it 'fails to decode with wrong key' do
       data = request.headers['Authorization']
 

ee/spec/models/geo_node_spec.rb

@@ -552,6 +552,12 @@ describe GeoNode, :request_store, :geo, type: :model do
     end
   end
 
+  describe '#node_api_url' do
+    it 'returns an api url based on the node uri and provided node id' do
+      expect(new_primary_node.node_api_url(new_node)).to eq("https://localhost:3000/gitlab/api/#{api_version}/geo_nodes/#{new_node.id}")
+    end
+  end
+
   describe '#snapshot_url' do
     let(:project) { create(:project) }
     let(:snapshot_url) { "https://localhost:3000/gitlab/api/#{api_version}/projects/#{project.id}/snapshot" }

ee/spec/requests/api/geo_nodes_spec.rb

@@ -280,6 +280,58 @@ describe API::GeoNodes, :request_store, :geo_fdw, :prometheus, api: true do
 
       expect(response).to have_gitlab_http_status(:bad_request)
     end
+
+    context 'auth with geo node token' do
+      let(:geo_base_request) { Gitlab::Geo::BaseRequest.new(scope: ::Gitlab::Geo::API_SCOPE) }
+
+      before do
+        stub_current_geo_node(primary)
+        allow(geo_base_request).to receive(:requesting_node) { secondary }
+      end
+
+      it 'enables the secondary node' do
+        secondary.update(enabled: false)
+
+        put api("/geo_nodes/#{secondary.id}"), params: { enabled: true }, headers: geo_base_request.headers
+
+        expect(response).to have_gitlab_http_status(:ok)
+        expect(secondary.reload).to be_enabled
+      end
+
+      it 'disables the secondary node' do
+        secondary.update(enabled: true)
+
+        put api("/geo_nodes/#{secondary.id}"), params: { enabled: false }, headers: geo_base_request.headers
+
+        expect(response).to have_gitlab_http_status(:ok)
+        expect(secondary.reload).not_to be_enabled
+      end
+
+      it 'returns bad request if you try to update the primary' do
+        put api("/geo_nodes/#{primary.id}"), params: { enabled: false }, headers: geo_base_request.headers
+
+        expect(response).to have_gitlab_http_status(:bad_request)
+        expect(primary.reload).to be_enabled
+      end
+
+      it 'responds with 401 when IP is not allowed' do
+        stub_application_setting(geo_node_allowed_ips: '192.34.34.34')
+
+        put api("/geo_nodes/#{secondary.id}"), params: {}, headers: geo_base_request.headers
+
+        expect(response).to have_gitlab_http_status(:unauthorized)
+      end
+
+      it 'responds 401 if auth header is bad' do
+        allow_next_instance_of(Gitlab::Geo::JwtRequestDecoder) do |instance|
+          allow(instance).to receive(:decode).and_raise(Gitlab::Geo::InvalidDecryptionKeyError)
+        end
+
+        put api("/geo_nodes/#{secondary.id}"), params: {}, headers: geo_base_request.headers
+
+        expect(response).to have_gitlab_http_status(:unauthorized)
+      end
+    end
   end
 
   describe 'DELETE /geo_nodes/:id' do

ee/spec/requests/api/geo_spec.rb

@@ -14,6 +14,9 @@ describe API::Geo do
   let(:geo_token_header) do
     { 'X-Gitlab-Token' => secondary_node.system_hook.token }
   end
+  let(:invalid_geo_auth_header) do
+    { Authorization: "#{::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE}...Test" }
+  end
 
   let(:not_found_req_header) do
     Gitlab::Geo::TransferRequest.new(transfer.request_data.merge(file_id: 100000)).headers
@@ -75,7 +78,7 @@ describe API::Geo do
 
     context 'invalid requests' do
       it 'responds with 401 with invalid auth header' do
-        get api("/geo/retrieve/#{replicator.replicable_name}/#{resource.id}"), headers: { Authorization: 'Test' }
+        get api("/geo/retrieve/#{replicator.replicable_name}/#{resource.id}"), headers: invalid_geo_auth_header
 
         expect(response).to have_gitlab_http_status(:unauthorized)
       end
@@ -135,7 +138,7 @@ describe API::Geo do
 
       it 'responds with 401 when an invalid auth header is provided' do
         path = File.join(api_path, resource.id.to_s)
-        get api(path), headers: { Authorization: 'Test' }
+        get api(path), headers: invalid_geo_auth_header
 
         expect(response).to have_gitlab_http_status(:unauthorized)
       end
@@ -284,7 +287,7 @@ describe API::Geo do
         end
 
         it 'responds with 401 when an invalid auth header is provided' do
-          get api("/geo/transfers/lfs/#{resource.id}"), headers: { Authorization: 'Test' }
+          get api("/geo/transfers/lfs/#{resource.id}"), headers: invalid_geo_auth_header
 
           expect(response).to have_gitlab_http_status(:unauthorized)
         end
@@ -372,7 +375,7 @@ describe API::Geo do
     subject(:request) { post api('/geo/status'), params: data, headers: geo_base_request.headers }
 
     it 'responds with 401 with invalid auth header' do
-      post api('/geo/status'), headers: { Authorization: 'Test' }
+      post api('/geo/status'), headers: invalid_geo_auth_header
 
       expect(response).to have_gitlab_http_status(:unauthorized)
     end