Merge pull request #5478 from avalonmediasystem/stream_token_optimized

Stream token optimization for media object IIIF manifest endpoint

app/controllers/media_objects_controller.rb

@@ -464,11 +464,8 @@ def manifest
     @media_object = SpeedyAF::Proxy::MediaObject.find(params[:id])
     authorize! :read, @media_object
 
-    master_files = master_file_presenters
-    canvas_presenters = master_files.collect do |mf|
-      stream_info = secure_streams(mf.stream_details, @media_object.id)
-      IiifCanvasPresenter.new(master_file: mf, stream_info: stream_info)
-    end
+    stream_info_hash = secure_stream_infos(master_file_presenters, @media_object.id)
+    canvas_presenters = master_file_presenters.collect { |mf| IiifCanvasPresenter.new(master_file: mf, stream_info: stream_info_hash[mf.id]) }
     presenter = IiifManifestPresenter.new(media_object: @media_object, master_files: canvas_presenters, lending_enabled: lending_enabled?(@media_object))
 
     manifest = IIIFManifest::V3::ManifestFactory.new(presenter).to_h

app/helpers/security_helper.rb

@@ -26,11 +26,34 @@ def secure_streams(stream_info, media_object_id)
     stream_info
   end
 
-  def add_stream_url(stream_info)
+  # Batch version of secure_streams
+  # media_object CDL check only happens once
+  # session tokens retrieved in batch then passed into add_stream_url
+  # Returns Hash[MasterFile.id, stream_info]
+  def secure_stream_infos(master_files, media_object_id)
+    stream_info_hash = {}
+    if not_checked_out?(media_object_id)
+      master_files.each { |mf| stream_info_hash[mf.id] = mf.stream_details }
+    else
+      stream_tokens = StreamToken.get_session_tokens_for(session: session, targets: master_files.map(&:id))
+      stream_token_hash = stream_tokens.pluck(:target, :token).to_h
+      master_files.each { |mf| stream_info_hash[mf.id] = secure_stream_info(mf.stream_details, stream_token_hash[mf.id]) }
+    end
+    stream_info_hash
+  end
+
+  # Same as secure_streams except without CDL checking
+  def secure_stream_info(stream_info, token)
+    add_stream_url(stream_info, token: token)
+    stream_info
+  end
+
+  # Optional token kwarg used if passed in
+  def add_stream_url(stream_info, token: nil)
     add_stream_cookies(id: stream_info[:id])
     [:stream_hls].each do |protocol|
       stream_info[protocol].each do |quality|
-        quality[:url] = SecurityHandler.secure_url(quality[:url], session: session, target: stream_info[:id], protocol: protocol)
+        quality[:url] = SecurityHandler.secure_url(quality[:url], session: session, target: stream_info[:id], protocol: protocol, token: token)
       end
     end
   end

app/models/stream_token.rb

@@ -34,15 +34,35 @@ def self.find_or_create_session_token(session, target)
     result.token
   end
 
+  def self.get_session_tokens_for(session: {}, targets: [])
+    purge_expired!(session)
+
+    token_attributes = targets.collect do |target|
+      hash_token = Digest::SHA1.new
+      hash_token << media_token(session) << target
+      {target: target, token: hash_token.to_s, expires: (Time.now.utc + Settings.streaming.stream_token_ttl.minutes)}
+    end
+    existing_token_hash = StreamToken.where(token: token_attributes.pluck(:token)).pluck(:token, :id).to_h
+    token_attributes.each do |attrs|
+      attrs[:id] = existing_token_hash[attrs[:token]] if existing_token_hash[attrs[:token]].present?
+    end
+    result = StreamToken.upsert_all(token_attributes)
+
+    tokens = StreamToken.where(id: result.to_a.pluck("id"))
+    session[:hash_tokens] += tokens.pluck(:token)
+    session[:hash_tokens].uniq! # Avoid duplicate entry
+    tokens.to_a
+  end
+
   def self.logout!(session)
     session[:hash_tokens].each do |sha|
       where(token: sha).find_each(&:delete)
     end unless session[:hash_tokens].nil?
   end
 
   def self.purge_expired!(session)
-    purged = expired.each(&:delete)
-    Array(session[:hash_tokens]).reject! { |token| !StreamToken.exists?(token: token) }
+    purged = expired.delete_all
+    session[:hash_tokens] = StreamToken.where(token: Array(session[:hash_tokens])).pluck(:token)
     purged
   end
 

app/services/security_service.rb

@@ -28,7 +28,7 @@ def rewrite_url(url, context)
       end
     else
       session = context[:session] || { media_token: nil }
-      token = StreamToken.find_or_create_session_token(session, context[:target])
+      token = context[:token] || StreamToken.find_or_create_session_token(session, context[:target])
       "#{url}?token=#{token}"
     end
   end

spec/controllers/media_objects_controller_spec.rb

@@ -1840,7 +1840,8 @@
 
   describe '#manifest' do
     context 'read from solr' do
-      let!(:media_object) { FactoryBot.create(:published_media_object, :with_master_file, visibility: 'public') }
+      let!(:master_file) { FactoryBot.create(:master_file, :with_derivative, media_object: media_object) }
+      let!(:media_object) { FactoryBot.create(:published_media_object, visibility: 'public') }
       it 'should not read from fedora' do
         perform_enqueued_jobs(only: MediaObjectIndexingJob)
         WebMock.reset_executed_requests!

spec/helpers/security_helper_spec.rb

@@ -103,7 +103,7 @@
 
   context 'when using non-AWS streaming server' do
     describe '#add_stream_cookies' do
-      it 'adds security tokens to cookies' do
+      it 'does not add security tokens to cookies' do
         expect { helper.add_stream_cookies(stream_info) }.not_to change { controller.cookies.sum {|k,v| 1} }
       end
     end
@@ -158,5 +158,60 @@
         end
       end
     end
+
+    describe '#secure_stream_infos' do
+      let(:master_file) { FactoryBot.create(:master_file, :with_derivative, media_object: media_object) }
+      let(:token) { 'dcba-4321' }
+      let(:secure_url_with_token) { secure_url + "?token=#{token}" }
+
+      before do
+        allow(SecurityHandler).to receive(:secure_url).with(String, token: token).and_return(secure_url_with_token)
+      end
+
+      context 'controlled digital lending is disabled' do
+        before { allow(Settings.controlled_digital_lending).to receive(:enable).and_return(false) }
+
+        it 'rewrites urls in the stream_infos' do
+          stream_info_hash = helper.secure_stream_infos([master_file], media_object.id)
+          stream_info = stream_info_hash[master_file.id]
+          [:stream_hls].each do |protocol|
+            stream_info[protocol].each do |quality|
+              expect(quality[:url]).to eq secure_url
+            end
+          end
+        end
+      end
+
+      context 'controlled digital lending is enabled' do
+        before { allow(Settings.controlled_digital_lending).to receive(:enable).and_return(true) }
+        before { allow(Settings.controlled_digital_lending).to receive(:collections_enabled).and_return(true) }
+
+        context 'the user has the item checked out' do
+          before { FactoryBot.create(:checkout, media_object_id: media_object.id, user_id: user.id)}
+
+          it 'rewrites urls in the stream_infos' do
+            stream_info_hash = helper.secure_stream_infos([master_file], media_object.id)
+            stream_info = stream_info_hash[master_file.id]
+            [:stream_hls].each do |protocol|
+              stream_info[protocol].each do |quality|
+                expect(quality[:url]).to eq secure_url
+              end
+            end
+          end
+        end
+
+        context 'the user does not have the item checked out' do
+          it 'does not rewrite urls in the stream_info' do
+            stream_info_hash = helper.secure_stream_infos([master_file], media_object.id)
+            stream_info = stream_info_hash[master_file.id]
+            [:stream_hls].each do |protocol|
+              stream_info[protocol].each do |quality|
+                expect(quality[:url]).not_to eq secure_url
+              end
+            end
+          end
+        end
+      end
+    end
   end
 end

spec/models/stream_token_spec.rb

@@ -37,6 +37,33 @@
       expect(token).to eq token2
       expect(session[:hash_tokens].count(token)).to eq 1
     end
+
+    context '#get_session_tokens_for' do
+      let(:targets) { ['D1452CB7-4DC2-4A26-AC2C-FBCE943C164C', target] }
+
+      it 'should create the tokens' do
+        expect(StreamToken.get_session_tokens_for(session: session, targets: targets)).to match_array([be_instance_of(StreamToken), be_instance_of(StreamToken)])
+      end
+
+      it 'stores the tokens in the session' do
+        tokens = StreamToken.get_session_tokens_for(session: session, targets: targets)
+        expect(session[:hash_tokens]).to include *tokens.pluck(:token)
+      end
+
+      it 'stores the tokens once in the session' do
+        tokens = StreamToken.get_session_tokens_for(session: session, targets: targets)
+        tokens2 = StreamToken.get_session_tokens_for(session: session, targets: targets)
+        expect(tokens).to eq tokens2
+        expect(session[:hash_tokens].count(tokens.first.token)).to eq 1
+        expect(session[:hash_tokens].count(tokens.second.token)).to eq 1
+      end
+
+      it 'updates expires' do
+        StreamToken.get_session_tokens_for(session: session, targets: targets)
+        expect { StreamToken.get_session_tokens_for(session: session, targets: targets) }.not_to change { StreamToken.count }
+        expect { StreamToken.get_session_tokens_for(session: session, targets: targets) }.to change { StreamToken.pluck(:expires) }
+      end
+    end
   end
 
   describe 'valid_token?' do

spec/services/security_service_spec.rb

@@ -71,6 +71,14 @@
       it 'adds a StreamToken param' do
         expect(subject.rewrite_url(url, context)).to start_with "http://example.com/streaming/id?token="
       end
+
+      context 'when token provided' do
+        let(:context) {{ session: {}, target: 'abcd1234', protocol: :stream_hls, token: 'dcba-4321' }}
+
+        it 'adds token param' do
+          expect(subject.rewrite_url(url, context)).to eq "http://example.com/streaming/id?token=dcba-4321"
+        end
+      end
     end
   end