trufflehog/pkg/sources/github/map_installations_test.go

package github

import (
	"encoding/json"
	"testing"

	"github.com/google/go-github/v67/github"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/trufflesecurity/trufflehog/v3/pkg/cache/simple"
	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
)

// This is meant to track APIUnitReporter.UnitOK in scanner/cmd/scanner/pipeline.go
// precisely
func newMarshalledRepoUnit(t *testing.T, su sources.SourceUnit) []byte {
	t.Helper()

	id, kind := su.SourceUnitID()

	unitData, err := json.Marshal(su)
	require.NoError(t, err)

	out, err := json.Marshal(map[string]any{
		"id":        id,
		"kind":      string(kind),
		"display":   su.Display(),
		"unit_data": unitData,
	})
	require.NoError(t, err)

	return out
}

func TestGitHub_UnmarshalSourceUnit(t *testing.T) {
	s := &Source{}

	t.Run("source unit with unit_data recovers installation id", func(t *testing.T) {
		data := newMarshalledRepoUnit(t, RepoUnit{
			Name:           "spotify/backstage",
			URL:            "https://github.com/spotify/backstage.git",
			InstallationID: 2448,
		})

		unit, err := s.UnmarshalSourceUnit(data)
		require.NoError(t, err)

		repoUnit, ok := unit.(RepoUnit)
		require.True(t, ok, "expected a RepoUnit so ChunkUnit can read the installation id")
		assert.Equal(t, "https://github.com/spotify/backstage.git", repoUnit.URL)
		assert.Equal(t, "spotify/backstage", repoUnit.Name)
		assert.Equal(t, int64(2448), repoUnit.InstallationID)

		id, kind := unit.SourceUnitID()
		assert.Equal(t, "https://github.com/spotify/backstage.git", id)
		assert.EqualValues(t, "repo", kind)
	})

	t.Run("legacy source unit without unit_data has no installation id", func(t *testing.T) {
		data, err := json.Marshal(map[string]any{
			"id":      "https://github.com/acme/widgets.git",
			"kind":    "repo",
			"display": "acme/widgets",
		})
		require.NoError(t, err)

		unit, err := s.UnmarshalSourceUnit(data)
		require.NoError(t, err)

		repoUnit, ok := unit.(RepoUnit)
		require.True(t, ok)
		assert.Equal(t, "https://github.com/acme/widgets.git", repoUnit.URL)
		assert.Equal(t, "acme/widgets", repoUnit.Name)
		assert.Zero(t, repoUnit.InstallationID, "legacy units fall back to deriving the installation")
	})

	t.Run("full enumeration payload", func(t *testing.T) {
		data, err := json.Marshal(RepoUnit{
			Name:           "acme/widgets",
			URL:            "https://github.com/acme/widgets.git",
			InstallationID: 111,
		})
		require.NoError(t, err)

		unit, err := s.UnmarshalSourceUnit(data)
		require.NoError(t, err)

		repoUnit, ok := unit.(RepoUnit)
		require.True(t, ok)
		assert.Equal(t, int64(111), repoUnit.InstallationID)
	})

	t.Run("gist unit round trips", func(t *testing.T) {
		data := newMarshalledRepoUnit(t, GistUnit{Name: "abc", URL: "https://gist.github.com/abc.git"})

		unit, err := s.UnmarshalSourceUnit(data)
		require.NoError(t, err)

		_, ok := unit.(GistUnit)
		require.True(t, ok)
		_, kind := unit.SourceUnitID()
		assert.EqualValues(t, "gist", kind)
	})

	t.Run("rejects unrecognized payload", func(t *testing.T) {
		_, err := s.UnmarshalSourceUnit([]byte(`{"unrelated":"value"}`))
		require.Error(t, err)
	})
}

// Ensure we store installation IDs when enumerating
func TestGitHub_ProcessReposEmitsInstallationID(t *testing.T) {
	ctx := context.Background()

	connector := &appConnector{
		installationID:      1,
		repoInstallationMap: make(map[string]int64),
	}
	s := &Source{
		conn:              &sourcespb.GitHub{ScanAllInstallations: true},
		connector:         connector,
		filteredRepoCache: &filteredRepoCache{Cache: simple.NewCache[string]()},
		repoInfoCache:     newRepoInfoCache(),
	}

	const installID int64 = 2448
	repo := &github.Repository{
		Name:     github.String("backstage"),
		FullName: github.String("spotify/backstage"),
		CloneURL: github.String("https://github.com/spotify/backstage.git"),
		Owner:    &github.User{Login: github.String("spotify"), Type: github.String("Organization")},
	}

	// Mirror enumerateAllInstallationRepos: the lister records the installation
	// for each repo in a page before processRepos emits the units for it.
	page := 0
	listRepos := func(ctx context.Context, _ string, _ repoListOptions) ([]*github.Repository, *github.Response, error) {
		page++
		if page > 1 {
			return nil, &github.Response{}, nil
		}
		connector.setRepoInstallationForRepoName(repo.GetCloneURL(), repo.GetName(), installID)
		return []*github.Repository{repo}, &github.Response{}, nil
	}

	var units []sources.SourceUnit
	reporter := sources.VisitorReporter{
		VisitUnit: func(_ context.Context, u sources.SourceUnit) error {
			units = append(units, u)
			return nil
		},
	}

	err := s.processRepos(ctx, "spotify", reporter, listRepos, &appListOptions{
		ListOptions: github.ListOptions{PerPage: 100},
	})
	require.NoError(t, err)

	require.Len(t, units, 1)
	repoUnit, ok := units[0].(RepoUnit)
	require.True(t, ok)
	assert.Equal(t, "https://github.com/spotify/backstage.git", repoUnit.URL)
	assert.Equal(t, installID, repoUnit.InstallationID)
}