132 lines
4.0 KiB
Go
132 lines
4.0 KiB
Go
package gateway
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"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) 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")
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+path, bytes.NewReader(payload))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.HTTPClient.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)
|