1 files changed, 264 insertions, 0 deletions
diff --git a/qemu/tests/qemu-iotests/nbd-fault-injector.py b/qemu/tests/qemu-iotests/nbd-fault-injector.py
new file mode 100755
index 000000000..6c07191a5
--- /dev/null
+++ b/qemu/tests/qemu-iotests/nbd-fault-injector.py
@@ -0,0 +1,264 @@
+#!/usr/bin/env python
+# NBD server - fault injection utility
+#
+# Configuration file syntax:
+#   [inject-error "disconnect-neg1"]
+#   event=neg1
+#   io=readwrite
+#   when=before
+#
+# Note that Python's ConfigParser squashes together all sections with the same
+# name, so give each [inject-error] a unique name.
+#
+# inject-error options:
+#   event - name of the trigger event
+#           "neg1" - first part of negotiation struct
+#           "export" - export struct
+#           "neg2" - second part of negotiation struct
+#           "request" - NBD request struct
+#           "reply" - NBD reply struct
+#           "data" - request/reply data
+#   io    - I/O direction that triggers this rule:
+#           "read", "write", or "readwrite"
+#           default: readwrite
+#   when  - after how many bytes to inject the fault
+#           -1 - inject error after I/O
+#           0 - inject error before I/O
+#           integer - inject error after integer bytes
+#           "before" - alias for 0
+#           "after" - alias for -1
+#           default: before
+#
+# Currently the only error injection action is to terminate the server process.
+# This resets the TCP connection and thus forces the client to handle
+# unexpected connection termination.
+#
+# Other error injection actions could be added in the future.
+#
+# Copyright Red Hat, Inc. 2014
+#
+# Authors:
+#   Stefan Hajnoczi <stefanha@redhat.com>
+#
+# This work is licensed under the terms of the GNU GPL, version 2 or later.
+# See the COPYING file in the top-level directory.
+
+import sys
+import socket
+import struct
+import collections
+import ConfigParser
+
+FAKE_DISK_SIZE = 8 * 1024 * 1024 * 1024 # 8 GB
+
+# Protocol constants
+NBD_CMD_READ = 0
+NBD_CMD_WRITE = 1
+NBD_CMD_DISC = 2
+NBD_REQUEST_MAGIC = 0x25609513
+NBD_REPLY_MAGIC = 0x67446698
+NBD_PASSWD = 0x4e42444d41474943
+NBD_OPTS_MAGIC = 0x49484156454F5054
+NBD_CLIENT_MAGIC = 0x0000420281861253
+NBD_OPT_EXPORT_NAME = 1 << 0
+
+# Protocol structs
+neg_classic_struct = struct.Struct('>QQQI124x')
+neg1_struct = struct.Struct('>QQH')
+export_tuple = collections.namedtuple('Export', 'reserved magic opt len')
+export_struct = struct.Struct('>IQII')
+neg2_struct = struct.Struct('>QH124x')
+request_tuple = collections.namedtuple('Request', 'magic type handle from_ len')
+request_struct = struct.Struct('>IIQQI')
+reply_struct = struct.Struct('>IIQ')
+
+def err(msg):
+    sys.stderr.write(msg + '\n')
+    sys.exit(1)
+
+def recvall(sock, bufsize):
+    received = 0
+    chunks = []
+    while received < bufsize:
+        chunk = sock.recv(bufsize - received)
+        if len(chunk) == 0:
+            raise Exception('unexpected disconnect')
+        chunks.append(chunk)
+        received += len(chunk)
+    return ''.join(chunks)
+
+class Rule(object):
+    def __init__(self, name, event, io, when):
+        self.name = name
+        self.event = event
+        self.io = io
+        self.when = when
+
+    def match(self, event, io):
+        if event != self.event:
+            return False
+        if io != self.io and self.io != 'readwrite':
+            return False
+        return True
+
+class FaultInjectionSocket(object):
+    def __init__(self, sock, rules):
+        self.sock = sock
+        self.rules = rules
+
+    def check(self, event, io, bufsize=None):
+        for rule in self.rules:
+            if rule.match(event, io):
+                if rule.when == 0 or bufsize is None:
+                    print 'Closing connection on rule match %s' % rule.name
+                    sys.exit(0)
+                if rule.when != -1:
+                    return rule.when
+        return bufsize
+
+    def send(self, buf, event):
+        bufsize = self.check(event, 'write', bufsize=len(buf))
+        self.sock.sendall(buf[:bufsize])
+        self.check(event, 'write')
+
+    def recv(self, bufsize, event):
+        bufsize = self.check(event, 'read', bufsize=bufsize)
+        data = recvall(self.sock, bufsize)
+        self.check(event, 'read')
+        return data
+
+    def close(self):
+        self.sock.close()
+
+def negotiate_classic(conn):
+    buf = neg_classic_struct.pack(NBD_PASSWD, NBD_CLIENT_MAGIC,
+                                  FAKE_DISK_SIZE, 0)
+    conn.send(buf, event='neg-classic')
+
+def negotiate_export(conn):
+    # Send negotiation part 1
+    buf = neg1_struct.pack(NBD_PASSWD, NBD_OPTS_MAGIC, 0)
+    conn.send(buf, event='neg1')
+
+    # Receive export option
+    buf = conn.recv(export_struct.size, event='export')
+    export = export_tuple._make(export_struct.unpack(buf))
+    assert export.magic == NBD_OPTS_MAGIC
+    assert export.opt == NBD_OPT_EXPORT_NAME
+    name = conn.recv(export.len, event='export-name')
+
+    # Send negotiation part 2
+    buf = neg2_struct.pack(FAKE_DISK_SIZE, 0)
+    conn.send(buf, event='neg2')
+
+def negotiate(conn, use_export):
+    '''Negotiate export with client'''
+    if use_export:
+        negotiate_export(conn)
+    else:
+        negotiate_classic(conn)
+
+def read_request(conn):
+    '''Parse NBD request from client'''
+    buf = conn.recv(request_struct.size, event='request')
+    req = request_tuple._make(request_struct.unpack(buf))
+    assert req.magic == NBD_REQUEST_MAGIC
+    return req
+
+def write_reply(conn, error, handle):
+    buf = reply_struct.pack(NBD_REPLY_MAGIC, error, handle)
+    conn.send(buf, event='reply')
+
+def handle_connection(conn, use_export):
+    negotiate(conn, use_export)
+    while True:
+        req = read_request(conn)
+        if req.type == NBD_CMD_READ:
+            write_reply(conn, 0, req.handle)
+            conn.send('\0' * req.len, event='data')
+        elif req.type == NBD_CMD_WRITE:
+            _ = conn.recv(req.len, event='data')
+            write_reply(conn, 0, req.handle)
+        elif req.type == NBD_CMD_DISC:
+            break
+        else:
+            print 'unrecognized command type %#02x' % req.type
+            break
+    conn.close()
+
+def run_server(sock, rules, use_export):
+    while True:
+        conn, _ = sock.accept()
+        handle_connection(FaultInjectionSocket(conn, rules), use_export)
+
+def parse_inject_error(name, options):
+    if 'event' not in options:
+        err('missing \"event\" option in %s' % name)
+    event = options['event']
+    if event not in ('neg-classic', 'neg1', 'export', 'neg2', 'request', 'reply', 'data'):
+        err('invalid \"event\" option value \"%s\" in %s' % (event, name))
+    io = options.get('io', 'readwrite')
+    if io not in ('read', 'write', 'readwrite'):
+        err('invalid \"io\" option value \"%s\" in %s' % (io, name))
+    when = options.get('when', 'before')
+    try:
+        when = int(when)
+    except ValueError:
+        if when == 'before':
+            when = 0
+        elif when == 'after':
+            when = -1
+        else:
+            err('invalid \"when\" option value \"%s\" in %s' % (when, name))
+    return Rule(name, event, io, when)
+
+def parse_config(config):
+    rules = []
+    for name in config.sections():
+        if name.startswith('inject-error'):
+            options = dict(config.items(name))
+            rules.append(parse_inject_error(name, options))
+        else:
+            err('invalid config section name: %s' % name)
+    return rules
+
+def load_rules(filename):
+    config = ConfigParser.RawConfigParser()
+    with open(filename, 'rt') as f:
+        config.readfp(f, filename)
+    return parse_config(config)
+
+def open_socket(path):
+    '''Open a TCP or UNIX domain listen socket'''
+    if ':' in path:
+        host, port = path.split(':', 1)
+        sock = socket.socket()
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind((host, int(port)))
+    else:
+        sock = socket.socket(socket.AF_UNIX)
+        sock.bind(path)
+    sock.listen(0)
+    print 'Listening on %s' % path
+    return sock
+
+def usage(args):
+    sys.stderr.write('usage: %s [--classic-negotiation] <tcp-port>|<unix-path> <config-file>\n' % args[0])
+    sys.stderr.write('Run an fault injector NBD server with rules defined in a config file.\n')
+    sys.exit(1)
+
+def main(args):
+    if len(args) != 3 and len(args) != 4:
+        usage(args)
+    use_export = True
+    if args[1] == '--classic-negotiation':
+        use_export = False
+    elif len(args) == 4:
+        usage(args)
+    sock = open_socket(args[1 if use_export else 2])
+    rules = load_rules(args[2 if use_export else 3])
+    run_server(sock, rules, use_export)
+    return 0
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))