Feat: Build secrets

pkg/abstractions/image/build.go

@@ -65,6 +65,7 @@ type BuildOpts struct {
 	BuildSteps         []BuildStep
 	ForceRebuild       bool
 	EnvVars            []string
+	BuildSecrets       []string
 }
 
 func (o *BuildOpts) String() string {
@@ -259,6 +260,7 @@ func (b *Builder) Build(ctx context.Context, opts *BuildOpts, outputChan chan co
 			SourceImageCreds: opts.BaseImageCreds,
 			Dockerfile:       dockerfile,
 			BuildCtxObject:   &opts.BuildCtxObject,
+			BuildSecrets:     opts.BuildSecrets,
 		},
 		ContainerId:  containerId,
 		Env:          opts.EnvVars,

pkg/abstractions/image/image.go

@@ -2,7 +2,9 @@ package image
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/beam-cloud/beta9/pkg/auth"
 	"github.com/beam-cloud/beta9/pkg/common"
 	"github.com/beam-cloud/beta9/pkg/network"
 	"github.com/beam-cloud/beta9/pkg/repository"
@@ -21,13 +23,15 @@ type ImageService interface {
 
 type RuncImageService struct {
 	pb.UnimplementedImageServiceServer
-	builder *Builder
-	config  types.AppConfig
+	builder     *Builder
+	config      types.AppConfig
+	backendRepo repository.BackendRepository
 }
 
 type ImageServiceOpts struct {
 	Config        types.AppConfig
 	ContainerRepo repository.ContainerRepository
+	BackendRepo   repository.BackendRepository
 	Scheduler     *scheduler.Scheduler
 	Tailscale     *network.Tailscale
 }
@@ -47,8 +51,9 @@ func NewRuncImageService(
 	}
 
 	return &RuncImageService{
-		builder: builder,
-		config:  opts.Config,
+		builder:     builder,
+		config:      opts.Config,
+		backendRepo: opts.BackendRepo,
 	}, nil
 }
 
@@ -60,6 +65,12 @@ func (is *RuncImageService) VerifyImageBuild(ctx context.Context, in *pb.VerifyI
 		return nil, errors.Errorf("Python version not supportted: %s", in.PythonVersion)
 	}
 
+	authInfo, _ := auth.AuthInfoFromContext(ctx)
+	buildSecrets, err := is.retrieveBuildSecrets(ctx, in.Secrets, authInfo)
+	if err != nil {
+		return nil, err
+	}
+
 	opts := &BuildOpts{
 		BaseImageTag:      baseImageTag,
 		BaseImageName:     is.config.ImageService.Runner.BaseImageName,
@@ -72,6 +83,7 @@ func (is *RuncImageService) VerifyImageBuild(ctx context.Context, in *pb.VerifyI
 		EnvVars:           in.EnvVars,
 		Dockerfile:        in.Dockerfile,
 		BuildCtxObject:    in.BuildCtxObject,
+		BuildSecrets:      buildSecrets,
 	}
 
 	if in.ExistingImageUri != "" {
@@ -97,6 +109,13 @@ func (is *RuncImageService) VerifyImageBuild(ctx context.Context, in *pb.VerifyI
 func (is *RuncImageService) BuildImage(in *pb.BuildImageRequest, stream pb.ImageService_BuildImageServer) error {
 	log.Info().Interface("request", in).Msg("incoming image build request")
 
+	authInfo, _ := auth.AuthInfoFromContext(stream.Context())
+
+	buildSecrets, err := is.retrieveBuildSecrets(stream.Context(), in.Secrets, authInfo)
+	if err != nil {
+		return err
+	}
+
 	buildOptions := &BuildOpts{
 		BaseImageTag:       is.config.ImageService.Runner.Tags[in.PythonVersion],
 		BaseImageName:      is.config.ImageService.Runner.BaseImageName,
@@ -110,6 +129,7 @@ func (is *RuncImageService) BuildImage(in *pb.BuildImageRequest, stream pb.Image
 		EnvVars:            in.EnvVars,
 		Dockerfile:         in.Dockerfile,
 		BuildCtxObject:     in.BuildCtxObject,
+		BuildSecrets:       buildSecrets,
 	}
 
 	ctx := stream.Context()
@@ -149,6 +169,21 @@ func (is *RuncImageService) BuildImage(in *pb.BuildImageRequest, stream pb.Image
 	return nil
 }
 
+func (is *RuncImageService) retrieveBuildSecrets(ctx context.Context, secrets []string, authInfo *auth.AuthInfo) ([]string, error) {
+	var buildSecrets []string
+	if secrets != nil {
+		secrets, err := is.backendRepo.GetSecretsByNameDecrypted(ctx, authInfo.Workspace, secrets)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, secret := range secrets {
+			buildSecrets = append(buildSecrets, fmt.Sprintf("%s=%s", secret.Name, secret.Value))
+		}
+	}
+	return buildSecrets, nil
+}
+
 func convertBuildSteps(buildSteps []*pb.BuildStep) []BuildStep {
 	steps := make([]BuildStep, len(buildSteps))
 	for i, s := range buildSteps {

pkg/abstractions/image/image.proto

@@ -25,6 +25,7 @@ message VerifyImageBuildRequest {
   repeated string env_vars = 7;
   string dockerfile = 8;
   string build_ctx_object = 9;
+  repeated string secrets = 10;
 }
 
 message VerifyImageBuildResponse {
@@ -48,6 +49,7 @@ message BuildImageRequest {
   repeated string env_vars = 7;
   string dockerfile = 8;
   string build_ctx_object = 9;
+  repeated string secrets = 10;
 }
 
 message BuildImageResponse {

pkg/gateway/gateway.go

@@ -249,6 +249,7 @@ func (g *Gateway) registerServices() error {
 		ContainerRepo: g.ContainerRepo,
 		Scheduler:     g.Scheduler,
 		Tailscale:     g.Tailscale,
+		BackendRepo:   g.BackendRepo,
 	})
 	if err != nil {
 		return err

pkg/repository/backend_postgres.go

@@ -1416,6 +1416,18 @@ func (r *PostgresBackendRepository) GetSecretByName(ctx context.Context, workspa
 	return &secret, nil
 }
 
+func (r *PostgresBackendRepository) GetSecretsByName(ctx context.Context, workspace *types.Workspace, names []string) ([]types.Secret, error) {
+	query := `SELECT id, external_id, name, value, workspace_id, last_updated_by, created_at, updated_at FROM workspace_secret WHERE name = ANY($1) AND workspace_id = $2;`
+
+	var secrets []types.Secret
+	err := r.client.SelectContext(ctx, &secrets, query, pq.Array(names), workspace.Id)
+	if err != nil {
+		return nil, err
+	}
+
+	return secrets, nil
+}
+
 func (r *PostgresBackendRepository) GetSecretByNameDecrypted(ctx context.Context, workspace *types.Workspace, name string) (*types.Secret, error) {
 	secret, err := r.GetSecretByName(ctx, workspace, name)
 	if err != nil {
@@ -1437,6 +1449,28 @@ func (r *PostgresBackendRepository) GetSecretByNameDecrypted(ctx context.Context
 	return secret, nil
 }
 
+func (r *PostgresBackendRepository) GetSecretsByNameDecrypted(ctx context.Context, workspace *types.Workspace, names []string) ([]types.Secret, error) {
+	secrets, err := r.GetSecretsByName(ctx, workspace, names)
+	if err != nil {
+		return nil, err
+	}
+
+	signingKey, err := pkgCommon.ParseSigningKey(*workspace.SigningKey)
+	if err != nil {
+		return nil, err
+	}
+
+	for i, secret := range secrets {
+		decryptedSecret, err := pkgCommon.Decrypt(signingKey, secret.Value)
+		if err != nil {
+			return nil, err
+		}
+		secrets[i].Value = string(decryptedSecret)
+	}
+
+	return secrets, nil
+}
+
 func (r *PostgresBackendRepository) ListSecrets(ctx context.Context, workspace *types.Workspace) ([]types.Secret, error) {
 	query := `SELECT id, external_id, name, workspace_id, last_updated_by, created_at, updated_at FROM workspace_secret WHERE workspace_id = $1;`
 

pkg/repository/base.go

@@ -124,7 +124,9 @@ type BackendRepository interface {
 	UpdateConcurrencyLimit(ctx context.Context, concurrencyLimitId uint, gpuLimit uint32, cpuMillicoreLimit uint32) (*types.ConcurrencyLimit, error)
 	CreateSecret(ctx context.Context, workspace *types.Workspace, tokenId uint, name string, value string) (*types.Secret, error)
 	GetSecretByName(ctx context.Context, workspace *types.Workspace, name string) (*types.Secret, error)
+	GetSecretsByName(ctx context.Context, workspace *types.Workspace, names []string) ([]types.Secret, error)
 	GetSecretByNameDecrypted(ctx context.Context, workspace *types.Workspace, name string) (*types.Secret, error)
+	GetSecretsByNameDecrypted(ctx context.Context, workspace *types.Workspace, names []string) ([]types.Secret, error)
 	ListSecrets(ctx context.Context, workspace *types.Workspace) ([]types.Secret, error)
 	UpdateSecret(ctx context.Context, workspace *types.Workspace, tokenId uint, secretId string, value string) (*types.Secret, error)
 	DeleteSecret(ctx context.Context, workspace *types.Workspace, secretName string) error

pkg/types/scheduler.go

@@ -80,10 +80,11 @@ type ContainerState struct {
 }
 
 type BuildOptions struct {
-	SourceImage      *string `json:"source_image"`
-	Dockerfile       *string `json:"dockerfile"`
-	BuildCtxObject   *string `json:"build_context"`
-	SourceImageCreds string  `json:"source_image_creds"`
+	SourceImage      *string  `json:"source_image"`
+	Dockerfile       *string  `json:"dockerfile"`
+	BuildCtxObject   *string  `json:"build_context"`
+	SourceImageCreds string   `json:"source_image_creds"`
+	BuildSecrets     []string `json:"build_secrets"`
 }
 
 type ContainerRequest struct {

pkg/worker/runc_server.go

@@ -111,6 +111,8 @@ func (s *RunCServer) RunCExec(ctx context.Context, in *pb.RunCExecRequest) (*pb.
 	if !exists {
 		return &pb.RunCExecResponse{Ok: false}, nil
 	}
+	process.Env = append(process.Env, instance.Spec.Process.Env...)
+	process.Env = append(process.Env, instance.Request.BuildOptions.BuildSecrets...)
 
 	err = s.runcHandle.Exec(ctx, in.ContainerId, *process, &runc.ExecOpts{
 		OutputWriter: instance.OutputWriter,