diff --git a/client/server-messages.go b/client/server-messages.go index 935e60b..7c70ceb 100644 --- a/client/server-messages.go +++ b/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()) } diff --git a/common/encoding.go b/common/encoding.go index 495aff7..87ade9e 100644 --- a/common/encoding.go +++ b/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. diff --git a/common/readers.go b/common/readers.go index d92f37c..e21e3f1 100644 --- a/common/readers.go +++ b/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) diff --git a/common/server-message.go b/common/server-message.go index c1b5c9a..71ac106 100644 --- a/common/server-message.go +++ b/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: diff --git a/encodings/enc-copy-rect.go b/encodings/enc-copy-rect.go index 05d372f..d812f7d 100644 --- a/encodings/enc-copy-rect.go +++ b/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() diff --git a/encodings/enc-corre.go b/encodings/enc-corre.go index bcce323..70ba424 100644 --- a/encodings/enc-corre.go +++ b/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); - } diff --git a/encodings/enc-hextile.go b/encodings/enc-hextile.go index b171d67..c928522 100644 --- a/encodings/enc-hextile.go +++ b/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 { diff --git a/encodings/enc-pseudo.go b/encodings/enc-pseudo.go index 32430b5..eaa873a 100644 --- a/encodings/enc-pseudo.go +++ b/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 } diff --git a/encodings/enc-raw.go b/encodings/enc-raw.go index a462ab0..7d843ef 100644 --- a/encodings/enc-raw.go +++ b/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 } diff --git a/encodings/enc-rre.go b/encodings/enc-rre.go index c1feeba..6acebae 100644 --- a/encodings/enc-rre.go +++ b/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 } diff --git a/encodings/enc-tight.go b/encodings/enc-tight.go index 0646aed..91ba74b 100644 --- a/encodings/enc-tight.go +++ b/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 diff --git a/encodings/enc-tightpng.go b/encodings/enc-tightpng.go index 3d20a2d..4894114 100644 --- a/encodings/enc-tightpng.go +++ b/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() diff --git a/encodings/enc-zlib.go b/encodings/enc-zlib.go index a7bfaab..0847af7 100644 --- a/encodings/enc-zlib.go +++ b/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 } diff --git a/encodings/enc-zrle.go b/encodings/enc-zrle.go index 58f4f9d..a87b721 100644 --- a/encodings/enc-zrle.go +++ b/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 } diff --git a/player/fbs-reader.go b/player/fbs-reader.go index 865fc6f..0e2ae66 100644 --- a/player/fbs-reader.go +++ b/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 diff --git a/player/player_test.go b/player/player_test.go index 07bbadc..e0837af 100644 --- a/player/player_test.go +++ b/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 }