github.com/munnerz/test-infra@v0.0.0-20190108210205-ce3d181dc989/testgrid/util/gcs/gcs.go (about)

     1  /*
     2  Copyright 2018 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package gcs
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"fmt"
    23  	"hash/crc32"
    24  	"log"
    25  	"net/url"
    26  	"strings"
    27  
    28  	"cloud.google.com/go/storage"
    29  	"google.golang.org/api/option"
    30  )
    31  
    32  // ClientWithCreds returns a storage client, optionally authenticated with the specified .json creds
    33  func ClientWithCreds(ctx context.Context, creds ...string) (*storage.Client, error) {
    34  	var options []option.ClientOption
    35  	switch l := len(creds); l {
    36  	case 0: // Do nothing
    37  	case 1:
    38  		options = append(options, option.WithCredentialsFile(creds[0]))
    39  	default:
    40  		return nil, fmt.Errorf("%d creds files unsupported (at most 1)", l)
    41  	}
    42  	return storage.NewClient(ctx, options...)
    43  }
    44  
    45  // Path parses gs://bucket/obj urls
    46  type Path struct {
    47  	url url.URL
    48  }
    49  
    50  // String returns the gs://bucket/obj url
    51  func (g Path) String() string {
    52  	return g.url.String()
    53  }
    54  
    55  // Set updates value from a gs://bucket/obj string, validating errors.
    56  func (g *Path) Set(v string) error {
    57  	u, err := url.Parse(v)
    58  	if err != nil {
    59  		return fmt.Errorf("invalid gs:// url %s: %v", v, err)
    60  	}
    61  	return g.SetURL(u)
    62  }
    63  
    64  // SetURL updates value to the passed in gs://bucket/obj url
    65  func (g *Path) SetURL(u *url.URL) error {
    66  	switch {
    67  	case u == nil:
    68  		return errors.New("nil url")
    69  	case u.Scheme != "gs":
    70  		return fmt.Errorf("must use a gs:// url: %s", u)
    71  	case strings.Contains(u.Host, ":"):
    72  		return fmt.Errorf("gs://bucket may not contain a port: %s", u)
    73  	case u.Opaque != "":
    74  		return fmt.Errorf("url must start with gs://: %s", u)
    75  	case u.User != nil:
    76  		return fmt.Errorf("gs://bucket may not contain an user@ prefix: %s", u)
    77  	case u.RawQuery != "":
    78  		return fmt.Errorf("gs:// url may not contain a ?query suffix: %s", u)
    79  	case u.Fragment != "":
    80  		return fmt.Errorf("gs:// url may not contain a #fragment suffix: %s", u)
    81  	}
    82  	g.url = *u
    83  	return nil
    84  }
    85  
    86  // ResolveReference returns the path relative to the current path
    87  func (g Path) ResolveReference(ref *url.URL) (*Path, error) {
    88  	var newP Path
    89  	if err := newP.SetURL(g.url.ResolveReference(ref)); err != nil {
    90  		return nil, err
    91  	}
    92  	return &newP, nil
    93  }
    94  
    95  // Bucket returns bucket in gs://bucket/obj
    96  func (g Path) Bucket() string {
    97  	return g.url.Host
    98  }
    99  
   100  // Object returns path/to/something in gs://bucket/path/to/something
   101  func (g Path) Object() string {
   102  	if g.url.Path == "" {
   103  		return g.url.Path
   104  	}
   105  	return g.url.Path[1:]
   106  }
   107  
   108  func calcCRC(buf []byte) uint32 {
   109  	return crc32.Checksum(buf, crc32.MakeTable(crc32.Castagnoli))
   110  }
   111  
   112  // Upload writes bytes to the specified Path
   113  func Upload(ctx context.Context, client *storage.Client, path Path, buf []byte) error {
   114  	crc := calcCRC(buf)
   115  	w := client.Bucket(path.Bucket()).Object(path.Object()).NewWriter(ctx)
   116  	w.SendCRC32C = true
   117  	// Send our CRC32 to ensure google received the same data we sent.
   118  	// See checksum example at:
   119  	// https://godoc.org/cloud.google.com/go/storage#Writer.Write
   120  	w.ObjectAttrs.CRC32C = crc
   121  	w.ProgressFunc = func(bytes int64) {
   122  		log.Printf("Uploading %s: %d/%d...", path, bytes, len(buf))
   123  	}
   124  	if n, err := w.Write(buf); err != nil {
   125  		return fmt.Errorf("writing %s failed: %v", path, err)
   126  	} else if n != len(buf) {
   127  		return fmt.Errorf("partial write of %s: %d < %d", path, n, len(buf))
   128  	}
   129  	if err := w.Close(); err != nil {
   130  		return fmt.Errorf("closing %s failed: %v", path, err)
   131  	}
   132  	return nil
   133  }