dissertation-2-code/tcp/flow.go

package tcp

import (
	"bufio"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"mpbl3p/proxy"
	"mpbl3p/shared"
	"net"
	"sync"
	"time"
)

type Conn interface {
	Read(b []byte) (n int, err error)
	Write(b []byte) (n int, err error)
	SetWriteDeadline(time.Time) error

	// For printing
	LocalAddr() net.Addr
	RemoteAddr() net.Addr
}

type InitiatedFlow struct {
	Local  func() string
	Remote string

	mu sync.RWMutex

	Flow
}

func (f *InitiatedFlow) String() string {
	return fmt.Sprintf("TcpOutbound{%v -> %v}", f.Local(), f.Remote)
}

type Flow struct {
	conn    Conn
	isAlive bool

	toConsume, produced          chan []byte
	consumeErrors, produceErrors chan error

	generators []proxy.MacGenerator
	verifiers  []proxy.MacVerifier
}

func NewFlow(vs []proxy.MacVerifier, gs []proxy.MacGenerator) Flow {
	return Flow{
		verifiers:  vs,
		generators: gs,

		toConsume:     make(chan []byte),
		produced:      make(chan []byte),
		consumeErrors: make(chan error),
		produceErrors: make(chan error),
	}
}

func NewFlowConn(ctx context.Context, conn Conn, vs []proxy.MacVerifier, gs []proxy.MacGenerator) Flow {
	f := Flow{
		conn:    conn,
		isAlive: true,

		generators: gs,
		verifiers:  vs,

		toConsume:     make(chan []byte),
		produced:      make(chan []byte),
		consumeErrors: make(chan error),
		produceErrors: make(chan error),
	}

	go f.produceMarshalled(ctx)
	go f.consumeMarshalled(ctx)

	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, vs []proxy.MacVerifier, gs []proxy.MacGenerator) (*InitiatedFlow, error) {
	f := InitiatedFlow{
		Local:  local,
		Remote: remote,

		Flow: NewFlow(vs, gs),
	}

	return &f, nil
}

func (f *InitiatedFlow) Reconnect(ctx context.Context) error {
	f.mu.Lock()
	defer f.mu.Unlock()

	if f.isAlive {
		return nil
	}

	localAddr, err := net.ResolveTCPAddr("tcp", f.Local())
	if err != nil {
		return err
	}

	remoteAddr, err := net.ResolveTCPAddr("tcp", f.Remote)
	if err != nil {
		return err
	}

	conn, err := net.DialTCP("tcp", localAddr, remoteAddr)
	if err != nil {
		return err
	}

	if err := conn.SetWriteBuffer(0); err != nil {
		return err
	}

	f.conn = conn
	f.isAlive = true

	go f.produceMarshalled(ctx)
	go f.consumeMarshalled(ctx)

	return nil
}

func (f *InitiatedFlow) Consume(ctx context.Context, p proxy.Packet) error {
	f.mu.RLock()
	defer f.mu.RUnlock()

	return f.Flow.Consume(ctx, p)
}

func (f *InitiatedFlow) Produce(ctx context.Context) (proxy.Packet, error) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	return f.Flow.Produce(ctx)
}

func (f *Flow) Consume(ctx context.Context, p proxy.Packet) error {
	if !f.isAlive {
		return shared.ErrDeadConnection
	}

	select {
	case err := <-f.consumeErrors:
		f.isAlive = false
		return err
	default:
	}

	data := p.Marshal()
	for _, g := range f.generators {
		data = proxy.AppendMac(data, g)
	}

	prefixedData := make([]byte, len(data)+4)
	binary.LittleEndian.PutUint32(prefixedData, uint32(len(data)))
	copy(prefixedData[4:], data)

	select {
	case f.toConsume <- prefixedData:
	case <-ctx.Done():
		return ctx.Err()
	}

	return nil
}

func (f *Flow) Produce(ctx context.Context) (proxy.Packet, error) {
	if !f.isAlive {
		return nil, shared.ErrDeadConnection
	}

	var data []byte

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case data = <-f.produced:
	case err := <-f.produceErrors:
		f.isAlive = false
		return nil, err
	}

	for _, v := range f.verifiers {
		var err error

		data, err = proxy.StripMac(data, v)
		if err != nil {
			return nil, err
		}
	}

	return proxy.SimplePacket(data), nil
}

func (f *Flow) consumeMarshalled(ctx context.Context) {
	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(ctx context.Context) {
	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
	}
}