From 1b557ac1ead92b39fd96606c2c2ba3882bbc216e Mon Sep 17 00:00:00 2001 From: Juan Font Date: Fri, 25 Nov 2022 17:50:12 +0000 Subject: [PATCH] Update protobuf definitions + support methods for the API Add more logging Updated protos with new routes API --- proto/headscale/v1/headscale.proto | 26 +++++++-- proto/headscale/v1/routes.proto | 55 ++++++++++++------ routes.go | 89 ++++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 22 deletions(-) diff --git a/proto/headscale/v1/headscale.proto b/proto/headscale/v1/headscale.proto index 40d9be2..4464cf2 100644 --- a/proto/headscale/v1/headscale.proto +++ b/proto/headscale/v1/headscale.proto @@ -126,17 +126,31 @@ service HeadscaleService { // --- Machine end --- // --- Route start --- - rpc GetMachineRoute(GetMachineRouteRequest) returns (GetMachineRouteResponse) { + rpc GetRoutes(GetRoutesRequest) returns (GetRoutesResponse) { + option (google.api.http) = { + get: "/api/v1/routes" + }; + } + + rpc EnableRoute(EnableRouteRequest) returns (EnableRouteResponse) { + option (google.api.http) = { + post: "/api/v1/routes/{route_id}/enable" + }; + } + + rpc DisableRoute(DisableRouteRequest) returns (DisableRouteResponse) { + option (google.api.http) = { + post: "/api/v1/routes/{route_id}/disable" + }; + } + + + rpc GetMachineRoutes(GetMachineRoutesRequest) returns (GetMachineRoutesResponse) { option (google.api.http) = { get: "/api/v1/machine/{machine_id}/routes" }; } - rpc EnableMachineRoutes(EnableMachineRoutesRequest) returns (EnableMachineRoutesResponse) { - option (google.api.http) = { - post: "/api/v1/machine/{machine_id}/routes" - }; - } // --- Route end --- // --- ApiKeys start --- diff --git a/proto/headscale/v1/routes.proto b/proto/headscale/v1/routes.proto index 353c429..e55e769 100644 --- a/proto/headscale/v1/routes.proto +++ b/proto/headscale/v1/routes.proto @@ -2,24 +2,47 @@ syntax = "proto3"; package headscale.v1; option go_package = "github.com/juanfont/headscale/gen/go/v1"; -message Routes { - repeated string advertised_routes = 1; - repeated string enabled_routes = 2; +import "google/protobuf/timestamp.proto"; +import "headscale/v1/machine.proto"; + +message Route { + uint64 id = 1; + Machine machine = 2; + string prefix = 3; + bool advertised = 4; + bool enabled = 5; + bool is_primary = 6; + + google.protobuf.Timestamp created_at = 7; + google.protobuf.Timestamp updated_at = 8; + google.protobuf.Timestamp deleted_at = 9; } -message GetMachineRouteRequest { +message GetRoutesRequest { +} + +message GetRoutesResponse { + repeated Route routes = 1; +} + +message EnableRouteRequest { + uint64 route_id = 1; +} + +message EnableRouteResponse { +} + +message DisableRouteRequest { + uint64 route_id = 1; +} + +message DisableRouteResponse { +} + +message GetMachineRoutesRequest { uint64 machine_id = 1; } -message GetMachineRouteResponse { - Routes routes = 1; -} - -message EnableMachineRoutesRequest { - uint64 machine_id = 1; - repeated string routes = 2; -} - -message EnableMachineRoutesResponse { - Routes routes = 1; -} +message GetMachineRoutesResponse { + repeated Route routes = 1; +} \ No newline at end of file diff --git a/routes.go b/routes.go index f1b1913..f59a603 100644 --- a/routes.go +++ b/routes.go @@ -5,7 +5,9 @@ import ( "fmt" "net/netip" + v1 "github.com/juanfont/headscale/gen/go/headscale/v1" "github.com/rs/zerolog/log" + "google.golang.org/protobuf/types/known/timestamppb" "gorm.io/gorm" ) @@ -49,6 +51,64 @@ func (rs Routes) toPrefixes() []netip.Prefix { return prefixes } +func (h *Headscale) GetRoutes() ([]Route, error) { + var routes []Route + err := h.db.Preload("Machine").Find(&routes).Error + if err != nil { + return nil, err + } + + return routes, nil +} + +func (h *Headscale) GetMachineRoutes(m *Machine) ([]Route, error) { + var routes []Route + err := h.db. + Preload("Machine"). + Where("machine_id = ?", m.ID). + Find(&routes).Error + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, err + } + + return routes, nil +} + +func (h *Headscale) GetRoute(id uint64) (*Route, error) { + var route Route + err := h.db.Preload("Machine").First(&route, id).Error + if err != nil { + return nil, err + } + + return &route, nil +} + +func (h *Headscale) EnableRoute(id uint64) error { + route, err := h.GetRoute(id) + if err != nil { + return err + } + + return h.EnableRoutes(&route.Machine, netip.Prefix(route.Prefix).String()) +} + +func (h *Headscale) DisableRoute(id uint64) error { + route, err := h.GetRoute(id) + if err != nil { + return err + } + + route.Enabled = false + route.IsPrimary = false + err = h.db.Save(route).Error + if err != nil { + return err + } + + return h.handlePrimarySubnetFailover() +} + // isUniquePrefix returns if there is another machine providing the same route already. func (h *Headscale) isUniquePrefix(route Route) bool { var count int64 @@ -163,6 +223,10 @@ func (h *Headscale) handlePrimarySubnetFailover() error { if !route.IsPrimary { _, err := h.getPrimaryRoute(netip.Prefix(route.Prefix)) if h.isUniquePrefix(route) || errors.Is(err, gorm.ErrRecordNotFound) { + log.Info(). + Str("prefix", netip.Prefix(route.Prefix).String()). + Str("machine", route.Machine.GivenName). + Msg("Setting primary route") routes[pos].IsPrimary = true err := h.db.Save(&routes[pos]).Error if err != nil { @@ -247,3 +311,28 @@ func (h *Headscale) handlePrimarySubnetFailover() error { return nil } + +func (rs Routes) toProto() []*v1.Route { + protoRoutes := []*v1.Route{} + + for _, route := range rs { + protoRoute := v1.Route{ + Id: uint64(route.ID), + Machine: route.Machine.toProto(), + Prefix: netip.Prefix(route.Prefix).String(), + Advertised: route.Advertised, + Enabled: route.Enabled, + IsPrimary: route.IsPrimary, + CreatedAt: timestamppb.New(route.CreatedAt), + UpdatedAt: timestamppb.New(route.UpdatedAt), + } + + if route.DeletedAt.Valid { + protoRoute.DeletedAt = timestamppb.New(route.DeletedAt.Time) + } + + protoRoutes = append(protoRoutes, &protoRoute) + } + + return protoRoutes +}