catch panics when reading / writing protobuf messages (#31)

libp2p · Apr 14, 2022 · 93c76c2 · 93c76c2
1 parent 05788ec
commit 93c76c2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/protoio/uvarint_reader.go b/protoio/uvarint_reader.go
@@ -34,6 +34,7 @@ package protoio
 
 import (
 	"bufio"
+	"fmt"
 	"io"
 
 	"github.com/gogo/protobuf/proto"
@@ -56,7 +57,13 @@ func NewDelimitedReader(r io.Reader, maxSize int) ReadCloser {
 	return &uvarintReader{bufio.NewReader(r), nil, maxSize, closer}
 }
 
-func (ur *uvarintReader) ReadMsg(msg proto.Message) error {
+func (ur *uvarintReader) ReadMsg(msg proto.Message) (err error) {
+	defer func() {
+		if rerr := recover(); rerr != nil {
+			err = fmt.Errorf("panic reading message: %s", rerr)
+		}
+	}()
+
 	length64, err := varint.ReadUvarint(ur.r)
 	if err != nil {
 		return err

diff --git a/protoio/uvarint_writer.go b/protoio/uvarint_writer.go
@@ -33,6 +33,7 @@
 package protoio
 
 import (
+	"fmt"
 	"io"
 
 	"github.com/gogo/protobuf/proto"
@@ -51,6 +52,12 @@ func NewDelimitedWriter(w io.Writer) WriteCloser {
 }
 
 func (uw *uvarintWriter) WriteMsg(msg proto.Message) (err error) {
+	defer func() {
+		if rerr := recover(); rerr != nil {
+			err = fmt.Errorf("panic reading message: %s", rerr)
+		}
+	}()
+
 	var data []byte
 	if m, ok := msg.(interface {
 		MarshalTo(data []byte) (n int, err error)