vantage-agent/internal/grpc/client.go

package grpcclient

import (
	"context"
	"crypto/tls"
	"strings"
	"time"

	"gitea.hostxtra.co.uk/mrhid6/vantage/agent/internal/grpc/pb"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/encoding"
	"google.golang.org/grpc/keepalive"
)

func init() {
	encoding.RegisterCodec(JSONCodec{})
}

type Client struct {
	conn   *grpc.ClientConn
	client pb.VantageClient
}

func New(serverURL string, useTLS bool) (*Client, error) {
	serverURL = strings.TrimPrefix(serverURL, "https://")
	serverURL = strings.TrimPrefix(serverURL, "http://")

	dialOpts := []grpc.DialOption{
		grpc.WithKeepaliveParams(keepalive.ClientParameters{
			Time:                30 * time.Second,
			Timeout:             10 * time.Second,
			PermitWithoutStream: false,
		}),
	}

	if useTLS {
		tlsCfg := &tls.Config{
			InsecureSkipVerify: false,
		}
		creds := credentials.NewTLS(tlsCfg)
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(creds))
	} else {
		dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	conn, err := grpc.DialContext(ctx, serverURL, dialOpts...)
	if err != nil {
		return nil, err
	}

	return &Client{
		conn:   conn,
		client: pb.NewVantageClient(conn),
	}, nil
}

func (c *Client) Close() error {
	return c.conn.Close()
}

func (c *Client) Register(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.Register(ctx, &pb.RegisterRequest{
		ServerId:    serverID,
		PreRegToken: preRegToken,
		Hostname:    hostname,
		IpAddress:   ipAddress,
		OsInfo:      osInfo,
	})
	if err != nil {
		return "", err
	}
	return resp.AgentToken, nil
}

// SyncKeys returns the whole response rather than just the keys: the poll now
// also carries CollectPackages, and a second RPC purely to learn one boolean
// would be a message every 30 seconds for a value that changes at most when a
// licence does.
func (c *Client) SyncKeys(serverID, agentToken, version string) (*pb.SyncResponse, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.SyncKeys(ctx, &pb.SyncRequest{
		ServerId:     serverID,
		AgentToken:   agentToken,
		AgentVersion: version,
	})
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// ReportPackages sends a package report and returns whether the server wants
// the full list. Given a longer deadline than the other unary calls because the
// full body is ~150KB on a slow link.
func (c *Client) ReportPackages(req *pb.ReportPackagesRequest) (bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	resp, err := c.client.ReportPackages(ctx, req)
	if err != nil {
		return false, err
	}
	return resp.NeedFull, nil
}

// ReportWorkloads sends a workload report and returns whether the server wants
// the full list.
func (c *Client) ReportWorkloads(req *pb.ReportWorkloadsRequest) (bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	resp, err := c.client.ReportWorkloads(ctx, req)
	if err != nil {
		return false, err
	}
	return resp.NeedFull, nil
}

func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, privateKey, label string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	resp, err := c.client.UploadGeneratedKey(ctx, &pb.UploadKeyRequest{
		ServerId:   serverID,
		AgentToken: agentToken,
		PublicKey:  publicKey,
		PrivateKey: privateKey,
		Label:      label,
	})
	if err != nil {
		return "", err
	}
	return resp.KeyId, nil
}

func (c *Client) ReportUpdates(serverID, agentToken string, updates []pb.PackageUpdate) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	_, err := c.client.ReportUpdates(ctx, &pb.ReportUpdatesRequest{
		ServerId:   serverID,
		AgentToken: agentToken,
		Updates:    updates,
	})
	return err
}

func (c *Client) ReportInventory(report *pb.InventoryReport) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, err := c.client.ReportInventory(ctx, report)
	return err
}

func (c *Client) SyncMonitors(serverID, agentToken string) ([]pb.MonitorSpec, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	resp, err := c.client.SyncMonitors(ctx, &pb.SyncMonitorsRequest{ServerId: serverID, AgentToken: agentToken})
	if err != nil {
		return nil, err
	}
	return resp.Monitors, nil
}

func (c *Client) ReportChecks(serverID, agentToken string, results []pb.CheckResult) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, err := c.client.ReportChecks(ctx, &pb.ReportChecksRequest{ServerId: serverID, AgentToken: agentToken, Results: results})
	return err
}

func (c *Client) CommandStream(ctx context.Context) (pb.Vantage_CommandStreamClient, error) {
	return c.client.CommandStream(ctx)
}

func (c *Client) ProxyStream(ctx context.Context) (pb.Vantage_ProxyStreamClient, error) {
	return c.client.ProxyStream(ctx)
}