#!/usr/bin/perl -w
# HTB Hashiz version 0.3
# Rate limiting htb with hashing filters by Dmitriy Kozlov <dimka@freelsd.net>
# Hints and idea by Michael Ulitskiy from LARTC list with subject "u32 clarification...limits on 2000>???"
# Reference http://mailman.ds9a.nl/pipermail/lartc/2003q2/009109.html
#
# Use:
# ./htb-hashiz.pl [options]
# . rc.qos.eth1
# 

use strict;
use Socket;
use Getopt::Long;


my $NETWORKS = '169.254.0.0/20 10.10.0.0/20 192.168.0.0/22';
my $EXCLUDE = '169.254.0.0 169.254.0.0/24 10.10.0.0-10.10.0.20 10.10.1.0 10.10.2.0/24';
my $IF = 'eth1';	# output interface
my $RATE = '512Kbit';   # default rate 
my $MTU = 'mtu 16500';  # for some eth cards where TSO is ON
my $DIRECT = 'dst';	# 'dst' or 'src'

my $opt_v;
my $opt_help;

usage() unless GetOptions(
  "networks|n=s"      => \$NETWORKS,
  "exclude|x=s"       => \$EXCLUDE,
  "interface|i=s"     => \$IF,
  "target|t=s"        => \$DIRECT,
  "rate|r=s"          => \$RATE,
  "mtu|m=s"           => \$MTU,
  "verbose|v"         => \$opt_v,
  "help|h|?"          => \$opt_help
);
usage() if $opt_help;


# ... no more changes =)

my $OUTFILE="rc.qos.$IF";
my $INFFILE="rc.qos.$IF.helper";
my $qhandle = 0x1000;	# generated qdisc handle started from
my $fhandle = 2;	# genereted filter handle startet from

my $DIRPOS;
if ($DIRECT eq 'dst') {
  $DIRPOS = 16 if ($DIRECT eq 'dst'); # dst. for hashkey mask
} elsif ($DIRECT eq 'src') {
  $DIRPOS = 12 if ($DIRECT eq 'src'); # dst. for hashkey mask
} else {
  die("Wrong direction $DIRECT");
}
print "$DIRECT at $DIRPOS\n" if ($opt_v);

sub init($);
sub exclude($);


open(FH,">$OUTFILE") or die("Can't create $OUTFILE: $!");
open(FH2, ">$INFFILE") or die("Can't create $INFFILE: $!");

print FH <<EOF;
#!/bin/bash

# show
if [ "\$1" = "show" ]; then
  tc -s qdisc show dev $IF
  tc -s class show dev $IF
  tc -s filter show dev $IF
  tc -s filter show dev $IF parent 12:
  exit;
fi

# clean
tc qdisc del dev $IF root 2>/dev/null

if [ "\$1" = "stop" ]; then
  echo "Stopped."
  exit
fi

echo "Started."
# root
tc qdisc add dev $IF root handle 1: prio bands 3

tc qdisc add dev $IF parent 1:1 handle 10: pfifo
tc qdisc add dev $IF parent 1:2 handle 11: pfifo
tc qdisc add dev $IF parent 1:3 handle 12: htb default 0


EOF

my @nets = split(' ', $NETWORKS);
my $count_ips = 0;
foreach (@nets) {
  print $_ . "\n";
  my ($addr, $addr_end, $mask, $prefix, $counts) = init($_);
  my ($divisor1, $divisor2);
  my $link1;
  my $link2;
  my $hashmask1;
  my $hashmask2;
  my $saddr;

  $count_ips += $counts;
  if ($count_ips > (0xffff - 0x1000 - 1)) {  # 0..0x1000 - reserved by me, 0xffff - reserved for ingress qdisc
    die("Too much networks.");
  }
  
  if ($prefix >= 24) {  # single hash
    $divisor1 = 0;
    $divisor2 = 2 ** (32-$prefix);
    $link1 = 0;
    $link2 = $fhandle++;
    $hashmask1 = 0;
    $hashmask2 = sprintf("0x%08x", ~$mask & 0x000000FF);
  } elsif ($prefix >= 17) {
    $divisor1 = 2 ** (24-$prefix);
    $divisor2 = 256;
    $link1 = $fhandle++;
    $link2 = $fhandle++;
    $hashmask1 = sprintf("0x%08x", ~$mask & 0x0000FF00);
    $hashmask2 = sprintf("0x%08x", ~$mask & 0x000000FF);
  } else {
    die("Bad prefix $prefix. Network must be not more then /17\n");
  }

  print "prefix = $prefix\n" if ($opt_v);
  print "divisor1 = $divisor1\n" if ($opt_v);
  print "divisor2 = $divisor2\n" if ($opt_v);
  print "link1 = $link1\n" if ($opt_v);
  print "link2 = $link2\n" if ($opt_v);
  print "hashmask1 = $hashmask1\n" if ($opt_v);
  print "hashmask2 = $hashmask2\n" if ($opt_v);
  $saddr = pack('N',$addr);
  $saddr = inet_ntoa($saddr);

  printf FH "# init hash filters for $saddr/$prefix network\n";
  printf FH "tc filter add dev $IF parent 12: prio 7 handle $link1: protocol ip u32 divisor $divisor1\n" if ($divisor1);
  printf FH "tc filter add dev $IF parent 12: prio 7 handle $link2: protocol ip u32 divisor $divisor2\n";
  printf FH "tc filter add dev $IF protocol ip parent 12: prio 7 u32 \\\n" if ($divisor1);
  printf FH "match ip $DIRECT $saddr/$prefix hashkey mask $hashmask1 at $DIRPOS link $link1:\n" if ($divisor1);
  printf FH "tc filter add dev $IF protocol ip parent 12: prio 7 u32 \\\n";
  printf FH "match ip $DIRECT $saddr/$prefix hashkey mask $hashmask2 at $DIRPOS link $link2:\n";
  printf FH "\n\n";

  printf FH "# fill out hash\n";
  for (my $i = $addr; $i <= $addr_end; $i++) {
    my $si = pack('N',$i);
    $si = inet_ntoa($si);

    if (exclude($si)) {
      print "$si excluded\n" if ($opt_v);
      next;
    }; 

    my $qh = $qhandle++;
    die("End of space for qdisc handles. Too big ip range or to much networks.") if ($qhandle >= 0xffff); 
    my $sqh = sprintf("%04x", $qh);
    my $key1 = sprintf("%02x", ($i & 0x000ff00) >> 8);
    my $key2 = sprintf("%02x", $i & 0x00000ff); 

    printf FH "tc class add dev $IF parent 12: classid 12:$sqh htb rate $RATE ceil $RATE quantum 1514 $MTU\n";
    printf FH "tc qdisc add dev $IF parent 12:$sqh handle $sqh: sfq perturb 5\n";
    if ($divisor1) {
      printf FH "tc filter add dev $IF protocol ip parent 12: prio 7 u32 ht $link1:$key1: ht $link2:$key2: \\\n";
    } else {
      printf FH "tc filter add dev $IF protocol ip parent 12: prio 7 u32 ht $link2:$key2: \\\n";
    }
    printf FH "match ip $DIRECT $si/32 flowid 12:$sqh\n";
    printf FH "\n";
    
    printf FH2 "$si $IF 12:$sqh \"tc class change dev $IF parent 12: classid 12:$sqh htb rate #RATE# ceil #RATE# quantum 1514 $MTU\"\n";
  }

  printf FH "# fill out htb root by packets (for this network)\n";
  printf FH "tc filter add dev $IF parent 1:0 protocol ip prio 4 u32 match ip $DIRECT $saddr/$prefix flowid 1:3\n";


}

close(FH);
close(FH2);



sub exclude($) {
  my $saddr = shift;
  my $addr = unpack('N',inet_aton($saddr));
  my ($exstart, $exend, $exprefix, $exmask, $exaddr);

#  print "check $saddr";
  return 0 if (!defined($EXCLUDE));

  foreach my $net (split(' ',$EXCLUDE)) {
#    printf "  $net\n";

    if ($net =~ /^(\d+\.\d+\.\d+\.\d+)\-(\d+\.\d+\.\d+\.\d+)$/) {   # range
      die("Bad exclude $net") if (!inet_aton($1) or !inet_aton($2));
      $exstart = unpack('N',inet_aton($1));
      $exend = unpack('N',inet_aton($2));
      die("Bad exclude range $net") if ($exstart >= $exend);
      return 1 if ($addr >= $exstart and $addr <= $exend);
    } elsif ($net =~ /^(\d+\.\d+\.\d+\.\d+)\/(\d+)$/) {	            # ip/prefix
      die("Bad exclude $net") if (!inet_aton($1));
      $exstart = unpack('N',inet_aton($1));
      $exprefix = $2;
      die("Bad prefix $exprefix") if ($exprefix < 0 or $exprefix > 32);
      $exmask = 0xFFFFFFFF << (32 - $exprefix);
      $exstart &= $exmask;
      $exend = $exstart | ~$exmask;
      return 1 if ($addr >= $exstart and $addr <= $exend);
    } elsif ($net =~ /^(\d+\.\d+\.\d+\.\d+)$/) {                    # ip
      die("Bad exclude $net") if (!inet_aton($1));
      $exaddr = unpack('N',inet_aton($1));
      return 1 if ($addr == $exaddr);
    } else {
      die("Bad exclude syntax $net\n");
    }
  }
  return 0;
}

sub init($) {
  my $net = shift;
  my $naddr;
  my $nmask;
  my $prefix; 
  my $addr;
  my $mask;
  my $naddr_end;
  my $addr_end;
   
  if ($net =~ /^(\d+\.\d+\.\d+\.\d+)\/(\d+)$/) {
    $prefix = $2;
    if ($prefix < 17 or $prefix > 30) {
      die("Bad prefix $prefix\n");
    }
    $mask = 0xFFFFFFFF << (32 - $prefix);
    $nmask = pack('N',$mask);

    $naddr = inet_aton($1);
    $addr = unpack('N', $naddr);
    $addr &= $mask;
    $addr_end = $addr | ~$mask;
    $naddr_end = pack('N',$addr_end);
    $naddr = pack('N',$addr);

    print "  " . inet_ntoa($naddr) . "\n" if ($opt_v);
    print "  " . inet_ntoa($nmask) . "\n" if ($opt_v);
    print "  " . inet_ntoa($naddr_end) . "\n" if ($opt_v);
    print "  " . "count: " . ($addr_end - $addr + 1) . "\n" if ($opt_v);
  } else {
    die("Bad format. Expected 1.2.3.4/20");
  }
  return ($addr, $addr_end, $mask, $prefix, ($addr_end - $addr + 1));
}

sub usage {
   print "
Usage: $0 [OPTIONS]

Options:
  -n string, --networks=string (default '$NETWORKS')
  -x string, --exclude=string (default '$EXCLUDE')
  -i string, --interface=string (default '$IF')
  -t string, --target=string (default '$DIRECT') 'src' or 'dst'
  -r string, --rate=string (default '$RATE') with Kbit, Mbit, bps
  -m string, --mtu=string (default '$MTU') where TSO is ON
  -v, --verbose           Verbose
  -h, --help              Help
\n";
   exit;
}