hswaw/laserproxy/proxy.go

package main

import (
	"context"
	"net"
	"strings"

	"github.com/golang/glog"
)

type packetFrom struct {
	addr net.Addr
	data []byte
}

func (s *service) runProxy(ctx context.Context, updates chan *lockUpdate, laserAddr net.Addr, laserNet, clientNet net.PacketConn) {
	glog.Infof("Proxy starting... (laser: %v, laser network: %v, client network: %v)", laserAddr, laserNet, clientNet)


	pipe := func(conn net.PacketConn, C chan *packetFrom) {
		for {
			buf := make([]byte, 1500)
			n, addr, err := conn.ReadFrom(buf)
			if err != nil {
				glog.Errorf("pipe failed: %v", err)
			}
			C <- &packetFrom{ addr, buf[:n] }
		}
	}

	laserC := make(chan *packetFrom)
	go pipe(laserNet, laserC)
	clientC := make(chan *packetFrom)
	go pipe(clientNet, clientC)

	var allowedClient string
	var curClient *net.Addr
	for {
		select {
		case <-ctx.Done():
			err := ctx.Err()
			glog.Errorf("Proxy stopped: %v", err)
			return
		case u := <-updates:
			allowedClient = u.addr
			glog.Infof("New allowed client: %q", allowedClient)
		case p := <-laserC:
			s.lockCtrl <- &lockCtrl{
				bump: &lockCtrlBump{},
			}
			if curClient == nil {
				glog.Warningf("Packet from laser without client connected, dropping.")
				break
			}
			clientNet.WriteTo(p.data, *curClient)
		case p := <-clientC:
			s.lockCtrl <- &lockCtrl{
				bump: &lockCtrlBump{},
			}
			if strings.Split(p.addr.String(), ":")[0] == allowedClient {
				curClient = &p.addr
				laserNet.WriteTo(p.data, laserAddr)
			} else {
				glog.Infof("Rejecting packet from %s", p.addr.String())
			}
		}
	}
}