VerseVDI-Data-Plane/gateway/control_plane.go

package gateway

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

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

type ControlPlaneClient struct {
	BaseURL    string
	HTTPClient *http.Client
}

func NewControlPlaneClient(baseURL string, client *http.Client) *ControlPlaneClient {
	if client == nil {
		client = &http.Client{}
	}
	return &ControlPlaneClient{BaseURL: strings.TrimRight(baseURL, "/"), HTTPClient: client}
}

func (c *ControlPlaneClient) Register(ctx context.Context, registration protocol.GatewayRegistration) (protocol.GatewayRegistration, error) {
	payload, err := protocol.EncodeGatewayRegistration(registration)
	if err != nil {
		return protocol.GatewayRegistration{}, err
	}
	response, err := c.post(ctx, "/api/v1/gateway/register", payload)
	if err != nil {
		return protocol.GatewayRegistration{}, err
	}
	return protocol.DecodeGatewayRegistration(response)
}

func (c *ControlPlaneClient) Heartbeat(ctx context.Context, heartbeat protocol.GatewayHeartbeat) error {
	payload, err := protocol.EncodeGatewayHeartbeat(heartbeat)
	if err != nil {
		return err
	}
	_, err = c.post(ctx, "/api/v1/gateway/heartbeat", payload)
	return err
}

func (c *ControlPlaneClient) Drain(ctx context.Context, drain protocol.GatewayDrain) error {
	payload, err := protocol.EncodeGatewayDrain(drain)
	if err != nil {
		return err
	}
	_, err = c.post(ctx, "/api/v1/gateway/drain", payload)
	return err
}

func (c *ControlPlaneClient) Admit(ctx context.Context, request protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error) {
	payload, err := protocol.EncodeTunnelAdmissionRequest(request)
	if err != nil {
		return protocol.SessionAuthority{}, err
	}
	response, err := c.post(ctx, "/api/v1/gateway/admit", payload)
	if err != nil {
		return protocol.SessionAuthority{}, err
	}
	return protocol.DecodeSessionAuthority(response)
}

func (c *ControlPlaneClient) Release(ctx context.Context, authority protocol.SessionAuthority) error {
	payload, err := protocol.EncodeSessionAuthority(authority)
	if err != nil {
		return err
	}
	_, err = c.post(ctx, "/api/v1/gateway/release", payload)
	return err
}

func (c *ControlPlaneClient) ProviderWork(ctx context.Context, authority protocol.SessionAuthority) (protocol.ProviderSessionWork, error) {
	payload, err := protocol.EncodeSessionAuthority(authority)
	if err != nil {
		return protocol.ProviderSessionWork{}, err
	}
	response, err := c.post(ctx, "/api/v1/gateway/provider-work", payload)
	if err != nil {
		return protocol.ProviderSessionWork{}, err
	}
	return protocol.DecodeProviderSessionWork(response)
}

func (c *ControlPlaneClient) ReportProviderState(ctx context.Context, state protocol.ProviderState) error {
	payload, err := protocol.EncodeProviderState(state)
	if err != nil {
		return err
	}
	_, err = c.post(ctx, "/api/v1/gateway/provider-state", payload)
	return err
}

func (c *ControlPlaneClient) ReportClipboardAudit(ctx context.Context, audit protocol.GatewayClipboardAudit) error {
	payload, err := protocol.EncodeGatewayClipboardAudit(audit)
	if err != nil {
		return err
	}
	_, err = c.post(ctx, "/api/v1/gateway/clipboard-audit", payload)
	return err
}

func (c *ControlPlaneClient) post(ctx context.Context, path string, payload []byte) ([]byte, error) {
	if c == nil || c.HTTPClient == nil || c.BaseURL == "" {
		return nil, errors.New("control-plane client is not configured")
	}
	baseURL, err := url.Parse(c.BaseURL)
	if err != nil || baseURL.Scheme != "https" || baseURL.Host == "" {
		return nil, errors.New("control-plane base URL must be absolute HTTPS")
	}
	requestURL := *baseURL
	requestURL.Path = strings.TrimRight(baseURL.Path, "/") + path
	requestURL.RawPath = ""
	requestURL.RawQuery = ""
	requestURL.Fragment = ""
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}
	request.Header.Set("Content-Type", "application/json")
	client := *c.HTTPClient
	client.CheckRedirect = func(*http.Request, []*http.Request) error {
		return errors.New("control-plane redirects are not permitted")
	}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	body, err := io.ReadAll(io.LimitReader(response.Body, defaultControlLimit+1))
	if err != nil {
		return nil, err
	}
	if len(body) > defaultControlLimit {
		return nil, ErrFrameSize
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		var stable protocol.StableError
		if json.Unmarshal(body, &stable) == nil && stable.Code != "" {
			return nil, fmt.Errorf("%s: %s", stable.Code, stable.Message)
		}
		return nil, fmt.Errorf("control-plane status %d", response.StatusCode)
	}
	return body, nil
}

var _ Admission = (*ControlPlaneClient)(nil)
var _ ClipboardAuditReporter = (*ControlPlaneClient)(nil)