Clean up docstrings in dulwich.protocol.
[jelmer/dulwich-libgit2.git] / dulwich / protocol.py
1 # protocol.py -- Shared parts of the git protocols
2 # Copyright (C) 2008 John Carr <john.carr@unrouted.co.uk>
3 # Copyright (C) 2008 Jelmer Vernooij <jelmer@samba.org>
4 #
5 # This program is free software; you can redistribute it and/or
6 # modify it under the terms of the GNU General Public License
7 # as published by the Free Software Foundation; version 2
8 # or (at your option) any later version of the License.
9 #
10 # This program is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 # GNU General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program; if not, write to the Free Software
17 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
18 # MA  02110-1301, USA.
19
20 """Generic functions for talking the git smart server protocol."""
21
22 from cStringIO import StringIO
23 import os
24 import socket
25
26 from dulwich.errors import (
27     HangupException,
28     GitProtocolError,
29     )
30 from dulwich.misc import (
31     SEEK_END,
32     )
33
34 TCP_GIT_PORT = 9418
35
36 ZERO_SHA = "0" * 40
37
38 SINGLE_ACK = 0
39 MULTI_ACK = 1
40 MULTI_ACK_DETAILED = 2
41
42
43 class ProtocolFile(object):
44     """A dummy file for network ops that expect file-like objects."""
45
46     def __init__(self, read, write):
47         self.read = read
48         self.write = write
49
50     def tell(self):
51         pass
52
53     def close(self):
54         pass
55
56
57 def pkt_line(data):
58     """Wrap data in a pkt-line.
59
60     :param data: The data to wrap, as a str or None.
61     :return: The data prefixed with its length in pkt-line format; if data was
62         None, returns the flush-pkt ('0000').
63     """
64     if data is None:
65         return '0000'
66     return '%04x%s' % (len(data) + 4, data)
67
68
69 class Protocol(object):
70     """Class for interacting with a remote git process over the wire.
71
72     Parts of the git wire protocol use 'pkt-lines' to communicate. A pkt-line
73     consists of the length of the line as a 4-byte hex string, followed by the
74     payload data. The length includes the 4-byte header. The special line '0000'
75     indicates the end of a section of input and is called a 'flush-pkt'.
76
77     For details on the pkt-line format, see the cgit distribution:
78         Documentation/technical/protocol-common.txt
79     """
80
81     def __init__(self, read, write, report_activity=None):
82         self.read = read
83         self.write = write
84         self.report_activity = report_activity
85
86     def read_pkt_line(self):
87         """Reads a pkt-line from the remote git process.
88
89         :return: The next string from the stream, without the length prefix, or
90             None for a flush-pkt ('0000').
91         """
92         try:
93             sizestr = self.read(4)
94             if not sizestr:
95                 raise HangupException()
96             size = int(sizestr, 16)
97             if size == 0:
98                 if self.report_activity:
99                     self.report_activity(4, 'read')
100                 return None
101             if self.report_activity:
102                 self.report_activity(size, 'read')
103             return self.read(size-4)
104         except socket.error, e:
105             raise GitProtocolError(e)
106
107     def read_pkt_seq(self):
108         """Read a sequence of pkt-lines from the remote git process.
109
110         :yield: Each line of data up to but not including the next flush-pkt.
111         """
112         pkt = self.read_pkt_line()
113         while pkt:
114             yield pkt
115             pkt = self.read_pkt_line()
116
117     def write_pkt_line(self, line):
118         """Sends a pkt-line to the remote git process.
119
120         :param line: A string containing the data to send, without the length
121             prefix.
122         """
123         try:
124             line = pkt_line(line)
125             self.write(line)
126             if self.report_activity:
127                 self.report_activity(len(line), 'write')
128         except socket.error, e:
129             raise GitProtocolError(e)
130
131     def write_file(self):
132         """Return a writable file-like object for this protocol."""
133
134         class ProtocolFile(object):
135
136             def __init__(self, proto):
137                 self._proto = proto
138                 self._offset = 0
139
140             def write(self, data):
141                 self._proto.write(data)
142                 self._offset += len(data)
143
144             def tell(self):
145                 return self._offset
146
147             def close(self):
148                 pass
149
150         return ProtocolFile(self)
151
152     def write_sideband(self, channel, blob):
153         """Write multiplexed data to the sideband.
154
155         :param channel: An int specifying the channel to write to.
156         :param blob: A blob of data (as a string) to send on this channel.
157         """
158         # a pktline can be a max of 65520. a sideband line can therefore be
159         # 65520-5 = 65515
160         # WTF: Why have the len in ASCII, but the channel in binary.
161         while blob:
162             self.write_pkt_line("%s%s" % (chr(channel), blob[:65515]))
163             blob = blob[65515:]
164
165     def send_cmd(self, cmd, *args):
166         """Send a command and some arguments to a git server.
167
168         Only used for the TCP git protocol (git://).
169
170         :param cmd: The remote service to access.
171         :param args: List of arguments to send to remove service.
172         """
173         self.write_pkt_line("%s %s" % (cmd, "".join(["%s\0" % a for a in args])))
174
175     def read_cmd(self):
176         """Read a command and some arguments from the git client
177
178         Only used for the TCP git protocol (git://).
179
180         :return: A tuple of (command, [list of arguments]).
181         """
182         line = self.read_pkt_line()
183         splice_at = line.find(" ")
184         cmd, args = line[:splice_at], line[splice_at+1:]
185         assert args[-1] == "\x00"
186         return cmd, args[:-1].split(chr(0))
187
188
189 _RBUFSIZE = 8192  # Default read buffer size.
190
191
192 class ReceivableProtocol(Protocol):
193     """Variant of Protocol that allows reading up to a size without blocking.
194
195     This class has a recv() method that behaves like socket.recv() in addition
196     to a read() method.
197
198     If you want to read n bytes from the wire and block until exactly n bytes
199     (or EOF) are read, use read(n). If you want to read at most n bytes from the
200     wire but don't care if you get less, use recv(n). Note that recv(n) will
201     still block until at least one byte is read.
202     """
203
204     def __init__(self, recv, write, report_activity=None, rbufsize=_RBUFSIZE):
205         super(ReceivableProtocol, self).__init__(self.read, write,
206                                                  report_activity)
207         self._recv = recv
208         self._rbuf = StringIO()
209         self._rbufsize = rbufsize
210
211     def read(self, size):
212         # From _fileobj.read in socket.py in the Python 2.6.5 standard library,
213         # with the following modifications:
214         #  - omit the size <= 0 branch
215         #  - seek back to start rather than 0 in case some buffer has been
216         #    consumed.
217         #  - use SEEK_END instead of the magic number.
218         # Copyright (c) 2001-2010 Python Software Foundation; All Rights Reserved
219         # Licensed under the Python Software Foundation License.
220         # TODO: see if buffer is more efficient than cStringIO.
221         assert size > 0
222
223         # Our use of StringIO rather than lists of string objects returned by
224         # recv() minimizes memory usage and fragmentation that occurs when
225         # rbufsize is large compared to the typical return value of recv().
226         buf = self._rbuf
227         start = buf.tell()
228         buf.seek(0, SEEK_END)
229         # buffer may have been partially consumed by recv()
230         buf_len = buf.tell() - start
231         if buf_len >= size:
232             # Already have size bytes in our buffer?  Extract and return.
233             buf.seek(start)
234             rv = buf.read(size)
235             self._rbuf = StringIO()
236             self._rbuf.write(buf.read())
237             self._rbuf.seek(0)
238             return rv
239
240         self._rbuf = StringIO()  # reset _rbuf.  we consume it via buf.
241         while True:
242             left = size - buf_len
243             # recv() will malloc the amount of memory given as its
244             # parameter even though it often returns much less data
245             # than that.  The returned data string is short lived
246             # as we copy it into a StringIO and free it.  This avoids
247             # fragmentation issues on many platforms.
248             data = self._recv(left)
249             if not data:
250                 break
251             n = len(data)
252             if n == size and not buf_len:
253                 # Shortcut.  Avoid buffer data copies when:
254                 # - We have no data in our buffer.
255                 # AND
256                 # - Our call to recv returned exactly the
257                 #   number of bytes we were asked to read.
258                 return data
259             if n == left:
260                 buf.write(data)
261                 del data  # explicit free
262                 break
263             assert n <= left, "_recv(%d) returned %d bytes" % (left, n)
264             buf.write(data)
265             buf_len += n
266             del data  # explicit free
267             #assert buf_len == buf.tell()
268         buf.seek(start)
269         return buf.read()
270
271     def recv(self, size):
272         assert size > 0
273
274         buf = self._rbuf
275         start = buf.tell()
276         buf.seek(0, SEEK_END)
277         buf_len = buf.tell()
278         buf.seek(start)
279
280         left = buf_len - start
281         if not left:
282             # only read from the wire if our read buffer is exhausted
283             data = self._recv(self._rbufsize)
284             if len(data) == size:
285                 # shortcut: skip the buffer if we read exactly size bytes
286                 return data
287             buf = StringIO()
288             buf.write(data)
289             buf.seek(0)
290             del data  # explicit free
291             self._rbuf = buf
292         return buf.read(size)
293
294
295 def extract_capabilities(text):
296     """Extract a capabilities list from a string, if present.
297
298     :param text: String to extract from
299     :return: Tuple with text with capabilities removed and list of capabilities
300     """
301     if not "\0" in text:
302         return text, []
303     text, capabilities = text.rstrip().split("\0")
304     return (text, capabilities.strip().split(" "))
305
306
307 def extract_want_line_capabilities(text):
308     """Extract a capabilities list from a want line, if present.
309
310     Note that want lines have capabilities separated from the rest of the line
311     by a space instead of a null byte. Thus want lines have the form:
312
313         want obj-id cap1 cap2 ...
314
315     :param text: Want line to extract from
316     :return: Tuple with text with capabilities removed and list of capabilities
317     """
318     split_text = text.rstrip().split(" ")
319     if len(split_text) < 3:
320         return text, []
321     return (" ".join(split_text[:2]), split_text[2:])
322
323
324 def ack_type(capabilities):
325     """Extract the ack type from a capabilities list."""
326     if 'multi_ack_detailed' in capabilities:
327         return MULTI_ACK_DETAILED
328     elif 'multi_ack' in capabilities:
329         return MULTI_ACK
330     return SINGLE_ACK
331
332
333 class BufferedPktLineWriter(object):
334     """Writer that wraps its data in pkt-lines and has an independent buffer.
335
336     Consecutive calls to write() wrap the data in a pkt-line and then buffers it
337     until enough lines have been written such that their total length (including
338     length prefix) reach the buffer size.
339     """
340
341     def __init__(self, write, bufsize=65515):
342         """Initialize the BufferedPktLineWriter.
343
344         :param write: A write callback for the underlying writer.
345         :param bufsize: The internal buffer size, including length prefixes.
346         """
347         self._write = write
348         self._bufsize = bufsize
349         self._wbuf = StringIO()
350         self._buflen = 0
351
352     def write(self, data):
353         """Write data, wrapping it in a pkt-line."""
354         line = pkt_line(data)
355         line_len = len(line)
356         over = self._buflen + line_len - self._bufsize
357         if over >= 0:
358             start = line_len - over
359             self._wbuf.write(line[:start])
360             self.flush()
361         else:
362             start = 0
363         saved = line[start:]
364         self._wbuf.write(saved)
365         self._buflen += len(saved)
366
367     def flush(self):
368         """Flush all data from the buffer."""
369         data = self._wbuf.getvalue()
370         if data:
371             self._write(data)
372         self._len = 0
373         self._wbuf = StringIO()