js.go

package casbinjetstream

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	tlsconfig "github.com/grepplabs/cert-source/config"
	tlsclientconfig "github.com/grepplabs/cert-source/tls/client/config"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
	"golang.org/x/sync/errgroup"
)

type TLSConfig struct {
	// Enable client-side TLS.
	Enable bool
	// Interval for refreshing client TLS certificates.
	Refresh time.Duration
	// Skip TLS verification on client side.
	InsecureSkipVerify bool
	// Client TLS key/cert/CA files.
	File TLSClientFiles
	// Optional password to decrypt RSA private key.
	KeyPassword string
	// Use system pool for root CAs.
	UseSystemPool bool
}

type TLSClientFiles struct {
	// Optional path to client TLS key file.
	Key string
	// Optional path to client TLS certificate file.
	Cert string
	// Optional path to client root CAs for server verification.
	RootCAs string
}

type Auth struct {
	// Basic auth
	User     string
	Password string
	// Token auth
	Token string
	// JWT credentials file generated by nsc
	CredentialsFile string // e.g. "/path/to/user.creds"
}

func newJSClient(config *Config) (jetstream.JetStream, error) {
	logger := config.Logger
	opts := []nats.Option{
		nats.MaxReconnects(-1),
		nats.DisconnectErrHandler(func(nc *nats.Conn, err error) {
			if err != nil {
				logger.Warn("disconnected from NATS", slog.String("error", err.Error()))
			} else {
				logger.Info("disconnected from NATS")
			}
		}),
		nats.ReconnectHandler(func(nc *nats.Conn) {
			logger.Info(fmt.Sprintf("reconnected to NATS at %s", nc.ConnectedUrl()))
		}),
	}
	if config.TLSConfig.Enable {
		tlsConfigFunc, err := tlsclientconfig.GetTLSClientConfigFunc(logger, &tlsconfig.TLSClientConfig{
			Enable:             config.TLSConfig.Enable,
			Refresh:            config.TLSConfig.Refresh,
			InsecureSkipVerify: config.TLSConfig.InsecureSkipVerify,
			File: tlsconfig.TLSClientFiles{
				Key:     config.TLSConfig.File.Key,
				Cert:    config.TLSConfig.File.Cert,
				RootCAs: config.TLSConfig.File.RootCAs,
			},
			KeyPassword:   config.TLSConfig.KeyPassword,
			UseSystemPool: config.TLSConfig.UseSystemPool,
		})
		if err != nil {
			return nil, fmt.Errorf("failed to setup NATS TLS client config: %w", err)
		}
		base := tlsConfigFunc()
		opts = append(opts, nats.Secure(&tls.Config{
			MinVersion:         tls.VersionTLS12,
			ServerName:         base.ServerName,
			InsecureSkipVerify: base.InsecureSkipVerify,
		}))
		certCB := func() (tls.Certificate, error) {
			cfg := tlsConfigFunc()
			if cfg.GetClientCertificate != nil {
				c, err := cfg.GetClientCertificate(&tls.CertificateRequestInfo{})
				if err != nil {
					return tls.Certificate{}, err
				}
				return *c, nil
			}
			if len(cfg.Certificates) > 0 {
				return cfg.Certificates[0], nil
			}
			return tls.Certificate{}, nil
		}
		rootCAsCB := func() (*x509.CertPool, error) {
			cfg := tlsConfigFunc()
			return cfg.RootCAs, nil
		}
		opts = append(opts, nats.ClientTLSConfig(certCB, rootCAsCB))
	}
	// creds file
	if config.Auth.CredentialsFile != "" {
		opts = append(opts, nats.UserCredentials(config.Auth.CredentialsFile))
	}
	// user/pass
	if config.Auth.User != "" || config.Auth.Password != "" {
		opts = append(opts, nats.UserInfo(config.Auth.User, config.Auth.Password))
	}
	// token
	if config.Auth.Token != "" {
		opts = append(opts, nats.Token(config.Auth.Token))
	}
	nc, err := nats.Connect(config.URL, opts...)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
	}
	js, err := jetstream.New(nc)
	if err != nil {
		return nil, fmt.Errorf("failed to setup JetStream: %w", err)
	}
	return js, nil
}

type store struct {
	js          jetstream.JetStream
	kv          atomic.Pointer[jetstream.KeyValue]
	logger      *slog.Logger
	concurrency int
	bucket      string
}

func cleanup(s *store) {
	if !s.js.Conn().IsClosed() {
		s.logger.Info("closing jetstream connection")
		s.js.Conn().Close()
	}
}

func newStore(ctx context.Context, config *Config) (*store, error) {
	js, err := newJSClient(config)
	if err != nil {
		return nil, fmt.Errorf("failed to create jetstream client: %w", err)
	}
	kv, err := js.KeyValue(ctx, config.Bucket)
	if err != nil {
		if !errors.Is(err, jetstream.ErrBucketNotFound) {
			return nil, fmt.Errorf("failed to get jetstream bucket '%s': %w", config.Bucket, err)
		}
		kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
			Bucket: config.Bucket,
		})
		if err != nil {
			return nil, fmt.Errorf("failed to create jetstream bucket '%s': %w", config.Bucket, err)
		}
	}
	s := &store{
		js:          js,
		logger:      config.Logger,
		concurrency: config.Concurrency,
		bucket:      config.Bucket,
	}
	s.setKV(kv)
	return s, nil
}

func (s *store) getKV() jetstream.KeyValue {
	p := s.kv.Load()
	if p == nil {
		return nil
	}
	return *p
}

func (s *store) setKV(kv jetstream.KeyValue) {
	s.kv.Store(&kv)
}

func keyFor(r CasbinRule) string {
	const delimiter = "\x1f" // unlikely delimiter \u001F
	parts := []string{r.PType, r.V0, r.V1, r.V2, r.V3, r.V4, r.V5}
	base := strings.Join(parts, delimiter)
	sum := sha256.Sum256([]byte(base))
	return "rule_" + hex.EncodeToString(sum[:])
}

func (s *store) RecreateBucket(ctx context.Context) error {
	s.logger.Info("recreating jetstream bucket", slog.String("bucket", s.bucket))

	if err := s.js.DeleteKeyValue(ctx, s.bucket); err != nil {
		if !errors.Is(err, jetstream.ErrBucketNotFound) {
			return fmt.Errorf("failed to delete jetstream bucket %q: %w", s.bucket, err)
		}
	}
	kv, err := s.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
		Bucket: s.bucket,
	})
	if err != nil {
		return fmt.Errorf("failed to recreate jetstream bucket %q: %w", s.bucket, err)
	}
	s.setKV(kv)
	s.logger.Info("jetstream bucket recreated", slog.String("bucket", s.bucket))
	return nil
}

func (s *store) CreatePolicy(ctx context.Context, r CasbinRule) error {
	data, err := json.Marshal(r)
	if err != nil {
		return err
	}
	_, err = s.getKV().Create(ctx, keyFor(r), data)
	if err != nil {
		if errors.Is(err, jetstream.ErrKeyExists) {
			return nil
		}
		return err
	}
	return nil
}

func (s *store) PurgePolicy(ctx context.Context, r CasbinRule) error {
	key := keyFor(r)
	s.logger.Debug("one: purging key", slog.String("key", key))
	err := s.getKV().Purge(ctx, key)
	return err
}

func (s *store) PurgeAllPolicies(ctx context.Context) error {
	lister, err := s.getKV().ListKeys(ctx)
	if err != nil {
		return err
	}
	defer lister.Stop()

	grp, gctx := errgroup.WithContext(ctx)
	grp.SetLimit(slices.Max([]int{s.concurrency, 1}))

	for k := range lister.Keys() {
		key := k // capture
		grp.Go(func() error {
			s.logger.Debug("all: purging key", slog.String("key", key))
			return s.getKV().Purge(gctx, key)
		})
	}
	err = grp.Wait()
	if err != nil {
		return err
	}
	// compact the bucket to remove deleted entries
	return s.getKV().PurgeDeletes(ctx, jetstream.DeleteMarkersOlderThan(-1))
}

func (s *store) PurgeFilteredPolicies(ctx context.Context, pattern CasbinRule) error {
	lister, err := s.ListEntries(ctx)
	if err != nil {
		return err
	}
	defer lister.Stop()

	grp, gctx := errgroup.WithContext(ctx)
	grp.SetLimit(slices.Max([]int{s.concurrency, 1}))

	for e := range lister.Entries() {
		var line CasbinRule
		if err := json.Unmarshal(e.Value(), &line); err != nil {
			return fmt.Errorf("failed to unmarshal entry key=%s: %w", e.Key(), err)
		}
		if !filteredMatch(pattern, line) {
			continue
		}
		key := e.Key() // capture
		grp.Go(func() error {
			s.logger.Debug("filtered: purging key", slog.String("key", key))
			return s.getKV().Purge(gctx, key)
		})
	}
	err = grp.Wait()
	if err != nil {
		return err
	}
	// compact the bucket to remove deleted entries
	return s.getKV().PurgeDeletes(ctx, jetstream.DeleteMarkersOlderThan(-1))
}

func filteredMatch(pattern, line CasbinRule) bool {
	if pattern.PType != "" && line.PType != pattern.PType {
		return false
	}
	if pattern.V0 != "" && line.V0 != pattern.V0 {
		return false
	}
	if pattern.V1 != "" && line.V1 != pattern.V1 {
		return false
	}
	if pattern.V2 != "" && line.V2 != pattern.V2 {
		return false
	}
	if pattern.V3 != "" && line.V3 != pattern.V3 {
		return false
	}
	if pattern.V4 != "" && line.V4 != pattern.V4 {
		return false
	}
	if pattern.V5 != "" && line.V5 != pattern.V5 {
		return false
	}
	return true
}

// GetAllPolicies retrieves all policies, de-duplicating by key and sorting by creation time and revision.
func (s *store) GetAllPolicies(ctx context.Context) ([]CasbinRule, error) {
	lister, err := s.ListEntries(ctx)
	if err != nil {
		return nil, err
	}
	defer lister.Stop()

	type record struct {
		line     CasbinRule
		created  time.Time
		revision uint64
	}
	recs := make(map[string]record)
	for e := range lister.Entries() {
		var line CasbinRule
		if err := json.Unmarshal(e.Value(), &line); err != nil {
			return nil, fmt.Errorf("failed to unmarshal entry key=%s: %w", e.Key(), err)
		}
		recs[e.Key()] = record{
			line:     line,
			created:  e.Created(),
			revision: e.Revision(),
		}
	}
	if err := lister.Err(); err != nil {
		return nil, fmt.Errorf("list entries failed: %w", err)
	}
	entries := make([]record, 0, len(recs))
	for _, r := range recs {
		entries = append(entries, r)
	}
	sort.Slice(entries, func(i, j int) bool {
		if entries[i].created.Equal(entries[j].created) {
			return entries[i].revision < entries[j].revision
		}
		return entries[i].created.Before(entries[j].created)
	})
	lines := make([]CasbinRule, 0, len(entries))
	for _, r := range entries {
		lines = append(lines, r.line)
	}
	return lines, nil
}

type entryLister struct {
	watcher jetstream.KeyWatcher
	entries chan jetstream.KeyValueEntry
	err     error
}

func (el *entryLister) Entries() <-chan jetstream.KeyValueEntry { return el.entries }
func (el *entryLister) Stop() error                             { return el.watcher.Stop() }
func (el *entryLister) Err() error                              { return el.err }

func (s *store) ListEntries(ctx context.Context, opts ...jetstream.WatchOpt) (*entryLister, error) {
	opts = append(opts, jetstream.IgnoreDeletes())
	watcher, err := s.getKV().WatchAll(ctx, opts...)
	if err != nil {
		return nil, err
	}
	el := &entryLister{watcher: watcher, entries: make(chan jetstream.KeyValueEntry, 256)}

	go func() {
		defer close(el.entries)
		defer watcher.Stop()
		for {
			select {
			case entry, ok := <-watcher.Updates():
				if !ok || entry == nil {
					return
				}
				el.entries <- entry
			case <-ctx.Done():
				el.err = ctx.Err()
				return
			}
		}
	}()
	return el, nil
}