265 lines
8.2 KiB
Python
Executable File
265 lines
8.2 KiB
Python
Executable File
#!/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))
|