From 4ceeeebeb2ce72e8e9b76b4e658a91f9346b98e3 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@selfip.ru>
Date: Tue, 20 Jun 2017 00:55:32 +0300
Subject: [PATCH] minor fixes and proper auth

Signed-off-by: Vasiliy Tolstov <v.tolstov@selfip.ru>
---
 client.go   |  8 --------
 encoding.go |  1 +
 handlers.go | 17 +++++++++++++++--
 server.go   | 27 +++++++--------------------
 4 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/client.go b/client.go
index a07764c..ddad1c7 100644
--- a/client.go
+++ b/client.go
@@ -66,13 +66,7 @@ func (c *ClientConn) SetEncodings(encs []EncodingType) error {
 	return msg.Write(c)
 }
 
-func (c *ClientConn) UnreadByte() error {
-	return c.br.UnreadByte()
-}
-
 func (c *ClientConn) Flush() error {
-	//	c.m.Lock()
-	//	defer c.m.Unlock()
 	return c.bw.Flush()
 }
 
@@ -85,8 +79,6 @@ func (c *ClientConn) Read(buf []byte) (int, error) {
 }
 
 func (c *ClientConn) Write(buf []byte) (int, error) {
-	//	c.m.Lock()
-	//	defer c.m.Unlock()
 	return c.bw.Write(buf)
 }
 
diff --git a/encoding.go b/encoding.go
index 21acf60..ff98b91 100644
--- a/encoding.go
+++ b/encoding.go
@@ -211,6 +211,7 @@ func (enc *TightPngEncoding) Write(c Conn, rect *Rectangle) error {
 		buf.Reset()
 		defer bPool.Put(buf)
 		pngEnc := &png.Encoder{CompressionLevel: png.BestSpeed}
+		//pngEnc := &png.Encoder{CompressionLevel: png.NoCompression}
 		if err := pngEnc.Encode(buf, enc.Image); err != nil {
 			return err
 		}
diff --git a/handlers.go b/handlers.go
index 8fa703f..c9ac065 100644
--- a/handlers.go
+++ b/handlers.go
@@ -68,8 +68,7 @@ func ClientVersionHandler(cfg *ClientConfig, c Conn) error {
 	if pv == ProtoVersionUnknown {
 		return fmt.Errorf("ProtocolVersion handshake failed; unsupported version '%v'", string(version[:]))
 	}
-
-	c.SetProtoVersion(pv)
+	c.SetProtoVersion(string(version[:]))
 
 	if err := binary.Write(c, binary.BigEndian, []byte(pv)); err != nil {
 		return err
@@ -120,6 +119,15 @@ func ClientSecurityHandler(cfg *ClientConfig, c Conn) error {
 		return err
 	}
 
+	var secType SecurityHandler
+	for _, st := range cfg.SecurityHandlers {
+		for _, sc := range secTypes {
+			if st.Type() == sc {
+				secType = st
+			}
+		}
+	}
+
 	if err := binary.Write(c, binary.BigEndian, cfg.SecurityHandlers[0].Type()); err != nil {
 		return err
 	}
@@ -128,6 +136,11 @@ func ClientSecurityHandler(cfg *ClientConfig, c Conn) error {
 		return err
 	}
 
+	err := secType.Auth(c)
+	if err != nil {
+		return err
+	}
+
 	var authCode uint32
 	if err := binary.Read(c, binary.BigEndian, &authCode); err != nil {
 		return err
diff --git a/server.go b/server.go
index 0ee18fb..f664902 100644
--- a/server.go
+++ b/server.go
@@ -27,10 +27,6 @@ type ServerInit struct {
 
 var _ Conn = (*ServerConn)(nil)
 
-func (c *ServerConn) UnreadByte() error {
-	return c.br.UnreadByte()
-}
-
 func (c *ServerConn) Conn() net.Conn {
 	return c.c
 }
@@ -53,8 +49,6 @@ func (c *ServerConn) SetProtoVersion(pv string) {
 }
 
 func (c *ServerConn) Flush() error {
-	//	c.m.Lock()
-	//	defer c.m.Unlock()
 	return c.bw.Flush()
 }
 
@@ -62,22 +56,11 @@ func (c *ServerConn) Close() error {
 	return c.c.Close()
 }
 
-/*
-func (c *ServerConn) Input() chan *ServerMessage {
-	return c.cfg.ServerMessageCh
-}
-
-func (c *ServerConn) Output() chan *ClientMessage {
-	return c.cfg.ClientMessageCh
-}
-*/
 func (c *ServerConn) Read(buf []byte) (int, error) {
 	return c.br.Read(buf)
 }
 
 func (c *ServerConn) Write(buf []byte) (int, error) {
-	//	c.m.Lock()
-	//	defer c.m.Unlock()
 	return c.bw.Write(buf)
 }
 
@@ -232,6 +215,7 @@ type ServerConfig struct {
 	ColorMap          *ColorMap
 	ClientMessageCh   chan ClientMessage
 	ServerMessageCh   chan ServerMessage
+	ErrorCh           chan error
 	ClientMessages    []ClientMessage
 	DesktopName       []byte
 	Height            uint16
@@ -402,7 +386,7 @@ func (msg *ServerCutText) Write(c Conn) error {
 	if err := binary.Write(c, binary.BigEndian, msg.Text); err != nil {
 		return err
 	}
-	return nil
+	return c.Flush()
 }
 
 type Bell struct{}
@@ -416,7 +400,10 @@ func (*Bell) Read(c Conn) (ServerMessage, error) {
 }
 
 func (msg *Bell) Write(c Conn) error {
-	return binary.Write(c, binary.BigEndian, msg.Type())
+	if err := binary.Write(c, binary.BigEndian, msg.Type()); err != nil {
+		return err
+	}
+	return c.Flush()
 }
 
 type SetColorMapEntries struct {
@@ -486,5 +473,5 @@ func (msg *SetColorMapEntries) Write(c Conn) error {
 		}
 	}
 
-	return nil
+	return c.Flush()
 }