]>
Commit | Line | Data |
---|---|---|
1e8ece0d SH |
1 | #!/usr/bin/env python |
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 | ||
f03868bd | 46 | from __future__ import print_function |
1e8ece0d SH |
47 | import sys |
48 | import socket | |
49 | import struct | |
50 | import collections | |
2d894bee HR |
51 | if sys.version_info.major >= 3: |
52 | import configparser | |
53 | else: | |
54 | import ConfigParser as configparser | |
1e8ece0d SH |
55 | |
56 | FAKE_DISK_SIZE = 8 * 1024 * 1024 * 1024 # 8 GB | |
57 | ||
58 | # Protocol constants | |
59 | NBD_CMD_READ = 0 | |
60 | NBD_CMD_WRITE = 1 | |
61 | NBD_CMD_DISC = 2 | |
62 | NBD_REQUEST_MAGIC = 0x25609513 | |
7b3158f9 | 63 | NBD_SIMPLE_REPLY_MAGIC = 0x67446698 |
1e8ece0d SH |
64 | NBD_PASSWD = 0x4e42444d41474943 |
65 | NBD_OPTS_MAGIC = 0x49484156454F5054 | |
66 | NBD_CLIENT_MAGIC = 0x0000420281861253 | |
67 | NBD_OPT_EXPORT_NAME = 1 << 0 | |
68 | ||
69 | # Protocol structs | |
70 | neg_classic_struct = struct.Struct('>QQQI124x') | |
71 | neg1_struct = struct.Struct('>QQH') | |
72 | export_tuple = collections.namedtuple('Export', 'reserved magic opt len') | |
73 | export_struct = struct.Struct('>IQII') | |
74 | neg2_struct = struct.Struct('>QH124x') | |
75 | request_tuple = collections.namedtuple('Request', 'magic type handle from_ len') | |
76 | request_struct = struct.Struct('>IIQQI') | |
77 | reply_struct = struct.Struct('>IIQ') | |
78 | ||
79 | def err(msg): | |
80 | sys.stderr.write(msg + '\n') | |
81 | sys.exit(1) | |
82 | ||
83 | def recvall(sock, bufsize): | |
84 | received = 0 | |
85 | chunks = [] | |
86 | while received < bufsize: | |
87 | chunk = sock.recv(bufsize - received) | |
88 | if len(chunk) == 0: | |
89 | raise Exception('unexpected disconnect') | |
90 | chunks.append(chunk) | |
91 | received += len(chunk) | |
8eb5e674 | 92 | return b''.join(chunks) |
1e8ece0d SH |
93 | |
94 | class Rule(object): | |
95 | def __init__(self, name, event, io, when): | |
96 | self.name = name | |
97 | self.event = event | |
98 | self.io = io | |
99 | self.when = when | |
100 | ||
101 | def match(self, event, io): | |
102 | if event != self.event: | |
103 | return False | |
104 | if io != self.io and self.io != 'readwrite': | |
105 | return False | |
106 | return True | |
107 | ||
108 | class FaultInjectionSocket(object): | |
109 | def __init__(self, sock, rules): | |
110 | self.sock = sock | |
111 | self.rules = rules | |
112 | ||
113 | def check(self, event, io, bufsize=None): | |
114 | for rule in self.rules: | |
115 | if rule.match(event, io): | |
116 | if rule.when == 0 or bufsize is None: | |
f03868bd | 117 | print('Closing connection on rule match %s' % rule.name) |
6d39db96 | 118 | self.sock.flush() |
1e8ece0d SH |
119 | sys.exit(0) |
120 | if rule.when != -1: | |
121 | return rule.when | |
122 | return bufsize | |
123 | ||
124 | def send(self, buf, event): | |
125 | bufsize = self.check(event, 'write', bufsize=len(buf)) | |
126 | self.sock.sendall(buf[:bufsize]) | |
127 | self.check(event, 'write') | |
128 | ||
129 | def recv(self, bufsize, event): | |
130 | bufsize = self.check(event, 'read', bufsize=bufsize) | |
131 | data = recvall(self.sock, bufsize) | |
132 | self.check(event, 'read') | |
133 | return data | |
134 | ||
135 | def close(self): | |
136 | self.sock.close() | |
137 | ||
138 | def negotiate_classic(conn): | |
139 | buf = neg_classic_struct.pack(NBD_PASSWD, NBD_CLIENT_MAGIC, | |
140 | FAKE_DISK_SIZE, 0) | |
141 | conn.send(buf, event='neg-classic') | |
142 | ||
143 | def negotiate_export(conn): | |
144 | # Send negotiation part 1 | |
145 | buf = neg1_struct.pack(NBD_PASSWD, NBD_OPTS_MAGIC, 0) | |
146 | conn.send(buf, event='neg1') | |
147 | ||
148 | # Receive export option | |
149 | buf = conn.recv(export_struct.size, event='export') | |
150 | export = export_tuple._make(export_struct.unpack(buf)) | |
151 | assert export.magic == NBD_OPTS_MAGIC | |
152 | assert export.opt == NBD_OPT_EXPORT_NAME | |
153 | name = conn.recv(export.len, event='export-name') | |
154 | ||
155 | # Send negotiation part 2 | |
156 | buf = neg2_struct.pack(FAKE_DISK_SIZE, 0) | |
157 | conn.send(buf, event='neg2') | |
158 | ||
159 | def negotiate(conn, use_export): | |
160 | '''Negotiate export with client''' | |
161 | if use_export: | |
162 | negotiate_export(conn) | |
163 | else: | |
164 | negotiate_classic(conn) | |
165 | ||
166 | def read_request(conn): | |
167 | '''Parse NBD request from client''' | |
168 | buf = conn.recv(request_struct.size, event='request') | |
169 | req = request_tuple._make(request_struct.unpack(buf)) | |
170 | assert req.magic == NBD_REQUEST_MAGIC | |
171 | return req | |
172 | ||
173 | def write_reply(conn, error, handle): | |
7b3158f9 | 174 | buf = reply_struct.pack(NBD_SIMPLE_REPLY_MAGIC, error, handle) |
1e8ece0d SH |
175 | conn.send(buf, event='reply') |
176 | ||
177 | def handle_connection(conn, use_export): | |
178 | negotiate(conn, use_export) | |
179 | while True: | |
180 | req = read_request(conn) | |
181 | if req.type == NBD_CMD_READ: | |
182 | write_reply(conn, 0, req.handle) | |
8eb5e674 | 183 | conn.send(b'\0' * req.len, event='data') |
1e8ece0d SH |
184 | elif req.type == NBD_CMD_WRITE: |
185 | _ = conn.recv(req.len, event='data') | |
186 | write_reply(conn, 0, req.handle) | |
187 | elif req.type == NBD_CMD_DISC: | |
188 | break | |
189 | else: | |
f03868bd | 190 | print('unrecognized command type %#02x' % req.type) |
1e8ece0d SH |
191 | break |
192 | conn.close() | |
193 | ||
194 | def run_server(sock, rules, use_export): | |
195 | while True: | |
196 | conn, _ = sock.accept() | |
197 | handle_connection(FaultInjectionSocket(conn, rules), use_export) | |
198 | ||
199 | def parse_inject_error(name, options): | |
200 | if 'event' not in options: | |
201 | err('missing \"event\" option in %s' % name) | |
202 | event = options['event'] | |
203 | if event not in ('neg-classic', 'neg1', 'export', 'neg2', 'request', 'reply', 'data'): | |
204 | err('invalid \"event\" option value \"%s\" in %s' % (event, name)) | |
205 | io = options.get('io', 'readwrite') | |
206 | if io not in ('read', 'write', 'readwrite'): | |
207 | err('invalid \"io\" option value \"%s\" in %s' % (io, name)) | |
208 | when = options.get('when', 'before') | |
209 | try: | |
210 | when = int(when) | |
211 | except ValueError: | |
212 | if when == 'before': | |
213 | when = 0 | |
214 | elif when == 'after': | |
215 | when = -1 | |
216 | else: | |
217 | err('invalid \"when\" option value \"%s\" in %s' % (when, name)) | |
218 | return Rule(name, event, io, when) | |
219 | ||
220 | def parse_config(config): | |
221 | rules = [] | |
222 | for name in config.sections(): | |
223 | if name.startswith('inject-error'): | |
224 | options = dict(config.items(name)) | |
225 | rules.append(parse_inject_error(name, options)) | |
226 | else: | |
227 | err('invalid config section name: %s' % name) | |
228 | return rules | |
229 | ||
230 | def load_rules(filename): | |
2d894bee | 231 | config = configparser.RawConfigParser() |
1e8ece0d SH |
232 | with open(filename, 'rt') as f: |
233 | config.readfp(f, filename) | |
234 | return parse_config(config) | |
235 | ||
236 | def open_socket(path): | |
237 | '''Open a TCP or UNIX domain listen socket''' | |
238 | if ':' in path: | |
239 | host, port = path.split(':', 1) | |
240 | sock = socket.socket() | |
241 | sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | |
242 | sock.bind((host, int(port))) | |
6e592fc9 SH |
243 | |
244 | # If given port was 0 the final port number is now available | |
245 | path = '%s:%d' % sock.getsockname() | |
1e8ece0d SH |
246 | else: |
247 | sock = socket.socket(socket.AF_UNIX) | |
248 | sock.bind(path) | |
249 | sock.listen(0) | |
f03868bd | 250 | print('Listening on %s' % path) |
6e592fc9 | 251 | sys.stdout.flush() # another process may be waiting, show message now |
1e8ece0d SH |
252 | return sock |
253 | ||
254 | def usage(args): | |
255 | sys.stderr.write('usage: %s [--classic-negotiation] <tcp-port>|<unix-path> <config-file>\n' % args[0]) | |
256 | sys.stderr.write('Run an fault injector NBD server with rules defined in a config file.\n') | |
257 | sys.exit(1) | |
258 | ||
259 | def main(args): | |
260 | if len(args) != 3 and len(args) != 4: | |
261 | usage(args) | |
262 | use_export = True | |
263 | if args[1] == '--classic-negotiation': | |
264 | use_export = False | |
265 | elif len(args) == 4: | |
266 | usage(args) | |
267 | sock = open_socket(args[1 if use_export else 2]) | |
268 | rules = load_rules(args[2 if use_export else 3]) | |
269 | run_server(sock, rules, use_export) | |
270 | return 0 | |
271 | ||
272 | if __name__ == '__main__': | |
273 | sys.exit(main(sys.argv)) |