From d1a4710aa29ecb9bcc116fb84e8cade73bcb639b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Oliveirinha?= <joliveirinha@cloudflare.com>
Date: Mon, 11 Apr 2022 11:44:42 +0100
Subject: [PATCH] TUN-6035: Reduce buffer size when proxying data

---
 cfio/copy.go           | 27 +++++++++++++++++++++++++++
 websocket/websocket.go |  4 +++-
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 cfio/copy.go

diff --git a/cfio/copy.go b/cfio/copy.go
new file mode 100644
index 00000000..6d837d5f
--- /dev/null
+++ b/cfio/copy.go
@@ -0,0 +1,27 @@
+package cfio
+
+import (
+	"io"
+	"sync"
+)
+
+const defaultBufferSize = 16 * 1024
+
+var bufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, defaultBufferSize)
+	},
+}
+
+func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
+	_, okWriteTo := src.(io.WriterTo)
+	_, okReadFrom := dst.(io.ReaderFrom)
+	var buffer []byte = nil
+
+	if !(okWriteTo || okReadFrom) {
+		buffer = bufferPool.Get().([]byte)
+		defer bufferPool.Put(buffer)
+	}
+
+	return io.CopyBuffer(dst, src, buffer)
+}
diff --git a/websocket/websocket.go b/websocket/websocket.go
index ee0bd225..0c879681 100644
--- a/websocket/websocket.go
+++ b/websocket/websocket.go
@@ -15,6 +15,8 @@ import (
 	"github.com/getsentry/raven-go"
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
+
+	"github.com/cloudflare/cloudflared/cfio"
 )
 
 // IsWebSocketUpgrade checks to see if the request is a WebSocket connection.
@@ -146,7 +148,7 @@ func copyData(dst io.Writer, src io.Reader, dir string) (written int64, err erro
 		}
 		return copyBuffer(dst, src, dir)
 	} else {
-		return io.Copy(dst, src)
+		return cfio.Copy(dst, src)
 	}
 }