From 297cd4cfe80b86ba7870d7d5ee64624ebd551ecb Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Wed, 27 Jun 2018 14:15:37 +0400
Subject: [PATCH] limit HTTP request body and WS read msg size to 1MB

---
 rpc/lib/server/handlers.go    |  8 +++++++-
 rpc/lib/server/http_server.go | 20 ++++++++++++++++++--
 types/params.go               |  1 +
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/rpc/lib/server/handlers.go b/rpc/lib/server/handlers.go
index 6cc03012a..dcacfb663 100644
--- a/rpc/lib/server/handlers.go
+++ b/rpc/lib/server/handlers.go
@@ -406,7 +406,13 @@ type wsConnection struct {
 // description of how to configure ping period and pong wait time. NOTE: if the
 // write buffer is full, pongs may be dropped, which may cause clients to
 // disconnect. see https://github.com/gorilla/websocket/issues/97
-func NewWSConnection(baseConn *websocket.Conn, funcMap map[string]*RPCFunc, cdc *amino.Codec, options ...func(*wsConnection)) *wsConnection {
+func NewWSConnection(
+	baseConn *websocket.Conn,
+	funcMap map[string]*RPCFunc,
+	cdc *amino.Codec,
+	options ...func(*wsConnection),
+) *wsConnection {
+	baseConn.SetReadLimit(maxBodyBytes)
 	wsc := &wsConnection{
 		remoteAddr:        baseConn.RemoteAddr().String(),
 		baseConn:          baseConn,
diff --git a/rpc/lib/server/http_server.go b/rpc/lib/server/http_server.go
index 9bdb4dffa..6223d205c 100644
--- a/rpc/lib/server/http_server.go
+++ b/rpc/lib/server/http_server.go
@@ -23,6 +23,12 @@ type Config struct {
 	MaxOpenConnections int
 }
 
+const (
+	// maxBodyBytes controls the maximum number of bytes the
+	// server will read parsing the request body.
+	maxBodyBytes = int64(1000000) // 1MB
+)
+
 // StartHTTPServer starts an HTTP server on listenAddr with the given handler.
 // It wraps handler with RecoverAndLogHandler.
 func StartHTTPServer(
@@ -53,7 +59,7 @@ func StartHTTPServer(
 	go func() {
 		err := http.Serve(
 			listener,
-			RecoverAndLogHandler(handler, logger),
+			RecoverAndLogHandler(maxBytesHandler{h: handler, n: maxBodyBytes}, logger),
 		)
 		logger.Error("RPC HTTP server stopped", "err", err)
 	}()
@@ -99,7 +105,7 @@ func StartHTTPAndTLSServer(
 	go func() {
 		err := http.ServeTLS(
 			listener,
-			RecoverAndLogHandler(handler, logger),
+			RecoverAndLogHandler(maxBytesHandler{h: handler, n: maxBodyBytes}, logger),
 			certFile,
 			keyFile,
 		)
@@ -202,3 +208,13 @@ func (w *ResponseWriterWrapper) WriteHeader(status int) {
 func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	return w.ResponseWriter.(http.Hijacker).Hijack()
 }
+
+type maxBytesHandler struct {
+	h http.Handler
+	n int64
+}
+
+func (h maxBytesHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	r.Body = http.MaxBytesReader(w, r.Body, h.n)
+	h.h.ServeHTTP(w, r)
+}
diff --git a/types/params.go b/types/params.go
index d068342c6..6cbac47a6 100644
--- a/types/params.go
+++ b/types/params.go
@@ -7,6 +7,7 @@ import (
 )
 
 const (
+	// MaxBlockSizeBytes is the maximum permitted size of the blocks.
 	MaxBlockSizeBytes = 104857600 // 100MB
 )