Fix daemon races when tunnels are opened and closed concurrently

internal/daemon/daemon.go

@@ -18,6 +18,7 @@ import (
 	"github.com/alebeck/boring/internal/ipc"
 	"github.com/alebeck/boring/internal/log"
 	"github.com/alebeck/boring/internal/tunnel"
+	"golang.org/x/sync/singleflight"
 )
 
 const (
@@ -49,6 +50,8 @@ type daemon struct {
 	// TODO: write proper concurrent map structure for this
 	tunnels map[string]*tunnel.Tunnel
 	mutex   sync.RWMutex
+	// Merges concurrent opens of the same tunnel into one attempt
+	opening singleflight.Group
 
 	once sync.Once
 	wg   sync.WaitGroup
@@ -157,33 +160,54 @@ func (d *daemon) openTunnel(conn net.Conn, desc *tunnel.Desc) {
 	var err error
 	defer func() { respond(conn, err, nil) }()
 
+	// Clients joining an open in progress report the tunnel as already running
+	leader := false
+	_, err, _ = d.opening.Do(desc.Name, func() (any, error) {
+		leader = true
+		return nil, d.open(desc)
+	})
+	if err == nil && !leader {
+		err = AlreadyRunning
+	}
+	if err != nil {
+		log.Errorf("%v: could not open: %v", desc.Name, err)
+	}
+}
+
+func (d *daemon) open(desc *tunnel.Desc) error {
 	d.mutex.RLock()
 	_, exists := d.tunnels[desc.Name]
 	d.mutex.RUnlock()
 	if exists {
-		err = AlreadyRunning
-		log.Errorf("%v: could not open: %v", desc.Name, err)
-		return
+		return AlreadyRunning
 	}
 
 	t := tunnel.FromDesc(desc)
-	if err = t.Open(); err != nil {
-		log.Errorf("%v: could not open: %v", t.Name, err)
-		return
+	if err := t.Open(); err != nil {
+		return err
 	}
 
 	d.mutex.Lock()
-	d.tunnels[t.Name] = t
+	d.tunnels[desc.Name] = t
 	d.mutex.Unlock()
 
 	// Register closing logic
 	go func() {
 		<-t.Closed
-		d.mutex.Lock()
-		delete(d.tunnels, t.Name)
-		d.mutex.Unlock()
-		log.Infof("Closed tunnel %s", t.Name)
+		d.removeTunnel(t)
+		log.Infof("Closed tunnel %s", desc.Name)
 	}()
+	return nil
+}
+
+// removeTunnel forgets about t, unless its name has meanwhile been taken
+// by a tunnel that was opened after it.
+func (d *daemon) removeTunnel(t *tunnel.Tunnel) {
+	d.mutex.Lock()
+	defer d.mutex.Unlock()
+	if d.tunnels[t.Name] == t {
+		delete(d.tunnels, t.Name)
+	}
 }
 
 func (d *daemon) closeTunnel(conn net.Conn, q *tunnel.Desc) {
@@ -204,13 +228,16 @@ func (d *daemon) closeTunnel(conn net.Conn, q *tunnel.Desc) {
 		return
 	}
 	<-t.Closed
+	// Also remove it here, so it is gone by the time the client gets the
+	// response, rather than whenever the closing goroutine gets to it.
+	d.removeTunnel(t)
 }
 
 func (d *daemon) listTunnels(conn net.Conn) {
 	d.mutex.RLock()
 	ts := make(map[string]tunnel.Desc, len(d.tunnels))
 	for n, t := range d.tunnels {
-		ts[n] = *t.Desc
+		ts[n] = t.Snapshot()
 	}
 	d.mutex.RUnlock()
 	respond(conn, nil, ts)

internal/daemon/daemon_test.go

@@ -0,0 +1,24 @@
+package daemon
+
+import (
+	"testing"
+
+	"github.com/alebeck/boring/internal/tunnel"
+)
+
+func TestRemoveTunnelKeepsNewer(t *testing.T) {
+	d := &daemon{tunnels: make(map[string]*tunnel.Tunnel)}
+	old := tunnel.FromDesc(&tunnel.Desc{Name: "a"})
+	cur := tunnel.FromDesc(&tunnel.Desc{Name: "a"})
+	d.tunnels["a"] = cur
+
+	d.removeTunnel(old)
+	if d.tunnels["a"] != cur {
+		t.Fatal("removing an old tunnel dropped the one that replaced it")
+	}
+
+	d.removeTunnel(cur)
+	if _, ok := d.tunnels["a"]; ok {
+		t.Fatal("tunnel not removed")
+	}
+}

internal/tunnel/tunnel.go

@@ -46,11 +46,15 @@ type Tunnel struct {
 	hops       []ssh_config.Hop
 	Closed     chan struct{}
 	stop       chan struct{}
+	stopOnce   sync.Once
 	listener   net.Listener
 	wg         sync.WaitGroup
 	client     *ssh.Client
 	localAddr  *address
 	remoteAddr *address
+	// mu guards Status and LastConn, which the tunnel's own goroutines
+	// update while the daemon may be reading them for a listing.
+	mu sync.Mutex
 	*Desc
 }
 
@@ -85,14 +89,31 @@ func (t *Tunnel) Open() (err error) {
 		t.Closed = make(chan struct{})
 	}
 
+	t.mu.Lock()
+	t.Status = Open
+	t.LastConn = time.Now()
+	t.mu.Unlock()
+
 	go t.run()
 
 	log.Infof("%v: opened tunnel", t.Name)
-	t.Status = Open
-	t.LastConn = time.Now()
 	return
 }
 
+// Snapshot returns a copy of the tunnel's description that is safe to take
+// while the tunnel is running.
+func (t *Tunnel) Snapshot() Desc {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return *t.Desc
+}
+
+func (t *Tunnel) setStatus(s Status) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	t.Status = s
+}
+
 func (t *Tunnel) prepare() error {
 	// We need to pass the user as it's needed for matching Match blocks
 	sc, err := ssh_config.ParseSSHConfig(t.Host, t.User)
@@ -175,7 +196,7 @@ func (t *Tunnel) makeClient() error {
 	}
 
 	// Wait for all wrapped clients to close in case of tunnel closing or reconnection
-	go t.waitFor(func() { wg.Wait() })
+	t.goWait(wg.Wait)
 
 	t.client = c
 	return nil
@@ -222,8 +243,8 @@ func (t *Tunnel) run() {
 		close(disconn)
 	}()
 
-	go t.waitFor(func() { t.keepAlive(disconn) })
-	go t.waitFor(func() { t.handleConns() })
+	t.goWait(func() { t.keepAlive(disconn) })
+	t.goWait(t.handleConns)
 
 	stopped := false
 	select {
@@ -243,7 +264,7 @@ func (t *Tunnel) run() {
 			return
 		}
 	}
-	t.Status = Closed
+	t.setStatus(Closed)
 	close(t.Closed)
 }
 
@@ -290,7 +311,7 @@ func (t *Tunnel) handleForward() {
 			log.Errorf("%v: could not accept: %v", t.Name, err)
 			return
 		}
-		go t.waitFor(func() {
+		t.goWait(func() {
 			addr := t.remoteAddr
 			if t.Mode == Remote || t.Mode == RemoteSocks {
 				addr = t.localAddr
@@ -335,12 +356,12 @@ func (t *Tunnel) handleSocks() {
 			log.Errorf("%v: could not accept: %v", t.Name, err)
 			return
 		}
-		go t.waitFor(func() { serv.ServeConn(conn) })
+		t.goWait(func() { serv.ServeConn(conn) })
 	}
 }
 
 func (t *Tunnel) reconnectLoop() error {
-	t.Status = Reconn
+	t.setStatus(Reconn)
 	timeout := time.After(reconnectTimeout)
 	wait := time.NewTimer(2 * time.Millisecond) // First time try (essent.) immediately
 	waitTime := initReconnectWait
@@ -368,20 +389,28 @@ func (t *Tunnel) reconnectLoop() error {
 	}
 }
 
+// Close signals the tunnel to stop. It is safe to call concurrently and
+// more than once; wait on Closed for the tunnel to actually shut down.
 func (t *Tunnel) Close() error {
-	if t.Status == Closed {
+	t.mu.Lock()
+	closed := t.Status == Closed
+	t.mu.Unlock()
+	if closed {
 		return fmt.Errorf("trying to close a closed tunnel")
 	}
-	close(t.stop)
+	t.stopOnce.Do(func() { close(t.stop) })
 	return nil
 }
 
-// Logic registered with waitFor will be waited for upon tunnel closing
-// and reconnecting.
-func (t *Tunnel) waitFor(f func()) {
+// goWait runs f in a new goroutine that will be waited for upon tunnel
+// closing and reconnecting. The wait group is incremented before the
+// goroutine starts, so a concurrent Wait cannot miss it.
+func (t *Tunnel) goWait(f func()) {
 	t.wg.Add(1)
-	defer t.wg.Done()
-	f()
+	go func() {
+		defer t.wg.Done()
+		f()
+	}()
 }
 
 func parseAddr(addr string, allowShort bool) (*address, error) {

internal/tunnel/tunnel_test.go

@@ -0,0 +1,73 @@
+package tunnel
+
+import (
+	"sync"
+	"testing"
+)
+
+// openTunnel returns a tunnel in the state Open leaves it in, without
+// connecting anywhere.
+func openTunnel() *Tunnel {
+	return &Tunnel{
+		Desc:   &Desc{Name: "test", Status: Open},
+		stop:   make(chan struct{}),
+		Closed: make(chan struct{}),
+	}
+}
+
+func TestCloseTwice(t *testing.T) {
+	tun := openTunnel()
+	if err := tun.Close(); err != nil {
+		t.Fatalf("first close: %v", err)
+	}
+	// The tunnel stays open until its run loop notices the stop signal,
+	// so a second close in that window must not panic.
+	if err := tun.Close(); err != nil {
+		t.Fatalf("second close: %v", err)
+	}
+	select {
+	case <-tun.stop:
+	default:
+		t.Fatal("stop channel not closed")
+	}
+}
+
+func TestCloseConcurrent(t *testing.T) {
+	tun := openTunnel()
+	var wg sync.WaitGroup
+	for range 10 {
+		wg.Go(func() {
+			if err := tun.Close(); err != nil {
+				t.Errorf("close: %v", err)
+			}
+		})
+	}
+	wg.Wait()
+}
+
+func TestCloseClosed(t *testing.T) {
+	tun := openTunnel()
+	tun.Status = Closed
+	if err := tun.Close(); err == nil {
+		t.Fatal("expected error when closing a closed tunnel")
+	}
+}
+
+// Run with -race: status updates from the tunnel's goroutines must not race
+// with the daemon taking snapshots for a listing.
+func TestSnapshotConcurrent(t *testing.T) {
+	tun := openTunnel()
+	var wg sync.WaitGroup
+	wg.Go(func() {
+		for range 1000 {
+			tun.setStatus(Reconn)
+			tun.setStatus(Open)
+		}
+	})
+	for range 1000 {
+		if s := tun.Snapshot().Status; s != Open && s != Reconn {
+			t.Fatalf("unexpected status %v", s)
+		}
+	}
+	wg.Wait()
+}