fix(shwap/shrex): enable recovery middleware in shrex server (#3897)

celestiaorg · Oct 30, 2024 · d5eea79 · d5eea79
1 parent 9ff5857
commit d5eea79
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/share/shwap/p2p/shrex/shrexeds/server.go b/share/shwap/p2p/shrex/shrexeds/server.go
@@ -31,7 +31,8 @@ type Server struct {
 
 	store *store.Store
 
-	params     *Parameters
+	params *Parameters
+	// TODO: decouple middleware metrics from shrex and remove middleware from Server
 	middleware *shrex.Middleware
 	metrics    *shrex.Metrics
 }
@@ -55,7 +56,10 @@ func (s *Server) Start(context.Context) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	s.cancel = cancel
 
-	s.host.SetStreamHandler(s.protocolID, s.middleware.RateLimitHandler(s.streamHandler(ctx)))
+	handler := s.streamHandler(ctx)
+	withRateLimit := s.middleware.RateLimitHandler(handler)
+	withRecovery := shrex.RecoveryMiddleware(withRateLimit)
+	s.host.SetStreamHandler(s.protocolID, withRecovery)
 	return nil
 }
 

diff --git a/share/shwap/p2p/shrex/shrexnd/server.go b/share/shwap/p2p/shrex/shrexnd/server.go
@@ -32,7 +32,8 @@ type Server struct {
 	handler network.StreamHandler
 	store   *store.Store
 
-	params     *Parameters
+	params *Parameters
+	// TODO: decouple middleware metrics from shrex and remove middleware from Server
 	middleware *shrex.Middleware
 	metrics    *shrex.Metrics
 }
@@ -54,7 +55,10 @@ func NewServer(params *Parameters, host host.Host, store *store.Store) (*Server,
 	ctx, cancel := context.WithCancel(context.Background())
 	srv.cancel = cancel
 
-	srv.handler = srv.middleware.RateLimitHandler(srv.streamHandler(ctx))
+	handler := srv.streamHandler(ctx)
+	withRateLimit := srv.middleware.RateLimitHandler(handler)
+	withRecovery := shrex.RecoveryMiddleware(withRateLimit)
+	srv.handler = withRecovery
 	return srv, nil
 }