Fix piping

kylecarbs · kylecarbs · commit 9c949f8ea485 · 2023-04-21T14:39:34.000Z
diff --git a/agent/agentssh/x11.go b/agent/agentssh/x11.go
@@ -61,6 +61,7 @@ func (s *Server) x11Handler(ctx ssh.Context, x11 ssh.X11) bool {
 
 	go func() {
 		defer s.trackListener(listener, false)
+		handledFirstConnection := false
 
 		for {
 			conn, err := listener.Accept()
@@ -71,6 +72,13 @@ func (s *Server) x11Handler(ctx ssh.Context, x11 ssh.X11) bool {
 				s.logger.Warn(ctx, "failed to accept X11 connection", slog.Error(err))
 				return
 			}
+			if x11.SingleConnection && handledFirstConnection {
+				s.logger.Warn(ctx, "X11 connection rejected because single connection is enabled")
+				_ = conn.Close()
+				continue
+			}
+			handledFirstConnection = true
+
 			unixConn, ok := conn.(*net.UnixConn)
 			if !ok {
 				s.logger.Warn(ctx, fmt.Sprintf("failed to cast connection to UnixConn. got: %T", conn))
diff --git a/agent/agentssh/x11_test.go b/agent/agentssh/x11_test.go
@@ -68,9 +68,12 @@ func TestServer_X11(t *testing.T) {
 	require.NoError(t, err)
 
 	x11Chans := c.HandleChannelOpen("x11")
+	payload := "hello world"
 	require.Eventually(t, func() bool {
 		conn, err := net.Dial("unix", filepath.Join(dir, "X0"))
 		if err == nil {
+			_, err = conn.Write([]byte(payload))
+			assert.NoError(t, err)
 			_ = conn.Close()
 		}
 		return err == nil
@@ -80,9 +83,12 @@ func TestServer_X11(t *testing.T) {
 	ch, reqs, err := x11.Accept()
 	require.NoError(t, err)
 	go gossh.DiscardRequests(reqs)
-	err = ch.Close()
+	got := make([]byte, len(payload))
+	_, err = ch.Read(got)
 	require.NoError(t, err)
-	s.Close()
+	assert.Equal(t, payload, string(got))
+	_ = ch.Close()
+	_ = s.Close()
 	<-done
 
 	// Ensure the Xauthority file was written!
