dissertation-2-code/udp/flow.go

package udp

import (
	"errors"
	"fmt"
	"log"
	"mpbl3p/config"
	"mpbl3p/proxy"
	"mpbl3p/shared"
	"net"
	"sync"
	"time"
)

type PacketWriter interface {
	WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
	LocalAddr() net.Addr
}

type PacketConn interface {
	PacketWriter
	ReadFromUDP(b []byte) (int, *net.UDPAddr, error)
}

type InitiatedFlow struct {
	Local  string
	Remote string

	g         proxy.MacGenerator
	keepalive time.Duration

	mu sync.RWMutex
	Flow
}

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

type Flow struct {
	writer PacketWriter
	raddr  net.UDPAddr

	isAlive    bool
	congestion Congestion

	inboundDatagrams chan []byte
}

func (f Flow) String() string {
	return fmt.Sprintf("UdpInbound{%v -> %v}", f.raddr, f.writer.LocalAddr())
}

func InitiateFlow(
	local, remote string,
	g proxy.MacGenerator,
	c Congestion,
	keepalive time.Duration,
) (*InitiatedFlow, error) {
	f := InitiatedFlow{
		Local:     local,
		Remote:    remote,
		Flow:      newFlow(c),
		g:         g,
		keepalive: keepalive,
	}

	return &f, nil
}

func newFlow(c Congestion) Flow {
	return Flow{
		inboundDatagrams: make(chan []byte),
		congestion:       c,
	}
}

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

	if f.isAlive {
		return nil
	}

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

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

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

	f.writer = conn
	f.isAlive = true

	go func() {
		for {
			buf := make([]byte, 6000)
			n, _, err := conn.ReadFromUDP(buf)
			if err != nil {
				panic(err)
			}

			f.inboundDatagrams <- buf[:n]
		}
	}()

	go func() {
		var err error
		for !errors.Is(err, shared.ErrDeadConnection) {
			f.congestion.AwaitEarlyUpdate(f.keepalive)
			err = f.Consume(proxy.NewSimplePacket(nil), f.g)
		}
	}()

	return nil
}

func (f *InitiatedFlow) Consume(p proxy.Packet, g proxy.MacGenerator) error {
	f.mu.RLock()
	defer f.mu.RUnlock()

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

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
}

func (f *Flow) Consume(pp proxy.Packet, g proxy.MacGenerator) error {
	if !f.isAlive {
		return shared.ErrDeadConnection
	}

	// Sequence is the congestion controllers opportunity to block
	p := Packet{
		seq:  f.congestion.Sequence(),
		data: pp,
	}

	p.ack = f.congestion.NextAck()
	p.nack = f.congestion.NextNack()

	b := p.Marshal()
	b = proxy.AppendMac(b, g)

	_, err := f.writer.WriteToUDP(b, &f.raddr)
	return err
}

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

	b, err := proxy.StripMac(<-f.inboundDatagrams, v)
	if err != nil {
		return nil, err
	}

	p, err := UnmarshalPacket(b)
	if err != nil {
		return nil, err
	}

	// schedule an ack for this sequence number
	f.congestion.ReceivedPacket(p.seq)

	// adjust our sending congestion control based on their acks
	if p.ack != 0 {
		f.congestion.ReceivedAck(p.ack)
	}
	// adjust our sending congestion control based on their nacks
	if p.nack != 0 {
		f.congestion.ReceivedNack(p.nack)
	}

	return p, nil
}

func (f *Flow) handleDatagram(p []byte) {
	// TODO: Fix with security
	// 12 bytes for header + the MAC + a timestamp
	if len(p) == 12+(config.UselessMac{}).CodeLength()+8 {
		b, err := proxy.StripMac(<-f.inboundDatagrams, config.UselessMac{})
		if err != nil {
			log.Println(err)
			return
		}

		p, err := UnmarshalPacket(b)
		if err != nil {
			log.Println(err)
			return
		}

		// TODO: Decide whether to use this line. It means an ACK loop will start, but also is a packet loss.
		f.congestion.ReceivedPacket(p.seq)
		if p.ack != 0 {
			f.congestion.ReceivedAck(p.ack)
		}
		if p.nack != 0 {
			f.congestion.ReceivedNack(p.nack)
		}
	} else {
		f.inboundDatagrams <- p
	}
}