ziti/controller/agent.go

package controller

import (
	"fmt"
	"github.com/openziti/ziti/common/pb/cmd_pb"
	"io"
	"net"
	"time"

	"github.com/michaelquigley/pfxlog"
	"github.com/openziti/channel/v2"
	"github.com/openziti/channel/v2/protobufs"
	"github.com/openziti/ziti/common/handler_common"
	"github.com/openziti/ziti/common/pb/mgmt_pb"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	AgentAppId byte = 1

	AgentIdHeader      = 10
	AgentAddrHeader    = 11
	AgentIsVoterHeader = 12
)

func (self *Controller) RegisterAgentBindHandler(bindHandler channel.BindHandler) {
	self.agentBindHandlers = append(self.agentBindHandlers, bindHandler)
}

func (self *Controller) bindAgentChannel(binding channel.Binding) error {
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_SnapshotDbRequestType), self.agentOpSnapshotDb)
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_RaftListMembersRequestType), self.agentOpRaftList)
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_RaftAddPeerRequestType), self.agentOpRaftAddPeer)
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_RaftRemovePeerRequestType), self.agentOpRaftRemovePeer)
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_RaftTransferLeadershipRequestType), self.agentOpRaftTransferLeadership)
	binding.AddReceiveHandlerF(int32(mgmt_pb.ContentType_RaftInitFromDb), self.agentOpInitFromDb)

	for _, bh := range self.agentBindHandlers {
		if err := binding.Bind(bh); err != nil {
			return err
		}
	}
	return nil
}

func (self *Controller) HandleCustomAgentAsyncOp(conn net.Conn) error {
	logrus.Debug("received agent operation request")

	appIdBuf := []byte{0}
	_, err := io.ReadFull(conn, appIdBuf)
	if err != nil {
		return err
	}
	appId := appIdBuf[0]

	if appId != AgentAppId {
		logrus.WithField("appId", appId).Debug("invalid app id on agent request")
		return errors.New("invalid operation for controller")
	}

	options := channel.DefaultOptions()
	options.ConnectTimeout = time.Second
	listener := channel.NewExistingConnListener(self.config.Id, conn, nil)
	_, err = channel.NewChannel("agent", listener, channel.BindHandlerF(self.bindAgentChannel), options)
	return err
}

func (self *Controller) agentOpSnapshotDb(m *channel.Message, ch channel.Channel) {
	log := pfxlog.Logger()
	if err := self.network.SnapshotDatabase(); err != nil {
		log.WithError(err).Error("failed to snapshot db")
		handler_common.SendOpResult(m, ch, "db.snapshot", err.Error(), false)
	} else {
		handler_common.SendOpResult(m, ch, "db.snapshot", "", true)
	}
}

func (self *Controller) agentOpRaftList(m *channel.Message, ch channel.Channel) {
	members, err := self.raftController.ListMembers()
	if err != nil {
		handler_common.SendOpResult(m, ch, "raft.list", err.Error(), false)
	}

	result := &mgmt_pb.RaftMemberListResponse{}
	for _, member := range members {
		result.Members = append(result.Members, &mgmt_pb.RaftMember{
			Id:          member.Id,
			Addr:        member.Addr,
			IsVoter:     member.Voter,
			IsLeader:    member.Leader,
			Version:     member.Version,
			IsConnected: member.Connected,
		})
	}

	if err = protobufs.MarshalTyped(result).ReplyTo(m).WithTimeout(time.Second).Send(ch); err != nil {
		pfxlog.Logger().WithError(err).Error("failure sending raft member list response")
	}
}

func (self *Controller) agentOpRaftAddPeer(m *channel.Message, ch channel.Channel) {
	addr, found := m.GetStringHeader(AgentAddrHeader)
	if !found {
		handler_common.SendOpResult(m, ch, "raft.join", "address not supplied", false)
		return
	}

	id, found := m.GetStringHeader(AgentIdHeader)
	if !found {
		peerId, peerAddr, err := self.raftController.Mesh.GetPeerInfo(addr, 15*time.Second)
		if err != nil {
			errMsg := fmt.Sprintf("id not supplied and unable to retrieve [%v]", err.Error())
			handler_common.SendOpResult(m, ch, "raft.join", errMsg, false)
			return
		}
		id = string(peerId)
		addr = string(peerAddr)
	}

	isVoter, found := m.GetBoolHeader(AgentIsVoterHeader)
	if !found {
		isVoter = true
	}

	req := &cmd_pb.AddPeerRequest{
		Addr:    addr,
		Id:      id,
		IsVoter: isVoter,
	}

	if err := self.raftController.Join(req); err != nil {
		handler_common.SendOpResult(m, ch, "cluster.add-peer", err.Error(), false)
		return
	}
	handler_common.SendOpResult(m, ch, "cluster.add-peer", fmt.Sprintf("success, added %v at %v to cluster", id, addr), true)
}

func (self *Controller) agentOpRaftRemovePeer(m *channel.Message, ch channel.Channel) {
	id, found := m.GetStringHeader(AgentIdHeader)
	if !found {
		handler_common.SendOpResult(m, ch, "cluster.remove-peer", "id not supplied", false)
		return
	}

	req := &cmd_pb.RemovePeerRequest{
		Id: id,
	}

	if err := self.raftController.HandleRemovePeer(req); err != nil {
		handler_common.SendOpResult(m, ch, "cluster.remove-peer", err.Error(), false)
		return
	}
	handler_common.SendOpResult(m, ch, "cluster.remove-peer", fmt.Sprintf("success, removed %v from cluster", id), true)
}

func (self *Controller) agentOpRaftTransferLeadership(m *channel.Message, ch channel.Channel) {
	id, _ := m.GetStringHeader(AgentIdHeader)
	req := &cmd_pb.TransferLeadershipRequest{
		Id: id,
	}

	if err := self.raftController.HandleTransferLeadership(req); err != nil {
		handler_common.SendOpResult(m, ch, "cluster.transfer-leadership", err.Error(), false)
		return
	}
	handler_common.SendOpResult(m, ch, "cluster.transfer-leadership", "success", true)
}

func (self *Controller) agentOpInitFromDb(m *channel.Message, ch channel.Channel) {
	sourceDbPath := string(m.Body)
	if len(sourceDbPath) == 0 {
		handler_common.SendOpResult(m, ch, "raft.initFromDb", "source db not supplied", false)
		return
	}

	if err := self.InitializeRaftFromBoltDb(sourceDbPath); err != nil {
		handler_common.SendOpResult(m, ch, "raft.initFromDb", err.Error(), false)
		return
	}
	handler_common.SendOpResult(m, ch, "raft.initFromDb", fmt.Sprintf("success, initialized from [%v]", sourceDbPath), true)
}