main.go

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"strings"

	storage "google.golang.org/api/storage/v1"
)

var (
	url    = flag.String("url", "", "URL to fetch")
	n      = flag.Int("n", 100, "Number of chunks to use")
	bucket = flag.String("bucket", "", "GCS bucket to write to")
	tok    = flag.String("tok", "", "OAuth token to write to GCS")
)

func main() {
	flag.Parse()

	// Get total size of file
	resp, err := http.Head(*url)
	if err != nil {
		log.Fatalf("HEAD %s: %v", *url, err)
	}
	size, err := strconv.Atoi(resp.Header.Get("Content-Length"))
	if err != nil {
		log.Fatalf("strconv.Atoi: %v", err)
	}
	log.Printf("Total file size: %d", size)

	chunkSize := size / *n
	log.Printf("Will fetch %d chunks of %d bytes each", *n, chunkSize)

	for i := 0; i < *n; i++ {
		start := i * chunkSize
		end := start + chunkSize - 1
		if end > size {
			end = size
		}
		if err := fetchChunk(i, start, end); err != nil {
			log.Fatalf("fetchChunk(%d): %v", i, err)
		}
	}
}

func fetchChunk(i, start, end int) error {
	log.Printf("Chunk %d: range %d-%d", i, start, end)

	req, err := http.NewRequest("GET", *url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status)
	}
	defer resp.Body.Close()

	svc, err := storage.New(&http.Client{
		Transport: tokTransport(*tok),
	})
	if err != nil {
		return err
	}
	log.Printf("Writing part-%d", i)
	if _, err := svc.Objects.Insert(*bucket, &storage.Object{}).
		Name(fmt.Sprintf("part-%d", i)).
		Media(resp.Body).
		Do(); err != nil {
		return err
	}
	log.Printf("Finished writing part-%d", i)

	// Check if there are more parts to fetch, and if not, compose objects
	// together.
	l, err := svc.Objects.List(*bucket).Prefix("part-").Do()
	if err != nil {
		return err
	}
	if len(l.Items) == *n {
		log.Println("That was the last part, composing...")

		if _, err := svc.Objects.Insert(*bucket, &storage.Object{}).Name("final").Media(ioutil.NopCloser(strings.NewReader(""))).Do(); err != nil {
			return err
		}

		components := []*storage.ComposeRequestSourceObjects{{
			Name: "final",
		}}
		for i = 0; i < *n; i++ {
			components = append(components, &storage.ComposeRequestSourceObjects{Name: fmt.Sprintf("part-%d", i)})
			if len(components) == 30 || i == *n-1 {
				log.Printf("Composing from %d objects", len(components))
				if _, err := svc.Objects.Compose(*bucket, "final", &storage.ComposeRequest{
					SourceObjects: components,
				}).Do(); err != nil {
					return err
				}
				components = []*storage.ComposeRequestSourceObjects{{
					Name: "final",
				}}
			}
		}
	}
	return nil
}

type tokTransport string

func (t tokTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	r.Header.Set("Authorization", "Bearer "+string(t))
	return http.DefaultTransport.RoundTrip(r)
}