diff --git a/rpc/server.go b/rpc/server.go index 2742adf07b..a97c927757 100644 --- a/rpc/server.go +++ b/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(), - codecs: make(map[ServerCodec]struct{}), + idgen: randomIDGenerator(), + 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 {