src/pkg/fakes/gcs.go - cos/tools - Git at Google

 // Copyright 2018 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package fakes

 import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
 	"mime"
 	"mime/multipart"
 	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 	"time"

 	"cloud.google.com/go/storage"
 	"google.golang.org/api/option"
 )

 type gcsObject struct{ Name, Bucket string }
 type gcsObjects struct {
 	Items    []gcsObject
 	Prefixes []string
 }

 func setTransportAddr(transport *http.Transport, addr string) {
 	transport.DialTLS = func(_, _ string) (net.Conn, error) {
 		return tls.Dial("tcp", addr, transport.TLSClientConfig)
 	}
 }

 // GCS contains data and functionality for a fake GCS server. It is intended to be constructed with NewGCSServer.
 //
 // The GCS struct represents the state of the fake GCS instance. Fields on this struct can be modified to influence the
 // return values of GCS API calls.
 //
 // The fake GCS server implements a small part of the API discussed here:
 // https://godoc.org/cloud.google.com/go/storage. Only the parts that we need for testing
 // are implemented here. Documentation for the GCS JSON API is here:
 // https://cloud.google.com/storage/docs/json_api/v1/
 //
 // This struct should not be considered concurrency safe.
 type GCS struct {
 	// Objects represents the collection of objects that exist in the fake GCS server.
 	// Keys are strings of the form "/<bucket>/<object path>". Values are data that belong
 	// in each object.
 	Objects map[string][]byte
 	// Client is the client to use when accessing the fake GCS server.
 	Client *storage.Client
 	// Server is the fake GCS server. It uses state from this struct for serving requests.
 	Server *httptest.Server
 	Delay  time.Duration
 	LogMissing bool
 }

 // NewGCSServer constructs a fake GCS implementation.
 func NewGCSServer(ctx context.Context) (*GCS, error) {
 	var err error
 	gcs := &GCS{make(map[string][]byte), nil, nil, time.Second * 0, true}
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", gcs.objectHandler)
 	mux.HandleFunc("/storage/v1/b/", gcs.bucketHandler)
 	mux.HandleFunc("/upload/storage/v1/b/", gcs.uploadHandler)
 	gcs.Server = httptest.NewTLSServer(mux)
 	httpClient := gcs.Server.Client()
 	setTransportAddr(httpClient.Transport.(*http.Transport), gcs.Server.Listener.Addr().String())
 	gcs.Client, err = storage.NewClient(ctx, option.WithHTTPClient(httpClient), option.WithoutAuthentication())
 	if err != nil {
 		gcs.Server.Close()
 		return nil, err
 	}
 	return gcs, nil
 }

 func (g *GCS) objectHandler(w http.ResponseWriter, r *http.Request) {
 	if g.Delay != 0 {
 		log.Printf("delaying object operation by %v", g.Delay)
 		time.Sleep(g.Delay)
 	}

 	data, ok := g.Objects[r.URL.Path]
 	if !ok {
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 	if _, err := w.Write(data); err != nil {
 		log.Printf("write %q failed: %v", r.URL.Path, err)
 	}
 }

 // list handles a `list` request.
 // See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `list` method.
 // Only handles the 'prefix', 'startOffset' and 'delimiter' optional parameters.
 func (g *GCS) list(w http.ResponseWriter, r *http.Request, bucket string) {
 	if err := r.ParseForm(); err != nil {
 		log.Printf("failed to parse form %q: %v", r.URL.Path, err)
 		return
 	}
 	bucketPrefix := fmt.Sprintf("/%s/", bucket)
 	prefix := bucketPrefix + r.Form.Get("prefix")
 	startOffset := bucketPrefix + r.Form.Get("startOffset")
 	delimiter := r.Form.Get("delimiter")
 	var all gcsObjects
 	allPrefixes := make(map[string]bool)
 	for k := range g.Objects {
 		if k < startOffset {
 			continue
 		}
 		if strings.HasPrefix(k, prefix) {
 			part := strings.TrimPrefix(k, prefix)
 			idx := strings.Index(part, delimiter)
 			if delimiter == "" || idx == -1 {
 				all.Items = append(all.Items, gcsObject{strings.TrimPrefix(k, bucketPrefix), bucket})
 			} else {
 				allPrefixes[strings.TrimPrefix(prefix+part[:idx+1], bucketPrefix)] = true
 			}
 		}
 	}
 	for k := range allPrefixes {
 		all.Prefixes = append(all.Prefixes, k)
 	}
 	bytes, err := json.Marshal(all)
 	if err != nil {
 		writeError(w, r, http.StatusInternalServerError)
 		return
 	}
 	if _, err := w.Write(bytes); err != nil {
 		log.Printf("write %q failed: %v", r.URL.Path, err)
 	}
 }

 // del handles a `delete` request.
 // See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `delete` method.
 // Doesn't handle any optional parameters.
 func (g *GCS) del(w http.ResponseWriter, r *http.Request, bucket, objectPath string) {
 	key := fmt.Sprintf("/%s/%s", bucket, objectPath)
 	if _, ok := g.Objects[key]; !ok {
 		if g.LogMissing {
 			log.Printf("delete failed: item %s does not exist", key)
 		}
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 	delete(g.Objects, key)
 }

 // get handles a `get` request.
 // See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `get` method.
 // Doesn't handle any optional parameters.
 func (g *GCS) get(w http.ResponseWriter, r *http.Request, bucket, objectPath string) {
 	key := fmt.Sprintf("/%s/%s", bucket, objectPath)
 	if _, ok := g.Objects[key]; !ok {
 		if g.LogMissing {
 			log.Printf("get failed: item %s does not exist", key)
 		}
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 	var result gcsObject
 	result.Name = key
 	result.Bucket = string(g.Objects[key])
 	bytes, err := json.Marshal(result)
 	if err != nil {
 		writeError(w, r, http.StatusInternalServerError)
 		return
 	}
 	if _, err := w.Write(bytes); err != nil {
 		log.Printf("write %q failed: %v", r.URL.Path, err)
 	}
 }

 func (g *GCS) bucketHandler(w http.ResponseWriter, r *http.Request) {
 	// Path looks like:
 	// - /storage/v1/b/<bucket>/o
 	// - /storage/v1/b/<bucket>/o/<object>
 	splitPath := strings.SplitN(r.URL.Path, "/", 7)
 	if len(splitPath) < 6 || splitPath[5] != "o" {
 		log.Printf("unrecognized path: %s", r.URL.Path)
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 	objectPath := ""
 	if len(splitPath) == 7 {
 		objectPath = splitPath[6]
 	}
 	bucket := splitPath[4]
 	switch {
 	case objectPath != "" && r.Method == "DELETE":
 		g.del(w, r, bucket, objectPath)
 	case objectPath != "" && r.Method == "GET":
 		g.get(w, r, bucket, objectPath)
 	case objectPath == "":
 		g.list(w, r, bucket)
 	default:
 		log.Printf("unrecognized path: %s", r.URL.Path)
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 }

 func (g *GCS) uploadHandler(w http.ResponseWriter, r *http.Request) {
 	splitPath := strings.Split(r.URL.Path, "/")
 	// Path looks like /upload/storage/v1/b/<bucket>/o
 	// (see: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `insert` method)
 	// This implementation only accepts the path parameter 'bucket'. It does not accept any optional parameters.
 	//
 	// GCS uses multipart HTTP messages to upload data. The first part contains object metadata (name, bucket, etc)
 	// in JSON format, and the second part contains the object data. Here, we extract the object metadata and data
 	// from the multipart message and store it
 	if len(splitPath) != 7 || splitPath[6] != "o" {
 		log.Printf("unrecognized path: %s", r.URL.Path)
 		writeError(w, r, http.StatusNotFound)
 		return
 	}
 	_, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
 	if err != nil {
 		log.Printf("failed to parse Content-Type: %s", r.Header.Get("Content-Type"))
 		writeError(w, r, http.StatusInternalServerError)
 		return
 	}
 	var parts [][]byte
 	mr := multipart.NewReader(r.Body, params["boundary"])
 	for {
 		part, err := mr.NextPart()
 		if err == io.EOF {
 			break
 		}
 		if err != nil {
 			log.Printf("failed to parse request: %v", r)
 			writeError(w, r, http.StatusInternalServerError)
 			return
 		}
 		partData, err := ioutil.ReadAll(part)
 		if err != nil {
 			log.Printf("failed to parse request: %v", r)
 			writeError(w, r, http.StatusInternalServerError)
 			return
 		}
 		parts = append(parts, partData)
 	}
 	objectMetadata := parts[0]
 	objectData := parts[1]
 	object := &gcsObject{}
 	if err := json.Unmarshal(objectMetadata, object); err != nil {
 		log.Printf("failed to parse object: %s", string(objectMetadata))
 		writeError(w, r, http.StatusInternalServerError)
 		return
 	}
 	g.Objects[fmt.Sprintf("/%s/%s", object.Bucket, object.Name)] = objectData
 	if _, err := w.Write(objectMetadata); err != nil {
 		log.Printf("write %q failed: %v", r.URL.Path, err)
 	}
 }

 // Close closes the fake GCS server and its client.
 func (g *GCS) Close() error {
 	defer g.Server.Close()
 	return g.Client.Close()
 }

 // GCSForTest encapsulates boilerplate for getting a GCS object in tests.
 func GCSForTest(t *testing.T) *GCS {
 	t.Helper()
 	gcs, err := NewGCSServer(context.Background())
 	if err != nil {
 		t.Fatal(err)
 	}
 	return gcs
 }
	// Copyright 2018 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package fakes

	import (
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"mime/multipart"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"cloud.google.com/go/storage"
	"google.golang.org/api/option"
	)

	type gcsObject struct{ Name, Bucket string }
	type gcsObjects struct {
	Items []gcsObject
	Prefixes []string
	}

	func setTransportAddr(transport *http.Transport, addr string) {
	transport.DialTLS = func(_, _ string) (net.Conn, error) {
	return tls.Dial("tcp", addr, transport.TLSClientConfig)
	}
	}

	// GCS contains data and functionality for a fake GCS server. It is intended to be constructed with NewGCSServer.
	//
	// The GCS struct represents the state of the fake GCS instance. Fields on this struct can be modified to influence the
	// return values of GCS API calls.
	//
	// The fake GCS server implements a small part of the API discussed here:
	// https://godoc.org/cloud.google.com/go/storage. Only the parts that we need for testing
	// are implemented here. Documentation for the GCS JSON API is here:
	// https://cloud.google.com/storage/docs/json_api/v1/
	//
	// This struct should not be considered concurrency safe.
	type GCS struct {
	// Objects represents the collection of objects that exist in the fake GCS server.
	// Keys are strings of the form "/<bucket>/<object path>". Values are data that belong
	// in each object.
	Objects map[string][]byte
	// Client is the client to use when accessing the fake GCS server.
	Client *storage.Client
	// Server is the fake GCS server. It uses state from this struct for serving requests.
	Server *httptest.Server
	Delay time.Duration
	LogMissing bool
	}

	// NewGCSServer constructs a fake GCS implementation.
	func NewGCSServer(ctx context.Context) (*GCS, error) {
	var err error
	gcs := &GCS{make(map[string][]byte), nil, nil, time.Second * 0, true}
	mux := http.NewServeMux()
	mux.HandleFunc("/", gcs.objectHandler)
	mux.HandleFunc("/storage/v1/b/", gcs.bucketHandler)
	mux.HandleFunc("/upload/storage/v1/b/", gcs.uploadHandler)
	gcs.Server = httptest.NewTLSServer(mux)
	httpClient := gcs.Server.Client()
	setTransportAddr(httpClient.Transport.(*http.Transport), gcs.Server.Listener.Addr().String())
	gcs.Client, err = storage.NewClient(ctx, option.WithHTTPClient(httpClient), option.WithoutAuthentication())
	if err != nil {
	gcs.Server.Close()
	return nil, err
	}
	return gcs, nil
	}

	func (g GCS) objectHandler(w http.ResponseWriter, r http.Request) {
	if g.Delay != 0 {
	log.Printf("delaying object operation by %v", g.Delay)
	time.Sleep(g.Delay)
	}

	data, ok := g.Objects[r.URL.Path]
	if !ok {
	writeError(w, r, http.StatusNotFound)
	return
	}
	if _, err := w.Write(data); err != nil {
	log.Printf("write %q failed: %v", r.URL.Path, err)
	}
	}

	// list handles a `list` request.
	// See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `list` method.
	// Only handles the 'prefix', 'startOffset' and 'delimiter' optional parameters.
	func (g GCS) list(w http.ResponseWriter, r http.Request, bucket string) {
	if err := r.ParseForm(); err != nil {
	log.Printf("failed to parse form %q: %v", r.URL.Path, err)
	return
	}
	bucketPrefix := fmt.Sprintf("/%s/", bucket)
	prefix := bucketPrefix + r.Form.Get("prefix")
	startOffset := bucketPrefix + r.Form.Get("startOffset")
	delimiter := r.Form.Get("delimiter")
	var all gcsObjects
	allPrefixes := make(map[string]bool)
	for k := range g.Objects {
	if k < startOffset {
	continue
	}
	if strings.HasPrefix(k, prefix) {
	part := strings.TrimPrefix(k, prefix)
	idx := strings.Index(part, delimiter)
	if delimiter == "" \|\| idx == -1 {
	all.Items = append(all.Items, gcsObject{strings.TrimPrefix(k, bucketPrefix), bucket})
	} else {
	allPrefixes[strings.TrimPrefix(prefix+part[:idx+1], bucketPrefix)] = true
	}
	}
	}
	for k := range allPrefixes {
	all.Prefixes = append(all.Prefixes, k)
	}
	bytes, err := json.Marshal(all)
	if err != nil {
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	if _, err := w.Write(bytes); err != nil {
	log.Printf("write %q failed: %v", r.URL.Path, err)
	}
	}

	// del handles a `delete` request.
	// See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `delete` method.
	// Doesn't handle any optional parameters.
	func (g GCS) del(w http.ResponseWriter, r http.Request, bucket, objectPath string) {
	key := fmt.Sprintf("/%s/%s", bucket, objectPath)
	if _, ok := g.Objects[key]; !ok {
	if g.LogMissing {
	log.Printf("delete failed: item %s does not exist", key)
	}
	writeError(w, r, http.StatusNotFound)
	return
	}
	delete(g.Objects, key)
	}

	// get handles a `get` request.
	// See: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `get` method.
	// Doesn't handle any optional parameters.
	func (g GCS) get(w http.ResponseWriter, r http.Request, bucket, objectPath string) {
	key := fmt.Sprintf("/%s/%s", bucket, objectPath)
	if _, ok := g.Objects[key]; !ok {
	if g.LogMissing {
	log.Printf("get failed: item %s does not exist", key)
	}
	writeError(w, r, http.StatusNotFound)
	return
	}
	var result gcsObject
	result.Name = key
	result.Bucket = string(g.Objects[key])
	bytes, err := json.Marshal(result)
	if err != nil {
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	if _, err := w.Write(bytes); err != nil {
	log.Printf("write %q failed: %v", r.URL.Path, err)
	}
	}

	func (g GCS) bucketHandler(w http.ResponseWriter, r http.Request) {
	// Path looks like:
	// - /storage/v1/b/<bucket>/o
	// - /storage/v1/b/<bucket>/o/<object>
	splitPath := strings.SplitN(r.URL.Path, "/", 7)
	if len(splitPath) < 6 \|\| splitPath[5] != "o" {
	log.Printf("unrecognized path: %s", r.URL.Path)
	writeError(w, r, http.StatusNotFound)
	return
	}
	objectPath := ""
	if len(splitPath) == 7 {
	objectPath = splitPath[6]
	}
	bucket := splitPath[4]
	switch {
	case objectPath != "" && r.Method == "DELETE":
	g.del(w, r, bucket, objectPath)
	case objectPath != "" && r.Method == "GET":
	g.get(w, r, bucket, objectPath)
	case objectPath == "":
	g.list(w, r, bucket)
	default:
	log.Printf("unrecognized path: %s", r.URL.Path)
	writeError(w, r, http.StatusNotFound)
	return
	}
	}

	func (g GCS) uploadHandler(w http.ResponseWriter, r http.Request) {
	splitPath := strings.Split(r.URL.Path, "/")
	// Path looks like /upload/storage/v1/b/<bucket>/o
	// (see: https://cloud.google.com/storage/docs/json_api/v1/#Objects, `insert` method)
	// This implementation only accepts the path parameter 'bucket'. It does not accept any optional parameters.
	//
	// GCS uses multipart HTTP messages to upload data. The first part contains object metadata (name, bucket, etc)
	// in JSON format, and the second part contains the object data. Here, we extract the object metadata and data
	// from the multipart message and store it
	if len(splitPath) != 7 \|\| splitPath[6] != "o" {
	log.Printf("unrecognized path: %s", r.URL.Path)
	writeError(w, r, http.StatusNotFound)
	return
	}
	_, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
	log.Printf("failed to parse Content-Type: %s", r.Header.Get("Content-Type"))
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	var parts [][]byte
	mr := multipart.NewReader(r.Body, params["boundary"])
	for {
	part, err := mr.NextPart()
	if err == io.EOF {
	break
	}
	if err != nil {
	log.Printf("failed to parse request: %v", r)
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	partData, err := ioutil.ReadAll(part)
	if err != nil {
	log.Printf("failed to parse request: %v", r)
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	parts = append(parts, partData)
	}
	objectMetadata := parts[0]
	objectData := parts[1]
	object := &gcsObject{}
	if err := json.Unmarshal(objectMetadata, object); err != nil {
	log.Printf("failed to parse object: %s", string(objectMetadata))
	writeError(w, r, http.StatusInternalServerError)
	return
	}
	g.Objects[fmt.Sprintf("/%s/%s", object.Bucket, object.Name)] = objectData
	if _, err := w.Write(objectMetadata); err != nil {
	log.Printf("write %q failed: %v", r.URL.Path, err)
	}
	}

	// Close closes the fake GCS server and its client.
	func (g *GCS) Close() error {
	defer g.Server.Close()
	return g.Client.Close()
	}

	// GCSForTest encapsulates boilerplate for getting a GCS object in tests.
	func GCSForTest(t testing.T) GCS {
	t.Helper()
	gcs, err := NewGCSServer(context.Background())
	if err != nil {
	t.Fatal(err)
	}
	return gcs
	}