File Coverage

blib/lib/Perlbal/Plugin/SessionAffinity.pm

Criterion	Covered	Total	%
statement	30	190	15.7
branch	0	60	0.0
condition	0	30	0.0
subroutine	9	25	36.0
pod	8	10	80.0
total	47	315	14.9

use strict;
use warnings;
package Perlbal::Plugin::SessionAffinity;
# ABSTRACT: Sane session affinity (sticky sessions) for Perlbal
$Perlbal::Plugin::SessionAffinity::VERSION = '0.100';
use Perlbal;
use Hash::Util;
use CGI::Cookie;
use MIME::Base64;
use Digest::MD5 'md5';
use Digest::SHA 'sha1_hex';

my $default_cookie_hdr = 'X-SERVERID';
my $cookie_hdr_sub     = sub { encode_base64( md5( $_[0] ) ) };
my $salt               = join q{}, map { $_ = rand 999; s/\.//; $_ } 1 .. 10;
my $use_salt           = 0;
my $use_domain         = 0;
my $use_dynamic_cookie = 0;

sub get_domain_from_req {
    my $req    = shift;
    my $domain = ref $req eq 'Perlbal::XS::HTTPHeaders' ?
                 $req->getHeader('host')                : # XS version
                 $req->{'headers'}{'host'};               # PP version

    return $domain;
}

# get the ip and port of the requested backend from the cookie
sub get_ip_port {
    my ( $svc, $req ) = @_;

    my $domain     = get_domain_from_req($req);
    my $cookie_hdr = $use_dynamic_cookie        ?
                     $cookie_hdr_sub->($domain) :
                     $default_cookie_hdr;

    my $cookie     = $req->header('Cookie');
    my %cookies    = ();

    if ( defined $cookie ) {
        %cookies = CGI::Cookie->parse($cookie);

        if ( defined $cookies{$cookie_hdr} ) {
            my $id      = $cookies{$cookie_hdr}->value || '';
            my $backend = find_backend_by_id( $svc, $id );

            ref $backend and return join ':', @{$backend};
        }
    }

    return;
}

# create a domain ID
sub create_domain_id {
    my $domain = shift || '';
    my @nodes  = @_;

    # the ID is determined by the specific server
    # that has the matching index for the domain
    my $index = domain_index( $domain, scalar @nodes );
    my $node  = join ':', @{ $nodes[$index] };
    return sha1_hex( $use_salt ? $salt . $node : $node );
}

# create an id from ip and optional port
sub create_id {
    my $ip   = shift;
    my $port = shift || '';
    my $str  = $use_salt ? $salt . "$ip:$port" : "$ip:$port";
    return sha1_hex($str);
}

# a nifty little trick:
# we create a numeric value of the domain name
# then we use that as a seed for the random function
# then create a random number which is predictable
# that is the index of the domain
sub domain_index {
    my $domain = shift;
    my $max    = shift;
    my $seed   = 0;

    foreach my $char ( split //, $domain ) {
        $seed += ord $char;
    }

    return ( $seed % $max);
}

# using an sha1 checksum id, find the matching backend
sub find_backend_by_id {
    my ( $svc, $id ) = @_;

    foreach my $backend ( @{ $svc->{'pool'}{'nodes'} } ) {
        my $backendid = create_id( @{$backend} );

        if ( $backendid eq $id ) {
            return $backend;
        }
    }

    return;
}

# TODO: refactor this function
sub find_backend_by_domain_id {
    my ( $svc, $id ) = @_;

    foreach my $backend ( @{ $svc->{'pool'}{'nodes'} } ) {
        my $backendid = create_id( @{$backend} );

        if ( $backendid eq $id ) {
            return $backend;
        }
    }

    return;
}

sub load {
    # the name of header in the cookie that stores the backend ID
    Perlbal::register_global_hook(
        'manage_command.affinity_cookie_header', sub {
            my $mc = shift->parse(qr/^\s*affinity_cookie_header\s+=\s+(.+)\s*$/,
                      "usage: AFFINITY_COOKIE_HEADER = ");

            ($default_cookie_hdr) = $mc->args;

            return $mc->ok;
        },
    );

    Perlbal::register_global_hook(
        'manage_command.affinity_salt', sub {
            my $mc = shift->parse(qr/^\s*affinity_salt\s+=\s+(.+)\s*$/,
                      "usage: AFFINITY_SALT = ");

            ($salt) = $mc->args;

            return $mc->ok;
        },
    );

    Perlbal::register_global_hook(
        'manage_command.affinity_use_salt', sub {
            my $mc = shift->parse(qr/^\s*affinity_use_salt\s+=\s+(.+)\s*$/,
                      "usage: AFFINITY_USE_SALT = ");

            my ($res) = $mc->args;
            if ( $res eq 'yes' || $res == 1 ) {
                $use_salt = 1;
            } elsif ( $res eq 'no' || $res == 0 ) {
                $use_salt = 0;
            } else {
                die qq"affinity_use_salt must be boolean (yes/no/1/0)";
            }

            return $mc->ok;
        },
    );

    Perlbal::register_global_hook(
        'manage_command.affinity_use_domain', sub {
            my $mc = shift->parse(qr/^\s*affinity_use_domain\s+=\s+(.+)\s*$/,
                      "usage: AFFINITY_USE_DOMAIN = ");

            my ($res) = $mc->args;
            if ( $res eq 'yes' || $res == 1 ) {
                $use_domain = 1;
            } elsif ( $res eq 'no' || $res == 0 ) {
                $use_domain = 0;
            } else {
                die qq"affinity_use_domain must be boolean (yes/no/1/0)";
            }

            return $mc->ok;
        },
    );

    Perlbal::register_global_hook(
        'manage_command.affinity_use_dynamic_cookie', sub {
            my $mc = shift->parse(qr/^\s*affinity_use_dynamic_cookie\s+=\s+(.+)\s*$/,
                      "usage: AFFINITY_USE_DYNAMIC_COOKIE = ");

            my ($res) = $mc->args;
            if ( $res eq 'yes' || $res == 1 ) {
                $use_dynamic_cookie = 1;
            } elsif ( $res eq 'no' || $res == 0 ) {
                $use_dynamic_cookie = 0;
            } else {
                die qq"affinity_use_dynamic_cookie must be boolean (yes/no/1/0)";
            }

            return $mc->ok;
        },
    );

    return 1;
}

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

    my $check_cookie = sub {
        my $client = shift;
        my $req    = $client->{'req_headers'} or return 0;
        my $svc    = $client->{'service'};
        my $pool   = $svc->{'pool'};

        # make sure all nodes in this service have their own pool
        foreach my $node ( @{ $pool->{'nodes'} } ) {
            my ( $ip, $port ) = @{$node};

            # pool
            my $poolid = create_id( $ip, $port );
            exists $Perlbal::pool{$poolid} and next;

            my $nodepool = Perlbal::Pool->new($poolid);
            $nodepool->add( $ip, $port );
            $Perlbal::pool{$poolid} = $nodepool;

            # service
            my $serviceid = "${poolid}_service";
            exists $Perlbal::service{$serviceid} and next;

            my $nodeservice = Perlbal->create_service($serviceid);
            my $svc_role    = $svc->{'role'};

            # role sets up constraints for the rest
            # so it goes first
            $nodeservice->set( role => $svc_role );

            foreach my $tunable_name ( keys %{$Perlbal::Service::tunables} ) {
                # skip role because we had already set it
                $tunable_name eq 'role' and next;

                # persist_client_timeout is DEPRECATED
                # but not marked anywhere as deprecated. :(
                # (well, nowhere we can actually predictably inspect)
                $tunable_name eq 'persist_client_timeout' and next;

                # we skip the pool because we're gonna set it to a specific one
                $tunable_name eq 'pool' and next;

                # make sure svc has value for this tunable
                defined $svc->{$tunable_name} or next;

                my $tunable = $Perlbal::Service::tunables->{$tunable_name};
                my $role    = $tunable->{'check_role'};

                if ( $role eq '*' || $role eq $svc_role ) {
                    $nodeservice->set( $tunable_name, $svc->{$tunable_name} );
                }
            }

            # restricted hashes are stupid
            # so we have to use the API to add them
            foreach my $hook_name ( keys %{ $svc->{'hooks'} } ) {
                foreach my $set ( @{ $svc->{'hooks'}{$hook_name} } ) {
                    my ( $plugin, $sub ) = @{$set};
                    $nodeservice->register_hook( $plugin, $hook_name, $sub );
                }
            }

            # add all the extra config and extra headers
            $nodeservice->{'extra_config'}  = $svc->{'extra_config'};
            $nodeservice->{'extra_headers'} = $svc->{'extra_headers'};

            $nodeservice->set( pool => $poolid );

            $Perlbal::service{$serviceid} = $nodeservice;
        }

        my $ip_port = get_ip_port( $svc, $req );

        if ( ! $ip_port ) {
            $use_domain or return 0;

            # we're going to override whatever Perlbal found
            # because we only care about the domain
            my $domain = get_domain_from_req($req);

            my @ordered_nodes = sort {
                ( join ':', @{$a} ) cmp ( join ':', @{$b} )
            } @{ $svc->{'pool'}{'nodes'} };

            my $id      = create_domain_id( $domain, @ordered_nodes );
            my $backend = find_backend_by_domain_id( $svc, $id );
            $ip_port = join ':', @{$backend};
        }

        my ( $ip, $port )    = split /:/, $ip_port;
        my $req_pool_id      = create_id( $ip, $port );
        my $req_svc          = $Perlbal::service{"${req_pool_id}_service"};
        $client->{'service'} = $req_svc;

        return 0;
    };

    my $set_cookie = sub {
        my $backend = shift; # Perlbal::BackendHTTP

        defined $backend or return 0;

        my $res        = $backend->{'res_headers'};
        my $req        = $backend->{'req_headers'};
        my $svc        = $backend->{'service'};
        my $backend_id = create_id( split /:/, $backend->{'ipport'} );
        my $domain     = get_domain_from_req($req);
        my $cookie_hdr = $use_dynamic_cookie        ?
                         $cookie_hdr_sub->($domain) :
                         $default_cookie_hdr;

        my %cookies = ();
        if ( my $cookie = $req->header('Cookie') ) {
            %cookies = CGI::Cookie->parse($cookie);
        }

        if ( ! defined $cookies{$cookie_hdr} ||
             $cookies{$cookie_hdr}->value ne $backend_id ) {

            my $backend_cookie = CGI::Cookie->new(
                -name  => $cookie_hdr,
                -value => $backend_id,
            );

            if ( defined $res->header('set-cookie') ) {
                my $value = $res->header('set-cookie') .
                            "\r\nSet-Cookie: "         .
                            $backend_cookie->as_string;

                $res->header( 'Set-Cookie' => $value );
            } else {
                $res->header( 'Set-Cookie' => $backend_cookie->as_string );
            }
        }

        return 0;
    };

    $gsvc->register_hook(
        'SessionAffinity', 'start_proxy_request', $check_cookie,
    );

    $gsvc->register_hook(
        'SessionAffinity', 'backend_response_received', $set_cookie,
    );

    return 1;
}

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

    # TODO: are we using setters?
    $svc->unregister_hooks('SessionAffinity');
    $svc->unregister_setters('SessionAffinity');

    return 1;
}

1;

__END__