From de17fb04b4511eec6b616137296ab0eee8a70b9b Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Sat, 19 Aug 2023 16:31:02 +0000
Subject: [PATCH] Add helper functions for marshaling and sealing

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 database/sql/instances.go | 35 ++++++++++++++++++++++++++++-------
 database/sql/util.go      |  8 ++------
 params/requests.go        | 13 +++++++------
 3 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/database/sql/instances.go b/database/sql/instances.go
index d5d2eb87..0c8e4a48 100644
--- a/database/sql/instances.go
+++ b/database/sql/instances.go
@@ -17,7 +17,6 @@ package sql
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 
 	runnerErrors "github.com/cloudbase/garm-provider-common/errors"
 	"github.com/cloudbase/garm-provider-common/util"
@@ -30,6 +29,25 @@ import (
 	"gorm.io/gorm/clause"
 )
 
+func (s *sqlDatabase) marshalAndSeal(data interface{}) ([]byte, error) {
+	enc, err := json.Marshal(data)
+	if err != nil {
+		return nil, errors.Wrap(err, "marshalling data")
+	}
+	return util.Seal(enc, []byte(s.cfg.Passphrase))
+}
+
+func (s *sqlDatabase) unsealAndUnmarshal(data []byte, target interface{}) error {
+	decrypted, err := util.Unseal(data, []byte(s.cfg.Passphrase))
+	if err != nil {
+		return errors.Wrap(err, "decrypting data")
+	}
+	if err := json.Unmarshal(decrypted, target); err != nil {
+		return errors.Wrap(err, "unmarshalling data")
+	}
+	return nil
+}
+
 func (s *sqlDatabase) CreateInstance(ctx context.Context, poolID string, param params.CreateInstanceParams) (params.Instance, error) {
 	pool, err := s.getPoolByID(ctx, poolID)
 	if err != nil {
@@ -46,15 +64,10 @@ func (s *sqlDatabase) CreateInstance(ctx context.Context, poolID string, param p
 
 	var secret []byte
 	if len(param.JitConfiguration) > 0 {
-		jitConfig, err := json.Marshal(param.JitConfiguration)
+		secret, err = s.marshalAndSeal(param.JitConfiguration)
 		if err != nil {
 			return params.Instance{}, errors.Wrap(err, "marshalling jit config")
 		}
-
-		secret, err = util.Seal(jitConfig, []byte(s.cfg.Passphrase))
-		if err != nil {
-			return params.Instance{}, fmt.Errorf("failed to encrypt jitconfig: %w", err)
-		}
 	}
 
 	newInstance := Instance{
@@ -251,6 +264,14 @@ func (s *sqlDatabase) UpdateInstance(ctx context.Context, instanceID string, par
 		instance.TokenFetched = *param.TokenFetched
 	}
 
+	if param.JitConfiguration != nil {
+		secret, err := s.marshalAndSeal(param.JitConfiguration)
+		if err != nil {
+			return params.Instance{}, errors.Wrap(err, "marshalling jit config")
+		}
+		instance.JitConfiguration = secret
+	}
+
 	instance.ProviderFault = param.ProviderFault
 
 	q := s.conn.Save(&instance)
diff --git a/database/sql/util.go b/database/sql/util.go
index 3930d9c4..ad0a4d8b 100644
--- a/database/sql/util.go
+++ b/database/sql/util.go
@@ -43,12 +43,8 @@ func (s *sqlDatabase) sqlToParamsInstance(instance Instance) (params.Instance, e
 
 	var jitConfig map[string]string
 	if len(instance.JitConfiguration) > 0 {
-		decrypted, err := util.Unseal(instance.JitConfiguration, []byte(s.cfg.Passphrase))
-		if err != nil {
-			return params.Instance{}, errors.Wrap(err, "decrypting jit config")
-		}
-		if err := json.Unmarshal(decrypted, &jitConfig); err != nil {
-			return params.Instance{}, errors.Wrap(err, "unmarshalling jit config")
+		if err := s.unsealAndUnmarshal(instance.JitConfiguration, &jitConfig); err != nil {
+			return params.Instance{}, errors.Wrap(err, "unmarshalling jit configuration")
 		}
 	}
 	ret := params.Instance{
diff --git a/params/requests.go b/params/requests.go
index 12a2acc3..0293223e 100644
--- a/params/requests.go
+++ b/params/requests.go
@@ -199,12 +199,13 @@ type UpdateInstanceParams struct {
 	// for this instance.
 	Addresses []commonParams.Address `json:"addresses,omitempty"`
 	// Status is the status of the instance inside the provider (eg: running, stopped, etc)
-	Status        commonParams.InstanceStatus `json:"status,omitempty"`
-	RunnerStatus  RunnerStatus                `json:"runner_status,omitempty"`
-	ProviderFault []byte                      `json:"provider_fault,omitempty"`
-	AgentID       int64                       `json:"-"`
-	CreateAttempt int                         `json:"-"`
-	TokenFetched  *bool                       `json:"-"`
+	Status           commonParams.InstanceStatus `json:"status,omitempty"`
+	RunnerStatus     RunnerStatus                `json:"runner_status,omitempty"`
+	ProviderFault    []byte                      `json:"provider_fault,omitempty"`
+	AgentID          int64                       `json:"-"`
+	CreateAttempt    int                         `json:"-"`
+	TokenFetched     *bool                       `json:"-"`
+	JitConfiguration map[string]string           `json:"-"`
 }
 
 type UpdateUserParams struct {