@@ -129,9 +129,10 @@ type Server struct {
129
129
// a lock on mu but protected by closing.
130
130
wg sync.WaitGroup
131
131
132
- Execer agentexec.Execer
133
- logger slog.Logger
134
- srv * ssh.Server
132
+ Execer agentexec.Execer
133
+ logger slog.Logger
134
+ srv * ssh.Server
135
+ x11Forwarder * x11Forwarder
135
136
136
137
config * Config
137
138
@@ -187,6 +188,14 @@ func NewServer(ctx context.Context, logger slog.Logger, prometheusRegistry *prom
187
188
config :config ,
188
189
189
190
metrics :metrics ,
191
+ x11Forwarder :& x11Forwarder {
192
+ logger :logger ,
193
+ x11HandlerErrors :metrics .x11HandlerErrors ,
194
+ fs :fs ,
195
+ displayOffset :* config .X11DisplayOffset ,
196
+ sessions :make (map [* x11Session ]struct {}),
197
+ connections :make (map [net.Conn ]struct {}),
198
+ },
190
199
}
191
200
192
201
srv := & ssh.Server {
@@ -454,7 +463,7 @@ func (s *Server) sessionHandler(session ssh.Session) {
454
463
455
464
x11 ,hasX11 := session .X11 ()
456
465
if hasX11 {
457
- display ,handled := s .x11Handler (ctx ,x11 )
466
+ display ,handled := s .x11Forwarder . x11Handler (ctx ,session )
458
467
if ! handled {
459
468
logger .Error (ctx ,"x11 handler failed" )
460
469
closeCause ("x11 handler failed" )
@@ -1096,6 +1105,9 @@ func (s *Server) Close() error {
1096
1105
1097
1106
s .mu .Unlock ()
1098
1107
1108
+ s .logger .Debug (ctx ,"closing X11 forwarding" )
1109
+ _ = s .x11Forwarder .Close ()
1110
+
1099
1111
s .logger .Debug (ctx ,"waiting for all goroutines to exit" )
1100
1112
s .wg .Wait ()// Wait for all goroutines to exit.
1101
1113