sdk/artifact/storage.go

package artifact

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"sort"
	"sync"
)

type PullResult struct {
	SHA256Sum string
	Size      int64
}

type PushResult struct {
	SHA256Sum string
	Size      int64
}

func httpClientOrDefault(c *http.Client) *http.Client {
	if c == nil {
		return http.DefaultClient
	}
	return c
}

// PushBlob uploads the data read from r to the Azure block blob at blobURL as a
// series of blocks which are committed in order. The blobURL is expected to be
// a pre-signed URL. It returns the SHA-256 sum and size of the uploaded data.
func PushBlob(ctx context.Context, client *http.Client, r io.Reader, blobURL string, chunkSize int64, parallelism int) (*PushResult, error) {
	client = httpClientOrDefault(client)
	if chunkSize <= 0 {
		chunkSize = 4 * 1024 * 1024 // 4MB default
	}
	if parallelism <= 0 {
		parallelism = 4
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	type chunk struct {
		index int
		data  []byte
	}

	wrkinp := make(chan chunk)

	var firstErr error
	var errOnce sync.Once
	fail := func(err error) {
		errOnce.Do(func() {
			firstErr = err
			cancel()
		})
	}

	var blockMu sync.Mutex
	blocks := map[int]string{}

	var wg sync.WaitGroup
	for i := 0; i < parallelism; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for c := range wrkinp {
				id := base64.StdEncoding.EncodeToString(fmt.Appendf(nil, "%06d", c.index))
				address := fmt.Sprintf("%s&comp=block&blockid=%s", blobURL, id)
				req, err := http.NewRequestWithContext(ctx, http.MethodPut, address, bytes.NewReader(c.data))
				if err != nil {
					fail(err)
					return
				}
				req.ContentLength = int64(len(c.data))
				resp, err := client.Do(req)
				if err != nil {
					fail(err)
					return
				}
				body, _ := io.ReadAll(resp.Body)
				resp.Body.Close()
				if resp.StatusCode != http.StatusCreated {
					fail(fmt.Errorf("chunk %d upload failed: %s body: %s", c.index, resp.Status, string(body)))
					return
				}

				blockMu.Lock()
				blocks[c.index] = id
				blockMu.Unlock()
			}
		}()
	}

	index := 0
	hash := sha256.New()
	var size int64
	for {
		buf := make([]byte, chunkSize)
		n, err := io.ReadFull(r, buf)
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			if n > 0 {
				size += int64(n)
				hash.Write(buf[:n])
				select {
				case wrkinp <- chunk{index: index, data: buf[:n]}:
					index++
				case <-ctx.Done():
				}
			}
			break
		}
		if err != nil {
			fail(err)
			break
		}
		size += int64(n)
		hash.Write(buf)
		select {
		case wrkinp <- chunk{index: index, data: buf}:
			index++
		case <-ctx.Done():
			close(wrkinp)
			wg.Wait()
			return nil, firstErr
		}
	}
	close(wrkinp)
	wg.Wait()

	if firstErr != nil {
		return nil, firstErr
	}
	if len(blocks) == 0 {
		return nil, fmt.Errorf("no blocks uploaded")
	}

	type Latest struct {
		Value string `xml:",chardata"`
	}
	type BlockList struct {
		Latest []Latest `xml:"Latest"`
	}

	// The order of the block list determines the order of the resulting blob, so
	// the blocks must be committed by ascending index.
	indexes := make([]int, 0, len(blocks))
	for i := range blocks {
		indexes = append(indexes, i)
	}
	sort.Ints(indexes)

	blockList := &BlockList{}
	for _, i := range indexes {
		blockList.Latest = append(blockList.Latest, Latest{Value: blocks[i]})
	}
	blockListData, err := xml.Marshal(blockList)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPut, fmt.Sprintf("%s&comp=blocklist", blobURL), bytes.NewReader(blockListData))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/xml")
	req.ContentLength = int64(len(blockListData))
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("blocklist commit failed: %s body: %s", resp.Status, string(body))
	}
	return &PushResult{SHA256Sum: hex.EncodeToString(hash.Sum(nil)), Size: size}, nil
}

// PullBlob downloads the blob at blobURL, writing its contents to w. The
// blobURL is expected to be a pre-signed URL. It returns the SHA-256 sum and
// size of the downloaded data.
func PullBlob(ctx context.Context, client *http.Client, w io.Writer, blobURL string) (*PullResult, error) {
	client = httpClientOrDefault(client)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, blobURL, nil)
	if err != nil {
		return nil, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 32*1024))
		return nil, fmt.Errorf("pull blob failed: %s body: %s", resp.Status, string(body))
	}

	hash := sha256.New()
	size, err := io.Copy(io.MultiWriter(w, hash), resp.Body)
	if err != nil {
		return nil, err
	}
	return &PullResult{SHA256Sum: hex.EncodeToString(hash.Sum(nil)), Size: size}, nil
}