From 73463cd208c92d180cd4c0977299d1e18eced54d Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Wed, 8 Jun 2016 16:12:06 -0700
Subject: [PATCH] respect contexts while reading messages in dht

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>
---
 routing/dht/dht_net.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/routing/dht/dht_net.go b/routing/dht/dht_net.go
index 0152dab4a91..8ad4286ce5f 100644
--- a/routing/dht/dht_net.go
+++ b/routing/dht/dht_net.go
@@ -214,7 +214,7 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
 	log.Event(ctx, "dhtSentMessage", ms.dht.self, ms.p, pmes)
 
 	mes := new(pb.Message)
-	if err := ms.r.ReadMsg(mes); err != nil {
+	if err := ms.ctxReadMsg(ctx, mes); err != nil {
 		ms.s.Close()
 		ms.s = nil
 		return nil, err
@@ -227,3 +227,17 @@ func (ms *messageSender) SendRequest(ctx context.Context, pmes *pb.Message) (*pb
 
 	return mes, nil
 }
+
+func (ms *messageSender) ctxReadMsg(ctx context.Context, mes *pb.Message) error {
+	errc := make(chan error, 1)
+	go func() {
+		errc <- ms.r.ReadMsg(mes)
+	}()
+
+	select {
+	case err := <-errc:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}