Merge pull request #5 from libp2p/use-context

close the underlying connection when the context is canceled

transport.go

@@ -43,23 +43,51 @@ var _ cs.Transport = &Transport{}
 // SecureInbound runs the TLS handshake as a server.
 func (t *Transport) SecureInbound(ctx context.Context, insecure net.Conn) (cs.Conn, error) {
 	serv := tls.Server(insecure, t.identity.Config)
-	// TODO: use the ctx
-	// see https://github.com/golang/go/issues/18482
-	if err := serv.Handshake(); err != nil {
-		return nil, err
-	}
-	return t.setupConn(serv)
+	return t.handshake(ctx, insecure, serv)
 }
 
 // SecureOutbound runs the TLS handshake as a client.
 func (t *Transport) SecureOutbound(ctx context.Context, insecure net.Conn, p peer.ID) (cs.Conn, error) {
 	cl := tls.Client(insecure, t.identity.ConfigForPeer(p))
-	// TODO: use the ctx
-	// see https://github.com/golang/go/issues/18482
-	if err := cl.Handshake(); err != nil {
+	return t.handshake(ctx, insecure, cl)
+}
+
+func (t *Transport) handshake(
+	ctx context.Context,
+	// in Go 1.10, we need to close the underlying net.Conn
+	// in Go 1.11 this was fixed, and tls.Conn.Close() works as well
+	insecure net.Conn,
+	tlsConn *tls.Conn,
+) (cs.Conn, error) {
+	// There's no way to pass a context to tls.Conn.Handshake().
+	// See https://github.com/golang/go/issues/18482.
+	// Close the connection instead.
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		select {
+		case <-done:
+		case <-ctx.Done():
+			insecure.Close()
+		}
+	}()
+
+	if err := tlsConn.Handshake(); err != nil {
+		// if the context was canceled, return the context error
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, ctxErr
+		}
+		return nil, err
+	}
+	conn, err := t.setupConn(tlsConn)
+	if err != nil {
+		// if the context was canceled, return the context error
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, ctxErr
+		}
 		return nil, err
 	}
-	return t.setupConn(cl)
+	return conn, nil
 }
 
 func (t *Transport) setupConn(tlsConn *tls.Conn) (cs.Conn, error) {

transport_test.go

@@ -96,6 +96,44 @@ var _ = Describe("Transport", func() {
 		Expect(string(b)).To(Equal("foobar"))
 	})
 
+	It("fails when the context of the outgoing connection is canceled", func() {
+		clientTransport, err := New(clientKey)
+		Expect(err).ToNot(HaveOccurred())
+		serverTransport, err := New(serverKey)
+		Expect(err).ToNot(HaveOccurred())
+
+		clientInsecureConn, serverInsecureConn := connect()
+
+		go func() {
+			defer GinkgoRecover()
+			_, err := serverTransport.SecureInbound(context.Background(), serverInsecureConn)
+			Expect(err).To(HaveOccurred())
+		}()
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		_, err = clientTransport.SecureOutbound(ctx, clientInsecureConn, serverID)
+		Expect(err).To(MatchError(context.Canceled))
+	})
+
+	It("fails when the context of the incoming connection is canceled", func() {
+		clientTransport, err := New(clientKey)
+		Expect(err).ToNot(HaveOccurred())
+		serverTransport, err := New(serverKey)
+		Expect(err).ToNot(HaveOccurred())
+
+		clientInsecureConn, serverInsecureConn := connect()
+
+		go func() {
+			defer GinkgoRecover()
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+			_, err := serverTransport.SecureInbound(ctx, serverInsecureConn)
+			Expect(err).To(MatchError(context.Canceled))
+		}()
+		_, err = clientTransport.SecureOutbound(context.Background(), clientInsecureConn, serverID)
+		Expect(err).To(HaveOccurred())
+	})
+
 	It("fails if the peer ID doesn't match", func() {
 		thirdPartyID, _ := createPeer()