improve changegroup.readbundle(), use it in hgweb

mercurial/changegroup.py

@@ -1,129 +1,141 b''
 """
 changegroup.py - Mercurial changegroup manipulation functions
 
  Copyright 2006 Matt Mackall <[e-mail]>
 
 This software may be used and distributed according to the terms
 of the GNU General Public License, incorporated herein by reference.
 """
 
 from i18n import _
 import struct, os, bz2, zlib, util, tempfile
 
 def getchunk(source):
     """get a chunk from a changegroup"""
     d = source.read(4)
     if not d:
         return ""
     l = struct.unpack(">l", d)[0]
     if l <= 4:
         return ""
     d = source.read(l - 4)
     if len(d) < l - 4:
         raise util.Abort(_("premature EOF reading chunk"
                            " (got %d bytes, expected %d)")
                           % (len(d), l - 4))
     return d
 
 def chunkiter(source):
     """iterate through the chunks in source"""
     while 1:
         c = getchunk(source)
         if not c:
             break
         yield c
 
 def chunkheader(length):
     """build a changegroup chunk header"""
     return struct.pack(">l", length + 4)
 
 def closechunk():
     return struct.pack(">l", 0)
 
 class nocompress(object):
     def compress(self, x):
         return x
     def flush(self):
         return ""
 
 bundletypes = {
     "": ("", nocompress),
     "HG10UN": ("HG10UN", nocompress),
     "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()),
     "HG10GZ": ("HG10GZ", lambda: zlib.compressobj()),
 }
 
 # hgweb uses this list to communicate it's preferred type
 bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN']
 
 def writebundle(cg, filename, bundletype):
     """Write a bundle file and return its filename.
 
     Existing files will not be overwritten.
     If no filename is specified, a temporary file is created.
     bz2 compression can be turned off.
     The bundle file will be deleted in case of errors.
     """
 
     fh = None
     cleanup = None
     try:
         if filename:
             fh = open(filename, "wb")
         else:
             fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")
             fh = os.fdopen(fd, "wb")
         cleanup = filename
 
         header, compressor = bundletypes[bundletype]
         fh.write(header)
         z = compressor()
 
         # parse the changegroup data, otherwise we will block
         # in case of sshrepo because we don't know the end of the stream
 
         # an empty chunkiter is the end of the changegroup
         # a changegroup has at least 2 chunkiters (changelog and manifest).
         # after that, an empty chunkiter is the end of the changegroup
         empty = False
         count = 0
         while not empty or count <= 2:
             empty = True
             count += 1
             for chunk in chunkiter(cg):
                 empty = False
                 fh.write(z.compress(chunkheader(len(chunk))))
                 pos = 0
                 while pos < len(chunk):
                     next = pos + 2**20
                     fh.write(z.compress(chunk[pos:next]))
                     pos = next
             fh.write(z.compress(closechunk()))
         fh.write(z.flush())
         cleanup = None
         return filename
     finally:
         if fh is not None:
             fh.close()
         if cleanup is not None:
             os.unlink(cleanup)
 
-def readbundle(fh, fname):
-    header = fh.read(6)
-    if not header.startswith("HG"):
-        raise util.Abort(_("%s: not a Mercurial bundle file") % fname)
-    elif not header.startswith("HG10"):
-        raise util.Abort(_("%s: unknown bundle version") % fname)
-
-    if header == "HG10BZ":
+def unbundle(header, fh):
+    if header == 'HG10UN':
+        return fh
+    elif not header.startswith('HG'):
+        # old client with uncompressed bundle
+        def generator(f):
+            yield header
+            for chunk in f:
+                yield chunk
+    elif header == 'HG10GZ':
+        def generator(f):
+            zd = zlib.decompressobj()
+            for chunk in f:
+                yield zd.decompress(chunk)
+    elif header == 'HG10BZ':
         def generator(f):
             zd = bz2.BZ2Decompressor()
             zd.decompress("BZ")
             for chunk in util.filechunkiter(f, 4096):
                 yield zd.decompress(chunk)
-        return util.chunkbuffer(generator(fh))
-    elif header == "HG10UN":
-        return fh
+    return util.chunkbuffer(generator(fh))
 
-    raise util.Abort(_("%s: unknown bundle compression type")
-                     % fname)
+def readbundle(fh, fname):
+    header = fh.read(6)
+    if not header.startswith('HG'):
+        raise util.Abort(_('%s: not a Mercurial bundle file') % fname)
+    if not header.startswith('HG10'):
+        raise util.Abort(_('%s: unknown bundle version') % fname)
+    elif header not in bundletypes:
+        raise util.Abort(_('%s: unknown bundle compression type') % fname)
+    return unbundle(header, fh)

mercurial/hgweb/protocol.py

@@ -1,243 +1,221 b''
 #
 # Copyright 21 May 2005 - (c) 2005 Jake Edge <[e-mail]>
 # Copyright 2005-2007 Matt Mackall <[e-mail]>
 #
 # This software may be used and distributed according to the terms
 # of the GNU General Public License, incorporated herein by reference.
 
 import cStringIO, zlib, bz2, tempfile, errno, os, sys
 from mercurial import util, streamclone
 from mercurial.i18n import gettext as _
 from mercurial.node import *
+from mercurial import changegroup as changegroupmod
 from common import HTTP_OK, HTTP_NOT_FOUND, HTTP_SERVER_ERROR
 
 # __all__ is populated with the allowed commands. Be sure to add to it if
 # you're adding a new command, or the new command won't work.
 
 __all__ = [
    'lookup', 'heads', 'branches', 'between', 'changegroup',
    'changegroupsubset', 'capabilities', 'unbundle', 'stream_out',
 ]
 
 HGTYPE = 'application/mercurial-0.1'
 
 def lookup(web, req):
     try:
         r = hex(web.repo.lookup(req.form['key'][0]))
         success = 1
     except Exception,inst:
         r = str(inst)
         success = 0
     resp = "%s %s\n" % (success, r)
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     req.write(resp)
 
 def heads(web, req):
     resp = " ".join(map(hex, web.repo.heads())) + "\n"
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     req.write(resp)
 
 def branches(web, req):
     nodes = []
     if 'nodes' in req.form:
         nodes = map(bin, req.form['nodes'][0].split(" "))
     resp = cStringIO.StringIO()
     for b in web.repo.branches(nodes):
         resp.write(" ".join(map(hex, b)) + "\n")
     resp = resp.getvalue()
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     req.write(resp)
 
 def between(web, req):
     if 'pairs' in req.form:
         pairs = [map(bin, p.split("-"))
                  for p in req.form['pairs'][0].split(" ")]
     resp = cStringIO.StringIO()
     for b in web.repo.between(pairs):
         resp.write(" ".join(map(hex, b)) + "\n")
     resp = resp.getvalue()
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     req.write(resp)
 
 def changegroup(web, req):
     req.respond(HTTP_OK, HGTYPE)
     nodes = []
     if not web.allowpull:
         return
 
     if 'roots' in req.form:
         nodes = map(bin, req.form['roots'][0].split(" "))
 
     z = zlib.compressobj()
     f = web.repo.changegroup(nodes, 'serve')
     while 1:
         chunk = f.read(4096)
         if not chunk:
             break
         req.write(z.compress(chunk))
 
     req.write(z.flush())
 
 def changegroupsubset(web, req):
     req.respond(HTTP_OK, HGTYPE)
     bases = []
     heads = []
     if not web.allowpull:
         return
 
     if 'bases' in req.form:
         bases = [bin(x) for x in req.form['bases'][0].split(' ')]
     if 'heads' in req.form:
         heads = [bin(x) for x in req.form['heads'][0].split(' ')]
 
     z = zlib.compressobj()
     f = web.repo.changegroupsubset(bases, heads, 'serve')
     while 1:
         chunk = f.read(4096)
         if not chunk:
             break
         req.write(z.compress(chunk))
 
     req.write(z.flush())
 
 def capabilities(web, req):
     resp = ' '.join(web.capabilities())
     req.respond(HTTP_OK, HGTYPE, length=len(resp))
     req.write(resp)
 
 def unbundle(web, req):
     def bail(response, headers={}):
         length = int(req.env['CONTENT_LENGTH'])
         for s in util.filechunkiter(req, limit=length):
             # drain incoming bundle, else client will not see
             # response when run outside cgi script
             pass
         req.header(headers.items())
         req.respond(HTTP_OK, HGTYPE)
         req.write('0\n')
         req.write(response)
 
     # require ssl by default, auth info cannot be sniffed and
     # replayed
     ssl_req = web.configbool('web', 'push_ssl', True)
     if ssl_req:
         if req.env.get('wsgi.url_scheme') != 'https':
             bail(_('ssl required\n'))
             return
         proto = 'https'
     else:
         proto = 'http'
 
     # do not allow push unless explicitly allowed
     if not web.check_perm(req, 'push', False):
         bail(_('push not authorized\n'),
              headers={'status': '401 Unauthorized'})
         return
 
     their_heads = req.form['heads'][0].split(' ')
 
     def check_heads():
         heads = map(hex, web.repo.heads())
         return their_heads == [hex('force')] or their_heads == heads
 
     # fail early if possible
     if not check_heads():
         bail(_('unsynced changes\n'))
         return
 
     req.respond(HTTP_OK, HGTYPE)
 
     # do not lock repo until all changegroup data is
     # streamed. save to temporary file.
 
     fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
     fp = os.fdopen(fd, 'wb+')
     try:
         length = int(req.env['CONTENT_LENGTH'])
         for s in util.filechunkiter(req, limit=length):
             fp.write(s)
 
         try:
             lock = web.repo.lock()
             try:
                 if not check_heads():
                     req.write('0\n')
                     req.write(_('unsynced changes\n'))
                     return
 
                 fp.seek(0)
                 header = fp.read(6)
-                if not header.startswith("HG"):
-                    # old client with uncompressed bundle
-                    def generator(f):
-                        yield header
-                        for chunk in f:
-                            yield chunk
-                elif not header.startswith("HG10"):
-                    req.write("0\n")
-                    req.write(_("unknown bundle version\n"))
-                    return
-                elif header == "HG10GZ":
-                    def generator(f):
-                        zd = zlib.decompressobj()
-                        for chunk in f:
-                            yield zd.decompress(chunk)
-                elif header == "HG10BZ":
-                    def generator(f):
-                        zd = bz2.BZ2Decompressor()
-                        zd.decompress("BZ")
-                        for chunk in f:
-                            yield zd.decompress(chunk)
-                elif header == "HG10UN":
-                    def generator(f):
-                        for chunk in f:
-                            yield chunk
-                else:
-                    req.write("0\n")
-                    req.write(_("unknown bundle compression type\n"))
-                    return
-                gen = generator(util.filechunkiter(fp, 4096))
+                if header.startswith('HG') and not header.startswith('HG10'):
+                    raise ValueError('unknown bundle version')
+                elif header not in changegroupmod.bundletypes:
+                    raise ValueError('unknown bundle compression type')
+                gen = changegroupmod.unbundle(header, fp)
 
                 # send addchangegroup output to client
 
                 old_stdout = sys.stdout
                 sys.stdout = cStringIO.StringIO()
 
                 try:
                     url = 'remote:%s:%s' % (proto,
                                             req.env.get('REMOTE_HOST', ''))
                     try:
-                        ret = web.repo.addchangegroup(
-                                    util.chunkbuffer(gen), 'serve', url)
+                        ret = web.repo.addchangegroup(gen, 'serve', url)
                     except util.Abort, inst:
                         sys.stdout.write("abort: %s\n" % inst)
                         ret = 0
                 finally:
                     val = sys.stdout.getvalue()
                     sys.stdout = old_stdout
                 req.write('%d\n' % ret)
                 req.write(val)
             finally:
                 del lock
+        except ValueError, inst:
+            req.write('0\n')
+            req.write(str(inst) + '\n')
         except (OSError, IOError), inst:
             req.write('0\n')
             filename = getattr(inst, 'filename', '')
             # Don't send our filesystem layout to the client
             if filename.startswith(web.repo.root):
                 filename = filename[len(web.repo.root)+1:]
             else:
                 filename = ''
             error = getattr(inst, 'strerror', 'Unknown error')
             if inst.errno == errno.ENOENT:
                 code = HTTP_NOT_FOUND
             else:
                 code = HTTP_SERVER_ERROR
             req.respond(code)
             req.write('%s: %s\n' % (error, filename))
     finally:
         fp.close()
         os.unlink(tempname)
 
 def stream_out(web, req):
     req.respond(HTTP_OK, HGTYPE)
     streamclone.stream_out(web.repo, req, untrusted=True)