add: support for VPC endpoint in download from s3

skit-ai · Jan 5, 2024 · c18aa5b · c18aa5b
1 parent 5fb646a
commit c18aa5b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/aws/s3.go b/aws/s3.go
@@ -26,6 +26,7 @@ const (
 		"(([-a-z0-9]+)\\.)?" + // region name, optional for us-east-1
 		"vpce\\." +
 		"(amazonaws\\.com|c2s\\.ic\\.gov|sc2s\\.sgov\\.gov)"
+	vpceUrlPatternHostIdx   = 0
 	vpceUrlPatternBucketIdx = 2
 	vpceUrlPatternRegionIdx = 5
 
@@ -46,6 +47,7 @@ var (
 // S3URL holds interesting pieces after parsing a s3 URL
 type S3URL struct {
 	IsPathStyle bool
+	EndPoint    string
 	Bucket      string
 	Key         string
 	Region      string
@@ -54,7 +56,8 @@ type S3URL struct {
 // DownloadFile downloads a file from s3 based on the key and writes it into WriteAt.
 func (u S3URL) DownloadFile(ctx context.Context, w io.WriterAt) error {
 	sess, err := session.NewSession(&aws.Config{
-		Region: aws.String(u.Region), // Specify the region where the bucket is located
+		Region:   aws.String(u.Region), // Specify the region where the bucket is located
+		Endpoint: aws.String(u.EndPoint),
 	})
 	if err != nil {
 		return errors.NewError("Error creating session", err, false)
@@ -89,12 +92,9 @@ func (u S3URL) DownloadFile(ctx context.Context, w io.WriterAt) error {
 //
 //	https://s3.us-east-1.amazonaws.com/mybucket/a/b/c
 func ParseAmazonS3URL(s3URL *url.URL) (S3URL, error) {
-	output, err := parseBucketAndRegionFromHost(s3URL.Host, vpceUrlRegex, vpceUrlPatternBucketIdx, vpceUrlPatternRegionIdx)
+	output, err := parseBucketAndRegionFromHost(s3URL.Host)
 	if err != nil {
-		output, err = parseBucketAndRegionFromHost(s3URL.Host, nonVpceUrlRegex, nonVpceUrlPatternBucketIdx, nonVpceUrlPatternRegionIdx)
-		if err != nil {
-			return S3URL{}, err
-		}
+		return S3URL{}, errors.NewError("parsing host failed", err, false)
 	}
 
 	output.IsPathStyle = output.Bucket == ""
@@ -143,15 +143,24 @@ func ParseAmazonS3URL(s3URL *url.URL) (S3URL, error) {
 	return output, nil
 }
 
-func parseBucketAndRegionFromHost(host string, re *regexp.Regexp, bucketIdx, regionIdx int) (S3URL, error) {
-	result := re.FindStringSubmatch(host)
-	if result != nil && len(result) > bucketIdx && len(result) > regionIdx {
+func parseBucketAndRegionFromHost(host string) (S3URL, error) {
+	result := vpceUrlRegex.FindStringSubmatch(host)
+	if result != nil && len(result) > vpceUrlPatternBucketIdx && len(result) > vpceUrlPatternRegionIdx {
 		return S3URL{
-			Bucket: result[bucketIdx],
-			Region: result[regionIdx],
+			EndPoint: result[vpceUrlPatternHostIdx],
+			Bucket:   result[vpceUrlPatternBucketIdx],
+			Region:   result[vpceUrlPatternRegionIdx],
 		}, nil
 	} else {
-		return S3URL{}, errors.NewError("no match", nil, false)
+		result = nonVpceUrlRegex.FindStringSubmatch(host)
+		if result != nil && len(result) > vpceUrlPatternBucketIdx && len(result) > vpceUrlPatternRegionIdx {
+			return S3URL{
+				Bucket: result[nonVpceUrlPatternBucketIdx],
+				Region: result[nonVpceUrlPatternRegionIdx],
+			}, nil
+		} else {
+			return S3URL{}, errors.NewError("failed to match URL", nil, false)
+		}
 	}
 }
 

diff --git a/go.sum b/go.sum
@@ -904,6 +904,7 @@ golang.org/x/sys v0.3.0 h1:w8ZOecv6NaNa/zC8944JTU3vz4u6Lagfk4RPQxv92NQ=
 golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
 golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
 golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
@@ -922,6 +923,7 @@ golang.org/x/text v0.5.0 h1:OLmvp0KP+FVG99Ct/qFiL/Fhk4zp4QQnZ7b2U+5piUM=
 golang.org/x/text v0.5.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
 golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=