server.rb in dragonfly-0.9.1

- old
+ new

@@ -1,11 +1,16 @@
 module Dragonfly
   class Server
 
+    # Exceptions
+    class JobNotAllowed < RuntimeError; end
+
     include Loggable
     include Configurable
     
+    configurable_attr :allow_fetch_file, false
+    configurable_attr :allow_fetch_url, false
     configurable_attr :dragonfly_url, '/dragonfly'
     configurable_attr :protect_from_dos_attacks, false
     configurable_attr :url_format, '/:job/:basename.:format'
     configurable_attr :url_host
 
@@ -26,10 +31,11 @@
     def call(env)
       if dragonfly_url == env["PATH_INFO"]
         dragonfly_response
       elsif (params = url_mapper.params_for(env["PATH_INFO"], env["QUERY_STRING"])) && params['job']
         job = Job.deserialize(params['job'], app)
+        validate_job!(job)
         job.validate_sha!(params['sha']) if protect_from_dos_attacks
         response = Response.new(job, env)
         catch(:halt) do
           if before_serve_callback && response.will_be_served?
             before_serve_callback.call(job, env)
@@ -37,17 +43,20 @@
           response.to_response
         end
       else
         [404, {'Content-Type' => 'text/plain', 'X-Cascade' => 'pass'}, ['Not found']]
       end
-    rescue Serializer::BadString, Job::InvalidArray => e
-      log.warn(e.message)
-      [404, {'Content-Type' => 'text/plain'}, ['Not found']]
     rescue Job::NoSHAGiven => e
       [400, {"Content-Type" => 'text/plain'}, ["You need to give a SHA parameter"]]
     rescue Job::IncorrectSHA => e
       [400, {"Content-Type" => 'text/plain'}, ["The SHA parameter you gave (#{e}) is incorrect"]]
+    rescue JobNotAllowed => e
+      log.warn(e.message)
+      [403, {"Content-Type" => 'text/plain'}, ["Forbidden"]]
+    rescue Serializer::BadString, Job::InvalidArray => e
+      log.warn(e.message)
+      [404, {'Content-Type' => 'text/plain'}, ['Not found']]
     end
 
     def url_for(job, opts={})
       opts = opts.dup
       host = opts.delete(:host) || url_host
@@ -95,9 +104,16 @@
         'Content-Type' => 'text/plain',
         'Content-Size' => body.bytesize.to_s
         },
         [body]
       ]
+    end
+
+    def validate_job!(job)
+      if job.fetch_file_step && !allow_fetch_file ||
+         job.fetch_url_step && !allow_fetch_url
+        raise JobNotAllowed, "Dragonfly Server doesn't allow requesting job with steps #{job.steps.inspect}"
+      end
     end
 
   end
 end