Refactor Geo signing data logic

Extract specific encoding/decoding and signing data logic
in a separate Gitlab::Geo::SignedData class to decouple
from request authorization logic.

This can be then used separately to generate JWT signed
tokens from an existing node that can encode arbitrary
data as needed.

ee/app/models/geo_node.rb

@@ -53,6 +53,7 @@ class GeoNode < ApplicationRecord
   scope :secondary_nodes, -> { where(primary: false) }
   scope :name_in, -> (names) { where(name: names) }
   scope :ordered, -> { order(:id) }
+  scope :enabled, -> { where(enabled: true) }
 
   attr_encrypted :secret_access_key,
                  key: Settings.attr_encrypted_db_key_base_32,

ee/lib/gitlab/geo/base_request.rb

@@ -29,14 +29,9 @@ module Gitlab
       private
 
       def geo_auth_token(message)
-        geo_node = requesting_node
-        raise GeoNodeNotFoundError unless geo_node
+        signed_data = Gitlab::Geo::SignedData.new(geo_node: requesting_node).sign_and_encode_data(message)
 
-        token = JSONWebToken::HMACToken.new(geo_node.secret_access_key)
-        token.expire_time = Time.now + expiration_time
-        token[:data] = message.to_json
-
-        "#{GITLAB_GEO_AUTH_TOKEN_TYPE} #{geo_node.access_key}:#{token.encoded}"
+        "#{GITLAB_GEO_AUTH_TOKEN_TYPE} #{signed_data}"
       end
 
       def requesting_node

ee/lib/gitlab/geo/jwt_request_decoder.rb

@@ -6,8 +6,6 @@ module Gitlab
       include LogHelpers
       include ::Gitlab::Utils::StrongMemoize
 
-      IAT_LEEWAY = 60.seconds.to_i
-
       def self.geo_auth_attempt?(header)
         token_type, = header&.split(' ', 2)
         token_type == ::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE
@@ -54,60 +52,16 @@ module Gitlab
         # For example:
         # JWT payload = { "data": { "oid": "12345" }, iat: 123456 }
         #
-        begin
-          data = decode_auth_header
+        signed_data = parse_auth_header
+
+        Gitlab::Geo::SignedData.new(include_disabled_nodes: @include_disabled).decode_data(signed_data)
       rescue OpenSSL::Cipher::CipherError
         message = 'Error decrypting the Geo secret from the database. Check that the primary and secondary have the same db_key_base.'
         log_error(message)
         raise InvalidDecryptionKeyError, message
       end
 
-        return unless data.present?
-
-        secret, encoded_message = data
-
-        begin
-          decoded = JWT.decode(
-            encoded_message,
-            secret,
-            true,
-            { leeway: IAT_LEEWAY, verify_iat: true, algorithm: 'HS256' }
-          )
-
-          message = decoded.first
-          data = Gitlab::Json.parse(message['data']) if message
-          data&.deep_symbolize_keys!
-          data
-        rescue JWT::ImmatureSignature, JWT::ExpiredSignature
-          message = "Signature not within leeway of #{IAT_LEEWAY} seconds. Check your system clocks!"
-          log_error(message)
-          raise InvalidSignatureTimeError, message
-        rescue JWT::DecodeError => e
-          log_error("Error decoding Geo request: #{e}")
-          nil
-        end
-      end
-
-      # rubocop: disable CodeReuse/ActiveRecord
-      def hmac_secret(access_key)
-        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
+      def parse_auth_header
         return unless auth_header.present?
 
         tokens = auth_header.split(' ')
@@ -115,18 +69,7 @@ module Gitlab
         return unless tokens.count == 2
         return unless tokens[0] == Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE
 
-        # Split at the first occurrence of a colon
-        geo_tokens = tokens[1].split(':', 2)
-
-        return unless geo_tokens.count == 2
-
-        access_key = geo_tokens[0]
-        encoded_message = geo_tokens[1]
-        secret = hmac_secret(access_key)
-
-        return unless secret.present?
-
-        [secret, encoded_message]
+        tokens[1]
       end
     end
   end

ee/lib/gitlab/geo/signed_data.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Geo
+    class SignedData
+      include LogHelpers
+
+      VALIDITY_PERIOD = 1.minute
+      IAT_LEEWAY = 60.seconds.to_i
+
+      attr_reader :geo_node
+
+      def initialize(geo_node: nil, validity_period: VALIDITY_PERIOD, include_disabled_nodes: false)
+        @geo_node = geo_node
+        @validity_period = validity_period
+        @include_disabled_nodes = include_disabled_nodes
+      end
+
+      def sign_and_encode_data(data)
+        raise ::Gitlab::Geo::GeoNodeNotFoundError unless @geo_node
+
+        token = JSONWebToken::HMACToken.new(@geo_node.secret_access_key)
+        token.expire_time = Time.zone.now + @validity_period
+        token[:data] = data.to_json
+
+        "#{@geo_node.access_key}:#{token.encoded}"
+      end
+
+      def decode_data(signed_data)
+        return unless signed_data.present?
+
+        parse_data(signed_data)
+
+        return unless @secret_access_key.present? && @encoded_message.present?
+
+        decoded = JWT.decode(
+          @encoded_message,
+          @secret_access_key,
+          true,
+          { leeway: IAT_LEEWAY, verify_iat: true, algorithm: 'HS256' }
+        )
+
+        message = decoded.first
+        data = Gitlab::Json.parse(message['data']) if message
+        data&.deep_symbolize_keys!
+        data
+      rescue JWT::ImmatureSignature, JWT::ExpiredSignature
+        message = "Signature not within leeway of #{IAT_LEEWAY} seconds. Check your system clocks!"
+        log_error(message)
+        raise InvalidSignatureTimeError, message
+      rescue JWT::DecodeError => e
+        log_error("Error decoding Geo request: #{e}")
+        nil
+      end
+
+      private
+
+      def parse_data(signed_data)
+        geo_tokens = signed_data.split(':', 2)
+
+        return unless geo_tokens.count == 2
+
+        @access_key = geo_tokens[0]
+        @encoded_message = geo_tokens[1]
+        @secret_access_key = hmac_secret
+      end
+
+      def hmac_secret
+        # 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_nodes` 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.
+        scoped_geo_nodes = if @include_disabled_nodes
+                             GeoNode.all
+                           else
+                             GeoNode.enabled
+                           end
+
+        @geo_node ||= scoped_geo_nodes.find_by_access_key(@access_key)
+        @geo_node&.secret_access_key
+      end
+    end
+  end
+end

ee/lib/system_check/geo/clocks_synchronization_check.rb

@@ -6,7 +6,7 @@ module SystemCheck
       set_name 'Machine clock is synchronized'
 
       def check?
-        Net::NTP.get.offset.abs < Gitlab::Geo::JwtRequestDecoder::IAT_LEEWAY
+        Net::NTP.get.offset.abs < Gitlab::Geo::SignedData::IAT_LEEWAY
       end
 
       def show_error

ee/spec/lib/ee/gitlab/auth/auth_finders_spec.rb

@@ -65,7 +65,7 @@ RSpec.describe EE::Gitlab::Auth::AuthFinders do
 
       it 'does not authenticate with invalid decryption key error' do
         allow_next_instance_of(::Gitlab::Geo::JwtRequestDecoder) do |instance|
-          expect(instance).to receive(:decode_auth_header).and_raise(Gitlab::Geo::InvalidDecryptionKeyError)
+          expect(instance).to receive(:decode).and_raise(Gitlab::Geo::InvalidDecryptionKeyError)
         end
 
         expect { subject }.to raise_error(::Gitlab::Auth::UnauthorizedError)

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

@@ -60,12 +60,23 @@ RSpec.describe Gitlab::Geo::JwtRequestDecoder do
       travel_to(2.minutes.ago) { expect { subject.decode }.to raise_error(Gitlab::Geo::InvalidSignatureTimeError) }
     end
 
-    it 'raises invalid decryption key error' do
-      allow_next_instance_of(described_class) do |instance|
-        allow(instance).to receive(:decode_auth_header).and_raise(Gitlab::Geo::InvalidDecryptionKeyError)
+    context 'surfaces raised errors' do
+      where(:raised_error, :expected_error) do
+        [
+          [Gitlab::Geo::InvalidDecryptionKeyError, Gitlab::Geo::InvalidDecryptionKeyError],
+          [OpenSSL::Cipher::CipherError, Gitlab::Geo::InvalidDecryptionKeyError]
+        ]
       end
 
-      expect { subject.decode }.to raise_error(Gitlab::Geo::InvalidDecryptionKeyError)
+      with_them do
+        it 'raises expected error' do
+          allow_next_instance_of(Gitlab::Geo::SignedData) do |instance|
+            allow(instance).to receive(:decode_data).and_raise(raised_error)
+          end
+
+          expect { subject.decode }.to raise_error(expected_error)
+        end
+      end
     end
   end
 

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Geo::SignedData do
+  using RSpec::Parameterized::TableSyntax
+  include EE::GeoHelpers
+
+  let_it_be_with_reload(:geo_node) { create(:geo_node) }
+
+  let(:validity_period) { 42.minutes }
+  let(:include_disabled_nodes) { false }
+  let(:klass_args) { {} }
+  let(:data) { { input: 123, other_input: 'string value' } }
+  let(:signed_data) { described_class.new(geo_node: geo_node).sign_and_encode_data(data) }
+
+  before do
+    stub_current_geo_node(geo_node)
+  end
+
+  describe '#sign_and_encode_data' do
+    subject(:returned_data) { described_class.new(**klass_args).sign_and_encode_data(data) }
+
+    let(:klass_args) { { geo_node: geo_node } }
+    let(:parsed_access_key) { returned_data.split(':').first }
+    let(:jwt) { JWT.decode(returned_data.split(':').second, geo_node.secret_access_key) }
+    let(:decoded_data) { Gitlab::Json.parse(jwt.first['data']) }
+
+    context 'when data is not set' do
+      let(:data) { nil }
+
+      it 'does not set the data attribute' do
+        expect(decoded_data).to be_nil
+      end
+    end
+
+    context 'when geo_node is not set' do
+      let(:geo_node) { nil }
+
+      it 'raises a GeoNodeNotFoundError error' do
+        expect { subject }.to raise_error(::Gitlab::Geo::GeoNodeNotFoundError)
+      end
+    end
+
+    it 'formats the signed data properly' do
+      expect(parsed_access_key).to eq(geo_node.access_key)
+
+      decoded_data.deep_symbolize_keys!
+      expect(decoded_data).to eq(data)
+    end
+
+    it 'defaults to 1-minute expiration time' do
+      freeze_time do
+        expect(jwt.first['exp']).to eq((Time.zone.now + 1.minute).to_i)
+      end
+    end
+
+    context 'with custom validity period' do
+      let(:klass_args) { { geo_node: geo_node, validity_period: 42.minutes } }
+
+      it 'uses that expiration time' do
+        freeze_time do
+          expect(jwt.first['exp']).to eq((Time.zone.now + 42.minutes).to_i)
+        end
+      end
+    end
+  end
+
+  describe '#decode_data' do
+    subject(:returned_data) { described_class.new(**klass_args).decode_data(signed_data) }
+
+    it { is_expected.to eq(data) }
+
+    context 'when data is not set' do
+      let(:signed_data) { nil }
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'for disabled nodes' do
+      before(:all) do
+        geo_node.update_attribute(:enabled, false)
+      end
+
+      after(:all) do
+        geo_node.update_attribute(:enabled, true)
+      end
+
+      context 'fails to decode for disabled nodes by default' do
+        it { is_expected.to be_nil }
+      end
+
+      context 'when include_disabled_nodes is set to false' do
+        let(:klass_args) { { include_disabled_nodes: false } }
+
+        it { is_expected.to be_nil }
+      end
+
+      context 'when include_disabled_nodes is set to true' do
+        let(:klass_args) { { include_disabled_nodes: true } }
+
+        it { is_expected.to eq(data) }
+      end
+    end
+
+    context 'with the wrong key' do
+      before do
+        # ensure this generated before changing the key
+        signed_data
+
+        geo_node.update_attribute(:secret_access_key, 'invalid')
+      end
+
+      it { is_expected.to be_nil }
+    end
+
+    context 'time checks' do
+      before do
+        signed_data
+      end
+
+      it 'successfully decodes when clocks are off by IAT leeway' do
+        travel_to(30.seconds.ago) { expect(subject).to eq(data) }
+      end
+
+      it 'raises an error after expiring' do
+        travel_to(2.minutes.from_now) { expect { subject }.to raise_error(Gitlab::Geo::InvalidSignatureTimeError) }
+      end
+
+      it 'raises an error when clocks are not in sync' do
+        travel_to(2.minutes.ago) { expect { subject }.to raise_error(Gitlab::Geo::InvalidSignatureTimeError) }
+      end
+    end
+
+    context 'JWT raised errors' do
+      context 'surfaces expected errors' do
+        where(:raised_error, :expected_error) do
+          JWT::ImmatureSignature | Gitlab::Geo::InvalidSignatureTimeError
+          JWT::ExpiredSignature  | Gitlab::Geo::InvalidSignatureTimeError
+        end
+
+        with_them do
+          it 'raises expected error' do
+            expect(JWT).to receive(:decode).and_raise(raised_error)
+
+            expect { subject }.to raise_error(expected_error)
+          end
+        end
+      end
+
+      context 'for a decoding error' do
+        before do
+          allow(JWT).to receive(:decode).and_raise(JWT::DecodeError)
+        end
+
+        it { is_expected.to be_nil }
+      end
+    end
+  end
+end