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 { type KernelCmd struct {
NoDownload bool `help:"do not download qemu image while kernel generation"` NoDownload bool `help:"do not download qemu image while kernel generation"`
UseHost bool `help:"also use host kernels"` UseHost bool `help:"also use host kernels"`
Force bool `help:"force reinstall kernel"` Force bool `help:"force reinstall kernel"`
NoHeaders bool `help:"do not install kernel headers"` NoHeaders bool `help:"do not install kernel headers"`
Shuffle bool `help:"randomize kernels installation order"` Shuffle bool `help:"randomize kernels installation order"`
Retries int64 `help:"amount of tries for each kernel" default:"10"` Retries int `help:"amount of tries for each kernel" default:"10"`
Update bool `help:"update container"` Threads int `help:"threads for parallel installation" default:"1"`
Update bool `help:"update container"`
List KernelListCmd `cmd:"" help:"list kernels"` List KernelListCmd `cmd:"" help:"list kernels"`
ListRemote KernelListRemoteCmd `cmd:"" help:"list remote 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 { 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) { 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.Registry,
g.Config.Docker.Commands, g.Config.Docker.Commands,
cmd.Max, kernelCmd.Retries, cmd.Max, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload, !kernelCmd.NoDownload,
kernelCmd.Force, kernelCmd.Force,
!kernelCmd.NoHeaders, !kernelCmd.NoHeaders,
@ -157,6 +159,7 @@ func (cmd *KernelGenallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
g.Config.Docker.Registry, g.Config.Docker.Registry,
g.Config.Docker.Commands, g.Config.Docker.Commands,
math.MaxUint32, kernelCmd.Retries, math.MaxUint32, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload, !kernelCmd.NoDownload,
kernelCmd.Force, kernelCmd.Force,
!kernelCmd.NoHeaders, !kernelCmd.NoHeaders,
@ -194,6 +197,7 @@ func (cmd *KernelInstallCmd) Run(kernelCmd *KernelCmd, g *Globals) (err error) {
g.Config.Docker.Registry, g.Config.Docker.Registry,
g.Config.Docker.Commands, g.Config.Docker.Commands,
math.MaxUint32, kernelCmd.Retries, math.MaxUint32, kernelCmd.Retries,
kernelCmd.Threads,
!kernelCmd.NoDownload, !kernelCmd.NoDownload,
kernelCmd.Force, kernelCmd.Force,
!kernelCmd.NoHeaders, !kernelCmd.NoHeaders,

View File

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

View File

@ -20,7 +20,7 @@ type PackCmd struct {
NoDownload bool `help:"do not download qemu image while kernel generation"` NoDownload bool `help:"do not download qemu image while kernel generation"`
ExploitRuns int64 `default:"4" help:"amount of runs of each exploit"` ExploitRuns int64 `default:"4" help:"amount of runs of each exploit"`
KernelRuns int64 `default:"1" help:"amount of runs of each kernel"` 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"` Threads int `help:"threads" default:"4"`