Only run X forwarding on Linux

coder · kylecarbs · Apr 21, 2023 · Apr 19, 2023 · Apr 20, 2023 · Apr 21, 2023
commit 5ad863394c95f504a82b786daae220dd150dd9d1
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -49,6 +49,7 @@ const (
 
 type Server struct {
 	mu        sync.RWMutex // Protects following.
+	fs        afero.Fs
 	listeners map[net.Listener]struct{}
 	conns     map[net.Conn]struct{}
 	sessions  map[ssh.Session]struct{}
@@ -85,16 +86,13 @@ func NewServer(ctx context.Context, logger slog.Logger, fs afero.Fs, maxTimeout
 	if x11SocketDir == "" {
 		x11SocketDir = filepath.Join(os.TempDir(), ".X11-unix")
 	}
-	err = fs.MkdirAll(x11SocketDir, 0700)
-	if err != nil {
-		return nil, err
-	}
 
 	forwardHandler := &ssh.ForwardedTCPHandler{}
 	unixForwardHandler := &forwardedUnixHandler{log: logger}
 
 	s := &Server{
 		listeners:    make(map[net.Listener]struct{}),
+		fs:           fs,
 		conns:        make(map[net.Conn]struct{}),
 		sessions:     make(map[ssh.Session]struct{}),
 		logger:       logger,
@@ -135,9 +133,7 @@ func NewServer(ctx context.Context, logger slog.Logger, fs afero.Fs, maxTimeout
 			"streamlocal-forward@openssh.com":        unixForwardHandler.HandleSSHRequest,
 			"cancel-streamlocal-forward@openssh.com": unixForwardHandler.HandleSSHRequest,
 		},
-		X11Callback: func(ctx ssh.Context, x11 ssh.X11) bool {
-			return x11Callback(logger, fs, ctx, x11)
-		},
+		X11Callback: s.x11Callback,
 		ServerConfigCallback: func(ctx ssh.Context) *gossh.ServerConfig {
 			return &gossh.ServerConfig{
 				NoClientAuth: true,

diff --git a/agent/agentssh/x11.go b/agent/agentssh/x11.go
@@ -1,6 +1,7 @@
 package agentssh
 
 import (
+	"context"
 	"encoding/binary"
 	"encoding/hex"
 	"errors"
@@ -9,8 +10,10 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"time"
 
 	"github.com/gliderlabs/ssh"
+	"github.com/gofrs/flock"
 	"github.com/spf13/afero"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
@@ -20,16 +23,22 @@ import (
 
 // x11Callback is called when the client requests X11 forwarding.
 // It adds an Xauthority entry to the Xauthority file.
-func x11Callback(logger slog.Logger, fs afero.Fs, ctx ssh.Context, x11 ssh.X11) bool {
+func (s *Server) x11Callback(ctx ssh.Context, x11 ssh.X11) bool {
 	hostname, err := os.Hostname()
 	if err != nil {
-		logger.Warn(ctx, "failed to get hostname", slog.Error(err))
+		s.logger.Warn(ctx, "failed to get hostname", slog.Error(err))
 		return false
 	}
 
-	err = addXauthEntry(fs, hostname, strconv.Itoa(int(x11.ScreenNumber)), x11.AuthProtocol, x11.AuthCookie)
+	err = s.fs.MkdirAll(s.x11SocketDir, 0o700)
 	if err != nil {
-		logger.Warn(ctx, "failed to add Xauthority entry", slog.Error(err))
+		s.logger.Warn(ctx, "failed to make the x11 socket dir", slog.F("dir", s.x11SocketDir), slog.Error(err))
+		return false
+	}
+
+	err = addXauthEntry(ctx, s.fs, hostname, strconv.Itoa(int(x11.ScreenNumber)), x11.AuthProtocol, x11.AuthCookie)
+	if err != nil {
+		s.logger.Warn(ctx, "failed to add Xauthority entry", slog.Error(err))
 		return false
 	}
 	return true
@@ -64,16 +73,16 @@ func (s *Server) x11Handler(ctx ssh.Context, x11 ssh.X11) bool {
 			}
 			unixConn, ok := conn.(*net.UnixConn)
 			if !ok {
-				s.logger.Warn(ctx, "failed to cast connection to UnixConn")
+				s.logger.Warn(ctx, fmt.Sprintf("failed to cast connection to UnixConn. got: %T", conn))
 				return
 			}
 			unixAddr, ok := unixConn.LocalAddr().(*net.UnixAddr)
 			if !ok {
-				s.logger.Warn(ctx, "failed to cast local address to UnixAddr")
+				s.logger.Warn(ctx, fmt.Sprintf("failed to cast local address to UnixAddr. got: %T", unixConn.LocalAddr()))
 				return
 			}
 
-			channel, _, err := serverConn.OpenChannel("x11", gossh.Marshal(struct {
+			channel, reqs, err := serverConn.OpenChannel("x11", gossh.Marshal(struct {
 				OriginatorAddress string
 				OriginatorPort    uint32
 			}{
@@ -84,7 +93,7 @@ func (s *Server) x11Handler(ctx ssh.Context, x11 ssh.X11) bool {
 				s.logger.Warn(ctx, "failed to open X11 channel", slog.Error(err))
 				return
 			}
-
+			go gossh.DiscardRequests(reqs)
 			go Bicopy(ctx, conn, channel)
 		}
 	}()
@@ -93,7 +102,7 @@ func (s *Server) x11Handler(ctx ssh.Context, x11 ssh.X11) bool {
 
 // addXauthEntry adds an Xauthority entry to the Xauthority file.
 // The Xauthority file is located at ~/.Xauthority.
-func addXauthEntry(fs afero.Fs, host string, display string, authProtocol string, authCookie string) error {
+func addXauthEntry(ctx context.Context, fs afero.Fs, host string, display string, authProtocol string, authCookie string) error {
 	// Get the Xauthority file path
 	homeDir, err := os.UserHomeDir()
 	if err != nil {
@@ -102,8 +111,15 @@ func addXauthEntry(fs afero.Fs, host string, display string, authProtocol string
 
 	xauthPath := filepath.Join(homeDir, ".Xauthority")
 
+	lock := flock.New(xauthPath)
+	ok, err := lock.TryLockContext(ctx, 100*time.Millisecond)
+	if !ok {
+		return xerrors.Errorf("failed to lock Xauthority file: %w", err)
+	}
+	defer lock.Close()
+
 	// Open or create the Xauthority file
-	file, err := fs.OpenFile(xauthPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
+	file, err := fs.OpenFile(xauthPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o600)
 	if err != nil {
 		return xerrors.Errorf("failed to open Xauthority file: %w", err)
 	}

diff --git a/agent/agentssh/x11_test.go b/agent/agentssh/x11_test.go
@@ -4,7 +4,9 @@ import (
 	"context"
 	"encoding/hex"
 	"net"
+	"os"
 	"path/filepath"
+	"runtime"
 	"testing"
 
 	"github.com/gliderlabs/ssh"
@@ -23,6 +25,9 @@ import (
 
 func TestServer_X11(t *testing.T) {
 	t.Parallel()
+	if runtime.GOOS != "linux" {
+		t.Skip("X11 forwarding is only supported on Linux")
+	}
 
 	ctx := context.Background()
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
@@ -79,4 +84,10 @@ func TestServer_X11(t *testing.T) {
 	require.NoError(t, err)
 	s.Close()
 	<-done
+
+	// Ensure the Xauthority file was written!
+	home, err := os.UserHomeDir()
+	require.NoError(t, err)
+	_, err = fs.Stat(filepath.Join(home, ".Xauthority"))
+	require.NoError(t, err)
 }