backup: accept canonicalized snapshot metadata uri

Signed-off-by: huanghaoyuanhhy <haoyuan.huang@zilliz.com>

Author: huanghaoyuanhhy

core/backup/snapshot_target.go

@@ -2,6 +2,8 @@ package backup
 
 import (
 	"fmt"
+	"net"
+	"net/url"
 	"path"
 	"strings"
 
@@ -52,10 +54,137 @@ func newSnapshotTarget(milvusCfg, backupCfg storage.Config, backupDir string) (s
 // the backup directory, which is what the backup meta records: an absolute uri would pin
 // the backup to the bucket and prefix it was written to.
 func (t snapshotTarget) metadataPath(metadataURI string) (string, error) {
-	root := strings.TrimSuffix(t.Path, "/") + "/"
-	if !strings.HasPrefix(metadataURI, root) {
+	target, err := parseSnapshotURI(t.Path)
+	if err != nil {
+		return "", fmt.Errorf("backup: parse snapshot target %s: %w", t.Path, err)
+	}
+	metadata, err := parseSnapshotURI(metadataURI)
+	if err != nil {
+		return "", fmt.Errorf("backup: parse exported metadata %s: %w", metadataURI, err)
+	}
+
+	if !target.sameStorage(metadata) {
 		return "", fmt.Errorf("backup: exported metadata %s is not under the target %s", metadataURI, t.Path)
 	}
 
-	return path.Join(t.Dir, strings.TrimPrefix(metadataURI, root)), nil
+	root := strings.TrimSuffix(target.key, "/") + "/"
+	if !strings.HasPrefix(metadata.key, root) {
+		return "", fmt.Errorf("backup: exported metadata %s is not under the target %s", metadataURI, t.Path)
+	}
+
+	return path.Join(t.Dir, strings.TrimPrefix(metadata.key, root)), nil
+}
+
+// snapshotURI is the storage identity and object key encoded by one of the URI
+// forms Milvus accepts. Endpoint-style forms put the bucket in the first path
+// segment, while provider-style forms put it in the host.
+type snapshotURI struct {
+	scheme        string
+	endpoint      string
+	bucket        string
+	key           string
+	endpointStyle bool
+}
+
+func parseSnapshotURI(raw string) (snapshotURI, error) {
+	u, err := url.Parse(raw)
+	if err != nil {
+		return snapshotURI{}, err
+	}
+	if u.Scheme == "" || u.Host == "" {
+		return snapshotURI{}, fmt.Errorf("uri needs a scheme and host")
+	}
+	if u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" {
+		return snapshotURI{}, fmt.Errorf("uri must not contain credentials, query parameters, or a fragment")
+	}
+
+	objectPath, err := url.PathUnescape(u.EscapedPath())
+	if err != nil {
+		return snapshotURI{}, fmt.Errorf("unescape object path: %w", err)
+	}
+	objectPath, err = cleanSnapshotObjectKey(objectPath)
+	if err != nil {
+		return snapshotURI{}, err
+	}
+
+	scheme := strings.ToLower(u.Scheme)
+	switch scheme {
+	case "minio", "http", "https", "az", "azure":
+		parts := strings.SplitN(objectPath, "/", 2)
+		if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
+			return snapshotURI{}, fmt.Errorf("endpoint-style uri needs a bucket and object key")
+		}
+		endpoint, err := snapshotEndpoint(u, scheme)
+		if err != nil {
+			return snapshotURI{}, err
+		}
+		return snapshotURI{
+			scheme:        scheme,
+			endpoint:      endpoint,
+			bucket:        parts[0],
+			key:           parts[1],
+			endpointStyle: true,
+		}, nil
+	case "s3", "gs", "gcs":
+		if u.Port() != "" {
+			return snapshotURI{}, fmt.Errorf("provider-style uri bucket must not contain a port")
+		}
+		return snapshotURI{
+			scheme: scheme,
+			bucket: u.Host,
+			key:    objectPath,
+		}, nil
+	default:
+		return snapshotURI{}, fmt.Errorf("unsupported snapshot uri scheme %q", scheme)
+	}
+}
+
+func cleanSnapshotObjectKey(objectPath string) (string, error) {
+	objectPath = strings.Trim(objectPath, "/")
+	if objectPath == "" {
+		return "", fmt.Errorf("uri needs an object key")
+	}
+	for _, part := range strings.Split(objectPath, "/") {
+		if part == "." || part == ".." {
+			return "", fmt.Errorf("object key must not contain path traversal")
+		}
+	}
+	return path.Clean(objectPath), nil
+}
+
+func snapshotEndpoint(u *url.URL, scheme string) (string, error) {
+	host := strings.ToLower(strings.TrimSuffix(u.Hostname(), "."))
+	if host == "" {
+		return "", fmt.Errorf("uri needs an endpoint host")
+	}
+	port := u.Port()
+	if port == "" {
+		switch scheme {
+		case "https":
+			port = "443"
+		case "http":
+			port = "80"
+		}
+	}
+	if port == "" {
+		return host, nil
+	}
+	return net.JoinHostPort(host, port), nil
+}
+
+func (u snapshotURI) sameStorage(other snapshotURI) bool {
+	if u.endpointStyle != other.endpointStyle || u.bucket != other.bucket {
+		return false
+	}
+	if u.endpointStyle {
+		return u.endpoint == other.endpoint
+	}
+	return canonicalSnapshotScheme(u.scheme) == canonicalSnapshotScheme(other.scheme)
+}
+
+func canonicalSnapshotScheme(scheme string) string {
+	if scheme == "gcs" {
+		return "gs"
+	}
+	return scheme
 }

core/backup/snapshot_target_test.go

@@ -53,10 +53,37 @@ func TestSnapshotTarget_MetadataPath(t *testing.T) {
 		assert.Equal(t, "bundle/snapshots/449577/metadata/449580.json", got)
 	})
 
+	// Milvus may report an endpoint-style URI using its transport scheme and omit
+	// the default port. That is the same storage location as the minio URI sent in
+	// the export request and must not make an otherwise successful backup fail.
+	t.Run("CanonicalizedEndpointURI", func(t *testing.T) {
+		target := snapshotTarget{
+			Path: "minio://s3.us-west-2.amazonaws.com:443/backup-bucket/backup/mybackup/bundle",
+			Dir:  "bundle",
+		}
+		got, err := target.metadataPath("https://s3.us-west-2.amazonaws.com/backup-bucket/backup/mybackup/bundle/snapshots/449577/metadata/449580.json")
+		require.NoError(t, err)
+		assert.Equal(t, "bundle/snapshots/449577/metadata/449580.json", got)
+	})
+
 	// Anything outside the target is not ours to record, and a path relative to the
 	// backup directory could not describe it anyway.
 	t.Run("OutsideTarget", func(t *testing.T) {
 		_, err := target.metadataPath("s3://other-bucket/snapshots/449577/metadata/449580.json")
 		assert.Error(t, err)
 	})
+
+	t.Run("SiblingPrefixIsOutsideTarget", func(t *testing.T) {
+		_, err := target.metadataPath("s3://backup-bucket/backup/mybackup/bundle-other/snapshots/449577/metadata/449580.json")
+		assert.Error(t, err)
+	})
+
+	t.Run("DifferentEndpointIsOutsideTarget", func(t *testing.T) {
+		target := snapshotTarget{
+			Path: "minio://s3.us-west-2.amazonaws.com:443/backup-bucket/backup/mybackup/bundle",
+			Dir:  "bundle",
+		}
+		_, err := target.metadataPath("https://s3.us-east-1.amazonaws.com/backup-bucket/backup/mybackup/bundle/snapshots/449577/metadata/449580.json")
+		assert.Error(t, err)
+	})
 }