Skip to content

Commit

Permalink
Update protobuf definitions + support methods for the API
Browse files Browse the repository at this point in the history
Add more logging

Updated protos with new routes API
  • Loading branch information
juanfont authored and kradalby committed Dec 6, 2022
1 parent 8170f5e commit 1b557ac
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 16 deletions.
22 changes: 18 additions & 4 deletions proto/headscale/v1/headscale.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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/machine/{machine_id}/routes"
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 EnableMachineRoutes(EnableMachineRoutesRequest) returns (EnableMachineRoutesResponse) {

rpc GetMachineRoutes(GetMachineRoutesRequest) returns (GetMachineRoutesResponse) {
option (google.api.http) = {
post: "/api/v1/machine/{machine_id}/routes"
get: "/api/v1/machine/{machine_id}/routes"
};
}

// --- Route end ---

// --- ApiKeys start ---
Expand Down
47 changes: 35 additions & 12 deletions proto/headscale/v1/routes.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
uint64 machine_id = 1;
message GetRoutesRequest {
}

message GetMachineRouteResponse {
Routes routes = 1;
message GetRoutesResponse {
repeated Route routes = 1;
}

message EnableMachineRoutesRequest {
uint64 machine_id = 1;
repeated string routes = 2;
message EnableRouteRequest {
uint64 route_id = 1;
}

message EnableMachineRoutesResponse {
Routes routes = 1;
message EnableRouteResponse {
}

message DisableRouteRequest {
uint64 route_id = 1;
}

message DisableRouteResponse {
}

message GetMachineRoutesRequest {
uint64 machine_id = 1;
}

message GetMachineRoutesResponse {
repeated Route routes = 1;
}
89 changes: 89 additions & 0 deletions routes.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}

0 comments on commit 1b557ac

Please sign in to comment.