From 902ef26b8460936d94e49b64ccdb2dccba8bf71a Mon Sep 17 00:00:00 2001 From: amit bezalel Date: Wed, 24 Jan 2018 09:24:36 +0200 Subject: [PATCH] added zrle --- README.md | 12 +- encoding_util.go | 30 +++++ encoding_zrle.go | 325 +++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 352 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 30a574b..a55f32d 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ # A **real wold** implementation of vnc client for go -After searching the web for an vnc client in golang which is not a toy & support more than handshake + RAW encoding, I came up blank, so, I set out to write it myself. +After searching the web for an vnc client in golang which is not a toy & support more than handshake + RAW encoding, I came up blank, so, I set out to write one myself. + +The video encoding part means that something can be viewed, and since I don't really feel like writing GTK UIs in 2018 (plus VNC viewers are a dime a dozen), a video file will do. +In actuality the images produced are go image structs and can easily be saved as JPEG, or displayed in any UI you want to create. + ## Encoding support: * Tight VNC * Hextile @@ -7,7 +11,7 @@ After searching the web for an vnc client in golang which is not a toy & support * CopyRect * Raw * RRE -* ZRLE [**TBD - coming soon**] +* ZRLE * Rich-cursor pseudo * Desktop Size Pseudo * Cursor pos Pseudo @@ -24,10 +28,10 @@ Since go has no good client UI library I chose to encode video instead, but the * This allows recording vnc without the cost of video encoding while retaining the ability to have video later if the vnc session is marked as important. ## About -It may seem strange that i didn't use my previous vncproxy code in order to create this client, but since that code is highly optimized to be a proxy (never hold a full message in buffer & introduce no lags), it is not best suited to be a client, so instead of spending the time reverting all the proxy-specific code, I just started from the most advanced go vnc-client code I found. +It may seem strange that I didn't use my previous vncproxy code in order to create this client, but since that code is highly optimized to be a proxy (never hold a full message in buffer & introduce no lags), it is not best suited to be a client, so instead of spending the time reverting all the proxy-specific code, I just started from the most advanced go vnc-client code I found. Most of what I added is the rfb-encoder & video encoding implementations, there are naturally some additional changes in order to get a global canvas (draw.Image) to render on by all encodings. -The code for the encodings was gathered by looking at several RFB source codes in cpp & some in java, reading the excellent documentation in [rfbproto](https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst), and **a lot** of gritty bit-plucking, pixel jogging & code cajoling until everything fell into place on screen. +The code for the encodings was gathered by peeking at several RFB source codes in cpp & some in java, reading the excellent documentation in [rfbproto](https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst), and **a lot** of gritty bit-plucking, pixel jogging & code cajoling until everything fell into place on screen. I did not include tightPng in the supported encoding list since I didn't find a server to test it with, so I can't vouch for the previous implementation, If you have such a server handy, please check and tell me if it works. \ No newline at end of file diff --git a/encoding_util.go b/encoding_util.go index 9686ebf..76b9003 100644 --- a/encoding_util.go +++ b/encoding_util.go @@ -7,8 +7,16 @@ import ( "image/color" "image/draw" "io" + "vnc2video/logger" ) +func Min(a, b int) int { + if a < b { + return a + } + return b +} + func FillRect(img draw.Image, rect *image.Rectangle, c color.Color) { for x := rect.Min.X; x < rect.Max.X; x++ { for y := rect.Min.Y; y < rect.Max.Y; y++ { @@ -17,6 +25,28 @@ func FillRect(img draw.Image, rect *image.Rectangle, c color.Color) { } } +func readRunLength(r io.Reader) (int, error) { + runLen := 1 + + mod, err := ReadUint8(r) + if err != nil { + logger.Errorf("renderZRLE: error while reading mod in plain RLE subencoding: %v", err) + return 0, err + } + runLen += int(mod) + + for mod == 255 { + //mod = fromZlib.read(); + mod, err = ReadUint8(r) + if err != nil { + logger.Errorf("renderZRLE: error while reading mod in-loop plain RLE subencoding: %v", err) + return 0, err + } + runLen += int(mod) + } + return runLen, nil +} + // Read unmarshal color from conn func ReadColor(c io.Reader, pf *PixelFormat) (*color.RGBA, error) { if pf.TrueColor == 0 { diff --git a/encoding_zrle.go b/encoding_zrle.go index 4df2a33..ad5c889 100644 --- a/encoding_zrle.go +++ b/encoding_zrle.go @@ -2,36 +2,339 @@ package vnc2video import ( "bytes" - "encoding/binary" + "compress/zlib" + "errors" + "image/color" + "image/draw" "io" + "vnc2video/logger" ) type ZRLEEncoding struct { - bytes []byte + bytes []byte + Image draw.Image + unzipper io.Reader + zippedBuff *bytes.Buffer } -func (z *ZRLEEncoding) Type() int32 { - return 16 +func (*ZRLEEncoding) Supported(Conn) bool { + return true } +func (enc *ZRLEEncoding) SetTargetImage(img draw.Image) { + enc.Image = img +} + +func (enc *ZRLEEncoding) Reset() error { + enc.unzipper = nil + return nil +} + +func (*ZRLEEncoding) Type() EncodingType { return EncZRLE } + func (z *ZRLEEncoding) WriteTo(w io.Writer) (n int, err error) { return w.Write(z.bytes) } -func (z *ZRLEEncoding) Read(r Conn, rect *Rectangle) error { - //func (z *ZRLEEncoding) Read(pixelFmt *PixelFormat, rect *Rectangle, r io.Reader) (Encoding, error) { - bytes := &bytes.Buffer{} +func (enc *ZRLEEncoding) Write(c Conn, rect *Rectangle) error { + return nil +} + +func IsCPixelSpecific(pf *PixelFormat) bool { + significant := int(pf.RedMax<= 2 && subEnc <= 16: + err = enc.handlePaletteTile(tileOffsetX, tileOffsetY, tileWidth, tileHeight, subEnc, pf, rect) + if err != nil { + return err + } + case subEnc == 128: + err = enc.handlePlainRLETile(tileOffsetX, tileOffsetY, tileWidth, tileHeight, pf, rect) + if err != nil { + return err + } + case subEnc >= 130 && subEnc <= 255: + err = enc.handlePaletteRLETile(tileOffsetX, tileOffsetY, tileWidth, tileHeight, subEnc, pf, rect) + if err != nil { + return err + } + default: + logger.Errorf("Unknown ZRLE subencoding: %v", subEnc) + } + } + } + return nil +} + +func (enc *ZRLEEncoding) handlePaletteRLETile(tileOffsetX, tileOffsetY, tileWidth, tileHeight int, subEnc uint8, pf *PixelFormat, rect *Rectangle) error { + // Palette RLE + paletteSize := subEnc - 128 + palette := make([]*color.RGBA, paletteSize) + var err error + + // Read RLE palette + for j := 0; j < int(paletteSize); j++ { + palette[j], err = readCPixel(enc.unzipper, pf) + if err != nil { + logger.Errorf("renderZRLE: error while reading color in palette RLE subencoding: %v", err) + return err + } + } + var index uint8 + runLen := 0 + for y := 0; y < tileHeight; y++ { + for x := 0; x < tileWidth; x++ { + + if runLen == 0 { + + // Read length and index + index, err = ReadUint8(enc.unzipper) + if err != nil { + logger.Errorf("renderZRLE: error while reading length and index in palette RLE subencoding: %v", err) + //return err + } + runLen = 1 + + // Run is represented by index | 0x80 + // Otherwise, single pixel + if (index & 0x80) != 0 { + + index -= 128 + + runLen, err = readRunLength(enc.unzipper) + if err != nil { + logger.Errorf("handlePlainRLETile: error while reading runlength in plain RLE subencoding: %v", err) + return err + } + + } + //logger.Debugf("renderZRLE: writing pixel: col=%v times=%d", palette[index], runLen) + } + + // Write pixel to image + enc.Image.Set(tileOffsetX+int(rect.X)+x, tileOffsetY+int(rect.Y)+y, palette[index]) + runLen-- + } + } + return nil +} + +func (enc *ZRLEEncoding) handlePaletteTile(tileOffsetX, tileOffsetY, tileWidth, tileHeight int, subEnc uint8, pf *PixelFormat, rect *Rectangle) error { + //subenc here is also palette size + paletteSize := subEnc + palette := make([]*color.RGBA, paletteSize) + var err error + // Read palette + for j := 0; j < int(paletteSize); j++ { + palette[j], err = readCPixel(enc.unzipper, pf) + if err != nil { + logger.Errorf("renderZRLE: error while reading CPixel for palette tile: %v", err) + return err + } + } + // Calculate index size + var indexBits, mask uint32 + if paletteSize == 2 { + indexBits = 1 + mask = 0x80 + } else if paletteSize <= 4 { + indexBits = 2 + mask = 0xC0 + } else { + indexBits = 4 + mask = 0xF0 + } + for y := 0; y < tileHeight; y++ { + + // Packing only occurs per-row + bitsAvailable := uint32(0) + buffer := uint32(0) + + for x := 0; x < tileWidth; x++ { + + // Buffer more bits if necessary + if bitsAvailable == 0 { + bits, err := ReadUint8(enc.unzipper) + if err != nil { + logger.Errorf("renderZRLE: error while reading first uint8 into buffer: %v", err) + return err + } + buffer = uint32(bits) + bitsAvailable = 8 + } + + // Read next pixel + index := (buffer & mask) >> (8 - indexBits) + buffer <<= indexBits + bitsAvailable -= indexBits + + // Write pixel to image + enc.Image.Set(tileOffsetX+int(rect.X), tileOffsetY+int(rect.Y), palette[index]) + } + } + return err +} + +func (enc *ZRLEEncoding) handlePlainRLETile(tileOffsetX int, tileOffsetY int, tileWidth int, tileHeight int, pf *PixelFormat, rect *Rectangle) error { + var col *color.RGBA + var err error + runLen := 0 + for y := 0; y < tileHeight; y++ { + for x := 0; x < tileWidth; x++ { + + if runLen == 0 { + + // Read length and color + col, err = readCPixel(enc.unzipper, pf) + if err != nil { + logger.Errorf("handlePlainRLETile: error while reading CPixel in plain RLE subencoding: %v", err) + return err + } + runLen, err = readRunLength(enc.unzipper) + if err != nil { + logger.Errorf("handlePlainRLETile: error while reading runlength in plain RLE subencoding: %v", err) + return err + } + + } + + // Write pixel to image + enc.Image.Set(tileOffsetX+int(rect.X)+x, tileOffsetY+int(rect.Y)+y, col) + runLen-- + } + } + return err +} + +// Read unmarshal color from conn +func readCPixel(c io.Reader, pf *PixelFormat) (*color.RGBA, error) { + if pf.TrueColor == 0 { + return nil, errors.New("support for non true color formats was not implemented") + } + + isZRLEFormat := IsCPixelSpecific(pf) + var col *color.RGBA + if isZRLEFormat { + //tbytes := make([]byte, 3) + tbytes, err := ReadBytes(3, c) + if err != nil { + return nil, err + } + + if pf.BigEndian != 1 { + col = &color.RGBA{ + B: uint8(tbytes[0]), + G: uint8(tbytes[1]), + R: uint8(tbytes[2]), + A: uint8(1), + } + } else { + col = &color.RGBA{ + R: uint8(tbytes[0]), + G: uint8(tbytes[1]), + B: uint8(tbytes[2]), + A: uint8(1), + } + } + return col, nil + } + + col, err := ReadColor(c, pf) + if err != nil { + logger.Errorf("readCPixel: Error while reading zrle: %v", err) + } + + return col, nil +}