File Coverage

blib/lib/Apache/Quota.pm

Criterion	Covered	Total	%
statement	4	6	66.6
branch			n/a
condition			n/a
subroutine	2	2	100.0
pod			n/a
total	6	8	75.0

package Apache::Quota;

use strict;

use Apache;
use Apache::Constants qw( OK DECLINED FORBIDDEN HTTP_OK );
use Fcntl qw(O_CREAT O_RDWR O_RDONLY);
use File::Path ();
use File::Temp ();

use vars qw ( $VERSION $DEBUG );

$VERSION = 0.04;

use constant KB => 1024;
use constant MB => 1024 ** 2;
use constant GB => 1024 ** 3;

use constant SECOND => 1;
use constant MINUTE => 60;
use constant HOUR   => 60 * 60;
use constant DAY    => 24 * 60 * 60;

sub handler
{
    my $r = shift;

    return DECLINED if $r->header_only;

    local $DEBUG = $r->dir_config('QuotaDebug') ? 1 : 0;
    $r->notes( 'Apache::Quota::DEBUG' => 1 ) if $DEBUG;

    my $main = $r->is_main ? $r : $r->main;

    unless ( $main->notes('Apache::Quota::initialized') )
    {
        my $file = $r->dir_config('QuotaFile');

        die "Cannot use Apache::Quota unless QuotaFile is specified"
            unless defined $file;

        $main->notes( 'Apache::Quota::file' => $file );

        my $locker;
        {
            my $class;

            if ( $r->dir_config('QuotaLocker') )
            {
                $class = $r->dir_config('QuotaLocker');
                $class = "Apache::Quota::$class"
                    unless $class =~ /^Apache::Quota::/;

                eval "require $class";
                die $@ if $@;
            }
            else
            {
                for my $c ( qw( Apache::Quota::BerkeleyDB Apache::Quota::DB_File::Lock ) )
                {
                    eval "require $_";
                    $class = $c unless $@;
                }

                die "Cannot load Apache::Quota::BerkeleyDB or Apache::Quota::DB_File::Lock"
                    unless $class;
            }

            _log( $r, "locker is $class" ) if $DEBUG;

            $locker = $class;
            $main->notes( 'Apache::Quota::locker' => $locker );
        }


        _log( $r, "will record usage in $file for request for " . $r->uri )
            if $DEBUG;

        my $type = lc $r->dir_config('QuotaType') || 'global';

        _log( $r, "limit type is $type" )
            if $DEBUG;

        my $key  = $r->dir_config('QuotaLocationKey') || 'Apache-Quota global key';

        _log( $r, "key is $key" )
            if $DEBUG;

        die "Your location key cannot contain a colon" if $key =~ /:/;

        my $full_key = ( $type eq 'client-ip'
                         ? join ':', 'client-ip', $r->connection->remote_ip, $key
                         : $type eq 'sub'
                         ? join ':', 'sub', _key_from_sub($r), $key
                         : $key
                       );

        _log( $r, "full key is $full_key" )
            if $DEBUG;

        $main->notes( 'Apache::Quota::full_key' => $full_key );

        my $limit = _get_limit($r);

        _log( $r, "limit is $limit bytes" )
            if $DEBUG;

        my $period = _get_period($r);

        _log( $r, "period is $period seconds" )
            if $DEBUG;

        my $db = $locker->_open_db( file => $file, mode => 'read' );

        my $expired = time - $period;
        $main->notes( 'Apache::Quota::expired' => $expired );

        if ( exists $db->{$full_key} )
        {
            _log( $r, "found key ($full_key) in DB file" )
                if $DEBUG;

            $main->notes( 'Apache::Quota::exceeds' =>
                          _check_against_limit( $r, $expired, $limit, $db->{$full_key} )
                        );
        }

        untie %$db;
    }

    $main->notes( 'Apache::Quota::initialized' => 1 );

    if ( $main->notes('Apache::Quota::exceeds') )
    {
        my $on_exceed = lc $r->dir_config('QuotaOnExceed') || 'deny';

        _log( $r, "on exceed is $on_exceed" )
            if $DEBUG;

        if ( $on_exceed eq 'notes' )
        {
            _log( $r, "setting \$r->notes() value for Apache::Quota::exceeded to 1" )
                if $DEBUG;

            $r->notes( 'Apache::Quota::exceeded' => 1 );
        }
        else
        {
            _log( $r, "returning FORBIDDEN" )
                if $DEBUG;

            return FORBIDDEN;
        }
    }
    else
    {
        _log( $r, "registering cleanup" )
            if $DEBUG;

        $r->register_cleanup( \&_record_usage );
    }

    return OK;
}

sub _key_from_sub
{
    my $r = shift;

    my $sub = $r->dir_config('QuotaSub');

    die "Cannot set QuotaType to sub if QuotaSub is not set"
        unless $sub;

    _log( $r, "sub is $sub" )
        if $DEBUG;

    my $key;
    {
        no strict 'refs';
        $key = &{$sub}($r);
    }

    return $key;
}

my %limit_letter_to_number = ( k => KB, m => MB, g => GB );
sub _get_limit
{
    my $r = shift;

    my $limit = $r->dir_config('QuotaLimit');

    die "Cannot use Apache::Quota unless QuotaLimit is specified"
        unless defined $limit;

    _log( $r, "limit set to $limit in config file" )
        if $DEBUG;

    my ( $num, $letter ) = $limit =~ /(\d+)([kmg])/i;
    $letter ||= 'k';

    die "Invalid limit in QuotaLimit: $limit"
        unless exists $limit_letter_to_number{lc $letter};

    return $num * $limit_letter_to_number{ lc $letter };
}

sub _get_period
{
    my $r = shift;

    my $period = $r->dir_config('QuotaPeriod');

    die "Cannot use Apache::Quota unless QuotaPeriod is specified"
        unless defined $period;

    _log( $r, "period set to $period in config file" )
        if $DEBUG;

    return _parse_period($period);
}

my %period_letter_to_number = ( s => SECOND, m => MINUTE, h => HOUR, d => DAY );
sub _parse_period
{
    my $period = shift;

    my ( $num, $letter ) = $period =~ /(\d+)([smhd])/i;
    $letter ||= 's';

    die "Invalid period: $period"
        unless exists $period_letter_to_number{ lc $letter };

    return $num * $period_letter_to_number{ lc $letter };
}

sub _check_against_limit
{
    my ( $r, $expired, $limit, $current ) = @_;

    my $total = 0;
    foreach my $record ( split /;/, $current )
    {
        my ( $time, $bytes ) = split /:/, $record;

        _log( $r, "key has record: $time - $bytes" )
            if $DEBUG;

        # These will be removed during cleanup, but for now we'll
        # avoid writing to the file to try to speed things up a bit.
        next unless $time > $expired;

        $total += $bytes;
    }

    _log( $r, "key has total bytes of $total" )
        if $DEBUG;

    if ( $total > $limit )
    {
        _log( $r, "total bytes ($total) exceeds limit ($limit)" )
            if $DEBUG;

        return 1;
    }

    return 0;
}

sub _record_usage
{
    my $r = shift;

    return 0 unless $r->status == HTTP_OK;

    return 0 unless $r->bytes_sent;

    _log( $r, "status was HTTP_OK, recording traffic" )
        if $DEBUG;

    my $db =
        $r->notes('Apache::Quota::locker')->_open_db
            ( file => $r->notes('Apache::Quota::file'), mode => 'write' );

    my $key = $r->notes('Apache::Quota::full_key');
    my $expired = $r->notes('Apache::Quota::expired');

    my @records;
    if ( exists $db->{$key} )
    {
        foreach my $record ( split /;/, $db->{$key} )
        {
            my ( $time, $bytes ) = split /:/, $record;

            next unless $time > $expired;

            push @records, "$time:$bytes";
        }
    }

    push @records, join ':', time, $r->bytes_sent;

    _log( $r, "adding record of " . $r->bytes_sent . " bytes sent" )
        if $r->notes('Apache::Quota::DEBUG');

    $db->{$key} = join ';', @records;

    untie %$db;

    return 1;
}

sub usage
{
    my %p = @_;

    my $locker = $p{locker};
    $locker = "Apache::Quota::$locker"
        unless $locker =~ /^Apache::Quota::/;

    eval "require $locker";
    die $@ if $@;

    my $db = $locker->_open_db( file => $p{file}, mode => 'read' );

    my $expired = 0;
    if ( defined $p{period} )
    {
        $expired = time - _parse_period( $p{period} );
    }

    my %vals;
    foreach my $key ( keys %$db )
    {
        my $total = 0;
        if ( defined $db->{$key} )
        {
            foreach my $record ( split /;/, $db->{$key} )
            {
                my ( $time, $bytes ) = split /:/, $record;

                next unless $time > $expired;

                $total += $bytes;
            }
        }

        next unless $total;

        my %extra;
        if ( $key =~ /^client-ip:([\d\.]+)/ )
        {
            $extra{ip} = $1;
        }
        elsif ( $key =~ /^sub:(.+):[^:]+$/ )
        {
            $extra{sub} = $1;
        }

        $vals{$key} = { %extra, bytes => $total };
    }

    untie %$db;

    return %vals;
}

sub set_key
{
    my %p = @_;

    die "Cannot call set_key without a key parameter"
        unless defined $p{key};

    my $locker = $p{locker};
    $locker = "Apache::Quota::$locker"
        unless $locker =~ /^Apache::Quota::/;

    eval "require $locker";
    die $@ if $@;

    my $db = $locker->_open_db( file => $p{file}, mode => 'write' );

    return 0 unless exists $db->{ $p{key} };

    if ( defined $p{value} )
    {
        die "Cannot set quota to a non-numeric value"
            unless $p{value} =~ /^\d+$/;

        $db->{ $p{key} } = time . ":$p{value}";
    }
    else
    {
        delete $db->{ $p{key} };
    }

    untie %$db;

    return 1;
}

sub reset_key { set_key( @_, value => undef ) }

sub _log
{
    my $r = shift;
    my $msg = '[Apache::Quota debug] ' . shift;
    $msg .= "\n";

    warn($msg);
}

1;

__END__