src/add-ons/kernel/drivers/disk/virtual/nbd/nbd-server.py

   1 #!/usr/bin/python
   2 # from http://lists.canonical.org/pipermail/kragen-hacks/2004-May/000397.html
   3 import struct, socket, sys
   4 # network block device server, substitute for nbd-server.  Probably slower.
   5 # But it works!  And it's probably a lot easier to improve the
   6 # performance of this Python version than of the C version.  This
   7 # Python version is 14% of the size and perhaps 20% of the features of
   8 # the C version.  Hmm, that's not so great after all...
   9 # Working:
  10 # - nbd protocol
  11 # - read/write serving up files
  12 # - error handling
  13 # - file size detection
  14 # - in theory, large file support... not really
  15 # - so_reuseaddr
  16 # - nonforking
  17 # Missing:
  18 # - reporting errors to client (in particular writing and reading past end)
  19 # - multiple clients (this probably requires copy-on-write or read-only)
  20 # - copy on write
  21 # - read-only
  22 # - permission tracking
  23 # - idle timeouts
  24 # - running from inetd
  25 # - filename substitution
  26 # - partial file exports
  27 # - exports of large files (bigger than 1/4 of RAM)
  28 # - manual exportsize specification
  29 # - so_keepalive
  30 # - that "split an export file into multiple files" thing that sticks the .0
  31 #   on the end of your filename
  32 # - backgrounding
  33 # - daemonizing
  34
  35 class Error(Exception): pass
  36
  37 class buffsock:
  38     "Buffered socket wrapper; always returns the amount of data you want."
  39     def __init__(self, sock): self.sock = sock
  40     def recv(self, nbytes):
  41         rv = ''
  42         while len(rv) < nbytes:
  43             more = self.sock.recv(nbytes - len(rv))
  44             if more == '': raise Error(nbytes)
  45             rv += more
  46         return rv
  47     def send(self, astring): self.sock.send(astring)
  48     def close(self): self.sock.close()
  49
  50
  51 class debugsock:
  52     "Debugging socket wrapper."
  53     def __init__(self, sock): self.sock = sock
  54     def recv(self, nbytes):
  55         print "recv(%d) =" % nbytes,
  56         rv = self.sock.recv(nbytes)
  57         print `rv`
  58         return rv
  59     def send(self, astring):
  60         print "send(%r) =" % astring,
  61         rv = self.sock.send(astring)
  62         print `rv`
  63         return rv
  64     def close(self):
  65         print "close()"
  66         self.sock.close()
  67
  68 def negotiation(exportsize):
  69     "Returns initial NBD negotiation sequence for exportsize in bytes."
  70     return ('NBDMAGIC' + '\x00\x00\x42\x02\x81\x86\x12\x53' +
  71         struct.pack('>Q', exportsize) + '\0' * 128);
  72
  73 def nbd_reply(error=0, handle=1, data=''):
  74     "Construct an NBD reply."
  75     assert type(handle) is type('') and len(handle) == 8
  76     return ('\x67\x44\x66\x98' + struct.pack('>L', error) + handle + data)
  77
  78 # possible request types
  79 read_request = 0
  80 write_request = 1
  81 disconnect_request = 2
  82
  83 class nbd_request:
  84     "Decodes an NBD request off the TCP socket."
  85     def __init__(self, conn):
  86         conn = buffsock(conn)
  87         template = '>LL8sQL'
  88         header = conn.recv(struct.calcsize(template))
  89         (self.magic, self.type, self.handle, self.offset,
  90             self.len) = struct.unpack(template, header)
  91         if self.magic != 0x25609513: raise Error(self.magic)
  92         if self.type == write_request:
  93             self.data = conn.recv(self.len)
  94             assert len(self.data) == self.len
  95     def reply(self, error, data=''):
  96         return nbd_reply(error=error, handle=self.handle, data=data)
  97     def range(self):
  98         return slice(self.offset, self.offset + self.len)
  99
 100 def serveclient(asock, afile):
 101     "Serves a single client until it exits."
 102     afile.seek(0)
 103     abuf = list(afile.read())
 104     asock.send(negotiation(len(abuf)))
 105     while 1:
 106         req = nbd_request(asock)
 107         if req.type == read_request:
 108             asock.send(req.reply(error=0,
 109                 data=''.join(abuf[req.range()])))
 110         elif req.type == write_request:
 111             abuf[req.range()] = req.data
 112             afile.seek(req.offset)
 113             afile.write(req.data)
 114             afile.flush()
 115             asock.send(req.reply(error=0))
 116         elif req.type == disconnect_request:
 117             asock.close()
 118             return
 119
 120 def mainloop(listensock, afile):
 121     "Serves clients forever."
 122     while 1:
 123         (sock, addr) = listensock.accept()
 124         print "got conn on", addr
 125         serveclient(sock, afile)
 126
 127 def main(argv):
 128     "Given a port and a filename, serves up the file."
 129     afile = file(argv[2], 'rb+')
 130     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 131     sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 132     sock.bind(('', int(argv[1])))
 133     sock.listen(5)
 134     mainloop(sock, afile)
 135
 136 if __name__ == '__main__': main(sys.argv)