From c8905c0ea947d85d7c1ad78fdf9b524699b385eb Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@selfip.ru>
Date: Sun, 18 Jun 2017 21:01:01 +0300
Subject: [PATCH] limit buffer usage

Signed-off-by: Vasiliy Tolstov <v.tolstov@selfip.ru>
---
 encoding.go           | 7 +------
 example/proxy/main.go | 6 ++++++
 image.go              | 6 ++----
 server.go             | 5 +----
 4 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/encoding.go b/encoding.go
index e5b70b2..384e926 100644
--- a/encoding.go
+++ b/encoding.go
@@ -244,16 +244,11 @@ func (enc *TightPngEncoding) Read(c Conn, rect *Rectangle) error {
 	cmp := enc.TightCC.Compression
 	switch cmp {
 	case TightCompressionPNG:
-		buf := bytes.NewBuffer(nil)
 		l, err := readTightLength(c)
 		if err != nil {
 			return err
 		}
-		_, err = io.CopyN(buf, c, int64(l))
-		if err != nil {
-			return err
-		}
-		enc.Image, err = png.Decode(buf)
+		enc.Image, err = png.Decode(io.LimitReader(c, int64(l)))
 		if err != nil {
 			return err
 		}
diff --git a/example/proxy/main.go b/example/proxy/main.go
index 8fc24e6..ed5a9f0 100644
--- a/example/proxy/main.go
+++ b/example/proxy/main.go
@@ -4,11 +4,17 @@ import (
 	"context"
 	"log"
 	"net"
+	"net/http"
+	_ "net/http/pprof"
 
 	vnc "github.com/vtolstov/go-vnc"
 )
 
 func main() {
+	go func() {
+		log.Println(http.ListenAndServe(":6060", nil))
+	}()
+
 	ln, err := net.Listen("tcp", ":5900")
 	if err != nil {
 		log.Fatalf("Error listen. %v", err)
diff --git a/image.go b/image.go
index c795a8b..6c1d806 100644
--- a/image.go
+++ b/image.go
@@ -114,6 +114,7 @@ func colorsToImage(x, y, width, height uint16, colors []Color) *image.RGBA64 {
 
 // Marshal implements the Marshaler interface.
 func (r *Rectangle) Write(c Conn) error {
+
 	var err error
 	if err = binary.Write(c, binary.BigEndian, r.X); err != nil {
 		return err
@@ -131,10 +132,7 @@ func (r *Rectangle) Write(c Conn) error {
 		return err
 	}
 
-	if err := r.Enc.Write(c, r); err != nil {
-		return err
-	}
-	return c.Flush()
+	return r.Enc.Write(c, r)
 }
 
 func (r *Rectangle) Read(c Conn) error {
diff --git a/server.go b/server.go
index a21849c..cd3dc9b 100644
--- a/server.go
+++ b/server.go
@@ -416,10 +416,7 @@ func (*Bell) Read(c Conn) (ServerMessage, error) {
 }
 
 func (msg *Bell) Write(c Conn) error {
-	if err := binary.Write(c, binary.BigEndian, msg.Type()); err != nil {
-		return err
-	}
-	return c.Flush()
+	return binary.Write(c, binary.BigEndian, msg.Type())
 }
 
 type SetColorMapEntries struct {