0d83bb814e41b11c61895724648792dfafc82c77
[jelmer/dulwich-libgit2.git] / dulwich / protocol.py
1 # protocol.py -- Shared parts of the git protocols
2 # Copryight (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
31 TCP_GIT_PORT = 9418
32
33 ZERO_SHA = "0" * 40
34
35 SINGLE_ACK = 0
36 MULTI_ACK = 1
37 MULTI_ACK_DETAILED = 2
38
39 class ProtocolFile(object):
40     """
41     Some network ops are like file ops. The file ops expect to operate on
42     file objects, so provide them with a dummy file.
43     """
44
45     def __init__(self, read, write):
46         self.read = read
47         self.write = write
48
49     def tell(self):
50         pass
51
52     def close(self):
53         pass
54
55
56 class Protocol(object):
57
58     def __init__(self, read, write, report_activity=None):
59         self.read = read
60         self.write = write
61         self.report_activity = report_activity
62
63     def read_pkt_line(self):
64         """
65         Reads a 'pkt line' from the remote git process
66
67         :return: The next string from the stream
68         """
69         try:
70             sizestr = self.read(4)
71             if not sizestr:
72                 raise HangupException()
73             size = int(sizestr, 16)
74             if size == 0:
75                 if self.report_activity:
76                     self.report_activity(4, 'read')
77                 return None
78             if self.report_activity:
79                 self.report_activity(size, 'read')
80             return self.read(size-4)
81         except socket.error, e:
82             raise GitProtocolError(e)
83
84     def read_pkt_seq(self):
85         pkt = self.read_pkt_line()
86         while pkt:
87             yield pkt
88             pkt = self.read_pkt_line()
89
90     def write_pkt_line(self, line):
91         """
92         Sends a 'pkt line' to the remote git process
93
94         :param line: A string containing the data to send
95         """
96         try:
97             if line is None:
98                 self.write("0000")
99                 if self.report_activity:
100                     self.report_activity(4, 'write')
101             else:
102                 self.write("%04x%s" % (len(line)+4, line))
103                 if self.report_activity:
104                     self.report_activity(4+len(line), 'write')
105         except socket.error, e:
106             raise GitProtocolError(e)
107
108     def write_file(self):
109         class ProtocolFile(object):
110
111             def __init__(self, proto):
112                 self._proto = proto
113                 self._offset = 0
114
115             def write(self, data):
116                 self._proto.write(data)
117                 self._offset += len(data)
118
119             def tell(self):
120                 return self._offset
121
122             def close(self):
123                 pass
124
125         return ProtocolFile(self)
126
127     def write_sideband(self, channel, blob):
128         """
129         Write data to the sideband (a git multiplexing method)
130
131         :param channel: int specifying which channel to write to
132         :param blob: a blob of data (as a string) to send on this channel
133         """
134         # a pktline can be a max of 65520. a sideband line can therefore be
135         # 65520-5 = 65515
136         # WTF: Why have the len in ASCII, but the channel in binary.
137         while blob:
138             self.write_pkt_line("%s%s" % (chr(channel), blob[:65515]))
139             blob = blob[65515:]
140
141     def send_cmd(self, cmd, *args):
142         """
143         Send a command and some arguments to a git server
144
145         Only used for git://
146
147         :param cmd: The remote service to access
148         :param args: List of arguments to send to remove service
149         """
150         self.write_pkt_line("%s %s" % (cmd, "".join(["%s\0" % a for a in args])))
151
152     def read_cmd(self):
153         """
154         Read a command and some arguments from the git client
155
156         Only used for git://
157
158         :return: A tuple of (command, [list of arguments])
159         """
160         line = self.read_pkt_line()
161         splice_at = line.find(" ")
162         cmd, args = line[:splice_at], line[splice_at+1:]
163         assert args[-1] == "\x00"
164         return cmd, args[:-1].split(chr(0))
165
166
167 _RBUFSIZE = 8192  # Default read buffer size.
168
169
170 class ReceivableProtocol(Protocol):
171     """Variant of Protocol that allows reading up to a size without blocking.
172
173     This class has a recv() method that behaves like socket.recv() in addition
174     to a read() method.
175
176     If you want to read n bytes from the wire and block until exactly n bytes
177     (or EOF) are read, use read(n). If you want to read at most n bytes from the
178     wire but don't care if you get less, use recv(n). Note that recv(n) will
179     still block until at least one byte is read.
180     """
181
182     def __init__(self, recv, write, report_activity=None, rbufsize=_RBUFSIZE):
183         super(ReceivableProtocol, self).__init__(self.read, write,
184                                                 report_activity)
185         self._recv = recv
186         self._rbuf = StringIO()
187         self._rbufsize = rbufsize
188
189     def read(self, size):
190         # From _fileobj.read in socket.py in the Python 2.6.5 standard library,
191         # with the following modifications:
192         #  - omit the size <= 0 branch
193         #  - seek back to start rather than 0 in case some buffer has been
194         #    consumed.
195         #  - use os.SEEK_END instead of the magic number.
196         # Copyright (c) 2001-2010 Python Software Foundation; All Rights Reserved
197         # Licensed under the Python Software Foundation License.
198         # TODO: see if buffer is more efficient than cStringIO.
199         assert size > 0
200
201         # Our use of StringIO rather than lists of string objects returned by
202         # recv() minimizes memory usage and fragmentation that occurs when
203         # rbufsize is large compared to the typical return value of recv().
204         buf = self._rbuf
205         start = buf.tell()
206         buf.seek(0, os.SEEK_END)
207         # buffer may have been partially consumed by recv()
208         buf_len = buf.tell() - start
209         if buf_len >= size:
210             # Already have size bytes in our buffer?  Extract and return.
211             buf.seek(start)
212             rv = buf.read(size)
213             self._rbuf = StringIO()
214             self._rbuf.write(buf.read())
215             self._rbuf.seek(0)
216             return rv
217
218         self._rbuf = StringIO()  # reset _rbuf.  we consume it via buf.
219         while True:
220             left = size - buf_len
221             # recv() will malloc the amount of memory given as its
222             # parameter even though it often returns much less data
223             # than that.  The returned data string is short lived
224             # as we copy it into a StringIO and free it.  This avoids
225             # fragmentation issues on many platforms.
226             data = self._recv(left)
227             if not data:
228                 break
229             n = len(data)
230             if n == size and not buf_len:
231                 # Shortcut.  Avoid buffer data copies when:
232                 # - We have no data in our buffer.
233                 # AND
234                 # - Our call to recv returned exactly the
235                 #   number of bytes we were asked to read.
236                 return data
237             if n == left:
238                 buf.write(data)
239                 del data  # explicit free
240                 break
241             assert n <= left, "_recv(%d) returned %d bytes" % (left, n)
242             buf.write(data)
243             buf_len += n
244             del data  # explicit free
245             #assert buf_len == buf.tell()
246         buf.seek(start)
247         return buf.read()
248
249     def recv(self, size):
250         assert size > 0
251
252         buf = self._rbuf
253         start = buf.tell()
254         buf.seek(0, os.SEEK_END)
255         buf_len = buf.tell()
256         buf.seek(start)
257
258         left = buf_len - start
259         if not left:
260             # only read from the wire if our read buffer is exhausted
261             data = self._recv(self._rbufsize)
262             if len(data) == size:
263                 # shortcut: skip the buffer if we read exactly size bytes
264                 return data
265             buf = StringIO()
266             buf.write(data)
267             buf.seek(0)
268             del data  # explicit free
269             self._rbuf = buf
270         return buf.read(size)
271
272
273 def extract_capabilities(text):
274     """Extract a capabilities list from a string, if present.
275
276     :param text: String to extract from
277     :return: Tuple with text with capabilities removed and list of capabilities
278     """
279     if not "\0" in text:
280         return text, []
281     text, capabilities = text.rstrip().split("\0")
282     return (text, capabilities.strip().split(" "))
283
284
285 def extract_want_line_capabilities(text):
286     """Extract a capabilities list from a want line, if present.
287
288     Note that want lines have capabilities separated from the rest of the line
289     by a space instead of a null byte. Thus want lines have the form:
290
291         want obj-id cap1 cap2 ...
292
293     :param text: Want line to extract from
294     :return: Tuple with text with capabilities removed and list of capabilities
295     """
296     split_text = text.rstrip().split(" ")
297     if len(split_text) < 3:
298         return text, []
299     return (" ".join(split_text[:2]), split_text[2:])
300
301
302 def ack_type(capabilities):
303     """Extract the ack type from a capabilities list."""
304     if 'multi_ack_detailed' in capabilities:
305       return MULTI_ACK_DETAILED
306     elif 'multi_ack' in capabilities:
307         return MULTI_ACK
308     return SINGLE_ACK