2 files changed, 43 insertions, 9 deletions
diff --git a/internal/shim/parent.go b/internal/shim/parent.go
index 9382db22..0e903838 100644
--- a/internal/shim/parent.go
+++ b/internal/shim/parent.go
@@ -14,19 +14,18 @@ import (
 
 // called in the parent process
 
-func ServeConfig(socket string, uid int, payload *Payload, wl string, done chan struct{}) (*net.UnixConn, error) {
-	var ws *net.UnixConn
+func ServeConfig(socket string, uid int, payload *Payload, wl *Wayland) error {
 	if payload.WL {
-		if f, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: wl, Net: "unix"}); err != nil {
-			return nil, err
+		if f, err := net.DialUnix("unix", nil, &net.UnixAddr{Name: wl.Path, Net: "unix"}); err != nil {
+			return err
 		} else {
 			verbose.Println("connected to wayland at", wl)
-			ws = f
+			wl.UnixConn = f
 		}
 	}
 
 	if c, err := net.ListenUnix("unix", &net.UnixAddr{Name: socket, Net: "unix"}); err != nil {
-		return nil, err
+		return err
 	} else {
 		verbose.Println("configuring shim on socket", socket)
 		if err = acl.UpdatePerm(socket, uid, acl.Read, acl.Write, acl.Execute); err != nil {
@@ -47,7 +46,7 @@ func ServeConfig(socket string, uid int, payload *Payload, wl string, done chan
 				if payload.WL {
 					// get raw connection
 					var rc syscall.RawConn
-					if rc, err = ws.SyscallConn(); err != nil {
+					if rc, err = wl.SyscallConn(); err != nil {
 						fmt.Println("fortify: cannot obtain raw wayland connection:", err)
 						return
 					} else {
@@ -61,7 +60,7 @@ func ServeConfig(socket string, uid int, payload *Payload, wl string, done chan
 								_ = conn.Close()
 
 								// block until shim exits
-								<-done
+								<-wl.done
 								verbose.Println("releasing wayland connection")
 							}); err != nil {
 								fmt.Println("fortify: cannot obtain wayland connection fd:", err)
@@ -79,6 +78,6 @@ func ServeConfig(socket string, uid int, payload *Payload, wl string, done chan
 				fmt.Println("fortify: cannot remove dangling shim socket:", err)
 			}
 		}()
-		return ws, nil
+		return nil
 	}
 }
diff --git a/internal/shim/wayland.go b/internal/shim/wayland.go
new file mode 100644
index 00000000..05f04278
--- /dev/null
+++ b/internal/shim/wayland.go
@@ -0,0 +1,35 @@
+package shim
+
+import (
+	"net"
+	"sync"
+)
+
+// Wayland implements wayland mediation.
+type Wayland struct {
+	// wayland socket path
+	Path string
+
+	// wayland connection
+	*net.UnixConn
+
+	connErr error
+	sync.Once
+	// wait for wayland client to exit
+	done chan struct{}
+}
+
+func (wl *Wayland) Close() error {
+	wl.Do(func() {
+		close(wl.done)
+		wl.connErr = wl.UnixConn.Close()
+	})
+
+	return wl.connErr
+}
+
+func NewWayland() *Wayland {
+	wl := new(Wayland)
+	wl.done = make(chan struct{})
+	return wl
+}