Repository: godbus/dbus
Pull request #171 (2019-07-29 00:36:15)

Add a context to connections
This PR attaches a context to D-Bus connections.  It adds a `WithContext(ctx)` option to the connection constructor to set a parent context for the connection (overriding a default of `context.Background()`), and adds a `Conn.Context` method to retrieve the connection's context.  The connection's context will be done when either the parent context is done, or a disconnect is detected.  In either case, the connection is closed when the context is done.

The use cases for this are:

1. by setting a parent context for a connection, it can automatically be cleaned up with related code.

2. by monitoring the context of the shared `SessionBus` or `SystemBus` connections, it is possible to have a daemon exit cleanly when the bus exits (e.g. on user session shutdown).

3. while it isn't implemented in this PR, the context could be used to derive one passed to server methods.

amenzhinsky: I think a context should only affect the dialing process, it's something that is passed along with each blocking call rather than assigned to a client entirely.
jhenstridge: The main uses I see for having a context attached to the connection itself are:

1. provide a way to signal that the D-Bus connection has been closed.  Depending on the application, you might want to exit on this condition, or try to reconnect to the bus and re-export services, etc.
2. provide a parent for contexts assigned to incoming method calls.  Contexts are useful for serving incoming calls, since (a) it can be used to assign deadlines and (b) they can be used as a context for outgoing method calls required to service the request (e.g. a server performing a call-out to polkitd to authenticate its client).

You could probably do a custom disconnect reporting API for (1), but when you add (2) to the mix it seems like `Context` is a natural fit.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 6c3d03c..dd67672 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,28 +8,10 @@ go:
   - 1.13.x
   - tip
 
-env:
-  global:
-  matrix:
-   - TARGET=amd64
-   - TARGET=arm64
-   - TARGET=arm
-   - TARGET=386
-   - TARGET=ppc64le
-
 matrix:
   fast_finish: true
   allow_failures:
     - go: tip
-  exclude:
-  - go: tip
-    env: TARGET=arm
-  - go: tip
-    env: TARGET=arm64
-  - go: tip
-    env: TARGET=386
-  - go: tip
-    env: TARGET=ppc64le
 
 addons:
   apt:
@@ -43,3 +25,26 @@ before_install:
 script:
   - go test -v -race -mod=readonly ./...     # Run all the tests with the race detector enabled
   - go vet ./...                             # go vet is the official Go static analyzer
+
+jobs:
+  include:
+    # The build matrix doesn't cover build stages, so manually expand
+    # the jobs with anchors
+    - &multiarch
+      stage: "Multiarch Test"
+      go: 1.11.x
+      env: TARGETS="386 arm arm64 ppc64le"
+      before_install:
+        - docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
+      script:
+        - |
+          set -e
+          for target in $TARGETS; do
+            printf "\e[1mRunning test suite under ${target}.\e[0m\n"
+            GOARCH="$target" go test -v ./...
+            printf "\n\n"
+          done
+    - <<: *multiarch
+      go: 1.12.x
+    - <<: *multiarch
+      go: 1.13.x
diff --git a/conn.go b/conn.go
index 0255fda..b55bc99 100644
--- a/conn.go
+++ b/conn.go
@@ -30,6 +30,12 @@ var ErrClosed = errors.New("dbus: connection closed by user")
 type Conn struct {
 	transport
 
+	ctx       context.Context
+	cancelCtx context.CancelFunc
+
+	closeOnce sync.Once
+	closeErr  error
+
 	busObj BusObject
 	unixFD bool
 	uuid   string
@@ -210,6 +216,14 @@ func WithOutgoingInterceptor(interceptor Interceptor) ConnOption {
 	}
 }
 
+// WithContext overrides  the default context for the connection.
+func WithContext(ctx context.Context) ConnOption {
+	return func(conn *Conn) error {
+		conn.ctx = ctx
+		return nil
+	}
+}
+
 // NewConn creates a new private *Conn from an already established connection.
 func NewConn(conn io.ReadWriteCloser, opts ...ConnOption) (*Conn, error) {
 	return newConn(genericTransport{conn}, opts...)
@@ -231,6 +245,15 @@ func newConn(tr transport, opts ...ConnOption) (*Conn, error) {
 			return nil, err
 		}
 	}
+	if conn.ctx == nil {
+		conn.ctx = context.Background()
+	}
+	conn.ctx, conn.cancelCtx = context.WithCancel(conn.ctx)
+	go func() {
+		<-conn.ctx.Done()
+		conn.Close()
+	}()
+
 	conn.calls = newCallTracker()
 	if conn.handler == nil {
 		conn.handler = NewDefaultHandler()
@@ -257,22 +280,33 @@ func (conn *Conn) BusObject() BusObject {
 // and the channels passed to Eavesdrop and Signal are closed. This method must
 // not be called on shared connections.
 func (conn *Conn) Close() error {
-	conn.outHandler.close()
-	if term, ok := conn.signalHandler.(Terminator); ok {
-		term.Terminate()
-	}
+	conn.closeOnce.Do(func() {
+		conn.outHandler.close()
+		if term, ok := conn.signalHandler.(Terminator); ok {
+			term.Terminate()
+		}
 
-	if term, ok := conn.handler.(Terminator); ok {
-		term.Terminate()
-	}
+		if term, ok := conn.handler.(Terminator); ok {
+			term.Terminate()
+		}
 
-	conn.eavesdroppedLck.Lock()
-	if conn.eavesdropped != nil {
-		close(conn.eavesdropped)
-	}
-	conn.eavesdroppedLck.Unlock()
+		conn.eavesdroppedLck.Lock()
+		if conn.eavesdropped != nil {
+			close(conn.eavesdropped)
+		}
+		conn.eavesdroppedLck.Unlock()
+
+		conn.cancelCtx()
+
+		conn.closeErr = conn.transport.Close()
+	})
+	return conn.closeErr
+}
 
-	return conn.transport.Close()
+// Context returns the context associated with the connection.  The
+// context will be cancelled when the connection is closed.
+func (conn *Conn) Context() context.Context {
+	return conn.ctx
 }
 
 // Eavesdrop causes conn to send all incoming messages to the given channel


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 13c15e5..ef3b8d9 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -1,6 +1,7 @@
 package dbus
 
 import (
+	"context"
 	"encoding/binary"
 	"io"
 	"io/ioutil"
@@ -453,3 +454,107 @@ func TestInterceptors(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestCloseCancelsConnectionContext(t *testing.T) {
+	bus, err := SessionBusPrivate()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer bus.Close()
+
+	if err = bus.Auth(nil); err != nil {
+		t.Fatal(err)
+	}
+	if err = bus.Hello(); err != nil {
+		t.Fatal(err)
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// The context is not done at this point
+	ctx := bus.Context()
+	select {
+	case <-ctx.Done():
+		t.Fatal("context should not be done")
+	default:
+	}
+
+	err = bus.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-ctx.Done():
+		// expected
+	case <-time.After(5 * time.Second):
+		t.Fatal("context is not done after connection closed")
+	}
+}
+
+func TestDisconnectCancelsConnectionContext(t *testing.T) {
+	reader, pipewriter := io.Pipe()
+	defer pipewriter.Close()
+	defer reader.Close()
+
+	bus, err := NewConn(rwc{Reader: reader, Writer: ioutil.Discard})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		_, err := pipewriter.Write([]byte("REJECTED name\r\nOK myuuid\r\n"))
+		if err != nil {
+			t.Errorf("error writing to pipe: %v", err)
+		}
+	}()
+	err = bus.Auth([]Auth{fakeAuth{}})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := bus.Context()
+
+	pipewriter.Close()
+	select {
+	case <-ctx.Done():
+		// expected
+	case <-time.After(5 * time.Second):
+		t.Fatal("context is not done after connection closed")
+	}
+}
+
+func TestCancellingContextClosesConnection(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	reader, pipewriter := io.Pipe()
+	defer pipewriter.Close()
+	defer reader.Close()
+
+	bus, err := NewConn(rwc{Reader: reader, Writer: ioutil.Discard}, WithContext(ctx))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		_, err := pipewriter.Write([]byte("REJECTED name\r\nOK myuuid\r\n"))
+		if err != nil {
+			t.Errorf("error writing to pipe: %v", err)
+		}
+	}()
+	err = bus.Auth([]Auth{fakeAuth{}})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Cancel the connection's parent context and give time for
+	// other goroutines to schedule.
+	cancel()
+	time.Sleep(50 * time.Millisecond)
+
+	err = bus.BusObject().Call("org.freedesktop.DBus.Peer.Ping", 0).Store()
+	if err != ErrClosed {
+		t.Errorf("expected connection to be closed, but got: %v", err)
+	}
+}