dissertation-2-code/tcp/flow.go

package tcp

import (
	"encoding/binary"
	"errors"
	"io"
	"mbpl3p/proxy"
	"net"
)

var ErrNotEnoughBytes = errors.New("not enough bytes")

type Flow struct {
	Local  string
	Remote string

	conn *net.TCPConn
}

func NewFlow(local, remote string) (Flow, error) {
	f := Flow{
		Local: local,
		Remote: remote,
	}

	if err := f.dial(); err != nil {
		return Flow{}, err
	}

	return f, nil
}

func newFlowConn(conn *net.TCPConn) Flow {
	return Flow{conn: conn}
}

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

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

	f.conn, err = net.DialTCP("tcp", localAddr, remoteAddr)

	if err != nil {
		return err
	}

	return nil
}

func (f *Flow) IsAlive() bool {
	// TODO: Implement this
	return true
}

func (f *Flow) Consume(p proxy.Packet, g proxy.MacGenerator) error {
	data := p.Marshal(g)
	return f.consumeMarshalled(data)
}

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.Write(prefixedData)
	return err
}

func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
	data, err := f.produceMarshalled()
	if err != nil {
		return proxy.Packet{}, err
	}

	return proxy.UnmarshalPacket(data, v)
}

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, ErrNotEnoughBytes
	}

	length := binary.LittleEndian.Uint32(lengthBytes)
	dataBytes := make([]byte, length)
	if n, err := io.LimitReader(f.conn, int64(length)).Read(dataBytes); err != nil {
		return nil, err
	} else if n != int(length) {
		return nil, ErrNotEnoughBytes
	}

	return dataBytes, nil
}