Merge tag 'v1.13.5-sei-11' into extended-tracer-sei

rpc/server.go

@@ -48,6 +48,7 @@ type Server struct {
 
 	mutex              sync.Mutex
 	codecs             map[ServerCodec]struct{}
+	denyList           map[string]struct{}
 	run                atomic.Bool
 	batchItemLimit     int
 	batchResponseLimit int
@@ -56,8 +57,9 @@ type Server struct {
 // NewServer creates a new server instance with no registered handlers.
 func NewServer() *Server {
 	server := &Server{
-		idgen:  randomIDGenerator(),
+		idgen:    randomIDGenerator(),
-		codecs: make(map[ServerCodec]struct{}),
+		codecs:   make(map[ServerCodec]struct{}),
+		denyList: make(map[string]struct{}),
 	}
 	server.run.Store(true)
 	// Register the default service providing meta information about the RPC service such
@@ -86,6 +88,12 @@ func (s *Server) RegisterName(name string, receiver interface{}) error {
 	return s.services.registerName(name, receiver)
 }
 
+// RegisterDenyList add given method name to the deny list so that RPC requests that matches
+// any of the methods in deny list will got rejected directly.
+func (s *Server) RegisterDenyList(methodName string) {
+	s.denyList[methodName] = struct{}{}
+}
+
 // ServeCodec reads incoming requests from codec, calls the appropriate callback and writes
 // the response back using the given codec. It will block until the codec is closed or the
 // server is stopped. In either case the codec is closed.
@@ -141,6 +149,7 @@ func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {
 	defer h.close(io.EOF, nil)
 
 	reqs, batch, err := codec.readBatch()
+
 	if err != nil {
 		if err != io.EOF {
 			resp := errorMessage(&invalidMessageError{"parse error"})
@@ -148,6 +157,15 @@ func (s *Server) serveSingleRequest(ctx context.Context, codec ServerCodec) {
 		}
 		return
 	}
+	// check deny list
+	for _, req := range reqs {
+		method := req.Method
+		if _, found := s.denyList[method]; found {
+			resp := errorMessage(&methodNotFoundError{method: method})
+			codec.writeJSON(ctx, resp, true)
+			return
+		}
+	}
 	if batch {
 		h.handleBatch(reqs)
 	} else {