6 files changed, 47 insertions, 49 deletions
diff --git a/lib/active_file/disk_controller.rb b/lib/active_file/disk_controller.rb
index 7d94b02f5c..b5a19fa5fc 100644
--- a/lib/active_file/disk_controller.rb
+++ b/lib/active_file/disk_controller.rb
@@ -1,17 +1,16 @@
-# FIXME: To be used by DiskSite#url
 class ActiveFile::DiskController < ActionController::Base
   def show
-    if verified_key.expired?
-      head :gone
-    else
-      blob = ActiveFile::Blob.find_by!(key: verified_key.to_s)
+    if key = decode_verified_key
+      blob = ActiveFile::Blob.find_by!(key: key)
       send_data blob.download, filename: blob.filename, type: blob.content_type, disposition: disposition_param
+    else
+      head :not_found
     end
   end
 
   private
-    def verified_key
-      ActiveFile::Sites::DiskSite::VerifiedKeyWithExpiration.new(params[:id])
+    def decode_verified_key
+      ActiveFile::Sites::DiskSite::VerifiedKeyWithExpiration.decode(params[:id])
     end
 
     def disposition_param
diff --git a/lib/active_file/sites/disk_site.rb b/lib/active_file/sites/disk_site.rb
index da1e69df03..be8a2437a1 100644
--- a/lib/active_file/sites/disk_site.rb
+++ b/lib/active_file/sites/disk_site.rb
@@ -2,42 +2,6 @@ require "fileutils"
 require "pathname"
 
 class ActiveFile::Sites::DiskSite < ActiveFile::Site
-  class_attribute :verifier, default: -> { Rails.application.message_verifier('ActiveFile::DiskSite') }
-
-  class << self
-    def generate_verifiable_key(key, expires_in:)
-      VerifiedKeyWithExpiration
-    end
-  end
-
-  class VerifiableKeyWithExpiration
-    def initialize(verifiable_key_with_expiration)
-      verified_key_with_expiration = ActiveFile::Sites::DiskSite.verify(verifiable_key_with_expiration)
-
-      @key        = verified_key_with_expiration[:key]
-      @expires_at = verified_key_with_expiration[:expires_at]
-    end
-
-    def expired?
-      @expires_at && Time.now.utc > @expires_at
-    end
-
-    def decoded
-      key
-    end
-  end
-
-  class VerifiedKeyWithExpiration
-    def initialize(key, expires_in: nil)
-      @key = key
-      @expires_at = Time.now.utc.advance(sec: expires_in)
-    end
-
-    def encoded
-      ActiveFile::Sites::DiskSite.verify.generate({ key: @key, expires_at: @expires_at })
-    end
-  end
-
   attr_reader :root
 
   def initialize(root:)
@@ -75,10 +39,12 @@ class ActiveFile::Sites::DiskSite < ActiveFile::Site
 
 
   def url(key, disposition:, expires_in: nil)
+    verified_key_with_expiration = ActiveFile::VerifiedKeyWithExpiration.encode(key, expires_in: expires_in)
+
     if defined?(Rails)
-      Rails.application.routes.url_helpers.rails_disk_blob_path(key)
+      Rails.application.routes.url_helpers.rails_disk_blob_path(verified_key_with_expiration)
     else
-      "/rails/blobs/#{key}"
+      "/rails/blobs/#{verified_key_with_expiration}"
     end
   end
 
diff --git a/lib/active_file/verified_key_with_expiration.rb b/lib/active_file/verified_key_with_expiration.rb
new file mode 100644
index 0000000000..601401278b
--- /dev/null
+++ b/lib/active_file/verified_key_with_expiration.rb
@@ -0,0 +1,15 @@
+class ActiveFile::VerifiedKeyWithExpiration
+  class_attribute :verifier, default: defined?(Rails) ? Rails.application.message_verifier('ActiveFile') : nil
+
+  def self.encode(key, expires_in: nil)
+    verifier.generate([ key, expires_in ? Time.now.utc.advance(sec: expires_in) : nil ])
+  end
+
+  def self.decode(encoded_key)
+    key, expires_at = verifier.verified(encoded_key)
+
+    if key
+      key if expires_at.nil? || Time.now.utc < expires_at
+    end
+  end
+end
diff --git a/test/blob_test.rb b/test/blob_test.rb
index 88b513c946..ac54e0f2ca 100644
--- a/test/blob_test.rb
+++ b/test/blob_test.rb
@@ -2,8 +2,6 @@ require "test_helper"
 require "database/setup"
 require "active_file/blob"
 
-ActiveFile::Blob.site = ActiveFile::Sites::DiskSite.new(root: File.join(Dir.tmpdir, "active_file"))
-
 class ActiveFile::BlobTest < ActiveSupport::TestCase
   test "create after upload sets byte size and checksum" do
     data = "Hello world!"
@@ -14,9 +12,12 @@ class ActiveFile::BlobTest < ActiveSupport::TestCase
     assert_equal Digest::MD5.hexdigest(data), blob.checksum
   end
 
-  test "url" do
+  test "url expiring in 5 minutes" do
     blob = create_blob
-    assert_equal "/rails/blobs/#{blob.key}", blob.url
+
+    travel_to Time.now do
+      assert_equal "/rails/blobs/#{ActiveFile::VerifiedKeyWithExpiration.encode(blob.key, expires_in: 5.minutes)}", blob.url
+    end
   end
 
   private
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 0964774e00..5be2631ceb 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -4,3 +4,9 @@ require "active_support/testing/autorun"
 require "byebug"
 
 require "active_file"
+
+require "active_file/site"
+ActiveFile::Blob.site = ActiveFile::Sites::DiskSite.new(root: File.join(Dir.tmpdir, "active_file"))
+
+require "active_file/verified_key_with_expiration"
+ActiveFile::VerifiedKeyWithExpiration.verifier = ActiveSupport::MessageVerifier.new("Testing")
diff --git a/test/verified_key_with_expiration_test.rb b/test/verified_key_with_expiration_test.rb
new file mode 100644
index 0000000000..ac605a95e9
--- /dev/null
+++ b/test/verified_key_with_expiration_test.rb
@@ -0,0 +1,11 @@
+require "test_helper"
+require "active_support/core_ext/securerandom"
+
+class ActiveFile::VerifiedKeyWithExpirationTest < ActiveSupport::TestCase
+  FIXTURE_KEY = SecureRandom.base58(24)
+
+  test "without expiration" do
+    encoded_key = ActiveFile::VerifiedKeyWithExpiration.encode(FIXTURE_KEY)
+    assert_equal FIXTURE_KEY, ActiveFile::VerifiedKeyWithExpiration.decode(encoded_key)
+  end
+end