VerseVDI-Data-Plane/cmd/verse-gateway/main.go

package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"net/http"

	"git.sechmachine.io.vn/sechmachine/VerseVDI-Data-Plane/gateway"
	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	var listen, advertiseAddress, controlPlane, certFile, keyFile, clientCAFile string
	var gatewayID, instanceIdentity, certificateIdentity, publicIdentity string
	var providerManagement, providerRTSPAddress, providerRTSPURL, providerIdentity string
	flag.StringVar(&listen, "listen", "0.0.0.0:443", "gateway QUIC listen address")
	flag.StringVar(&advertiseAddress, "advertise-address", "", "client-visible gateway address host:port")
	flag.StringVar(&controlPlane, "control-plane", "", "Connection Server HTTPS base URL")
	flag.StringVar(&certFile, "cert", "", "gateway certificate PEM")
	flag.StringVar(&keyFile, "key", "", "gateway private key PEM")
	flag.StringVar(&clientCAFile, "client-ca", "", "Connection Server/client CA PEM")
	flag.StringVar(&gatewayID, "gateway-id", "", "stable gateway identifier")
	flag.StringVar(&instanceIdentity, "instance-identity", "", "gateway instance identity")
	flag.StringVar(&certificateIdentity, "certificate-identity", "", "gateway certificate identity")
	flag.StringVar(&publicIdentity, "public-identity", "gateway", "gateway public identity")
	flag.StringVar(&providerManagement, "provider-management", "", "internal Apollo management URL")
	flag.StringVar(&providerRTSPAddress, "provider-rtsp-address", "", "internal Apollo RTSP address")
	flag.StringVar(&providerRTSPURL, "provider-rtsp-url", "", "internal Apollo RTSP URL")
	flag.StringVar(&providerIdentity, "provider-identity", "", "enrolled Apollo identity unique-id#fingerprint")
	flag.Parse()
	for name, value := range map[string]string{"control-plane": controlPlane, "advertise-address": advertiseAddress, "cert": certFile, "key": keyFile, "client-ca": clientCAFile, "gateway-id": gatewayID, "instance-identity": instanceIdentity, "certificate-identity": certificateIdentity, "provider-management": providerManagement, "provider-rtsp-address": providerRTSPAddress, "provider-rtsp-url": providerRTSPURL, "provider-identity": providerIdentity} {
		if value == "" {
			return fmt.Errorf("-%s is required", name)
		}
	}
	if err := validateAdvertisedAddress(advertiseAddress); err != nil {
		return err
	}
	expectedIdentity, err := parseProviderIdentity(providerIdentity)
	if err != nil {
		return err
	}
	serverTLS, clientTLS, err := loadTLS(certFile, keyFile, clientCAFile)
	if err != nil {
		return err
	}
	transport := &http.Transport{TLSClientConfig: clientTLS}
	controlPlaneClient := gateway.NewControlPlaneClient(controlPlane, &http.Client{Transport: transport, Timeout: 5 * time.Second})
	providerBackend := gateway.NewNativeApolloBackend(providerManagement, providerRTSPAddress, providerRTSPURL, &http.Client{Timeout: 5 * time.Second})
	provider := gateway.NewApolloAdapter(providerBackend, expectedIdentity)
	capabilities := gateway.DefaultCapabilities()
	server, err := gateway.NewServer(gateway.ServerConfig{ListenAddress: listen, TLSConfig: serverTLS, GatewayID: gatewayID, Capabilities: capabilities, ProviderCapabilities: capabilities, Admission: controlPlaneClient, ProviderStateReporter: controlPlaneClient, Provider: provider, PacerKbps: 100000})
	if err != nil {
		return err
	}
	registration := protocol.GatewayRegistration{Version: "1", GatewayID: gatewayID, InstanceIdentity: instanceIdentity, CertificateIdentity: certificateIdentity, PublicIdentity: publicIdentity, Address: advertiseAddress, ProviderIdentity: providerIdentity, ProtocolMinVersion: 1, ProtocolMaxVersion: 1, ConnectionCapacity: 8, BandwidthCapacityKbps: 100000, Features: []string{"quic-tls13", "datagram.media", "apollo"}, Capabilities: capabilities}
	if _, err := controlPlaneClient.Register(context.Background(), registration); err != nil {
		_ = server.Close()
		return err
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	go heartbeatLoop(ctx, controlPlaneClient, server, registration)
	return server.Serve(ctx)
}

func validateAdvertisedAddress(address string) error {
	host, port, err := net.SplitHostPort(address)
	if err != nil || host == "" || port == "" {
		return fmt.Errorf("-advertise-address must be a host:port address")
	}
	if parsed := net.ParseIP(host); parsed != nil && parsed.IsUnspecified() {
		return fmt.Errorf("-advertise-address must not be an unspecified address")
	}
	if net.JoinHostPort(host, port) != address {
		return fmt.Errorf("-advertise-address must use canonical host:port syntax")
	}
	return nil
}

func heartbeatLoop(ctx context.Context, client *gateway.ControlPlaneClient, server *gateway.Server, registration protocol.GatewayRegistration) {
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	var sequence int64
	for {
		select {
		case <-ctx.Done():
			server.BeginDrain()
			deadline := time.Now().Add(5 * time.Second).UTC().Format(time.RFC3339Nano)
			_ = client.Drain(context.Background(), protocol.GatewayDrain{Version: "1", GatewayID: registration.GatewayID, Sequence: sequence + 1, Reason: "shutdown", Deadline: deadline})
			return
		case <-ticker.C:
			sequence++
			state := "ready"
			if server.Draining() {
				state = "draining"
			}
			metrics := server.Metrics()
			_ = client.Heartbeat(ctx, protocol.GatewayHeartbeat{Version: "1", GatewayID: registration.GatewayID, Sequence: sequence, ObservedAt: time.Now().UTC().Format(time.RFC3339Nano), ActiveConnections: metrics.ActiveSessions, EgressKbps: registration.BandwidthCapacityKbps, State: state})
		}
	}
}

func loadTLS(certFile, keyFile, clientCAFile string) (*tls.Config, *tls.Config, error) {
	certificate, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, nil, err
	}
	caBytes, err := os.ReadFile(clientCAFile)
	if err != nil {
		return nil, nil, err
	}
	clientCAs := x509.NewCertPool()
	if !clientCAs.AppendCertsFromPEM(caBytes) {
		return nil, nil, errors.New("client CA PEM contains no certificate")
	}
	return &tls.Config{MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{certificate}, ClientAuth: tls.RequireAndVerifyClientCert, ClientCAs: clientCAs}, &tls.Config{MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{certificate}, RootCAs: clientCAs}, nil
}

func parseProviderIdentity(value string) (gateway.ProviderIdentity, error) {
	uniqueID, fingerprint, ok := strings.Cut(value, "#")
	if !ok || uniqueID == "" || fingerprint == "" {
		return gateway.ProviderIdentity{}, errors.New("provider identity must be unique-id#fingerprint")
	}
	return gateway.ProviderIdentity{UniqueID: uniqueID, Fingerprint: fingerprint}, nil
}