Check integrity after uploads

7 files changed, 34 insertions, 9 deletions

diff --git a/lib/active_storage/blob.rb b/lib/active_storage/blob.rb
index b10dc2c771..26c116712b 100644
--- a/lib/active_storage/blob.rb
+++ b/lib/active_storage/blob.rb
@@ -45,7 +45,7 @@ class ActiveStorage::Blob < ActiveRecord::Base
     self.checksum  = compute_checksum_in_chunks(io)
     self.byte_size = io.size
 
-    service.upload(key, io)
+    service.upload(key, io, checksum: checksum)
   end
 
   def download
diff --git a/lib/active_storage/service.rb b/lib/active_storage/service.rb
index 0d0ebf6010..9aab654d80 100644
--- a/lib/active_storage/service.rb
+++ b/lib/active_storage/service.rb
@@ -1,5 +1,7 @@
 # Abstract class serving as an interface for concrete services.
 class ActiveStorage::Service
+  class ActiveStorage::IntegrityError < StandardError; end
+
   def self.configure(service, **options)
     begin
       require "active_storage/service/#{service.to_s.downcase}_service"
@@ -10,7 +12,7 @@ class ActiveStorage::Service
   end
 
 
-  def upload(key, io)
+  def upload(key, io, checksum: nil)
     raise NotImplementedError
   end
 
diff --git a/lib/active_storage/service/disk_service.rb b/lib/active_storage/service/disk_service.rb
index 668a837663..6164caf86c 100644
--- a/lib/active_storage/service/disk_service.rb
+++ b/lib/active_storage/service/disk_service.rb
@@ -9,12 +9,14 @@ class ActiveStorage::Service::DiskService < ActiveStorage::Service
     @root = root
   end
 
-  def upload(key, io)
+  def upload(key, io, checksum: nil)
     File.open(make_path_for(key), "wb") do |file|
       while chunk = io.read(64.kilobytes)
         file.write(chunk)
       end
     end
+
+    ensure_integrity_of(key, checksum) if checksum
   end
 
   def download(key)
@@ -59,4 +61,10 @@ class ActiveStorage::Service::DiskService < ActiveStorage::Service
     def make_path_for(key)
       path_for(key).tap { |path| FileUtils.mkdir_p File.dirname(path) }
     end
+
+    def ensure_integrity_of(key, checksum)
+      unless Digest::MD5.file(path_for(key)).base64digest == checksum
+        raise ActiveStorage::IntegrityError
+      end
+    end
 end
diff --git a/lib/active_storage/service/gcs_service.rb b/lib/active_storage/service/gcs_service.rb
index c725afb35c..a558791d89 100644
--- a/lib/active_storage/service/gcs_service.rb
+++ b/lib/active_storage/service/gcs_service.rb
@@ -9,7 +9,8 @@ class ActiveStorage::Service::GCSService < ActiveStorage::Service
     @bucket = @client.bucket(bucket)
   end
 
-  def upload(key, io)
+  def upload(key, io, checksum: nil)
+    # FIXME: Ensure integrity by sending the checksum for service side verification
     bucket.create_file(io, key)
   end
 
diff --git a/lib/active_storage/service/mirror_service.rb b/lib/active_storage/service/mirror_service.rb
index 0d37ad96a3..59ad3fc472 100644
--- a/lib/active_storage/service/mirror_service.rb
+++ b/lib/active_storage/service/mirror_service.rb
@@ -9,9 +9,9 @@ class ActiveStorage::Service::MirrorService < ActiveStorage::Service
     @services = services
   end
 
-  def upload(key, io)
+  def upload(key, io, checksum: nil)
     services.collect do |service|
-      service.upload key, io
+      service.upload key, io, checksum: checksum
       io.rewind
     end
   end
diff --git a/lib/active_storage/service/s3_service.rb b/lib/active_storage/service/s3_service.rb
index 746a636912..fd8ef6e9a6 100644
--- a/lib/active_storage/service/s3_service.rb
+++ b/lib/active_storage/service/s3_service.rb
@@ -9,7 +9,8 @@ class ActiveStorage::Service::S3Service < ActiveStorage::Service
     @bucket = @client.bucket(bucket)
   end
 
-  def upload(key, io)
+  def upload(key, io, checksum: nil)
+    # FIXME: Ensure integrity by sending the checksum for service side verification
     object_for(key).put(body: io)
   end
 
diff --git a/test/service/shared_service_tests.rb b/test/service/shared_service_tests.rb
index 3676272e27..b4c888e77c 100644
--- a/test/service/shared_service_tests.rb
+++ b/test/service/shared_service_tests.rb
@@ -26,11 +26,11 @@ module ActiveStorage::Service::SharedServiceTests
       FIXTURE_FILE.rewind
     end
 
-    test "uploading" do
+    test "uploading with integrity" do
       begin
         key  = SecureRandom.base58(24)
         data = "Something else entirely!"
-        @service.upload(key, StringIO.new(data))
+        @service.upload(key, StringIO.new(data), checksum: Digest::MD5.base64digest(data))
 
         assert_equal data, @service.download(key)
       ensure
@@ -38,6 +38,19 @@ module ActiveStorage::Service::SharedServiceTests
       end
     end
 
+    test "upload without integrity" do
+      begin
+        key  = SecureRandom.base58(24)
+        data = "Something else entirely!"
+
+        assert_raises(ActiveStorage::IntegrityError) do
+          @service.upload(key, StringIO.new(data), checksum: "BAD_CHECKSUM")
+        end
+      ensure
+        @service.delete key
+      end
+    end
+
     test "downloading" do
       assert_equal FIXTURE_FILE.read, @service.download(FIXTURE_KEY)
     end