✨ Add init container for check middleware (#405)

✨ Add init container for check middlewares
go-sigma · Sep 24, 2024 · 2537949 · 2537949
1 parent bba9c5a
commit 2537949
Show file tree

Hide file tree

Showing 11 changed files with 172 additions and 52 deletions.
diff --git a/build/local.Dockerfile b/build/local.Dockerfile
@@ -46,7 +46,9 @@ VOLUME /var/lib/sigma
 VOLUME /etc/sigma
 
 RUN adduser --disabled-password -h /home/sigma -s /bin/sh -u 1001 sigma && \
-  chown -R 1001:1001 /opt/trivy/
+  chown -R 1001:1001 /opt/trivy && \
+  mkdir -p /var/lib/sigma && \
+  chown -R 1001:1001 /var/lib/sigma
 
 WORKDIR /home/sigma
 

diff --git a/cmd/root.go b/cmd/root.go
@@ -16,14 +16,14 @@ package cmd
 
 import (
 	"os"
-	"path"
 	"strings"
 
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 
 	"github.com/go-sigma/sigma/pkg/configs"
+	"github.com/go-sigma/sigma/pkg/consts"
 
 	_ "github.com/go-sigma/sigma/cmd/imports"
 )
@@ -55,16 +55,13 @@ func initConfig() {
 	if cfgFile != "" {
 		viper.SetConfigFile(cfgFile)
 	} else {
-		pwd, err := os.Getwd()
-		cobra.CheckErr(err)
 		viper.AddConfigPath("/etc/sigma")
-		viper.AddConfigPath("$HOME/.sigma")
-		viper.AddConfigPath(path.Join(pwd, "conf"))
 		viper.SetConfigType("yaml")
 		viper.SetConfigName("config.yaml")
 	}
 
 	viper.AutomaticEnv()
+	viper.SetEnvPrefix(consts.AppName)
 	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
 
 	cobra.CheckErr(viper.ReadInConfig())

diff --git a/cmd/tools.go b/cmd/tools.go
@@ -43,61 +43,69 @@ import (
 	"github.com/go-sigma/sigma/pkg/utils/token"
 )
 
+func init() {
+	toolsCmd.AddCommand(
+		toolsMiddlewareCheckerCmd(),
+		toolsForPushBuilderImageCmd(),
+	)
+
+	rootCmd.AddCommand(toolsCmd)
+}
+
 // toolsCmd represents the tools command
 var toolsCmd = &cobra.Command{
 	Use:   "tools",
 	Short: "Tools for sigma",
 }
 
-var toolsForPushBuilderImageCmd = &cobra.Command{
-	Use:   "push-builder-images",
-	Short: "Push builder images to distribution",
-	PersistentPreRun: func(_ *cobra.Command, _ []string) {
-		initConfig()
-		logger.SetLevel(viper.GetString("log.level"))
-	},
-	RunE: func(_ *cobra.Command, _ []string) error {
-		err := configs.Initialize()
-		if err != nil {
-			log.Error().Err(err).Msg("initialize configs with error")
-			return err
-		}
+func toolsForPushBuilderImageCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "push-builder-images",
+		Short: "Push builder images to distribution",
+		PersistentPreRun: func(_ *cobra.Command, _ []string) {
+			initConfig()
+			logger.SetLevel(viper.GetString("log.level"))
+		},
+		RunE: func(_ *cobra.Command, _ []string) error {
+			err := configs.Initialize()
+			if err != nil {
+				log.Error().Err(err).Msg("initialize configs with error")
+				return err
+			}
 
-		config := ptr.To(configs.GetConfiguration())
+			config := ptr.To(configs.GetConfiguration())
 
-		err = badger.Initialize(context.Background(), config)
-		if err != nil {
-			log.Error().Err(err).Msg("initialize badger with error")
-			return err
-		}
+			err = badger.Initialize(context.Background(), config)
+			if err != nil {
+				log.Error().Err(err).Msg("initialize badger with error")
+				return err
+			}
 
-		err = locker.Initialize(config)
-		if err != nil {
-			log.Error().Err(err).Msg("initialize locker with error")
-			return err
-		}
+			err = locker.Initialize(config)
+			if err != nil {
+				log.Error().Err(err).Msg("initialize locker with error")
+				return err
+			}
 
-		err = dal.Initialize(config)
-		if err != nil {
-			log.Error().Err(err).Msg("initialize database with error")
-			return err
-		}
+			err = dal.Initialize(config)
+			if err != nil {
+				log.Error().Err(err).Msg("initialize database with error")
+				return err
+			}
 
-		err = initBaseimage(config)
-		if err != nil {
-			log.Error().Err(err).Msg("push builder image with error")
-			return err
-		}
+			err = initBaseimage(config)
+			if err != nil {
+				log.Error().Err(err).Msg("push builder image with error")
+				return err
+			}
 
-		return nil
-	},
-}
+			return nil
+		},
+	}
 
-func init() {
-	toolsForPushBuilderImageCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is /etc/sigma/sigma.yaml)")
+	cmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is /etc/sigma/sigma.yaml)")
 
-	toolsCmd.AddCommand(toolsForPushBuilderImageCmd)
-	rootCmd.AddCommand(toolsCmd)
+	return cmd
 }
 
 func initBaseimage(config configs.Configuration) error {
@@ -216,3 +224,48 @@ func pushImage(config configs.Configuration, path, name, version string) error {
 	}
 	return nil
 }
+
+func toolsMiddlewareCheckerCmd() *cobra.Command {
+	var waitTimeout time.Duration
+	cmd := &cobra.Command{
+		Use:   "middleware-checker",
+		Short: "Check all of middleware status all ready",
+		PersistentPreRun: func(_ *cobra.Command, _ []string) {
+			initConfig()
+			logger.SetLevel(viper.GetString("log.level"))
+		},
+		RunE: func(_ *cobra.Command, _ []string) error {
+			err := configs.Initialize()
+			if err != nil {
+				log.Error().Err(err).Msg("initialize configs with error")
+				return err
+			}
+
+			if waitTimeout == 0 {
+				waitTimeout = time.Second * 120
+			}
+
+			ctx, cancel := context.WithTimeout(context.Background(), waitTimeout)
+			defer cancel()
+
+			for {
+				select {
+				case <-ctx.Done():
+					return fmt.Errorf("middleware checker timeout, not all of middleware ready")
+				case <-time.After(time.Second * 3):
+					err = configs.CheckMiddleware()
+					if err != nil {
+						log.Error().Err(err).Msg("check middleware with error")
+					} else {
+						return nil
+					}
+				}
+			}
+		},
+	}
+
+	cmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is /etc/sigma/sigma.yaml)")
+	cmd.PersistentFlags().DurationVar(&waitTimeout, "wait-timeout", time.Second*120, "wait middleware timeout")
+
+	return cmd
+}
diff --git a/deploy/sigma/templates/distribution/deployment.yaml b/deploy/sigma/templates/distribution/deployment.yaml
@@ -28,6 +28,20 @@ spec:
       {{- end }}
       securityContext:
         {{- toYaml .Values.podSecurityContext | nindent 8 }}
+      initContainers:
+        - name: check-middlewares
+          image: {{ printf "%s/%s:%s" .Values.image.registry .Values.image.repository .Values.image.tag | quote }}
+          imagePullPolicy: {{ .Values.image.pullPolicy }}
+          command:
+            - "/bin/bash"
+            - "-c"
+            - |-
+              set -e;
+              echo "Waiting for sigma middlewares to be ready...";
+              sigma \
+                tools \
+                middleware-checker;
+              echo "sigma middlewares are all ready";
       containers:
         - name: {{ printf "%s-distribution" ( include "sigma.fullname" . ) | quote }}
           securityContext:

diff --git a/deploy/sigma/templates/server/deployment.yaml b/deploy/sigma/templates/server/deployment.yaml
@@ -29,6 +29,20 @@ spec:
       securityContext:
         {{- toYaml .Values.podSecurityContext | nindent 8 }}
       serviceAccountName: {{ include "sigma.fullname" . | quote }}
+      initContainers:
+        - name: check-middlewares
+          image: {{ printf "%s/%s:%s" .Values.image.registry .Values.image.repository .Values.image.tag | quote }}
+          imagePullPolicy: {{ .Values.image.pullPolicy }}
+          command:
+            - "/bin/bash"
+            - "-c"
+            - |-
+              set -e;
+              echo "Waiting for sigma middlewares to be ready...";
+              sigma \
+                tools \
+                middleware-checker;
+              echo "sigma middlewares are all ready";
       containers:
         - name: {{ printf "%s-server" .Chart.Name }}
           securityContext:

diff --git a/deploy/sigma/templates/worker/deployment.yaml b/deploy/sigma/templates/worker/deployment.yaml
@@ -29,6 +29,20 @@ spec:
       securityContext:
         {{- toYaml .Values.podSecurityContext | nindent 8 }}
       serviceAccountName: {{ include "sigma.fullname" . | quote }}
+      initContainers:
+        - name: check-middlewares
+          image: {{ printf "%s/%s:%s" .Values.image.registry .Values.image.repository .Values.image.tag | quote }}
+          imagePullPolicy: {{ .Values.image.pullPolicy }}
+          command:
+            - "/bin/bash"
+            - "-c"
+            - |-
+              set -e;
+              echo "Waiting for sigma middlewares to be ready...";
+              sigma \
+                tools \
+                middleware-checker;
+              echo "sigma middlewares are all ready";
       containers:
         - name: {{ include "sigma.worker" . | quote }}
           securityContext:

diff --git a/pkg/cmds/distribution/distribution.go b/pkg/cmds/distribution/distribution.go
@@ -80,7 +80,7 @@ func Serve() error {
 		pprof.Register(e, consts.PprofPath)
 	}
 
-	err := workq.Initialize(config)
+	err := workq.InitProducer(config)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/configs/configs.go b/pkg/configs/configs.go
@@ -24,12 +24,16 @@ var checkers []checker
 func Initialize() error {
 	defaultSettings()
 
+	return nil
+}
+
+// CheckMiddleware ...
+func CheckMiddleware() error {
 	for _, checker := range checkers {
 		err := checker(ptr.To(configuration))
 		if err != nil {
 			return err
 		}
 	}
-
 	return nil
 }
diff --git a/pkg/configs/configs_test.go b/pkg/configs/configs_test.go
@@ -33,7 +33,8 @@ func TestInitialize(t *testing.T) {
 	checkers = make([]checker, 0)
 	err := Initialize()
 	assert.NoError(t, err)
+
 	checkers = append(checkers, noErrChecker, errChecker)
-	err = Initialize()
+	err = CheckMiddleware()
 	assert.Error(t, err)
 }
diff --git a/pkg/daemon/scan/decorator.go b/pkg/daemon/scan/decorator.go
@@ -62,7 +62,7 @@ func decorator(runner func(context.Context, *models.Artifact, chan decoratorArti
 			for status := range statusChan {
 				switch status.Daemon {
 				case enums.DaemonVulnerability:
-					err = artifactService.UpdateVulnerability(ctx, id,
+					err = artifactService.UpdateVulnerability(context.Background(), id,
 						map[string]any{
 							query.ArtifactVulnerability.Raw.ColumnName().String():     status.Raw,
 							query.ArtifactVulnerability.Result.ColumnName().String():  status.Result,
@@ -73,7 +73,7 @@ func decorator(runner func(context.Context, *models.Artifact, chan decoratorArti
 						},
 					)
 				case enums.DaemonSbom:
-					err = artifactService.UpdateSbom(ctx,
+					err = artifactService.UpdateSbom(context.Background(),
 						id,
 						map[string]any{
 							query.ArtifactSbom.Raw.ColumnName().String():     status.Raw,

diff --git a/pkg/modules/workq/workq.go b/pkg/modules/workq/workq.go
@@ -37,6 +37,27 @@ var TopicHandlers = make(map[enums.Daemon]definition.Consumer)
 // ProducerClient ...
 var ProducerClient definition.WorkQueueProducer
 
+// InitProducer ...
+func InitProducer(config configs.Configuration) error {
+	var err error
+	switch config.WorkQueue.Type {
+	case enums.WorkQueueTypeDatabase:
+		ProducerClient, err = database.NewWorkQueueProducer(config, TopicHandlers)
+	case enums.WorkQueueTypeKafka:
+		ProducerClient, err = kafka.NewWorkQueueProducer(config, TopicHandlers)
+	case enums.WorkQueueTypeRedis:
+		ProducerClient, err = redis.NewWorkQueueProducer(config, TopicHandlers)
+	case enums.WorkQueueTypeInmemory:
+		ProducerClient, err = inmemory.NewWorkQueueProducer(config, TopicHandlers)
+	default:
+		return fmt.Errorf("Workq %s not support", config.WorkQueue.Type.String())
+	}
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // Initialize ...
 func Initialize(config configs.Configuration) error {
 	var err error