riverencrypt

package main

import (
	"context"
	"encoding/base64"
	"fmt"
	"log/slog"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/riverqueue/river"
	"github.com/riverqueue/river/rivershared/util/slogutil"
	"github.com/riverqueue/river/rivertype"

	"riverqueue.com/riverpro"
	"riverqueue.com/riverpro/driver/riverpropgxv5"
	"riverqueue.com/riverpro/internal/internaltest"
	"riverqueue.com/riverpro/riverencrypt"
	"riverqueue.com/riverpro/riverencrypt/riversecretbox"
)

type EncryptArgs struct {
	// Message is a secret message encrypted in the database, but available to
	// workers.
	Message string `json:"message"`
}

func (EncryptArgs) Kind() string { return "encrypt" }

type EncryptWorker struct {
	river.WorkerDefaults[EncryptArgs]
}

func (w *EncryptWorker) Work(ctx context.Context, job *river.Job[EncryptArgs]) error {
	fmt.Printf("Secret message: %+v\n", job.Args.Message)
	return nil
}

// Example_encryptHook demonstrates the use of encryption hook that encrypts and
// decrypts job args.
func main() {
	ctx := context.Background()

	dbPool, err := pgxpool.NewWithConfig(ctx, internaltest.DatabaseConfig(""))
	if err != nil {
		panic(err)
	}
	defer dbPool.Close()

	// Required for the purpose of this test, but not necessary in real usage.
	if err := internaltest.TruncateRiverTables(ctx, dbPool); err != nil {
		panic(err)
	}

	workers := river.NewWorkers()
	river.AddWorker(workers, &EncryptWorker{})

	// Generate your own key with something like:
	//
	// 	var key [32]byte
	// 	if _, err := rand.Reader.Read(key[:]); err != nil {
	// 		panic(err)
	// 	}
	// 	encodedKey := base64.StdEncoding.EncodeToString(key[:])
	// 	fmt.Printf("encoded key: %s\n", encodedKey)
	//
	key := mustDecodeBase64EncodedKey("iRmwTuVGl2BAwTUPRTJbP/iA2EKpTrzXpEcNIXG2BI0=")

	riverClient, err := riverpro.NewClient(riverpropgxv5.New(dbPool), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHook(riversecretbox.NewEncryptor(key)),
			},
			Logger: slog.New(&slogutil.SlogMessageOnlyHandler{Level: slog.LevelWarn}),
			Queues: map[string]river.QueueConfig{
				river.QueueDefault: {MaxWorkers: 100},
			},
			TestOnly: true, // suitable only for use in tests; remove for live environments
			Workers:  workers,
		},
	})
	if err != nil {
		panic(err)
	}

	// Out of example scope, but used to wait until a job is worked.
	subscribeChan, subscribeCancel := riverClient.Subscribe(river.EventKindJobCompleted)
	defer subscribeCancel()

	if err := riverClient.Start(ctx); err != nil {
		panic(err)
	}

	_, err = riverClient.Insert(ctx, EncryptArgs{
		Message: "This message is encrypted in the database, but plaintext in workers.",
	}, nil)
	if err != nil {
		panic(err)
	}

	waitForNJobs(subscribeChan, 1)

	if err := riverClient.Stop(ctx); err != nil {
		panic(err)
	}

}

// Decodes a base64-encoded key to a 32-byte array for use with riverencrypt.
// This exactly methodology isn't necessary, but it demonstrates how a key might
// be easily stored to an env var or vault.
func mustDecodeBase64EncodedKey(encodedKey string) [32]byte {
	decodedBytes, err := base64.StdEncoding.DecodeString(encodedKey)
	if err != nil {
		panic(err)
	}

	var decodedKey [32]byte
	if copy(decodedKey[:], decodedBytes) != 32 {
		panic("expected to copy exactly 32 bytes")
	}

	return decodedKey
}

Output:

Secret message: This message is encrypted in the database, but plaintext in workers.

var (
	keyNew = mustDecodeBase64EncodedKey("T8sUAPOQNSDDMAiMyfrK8EaLOlY/cJ21PPNn1InCqIQ=")
	keyOld = mustDecodeBase64EncodedKey("fdnQ7+v/5Pb28rYqpynRSdzWfqs1gD6/J/0I9IUh65s=")
)

// Step 0: Client starts with the old key (i.e. the one used before starting
// the key rotation).
{
	_, err := riverpro.NewClient(riverpropgxv5.New(nil), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHook(riversecretbox.NewEncryptor(
					keyOld,
				)),
			},
		},
	})
	if err != nil {
		panic(err)
	}
}

// Step 1: Generate a new key and add it to the encryptor (in this case
// using NaCl NewEncryptor).
//
// The new key should be in the first key position so that all new jobs will
// be encrypted with it.
{
	_, err := riverpro.NewClient(riverpropgxv5.New(nil), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHook(riversecretbox.NewEncryptor(
					keyNew,
					keyOld,
				)),
			},
		},
	})
	if err != nil {
		panic(err)
	}
}

// Step 2: After all jobs using the original key have rotated through the
// queue, remove the old key.
//
// Remember that there may be jobs encrypted using the original key that
// errored and are queued for future retry. The original key shouldn't be
// removed until all those jobs have gotten all the way to their last retry.
{
	_, err := riverpro.NewClient(riverpropgxv5.New(nil), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHook(riversecretbox.NewEncryptor(
					keyNew,
				)),
			},
		},
	})
	if err != nil {
		panic(err)
	}
}

// Output:

key := mustDecodeBase64EncodedKey("iRmwTuVGl2BAwTUPRTJbP/iA2EKpTrzXpEcNIXG2BI0=")

// Demonstrates EncryptHook.JobKindsExclude.
{
	_, err := riverpro.NewClient(riverpropgxv5.New(nil), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHookConfig(&riverencrypt.EncryptHookConfig{
					Encryptor: riversecretbox.NewEncryptor(key),
					// Encrypt/decrypt all job args except those in this list.
					JobKindsExclude: []string{
						(EncryptArgs{}).Kind(),
					},
				}),
			},
		},
	})
	if err != nil {
		panic(err)
	}
}

// Demonstrates EncryptHook.JobKindsInclude.
{
	_, err := riverpro.NewClient(riverpropgxv5.New(nil), &riverpro.Config{
		Config: river.Config{
			Hooks: []rivertype.Hook{
				riverencrypt.NewEncryptHookConfig(&riverencrypt.EncryptHookConfig{
					Encryptor: riversecretbox.NewEncryptor(key),
					// Only encrypt/decrypt job args included in this list.
					JobKindsInclude: []string{
						(EncryptArgs{}).Kind(),
					},
				}),
			},
		},
	})
	if err != nil {
		panic(err)
	}
}

// Output: