Skip to content
This repository has been archived by the owner on Jul 12, 2023. It is now read-only.

Commit

Permalink
Browse files Browse the repository at this point in the history
Added fallback to local storage
  • Loading branch information
Noam committed May 19, 2011
1 parent 86b89b3 commit 84faa83
Show file tree
Hide file tree
Showing 4 changed files with 287 additions and 37 deletions.
Empty file.
90 changes: 60 additions & 30 deletions lib/rack/contrib/deflect.rb
Expand Up @@ -49,6 +49,7 @@ class Deflect

def initialize app, options = {}
@mutex = Mutex.new
@local_storage_map = {}
@app, @options = app, {
:log => false,
:log_format => 'deflect(%s): %s',
Expand All @@ -63,7 +64,12 @@ def initialize app, options = {}
:notifier_callback => nil
}.merge(options)

@redis_storage = Redis.new(@options[:redis_interface])
@redis_storage = nil
begin
Redis.new(@options[:redis_interface]) unless @options[:redis_interface].blank?
rescue Timeout::Error
# No redis.
end

unless @options[:reset_for].nil?
@options[:reset_for].each do |addr|
Expand All @@ -73,7 +79,7 @@ def initialize app, options = {}
end

def call env
if @options[:ignore_agents].any? {|word| env["HTTP_USER_AGENT"].to_s.downcase.include?(word) }
if options[:ignore_agents].any? {|word| env["HTTP_USER_AGENT"].to_s.downcase.include?(word) }
log "Skipping user agent #{env["HTTP_USER_AGENT"]}"
status, headers, body = @app.call env
[status, headers, body]
Expand All @@ -91,14 +97,14 @@ def deflect!
def deflect? env
@env = env
@remote_addr = env['REMOTE_ADDR']
return false if @options[:whitelist].include? @remote_addr
return true if @options[:blacklist].include? @remote_addr
return false if options[:whitelist].include? @remote_addr
return true if options[:blacklist].include? @remote_addr
sync { watch }
end

def log message
return unless @options[:log]
@options[:log].puts(@options[:log_format] % [Time.now.strftime(@options[:log_date_format]), message])
return unless options[:log]
options[:log].puts(options[:log_format] % [Time.now.strftime(options[:log_date_format]), message])
end

def sync &block
Expand All @@ -107,6 +113,8 @@ def sync &block

def watch
increment_requests
init_local_storage_map(@remote_addr)
set_key("expires", Time.now + options[:interval])
clear! if watch_expired? and not blocked?
clear! if blocked? and block_expired?
block! if watching? and exceeded_request_threshold?
Expand All @@ -116,23 +124,23 @@ def watch
def block!
return if blocked?
log "blocked #{@remote_addr}"
rset("block_expires", Time.now + @options[:block_duration])
notifier = @options[:notifier_callback]
blocked_uris = rget("requested_uris")
set_key("block_expires", (Time.now + options[:block_duration]).to_s)
notifier = options[:notifier_callback]
blocked_uris = get_key("requested_uris")
notifier.call(@remote_addr, Socket.gethostname, blocked_uris) unless notifier.nil?
end

def blocked?
!(rget("block_expires").nil?)
!(get_key("block_expires").nil?)
end

def block_expired?
block_expires_str = rget("block_expires")
block_expires_str = get_key("block_expires")
!(block_expires_str.nil?) && (Time.parse(block_expires_str) < Time.now) rescue false
end

def watching?
Integer(rget('requests')) > 0
Integer(get_key('requests')) > 0
end

def clear!
Expand All @@ -142,49 +150,71 @@ def clear!
end

def clear_for_address(address)
rdel("block_expires", address)
rdel("requests", address)
rdel("request_uris", address)
del_key("expires", address)
del_key("block_expires", address)
del_key("requests", address)
del_key("request_uris", address)
end

def increment_requests
rincr("requests")
incr_key("requests")
log "Current Request: #{@env["REQUEST_URI"]}"
rrpush("request_uris", "#{@env["REQUEST_METHOD"]} '#{@env["HTTP_HOST"]}#{@env["REQUEST_URI"]}' #{} => #{@env["HTTP_USER_AGENT"]}")
end

def exceeded_request_threshold?
Integer(rget("requests")) > @options[:request_threshold]
Integer(get_key("requests")) > options[:request_threshold]
end

def watch_expired?
expires_str = rget("block_expires")
!(expires_str.nil?) && (Time.parse(rget("block_expires")) <= Time.now) rescue false
expires_str = get_key("expires")
!(expires_str.blank?) && (Time.parse(expires_str) <= Time.now) rescue false
end

def init_local_storage_map(addr)
@local_storage_map[addr] = {
:expires => Time.now + options[:interval],
:requests => 0,
:request_uris => []
}
end

# Redis-related functions:
def rkey(key, addr=@remote_addr)
def redis_key(key, addr=@remote_addr)
"#{addr}:#{key}"
end

def rset(key, val, addr=@remote_addr)
@redis_storage.set(rkey(key, addr), val)
def set_key(key, val, addr=@remote_addr)
@local_storage_map[addr.to_sym] = {} if @local_storage_map[addr.to_sym].nil?
@local_storage_map[addr.to_sym][key.to_sym] = val
@redis_storage.set(redis_key(key, addr), val) unless @redis_storage.nil?
end

def rget(key, addr=@remote_addr)
@redis_storage.get(rkey(key, addr))
def get_key(key, addr=@remote_addr)
val = @redis_storage.get(redis_key(key, addr)) unless @redis_storage.nil?
val = @local_storage_map[addr.to_sym][key.to_sym] if val.nil? && !(@local_storage_map[addr.to_sym].nil?)
val
end

def rdel(key, addr=@remote_addr)
@redis_storage.del(rkey(key, addr))
def del_key(key, addr=@remote_addr)
@local_storage_map[addr.to_sym].delete(key.to_sym) unless @local_storage_map[addr.to_sym].nil?
@redis_storage.del(redis_key(key, addr)) unless @redis_storage.nil?
end

def rincr(key, addr=@remote_addr)
@redis_storage.incr(rkey(key, addr))
def incr_key(key, addr=@remote_addr)
@local_storage_map[addr.to_sym] = {} if @local_storage_map[addr.to_sym].nil?
@local_storage_map[addr.to_sym][key.to_sym] = (@local_storage_map[addr.to_sym][key.to_sym].nil? ? 0 : @local_storage_map[addr.to_sym][key.to_sym]) + 1
@redis_storage.incr(redis_key(key, addr)) unless @redis_storage.nil?
end

def rrpush(key, val, addr=@remote_addr)
@redis_storage.rpush(rkey(key, addr), val)
@local_storage_map[addr.to_sym] = {} if @local_storage_map[addr.to_sym].nil?
if @local_storage_map[addr.to_sym][key.to_sym]
@local_storage_map[addr.to_sym][key.to_sym] << val
else
@local_storage_map[addr.to_sym][key.to_sym] = [val]
end

@redis_storage.rpush(redis_key(key, addr), val) unless @redis_storage.nil?
end
end
end
155 changes: 155 additions & 0 deletions lib/rack/contrib/deflect_orig.rb
@@ -0,0 +1,155 @@
require 'thread'
require 'socket'
# TODO: optional stats
# TODO: performance
# TODO: clean up tests

module Rack

##
# Rack middleware for protecting against Denial-of-service attacks
# http://en.wikipedia.org/wiki/Denial-of-service_attack.
#
# This middleware is designed for small deployments, which most likely
# are not utilizing load balancing from other software or hardware. Deflect
# current supports the following functionality:
#
# * Saturation prevention (small DoS attacks, or request abuse)
# * Blacklisting of remote addresses
# * Whitelisting of remote addresses
# * Logging
#
# === Options:
#
# :log When false logging will be bypassed, otherwise pass an object responding to #puts
# :log_format Alter the logging format
# :log_date_format Alter the logging date format
# :request_threshold Number of requests allowed within the set :interval. Defaults to 100
# :interval Duration in seconds until the request counter is reset. Defaults to 5
# :block_duration Duration in seconds that a remote address will be blocked. Defaults to 900 (15 minutes)
# :whitelist Array of remote addresses which bypass Deflect. NOTE: this does not block others
# :blacklist Array of remote addresses immediately considered malicious
# :ignore_agents a list of words from user agents allow in.
#
# === Examples:
#
# use Rack::Deflect, :log => $stdout, :request_threshold => 20, :interval => 2, :block_duration => 60
#
# CREDIT: TJ Holowaychuk <tj@vision-media.ca>
#

class Deflect

attr_reader :options

def initialize app, options = {}
@mutex = Mutex.new
@remote_addr_map = {}
@app, @options = app, {
:log => false,
:log_format => 'deflect(%s): %s',
:log_date_format => '%m/%d/%Y',
:request_threshold => 100,
:interval => 5,
:block_duration => 900,
:whitelist => [],
:blacklist => [],
:ignore_agents => []
}.merge(options)
end

def call env
if options[:ignore_agents].any? {|word| env["HTTP_USER_AGENT"].to_s.downcase.include?(word) }
rails_logger "Skipping user agent #{env["HTTP_USER_AGENT"]}"
status, headers, body = @app.call env
[status, headers, body]
else
return deflect! if deflect? env
status, headers, body = @app.call env
[status, headers, body]
end
end

def deflect!
[403, { 'Content-Type' => 'text/html', 'Content-Length' => '0' }, []]
end

def deflect? env
@env = env
@remote_addr = env['REMOTE_ADDR']
return false if options[:whitelist].include? @remote_addr
return true if options[:blacklist].include? @remote_addr
sync { watch }
end

def log message
return unless options[:log]
options[:log].puts(options[:log_format] % [Time.now.strftime(options[:log_date_format]), message])
end

def sync &block
@mutex.synchronize(&block)
end

def map
@remote_addr_map[@remote_addr] ||= {
:expires => Time.now + options[:interval],
:requests => 0,
:request_uris => []
}
end

def watch
increment_requests
clear! if watch_expired? and not blocked?
clear! if blocked? and block_expired?
block! if watching? and exceeded_request_threshold?
blocked?
end

def block!
return if blocked?
log "blocked #{@remote_addr}"
Notifier.deliver_ip_blocked(@remote_addr, Socket.gethostname, map[:request_uris])
map[:block_expires] = Time.now + options[:block_duration]
end

def blocked?
map.has_key? :block_expires
end

def block_expired?
map[:block_expires] < Time.now rescue false
end

def watching?
@remote_addr_map.has_key? @remote_addr
end

def clear!
return unless watching?
log "released #{@remote_addr}" if blocked?
@remote_addr_map.delete @remote_addr
end

def rails_logger(message)
RAILS_DEFAULT_LOGGER.info "Rack::Deflect = #{message}"
end

def increment_requests
map[:requests] += 1
rails_logger "Current Request: #{@env["REQUEST_URI"]}"
map[:request_uris] = map[:request_uris] << "#{@env["REQUEST_METHOD"]} '#{@env["HTTP_HOST"]}#{@env["REQUEST_URI"]}' #{} => #{@env["HTTP_USER_AGENT"]}"
end

def exceeded_request_threshold?
map[:requests] > options[:request_threshold]
end

def watch_expired?
map[:expires] <= Time.now
end

end
end

0 comments on commit 84faa83

Please sign in to comment.