From 5ffd3453580a018e8cb1423f9fc4f1d5042e66d9 Mon Sep 17 00:00:00 2001
From: HAOYUatHZ <37070449+HAOYUatHZ@users.noreply.github.com>
Date: Mon, 20 Nov 2023 16:39:12 +0800
Subject: [PATCH] feat: http&ws compression (#569)

---
 rpc/http.go      | 28 +++++++++++++++++++++++++++-
 rpc/server.go    | 14 ++++++++++++++
 rpc/websocket.go | 14 ++++++++++----
 3 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/rpc/http.go b/rpc/http.go
index 741fa1c0eb..b1a1cac3b6 100644
--- a/rpc/http.go
+++ b/rpc/http.go
@@ -18,6 +18,8 @@ package rpc
 
 import (
 	"bytes"
+	"compress/gzip"
+	"compress/zlib"
 	"context"
 	"encoding/json"
 	"errors"
@@ -28,6 +30,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 )
@@ -243,7 +246,8 @@ func (hc *httpConn) doRequest(ctx context.Context, msg interface{}) (io.ReadClos
 			Body:       body,
 		}
 	}
-	return resp.Body, nil
+	// use encoding if set
+	return newDecodeCompression(resp.Header.Get("Content-Encoding"), resp.Body)
 }
 
 // httpServerConn turns a HTTP connection into a Conn.
@@ -253,6 +257,28 @@ type httpServerConn struct {
 	r *http.Request
 }
 
+func newDecodeCompression(decoding string, rc io.ReadCloser) (io.ReadCloser, error) {
+	tps := strings.Split(strings.TrimSpace(strings.ToLower(decoding)), ",")
+	var res io.ReadCloser
+	switch tps[0] {
+	case "gzip":
+		gz, err := gzip.NewReader(rc)
+		if err != nil {
+			return nil, err
+		}
+		res = gz
+	case "deflate":
+		zl, err := zlib.NewReader(rc)
+		if err != nil {
+			return nil, err
+		}
+		res = zl
+	default:
+		res = rc
+	}
+	return res, nil
+}
+
 func newHTTPServerConn(r *http.Request, w http.ResponseWriter) ServerCodec {
 	body := io.LimitReader(r.Body, maxRequestContentLength)
 	conn := &httpServerConn{Reader: body, Writer: w, r: r}
diff --git a/rpc/server.go b/rpc/server.go
index 2742adf07b..3e4c8f4a90 100644
--- a/rpc/server.go
+++ b/rpc/server.go
@@ -17,7 +17,9 @@
 package rpc
 
 import (
+	"compress/flate"
 	"context"
+	"errors"
 	"io"
 	"sync"
 	"sync/atomic"
@@ -51,6 +53,9 @@ type Server struct {
 	run                atomic.Bool
 	batchItemLimit     int
 	batchResponseLimit int
+
+	// Add compressionLevel inorder to enable set it when open websocket server.
+	compressionLevel int
 }
 
 // NewServer creates a new server instance with no registered handlers.
@@ -127,6 +132,15 @@ func (s *Server) untrackCodec(codec ServerCodec) {
 	delete(s.codecs, codec)
 }
 
+// SetCompressionLevel set compression level (-2 ~ 9), this function only works on websocket.
+func (s *Server) SetCompressionLevel(level int) error {
+	if !(flate.HuffmanOnly <= level && level <= flate.BestCompression) {
+		return errors.New("websocket: invalid compression level")
+	}
+	s.compressionLevel = level
+	return nil
+}
+
 // serveSingleRequest reads and processes a single RPC request from the given codec. This
 // is used to serve HTTP connections. Subscriptions and reverse calls are not allowed in
 // this mode.
diff --git a/rpc/websocket.go b/rpc/websocket.go
index 23c97c89f7..913a67d912 100644
--- a/rpc/websocket.go
+++ b/rpc/websocket.go
@@ -17,6 +17,7 @@
 package rpc
 
 import (
+	"compress/flate"
 	"context"
 	"encoding/base64"
 	"fmt"
@@ -48,11 +49,13 @@ var wsBufferPool = new(sync.Pool)
 // allowedOrigins should be a comma-separated list of allowed origin URLs.
 // To allow connections with any origin, pass "*".
 func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
+	enableCompression := s.compressionLevel != flate.NoCompression
 	var upgrader = websocket.Upgrader{
-		ReadBufferSize:  wsReadBuffer,
-		WriteBufferSize: wsWriteBuffer,
-		WriteBufferPool: wsBufferPool,
-		CheckOrigin:     wsHandshakeValidator(allowedOrigins),
+		EnableCompression: enableCompression,
+		ReadBufferSize:    wsReadBuffer,
+		WriteBufferSize:   wsWriteBuffer,
+		WriteBufferPool:   wsBufferPool,
+		CheckOrigin:       wsHandshakeValidator(allowedOrigins),
 	}
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		conn, err := upgrader.Upgrade(w, r, nil)
@@ -60,6 +63,9 @@ func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
 			log.Debug("WebSocket upgrade failed", "err", err)
 			return
 		}
+		if enableCompression {
+			_ = conn.SetCompressionLevel(s.compressionLevel)
+		}
 		codec := newWebsocketCodec(conn, r.Host, r.Header, wsDefaultReadLimit)
 		s.ServeCodec(codec, 0)
 	})