go-playit/proto/control_message.go

package proto

import (
	"fmt"
	"io"
	"net/netip"
	"time"

	"sirherobrine23.com.br/go-bds/go-playit/xcode"
)

var (
	_ xcode.DecodeMessage = (*ControlRequest)(nil)
	_ xcode.EncodeMessage = (*ControlRequest)(nil)
)

type ControlRequest struct {
	Ping                  *Ping
	AgentRegister         *AgentRegister
	AgentKeepAlive        *AgentSessionId
	SetupUdpChannel       *AgentSessionId
	AgentCheckPortMapping *AgentCheckPortMapping
}

func (cr ControlRequest) WriteMessage(_ io.Writer, DefaultWrite func(any) error) (err error) {
	switch {
	default:
		return fmt.Errorf("no field set in ControlRequest")
	case cr.Ping != nil:
		if err = DefaultWrite(ControlRequestIdPingV2); err != nil {
			return
		}
		return DefaultWrite(cr.Ping)
	case cr.AgentRegister != nil:
		if err = DefaultWrite(ControlRequestIdAgentRegisterV2); err != nil {
			return
		}
		return DefaultWrite(cr.AgentRegister)
	case cr.AgentKeepAlive != nil:
		if err = DefaultWrite(ControlRequestIdAgentKeepAliveV1); err != nil {
			return
		}
		return DefaultWrite(cr.AgentKeepAlive)
	case cr.SetupUdpChannel != nil:
		if err = DefaultWrite(ControlRequestIdSetupUdpChannelV1); err != nil {
			return
		}
		return DefaultWrite(cr.SetupUdpChannel)
	case cr.AgentCheckPortMapping != nil:
		if err = DefaultWrite(ControlRequestIdAgentCheckPortMappingV1); err != nil {
			return
		}
		return DefaultWrite(cr.AgentCheckPortMapping)
	}
}

func (cr *ControlRequest) ReadMessage(_ io.Reader, Decode *xcode.DecodeReader) error {
	var id ControlRequestId
	if err := Decode.Decode(&id); err != nil {
		return err
	}

	switch id {
	case ControlRequestIdPingV1:
		cr.Ping = &Ping{}
		now, err := Decode.ReadUint64()
		if err != nil {
			return err
		}
		cr.Ping.Now = time.UnixMilli(int64(now))
		return nil
	case ControlRequestIdPingV2:
		cr.Ping = &Ping{}
		return Decode.Decode(cr.Ping)
	case ControlRequestIdAgentRegisterV1:
		var v1 = AgentRegisterV1{}
		if err := Decode.Decode(&v1); err != nil {
			return err
		}
		cr.AgentRegister = &AgentRegister{
			ProtoVersion: 1,
			AccountId:    v1.AccountId,
			AgentId:      v1.AgentId,
			AgentVersion: v1.AgentVersion,
			Timestamp:    v1.Timestamp,
			ClientAddr:   v1.ClientAddr,
			TunnelAddr:   v1.TunnelAddr,
			Signature:    v1.Signature,
		}
		return nil
	case ControlRequestIdAgentRegisterV2:
		cr.AgentRegister = &AgentRegister{}
		return Decode.Decode(cr.AgentRegister)
	case ControlRequestIdAgentKeepAliveV1:
		cr.AgentKeepAlive = &AgentSessionId{}
		return Decode.Decode(cr.AgentKeepAlive)
	case ControlRequestIdSetupUdpChannelV1:
		cr.SetupUdpChannel = &AgentSessionId{}
		return Decode.Decode(cr.SetupUdpChannel)
	case ControlRequestIdAgentCheckPortMappingV1:
		cr.AgentCheckPortMapping = &AgentCheckPortMapping{}
		return Decode.Decode(cr.AgentCheckPortMapping)
	default:
		return fmt.Errorf("unknown ControlRequestId: %d", id)
	}
}

var (
	_ xcode.DecodeMessage = (*ControlRequestId)(nil)
	_ xcode.EncodeMessage = (*ControlRequestId)(nil)
)

type ControlRequestId uint32

const (
	ControlRequestIdPingV1 ControlRequestId = iota + 1
	ControlRequestIdAgentRegisterV1
	ControlRequestIdAgentKeepAliveV1
	ControlRequestIdSetupUdpChannelV1
	ControlRequestIdAgentCheckPortMappingV1
	ControlRequestIdPingV2
	ControlRequestIdAgentRegisterV2
	ControlRequestIdEND
)

func (cr ControlRequestId) WriteMessage(_ io.Writer, DefaultWrite func(any) error) (err error) {
	return DefaultWrite(uint32(cr))
}
func (cr *ControlRequestId) ReadMessage(_ io.Reader, Decode *xcode.DecodeReader) error {
	v, err := Decode.ReadUint32()
	if err != nil {
		return err
	}
	*cr = ControlRequestId(v)
	return nil
}

var (
	_ xcode.DecodeMessage = (*ControlResponse)(nil)
	_ xcode.EncodeMessage = (*ControlResponse)(nil)
)

type ControlResponse struct {
	Pong              *Pong
	InvalidSignature  *struct{}
	Unauthorized      *struct{}
	RequestQueued     *struct{}
	TryAgainLater     *struct{}
	AgentRegistered   *AgentRegistered
	AgentPortMapping  *AgentPortMapping
	UdpChannelDetails *UdpChannelDetails
}

func (cr ControlResponse) WriteMessage(_ io.Writer, DefaultWrite func(any) error) (err error) {
	typeof := uint32(0)
	data := any(nil)

	switch {
	case cr.Pong != nil:
		typeof = 1
		data = cr.Pong
	case cr.InvalidSignature != nil:
		typeof = 2
	case cr.Unauthorized != nil:
		typeof = 3
	case cr.RequestQueued != nil:
		typeof = 4
	case cr.TryAgainLater != nil:
		typeof = 5
	case cr.AgentRegistered != nil:
		typeof = 6
		data = cr.AgentRegistered
	case cr.AgentPortMapping != nil:
		typeof = 7
		data = cr.AgentPortMapping
	case cr.UdpChannelDetails != nil:
		typeof = 8
		data = cr.UdpChannelDetails
	default:
		return fmt.Errorf("no field set in ControlResponse")
	}

	if err = DefaultWrite(typeof); err != nil {
		return
	}
	return DefaultWrite(data)
}

func (cr *ControlResponse) ReadMessage(_ io.Reader, Decode *xcode.DecodeReader) error {
	typeof, err := Decode.ReadUint32()
	if err != nil {
		return err
	}

	switch typeof {
	case 1:
		cr.Pong = &Pong{}
		err = Decode.Decode(cr.Pong)
	case 2:
		cr.InvalidSignature = &struct{}{}
	case 3:
		cr.Unauthorized = &struct{}{}
	case 4:
		cr.RequestQueued = &struct{}{}
	case 5:
		cr.TryAgainLater = &struct{}{}
	case 6:
		cr.AgentRegistered = &AgentRegistered{}
		err = Decode.Decode(cr.AgentRegistered)
	case 7:
		cr.AgentPortMapping = &AgentPortMapping{}
		err = Decode.Decode(cr.AgentPortMapping)
	case 8:
		cr.UdpChannelDetails = &UdpChannelDetails{}
		err = Decode.Decode(cr.UdpChannelDetails)
	}

	return err
}

type Ping struct {
	Now         time.Time
	CurrentPing xcode.Option[uint32]
	SessionID   xcode.Option[AgentSessionId]
}

type Pong struct {
	RequestNow             time.Time
	ServerNow              time.Time
	ServerId               uint64
	DataCenterId           uint32
	ClientAddr, TunnelAddr netip.AddrPort
	SessionExpireAt        xcode.Option[time.Time]
}

type AgentRegistered struct {
	ID        AgentSessionId
	ExpiresAt time.Time
}

type AgentPortMapping struct {
	Range PortRange
	Found xcode.Option[AgentPortMappingFound]
}

type AgentPortMappingFound struct {
	ToAgent xcode.Option[AgentSessionId]
}

type UdpChannelDetails struct {
	TunnelAddr netip.AddrPort
	Token      []byte
}

type AgentCheckPortMapping struct {
	AgentSessionId AgentSessionId
	PortRange      PortRange
}

type AgentRegister struct {
	ProtoVersion           uint64
	AccountId              uint64
	AgentId                uint64
	AgentVersion           uint64
	Timestamp              uint64
	ClientAddr, TunnelAddr netip.AddrPort
	Signature              [32]byte
}

type AgentRegisterV1 struct {
	AccountId              uint64
	AgentId                uint64
	AgentVersion           uint64
	Timestamp              uint64
	ClientAddr, TunnelAddr netip.AddrPort
	Signature              [32]byte
}