fix Client.Exchange not canceling when context was canceled

Fixes layeh#26

client.go

@@ -61,37 +61,42 @@ func (c *Client) Exchange(ctx context.Context, packet *Packet, addr string) (*Pa
 	}
 	defer conn.Close()
 
-	if deadline, deadlineSet := ctx.Deadline(); deadlineSet {
-		conn.SetDeadline(deadline)
-	}
-
 	conn.Write(wire)
 
+	var cancel context.CancelFunc
+	ctx, cancel = context.WithCancel(ctx)
+	defer cancel()
+
+	var retryTimer <-chan time.Time
 	if c.Retry > 0 {
 		retry := time.NewTicker(c.Retry)
 		defer retry.Stop()
-		end := make(chan struct{})
-		defer close(end)
-		go func() {
-			for {
-				select {
-				case <-retry.C:
-					conn.Write(wire)
-				case <-ctx.Done():
-					return
-				case <-end:
-					return
-				}
-			}
-		}()
+		retryTimer = retry.C
 	}
 
+	go func() {
+		defer conn.Close()
+		for {
+			select {
+			case <-retryTimer:
+				conn.Write(wire)
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
 	var packetErrorCount int
 
 	var incoming [MaxPacketLength]byte
 	for {
 		n, err := conn.Read(incoming[:])
 		if err != nil {
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			default:
+			}
 			return nil, err
 		}
 

client_test.go

@@ -64,3 +64,31 @@ func TestClient_Exchange_retry(t *testing.T) {
 		t.Fatalf("response received in %d attemps; expecting 4", attempts)
 	}
 }
+
+func TestClient_Exchange_cancelled(t *testing.T) {
+	secret := []byte(`12345`)
+	handler := HandlerFunc(func(w ResponseWriter, r *Request) {
+		// ignore
+	})
+	server := newTestServer(handler, StaticSecretSource(secret))
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(time.Millisecond * 50)
+		cancel()
+	}()
+
+	req := New(CodeAccessRequest, secret)
+
+	client := Client{
+		Retry: time.Millisecond * 5,
+	}
+	resp, err := client.Exchange(ctx, req, server.Addr)
+	if resp != nil {
+		t.Fatalf("got non-nil response (%v); expected nil", resp)
+	}
+	if err != context.Canceled {
+		t.Fatalf("got error = %v; expecting context.Canceled", err)
+	}
+}