package vnc import ( "bufio" "context" "encoding/binary" "fmt" "net" "sync" ) var ( // DefaultClientHandlers represents default client handlers DefaultClientHandlers = []Handler{ &DefaultClientVersionHandler{}, &DefaultClientSecurityHandler{}, &DefaultClientClientInitHandler{}, &DefaultClientServerInitHandler{}, &DefaultClientMessageHandler{}, } ) // Connect handshake with remote server using underlining net.Conn func Connect(ctx context.Context, c net.Conn, cfg *ClientConfig) (*ClientConn, error) { conn, err := NewClientConn(c, cfg) if err != nil { conn.Close() cfg.ErrorCh <- err return nil, err } if len(cfg.Handlers) == 0 { cfg.Handlers = DefaultClientHandlers } for _, h := range cfg.Handlers { if err := h.Handle(conn); err != nil { conn.Close() cfg.ErrorCh <- err return nil, err } } return conn, nil } var _ Conn = (*ClientConn)(nil) // Config returns connection config func (c *ClientConn) Config() interface{} { return c.cfg } // Wait waiting for connection close func (c *ClientConn) Wait() { <-c.quit } // Conn return underlining net.Conn func (c *ClientConn) Conn() net.Conn { return c.c } // SetProtoVersion sets proto version func (c *ClientConn) SetProtoVersion(pv string) { c.protocol = pv } // SetEncodings write SetEncodings message func (c *ClientConn) SetEncodings(encs []EncodingType) error { msg := &SetEncodings{ EncNum: uint16(len(encs)), Encodings: encs, } return msg.Write(c) } // Flush flushes data to conn func (c *ClientConn) Flush() error { return c.bw.Flush() } // Close closing conn func (c *ClientConn) Close() error { if c.quit != nil { close(c.quit) c.quit = nil } if c.quitCh != nil { close(c.quitCh) } return c.c.Close() } // Read reads data from conn func (c *ClientConn) Read(buf []byte) (int, error) { return c.br.Read(buf) } // Write data to conn must be Flushed func (c *ClientConn) Write(buf []byte) (int, error) { return c.bw.Write(buf) } // ColorMap returns color map func (c *ClientConn) ColorMap() ColorMap { return c.colorMap } // SetColorMap sets color map func (c *ClientConn) SetColorMap(cm ColorMap) { c.colorMap = cm } // DesktopName returns connection desktop name func (c *ClientConn) DesktopName() []byte { return c.desktopName } // PixelFormat returns connection pixel format func (c *ClientConn) PixelFormat() PixelFormat { return c.pixelFormat } // SetDesktopName sets desktop name func (c *ClientConn) SetDesktopName(name []byte) { c.desktopName = name } // SetPixelFormat sets pixel format func (c *ClientConn) SetPixelFormat(pf PixelFormat) error { c.pixelFormat = pf return nil } // Encodings returns client encodings func (c *ClientConn) Encodings() []Encoding { return c.encodings } // Width returns width func (c *ClientConn) Width() uint16 { return c.fbWidth } // Height returns height func (c *ClientConn) Height() uint16 { return c.fbHeight } // Protocol returns protocol func (c *ClientConn) Protocol() string { return c.protocol } // SetWidth sets width of client conn func (c *ClientConn) SetWidth(width uint16) { c.fbWidth = width } // SetHeight sets height of client conn func (c *ClientConn) SetHeight(height uint16) { c.fbHeight = height } // The ClientConn type holds client connection information type ClientConn struct { c net.Conn br *bufio.Reader bw *bufio.Writer cfg *ClientConfig protocol string // If the pixel format uses a color map, then this is the color // map that is used. This should not be modified directly, since // the data comes from the server. // Definition in ยง5 - Representation of Pixel Data. colorMap ColorMap // Name associated with the desktop, sent from the server. desktopName []byte // Encodings supported by the client. This should not be modified // directly. Instead, SetEncodings() should be used. encodings []Encoding // Height of the frame buffer in pixels, sent from the server. fbHeight uint16 // Width of the frame buffer in pixels, sent from the server. fbWidth uint16 // The pixel format associated with the connection. This shouldn't // be modified. If you wish to set a new pixel format, use the // SetPixelFormat method. pixelFormat PixelFormat quitCh chan struct{} quit chan struct{} errorCh chan error } // NewClientConn creates new client conn using config func NewClientConn(c net.Conn, cfg *ClientConfig) (*ClientConn, error) { if len(cfg.Encodings) == 0 { return nil, fmt.Errorf("client can't handle encodings") } return &ClientConn{ c: c, cfg: cfg, br: bufio.NewReader(c), bw: bufio.NewWriter(c), encodings: cfg.Encodings, quitCh: cfg.QuitCh, errorCh: cfg.ErrorCh, pixelFormat: cfg.PixelFormat, quit: make(chan struct{}), }, nil } // DefaultClientMessageHandler represents default client message handler type DefaultClientMessageHandler struct{} // Handle handles server messages. func (*DefaultClientMessageHandler) Handle(c Conn) error { cfg := c.Config().(*ClientConfig) var err error var wg sync.WaitGroup wg.Add(2) defer c.Close() serverMessages := make(map[ServerMessageType]ServerMessage) for _, m := range cfg.Messages { serverMessages[m.Type()] = m } go func() { defer wg.Done() for { select { case msg := <-cfg.ClientMessageCh: if err = msg.Write(c); err != nil { cfg.ErrorCh <- err return } } } }() go func() { defer wg.Done() for { select { default: var messageType ServerMessageType if err = binary.Read(c, binary.BigEndian, &messageType); err != nil { cfg.ErrorCh <- err return } msg, ok := serverMessages[messageType] if !ok { err = fmt.Errorf("unknown message-type: %v", messageType) cfg.ErrorCh <- err return } parsedMsg, err := msg.Read(c) if err != nil { cfg.ErrorCh <- err return } cfg.ServerMessageCh <- parsedMsg } } }() wg.Wait() return nil } // A ClientConfig structure is used to configure a ClientConn. After // one has been passed to initialize a connection, it must not be modified. type ClientConfig struct { Handlers []Handler SecurityHandlers []SecurityHandler Encodings []Encoding PixelFormat PixelFormat ColorMap ColorMap ClientMessageCh chan ClientMessage ServerMessageCh chan ServerMessage Exclusive bool Messages []ServerMessage QuitCh chan struct{} ErrorCh chan error quit chan struct{} }