ziti/controller/model/controller_manager.go

/*
	Copyright NetFoundry Inc.

	Licensed under the Apache License, Version 2.0 (the "License");
	you may not use this file except in compliance with the License.
	You may obtain a copy of the License at

	https://www.apache.org/licenses/LICENSE-2.0

	Unless required by applicable law or agreed to in writing, software
	distributed under the License is distributed on an "AS IS" BASIS,
	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	See the License for the specific language governing permissions and
	limitations under the License.
*/

package model

import (
	"crypto/x509"
	"fmt"
	"time"

	"github.com/michaelquigley/pfxlog"
	nfpem "github.com/openziti/foundation/v2/pem"
	"github.com/openziti/ziti/v2/common/pb/edge_cmd_pb"
	"github.com/openziti/ziti/v2/controller/change"
	"github.com/openziti/ziti/v2/controller/command"
	"github.com/openziti/ziti/v2/controller/db"
	"github.com/openziti/ziti/v2/controller/event"
	"github.com/openziti/ziti/v2/controller/fields"
	"github.com/openziti/ziti/v2/controller/models"
	"github.com/openziti/ziti/v2/controller/storage/boltz"
	"google.golang.org/protobuf/proto"
)

func NewControllerManager(env Env) *ControllerManager {
	manager := &ControllerManager{
		baseEntityManager: newBaseEntityManager[*Controller, *db.Controller](env, env.GetStores().Controller),
	}
	manager.impl = manager

	RegisterManagerDecoder[*Controller](env, manager)

	return manager
}

type ControllerManager struct {
	baseEntityManager[*Controller, *db.Controller]
}

func (self *ControllerManager) NewModelEntity() *Controller {
	return &Controller{}
}

func (self *ControllerManager) Create(entity *Controller, ctx *change.Context) error {
	return DispatchCreate[*Controller](self, entity, ctx)
}

func (self *ControllerManager) ApplyCreate(cmd *command.CreateEntityCommand[*Controller], ctx boltz.MutateContext) error {
	_, err := self.createEntity(cmd.Entity, ctx)
	return err
}

func (self *ControllerManager) Update(entity *Controller, checker fields.UpdatedFields, ctx *change.Context) error {
	return DispatchUpdate[*Controller](self, entity, checker, ctx)
}

func (self *ControllerManager) ApplyUpdate(cmd *command.UpdateEntityCommand[*Controller], ctx boltz.MutateContext) error {
	return self.updateEntity(cmd.Entity, cmd.UpdatedFields, ctx)
}

func (self *ControllerManager) Read(id string) (*Controller, error) {
	modelEntity := &Controller{}
	if err := self.readEntity(id, modelEntity); err != nil {
		return nil, err
	}
	return modelEntity, nil
}

// ReadAll returns all controllers from the database. If no controllers exist (single-controller
// non-HA mode), it returns a synthetic controller representing the running instance so that
// token verification via kid lookup works without special-case fallback paths. This runs during
// startup (token issuer cache load), so it must stay startup-safe; for the client-facing list with
// API addresses, use ReadAllForClient at request time.
func (self *ControllerManager) ReadAll() ([]*Controller, error) {
	result, err := self.BaseList("true limit none")
	if err != nil {
		return nil, err
	}

	if len(result.Entities) > 0 {
		return result.Entities, nil
	}

	// No controllers in the DB — single-controller mode. Build a representation of self
	// so that the token issuer cache can verify controller-issued JWTs via kid lookup.
	selfController, err := self.buildSelfController()
	if err != nil {
		return nil, fmt.Errorf("no controllers in database and could not build self controller: %w", err)
	}

	return []*Controller{selfController}, nil
}

// buildSelfController constructs a minimal Controller model representing the running instance from
// the root TLS JWT signer certificate. It is only called in single-controller (non-HA) deployments
// where no controller records exist in the database. It must stay free of dependencies that are not
// ready during controller startup: the token issuer cache calls ReadAll while loading, before xweb
// is initialized, so this must NOT read API addresses (see ReadAllForClient for the request-time,
// address-populated variant).
func (self *ControllerManager) buildSelfController() (*Controller, error) {
	signer := self.env.GetRootTlsJwtSigner()
	if signer == nil || signer.TlsCerts == nil || len(signer.TlsCerts.Certificate) == 0 {
		return nil, fmt.Errorf("root TLS JWT signer has no certificates")
	}

	cert, err := x509.ParseCertificate(signer.TlsCerts.Certificate[0])
	if err != nil {
		return nil, fmt.Errorf("could not parse root TLS certificate: %w", err)
	}

	return &Controller{
		BaseEntity: models.BaseEntity{
			Id: "single-controller",
		},
		Name:        cert.Subject.CommonName,
		CertPem:     nfpem.EncodeToString(cert),
		Fingerprint: nfpem.FingerprintFromCertificate(cert),
		IsOnline:    true,
	}, nil
}

// ReadAllForClient returns the controllers to advertise to clients, e.g. in enrollment responses.
// It behaves like ReadAll, but in single-controller (non-HA) mode the synthetic self additionally
// carries the running controller's real id and its API addresses from the live xweb configuration,
// so the enrolling client receives connectable endpoints (matching what an HA controller would have
// stored for itself). It MUST be called only at request time: it reads GetApiAddresses, which blocks
// until xweb is initialized and would deadlock if called during controller startup.
func (self *ControllerManager) ReadAllForClient() ([]*Controller, error) {
	result, err := self.BaseList("true limit none")
	if err != nil {
		return nil, err
	}

	if len(result.Entities) > 0 {
		return result.Entities, nil
	}

	selfController, err := self.buildSelfController()
	if err != nil {
		return nil, fmt.Errorf("no controllers in database and could not build self controller: %w", err)
	}

	selfController.Id = self.env.GetId()

	selfController.ApiAddresses = map[string][]ApiAddress{}
	envApiAddresses, _ := self.env.GetApiAddresses()
	for apiKey, instances := range envApiAddresses {
		for _, instance := range instances {
			selfController.ApiAddresses[apiKey] = append(selfController.ApiAddresses[apiKey], ApiAddress{
				Url:     instance.Url,
				Version: instance.Version,
			})
		}
	}

	return []*Controller{selfController}, nil
}

func (self *ControllerManager) ReadByName(name string) (*Controller, error) {
	modelEntity := &Controller{}
	nameIndex := self.env.GetStores().Controller.GetNameIndex()
	if err := self.readEntityWithIndex("name", []byte(name), nameIndex, modelEntity); err != nil {
		return nil, err
	}
	return modelEntity, nil
}

func (self *ControllerManager) Marshall(entity *Controller) ([]byte, error) {
	msg := &edge_cmd_pb.Controller{
		Id:                entity.Id,
		Name:              entity.Name,
		Address:           entity.CtrlAddress,
		CertPem:           entity.CertPem,
		Fingerprint:       entity.Fingerprint,
		IsOnline:          entity.IsOnline,
		LastJoinedAt:      timePtrToPb(&entity.LastJoinedAt),
		IsPreferredLeader: entity.IsPreferredLeader,
		ApiAddresses:      map[string]*edge_cmd_pb.ApiAddressList{},
	}

	for apiKey, instances := range entity.ApiAddresses {
		msg.ApiAddresses[apiKey] = &edge_cmd_pb.ApiAddressList{}
		for _, instance := range instances {
			msg.ApiAddresses[apiKey].Addresses = append(msg.ApiAddresses[apiKey].Addresses, &edge_cmd_pb.ApiAddress{
				Url:     instance.Url,
				Version: instance.Version,
			})
		}
	}

	return proto.Marshal(msg)
}

func (self *ControllerManager) Unmarshall(bytes []byte) (*Controller, error) {
	msg := &edge_cmd_pb.Controller{}

	if err := proto.Unmarshal(bytes, msg); err != nil {
		return nil, err
	}

	lastJoinedAt := time.Time{}
	if msg.LastJoinedAt != nil {
		lastJoinedAt = *pbTimeToTimePtr(msg.LastJoinedAt)
	}

	controller := &Controller{
		BaseEntity: models.BaseEntity{
			Id:   msg.Id,
			Tags: edge_cmd_pb.DecodeTags(msg.Tags),
		},
		Name:              msg.Name,
		CtrlAddress:       msg.Address,
		CertPem:           msg.CertPem,
		Fingerprint:       msg.Fingerprint,
		IsOnline:          msg.IsOnline,
		LastJoinedAt:      lastJoinedAt,
		IsPreferredLeader: msg.IsPreferredLeader,
		ApiAddresses:      map[string][]ApiAddress{},
	}

	for apiKey, instanceList := range msg.ApiAddresses {
		controller.ApiAddresses[apiKey] = nil
		if instanceList != nil {
			for _, instance := range instanceList.Addresses {
				controller.ApiAddresses[apiKey] = append(controller.ApiAddresses[apiKey], ApiAddress{
					Url:     instance.Url,
					Version: instance.Version,
				})
			}
		}
	}

	return controller, nil
}

func (self *ControllerManager) getCurrentAsClusterPeer() *event.ClusterPeer {
	addr, id, version := self.env.GetRaftInfo()
	clientApiCert := self.env.GetRootTlsJwtSigner()
	var leaderCerts []*x509.Certificate

	for _, certBytes := range clientApiCert.TlsCerts.Certificate {
		if cert, err := x509.ParseCertificate(certBytes); err == nil {
			leaderCerts = append(leaderCerts, cert)
		}
	}

	apiAddresses, _ := self.env.GetApiAddresses()

	isPreferredLeader := false
	if raftConfig := self.env.GetConfig().Raft; raftConfig != nil {
		isPreferredLeader = raftConfig.PreferredLeader
	}

	return &event.ClusterPeer{
		Id:                id,
		Addr:              addr,
		Version:           version,
		ServerCert:        leaderCerts,
		ApiAddresses:      apiAddresses,
		IsPreferredLeader: isPreferredLeader,
	}
}

func (self *ControllerManager) UpdateControllerState(peers []*event.ClusterPeer, peerConnectedEvent bool) {
	controllers := map[string]*Controller{}

	result, err := self.BaseList("true limit none")
	if err != nil {
		pfxlog.Logger().WithError(err).Error("failed to list controllers")
		return
	}

	for _, ctrl := range result.Entities {
		controllers[ctrl.Id] = ctrl
	}

	changeCtx := change.New()
	if peerConnectedEvent {
		changeCtx.SetSourceType("raft.peers.connected").
			SetChangeAuthorType(change.AuthorTypeController)
	} else {
		changeCtx.SetSourceType("raft.leadership.gained").
			SetChangeAuthorType(change.AuthorTypeController)
	}

	selfAsPeer := self.getCurrentAsClusterPeer()
	peerFingerprints := ""
	for _, peer := range peers {
		if len(peer.ServerCert) > 0 {
			fingerprint := nfpem.FingerprintFromCertificate(peer.ServerCert[0])

			if peerFingerprints == "" {
				peerFingerprints = fingerprint
			} else {
				peerFingerprints = peerFingerprints + ", " + fingerprint
			}
		}
	}

	pfxlog.Logger().Infof("acting as leader, updating controllers from peers, connectEvt? %v, self: %s, peer count: %d, peers: %s",
		peerConnectedEvent, nfpem.FingerprintFromCertificate(selfAsPeer.ServerCert[0]), len(peers), peerFingerprints)

	if !peerConnectedEvent {
		// add this controller as a "peer" when leadership is gained
		peers = append(peers, selfAsPeer)
	}

	for _, peer := range peers {
		// Use our locally built peer instance to represent ourselves in the list
		if peer.Id == selfAsPeer.Id && peer != selfAsPeer {
			continue
		}

		if len(peer.ServerCert) < 1 {
			pfxlog.Logger().Errorf("peer %s has no certificate", peer.Id)
			continue
		}

		newController := &Controller{
			BaseEntity: models.BaseEntity{
				Id: peer.Id,
			},
			Name:              peer.ServerCert[0].Subject.CommonName,
			CertPem:           nfpem.EncodeToString(peer.ServerCert[0]),
			Fingerprint:       nfpem.FingerprintFromCertificate(peer.ServerCert[0]),
			CtrlAddress:       peer.Addr,
			IsOnline:          true,
			LastJoinedAt:      time.Now(),
			IsPreferredLeader: peer.IsPreferredLeader,
			ApiAddresses:      apiAddressesFromPeer(peer),
		}

		existing := controllers[peer.Id]
		if existing == nil {
			if err = self.Create(newController, changeCtx); err != nil {
				pfxlog.Logger().WithError(err).WithField("ctrlId", peer.Id).
					Error("could not create controller during peer(s) connection")
			}
		} else if peerConnectedEvent || existing.IsChanged(newController) {
			if err = self.Update(newController, nil, changeCtx); err != nil {
				pfxlog.Logger().WithError(err).WithField("ctrlId", peer.Id).
					Error("could not update controller during peer(s) connection")
			}
		}
	}

	// If we're the new leader, marking any controller not connected to us as offline
	if !peerConnectedEvent {
		connectedPeers := map[string]struct{}{}
		for _, peer := range self.env.GetCommandDispatcher().GetPeers() {
			connectedPeers[peer.Id()] = struct{}{}
		}

		disconnectFields := fields.UpdatedFieldsMap{
			db.FieldControllerIsOnline: struct{}{},
		}

		for _, controller := range controllers {
			if controller.IsOnline && controller.Id != selfAsPeer.Id {
				if _, ok := connectedPeers[controller.Id]; !ok {
					controller.IsOnline = false

					if err := self.Update(controller, disconnectFields, changeCtx); err != nil {
						pfxlog.Logger().WithError(err).Error("could not update controller marking peer disconnected")
					}
				}
			}
		}
	}
}

func (self *ControllerManager) DeleteRemovedPeers(peers []*event.ClusterPeer) {
	controllers := map[string]*Controller{}

	result, err := self.BaseList("true limit none")
	if err != nil {
		pfxlog.Logger().WithError(err).Error("failed to list controllers")
		return
	}

	for _, ctrl := range result.Entities {
		controllers[ctrl.Id] = ctrl
	}

	changeCtx := change.New()
	changeCtx.SetSourceType("raft.peers.changed").
		SetChangeAuthorType(change.AuthorTypeController)

	peers = append(peers, self.getCurrentAsClusterPeer())

	for _, peer := range peers {
		delete(controllers, peer.Id)
	}

	for _, controller := range controllers {
		if err = self.Delete(controller.Id, changeCtx); err != nil {
			pfxlog.Logger().WithError(err).WithField("ctrlId", controller.Id).
				Error("could not delete controller during cluster member change")
		}
	}
}

func (self *ControllerManager) PeersDisconnected(peers []*event.ClusterPeer) {
	changeCtx := change.New()
	changeCtx.SetSourceType("raft.peers.disconnected").
		SetChangeAuthorType(change.AuthorTypeController)

	disconnectFields := fields.UpdatedFieldsMap{
		db.FieldControllerIsOnline: struct{}{},
	}
	for _, peer := range peers {
		controller := &Controller{
			BaseEntity: models.BaseEntity{
				Id: peer.Id,
			},
			IsOnline: false,
		}

		if err := self.Update(controller, disconnectFields, changeCtx); err != nil {
			pfxlog.Logger().WithError(err).Error("could not update controller during peer(s) disconnection")
		}
	}
}

// UpdateSelfOnNewLeader has each controller update itself with its current configuration. This is then
// forwarded to the current leader for updates. Ensure that new leaders always have the most update-to-date configuration
// from each controller and avoid stale or out-of-date information.
func (self *ControllerManager) UpdateSelfOnNewLeader() {
	peer := self.getCurrentAsClusterPeer()

	newController := &Controller{
		BaseEntity: models.BaseEntity{
			Id: peer.Id,
		},
		Name:              peer.ServerCert[0].Subject.CommonName,
		CertPem:           nfpem.EncodeToString(peer.ServerCert[0]),
		Fingerprint:       nfpem.FingerprintFromCertificate(peer.ServerCert[0]),
		CtrlAddress:       peer.Addr,
		IsOnline:          true,
		IsPreferredLeader: peer.IsPreferredLeader,
		ApiAddresses:      apiAddressesFromPeer(peer),
	}
	disconnectFields := fields.UpdatedFieldsMap{
		db.FieldControllerIsOnline:          struct{}{},
		db.FieldControllerCertPem:           struct{}{},
		db.FieldControllerFingerprint:       struct{}{},
		db.FieldControllerCtrlAddress:       struct{}{},
		db.FieldControllerApiAddresses:      struct{}{},
		db.FieldControllerApiAddressUrl:     struct{}{},
		db.FieldControllerApiAddressVersion: struct{}{},
		db.FieldControllerIsPreferredLeader: struct{}{},
		db.FieldName:                        struct{}{},
	}

	changeCtx := change.New()
	changeCtx.SetSourceType("raft.state.has_leader").
		SetChangeAuthorType(change.AuthorTypeController)

	if err := self.Update(newController, disconnectFields, changeCtx); err != nil {
		pfxlog.Logger().WithError(err).Error("could not update self")
	}
}

// apiAddressFromPeer converts event.ClusterPeer API Addresses to model API Addresses
func apiAddressesFromPeer(peer *event.ClusterPeer) map[string][]ApiAddress {
	result := map[string][]ApiAddress{}

	for apiKey, instances := range peer.ApiAddresses {
		result[apiKey] = nil

		for _, instance := range instances {
			result[apiKey] = append(result[apiKey], ApiAddress{
				Url:     instance.Url,
				Version: instance.Version,
			})
		}
	}

	return result
}