From ef4c4223d1f47a7512df74718f71227ec6bccdf5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Kripalani?= <raul@protocol.ai>
Date: Thu, 23 Apr 2020 10:44:44 +0100
Subject: [PATCH] propagate context deadlines to handshake.

Fixes https://github.com/libp2p/go-libp2p-noise/issues/75.
---
 p2p/security/noise/handshake.go      |  8 ++++++++
 p2p/security/noise/transport_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/p2p/security/noise/handshake.go b/p2p/security/noise/handshake.go
index 6231f8c7..92be729a 100644
--- a/p2p/security/noise/handshake.go
+++ b/p2p/security/noise/handshake.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"fmt"
+	"time"
 
 	"github.com/flynn/noise"
 	"github.com/gogo/protobuf/proto"
@@ -46,6 +47,13 @@ func (s *secureSession) runHandshake(ctx context.Context) error {
 		return err
 	}
 
+	// set a deadline to complete the handshake, if one has been supplied.
+	// clear it after we're done.
+	if deadline, ok := ctx.Deadline(); ok {
+		s.SetDeadline(deadline)
+		defer s.SetDeadline(time.Time{})
+	}
+
 	if s.initiator {
 		// stage 0 //
 		// do not send the payload just yet, as it would be plaintext; not secret.
diff --git a/p2p/security/noise/transport_test.go b/p2p/security/noise/transport_test.go
index 6e51e2a0..0bed8759 100644
--- a/p2p/security/noise/transport_test.go
+++ b/p2p/security/noise/transport_test.go
@@ -3,9 +3,11 @@ package noise
 import (
 	"bytes"
 	"context"
+	"errors"
 	"math/rand"
 	"net"
 	"testing"
+	"time"
 
 	crypto "github.com/libp2p/go-libp2p-core/crypto"
 	"github.com/libp2p/go-libp2p-core/peer"
@@ -86,6 +88,28 @@ func connect(t *testing.T, initTransport, respTransport *Transport) (*secureSess
 	return initConn.(*secureSession), respConn.(*secureSession)
 }
 
+func TestDeadlines(t *testing.T) {
+	initTransport := newTestTransport(t, crypto.Ed25519, 2048)
+	respTransport := newTestTransport(t, crypto.Ed25519, 2048)
+
+	init, resp := newConnPair(t)
+	defer init.Close()
+	defer resp.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	_, err := initTransport.SecureOutbound(ctx, init, respTransport.localID)
+	if err == nil {
+		t.Fatalf("expected i/o timeout err; got: %s", err)
+	}
+
+	var neterr net.Error
+	if ok := errors.As(err, &neterr); !ok || !neterr.Timeout() {
+		t.Fatalf("expected i/o timeout err; got: %s", err)
+	}
+}
+
 func TestIDs(t *testing.T) {
 	initTransport := newTestTransport(t, crypto.Ed25519, 2048)
 	respTransport := newTestTransport(t, crypto.Ed25519, 2048)