added RfbReaderHelper which wraps the connection reader and allows to add listeners like the recorder or proxy

added fbs recorder, and tested it against vine vnc server 5.0.1 and rfbplayer 1.4 —> works! (there is still some problem with tight VNC server)

common/encoding.go

@@ -1,9 +1,5 @@
 package common
 
-import (
-	"io"
-)
-
 // An Encoding implements a method for encoding pixel data that is
 // sent by the server to the client.
 type Encoding interface {
@@ -13,7 +9,7 @@ type Encoding interface {
 	// Read reads the contents of the encoded pixel data from the reader.
 	// This should return a new Encoding implementation that contains
 	// the proper data.
-	Read(*PixelFormat, *Rectangle, io.Reader) (Encoding, error)
+	Read(*PixelFormat, *Rectangle, *RfbReadHelper) (Encoding, error)
 }
 
 const (

common/readers.go

@@ -1,45 +1,85 @@
-// Package vnc implements a VNC client.
-//
-// References:
-//   [PROTOCOL]: http://tools.ietf.org/html/rfc6143
 package common
 
 import (
-	"bytes"
 	"encoding/binary"
 	"fmt"
 	"io"
 )
 
-// type DataSource struct {
-// 	conn        io.Reader
-// 	output      io.Writer
-// 	passThrough bool
-// 	PixelFormat PixelFormat
-// }Color
-type RfbReader struct {
-	reader    io.Reader
-	saveBytes bool
-	savedBuff bytes.Buffer
-}
+var TightMinToCompress = 12
 
-func (r *RfbReader) Read(p []byte) (n int, err error) {
-	readLen, err := r.reader.Read(p)
-	r.savedBuff.Write(p)
-	return readLen, err
+const (
+	SegmentBytes SegmentType = iota
+	SegmentMessageSeparator
+	SegmentRectSeparator
+)
+
+type SegmentType int
+
+type RfbSegment struct {
+	Bytes              []byte
+	SegmentType        SegmentType
+	UpcomingObjectType int
 }
-func (r *RfbReader) SavedBuff() bytes.Buffer {
-	return r.savedBuff
+type SegmentConsumer interface {
+	Consume(*RfbSegment) error
 }
 
 type RfbReadHelper struct {
 	io.Reader
+	Listener SegmentConsumer
 }
 
-func (d *RfbReadHelper) ReadBytes(count int) ([]byte, error) {
+func (r *RfbReadHelper) ReadDiscrete(p []byte) (int, error) {
+	return r.Read(p)
+}
+
+func (r *RfbReadHelper) SendRectSeparator(upcomingRectType int) error {
+	seg := &RfbSegment{SegmentType: SegmentRectSeparator, UpcomingObjectType: upcomingRectType}
+	if r.Listener != nil {
+		return nil
+	}
+	return r.Listener.Consume(seg)
+
+}
+
+func (r *RfbReadHelper) SendMessageSeparator(upcomingMessageType int) error {
+	seg := &RfbSegment{SegmentType: SegmentMessageSeparator, UpcomingObjectType: upcomingMessageType}
+	if r.Listener == nil {
+		return nil
+	}
+	return r.Listener.Consume(seg)
+}
+
+func (r *RfbReadHelper) PublishBytes(p []byte) error {
+	seg := &RfbSegment{Bytes: p, SegmentType: SegmentBytes}
+	if r.Listener == nil {
+		return nil
+	}
+	return r.Listener.Consume(seg)
+}
+
+func (r *RfbReadHelper) Read(p []byte) (n int, err error) {
+	readLen, err := r.Reader.Read(p)
+	if err != nil {
+		return 0, err
+	}
+	//write the bytes to the Listener for further processing
+	seg := &RfbSegment{Bytes: p, SegmentType: SegmentBytes}
+	if r.Listener == nil {
+		return 0,nil
+	}
+	r.Listener.Consume(seg)
+	if err != nil {
+		return 0, err
+	}
+	return readLen, err
+}
+
+func (r *RfbReadHelper) ReadBytes(count int) ([]byte, error) {
 	buff := make([]byte, count)
 
-	_, err := io.ReadFull(d.Reader, buff)
+	_, err := io.ReadFull(r, buff)
 	if err != nil {
 		//if err := binary.Read(d.conn, binary.BigEndian, &buff); err != nil {
 		return nil, err
@@ -47,41 +87,41 @@ func (d *RfbReadHelper) ReadBytes(count int) ([]byte, error) {
 	return buff, nil
 }
 
-func (d *RfbReadHelper) ReadUint8() (uint8, error) {
+func (r *RfbReadHelper) ReadUint8() (uint8, error) {
 	var myUint uint8
-	if err := binary.Read(d.Reader, binary.BigEndian, &myUint); err != nil {
+	if err := binary.Read(r, binary.BigEndian, &myUint); err != nil {
 		return 0, err
 	}
 	//fmt.Printf("myUint=%d", myUint)
 	return myUint, nil
 }
-func (d *RfbReadHelper) ReadUint16() (uint16, error) {
+func (r *RfbReadHelper) ReadUint16() (uint16, error) {
 	var myUint uint16
-	if err := binary.Read(d.Reader, binary.BigEndian, &myUint); err != nil {
+	if err := binary.Read(r, binary.BigEndian, &myUint); err != nil {
 		return 0, err
 	}
 	//fmt.Printf("myUint=%d", myUint)
 	return myUint, nil
 }
-func (d *RfbReadHelper) ReadUint32() (uint32, error) {
+func (r *RfbReadHelper) ReadUint32() (uint32, error) {
 	var myUint uint32
-	if err := binary.Read(d.Reader, binary.BigEndian, &myUint); err != nil {
+	if err := binary.Read(r, binary.BigEndian, &myUint); err != nil {
 		return 0, err
 	}
 	//fmt.Printf("myUint=%d", myUint)
 	return myUint, nil
 }
-func (d *RfbReadHelper) ReadCompactLen() (int, error) {
+func (r *RfbReadHelper) ReadCompactLen() (int, error) {
 	var err error
-	part, err := d.ReadUint8()
+	part, err := r.ReadUint8()
 	//byteCount := 1
 	len := uint32(part & 0x7F)
 	if (part & 0x80) != 0 {
-		part, err = d.ReadUint8()
+		part, err = r.ReadUint8()
 		//byteCount++
 		len |= uint32(part&0x7F) << 7
 		if (part & 0x80) != 0 {
-			part, err = d.ReadUint8()
+			part, err = r.ReadUint8()
 			//byteCount++
 			len |= uint32(part&0xFF) << 14
 		}
@@ -94,8 +134,6 @@ func (d *RfbReadHelper) ReadCompactLen() (int, error) {
 	return int(len), err
 }
 
-var TightMinToCompress int = 12
-
 func (r *RfbReadHelper) ReadTightData(dataSize int) ([]byte, error) {
 	if int(dataSize) < TightMinToCompress {
 		return r.ReadBytes(int(dataSize))

encodings/enc-copy-rect.go

@@ -1,9 +1,6 @@
 package encodings
 
-import (
-	"io"
-	"vncproxy/common"
-)
+import "vncproxy/common"
 
 type CopyRectEncoding struct {
 	//Colors       []Color
@@ -14,12 +11,9 @@ type CopyRectEncoding struct {
 func (z *CopyRectEncoding) Type() int32 {
 	return 1
 }
-func (z *CopyRectEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
-	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
-	//bytesPerPixel := c.PixelFormat.BPP / 8
-	z.copyRectSrcX, _ = conn.ReadUint16()
-	z.copyRectSrcY, _ = conn.ReadUint16()
+func (z *CopyRectEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
+	z.copyRectSrcX, _ = r.ReadUint16()
+	z.copyRectSrcY, _ = r.ReadUint16()
 	return z, nil
 }
 

encodings/enc-corre.go

@@ -1,7 +1,6 @@
 package encodings
 
 import (
-	"io"
 	"vncproxy/common"
 )
 
@@ -13,16 +12,15 @@ func (z *CoRREEncoding) Type() int32 {
 	return 4
 }
 
-func (z *CoRREEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
+func (z *CoRREEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPerPixel := int(pixelFmt.BPP / 8)
-	numOfSubrectangles, _ := conn.ReadUint32()
+	numOfSubrectangles, _ := r.ReadUint32()
 
 	//read whole rect background color
-	conn.ReadBytes(bytesPerPixel)
+	r.ReadBytes(bytesPerPixel)
 
 	//read all individual rects (color=BPP + x=16b + y=16b + w=16b + h=16b)
-	_, err := conn.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 4))
+	_, err := r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 4))
 
 	if err != nil {
 		return nil, err

encodings/enc-hextile.go

@@ -1,9 +1,6 @@
 package encodings
 
-import (
-	"io"
-	"vncproxy/common"
-)
+import "vncproxy/common"
 
 const (
 	HextileRaw                 = 1
@@ -20,11 +17,9 @@ type HextileEncoding struct {
 func (z *HextileEncoding) Type() int32 {
 	return 5
 }
-func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
-	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
+func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPerPixel := int(pixelFmt.BPP) / 8
-	//buf := make([]byte, bytesPerPixel)
+
 	for ty := rect.Y; ty < rect.Y+rect.Height; ty += 16 {
 		th := 16
 		if rect.Y+rect.Height-ty < 16 {
@@ -38,7 +33,7 @@ func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectan
 			}
 
 			//handle Hextile Subrect(tx, ty, tw, th):
-			subencoding, err := conn.ReadUint8()
+			subencoding, err := r.ReadUint8()
 			//fmt.Printf("hextile reader tile: (%d,%d) subenc=%d\n", ty, tx, subencoding)
 			if err != nil {
 				//fmt.Printf("error in hextile reader: %v\n", err)
@@ -47,22 +42,22 @@ func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectan
 
 			if (subencoding & HextileRaw) != 0 {
 				//ReadRawRect(c, rect, r)
-				conn.ReadBytes(tw * th * bytesPerPixel)
+				r.ReadBytes(tw * th * bytesPerPixel)
 				//fmt.Printf("hextile reader: HextileRaw\n")
 				continue
 			}
 			if (subencoding & HextileBackgroundSpecified) != 0 {
-				conn.ReadBytes(int(bytesPerPixel))
+				r.ReadBytes(int(bytesPerPixel))
 			}
 			if (subencoding & HextileForegroundSpecified) != 0 {
-				conn.ReadBytes(int(bytesPerPixel))
+				r.ReadBytes(int(bytesPerPixel))
 			}
 			if (subencoding & HextileAnySubrects) == 0 {
 				//fmt.Printf("hextile reader: no Subrects\n")
 				continue
 			}
 			//fmt.Printf("hextile reader: handling Subrects\n")
-			nSubrects, err := conn.ReadUint8()
+			nSubrects, err := r.ReadUint8()
 			if err != nil {
 				return nil, err
 			}
@@ -71,7 +66,7 @@ func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectan
 				bufsize += int(nSubrects) * int(bytesPerPixel)
 			}
 			//byte[] buf = new byte[bufsize];
-			conn.ReadBytes(bufsize)
+			r.ReadBytes(bufsize)
 		}
 	}
 

encodings/enc-raw.go

@@ -1,7 +1,6 @@
 package encodings
 
 import (
-	"io"
 	"vncproxy/common"
 )
 
@@ -16,9 +15,9 @@ func (*RawEncoding) Type() int32 {
 	return 0
 }
 
-func (*RawEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
+func (*RawEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
-	conn := common.RfbReadHelper{r}
+	//conn := common.RfbReadHelper{Reader:r}
 	bytesPerPixel := int(pixelFmt.BPP / 8)
 	//pixelBytes := make([]uint8, bytesPerPixel)
 
@@ -31,7 +30,7 @@ func (*RawEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r
 
 	for y := uint16(0); y < rect.Height; y++ {
 		for x := uint16(0); x < rect.Width; x++ {
-			if _, err := conn.ReadBytes(bytesPerPixel); err != nil {
+			if _, err := r.ReadBytes(bytesPerPixel); err != nil {
 				return nil, err
 			}
 

encodings/enc-rre.go

@@ -1,6 +1,5 @@
 package encodings
 
-import "io"
 import "vncproxy/common"
 
 type RREEncoding struct {
@@ -10,16 +9,16 @@ type RREEncoding struct {
 func (z *RREEncoding) Type() int32 {
 	return 2
 }
-func (z *RREEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
+func (z *RREEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
+	//conn := common.RfbReadHelper{Reader:r}
 	bytesPerPixel := int(pixelFmt.BPP / 8)
-	numOfSubrectangles, _ := conn.ReadUint32()
+	numOfSubrectangles, _ := r.ReadUint32()
 
 	//read whole rect background color
-	conn.ReadBytes(bytesPerPixel)
+	r.ReadBytes(bytesPerPixel)
 
 	//read all individual rects (color=BPP + x=16b + y=16b + w=16b + h=16b)
-	_, err := conn.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 8))
+	_, err := r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 8))
 
 	if err != nil {
 		return nil, err

encodings/enc-tight.go

@@ -69,13 +69,13 @@ func calcTightBytePerPixel(pf *common.PixelFormat) int {
 	return bytesPerPixelTight
 }
 
-func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, reader io.Reader) (common.Encoding, error) {
+func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPixel := calcTightBytePerPixel(pixelFmt)
-	conn := common.RfbReadHelper{reader}
+	//conn := common.RfbReadHelper{Reader:reader}
 	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
 
 	//var subencoding uint8
-	subencoding, err := conn.ReadUint8()
+	subencoding, err := r.ReadUint8()
 	if err != nil {
 		fmt.Printf("error in handling tight encoding: %v\n", err)
 		return nil, err
@@ -99,19 +99,19 @@ func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangl
 	case TightFill:
 		fmt.Printf("reading fill size=%d\n", bytesPixel)
 		//read color
-		conn.ReadBytes(int(bytesPixel))
+		r.ReadBytes(int(bytesPixel))
 		return t, nil
 	case TightJpeg:
 		if pixelFmt.BPP == 8 {
 			return nil, errors.New("Tight encoding: JPEG is not supported in 8 bpp mode")
 		}
 
-		len, err := conn.ReadCompactLen()
+		len, err := r.ReadCompactLen()
 		if err != nil {
 			return nil, err
 		}
 		fmt.Printf("reading jpeg size=%d\n", len)
-		conn.ReadBytes(len)
+		r.ReadBytes(len)
 		return t, nil
 	default:
 
@@ -119,13 +119,13 @@ func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangl
 			fmt.Println("Compression control byte is incorrect!")
 		}
 
-		handleTightFilters(subencoding, pixelFmt, rect, reader)
+		handleTightFilters(subencoding, pixelFmt, rect, r)
 		return t, nil
 	}
 }
 
-func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *common.Rectangle, reader io.Reader) {
-	conn := common.RfbReadHelper{reader}
+func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) {
+	//conn := common.RfbReadHelper{Reader:reader}
 	var FILTER_ID_MASK uint8 = 0x40
 	//var STREAM_ID_MASK uint8 = 0x30
 
@@ -134,7 +134,7 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 	var err error
 
 	if (subencoding & FILTER_ID_MASK) > 0 { // filter byte presence
-		filterid, err = conn.ReadUint8()
+		filterid, err = r.ReadUint8()
 		if err != nil {
 			fmt.Printf("error in handling tight encoding, reading filterid: %v\n", err)
 			return
@@ -154,11 +154,11 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 	switch filterid {
 	case TightFilterPalette: //PALETTE_FILTER
 
-		colorCount, err := conn.ReadUint8()
+		colorCount, err := r.ReadUint8()
 		paletteSize := colorCount + 1 // add one more
 		fmt.Printf("----PALETTE_FILTER: paletteSize=%d bytesPixel=%d\n", paletteSize, bytesPixel)
 		//complete palette
-		conn.ReadBytes(int(paletteSize) * bytesPixel)
+		r.ReadBytes(int(paletteSize) * bytesPixel)
 
 		var dataLength int
 		if paletteSize == 2 {
@@ -166,7 +166,7 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 		} else {
 			dataLength = int(rect.Width * rect.Height)
 		}
-		_, err = conn.ReadTightData(dataLength)
+		_, err = r.ReadTightData(dataLength)
 		if err != nil {
 			fmt.Printf("error in handling tight encoding, Reading Palette: %v\n", err)
 			return
@@ -175,10 +175,10 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 		fmt.Printf("----GRADIENT_FILTER: bytesPixel=%d\n", bytesPixel)
 		//useGradient = true
 		fmt.Printf("usegrad: %d\n", filterid)
-		conn.ReadTightData(lengthCurrentbpp)
+		r.ReadTightData(lengthCurrentbpp)
 	case TightFilterCopy: //BASIC_FILTER
 		fmt.Printf("----BASIC_FILTER: bytesPixel=%d\n", bytesPixel)
-		conn.ReadTightData(lengthCurrentbpp)
+		r.ReadTightData(lengthCurrentbpp)
 	default:
 		fmt.Printf("Bad tight filter id: %d\n", filterid)
 		return

encodings/enc-zlib.go

@@ -1,6 +1,5 @@
 package encodings
 
-import "io"
 import "vncproxy/common"
 
 type ZLibEncoding struct {
@@ -10,12 +9,12 @@ type ZLibEncoding struct {
 func (z *ZLibEncoding) Type() int32 {
 	return 6
 }
-func (z *ZLibEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
+func (z *ZLibEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
+	//conn := common.RfbReadHelper{Reader:r}
 	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
 	//bytesPerPixel := c.PixelFormat.BPP / 8
-	len, _ := conn.ReadUint32()
-	_, err := conn.ReadBytes(int(len))
+	len, _ := r.ReadUint32()
+	_, err := r.ReadBytes(int(len))
 
 	if err != nil {
 		return nil, err

encodings/enc-zrle.go

@@ -1,6 +1,5 @@
 package encodings
 
-import "io"
 import "vncproxy/common"
 
 type ZRLEEncoding struct {
@@ -10,12 +9,12 @@ type ZRLEEncoding struct {
 func (z *ZRLEEncoding) Type() int32 {
 	return 16
 }
-func (z *ZRLEEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r io.Reader) (common.Encoding, error) {
-	conn := common.RfbReadHelper{r}
+func (z *ZRLEEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
+	//conn := common.RfbReadHelper{Reader: r}
 	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
 	//bytesPerPixel := c.PixelFormat.BPP / 8
-	len, _ := conn.ReadUint32()
-	_, err := conn.ReadBytes(int(len))
+	len, _ := r.ReadUint32()
+	_, err := r.ReadBytes(int(len))
 
 	if err != nil {
 		return nil, err

listeners/multiListener.go

@@ -0,0 +1,22 @@
+package listeners
+
+import "vncproxy/common"
+
+type MultiListener struct {
+	listeners []common.SegmentConsumer
+}
+
+func (m *MultiListener) AddListener(listener common.SegmentConsumer) {
+	m.listeners = append(m.listeners, listener)
+}
+
+func (m *MultiListener) Consume(seg *common.RfbSegment) error {
+	for _, li := range m.listeners {
+		//fmt.Println(li)
+		err := li.Consume(seg)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

listeners/pass-to.go

@@ -0,0 +1,12 @@
+package listeners
+
+import "vncproxy/common"
+import "io"
+
+type PassListener struct {
+	io.Writer
+}
+
+func (*PassListener) Consume(seg *common.RfbSegment) error {
+	return nil
+}

listeners/recorder.go

@@ -0,0 +1,155 @@
+package listeners
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"os"
+	"time"
+	"vncproxy/common"
+)
+
+type Recorder struct {
+	//common.BytesListener
+	RBSFileName string
+	writer      *os.File
+	logger      common.Logger
+	startTime   int
+	buffer      bytes.Buffer
+}
+
+func getNowMillisec() int {
+	return int(time.Now().UnixNano() / int64(time.Millisecond))
+}
+
+func NewRecorder(saveFilePath string, desktopName string, fbWidth uint16, fbHeight uint16) *Recorder {
+	//delete file if it exists
+	if _, err := os.Stat(saveFilePath); err == nil {
+		os.Remove(saveFilePath)
+	}
+
+	rec := Recorder{RBSFileName: saveFilePath, startTime: getNowMillisec()}
+	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
+}
+
+const (
+	FramebufferUpdate   = 0
+	SetColourMapEntries = 1
+	Bell                = 2
+	ServerCutText       = 3
+)
+const versionMsg_3_3 = "RFB 003.003\n"
+const versionMsg_3_7 = "RFB 003.007\n"
+const versionMsg_3_8 = "RFB 003.008\n"
+
+// Security types
+const (
+	SecTypeInvalid = 0
+	SecTypeNone    = 1
+	SecTypeVncAuth = 2
+	SecTypeTight   = 16
+)
+
+// func (r *Recorder) writeHeader() error {
+// 	_, err := r.writer.WriteString("FBS 001.000\n")
+// 	return err
+// 	// df.write("FBS 001.000\n".getBytes());
+// }
+
+func (r *Recorder) writeStartSession(desktopName string, framebufferWidth uint16, framebufferHeight uint16) error {
+
+	//write rfb header information (the only part done without the [size|data|timestamp] block wrapper)
+	r.buffer.WriteString("FBS 001.000\n")
+	r.buffer.WriteTo(r.writer)
+	r.buffer.Reset()
+
+	//push the version message into the buffer so it will be written in the first rbs block
+	r.buffer.WriteString(versionMsg_3_3)
+
+	//push sec type and fb dimensions
+	binary.Write(&r.buffer, binary.BigEndian, int32(SecTypeNone))
+	binary.Write(&r.buffer, binary.BigEndian, int16(framebufferWidth))
+	binary.Write(&r.buffer, binary.BigEndian, int16(framebufferHeight))
+
+	var fbsServerInitMsg = []byte{32, 24, 0, 1, 0, byte(0xFF), 0, byte(0xFF), 0, byte(0xFF), 16, 8, 0, 0, 0, 0}
+	r.buffer.Write(fbsServerInitMsg)
+
+	binary.Write(&r.buffer, binary.BigEndian, uint32(len(desktopName)+1))
+
+	r.buffer.WriteString(desktopName)
+	binary.Write(&r.buffer, binary.BigEndian, byte(0)) // add null termination for desktop string
+
+	return nil
+}
+
+func (r *Recorder) Consume(data *common.RfbSegment) error {
+	switch data.SegmentType {
+	case common.SegmentMessageSeparator:
+		switch data.UpcomingObjectType {
+		case FramebufferUpdate:
+			r.writeToDisk()
+		case SetColourMapEntries:
+		case Bell:
+		case ServerCutText:
+		default:
+			return errors.New("unknown message type:" + string(data.UpcomingObjectType))
+		}
+
+	case common.SegmentRectSeparator:
+		r.writeToDisk()
+	case common.SegmentBytes:
+		_, err := r.buffer.Write(data.Bytes)
+		return err
+
+	default:
+		return errors.New("undefined RfbSegment type")
+	}
+	return nil
+}
+
+func (r *Recorder) writeToDisk() error {
+	timeSinceStart := getNowMillisec() - r.startTime
+	if r.buffer.Len() == 0 {
+		return nil
+	}
+
+	//write buff length
+	bytesLen := r.buffer.Len()
+	binary.Write(r.writer, binary.BigEndian, uint32(bytesLen))
+	paddedSize := (bytesLen + 3) & 0x7FFFFFFC
+	paddingSize := paddedSize - bytesLen
+
+	fmt.Printf("paddedSize=%d paddingSize=%d bytesLen=%d", paddedSize, paddingSize, bytesLen)
+	//write buffer padded to 32bit
+	_, err := r.buffer.WriteTo(r.writer)
+	padding := make([]byte, paddingSize)
+	fmt.Printf("padding=%v ", padding)
+
+	binary.Write(r.writer, binary.BigEndian, padding)
+
+	//write timestamp
+	binary.Write(r.writer, binary.BigEndian, uint32(timeSinceStart))
+	r.buffer.Reset()
+	return err
+}
+
+// func (r *Recorder) WriteUInt8(data uint8) error {
+// 	buf := make([]byte, 1)
+// 	buf[0] = byte(data) // cast int8 to byte
+// 	return r.Write(buf)
+// }
+
+func (r *Recorder) Close() {
+	r.writer.Close()
+}

main.go

@@ -16,7 +16,7 @@ func main() {
 	nc, err := net.Dial("tcp", "localhost:5903")
 
 	if err != nil {
-		fmt.Printf(";error connecting to vnc server: %s", err)
+		fmt.Printf("error connecting to vnc server: %s", err)
 	}
 	var noauth vnc.ClientAuthNone
 	authArr := []vnc.ClientAuth{&vnc.PasswordAuth{Password: "Ch_#!T@8"}, &noauth}

vnc/client-conn.go

@@ -8,6 +8,7 @@ import (
 	"net"
 	"unicode"
 	"vncproxy/common"
+	"vncproxy/listeners"
 )
 
 // A ServerMessage implements a message sent from the server to the client.
@@ -18,7 +19,7 @@ type ServerMessage interface {
 	// Read reads the contents of the message from the reader. At the point
 	// this is called, the message type has already been read from the reader.
 	// This should return a new ServerMessage that is the appropriate type.
-	Read(*ClientConn, io.Reader) (ServerMessage, error)
+	Read(*ClientConn, *common.RfbReadHelper) (ServerMessage, error)
 }
 
 // A ClientAuth implements a method of authenticating with a remote server.
@@ -442,7 +443,9 @@ 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}
 	// Build the map of available server messages
 	typeMap := make(map[uint8]ServerMessage)
 
@@ -474,8 +477,10 @@ func (c *ClientConn) mainLoop() {
 			// Unsupported message type! Bad!
 			break
 		}
+		reader.SendMessageSeparator(int(messageType))
+		reader.PublishBytes([]byte{byte(messageType)})
 
-		parsedMsg, err := msg.Read(c, c.conn)
+		parsedMsg, err := msg.Read(c, reader)
 		if err != nil {
 			break
 		}

vnc/recorder.go

@@ -1,38 +0,0 @@
-package vnc
-
-import (
-	"os"
-	"vncproxy/common"
-)
-
-type Recorder struct {
-	RBSFileName string
-	fileHandle  *os.File
-	logger      common.Logger
-}
-
-func NewRecorder(saveFilePath string, logger common.Logger) *Recorder {
-	rec := Recorder{RBSFileName: saveFilePath}
-	var err error
-	rec.fileHandle, err = os.OpenFile(saveFilePath, os.O_RDWR|os.O_CREATE, 0755)
-	if err != nil {
-		logger.Errorf("unable to open file: %s, error: %v", saveFilePath, err)
-		return nil
-	}
-	return &rec
-}
-
-func (r *Recorder) Write(data []byte) error {
-	_, err := r.fileHandle.Write(data)
-	return err
-}
-
-// func (r *Recorder) WriteUInt8(data uint8) error {
-// 	buf := make([]byte, 1)
-// 	buf[0] = byte(data) // cast int8 to byte
-// 	return r.Write(buf)
-// }
-
-func (r *Recorder) Close() {
-	r.fileHandle.Close()
-}

vnc/server_messages.go

@@ -27,7 +27,8 @@ func (*FramebufferUpdateMessage) Type() uint8 {
 	return 0
 }
 
-func (*FramebufferUpdateMessage) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
+func (fbm *FramebufferUpdateMessage) Read(c *ClientConn, r *common.RfbReadHelper) (ServerMessage, error) {
+
 	// Read off the padding
 	var padding [1]byte
 	if _, err := io.ReadFull(r, padding[:]); err != nil {
@@ -53,8 +54,9 @@ func (*FramebufferUpdateMessage) Read(c *ClientConn, r io.Reader) (ServerMessage
 	rects := make([]common.Rectangle, numRects)
 	for i := uint16(0); i < numRects; i++ {
 		fmt.Printf("###############rect################: %d\n", i)
-		var encodingType int32
 
+		var encodingType int32
+		r.SendRectSeparator(-1)
 		rect := &rects[i]
 		data := []interface{}{
 			&rect.X,
@@ -108,7 +110,7 @@ func (*SetColorMapEntriesMessage) Type() uint8 {
 	return 1
 }
 
-func (*SetColorMapEntriesMessage) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
+func (*SetColorMapEntriesMessage) Read(c *ClientConn, r *common.RfbReadHelper) (ServerMessage, error) {
 	// Read off the padding
 	var padding [1]byte
 	if _, err := io.ReadFull(r, padding[:]); err != nil {
@@ -161,7 +163,7 @@ func (*BellMessage) Type() uint8 {
 	return 2
 }
 
-func (*BellMessage) Read(*ClientConn, io.Reader) (ServerMessage, error) {
+func (*BellMessage) Read(*ClientConn, *common.RfbReadHelper) (ServerMessage, error) {
 	return new(BellMessage), nil
 }
 
@@ -180,22 +182,31 @@ func (*ServerCutTextMessage) Type() uint8 {
 	return 3
 }
 
-func (*ServerCutTextMessage) Read(conn *ClientConn, r io.Reader) (ServerMessage, error) {
+func (*ServerCutTextMessage) Read(conn *ClientConn, r *common.RfbReadHelper) (ServerMessage, error) {
+	//reader := common.RfbReadHelper{Reader: r}
+
 	// Read off the padding
-	var padding [1]byte
+	var padding [3]byte
 	if _, err := io.ReadFull(r, padding[:]); err != nil {
 		return nil, err
 	}
-
-	var textLength uint32
-	if err := binary.Read(r, binary.BigEndian, &textLength); err != nil {
+	textLength, err := r.ReadUint32()
+	if err != nil {
 		return nil, err
 	}
-
-	textBytes := make([]uint8, textLength)
-	if err := binary.Read(r, binary.BigEndian, &textBytes); err != nil {
+	textBytes, err := r.ReadBytes(int(textLength))
+	if err != nil {
 		return nil, err
 	}
+	//var textLength uint32
+	// if err := binary.Read(r, binary.BigEndian, &textLength); err != nil {
+	// 	return nil, err
+	// }
+
+	// textBytes := make([]uint8, textLength)
+	// if err := binary.Read(r, binary.BigEndian, &textBytes); err != nil {
+	// 	return nil, err
+	// }
 
 	return &ServerCutTextMessage{string(textBytes)}, nil
 }