From edfd5a900979c46a2f3295c1d0189c81d8af92a4 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Tue, 11 May 2021 21:43:47 +0100
Subject: [PATCH] tcp context fixes

---
 .gitignore      |  2 +-
 tcp/flow.go     | 16 ++++++++--------
 tcp/listener.go |  2 +-
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index 7d235d7..42e21ef 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,4 @@
-config.ini
+*.conf
 logs/
 
 # Created by https://www.toptal.com/developers/gitignore/api/intellij+all,go
diff --git a/tcp/flow.go b/tcp/flow.go
index cf59fde..bf54675 100644
--- a/tcp/flow.go
+++ b/tcp/flow.go
@@ -53,7 +53,7 @@ func NewFlow() Flow {
 	}
 }
 
-func NewFlowConn(conn Conn) Flow {
+func NewFlowConn(ctx context.Context, conn Conn) Flow {
 	f := Flow{
 		conn:    conn,
 		isAlive: true,
@@ -64,8 +64,8 @@ func NewFlowConn(conn Conn) Flow {
 		produceErrors: make(chan error),
 	}
 
-	go f.produceMarshalled()
-	go f.consumeMarshalled()
+	go f.produceMarshalled(ctx)
+	go f.consumeMarshalled(ctx)
 
 	return f
 }
@@ -89,7 +89,7 @@ func InitiateFlow(local func() string, remote string) (*InitiatedFlow, error) {
 	return &f, nil
 }
 
-func (f *InitiatedFlow) Reconnect() error {
+func (f *InitiatedFlow) Reconnect(ctx context.Context) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
@@ -119,8 +119,8 @@ func (f *InitiatedFlow) Reconnect() error {
 	f.conn = conn
 	f.isAlive = true
 
-	go f.produceMarshalled()
-	go f.consumeMarshalled()
+	go f.produceMarshalled(ctx)
+	go f.consumeMarshalled(ctx)
 
 	return nil
 }
@@ -191,7 +191,7 @@ func (f *Flow) Produce(ctx context.Context, v proxy.MacVerifier) (proxy.Packet,
 	return proxy.SimplePacket(b), nil
 }
 
-func (f *Flow) consumeMarshalled() {
+func (f *Flow) consumeMarshalled(ctx context.Context) {
 	for {
 		data := <-f.toConsume
 
@@ -208,7 +208,7 @@ func (f *Flow) consumeMarshalled() {
 	}
 }
 
-func (f *Flow) produceMarshalled() {
+func (f *Flow) produceMarshalled(ctx context.Context) {
 	buf := bufio.NewReader(f.conn)
 
 	for {
diff --git a/tcp/listener.go b/tcp/listener.go
index ba44492..cf87400 100644
--- a/tcp/listener.go
+++ b/tcp/listener.go
@@ -29,7 +29,7 @@ func NewListener(ctx context.Context, p *proxy.Proxy, local string, v func() pro
 				panic(err)
 			}
 
-			f := NewFlowConn(conn)
+			f := NewFlowConn(ctx, conn)
 
 			log.Printf("received new tcp connection: %v\n", f)