#!/usr/bin/perl -w
#
# Copyright (c) 2007-2016 University of Utah and the Flux Group.
# 
# {{{EMULAB-LICENSE
# 
# This file is part of the Emulab network testbed software.
# 
# This file is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
# 
# This file is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public
# License for more details.
# 
# You should have received a copy of the GNU Affero General Public License
# along with this file.  If not, see <http://www.gnu.org/licenses/>.
# 
# }}}
#
use English;
use Getopt::Std;

# Drag in path stuff so we can find emulab stuff.
BEGIN { require "/etc/emulab/paths.pm"; import emulabpaths; }

sub usage()
{
    print "Usage: fixup-fstab-swaps [-U] [-E]\n";
    exit(1);
}

my $optlist = "UE";

my $FSTAB = "/etc/fstab";
my $SWAPS = "/proc/swaps";
my $FIRST_BOOT_FLAG = "$BOOTDIR/added_swaps";
my $script = $0;

my $noupdate = 0;
my $noenable = 0;

sub removeswaps();
sub findswaps();
sub swapenable();

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

#
# Parse command arguments. Once we return from getopts, all that should be
# left are the required arguments.
#
%options = ();
if (! getopts($optlist, \%options)) {
    usage();
}
if (defined($options{"U"})) {
    $noupdate = 1;
}
if (defined($options{"E"})) {
    $noenable = 1;
}
if (@ARGV) {
    usage();
}

# Remove any existing swap devices from fstab; we can't ensure that they'll
# work on this node.  This only runs the first time the node is booted after
# being imaged.  Then look for linux swap partitions on the root disk and add
# them to fstab.  This way we have sane defaults for swap.  The user is, of
# course, free to change the swap configration later if necessary.
if (! -f $FIRST_BOOT_FLAG && !$noupdate) {

	removeswaps();
	findswaps();

	open FD, ">$FIRST_BOOT_FLAG";
	close FD;
}

#
# Figure out if there is a swap device, add it to fstab and intialize if
# necessary.
#
swapenable()
    if (!$noenable);

exit(0);

sub removeswaps()
{
    my @buffer;

    if (!open(FD, "+<$FSTAB")) {
	print STDERR "*** WARNING: could not open $FSTAB for writing: $!\n";
	return;
    }
    while (<FD>) {
	next if (/^\s*$/ || /^#/);
	my ($fs,$mpt,$type,$opt) = split;
	if ($type eq "swap") {
	    print "Removing old swap device $fs\n";
	    next;
	}

	push @buffer, $_;
    }

    truncate FD, 0;
    seek(FD, 0, 0);
    print FD $_ for (@buffer);
    close(FD);
}

sub findswaps()
{
    #
    # No swap device.
    # Identify the root disk and see if we can locate a Linux swap partition
    # on it.
    #
    my @swapdevs;
    my $rdisk;
    my $rootfs = `df / | grep /dev/`;
    if ($rootfs =~ /^(\/dev\/nvme\d+n\d+)p\d+/) {
	$rdisk = $1;
    } elsif ($rootfs =~ /^(\/dev\/[a-z]+)\d+/) {
	$rdisk = $1;
    }
    elsif ($rootfs =~ /^(\/dev\/[^\s]+)\s/) {
	$rootfs = `readlink -f $1`;
	if ($rootfs =~ /^(\/dev\/[a-z]+)\d+/) {
	    $rdisk = $1;
	}
	else {
	    print STDERR "*** WARNING: could not identify root disk, ".
		"no swap enabled\n";
	    return;
	}
    }
    else {
	print STDERR "*** WARNING: could not identify root disk, ".
	    "no swap enabled\n";
	return;
    }

    @lines = `fdisk -l $rdisk 2>/dev/null | grep 'Linux swap'`;
    if ($? != 0) {
	print STDERR "*** WARNING: could not read MBR of $rdisk, ".
	    "no swap enabled\n";
	return;
    }

    chomp(@lines);
    foreach $line (@lines) {
	if ($line =~ /^(\/dev\/\S+)\s+\d+\s+\d+\s+(\d+)\+?/) {
	    my $dev = $1;
	    my $size = $2;
	    print "Using $dev ($size sectors) for swap\n";
	    push @swapdevs, $dev;
	}
    }

    if (!@swapdevs) {
        print STDERR "*** WARNING: could not locate a suitable swap device, ".
	    "no swap enabled\n";
	return;
    }

    if (!open(FD, ">>$FSTAB")) {
	print STDERR "*** WARNING: could not add swap devices to $FSTAB: $!\n";
	return;
    }

    print FD "# the following swap devices added by $script\n";
    for (@swapdevs) {
	print FD "$_\t\tswap\tswap\tdefaults\t0\t0\n";
    }
    close(FD);
}

sub swapenable()
{
    my %curswaps;

    if (!open(FD, "<$SWAPS")) {
	print STDERR "*** WARNING: could not open $SWAPS for reading: $!\n";
	return;
    }

    <FD>; # Throw away the header
    while(<FD>) {
	    @_ = split;
	    $curswaps{$_[0]} = 1;
    }
    close FD;

    if (!open(FD, "<$FSTAB")) {
	print STDERR "*** WARNING: could not open $FSTAB for reading, ".
	    "no swap enabled\n";
	return;
    }
    while (<FD>) {
	next if (/^\s*$/ || /^#/);
	my ($fs,$mpt,$type,$opt) = split;
	next if ($type ne "swap" || ! -b $fs || exists $curswaps{$fs});
	next if ($opt =~ /\bnoauto\b/);

	if (system("mkswap $fs")) {
		print STDERR "*** WARNING: could not initialize swap on $fs\n";
		next;
	}

	if (system("swapon $fs")) {
	    print STDERR "*** WARNING: could not enable swap on $fs\n";
	    next;
	}
    }
    close(FD);

}
