feat: add support for X11 forwarding (#2)

kylecarbs · web-flow · commit 04bb837133e1 · 2023-04-21T09:02:25.000-05:00
diff --git a/server.go b/server.go
@@ -42,6 +42,7 @@ type Server struct {
 	PasswordHandler               PasswordHandler               // password authentication handler
 	PublicKeyHandler              PublicKeyHandler              // public key authentication handler
 	PtyCallback                   PtyCallback                   // callback for allowing PTY sessions, allows all if nil
+	X11Callback                   X11Callback                   // callback for allowing X11 forwarding, denies all if nil
 	ConnCallback                  ConnCallback                  // optional callback for wrapping net.Conn before handling
 	LocalPortForwardingCallback   LocalPortForwardingCallback   // callback for allowing local port forwarding, denies all if nil
 	ReversePortForwardingCallback ReversePortForwardingCallback // callback for allowing reverse port forwarding, denies all if nil
diff --git a/session.go b/session.go
@@ -69,6 +69,10 @@ type Session interface {
 	// of whether or not a PTY was accepted for this session.
 	Pty() (Pty, <-chan Window, bool)
 
+	// X11 returns X11 forwarding information and a boolean of whether or not X11
+	// forwarding was accepted for this session.
+	X11() (X11, bool)
+
 	// Signals registers a channel to receive signals sent from the client. The
 	// channel must handle signal sends or it will block the SSH request loop.
 	// Registering nil will unregister the channel from signal sends. During the
@@ -106,6 +110,7 @@ func DefaultSessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.Ne
 		conn:              conn,
 		handler:           srv.Handler,
 		ptyCb:             srv.PtyCallback,
+		x11Cb:             srv.X11Callback,
 		sessReqCb:         srv.SessionRequestCallback,
 		subsystemHandlers: srv.SubsystemHandlers,
 		ctx:               ctx,
@@ -122,9 +127,11 @@ type session struct {
 	handled             bool
 	exited              bool
 	pty                 *Pty
+	x11                 *X11
 	winch               chan Window
 	env                 []string
 	ptyCb               PtyCallback
+	x11Cb               X11Callback
 	sessReqCb           SessionRequestCallback
 	rawCmd              string
 	subsystem           string
@@ -226,6 +233,13 @@ func (sess *session) Pty() (Pty, <-chan Window, bool) {
 	return Pty{}, sess.winch, false
 }
 
+func (sess *session) X11() (X11, bool) {
+	if sess.x11 != nil {
+		return *sess.x11, true
+	}
+	return X11{}, false
+}
+
 func (sess *session) Signals(c chan<- Signal) {
 	sess.Lock()
 	defer sess.Unlock()
@@ -353,6 +367,23 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 				close(sess.winch)
 			}()
 			req.Reply(ok, nil)
+		case "x11-req":
+			if sess.handled || sess.x11 != nil {
+				req.Reply(false, nil)
+				continue
+			}
+			x11Req, ok := parseX11Request(req.Payload)
+			if !ok {
+				req.Reply(false, nil)
+				continue
+			}
+			sess.x11 = &x11Req
+			if sess.x11Cb != nil {
+				ok := sess.x11Cb(sess.ctx, x11Req)
+				req.Reply(ok, nil)
+				continue
+			}
+			req.Reply(false, nil)
 		case "window-change":
 			if sess.pty == nil {
 				req.Reply(false, nil)
diff --git a/session_test.go b/session_test.go
@@ -228,11 +228,51 @@ func TestPty(t *testing.T) {
 	<-done
 }
 
+func TestX11(t *testing.T) {
+	t.Parallel()
+	done := make(chan struct{})
+	session, _, cleanup := newTestSession(t, &Server{
+		X11Callback: func(ctx Context, x11 X11) bool {
+			return true
+		},
+		Handler: func(s Session) {
+			x11Req, isX11 := s.X11()
+			if !isX11 {
+				t.Fatalf("expected x11 but none requested")
+			}
+			if !x11Req.SingleConnection {
+				t.Fatalf("expected single connection but got %#v", x11Req.SingleConnection)
+			}
+			close(done)
+		},
+	}, nil)
+	defer cleanup()
+
+	reply, err := session.SendRequest("x11-req", true, gossh.Marshal(X11{
+		SingleConnection: true,
+		AuthProtocol:     "MIT-MAGIC-COOKIE-1",
+		AuthCookie:       "deadbeef",
+		ScreenNumber:     1,
+	}))
+	if err != nil {
+		t.Fatalf("expected nil but got %v", err)
+	}
+	if !reply {
+		t.Fatalf("expected true but got %v", reply)
+	}
+	err = session.Shell()
+	if err != nil {
+		t.Fatalf("expected nil but got %v", err)
+	}
+	session.Close()
+	<-done
+}
+
 func TestPtyResize(t *testing.T) {
 	t.Parallel()
-	winch0 := Window{40, 80}
-	winch1 := Window{80, 160}
-	winch2 := Window{20, 40}
+	winch0 := Window{40, 80, 0, 0}
+	winch1 := Window{80, 160, 0, 0}
+	winch2 := Window{20, 40, 0, 0}
 	winches := make(chan Window)
 	done := make(chan bool)
 	session, _, cleanup := newTestSession(t, &Server{
diff --git a/ssh.go b/ssh.go
@@ -47,6 +47,9 @@ type KeyboardInteractiveHandler func(ctx Context, challenger gossh.KeyboardInter
 // PtyCallback is a hook for allowing PTY sessions.
 type PtyCallback func(ctx Context, pty Pty) bool
 
+// X11Callback is a hook for allowing X11 forwarding.
+type X11Callback func(ctx Context, x11 X11) bool
+
 // SessionRequestCallback is a callback for allowing or denying SSH sessions.
 type SessionRequestCallback func(sess Session, requestType string) bool
 
@@ -106,6 +109,16 @@ type Pty struct {
 	Modes gossh.TerminalModes
 }
 
+// X11 represents a X11 forwarding request.
+type X11 struct {
+	// SingleConnection is whether the X11 connection should be closed after
+	// the first use.
+	SingleConnection bool
+	AuthProtocol     string
+	AuthCookie       string
+	ScreenNumber     uint32
+}
+
 // Serve accepts incoming SSH connections on the listener l, creating a new
 // connection goroutine for each. The connection goroutines read requests and
 // then calls handler to handle sessions. Handler is typically nil, in which
diff --git a/util.go b/util.go
@@ -53,6 +53,12 @@ func parsePtyRequest(payload []byte) (pty Pty, ok bool) {
 	return
 }
 
+// parseX11Request parses an X11 forwarding request.
+// See https://www.rfc-editor.org/rfc/rfc4254#section-6.3
+func parseX11Request(payload []byte) (x11 X11, ok bool) {
+	return x11, ssh.Unmarshal(payload, &x11) == nil
+}
+
 func parseTerminalModes(in []byte) (modes ssh.TerminalModes, ok bool) {
 	// See https://datatracker.ietf.org/doc/html/rfc4254#section-8
 	// 8.  Encoding of Terminal Modes
