main
1// Package proxy provides VNC proxy functionality for multiplexing clients.
2package proxy
3
4import (
5 "context"
6 "fmt"
7 "log/slog"
8 "sync"
9
10 "goVNC/pkg/rfb"
11 "goVNC/pkg/transport"
12)
13
14// SessionMode determines how clients interact with the VNC session.
15type SessionMode int
16
17const (
18 // SharedMode - all clients share the same VNC session view and can send input.
19 SharedMode SessionMode = iota
20
21 // IsolatedMode - input is queued and processed sequentially.
22 IsolatedMode
23)
24
25// Multiplexer manages multiple WebSocket clients connected to a single VNC session.
26type Multiplexer struct {
27 vnc *rfb.Client
28 mode SessionMode
29
30 clients map[string]*ClientConn
31 clientMu sync.RWMutex
32
33 // For broadcasting server messages to all clients
34 broadcast chan []byte
35
36 closed bool
37 closeMu sync.RWMutex
38}
39
40// ClientConn represents a connected client.
41type ClientConn struct {
42 ID string
43 Transport transport.Transport
44 send chan []byte
45 done chan struct{}
46}
47
48// NewMultiplexer creates a new multiplexer for the given VNC client.
49func NewMultiplexer(vnc *rfb.Client, mode SessionMode) *Multiplexer {
50 return &Multiplexer{
51 vnc: vnc,
52 mode: mode,
53 clients: make(map[string]*ClientConn),
54 broadcast: make(chan []byte, 100),
55 }
56}
57
58// AddClient registers a new client connection.
59func (m *Multiplexer) AddClient(id string, t transport.Transport) *ClientConn {
60 client := &ClientConn{
61 ID: id,
62 Transport: t,
63 send: make(chan []byte, 100),
64 done: make(chan struct{}),
65 }
66
67 m.clientMu.Lock()
68 m.clients[id] = client
69 m.clientMu.Unlock()
70
71 slog.Info("client connected", slog.String("id", id))
72 return client
73}
74
75// RemoveClient unregisters a client.
76func (m *Multiplexer) RemoveClient(id string) {
77 m.clientMu.Lock()
78 client, ok := m.clients[id]
79 if ok {
80 delete(m.clients, id)
81 close(client.done)
82 }
83 m.clientMu.Unlock()
84
85 if ok {
86 slog.Info("client disconnected", slog.String("id", id))
87 }
88}
89
90// GetClient returns a client by ID.
91func (m *Multiplexer) GetClient(id string) (*ClientConn, bool) {
92 m.clientMu.RLock()
93 defer m.clientMu.RUnlock()
94 client, ok := m.clients[id]
95 return client, ok
96}
97
98// ClientCount returns the number of connected clients.
99func (m *Multiplexer) ClientCount() int {
100 m.clientMu.RLock()
101 defer m.clientMu.RUnlock()
102 return len(m.clients)
103}
104
105// Broadcast sends a message to all connected clients.
106func (m *Multiplexer) Broadcast(data []byte) {
107 m.clientMu.RLock()
108 defer m.clientMu.RUnlock()
109
110 for _, client := range m.clients {
111 select {
112 case client.send <- data:
113 default:
114 slog.Warn("client send buffer full", slog.String("id", client.ID))
115 }
116 }
117}
118
119// HandleClientInput processes input from a client and forwards to VNC.
120func (m *Multiplexer) HandleClientInput(ctx context.Context, clientID string, data []byte) error {
121 if len(data) == 0 {
122 return nil
123 }
124
125 // Parse the message type
126 msgType := data[0]
127
128 switch msgType {
129 case rfb.MsgTypeKeyEvent:
130 // Forward key events to VNC
131 return m.vnc.Transport().Write(ctx, data)
132
133 case rfb.MsgTypePointerEvent:
134 // Forward pointer events to VNC
135 return m.vnc.Transport().Write(ctx, data)
136
137 case rfb.MsgTypeClientCutText:
138 // Forward clipboard to VNC
139 return m.vnc.Transport().Write(ctx, data)
140
141 case rfb.MsgTypeFramebufferUpdateRequest:
142 // Don't forward FBU requests - the main client handles these
143 return nil
144
145 case rfb.MsgTypeSetEncodings:
146 // Don't forward encoding changes
147 return nil
148
149 case rfb.MsgTypeSetPixelFormat:
150 // Don't forward pixel format changes
151 return nil
152
153 default:
154 slog.Debug("unknown client message type",
155 slog.String("client", clientID),
156 slog.Int("type", int(msgType)))
157 }
158
159 return nil
160}
161
162// RunClientWriter writes messages from the send channel to a client.
163func (m *Multiplexer) RunClientWriter(ctx context.Context, client *ClientConn) error {
164 for {
165 select {
166 case <-ctx.Done():
167 return ctx.Err()
168 case <-client.done:
169 return nil
170 case data := <-client.send:
171 if err := client.Transport.Write(ctx, data); err != nil {
172 return fmt.Errorf("writing to client: %w", err)
173 }
174 }
175 }
176}
177
178// RunClientReader reads messages from a client and processes them.
179func (m *Multiplexer) RunClientReader(ctx context.Context, client *ClientConn) error {
180 for {
181 select {
182 case <-ctx.Done():
183 return ctx.Err()
184 case <-client.done:
185 return nil
186 default:
187 }
188
189 data, err := client.Transport.Read(ctx)
190 if err != nil {
191 return fmt.Errorf("reading from client: %w", err)
192 }
193
194 if err := m.HandleClientInput(ctx, client.ID, data); err != nil {
195 slog.Warn("handling client input",
196 slog.String("client", client.ID),
197 slog.String("error", err.Error()))
198 }
199 }
200}
201
202// Close shuts down the multiplexer.
203func (m *Multiplexer) Close() error {
204 m.closeMu.Lock()
205 defer m.closeMu.Unlock()
206
207 if m.closed {
208 return nil
209 }
210 m.closed = true
211
212 // Close all client connections
213 m.clientMu.Lock()
214 for id, client := range m.clients {
215 close(client.done)
216 client.Transport.Close()
217 delete(m.clients, id)
218 }
219 m.clientMu.Unlock()
220
221 close(m.broadcast)
222 return nil
223}
224
225// VNCClient returns the underlying VNC client.
226func (m *Multiplexer) VNCClient() *rfb.Client {
227 return m.vnc
228}