artifacts/storage.go

package artifacts

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

type PullResult struct {
	SHA256Sum string
	Size      int64
}
type PushResult struct {
	SHA256Sum string
	Size      int64
}

func PushBlob(ctx context.Context, r io.Reader, blobURL string, chunkSize int64, parallelism int) (*PushResult, error) {
	if chunkSize <= 0 {
		chunkSize = 4 * 1024 * 1024 // 4MB default
	}
	if parallelism <= 0 {
		parallelism = 4
	}

	type chunk struct {
		index int
		data  []byte
	}
	wrkinp := make(chan chunk)
	wrkerr := make(chan error, parallelism)

	var blockMu sync.Mutex
	blockIDs := make([]string, 0)

	var wg sync.WaitGroup

	// Worker pool for parallel uploading
	for i := 0; i < parallelism; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for c := range wrkinp {
				id := base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%06d", c.index)))
				address := fmt.Sprintf("%s&comp=block&blockid=%s", blobURL, id)
				req, err := http.NewRequestWithContext(ctx, "PUT", address, bytes.NewReader(c.data))
				if err != nil {
					wrkerr <- err
					return
				}
				req.ContentLength = int64(len(c.data))
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					wrkerr <- err
					return
				}
				defer resp.Body.Close()
				if resp.StatusCode != http.StatusCreated {
					body, _ := io.ReadAll(resp.Body)
					wrkerr <- fmt.Errorf("chunk %d upload failed: %s body: %s", c.index, resp.Status, string(body))
					return
				}

				blockMu.Lock()
				defer blockMu.Unlock()
				blockIDs = append(blockIDs, id)
			}
		}()
	}

	// Chunk the data and send to workers
	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])
				wrkinp <- chunk{index: index, data: buf[:n]}
				index++
			}
			break
		}
		if err != nil {
			close(wrkinp)
			wg.Wait()
			return nil, err
		}
		hash.Write(buf)
		wrkinp <- chunk{index: index, data: buf}
		index++
	}
	close(wrkinp)
	wg.Wait()

	select {
	case err := <-wrkerr:
		return nil, err
	default:
	}

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

	// Sort block IDs to ensure correct order
	blockMu.Lock()
	defer blockMu.Unlock()
	if len(blockIDs) == 0 {
		return nil, fmt.Errorf("no blocks uploaded")
	}
	blockList := &BlockList{}
	for i := 0; i < len(blockIDs); i++ {
		blockList.Latest = append(blockList.Latest, Latest{Value: blockIDs[i]})
	}
	blockListData, err := xml.Marshal(blockList)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, "PUT", 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 := http.DefaultClient.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
}

func PullBlob(ctx context.Context, w io.Writer, blobURL string, chunkSize int64, parallelism int) (*PullResult, error) {
	return nil, nil
}