feat: artifacts client

artifact/storage.go

@@ -0,0 +1,212 @@
+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
+}