cachepc-qemu

Fork of AMDESE/qemu with changes for cachepc side-channel attack
git clone https://git.sinitax.com/sinitax/cachepc-qemu
Log | Files | Refs | Submodules | LICENSE | sfeed.txt

nbd-fault-injector.py (8638B)

      1#!/usr/bin/env python3
      2# NBD server - fault injection utility
      3#
      4# Configuration file syntax:
      5#   [inject-error "disconnect-neg1"]
      6#   event=neg1
      7#   io=readwrite
      8#   when=before
      9#
     10# Note that Python's ConfigParser squashes together all sections with the same
     11# name, so give each [inject-error] a unique name.
     12#
     13# inject-error options:
     14#   event - name of the trigger event
     15#           "neg1" - first part of negotiation struct
     16#           "export" - export struct
     17#           "neg2" - second part of negotiation struct
     18#           "request" - NBD request struct
     19#           "reply" - NBD reply struct
     20#           "data" - request/reply data
     21#   io    - I/O direction that triggers this rule:
     22#           "read", "write", or "readwrite"
     23#           default: readwrite
     24#   when  - after how many bytes to inject the fault
     25#           -1 - inject error after I/O
     26#           0 - inject error before I/O
     27#           integer - inject error after integer bytes
     28#           "before" - alias for 0
     29#           "after" - alias for -1
     30#           default: before
     31#
     32# Currently the only error injection action is to terminate the server process.
     33# This resets the TCP connection and thus forces the client to handle
     34# unexpected connection termination.
     35#
     36# Other error injection actions could be added in the future.
     37#
     38# Copyright Red Hat, Inc. 2014
     39#
     40# Authors:
     41#   Stefan Hajnoczi <stefanha@redhat.com>
     42#
     43# This work is licensed under the terms of the GNU GPL, version 2 or later.
     44# See the COPYING file in the top-level directory.
     45
     46import sys
     47import socket
     48import struct
     49import collections
     50import configparser
     51
     52FAKE_DISK_SIZE = 8 * 1024 * 1024 * 1024 # 8 GB
     53
     54# Protocol constants
     55NBD_CMD_READ = 0
     56NBD_CMD_WRITE = 1
     57NBD_CMD_DISC = 2
     58NBD_REQUEST_MAGIC = 0x25609513
     59NBD_SIMPLE_REPLY_MAGIC = 0x67446698
     60NBD_PASSWD = 0x4e42444d41474943
     61NBD_OPTS_MAGIC = 0x49484156454F5054
     62NBD_CLIENT_MAGIC = 0x0000420281861253
     63NBD_OPT_EXPORT_NAME = 1 << 0
     64
     65# Protocol structs
     66neg_classic_struct = struct.Struct('>QQQI124x')
     67neg1_struct = struct.Struct('>QQH')
     68export_tuple = collections.namedtuple('Export', 'reserved magic opt len')
     69export_struct = struct.Struct('>IQII')
     70neg2_struct = struct.Struct('>QH124x')
     71request_tuple = collections.namedtuple('Request', 'magic type handle from_ len')
     72request_struct = struct.Struct('>IIQQI')
     73reply_struct = struct.Struct('>IIQ')
     74
     75def err(msg):
     76    sys.stderr.write(msg + '\n')
     77    sys.exit(1)
     78
     79def recvall(sock, bufsize):
     80    received = 0
     81    chunks = []
     82    while received < bufsize:
     83        chunk = sock.recv(bufsize - received)
     84        if len(chunk) == 0:
     85            raise Exception('unexpected disconnect')
     86        chunks.append(chunk)
     87        received += len(chunk)
     88    return b''.join(chunks)
     89
     90class Rule(object):
     91    def __init__(self, name, event, io, when):
     92        self.name = name
     93        self.event = event
     94        self.io = io
     95        self.when = when
     96
     97    def match(self, event, io):
     98        if event != self.event:
     99            return False
    100        if io != self.io and self.io != 'readwrite':
    101            return False
    102        return True
    103
    104class FaultInjectionSocket(object):
    105    def __init__(self, sock, rules):
    106        self.sock = sock
    107        self.rules = rules
    108
    109    def check(self, event, io, bufsize=None):
    110        for rule in self.rules:
    111            if rule.match(event, io):
    112                if rule.when == 0 or bufsize is None:
    113                    print('Closing connection on rule match %s' % rule.name)
    114                    self.sock.close()
    115                    sys.stdout.flush()
    116                    sys.exit(0)
    117                if rule.when != -1:
    118                    return rule.when
    119        return bufsize
    120
    121    def send(self, buf, event):
    122        bufsize = self.check(event, 'write', bufsize=len(buf))
    123        self.sock.sendall(buf[:bufsize])
    124        self.check(event, 'write')
    125
    126    def recv(self, bufsize, event):
    127        bufsize = self.check(event, 'read', bufsize=bufsize)
    128        data = recvall(self.sock, bufsize)
    129        self.check(event, 'read')
    130        return data
    131
    132    def close(self):
    133        self.sock.close()
    134
    135def negotiate_classic(conn):
    136    buf = neg_classic_struct.pack(NBD_PASSWD, NBD_CLIENT_MAGIC,
    137                                  FAKE_DISK_SIZE, 0)
    138    conn.send(buf, event='neg-classic')
    139
    140def negotiate_export(conn):
    141    # Send negotiation part 1
    142    buf = neg1_struct.pack(NBD_PASSWD, NBD_OPTS_MAGIC, 0)
    143    conn.send(buf, event='neg1')
    144
    145    # Receive export option
    146    buf = conn.recv(export_struct.size, event='export')
    147    export = export_tuple._make(export_struct.unpack(buf))
    148    assert export.magic == NBD_OPTS_MAGIC
    149    assert export.opt == NBD_OPT_EXPORT_NAME
    150    name = conn.recv(export.len, event='export-name')
    151
    152    # Send negotiation part 2
    153    buf = neg2_struct.pack(FAKE_DISK_SIZE, 0)
    154    conn.send(buf, event='neg2')
    155
    156def negotiate(conn, use_export):
    157    '''Negotiate export with client'''
    158    if use_export:
    159        negotiate_export(conn)
    160    else:
    161        negotiate_classic(conn)
    162
    163def read_request(conn):
    164    '''Parse NBD request from client'''
    165    buf = conn.recv(request_struct.size, event='request')
    166    req = request_tuple._make(request_struct.unpack(buf))
    167    assert req.magic == NBD_REQUEST_MAGIC
    168    return req
    169
    170def write_reply(conn, error, handle):
    171    buf = reply_struct.pack(NBD_SIMPLE_REPLY_MAGIC, error, handle)
    172    conn.send(buf, event='reply')
    173
    174def handle_connection(conn, use_export):
    175    negotiate(conn, use_export)
    176    while True:
    177        req = read_request(conn)
    178        if req.type == NBD_CMD_READ:
    179            write_reply(conn, 0, req.handle)
    180            conn.send(b'\0' * req.len, event='data')
    181        elif req.type == NBD_CMD_WRITE:
    182            _ = conn.recv(req.len, event='data')
    183            write_reply(conn, 0, req.handle)
    184        elif req.type == NBD_CMD_DISC:
    185            break
    186        else:
    187            print('unrecognized command type %#02x' % req.type)
    188            break
    189    conn.close()
    190
    191def run_server(sock, rules, use_export):
    192    while True:
    193        conn, _ = sock.accept()
    194        handle_connection(FaultInjectionSocket(conn, rules), use_export)
    195
    196def parse_inject_error(name, options):
    197    if 'event' not in options:
    198        err('missing \"event\" option in %s' % name)
    199    event = options['event']
    200    if event not in ('neg-classic', 'neg1', 'export', 'neg2', 'request', 'reply', 'data'):
    201        err('invalid \"event\" option value \"%s\" in %s' % (event, name))
    202    io = options.get('io', 'readwrite')
    203    if io not in ('read', 'write', 'readwrite'):
    204        err('invalid \"io\" option value \"%s\" in %s' % (io, name))
    205    when = options.get('when', 'before')
    206    try:
    207        when = int(when)
    208    except ValueError:
    209        if when == 'before':
    210            when = 0
    211        elif when == 'after':
    212            when = -1
    213        else:
    214            err('invalid \"when\" option value \"%s\" in %s' % (when, name))
    215    return Rule(name, event, io, when)
    216
    217def parse_config(config):
    218    rules = []
    219    for name in config.sections():
    220        if name.startswith('inject-error'):
    221            options = dict(config.items(name))
    222            rules.append(parse_inject_error(name, options))
    223        else:
    224            err('invalid config section name: %s' % name)
    225    return rules
    226
    227def load_rules(filename):
    228    config = configparser.RawConfigParser()
    229    with open(filename, 'rt') as f:
    230        config.read_file(f, filename)
    231    return parse_config(config)
    232
    233def open_socket(path):
    234    '''Open a TCP or UNIX domain listen socket'''
    235    if ':' in path:
    236        host, port = path.split(':', 1)
    237        sock = socket.socket()
    238        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    239        sock.bind((host, int(port)))
    240
    241        # If given port was 0 the final port number is now available
    242        path = '%s:%d' % sock.getsockname()
    243    else:
    244        sock = socket.socket(socket.AF_UNIX)
    245        sock.bind(path)
    246    sock.listen(0)
    247    print('Listening on %s' % path)
    248    sys.stdout.flush() # another process may be waiting, show message now
    249    return sock
    250
    251def usage(args):
    252    sys.stderr.write('usage: %s [--classic-negotiation] <tcp-port>|<unix-path> <config-file>\n' % args[0])
    253    sys.stderr.write('Run an fault injector NBD server with rules defined in a config file.\n')
    254    sys.exit(1)
    255
    256def main(args):
    257    if len(args) != 3 and len(args) != 4:
    258        usage(args)
    259    use_export = True
    260    if args[1] == '--classic-negotiation':
    261        use_export = False
    262    elif len(args) == 4:
    263        usage(args)
    264    sock = open_socket(args[1 if use_export else 2])
    265    rules = load_rules(args[2 if use_export else 3])
    266    run_server(sock, rules, use_export)
    267    return 0
    268
    269if __name__ == '__main__':
    270    sys.exit(main(sys.argv))