tcp #10

Merged
JakeHillion merged 17 commits from tcp into develop 2021-03-23 19:40:56 +00:00
3 changed files with 103 additions and 47 deletions

View File

@ -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,12 +183,34 @@ func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
return proxy.SimplePacket(b), nil
}
func (f *Flow) produceMarshalled() ([]byte, error) {
func (f *Flow) consumeMarshalled() {
for {
data := <-f.toConsume
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
}
}
}
func (f *Flow) produceMarshalled() {
buf := bufio.NewReader(f.conn)
for {
lengthBytes := make([]byte, 4)
if n, err := io.LimitReader(f.conn, 4).Read(lengthBytes); err != nil {
return nil, err
if n, err := io.LimitReader(buf, 4).Read(lengthBytes); err != nil {
f.produceErrors <- err
return
} else if n != 4 {
return nil, shared.ErrNotEnoughBytes
f.produceErrors <- shared.ErrNotEnoughBytes
return
}
length := binary.LittleEndian.Uint32(lengthBytes)
@ -163,12 +218,14 @@ func (f *Flow) produceMarshalled() ([]byte, error) {
var read uint32
for read < length {
if n, err := io.LimitReader(f.conn, int64(length-read)).Read(dataBytes[read:]); err != nil {
return nil, err
if n, err := io.LimitReader(buf, int64(length-read)).Read(dataBytes[read:]); err != nil {
f.produceErrors <- err
return
} else {
read += uint32(n)
}
}
return dataBytes, nil
f.produced <- dataBytes
}
}

View File

@ -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)

View File

@ -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)