From 460f205112c933692f31eee6e5e5296347ba33fc Mon Sep 17 00:00:00 2001
From: sukun <sukunrt@gmail.com>
Date: Wed, 2 Oct 2024 16:41:39 +0530
Subject: [PATCH] autonatv2: catch panics

---
 p2p/protocol/autonatv2/autonat.go      |  4 ++--
 p2p/protocol/autonatv2/autonat_test.go |  1 -
 p2p/protocol/autonatv2/client.go       |  8 ++++++++
 p2p/protocol/autonatv2/server.go       | 11 ++++++++++-
 4 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/p2p/protocol/autonatv2/autonat.go b/p2p/protocol/autonatv2/autonat.go
index 3dc3bc166a..888a95b7ae 100644
--- a/p2p/protocol/autonatv2/autonat.go
+++ b/p2p/protocol/autonatv2/autonat.go
@@ -25,9 +25,9 @@ const (
 	DialProtocol     = "/libp2p/autonat/2/dial-request"
 
 	maxMsgSize            = 8192
-	streamTimeout         = time.Minute
+	streamTimeout         = 15 * time.Second
 	dialBackStreamTimeout = 5 * time.Second
-	dialBackDialTimeout   = 30 * time.Second
+	dialBackDialTimeout   = 10 * time.Second
 	dialBackMaxMsgSize    = 1024
 	minHandshakeSizeBytes = 30_000 // for amplification attack prevention
 	maxHandshakeSizeBytes = 100_000
diff --git a/p2p/protocol/autonatv2/autonat_test.go b/p2p/protocol/autonatv2/autonat_test.go
index ee97ccc695..11c8f02195 100644
--- a/p2p/protocol/autonatv2/autonat_test.go
+++ b/p2p/protocol/autonatv2/autonat_test.go
@@ -657,5 +657,4 @@ func TestAreAddrsConsistency(t *testing.T) {
 			}
 		})
 	}
-
 }
diff --git a/p2p/protocol/autonatv2/client.go b/p2p/protocol/autonatv2/client.go
index c1d82b7f3a..b9e88ec414 100644
--- a/p2p/protocol/autonatv2/client.go
+++ b/p2p/protocol/autonatv2/client.go
@@ -3,6 +3,8 @@ package autonatv2
 import (
 	"context"
 	"fmt"
+	"os"
+	"runtime/debug"
 	"sync"
 	"time"
 
@@ -248,6 +250,12 @@ func newDialRequest(reqs []Request, nonce uint64) pb.Message {
 
 // handleDialBack receives the nonce on the dial-back stream
 func (ac *client) handleDialBack(s network.Stream) {
+	defer func() {
+		if rerr := recover(); rerr != nil {
+			fmt.Fprintf(os.Stderr, "caught panic: %s\n%s\n", rerr, debug.Stack())
+		}
+	}()
+
 	if err := s.Scope().SetService(ServiceName); err != nil {
 		log.Debugf("failed to attach stream to service %s: %w", ServiceName, err)
 		s.Reset()
diff --git a/p2p/protocol/autonatv2/server.go b/p2p/protocol/autonatv2/server.go
index c5dd2dce3e..f2b9bb6429 100644
--- a/p2p/protocol/autonatv2/server.go
+++ b/p2p/protocol/autonatv2/server.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
+	"runtime/debug"
 	"sync"
 	"time"
 
@@ -88,6 +90,13 @@ func (as *server) Close() {
 
 // handleDialRequest is the dial-request protocol stream handler
 func (as *server) handleDialRequest(s network.Stream) {
+	defer func() {
+		if rerr := recover(); rerr != nil {
+			fmt.Fprintf(os.Stderr, "caught panic: %s\n%s\n", rerr, debug.Stack())
+		}
+	}()
+
+	log.Debugf("received dial-request from: %s, addr: %s", s.Conn().RemotePeer(), s.Conn().RemoteMultiaddr())
 	evt := as.serveDialRequest(s)
 	log.Debugf("completed dial-request from %s, response status: %s, dial status: %s, err: %s",
 		s.Conn().RemotePeer(), evt.ResponseStatus, evt.DialStatus, evt.Error)
@@ -96,7 +105,7 @@ func (as *server) handleDialRequest(s network.Stream) {
 	}
 }
 
-func (as *server) serveDialRequest(s network.Stream) EventDialRequestCompleted {
+func (as *server) serveDialRequest(s network.Stream) (result EventDialRequestCompleted) {
 	if err := s.Scope().SetService(ServiceName); err != nil {
 		s.Reset()
 		log.Debugf("failed to attach stream to %s service: %w", ServiceName, err)