open-vault/vault/activity_log_testing_util.go

package vault

import (
	"context"
	"testing"

	"github.com/hashicorp/vault/helper/constants"
	"github.com/hashicorp/vault/sdk/logical"
	"github.com/hashicorp/vault/vault/activity"
)

// InjectActivityLogDataThisMonth populates the in-memory client store
// with some entities and tokens, overriding what was already there
// It is currently used for API integration tests
func (c *Core) InjectActivityLogDataThisMonth(t *testing.T) (map[string]uint64, map[string]uint64) {
	t.Helper()
	tokens := make(map[string]uint64, 0)
	entitiesByNS := make(map[string]uint64, 0)
	tokens["root"] = 5
	entitiesByNS["root"] = 5

	if constants.IsEnterprise {
		tokens["ns0"] = 5
		tokens["ns1"] = 1
		tokens["ns2"] = 1
		entitiesByNS["ns0"] = 1
		entitiesByNS["ns1"] = 1
		entitiesByNS["ns2"] = 1
	}

	c.activityLog.l.Lock()
	defer c.activityLog.l.Unlock()
	c.activityLog.fragmentLock.Lock()
	defer c.activityLog.fragmentLock.Unlock()

	c.activityLog.currentSegment.tokenCount.CountByNamespaceID = tokens
	c.activityLog.entityTracker.entityCountByNamespaceID = entitiesByNS
	return entitiesByNS, tokens
}

// Return the in-memory activeEntities from an activity log
func (c *Core) GetActiveEntities() map[string]struct{} {
	out := make(map[string]struct{})

	c.stateLock.RLock()
	c.activityLog.fragmentLock.RLock()
	for k, v := range c.activityLog.entityTracker.activeEntities {
		out[k] = v
	}
	c.activityLog.fragmentLock.RUnlock()
	c.stateLock.RUnlock()

	return out
}

// GetCurrentEntities returns the current entity activity log
func (a *ActivityLog) GetCurrentEntities() *activity.EntityActivityLog {
	a.l.RLock()
	defer a.l.RUnlock()
	return a.currentSegment.currentEntities
}

// WriteToStorage is used to put entity data in storage
// `path` should be the complete path (not relative to the view)
func WriteToStorage(t *testing.T, c *Core, path string, data []byte) {
	t.Helper()
	err := c.barrier.Put(context.Background(), &logical.StorageEntry{
		Key:   path,
		Value: data,
	})
	if err != nil {
		t.Fatalf("Failed to write %s\nto %s\nerror: %v", data, path, err)
	}
}

// SetStandbyEnable sets enabled on a performance standby (using config)
func (a *ActivityLog) SetStandbyEnable(ctx context.Context, enabled bool) {
	var enableStr string
	if enabled {
		enableStr = "enable"
	} else {
		enableStr = "disable"
	}

	// TODO only patch enabled?
	a.SetConfigStandby(ctx, activityConfig{
		DefaultReportMonths: 12,
		RetentionMonths:     24,
		Enabled:             enableStr,
	})
}

// ExpectCurrentSegmentRefreshed verifies that the current segment has been refreshed
// non-nil empty components and updated with the `expectedStart` timestamp
// Note: if `verifyTimeNotZero` is true, ignore `expectedStart` and just make sure the timestamp isn't 0
func (a *ActivityLog) ExpectCurrentSegmentRefreshed(t *testing.T, expectedStart int64, verifyTimeNotZero bool) {
	t.Helper()

	a.l.RLock()
	defer a.l.RUnlock()
	a.fragmentLock.RLock()
	defer a.fragmentLock.RUnlock()
	if a.currentSegment.currentEntities == nil {
		t.Fatalf("expected non-nil currentSegment.currentEntities")
	}
	if a.currentSegment.currentEntities.Entities == nil {
		t.Errorf("expected non-nil currentSegment.currentEntities.Entities")
	}
	if a.entityTracker.activeEntities == nil {
		t.Errorf("expected non-nil activeEntities")
	}
	if a.currentSegment.tokenCount == nil {
		t.Fatalf("expected non-nil currentSegment.tokenCount")
	}
	if a.currentSegment.tokenCount.CountByNamespaceID == nil {
		t.Errorf("expected non-nil currentSegment.tokenCount.CountByNamespaceID")
	}

	if len(a.currentSegment.currentEntities.Entities) > 0 {
		t.Errorf("expected no current entity segment to be loaded. got: %v", a.currentSegment.currentEntities)
	}
	if len(a.entityTracker.activeEntities) > 0 {
		t.Errorf("expected no active entity segment to be loaded. got: %v", a.entityTracker.activeEntities)
	}
	if len(a.currentSegment.tokenCount.CountByNamespaceID) > 0 {
		t.Errorf("expected no token counts to be loaded. got: %v", a.currentSegment.tokenCount.CountByNamespaceID)
	}

	if verifyTimeNotZero {
		if a.currentSegment.startTimestamp == 0 {
			t.Error("bad start timestamp. expected no reset but timestamp was reset")
		}
	} else if a.currentSegment.startTimestamp != expectedStart {
		t.Errorf("bad start timestamp. expected: %v got: %v", expectedStart, a.currentSegment.startTimestamp)
	}
}

// ActiveEntitiesEqual checks that only the set of `test` exists in `active`
func ActiveEntitiesEqual(active map[string]struct{}, test []*activity.EntityRecord) bool {
	if len(active) != len(test) {
		return false
	}

	for _, ent := range test {
		if _, ok := active[ent.EntityID]; !ok {
			return false
		}
	}

	return true
}

// GetStartTimestamp returns the start timestamp on an activity log
func (a *ActivityLog) GetStartTimestamp() int64 {
	a.l.RLock()
	defer a.l.RUnlock()
	return a.currentSegment.startTimestamp
}

// SetStartTimestamp sets the start timestamp on an activity log
func (a *ActivityLog) SetStartTimestamp(timestamp int64) {
	a.l.Lock()
	defer a.l.Unlock()
	a.currentSegment.startTimestamp = timestamp
}

// SetTokenCount sets the tokenCount on an activity log
func (a *ActivityLog) SetTokenCount(tokenCount *activity.TokenCount) {
	a.l.Lock()
	defer a.l.Unlock()
	a.currentSegment.tokenCount = tokenCount
}

// GetCountByNamespaceID returns the count of tokens by namespace ID
func (a *ActivityLog) GetCountByNamespaceID() map[string]uint64 {
	a.l.RLock()
	defer a.l.RUnlock()
	return a.currentSegment.tokenCount.CountByNamespaceID
}

// GetEntitySequenceNumber returns the current entity sequence number
func (a *ActivityLog) GetEntitySequenceNumber() uint64 {
	a.l.RLock()
	defer a.l.RUnlock()
	return a.currentSegment.entitySequenceNumber
}

// SetEnable sets the enabled flag on the activity log
func (a *ActivityLog) SetEnable(enabled bool) {
	a.l.Lock()
	defer a.l.Unlock()
	a.fragmentLock.Lock()
	defer a.fragmentLock.Unlock()
	a.enabled = enabled
}

// GetEnabled returns the enabled flag on an activity log
func (a *ActivityLog) GetEnabled() bool {
	a.fragmentLock.RLock()
	defer a.fragmentLock.RUnlock()
	return a.enabled
}

// GetActivityLog returns a pointer to the (private) activity log on a core
// Note: you must do the usual locking scheme when modifying the ActivityLog
func (c *Core) GetActivityLog() *ActivityLog {
	return c.activityLog
}