Merge pull request #10 from kward/fix-protovers-parsing

Fix ProtocolVersion message parsing to support OS X 10.10.3.

Author: mitchellh

client.go

@@ -279,24 +279,41 @@ func (c *ClientConn) SetPixelFormat(format *PixelFormat) error {
 	return nil
 }
 
+const pvLen = 12 // ProtocolVersion message length.
+
+func parseProtocolVersion(pv []byte) (uint, uint, error) {
+	var major, minor uint
+
+	if len(pv) < pvLen {
+		return 0, 0, fmt.Errorf("ProtocolVersion message too short (%v < %v)", len(pv), pvLen)
+	}
+
+	l, err := fmt.Sscanf(string(pv), "RFB %d.%d\n", &major, &minor)
+	if l != 2 {
+		return 0, 0, fmt.Errorf("error parsing ProtocolVersion.")
+	}
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return major, minor, nil
+}
+
 func (c *ClientConn) handshake() error {
-	var protocolVersion [12]byte
+	var protocolVersion [pvLen]byte
 
 	// 7.1.1, read the ProtocolVersion message sent by the server.
 	if _, err := io.ReadFull(c.c, protocolVersion[:]); err != nil {
 		return err
 	}
 
-	var maxMajor, maxMinor uint8
-	_, err := fmt.Sscanf(string(protocolVersion[:]), "RFB %d.%d\n", &maxMajor, &maxMinor)
+	maxMajor, maxMinor, err := parseProtocolVersion(protocolVersion[:])
 	if err != nil {
 		return err
 	}
-
 	if maxMajor < 3 {
 		return fmt.Errorf("unsupported major version, less than 3: %d", maxMajor)
 	}
-
 	if maxMinor < 8 {
 		return fmt.Errorf("unsupported minor version, less than 8: %d", maxMinor)
 	}

client_test.go

@@ -60,3 +60,36 @@ func TestClient_LowMinorVersion(t *testing.T) {
 		t.Fatalf("unexpected error: %s", err)
 	}
 }
+
+func TestParseProtocolVersion(t *testing.T) {
+	tests := []struct {
+		proto        []byte
+		major, minor uint
+		isErr        bool
+	}{
+		// Valid ProtocolVersion messages.
+		{[]byte{82, 70, 66, 32, 48, 48, 51, 46, 48, 48, 56, 10}, 3, 8, false},   // RFB 003.008\n
+		{[]byte{82, 70, 66, 32, 48, 48, 51, 46, 56, 56, 57, 10}, 3, 889, false}, // RFB 003.889\n -- OS X 10.10.3
+		{[]byte{82, 70, 66, 32, 48, 48, 48, 46, 48, 48, 48, 10}, 0, 0, false},   // RFB 000.0000\n
+		// Invalid messages.
+		{[]byte{82, 70, 66, 32, 51, 46, 56, 10}, 0, 0, true}, // RFB 3.8\n -- too short; not zero padded
+		{[]byte{82, 70, 66, 10}, 0, 0, true},                 // RFB\n -- too short
+		{[]byte{}, 0, 0, true},                               // (empty) -- too short
+	}
+
+	for _, tt := range tests {
+		major, minor, err := parseProtocolVersion(tt.proto)
+		if err != nil && !tt.isErr {
+			t.Fatalf("parseProtocolVersion(%v) unexpected error %v", tt.proto, err)
+		}
+		if err == nil && tt.isErr {
+			t.Fatalf("parseProtocolVersion(%v) expected error", tt.proto)
+		}
+		if major != tt.major {
+			t.Errorf("parseProtocolVersion(%v) major = %v, want %v", tt.proto, major, tt.major)
+		}
+		if major != tt.major {
+			t.Errorf("parseProtocolVersion(%v) minor = %v, want %v", tt.proto, minor, tt.minor)
+		}
+	}
+}