| #!/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)) |
