History

Attempted to improve handling of different format fields, treating certain types as special cases and allowing lengths to refer to earlier fields in a format list.

David Boddie committed 8 years ago

Commit 6add7821 parent 4da8c8f1 file · +264 −138Browse files

Unified · Split

styx.py+264−138
@@ -76,6 +76,67 @@ def encode_qid(qid): · 75 unmodified lines
76 76
77 return struct.pack("<BIQ", *qid)77 return struct.pack("<BIQ", *qid)
78 78
79def decode_format(stream, obj):
80
81 for name, length in obj.format:
82
83 if length == "s":
84 # UTF-8 string
85 value = decode_string(stream)
86 elif length == "n":
87 # n bytes of raw data
88 value = decode_data(stream)
89 elif length == 1:
90 value = struct.unpack("<B", stream.recv(1))[0]
91 elif length == 2:
92 value = struct.unpack("<H", stream.recv(2))[0]
93 elif length == 4:
94 value = struct.unpack("<I", stream.recv(4))[0]
95 elif length == 8:
96 value = struct.unpack("<Q", stream.recv(8))[0]
97 elif name == "qid":
98 value = decode_qid(stream)
99 elif name == "stat":
100 value = Stat()
101 value.decode(stream)
102 elif length in obj.__dict__:
103 value = stream.recv(obj.__dict__[length])
104 else:
105 raise StyxError("decode_format: Unknown field length specifier: %s" % length)
106
107 obj.__dict__[name] = value
108
109def encode_format(stream, obj):
110
111 data = ""
112
113 for name, length in obj.format:
114
115 value = obj.__dict__[name]
116
117 if length == "s":
118 # UTF-8 string
119 data += encode_string(value)
120 elif length == "n":
121 # n bytes of raw data
122 data += encode_data(value)
123 elif length == 1:
124 data += struct.pack("<B", value)
125 elif length == 2:
126 data += struct.pack("<H", value)
127 elif length == 4:
128 data += struct.pack("<I", value)
129 elif length == 8:
130 data += struct.pack("<Q", value)
131 elif name == "qid":
132 data += encode_qid(value)
133 elif name == "stat":
134 data += value.encode()
135 else:
136 data += value
137
138 return data
139
79 140
80class StyxError(Exception):141class StyxError(Exception):
81 pass142 pass
@@ -88,33 +149,6 @@ class StyxMessage: · 6 unmodified lines
88 self.size = size149 self.size = size
89 self.tag = tag150 self.tag = tag
90 151
91 def decode(self, size, tag, stream):
92
93 self.init(size, tag)
94
95 for name, length in self.format:
96
97 if length == "s":
98 # UTF-8 string
99 value = decode_string(stream)
100 elif length == "n":
101 # n bytes of raw data
102 value = decode_data(stream)
103 elif length == 2:
104 value = struct.unpack("<H", stream.recv(2))[0]
105 elif length == 4:
106 value = struct.unpack("<I", stream.recv(4))[0]
107 elif length == 8:
108 value = struct.unpack("<Q", stream.recv(8))[0]
109 elif name == "qid":
110 value = decode_qid(stream)
111 else:
112 raise StyxError("Unknown field length specifier: %s" % length)
113
114 self.__dict__[name] = value
115
116 return self
117
118 def __repr__(self):152 def __repr__(self):
119 153
120 s = self.msg_name + "(tag=" + repr(self.tag)154 s = self.msg_name + "(tag=" + repr(self.tag)
@@ -124,41 +158,21 @@ class StyxMessage: · 3 unmodified lines
124 158
125 return s + ")"159 return s + ")"
126 160
161 def decode(self, size, tag, stream):
162
163 self.init(size, tag)
164 decode_format(stream, self)
165 return self
166
127 def encode(self, stream):167 def encode(self, stream):
128 168
129 # Start with the message type and tag before adding all the other pieces.169 # Start with the message type and tag before adding all the other pieces.
130 data = struct.pack("<bH", self.code, self.tag)170 data = struct.pack("<BH", self.code, self.tag)
131 171 data += encode_format(stream, self)
132 data += self._encode_format(stream)
133 172
134 # Finally, prepend the length to the byte string and send it.173 # Finally, prepend the length to the byte string and send it.
135 data = struct.pack("<I", len(data) + 4) + data174 data = struct.pack("<I", len(data) + 4) + data
136 stream.sendall(data)175 stream.sendall(data)
137
138 def _encode_format(self, stream):
139
140 data = ""
141
142 for name, length in self.format:
143
144 value = self.__dict__[name]
145
146 if length == "s":
147 # UTF-8 string
148 data += encode_string(value)
149 elif length == "n":
150 # n bytes of raw data
151 data += encode_data(value)
152 elif length == 2:
153 data += struct.pack("<H", value)
154 elif length == 4:
155 data += struct.pack("<I", value)
156 elif length == 8:
157 data += struct.pack("<Q", value)
158 elif name == "qid":
159 data += encode_qid(value)
160
161 return data
162 176
163 177
164class Tversion(StyxMessage):178class Tversion(StyxMessage):
@@ -212,77 +226,16 @@ class Rattach(StyxMessage): · 47 unmodified lines
212 self.qid = qid226 self.qid = qid
213 227
214 228
215class Tstat(StyxMessage):229class Rerror(StyxMessage):
216
217 msg_name = "Tstat"
218 code = 124
219 format = [("fid", 4)]
220
221 def __init__(self, tag = None, fid = None):
222
223 self.tag = tag
224 self.fid = fid
225
226class Rstat(StyxMessage):
227 230
228 msg_name = "Rstat"231 msg_name = "Rerror"
229 code = 125232 code = 107
230 # Because the format of this message describes the stat information as233 format = [("ename", "s")]
231 # stat[n] it includes a 16-bit length, but the stat object itself
232 # includes its own length field.
233 format = [
234 ("stat_size", 2), ("size", 2), ("type", 2), ("dev", 4), ("qid", 13),
235 ("mode", 4), ("atime", 4), ("mtime", 4), ("length", 8), ("name", "s"),
236 ("uid", "s"), ("gid", "s"), ("muid", "s")
237 ]
238 234
239 def __init__(self, tag = None, stat = None):235 def __init__(self, tag = None, ename = None):
240 236
241 self.tag = tag237 self.tag = tag
242 self.stat = stat238 self.ename = ename
243
244 def encode(self, stream):
245
246 data = struct.pack("<bH", self.code, self.tag)
247
248 # Encode the stat structure.
249 data += self.stat.encode()
250
251 data = struct.pack("<I", len(data) + 4) + data
252 stream.sendall(data)
253
254
255class Stat:
256
257 def __init__(self, type, dev, qid, mode, atime, mtime, length, name, uid,
258 gid, muid):
259
260 self.type = type
261 self.dev = dev
262 self.qid = qid
263 self.mode = mode
264 self.atime = atime
265 self.mtime = mtime
266 self.length = length
267 self.name = name
268 self.uid = uid
269 self.gid = gid
270 self.muid = muid
271
272 def encode(self):
273
274 args = (self.type, self.dev) + self.qid + (self.mode, self.atime, self.mtime, self.length)
275 data = struct.pack("<HIBIQIIIQ", *args)
276 data += encode_string(self.name)
277 data += encode_string(self.uid)
278 data += encode_string(self.gid)
279 data += encode_string(self.muid)
280
281 # Prepend the size of the fields within the stat structure.
282 data = encode_data(data)
283
284 # Prepend the size of the stat structure itself.
285 return encode_data(data)
286 239
287 240
288class Twalk(StyxMessage):241class Twalk(StyxMessage):
@@ -291,12 +244,13 @@ class Twalk(StyxMessage): · 2 unmodified lines
291 code = 110244 code = 110
292 format = [("fid", 4), ("newfid", 4), ("nwname", 2)] # nwname of wname[s]245 format = [("fid", 4), ("newfid", 4), ("nwname", 2)] # nwname of wname[s]
293 246
294 def __init__(self, tag = None, fid = None, newfid = None, wname = None):247 def __init__(self, tag = None, fid = None, newfid = None, wname = []):
295 248
296 self.tag = tag249 self.tag = tag
297 self.fid = fid250 self.fid = fid
298 self.newfid = newfid251 self.newfid = newfid
299 self.wname = wname252 self.wname = wname
253 self.nwname = len(wname)
300 254
301 def decode(self, size, tag, stream):255 def decode(self, size, tag, stream):
302 256
@@ -309,14 +263,16 @@ class Twalk(StyxMessage): · 6 unmodified lines
309 self.wname.append(decode_string(stream))263 self.wname.append(decode_string(stream))
310 i += 1264 i += 1
311 265
266 self.format.append(("wname", "nwname"))
267
312 return self268 return self
313 269
314 def encode(self, stream):270 def encode(self, stream):
315 271
316 data = struct.pack("<bH", self.code, self.tag)272 data = struct.pack("<BH", self.code, self.tag)
317 273
318 self.nwname = len(self.wname)274 self.nwname = len(self.wname)
319 data += self._encode_format(stream)275 data += encode_format(stream, self)
320 276
321 # Encode the names separately.277 # Encode the names separately.
322 for name in self.wname:278 for name in self.wname:
@@ -325,17 +281,17 @@ class Twalk(StyxMessage): · 2 unmodified lines
325 data = struct.pack("<I", len(data) + 4) + data281 data = struct.pack("<I", len(data) + 4) + data
326 stream.sendall(data)282 stream.sendall(data)
327 283
328
329class Rwalk(StyxMessage):284class Rwalk(StyxMessage):
330 285
331 msg_name = "Rwalk"286 msg_name = "Rwalk"
332 code = 111287 code = 111
333 format = [("nwqid", 2)] # nwqid of wqid[13]288 format = [("nwqid", 2)] # nwqid of wqid[13]
334 289
335 def __init__(self, tag = None, wqid = None):290 def __init__(self, tag = None, wqid = []):
336 291
337 self.tag = tag292 self.tag = tag
338 self.wqid = wqid293 self.wqid = wqid
294 self.nwqid = len(wqid)
339 295
340 def decode(self, size, tag, stream):296 def decode(self, size, tag, stream):
341 297
@@ -348,14 +304,16 @@ class Rwalk(StyxMessage): · 6 unmodified lines
348 self.wqid.append(decode_qid(stream))304 self.wqid.append(decode_qid(stream))
349 i += 1305 i += 1
350 306
307 self.format.append(("wqid", "nwqid"))
308
351 return self309 return self
352 310
353 def encode(self, stream):311 def encode(self, stream):
354 312
355 data = struct.pack("<bH", self.code, self.tag)313 data = struct.pack("<BH", self.code, self.tag)
356 314
357 self.nwqid = len(self.wqid)315 self.nwqid = len(self.wqid)
358 data += self._encode_format(stream)316 data += encode_format(stream, self)
359 317
360 # Encode the qids separately.318 # Encode the qids separately.
361 for qid in self.wqid:319 for qid in self.wqid:
@@ -365,6 +323,171 @@ class Rwalk(StyxMessage): · 3 unmodified lines
365 stream.sendall(data)323 stream.sendall(data)
366 324
367 325
326class Topen(StyxMessage):
327
328 msg_name = "Topen"
329 code = 112
330 format = [("fid", 4), ("mode", 1)]
331
332 def __init__(self, tag = None, fid = None, mode = None):
333
334 self.tag = tag
335 self.fid = fid
336 self.mode = mode
337
338class Ropen(StyxMessage):
339
340 msg_name = "Ropen"
341 code = 113
342 format = [("qid", 13), ("iounit", 4)]
343
344 def __init__(self, tag = None, qid = None, iounit = None):
345
346 self.tag = tag
347 self.qid = qid
348 self.iounit = iounit
349
350
351class Tread(StyxMessage):
352
353 msg_name = "Tread"
354 code = 116
355 format = [("fid", 4), ("offset", 8), ("count", 4)]
356
357 def __init__(self, tag = None, fid = None, offset = 0, count = 0):
358
359 self.tag = tag
360 self.fid = fid
361 self.offset = offset
362 self.count = count
363
364class Rread(StyxMessage):
365
366 msg_name = "Rread"
367 code = 117
368 format = [("count", 4), ("data", "count")]
369
370 def __init__(self, tag = None, data = ""):
371
372 self.tag = tag
373 self.data = data
374 self.count = len(data)
375
376
377class Tclunk(StyxMessage):
378
379 msg_name = "Tclunk"
380 code = 120
381 format = [("fid", 4)]
382
383 def __init__(self, tag = None, fid = None):
384
385 self.tag = tag
386 self.fid = fid
387
388class Rclunk(StyxMessage):
389
390 msg_name = "Rclunk"
391 code = 121
392 format = []
393
394 def __init__(self, tag = None):
395
396 self.tag = tag
397
398
399class Tstat(StyxMessage):
400
401 msg_name = "Tstat"
402 code = 124
403 format = [("fid", 4)]
404
405 def __init__(self, tag = None, fid = None):
406
407 self.tag = tag
408 self.fid = fid
409
410class Rstat(StyxMessage):
411
412 msg_name = "Rstat"
413 code = 125
414 # Because the format of this message describes the stat information as
415 # stat[n] it includes a 16-bit length, but the stat object itself
416 # includes its own length field.
417 format = [("stat_size", 2), ("stat", "stat")]
418
419 def __init__(self, tag = None, stat = None):
420
421 self.tag = tag
422 self.stat = stat
423
424 def __repr__(self):
425
426 return self.msg_name + "(tag=%s, stat=%s)" % (repr(self.tag), repr(self.stat))
427
428 def encode(self, stream):
429
430 # Encode the stat structure.
431 stat_data = self.stat.encode()
432
433 # Prepend the size of the stat structure itself.
434 stat_data = encode_data(stat_data)
435
436 data = struct.pack("<BH", self.code, self.tag) + stat_data
437
438 data = struct.pack("<I", len(data) + 4) + data
439 stream.sendall(data)
440
441
442class Stat:
443
444 format = [
445 ("size", 2), ("type", 2), ("dev", 4), ("qid", 13), ("mode", 4),
446 ("atime", 4), ("mtime", 4), ("length", 8), ("name", "s"), ("uid", "s"),
447 ("gid", "s"), ("muid", "s")
448 ]
449
450 def __init__(self, type = 0, dev = 0, qid = (0, 0, 0), mode = 0,
451 atime = 0, mtime = 0, length = 0, name = u"", uid = u"",
452 gid = "", muid = u""):
453
454 self.type = type
455 self.dev = dev
456 self.qid = qid
457 self.mode = mode
458 self.atime = atime
459 self.mtime = mtime
460 self.length = length
461 self.name = name
462 self.uid = uid
463 self.gid = gid
464 self.muid = muid
465
466 def __repr__(self):
467
468 pieces = []
469 for name, length in self.format:
470 pieces.append("%s=%s" % (name, repr(self.__dict__[name])))
471
472 return "Stat(" + ", ".join(pieces) + ")"
473
474 def decode(self, stream):
475
476 decode_format(stream, self)
477
478 def encode(self):
479
480 args = (self.type, self.dev) + self.qid + (self.mode, self.atime, self.mtime, self.length)
481 data = struct.pack("<HIBIQIIIQ", *args)
482 data += encode_string(self.name)
483 data += encode_string(self.uid)
484 data += encode_string(self.gid)
485 data += encode_string(self.muid)
486
487 # Prepend the size of the fields within the stat structure.
488 return encode_data(data)
489
490
368MessageTypes = {491MessageTypes = {
369 Tversion.code: Tversion, 492 Tversion.code: Tversion,
370 Rversion.code: Rversion,493 Rversion.code: Rversion,
@@ -373,21 +496,21 @@ MessageTypes = { · 2 unmodified lines
373 Tattach.code: Tattach,496 Tattach.code: Tattach,
374 Rattach.code: Rattach,497 Rattach.code: Rattach,
375# Terror.code: Terror,498# Terror.code: Terror,
376# Rerror.code: Rerror,499 Rerror.code: Rerror,
377# Tflush.code: Tflush,500# Tflush.code: Tflush,
378# Rflush.code: Rflush,501# Rflush.code: Rflush,
379 Twalk.code: Twalk,502 Twalk.code: Twalk,
380 Rwalk.code: Rwalk,503 Rwalk.code: Rwalk,
381# Topen.code: Topen,504 Topen.code: Topen,
382# Ropen.code: Ropen,505 Ropen.code: Ropen,
383# Tcreate.code: Tcreate,506# Tcreate.code: Tcreate,
384# Rcreate.code: Rcreate,507# Rcreate.code: Rcreate,
385# Tread.code: Tread,508 Tread.code: Tread,
386# Rread.code: Rread,509 Rread.code: Rread,
387# Twrite.code: Twrite,510# Twrite.code: Twrite,
388# Rwrite.code: Rwrite,511# Rwrite.code: Rwrite,
389# Tclunk.code: Tclunk,512 Tclunk.code: Tclunk,
390# Rclunk.code: Rclunk,513 Rclunk.code: Rclunk,
391# Tremove.code: Tremove,514# Tremove.code: Tremove,
392# Rremove.code: Rremove,515# Rremove.code: Rremove,
393 Tstat.code: Tstat,516 Tstat.code: Tstat,
@@ -412,5 +535,8 @@ def decode(sock = None, data = None): · 18 unmodified lines
412 535
413 # Find the relevant message class to handle this type and create an536 # Find the relevant message class to handle this type and create an
414 # instance of it to parse the message data.537 # instance of it to parse the message data.
415 Message = MessageTypes[message_type]538 try:
416 return Message().decode(size, tag, stream)539 Message = MessageTypes[message_type]
540 return Message().decode(size, tag, stream)
541 except:
542 return size, message_type, tag, stream