From 88b99c68a5f2c40395bfd97584114fa0d1ed6a3f Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@selfip.ru>
Date: Wed, 5 Jul 2017 16:08:27 +0300
Subject: [PATCH] add security handler info to connection

Signed-off-by: Vasiliy Tolstov <v.tolstov@selfip.ru>
---
 client.go   | 13 +++++++++++++
 conn.go     |  2 ++
 handlers.go |  3 ++-
 server.go   | 13 +++++++++++++
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/client.go b/client.go
index 277be70..5a5aa02 100644
--- a/client.go
+++ b/client.go
@@ -165,6 +165,17 @@ func (c *ClientConn) SetHeight(height uint16) {
 	c.fbHeight = height
 }
 
+// SecurityHandler returns security handler
+func (c *ClientConn) SecurityHandler() SecurityHandler {
+	return c.securityHandler
+}
+
+// SetSecurityHandler sets security handler
+func (c *ClientConn) SetSecurityHandler(sechandler SecurityHandler) error {
+	c.securityHandler = sechandler
+	return nil
+}
+
 // The ClientConn type holds client connection information
 type ClientConn struct {
 	c        net.Conn
@@ -186,6 +197,8 @@ type ClientConn struct {
 	// directly. Instead, SetEncodings() should be used.
 	encodings []Encoding
 
+	securityHandler SecurityHandler
+
 	// Height of the frame buffer in pixels, sent from the server.
 	fbHeight uint16
 
diff --git a/conn.go b/conn.go
index 853dc08..7427cbf 100644
--- a/conn.go
+++ b/conn.go
@@ -26,4 +26,6 @@ type Conn interface {
 	Flush() error
 	Wait()
 	SetProtoVersion(string)
+	SetSecurityHandler(SecurityHandler) error
+	SecurityHandler() SecurityHandler
 }
diff --git a/handlers.go b/handlers.go
index 37d4617..014f638 100644
--- a/handlers.go
+++ b/handlers.go
@@ -165,7 +165,7 @@ func (*DefaultClientSecurityHandler) Handle(c Conn) error {
 		}
 		return fmt.Errorf("%s", reasonText)
 	}
-
+	c.SetSecurityHandler(secType)
 	return nil
 }
 
@@ -227,6 +227,7 @@ func (*DefaultServerSecurityHandler) Handle(c Conn) error {
 		}
 		return authErr
 	}
+	c.SetSecurityHandler(sType)
 	return nil
 }
 
diff --git a/server.go b/server.go
index dbd958d..c600c28 100644
--- a/server.go
+++ b/server.go
@@ -120,6 +120,17 @@ func (c *ServerConn) Protocol() string {
 	return c.protocol
 }
 
+// SecurityHandler returns security handler
+func (c *ServerConn) SecurityHandler() SecurityHandler {
+	return c.securityHandler
+}
+
+// SetSecurityHandler sets security handler
+func (c *ServerConn) SetSecurityHandler(sechandler SecurityHandler) error {
+	c.securityHandler = sechandler
+	return nil
+}
+
 // SetWidth sets framebuffer width
 func (c *ServerConn) SetWidth(w uint16) {
 	// TODO send desktopsize pseudo encoding
@@ -152,6 +163,8 @@ type ServerConn struct {
 	// directly. Instead, SetEncodings() should be used.
 	encodings []Encoding
 
+	securityHandler SecurityHandler
+
 	// Height of the frame buffer in pixels, sent to the client.
 	fbHeight uint16