#!/usr/bin/env python3
#
# Test parallels dirty bitmaps
#
# Copyright (c) 2021 Virtuozzo International GmbH.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

import hashlib
import struct
import uuid

from qemu.utils import VerboseProcessError

import iotests
from iotests import qemu_nbd_popen, qemu_img, qemu_img_create, qemu_img_map, \
    qemu_io, log, file_path, QemuStorageDaemon

iotests.script_initialize(supported_fmts=['parallels'])

nbd_sock = file_path('nbd-sock', base_dir=iotests.sock_dir)
nbd_sock2 = file_path('nbd-sock2', base_dir=iotests.sock_dir)
disk = iotests.file_path('disk')
bitmap = 'e4f2eed0-37fe-4539-b50b-85d2e7fd235f'
MULTI_L1_UUID = '5f2e1c00-0000-4000-8000-0123456789ab'
nbd_opts = f'driver=nbd,server.type=unix,server.path={nbd_sock}' \
        f',x-dirty-bitmap=qemu:dirty-bitmap:{bitmap}'


iotests.unarchive_sample_image('parallels-with-bitmap', disk)


with qemu_nbd_popen('--read-only', f'--socket={nbd_sock}',
                    f'--bitmap={bitmap}', '-f', iotests.imgfmt, disk):
    chunks = qemu_img_map('--image-opts', nbd_opts)
    cluster = 64 * 1024

    log('dirty clusters (cluster size is 64K):')
    for c in chunks:
        assert c['start'] % cluster == 0
        assert c['length'] % cluster == 0
        if c['data']:
            continue

        a = c['start'] // cluster
        b = (c['start'] + c['length']) // cluster
        if b - a > 1:
            log(f'{a}-{b-1}')
        else:
            log(a)


# Hand crafted Format Extension clusters.

EXT_MAGIC = 0xAB234CEF23DCEA87
BITMAP_MAGIC = 0x20385FAE252CB34A

CLUSTER = 512           # one sector per track
SECTORS = 8             # 4 KiB disk
EXT_SECTOR = 1          # the extension cluster follows the header
EH_SIZE = 24            # ParallelsFormatExtensionHeader
FH_SIZE = 24            # ParallelsFeatureHeader
BF_SIZE = 32            # ParallelsDirtyBitmapFeature

crafted = file_path('crafted')


def feature(magic, data_size, flags=0):
    return struct.pack('<QQII', magic, flags, data_size, 0)


def bitmap_feature(l1, granularity=1, nb_sectors=SECTORS, bid=None):
    """ParallelsDirtyBitmapFeature followed by its L1 table"""
    bid = uuid.UUID(bid).bytes if bid else uuid.uuid4().bytes
    return (struct.pack('<Q16sII', nb_sectors, bid, granularity, len(l1))
            + b''.join(struct.pack('<Q', entry) for entry in l1))


def extension(body, magic=EXT_MAGIC, checksum=True):
    body = body.ljust(CLUSTER - EH_SIZE, b'\0')
    csum = hashlib.md5(body).digest() if checksum else bytes(16)
    return struct.pack('<Q16s', magic, csum) + body


def parallels_header(tracks, bat_entries, nb_sectors, ext_sector, data_off=1):
    return struct.pack('<16sIIIIIQIIIQ', b'WithouFreSpacExt', 2, 16, 0,
                       tracks, bat_entries, nb_sectors, 0, data_off, 0,
                       ext_sector)


def write_sparse_image(body, tracks, nb_sectors, bat_entries, ext_sector):
    """Image whose extension cluster is the given body, zero padded."""
    cluster_size = tracks * 512
    md5 = hashlib.md5()
    md5.update(body)
    md5.update(bytes(cluster_size - EH_SIZE - len(body)))

    with open(crafted, 'wb') as f:
        f.write(parallels_header(tracks, bat_entries, nb_sectors, ext_sector))
        f.write(bytes(ext_sector * 512 - 64))            # BAT, unallocated
        f.write(struct.pack('<Q16s', EXT_MAGIC, md5.digest()) + body)
        f.truncate(ext_sector * 512 + cluster_size)      # sparse zero tail


def report(name):
    log(f'--- {name}')
    try:
        res = qemu_img('info', '-f', iotests.imgfmt, crafted, check=False)
    except VerboseProcessError as exc:
        # check=False does not cover termination by a signal
        log(f'qemu-img died with signal {-exc.returncode}')
        return

    out = iotests.filter_generated_node_ids(
        iotests.filter_testfiles(res.stdout))
    if res.returncode != 0:
        log(out.strip())
        return

    log('image opened')
    for line in out.splitlines():
        if 'warning' in line:
            log(line.strip())


def check(name, ext, tracks=1, data_off=1, pad=0):
    header = parallels_header(tracks, SECTORS, SECTORS, EXT_SECTOR, data_off)
    with open(crafted, 'wb') as f:
        f.write(header)
        f.write(bytes(EXT_SECTOR * 512 - len(header)))   # BAT, unallocated
        f.write(ext)
        if pad:
            f.truncate(EXT_SECTOR * 512 + tracks * 512 * (1 + pad))

    report(name)


log('')

# Control case, must keep working.
good = bitmap_feature([0])
check('well-formed extension',
      extension(feature(BITMAP_MAGIC, len(good)) + good + feature(0, 0)))

check('wrong extension magic',
      extension(feature(0, 0), magic=EXT_MAGIC ^ 1))

check('wrong extension checksum',
      extension(feature(0, 0), checksum=False))

# An older qemu truncates the extension cluster away on close.
with open(crafted, 'wb') as img:
    img.write(parallels_header(1, SECTORS, SECTORS, EXT_SECTOR))
    img.truncate(EXT_SECTOR * 512)
report('extension cut off the end of the file')

check('unknown feature', extension(feature(BITMAP_MAGIC ^ 1, 0)))

check('feature flags set', extension(feature(0, 0, flags=1)))

check('feature data_size beyond the cluster',
      extension(feature(BITMAP_MAGIC, CLUSTER)))

payload = CLUSTER - EH_SIZE - FH_SIZE
check('feature payload consumes the cluster',
      extension(feature(BITMAP_MAGIC, payload)
                + good.ljust(payload, b'\0')))

# bf.size must match the disk size or the compiler elides the OOB loads.
data1 = CLUSTER - 40 - EH_SIZE - FH_SIZE
check('second feature payload beyond the cluster',
      extension(feature(BITMAP_MAGIC, data1) + good.ljust(data1, b'\0')
                + feature(BITMAP_MAGIC, BF_SIZE + 8)
                + struct.pack('<Q', SECTORS)))

# Zero, not a power of two, and wrapping when shifted: all three assert.
for gran in (0, 3, 1 << 23):
    bf = bitmap_feature([0], granularity=gran)
    check(f'bitmap granularity {gran}',
          extension(feature(BITMAP_MAGIC, len(bf)) + bf + feature(0, 0)))

# entry << BDRV_SECTOR_BITS must not wrap into a negative offset.
bf = bitmap_feature([0xffffffffffffffff])
check('bitmap L1 entry overflows',
      extension(feature(BITMAP_MAGIC, len(bf)) + bf + feature(0, 0)))

# A cluster which is not there reads back as zeroes, as short reads on the
# protocol node are zero filled, so the bitmap would silently load as clean.
bf = bitmap_feature([8])
check('bitmap L1 entry past the image file',
      extension(feature(BITMAP_MAGIC, len(bf)) + bf + feature(0, 0)))

# An entry pointing at the header would deserialize the header as bitmap data.
bf = bitmap_feature([4])
check('bitmap L1 entry inside the header',
      extension(feature(BITMAP_MAGIC, len(bf)) + bf + feature(0, 0)),
      data_off=8, pad=16)

# Largest cluster_size parallels_open() accepts, about 2 GiB.
check('cluster_size beyond the file size',
      extension(feature(0, 0)), tracks=0x7fffffff // 513)


def check_big_cluster():
    """Reaches the serialization coverage overflow, costs about 550 MiB."""
    tracks = 1 << 19                    # cluster_size = 2^28
    nb_sectors = 1 << 32                # 2 TiB, needs two L1 entries
    ext_sector = 256                    # past the BAT
    l1 = [0, 1]                         # entry 1 deserializes as all ones

    bmap = bitmap_feature(l1, nb_sectors=nb_sectors)
    body = feature(BITMAP_MAGIC, len(bmap)) + bmap + feature(0, 0)

    write_sparse_image(body, tracks, nb_sectors, nb_sectors // tracks,
                       ext_sector)
    report('bitmap serialization coverage overflow')


check_big_cluster()


def check_multi_l1():
    """Bitmap spanning two L1 entries with an ordinary 64 KiB cluster."""
    tracks = 128                        # cluster_size = 64 KiB
    nb_sectors = 1 << 27                # 64 GiB, needs two L1 entries
    gran_sectors = 128                  # 64 KiB, as in the sample image
    bat_entries = nb_sectors // tracks
    ext_sector = (64 + 4 * bat_entries) // 512 + 8   # past the BAT
    l1 = [0, 1]                         # first half clean, second all ones

    bmap = bitmap_feature(l1, gran_sectors, nb_sectors, MULTI_L1_UUID)
    body = feature(BITMAP_MAGIC, len(bmap)) + bmap + feature(0, 0)

    write_sparse_image(body, tracks, nb_sectors, bat_entries, ext_sector)

    log('--- bitmap spanning two L1 entries')
    opts = f'driver=nbd,server.type=unix,server.path={nbd_sock2}' \
           f',x-dirty-bitmap=qemu:dirty-bitmap:{MULTI_L1_UUID}'
    with qemu_nbd_popen('--read-only', f'--socket={nbd_sock2}',
                        f'--bitmap={MULTI_L1_UUID}', '-f', iotests.imgfmt,
                        crafted):
        for chunk in qemu_img_map('--image-opts', opts):
            if not chunk['data']:
                first = chunk['start'] >> 30
                last = (chunk['start'] + chunk['length']) >> 30
                log(f'dirty {first}-{last} GiB')


check_multi_l1()


# Storing a bitmap. The driver has to produce the extension itself here,
# so the images are real ones rather than hand crafted clusters.

INUSE_MAGIC = 0x746F6E59
INUSE_OFF = 0x2c
EXT_OFF = 0x38
TRACKS_OFF = 0x1c

stored = file_path('stored')
copy = file_path('copy')
STORED = '9b7b8b2e-1c4a-4d3f-8e6a-1f2d3c4b5a69'
OTHER = '11111111-2222-3333-4444-555555555555'
MERGED = '22222222-3333-4444-5555-666666666666'


def read_header(off, fmt):
    with open(stored, 'rb') as f:
        return struct.unpack_from(fmt, f.read(0x40), off)[0]


def in_use():
    return read_header(INUSE_OFF, '<I') == INUSE_MAGIC


def bitmaps_in_file():
    """The bitmaps the image file holds right now, node open or not."""
    ext_off = read_header(EXT_OFF, '<Q') * 512
    if ext_off == 0:
        return []
    with open(stored, 'rb') as f:
        f.seek(ext_off)
        ext = f.read(read_header(TRACKS_OFF, '<I') * 512)

    names, pos = [], EH_SIZE
    while pos + FH_SIZE <= len(ext):
        magic, _, data_size, _ = struct.unpack_from('<QQII', ext, pos)
        if magic == 0:
            break
        if magic == BITMAP_MAGIC:
            names.append(struct.unpack_from('<16s', ext, pos + FH_SIZE + 8)[0])
        pos += FH_SIZE + ((data_size + 7) & ~7)
    return names


def bitmaps_reported():
    info = iotests.qemu_img_info('-f', iotests.imgfmt, stored)
    data = info.get('format-specific', {}).get('data', {})
    return [b['name'] for b in data.get('bitmaps', [])]


def make_image(*bitmaps, size='8M', cluster_size=None):
    opts = ['-o', f'cluster_size={cluster_size}'] if cluster_size else []
    qemu_img_create('-f', iotests.imgfmt, *opts, stored, size)
    for name in bitmaps:
        qemu_img('bitmap', '--add', '-f', iotests.imgfmt, stored, name)


def refuse(reason, size):
    """Add a fine grained bitmap which the Format Extension can not hold."""
    log(f'--- {reason}')
    make_image(size=size, cluster_size=512)
    res = qemu_img('bitmap', '--add', '-g', '512', '-f', iotests.imgfmt,
                   stored, STORED, check=False)
    log(iotests.filter_testfiles(res.stdout).strip())


def open_node():
    daemon = QemuStorageDaemon(qmp=True)
    daemon.cmd('blockdev-add',
               {'driver': iotests.imgfmt, 'node-name': 'n0',
                'file': {'driver': 'file', 'filename': stored}})
    return daemon


def stored_dirty(name, size, *writes):
    """Dirty the given ranges, then read the stored bitmap back over NBD."""
    log(f'--- {name}')
    make_image(STORED, size=size)
    qemu_io('-f', iotests.imgfmt,
            *[arg for w in writes for arg in ('-c', w)], stored)

    opts = f'driver=nbd,server.type=unix,server.path={nbd_sock}' \
           f',x-dirty-bitmap=qemu:dirty-bitmap:{STORED}'
    with qemu_nbd_popen('--read-only', f'--socket={nbd_sock}',
                        f'--bitmap={STORED}', '-f', iotests.imgfmt, stored):
        mapping = qemu_img_map('--image-opts', opts)

    for chunk in mapping:
        if not chunk['data']:
            log('    dirty %d MiB + %d KiB' % (chunk['start'] // (1 << 20),
                                               chunk['length'] // 1024))


# A bitmap is stored as an L1 table of clusters holding its serialized
# bits, with an entry of one standing for a chunk which is dirty as a
# whole. Both of those have to describe exactly what was written.
stored_dirty('three separate ranges', '64M',
             'write -P 0x11 0 64k', 'write -P 0x22 4M 128k',
             'write -P 0x33 40M 64k')
stored_dirty('the whole image', '8M', 'write -P 0x11 0 8M')
stored_dirty('less than a cluster', '4M', 'write -P 0x11 0 1M')
stored_dirty('the tail of the image', '8M', 'write -P 0x11 6M 2M')

# The L1 table shares the Format Extension cluster with the feature
# headers and the end of features marker, and a small cluster with a fine
# granularity does not leave room for it. The marker needs room of its own
# behind the L1 table, so a bitmap which fills the cluster to within a
# header of the end can not be stored either.
refuse('a bitmap too large for the Format Extension cluster', '4G')
refuse('a bitmap leaving no room for the end of features marker', '104M')

# The cluster holds every bitmap of the image at once, so what still fits
# depends on what is stored there already.
log('--- a bitmap which does not fit next to the ones already stored')
make_image(size='64M', cluster_size=512)
for i in range(1, 9):
    bm_uuid = '9b7b8b2e-1c4a-4d3f-8e6a-1f2d3c4b5a%02d' % i
    add_out = qemu_img('bitmap', '--add', '-f', iotests.imgfmt, stored,
                       bm_uuid, check=False)
    if add_out.returncode != 0:
        log('    bitmap %d refused' % i)
        log('    ' + iotests.filter_testfiles(add_out.stdout).strip())
        break

log('--- the ones which did fit are all still there')
log('    %d bitmaps stored' % len(bitmaps_reported()))

log('--- removing a bitmap reaches the image right away')
make_image(STORED, OTHER)
qsd = open_node()
log(f'    stored before: {len(bitmaps_in_file())} bitmaps')
qsd.cmd('block-dirty-bitmap-remove', {'node': 'n0', 'name': STORED})
log(f'    stored after:  {len(bitmaps_in_file())} bitmaps, '
    'without a clean close')
qsd.stop()

log('--- the in use flag comes back when a node is activated again')
make_image(STORED)
qsd = open_node()
log(f'    in use while open:         {in_use()}')
qsd.cmd('blockdev-set-active', {'node-name': 'n0', 'active': False})
log(f'    in use after deactivation: {in_use()}')
qsd.cmd('blockdev-set-active', {'node-name': 'n0', 'active': True})
log(f'    in use after activation:   {in_use()}')
qsd.stop()

# The bitmap data cluster is allocated right behind the only data cluster,
# so it starts at sector 4096 of a 4M image with the default cluster size.
log('--- a bitmap which can not be written is reported, not lost in silence')
make_image(STORED, size='4M')
err_opts = (f'driver={iotests.imgfmt},file.driver=blkdebug,'
            f'file.image.driver=file,file.image.filename={stored},'
            f'file.inject-error.0.event=none,'
            f'file.inject-error.0.iotype=write,'
            f'file.inject-error.0.sector=4096')
err_out = qemu_io('--image-opts', err_opts, '-c', 'write -P 0x11 0 64k',
                  check=False)
for err_line in iotests.filter_generated_node_ids(err_out.stdout).splitlines():
    if 'Lost persistent bitmaps' in err_line:
        log('    ' + err_line.strip())
log(f'    in use after the failure:  {in_use()}')

log('--- but the guest data went through, and the image says it is in use')
read_out = qemu_io('-f', iotests.imgfmt, '-r', '-c', 'read -P 0x11 0 64k',
                   stored)
log(read_out.stdout, filters=[iotests.filter_generated_node_ids,
                              iotests.filter_testfiles,
                              iotests.filter_qemu_io])
log('    %d bitmaps stored' % len(bitmaps_reported()))

log('--- and the image is still readable')
info_out = qemu_img('info', '-f', iotests.imgfmt, stored, check=False)
log('image opened' if info_out.returncode == 0
    else iotests.filter_testfiles(info_out.stdout).strip())


# qemu-img-bitmaps drives these for qcow2 alone, as it is built on a
# backing chain and a rebase, which parallels has neither of.

log('--- bitmaps can not be preserved into raw')
make_image(STORED)
raw_out = qemu_img('convert', '--bitmaps', '-f', iotests.imgfmt, '-O',
                   'raw', stored, copy, check=False)
log('    ' + iotests.filter_testfiles(raw_out.stdout).strip())

log('--- convert --bitmaps carries the bitmaps into the copy')
make_image(STORED, OTHER)
qemu_io('-f', iotests.imgfmt, '-c', 'write -P 0x11 1M 64k', stored)
qemu_img('convert', '--bitmaps', '-f', iotests.imgfmt, '-O', iotests.imgfmt,
         stored, copy)
copy_info = iotests.qemu_img_info('-f', iotests.imgfmt, copy)
copy_data = copy_info.get('format-specific', {}).get('data', {})
log('    %d bitmaps in the copy' % len(copy_data.get('bitmaps', [])))

log('--- merging one bitmap into a new one')
qemu_img('bitmap', '--add', '--merge', STORED, '-f', iotests.imgfmt,
         stored, MERGED)
log('    %d bitmaps stored' % len(bitmaps_reported()))

log('--- clearing a bitmap keeps it in the image')
qemu_img('bitmap', '--clear', '-f', iotests.imgfmt, stored, STORED)
log('    %d bitmaps stored' % len(bitmaps_reported()))

log('--- remove-all sweeps every bitmap of the image')
qemu_img('bitmap', '--remove-all', '-f', iotests.imgfmt, stored)
log('    %d bitmaps after remove-all' % len(bitmaps_reported()))

log('--- remove-all with --add leaves only the new bitmap')
make_image(STORED, OTHER)
qemu_img('bitmap', '--remove-all', '--add', '-f', iotests.imgfmt,
         stored, MERGED)
log('    ' + ' '.join(bitmaps_reported()))
