made vnc replay-to-client work!

client/server-messages.go

@@ -9,6 +9,7 @@ import (
 	"vncproxy/common"
 	"vncproxy/encodings"
 	"vncproxy/logger"
+	listeners "vncproxy/tee-listeners"
 )
 
 // FramebufferUpdateMessage consists of a sequence of rectangles of
@@ -29,6 +30,14 @@ func (*FramebufferUpdateMessage) Type() uint8 {
 	return 0
 }
 
+func (fbm *FramebufferUpdateMessage) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
+	reader := common.NewRfbReadHelper(r)
+	writeTo := &listeners.WriteTo{w, "FramebufferUpdateMessage.CopyTo"}
+	reader.Listeners.AddListener(writeTo)
+	_, err := fbm.Read(c, reader)
+	return err
+}
+
 func (fbm *FramebufferUpdateMessage) Read(c common.IClientConn, r *common.RfbReadHelper) (common.ServerMessage, error) {
 
 	// Read off the padding
@@ -112,6 +121,13 @@ type SetColorMapEntriesMessage struct {
 	Colors     []common.Color
 }
 
+func (fbm *SetColorMapEntriesMessage) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
+	reader := &common.RfbReadHelper{Reader: r}
+	writeTo := &listeners.WriteTo{w, "SetColorMapEntriesMessage.CopyTo"}
+	reader.Listeners.AddListener(writeTo)
+	_, err := fbm.Read(c, reader)
+	return err
+}
 func (m *SetColorMapEntriesMessage) String() string {
 	return fmt.Sprintf("SetColorMapEntriesMessage (type=%d) first:%d colors: %v: ", m.Type(), m.FirstColor, m.Colors)
 }
@@ -165,6 +181,9 @@ func (*SetColorMapEntriesMessage) Read(c common.IClientConn, r *common.RfbReadHe
 // See RFC 6143 Section 7.6.3
 type BellMessage byte
 
+func (fbm *BellMessage) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
+	return nil
+}
 func (m *BellMessage) String() string {
 	return fmt.Sprintf("BellMessage (type=%d)", m.Type())
 }
@@ -184,6 +203,13 @@ type ServerCutTextMessage struct {
 	Text string
 }
 
+func (fbm *ServerCutTextMessage) CopyTo(r io.Reader, w io.Writer, c common.IClientConn) error {
+	reader := &common.RfbReadHelper{Reader: r}
+	writeTo := &listeners.WriteTo{w, "ServerCutTextMessage.CopyTo"}
+	reader.Listeners.AddListener(writeTo)
+	_, err := fbm.Read(c, reader)
+	return err
+}
 func (m *ServerCutTextMessage) String() string {
 	return fmt.Sprintf("ServerCutTextMessage (type=%d)", m.Type())
 }

common/encoding.go

@@ -11,7 +11,7 @@ import (
 type Encoding interface {
 	// The number that uniquely identifies this encoding type.
 	Type() int32
-
+	WriteTo(w io.Writer) (n int, err error)
 	// Read reads the contents of the encoded pixel data from the reader.
 	// This should return a new Encoding implementation that contains
 	// the proper data.

common/readers.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"bytes"
 	"encoding/binary"
 	"io"
 	"vncproxy/logger"
@@ -18,6 +19,14 @@ const (
 	SegmentConnectionClosed
 )
 
+// type ListenerType int
+
+// const (
+// 	ListenerTypeRawBytes ListenerType = iota
+// 	ListenerTypeParsedBytes
+// 	ListenerTypeClientMessages
+// )
+
 type SegmentType int
 
 func (seg SegmentType) String() string {
@@ -34,7 +43,10 @@ func (seg SegmentType) String() string {
 		return "SegmentFullyParsedServerMessage"
 	case SegmentServerInitMessage:
 		return "SegmentServerInitMessage"
+	case SegmentConnectionClosed:
+		return "SegmentConnectionClosed"
 	}
+
 	return ""
 }
 
@@ -51,7 +63,22 @@ type SegmentConsumer interface {
 
 type RfbReadHelper struct {
 	io.Reader
-	Listeners *MultiListener
+	Listeners  *MultiListener
+	savedBytes *bytes.Buffer
+}
+
+func NewRfbReadHelper(r io.Reader) *RfbReadHelper {
+	return &RfbReadHelper{Reader: r, Listeners: &MultiListener{}}
+}
+
+func (r *RfbReadHelper) StartByteCollection() {
+	r.savedBytes = &bytes.Buffer{}
+}
+
+func (r *RfbReadHelper) EndByteCollection() []byte {
+	bts := r.savedBytes.Bytes()
+	r.savedBytes = nil
+	return bts
 }
 
 func (r *RfbReadHelper) ReadDiscrete(p []byte) (int, error) {
@@ -78,6 +105,14 @@ func (r *RfbReadHelper) Read(p []byte) (n int, err error) {
 	if err != nil {
 		return 0, err
 	}
+	//if saving up our bytes, write them into the predefined buffer
+	if r.savedBytes != nil {
+		_, err := r.savedBytes.Write(p)
+		if err != nil {
+			logger.Warn("RfbReadHelper.Read: failed to collect bytes in mem buffer:", err)
+		}
+	}
+
 	//write the bytes to the Listener for further processing
 	seg := &RfbSegment{Bytes: p, SegmentType: SegmentBytes}
 	err = r.Listeners.Consume(seg)

common/server-message.go

@@ -1,5 +1,9 @@
 package common
 
+import (
+	"io"
+)
+
 type IClientConn interface {
 	CurrentPixelFormat() *PixelFormat
 	CurrentColorMap() *ColorMap
@@ -10,6 +14,7 @@ type ServerMessage interface {
 	// The type of the message that is sent down on the wire.
 	Type() uint8
 	String() string
+	CopyTo(r io.Reader, w io.Writer, c IClientConn) error
 	// 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.
@@ -18,14 +23,14 @@ type ServerMessage interface {
 type ServerMessageType int8
 
 const (
-	FramebufferUpdate   ServerMessageType = iota
+	FramebufferUpdate ServerMessageType = iota
 	SetColourMapEntries
 	Bell
 	ServerCutText
 )
 
 func (typ ServerMessageType) String() string {
-	switch    typ {
+	switch typ {
 	case FramebufferUpdate:
 		return "FramebufferUpdate"
 	case SetColourMapEntries:

encodings/enc-copy-rect.go

@@ -1,6 +1,10 @@
 package encodings
 
-import "vncproxy/common"
+import (
+	"encoding/binary"
+	"io"
+	"vncproxy/common"
+)
 
 type CopyRectEncoding struct {
 	//Colors       []Color
@@ -11,6 +15,18 @@ type CopyRectEncoding struct {
 func (z *CopyRectEncoding) Type() int32 {
 	return 1
 }
+func (z *CopyRectEncoding) WriteTo(w io.Writer) (n int, err error) {
+	binary.Write(w, binary.BigEndian, z.copyRectSrcX)
+	if err != nil {
+		return 0, err
+	}
+	binary.Write(w, binary.BigEndian, z.copyRectSrcY)
+	if err != nil {
+		return 0, err
+	}
+	return 4, nil
+}
+
 func (z *CopyRectEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	z.copyRectSrcX, _ = r.ReadUint16()
 	z.copyRectSrcY, _ = r.ReadUint16()

encodings/enc-corre.go

@@ -1,10 +1,15 @@
 package encodings
 
 import (
+	"encoding/binary"
+	"io"
 	"vncproxy/common"
 )
 
 type CoRREEncoding struct {
+	numSubRects     uint32
+	backgroundColor []byte
+	subRectData     []byte
 	//Colors []Color
 }
 
@@ -12,35 +17,45 @@ func (z *CoRREEncoding) Type() int32 {
 	return 4
 }
 
+func (z *CoRREEncoding) WriteTo(w io.Writer) (n int, err error) {
+	binary.Write(w, binary.BigEndian, z.numSubRects)
+	if err != nil {
+		return 0, err
+	}
+
+	w.Write(z.backgroundColor)
+	if err != nil {
+		return 0, err
+	}
+
+	w.Write(z.subRectData)
+
+	if err != nil {
+		return 0, err
+	}
+	b := len(z.backgroundColor) + len(z.subRectData) + 4
+	return b, nil
+}
+
 func (z *CoRREEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPerPixel := int(pixelFmt.BPP / 8)
-	numOfSubrectangles, _ := r.ReadUint32()
-
-	//read whole rect background color
-	r.ReadBytes(bytesPerPixel)
-
-	//read all individual rects (color=BPP + x=16b + y=16b + w=16b + h=16b)
-	_, err := r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 4))
-
+	numOfSubrectangles, err := r.ReadUint32()
 	if err != nil {
 		return nil, err
 	}
+	z.numSubRects = numOfSubrectangles
+
+	//read whole-rect background color
+	z.backgroundColor, err = r.ReadBytes(bytesPerPixel)
+	if err != nil {
+		return nil, err
+	}
+
+	//read all individual rects (color=BPP + x=16b + y=16b + w=16b + h=16b)
+	z.subRectData, err = r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 4))
+	if err != nil {
+		return nil, err
+	}
+
 	return z, nil
-
-	//int nSubrects = rfb.readU32();
-
-	//byte[] bg_buf = new byte[bytesPerPixel];
-	//rfb.readFully(bytesPerPixel);
-	//Color pixel;
-	// if (bytesPixel == 1) {
-	//   pixel = colors[bg_buf[0] & 0xFF];
-	// } else {
-	//   pixel = new Color(bg_buf[2] & 0xFF, bg_buf[1] & 0xFF, bg_buf[0] & 0xFF);
-	// }
-	// memGraphics.setColor(pixel);
-	// memGraphics.fillRect(x, y, w, h);
-
-	// byte[] buf = new byte[nSubrects * (bytesPixel + 4)];
-	// rfb.readFully(buf);
-
 }

encodings/enc-hextile.go

@@ -1,6 +1,7 @@
 package encodings
 
 import (
+	"io"
 	"vncproxy/common"
 	"vncproxy/logger"
 )
@@ -15,14 +16,23 @@ const (
 
 type HextileEncoding struct {
 	//Colors []Color
+	bytes []byte
 }
 
 func (z *HextileEncoding) Type() int32 {
 	return 5
 }
+func (z *HextileEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
+}
 func (z *HextileEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPerPixel := int(pixelFmt.BPP) / 8
 
+	r.StartByteCollection()
+	defer func() {
+		z.bytes = r.EndByteCollection()
+	}()
+
 	for ty := rect.Y; ty < rect.Y+rect.Height; ty += 16 {
 		th := 16
 		if rect.Y+rect.Height-ty < 16 {

encodings/enc-pseudo.go

@@ -1,6 +1,9 @@
 package encodings
 
-import "vncproxy/common"
+import (
+	"io"
+	"vncproxy/common"
+)
 
 type PseudoEncoding struct {
 	Typ int32
@@ -9,7 +12,9 @@ type PseudoEncoding struct {
 func (pe *PseudoEncoding) Type() int32 {
 	return pe.Typ
 }
-
+func (z *PseudoEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return 0, nil
+}
 func (pe *PseudoEncoding) Read(*common.PixelFormat, *common.Rectangle, *common.RfbReadHelper) (common.Encoding, error) {
 	return pe, nil
 }

encodings/enc-raw.go

@@ -1,6 +1,8 @@
 package encodings
 
 import (
+	"bytes"
+	"io"
 	"vncproxy/common"
 )
 
@@ -9,12 +11,15 @@ import (
 // See RFC 6143 Section 7.7.1
 type RawEncoding struct {
 	//Colors []Color
+	bytes []byte
 }
 
 func (*RawEncoding) Type() int32 {
 	return 0
 }
-
+func (z *RawEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
+}
 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{Reader:r}
@@ -27,10 +32,11 @@ func (*RawEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r
 	// }
 
 	//colors := make([]vnc.Color, int(rect.Height)*int(rect.Width))
-
+	bytes := &bytes.Buffer{}
 	for y := uint16(0); y < rect.Height; y++ {
 		for x := uint16(0); x < rect.Width; x++ {
-			if _, err := r.ReadBytes(bytesPerPixel); err != nil {
+			if bts, err := r.ReadBytes(bytesPerPixel); err != nil {
+				StoreBytes(bytes, bts)
 				return nil, err
 			}
 
@@ -54,5 +60,5 @@ func (*RawEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r
 		}
 	}
 
-	return &RawEncoding{}, nil
+	return &RawEncoding{bytes.Bytes()}, nil
 }

encodings/enc-rre.go

@@ -1,25 +1,57 @@
 package encodings
 
-import "vncproxy/common"
+import (
+	"encoding/binary"
+	"io"
+	"vncproxy/common"
+)
 
 type RREEncoding struct {
 	//Colors []Color
+	numSubRects     uint32
+	backgroundColor []byte
+	subRectData     []byte
+}
+
+func (z *RREEncoding) WriteTo(w io.Writer) (n int, err error) {
+	binary.Write(w, binary.BigEndian, z.numSubRects)
+	if err != nil {
+		return 0, err
+	}
+
+	w.Write(z.backgroundColor)
+	if err != nil {
+		return 0, err
+	}
+
+	w.Write(z.subRectData)
+
+	if err != nil {
+		return 0, err
+	}
+	b := len(z.backgroundColor) + len(z.subRectData) + 4
+	return b, nil
 }
 
 func (z *RREEncoding) Type() int32 {
 	return 2
 }
 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, _ := r.ReadUint32()
+	numOfSubrectangles, err := r.ReadUint32()
+	if err != nil {
+		return nil, err
+	}
+	z.numSubRects = numOfSubrectangles
 
-	//read whole rect background color
-	r.ReadBytes(bytesPerPixel)
+	//read whole-rect background color
+	z.backgroundColor, err = r.ReadBytes(bytesPerPixel)
+	if err != nil {
+		return nil, err
+	}
 
 	//read all individual rects (color=bytesPerPixel + x=16b + y=16b + w=16b + h=16b)
-	_, err := r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 8)) // x+y+w+h=8 bytes
-
+	z.subRectData, err = r.ReadBytes(int(numOfSubrectangles) * (bytesPerPixel + 8)) // x+y+w+h=8 bytes
 	if err != nil {
 		return nil, err
 	}

encodings/enc-tight.go

@@ -1,9 +1,12 @@
 package encodings
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"vncproxy/common"
+	"vncproxy/logger"
 )
 
 var TightMinToCompress int = 12
@@ -20,6 +23,7 @@ const (
 )
 
 type TightEncoding struct {
+	bytes []byte
 	//output io.Writer
 	//logger common.Logger
 }
@@ -67,13 +71,29 @@ func calcTightBytePerPixel(pf *common.PixelFormat) int {
 	return bytesPerPixelTight
 }
 
+func (z *TightEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
+}
+
+func StoreBytes(bytes *bytes.Buffer, data []byte) {
+	_, err := bytes.Write(data)
+	if err != nil {
+		logger.Error("Error in encoding while saving bytes: ", err)
+	}
+}
+
 func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPixel := calcTightBytePerPixel(pixelFmt)
 	//conn := common.RfbReadHelper{Reader:reader}
 	//conn := &DataSource{conn: conn.c, PixelFormat: conn.PixelFormat}
-
+	r.StartByteCollection()
+	defer func() {
+		t.bytes = r.EndByteCollection()
+	}()
 	//var subencoding uint8
 	compctl, err := r.ReadUint8()
+	//binary.Write(bytes, binary.BigEndian, compctl)
+
 	if err != nil {
 		fmt.Printf("error in handling tight encoding: %v\n", err)
 		return nil, err
@@ -88,7 +108,13 @@ func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangl
 	case TightFill:
 		fmt.Printf("reading fill size=%d\n", bytesPixel)
 		//read color
-		r.ReadBytes(int(bytesPixel))
+		_, err := r.ReadBytes(int(bytesPixel))
+		if err != nil {
+			fmt.Printf("error in handling tight encoding: %v\n", err)
+			return nil, err
+		}
+		//StoreBytes(bytes, bts)
+
 		//byt, _ := r.ReadBytes(3)
 		//fmt.Printf(">>>>>>>>>TightFillBytes=%v", byt)
 		return t, nil
@@ -98,11 +124,16 @@ func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangl
 		}
 
 		len, err := r.ReadCompactLen()
+		//binary.Write(bytes, binary.BigEndian, len)
 		if err != nil {
 			return nil, err
 		}
-		fmt.Printf("reading jpeg size=%d\n", len)
-		r.ReadBytes(len)
+		fmt.Printf("reading jpeg, size=%d\n", len)
+		_, err = r.ReadBytes(len)
+		if err != nil {
+			return nil, err
+		}
+		//StoreBytes(bytes, bts)
 		return t, nil
 	default:
 
@@ -111,6 +142,7 @@ func (t *TightEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangl
 		}
 
 		handleTightFilters(compctl, pixelFmt, rect, r)
+		//t.bytes = bytes.Bytes()
 		return t, nil
 	}
 }
@@ -126,6 +158,7 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 
 	if (subencoding & FILTER_ID_MASK) > 0 { // filter byte presence
 		filterid, err = r.ReadUint8()
+		//binary.Write(bytes, binary.BigEndian, filterid)
 		if err != nil {
 			fmt.Printf("error in handling tight encoding, reading filterid: %v\n", err)
 			return
@@ -146,10 +179,12 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 	case TightFilterPalette: //PALETTE_FILTER
 
 		colorCount, err := r.ReadUint8()
+		//binary.Write(bytes, binary.BigEndian, colorCount)
 		paletteSize := colorCount + 1 // add one more
 		fmt.Printf("----PALETTE_FILTER: paletteSize=%d bytesPixel=%d\n", paletteSize, bytesPixel)
 		//complete palette
-		r.ReadBytes(int(paletteSize) * bytesPixel)
+		_, err = r.ReadBytes(int(paletteSize) * bytesPixel)
+		//StoreBytes(bytes, bts)
 
 		var dataLength int
 		if paletteSize == 2 {
@@ -162,14 +197,26 @@ func handleTightFilters(subencoding uint8, pixelFmt *common.PixelFormat, rect *c
 			fmt.Printf("error in handling tight encoding, Reading Palette: %v\n", err)
 			return
 		}
+		//StoreBytes(bytes, bts)
+
 	case TightFilterGradient: //GRADIENT_FILTER
 		fmt.Printf("----GRADIENT_FILTER: bytesPixel=%d\n", bytesPixel)
 		//useGradient = true
 		fmt.Printf("usegrad: %d\n", filterid)
-		r.ReadTightData(lengthCurrentbpp)
+		_, err := r.ReadTightData(lengthCurrentbpp)
+		if err != nil {
+			fmt.Printf("error in handling tight encoding, Reading GRADIENT_FILTER: %v\n", err)
+			return
+		}
+		//StoreBytes(bytes, bts)
 	case TightFilterCopy: //BASIC_FILTER
 		fmt.Printf("----BASIC_FILTER: bytesPixel=%d\n", bytesPixel)
-		r.ReadTightData(lengthCurrentbpp)
+		_, err := r.ReadTightData(lengthCurrentbpp)
+		if err != nil {
+			fmt.Printf("error in handling tight encoding, Reading BASIC_FILTER: %v\n", err)
+			return
+		}
+		//StoreBytes(bytes, bts)
 	default:
 		fmt.Printf("Bad tight filter id: %d\n", filterid)
 		return

encodings/enc-tightpng.go

@@ -2,17 +2,27 @@ package encodings
 
 import (
 	"fmt"
+	"io"
 	"vncproxy/common"
 	"vncproxy/logger"
 )
 
 type TightPngEncoding struct {
+	bytes []byte
+}
+
+func (z *TightPngEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
 }
 
 func (*TightPngEncoding) Type() int32 { return int32(common.EncTightPng) }
 
 func (t *TightPngEncoding) Read(pixelFmt *common.PixelFormat, rect *common.Rectangle, r *common.RfbReadHelper) (common.Encoding, error) {
 	bytesPixel := calcTightBytePerPixel(pixelFmt)
+	r.StartByteCollection()
+	defer func() {
+		t.bytes = r.EndByteCollection()
+	}()
 
 	//var subencoding uint8
 	compctl, err := r.ReadUint8()

encodings/enc-zlib.go

@@ -1,23 +1,39 @@
 package encodings
 
-import "vncproxy/common"
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"vncproxy/common"
+)
 
 type ZLibEncoding struct {
 	//Colors []Color
+	bytes []byte
 }
 
 func (z *ZLibEncoding) Type() int32 {
 	return 6
 }
+func (z *ZLibEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
+}
 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, _ := r.ReadUint32()
-	_, err := r.ReadBytes(int(len))
-
+	bytes := &bytes.Buffer{}
+	len, err := r.ReadUint32()
 	if err != nil {
 		return nil, err
 	}
+
+	binary.Write(bytes, binary.BigEndian, len)
+	bts, err := r.ReadBytes(int(len))
+	if err != nil {
+		return nil, err
+	}
+	StoreBytes(bytes, bts)
+	z.bytes = bytes.Bytes()
 	return z, nil
 }

encodings/enc-zrle.go

@@ -1,23 +1,41 @@
 package encodings
 
-import "vncproxy/common"
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"vncproxy/common"
+)
 
 type ZRLEEncoding struct {
 	//Colors []Color
+	bytes []byte
 }
 
 func (z *ZRLEEncoding) Type() int32 {
 	return 16
 }
+
+func (z *ZRLEEncoding) WriteTo(w io.Writer) (n int, err error) {
+	return w.Write(z.bytes)
+}
+
 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, _ := r.ReadUint32()
-	_, err := r.ReadBytes(int(len))
-
+	bytes := &bytes.Buffer{}
+	len, err := r.ReadUint32()
 	if err != nil {
 		return nil, err
 	}
+
+	binary.Write(bytes, binary.BigEndian, len)
+	bts, err := r.ReadBytes(int(len))
+	if err != nil {
+		return nil, err
+	}
+	StoreBytes(bytes, bts)
+	z.bytes = bytes.Bytes()
 	return z, nil
 }

player/fbs-reader.go

@@ -1,25 +1,28 @@
 package player
 
 import (
+	"bytes"
 	"encoding/binary"
 	"io"
 	"os"
 	"vncproxy/common"
+	"vncproxy/encodings"
 	"vncproxy/logger"
-	"bytes"
 )
 
 type FbsReader struct {
-	reader io.Reader
-	buffer bytes.Buffer
+	reader           io.Reader
+	buffer           bytes.Buffer
 	currentTimestamp uint32
+	pixelFormat      *common.PixelFormat
+	encodings        []common.Encoding
 }
 
 func (fbs *FbsReader) Read(p []byte) (n int, err error) {
 	if fbs.buffer.Len() < len(p) {
 		seg, err := fbs.ReadSegment()
 		if err != nil {
-			logger.Error("FBSReader.Read: error reading FBSsegment: ",err)
+			logger.Error("FBSReader.Read: error reading FBSsegment: ", err)
 			return 0, err
 		}
 		fbs.buffer.Write(seg.bytes)
@@ -28,6 +31,10 @@ func (fbs *FbsReader) Read(p []byte) (n int, err error) {
 	return fbs.buffer.Read(p)
 }
 
+func (fbs *FbsReader) CurrentPixelFormat() *common.PixelFormat { return fbs.pixelFormat }
+func (fbs *FbsReader) CurrentColorMap() *common.ColorMap       { return &common.ColorMap{} }
+func (fbs *FbsReader) Encodings() []common.Encoding            { return fbs.encodings }
+
 func NewFbsReader(fbsFile string) (*FbsReader, error) {
 
 	reader, err := os.OpenFile(fbsFile, os.O_RDONLY, 0644)
@@ -35,7 +42,20 @@ func NewFbsReader(fbsFile string) (*FbsReader, error) {
 		logger.Error("NewFbsReader: can't open fbs file: ", fbsFile)
 		return nil, err
 	}
-	return &FbsReader{reader: reader}, nil
+	return &FbsReader{reader: reader,
+		encodings: []common.Encoding{
+			&encodings.CopyRectEncoding{},
+			&encodings.ZLibEncoding{},
+			&encodings.ZRLEEncoding{},
+			&encodings.CoRREEncoding{},
+			&encodings.HextileEncoding{},
+			&encodings.TightEncoding{},
+			&encodings.TightPngEncoding{},
+			&encodings.RawEncoding{},
+			&encodings.RREEncoding{},
+		},
+	}, nil
+
 }
 
 func (fbs *FbsReader) ReadStartSession() (*common.ServerInit, error) {
@@ -96,6 +116,7 @@ func (fbs *FbsReader) ReadStartSession() (*common.ServerInit, error) {
 	//read padding
 	bytes = make([]byte, 3)
 	fbs.Read(bytes)
+	fbs.pixelFormat = pixelFormat
 
 	//read desktop name
 	var desknameLen uint32

player/player_test.go

@@ -1,23 +1,36 @@
 package player
 
 import (
+	"encoding/binary"
 	"log"
 	"testing"
 	"time"
+	"vncproxy/client"
 	"vncproxy/common"
 	"vncproxy/encodings"
 	"vncproxy/logger"
 	"vncproxy/server"
-	"encoding/binary"
 )
 
 type ServerMessageHandler struct {
-	Conn         *server.ServerConn
-	Fbs          *FbsReader
-	firstSegDone bool
-	startTime    int
+	Conn             *server.ServerConn
+	Fbs              *FbsReader
+	serverMessageMap map[uint8]common.ServerMessage
+	firstSegDone     bool
+	startTime        int
 }
 
+func NewServerMessageHandler(conn *server.ServerConn, r *FbsReader) *ServerMessageHandler {
+	h := &ServerMessageHandler{Conn: conn, Fbs: r}
+	cm := client.BellMessage(0)
+	h.serverMessageMap = make(map[uint8]common.ServerMessage)
+	h.serverMessageMap[0] = &client.FramebufferUpdateMessage{}
+	h.serverMessageMap[1] = &client.SetColorMapEntriesMessage{}
+	h.serverMessageMap[2] = &cm
+	h.serverMessageMap[3] = &client.ServerCutTextMessage{}
+
+	return h
+}
 func (handler *ServerMessageHandler) Consume(seg *common.RfbSegment) error {
 
 	switch seg.SegmentType {
@@ -40,29 +53,27 @@ func (handler *ServerMessageHandler) Consume(seg *common.RfbSegment) error {
 
 func (h *ServerMessageHandler) sendFbsMessage() {
 	var messageType uint8
+	//messages := make(map[uint8]common.ServerMessage)
 	fbs := h.Fbs
 	//conn := h.Conn
-	binary.Read(fbs,binary.BigEndian,&messageType)
-	bytes := messages[messageType].Read(fbs)
-	h.Conn.Write(bytes)
-
-	//seg, err := fbs.ReadSegment()
-	//
-	//now := int(time.Now().UnixNano() / int64(time.Millisecond))
-	//if err != nil {
-	//	logger.Error("TestServer.NewConnHandler: Error in reading FBS segment: ", err)
-	//	return
-	//}
-	//timeSinceStart := now - h.startTime
-	//
-	//timeToWait := timeSinceStart - int(seg.timestamp)
-	//
-	//if timeToWait > 0 {
-	//	time.Sleep(time.Duration(timeToWait) * time.Millisecond)
-	//}
-	//fmt.Printf("bytes: %v", seg.bytes)
-	//conn.Write(seg.bytes)
-}	
+	err := binary.Read(fbs, binary.BigEndian, &messageType)
+	if err != nil {
+		logger.Error("TestServer.NewConnHandler: Error in reading FBS segment: ", err)
+		return
+	}
+	//common.IClientConn{}
+	binary.Write(h.Conn, binary.BigEndian, messageType)
+	msg := h.serverMessageMap[messageType]
+	if msg == nil {
+		logger.Error("TestServer.NewConnHandler: Error unknown message type: ", messageType)
+		return
+	}
+	err = msg.CopyTo(fbs, h.Conn, fbs)
+	if err != nil {
+		logger.Error("TestServer.NewConnHandler: Error in reading FBS segment: ", err)
+		return
+	}
+}
 
 func loadFbsFile(filename string, conn *server.ServerConn) (*FbsReader, error) {
 	fbs, err := NewFbsReader(filename)
@@ -109,7 +120,7 @@ func TestServer(t *testing.T) {
 			logger.Error("TestServer.NewConnHandler: Error in loading FBS: ", err)
 			return err
 		}
-		conn.Listeners.AddListener(&ServerMessageHandler{conn, fbs, false, 0})
+		conn.Listeners.AddListener(NewServerMessageHandler(conn, fbs))
 		return nil
 	}