[teleport-update] Add unlink-package command

lib/autoupdate/agent/installer.go

@@ -39,6 +39,7 @@ import (
 	"github.com/google/renameio/v2"
 	"github.com/gravitational/trace"
 
+	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/lib/utils"
 )
 
@@ -443,6 +444,46 @@ func (li *LocalInstaller) LinkSystem(ctx context.Context) (revert func(context.C
 	return revert, trace.Wrap(err)
 }
 
+// TryLink links the specified version, but only in the case that
+// no installation of Teleport is already linked or partially linked.
+// See Installer interface for additional specs.
+func (li *LocalInstaller) TryLink(ctx context.Context, version string) error {
+	versionDir, err := li.versionDir(version)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	return trace.Wrap(li.tryLinks(ctx,
+		filepath.Join(versionDir, "bin"),
+		filepath.Join(versionDir, serviceDir),
+	))
+}
+
+// TryLinkSystem links the system installation, but only in the case that
+// no installation of Teleport is already linked or partially linked.
+// See Installer interface for additional specs.
+func (li *LocalInstaller) TryLinkSystem(ctx context.Context) error {
+	return trace.Wrap(li.tryLinks(ctx, li.SystemBinDir, li.SystemServiceDir))
+}
+
+// Unlink unlinks a version from LinkBinDir and LinkServiceDir.
+// See Installer interface for additional specs.
+func (li *LocalInstaller) Unlink(ctx context.Context, version string) error {
+	versionDir, err := li.versionDir(version)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	return trace.Wrap(li.removeLinks(ctx,
+		filepath.Join(versionDir, "bin"),
+		filepath.Join(versionDir, serviceDir),
+	))
+}
+
+// UnlinkSystem unlinks the system (package) version from LinkBinDir and LinkServiceDir.
+// See Installer interface for additional specs.
+func (li *LocalInstaller) UnlinkSystem(ctx context.Context) error {
+	return trace.Wrap(li.removeLinks(ctx, li.SystemBinDir, li.SystemServiceDir))
+}
+
 // symlink from oldname to newname
 type symlink struct {
 	oldname, newname string
@@ -640,25 +681,67 @@ func readFileN(name string, n int64) ([]byte, error) {
 	return data, trace.Wrap(err)
 }
 
-// TryLink links the specified version, but only in the case that
-// no installation of Teleport is already linked or partially linked.
-// See Installer interface for additional specs.
-func (li *LocalInstaller) TryLink(ctx context.Context, version string) error {
-	versionDir, err := li.versionDir(version)
+func (li *LocalInstaller) removeLinks(ctx context.Context, binDir, svcDir string) error {
+	removeService := false
+	entries, err := os.ReadDir(binDir)
+	if err != nil {
+		return trace.Errorf("failed to find Teleport binary directory: %w", err)
+	}
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+		oldname := filepath.Join(binDir, entry.Name())
+		newname := filepath.Join(li.LinkBinDir, entry.Name())
+		v, err := os.Readlink(newname)
+		if errors.Is(err, os.ErrNotExist) ||
+			errors.Is(err, os.ErrInvalid) ||
+			errors.Is(err, syscall.EINVAL) {
+			li.Log.DebugContext(ctx, "Link not present.", "oldname", oldname, "newname", newname)
+			continue
+		}
+		if err != nil {
+			return trace.Errorf("error reading link for %s: %w", filepath.Base(newname), err)
+		}
+		if v != oldname {
+			li.Log.DebugContext(ctx, "Skipping link to different binary.", "oldname", oldname, "newname", newname)
+			continue
+		}
+		if err := os.Remove(newname); err != nil {
+			li.Log.ErrorContext(ctx, "Unable to remove link.", "oldname", oldname, "newname", newname, errorKey, err)
+			continue
+		}
+		if filepath.Base(newname) == teleport.ComponentTeleport {
+			removeService = true
+		}
+	}
+	// only remove service if teleport was removed
+	if !removeService {
+		li.Log.DebugContext(ctx, "Teleport binary not unlinked. Skipping removal of teleport.service.")
+		return nil
+	}
+	src := filepath.Join(svcDir, serviceName)
+	srcBytes, err := readFileN(src, maxServiceFileSize)
 	if err != nil {
 		return trace.Wrap(err)
 	}
-	return trace.Wrap(li.tryLinks(ctx,
-		filepath.Join(versionDir, "bin"),
-		filepath.Join(versionDir, serviceDir),
-	))
-}
-
-// TryLinkSystem links the system installation, but only in the case that
-// no installation of Teleport is already linked or partially linked.
-// See Installer interface for additional specs.
-func (li *LocalInstaller) TryLinkSystem(ctx context.Context) error {
-	return trace.Wrap(li.tryLinks(ctx, li.SystemBinDir, li.SystemServiceDir))
+	dst := filepath.Join(li.LinkServiceDir, serviceName)
+	dstBytes, err := readFileN(dst, maxServiceFileSize)
+	if errors.Is(err, os.ErrNotExist) {
+		li.Log.DebugContext(ctx, "Service not present.", "path", dst)
+		return nil
+	}
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	if !bytes.Equal(srcBytes, dstBytes) {
+		li.Log.WarnContext(ctx, "Removed teleport binary link, but skipping removal of custom teleport.service.")
+		return nil
+	}
+	if err := os.Remove(dst); err != nil {
+		return trace.Errorf("error removing copy of %s: %w", filepath.Base(dst), err)
+	}
+	return nil
 }
 
 // tryLinks create binary and service links for files in binDir and svcDir if links are not already present.

lib/autoupdate/agent/installer_test.go

@@ -32,6 +32,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"slices"
 	"strconv"
 	"strings"
 	"testing"
@@ -393,7 +394,7 @@ func TestLocalInstaller_Link(t *testing.T) {
 			installer := &LocalInstaller{
 				InstallDir:     versionsDir,
 				LinkBinDir:     filepath.Join(linkDir, "bin"),
-				LinkServiceDir: filepath.Join(linkDir, "lib/systemd/system"),
+				LinkServiceDir: filepath.Join(linkDir, serviceDir),
 				Log:            slog.Default(),
 			}
 			ctx := context.Background()
@@ -635,7 +636,7 @@ func TestLocalInstaller_TryLink(t *testing.T) {
 			installer := &LocalInstaller{
 				InstallDir:     versionsDir,
 				LinkBinDir:     filepath.Join(linkDir, "bin"),
-				LinkServiceDir: filepath.Join(linkDir, "lib/systemd/system"),
+				LinkServiceDir: filepath.Join(linkDir, serviceDir),
 				Log:            slog.Default(),
 			}
 			ctx := context.Background()
@@ -773,8 +774,8 @@ func TestLocalInstaller_Remove(t *testing.T) {
 
 			installer := &LocalInstaller{
 				InstallDir:     versionsDir,
-				LinkBinDir:     linkDir,
-				LinkServiceDir: linkDir,
+				LinkBinDir:     filepath.Join(linkDir, "bin"),
+				LinkServiceDir: filepath.Join(linkDir, serviceDir),
 				Log:            slog.Default(),
 			}
 			ctx := context.Background()
@@ -796,6 +797,185 @@ func TestLocalInstaller_Remove(t *testing.T) {
 	}
 }
 
+func TestLocalInstaller_Unlink(t *testing.T) {
+	t.Parallel()
+	const version = "existing-version"
+	servicePath := filepath.Join(serviceDir, serviceName)
+
+	tests := []struct {
+		name    string
+		bins    []string
+		svcOrig []byte
+
+		links   []symlink
+		svcCopy []byte
+
+		remaining []string
+		errMatch  string
+	}{
+		{
+			name:    "normal",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+			svcCopy: []byte("orig"),
+		},
+		{
+			name:    "different services",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+			svcCopy:   []byte("custom"),
+			remaining: []string{servicePath},
+		},
+		{
+			name:    "missing target service",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+		},
+		{
+			name: "missing source service",
+			bins: []string{"teleport", "tsh"},
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+			svcCopy:   []byte("custom"),
+			remaining: []string{servicePath},
+			errMatch:  "no such",
+		},
+		{
+			name:    "missing teleport link",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+			svcCopy:   []byte("orig"),
+			remaining: []string{servicePath},
+		},
+		{
+			name:    "missing other link",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+			},
+			svcCopy: []byte("orig"),
+		},
+		{
+			name:    "wrong teleport link",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "other", newname: "bin/teleport"},
+				{oldname: "bin/tsh", newname: "bin/tsh"},
+			},
+			svcCopy:   []byte("orig"),
+			remaining: []string{servicePath, "bin/teleport"},
+		},
+		{
+			name:    "wrong other link",
+			bins:    []string{"teleport", "tsh"},
+			svcOrig: []byte("orig"),
+			links: []symlink{
+				{oldname: "bin/teleport", newname: "bin/teleport"},
+				{oldname: "wrong", newname: "bin/tsh"},
+			},
+			svcCopy:   []byte("orig"),
+			remaining: []string{"bin/tsh"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			versionsDir := t.TempDir()
+			versionDir := filepath.Join(versionsDir, version)
+			err := os.MkdirAll(versionDir, 0o755)
+			require.NoError(t, err)
+			linkDir := t.TempDir()
+
+			var files []smallFile
+			for _, n := range tt.bins {
+				files = append(files, smallFile{
+					name: filepath.Join(versionDir, "bin", n),
+					data: []byte("binary"),
+					mode: os.ModePerm,
+				})
+			}
+			if tt.svcOrig != nil {
+				files = append(files, smallFile{
+					name: filepath.Join(versionDir, servicePath),
+					data: tt.svcOrig,
+					mode: os.ModePerm,
+				})
+			}
+			if tt.svcCopy != nil {
+				files = append(files, smallFile{
+					name: filepath.Join(linkDir, servicePath),
+					data: tt.svcCopy,
+					mode: os.ModePerm,
+				})
+			}
+
+			for _, n := range files {
+				err = os.MkdirAll(filepath.Dir(n.name), os.ModePerm)
+				require.NoError(t, err)
+				err = os.WriteFile(n.name, n.data, n.mode)
+				require.NoError(t, err)
+			}
+			for _, n := range tt.links {
+				newname := filepath.Join(linkDir, n.newname)
+				oldname := filepath.Join(versionDir, n.oldname)
+				err = os.MkdirAll(filepath.Dir(newname), os.ModePerm)
+				require.NoError(t, err)
+				err = os.Symlink(oldname, newname)
+				require.NoError(t, err)
+			}
+
+			installer := &LocalInstaller{
+				InstallDir:     versionsDir,
+				LinkBinDir:     filepath.Join(linkDir, "bin"),
+				LinkServiceDir: filepath.Join(linkDir, serviceDir),
+				Log:            slog.Default(),
+			}
+			ctx := context.Background()
+			err = installer.Unlink(ctx, version)
+			if tt.errMatch != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.errMatch)
+			} else {
+				require.NoError(t, err)
+			}
+			for _, n := range tt.remaining {
+				_, err = os.Lstat(filepath.Join(linkDir, n))
+				require.NoError(t, err)
+			}
+			for _, n := range tt.links {
+				if slices.Contains(tt.remaining, n.newname) {
+					continue
+				}
+				_, err = os.Lstat(filepath.Join(linkDir, n.newname))
+				require.ErrorIs(t, err, os.ErrNotExist)
+			}
+			if !slices.Contains(tt.remaining, servicePath) {
+				_, err = os.Lstat(filepath.Join(linkDir, servicePath))
+				require.ErrorIs(t, err, os.ErrNotExist)
+			}
+		})
+	}
+}
+
 func TestLocalInstaller_List(t *testing.T) {
 	installDir := t.TempDir()
 	versions := []string{"v1", "v2"}