address comments

tikv · Jun 23, 2017 · 7f39fc3 · 7f39fc3
1 parent e2f56dd
commit 7f39fc3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/server/api/history_test.go b/server/api/history_test.go
@@ -14,14 +14,14 @@
 package api
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/pd/server"
+	"golang.org/x/net/context"
 )
 
 var _ = Suite(&testHistorySuite{})

diff --git a/server/grpc_service.go b/server/grpc_service.go
@@ -244,7 +244,10 @@ func (s *Server) RegionHeartbeat(server pdpb.PD_RegionHeartbeatServer) error {
 			return errors.Trace(err)
 		}
 
-		// TODO: should we check headers here?
+		if err := s.validateRequest(request.GetHeader()); err != nil {
+			// TODO: How to close this stream?
+			return errors.Trace(err)
+		}
 
 		cluster := s.GetRaftCluster()
 		if cluster == nil {