108 lines
2.0 KiB
Go
108 lines
2.0 KiB
Go
package tcp
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"mpbl3p/proxy"
|
|
"net"
|
|
)
|
|
|
|
var ErrNotEnoughBytes = errors.New("not enough bytes")
|
|
|
|
type Conn interface {
|
|
Read(b []byte) (n int, err error)
|
|
Write(b []byte) (n int, err error)
|
|
}
|
|
|
|
type Flow struct {
|
|
Local string
|
|
Remote string
|
|
|
|
conn Conn
|
|
}
|
|
|
|
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
|
|
}
|