add request body size limit (#913)

ontio · May 21, 2019 · 75960f1 · 75960f1
1 parent 6686600
commit 75960f1
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 15 deletions.
diff --git a/http/base/common/common.go b/http/base/common/common.go
@@ -44,6 +44,7 @@ import (
 )
 
 const MAX_SEARCH_HEIGHT uint32 = 100
+const MAX_REQUEST_BODY_SIZE = 1 << 20
 
 type BalanceOfRsp struct {
 	Ont string `json:"ont"`

diff --git a/http/base/rpc/rpc.go b/http/base/rpc/rpc.go
@@ -23,7 +23,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/ontio/ontology/common/log"
+	"github.com/ontio/ontology/http/base/common"
 	berr "github.com/ontio/ontology/http/base/error"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"os"
@@ -79,7 +81,6 @@ func Handle(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
-
 	//check if there is Request Body to read
 	if r.Body == nil {
 		if mainMux.defaultFunction != nil {
@@ -91,15 +92,10 @@ func Handle(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 	}
-
-	//read the body of the request
-	body, err := ioutil.ReadAll(r.Body)
-	if err != nil {
-		log.Error("HTTP JSON RPC Handle - ioutil.ReadAll: ", err)
-		return
-	}
 	request := make(map[string]interface{})
-	err = json.Unmarshal(body, &request)
+	defer r.Body.Close()
+	decoder := json.NewDecoder(io.LimitReader(r.Body, common.MAX_REQUEST_BODY_SIZE))
+	err := decoder.Decode(&request)
 	if err != nil {
 		log.Error("HTTP JSON RPC Handle - json.Unmarshal: ", err)
 		return

diff --git a/http/restful/restful/server.go b/http/restful/restful/server.go
@@ -25,10 +25,11 @@ import (
 	"encoding/json"
 	cfg "github.com/ontio/ontology/common/config"
 	"github.com/ontio/ontology/common/log"
+	"github.com/ontio/ontology/http/base/common"
 	berr "github.com/ontio/ontology/http/base/error"
 	"github.com/ontio/ontology/http/base/rest"
 	"golang.org/x/net/netutil"
-	"io/ioutil"
+	"io"
 	"net"
 	"net/http"
 	"strconv"
@@ -271,16 +272,14 @@ func (this *restServer) initGetHandler() {
 func (this *restServer) initPostHandler() {
 	for k, _ := range this.postMap {
 		this.router.Post(k, func(w http.ResponseWriter, r *http.Request) {
-
-			body, _ := ioutil.ReadAll(r.Body)
+			decoder := json.NewDecoder(io.LimitReader(r.Body, common.MAX_REQUEST_BODY_SIZE))
 			defer r.Body.Close()
-
 			var req = make(map[string]interface{})
 			var resp map[string]interface{}
 
 			url := this.getPath(r.URL.Path)
 			if h, ok := this.postMap[url]; ok {
-				if err := json.Unmarshal(body, &req); err == nil {
+				if err := decoder.Decode(&req); err == nil {
 					req = this.getParams(r, url, req)
 					resp = h.handler(req)
 					resp["Action"] = h.name

diff --git a/http/websocket/websocket/server.go b/http/websocket/websocket/server.go
@@ -252,7 +252,7 @@ func (self *WsServer) checkSessionsTimeout(done chan bool) {
 
 func (self *WsServer) webSocketHandler(w http.ResponseWriter, r *http.Request) {
 	wsConn, err := self.Upgrader.Upgrade(w, r, nil)
-
+	wsConn.SetReadLimit(1024 * 1024)
 	if err != nil {
 		log.Error("websocket Upgrader: ", err)
 		return