From 1bce3011250ec59e9cad7572bf4cb536fb5c0a11 Mon Sep 17 00:00:00 2001 From: Bezalel Date: Wed, 1 Aug 2018 02:04:30 +0300 Subject: [PATCH] added keyframes & an rfbFile writer implementation (initial) --- logger/logger.go | 2 +- player/rfb-reader.go | 199 ++++++++++++++++++++++++++++++++++++++ recorder/cmd/main.go | 6 +- recorder/recorder.go | 5 + recorder/rfb-requester.go | 28 ++++-- 5 files changed, 232 insertions(+), 8 deletions(-) create mode 100644 player/rfb-reader.go diff --git a/logger/logger.go b/logger/logger.go index 5b3fecf..7c4a3a5 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -2,7 +2,7 @@ package logger import "fmt" -var simpleLogger = SimpleLogger{LogLevelWarn} +var simpleLogger = SimpleLogger{LogLevelInfo} type Logger interface { Debug(v ...interface{}) diff --git a/player/rfb-reader.go b/player/rfb-reader.go new file mode 100644 index 0000000..39a41fd --- /dev/null +++ b/player/rfb-reader.go @@ -0,0 +1,199 @@ +package player + +import ( + "bytes" + "encoding/binary" + "io" + "os" + "vncproxy/common" + "vncproxy/encodings" + "vncproxy/logger" +) + +type RfbReader struct { + reader io.Reader + buffer bytes.Buffer + currentTimestamp int + pixelFormat *common.PixelFormat + encodings []common.IEncoding +} + +/************************************************************** +** RFB File documentation: +** Sections: +** 0. header: + * index seek position + * +** 1. init message +** 2. content + * frame message: + * size, timestamp, type, content +** 3. index: + * each frame message start position, full/incremental, timestamp + * +***************************************************************/ + +func (rfb *RfbReader) CurrentTimestamp() int { + return rfb.currentTimestamp +} + +func (rfb *RfbReader) Read(p []byte) (n int, err error) { + if rfb.buffer.Len() < len(p) { + seg, err := rfb.ReadSegment() + + if err != nil { + logger.Error("rfbReader.Read: error reading rfbsegment: ", err) + return 0, err + } + rfb.buffer.Write(seg.bytes) + rfb.currentTimestamp = int(seg.timestamp) + } + return rfb.buffer.Read(p) +} + +func (rfb *RfbReader) CurrentPixelFormat() *common.PixelFormat { return rfb.pixelFormat } + +//func (rfb *rfbReader) CurrentColorMap() *common.ColorMap { return &common.ColorMap{} } +func (rfb *RfbReader) Encodings() []common.IEncoding { return rfb.encodings } + +func NewRfbReader(rfbFile string) (*RfbReader, error) { + + reader, err := os.OpenFile(rfbFile, os.O_RDONLY, 0644) + if err != nil { + logger.Error("NewrfbReader: can't open rfb file: ", rfbFile) + return nil, err + } + return &RfbReader{reader: reader, + encodings: []common.IEncoding{ + &encodings.CopyRectEncoding{}, + &encodings.ZLibEncoding{}, + &encodings.ZRLEEncoding{}, + &encodings.CoRREEncoding{}, + &encodings.HextileEncoding{}, + &encodings.TightEncoding{}, + &encodings.TightPngEncoding{}, + &encodings.EncCursorPseudo{}, + &encodings.RawEncoding{}, + &encodings.RREEncoding{}, + }, + }, nil + +} + +func (rfb *RfbReader) ReadStartSession() (*common.ServerInit, error) { + + initMsg := common.ServerInit{} + reader := rfb.reader + + var framebufferWidth uint16 + var framebufferHeight uint16 + var SecTypeNone uint32 + //read rfb header information (the only part done without the [size|data|timestamp] block wrapper) + //.("rfb 001.000\n") + bytes := make([]byte, 12) + _, err := reader.Read(bytes) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init message - rfb file Version:", err) + return nil, err + } + + //read the version message into the buffer so it will be written in the first rbs block + //RFB 003.008\n + bytes = make([]byte, 12) + _, err = rfb.Read(bytes) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - RFB Version: ", err) + return nil, err + } + + //push sec type and fb dimensions + binary.Read(rfb, binary.BigEndian, &SecTypeNone) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - SecType: ", err) + } + + //read frame buffer width, height + binary.Read(rfb, binary.BigEndian, &framebufferWidth) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - FBWidth: ", err) + return nil, err + } + initMsg.FBWidth = framebufferWidth + + binary.Read(rfb, binary.BigEndian, &framebufferHeight) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - FBHeight: ", err) + return nil, err + } + initMsg.FBHeight = framebufferHeight + + //read pixel format + pixelFormat := &common.PixelFormat{} + binary.Read(rfb, binary.BigEndian, pixelFormat) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - Pixelformat: ", err) + return nil, err + } + initMsg.PixelFormat = *pixelFormat + //read padding + bytes = make([]byte, 3) + rfb.Read(bytes) + rfb.pixelFormat = pixelFormat + + //read desktop name + var desknameLen uint32 + binary.Read(rfb, binary.BigEndian, &desknameLen) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - deskname Len: ", err) + return nil, err + } + initMsg.NameLength = desknameLen + + bytes = make([]byte, desknameLen) + rfb.Read(bytes) + if err != nil { + logger.Error("rfbReader.ReadStartSession: error reading rbs init - desktopName: ", err) + return nil, err + } + + initMsg.NameText = bytes + + return &initMsg, nil +} + +func (rfb *RfbReader) ReadSegment() (*FbsSegment, error) { + reader := rfb.reader + var bytesLen uint32 + + //read length + err := binary.Read(reader, binary.BigEndian, &bytesLen) + if err != nil { + logger.Error("rfbReader.ReadStartSession: read len, error reading rbs file: ", err) + return nil, err + } + + paddedSize := (bytesLen + 3) & 0x7FFFFFFC + + //read bytes + bytes := make([]byte, paddedSize) + _, err = reader.Read(bytes) + if err != nil { + logger.Error("rfbReader.ReadSegment: read bytes, error reading rbs file: ", err) + return nil, err + } + + //remove padding + actualBytes := bytes[:bytesLen] + + //read timestamp + var timeSinceStart uint32 + binary.Read(reader, binary.BigEndian, &timeSinceStart) + if err != nil { + logger.Error("rfbReader.ReadSegment: read timestamp, error reading rbs file: ", err) + return nil, err + } + + //timeStamp := time.Unix(timeSinceStart, 0) + seg := &FbsSegment{bytes: actualBytes, timestamp: timeSinceStart} + return seg, nil +} diff --git a/recorder/cmd/main.go b/recorder/cmd/main.go index f2c64f5..ea01e53 100644 --- a/recorder/cmd/main.go +++ b/recorder/cmd/main.go @@ -67,7 +67,11 @@ func main() { }) clientConn.Listeners.AddListener(rec) - clientConn.Listeners.AddListener(&recorder.RfbRequester{Conn: clientConn, Name: "Rfb Requester"}) + clientConn.Listeners.AddListener(&recorder.RfbRequester{ + Conn: clientConn, + Name: "Rfb Requester", + FullScreenRefreshInSec: 30, //create a full refresh key frame every 30sec for seeking + }) clientConn.Connect() if err != nil { diff --git a/recorder/recorder.go b/recorder/recorder.go index 28dfd0e..3a44f22 100644 --- a/recorder/recorder.go +++ b/recorder/recorder.go @@ -189,6 +189,11 @@ func (r *Recorder) writeToDisk() error { paddedSize := (bytesLen + 3) & 0x7FFFFFFC paddingSize := paddedSize - bytesLen + /// KeyFramePos, _ := r.writer.Seek(0, os.SEEK_CUR) + /// fi, err := r.writer.Stat() + /// KeyFramePos := fi.Size() + KeyFramePosInBuffer + // now save the KF pos in some file + //logger.Debugf("paddedSize=%d paddingSize=%d bytesLen=%d", paddedSize, paddingSize, bytesLen) //write buffer padded to 32bit _, err := r.buffer.WriteTo(r.writer) diff --git a/recorder/rfb-requester.go b/recorder/rfb-requester.go index b72806c..6bdf2ee 100644 --- a/recorder/rfb-requester.go +++ b/recorder/rfb-requester.go @@ -8,11 +8,13 @@ import ( ) type RfbRequester struct { - Conn *client.ClientConn - Name string - Width uint16 - Height uint16 - lastRequestTime time.Time + Conn *client.ClientConn + Name string + Width uint16 + Height uint16 + lastRequestTime time.Time + nextFullScreenRefresh time.Time + FullScreenRefreshInSec int // refresh interval (creates keyframes) if 0, disables keyframe creation } func (p *RfbRequester) Consume(seg *common.RfbSegment) error { @@ -29,6 +31,7 @@ func (p *RfbRequester) Consume(seg *common.RfbSegment) error { p.Height = serverInitMessage.FBHeight p.lastRequestTime = time.Now() p.Conn.FramebufferUpdateRequest(false, 0, 0, p.Width, p.Height) + p.nextFullScreenRefresh = time.Now().Add(time.Duration(p.FullScreenRefreshInSec) * time.Second) case common.SegmentMessageStart: case common.SegmentRectSeparator: @@ -39,7 +42,20 @@ func (p *RfbRequester) Consume(seg *common.RfbSegment) error { // timeForNextReq := p.lastRequestTime.Unix() + minTimeBetweenReq.Nanoseconds()/1000 // if seg.UpcomingObjectType == int(common.FramebufferUpdate) && time.Now().Unix() > timeForNextReq { //time.Sleep(300 * time.Millisecond) - p.Conn.FramebufferUpdateRequest(true, 0, 0, p.Width, p.Height) + p.lastRequestTime = time.Now() + incremental := true + + if p.FullScreenRefreshInSec > 0 { + // if p.nextFullScreenRefresh.IsZero() { + // p.nextFullScreenRefresh = time.Now().Add(time.Duration(p.FullScreenRefreshInSec) * time.Second) + // } + if time.Now().Sub(p.nextFullScreenRefresh) <= 0 { + logger.Warn(">>Creating keyframe") + p.nextFullScreenRefresh = time.Now().Add(time.Duration(p.FullScreenRefreshInSec) * time.Second) + incremental = false + } + } + p.Conn.FramebufferUpdateRequest(incremental, 0, 0, p.Width, p.Height) //} default: }