les: clean up server handler (#22357)

les/server_handler.go

@@ -204,59 +204,27 @@ func (h *serverHandler) handle(p *clientPeer) error {
 	}
 }
 
-// handleMsg is invoked whenever an inbound message is received from a remote
+// beforeHandle will do a series of prechecks before handling message.
-// peer. The remote connection is torn down upon returning any error.
+func (h *serverHandler) beforeHandle(p *clientPeer, reqID, responseCount uint64, msg p2p.Msg, reqCnt uint64, maxCount uint64) (*servingTask, uint64) {
-func (h *serverHandler) handleMsg(p *clientPeer, wg *sync.WaitGroup) error {
+	// Ensure that the request sent by client peer is valid
-	// Read the next message from the remote peer, and ensure it's fully consumed
-	msg, err := p.rw.ReadMsg()
-	if err != nil {
-		return err
-	}
-	p.Log().Trace("Light Ethereum message arrived", "code", msg.Code, "bytes", msg.Size)
-
-	// Discard large message which exceeds the limitation.
-	if msg.Size > ProtocolMaxMsgSize {
-		clientErrorMeter.Mark(1)
-		return errResp(ErrMsgTooLarge, "%v > %v", msg.Size, ProtocolMaxMsgSize)
-	}
-	defer msg.Discard()
-
-	p.responseCount++
-	responseCount := p.responseCount
-
-	req, ok := Les3[msg.Code]
-	if !ok {
-		p.Log().Trace("Received invalid message", "code", msg.Code)
-		clientErrorMeter.Mark(1)
-		return errResp(ErrInvalidMsgCode, "%v", msg.Code)
-	}
-	p.Log().Trace("Received " + req.Name)
-
-	serve, reqID, reqCnt, err := req.Handle(msg)
-	if err != nil {
-		clientErrorMeter.Mark(1)
-		return errResp(ErrDecode, "%v: %v", msg, err)
-	}
-
-	if metrics.EnabledExpensive {
-		req.InPacketsMeter.Mark(1)
-		req.InTrafficMeter.Mark(int64(msg.Size))
-	}
-
-	// Short circuit if the peer is already frozen or the request is invalid.
 	inSizeCost := h.server.costTracker.realCost(0, msg.Size, 0)
-	if p.isFrozen() || reqCnt == 0 || reqCnt > req.MaxCount {
+	if reqCnt == 0 || reqCnt > maxCount {
 		p.fcClient.OneTimeCost(inSizeCost)
-		return nil
+		return nil, 0
+	}
+	// Ensure that the client peer complies with the flow control
+	// rules agreed by both sides.
+	if p.isFrozen() {
+		p.fcClient.OneTimeCost(inSizeCost)
+		return nil, 0
 	}
-	// Prepaid max cost units before request been serving.
 	maxCost := p.fcCosts.getMaxCost(msg.Code, reqCnt)
 	accepted, bufShort, priority := p.fcClient.AcceptRequest(reqID, responseCount, maxCost)
 	if !accepted {
 		p.freeze()
 		p.Log().Error("Request came too early", "remaining", common.PrettyDuration(time.Duration(bufShort*1000000/p.fcParams.MinRecharge)))
 		p.fcClient.OneTimeCost(inSizeCost)
-		return nil
+		return nil, 0
 	}
 	// Create a multi-stage task, estimate the time it takes for the task to
 	// execute, and cache it in the request service queue.
@@ -267,15 +235,19 @@ func (h *serverHandler) handleMsg(p *clientPeer, wg *sync.WaitGroup) error {
 	}
 	maxTime := uint64(float64(maxCost) / factor)
 	task := h.server.servingQueue.newTask(p, maxTime, priority)
-	if task.start() {
+	if !task.start() {
-		wg.Add(1)
+		p.fcClient.RequestProcessed(reqID, responseCount, maxCost, inSizeCost)
-		go func() {
+		return nil, 0
-			defer wg.Done()
+	}
-			reply := serve(h, p, task.waitOrStop)
+	return task, maxCost
+}
+
+// Afterhandle will perform a series of operations after message handling,
+// such as updating flow control data, sending reply, etc.
+func (h *serverHandler) afterHandle(p *clientPeer, reqID, responseCount uint64, msg p2p.Msg, maxCost uint64, reqCnt uint64, task *servingTask, reply *reply) {
 	if reply != nil {
 		task.done()
 	}
-
 	p.responseLock.Lock()
 	defer p.responseLock.Unlock()
 
@@ -313,17 +285,72 @@ func (h *serverHandler) handleMsg(p *clientPeer, wg *sync.WaitGroup) error {
 				}
 			}
 		})
+	}
+}
+
+// handleMsg is invoked whenever an inbound message is received from a remote
+// peer. The remote connection is torn down upon returning any error.
+func (h *serverHandler) handleMsg(p *clientPeer, wg *sync.WaitGroup) error {
+	// Read the next message from the remote peer, and ensure it's fully consumed
+	msg, err := p.rw.ReadMsg()
+	if err != nil {
+		return err
+	}
+	p.Log().Trace("Light Ethereum message arrived", "code", msg.Code, "bytes", msg.Size)
+
+	// Discard large message which exceeds the limitation.
+	if msg.Size > ProtocolMaxMsgSize {
+		clientErrorMeter.Mark(1)
+		return errResp(ErrMsgTooLarge, "%v > %v", msg.Size, ProtocolMaxMsgSize)
+	}
+	defer msg.Discard()
+
+	// Lookup the request handler table, ensure it's supported
+	// message type by the protocol.
+	req, ok := Les3[msg.Code]
+	if !ok {
+		p.Log().Trace("Received invalid message", "code", msg.Code)
+		clientErrorMeter.Mark(1)
+		return errResp(ErrInvalidMsgCode, "%v", msg.Code)
+	}
+	p.Log().Trace("Received " + req.Name)
+
+	// Decode the p2p message, resolve the concrete handler for it.
+	serve, reqID, reqCnt, err := req.Handle(msg)
+	if err != nil {
+		clientErrorMeter.Mark(1)
+		return errResp(ErrDecode, "%v: %v", msg, err)
+	}
 	if metrics.EnabledExpensive {
+		req.InPacketsMeter.Mark(1)
+		req.InTrafficMeter.Mark(int64(msg.Size))
+	}
+	p.responseCount++
+	responseCount := p.responseCount
+
+	// First check this client message complies all rules before
+	// handling it and return a processor if all checks are passed.
+	task, maxCost := h.beforeHandle(p, reqID, responseCount, msg, reqCnt, req.MaxCount)
+	if task == nil {
+		return nil
+	}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		reply := serve(h, p, task.waitOrStop)
+		h.afterHandle(p, reqID, responseCount, msg, maxCost, reqCnt, task, reply)
+
+		if metrics.EnabledExpensive {
+			size := uint32(0)
+			if reply != nil {
+				size = reply.size()
+			}
 			req.OutPacketsMeter.Mark(1)
-					req.OutTrafficMeter.Mark(int64(replySize))
+			req.OutTrafficMeter.Mark(int64(size))
 			req.ServingTimeMeter.Update(time.Duration(task.servingTime))
 		}
-			}
 	}()
-	} else {
-		p.fcClient.RequestProcessed(reqID, responseCount, maxCost, inSizeCost)
-	}
-
 	// If the client has made too much invalid request(e.g. request a non-existent data),
 	// reject them to prevent SPAM attack.
 	if p.getInvalid() > maxRequestErrors {

les/server_requests.go

@@ -65,7 +65,7 @@ type serveRequestFn func(backend serverBackend, peer *clientPeer, waitOrStop fun
 
 // Les3 contains the request types supported by les/2 and les/3
 var Les3 = map[uint64]RequestType{
-	GetBlockHeadersMsg: RequestType{
+	GetBlockHeadersMsg: {
 		Name:             "block header request",
 		MaxCount:         MaxHeaderFetch,
 		InPacketsMeter:   miscInHeaderPacketsMeter,
@@ -75,7 +75,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeHeaderTimer,
 		Handle:           handleGetBlockHeaders,
 	},
-	GetBlockBodiesMsg: RequestType{
+	GetBlockBodiesMsg: {
 		Name:             "block bodies request",
 		MaxCount:         MaxBodyFetch,
 		InPacketsMeter:   miscInBodyPacketsMeter,
@@ -85,7 +85,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeBodyTimer,
 		Handle:           handleGetBlockBodies,
 	},
-	GetCodeMsg: RequestType{
+	GetCodeMsg: {
 		Name:             "code request",
 		MaxCount:         MaxCodeFetch,
 		InPacketsMeter:   miscInCodePacketsMeter,
@@ -95,7 +95,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeCodeTimer,
 		Handle:           handleGetCode,
 	},
-	GetReceiptsMsg: RequestType{
+	GetReceiptsMsg: {
 		Name:             "receipts request",
 		MaxCount:         MaxReceiptFetch,
 		InPacketsMeter:   miscInReceiptPacketsMeter,
@@ -105,7 +105,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeReceiptTimer,
 		Handle:           handleGetReceipts,
 	},
-	GetProofsV2Msg: RequestType{
+	GetProofsV2Msg: {
 		Name:             "les/2 proofs request",
 		MaxCount:         MaxProofsFetch,
 		InPacketsMeter:   miscInTrieProofPacketsMeter,
@@ -115,7 +115,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeTrieProofTimer,
 		Handle:           handleGetProofs,
 	},
-	GetHelperTrieProofsMsg: RequestType{
+	GetHelperTrieProofsMsg: {
 		Name:             "helper trie proof request",
 		MaxCount:         MaxHelperTrieProofsFetch,
 		InPacketsMeter:   miscInHelperTriePacketsMeter,
@@ -125,7 +125,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeHelperTrieTimer,
 		Handle:           handleGetHelperTrieProofs,
 	},
-	SendTxV2Msg: RequestType{
+	SendTxV2Msg: {
 		Name:             "new transactions",
 		MaxCount:         MaxTxSend,
 		InPacketsMeter:   miscInTxsPacketsMeter,
@@ -135,7 +135,7 @@ var Les3 = map[uint64]RequestType{
 		ServingTimeMeter: miscServingTimeTxTimer,
 		Handle:           handleSendTx,
 	},
-	GetTxStatusMsg: RequestType{
+	GetTxStatusMsg: {
 		Name:             "transaction status query request",
 		MaxCount:         MaxTxStatus,
 		InPacketsMeter:   miscInTxStatusPacketsMeter,