75 lines
1.5 KiB
Go
75 lines
1.5 KiB
Go
package tcp
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"mbpl3p/proxy"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
var ErrNotEnoughBytes = errors.New("not enough bytes")
|
|
|
|
type Flow struct {
|
|
Remote net.TCPAddr
|
|
|
|
conn *net.TCPConn
|
|
}
|
|
|
|
func NewFlow(remote net.TCPAddr) (Flow, error) {
|
|
f := Flow{Remote: remote}
|
|
|
|
var err error
|
|
f.conn, err = net.DialTCP("", nil, &f.Remote)
|
|
|
|
if err != nil {
|
|
return Flow{}, err
|
|
}
|
|
|
|
if err := f.conn.SetKeepAlive(true); err != nil {
|
|
return Flow{}, err
|
|
}
|
|
|
|
if err := f.conn.SetKeepAlivePeriod(30 * time.Second); err != nil {
|
|
return Flow{}, err
|
|
}
|
|
|
|
return f, 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)
|
|
|
|
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) {
|
|
lengthBytes := make([]byte, 4)
|
|
if n, err := io.LimitReader(f.conn, 4).Read(lengthBytes); err != nil {
|
|
return proxy.Packet{}, err
|
|
} else if n != 4 {
|
|
return proxy.Packet{}, ErrNotEnoughBytes
|
|
}
|
|
|
|
length := binary.LittleEndian.Uint32(lengthBytes)
|
|
dataBytes := make([]byte, length)
|
|
if n, err := io.LimitReader(f.conn, int64(length)).Read(dataBytes); err != nil {
|
|
return proxy.Packet{}, err
|
|
} else if n != int(length) {
|
|
return proxy.Packet{}, ErrNotEnoughBytes
|
|
}
|
|
|
|
return proxy.UnmarshalPacket(dataBytes, v)
|
|
}
|