diff --git a/README.md b/README.md index 4bfa51c..65a6439 100644 --- a/README.md +++ b/README.md @@ -3,8 +3,9 @@ Core of [Opifica](https://github.com/opifica), an open-source CI/CD system built in the EU where every job runs in a fresh virtual machine. -> **Status:** early development. The pipeline language, the API definitions and the worker kit's warm pool -> work and are tested. The controller, agent and signer are skeletons. +> **Status:** early development. Working and tested: the pipeline language, the worker kit's warm pool, and +> workers joining and staying connected to the controller. Scheduling and the agent come next; see +> [docs/ROADMAP.md](docs/ROADMAP.md). > See the [organisation page](https://github.com/opifica) for the architecture. ## What lives here @@ -15,9 +16,10 @@ where every job runs in a fresh virtual machine. | `schema/` | JSON Schema for pipeline files, for editor autocompletion | Working | | `cmd/opifica` | CLI: `validate`, `plan`, `schema` | Working | | `proto/` | API definitions: worker ↔ controller, worker ↔ guest agent (Protobuf, ConnectRPC) | Draft | -| `pkg/worker` | Shared worker kit: `Hypervisor` interface, warm VM pool, image verification | Working | +| `pkg/worker` | Shared worker kit: `Hypervisor` interface, warm VM pool, image verification, controller client | Working | | `pkg/worker/workertest` | Fake hypervisor for testing worker code without VMs | Working | -| `cmd/controller` | API, Git providers, pipeline engine, scheduler, secrets broker | Skeleton | +| `cmd/controller` | Join tokens, worker joining over TLS, the Connect stream, worker revocation | Working | +| | Scheduler, runs, logs, secrets, Git providers | Not started | | `cmd/agent` | Runs inside each VM, executes steps, streams logs and results | Skeleton | | `cmd/signer` | Signing service for release artifacts | Not started | @@ -111,9 +113,22 @@ broken.yaml:11:5: unknown field "secret" Requires Go 1.27 or later. ```sh -make tools # install buf and the protoc plugins (once) +make tools # install buf, the protoc plugins and sqlc (once) make # lint, test and build into bin/ -make generate # after changing proto/; generated code in gen/ is committed +make generate # after changing proto/ or SQL queries; generated code is committed +``` + +Database tests need PostgreSQL. Either set `OPIFICA_TEST_DATABASE_URL` to a server the tests may create databases +on, or install the PostgreSQL server binaries (e.g. `apt install postgresql-17`): the tests then start their own +throwaway server. Without either, those tests are skipped. + +### Trying the controller + +```sh +export OPIFICA_DATABASE_URL=postgres://localhost/opifica +go run ./cmd/controller token create -pool trusted # prints a join token and the CA fingerprint +go run ./cmd/controller serve # TLS on :8443 +go run ./cmd/controller worker list ``` ## Stack diff --git a/cmd/controller/main.go b/cmd/controller/main.go index b11d822..90ba172 100644 --- a/cmd/controller/main.go +++ b/cmd/controller/main.go @@ -2,28 +2,27 @@ // integration, pipeline engine, scheduler and secrets broker. // // Implemented so far: the database, the built-in certificate authority, join -// tokens and Join. Every other RPC returns "unimplemented". +// tokens, Join and the Connect stream (sessions, heartbeats, revocation). +// Scheduling comes next. package main import ( "context" - "crypto/tls" "errors" "flag" "fmt" "io" "log/slog" - "net/http" "os" "os/signal" "strings" "syscall" + "text/tabwriter" "time" - "connectrpc.com/connect" + "github.com/google/uuid" - opificav1 "github.com/opifica/opifica/gen/opifica/v1" - "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/controller/ca" "github.com/opifica/opifica/internal/controller/store" "github.com/opifica/opifica/internal/controller/workers" @@ -36,6 +35,8 @@ const usage = `Usage: controller [options] Commands: serve Run the controller token create Create a one-time join token for a new build machine + worker list List workers + worker revoke Revoke a worker; it is disconnected and can't reconnect version Print the version Common options (or environment variables): @@ -57,6 +58,10 @@ func run(args []string, stdout, stderr io.Writer) int { return serve(args[1:], stderr) case args[0] == "token" && len(args) > 1 && args[1] == "create": return tokenCreate(args[2:], stdout, stderr) + case args[0] == "worker" && len(args) > 1 && args[1] == "list": + return workerList(args[2:], stdout, stderr) + case args[0] == "worker" && len(args) > 1 && args[1] == "revoke": + return workerRevoke(args[2:], stdout, stderr) case args[0] == "version": fmt.Fprintln(stdout, version.Version) return 0 @@ -121,42 +126,19 @@ func serve(args []string, stderr io.Writer) int { defer closeFn() svc.Log = log - tlsCfg, err := tlsConfig(svc.CA, splitList(*hosts)) + tlsCfg, err := api.TLSConfig(svc.CA, splitList(*hosts)) if err != nil { log.Error("issuing server certificate", "err", err) return 1 } log.Info("certificate authority ready", "fingerprint", svc.CA.Fingerprint()) - if err := server.RunTLS(ctx, *addr, newHandler(svc), tlsCfg, log); err != nil { + if err := server.RunTLS(ctx, *addr, api.NewHandler(svc), tlsCfg, log); err != nil { log.Error("server stopped", "err", err) return 1 } return 0 } -// tlsConfig issues a fresh server certificate and accepts, but doesn't yet -// require, worker client certificates: Join is called before a worker has -// one. RPCs that need a worker identity check it themselves. -func tlsConfig(authority *ca.CA, hosts []string) (*tls.Config, error) { - certPEM, keyPEM, err := authority.IssueServerCert(hosts) - if err != nil { - return nil, err - } - // Send the CA too, so a joining worker can check it against the pinned - // fingerprint before it has the CA bundle. - certPEM = append(certPEM, authority.CertPEM()...) - pair, err := tls.X509KeyPair(certPEM, keyPEM) - if err != nil { - return nil, err - } - return &tls.Config{ - Certificates: []tls.Certificate{pair}, - ClientCAs: authority.Pool(), - ClientAuth: tls.VerifyClientCertIfGiven, - MinVersion: tls.VersionTLS13, - }, nil -} - func defaultHosts() string { hosts := []string{"localhost", "127.0.0.1", "::1"} if h, err := os.Hostname(); err == nil && h != "" { @@ -175,27 +157,6 @@ func splitList(s string) []string { return out } -func newHandler(svc *workers.Service) http.Handler { - mux := http.NewServeMux() - mux.HandleFunc("GET /healthz", server.Health) - mux.Handle(opificav1connect.NewWorkerServiceHandler(&workerHandler{svc: svc})) - return mux -} - -// workerHandler implements the API workers connect to. -type workerHandler struct { - opificav1connect.UnimplementedWorkerServiceHandler - svc *workers.Service -} - -func (h *workerHandler) Join(ctx context.Context, req *connect.Request[opificav1.JoinRequest]) (*connect.Response[opificav1.JoinResponse], error) { - resp, err := h.svc.Join(ctx, req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - type labelFlag map[string]string func (l labelFlag) String() string { return "" } @@ -252,3 +213,74 @@ The token is shown only once. `, *pool, tok.ExpiresAt.Format(time.RFC3339), tok.Secret, tok.CAFingerprint) return 0 } + +func workerList(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("worker list", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + if err := fs.Parse(args); err != nil { + return 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + svc, closeFn, err := c.open(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + defer closeFn() + ws, err := svc.Store.ListWorkers(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + tw := tabwriter.NewWriter(stdout, 0, 4, 2, ' ', 0) + fmt.Fprintln(tw, "ID\tHOSTNAME\tOS\tARCH\tPOOL\tJOBS\tLAST SEEN\tSTATUS") + for _, w := range ws { + seen, status := "never", "active" + if w.LastSeenAt != nil { + seen = time.Since(*w.LastSeenAt).Round(time.Second).String() + " ago" + } + if w.RevokedAt != nil { + status = "revoked" + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%d\t%s\t%s\n", w.ID, w.Hostname, w.Os, w.Arch, w.Pool, w.MaxConcurrentJobs, seen, status) + } + tw.Flush() + return 0 +} + +func workerRevoke(args []string, stdout, stderr io.Writer) int { + fs := flag.NewFlagSet("worker revoke", flag.ContinueOnError) + fs.SetOutput(stderr) + var c common + c.register(fs) + fs.Usage = func() { fmt.Fprintln(stderr, "Usage: controller worker revoke [options] ") } + if err := fs.Parse(args); err != nil { + return 2 + } + if fs.NArg() != 1 { + fs.Usage() + return 2 + } + id, err := uuid.Parse(fs.Arg(0)) + if err != nil { + fmt.Fprintf(stderr, "controller: invalid worker id %q\n", fs.Arg(0)) + return 2 + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + svc, closeFn, err := c.open(ctx) + if err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + defer closeFn() + if err := svc.Revoke(ctx, id); err != nil { + fmt.Fprintf(stderr, "controller: %v\n", err) + return 1 + } + fmt.Fprintf(stdout, "Worker %s revoked. A connected controller disconnects it within %v.\n", id, workers.DefaultTouchInterval) + return 0 +} diff --git a/cmd/controller/main_test.go b/cmd/controller/main_test.go index 9399a12..3648b9d 100644 --- a/cmd/controller/main_test.go +++ b/cmd/controller/main_test.go @@ -22,6 +22,7 @@ import ( opificav1 "github.com/opifica/opifica/gen/opifica/v1" "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" "github.com/opifica/opifica/internal/server" "github.com/opifica/opifica/internal/testutil/pgtest" "github.com/opifica/opifica/pkg/worker" @@ -56,7 +57,7 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { } defer closeFn() svc.Log = slog.New(slog.NewTextHandler(io.Discard, nil)) - tlsCfg, err := tlsConfig(svc.CA, []string{"127.0.0.1"}) + tlsCfg, err := api.TLSConfig(svc.CA, []string{"127.0.0.1"}) if err != nil { t.Fatal(err) } @@ -64,7 +65,7 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { if err != nil { t.Fatal(err) } - go server.ServeTLS(ctx, ln, newHandler(svc), tlsCfg, svc.Log) + go server.ServeTLS(ctx, ln, api.NewHandler(svc), tlsCfg, svc.Log) // The worker has no CA yet: it accepts the controller only if the // certificate chains to a CA with the pinned fingerprint. @@ -94,11 +95,11 @@ func TestTokenCreateAndJoinOverTLS(t *testing.T) { t.Errorf("Join response = %+v", resp.Msg) } - // Still unimplemented: the stream comes next. + // Connect without a client certificate is refused. stream := wc.Connect(ctx) stream.Send(&opificav1.ConnectRequest{}) - if _, err := stream.Receive(); connect.CodeOf(err) != connect.CodeUnimplemented { - t.Errorf("Connect: err = %v, want unimplemented", err) + if _, err := stream.Receive(); connect.CodeOf(err) != connect.CodeUnauthenticated { + t.Errorf("Connect without a certificate: err = %v, want unauthenticated", err) } stream.CloseRequest() stream.CloseResponse() diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 2c63900..4fdf977 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -29,7 +29,8 @@ result are stored. Tested in ordinary CI, using the fake hypervisor with a real - [x] **Certificate authority:** built-in CA that issues worker client certificates and the server certificate - [x] **Join:** admins create one-time join tokens (`controller token create`); `Join` checks the token and signs the worker's CSR; first contact is verified against the CA fingerprint printed with the token -- [ ] **Connect stream:** Hello, heartbeats, worker registry, disconnect and reconnect handling +- [x] **Connect stream:** Hello, heartbeats, worker registry, disconnect and reconnect handling, revocation + (`controller worker list`, `controller worker revoke`) - [ ] **Leases:** job assignment with a lease renewed by heartbeats; expired leases put the job back in the queue - [ ] **Scheduler:** match planned jobs to workers by OS, architecture, image, labels and pool; respect `needs` - [ ] **Runs:** create a run from a pipeline file and an event; track job and run state @@ -37,12 +38,14 @@ result are stored. Tested in ordinary CI, using the fake hypervisor with a real - [ ] **Secrets:** encrypted store; requests checked against the job token and policy - [ ] **Admin config:** static file for pools, routing (event → pool) and secret policy - [x] **TLS** on the worker API, with worker client certificates accepted -- [ ] **mTLS required** for `Connect`, and certificate renewal +- [x] **mTLS required** for `Connect`; the certificate must be the worker's current one +- [ ] Certificate renewal before the 30-day expiry - [ ] **Admin API + CLI:** `opifica worker token create`, `opifica runs list`, `opifica logs ` ### Worker kit and agent -- [ ] Controller client: join, connect, reconnect with backoff, heartbeats from pool stats +- [x] Controller client: join (saving the identity to disk), connect, reconnect with backoff, heartbeats +- [ ] Heartbeats filled from the warm pool and job runner - [ ] Job runner: lease a VM → checkout → run steps through the agent → stream masked logs → collect reports and artifacts → destroy the VM - [ ] Local Git mirror per repository diff --git a/internal/controller/api/api.go b/internal/controller/api/api.go new file mode 100644 index 0000000..a2dba33 --- /dev/null +++ b/internal/controller/api/api.go @@ -0,0 +1,66 @@ +// Package api assembles the controller's HTTP handler and TLS settings. +package api + +import ( + "context" + "crypto/tls" + "net/http" + + "connectrpc.com/connect" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/server" +) + +// NewHandler returns the controller's HTTP handler. +func NewHandler(svc *workers.Service) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("GET /healthz", server.Health) + mux.Handle(opificav1connect.NewWorkerServiceHandler(&workerHandler{svc: svc})) + return workers.WithPeerCert(mux) +} + +// TLSConfig issues a fresh server certificate for hosts. Worker client +// certificates are verified when presented but not required at the TLS +// layer, because Join is called before a worker has one; Connect requires +// one itself. +func TLSConfig(authority *ca.CA, hosts []string) (*tls.Config, error) { + certPEM, keyPEM, err := authority.IssueServerCert(hosts) + if err != nil { + return nil, err + } + // Send the CA too, so a joining worker can check it against the pinned + // fingerprint before it has the CA bundle. + certPEM = append(certPEM, authority.CertPEM()...) + pair, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + return nil, err + } + return &tls.Config{ + Certificates: []tls.Certificate{pair}, + ClientCAs: authority.Pool(), + ClientAuth: tls.VerifyClientCertIfGiven, + MinVersion: tls.VersionTLS13, + }, nil +} + +// workerHandler implements the API workers connect to. +type workerHandler struct { + opificav1connect.UnimplementedWorkerServiceHandler + svc *workers.Service +} + +func (h *workerHandler) Join(ctx context.Context, req *connect.Request[opificav1.JoinRequest]) (*connect.Response[opificav1.JoinResponse], error) { + resp, err := h.svc.Join(ctx, req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (h *workerHandler) Connect(ctx context.Context, stream *connect.BidiStream[opificav1.ConnectRequest, opificav1.ConnectResponse]) error { + return h.svc.Connect(ctx, workers.PeerCert(ctx), stream) +} diff --git a/internal/controller/store/db/workers.sql.go b/internal/controller/store/db/workers.sql.go index dfc2409..fed158f 100644 --- a/internal/controller/store/db/workers.sql.go +++ b/internal/controller/store/db/workers.sql.go @@ -267,11 +267,39 @@ func (q *Queries) SetJoinTokenWorker(ctx context.Context, arg SetJoinTokenWorker return err } -const touchWorker = `-- name: TouchWorker :exec -UPDATE workers SET last_seen_at = now() WHERE id = $1 +const touchWorker = `-- name: TouchWorker :one +UPDATE workers SET last_seen_at = now() WHERE id = $1 RETURNING revoked_at ` -func (q *Queries) TouchWorker(ctx context.Context, id uuid.UUID) error { - _, err := q.db.Exec(ctx, touchWorker, id) +// Records that the worker is alive and returns its revocation time, so a +// connected worker that has been revoked can be disconnected. +func (q *Queries) TouchWorker(ctx context.Context, id uuid.UUID) (*time.Time, error) { + row := q.db.QueryRow(ctx, touchWorker, id) + var revoked_at *time.Time + err := row.Scan(&revoked_at) + return revoked_at, err +} + +const updateWorkerHello = `-- name: UpdateWorkerHello :exec +UPDATE workers +SET hostname = $2, version = $3, max_concurrent_jobs = $4, last_seen_at = now() +WHERE id = $1 +` + +type UpdateWorkerHelloParams struct { + ID uuid.UUID + Hostname string + Version string + MaxConcurrentJobs int32 +} + +// Facts a worker reports each time it connects. +func (q *Queries) UpdateWorkerHello(ctx context.Context, arg UpdateWorkerHelloParams) error { + _, err := q.db.Exec(ctx, updateWorkerHello, + arg.ID, + arg.Hostname, + arg.Version, + arg.MaxConcurrentJobs, + ) return err } diff --git a/internal/controller/store/queries/workers.sql b/internal/controller/store/queries/workers.sql index 5419fbf..83a4031 100644 --- a/internal/controller/store/queries/workers.sql +++ b/internal/controller/store/queries/workers.sql @@ -30,8 +30,16 @@ SELECT * FROM workers WHERE id = $1; -- name: ListWorkers :many SELECT * FROM workers ORDER BY created_at; --- name: TouchWorker :exec -UPDATE workers SET last_seen_at = now() WHERE id = $1; +-- name: TouchWorker :one +-- Records that the worker is alive and returns its revocation time, so a +-- connected worker that has been revoked can be disconnected. +UPDATE workers SET last_seen_at = now() WHERE id = $1 RETURNING revoked_at; + +-- name: UpdateWorkerHello :exec +-- Facts a worker reports each time it connects. +UPDATE workers +SET hostname = $2, version = $3, max_concurrent_jobs = $4, last_seen_at = now() +WHERE id = $1; -- name: RevokeWorker :one UPDATE workers SET revoked_at = now() WHERE id = $1 AND revoked_at IS NULL RETURNING *; diff --git a/internal/controller/workers/connect.go b/internal/controller/workers/connect.go new file mode 100644 index 0000000..99151d1 --- /dev/null +++ b/internal/controller/workers/connect.go @@ -0,0 +1,300 @@ +package workers + +import ( + "context" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sync" + "time" + + "connectrpc.com/connect" + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/store" + "github.com/opifica/opifica/internal/controller/store/db" +) + +// Defaults for ConnectConfig. +const ( + DefaultHelloTimeout = 10 * time.Second + DefaultHeartbeatTimeout = 30 * time.Second + DefaultTouchInterval = 30 * time.Second +) + +// ConnectConfig tunes the Connect stream. Zero values use the defaults. +type ConnectConfig struct { + // HelloTimeout is how long a new stream may take to send Hello. + HelloTimeout time.Duration + // HeartbeatTimeout ends a session that has sent nothing for this long. + // Workers send heartbeats every 10s by default. + HeartbeatTimeout time.Duration + // TouchInterval is how often last_seen is written to the database and + // revocation is checked. + TouchInterval time.Duration +} + +func (c ConnectConfig) withDefaults() ConnectConfig { + if c.HelloTimeout <= 0 { + c.HelloTimeout = DefaultHelloTimeout + } + if c.HeartbeatTimeout <= 0 { + c.HeartbeatTimeout = DefaultHeartbeatTimeout + } + if c.TouchInterval <= 0 { + c.TouchInterval = DefaultTouchInterval + } + return c +} + +// Stream is the server side of a Connect stream. +type Stream interface { + Receive() (*opificav1.ConnectRequest, error) + Send(*opificav1.ConnectResponse) error +} + +// MessageHandler receives what workers send about jobs. It is nil until the +// scheduler exists; until then those messages are logged and dropped. +type MessageHandler interface { + HandleWorkerMessage(ctx context.Context, worker uuid.UUID, msg *opificav1.ConnectRequest) error +} + +type peerCertKey struct{} + +// WithPeerCert stores the verified TLS client certificate, if any, in the +// request context so RPC handlers can identify the worker. +func WithPeerCert(h http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.TLS != nil && len(r.TLS.VerifiedChains) > 0 { + r = r.WithContext(context.WithValue(r.Context(), peerCertKey{}, r.TLS.VerifiedChains[0][0])) + } + h.ServeHTTP(w, r) + }) +} + +// PeerCert returns the verified client certificate stored by WithPeerCert. +func PeerCert(ctx context.Context) *x509.Certificate { + c, _ := ctx.Value(peerCertKey{}).(*x509.Certificate) + return c +} + +var ( + errNoCert = connect.NewError(connect.CodeUnauthenticated, errors.New("a worker client certificate is required")) + errRevoked = connect.NewError(connect.CodePermissionDenied, errors.New("worker has been revoked")) + errUnknown = connect.NewError(connect.CodeUnauthenticated, errors.New("unknown worker")) + errSuperseded = connect.NewError(connect.CodeUnauthenticated, errors.New("certificate has been superseded")) + + errRevokedCause = errors.New("revoked") +) + +// Connect runs one worker's stream until it ends. The worker is identified +// by its client certificate, never by what it says in Hello. +func (s *Service) Connect(ctx context.Context, cert *x509.Certificate, stream Stream) error { + cfg := s.StreamConfig.withDefaults() + if cert == nil { + return errNoCert + } + idStr, err := ca.WorkerID(cert) + if err != nil { + return errNoCert + } + id, err := uuid.Parse(idStr) + if err != nil { + return errNoCert + } + w, err := s.Store.GetWorker(ctx, id) + switch { + case store.IsNotFound(err): + return errUnknown + case err != nil: + s.log().Error("loading worker", "worker", id, "err", err) + return connect.NewError(connect.CodeUnavailable, errors.New("try again later")) + case w.RevokedAt != nil: + return errRevoked + case w.CertSerial != cert.SerialNumber.Text(16): + return errSuperseded + } + + // Deferred first so it runs last: the writer must be done with the stream + // before the handler returns. + var writer sync.WaitGroup + defer writer.Wait() + ctx, cancel := context.WithCancelCause(ctx) + defer cancel(nil) + + // Receive in the background so timeouts and cancellation can interrupt + // waiting for the next message. + msgs := make(chan *opificav1.ConnectRequest) + recvErr := make(chan error, 1) + go func() { + for { + m, err := stream.Receive() + if err != nil { + recvErr <- err + return + } + select { + case msgs <- m: + case <-ctx.Done(): + return + } + } + }() + + hello, err := s.awaitHello(ctx, msgs, recvErr, cfg.HelloTimeout) + if err != nil { + return err + } + if hello.GetWorkerId() != id.String() { + return connect.NewError(connect.CodeInvalidArgument, + fmt.Errorf("hello names worker %q but the certificate is for %s", hello.GetWorkerId(), id)) + } + info := hello.GetInfo() + maxJobs := int(info.GetMaxConcurrentJobs()) + if maxJobs < 1 || maxJobs > maxJobsPerHost { + return connect.NewError(connect.CodeInvalidArgument, + fmt.Errorf("max_concurrent_jobs must be between 1 and %d", maxJobsPerHost)) + } + if err := s.Store.UpdateWorkerHello(ctx, db.UpdateWorkerHelloParams{ + ID: id, Hostname: info.GetHostname(), Version: info.GetVersion(), MaxConcurrentJobs: int32(maxJobs), + }); err != nil { + s.log().Error("recording hello", "worker", id, "err", err) + return connect.NewError(connect.CodeUnavailable, errors.New("try again later")) + } + + var labels map[string]string + if err := json.Unmarshal(w.Labels, &labels); err != nil { + return connect.NewError(connect.CodeInternal, errors.New("stored worker labels are invalid")) + } + now := s.now() + sess := &Session{ + WorkerID: id, Pool: w.Pool, OS: w.Os, Arch: w.Arch, Labels: labels, MaxJobs: maxJobs, + ConnectedAt: now, lastSeen: now, running: hello.GetRunningJobIds(), + out: make(chan *opificav1.ConnectResponse, 16), ctx: ctx, cancel: cancel, + } + s.Sessions().register(sess) + defer s.Sessions().unregister(sess) + log := s.log().With("worker", id, "hostname", info.GetHostname()) + log.Info("worker connected", "pool", w.Pool, "os", w.Os, "arch", w.Arch, "max_jobs", maxJobs) + + // Writer: the only goroutine that sends on the stream. + writer.Add(1) + go func() { + defer writer.Done() + for { + select { + case m := <-sess.out: + if err := stream.Send(m); err != nil { + cancel(err) + return + } + case <-ctx.Done(): + return + } + } + }() + + idle := time.NewTimer(cfg.HeartbeatTimeout) + defer idle.Stop() + touch := time.NewTicker(cfg.TouchInterval) + defer touch.Stop() + + for { + select { + case m := <-msgs: + idle.Reset(cfg.HeartbeatTimeout) + sess.seen(s.now()) + if err := s.handle(ctx, sess, m); err != nil { + log.Info("worker disconnected", "reason", err) + return err + } + case err := <-recvErr: + if errors.Is(err, io.EOF) { + log.Info("worker disconnected", "reason", "closed by worker") + return nil + } + log.Info("worker disconnected", "reason", err) + return err + case <-idle.C: + log.Warn("worker disconnected", "reason", "no heartbeat") + return connect.NewError(connect.CodeDeadlineExceeded, + fmt.Errorf("no message for %v", cfg.HeartbeatTimeout)) + case <-touch.C: + revoked, err := s.Store.TouchWorker(ctx, id) + if err != nil && ctx.Err() == nil { + log.Error("recording heartbeat", "err", err) + continue + } + if revoked != nil { + log.Warn("worker disconnected", "reason", "revoked") + return errRevoked + } + case <-ctx.Done(): + cause := context.Cause(ctx) + log.Info("worker disconnected", "reason", cause) + if errors.Is(cause, errReplaced) { + return connect.NewError(connect.CodeAborted, errReplaced) + } + if errors.Is(cause, errRevokedCause) { + return errRevoked + } + return connect.NewError(connect.CodeCanceled, cause) + } + } +} + +func (s *Service) awaitHello(ctx context.Context, msgs <-chan *opificav1.ConnectRequest, recvErr <-chan error, timeout time.Duration) (*opificav1.Hello, error) { + t := time.NewTimer(timeout) + defer t.Stop() + select { + case m := <-msgs: + h := m.GetHello() + if h == nil { + return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("the first message must be hello")) + } + return h, nil + case err := <-recvErr: + return nil, err + case <-t.C: + return nil, connect.NewError(connect.CodeDeadlineExceeded, errors.New("no hello received")) + case <-ctx.Done(): + return nil, context.Cause(ctx) + } +} + +func (s *Service) handle(ctx context.Context, sess *Session, m *opificav1.ConnectRequest) error { + switch msg := m.GetMessage().(type) { + case *opificav1.ConnectRequest_Heartbeat: + sess.heartbeat(msg.Heartbeat, s.now()) + return nil + case *opificav1.ConnectRequest_Hello: + return connect.NewError(connect.CodeInvalidArgument, errors.New("hello may only be sent once")) + case *opificav1.ConnectRequest_JobUpdate, *opificav1.ConnectRequest_LogChunk, *opificav1.ConnectRequest_SecretRequest: + if s.Messages == nil { + s.log().Debug("dropping job message: no scheduler yet", "worker", sess.WorkerID) + return nil + } + return s.Messages.HandleWorkerMessage(ctx, sess.WorkerID, m) + case nil: + return connect.NewError(connect.CodeInvalidArgument, errors.New("empty message")) + } + return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("unknown message %T", m.GetMessage())) +} + +// Revoke marks a worker as revoked and disconnects it if this controller +// holds its session. Other controllers disconnect it at their next touch. +func (s *Service) Revoke(ctx context.Context, id uuid.UUID) error { + if _, err := s.Store.RevokeWorker(ctx, id); err != nil { + if store.IsNotFound(err) { + return fmt.Errorf("worker %s not found or already revoked", id) + } + return err + } + s.Sessions().Disconnect(id, errRevokedCause) + return nil +} diff --git a/internal/controller/workers/join.go b/internal/controller/workers/join.go index 256b5ed..06b60b7 100644 --- a/internal/controller/workers/join.go +++ b/internal/controller/workers/join.go @@ -14,6 +14,7 @@ import ( "maps" "regexp" "strings" + "sync" "time" "connectrpc.com/connect" @@ -49,6 +50,20 @@ type Service struct { CA *ca.CA Log *slog.Logger Now func() time.Time + + // StreamConfig tunes the Connect stream. + StreamConfig ConnectConfig + // Messages receives job messages from workers; nil drops them. + Messages MessageHandler + + sessionsOnce sync.Once + sessions *Registry +} + +// Sessions returns the registry of connected workers. +func (s *Service) Sessions() *Registry { + s.sessionsOnce.Do(func() { s.sessions = NewRegistry() }) + return s.sessions } func (s *Service) now() time.Time { diff --git a/internal/controller/workers/registry.go b/internal/controller/workers/registry.go new file mode 100644 index 0000000..f41863c --- /dev/null +++ b/internal/controller/workers/registry.go @@ -0,0 +1,169 @@ +package workers + +import ( + "context" + "errors" + "maps" + "slices" + "sync" + "time" + + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" +) + +// ErrNotConnected is returned when sending to a worker without a live session. +var ErrNotConnected = errors.New("worker is not connected") + +// Session is one live Connect stream from a worker. +type Session struct { + WorkerID uuid.UUID + Pool string + OS string + Arch string + Labels map[string]string + MaxJobs int + ConnectedAt time.Time + + out chan *opificav1.ConnectResponse + ctx context.Context + cancel context.CancelCauseFunc + + mu sync.Mutex + lastSeen time.Time + warmVMs uint32 + running []string +} + +// Send queues a message for the worker. It fails if the session ends first. +func (s *Session) Send(ctx context.Context, msg *opificav1.ConnectResponse) error { + select { + case s.out <- msg: + return nil + case <-s.ctx.Done(): + return ErrNotConnected + case <-ctx.Done(): + return ctx.Err() + } +} + +func (s *Session) heartbeat(hb *opificav1.Heartbeat, now time.Time) { + s.mu.Lock() + defer s.mu.Unlock() + s.lastSeen = now + s.warmVMs = hb.GetWarmVms() + s.running = slices.Clone(hb.GetRunningJobIds()) +} + +func (s *Session) seen(now time.Time) { + s.mu.Lock() + s.lastSeen = now + s.mu.Unlock() +} + +// SessionInfo is a snapshot of a session, for the scheduler and admin views. +type SessionInfo struct { + WorkerID uuid.UUID + Pool string + OS string + Arch string + Labels map[string]string + MaxJobs int + ConnectedAt time.Time + LastSeen time.Time + WarmVMs uint32 + RunningJobs []string +} + +func (s *Session) info() SessionInfo { + s.mu.Lock() + defer s.mu.Unlock() + return SessionInfo{ + WorkerID: s.WorkerID, Pool: s.Pool, OS: s.OS, Arch: s.Arch, + Labels: maps.Clone(s.Labels), MaxJobs: s.MaxJobs, ConnectedAt: s.ConnectedAt, + LastSeen: s.lastSeen, WarmVMs: s.warmVMs, RunningJobs: slices.Clone(s.running), + } +} + +// Registry holds the live session of every connected worker. +type Registry struct { + mu sync.Mutex + sessions map[uuid.UUID]*Session +} + +// NewRegistry returns an empty registry. +func NewRegistry() *Registry { + return &Registry{sessions: map[uuid.UUID]*Session{}} +} + +// errReplaced ends a session when the same worker connects again, e.g. after +// a network change left the old stream half-open. +var errReplaced = errors.New("replaced by a newer connection from the same worker") + +func (r *Registry) register(s *Session) { + r.mu.Lock() + old := r.sessions[s.WorkerID] + r.sessions[s.WorkerID] = s + r.mu.Unlock() + if old != nil { + old.cancel(errReplaced) + } +} + +func (r *Registry) unregister(s *Session) { + r.mu.Lock() + defer r.mu.Unlock() + if r.sessions[s.WorkerID] == s { + delete(r.sessions, s.WorkerID) + } +} + +// Get returns the live session for a worker. +func (r *Registry) Get(id uuid.UUID) (*Session, bool) { + r.mu.Lock() + defer r.mu.Unlock() + s, ok := r.sessions[id] + return s, ok +} + +// Send queues a message for a connected worker. +func (r *Registry) Send(ctx context.Context, id uuid.UUID, msg *opificav1.ConnectResponse) error { + s, ok := r.Get(id) + if !ok { + return ErrNotConnected + } + return s.Send(ctx, msg) +} + +// Disconnect ends a worker's session, if it has one. +func (r *Registry) Disconnect(id uuid.UUID, cause error) { + if s, ok := r.Get(id); ok { + s.cancel(cause) + } +} + +// List returns a snapshot of every connected worker, ordered by ID. +func (r *Registry) List() []SessionInfo { + r.mu.Lock() + ss := make([]*Session, 0, len(r.sessions)) + for _, s := range r.sessions { + ss = append(ss, s) + } + r.mu.Unlock() + out := make([]SessionInfo, len(ss)) + for i, s := range ss { + out[i] = s.info() + } + slices.SortFunc(out, func(a, b SessionInfo) int { return compareUUID(a.WorkerID, b.WorkerID) }) + return out +} + +func compareUUID(a, b uuid.UUID) int { + for i := range a { + if a[i] != b[i] { + return int(a[i]) - int(b[i]) + } + } + return 0 +} diff --git a/internal/controller/workers/registry_test.go b/internal/controller/workers/registry_test.go new file mode 100644 index 0000000..6fd45f5 --- /dev/null +++ b/internal/controller/workers/registry_test.go @@ -0,0 +1,72 @@ +package workers + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" +) + +func newSession(id uuid.UUID) *Session { + ctx, cancel := context.WithCancelCause(context.Background()) + return &Session{WorkerID: id, out: make(chan *opificav1.ConnectResponse, 1), ctx: ctx, cancel: cancel} +} + +func TestRegistryReplacesOldSession(t *testing.T) { + r := NewRegistry() + id := uuid.New() + first, second := newSession(id), newSession(id) + + r.register(first) + r.register(second) + if !errors.Is(context.Cause(first.ctx), errReplaced) { + t.Errorf("old session cause = %v, want errReplaced", context.Cause(first.ctx)) + } + if second.ctx.Err() != nil { + t.Error("new session was cancelled") + } + + // The old session's handler unregisters on its way out; that must not + // remove the new session. + r.unregister(first) + if s, ok := r.Get(id); !ok || s != second { + t.Error("unregistering the old session removed the new one") + } + r.unregister(second) + if _, ok := r.Get(id); ok { + t.Error("session still registered") + } +} + +func TestSessionSend(t *testing.T) { + r := NewRegistry() + id := uuid.New() + s := newSession(id) + r.register(s) + + msg := &opificav1.ConnectResponse{} + if err := r.Send(context.Background(), id, msg); err != nil { + t.Fatal(err) + } + if got := <-s.out; got != msg { + t.Error("wrong message queued") + } + + // Queue full and the caller gives up. + s.out <- msg + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + if err := r.Send(ctx, id, msg); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("send to full queue: err = %v", err) + } + + // Session ended while the sender waits. + r.Disconnect(id, errors.New("gone")) + if err := s.Send(context.Background(), msg); !errors.Is(err, ErrNotConnected) { + t.Errorf("send to ended session: err = %v", err) + } +} diff --git a/internal/integration/connect_test.go b/internal/integration/connect_test.go new file mode 100644 index 0000000..aa8060f --- /dev/null +++ b/internal/integration/connect_test.go @@ -0,0 +1,361 @@ +// Package integration runs the controller and the worker kit against each +// other over real TLS connections and a real PostgreSQL database. +package integration + +import ( + "context" + "errors" + "io" + "log/slog" + "net" + "net/http" + "os" + "sync" + "testing" + "time" + + "connectrpc.com/connect" + "github.com/google/uuid" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" + "github.com/opifica/opifica/internal/controller/api" + "github.com/opifica/opifica/internal/controller/ca" + "github.com/opifica/opifica/internal/controller/store" + "github.com/opifica/opifica/internal/controller/workers" + "github.com/opifica/opifica/internal/server" + "github.com/opifica/opifica/internal/testutil/pgtest" + "github.com/opifica/opifica/pkg/worker" +) + +func TestMain(m *testing.M) { os.Exit(pgtest.Run(m)) } + +func quiet() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) } + +type env struct { + t *testing.T + svc *workers.Service + url string +} + +// newEnv starts a controller with short timeouts on a random port. +func newEnv(t *testing.T) *env { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + st, err := store.Open(ctx, pgtest.New(t)) + if err != nil { + t.Fatal(err) + } + authority, err := ca.LoadOrCreate(t.TempDir()) + if err != nil { + t.Fatal(err) + } + svc := &workers.Service{ + Store: st, CA: authority, Log: quiet(), + StreamConfig: workers.ConnectConfig{ + HelloTimeout: 2 * time.Second, + HeartbeatTimeout: 500 * time.Millisecond, + TouchInterval: 50 * time.Millisecond, + }, + } + tlsCfg, err := api.TLSConfig(authority, []string{"127.0.0.1"}) + if err != nil { + t.Fatal(err) + } + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan struct{}) + go func() { + server.ServeTLS(ctx, ln, api.NewHandler(svc), tlsCfg, quiet()) + close(done) + }() + t.Cleanup(func() { + cancel() + <-done + st.Close() + }) + return &env{t: t, svc: svc, url: "https://" + ln.Addr().String()} +} + +func info() *opificav1.WorkerInfo { + return &opificav1.WorkerInfo{ + Hostname: "build-1", Os: opificav1.OperatingSystem_OPERATING_SYSTEM_LINUX, + Arch: opificav1.Architecture_ARCHITECTURE_AMD64, MaxConcurrentJobs: 4, Version: "test", + Labels: map[string]string{"kvm": "true"}, + } +} + +// join creates a token for pool and joins a worker with it. +func (e *env) join(pool string) *worker.Identity { + e.t.Helper() + tok, err := e.svc.CreateToken(context.Background(), workers.TokenRequest{Pool: pool}) + if err != nil { + e.t.Fatal(err) + } + id, err := worker.Join(context.Background(), worker.JoinConfig{ + ControllerURL: e.url, Token: tok.Secret, CAFingerprint: tok.CAFingerprint, + StateDir: e.t.TempDir(), Info: info(), + }) + if err != nil { + e.t.Fatalf("Join: %v", err) + } + return id +} + +type runningClient struct { + *worker.Client + cancel context.CancelFunc + done chan error + msgs chan *opificav1.ConnectResponse +} + +func (e *env) start(id *worker.Identity, mutate func(*worker.ClientConfig)) *runningClient { + e.t.Helper() + rc := &runningClient{done: make(chan error, 1), msgs: make(chan *opificav1.ConnectResponse, 16)} + cfg := worker.ClientConfig{ + ControllerURL: e.url, Identity: id, Info: info(), + HeartbeatInterval: 50 * time.Millisecond, MaxBackoff: 100 * time.Millisecond, Logger: quiet(), + Status: func() (uint32, []string) { return 3, []string{"job-1"} }, + OnMessage: func(_ context.Context, m *opificav1.ConnectResponse) { rc.msgs <- m }, + } + if mutate != nil { + mutate(&cfg) + } + c, err := worker.NewClient(cfg) + if err != nil { + e.t.Fatal(err) + } + rc.Client = c + ctx, cancel := context.WithCancel(context.Background()) + rc.cancel = cancel + go func() { rc.done <- c.Run(ctx) }() + e.t.Cleanup(func() { + cancel() + <-rc.done + }) + return rc +} + +func workerHTTPClient(id *worker.Identity) *http.Client { + return &http.Client{Transport: &http.Transport{TLSClientConfig: id.TLSConfig(), ForceAttemptHTTP2: true}} +} + +func eventually(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for !cond() { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %s", what) + } + time.Sleep(5 * time.Millisecond) + } +} + +func (e *env) session(id *worker.Identity) (workers.SessionInfo, bool) { + for _, s := range e.svc.Sessions().List() { + if s.WorkerID.String() == id.WorkerID { + return s, true + } + } + return workers.SessionInfo{}, false +} + +func TestWorkerConnectsAndReportsStatus(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + e.start(id, nil) + + eventually(t, "a session with heartbeat status", func() bool { + s, ok := e.session(id) + return ok && s.WarmVMs == 3 && len(s.RunningJobs) == 1 + }) + s, _ := e.session(id) + if s.Pool != "trusted" || s.OS != "linux" || s.Arch != "amd64" || s.MaxJobs != 4 || s.Labels["kvm"] != "true" { + t.Errorf("session = %+v", s) + } + + eventually(t, "last_seen recorded", func() bool { + w, err := e.svc.Store.GetWorker(context.Background(), uuid.MustParse(id.WorkerID)) + return err == nil && w.LastSeenAt != nil && w.Version == "test" + }) +} + +func TestControllerSendsToWorker(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + rc := e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + + err := e.svc.Sessions().Send(context.Background(), uuid.MustParse(id.WorkerID), &opificav1.ConnectResponse{ + Message: &opificav1.ConnectResponse_Drain{Drain: &opificav1.Drain{Reason: "maintenance"}}, + }) + if err != nil { + t.Fatal(err) + } + select { + case m := <-rc.msgs: + if m.GetDrain().GetReason() != "maintenance" { + t.Errorf("worker got %v", m) + } + case <-time.After(5 * time.Second): + t.Fatal("worker never received the message") + } + + if err := e.svc.Sessions().Send(context.Background(), uuid.New(), &opificav1.ConnectResponse{}); !errors.Is(err, workers.ErrNotConnected) { + t.Errorf("send to unknown worker: err = %v", err) + } +} + +func TestWorkerReconnects(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + first, _ := e.session(id) + + e.svc.Sessions().Disconnect(first.WorkerID, errors.New("network blip")) + eventually(t, "a new session", func() bool { + s, ok := e.session(id) + return ok && s.ConnectedAt.After(first.ConnectedAt) + }) +} + +func TestSilentWorkerIsDropped(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + // Heartbeats far apart: the controller gives up after 500ms of silence. + e.start(id, func(c *worker.ClientConfig) { c.HeartbeatInterval = time.Hour }) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + first, _ := e.session(id) + + // The worker reconnects, is dropped again, and so on: each time a new session. + eventually(t, "the silent session to be replaced", func() bool { + s, ok := e.session(id) + return ok && s.ConnectedAt.After(first.ConnectedAt) + }) +} + +func TestRevokedWorkerIsDisconnectedAndStaysOut(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + rc := e.start(id, nil) + eventually(t, "session", func() bool { _, ok := e.session(id); return ok }) + + if err := e.svc.Revoke(context.Background(), uuid.MustParse(id.WorkerID)); err != nil { + t.Fatal(err) + } + select { + case err := <-rc.done: + if connect.CodeOf(err) != connect.CodePermissionDenied || !worker.IsPermanent(err) { + t.Errorf("Run returned %v, want permission denied", err) + } + rc.done <- err // for the cleanup + case <-time.After(5 * time.Second): + t.Fatal("revoked worker is still running") + } + if _, ok := e.session(id); ok { + t.Error("revoked worker still has a session") + } + + // Revoked on another controller (directly in the database): the + // periodic touch finds it. + id2 := e.join("trusted") + rc2 := e.start(id2, nil) + eventually(t, "second session", func() bool { _, ok := e.session(id2); return ok }) + if _, err := e.svc.Store.RevokeWorker(context.Background(), uuid.MustParse(id2.WorkerID)); err != nil { + t.Fatal(err) + } + select { + case err := <-rc2.done: + if !worker.IsPermanent(err) { + t.Errorf("Run returned %v, want a permanent error", err) + } + rc2.done <- err + case <-time.After(5 * time.Second): + t.Fatal("worker revoked in the database is still connected") + } +} + +func TestConnectRejects(t *testing.T) { + e := newEnv(t) + id := e.join("trusted") + ctx := context.Background() + + connectWith := func(t *testing.T, hello *opificav1.Hello) error { + t.Helper() + client := opificav1connect.NewWorkerServiceClient(workerHTTPClient(id), e.url) + stream := client.Connect(ctx) + defer stream.CloseResponse() + stream.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Hello{Hello: hello}}) + _, err := stream.Receive() + return err + } + + t.Run("hello for another worker", func(t *testing.T) { + err := connectWith(t, &opificav1.Hello{WorkerId: uuid.NewString(), Info: info()}) + if connect.CodeOf(err) != connect.CodeInvalidArgument { + t.Errorf("err = %v, want invalid argument", err) + } + }) + t.Run("bad capacity", func(t *testing.T) { + bad := info() + bad.MaxConcurrentJobs = 0 + err := connectWith(t, &opificav1.Hello{WorkerId: id.WorkerID, Info: bad}) + if connect.CodeOf(err) != connect.CodeInvalidArgument { + t.Errorf("err = %v, want invalid argument", err) + } + }) + t.Run("superseded certificate", func(t *testing.T) { + if _, err := e.svc.Store.Pool.Exec(ctx, "UPDATE workers SET cert_serial = 'other' WHERE id = $1", id.WorkerID); err != nil { + t.Fatal(err) + } + err := connectWith(t, &opificav1.Hello{WorkerId: id.WorkerID, Info: info()}) + if connect.CodeOf(err) != connect.CodeUnauthenticated { + t.Errorf("err = %v, want unauthenticated", err) + } + }) +} + +// A worker whose certificate is from a different CA can't connect. +func TestForeignCertificateRejected(t *testing.T) { + e := newEnv(t) + other := newEnv(t) + foreign := other.join("trusted") + + c, err := worker.NewClient(worker.ClientConfig{ + ControllerURL: e.url, Identity: foreign, Info: info(), Logger: quiet(), MaxBackoff: 50 * time.Millisecond, + }) + if err != nil { + t.Fatal(err) + } + // Neither side trusts the other, so the TLS handshake fails on every + // attempt and Run keeps retrying until the deadline. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := c.Run(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("Run returned %v, want it to keep retrying until the deadline", err) + } + if n := len(e.svc.Sessions().List()); n != 0 { + t.Errorf("%d sessions, want 0", n) + } +} + +func TestManyWorkers(t *testing.T) { + e := newEnv(t) + var wg sync.WaitGroup + ids := make([]*worker.Identity, 10) + for i := range ids { + wg.Add(1) + go func() { + defer wg.Done() + ids[i] = e.join("untrusted") + }() + } + wg.Wait() + for _, id := range ids { + e.start(id, nil) + } + eventually(t, "10 sessions", func() bool { return len(e.svc.Sessions().List()) == 10 }) +} diff --git a/pkg/worker/client.go b/pkg/worker/client.go new file mode 100644 index 0000000..c3c09b8 --- /dev/null +++ b/pkg/worker/client.go @@ -0,0 +1,343 @@ +package worker + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "encoding/pem" + "errors" + "fmt" + "log/slog" + "math" + mrand "math/rand/v2" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "connectrpc.com/connect" + + opificav1 "github.com/opifica/opifica/gen/opifica/v1" + "github.com/opifica/opifica/gen/opifica/v1/opificav1connect" +) + +// Files in the worker's state directory. +const ( + idFile = "worker.id" + keyFile = "worker.key" + certFile = "worker.crt" + caFile = "ca.crt" +) + +// Identity is what a worker keeps after joining: its ID, private key, +// certificate and the controller's CA. +type Identity struct { + WorkerID string + cert tls.Certificate + caPool *x509.CertPool +} + +// JoinConfig is what a worker needs to join a controller. +type JoinConfig struct { + // ControllerURL is e.g. "https://controller.example:8443". + ControllerURL string + // Token and CAFingerprint are printed by `controller token create`. + Token string + CAFingerprint string + // StateDir is where the identity is stored. It must be private to the + // worker: it holds the worker's private key. + StateDir string + Info *opificav1.WorkerInfo +} + +// Join exchanges a join token for a worker identity and saves it in +// cfg.StateDir. The private key is generated here and never leaves the +// machine; only a certificate request is sent. +func Join(ctx context.Context, cfg JoinConfig) (*Identity, error) { + if _, err := os.Stat(filepath.Join(cfg.StateDir, idFile)); err == nil { + return nil, fmt.Errorf("%s already holds a worker identity; remove it to join again", cfg.StateDir) + } + u, err := url.Parse(cfg.ControllerURL) + if err != nil || u.Scheme != "https" || u.Hostname() == "" { + return nil, fmt.Errorf("controller URL must be https://host[:port], got %q", cfg.ControllerURL) + } + tlsCfg, err := PinnedCAConfig(cfg.CAFingerprint, u.Hostname()) + if err != nil { + return nil, err + } + + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return nil, err + } + csrDER, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{}, key) + if err != nil { + return nil, err + } + client := opificav1connect.NewWorkerServiceClient(httpClient(tlsCfg), cfg.ControllerURL) + resp, err := client.Join(ctx, connect.NewRequest(&opificav1.JoinRequest{ + JoinToken: cfg.Token, + CsrPem: pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrDER}), + Info: cfg.Info, + })) + if err != nil { + return nil, fmt.Errorf("joining: %w", err) + } + + keyDER, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + return nil, err + } + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}) + if err := os.MkdirAll(cfg.StateDir, 0o700); err != nil { + return nil, err + } + // Write the ID last: its presence marks a complete identity. + for _, f := range []struct { + name string + data []byte + perm os.FileMode + }{ + {keyFile, keyPEM, 0o600}, + {certFile, resp.Msg.GetCertificatePem(), 0o644}, + {caFile, resp.Msg.GetCaBundlePem(), 0o644}, + {idFile, []byte(resp.Msg.GetWorkerId() + "\n"), 0o644}, + } { + if err := os.WriteFile(filepath.Join(cfg.StateDir, f.name), f.data, f.perm); err != nil { + return nil, err + } + } + return LoadIdentity(cfg.StateDir) +} + +// LoadIdentity reads an identity saved by Join. +func LoadIdentity(dir string) (*Identity, error) { + id, err := os.ReadFile(filepath.Join(dir, idFile)) + if err != nil { + return nil, fmt.Errorf("no worker identity in %s (join first): %w", dir, err) + } + cert, err := tls.LoadX509KeyPair(filepath.Join(dir, certFile), filepath.Join(dir, keyFile)) + if err != nil { + return nil, fmt.Errorf("loading worker certificate: %w", err) + } + caPEM, err := os.ReadFile(filepath.Join(dir, caFile)) + if err != nil { + return nil, err + } + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(caPEM) { + return nil, errors.New("ca.crt holds no certificates") + } + return &Identity{WorkerID: strings.TrimSpace(string(id)), cert: cert, caPool: pool}, nil +} + +// TLSConfig returns the mTLS config for talking to the controller. +func (id *Identity) TLSConfig() *tls.Config { + return &tls.Config{ + Certificates: []tls.Certificate{id.cert}, + RootCAs: id.caPool, + MinVersion: tls.VersionTLS13, + } +} + +func httpClient(cfg *tls.Config) *http.Client { + return &http.Client{Transport: &http.Transport{ + TLSClientConfig: cfg, + ForceAttemptHTTP2: true, + IdleConnTimeout: 90 * time.Second, + }} +} + +// ClientConfig configures a Client. +type ClientConfig struct { + ControllerURL string + Identity *Identity + Info *opificav1.WorkerInfo + + // HeartbeatInterval defaults to 10s; the controller drops workers it + // hasn't heard from for 30s. + HeartbeatInterval time.Duration + // Status reports what goes into heartbeats. Optional. + Status func() (warmVMs uint32, runningJobs []string) + // OnMessage receives every message from the controller, in order, on + // one goroutine. It should return quickly. + OnMessage func(ctx context.Context, msg *opificav1.ConnectResponse) + // MaxBackoff caps the wait between reconnect attempts. Defaults to 1m. + MaxBackoff time.Duration + Logger *slog.Logger +} + +// Client keeps a worker connected to the controller. +type Client struct { + cfg ClientConfig + rpc opificav1connect.WorkerServiceClient + log *slog.Logger + + mu sync.Mutex + stream *connect.BidiStreamForClient[opificav1.ConnectRequest, opificav1.ConnectResponse] + ready chan struct{} // closed while connected +} + +// NewClient creates a client. Call Run to connect. +func NewClient(cfg ClientConfig) (*Client, error) { + if cfg.Identity == nil { + return nil, errors.New("identity is required") + } + if cfg.Info == nil { + return nil, errors.New("info is required") + } + if cfg.HeartbeatInterval <= 0 { + cfg.HeartbeatInterval = 10 * time.Second + } + if cfg.MaxBackoff <= 0 { + cfg.MaxBackoff = time.Minute + } + log := cfg.Logger + if log == nil { + log = slog.Default() + } + return &Client{ + cfg: cfg, + rpc: opificav1connect.NewWorkerServiceClient(httpClient(cfg.Identity.TLSConfig()), cfg.ControllerURL), + log: log.With("worker", cfg.Identity.WorkerID), + ready: make(chan struct{}), + }, nil +} + +// ErrNotConnected is returned by Send while the client is between connections. +var ErrNotConnected = errors.New("not connected to the controller") + +// Send sends a message on the current stream. +func (c *Client) Send(msg *opificav1.ConnectRequest) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.stream == nil { + return ErrNotConnected + } + return c.stream.Send(msg) +} + +// Connected returns a channel that is closed while the client is connected. +// After a disconnect a new channel is returned. +func (c *Client) Connected() <-chan struct{} { + c.mu.Lock() + defer c.mu.Unlock() + return c.ready +} + +// IsPermanent reports whether a Run error means retrying is pointless: the +// worker was revoked or its identity is no longer accepted. +func IsPermanent(err error) bool { + switch connect.CodeOf(err) { + case connect.CodePermissionDenied, connect.CodeUnauthenticated: + return true + } + return false +} + +// Run connects and stays connected, reconnecting with backoff, until ctx is +// cancelled or the controller rejects the worker permanently. +func (c *Client) Run(ctx context.Context) error { + attempt := 0 + for { + start := time.Now() + err := c.session(ctx) + if ctx.Err() != nil { + return ctx.Err() + } + if IsPermanent(err) { + c.log.Error("controller rejected this worker", "err", err) + return err + } + if time.Since(start) > c.cfg.MaxBackoff { + attempt = 0 // the last connection was healthy + } + wait := backoff(attempt, c.cfg.MaxBackoff) + attempt++ + c.log.Warn("disconnected from controller", "err", err, "retry_in", wait.Round(time.Millisecond)) + select { + case <-time.After(wait): + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// backoff is exponential from 500ms with ±25% jitter, capped at max. +func backoff(attempt int, max time.Duration) time.Duration { + d := time.Duration(float64(500*time.Millisecond) * math.Pow(2, float64(min(attempt, 20)))) + d = min(d, max) + return time.Duration(float64(d) * (0.75 + 0.5*mrand.Float64())) +} + +func (c *Client) status() (uint32, []string) { + if c.cfg.Status == nil { + return 0, nil + } + return c.cfg.Status() +} + +func (c *Client) session(ctx context.Context) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + stream := c.rpc.Connect(ctx) + + _, running := c.status() + if err := stream.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Hello{Hello: &opificav1.Hello{ + WorkerId: c.cfg.Identity.WorkerID, Info: c.cfg.Info, RunningJobIds: running, + }}}); err != nil { + // The real error, e.g. a rejected certificate, comes from Receive. + _, rerr := stream.Receive() + return errors.Join(err, rerr) + } + + c.mu.Lock() + c.stream = stream + close(c.ready) + c.mu.Unlock() + defer func() { + c.mu.Lock() + c.stream = nil + c.ready = make(chan struct{}) + c.mu.Unlock() + stream.CloseRequest() + stream.CloseResponse() + }() + + go func() { + t := time.NewTicker(c.cfg.HeartbeatInterval) + defer t.Stop() + for { + select { + case <-t.C: + warm, running := c.status() + err := c.Send(&opificav1.ConnectRequest{Message: &opificav1.ConnectRequest_Heartbeat{Heartbeat: &opificav1.Heartbeat{ + WarmVms: warm, RunningJobIds: running, + }}}) + if err != nil { + cancel() + return + } + case <-ctx.Done(): + return + } + } + }() + + c.log.Info("connected to controller") + for { + msg, err := stream.Receive() + if err != nil { + return err + } + if c.cfg.OnMessage != nil { + c.cfg.OnMessage(ctx, msg) + } + } +}