github.com/amitbet/vnc2video@v0.0.0-20190616012314-9d50b9dab1d9/server.go (about) 1 package vnc2video 2 3 import ( 4 "bufio" 5 "context" 6 "encoding/binary" 7 "fmt" 8 "net" 9 "sync" 10 ) 11 12 var _ Conn = (*ServerConn)(nil) 13 14 // Config returns config for server conn 15 func (c *ServerConn) Config() interface{} { 16 return c.cfg 17 } 18 func (c *ServerConn) GetEncInstance(typ EncodingType) Encoding { 19 for _, enc := range c.encodings { 20 if enc.Type() == typ { 21 return enc 22 } 23 } 24 return nil 25 } 26 27 // Conn returns underlining server net.Conn 28 func (c *ServerConn) Conn() net.Conn { 29 return c.c 30 } 31 32 // Wait waits connection to close 33 func (c *ServerConn) Wait() { 34 <-c.quit 35 } 36 37 // SetEncodings ??? sets server connection encodings 38 func (c *ServerConn) SetEncodings(encs []EncodingType) error { 39 encodings := make(map[EncodingType]Encoding) 40 for _, enc := range c.cfg.Encodings { 41 encodings[enc.Type()] = enc 42 } 43 for _, encType := range encs { 44 if enc, ok := encodings[encType]; ok { 45 c.encodings = append(c.encodings, enc) 46 } 47 } 48 return nil 49 } 50 51 // SetProtoVersion ??? sets proto version 52 func (c *ServerConn) SetProtoVersion(pv string) { 53 c.protocol = pv 54 } 55 56 // Flush buffered data to server conn 57 func (c *ServerConn) Flush() error { 58 return c.bw.Flush() 59 } 60 61 // Close closing server conn 62 func (c *ServerConn) Close() error { 63 if c.quit != nil { 64 close(c.quit) 65 c.quit = nil 66 } 67 return c.c.Close() 68 } 69 70 // Read reads data from net.Conn 71 func (c *ServerConn) Read(buf []byte) (int, error) { 72 return c.br.Read(buf) 73 } 74 75 // Write writes data to net.Conn, must be Flashed 76 func (c *ServerConn) Write(buf []byte) (int, error) { 77 return c.bw.Write(buf) 78 } 79 80 // ColorMap returns server connection color map 81 func (c *ServerConn) ColorMap() ColorMap { 82 return c.colorMap 83 } 84 85 // SetColorMap sets connection color map 86 func (c *ServerConn) SetColorMap(cm ColorMap) { 87 c.colorMap = cm 88 } 89 90 // DesktopName returns connection desktop name 91 func (c *ServerConn) DesktopName() []byte { 92 return c.desktopName 93 } 94 95 // PixelFormat return connection pixel format 96 func (c *ServerConn) PixelFormat() PixelFormat { 97 return c.pixelFormat 98 } 99 100 // SetDesktopName sets connection desktop name 101 func (c *ServerConn) SetDesktopName(name []byte) { 102 c.desktopName = name 103 } 104 105 // SetPixelFormat sets pixel format for server conn 106 func (c *ServerConn) SetPixelFormat(pf PixelFormat) error { 107 c.pixelFormat = pf 108 return nil 109 } 110 111 // Encodings returns connection encodings 112 func (c *ServerConn) Encodings() []Encoding { 113 return c.encodings 114 } 115 116 // Width returns framebuffer width 117 func (c *ServerConn) Width() uint16 { 118 return c.fbWidth 119 } 120 121 // Height returns framebuffer height 122 func (c *ServerConn) Height() uint16 { 123 return c.fbHeight 124 } 125 126 // Protocol returns protocol 127 func (c *ServerConn) Protocol() string { 128 return c.protocol 129 } 130 131 // SecurityHandler returns security handler 132 func (c *ServerConn) SecurityHandler() SecurityHandler { 133 return c.securityHandler 134 } 135 136 // SetSecurityHandler sets security handler 137 func (c *ServerConn) SetSecurityHandler(sechandler SecurityHandler) error { 138 c.securityHandler = sechandler 139 return nil 140 } 141 142 // SetWidth sets framebuffer width 143 func (c *ServerConn) SetWidth(w uint16) { 144 // TODO send desktopsize pseudo encoding 145 c.fbWidth = w 146 } 147 148 // SetHeight sets framebuffer height 149 func (c *ServerConn) SetHeight(h uint16) { 150 // TODO send desktopsize pseudo encoding 151 c.fbHeight = h 152 } 153 154 // ServerConn underlining server conn 155 type ServerConn struct { 156 c net.Conn 157 cfg *ServerConfig 158 br *bufio.Reader 159 bw *bufio.Writer 160 protocol string 161 // If the pixel format uses a color map, then this is the color 162 // map that is used. This should not be modified directly, since 163 // the data comes from the server. 164 // Definition in ยง5 - Representation of Pixel Data. 165 colorMap ColorMap 166 167 // Name associated with the desktop, sent from the server. 168 desktopName []byte 169 170 // Encodings supported by the client. This should not be modified 171 // directly. Instead, SetEncodings() should be used. 172 encodings []Encoding 173 174 securityHandler SecurityHandler 175 176 // Height of the frame buffer in pixels, sent to the client. 177 fbHeight uint16 178 179 // Width of the frame buffer in pixels, sent to the client. 180 fbWidth uint16 181 182 // The pixel format associated with the connection. This shouldn't 183 // be modified. If you wish to set a new pixel format, use the 184 // SetPixelFormat method. 185 pixelFormat PixelFormat 186 187 quit chan struct{} 188 } 189 190 var ( 191 // DefaultServerHandlers uses default handlers for hanshake 192 DefaultServerHandlers = []Handler{ 193 &DefaultServerVersionHandler{}, 194 &DefaultServerSecurityHandler{}, 195 &DefaultServerClientInitHandler{}, 196 &DefaultServerServerInitHandler{}, 197 &DefaultServerMessageHandler{}, 198 } 199 ) 200 201 // ServerConfig config struct 202 type ServerConfig struct { 203 Handlers []Handler 204 SecurityHandlers []SecurityHandler 205 Encodings []Encoding 206 PixelFormat PixelFormat 207 ColorMap ColorMap 208 ClientMessageCh chan ClientMessage 209 ServerMessageCh chan ServerMessage 210 Messages []ClientMessage 211 DesktopName []byte 212 Height uint16 213 Width uint16 214 ErrorCh chan error 215 } 216 217 // NewServerConn returns new Server connection fron net.Conn 218 func NewServerConn(c net.Conn, cfg *ServerConfig) (*ServerConn, error) { 219 return &ServerConn{ 220 c: c, 221 br: bufio.NewReader(c), 222 bw: bufio.NewWriter(c), 223 cfg: cfg, 224 desktopName: cfg.DesktopName, 225 encodings: cfg.Encodings, 226 pixelFormat: cfg.PixelFormat, 227 fbWidth: cfg.Width, 228 fbHeight: cfg.Height, 229 quit: make(chan struct{}), 230 }, nil 231 } 232 233 // Serve serves requests from net.Listener using ServerConfig 234 func Serve(ctx context.Context, ln net.Listener, cfg *ServerConfig) error { 235 for { 236 237 c, err := ln.Accept() 238 if err != nil { 239 continue 240 } 241 242 conn, err := NewServerConn(c, cfg) 243 if err != nil { 244 cfg.ErrorCh <- err 245 continue 246 } 247 248 if len(cfg.Handlers) == 0 { 249 cfg.Handlers = DefaultServerHandlers 250 } 251 252 handlerLoop: 253 for _, h := range cfg.Handlers { 254 if err := h.Handle(conn); err != nil { 255 if cfg.ErrorCh != nil { 256 cfg.ErrorCh <- err 257 } 258 conn.Close() 259 break handlerLoop 260 } 261 } 262 } 263 } 264 265 // DefaultServerMessageHandler default package handler 266 type DefaultServerMessageHandler struct{} 267 268 // Handle handles messages from clients 269 func (*DefaultServerMessageHandler) Handle(c Conn) error { 270 cfg := c.Config().(*ServerConfig) 271 var err error 272 var wg sync.WaitGroup 273 274 defer c.Close() 275 clientMessages := make(map[ClientMessageType]ClientMessage) 276 for _, m := range cfg.Messages { 277 clientMessages[m.Type()] = m 278 } 279 wg.Add(2) 280 281 quit := make(chan struct{}) 282 283 // server 284 go func() { 285 defer wg.Done() 286 for { 287 select { 288 case <-quit: 289 return 290 case msg := <-cfg.ServerMessageCh: 291 if err = msg.Write(c); err != nil { 292 cfg.ErrorCh <- err 293 if quit != nil { 294 close(quit) 295 quit = nil 296 } 297 return 298 } 299 } 300 } 301 }() 302 303 // client 304 go func() { 305 defer wg.Done() 306 for { 307 select { 308 case <-quit: 309 return 310 default: 311 var messageType ClientMessageType 312 if err := binary.Read(c, binary.BigEndian, &messageType); err != nil { 313 cfg.ErrorCh <- err 314 if quit != nil { 315 close(quit) 316 quit = nil 317 } 318 return 319 } 320 msg, ok := clientMessages[messageType] 321 if !ok { 322 cfg.ErrorCh <- fmt.Errorf("unsupported message-type: %v", messageType) 323 close(quit) 324 return 325 } 326 parsedMsg, err := msg.Read(c) 327 if err != nil { 328 cfg.ErrorCh <- err 329 if quit != nil { 330 close(quit) 331 quit = nil 332 } 333 return 334 } 335 cfg.ClientMessageCh <- parsedMsg 336 } 337 } 338 }() 339 340 wg.Wait() 341 return nil 342 }