1#!/usr/bin/python
2# from http://lists.canonical.org/pipermail/kragen-hacks/2004-May/000397.html
3import 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
35class Error(Exception): pass
36
37class 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
51class 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
68def 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
73def 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
79read_request = 0
80write_request = 1
81disconnect_request = 2
82
83class 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
100def 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
120def 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
127def 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
136if __name__ == '__main__': main(sys.argv)