gcs-s3-compat-client-with-sse-c.go

package gcs

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

// PresignedDownloadURL contains the details for downloading an object.
type PresignedDownloadURL struct {
	Method    string      `json:"method"`
	URL       string      `json:"url"`
	Headers   http.Header `json:"headers"`
	ExpiresAt time.Time   `json:"expiresAt"`
}

// PresignedUploadURL contains the details for uploading an object or part.
type PresignedUploadURL struct {
	Method    string      `json:"method"`
	URL       string      `json:"url"`
	Headers   http.Header `json:"headers"`
	ExpiresAt time.Time   `json:"expiresAt"`
}

// CompletedPart represents a completed part in a multipart upload.
type CompletedPart struct {
	PartNumber int32
	ETag       string
}

type Client struct {
	endpoint        string
	accessKeyID     string
	secretAccessKey string
	httpClient      *http.Client
}

// ClientOptions contains configuration for creating a GCS client.
type ClientOptions struct {
	Endpoint string

	// AccessKeyID is the GCS HMAC access key ID
	AccessKeyID string
	// SecretAccessKey is the GCS HMAC secret key
	SecretAccessKey string
}

// NewClient creates a new GCS client using HMAC keys with the XML API.
func NewClient(ctx context.Context, opts ClientOptions) (*Client, error) {
	if opts.AccessKeyID == "" || opts.SecretAccessKey == "" {
		return nil, fmt.Errorf("HMAC access key ID and secret are required for GCS XML API")
	}

	return &Client{
		endpoint:        opts.Endpoint,
		accessKeyID:     opts.AccessKeyID,
		secretAccessKey: opts.SecretAccessKey,
		httpClient:      &http.Client{Timeout: 60 * time.Second},
	}, nil
}

// PresignGetObject creates a presigned GET URL for downloading an object.
func (c *Client) PresignGetObject(
	ctx context.Context,
	bucket,
	key string,
	sseKey []byte,
	expiresIn time.Duration,
) (*PresignedDownloadURL, error) {
	now := time.Now().UTC()
	expiresAt := now.Add(expiresIn)
	expiration := int(expiresIn.Seconds())

	// Build CSEK headers if provided - these will be signed into the URL
	var headers http.Header
	if len(sseKey) > 0 {
		headers = make(http.Header)
		encodedKey := base64.StdEncoding.EncodeToString(sseKey)
		hash := sha256.Sum256(sseKey)
		encodedHash := base64.StdEncoding.EncodeToString(hash[:])

		headers.Set("x-goog-encryption-algorithm", "AES256")
		headers.Set("x-goog-encryption-key", encodedKey)
		headers.Set("x-goog-encryption-key-sha256", encodedHash)
	}

	signedURL, err := c.generateSignedURL("GET", bucket, key, nil, headers, now, expiration)
	if err != nil {
		return nil, fmt.Errorf("generate signed URL: %w", err)
	}

	return &PresignedDownloadURL{
		Method:    "GET",
		URL:       signedURL,
		Headers:   headers,
		ExpiresAt: expiresAt,
	}, nil
}

// CreateMultipartUpload initiates a multipart upload and returns the upload ID.
func (c *Client) CreateMultipartUpload(
	ctx context.Context,
	bucket,
	key string,
	sseKey []byte,
) (string, error) {
	endpoint := fmt.Sprintf("%s/%s/%s?uploads", c.endpoint, bucket, url.PathEscape(key))
	now := time.Now().UTC()

	req, err := http.NewRequestWithContext(ctx, "POST", endpoint, nil)
	if err != nil {
		return "", fmt.Errorf("create request: %w", err)
	}

	// Set headers before signing
	req.Header.Set("Content-Length", "0")
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("x-goog-content-sha256", sha256Hash(""))

	// Add CSEK headers if provided
	addCSEKHeaders(req.Header, sseKey)

	// Sign the request
	if err := c.signRequest(req, bucket, key, "uploads=", "", now); err != nil {
		return "", fmt.Errorf("sign request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("initiate multipart upload failed: status %d: %s", resp.StatusCode, string(body))
	}

	// Parse XML response
	var result struct {
		UploadID string `xml:"UploadId"`
	}
	if err := xml.Unmarshal(body, &result); err != nil {
		return "", fmt.Errorf("parse response: %w", err)
	}

	return result.UploadID, nil
}

// PresignUploadPart creates a presigned URL for uploading a part.
func (c *Client) PresignUploadPart(
	ctx context.Context,
	bucket,
	key, uploadID string,
	sseKey []byte,
	partNumber int32,
	expiresIn time.Duration,
) (*PresignedUploadURL, error) {
	now := time.Now().UTC()
	expiresAt := now.Add(expiresIn)
	expiration := int(expiresIn.Seconds())

	queryParams := url.Values{}
	queryParams.Set("partNumber", fmt.Sprintf("%d", partNumber))
	queryParams.Set("uploadId", uploadID)

	// Build CSEK headers if provided - these will be signed into the URL
	var headers http.Header
	if len(sseKey) > 0 {
		headers = make(http.Header)
		encodedKey := base64.StdEncoding.EncodeToString(sseKey)
		hash := sha256.Sum256(sseKey)
		encodedHash := base64.StdEncoding.EncodeToString(hash[:])

		headers.Set("x-goog-encryption-algorithm", "AES256")
		headers.Set("x-goog-encryption-key", encodedKey)
		headers.Set("x-goog-encryption-key-sha256", encodedHash)
	}

	signedURL, err := c.generateSignedURL("PUT", bucket, key, queryParams, headers, now, expiration)
	if err != nil {
		return nil, fmt.Errorf("generate signed URL: %w", err)
	}

	return &PresignedUploadURL{
		Method:    "PUT",
		URL:       signedURL,
		Headers:   headers,
		ExpiresAt: expiresAt,
	}, nil
}

// CompleteMultipartUpload finalizes a multipart upload.
func (c *Client) CompleteMultipartUpload(
	ctx context.Context,
	bucket string,
	key string,
	uploadID string,
	parts []CompletedPart,
	sseKey []byte,
) error {
	// Build XML body
	type Part struct {
		PartNumber int32  `xml:"PartNumber"`
		ETag       string `xml:"ETag"`
	}
	requestBody := struct {
		XMLName xml.Name `xml:"CompleteMultipartUpload"`
		Parts   []Part   `xml:"Part"`
	}{}

	for _, part := range parts {
		requestBody.Parts = append(requestBody.Parts, Part{
			PartNumber: part.PartNumber,
			ETag:       part.ETag,
		})
	}

	xmlBody, err := xml.Marshal(requestBody)
	if err != nil {
		return fmt.Errorf("marshal XML: %w", err)
	}

	endpoint := fmt.Sprintf(
		"%s/%s/%s?uploadId=%s", c.endpoint, bucket, url.PathEscape(key), url.QueryEscape(uploadID),
	)
	now := time.Now().UTC()

	req, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(xmlBody))
	if err != nil {
		return fmt.Errorf("create request: %w", err)
	}

	// Set headers before signing
	req.Header.Set("Content-Length", fmt.Sprintf("%d", len(xmlBody)))
	req.Header.Set("Content-Type", "application/xml")
	req.Header.Set("x-goog-content-sha256", sha256Hash(string(xmlBody)))

	// Add CSEK headers if provided
	addCSEKHeaders(req.Header, sseKey)

	// Sign the request
	queryString := fmt.Sprintf("uploadId=%s", url.QueryEscape(uploadID))
	if err := c.signRequest(req, bucket, key, queryString, string(xmlBody), now); err != nil {
		return fmt.Errorf("sign request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("complete multipart upload failed: status %d: %s", resp.StatusCode, string(body))
	}

	return nil
}

// AbortMultipartUpload cancels a multipart upload.
func (c *Client) AbortMultipartUpload(
	ctx context.Context,
	bucket string,
	key string,
	uploadID string,
	sseKey []byte,
) error {
	endpoint := fmt.Sprintf(
		"%s/%s/%s?uploadId=%s", c.endpoint, bucket, url.PathEscape(key), url.QueryEscape(uploadID),
	)
	now := time.Now().UTC()

	req, err := http.NewRequestWithContext(ctx, "DELETE", endpoint, nil)
	if err != nil {
		return fmt.Errorf("create request: %w", err)
	}

	// Set headers before signing
	req.Header.Set("x-goog-content-sha256", sha256Hash(""))

	// Add CSEK headers if provided
	addCSEKHeaders(req.Header, sseKey)

	// Sign the request
	queryString := fmt.Sprintf("uploadId=%s", url.QueryEscape(uploadID))
	if err := c.signRequest(req, bucket, key, queryString, "", now); err != nil {
		return fmt.Errorf("sign request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("abort multipart upload failed: status %d: %s", resp.StatusCode, string(body))
	}

	return nil
}

// DeleteObject deletes an object from storage.
func (c *Client) DeleteObject(
	ctx context.Context,
	bucket string,
	key string,
) error {
	endpoint := fmt.Sprintf("%s/%s/%s", c.endpoint, bucket, url.PathEscape(key))
	now := time.Now().UTC()

	req, err := http.NewRequestWithContext(ctx, "DELETE", endpoint, nil)
	if err != nil {
		return fmt.Errorf("create request: %w", err)
	}

	// Set headers before signing
	req.Header.Set("x-goog-content-sha256", sha256Hash(""))

	// Sign the request
	if err := c.signRequest(req, bucket, key, "", "", now); err != nil {
		return fmt.Errorf("sign request: %w", err)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("execute request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("delete object failed: status %d: %s", resp.StatusCode, string(body))
	}

	return nil
}

// signRequest signs an HTTP request using GOOG4-HMAC-SHA256
func (c *Client) signRequest(req *http.Request, bucket, objectKey, queryString, payload string, now time.Time) error {
	// Set required headers
	req.Header.Set("Host", hostFromEndpoint(c.endpoint))
	req.Header.Set("X-Goog-Date", now.Format("20060102T150405Z"))

	// Create canonical request
	canonicalHeaders, signedHeaders := c.buildCanonicalHeaders(req.Header)
	payloadHash := sha256Hash(payload)

	// Build canonical URI: /{bucket}/{object}
	canonicalURI := fmt.Sprintf("/%s/%s", bucket, objectKey)

	// Build canonical query string (sorted alphabetically)
	canonicalQueryString := ""
	if queryString != "" {
		canonicalQueryString = queryString
	}

	canonicalRequest := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s",
		req.Method,
		canonicalURI,
		canonicalQueryString,
		canonicalHeaders,
		signedHeaders,
		payloadHash,
	)

	// Create string to sign
	datestamp := now.Format("20060102")
	credentialScope := fmt.Sprintf("%s/auto/storage/goog4_request", datestamp)
	stringToSign := fmt.Sprintf("GOOG4-HMAC-SHA256\n%s\n%s\n%s",
		now.Format("20060102T150405Z"),
		credentialScope,
		sha256Hash(canonicalRequest),
	)

	// Calculate signature
	signature := c.calculateSignature(c.secretAccessKey, datestamp, stringToSign)

	// Set authorization header
	authHeader := fmt.Sprintf("GOOG4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
		c.accessKeyID,
		credentialScope,
		signedHeaders,
		signature,
	)
	req.Header.Set("Authorization", authHeader)

	return nil
}

// generateSignedURL creates a presigned URL with optional additional headers
func (c *Client) generateSignedURL(method, bucket, key string, queryParams url.Values, additionalHeaders http.Header, now time.Time, expiration int) (string, error) {
	datestamp := now.Format("20060102")
	timestamp := now.Format("20060102T150405Z")
	credentialScope := fmt.Sprintf("%s/auto/storage/goog4_request", datestamp)

	if queryParams == nil {
		queryParams = url.Values{}
	}

	// Build canonical headers and signed headers list
	headers := make(http.Header)
	headers.Set("host", hostFromEndpoint(c.endpoint))

	// Add any additional headers (e.g., CSEK headers)
	for k, v := range additionalHeaders {
		if len(v) > 0 {
			headers.Set(k, v[0])
		}
	}

	canonicalHeaders, signedHeaders := c.buildCanonicalHeaders(headers)

	queryParams.Set("X-Goog-Algorithm", "GOOG4-HMAC-SHA256")
	queryParams.Set("X-Goog-Credential", c.accessKeyID+"/"+credentialScope)
	queryParams.Set("X-Goog-Date", timestamp)
	queryParams.Set("X-Goog-Expires", fmt.Sprintf("%d", expiration))
	queryParams.Set("X-Goog-SignedHeaders", signedHeaders)

	canonicalQueryString := queryParams.Encode()
	canonicalURI := "/" + bucket + "/" + key

	canonicalRequest := fmt.Sprintf("%s\n%s\n%s\n%s\n%s\nUNSIGNED-PAYLOAD",
		method,
		canonicalURI,
		canonicalQueryString,
		canonicalHeaders,
		signedHeaders,
	)

	stringToSign := fmt.Sprintf("GOOG4-HMAC-SHA256\n%s\n%s\n%s",
		timestamp,
		credentialScope,
		sha256Hash(canonicalRequest),
	)

	signature := c.calculateSignature(c.secretAccessKey, datestamp, stringToSign)

	finalURL := fmt.Sprintf("%s%s?%s&X-Goog-Signature=%s", c.endpoint,
		canonicalURI,
		canonicalQueryString,
		signature,
	)

	return finalURL, nil
}

// calculateSignature calculates GOOG4-HMAC-SHA256 signature
func (c *Client) calculateSignature(secretKey, datestamp, stringToSign string) string {
	dateKey := hmacSHA256([]byte("GOOG4"+secretKey), datestamp)
	regionKey := hmacSHA256(dateKey, "auto")
	serviceKey := hmacSHA256(regionKey, "storage")
	signingKey := hmacSHA256(serviceKey, "goog4_request")
	signature := hmacSHA256(signingKey, stringToSign)
	return hex.EncodeToString(signature)
}

// buildCanonicalHeaders builds canonical headers and signed headers list
func (c *Client) buildCanonicalHeaders(headers http.Header) (string, string) {
	var headerKeys []string
	headerMap := make(map[string]string)

	for k, v := range headers {
		lowerKey := strings.ToLower(k)
		if lowerKey == "host" || strings.HasPrefix(lowerKey, "x-goog-") {
			headerKeys = append(headerKeys, lowerKey)
			headerMap[lowerKey] = strings.TrimSpace(v[0])
		}
	}

	sort.Strings(headerKeys)

	var canonical strings.Builder
	for _, key := range headerKeys {
		canonical.WriteString(key)
		canonical.WriteString(":")
		canonical.WriteString(headerMap[key])
		canonical.WriteString("\n")
	}

	return canonical.String(), strings.Join(headerKeys, ";")
}

// Helper functions
func sha256Hash(data string) string {
	h := sha256.New()
	h.Write([]byte(data))
	return hex.EncodeToString(h.Sum(nil))
}

func hmacSHA256(key []byte, data string) []byte {
	h := hmac.New(sha256.New, key)
	h.Write([]byte(data))
	return h.Sum(nil)
}

func hostFromEndpoint(endpoint string) string {
	// https://storage.googleapis.com -> storage.googleapis.com
	url, err := url.Parse(endpoint)
	if err != nil {
		return ""
	}
	return url.Hostname()
}

// addCSEKHeaders adds Customer-Supplied Encryption Key headers to the request if sseKey is provided
func addCSEKHeaders(headers http.Header, sseKey []byte) {
	if len(sseKey) == 0 {
		return
	}

	// Encode the key in base64
	encodedKey := base64.StdEncoding.EncodeToString(sseKey)

	// Calculate SHA256 hash of the key and encode in base64
	hash := sha256.Sum256(sseKey)
	encodedHash := base64.StdEncoding.EncodeToString(hash[:])

	// Set CSEK headers
	headers.Set("x-goog-encryption-algorithm", "AES256")
	headers.Set("x-goog-encryption-key", encodedKey)
	headers.Set("x-goog-encryption-key-sha256", encodedHash)
}