Add range checking

4 files changed, 57 insertions, 11 deletions

diff --git a/app/models/ci/build.rb b/app/models/ci/build.rb
index 16abdb143bb..9eee4b423cc 100644
--- a/app/models/ci/build.rb
+++ b/app/models/ci/build.rb
@@ -230,6 +230,14 @@ module Ci
       end
     end
 
+    def trace_length
+      unless trace.present?
+        0
+      else
+        trace.length
+      end
+    end
+
     def trace=(trace)
       recreate_trace_dir
       File.write(path_to_trace, trace)
diff --git a/lib/ci/api/api.rb b/lib/ci/api/api.rb
index 4e85d2c3c74..353c4ddebf8 100644
--- a/lib/ci/api/api.rb
+++ b/lib/ci/api/api.rb
@@ -23,6 +23,8 @@ module Ci
         rack_response({ 'message' => '500 Internal Server Error' }, 500)
       end
 
+      content_type :txt,  'text/plain'
+      content_type :json, 'application/json'
       format :json
 
       helpers ::Ci::API::Helpers
diff --git a/lib/ci/api/builds.rb b/lib/ci/api/builds.rb
index 61e15675535..decee3739c8 100644
--- a/lib/ci/api/builds.rb
+++ b/lib/ci/api/builds.rb
@@ -51,12 +51,24 @@ module Ci
         end
 
         patch ":id/trace.txt" do
-          authenticate_runner!
-          update_runner_last_contact
-          build = Ci::Build.where(runner_id: current_runner.id).running.find(params[:id])
+          build = Ci::Build.find_by_id(params[:id])
+          not_found! unless build
+          authenticate_build_token!(build)
           forbidden!('Build has been erased!') if build.erased?
 
-          build.append_trace(params[:trace_part])
+          error!('400 Missing header Content-Range', 400) unless request.headers.has_key?('Content-Range')
+          content_range = request.headers['Content-Range']
+          content_range = content_range.split('-')
+
+          unless build.trace_length == content_range[0].to_i
+            return error!('416 Range Not Satisfiable', 416, { 'Range' => "0-#{build.trace_length}" })
+          end
+
+          build.append_trace(request.body.read)
+
+          status 202
+          header 'Build-Status', build.status
+          header 'Range', "0-#{build.trace_length}"
         end
 
         # Authorize artifacts uploading for build - Runners only
diff --git a/spec/requests/ci/api/builds_spec.rb b/spec/requests/ci/api/builds_spec.rb
index 4da07fcd116..eb8c476c882 100644
--- a/spec/requests/ci/api/builds_spec.rb
+++ b/spec/requests/ci/api/builds_spec.rb
@@ -158,23 +158,47 @@ describe Ci::API::API do
 
     describe 'PATCH /builds/:id/trace.txt' do
       let(:build) { create(:ci_build, :trace, runner_id: runner.id) }
+      let(:headers) { { Ci::API::Helpers::BUILD_TOKEN_HEADER => build.token, 'Content-Type' => 'text/plain' } }
+      let(:headers_with_range) { headers.merge({ 'Content-Range' => '11-20' }) }
 
       before do
         build.run!
-        patch ci_api("/builds/#{build.id}/trace.txt"), trace_part: ' appended', token: runner.token
+        patch ci_api("/builds/#{build.id}/trace.txt"), ' appended', headers_with_range
       end
 
-      it 'should append trace part to the trace' do
-        expect(response.status).to eq 200
-        expect(build.reload.trace).to eq 'BUILD TRACE appended'
+      context 'when request is valid' do
+        it { expect(response.status).to eq 202 }
+        it { expect(build.reload.trace).to eq 'BUILD TRACE appended' }
+        it { expect(response.header).to have_key 'Range' }
+        it { expect(response.header).to have_key 'Build-Status' }
+      end
+
+      context 'when content-range start is too big' do
+        let(:headers_with_range) { headers.merge({ 'Content-Range' => '15-20' }) }
+
+        it { expect(response.status).to eq 416 }
+        it { expect(response.header).to have_key 'Range' }
+        it { expect(response.header['Range']).to eq '0-11' }
+      end
+
+      context 'when content-range start is too small' do
+        let(:headers_with_range) { headers.merge({ 'Content-Range' => '8-20' }) }
+
+        it { expect(response.status).to eq 416 }
+        it { expect(response.header).to have_key 'Range' }
+        it { expect(response.header['Range']).to eq '0-11' }
+      end
+
+      context 'when Content-Range header is missing' do
+        let(:headers_with_range) { headers.merge({}) }
+
+        it { expect(response.status).to eq 400 }
       end
 
       context 'when build has been erased' do
         let(:build) { create(:ci_build, runner_id: runner.id, erased_at: Time.now) }
 
-        it 'should respond with forbidden' do
-          expect(response.status).to eq 403
-        end
+        it { expect(response.status).to eq 403 }
       end
     end