diff --git a/components/workspacekit/cmd/rings.go b/components/workspacekit/cmd/rings.go index 905ca59cc0d094..8d676f692206ff 100644 --- a/components/workspacekit/cmd/rings.go +++ b/components/workspacekit/cmd/rings.go @@ -66,23 +66,30 @@ var ring0Cmd = &cobra.Command{ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - client, conn, err := connectToInWorkspaceDaemonService(ctx) + client, err := connectToInWorkspaceDaemonService(ctx) if err != nil { log.WithError(err).Error("cannot connect to daemon") return } - defer conn.Close() prep, err := client.PrepareForUserNS(ctx, &daemonapi.PrepareForUserNSRequest{}) if err != nil { log.WithError(err).Fatal("cannot prepare for user namespaces") return } + client.Close() defer func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() + client, err := connectToInWorkspaceDaemonService(ctx) + if err != nil { + log.WithError(err).Error("cannot connect to daemon") + return + } + defer client.Close() + _, err = client.Teardown(ctx, &daemonapi.TeardownRequest{}) if err != nil { log.WithError(err).Error("cannot trigger teardown") @@ -178,18 +185,18 @@ var ring1Cmd = &cobra.Command{ ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - client, conn, err := connectToInWorkspaceDaemonService(ctx) - if err != nil { - log.WithError(err).Error("cannot connect to daemon") - return - } - defer conn.Close() - mapping := []*daemonapi.WriteIDMappingRequest_Mapping{ {ContainerId: 0, HostId: 33333, Size: 1}, {ContainerId: 1, HostId: 100000, Size: 65534}, } if !ring1Opts.MappingEstablished { + client, err := connectToInWorkspaceDaemonService(ctx) + if err != nil { + log.WithError(err).Error("cannot connect to daemon") + return + } + defer client.Close() + _, err = client.WriteIDMapping(ctx, &daemonapi.WriteIDMappingRequest{Pid: int64(os.Getpid()), Gid: false, Mapping: mapping}) if err != nil { log.WithError(err).Error("cannot establish UID mapping") @@ -205,6 +212,7 @@ var ring1Cmd = &cobra.Command{ log.WithError(err).Error("cannot exec /proc/self/exe") return } + return } @@ -348,10 +356,19 @@ var ring1Cmd = &cobra.Command{ return } + client, err := connectToInWorkspaceDaemonService(ctx) + if err != nil { + log.WithError(err).Error("cannot connect to daemon") + return + } + _, err = client.MountProc(ctx, &daemonapi.MountProcRequest{ Target: procLoc, Pid: int64(cmd.Process.Pid), }) + + client.Close() + if err != nil { log.WithError(err).Error("cannot mount proc") return @@ -760,8 +777,22 @@ type ringSyncMsg struct { FSShift api.FSShiftMethod `json:"fsshift"` } +type inWorkspaceServiceClient struct { + daemonapi.InWorkspaceServiceClient + + conn *grpc.ClientConn +} + +func (iwsc *inWorkspaceServiceClient) Close() error { + if iwsc.conn == nil { + return nil + } + + return iwsc.conn.Close() +} + // ConnectToInWorkspaceDaemonService attempts to connect to the InWorkspaceService offered by the ws-daemon. -func connectToInWorkspaceDaemonService(ctx context.Context) (daemonapi.InWorkspaceServiceClient, *grpc.ClientConn, error) { +func connectToInWorkspaceDaemonService(ctx context.Context) (*inWorkspaceServiceClient, error) { const socketFN = "/.workspace/daemon.sock" t := time.NewTicker(500 * time.Millisecond) @@ -775,16 +806,19 @@ func connectToInWorkspaceDaemonService(ctx context.Context) (daemonapi.InWorkspa case <-t.C: continue case <-ctx.Done(): - return nil, nil, fmt.Errorf("socket did not appear before context was canceled") + return nil, fmt.Errorf("socket did not appear before context was canceled") } } conn, err := grpc.DialContext(ctx, "unix://"+socketFN, grpc.WithInsecure()) if err != nil { - return nil, nil, err + return nil, err } - return daemonapi.NewInWorkspaceServiceClient(conn), conn, nil + return &inWorkspaceServiceClient{ + InWorkspaceServiceClient: daemonapi.NewInWorkspaceServiceClient(conn), + conn: conn, + }, nil } func init() {