File Coverage

blib/lib/Perlbal/Plugin/Throttle.pm

Criterion	Covered	Total	%
statement	75	313	23.9
branch	0	118	0.0
condition	1	51	1.9
subroutine	25	40	62.5
pod	0	3	0.0
total	101	525	19.2

package Perlbal::Plugin::Throttle;

use strict;
use warnings;

our $VERSION = '1.20';

use List::Util 'min';
use Danga::Socket 1.59;
use Perlbal 1.70;
use Perlbal::ClientProxy ();
use Perlbal::HTTPHeaders ();
use Time::HiRes ();

# Debugging flag
use constant VERBOSE => $ENV{THROTTLE_VERBOSE} || 0;

sub load {
    # behavior
    Perlbal::Service::add_tunable(
        whitelist_file => {
            check_role => '*',
            des => "File containing CIDRs which are never throttled. (Net::CIDR::Lite must be installed.)",
            check_type => 'file_or_none',
        }
    );
    Perlbal::Service::add_tunable(
        blacklist_file => {
            check_role => '*',
            des => "File containing CIDRs which are always denied outright. (Net::CIDR::Lite must be installed.)",
            check_type => 'file_or_none',
        }
    );
    Perlbal::Service::add_tunable(
        default_action => {
            check_role => '*',
            des => "Whether to throttle or allow new connections from clients on neither the whitelist nor blacklist.",
            check_type => [enum => [qw( allow throttle )]],
            default => 'throttle',
        }
    );
    Perlbal::Service::add_tunable(
        blacklist_action => {
            check_role => '*',
            des => "Whether to deny or throttle connections from blacklisted IPs.",
            check_type => [enum => [qw( deny throttle )]],
            default => 'deny',
        }
    );

    # filters
    Perlbal::Service::add_tunable(
        path_regex => {
            check_role => '*',
            des => "Regex which path portion of URI must match for throttling to be in effect.",
        }
    );
    Perlbal::Service::add_tunable(
        method_regex => {
            check_role => '*',
            des => "Regex which HTTP method must match for throttling to be in effect.",
        }
    );

    # logging
    Perlbal::Service::add_tunable(
        log_events => {
            check_role => '*',
            des => q{Comma-separated list of events to log (ban, unban, whitelisted, blacklisted, concurrent, throttled, banned; all; none). If this is changed after the plugin is registered, the "throttle reload config" command must be issued.},
            check_type => [regexp => qr/^(ban|unban|whitelisted|blacklisted|concurrent|throttled|banned|all|none| |,)+$/, "log_events is a comma-separated list of loggable events"],
            default => 'all',
        }
    );
    Perlbal::Service::add_tunable(
        log_only => {
            check_role => '*',
            des => "Perform the full throttling calculation, but don't actually throttle or deny connections.",
            check_type => 'bool',
            default => 0,
        }
    );

    # throttler parameters
    Perlbal::Service::add_tunable(
        throttle_threshold_seconds => {
            check_role => '*',
            des => "Minimum allowable time between requests. If a non-white/-blacklisted client makes another connection within this interval, it will be throttled for initial_delay seconds. Further connections will double the delay time.",
            check_type => 'int',
            default => 60,
        }
    );
    Perlbal::Service::add_tunable(
        initial_delay => {
            check_role => '*',
            des => "Minimum time for a connection to be throttled if occurring within throttle_threshold_seconds of last attempt.",
            check_type => 'int',
            default => 3,
        }
    );
    Perlbal::Service::add_tunable(
        max_delay => {
            check_role => '*',
            des => "Maximum time for a connection to be throttled after exponential increase from initial_delay.",
            check_type => 'int',
            default => 300,
        }
    );
    Perlbal::Service::add_tunable(
        max_concurrent => {
            check_role => '*',
            des => "Maximum number of connections accepted at a time from a single IP, per perlbal instance.",
            check_type => 'int',
            default => 2,
        }
    );
    Perlbal::Service::add_tunable(
        ban_threshold => {
            check_role => '*',
            des => "Number of accumulated violations required to temporarily ban the source IP.",
            check_type => 'int',
            default => 0,
        }
    );
    Perlbal::Service::add_tunable(
        ban_expiration => {
            check_role => '*',
            des => "Number of seconds after which banned IP is unbanned.",
            check_type => 'int',
            default => 60,
        }
    );

    # memcached
    Perlbal::Service::add_tunable(
        memcached_servers => {
            check_role => '*',
            des => "List of memcached servers to share state in, if desired. (Cache::Memcached::Async must be installed.)",
        }
    );
    Perlbal::Service::add_tunable(
        memcached_async_clients => {
            check_role => '*',
            des => "Number of parallel Cache::Memcached::Async objects to use.",
            check_type => 'int',
            default => 10,
        }
    );
    Perlbal::Service::add_tunable(
        instance_name => {
            check_role => '*',
            des => "Name of throttler instance; instances with the same name will share knowledge of IPs.",
            default => 'Throttle',
        }
    );

    Perlbal::register_global_hook('manage_command.throttle', sub {
        my $mc = shift->parse(qr/^
                              throttle\s+
                              (reload)\s+ # command
                              (whitelist|blacklist|config)
                              $/xi,
                              "usage: throttle reload ");
        my ($cmd, $what) = $mc->args;

        my $svcname = $mc->{ctx}{last_created};
        unless ($svcname) {
            return $mc->err("No service name set. This command must be used after CREATE SERVICE  or USE ");
        }

        my $ss = Perlbal->service($svcname);
        return $mc->err("Non-existent service '$svcname'") unless $ss;

        my $cfg = $ss->{extra_config} ||= {};
        my $stash = $cfg->{_throttle_stash} ||= {};

        if ($cmd eq 'reload') {
            if ($what eq 'whitelist') {
                if (my $whitelist = $cfg->{whitelist_file}) {
                    eval { $stash->{whitelist} = load_cidr_list($whitelist); };
                    return $mc->err("Couldn't load $whitelist: $@")
                        if $@ || !$stash->{whitelist};
                }
                else {
                    return $mc->err("no whitelist file configured");
                }
            }
            elsif ($what eq 'blacklist') {
                if (my $blacklist = $cfg->{blacklist_file}) {
                    eval { $stash->{blacklist} = load_cidr_list($blacklist); };
                    return $mc->err("Couldn't load $blacklist: $@")
                        if $@ || !$stash->{blacklist};
                }
                else {
                    return $mc->err("no blacklist file configured");
                }
            }
            elsif ($what eq 'config') {
                $stash->{config_reloader}->();
            }
            else {
                return $mc->err("unknown object to reload: $what");
            }
        }
        else {
            return $mc->err("unknown command $cmd");
        }

        return $mc->ok;
    });
}

# magical Perlbal hook return value constants
use constant HANDLE_REQUEST             => 0;
use constant IGNORE_REQUEST             => 1;

# indexes into logging flag list
use constant LOG_BAN_ADDED              => 0;
use constant LOG_BAN_REMOVED            => 1;
use constant LOG_ALLOW_WHITELISTED      => 2;
use constant LOG_ALLOW_DEFAULT          => 3;
use constant LOG_DENY_BANNED            => 4;
use constant LOG_DENY_BLACKLISTED       => 5;
use constant LOG_DENY_CONCURRENT        => 6;
use constant LOG_THROTTLE_BLACKLISTED   => 7;
use constant LOG_THROTTLE_DEFAULT       => 8;
use constant NUM_LOG_FLAGS              => 9;

use constant RESULT_ALLOW               => 0;
use constant RESULT_THROTTLE            => 1;
use constant RESULT_DENY                => 2;

# localized variable to track if a connection has already been throttled
our $DELAYED = undef;

sub register {
    my ($class, $svc) = @_;

    VERBOSE and Perlbal::log(info => "Registering Throttle plugin on service $svc->{name}");

    my $cfg   = $svc->{extra_config}    ||= {};
    my $stash = $cfg->{_throttle_stash} ||= {};

    # these are allowed to die at register time
    $stash->{whitelist} = load_cidr_list($cfg->{whitelist_file}) if $cfg->{whitelist_file};
    $stash->{blacklist} = load_cidr_list($cfg->{blacklist_file}) if $cfg->{blacklist_file};

    # several service tunables are cached in lexicals for efficiency. if these
    # are changed, the "throttle reload config" command must be issued to
    # update the cache. this implements the reloading (and initial loading).
    my ($log, $path_regex, $method_regex);
    my $loader = $stash->{config_reloader} = sub {
        my @log_on_cfg = grep {length} split /[, ]+/, lc $cfg->{log_events};
        my @log_events = (0) x NUM_LOG_FLAGS;
        for (@log_on_cfg) {
            $log_events[LOG_BAN_ADDED]              = 1 if $_ eq 'ban';
            $log_events[LOG_BAN_REMOVED]            = 1 if $_ eq 'unban';
            $log_events[LOG_ALLOW_WHITELISTED]      = 1 if $_ eq 'whitelisted';
            $log_events[LOG_DENY_BANNED]            = 1 if $_ eq 'banned';
            $log_events[LOG_DENY_BLACKLISTED]       =
            $log_events[LOG_THROTTLE_BLACKLISTED]   = 1 if $_ eq 'blacklisted';
            $log_events[LOG_DENY_CONCURRENT]        = 1 if $_ eq 'concurrent';
            $log_events[LOG_THROTTLE_DEFAULT]       = 1 if $_ eq 'throttled';
            @log_events = (1) x NUM_LOG_FLAGS           if $_ eq 'all';
            @log_events = (0) x NUM_LOG_FLAGS           if $_ eq 'none';
        }

        $log = sub {};
        if (grep {$_} @log_events) {
            my $has_syslogger = eval { require Perlbal::Plugin::Syslogger; 1 };
            if ($has_syslogger && $cfg->{syslog_host}) {
                VERBOSE and Perlbal::log(info => "Using Perlbal::Plugin::Syslogger");
                $log = sub {
                    my $action = shift;
                    return unless $log_events[$action];
                    Perlbal::Plugin::Syslogger::send_syslog_msg($svc, @_);
                };
            }
            else {
                VERBOSE and Perlbal::log(warn => "Syslogger plugin unavailable, using Perlbal::log");
                $log = sub {
                    my $action = shift;
                    return unless $log_events[$action];
                    Perlbal::log(info => @_);
                };
            }
        }

        $path_regex   = $cfg->{path_regex}   ? qr/$cfg->{path_regex}/   : undef;
        $method_regex = $cfg->{method_regex} ? qr/$cfg->{method_regex}/ : undef;
    };
    $loader->();

    # structures for tracking IP states
    my %throttled;
    my %banned;
    my $store = Perlbal::Plugin::Throttle::Store->new($cfg);

    my $start_handler = sub {
        my $retval = eval {
            my $request_start = Time::HiRes::time;

            VERBOSE and Perlbal::log(info => "In Throttle (%s)",
                defined $DELAYED ? sprintf 'back after %.2fs', $DELAYED : 'initial'
            );

            my Perlbal::ClientProxy $cp = shift;
            unless ($cp) {
                VERBOSE and Perlbal::log(error => "Missing ClientProxy");
                return HANDLE_REQUEST;
            }

            my $headers = $cp->{req_headers};
            unless ($headers) {
                VERBOSE and Perlbal::log(info => "Missing headers");
                return HANDLE_REQUEST;
            }
            my $uri    = $headers->request_uri;
            my $method = $headers->request_method;

            my $ip = $cp->observed_ip_string() || $cp->peer_ip_string;
            unless (defined $ip) {
                # happens if client goes away
                VERBOSE and Perlbal::log(warn => "Client went away");
                $cp->send_response(500, "Internal server error.\n");
                return IGNORE_REQUEST;
            }

            # back from throttling, all later checks were already passed
            return HANDLE_REQUEST if defined $DELAYED;

            # increment the count of throttled conns
            $throttled{$ip}++;

            my $result = sub {
                # immediately passthrough whitelistees
                if ($stash->{whitelist} && $stash->{whitelist}->find($ip)) {
                    $log->(LOG_ALLOW_WHITELISTED, "Letting whitelisted ip $ip through");
                    return RESULT_ALLOW;
                }

                # drop conns from banned IPs
                if ($banned{$ip}) {
                    $log->(LOG_DENY_BANNED, "Denying banned IP $ip");
                    return RESULT_DENY;
                }

                # drop conns from banned/blacklisted IPs
                if ($stash->{blacklist} && $stash->{blacklist}->find($ip)) {
                    if ($cfg->{blacklist_action} eq 'deny') {
                        $log->(LOG_DENY_BLACKLISTED, "Denying blacklisted IP $ip");
                        return RESULT_DENY;
                    }
                    else {
                        $log->(LOG_THROTTLE_BLACKLISTED, "Throttling blacklisted IP $ip");
                        return RESULT_THROTTLE;
                    }
                }

                if (exists $throttled{$ip} && $throttled{$ip} > $cfg->{max_concurrent}) {
                    $log->(LOG_DENY_CONCURRENT, "Too many concurrent connections from $ip");
                    return RESULT_DENY;
                }

                # only throttle matching requests
                if (defined $path_regex && $uri !~ $path_regex) {
                    VERBOSE && Perlbal::log(info => "This isn't a throttled URL: %s", $uri);
                    return RESULT_ALLOW;
                }
                if (defined $method_regex && $method !~ $method_regex) {
                    VERBOSE && Perlbal::log(info => "This isn't a throttled method: %s", $method);
                    return RESULT_ALLOW;
                }

                return $cfg->{default_action} eq 'allow' ? RESULT_ALLOW : RESULT_THROTTLE;
            }->();

            if ($result == RESULT_DENY) {
                unless ($cfg->{log_only}) {
                    $cp->send_response(403, "Forbidden.\n");
                    return IGNORE_REQUEST;
                }
            }
            elsif ($result == RESULT_ALLOW) {
                return HANDLE_REQUEST;
            }

            # now entering throttle path...

            # check if we've seen this IP lately.
            my $key = $cfg->{instance_name} . $ip;
            $store->get(key => $key, timeout => $cfg->{initial_delay}, callback => sub {
                my ($last_request_time, $violations) = @_;
                $violations ||= 0;

                # do an early set to update the last_request_time and
                # expiration in case of early exit
                $store->set(
                    key     => $key,
                    start   => $request_start,
                    count   => $violations,
                    exptime => $cfg->{throttle_threshold_seconds},
                    timeout => $cfg->{initial_delay},
                );

                my $time_since_last_request;
                if (defined $last_request_time) {
                    $time_since_last_request = $request_start - $last_request_time;
                }

                VERBOSE and Perlbal::log(
                    info => "%s; this request at %.3f; last at %s; interval is %s",
                    $ip, $request_start,
                    $last_request_time || 'n/a', $time_since_last_request || 'n/a'
                );

                my $handle_after = sub {
                    my $delay = shift;
                    $delay = 0 if $cfg->{log_only};

                    # put request on the backburner
                    $cp->watch_read(0);
                    Danga::Socket->AddTimer($delay, sub {
                        # we're now executing in a timer callback after
                        # perlbal has been told to ignore the request. so if we
                        # now want it handled it needs to be re-adopted.
                        local $DELAYED = $delay; # to short-circuit throttling logic on the next pass through
                        $cp->watch_read(1);
                        $svc->adopt_base_client($cp);
                    });

                    return IGNORE_REQUEST;
                };

                # can we let it through immediately?
                unless (defined $time_since_last_request) {
                    # forgotten or haven't seen ip before
                    $log->(LOG_ALLOW_DEFAULT, "Allowed unseen $ip");
                    return $handle_after->(0);
                }
                if ($time_since_last_request >= $cfg->{throttle_threshold_seconds}) {
                    # waited long enough
                    $log->(LOG_ALLOW_DEFAULT, "Allowed reformed $ip");
                    return $handle_after->(0);
                }

                # need to throttle, now figure out by how much. at least
                # initial_delay, at most max_delay, exponentially increasing in
                # between
                my $delay = min($cfg->{initial_delay} * 2**$violations, $cfg->{max_delay});

                $violations++;

                # banhammer for great justice
                if ($cfg->{ban_threshold} && $violations >= $cfg->{ban_threshold}) {
                    $log->(LOG_BAN_ADDED, "Banning $ip for $cfg->{ban_expiration}s: %s", $uri);
                    $banned{$ip}++ unless $cfg->{log_only};
                    Danga::Socket->AddTimer($cfg->{ban_expiration}, sub {
                        $log->(LOG_BAN_REMOVED, "Unbanning $ip");
                        delete $banned{$ip};
                    });
                    $cp->close;
                    return IGNORE_REQUEST;
                }

                $store->set(
                    key     => $key,
                    start   => $request_start,
                    count   => $violations,
                    exptime => $delay,
                    timeout => $cfg->{initial_delay},
                );

                $log->(LOG_THROTTLE_DEFAULT, "Throttling $ip for $delay: %s", $uri);

                # schedule request to be re-handled
                return $handle_after->($delay);
            });

            # make sure we don't take up reading until readoption
            $cp->watch_read(0);
            return IGNORE_REQUEST;
        };
        if ($@) {
            # if something horrible should happen internally, don't take out perlbal
            Perlbal::log(err => "Throttle failed: '%s'", $@);
            return HANDLE_REQUEST;
        }
        else {
            return $retval;
        }
    };

    my $end_handler = sub {
        my Perlbal::ClientProxy $cp = shift;

        my $ip = $cp->observed_ip_string() || $cp->peer_ip_string;
        return unless $ip;

        delete $throttled{$ip} unless --$throttled{$ip} > 0;
    };

    $svc->register_hook(Throttle => start_proxy_request => $start_handler);
    $svc->register_hook(Throttle => end_proxy_request   => $end_handler);
}

sub load_cidr_list {
    my $file = shift;

    require Net::CIDR::Lite;

    my $empty = 1;
    my $list = Net::CIDR::Lite->new;

    open my $fh, '<', $file or die "Unable to open file $file: $!";
    while (my $line = <$fh>) {
        $line =~ s/#.*//; # comments
        if ($line =~ /([0-9\/\.]+)/) {
            my $cidr = $1;
            if (index($cidr, "/") < 0) {
                # slash-less specifications are assumed to be singular IPs
                $list->add_ip($cidr);
            }
            else {
                $list->add($cidr);
            }
            $empty = 0;
        }
    }

    die "$file contains no recognizable CIDRs\n" if $empty;

    return $list;
}

package Perlbal::Plugin::Throttle::Store;

sub new {
    my $class = shift;
    my $cfg = shift;

    my $want_memcached = $cfg->{memcached_servers};
    my $has_memcached = eval { require Cache::Memcached::Async; 1 };

    if ($want_memcached && !$has_memcached) {
        die "memcached support requested but Cache::Memcached::Async failed to load: $@\n";
    }
    return $want_memcached
        ? Perlbal::Plugin::Throttle::Store::Memcached->new($cfg)
        : Perlbal::Plugin::Throttle::Store::Memory->new($cfg);
}

package Perlbal::Plugin::Throttle::Store::Memcached;

sub new {
    my $class = shift;
    my $cfg = shift;

    my @servers = split /[,\s]+/, $cfg->{memcached_servers};
    my @cxns = map {
        Cache::Memcached::Async->new({ servers => \@servers })
    } 1 .. $cfg->{memcached_async_clients};

    return bless \@cxns, $class;
}

sub get {
    my $self = shift;
    my %p = @_;
    $self->[rand @$self]->get(
        $p{key},
        timeout => $p{timeout},
        callback => sub {
            my $value = shift;
            return $p{callback}->() unless $value;
            return $p{callback}->(unpack('FS', $value));
        },
    );
    return;
}

sub set {
    my $self = shift;
    my %p = @_;

    $self->[rand @$self]->set(
        $p{key} => pack('FS', $p{start}, $p{count}),
        exptime => $p{exptime},
        timeout => $p{timeout},
    );
}

package Perlbal::Plugin::Throttle::Store::Memory;

use Time::HiRes 'time';

sub new {
    my $class = shift;
    my $cfg = shift;
    return bless {}, $class;
}

sub get {
    my $self = shift;
    my %p = @_;
    my $entry = $self->{$p{key}};

    return $p{callback}->($entry->[1], $entry->[2]) if $entry && time < $entry->[0];
    return $p{callback}->();
}

sub set {
    my $self = shift;
    my %p = @_;
    $self->{$p{key}} = [time + $p{exptime}, $p{start}, $p{count}];
    return;
}

1;

__END__