<|start_filename|>internal/core/bigsegments/sync_test.go<|end_filename|>
package bigsegments

import (
	"net/http/httptest"
	"sort"
	"sync"
	"testing"
	"time"

	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldtime"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type bigSegmentStoreMock struct {
	cursor     string
	lock       sync.Mutex
	patchCh    chan bigSegmentPatch
	syncTimeCh chan ldtime.UnixMillisecondTime
}

func (s *bigSegmentStoreMock) applyPatch(patch bigSegmentPatch) (bool, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.cursor != patch.PreviousVersion {
		return false, nil
	}
	s.cursor = patch.Version

	s.patchCh <- patch

	return true, nil
}

func (s *bigSegmentStoreMock) getCursor() (string, error) {
	s.lock.Lock()
	defer s.lock.Unlock()

	return s.cursor, nil
}

func (s *bigSegmentStoreMock) setSynchronizedOn(synchronizedOn ldtime.UnixMillisecondTime) error {
	s.syncTimeCh <- synchronizedOn

	return nil
}

func (s *bigSegmentStoreMock) GetSynchronizedOn() (ldtime.UnixMillisecondTime, error) {
	return 0, nil
}

func (s *bigSegmentStoreMock) Close() error {
	return nil
}

func newBigSegmentStoreMock() *bigSegmentStoreMock {
	return &bigSegmentStoreMock{
		patchCh:    make(chan bigSegmentPatch, 100),
		syncTimeCh: make(chan ldtime.UnixMillisecondTime, 100),
	}
}

func assertPollRequest(t *testing.T, req httphelpers.HTTPRequestInfo, afterVersion string) {
	assert.Equal(t, string(testSDKKey), req.Request.Header.Get("Authorization"))
	assert.Equal(t, unboundedPollPath, req.Request.URL.Path)
	if afterVersion == "" {
		assert.Equal(t, "", req.Request.URL.RawQuery)
	} else {
		assert.Equal(t, "after="+afterVersion, req.Request.URL.RawQuery)
	}
}

func assertStreamRequest(t *testing.T, req httphelpers.HTTPRequestInfo) {
	assert.Equal(t, string(testSDKKey), req.Request.Header.Get("Authorization"))
	assert.Equal(t, unboundedStreamPath, req.Request.URL.Path)
}

func requirePatch(t *testing.T, s *bigSegmentStoreMock, expectedPatch bigSegmentPatch) {
	select {
	case patch := <-s.patchCh:
		require.Equal(t, expectedPatch, patch)
	case <-time.After(time.Second):
		require.Fail(t, "timed out waiting for patch")
	}
}

func requireUpdates(t *testing.T, ch <-chan UpdatesSummary, expectedKeys []string) {
	select {
	case u := <-ch:
		sort.Strings(u.SegmentKeysUpdated)
		sort.Strings(expectedKeys)
		require.Equal(t, expectedKeys, u.SegmentKeysUpdated)
	case <-time.After(time.Second):
		require.Fail(t, "timed out waiting for updates")
	}
}

func TestBasicSync(t *testing.T) {
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	patch1 := newPatchBuilder("segment.g1", "1", "").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	patch2 := newPatchBuilder("segment.g1", "2", "1").
		removeIncludes("included1").removeExcludes("excluded1").build()

	pollHandler, requestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{patch1}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
		),
	)

	sseHandler, _ := httphelpers.SSEHandler(makePatchEvent(patch2))
	streamHandler, streamRequestsCh := httphelpers.RecordingHandler(sseHandler)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamHandler, func(streamServer *httptest.Server) {
			startTime := ldtime.UnixMillisNow()

			storeMock := newBigSegmentStoreMock()
			defer storeMock.Close()

			segmentSync := newDefaultBigSegmentSynchronizer(sharedtest.MakeBasicHTTPConfig(), storeMock,
				pollServer.URL, streamServer.URL, config.EnvironmentID("env-xyz"), testSDKKey, mockLog.Loggers, "")
			defer segmentSync.Close()
			segmentSync.Start()

			updatesCh := segmentSync.SegmentUpdatesCh()
			go func() {
				for range updatesCh {
				}
			}() // just ensures that the synchronizer won't be blocked by the channel

			pollReq1 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq1, "")
			requirePatch(t, storeMock, patch1)
			require.Equal(t, 0, len(storeMock.syncTimeCh))

			pollReq2 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq2, patch1.Version)

			pollReq3 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq3, patch1.Version)

			require.Equal(t, 0, len(storeMock.patchCh))

			sharedtest.ExpectNoTestRequests(t, requestsCh, time.Millisecond*50)

			syncTime := <-storeMock.syncTimeCh
			assert.True(t, syncTime >= startTime)
			assert.True(t, syncTime <= ldtime.UnixMillisNow())

			streamReq1 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq1)
			requirePatch(t, storeMock, patch2)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			assert.Equal(t, []string{
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
			}, mockLog.GetOutput(ldlog.Info))
			assert.Len(t, mockLog.GetOutput(ldlog.Warn), 0)
		})
	})
}

func TestSyncSendsUpdates(t *testing.T) {
	// Scenario:
	// - Polling returns 3 patches (in 2 poll responses); these are aggregated into one UpdatesSummary
	// - Then the stream returns 1 more patch which generates another UpdatesSummary
	// We're also testing that segment IDs are aggregated into segment keys, i.e. "segment1.g1" and
	// "segment1.g2" together are reported as one update to "segment1".
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	poll1Patch1 := newPatchBuilder("segment1.g1", "1", "").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	poll2Patch1 := newPatchBuilder("segment1.g2", "2", "1").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	poll2Patch2 := newPatchBuilder("segment2.g3", "3", "2").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	streamPatch := newPatchBuilder("segment2.g3", "4", "3").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()

	pollHandler, requestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{poll1Patch1}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{poll2Patch1, poll2Patch2}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
		),
	)

	sseHandler, _ := httphelpers.SSEHandler(makePatchEvent(streamPatch))
	streamHandler, streamRequestsCh := httphelpers.RecordingHandler(sseHandler)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamHandler, func(streamServer *httptest.Server) {
			storeMock := newBigSegmentStoreMock()
			defer storeMock.Close()

			segmentSync := newDefaultBigSegmentSynchronizer(sharedtest.MakeBasicHTTPConfig(), storeMock,
				pollServer.URL, streamServer.URL, config.EnvironmentID("env-xyz"), testSDKKey, mockLog.Loggers, "")
			defer segmentSync.Close()
			segmentSync.Start()

			updatesCh := segmentSync.SegmentUpdatesCh()

			pollReq1 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq1, "")
			requirePatch(t, storeMock, poll1Patch1)
			require.Equal(t, 0, len(storeMock.syncTimeCh))

			pollReq2 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq2, poll1Patch1.Version)
			requirePatch(t, storeMock, poll2Patch1)
			requirePatch(t, storeMock, poll2Patch2)

			pollReq3 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq3, poll2Patch2.Version)

			pollReq4 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq4, poll2Patch2.Version)

			require.Equal(t, 0, len(storeMock.patchCh))

			requireUpdates(t, updatesCh, []string{"segment1", "segment2"})

			sharedtest.ExpectNoTestRequests(t, requestsCh, time.Millisecond*50)

			streamReq1 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq1)
			requirePatch(t, storeMock, streamPatch)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			requireUpdates(t, updatesCh, []string{"segment2"})
		})
	})
}

func TestSyncSkipsOutOfOrderUpdateFromPoll(t *testing.T) {
	// Scenario:
	// - Poll returns 3 patches: first patch is valid, second patch is non-matching, third is matching
	// - We apply the first patch
	// - Second patch causes a warning and causes remainder of list to be skipped
	// - Then we proceed with stream request as usual
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	patch1 := newPatchBuilder("segment.g1", "1", "").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	patch1x := newPatchBuilder("segment.g1", "1x", "non-matching-previous-version").
		addIncludes("includedx").addExcludes("excludedx").build()
	patch1y := newPatchBuilder("segment.g1", "2", "1").
		addIncludes("includedy").addExcludes("excludedy").build()
	patch2 := newPatchBuilder("segment.g1", "2", "1").
		removeIncludes("included1").removeExcludes("excluded1").build()

	pollHandler, requestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{patch1, patch1x, patch1y}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
		),
	)

	sseHandler, _ := httphelpers.SSEHandler(makePatchEvent(patch2))
	streamHandler, streamRequestsCh := httphelpers.RecordingHandler(sseHandler)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamHandler, func(streamServer *httptest.Server) {
			startTime := ldtime.UnixMillisNow()

			storeMock := newBigSegmentStoreMock()
			defer storeMock.Close()

			segmentSync := newDefaultBigSegmentSynchronizer(sharedtest.MakeBasicHTTPConfig(), storeMock,
				pollServer.URL, streamServer.URL, config.EnvironmentID("env-xyz"), testSDKKey, mockLog.Loggers, "")
			defer segmentSync.Close()
			segmentSync.Start()

			pollReq1 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq1, "")
			requirePatch(t, storeMock, patch1)
			require.Equal(t, 0, len(storeMock.syncTimeCh))

			pollReq2 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq2, patch1.Version)

			pollReq3 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq3, patch1.Version)

			require.Equal(t, 0, len(storeMock.patchCh))
			sharedtest.ExpectNoTestRequests(t, requestsCh, time.Millisecond*50)

			syncTime := <-storeMock.syncTimeCh
			assert.True(t, syncTime >= startTime)
			assert.True(t, syncTime <= ldtime.UnixMillisNow())

			streamReq1 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq1)
			requirePatch(t, storeMock, patch2)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			assert.Equal(t, []string{
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
			}, mockLog.GetOutput(ldlog.Info))
			mockLog.AssertMessageMatch(t, true, ldlog.Warn, `"non-matching-previous-version" which was not the latest`)
		})
	})
}

func TestSyncSkipsOutOfOrderUpdateFromStreamAndRestartsStream(t *testing.T) {
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	patch1 := newPatchBuilder("segment.g1", "1", "").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	patch2x := newPatchBuilder("segment.g1", "2", "non-matching-previous-version").
		removeIncludes("included1").removeExcludes("excluded1").build()
	patch2 := newPatchBuilder("segment.g1", "2", "1").
		removeIncludes("included1").removeExcludes("excluded1").build()

	pollHandler, requestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{patch1}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),
		),
	)

	firstStream, _ := httphelpers.SSEHandler(makePatchEvent(patch2x))
	secondStream, _ := httphelpers.SSEHandler(makePatchEvent(patch2))
	streamHandler, streamRequestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(firstStream, secondStream),
	)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamHandler, func(streamServer *httptest.Server) {
			startTime := ldtime.UnixMillisNow()

			storeMock := newBigSegmentStoreMock()
			defer storeMock.Close()

			segmentSync := newDefaultBigSegmentSynchronizer(sharedtest.MakeBasicHTTPConfig(), storeMock,
				pollServer.URL, streamServer.URL, config.EnvironmentID("env-xyz"), testSDKKey, mockLog.Loggers, "")
			segmentSync.streamRetryInterval = time.Millisecond
			defer segmentSync.Close()
			segmentSync.Start()

			pollReq1 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq1, "")
			requirePatch(t, storeMock, patch1)
			require.Equal(t, 0, len(storeMock.syncTimeCh))

			pollReq2 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq2, patch1.Version)

			pollReq3 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq3, patch1.Version)

			require.Equal(t, 0, len(storeMock.patchCh))

			syncTime := <-storeMock.syncTimeCh
			assert.True(t, syncTime >= startTime)
			assert.True(t, syncTime <= ldtime.UnixMillisNow())

			streamReq1 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq1)

			pollReq4 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq4, patch1.Version)

			streamReq2 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq2)

			pollReq5 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq5, patch1.Version)
			sharedtest.ExpectNoTestRequests(t, requestsCh, time.Millisecond*50)

			requirePatch(t, storeMock, patch2)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			assert.Equal(t, []string{
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
			}, mockLog.GetOutput(ldlog.Info))
			mockLog.AssertMessageMatch(t, true, ldlog.Warn, `"non-matching-previous-version" which was not the latest`)
		})
	})
}

func TestSyncRetryIfStreamFails(t *testing.T) {
	// In this test, we set up a successful poll and stream. Then we force the stream to close.
	// The synchronizer should start over with a new poll and stream.
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	patch1 := newPatchBuilder("segment.g1", "1", "").build()
	patch2 := newPatchBuilder("segment.g1", "2", "1").build()
	patch3 := newPatchBuilder("segment.g1", "3", "2").build()
	patch4 := newPatchBuilder("segment.g1", "4", "3").build()

	pollHandler, requestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{patch1}, nil), // poll 1: initial connection
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),       // poll 2: completion of poll 1
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),       // poll 3: done in conjunction with stream 1
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{patch3}, nil), // poll 4: retry after stream fails
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),       // poll 5: completion of poll 4
			httphelpers.HandlerWithJSONResponse([]bigSegmentPatch{}, nil),       // poll 6: done in conjunction with stream 2
		),
	)

	sseHandler1, sseControl1 := httphelpers.SSEHandler(makePatchEvent(patch2))
	sseHandler2, _ := httphelpers.SSEHandler(makePatchEvent(patch4))
	streamsHandler, streamRequestsCh := httphelpers.RecordingHandler(
		httphelpers.SequentialHandler(sseHandler1, sseHandler2),
	)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamsHandler, func(streamServer *httptest.Server) {
			startTime := ldtime.UnixMillisNow()

			storeMock := newBigSegmentStoreMock()
			defer storeMock.Close()

			segmentSync := newDefaultBigSegmentSynchronizer(sharedtest.MakeBasicHTTPConfig(), storeMock,
				pollServer.URL, streamServer.URL, config.EnvironmentID("env-xyz"), testSDKKey, mockLog.Loggers, "")
			segmentSync.streamRetryInterval = time.Millisecond
			defer segmentSync.Close()
			segmentSync.Start()

			pollReq1 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq1, "")
			requirePatch(t, storeMock, patch1)
			require.Equal(t, 0, len(storeMock.syncTimeCh))

			pollReq2 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq2, patch1.Version)

			pollReq3 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq3, patch1.Version)

			require.Equal(t, 0, len(storeMock.patchCh))

			sharedtest.ExpectNoTestRequests(t, requestsCh, time.Millisecond*50)

			syncTime := <-storeMock.syncTimeCh
			assert.True(t, syncTime >= startTime)
			assert.True(t, syncTime <= ldtime.UnixMillisNow())

			streamReq1 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq1)
			requirePatch(t, storeMock, patch2)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			// Now cause stream 1 to close
			sseControl1.Close()

			// Expect another poll+stream cycle; this time we get patch3 from the poll, patch4 from the stream
			pollReq4 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq4, patch2.Version)
			requirePatch(t, storeMock, patch3)

			pollReq5 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq5, patch3.Version)

			pollReq6 := sharedtest.ExpectTestRequest(t, requestsCh, time.Second)
			assertPollRequest(t, pollReq6, patch3.Version)

			streamReq2 := sharedtest.ExpectTestRequest(t, streamRequestsCh, time.Second)
			assertStreamRequest(t, streamReq2)
			requirePatch(t, storeMock, patch4)

			sharedtest.ExpectNoTestRequests(t, streamRequestsCh, time.Millisecond*50)

			assert.Equal(t, []string{
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
				"BigSegmentSynchronizer: Applied 1 update",
			}, mockLog.GetOutput(ldlog.Info))
			assert.Equal(t, []string{
				"BigSegmentSynchronizer: Stream connection failed: EOF",
				"BigSegmentSynchronizer: Will retry",
				"BigSegmentSynchronizer: Re-established connection",
			}, mockLog.GetOutput(ldlog.Warn))
			assert.Len(t, mockLog.GetOutput(ldlog.Error), 0)
		})
	})
}


<|start_filename|>internal/basictypes/package_info.go<|end_filename|>
// Package basictypes contains types and constants that are used by multiple Relay packages
// and do not have any testable logic of their own.
//
// Putting such things here, instead of in one of the packages that use them, allows us to
// reference them in shared test code as well without causing import cycles if those other
// packages also use the shared test code. It also provides a convenient way to see basic
// characteristics of Relay's internal model.
package basictypes


<|start_filename|>integrationtests/standard_mode_test.go<|end_filename|>
// +build integrationtests

package integrationtests

import (
	"testing"

	"github.com/stretchr/testify/require"
)

// See package_info.go for how these integration tests are configured and run.

type standardModeTestData struct {
	projsAndEnvs projsAndEnvs
}

func testStandardMode(t *testing.T, manager *integrationTestManager) {
	withStandardModeTestData(t, manager, func(testData standardModeTestData) {
		envVars := make(map[string]string)
		testData.projsAndEnvs.enumerateEnvs(func(proj projectInfo, env environmentInfo) {
			envVars["LD_ENV_"+string(env.name)] = string(env.sdkKey)
			envVars["LD_MOBILE_KEY_"+string(env.name)] = string(env.mobileKey)
			envVars["LD_CLIENT_SIDE_ID_"+string(env.name)] = string(env.id)
		})
		manager.startRelay(t, envVars)
		defer manager.stopRelay(t)

		manager.awaitEnvironments(t, testData.projsAndEnvs, false, func(proj projectInfo, env environmentInfo) string {
			return string(env.name)
		})
		manager.verifyFlagValues(t, testData.projsAndEnvs)
	})
}

func withStandardModeTestData(t *testing.T, manager *integrationTestManager, fn func(standardModeTestData)) {
	projsAndEnvs, err := manager.apiHelper.createProjectsAndEnvironments(2, 2)
	require.NoError(t, err)
	defer manager.apiHelper.deleteProjects(projsAndEnvs)

	require.NoError(t, manager.apiHelper.createFlags(projsAndEnvs))

	testData := standardModeTestData{
		projsAndEnvs: projsAndEnvs,
	}
	fn(testData)
}


<|start_filename|>internal/core/streams/env_streams.go<|end_filename|>
package streams

import (
	"sync"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
)

// EnvStreamUpdates is an interface representing the kinds of updates we can publish to streams. Other
// components that publish updates to EnvStreams should use this interface rather than the implementation
// type, both to clarify that they don't need other EnvStreams functionality and to simplify testing.
type EnvStreamUpdates interface {
	// SendAllDataUpdate signals an update to the entire SDK data set. The specified data items will be
	// broadcast to all connected server-side SDKs in a "put" event, and all connected client-side SDKs
	// will receive a "ping" event to refresh their state.
	SendAllDataUpdate(allData []ldstoretypes.Collection)

	// SendSingleItemUpdate signals an update to an individual SDK data item (flag or segment). The
	// specified data item will be broadcast to all connected server-side SDKs in a "patch" event, and
	// all connected client-side SDKs will receive a "ping" event to refresh their state.
	SendSingleItemUpdate(kind ldstoretypes.DataKind, key string, item ldstoretypes.ItemDescriptor)

	// InvalidateClientSideState signals an update that could affect client-side evaluations, but does
	// not change any server-side data item. Updating big segment data is such an event, since the
	// segment in the regular server-side SDK data does not change but evaluating the segment could
	// now produce a different result. Nothing is broadcast to the server-side SDKs (since they have
	// their own mechanisms for detecting big segment updates), but all connected client-side SDKs will
	// receive a "ping" event to refresh their state.
	InvalidateClientSideState()
}

// EnvStreams encapsulates streaming behavior for a specific environment.
//
// EnvStreams itself does not know anything about what kind of streams are available; those are
// determined by the StreamProvider instances that are passed in the constructor, and the credentials
// that are registered with AddCredential. For each combination of a credential and a StreamProvider
// that can handle that credential, a stream is available, and data updates that are sent with the
// EnvStreamUpdates methods will be rebroadcast to all of those streams, in a format that is
// determined by each StreamProvider.
type EnvStreams struct {
	streamProviders []StreamProvider
	storeQueries    EnvStoreQueries
	activeStreams   []streamInfo
	loggers         ldlog.Loggers
	lock            sync.RWMutex
	closeCh         chan struct{}
	heartbeatsDone  chan struct{} // used in testing only
}

type streamInfo struct {
	credential        config.SDKCredential
	envStreamProvider EnvStreamProvider
}

// EnvStoreQueries is a subset of DataStore methods that are used by EnvStreams to query existing
// data from the store, for generating "put" events.
type EnvStoreQueries interface {
	IsInitialized() bool
	GetAll(ldstoretypes.DataKind) ([]ldstoretypes.KeyedItemDescriptor, error)
}

// NewEnvStreams creates an instance of EnvStreams.
func NewEnvStreams(
	streamProviders []StreamProvider,
	storeQueries EnvStoreQueries,
	heartbeatInterval time.Duration,
	loggers ldlog.Loggers,
) *EnvStreams {
	es := &EnvStreams{
		streamProviders: streamProviders,
		storeQueries:    storeQueries,
		loggers:         loggers,
		closeCh:         make(chan struct{}),
	}

	if heartbeatInterval > 0 {
		heartbeats := time.NewTicker(heartbeatInterval)
		es.heartbeatsDone = make(chan struct{})
		go func() {
			for {
				select {
				case <-heartbeats.C:
					for _, esp := range es.getEnvStreamProviders() {
						esp.SendHeartbeat()
					}
				case <-es.closeCh:
					heartbeats.Stop()
					close(es.heartbeatsDone)
					return
				}
			}
		}()
	}

	return es
}

// AddCredential adds an environment credential and creates a corresponding EnvStreamProvider.
func (es *EnvStreams) AddCredential(credential config.SDKCredential) {
	if credential == nil {
		return
	}
	for _, sp := range es.streamProviders {
		if esp := sp.Register(credential, es.storeQueries, es.loggers); esp != nil {
			es.lock.Lock()
			es.activeStreams = append(es.activeStreams, streamInfo{credential, esp})
			es.lock.Unlock()
		}
	}
}

// RemoveCredential shuts down the EnvStreamProvider, if any, for the specified credential.
func (es *EnvStreams) RemoveCredential(credential config.SDKCredential) {
	var retained []streamInfo
	var removed []EnvStreamProvider

	es.lock.Lock()
	for _, s := range es.activeStreams {
		if s.credential == credential {
			removed = append(removed, s.envStreamProvider)
		} else {
			retained = append(retained, s)
		}
	}
	es.activeStreams = retained
	es.lock.Unlock()

	for _, s := range removed {
		s.Close()
	}
}

// SendAllDataUpdate sends all appropriate stream updates for when the full data set has been refreshed.
func (es *EnvStreams) SendAllDataUpdate(
	allData []ldstoretypes.Collection,
) {
	for _, esp := range es.getEnvStreamProviders() {
		esp.SendAllDataUpdate(allData)
	}
}

// SendSingleItemUpdate sends all appropriate stream updates for when an individual item has been updated.
func (es *EnvStreams) SendSingleItemUpdate(
	kind ldstoretypes.DataKind,
	key string,
	item ldstoretypes.ItemDescriptor,
) {
	for _, esp := range es.getEnvStreamProviders() {
		esp.SendSingleItemUpdate(kind, key, item)
	}
}

// InvalidateClientSideState sends all appropriate stream updates for when client-side state should be refreshed.
func (es *EnvStreams) InvalidateClientSideState() {
	for _, esp := range es.getEnvStreamProviders() {
		esp.InvalidateClientSideState()
	}
}

// Close shuts down all currently active streams for this environment and releases its resources.
func (es *EnvStreams) Close() error {
	close(es.closeCh)
	for _, esp := range es.getEnvStreamProviders() {
		esp.Close()
	}
	return nil
}

func (es *EnvStreams) getEnvStreamProviders() []EnvStreamProvider {
	es.lock.RLock()
	ret := make([]EnvStreamProvider, 0, len(es.activeStreams))
	for _, s := range es.activeStreams {
		ret = append(ret, s.envStreamProvider)
	}
	es.lock.RUnlock()
	return ret
}


<|start_filename|>internal/filedata/archive_manager_test.go<|end_filename|>
package filedata

import (
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	helpers "github.com/launchdarkly/go-test-helpers/v2"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
)

func TestStartWithValidFile(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, allTestEnvs...)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(allTestEnvs...)
	})
}

func TestStartWithMissingFile(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {}, func(p archiveManagerTestParams) {
		require.Error(t, p.archiveManagerError)
		assert.Contains(t, p.archiveManagerError.Error(), "unable to read file data source")
		assert.Contains(t, p.archiveManagerError.Error(), "no such file")
	})
}

func TestStartWithMalformedFile(t *testing.T) {
	archiveManagerTest(t, writeMalformedArchive, func(p archiveManagerTestParams) {
		require.Error(t, p.archiveManagerError)
	})
}

func TestStartWithFileWithBadChecksum(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, makeChecksumFileInvalidInArchive, allTestEnvs...)
	}, func(p archiveManagerTestParams) {
		require.Error(t, p.archiveManagerError)
		assert.Contains(t, p.archiveManagerError.Error(), "checksum of environments did not match")
	})
}

func TestStartWithFileWithNoEnvironments(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)
		p.mockLog.AssertMessageMatch(t, true, ldlog.Warn, logMsgNoEnvs)
	})
}

func TestStartWithFileWithOneBadEnvironment(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, func(dirPath string) {
			badData := []byte("that doesn't look like anything to me")
			require.NoError(t, ioutil.WriteFile(envMetadataFilePath(dirPath, testEnv1.rep.EnvID), badData, 0660))
			rehash(dirPath, testEnv1.rep.EnvID, testEnv2.rep.EnvID)
		}, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)
		p.mockLog.AssertMessageMatch(t, true, ldlog.Error, fmt.Sprintf(logMsgBadEnvData, testEnv1.rep.EnvID))
		p.expectEnvironmentsAdded(testEnv2)
	})
}

func TestDefaultRetryInterval(t *testing.T) {
	helpers.WithTempFile(func(filePath string) {
		writeArchive(t, filePath, false, nil)

		mockLog := ldlogtest.NewMockLog()
		defer mockLog.DumpIfTestFailed(t)

		messageHandler := newTestMessageHandler()

		archiveManager, err := NewArchiveManager(
			filePath,
			messageHandler,
			0,
			mockLog.Loggers,
		)
		require.NoError(t, err)
		defer archiveManager.Close()

		assert.Equal(t, defaultRetryInterval, archiveManager.retryInterval)
	})
}

func TestFileUpdatedWithValidDataAddedEnvironment(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1)

		writeArchive(t, p.filePath, false, nil, testEnv1, testEnv2)

		p.expectEnvironmentsAdded(testEnv2)
		p.expectReloaded()
	})
}

func TestFileUpdatedWithValidDataUpdatedEnvironmentMetadataOnly(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		testEnv1a := testEnv1.withMetadataChange()
		writeArchive(t, p.filePath, false, nil, testEnv1a, testEnv2)

		p.expectEnvironmentsUpdated(testEnv1a.withoutSDKData())
		p.expectReloaded()
	})
}

func TestFileUpdatedWithValidDataUpdatedEnvironmentSDKData(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		testEnv1a := testEnv1.withSDKDataChange()
		writeArchive(t, p.filePath, false, nil, testEnv1a, testEnv2)

		p.expectEnvironmentsUpdated(testEnv1a)
		p.expectReloaded()
	})
}

func TestFileUpdatedWithValidDataDeletedEnvironment(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		writeArchive(t, p.filePath, false, nil, testEnv2)

		p.expectEnvironmentsDeleted(testEnv1.rep.EnvID)
		p.expectReloaded()
	})
}

func TestFileUpdatedWithInvalidDataAndDoesNotBecomeValid(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		writeMalformedArchive(p.filePath)

		requireLogMessage(t, p.mockLog, ldlog.Error, "giving up until next change")

		p.requireNoMoreMessages()
	})
}

func TestFileUpdatedWithInvalidDataAndThenValidData(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		writeMalformedArchive(p.filePath)

		requireLogMessage(t, p.mockLog, ldlog.Warn, "file is invalid")

		testEnv1a := testEnv1.withMetadataChange().withSDKDataChange()
		writeArchive(t, p.filePath, false, nil, testEnv1a, testEnv2)

		// Because writeArchive updates the file in-place, we expect several file watch events to be triggered
		// in a row as the data is written incrementally, so this should exercise our logic for retrying after
		// consecutive errors.

		p.expectEnvironmentsUpdated(testEnv1a)
		p.expectReloaded()
	})
}

func TestFileDeletedAndThenRecreatedWithValidData(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		require.NoError(t, os.Remove(p.filePath))

		requireLogMessage(t, p.mockLog, ldlog.Warn, "file not found")

		testEnv1a := testEnv1.withMetadataChange().withSDKDataChange()
		writeArchive(t, p.filePath, false, nil, testEnv1a, testEnv2)

		p.expectEnvironmentsUpdated(testEnv1a)
		p.expectReloaded()
	})
}

func TestFileDeletedAndThenRecreatedWithInvalidDataAndThenValidData(t *testing.T) {
	archiveManagerTest(t, func(filePath string) {
		writeArchive(t, filePath, false, nil, testEnv1, testEnv2)
	}, func(p archiveManagerTestParams) {
		require.NoError(t, p.archiveManagerError)

		p.expectEnvironmentsAdded(testEnv1, testEnv2)

		require.NoError(t, os.Remove(p.filePath))

		requireLogMessage(t, p.mockLog, ldlog.Warn, "file not found")

		writeMalformedArchive(p.filePath)

		requireLogMessage(t, p.mockLog, ldlog.Warn, "file is invalid")

		testEnv1a := testEnv1.withMetadataChange().withSDKDataChange()
		writeArchive(t, p.filePath, false, nil, testEnv1a, testEnv2)

		p.expectEnvironmentsUpdated(testEnv1a)
		p.expectReloaded()
	})
}

func requireLogMessage(t *testing.T, mockLog *ldlogtest.MockLog, level ldlog.LogLevel, expectedSubstring string) {
	require.Eventuallyf(t, func() bool {
		warnings := mockLog.GetOutput(level)
		for _, w := range warnings {
			if strings.Contains(w, expectedSubstring) {
				return true
			}
		}
		return false
	}, time.Second*5, time.Millisecond*100, "wanted log message (%s) containing %q", level, expectedSubstring)
}


<|start_filename|>internal/core/streams/env_streams_test.go<|end_filename|>
package streams

import (
	"net/http"
	"reflect"
	"sync"
	"testing"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type mockStreamProvider struct {
	credentialOfDesiredType config.SDKCredential
	createdStreams          []*mockEnvStreamProvider
}

type mockEnvStreamProvider struct {
	parent         *mockStreamProvider
	credential     config.SDKCredential
	store          EnvStoreQueries
	allDataUpdates [][]ldstoretypes.Collection
	itemUpdates    []sharedtest.ReceivedItemUpdate
	clientSideUps  int
	numHeartbeats  int
	closed         bool
	lock           sync.Mutex
}

func (p *mockStreamProvider) Handler(credential config.SDKCredential) http.HandlerFunc {
	return nil
}

func (p *mockStreamProvider) Register(
	credential config.SDKCredential,
	store EnvStoreQueries,
	loggers ldlog.Loggers,
) EnvStreamProvider {
	if reflect.TypeOf(credential) != reflect.TypeOf(p.credentialOfDesiredType) {
		return nil
	}
	esp := &mockEnvStreamProvider{parent: p, credential: credential, store: store}
	p.createdStreams = append(p.createdStreams, esp)
	return esp
}

func (p *mockStreamProvider) Close() {}

func (e *mockEnvStreamProvider) SendAllDataUpdate(allData []ldstoretypes.Collection) {
	e.allDataUpdates = append(e.allDataUpdates, allData)
}

func (e *mockEnvStreamProvider) SendSingleItemUpdate(kind ldstoretypes.DataKind, key string, item ldstoretypes.ItemDescriptor) {
	e.itemUpdates = append(e.itemUpdates, sharedtest.ReceivedItemUpdate{kind, key, item})
}

func (e *mockEnvStreamProvider) InvalidateClientSideState() {
	e.clientSideUps++
}

func (e *mockEnvStreamProvider) SendHeartbeat() {
	e.lock.Lock()
	defer e.lock.Unlock()
	e.numHeartbeats++
}

func (e *mockEnvStreamProvider) Close() {
	e.closed = true
}

func (e *mockEnvStreamProvider) getNumHeartbeats() int {
	e.lock.Lock()
	defer e.lock.Unlock()
	return e.numHeartbeats
}

func TestAddCredential(t *testing.T) {
	sp1 := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}
	sp2 := &mockStreamProvider{credentialOfDesiredType: config.MobileKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp1, sp2}, store, 0, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key1")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)

	mobileKey := config.MobileKey("mobile-key")
	es.AddCredential(mobileKey)

	unsupportedKey := config.EnvironmentID("x")
	es.AddCredential(unsupportedKey)

	es.AddCredential(nil)

	require.Len(t, sp1.createdStreams, 2)
	esp1, esp2 := sp1.createdStreams[0], sp1.createdStreams[1]
	assert.Equal(t, sdkKey1, esp1.credential)
	assert.Equal(t, sdkKey2, esp2.credential)
	assert.Equal(t, store, esp1.store)
	assert.Equal(t, store, esp2.store)

	require.Len(t, sp2.createdStreams, 1)
	esp3 := sp2.createdStreams[0]
	assert.Equal(t, mobileKey, esp3.credential)
	assert.Equal(t, store, esp3.store)
}

func TestRemoveCredential(t *testing.T) {
	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, 0, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)

	require.Len(t, sp.createdStreams, 2)
	esp1, esp2 := sp.createdStreams[0], sp.createdStreams[1]
	assert.Equal(t, sdkKey1, esp1.credential)
	assert.Equal(t, sdkKey2, esp2.credential)
	assert.False(t, esp1.closed)
	assert.False(t, esp2.closed)

	es.RemoveCredential(sdkKey2)
	assert.False(t, esp1.closed)
	assert.True(t, esp2.closed)
}

func TestCloseEnvStreamsClosesAll(t *testing.T) {
	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, 0, ldlog.NewDisabledLoggers())

	sdkKey1, sdkKey2, sdkKey3 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2"), config.SDKKey("sdk-key3")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)
	es.AddCredential(sdkKey3)

	require.Len(t, sp.createdStreams, 3)
	esp1, esp2, esp3 := sp.createdStreams[0], sp.createdStreams[1], sp.createdStreams[2]

	es.RemoveCredential(sdkKey2)
	esp2.closed = false
	assert.False(t, esp1.closed)
	assert.False(t, esp3.closed)

	es.Close()

	assert.True(t, esp1.closed)
	assert.True(t, esp3.closed)
	assert.False(t, esp2.closed)
}

func TestSendAllDataUpdateGoesToAllStreams(t *testing.T) {
	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, 0, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2, sdkKey3 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2"), config.SDKKey("sdk-key3")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)
	es.AddCredential(sdkKey3)

	require.Len(t, sp.createdStreams, 3)
	esp1, esp2, esp3 := sp.createdStreams[0], sp.createdStreams[1], sp.createdStreams[2]

	es.RemoveCredential(sdkKey2)

	es.SendAllDataUpdate(allData)
	expected := [][]ldstoretypes.Collection{allData}

	assert.Equal(t, expected, esp1.allDataUpdates)
	assert.Len(t, esp2.allDataUpdates, 0)
	assert.Equal(t, expected, esp3.allDataUpdates)
}

func TestSendSingleItemUpdateGoesToAllStreams(t *testing.T) {
	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, 0, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2, sdkKey3 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2"), config.SDKKey("sdk-key3")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)
	es.AddCredential(sdkKey3)

	require.Len(t, sp.createdStreams, 3)
	esp1, esp2, esp3 := sp.createdStreams[0], sp.createdStreams[1], sp.createdStreams[2]

	es.RemoveCredential(sdkKey2)

	es.SendSingleItemUpdate(ldstoreimpl.Features(), testFlag1.Key, sharedtest.FlagDesc(testFlag1))
	expected := []sharedtest.ReceivedItemUpdate{{ldstoreimpl.Features(), testFlag1.Key, sharedtest.FlagDesc(testFlag1)}}

	assert.Equal(t, expected, esp1.itemUpdates)
	assert.Len(t, esp2.itemUpdates, 0)
	assert.Equal(t, expected, esp3.itemUpdates)
}

func TestInvalidateClientSideStateGoesToAllStreams(t *testing.T) {
	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, 0, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2, sdkKey3 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2"), config.SDKKey("sdk-key3")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)
	es.AddCredential(sdkKey3)

	require.Len(t, sp.createdStreams, 3)
	esp1, esp2, esp3 := sp.createdStreams[0], sp.createdStreams[1], sp.createdStreams[2]

	es.RemoveCredential(sdkKey2)

	es.InvalidateClientSideState()

	assert.Equal(t, 1, esp1.clientSideUps)
	assert.Equal(t, 0, esp2.clientSideUps)
	assert.Equal(t, 1, esp3.clientSideUps)
}

func TestHeartbeatsGoToAllStreams(t *testing.T) {
	heartbeatInterval := time.Millisecond * 20

	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, heartbeatInterval, ldlog.NewDisabledLoggers())
	defer es.Close()

	sdkKey1, sdkKey2 := config.SDKKey("sdk-key1"), config.SDKKey("sdk-key2")
	es.AddCredential(sdkKey1)
	es.AddCredential(sdkKey2)

	require.Len(t, sp.createdStreams, 2)
	esp1, esp2 := sp.createdStreams[0], sp.createdStreams[1]

	<-time.After(heartbeatInterval * 4)

	assert.GreaterOrEqual(t, esp1.getNumHeartbeats(), 2)
	assert.GreaterOrEqual(t, esp2.getNumHeartbeats(), 2)
}

func TestHeartbeatsAreStopped(t *testing.T) {
	heartbeatInterval := time.Millisecond * 20

	sp := &mockStreamProvider{credentialOfDesiredType: config.SDKKey("")}

	store := makeMockStore(nil, nil)
	es := NewEnvStreams([]StreamProvider{sp}, store, heartbeatInterval, ldlog.NewDisabledLoggers())

	es.AddCredential(config.SDKKey("sdk-key1"))

	require.Len(t, sp.createdStreams, 1)
	esp1 := sp.createdStreams[0]

	<-time.After(heartbeatInterval * 2) // just to make sure the heartbeat goroutine has already started
	assert.GreaterOrEqual(t, esp1.getNumHeartbeats(), 1)

	es.Close()

	select {
	case _, ok := <-es.heartbeatsDone:
		assert.False(t, ok, "heartbeatsDone channel should have been closed")
	case <-time.After(time.Second):
		assert.Fail(t, "heartbeatsDone channel should have been closed")
	}
}


<|start_filename|>internal/core/sharedtest/configs.go<|end_filename|>
package sharedtest

import (
	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
)

func MakeBasicHTTPConfig() httpconfig.HTTPConfig {
	ret, err := httpconfig.NewHTTPConfig(config.ProxyConfig{}, nil, "", ldlog.NewDisabledLoggers())
	if err != nil {
		panic(err)
	}
	return ret
}


<|start_filename|>internal/core/bigsegments/package_info.go<|end_filename|>
// Package bigsegments contains logic for synchronizing big segments.
package bigsegments


<|start_filename|>integrationtests/all_test.go<|end_filename|>
// +build integrationtests

package integrationtests

import (
	"testing"

	"github.com/stretchr/testify/require"
)

// See package_info.go for how these integration tests are configured and run.

func TestEndToEnd(t *testing.T) {
	// All of the end-to-end tests are grouped together here so that we only have to do the
	// newIntegrationTestManager step once; that step is slow because it has to create the Relay
	// Docker image from scratch.
	manager, err := newIntegrationTestManager()
	require.NoError(t, err)

	defer manager.close(t)

	t.Run("standard mode", func(t *testing.T) {
		testStandardMode(t, manager)
	})

	t.Run("auto-configuration mode", func(t *testing.T) {
		testAutoConfig(t, manager)
	})

	t.Run("offline mode", func(t *testing.T) {
		testOfflineMode(t, manager)
	})

	t.Run("database integrations", func(t *testing.T) {
		testDatabaseIntegrations(t, manager)
	})

	t.Run("big segments", func(t *testing.T) {
		testBigSegments(t, manager)
	})
}


<|start_filename|>internal/basictypes/sdk_kinds.go<|end_filename|>
package basictypes

// SDKKind represents any of the supported SDK categories that has distinct behavior from the others.
type SDKKind string

const (
	// ServerSDK represents server-side SDKs, which use server-side endpoints and authenticate their requests
	// with an SDK key.
	ServerSDK SDKKind = "server"

	// MobileSDK represents mobile SDKs, which use mobile endpoints and authenticate their requests with a
	// mobile key.
	MobileSDK SDKKind = "mobile"

	// JSClientSDK represents client-side JavaScript-based SDKs, which use client-side endpoints and
	// authenticate their requests insecurely with an environment ID.
	JSClientSDK SDKKind = "js"
)


<|start_filename|>internal/core/sharedtest/endpoints.go<|end_filename|>
package sharedtest

import (
	"encoding/base64"
	"fmt"
	"net/http"

	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
)

// These methods should be used by all test code that makes HTTP requests to Relay endpoints.
// We are deliberately not sharing the same URL path strings between this test code and the non-test code
// that configures Relay's endpoint routing, so that if we accidentally change the routing, the tests
// will fail, rather than succeed based on incorrect paths.

// SimpleUserJSON is a basic user.
const SimpleUserJSON = `{"key":"userkey"}`

// ToBase64 is a shortcut for base64 encoding.
func ToBase64(s string) string {
	return base64.StdEncoding.EncodeToString([]byte(s))
}

// SDKRequestVariant represents distinctions between endpoints that are not described in full by
// basictypes.SDKKind or basictypes.StreamKind.
type SDKRequestVariant int

const (
	// Evalx indicates that a client-side request wants the "evalx" format rather than the older
	// value-only format.
	Evalx SDKRequestVariant = 2

	// ReportMode means a client-side request should be done with REPORT rather than GET.
	ReportMode SDKRequestVariant = 4
)

// MakeSDKStreamEndpointRequest creates a request to one of the streaming SDK endpoints.
func MakeSDKStreamEndpointRequest(
	baseURL string,
	kind basictypes.StreamKind,
	testEnv TestEnv,
	userJSON string,
	variant SDKRequestVariant,
) *http.Request {
	switch {
	case kind == basictypes.ServerSideStream:
		return BuildRequestWithAuth("GET", fmt.Sprintf("%s/all", baseURL), testEnv.Config.SDKKey, nil)
	case kind == basictypes.ServerSideFlagsOnlyStream:
		return BuildRequestWithAuth("GET", fmt.Sprintf("%s/flags", baseURL), testEnv.Config.SDKKey, nil)
	case kind == basictypes.MobilePingStream && (variant&ReportMode == 0):
		return BuildRequestWithAuth("GET",
			fmt.Sprintf("%s/m%s/%s", baseURL, evalOrEvalx(variant), ToBase64(userJSON)),
			testEnv.Config.MobileKey, nil)
	case kind == basictypes.MobilePingStream && (variant&ReportMode != 0):
		return BuildRequestWithAuth("REPORT",
			fmt.Sprintf("%s/m%s", baseURL, evalOrEvalx(variant)),
			testEnv.Config.MobileKey, []byte(ToBase64(userJSON)))
	case kind == basictypes.JSClientPingStream && (variant&ReportMode == 0):
		return BuildRequest("GET",
			fmt.Sprintf("%s/%s/%s/%s", baseURL, evalOrEvalx(variant), testEnv.Config.EnvID, ToBase64(userJSON)),
			nil, nil)
	case kind == basictypes.JSClientPingStream && (variant&ReportMode != 0):
		return BuildRequest("REPORT",
			fmt.Sprintf("%s/%s/%s", baseURL, evalOrEvalx(variant), testEnv.Config.EnvID),
			[]byte(ToBase64(userJSON)), nil)
	default:
		panic("invalid StreamKind value")
	}
}

func MakeSDKEvalEndpointRequest(baseURL string, kind basictypes.SDKKind, testEnv TestEnv, userJSON string, variant SDKRequestVariant) *http.Request {
	switch {
	case kind == basictypes.MobileSDK && (variant&ReportMode == 0):
		return BuildRequestWithAuth("GET",
			fmt.Sprintf("%s/msdk/%s/users/%s", baseURL, evalOrEvalx(variant), ToBase64(userJSON)),
			testEnv.Config.MobileKey, nil)
	case kind == basictypes.MobileSDK && (variant&ReportMode != 0):
		return BuildRequestWithAuth("REPORT",
			fmt.Sprintf("%s/msdk/%s/user", baseURL, evalOrEvalx(variant)),
			testEnv.Config.MobileKey, []byte(ToBase64(userJSON)))
	case kind == basictypes.JSClientSDK && (variant&ReportMode == 0):
		return BuildRequestWithAuth("GET",
			fmt.Sprintf("%s/sdk/%s/%s/users/%s", baseURL, evalOrEvalx(variant), testEnv.Config.EnvID, ToBase64(userJSON)),
			nil, nil)
	case kind == basictypes.JSClientSDK && (variant&ReportMode != 0):
		return BuildRequestWithAuth("REPORT",
			fmt.Sprintf("%s/sdk/%s/%s/user", baseURL, evalOrEvalx(variant), testEnv.Config.EnvID),
			nil, []byte(ToBase64(userJSON)))
	default:
		panic("invalid SDKKind value")
	}
}

func evalOrEvalx(variant SDKRequestVariant) string {
	if variant&Evalx == 0 {
		return "eval"
	}
	return "evalx"
}


<|start_filename|>internal/core/relay_core_end_to_end_test.go<|end_filename|>
package core

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/launchdarkly/eventsource"
	c "github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/relayenv"
	st "github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"

	"github.com/launchdarkly/go-configtypes"
	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
	"github.com/launchdarkly/go-test-helpers/v2/ldservices"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldvalue"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldbuilders"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// These tests validate the processing of requests through a full HTTP request-response cycle, using a
// real embedded HTTP server to host the Relay instance, but without actually connecting to LaunchDarkly.

var testFlag = ldbuilders.NewFlagBuilder("test-flag").Version(1).
	ClientSideUsingMobileKey(true).ClientSideUsingEnvironmentID(true). // ensures that it'll show up in all endpoints
	Build()

type relayCoreEndToEndTestParams struct {
	t            *testing.T
	requestsCh   <-chan httphelpers.HTTPRequestInfo
	relayURL     string
	relayMockLog *ldlogtest.MockLog
	loggers      ldlog.Loggers
	singleLogger ldlog.BaseLogger
}

func relayCoreEndToEndTest(t *testing.T, config c.Config, ldStreamHandler http.Handler, action func(p relayCoreEndToEndTestParams)) {
	relayMockLog := ldlogtest.NewMockLog()
	defer relayMockLog.DumpIfTestFailed(t)

	pollHandler := httphelpers.HandlerWithStatus(404)
	streamHandler, requestsCh := httphelpers.RecordingHandler(ldStreamHandler)
	eventsHandler := httphelpers.HandlerWithStatus(202)

	httphelpers.WithServer(pollHandler, func(pollServer *httptest.Server) {
		httphelpers.WithServer(streamHandler, func(streamServer *httptest.Server) {
			httphelpers.WithServer(eventsHandler, func(eventsServer *httptest.Server) {
				config.Main.BaseURI, _ = configtypes.NewOptURLAbsoluteFromString(pollServer.URL)
				config.Main.StreamURI, _ = configtypes.NewOptURLAbsoluteFromString(streamServer.URL)
				config.Events.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(eventsServer.URL)
				core, err := NewRelayCore(
					config,
					relayMockLog.Loggers,
					nil,
					"1.2.3",
					"FakeRelay",
					relayenv.LogNameIsEnvID,
				)
				require.NoError(t, err)
				defer core.Close()

				for _, env := range config.Environment {
					streamReq := st.ExpectTestRequest(t, requestsCh, time.Second*5)
					assert.Equal(t, string(env.SDKKey), streamReq.Request.Header.Get("Authorization"))
				}

				httphelpers.WithServer(core.MakeRouter(), func(relayServer *httptest.Server) {
					mockLog := ldlogtest.NewMockLog()
					mockLog.Loggers.SetPrefix("TestClient:")
					defer mockLog.DumpIfTestFailed(t)
					p := relayCoreEndToEndTestParams{
						t:            t,
						requestsCh:   requestsCh,
						relayURL:     relayServer.URL,
						relayMockLog: relayMockLog,
						loggers:      mockLog.Loggers,
						singleLogger: mockLog.Loggers.ForLevel(ldlog.Info),
					}
					action(p)
				})
			})
		})
	})
}

func (p relayCoreEndToEndTestParams) waitForSuccessfulInit() {
	p.waitForLogMessage(ldlog.Info, "Initialized LaunchDarkly client for", "Relay initialization")
}

func (p relayCoreEndToEndTestParams) waitForLogMessage(level ldlog.LogLevel, pattern, conditionDesc string) {
	require.Eventually(p.t, func() bool {
		return p.relayMockLog.HasMessageMatch(level, pattern)
	}, time.Second*2, time.Millisecond*50, "timed out waiting for %s", conditionDesc)
}

func (p relayCoreEndToEndTestParams) subscribeStream(testEnv st.TestEnv, kind basictypes.StreamKind) (
	*eventsource.Stream, *eventsource.SubscriptionError) {
	req := st.MakeSDKStreamEndpointRequest(p.relayURL, kind, testEnv, st.SimpleUserJSON, 0)
	stream, err := eventsource.SubscribeWithRequestAndOptions(req, eventsource.StreamOptionLogger(p.singleLogger))
	if err != nil {
		require.IsType(p.t, eventsource.SubscriptionError{}, err)
		se := err.(eventsource.SubscriptionError)
		return nil, &se
	}
	return stream, nil
}

func (p relayCoreEndToEndTestParams) expectStreamEvent(testEnv st.TestEnv, kind basictypes.StreamKind) eventsource.Event {
	stream, err := p.subscribeStream(testEnv, kind)
	require.Nil(p.t, err)
	require.NotNil(p.t, stream)
	defer stream.Close()
	return st.ExpectStreamEvent(p.t, stream, time.Second*5)
}

func (p relayCoreEndToEndTestParams) expectStreamWithNoEvent(testEnv st.TestEnv, kind basictypes.StreamKind) {
	stream, err := p.subscribeStream(testEnv, kind)
	require.Nil(p.t, err)
	require.NotNil(p.t, stream)
	defer stream.Close()
	st.ExpectNoStreamEvent(p.t, stream, time.Millisecond*100)
}

func (p relayCoreEndToEndTestParams) expectStreamError(testEnv st.TestEnv, kind basictypes.StreamKind, status int) {
	_, err := p.subscribeStream(testEnv, kind)
	require.NotNil(p.t, err)
	assert.Equal(p.t, status, err.Code)
}

func (p relayCoreEndToEndTestParams) expectEvalResult(testEnv st.TestEnv, kind basictypes.SDKKind) ldvalue.Value {
	req := st.MakeSDKEvalEndpointRequest(p.relayURL, kind, testEnv, st.SimpleUserJSON, 0)
	resp, err := http.DefaultClient.Do(req)
	require.NoError(p.t, err)
	defer resp.Body.Close()
	require.Equal(p.t, 200, resp.StatusCode)
	body, err := ioutil.ReadAll(resp.Body)
	require.NoError(p.t, err)
	return ldvalue.Parse(body)
}

func (p relayCoreEndToEndTestParams) expectEvalError(testEnv st.TestEnv, kind basictypes.SDKKind, status int) {
	req := st.MakeSDKEvalEndpointRequest(p.relayURL, kind, testEnv, st.SimpleUserJSON, 0)
	resp, err := http.DefaultClient.Do(req)
	require.NoError(p.t, err)
	require.Equal(p.t, status, resp.StatusCode)
}

func (p relayCoreEndToEndTestParams) expectSuccessFromAllEndpoints(testEnv st.TestEnv) {
	serverSideEvent := p.expectStreamEvent(testEnv, basictypes.ServerSideStream)
	assert.Equal(p.t, "put", serverSideEvent.Event())
	jsonData := ldvalue.Parse([]byte(serverSideEvent.Data()))
	assert.Equal(p.t, []string{testFlag.Key}, jsonData.GetByKey("data").GetByKey("flags").Keys())

	serverSideFlagsEvent := p.expectStreamEvent(testEnv, basictypes.ServerSideFlagsOnlyStream)
	assert.Equal(p.t, "put", serverSideFlagsEvent.Event())
	jsonFlagsData := ldvalue.Parse([]byte(serverSideFlagsEvent.Data()))
	assert.Equal(p.t, []string{testFlag.Key}, jsonFlagsData.Keys())

	mobileEvent := p.expectStreamEvent(testEnv, basictypes.MobilePingStream)
	assert.Equal(p.t, "ping", mobileEvent.Event())

	mobileEval := p.expectEvalResult(testEnv, basictypes.MobileSDK)
	assert.Equal(p.t, []string{testFlag.Key}, mobileEval.Keys())

	jsClientEval := p.expectEvalResult(testEnv, basictypes.JSClientSDK)
	assert.Equal(p.t, []string{testFlag.Key}, jsClientEval.Keys())
}

func TestRelayCoreEndToEndSuccess(t *testing.T) {
	putEvent := ldservices.NewServerSDKData().Flags(&testFlag).ToPutEvent()
	streamHandler, _ := ldservices.ServerSideStreamingServiceHandler(putEvent)
	testEnv := st.EnvWithAllCredentials

	config := c.Config{Environment: st.MakeEnvConfigs(testEnv)}
	relayCoreEndToEndTest(t, config, streamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForSuccessfulInit()
		p.expectSuccessFromAllEndpoints(testEnv)
	})
}

func TestRelayCoreEndToEndPermanentFailure(t *testing.T) {
	streamHandler := httphelpers.HandlerWithStatus(401)
	testEnv := st.EnvWithAllCredentials

	config := c.Config{Environment: st.MakeEnvConfigs(testEnv)}
	relayCoreEndToEndTest(t, config, streamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForLogMessage(ldlog.Error, "Error initializing LaunchDarkly client for", "initialization failure")
		p.expectStreamError(testEnv, basictypes.ServerSideStream, 401)
		p.expectStreamError(testEnv, basictypes.ServerSideFlagsOnlyStream, 401)
		p.expectStreamError(testEnv, basictypes.MobilePingStream, 401)
		p.expectStreamError(testEnv, basictypes.JSClientPingStream, 404)
		p.expectEvalError(testEnv, basictypes.MobileSDK, 401)
		p.expectEvalError(testEnv, basictypes.JSClientSDK, 404)
	})
}

func TestRelayCoreEndToEndInitTimeoutWithUninitializedDataStore(t *testing.T) {
	// For the "timed out, but we *do* have previous flag data" condition, see
	// TestRelayCoreEndToEndRedisInitTimeoutWithInitializedDataStore.
	putEvent := ldservices.NewServerSDKData().Flags(&testFlag).ToPutEvent()
	streamHandler, _ := ldservices.ServerSideStreamingServiceHandler(putEvent)
	delayBeforeStreamHandler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		select {
		case <-time.After(time.Second * 10):
			break
		case <-req.Context().Done(): // so the test won't be delayed when closing the test server
			break
		}
		streamHandler.ServeHTTP(w, req)
	})
	testEnv := st.EnvWithAllCredentials

	config := c.Config{
		Main: c.MainConfig{
			InitTimeout: configtypes.NewOptDuration(time.Millisecond),
		},
		Environment: st.MakeEnvConfigs(testEnv),
	}
	relayCoreEndToEndTest(t, config, delayBeforeStreamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForLogMessage(ldlog.Error, "timeout encountered waiting for LaunchDarkly client initialization",
			"initialization timeout")
		p.expectStreamWithNoEvent(testEnv, basictypes.ServerSideStream)
		p.expectStreamWithNoEvent(testEnv, basictypes.ServerSideFlagsOnlyStream)
		p.expectStreamWithNoEvent(testEnv, basictypes.MobilePingStream)
		p.expectStreamWithNoEvent(testEnv, basictypes.JSClientPingStream)
		p.expectEvalError(testEnv, basictypes.MobileSDK, 503)
		p.expectEvalError(testEnv, basictypes.JSClientSDK, 503)
	})
}


<|start_filename|>internal/core/bigsegments/store_common_test.go<|end_filename|>
// +build big_segment_external_store_tests

package bigsegments

import (
	"strconv"
	"testing"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldtime"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type bigSegmentOperations struct {
	isUserIncluded func(segmentKey string, userKey string) (bool, error)
	isUserExcluded func(segmentKey string, userKey string) (bool, error)
}

func testGenericAll(
	t *testing.T,
	withBigSegmentStore func(t *testing.T, action func(BigSegmentStore, bigSegmentOperations)),
) {
	patch1 := newPatchBuilder("segment.g1", "1", "").
		addIncludes("included1", "included2").addExcludes("excluded1", "excluded2").build()
	patch2 := newPatchBuilder("segment.g1", "2", "1").
		removeIncludes("included1").removeExcludes("excluded1").build()

	t.Run("synchronizedOn", func(t *testing.T) {
		withBigSegmentStore(t, func(store BigSegmentStore, operations bigSegmentOperations) {
			sync1, err := store.GetSynchronizedOn()
			require.NoError(t, err)
			assert.False(t, sync1.IsDefined())

			now := ldtime.UnixMillisNow()
			err = store.setSynchronizedOn(now)
			require.NoError(t, err)

			sync2, err := store.GetSynchronizedOn()
			require.NoError(t, err)
			require.Equal(t, now, sync2)
		})
	})

	t.Run("applyPatchSequence", func(t *testing.T) {
		withBigSegmentStore(t, func(store BigSegmentStore, operations bigSegmentOperations) {
			// first set synchronizedOn, so we can verify that applying a patch does *not* change that value
			initialSyncTime := ldtime.UnixMillisecondTime(99999)
			require.NoError(t, store.setSynchronizedOn(initialSyncTime))

			// apply initial patch that adds users
			success, err := store.applyPatch(patch1)
			require.NoError(t, err)
			require.True(t, success)

			cursor, err := store.getCursor()
			require.NoError(t, err)
			require.Equal(t, patch1.Version, cursor)

			membership, err := operations.isUserIncluded(patch1.SegmentID, patch1.Changes.Included.Add[0])
			require.NoError(t, err)
			assert.Equal(t, true, membership)

			membership, err = operations.isUserExcluded(patch1.SegmentID, patch1.Changes.Excluded.Add[0])
			require.NoError(t, err)
			assert.Equal(t, true, membership)

			// apply second patch in sequence that removes users
			success, err = store.applyPatch(patch2)
			require.NoError(t, err)
			require.True(t, success)

			cursor, err = store.getCursor()
			require.NoError(t, err)
			assert.Equal(t, patch2.Version, cursor)

			membership, err = operations.isUserIncluded(patch1.SegmentID, patch1.Changes.Included.Add[0])
			require.NoError(t, err)
			assert.Equal(t, false, membership)

			membership, err = operations.isUserExcluded(patch1.SegmentID, patch1.Changes.Excluded.Add[0])
			require.NoError(t, err)
			assert.Equal(t, false, membership)

			// apply old patch
			success, err = store.applyPatch(patch1)
			require.NoError(t, err)
			require.False(t, success)

			// verify that the stored cursor was updated
			cursor, err = store.getCursor()
			require.NoError(t, err)
			assert.Equal(t, patch2.Version, cursor)

			// verify that the sync time is still there
			syncTime, err := store.GetSynchronizedOn()
			require.NoError(t, err)
			assert.Equal(t, initialSyncTime, syncTime)

			// now update the sync time and verify that that doesn't affect the cursor
			newSyncTime := initialSyncTime + 1
			require.NoError(t, store.setSynchronizedOn(newSyncTime))
			syncTime, err = store.GetSynchronizedOn()
			require.NoError(t, err)
			assert.Equal(t, newSyncTime, syncTime)
			cursor, err = store.getCursor()
			require.NoError(t, err)
			assert.Equal(t, patch2.Version, cursor)
		})
	})

	t.Run("patchLarge", func(t *testing.T) {
		withBigSegmentStore(t, func(store BigSegmentStore, operations bigSegmentOperations) {
			userCount := 50
			var users []string = make([]string, 0, userCount)
			for i := 0; i < userCount; i++ {
				users = append(users, strconv.FormatUint(uint64(i), 10))
			}

			patch := bigSegmentPatch{
				EnvironmentID:   "abc",
				SegmentID:       "segment.g1",
				Version:         "1",
				PreviousVersion: "",
				Changes: bigSegmentPatchChanges{
					Included: bigSegmentPatchChangesMutations{
						Add: users,
					},
				},
			}

			success, err := store.applyPatch(patch)
			require.NoError(t, err)
			require.True(t, success)

			membership, err := operations.isUserIncluded(patch.SegmentID, strconv.FormatUint(uint64(userCount-1), 10))
			assert.Equal(t, true, membership)
			require.NoError(t, err)

			cursor, err := store.getCursor()
			require.NoError(t, err)
			assert.Equal(t, patch.Version, cursor)
		})
	})
}


<|start_filename|>internal/core/sdks/big_segments_test.go<|end_filename|>
package sdks

import (
	"testing"

	lddynamodb "github.com/launchdarkly/go-server-sdk-dynamodb"
	"github.com/launchdarkly/ld-relay/v6/config"

	"github.com/launchdarkly/go-configtypes"
	ldredis "github.com/launchdarkly/go-server-sdk-redis-redigo"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func assertBigSegmentsConfigured(
	t *testing.T,
	expected interfaces.BigSegmentsConfigurationFactory,
	c config.Config,
	ec config.EnvConfig,
) *ldlogtest.MockLog {
	mockLog := ldlogtest.NewMockLog()
	factory, err := ConfigureBigSegments(c, ec, mockLog.Loggers)
	require.NoError(t, err)
	assert.Equal(t, expected, factory)
	return mockLog
}

func TestBigSegmentsDefault(t *testing.T) {
	log := assertBigSegmentsConfigured(t, nil, config.Config{}, config.EnvConfig{})
	assert.Len(t, log.GetAllOutput(), 0)
}

func TestBigSegmentsRedis(t *testing.T) {
	redisURL := "redis://redishost:3000"
	redisSecureURL := "rediss://redishost:3000"
	optRedisURL, _ := configtypes.NewOptURLAbsoluteFromString(redisURL)

	t.Run("basic properties", func(t *testing.T) {
		c := config.Config{
			Redis: config.RedisConfig{
				URL: optRedisURL,
			},
		}
		expected := ldcomponents.BigSegments(ldredis.DataStore().URL(redisURL))
		log := assertBigSegmentsConfigured(t, expected, c, config.EnvConfig{})
		log.AssertMessageMatch(t, true, ldlog.Info, "Using Redis big segment store: "+redisURL)
	})

	t.Run("prefix", func(t *testing.T) {
		c := config.Config{
			Redis: config.RedisConfig{
				URL: optRedisURL,
			},
		}
		ec := config.EnvConfig{Prefix: "abc"}
		expected := ldcomponents.BigSegments(ldredis.DataStore().URL(redisURL).Prefix("abc"))
		log := assertBigSegmentsConfigured(t, expected, c, ec)
		log.AssertMessageMatch(t, true, ldlog.Info, "Using Redis big segment store: "+redisURL+" with prefix: abc")
	})

	t.Run("TLS", func(t *testing.T) {
		c := config.Config{
			Redis: config.RedisConfig{
				URL: optRedisURL,
				TLS: true,
			},
		}
		expected := ldcomponents.BigSegments(ldredis.DataStore().URL(redisSecureURL))
		log := assertBigSegmentsConfigured(t, expected, c, config.EnvConfig{})
		log.AssertMessageMatch(t, true, ldlog.Info, "Using Redis big segment store: "+redisSecureURL)
	})
}

func TestBigSegmentsDynamoDB(t *testing.T) {
	table := "my-table"

	t.Run("basic properties - global table name", func(t *testing.T) {
		c := config.Config{
			DynamoDB: config.DynamoDBConfig{
				Enabled:   true,
				TableName: table,
			},
		}
		expected := ldcomponents.BigSegments(lddynamodb.DataStore(table))
		log := assertBigSegmentsConfigured(t, expected, c, config.EnvConfig{})
		log.AssertMessageMatch(t, true, ldlog.Info, "Using DynamoDB big segment store: "+table)
	})

	t.Run("basic properties - per-environment table name", func(t *testing.T) {
		c := config.Config{
			DynamoDB: config.DynamoDBConfig{
				Enabled: true,
			},
		}
		ec := config.EnvConfig{TableName: table}
		expected := ldcomponents.BigSegments(lddynamodb.DataStore(table))
		log := assertBigSegmentsConfigured(t, expected, c, ec)
		log.AssertMessageMatch(t, true, ldlog.Info, "Using DynamoDB big segment store: "+table)
	})

	t.Run("prefix", func(t *testing.T) {
		c := config.Config{
			DynamoDB: config.DynamoDBConfig{
				Enabled:   true,
				TableName: table,
			},
		}
		ec := config.EnvConfig{Prefix: "abc"}
		expected := ldcomponents.BigSegments(lddynamodb.DataStore(table).Prefix("abc"))
		log := assertBigSegmentsConfigured(t, expected, c, ec)
		log.AssertMessageMatch(t, true, ldlog.Info, "Using DynamoDB big segment store: "+table+" with prefix: abc")
	})
}


<|start_filename|>internal/core/relay_core_endpoints.go<|end_filename|>
package core

import (
	"crypto/sha1" //nolint:gosec // we're not using SHA1 for encryption, just for generating an insecure hash
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"sort"
	"strconv"
	"time"

	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/logging"
	"github.com/launchdarkly/ld-relay/v6/internal/core/middleware"
	"github.com/launchdarkly/ld-relay/v6/internal/core/relayenv"
	"github.com/launchdarkly/ld-relay/v6/internal/core/streams"
	"github.com/launchdarkly/ld-relay/v6/internal/util"

	"gopkg.in/launchdarkly/go-jsonstream.v1/jwriter"
	"gopkg.in/launchdarkly/go-sdk-common.v2/lduser"
	ldevents "gopkg.in/launchdarkly/go-sdk-events.v1"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldmodel"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"

	"github.com/gorilla/mux"
)

func getClientSideUserProperties(
	clientCtx relayenv.EnvContext,
	sdkKind basictypes.SDKKind,
	req *http.Request,
	w http.ResponseWriter,
) (lduser.User, bool) {
	var user lduser.User
	var userDecodeErr error

	if req.Method == "REPORT" {
		if req.Header.Get("Content-Type") != "application/json" {
			w.WriteHeader(http.StatusUnsupportedMediaType)
			_, _ = w.Write([]byte("Content-Type must be application/json."))
			return user, false
		}
		body, _ := ioutil.ReadAll(req.Body)
		userDecodeErr = json.Unmarshal(body, &user)
	} else {
		base64User := mux.Vars(req)["user"]
		user, userDecodeErr = middleware.UserFromBase64(base64User)
	}
	if userDecodeErr != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write(util.ErrorJSONMsg(userDecodeErr.Error()))
		return user, false
	}

	if clientCtx.IsSecureMode() && sdkKind == basictypes.JSClientSDK {
		hash := req.URL.Query().Get("h")
		valid := false
		if hash != "" {
			validHash := clientCtx.GetClient().SecureModeHash(user)
			valid = hash == validHash
		}
		if !valid {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write(util.ErrorJSONMsg("Environment is in secure mode, and user hash does not match."))
			return user, false
		}
	}

	return user, true
}

// Old stream endpoint that just sends "ping" events: clientstream.ld.com/mping (mobile)
// or clientstream.ld.com/ping/{envId} (JS)
func pingStreamHandler(streamProvider streams.StreamProvider) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		clientCtx := middleware.GetEnvContextInfo(req.Context())
		clientCtx.Env.GetLoggers().Debug("Application requested client-side ping stream")
		clientCtx.Env.GetStreamHandler(streamProvider, clientCtx.Credential).ServeHTTP(w, req)
	})
}

// This handler is used for client-side streaming endpoints that require user properties. Currently it is
// implemented the same as the ping stream once we have validated the user.
func pingStreamHandlerWithUser(sdkKind basictypes.SDKKind, streamProvider streams.StreamProvider) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		clientCtx := middleware.GetEnvContextInfo(req.Context())
		clientCtx.Env.GetLoggers().Debug("Application requested client-side ping stream")

		if _, ok := getClientSideUserProperties(clientCtx.Env, sdkKind, req, w); ok {
			clientCtx.Env.GetStreamHandler(streamProvider, clientCtx.Credential).ServeHTTP(w, req)
		}
	})
}

// Multi-purpose streaming handler; all details of the behavior of the particular type of stream are
// abstracted in StreamProvider and EnvStreams
func streamHandler(streamProvider streams.StreamProvider, logMessage string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		clientCtx := middleware.GetEnvContextInfo(req.Context())
		clientCtx.Env.GetLoggers().Debug(logMessage)
		clientCtx.Env.GetStreamHandler(streamProvider, clientCtx.Credential).ServeHTTP(w, req)
	})
}

// PHP SDK polling endpoint for all flags: app.ld.com/sdk/flags
func pollAllFlagsHandler(w http.ResponseWriter, req *http.Request) {
	clientCtx := middleware.GetEnvContextInfo(req.Context())
	data, err := clientCtx.Env.GetStore().GetAll(ldstoreimpl.Features())
	if err != nil {
		clientCtx.Env.GetLoggers().Errorf("Error reading feature store: %s", err)
		w.WriteHeader(500)
		return
	}
	respData := serializeFlagsAsMap(data)
	// Compute an overall Etag for the data set by hashing flag keys and versions
	hash := sha1.New()                                                         // nolint:gas // just used for insecure hashing
	sort.Slice(data, func(i, j int) bool { return data[i].Key < data[j].Key }) // makes the hash deterministic
	for _, item := range data {
		_, _ = io.WriteString(hash, fmt.Sprintf("%s:%d", item.Key, item.Item.Version))
	}
	etag := hex.EncodeToString(hash.Sum(nil))[:15]
	writeCacheableJSONResponse(w, req, clientCtx.Env, respData, etag)
}

// PHP SDK polling endpoint for a flag: app.ld.com/sdk/flags/{key}
func pollFlagHandler(w http.ResponseWriter, req *http.Request) {
	pollFlagOrSegment(middleware.GetEnvContextInfo(req.Context()).Env, ldstoreimpl.Features())(w, req)
}

// PHP SDK polling endpoint for a segment: app.ld.com/sdk/segments/{key}
func pollSegmentHandler(w http.ResponseWriter, req *http.Request) {
	pollFlagOrSegment(middleware.GetEnvContextInfo(req.Context()).Env, ldstoreimpl.Segments())(w, req)
}

// Event-recorder endpoints:
// events.ld.com/bulk (server-side)
// events.ld.com/diagnostic (server-side diagnostic)
// events.ld.com/mobile, events.ld.com/mobile/events, events.ld.com/mobileevents/bulk (mobile)
// events.ld.com/mobile/events/diagnostic (mobile diagnostic)
// events.ld.com/events/bulk/{envId} (JS)
// events.ld.com/events/diagnostic/{envId} (JS)
func bulkEventHandler(sdkKind basictypes.SDKKind, eventsKind ldevents.EventDataKind, offline bool) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if offline {
			w.WriteHeader(http.StatusAccepted)
			if req.Body != nil {
				_ = req.Body.Close()
			}
			return
		}

		clientCtx := middleware.GetEnvContextInfo(req.Context())
		dispatcher := clientCtx.Env.GetEventDispatcher()
		if dispatcher == nil {
			w.WriteHeader(http.StatusServiceUnavailable)
			_, _ = w.Write(util.ErrorJSONMsg("Event proxy is not enabled for this environment"))
			return
		}
		handler := dispatcher.GetHandler(sdkKind, eventsKind)
		if handler == nil {
			// Note, if this ever happens, it is a programming error since we are only supposed to
			// be using a fixed set of Endpoint values that the dispatcher knows about.
			w.WriteHeader(http.StatusServiceUnavailable)
			_, _ = w.Write(util.ErrorJSONMsg("Internal error in event proxy"))
			logging.GetGlobalContextLoggers(req.Context()).Errorf("Tried to proxy %s events for %s but no handler was defined",
				eventsKind, sdkKind)
			return
		}
		handler(w, req)
	})
}

// Client-side evaluation endpoint, new schema with metadata:
// app.ld.com/sdk/evalx/{envId}/users/{user} (GET)
// app.ld.com/sdk/evalx/{envId}/user (REPORT)
// app.ld/com/sdk/evalx/users/{user} (GET - with SDK key auth)
// app.ld/com/sdk/evalx/user (REPORT - with SDK key auth)
func evaluateAllFeatureFlags(sdkKind basictypes.SDKKind) func(w http.ResponseWriter, req *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		evaluateAllShared(w, req, false, sdkKind)
	}
}

// Client-side evaluation endpoint, old schema with only values:
// app.ld.com/sdk/eval/{envId}/users/{user} (GET)
// app.ld.com/sdk/eval/{envId}/user (REPORT)
// app.ld/com/sdk/eval/users/{user} (GET - with SDK key auth)
// app.ld/com/sdk/eval/user (REPORT - with SDK key auth)
func evaluateAllFeatureFlagsValueOnly(sdkKind basictypes.SDKKind) func(w http.ResponseWriter, req *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		evaluateAllShared(w, req, true, sdkKind)
	}
}

func evaluateAllShared(w http.ResponseWriter, req *http.Request, valueOnly bool, sdkKind basictypes.SDKKind) {
	clientCtx := middleware.GetEnvContextInfo(req.Context())
	client := clientCtx.Env.GetClient()
	store := clientCtx.Env.GetStore()
	loggers := clientCtx.Env.GetLoggers()

	user, ok := getClientSideUserProperties(clientCtx.Env, sdkKind, req, w)
	if !ok {
		return
	}

	withReasons := req.URL.Query().Get("withReasons") == "true"

	w.Header().Set("Content-Type", "application/json")

	if !client.Initialized() {
		if store.IsInitialized() {
			loggers.Warn("Called before client initialization; using last known values from feature store")
		} else {
			w.WriteHeader(http.StatusServiceUnavailable)
			loggers.Warn("Called before client initialization. Feature store not available")
			_, _ = w.Write(util.ErrorJSONMsg("Service not initialized"))
			return
		}
	}

	if user.GetKey() == "" {
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write(util.ErrorJSONMsg("User must have a 'key' attribute"))
		return
	}

	loggers.Debugf("Application requested client-side flags (%s) for user: %s", sdkKind, user.GetKey())

	items, err := store.GetAll(ldstoreimpl.Features())
	if err != nil {
		loggers.Warnf("Unable to fetch flags from feature store. Returning nil map. Error: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write(util.ErrorJSONMsgf("Error fetching flags from feature store: %s", err))
		return
	}

	evaluator := clientCtx.Env.GetEvaluator()

	responseWriter := jwriter.NewWriter()
	responseObj := responseWriter.Object()
	for _, item := range items {
		if flag, ok := item.Item.Item.(*ldmodel.FeatureFlag); ok {
			switch sdkKind {
			case basictypes.JSClientSDK:
				if !flag.ClientSideAvailability.UsingEnvironmentID {
					continue
				}
			case basictypes.MobileSDK:
				if !flag.ClientSideAvailability.UsingMobileKey {
					continue
				}
			}
			detail := evaluator.Evaluate(flag, user, nil)
			if valueOnly {
				detail.Value.WriteToJSONWriter(responseObj.Name(flag.Key))
			} else {
				isExperiment := flag.IsExperimentationEnabled(detail.Reason)
				valueObj := responseObj.Name(flag.Key).Object()
				detail.Value.WriteToJSONWriter(valueObj.Name("value"))
				detail.VariationIndex.WriteToJSONWriter(valueObj.Name("variation"))
				valueObj.Name("version").Int(flag.Version)
				valueObj.Maybe("trackEvents", flag.TrackEvents || isExperiment).Bool(true)
				valueObj.Maybe("trackReason", isExperiment).Bool(true)
				if withReasons || isExperiment {
					detail.Reason.WriteToJSONWriter(valueObj.Name("reason"))
				}
				valueObj.Maybe("debugEventsUntilDate", flag.DebugEventsUntilDate != 0).
					Float64(float64(flag.DebugEventsUntilDate))
				valueObj.End()
			}
		}
	}
	responseObj.End()
	result := responseWriter.Bytes()

	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(result)
}

func pollFlagOrSegment(clientContext relayenv.EnvContext, kind ldstoretypes.DataKind) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		key := mux.Vars(req)["key"]
		item, err := clientContext.GetStore().Get(kind, key)
		if err != nil {
			clientContext.GetLoggers().Errorf("Error reading feature store: %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if item.Item == nil {
			w.WriteHeader(http.StatusNotFound)
		} else {
			bytes, err := json.Marshal(item.Item)
			if err == nil {
				writeCacheableJSONResponse(w, req, clientContext, bytes, strconv.Itoa(item.Version))
			} else {
				clientContext.GetLoggers().Errorf("Error marshaling JSON: %s", err)
				w.WriteHeader(http.StatusInternalServerError)
			}
		}
	}
}

func writeCacheableJSONResponse(w http.ResponseWriter, req *http.Request, clientContext relayenv.EnvContext,
	bytes []byte, etagValue string) {
	etag := fmt.Sprintf("relay-%s", etagValue) // just to make it extra clear that these are relay-specific etags
	if cachedEtag := req.Header.Get("If-None-Match"); cachedEtag != "" {
		if cachedEtag == etag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
	}
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Etag", etag)
	ttl := clientContext.GetTTL()
	if ttl > 0 {
		w.Header().Set("Vary", "Authorization")
		expiresAt := time.Now().UTC().Add(ttl)
		w.Header().Set("Expires", expiresAt.Format(http.TimeFormat))
		// We're setting "Expires:" instead of "Cache-Control:max-age=" so that if someone puts an
		// HTTP cache in front of ld-relay, multiple clients hitting the cache at different times
		// will all see the same expiration time.
	}
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(bytes)
}

func serializeFlagsAsMap(coll []ldstoretypes.KeyedItemDescriptor) []byte {
	w := jwriter.NewWriter()
	obj := w.Object()
	for _, item := range coll {
		if item.Item.Item != nil {
			ldmodel.MarshalFeatureFlagToJSONWriter(*item.Item.Item.(*ldmodel.FeatureFlag), obj.Name(item.Key))
		}
	}
	obj.End()
	return w.Bytes()
}


<|start_filename|>internal/core/sdks/sdk_kinds.go<|end_filename|>
package sdks

import (
	"errors"
	"net/http"
	"strings"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"

	"github.com/gorilla/mux"
)

var (
	errNoAuthToken    = errors.New("no valid token found")
	errNoEnvID        = errors.New("environment ID not found in URL")
	errUnknownSDKKind = errors.New("unknown SDK kind")
)

// GetCredential attempts to get the appropriate kind of authentication credential for this SDK kind
// from an HTTP request. For Server and Mobile, this uses the Authorization header; for JSClient, it
// is in a path parameter.
func GetCredential(k basictypes.SDKKind, req *http.Request) (config.SDKCredential, error) {
	switch k {
	case basictypes.ServerSDK:
		value, err := fetchAuthToken(req)
		if err == nil {
			return config.SDKKey(value), nil
		}
		return nil, err
	case basictypes.MobileSDK:
		value, err := fetchAuthToken(req)
		if err == nil {
			return config.MobileKey(value), nil
		}
		return nil, err
	case basictypes.JSClientSDK:
		value := mux.Vars(req)["envId"]
		if value == "" {
			return nil, errNoEnvID
		}
		return config.EnvironmentID(value), nil
	}
	return nil, errUnknownSDKKind
}

func fetchAuthToken(req *http.Request) (string, error) {
	authHdr := req.Header.Get("Authorization")
	if strings.HasPrefix(authHdr, "api_key ") {
		authHdr = strings.TrimSpace(strings.TrimPrefix(authHdr, "api_key "))
	}
	if authHdr == "" || strings.Contains(authHdr, " ") {
		return "", errNoAuthToken
	}
	return authHdr, nil
}


<|start_filename|>internal/core/sharedtest/requests.go<|end_filename|>
package sharedtest

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/launchdarkly/eventsource"
	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
	"github.com/launchdarkly/ld-relay/v6/config"

	"github.com/stretchr/testify/require"
)

// BuildRequest is a simple shortcut for creating a request that may or may not have a body.
func BuildRequest(method, url string, body []byte, headers http.Header) *http.Request {
	var bodyBuffer io.Reader
	if body != nil {
		bodyBuffer = bytes.NewBuffer(body)
	}
	r, err := http.NewRequest(method, url, bodyBuffer)
	if err != nil {
		panic(err)
	}
	if headers != nil {
		r.Header = headers
	}
	return r
}

// BuildRequestWithAuth creates a GET request with an Authorization header.
func BuildRequestWithAuth(method, url string, authKey config.SDKCredential, body []byte) *http.Request {
	h := make(http.Header)
	if authKey != nil {
		h.Add("Authorization", authKey.GetAuthorizationHeaderValue())
	}
	return BuildRequest(method, url, body, h)
}

// AddQueryParam is a shortcut for concatenating a query string to a URL that may or may not have one.
func AddQueryParam(url, query string) string {
	if strings.Contains(url, "?") {
		return url + "&" + query
	}
	return url + "?" + query
}

// DoRequest is a shortcut for executing an endpoint handler against a request and getting the response.
func DoRequest(req *http.Request, handler http.Handler) (*http.Response, []byte) {
	w := httptest.NewRecorder()
	handler.ServeHTTP(w, req)
	result := w.Result()
	var body []byte
	if result.Body != nil {
		body, _ = ioutil.ReadAll(result.Body)
		_ = result.Body.Close()
	}

	return result, body
}

// ExpectTestRequest is a shortcut for reading from an httphelpers request-capturing channel with a timeout.
func ExpectTestRequest(t *testing.T, ch <-chan httphelpers.HTTPRequestInfo, timeout time.Duration) httphelpers.HTTPRequestInfo {
	select {
	case r := <-ch:
		return r
	case <-time.After(timeout):
		require.Fail(t, "timed out waiting for request")
		return httphelpers.HTTPRequestInfo{}
	}
}

// ExpectNoTestRequests causes a test failure if an httphelpers request-capturing channel is not empty.
func ExpectNoTestRequests(t *testing.T, ch <-chan httphelpers.HTTPRequestInfo, timeout time.Duration) {
	select {
	case <-ch:
		require.Fail(t, "received unexpected request")
	case <-time.After(timeout):
		break
	}
}

// ExpectStreamEvent is a shortcut for reading from an SSE stream with a timeout.
func ExpectStreamEvent(t *testing.T, stream *eventsource.Stream, timeout time.Duration) eventsource.Event {
	return ExpectStreamChEvent(t, stream.Events, timeout)
}

// ExpectStreamChEvent is a shortcut for reading from an SSE stream channel with a timeout.
func ExpectStreamChEvent(t *testing.T, ch <-chan eventsource.Event, timeout time.Duration) eventsource.Event {
	select {
	case e := <-ch:
		require.NotNil(t, e)
		return e
	case <-time.After(timeout):
		require.Fail(t, "timed out waiting for stream event")
		return nil
	}
}

// ExpectNoStreamEvent causes a test failure if an event is seen on an SSE stream.
func ExpectNoStreamEvent(t *testing.T, stream *eventsource.Stream, timeout time.Duration) {
	ExpectNoStreamChEvent(t, stream.Events, timeout)
}

// ExpectNoStreamChEvent causes a test failure if an event is seen on an SSE stream channel.
func ExpectNoStreamChEvent(t *testing.T, ch <-chan eventsource.Event, timeout time.Duration) {
	select {
	case <-ch:
		require.Fail(t, "received unexpected stream event")
	case <-time.After(timeout):
	}
}

// CallHandlerAndAwaitStatus calls an HTTP handler directly with a request and then blocks
// until the handler has started a response, returning the response status (and cancelling
// the request). We use this when we don't need to wait for a complete response (or when there's
// no such thing as a complete response, as in the case of streaming endpoints). It raises a
// fatal test failure if the timeout elapses before receiving a status.
func CallHandlerAndAwaitStatus(t *testing.T, handler http.Handler, req *http.Request, timeout time.Duration) int {
	var s simpleResponseSink
	go handler.ServeHTTP(&s, req)
	require.True(t, s.awaitResponseStarted(timeout), "timed out waiting for HTTP handler to start response")
	return s.getStatus()
}

// simpleResponseSink is used by CallHandlerAndAwaitStatus instead of httptest.ResponseRecorder
// because ResponseRecorder requires you to wait for a complete response.
type simpleResponseSink struct {
	status    int
	header    http.Header
	body      []byte
	started   bool
	startedCh chan struct{}
	lock      sync.Mutex
}

func (s *simpleResponseSink) awaitResponseStarted(timeout time.Duration) bool {
	s.lock.Lock()
	if s.started {
		s.lock.Unlock()
		return true
	}
	if s.startedCh == nil {
		s.startedCh = make(chan struct{}, 1)
	}
	ch := s.startedCh
	s.lock.Unlock()
	select {
	case <-ch:
		return true
	case <-time.After(timeout):
		return false
	}
}

func (s *simpleResponseSink) getStatus() int {
	s.lock.Lock()
	defer s.lock.Unlock()
	return s.status
}

func (s *simpleResponseSink) Header() http.Header {
	s.lock.Lock()
	defer s.lock.Unlock()
	if s.header == nil {
		s.header = make(http.Header)
	}
	return s.header
}

func (s *simpleResponseSink) Write(data []byte) (int, error) {
	s.start()
	s.lock.Lock()
	defer s.lock.Unlock()
	s.body = append(s.body, data...)
	return len(data), nil
}

func (s *simpleResponseSink) WriteHeader(status int) {
	s.lock.Lock()
	s.status = status
	s.lock.Unlock()
	s.start()
}

func (s *simpleResponseSink) Flush() {}

func (s *simpleResponseSink) start() {
	s.lock.Lock()
	defer s.lock.Unlock()
	if s.started {
		return
	}
	if s.status == 0 {
		s.status = 200
	}
	s.started = true
	if s.startedCh == nil {
		s.startedCh = make(chan struct{}, 1)
	}
	s.startedCh <- struct{}{}
}


<|start_filename|>Makefile<|end_filename|>

GOLANGCI_LINT_VERSION=v1.27.0

LINTER=./bin/golangci-lint
LINTER_VERSION_FILE=./bin/.golangci-lint-version-$(GOLANGCI_LINT_VERSION)

GORELEASER_VERSION=v0.141.0

SHELL=/bin/bash

LINTER=./bin/golangci-lint

TEST_COVERAGE_REPORT_FILE ?= coverage.txt

ALL_SOURCES := $(shell find * -type f -name "*.go")
COVERAGE_PROFILE_RAW=./build/coverage_raw.out
COVERAGE_PROFILE_RAW_HTML=./build/coverage_raw.html
COVERAGE_PROFILE_FILTERED=./build/coverage.out
COVERAGE_PROFILE_FILTERED_HTML=./build/coverage.html
COVERAGE_ENFORCER_FLAGS=\
  	-skipfiles 'internal/core/sharedtest/' \
	-skipcode "// COVERAGE" -packagestats -filestats -showcode

OPTIONAL_TAGS_PARAM=$(if ${TAGS},-tags ${TAGS},)
ALL_TEST_TAGS=big_segment_external_store_tests,integrationtests,redis_unit_tests

build:
	go build .

test:
	go test -run=not-a-real-test -tags $(ALL_TEST_TAGS) ./...  # just ensures that the tests compile
	go test -race -v $(OPTIONAL_TAGS_PARAM) ./...

test-coverage: $(COVERAGE_PROFILE_RAW)
	if [ ! -x "$(GOPATH)/bin/go-coverage-enforcer)" ]; then go get -u github.com/launchdarkly-labs/go-coverage-enforcer; fi
	$(GOPATH)/bin/go-coverage-enforcer $(COVERAGE_ENFORCER_FLAGS) -outprofile $(COVERAGE_PROFILE_FILTERED) $(COVERAGE_PROFILE_RAW) || true
	@# added || true because we don't currently want go-coverage-enforcer to stop the build due to coverage gaps
	go tool cover -html $(COVERAGE_PROFILE_FILTERED) -o $(COVERAGE_PROFILE_FILTERED_HTML)
	go tool cover -html $(COVERAGE_PROFILE_RAW) -o $(COVERAGE_PROFILE_RAW_HTML)

integration-test:
	go test -v -tags integrationtests ./integrationtests

benchmarks: build
	go test -benchmem '-run=^$$' '-bench=.*' ./...

$(COVERAGE_PROFILE_RAW): $(ALL_SOURCES)
	@mkdir -p ./build
	go test -run=not-a-real-test -tags $(ALL_TEST_TAGS) ./...  # just ensures that the tests compile
	go test $(OPTIONAL_TAGS_PARAM) -coverprofile $(COVERAGE_PROFILE_RAW) -coverpkg=./... ./...

$(LINTER_VERSION_FILE):
	rm -f $(LINTER)
	curl -sfL https://install.goreleaser.com/github.com/golangci/golangci-lint.sh | bash -s $(GOLANGCI_LINT_VERSION)
	touch $(LINTER_VERSION_FILE)

lint: $(LINTER_VERSION_FILE)
	$(LINTER) run ./...

# Get the lines added to the most recent changelog update (minus the first 2 lines)
RELEASE_NOTES=<(GIT_EXTERNAL_DIFF='bash -c "diff --unchanged-line-format=\"\" $$2 $$5" || true' git log --ext-diff -1 --pretty= -p CHANGELOG.md)

echo-release-notes:
	@cat $(RELEASE_NOTES)

RELEASE_CMD=curl -sL https://git.io/goreleaser | GOPATH=$(mktemp -d) VERSION=$(GORELEASER_VERSION) bash -s -- --rm-dist --release-notes $(RELEASE_NOTES)

# Note that we're setting GOPATH to a temporary location when running goreleaser, because
# we want it to start from a clean state even if we've previously run a build - and also
# because during a release, we may need to run this command under another account and we
# don't want to mess up file permissions in the regular GOPATH.
publish:
	$(RELEASE_CMD)

products-for-release:
	$(RELEASE_CMD) --skip-publish --skip-validate

DOCKER_COMPOSE_TEST=docker-compose -f docker-compose.test.yml

test-centos test-debian test-docker test-docker-standalone: products-for-release
	$(DOCKER_COMPOSE_TEST) up --force-recreate -d $(subst test,relay,$@)
	trap "$(DOCKER_COMPOSE_TEST) logs && $(DOCKER_COMPOSE_TEST) rm -f" EXIT; $(DOCKER_COMPOSE_TEST) run --rm $@

docker-smoke-test: test-centos test-debian test-docker test-docker-standalone

.PHONY: docker build lint publish products-for-release test test-centos test-debian test-docker test-all test-docker-standalone integration-test benchmarks


<|start_filename|>integrationtests/projects_and_environments_test.go<|end_filename|>
// +build integrationtests

package integrationtests

import (
	"github.com/launchdarkly/ld-relay/v6/config"
)

type projectInfo struct {
	key  string
	name string
}

type environmentInfo struct {
	id        config.EnvironmentID
	key       string
	name      string
	sdkKey    config.SDKKey
	mobileKey config.MobileKey
	prefix    string
}

type projsAndEnvs map[projectInfo][]environmentInfo

func (pe projsAndEnvs) enumerateEnvs(fn func(projectInfo, environmentInfo)) {
	for proj, envs := range pe {
		for _, env := range envs {
			fn(proj, env)
		}
	}
}

func (pe projsAndEnvs) countEnvs() int {
	n := 0
	pe.enumerateEnvs(func(projectInfo, environmentInfo) { n++ })
	return n
}


<|start_filename|>relay/testutils_test.go<|end_filename|>
package relay

import (
	"context"
	"net/http"
	"reflect"
	"testing"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	c "github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/relayenv"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"github.com/launchdarkly/ld-relay/v6/internal/envfactory"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type relayTestHelper struct {
	t     *testing.T
	relay *Relay
}

func (h relayTestHelper) awaitEnvironment(envID c.EnvironmentID) relayenv.EnvContext {
	var e relayenv.EnvContext
	require.Eventually(h.t, func() bool {
		e, _ = h.relay.core.GetEnvironment(envID)
		return e != nil
	}, time.Second, time.Millisecond*5)
	return e
}

func (h relayTestHelper) shouldNotHaveEnvironment(envID c.EnvironmentID, timeout time.Duration) {
	require.Eventually(h.t, func() bool {
		e, _ := h.relay.core.GetEnvironment(envID)
		return e == nil
	}, timeout, time.Millisecond*5)
}

func (h relayTestHelper) assertEnvLookup(env relayenv.EnvContext, expected envfactory.EnvironmentParams) {
	foundEnv, _ := h.relay.core.GetEnvironment(expected.EnvID)
	assert.Equal(h.t, env, foundEnv)
	foundEnv, _ = h.relay.core.GetEnvironment(expected.MobileKey)
	assert.Equal(h.t, env, foundEnv)
	foundEnv, _ = h.relay.core.GetEnvironment(expected.SDKKey)
	assert.Equal(h.t, env, foundEnv)

	h.assertSDKEndpointsAvailability(true, expected.SDKKey, expected.MobileKey, expected.EnvID)
}

func (h relayTestHelper) assertSDKEndpointsAvailability(
	shouldBeAvailable bool,
	sdkKey config.SDKKey,
	mobileKey config.MobileKey,
	envID config.EnvironmentID,
) {
	// Here we're making sure that all of the SDK endpoints are properly recognizing the appropriate
	// authorized credentials (if shouldBeAvailable is true) or rejecting unauthorized credentials (if
	// shouldBeAvailable is false). We're not checking the response body, just that we get the
	// appropriate HTTP status. These are tested more thoroughly in test suites like DoStreamEndpointsTests,
	// but we use this simpler test when we're dynamically changing what the credentials are.

	simpleUserJSON := []byte(`{"key":"userkey"}`)
	simpleUserBase64 := "<KEY>
	status200Or401, status200Or404 := 200, 200
	if !shouldBeAvailable {
		status200Or401 = 401
		status200Or404 = 404
	}
	if sdkKey != "" {
		h.assertEndpointStatus(status200Or401, "GET", "/all", sdkKey, nil)
		h.assertEndpointStatus(status200Or401, "GET", "/flags", sdkKey, nil)
		h.assertEndpointStatus(status200Or401, "REPORT", "/sdk/eval/user", sdkKey, simpleUserJSON)
		h.assertEndpointStatus(status200Or401, "REPORT", "/sdk/evalx/user", sdkKey, simpleUserJSON)
	}
	if mobileKey != "" {
		h.assertEndpointStatus(status200Or401, "GET", "/mping", mobileKey, nil)
		h.assertEndpointStatus(status200Or401, "GET", "/meval/"+simpleUserBase64, mobileKey, nil)
		h.assertEndpointStatus(status200Or401, "REPORT", "/meval", mobileKey, simpleUserJSON)
	}
	if envID != "" {
		h.assertEndpointStatus(status200Or404, "GET", "/ping/"+string(envID), nil, nil)
		h.assertEndpointStatus(status200Or404, "GET", "/eval/"+string(envID)+"/"+simpleUserBase64, nil, nil)
		h.assertEndpointStatus(status200Or404, "REPORT", "/eval/"+string(envID), nil, simpleUserJSON)
	}
}

func (h relayTestHelper) assertEndpointStatus(
	expectedStatus int,
	method, path string,
	authKey config.SDKCredential,
	body []byte,
) {
	ctx, cancelFn := context.WithCancel(context.Background())
	defer cancelFn()
	headers := make(http.Header)
	var authValue string
	if authKey != nil {
		authValue = authKey.GetAuthorizationHeaderValue()
	}
	if authValue != "" {
		headers.Add("Authorization", authValue)
	}
	if body != nil {
		headers.Add("Content-Type", "application/json")
	}
	req := sharedtest.BuildRequest(method, path, body, headers).WithContext(ctx)
	status := sharedtest.CallHandlerAndAwaitStatus(h.t, h.relay, req, time.Second)
	require.Equal(h.t, expectedStatus, status, "expected status %d but got %d from %s %s with auth key %s",
		expectedStatus, status, method, path, authValue)
}

func (h relayTestHelper) awaitCredentialsUpdated(env relayenv.EnvContext, expected envfactory.EnvironmentParams) {
	expectedCredentials := credentialsAsSet(expected.EnvID, expected.MobileKey, expected.SDKKey)
	isChanged := func() bool {
		return reflect.DeepEqual(credentialsAsSet(env.GetCredentials()...), expectedCredentials)
	}
	require.Eventually(h.t, isChanged, time.Second, time.Millisecond*5)
	h.assertEnvLookup(env, expected)
}

func assertEnvProps(t *testing.T, expected envfactory.EnvironmentParams, env relayenv.EnvContext) {
	assert.Equal(t, credentialsAsSet(expected.EnvID, expected.MobileKey, expected.SDKKey),
		credentialsAsSet(env.GetCredentials()...))
	assert.Equal(t, expected.Identifiers, env.GetIdentifiers())
	assert.Equal(t, expected.Identifiers.ProjName+" "+expected.Identifiers.EnvName,
		env.GetIdentifiers().GetDisplayName())
}

func credentialsAsSet(cs ...c.SDKCredential) map[c.SDKCredential]struct{} {
	ret := make(map[c.SDKCredential]struct{}, len(cs))
	for _, c := range cs {
		ret[c] = struct{}{}
	}
	return ret
}


<|start_filename|>internal/core/relay_core_test.go<|end_filename|>
package core

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	c "github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/relayenv"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sdks"
	st "github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest/testclient"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest/testenv"

	"github.com/launchdarkly/eventsource"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	ld "gopkg.in/launchdarkly/go-server-sdk.v5"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func makeBasicCore(config c.Config) (*RelayCore, error) {
	return NewRelayCore(config, ldlog.NewDisabledLoggers(), testclient.FakeLDClientFactory(true), "", "", false)
}

func TestNewRelayCoreRejectsConfigWithContradictoryProperties(t *testing.T) {
	// it is an error to enable TLS but not provide a cert or key
	config := c.Config{Main: c.MainConfig{TLSEnabled: true}}
	core, err := makeBasicCore(config)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "TLS cert")
	assert.Nil(t, core)
}

func TestRelayCoreGetEnvironment(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile, st.EnvClientSide),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	env, inited := core.GetEnvironment(st.EnvMain.Config.SDKKey)
	require.NotNil(t, env)
	assert.True(t, inited)
	assert.Equal(t, st.EnvMain.Name, env.GetIdentifiers().ConfiguredName)

	env, inited = core.GetEnvironment(st.EnvMobile.Config.SDKKey)
	require.NotNil(t, env)
	assert.True(t, inited)
	assert.Equal(t, st.EnvMobile.Name, env.GetIdentifiers().ConfiguredName)

	env, inited = core.GetEnvironment(st.EnvClientSide.Config.SDKKey)
	require.NotNil(t, env)
	assert.True(t, inited)
	assert.Equal(t, st.EnvClientSide.Name, env.GetIdentifiers().ConfiguredName)

	env, inited = core.GetEnvironment(st.EnvMobile.Config.MobileKey)
	require.NotNil(t, env)
	assert.True(t, inited)
	assert.Equal(t, st.EnvMobile.Name, env.GetIdentifiers().ConfiguredName)

	env, inited = core.GetEnvironment(st.EnvClientSide.Config.EnvID)
	require.NotNil(t, env)
	assert.True(t, inited)
	assert.Equal(t, st.EnvClientSide.Name, env.GetIdentifiers().ConfiguredName)

	env, inited = core.GetEnvironment(st.UndefinedSDKKey)
	assert.Nil(t, env)
	assert.True(t, inited)

	env, inited = core.GetEnvironment(st.UnsupportedSDKCredential{})
	assert.Nil(t, env)
	assert.True(t, inited)
}

func TestRelayCoreGetAllEnvironments(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile, st.EnvClientSide),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	envs := core.GetAllEnvironments()
	assert.Len(t, envs, 3)
	var names []string
	for _, e := range envs {
		names = append(names, e.GetIdentifiers().ConfiguredName)
	}
	assert.Contains(t, names, st.EnvMain.Name)
	assert.Contains(t, names, st.EnvMobile.Name)
	assert.Contains(t, names, st.EnvClientSide.Name)
}

func TestRelayCoreAddEnvironment(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	env, resultCh, err := core.AddEnvironment(relayenv.EnvIdentifiers{ConfiguredName: st.EnvMobile.Name}, st.EnvMobile.Config, nil)
	require.NoError(t, err)
	require.NotNil(t, env)
	require.NotNil(t, resultCh)
	assert.Equal(t, st.EnvMobile.Name, env.GetIdentifiers().ConfiguredName)

	env1, _ := core.GetEnvironment(st.EnvMobile.Config.SDKKey)
	assert.Equal(t, env, env1)

	select {
	case env2 := <-resultCh:
		assert.Equal(t, env, env2)
	case <-time.After(time.Second):
		assert.Fail(t, "timed out waiting for new environment to initialize")
	}
}

func TestRelayCoreAddEnvironmentAfterClosed(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	core.Close()

	env, resultCh, err := core.AddEnvironment(relayenv.EnvIdentifiers{ConfiguredName: st.EnvMobile.Name}, st.EnvMobile.Config, nil)
	assert.Error(t, err)
	assert.Nil(t, env)
	assert.Nil(t, resultCh)
}

func TestRelayCoreRemoveEnvironment(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	env, _ := core.GetEnvironment(st.EnvMobile.Config.SDKKey)
	require.NotNil(t, env)
	assert.Equal(t, st.EnvMobile.Name, env.GetIdentifiers().ConfiguredName)

	removed := core.RemoveEnvironment(env)
	assert.True(t, removed)

	noEnv, _ := core.GetEnvironment(st.EnvMobile.Config.SDKKey)
	assert.Nil(t, noEnv)
}

func TestRelayCoreRemoveUnknownEnvironment(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	env := testenv.NewTestEnvContext("unknown", true, st.NewInMemoryStore())

	assert.False(t, core.RemoveEnvironment(env))
}

func TestRelayCoreAddedEnvironmentCredential(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	env, _ := core.GetEnvironment(st.EnvMain.Config.SDKKey)
	require.NotNil(t, env)
	assert.Equal(t, st.EnvMain.Name, env.GetIdentifiers().ConfiguredName)

	extraKey := c.SDKKey(string(st.EnvMain.Config.SDKKey) + "-extra")
	noEnv, _ := core.GetEnvironment(extraKey)
	assert.Nil(t, noEnv)

	core.AddedEnvironmentCredential(env, extraKey)

	env1, _ := core.GetEnvironment(extraKey)
	assert.Equal(t, env, env1)
}

func TestRelayCoreRemovingEnvironmentCredential(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile),
	}
	core, err := makeBasicCore(config)
	require.NoError(t, err)
	defer core.Close()

	core.RemovingEnvironmentCredential(st.EnvMain.Config.SDKKey)

	noEnv, _ := core.GetEnvironment(st.EnvMain.Config.SDKKey)
	assert.Nil(t, noEnv)

	env, _ := core.GetEnvironment(st.EnvMobile.Config.SDKKey)
	require.NotNil(t, env)
	assert.Equal(t, st.EnvMobile.Name, env.GetIdentifiers().ConfiguredName)

	assert.Len(t, core.GetAllEnvironments(), 2) // EnvMain is not removed from this list
}

func TestRelayCoreWaitForAllEnvironments(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile),
	}

	t.Run("returns nil if all environments initialize successfully", func(t *testing.T) {
		core, err := NewRelayCore(config, ldlog.NewDisabledLoggers(), testclient.FakeLDClientFactory(true), "", "", false)
		require.NoError(t, err)
		defer core.Close()

		err = core.WaitForAllClients(time.Second)
		assert.NoError(t, err)
	})

	t.Run("returns error if any environment does not initialize successfully", func(t *testing.T) {
		core, err := NewRelayCore(config, ldlog.NewDisabledLoggers(),
			oneEnvFails(st.EnvMobile.Config.SDKKey, false, nil), "", "", false)
		require.NoError(t, err)
		defer core.Close()

		err = core.WaitForAllClients(time.Second)
		assert.Error(t, err)
	})
}

func TestRelayCoreUninitializedEnvironment(t *testing.T) {
	config := c.Config{
		Environment: st.MakeEnvConfigs(st.EnvMain, st.EnvMobile),
	}
	problemEnv := st.EnvMobile

	t.Run("handlers return 503 for environment that is still initializing", func(t *testing.T) {
		gateCh := make(chan struct{})
		defer close(gateCh)

		core, err := NewRelayCore(config, ldlog.NewDisabledLoggers(),
			oneEnvFails(problemEnv.Config.SDKKey, true, gateCh), "", "", false)
		require.NoError(t, err)
		defer core.Close()
		router := core.MakeRouter()

		req1 := st.MakeSDKStreamEndpointRequest("", basictypes.ServerSideStream, problemEnv, st.SimpleUserJSON, 0)
		rr1 := httptest.NewRecorder()
		router.ServeHTTP(rr1, req1)
		assert.Equal(t, http.StatusServiceUnavailable, rr1.Result().StatusCode)

		req2 := st.MakeSDKStreamEndpointRequest("", basictypes.MobilePingStream, problemEnv, st.SimpleUserJSON, 0)
		rr2 := httptest.NewRecorder()
		router.ServeHTTP(rr2, req2)
		assert.Equal(t, http.StatusServiceUnavailable, rr2.Result().StatusCode)
	})

	t.Run("handlers accept requests for environment that failed to initialize", func(t *testing.T) {
		core, err := NewRelayCore(config, ldlog.NewDisabledLoggers(),
			oneEnvFails(problemEnv.Config.SDKKey, true, nil), "", "", false)
		require.NoError(t, err)
		defer core.Close()
		router := core.MakeRouter()

		err = core.WaitForAllClients(time.Millisecond * 100)
		assert.Error(t, err)

		env, _ := core.GetEnvironment(problemEnv.Config.SDKKey)
		assert.NotNil(t, env)
		store := env.GetStore()
		assert.NotNil(t, store)
		store.Init(nil)

		req1 := st.MakeSDKStreamEndpointRequest("", basictypes.ServerSideStream, problemEnv, "", 0)
		resp1 := st.WithStreamRequest(t, req1, router, func(ch <-chan eventsource.Event) { <-ch })
		assert.Equal(t, http.StatusOK, resp1.StatusCode)

		req2 := st.MakeSDKStreamEndpointRequest("", basictypes.MobilePingStream, problemEnv, st.SimpleUserJSON, 0)
		resp2 := st.WithStreamRequest(t, req2, router, func(ch <-chan eventsource.Event) { <-ch })
		assert.Equal(t, http.StatusOK, resp2.StatusCode)
	})
}

func oneEnvFails(
	badSDKKey c.SDKKey,
	returnClientInstanceAnyway bool,
	gateCh <-chan struct{},
) sdks.ClientFactoryFunc {
	return func(sdkKey c.SDKKey, config ld.Config, timeout time.Duration) (sdks.LDClientContext, error) {
		client, _ := testclient.FakeLDClientFactory(true)(sdkKey, config, timeout)
		if sdkKey == badSDKKey {
			if gateCh != nil {
				<-gateCh
			}
			err := errors.New("sorry")
			if returnClientInstanceAnyway {
				return client, err
			}
			return nil, err
		}
		return client, nil
	}
}


<|start_filename|>internal/core/sharedtest/big_segments_fixtures.go<|end_filename|>
package sharedtest

import "gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"

// NoOpSDKBigSegmentStore is a stub implementation of the SDK's BigSegmentStore (not the type
// of the same name that Relay uses internally).
type NoOpSDKBigSegmentStore struct{}

func (m *NoOpSDKBigSegmentStore) Close() error {
	return nil
}

func (m *NoOpSDKBigSegmentStore) GetMetadata() (interfaces.BigSegmentStoreMetadata, error) {
	return interfaces.BigSegmentStoreMetadata{}, nil
}

func (m *NoOpSDKBigSegmentStore) GetUserMembership(
	userHash string,
) (interfaces.BigSegmentMembership, error) {
	return nil, nil
}


<|start_filename|>internal/core/internal/events/event-relay.go<|end_filename|>
package events

import (
	"encoding/json"
	"io/ioutil"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"sync"

	c "github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/store"
	"github.com/launchdarkly/ld-relay/v6/internal/util"

	"github.com/launchdarkly/go-configtypes"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	ldevents "gopkg.in/launchdarkly/go-sdk-events.v1"
)

type eventVerbatimRelay struct {
	config    c.EventsConfig
	publisher EventPublisher
}

const (
	// SummaryEventsSchemaVersion is the minimum event schema that supports summary events
	SummaryEventsSchemaVersion = 3

	// EventSchemaHeader is an HTTP header that describes the schema version for event requests
	EventSchemaHeader = "X-LaunchDarkly-Event-Schema"
)

// EventDispatcher relays events to LaunchDarkly for an environment
type EventDispatcher struct {
	analyticsEndpoints  map[basictypes.SDKKind]*analyticsEventEndpointDispatcher
	diagnosticEndpoints map[basictypes.SDKKind]*diagnosticEventEndpointDispatcher
}

type analyticsEventEndpointDispatcher struct {
	config           c.EventsConfig
	httpClient       *http.Client
	httpConfig       httpconfig.HTTPConfig
	authKey          c.SDKCredential
	remotePath       string
	verbatimRelay    *eventVerbatimRelay
	summarizingRelay *eventSummarizingRelay
	storeAdapter     *store.SSERelayDataStoreAdapter
	loggers          ldlog.Loggers
	mu               sync.Mutex
}

type diagnosticEventEndpointDispatcher struct {
	httpClient        *http.Client
	httpConfig        httpconfig.HTTPConfig
	remoteEndpointURI string
	loggers           ldlog.Loggers
}

// GetHandler returns the HTTP handler for an endpoint, or nil if none is defined
func (r *EventDispatcher) GetHandler(sdkKind basictypes.SDKKind, eventsKind ldevents.EventDataKind) func(w http.ResponseWriter, req *http.Request) {
	if eventsKind == ldevents.DiagnosticEventDataKind {
		if e, ok := r.diagnosticEndpoints[sdkKind]; ok {
			return e.dispatch
		}
	} else {
		if e, ok := r.analyticsEndpoints[sdkKind]; ok {
			return e.dispatch
		}
	}
	return nil
}

func (r *analyticsEventEndpointDispatcher) dispatch(w http.ResponseWriter, req *http.Request) {
	consumeEvents(w, req, r.loggers, func(body []byte) {
		evts := make([]json.RawMessage, 0)
		err := json.Unmarshal(body, &evts)
		if err != nil {
			r.loggers.Errorf("Error unmarshaling event post body: %+v", err)
			return
		}

		payloadVersion, _ := strconv.Atoi(req.Header.Get(EventSchemaHeader))
		if payloadVersion == 0 {
			payloadVersion = 1
		}
		r.loggers.Debugf("Received %d events (v%d) to be proxied to %s", len(evts), payloadVersion, r.remotePath)
		if payloadVersion >= SummaryEventsSchemaVersion {
			// New-style events that have already gone through summarization - deliver them as-is
			r.getVerbatimRelay().enqueue(evts)
		} else {
			r.getSummarizingRelay().enqueue(evts, payloadVersion)
		}
	})
}

func (r *analyticsEventEndpointDispatcher) replaceCredential(newCredential c.SDKCredential) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if reflect.TypeOf(r.authKey) == reflect.TypeOf(newCredential) {
		r.authKey = newCredential
		if r.summarizingRelay != nil {
			r.summarizingRelay.eventSender.replaceCredential(newCredential)
		}
		if r.verbatimRelay != nil {
			r.verbatimRelay.publisher.ReplaceCredential(newCredential)
		}
	}
}

func (r *analyticsEventEndpointDispatcher) close() {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.summarizingRelay != nil {
		r.summarizingRelay.close()
	}
	if r.verbatimRelay != nil {
		r.verbatimRelay.close()
	}
}

func (d *diagnosticEventEndpointDispatcher) dispatch(w http.ResponseWriter, req *http.Request) {
	consumeEvents(w, req, d.loggers, func(body []byte) {
		// We are just operating as a reverse proxy and passing the request on verbatim to LD; we do not
		// need to parse the JSON.
		d.loggers.Debugf("Received diagnostic event to be proxied to %s", d.remoteEndpointURI)

		// We use the default EventSender from ldevents, which provides the standard retry logic and logging.
		// Since we don't want to use a fixed set of headers, but instead pass along the same headers we got
		// from the request, we're creating a new EventSender each time; that's a little inefficient, but
		// diagnostic events are relatively infrequent.
		sender := ldevents.NewDefaultEventSender(d.httpClient, "", d.remoteEndpointURI, req.Header, d.loggers)
		_ = sender.SendEventData(ldevents.DiagnosticEventDataKind, body, 1)
	})
}

func consumeEvents(w http.ResponseWriter, req *http.Request, loggers ldlog.Loggers, thenExecute func([]byte)) {
	body, bodyErr := ioutil.ReadAll(req.Body)

	if bodyErr != nil { // COVERAGE: can't make this happen in unit tests
		loggers.Errorf("Error reading event post body: %+v", bodyErr)
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write(util.ErrorJSONMsg("unable to read request body"))
		return
	}

	if len(body) == 0 {
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write(util.ErrorJSONMsg("body may not be empty"))
		return
	}

	// Always accept the data
	w.WriteHeader(http.StatusAccepted)

	defer func() {
		if err := recover(); err != nil { // COVERAGE: can't make this happen in unit tests
			loggers.Errorf("Unexpected panic in event relay: %+v", err)
		}
	}()
	thenExecute(body)
}

func (r *analyticsEventEndpointDispatcher) getVerbatimRelay() *eventVerbatimRelay {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.verbatimRelay == nil {
		r.verbatimRelay = newEventVerbatimRelay(r.authKey, r.config, r.httpConfig, r.loggers, r.remotePath)
	}
	return r.verbatimRelay
}

func (r *analyticsEventEndpointDispatcher) getSummarizingRelay() *eventSummarizingRelay {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.summarizingRelay == nil {
		r.summarizingRelay = newEventSummarizingRelay(r.config, r.httpConfig, r.authKey, r.storeAdapter, r.loggers, r.remotePath)
	}
	return r.summarizingRelay
}

func (r *analyticsEventEndpointDispatcher) flush() { //nolint:unused // used only in tests
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.verbatimRelay != nil {
		r.verbatimRelay.publisher.Flush()
	}
	if r.summarizingRelay != nil {
		r.summarizingRelay.eventProcessor.Flush()
	}
}

// NewEventDispatcher creates a handler for relaying events to LaunchDarkly for an environment
func NewEventDispatcher(
	sdkKey c.SDKKey,
	mobileKey c.MobileKey,
	envID c.EnvironmentID,
	loggers ldlog.Loggers,
	config c.EventsConfig,
	httpConfig httpconfig.HTTPConfig,
	storeAdapter *store.SSERelayDataStoreAdapter,
) *EventDispatcher {
	ep := &EventDispatcher{
		analyticsEndpoints: map[basictypes.SDKKind]*analyticsEventEndpointDispatcher{
			basictypes.ServerSDK: newAnalyticsEventEndpointDispatcher(sdkKey,
				config, httpConfig, storeAdapter, loggers, "/bulk"),
		},
		diagnosticEndpoints: map[basictypes.SDKKind]*diagnosticEventEndpointDispatcher{
			basictypes.ServerSDK: newDiagnosticEventEndpointDispatcher(config, httpConfig, loggers, "/diagnostic"),
		},
	}
	if mobileKey != "" {
		ep.analyticsEndpoints[basictypes.MobileSDK] = newAnalyticsEventEndpointDispatcher(mobileKey,
			config, httpConfig, storeAdapter, loggers, "/mobile")
		ep.diagnosticEndpoints[basictypes.MobileSDK] = newDiagnosticEventEndpointDispatcher(config, httpConfig, loggers, "/mobile/events/diagnostic")
	}
	if envID != "" {
		ep.analyticsEndpoints[basictypes.JSClientSDK] = newAnalyticsEventEndpointDispatcher(envID, config, httpConfig, storeAdapter, loggers,
			"/events/bulk/"+string(envID))
		ep.diagnosticEndpoints[basictypes.JSClientSDK] = newDiagnosticEventEndpointDispatcher(config, httpConfig, loggers,
			"/events/diagnostic/"+string(envID))
	}
	return ep
}

// Close shuts down any goroutines/channels being used by the EventDispatcher.
func (r *EventDispatcher) Close() {
	for _, e := range r.analyticsEndpoints {
		e.close()
	}
	// diagnosticEventEndpointDispatcher doesn't currently need to be closed, because it doesn't maintain any
	// goroutines or channels
}

func (r *EventDispatcher) flush() { //nolint:unused // used only in tests
	for _, e := range r.analyticsEndpoints {
		e.flush()
	}
}

// ReplaceCredential changes the authorization credentail that is used when forwarding events to any
// endpoints that use that type of credential. For instance, if newCredential is a MobileKey, this
// affects only endpoints that use a mobile key.
func (r *EventDispatcher) ReplaceCredential(newCredential c.SDKCredential) {
	for _, d := range r.analyticsEndpoints {
		d.replaceCredential(newCredential)
	}
	// diagnosticEventEndpointDispatcher doesn't need to be updated for this, because it always uses whatever
	// credential was present on the incoming request
}

func newDiagnosticEventEndpointDispatcher(
	config c.EventsConfig,
	httpConfig httpconfig.HTTPConfig,
	loggers ldlog.Loggers,
	remotePath string,
) *diagnosticEventEndpointDispatcher {
	eventsURI := getEventsURI(config)
	return &diagnosticEventEndpointDispatcher{
		httpClient:        httpConfig.Client(),
		httpConfig:        httpConfig,
		remoteEndpointURI: strings.TrimRight(eventsURI, "/") + remotePath,
		loggers:           loggers,
	}
}

func newAnalyticsEventEndpointDispatcher(
	authKey c.SDKCredential,
	config c.EventsConfig,
	httpConfig httpconfig.HTTPConfig,
	storeAdapter *store.SSERelayDataStoreAdapter,
	loggers ldlog.Loggers,
	remotePath string,
) *analyticsEventEndpointDispatcher {
	return &analyticsEventEndpointDispatcher{
		authKey:      authKey,
		config:       config,
		httpClient:   httpConfig.Client(),
		httpConfig:   httpConfig,
		storeAdapter: storeAdapter,
		loggers:      loggers,
		remotePath:   remotePath,
	}
}

func newEventVerbatimRelay(
	authKey c.SDKCredential,
	config c.EventsConfig,
	httpConfig httpconfig.HTTPConfig,
	loggers ldlog.Loggers,
	remotePath string,
) *eventVerbatimRelay {
	eventsURI := getEventsURI(config)
	opts := []OptionType{
		OptionCapacity(config.Capacity.GetOrElse(c.DefaultEventCapacity)),
		OptionEndpointURI(strings.TrimRight(eventsURI, "/") + remotePath),
	}

	opts = append(opts, OptionFlushInterval(config.FlushInterval.GetOrElse(c.DefaultEventsFlushInterval)))

	publisher, _ := NewHTTPEventPublisher(authKey, httpConfig, loggers, opts...)

	res := &eventVerbatimRelay{
		config:    config,
		publisher: publisher,
	}

	return res
}

func (er *eventVerbatimRelay) enqueue(evts []json.RawMessage) {
	er.publisher.Publish(evts...)
}

func (er *eventVerbatimRelay) close() {
	er.publisher.Close()
}

func getEventsURI(config c.EventsConfig) string {
	return configtypes.NewOptStringNonEmpty(config.EventsURI.String()).GetOrElse(c.DefaultEventsURI)
}


<|start_filename|>internal/core/relay_core_endpoints_benchmark_test.go<|end_filename|>
package core

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest/testenv"

	"gopkg.in/launchdarkly/go-sdk-common.v2/lduser"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldvalue"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldbuilders"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"
)

func BenchmarkEvaluateAllFlags(b *testing.B) {
	numFlags := 50

	allData := []ldstoretypes.Collection{
		{
			Kind:  ldstoreimpl.Features(),
			Items: []ldstoretypes.KeyedItemDescriptor{},
		},
		{
			Kind:  ldstoreimpl.Segments(),
			Items: []ldstoretypes.KeyedItemDescriptor{},
		},
	}

	for i := 0; i < numFlags; i++ {
		flag := ldbuilders.NewFlagBuilder(fmt.Sprintf("flag%d", i)).Version(1).
			SingleVariation(ldvalue.String(fmt.Sprintf("value%d", i))).
			ClientSideUsingEnvironmentID(true).
			Build()
		allData[0].Items = append(allData[0].Items, ldstoretypes.KeyedItemDescriptor{flag.Key, sharedtest.FlagDesc(flag)})
	}

	user := lduser.NewUserBuilder("user-key").Name("name").Email("email").Custom("a", ldvalue.String("b")).Build()

	store := sharedtest.NewInMemoryStore()
	store.Init(allData)
	ctx := testenv.NewTestEnvContext("", false, store)
	userData := []byte(user.String())
	headers := make(http.Header)
	headers.Set("Content-Type", "application/json")

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		req := buildPreRoutedRequest("REPORT", userData, headers, nil, ctx)
		resp := httptest.NewRecorder()
		evaluateAllFeatureFlags(basictypes.JSClientSDK)(resp, req)
	}
}


<|start_filename|>internal/core/bigsegments/testdata_test.go<|end_filename|>
package bigsegments

import (
	"encoding/json"

	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
)

const testSDKKey = "sdk-key"
const testEnvironmentID = "abc"

func makePatchEvent(patches ...bigSegmentPatch) *httphelpers.SSEEvent {
	bytes, err := json.Marshal(patches)
	if err != nil {
		panic(err)
	}
	return &httphelpers.SSEEvent{Data: string(bytes)}
}

type patchBuilder struct {
	patch bigSegmentPatch
}

func newPatchBuilder(segmentID, version, previous string) *patchBuilder {
	return &patchBuilder{
		patch: bigSegmentPatch{
			EnvironmentID:   testEnvironmentID,
			SegmentID:       segmentID,
			Version:         version,
			PreviousVersion: previous,
		},
	}
}

func (b *patchBuilder) addIncludes(keys ...string) *patchBuilder {
	b.patch.Changes.Included.Add = append(b.patch.Changes.Included.Add, keys...)
	return b
}

func (b *patchBuilder) addExcludes(keys ...string) *patchBuilder {
	b.patch.Changes.Excluded.Add = append(b.patch.Changes.Excluded.Add, keys...)
	return b
}

func (b *patchBuilder) removeIncludes(keys ...string) *patchBuilder {
	b.patch.Changes.Included.Remove = append(b.patch.Changes.Included.Remove, keys...)
	return b
}

func (b *patchBuilder) removeExcludes(keys ...string) *patchBuilder {
	b.patch.Changes.Excluded.Remove = append(b.patch.Changes.Excluded.Remove, keys...)
	return b
}

func (b *patchBuilder) build() bigSegmentPatch {
	return b.patch
}


<|start_filename|>internal/core/client-side.go<|end_filename|>
package core

import (
	"bytes"
	"encoding/base64"
	"net/http"
	"net/http/httptest"
	"strconv"

	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/browser"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/events"
	"github.com/launchdarkly/ld-relay/v6/internal/core/middleware"
	"github.com/launchdarkly/ld-relay/v6/internal/util"

	ldevents "gopkg.in/launchdarkly/go-sdk-events.v1"
)

func getEventsImage(w http.ResponseWriter, req *http.Request) {
	clientCtx := middleware.GetEnvContextInfo(req.Context())

	if clientCtx.Env.GetEventDispatcher() == nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = w.Write(util.ErrorJSONMsg("Event proxy is not enabled for this environment"))
		return
	}
	handler := clientCtx.Env.GetEventDispatcher().GetHandler(basictypes.JSClientSDK, ldevents.AnalyticsEventDataKind)
	if handler == nil { // COVERAGE: abnormal condition that can't be caused in unit tests
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = w.Write(util.ErrorJSONMsg("Event proxy for browser clients is not enabled for this environment"))
		return
	}

	d := req.URL.Query().Get("d")
	if d != "" {
		go func() {
			nullW := httptest.NewRecorder()
			eventData, _ := base64.StdEncoding.DecodeString(d)
			eventsReq, _ := http.NewRequest("POST", "", bytes.NewBuffer(eventData))
			eventsReq.Header.Add("Content-Type", "application/json")
			eventsReq.Header.Add("X-LaunchDarkly-User-Agent", eventsReq.Header.Get("X-LaunchDarkly-User-Agent"))
			eventsReq.Header.Add(events.EventSchemaHeader, strconv.Itoa(events.SummaryEventsSchemaVersion))
			handler(nullW, eventsReq)
		}()
	}

	w.Header().Set("Content-Type", "image/gif")
	_, _ = w.Write(browser.Transparent1PixelImageData)
}

func getGoals(w http.ResponseWriter, req *http.Request) {
	clientCtx := middleware.GetEnvContextInfo(req.Context())
	clientCtx.Env.GetJSClientContext().Proxy.ServeHTTP(w, req)
}


<|start_filename|>internal/core/internal/events/event-relay_test.go<|end_filename|>
package events

import (
	"net/http"
	"net/http/httptest"
	"strconv"
	"testing"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"
	st "github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"

	"github.com/launchdarkly/go-configtypes"
	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
	ldevents "gopkg.in/launchdarkly/go-sdk-events.v1"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	eventPayloadForVerbatimOnly = `["fake-event-1","fake-event-2","fake-event-3"]`
	// when events are proxied verbatim, we don't care about the schema, it just has to be valid JSON
)

type testEndpointInfo struct {
	sdkKind        basictypes.SDKKind
	analyticsPath  string
	diagnosticPath string
	authKey        string
	newCredential  config.SDKCredential
}

var testServerEndpointInfo = testEndpointInfo{
	sdkKind:        basictypes.ServerSDK,
	analyticsPath:  "/bulk",
	diagnosticPath: "/diagnostic",
	authKey:        string(st.EnvWithAllCredentials.Config.SDKKey),
	newCredential:  config.SDKKey(string(st.EnvWithAllCredentials.Config.SDKKey) + "-new"),
}

var testMobileEndpointInfo = testEndpointInfo{
	sdkKind:        basictypes.MobileSDK,
	analyticsPath:  "/mobile",
	diagnosticPath: "/mobile/events/diagnostic",
	authKey:        string(st.EnvWithAllCredentials.Config.MobileKey),
	newCredential:  config.MobileKey(string(st.EnvWithAllCredentials.Config.MobileKey) + "-new"),
}

var testJSClientEndpointInfo = testEndpointInfo{
	sdkKind:        basictypes.JSClientSDK,
	analyticsPath:  "/events/bulk/" + string(st.EnvWithAllCredentials.Config.EnvID),
	diagnosticPath: "/events/diagnostic/" + string(st.EnvWithAllCredentials.Config.EnvID),
}
var allTestEndpoints = []testEndpointInfo{testServerEndpointInfo, testMobileEndpointInfo, testJSClientEndpointInfo}

type eventRelayTestParams struct {
	dispatcher *EventDispatcher
	requestsCh <-chan httphelpers.HTTPRequestInfo
	dataStore  interfaces.DataStore
	mockLog    *ldlogtest.MockLog
}

func eventRelayTest(
	t *testing.T,
	testEnv st.TestEnv,
	eventsConfig config.EventsConfig,
	fn func(eventRelayTestParams),
) {
	mockLog := ldlogtest.NewMockLog()
	mockLog.Loggers.SetMinLevel(ldlog.Debug)
	defer mockLog.DumpIfTestFailed(t)

	httpConfig, _ := httpconfig.NewHTTPConfig(config.ProxyConfig{}, nil, "", mockLog.Loggers)

	store := st.NewInMemoryStore()

	handler, requestsCh := httphelpers.RecordingHandler(httphelpers.HandlerWithStatus(202))
	httphelpers.WithServer(handler, func(server *httptest.Server) {
		eventsConfig.SendEvents = true
		if !eventsConfig.FlushInterval.IsDefined() {
			// by default, set a long flush interval so tests can control when a flush will happen
			eventsConfig.FlushInterval = configtypes.NewOptDuration(time.Hour)
		}
		eventsConfig.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(server.URL)

		dispatcher := NewEventDispatcher(
			testEnv.Config.SDKKey,
			testEnv.Config.MobileKey,
			testEnv.Config.EnvID,
			mockLog.Loggers,
			eventsConfig,
			httpConfig,
			makeStoreAdapterWithExistingStore(store),
		)
		defer dispatcher.Close()

		p := eventRelayTestParams{
			dispatcher: dispatcher,
			requestsCh: requestsCh,
			dataStore:  store,
			mockLog:    mockLog,
		}
		fn(p)
	})
}

func TestEventDispatcherCreatesHandlersOnlyForConfiguredCredentials(t *testing.T) {
	eventRelayTest(t, st.EnvMain, config.EventsConfig{}, func(p eventRelayTestParams) {
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.DiagnosticEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.AnalyticsEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.DiagnosticEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.AnalyticsEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.DiagnosticEventDataKind))
	})
	eventRelayTest(t, st.EnvMobile, config.EventsConfig{}, func(p eventRelayTestParams) {
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.DiagnosticEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.DiagnosticEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.AnalyticsEventDataKind))
		assert.Nil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.DiagnosticEventDataKind))
	})
	eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.ServerSDK, ldevents.DiagnosticEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.MobileSDK, ldevents.DiagnosticEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.AnalyticsEventDataKind))
		assert.NotNil(t, p.dispatcher.GetHandler(basictypes.JSClientSDK, ldevents.DiagnosticEventDataKind))
	})
}

func TestVerbatimEventHandlers(t *testing.T) {
	for _, e := range allTestEndpoints {
		t.Run(string(e.sdkKind), func(t *testing.T) {
			eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
				req := st.BuildRequest("POST", "/", []byte(eventPayloadForVerbatimOnly),
					headersWithEventSchema(SummaryEventsSchemaVersion))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				p.dispatcher.flush()

				r := st.ExpectTestRequest(t, p.requestsCh, time.Second)
				assert.Equal(t, "POST", r.Request.Method)
				assert.Equal(t, e.analyticsPath, r.Request.URL.Path)
				assert.Equal(t, e.authKey, r.Request.Header.Get("Authorization"))
				assert.Equal(t, strconv.Itoa(SummaryEventsSchemaVersion), r.Request.Header.Get(EventSchemaHeader))
				assert.Equal(t, eventPayloadForVerbatimOnly, string(r.Body))
			})
		})
	}
}

func TestSummarizingEventHandlers(t *testing.T) {
	// The summarizing relay logic is tested in more detail in summarizing-relay_test.go. The test here
	// just verifies that we are indeed using the summarizing relay for these endpoints.
	for _, e := range allTestEndpoints {
		t.Run(string(e.sdkKind), func(t *testing.T) {
			eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
				req := st.BuildRequest("POST", "/", []byte(summarizableFeatureEvents), headersWithEventSchema(0))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				p.dispatcher.flush()

				r := st.ExpectTestRequest(t, p.requestsCh, time.Second)
				assert.Equal(t, "POST", r.Request.Method)
				assert.Equal(t, e.analyticsPath, r.Request.URL.Path)
				assert.Equal(t, e.authKey, r.Request.Header.Get("Authorization"))
				assert.Equal(t, strconv.Itoa(SummaryEventsSchemaVersion), r.Request.Header.Get(EventSchemaHeader))
				assert.JSONEq(t, expectedSummarizedFeatureEventsOutputUnknownFlagWithVersion, string(r.Body))
			})
		})
	}
}

func TestDiagnosticEventForwarding(t *testing.T) {
	for _, e := range allTestEndpoints {
		t.Run(string(e.sdkKind), func(t *testing.T) {
			eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
				req := st.BuildRequest("POST", "/", []byte(eventPayloadForVerbatimOnly), headersWithEventSchema(0))
				req.Header.Add("Authorization", "fake-auth")
				req.Header.Add("User-Agent", "fake-user-agent")
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.DiagnosticEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				r := st.ExpectTestRequest(t, p.requestsCh, time.Second)
				assert.Equal(t, "POST", r.Request.Method)
				assert.Equal(t, e.diagnosticPath, r.Request.URL.Path)
				assert.Equal(t, "fake-auth", r.Request.Header.Get("Authorization"))
				assert.Equal(t, "fake-user-agent", r.Request.Header.Get("User-Agent"))
				assert.Equal(t, eventPayloadForVerbatimOnly, string(r.Body))
			})
		})
	}
}

func TestEventDispatcherReplaceCredential(t *testing.T) {
	eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
		// First, just post some events to all the dispatchers to make sure they've been lazily created.
		// We don't need to check for the original credential in the request headers, because we already
		// have other tests verifying that behavior.
		for _, e := range allTestEndpoints {
			if e.newCredential == nil {
				continue
			}
			for _, schemaVersion := range []int{0, 2} {
				req := st.BuildRequest("POST", "/", []byte(summarizableFeatureEvents), headersWithEventSchema(schemaVersion))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				p.dispatcher.flush()
				_ = st.ExpectTestRequest(t, p.requestsCh, time.Second)
			}
		}

		// Now change both the SDK key and the mobile key (the environment ID can't change)
		p.dispatcher.ReplaceCredential(testServerEndpointInfo.newCredential)
		p.dispatcher.ReplaceCredential(testMobileEndpointInfo.newCredential)

		// Verify that proxied events now use the new credentials
		for _, e := range allTestEndpoints {
			if e.newCredential == nil {
				continue
			}
			for _, schemaVersion := range []int{0, 2} {
				req := st.BuildRequest("POST", "/", []byte(summarizableFeatureEvents), headersWithEventSchema(schemaVersion))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				p.dispatcher.flush()
				r := st.ExpectTestRequest(t, p.requestsCh, time.Second)
				assert.Equal(t, e.newCredential.GetAuthorizationHeaderValue(), r.Request.Header.Get("Authorization"))
			}
		}
	})
}

func TestEventHandlersRejectMalformedJSON(t *testing.T) {
	malformedInput := `[{"no`
	eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
		for _, e := range allTestEndpoints {
			for _, schemaVersion := range []int{0, 2} {
				req := st.BuildRequest("POST", "/", []byte(malformedInput), headersWithEventSchema(schemaVersion))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusAccepted, w.Result().StatusCode)

				p.dispatcher.flush()
				st.ExpectNoTestRequests(t, p.requestsCh, time.Millisecond*20)
			}
		}
	})
}

func TestEventHandlersRejectEmptyBody(t *testing.T) {
	eventRelayTest(t, st.EnvWithAllCredentials, config.EventsConfig{}, func(p eventRelayTestParams) {
		for _, e := range allTestEndpoints {
			for _, schemaVersion := range []int{0, 2} {
				req := st.BuildRequest("POST", "/", []byte{}, headersWithEventSchema(schemaVersion))
				handler := p.dispatcher.GetHandler(e.sdkKind, ldevents.AnalyticsEventDataKind)
				require.NotNil(t, handler)
				w := httptest.NewRecorder()
				handler(w, req)
				assert.Equal(t, http.StatusBadRequest, w.Result().StatusCode)
			}
		}
	})
}

func headersWithEventSchema(schemaVersion int) http.Header {
	headers := make(http.Header)
	headers.Set("Content-Type", "application/json")
	if schemaVersion > 0 {
		headers.Set(EventSchemaHeader, strconv.Itoa(schemaVersion))
	}
	return headers
}


<|start_filename|>internal/core/bigsegments/sync.go<|end_filename|>
package bigsegments

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"

	es "github.com/launchdarkly/eventsource"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldtime"
)

const (
	unboundedPollPath          = "/sdk/big-segments/revisions"
	unboundedStreamPath        = "/big-segments"
	streamReadTimeout          = 5 * time.Minute
	defaultStreamRetryInterval = 10 * time.Second
	synchronizedOnInterval     = 30 * time.Second

	segmentUpdatesChannelBufferSize = 20
)

// BigSegmentSynchronizer synchronizes big segment state for a given environment.
type BigSegmentSynchronizer interface {
	// Start begins synchronization of an environment.
	//
	// This method does not block.
	//
	// If the BigSegmentSynchronizer has already been started, or has been closed, this has no effect.
	Start()

	// HasSynced returns true if the synchronizer has ever successfully synced the data.
	//
	// We use this to determine whether Relay's internal SDK instances should bother trying to query
	// big segments metadata. If we haven't yet written any metadata, then trying to do so would
	// produce useless errors.
	HasSynced() bool

	// SegmentUpdatesCh returns a channel for notifications about segment data updates.
	//
	// Each value posted to this channel represents a batch of updates that the synchronizer has
	// applied. The caller is responsible for reading the channel to avoid blocking the
	// synchronizer.
	SegmentUpdatesCh() <-chan UpdatesSummary

	// Close ends synchronization of an environment.
	//
	// This method does not block.
	//
	// The BigSegmentSynchronizer cannot be restarted after calling Close.
	Close()
}

// UpdatesSummary describes a batch of updates that the synchronizer has applied.
type UpdatesSummary struct {
	// SegmentKeysUpdated is a slice of segment keys (plain keys as used by the SDK-- not segment
	// IDs, i.e. there is no generation suffix).
	SegmentKeysUpdated []string
}

// BigSegmentSynchronizerFactory creates an implementation of BigSegmentSynchronizer. We
// only use a single implementation in real life, but this allows us to use a mock one
// in tests. Calling the factory does not automatically start the synchronizer.
type BigSegmentSynchronizerFactory func(
	httpConfig httpconfig.HTTPConfig,
	store BigSegmentStore,
	pollURI string,
	streamURI string,
	envID config.EnvironmentID,
	sdkKey config.SDKKey,
	loggers ldlog.Loggers,
	logPrefix string,
) BigSegmentSynchronizer

// defaultBigSegmentSynchronizer is the standard implementation of BigSegmentSynchronizer.
type defaultBigSegmentSynchronizer struct {
	httpConfig          httpconfig.HTTPConfig
	store               BigSegmentStore
	pollURI             string
	streamURI           string
	envID               config.EnvironmentID
	sdkKey              config.SDKKey
	streamRetryInterval time.Duration
	segmentUpdatesChan  chan UpdatesSummary
	hasSynced           bool
	syncedLock          sync.RWMutex
	startOnce           sync.Once
	closeChan           chan struct{}
	closeOnce           sync.Once
	loggers             ldlog.Loggers
}

type segmentChangesSummary map[string]struct{}

type applyPatchesResult struct {
	totalPatchesCount   int
	patchesAppliedCount int
	segmentsUpdated     segmentChangesSummary
}

// DefaultBigSegmentSynchronizerFactory creates the default implementation of BigSegmentSynchronizer.
func DefaultBigSegmentSynchronizerFactory(
	httpConfig httpconfig.HTTPConfig,
	store BigSegmentStore,
	pollURI string,
	streamURI string,
	envID config.EnvironmentID,
	sdkKey config.SDKKey,
	loggers ldlog.Loggers,
	logPrefix string,
) BigSegmentSynchronizer {
	return newDefaultBigSegmentSynchronizer(httpConfig, store, pollURI, streamURI, envID, sdkKey, loggers, logPrefix)
}

func newDefaultBigSegmentSynchronizer(
	httpConfig httpconfig.HTTPConfig,
	store BigSegmentStore,
	pollURI string,
	streamURI string,
	envID config.EnvironmentID,
	sdkKey config.SDKKey,
	loggers ldlog.Loggers,
	logPrefix string,
) *defaultBigSegmentSynchronizer {
	s := defaultBigSegmentSynchronizer{
		httpConfig:          httpConfig,
		store:               store,
		pollURI:             strings.TrimSuffix(pollURI, "/") + unboundedPollPath,
		streamURI:           strings.TrimSuffix(streamURI, "/") + unboundedStreamPath,
		envID:               envID,
		sdkKey:              sdkKey,
		streamRetryInterval: defaultStreamRetryInterval,
		segmentUpdatesChan:  make(chan UpdatesSummary, segmentUpdatesChannelBufferSize),
		closeChan:           make(chan struct{}),
		loggers:             loggers,
	}

	if logPrefix != "" {
		logPrefix += " "
	}
	logPrefix += "BigSegmentSynchronizer:"
	s.loggers.SetPrefix(logPrefix)

	return &s
}

type httpStatusError struct {
	statusCode int
}

func (m httpStatusError) Error() string {
	return fmt.Sprintf("HTTP error %d", m.statusCode)
}

func (s *defaultBigSegmentSynchronizer) Start() {
	s.startOnce.Do(func() {
		go s.syncSupervisor()
	})
}

func (s *defaultBigSegmentSynchronizer) HasSynced() bool {
	s.syncedLock.RLock()
	ret := s.hasSynced
	s.syncedLock.RUnlock()
	return ret
}

func (s *defaultBigSegmentSynchronizer) SegmentUpdatesCh() <-chan UpdatesSummary {
	return s.segmentUpdatesChan
}

func (s *defaultBigSegmentSynchronizer) Close() {
	// If we haven't yet started, we still need to close the updates channel; calling
	// startOnce.Do also ensures that Start() will have no effect after this
	s.startOnce.Do(func() {
		close(s.segmentUpdatesChan)
	})
	// If we had already started, then there's a goroutine which will detect the closing
	// of closeChan, and that goroutine will take care of closing the updates channel
	s.closeOnce.Do(func() {
		close(s.closeChan)
	})
}

func (s *defaultBigSegmentSynchronizer) syncSupervisor() {
	isRetry := false
	for {
		err := s.sync(isRetry)
		if err != nil {
			s.loggers.Error("Synchronization failed:", err)
			if statusError, ok := err.(httpStatusError); ok {
				if !isHTTPErrorRecoverable(statusError.statusCode) {
					return
				}
			}
		}
		s.loggers.Warn("Will retry")
		timer := time.NewTimer(s.streamRetryInterval)
		defer timer.Stop()
		select {
		case <-s.closeChan:
			close(s.segmentUpdatesChan)
			return
		case <-timer.C:
		}
		isRetry = true
	}
}

func (s *defaultBigSegmentSynchronizer) sync(isRetry bool) error {
	s.loggers.Debug("Polling for big segment updates")
	segmentsUpdated := make(segmentChangesSummary)
	for {
	SyncLoop:
		for {
			select {
			case <-s.closeChan:
				return nil
			default:
				done, updates, err := s.poll()
				if err != nil {
					return err
				}
				if isRetry {
					s.loggers.Warn("Re-established connection")
					isRetry = false
				}
				segmentsUpdated.addAll(updates)
				if done {
					break SyncLoop
				}
			}
		}

		stream, err := s.connectStream()
		if err != nil {
			return err
		}
		defer stream.Close()

		done, updates, err := s.poll()
		if err != nil {
			return err
		}
		segmentsUpdated.addAll(updates)
		if !done {
			continue
		}

		s.loggers.Debug("Marking store as synchronized")
		err = s.setSynced()
		if err != nil {
			s.loggers.Error("Updating store timestamp failed:", err)
			return err
		}

		s.notifySegmentsUpdated(segmentsUpdated)

		return s.consumeStream(stream)
	}
}

func (s *defaultBigSegmentSynchronizer) setSynced() error {
	err := s.store.setSynchronizedOn(ldtime.UnixMillisNow())
	if err != nil {
		return err
	}
	s.syncedLock.Lock()
	s.hasSynced = true
	s.syncedLock.Unlock()
	return nil
}

// Tests whether an HTTP error status represents a condition that might resolve
// on its own if we retry, or at least should not make us permanently stop
// sending requests.
func isHTTPErrorRecoverable(statusCode int) bool {
	if statusCode >= 400 && statusCode < 500 {
		switch statusCode {
		case 400: // bad request
			return true
		case 408: // request timeout
			return true
		case 429: // too many requests
			return true
		default:
			return false // all other 4xx errors are unrecoverable
		}
	}
	return true
}

func (s *defaultBigSegmentSynchronizer) poll() (bool, segmentChangesSummary, error) {
	client := s.httpConfig.Client()

	request, err := http.NewRequest("GET", s.pollURI, nil)
	if err != nil {
		return false, nil, err
	}

	request.Header.Set("Authorization", string(s.sdkKey))

	cursor, err := s.store.getCursor()
	if err != nil {
		return false, segmentChangesSummary{}, err
	}

	if cursor != "" {
		query := request.URL.Query()
		query.Add("after", cursor)
		request.URL.RawQuery = query.Encode()
	}

	s.loggers.Debugf("Polling %s", request.URL)
	response, err := client.Do(request)
	if err != nil {
		return false, segmentChangesSummary{}, err
	}
	defer response.Body.Close() //nolint:errcheck

	if response.StatusCode != 200 {
		return false, segmentChangesSummary{}, &httpStatusError{response.StatusCode}
	}

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return false, segmentChangesSummary{}, err
	}

	applyPatchResult, err := s.applyPatches(responseBody)

	return applyPatchResult.totalPatchesCount == 0, applyPatchResult.segmentsUpdated, err
}

func (s *defaultBigSegmentSynchronizer) connectStream() (*es.Stream, error) {
	s.loggers.Debugf("Making stream request to %s", s.streamURI)
	request, err := http.NewRequest("GET", s.streamURI, nil)
	if err != nil {
		return nil, err
	}

	request.Header.Set("Authorization", string(s.sdkKey))

	client := s.httpConfig.Client()
	client.Timeout = 0

	stream, err := es.SubscribeWithRequestAndOptions(request,
		es.StreamOptionHTTPClient(client),
		es.StreamOptionReadTimeout(streamReadTimeout),
		es.StreamOptionLogger(s.loggers.ForLevel(ldlog.Info)),
		es.StreamOptionErrorHandler(func(err error) es.StreamErrorHandlerResult {
			s.loggers.Warnf("Stream connection failed: %s", err)
			return es.StreamErrorHandlerResult{CloseNow: true}
		}),
	)
	if err != nil {
		if se, ok := err.(es.SubscriptionError); ok {
			return nil, &httpStatusError{se.Code}
		}
		return nil, err
	}

	return stream, nil
}

func (s *defaultBigSegmentSynchronizer) consumeStream(stream *es.Stream) error {
	for {
		timer := time.NewTimer(synchronizedOnInterval)
		select {
		case event, ok := <-stream.Events:
			timer.Stop()
			if !ok {
				s.loggers.Debug("Stream ended")
				return nil
			}

			s.loggers.Debug("Received update(s) from stream")
			applyPatchResult, err := s.applyPatches([]byte(event.Data()))
			if err != nil {
				return err
			}
			s.notifySegmentsUpdated(applyPatchResult.segmentsUpdated)
			if applyPatchResult.patchesAppliedCount < applyPatchResult.totalPatchesCount {
				return nil // forces a restart if we got an out-of-order patch
			}

			if err := s.setSynced(); err != nil {
				return err
			}
		case <-timer.C:
			err := s.setSynced()
			if err != nil {
				return err
			}
		case <-s.closeChan:
			timer.Stop()
			return nil
		}
	}
}

// Returns total number of patches, number of patches applied, raw segment IDs, error
func (s *defaultBigSegmentSynchronizer) applyPatches(jsonData []byte) (applyPatchesResult, error) {
	var patches []bigSegmentPatch
	err := json.Unmarshal(jsonData, &patches)
	if err != nil {
		return applyPatchesResult{}, err
	}

	ret := applyPatchesResult{
		totalPatchesCount: len(patches),
		segmentsUpdated:   make(segmentChangesSummary),
	}
	for _, patch := range patches {
		s.loggers.Debugf("Received patch for version %q (from previous version %q)", patch.Version, patch.PreviousVersion)
		success, err := s.store.applyPatch(patch)
		if err != nil {
			return ret, err
		}
		if !success {
			s.loggers.Warnf("Received a patch to previous version %q which was not the latest known version; skipping", patch.PreviousVersion)
			break
		}
		ret.patchesAppliedCount++
		ret.segmentsUpdated.addSegmentID(patch.SegmentID)
	}
	if ret.patchesAppliedCount > 0 {
		updatesDesc := "updates"
		if ret.patchesAppliedCount == 1 {
			updatesDesc = "update"
		}
		s.loggers.Infof("Applied %d %s", ret.patchesAppliedCount, updatesDesc)
	}
	return ret, nil
}

func (s *defaultBigSegmentSynchronizer) notifySegmentsUpdated(segmentsUpdated segmentChangesSummary) {
	keys := segmentsUpdated.getUpdatedSegmentKeys()
	if len(keys) != 0 {
		s.segmentUpdatesChan <- UpdatesSummary{SegmentKeysUpdated: keys}
	}
}

func segmentIDToSegmentKey(segmentID string) string {
	if p := strings.LastIndexByte(segmentID, '.'); p > 0 {
		return segmentID[0:p]
	}
	return segmentID
}

func (s *segmentChangesSummary) addSegmentID(segmentID string) {
	(*s)[segmentIDToSegmentKey(segmentID)] = struct{}{}
}

func (s *segmentChangesSummary) addAll(other segmentChangesSummary) {
	for key := range other {
		(*s)[key] = struct{}{}
	}
}

func (s *segmentChangesSummary) getUpdatedSegmentKeys() []string {
	if len(*s) == 0 {
		return nil
	}
	ret := make([]string, 0, len(*s))
	for key := range *s {
		ret = append(ret, key)
	}
	return ret
}


<|start_filename|>internal/core/bigsegments/store_dynamodb.go<|end_filename|>
package bigsegments

import (
	"strconv"

	ct "github.com/launchdarkly/go-configtypes"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldtime"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
)

const (
	tablePartitionKey         = "namespace"
	tableSortKey              = "key"
	dynamoDBCursorAttr        = "lastVersion"
	dynamoDBIncludedAttr      = "included"
	dynamoDBExcludedAttr      = "excluded"
	dynamoDBSyncTimeAttr      = "synchronizedOn"
	updateExpressionAdd       = "ADD #0 :0"
	updateExpressionRemove    = "DELETE #0 :0"
	dynamoTransactionMaxItems = 25
)

// dynamoDBBigSegmentStore implements BigSegmentStore for DynamoDB.
type dynamoDBBigSegmentStore struct {
	client  dynamodbiface.DynamoDBAPI
	loggers ldlog.Loggers
	table   string
	prefix  string
}

func dynamoDBMetadataKey(prefix string) string {
	return dynamoDBPrefixedKey(prefix, "big_segments_metadata")
}

func dynamoDBUserDataKey(prefix string) string {
	return dynamoDBPrefixedKey(prefix, "big_segments_user")
}

func dynamoDBPrefixedKey(prefix, key string) string {
	if prefix == "" {
		return key
	}
	return prefix + ":" + key
}

func newDynamoDBBigSegmentStore(
	url ct.OptURLAbsolute,
	config aws.Config,
	loggers ldlog.Loggers,
	table string,
	prefix string,
) (*dynamoDBBigSegmentStore, error) {
	if url.IsDefined() {
		config.Endpoint = aws.String(url.String())
	}

	sess, err := session.NewSession(&config)
	if err != nil {
		return nil, err
	}

	store := dynamoDBBigSegmentStore{
		table:   table,
		loggers: loggers,
		prefix:  prefix,
		client:  dynamodb.New(sess),
	}

	store.loggers.SetPrefix("DynamoDBBigSegmentStore:")
	store.loggers.Infof(`Using DynamoDB table %s`, store.table)

	return &store, nil
}

func (store *dynamoDBBigSegmentStore) makeTransactionItem(updateExpression, attribute, segmentID, userKey string) *dynamodb.TransactWriteItem {
	return &dynamodb.TransactWriteItem{
		Update: &dynamodb.Update{
			TableName:        aws.String(store.table),
			UpdateExpression: aws.String(updateExpression),
			ExpressionAttributeNames: map[string]*string{
				"#0": aws.String(attribute),
			},
			ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
				":0": {SS: []*string{aws.String(segmentID)}},
			},
			Key: map[string]*dynamodb.AttributeValue{
				tablePartitionKey: {S: aws.String(dynamoDBUserDataKey(store.prefix))},
				tableSortKey:      {S: aws.String(userKey)},
			},
		},
	}
}

func makeCursorUpdateCondition(previousVersion string) (string, map[string]*string, map[string]*dynamodb.AttributeValue) {
	names := map[string]*string{"#0": aws.String(dynamoDBCursorAttr)}
	if previousVersion == "" {
		return "attribute_not_exists(#0)", names, nil
	}
	return "#0 = :0", names, map[string]*dynamodb.AttributeValue{
		":0": {S: aws.String(previousVersion)},
	}
}

func (store *dynamoDBBigSegmentStore) applyPatch(patch bigSegmentPatch) (bool, error) {
	bigSegmentsMetadataKeyWithPrefix := dynamoDBMetadataKey(store.prefix)

	txConditionExpression, txExprAttrNames, txExprAttrValues := makeCursorUpdateCondition(patch.PreviousVersion)

	conditionCheckItem := &dynamodb.TransactWriteItem{
		ConditionCheck: &dynamodb.ConditionCheck{
			ConditionExpression:       aws.String(txConditionExpression),
			TableName:                 aws.String(store.table),
			ExpressionAttributeNames:  txExprAttrNames,
			ExpressionAttributeValues: txExprAttrValues,
			Key: map[string]*dynamodb.AttributeValue{
				tablePartitionKey: {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
				tableSortKey:      {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
			},
			ReturnValuesOnConditionCheckFailure: aws.String(dynamodb.ReturnValuesOnConditionCheckFailureNone),
		},
	}

	totalItems := len(patch.Changes.Included.Add) + len(patch.Changes.Included.Remove) +
		len(patch.Changes.Excluded.Add) + len(patch.Changes.Excluded.Remove)
	transactionItems := make([]*dynamodb.TransactWriteItem, 0, totalItems)

	for _, user := range patch.Changes.Included.Add {
		item := store.makeTransactionItem(updateExpressionAdd, dynamoDBIncludedAttr, patch.SegmentID, user)
		transactionItems = append(transactionItems, item)
	}

	for _, user := range patch.Changes.Excluded.Add {
		item := store.makeTransactionItem(updateExpressionAdd, dynamoDBExcludedAttr, patch.SegmentID, user)
		transactionItems = append(transactionItems, item)
	}

	for _, user := range patch.Changes.Included.Remove {
		item := store.makeTransactionItem(updateExpressionRemove, dynamoDBIncludedAttr, patch.SegmentID, user)
		transactionItems = append(transactionItems, item)
	}

	for _, user := range patch.Changes.Excluded.Remove {
		item := store.makeTransactionItem(updateExpressionRemove, dynamoDBExcludedAttr, patch.SegmentID, user)
		transactionItems = append(transactionItems, item)
	}

	transactionBatch := make([]*dynamodb.TransactWriteItem, 0, dynamoTransactionMaxItems)

	for batchStart := 0; batchStart < len(transactionItems); batchStart += dynamoTransactionMaxItems - 1 {
		batchEnd := batchStart + dynamoTransactionMaxItems - 1
		if batchEnd > len(transactionItems) {
			batchEnd = len(transactionItems)
		}
		transactionBatch = append(transactionBatch, conditionCheckItem)
		transactionBatch = append(transactionBatch, transactionItems[batchStart:batchEnd]...)
		_, err := store.client.TransactWriteItems(&dynamodb.TransactWriteItemsInput{
			TransactItems: transactionBatch,
		})
		if err != nil {
			// DynamoDB doesn't seem to provide a more convenient programmatic way to distinguish
			// "transaction was cancelled due to the condition check" from other errors here; we
			// need to go to this trouble because we want the synchronizer to be able to log an
			// out-of-order update in a clear way that doesn't look like a random database error.
			if tce, ok := err.(*dynamodb.TransactionCanceledException); ok {
				for _, reason := range tce.CancellationReasons {
					if reason.Code != nil && *reason.Code == "ConditionalCheckFailed" {
						return false, nil
					}
				}
			}
			return false, err
		}
		transactionBatch = transactionBatch[:0]
	}

	updateConditionExpression, updateExprAttrNames, updateExprAttrValues := makeCursorUpdateCondition(patch.PreviousVersion)
	if updateExprAttrValues == nil {
		updateExprAttrValues = map[string]*dynamodb.AttributeValue{}
	}
	updateExprAttrValues[":1"] = &dynamodb.AttributeValue{
		S: aws.String(patch.Version),
	}
	updateCursorInput := dynamodb.UpdateItemInput{
		ConditionExpression:       aws.String(updateConditionExpression),
		TableName:                 aws.String(store.table),
		ExpressionAttributeNames:  updateExprAttrNames,
		ExpressionAttributeValues: updateExprAttrValues,
		Key: map[string]*dynamodb.AttributeValue{
			tablePartitionKey: {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
			tableSortKey:      {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
		},
		UpdateExpression: aws.String("SET #0 = :1"),
	}

	_, err := store.client.UpdateItem(&updateCursorInput)
	if err == nil {
		return true, nil
	}
	return false, err
}

func (store *dynamoDBBigSegmentStore) getCursor() (string, error) {
	metadataKey := dynamoDBMetadataKey(store.prefix)
	result, err := store.client.GetItem(&dynamodb.GetItemInput{
		TableName:      aws.String(store.table),
		ConsistentRead: aws.Bool(true),
		ExpressionAttributeNames: map[string]*string{
			"#0": aws.String(dynamoDBCursorAttr),
		},
		Key: map[string]*dynamodb.AttributeValue{
			tablePartitionKey: {S: aws.String(metadataKey)},
			tableSortKey:      {S: aws.String(metadataKey)},
		},
		ProjectionExpression: aws.String("#0"),
	})
	if err != nil || len(result.Item) == 0 {
		return "", err
	}
	item := result.Item[dynamoDBCursorAttr]
	if item == nil || item.S == nil {
		return "", nil
	}
	return *item.S, nil
}

func (store *dynamoDBBigSegmentStore) setSynchronizedOn(synchronizedOn ldtime.UnixMillisecondTime) error {
	bigSegmentsMetadataKeyWithPrefix := dynamoDBMetadataKey(store.prefix)
	unixMilliseconds := strconv.FormatUint(uint64(synchronizedOn), 10)
	_, err := store.client.UpdateItem(&dynamodb.UpdateItemInput{
		TableName: aws.String(store.table),
		Key: map[string]*dynamodb.AttributeValue{
			tablePartitionKey: {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
			tableSortKey:      {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
		},
		UpdateExpression:         aws.String("SET #0 = :0"),
		ExpressionAttributeNames: map[string]*string{"#0": aws.String(dynamoDBSyncTimeAttr)},
		ExpressionAttributeValues: map[string]*dynamodb.AttributeValue{
			":0": {
				N: aws.String(unixMilliseconds),
			},
		},
	})
	return err
}

func (store *dynamoDBBigSegmentStore) GetSynchronizedOn() (ldtime.UnixMillisecondTime, error) {
	bigSegmentsMetadataKeyWithPrefix := dynamoDBMetadataKey(store.prefix)
	result, err := store.client.GetItem(&dynamodb.GetItemInput{
		TableName:      aws.String(store.table),
		ConsistentRead: aws.Bool(true),
		Key: map[string]*dynamodb.AttributeValue{
			tablePartitionKey: {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
			tableSortKey:      {S: aws.String(bigSegmentsMetadataKeyWithPrefix)},
		},
		ProjectionExpression: aws.String(dynamoDBSyncTimeAttr),
	})
	if err != nil || len(result.Item) == 0 {
		return 0, err
	}
	item := result.Item[dynamoDBSyncTimeAttr]
	if item == nil || item.N == nil {
		return 0, nil
	}
	value, err := strconv.Atoi(*item.N)
	if err != nil {
		return 0, nil
	}
	return ldtime.UnixMillisecondTime(value), nil
}

func (store *dynamoDBBigSegmentStore) Close() error {
	return nil
}


<|start_filename|>internal/core/streams/stream_provider.go<|end_filename|>
package streams

import (
	"net/http"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"

	"github.com/launchdarkly/eventsource"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
)

// StreamProvider is an abstraction of a specific kind of SSE event stream, such as the server-side SDK
// "/all" stream. The streams package provides default implementations of this interface for the streams
// that are supported by the standard Relay Proxy.
//
// Each StreamProvider expects a specific kind of SDKCredential, e.g. config.SDKKey for the server-side
// streams. If the wrong kind of credential is passed, it should behave as it would for an unrecognized
// key. It is important that there can be more than one StreamProvider for a given credential.
type StreamProvider interface {
	// Handler returns an HTTP request handler for the given credential. It can return nil if it does
	// not support this type of credential.
	Handler(credential config.SDKCredential) http.HandlerFunc

	// Register tells the StreamProvider about an environment that it should support, and returns an
	// implementation of EnvStreamsUpdates for pushing updates related to that environment. It can
	// return nil if it does not support this type of credential.
	Register(credential config.SDKCredential, store EnvStoreQueries, loggers ldlog.Loggers) EnvStreamProvider

	// Close tells the StreamProvider to release all of its resources and close all connections.
	Close()
}

// EnvStreamProvider is an abstraction of publishing events to a stream for a specific environment.
// Implementations of this interface are created by StreamProvider.Register().
type EnvStreamProvider interface {
	EnvStreamUpdates // SendAllDataUpdate, SendSingleItemUpdate

	// SendHeartbeat sends keep-alive data on the stream.
	SendHeartbeat()

	// Close releases all resources for this EnvStreamProvider and closes all connections to it.
	Close()
}

// NewStreamProvider creates a StreamProvider implementation for the specified kind of stream endpoint.
func NewStreamProvider(kind basictypes.StreamKind, maxConnTime time.Duration) StreamProvider {
	switch kind {
	case basictypes.ServerSideFlagsOnlyStream:
		return &serverSideFlagsOnlyStreamProvider{
			server: newSSEServer(maxConnTime),
		}
	case basictypes.MobilePingStream:
		return &clientSidePingStreamProvider{
			server:     newSSEServer(maxConnTime),
			isJSClient: false,
		}
	case basictypes.JSClientPingStream:
		return &clientSidePingStreamProvider{
			server:     newSSEServer(maxConnTime),
			isJSClient: true,
		}
	default:
		return &serverSideStreamProvider{
			server: newSSEServer(maxConnTime),
		}
	}
}

func newSSEServer(maxConnTime time.Duration) *eventsource.Server {
	s := eventsource.NewServer()
	s.Gzip = false
	s.AllowCORS = true
	s.ReplayAll = true
	s.MaxConnTime = maxConnTime
	return s
}

func removeDeleted(items []ldstoretypes.KeyedItemDescriptor) []ldstoretypes.KeyedItemDescriptor {
	var ret []ldstoretypes.KeyedItemDescriptor
	for i, keyedItem := range items {
		if keyedItem.Item.Item == nil {
			if ret == nil {
				ret = make([]ldstoretypes.KeyedItemDescriptor, i)
				copy(ret, items[0:i])
			}
		} else {
			if ret != nil {
				ret = append(ret, keyedItem)
			}
		}
	}
	if ret == nil {
		return items
	}
	return ret
}


<|start_filename|>internal/core/sdks/sdk_context.go<|end_filename|>
package sdks

import (
	ld "gopkg.in/launchdarkly/go-server-sdk.v5"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"
)

// NewSimpleClientContext creates a simple implementation of the SDK's ClientContext interface for
// initializing SDK components. The SDK doesn't surface a way to do this because components aren't
// normally created outside of its own constructor.
func NewSimpleClientContext(sdkKey string, sdkConfig ld.Config) interfaces.ClientContext {
	basic := interfaces.BasicConfiguration{SDKKey: sdkKey}
	if sdkConfig.HTTP == nil {
		sdkConfig.HTTP = ldcomponents.HTTPConfiguration()
	}
	http, _ := sdkConfig.HTTP.CreateHTTPConfiguration(basic)
	if sdkConfig.Logging == nil {
		sdkConfig.Logging = ldcomponents.Logging()
	}
	logging, _ := sdkConfig.Logging.CreateLoggingConfiguration(basic)
	return simpleClientContext{basic: basic, http: http, logging: logging}
}

type simpleClientContext struct {
	basic   interfaces.BasicConfiguration
	http    interfaces.HTTPConfiguration
	logging interfaces.LoggingConfiguration
}

func (s simpleClientContext) GetBasic() interfaces.BasicConfiguration     { return s.basic }
func (s simpleClientContext) GetHTTP() interfaces.HTTPConfiguration       { return s.http }
func (s simpleClientContext) GetLogging() interfaces.LoggingConfiguration { return s.logging }


<|start_filename|>internal/core/relayenv/env_context.go<|end_filename|>
package relayenv

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/bigsegments"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/events"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sdks"
	"github.com/launchdarkly/ld-relay/v6/internal/core/streams"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	ldeval "gopkg.in/launchdarkly/go-server-sdk-evaluation.v1"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
)

// EnvContext is the interface for all Relay operations that are specific to one configured LD environment.
//
// The EnvContext is normally associated with an LDClient instance from the Go SDK, and allows direct access
// to the DataStore that is used by the SDK client. However, these are created asynchronously since the client
// connection may take a while, so it is possible for the client and store references to be nil if initialization
// is not yet complete.
type EnvContext interface {
	io.Closer

	// GetIdentifiers returns information about the environment and project names and keys.
	GetIdentifiers() EnvIdentifiers

	// SetIdentifiers updates the environment and project names and keys.
	SetIdentifiers(EnvIdentifiers)

	// GetCredentials returns all currently enabled and non-deprecated credentials for the environment.
	GetCredentials() []config.SDKCredential

	// GetDeprecatedCredentials returns all deprecated and not-yet-removed credentials for the environment.
	GetDeprecatedCredentials() []config.SDKCredential

	// AddCredential adds a new credential for the environment.
	//
	// If the credential is an SDK key, then a new SDK client is started with that SDK key, and event forwarding
	// to server-side endpoints is switched to use the new key.
	AddCredential(config.SDKCredential)

	// RemoveCredential removes a credential from the environment. Any active stream connections using that
	// credential are immediately dropped.
	//
	// If the credential is an SDK key, then the SDK client that we started with that SDK key is disposed of.
	RemoveCredential(config.SDKCredential)

	// DeprecateCredential marks an existing credential as not being a preferred one, without removing it or
	// dropping any connections. It will no longer be included in the return value of GetCredentials(). This is
	// used in Relay Proxy Enterprise when an SDK key is being changed but the old key has not expired yet.
	DeprecateCredential(config.SDKCredential)

	// GetClient returns the SDK client instance for this environment. This is nil if initialization is not yet
	// complete. Rather than providing the full client object, we use the simpler sdks.LDClientContext which
	// includes only the operations Relay needs to do.
	GetClient() sdks.LDClientContext

	// GetStore returns the SDK data store instance for this environment. This is nil if initialization is not
	// yet complete.
	GetStore() interfaces.DataStore

	// GetEvaluator returns an instance of the evaluation engine for evaluating feature flags in this environment.
	// This is nil if initialization is not yet complete.
	GetEvaluator() ldeval.Evaluator

	// GetBigSegmentStore returns the big segment data store instance for this environment. If a big
	// segment store is not configured this returns nil.
	GetBigSegmentStore() bigsegments.BigSegmentStore

	// GetLoggers returns a Loggers instance that is specific to this environment. We configure each of these to
	// have its own prefix string and, optionally, its own log level.
	GetLoggers() ldlog.Loggers

	// GetHandler returns the HTTP handler for the specified kind of stream requests and credential for this
	// environment. If there is none, it returns a handler for a 404 status (not nil).
	GetStreamHandler(streams.StreamProvider, config.SDKCredential) http.Handler

	// GetEventDispatcher returns the object that proxies events for this environment.
	GetEventDispatcher() *events.EventDispatcher

	// GetJSClientContext returns the JSClientContext that is used for browser endpoints.
	GetJSClientContext() JSClientContext

	// GetMetricsContext returns the Context that should be used for OpenCensus operations related to this
	// environment.
	GetMetricsContext() context.Context

	// GetTTL returns the configured cache TTL for PHP SDK endpoints for this environment.
	GetTTL() time.Duration

	// SetTTL changes the configured cache TTL for PHP SDK endpoints for this environment.
	SetTTL(time.Duration)

	// GetInitError returns an error if initialization has failed, or nil otherwise.
	GetInitError() error

	// IsSecureMode returns true if client-side evaluation requests for this environment must have a valid
	// secure mode hash.
	IsSecureMode() bool

	// SetSecureMode changes the secure mode setting.
	SetSecureMode(bool)

	// GetCreationTime returns the time that this EnvContext was created.
	GetCreationTime() time.Time

	// GetDataStoreInfo returns information about the environment's data store.
	GetDataStoreInfo() sdks.DataStoreEnvironmentInfo

	// FlushMetricsEvents is used in testing to ensure that metrics events are delivered promptly.
	FlushMetricsEvents()
}

// EnvIdentifiers contains environment and project name and key properties.
//
// When running in Relay Proxy Enterprise's auto-configuration mode, EnvKey, EnvName, ProjKey, and ProjName are
// copied from the LaunchDarkly dashboard settings. Otherwise, those are all blank and ConfiguredName is set in
// the local configuration.
type EnvIdentifiers struct {
	// EnvKey is the environment key (normally a lowercase string like "production").
	EnvKey string

	// EnvName is the environment name (normally a title-cased string like "Production").
	EnvName string

	// ProjKey is the project key (normally a lowercase string like "my-application").
	ProjKey string

	// ProjName is the project name (normally a title-cased string like "My Application").
	ProjName string

	// ConfiguredName is a human-readable unique name for this environment, if the user specified one. When
	// using a local configuration, this is always set; in auto-configuration mode, it is always empty (but
	// EnvIdentifiers.GetDisplayName() will compute one).
	ConfiguredName string
}

// GetDisplayName returns a human-readable unique name for this environment. If none was set in the
// configuration, it computes one in the format "ProjName EnvName".
func (ei EnvIdentifiers) GetDisplayName() string {
	if ei.ConfiguredName == "" {
		return fmt.Sprintf("%s %s", ei.ProjName, ei.EnvName)
	}
	return ei.ConfiguredName
}

// GetEnvironmentID is a helper for extracting the EnvironmentID, if any, from the set of credentials.
func GetEnvironmentID(env EnvContext) config.EnvironmentID {
	for _, c := range env.GetCredentials() {
		if e, ok := c.(config.EnvironmentID); ok {
			return e
		}
	}
	return ""
}


<|start_filename|>internal/core/bigsegments/model.go<|end_filename|>
package bigsegments

// bigSegmentPatchChangesMutations lists users that should be added or removed
// to either the included or excluded set of a big segment.
type bigSegmentPatchChangesMutations struct {
	Add    []string `json:"add"`
	Remove []string `json:"remove"`
}

// bigSegmentPatchChanges represents changes to the included and excluded sets
// of a segment.
type bigSegmentPatchChanges struct {
	Included bigSegmentPatchChangesMutations `json:"included"`
	Excluded bigSegmentPatchChangesMutations `json:"excluded"`
}

// bigSegmentPatch represents a patch of of a big segment in an environment.
type bigSegmentPatch struct {
	EnvironmentID   string                 `json:"environmentId"`
	SegmentID       string                 `json:"segmentId"`
	Version         string                 `json:"version"`
	PreviousVersion string                 `json:"previousVersion"`
	Changes         bigSegmentPatchChanges `json:"changes"`
}


<|start_filename|>internal/core/sharedtest/testenv/test_env_context.go<|end_filename|>
// Package testenv contains test helpers that reference the relayenv package. These are in sharedtest/testenv
// rather than just sharedtest so that sharedtest can be used by relayenv itself without a circular reference.
package testenv

import (
	"time"

	"github.com/launchdarkly/ld-relay/v6/internal/core/relayenv"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sdks"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest/testclient"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"
)

func NewTestEnvContext(name string, shouldBeInitialized bool, store interfaces.DataStore) relayenv.EnvContext {
	return NewTestEnvContextWithClientFactory(name, testclient.FakeLDClientFactory(shouldBeInitialized), store)
}

func NewTestEnvContextWithClientFactory(
	name string,
	f sdks.ClientFactoryFunc,
	store interfaces.DataStore,
) relayenv.EnvContext {
	dataStoreFactory := ldcomponents.InMemoryDataStore()
	if store != nil {
		dataStoreFactory = sharedtest.ExistingDataStoreFactory{Instance: store}
	}
	readyCh := make(chan relayenv.EnvContext)
	c, err := relayenv.NewEnvContext(relayenv.EnvContextImplParams{
		Identifiers:      relayenv.EnvIdentifiers{ConfiguredName: name},
		ClientFactory:    f,
		DataStoreFactory: dataStoreFactory,
		UserAgent:        "fake-user-agent",
		Loggers:          ldlog.NewDisabledLoggers(),
	}, readyCh)
	if err != nil {
		panic(err)
	}
	select {
	case <-readyCh:
		return c
	case <-time.After(time.Second):
		panic("timed out waiting for client initialization")
	}
}

func MakeTestContextWithData() relayenv.EnvContext {
	return NewTestEnvContext("", true, sharedtest.MakeStoreWithData(true))
}


<|start_filename|>internal/core/relayenv/env_context_impl_test.go<|end_filename|>
package relayenv

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"regexp"
	"sync"
	"testing"
	"time"

	"github.com/launchdarkly/eventsource"
	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/bigsegments"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/metrics"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sdks"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	st "github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest/testclient"
	"github.com/launchdarkly/ld-relay/v6/internal/core/streams"

	"github.com/launchdarkly/go-configtypes"
	"github.com/launchdarkly/go-test-helpers/v2/httphelpers"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlogtest"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldvalue"
	ldevents "gopkg.in/launchdarkly/go-sdk-events.v1"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldbuilders"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
)

const envName = "envname"

func requireEnvReady(t *testing.T, readyCh <-chan EnvContext) EnvContext {
	select {
	case e := <-readyCh:
		return e
	case <-time.After(time.Second):
		require.Fail(t, "timed out waiting for environment")
		return nil
	}
}

func requireClientReady(t *testing.T, clientCh chan *testclient.FakeLDClient) *testclient.FakeLDClient {
	select {
	case c := <-clientCh:
		return c
	case <-time.After(time.Second):
		require.Fail(t, "timed out waiting for client")
		return nil
	}
}

func makeBasicEnv(t *testing.T, envConfig config.EnvConfig, clientFactory sdks.ClientFactoryFunc,
	loggers ldlog.Loggers, readyCh chan EnvContext) EnvContext {
	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:   EnvIdentifiers{ConfiguredName: envName},
		EnvConfig:     envConfig,
		ClientFactory: clientFactory,
		Loggers:       loggers,
	}, readyCh)
	require.NoError(t, err)
	return env
}

func TestConstructorBasicProperties(t *testing.T) {
	envConfig := st.EnvWithAllCredentials.Config
	envConfig.TTL = configtypes.NewOptDuration(time.Hour)
	envConfig.SecureMode = true
	readyCh := make(chan EnvContext, 1)

	clientCh := make(chan *testclient.FakeLDClient, 1)
	clientFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh)

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, clientFactory, mockLog.Loggers, readyCh)
	defer env.Close()

	assert.Equal(t, envName, env.GetIdentifiers().ConfiguredName)
	assert.Equal(t, time.Hour, env.GetTTL())
	assert.True(t, env.IsSecureMode())
	assert.Nil(t, env.GetEventDispatcher())                        // events were not enabled
	assert.Equal(t, context.Background(), env.GetMetricsContext()) // metrics aren't being used

	creds := env.GetCredentials()
	assert.Len(t, creds, 3)
	assert.Contains(t, creds, envConfig.SDKKey)
	assert.Contains(t, creds, envConfig.MobileKey)
	assert.Contains(t, creds, envConfig.EnvID)

	assert.Equal(t, env, requireEnvReady(t, readyCh))
	assert.Equal(t, env.GetClient(), requireClientReady(t, clientCh))
	assert.Nil(t, env.GetInitError())

	assert.NotNil(t, env.GetStore())
}

func TestConstructorWithOnlySDKKey(t *testing.T) {
	envConfig := st.EnvMain.Config
	readyCh := make(chan EnvContext, 1)

	clientCh := make(chan *testclient.FakeLDClient, 1)
	clientFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh)

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, clientFactory, mockLog.Loggers, readyCh)
	defer env.Close()

	assert.Equal(t, []config.SDKCredential{envConfig.SDKKey}, env.GetCredentials())

	assert.Equal(t, env, requireEnvReady(t, readyCh))
	assert.Equal(t, env.GetClient(), requireClientReady(t, clientCh))
	assert.Nil(t, env.GetInitError())
}

func TestConstructorWithJSClientContext(t *testing.T) {
	envConfig := st.EnvWithAllCredentials.Config
	jsClientContext := JSClientContext{Origins: []string{"origin"}}
	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:     EnvIdentifiers{ConfiguredName: envName},
		EnvConfig:       envConfig,
		ClientFactory:   testclient.FakeLDClientFactory(true),
		JSClientContext: jsClientContext,
		Loggers:         ldlog.NewDisabledLoggers(),
	}, nil)
	require.NoError(t, err)
	defer env.Close()

	assert.Equal(t, jsClientContext, env.GetJSClientContext())
}

func TestLogPrefix(t *testing.T) {
	testPrefix := func(desc string, mode LogNameMode, sdkKey config.SDKKey, envID config.EnvironmentID, expected string) {
		t.Run(desc, func(t *testing.T) {
			envConfig := config.EnvConfig{SDKKey: sdkKey, EnvID: envID}
			mockLog := ldlogtest.NewMockLog()
			env, err := NewEnvContext(EnvContextImplParams{
				Identifiers:   EnvIdentifiers{ConfiguredName: "name"},
				EnvConfig:     envConfig,
				ClientFactory: testclient.FakeLDClientFactory(true),
				UserAgent:     "user-agent",
				LogNameMode:   mode,
				Loggers:       mockLog.Loggers,
			}, nil)
			require.NoError(t, err)
			defer env.Close()
			envImpl := env.(*envContextImpl)
			envImpl.loggers.Error("message")
			mockLog.AssertMessageMatch(t, true, ldlog.Error, "^"+regexp.QuoteMeta(expected)+" message")
		})
	}

	testPrefix("SDK key", LogNameIsSDKKey, config.SDKKey("1234567890"), config.EnvironmentID("abcdefghij"), "[env: ...7890]")
	testPrefix("env ID", LogNameIsEnvID, config.SDKKey("1234567890"), config.EnvironmentID("abcdefghij"), "[env: ...ghij]")
	testPrefix("env ID not set", LogNameIsEnvID, config.SDKKey("1234567890"), "", "[env: ...7890]")
	testPrefix("impossibly short SDK key", LogNameIsSDKKey, config.SDKKey("890"), config.EnvironmentID("abcdefghij"), "[env: 890]")
	testPrefix("impossibly short env ID", LogNameIsEnvID, config.SDKKey("1234567890"), config.EnvironmentID("hij"), "[env: hij]")
}

func TestAddRemoveCredential(t *testing.T) {
	envConfig := st.EnvMain.Config

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactory(true), mockLog.Loggers, nil)
	defer env.Close()

	assert.Equal(t, []config.SDKCredential{envConfig.SDKKey}, env.GetCredentials())

	env.AddCredential(st.EnvWithAllCredentials.Config.MobileKey)
	env.AddCredential(st.EnvWithAllCredentials.Config.EnvID)

	creds := env.GetCredentials()
	assert.Len(t, creds, 3)
	assert.Contains(t, creds, envConfig.SDKKey)
	assert.Contains(t, creds, st.EnvWithAllCredentials.Config.MobileKey)
	assert.Contains(t, creds, st.EnvWithAllCredentials.Config.EnvID)

	env.RemoveCredential(st.EnvWithAllCredentials.Config.MobileKey)

	creds = env.GetCredentials()
	assert.Len(t, creds, 2)
	assert.Contains(t, creds, envConfig.SDKKey)
	assert.Contains(t, creds, st.EnvWithAllCredentials.Config.EnvID)
}

func TestAddExistingCredentialDoesNothing(t *testing.T) {
	envConfig := st.EnvMain.Config

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactory(true), mockLog.Loggers, nil)
	defer env.Close()

	assert.Equal(t, []config.SDKCredential{envConfig.SDKKey}, env.GetCredentials())

	env.AddCredential(st.EnvWithAllCredentials.Config.MobileKey)

	creds := env.GetCredentials()
	assert.Len(t, creds, 2)
	assert.Contains(t, creds, envConfig.SDKKey)
	assert.Contains(t, creds, st.EnvWithAllCredentials.Config.MobileKey)

	env.AddCredential(st.EnvWithAllCredentials.Config.MobileKey)

	creds = env.GetCredentials()
	assert.Len(t, creds, 2)
	assert.Contains(t, creds, envConfig.SDKKey)
	assert.Contains(t, creds, st.EnvWithAllCredentials.Config.MobileKey)
}

func TestChangeSDKKey(t *testing.T) {
	envConfig := st.EnvMain.Config
	readyCh := make(chan EnvContext, 1)
	newKey := config.SDKKey("new-key")

	clientCh := make(chan *testclient.FakeLDClient, 1)
	clientFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh)

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, clientFactory, mockLog.Loggers, readyCh)
	defer env.Close()

	assert.Equal(t, env, requireEnvReady(t, readyCh))
	client1 := requireClientReady(t, clientCh)
	assert.Equal(t, env.GetClient(), client1)
	assert.Nil(t, env.GetInitError())

	env.AddCredential(newKey)
	env.DeprecateCredential(envConfig.SDKKey)

	assert.Equal(t, []config.SDKCredential{newKey}, env.GetCredentials())

	client2 := requireClientReady(t, clientCh)
	assert.NotEqual(t, client1, client2)
	assert.Equal(t, env.GetClient(), client2)

	select {
	case <-client1.CloseCh:
		require.Fail(t, "client for deprecated key should not have been closed")
	case <-time.After(time.Millisecond * 20):
		break
	}

	env.RemoveCredential(envConfig.SDKKey)

	assert.Equal(t, []config.SDKCredential{newKey}, env.GetCredentials())

	client1.AwaitClose(t, time.Millisecond*20)
}

func TestSDKClientCreationFails(t *testing.T) {
	envConfig := st.EnvWithAllCredentials.Config
	envConfig.TTL = configtypes.NewOptDuration(time.Hour)
	envConfig.SecureMode = true
	readyCh := make(chan EnvContext, 1)

	fakeError := errors.New("sorry")

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env := makeBasicEnv(t, envConfig, testclient.ClientFactoryThatFails(fakeError), mockLog.Loggers, readyCh)
	defer env.Close()

	assert.Equal(t, env, requireEnvReady(t, readyCh))
	assert.Equal(t, fakeError, env.GetInitError())
	assert.Nil(t, env.GetStore())
}

func TestDisplayName(t *testing.T) {
	ei1 := EnvIdentifiers{ProjName: "a", EnvName: "b", ConfiguredName: "thing"}
	assert.Equal(t, "thing", ei1.GetDisplayName())

	ei2 := EnvIdentifiers{ProjName: "a", EnvName: "b"}
	assert.Equal(t, "a b", ei2.GetDisplayName())
}

func TestMetricsAreExportedForEnvironment(t *testing.T) {
	// We already have tests for openCensusEventsExporter in the metrics package, but this test verifies that
	// exporting is configured automatically for every environment that we add (if not disabled).
	view.SetReportingPeriod(time.Millisecond * 10)
	trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)
	fakeUserAgent := "fake-user-agent"

	handler, requestsCh := httphelpers.RecordingHandler(httphelpers.HandlerWithStatus(202))
	httphelpers.WithServer(handler, func(server *httptest.Server) {
		var allConfig config.Config
		allConfig.Events.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(server.URL)
		metricsManager, err := metrics.NewManager(config.MetricsConfig{}, time.Minute, mockLog.Loggers)
		require.NoError(t, err)
		env, err := NewEnvContext(EnvContextImplParams{
			Identifiers:    EnvIdentifiers{ConfiguredName: envName},
			EnvConfig:      st.EnvMain.Config,
			AllConfig:      allConfig,
			ClientFactory:  testclient.FakeLDClientFactory(true),
			MetricsManager: metricsManager,
			UserAgent:      fakeUserAgent,
			Loggers:        mockLog.Loggers,
		}, nil)
		require.NoError(t, err)
		defer env.Close()
		envImpl := env.(*envContextImpl)
		metrics.WithCount(env.GetMetricsContext(), fakeUserAgent, func() {
			require.Eventually(t, func() bool {
				flushMetricsEvents(envImpl)
				select {
				case req := <-requestsCh:
					mockLog.Loggers.Infof("received metrics events: %s", req.Body)
					data := ldvalue.Parse(req.Body)
					event := data.GetByIndex(0)
					if !event.IsNull() {
						conns := event.GetByKey("connections")
						return event.GetByKey("kind").StringValue() == "relayMetrics" &&
							conns.Count() == 1 &&
							conns.GetByIndex(0).GetByKey("userAgent").StringValue() == fakeUserAgent &&
							conns.GetByIndex(0).GetByKey("current").IntValue() == 1
					}
				default:
					break
				}
				return false
			}, time.Second, time.Millisecond*10, "timed out waiting for metrics event with counter")
		}, metrics.BrowserConns)
	})
}

func TestMetricsAreNotExportedForEnvironmentIfDisabled(t *testing.T) {
	var allConfig config.Config
	allConfig.Main.DisableInternalUsageMetrics = true
	testMetricsDisabled(t, allConfig)
}

func TestMetricsAreNotExportedForEnvironmentInOfflineMode(t *testing.T) {
	var allConfig config.Config
	allConfig.OfflineMode.FileDataSource = "fake-file-path"
	testMetricsDisabled(t, allConfig)
}

func testMetricsDisabled(t *testing.T, allConfig config.Config) {
	view.SetReportingPeriod(time.Millisecond * 10)
	trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)
	fakeUserAgent := "fake-user-agent"

	handler, requestsCh := httphelpers.RecordingHandler(httphelpers.HandlerWithStatus(202))
	httphelpers.WithServer(handler, func(server *httptest.Server) {
		allConfig.Events.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(server.URL)
		metricsManager, err := metrics.NewManager(config.MetricsConfig{}, time.Minute, mockLog.Loggers)
		require.NoError(t, err)
		env, err := NewEnvContext(EnvContextImplParams{
			Identifiers:    EnvIdentifiers{ConfiguredName: envName},
			EnvConfig:      st.EnvMain.Config,
			AllConfig:      allConfig,
			ClientFactory:  testclient.FakeLDClientFactory(true),
			MetricsManager: metricsManager,
			Loggers:        mockLog.Loggers,
		}, nil)
		require.NoError(t, err)
		defer env.Close()
		envImpl := env.(*envContextImpl)
		metrics.WithCount(env.GetMetricsContext(), fakeUserAgent, func() {
			require.Never(t, func() bool {
				flushMetricsEvents(envImpl)
				select {
				case <-requestsCh:
					return true
				default:
					break
				}
				return false
			}, time.Millisecond*100, time.Millisecond*10, "received unexpected metrics event")
		}, metrics.BrowserConns)
	})
}

func TestEventDispatcherIsCreatedIfSendEventsIsTrueAndNotInOfflineMode(t *testing.T) {
	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	eventRecorderHandler, requestsCh := httphelpers.RecordingHandler(httphelpers.HandlerWithStatus(202))
	httphelpers.WithServer(eventRecorderHandler, func(server *httptest.Server) {
		var allConfig config.Config
		allConfig.Main.DisableInternalUsageMetrics = true
		allConfig.Events.SendEvents = true
		allConfig.Events.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(server.URL)
		allConfig.Events.FlushInterval = configtypes.NewOptDuration(time.Millisecond * 10)
		env, err := NewEnvContext(EnvContextImplParams{
			Identifiers:   EnvIdentifiers{ConfiguredName: envName},
			EnvConfig:     st.EnvMain.Config,
			AllConfig:     allConfig,
			ClientFactory: testclient.FakeLDClientFactory(true),
			Loggers:       mockLog.Loggers,
		}, nil)
		require.NoError(t, err)
		defer env.Close()
		envImpl := env.(*envContextImpl)

		ed := envImpl.GetEventDispatcher()
		require.NotNil(t, ed)
		eventDispatchHandler := ed.GetHandler(basictypes.ServerSDK, ldevents.AnalyticsEventDataKind)
		require.NotNil(t, eventDispatchHandler)

		rr := httptest.NewRecorder()
		headers := make(http.Header)
		headers.Add("Content-Type", "application/json")
		headers.Add("Authorization", string(st.EnvMain.Config.SDKKey))
		body := `[{"kind":"identify","creationDate":1000,"key":"userkey","user":{"key":"userkey"}}]`
		req := st.BuildRequest("POST", server.URL+"/bulk", []byte(body), headers)
		eventDispatchHandler(rr, req)
		require.Equal(t, 202, rr.Result().StatusCode)

		eventPost := st.ExpectTestRequest(t, requestsCh, time.Second)
		require.Equal(t, string(st.EnvMain.Config.SDKKey), eventPost.Request.Header.Get("Authorization"))
		require.Equal(t, string(body), string(eventPost.Body))
	})
}

func TestEventDispatcherIsNotCreatedIfSendEventsIsTrueAndNotInOfflineMode(t *testing.T) {
	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	eventRecorderHandler, _ := httphelpers.RecordingHandler(httphelpers.HandlerWithStatus(202))
	httphelpers.WithServer(eventRecorderHandler, func(server *httptest.Server) {
		var allConfig config.Config
		allConfig.Main.DisableInternalUsageMetrics = true
		allConfig.OfflineMode.FileDataSource = "fake-file-path"
		allConfig.Events.SendEvents = true
		allConfig.Events.EventsURI, _ = configtypes.NewOptURLAbsoluteFromString(server.URL)
		allConfig.Events.FlushInterval = configtypes.NewOptDuration(time.Millisecond * 10)
		env, err := NewEnvContext(EnvContextImplParams{
			Identifiers:   EnvIdentifiers{ConfiguredName: envName},
			EnvConfig:     st.EnvMain.Config,
			AllConfig:     allConfig,
			ClientFactory: testclient.FakeLDClientFactory(true),
			Loggers:       mockLog.Loggers,
		}, nil)
		require.NoError(t, err)
		defer env.Close()
		envImpl := env.(*envContextImpl)

		ed := envImpl.GetEventDispatcher()
		require.Nil(t, ed)
	})
}

func TestBigSegmentsSynchronizerIsCreatedIfBigSegmentStoreExists(t *testing.T) {
	envConfig := st.EnvMain.Config
	allConfig := config.Config{}

	fakeBigSegmentStoreFactory := func(config.EnvConfig, config.Config, ldlog.Loggers) (bigsegments.BigSegmentStore, error) {
		return bigsegments.NewNullBigSegmentStore(), nil
	}
	fakeSynchronizerFactory := &mockBigSegmentSynchronizerFactory{}

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:                   EnvIdentifiers{ConfiguredName: st.EnvMain.Name},
		EnvConfig:                     envConfig,
		AllConfig:                     allConfig,
		BigSegmentStoreFactory:        fakeBigSegmentStoreFactory,
		BigSegmentSynchronizerFactory: fakeSynchronizerFactory.create,
		ClientFactory:                 testclient.FakeLDClientFactory(true),
		SDKBigSegmentsConfigFactory: ldcomponents.BigSegments(
			mockSDKBigSegmentStoreFactory{&sharedtest.NoOpSDKBigSegmentStore{}},
		),
		Loggers: mockLog.Loggers,
	}, nil)
	require.NoError(t, err)

	if assert.NotNil(t, fakeSynchronizerFactory.synchronizer) {
		assert.False(t, fakeSynchronizerFactory.synchronizer.isStarted())
		assert.False(t, fakeSynchronizerFactory.synchronizer.isClosed())
	}

	env.Close()

	assert.True(t, fakeSynchronizerFactory.synchronizer.isClosed())
}

func TestBigSegmentsSynchronizerIsStartedByFullDataUpdateWithBigSegment(t *testing.T) {
	envConfig := st.EnvMain.Config
	allConfig := config.Config{}

	fakeBigSegmentStoreFactory := func(config.EnvConfig, config.Config, ldlog.Loggers) (bigsegments.BigSegmentStore, error) {
		return bigsegments.NewNullBigSegmentStore(), nil
	}
	fakeSynchronizerFactory := &mockBigSegmentSynchronizerFactory{}

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:                   EnvIdentifiers{ConfiguredName: st.EnvMain.Name},
		EnvConfig:                     envConfig,
		AllConfig:                     allConfig,
		BigSegmentStoreFactory:        fakeBigSegmentStoreFactory,
		BigSegmentSynchronizerFactory: fakeSynchronizerFactory.create,
		ClientFactory:                 testclient.FakeLDClientFactory(true),
		SDKBigSegmentsConfigFactory: ldcomponents.BigSegments(
			mockSDKBigSegmentStoreFactory{&sharedtest.NoOpSDKBigSegmentStore{}},
		),
		Loggers: mockLog.Loggers,
	}, nil)
	require.NoError(t, err)
	defer env.Close()

	synchronizer := fakeSynchronizerFactory.synchronizer
	require.NotNil(t, synchronizer)
	assert.False(t, synchronizer.isStarted())

	// Simulate receiving some data
	updates := env.(*envContextImpl).storeAdapter.GetUpdates()

	s1 := ldbuilders.NewSegmentBuilder("s1").Build()
	dataWithNoBigSegment := []ldstoretypes.Collection{
		{
			Kind: ldstoreimpl.Segments(),
			Items: []ldstoretypes.KeyedItemDescriptor{
				{Key: "s1", Item: sharedtest.SegmentDesc(s1)},
			},
		},
	}
	updates.SendAllDataUpdate(dataWithNoBigSegment)

	assert.False(t, synchronizer.isStarted())

	s2 := ldbuilders.NewSegmentBuilder("s2").Unbounded(true).Generation(1).Build()
	dataWithBigSegment := []ldstoretypes.Collection{
		{
			Kind: ldstoreimpl.Segments(),
			Items: []ldstoretypes.KeyedItemDescriptor{
				{Key: "s1", Item: sharedtest.SegmentDesc(s1)},
				{Key: "s2", Item: sharedtest.SegmentDesc(s2)},
			},
		},
	}
	updates.SendAllDataUpdate(dataWithBigSegment)

	assert.True(t, synchronizer.isStarted())
}

func TestBigSegmentsSynchronizerIsStartedBySingleItemUpdateWithBigSegment(t *testing.T) {
	envConfig := st.EnvMain.Config
	allConfig := config.Config{}

	fakeBigSegmentStoreFactory := func(config.EnvConfig, config.Config, ldlog.Loggers) (bigsegments.BigSegmentStore, error) {
		return bigsegments.NewNullBigSegmentStore(), nil
	}
	fakeSynchronizerFactory := &mockBigSegmentSynchronizerFactory{}

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:                   EnvIdentifiers{ConfiguredName: st.EnvMain.Name},
		EnvConfig:                     envConfig,
		AllConfig:                     allConfig,
		BigSegmentStoreFactory:        fakeBigSegmentStoreFactory,
		BigSegmentSynchronizerFactory: fakeSynchronizerFactory.create,
		ClientFactory:                 testclient.FakeLDClientFactory(true),
		SDKBigSegmentsConfigFactory: ldcomponents.BigSegments(
			mockSDKBigSegmentStoreFactory{&sharedtest.NoOpSDKBigSegmentStore{}},
		),
		Loggers: mockLog.Loggers,
	}, nil)
	require.NoError(t, err)
	defer env.Close()

	synchronizer := fakeSynchronizerFactory.synchronizer
	require.NotNil(t, synchronizer)
	assert.False(t, synchronizer.isStarted())

	// Simulate receiving some data
	updates := env.(*envContextImpl).storeAdapter.GetUpdates()

	f1 := ldbuilders.NewFlagBuilder("f1").Build()
	updates.SendSingleItemUpdate(ldstoreimpl.Features(), f1.Key, sharedtest.FlagDesc(f1))

	assert.False(t, synchronizer.isStarted())

	s1 := ldbuilders.NewSegmentBuilder("s1").Build()
	updates.SendSingleItemUpdate(ldstoreimpl.Segments(), s1.Key, sharedtest.SegmentDesc(s1))

	assert.False(t, synchronizer.isStarted())

	s2 := ldbuilders.NewSegmentBuilder("s2").Unbounded(true).Generation(1).Build()
	updates.SendSingleItemUpdate(ldstoreimpl.Segments(), s2.Key, sharedtest.SegmentDesc(s2))

	assert.True(t, synchronizer.isStarted())
}

func TestReceivingBigSegmentsUpdateCausesClientSideInvalidationEvent(t *testing.T) {
	envConfig := st.EnvClientSide.Config
	allConfig := config.Config{}

	fakeBigSegmentStoreFactory := func(config.EnvConfig, config.Config, ldlog.Loggers) (bigsegments.BigSegmentStore, error) {
		return bigsegments.NewNullBigSegmentStore(), nil
	}
	fakeSynchronizerFactory := &mockBigSegmentSynchronizerFactory{}

	mockLog := ldlogtest.NewMockLog()
	defer mockLog.DumpIfTestFailed(t)

	jsClientStreams := streams.NewStreamProvider(basictypes.JSClientPingStream, time.Hour)
	sdkStartedCh := make(chan EnvContext)
	env, err := NewEnvContext(EnvContextImplParams{
		Identifiers:                   EnvIdentifiers{ConfiguredName: st.EnvMain.Name},
		EnvConfig:                     envConfig,
		AllConfig:                     allConfig,
		BigSegmentStoreFactory:        fakeBigSegmentStoreFactory,
		BigSegmentSynchronizerFactory: fakeSynchronizerFactory.create,
		ClientFactory:                 testclient.FakeLDClientFactory(true),
		SDKBigSegmentsConfigFactory: ldcomponents.BigSegments(
			mockSDKBigSegmentStoreFactory{&st.NoOpSDKBigSegmentStore{}},
		),
		StreamProviders: []streams.StreamProvider{jsClientStreams},
		Loggers:         mockLog.Loggers,
	}, sdkStartedCh)
	require.NoError(t, err)
	defer env.Close()

	synchronizer := fakeSynchronizerFactory.synchronizer
	require.NotNil(t, synchronizer)

	streamHandler := env.GetStreamHandler(jsClientStreams, envConfig.EnvID)

	// Make sure the data store is initialized, otherwise the client-side endpoint won't broadcast a ping
	<-sdkStartedCh
	_ = env.GetStore().Init(nil)

	req, _ := http.NewRequest("GET", "", nil)
	sharedtest.WithStreamRequest(t, req, streamHandler, func(eventCh <-chan eventsource.Event) {
		initEvent := sharedtest.ExpectStreamChEvent(t, eventCh, time.Minute)
		assert.Equal(t, "ping", initEvent.Event())

		sharedtest.ExpectNoStreamChEvent(t, eventCh, time.Millisecond*100)

		synchronizer.updateCh <- bigsegments.UpdatesSummary{SegmentKeysUpdated: []string{"fake-segment-key"}}

		pingEvent := sharedtest.ExpectStreamChEvent(t, eventCh, time.Second)
		assert.Equal(t, "ping", pingEvent.Event())
	})
}

// This method forces the metrics events exporter to post an event to the event publisher, and then triggers a
// flush of the event publisher. Because both of those actions are asynchronous, it may be necessary to call it
// more than once to ensure that the newly posted event is included in the flush.
func flushMetricsEvents(c *envContextImpl) {
	if c.metricsEventPub != nil {
		c.metricsEnv.FlushEventsExporter()
		c.metricsEventPub.Flush()
	}
}

type mockBigSegmentSynchronizerFactory struct {
	synchronizer *mockBigSegmentSynchronizer
}

func (f *mockBigSegmentSynchronizerFactory) create(
	httpConfig httpconfig.HTTPConfig,
	store bigsegments.BigSegmentStore,
	pollURI string,
	streamURI string,
	envID config.EnvironmentID,
	sdkKey config.SDKKey,
	loggers ldlog.Loggers,
	logPrefix string,
) bigsegments.BigSegmentSynchronizer {
	f.synchronizer = &mockBigSegmentSynchronizer{updateCh: make(chan bigsegments.UpdatesSummary)}
	return f.synchronizer
}

type mockBigSegmentSynchronizer struct {
	started  bool
	closed   bool
	updateCh chan bigsegments.UpdatesSummary
	lock     sync.Mutex
}

func (s *mockBigSegmentSynchronizer) Start() {
	s.lock.Lock()
	s.started = true
	s.lock.Unlock()
}

func (s *mockBigSegmentSynchronizer) HasSynced() bool {
	return true
}

func (s *mockBigSegmentSynchronizer) SegmentUpdatesCh() <-chan bigsegments.UpdatesSummary {
	return s.updateCh
}

func (s *mockBigSegmentSynchronizer) Close() {
	s.lock.Lock()
	s.closed = true
	s.lock.Unlock()
}

func (s *mockBigSegmentSynchronizer) isStarted() bool {
	s.lock.Lock()
	defer s.lock.Unlock()
	return s.started
}

func (s *mockBigSegmentSynchronizer) isClosed() bool {
	s.lock.Lock()
	defer s.lock.Unlock()
	return s.closed
}

type mockSDKBigSegmentStoreFactory struct {
	store interfaces.BigSegmentStore
}

func (m mockSDKBigSegmentStoreFactory) CreateBigSegmentStore(c interfaces.ClientContext) (interfaces.BigSegmentStore, error) {
	return m.store, nil
}


<|start_filename|>internal/core/relayenv/env_context_impl.go<|end_filename|>
package relayenv

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"
	"github.com/launchdarkly/ld-relay/v6/internal/core/bigsegments"
	"github.com/launchdarkly/ld-relay/v6/internal/core/httpconfig"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/events"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/metrics"
	"github.com/launchdarkly/ld-relay/v6/internal/core/internal/store"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sdks"
	"github.com/launchdarkly/ld-relay/v6/internal/core/streams"
	"github.com/launchdarkly/ld-relay/v6/internal/util"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	ldeval "gopkg.in/launchdarkly/go-server-sdk-evaluation.v1"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldmodel"
	ld "gopkg.in/launchdarkly/go-server-sdk.v5"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"
)

// LogNameMode is used in NewEnvContext to determine whether the environment's log messages should be
// tagged by SDK key or by environment ID.
type LogNameMode bool

const (
	// LogNameIsSDKKey means the log messages should be tagged with the last 4 characters of the SDK key.
	// This is the default behavior for the Relay Proxy.
	LogNameIsSDKKey LogNameMode = false

	// LogNameIsEnvID means the log messages should be tagged with the last 4 characters of the environment
	// ID. This is the default behavior for Relay Proxy Enterprise when running in auto-configuration mode,
	// where we always know the environment ID but the SDK key is subject to change.
	LogNameIsEnvID LogNameMode = true
)

func errInitPublisher(err error) error {
	return fmt.Errorf("failed to initialize event publisher: %w", err)
}

func errInitMetrics(err error) error {
	return fmt.Errorf("failed to initialize metrics for environment: %w", err)
}

// EnvContextImplParams contains the constructor parameters for NewEnvContextImpl. These have their
// own type because there are a lot of them, and many are irrelevant in tests.
type EnvContextImplParams struct {
	Identifiers                   EnvIdentifiers
	EnvConfig                     config.EnvConfig
	AllConfig                     config.Config
	ClientFactory                 sdks.ClientFactoryFunc
	DataStoreFactory              interfaces.DataStoreFactory
	DataStoreInfo                 sdks.DataStoreEnvironmentInfo
	StreamProviders               []streams.StreamProvider
	JSClientContext               JSClientContext
	MetricsManager                *metrics.Manager
	BigSegmentStoreFactory        bigsegments.BigSegmentStoreFactory
	BigSegmentSynchronizerFactory bigsegments.BigSegmentSynchronizerFactory
	SDKBigSegmentsConfigFactory   interfaces.BigSegmentsConfigurationFactory // set only in tests
	UserAgent                     string
	LogNameMode                   LogNameMode
	Loggers                       ldlog.Loggers
}

type envContextImpl struct {
	mu               sync.RWMutex
	clients          map[config.SDKKey]sdks.LDClientContext
	storeAdapter     *store.SSERelayDataStoreAdapter
	loggers          ldlog.Loggers
	credentials      map[config.SDKCredential]bool // true if not deprecated
	identifiers      EnvIdentifiers
	secureMode       bool
	envStreams       *streams.EnvStreams
	streamProviders  []streams.StreamProvider
	handlers         map[streams.StreamProvider]map[config.SDKCredential]http.Handler
	jsContext        JSClientContext
	evaluator        ldeval.Evaluator
	eventDispatcher  *events.EventDispatcher
	bigSegmentSync   bigsegments.BigSegmentSynchronizer
	bigSegmentStore  bigsegments.BigSegmentStore
	sdkBigSegments   *ldstoreimpl.BigSegmentStoreWrapper
	sdkConfig        ld.Config
	sdkClientFactory sdks.ClientFactoryFunc
	sdkInitTimeout   time.Duration
	metricsManager   *metrics.Manager
	metricsEnv       *metrics.EnvironmentManager
	metricsEventPub  events.EventPublisher
	dataStoreInfo    sdks.DataStoreEnvironmentInfo
	globalLoggers    ldlog.Loggers
	ttl              time.Duration
	initErr          error
	creationTime     time.Time
}

// Implementation of the DataStoreQueries interface that the streams package uses as an abstraction of
// accessing our data store.
type envContextStoreQueries struct {
	context *envContextImpl
}

// Implementation of the EnvStreamUpdates interface that intercepts all updates from the SDK to the
// data store.
type envContextStreamUpdates struct {
	context *envContextImpl
}

// NewEnvContext creates the internal implementation of EnvContext.
//
// It immediately begins trying to initialize the SDK client for this environment. Since that might
// take a while, it is done on a separate goroutine. The EnvContext instance is returned immediately
// in an uninitialized state, and once the SDK client initialization has either succeeded or failed,
// the same EnvContext will be pushed to the channel readyCh.
//
// NewEnvContext can also immediately return an error, with a nil EnvContext, if the configuration is
// invalid.
func NewEnvContext( //nolint:gocyclo
	params EnvContextImplParams,
	readyCh chan<- EnvContext,
	// readyCh is a separate parameter because it's not a property of the environment itself, but
	// just part of the semantics of the constructor
) (EnvContext, error) {
	var thingsToCleanUp util.CleanupTasks // keeps track of partially constructed things in case we exit early
	defer thingsToCleanUp.Run()

	offlineMode := params.AllConfig.OfflineMode.FileDataSource != ""
	envConfig := params.EnvConfig
	allConfig := params.AllConfig

	envLoggers := params.Loggers
	logPrefix := makeLogPrefix(params.LogNameMode, envConfig.SDKKey, envConfig.EnvID)
	envLoggers.SetPrefix(logPrefix)
	envLoggers.SetMinLevel(
		envConfig.LogLevel.GetOrElse(
			allConfig.Main.LogLevel.GetOrElse(ldlog.Info),
		),
	)

	httpConfig, err := httpconfig.NewHTTPConfig(allConfig.Proxy, envConfig.SDKKey, params.UserAgent, params.Loggers)
	if err != nil {
		return nil, err
	}

	credentials := make(map[config.SDKCredential]bool, 3)
	credentials[envConfig.SDKKey] = true
	if envConfig.MobileKey != "" {
		credentials[envConfig.MobileKey] = true
	}
	if envConfig.EnvID != "" {
		credentials[envConfig.EnvID] = true
	}

	envContext := &envContextImpl{
		identifiers:      params.Identifiers,
		clients:          make(map[config.SDKKey]sdks.LDClientContext),
		credentials:      credentials,
		loggers:          envLoggers,
		secureMode:       envConfig.SecureMode,
		streamProviders:  params.StreamProviders,
		handlers:         make(map[streams.StreamProvider]map[config.SDKCredential]http.Handler),
		jsContext:        params.JSClientContext,
		sdkClientFactory: params.ClientFactory,
		sdkInitTimeout:   allConfig.Main.InitTimeout.GetOrElse(config.DefaultInitTimeout),
		metricsManager:   params.MetricsManager,
		globalLoggers:    params.Loggers,
		ttl:              envConfig.TTL.GetOrElse(0),
		dataStoreInfo:    params.DataStoreInfo,
		creationTime:     time.Now(),
	}

	bigSegmentStoreFactory := params.BigSegmentStoreFactory
	if bigSegmentStoreFactory == nil {
		bigSegmentStoreFactory = bigsegments.DefaultBigSegmentStoreFactory
	}
	bigSegmentStore, err := bigSegmentStoreFactory(envConfig, allConfig, envLoggers)
	if err != nil {
		return nil, err
	}
	if bigSegmentStore != nil {
		thingsToCleanUp.AddCloser(bigSegmentStore)
		envContext.bigSegmentStore = bigSegmentStore

		factory := params.BigSegmentSynchronizerFactory
		if factory == nil {
			factory = bigsegments.DefaultBigSegmentSynchronizerFactory
		}
		envContext.bigSegmentSync = factory(
			httpConfig, bigSegmentStore, allConfig.Main.BaseURI.String(), allConfig.Main.StreamURI.String(),
			envConfig.EnvID, envConfig.SDKKey, envLoggers, logPrefix)
		thingsToCleanUp.AddFunc(envContext.bigSegmentSync.Close)
		segmentUpdateCh := envContext.bigSegmentSync.SegmentUpdatesCh()
		if segmentUpdateCh != nil {
			go func() {
				for range segmentUpdateCh {
					// BigSegmentSynchronizer sends to this channel after processing a batch of
					// big segment updates. The value it sends is a list of segment keys, but in
					// the current implementation, we don't care what those keys are because we'll
					// just be broadcasting a "ping" to all connected client-side SDKs. In the future
					// if we have real evaluation streams, we'll need to determine which flags should
					// be re-evaluated based on the segments.
					if envContext.sdkBigSegments != nil {
						envContext.sdkBigSegments.ClearCache()
					}
					if envContext.envStreams != nil {
						envContext.envStreams.InvalidateClientSideState()
					}
					// If we shut down the environment, the BigSegmentSynchronizer will be closed which
					// will also cause this channel to be closed, exiting this goroutine.
				}
			}()
		}
		// We deliberate do not call bigSegmentSync.Start() here because we don't want the synchronizer to
		// start until we know that at least one big segment exists. That's implemented by the
		// envContextStreamUpdates methods.
	}

	envStreams := streams.NewEnvStreams(
		params.StreamProviders,
		envContextStoreQueries{envContext},
		allConfig.Main.HeartbeatInterval.GetOrElse(config.DefaultHeartbeatInterval),
		envLoggers,
	)
	envContext.envStreams = envStreams
	thingsToCleanUp.AddCloser(envStreams)

	envStreamUpdates := &envContextStreamUpdates{
		context: envContext,
	}

	for c := range credentials {
		envStreams.AddCredential(c)
	}
	for _, sp := range params.StreamProviders {
		handlers := make(map[config.SDKCredential]http.Handler)
		for c := range credentials {
			h := sp.Handler(c)
			if h != nil {
				handlers[c] = h
			}
		}
		envContext.handlers[sp] = handlers
	}

	dataStoreFactory := params.DataStoreFactory
	if dataStoreFactory == nil {
		dataStoreFactory = ldcomponents.InMemoryDataStore()
	}
	storeAdapter := store.NewSSERelayDataStoreAdapter(dataStoreFactory, envStreamUpdates)
	envContext.storeAdapter = storeAdapter

	var eventDispatcher *events.EventDispatcher
	if allConfig.Events.SendEvents {
		if offlineMode {
			envLoggers.Info("Events will be accepted for this environment, but will be discarded, since offline mode is enabled")
		} else {
			envLoggers.Info("Proxying events for this environment")
			eventLoggers := envLoggers
			eventLoggers.SetPrefix(logPrefix + " (event proxy)")
			eventDispatcher = events.NewEventDispatcher(envConfig.SDKKey, envConfig.MobileKey, envConfig.EnvID,
				envLoggers, allConfig.Events, httpConfig, storeAdapter)
		}
	}
	envContext.eventDispatcher = eventDispatcher

	streamURI := allConfig.Main.StreamURI.String()
	if streamURI == "" {
		streamURI = config.DefaultStreamURI
	}
	eventsURI := allConfig.Events.EventsURI.String()
	if eventsURI == "" {
		eventsURI = config.DefaultEventsURI
	}

	enableDiagnostics := !allConfig.Main.DisableInternalUsageMetrics && !offlineMode
	var em *metrics.EnvironmentManager
	if params.MetricsManager != nil {
		if enableDiagnostics {
			pubLoggers := envLoggers
			pubLoggers.SetPrefix(logPrefix + " (usage metrics)")
			eventsPublisher, err := events.NewHTTPEventPublisher(envConfig.SDKKey, httpConfig, pubLoggers,
				events.OptionURI(eventsURI))
			if err != nil {
				return nil, errInitPublisher(err)
			}
			thingsToCleanUp.AddFunc(eventsPublisher.Close)
			envContext.metricsEventPub = eventsPublisher
		}

		em, err = params.MetricsManager.AddEnvironment(params.Identifiers.GetDisplayName(), envContext.metricsEventPub)
		if err != nil {
			return nil, errInitMetrics(err)
		}
		thingsToCleanUp.AddFunc(func() { params.MetricsManager.RemoveEnvironment(em) })
	}
	envContext.metricsEnv = em

	disconnectedStatusTime := allConfig.Main.DisconnectedStatusTime.GetOrElse(config.DefaultDisconnectedStatusTime)

	envContext.sdkConfig = ld.Config{
		DataSource:       ldcomponents.StreamingDataSource().BaseURI(streamURI),
		DataStore:        storeAdapter,
		DiagnosticOptOut: !enableDiagnostics,
		Events:           ldcomponents.SendEvents().BaseURI(eventsURI),
		HTTP:             httpConfig.SDKHTTPConfigFactory,
		Logging: ldcomponents.Logging().
			Loggers(envLoggers).
			LogDataSourceOutageAsErrorAfter(disconnectedStatusTime),
	}

	// If appropriate, create the SDK subcomponent that will be used for flag evaluations. We're
	// creating and managing it separately from the full SDK instance that we'll be creating (in
	// startSDKClient) - we use the SDK instance only for talking to LaunchDarkly and populating
	// the data store, not for evaluating flags, because Relay needs to customize the evaluation
	// behavior. The other component we need for evaluations is the Evaluator, but we can't create
	// that one we get to startSDKClient because it has to be hooked up to the SDK's data store.
	if bigSegmentStore != nil {
		configFactory := params.SDKBigSegmentsConfigFactory
		if configFactory == nil {
			configFactory, err = sdks.ConfigureBigSegments(allConfig, envConfig, params.Loggers)
			if err != nil {
				return nil, err
			}
		}
		bigSegConfig, err := configFactory.CreateBigSegmentsConfiguration(
			sdks.NewSimpleClientContext(string(envConfig.SDKKey), envContext.sdkConfig))
		if err != nil {
			return nil, err
		}
		if bigSegConfig != nil {
			envContext.sdkBigSegments = ldstoreimpl.NewBigSegmentStoreWrapperWithConfig(
				ldstoreimpl.BigSegmentsConfigurationProperties{
					Store:              bigSegConfig.GetStore(),
					StatusPollInterval: bigSegConfig.GetStatusPollInterval(),
					StaleAfter:         bigSegConfig.GetStaleAfter(),
					UserCacheSize:      bigSegConfig.GetUserCacheSize(),
					UserCacheTime:      bigSegConfig.GetUserCacheTime(),
					StartPolling:       false, // we will start it later if we see a big segment
				},
				nil,
				envLoggers,
			)
			thingsToCleanUp.AddFunc(envContext.sdkBigSegments.Close)
		}
	}

	// Connecting may take time, so do this in parallel
	go envContext.startSDKClient(envConfig.SDKKey, readyCh, allConfig.Main.IgnoreConnectionErrors)

	thingsToCleanUp.Clear() // we've succeeded so we do not want to throw away these things

	return envContext, nil
}

func (c *envContextImpl) startSDKClient(sdkKey config.SDKKey, readyCh chan<- EnvContext, suppressErrors bool) {
	client, err := c.sdkClientFactory(sdkKey, c.sdkConfig, c.sdkInitTimeout)
	c.mu.Lock()
	name := c.identifiers.GetDisplayName()
	if client != nil {
		c.clients[sdkKey] = client

		// The data store instance is created by the SDK when it creates the client. Now that
		// we have a data store, we can finish setting up the Evaluator that we'll use for this
		// environment.
		store := c.storeAdapter.GetStore()
		dataProvider := ldstoreimpl.NewDataStoreEvaluatorDataProvider(store, c.loggers)
		if c.sdkBigSegments == nil {
			c.evaluator = ldeval.NewEvaluator(dataProvider)
		} else {
			c.evaluator = ldeval.NewEvaluatorWithBigSegments(dataProvider, c.sdkBigSegments)
		}
	}
	c.initErr = err
	c.mu.Unlock()

	if err != nil {
		if suppressErrors {
			c.globalLoggers.Warnf("Ignoring error initializing LaunchDarkly client for %q: %+v",
				name, err)
		} else {
			c.globalLoggers.Errorf("Error initializing LaunchDarkly client for %q: %+v",
				name, err)
			if readyCh != nil {
				readyCh <- c
			}
			return
		}
	} else {
		c.globalLoggers.Infof("Initialized LaunchDarkly client for %q", name)
	}
	if readyCh != nil {
		readyCh <- c
	}
}

func (c *envContextImpl) GetIdentifiers() EnvIdentifiers {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return c.identifiers
}

func (c *envContextImpl) SetIdentifiers(ei EnvIdentifiers) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.identifiers = ei
}

func (c *envContextImpl) GetCredentials() []config.SDKCredential {
	return c.getCredentialsInternal(true)
}

func (c *envContextImpl) GetDeprecatedCredentials() []config.SDKCredential {
	return c.getCredentialsInternal(false)
}

func (c *envContextImpl) getCredentialsInternal(preferred bool) []config.SDKCredential {
	c.mu.RLock()
	defer c.mu.RUnlock()

	ret := make([]config.SDKCredential, 0, len(c.credentials))
	for c, nonDeprecated := range c.credentials {
		if nonDeprecated == preferred {
			ret = append(ret, c)
		}
	}
	return ret
}

func (c *envContextImpl) AddCredential(newCredential config.SDKCredential) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, found := c.credentials[newCredential]; found {
		return
	}
	c.credentials[newCredential] = true
	c.envStreams.AddCredential(newCredential)
	for streamProvider, handlers := range c.handlers {
		if h := streamProvider.Handler(newCredential); h != nil {
			handlers[newCredential] = h
		}
	}

	// A new SDK key means 1. we should start a new SDK client, 2. we should tell all event forwarding
	// components that use an SDK key to use the new one. A new mobile key does not require starting a
	// new SDK client, but does requiring updating any event forwarding components that use a mobile key.
	switch key := newCredential.(type) {
	case config.SDKKey:
		go c.startSDKClient(key, nil, false)
		if c.metricsEventPub != nil { // metrics event publisher always uses SDK key
			c.metricsEventPub.ReplaceCredential(key)
		}
		if c.eventDispatcher != nil {
			c.eventDispatcher.ReplaceCredential(key)
		}
	case config.MobileKey:
		if c.eventDispatcher != nil {
			c.eventDispatcher.ReplaceCredential(key)
		}
	}
}

func (c *envContextImpl) RemoveCredential(oldCredential config.SDKCredential) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, found := c.credentials[oldCredential]; found {
		delete(c.credentials, oldCredential)
		c.envStreams.RemoveCredential(oldCredential)
		for _, handlers := range c.handlers {
			delete(handlers, oldCredential)
		}
		if sdkKey, ok := oldCredential.(config.SDKKey); ok {
			// The SDK client instance is tied to the SDK key, so get rid of it
			if client := c.clients[sdkKey]; client != nil {
				delete(c.clients, sdkKey)
				_ = client.Close()
			}
		}
	}
}

func (c *envContextImpl) DeprecateCredential(credential config.SDKCredential) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, found := c.credentials[credential]; found {
		c.credentials[credential] = false
	}
}

func (c *envContextImpl) GetClient() sdks.LDClientContext {
	c.mu.RLock()
	defer c.mu.RUnlock()
	// There might be multiple clients if there's an expiring SDK key. Find the SDK key that has a true
	// value in our map (meaning it's not deprecated) and return that client.
	for cred, valid := range c.credentials {
		if sdkKey, ok := cred.(config.SDKKey); ok && valid {
			return c.clients[sdkKey]
		}
	}
	return nil
}

func (c *envContextImpl) GetStore() interfaces.DataStore {
	return c.storeAdapter.GetStore()
}

func (c *envContextImpl) GetEvaluator() ldeval.Evaluator {
	c.mu.RLock()
	ret := c.evaluator
	c.mu.RUnlock()
	return ret
}

func (c *envContextImpl) GetBigSegmentStore() bigsegments.BigSegmentStore {
	return c.bigSegmentStore
}

func (c *envContextImpl) GetLoggers() ldlog.Loggers {
	return c.loggers
}

func (c *envContextImpl) GetStreamHandler(streamProvider streams.StreamProvider, credential config.SDKCredential) http.Handler {
	c.mu.RLock()
	defer c.mu.RUnlock()
	h := c.handlers[streamProvider][credential]
	if h == nil {
		return http.HandlerFunc(invalidStreamHandler)
	}
	return h
}

func invalidStreamHandler(w http.ResponseWriter, req *http.Request) {
	w.WriteHeader(http.StatusNotFound)
}

func (c *envContextImpl) GetEventDispatcher() *events.EventDispatcher {
	return c.eventDispatcher
}

func (c *envContextImpl) GetJSClientContext() JSClientContext {
	return c.jsContext
}

func (c *envContextImpl) GetMetricsContext() context.Context {
	if c.metricsEnv == nil {
		return context.Background()
	}
	return c.metricsEnv.GetOpenCensusContext()
}

func (c *envContextImpl) GetTTL() time.Duration {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return c.ttl
}

func (c *envContextImpl) SetTTL(newTTL time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.ttl = newTTL
}

func (c *envContextImpl) GetInitError() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.initErr
}

func (c *envContextImpl) IsSecureMode() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()

	return c.secureMode
}

func (c *envContextImpl) SetSecureMode(secureMode bool) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.secureMode = secureMode
}

func (c *envContextImpl) GetDataStoreInfo() sdks.DataStoreEnvironmentInfo {
	return c.dataStoreInfo
}

func (c *envContextImpl) GetCreationTime() time.Time {
	return c.creationTime
}

func (c *envContextImpl) FlushMetricsEvents() {
	if c.metricsEnv != nil && c.metricsEventPub != nil {
		c.metricsEnv.FlushEventsExporter()
		c.metricsEventPub.Flush()
	}
}

func (c *envContextImpl) Close() error {
	c.mu.Lock()
	for _, client := range c.clients {
		_ = client.Close()
	}
	c.clients = make(map[config.SDKKey]sdks.LDClientContext)
	c.mu.Unlock()
	_ = c.envStreams.Close()
	if c.metricsManager != nil && c.metricsEnv != nil {
		c.metricsManager.RemoveEnvironment(c.metricsEnv)
	}
	if c.metricsEventPub != nil {
		c.metricsEventPub.Close()
	}
	if c.eventDispatcher != nil {
		c.eventDispatcher.Close()
	}
	if c.bigSegmentSync != nil {
		c.bigSegmentSync.Close()
	}
	if c.bigSegmentStore != nil {
		_ = c.bigSegmentStore.Close()
	}
	if c.sdkBigSegments != nil {
		c.sdkBigSegments.Close()
	}
	return nil
}

func (q envContextStoreQueries) IsInitialized() bool {
	if s := q.context.storeAdapter.GetStore(); s != nil {
		return s.IsInitialized()
	}
	return false
}

func (q envContextStoreQueries) GetAll(kind ldstoretypes.DataKind) ([]ldstoretypes.KeyedItemDescriptor, error) {
	if s := q.context.storeAdapter.GetStore(); s != nil {
		return s.GetAll(kind)
	}
	return nil, nil
}

func (u *envContextStreamUpdates) SendAllDataUpdate(allData []ldstoretypes.Collection) {
	// We use this delegator, rather than sending updates directory to context.envStreams, so that we
	// can detect the presence of a big segment and turn on the big segment synchronizer as needed.
	u.context.envStreams.SendAllDataUpdate(allData)
	if u.context.bigSegmentSync == nil {
		return
	}
	hasBigSegment := false
	for _, coll := range allData {
		if coll.Kind == ldstoreimpl.Segments() {
			for _, keyedItem := range coll.Items {
				if s, ok := keyedItem.Item.Item.(*ldmodel.Segment); ok && s.Unbounded {
					hasBigSegment = true
					break
				}
			}
		}
	}
	if hasBigSegment {
		u.context.bigSegmentSync.Start() // has no effect if already started
	}
}

func (u *envContextStreamUpdates) SendSingleItemUpdate(kind ldstoretypes.DataKind, key string, item ldstoretypes.ItemDescriptor) {
	// See comments in SendAllDataUpdate.
	u.context.envStreams.SendSingleItemUpdate(kind, key, item)
	if u.context.bigSegmentSync == nil {
		return
	}
	hasBigSegment := false
	if kind == ldstoreimpl.Segments() {
		if s, ok := item.Item.(*ldmodel.Segment); ok && s.Unbounded {
			hasBigSegment = true
		}
	}
	if hasBigSegment {
		u.context.bigSegmentSync.Start()                // has no effect if already started
		u.context.sdkBigSegments.SetPollingActive(true) // has no effect if already active
	}
}

func (u *envContextStreamUpdates) InvalidateClientSideState() {
	u.context.envStreams.InvalidateClientSideState()
}

func makeLogPrefix(logNameMode LogNameMode, sdkKey config.SDKKey, envID config.EnvironmentID) string {
	name := string(sdkKey)
	if logNameMode == LogNameIsEnvID && envID != "" {
		name = string(envID)
	}
	if len(name) > 4 { // real keys are always longer than this
		name = "..." + name[len(name)-4:]
	}
	return fmt.Sprintf("[env: %s]", name)
}


<|start_filename|>internal/core/streams/stream_provider_client_side_ping_test.go<|end_filename|>
package streams

import (
	"testing"
	"time"

	"github.com/launchdarkly/ld-relay/v6/internal/basictypes"
	"github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk-evaluation.v1/ldmodel"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents/ldstoreimpl"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// The only difference between the mobile ping stream and the JS client ping stream is which kind of
// authorization credential they support.

func TestStreamProviderMobilePing(t *testing.T) {
	validCredential := testMobileKey
	invalidCredential1 := testSDKKey
	invalidCredential2 := testEnvID

	withStreamProvider := func(t *testing.T, maxConnTime time.Duration, action func(StreamProvider)) {
		sp := NewStreamProvider(basictypes.MobilePingStream, maxConnTime)
		require.NotNil(t, sp)
		defer sp.Close()
		action(sp)
	}

	t.Run("constructor", func(t *testing.T) {
		maxConnTime := time.Hour
		withStreamProvider(t, maxConnTime, func(sp StreamProvider) {
			require.IsType(t, &clientSidePingStreamProvider{}, sp)
			assert.False(t, sp.(*clientSidePingStreamProvider).isJSClient)
			verifyServerProperties(t, sp.(*clientSidePingStreamProvider).server, maxConnTime)
		})
	})

	t.Run("Handler", func(t *testing.T) {
		withStreamProvider(t, 0, func(sp StreamProvider) {
			assert.NotNil(t, sp.Handler(validCredential))
			assert.Nil(t, sp.Handler(invalidCredential1))
			assert.Nil(t, sp.Handler(invalidCredential2))
		})
	})

	t.Run("Register", func(t *testing.T) {
		store := makeMockStore(nil, nil)
		withStreamProvider(t, 0, func(sp StreamProvider) {
			assert.Nil(t, sp.Register(invalidCredential1, store, ldlog.NewDisabledLoggers()))
			assert.Nil(t, sp.Register(invalidCredential2, store, ldlog.NewDisabledLoggers()))

			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()
			require.IsType(t, &clientSidePingEnvStreamProvider{}, esp)
		})
	})
}

func TestStreamProviderJSClientPing(t *testing.T) {
	validCredential := testEnvID
	invalidCredential1 := testSDKKey
	invalidCredential2 := testMobileKey

	withStreamProvider := func(t *testing.T, maxConnTime time.Duration, action func(StreamProvider)) {
		sp := NewStreamProvider(basictypes.JSClientPingStream, maxConnTime)
		require.NotNil(t, sp)
		defer sp.Close()
		action(sp)
	}

	t.Run("constructor", func(t *testing.T) {
		maxConnTime := time.Hour
		withStreamProvider(t, maxConnTime, func(sp StreamProvider) {
			require.IsType(t, &clientSidePingStreamProvider{}, sp)
			assert.True(t, sp.(*clientSidePingStreamProvider).isJSClient)
			verifyServerProperties(t, sp.(*clientSidePingStreamProvider).server, maxConnTime)
		})
	})

	t.Run("Handler", func(t *testing.T) {
		withStreamProvider(t, 0, func(sp StreamProvider) {
			assert.NotNil(t, sp.Handler(validCredential))
			assert.Nil(t, sp.Handler(invalidCredential1))
			assert.Nil(t, sp.Handler(invalidCredential2))
		})
	})

	t.Run("Register", func(t *testing.T) {
		store := makeMockStore(nil, nil)
		withStreamProvider(t, 0, func(sp StreamProvider) {
			assert.Nil(t, sp.Register(invalidCredential1, store, ldlog.NewDisabledLoggers()))
			assert.Nil(t, sp.Register(invalidCredential2, store, ldlog.NewDisabledLoggers()))

			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()
			require.IsType(t, &clientSidePingEnvStreamProvider{}, esp)
		})
	})
}

func TestStreamProviderAllClientSidePing(t *testing.T) {
	// This uses only the mobile ping stream to test the event behavior, because we are using the same
	// implementation type for both mobile and JS client and we've already tested the individual
	// constructors above.

	validCredential := testMobileKey
	withStreamProvider := func(t *testing.T, maxConnTime time.Duration, action func(StreamProvider)) {
		sp := NewStreamProvider(basictypes.MobilePingStream, maxConnTime)
		require.NotNil(t, sp)
		defer sp.Close()
		action(sp)
	}

	t.Run("initial event", func(t *testing.T) {
		store := makeMockStore([]ldmodel.FeatureFlag{testFlag1, testFlag2}, []ldmodel.Segment{testSegment1})

		withStreamProvider(t, 0, func(sp StreamProvider) {
			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()

			verifyHandlerInitialEvent(t, sp, validCredential, MakePingEvent())
		})
	})

	t.Run("initial event - store not initialized", func(t *testing.T) {
		store := makeMockStore([]ldmodel.FeatureFlag{testFlag1, testFlag2}, []ldmodel.Segment{testSegment1})
		store.initialized = false

		withStreamProvider(t, 0, func(sp StreamProvider) {
			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()

			verifyHandlerInitialEvent(t, sp, validCredential, nil)
		})
	})

	t.Run("SendAllDataUpdate", func(t *testing.T) {
		store := makeMockStore(nil, nil)

		withStreamProvider(t, 0, func(sp StreamProvider) {
			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()

			newData := []ldstoretypes.Collection{
				{Kind: ldstoreimpl.Features(), Items: store.flags},
				{Kind: ldstoreimpl.Segments(), Items: store.segments},
			}

			verifyHandlerUpdateEvent(t, sp, validCredential, MakePingEvent(),
				func() {
					esp.SendAllDataUpdate(newData)
				},
				MakePingEvent(),
			)
		})
	})

	t.Run("SendSingleItemUpdate", func(t *testing.T) {
		store := makeMockStore(nil, nil)

		withStreamProvider(t, 0, func(sp StreamProvider) {
			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()

			verifyHandlerUpdateEvent(t, sp, validCredential, MakePingEvent(),
				func() {
					esp.SendSingleItemUpdate(ldstoreimpl.Features(), testFlag1.Key, sharedtest.FlagDesc(testFlag1))
				},
				MakePingEvent(),
			)

			verifyHandlerUpdateEvent(t, sp, validCredential, MakePingEvent(),
				func() {
					esp.SendSingleItemUpdate(ldstoreimpl.Segments(), testSegment1.Key, sharedtest.SegmentDesc(testSegment1))
				},
				MakePingEvent(),
			)
		})
	})

	t.Run("Heartbeat", func(t *testing.T) {
		store := makeMockStore(nil, nil)

		withStreamProvider(t, 0, func(sp StreamProvider) {
			esp := sp.Register(validCredential, store, ldlog.NewDisabledLoggers())
			require.NotNil(t, esp)
			defer esp.Close()

			verifyHandlerHeartbeat(t, sp, esp, validCredential)
		})
	})
}


<|start_filename|>internal/core/bigsegments/store_dynamodb_test.go<|end_filename|>
// +build big_segment_external_store_tests

package bigsegments

import (
	"fmt"
	"math"
	"strings"
	"testing"
	"time"

	ct "github.com/launchdarkly/go-configtypes"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/stretchr/testify/require"
)

const (
	testTableName = "LD_DYNAMODB_TEST_TABLE"
	localEndpoint = "http://localhost:8000"
)

var (
	testDynamoDBConfig = aws.Config{
		Region:      aws.String("us-west-2"),
		Endpoint:    aws.String(localEndpoint),
		Credentials: credentials.NewStaticCredentials("dummy", "not", "used"),
	}
)

func TestDynamoDBGenericAll(t *testing.T) {
	require.NoError(t, createTableIfNecessary())

	t.Run("without prefix", func(t *testing.T) { testGenericAll(t, withDynamoDBStoreGeneric("")) })
	t.Run("with prefix", func(t *testing.T) { testGenericAll(t, withDynamoDBStoreGeneric("testprefix")) })
}

func (store *dynamoDBBigSegmentStore) checkSetIncludes(attribute, segmentKey, userKey string) (bool, error) {
	bigSegmentsUserDataKeyWithPrefix := dynamoDBUserDataKey(store.prefix)

	result, err := store.client.GetItem(&dynamodb.GetItemInput{
		TableName:      aws.String(store.table),
		ConsistentRead: aws.Bool(true),
		Key: map[string]*dynamodb.AttributeValue{
			tablePartitionKey: {S: aws.String(bigSegmentsUserDataKeyWithPrefix)},
			tableSortKey:      {S: aws.String(userKey)},
		},
	})
	if err != nil || len(result.Item) == 0 {
		return false, err
	}
	item := result.Item[attribute]
	if item == nil || item.SS == nil {
		return false, nil
	}
	for _, v := range item.SS {
		if v == nil {
			continue
		}
		if *v == segmentKey {
			return true, nil
		}
	}
	return false, nil
}

func dynamoDBMakeOperations(store *dynamoDBBigSegmentStore) bigSegmentOperations {
	return bigSegmentOperations{
		isUserIncluded: func(segmentKey string, userKey string) (bool, error) {
			return store.checkSetIncludes(dynamoDBIncludedAttr, segmentKey, userKey)
		},
		isUserExcluded: func(segmentKey string, userKey string) (bool, error) {
			return store.checkSetIncludes(dynamoDBExcludedAttr, segmentKey, userKey)
		},
	}
}

func withDynamoDBStoreGeneric(prefix string) func(*testing.T, func(BigSegmentStore, bigSegmentOperations)) {
	return func(t *testing.T, action func(BigSegmentStore, bigSegmentOperations)) {
		require.NoError(t, clearTestData(prefix))
		url, err := ct.NewOptURLAbsoluteFromString(localEndpoint)
		require.NoError(t, err)
		store, err := newDynamoDBBigSegmentStore(url, testDynamoDBConfig, ldlog.NewDisabledLoggers(), testTableName, prefix)
		require.NoError(t, err)
		require.NotNil(t, store)
		defer store.Close()
		action(store, dynamoDBMakeOperations(store))
	}
}

func createTestClient() (*dynamodb.DynamoDB, error) {
	sess, err := session.NewSession(&testDynamoDBConfig)
	if err != nil {
		return nil, err
	}

	return dynamodb.New(sess), nil
}

func clearTestData(prefix string) error {
	if prefix != "" {
		prefix += ":"
	}

	client, err := createTestClient()
	if err != nil {
		return err
	}
	var items []map[string]*dynamodb.AttributeValue

	err = client.ScanPages(&dynamodb.ScanInput{
		TableName:            aws.String(testTableName),
		ConsistentRead:       aws.Bool(true),
		ProjectionExpression: aws.String("#namespace, #key"),
		ExpressionAttributeNames: map[string]*string{
			"#namespace": aws.String(tablePartitionKey),
			"#key":       aws.String(tableSortKey),
		},
	}, func(out *dynamodb.ScanOutput, lastPage bool) bool {
		items = append(items, out.Items...)
		return !lastPage
	})
	if err != nil {
		return err
	}

	var requests []*dynamodb.WriteRequest
	for _, item := range items {
		if strings.HasPrefix(*item[tablePartitionKey].S, prefix) {
			requests = append(requests, &dynamodb.WriteRequest{
				DeleteRequest: &dynamodb.DeleteRequest{Key: item},
			})
		}
	}
	return batchWriteRequests(client, testTableName, requests)
}

func createTableIfNecessary() error {
	client, err := createTestClient()
	if err != nil {
		return err
	}
	_, err = client.DescribeTable(&dynamodb.DescribeTableInput{TableName: aws.String(testTableName)})
	if err == nil {
		return nil
	}
	if e, ok := err.(awserr.Error); !ok || e.Code() != dynamodb.ErrCodeResourceNotFoundException {
		return err
	}
	createParams := dynamodb.CreateTableInput{
		AttributeDefinitions: []*dynamodb.AttributeDefinition{
			&dynamodb.AttributeDefinition{
				AttributeName: aws.String(tablePartitionKey),
				AttributeType: aws.String(dynamodb.ScalarAttributeTypeS),
			},
			&dynamodb.AttributeDefinition{
				AttributeName: aws.String(tableSortKey),
				AttributeType: aws.String(dynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*dynamodb.KeySchemaElement{
			&dynamodb.KeySchemaElement{
				AttributeName: aws.String(tablePartitionKey),
				KeyType:       aws.String(dynamodb.KeyTypeHash),
			},
			&dynamodb.KeySchemaElement{
				AttributeName: aws.String(tableSortKey),
				KeyType:       aws.String(dynamodb.KeyTypeRange),
			},
		},
		ProvisionedThroughput: &dynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(1),
			WriteCapacityUnits: aws.Int64(1),
		},
		TableName: aws.String(testTableName),
	}
	_, err = client.CreateTable(&createParams)
	if err != nil {
		return err
	}
	// When DynamoDB creates a table, it may not be ready to use immediately
	deadline := time.After(10 * time.Second)
	retry := time.NewTicker(100 * time.Millisecond)
	defer retry.Stop()
	for {
		select {
		case <-deadline:
			return fmt.Errorf("timed out waiting for new table to be ready")
		case <-retry.C:
			tableInfo, err := client.DescribeTable(&dynamodb.DescribeTableInput{TableName: aws.String(testTableName)})
			if err == nil && *tableInfo.Table.TableStatus == dynamodb.TableStatusActive {
				return nil
			}
		}
	}
}

// batchWriteRequests executes a list of write requests (PutItem or DeleteItem)
// in batches of 25, which is the maximum BatchWriteItem can handle.
func batchWriteRequests(
	client dynamodbiface.DynamoDBAPI,
	table string,
	requests []*dynamodb.WriteRequest,
) error {
	for len(requests) > 0 {
		batchSize := int(math.Min(float64(len(requests)), 25))
		batch := requests[:batchSize]
		requests = requests[batchSize:]

		_, err := client.BatchWriteItem(&dynamodb.BatchWriteItemInput{
			RequestItems: map[string][]*dynamodb.WriteRequest{table: batch},
		})
		if err != nil {
			// COVERAGE: can't simulate this condition in unit tests because we will only get this
			// far if the initial query in Init() already succeeded, and we don't have the ability
			// to make DynamoDB fail *selectively* within a single test
			return err
		}
	}
	return nil
}


<|start_filename|>internal/core/streams/stream_provider_client_side_ping.go<|end_filename|>
package streams

import (
	"net/http"
	"sync"

	"github.com/launchdarkly/ld-relay/v6/config"

	"github.com/launchdarkly/eventsource"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces/ldstoretypes"
)

// This is the standard implementation of a stream for client-side/mobile SDKs that sends only "ping" events,
// and does not do flag evaluations for specific users. The behavior of this stream is that it sends one "ping"
// event on initial connection, and another "ping" every time there is a data update of any kind.

type clientSidePingStreamProvider struct {
	server     *eventsource.Server
	isJSClient bool
	closeOnce  sync.Once
}

type clientSidePingEnvStreamProvider struct {
	server   *eventsource.Server
	channels []string
}

type clientSidePingEnvStreamRepository struct {
	store EnvStoreQueries
}

func (s *clientSidePingStreamProvider) validateCredential(credential config.SDKCredential) string {
	if s.isJSClient {
		if key, ok := credential.(config.EnvironmentID); ok {
			return string(key)
		}
	} else {
		if key, ok := credential.(config.MobileKey); ok {
			return string(key)
		}
	}
	return ""
}

func (s *clientSidePingStreamProvider) Handler(credential config.SDKCredential) http.HandlerFunc {
	if key := s.validateCredential(credential); key != "" {
		return s.server.Handler(key)
	}
	return nil
}

func (s *clientSidePingStreamProvider) Register(
	credential config.SDKCredential,
	store EnvStoreQueries,
	loggers ldlog.Loggers,
) EnvStreamProvider {
	if key := s.validateCredential(credential); key != "" {
		repo := &clientSidePingEnvStreamRepository{store: store}
		s.server.Register(key, repo)
		envStream := &clientSidePingEnvStreamProvider{server: s.server, channels: []string{key}}
		return envStream
	}
	return nil
}

func (s *clientSidePingStreamProvider) Close() {
	s.closeOnce.Do(func() {
		s.server.Close()
	})
}

func (e *clientSidePingEnvStreamProvider) SendAllDataUpdate(allData []ldstoretypes.Collection) {
	e.server.Publish(e.channels, MakePingEvent())
}

func (e *clientSidePingEnvStreamProvider) SendSingleItemUpdate(kind ldstoretypes.DataKind, key string, item ldstoretypes.ItemDescriptor) {
	e.server.Publish(e.channels, MakePingEvent())
}

func (e *clientSidePingEnvStreamProvider) InvalidateClientSideState() {
	e.server.Publish(e.channels, MakePingEvent())
}

func (e *clientSidePingEnvStreamProvider) SendHeartbeat() {
	e.server.PublishComment(e.channels, "")
}

func (e *clientSidePingEnvStreamProvider) Close() {
	for _, key := range e.channels {
		e.server.Unregister(key, true)
	}
}

func (r *clientSidePingEnvStreamRepository) Replay(channel, id string) chan eventsource.Event {
	out := make(chan eventsource.Event, 1)
	if r.store.IsInitialized() { // See serverSideEnvStreamRepository.Replay
		out <- MakePingEvent()
	}
	close(out)
	return out
}


<|start_filename|>internal/core/relay_core_end_to_end_redis_test.go<|end_filename|>
// +build redis_unit_tests

package core

// Continuation of relay_core_end_to_end_test.go that includes persistent storage behavior. A Redis server
// must be running on localhost for these tests.

import (
	"net/http"
	"testing"
	"time"

	c "github.com/launchdarkly/ld-relay/v6/config"
	st "github.com/launchdarkly/ld-relay/v6/internal/core/sharedtest"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"

	"github.com/launchdarkly/go-configtypes"
	"github.com/launchdarkly/go-test-helpers/v2/ldservices"
)

var basicRedisConfig = c.RedisConfig{Host: "localhost", LocalTTL: configtypes.NewOptDuration(time.Minute)}
var uncachedRedisConfig = c.RedisConfig{Host: "localhost", LocalTTL: configtypes.NewOptDuration(0)}

func TestRelayCoreEndToEndRedisSuccessWithCache(t *testing.T) {
	putEvent := ldservices.NewServerSDKData().Flags(&testFlag).ToPutEvent()
	streamHandler, _ := ldservices.ServerSideStreamingServiceHandler(putEvent)
	testEnv := st.EnvWithAllCredentials

	config := c.Config{Environment: st.MakeEnvConfigs(testEnv), Redis: basicRedisConfig}
	relayCoreEndToEndTest(t, config, streamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForSuccessfulInit()
		p.expectSuccessFromAllEndpoints(testEnv)
	})
}

func TestRelayCoreEndToEndRedisSuccessWithoutCache(t *testing.T) {
	// Turning off the cache isn't something that would be done in normal usage, but it lets us verify
	// that Relay will read flags from the database as needed when servicing requests.
	putEvent := ldservices.NewServerSDKData().Flags(&testFlag).ToPutEvent()
	streamHandler, _ := ldservices.ServerSideStreamingServiceHandler(putEvent)
	testEnv := st.EnvWithAllCredentials

	config := c.Config{Environment: st.MakeEnvConfigs(testEnv), Redis: uncachedRedisConfig}
	relayCoreEndToEndTest(t, config, streamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForSuccessfulInit()
		p.expectSuccessFromAllEndpoints(testEnv)
	})
}

func TestRelayCoreEndToEndRedisInitTimeoutWithInitializedDataStore(t *testing.T) {
	putEvent := ldservices.NewServerSDKData().Flags(&testFlag).ToPutEvent()
	streamHandler, _ := ldservices.ServerSideStreamingServiceHandler(putEvent)
	hangingHandler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		<-req.Context().Done() // hang until the request is cancelled by the client or the server
	})
	testEnv := st.EnvWithAllCredentials

	// First, run Relay with a successful connection to fake-LD, just to populate the database.
	preliminaryConfig := c.Config{Environment: st.MakeEnvConfigs(testEnv), Redis: basicRedisConfig}
	relayCoreEndToEndTest(t, preliminaryConfig, streamHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForSuccessfulInit()
		p.expectSuccessFromAllEndpoints(testEnv)
	})

	// Now, run Relay again against a fake-LD endpoint that hangs without returning any data, and a
	// short initTimeout. Clients should receive the data that's in the database from the previous run.
	config := c.Config{
		Main: c.MainConfig{
			InitTimeout: configtypes.NewOptDuration(time.Millisecond),
		},
		Environment: st.MakeEnvConfigs(testEnv),
		Redis:       basicRedisConfig,
	}
	relayCoreEndToEndTest(t, config, hangingHandler, func(p relayCoreEndToEndTestParams) {
		p.waitForLogMessage(ldlog.Error, "timeout encountered waiting for LaunchDarkly client initialization",
			"initialization timeout")
		p.expectSuccessFromAllEndpoints(testEnv)
	})
}


<|start_filename|>internal/basictypes/stream_kinds.go<|end_filename|>
package basictypes

// StreamKind is an enum-like type of the supported kinds of SDK streaming endpoint. This is different
// from basictypes.SDKKind because there are more kinds of streams than there are kinds of SDKs (server-side
// SDKs have two streams).
type StreamKind string

const (
	// ServerSideStream represents the server-side SDK "/all" endpoint, which is used by all current
	// server-side SDK versions.
	ServerSideStream StreamKind = "server"

	// ServerSideFlagsOnlyStream represents the server-side SDK "/flags" endpoint, which is used only
	// by old SDKs that do not support segments.
	ServerSideFlagsOnlyStream StreamKind = "server-flags"

	// MobilePingStream represents the mobile streaming endpoints, which will generate only "ping" events.
	MobilePingStream StreamKind = "mobile-ping"

	// JSClientPingStream represents the JS client-side streaming endpoints, which will generate only
	// "ping" events. This is identical to MobilePingStream except that it only handles requests
	// authenticated with an environment ID.
	JSClientPingStream StreamKind = "js-ping"
)


<|start_filename|>integrationtests/api_helpers_test.go<|end_filename|>
// +build integrationtests

package integrationtests

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/launchdarkly/ld-relay/v6/config"

	ldapi "github.com/launchdarkly/api-client-go"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldtime"
	"gopkg.in/launchdarkly/go-sdk-common.v2/ldvalue"

	"github.com/antihax/optional"
	"github.com/pborman/uuid"
)

type apiHelper struct {
	apiClient  *ldapi.APIClient
	apiContext context.Context
	apiBaseURL string
	loggers    ldlog.Loggers
}

func (a *apiHelper) logResult(desc string, err error) error {
	if err == nil {
		a.loggers.Infof("%s: OK", desc)
		return nil
	}
	addInfo := ""
	if gse, ok := err.(ldapi.GenericSwaggerError); ok {
		body := string(gse.Body())
		addInfo = " - " + string(body)
	}
	a.loggers.Errorf("%s: FAILED - %s%s", desc, err, addInfo)
	return err
}

func (a *apiHelper) createProjectsAndEnvironments(numProjects, numEnvironments int) (projsAndEnvs, error) {
	ret := make(projsAndEnvs)
	for i := 0; i < numProjects; i++ {
		proj, envs, err := a.createProject(numEnvironments)
		if err != nil {
			_ = a.deleteProjects(ret)
			return nil, err
		}
		ret[proj] = envs
	}
	return ret, nil
}

func (a *apiHelper) deleteProjects(projsAndEnvs projsAndEnvs) error {
	for p := range projsAndEnvs {
		if err := a.deleteProject(p); err != nil {
			return err
		}
	}
	return nil
}

func (a *apiHelper) createProject(numEnvironments int) (projectInfo, []environmentInfo, error) {
	projKey := randomApiKey("relayi9n-")
	projName := projKey
	projectBody := ldapi.ProjectBody{
		Name: projName,
		Key:  projKey,
	}
	for i := 0; i < numEnvironments; i++ {
		envKey := randomApiKey("env-")
		envName := envKey
		projectBody.Environments = append(projectBody.Environments, ldapi.EnvironmentPost{
			Name:  envName,
			Key:   envKey,
			Color: "000000",
		})
	}
	project, _, err := a.apiClient.ProjectsApi.PostProject(a.apiContext, projectBody)
	if err != nil {
		return projectInfo{}, nil, a.logResult("Create project", err)
	}
	var envInfos []environmentInfo
	for _, env := range project.Environments {
		envInfos = append(envInfos, environmentInfo{
			id:        config.EnvironmentID(env.Id),
			key:       env.Key,
			name:      env.Name,
			sdkKey:    config.SDKKey(env.ApiKey),
			mobileKey: config.MobileKey(env.MobileKey),
		})
	}
	a.loggers.Infof("Created project %q\n", projKey)
	return projectInfo{key: projKey, name: projName}, envInfos, nil
}

func randomApiKey(prefix string) string {
	return (prefix + uuid.New())[0:20]
}

func (a *apiHelper) deleteProject(project projectInfo) error {
	_, err := a.apiClient.ProjectsApi.DeleteProject(a.apiContext, project.key)
	return a.logResult(fmt.Sprintf("Delete project %q", project.key), err)
}

func (a *apiHelper) addEnvironment(project projectInfo) (environmentInfo, error) {
	envKey := randomApiKey("env-")
	envName := envKey
	envBody := ldapi.EnvironmentPost{
		Key:   envKey,
		Name:  envName,
		Color: "000000",
	}
	env, _, err := a.apiClient.EnvironmentsApi.PostEnvironment(a.apiContext, project.key, envBody)
	if err != nil {
		return environmentInfo{}, a.logResult("Create environment", err)
	}
	a.loggers.Infof("created environment %q\n", envKey)
	return environmentInfo{
		id:        config.EnvironmentID(env.Id),
		key:       env.Key,
		name:      env.Name,
		sdkKey:    config.SDKKey(env.ApiKey),
		mobileKey: config.MobileKey(env.MobileKey),
	}, nil
}

func (a *apiHelper) deleteEnvironment(project projectInfo, env environmentInfo) error {
	_, err := a.apiClient.EnvironmentsApi.DeleteEnvironment(a.apiContext, project.key, env.key)
	return a.logResult(fmt.Sprintf("Delete environment %q", env.key), err)
}

func (a *apiHelper) rotateSDKKey(project projectInfo, env environmentInfo, expirationTime time.Time) (
	config.SDKKey, error) {
	var apiOptions *ldapi.EnvironmentsApiResetEnvironmentSDKKeyOpts
	if !expirationTime.IsZero() {
		apiOptions = &ldapi.EnvironmentsApiResetEnvironmentSDKKeyOpts{Expiry: optional.NewInt64(int64(ldtime.UnixMillisFromTime(expirationTime)))}
	}
	envResult, _, err := a.apiClient.EnvironmentsApi.ResetEnvironmentSDKKey(a.apiContext, project.key, env.key, apiOptions)
	var newKey config.SDKKey
	if err == nil {
		newKey = config.SDKKey(envResult.ApiKey)
	}
	return newKey, a.logResult(fmt.Sprintf("Change SDK key for environment %q", env.key), err)
}

func (a *apiHelper) rotateMobileKey(project projectInfo, env environmentInfo) (config.MobileKey, error) {
	envResult, _, err := a.apiClient.EnvironmentsApi.ResetEnvironmentMobileKey(a.apiContext, project.key, env.key, nil)
	var newKey config.MobileKey
	if err == nil {
		newKey = config.MobileKey(envResult.MobileKey)
	}
	return newKey, a.logResult(fmt.Sprintf("Change mobile key for environment %q", env.key), err)
}

func (a *apiHelper) createAutoConfigKey(policyResources []string) (autoConfigID, config.AutoConfigKey, error) {
	body := ldapi.RelayProxyConfigBody{
		Name: uuid.New(),
		Policy: []ldapi.Policy{
			{
				Resources: policyResources,
				Actions:   []string{"*"},
				Effect:    "allow",
			},
		},
	}
	entity, _, err := a.apiClient.RelayProxyConfigurationsApi.PostRelayAutoConfig(a.apiContext, body)
	return autoConfigID(entity.Id), config.AutoConfigKey(entity.FullKey), a.logResult("Create auto-config key", err)
}

func (a *apiHelper) updateAutoConfigPolicy(id autoConfigID, newPolicyResources []string) error {
	var patchValue interface{} = newPolicyResources
	patchOps := []ldapi.PatchOperation{
		{Op: "replace", Path: "/policy/0/resources", Value: &patchValue},
	}
	_, _, err := a.apiClient.RelayProxyConfigurationsApi.PatchRelayProxyConfig(a.apiContext, string(id), patchOps)
	return a.logResult("Update auto-config policy", err)
}

func (a *apiHelper) deleteAutoConfigKey(id autoConfigID) error {
	_, err := a.apiClient.RelayProxyConfigurationsApi.DeleteRelayProxyConfig(a.apiContext, string(id))
	return a.logResult("Delete auto-config key", err)
}

// createFlag creates a flag in the specified project, and configures each environment to return a specific
// value for that flag in that environment which we'll check for later in verifyFlagValues.
func (a *apiHelper) createFlag(
	proj projectInfo,
	envs []environmentInfo,
	flagKey string,
	valueForEnv func(environmentInfo) ldvalue.Value,
) error {
	flagPost := ldapi.FeatureFlagBody{
		Name: flagKey,
		Key:  flagKey,
	}
	for _, env := range envs {
		valueAsInterface := valueForEnv(env).AsArbitraryValue()
		flagPost.Variations = append(flagPost.Variations, ldapi.Variation{Value: &valueAsInterface})
	}

	_, _, err := a.apiClient.FeatureFlagsApi.PostFeatureFlag(a.apiContext, proj.key, flagPost, nil)
	err = a.logResult("Create flag "+flagKey+" in "+proj.key, err)
	if err != nil {
		return err
	}

	for i, env := range envs {
		envPrefix := fmt.Sprintf("/environments/%s", env.key)
		patches := []ldapi.PatchOperation{
			makePatch("replace", envPrefix+"/offVariation", i),
		}
		_, _, err = a.apiClient.FeatureFlagsApi.PatchFeatureFlag(a.apiContext, proj.key, flagKey,
			ldapi.PatchComment{Patch: patches})
		err = a.logResult("Configure flag "+flagKey+" for "+env.key, err)
		if err != nil {
			return err
		}
	}

	return nil
}

func (a *apiHelper) createFlags(projsAndEnvs projsAndEnvs) error {
	for proj, envs := range projsAndEnvs {
		err := a.createFlag(proj, envs, flagKeyForProj(proj), flagValueForEnv)
		if err != nil {
			return err
		}
	}
	return nil
}

func (a *apiHelper) createBooleanFlagThatUsesSegment(
	flagKey string,
	project projectInfo,
	envs []environmentInfo,
	segmentKey string,
) error {
	flagPost := ldapi.FeatureFlagBody{
		Name: flagKey,
		Key:  flagKey,
	}
	for _, value := range []bool{true, false} {
		var valueAsInterface interface{} = value
		flagPost.Variations = append(flagPost.Variations, ldapi.Variation{Value: &valueAsInterface})
	}

	_, _, err := a.apiClient.FeatureFlagsApi.PostFeatureFlag(a.apiContext, project.key, flagPost, nil)
	err = a.logResult("Create flag "+flagKey+" in "+project.key, err)
	if err != nil {
		return err
	}

	for _, env := range envs {
		envPrefix := fmt.Sprintf("/environments/%s", env.key)
		rulesJSON := fmt.Sprintf(`[
			{"variation": 0, "clauses": [{"attribute": "", "op": "segmentMatch", "values":["%s"]}]}
		]`, segmentKey)
		patches := []ldapi.PatchOperation{
			makePatch("replace", envPrefix+"/on", true),
			makePatch("replace", envPrefix+"/rules", json.RawMessage(rulesJSON)),
			makePatch("replace", envPrefix+"/fallthrough/variation", 1),
		}
		_, _, err = a.apiClient.FeatureFlagsApi.PatchFeatureFlag(a.apiContext, project.key, flagKey,
			ldapi.PatchComment{Patch: patches})
		err = a.logResult("Configure flag "+flagKey+" for "+env.key, err)
		if err != nil {
			return err
		}
	}
	return nil
}

func (a *apiHelper) createBigSegment(
	project projectInfo,
	env environmentInfo,
	segmentKey string,
	userKeysIncluded []string,
	userKeysExcluded []string,
	userKeysIncludedViaRule []string,
) error {
	userSegmentBody := ldapi.UserSegmentBody{
		Key:       segmentKey,
		Name:      segmentKey,
		Unbounded: true,
	}
	_, _, err := a.apiClient.UserSegmentsApi.PostUserSegment(a.apiContext, project.key, env.key, userSegmentBody)
	if err = a.logResult("Create "+segmentKey, err); err != nil {
		return err
	}

	if len(userKeysIncludedViaRule) > 0 {
		rulesJSON := fmt.Sprintf(`[
			{"clauses": [{"attribute": "key", "op": "in", "values":%v}]}
		]`, ldvalue.CopyArbitraryValue(userKeysIncludedViaRule))
		patches := []ldapi.PatchOperation{
			makePatch("replace", "/rules", json.RawMessage(rulesJSON)),
		}
		_, _, err := a.apiClient.UserSegmentsApi.PatchUserSegment(a.apiContext, project.key, env.key,
			segmentKey, patches)
		if err = a.logResult("Add rule to "+segmentKey, err); err != nil {
			return err
		}
	}

	updates := ldapi.BigSegmentTargetsBody{
		Included: &ldapi.BigSegmentTargetChanges{Add: userKeysIncluded},
		Excluded: &ldapi.BigSegmentTargetChanges{Add: userKeysExcluded},
	}
	_, err = a.apiClient.UserSegmentsApi.UpdatedBigSegmentTargets(a.apiContext, project.key, env.key,
		segmentKey, updates)
	return a.logResult("Update "+segmentKey, err)
}

func makePatch(op, path string, value interface{}) ldapi.PatchOperation {
	return ldapi.PatchOperation{Op: op, Path: path, Value: &value}
}


<|start_filename|>internal/core/sdks/big_segments.go<|end_filename|>
package sdks

import (
	"github.com/launchdarkly/ld-relay/v6/config"

	"gopkg.in/launchdarkly/go-sdk-common.v2/ldlog"
	"gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"
	"gopkg.in/launchdarkly/go-server-sdk.v5/ldcomponents"
)

// ConfigureBigSegments provides the appropriate Go SDK big segments configuration based on the Relay
// configuration, or nil if big segments are not enabled. The big segments stores in Relay's SDK
// instances are used for client-side evaluations; server-side SDKs will read from the same database
// via their own big segments stores, which will need to be configured similarly to what's here.
func ConfigureBigSegments(
	allConfig config.Config,
	envConfig config.EnvConfig,
	loggers ldlog.Loggers,
) (interfaces.BigSegmentsConfigurationFactory, error) {
	var storeFactory interfaces.BigSegmentStoreFactory

	if allConfig.Redis.URL.IsDefined() {
		redisBuilder, redisURL := makeRedisDataStoreBuilder(allConfig, envConfig)
		loggers.Infof("Using Redis big segment store: %s with prefix: %s", redisURL, envConfig.Prefix)
		storeFactory = redisBuilder
	} else if allConfig.DynamoDB.Enabled {
		dynamoDBBuilder, tableName, err := makeDynamoDBDataStoreBuilder(allConfig, envConfig)
		if err != nil {
			return nil, err
		}
		loggers.Infof("Using DynamoDB big segment store: %s with prefix: %s", tableName, envConfig.Prefix)
		storeFactory = dynamoDBBuilder
	}

	if storeFactory != nil {
		return ldcomponents.BigSegments(storeFactory), nil
	}
	return nil, nil
}


<|start_filename|>internal/core/sharedtest/data_source_helpers.go<|end_filename|>
package sharedtest

import "gopkg.in/launchdarkly/go-server-sdk.v5/interfaces"

type ExistingDataSourceFactory struct{ Instance interfaces.DataSource }
type DataSourceThatNeverStarts struct{}
type DataSourceThatStartsWithoutInitializing struct{}

func (f ExistingDataSourceFactory) CreateDataSource(
	context interfaces.ClientContext,
	updates interfaces.DataSourceUpdates,
) (interfaces.DataSource, error) {
	return f.Instance, nil
}

func (d DataSourceThatNeverStarts) Start(chan<- struct{}) {}
func (d DataSourceThatNeverStarts) IsInitialized() bool   { return false }
func (d DataSourceThatNeverStarts) Close() error          { return nil }

func (d DataSourceThatStartsWithoutInitializing) Start(closeWhenReady chan<- struct{}) {
	go close(closeWhenReady)
}
func (d DataSourceThatStartsWithoutInitializing) IsInitialized() bool { return false }
func (d DataSourceThatStartsWithoutInitializing) Close() error        { return nil }
