LG-8604 - Optimize s3 by streaming data w/ feature flag

app/controllers/api/irs_attempts_api_controller.rb

@@ -7,6 +7,7 @@
 module Api
   class IrsAttemptsApiController < ApplicationController
     include RenderConditionConcern
+    include ActionController::Live
 
     check_or_render_not_found -> { IdentityConfig.store.irs_attempt_api_enabled }
 
@@ -25,16 +26,11 @@ def create
             log_file_record = IrsAttemptApiLogFile.find_by(
               requested_time: timestamp_key(key: timestamp),
             )
-
             headers['X-Payload-Key'] = log_file_record.encrypted_key
             headers['X-Payload-IV'] = log_file_record.iv
 
-            requested_data = s3_helper.s3_client.get_object(
-              bucket: s3_helper.attempts_bucket_name,
-              key: log_file_record.filename,
-            )
+            serve_s3_response(log_file_record: log_file_record)
 
-            send_data requested_data.body.read, disposition: "filename=#{log_file_record.filename}"
           else
             render json: { status: :not_found, description: 'File not found for Timestamp' },
                    status: :not_found
@@ -62,6 +58,52 @@ def create
 
     private
 
+    def buffer_range(current_buffer_index:, buffer_size:, file_size:)
+      buffer_end = [current_buffer_index + buffer_size, file_size].min
+      "bytes=#{current_buffer_index}-#{buffer_end}"
+    end
+
+    def serve_s3_response(log_file_record:)
+      if IdentityConfig.store.irs_attempt_api_aws_s3_stream_enabled
+        response = s3_helper.s3_client.head_object(
+          bucket: s3_helper.attempts_bucket_name,
+          key: log_file_record.filename,
+        )
+
+        requested_data_size = response.content_length
+
+        buffer_index = 0
+        buffer_size = IdentityConfig.store.irs_attempt_api_aws_s3_stream_buffer_size
+
+        send_stream(
+          type: response.content_type,
+          filename: log_file_record.filename,
+        ) do |stream|
+          while buffer_index < requested_data_size
+            requested_data = s3_helper.s3_client.get_object(
+              bucket: s3_helper.attempts_bucket_name,
+              key: log_file_record.filename,
+              range: buffer_range(
+                current_buffer_index: buffer_index,
+                buffer_size: buffer_size,
+                file_size: requested_data_size,
+              ),
+            )
+            buffer_index += buffer_size + 1
+            stream.write(requested_data.body.read)
+          end
+        end
+      else
+        requested_data = s3_helper.s3_client.get_object(
+          bucket: s3_helper.attempts_bucket_name,
+          key: log_file_record.filename,
+        )
+
+        send_data requested_data.body.read,
+                  disposition: "filename=#{log_file_record.filename}"
+      end
+    end
+
     def authenticate_client
       bearer, csp_id, token = request.authorization&.split(' ', 3)
       if bearer != 'Bearer' || !valid_auth_tokens.include?(token) ||

config/application.yml.default

@@ -151,6 +151,8 @@ irs_attempt_api_auth_tokens: ''
 irs_attempt_api_csp_id: 'LOGIN.gov'
 irs_attempt_api_enabled: false
 irs_attempt_api_aws_s3_enabled: false
+irs_attempt_api_aws_s3_stream_enabled: false
+irs_attempt_api_aws_s3_stream_buffer_size: 16_777_216
 irs_attempt_api_event_ttl_seconds: 86400
 irs_attempt_api_event_count_default: 1000
 irs_attempt_api_event_count_max: 10000

lib/identity_config.rb

@@ -231,6 +231,8 @@ def self.build_store(config_map)
     config.add(:irs_attempt_api_csp_id)
     config.add(:irs_attempt_api_enabled, type: :boolean)
     config.add(:irs_attempt_api_aws_s3_enabled, type: :boolean)
+    config.add(:irs_attempt_api_aws_s3_stream_enabled, type: :boolean)
+    config.add(:irs_attempt_api_aws_s3_stream_buffer_size, type: :integer)
     config.add(:irs_attempt_api_event_ttl_seconds, type: :integer)
     config.add(:irs_attempt_api_event_count_default, type: :integer)
     config.add(:irs_attempt_api_event_count_max, type: :integer)

spec/controllers/api/irs_attempts_api_controller_spec.rb

@@ -50,34 +50,6 @@
       }
     end
 
-    context 'with timestamp problems' do
-      it 'returns unprocessable_entity with no timestamp' do
-        post :create, params: { timestamp: nil }
-
-        expect(response.status).to eq(422)
-      end
-
-      it 'returns unprocessable_entity with invalid timestamp' do
-        post :create, params: { timestamp: 'INVALID*TIME' }
-
-        expect(response.status).to eq(422)
-      end
-    end
-
-    context 'with aws_s3 disabled' do
-      let(:timestamp) { '2022-11-08T18:00:00.000Z' }
-      it 'should bypass s3 retrieval' do
-        expect_any_instance_of(Aws::S3::Client).not_to receive(:get_object)
-
-        post :create, params: { timestamp: timestamp }
-
-        expect(response).to be_ok
-        expect(Base64.strict_decode64(response.headers['X-Payload-IV'])).to be_present
-        expect(Base64.strict_decode64(response.headers['X-Payload-Key'])).to be_present
-        expect(Base64.strict_decode64(response.body)).to be_present
-      end
-    end
-
     context 'with aws_s3 enabled' do
       let(:wrong_time) { time - 1.year }
       let(:timestamp) { '2022-11-08T18:00:00.000Z' }
@@ -107,6 +79,72 @@
         expect(Base64.strict_decode64(response.headers['X-Payload-Key'])).to be_present
         expect(response.body).to eq(test_object)
       end
+
+      context 'with aws_s3_stream enabled' do
+        let(:test_object) { '{test: "1234567890 12345"}' }
+        before do
+          allow(IdentityConfig.store).to receive(:irs_attempt_api_aws_s3_stream_enabled).
+            and_return(true)
+          allow(IdentityConfig.store).to receive(:irs_attempt_api_aws_s3_stream_buffer_size).
+            and_return(10)
+
+          Aws.config[:s3] = {
+            stub_responses: {
+              head_object: { content_length: test_object.bytesize },
+              get_object: proc do |context|
+                range_string = context.params[:range]
+                _, byte_string = range_string.split('=')
+                start_byte, _ = byte_string.split('-')
+                { body: test_object.byteslice(
+                  start_byte.to_i,
+                  IdentityConfig.store.irs_attempt_api_aws_s3_stream_buffer_size + 1,
+                ) }
+              end,
+            },
+          }
+        end
+
+        it 'should render data streamed from s3 correctly' do
+          post :create, params: { timestamp: time.iso8601 }
+
+          expect(response).to be_ok
+          expect(Base64.strict_decode64(response.headers['X-Payload-IV'])).to be_present
+          expect(Base64.strict_decode64(response.headers['X-Payload-Key'])).to be_present
+          expect(response.content_type).to eq('application/octet-stream')
+          expect(response['Content-Disposition']).
+            to eq("attachment; filename=\"test_filename\"; filename*=UTF-8''test_filename")
+
+          expect(response.stream.body).to eq(test_object)
+        end
+      end
+    end
+
+    context 'with timestamp problems' do
+      it 'returns unprocessable_entity when given no timestamp' do
+        post :create, params: { timestamp: nil }
+
+        expect(response.status).to eq(422)
+      end
+
+      it 'returns unprocessable_entity when timestamp is invalid' do
+        post :create, params: { timestamp: 'INVALID*TIME' }
+
+        expect(response.status).to eq(422)
+      end
+    end
+
+    context 'with aws_s3 disabled' do
+      let(:timestamp) { '2022-11-08T18:00:00.000Z' }
+      it 'should bypass s3 retrieval' do
+        expect_any_instance_of(Aws::S3::Client).not_to receive(:get_object)
+
+        post :create, params: { timestamp: timestamp }
+
+        expect(response).to be_ok
+        expect(Base64.strict_decode64(response.headers['X-Payload-IV'])).to be_present
+        expect(Base64.strict_decode64(response.headers['X-Payload-Key'])).to be_present
+        expect(Base64.strict_decode64(response.body)).to be_present
+      end
     end
 
     context 'with CSRF protection enabled' do
@@ -144,17 +182,17 @@
       expect(response.status).to eq(404)
     end
 
-    it 'returns an error without required timestamp parameter' do
+    it 'returns an error when required timestamp parameter is missing' do
       post :create, params: {}
       expect(response.status).to eq 422
     end
 
-    it 'returns an error with empty timestamp parameter' do
+    it 'returns an error when timestamp parameter is empty' do
       post :create, params: { timestamp: '' }
       expect(response.status).to eq 422
     end
 
-    it 'returns an error with invalid timestamp parameter' do
+    it 'returns an error when timestamp parameter is invalid' do
       post :create, params: { timestamp: 'abc' }
       expect(response.status).to eq 422