mc-proxy/client/mcproxy/client.go

// Package mcproxy provides a client for the mc-proxy gRPC admin API.
package mcproxy

import (
	"context"
	"fmt"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"

	pb "git.wntrmute.dev/kyle/mc-proxy/gen/mc_proxy/v1"
)

// Client provides access to the mc-proxy admin API.
type Client struct {
	conn   *grpc.ClientConn
	admin  pb.ProxyAdminServiceClient
	health healthpb.HealthClient
}

// Dial connects to the mc-proxy admin API via Unix socket.
func Dial(socketPath string) (*Client, error) {
	conn, err := grpc.NewClient("unix://"+socketPath,
		grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, fmt.Errorf("connecting to %s: %w", socketPath, err)
	}

	return &Client{
		conn:   conn,
		admin:  pb.NewProxyAdminServiceClient(conn),
		health: healthpb.NewHealthClient(conn),
	}, nil
}

// Close closes the connection to the server.
func (c *Client) Close() error {
	return c.conn.Close()
}

// Route represents a hostname to backend mapping.
type Route struct {
	Hostname string
	Backend  string
}

// ListRoutes returns all routes for the given listener address.
func (c *Client) ListRoutes(ctx context.Context, listenerAddr string) ([]Route, error) {
	resp, err := c.admin.ListRoutes(ctx, &pb.ListRoutesRequest{
		ListenerAddr: listenerAddr,
	})
	if err != nil {
		return nil, err
	}

	routes := make([]Route, len(resp.Routes))
	for i, r := range resp.Routes {
		routes[i] = Route{
			Hostname: r.Hostname,
			Backend:  r.Backend,
		}
	}
	return routes, nil
}

// AddRoute adds a route to the given listener.
func (c *Client) AddRoute(ctx context.Context, listenerAddr, hostname, backend string) error {
	_, err := c.admin.AddRoute(ctx, &pb.AddRouteRequest{
		ListenerAddr: listenerAddr,
		Route: &pb.Route{
			Hostname: hostname,
			Backend:  backend,
		},
	})
	return err
}

// RemoveRoute removes a route from the given listener.
func (c *Client) RemoveRoute(ctx context.Context, listenerAddr, hostname string) error {
	_, err := c.admin.RemoveRoute(ctx, &pb.RemoveRouteRequest{
		ListenerAddr: listenerAddr,
		Hostname:     hostname,
	})
	return err
}

// FirewallRuleType represents the type of firewall rule.
type FirewallRuleType string

const (
	FirewallRuleIP      FirewallRuleType = "ip"
	FirewallRuleCIDR    FirewallRuleType = "cidr"
	FirewallRuleCountry FirewallRuleType = "country"
)

// FirewallRule represents a firewall block rule.
type FirewallRule struct {
	Type  FirewallRuleType
	Value string
}

// GetFirewallRules returns all firewall rules.
func (c *Client) GetFirewallRules(ctx context.Context) ([]FirewallRule, error) {
	resp, err := c.admin.GetFirewallRules(ctx, &pb.GetFirewallRulesRequest{})
	if err != nil {
		return nil, err
	}

	rules := make([]FirewallRule, len(resp.Rules))
	for i, r := range resp.Rules {
		rules[i] = FirewallRule{
			Type:  protoToRuleType(r.Type),
			Value: r.Value,
		}
	}
	return rules, nil
}

// AddFirewallRule adds a firewall rule.
func (c *Client) AddFirewallRule(ctx context.Context, ruleType FirewallRuleType, value string) error {
	_, err := c.admin.AddFirewallRule(ctx, &pb.AddFirewallRuleRequest{
		Rule: &pb.FirewallRule{
			Type:  ruleTypeToProto(ruleType),
			Value: value,
		},
	})
	return err
}

// RemoveFirewallRule removes a firewall rule.
func (c *Client) RemoveFirewallRule(ctx context.Context, ruleType FirewallRuleType, value string) error {
	_, err := c.admin.RemoveFirewallRule(ctx, &pb.RemoveFirewallRuleRequest{
		Rule: &pb.FirewallRule{
			Type:  ruleTypeToProto(ruleType),
			Value: value,
		},
	})
	return err
}

// ListenerStatus contains status information for a single listener.
type ListenerStatus struct {
	Addr              string
	RouteCount        int
	ActiveConnections int64
}

// Status contains the server's current status.
type Status struct {
	Version          string
	StartedAt        time.Time
	TotalConnections int64
	Listeners        []ListenerStatus
}

// GetStatus returns the server's current status.
func (c *Client) GetStatus(ctx context.Context) (*Status, error) {
	resp, err := c.admin.GetStatus(ctx, &pb.GetStatusRequest{})
	if err != nil {
		return nil, err
	}

	status := &Status{
		Version:          resp.Version,
		TotalConnections: resp.TotalConnections,
	}
	if resp.StartedAt != nil {
		status.StartedAt = resp.StartedAt.AsTime()
	}

	status.Listeners = make([]ListenerStatus, len(resp.Listeners))
	for i, ls := range resp.Listeners {
		status.Listeners[i] = ListenerStatus{
			Addr:              ls.Addr,
			RouteCount:        int(ls.RouteCount),
			ActiveConnections: ls.ActiveConnections,
		}
	}

	return status, nil
}

// HealthStatus represents the health of the server.
type HealthStatus int

const (
	HealthUnknown    HealthStatus = 0
	HealthServing    HealthStatus = 1
	HealthNotServing HealthStatus = 2
)

func (h HealthStatus) String() string {
	switch h {
	case HealthServing:
		return "SERVING"
	case HealthNotServing:
		return "NOT_SERVING"
	default:
		return "UNKNOWN"
	}
}

// CheckHealth checks the health of the server.
func (c *Client) CheckHealth(ctx context.Context) (HealthStatus, error) {
	resp, err := c.health.Check(ctx, &healthpb.HealthCheckRequest{})
	if err != nil {
		return HealthUnknown, err
	}
	return HealthStatus(resp.Status), nil
}

func protoToRuleType(t pb.FirewallRuleType) FirewallRuleType {
	switch t {
	case pb.FirewallRuleType_FIREWALL_RULE_TYPE_IP:
		return FirewallRuleIP
	case pb.FirewallRuleType_FIREWALL_RULE_TYPE_CIDR:
		return FirewallRuleCIDR
	case pb.FirewallRuleType_FIREWALL_RULE_TYPE_COUNTRY:
		return FirewallRuleCountry
	default:
		return ""
	}
}

func ruleTypeToProto(t FirewallRuleType) pb.FirewallRuleType {
	switch t {
	case FirewallRuleIP:
		return pb.FirewallRuleType_FIREWALL_RULE_TYPE_IP
	case FirewallRuleCIDR:
		return pb.FirewallRuleType_FIREWALL_RULE_TYPE_CIDR
	case FirewallRuleCountry:
		return pb.FirewallRuleType_FIREWALL_RULE_TYPE_COUNTRY
	default:
		return pb.FirewallRuleType_FIREWALL_RULE_TYPE_UNSPECIFIED
	}
}