Merge pull request #17 from sei-protocol/yzang/SEI-6785

Add method deny list in geth http server
This commit is contained in:
Yiming Zang 2024-03-12 16:33:10 +08:00 committed by GitHub
commit db3a09495b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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
@ -58,6 +59,7 @@ func NewServer() *Server {
server := &Server{
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 {