tcp/flow.go

@@ -1,6 +1,7 @@
 package tcp
 
 import (
+	"bufio"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -37,16 +38,51 @@ func (f *InitiatedFlow) String() string {
 type Flow struct {
 	conn    Conn
 	isAlive bool
+
+	toConsume, produced          chan []byte
+	consumeErrors, produceErrors chan error
+}
+
+func NewFlow() Flow {
+	return Flow{
+		toConsume:     make(chan []byte),
+		produced:      make(chan []byte),
+		consumeErrors: make(chan error),
+		produceErrors: make(chan error),
+	}
+}
+
+func NewFlowConn(conn Conn) Flow {
+	f := Flow{
+		conn:    conn,
+		isAlive: true,
+
+		toConsume:     make(chan []byte),
+		produced:      make(chan []byte),
+		consumeErrors: make(chan error),
+		produceErrors: make(chan error),
+	}
+
+	go f.produceMarshalled()
+	go f.consumeMarshalled()
+
+	return f
 }
 
 func (f Flow) String() string {
 	return fmt.Sprintf("TcpInbound{%v -> %v}", f.conn.RemoteAddr(), f.conn.LocalAddr())
 }
 
+func (f *Flow) IsAlive() bool {
+	return f.isAlive
+}
+
 func InitiateFlow(local func() string, remote string) (*InitiatedFlow, error) {
 	f := InitiatedFlow{
 		Local:  local,
 		Remote: remote,
+
+		Flow: NewFlow(),
 	}
 
 	return &f, nil
@@ -75,13 +111,16 @@ func (f *InitiatedFlow) Reconnect() error {
 		return err
 	}
 
-	err = conn.SetWriteBuffer(0)
-	if err != nil {
+	if err := conn.SetWriteBuffer(0); err != nil {
 		return err
 	}
 
 	f.conn = conn
 	f.isAlive = true
+
+	go f.produceMarshalled()
+	go f.consumeMarshalled()
+
 	return nil
 }
 
@@ -99,36 +138,27 @@ func (f *InitiatedFlow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
 	return f.Flow.Produce(v)
 }
 
-func (f *Flow) IsAlive() bool {
-	return f.isAlive
-}
-
-func (f *Flow) Consume(p proxy.Packet, g proxy.MacGenerator) (err error) {
+func (f *Flow) Consume(p proxy.Packet, g proxy.MacGenerator) error {
 	if !f.isAlive {
 		return shared.ErrDeadConnection
 	}
 
+	select {
+	case err := <-f.consumeErrors:
+		f.isAlive = false
+		return err
+	default:
+	}
+
 	marshalled := p.Marshal()
 	data := proxy.AppendMac(marshalled, g)
 
-	err = f.consumeMarshalled(data)
-	if err != nil {
-		f.isAlive = false
-	}
-	return
-}
-
-func (f *Flow) consumeMarshalled(data []byte) error {
 	prefixedData := make([]byte, len(data)+4)
 	binary.LittleEndian.PutUint32(prefixedData, uint32(len(data)))
 	copy(prefixedData[4:], data)
 
-	err := f.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
-	if err != nil {
-		return err
-	}
-	_, err = f.conn.Write(prefixedData)
-	return err
+	f.toConsume <- prefixedData
+	return nil
 }
 
 func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
@@ -136,8 +166,11 @@ func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
 		return nil, shared.ErrDeadConnection
 	}
 
-	data, err := f.produceMarshalled()
-	if err != nil {
+	var data []byte
+
+	select {
+	case data = <-f.produced:
+	case err := <-f.produceErrors:
 		f.isAlive = false
 		return nil, err
 	}
@@ -150,25 +183,49 @@ func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
 	return proxy.SimplePacket(b), nil
 }
 
-func (f *Flow) produceMarshalled() ([]byte, error) {
-	lengthBytes := make([]byte, 4)
-	if n, err := io.LimitReader(f.conn, 4).Read(lengthBytes); err != nil {
-		return nil, err
-	} else if n != 4 {
-		return nil, shared.ErrNotEnoughBytes
-	}
+func (f *Flow) consumeMarshalled() {
+	for {
+		data := <-f.toConsume
 
-	length := binary.LittleEndian.Uint32(lengthBytes)
-	dataBytes := make([]byte, length)
-
-	var read uint32
-	for read < length {
-		if n, err := io.LimitReader(f.conn, int64(length-read)).Read(dataBytes[read:]); err != nil {
-			return nil, err
-		} else {
-			read += uint32(n)
+		err := f.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+		if err != nil {
+			f.consumeErrors <- err
+			return
+		}
+		_, err = f.conn.Write(data)
+		if err != nil {
+			f.consumeErrors <- err
+			return
 		}
 	}
-
-	return dataBytes, nil
+}
+
+func (f *Flow) produceMarshalled() {
+	buf := bufio.NewReader(f.conn)
+
+	for {
+		lengthBytes := make([]byte, 4)
+		if n, err := io.LimitReader(buf, 4).Read(lengthBytes); err != nil {
+			f.produceErrors <- err
+			return
+		} else if n != 4 {
+			f.produceErrors <- shared.ErrNotEnoughBytes
+			return
+		}
+
+		length := binary.LittleEndian.Uint32(lengthBytes)
+		dataBytes := make([]byte, length)
+
+		var read uint32
+		for read < length {
+			if n, err := io.LimitReader(buf, int64(length-read)).Read(dataBytes[read:]); err != nil {
+				f.produceErrors <- err
+				return
+			} else {
+				read += uint32(n)
+			}
+		}
+
+		f.produced <- dataBytes
+	}
 }

tcp/flow_test.go

@@ -17,7 +17,7 @@ func TestFlow_Consume(t *testing.T) {
 	t.Run("Length", func(t *testing.T) {
 		testConn := mocks.NewMockPerfectBiStreamConn(100)
 
-		flowA := Flow{conn: testConn.SideA(), isAlive: true}
+		flowA := NewFlowConn(testConn.SideA())
 
 		err := flowA.Consume(testPacket, testMac)
 		require.Nil(t, err)
@@ -41,7 +41,7 @@ func TestFlow_Produce(t *testing.T) {
 	t.Run("Length", func(t *testing.T) {
 		testConn := mocks.NewMockPerfectBiStreamConn(100)
 
-		flowA := Flow{conn: testConn.SideA(), isAlive: true}
+		flowA := NewFlowConn(testConn.SideA())
 
 		_, err := testConn.SideB().Write(testMarshalled)
 		require.Nil(t, err)
@@ -54,7 +54,7 @@ func TestFlow_Produce(t *testing.T) {
 	t.Run("Value", func(t *testing.T) {
 		testConn := mocks.NewMockPerfectBiStreamConn(100)
 
-		flowA := Flow{conn: testConn.SideA(), isAlive: true}
+		flowA := NewFlowConn(testConn.SideA())
 
 		_, err := testConn.SideB().Write(testMarshalled)
 		require.Nil(t, err)

tcp/listener.go

@@ -24,12 +24,11 @@ func NewListener(p *proxy.Proxy, local string, v func() proxy.MacVerifier, g fun
 				panic(err)
 			}
 
-			err = conn.SetWriteBuffer(0)
-			if err != nil {
+			if err := conn.SetWriteBuffer(0); err != nil {
 				panic(err)
 			}
 
-			f := Flow{conn: conn, isAlive: true}
+			f := NewFlowConn(conn)
 
 			log.Printf("received new tcp connection: %v\n", f)