#!/usr/bin/env perl
# vim:ts=4:sw=4:et
#
# © 2017 Michael Stapelberg <stapelberg@debian.org>
#
# Extends the root partition and root file system to cover the entire remainder
# of the device. Requires the root file system to be a block device matching
# /.+p[0-9]$/, e.g. /dev/mmcblk0p2.
#
# Requires only perl-base (present on all Debian installations)

use strict;
use Fcntl qw(:DEFAULT :seek);

################################################################################
# Find the root device by looking at what is mounted on /
################################################################################

sub slurp {
    my ($fn) = @_;
    open(my $fh, '<', $fn)
        or die qq|open($fn): $!|;
    local $/;
    <$fh>;
}

my ($rootpart) = grep { $_ ne '' } map { /([^ ]+) \/ / && $1 } split("\n", slurp(</proc/mounts>));
my $rootdev;
my $pno;
if ($rootpart =~ /^(.+)p([0-9])$/) {
    $rootdev = $1;
    $pno = int($2);
} else {
    die qq|root partition "$rootpart" unexpectedly does not end in p[0-9]|;
}

################################################################################
# Get the size of the root block device in bytes
################################################################################

sub SYS_ioctl { 29; } # aarch64-specific
sub BLKGETSIZE64 { 2148012658; }

sysopen(my $root, $rootdev, O_RDWR)
    or die qq|sysopen($rootdev): $!|;

my $devsizep = pack("Q", ());
$! = 0;
syscall(SYS_ioctl(), fileno($root), BLKGETSIZE64(), $devsizep) >= 0
    or die qq|ioctl(BLKGETSIZE64): $!|;
my ($devsize) = unpack("Q", $devsizep);

################################################################################
# Read the partition table entry
################################################################################

sub boot_code { 446; }
sub partition_table_entry { 16; }
sub partition_table_start_offset { 8; }
sub partition_table_length_offset { 12; }
sub partition_count { 4; }
sub partition_signature_len { 2; }
sub tablelength { partition_table_entry() * partition_count() + partition_signature_len(); }

# The contents of $original_partition_table must be updated whenever the
# partition layout on the SD card image changes.
# This layout matches
# https://people.debian.org/~stapelberg/raspberrypi3/2017-10-08/2017-10-08-raspberry-pi-3-buster-PREVIEW.img.bz2
my $original_partition_table = pack('H4' x (tablelength()/2),
    "0020", "2100", "0c3e", "1826", "0008", "0000", "0058", "0900",
    "003e", "1926", "833a", "2e8c", "0060", "0900", "0000", "1900",
    "0000", "0000", "0000", "0000", "0000", "0000", "0000", "0000",
    "0000", "0000", "0000", "0000", "0000", "0000", "0000", "0000",
    "55aa");

sub check_orig_partition_table {
    my ($offset, $len, $bytes) = @_;
    sysseek($root, $offset, SEEK_SET)
        or die qq|sysseek($offset): $!|;
    my $buf;
    sysread($root, $buf, $len)
        or die qq|sysread(): $!|;
    if ($buf ne $bytes) {
        print "Partition table already modified\n";
        exit 0;
    }
}

sub read_uint32_le {
    my ($offset) = @_;
    sysseek($root, $offset, SEEK_SET)
        or die qq|sysseek($offset): $!|;
    my $buf;
    sysread($root, $buf, 4)
        or die qq|sysread(): $!|;
    my ($val) = unpack("V", $buf);
    return $val;
}

# Ensure partition table has not been modified by the user.
check_orig_partition_table(boot_code(), tablelength(), $original_partition_table);

my $entry_offset = boot_code() + (partition_table_entry() * ($pno - 1));
my $start     = 512 * read_uint32_le($entry_offset + partition_table_start_offset());
my $oldlength = 512 * read_uint32_le($entry_offset + partition_table_length_offset());
my $newlength = ($devsize - $start);
if ($oldlength == $newlength) {
    print "Partition $rootpart already at maximum size $newlength\n";
    exit 0;
}

################################################################################
# Change the partition length
################################################################################

sub write_uint32_le {
    my ($offset, $val) = @_;
    sysseek($root, $offset, SEEK_SET)
        or die qq|sysseek($offset): $!|;
    my $buf = pack("V", $val);
    syswrite($root, $buf, 4)
        or die qq|syswrite: $!|;
}

print "Resizing partition $rootpart from $oldlength to $newlength bytes\n";
write_uint32_le($entry_offset + partition_table_length_offset(), $newlength / 512);

################################################################################
# Tell linux about the new partition size using the BLKPG ioctl (BLKRRPART
# cannot be used when the device is mounted, even read-only).
################################################################################

sub BLKPG { 0x1269; }
sub BLKPG_RESIZE_PARTITION { 3; }

my $part = pack("q q i Z64 Z64", $start, $newlength, $pno, "devname", "volname");
my $partb = "\x00" x length($part);
my $op = BLKPG_RESIZE_PARTITION();
my $ioctl_arg = pack("i i i x4 p", $op, 0, length($part), $part);
$! = 0;
syscall(SYS_ioctl(), fileno($root), BLKPG(), $ioctl_arg) >= 0
    or die "ioctl(BLKPG): $!";

################################################################################
# Run resize2fs to enlarge the file system
################################################################################

# resize2fs requires the file system to be writeable.
my @remountcmd = ('mount', '-o', 'remount,rw', $rootpart);
system(@remountcmd) == 0
    or die qq|system(| . join(" ", @remountcmd) . qq|): $!|;

my @resizecmd = ('resize2fs', "${rootpart}");
exec { $resizecmd[0] } @resizecmd
    or die qq|exec(| . join(" ", @resizecmd) . qq|): $!|;