<commit_before>package message

import (
	"bytes"
	"io"
	"testing"
)

type sampleMessage struct {
	Message string
	Count   int
}

func runStreamClean(t *testing.T, compressed bool) {
	// Mark this as a helper function.
	t.Helper()

	// Create a byte buffer to use as a transport.
	buffer := &bytes.Buffer{}

	// Wrap it in a message stream.
	stream := NewStream(buffer, compressed)

	// Write a message to the stream.
	if err := stream.Encode(sampleMessage{"content", 100}); err != nil {
		t.Fatal("unable to encode sample message:", err)
	}

	// Read it from the stream.
	var message sampleMessage
	if err := stream.Decode(&message); err != nil {
		t.Fatal("unable to decode sample message:", err)
	}

	// Verify that the transport is left clean.
	if buffer.Len() != 0 {
		t.Error("transport left dirty")
	}
}

func TestStreamClean(t *testing.T) {
	runStreamClean(t, false)
}

func TestCompressedStreamClean(t *testing.T) {
	runStreamClean(t, true)
}

type badReadWriter struct{}

func (r *badReadWriter) Read(_ []byte) (int, error) {
	return 0, io.EOF
}

func (w *badReadWriter) Write(_ []byte) (int, error) {
	return 0, io.ErrClosedPipe
}

func TestEOFDecode(t *testing.T) {
	// Create a bad connection.
	connection := &badReadWriter{}

	// Create a message stream.
	stream := NewStream(connection, false)

	// Ensure that the decode fails and returns the EOF error unwrapped.
	var message sampleMessage
	if err := stream.Decode(&message); err != io.EOF {
		t.Error("expected EOF from connection at EOF")
	}
}

func TestEncodeFail(t *testing.T) {
	// Create a bad connection.
	connection := &badReadWriter{}

	// Create a message stream without compression so that we can check for
	// failures encoding to the underlying connection.
	stream := NewStream(connection, false)

	// Ensure that the encode fails and returns the ErrClosedPipe error
	// unwrapped.
	var message sampleMessage
	if err := stream.Encode(&message); err != io.ErrClosedPipe {
		t.Error("expected ErrClosedPipe from closed connection")
	}
}

func TestEncodeFailOnFlush(t *testing.T) {
	// Create a bad connection.
	connection := &badReadWriter{}

	// Create a message stream with compression so that we can check for
	// failures flushing to the underlying connection. This assumes that the
	// compressors buffer is large enough not to flush during the encode.
	stream := NewStream(connection, true)

	// Ensure that the encode fails and returns the ErrClosedPipe error
	// unwrapped.
	var message sampleMessage
	if err := stream.Encode(&message); err != io.ErrClosedPipe {
		t.Error("expected ErrClosedPipe from closed connection")
	}
}
<commit_msg>Fixed message package tests.<commit_after>package message

import (
	"bytes"
	"io"
	"testing"
)

type sampleMessage struct {
	Message string
	Count   int
}

func TestStreamClean(t *testing.T) {
	// Create a byte buffer to use as a transport.
	buffer := &bytes.Buffer{}

	// Wrap it in a message stream.
	stream := NewStream(buffer)

	// Write a message to the stream.
	if err := stream.Encode(sampleMessage{"content", 100}); err != nil {
		t.Fatal("unable to encode sample message:", err)
	}

	// Read it from the stream.
	var message sampleMessage
	if err := stream.Decode(&message); err != nil {
		t.Fatal("unable to decode sample message:", err)
	}

	// Verify that the transport is left clean.
	if buffer.Len() != 0 {
		t.Error("transport left dirty")
	}
}

type badReadWriter struct{}

func (r *badReadWriter) Read(_ []byte) (int, error) {
	return 0, io.EOF
}

func (w *badReadWriter) Write(_ []byte) (int, error) {
	return 0, io.ErrClosedPipe
}

func TestEOFDecode(t *testing.T) {
	// Create a bad connection.
	connection := &badReadWriter{}

	// Create a message stream.
	stream := NewStream(connection)

	// Ensure that the decode fails and returns the EOF error unwrapped.
	var message sampleMessage
	if err := stream.Decode(&message); err != io.EOF {
		t.Error("expected EOF from connection at EOF")
	}
}

func TestEncodeFail(t *testing.T) {
	// Create a bad connection.
	connection := &badReadWriter{}

	// Create a message stream.
	stream := NewStream(connection)

	// Ensure that the encode fails and returns the ErrClosedPipe error
	// unwrapped.
	var message sampleMessage
	if err := stream.Encode(&message); err != io.ErrClosedPipe {
		t.Error("expected ErrClosedPipe from closed connection")
	}
}
<|endoftext|>


<commit_before>package repository

import (
	"fmt"
	"github.com/timeredbull/tsuru/api/unit"
	"github.com/timeredbull/tsuru/log"
)

const gitServer = "tsuru.plataformas.glb.com"

func Clone(app string, machine int) (err error) {
	u := unit.Unit{Name: app, Machine: machine}
	cmd := fmt.Sprintf("git clone %s /home/application/current --depth 1", GetReadOnlyUrl(app))
	output, err := u.Command(cmd)
	log.Printf("Command output: " + string(output))
	if err != nil {
		return
	}
	return u.ExecuteHook("restart")
}

func Pull(app string, machine int) (err error) {
	u := unit.Unit{Name: app, Machine: machine}
	cmd := fmt.Sprintf("git --git-dir=/home/application/current/.git --work-tree=/home/application/current pull origin master")
	output, err := u.Command(cmd)
	log.Printf("Command output: " + string(output))
	if err != nil {
		return
	}
	return u.ExecuteHook("restart")
}

func GetUrl(app string) string {
	return fmt.Sprintf("git@%s:%s.git", gitServer, app)
}

func GetReadOnlyUrl(app string) string {
	return fmt.Sprintf("git://%s/%s.git", gitServer, app)
}
<commit_msg>Changed hook call when cloning/pulling an app<commit_after>package repository

import (
	"fmt"
	"github.com/timeredbull/tsuru/api/unit"
	"github.com/timeredbull/tsuru/log"
)

const gitServer = "tsuru.plataformas.glb.com"

func Clone(app string, machine int) (err error) {
	u := unit.Unit{Name: app, Machine: machine}
	cmd := fmt.Sprintf("git clone %s /home/application/current --depth 1", GetReadOnlyUrl(app))
	output, err := u.Command(cmd)
	log.Printf("Command output: " + string(output))
	if err != nil {
		return
	}
	return u.ExecuteHook("reload-gunicorn")
}

func Pull(app string, machine int) (err error) {
	u := unit.Unit{Name: app, Machine: machine}
	cmd := fmt.Sprintf("git --git-dir=/home/application/current/.git --work-tree=/home/application/current pull origin master")
	output, err := u.Command(cmd)
	log.Printf("Command output: " + string(output))
	if err != nil {
		return
	}
	return u.ExecuteHook("reload-gunicorn")
}

func GetUrl(app string) string {
	return fmt.Sprintf("git@%s:%s.git", gitServer, app)
}

func GetReadOnlyUrl(app string) string {
	return fmt.Sprintf("git://%s/%s.git", gitServer, app)
}
<|endoftext|>


<commit_before>// Copyright 2012 tsuru 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 repository

import (
	"github.com/globocom/commandmocker"
	"github.com/globocom/config"
	//"github.com/globocom/tsuru/log"
	. "launchpad.net/gocheck"
	//stdlog "log"
	//"os"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	//logFile     *os.File
	tmpdir      string
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	var err error
	s.tmpdir, err = commandmocker.Add("juju", "")
	c.Assert(err, IsNil)
	err = config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	//log.SetLogger(stdlog.New(s.logFile, "[tsuru-tests]", stdlog.LstdFlags|stdlog.Llongfile))
}

func (s *S) TearDownSuite(c *C) {
	commandmocker.Remove(s.tmpdir)
	//defer s.logFile.Close()
}

func (s *S) TestGetGitServerPanicsIfTheConfigFileHasNoServer(c *C) {
	oldConfig, err := config.Get("git")
	c.Assert(err, IsNil)
	err = config.Unset("git")
	c.Assert(err, IsNil)
	defer config.Set("git", oldConfig)
	c.Assert(getGitServer, PanicMatches, "key git:server not found")
}
<commit_msg>repository: gofmt -s -w<commit_after>// Copyright 2012 tsuru 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 repository

import (
	"github.com/globocom/commandmocker"
	"github.com/globocom/config"
	//"github.com/globocom/tsuru/log"
	. "launchpad.net/gocheck"
	//stdlog "log"
	//"os"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	//logFile     *os.File
	tmpdir string
}

var _ = Suite(&S{})

func (s *S) SetUpSuite(c *C) {
	var err error
	s.tmpdir, err = commandmocker.Add("juju", "")
	c.Assert(err, IsNil)
	err = config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	//log.SetLogger(stdlog.New(s.logFile, "[tsuru-tests]", stdlog.LstdFlags|stdlog.Llongfile))
}

func (s *S) TearDownSuite(c *C) {
	commandmocker.Remove(s.tmpdir)
	//defer s.logFile.Close()
}

func (s *S) TestGetGitServerPanicsIfTheConfigFileHasNoServer(c *C) {
	oldConfig, err := config.Get("git")
	c.Assert(err, IsNil)
	err = config.Unset("git")
	c.Assert(err, IsNil)
	defer config.Set("git", oldConfig)
	c.Assert(getGitServer, PanicMatches, "key git:server not found")
}
<|endoftext|>


<commit_before>// Copyright Project Harbor 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.

package replication

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/docker/distribution"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/docker/distribution/manifest/schema2"
	common_http "github.com/goharbor/harbor/src/common/http"
	"github.com/goharbor/harbor/src/common/http/modifier"
	httpauth "github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/utils"
	reg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/common/utils/registry/auth"
	"github.com/goharbor/harbor/src/jobservice/env"
	job_utils "github.com/goharbor/harbor/src/jobservice/job/impl/utils"
	"github.com/goharbor/harbor/src/jobservice/logger"
)

var (
	errCanceled = errors.New("the job is canceled")
)

// Transfer images from source registry to the destination one
type Transfer struct {
	ctx         env.JobContext
	repository  *repository
	srcRegistry *registry
	dstRegistry *registry
	logger      logger.Interface
	retry       bool
}

// ShouldRetry : retry if the error is network error
func (t *Transfer) ShouldRetry() bool {
	return t.retry
}

// MaxFails ...
func (t *Transfer) MaxFails() uint {
	return 3
}

// Validate ....
func (t *Transfer) Validate(params map[string]interface{}) error {
	return nil
}

// Run ...
func (t *Transfer) Run(ctx env.JobContext, params map[string]interface{}) error {
	err := t.run(ctx, params)
	t.retry = retry(err)
	return err
}

func (t *Transfer) run(ctx env.JobContext, params map[string]interface{}) error {
	// initialize
	if err := t.init(ctx, params); err != nil {
		return err
	}
	// try to create project on destination registry
	if err := t.createProject(); err != nil {
		return err
	}
	// replicate the images
	for _, tag := range t.repository.tags {
		digest, manifest, err := t.pullManifest(tag)
		if err != nil {
			return err
		}
		if err := t.transferLayers(tag, manifest.References()); err != nil {
			return err
		}
		if err := t.pushManifest(tag, digest, manifest); err != nil {
			return err
		}
	}

	return nil
}

func (t *Transfer) init(ctx env.JobContext, params map[string]interface{}) error {
	t.logger = ctx.GetLogger()
	t.ctx = ctx

	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}

	// init images that need to be replicated
	t.repository = &repository{
		name: params["repository"].(string),
	}
	if tags, ok := params["tags"]; ok {
		tgs := tags.([]interface{})
		for _, tg := range tgs {
			t.repository.tags = append(t.repository.tags, tg.(string))
		}
	}

	var err error
	// init source registry client
	srcURL := params["src_registry_url"].(string)
	srcInsecure := params["src_registry_insecure"].(bool)
	srcCred := httpauth.NewSecretAuthorizer(secret())
	srcTokenServiceURL := ""
	if stsu, ok := params["src_token_service_url"]; ok {
		srcTokenServiceURL = stsu.(string)
	}

	if len(srcTokenServiceURL) > 0 {
		t.srcRegistry, err = initRegistry(srcURL, srcInsecure, srcCred, t.repository.name, srcTokenServiceURL)
	} else {
		t.srcRegistry, err = initRegistry(srcURL, srcInsecure, srcCred, t.repository.name)
	}
	if err != nil {
		t.logger.Errorf("failed to create client for source registry: %v", err)
		return err
	}

	// init destination registry client
	dstURL := params["dst_registry_url"].(string)
	dstInsecure := params["dst_registry_insecure"].(bool)
	dstCred := auth.NewBasicAuthCredential(
		params["dst_registry_username"].(string),
		params["dst_registry_password"].(string))
	t.dstRegistry, err = initRegistry(dstURL, dstInsecure, dstCred, t.repository.name)
	if err != nil {
		t.logger.Errorf("failed to create client for destination registry: %v", err)
		return err
	}

	// get the tag list first if it is null
	if len(t.repository.tags) == 0 {
		tags, err := t.srcRegistry.ListTag()
		if err != nil {
			t.logger.Errorf("an error occurred while listing tags for the source repository: %v", err)
			return err
		}

		if len(tags) == 0 {
			err = fmt.Errorf("empty tag list for repository %s", t.repository.name)
			t.logger.Error(err)
			return err
		}
		t.repository.tags = tags
	}

	t.logger.Infof("initialization completed: repository: %s, tags: %v, source registry: URL-%s insecure-%v, destination registry: URL-%s insecure-%v",
		t.repository.name, t.repository.tags, t.srcRegistry.url, t.srcRegistry.insecure, t.dstRegistry.url, t.dstRegistry.insecure)

	return nil
}

func initRegistry(url string, insecure bool, credential modifier.Modifier,
	repository string, tokenServiceURL ...string) (*registry, error) {
	registry := &registry{
		url:      url,
		insecure: insecure,
	}

	// use the same transport for clients connecting to docker registry and Harbor UI
	transport := reg.GetHTTPTransport(insecure)

	authorizer := auth.NewStandardTokenAuthorizer(&http.Client{
		Transport: transport,
	}, credential, tokenServiceURL...)
	uam := &job_utils.UserAgentModifier{
		UserAgent: "harbor-registry-client",
	}
	repositoryClient, err := reg.NewRepository(repository, url,
		&http.Client{
			Transport: reg.NewTransport(transport, authorizer, uam),
		})
	if err != nil {
		return nil, err
	}
	registry.Repository = *repositoryClient

	registry.client = common_http.NewClient(
		&http.Client{
			Transport: transport,
		}, credential)
	return registry, nil
}

func (t *Transfer) createProject() error {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}
	p, _ := utils.ParseRepository(t.repository.name)
	project, err := t.srcRegistry.GetProject(p)
	if err != nil {
		t.logger.Errorf("failed to get project %s from source registry: %v", p, err)
		return err
	}

	if err = t.dstRegistry.CreateProject(project); err != nil {
		// other jobs may be also doing the same thing when the current job
		// is creating project or the project has already exist, so when the
		// response code is 409, continue to do next step
		if e, ok := err.(*common_http.Error); ok && e.Code == http.StatusConflict {
			t.logger.Warningf("the status code is 409 when creating project %s on destination registry, try to do next step", p)
			return nil
		}

		t.logger.Errorf("an error occurred while creating project %s on destination registry: %v", p, err)
		return err
	}
	t.logger.Infof("project %s is created on destination registry", p)
	return nil
}

func (t *Transfer) pullManifest(tag string) (string, distribution.Manifest, error) {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return "", nil, errCanceled
	}

	acceptMediaTypes := []string{schema1.MediaTypeManifest, schema2.MediaTypeManifest}
	digest, mediaType, payload, err := t.srcRegistry.PullManifest(tag, acceptMediaTypes)
	if err != nil {
		t.logger.Errorf("an error occurred while pulling manifest of %s:%s from source registry: %v",
			t.repository.name, tag, err)
		return "", nil, err
	}
	t.logger.Infof("manifest of %s:%s pulled successfully from source registry: %s",
		t.repository.name, tag, digest)

	if strings.Contains(mediaType, "application/json") {
		mediaType = schema1.MediaTypeManifest
	}

	manifest, _, err := reg.UnMarshal(mediaType, payload)
	if err != nil {
		t.logger.Errorf("an error occurred while parsing manifest: %v", err)
		return "", nil, err
	}

	return digest, manifest, nil
}

func (t *Transfer) transferLayers(tag string, blobs []distribution.Descriptor) error {
	repository := t.repository.name

	// all blobs(layers and config)
	for _, blob := range blobs {
		if canceled(t.ctx) {
			t.logger.Warning(errCanceled.Error())
			return errCanceled
		}

		digest := blob.Digest.String()
		exist, err := t.dstRegistry.BlobExist(digest)
		if err != nil {
			t.logger.Errorf("an error occurred while checking existence of blob %s of %s:%s on destination registry: %v",
				digest, repository, tag, err)
			return err
		}
		if exist {
			t.logger.Infof("blob %s of %s:%s already exists on the destination registry, skip",
				digest, repository, tag)
			continue
		}

		t.logger.Infof("transferring blob %s of %s:%s to the destination registry ...",
			digest, repository, tag)
		size, data, err := t.srcRegistry.PullBlob(digest)
		if err != nil {
			t.logger.Errorf("an error occurred while pulling blob %s of %s:%s from the source registry: %v",
				digest, repository, tag, err)
			return err
		}
		if data != nil {
			defer data.Close()
		}
		if err = t.dstRegistry.PushBlob(digest, size, data); err != nil {
			t.logger.Errorf("an error occurred while pushing blob %s of %s:%s to the distination registry: %v",
				digest, repository, tag, err)
			return err
		}
		t.logger.Infof("blob %s of %s:%s transferred to the destination registry completed",
			digest, repository, tag)
	}

	return nil
}

func (t *Transfer) pushManifest(tag, digest string, manifest distribution.Manifest) error {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}

	repository := t.repository.name
	dgt, exist, err := t.dstRegistry.ManifestExist(tag)
	if err != nil {
		t.logger.Warningf("an error occurred while checking the existence of manifest of %s:%s on the destination registry: %v, try to push manifest",
			repository, tag, err)
	} else {
		if exist && dgt == digest {
			t.logger.Infof("manifest of %s:%s exists on the destination registry, skip manifest pushing",
				repository, tag)
			return nil
		}
	}

	mediaType, data, err := manifest.Payload()
	if err != nil {
		t.logger.Errorf("an error occurred while getting payload of manifest for %s:%s : %v",
			repository, tag, err)
		return err
	}

	if _, err = t.dstRegistry.PushManifest(tag, mediaType, data); err != nil {
		t.logger.Errorf("an error occurred while pushing manifest of %s:%s to the destination registry: %v",
			repository, tag, err)
		return err
	}
	t.logger.Infof("manifest of %s:%s has been pushed to the destination registry",
		repository, tag)

	return nil
}

func canceled(ctx env.JobContext) bool {
	_, canceled := ctx.OPCommand()
	return canceled
}

func retry(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(net.Error)
	return ok
}

func secret() string {
	return os.Getenv("JOBSERVICE_SECRET")
}
<commit_msg>Skip transfer blob for foreign layers<commit_after>// Copyright Project Harbor 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.

package replication

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/docker/distribution"
	"github.com/docker/distribution/manifest/schema1"
	"github.com/docker/distribution/manifest/schema2"
	common_http "github.com/goharbor/harbor/src/common/http"
	"github.com/goharbor/harbor/src/common/http/modifier"
	httpauth "github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/utils"
	reg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/common/utils/registry/auth"
	"github.com/goharbor/harbor/src/jobservice/env"
	job_utils "github.com/goharbor/harbor/src/jobservice/job/impl/utils"
	"github.com/goharbor/harbor/src/jobservice/logger"
)

var (
	errCanceled = errors.New("the job is canceled")
)

// Transfer images from source registry to the destination one
type Transfer struct {
	ctx         env.JobContext
	repository  *repository
	srcRegistry *registry
	dstRegistry *registry
	logger      logger.Interface
	retry       bool
}

// ShouldRetry : retry if the error is network error
func (t *Transfer) ShouldRetry() bool {
	return t.retry
}

// MaxFails ...
func (t *Transfer) MaxFails() uint {
	return 3
}

// Validate ....
func (t *Transfer) Validate(params map[string]interface{}) error {
	return nil
}

// Run ...
func (t *Transfer) Run(ctx env.JobContext, params map[string]interface{}) error {
	err := t.run(ctx, params)
	t.retry = retry(err)
	return err
}

func (t *Transfer) run(ctx env.JobContext, params map[string]interface{}) error {
	// initialize
	if err := t.init(ctx, params); err != nil {
		return err
	}
	// try to create project on destination registry
	if err := t.createProject(); err != nil {
		return err
	}
	// replicate the images
	for _, tag := range t.repository.tags {
		digest, manifest, err := t.pullManifest(tag)
		if err != nil {
			return err
		}
		if err := t.transferLayers(tag, manifest.References()); err != nil {
			return err
		}
		if err := t.pushManifest(tag, digest, manifest); err != nil {
			return err
		}
	}

	return nil
}

func (t *Transfer) init(ctx env.JobContext, params map[string]interface{}) error {
	t.logger = ctx.GetLogger()
	t.ctx = ctx

	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}

	// init images that need to be replicated
	t.repository = &repository{
		name: params["repository"].(string),
	}
	if tags, ok := params["tags"]; ok {
		tgs := tags.([]interface{})
		for _, tg := range tgs {
			t.repository.tags = append(t.repository.tags, tg.(string))
		}
	}

	var err error
	// init source registry client
	srcURL := params["src_registry_url"].(string)
	srcInsecure := params["src_registry_insecure"].(bool)
	srcCred := httpauth.NewSecretAuthorizer(secret())
	srcTokenServiceURL := ""
	if stsu, ok := params["src_token_service_url"]; ok {
		srcTokenServiceURL = stsu.(string)
	}

	if len(srcTokenServiceURL) > 0 {
		t.srcRegistry, err = initRegistry(srcURL, srcInsecure, srcCred, t.repository.name, srcTokenServiceURL)
	} else {
		t.srcRegistry, err = initRegistry(srcURL, srcInsecure, srcCred, t.repository.name)
	}
	if err != nil {
		t.logger.Errorf("failed to create client for source registry: %v", err)
		return err
	}

	// init destination registry client
	dstURL := params["dst_registry_url"].(string)
	dstInsecure := params["dst_registry_insecure"].(bool)
	dstCred := auth.NewBasicAuthCredential(
		params["dst_registry_username"].(string),
		params["dst_registry_password"].(string))
	t.dstRegistry, err = initRegistry(dstURL, dstInsecure, dstCred, t.repository.name)
	if err != nil {
		t.logger.Errorf("failed to create client for destination registry: %v", err)
		return err
	}

	// get the tag list first if it is null
	if len(t.repository.tags) == 0 {
		tags, err := t.srcRegistry.ListTag()
		if err != nil {
			t.logger.Errorf("an error occurred while listing tags for the source repository: %v", err)
			return err
		}

		if len(tags) == 0 {
			err = fmt.Errorf("empty tag list for repository %s", t.repository.name)
			t.logger.Error(err)
			return err
		}
		t.repository.tags = tags
	}

	t.logger.Infof("initialization completed: repository: %s, tags: %v, source registry: URL-%s insecure-%v, destination registry: URL-%s insecure-%v",
		t.repository.name, t.repository.tags, t.srcRegistry.url, t.srcRegistry.insecure, t.dstRegistry.url, t.dstRegistry.insecure)

	return nil
}

func initRegistry(url string, insecure bool, credential modifier.Modifier,
	repository string, tokenServiceURL ...string) (*registry, error) {
	registry := &registry{
		url:      url,
		insecure: insecure,
	}

	// use the same transport for clients connecting to docker registry and Harbor UI
	transport := reg.GetHTTPTransport(insecure)

	authorizer := auth.NewStandardTokenAuthorizer(&http.Client{
		Transport: transport,
	}, credential, tokenServiceURL...)
	uam := &job_utils.UserAgentModifier{
		UserAgent: "harbor-registry-client",
	}
	repositoryClient, err := reg.NewRepository(repository, url,
		&http.Client{
			Transport: reg.NewTransport(transport, authorizer, uam),
		})
	if err != nil {
		return nil, err
	}
	registry.Repository = *repositoryClient

	registry.client = common_http.NewClient(
		&http.Client{
			Transport: transport,
		}, credential)
	return registry, nil
}

func (t *Transfer) createProject() error {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}
	p, _ := utils.ParseRepository(t.repository.name)
	project, err := t.srcRegistry.GetProject(p)
	if err != nil {
		t.logger.Errorf("failed to get project %s from source registry: %v", p, err)
		return err
	}

	if err = t.dstRegistry.CreateProject(project); err != nil {
		// other jobs may be also doing the same thing when the current job
		// is creating project or the project has already exist, so when the
		// response code is 409, continue to do next step
		if e, ok := err.(*common_http.Error); ok && e.Code == http.StatusConflict {
			t.logger.Warningf("the status code is 409 when creating project %s on destination registry, try to do next step", p)
			return nil
		}

		t.logger.Errorf("an error occurred while creating project %s on destination registry: %v", p, err)
		return err
	}
	t.logger.Infof("project %s is created on destination registry", p)
	return nil
}

func (t *Transfer) pullManifest(tag string) (string, distribution.Manifest, error) {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return "", nil, errCanceled
	}

	acceptMediaTypes := []string{schema1.MediaTypeManifest, schema2.MediaTypeManifest}
	digest, mediaType, payload, err := t.srcRegistry.PullManifest(tag, acceptMediaTypes)
	if err != nil {
		t.logger.Errorf("an error occurred while pulling manifest of %s:%s from source registry: %v",
			t.repository.name, tag, err)
		return "", nil, err
	}
	t.logger.Infof("manifest of %s:%s pulled successfully from source registry: %s",
		t.repository.name, tag, digest)

	if strings.Contains(mediaType, "application/json") {
		mediaType = schema1.MediaTypeManifest
	}

	manifest, _, err := reg.UnMarshal(mediaType, payload)
	if err != nil {
		t.logger.Errorf("an error occurred while parsing manifest: %v", err)
		return "", nil, err
	}

	return digest, manifest, nil
}

func (t *Transfer) transferLayers(tag string, blobs []distribution.Descriptor) error {
	repository := t.repository.name

	// all blobs(layers and config)
	for _, blob := range blobs {
		if canceled(t.ctx) {
			t.logger.Warning(errCanceled.Error())
			return errCanceled
		}

		digest := blob.Digest.String()

		if blob.MediaType == schema2.MediaTypeForeignLayer {
			t.logger.Infof("blob %s of %s:%s is an foreign layer, skip", digest, repository, tag)
			continue
		}

		exist, err := t.dstRegistry.BlobExist(digest)
		if err != nil {
			t.logger.Errorf("an error occurred while checking existence of blob %s of %s:%s on destination registry: %v",
				digest, repository, tag, err)
			return err
		}
		if exist {
			t.logger.Infof("blob %s of %s:%s already exists on the destination registry, skip",
				digest, repository, tag)
			continue
		}

		t.logger.Infof("transferring blob %s of %s:%s to the destination registry ...",
			digest, repository, tag)
		size, data, err := t.srcRegistry.PullBlob(digest)
		if err != nil {
			t.logger.Errorf("an error occurred while pulling blob %s of %s:%s from the source registry: %v",
				digest, repository, tag, err)
			return err
		}
		if data != nil {
			defer data.Close()
		}
		if err = t.dstRegistry.PushBlob(digest, size, data); err != nil {
			t.logger.Errorf("an error occurred while pushing blob %s of %s:%s to the distination registry: %v",
				digest, repository, tag, err)
			return err
		}
		t.logger.Infof("blob %s of %s:%s transferred to the destination registry completed",
			digest, repository, tag)
	}

	return nil
}

func (t *Transfer) pushManifest(tag, digest string, manifest distribution.Manifest) error {
	if canceled(t.ctx) {
		t.logger.Warning(errCanceled.Error())
		return errCanceled
	}

	repository := t.repository.name
	dgt, exist, err := t.dstRegistry.ManifestExist(tag)
	if err != nil {
		t.logger.Warningf("an error occurred while checking the existence of manifest of %s:%s on the destination registry: %v, try to push manifest",
			repository, tag, err)
	} else {
		if exist && dgt == digest {
			t.logger.Infof("manifest of %s:%s exists on the destination registry, skip manifest pushing",
				repository, tag)
			return nil
		}
	}

	mediaType, data, err := manifest.Payload()
	if err != nil {
		t.logger.Errorf("an error occurred while getting payload of manifest for %s:%s : %v",
			repository, tag, err)
		return err
	}

	if _, err = t.dstRegistry.PushManifest(tag, mediaType, data); err != nil {
		t.logger.Errorf("an error occurred while pushing manifest of %s:%s to the destination registry: %v",
			repository, tag, err)
		return err
	}
	t.logger.Infof("manifest of %s:%s has been pushed to the destination registry",
		repository, tag)

	return nil
}

func canceled(ctx env.JobContext) bool {
	_, canceled := ctx.OPCommand()
	return canceled
}

func retry(err error) bool {
	if err == nil {
		return false
	}
	_, ok := err.(net.Error)
	return ok
}

func secret() string {
	return os.Getenv("JOBSERVICE_SECRET")
}
<|endoftext|>


<commit_before>package responder

import (
	"github.com/stretchr/goweb"
	"github.com/stretchr/goweb/context"
	"net/http"
)

// The standard API response object
type standardResponse struct {
	S int         `json:"status"`
	D interface{} `json:"data"`
	E []string    `json:"error"`
}

// The standard API response object
type paginatedResponse struct {
	S      int         `json:"status"`
	D      interface{} `json:"data"`
	E      []string    `json:"error"`
	Limit  int         `json:"limit"`
	Offset int         `json:"offset"`
	Count  int         `json:"total_count"`
}

func RespondOK(ctx context.Context) error {
	response := new(standardResponse)
	response.S = http.StatusOK
	response.D = nil
	response.E = nil
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func WriteResponseObject(ctx context.Context, status int, responseObject interface{}) error {
	return goweb.API.WriteResponseObject(ctx, status, responseObject)
}

func RespondWithData(ctx context.Context, data interface{}) error {
	response := new(standardResponse)
	response.S = http.StatusOK
	response.D = data
	response.E = nil
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func RespondWithError(ctx context.Context, status int, err string) error {
	response := new(standardResponse)
	response.S = status
	response.D = nil
	response.E = append(response.E, err)
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func RespondWithPaginatedData(ctx context.Context, data interface{}, limit, offset, count int) error {
	// make the standard response object
	response := new(paginatedResponse)
	response.S = http.StatusOK
	response.D = data
	response.E = nil
	response.Limit = limit
	response.Offset = offset
	response.Count = count

	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}
<commit_msg>Adding necessary headers to response.<commit_after>package responder

import (
	"github.com/stretchr/goweb"
	"github.com/stretchr/goweb/context"
	"net/http"
)

// The standard API response object
type standardResponse struct {
	S int         `json:"status"`
	D interface{} `json:"data"`
	E []string    `json:"error"`
}

// The standard API response object
type paginatedResponse struct {
	S      int         `json:"status"`
	D      interface{} `json:"data"`
	E      []string    `json:"error"`
	Limit  int         `json:"limit"`
	Offset int         `json:"offset"`
	Count  int         `json:"total_count"`
}

func RespondOK(ctx context.Context) error {
	addResponseHeaders(ctx)
	response := new(standardResponse)
	response.S = http.StatusOK
	response.D = nil
	response.E = nil
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func WriteResponseObject(ctx context.Context, status int, responseObject interface{}) error {
	addResponseHeaders(ctx)
	return goweb.API.WriteResponseObject(ctx, status, responseObject)
}

func RespondWithData(ctx context.Context, data interface{}) error {
	addResponseHeaders(ctx)
	response := new(standardResponse)
	response.S = http.StatusOK
	response.D = data
	response.E = nil
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func RespondWithError(ctx context.Context, status int, err string) error {
	addResponseHeaders(ctx)
	response := new(standardResponse)
	response.S = status
	response.D = nil
	response.E = append(response.E, err)
	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func RespondWithPaginatedData(ctx context.Context, data interface{}, limit, offset, count int) error {
	addResponseHeaders(ctx)
	response := new(paginatedResponse)
	response.S = http.StatusOK
	response.D = data
	response.E = nil
	response.Limit = limit
	response.Offset = offset
	response.Count = count

	return goweb.API.WriteResponseObject(ctx, http.StatusOK, response)
}

func addResponseHeaders(ctx context.Context) {
	ctx.HttpResponseWriter().Header().Set("Connection", "close")
	ctx.HttpResponseWriter().Header().Set("Access-Control-Allow-Headers", "Authorization")
	ctx.HttpResponseWriter().Header().Set("Access-Control-Allow-Methods", "POST, GET, PUT, DELETE, OPTIONS")
	ctx.HttpResponseWriter().Header().Set("Access-Control-Allow-Origin", "*")
}
<|endoftext|>


<commit_before>package mssqlclrgeo

type FIGURE uint8

const (
	FIGURE_INTERIOR_RING      FIGURE = 0x00
	FIGURE_STROKE             FIGURE = 0x01
	FIGURE_EXTERIOR_RING      FIGURE = 0x02
	FIGURE_V2_POINT           FIGURE = 0x00
	FIGURE_V2_LINE            FIGURE = 0x01
	FIGURE_V2_ARC             FIGURE = 0x02
	FIGURE_V2_COMPOSITE_CURVE FIGURE = 0x03
)

type SHAPE uint8

const (
	SHAPE_POINT               SHAPE = 0x01
	SHAPE_LINESTRING          SHAPE = 0x02
	SHAPE_POLYGON             SHAPE = 0x03
	SHAPE_MULTIPOINT          SHAPE = 0x04
	SHAPE_MULTILINESTRING     SHAPE = 0x05
	SHAPE_MULTIPOLYGON        SHAPE = 0x06
	SHAPE_GEOMETRY_COLLECTION SHAPE = 0x07
	//V2
	SHAPE_CIRCULAR_STRING SHAPE = 0x08
	SHAPE_COMPOUND_CURVE  SHAPE = 0x09
	SHAPE_CURVE_POLYGON   SHAPE = 0x0A
	SHAPE_FULL_GLOBE      SHAPE = 0x0B
)

type SEGMENT uint8

const (
	SEGMENT_LINE       SEGMENT = 0x00
	SEGMENT_ARC        SEGMENT = 0x01
	SEGMENT_FIRST_LINE SEGMENT = 0x02
	SEGMENT_FIRST_ARC  SEGMENT = 0x03
)

type Figure struct {
	Attribute FIGURE
	Offset    uint32
}

type Shape struct {
	ParentOffset int32
	FigureOffset int32
	OpenGisType  SHAPE

	index int
}

type Point struct {
	X float64
	Y float64
	Z float64
	M float64
}
type Segment struct {
	Type SEGMENT
}

type SerializationProperties struct {
	H bool //V2
	L bool
	P bool
	V bool
	M bool
	Z bool
}

type Geometry struct {
	SRID       int32
	Version    uint8
	Properties SerializationProperties

	Points   []Point
	Figures  []Figure
	Shapes   []Shape
	Segments []Segment
}

func (s *Shape) getFigures(g Geometry) (figures []Figure) {
	offset := s.FigureOffset

	if len(g.Shapes) > s.index+1 {
		offsetEnd := g.Shapes[s.index+1].FigureOffset
		figures = g.Figures[offset:offsetEnd]

	} else {
		figures = g.Figures[offset:]
	}
	return
}
func (s *Shape) getParent(g Geometry) (parent *Shape) {
	offset := s.ParentOffset

	if offset == 0 {
		return nil
	}

	parent = &g.Shapes[offset]

	return
}
func (g *Geometry) getFiguresFromShape(shapeIdx int) (figures []Figure) {
	offset := g.Shapes[shapeIdx].FigureOffset

	if len(g.Shapes) > shapeIdx+1 {
		offsetEnd := g.Shapes[shapeIdx+1].FigureOffset
		figures = g.Figures[offset:offsetEnd]

	} else {
		figures = g.Figures[offset:]
	}
	return
}
func (g *Geometry) getPointsFromFigure(figureIdx int) (points []Point) {
	offset := g.Figures[figureIdx].Offset

	if len(g.Figures) > figureIdx+1 {
		offsetEnd := g.Figures[figureIdx+1].Offset
		points = g.Points[offset:offsetEnd]
	} else {
		points = g.Points[offset:]
	}
	return
}
<commit_msg>GetFigures<commit_after>package mssqlclrgeo

type FIGURE uint8

const (
	FIGURE_INTERIOR_RING      FIGURE = 0x00
	FIGURE_STROKE             FIGURE = 0x01
	FIGURE_EXTERIOR_RING      FIGURE = 0x02
	FIGURE_V2_POINT           FIGURE = 0x00
	FIGURE_V2_LINE            FIGURE = 0x01
	FIGURE_V2_ARC             FIGURE = 0x02
	FIGURE_V2_COMPOSITE_CURVE FIGURE = 0x03
)

type SHAPE uint8

const (
	SHAPE_POINT               SHAPE = 0x01
	SHAPE_LINESTRING          SHAPE = 0x02
	SHAPE_POLYGON             SHAPE = 0x03
	SHAPE_MULTIPOINT          SHAPE = 0x04
	SHAPE_MULTILINESTRING     SHAPE = 0x05
	SHAPE_MULTIPOLYGON        SHAPE = 0x06
	SHAPE_GEOMETRY_COLLECTION SHAPE = 0x07
	//V2
	SHAPE_CIRCULAR_STRING SHAPE = 0x08
	SHAPE_COMPOUND_CURVE  SHAPE = 0x09
	SHAPE_CURVE_POLYGON   SHAPE = 0x0A
	SHAPE_FULL_GLOBE      SHAPE = 0x0B
)

type SEGMENT uint8

const (
	SEGMENT_LINE       SEGMENT = 0x00
	SEGMENT_ARC        SEGMENT = 0x01
	SEGMENT_FIRST_LINE SEGMENT = 0x02
	SEGMENT_FIRST_ARC  SEGMENT = 0x03
)

type Figure struct {
	Attribute FIGURE
	Offset    uint32
}

type Shape struct {
	ParentOffset int32
	FigureOffset int32
	OpenGisType  SHAPE

	index int
}

type Point struct {
	X float64
	Y float64
	Z float64
	M float64
}
type Segment struct {
	Type SEGMENT
}

type SerializationProperties struct {
	H bool //V2
	L bool
	P bool
	V bool
	M bool
	Z bool
}

type Geometry struct {
	SRID       int32
	Version    uint8
	Properties SerializationProperties

	Points   []Point
	Figures  []Figure
	Shapes   []Shape
	Segments []Segment
}

func (s *Shape) GetFigures(g Geometry) (figures []Figure) {
	offset := s.FigureOffset

	if len(g.Shapes) > s.index+1 {
		offsetEnd := g.Shapes[s.index+1].FigureOffset
		figures = g.Figures[offset:offsetEnd]

	} else {
		figures = g.Figures[offset:]
	}
	return
}

func (g *Geometry) GetShapes() (shapes []Shape) {
	return g.Shapes
}

func (g *Figure) getPoints() (points []Point) {

	return
}

func (s *Shape) getParent(g Geometry) (parent *Shape) {
	offset := s.ParentOffset

	if offset == 0 {
		return nil
	}

	parent = &g.Shapes[offset]

	return
}
func (g *Geometry) getFiguresFromShape(shapeIdx int) (figures []Figure) {
	offset := g.Shapes[shapeIdx].FigureOffset

	if len(g.Shapes) > shapeIdx+1 {
		offsetEnd := g.Shapes[shapeIdx+1].FigureOffset
		figures = g.Figures[offset:offsetEnd]

	} else {
		figures = g.Figures[offset:]
	}
	return
}
func (g *Geometry) getPointsFromFigure(figureIdx int) (points []Point) {
	offset := g.Figures[figureIdx].Offset

	if len(g.Figures) > figureIdx+1 {
		offsetEnd := g.Figures[figureIdx+1].Offset
		points = g.Points[offset:offsetEnd]
	} else {
		points = g.Points[offset:]
	}
	return
}
<|endoftext|>


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

package ui

import (
	"log"
	"runtime"
)

func Version() string {
	return "0.1.1"
}

func Async(fn func()) {
	tasks.Append(fn)
	_DirectQtDrv(nil, 1, 200, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
}

func Run(fn func()) int32 {
	runtime.LockOSThread()
	if qtdrv_init_error != nil {
		log.Println(qtdrv_init_error)
		return -2
	}
	app := NewApplication(nil)
	Async(func() {
		fn()
	})
	return app.Exec()
}

func RunEx(args []string, fn func()) int32 {
	runtime.LockOSThread()
	if qtdrv_init_error != nil {
		log.Println(qtdrv_init_error)
		return -2
	}
	app := NewApplication(args)
	Async(func() {
		fn()
	})
	return app.Exec()
}

func Application() *QApplication {
	return QApplicationInstance()
}
<commit_msg>Added call to runtime.LockOSThread() to ensure that QApplication is created in the main thread<commit_after>// Copyright 2015-2016 visualfc <visualfc@gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ui

import (
	"log"
	"runtime"
)

func init() {
	runtime.LockOSThread()
}

func mainLoop() int32 {
	for f := range mainfunc {
		return f()
	}
	return 0
}

var mainfunc = make(chan func() int32)

func runInOsMainThread(f func() int32) {
	done := make(chan bool, 1)
	mainfunc <- func() int32 {
		res := f()
		done <- true
		return res
	}
	<-done
}

func Version() string {
	return "0.1.1"
}

func Async(fn func()) {
	tasks.Append(fn)
	_DirectQtDrv(nil, 1, 200, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
}

func Run(fn func()) int32 {
	go runInOsMainThread(func() int32 {
		if qtdrv_init_error != nil {
			log.Println(qtdrv_init_error)
			return -2
		}
		app := NewApplication(nil)
		Async(func() {
			fn()
		})
		return app.Exec()
	})
	return mainLoop()
}

func RunEx(args []string, fn func()) int32 {
	go runInOsMainThread(func() int32 {
		if qtdrv_init_error != nil {
			log.Println(qtdrv_init_error)
			return -2
		}
		app := NewApplication(args)
		Async(func() {
			fn()
		})
		return app.Exec()
	})
	return mainLoop()
}

func Application() *QApplication {
	return QApplicationInstance()
}
<|endoftext|>


<commit_before>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		action := fmt.Sprintf("reopen %s", u.target)
		u.handleError(action, e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write to log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		select {
		case e := <-u.errs:
			if e != nil && e != io.EOF {
				panic(e)
			} else {
				return
			}
		case <-u.sigReopen:
			u.reopen()
		case <-u.sigTerm:
			if u.shutdown != nil {
				close(u.shutdown)
				u.shutdown = nil
			}
		case line, ok := <-u.lines:
			if !ok {
				return
			}
			u.logLine(line)
		}
	}
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	message := fmt.Sprintf("Could not %s: %s", action, e.Error())

	if terminal.IsTerminal(1) {
		fmt.Printf("%s\n", message)
		return
	}

	if Stats != nil {
		emsg := strings.Replace(e.Error(), " ", "_", -1)
		Stats.Count("unilog.errors.error_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("FileName:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<commit_msg>Fixing tag usage<commit_after>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		u.handleError("reopen_file", e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write_to_log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		select {
		case e := <-u.errs:
			if e != nil && e != io.EOF {
				panic(e)
			} else {
				return
			}
		case <-u.sigReopen:
			u.reopen()
		case <-u.sigTerm:
			if u.shutdown != nil {
				close(u.shutdown)
				u.shutdown = nil
			}
		case line, ok := <-u.lines:
			if !ok {
				return
			}
			u.logLine(line)
		}
	}
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	message := fmt.Sprintf("Could not %s: %s", action, e.Error())

	if terminal.IsTerminal(1) {
		fmt.Printf("%s\n", message)
		return
	}

	if Stats != nil {
		emsg := fmt.Sprintf("err_action:%s", action)
		Stats.Count("unilog.errors.error_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("file_name:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		manifest := getUpdateManifest(channel)
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		Errf("updating to %s (%s)... ", manifest.Version, manifest.Channel)
		update(build.URL, build.Sha1)
		Errln("done")
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest(Channel)
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		Errln("WARNING: Error autoupdating. The CLI will not be able to update itself.")
		return false
	}
	// TODO: Increase the autoupdate time later
	return f.ModTime().Add(2 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("http://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<commit_msg>add .exe for windows<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		manifest := getUpdateManifest(channel)
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		Errf("updating to %s (%s)... ", manifest.Version, manifest.Channel)
		update(build.URL, build.Sha1)
		Errln("done")
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	// TODO: update plugins
	manifest := getUpdateManifest(Channel)
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		Errln("WARNING: Error autoupdating. The CLI will not be able to update itself.")
		return false
	}
	// TODO: Increase the autoupdate time later
	return f.ModTime().Add(2 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("http://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

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

const (
	LENGTH         = 6
	PORT           = ":8080"
	UPDIRECTORY    = "/tmp/"
	GRILLDIRECTORY = "pomf/build/img/"
	POMFDIRECTORY  = "pomf/build"
	UPADDRESS      = "http://localhost"
	dbUSERNAME     = ""
	dbNAME         = ""
	dbPASSWORD     = ""
	DATABASE       = dbUSERNAME + ":" + dbPASSWORD + "@/" + dbNAME + "?charset=utf8"
	MAXSIZE        = 10 * 1024 * 1024
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `json:"size"`
}

type Response struct {
	Success     bool     `json:"success"`
	ErrorCode   int      `json:"errorcode,omitempty"`
	Description string   `json:"description,omitempty"`
	Files       []Result `json:"files,omitempty"`
}

func generateName() (string, error) {
	name := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	if err != nil {
		return "", err
	}
	var id string
	err = db.QueryRow("select id from files where id=?", name).Scan(&id)
	if err != sql.ErrNoRows {
		generateName()
	}
	db.Close()

	return name, nil
}
func respond(w http.ResponseWriter, output string, resp Response) {
	if resp.ErrorCode != 0 {
		resp.Files = []Result{}
		resp.Success = false
	} else {
		resp.Success = true
	}

	switch output {
	case "json":

		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	case "xml":
		x, err := xml.MarshalIndent(resp, "", "  ")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/xml")
		w.Write(x)

	case "html":
		w.Header().Set("Content-Type", "text/html")
		for _, file := range resp.Files {
			io.WriteString(w, ""+file.URL+"
")
		}

	case "gyazo", "text":
		w.Header().Set("Content-Type", "plain/text")
		for _, file := range resp.Files {
			io.WriteString(w, file.URL+"\n")
		}

	case "csv":
		w.Header().Set("Content-Type", "text/csv")
		io.WriteString(w, "name, url, hash, size\n")
		for _, file := range resp.Files {
			io.WriteString(w, file.Name+","+file.URL+","+file.Hash+","+strconv.FormatInt(file.Size, 10)+"\n")
		}

	default:
		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

	}

}
func grillHandler(w http.ResponseWriter, r *http.Request) {
	kawaii, err := ioutil.ReadDir(GRILLDIRECTORY)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	http.Redirect(w, r, "/img/"+kawaii[rand.Intn(len(kawaii))].Name(), http.StatusFound)
}
func uploadHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	output := r.FormValue("output")
	reader, err := r.MultipartReader()
	resp := Response{Files: []Result{}}
	if err != nil {
		resp.ErrorCode = http.StatusInternalServerError
		resp.Description = err.Error()
		respond(w, output, resp)
		return
	}

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

	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}

		if part.FileName() == "" {
			continue
		}

		s, err := generateName()
		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			break
		}
		extName := filepath.Ext(part.FileName())
		filename := s + extName
		dst, err := os.Create(UPDIRECTORY + filename)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			break
		}
		defer dst.Close()

		h := sha1.New()
		t := io.TeeReader(part, h)
		_, err = io.Copy(dst, t)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			break
		}
		hash := h.Sum(nil)
		sha1 := base64.URLEncoding.EncodeToString(hash)
		stat, _ := dst.Stat()
		size := stat.Size()
		if size > MAXSIZE {
			resp.ErrorCode = http.StatusRequestEntityTooLarge
			resp.Description = err.Error()
			break
		}

		originalname := part.FileName()
		err = db.QueryRow("select originalname, filename, size from files where hash=?", sha1).Scan(&originalname, &filename, &size)
		res := Result{
			URL:  UPADDRESS + "/" + filename,
			Name: originalname,
			Hash: sha1,
			Size: size,
		}
		if err == sql.ErrNoRows {
			query, err := db.Prepare("INSERT into files(hash, originalname, filename, size, date) values(?, ?, ?, ?, ?)")
			if err != nil {
				resp.ErrorCode = http.StatusInternalServerError
				resp.Description = err.Error()
				break
			}
			_, err = query.Exec(res.Hash, res.Name, filename, res.Size, time.Now().Format("2016-01-02"))
			if err != nil {
				resp.ErrorCode = http.StatusInternalServerError
				resp.Description = err.Error()
				break
			}
		}
		resp.Files = append(resp.Files, res)
	}
	respond(w, output, resp)
}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	http.HandleFunc("/grill.php", grillHandler)
	http.Handle("/", http.StripPrefix("/", http.FileServer(http.Dir(POMFDIRECTORY))))
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<commit_msg>Add further error handling<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

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

const (
	LENGTH         = 6
	PORT           = ":8080"
	UPDIRECTORY    = "/tmp/"
	GRILLDIRECTORY = "pomf/build/img/"
	POMFDIRECTORY  = "pomf/build"
	UPADDRESS      = "http://localhost"
	dbUSERNAME     = ""
	dbNAME         = ""
	dbPASSWORD     = ""
	DATABASE       = dbUSERNAME + ":" + dbPASSWORD + "@/" + dbNAME + "?charset=utf8"
	MAXSIZE        = 10 * 1024 * 1024
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `json:"size"`
}

type Response struct {
	Success     bool     `json:"success"`
	ErrorCode   int      `json:"errorcode,omitempty"`
	Description string   `json:"description,omitempty"`
	Files       []Result `json:"files,omitempty"`
}

func generateName() (string, error) {
	name := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	if err != nil {
		return "", err
	}
	var id string
	err = db.QueryRow("select id from files where id=?", name).Scan(&id)
	if err != sql.ErrNoRows {
		generateName()
	}
	db.Close()

	return name, nil
}
func respond(w http.ResponseWriter, output string, resp Response) {
	if resp.ErrorCode != 0 {
		resp.Files = []Result{}
		resp.Success = false
	} else {
		resp.Success = true
	}

	switch output {
	case "json":

		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	case "xml":
		x, err := xml.MarshalIndent(resp, "", "  ")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/xml")
		w.Write(x)

	case "html":
		w.Header().Set("Content-Type", "text/html")
		for _, file := range resp.Files {
			io.WriteString(w, ""+file.URL+"
")
		}

	case "gyazo", "text":
		w.Header().Set("Content-Type", "plain/text")
		for _, file := range resp.Files {
			io.WriteString(w, file.URL+"\n")
		}

	case "csv":
		w.Header().Set("Content-Type", "text/csv")
		io.WriteString(w, "name, url, hash, size\n")
		for _, file := range resp.Files {
			io.WriteString(w, file.Name+","+file.URL+","+file.Hash+","+strconv.FormatInt(file.Size, 10)+"\n")
		}

	default:
		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(resp)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

	}

}

func grillHandler(w http.ResponseWriter, r *http.Request) {
	kawaii, err := ioutil.ReadDir(GRILLDIRECTORY)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	http.Redirect(w, r, "/img/"+kawaii[rand.Intn(len(kawaii))].Name(), http.StatusFound)
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	r.ParseForm()
	output := r.FormValue("output")
	reader, err := r.MultipartReader()
	resp := Response{Files: []Result{}}
	if err != nil {
		resp.ErrorCode = http.StatusInternalServerError
		resp.Description = err.Error()
		respond(w, output, resp)
		return
	}

	db, err := sql.Open("mysql", DATABASE)
	if err != nil {
		resp.ErrorCode = http.StatusInternalServerError
		resp.Description = err.Error()
		return
	}

	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}

		if part.FileName() == "" {
			continue
		}

		s, err := generateName()
		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			break
		}
		extName := filepath.Ext(part.FileName())
		filename := s + extName
		dst, err := os.Create(UPDIRECTORY + filename)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			break
		}
		defer dst.Close()

		h := sha1.New()
		t := io.TeeReader(part, h)
		_, err = io.Copy(dst, t)

		if err != nil {
			resp.ErrorCode = http.StatusInternalServerError
			resp.Description = err.Error()
			respond(w, output, resp)
			break
		}
		hash := h.Sum(nil)
		sha1 := base64.URLEncoding.EncodeToString(hash)
		stat, _ := dst.Stat()
		size := stat.Size()
		if size > MAXSIZE {
			resp.ErrorCode = http.StatusRequestEntityTooLarge
			resp.Description = err.Error()
			break
		}

		originalname := part.FileName()
		err = db.QueryRow("select originalname, filename, size from files where hash=?", sha1).Scan(&originalname, &filename, &size)
		res := Result{
			URL:  UPADDRESS + "/" + filename,
			Name: originalname,
			Hash: sha1,
			Size: size,
		}
		if err == sql.ErrNoRows {
			query, err := db.Prepare("INSERT into files(hash, originalname, filename, size, date) values(?, ?, ?, ?, ?)")
			if err != nil {
				resp.ErrorCode = http.StatusInternalServerError
				resp.Description = err.Error()
				break
			}
			_, err = query.Exec(res.Hash, res.Name, filename, res.Size, time.Now().Format("2016-01-02"))
			if err != nil {
				resp.ErrorCode = http.StatusInternalServerError
				resp.Description = err.Error()
				break
			}
		}
		resp.Files = append(resp.Files, res)
	}
	respond(w, output, resp)
}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	http.HandleFunc("/grill.php", grillHandler)
	http.Handle("/", http.StripPrefix("/", http.FileServer(http.Dir(POMFDIRECTORY))))
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<|endoftext|>


<commit_before>package web

import "embed"

//go:embed static/*
var StaticFiles embed.FS

//go:embed template/*
var TemplateFiles embed.FS
<commit_msg>web: embed static files form dist<commit_after>package web

import "embed"

//go:embed dist/*
var StaticFiles embed.FS

//go:embed template/*
var TemplateFiles embed.FS
<|endoftext|>


<commit_before>// Copyright GoFrame Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gdb_test

import (
	"fmt"
	"testing"

	"github.com/gogf/gf/frame/g"

	"github.com/gogf/gf/test/gtest"
)

func Test_Types(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		if _, err := db.Exec(fmt.Sprintf(`
    CREATE TABLE IF NOT EXISTS types (
        id int(10) unsigned NOT NULL AUTO_INCREMENT,
        %s blob NOT NULL,
        %s binary(8) NOT NULL,
        %s date NOT NULL,
        %s time NOT NULL,
        %s decimal(5,2) NOT NULL,
        %s double NOT NULL,
        %s bit(2) NOT NULL,
        %s tinyint(1) NOT NULL,
        %s bool NOT NULL,
        PRIMARY KEY (id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8;
    `,
			"`blob`",
			"`binary`",
			"`date`",
			"`time`",
			"`decimal`",
			"`double`",
			"`bit`",
			"`tinyint`",
			"`bool`")); err != nil {
			gtest.Error(err)
		}
		defer dropTable("types")
		data := g.Map{
			"id":      1,
			"blob":    "i love gf",
			"binary":  []byte("abcdefgh"),
			"date":    "1880-10-24",
			"time":    "10:00:01",
			"decimal": -123.456,
			"double":  -123.456,
			"bit":     2,
			"tinyint": true,
			"bool":    false,
		}
		r, err := db.Table("types").Data(data).Insert()
		t.Assert(err, nil)
		n, _ := r.RowsAffected()
		t.Assert(n, 1)

		one, err := db.Table("types").One()
		t.Assert(err, nil)
		t.Assert(one["id"].Int(), 1)
		t.Assert(one["blob"].String(), data["blob"])
		t.Assert(one["binary"].String(), data["binary"])
		t.Assert(one["date"].String(), data["date"])
		t.Assert(one["time"].String(), data["time"])
		t.Assert(one["decimal"].String(), -123.46)
		t.Assert(one["double"].String(), data["double"])
		t.Assert(one["bit"].Int(), data["bit"])
		t.Assert(one["tinyint"].Bool(), data["tinyint"])
		t.Assert(one["tinyint"].Bool(), data["tinyint"])
	})
}
<commit_msg>fix issue in uint testing case for package gdb<commit_after>// Copyright GoFrame Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package gdb_test

import (
	"fmt"
	"testing"

	"github.com/gogf/gf/frame/g"

	"github.com/gogf/gf/test/gtest"
)

func Test_Types(t *testing.T) {
	gtest.C(t, func(t *gtest.T) {
		if _, err := db.Exec(fmt.Sprintf(`
    CREATE TABLE IF NOT EXISTS types (
        id int(10) unsigned NOT NULL AUTO_INCREMENT,
        %s blob NOT NULL,
        %s binary(8) NOT NULL,
        %s date NOT NULL,
        %s time NOT NULL,
        %s decimal(5,2) NOT NULL,
        %s double NOT NULL,
        %s bit(2) NOT NULL,
        %s tinyint(1) NOT NULL,
        %s bool NOT NULL,
        PRIMARY KEY (id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8;
    `,
			"`blob`",
			"`binary`",
			"`date`",
			"`time`",
			"`decimal`",
			"`double`",
			"`bit`",
			"`tinyint`",
			"`bool`")); err != nil {
			gtest.Error(err)
		}
		defer dropTable("types")
		data := g.Map{
			"id":      1,
			"blob":    "i love gf",
			"binary":  []byte("abcdefgh"),
			"date":    "1880-10-24",
			"time":    "10:00:01",
			"decimal": -123.456,
			"double":  -123.456,
			"bit":     2,
			"tinyint": true,
			"bool":    false,
		}
		r, err := db.Table("types").Data(data).Insert()
		t.Assert(err, nil)
		n, _ := r.RowsAffected()
		t.Assert(n, 1)

		one, err := db.Table("types").One()
		t.Assert(err, nil)
		t.Assert(one["id"].Int(), 1)
		t.Assert(one["blob"].String(), data["blob"])
		t.Assert(one["binary"].String(), data["binary"])
		t.Assert(one["date"].String(), data["date"])
		t.Assert(one["time"].String(), `0000-01-01 10:00:01`)
		t.Assert(one["decimal"].String(), -123.46)
		t.Assert(one["double"].String(), data["double"])
		t.Assert(one["bit"].Int(), data["bit"])
		t.Assert(one["tinyint"].Bool(), data["tinyint"])
		t.Assert(one["tinyint"].Bool(), data["tinyint"])
	})
}
<|endoftext|>


<commit_before>// Copyright ©2013 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 unit

import (
	"fmt"
	"math"
	"testing"
)

var formatTests = []struct {
	unit   Uniter
	format string
	expect string
}{
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%f", "9.81 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%1.f", "10 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%.1f", "9.8 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2, LengthDim: 0}), "%f", "9.81 kg s^-2"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%e", "6.62606957e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%.3e", "6.626e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%#v", `&unit.Unit{dimensions:unit.Dimensions{4:2, 6:-1}, formatted:"", value:6.62606957e-34}`},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%v", "6.62606957e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%s", "%!s(*Unit=6.62606957e-34 kg^2 s^-1)"},
	{Dimless(math.E), "%v", "2.718281828459045"},
	{Dimless(math.E), "%#v", "unit.Dimless(2.718281828459045)"},
	{Dimless(math.E), "%s", "%!s(unit.Dimless=2.718281828459045)"},
	{Mass(1), "%v", "1 kg"},
	{Mass(1), "%#v", "unit.Mass(1)"},
	{Mass(1), "%s", "%!s(unit.Mass=1 kg)"},
	{Length(1.61619926e-35), "%v", "1.61619926e-35 m"},
	{Length(1.61619926e-35), "%#v", "unit.Length(1.61619926e-35)"},
	{Length(1.61619926e-35), "%s", "%!s(unit.Length=1.61619926e-35 m)"},
	{Time(15.2), "%v", "15.2 s"},
	{Time(15.2), "%#v", "unit.Time(15.2)"},
	{Time(15.2), "%s", "%!s(unit.Time=15.2 s)"},
}

func TestFormat(t *testing.T) {
	for _, ts := range formatTests {
		if r := fmt.Sprintf(ts.format, ts.unit); r != ts.expect {
			t.Errorf("Format %q: got: %q expected: %q", ts.format, r, ts.expect)
		}
	}
}

var initializationTests = []struct {
	unit     *Unit
	expValue float64
	expMap   map[Dimension]int
}{
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), 9.81, Dimensions{MassDim: 1, TimeDim: -2}},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2, LengthDim: 0, CurrentDim: 0}), 9.81, Dimensions{MassDim: 1, TimeDim: -2}},
}

func TestInitialization(t *testing.T) {
	for _, ts := range initializationTests {
		if ts.expValue != ts.unit.value {
			t.Errorf("Value wrong on initialization: got: %v expected: %v", ts.unit.value, ts.expValue)
		}
		if len(ts.expMap) != len(ts.unit.dimensions) {
			t.Errorf("Map mismatch: got: %#v expected: %#v", ts.unit.dimensions, ts.expMap)
		}
		for key, val := range ts.expMap {
			if ts.unit.dimensions[key] != val {
				t.Errorf("Map mismatch: got: %#v expected: %#v", ts.unit.dimensions, ts.expMap)
			}
		}
	}
}

var dimensionEqualityTests = []struct {
	name        string
	a           Uniter
	b           Uniter
	shouldMatch bool
}{
	{"same_empty", New(1.0, Dimensions{}), New(1.0, Dimensions{}), true},
	{"same_one", New(1.0, Dimensions{TimeDim: 1}), New(1.0, Dimensions{TimeDim: 1}), true},
	{"same_mult", New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), true},
	{"diff_one_empty", New(1.0, Dimensions{}), New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), false},
	{"diff_same_dim", New(1.0, Dimensions{TimeDim: 1}), New(1.0, Dimensions{TimeDim: 2}), false},
	{"diff_same_pow", New(1.0, Dimensions{LengthDim: 1}), New(1.0, Dimensions{TimeDim: 1}), false},
	{"diff_numdim", New(1.0, Dimensions{TimeDim: 1, LengthDim: 2}), New(1.0, Dimensions{TimeDim: 2}), false},
	{"diff_one_same_dim", New(1.0, Dimensions{LengthDim: 1, TimeDim: 1}), New(1.0, Dimensions{LengthDim: 1, TimeDim: 2}), false},
}

func TestDimensionEquality(t *testing.T) {
	for _, ts := range dimensionEqualityTests {
		if DimensionsMatch(ts.a, ts.b) != ts.shouldMatch {
			t.Errorf("Dimension comparison incorrect for case %s. got: %v, expected: %v", ts.name, !ts.shouldMatch, ts.shouldMatch)
		}
	}
}

type UnitStructer interface {
	UnitStruct() *UnitStruct
}

type UnitStruct struct {
	current     int
	length      int
	luminosity  int
	mass        int
	temperature int
	time        int
	chemamt     int // For mol
	value       float64
}

// Check if the dimensions of two units are the same
func DimensionsMatchStruct(aU, bU UnitStructer) bool {
	a := aU.UnitStruct()
	b := bU.UnitStruct()
	if a.length != b.length {
		return false
	}
	if a.time != b.time {
		return false
	}
	if a.mass != b.mass {
		return false
	}
	if a.current != b.current {
		return false
	}
	if a.temperature != b.temperature {
		return false
	}
	if a.luminosity != b.luminosity {
		return false
	}
	if a.chemamt != b.chemamt {
		return false
	}
	return true
}

func (u *UnitStruct) UnitStruct() *UnitStruct {
	return u
}

func (u *UnitStruct) Add(aU UnitStructer) *UnitStruct {
	a := aU.UnitStruct()
	if !DimensionsMatchStruct(a, u) {
		panic("dimension mismatch")
	}
	u.value += a.value
	return u
}

func (u *UnitStruct) Mul(aU UnitStructer) *UnitStruct {
	a := aU.UnitStruct()
	u.length += a.length
	u.time += a.time
	u.mass += a.mass
	u.current += a.current
	u.temperature += a.temperature
	u.luminosity += a.luminosity
	u.chemamt += a.chemamt
	u.value *= a.value
	return u
}

var u3 *UnitStruct

func BenchmarkAddStruct(b *testing.B) {
	u1 := &UnitStruct{current: 1, chemamt: 5, value: 10}
	u2 := &UnitStruct{current: 1, chemamt: 5, value: 100}
	for i := 0; i < b.N; i++ {
		u2.Add(u1)
	}
}

func BenchmarkMulStruct(b *testing.B) {
	u1 := &UnitStruct{current: 1, chemamt: 5, value: 10}
	u2 := &UnitStruct{mass: 1, time: 1, value: 100}
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

type UnitMapper interface {
	UnitMap() *UnitMap
}

type dimensionMap int

const (
	LengthB dimensionMap = iota
	TimeB
	MassB
	CurrentB
	TemperatureB
	LuminosityB
	ChemAmtB
)

type UnitMap struct {
	dimension map[dimensionMap]int
	value     float64
}

// Check if the dimensions of two units are the same
func DimensionsMatchMap(aU, bU UnitMapper) bool {
	a := aU.UnitMap()
	b := bU.UnitMap()
	if len(a.dimension) != len(b.dimension) {
		panic("Unequal dimension")
	}
	for key, dimA := range a.dimension {
		dimB, ok := b.dimension[key]
		if !ok || dimA != dimB {
			panic("Unequal dimension")
		}
	}
	return true
}

func (u *UnitMap) UnitMap() *UnitMap {
	return u
}

func (u *UnitMap) Add(aU UnitMapper) *UnitMap {
	a := aU.UnitMap()
	if !DimensionsMatchMap(a, u) {
		panic("dimension mismatch")
	}
	u.value += a.value
	return u
}

func (u *UnitMap) Mul(aU UnitMapper) *UnitMap {
	a := aU.UnitMap()
	for key, val := range a.dimension {
		u.dimension[key] += val
	}
	u.value *= a.value
	return u
}

func BenchmarkAddFloat(b *testing.B) {
	a := 0.0
	c := 10.0
	for i := 0; i < b.N; i++ {
		a += c
	}
}

func BenchmarkMulFloat(b *testing.B) {
	a := 0.0
	c := 10.0
	for i := 0; i < b.N; i++ {
		a *= c
	}
}

func BenchmarkAddMapSmall(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[CurrentB] = 1
	u1.dimension[ChemAmtB] = 5

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[CurrentB] = 1
	u2.dimension[ChemAmtB] = 5
	for i := 0; i < b.N; i++ {
		u2.Add(u1)
	}
}

func BenchmarkMulMapSmallDiff(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[MassB] = 1
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapSmallSame(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[LengthB] = 2
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapLargeDiff(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1
	u1.dimension[MassB] = 1
	u1.dimension[ChemAmtB] = 1
	u1.dimension[TemperatureB] = 1
	u1.dimension[LuminosityB] = 1
	u1.dimension[TimeB] = 1
	u1.dimension[CurrentB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[MassB] = 1
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapLargeSame(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 2
	u1.dimension[MassB] = 2
	u1.dimension[ChemAmtB] = 2
	u1.dimension[TemperatureB] = 2
	u1.dimension[LuminosityB] = 2
	u1.dimension[TimeB] = 2
	u1.dimension[CurrentB] = 2

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[LengthB] = 3
	u2.dimension[MassB] = 3
	u2.dimension[ChemAmtB] = 3
	u2.dimension[TemperatureB] = 3
	u2.dimension[LuminosityB] = 3
	u2.dimension[TimeB] = 3
	u2.dimension[CurrentB] = 3
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}
<commit_msg>split out gesturing test into new test function<commit_after>// Copyright ©2013 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 unit

import (
	"fmt"
	"math"
	"testing"
)

var formatTests = []struct {
	unit   Uniter
	format string
	expect string
}{
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%f", "9.81 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%1.f", "10 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), "%.1f", "9.8 kg s^-2"},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2, LengthDim: 0}), "%f", "9.81 kg s^-2"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%e", "6.62606957e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%.3e", "6.626e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%v", "6.62606957e-34 kg^2 s^-1"},
	{New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1}), "%s", "%!s(*Unit=6.62606957e-34 kg^2 s^-1)"},
	{Dimless(math.E), "%v", "2.718281828459045"},
	{Dimless(math.E), "%#v", "unit.Dimless(2.718281828459045)"},
	{Dimless(math.E), "%s", "%!s(unit.Dimless=2.718281828459045)"},
	{Mass(1), "%v", "1 kg"},
	{Mass(1), "%#v", "unit.Mass(1)"},
	{Mass(1), "%s", "%!s(unit.Mass=1 kg)"},
	{Length(1.61619926e-35), "%v", "1.61619926e-35 m"},
	{Length(1.61619926e-35), "%#v", "unit.Length(1.61619926e-35)"},
	{Length(1.61619926e-35), "%s", "%!s(unit.Length=1.61619926e-35 m)"},
	{Time(15.2), "%v", "15.2 s"},
	{Time(15.2), "%#v", "unit.Time(15.2)"},
	{Time(15.2), "%s", "%!s(unit.Time=15.2 s)"},
}

func TestFormat(t *testing.T) {
	for _, ts := range formatTests {
		if r := fmt.Sprintf(ts.format, ts.unit); r != ts.expect {
			t.Errorf("Format %q: got: %q expected: %q", ts.format, r, ts.expect)
		}
	}
}

func TestGoStringFormat(t *testing.T) {
	expect1 := `&unit.Unit{dimensions:unit.Dimensions{4:2, 6:-1}, formatted:"", value:6.62606957e-34}`
	expect2 := `&unit.Unit{dimensions:unit.Dimensions{6:-1, 4:2}, formatted:"", value:6.62606957e-34}`
	if r := fmt.Sprintf("%#v", New(6.62606957e-34, Dimensions{MassDim: 2, TimeDim: -1})); r != expect1 && r != expect2 {
		t.Errorf("Format %q: got: %q expected: %q", "%#v", r, expect1)
	}
}

var initializationTests = []struct {
	unit     *Unit
	expValue float64
	expMap   map[Dimension]int
}{
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2}), 9.81, Dimensions{MassDim: 1, TimeDim: -2}},
	{New(9.81, Dimensions{MassDim: 1, TimeDim: -2, LengthDim: 0, CurrentDim: 0}), 9.81, Dimensions{MassDim: 1, TimeDim: -2}},
}

func TestInitialization(t *testing.T) {
	for _, ts := range initializationTests {
		if ts.expValue != ts.unit.value {
			t.Errorf("Value wrong on initialization: got: %v expected: %v", ts.unit.value, ts.expValue)
		}
		if len(ts.expMap) != len(ts.unit.dimensions) {
			t.Errorf("Map mismatch: got: %#v expected: %#v", ts.unit.dimensions, ts.expMap)
		}
		for key, val := range ts.expMap {
			if ts.unit.dimensions[key] != val {
				t.Errorf("Map mismatch: got: %#v expected: %#v", ts.unit.dimensions, ts.expMap)
			}
		}
	}
}

var dimensionEqualityTests = []struct {
	name        string
	a           Uniter
	b           Uniter
	shouldMatch bool
}{
	{"same_empty", New(1.0, Dimensions{}), New(1.0, Dimensions{}), true},
	{"same_one", New(1.0, Dimensions{TimeDim: 1}), New(1.0, Dimensions{TimeDim: 1}), true},
	{"same_mult", New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), true},
	{"diff_one_empty", New(1.0, Dimensions{}), New(1.0, Dimensions{TimeDim: 1, LengthDim: -2}), false},
	{"diff_same_dim", New(1.0, Dimensions{TimeDim: 1}), New(1.0, Dimensions{TimeDim: 2}), false},
	{"diff_same_pow", New(1.0, Dimensions{LengthDim: 1}), New(1.0, Dimensions{TimeDim: 1}), false},
	{"diff_numdim", New(1.0, Dimensions{TimeDim: 1, LengthDim: 2}), New(1.0, Dimensions{TimeDim: 2}), false},
	{"diff_one_same_dim", New(1.0, Dimensions{LengthDim: 1, TimeDim: 1}), New(1.0, Dimensions{LengthDim: 1, TimeDim: 2}), false},
}

func TestDimensionEquality(t *testing.T) {
	for _, ts := range dimensionEqualityTests {
		if DimensionsMatch(ts.a, ts.b) != ts.shouldMatch {
			t.Errorf("Dimension comparison incorrect for case %s. got: %v, expected: %v", ts.name, !ts.shouldMatch, ts.shouldMatch)
		}
	}
}

type UnitStructer interface {
	UnitStruct() *UnitStruct
}

type UnitStruct struct {
	current     int
	length      int
	luminosity  int
	mass        int
	temperature int
	time        int
	chemamt     int // For mol
	value       float64
}

// Check if the dimensions of two units are the same
func DimensionsMatchStruct(aU, bU UnitStructer) bool {
	a := aU.UnitStruct()
	b := bU.UnitStruct()
	if a.length != b.length {
		return false
	}
	if a.time != b.time {
		return false
	}
	if a.mass != b.mass {
		return false
	}
	if a.current != b.current {
		return false
	}
	if a.temperature != b.temperature {
		return false
	}
	if a.luminosity != b.luminosity {
		return false
	}
	if a.chemamt != b.chemamt {
		return false
	}
	return true
}

func (u *UnitStruct) UnitStruct() *UnitStruct {
	return u
}

func (u *UnitStruct) Add(aU UnitStructer) *UnitStruct {
	a := aU.UnitStruct()
	if !DimensionsMatchStruct(a, u) {
		panic("dimension mismatch")
	}
	u.value += a.value
	return u
}

func (u *UnitStruct) Mul(aU UnitStructer) *UnitStruct {
	a := aU.UnitStruct()
	u.length += a.length
	u.time += a.time
	u.mass += a.mass
	u.current += a.current
	u.temperature += a.temperature
	u.luminosity += a.luminosity
	u.chemamt += a.chemamt
	u.value *= a.value
	return u
}

var u3 *UnitStruct

func BenchmarkAddStruct(b *testing.B) {
	u1 := &UnitStruct{current: 1, chemamt: 5, value: 10}
	u2 := &UnitStruct{current: 1, chemamt: 5, value: 100}
	for i := 0; i < b.N; i++ {
		u2.Add(u1)
	}
}

func BenchmarkMulStruct(b *testing.B) {
	u1 := &UnitStruct{current: 1, chemamt: 5, value: 10}
	u2 := &UnitStruct{mass: 1, time: 1, value: 100}
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

type UnitMapper interface {
	UnitMap() *UnitMap
}

type dimensionMap int

const (
	LengthB dimensionMap = iota
	TimeB
	MassB
	CurrentB
	TemperatureB
	LuminosityB
	ChemAmtB
)

type UnitMap struct {
	dimension map[dimensionMap]int
	value     float64
}

// Check if the dimensions of two units are the same
func DimensionsMatchMap(aU, bU UnitMapper) bool {
	a := aU.UnitMap()
	b := bU.UnitMap()
	if len(a.dimension) != len(b.dimension) {
		panic("Unequal dimension")
	}
	for key, dimA := range a.dimension {
		dimB, ok := b.dimension[key]
		if !ok || dimA != dimB {
			panic("Unequal dimension")
		}
	}
	return true
}

func (u *UnitMap) UnitMap() *UnitMap {
	return u
}

func (u *UnitMap) Add(aU UnitMapper) *UnitMap {
	a := aU.UnitMap()
	if !DimensionsMatchMap(a, u) {
		panic("dimension mismatch")
	}
	u.value += a.value
	return u
}

func (u *UnitMap) Mul(aU UnitMapper) *UnitMap {
	a := aU.UnitMap()
	for key, val := range a.dimension {
		u.dimension[key] += val
	}
	u.value *= a.value
	return u
}

func BenchmarkAddFloat(b *testing.B) {
	a := 0.0
	c := 10.0
	for i := 0; i < b.N; i++ {
		a += c
	}
}

func BenchmarkMulFloat(b *testing.B) {
	a := 0.0
	c := 10.0
	for i := 0; i < b.N; i++ {
		a *= c
	}
}

func BenchmarkAddMapSmall(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[CurrentB] = 1
	u1.dimension[ChemAmtB] = 5

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[CurrentB] = 1
	u2.dimension[ChemAmtB] = 5
	for i := 0; i < b.N; i++ {
		u2.Add(u1)
	}
}

func BenchmarkMulMapSmallDiff(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[MassB] = 1
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapSmallSame(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[LengthB] = 2
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapLargeDiff(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 1
	u1.dimension[MassB] = 1
	u1.dimension[ChemAmtB] = 1
	u1.dimension[TemperatureB] = 1
	u1.dimension[LuminosityB] = 1
	u1.dimension[TimeB] = 1
	u1.dimension[CurrentB] = 1

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[MassB] = 1
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}

func BenchmarkMulMapLargeSame(b *testing.B) {
	u1 := &UnitMap{value: 10}
	u1.dimension = make(map[dimensionMap]int)
	u1.dimension[LengthB] = 2
	u1.dimension[MassB] = 2
	u1.dimension[ChemAmtB] = 2
	u1.dimension[TemperatureB] = 2
	u1.dimension[LuminosityB] = 2
	u1.dimension[TimeB] = 2
	u1.dimension[CurrentB] = 2

	u2 := &UnitMap{value: 10}
	u2.dimension = make(map[dimensionMap]int)
	u2.dimension[LengthB] = 3
	u2.dimension[MassB] = 3
	u2.dimension[ChemAmtB] = 3
	u2.dimension[TemperatureB] = 3
	u2.dimension[LuminosityB] = 3
	u2.dimension[TimeB] = 3
	u2.dimension[CurrentB] = 3
	for i := 0; i < b.N; i++ {
		u2.Mul(u1)
	}
}
<|endoftext|>