tus-server 0.2.0 → 0.9.0

data/lib/tus/storage/gridfs.rb

@@ -1,72 +1,198 @@
 require "mongo"
-require "stringio"
-require "tempfile"
+
+require "tus/info"
+require "tus/errors"
+
+require "digest"
 
 module Tus
   module Storage
     class Gridfs
-      attr_reader :client, :prefix, :bucket
+      attr_reader :client, :prefix, :bucket, :chunk_size
 
-      def initialize(client:, prefix: "fs")
+      def initialize(client:, prefix: "fs", chunk_size: nil)
         @client = client
         @prefix = prefix
         @bucket = @client.database.fs(bucket_name: @prefix)
         @bucket.send(:ensure_indexes!)
+        @chunk_size = chunk_size
       end
 
-      def create_file(uid, metadata = {})
-        file = Mongo::Grid::File.new("", filename: uid, metadata: metadata)
+      def create_file(uid, info = {})
+        tus_info     = Tus::Info.new(info)
+        content_type = tus_info.metadata["content_type"]
+
+        file = Mongo::Grid::File.new("",
+          filename:     uid,
+          metadata:     {},
+          chunk_size:   chunk_size,
+          content_type: content_type,
+        )
+
         bucket.insert_one(file)
       end
 
-      def file_exists?(uid)
-        !!bucket.files_collection.find(filename: uid).first
-      end
+      def concatenate(uid, part_uids, info = {})
+        file_infos = bucket.files_collection.find(filename: {"$in" => part_uids}).to_a
+        file_infos.sort_by! { |file_info| part_uids.index(file_info[:filename]) }
+
+        if file_infos.count != part_uids.count
+          raise Tus::Error, "some parts for concatenation are missing"
+        end
+
+        chunk_sizes = file_infos.map { |file_info| file_info[:chunkSize] }
+        if chunk_sizes[0..-2].uniq.count > 1
+          raise Tus::Error, "some parts have different chunk sizes, so they cannot be concatenated"
+        end
+
+        if chunk_sizes.uniq != [chunk_sizes.last] && bucket.chunks_collection.find(files_id: file_infos.last[:_id]).count > 1
+          raise Tus::Error, "last part has different chunk size and is composed of more than one chunk"
+        end
+
+        length       = file_infos.inject(0) { |sum, file_info| sum + file_info[:length] }
+        chunk_size   = file_infos.first[:chunkSize]
+        tus_info     = Tus::Info.new(info)
+        content_type = tus_info.metadata["content_type"]
+
+        file = Mongo::Grid::File.new("",
+          filename:     uid,
+          metadata:     {},
+          chunk_size:   chunk_size,
+          length:       length,
+          content_type: content_type,
+        )
+
+        bucket.insert_one(file)
+
+        file_infos.inject(0) do |offset, file_info|
+          result = bucket.chunks_collection
+            .find(files_id: file_info[:_id])
+            .update_many("$set" => {files_id: file.id}, "$inc" => {n: offset})
+
+          offset += result.modified_count
+        end
 
-      def read_file(uid)
-        file = bucket.find_one(filename: uid)
-        file.data
+        bucket.files_collection.delete_many(filename: {"$in" => part_uids})
+
+        # server requires us to return the size of the concatenated file
+        length
       end
 
-      def patch_file(uid, content)
+      def patch_file(uid, io, info = {})
         file_info = bucket.files_collection.find(filename: uid).first
-        file_info["md5"] = Digest::MD5.new # hack around not able to update digest
-        file_info = Mongo::Grid::File::Info.new(file_info)
+        raise Tus::NotFound if file_info.nil?
+
+        file_info[:md5] = Digest::MD5.new # hack for `Chunk.split` updating MD5
+        file_info[:chunkSize] ||= io.size
+        file_info = Mongo::Grid::File::Info.new(Mongo::Options::Mapper.transform(file_info, Mongo::Grid::File::Info::MAPPINGS.invert))
+
+        tus_info = Tus::Info.new(info)
+        last_chunk = (tus_info.length && io.size == tus_info.remaining_length)
+
+        if io.size % file_info.chunk_size != 0 && !last_chunk
+          raise Tus::Error,
+            "Input has length #{io.size} but expected it to be a multiple of " \
+            "chunk size #{file_info.chunk_size} or for it to be the last chunk"
+        end
+
         offset = bucket.chunks_collection.find(files_id: file_info.id).count
-        chunks = Mongo::Grid::File::Chunk.split(content, file_info, offset)
+        chunks = Mongo::Grid::File::Chunk.split(io, file_info, offset)
+
         bucket.chunks_collection.insert_many(chunks)
-      end
+        chunks.each { |chunk| chunk.data.data.clear } # deallocate strings
 
-      def download_file(uid)
-        tempfile = Tempfile.new("tus", binmode: true)
-        tempfile.sync = true
-        bucket.download_to_stream_by_name(uid, tempfile)
-        tempfile.path
+        bucket.files_collection.find(filename: uid).update_one("$set" => {
+          length:     file_info.length + io.size,
+          uploadDate: Time.now.utc,
+          chunkSize:  file_info.chunk_size,
+        })
       end
 
-      def delete_file(uid)
+      def read_info(uid)
         file_info = bucket.files_collection.find(filename: uid).first
-        bucket.delete(file_info.fetch("_id")) if file_info
-      end
+        raise Tus::NotFound if file_info.nil?
 
-      def read_info(uid)
-        info = bucket.files_collection.find(filename: uid).first
-        info.fetch("metadata")
+        file_info.fetch("metadata")
       end
 
       def update_info(uid, info)
-        bucket.files_collection.find(filename: uid).update_one("$set" => {metadata: info})
+        bucket.files_collection.find(filename: uid)
+          .update_one("$set" => {metadata: info})
       end
 
-      def list_files
-        infos = bucket.files_collection.find.to_a
-        infos.map { |info| info.fetch("filename") }
+      def get_file(uid, info = {}, range: nil)
+        file_info = bucket.files_collection.find(filename: uid).first
+        raise Tus::NotFound if file_info.nil?
+
+        filter = {files_id: file_info[:_id]}
+
+        if range
+          chunk_start = range.begin / file_info[:chunkSize] if range.begin
+          chunk_stop  = range.end   / file_info[:chunkSize] if range.end
+
+          filter[:n] = {}
+          filter[:n].update("$gte" => chunk_start) if chunk_start
+          filter[:n].update("$lte" => chunk_stop) if chunk_stop
+        end
+
+        chunks_view = bucket.chunks_collection.find(filter).read(bucket.read_preference).sort(n: 1)
+
+        chunks = Enumerator.new do |yielder|
+          chunks_view.each do |document|
+            data = document[:data].data
+
+            if document[:n] == chunk_start && document[:n] == chunk_stop
+              byte_start = range.begin % file_info[:chunkSize]
+              byte_stop  = range.end   % file_info[:chunkSize]
+            elsif document[:n] == chunk_start
+              byte_start = range.begin % file_info[:chunkSize]
+              byte_stop  = file_info[:chunkSize] - 1
+            elsif document[:n] == chunk_stop
+              byte_start = 0
+              byte_stop  = range.end % file_info[:chunkSize]
+            end
+
+            if byte_start && byte_stop
+              partial_data = data[byte_start..byte_stop]
+              yielder << partial_data
+              partial_data.clear # deallocate chunk string
+            else
+              yielder << data
+            end
+
+            data.clear # deallocate chunk string
+          end
+        end
+
+        Response.new(chunks: chunks, close: ->{chunks_view.close_query})
+      end
+
+      def delete_file(uid, info = {})
+        file_info = bucket.files_collection.find(filename: uid).first
+        bucket.delete(file_info.fetch("_id")) if file_info
       end
 
-      private
+      def expire_files(expiration_date)
+        file_infos = bucket.files_collection.find(uploadDate: {"$lte" => expiration_date}).to_a
+        file_info_ids = file_infos.map { |info| info[:_id] }
+
+        bucket.files_collection.delete_many(_id: {"$in" => file_info_ids})
+        bucket.chunks_collection.delete_many(files_id: {"$in" => file_info_ids})
+      end
+
+      class Response
+        def initialize(chunks:, close:)
+          @chunks = chunks
+          @close  = close
+        end
+
+        def each(&block)
+          @chunks.each(&block)
+        end
 
-      def bson_id(uid)
-        BSON::ObjectId(uid)
+        def close
+          @close.call
+        end
       end
     end
   end

data/lib/tus/storage/s3.rb

@@ -0,0 +1,242 @@
+require "aws-sdk"
+
+require "tus/info"
+require "tus/checksum"
+require "tus/errors"
+
+require "json"
+require "cgi/util"
+
+Aws.eager_autoload!(services: ["S3"])
+
+module Tus
+  module Storage
+    class S3
+      MIN_PART_SIZE = 5 * 1024 * 1024
+
+      attr_reader :client, :bucket, :prefix, :upload_options
+
+      def initialize(bucket:, prefix: nil, upload_options: {}, **client_options)
+        resource = Aws::S3::Resource.new(**client_options)
+
+        @client = resource.client
+        @bucket = resource.bucket(bucket)
+        @prefix = prefix
+        @upload_options = upload_options
+      end
+
+      def create_file(uid, info = {})
+        tus_info = Tus::Info.new(info)
+
+        options = upload_options.dup
+        options[:content_type] = tus_info.metadata["content_type"]
+
+        if filename = tus_info.metadata["filename"]
+          options[:content_disposition] ||= "inline"
+          options[:content_disposition]  += "; filename=\"#{CGI.escape(filename).gsub("+", " ")}\""
+        end
+
+        multipart_upload = object(uid).initiate_multipart_upload(options)
+
+        info["multipart_id"]    = multipart_upload.id
+        info["multipart_parts"] = []
+      end
+
+      def concatenate(uid, part_uids, info = {})
+        create_file(uid, info)
+
+        multipart_upload = object(uid).multipart_upload(info["multipart_id"])
+
+        queue = Queue.new
+        part_uids.each_with_index do |part_uid, idx|
+          queue << {
+            copy_source: [bucket.name, object(part_uid).key].join("/"),
+            part_number: idx + 1
+          }
+        end
+
+        threads = 10.times.map do
+          Thread.new do
+            Thread.current.abort_on_exception = true
+            completed = []
+
+            begin
+              loop do
+                multipart_copy_task = queue.deq(true) rescue break
+
+                part_number = multipart_copy_task[:part_number]
+                copy_source = multipart_copy_task[:copy_source]
+
+                part = multipart_upload.part(part_number)
+                response = part.copy_from(copy_source: copy_source)
+
+                completed << {
+                  part_number: part_number,
+                  etag: response.copy_part_result.etag,
+                }
+              end
+
+              completed
+            rescue
+              queue.clear
+              raise
+            end
+          end
+        end
+
+        parts = threads.flat_map(&:value).sort_by { |part| part[:part_number] }
+
+        multipart_upload.complete(multipart_upload: {parts: parts})
+
+        delete(part_uids.flat_map { |part_uid| [object(part_uid), object("#{part_uid}.info")] })
+
+        info.delete("multipart_id")
+        info.delete("multipart_parts")
+      rescue
+        abort_multipart_upload(multipart_upload) if multipart_upload
+        raise
+      end
+
+      def patch_file(uid, io, info = {})
+        tus_info = Tus::Info.new(info)
+        last_chunk = (tus_info.length && io.size == tus_info.remaining_length)
+
+        if io.size < MIN_PART_SIZE && !last_chunk
+          raise Tus::Error, "Chunk size cannot be smaller than 5MB"
+        end
+
+        upload_id   = info["multipart_id"]
+        part_number = info["multipart_parts"].count + 1
+
+        multipart_upload = object(uid).multipart_upload(upload_id)
+        multipart_part   = multipart_upload.part(part_number)
+        md5              = Tus::Checksum.new("md5").generate(io)
+
+        begin
+          response = multipart_part.upload(body: io, content_md5: md5)
+        rescue Aws::S3::Errors::NoSuchUpload
+          raise Tus::NotFound
+        end
+
+        info["multipart_parts"] << {
+          "part_number" => part_number,
+          "etag"        => response.etag[/"(.+)"/, 1],
+        }
+
+        # finalize the multipart upload if this chunk was the last part
+        if last_chunk
+          multipart_upload.complete(
+            multipart_upload: {
+              parts: info["multipart_parts"].map do |part|
+                {part_number: part["part_number"], etag: part["etag"]}
+              end
+            }
+          )
+
+          info.delete("multipart_id")
+          info.delete("multipart_parts")
+        end
+      end
+
+      def read_info(uid)
+        response = object("#{uid}.info").get
+        JSON.parse(response.body.string)
+      rescue Aws::S3::Errors::NoSuchKey
+        raise Tus::NotFound
+      end
+
+      def update_info(uid, info)
+        object("#{uid}.info").put(body: info.to_json)
+      end
+
+      def get_file(uid, info = {}, range: nil)
+        if range
+          range = "bytes=#{range.begin}-#{range.end}"
+        end
+
+        raw_chunks = Enumerator.new do |yielder|
+          object(uid).get(range: range) do |chunk|
+            yielder << chunk
+            chunk.clear # deallocate string
+          end
+        end
+
+        begin
+          first_chunk = raw_chunks.next
+        rescue Aws::S3::Errors::NoSuchKey
+          raise Tus::NotFound
+        end
+
+        chunks = Enumerator.new do |yielder|
+          yielder << first_chunk
+          loop { yielder << raw_chunks.next }
+        end
+
+        Response.new(chunks: chunks)
+      end
+
+      def delete_file(uid, info = {})
+        if info["multipart_id"]
+          multipart_upload = object(uid).multipart_upload(info["multipart_id"])
+          abort_multipart_upload(multipart_upload)
+
+          delete [object("#{uid}.info")]
+        else
+          delete [object(uid), object("#{uid}.info")]
+        end
+      end
+
+      def expire_files(expiration_date)
+        old_objects = bucket.objects.select do |object|
+          object.last_modified <= expiration_date
+        end
+
+        delete(old_objects)
+
+        bucket.multipart_uploads.each do |multipart_upload|
+          next unless multipart_upload.initiated <= expiration_date
+          most_recent_part = multipart_upload.parts.sort_by(&:last_modified).last
+          if most_recent_part.nil? || most_recent_part.last_modified <= expiration_date
+            abort_multipart_upload(multipart_upload)
+          end
+        end
+      end
+
+      private
+
+      def delete(objects)
+        # S3 can delete maximum of 1000 objects in a single request
+        objects.each_slice(1000) do |objects_batch|
+          delete_params = {objects: objects_batch.map { |object| {key: object.key} }}
+          bucket.delete_objects(delete: delete_params)
+        end
+      end
+
+      # In order to ensure the multipart upload was successfully aborted,
+      # we need to check whether all parts have been deleted, and retry
+      # the abort if the list is nonempty.
+      def abort_multipart_upload(multipart_upload)
+        loop do
+          multipart_upload.abort
+          break unless multipart_upload.parts.any?
+        end
+      rescue Aws::S3::Errors::NoSuchUpload
+        # multipart upload was successfully aborted or doesn't exist
+      end
+
+      def object(key)
+        bucket.object([*prefix, key].join("/"))
+      end
+
+      class Response
+        def initialize(chunks:)
+          @chunks = chunks
+        end
+
+        def each(&block)
+          @chunks.each(&block)
+        end
+      end
+    end
+  end
+end