links.py

#!/usr/bin/env python3

import collections
import multiprocessing
import multiprocessing.connection
import os
import subprocess
import sys


TargetFile = collections.namedtuple('TargetFile', ['filename', 'content'])

TARGETS = [TargetFile('target-1', b'123'),
           TargetFile('target-2', b'456')]


LINK_FILENAME = 'link'


def get_test_dir():
    return f"{os.environ['XDG_RUNTIME_DIR']}/links-test"


READY_MESSAGE = b'ready'

def wait_for_ready_message(conn):
    message = conn.recv()
    assert message == READY_MESSAGE

def send_ready_message(conn):
    conn.send(READY_MESSAGE)


def parent():
    parent, child = multiprocessing.Pipe(duplex=True)

    p = subprocess.Popen(['bwrap', '--bind', '/', '/', '--tmpfs', get_test_dir(),
                          '--unshare-all',
                          'python3', __file__, 'child', str(child.fileno())],
                          pass_fds=(child.fileno(),))
    child.close()

    fd = multiprocessing.reduction.recv_handle(parent)

    print('[parent] Setting up child environment')

    for target in TARGETS:
        target_fd = os.open(target.filename, os.O_WRONLY | os.O_CREAT, dir_fd=fd)
        os.write(target_fd, target.content)
        os.close(target_fd)

    os.symlink(f'{get_test_dir()}/{TARGETS[0].filename}', LINK_FILENAME, dir_fd=fd)

    print('[parent] Telling child parent is ready')
    send_ready_message(parent)

    print('[parent] Waiting for child to be ready')
    wait_for_ready_message(parent)

    print('[parent] Adjusting child environment')

    os.symlink(f'{get_test_dir()}/{TARGETS[1].filename}', LINK_FILENAME + '.tmp',
               dir_fd=fd)
    os.rename(LINK_FILENAME + '.tmp', LINK_FILENAME, src_dir_fd=fd, dst_dir_fd=fd)

    print('[parent] Telling child parent is ready again')
    send_ready_message(parent)

    print('[parent] Waiting for child')
    p.wait()


def child(conn_fd):
    fd = os.open(get_test_dir(), os.O_DIRECTORY)

    # XXX: using internal functions here
    conn = multiprocessing.connection.rebuild_connection(
                multiprocessing.reduction.DupFd(conn_fd), True, True)
    multiprocessing.reduction.send_handle(conn, fd, 0)

    print('[child] Waiting for parent to be ready')
    wait_for_ready_message(conn)

    with open(f'{get_test_dir()}/{LINK_FILENAME}') as fp:
        print('[child] Inspection #1', fp.read())

    print('[child] Telling parent we are ready')
    send_ready_message(conn)

    print('[child] Waiting for parent to perform adjustments')
    wait_for_ready_message(conn)

    with open(f'{get_test_dir()}/{LINK_FILENAME}') as fp:
        print('[child] Inspection #2:', fp.read())

    send_ready_message(conn)

def main():
    if len(sys.argv) == 1:
        parent()
    elif len(sys.argv) == 3:
        assert sys.argv[1] == 'child'
        child(int(sys.argv[2]))
    else:
        assert False

if __name__ == '__main__':
    main()