Fix path traversal issue with storage key

core/common/utils.go

@@ -17,9 +17,11 @@ package common
 import (
 	"context"
 	"crypto/md5"
+	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
-	"path/filepath"
+	"io/fs"
+	"net/url"
 	"sort"
 	"strings"
 
@@ -46,11 +48,17 @@ const (
 	DefaultMetadataMapChunkSize = 50_000
 )
 
-// ToShortRemote returns the short remote name given a git ssh remote string.
-// For example, "git@github:uber/tango" will return "uber/tango".
+// ToShortRemote returns the repository path from an SSH or URL remote.
+// For example, "git@github:uber/tango" and
+// "https://github.com/uber/tango" both return "uber/tango".
 func ToShortRemote(remote string) string {
-	strs := strings.Split(remote, ":")
-	return strs[len(strs)-1]
+	if parsed, err := url.Parse(remote); err == nil && parsed.Scheme != "" && parsed.Host != "" {
+		return strings.TrimLeft(parsed.Path, "/")
+	}
+	if idx := strings.LastIndex(remote, ":"); idx >= 0 {
+		remote = remote[idx+1:]
+	}
+	return strings.TrimLeft(remote, "/")
 }
 
 // GetGraphByTreeHash returns the cache path for the target graph by treehash.
@@ -59,23 +67,32 @@ func ToShortRemote(remote string) string {
 // requestOptions is folded into the key when any of its fields affect computation
 // (today: extra_exclude_files_regex). Empty/nil ⇒ legacy path unchanged.
 func GetGraphByTreeHash(remote, treehash string, strategy tangopb.ComputationStrategy, requestOptions *tangopb.RequestOptions) string {
-	path := filepath.Join(ToShortRemote(remote), "graphs", treehash, strategy.String())
+	key := strings.Join([]string{
+		safeCachePath(ToShortRemote(remote)),
+		"graphs",
+		safeCacheSegment(treehash),
+		strategy.String(),
+	}, "/")
 	if hash := HashRequestOptions(requestOptions); hash != "" {
-		path += "_requests-options-" + hash
+		key += "_requests-options-" + hash
 	}
-	return path
+	return key
 }
 
 // GetTreehashCachePath returns the cache path for the treehash mapping.
 // The git treehash is purely a function of git state (base SHA + applied
 // requests), so neither requestOptions nor the computation strategy is part
 // of this key.
 func GetTreehashCachePath(buildDescription *tangopb.BuildDescription) string {
-	path := filepath.Join(ToShortRemote(buildDescription.Remote), "treehashes", fmt.Sprintf("base-sha-%s", buildDescription.BaseSha))
+	key := strings.Join([]string{
+		safeCachePath(ToShortRemote(buildDescription.Remote)),
+		"treehashes",
+		"base-sha-" + safeCacheSegment(buildDescription.BaseSha),
+	}, "/")
 	if len(buildDescription.Requests) > 0 {
-		path += "_request-urls-" + GetReqURLsHash(buildDescription.Requests)
+		key += "_request-urls-" + GetReqURLsHash(buildDescription.Requests)
 	}
-	return path
+	return key
 }
 
 // GetComparedTargetsCachePath returns the cache path for a compared target graph result.
@@ -84,11 +101,36 @@ func GetTreehashCachePath(buildDescription *tangopb.BuildDescription) string {
 // requestOptions is folded into the key when any of its fields affect computation.
 // Empty/nil ⇒ legacy path unchanged.
 func GetComparedTargetsCachePath(remote, treehash1, treehash2 string, requestOptions *tangopb.RequestOptions) string {
-	path := filepath.Join(ToShortRemote(remote), "compared-targets", treehash1+"_"+treehash2)
+	key := strings.Join([]string{
+		safeCachePath(ToShortRemote(remote)),
+		"compared-targets",
+		safeCacheSegment(treehash1) + "_" + safeCacheSegment(treehash2),
+	}, "/")
 	if hash := HashRequestOptions(requestOptions); hash != "" {
-		path += "_requests-options-" + hash
+		key += "_requests-options-" + hash
+	}
+	return key
+}
+
+func safeCachePath(value string) string {
+	if value != "" && value != "." && fs.ValidPath(value) && !strings.ContainsAny(value, "\\:\x00") {
+		return value
 	}
-	return path
+	return unsafeCacheComponent(value)
+}
+
+func safeCacheSegment(value string) string {
+	if !strings.Contains(value, "/") {
+		if safe := safeCachePath(value); safe == value {
+			return value
+		}
+	}
+	return unsafeCacheComponent(value)
+}
+
+func unsafeCacheComponent(value string) string {
+	hash := sha256.Sum256([]byte(value))
+	return "unsafe-" + hex.EncodeToString(hash[:])
 }
 
 // GetReqURLsHash returns a fixed-length MD5 hash of the sorted request URLs.

core/common/utils_test.go

@@ -18,7 +18,9 @@ import (
 	"context"
 	"crypto/md5"
 	"fmt"
-	"path/filepath"
+	"io/fs"
+	"path"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -49,6 +51,16 @@ func TestToShortRemote(t *testing.T) {
 			remote: "git@github:org/project/sub",
 			want:   "org/project/sub",
 		},
+		{
+			name:   "HTTPS remote",
+			remote: "https://github.com/uber/tango.git",
+			want:   "uber/tango.git",
+		},
+		{
+			name:   "SSH URL remote",
+			remote: "ssh://git@github.com/uber/tango.git",
+			want:   "uber/tango.git",
+		},
 	}
 	for _, tt := range tests {
 		tt := tt
@@ -66,7 +78,7 @@ func TestGetGraphByTreeHash(t *testing.T) {
 
 	// Nil/empty options ⇒ no request-options suffix.
 	got := GetGraphByTreeHash(remote, treehash, strategy, nil)
-	assert.Equal(t, filepath.Join("uber/tango", "graphs", treehash, strategy.String()), got)
+	assert.Equal(t, path.Join("uber/tango", "graphs", treehash, strategy.String()), got)
 	assert.Equal(t, got, GetGraphByTreeHash(remote, treehash, strategy, &pb.RequestOptions{}))
 
 	// Different strategies ⇒ different keys.
@@ -99,7 +111,7 @@ func TestGetTreehashCachePath(t *testing.T) {
 	h := md5.New()
 	h.Write([]byte("custom://foo/bar"))
 	h.Write([]byte("github://org/repo/pull/1"))
-	want := filepath.Join("uber/tango", "treehashes", "base-sha-deadbeef") + "_request-urls-" + fmt.Sprintf("%x", h.Sum(nil))
+	want := path.Join("uber/tango", "treehashes", "base-sha-deadbeef") + "_request-urls-" + fmt.Sprintf("%x", h.Sum(nil))
 	assert.Equal(t, want, got)
 }
 
@@ -149,7 +161,7 @@ func TestGetReqURLsHash(t *testing.T) {
 func TestGetComparedTargetsCachePath(t *testing.T) {
 	t.Parallel()
 	got := GetComparedTargetsCachePath("git@github:uber/tango", "abc", "def", nil)
-	assert.Equal(t, filepath.Join("uber/tango", "compared-targets", "abc_def"), got)
+	assert.Equal(t, path.Join("uber/tango", "compared-targets", "abc_def"), got)
 
 	// Nil/empty options ⇒ legacy path.
 	assert.Equal(t, got, GetComparedTargetsCachePath("git@github:uber/tango", "abc", "def", &pb.RequestOptions{}))
@@ -158,6 +170,32 @@ func TestGetComparedTargetsCachePath(t *testing.T) {
 	assert.NotEqual(t, got, GetComparedTargetsCachePath("git@github:uber/tango", "abc", "def", &pb.RequestOptions{ExtraExcludeFilesRegex: []string{"foo.*"}}))
 }
 
+func TestCachePathsEncodeUnsafeComponents(t *testing.T) {
+	t.Parallel()
+
+	graphKey := GetGraphByTreeHash("../../outside", "../tree", pb.COMPUTATION_STRATEGY_NATIVE, nil)
+	assertSafeCacheKey(t, graphKey)
+	assert.Equal(t, graphKey, GetGraphByTreeHash("../../outside", "../tree", pb.COMPUTATION_STRATEGY_NATIVE, nil))
+
+	treehashKey := GetTreehashCachePath(&pb.BuildDescription{
+		Remote:  `..\..\outside`,
+		BaseSha: "../../base",
+	})
+	assertSafeCacheKey(t, treehashKey)
+
+	comparedKey := GetComparedTargetsCachePath("/outside", "../one", "two/../../three", nil)
+	assertSafeCacheKey(t, comparedKey)
+}
+
+func assertSafeCacheKey(t *testing.T, key string) {
+	t.Helper()
+	assert.True(t, fs.ValidPath(key), key)
+	assert.NotEqual(t, ".", key)
+	assert.NotContains(t, key, "\\")
+	assert.NotContains(t, key, ":")
+	assert.False(t, strings.HasPrefix(key, "/"), key)
+}
+
 func TestChunkTargets(t *testing.T) {
 	t.Parallel()
 

core/itg/cache/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_go//go:def.bzl", "go_library")
+load("@rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "cache",
@@ -10,3 +10,15 @@ go_library(
         "//core/storage",
     ],
 )
+
+go_test(
+    name = "cache_test",
+    srcs = ["cache_test.go"],
+    embed = [":cache"],
+    deps = [
+        "//core/itg/graph",
+        "//core/storage",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
+    ],
+)

core/itg/cache/cache.go

@@ -19,7 +19,8 @@ import (
 	"context"
 	"encoding/gob"
 	"fmt"
-	"path/filepath"
+	"io/fs"
+	"path"
 	"slices"
 	"strconv"
 	"strings"
@@ -64,10 +65,16 @@ var CompareKeyFunc = func(a Key, b Key) int {
 // EmptyKey means no cache found.
 var EmptyKey = Key{}
 
-// toStorageKey converts a cache key to its storage key: itg/{remote}/{date}/{committime}_{sha}
-func (k *Key) toStorageKey() string {
+// toStorageKey converts a cache key to its storage key: itg/{remote}/{date}/{committime}_{sha}.
+func (k *Key) toStorageKey() (string, error) {
+	if err := validateRemote(k.Remote); err != nil {
+		return "", err
+	}
+	if err := validateBaseSHA(k.BaseSha); err != nil {
+		return "", err
+	}
 	date := time.Unix(k.BaseCommitTimeSecond, 0).UTC().Format("2006-01-02")
-	return filepath.Join(keyPrefix, k.Remote, date, fmt.Sprintf("%d_%s", k.BaseCommitTimeSecond, k.BaseSha))
+	return path.Join(keyPrefix, k.Remote, date, fmt.Sprintf("%d_%s", k.BaseCommitTimeSecond, k.BaseSha)), nil
 }
 
 // NewStorageCache creates a new cache backed by a storage.Storage implementation.
@@ -82,7 +89,11 @@ type storageCache struct {
 }
 
 func (c *storageCache) Put(ctx context.Context, optimizedGraph *graph.OptimizedGraph, key Key) error {
-	exists, err := c.storage.Exists(ctx, key.toStorageKey())
+	storageKey, err := key.toStorageKey()
+	if err != nil {
+		return err
+	}
+	exists, err := c.storage.Exists(ctx, storageKey)
 	if err != nil {
 		return err
 	}
@@ -94,19 +105,23 @@ func (c *storageCache) Put(ctx context.Context, optimizedGraph *graph.OptimizedG
 	if err := gob.NewEncoder(&buf).Encode(optimizedGraph); err != nil {
 		return err
 	}
-	return c.storage.Put(ctx, storage.UploadRequest{Key: key.toStorageKey(), Reader: &buf})
+	return c.storage.Put(ctx, storage.UploadRequest{Key: storageKey, Reader: &buf})
 }
 
 func (c *storageCache) Get(ctx context.Context, key Key) (*graph.OptimizedGraph, error) {
-	resp, err := c.storage.Get(ctx, storage.DownloadRequest{Key: key.toStorageKey()})
+	storageKey, err := key.toStorageKey()
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.storage.Get(ctx, storage.DownloadRequest{Key: storageKey})
 	if err != nil {
-		return nil, fmt.Errorf("download graph %s: %w", key.toStorageKey(), err)
+		return nil, fmt.Errorf("download graph %s: %w", storageKey, err)
 	}
 	defer resp.ReadCloser.Close()
 
 	var optimizedGraph graph.OptimizedGraph
 	if err := gob.NewDecoder(resp.ReadCloser).Decode(&optimizedGraph); err != nil {
-		return nil, fmt.Errorf("decode graph %s: %w", key.toStorageKey(), err)
+		return nil, fmt.Errorf("decode graph %s: %w", storageKey, err)
 	}
 	for _, t := range optimizedGraph.OptimizedTargets {
 		if t.Hash == nil {
@@ -117,7 +132,10 @@ func (c *storageCache) Get(ctx context.Context, key Key) (*graph.OptimizedGraph,
 }
 
 func (c *storageCache) FloorKey(ctx context.Context, remote string, targetTimeSecond int64) (Key, error) {
-	remotePrefix := keyPrefix + remote + "/"
+	if err := validateRemote(remote); err != nil {
+		return EmptyKey, err
+	}
+	remotePrefix := path.Join(keyPrefix, remote) + "/"
 	allKeys, err := c.storage.List(ctx, remotePrefix)
 	if err != nil {
 		return EmptyKey, err
@@ -155,6 +173,20 @@ func (c *storageCache) FloorKey(ctx context.Context, remote string, targetTimeSe
 	return cacheKeys[idx], nil
 }
 
+func validateRemote(remote string) error {
+	if remote == "" || remote == "." || !fs.ValidPath(remote) || strings.ContainsAny(remote, "\\:\x00") {
+		return fmt.Errorf("invalid remote %q for cache key", remote)
+	}
+	return nil
+}
+
+func validateBaseSHA(baseSHA string) error {
+	if baseSHA == "" || strings.Contains(baseSHA, "/") || !fs.ValidPath(baseSHA) || strings.ContainsAny(baseSHA, "\\:\x00") {
+		return fmt.Errorf("invalid base SHA %q for cache key", baseSHA)
+	}
+	return nil
+}
+
 func binarySearch(cacheKeys []Key, targetTimeSecond int64) (int, bool) {
 	return slices.BinarySearchFunc(cacheKeys, targetTimeSecond, func(a Key, b int64) int {
 		switch {