File Coverage

blib/lib/RateLimitations.pm

Criterion	Covered	Total	%
statement	70	110	63.6
branch	5	16	31.2
condition	5	17	29.4
subroutine	14	20	70.0
pod	6	6	100.0
total	100	169	59.1

package RateLimitations;

use strict;
use warnings;
our $VERSION = '0.04';

use Carp;

use Cache::RedisDB;
use Config::Onion;
use List::Util qw(all);
use Time::Duration::Concise;

use base qw( Exporter );
our @EXPORT_OK = qw(
    all_service_consumers
    flush_all_service_consumers
    rate_limited_services
    rate_limits_for_service
    verify_rate_limitations_config
    within_rate_limits
);
use constant KEYSPACE  => 'RATELIMITATIONS';    # Everything will fall under this
use constant SEPARATOR => '::';                 # How to join strings together

my $rates_file_content;
my %limits;

BEGIN {
    my $cfg = Config::Onion->new;
    $cfg->set_default(
        rl_internal_testing => {
            '10s' => 2,
            '5m'  => 6
        });
    $cfg->load('/etc/perl_rate_limitations', '/etc/rmg/perl_rate_limitations');
    $rates_file_content = $cfg->get;
    foreach my $svc (sort keys %$rates_file_content) {
        my @service_limits;
        foreach my $time (keys %{$rates_file_content->{$svc}}) {
            my $ti = Time::Duration::Concise->new(interval => $time);
            my $count = $rates_file_content->{$svc}->{$time};
            push @service_limits, [$ti->seconds, $count];
        }
        @service_limits = sort { $a->[0] <=> $b->[0] } @service_limits;
        $limits{$svc} = {
            rates   => \@service_limits,
            seconds => $service_limits[-1]->[0],
            entries => $service_limits[-1]->[1],
        };
    }
}

sub verify_rate_limitations_config {
    my $proper = 1;    # Assume it is proper until we find a bad entry
    foreach my $svc (sort keys %$rates_file_content) {
        my @service_limits;
        foreach my $time (keys %{$rates_file_content->{$svc}}) {
            my $ti = Time::Duration::Concise->new(interval => $time);
            my $count = $rates_file_content->{$svc}->{$time};
            push @service_limits, [$ti->seconds, $count, $count / $ti->seconds, undef, $time];
        }
        @service_limits = sort { $a->[0] <=> $b->[0] } @service_limits;

        my @final_limits;
        while (my $this_limit = shift @service_limits) {
            my ($improper, $index) = ($this_limit->[3], $#service_limits);
            while (not $improper and $index > -1) {
                my $that_limit = $service_limits[$index];
                # This one is improper if that longer period has the same or smaller count
                $improper = 'count should be lower than ' . $that_limit->[4] . ' count' if ($that_limit->[1] <= $this_limit->[1]);
                # That one is improper if this shorter period has the smaller rate
                $service_limits[$index]->[3] = 'rate should be lower than ' . $this_limit->[4] . ' rate'
                    if (not $improper and $this_limit->[2] < $that_limit->[2]);
                $index--;
            }
            if ($improper) {
                # If any entry is improper we will fail and warn.
                # We still check the rest for completeness
                $proper = 0;
                carp $svc . ' - ' . $this_limit->[4] . ' entry improper: ' . $improper;
            }
        }
    }
    return $proper;
}

sub within_rate_limits {
    my $args = shift;

    croak 'Must supply args as a hash reference' unless ref $args eq 'HASH';
    my ($service, $consumer) = @{$args}{'service', 'consumer'};
    croak 'Must supply both "service" and "consumer" arguments' unless all { defined } ($service, $consumer);
    my $limit = $limits{$service};
    croak 'Unknown service supplied: ' . $service unless $limit;

    my $redis         = Cache::RedisDB->redis;
    my $key           = _make_key($service, $consumer);
    my $within_limits = 1;
    my $now           = time;
    $redis->lpush($key, $now);    # We push first so that we hit limits more often in heavy (DoS) conditions
    $redis->ltrim($key, 0, $limit->{entries});    # Our new entry is now in index 0.. we keep 1 extra entry.
    $redis->expire($key, $limit->{seconds});
    foreach my $rate (@{$limit->{rates}}) {
        if (($redis->lindex($key, $rate->[1]) // 0) > $now - $rate->[0]) {
            $within_limits = 0;
            last;
        }
    }

    return $within_limits;
}

sub flush_all_service_consumers {
    my $redis = Cache::RedisDB->redis;
    my $count = 0;

    foreach my $key (_all_keys($redis)) {
        $count += $redis->del($key);
    }

    return $count;
}

sub _all_keys { my $redis = shift // Cache::RedisDB->redis; return @{$redis->keys(_make_key('*', '*')) // []}; }

sub rate_limited_services { return (sort keys %limits); }

sub rate_limits_for_service {
    my $service = shift // 'undef';
    my $svc_limits = $limits{$service};
    croak 'Unknown service supplied: ' . $service unless $svc_limits;

    return @{$svc_limits->{rates}};
}

sub all_service_consumers {

    my %consumers;

    foreach my $pair (map { [(split SEPARATOR, $_)[-2, -1]] } _all_keys()) {
        $consumers{$pair->[0]} //= [];
        push @{$consumers{$pair->[0]}}, $pair->[1];
    }

    return \%consumers;
}

sub _make_key {
    my ($service, $consumer) = @_;

    return join(SEPARATOR, KEYSPACE, $service, $consumer);
}

1;
__END__