os_setup.in

#!/usr/bin/perl -wT
use English;
require 'ctime.pl';

#
# TODO: Delta stuff.
#       Reload disk images.
#       Finish up local OS spec stuff. Kinda hacky right now. Does not deal
#       with command lines or disk images. Only does neboot type stuff and
#	hardwires the path to /tftpboot/proj/$pid/$imagepaths{$os}
#

#
# Parse an IR file and determine what OS has been requested on each node.
# Do the database magic to make it so. Only root, admin types, or the
# owners of the nodes in an experiments may do this. The machines are then
# rebooted (or power cycled).
# 
# usage: os_setup <pid> <eid> <ir_file>
#

#
# Configure variables
#
my $TB		= "@prefix@";
my $DBNAME	= "@TBDBNAME@";

my $ssh		= "ssh -n -q";
my $power	= "$TB/bin/power";
my $ping	= "/sbin/ping";
my $mail        = "/usr/bin/mail";
my $tbops       = "testbed-ops\@flux.cs.utah.edu";
my $dbg		= 0;
my %imagepaths  = ();
my %imageparts  = ();
my %nodeos      = ();
my %nodepath    = ();
my %nodepart    = ();
my %waitfor     = ();
my %pids	= ();
my $SAVEUID	= $UID;
my @row;

#
# This stuff is BOGUS! Quick hack for paper deadline to make Jay happy.
#
my $doreloading = 0;
my $NETDISK     = "/tftpboot/netdisk";
my $PAPERADDR	= "boss.emulab.net";
my $IMAGE       = "/usr/testbed/images/wd0-all.ndz";
my $RELOADCMD   = "${PAPERADDR}:${IMAGE} wd0";
my %reload      = ();
    
# un-taint path
$ENV{'PATH'} = '/bin:/usr/bin:/usr/local/bin';
delete @ENV{'IFS', 'CDPATH', 'ENV', 'BASH_ENV'};

$| = 1; #Turn off line buffering on output

#
# Set up for querying the database.
# 
use Mysql;
my $DB = Mysql->connect("localhost", $DBNAME, "script", "none");

if ( $#ARGV != 2) {
    die("Usage: os_setup <pid> <eid> <ir_file>\n". 	
	"Sets node OS configuration from a .ir file.\n");
}
my $pid = $ARGV[0];
my $eid = $ARGV[1];
my $ir  = $ARGV[2];

#
# Figure out who called us. Only root, tbroot, people with admin status
# in the DB, or the owner of the experiment can run this script.
#

$db_result = $DB->query("select expt_head_uid from experiments ".
			"where eid='$eid' and pid='$pid'");
if ($db_result->numrows < 1) {	
  die("There is no experiment '$eid' in project '$pid'.\n");
}

if ($UID != 0) {
    my ($me) = getpwuid($UID)
	or die "$UID not in passwd file";

    @row = $db_result->fetchrow_array();
    if ($row[0] ne "$me") {
	print STDERR "Checking for admin status ...\n" if $dbg;
	$db_result = $DB->query("select admin from users where uid='$me'");
        @row = $db_result->fetchrow_array();
	if ($row[0] != 1) {
	    die("os_setup: You must be root or a TB administrator\n");
	}
    }
}

#
# Open up the ir file.
#
if (-e "$ir") {
    open(IN, $ir) || die("Couldn't open $ir\n");
}
else {
    die("Couldn't open $ir\n");
}

#
# Look for the start of the OS section. Exit if not found
#
my $ossection=0;

while (<IN>) {
    if ( /^start os/i ) { 
	$ossection=1; 
	print STDERR "Start OS section...\n" if $dbg;
	last;
    }
}
if ($ossection == 0) {
    die("No OS section in $ir\n");
}

#
# Parse the images table to see what the local images are.
# 
while (<IN>) {
    if ( /^start images/i ) { 
	print STDERR "Start IMAGES section...\n" if $dbg;
	next;
    }
    elsif ( /^end images/i ) {
	print STDERR "End IMAGES section...\n" if $dbg;
	last;
    }
    my ($id,$path,$part) = split();
    print STDERR "$id $path $part\n" if $dbg;
    $imagepaths{$id} = $path;
    $imageparts{$id} = $part;
}

#
# Okay, parse the OS section.
# 
# Search the NODES section for OS labels. Consult the database to make
# sure the node specified in the IR file is really in the pid/eid that
# was given on the command line. 
# 
while (<IN>) {
    if ( /^start nodes/i ) { 
	print STDERR "Start NODES section...\n" if $dbg;
	next;
    }
    elsif ( /^end nodes/i ) {
	print STDERR "End NODES section...\n" if $dbg;
	last;
    }
    my ($node,$os) = split();
    print STDERR "$node $os\n" if $dbg;
    $nodeos{$node} = $os;

    $db_result = $DB->query("select pid,eid from reserved ".
			    "where node_id='$node'");
    if ($db_result->numrows < 1) {	
	die("There is no node '$node' reserved in the DB.\n");
    }
    @row = $db_result->fetchrow_array();
    if ($row[0] ne "$pid" ||
	$row[1] ne "$eid") {
	die("Node '$node' pid/eid mismatch: $pid/$eid ... $row[0]/$row[1]\n");
    }

    #
    # BOGUS
    # 
    $db_result = $DB->query("select pid from last_reservation ".
			    "where node_id='$node'");
    $reload{$node} = 0;
    if ($doreloading) {
	if ($db_result->numrows) {
	    @row = $db_result->fetchrow_array();
	    if ($row[0] ne $pid) {
		$reload{$node} = 1;
	    }
	}
    }

    #
    # First see if the OS spec is a local one from the images table. We
    # leave it up to the user to make sure the OS is capable of running
    # on the node! This local image stuff is pretty hacky right now. I'm
    # only going to allow netboot type stuff and stuff on partition 4.
    #
    if ($imagepaths{$os}) {
	if ($imageparts{$os} eq "4") {
	    $nodepart{$node} = 4;
	    $waitfor{$node}  = 0;
	    next;
	}
	if ($imageparts{$os} eq "mb") {
	    #
	    # Okay, now it gets really bad. I'm going to form the tftpboot
	    # path right here, since local OS specs are not in the database
	    # disk_images table.
	    # 
	    $nodepath{$node} = "/tftpboot/proj/$pid/$imagepaths{$os}";
	    $waitfor{$node}  = 0;
	    next;
	}
	die("Improper local OS spec: $os. ".
	    "Partition can currently only be 4 or 'mb' (multiboot)");
    }

    #
    # Check to make sure that the OS spec is valid by checking the database.
    # Cross check type from nodes table against image_ids for that type in
    # the disk_images table. 
    #
    $db_result = $DB->query("select ".
			    "disk_images.image_id,disk_images.img_path ".
			    "from disk_images left join nodes ".
			    "on nodes.type=disk_images.type ".
			    "and disk_images.image_id='$os' ".
                            "where nodes.node_id='$node'");
    if ($db_result->numrows < 1) {	
	die("Improper image specification $node:$os in IR file.\n");
    }
    @row = $db_result->fetchrow_array();
    $nodepath{$node} = $row[1];

    #
    # See if the OS that is going to be booted supports a ping feature.
    # This is the only way we can tell if the machine has come back alive.
    # If not, then we effectively do not support the os for anything at all.
    #
    if (OSFeatureSupported($os, "ping")) {
	$waitfor{$node} = 1;
    }
    else {
	$waitfor{$node} = 0;
    }
}

#
# Lifted right out of delay_setup.
# 
foreach my $node ( keys %nodeos ) {
    my $pc	= $node;
    my $os	= $nodeos{$node};
    my $path    = $nodepath{$node};

    print STDOUT "Changing default OS for $pc to $os:$path ".
	         "and rebooting ...\n";

    #
    # database goo. Reset the OS stuff.
    #
    $sth = $DB->query("update nodes set ".
		      "def_boot_image_id='$os',def_boot_path='$path' ".
		      "where node_id='$pc'");
    if ($sth == 0) {
	die("Database update failed. Aborted...\n");
    }

    if ($nodepart{$node}) {
	#
	# At this point we would perhaps want to do an os_load, but I'm
	# not ready to do that. Just skip the node and let the user deal
	# with it.
	#
	printf STDOUT "Skipping $node. You will need to load the OS.\n";
	next;
    }

    #
    # BOGUS!
    #
    if ($reload{$pc}) {
	$sth = $DB->query("update nodes set ".
			  "next_boot_path='$NETDISK',".
			  "next_boot_cmd_line='$RELOADCMD' ".
			  "where node_id='$pc'");
    }

    #
    # Fire off a reboot process so that we can overlap them all.
    # We need the pid so we can wait for them all before preceeding.
    #
    $mypid = RebootNode($pc);
    $pids{$pc} = $mypid;
}

#
# Wait for all the reboot children to exit before continuing.
#
foreach my $node ( keys %nodeos ) {
    my $pc	= $node;
    my $mypid     = $pids{$pc};

    waitpid($mypid, 0);
    if ($?) {
	die("Reboot of node $pc failed!");
    }
    my $t = ctime(time);
    print STDOUT "$t";
    print STDOUT "$pc rebooting ...\n"
}

my $t = ctime(time);
print STDOUT "$t";
	
print STDOUT "Waiting for testbed nodes to finish rebooting ...\n";

my $waitstart = time;

#
# Now lets wait for them to come back alive.
#
foreach my $node ( keys %nodeos ) {
    $node =~ /^([a-zA-Z0-9_\-]*)$/;
    my $pc = $1;

    #
    # Don't bother to wait for nodes that are running foreign OSs since
    # we are not going to deal with them anyway later in the process.
    #
    if ($waitfor{$pc} == 0) {
	print STDOUT "Not waiting for $pc to come alive. Foreign OS.\n";
	next;
    }	

    if (WaitTillAlive($pc) == 0) {
	my $t = ctime(time);
	print STDOUT "$t";
	
	print STDOUT "$pc is alive and well\n";
	next;
    }

    print STDOUT "$pc may be down. This has been reported to testbed-ops.\n";
    print STDOUT "Please end this experiment, and try again.\n";
    
    # Reserve it to testbed down

    $cmd = "update reserved set pid='testbed',eid='down' ".
      "where eid='$eid' and pid='$pid' and node_id='$pc'";
    print "Using '$cmd'\n" if $dbg;
    $db_result = $DB->query($cmd) 
      || print STDERR "WARNING: Couldn't change reservation:".
	$DB->errmsg."\n";
    if ($db_result->num_rows < 1 ) {
      print STDERR "WARNING: Couldn't change reservation!\n";
    }

    # Send mail to testbed-ops about it
    open(MAIL,"| $mail -s \"TESTBED: $pc down?\" $tbops");
    print MAIL "User ".getpwuid($SAVEUID)." was running expt. $eid\n";
    print MAIL "in proj. $pid using ir file /proj/$pid/exp/$eid/tbdata/$ir\n";
    print MAIL "but $pc appears to be unresponsive.\n";
    print MAIL "\nPlease look into this matter. $pc has been reserved to\n";
    print MAIL "the testbed/down experiment until this has been resolved.\n\n";
    print MAIL "Thanks,\nTestbed Operations\ntestbed-ops\@flux.cs.utah.edu\n";
    close(MAIL);
    
    die("Oops, $pc did not come back alive!\n");
}

print STDOUT "OS Setup Done!\n";
exit 0;

#
# Power cycle a PC using the testbed power program.
#
sub PowerCycle {
    local($pc) = @_;

    if (system("$power cycle $pc") != 0) {
	print STDERR "WARNING: Could not power cycle $pc. Skipping ...\n";
    }
}

sub WaitTillAlive {
    my ($pc) = @_;

    my $maxwait = 150;
    if ($reload{$pc}) {
	$maxwait += 150;
    }

    print STDERR "Waiting for $pc to come alive\n" if $dbg;
    #
    # Sigh, a long ping results in the script waiting until all the
    # packets are sent from all the pings, before it will exit. So,
    # loop doing a bunch of shorter pings.
    #
    my $lasttime = ( (time - $waitstart) > 60 ? 61 : (time - $waitstart));
    for ($i = 0; $i < 120; $i++) {
	open(PING, "$ping -c 3 -t 4 $pc 2>&1 |");
	do {
	    $_ = <PING>;
	    if ( $_ =~ /bytes from/ ) {
		print STDERR "Yep, $pc alive and well\n" if $dbg;
		return 0;
	    }
	}
	until ( $_ =~ /transmitted, (\d*) packets received/ );
	my $curtime = time - $waitstart;
	print "Waited ",$curtime," seconds...\n" if $dbg;
	if ( $curtime % 60 < $lasttime % 60 ) { 
	  print STDERR "Still waiting for $pc - its been ",
	  (int ($curtime/60))," min.\n";
	}
	$lasttime = $curtime;
	if ($i > 3 && $curtime > $maxwait) { last; }
    }
    close(PING);
    print STDERR "$pc is not responding. Better check into it.\n" if $dbg;
    return 1;
}

#
# Reboot a node in a child process. Return the pid to the parent so
# that it can wait on all the children later.
# 
sub RebootNode {
    local($pc) = @_;

    my $t = ctime(time);
    print STDOUT "$t";
	
    print STDOUT "Rebooting $pc ...\n";

    $mypid = fork();
    if ($mypid) {
	return $mypid;
    }

    #
    # See if the machine is pingable. If its not pingable, then
    # we just power cycle the machine rather than wait for a bunch
    # of ssh/rsh commands to time out.
    #
    print STDERR "Pinging $pc ... \n" if $dbg;
    if (-e $ping) {
	open(PING, "$ping -c 4 -t 4 $pc 2>&1 |");
    }
    else {
	die("PING command $ping not found!\n");
    }
    do {
    }
    until ( <PING> =~ /transmitted, (\d*) packets received/ );
    close(PING);
    print STDERR "Got back $1 ping packets from $pc.\n" if $dbg;

    #
    # Power cycle if the machine is dead. It will come back up with the
    # proper OS, cause we modified the database above.
    #
    if ( $1 == 0 ) {
	print STDERR "$pc appears to be dead. Power cycling ...\n" if $dbg;
	PowerCycle($pc);
	exit(0);
    }

    #
    # Machine is pingable at least. Try to reboot it gracefully,
    # or power cycle anyway if that does not work. To this, we must
    # change our real UID to root so that ssh will work.
    #
    print STDERR "Rebooting $pc with ssh command ...\n" if $dbg;

    #
    # Run an ssh command in a child process, protected by an alarm to
    # ensure that the ssh is not hung up forever if the machine is in
    # some funky state.
    # 
    $syspid = fork();
    if ($syspid) {
	local $SIG{ALRM} = sub { kill("TERM", $syspid); };
	alarm 60;
	waitpid($syspid, 0);
	alarm 0;

	#
	# If ssh times out, just punch the button.
	# 
	if ($? == 15) {
	    print STDERR "$pc appears to be wedged. Power cycling ...\n"
		if $dbg;
	    PowerCycle($pc);
	    exit(0);
	}
    }
    else {
	$UID = 0;
	exec("$ssh $pc /sbin/reboot");
	exit(0);
    }
    
    #
    # Okay, before we power cycle lets really make sure. On FreeBSD, it might
    # have rebooted, but since the connection is terminated, system returns
    # an error status. So, lets ping it again and if its pingable, the
    # reboot must have failed. If it is not pingable, I assume that the
    # reboot really worked, and the exit value can be ignored.
    #
    my $exit_value = $? >> 8;
    print STDERR "reboot returned $exit_value. Lets make sure it dies\n"
	if $dbg;

    if (WaitTillDead($pc) == 0) {
	exit(0);
    }
	
    print STDERR "$pc appears to still be running. Power cycling ...\n"
	if $dbg;
    PowerCycle($pc);
    exit(0);
}

sub WaitTillDead {
    local($pc) = @_;

    print STDERR "Waiting for $pc to die off\n" if $dbg;
    #
    # Sigh, a long ping results in the script waiting until all the
    # packets are sent from all the pings, before it will exit. So,
    # loop doing a bunch of shorter pings.
    #
    for ($i = 0; $i < 15; $i++) {
	open(PING, "$ping -c 4 -t 4 $pc 2>&1 |");
	do {
	}
	until ( <PING> =~ /transmitted, (\d*) packets received/ );

	if ( $1 == 0 ) {
	    print STDERR "Good, $pc must have rebooted.\n" if $dbg;
	    return 0;
	}
    }
    close(PING);
    print STDERR "$pc is still alive.\n" if $dbg;
    return 1;
}

sub OSFeatureSupported {
    local($os)      = $_[0];
    local($feature) = $_[1];

    $db_result = $DB->query("select osfeatures from disk_images ".
			    "where image_id='$os'");

    if ($db_result->numrows < 1) {
	return 0;
    }
    foreach $osfeature (split(',', $db_result->fetchrow_array())) {
	if ($feature eq $osfeature) {
	    return 1;
	}
    }
    return 0;
}