Add full support for RFB 3.3, with tests.

client.go

@@ -1,7 +1,9 @@
-// Package vnc implements a VNC client.
-//
-// References:
-//   [PROTOCOL]: http://tools.ietf.org/html/rfc6143
+/*
+Package vnc implements a VNC client.
+
+References:
+  [PROTOCOL]: http://tools.ietf.org/html/rfc6143
+*/
 package vnc
 
 import (
@@ -13,9 +15,11 @@ import (
 	"unicode"
 )
 
+// The ClientConn type holds client connection information.
 type ClientConn struct {
-	c      net.Conn
-	config *ClientConfig
+	c               net.Conn
+	config          *ClientConfig
+	protocolVersion 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
@@ -48,6 +52,9 @@ type ClientConfig struct {
 	// suitable by the server will be used to authenticate.
 	Auth []ClientAuth
 
+	// Password for servers that require authentication.
+	Password string
+
 	// Exclusive determines whether the connection is shared with other
 	// clients. If true, then all other clients connected will be
 	// disconnected when a connection is established to the VNC server.
@@ -66,13 +73,36 @@ type ClientConfig struct {
 	ServerMessages []ServerMessage
 }
 
+func NewClientConfig(p string) *ClientConfig {
+	return &ClientConfig{
+		Auth:     []ClientAuth{&ClientAuthNone{}, &ClientAuthVNC{p}},
+		Password: p,
+	}
+}
+
 func Client(c net.Conn, cfg *ClientConfig) (*ClientConn, error) {
 	conn := &ClientConn{
 		c:      c,
 		config: cfg,
 	}
 
-	if err := conn.handshake(); err != nil {
+	if err := conn.protocolVersionHandshake(); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if err := conn.securityHandshake(); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if err := conn.securityResultHandshake(); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if err := conn.clientInit(); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	if err := conn.serverInit(); err != nil {
 		conn.Close()
 		return nil, err
 	}
@@ -152,15 +182,14 @@ func (c *ClientConn) FramebufferUpdateRequest(incremental bool, x, y, width, hei
 			return err
 		}
 	}
-
 	if _, err := c.c.Write(buf.Bytes()[0:10]); err != nil {
 		return err
 	}
 
 	return nil
 }
 
-// KeyEvent indiciates a key press or release and sends it to the server.
+// KeyEvent indicates a key press or release and sends it to the server.
 // The key is indicated using the X Window System "keysym" value. Use
 // Google to find a reference of these values. To simulate a key press,
 // you must send a key with both a down event, and a non-down event.
@@ -299,119 +328,194 @@ func parseProtocolVersion(pv []byte) (uint, uint, error) {
 	return major, minor, nil
 }
 
-func (c *ClientConn) handshake() error {
+const (
+	// Client ProtocolVersions.
+	PROTO_VERS_UNSUP = "UNSUPPORTED"
+	PROTO_VERS_3_3   = "RFB 003.003\n"
+	PROTO_VERS_3_8   = "RFB 003.008\n"
+)
+
+// protocolVersionHandshake implements §7.1.1 ProtocolVersion Handshake.
+func (c *ClientConn) protocolVersionHandshake() error {
 	var protocolVersion [pvLen]byte
 
-	// 7.1.1, read the ProtocolVersion message sent by the server.
+	// Read the ProtocolVersion message sent by the server.
 	if _, err := io.ReadFull(c.c, protocolVersion[:]); err != nil {
 		return err
 	}
 
-	maxMajor, maxMinor, err := parseProtocolVersion(protocolVersion[:])
+	major, minor, err := parseProtocolVersion(protocolVersion[:])
 	if err != nil {
 		return err
 	}
-	if maxMajor < 3 {
-		return fmt.Errorf("unsupported major version, less than 3: %d", maxMajor)
+	pv := PROTO_VERS_UNSUP
+	if major == 3 {
+		if minor >= 8 {
+			pv = PROTO_VERS_3_8
+		} else if minor >= 3 {
+			pv = PROTO_VERS_3_3
+		}
 	}
-	if maxMinor < 8 {
-		return fmt.Errorf("unsupported minor version, less than 8: %d", maxMinor)
+	if pv == PROTO_VERS_UNSUP {
+		return NewVNCError(fmt.Sprintf("ProtocolVersion handshake failed; unsupported version '%v'", string(protocolVersion[:])))
 	}
+	c.protocolVersion = pv
 
 	// Respond with the version we will support
-	if _, err = c.c.Write([]byte("RFB 003.008\n")); err != nil {
+	if _, err = c.c.Write([]byte(pv)); err != nil {
 		return err
 	}
 
-	// 7.1.2 Security Handshake from server
-	var numSecurityTypes uint8
-	if err = binary.Read(c.c, binary.BigEndian, &numSecurityTypes); err != nil {
-		return err
+	return nil
+}
+
+// securityHandshake implements §7.1.2 Security Handshake.
+func (c *ClientConn) securityHandshake() error {
+
+	switch c.protocolVersion {
+	case PROTO_VERS_3_3:
+		err := c.securityHandshake33()
+		if err != nil {
+			return err
+		}
+	case PROTO_VERS_3_8:
+		err := c.securityHandshake38()
+		if err != nil {
+			return err
+		}
+	default:
+		return NewVNCError(fmt.Sprintf("Security handshake failed; unsupported protocol"))
 	}
+	return nil
+}
 
-	if numSecurityTypes == 0 {
-		return fmt.Errorf("no security types: %s", c.readErrorReason())
+func (c *ClientConn) securityHandshake33() error {
+	var secType uint32
+	if err := binary.Read(c.c, binary.BigEndian, &secType); err != nil {
+		return err
 	}
 
-	securityTypes := make([]uint8, numSecurityTypes)
-	if err = binary.Read(c.c, binary.BigEndian, &securityTypes); err != nil {
+	var auth ClientAuth
+	switch secType {
+	case secTypeInvalid: // Connection failed.
+		reason, err := c.readErrorReason()
+		if err != nil {
+			return err
+		}
+		return NewVNCError(fmt.Sprintf("Security handshake failed; connection failed: %s", reason))
+	case secTypeNone:
+		auth = &ClientAuthNone{}
+	case secTypeVNCAuth:
+		auth = &ClientAuthVNC{c.config.Password}
+	default:
+		return NewVNCError(fmt.Sprintf("Security handshake failed; invalid security type: %v", secType))
+	}
+	if err := auth.Handshake(c.c); err != nil {
 		return err
 	}
 
-	clientSecurityTypes := c.config.Auth
-	if clientSecurityTypes == nil {
-		clientSecurityTypes = []ClientAuth{new(ClientAuthNone)}
+	return nil
+}
+
+func (c *ClientConn) securityHandshake38() error {
+	// Determine server supported security types.
+	var numSecurityTypes uint8
+	if err := binary.Read(c.c, binary.BigEndian, &numSecurityTypes); err != nil {
+		return err
+	}
+	if numSecurityTypes == 0 {
+		reason, err := c.readErrorReason()
+		if err != nil {
+			return err
+		}
+		return NewVNCError(fmt.Sprintf("Security handshake failed; no security types: %v", reason))
+	}
+	securityTypes := make([]uint8, numSecurityTypes)
+	if err := binary.Read(c.c, binary.BigEndian, &securityTypes); err != nil {
+		return err
 	}
 
+	// Choose client security type.
+	// TODO(kward): try "better" security types first.
 	var auth ClientAuth
 FindAuth:
-	for _, curAuth := range clientSecurityTypes {
-		for _, securityType := range securityTypes {
-			if curAuth.SecurityType() == securityType {
-				// We use the first matching supported authentication
-				auth = curAuth
+	for _, securityType := range securityTypes {
+		for _, a := range c.config.Auth {
+			if a.SecurityType() == securityType {
+				// We use the first matching supported authentication.
+				auth = a
 				break FindAuth
 			}
 		}
 	}
-
 	if auth == nil {
-		return fmt.Errorf("no suitable auth schemes found. server supported: %#v", securityTypes)
+		return NewVNCError(fmt.Sprintf("Security handshake failed; no suitable auth schemes found; server supports: %#v", securityTypes))
 	}
-
-	// Respond back with the security type we'll use
-	if err = binary.Write(c.c, binary.BigEndian, auth.SecurityType()); err != nil {
+	if err := binary.Write(c.c, binary.BigEndian, auth.SecurityType()); err != nil {
 		return err
 	}
 
-	if err = auth.Handshake(c.c); err != nil {
+	if err := auth.Handshake(c.c); err != nil {
 		return err
 	}
+	return nil
+}
+
+// securityResultHandshake implements §7.1.3 SecurityResult Handshake.
+func (c *ClientConn) securityResultHandshake() error {
+	if c.protocolVersion == PROTO_VERS_3_3 {
+		// Not required for 3.3.
+		return nil
+	}
 
-	// 7.1.3 SecurityResult Handshake
 	var securityResult uint32
-	if err = binary.Read(c.c, binary.BigEndian, &securityResult); err != nil {
+	if err := binary.Read(c.c, binary.BigEndian, &securityResult); err != nil {
 		return err
 	}
-
 	if securityResult == 1 {
-		return fmt.Errorf("security handshake failed: %s", c.readErrorReason())
+		reason, err := c.readErrorReason()
+		if err != nil {
+			return err
+		}
+		return NewVNCError(fmt.Sprintf("SecurityResult handshake failed: %s", reason))
 	}
+	return nil
+}
 
-	// 7.3.1 ClientInit
-	var sharedFlag uint8 = 1
-	if c.config.Exclusive {
-		sharedFlag = 0
-	}
+// clientInit implements §7.3.1 ClientInit.
+func (c *ClientConn) clientInit() error {
+	var sharedFlag uint8
 
-	if err = binary.Write(c.c, binary.BigEndian, sharedFlag); err != nil {
+	if !c.config.Exclusive {
+		sharedFlag = 1
+	}
+	if err := binary.Write(c.c, binary.BigEndian, sharedFlag); err != nil {
 		return err
 	}
 
-	// 7.3.2 ServerInit
-	if err = binary.Read(c.c, binary.BigEndian, &c.FrameBufferWidth); err != nil {
+	return nil
+}
+
+// serverInit implements §7.3.2 ServerInit.
+func (c *ClientConn) serverInit() error {
+	if err := binary.Read(c.c, binary.BigEndian, &c.FrameBufferWidth); err != nil {
 		return err
 	}
-
-	if err = binary.Read(c.c, binary.BigEndian, &c.FrameBufferHeight); err != nil {
+	if err := binary.Read(c.c, binary.BigEndian, &c.FrameBufferHeight); err != nil {
 		return err
 	}
-
-	// Read the pixel format
-	if err = readPixelFormat(c.c, &c.PixelFormat); err != nil {
+	if err := readPixelFormat(c.c, &c.PixelFormat); err != nil {
 		return err
 	}
 
 	var nameLength uint32
-	if err = binary.Read(c.c, binary.BigEndian, &nameLength); err != nil {
+	if err := binary.Read(c.c, binary.BigEndian, &nameLength); err != nil {
 		return err
 	}
-
 	nameBytes := make([]uint8, nameLength)
-	if err = binary.Read(c.c, binary.BigEndian, &nameBytes); err != nil {
+	if err := binary.Read(c.c, binary.BigEndian, &nameBytes); err != nil {
 		return err
 	}
-
 	c.DesktopName = string(nameBytes)
 
 	return nil
@@ -467,16 +571,31 @@ func (c *ClientConn) mainLoop() {
 	}
 }
 
-func (c *ClientConn) readErrorReason() string {
+// TODO(kward): need a context for timeout
+func (c *ClientConn) readErrorReason() (string, error) {
 	var reasonLen uint32
 	if err := binary.Read(c.c, binary.BigEndian, &reasonLen); err != nil {
-		return "<error>"
+		return "", err
 	}
 
 	reason := make([]uint8, reasonLen)
 	if err := binary.Read(c.c, binary.BigEndian, &reason); err != nil {
-		return "<error>"
+		return "", err
 	}
 
-	return string(reason)
+	return string(reason), nil
+}
+
+// VNCError implements error interface.
+type VNCError struct {
+	s string
+}
+
+// NewVNCError returns a custom VNCError error.
+func NewVNCError(s string) error {
+	return &VNCError{s}
+}
+
+func (e VNCError) Error() string {
+	return e.s
 }