1
0

feat: parallel kernel installation

This commit is contained in:
dump_stack() 2023-05-21 21:43:18 +00:00
parent b02da8adeb
commit 632e4f5ffc
Signed by: dump_stack
GPG Key ID: BE44DA8C062D87DC
3 changed files with 43 additions and 30 deletions

View File

@ -18,13 +18,14 @@ import (
)
type KernelCmd struct {
NoDownload bool `help:"do not download qemu image while kernel generation"`
UseHost bool `help:"also use host kernels"`
Force bool `help:"force reinstall kernel"`
NoHeaders bool `help:"do not install kernel headers"`
Shuffle bool `help:"randomize kernels installation order"`
Retries int64 `help:"amount of tries for each kernel" default:"10"`
Update bool `help:"update container"`
NoDownload bool `help:"do not download qemu image while kernel generation"`
UseHost bool `help:"also use host kernels"`
Force bool `help:"force reinstall kernel"`
NoHeaders bool `help:"do not install kernel headers"`
Shuffle bool `help:"randomize kernels installation order"`
Retries int `help:"amount of tries for each kernel" default:"10"`
Threads int `help:"threads for parallel installation" default:"1"`
Update bool `help:"update container"`
List KernelListCmd `cmd:"" help:"list kernels"`
ListRemote KernelListRemoteCmd `cmd:"" help:"list remote kernels"`
@ -95,7 +96,7 @@ func (cmd *KernelListRemoteCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error
}
type KernelAutogenCmd struct {
Max int64 `help:"download kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"100500"`
Max int `help:"download kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"100500"`
}
func (cmd KernelAutogenCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
@ -117,6 +118,7 @@ func (cmd KernelAutogenCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
g.Config.Docker.Registry,
g.Config.Docker.Commands,
cmd.Max, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload,
kernelCmd.Force,
!kernelCmd.NoHeaders,
@ -157,6 +159,7 @@ func (cmd *KernelGenallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
g.Config.Docker.Registry,
g.Config.Docker.Commands,
math.MaxUint32, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload,
kernelCmd.Force,
!kernelCmd.NoHeaders,
@ -194,6 +197,7 @@ func (cmd *KernelInstallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
g.Config.Docker.Registry,
g.Config.Docker.Commands,
math.MaxUint32, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload,
kernelCmd.Force,
!kernelCmd.NoHeaders,

View File

@ -19,6 +19,7 @@ import (
"time"
"github.com/naoina/toml"
"github.com/remeh/sizedwaitgroup"
"github.com/rs/zerolog/log"
"code.dumpstack.io/tools/out-of-tree/cache"
@ -505,7 +506,7 @@ func SetSigintHandler(variable *bool) {
// FIXME too many parameters
func GenerateKernels(km config.Target, registry string,
commands []config.DockerCommand, max, retries int64,
commands []config.DockerCommand, max, retries, threads int,
download, force, headers, shuffle, update bool,
shutdown *bool) (err error) {
@ -530,6 +531,9 @@ func GenerateKernels(km config.Target, registry string,
if shuffle {
pkgs = shuffleStrings(pkgs)
}
swg := sizedwaitgroup.New(threads)
for i, pkg := range pkgs {
if max <= 0 {
log.Print("Max is reached")
@ -542,30 +546,35 @@ func GenerateKernels(km config.Target, registry string,
}
log.Info().Msgf("%d/%d %s", i+1, len(pkgs), pkg)
var attempt int64
for {
attempt++
swg.Add()
go func() {
defer swg.Done()
var attempt int
for {
attempt++
if *shutdown {
err = nil
return
}
if *shutdown {
err = nil
return
}
err = installKernel(km, pkg, force, headers)
if err == nil {
max--
break
} else if attempt >= retries {
log.Error().Err(err).Msg("install kernel")
log.Debug().Msg("skip")
break
} else {
log.Warn().Err(err).Msg("install kernel")
time.Sleep(time.Second)
log.Info().Msg("retry")
err = installKernel(km, pkg, force, headers)
if err == nil {
max--
break
} else if attempt >= retries {
log.Error().Err(err).Msg("install kernel")
log.Debug().Msg("skip")
break
} else {
log.Warn().Err(err).Msg("install kernel")
time.Sleep(time.Second)
log.Info().Msg("retry")
}
}
}
}()
}
swg.Wait()
return
}

View File

@ -20,7 +20,7 @@ type PackCmd struct {
NoDownload bool `help:"do not download qemu image while kernel generation"`
ExploitRuns int64 `default:"4" help:"amount of runs of each exploit"`
KernelRuns int64 `default:"1" help:"amount of runs of each kernel"`
Max int64 `help:"download random kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"1"`
Max int `help:"download random kernels from set defined by regex in release_mask, but no more than X for each of release_mask" default:"1"`
Threads int `help:"threads" default:"4"`