fix stream deadlines

client.go

@@ -3,23 +3,17 @@ package autonat
 import (
 	"context"
 	"fmt"
+	"time"
 
+	pb "github.com/libp2p/go-libp2p-autonat/pb"
 	"github.com/libp2p/go-libp2p-core/host"
 	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-msgio/protoio"
 
-	pb "github.com/libp2p/go-libp2p-autonat/pb"
-
-	protoio "github.com/libp2p/go-msgio/protoio"
 	ma "github.com/multiformats/go-multiaddr"
 )
 
-// Error wraps errors signalled by AutoNAT services
-type Error struct {
-	Status pb.Message_ResponseStatus
-	Text   string
-}
-
 // NewAutoNATClient creates a fresh instance of an AutoNATClient
 // If addrFunc is nil, h.Addrs will be used
 func NewAutoNATClient(h host.Host, addrFunc AddrFunc) Client {
@@ -34,11 +28,14 @@ type client struct {
 	addrFunc AddrFunc
 }
 
+// DialBack asks peer p to dial us back on all addresses returned by the addrFunc.
+// It blocks until we've received a response from the peer.
 func (c *client) DialBack(ctx context.Context, p peer.ID) (ma.Multiaddr, error) {
 	s, err := c.h.NewStream(ctx, p, AutoNATProto)
 	if err != nil {
 		return nil, err
 	}
+	s.SetDeadline(time.Now().Add(streamTimeout))
 	// Might as well just reset the stream. Once we get to this point, we
 	// don't care about being nice.
 	defer s.Close()
@@ -47,20 +44,18 @@ func (c *client) DialBack(ctx context.Context, p peer.ID) (ma.Multiaddr, error)
 	w := protoio.NewDelimitedWriter(s)
 
 	req := newDialMessage(peer.AddrInfo{ID: c.h.ID(), Addrs: c.addrFunc()})
-	err = w.WriteMsg(req)
-	if err != nil {
+	if err := w.WriteMsg(req); err != nil {
 		s.Reset()
 		return nil, err
 	}
 
 	var res pb.Message
-	err = r.ReadMsg(&res)
-	if err != nil {
+	if err := r.ReadMsg(&res); err != nil {
 		s.Reset()
 		return nil, err
 	}
-
 	if res.GetType() != pb.Message_DIAL_RESPONSE {
+		s.Reset()
 		return nil, fmt.Errorf("unexpected response: %s", res.GetType().String())
 	}
 
@@ -69,12 +64,17 @@ func (c *client) DialBack(ctx context.Context, p peer.ID) (ma.Multiaddr, error)
 	case pb.Message_OK:
 		addr := res.GetDialResponse().GetAddr()
 		return ma.NewMultiaddrBytes(addr)
-
 	default:
 		return nil, Error{Status: status, Text: res.GetDialResponse().GetStatusText()}
 	}
 }
 
+// Error wraps errors signalled by AutoNAT services
+type Error struct {
+	Status pb.Message_ResponseStatus
+	Text   string
+}
+
 func (e Error) Error() string {
 	return fmt.Sprintf("AutoNAT error: %s (%s)", e.Text, e.Status.String())
 }

svc.go

@@ -7,17 +7,16 @@ import (
 	"sync"
 	"time"
 
+	pb "github.com/libp2p/go-libp2p-autonat/pb"
 	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
 
-	pb "github.com/libp2p/go-libp2p-autonat/pb"
-
 	"github.com/libp2p/go-msgio/protoio"
 	ma "github.com/multiformats/go-multiaddr"
 )
 
-var streamReadTimeout = 60 * time.Second
+var streamTimeout = 60 * time.Second
 
 // AutoNATService provides NAT autodetection services to other peers
 type autoNATService struct {
@@ -49,8 +48,7 @@ func newAutoNATService(ctx context.Context, c *config) (*autoNATService, error)
 }
 
 func (as *autoNATService) handleStream(s network.Stream) {
-	s.SetReadDeadline(time.Now().Add(streamReadTimeout))
-
+	s.SetDeadline(time.Now().Add(streamTimeout))
 	defer s.Close()
 
 	pid := s.Conn().RemotePeer()