added serverInit listening and recorder adaptation for being created outside of the client

.idea/misc.xml

@@ -1,30 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="GOROOT" path="/usr/local/Cellar/go/1.8.3/libexec" />
-  <component name="MavenImportPreferences">
-    <option name="generalSettings">
-      <MavenGeneralSettings>
-        <option name="mavenHome" value="Bundled (Maven 3)" />
-      </MavenGeneralSettings>
-    </option>
-  </component>
   <component name="ProjectRootManager" version="2" languageLevel="JDK_1_6" default="true">
     <output url="file://$PROJECT_DIR$/out" />
   </component>
-  <component name="masterDetails">
-    <states>
-      <state key="ProjectJDKs.UI">
-        <settings>
-          <last-edited>1.8</last-edited>
-          <splitter-proportions>
-            <option name="proportions">
-              <list>
-                <option value="0.2" />
-              </list>
-            </option>
-          </splitter-proportions>
-        </settings>
-      </state>
-    </states>
-  </component>
 </project>

client/client-conn.go

@@ -8,7 +8,6 @@ import (
 	"net"
 	"unicode"
 	"vncproxy/common"
-	"vncproxy/tee-listeners"
 )
 
 // A ServerMessage implements a message sent from the server to the client.
@@ -21,11 +20,11 @@ type ClientAuth interface {
 
 	// Handshake is called when the authentication handshake should be
 	// performed, as part of the general RFB handshake. (see 7.2.1)
-	Handshake(net.Conn) error
+	Handshake(io.ReadWriteCloser) error
 }
 
 type ClientConn struct {
-	conn net.Conn
+	conn io.ReadWriteCloser
 
 	//c      net.Conn
 	config *ClientConfig
@@ -77,6 +76,7 @@ type ClientConfig struct {
 	// This only needs to contain NEW server messages, and doesn't
 	// need to explicitly contain the RFC-required messages.
 	ServerMessages []common.ServerMessage
+	Listener       common.SegmentConsumer
 }
 
 func Client(c net.Conn, cfg *ClientConfig) (*ClientConn, error) {
@@ -438,6 +438,15 @@ FindAuth:
 	}
 
 	c.DesktopName = string(nameBytes)
+	srvInit := common.ServerInit{
+		NameLength:  nameLength,
+		NameText:    nameBytes,
+		FBHeight:    c.FrameBufferHeight,
+		FBWidth:     c.FrameBufferWidth,
+		PixelFormat: c.PixelFormat,
+	}
+	rfbSeg := &common.RfbSegment{SegmentType: common.SegmentServerInitMessage, Message: &srvInit}
+	c.config.Listener.Consume(rfbSeg)
 
 	return nil
 }
@@ -446,9 +455,8 @@ FindAuth:
 // proper channels for users of the client to read.
 func (c *ClientConn) mainLoop() {
 	defer c.Close()
-	rec := listeners.NewRecorder("/Users/amitbet/recording.rbs", c.DesktopName, c.FrameBufferWidth, c.FrameBufferHeight)
 
-	reader := &common.RfbReadHelper{Reader: c.conn, Listener: rec}
+	reader := &common.RfbReadHelper{Reader: c.conn, Listener: c.config.Listener}
 	// Build the map of available server messages
 	typeMap := make(map[uint8]common.ServerMessage)
 

client/client_auth.go

@@ -1,13 +1,11 @@
 package client
 
 import (
-	"net"
-
 	"crypto/des"
 	"encoding/binary"
+	"io"
 )
 
-
 // ClientAuthNone is the "none" authentication. See 7.2.1
 type ClientAuthNone byte
 
@@ -15,7 +13,7 @@ func (*ClientAuthNone) SecurityType() uint8 {
 	return 1
 }
 
-func (*ClientAuthNone) Handshake(net.Conn) error {
+func (*ClientAuthNone) Handshake(closer io.ReadWriteCloser) error {
 	return nil
 }
 
@@ -28,7 +26,7 @@ func (p *PasswordAuth) SecurityType() uint8 {
 	return 2
 }
 
-func (p *PasswordAuth) Handshake(c net.Conn) error {
+func (p *PasswordAuth) Handshake(c io.ReadWriteCloser) error {
 	randomValue := make([]uint8, 16)
 	if err := binary.Read(c, binary.BigEndian, &randomValue); err != nil {
 		return err
@@ -36,7 +34,7 @@ func (p *PasswordAuth) Handshake(c net.Conn) error {
 
 	crypted, err := p.encrypt(p.Password, randomValue)
 
-	if (err != nil) {
+	if err != nil {
 		return err
 	}
 

common/readers.go

@@ -12,6 +12,9 @@ const (
 	SegmentBytes SegmentType = iota
 	SegmentMessageSeparator
 	SegmentRectSeparator
+	SegmentFullyParsedClientMessage
+	SegmentFullyParsedServerMessage
+	SegmentServerInitMessage
 )
 
 type SegmentType int
@@ -20,7 +23,9 @@ type RfbSegment struct {
 	Bytes              []byte
 	SegmentType        SegmentType
 	UpcomingObjectType int
+	Message            interface{}
 }
+
 type SegmentConsumer interface {
 	Consume(*RfbSegment) error
 }

common/server-message.go

@@ -23,3 +23,10 @@ const (
 	Bell
 	ServerCutText
 )
+
+type ServerInit struct {
+	FBWidth, FBHeight uint16
+	PixelFormat       PixelFormat
+	NameLength        uint32
+	NameText          []byte
+}

encodings/enc-tight.go

@@ -28,9 +28,7 @@ type TightEncoding struct {
 // 	t.output = output
 // }
 
-func (*TightEncoding) Type() int32 {
+func (*TightEncoding) Type() int32 { return int32(common.EncTight) }
-	return 7
-}
 
 // func ReadAndRecBytes(conn io.Reader, rec io.Writer, count int) ([]byte, error) {
 // 	buf, err := readBytes(conn, count)

main.go

@@ -4,9 +4,10 @@ import (
 	"fmt"
 	"net"
 	"time"
+	"vncproxy/client"
 	"vncproxy/common"
 	"vncproxy/encodings"
-	"vncproxy/client"
+	listeners "vncproxy/tee-listeners"
 )
 
 func main() {
@@ -21,11 +22,18 @@ func main() {
 	authArr := []client.ClientAuth{&client.PasswordAuth{Password: "Ch_#!T@8"}, &noauth}
 
 	vncSrvMessagesChan := make(chan common.ServerMessage)
+
+	rec := listeners.NewRecorder("c:/Users/betzalel/recording.rbs")
+
+	split := &listeners.MultiListener{}
+	split.AddListener(rec)
+
 	clientConn, err := client.Client(nc,
 		&client.ClientConfig{
 			Auth:            authArr,
 			ServerMessageCh: vncSrvMessagesChan,
 			Exclusive:       true,
+			Listener:        split,
 		})
 
 	if err != nil {
@@ -37,6 +45,7 @@ func main() {
 	// }
 
 	tight := encodings.TightEncoding{}
+	tightPng := encodings.TightPngEncoding{}
 	//rre := encodings.RREEncoding{}
 	//zlib := encodings.ZLibEncoding{}
 	//zrle := encodings.ZRLEEncoding{}
@@ -47,7 +56,7 @@ func main() {
 	// defer file.Close()
 
 	//tight.SetOutput(file)
-	clientConn.SetEncodings([]common.Encoding{&cpyRect, &tight})
+	clientConn.SetEncodings([]common.Encoding{&cpyRect, &tightPng, &tight})
 
 	go func() {
 		for {

server/handlers.go

@@ -3,6 +3,7 @@ package server
 import (
 	"encoding/binary"
 	"fmt"
+	"vncproxy/common"
 
 	"io"
 )
@@ -126,7 +127,7 @@ func ServerSecurityHandler(cfg *ServerConfig, c *ServerConn) error {
 }
 
 func ServerServerInitHandler(cfg *ServerConfig, c *ServerConn) error {
-	srvInit := &ServerInit{
+	srvInit := &common.ServerInit{
 		FBWidth:     c.Width(),
 		FBHeight:    c.Height(),
 		PixelFormat: *c.PixelFormat(),

server/server.go

@@ -19,12 +19,7 @@ var DefaultClientMessages = []common.ClientMessage{
 	&ClientCutText{},
 }
 
-type ServerInit struct {
+
-	FBWidth, FBHeight uint16
-	PixelFormat       common.PixelFormat
-	NameLength        uint32
-	NameText          []byte
-}
 
 //var _ Conn = (*ServerConn)(nil)
 

tee-listeners/recorder.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"time"
 	"vncproxy/common"
+	"vncproxy/server"
 )
 
 type Recorder struct {
@@ -17,13 +18,15 @@ type Recorder struct {
 	logger              common.Logger
 	startTime           int
 	buffer              bytes.Buffer
+	serverInitMessage   *common.ServerInit
+	sessionStartWritten bool
 }
 
 func getNowMillisec() int {
 	return int(time.Now().UnixNano() / int64(time.Millisecond))
 }
 
-func NewRecorder(saveFilePath string, desktopName string, fbWidth uint16, fbHeight uint16) *Recorder {
+func NewRecorder(saveFilePath string) *Recorder {
 	//delete file if it exists
 	if _, err := os.Stat(saveFilePath); err == nil {
 		os.Remove(saveFilePath)
@@ -33,16 +36,25 @@ func NewRecorder(saveFilePath string, desktopName string, fbWidth uint16, fbHeig
 	var err error
 
 	rec.writer, err = os.OpenFile(saveFilePath, os.O_RDWR|os.O_CREATE, 0755)
-	rec.writeStartSession(desktopName, fbWidth, fbHeight)
-
 	if err != nil {
 		fmt.Printf("unable to open file: %s, error: %v", saveFilePath, err)
 		return nil
 	}
-
 	return &rec
 }
 
+// func (rec *Recorder) startSession(desktopName string, fbWidth uint16, fbHeight uint16) error {
+
+// 	err := rec.writeStartSession(desktopName, fbWidth, fbHeight)
+
+// 	if err != nil {
+// 		fmt.Printf("Recorder was unable to write StartSession to file error: %v", err)
+// 		return nil
+// 	}
+
+// 	return nil
+// }
+
 const versionMsg_3_3 = "RFB 003.003\n"
 const versionMsg_3_7 = "RFB 003.007\n"
 const versionMsg_3_8 = "RFB 003.008\n"
@@ -61,7 +73,11 @@ const (
 // 	// df.write("FBS 001.000\n".getBytes());
 // }
 
-func (r *Recorder) writeStartSession(desktopName string, framebufferWidth uint16, framebufferHeight uint16) error {
+func (r *Recorder) writeStartSession(initMsg *common.ServerInit) error {
+
+	desktopName := string(initMsg.NameText)
+	framebufferWidth := initMsg.FBWidth
+	framebufferHeight := initMsg.FBHeight
 
 	//write rfb header information (the only part done without the [size|data|timestamp] block wrapper)
 	r.buffer.WriteString("FBS 001.000\n")
@@ -90,6 +106,9 @@ func (r *Recorder) writeStartSession(desktopName string, framebufferWidth uint16
 func (r *Recorder) Consume(data *common.RfbSegment) error {
 	switch data.SegmentType {
 	case common.SegmentMessageSeparator:
+		if !r.sessionStartWritten {
+			r.writeStartSession(r.serverInitMessage)
+		}
 		switch common.ServerMessageType(data.UpcomingObjectType) {
 		case common.FramebufferUpdate:
 			r.writeToDisk()
@@ -105,6 +124,18 @@ func (r *Recorder) Consume(data *common.RfbSegment) error {
 	case common.SegmentBytes:
 		_, err := r.buffer.Write(data.Bytes)
 		return err
+	case common.SegmentServerInitMessage:
+		r.serverInitMessage = data.Message.(*common.ServerInit)
+	case common.SegmentFullyParsedClientMessage:
+		clientMsg := data.Message.(common.ClientMessage)
+		switch clientMsg.Type() {
+
+		case common.SetPixelFormatMsgType:
+			clientMsg := data.Message.(*server.SetPixelFormat)
+			r.serverInitMessage.PixelFormat = clientMsg.PF
+		default:
+			return errors.New("unknown client message type:" + string(data.UpcomingObjectType))
+		}
 
 	default:
 		return errors.New("undefined RfbSegment type")

tee-listeners/write-to.go

@@ -6,11 +6,11 @@ import (
 	"vncproxy/common"
 )
 
-type PassListener struct {
+type WriteToListener struct {
 	io.Writer
 }
 
-func (p *PassListener) Consume(seg *common.RfbSegment) error {
+func (p *WriteToListener) Consume(seg *common.RfbSegment) error {
 	switch seg.SegmentType {
 	case common.SegmentMessageSeparator:
 	case common.SegmentRectSeparator: