client.py

   1 import re, struct, errno
   2 import git
   3 from helpers import *
   4 from subprocess import Popen, PIPE
   5
   6 class ClientError(Exception):
   7     pass
   8
   9 class Client:
  10     def __init__(self, remote, create=False):
  11         self._busy = None
  12         self._indexes_synced = 0
  13         self.p = None
  14         self.conn = None
  15         rs = remote.split(':', 1)
  16         nicedir = os.path.split(os.path.abspath(sys.argv[0]))[0]
  17         nicedir = re.sub(r':', "_", nicedir)
  18         if len(rs) == 1:
  19             (host, dir) = ('NONE', remote)
  20             def fixenv():
  21                 os.environ['PATH'] = ':'.join([nicedir,
  22                                                os.environ.get('PATH', '')])
  23             argv = ['bup', 'server']
  24         else:
  25             (host, dir) = rs
  26             fixenv = None
  27             # WARNING: shell quoting security holes are possible here, so we
  28             # have to be super careful.  We have to use 'sh -c' because
  29             # csh-derived shells can't handle PATH= notation.  We can't
  30             # set PATH in advance, because ssh probably replaces it.  We
  31             # can't exec *safely* using argv, because *both* ssh and 'sh -c'
  32             # allow shellquoting.  So we end up having to double-shellquote
  33             # stuff here.
  34             escapedir = re.sub(r'([^\w/])', r'\\\\\\\1', nicedir)
  35             cmd = r"""
  36                        sh -c PATH=%s:'$PATH bup server'
  37                    """ % escapedir
  38             argv = ['ssh', host, '--', cmd.strip()]
  39             #log('argv is: %r\n' % argv)
  40         (self.host, self.dir) = (host, dir)
  41         self.cachedir = git.repo('index-cache/%s'
  42                                  % re.sub(r'[^@:\w]', '_',
  43                                           "%s:%s" % (host, dir)))
  44         try:
  45             self.p = p = Popen(argv, stdin=PIPE, stdout=PIPE, preexec_fn=fixenv)
  46         except OSError, e:
  47             raise ClientError, 'exec %r: %s' % (argv[0], e), sys.exc_info()[2]
  48         self.conn = conn = Conn(p.stdout, p.stdin)
  49         if dir:
  50             dir = re.sub(r'[\r\n]', ' ', dir)
  51             if create:
  52                 conn.write('init-dir %s\n' % dir)
  53             else:
  54                 conn.write('set-dir %s\n' % dir)
  55             self.check_ok()
  56
  57     def __del__(self):
  58         try:
  59             self.close()
  60         except IOError, e:
  61             if e.errno == errno.EPIPE:
  62                 pass
  63             else:
  64                 raise
  65
  66     def close(self):
  67         if self.conn and not self._busy:
  68             self.conn.write('quit\n')
  69         if self.p:
  70             self.p.stdin.close()
  71             while self.p.stdout.read(65536):
  72                 pass
  73             self.p.stdout.close()
  74             self.p.wait()
  75             rv = self.p.wait()
  76             if rv:
  77                 raise ClientError('server tunnel returned exit code %d' % rv)
  78         self.conn = None
  79         self.p = None
  80
  81     def check_ok(self):
  82         rv = self.p.poll()
  83         if rv != None:
  84             raise ClientError('server exited unexpectedly with code %r' % rv)
  85         try:
  86             self.conn.check_ok()
  87         except Exception, e:
  88             raise ClientError, e, sys.exc_info()[2]
  89
  90     def check_busy(self):
  91         if self._busy:
  92             raise ClientError('already busy with command %r' % self._busy)
  93
  94     def _not_busy(self):
  95         self._busy = None
  96
  97     def sync_indexes(self):
  98         self.check_busy()
  99         conn = self.conn
 100         conn.write('list-indexes\n')
 101         packdir = git.repo('objects/pack')
 102         mkdirp(self.cachedir)
 103         all = {}
 104         needed = {}
 105         for line in linereader(conn):
 106             if not line:
 107                 break
 108             all[line] = 1
 109             assert(line.find('/') < 0)
 110             if not os.path.exists(os.path.join(self.cachedir, line)):
 111                 needed[line] = 1
 112         self.check_ok()
 113
 114         for f in os.listdir(self.cachedir):
 115             if f.endswith('.idx') and not f in all:
 116                 log('pruning old index: %r\n' % f)
 117                 os.unlink(os.path.join(self.cachedir, f))
 118
 119         # FIXME this should be pipelined: request multiple indexes at a time, or
 120         # we waste lots of network turnarounds.
 121         for name in needed.keys():
 122             log('requesting %r\n' % name)
 123             conn.write('send-index %s\n' % name)
 124             n = struct.unpack('!I', conn.read(4))[0]
 125             assert(n)
 126             log('   expect %d bytes\n' % n)
 127             fn = os.path.join(self.cachedir, name)
 128             f = open(fn + '.tmp', 'w')
 129             for b in chunkyreader(conn, n):
 130                 f.write(b)
 131             self.check_ok()
 132             f.close()
 133             os.rename(fn + '.tmp', fn)
 134
 135         self._indexes_synced = 1
 136
 137     def _make_objcache(self):
 138         ob = self._busy
 139         self._busy = None
 140         self.sync_indexes()
 141         self._busy = ob
 142         return git.MultiPackIndex(self.cachedir)
 143
 144     def new_packwriter(self):
 145         self.check_busy()
 146         self._busy = 'receive-objects'
 147         return PackWriter_Remote(self.conn,
 148                                  objcache_maker = self._make_objcache,
 149                                  onclose = self._not_busy)
 150
 151     def read_ref(self, refname):
 152         self.check_busy()
 153         self.conn.write('read-ref %s\n' % refname)
 154         r = self.conn.readline().strip()
 155         self.check_ok()
 156         if r:
 157             assert(len(r) == 40)   # hexified sha
 158             return r.decode('hex')
 159         else:
 160             return None   # nonexistent ref
 161
 162     def update_ref(self, refname, newval, oldval):
 163         self.check_busy()
 164         self.conn.write('update-ref %s\n%s\n%s\n'
 165                         % (refname, newval.encode('hex'),
 166                            (oldval or '').encode('hex')))
 167         self.check_ok()
 168
 169     def cat(self, id):
 170         self.check_busy()
 171         self._busy = 'cat'
 172         self.conn.write('cat %s\n' % re.sub(r'[\n\r]', '_', id))
 173         while 1:
 174             sz = struct.unpack('!I', self.conn.read(4))[0]
 175             if not sz: break
 176             yield self.conn.read(sz)
 177         self.check_ok()
 178         self._not_busy()
 179
 180
 181 class PackWriter_Remote(git.PackWriter):
 182     def __init__(self, conn, objcache_maker=None, onclose=None):
 183         git.PackWriter.__init__(self, objcache_maker)
 184         self.file = conn
 185         self.filename = 'remote socket'
 186         self.onclose = onclose
 187         self._packopen = False
 188
 189     def _open(self):
 190         if not self._packopen:
 191             self._make_objcache()
 192             self.file.write('receive-objects\n')
 193             self._packopen = True
 194
 195     def _end(self):
 196         if self._packopen and self.file:
 197             self.file.write('\0\0\0\0')
 198             self._packopen = False
 199             id = self.file.readline().strip()
 200             self.file.check_ok()
 201             self.objcache = None
 202             if self.onclose:
 203                 self.onclose()
 204             return id
 205
 206     def close(self):
 207         id = self._end()
 208         self.file = None
 209         return id
 210
 211     def abort(self):
 212         raise GitError("don't know how to abort remote pack writing")
 213
 214     def _raw_write(self, datalist):
 215         assert(self.file)
 216         if not self._packopen:
 217             self._open()
 218         data = ''.join(datalist)
 219         assert(len(data))
 220         self.file.write(struct.pack('!I', len(data)) + data)
 221         self.outbytes += len(data)
 222         self.count += 1
 223
 224