--- a/states.py	Mon May 23 15:02:16 2011 +0200
+++ b/states.py	Mon May 23 15:41:20 2011 +0200
@@ -29,6 +29,7 @@
 from mercurial.node import nullid
 from mercurial import discovery
 from mercurial import extensions
+from mercurial import wireproto
 
 
 _NOPULLPUSH=2
@@ -106,6 +107,19 @@
     extensions.wrapfunction(discovery, 'findcommonoutgoing', filterprivateout)
     extensions.wrapfunction(discovery, 'findcommonincoming', filterprivatein)
 
+    # Write protocols
+    ####################
+    def heads(repo, proto):
+        h = repo._publicheads
+        return wireproto.encodelist(h) + "\n"
+
+    def _reducehead(wirerepo, heads):
+        """heads filtering is done repo side"""
+        return heads
+
+    wireproto.wirerepository._reducehead = _reducehead
+    wireproto.commands['heads'] = (heads, '')
+
 def reposetup(ui, repo):
 
     if not repo.local():

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/killdaemons.py	Mon May 23 15:41:20 2011 +0200
@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+
+import os, time, errno, signal
+
+# Kill off any leftover daemon processes
+try:
+    fp = open(os.environ['DAEMON_PIDS'])
+    for line in fp:
+        try:
+            pid = int(line)
+        except ValueError:
+            continue
+        try:
+            os.kill(pid, 0)
+            os.kill(pid, signal.SIGTERM)
+            for i in range(10):
+                time.sleep(0.05)
+                os.kill(pid, 0)
+            os.kill(pid, signal.SIGKILL)
+        except OSError, err:
+            if err.errno != errno.ESRCH:
+                raise
+    fp.close()
+except IOError:
+    pass

--- a/tests/test-private.t	Mon May 23 15:02:16 2011 +0200
+++ b/tests/test-private.t	Mon May 23 15:41:20 2011 +0200
@@ -1,4 +1,7 @@
   $ cat >> $HGRCPATH <<EOF
+  > [web]
+  > push_ssl = false
+  > allow_push = *
   > [extensions]
   > EOF
   $ echo "states=$(echo $(dirname $TESTDIR))/states.py" >> $HGRCPATH
@@ -123,8 +126,81 @@
 
   $ hg -R local serve -p $HGPORT -d --pid-file=local.pid
   $ cat local.pid >> "$DAEMON_PIDS"
-  $ hg clone http://localhost:$HGPORT/ other3
-  $ hg -R other2 log --template='{rev}:{node|short}\n'
+  $ hg clone http://localhost:$HGPORT/ fromhttp
+  requesting all changes
+  adding changesets
+  adding manifests
+  adding file changes
+  added 2 changesets with 2 changes to 1 files
+  updating to branch default
+  1 files updated, 0 files merged, 0 files removed, 0 files unresolved
+  $ hg -R fromhttp log --template='{rev}:{node|short}\n'
+  1:710fe444b3b0
+  0:5caa672bac26
+
+  $ hg init fromhttp2
+  $ cd fromhttp2
+  $ hg inc http://localhost:$HGPORT/ --template='{rev}:{node|short}\n'
+  comparing with http://localhost:$HGPORT/
+  0:5caa672bac26
+  1:710fe444b3b0
+  $ hg pull http://localhost:$HGPORT/
+  pulling from http://localhost:$HGPORT/
+  requesting all changes
+  adding changesets
+  adding manifests
+  adding file changes
+  added 2 changesets with 2 changes to 1 files
+  (run 'hg update' to get a working copy)
+  $ hg log --template='{rev}:{node|short}\n'
   1:710fe444b3b0
   0:5caa672bac26
-  $ kill `cat local.pid`
+  $ hg inc http://localhost:$HGPORT/ --template='{rev}:{node|short}\n'
+  comparing with http://localhost:$HGPORT/
+  searching for changes
+  no changes found
+  [1]
+
+turn private off again (repo side)
+  $ cd ..
+  $ "$TESTDIR/killdaemons.py"
+  $ sed -i 's/^private=.*$/private=off/' ./local/.hg/hgrc
+  $ hg -R local serve -p $HGPORT -d --pid-file=local.pid
+  $ cat local.pid >> "$DAEMON_PIDS"
+  $ cd fromhttp2
+
+  $ hg inc http://localhost:$HGPORT/ --template='{rev}:{node|short}\n'
+  comparing with http://localhost:$HGPORT/
+  searching for changes
+  2:3c8695235a32
+  3:73585b17392a
+  $ hg pull http://localhost:$HGPORT/
+  pulling from http://localhost:$HGPORT/
+  searching for changes
+  adding changesets
+  adding manifests
+  adding file changes
+  added 2 changesets with 2 changes to 2 files
+  (run 'hg update' to get a working copy)
+  $ cd ..
+  $ "$TESTDIR/killdaemons.py"
+
+turn private on again (repo side)
+  $ sed -i 's/^private=.*$/private=on/' local/.hg/hgrc
+  $ hg init httpto
+  $ hg -R httpto serve -p $HGPORT -d --pid-file=remote.pid
+  $ cat remote.pid >> "$DAEMON_PIDS"
+  $ cd local
+  $ hg out http://localhost:$HGPORT/  --template='{rev}:{node|short}\n'
+  comparing with http://localhost:$HGPORT/
+  searching for changes
+  0:5caa672bac26
+  1:710fe444b3b0
+  $ hg push http://localhost:$HGPORT/
+  pushing to http://localhost:$HGPORT/
+  searching for changes
+  remote: adding changesets
+  remote: adding manifests
+  remote: adding file changes
+  remote: added 2 changesets with 2 changes to 1 files
+  $ "$TESTDIR/killdaemons.py"