github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/providers/agent/mcorpc/golang/provision/provision.go (about)

     1  // Copyright (c) 2021, R.I. Pienaar and the Choria Project contributors
     2  //
     3  // SPDX-License-Identifier: Apache-2.0
     4  
     5  package provision
     6  
     7  import (
     8  	"context"
     9  	"runtime"
    10  	"sync"
    11  	"time"
    12  
    13  	"github.com/choria-io/go-choria/build"
    14  	"github.com/choria-io/go-choria/inter"
    15  	"github.com/choria-io/go-choria/plugin"
    16  	"github.com/choria-io/go-choria/providers/agent/mcorpc"
    17  	"github.com/choria-io/go-choria/server"
    18  	"github.com/choria-io/go-choria/server/agents"
    19  	"github.com/sirupsen/logrus"
    20  )
    21  
    22  // Reply is a generic reply used by most actions
    23  type Reply struct {
    24  	Message string `json:"message"`
    25  }
    26  
    27  var (
    28  	// LockedError is the error that will be produced when another provisioning is busy provisioning this server
    29  	LockedError = "Could not obtain provisioning lock"
    30  	// LockWindow is how long provisioning locks will be held or extended for
    31  	LockWindow = 2 * time.Minute
    32  
    33  	mu           = &sync.Mutex{}
    34  	allowRestart = true
    35  	ecdhPublic   []byte
    36  	ecdhPrivate  []byte
    37  	log          *logrus.Entry
    38  	lockMu       sync.Mutex
    39  	lockedUntil  time.Time
    40  	lockedBy     string
    41  )
    42  
    43  var metadata = &agents.Metadata{
    44  	Name:        "choria_provision",
    45  	Description: "Choria Provisioner",
    46  	Author:      "R.I.Pienaar <rip@devco.net>",
    47  	Version:     build.Version,
    48  	License:     build.License,
    49  	Timeout:     20,
    50  	URL:         "https://choria.io",
    51  }
    52  
    53  var restartCb = restart
    54  var shutdownCb = restartViaExit
    55  
    56  func init() {
    57  	switch {
    58  	case runtime.GOOS == "windows":
    59  		SetRestartAction(restartViaExit)
    60  	default:
    61  		SetRestartAction(restart)
    62  	}
    63  }
    64  
    65  // New creates a new instance of the agent
    66  func New(mgr server.AgentManager) (agents.Agent, error) {
    67  	log = mgr.Logger()
    68  
    69  	agent := mcorpc.New("choria_provision", metadata, mgr.Choria(), log)
    70  
    71  	agent.SetActivationChecker(func() bool {
    72  		return mgr.Choria().SupportsProvisioning()
    73  	})
    74  
    75  	agent.MustRegisterAction("configure", lockedAction(configureAction))
    76  	agent.MustRegisterAction("gen25519", lockedAction(ed25519Action))
    77  	agent.MustRegisterAction("gencsr", lockedAction(csrAction))
    78  	agent.MustRegisterAction("jwt", lockedAction(jwtAction))
    79  	agent.MustRegisterAction("release_update", lockedAction(releaseUpdateAction))
    80  	agent.MustRegisterAction("reprovision", reprovisionAction)
    81  	agent.MustRegisterAction("restart", restartAction)
    82  	agent.MustRegisterAction("shutdown", shutdownAction)
    83  
    84  	return agent, nil
    85  }
    86  
    87  // middleware to ensure actions will only be run by a caller who are currently provisioning the server
    88  func lockedAction(next mcorpc.Action) mcorpc.Action {
    89  	return func(ctx context.Context, req *mcorpc.Request, reply *mcorpc.Reply, agent *mcorpc.Agent, conn inter.ConnectorInfo) {
    90  		if locked, _ := lockFor(req.SenderID, req.CallerID); !locked {
    91  			abort(LockedError, reply)
    92  			return
    93  		}
    94  
    95  		next(ctx, req, reply, agent, conn)
    96  	}
    97  }
    98  
    99  // ChoriaPlugin creates the choria plugin hooks
   100  func ChoriaPlugin() plugin.Pluggable {
   101  	return mcorpc.NewChoriaAgentPlugin(metadata, New)
   102  }
   103  
   104  // SetRestartAction sets a custom restart function to call than the default that
   105  // causes a os.Exec() to be issued replacing the running instance with a new
   106  // process on the old pid
   107  func SetRestartAction(f func(splay time.Duration, si agents.ServerInfoSource, log *logrus.Entry)) {
   108  	mu.Lock()
   109  	restartCb = f
   110  	mu.Unlock()
   111  }
   112  
   113  // SetShutdownAction sets a custom shutdown function to call than the default that
   114  // causes a os.Exit(0) to be called
   115  func SetShutdownAction(f func(splay time.Duration, si agents.ServerInfoSource, log *logrus.Entry)) {
   116  	mu.Lock()
   117  	shutdownCb = f
   118  	mu.Unlock()
   119  }