go-sip/internal/agent/upload.go

package agent

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	agentv1 "git.ipao.vip/rogee/go-sip/gen/agent/v1"
)

// UploadClient performs the Agent-direct data-plane upload using a restricted
// Dispatcher grant. It never sends file bytes through Dispatcher or writes an
// OSS credential to logs. It never deletes or moves the source asset; the
// lifecycle owner retains it until upload-fact delivery to MQ is durably recorded.
type UploadClient struct {
	HTTPClient          *http.Client
	AllowedHosts        map[string]struct{}
	AllowInsecureHTTP   bool
	MaxResponseBodySize int64
	Now                 func() time.Time
}

var ErrUploadGrantExpired = errors.New("upload grant is expired")

type UploadResult struct {
	StatusCode int
	SizeBytes  int64
	SHA256     string
	ETag       string
}

func (c UploadClient) UploadFile(ctx context.Context, grant *agentv1.UploadGrant, path string) (result UploadResult, err error) {
	if grant == nil {
		return UploadResult{}, errors.New("upload grant is required")
	}
	if grant.TargetUrl == "" || grant.UploadId == "" || grant.ObjectKey == "" {
		return UploadResult{}, errors.New("upload URL, ID and object key are required")
	}
	if grant.ExpiresAtUnixMs <= 0 {
		return UploadResult{}, errors.New("upload grant expiry is required")
	}
	now := time.Now
	if c.Now != nil {
		now = c.Now
	}
	if !now().Before(time.UnixMilli(grant.ExpiresAtUnixMs)) {
		return UploadResult{}, ErrUploadGrantExpired
	}
	parsed, err := url.Parse(grant.TargetUrl)
	if err != nil || parsed.Host == "" {
		return UploadResult{}, errors.New("upload URL is invalid")
	}
	if parsed.Scheme != "https" && !(c.AllowInsecureHTTP && parsed.Scheme == "http") {
		return UploadResult{}, errors.New("upload URL must use HTTPS")
	}
	if len(c.AllowedHosts) > 0 {
		if _, ok := c.AllowedHosts[strings.ToLower(parsed.Host)]; !ok {
			return UploadResult{}, fmt.Errorf("upload host %q is not allowed", parsed.Host)
		}
	}
	if err := ctx.Err(); err != nil {
		return UploadResult{}, err
	}
	file, err := os.Open(path)
	if err != nil {
		return UploadResult{}, err
	}
	stat, err := file.Stat()
	if err != nil {
		_ = file.Close()
		return UploadResult{}, err
	}
	if stat.IsDir() {
		_ = file.Close()
		return UploadResult{}, errors.New("upload path is a directory")
	}
	if grant.MaxBytes > 0 && stat.Size() > grant.MaxBytes {
		_ = file.Close()
		return UploadResult{}, fmt.Errorf("asset exceeds grant limit: %d > %d", stat.Size(), grant.MaxBytes)
	}
	defer func() {
		if closeErr := file.Close(); closeErr != nil && !errors.Is(closeErr, os.ErrClosed) {
			result = UploadResult{}
			err = errors.Join(err, closeErr)
		}
	}()
	digest, err := digestFile(file)
	if err != nil {
		return UploadResult{}, err
	}
	if grant.RequiredChecksumSha256 != "" && !strings.EqualFold(grant.RequiredChecksumSha256, digest) {
		return UploadResult{}, errors.New("asset checksum does not match upload grant")
	}

	if _, err := file.Seek(0, io.SeekStart); err != nil {
		return UploadResult{}, err
	}
	transmitted := &uploadChecksum{hash: sha256.New()}
	body := io.TeeReader(io.LimitReader(file, stat.Size()), transmitted)
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, parsed.String(), body)
	if err != nil {
		return UploadResult{}, err
	}
	req.ContentLength = stat.Size()
	for _, header := range grant.Headers {
		if strings.EqualFold(header.Name, "host") || strings.EqualFold(header.Name, "content-length") {
			return UploadResult{}, errors.New("upload grant contains a forbidden header")
		}
		req.Header.Set(header.Name, header.Value)
	}
	client := c.HTTPClient
	if client == nil {
		client = &http.Client{}
	}
	copyClient := *client
	copyClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }
	resp, err := copyClient.Do(req)
	if err != nil {
		// net/http includes the entire signed URL in *url.Error. Retain the
		// underlying transport cause without exposing the temporary token.
		var requestError *url.Error
		if errors.As(err, &requestError) {
			return UploadResult{}, fmt.Errorf("upload PUT transport failure: %w", requestError.Err)
		}
		return UploadResult{}, err
	}
	defer resp.Body.Close()
	maxResponse := c.MaxResponseBodySize
	if maxResponse <= 0 {
		maxResponse = 64 << 10
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxResponse))
		return UploadResult{}, fmt.Errorf("upload returned HTTP %d", resp.StatusCode)
	}
	if _, err := io.Copy(io.Discard, io.LimitReader(resp.Body, maxResponse)); err != nil {
		return UploadResult{}, fmt.Errorf("read upload response: %w", err)
	}
	sentDigest, sentBytes := transmitted.result()
	if sentBytes != stat.Size() || sentDigest != digest {
		return UploadResult{}, errors.New("transmitted upload bytes do not match the validated asset")
	}
	return UploadResult{StatusCode: resp.StatusCode, SizeBytes: sentBytes, SHA256: sentDigest, ETag: resp.Header.Get("ETag")}, nil
}

// HTTP transports may still be writing the request when response headers arrive.
// Synchronize observation so early responses cannot race checksum calculation.
type uploadChecksum struct {
	mu    sync.Mutex
	hash  hash.Hash
	bytes int64
}

func (c *uploadChecksum) Write(p []byte) (int, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	n, err := c.hash.Write(p)
	c.bytes += int64(n)
	return n, err
}
func (c *uploadChecksum) result() (string, int64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return hex.EncodeToString(c.hash.Sum(nil)), c.bytes
}

func digestFile(file *os.File) (string, error) {
	if _, err := file.Seek(0, io.SeekStart); err != nil {
		return "", err
	}
	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}