dissertation-2-code/config/builder.go

package config

import (
	"context"
	"encoding/base64"
	"fmt"
	"mpbl3p/crypto"
	"mpbl3p/crypto/sharedkey"
	"mpbl3p/proxy"
	"mpbl3p/replay"
	"mpbl3p/tcp"
	"mpbl3p/udp"
	"mpbl3p/udp/congestion"
	"mpbl3p/udp/congestion/newreno"
	"time"
)

func (c Configuration) Build(ctx context.Context, source proxy.Source, sink proxy.Sink) (*proxy.Proxy, error) {
	p := proxy.NewProxy(0)

	var gs []func() proxy.MacGenerator
	var vs []func() proxy.MacVerifier

	if c.Host.ReplayProtection {
		rp := replay.NewAntiReplay()
		gs = append(gs, func() proxy.MacGenerator { return rp })
		vs = append(vs, func() proxy.MacVerifier { return rp })
	}

	switch c.Host.Crypto {
	case "None":
		gs = append(gs, func() proxy.MacGenerator { return crypto.None{} })
		vs = append(vs, func() proxy.MacVerifier { return crypto.None{} })
	case "Blake2s":
		key, err := base64.StdEncoding.DecodeString(c.Host.SharedKey)
		if err != nil {
			return nil, err
		}
		if _, err := sharedkey.NewBlake2s(key); err != nil {
			return nil, err
		}
		gs = append(gs, func() proxy.MacGenerator {
			g, _ := sharedkey.NewBlake2s(key)
			return g
		})
		vs = append(vs, func() proxy.MacVerifier {
			v, _ := sharedkey.NewBlake2s(key)
			return v
		})
	}

	p.Source = source
	p.Sink = sink

	for _, peer := range c.Peers {
		switch peer.Method {
		case "TCP":
			if err := buildTcp(ctx, p, peer, gs, vs); err != nil {
				return nil, err
			}
		case "UDP":
			if err := buildUdp(ctx, p, peer, gs, vs); err != nil {
				return nil, err
			}
		}
	}

	return p, nil
}

func buildTcp(
	ctx context.Context,
	p *proxy.Proxy,
	peer Peer,
	gs []func() proxy.MacGenerator,
	vs []func() proxy.MacVerifier,
) error {
	var laddr func() string
	if peer.LocalPort == 0 {
		laddr = func() string { return fmt.Sprintf("%s:", peer.GetLocalHost()) }
	} else {
		laddr = func() string { return fmt.Sprintf("%s:%d", peer.GetLocalHost(), peer.LocalPort) }
	}

	if peer.RemoteHost != "" {
		f, err := tcp.InitiateFlow(laddr, fmt.Sprintf("%s:%d", peer.RemoteHost, peer.RemotePort), initialiseVerifiers(vs), initialiseGenerators(gs))

		if err != nil {
			return err
		}

		if !peer.DisableConsumer {
			p.AddConsumer(ctx, f)
		}
		if !peer.DisableProducer {
			p.AddProducer(ctx, f)
		}

		return nil
	}

	err := tcp.NewListener(ctx, p, laddr(), vs, gs, !peer.DisableConsumer, !peer.DisableProducer)
	if err != nil {
		return err
	}

	return nil
}

func buildUdp(
	ctx context.Context,
	p *proxy.Proxy,
	peer Peer,
	gs []func() proxy.MacGenerator,
	vs []func() proxy.MacVerifier,
) error {
	var laddr func() string
	if peer.LocalPort == 0 {
		laddr = func() string { return fmt.Sprintf("%s:", peer.GetLocalHost()) }
	} else {
		laddr = func() string { return fmt.Sprintf("%s:%d", peer.GetLocalHost(), peer.LocalPort) }
	}

	var c func() udp.Congestion
	switch peer.Congestion {
	case "None":
		c = func() udp.Congestion { return congestion.NewNone() }
	default:
		fallthrough
	case "NewReno":
		c = func() udp.Congestion { return newreno.NewNewReno() }
	}

	if peer.RemoteHost != "" {
		f, err := udp.InitiateFlow(
			laddr,
			fmt.Sprintf("%s:%d", peer.RemoteHost, peer.RemotePort),
			initialiseVerifiers(vs),
			initialiseGenerators(gs),
			c(),
			time.Duration(peer.KeepAlive)*time.Second,
		)

		if err != nil {
			return err
		}

		if !peer.DisableConsumer {
			p.AddConsumer(ctx, f)
		}
		if !peer.DisableProducer {
			p.AddProducer(ctx, f)
		}

		return nil
	}

	err := udp.NewListener(ctx, p, laddr(), vs, gs, c, !peer.DisableConsumer, !peer.DisableProducer)
	if err != nil {
		return err
	}

	return nil
}

func initialiseVerifiers(vs []func() proxy.MacVerifier) (out []proxy.MacVerifier) {
	out = make([]proxy.MacVerifier, len(vs))
	for i, v := range vs {
		out[i] = v()
	}
	return
}

func initialiseGenerators(gs []func() proxy.MacGenerator) (out []proxy.MacGenerator) {
	out = make([]proxy.MacGenerator, len(gs))
	for i, g := range gs {
		out[i] = g()
	}
	return
}