dissertation-2-code/tcp/flow.go

175 lines
3.1 KiB
Go
Raw Normal View History

package tcp
import (
"encoding/binary"
2020-11-26 19:18:07 +00:00
"fmt"
"io"
"mpbl3p/proxy"
2020-11-01 18:18:37 +00:00
"mpbl3p/shared"
"net"
2020-11-01 17:16:43 +00:00
"sync"
2020-11-02 17:24:15 +00:00
"time"
)
2020-10-25 21:44:56 +00:00
type Conn interface {
Read(b []byte) (n int, err error)
Write(b []byte) (n int, err error)
2020-11-02 17:24:15 +00:00
SetWriteDeadline(time.Time) error
2020-11-26 19:18:07 +00:00
// For printing
LocalAddr() net.Addr
RemoteAddr() net.Addr
2020-10-25 21:44:56 +00:00
}
2020-11-01 16:50:02 +00:00
type InitiatedFlow struct {
2020-10-24 17:44:14 +01:00
Local string
Remote string
2020-11-01 17:16:43 +00:00
mu sync.RWMutex
2020-11-01 16:50:02 +00:00
Flow
}
2020-11-26 19:18:07 +00:00
func (f *InitiatedFlow) String() string {
return fmt.Sprintf("TcpOutbound{%v -> %v}", f.Local, f.Remote)
}
2020-11-01 16:50:02 +00:00
type Flow struct {
2020-11-01 18:18:37 +00:00
conn Conn
isAlive bool
}
2020-11-26 19:18:07 +00:00
func (f Flow) String() string {
return fmt.Sprintf("TcpInbound{%v -> %v}", f.conn.RemoteAddr(), f.conn.LocalAddr())
}
2020-11-01 17:16:43 +00:00
func InitiateFlow(local, remote string) (*InitiatedFlow, error) {
2020-11-01 16:50:02 +00:00
f := InitiatedFlow{
2020-10-25 15:36:34 +00:00
Local: local,
2020-10-24 17:44:14 +01:00
Remote: remote,
}
2020-11-01 17:16:43 +00:00
return &f, nil
2020-10-24 17:44:14 +01:00
}
2020-11-01 16:50:02 +00:00
func (f *InitiatedFlow) Reconnect() error {
2020-11-01 17:16:43 +00:00
f.mu.Lock()
defer f.mu.Unlock()
2020-11-01 18:18:37 +00:00
if f.isAlive {
2020-11-01 17:16:43 +00:00
return nil
}
2020-10-24 17:44:14 +01:00
localAddr, err := net.ResolveTCPAddr("tcp", f.Local)
if err != nil {
2020-10-24 17:44:14 +01:00
return err
}
2020-10-24 17:44:14 +01:00
remoteAddr, err := net.ResolveTCPAddr("tcp", f.Remote)
if err != nil {
return err
}
2020-11-03 20:56:51 +00:00
conn, err := net.DialTCP("tcp", localAddr, remoteAddr)
if err != nil {
return err
}
2020-10-24 17:44:14 +01:00
2020-11-03 20:56:51 +00:00
err = conn.SetWriteBuffer(0)
2020-10-24 17:44:14 +01:00
if err != nil {
return err
}
2020-11-05 13:32:47 +00:00
f.conn = conn
2020-11-01 18:18:37 +00:00
f.isAlive = true
2020-10-24 17:44:14 +01:00
return nil
}
2020-11-01 17:16:43 +00:00
func (f *InitiatedFlow) Consume(p proxy.Packet, g proxy.MacGenerator) error {
f.mu.RLock()
defer f.mu.RUnlock()
return f.Flow.Consume(p, g)
}
2020-11-26 22:10:37 +00:00
func (f *InitiatedFlow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
f.mu.RLock()
defer f.mu.RUnlock()
return f.Flow.Produce(v)
}
func (f *Flow) IsAlive() bool {
return f.isAlive
}
2020-11-01 17:16:43 +00:00
func (f *Flow) Consume(p proxy.Packet, g proxy.MacGenerator) (err error) {
2020-11-01 18:18:37 +00:00
if !f.isAlive {
return shared.ErrDeadConnection
}
2020-11-26 18:55:29 +00:00
marshalled := p.Marshal()
data := proxy.AppendMac(marshalled, g)
2020-11-01 17:16:43 +00:00
err = f.consumeMarshalled(data)
if err != nil {
2020-11-01 18:18:37 +00:00
f.isAlive = false
2020-11-01 17:16:43 +00:00
}
return
2020-10-24 17:44:14 +01:00
}
2020-10-24 17:44:14 +01:00
func (f *Flow) consumeMarshalled(data []byte) error {
prefixedData := make([]byte, len(data)+4)
binary.LittleEndian.PutUint32(prefixedData, uint32(len(data)))
2020-10-25 21:44:56 +00:00
copy(prefixedData[4:], data)
2020-11-02 17:29:02 +00:00
err := f.conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
2020-11-02 17:24:15 +00:00
if err != nil {
return err
}
_, err = f.conn.Write(prefixedData)
return err
}
func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) {
2020-11-01 18:18:37 +00:00
if !f.isAlive {
2020-11-25 19:35:31 +00:00
return nil, shared.ErrDeadConnection
2020-11-01 18:18:37 +00:00
}
2020-10-24 17:44:14 +01:00
data, err := f.produceMarshalled()
if err != nil {
2020-11-01 18:18:37 +00:00
f.isAlive = false
2020-11-25 19:35:31 +00:00
return nil, err
2020-10-24 17:44:14 +01:00
}
2020-11-26 18:55:29 +00:00
b, err := proxy.StripMac(data, v)
if err != nil {
return nil, err
}
2020-11-28 17:15:56 +00:00
return proxy.SimplePacket(b), nil
2020-10-24 17:44:14 +01:00
}
2020-10-25 15:36:34 +00:00
func (f *Flow) produceMarshalled() ([]byte, error) {
lengthBytes := make([]byte, 4)
if n, err := io.LimitReader(f.conn, 4).Read(lengthBytes); err != nil {
2020-10-24 17:44:14 +01:00
return nil, err
} else if n != 4 {
2020-11-28 17:15:56 +00:00
return nil, shared.ErrNotEnoughBytes
}
length := binary.LittleEndian.Uint32(lengthBytes)
dataBytes := make([]byte, length)
2020-11-01 16:50:02 +00:00
var read uint32
for read < length {
2020-11-01 18:18:37 +00:00
if n, err := io.LimitReader(f.conn, int64(length-read)).Read(dataBytes[read:]); err != nil {
2020-11-01 16:50:02 +00:00
return nil, err
} else {
read += uint32(n)
}
}
2020-10-24 17:44:14 +01:00
return dataBytes, nil
2020-10-25 15:36:34 +00:00
}