From 88fba8004dd8b40c6abec5b61c6d595e8f246a5d Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Thu, 30 Jan 2025 17:33:32 +0400 Subject: [PATCH] Introduce WebSocket-based RPC v2 (#239) * Introduce WebSocket-based RPC v2 * go test: add -ldflags="-B gobuildid" * No need to change the "controller.workerNotifier.Notify()" error message * No need to modify Protocol Buffers/gRPC generated code * rpcWatch(): explain that connection shouldn't be normally be closed * Avoid "port forwarding failed: " repetition in error messages * Improve comments and avoid repetition in IP resolution errors --- .cirrus.yml | 2 +- api/openapi.yaml | 5 + internal/command/controller/run.go | 7 + internal/command/dev/dev.go | 11 +- internal/controller/api.go | 11 ++ internal/controller/api_controller.go | 13 +- internal/controller/api_rpc_portforward.go | 48 ++++++ internal/controller/api_rpc_resolve_ip.go | 31 ++++ internal/controller/api_rpc_watch.go | 84 +++++++++++ internal/controller/api_vms_ip.go | 10 +- internal/controller/api_vms_portforward.go | 10 +- internal/controller/controller.go | 9 +- internal/controller/option.go | 6 + internal/controller/rendezvous/rendezvous.go | 5 + internal/controller/rpc.go | 12 +- internal/controller/sshserver/sshserver.go | 20 ++- internal/controller/wserror.go | 26 ++++ internal/tests/devcontroller/devcontroller.go | 3 + internal/worker/rpcv2.go | 141 ++++++++++++++++++ internal/worker/worker.go | 35 ++++- pkg/client/client.go | 6 + pkg/client/rpc.go | 74 +++++++++ pkg/resource/v1/v1.go | 24 ++- pkg/resource/v1/watch_instruction.go | 22 +++ 24 files changed, 588 insertions(+), 27 deletions(-) create mode 100644 internal/controller/api_rpc_portforward.go create mode 100644 internal/controller/api_rpc_resolve_ip.go create mode 100644 internal/controller/api_rpc_watch.go create mode 100644 internal/controller/wserror.go create mode 100644 internal/worker/rpcv2.go create mode 100644 pkg/client/rpc.go create mode 100644 pkg/resource/v1/watch_instruction.go diff --git a/.cirrus.yml b/.cirrus.yml index 43e30a2..665284f 100644 --- a/.cirrus.yml +++ b/.cirrus.yml @@ -8,7 +8,7 @@ task: pull_script: - tart pull ghcr.io/cirruslabs/macos-sonoma-base:latest test_script: - - go test -v -count=1 ./... + - go test -ldflags="-B gobuildid" -v -count=1 ./... always: cleanup_script: - tart list diff --git a/api/openapi.yaml b/api/openapi.yaml index f3821dd..669d2f3 100644 --- a/api/openapi.yaml +++ b/api/openapi.yaml @@ -488,6 +488,11 @@ components: commit: type: string description: Commit hash + capabilities: + type: array + items: + type: string + description: Supported capabilities ClusterSettings: title: Cluster settings type: object diff --git a/internal/command/controller/run.go b/internal/command/controller/run.go index afbf076..5120385 100644 --- a/internal/command/controller/run.go +++ b/internal/command/controller/run.go @@ -19,6 +19,7 @@ var address string var addressSSH string var debug bool var sshNoClientAuth bool +var experimentalRPCV2 bool func newRunCommand() *cobra.Command { cmd := &cobra.Command{ @@ -51,6 +52,8 @@ func newRunCommand() *cobra.Command { cmd.PersistentFlags().BoolVar(&sshNoClientAuth, "insecure-ssh-no-client-auth", false, "allow SSH clients to connect to the controller's SSH server without authentication, "+ "thus only authenticating on the target worker/VM's SSH server") + cmd.PersistentFlags().BoolVar(&experimentalRPCV2, "experimental-rpc-v2", false, + "enable experimental RPC v2 (https://github.com/cirruslabs/orchard/issues/235)") return cmd } @@ -112,6 +115,10 @@ func runController(cmd *cobra.Command, args []string) (err error) { controllerOpts = append(controllerOpts, controller.WithSSHServer(addressSSH, signer, sshNoClientAuth)) } + if experimentalRPCV2 { + controllerOpts = append(controllerOpts, controller.WithExperimentalRPCV2()) + } + controllerInstance, err := controller.New(controllerOpts...) if err != nil { return err diff --git a/internal/command/dev/dev.go b/internal/command/dev/dev.go index fa2fb0e..5fdcb9f 100644 --- a/internal/command/dev/dev.go +++ b/internal/command/dev/dev.go @@ -20,6 +20,7 @@ var ErrFailed = errors.New("failed to run development controller and worker") var devDataDirPath string var stringToStringResources map[string]string +var experimentalRPCV2 bool func NewCommand() *cobra.Command { command := &cobra.Command{ @@ -32,6 +33,8 @@ func NewCommand() *cobra.Command { "path to persist data between runs") command.PersistentFlags().StringToStringVar(&stringToStringResources, "resources", map[string]string{}, "resources that the development worker will provide") + command.PersistentFlags().BoolVar(&experimentalRPCV2, "experimental-rpc-v2", false, + "enable experimental RPC v2 (https://github.com/cirruslabs/orchard/issues/235)") return command } @@ -51,9 +54,15 @@ func runDev(cmd *cobra.Command, args []string) error { return fmt.Errorf("%w: %v", ErrFailed, err) } + var additionalControllerOpts []controller.Option + + if experimentalRPCV2 { + additionalControllerOpts = append(additionalControllerOpts, controller.WithExperimentalRPCV2()) + } + devController, devWorker, err := CreateDevControllerAndWorker(devDataDirPath, fmt.Sprintf(":%d", netconstants.DefaultControllerPort), resources, - nil, nil) + additionalControllerOpts, nil) if err != nil { return err } diff --git a/internal/controller/api.go b/internal/controller/api.go index 9cea5ce..774487c 100644 --- a/internal/controller/api.go +++ b/internal/controller/api.go @@ -111,6 +111,17 @@ func (controller *Controller) initAPI() *gin.Engine { controller.deleteWorker(c).Respond(c) }) + // RPC v2 + v1.GET("/rpc/watch", func(c *gin.Context) { + controller.rpcWatch(c).Respond(c) + }) + v1.GET("/rpc/port-forward", func(c *gin.Context) { + controller.rpcPortForward(c).Respond(c) + }) + v1.POST("/rpc/resolve-ip", func(c *gin.Context) { + controller.rpcResolveIP(c).Respond(c) + }) + // VMs v1.POST("/vms", func(c *gin.Context) { controller.createVM(c).Respond(c) diff --git a/internal/controller/api_controller.go b/internal/controller/api_controller.go index 3ab065a..1fdc2cc 100644 --- a/internal/controller/api_controller.go +++ b/internal/controller/api_controller.go @@ -15,8 +15,17 @@ func (controller *Controller) controllerInfo(ctx *gin.Context) responder.Respond return responder } + capabilities := []v1pkg.ControllerCapability{ + v1pkg.ControllerCapabilityRPCV1, + } + + if controller.experimentalRPCV2 { + capabilities = append(capabilities, v1pkg.ControllerCapabilityRPCV2) + } + return responder.JSON(http.StatusOK, &v1pkg.ControllerInfo{ - Version: version.Version, - Commit: version.Commit, + Version: version.Version, + Commit: version.Commit, + Capabilities: capabilities, }) } diff --git a/internal/controller/api_rpc_portforward.go b/internal/controller/api_rpc_portforward.go new file mode 100644 index 0000000..54e21f3 --- /dev/null +++ b/internal/controller/api_rpc_portforward.go @@ -0,0 +1,48 @@ +package controller + +import ( + "github.com/cirruslabs/orchard/internal/controller/rendezvous" + "github.com/cirruslabs/orchard/internal/responder" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/gin-gonic/gin" + "net" + "nhooyr.io/websocket" +) + +func (controller *Controller) rpcPortForward(ctx *gin.Context) responder.Responder { + if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil { + return responder + } + + // Retrieve and parse path and query parameters + session := ctx.Query("session") + errorMessage := ctx.Query("errorMessage") + + // Perform WebSocket protocol upgrade + wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{ + OriginPatterns: []string{"*"}, + }) + if err != nil { + return responder.Error(err) + } + + // Respond with the established connection + proxyCtx, err := controller.connRendezvous.Respond(session, rendezvous.ResultWithErrorMessage[net.Conn]{ + Result: websocket.NetConn(ctx, wsConn, websocket.MessageBinary), + ErrorMessage: errorMessage, + }) + if err != nil { + return controller.wsError(wsConn, websocket.StatusInternalError, "port forwarding RPC", + "failure to respond with the established WebSocket connection", err) + } + + select { + case <-proxyCtx.Done(): + // Do not close the WebSocket connection as it should be already closed by our rendezvous party + return responder.Empty() + case <-ctx.Done(): + // Connection should normally be closed by our rendezvous party, not by the worker + return controller.wsError(wsConn, websocket.StatusAbnormalClosure, "port forwarding RPC", + "unexpectedly disconnected worker", err) + } +} diff --git a/internal/controller/api_rpc_resolve_ip.go b/internal/controller/api_rpc_resolve_ip.go new file mode 100644 index 0000000..0b152b5 --- /dev/null +++ b/internal/controller/api_rpc_resolve_ip.go @@ -0,0 +1,31 @@ +package controller + +import ( + "github.com/cirruslabs/orchard/internal/controller/rendezvous" + "github.com/cirruslabs/orchard/internal/responder" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/gin-gonic/gin" + "net/http" +) + +func (controller *Controller) rpcResolveIP(ctx *gin.Context) responder.Responder { + if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeWrite); responder != nil { + return responder + } + + // Retrieve and parse path and query parameters + session := ctx.Query("session") + ip := ctx.Query("ip") + errorMessage := ctx.Query("errorMessage") + + // Respond with the resolved IP address + _, err := controller.ipRendezvous.Respond(session, rendezvous.ResultWithErrorMessage[string]{ + Result: ip, + ErrorMessage: errorMessage, + }) + if err != nil { + return responder.Error(err) + } + + return responder.Code(http.StatusOK) +} diff --git a/internal/controller/api_rpc_watch.go b/internal/controller/api_rpc_watch.go new file mode 100644 index 0000000..47ef667 --- /dev/null +++ b/internal/controller/api_rpc_watch.go @@ -0,0 +1,84 @@ +package controller + +import ( + "encoding/json" + "errors" + "github.com/cirruslabs/orchard/internal/responder" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/cirruslabs/orchard/rpc" + "github.com/gin-gonic/gin" + "nhooyr.io/websocket" + "time" +) + +func (controller *Controller) rpcWatch(ctx *gin.Context) responder.Responder { + if responder := controller.authorize(ctx, v1.ServiceAccountRoleComputeRead); responder != nil { + return responder + } + + // Retrieve and parse path and query parameters + workerName := ctx.Query("workerName") + + if workerName == "" { + return responder.Error(errors.New("worker name cannot be empty")) + } + + // Register with the worker notifier to forward requests from other + // parts of the Orchard Controller destined to this specific worker + workerCh, cancel := controller.workerNotifier.Register(ctx, workerName) + defer cancel() + + // Perform WebSocket protocol upgrade + wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{ + OriginPatterns: []string{"*"}, + }) + if err != nil { + return responder.Error(err) + } + + for { + select { + case msg := <-workerCh: + var watchInstruction v1.WatchInstruction + + // Perform Protocol Buffers to v1 API data structure conversion + switch typedAction := msg.Action.(type) { + case *rpc.WatchInstruction_PortForwardAction: + watchInstruction.PortForwardAction = &v1.PortForwardAction{ + Session: typedAction.PortForwardAction.Session, + VMUID: typedAction.PortForwardAction.VmUid, + Port: uint16(typedAction.PortForwardAction.Port), + } + case *rpc.WatchInstruction_SyncVmsAction: + watchInstruction.SyncVMsAction = &v1.SyncVMsAction{} + case *rpc.WatchInstruction_ResolveIpAction: + watchInstruction.ResolveIPAction = &v1.ResolveIPAction{ + Session: typedAction.ResolveIpAction.Session, + VMUID: typedAction.ResolveIpAction.VmUid, + } + default: + continue + } + + watchInstructionJSONBytes, err := json.Marshal(&watchInstruction) + if err != nil { + return controller.wsError(wsConn, websocket.StatusInternalError, "watch RPC", + "failure to marshal the watch instruction", err) + } + + if err := wsConn.Write(ctx, websocket.MessageBinary, watchInstructionJSONBytes); err != nil { + return controller.wsError(wsConn, websocket.StatusInternalError, "watch RPC", + "failure to write the watch instruction", err) + } + case <-time.After(30 * time.Second): + if err := wsConn.Ping(ctx); err != nil { + controller.logger.Warnf("watch RPC: failed to ping the worker %s, "+ + "connection might time out: %v", workerName, err) + } + case <-ctx.Done(): + // Connection shouldn't be normally closed by the worker + return controller.wsError(wsConn, websocket.StatusAbnormalClosure, "watch RPC", + "unexpectedly disconnected worker", err) + } + } +} diff --git a/internal/controller/api_vms_ip.go b/internal/controller/api_vms_ip.go index 1f612ec..ac1c6a4 100644 --- a/internal/controller/api_vms_ip.go +++ b/internal/controller/api_vms_ip.go @@ -2,6 +2,7 @@ package controller import ( "context" + "fmt" "github.com/cirruslabs/orchard/internal/responder" v1 "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/cirruslabs/orchard/rpc" @@ -55,11 +56,16 @@ func (controller *Controller) ip(ctx *gin.Context) responder.Responder { } select { - case ip := <-boomerangConnCh: + case rendezvousResponse := <-boomerangConnCh: + if rendezvousResponse.ErrorMessage != "" { + return responder.Error(fmt.Errorf("VM's IP resolution on the worker %s failed: %s", + vm.Worker, rendezvousResponse.ErrorMessage)) + } + result := struct { IP string `json:"ip"` }{ - IP: ip, + IP: rendezvousResponse.Result, } return responder.JSON(http.StatusOK, &result) diff --git a/internal/controller/api_vms_portforward.go b/internal/controller/api_vms_portforward.go index 6a0b3ee..47b8d03 100644 --- a/internal/controller/api_vms_portforward.go +++ b/internal/controller/api_vms_portforward.go @@ -2,6 +2,7 @@ package controller import ( "context" + "fmt" storepkg "github.com/cirruslabs/orchard/internal/controller/store" "github.com/cirruslabs/orchard/internal/netconncancel" "github.com/cirruslabs/orchard/internal/proxy" @@ -88,7 +89,12 @@ func (controller *Controller) portForward( // worker will asynchronously start port-forwarding so we wait select { - case fromWorkerConnection := <-boomerangConnCh: + case rendezvousResponse := <-boomerangConnCh: + if rendezvousResponse.ErrorMessage != "" { + return responder.Error(fmt.Errorf("failed to establish port forwarding session on the worker: %s", + rendezvousResponse.ErrorMessage)) + } + wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{ OriginPatterns: []string{"*"}, }) @@ -105,7 +111,7 @@ func (controller *Controller) portForward( } wsConnAsNetConn := websocket.NetConn(ctx, wsConn, expectedMsgType) - fromWorkerConnectionWithCancel := netconncancel.New(fromWorkerConnection, rendezvousCtxCancel) + fromWorkerConnectionWithCancel := netconncancel.New(rendezvousResponse.Result, rendezvousCtxCancel) if err := proxy.Connections(wsConnAsNetConn, fromWorkerConnectionWithCancel); err != nil { var websocketCloseError websocket.CloseError diff --git a/internal/controller/controller.go b/internal/controller/controller.go index c62c509..f8e6dd0 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -54,11 +54,12 @@ type Controller struct { logger *zap.SugaredLogger grpcServer *grpc.Server workerNotifier *notifier.Notifier - connRendezvous *rendezvous.Rendezvous[net.Conn] - ipRendezvous *rendezvous.Rendezvous[string] + connRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]] + ipRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[string]] enableSwaggerDocs bool workerOfflineTimeout time.Duration maxWorkersPerLicense uint + experimentalRPCV2 bool sshListenAddr string sshSigner ssh.Signer @@ -70,8 +71,8 @@ type Controller struct { func New(opts ...Option) (*Controller, error) { controller := &Controller{ - connRendezvous: rendezvous.New[net.Conn](), - ipRendezvous: rendezvous.New[string](), + connRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[net.Conn]](), + ipRendezvous: rendezvous.New[rendezvous.ResultWithErrorMessage[string]](), workerOfflineTimeout: 3 * time.Minute, maxWorkersPerLicense: maxWorkersPerDefaultLicense, } diff --git a/internal/controller/option.go b/internal/controller/option.go index bfa65a7..e49bb0c 100644 --- a/internal/controller/option.go +++ b/internal/controller/option.go @@ -53,6 +53,12 @@ func WithWorkerOfflineTimeout(workerOfflineTimeout time.Duration) Option { } } +func WithExperimentalRPCV2() Option { + return func(controller *Controller) { + controller.experimentalRPCV2 = true + } +} + func WithLogger(logger *zap.Logger) Option { return func(controller *Controller) { controller.logger = logger.Sugar() diff --git a/internal/controller/rendezvous/rendezvous.go b/internal/controller/rendezvous/rendezvous.go index dd175d8..2cbcc18 100644 --- a/internal/controller/rendezvous/rendezvous.go +++ b/internal/controller/rendezvous/rendezvous.go @@ -12,6 +12,11 @@ type Rendezvous[T any] struct { sessions *concurrentmap.ConcurrentMap[*TokenSlot[T]] } +type ResultWithErrorMessage[T any] struct { + Result T + ErrorMessage string +} + type TokenSlot[T any] struct { ctx context.Context ch chan T diff --git a/internal/controller/rpc.go b/internal/controller/rpc.go index 880b848..d3edc33 100644 --- a/internal/controller/rpc.go +++ b/internal/controller/rpc.go @@ -2,10 +2,12 @@ package controller import ( "context" + "github.com/cirruslabs/orchard/internal/controller/rendezvous" v1pkg "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/cirruslabs/orchard/rpc" "google.golang.org/grpc/metadata" "google.golang.org/protobuf/types/known/emptypb" + "net" //nolint:staticcheck // https://github.com/mitchellh/go-grpc-net-conn/pull/1 "github.com/golang/protobuf/proto" @@ -63,7 +65,11 @@ func (controller *Controller) PortForward(stream rpc.Controller_PortForwardServe } // make connection rendezvous aware of the connection - proxyCtx, err := controller.connRendezvous.Respond(sessionMetadataValue[0], conn) + proxyCtx, err := controller.connRendezvous.Respond(sessionMetadataValue[0], + rendezvous.ResultWithErrorMessage[net.Conn]{ + Result: conn, + }, + ) if err != nil { return err } @@ -87,7 +93,9 @@ func (controller *Controller) ResolveIP(ctx context.Context, request *rpc.Resolv } // Respond with the resolved IP address - _, err := controller.ipRendezvous.Respond(sessionMetadataValue[0], request.Ip) + _, err := controller.ipRendezvous.Respond(sessionMetadataValue[0], rendezvous.ResultWithErrorMessage[string]{ + Result: request.Ip, + }) if err != nil { return nil, err } diff --git a/internal/controller/sshserver/sshserver.go b/internal/controller/sshserver/sshserver.go index 2728cf4..33bd3b9 100644 --- a/internal/controller/sshserver/sshserver.go +++ b/internal/controller/sshserver/sshserver.go @@ -31,7 +31,7 @@ type SSHServer struct { listener net.Listener serverConfig *ssh.ServerConfig store storepkg.Store - connRendezvous *rendezvous.Rendezvous[net.Conn] + connRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]] workerNotifier *notifier.Notifier logger *zap.SugaredLogger } @@ -40,7 +40,7 @@ func NewSSHServer( address string, signer ssh.Signer, store storepkg.Store, - connRendezvous *rendezvous.Rendezvous[net.Conn], + connRendezvous *rendezvous.Rendezvous[rendezvous.ResultWithErrorMessage[net.Conn]], workerNotifier *notifier.Notifier, noClientAuth bool, logger *zap.SugaredLogger, @@ -253,7 +253,19 @@ func (server *SSHServer) handleDirectTCPIP(ctx context.Context, newChannel ssh.N // Wait for the connection from worker and commence port forwarding select { - case fromWorkerConnection := <-boomerangConnCh: + case rendezvousResponse := <-boomerangConnCh: + if rendezvousResponse.ErrorMessage != "" { + message := fmt.Sprintf("failed to establish port forwarding session on the worker: %s", + rendezvousResponse.ErrorMessage) + + if err := newChannel.Reject(ssh.ConnectionFailed, message); err != nil { + server.logger.Warnf("failed to reject new channel due to "+ + "failure establishing port forwarding session on the worker: %v", err) + } + + return + } + // Now that we have the connection from worker we can accept the channel acceptedChannel, acceptedChannelRequests, err := newChannel.Accept() if err != nil { @@ -277,7 +289,7 @@ func (server *SSHServer) handleDirectTCPIP(ctx context.Context, newChannel ssh.N }() // Commence port forwarding - if err := proxy.Connections(acceptedChannel, fromWorkerConnection); err != nil { + if err := proxy.Connections(acceptedChannel, rendezvousResponse.Result); err != nil { server.logger.Warnf("failed to port-forward: %v", err) } case <-ctx.Done(): diff --git a/internal/controller/wserror.go b/internal/controller/wserror.go new file mode 100644 index 0000000..e8d3779 --- /dev/null +++ b/internal/controller/wserror.go @@ -0,0 +1,26 @@ +package controller + +import ( + "fmt" + "github.com/cirruslabs/orchard/internal/responder" + "nhooyr.io/websocket" +) + +func (controller *Controller) wsError( + wsConn *websocket.Conn, + code websocket.StatusCode, + component string, + reason string, + err error, +) responder.Responder { + message := fmt.Sprintf("%s: %v", reason, err) + + controller.logger.Warn(message) + + if err := wsConn.Close(code, message); err != nil { + controller.logger.Warnf("%s: failed to close the WebSocket connection that entered error state"+ + " due to %s: %v", component, reason, err) + } + + return responder.Empty() +} diff --git a/internal/tests/devcontroller/devcontroller.go b/internal/tests/devcontroller/devcontroller.go index 22f0f2b..eb00031 100644 --- a/internal/tests/devcontroller/devcontroller.go +++ b/internal/tests/devcontroller/devcontroller.go @@ -29,6 +29,9 @@ func StartIntegrationTestEnvironmentWithAdditionalOpts( ) (*client.Client, *controller.Controller, *worker.Worker) { t.Setenv("ORCHARD_HOME", t.TempDir()) + // Enable experimental RPC v2 by default in tests + additionalControllerOpts = append(additionalControllerOpts, controller.WithExperimentalRPCV2()) + devController, devWorker, err := dev.CreateDevControllerAndWorker(t.TempDir(), ":0", nil, additionalControllerOpts, additionalWorkerOpts) require.NoError(t, err) diff --git a/internal/worker/rpcv2.go b/internal/worker/rpcv2.go new file mode 100644 index 0000000..0bfb397 --- /dev/null +++ b/internal/worker/rpcv2.go @@ -0,0 +1,141 @@ +package worker + +import ( + "context" + "fmt" + "github.com/cirruslabs/orchard/internal/proxy" + "github.com/cirruslabs/orchard/internal/worker/vmmanager" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/samber/lo" + "net" +) + +func (worker *Worker) watchRPCV2(ctx context.Context) error { + watchInstructionCh, watchErrCh, err := worker.client.RPC().Watch(ctx, worker.name) + if err != nil { + return err + } + + for { + select { + case watchInstruction := <-watchInstructionCh: + if portForwardAction := watchInstruction.PortForwardAction; portForwardAction != nil { + worker.handlePortForwardV2(ctx, portForwardAction) + } else if syncVMsAction := watchInstruction.SyncVMsAction; syncVMsAction != nil { + worker.requestVMSyncing() + } else if resolveIPAction := watchInstruction.ResolveIPAction; resolveIPAction != nil { + worker.handleGetIPV2(ctx, resolveIPAction) + } + case watchErr := <-watchErrCh: + return watchErr + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (worker *Worker) handlePortForwardV2(ctx context.Context, portForward *v1.PortForwardAction) { + var errorMessage string + + worker.logger.Debugf("received port-forwarding request to VM UID %s, port %d", + portForward.VMUID, portForward.Port) + + // Establish a connection with the VM + vmConn, err := worker.handlePortForwardV2Inner(ctx, portForward) + if err != nil { + errorMessage = fmt.Sprintf("port-forwarding failed: %v", err) + + worker.logger.Warn(errorMessage) + } + + // Respond + netConn, err := worker.client.RPC().RespondPortForward(ctx, portForward.Session, errorMessage) + if err != nil { + worker.logger.Warnf("port forwarding failed: failed to call API: %v", err) + + return + } + + // Proxy bytes if the connection was established without errors + if errorMessage == "" { + _ = proxy.Connections(vmConn, netConn) + } +} + +func (worker *Worker) handlePortForwardV2Inner( + ctx context.Context, + portForward *v1.PortForwardAction, +) (net.Conn, error) { + var host string + var err error + + if portForward.VMUID == "" { + // Port-forwarding request to a worker + host = "localhost" + } else { + // Port-forwarding request to a VM, find that VM + vm, ok := lo.Find(worker.vmm.List(), func(item *vmmanager.VM) bool { + return item.Resource.UID == portForward.VMUID + }) + if !ok { + return nil, fmt.Errorf("failed to get the VM: %v", err) + } + + // Obtain VM's IP address + host, err = vm.IP(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get VM's IP: %v", err) + } + } + + // Connect to the VM's port + vmConn, err := net.Dial("tcp", fmt.Sprintf("%s:%d", host, portForward.Port)) + if err != nil { + return nil, fmt.Errorf("failed to connect to the VM: %v", err) + } + + return vmConn, nil +} + +func (worker *Worker) handleGetIPV2(ctx context.Context, resolveIP *v1.ResolveIPAction) { + var errorMessage string + + worker.logger.Debugf("received IP resolution request to VM UID %s", resolveIP.VMUID) + + // Retrieve the VM's IP + ip, err := worker.handleGetIPV2Inner(ctx, resolveIP) + if err != nil { + errorMessage = fmt.Sprintf("failed to resolve VM's IP: %v", err) + + worker.logger.Warn(errorMessage) + } + + // Report results + if err := worker.client.RPC().RespondIP(ctx, resolveIP.Session, ip, errorMessage); err != nil { + worker.logger.Warnf("failed to resolve IP for the VM with UID %q: "+ + "failed to call back to the controller: %v", resolveIP.VMUID, err) + + return + } +} + +func (worker *Worker) handleGetIPV2Inner( + ctx context.Context, + resolveIP *v1.ResolveIPAction, +) (string, error) { + // Find the desired VM + vm, ok := lo.Find(worker.vmm.List(), func(item *vmmanager.VM) bool { + return item.Resource.UID == resolveIP.VMUID + }) + if !ok { + return "", fmt.Errorf("VM %q not found", resolveIP.VMUID) + } + + // Obtain VM's IP address + ip, err := vm.IP(ctx) + if err != nil { + return "", fmt.Errorf("\"tart ip\" failed for VM %q: %v", resolveIP.VMUID, err) + } + + return ip, nil +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 5814180..2e0b15a 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -122,13 +122,34 @@ func (worker *Worker) runNewSession(ctx context.Context) error { return nil } - go func() { - _ = retry.Do(func() error { - return worker.watchRPC(subCtx) - }, retry.OnRetry(func(n uint, err error) { - worker.logger.Warnf("failed to watch RPC: %v", err) - }), retry.Context(subCtx), retry.Attempts(0)) - }() + info, err := worker.client.Controller().Info(ctx) + if err != nil { + worker.logger.Warnf("failed to retrieve controller info: %v", err) + + return nil + } + + if info.Capabilities.Has(v1.ControllerCapabilityRPCV2) { + worker.logger.Infof("using WebSocket-based v2 RPC") + + go func() { + _ = retry.Do(func() error { + return worker.watchRPCV2(subCtx) + }, retry.OnRetry(func(n uint, err error) { + worker.logger.Warnf("failed to watch RPC v2: %v", err) + }), retry.Context(subCtx), retry.Attempts(0)) + }() + } else { + worker.logger.Infof("using gRPC-based v1 RPC") + + go func() { + _ = retry.Do(func() error { + return worker.watchRPC(subCtx) + }, retry.OnRetry(func(n uint, err error) { + worker.logger.Warnf("failed to watch RPC v1: %v", err) + }), retry.Context(subCtx), retry.Attempts(0)) + }() + } // Sync on-disk VMs if err := worker.syncOnDiskVMs(ctx); err != nil { diff --git a/pkg/client/client.go b/pkg/client/client.go index 51a1d01..a4c32b6 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -338,3 +338,9 @@ func (client *Client) ClusterSettings() *ClusterSettingsService { client: client, } } + +func (client *Client) RPC() *RPCService { + return &RPCService{ + client: client, + } +} diff --git a/pkg/client/rpc.go b/pkg/client/rpc.go new file mode 100644 index 0000000..fa587d3 --- /dev/null +++ b/pkg/client/rpc.go @@ -0,0 +1,74 @@ +package client + +import ( + "context" + "encoding/json" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "net" + "net/http" +) + +type RPCService struct { + client *Client +} + +func (service *RPCService) Watch( + ctx context.Context, + workerName string, +) (<-chan *v1.WatchInstruction, <-chan error, error) { + wsConn, err := service.client.wsRequest(ctx, "rpc/watch", map[string]string{ + "workerName": workerName, + }) + if err != nil { + return nil, nil, err + } + + watchEventCh := make(chan *v1.WatchInstruction) + errCh := make(chan error, 1) + + go func() { + defer func() { + _ = wsConn.Close() + }() + + decoder := json.NewDecoder(wsConn) + + for { + var watchInstruction v1.WatchInstruction + + if err := decoder.Decode(&watchInstruction); err != nil { + errCh <- err + + return + } + + watchEventCh <- &watchInstruction + } + }() + + return watchEventCh, errCh, nil +} + +func (service *RPCService) RespondPortForward( + ctx context.Context, + session string, + errorMessage string, +) (net.Conn, error) { + return service.client.wsRequest(ctx, "rpc/port-forward", map[string]string{ + "session": session, + "errorMessage": errorMessage, + }) +} + +func (service *RPCService) RespondIP( + ctx context.Context, + session string, + ip string, + errorMessage string, +) error { + return service.client.request(ctx, http.MethodPost, "rpc/resolve-ip", nil, nil, map[string]string{ + "session": session, + "ip": ip, + "errorMessage": errorMessage, + }) +} diff --git a/pkg/resource/v1/v1.go b/pkg/resource/v1/v1.go index d8918f6..5425411 100644 --- a/pkg/resource/v1/v1.go +++ b/pkg/resource/v1/v1.go @@ -103,7 +103,27 @@ const ( VMStatusFailed VMStatus = "failed" ) +type ControllerCapability string + +const ( + ControllerCapabilityRPCV1 ControllerCapability = "rpc-v1" + ControllerCapabilityRPCV2 ControllerCapability = "rpc-v2" +) + +type ControllerCapabilities []ControllerCapability + +func (controllerCapabilities ControllerCapabilities) Has(capability ControllerCapability) bool { + for _, controllerCapability := range controllerCapabilities { + if controllerCapability == capability { + return true + } + } + + return false +} + type ControllerInfo struct { - Version string `json:"version,omitempty"` - Commit string `json:"commit,omitempty"` + Version string `json:"version,omitempty"` + Commit string `json:"commit,omitempty"` + Capabilities ControllerCapabilities `json:"capabilities,omitempty"` } diff --git a/pkg/resource/v1/watch_instruction.go b/pkg/resource/v1/watch_instruction.go new file mode 100644 index 0000000..f7b08fd --- /dev/null +++ b/pkg/resource/v1/watch_instruction.go @@ -0,0 +1,22 @@ +package v1 + +type WatchInstruction struct { + PortForwardAction *PortForwardAction `json:"portForwardAction,omitempty"` + SyncVMsAction *SyncVMsAction `json:"syncVMsAction,omitempty"` + ResolveIPAction *ResolveIPAction `json:"resolveIPAction,omitempty"` +} + +type PortForwardAction struct { + Session string `json:"session"` + VMUID string `json:"vmUID"` + Port uint16 `json:"port"` +} + +type SyncVMsAction struct { + // nothing for now +} + +type ResolveIPAction struct { + Session string `json:"session"` + VMUID string `json:"vmUID"` +}