feat: add support for X11 forwarding

kylecarbs · kylecarbs · commit 5d1ffcdbf799 · 2023-04-19T16:52:38.000Z
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -84,6 +84,8 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 	forwardHandler := &ssh.ForwardedTCPHandler{}
 	unixForwardHandler := &forwardedUnixHandler{log: logger}
 
+	x11Handler := &x11Handler{log: logger}
+
 	s := &Server{
 		listeners: make(map[net.Listener]struct{}),
 		conns:     make(map[net.Conn]struct{}),
@@ -96,6 +98,7 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 			"direct-tcpip":                   ssh.DirectTCPIPHandler,
 			"direct-streamlocal@openssh.com": directStreamLocalHandler,
 			"session":                        ssh.DefaultSessionHandler,
+			"x11":                            x11Handler.channel,
 		},
 		ConnectionFailedCallback: func(_ net.Conn, err error) {
 			s.logger.Info(ctx, "ssh connection ended", slog.Error(err))
@@ -124,6 +127,7 @@ func NewServer(ctx context.Context, logger slog.Logger, maxTimeout time.Duration
 			"cancel-tcpip-forward":                   forwardHandler.HandleSSHRequest,
 			"streamlocal-forward@openssh.com":        unixForwardHandler.HandleSSHRequest,
 			"cancel-streamlocal-forward@openssh.com": unixForwardHandler.HandleSSHRequest,
+			"x11-req":                                x11Handler.request,
 		},
 		ServerConfigCallback: func(ctx ssh.Context) *gossh.ServerConfig {
 			return &gossh.ServerConfig{
diff --git a/agent/agentssh/x11.go b/agent/agentssh/x11.go
@@ -0,0 +1,138 @@
+package agentssh
+
+import (
+	"encoding/binary"
+	"encoding/hex"
+	"fmt"
+	"net"
+	"os"
+	"path/filepath"
+	"sync"
+
+	"github.com/gliderlabs/ssh"
+	gossh "golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+)
+
+// x11ReqPayload describes the extra data sent in a x11-req request.
+// https://www.rfc-editor.org/rfc/rfc4254#section-6.3
+type x11ReqPayload struct {
+	SingleConnection bool
+	AuthProtocol     string
+	AuthCookie       string
+	ScreenNumber     uint32
+}
+
+type x11Handler struct {
+	mutex sync.Mutex
+	log   slog.Logger
+
+	forwardingEnabled bool
+	screenNumber      uint32
+}
+
+func (h *x11Handler) request(ctx ssh.Context, _ *ssh.Server, req *gossh.Request) (bool, []byte) {
+	var reqPayload x11ReqPayload
+	err := gossh.Unmarshal(req.Payload, &reqPayload)
+	if err != nil {
+		h.log.Warn(ctx, "parse x11-req request payload from client", slog.Error(err))
+		return false, nil
+	}
+
+	err = addXauthEntry(fmt.Sprintf("localhost:%d.0", reqPayload.ScreenNumber), reqPayload.AuthProtocol, reqPayload.AuthCookie)
+	if err != nil {
+		h.log.Warn(ctx, "failed to add Xauthority entry", slog.Error(err))
+		return false, nil
+	}
+
+	h.mutex.Lock()
+	h.forwardingEnabled = true
+	h.screenNumber = reqPayload.ScreenNumber
+	h.mutex.Unlock()
+	return true, nil
+}
+
+func (h *x11Handler) channel(_ *ssh.Server, _ *gossh.ServerConn, newChan gossh.NewChannel, ctx ssh.Context) {
+	h.mutex.Lock()
+	enabled := h.forwardingEnabled
+	h.mutex.Unlock()
+	if !enabled {
+		_ = newChan.Reject(gossh.Prohibited, "X11 forwarding not requested")
+		return
+	}
+
+	x11Conn, err := net.Dial("unix", fmt.Sprintf(filepath.Join(os.TempDir(), ".X11-unix", "X%d"), h.screenNumber))
+	if err != nil {
+		_ = newChan.Reject(gossh.ConnectionFailed, "Failed to connect to local X11 server")
+		return
+	}
+
+	channel, requests, err := newChan.Accept()
+	if err != nil {
+		_ = x11Conn.Close()
+		return
+	}
+	go gossh.DiscardRequests(requests)
+	Bicopy(ctx, channel, x11Conn)
+}
+
+func addXauthEntry(display string, authProtocol string, authCookie string) error {
+	// Get the Xauthority file path
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return xerrors.Errorf("failed to get user home directory: %w", err)
+	}
+
+	xauthPath := filepath.Join(homeDir, ".Xauthority")
+
+	// Open or create the Xauthority file
+	file, err := os.OpenFile(xauthPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
+	if err != nil {
+		return xerrors.Errorf("failed to open Xauthority file: %w", err)
+	}
+	defer file.Close()
+
+	// Convert the authCookie from hex string to byte slice
+	authCookieBytes, err := hex.DecodeString(authCookie)
+	if err != nil {
+		return xerrors.Errorf("failed to decode auth cookie: %w", err)
+	}
+
+	// Write Xauthority entry
+	family := uint16(0x0100) // FamilyLocal
+	err = binary.Write(file, binary.BigEndian, family)
+	if err != nil {
+		return xerrors.Errorf("failed to write family: %w", err)
+	}
+
+	err = binary.Write(file, binary.BigEndian, uint16(len(display)))
+	if err != nil {
+		return xerrors.Errorf("failed to write display length: %w", err)
+	}
+	_, err = file.WriteString(display)
+	if err != nil {
+		return xerrors.Errorf("failed to write display: %w", err)
+	}
+
+	err = binary.Write(file, binary.BigEndian, uint16(len(authProtocol)))
+	if err != nil {
+		return xerrors.Errorf("failed to write auth protocol length: %w", err)
+	}
+	_, err = file.WriteString(authProtocol)
+	if err != nil {
+		return xerrors.Errorf("failed to write auth protocol: %w", err)
+	}
+
+	err = binary.Write(file, binary.BigEndian, uint16(len(authCookieBytes)))
+	if err != nil {
+		return xerrors.Errorf("failed to write auth cookie length: %w", err)
+	}
+	_, err = file.Write(authCookieBytes)
+	if err != nil {
+		return xerrors.Errorf("failed to write auth cookie: %w", err)
+	}
+
+	return nil
+}
