From c760fe17c29462ee32eaa091a28bc361eba11107 Mon Sep 17 00:00:00 2001 From: Amin Assaid Date: Tue, 5 Nov 2024 05:06:00 -0500 Subject: [PATCH] Implemented Artifactory upload, download and delete handling - Added GRABIT_ARTIFACTORY_TOKEN support - Implemented cache download priority - Added fallback to source URL - Added validation at each step - Improved error handling and logging - Integrated cache upload functionality Co-authored-by: Amin Assaid Co-authored-by: Amy Druham --- cmd/add.go | 18 +++- cmd/add_test.go | 27 +++++- cmd/delete_test.go | 1 + internal/artifactory_cache_test.go | 139 +++++++++++++++++++++++++++++ internal/lock.go | 126 +++++++++++++++++++++++++- internal/lock_test.go | 4 +- internal/resource.go | 71 +++++++++++---- internal/resource_test.go | 5 +- 8 files changed, 367 insertions(+), 24 deletions(-) create mode 100644 internal/artifactory_cache_test.go diff --git a/cmd/add.go b/cmd/add.go index 28c3444..5448f76 100644 --- a/cmd/add.go +++ b/cmd/add.go @@ -4,6 +4,9 @@ package cmd import ( + "fmt" + "os" + "github.com/cisco-open/grabit/internal" "github.com/spf13/cobra" ) @@ -18,6 +21,7 @@ func addAdd(cmd *cobra.Command) { addCmd.Flags().String("algo", internal.RecommendedAlgo, "Integrity algorithm") addCmd.Flags().String("filename", "", "Target file name to use when downloading the resource") addCmd.Flags().StringArray("tag", []string{}, "Resource tags") + addCmd.Flags().String("cache", "", "Artifactory cache URL") cmd.AddCommand(addCmd) } @@ -26,6 +30,18 @@ func runAdd(cmd *cobra.Command, args []string) error { if err != nil { return err } + // Get cache URL + cacheURL, err := cmd.Flags().GetString("cache") + if err != nil { + return err + } + // Check token if cache is requested + if cacheURL != "" { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token == "" { + return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") + } + } lock, err := internal.NewLock(lockFile, true) if err != nil { return err @@ -42,7 +58,7 @@ func runAdd(cmd *cobra.Command, args []string) error { if err != nil { return err } - err = lock.AddResource(args, algo, tags, filename) + err = lock.AddResource(args, algo, tags, filename, cacheURL) if err != nil { return err } diff --git a/cmd/add_test.go b/cmd/add_test.go index a7ba460..8ec01b2 100644 --- a/cmd/add_test.go +++ b/cmd/add_test.go @@ -10,6 +10,10 @@ import ( ) func TestRunAdd(t *testing.T) { + // Set the GRABIT_ARTIFACTORY_TOKEN environment variable. + t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token") + + // Setup HTTP handler for the resource handler := func(w http.ResponseWriter, r *http.Request) { _, err := w.Write([]byte(`abcdef`)) if err != nil { @@ -18,8 +22,29 @@ func TestRunAdd(t *testing.T) { } port, server := test.HttpHandler(handler) defer server.Close() + + // Setup dummy cache server + cacheHandler := func(w http.ResponseWriter, r *http.Request) { + if r.Method == "PUT" { + w.WriteHeader(http.StatusCreated) + } + } + cachePort, cacheServer := test.HttpHandler(cacheHandler) + defer cacheServer.Close() + + // Create empty lockfile + lockFile := test.TmpFile(t, "") + cmd := NewRootCmd() - cmd.SetArgs([]string{"-f", test.TmpFile(t, ""), "add", fmt.Sprintf("http://localhost:%d/test.html", port)}) + // Add cache URL to the command + cacheURL := fmt.Sprintf("http://localhost:%d", cachePort) + cmd.SetArgs([]string{ + "-f", lockFile, + "add", + fmt.Sprintf("http://localhost:%d/test.html", port), + "--cache", cacheURL, + }) + err := cmd.Execute() assert.Nil(t, err) } diff --git a/cmd/delete_test.go b/cmd/delete_test.go index 7f312be..496820c 100644 --- a/cmd/delete_test.go +++ b/cmd/delete_test.go @@ -15,6 +15,7 @@ func TestRunDelete(t *testing.T) { Tags = ['tag1', 'tag2'] `) cmd := NewRootCmd() + cmd.Flags().String("cache", "", "Artifactory URL for caching") cmd.SetArgs([]string{"-f", testfilepath, "delete", "http://localhost:123456/test.html"}) err := cmd.Execute() assert.Nil(t, err) diff --git a/internal/artifactory_cache_test.go b/internal/artifactory_cache_test.go new file mode 100644 index 0000000..38fa910 --- /dev/null +++ b/internal/artifactory_cache_test.go @@ -0,0 +1,139 @@ +package internal + +import ( + "crypto/sha256" + "encoding/base64" + "fmt" + "net/http" + "testing" + + "github.com/cisco-open/grabit/test" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestAddWithArtifactoryCache verifies adding a resource with caching enabled. +func TestAddWithArtifactoryCache(t *testing.T) { + // Sub-test to verify behavior when the token is not set. + t.Run("TokenNotSet", func(t *testing.T) { + // Clear the GRABIT_ARTIFACTORY_TOKEN environment variable. + t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "") + + // Setup a simple HTTP handler that always returns "test content". + handler := func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`test content`)) + } + // Start the HTTP server and get the port it runs on. + port, server := test.HttpHandler(handler) + defer server.Close() // Ensure the server is stopped after the test. + + // Create a temporary lock file for testing. + path := test.TmpFile(t, "") + lock, err := NewLock(path, true) + require.NoError(t, err) // Fail the test if the lock cannot be created. + + // Set up URLs for the source file and cache. + sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port) + cacheURL := fmt.Sprintf("http://localhost:%d", port) + + // Attempt to add a resource to the lock file. + err = lock.AddResource([]string{sourceURL}, "sha256", []string{}, "", cacheURL) + // Verify that the error message indicates the token is not set. + assert.Contains(t, err.Error(), "GRABIT_ARTIFACTORY_TOKEN environment variable is not set") + }) +} + +// TestDownloadWithArtifactoryCache verifies downloading resources with caching enabled. +func TestDownloadWithArtifactoryCache(t *testing.T) { + // Sub-test to verify behavior when NO_CACHE_UPLOAD is set. + t.Run("NO_CACHE_UPLOAD", func(t *testing.T) { + // Set the NO_CACHE_UPLOAD environment variable. + t.Setenv("NO_CACHE_UPLOAD", "1") + + // Prepare the content and calculate its hash. + testContent := []byte("test content") + hash := sha256.Sum256(testContent) + expectedHash := "sha256-" + base64.StdEncoding.EncodeToString(hash[:]) + + // Track if an upload is attempted. + uploadAttempted := false + // Setup an HTTP handler to serve the content and log upload attempts. + handler := func(w http.ResponseWriter, r *http.Request) { + if r.Method == "PUT" { // Check for upload attempts. + uploadAttempted = true + t.Error("Should not attempt upload when NO_CACHE_UPLOAD is set") + } + w.Write(testContent) + } + // Start the HTTP server and get the port it runs on. + port, server := test.HttpHandler(handler) + defer server.Close() // Ensure the server is stopped after the test. + + // Create a lock file with the resource and cache information. + lockContent := fmt.Sprintf(`[[Resource]] + Urls = ['http://localhost:%d/test.txt'] + Integrity = '%s' + CacheUri = 'http://localhost:%d/cache'`, port, expectedHash, port) + + lockPath := test.TmpFile(t, lockContent) + lock, err := NewLock(lockPath, false) + require.NoError(t, err) // Fail the test if the lock cannot be created. + + // Download the resource into a temporary directory. + tmpDir := test.TmpDir(t) + err = lock.Download(tmpDir, []string{}, []string{}, "", false) + assert.NoError(t, err) // Verify the download succeeded. + + // Ensure no upload was attempted. + assert.False(t, uploadAttempted) + }) +} + +// TestDeleteWithArtifactoryCache verifies deleting a resource with caching enabled. +func TestDeleteWithArtifactoryCache(t *testing.T) { + // Sub-test to verify successful deletion of a resource. + t.Run("SuccessfulDelete", func(t *testing.T) { + // Set the GRABIT_ARTIFACTORY_TOKEN environment variable. + t.Setenv("GRABIT_ARTIFACTORY_TOKEN", "test-token") + + // Setup an HTTP handler to handle DELETE requests. + handler := func(w http.ResponseWriter, r *http.Request) { + if r.Method == "DELETE" { // Respond with OK for DELETE requests. + w.WriteHeader(http.StatusOK) + } + } + // Start the HTTP server and get the port it runs on. + port, server := test.HttpHandler(handler) + defer server.Close() // Ensure the server is stopped after the test. + + // Set up URLs for the source file and cache. + sourceURL := fmt.Sprintf("http://localhost:%d/test.txt", port) + cacheURL := fmt.Sprintf("http://localhost:%d", port) + + // Create a lock file with the resource and cache information. + lockContent := fmt.Sprintf(`[[Resource]] + Urls = ['%s'] + Integrity = 'sha256-test' + CacheUri = '%s'`, sourceURL, cacheURL) + + lockPath := test.TmpFile(t, lockContent) + lock, err := NewLock(lockPath, false) + require.NoError(t, err) // Fail the test if the lock cannot be created. + + // Save the lock file before modifying it. + err = lock.Save() + require.NoError(t, err) // Fail the test if saving fails. + + // Delete the resource from the lock file. + lock.DeleteResource(sourceURL) + + // Save the lock file again after deletion. + err = lock.Save() + require.NoError(t, err) // Fail the test if saving fails. + + // Reload the lock file and verify the resource is gone. + newLock, err := NewLock(lockPath, false) + require.NoError(t, err) + assert.Equal(t, 0, len(newLock.conf.Resource)) // Ensure no resources remain. + }) +} diff --git a/internal/lock.go b/internal/lock.go index f525507..9f47acb 100644 --- a/internal/lock.go +++ b/internal/lock.go @@ -8,14 +8,26 @@ import ( "context" "errors" "fmt" + "net/http" "os" + "path" + "path/filepath" "strconv" + "strings" + "github.com/carlmjohnson/requests" toml "github.com/pelletier/go-toml/v2" + "github.com/rs/zerolog/log" ) var COMMENT_PREFIX = "//" +// getArtifactoryURL constructs the URL for an artifact in Artifactory +func getArtifactoryURL(baseURL, integrity string) string { + return fmt.Sprintf("%s/%s", baseURL, integrity) + +} + // Lock represents a grabit lockfile. type Lock struct { path string @@ -49,30 +61,99 @@ func NewLock(path string, newOk bool) (*Lock, error) { return &Lock{path: path, conf: conf}, nil } -func (l *Lock) AddResource(paths []string, algo string, tags []string, filename string) error { +func (l *Lock) AddResource(paths []string, algo string, tags []string, filename string, cacheURL string) error { for _, u := range paths { if l.Contains(u) { return fmt.Errorf("resource '%s' is already present", u) } } - r, err := NewResourceFromUrl(paths, algo, tags, filename) + r, err := NewResourceFromUrl(paths, algo, tags, filename, cacheURL) if err != nil { return err } + + // If cache URL is provided, handles Artifactory upload + if cacheURL != "" { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token == "" { + return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") + } + + // Add context here + ctx := context.Background() + path, err := GetUrltoTempFile(paths[0], token, ctx) + if err != nil { + return fmt.Errorf("failed to get file for cache: %s", err) + } + defer os.Remove(path) + + // Upload to Artifactory using hash as filename + err = uploadToArtifactory(path, cacheURL, r.Integrity) + if err != nil { + return fmt.Errorf("failed to upload to cache: %v", err) + } + } + l.conf.Resource = append(l.conf.Resource, *r) return nil } +func uploadToArtifactory(filePath, cacheURL, integrity string) error { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token == "" { + return fmt.Errorf("GRABIT_ARTIFACTORY_TOKEN environment variable is not set") + } + + // Use the integrity value directly for the URL + artifactoryURL := getArtifactoryURL(cacheURL, integrity) + + // Upload the file using the requests package + err := requests. + URL(artifactoryURL). + Method(http.MethodPut). + Header("Authorization", fmt.Sprintf("Bearer %s", token)). + BodyFile(filePath). // Using BodyFile instead of ReadFile + Fetch(context.Background()) + + if err != nil { + return fmt.Errorf("upload failed: %w", err) + } + + return nil +} func (l *Lock) DeleteResource(path string) { newStatements := []Resource{} for _, r := range l.conf.Resource { if !r.Contains(path) { newStatements = append(newStatements, r) + } else if r.Contains(path) && r.CacheUri != "" { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token == "" { + log.Warn().Msg("Warning: Unable to delete from Artifactory: GRABIT_ARTIFACTORY_TOKEN not set.") + + continue + } + + artifactoryURL := getArtifactoryURL(r.CacheUri, r.Integrity) + + err := deleteCache(artifactoryURL, token) + if err != nil { + log.Warn().Msg("Warning: Unable to delete from Artifactory") + } } } l.conf.Resource = newStatements } +func deleteCache(url, token string) error { + // Create and send a DELETE request with an Authorization header. + return requests. + URL(url). + Method(http.MethodDelete). + Header("Authorization", fmt.Sprintf("Bearer %s", token)). + Fetch(context.Background()) +} + const NoFileMode = os.FileMode(0) // strToFileMode converts a string to a os.FileMode. @@ -165,6 +246,47 @@ func (l *Lock) Download(dir string, tags []string, notags []string, perm string, for i, r := range filteredResources { resource := r go func() { + // Try Artifactory first if available + if resource.CacheUri != "" { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token != "" { + artifactoryURL := getArtifactoryURL(resource.CacheUri, resource.Integrity) + filename := resource.Filename + if filename == "" { + filename = path.Base(resource.Urls[0]) + } + fullPath := filepath.Join(dir, filename) + + // Use getUrl with bearer token + tmpPath, err := getUrl(artifactoryURL, fullPath, token, ctx) + if err == nil { + // integrity check + algo, err := getAlgoFromIntegrity(resource.Integrity) + if err != nil { + errorCh <- err + return + } + err = checkIntegrityFromFile(tmpPath, algo, resource.Integrity, artifactoryURL) + if err != nil { + errorCh <- err + return + } + if mode != NoFileMode { + err = os.Chmod(tmpPath, mode.Perm()) + } + if err == nil { + errorCh <- nil + if statusLine != nil { + statusLine.Increment(i) + } + return + } + } + if strings.Contains(err.Error(), "lookup invalid") || strings.Contains(err.Error(), "dial tcp") { + fmt.Printf("Failed to download from Artifactory, falling back to original URL: %v\n", err) + } + } + } err := resource.Download(dir, mode, ctx) errorCh <- err diff --git a/internal/lock_test.go b/internal/lock_test.go index eb6c810..f7ed791 100644 --- a/internal/lock_test.go +++ b/internal/lock_test.go @@ -51,7 +51,7 @@ func TestLockManipulations(t *testing.T) { port, server := test.HttpHandler(handler) defer server.Close() resource := fmt.Sprintf("http://localhost:%d/test2.html", port) - err = lock.AddResource([]string{resource}, "sha512", []string{}, "") + err = lock.AddResource([]string{resource}, "sha512", []string{}, "", "") assert.Nil(t, err) assert.Equal(t, 2, len(lock.conf.Resource)) err = lock.Save() @@ -68,7 +68,7 @@ func TestDuplicateResource(t *testing.T) { Integrity = 'sha256-asdasdasd'`, url)) lock, err := NewLock(path, false) assert.Nil(t, err) - err = lock.AddResource([]string{url}, "sha512", []string{}, "") + err = lock.AddResource([]string{url}, "sha512", []string{}, "", "") assert.NotNil(t, err) assert.Contains(t, err.Error(), "already present") } diff --git a/internal/resource.go b/internal/resource.go index a377e87..070941b 100644 --- a/internal/resource.go +++ b/internal/resource.go @@ -24,15 +24,16 @@ type Resource struct { Integrity string Tags []string `toml:",omitempty"` Filename string `toml:",omitempty"` + CacheUri string `toml:",omitempty"` } -func NewResourceFromUrl(urls []string, algo string, tags []string, filename string) (*Resource, error) { +func NewResourceFromUrl(urls []string, algo string, tags []string, filename string, cacheURL string) (*Resource, error) { if len(urls) < 1 { return nil, fmt.Errorf("empty url list") } url := urls[0] ctx := context.Background() - path, err := GetUrltoTempFile(url, ctx) + path, err := GetUrltoTempFile(url, "", ctx) if err != nil { return nil, fmt.Errorf("failed to get url: %s", err) } @@ -41,48 +42,79 @@ func NewResourceFromUrl(urls []string, algo string, tags []string, filename stri if err != nil { return nil, fmt.Errorf("failed to compute ressource integrity: %s", err) } - return &Resource{Urls: urls, Integrity: integrity, Tags: tags, Filename: filename}, nil + + return &Resource{Urls: urls, Integrity: integrity, Tags: tags, Filename: filename, CacheUri: cacheURL}, nil } // getUrl downloads the given resource and returns the path to it. -func getUrl(u string, fileName string, ctx context.Context) (string, error) { +func getUrl(u string, fileName string, bearer string, ctx context.Context) (string, error) { _, err := url.Parse(u) if err != nil { return "", fmt.Errorf("invalid url '%s': %s", u, err) } log.Debug().Str("URL", u).Msg("Downloading") - err = requests. + + req := requests. URL(u). Header("Accept", "*/*"). - ToFile(fileName). - Fetch(ctx) + ToFile(fileName) + + if bearer != "" { + req.Header("Authorization", fmt.Sprintf("Bearer %s", bearer)) + } + + err = req.Fetch(ctx) if err != nil { return "", fmt.Errorf("failed to download '%s': %s", u, err) } - log.Debug().Str("URL", u).Msg("Downloaded") + return fileName, nil } // GetUrlToDir downloads the given resource to the given directory and returns the path to it. -func GetUrlToDir(u string, targetDir string, ctx context.Context) (string, error) { +func GetUrlToDir(u string, targetDir string, bearer string, ctx context.Context) (string, error) { // create temporary name in the target directory. h := sha256.New() h.Write([]byte(u)) fileName := filepath.Join(targetDir, fmt.Sprintf(".%s", hex.EncodeToString(h.Sum(nil)))) - return getUrl(u, fileName, ctx) + return getUrl(u, fileName, bearer, ctx) } // GetUrlWithDir downloads the given resource to a temporary file and returns the path to it. -func GetUrltoTempFile(u string, ctx context.Context) (string, error) { +func GetUrltoTempFile(u string, bearer string, ctx context.Context) (string, error) { file, err := os.CreateTemp("", "prefix") if err != nil { log.Fatal().Err(err) } fileName := file.Name() - return getUrl(u, fileName, ctx) + return getUrl(u, fileName, "", ctx) } func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) error { + // Check if a cache URL exists to use Artifactory first + if l.CacheUri != "" { + token := os.Getenv("GRABIT_ARTIFACTORY_TOKEN") + if token != "" { + artifactoryURL := fmt.Sprintf("%s/%s", l.CacheUri, l.Integrity) + localName := l.Filename + if localName == "" { + localName = path.Base(l.Urls[0]) + } + resPath := filepath.Join(dir, localName) + + // Use getUrl directly with bearer token for Artifactory + tmpPath, err := getUrl(artifactoryURL, resPath, token, ctx) + if err == nil { + if mode != NoFileMode { + err = os.Chmod(tmpPath, mode.Perm()) + } + if err == nil { + return nil // Success + } + } + fmt.Printf("Failed to download from Artifactory, falling back to original URL: %v\n", err) + } + } ok := false algo, err := getAlgoFromIntegrity(l.Integrity) if err != nil { @@ -92,7 +124,7 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e for _, u := range l.Urls { // Download file in the target directory so that the call to // os.Rename is atomic. - lpath, err := GetUrlToDir(u, dir, ctx) + lpath, err := GetUrlToDir(u, dir, "", ctx) if err != nil { downloadError = err continue @@ -120,10 +152,19 @@ func (l *Resource) Download(dir string, mode os.FileMode, ctx context.Context) e } } ok = true + if l.CacheUri != "" && os.Getenv("NO_CACHE_UPLOAD") != "1" { + if uploadErr := uploadToArtifactory(resPath, l.CacheUri, l.Integrity); uploadErr != nil { + fmt.Printf("Warning: Failed to upload to cache: %v\n", uploadErr) + } + } } if !ok { - if downloadError != nil { - return downloadError + if err == nil { + if downloadError != nil { + return downloadError + } else { + panic("no error but no file downloaded") + } } return err } diff --git a/internal/resource_test.go b/internal/resource_test.go index 80a11e9..ea0f8f0 100644 --- a/internal/resource_test.go +++ b/internal/resource_test.go @@ -31,7 +31,7 @@ func TestNewResourceFromUrl(t *testing.T) { { urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, valid: true, - res: Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, Integrity: fmt.Sprintf("%s-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=", algo), Tags: []string{}, Filename: ""}, + res: Resource{Urls: []string{fmt.Sprintf("http://localhost:%d/test.html", port)}, Integrity: fmt.Sprintf("%s-vvV+x/U6bUC+tkCngKY5yDvCmsipgW8fxsXG3Nk8RyE=", algo), Tags: []string{}, Filename: "", CacheUri: ""}, }, { urls: []string{"invalid url"}, @@ -39,9 +39,8 @@ func TestNewResourceFromUrl(t *testing.T) { errorContains: "failed to download", }, } - for _, data := range tests { - resource, err := NewResourceFromUrl(data.urls, algo, []string{}, "") + resource, err := NewResourceFromUrl(data.urls, algo, []string{}, "", "") assert.Equal(t, data.valid, err == nil) if err != nil { assert.Contains(t, err.Error(), data.errorContains)