sshrepo: be more careful while reading data

This should avoid some tracebacks when the server sends some garbage.

mercurial/sshrepo.py

@@ -1,228 +1,239 b''
 # sshrepo.py - ssh repository proxy class for mercurial
 #
 # Copyright 2005, 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 node import *
 from remoterepo import *
 from i18n import _
 import repo, os, re, stat, util
 
 class sshrepository(remoterepository):
     def __init__(self, ui, path, create=0):
         self._url = path
         self.ui = ui
 
         m = re.match(r'^ssh://(([^@]+)@)?([^:/]+)(:(\d+))?(/(.*))?$', path)
         if not m:
             self.raise_(repo.RepoError(_("couldn't parse location %s") % path))
 
         self.user = m.group(2)
         self.host = m.group(3)
         self.port = m.group(5)
         self.path = m.group(7) or "."
 
         args = self.user and ("%s@%s" % (self.user, self.host)) or self.host
         args = self.port and ("%s -p %s") % (args, self.port) or args
 
         sshcmd = self.ui.config("ui", "ssh", "ssh")
         remotecmd = self.ui.config("ui", "remotecmd", "hg")
 
         if create:
             cmd = '%s %s "%s init %s"'
             cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
             ui.note('running %s\n' % cmd)
             res = util.system(cmd)
             if res != 0:
                 self.raise_(repo.RepoError(_("could not create remote repo")))
 
         self.validate_repo(ui, sshcmd, args, remotecmd)
 
     def url(self):
         return self._url
 
     def validate_repo(self, ui, sshcmd, args, remotecmd):
         # cleanup up previous run
         self.cleanup()
 
         cmd = '%s %s "%s -R %s serve --stdio"'
         cmd = cmd % (sshcmd, args, remotecmd, self.path)
 
         cmd = util.quotecommand(cmd)
         ui.note('running %s\n' % cmd)
         self.pipeo, self.pipei, self.pipee = os.popen3(cmd, 'b')
 
         # skip any noise generated by remote shell
         self.do_cmd("hello")
         r = self.do_cmd("between", pairs=("%s-%s" % ("0"*40, "0"*40)))
         lines = ["", "dummy"]
         max_noise = 500
         while lines[-1] and max_noise:
             l = r.readline()
             self.readerr()
             if lines[-1] == "1\n" and l == "\n":
                 break
             if l:
                 ui.debug(_("remote: "), l)
             lines.append(l)
             max_noise -= 1
         else:
             self.raise_(repo.RepoError(_("no suitable response from remote hg")))
 
         self.capabilities = util.set()
         lines.reverse()
         for l in lines:
             if l.startswith("capabilities:"):
                 self.capabilities.update(l[:-1].split(":")[1].split())
                 break
 
     def readerr(self):
         while 1:
             size = util.fstat(self.pipee).st_size
             if size == 0: break
             l = self.pipee.readline()
             if not l: break
             self.ui.status(_("remote: "), l)
 
     def raise_(self, exception):
         self.cleanup()
         raise exception
 
     def cleanup(self):
         try:
             self.pipeo.close()
             self.pipei.close()
             # read the error descriptor until EOF
             for l in self.pipee:
                 self.ui.status(_("remote: "), l)
             self.pipee.close()
         except:
             pass
 
     __del__ = cleanup
 
     def do_cmd(self, cmd, **args):
         self.ui.debug(_("sending %s command\n") % cmd)
         self.pipeo.write("%s\n" % cmd)
         for k, v in args.items():
             self.pipeo.write("%s %d\n" % (k, len(v)))
             self.pipeo.write(v)
         self.pipeo.flush()
 
         return self.pipei
 
     def call(self, cmd, **args):
-        r = self.do_cmd(cmd, **args)
-        l = r.readline()
+        self.do_cmd(cmd, **args)
+        return self._recv()
+
+    def _recv(self):
+        l = self.pipei.readline()
         self.readerr()
         try:
             l = int(l)
         except:
             self.raise_(util.UnexpectedOutput(_("unexpected response:"), l))
-        return r.read(l)
+        return self.pipei.read(l)
+
+    def _send(self, data, flush=False):
+        self.pipeo.write("%d\n" % len(data))
+        if data:
+            self.pipeo.write(data)
+        if flush:
+            self.pipeo.flush()
+        self.readerr()
 
     def lock(self):
         self.call("lock")
         return remotelock(self)
 
     def unlock(self):
         self.call("unlock")
 
     def lookup(self, key):
         self.requirecap('lookup', _('look up remote revision'))
         d = self.call("lookup", key=key)
         success, data = d[:-1].split(" ", 1)
         if int(success):
             return bin(data)
         else:
             self.raise_(repo.RepoError(data))
 
     def heads(self):
         d = self.call("heads")
         try:
             return map(bin, d[:-1].split(" "))
         except:
             self.raise_(util.UnexpectedOutput(_("unexpected response:"), d))
 
     def branches(self, nodes):
         n = " ".join(map(hex, nodes))
         d = self.call("branches", nodes=n)
         try:
             br = [ tuple(map(bin, b.split(" "))) for b in d.splitlines() ]
             return br
         except:
             self.raise_(util.UnexpectedOutput(_("unexpected response:"), d))
 
     def between(self, pairs):
         n = "\n".join(["-".join(map(hex, p)) for p in pairs])
         d = self.call("between", pairs=n)
         try:
             p = [ l and map(bin, l.split(" ")) or [] for l in d.splitlines() ]
             return p
         except:
             self.raise_(util.UnexpectedOutput(_("unexpected response:"), d))
 
     def changegroup(self, nodes, kind):
         n = " ".join(map(hex, nodes))
         return self.do_cmd("changegroup", roots=n)
 
     def changegroupsubset(self, bases, heads, kind):
         self.requirecap('changegroupsubset', _('look up remote changes'))
         bases = " ".join(map(hex, bases))
         heads = " ".join(map(hex, heads))
         return self.do_cmd("changegroupsubset", bases=bases, heads=heads)
 
     def unbundle(self, cg, heads, source):
         d = self.call("unbundle", heads=' '.join(map(hex, heads)))
         if d:
             # remote may send "unsynced changes"
             self.raise_(repo.RepoError(_("push refused: %s") % d))
 
         while 1:
             d = cg.read(4096)
-            if not d: break
-            self.pipeo.write(str(len(d)) + '\n')
-            self.pipeo.write(d)
-            self.readerr()
+            if not d:
+                break
+            self._send(d)
 
-        self.pipeo.write('0\n')
-        self.pipeo.flush()
+        self._send("", flush=True)
 
-        self.readerr()
-        l = int(self.pipei.readline())
-        r = self.pipei.read(l)
+        r = self._recv()
         if r:
             # remote may send "unsynced changes"
             self.raise_(hg.RepoError(_("push failed: %s") % r))
 
-        self.readerr()
-        l = int(self.pipei.readline())
-        r = self.pipei.read(l)
-        return int(r)
+        r = self._recv()
+        try:
+            return int(r)
+        except:
+            self.raise_(util.UnexpectedOutput(_("unexpected response:"), r))
 
     def addchangegroup(self, cg, source, url):
         d = self.call("addchangegroup")
         if d:
             self.raise_(repo.RepoError(_("push refused: %s") % d))
         while 1:
             d = cg.read(4096)
-            if not d: break
+            if not d:
+                break
             self.pipeo.write(d)
             self.readerr()
 
         self.pipeo.flush()
 
         self.readerr()
-        l = int(self.pipei.readline())
-        r = self.pipei.read(l)
+        r = self._recv()
         if not r:
             return 1
-        return int(r)
+        try:
+            return int(r)
+        except:
+            self.raise_(util.UnexpectedOutput(_("unexpected response:"), r))
 
     def stream_out(self):
         return self.do_cmd('stream_out')
 
 instance = sshrepository

tests/test-ssh

@@ -1,103 +1,111 b''
 #!/bin/sh
 
 cp "$TESTDIR"/printenv.py .
 
 # This test tries to exercise the ssh functionality with a dummy script
 
 cat <<EOF > dummyssh
 import sys
 import os
 
 os.chdir(os.path.dirname(sys.argv[0]))
 if sys.argv[1] != "user@dummy":
     sys.exit(-1)
 
 if not os.path.exists("dummyssh"):
     sys.exit(-1)
 
 os.environ["SSH_CLIENT"] = "127.0.0.1 1 2"
 
 log = open("dummylog", "ab")
 log.write("Got arguments")
 for i, arg in enumerate(sys.argv[1:]):
     log.write(" %d:%s" % (i+1, arg))
 log.write("\n")
 log.close()
 r = os.system(sys.argv[2])
 sys.exit(bool(r))
 EOF
 
+cat <<EOF > badhook
+import sys
+sys.stdout.write("KABOOM")
+EOF
+
 echo "# creating 'remote'"
 hg init remote
 cd remote
 echo this > foo
 echo this > fooO
 hg ci -A -m "init" -d "1000000 0" foo fooO
 echo '[server]' > .hg/hgrc
 echo 'uncompressed = True' >> .hg/hgrc
 echo '[hooks]' >> .hg/hgrc
 echo 'changegroup = python ../printenv.py changegroup-in-remote 0 ../dummylog' >> .hg/hgrc
 
 cd ..
 
 echo "# repo not found error"
 hg clone -e "python ./dummyssh" ssh://user@dummy/nonexistent local
 
 echo "# clone remote via stream"
 hg clone -e "python ./dummyssh" --uncompressed ssh://user@dummy/remote local-stream 2>&1 | \
   sed -e 's/[0-9][0-9.]*/XXX/g' -e 's/[KM]\(B\/sec\)/X\1/'
 cd local-stream
 hg verify
 cd ..
 
 echo "# clone remote via pull"
 hg clone -e "python ./dummyssh" ssh://user@dummy/remote local
 
 echo "# verify"
 cd local
 hg verify
 
 echo '[hooks]' >> .hg/hgrc
 echo 'changegroup = python ../printenv.py changegroup-in-local 0 ../dummylog' >> .hg/hgrc
 
 echo "# empty default pull"
 hg paths
 hg pull -e "python ../dummyssh"
 
 echo "# local change"
 echo bleah > foo
 hg ci -m "add" -d "1000000 0"
 
 echo "# updating rc"
 echo "default-push = ssh://user@dummy/remote" >> .hg/hgrc
 echo "[ui]" >> .hg/hgrc
 echo "ssh = python ../dummyssh" >> .hg/hgrc
 
 echo "# find outgoing"
 hg out ssh://user@dummy/remote
 
 echo "# find incoming on the remote side"
 hg incoming -R ../remote -e "python ../dummyssh" ssh://user@dummy/local
 
 echo "# push"
 hg push
 
 cd ../remote
 
 echo "# check remote tip"
 hg tip
 hg verify
 hg cat -r tip foo
 
 echo z > z
 hg ci -A -m z -d '1000001 0' z
+# a bad, evil hook that prints to stdout
+echo 'changegroup.stdout = python ../badhook' >> .hg/hgrc
 
 cd ../local
 echo r > r
 hg ci -A -m z -d '1000002 0' r
 
-echo "# push should succeed"
+echo "# push should succeed even though it has an unexpected response"
 hg push
+hg -R ../remote heads
 
 cd ..
 cat dummylog

tests/test-ssh.out

@@ -1,90 +1,105 b''
 # creating 'remote'
 # repo not found error
 remote: abort: There is no Mercurial repository here (.hg not found)!
 abort: no suitable response from remote hg!
 # clone remote via stream
 streaming all changes
 XXX files to transfer, XXX bytes of data
 transferred XXX bytes in XXX seconds (XXX XB/sec)
 XXX files updated, XXX files merged, XXX files removed, XXX files unresolved
 checking changesets
 checking manifests
 crosschecking files in changesets and manifests
 checking files
 2 files, 1 changesets, 2 total revisions
 # clone remote via pull
 requesting all changes
 adding changesets
 adding manifests
 adding file changes
 added 1 changesets with 2 changes to 2 files
 2 files updated, 0 files merged, 0 files removed, 0 files unresolved
 # verify
 checking changesets
 checking manifests
 crosschecking files in changesets and manifests
 checking files
 2 files, 1 changesets, 2 total revisions
 # empty default pull
 default = ssh://user@dummy/remote
 pulling from ssh://user@dummy/remote
 searching for changes
 no changes found
 # local change
 # updating rc
 # find outgoing
 comparing with ssh://user@dummy/remote
 searching for changes
 changeset:   1:572896fe480d
 tag:         tip
 user:        test
 date:        Mon Jan 12 13:46:40 1970 +0000
 summary:     add
 
 # find incoming on the remote side
 comparing with ssh://user@dummy/local
 searching for changes
 changeset:   1:572896fe480d
 tag:         tip
 user:        test
 date:        Mon Jan 12 13:46:40 1970 +0000
 summary:     add
 
 # push
 pushing to ssh://user@dummy/remote
 searching for changes
 remote: adding changesets
 remote: adding manifests
 remote: adding file changes
 remote: added 1 changesets with 1 changes to 1 files
 # check remote tip
 changeset:   1:572896fe480d
 tag:         tip
 user:        test
 date:        Mon Jan 12 13:46:40 1970 +0000
 summary:     add
 
 checking changesets
 checking manifests
 crosschecking files in changesets and manifests
 checking files
 2 files, 2 changesets, 3 total revisions
 bleah
-# push should succeed
+# push should succeed even though it has an unexpected response
 pushing to ssh://user@dummy/remote
 searching for changes
 note: unsynced remote changes!
 remote: adding changesets
 remote: adding manifests
 remote: adding file changes
 remote: added 1 changesets with 1 changes to 1 files
+abort: unexpected response:
+'KABOOM1\n'
+changeset:   3:ac7448082955
+tag:         tip
+parent:      1:572896fe480d
+user:        test
+date:        Mon Jan 12 13:46:42 1970 +0000
+summary:     z
+
+changeset:   2:187c6caa0d1e
+parent:      0:e34318c26897
+user:        test
+date:        Mon Jan 12 13:46:41 1970 +0000
+summary:     z
+
 Got arguments 1:user@dummy 2:hg -R nonexistent serve --stdio
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 Got arguments 1:user@dummy 2:hg -R local serve --stdio
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 changegroup-in-remote hook: HG_NODE=572896fe480d7581849806ee402175c49cb20037 HG_SOURCE=serve HG_URL=remote:ssh:127.0.0.1 
 Got arguments 1:user@dummy 2:hg -R remote serve --stdio
 changegroup-in-remote hook: HG_NODE=ac7448082955a0b2ff5cb4512c1e061c779bbc79 HG_SOURCE=serve HG_URL=remote:ssh:127.0.0.1