<commit_before>package libocit

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"path"
)

type OS struct {
	ID           string
	Distribution string
	Version      string
	Arch         string
	CPU          int64
	Memory       int64
	IP           string
	Status       string
}

//TODO add a 'casename' ?
type Deploy struct {
	Object     string
	Class      string
	Cmd        string
	Files      []string
	Containers []Container

	ResourceID string
}

//FIXME: the type is not consistent
type OSResource struct {
	CPU int
	Memory string
	Disk string
}

type Require struct {
	Class        string
	Type         string
	Distribution string
	Version      int
	Resource     OSResource
	Files        []string
}

type Container struct {
	Object       string
	Class        string
	Cmd          string
	Files        []string
	Distribution string
	Version      int
}

type Collect struct {
	Object string
	Files  []string

	ResourceID string
}

type Resource struct {
	//TODO: put following to a struct and make a hash?
	ID     string //returned
	Status bool   //whether it is available
	Msg    string //return value from server

	Req  Require
	Used bool
}

type TestCase struct {
	Name     string
	Version  string
	License  string
	Group    string
	Owner    string
	Description string
	Sources  []string
	Requires []Require
	Deploys  []Deploy
	Collects []Collect
}

func PreparePath(cachename string, filename string) (realurl string) {
	realurl = path.Join(cachename, filename)
	dir := path.Dir(realurl)
	p, err := os.Stat(dir)
	if err != nil {
		if !os.IsExist(err) {
			os.MkdirAll(dir, 0777)
		}
	} else {
		if p.IsDir() {
			return realurl
		} else {
			os.Remove(dir)
			os.MkdirAll(dir, 0777)
		}
	}
	return realurl
}

func SendFile(post_url string, file_url string, filename string) {
	bodyBuf := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(bodyBuf)

	//'tcfile': testcase file
	fileWriter, err := bodyWriter.CreateFormFile("tcfile", filename)
	if err != nil {
		fmt.Println("error writing to buffer")
		return
	}
	fh, err := os.Open(file_url)
	if err != nil {
		fmt.Println("error opening file")
		return
	}
	_, err = io.Copy(fileWriter, fh)
	if err != nil {
		return
	}
	contentType := bodyWriter.FormDataContentType()
	bodyWriter.Close()
	resp, err := http.Post(post_url, contentType, bodyBuf)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	resp_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	fmt.Println(resp.Status)
	fmt.Println(resp_body)
}

func SendCommand(apiurl string, b []byte) {
	body := bytes.NewBuffer(b)
	resp, perr := http.Post(apiurl, "application/json;charset=utf-8", body)
	defer resp.Body.Close()
	if perr != nil {
		// handle error
		fmt.Println("err in post:", perr)
		return
	} else {
		result, berr := ioutil.ReadAll(resp.Body)
		if berr != nil {
		} else {
			//TODO: write back the info
			fmt.Println(result)
		}
	}
}

func ReadFile(file_url string) (content string) {
	file, err := os.Open(file_url)
	defer file.Close()
	if err != nil {
		fmt.Println(file_url, err)
		return content
	}
	buf := bytes.NewBufferString("")
	buf.ReadFrom(file)
	content = buf.String()

	return content
}

func ReceiveFile(w http.ResponseWriter, r *http.Request, cache_url string) (real_url string, handle_name string) {
	r.ParseMultipartForm(32 << 20)
	file, handler, err := r.FormFile("tcfile")
	fmt.Println(handler.Filename)
	if err != nil {
		fmt.Println("Cannot find the tc file")
		return real_url, handle_name
	}
	defer file.Close()

	real_url = PreparePath(cache_url, handler.Filename)
	f, err := os.Create(real_url)
	if err != nil {
		fmt.Println("Cannot create the file ", real_url)
		//TODO: better system error
		http.Error(w, err.Error(), 500)
		return real_url, handle_name
	}
	defer f.Close()
	io.Copy(f, file)

	handle_name = handler.Filename
	return real_url, handle_name
}

func TarFilelist(filelist []string, case_dir string, object_name string) (tar_url string) {
	tar_url = path.Join(case_dir, object_name) + ".tar.gz"
	fw, err := os.Create(tar_url)
	if err != nil {
		fmt.Println("Failed in create tar file ", err)
		return tar_url
	}
	defer fw.Close()
	gw := gzip.NewWriter(fw)
	defer gw.Close()
	tw := tar.NewWriter(gw)
	defer tw.Close()

	for index := 0; index < len(filelist); index++ {
		source_file := filelist[index]
		fi, err := os.Stat(path.Join(case_dir, source_file))
		if err != nil {
			fmt.Println(err)
			continue
		}
		fr, err := os.Open(path.Join(case_dir, source_file))
		if err != nil {
			fmt.Println(err)
			continue
		}
		h := new(tar.Header)
		h.Name = source_file
		h.Size = fi.Size()
		h.Mode = int64(fi.Mode())
		h.ModTime = fi.ModTime()
		err = tw.WriteHeader(h)
		_, err = io.Copy(tw, fr)
	}
	return tar_url
}

func UntarFile(cache_url string, filename string) {
	fr, err := os.Open(filename)
	if err != nil {
		fmt.Println("fail in open file ", filename)
		return
	}
	defer fr.Close()
	gr, err := gzip.NewReader(fr)
	if err != nil {
		fmt.Println("fail in using gzip")
		return
	}
	defer gr.Close()

	tr := tar.NewReader(gr)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		target_url := PreparePath(cache_url, h.Name)
		fw, _ := os.OpenFile(target_url, os.O_CREATE|os.O_WRONLY, os.FileMode(h.Mode))
		defer fw.Close()

		io.Copy(fw, tr)
	}
}
<commit_msg>Update libocit.go<commit_after>package libocit

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"
	"path"
)

type OS struct {
	ID           string
	Distribution string
	Version      string
	Arch         string
	CPU          int64
	Memory       int64
	IP           string
	Status       string
}

//TODO add a 'casename' ?
type Deploy struct {
	Object     string
	Class      string
	Cmd        string
	Files      []string
	Containers []Container

	ResourceID string
}

//FIXME: the type is not consistent
type OSResource struct {
	CPU int
	Memory string
	Disk string
}

type Require struct {
	Class        string
	Type         string
	Distribution string
	Version      int
	Resource     OSResource
	Files        []string
}

type Container struct {
	Object       string
	Class        string
	Cmd          string
	Files        []string
	Distribution string
	Version      int
}

type Collect struct {
	Object string
	Files  []string

	ResourceID string
}

type Resource struct {
	//TODO: put following to a struct and make a hash?
	ID     string //returned
	Status bool   //whether it is available
	Msg    string //return value from server

	Req  Require
	Used bool
}

type TestCase struct {
	Name     string
	Version  string
	License  string
	Group    string
	Owner    string
	Description string
	Sources  []string
	Requires []Require
	Deploys  []Deploy
	Collects []Collect
}

func PreparePath(cachename string, filename string) (realurl string) {
	realurl = path.Join(cachename, filename)
	dir := path.Dir(realurl)
	p, err := os.Stat(dir)
	if err != nil {
		if !os.IsExist(err) {
			os.MkdirAll(dir, 0777)
		}
	} else {
		if p.IsDir() {
			return realurl
		} else {
			os.Remove(dir)
			os.MkdirAll(dir, 0777)
		}
	}
	return realurl
}

func SendFile(post_url string, file_url string, filename string) {
	bodyBuf := &bytes.Buffer{}
	bodyWriter := multipart.NewWriter(bodyBuf)

	//'tcfile': testcase file
	fileWriter, err := bodyWriter.CreateFormFile("tcfile", filename)
	if err != nil {
		fmt.Println("error writing to buffer")
		return
	}
	fh, err := os.Open(file_url)
	if err != nil {
		fmt.Println("error opening file")
		return
	}
	_, err = io.Copy(fileWriter, fh)
	if err != nil {
		return
	}
	contentType := bodyWriter.FormDataContentType()
	bodyWriter.Close()
	resp, err := http.Post(post_url, contentType, bodyBuf)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	resp_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	fmt.Println(resp.Status)
	fmt.Println(resp_body)
}

func SendCommand(apiurl string, b []byte) {
	body := bytes.NewBuffer(b)
	resp, perr := http.Post(apiurl, "application/json;charset=utf-8", body)
	defer resp.Body.Close()
	if perr != nil {
		// handle error
		fmt.Println("err in post:", perr)
		return
	} else {
		result, berr := ioutil.ReadAll(resp.Body)
		if berr != nil {
		} else {
			//TODO: write back the info
			fmt.Println(result)
		}
	}
}

func ReadFile(file_url string) (content string) {
	_, err := os.Stat(file_url)
	if err != nil {
		fmt.Println("cannot find the file ", file_url)
		return content
	}
	file, err := os.Open(file_url)
	defer file.Close()
	if err != nil {
		fmt.Println("cannot open the file ", file_url)
		return content
	}
	buf := bytes.NewBufferString("")
	buf.ReadFrom(file)
	content = buf.String()

	return content
}

func ReceiveFile(w http.ResponseWriter, r *http.Request, cache_url string) (real_url string, handle_name string) {
	r.ParseMultipartForm(32 << 20)
	file, handler, err := r.FormFile("tcfile")
	fmt.Println(handler.Filename)
	if err != nil {
		fmt.Println("Cannot find the tc file")
		return real_url, handle_name
	}
	defer file.Close()

	real_url = PreparePath(cache_url, handler.Filename)
	f, err := os.Create(real_url)
	if err != nil {
		fmt.Println("Cannot create the file ", real_url)
		//TODO: better system error
		http.Error(w, err.Error(), 500)
		return real_url, handle_name
	}
	defer f.Close()
	io.Copy(f, file)

	handle_name = handler.Filename
	return real_url, handle_name
}

func TarFilelist(filelist []string, case_dir string, object_name string) (tar_url string) {
	tar_url = path.Join(case_dir, object_name) + ".tar.gz"
	fw, err := os.Create(tar_url)
	if err != nil {
		fmt.Println("Failed in create tar file ", err)
		return tar_url
	}
	defer fw.Close()
	gw := gzip.NewWriter(fw)
	defer gw.Close()
	tw := tar.NewWriter(gw)
	defer tw.Close()

	for index := 0; index < len(filelist); index++ {
		source_file := filelist[index]
		fi, err := os.Stat(path.Join(case_dir, source_file))
		if err != nil {
			fmt.Println(err)
			continue
		}
		fr, err := os.Open(path.Join(case_dir, source_file))
		if err != nil {
			fmt.Println(err)
			continue
		}
		h := new(tar.Header)
		h.Name = source_file
		h.Size = fi.Size()
		h.Mode = int64(fi.Mode())
		h.ModTime = fi.ModTime()
		err = tw.WriteHeader(h)
		_, err = io.Copy(tw, fr)
	}
	return tar_url
}

func UntarFile(cache_url string, filename string) {
	fr, err := os.Open(filename)
	if err != nil {
		fmt.Println("fail in open file ", filename)
		return
	}
	defer fr.Close()
	gr, err := gzip.NewReader(fr)
	if err != nil {
		fmt.Println("fail in using gzip")
		return
	}
	defer gr.Close()

	tr := tar.NewReader(gr)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		target_url := PreparePath(cache_url, h.Name)
		fw, _ := os.OpenFile(target_url, os.O_CREATE|os.O_WRONLY, os.FileMode(h.Mode))
		defer fw.Close()

		io.Copy(fw, tr)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"io"
	"io/ioutil"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/go-connections/nat"
	"github.com/moby/moby/client"

	"golang.org/x/net/context"

	"github.com/ivanilves/lstags/docker"
	"github.com/ivanilves/lstags/docker/config"
)

// DockerSocket is a socket we use to connect to the Docker daemon
var DockerSocket = "/var/run/docker.sock"

// RetryPulls is a number of retries we do in case of Docker pull failure
var RetryPulls = 0

// RetryDelay is a delay between retries of the failed Docker pulls
var RetryDelay = 5 * time.Second

// DockerClient is a raw Docker client convenience wrapper
type DockerClient struct {
	cli *client.Client
	cnf *config.Config
}

// New creates new instance of DockerClient (our Docker client wrapper)
// Use DOCKER_HOST to set the url to the docker server.
// This Depends on the operating system for Linux unix:///var/run/docker.sock and for windows npipe:////./pipe/docker_engine
// Use DOCKER_API_VERSION to set the version of the API to reach, leave empty for latest.
// API_VERSION is by Default 1.27
// Use DOCKER_CERT_PATH to load the TLS certificates from.
// DOCKER_CERT_PATH/ca.pem
// DOCKER_CERT_PATH/cert.pem
// DOCKER_CERT_PATH/key.pem
// Use DOCKER_TLS_VERIFY to enable or disable TLS verification, off by default.
func New(cnf *config.Config) (*DockerClient, error) {
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}

	return &DockerClient{cli: cli, cnf: cnf}, nil
}

// ListImagesForRepo lists images present locally for the repo specified
func (dc *DockerClient) ListImagesForRepo(repo string) ([]types.ImageSummary, error) {
	listOptions, err := buildImageListOptions(repo)
	if err != nil {
		return nil, err
	}

	return dc.cli.ImageList(context.Background(), listOptions)
}

func buildImageListOptions(repo string) (types.ImageListOptions, error) {
	repoFilter := "reference=" + repo
	filterArgs := filters.NewArgs()

	filterArgs, err := filters.ParseFlag(repoFilter, filterArgs)
	if err != nil {
		return types.ImageListOptions{}, err
	}

	return types.ImageListOptions{Filters: filterArgs}, nil
}

// Pull pulls Docker image specified
func (dc *DockerClient) Pull(ref string) error {
	registryAuth := dc.cnf.GetRegistryAuth(
		docker.GetRegistry(ref),
	)

	pullOptions := types.ImagePullOptions{RegistryAuth: registryAuth}
	if registryAuth == "" {
		pullOptions = types.ImagePullOptions{}
	}

	tries := 1

	if RetryPulls > 0 {
		tries = tries + RetryPulls
	}

	var resp io.ReadCloser
	var err error

	for try := 1; try <= tries; try++ {
		resp, err = dc.cli.ImagePull(context.Background(), ref, pullOptions)

		if err == nil {
			break
		}

		time.Sleep(RetryDelay)

		RetryDelay += RetryDelay
	}

	if err != nil {
		return err
	}

	_, err = ioutil.ReadAll(resp)

	return err
}

// Push pushes Docker image specified
func (dc *DockerClient) Push(ref string) error {
	registryAuth := dc.cnf.GetRegistryAuth(
		docker.GetRegistry(ref),
	)

	pushOptions := types.ImagePushOptions{RegistryAuth: registryAuth}
	if registryAuth == "" {
		pushOptions = types.ImagePushOptions{RegistryAuth: "IA=="}
	}

	resp, err := dc.cli.ImagePush(context.Background(), ref, pushOptions)
	if err != nil {
		return err
	}

	_, err = ioutil.ReadAll(resp)

	return err
}

// Tag puts a "dst" tag on "src" Docker image
func (dc *DockerClient) Tag(src, dst string) error {
	return dc.cli.ImageTag(context.Background(), src, dst)
}

// Run runs Docker container from the image specified (like "docker run")
func (dc *DockerClient) Run(ref, name string, portSpecs []string) (string, error) {
	exposedPorts, portBindings, err := nat.ParsePortSpecs(portSpecs)
	if err != nil {
		return "", err
	}

	ctx := context.Background()

	if err := dc.Pull(ref); err != nil {
		return "", err
	}

	resp, err := dc.cli.ContainerCreate(
		ctx,
		&container.Config{Image: ref, ExposedPorts: exposedPorts},
		&container.HostConfig{PortBindings: portBindings},
		nil,
		name,
	)
	if err != nil {
		return "", err
	}

	if err := dc.cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		return "", err
	}

	return resp.ID, nil
}

// ForceRemove kills & removes Docker container having the ID specified (like "docker rm -f")
func (dc *DockerClient) ForceRemove(id string) error {
	return dc.cli.ContainerRemove(
		context.Background(),
		id,
		types.ContainerRemoveOptions{Force: true},
	)
}
<commit_msg>NORELEASE Update comments because I am crazy...<commit_after>package client

import (
	"io"
	"io/ioutil"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/go-connections/nat"
	"github.com/moby/moby/client"

	"golang.org/x/net/context"

	"github.com/ivanilves/lstags/docker"
	"github.com/ivanilves/lstags/docker/config"
)

// DockerSocket is a socket we use to connect to the Docker daemon
var DockerSocket = "/var/run/docker.sock"

// RetryPulls is a number of retries we do in case of Docker pull failure
var RetryPulls = 0

// RetryDelay is a delay between retries of the failed Docker pulls
var RetryDelay = 5 * time.Second

// DockerClient is a raw Docker client convenience wrapper
type DockerClient struct {
	cli *client.Client
	cnf *config.Config
}

// New creates new instance of DockerClient (our Docker client wrapper)
// Use DOCKER_HOST to set the URL to the Docker server.
// This depends on the operating system: for Linux unix:///var/run/docker.sock and for Windows npipe:////./pipe/docker_engine
// Use DOCKER_API_VERSION to set the version of the API to reach, leave empty for latest.
// API_VERSION is by default 1.27 (this may change)
// Use DOCKER_CERT_PATH to load the TLS certificates from.
// DOCKER_CERT_PATH/ca.pem
// DOCKER_CERT_PATH/cert.pem
// DOCKER_CERT_PATH/key.pem
// Use DOCKER_TLS_VERIFY to enable or disable TLS verification, off by default.
func New(cnf *config.Config) (*DockerClient, error) {
	cli, err := client.NewEnvClient()
	if err != nil {
		return nil, err
	}

	return &DockerClient{cli: cli, cnf: cnf}, nil
}

// ListImagesForRepo lists images present locally for the repo specified
func (dc *DockerClient) ListImagesForRepo(repo string) ([]types.ImageSummary, error) {
	listOptions, err := buildImageListOptions(repo)
	if err != nil {
		return nil, err
	}

	return dc.cli.ImageList(context.Background(), listOptions)
}

func buildImageListOptions(repo string) (types.ImageListOptions, error) {
	repoFilter := "reference=" + repo
	filterArgs := filters.NewArgs()

	filterArgs, err := filters.ParseFlag(repoFilter, filterArgs)
	if err != nil {
		return types.ImageListOptions{}, err
	}

	return types.ImageListOptions{Filters: filterArgs}, nil
}

// Pull pulls Docker image specified
func (dc *DockerClient) Pull(ref string) error {
	registryAuth := dc.cnf.GetRegistryAuth(
		docker.GetRegistry(ref),
	)

	pullOptions := types.ImagePullOptions{RegistryAuth: registryAuth}
	if registryAuth == "" {
		pullOptions = types.ImagePullOptions{}
	}

	tries := 1

	if RetryPulls > 0 {
		tries = tries + RetryPulls
	}

	var resp io.ReadCloser
	var err error

	for try := 1; try <= tries; try++ {
		resp, err = dc.cli.ImagePull(context.Background(), ref, pullOptions)

		if err == nil {
			break
		}

		time.Sleep(RetryDelay)

		RetryDelay += RetryDelay
	}

	if err != nil {
		return err
	}

	_, err = ioutil.ReadAll(resp)

	return err
}

// Push pushes Docker image specified
func (dc *DockerClient) Push(ref string) error {
	registryAuth := dc.cnf.GetRegistryAuth(
		docker.GetRegistry(ref),
	)

	pushOptions := types.ImagePushOptions{RegistryAuth: registryAuth}
	if registryAuth == "" {
		pushOptions = types.ImagePushOptions{RegistryAuth: "IA=="}
	}

	resp, err := dc.cli.ImagePush(context.Background(), ref, pushOptions)
	if err != nil {
		return err
	}

	_, err = ioutil.ReadAll(resp)

	return err
}

// Tag puts a "dst" tag on "src" Docker image
func (dc *DockerClient) Tag(src, dst string) error {
	return dc.cli.ImageTag(context.Background(), src, dst)
}

// Run runs Docker container from the image specified (like "docker run")
func (dc *DockerClient) Run(ref, name string, portSpecs []string) (string, error) {
	exposedPorts, portBindings, err := nat.ParsePortSpecs(portSpecs)
	if err != nil {
		return "", err
	}

	ctx := context.Background()

	if err := dc.Pull(ref); err != nil {
		return "", err
	}

	resp, err := dc.cli.ContainerCreate(
		ctx,
		&container.Config{Image: ref, ExposedPorts: exposedPorts},
		&container.HostConfig{PortBindings: portBindings},
		nil,
		name,
	)
	if err != nil {
		return "", err
	}

	if err := dc.cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
		return "", err
	}

	return resp.ID, nil
}

// ForceRemove kills & removes Docker container having the ID specified (like "docker rm -f")
func (dc *DockerClient) ForceRemove(id string) error {
	return dc.cli.ContainerRemove(
		context.Background(),
		id,
		types.ContainerRemoveOptions{Force: true},
	)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"encoding/xml"
	"fmt"
	update "github.com/coreos/updatectl/client/update/v1"
	"github.com/coreos/go-omaha/omaha"
	"log"
	"math/rand"
	"net/http"
	"text/tabwriter"
	"time"
)

var (
	instanceFlags struct {
		groupId       StringFlag
		appId         StringFlag
		start         int64
		end           int64
		verbose       bool
		clientsPerApp int
		minSleep      int
		maxSleep      int
		errorRate     int
		OEM           string
		pingOnly      int
		version       string
	}

	cmdInstance = &Command{
		Name:        "instance",
		Usage:       "[OPTION]...",
		Summary: "Operations to view instances.",
		Subcommands: []*Command{
			cmdInstanceListUpdates,
			cmdInstanceListAppVersions,
			cmdInstanceFake,
		},
	}

	cmdInstanceListUpdates = &Command{
		Name:        "instance list-updates",
		Usage:       "[OPTION]...",
		Description: "Generates a list of instance updates.",
		Run:         instanceListUpdates,
	}

	cmdInstanceListAppVersions = &Command{
		Name:        "instance list-app-versions",
		Usage:       "[OPTION]...",
		Description: "Generates a list of apps/versions with instance count.",
		Run:         instanceListAppVersions,
	}

	cmdInstanceFake = &Command{
		Name:        "instance fake",
		Usage:       "[OPTION]...",
		Description: "Simulate multiple fake instances.",
		Run:         instanceFake,
	}
)

func init() {
	cmdInstanceListUpdates.Flags.Var(&instanceFlags.groupId, "group-id", "Group id")
	cmdInstanceListUpdates.Flags.Var(&instanceFlags.appId, "app-id", "App id")
	cmdInstanceListUpdates.Flags.Int64Var(&instanceFlags.start, "start", 0, "Start date filter")
	cmdInstanceListUpdates.Flags.Int64Var(&instanceFlags.end, "end", 0, "End date filter")

	cmdInstanceListAppVersions.Flags.Var(&instanceFlags.groupId, "group-id", "Group id")
	cmdInstanceListAppVersions.Flags.Var(&instanceFlags.appId, "app-id", "App id")
	cmdInstanceListAppVersions.Flags.Int64Var(&instanceFlags.start, "start", 0, "Start date filter")
	cmdInstanceListAppVersions.Flags.Int64Var(&instanceFlags.end, "end", 0, "End date filter")

	cmdInstanceFake.Flags.BoolVar(&instanceFlags.verbose, "verbose", false, "Print out the request bodies")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.clientsPerApp, "clients-per-app", 20, "Number of fake fents per appid.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.minSleep, "min-sleep", 1, "Minimum time between update checks.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.maxSleep, "max-sleep", 10, "Maximum time between update checks.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.errorRate, "errorrate", 1, "Chance of error (0-100)%.")
	cmdInstanceFake.Flags.StringVar(&instanceFlags.OEM, "oem", "fakeclient", "oem to report")
	// simulate reboot lock.
	cmdInstanceFake.Flags.IntVar(&instanceFlags.pingOnly, "ping-only", 0, "halt update and just send ping requests this many times.")
	cmdInstanceFake.Flags.Var(&instanceFlags.appId, "app-id", "Application ID to update.")
	instanceFlags.appId.required = true
	cmdInstanceFake.Flags.Var(&instanceFlags.groupId, "group-id", "Group ID to update.")
	instanceFlags.groupId.required = true
	cmdInstanceFake.Flags.StringVar(&instanceFlags.version, "version", "0.0.0", "Version to report.")
}

func instanceListUpdates(args []string, service *update.Service, out *tabwriter.Writer) int {
	call := service.Clientupdate.List()
	call.DateStart(instanceFlags.start)
	call.DateEnd(instanceFlags.end)
	if instanceFlags.groupId.Get() != nil {
		call.GroupId(instanceFlags.groupId.String())
	}
	if instanceFlags.groupId.Get() != nil {
		call.AppId(instanceFlags.appId.String())
	}
	list, err := call.Do()

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintln(out, "AppID\tClientID\tVersion\tLastSeen\tGroup\tStatus\tOEM")
	for _, cl := range list.Items {
		fmt.Fprintf(out, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", cl.AppId,
			cl.ClientId, cl.Version, cl.LastSeen, cl.GroupId,
			cl.Status, cl.Oem)
	}
	out.Flush()
	return OK
}

func instanceListAppVersions(args []string, service *update.Service, out *tabwriter.Writer) int {
	call := service.Appversion.List()

	if instanceFlags.groupId.Get() != nil {
		call.GroupId(instanceFlags.groupId.String())
	}
	if instanceFlags.appId.Get() != nil {
		call.AppId(instanceFlags.appId.String())
	}
	if instanceFlags.start != 0 {
		call.DateStart(instanceFlags.start)
	}

	if instanceFlags.end != 0 {
		call.DateEnd(instanceFlags.end)
	}

	list, err := call.Do()

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintln(out, "AppID\tGroupID\tVersion\tClients")
	for _, cl := range list.Items {
		fmt.Fprintf(out, "%s\t%s\t%s\t%d\n", cl.AppId, cl.GroupId, cl.Version, cl.Count)
	}
	out.Flush()
	return OK
}

type serverConfig struct {
	server string
}

type Client struct {
	Id             string
	SessionId      string
	Version        string
	AppId          string
	Track          string
	config         *serverConfig
	errorRate      int
	pingsRemaining int
}

func (c *Client) Log(format string, v ...interface{}) {
	format = c.Id + ": " + format
	fmt.Printf(format, v...)
}

func (c *Client) OmahaRequest(otype, result string, updateCheck, isPing bool) *omaha.Request {
	req := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := req.AddApp(c.AppId, c.Version)
	app.MachineID = c.Id
	app.BootId = c.SessionId
	app.Track = c.Track
	app.OEM = instanceFlags.OEM

	if updateCheck {
		app.AddUpdateCheck()
	}

	if isPing {
		app.AddPing()
		app.Ping.LastReportDays = "1"
		app.Ping.Status = "1"
	}

	if otype != "" {
		event := app.AddEvent()
		event.Type = otype
		event.Result = result
		if result == "0" {
			event.ErrorCode = "2000"
		} else {
			event.ErrorCode = ""
		}
	}

	return req
}

func (c *Client) MakeRequest(otype, result string, updateCheck, isPing bool) (*omaha.Response, error) {
	client := &http.Client{}
	req := c.OmahaRequest(otype, result, updateCheck, isPing)
	raw, err := xml.MarshalIndent(req, "", " ")
	if err != nil {
		return nil, err
	}

	resp, err := client.Post(c.config.server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	oresp := new(omaha.Response)
	err = xml.NewDecoder(resp.Body).Decode(oresp)
	if err != nil {
		return nil, err
	}

	if instanceFlags.verbose {
		raw, _ := xml.MarshalIndent(req, "", " ")
		c.Log("request: %s\n", string(raw))
		raw, _ = xml.MarshalIndent(oresp, "", " ")
		c.Log("response: %s\n", string(raw))
	}

	return oresp, nil
}

func (c *Client) SetVersion(resp *omaha.Response) {
	// A field can potentially be nil.
	defer func() {
		if err := recover(); err != nil {
			c.Log("%s: error setting version: %v", c.Id, err)
		}
	}()

	uc := resp.Apps[0].UpdateCheck
	if uc.Status != "ok" {
		c.Log("%s\n", uc.Status)
		return
	}

	randFailRequest := func(eventType, eventResult string) (failed bool, err error) {
		if rand.Intn(100) <= c.errorRate {
			eventType = "3"
			eventResult = "0"
			failed = true
		}
		_, err = c.MakeRequest(eventType, eventResult, false, false)
		return
	}

	requests := [][]string{
		[]string{"13", "1"}, // downloading
		[]string{"14", "1"}, // downloaded
		[]string{"3", "1"},  // installed
	}

	for i, r := range requests {
		if i > 0 {
			time.Sleep(1 * time.Second)
		}
		failed, err := randFailRequest(r[0], r[1])
		if failed {
			log.Printf("failed to update in eventType: %s, eventResult: %s. Retrying.", r[0], r[1])
			time.Sleep(time.Second * time.Duration(instanceFlags.minSleep))
			c.MakeRequest(r[0], r[1], false, false)
			return
		}
		if err != nil {
			log.Println(err)
			return
		}
	}

	// simulate reboot lock for a while
	for c.pingsRemaining > 0 {
		c.MakeRequest("", "", false, true)
		c.pingsRemaining--
		time.Sleep(1 * time.Second)
	}

	c.Log("updated from %s to %s\n", c.Version, uc.Manifest.Version)

	c.Version = uc.Manifest.Version

	_, err := c.MakeRequest("3", "2", false, false) // Send complete with new version.
	if err != nil {
		log.Println(err)
	}

	c.SessionId = uuid.New()
}

// Sleep between n and m seconds
func (c *Client) Loop(n, m int) {
	for {
		randSleep(n, m)

		resp, err := c.MakeRequest("3", "2", true, false)
		if err != nil {
			log.Println(err)
			continue
		}
		c.SetVersion(resp)
	}
}

// Sleeps randomly between n and m seconds.
func randSleep(n, m int) {
	r := m
	if m-n > 0 {
		r = rand.Intn(m-n) + n
	}
	time.Sleep(time.Duration(r) * time.Second)
}

func instanceFake(args []string, service *update.Service, out *tabwriter.Writer) int {
	if instanceFlags.appId.Get() == nil || instanceFlags.groupId.Get() == nil {
		return ERROR_USAGE
	}

	conf := &serverConfig{
		server: globalFlags.Server,
	}

	for i := 0; i < instanceFlags.clientsPerApp; i++ {
		c := &Client{
			Id:             fmt.Sprintf("{fake-client-%03d}", i),
			SessionId:      uuid.New(),
			Version:        instanceFlags.version,
			AppId:          instanceFlags.appId.String(),
			Track:          instanceFlags.groupId.String(),
			config:         conf,
			errorRate:      instanceFlags.errorRate,
			pingsRemaining: instanceFlags.pingOnly,
		}
		go c.Loop(instanceFlags.minSleep, instanceFlags.maxSleep)
	}

	// run forever
	wait := make(chan bool)
	<-wait
	return OK
}
<commit_msg>instance: add newline to the log entry<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go-uuid/uuid"
	"encoding/xml"
	"fmt"
	update "github.com/coreos/updatectl/client/update/v1"
	"github.com/coreos/go-omaha/omaha"
	"log"
	"math/rand"
	"net/http"
	"text/tabwriter"
	"time"
)

var (
	instanceFlags struct {
		groupId       StringFlag
		appId         StringFlag
		start         int64
		end           int64
		verbose       bool
		clientsPerApp int
		minSleep      int
		maxSleep      int
		errorRate     int
		OEM           string
		pingOnly      int
		version       string
	}

	cmdInstance = &Command{
		Name:        "instance",
		Usage:       "[OPTION]...",
		Summary: "Operations to view instances.",
		Subcommands: []*Command{
			cmdInstanceListUpdates,
			cmdInstanceListAppVersions,
			cmdInstanceFake,
		},
	}

	cmdInstanceListUpdates = &Command{
		Name:        "instance list-updates",
		Usage:       "[OPTION]...",
		Description: "Generates a list of instance updates.",
		Run:         instanceListUpdates,
	}

	cmdInstanceListAppVersions = &Command{
		Name:        "instance list-app-versions",
		Usage:       "[OPTION]...",
		Description: "Generates a list of apps/versions with instance count.",
		Run:         instanceListAppVersions,
	}

	cmdInstanceFake = &Command{
		Name:        "instance fake",
		Usage:       "[OPTION]...",
		Description: "Simulate multiple fake instances.",
		Run:         instanceFake,
	}
)

func init() {
	cmdInstanceListUpdates.Flags.Var(&instanceFlags.groupId, "group-id", "Group id")
	cmdInstanceListUpdates.Flags.Var(&instanceFlags.appId, "app-id", "App id")
	cmdInstanceListUpdates.Flags.Int64Var(&instanceFlags.start, "start", 0, "Start date filter")
	cmdInstanceListUpdates.Flags.Int64Var(&instanceFlags.end, "end", 0, "End date filter")

	cmdInstanceListAppVersions.Flags.Var(&instanceFlags.groupId, "group-id", "Group id")
	cmdInstanceListAppVersions.Flags.Var(&instanceFlags.appId, "app-id", "App id")
	cmdInstanceListAppVersions.Flags.Int64Var(&instanceFlags.start, "start", 0, "Start date filter")
	cmdInstanceListAppVersions.Flags.Int64Var(&instanceFlags.end, "end", 0, "End date filter")

	cmdInstanceFake.Flags.BoolVar(&instanceFlags.verbose, "verbose", false, "Print out the request bodies")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.clientsPerApp, "clients-per-app", 20, "Number of fake fents per appid.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.minSleep, "min-sleep", 1, "Minimum time between update checks.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.maxSleep, "max-sleep", 10, "Maximum time between update checks.")
	cmdInstanceFake.Flags.IntVar(&instanceFlags.errorRate, "errorrate", 1, "Chance of error (0-100)%.")
	cmdInstanceFake.Flags.StringVar(&instanceFlags.OEM, "oem", "fakeclient", "oem to report")
	// simulate reboot lock.
	cmdInstanceFake.Flags.IntVar(&instanceFlags.pingOnly, "ping-only", 0, "halt update and just send ping requests this many times.")
	cmdInstanceFake.Flags.Var(&instanceFlags.appId, "app-id", "Application ID to update.")
	instanceFlags.appId.required = true
	cmdInstanceFake.Flags.Var(&instanceFlags.groupId, "group-id", "Group ID to update.")
	instanceFlags.groupId.required = true
	cmdInstanceFake.Flags.StringVar(&instanceFlags.version, "version", "0.0.0", "Version to report.")
}

func instanceListUpdates(args []string, service *update.Service, out *tabwriter.Writer) int {
	call := service.Clientupdate.List()
	call.DateStart(instanceFlags.start)
	call.DateEnd(instanceFlags.end)
	if instanceFlags.groupId.Get() != nil {
		call.GroupId(instanceFlags.groupId.String())
	}
	if instanceFlags.groupId.Get() != nil {
		call.AppId(instanceFlags.appId.String())
	}
	list, err := call.Do()

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintln(out, "AppID\tClientID\tVersion\tLastSeen\tGroup\tStatus\tOEM")
	for _, cl := range list.Items {
		fmt.Fprintf(out, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", cl.AppId,
			cl.ClientId, cl.Version, cl.LastSeen, cl.GroupId,
			cl.Status, cl.Oem)
	}
	out.Flush()
	return OK
}

func instanceListAppVersions(args []string, service *update.Service, out *tabwriter.Writer) int {
	call := service.Appversion.List()

	if instanceFlags.groupId.Get() != nil {
		call.GroupId(instanceFlags.groupId.String())
	}
	if instanceFlags.appId.Get() != nil {
		call.AppId(instanceFlags.appId.String())
	}
	if instanceFlags.start != 0 {
		call.DateStart(instanceFlags.start)
	}

	if instanceFlags.end != 0 {
		call.DateEnd(instanceFlags.end)
	}

	list, err := call.Do()

	if err != nil {
		log.Fatal(err)
	}

	fmt.Fprintln(out, "AppID\tGroupID\tVersion\tClients")
	for _, cl := range list.Items {
		fmt.Fprintf(out, "%s\t%s\t%s\t%d\n", cl.AppId, cl.GroupId, cl.Version, cl.Count)
	}
	out.Flush()
	return OK
}

type serverConfig struct {
	server string
}

type Client struct {
	Id             string
	SessionId      string
	Version        string
	AppId          string
	Track          string
	config         *serverConfig
	errorRate      int
	pingsRemaining int
}

func (c *Client) Log(format string, v ...interface{}) {
	format = c.Id + ": " + format
	fmt.Printf(format, v...)
}

func (c *Client) OmahaRequest(otype, result string, updateCheck, isPing bool) *omaha.Request {
	req := omaha.NewRequest("lsb", "CoreOS", "", "")
	app := req.AddApp(c.AppId, c.Version)
	app.MachineID = c.Id
	app.BootId = c.SessionId
	app.Track = c.Track
	app.OEM = instanceFlags.OEM

	if updateCheck {
		app.AddUpdateCheck()
	}

	if isPing {
		app.AddPing()
		app.Ping.LastReportDays = "1"
		app.Ping.Status = "1"
	}

	if otype != "" {
		event := app.AddEvent()
		event.Type = otype
		event.Result = result
		if result == "0" {
			event.ErrorCode = "2000"
		} else {
			event.ErrorCode = ""
		}
	}

	return req
}

func (c *Client) MakeRequest(otype, result string, updateCheck, isPing bool) (*omaha.Response, error) {
	client := &http.Client{}
	req := c.OmahaRequest(otype, result, updateCheck, isPing)
	raw, err := xml.MarshalIndent(req, "", " ")
	if err != nil {
		return nil, err
	}

	resp, err := client.Post(c.config.server+"/v1/update/", "text/xml", bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	oresp := new(omaha.Response)
	err = xml.NewDecoder(resp.Body).Decode(oresp)
	if err != nil {
		return nil, err
	}

	if instanceFlags.verbose {
		raw, _ := xml.MarshalIndent(req, "", " ")
		c.Log("request: %s\n", string(raw))
		raw, _ = xml.MarshalIndent(oresp, "", " ")
		c.Log("response: %s\n", string(raw))
	}

	return oresp, nil
}

func (c *Client) SetVersion(resp *omaha.Response) {
	// A field can potentially be nil.
	defer func() {
		if err := recover(); err != nil {
			c.Log("%s: error setting version: %v\n", c.Id, err)
		}
	}()

	uc := resp.Apps[0].UpdateCheck
	if uc.Status != "ok" {
		c.Log("%s\n", uc.Status)
		return
	}

	randFailRequest := func(eventType, eventResult string) (failed bool, err error) {
		if rand.Intn(100) <= c.errorRate {
			eventType = "3"
			eventResult = "0"
			failed = true
		}
		_, err = c.MakeRequest(eventType, eventResult, false, false)
		return
	}

	requests := [][]string{
		[]string{"13", "1"}, // downloading
		[]string{"14", "1"}, // downloaded
		[]string{"3", "1"},  // installed
	}

	for i, r := range requests {
		if i > 0 {
			time.Sleep(1 * time.Second)
		}
		failed, err := randFailRequest(r[0], r[1])
		if failed {
			log.Printf("failed to update in eventType: %s, eventResult: %s. Retrying.", r[0], r[1])
			time.Sleep(time.Second * time.Duration(instanceFlags.minSleep))
			c.MakeRequest(r[0], r[1], false, false)
			return
		}
		if err != nil {
			log.Println(err)
			return
		}
	}

	// simulate reboot lock for a while
	for c.pingsRemaining > 0 {
		c.MakeRequest("", "", false, true)
		c.pingsRemaining--
		time.Sleep(1 * time.Second)
	}

	c.Log("updated from %s to %s\n", c.Version, uc.Manifest.Version)

	c.Version = uc.Manifest.Version

	_, err := c.MakeRequest("3", "2", false, false) // Send complete with new version.
	if err != nil {
		log.Println(err)
	}

	c.SessionId = uuid.New()
}

// Sleep between n and m seconds
func (c *Client) Loop(n, m int) {
	for {
		randSleep(n, m)

		resp, err := c.MakeRequest("3", "2", true, false)
		if err != nil {
			log.Println(err)
			continue
		}
		c.SetVersion(resp)
	}
}

// Sleeps randomly between n and m seconds.
func randSleep(n, m int) {
	r := m
	if m-n > 0 {
		r = rand.Intn(m-n) + n
	}
	time.Sleep(time.Duration(r) * time.Second)
}

func instanceFake(args []string, service *update.Service, out *tabwriter.Writer) int {
	if instanceFlags.appId.Get() == nil || instanceFlags.groupId.Get() == nil {
		return ERROR_USAGE
	}

	conf := &serverConfig{
		server: globalFlags.Server,
	}

	for i := 0; i < instanceFlags.clientsPerApp; i++ {
		c := &Client{
			Id:             fmt.Sprintf("{fake-client-%03d}", i),
			SessionId:      uuid.New(),
			Version:        instanceFlags.version,
			AppId:          instanceFlags.appId.String(),
			Track:          instanceFlags.groupId.String(),
			config:         conf,
			errorRate:      instanceFlags.errorRate,
			pingsRemaining: instanceFlags.pingOnly,
		}
		go c.Loop(instanceFlags.minSleep, instanceFlags.maxSleep)
	}

	// run forever
	wait := make(chan bool)
	<-wait
	return OK
}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	pb "github.com/libp2p/go-libp2p-circuit/pb"

	ggio "github.com/gogo/protobuf/io"
	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
)

var log = logging.Logger("relay")

const ProtoID = "/libp2p/circuit/relay/0.1.0"

const maxMessageSize = 4096

var RelayAcceptTimeout = time.Minute

type Relay struct {
	host host.Host
	ctx  context.Context
	self peer.ID

	active bool
	hop    bool

	incoming chan *Conn

	arLk         sync.Mutex
	activeRelays []*Conn
}

type RelayOpt int

var (
	OptActive = RelayOpt(0)
	OptHop    = RelayOpt(1)
)

type RelayError struct {
	Code pb.CircuitRelay_Status
}

func (e RelayError) Error() string {
	return fmt.Sprintf("error opening relay circuit: %s (%d)", pb.CircuitRelay_Status_name[int32(e.Code)], e.Code)
}

func NewRelay(ctx context.Context, h host.Host, opts ...RelayOpt) (*Relay, error) {
	r := &Relay{
		host:     h,
		ctx:      ctx,
		self:     h.ID(),
		incoming: make(chan *Conn),
	}

	for _, opt := range opts {
		switch opt {
		case OptActive:
			r.active = true
		case OptHop:
			r.hop = true
		default:
			return nil, fmt.Errorf("unrecognized option: %d", opt)
		}
	}

	h.SetStreamHandler(ProtoID, r.handleNewStream)

	return r, nil
}

func (r *Relay) Dial(ctx context.Context, relay pstore.PeerInfo, dest pstore.PeerInfo) (*Conn, error) {
	err := r.host.Connect(ctx, relay)
	if err != nil {
		return nil, err
	}

	s, err := r.host.NewStream(ctx, relay.ID, ProtoID)
	if err != nil {
		return nil, err
	}

	rd := ggio.NewDelimitedReader(s, maxMessageSize)
	wr := ggio.NewDelimitedWriter(s)

	var msg pb.CircuitRelay

	msg.Type = pb.CircuitRelay_HOP.Enum()
	msg.SrcPeer = peerInfoToPeer(pstore.PeerInfo{r.self, r.host.Addrs()})
	msg.DstPeer = peerInfoToPeer(dest)

	err = wr.WriteMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	msg.Reset()

	err = rd.ReadMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		s.Close()
		return nil, fmt.Errorf("unexpected relay response; not a status message (%d)", msg.GetType())
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		s.Close()
		return nil, RelayError{msg.GetCode()}
	}

	return &Conn{Stream: s}, nil
}

func (r *Relay) handleNewStream(s inet.Stream) {
	log.Infof("new relay stream from: %s", s.Conn().RemotePeer())

	rd := ggio.NewDelimitedReader(s, maxMessageSize)

	var msg pb.CircuitRelay

	err := rd.ReadMsg(&msg)
	if err != nil {
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
		return
	}

	switch msg.GetType() {
	case pb.CircuitRelay_HOP:
		r.handleHopStream(s, &msg)
	case pb.CircuitRelay_STOP:
		r.handleStopStream(s, &msg)
	case pb.CircuitRelay_CAN_HOP:
		r.handleCanHop(s, &msg)
	default:
		log.Warningf("unexpected relay handshake: %d", msg.GetType())
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
	}
}

func (r *Relay) handleHopStream(s inet.Stream, msg *pb.CircuitRelay) {
	if !r.hop {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
		return
	}

	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	if src.ID != s.Conn().RemotePeer() {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_DST_MULTIADDR_INVALID)
		return
	}

	if dst.ID == r.self {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_RELAY_TO_SELF)
		return
	}

	// open stream
	ctp := r.host.Network().ConnsToPeer(dst.ID)
	if len(ctp) == 0 {
		if !r.active {
			r.handleError(s, pb.CircuitRelay_HOP_NO_CONN_TO_DST)
			return
		}

		ctx, cancel := context.WithTimeout(r.ctx, time.Second*10)
		defer cancel()
		err = r.host.Connect(ctx, dst)
		if err != nil {
			log.Debugf("error opening relay connection to %s: %s", dst.ID.Pretty(), err.Error())
			r.handleError(s, pb.CircuitRelay_HOP_CANT_DIAL_DST)
			return
		}
	}

	bs, err := r.host.NewStream(r.ctx, dst.ID, ProtoID)
	if err != nil {
		log.Debugf("error opening relay stream to %s: %s", dst.ID.Pretty(), err.Error())
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	// stop handshake
	rd := ggio.NewDelimitedReader(bs, maxMessageSize)
	wr := ggio.NewDelimitedWriter(bs)

	msg.Type = pb.CircuitRelay_STOP.Enum()

	err = wr.WriteMsg(msg)
	if err != nil {
		log.Debugf("error writing stop handshake: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	msg.Reset()

	err = rd.ReadMsg(msg)
	if err != nil {
		log.Debugf("error reading stop response: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		log.Debugf("unexpected relay stop response: not a status message (%d)", msg.GetType())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		log.Debugf("relay stop failure: %d", msg.GetCode())
		bs.Close()
		r.handleError(s, msg.GetCode())
		return
	}

	err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
		s.Close()
		return
	}

	// relay connection
	log.Infof("relaying connection between %s and %s", src.ID.Pretty(), dst.ID.Pretty())

	go func() {
		_, err := io.Copy(s, bs)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		s.Close()
	}()

	go func() {
		_, err := io.Copy(bs, s)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		bs.Close()
	}()
}

func (r *Relay) handleStopStream(s inet.Stream, msg *pb.CircuitRelay) {
	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil || len(src.Addrs) == 0 {
		r.handleError(s, pb.CircuitRelay_STOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil || dst.ID != r.self {
		r.handleError(s, pb.CircuitRelay_STOP_DST_MULTIADDR_INVALID)
		return
	}

	log.Infof("relay connection from: %s", src.ID)

	r.host.Peerstore().AddAddrs(src.ID, src.Addrs, pstore.TempAddrTTL)

	select {
	case r.incoming <- &Conn{Stream: s, remote: src}:
	case <-time.After(RelayAcceptTimeout):
		r.handleError(s, pb.CircuitRelay_STOP_RELAY_REFUSED)
	}
}

func (r *Relay) handleCanHop(s inet.Stream, msg *pb.CircuitRelay) {
	var err error

	if r.hop {
		err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	} else {
		err = r.writeResponse(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
	}

	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}

	s.Close()
}

func (r *Relay) handleError(s inet.Stream, code pb.CircuitRelay_Status) {
	log.Warningf("relay error: %s (%d)", pb.CircuitRelay_Status_name[int32(code)], code)
	err := r.writeResponse(s, code)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}
	s.Close()
}

func (r *Relay) writeResponse(s inet.Stream, code pb.CircuitRelay_Status) error {
	wr := ggio.NewDelimitedWriter(s)

	var msg pb.CircuitRelay
	msg.Type = pb.CircuitRelay_STATUS.Enum()
	msg.Code = code.Enum()

	return wr.WriteMsg(&msg)
}
<commit_msg>hop relay: close bs on status response error<commit_after>package relay

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	pb "github.com/libp2p/go-libp2p-circuit/pb"

	ggio "github.com/gogo/protobuf/io"
	logging "github.com/ipfs/go-log"
	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
)

var log = logging.Logger("relay")

const ProtoID = "/libp2p/circuit/relay/0.1.0"

const maxMessageSize = 4096

var RelayAcceptTimeout = time.Minute

type Relay struct {
	host host.Host
	ctx  context.Context
	self peer.ID

	active bool
	hop    bool

	incoming chan *Conn

	arLk         sync.Mutex
	activeRelays []*Conn
}

type RelayOpt int

var (
	OptActive = RelayOpt(0)
	OptHop    = RelayOpt(1)
)

type RelayError struct {
	Code pb.CircuitRelay_Status
}

func (e RelayError) Error() string {
	return fmt.Sprintf("error opening relay circuit: %s (%d)", pb.CircuitRelay_Status_name[int32(e.Code)], e.Code)
}

func NewRelay(ctx context.Context, h host.Host, opts ...RelayOpt) (*Relay, error) {
	r := &Relay{
		host:     h,
		ctx:      ctx,
		self:     h.ID(),
		incoming: make(chan *Conn),
	}

	for _, opt := range opts {
		switch opt {
		case OptActive:
			r.active = true
		case OptHop:
			r.hop = true
		default:
			return nil, fmt.Errorf("unrecognized option: %d", opt)
		}
	}

	h.SetStreamHandler(ProtoID, r.handleNewStream)

	return r, nil
}

func (r *Relay) Dial(ctx context.Context, relay pstore.PeerInfo, dest pstore.PeerInfo) (*Conn, error) {
	err := r.host.Connect(ctx, relay)
	if err != nil {
		return nil, err
	}

	s, err := r.host.NewStream(ctx, relay.ID, ProtoID)
	if err != nil {
		return nil, err
	}

	rd := ggio.NewDelimitedReader(s, maxMessageSize)
	wr := ggio.NewDelimitedWriter(s)

	var msg pb.CircuitRelay

	msg.Type = pb.CircuitRelay_HOP.Enum()
	msg.SrcPeer = peerInfoToPeer(pstore.PeerInfo{r.self, r.host.Addrs()})
	msg.DstPeer = peerInfoToPeer(dest)

	err = wr.WriteMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	msg.Reset()

	err = rd.ReadMsg(&msg)
	if err != nil {
		s.Close()
		return nil, err
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		s.Close()
		return nil, fmt.Errorf("unexpected relay response; not a status message (%d)", msg.GetType())
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		s.Close()
		return nil, RelayError{msg.GetCode()}
	}

	return &Conn{Stream: s}, nil
}

func (r *Relay) handleNewStream(s inet.Stream) {
	log.Infof("new relay stream from: %s", s.Conn().RemotePeer())

	rd := ggio.NewDelimitedReader(s, maxMessageSize)

	var msg pb.CircuitRelay

	err := rd.ReadMsg(&msg)
	if err != nil {
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
		return
	}

	switch msg.GetType() {
	case pb.CircuitRelay_HOP:
		r.handleHopStream(s, &msg)
	case pb.CircuitRelay_STOP:
		r.handleStopStream(s, &msg)
	case pb.CircuitRelay_CAN_HOP:
		r.handleCanHop(s, &msg)
	default:
		log.Warningf("unexpected relay handshake: %d", msg.GetType())
		r.handleError(s, pb.CircuitRelay_MALFORMED_MESSAGE)
	}
}

func (r *Relay) handleHopStream(s inet.Stream, msg *pb.CircuitRelay) {
	if !r.hop {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
		return
	}

	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	if src.ID != s.Conn().RemotePeer() {
		r.handleError(s, pb.CircuitRelay_HOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil {
		r.handleError(s, pb.CircuitRelay_HOP_DST_MULTIADDR_INVALID)
		return
	}

	if dst.ID == r.self {
		r.handleError(s, pb.CircuitRelay_HOP_CANT_RELAY_TO_SELF)
		return
	}

	// open stream
	ctp := r.host.Network().ConnsToPeer(dst.ID)
	if len(ctp) == 0 {
		if !r.active {
			r.handleError(s, pb.CircuitRelay_HOP_NO_CONN_TO_DST)
			return
		}

		ctx, cancel := context.WithTimeout(r.ctx, time.Second*10)
		defer cancel()
		err = r.host.Connect(ctx, dst)
		if err != nil {
			log.Debugf("error opening relay connection to %s: %s", dst.ID.Pretty(), err.Error())
			r.handleError(s, pb.CircuitRelay_HOP_CANT_DIAL_DST)
			return
		}
	}

	bs, err := r.host.NewStream(r.ctx, dst.ID, ProtoID)
	if err != nil {
		log.Debugf("error opening relay stream to %s: %s", dst.ID.Pretty(), err.Error())
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	// stop handshake
	rd := ggio.NewDelimitedReader(bs, maxMessageSize)
	wr := ggio.NewDelimitedWriter(bs)

	msg.Type = pb.CircuitRelay_STOP.Enum()

	err = wr.WriteMsg(msg)
	if err != nil {
		log.Debugf("error writing stop handshake: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	msg.Reset()

	err = rd.ReadMsg(msg)
	if err != nil {
		log.Debugf("error reading stop response: %s", err.Error())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetType() != pb.CircuitRelay_STATUS {
		log.Debugf("unexpected relay stop response: not a status message (%d)", msg.GetType())
		bs.Close()
		r.handleError(s, pb.CircuitRelay_HOP_CANT_OPEN_DST_STREAM)
		return
	}

	if msg.GetCode() != pb.CircuitRelay_SUCCESS {
		log.Debugf("relay stop failure: %d", msg.GetCode())
		bs.Close()
		r.handleError(s, msg.GetCode())
		return
	}

	err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
		bs.Close()
		s.Close()
		return
	}

	// relay connection
	log.Infof("relaying connection between %s and %s", src.ID.Pretty(), dst.ID.Pretty())

	go func() {
		_, err := io.Copy(s, bs)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		s.Close()
	}()

	go func() {
		_, err := io.Copy(bs, s)
		if err != io.EOF && err != nil {
			log.Debugf("relay copy error: %s", err)
		}
		bs.Close()
	}()
}

func (r *Relay) handleStopStream(s inet.Stream, msg *pb.CircuitRelay) {
	src, err := peerToPeerInfo(msg.GetSrcPeer())
	if err != nil || len(src.Addrs) == 0 {
		r.handleError(s, pb.CircuitRelay_STOP_SRC_MULTIADDR_INVALID)
		return
	}

	dst, err := peerToPeerInfo(msg.GetDstPeer())
	if err != nil || dst.ID != r.self {
		r.handleError(s, pb.CircuitRelay_STOP_DST_MULTIADDR_INVALID)
		return
	}

	log.Infof("relay connection from: %s", src.ID)

	r.host.Peerstore().AddAddrs(src.ID, src.Addrs, pstore.TempAddrTTL)

	select {
	case r.incoming <- &Conn{Stream: s, remote: src}:
	case <-time.After(RelayAcceptTimeout):
		r.handleError(s, pb.CircuitRelay_STOP_RELAY_REFUSED)
	}
}

func (r *Relay) handleCanHop(s inet.Stream, msg *pb.CircuitRelay) {
	var err error

	if r.hop {
		err = r.writeResponse(s, pb.CircuitRelay_SUCCESS)
	} else {
		err = r.writeResponse(s, pb.CircuitRelay_HOP_CANT_SPEAK_RELAY)
	}

	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}

	s.Close()
}

func (r *Relay) handleError(s inet.Stream, code pb.CircuitRelay_Status) {
	log.Warningf("relay error: %s (%d)", pb.CircuitRelay_Status_name[int32(code)], code)
	err := r.writeResponse(s, code)
	if err != nil {
		log.Debugf("error writing relay response: %s", err.Error())
	}
	s.Close()
}

func (r *Relay) writeResponse(s inet.Stream, code pb.CircuitRelay_Status) error {
	wr := ggio.NewDelimitedWriter(s)

	var msg pb.CircuitRelay
	msg.Type = pb.CircuitRelay_STATUS.Enum()
	msg.Code = code.Enum()

	return wr.WriteMsg(&msg)
}
<|endoftext|>


<commit_before>package content

import (
	"compress/gzip"
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	"github.com/kopia/kopia/internal/clock"
	"github.com/kopia/kopia/internal/gather"
	"github.com/kopia/kopia/internal/zaplogutil"
	"github.com/kopia/kopia/repo/blob"
)

const blobLoggerFlushThreshold = 4 << 20

// TextLogBlobPrefix is a prefix given to text logs stored in repositor.
const TextLogBlobPrefix = "_log_"

type internalLogManager struct {
	// +checkatomic
	enabled int32 // set by enable(), logger is ineffective until called

	// internalLogManager implements io.Writer and we must be able to write to the
	// repository asynchronously when the context is not provided.
	ctx context.Context // nolint:containedctx

	st             blob.Storage
	bc             *Crypter
	wg             sync.WaitGroup
	timeFunc       func() time.Time
	flushThreshold int
}

// Close closes the log manager.
func (m *internalLogManager) Close(ctx context.Context) {
	m.wg.Wait()
}

func (m *internalLogManager) encryptAndWriteLogBlob(prefix blob.ID, data gather.Bytes, closeFunc func()) {
	encrypted := gather.NewWriteBuffer()
	// Close happens in a goroutine

	blobID, err := m.bc.EncryptBLOB(data, prefix, "", encrypted)
	if err != nil {
		encrypted.Close()

		// this should not happen, also nothing can be done about this, we're not in a place where we can return error, log it.
		return
	}

	b := encrypted.Bytes()

	m.wg.Add(1)

	go func() {
		defer m.wg.Done()
		defer encrypted.Close()
		defer closeFunc()

		if err := m.st.PutBlob(m.ctx, blobID, b, blob.PutOptions{}); err != nil {
			// nothing can be done about this, we're not in a place where we can return error, log it.
			return
		}
	}()
}

// NewLogger creates new logger.
func (m *internalLogManager) NewLogger() *zap.SugaredLogger {
	var rnd [2]byte

	rand.Read(rnd[:]) // nolint:errcheck

	w := &internalLogger{
		m:      m,
		prefix: blob.ID(fmt.Sprintf("%v%v_%x", TextLogBlobPrefix, clock.Now().Local().Format("20060102150405"), rnd)),
	}

	return zap.New(zapcore.NewCore(
		zapcore.NewConsoleEncoder(zapcore.EncoderConfig{
			TimeKey:          "t",
			MessageKey:       "m",
			NameKey:          "n",
			EncodeTime:       zaplogutil.TimezoneAdjust(zaplogutil.PreciseTimeEncoder(), false),
			EncodeDuration:   zapcore.StringDurationEncoder,
			ConsoleSeparator: " ",
		}),
		w, zap.DebugLevel), zap.WithClock(zaplogutil.Clock())).Sugar()
}

// internalLogger represents a single log session that saves log files as blobs in the repository.
// The logger starts disabled and to actually persist logs enable() must be called.
type internalLogger struct {
	// +checkatomic
	nextChunkNumber int32 // chunk number incremented using atomic.AddInt32()

	m  *internalLogManager
	mu sync.Mutex

	// +checklocks:mu
	buf *gather.WriteBuffer
	// +checklocks:mu
	gzw *gzip.Writer

	startTime int64 // unix timestamp of the first log

	prefix blob.ID // +checklocksignore
}

func (m *internalLogManager) enable() {
	if m == nil {
		return
	}

	atomic.StoreInt32(&m.enabled, 1)
}

func (l *internalLogger) Write(b []byte) (int, error) {
	l.maybeEncryptAndWriteChunkUnlocked(l.addAndMaybeFlush(b))
	return len(b), nil
}

func (l *internalLogger) maybeEncryptAndWriteChunkUnlocked(data gather.Bytes, closeFunc func()) {
	if data.Length() == 0 {
		closeFunc()
		return
	}

	if atomic.LoadInt32(&l.m.enabled) == 0 {
		closeFunc()
		return
	}

	endTime := l.m.timeFunc().Unix()

	prefix := blob.ID(fmt.Sprintf("%v_%v_%v_%v_", l.prefix, l.startTime, endTime, atomic.AddInt32(&l.nextChunkNumber, 1)))

	l.m.encryptAndWriteLogBlob(prefix, data, closeFunc)
}

func (l *internalLogger) addAndMaybeFlush(b []byte) (payload gather.Bytes, closeFunc func()) {
	l.mu.Lock()
	defer l.mu.Unlock()

	w := l.ensureWriterInitializedLocked()

	_, err := w.Write(b)
	l.logUnexpectedError(err)

	if l.buf.Length() < l.m.flushThreshold {
		return gather.Bytes{}, func() {}
	}

	return l.flushAndResetLocked()
}

// +checklocks:l.mu
func (l *internalLogger) ensureWriterInitializedLocked() io.Writer {
	if l.gzw == nil {
		l.buf = gather.NewWriteBuffer()
		l.gzw = gzip.NewWriter(l.buf)
		l.startTime = l.m.timeFunc().Unix()
	}

	return l.gzw
}

// +checklocks:l.mu
func (l *internalLogger) flushAndResetLocked() (payload gather.Bytes, closeFunc func()) {
	if l.gzw == nil {
		return gather.Bytes{}, func() {}
	}

	l.logUnexpectedError(l.gzw.Flush())
	l.logUnexpectedError(l.gzw.Close())

	closeBuf := l.buf.Close
	res := l.buf.Bytes()

	l.buf = nil
	l.gzw = nil

	return res, closeBuf
}

func (l *internalLogger) logUnexpectedError(err error) {
	if err == nil {
		return
	}
}

func (l *internalLogger) Sync() error {
	l.mu.Lock()
	data, closeFunc := l.flushAndResetLocked()
	l.mu.Unlock()

	l.maybeEncryptAndWriteChunkUnlocked(data, closeFunc)

	return nil
}

// newInternalLogManager creates a new blobLogManager that will emit logs as repository blobs with a given prefix.
func newInternalLogManager(ctx context.Context, st blob.Storage, bc *Crypter) *internalLogManager {
	return &internalLogManager{
		ctx:            ctx,
		st:             st,
		bc:             bc,
		flushThreshold: blobLoggerFlushThreshold,
		timeFunc:       clock.Now,
	}
}
<commit_msg>chore(repository): fixed minor logger data race (#2025)<commit_after>package content

import (
	"compress/gzip"
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	"github.com/kopia/kopia/internal/clock"
	"github.com/kopia/kopia/internal/gather"
	"github.com/kopia/kopia/internal/zaplogutil"
	"github.com/kopia/kopia/repo/blob"
)

const blobLoggerFlushThreshold = 4 << 20

// TextLogBlobPrefix is a prefix given to text logs stored in repositor.
const TextLogBlobPrefix = "_log_"

type internalLogManager struct {
	// +checkatomic
	enabled int32 // set by enable(), logger is ineffective until called

	// internalLogManager implements io.Writer and we must be able to write to the
	// repository asynchronously when the context is not provided.
	ctx context.Context // nolint:containedctx

	st             blob.Storage
	bc             *Crypter
	wg             sync.WaitGroup
	timeFunc       func() time.Time
	flushThreshold int
}

// Close closes the log manager.
func (m *internalLogManager) Close(ctx context.Context) {
	m.wg.Wait()
}

func (m *internalLogManager) encryptAndWriteLogBlob(prefix blob.ID, data gather.Bytes, closeFunc func()) {
	encrypted := gather.NewWriteBuffer()
	// Close happens in a goroutine

	blobID, err := m.bc.EncryptBLOB(data, prefix, "", encrypted)
	if err != nil {
		encrypted.Close()

		// this should not happen, also nothing can be done about this, we're not in a place where we can return error, log it.
		return
	}

	b := encrypted.Bytes()

	m.wg.Add(1)

	go func() {
		defer m.wg.Done()
		defer encrypted.Close()
		defer closeFunc()

		if err := m.st.PutBlob(m.ctx, blobID, b, blob.PutOptions{}); err != nil {
			// nothing can be done about this, we're not in a place where we can return error, log it.
			return
		}
	}()
}

// NewLogger creates new logger.
func (m *internalLogManager) NewLogger() *zap.SugaredLogger {
	var rnd [2]byte

	rand.Read(rnd[:]) // nolint:errcheck

	w := &internalLogger{
		m:      m,
		prefix: blob.ID(fmt.Sprintf("%v%v_%x", TextLogBlobPrefix, clock.Now().Local().Format("20060102150405"), rnd)),
	}

	return zap.New(zapcore.NewCore(
		zapcore.NewConsoleEncoder(zapcore.EncoderConfig{
			TimeKey:          "t",
			MessageKey:       "m",
			NameKey:          "n",
			EncodeTime:       zaplogutil.TimezoneAdjust(zaplogutil.PreciseTimeEncoder(), false),
			EncodeDuration:   zapcore.StringDurationEncoder,
			ConsoleSeparator: " ",
		}),
		w, zap.DebugLevel), zap.WithClock(zaplogutil.Clock())).Sugar()
}

// internalLogger represents a single log session that saves log files as blobs in the repository.
// The logger starts disabled and to actually persist logs enable() must be called.
type internalLogger struct {
	// +checkatomic
	nextChunkNumber int32 // chunk number incremented using atomic.AddInt32()

	m  *internalLogManager
	mu sync.Mutex

	// +checklocks:mu
	buf *gather.WriteBuffer
	// +checklocks:mu
	gzw *gzip.Writer

	// +checklocks:mu
	startTime int64 // unix timestamp of the first log

	prefix blob.ID // +checklocksignore
}

func (m *internalLogManager) enable() {
	if m == nil {
		return
	}

	atomic.StoreInt32(&m.enabled, 1)
}

func (l *internalLogger) Write(b []byte) (int, error) {
	l.maybeEncryptAndWriteChunkUnlocked(l.addAndMaybeFlush(b))
	return len(b), nil
}

func (l *internalLogger) maybeEncryptAndWriteChunkUnlocked(data gather.Bytes, closeFunc func()) {
	if data.Length() == 0 {
		closeFunc()
		return
	}

	if atomic.LoadInt32(&l.m.enabled) == 0 {
		closeFunc()
		return
	}

	endTime := l.m.timeFunc().Unix()

	l.mu.Lock()
	st := l.startTime
	l.mu.Unlock()

	prefix := blob.ID(fmt.Sprintf("%v_%v_%v_%v_", l.prefix, st, endTime, atomic.AddInt32(&l.nextChunkNumber, 1)))

	l.m.encryptAndWriteLogBlob(prefix, data, closeFunc)
}

func (l *internalLogger) addAndMaybeFlush(b []byte) (payload gather.Bytes, closeFunc func()) {
	l.mu.Lock()
	defer l.mu.Unlock()

	w := l.ensureWriterInitializedLocked()

	_, err := w.Write(b)
	l.logUnexpectedError(err)

	if l.buf.Length() < l.m.flushThreshold {
		return gather.Bytes{}, func() {}
	}

	return l.flushAndResetLocked()
}

// +checklocks:l.mu
func (l *internalLogger) ensureWriterInitializedLocked() io.Writer {
	if l.gzw == nil {
		l.buf = gather.NewWriteBuffer()
		l.gzw = gzip.NewWriter(l.buf)
		l.startTime = l.m.timeFunc().Unix()
	}

	return l.gzw
}

// +checklocks:l.mu
func (l *internalLogger) flushAndResetLocked() (payload gather.Bytes, closeFunc func()) {
	if l.gzw == nil {
		return gather.Bytes{}, func() {}
	}

	l.logUnexpectedError(l.gzw.Flush())
	l.logUnexpectedError(l.gzw.Close())

	closeBuf := l.buf.Close
	res := l.buf.Bytes()

	l.buf = nil
	l.gzw = nil

	return res, closeBuf
}

func (l *internalLogger) logUnexpectedError(err error) {
	if err == nil {
		return
	}
}

func (l *internalLogger) Sync() error {
	l.mu.Lock()
	data, closeFunc := l.flushAndResetLocked()
	l.mu.Unlock()

	l.maybeEncryptAndWriteChunkUnlocked(data, closeFunc)

	return nil
}

// newInternalLogManager creates a new blobLogManager that will emit logs as repository blobs with a given prefix.
func newInternalLogManager(ctx context.Context, st blob.Storage, bc *Crypter) *internalLogManager {
	return &internalLogManager{
		ctx:            ctx,
		st:             st,
		bc:             bc,
		flushThreshold: blobLoggerFlushThreshold,
		timeFunc:       clock.Now,
	}
}
<|endoftext|>


<commit_before>package request_logger_mw

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/tilteng/go-api-router/api_router"
	"github.com/tilteng/go-logger/logger"
)

type LogBodyFilterFn func(context.Context, []byte) []byte

func (self LogBodyFilterFn) FilterBody(ctx context.Context, bytes []byte) []byte {
	return self(ctx, bytes)
}

type LogBodyFilter interface {
	FilterBody(context.Context, []byte) []byte
}

type LogHeadersFilterFn func(context.Context, http.Header) http.Header

func (self LogHeadersFilterFn) FilterHeaders(ctx context.Context, hdrs http.Header) http.Header {
	return self(ctx, hdrs)
}

type LogHeadersFilter interface {
	FilterHeaders(context.Context, http.Header) http.Header
}

type RequestLoggerOpts struct {
	LogBodyFilter    LogBodyFilter
	LogHeadersFilter LogHeadersFilter
	Logger           logger.CtxLogger
	Disable          bool
}

type RequestLoggerMiddleware struct {
	opts *RequestLoggerOpts
}

func (self *RequestLoggerMiddleware) NewWrapper(ctx context.Context, opts ...interface{}) *RequestLoggerWrapper {
	var opt *RequestLoggerOpts

	for _, opt_map_i := range opts {
		var ok bool
		opt, ok = opt_map_i.(*RequestLoggerOpts)
		if ok {
			break
		}
	}

	if opt == nil {
		opt = &RequestLoggerOpts{
			Logger: self.opts.Logger,
		}
	}

	if opt.Disable {
		return nil
	}

	return &RequestLoggerWrapper{
		base_opts: self.opts,
		opts:      opt,
	}
}

type RequestLoggerWrapper struct {
	opts      *RequestLoggerOpts
	base_opts *RequestLoggerOpts
}

func (self *RequestLoggerMiddleware) SetLogger(logger logger.CtxLogger) *RequestLoggerMiddleware {
	self.opts.Logger = logger
	return self
}

func (self *RequestLoggerWrapper) formatBody(ctx context.Context, body []byte) interface{} {
	if self.opts.LogBodyFilter != nil {
		body = self.opts.LogBodyFilter.FilterBody(ctx, body)
	}
	if self.base_opts.LogBodyFilter != nil {
		body = self.opts.LogBodyFilter.FilterBody(ctx, body)
	}
	// This is all hacky... just to get something logged for now
	// in a reasonable format.
	var body_obj interface{}

	if err := json.Unmarshal(body, &body_obj); err == nil {
		return body
	}

	body_str := string(body)

	if json_body, err := json.Marshal(body_str); err == nil {
		return json_body
	}

	return `"` + body_str + `"`
}

func (self *RequestLoggerWrapper) formatHeaders(ctx context.Context, hdrs http.Header) interface{} {
	if self.opts.LogHeadersFilter != nil {
		hdrs = self.opts.LogHeadersFilter.FilterHeaders(ctx, hdrs)
	}

	if self.base_opts.LogHeadersFilter != nil {
		hdrs = self.opts.LogHeadersFilter.FilterHeaders(ctx, hdrs)
	}

	json_hdrs, _ := json.Marshal(hdrs)
	return json_hdrs
}

func (self *RequestLoggerWrapper) Wrap(next api_router.RouteFn) api_router.RouteFn {
	return func(ctx context.Context) {
		rctx := api_router.RequestContextFromContext(ctx)
		rt := rctx.CurrentRoute()
		http_req := rctx.HTTPRequest()
		method := http_req.Method
		path := http_req.URL.String()

		body, err := rctx.BodyCopy()
		if err != nil {
			panic(fmt.Sprintf("Couldn't read body: %+v", err))
		}
		if self.opts.Logger != nil && !self.opts.Disable {
			self.opts.Logger.LogDebugf(
				ctx,
				`Received request: route:{"method":"%s","route":"%s","path":"%s"} headers:%s body:%s`,
				method,
				rt.FullPath(),
				path,
				self.formatHeaders(ctx, rctx.HTTPRequest().Header),
				self.formatBody(ctx, body),
			)
		}

		// Do something with body
		next(ctx)

		writer := rctx.ResponseWriter()
		body = writer.ResponseCopy()

		self.opts.Logger.LogDebugf(
			ctx,
			`Sent response: route:{"status":%d,"method":"%s","route":"%s","path":"%s"} headers:%s body:%s`,
			writer.Status(),
			method,
			rt.FullPath(),
			path,
			self.formatHeaders(ctx, rctx.ResponseWriter().Header()),
			self.formatBody(ctx, body),
		)
	}
}

func NewMiddleware(opts *RequestLoggerOpts) *RequestLoggerMiddleware {
	if opts == nil {
		opts = &RequestLoggerOpts{}
	}
	if opts.Logger == nil {
		opts.Logger = logger.DefaultStdoutCtxLogger()
	}
	return &RequestLoggerMiddleware{
		opts: opts,
	}
}
<commit_msg>Log query params<commit_after>package request_logger_mw

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"

	"github.com/tilteng/go-api-router/api_router"
	"github.com/tilteng/go-logger/logger"
)

type LogBodyFilterFn func(context.Context, []byte) []byte

func (self LogBodyFilterFn) FilterBody(ctx context.Context, bytes []byte) []byte {
	return self(ctx, bytes)
}

type LogBodyFilter interface {
	FilterBody(context.Context, []byte) []byte
}

type LogHeadersFilterFn func(context.Context, http.Header) http.Header

func (self LogHeadersFilterFn) FilterHeaders(ctx context.Context, hdrs http.Header) http.Header {
	return self(ctx, hdrs)
}

type LogHeadersFilter interface {
	FilterHeaders(context.Context, http.Header) http.Header
}

type RequestLoggerOpts struct {
	LogBodyFilter    LogBodyFilter
	LogHeadersFilter LogHeadersFilter
	Logger           logger.CtxLogger
	Disable          bool
}

type RequestLoggerMiddleware struct {
	opts *RequestLoggerOpts
}

func (self *RequestLoggerMiddleware) NewWrapper(ctx context.Context, opts ...interface{}) *RequestLoggerWrapper {
	var opt *RequestLoggerOpts

	for _, opt_map_i := range opts {
		var ok bool
		opt, ok = opt_map_i.(*RequestLoggerOpts)
		if ok {
			break
		}
	}

	if opt == nil {
		opt = &RequestLoggerOpts{
			Logger: self.opts.Logger,
		}
	}

	if opt.Disable {
		return nil
	}

	return &RequestLoggerWrapper{
		base_opts: self.opts,
		opts:      opt,
	}
}

type RequestLoggerWrapper struct {
	opts      *RequestLoggerOpts
	base_opts *RequestLoggerOpts
}

func (self *RequestLoggerMiddleware) SetLogger(logger logger.CtxLogger) *RequestLoggerMiddleware {
	self.opts.Logger = logger
	return self
}

func (self *RequestLoggerWrapper) formatBody(ctx context.Context, body []byte) interface{} {
	if self.opts.LogBodyFilter != nil {
		body = self.opts.LogBodyFilter.FilterBody(ctx, body)
	}
	if self.base_opts.LogBodyFilter != nil {
		body = self.opts.LogBodyFilter.FilterBody(ctx, body)
	}
	// This is all hacky... just to get something logged for now
	// in a reasonable format.
	var body_obj interface{}

	if err := json.Unmarshal(body, &body_obj); err == nil {
		return body
	}

	body_str := string(body)

	if json_body, err := json.Marshal(body_str); err == nil {
		return json_body
	}

	return `"` + body_str + `"`
}

func (self *RequestLoggerWrapper) formatQuery(ctx context.Context, v url.Values) string {
	j, _ := json.Marshal(v)
	return string(j)
}

func (self *RequestLoggerWrapper) formatHeaders(ctx context.Context, hdrs http.Header) interface{} {
	if self.opts.LogHeadersFilter != nil {
		hdrs = self.opts.LogHeadersFilter.FilterHeaders(ctx, hdrs)
	}

	if self.base_opts.LogHeadersFilter != nil {
		hdrs = self.opts.LogHeadersFilter.FilterHeaders(ctx, hdrs)
	}

	json_hdrs, _ := json.Marshal(hdrs)
	return json_hdrs
}

func (self *RequestLoggerWrapper) Wrap(next api_router.RouteFn) api_router.RouteFn {
	return func(ctx context.Context) {
		rctx := api_router.RequestContextFromContext(ctx)
		rt := rctx.CurrentRoute()
		http_req := rctx.HTTPRequest()
		method := http_req.Method
		path := http_req.URL.EscapedPath()

		body, err := rctx.BodyCopy()
		if err != nil {
			panic(fmt.Sprintf("Couldn't read body: %+v", err))
		}
		if self.opts.Logger != nil && !self.opts.Disable {
			self.opts.Logger.LogDebugf(
				ctx,
				`Received request: {"route":{"method":"%s","route":"%s","path":"%s","query":%s},"headers":%s,"body":%s}`,
				method,
				rt.FullPath(),
				path,
				self.formatQuery(ctx, http_req.URL.Query()),
				self.formatHeaders(ctx, rctx.HTTPRequest().Header),
				self.formatBody(ctx, body),
			)
		}

		// Do something with body
		next(ctx)

		writer := rctx.ResponseWriter()
		body = writer.ResponseCopy()

		self.opts.Logger.LogDebugf(
			ctx,
			`Sent response: {"status":%d,"route":{"method":"%s","route":"%s","path":"%s"},"headers":%s,"body":%s}`,
			writer.Status(),
			method,
			rt.FullPath(),
			path,
			self.formatHeaders(ctx, rctx.ResponseWriter().Header()),
			self.formatBody(ctx, body),
		)
	}
}

func NewMiddleware(opts *RequestLoggerOpts) *RequestLoggerMiddleware {
	if opts == nil {
		opts = &RequestLoggerOpts{}
	}
	if opts.Logger == nil {
		opts.Logger = logger.DefaultStdoutCtxLogger()
	}
	return &RequestLoggerMiddleware{
		opts: opts,
	}
}
<|endoftext|>


<commit_before>package google

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"

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

func TestAccStorage_basic(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderDefaults(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "US"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "false"),
				),
			},
		},
	})
}

func TestAccStorageCustomAttributes(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "EU"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "true"),
				),
			},
		},
	})
}

func TestAccStorageStorageClass(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acc-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "STANDARD"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "STANDARD"),
				),
			},
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "NEARLINE"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "NEARLINE"),
				),
			},
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "DURABLE_REDUCED_AVAILABILITY"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "DURABLE_REDUCED_AVAILABILITY"),
				),
			},
		},
	})
}

func TestAccStorageBucketUpdate(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderDefaults(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "US"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "false"),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "predefined_acl", "publicReadWrite"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "EU"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "true"),
				),
			},
		},
	})
}

func TestAccStorageForceDestroy(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketPutItem(bucketName),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes("idontexist"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketMissing(bucketName),
				),
			},
		},
	})
}

func testAccCheckCloudStorageBucketExists(n string, bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Project_ID is set")
		}

		config := testAccProvider.Meta().(*Config)

		found, err := config.clientStorage.Buckets.Get(rs.Primary.ID).Do()
		if err != nil {
			return err
		}

		if found.Id != rs.Primary.ID {
			return fmt.Errorf("Bucket not found")
		}

		if found.Name != bucketName {
			return fmt.Errorf("expected name %s, got %s", bucketName, found.Name)
		}
		return nil
	}
}

func testAccCheckCloudStorageBucketPutItem(bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		data := bytes.NewBufferString("test")
		dataReader := bytes.NewReader(data.Bytes())
		object := &storage.Object{Name: "bucketDestroyTestFile"}

		// This needs to use Media(io.Reader) call, otherwise it does not go to /upload API and fails
		if res, err := config.clientStorage.Objects.Insert(bucketName, object).Media(dataReader).Do(); err == nil {
			fmt.Printf("Created object %v at location %v\n\n", res.Name, res.SelfLink)
		} else {
			return fmt.Errorf("Objects.Insert failed: %v", err)
		}

		return nil
	}
}

func testAccCheckCloudStorageBucketMissing(bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		_, err := config.clientStorage.Buckets.Get(bucketName).Do()
		if err == nil {
			return fmt.Errorf("Found %s", bucketName)
		}

		if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
			return nil
		}

		return err
	}
}

func testAccGoogleStorageDestroy(s *terraform.State) error {
	config := testAccProvider.Meta().(*Config)

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "google_storage_bucket" {
			continue
		}

		_, err := config.clientStorage.Buckets.Get(rs.Primary.ID).Do()
		if err == nil {
			return fmt.Errorf("Bucket still exists")
		}
	}

	return nil
}

func testGoogleStorageBucketsReaderDefaults(bucketName string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
}
`, bucketName)
}

func testGoogleStorageBucketsReaderCustomAttributes(bucketName string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
	predefined_acl = "publicReadWrite"
	location = "EU"
	force_destroy = "true"
}
`, bucketName)
}

func testGoogleStorageBucketsReaderStorageClass(bucketName string, storageClass string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
	storage_class = "%s"
}
`, bucketName, storageClass)
}
<commit_msg>doc: gcs - Update storage_class documentation and tests<commit_after>package google

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"

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

func TestAccStorage_basic(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderDefaults(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "US"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "false"),
				),
			},
		},
	})
}

func TestAccStorageCustomAttributes(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "EU"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "true"),
				),
			},
		},
	})
}

func TestAccStorageStorageClass(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acc-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "MULTI_REGIONAL"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "MULTI_REGIONAL"),
				),
			},
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "NEARLINE"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "NEARLINE"),
				),
			},
			{
				Config: testGoogleStorageBucketsReaderStorageClass(bucketName, "REGIONAL"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "storage_class", "REGIONAL"),
				),
			},
		},
	})
}

func TestAccStorageBucketUpdate(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderDefaults(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "US"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "false"),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "predefined_acl", "publicReadWrite"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "location", "EU"),
					resource.TestCheckResourceAttr(
						"google_storage_bucket.bucket", "force_destroy", "true"),
				),
			},
		},
	})
}

func TestAccStorageForceDestroy(t *testing.T) {
	bucketName := fmt.Sprintf("tf-test-acl-bucket-%d", acctest.RandInt())

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccGoogleStorageDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketExists(
						"google_storage_bucket.bucket", bucketName),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes(bucketName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketPutItem(bucketName),
				),
			},
			resource.TestStep{
				Config: testGoogleStorageBucketsReaderCustomAttributes("idontexist"),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckCloudStorageBucketMissing(bucketName),
				),
			},
		},
	})
}

func testAccCheckCloudStorageBucketExists(n string, bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No Project_ID is set")
		}

		config := testAccProvider.Meta().(*Config)

		found, err := config.clientStorage.Buckets.Get(rs.Primary.ID).Do()
		if err != nil {
			return err
		}

		if found.Id != rs.Primary.ID {
			return fmt.Errorf("Bucket not found")
		}

		if found.Name != bucketName {
			return fmt.Errorf("expected name %s, got %s", bucketName, found.Name)
		}
		return nil
	}
}

func testAccCheckCloudStorageBucketPutItem(bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		data := bytes.NewBufferString("test")
		dataReader := bytes.NewReader(data.Bytes())
		object := &storage.Object{Name: "bucketDestroyTestFile"}

		// This needs to use Media(io.Reader) call, otherwise it does not go to /upload API and fails
		if res, err := config.clientStorage.Objects.Insert(bucketName, object).Media(dataReader).Do(); err == nil {
			fmt.Printf("Created object %v at location %v\n\n", res.Name, res.SelfLink)
		} else {
			return fmt.Errorf("Objects.Insert failed: %v", err)
		}

		return nil
	}
}

func testAccCheckCloudStorageBucketMissing(bucketName string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		_, err := config.clientStorage.Buckets.Get(bucketName).Do()
		if err == nil {
			return fmt.Errorf("Found %s", bucketName)
		}

		if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
			return nil
		}

		return err
	}
}

func testAccGoogleStorageDestroy(s *terraform.State) error {
	config := testAccProvider.Meta().(*Config)

	for _, rs := range s.RootModule().Resources {
		if rs.Type != "google_storage_bucket" {
			continue
		}

		_, err := config.clientStorage.Buckets.Get(rs.Primary.ID).Do()
		if err == nil {
			return fmt.Errorf("Bucket still exists")
		}
	}

	return nil
}

func testGoogleStorageBucketsReaderDefaults(bucketName string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
}
`, bucketName)
}

func testGoogleStorageBucketsReaderCustomAttributes(bucketName string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
	predefined_acl = "publicReadWrite"
	location = "EU"
	force_destroy = "true"
}
`, bucketName)
}

func testGoogleStorageBucketsReaderStorageClass(bucketName string, storageClass string) string {
	return fmt.Sprintf(`
resource "google_storage_bucket" "bucket" {
	name = "%s"
	storage_class = "%s"
}
`, bucketName, storageClass)
}
<|endoftext|>