vppn/node/peer-supervisor.go

package node

import (
	"net/netip"
	"sync/atomic"
	"time"
	"vppn/m"
)

const (
	connectTimeout  = 6 * time.Second
	pingInterval    = 6 * time.Second
	timeoutInterval = 20 * time.Second
)

type stateFunc func() stateFunc

type peerSuper struct {
	*remotePeer

	peer         *m.Peer
	remotePublic bool
	peerData     peerData

	pktBuf []byte
	encBuf []byte
}

func newPeerSuper(rp *remotePeer) *peerSuper {
	return &peerSuper{
		remotePeer: rp,
		peer:       nil,
		pktBuf:     make([]byte, bufferSize),
		encBuf:     make([]byte, bufferSize),
	}
}

func (rp *peerSuper) Run() {
	defer panicHandler()
	state := rp.stateInit
	for {
		state = state()
	}
}

// ----------------------------------------------------------------------------

func (rp *peerSuper) stateInit() stateFunc {
	//rp.logf("STATE: Init")
	x := peerData{}
	rp.shared.Store(&x)

	rp.peerData.controlCipher = nil
	rp.peerData.dataCipher = nil
	rp.peerData.remoteAddr = zeroAddrPort

	if rp.peer == nil {
		return rp.stateDisconnected
	}

	var addr netip.Addr
	addr, rp.remotePublic = netip.AddrFromSlice(rp.peer.PublicIP)
	if rp.remotePublic {
		rp.peerData.remoteAddr = netip.AddrPortFrom(addr, rp.peer.Port)
	}

	rp.peerData.controlCipher = newControlCipher(rp.privKey, rp.peer.EncPubKey)

	return rp.stateSelectRole()
}

// ----------------------------------------------------------------------------

func (rp *peerSuper) stateDisconnected() stateFunc {
	//rp.logf("STATE: Disconnected")
	for {
		select {
		case <-rp.controlPackets:
			// Drop
		case rp.peer = <-rp.peerUpdates:
			return rp.stateInit
		}
	}
}

// ----------------------------------------------------------------------------

func (rp *peerSuper) stateSelectRole() stateFunc {
	rp.logf("STATE: SelectRole")

	if !rp.localPublic && !rp.remotePublic {
		// TODO!
		return rp.stateDisconnected
	}

	if !rp.localPublic {
		return rp.stateServer
	} else if !rp.remotePublic {
		return rp.stateClient
	}

	if rp.localIP < rp.peer.PeerIP {
		return rp.stateClient
	}
	return rp.stateServer
}

// The remote is a server.
func (rp *peerSuper) stateServer() stateFunc {
	rp.logf("STATE: Server")
	rp.peerData.dataCipher = newDataCipher()
	rp.updateShared()

	var (
		pingTimer = time.NewTimer(pingInterval)
		ping      = pingPacket{SharedKey: ([32]byte)(rp.peerData.dataCipher.Key())}
	)
	defer pingTimer.Stop()

	ping.SentAt = time.Now().UnixMilli()
	rp.sendControlPacket(ping)

	for {
		select {
		case <-pingTimer.C:
			ping.SentAt = time.Now().UnixMilli()
			rp.sendControlPacket(ping)
			pingTimer.Reset(pingInterval)

		case <-rp.controlPackets:
			// Ignore

		case rp.peer = <-rp.peerUpdates:
			return rp.stateInit
		}
	}
}

// ----------------------------------------------------------------------------

// The remote is a client.
func (rp *peerSuper) stateClient() stateFunc {
	rp.logf("STATE: Client")
	rp.updateShared()

	// TODO: Could use timeout to set dataCipher to nil.
	var currentKey = [32]byte{}

	for {
		select {
		case cPkt := <-rp.controlPackets:
			if cPkt.RemoteAddr != rp.peerData.remoteAddr {
				rp.peerData.remoteAddr = cPkt.RemoteAddr
				rp.logf("Got new remote address: %v", cPkt.RemoteAddr)
				rp.updateShared()
			}

			ping, ok := cPkt.Payload.(pingPacket)
			if !ok {
				continue
			}

			if ping.SharedKey != currentKey {
				rp.logf("Connected with new shared key")
				currentKey = ping.SharedKey
				rp.peerData.dataCipher = newDataCipherFromKey(currentKey)
				rp.updateShared()
			}

			rp.sendControlPacket(newPongPacket(ping.SentAt))

		case rp.peer = <-rp.peerUpdates:
			return rp.stateInit
		}
	}
}

// ----------------------------------------------------------------------------

func (rp *peerSuper) updateShared() {
	data := rp.peerData
	rp.shared.Store(&data)
}

// ----------------------------------------------------------------------------

func (rp *peerSuper) sendControlPacket(pkt interface{ Marshal([]byte) []byte }) {
	buf := pkt.Marshal(rp.pktBuf)
	h := xHeader{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(&rp.counter, 1),
		SourceIP: rp.localIP,
		DestIP:   rp.remoteIP,
	}
	buf = rp.peerData.controlCipher.Encrypt(h, buf, rp.encBuf)
	rp.conn.WriteTo(buf, rp.peerData.remoteAddr)
}