Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add context #581

Merged
merged 5 commits into from
Nov 17, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 40 additions & 44 deletions distrobuilder/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ import "C"
import (
"bufio"
"bytes"
"context"
"errors"
"fmt"
"io"
Expand Down Expand Up @@ -110,6 +111,8 @@ type cmdGlobal struct {
interrupt chan os.Signal
logger *zap.SugaredLogger
overlayCleanup func()
ctx context.Context
cancel context.CancelFunc
}

func main() {
Expand All @@ -126,25 +129,34 @@ func main() {
os.Exit(1)
}

// Timeout handler
go func() {
// No timeout set
if globalCmd.flagTimeout == 0 {
return
}
var err error

globalCmd.logger, err = shared.GetLogger(globalCmd.flagDebug)
if err != nil {
fmt.Fprintf(os.Stderr, "Failed to get logger: %s\n", err)
os.Exit(1)
}

time.Sleep(time.Duration(globalCmd.flagTimeout) * time.Second)
if globalCmd.flagTimeout == 0 {
globalCmd.ctx, globalCmd.cancel = context.WithCancel(context.Background())
} else {
globalCmd.ctx, globalCmd.cancel = context.WithTimeout(context.Background(), time.Duration(globalCmd.flagTimeout)*time.Second)
}

// exit all chroots otherwise we cannot remove the cache directory
for _, exit := range shared.ActiveChroots {
if exit != nil {
exit()
go func() {
for {
select {
case <-globalCmd.interrupt:
globalCmd.cancel()
globalCmd.logger.Info("Interrupted")
return
case <-globalCmd.ctx.Done():
if globalCmd.flagTimeout > 0 {
globalCmd.logger.Info("Timed out")
}
return
}
}

globalCmd.postRun(nil, nil)
fmt.Println("Timed out")
os.Exit(1)
}()

// Create temp directory if the cache directory isn't explicitly set
Expand All @@ -157,14 +169,6 @@ func main() {

globalCmd.flagCacheDir = dir
}

var err error

globalCmd.logger, err = shared.GetLogger(globalCmd.flagDebug)
if err != nil {
fmt.Fprintf(os.Stderr, "Failed to get logger: %s\n", err)
os.Exit(1)
}
},
PersistentPostRunE: globalCmd.postRun,
CompletionOptions: cobra.CompletionOptions{DisableDefaultCmd: true},
Expand Down Expand Up @@ -204,21 +208,6 @@ func main() {
repackWindowsCmd := cmdRepackWindows{global: &globalCmd}
app.AddCommand(repackWindowsCmd.command())

go func() {
<-globalCmd.interrupt

// exit all chroots otherwise we cannot remove the cache directory
for _, exit := range shared.ActiveChroots {
if exit != nil {
exit()
}
}

globalCmd.postRun(nil, nil)
fmt.Println("Interrupted")
os.Exit(1)
}()

globalCmd.interrupt = make(chan os.Signal, 1)
signal.Notify(globalCmd.interrupt, os.Interrupt)

Expand Down Expand Up @@ -296,7 +285,7 @@ func (c *cmdGlobal) preRunBuild(cmd *cobra.Command, args []string) error {
}

// Load and run downloader
downloader, err := sources.Load(c.definition.Source.Downloader, c.logger, *c.definition, c.sourceDir, c.flagCacheDir, c.flagSourcesDir)
downloader, err := sources.Load(c.ctx, c.definition.Source.Downloader, c.logger, *c.definition, c.sourceDir, c.flagCacheDir, c.flagSourcesDir)
if err != nil {
return fmt.Errorf("Failed to load downloader %q: %w", c.definition.Source.Downloader, err)
}
Expand Down Expand Up @@ -346,7 +335,7 @@ func (c *cmdGlobal) preRunBuild(cmd *cobra.Command, args []string) error {
}
}

manager, err := managers.Load(c.definition.Packages.Manager, c.logger, *c.definition)
manager, err := managers.Load(c.ctx, c.definition.Packages.Manager, c.logger, *c.definition)
if err != nil {
return fmt.Errorf("Failed to load manager %q: %w", c.definition.Packages.Manager, err)
}
Expand All @@ -362,7 +351,7 @@ func (c *cmdGlobal) preRunBuild(cmd *cobra.Command, args []string) error {

// Run post unpack hook
for _, hook := range c.definition.GetRunnableActions("post-unpack", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-unpack: %w", err)
}
Expand All @@ -380,7 +369,7 @@ func (c *cmdGlobal) preRunBuild(cmd *cobra.Command, args []string) error {

// Run post packages hook
for _, hook := range c.definition.GetRunnableActions("post-packages", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-packages: %w", err)
}
Expand Down Expand Up @@ -426,6 +415,13 @@ func (c *cmdGlobal) postRun(cmd *cobra.Command, args []string) error {
defer c.logger.Sync()
}

// exit all chroots otherwise we cannot remove the cache directory
for _, exit := range shared.ActiveChroots {
if exit != nil {
exit()
}
}

// Clean up overlay
if c.overlayCleanup != nil {
if hasLogger {
Expand Down Expand Up @@ -467,7 +463,7 @@ func (c *cmdGlobal) getOverlayDir() (string, func(), error) {
overlayDir = filepath.Join(c.flagCacheDir, "overlay")

// Use rsync if overlay doesn't work
err = shared.RsyncLocal(c.sourceDir+"/", overlayDir)
err = shared.RsyncLocal(c.ctx, c.sourceDir+"/", overlayDir)
if err != nil {
return "", nil, fmt.Errorf("Failed to copy image content: %w", err)
}
Expand All @@ -479,7 +475,7 @@ func (c *cmdGlobal) getOverlayDir() (string, func(), error) {
overlayDir = filepath.Join(c.flagCacheDir, "overlay")

// Use rsync if overlay doesn't work
err = shared.RsyncLocal(c.sourceDir+"/", overlayDir)
err = shared.RsyncLocal(c.ctx, c.sourceDir+"/", overlayDir)
if err != nil {
return "", nil, fmt.Errorf("Failed to copy image content: %w", err)
}
Expand Down
10 changes: 5 additions & 5 deletions distrobuilder/main_lxc.go
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ func (c *cmdLXC) runPack(cmd *cobra.Command, args []string, overlayDir string) e

imageTargets := shared.ImageTargetAll | shared.ImageTargetContainer

manager, err := managers.Load(c.global.definition.Packages.Manager, c.global.logger, *c.global.definition)
manager, err := managers.Load(c.global.ctx, c.global.definition.Packages.Manager, c.global.logger, *c.global.definition)
if err != nil {
return fmt.Errorf("Failed to load manager %q: %w", c.global.definition.Packages.Manager, err)
}
Expand All @@ -123,7 +123,7 @@ func (c *cmdLXC) runPack(cmd *cobra.Command, args []string, overlayDir string) e

// Run post unpack hook
for _, hook := range c.global.definition.GetRunnableActions("post-unpack", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.global.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-unpack: %w", err)
}
Expand All @@ -141,7 +141,7 @@ func (c *cmdLXC) runPack(cmd *cobra.Command, args []string, overlayDir string) e

// Run post packages hook
for _, hook := range c.global.definition.GetRunnableActions("post-packages", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.global.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-packages: %w", err)
}
Expand All @@ -151,7 +151,7 @@ func (c *cmdLXC) runPack(cmd *cobra.Command, args []string, overlayDir string) e
}

func (c *cmdLXC) run(cmd *cobra.Command, args []string, overlayDir string) error {
img := image.NewLXCImage(overlayDir, c.global.targetDir,
img := image.NewLXCImage(c.global.ctx, overlayDir, c.global.targetDir,
c.global.flagCacheDir, *c.global.definition)

for _, file := range c.global.definition.Files {
Expand Down Expand Up @@ -185,7 +185,7 @@ func (c *cmdLXC) run(cmd *cobra.Command, args []string, overlayDir string) error

// Run post files hook
for _, action := range c.global.definition.GetRunnableActions("post-files", shared.ImageTargetUndefined|shared.ImageTargetAll|shared.ImageTargetContainer) {
err := shared.RunScript(action.Action)
err := shared.RunScript(c.global.ctx, action.Action)
if err != nil {
exitChroot()
return fmt.Errorf("Failed to run post-files: %w", err)
Expand Down
18 changes: 9 additions & 9 deletions distrobuilder/main_lxd.go
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ func (c *cmdLXD) runPack(cmd *cobra.Command, args []string, overlayDir string) e
imageTargets = shared.ImageTargetContainer
}

manager, err := managers.Load(c.global.definition.Packages.Manager, c.global.logger, *c.global.definition)
manager, err := managers.Load(c.global.ctx, c.global.definition.Packages.Manager, c.global.logger, *c.global.definition)
if err != nil {
return fmt.Errorf("Failed to load manager %q: %w", c.global.definition.Packages.Manager, err)
}
Expand All @@ -183,7 +183,7 @@ func (c *cmdLXD) runPack(cmd *cobra.Command, args []string, overlayDir string) e

// Run post unpack hook
for _, hook := range c.global.definition.GetRunnableActions("post-unpack", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.global.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-unpack: %w", err)
}
Expand All @@ -201,7 +201,7 @@ func (c *cmdLXD) runPack(cmd *cobra.Command, args []string, overlayDir string) e

// Run post packages hook
for _, hook := range c.global.definition.GetRunnableActions("post-packages", imageTargets) {
err := shared.RunScript(hook.Action)
err := shared.RunScript(c.global.ctx, hook.Action)
if err != nil {
return fmt.Errorf("Failed to run post-packages: %w", err)
}
Expand All @@ -211,7 +211,7 @@ func (c *cmdLXD) runPack(cmd *cobra.Command, args []string, overlayDir string) e
}

func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error {
img := image.NewLXDImage(overlayDir, c.global.targetDir,
img := image.NewLXDImage(c.global.ctx, overlayDir, c.global.targetDir,
c.global.flagCacheDir, *c.global.definition)

imageTargets := shared.ImageTargetUndefined | shared.ImageTargetAll
Expand Down Expand Up @@ -260,7 +260,7 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error

imgFile := filepath.Join(c.global.flagCacheDir, imgFilename)

vm, err = newVM(imgFile, vmDir, c.global.definition.Targets.LXD.VM.Filesystem, c.global.definition.Targets.LXD.VM.Size)
vm, err = newVM(c.global.ctx, imgFile, vmDir, c.global.definition.Targets.LXD.VM.Filesystem, c.global.definition.Targets.LXD.VM.Size)
if err != nil {
return fmt.Errorf("Failed to instanciate VM: %w", err)
}
Expand Down Expand Up @@ -290,7 +290,7 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error
if err != nil {
return fmt.Errorf("failed to mount root partion: %w", err)
}
defer lxd.RunCommand("umount", "-R", vmDir)
defer shared.RunCommand(vm.ctx, nil, nil, "umount", "-R", vmDir)

err = vm.createUEFIFS()
if err != nil {
Expand All @@ -304,7 +304,7 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error

// We cannot use LXD's rsync package as that uses the --delete flag which
// causes an issue due to the boot/efi directory being present.
err = shared.RsyncLocal(overlayDir+"/", vmDir)
err = shared.RsyncLocal(c.global.ctx, overlayDir+"/", vmDir)
if err != nil {
return fmt.Errorf("Failed to copy rootfs: %w", err)
}
Expand Down Expand Up @@ -350,7 +350,7 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error

// Run post files hook
for _, action := range c.global.definition.GetRunnableActions("post-files", imageTargets) {
err := shared.RunScript(action.Action)
err := shared.RunScript(c.global.ctx, action.Action)
if err != nil {
exitChroot()
return fmt.Errorf("Failed to run post-files: %w", err)
Expand All @@ -361,7 +361,7 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string, overlayDir string) error

// Unmount VM directory and loop device before creating the image.
if c.flagVM {
_, err := lxd.RunCommand("umount", "-R", vmDir)
err := shared.RunCommand(vm.ctx, nil, nil, "umount", "-R", vmDir)
if err != nil {
return fmt.Errorf("Failed to unmount %q: %w", vmDir, err)
}
Expand Down
27 changes: 14 additions & 13 deletions distrobuilder/main_repack-windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ func (c *cmdRepackWindows) preRun(cmd *cobra.Command, args []string) error {
logger.Info("Mounting Windows ISO")

// Mount ISO
_, err = lxd.RunCommand("mount", "-o", "loop", args[0], c.global.sourceDir)
err = shared.RunCommand(c.global.ctx, nil, nil, "mount", "-o", "loop", args[0], c.global.sourceDir)
if err != nil {
return fmt.Errorf("Failed to mount %q at %q: %w", args[0], c.global.sourceDir, err)
}
Expand Down Expand Up @@ -216,7 +216,7 @@ func (c *cmdRepackWindows) run(cmd *cobra.Command, args []string, overlayDir str
logger.Info("Mounting driver ISO")

// Mount driver ISO
_, err := lxd.RunCommand("mount", "-o", "loop", virtioISOPath, driverPath)
err := shared.RunCommand(c.global.ctx, nil, nil, "mount", "-o", "loop", virtioISOPath, driverPath)
if err != nil {
return fmt.Errorf("Failed to mount %q at %q: %w", virtioISOPath, driverPath, err)
}
Expand Down Expand Up @@ -268,7 +268,7 @@ func (c *cmdRepackWindows) run(cmd *cobra.Command, args []string, overlayDir str

var buf bytes.Buffer

err = lxd.RunCommandWithFds(nil, &buf, "wimlib-imagex", "info", installWim)
err = shared.RunCommand(c.global.ctx, nil, &buf, "wimlib-imagex", "info", installWim)
if err != nil {
return fmt.Errorf("Failed to retrieve wim file information: %w", err)
}
Expand Down Expand Up @@ -305,18 +305,19 @@ func (c *cmdRepackWindows) run(cmd *cobra.Command, args []string, overlayDir str
}

logger.Info("Generating new ISO")
var stdout strings.Builder

stdout, err := lxd.RunCommand("genisoimage", "--version")
err = shared.RunCommand(c.global.ctx, nil, &stdout, "genisoimage", "--version")
if err != nil {
return fmt.Errorf("Failed to determine version of genisoimage: %w", err)
}

version := strings.Split(stdout, "\n")[0]
version := strings.Split(stdout.String(), "\n")[0]

if strings.HasPrefix(version, "mkisofs") {
_, err = lxd.RunCommand("genisoimage", "-iso-level", "3", "-l", "-no-emul-boot", "-b", "efi/microsoft/boot/efisys.bin", "-o", args[1], overlayDir)
err = shared.RunCommand(c.global.ctx, nil, nil, "genisoimage", "-iso-level", "3", "-l", "-no-emul-boot", "-b", "efi/microsoft/boot/efisys.bin", "-o", args[1], overlayDir)
} else {
_, err = lxd.RunCommand("genisoimage", "--allow-limited-size", "-l", "-no-emul-boot", "-b", "efi/microsoft/boot/efisys.bin", "-o", args[1], overlayDir)
err = shared.RunCommand(c.global.ctx, nil, nil, "genisoimage", "--allow-limited-size", "-l", "-no-emul-boot", "-b", "efi/microsoft/boot/efisys.bin", "-o", args[1], overlayDir)
}
if err != nil {
return fmt.Errorf("Failed to generate ISO: %w", err)
Expand All @@ -341,13 +342,13 @@ func (c *cmdRepackWindows) modifyWim(path string, index int) error {

success := false

_, err := lxd.RunCommand("wimlib-imagex", "mountrw", wimFile, strconv.Itoa(index), wimPath, "--allow-other")
err := shared.RunCommand(c.global.ctx, nil, nil, "wimlib-imagex", "mountrw", wimFile, strconv.Itoa(index), wimPath, "--allow-other")
if err != nil {
return fmt.Errorf("Failed to mount %q: %w", filepath.Base(wimFile), err)
}
defer func() {
if !success {
lxd.RunCommand("wimlib-imagex", "unmount", wimPath)
shared.RunCommand(c.global.ctx, nil, nil, "wimlib-imagex", "unmount", wimPath)
}
}()

Expand Down Expand Up @@ -380,7 +381,7 @@ func (c *cmdRepackWindows) modifyWim(path string, index int) error {
return fmt.Errorf("Failed to inject drivers: %w", err)
}

_, err = lxd.RunCommand("wimlib-imagex", "unmount", wimPath, "--commit")
err = shared.RunCommand(c.global.ctx, nil, nil, "wimlib-imagex", "unmount", wimPath, "--commit")
if err != nil {
return fmt.Errorf("Failed to unmount WIM image: %w", err)
}
Expand Down Expand Up @@ -642,21 +643,21 @@ func (c *cmdRepackWindows) injectDrivers(dirs map[string]string) error {

logger.Debugw("Updating Windows registry", "hivefile", "DRIVERS")

err := lxd.RunCommandWithFds(strings.NewReader(driversRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\DRIVERS'", filepath.Join(dirs["config"], "DRIVERS"))
err := shared.RunCommand(c.global.ctx, strings.NewReader(driversRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\DRIVERS'", filepath.Join(dirs["config"], "DRIVERS"))
if err != nil {
return fmt.Errorf("Failed to edit Windows DRIVERS registry: %w", err)
}

logger.Debugw("Updating Windows registry", "hivefile", "SYSTEM")

err = lxd.RunCommandWithFds(strings.NewReader(systemRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\SYSTEM'", filepath.Join(dirs["config"], "SYSTEM"))
err = shared.RunCommand(c.global.ctx, strings.NewReader(systemRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\SYSTEM'", filepath.Join(dirs["config"], "SYSTEM"))
if err != nil {
return fmt.Errorf("Failed to edit Windows SYSTEM registry: %w", err)
}

logger.Debugw("Updating Windows registry", "hivefile", "SOFTWARE")

err = lxd.RunCommandWithFds(strings.NewReader(softwareRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\SOFTWARE'", filepath.Join(dirs["config"], "SOFTWARE"))
err = shared.RunCommand(c.global.ctx, strings.NewReader(softwareRegistry), nil, "hivexregedit", "--merge", "--prefix='HKEY_LOCAL_MACHINE\\SOFTWARE'", filepath.Join(dirs["config"], "SOFTWARE"))
if err != nil {
return fmt.Errorf("Failed to edit Windows SOFTWARE registry: %w", err)
}
Expand Down
Loading