snapshot of https://github.com/mattermost/mattermost.git (master, 2026-09-16; blobs > 131072B omitted)

server/public/pluginapi/kv_memory.go

@@ -0,0 +1,283 @@
+package pluginapi
+
+import (
+	"bytes"
+	"encoding/json"
+	"slices"
+	"strings"
+	"sync"
+	"time"
+	"unicode/utf8"
+
+	"github.com/mattermost/mattermost/server/public/model"
+
+	"github.com/pkg/errors"
+)
+
+// MemoryStore is an implementation of the plugin KV store API for testing.
+// It's not meant for production use.
+// It's safe for concurrent use by multiple goroutines.
+type MemoryStore struct {
+	mux   sync.RWMutex
+	elems map[string]kvElem
+}
+
+type kvElem struct {
+	value     []byte
+	expiresAt *time.Time
+}
+
+func (e kvElem) isExpired() bool {
+	return e.expiresAt != nil && e.expiresAt.Before(time.Now())
+}
+
+// Set stores a key-value pair, unique per plugin.
+// Keys prefixed with `mmi_` are reserved for internal use and will fail to be set.
+//
+// Returns (false, err) if DB error occurred
+// Returns (false, nil) if the value was not set
+// Returns (true, nil) if the value was set
+func (s *MemoryStore) Set(key string, value any, options ...KVSetOption) (bool, error) {
+	if key == "" {
+		return false, errors.New("key must not be empty")
+	}
+
+	if strings.HasPrefix(key, internalKeyPrefix) {
+		return false, errors.Errorf("'%s' prefix is not allowed for keys", internalKeyPrefix)
+	}
+
+	if utf8.RuneCountInString(key) > model.KeyValueKeyMaxRunes {
+		return false, errors.Errorf("key must not be longer then %d", model.KeyValueKeyMaxRunes)
+	}
+
+	opts := KVSetOptions{}
+	for _, o := range options {
+		if o != nil {
+			o(&opts)
+		}
+	}
+
+	var valueBytes []byte
+	if value != nil {
+		// Assume JSON encoding, unless explicitly given a byte slice.
+		var isValueInBytes bool
+		valueBytes, isValueInBytes = value.([]byte)
+		if !isValueInBytes {
+			var err error
+			valueBytes, err = json.Marshal(value)
+			if err != nil {
+				return false, errors.Wrapf(err, "failed to marshal value %v", value)
+			}
+		}
+	}
+
+	downstreamOpts := model.PluginKVSetOptions{
+		Atomic:          opts.Atomic,
+		ExpireInSeconds: opts.ExpireInSeconds,
+	}
+
+	if opts.oldValue != nil {
+		oldValueBytes, isOldValueInBytes := opts.oldValue.([]byte)
+		if isOldValueInBytes {
+			downstreamOpts.OldValue = oldValueBytes
+		} else {
+			data, err := json.Marshal(opts.oldValue)
+			if err != nil {
+				return false, errors.Wrapf(err, "failed to marshal value %v", opts.oldValue)
+			}
+
+			downstreamOpts.OldValue = data
+		}
+	}
+
+	if err := downstreamOpts.IsValid(); err != nil {
+		return false, err
+	}
+
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	if s.elems == nil {
+		s.elems = make(map[string]kvElem)
+	}
+
+	if !opts.Atomic {
+		if value == nil {
+			delete(s.elems, key)
+		} else {
+			s.elems[key] = kvElem{
+				value:     valueBytes,
+				expiresAt: expireTime(downstreamOpts.ExpireInSeconds),
+			}
+		}
+
+		return true, nil
+	}
+
+	oldElem := s.elems[key]
+	if !oldElem.isExpired() && !bytes.Equal(oldElem.value, downstreamOpts.OldValue) {
+		return false, nil
+	}
+
+	if value == nil {
+		delete(s.elems, key)
+	} else {
+		s.elems[key] = kvElem{
+			value:     valueBytes,
+			expiresAt: expireTime(downstreamOpts.ExpireInSeconds),
+		}
+	}
+
+	return true, nil
+}
+
+func (s *MemoryStore) SetAtomicWithRetries(key string, valueFunc func(oldValue []byte) (newValue any, err error)) error {
+	if valueFunc == nil {
+		return errors.New("function must not be nil")
+	}
+
+	for range numRetries {
+		var oldVal []byte
+		if err := s.Get(key, &oldVal); err != nil {
+			return errors.Wrapf(err, "failed to get value for key %s", key)
+		}
+
+		newVal, err := valueFunc(oldVal)
+		if err != nil {
+			return errors.Wrap(err, "valueFunc failed")
+		}
+
+		if saved, err := s.Set(key, newVal, SetAtomic(oldVal)); err != nil {
+			return errors.Wrapf(err, "DB failed to set value for key %s", key)
+		} else if saved {
+			return nil
+		}
+
+		// small delay to allow cooperative scheduling to do its thing
+		time.Sleep(10 * time.Millisecond)
+	}
+	return errors.Errorf("failed to set value after %d retries", numRetries)
+}
+
+func (s *MemoryStore) ListKeys(page int, count int, options ...ListKeysOption) ([]string, error) {
+	if page < 0 {
+		return nil, errors.New("page number must not be negative")
+	}
+
+	if count < 0 {
+		return nil, errors.New("count must not be negative")
+	}
+
+	if count == 0 {
+		return []string{}, nil
+	}
+
+	opt := listKeysOptions{}
+	for _, o := range options {
+		if o != nil {
+			o(&opt)
+		}
+	}
+
+	allKeys := make([]string, 0)
+	s.mux.RLock()
+	for k, e := range s.elems {
+		if e.isExpired() {
+			continue
+		}
+		allKeys = append(allKeys, k)
+	}
+	s.mux.RUnlock()
+
+	if len(allKeys) == 0 {
+		return []string{}, nil
+	}
+
+	slices.Sort(allKeys)
+
+	pageKeys := paginateSlice(allKeys, page, count)
+
+	if len(opt.checkers) == 0 {
+		return pageKeys, nil
+	}
+
+	n := 0
+	for _, k := range pageKeys {
+		keep := true
+		for _, c := range opt.checkers {
+			ok, err := c(k)
+			if err != nil {
+				return nil, err
+			}
+			if !ok {
+				keep = false
+				break
+			}
+		}
+
+		if keep {
+			pageKeys[n] = k
+			n++
+		}
+	}
+
+	return pageKeys[:n], nil
+}
+
+func (s *MemoryStore) Get(key string, o any) error {
+	s.mux.RLock()
+	e, ok := s.elems[key]
+	s.mux.RUnlock()
+	if !ok || len(e.value) == 0 || e.isExpired() {
+		return nil
+	}
+
+	if bytesOut, ok := o.(*[]byte); ok {
+		*bytesOut = e.value
+		return nil
+	}
+
+	if err := json.Unmarshal(e.value, o); err != nil {
+		return errors.Wrapf(err, "failed to unmarshal value for key %s", key)
+	}
+
+	return nil
+}
+
+func (s *MemoryStore) Delete(key string) error {
+	s.mux.Lock()
+	delete(s.elems, key)
+	s.mux.Unlock()
+
+	return nil
+}
+
+// DeleteAll removes all key-value pairs.
+func (s *MemoryStore) DeleteAll() error {
+	s.mux.Lock()
+	s.elems = make(map[string]kvElem)
+	s.mux.Unlock()
+
+	return nil
+}
+
+func expireTime(expireInSeconds int64) *time.Time {
+	if expireInSeconds == 0 {
+		return nil
+	}
+	t := time.Now().Add(time.Second * time.Duration(expireInSeconds))
+	return &t
+}
+
+func paginateSlice[T any](list []T, page int, perPage int) []T {
+	i := page * perPage
+	j := (page + 1) * perPage
+	l := len(list)
+	if j > l {
+		j = l
+	}
+	if i > l {
+		i = l
+	}
+	return list[i:j]
+}