# HG changeset patch
# User Alexis S. L. Carvalho <alexis@cecm.usp.br>
# Date 2008-02-03 23:03:46
# Node ID 7939c71f3132ec79f420e79302c0a81cb447f20d
# Parent  e82dd2dd62769fc6165f630b6cc30c055be2837e

sshrepo: be more careful while reading data

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

diff --git a/mercurial/sshrepo.py b/mercurial/sshrepo.py
--- a/mercurial/sshrepo.py
+++ b/mercurial/sshrepo.py
@@ -115,14 +115,25 @@ class sshrepository(remoterepository):
         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")
@@ -183,25 +194,22 @@ class sshrepository(remoterepository):
 
         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")
@@ -209,18 +217,21 @@ class sshrepository(remoterepository):
             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')
diff --git a/tests/test-ssh b/tests/test-ssh
--- a/tests/test-ssh
+++ b/tests/test-ssh
@@ -27,6 +27,11 @@ 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
@@ -91,13 +96,16 @@ 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
diff --git a/tests/test-ssh.out b/tests/test-ssh.out
--- a/tests/test-ssh.out
+++ b/tests/test-ssh.out
@@ -70,7 +70,7 @@ crosschecking files in changesets and ma
 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!
@@ -78,6 +78,21 @@ 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