<commit_before>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 command

import (
	"github.com/mitchellh/cli"
	"github.com/pquerna/distsync/common"
	"github.com/pquerna/distsync/crypto"
	"github.com/pquerna/distsync/storage"

	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

type Download struct {
	Ui   cli.Ui
	stop error
	dl   storage.PersistentDownloader
	conf *common.Conf
}

func (c *Download) Help() string {
	helpText := `
Usage: distsync download [options] file

  Downloads specified files to your local directory.

Options:

  -conf=~/.distsyncd         Read specific configuration file.
`
	return strings.TrimSpace(helpText)
}

func (c *Download) Run(args []string) int {
	var confFile string

	cmdFlags := flag.NewFlagSet("download", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	cmdFlags.StringVar(&confFile, "conf", "~/.distsyncd", "Configuration path.")

	err := cmdFlags.Parse(args)
	if err != nil {
		return 1
	}

	c.conf, err = common.ConfFromFile(confFile)
	if err != nil {
		c.Ui.Error("Configuration failure: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	files := cmdFlags.Args()
	if len(files) == 0 {
		c.Ui.Error("At least one file to download must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	}

	download, err := c.getFilesToDownload(files)
	if err != nil {
		c.Ui.Error("Error Getting files to download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	if len(download) == 0 {
		c.Ui.Error(fmt.Sprintf("Files not found: %v", files))
		c.Ui.Error("")
		return 1
	}

	c.dl, err = storage.NewPersistentDownloader(c.conf)
	if err != nil {
		c.Ui.Error("Error configuring download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	err = c.dl.Start()
	if err != nil {
		c.Ui.Error("Error starting download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	var wg sync.WaitGroup

	for _, file := range download {
		wg.Add(1)
		go c.downloadFile(&wg, file)
	}
	wg.Wait()

	if c.stop != nil {
		c.Ui.Error("Download failed: " + c.stop.Error())
		c.Ui.Error("")
		return 1
	}
	return 0
}

func (c *Download) getFilesToDownload(fnames []string) ([]storage.FileInfo, error) {
	ec, err := crypto.NewFromConf(c.conf)
	if err != nil {
		return nil, err
	}

	s, err := storage.NewFromConf(c.conf)
	if err != nil {
		return nil, err
	}

	storedFiles, err := s.List(ec)
	if err != nil {
		return nil, err
	}

	download := make([]storage.FileInfo, 0, 1)
	for _, file := range storedFiles {
		for _, fname := range fnames {
			// TODO: meh.
			if file.Name == fname {
				download = append(download, file)
			}
		}
	}

	return download, nil
}

func (c *Download) downloadFile(wg *sync.WaitGroup, fi storage.FileInfo) {
	defer wg.Done()

	err := c._downloadFile(fi)
	if err != nil {
		_, ok := err.(*stopError)

		if !ok {
			c.stop = err
		}
	}
}

func (c *Download) _downloadFile(fi storage.FileInfo) error {
	ec, err := crypto.NewFromConf(c.conf)

	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	// TODO: consider io.Pipe() ?
	tmpFileEnc, err := ioutil.TempFile(cwd, ".distsync-e")
	if err != nil {
		return err
	}
	defer func() {
		tmpFileEnc.Close()
		os.Remove(tmpFileEnc.Name())
	}()
	if c.stop != nil {
		return &stopError{}
	}

	c.Ui.Info("Downloading " + fi.EncryptedName + " -> " + fi.Name)

	err = c.dl.Download(fi.EncryptedName, tmpFileEnc)
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	tmpFile, err := ioutil.TempFile(cwd, ".distsync")
	if err != nil {
		return err
	}
	defer func() {
		tmpFile.Close()
		os.Remove(tmpFile.Name())
	}()
	if c.stop != nil {
		return &stopError{}
	}

	_, err = tmpFileEnc.Seek(0, 0)
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	err = ec.Decrypt(tmpFileEnc, tmpFile)

	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	err = os.Rename(tmpFile.Name(), fi.Name)
	if err != nil {
		return err
	}

	err = os.Chtimes(fi.Name, fi.LastModified, fi.LastModified)
	if err != nil {
		return err
	}

	return nil
}

func (c *Download) Synopsis() string {
	return "Downloads files from distsync"
}
<commit_msg>print what files cant be found<commit_after>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 command

import (
	"github.com/mitchellh/cli"
	"github.com/pquerna/distsync/common"
	"github.com/pquerna/distsync/crypto"
	"github.com/pquerna/distsync/storage"

	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"sync"
)

type Download struct {
	Ui   cli.Ui
	stop error
	dl   storage.PersistentDownloader
	conf *common.Conf
}

func (c *Download) Help() string {
	helpText := `
Usage: distsync download [options] file

  Downloads specified files to your local directory.

Options:

  -conf=~/.distsyncd         Read specific configuration file.
`
	return strings.TrimSpace(helpText)
}

func diffslice(a []string, b []string) []string {
	rv := make([]string, 0)
	tmp := make(map[string]int)

	for _, v := range a {
		tmp[v] = 1
	}

	for _, v := range b {
		tmp[v] += 1
	}

	for k, v := range tmp {
		if v == 1 {
			rv = append(rv, k)
		}
	}

	return rv
}

func (c *Download) Run(args []string) int {
	var confFile string

	cmdFlags := flag.NewFlagSet("download", flag.ContinueOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
	cmdFlags.StringVar(&confFile, "conf", "~/.distsyncd", "Configuration path.")

	err := cmdFlags.Parse(args)
	if err != nil {
		return 1
	}

	c.conf, err = common.ConfFromFile(confFile)
	if err != nil {
		c.Ui.Error("Configuration failure: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	files := cmdFlags.Args()
	if len(files) == 0 {
		c.Ui.Error("At least one file to download must be specified.")
		c.Ui.Error("")
		c.Ui.Error(c.Help())
		return 1
	}

	download, err := c.getFilesToDownload(files)
	if err != nil {
		c.Ui.Error("Error Getting files to download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	if len(download) != len(files) {
		dlfiles := make([]string, 0, len(download))
		for _, dl := range download {
			dlfiles = append(dlfiles, dl.Name)
		}
		diff := diffslice(files, dlfiles)
		c.Ui.Error(fmt.Sprintf("Files not found: %v", diff))
		c.Ui.Error("")
		return 1
	}

	c.dl, err = storage.NewPersistentDownloader(c.conf)
	if err != nil {
		c.Ui.Error("Error configuring download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	err = c.dl.Start()
	if err != nil {
		c.Ui.Error("Error starting download: " + err.Error())
		c.Ui.Error("")
		return 1
	}

	var wg sync.WaitGroup

	for _, file := range download {
		wg.Add(1)
		go c.downloadFile(&wg, file)
	}
	wg.Wait()

	if c.stop != nil {
		c.Ui.Error("Download failed: " + c.stop.Error())
		c.Ui.Error("")
		return 1
	}
	return 0
}

func (c *Download) getFilesToDownload(fnames []string) ([]storage.FileInfo, error) {
	ec, err := crypto.NewFromConf(c.conf)
	if err != nil {
		return nil, err
	}

	s, err := storage.NewFromConf(c.conf)
	if err != nil {
		return nil, err
	}

	storedFiles, err := s.List(ec)
	if err != nil {
		return nil, err
	}

	download := make([]storage.FileInfo, 0, 1)
	for _, file := range storedFiles {
		for _, fname := range fnames {
			// TODO: meh.
			if file.Name == fname {
				download = append(download, file)
			}
		}
	}

	return download, nil
}

func (c *Download) downloadFile(wg *sync.WaitGroup, fi storage.FileInfo) {
	defer wg.Done()

	err := c._downloadFile(fi)
	if err != nil {
		_, ok := err.(*stopError)

		if !ok {
			c.stop = err
		}
	}
}

func (c *Download) _downloadFile(fi storage.FileInfo) error {
	ec, err := crypto.NewFromConf(c.conf)

	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	// TODO: consider io.Pipe() ?
	tmpFileEnc, err := ioutil.TempFile(cwd, ".distsync-e")
	if err != nil {
		return err
	}
	defer func() {
		tmpFileEnc.Close()
		os.Remove(tmpFileEnc.Name())
	}()
	if c.stop != nil {
		return &stopError{}
	}

	c.Ui.Info("Downloading " + fi.EncryptedName + " -> " + fi.Name)

	err = c.dl.Download(fi.EncryptedName, tmpFileEnc)
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	tmpFile, err := ioutil.TempFile(cwd, ".distsync")
	if err != nil {
		return err
	}
	defer func() {
		tmpFile.Close()
		os.Remove(tmpFile.Name())
	}()
	if c.stop != nil {
		return &stopError{}
	}

	_, err = tmpFileEnc.Seek(0, 0)
	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	err = ec.Decrypt(tmpFileEnc, tmpFile)

	if err != nil {
		return err
	}
	if c.stop != nil {
		return &stopError{}
	}

	err = os.Rename(tmpFile.Name(), fi.Name)
	if err != nil {
		return err
	}

	err = os.Chtimes(fi.Name, fi.LastModified, fi.LastModified)
	if err != nil {
		return err
	}

	return nil
}

func (c *Download) Synopsis() string {
	return "Downloads files from distsync"
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package api_test

import (
	"net/http"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api"
	apihttp "github.com/juju/juju/api/http"
	apihttptesting "github.com/juju/juju/api/http/testing"
	jujutesting "github.com/juju/juju/juju/testing"
)

type httpSuite struct {
	apihttptesting.HTTPSuite
	jujutesting.JujuConnSuite
}

var _ = gc.Suite(&httpSuite{})

func (s *httpSuite) SetUpTest(c *gc.C) {
	s.HTTPSuite.SetUpTest(c)
	s.JujuConnSuite.SetUpTest(c)

	// This determines the client used in SendHTTPRequest().
	s.PatchValue(api.NewHTTPClient,
		func(*api.State) apihttp.HTTPClient {
			return s.Fake
		},
	)
}

func (s *httpSuite) TestNewHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, jc.ErrorIsNil)

	s.CheckRequest(c, req, "GET", "somefacade")
}

func (s *httpSuite) TestNewHTTPClientCorrectTransport(c *gc.C) {
	httpClient := s.APIState.NewHTTPClient()

	c.Assert(httpClient.Transport, gc.NotNil)
	c.Assert(httpClient.Transport, gc.FitsTypeOf, (*http.Transport)(nil))
	config := httpClient.Transport.(*http.Transport).TLSClientConfig

	c.Check(config.RootCAs, gc.NotNil)
}

func (s *httpSuite) TestNewHTTPClientValidatesCert(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	httpClient := s.APIState.NewHTTPClient()
	resp, err := httpClient.Do(req)
	c.Assert(err, jc.ErrorIsNil)

	c.Check(resp.StatusCode, gc.Equals, http.StatusNotFound)
}

func (s *httpSuite) TestSendHTTPRequestSuccess(c *gc.C) {
	req, resp, err := s.APIState.SendHTTPRequest("somefacade", nil)
	c.Assert(err, jc.ErrorIsNil)

	s.Fake.CheckCalled(c, req, resp)
}
<commit_msg>Add missing methods to a test suite.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package api_test

import (
	"net/http"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api"
	apihttp "github.com/juju/juju/api/http"
	apihttptesting "github.com/juju/juju/api/http/testing"
	jujutesting "github.com/juju/juju/juju/testing"
)

type httpSuite struct {
	apihttptesting.HTTPSuite
	jujutesting.JujuConnSuite
}

var _ = gc.Suite(&httpSuite{})

func (s *httpSuite) SetUpSuite(c *gc.C) {
	s.HTTPSuite.SetUpSuite(c)
	s.JujuConnSuite.SetUpSuite(c)
}

func (s *httpSuite) TearDownSuite(c *gc.C) {
	s.HTTPSuite.TearDownSuite(c)
	s.JujuConnSuite.TearDownSuite(c)
}

func (s *httpSuite) SetUpTest(c *gc.C) {
	s.HTTPSuite.SetUpTest(c)
	s.JujuConnSuite.SetUpTest(c)

	// This determines the client used in SendHTTPRequest().
	s.PatchValue(api.NewHTTPClient,
		func(*api.State) apihttp.HTTPClient {
			return s.Fake
		},
	)
}

func (s *httpSuite) TearDownTest(c *gc.C) {
	s.HTTPSuite.TearDownTest(c)
	s.JujuConnSuite.TearDownTest(c)
}

func (s *httpSuite) TestNewHTTPRequestSuccess(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	c.Assert(err, jc.ErrorIsNil)

	s.CheckRequest(c, req, "GET", "somefacade")
}

func (s *httpSuite) TestNewHTTPClientCorrectTransport(c *gc.C) {
	httpClient := s.APIState.NewHTTPClient()

	c.Assert(httpClient.Transport, gc.NotNil)
	c.Assert(httpClient.Transport, gc.FitsTypeOf, (*http.Transport)(nil))
	config := httpClient.Transport.(*http.Transport).TLSClientConfig

	c.Check(config.RootCAs, gc.NotNil)
}

func (s *httpSuite) TestNewHTTPClientValidatesCert(c *gc.C) {
	req, err := s.APIState.NewHTTPRequest("GET", "somefacade")
	httpClient := s.APIState.NewHTTPClient()
	resp, err := httpClient.Do(req)
	c.Assert(err, jc.ErrorIsNil)

	c.Check(resp.StatusCode, gc.Equals, http.StatusNotFound)
}

func (s *httpSuite) TestSendHTTPRequestSuccess(c *gc.C) {
	req, resp, err := s.APIState.SendHTTPRequest("somefacade", nil)
	c.Assert(err, jc.ErrorIsNil)

	s.Fake.CheckCalled(c, req, resp)
}
<|endoftext|>


<commit_before>// Copyright ©2018 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mds

import (
	"math"

	"gonum.org/v1/gonum/blas/blas64"
	"gonum.org/v1/gonum/mat"
)

// TorgersonScaling converts a dissimilarity matrix to a matrix containing
// Euclidean coordinates. TorgersonScaling places the coordinates in dst and
// returns it and the number of positive Eigenvalues if successful.
// If the scaling is not successful, dst will be empty upon return.
// When the scaling is successful, dst will be resized to k columns wide.
// Eigenvalues will be copied into eigdst and returned as eig if it is provided.
//
// TorgersonScaling will panic if dst is not empty.
func TorgersonScaling(dst *mat.Dense, eigdst []float64, dis mat.Symmetric) (k int, eig []float64) {
	// https://doi.org/10.1007/0-387-28981-X_12

	n := dis.Symmetric()
	if dst.IsEmpty() {
		dst.ReuseAs(n, n)
	} else {
		panic("mds: receiver matrix not empty")
	}

	b := mat.NewSymDense(n, nil)
	for i := 0; i < n; i++ {
		for j := i; j < n; j++ {
			v := dis.At(i, j)
			v *= v
			b.SetSym(i, j, v)
		}
	}
	c := mat.NewSymDense(n, nil)
	s := -1 / float64(n)
	for i := 0; i < n; i++ {
		c.SetSym(i, i, 1+s)
		for j := i + 1; j < n; j++ {
			c.SetSym(i, j, s)
		}
	}
	dst.Product(c, b, c)
	for i := 0; i < n; i++ {
		for j := i; j < n; j++ {
			b.SetSym(i, j, -0.5*dst.At(i, j))
		}
	}

	var ed mat.EigenSym
	ok := ed.Factorize(b, true)
	if !ok {
		return 0, eigdst
	}
	ed.VectorsTo(dst)
	vals := ed.Values(nil)
	reverse(vals, dst.RawMatrix())
	copy(eigdst, vals)

	for i, v := range vals {
		if v < 0 {
			vals[i] = 0
			continue
		}
		k = i + 1
		vals[i] = math.Sqrt(v)
	}

	var tmp mat.Dense
	tmp.Mul(dst, mat.NewDiagonalRect(n, k, vals[:k]))
	*dst = *dst.Slice(0, n, 0, k).(*mat.Dense)
	dst.Copy(&tmp)

	return k, eigdst
}

func reverse(values []float64, vectors blas64.General) {
	for i, j := 0, len(values)-1; i < j; i, j = i+1, j-1 {
		values[i], values[j] = values[j], values[i]
		blas64.Swap(blas64.Vector{N: vectors.Rows, Inc: vectors.Stride, Data: vectors.Data[i:]},
			blas64.Vector{N: vectors.Rows, Inc: vectors.Stride, Data: vectors.Data[j:]})
	}
}
<commit_msg>stat/mds: add warning about ED numerical stability<commit_after>// Copyright ©2018 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mds

import (
	"math"

	"gonum.org/v1/gonum/blas/blas64"
	"gonum.org/v1/gonum/mat"
)

// TorgersonScaling converts a dissimilarity matrix to a matrix containing
// Euclidean coordinates. TorgersonScaling places the coordinates in dst and
// returns it and the number of positive Eigenvalues if successful.
// Note that Eigen Decomposition is numerically unstable and so Eigenvalues
// near zero should be examined and the value returned for k is advisory only.
// If the scaling is not successful, dst will be empty upon return.
// When the scaling is successful, dst will be resized to k columns wide.
// Eigenvalues will be copied into eigdst and returned as eig if it is provided.
//
// TorgersonScaling will panic if dst is not empty.
func TorgersonScaling(dst *mat.Dense, eigdst []float64, dis mat.Symmetric) (k int, eig []float64) {
	// https://doi.org/10.1007/0-387-28981-X_12

	n := dis.Symmetric()
	if dst.IsEmpty() {
		dst.ReuseAs(n, n)
	} else {
		panic("mds: receiver matrix not empty")
	}

	b := mat.NewSymDense(n, nil)
	for i := 0; i < n; i++ {
		for j := i; j < n; j++ {
			v := dis.At(i, j)
			v *= v
			b.SetSym(i, j, v)
		}
	}
	c := mat.NewSymDense(n, nil)
	s := -1 / float64(n)
	for i := 0; i < n; i++ {
		c.SetSym(i, i, 1+s)
		for j := i + 1; j < n; j++ {
			c.SetSym(i, j, s)
		}
	}
	dst.Product(c, b, c)
	for i := 0; i < n; i++ {
		for j := i; j < n; j++ {
			b.SetSym(i, j, -0.5*dst.At(i, j))
		}
	}

	var ed mat.EigenSym
	ok := ed.Factorize(b, true)
	if !ok {
		return 0, eigdst
	}
	ed.VectorsTo(dst)
	vals := ed.Values(nil)
	reverse(vals, dst.RawMatrix())
	copy(eigdst, vals)

	for i, v := range vals {
		if v < 0 {
			vals[i] = 0
			continue
		}
		k = i + 1
		vals[i] = math.Sqrt(v)
	}

	var tmp mat.Dense
	tmp.Mul(dst, mat.NewDiagonalRect(n, k, vals[:k]))
	*dst = *dst.Slice(0, n, 0, k).(*mat.Dense)
	dst.Copy(&tmp)

	return k, eigdst
}

func reverse(values []float64, vectors blas64.General) {
	for i, j := 0, len(values)-1; i < j; i, j = i+1, j-1 {
		values[i], values[j] = values[j], values[i]
		blas64.Swap(blas64.Vector{N: vectors.Rows, Inc: vectors.Stride, Data: vectors.Data[i:]},
			blas64.Vector{N: vectors.Rows, Inc: vectors.Stride, Data: vectors.Data[j:]})
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Pseudo Request
//
// Ht allows to make several types of pseudo request: A request which is not a
// HTTP1.1 request but generates output which can be checked via the existing
// checks. The schema of the Test.Request.URL determines the request type.
// Normal HTTP request are made with the two schemas "http" and "https".
// Additionaly the following types of pseudo request are available:
//
// "file" pseudo request
//     This type of pseudo request can be used to read, write and delete a file
//     on the filesystem:
//        - The GET request method tries to read the file given by the URL.Path
//          and returns the content as the response body.
//        - The PUT requets method tries to store the Request.Body under the
//          path given by URL.Path.
//        - The DELETE request method tries to delete the file given by the
//          URL.Path.
//        - The returned HTTP status code is 200 except if any file operation
//          fails in which the Test has status Error.
//
// "bash" pseudo request
//     This type of pseudo request executes a bash script and captures its
//     output as the response.
//        - The script is providen in the Request.Body
//        - The working directory is taken to be URL.Path
//        - Environment is populated from the Request.Params
//        - The Request.Method and the Request.Header are ignored.
//        - The script execution is caneceld after Request.Timout (or the
//          default timeout).
//     The outcome is encoded as follows:
//        - The combined output (stdout and stderr) of the script is returned
//          as the Response.BodyStr.
//        - The HTTP status code is
//             200 if the script's exit code is 0.
//             408 if the script was canceled due to timeout
//             500 if the exit code is != 0.
//        - The Response.Header["Exit-Status"] is used to return the exit
//          status in case of 200 and 500 (sucess and failure).
//
// "sql" pseudo request
//     This type of pseudo request executes a database query (using package
//     database/sql.
//        - The database driver is selected via URL.Host
//        - The data source name is taken from Params["DSN"]
//        - The SQL query is read from the Request.Body
//        - For a GET Request.Method the SQL query is passed to sql.Query,
//          for a POST method the SQL query is passed to sql.Execute.
//     The result if the query is returned in the Response.BodyStr
//
package ht

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

func (t *Test) executeFile() error {
	t.infof("%s %q", t.Request.Request.Method, t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host != "" {
		if u.Host != "localhost" && u.Host != "127.0.0.1" { // TODO IPv6
			return fmt.Errorf("file:// on remote host not implemented")
		}
	}

	switch t.Request.Method {
	case "GET":
		file, err := os.Open(u.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		body, err := ioutil.ReadAll(file)
		t.Response.BodyStr = string(body)
		t.Response.BodyErr = err
	case "DELETE":
		err := os.Remove(u.Path)
		if err != nil {
			return err
		}
		t.Response.BodyStr = fmt.Sprintf("Successfully deleted %s", u)
		t.Response.BodyErr = nil
	case "PUT":
		err := ioutil.WriteFile(u.Path, []byte(t.Request.Body), 0666)
		if err != nil {
			return err
		}
		t.Response.BodyStr = fmt.Sprintf("Successfully wrote %s", u)
		t.Response.BodyErr = nil

	default:
		return fmt.Errorf("method %s not supported on file:// URL", t.Request.Method)
	}

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}

	return nil
}

// executeBash executes a bash script:
func (t *Test) executeBash() error {
	t.infof("Bash script in %q", t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host != "" {
		if u.Host != "localhost" && u.Host != "127.0.0.1" { // TODO IPv6
			return fmt.Errorf("bash:// on remote host not implemented")
		}
	}

	workDir := t.Request.Request.URL.Path

	// Save the request body to a temporary file in the working directory.
	temp, err := ioutil.TempFile(workDir, "bashscript")
	if err != nil {
		return err
	}
	name := temp.Name()
	defer os.Remove(name)
	_, err = temp.WriteString(t.Request.SentBody)
	cerr := temp.Close()
	if err != nil {
		return err
	}
	if cerr != nil {
		return cerr
	}

	ctx, cancel := context.WithTimeout(context.Background(), t.Request.Timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, "/bin/bash", name)
	cmd.Dir = workDir
	for k, v := range t.Request.Params {
		if strings.Contains(k, "=") {
			t.errorf("Environment variable %q from Params contains =; dropped.", k)
			continue
		}
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v[0]))
	}
	output, err := cmd.CombinedOutput()

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}
	t.Response.BodyStr = string(output)

	if ctx.Err() == context.DeadlineExceeded {
		t.Response.Response.StatusCode = http.StatusRequestTimeout
		t.Response.Response.Status = "408 Timeout" // TODO check!
	} else if err != nil {
		t.Response.Response.Status = "500 Internal Server Error"
		t.Response.Response.StatusCode = 500
		t.Response.Response.Header.Set("Exit-Status", err.Error())
	} else {
		t.Response.Response.Header.Set("Exit-Status", "exit status 0")
	}

	return nil
}

// executeSQL executes a SQL query:
func (t *Test) executeSQL() error {
	t.infof("SQL query in %q", t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host == "" {
		return fmt.Errorf("ht: missing database driver name (host of URL) in sql pseudo query")
	}
	dsn := t.Request.Params.Get("DSN")
	if dsn == "" {
		return fmt.Errorf("sql:// missing data source name (DSN parameter) in sql pseudo query")
	}

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return err
	}

	if t.Request.Body == "" {
		return fmt.Errorf("ht: missing query (reqquest body) in sql pseudo query")
	}

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}
	// TODO: set content type to JSON

	switch t.Request.Method {
	case "GET":
		t.Response.BodyStr, err = sqlQuery(db, t.Request.Body)
		if err != nil {
			return err
		}
	case "POST":
		t.Response.BodyStr, err = sqlExecute(db, t.Request.Body)
		if err != nil {
			return err
		}

	default:
		// TODO: this should render the Test into state Bogus, not just Error.
		return fmt.Errorf("ht: illegal method %s for sql pseuo query", t.Request.Method)
	}

	return nil
}

// Returns a json like
//    {
//        "LastInsertId": { "Value": 1234 },
//        "RowsAffected": {
//            "Value": 0,
//            "Error": "something went wrong"
//        }
//    }
func sqlExecute(db *sql.DB, query string) (string, error) {
	result, err := db.Exec(query)
	if err != nil {
		return "", err
	}

	tmp := struct {
		LastInsertId struct {
			Value int64
			Error string `json:",omitempty"`
		}
		RowsAffected struct {
			Value int64
			Error string `json:",omitempty"`
		}
	}{}

	lii, liiErr := result.LastInsertId()
	tmp.LastInsertId.Value = lii
	if liiErr != nil {
		tmp.LastInsertId.Error = liiErr.Error()
	}

	ra, raErr := result.RowsAffected()
	tmp.RowsAffected.Value = ra
	if raErr != nil {
		tmp.RowsAffected.Error = raErr.Error()
	}
	body, err := json.MarshalIndent(tmp, "", "    ")
	if err != nil {
		return "", err
	}

	return string(body), nil
}

func sqlQuery(db *sql.DB, query string) (string, error) {
	rows, err := db.Query(query)
	if err != nil {
		return "", err
	}
	defer rows.Close()
	columns, err := rows.Columns()
	if err != nil {
		return "", err
	}

	buf := bytes.Buffer{}
	buf.WriteString("[\n  ")

	values := make([]string, len(columns))
	ptrs := make([]interface{}, len(columns))
	for i := range values {
		ptrs[i] = &values[i]
	}
	tmp := make(map[string]string)
	sep := ""
	for rows.Next() {
		err = rows.Scan(ptrs...)
		if err != nil {
			return "", err
		}
		for i, col := range columns {
			tmp[col] = values[i]
		}
		record, err := json.MarshalIndent(tmp, "  ", "  ")
		if err != nil {
			return buf.String(), err
		}
		buf.WriteString(sep)
		buf.Write(record)
		sep = ",\n  "
	}
	buf.WriteString("\n]")
	err = rows.Err() // get any error encountered during iteration
	if err != nil {
		return buf.String(), err
	}

	return buf.String(), nil
}
<commit_msg>ht: refactor query in sql pseudo request<commit_after>// Copyright 2016 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Pseudo Request
//
// Ht allows to make several types of pseudo request: A request which is not a
// HTTP1.1 request but generates output which can be checked via the existing
// checks. The schema of the Test.Request.URL determines the request type.
// Normal HTTP request are made with the two schemas "http" and "https".
// Additionaly the following types of pseudo request are available:
//
// "file" pseudo request
//     This type of pseudo request can be used to read, write and delete a file
//     on the filesystem:
//        - The GET request method tries to read the file given by the URL.Path
//          and returns the content as the response body.
//        - The PUT requets method tries to store the Request.Body under the
//          path given by URL.Path.
//        - The DELETE request method tries to delete the file given by the
//          URL.Path.
//        - The returned HTTP status code is 200 except if any file operation
//          fails in which the Test has status Error.
//
// "bash" pseudo request
//     This type of pseudo request executes a bash script and captures its
//     output as the response.
//        - The script is providen in the Request.Body
//        - The working directory is taken to be URL.Path
//        - Environment is populated from the Request.Params
//        - The Request.Method and the Request.Header are ignored.
//        - The script execution is caneceld after Request.Timout (or the
//          default timeout).
//     The outcome is encoded as follows:
//        - The combined output (stdout and stderr) of the script is returned
//          as the Response.BodyStr.
//        - The HTTP status code is
//             200 if the script's exit code is 0.
//             408 if the script was canceled due to timeout
//             500 if the exit code is != 0.
//        - The Response.Header["Exit-Status"] is used to return the exit
//          status in case of 200 and 500 (sucess and failure).
//
// "sql" pseudo request
//     This type of pseudo request executes a database query (using package
//     database/sql.
//        - The database driver is selected via URL.Host
//        - The data source name is taken from Params["DSN"]
//        - The SQL query is read from the Request.Body
//        - For a POST method the SQL query is passed to sql.Execute
//          and the response body is a JSON with LastInsertId and RowsAffected.
//        - For a GET Request.Method the SQL query is passed to sql.Query
//          and the resulting rows are returned as the response body.
//          The format of the response body is determined by the Accept
//          header:
//             "application/json" (the default): a JSON array with the rows
//                 as objects
//             "text/csv": as a csv file
//     The result if the query is returned in the Response.BodyStr
//
package ht

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

func (t *Test) executeFile() error {
	t.infof("%s %q", t.Request.Request.Method, t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host != "" {
		if u.Host != "localhost" && u.Host != "127.0.0.1" { // TODO IPv6
			return fmt.Errorf("file:// on remote host not implemented")
		}
	}

	switch t.Request.Method {
	case "GET":
		file, err := os.Open(u.Path)
		if err != nil {
			return err
		}
		defer file.Close()
		body, err := ioutil.ReadAll(file)
		t.Response.BodyStr = string(body)
		t.Response.BodyErr = err
	case "DELETE":
		err := os.Remove(u.Path)
		if err != nil {
			return err
		}
		t.Response.BodyStr = fmt.Sprintf("Successfully deleted %s", u)
		t.Response.BodyErr = nil
	case "PUT":
		err := ioutil.WriteFile(u.Path, []byte(t.Request.Body), 0666)
		if err != nil {
			return err
		}
		t.Response.BodyStr = fmt.Sprintf("Successfully wrote %s", u)
		t.Response.BodyErr = nil

	default:
		return fmt.Errorf("method %s not supported on file:// URL", t.Request.Method)
	}

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}

	return nil
}

// executeBash executes a bash script:
func (t *Test) executeBash() error {
	t.infof("Bash script in %q", t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host != "" {
		if u.Host != "localhost" && u.Host != "127.0.0.1" { // TODO IPv6
			return fmt.Errorf("bash:// on remote host not implemented")
		}
	}

	workDir := t.Request.Request.URL.Path

	// Save the request body to a temporary file in the working directory.
	temp, err := ioutil.TempFile(workDir, "bashscript")
	if err != nil {
		return err
	}
	name := temp.Name()
	defer os.Remove(name)
	_, err = temp.WriteString(t.Request.SentBody)
	cerr := temp.Close()
	if err != nil {
		return err
	}
	if cerr != nil {
		return cerr
	}

	ctx, cancel := context.WithTimeout(context.Background(), t.Request.Timeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, "/bin/bash", name)
	cmd.Dir = workDir
	for k, v := range t.Request.Params {
		if strings.Contains(k, "=") {
			t.errorf("Environment variable %q from Params contains =; dropped.", k)
			continue
		}
		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v[0]))
	}
	output, err := cmd.CombinedOutput()

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}
	t.Response.BodyStr = string(output)

	if ctx.Err() == context.DeadlineExceeded {
		t.Response.Response.StatusCode = http.StatusRequestTimeout
		t.Response.Response.Status = "408 Timeout" // TODO check!
	} else if err != nil {
		t.Response.Response.Status = "500 Internal Server Error"
		t.Response.Response.StatusCode = 500
		t.Response.Response.Header.Set("Exit-Status", err.Error())
	} else {
		t.Response.Response.Header.Set("Exit-Status", "exit status 0")
	}

	return nil
}

// executeSQL executes a SQL query:
func (t *Test) executeSQL() error {
	t.infof("SQL query in %q", t.Request.Request.URL.String())

	start := time.Now()
	defer func() {
		t.Response.Duration = time.Since(start)
	}()

	u := t.Request.Request.URL
	if u.Host == "" {
		return fmt.Errorf("ht: missing database driver name (host of URL) in sql pseudo query")
	}
	dsn := t.Request.Params.Get("DSN")
	if dsn == "" {
		return fmt.Errorf("sql:// missing data source name (DSN parameter) in sql pseudo query")
	}

	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return err
	}

	if t.Request.Body == "" {
		return fmt.Errorf("ht: missing query (reqquest body) in sql pseudo query")
	}

	// Fake a http.Response
	t.Response.Response = &http.Response{
		Status:     "200 OK",
		StatusCode: 200,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Header:     make(http.Header),
		Body:       nil, // already close and consumed
		Trailer:    make(http.Header),
		Request:    t.Request.Request,
	}
	// TODO: set content type to JSON

	switch t.Request.Method {
	case "GET":
		t.Response.BodyStr, err = sqlQuery(db, t.Request.Body)
		if err != nil {
			return err
		}
	case "POST":
		t.Response.BodyStr, err = sqlExecute(db, t.Request.Body)
		if err != nil {
			return err
		}

	default:
		// TODO: this should render the Test into state Bogus, not just Error.
		return fmt.Errorf("ht: illegal method %s for sql pseuo query", t.Request.Method)
	}

	return nil
}

// Returns a json like
//    {
//        "LastInsertId": { "Value": 1234 },
//        "RowsAffected": {
//            "Value": 0,
//            "Error": "something went wrong"
//        }
//    }
func sqlExecute(db *sql.DB, query string) (string, error) {
	result, err := db.Exec(query)
	if err != nil {
		return "", err
	}

	tmp := struct {
		LastInsertId struct {
			Value int64
			Error string `json:",omitempty"`
		}
		RowsAffected struct {
			Value int64
			Error string `json:",omitempty"`
		}
	}{}

	lii, liiErr := result.LastInsertId()
	tmp.LastInsertId.Value = lii
	if liiErr != nil {
		tmp.LastInsertId.Error = liiErr.Error()
	}

	ra, raErr := result.RowsAffected()
	tmp.RowsAffected.Value = ra
	if raErr != nil {
		tmp.RowsAffected.Error = raErr.Error()
	}
	body, err := json.MarshalIndent(tmp, "", "    ")
	if err != nil {
		return "", err
	}

	return string(body), nil
}

type recordWriter interface {
	WriteRecord([]string) error
	Close() (string, error)
}

type jsonRecorder struct {
	cols  []string
	buf   *bytes.Buffer
	first bool
	tmp   map[string]string
	err   error
}

func newJsonRecorder(buf *bytes.Buffer, cols []string) *jsonRecorder {
	buf.WriteString("[\n  ")
	return &jsonRecorder{
		cols:  cols,
		buf:   buf,
		first: true,
		tmp:   make(map[string]string),
	}
}

func (jr *jsonRecorder) Close() (string, error) {
	_, err := jr.buf.WriteString("\n]")
	if err != nil {
		jr.err = err
	}
	return jr.buf.String(), jr.err
}
func (jr *jsonRecorder) WriteRecord(values []string) {
	if jr.err != nil {
		return
	}
	for i, col := range jr.cols {
		jr.tmp[col] = values[i]
	}
	record, err := json.MarshalIndent(jr.tmp, "  ", "  ")
	if err != nil {
		jr.err = err
		return
	}
	if jr.first {
		jr.first = false
	} else {
		_, err = jr.buf.WriteString(",\n  ")
		if err != nil {
			jr.err = err
			return
		}
	}
	_, err = jr.buf.Write(record)
	if err != nil {
		jr.err = err
	}
}

func sqlQuery(db *sql.DB, query string) (string, error) {
	rows, err := db.Query(query)
	if err != nil {
		return "", err
	}
	defer rows.Close()
	columns, err := rows.Columns()
	if err != nil {
		return "", err
	}

	buf := &bytes.Buffer{}
	recorder := newJsonRecorder(buf, columns)

	values := make([]string, len(columns))
	ptrs := make([]interface{}, len(columns))
	for i := range values {
		ptrs[i] = &values[i]
	}
	for rows.Next() {
		err = rows.Scan(ptrs...)
		if err != nil {
			return "", err
		}
		recorder.WriteRecord(values)
	}
	err = rows.Err() // get any error encountered during iteration
	body, _ := recorder.Close()
	if err != nil {
		return body, err
	}

	return body, nil
}
<|endoftext|>


<commit_before>package tftp

import (
	"net"
	"fmt"
	"io"
	"time"
	"log"
)

type receiver struct {
	remoteAddr *net.UDPAddr
	conn *net.UDPConn
	writer *io.PipeWriter
	Log *log.Logger
}

func (r *receiver) Run() {
	var blockNumber uint16
	blockNumber = 1
	var buffer []byte
	buffer = make([]byte, 50000)

	for {
		last, e := r.receiveBlock(buffer, blockNumber)
		if e != nil {
			if r.Log != nil {
				r.Log.Printf("Error receiving block %d: %v", blockNumber, e)
			}
			r.writer.CloseWithError(e)
			return
		}
		if last {
			break
		}
		blockNumber++
	}
	r.writer.Close()
	r.terminate(buffer, blockNumber, false)
	return
}

func (r *receiver) receiveBlock(b []byte, n uint16) (last bool, e error) {
	for i := 0; i < 3; i++ {
		ackPacket := ACK{n - 1}
		r.conn.WriteToUDP(ackPacket.Pack(), r.remoteAddr)
		setDeadlineError := r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		if setDeadlineError != nil {
			return false, fmt.Errorf("Could not set UDP timeout: %v", setDeadlineError)
		}
		for {
			c, _, readError := r.conn.ReadFromUDP(b)
			if networkError, ok := readError.(net.Error); ok && networkError.Timeout() {
				break
			} else if e != nil {
				return false, fmt.Errorf("Error reading UDP packet: %v", e)
			}
			packet, e := ParsePacket(b[:c])
			if e != nil {
				continue
			}
			switch p := Packet(*packet).(type) {
				case *DATA:
					if n == p.BlockNumber {
						_, e := r.writer.Write(p.Data)
						if e == nil {
							return len(p.Data) < 512, nil
						} else {
							errorPacket := ERROR{1, e.Error()}
							r.conn.WriteToUDP(errorPacket.Pack(), r.remoteAddr)
							return false, fmt.Errorf("Handler error: %v", e)
						}
					}
				case *ERROR:
					return false, fmt.Errorf("Transmission error %d: %s", p.ErrorCode, p.ErrorMessage)
			}
		}
	}
	return false, fmt.Errorf("Receive timeout")
}

func (r *receiver) terminate(b []byte, n uint16, dallying bool) (e error) {
	for i := 0; i < 3; i++ {
		ackPacket := ACK{n}
		_, e := r.conn.WriteToUDP(ackPacket.Pack(), r.remoteAddr)
		if !dallying {
			return e
		}
		setDeadlineError := r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		if setDeadlineError != nil {
			return fmt.Errorf("Could not set UDP timeout: %v", setDeadlineError)
		}
		l1:
		for {
			c, _, readError := r.conn.ReadFromUDP(b)
			if networkError, ok := readError.(net.Error); ok && networkError.Timeout() {
				return nil
			} else if e != nil {
				return fmt.Errorf("Error reading UDP packet: %v", e)
			}
			packet, e := ParsePacket(b[:c])
			if e != nil {
				continue
			}
			switch p := Packet(*packet).(type) {
				case *DATA:
					if n == p.BlockNumber {
						break l1
					}
				case *ERROR:
					fmt.Errorf("Transmission error %d: %s", p.ErrorCode, p.ErrorMessage)
			}
		}
	}
	return fmt.Errorf("Termination error")
}<commit_msg>support for client-side operation<commit_after>package tftp

import (
	"net"
	"fmt"
	"io"
	"time"
	"log"
)

type receiver struct {
	remoteAddr *net.UDPAddr
	conn *net.UDPConn
	writer *io.PipeWriter
	Log *log.Logger
}

func (r *receiver) Run(isServerMode bool) (error) {
	var blockNumber uint16
	blockNumber = 1
	var buffer []byte
	buffer = make([]byte, 50000)
	firstBlock := true
	for {
		last, e := r.receiveBlock(buffer, blockNumber, firstBlock && !isServerMode)
		if e != nil {
			if r.Log != nil {
				r.Log.Printf("Error receiving block %d: %v", blockNumber, e)
			}
			r.writer.CloseWithError(e)
			if firstBlock {
				return e
			} else {
				return nil
			}
		}
		firstBlock = false
		if last {
			break
		}
		blockNumber++
	}
	r.writer.Close()
	r.terminate(buffer, blockNumber, false)
	return nil
}

func (r *receiver) receiveBlock(b []byte, n uint16, firstBlockOnClient bool) (last bool, e error) {
	for i := 0; i < 3; i++ {
		if !firstBlockOnClient {
			ackPacket := ACK{n - 1}
			r.conn.WriteToUDP(ackPacket.Pack(), r.remoteAddr)
		}
		setDeadlineError := r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		if setDeadlineError != nil {
			return false, fmt.Errorf("Could not set UDP timeout: %v", setDeadlineError)
		}
		for {
			c, remoteAddr, readError := r.conn.ReadFromUDP(b)
			if firstBlockOnClient {
				r.remoteAddr = remoteAddr
			}
			if networkError, ok := readError.(net.Error); ok && networkError.Timeout() {
				if firstBlockOnClient {
					return false, fmt.Errorf("Receive timeout")
				} else {
					break
				}
			} else if e != nil {
				return false, fmt.Errorf("Error reading UDP packet: %v", e)
			}
			packet, e := ParsePacket(b[:c])
			if e != nil {
				continue
			}
			switch p := Packet(*packet).(type) {
				case *DATA:
					if n == p.BlockNumber {
						_, e := r.writer.Write(p.Data)
						if e == nil {
							return len(p.Data) < 512, nil
						} else {
							errorPacket := ERROR{1, e.Error()}
							r.conn.WriteToUDP(errorPacket.Pack(), r.remoteAddr)
							return false, fmt.Errorf("Handler error: %v", e)
						}
					}
				case *ERROR:
					return false, fmt.Errorf("Transmission error %d: %s", p.ErrorCode, p.ErrorMessage)
			}
		}
	}
	return false, fmt.Errorf("Receive timeout")
}

func (r *receiver) terminate(b []byte, n uint16, dallying bool) (e error) {
	for i := 0; i < 3; i++ {
		ackPacket := ACK{n}
		_, e := r.conn.WriteToUDP(ackPacket.Pack(), r.remoteAddr)
		if !dallying {
			return e
		}
		setDeadlineError := r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
		if setDeadlineError != nil {
			return fmt.Errorf("Could not set UDP timeout: %v", setDeadlineError)
		}
		l1:
		for {
			c, _, readError := r.conn.ReadFromUDP(b)
			if networkError, ok := readError.(net.Error); ok && networkError.Timeout() {
				return nil
			} else if e != nil {
				return fmt.Errorf("Error reading UDP packet: %v", e)
			}
			packet, e := ParsePacket(b[:c])
			if e != nil {
				continue
			}
			switch p := Packet(*packet).(type) {
				case *DATA:
					if n == p.BlockNumber {
						break l1
					}
				case *ERROR:
					fmt.Errorf("Transmission error %d: %s", p.ErrorCode, p.ErrorMessage)
			}
		}
	}
	return fmt.Errorf("Termination error")
}<|endoftext|>


<commit_before>package api

const (
	settings_yml = `denvhome: ~/.denv
infofile: .denvinfo
ignorefile: .denvignore
`
	_default_denvignore = `.bash_history
.git
.denv*
.ssh
.viminfo
`
)
<commit_msg>debugging travis error<commit_after>package api

const (
settings_yml = `denvhome: ~/.denv
infofile: .denvinfo
ignorefile: .denvignore
`
_default_denvignore = `.bash_history
.gimme
.cache
.gem
.bundle
.erlang.cookie
.git
.denv*
.ssh
.viminfo
`
)
<|endoftext|>


<commit_before><commit_msg>Fixed deadlock occuring when sending a generated field<commit_after><|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io/ioutil"
	"strings"
)

// Command is a subcommand for a cli.App.
type Command struct {
	// The name of the command
	Name string
	// short name of the command. Typically one character (deprecated, use `Aliases`)
	ShortName string
	// A list of aliases for the command
	Aliases []string
	// A short description of the usage of this command
	Usage string
	// A longer explanation of how the command works
	Description string
	// Example usage
	Examples string
	// The function to call when checking for bash command completions
	BashComplete func(context *Context)
	// An action to execute before any sub-subcommands are run, but after the context is ready
	// If a non-nil error is returned, no sub-subcommands are run
	Before func(context *Context) error
	// An action to execute after any subcommands are run, but after the subcommand has finished
	// It is run even if Action() panics
	After func(context *Context) error
	// The function to call when this command is invoked
	Action func(context *Context)
	// List of child commands
	Subcommands []Command
	// List of flags to parse
	Flags []Flag
	// Treat all flags as normal arguments if true
	SkipFlagParsing bool
	// Boolean to hide built-in help command
	HideHelp bool
	// for non-flag argument help in USAGE section
	// For `app command <path1> <path2>`, ArgumentHelp would be "<path1> <path2>"
	ArgumentHelp string

	commandNamePath []string
}

// Returns the full name of the command.
// For subcommands this ensures that parent commands are part of the command path
func (c Command) FullName() string {
	if c.commandNamePath == nil {
		return c.Name
	}
	return strings.Join(c.commandNamePath, " ")
}

// Invokes the command given the context, parses ctx.Args() to generate command-specific flags
func (c Command) Run(ctx *Context) error {
	if len(c.Subcommands) > 0 || c.Before != nil || c.After != nil {
		return c.startApp(ctx)
	}

	if !c.HideHelp && (HelpFlag != BoolFlag{}) {
		// append help to flags
		c.Flags = append(
			c.Flags,
			HelpFlag,
		)
	}

	if ctx.App.EnableBashCompletion {
		c.Flags = append(c.Flags, BashCompletionFlag)
	}

	set := flagSet(c.Name, c.Flags)
	set.SetOutput(ioutil.Discard)

	firstFlagIndex := -1
	terminatorIndex := -1
	for index, arg := range ctx.Args() {
		if arg == "--" {
			terminatorIndex = index
			break
		} else if strings.HasPrefix(arg, "-") && firstFlagIndex == -1 {
			firstFlagIndex = index
		}
	}

	var err error
	if firstFlagIndex > -1 && !c.SkipFlagParsing {
		args := ctx.Args()
		regularArgs := make([]string, len(args[1:firstFlagIndex]))
		copy(regularArgs, args[1:firstFlagIndex])

		var flagArgs []string
		if terminatorIndex > -1 {
			flagArgs = args[firstFlagIndex:terminatorIndex]
			regularArgs = append(regularArgs, args[terminatorIndex:]...)
		} else {
			flagArgs = args[firstFlagIndex:]
		}

		err = set.Parse(append(flagArgs, regularArgs...))
	} else {
		err = set.Parse(ctx.Args().Tail())
	}

	if err != nil {
		fmt.Fprintln(ctx.App.Writer, "Incorrect Usage.")
		fmt.Fprintln(ctx.App.Writer)
		ShowCommandHelp(ctx, c.Name)
		return err
	}

	nerr := normalizeFlags(c.Flags, set)
	if nerr != nil {
		fmt.Fprintln(ctx.App.Writer, nerr)
		fmt.Fprintln(ctx.App.Writer)
		ShowCommandHelp(ctx, c.Name)
		return nerr
	}
	context := NewContext(ctx.App, set, ctx)

	if checkCommandCompletions(context, c.Name) {
		return nil
	}

	if checkCommandHelp(context, c.Name) {
		return nil
	}
	context.Command = c
	c.Action(context)
	return nil
}

func (c Command) Names() []string {
	names := []string{c.Name}

	if c.ShortName != "" {
		names = append(names, c.ShortName)
	}

	return append(names, c.Aliases...)
}

// Returns true if Command.Name or Command.ShortName matches given name
func (c Command) HasName(name string) bool {
	for _, n := range c.Names() {
		if n == name {
			return true
		}
	}
	return false
}

// Strips and indents with 3 spaces examples.
func (c Command) ExamplesFormatted() string {
	prefix := "   "
	return prefix + strings.Join(strings.Split(strings.TrimSpace(c.Examples), "\n"), "\n"+prefix)
}

func (c Command) startApp(ctx *Context) error {
	app := NewApp()

	// set the name and usage
	app.Name = fmt.Sprintf("%s %s", ctx.App.Name, c.Name)
	if c.Description != "" {
		app.Usage = c.Description
	} else {
		app.Usage = c.Usage
	}

	// set CommandNotFound
	app.CommandNotFound = ctx.App.CommandNotFound

	// set the flags and commands
	app.Commands = c.Subcommands
	app.Flags = c.Flags
	app.HideHelp = c.HideHelp

	app.Version = ctx.App.Version
	app.HideVersion = ctx.App.HideVersion
	app.Compiled = ctx.App.Compiled
	app.Author = ctx.App.Author
	app.Email = ctx.App.Email
	app.Writer = ctx.App.Writer

	// bash completion
	app.EnableBashCompletion = ctx.App.EnableBashCompletion
	if c.BashComplete != nil {
		app.BashComplete = c.BashComplete
	}

	// set the actions
	app.Before = c.Before
	app.After = c.After
	if c.Action != nil {
		app.Action = c.Action
	} else {
		app.Action = helpSubcommand.Action
	}

	var newCmds []Command
	for _, cc := range app.Commands {
		cc.commandNamePath = []string{c.Name, cc.Name}
		newCmds = append(newCmds, cc)
	}
	app.Commands = newCmds

	return app.RunAsSubcommand(ctx)
}

type ByName []Command

func (c ByName) Len() int {
	return len(c)
}

func (c ByName) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

func (c ByName) Less(i, j int) bool {
	return c[i].Name < c[j].Name
}
<commit_msg>command: allow deeply-nested subcommands to have right full name<commit_after>package cli

import (
	"fmt"
	"io/ioutil"
	"strings"
)

// Command is a subcommand for a cli.App.
type Command struct {
	// The name of the command
	Name string
	// short name of the command. Typically one character (deprecated, use `Aliases`)
	ShortName string
	// A list of aliases for the command
	Aliases []string
	// A short description of the usage of this command
	Usage string
	// A longer explanation of how the command works
	Description string
	// Example usage
	Examples string
	// The function to call when checking for bash command completions
	BashComplete func(context *Context)
	// An action to execute before any sub-subcommands are run, but after the context is ready
	// If a non-nil error is returned, no sub-subcommands are run
	Before func(context *Context) error
	// An action to execute after any subcommands are run, but after the subcommand has finished
	// It is run even if Action() panics
	After func(context *Context) error
	// The function to call when this command is invoked
	Action func(context *Context)
	// List of child commands
	Subcommands []Command
	// List of flags to parse
	Flags []Flag
	// Treat all flags as normal arguments if true
	SkipFlagParsing bool
	// Boolean to hide built-in help command
	HideHelp bool
	// for non-flag argument help in USAGE section
	// For `app command <path1> <path2>`, ArgumentHelp would be "<path1> <path2>"
	ArgumentHelp string

	commandNamePath []string
}

// Returns the full name of the command.
// For subcommands this ensures that parent commands are part of the command path
func (c Command) FullName() string {
	if c.commandNamePath == nil {
		return c.Name
	}
	return strings.Join(c.commandNamePath, " ")
}

// Invokes the command given the context, parses ctx.Args() to generate command-specific flags
func (c Command) Run(ctx *Context) error {
	if len(c.Subcommands) > 0 || c.Before != nil || c.After != nil {
		return c.startApp(ctx)
	}

	if !c.HideHelp && (HelpFlag != BoolFlag{}) {
		// append help to flags
		c.Flags = append(
			c.Flags,
			HelpFlag,
		)
	}

	if ctx.App.EnableBashCompletion {
		c.Flags = append(c.Flags, BashCompletionFlag)
	}

	set := flagSet(c.Name, c.Flags)
	set.SetOutput(ioutil.Discard)

	firstFlagIndex := -1
	terminatorIndex := -1
	for index, arg := range ctx.Args() {
		if arg == "--" {
			terminatorIndex = index
			break
		} else if strings.HasPrefix(arg, "-") && firstFlagIndex == -1 {
			firstFlagIndex = index
		}
	}

	var err error
	if firstFlagIndex > -1 && !c.SkipFlagParsing {
		args := ctx.Args()
		regularArgs := make([]string, len(args[1:firstFlagIndex]))
		copy(regularArgs, args[1:firstFlagIndex])

		var flagArgs []string
		if terminatorIndex > -1 {
			flagArgs = args[firstFlagIndex:terminatorIndex]
			regularArgs = append(regularArgs, args[terminatorIndex:]...)
		} else {
			flagArgs = args[firstFlagIndex:]
		}

		err = set.Parse(append(flagArgs, regularArgs...))
	} else {
		err = set.Parse(ctx.Args().Tail())
	}

	if err != nil {
		fmt.Fprintln(ctx.App.Writer, "Incorrect Usage.")
		fmt.Fprintln(ctx.App.Writer)
		ShowCommandHelp(ctx, c.Name)
		return err
	}

	nerr := normalizeFlags(c.Flags, set)
	if nerr != nil {
		fmt.Fprintln(ctx.App.Writer, nerr)
		fmt.Fprintln(ctx.App.Writer)
		ShowCommandHelp(ctx, c.Name)
		return nerr
	}
	context := NewContext(ctx.App, set, ctx)

	if checkCommandCompletions(context, c.Name) {
		return nil
	}

	if checkCommandHelp(context, c.Name) {
		return nil
	}
	context.Command = c
	c.Action(context)
	return nil
}

func (c Command) Names() []string {
	names := []string{c.Name}

	if c.ShortName != "" {
		names = append(names, c.ShortName)
	}

	return append(names, c.Aliases...)
}

// Returns true if Command.Name or Command.ShortName matches given name
func (c Command) HasName(name string) bool {
	for _, n := range c.Names() {
		if n == name {
			return true
		}
	}
	return false
}

// Strips and indents with 3 spaces examples.
func (c Command) ExamplesFormatted() string {
	prefix := "   "
	return prefix + strings.Join(strings.Split(strings.TrimSpace(c.Examples), "\n"), "\n"+prefix)
}

func (c Command) startApp(ctx *Context) error {
	app := NewApp()

	// set the name and usage
	app.Name = fmt.Sprintf("%s %s", ctx.App.Name, c.Name)
	if c.Description != "" {
		app.Usage = c.Description
	} else {
		app.Usage = c.Usage
	}

	// set CommandNotFound
	app.CommandNotFound = ctx.App.CommandNotFound

	// set the flags and commands
	app.Commands = c.Subcommands
	app.Flags = c.Flags
	app.HideHelp = c.HideHelp

	app.Version = ctx.App.Version
	app.HideVersion = ctx.App.HideVersion
	app.Compiled = ctx.App.Compiled
	app.Author = ctx.App.Author
	app.Email = ctx.App.Email
	app.Writer = ctx.App.Writer

	// bash completion
	app.EnableBashCompletion = ctx.App.EnableBashCompletion
	if c.BashComplete != nil {
		app.BashComplete = c.BashComplete
	}

	// set the actions
	app.Before = c.Before
	app.After = c.After
	if c.Action != nil {
		app.Action = c.Action
	} else {
		app.Action = helpSubcommand.Action
	}

	var newCmds []Command
	for _, cc := range app.Commands {
		cc.commandNamePath = []string{c.FullName(), cc.Name}
		newCmds = append(newCmds, cc)
	}
	app.Commands = newCmds

	return app.RunAsSubcommand(ctx)
}

type ByName []Command

func (c ByName) Len() int {
	return len(c)
}

func (c ByName) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

func (c ByName) Less(i, j int) bool {
	return c[i].Name < c[j].Name
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach Authors.
//
// 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package sql_test

import (
	"reflect"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/protoutil"
)

var configID = sqlbase.ID(1)
var configDescKey = sqlbase.MakeDescMetadataKey(keys.MaxReservedDescID)

// forceNewConfig forces a system config update by writing a bogus descriptor with an
// incremented value inside. It then repeatedly fetches the gossip config until the
// just-written descriptor is found.
func forceNewConfig(t *testing.T, s *testServer) config.SystemConfig {
	configID++
	configDesc := &sqlbase.Descriptor{
		Union: &sqlbase.Descriptor_Database{
			Database: &sqlbase.DatabaseDescriptor{
				Name:       "sentinel",
				ID:         configID,
				Privileges: &sqlbase.PrivilegeDescriptor{},
			},
		},
	}

	// This needs to be done in a transaction with the system trigger set.
	if err := s.DB().Txn(func(txn *client.Txn) error {
		txn.SetSystemConfigTrigger()
		return txn.Put(configDescKey, configDesc)
	}); err != nil {
		t.Fatal(err)
	}
	return waitForConfigChange(t, s)
}

func waitForConfigChange(t *testing.T, s *testServer) config.SystemConfig {
	var foundDesc sqlbase.Descriptor
	var cfg config.SystemConfig
	util.SucceedsSoon(t, func() error {
		var ok bool
		if cfg, ok = s.Gossip().GetSystemConfig(); ok {
			if val := cfg.GetValue(configDescKey); val != nil {
				if err := val.GetProto(&foundDesc); err != nil {
					t.Fatal(err)
				}
				if id := foundDesc.GetDatabase().GetID(); id != configID {
					return util.Errorf("expected database id %d; got %d", configID, id)
				}
				return nil
			}
		}
		return util.Errorf("got nil system config")
	})
	return cfg
}

// TestGetZoneConfig exercises config.GetZoneConfig and the sql hook for it.
func TestGetZoneConfig(t *testing.T) {
	defer leaktest.AfterTest(t)()
	// Disable splitting. We're using bad attributes in zone configs
	// to be able to match.
	defer config.TestingDisableTableSplits()()
	s, sqlDB, _ := setup(t)
	defer cleanup(s, sqlDB)

	expectedCounter := uint32(keys.MaxReservedDescID + 1)

	defaultZoneConfig := config.DefaultZoneConfig()
	defaultZoneConfig.RangeMinBytes = 1 << 20
	defaultZoneConfig.RangeMaxBytes = 1 << 20
	defaultZoneConfig.GC.TTLSeconds = 60

	{
		buf, err := protoutil.Marshal(&defaultZoneConfig)
		if err != nil {
			t.Fatal(err)
		}
		objID := keys.RootNamespaceID
		if _, err = sqlDB.Exec(`UPDATE system.zones SET config = $2 WHERE id = $1`, objID, buf); err != nil {
			t.Fatalf("problem writing zone %+v: %s", defaultZoneConfig, err)
		}
	}

	// Naming scheme for database and tables:
	// db1 has tables tb11 and tb12
	// db2 has tables tb21 and tb22

	db1 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE DATABASE db1`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	db2 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE DATABASE db2`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb11 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db1.tb1 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb12 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db1.tb2 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb21 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db2.tb1 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb22 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db2.tb2 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	cfg := forceNewConfig(t, s)

	// We have no custom zone configs.
	testCases := []struct {
		key     roachpb.RKey
		zoneCfg *config.ZoneConfig
	}{
		{roachpb.RKeyMin, &defaultZoneConfig},
		{keys.MakeTablePrefix(0), &defaultZoneConfig},
		{keys.MakeTablePrefix(1), &defaultZoneConfig},
		{keys.MakeTablePrefix(keys.MaxReservedDescID), &defaultZoneConfig},
		{keys.MakeTablePrefix(db1), &defaultZoneConfig},
		{keys.MakeTablePrefix(db2), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb11), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb12), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb21), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb22), &defaultZoneConfig},
	}

	for tcNum, tc := range testCases {
		zoneCfg, err := cfg.GetZoneConfigForKey(tc.key)
		if err != nil {
			t.Fatalf("#%d: err=%s", tcNum, err)
		}

		if !reflect.DeepEqual(zoneCfg, tc.zoneCfg) {
			t.Errorf("#%d: bad zone config.\nexpected: %+v\ngot: %+v", tcNum, tc.zoneCfg, zoneCfg)
		}
	}

	// Now set some zone configs. We don't have a nice way of using table
	// names for this, so we do raw puts.
	// Here is the list of dbs/tables and whether they have a custom zone config:
	// db1: true
	//   tb1: true
	//   tb2: false
	// db1: false
	//   tb1: true
	//   tb2: false
	db1Cfg := config.ZoneConfig{ReplicaAttrs: []roachpb.Attributes{{[]string{"db1"}}}}
	tb11Cfg := config.ZoneConfig{ReplicaAttrs: []roachpb.Attributes{{[]string{"db1.tb1"}}}}
	tb21Cfg := config.ZoneConfig{ReplicaAttrs: []roachpb.Attributes{{[]string{"db2.tb1"}}}}
	for objID, objZone := range map[uint32]config.ZoneConfig{
		db1:  db1Cfg,
		tb11: tb11Cfg,
		tb21: tb21Cfg,
	} {
		buf, err := protoutil.Marshal(&objZone)
		if err != nil {
			t.Fatal(err)
		}
		if _, err = sqlDB.Exec(`INSERT INTO system.zones VALUES ($1, $2)`, objID, buf); err != nil {
			t.Fatalf("problem writing zone %+v: %s", objZone, err)
		}
	}

	cfg = forceNewConfig(t, s)

	testCases = []struct {
		key     roachpb.RKey
		zoneCfg *config.ZoneConfig
	}{
		{roachpb.RKeyMin, &defaultZoneConfig},
		{keys.MakeTablePrefix(0), &defaultZoneConfig},
		{keys.MakeTablePrefix(1), &defaultZoneConfig},
		{keys.MakeTablePrefix(keys.MaxReservedDescID), &defaultZoneConfig},
		{keys.MakeTablePrefix(db1), &db1Cfg},
		{keys.MakeTablePrefix(db2), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb11), &tb11Cfg},
		{keys.MakeTablePrefix(tb12), &db1Cfg},
		{keys.MakeTablePrefix(tb21), &tb21Cfg},
		{keys.MakeTablePrefix(tb22), &defaultZoneConfig},
	}

	for tcNum, tc := range testCases {
		zoneCfg, err := cfg.GetZoneConfigForKey(tc.key)
		if err != nil {
			t.Fatalf("#%d: err=%s", tcNum, err)
		}

		if !reflect.DeepEqual(zoneCfg, tc.zoneCfg) {
			t.Errorf("#%d: bad zone config.\nexpected: %+v\ngot: %+v", tcNum, tc.zoneCfg, zoneCfg)
		}
	}
}
<commit_msg>sql: avoid go1.7 vet warning about using an unkeyed field in a composite literal<commit_after>// Copyright 2015 The Cockroach Authors.
//
// 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.
//
// Author: Marc Berhault (marc@cockroachlabs.com)

package sql_test

import (
	"reflect"
	"testing"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/keys"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/sql/sqlbase"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	"github.com/cockroachdb/cockroach/util/protoutil"
)

var configID = sqlbase.ID(1)
var configDescKey = sqlbase.MakeDescMetadataKey(keys.MaxReservedDescID)

// forceNewConfig forces a system config update by writing a bogus descriptor with an
// incremented value inside. It then repeatedly fetches the gossip config until the
// just-written descriptor is found.
func forceNewConfig(t *testing.T, s *testServer) config.SystemConfig {
	configID++
	configDesc := &sqlbase.Descriptor{
		Union: &sqlbase.Descriptor_Database{
			Database: &sqlbase.DatabaseDescriptor{
				Name:       "sentinel",
				ID:         configID,
				Privileges: &sqlbase.PrivilegeDescriptor{},
			},
		},
	}

	// This needs to be done in a transaction with the system trigger set.
	if err := s.DB().Txn(func(txn *client.Txn) error {
		txn.SetSystemConfigTrigger()
		return txn.Put(configDescKey, configDesc)
	}); err != nil {
		t.Fatal(err)
	}
	return waitForConfigChange(t, s)
}

func waitForConfigChange(t *testing.T, s *testServer) config.SystemConfig {
	var foundDesc sqlbase.Descriptor
	var cfg config.SystemConfig
	util.SucceedsSoon(t, func() error {
		var ok bool
		if cfg, ok = s.Gossip().GetSystemConfig(); ok {
			if val := cfg.GetValue(configDescKey); val != nil {
				if err := val.GetProto(&foundDesc); err != nil {
					t.Fatal(err)
				}
				if id := foundDesc.GetDatabase().GetID(); id != configID {
					return util.Errorf("expected database id %d; got %d", configID, id)
				}
				return nil
			}
		}
		return util.Errorf("got nil system config")
	})
	return cfg
}

// TestGetZoneConfig exercises config.GetZoneConfig and the sql hook for it.
func TestGetZoneConfig(t *testing.T) {
	defer leaktest.AfterTest(t)()
	// Disable splitting. We're using bad attributes in zone configs
	// to be able to match.
	defer config.TestingDisableTableSplits()()
	s, sqlDB, _ := setup(t)
	defer cleanup(s, sqlDB)

	expectedCounter := uint32(keys.MaxReservedDescID + 1)

	defaultZoneConfig := config.DefaultZoneConfig()
	defaultZoneConfig.RangeMinBytes = 1 << 20
	defaultZoneConfig.RangeMaxBytes = 1 << 20
	defaultZoneConfig.GC.TTLSeconds = 60

	{
		buf, err := protoutil.Marshal(&defaultZoneConfig)
		if err != nil {
			t.Fatal(err)
		}
		objID := keys.RootNamespaceID
		if _, err = sqlDB.Exec(`UPDATE system.zones SET config = $2 WHERE id = $1`, objID, buf); err != nil {
			t.Fatalf("problem writing zone %+v: %s", defaultZoneConfig, err)
		}
	}

	// Naming scheme for database and tables:
	// db1 has tables tb11 and tb12
	// db2 has tables tb21 and tb22

	db1 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE DATABASE db1`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	db2 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE DATABASE db2`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb11 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db1.tb1 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb12 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db1.tb2 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb21 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db2.tb1 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	tb22 := expectedCounter
	if _, err := sqlDB.Exec(`CREATE TABLE db2.tb2 (k INT PRIMARY KEY, v INT)`); err != nil {
		t.Fatal(err)
	}
	expectedCounter++

	cfg := forceNewConfig(t, s)

	// We have no custom zone configs.
	testCases := []struct {
		key     roachpb.RKey
		zoneCfg *config.ZoneConfig
	}{
		{roachpb.RKeyMin, &defaultZoneConfig},
		{keys.MakeTablePrefix(0), &defaultZoneConfig},
		{keys.MakeTablePrefix(1), &defaultZoneConfig},
		{keys.MakeTablePrefix(keys.MaxReservedDescID), &defaultZoneConfig},
		{keys.MakeTablePrefix(db1), &defaultZoneConfig},
		{keys.MakeTablePrefix(db2), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb11), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb12), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb21), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb22), &defaultZoneConfig},
	}

	for tcNum, tc := range testCases {
		zoneCfg, err := cfg.GetZoneConfigForKey(tc.key)
		if err != nil {
			t.Fatalf("#%d: err=%s", tcNum, err)
		}

		if !reflect.DeepEqual(zoneCfg, tc.zoneCfg) {
			t.Errorf("#%d: bad zone config.\nexpected: %+v\ngot: %+v", tcNum, tc.zoneCfg, zoneCfg)
		}
	}

	// Now set some zone configs. We don't have a nice way of using table
	// names for this, so we do raw puts.
	// Here is the list of dbs/tables and whether they have a custom zone config:
	// db1: true
	//   tb1: true
	//   tb2: false
	// db1: false
	//   tb1: true
	//   tb2: false
	db1Cfg := config.ZoneConfig{
		ReplicaAttrs: []roachpb.Attributes{{Attrs: []string{"db1"}}},
	}
	tb11Cfg := config.ZoneConfig{
		ReplicaAttrs: []roachpb.Attributes{{Attrs: []string{"db1.tb1"}}},
	}
	tb21Cfg := config.ZoneConfig{
		ReplicaAttrs: []roachpb.Attributes{{Attrs: []string{"db2.tb1"}}},
	}
	for objID, objZone := range map[uint32]config.ZoneConfig{
		db1:  db1Cfg,
		tb11: tb11Cfg,
		tb21: tb21Cfg,
	} {
		buf, err := protoutil.Marshal(&objZone)
		if err != nil {
			t.Fatal(err)
		}
		if _, err = sqlDB.Exec(`INSERT INTO system.zones VALUES ($1, $2)`, objID, buf); err != nil {
			t.Fatalf("problem writing zone %+v: %s", objZone, err)
		}
	}

	cfg = forceNewConfig(t, s)

	testCases = []struct {
		key     roachpb.RKey
		zoneCfg *config.ZoneConfig
	}{
		{roachpb.RKeyMin, &defaultZoneConfig},
		{keys.MakeTablePrefix(0), &defaultZoneConfig},
		{keys.MakeTablePrefix(1), &defaultZoneConfig},
		{keys.MakeTablePrefix(keys.MaxReservedDescID), &defaultZoneConfig},
		{keys.MakeTablePrefix(db1), &db1Cfg},
		{keys.MakeTablePrefix(db2), &defaultZoneConfig},
		{keys.MakeTablePrefix(tb11), &tb11Cfg},
		{keys.MakeTablePrefix(tb12), &db1Cfg},
		{keys.MakeTablePrefix(tb21), &tb21Cfg},
		{keys.MakeTablePrefix(tb22), &defaultZoneConfig},
	}

	for tcNum, tc := range testCases {
		zoneCfg, err := cfg.GetZoneConfigForKey(tc.key)
		if err != nil {
			t.Fatalf("#%d: err=%s", tcNum, err)
		}

		if !reflect.DeepEqual(zoneCfg, tc.zoneCfg) {
			t.Errorf("#%d: bad zone config.\nexpected: %+v\ngot: %+v", tcNum, tc.zoneCfg, zoneCfg)
		}
	}
}
<|endoftext|>


<commit_before>// dedupe - gets rid of identical files remotes which can have duplicate file names (drive, mega)

package operations

import (
	"context"
	"fmt"
	"log"
	"path"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/fs/hash"
	"github.com/rclone/rclone/fs/walk"
	"github.com/spf13/pflag"
)

// dedupeRename renames the objs slice to different names
func dedupeRename(ctx context.Context, f fs.Fs, remote string, objs []fs.Object) {
	doMove := f.Features().Move
	if doMove == nil {
		log.Fatalf("Fs %v doesn't support Move", f)
	}
	ext := path.Ext(remote)
	base := remote[:len(remote)-len(ext)]

outer:
	for i, o := range objs {
		suffix := 1
		newName := fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
		_, err := f.NewObject(ctx, newName)
		for ; err != fs.ErrorObjectNotFound; suffix++ {
			if err != nil {
				fs.CountError(err)
				fs.Errorf(o, "Failed to check for existing object: %v", err)
				continue outer
			}
			if suffix > 100 {
				fs.Errorf(o, "Could not find an available new name")
				continue outer
			}
			newName = fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
			_, err = f.NewObject(ctx, newName)
		}
		if !fs.Config.DryRun {
			newObj, err := doMove(ctx, o, newName)
			if err != nil {
				fs.CountError(err)
				fs.Errorf(o, "Failed to rename: %v", err)
				continue
			}
			fs.Infof(newObj, "renamed from: %v", o)
		} else {
			fs.Logf(remote, "Not renaming to %q as --dry-run", newName)
		}
	}
}

// dedupeDeleteAllButOne deletes all but the one in keep
func dedupeDeleteAllButOne(ctx context.Context, keep int, remote string, objs []fs.Object) {
	for i, o := range objs {
		if i == keep {
			continue
		}
		_ = DeleteFile(ctx, o)
	}
	fs.Logf(remote, "Deleted %d extra copies", len(objs)-1)
}

// dedupeDeleteIdentical deletes all but one of identical (by hash) copies
func dedupeDeleteIdentical(ctx context.Context, ht hash.Type, remote string, objs []fs.Object) (remainingObjs []fs.Object) {
	// See how many of these duplicates are identical
	byHash := make(map[string][]fs.Object, len(objs))
	for _, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil || md5sum == "" {
			remainingObjs = append(remainingObjs, o)
		} else {
			byHash[md5sum] = append(byHash[md5sum], o)
		}
	}

	// Delete identical duplicates, filling remainingObjs with the ones remaining
	for md5sum, hashObjs := range byHash {
		if len(hashObjs) > 1 {
			fs.Logf(remote, "Deleting %d/%d identical duplicates (%v %q)", len(hashObjs)-1, len(hashObjs), ht, md5sum)
			for _, o := range hashObjs[1:] {
				_ = DeleteFile(ctx, o)
			}
		}
		remainingObjs = append(remainingObjs, hashObjs[0])
	}

	return remainingObjs
}

// dedupeInteractive interactively dedupes the slice of objects
func dedupeInteractive(ctx context.Context, f fs.Fs, ht hash.Type, remote string, objs []fs.Object) {
	fmt.Printf("%s: %d duplicates remain\n", remote, len(objs))
	for i, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil {
			md5sum = err.Error()
		}
		fmt.Printf("  %d: %12d bytes, %s, %v %32s\n", i+1, o.Size(), o.ModTime(ctx).Local().Format("2006-01-02 15:04:05.000000000"), ht, md5sum)
	}
	switch config.Command([]string{"sSkip and do nothing", "kKeep just one (choose which in next step)", "rRename all to be different (by changing file.jpg to file-1.jpg)"}) {
	case 's':
	case 'k':
		keep := config.ChooseNumber("Enter the number of the file to keep", 1, len(objs))
		dedupeDeleteAllButOne(ctx, keep-1, remote, objs)
	case 'r':
		dedupeRename(ctx, f, remote, objs)
	}
}

type objectsSortedByModTime []fs.Object

func (objs objectsSortedByModTime) Len() int      { return len(objs) }
func (objs objectsSortedByModTime) Swap(i, j int) { objs[i], objs[j] = objs[j], objs[i] }
func (objs objectsSortedByModTime) Less(i, j int) bool {
	return objs[i].ModTime(context.TODO()).Before(objs[j].ModTime(context.TODO()))
}

// DeduplicateMode is how the dedupe command chooses what to do
type DeduplicateMode int

// Deduplicate modes
const (
	DeduplicateInteractive DeduplicateMode = iota // interactively ask the user
	DeduplicateSkip                               // skip all conflicts
	DeduplicateFirst                              // choose the first object
	DeduplicateNewest                             // choose the newest object
	DeduplicateOldest                             // choose the oldest object
	DeduplicateRename                             // rename the objects
	DeduplicateLargest                            // choose the largest object
)

func (x DeduplicateMode) String() string {
	switch x {
	case DeduplicateInteractive:
		return "interactive"
	case DeduplicateSkip:
		return "skip"
	case DeduplicateFirst:
		return "first"
	case DeduplicateNewest:
		return "newest"
	case DeduplicateOldest:
		return "oldest"
	case DeduplicateRename:
		return "rename"
	case DeduplicateLargest:
		return "largest"
	}
	return "unknown"
}

// Set a DeduplicateMode from a string
func (x *DeduplicateMode) Set(s string) error {
	switch strings.ToLower(s) {
	case "interactive":
		*x = DeduplicateInteractive
	case "skip":
		*x = DeduplicateSkip
	case "first":
		*x = DeduplicateFirst
	case "newest":
		*x = DeduplicateNewest
	case "oldest":
		*x = DeduplicateOldest
	case "rename":
		*x = DeduplicateRename
	case "largest":
		*x = DeduplicateLargest
	default:
		return errors.Errorf("Unknown mode for dedupe %q.", s)
	}
	return nil
}

// Type of the value
func (x *DeduplicateMode) Type() string {
	return "string"
}

// Check it satisfies the interface
var _ pflag.Value = (*DeduplicateMode)(nil)

// dedupeFindDuplicateDirs scans f for duplicate directories
func dedupeFindDuplicateDirs(ctx context.Context, f fs.Fs) ([][]fs.Directory, error) {
	dirs := map[string][]fs.Directory{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListDirs, func(entries fs.DirEntries) error {
		entries.ForDir(func(d fs.Directory) {
			dirs[d.Remote()] = append(dirs[d.Remote()], d)
		})
		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, "find duplicate dirs")
	}
	duplicateDirs := [][]fs.Directory{}
	for _, ds := range dirs {
		if len(ds) > 1 {
			duplicateDirs = append(duplicateDirs, ds)
		}
	}
	return duplicateDirs, nil
}

// dedupeMergeDuplicateDirs merges all the duplicate directories found
func dedupeMergeDuplicateDirs(ctx context.Context, f fs.Fs, duplicateDirs [][]fs.Directory) error {
	mergeDirs := f.Features().MergeDirs
	if mergeDirs == nil {
		return errors.Errorf("%v: can't merge directories", f)
	}
	dirCacheFlush := f.Features().DirCacheFlush
	if dirCacheFlush == nil {
		return errors.Errorf("%v: can't flush dir cache", f)
	}
	for _, dirs := range duplicateDirs {
		if !fs.Config.DryRun {
			fs.Infof(dirs[0], "Merging contents of duplicate directories")
			err := mergeDirs(ctx, dirs)
			if err != nil {
				return errors.Wrap(err, "merge duplicate dirs")
			}
		} else {
			fs.Infof(dirs[0], "NOT Merging contents of duplicate directories as --dry-run")
		}
	}
	dirCacheFlush()
	return nil
}

// Deduplicate interactively finds duplicate files and offers to
// delete all but one or rename them to be different. Only useful with
// Google Drive which can have duplicate file names.
func Deduplicate(ctx context.Context, f fs.Fs, mode DeduplicateMode) error {
	fs.Infof(f, "Looking for duplicates using %v mode.", mode)

	// Find duplicate directories first and fix them - repeat
	// until all fixed
	for {
		duplicateDirs, err := dedupeFindDuplicateDirs(ctx, f)
		if err != nil {
			return err
		}
		if len(duplicateDirs) == 0 {
			break
		}
		err = dedupeMergeDuplicateDirs(ctx, f, duplicateDirs)
		if err != nil {
			return err
		}
		if fs.Config.DryRun {
			break
		}
	}

	// find a hash to use
	ht := f.Hashes().GetOne()

	// Now find duplicate files
	files := map[string][]fs.Object{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListObjects, func(entries fs.DirEntries) error {
		entries.ForObject(func(o fs.Object) {
			remote := o.Remote()
			files[remote] = append(files[remote], o)
		})
		return nil
	})
	if err != nil {
		return err
	}

	for remote, objs := range files {
		if len(objs) > 1 {
			fs.Logf(remote, "Found %d duplicates - deleting identical copies", len(objs))
			objs = dedupeDeleteIdentical(ctx, ht, remote, objs)
			if len(objs) <= 1 {
				fs.Logf(remote, "All duplicates removed")
				continue
			}
			switch mode {
			case DeduplicateInteractive:
				dedupeInteractive(ctx, f, ht, remote, objs)
			case DeduplicateFirst:
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateNewest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, len(objs)-1, remote, objs)
			case DeduplicateOldest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateRename:
				dedupeRename(ctx, f, remote, objs)
			case DeduplicateLargest:
				largest, largestIndex := int64(-1), -1
				for i, obj := range objs {
					size := obj.Size()
					if size > largest {
						largest, largestIndex = size, i
					}
				}
				if largestIndex > -1 {
					dedupeDeleteAllButOne(ctx, largestIndex, remote, objs)
				}
			case DeduplicateSkip:
				// skip
			default:
				//skip
			}
		}
	}
	return nil
}
<commit_msg>operations: display 'All duplicates removed' only if dedupe successful -fixes rclone#3550<commit_after>// dedupe - gets rid of identical files remotes which can have duplicate file names (drive, mega)

package operations

import (
	"context"
	"fmt"
	"log"
	"path"
	"sort"
	"strings"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/config"
	"github.com/rclone/rclone/fs/hash"
	"github.com/rclone/rclone/fs/walk"
	"github.com/spf13/pflag"
)

// dedupeRename renames the objs slice to different names
func dedupeRename(ctx context.Context, f fs.Fs, remote string, objs []fs.Object) {
	doMove := f.Features().Move
	if doMove == nil {
		log.Fatalf("Fs %v doesn't support Move", f)
	}
	ext := path.Ext(remote)
	base := remote[:len(remote)-len(ext)]

outer:
	for i, o := range objs {
		suffix := 1
		newName := fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
		_, err := f.NewObject(ctx, newName)
		for ; err != fs.ErrorObjectNotFound; suffix++ {
			if err != nil {
				fs.CountError(err)
				fs.Errorf(o, "Failed to check for existing object: %v", err)
				continue outer
			}
			if suffix > 100 {
				fs.Errorf(o, "Could not find an available new name")
				continue outer
			}
			newName = fmt.Sprintf("%s-%d%s", base, i+suffix, ext)
			_, err = f.NewObject(ctx, newName)
		}
		if !fs.Config.DryRun {
			newObj, err := doMove(ctx, o, newName)
			if err != nil {
				fs.CountError(err)
				fs.Errorf(o, "Failed to rename: %v", err)
				continue
			}
			fs.Infof(newObj, "renamed from: %v", o)
		} else {
			fs.Logf(remote, "Not renaming to %q as --dry-run", newName)
		}
	}
}

// dedupeDeleteAllButOne deletes all but the one in keep
func dedupeDeleteAllButOne(ctx context.Context, keep int, remote string, objs []fs.Object) {
	for i, o := range objs {
		if i == keep {
			continue
		}
		_ = DeleteFile(ctx, o)
	}
	fs.Logf(remote, "Deleted %d extra copies", len(objs)-1)
}

// dedupeDeleteIdentical deletes all but one of identical (by hash) copies
func dedupeDeleteIdentical(ctx context.Context, ht hash.Type, remote string, objs []fs.Object) (remainingObjs []fs.Object) {
	// See how many of these duplicates are identical
	byHash := make(map[string][]fs.Object, len(objs))
	for _, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil || md5sum == "" {
			remainingObjs = append(remainingObjs, o)
		} else {
			byHash[md5sum] = append(byHash[md5sum], o)
		}
	}

	// Delete identical duplicates, filling remainingObjs with the ones remaining
	for md5sum, hashObjs := range byHash {
		remainingObjs = append(remainingObjs, hashObjs[0])
		if len(hashObjs) > 1 {
			fs.Logf(remote, "Deleting %d/%d identical duplicates (%v %q)", len(hashObjs)-1, len(hashObjs), ht, md5sum)
			for _, o := range hashObjs[1:] {
				err := DeleteFile(ctx, o)
				if err != nil {
					remainingObjs = append(remainingObjs, o)
				}
			}
		}
	}

	return remainingObjs
}

// dedupeInteractive interactively dedupes the slice of objects
func dedupeInteractive(ctx context.Context, f fs.Fs, ht hash.Type, remote string, objs []fs.Object) {
	fmt.Printf("%s: %d duplicates remain\n", remote, len(objs))
	for i, o := range objs {
		md5sum, err := o.Hash(ctx, ht)
		if err != nil {
			md5sum = err.Error()
		}
		fmt.Printf("  %d: %12d bytes, %s, %v %32s\n", i+1, o.Size(), o.ModTime(ctx).Local().Format("2006-01-02 15:04:05.000000000"), ht, md5sum)
	}
	switch config.Command([]string{"sSkip and do nothing", "kKeep just one (choose which in next step)", "rRename all to be different (by changing file.jpg to file-1.jpg)"}) {
	case 's':
	case 'k':
		keep := config.ChooseNumber("Enter the number of the file to keep", 1, len(objs))
		dedupeDeleteAllButOne(ctx, keep-1, remote, objs)
	case 'r':
		dedupeRename(ctx, f, remote, objs)
	}
}

type objectsSortedByModTime []fs.Object

func (objs objectsSortedByModTime) Len() int      { return len(objs) }
func (objs objectsSortedByModTime) Swap(i, j int) { objs[i], objs[j] = objs[j], objs[i] }
func (objs objectsSortedByModTime) Less(i, j int) bool {
	return objs[i].ModTime(context.TODO()).Before(objs[j].ModTime(context.TODO()))
}

// DeduplicateMode is how the dedupe command chooses what to do
type DeduplicateMode int

// Deduplicate modes
const (
	DeduplicateInteractive DeduplicateMode = iota // interactively ask the user
	DeduplicateSkip                               // skip all conflicts
	DeduplicateFirst                              // choose the first object
	DeduplicateNewest                             // choose the newest object
	DeduplicateOldest                             // choose the oldest object
	DeduplicateRename                             // rename the objects
	DeduplicateLargest                            // choose the largest object
)

func (x DeduplicateMode) String() string {
	switch x {
	case DeduplicateInteractive:
		return "interactive"
	case DeduplicateSkip:
		return "skip"
	case DeduplicateFirst:
		return "first"
	case DeduplicateNewest:
		return "newest"
	case DeduplicateOldest:
		return "oldest"
	case DeduplicateRename:
		return "rename"
	case DeduplicateLargest:
		return "largest"
	}
	return "unknown"
}

// Set a DeduplicateMode from a string
func (x *DeduplicateMode) Set(s string) error {
	switch strings.ToLower(s) {
	case "interactive":
		*x = DeduplicateInteractive
	case "skip":
		*x = DeduplicateSkip
	case "first":
		*x = DeduplicateFirst
	case "newest":
		*x = DeduplicateNewest
	case "oldest":
		*x = DeduplicateOldest
	case "rename":
		*x = DeduplicateRename
	case "largest":
		*x = DeduplicateLargest
	default:
		return errors.Errorf("Unknown mode for dedupe %q.", s)
	}
	return nil
}

// Type of the value
func (x *DeduplicateMode) Type() string {
	return "string"
}

// Check it satisfies the interface
var _ pflag.Value = (*DeduplicateMode)(nil)

// dedupeFindDuplicateDirs scans f for duplicate directories
func dedupeFindDuplicateDirs(ctx context.Context, f fs.Fs) ([][]fs.Directory, error) {
	dirs := map[string][]fs.Directory{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListDirs, func(entries fs.DirEntries) error {
		entries.ForDir(func(d fs.Directory) {
			dirs[d.Remote()] = append(dirs[d.Remote()], d)
		})
		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, "find duplicate dirs")
	}
	duplicateDirs := [][]fs.Directory{}
	for _, ds := range dirs {
		if len(ds) > 1 {
			duplicateDirs = append(duplicateDirs, ds)
		}
	}
	return duplicateDirs, nil
}

// dedupeMergeDuplicateDirs merges all the duplicate directories found
func dedupeMergeDuplicateDirs(ctx context.Context, f fs.Fs, duplicateDirs [][]fs.Directory) error {
	mergeDirs := f.Features().MergeDirs
	if mergeDirs == nil {
		return errors.Errorf("%v: can't merge directories", f)
	}
	dirCacheFlush := f.Features().DirCacheFlush
	if dirCacheFlush == nil {
		return errors.Errorf("%v: can't flush dir cache", f)
	}
	for _, dirs := range duplicateDirs {
		if !fs.Config.DryRun {
			fs.Infof(dirs[0], "Merging contents of duplicate directories")
			err := mergeDirs(ctx, dirs)
			if err != nil {
				return errors.Wrap(err, "merge duplicate dirs")
			}
		} else {
			fs.Infof(dirs[0], "NOT Merging contents of duplicate directories as --dry-run")
		}
	}
	dirCacheFlush()
	return nil
}

// Deduplicate interactively finds duplicate files and offers to
// delete all but one or rename them to be different. Only useful with
// Google Drive which can have duplicate file names.
func Deduplicate(ctx context.Context, f fs.Fs, mode DeduplicateMode) error {
	fs.Infof(f, "Looking for duplicates using %v mode.", mode)

	// Find duplicate directories first and fix them - repeat
	// until all fixed
	for {
		duplicateDirs, err := dedupeFindDuplicateDirs(ctx, f)
		if err != nil {
			return err
		}
		if len(duplicateDirs) == 0 {
			break
		}
		err = dedupeMergeDuplicateDirs(ctx, f, duplicateDirs)
		if err != nil {
			return err
		}
		if fs.Config.DryRun {
			break
		}
	}

	// find a hash to use
	ht := f.Hashes().GetOne()

	// Now find duplicate files
	files := map[string][]fs.Object{}
	err := walk.ListR(ctx, f, "", true, fs.Config.MaxDepth, walk.ListObjects, func(entries fs.DirEntries) error {
		entries.ForObject(func(o fs.Object) {
			remote := o.Remote()
			files[remote] = append(files[remote], o)
		})
		return nil
	})
	if err != nil {
		return err
	}

	for remote, objs := range files {
		if len(objs) > 1 {
			fs.Logf(remote, "Found %d duplicates - deleting identical copies", len(objs))
			objs = dedupeDeleteIdentical(ctx, ht, remote, objs)
			if len(objs) <= 1 {
				fs.Logf(remote, "All duplicates removed")
				continue
			}
			switch mode {
			case DeduplicateInteractive:
				dedupeInteractive(ctx, f, ht, remote, objs)
			case DeduplicateFirst:
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateNewest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, len(objs)-1, remote, objs)
			case DeduplicateOldest:
				sort.Sort(objectsSortedByModTime(objs)) // sort oldest first
				dedupeDeleteAllButOne(ctx, 0, remote, objs)
			case DeduplicateRename:
				dedupeRename(ctx, f, remote, objs)
			case DeduplicateLargest:
				largest, largestIndex := int64(-1), -1
				for i, obj := range objs {
					size := obj.Size()
					if size > largest {
						largest, largestIndex = size, i
					}
				}
				if largestIndex > -1 {
					dedupeDeleteAllButOne(ctx, largestIndex, remote, objs)
				}
			case DeduplicateSkip:
				// skip
			default:
				//skip
			}
		}
	}
	return nil
}
<|endoftext|>