vppn/node/peer-states.go

package node

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

type peerState interface {
	Name() string
	OnPeerUpdate(*m.Peer) peerState
	OnPing(netip.AddrPort, pingPacket) peerState
	OnPong(netip.AddrPort, pongPacket) peerState
	OnPingTimer() peerState
	OnTimeoutTimer() peerState
}

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

type stateBase struct {
	// The purpose of this state machine is to manage this published data.
	published *atomic.Pointer[peerData]

	// The other remote peers.
	peers *remotePeers

	// Immutable data.
	localIP  byte
	localPub bool
	remoteIP byte
	privKey  []byte
	conn     *connWriter

	// For sending to peer.
	counter *uint64

	// Mutable peer data.
	peer      *m.Peer
	remotePub bool
	data      peerData // Local copy of shared data. See publish().

	// Timers
	pingTimer    *time.Timer
	timeoutTimer *time.Timer

	buf    []byte
	encBuf []byte
}

func (sb *stateBase) Name() string { return "idle" }

func (s *stateBase) OnPeerUpdate(peer *m.Peer) peerState {
	// Both nil: no change.
	if peer == nil && s.peer == nil {
		return nil
	}

	// No change.
	if peer != nil && s.peer != nil && s.peer.Version == peer.Version {
		return nil
	}

	s.peer = peer

	s.data = peerData{}
	s.data.controlCipher = newControlCipher(s.privKey, peer.EncPubKey)

	ip, isValid := netip.AddrFromSlice(peer.PublicIP)
	if isValid {
		s.remotePub = true
		s.data.remoteAddr = netip.AddrPortFrom(ip, peer.Port)
		s.data.relay = peer.Mediator

		if s.localPub && s.localIP < s.remoteIP {
			return newStateServer(s)
		}
		return newStateClient(s)
	}

	if s.localPub {
		return newStateServer(s)
	}

	// TODO: return newStateMediated(a/b)

	return nil
}

func (s *stateBase) OnPing(rAddr netip.AddrPort, p pingPacket) peerState { return nil }
func (s *stateBase) OnPong(rAddr netip.AddrPort, p pongPacket) peerState { return nil }
func (s *stateBase) OnPingTimer() peerState                              { return nil }
func (s *stateBase) OnTimeoutTimer() peerState                           { return nil }

// Helpers.

func (s *stateBase) resetPingTimer()    { s.pingTimer.Reset(pingInterval) }
func (s *stateBase) resetTimeoutTimer() { s.timeoutTimer.Reset(timeoutInterval) }
func (s *stateBase) stopPingTimer()     { s.pingTimer.Stop() }
func (s *stateBase) stopTimeoutTimer()  { s.timeoutTimer.Stop() }

func (s *stateBase) logf(msg string, args ...any) {
	log.Printf(fmt.Sprintf("[%03d] ", s.remoteIP)+msg, args...)
}

func (s *stateBase) publish() {
	data := s.data
	s.published.Store(&data)
}

func (s *stateBase) sendPing(sharedKey [32]byte) {
	s.sendControlPacket(newPingPacket(sharedKey))
}

func (s *stateBase) sendPong(ping pingPacket) {
	s.sendControlPacket(newPongPacket(ping.SentAt))
}

func (s *stateBase) sendControlPacket(pkt interface{ Marshal([]byte) []byte }) {
	buf := pkt.Marshal(s.buf)
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(s.counter, 1),
		SourceIP: s.localIP,
		DestIP:   s.remoteIP,
	}

	buf = s.data.controlCipher.Encrypt(h, buf, s.encBuf)
	if s.data.relayIP == 0 {
		s.conn.WriteTo(buf, s.data.remoteAddr)
		return
	}

	// TODO: Relay!
}

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

type stateClient struct {
	sharedKey [32]byte
	*stateBase
}

func newStateClient(b *stateBase) peerState {
	s := &stateClient{stateBase: b}
	s.publish()

	s.data.dataCipher = newDataCipher()
	s.sharedKey = s.data.dataCipher.Key()

	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	s.resetTimeoutTimer()
	return s
}

func (s *stateClient) Name() string { return "client" }

func (s *stateClient) OnPong(addr netip.AddrPort, p pongPacket) peerState {
	if !s.data.up {
		s.data.up = true
		s.publish()
	}
	s.resetTimeoutTimer()
	return nil
}

func (s *stateClient) OnPingTimer() peerState {
	s.sendPing(s.sharedKey)
	s.resetPingTimer()
	return nil
}

func (s *stateClient) OnTimeoutTimer() peerState {
	s.data.up = false
	s.publish()
	return nil
}

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

type stateServer struct {
	*stateBase
}

func newStateServer(b *stateBase) peerState {
	s := &stateServer{b}
	s.publish()
	s.stopPingTimer()
	s.stopTimeoutTimer()
	return s
}

func (s *stateServer) Name() string { return "server" }

func (s *stateServer) OnPing(addr netip.AddrPort, p pingPacket) peerState {
	if addr != s.data.remoteAddr {
		s.logf("Got new peer address: %v", addr)
		s.data.remoteAddr = addr
		s.data.up = true
		s.publish()
	}

	if s.data.dataCipher == nil || p.SharedKey != s.data.dataCipher.Key() {
		s.logf("Got new shared key.")
		s.data.dataCipher = newDataCipherFromKey(p.SharedKey)
		s.publish()
	}

	s.sendPong(p)
	return nil
}