Added context to reader

astilectron.go

@@ -248,7 +248,8 @@ func (a *Astilectron) acceptTCP(chanAccepted chan bool) {
 
 		// Create reader and writer
 		a.writer = newWriter(conn)
-		a.reader = newReader(a.dispatcher, conn)
+		ctx, _ := a.canceller.NewContext()
+		a.reader = newReader(ctx, a.dispatcher, conn)
 		go a.reader.read()
 	}
 }

reader.go

@@ -3,6 +3,7 @@ package astilectron
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"strings"
@@ -12,15 +13,17 @@ import (
 
 // reader represents an object capable of reading in the TCP server
 type reader struct {
-	d *dispatcher
-	r io.ReadCloser
+	ctx context.Context
+	d   *dispatcher
+	r   io.ReadCloser
 }
 
 // newReader creates a new reader
-func newReader(d *dispatcher, r io.ReadCloser) *reader {
+func newReader(ctx context.Context, d *dispatcher, r io.ReadCloser) *reader {
 	return &reader{
-		d: d,
-		r: r,
+		ctx: ctx,
+		d:   d,
+		r:   r,
 	}
 }
 
@@ -39,6 +42,11 @@ func (r *reader) isEOFErr(err error) bool {
 func (r *reader) read() {
 	var reader = bufio.NewReader(r.r)
 	for {
+		// Check context error
+		if r.ctx.Err() != nil {
+			return
+		}
+
 		// Read next line
 		var b []byte
 		var err error

reader_test.go

@@ -8,6 +8,8 @@ import (
 
 	"io/ioutil"
 
+	"context"
+
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 )
@@ -25,7 +27,7 @@ func (r *mockedReader) Close() error {
 }
 
 func TestReader_IsEOFErr(t *testing.T) {
-	var r = newReader(&dispatcher{}, ioutil.NopCloser(&bytes.Buffer{}))
+	var r = newReader(context.Background(), &dispatcher{}, ioutil.NopCloser(&bytes.Buffer{}))
 	assert.True(t, r.isEOFErr(io.EOF))
 	assert.True(t, r.isEOFErr(errors.New("read tcp 127.0.0.1:56093->127.0.0.1:56092: wsarecv: An existing connection was forcibly closed by the remote host.")))
 	assert.False(t, r.isEOFErr(errors.New("random error")))
@@ -53,7 +55,7 @@ func TestReader(t *testing.T) {
 		return
 	})
 	wg.Add(2)
-	var r = newReader(d, mr)
+	var r = newReader(context.Background(), d, mr)
 
 	// Test read
 	go r.read()