<commit_before>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) ([]byte, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}
	return b, json.Unmarshal(b, i)
}
<commit_msg>add jsonutil.MarshalSimple()<commit_after>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MarshalSimple(v interface{}, prefix, indent string) ([]byte, error) {
	if prefix == "" && indent == "" {
		return json.Marshal(v)
	}
	return json.MarshalIndent(v, prefix, indent)
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) ([]byte, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}
	return b, json.Unmarshal(b, i)
}
<|endoftext|>


<commit_before>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) error {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, i)
}
<commit_msg>update jsonutil.UnmarshalIoReader<commit_after>package jsonutil

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/ioutil"
)

var (
	MarshalPrefix = ""
	MarshalIndent = "    "
)

type mustMarhshalError struct {
	MustMarhshalError string `json:"must_marshal_error"`
}

func MustMarshal(i interface{}, embedError bool) []byte {
	bytes, err := json.Marshal(i)
	if err != nil {
		if embedError {
			e := mustMarhshalError{
				MustMarhshalError: err.Error(),
			}
			bytes, err := json.Marshal(e)
			if err != nil {
				panic(err)
			}
			return bytes
		}
		panic(err)
	}
	return bytes
}

func MustMarshalString(i interface{}, embedError bool) string {
	return string(MustMarshal(i, embedError))
}

func MustMarshalIndent(i interface{}, prefix, indent string, embedError bool) []byte {
	bytes, err := json.MarshalIndent(i, prefix, indent)
	if err != nil {
		panic(err)
	}
	return bytes
}

func PrettyPrint(b []byte) ([]byte, error) {
	var out bytes.Buffer
	err := json.Indent(&out, b, MarshalPrefix, MarshalIndent)
	return out.Bytes(), err
}

func MarshalBase64(i interface{}) (string, error) {
	data, err := json.Marshal(i)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

func UnmarshalIoReader(r io.Reader, i interface{}) ([]byte, error) {
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return b, err
	}
	err = json.Unmarshal(b, i)
	return b, err
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2016, 2017 All Rights Reserved.

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 statecouchdb

import (
	"os"
	"testing"
	"time"

	"github.com/hyperledger/fabric/common/ledger/testutil"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/statedb"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/statedb/commontests"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/version"
	"github.com/hyperledger/fabric/core/ledger/ledgerconfig"
	ledgertestutil "github.com/hyperledger/fabric/core/ledger/testutil"
	"github.com/spf13/viper"
)

func TestMain(m *testing.M) {

	// Read the core.yaml file for default config.
	ledgertestutil.SetupCoreYAMLConfig()
	viper.Set("peer.fileSystemPath", "/tmp/fabric/ledgertests/kvledger/txmgmt/statedb/statecouchdb")

	// Switch to CouchDB
	viper.Set("ledger.state.stateDatabase", "CouchDB")

	// both vagrant and CI have couchdb configured at host "couchdb"
	viper.Set("ledger.state.couchDBConfig.couchDBAddress", "couchdb:5984")
	// Replace with correct username/password such as
	// admin/admin if user security is enabled on couchdb.
	viper.Set("ledger.state.couchDBConfig.username", "")
	viper.Set("ledger.state.couchDBConfig.password", "")
	viper.Set("ledger.state.couchDBConfig.maxRetries", 3)
	viper.Set("ledger.state.couchDBConfig.maxRetriesOnStartup", 10)
	viper.Set("ledger.state.couchDBConfig.requestTimeout", time.Second*35)

	//run the actual test
	result := m.Run()

	//revert to default goleveldb
	viper.Set("ledger.state.stateDatabase", "goleveldb")
	os.Exit(result)
}

func TestBasicRW(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testbasicrw")
		defer env.Cleanup("testbasicrw")
		commontests.TestBasicRW(t, env.DBProvider)

	}
}

func TestMultiDBBasicRW(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testmultidbbasicrw")
		env.Cleanup("testmultidbbasicrw2")
		defer env.Cleanup("testmultidbbasicrw")
		defer env.Cleanup("testmultidbbasicrw2")
		commontests.TestMultiDBBasicRW(t, env.DBProvider)

	}
}

func TestDeletes(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {
		env := NewTestVDBEnv(t)
		env.Cleanup("testdeletes")
		defer env.Cleanup("testdeletes")
		commontests.TestDeletes(t, env.DBProvider)
	}
}

func TestIterator(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testiterator")
		defer env.Cleanup("testiterator")
		commontests.TestIterator(t, env.DBProvider)

	}
}

func TestEncodeDecodeValueAndVersion(t *testing.T) {
	testValueAndVersionEncoding(t, []byte("value1"), version.NewHeight(1, 2))
	testValueAndVersionEncoding(t, []byte{}, version.NewHeight(50, 50))
}

func testValueAndVersionEncoding(t *testing.T, value []byte, version *version.Height) {
	encodedValue := statedb.EncodeValue(value, version)
	val, ver := statedb.DecodeValue(encodedValue)
	testutil.AssertEquals(t, val, value)
	testutil.AssertEquals(t, ver, version)
}

func TestCompositeKey(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		testCompositeKey(t, "ns", "key")
		testCompositeKey(t, "ns", "")

	}
}

func testCompositeKey(t *testing.T, ns string, key string) {
	compositeKey := constructCompositeKey(ns, key)
	t.Logf("compositeKey=%#v", compositeKey)
	ns1, key1 := splitCompositeKey(compositeKey)
	testutil.AssertEquals(t, ns1, ns)
	testutil.AssertEquals(t, key1, key)
}

// The following tests are unique to couchdb, they are not used in leveldb
//  query test
func TestQuery(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testquery")
		defer env.Cleanup("testquery")
		commontests.TestQuery(t, env.DBProvider)

	}
}
<commit_msg>[FAB-3405] Improve unit tests for statecouchdb<commit_after>/*
Copyright IBM Corp. 2016, 2017 All Rights Reserved.

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 statecouchdb

import (
	"os"
	"testing"
	"time"

	"github.com/hyperledger/fabric/common/ledger/testutil"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/statedb"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/statedb/commontests"
	"github.com/hyperledger/fabric/core/ledger/kvledger/txmgmt/version"
	"github.com/hyperledger/fabric/core/ledger/ledgerconfig"
	ledgertestutil "github.com/hyperledger/fabric/core/ledger/testutil"
	"github.com/spf13/viper"
)

func TestMain(m *testing.M) {

	// Read the core.yaml file for default config.
	ledgertestutil.SetupCoreYAMLConfig()
	viper.Set("peer.fileSystemPath", "/tmp/fabric/ledgertests/kvledger/txmgmt/statedb/statecouchdb")

	// Switch to CouchDB
	viper.Set("ledger.state.stateDatabase", "CouchDB")

	// both vagrant and CI have couchdb configured at host "couchdb"
	viper.Set("ledger.state.couchDBConfig.couchDBAddress", "couchdb:5984")
	// Replace with correct username/password such as
	// admin/admin if user security is enabled on couchdb.
	viper.Set("ledger.state.couchDBConfig.username", "")
	viper.Set("ledger.state.couchDBConfig.password", "")
	viper.Set("ledger.state.couchDBConfig.maxRetries", 3)
	viper.Set("ledger.state.couchDBConfig.maxRetriesOnStartup", 10)
	viper.Set("ledger.state.couchDBConfig.requestTimeout", time.Second*35)

	//run the actual test
	result := m.Run()

	//revert to default goleveldb
	viper.Set("ledger.state.stateDatabase", "goleveldb")
	os.Exit(result)
}

func TestBasicRW(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testbasicrw")
		defer env.Cleanup("testbasicrw")
		commontests.TestBasicRW(t, env.DBProvider)

	}
}

func TestMultiDBBasicRW(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testmultidbbasicrw")
		env.Cleanup("testmultidbbasicrw2")
		defer env.Cleanup("testmultidbbasicrw")
		defer env.Cleanup("testmultidbbasicrw2")
		commontests.TestMultiDBBasicRW(t, env.DBProvider)

	}
}

func TestDeletes(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {
		env := NewTestVDBEnv(t)
		env.Cleanup("testdeletes")
		defer env.Cleanup("testdeletes")
		commontests.TestDeletes(t, env.DBProvider)
	}
}

func TestIterator(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testiterator")
		defer env.Cleanup("testiterator")
		commontests.TestIterator(t, env.DBProvider)

	}
}

func TestEncodeDecodeValueAndVersion(t *testing.T) {
	testValueAndVersionEncoding(t, []byte("value1"), version.NewHeight(1, 2))
	testValueAndVersionEncoding(t, []byte{}, version.NewHeight(50, 50))
}

func testValueAndVersionEncoding(t *testing.T, value []byte, version *version.Height) {
	encodedValue := statedb.EncodeValue(value, version)
	val, ver := statedb.DecodeValue(encodedValue)
	testutil.AssertEquals(t, val, value)
	testutil.AssertEquals(t, ver, version)
}

func TestCompositeKey(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		testCompositeKey(t, "ns", "key")
		testCompositeKey(t, "ns", "")

	}
}

func testCompositeKey(t *testing.T, ns string, key string) {
	compositeKey := constructCompositeKey(ns, key)
	t.Logf("compositeKey=%#v", compositeKey)
	ns1, key1 := splitCompositeKey(compositeKey)
	testutil.AssertEquals(t, ns1, ns)
	testutil.AssertEquals(t, key1, key)
}

// The following tests are unique to couchdb, they are not used in leveldb
//  query test
func TestQuery(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {

		env := NewTestVDBEnv(t)
		env.Cleanup("testquery")
		defer env.Cleanup("testquery")
		commontests.TestQuery(t, env.DBProvider)

	}
}

func TestGetStateMultipleKeys(t *testing.T) {
	if ledgerconfig.IsCouchDBEnabled() == true {
		env := NewTestVDBEnv(t)
		env.Cleanup("testgetmultiplekeys")
		defer env.Cleanup("testgetmultiplekeys")
		commontests.TestGetStateMultipleKeys(t, env.DBProvider)
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"errors"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/docker/libmachete/provisioners/api"
	"sort"
	"time"
)

// Builder is a ProvisionerBuilder for AWS.
type Builder struct {
}

// Build creates an AWS provisioner.
func (a Builder) Build(params map[string]string) (api.Provisioner, error) {
	region := params["REGION"]
	if region == "" {
		return nil, errors.New("REGION must be specified")
	}

	accessKey := params["ACCESS_KEY"]
	secretKey := params["SECRET_KEY"]
	sessionToken := params["SESSION_TOKEN"]
	var awsCredentials *credentials.Credentials
	if (accessKey == "" || secretKey == "") && sessionToken == "" {
		// Fall back to shared credentials
		awsCredentials = credentials.NewSharedCredentials("", "")
	} else {
		awsCredentials =
			credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)
	}

	client := CreateClient(region, awsCredentials, 5)

	return New(client), nil
}

type provisioner struct {
	client        ec2iface.EC2API
	sleepFunction func(time.Duration)
}

// New creates a new AWS provisioner that will use the provided EC2 API implementation.
func New(client ec2iface.EC2API) api.Provisioner {
	return &provisioner{client: client, sleepFunction: time.Sleep}
}

// CreateClient creates the actual EC2 API client.
func CreateClient(region string, awsCredentials *credentials.Credentials, retryCount int) ec2iface.EC2API {
	return ec2.New(session.New(aws.NewConfig().
		WithRegion(region).
		WithCredentials(awsCredentials).
		WithLogger(getLogger()).
		WithLogLevel(aws.LogDebugWithHTTPBody).
		WithMaxRetries(retryCount)))
}

func getInstanceSync(client ec2iface.EC2API, instanceID string) (*ec2.Instance, error) {
	result, err := client.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{&instanceID},
	})
	if err != nil {
		return nil, err
	}
	if len(result.Reservations) != 1 || len(result.Reservations[0].Instances) != 1 {
		return nil, &ErrUnexpectedResponse{}
	}
	return result.Reservations[0].Instances[0], nil
}

func tagSync(client ec2iface.EC2API, request CreateInstanceRequest, instance *ec2.Instance) error {
	tags := []*ec2.Tag{}

	// Gather the tag keys in sorted order, to provide predictable tag order.  This is
	// particularly useful for tests.
	var keys []string
	for k := range request.Tags {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		key := k
		value := request.Tags[key]
		tags = append(tags, &ec2.Tag{
			Key:   &key,
			Value: &value,
		})
	}

	_, err := client.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{instance.InstanceId},
		Tags:      tags,
	})
	return err
}

func createInstanceSync(
	client ec2iface.EC2API,
	request CreateInstanceRequest) (*ec2.Instance, error) {

	reservation, err := client.RunInstances(&ec2.RunInstancesInput{
		ImageId:  &request.ImageID,
		MinCount: aws.Int64(1),
		MaxCount: aws.Int64(1),
		Placement: &ec2.Placement{
			AvailabilityZone: &request.AvailabilityZone,
		},
		KeyName:      &request.KeyName,
		InstanceType: &request.InstanceType,
		NetworkInterfaces: []*ec2.InstanceNetworkInterfaceSpecification{{
			DeviceIndex:              aws.Int64(0), // eth0
			Groups:                   makePointerSlice(request.SecurityGroupIds),
			SubnetId:                 &request.SubnetID,
			AssociatePublicIpAddress: &request.AssociatePublicIPAddress,
			DeleteOnTermination:      &request.DeleteOnTermination,
		}},
		Monitoring: &ec2.RunInstancesMonitoringEnabled{
			Enabled: &request.Monitoring,
		},
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Name: &request.IamInstanceProfile,
		},
		EbsOptimized: &request.EbsOptimized,
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			{
				DeviceName: &request.BlockDeviceName,
				Ebs: &ec2.EbsBlockDevice{
					VolumeSize:          &request.RootSize,
					VolumeType:          &request.VolumeType,
					DeleteOnTermination: &request.DeleteOnTermination,
				},
			},
		},
	})

	if err != nil {
		return nil, err
	}

	if reservation == nil || len(reservation.Instances) != 1 {
		return nil, &ErrUnexpectedResponse{}
	}
	return reservation.Instances[0], nil
}

func (p *provisioner) blockUntilInstanceInState(instanceID string, instanceState string) error {
	return WaitUntil(p.sleepFunction, 30, 10*time.Second,
		func() (bool, error) {
			inst, err := getInstanceSync(p.client, instanceID)
			return inst != nil && *inst.State.Name == instanceState, err
		})
}

func (p *provisioner) NewRequestInstance() api.MachineRequest {
	return new(CreateInstanceRequest)
}

func (p *provisioner) CreateInstance(
	req api.MachineRequest) (<-chan api.CreateInstanceEvent, error) {

	request, is := req.(*CreateInstanceRequest)
	if !is {
		return nil, &ErrInvalidRequest{}
	}

	if err := request.Validate(); err != nil {
		return nil, err
	}

	events := make(chan api.CreateInstanceEvent)
	go func() {
		defer close(events)

		events <- api.CreateInstanceEvent{Type: api.CreateInstanceStarted}

		instance, err := createInstanceSync(p.client, *request)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}

		err = p.blockUntilInstanceInState(*instance.InstanceId, ec2.InstanceStateNameRunning)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}

		err = tagSync(p.client, *request, instance)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}
		events <- api.CreateInstanceEvent{
			Type:       api.CreateInstanceCompleted,
			InstanceID: *instance.InstanceId,
		}
	}()

	return events, nil
}

func destroyInstanceSync(client ec2iface.EC2API, instanceID string) error {
	result, err := client.TerminateInstances(&ec2.TerminateInstancesInput{
		InstanceIds: []*string{&instanceID}})

	if err != nil {
		return err
	}

	if len(result.TerminatingInstances) == 1 {
		return nil
	}

	// There was no match for the instance ID.
	return &ErrInvalidRequest{}
}

func (p *provisioner) DestroyInstance(instanceID string) (<-chan api.DestroyInstanceEvent, error) {
	events := make(chan api.DestroyInstanceEvent)

	go func() {
		defer close(events)

		events <- api.DestroyInstanceEvent{Type: api.DestroyInstanceStarted}

		err := destroyInstanceSync(p.client, instanceID)
		if err != nil {
			events <- api.DestroyInstanceEvent{
				Type:  api.DestroyInstanceError,
				Error: err}
			return
		}

		err = p.blockUntilInstanceInState(instanceID, ec2.InstanceStateNameTerminated)
		if err != nil {
			events <- api.DestroyInstanceEvent{
				Error: err,
				Type:  api.DestroyInstanceError,
			}
			return
		}

		events <- api.DestroyInstanceEvent{Type: api.DestroyInstanceCompleted}
	}()

	return events, nil
}
<commit_msg>Use chained credentials for AWS. (#40)<commit_after>package aws

import (
	"errors"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
	"github.com/docker/libmachete/provisioners/api"
	"sort"
	"time"
)

// Builder is a ProvisionerBuilder for AWS.
type Builder struct {
}

// Build creates an AWS provisioner.
func (a Builder) Build(params map[string]string) (api.Provisioner, error) {
	region := params["REGION"]
	if region == "" {
		return nil, errors.New("REGION must be specified")
	}

	accessKey := params["ACCESS_KEY"]
	secretKey := params["SECRET_KEY"]
	sessionToken := params["SESSION_TOKEN"]

	awsCredentials := credentials.NewChainCredentials([]credentials.Provider{
		&credentials.StaticProvider{Value: credentials.Value{
			AccessKeyID:     accessKey,
			SecretAccessKey: secretKey,
			SessionToken:    sessionToken,
		}},
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{},
		&ec2rolecreds.EC2RoleProvider{Client: ec2metadata.New(session.New())},
	})

	client := CreateClient(region, awsCredentials, 5)

	return New(client), nil
}

type provisioner struct {
	client        ec2iface.EC2API
	sleepFunction func(time.Duration)
}

// New creates a new AWS provisioner that will use the provided EC2 API implementation.
func New(client ec2iface.EC2API) api.Provisioner {
	return &provisioner{client: client, sleepFunction: time.Sleep}
}

// CreateClient creates the actual EC2 API client.
func CreateClient(region string, awsCredentials *credentials.Credentials, retryCount int) ec2iface.EC2API {
	return ec2.New(session.New(aws.NewConfig().
		WithRegion(region).
		WithCredentials(awsCredentials).
		WithLogger(getLogger()).
		WithLogLevel(aws.LogDebugWithHTTPBody).
		WithMaxRetries(retryCount)))
}

func getInstanceSync(client ec2iface.EC2API, instanceID string) (*ec2.Instance, error) {
	result, err := client.DescribeInstances(&ec2.DescribeInstancesInput{
		InstanceIds: []*string{&instanceID},
	})
	if err != nil {
		return nil, err
	}
	if len(result.Reservations) != 1 || len(result.Reservations[0].Instances) != 1 {
		return nil, &ErrUnexpectedResponse{}
	}
	return result.Reservations[0].Instances[0], nil
}

func tagSync(client ec2iface.EC2API, request CreateInstanceRequest, instance *ec2.Instance) error {
	tags := []*ec2.Tag{}

	// Gather the tag keys in sorted order, to provide predictable tag order.  This is
	// particularly useful for tests.
	var keys []string
	for k := range request.Tags {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	for _, k := range keys {
		key := k
		value := request.Tags[key]
		tags = append(tags, &ec2.Tag{
			Key:   &key,
			Value: &value,
		})
	}

	_, err := client.CreateTags(&ec2.CreateTagsInput{
		Resources: []*string{instance.InstanceId},
		Tags:      tags,
	})
	return err
}

func createInstanceSync(
	client ec2iface.EC2API,
	request CreateInstanceRequest) (*ec2.Instance, error) {

	reservation, err := client.RunInstances(&ec2.RunInstancesInput{
		ImageId:  &request.ImageID,
		MinCount: aws.Int64(1),
		MaxCount: aws.Int64(1),
		Placement: &ec2.Placement{
			AvailabilityZone: &request.AvailabilityZone,
		},
		KeyName:      &request.KeyName,
		InstanceType: &request.InstanceType,
		NetworkInterfaces: []*ec2.InstanceNetworkInterfaceSpecification{{
			DeviceIndex:              aws.Int64(0), // eth0
			Groups:                   makePointerSlice(request.SecurityGroupIds),
			SubnetId:                 &request.SubnetID,
			AssociatePublicIpAddress: &request.AssociatePublicIPAddress,
			DeleteOnTermination:      &request.DeleteOnTermination,
		}},
		Monitoring: &ec2.RunInstancesMonitoringEnabled{
			Enabled: &request.Monitoring,
		},
		IamInstanceProfile: &ec2.IamInstanceProfileSpecification{
			Name: &request.IamInstanceProfile,
		},
		EbsOptimized: &request.EbsOptimized,
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			{
				DeviceName: &request.BlockDeviceName,
				Ebs: &ec2.EbsBlockDevice{
					VolumeSize:          &request.RootSize,
					VolumeType:          &request.VolumeType,
					DeleteOnTermination: &request.DeleteOnTermination,
				},
			},
		},
	})

	if err != nil {
		return nil, err
	}

	if reservation == nil || len(reservation.Instances) != 1 {
		return nil, &ErrUnexpectedResponse{}
	}
	return reservation.Instances[0], nil
}

func (p *provisioner) blockUntilInstanceInState(instanceID string, instanceState string) error {
	return WaitUntil(p.sleepFunction, 30, 10*time.Second,
		func() (bool, error) {
			inst, err := getInstanceSync(p.client, instanceID)
			return inst != nil && *inst.State.Name == instanceState, err
		})
}

func (p *provisioner) NewRequestInstance() api.MachineRequest {
	return new(CreateInstanceRequest)
}

func (p *provisioner) CreateInstance(
	req api.MachineRequest) (<-chan api.CreateInstanceEvent, error) {

	request, is := req.(*CreateInstanceRequest)
	if !is {
		return nil, &ErrInvalidRequest{}
	}

	if err := request.Validate(); err != nil {
		return nil, err
	}

	events := make(chan api.CreateInstanceEvent)
	go func() {
		defer close(events)

		events <- api.CreateInstanceEvent{Type: api.CreateInstanceStarted}

		instance, err := createInstanceSync(p.client, *request)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}

		err = p.blockUntilInstanceInState(*instance.InstanceId, ec2.InstanceStateNameRunning)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}

		err = tagSync(p.client, *request, instance)
		if err != nil {
			events <- api.CreateInstanceEvent{
				Error: err,
				Type:  api.CreateInstanceError,
			}
			return
		}
		events <- api.CreateInstanceEvent{
			Type:       api.CreateInstanceCompleted,
			InstanceID: *instance.InstanceId,
		}
	}()

	return events, nil
}

func destroyInstanceSync(client ec2iface.EC2API, instanceID string) error {
	result, err := client.TerminateInstances(&ec2.TerminateInstancesInput{
		InstanceIds: []*string{&instanceID}})

	if err != nil {
		return err
	}

	if len(result.TerminatingInstances) == 1 {
		return nil
	}

	// There was no match for the instance ID.
	return &ErrInvalidRequest{}
}

func (p *provisioner) DestroyInstance(instanceID string) (<-chan api.DestroyInstanceEvent, error) {
	events := make(chan api.DestroyInstanceEvent)

	go func() {
		defer close(events)

		events <- api.DestroyInstanceEvent{Type: api.DestroyInstanceStarted}

		err := destroyInstanceSync(p.client, instanceID)
		if err != nil {
			events <- api.DestroyInstanceEvent{
				Type:  api.DestroyInstanceError,
				Error: err}
			return
		}

		err = p.blockUntilInstanceInState(instanceID, ec2.InstanceStateNameTerminated)
		if err != nil {
			events <- api.DestroyInstanceEvent{
				Error: err,
				Type:  api.DestroyInstanceError,
			}
			return
		}

		events <- api.DestroyInstanceEvent{Type: api.DestroyInstanceCompleted}
	}()

	return events, nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package format

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.5.0"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(color.Error, chr)
		}
	}

	fmt.Fprintln(color.Error)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(color.Error, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(color.Error, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(color.Error, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(color.Error, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(color.Error, ", ")
		}
		num++
	}
	fmt.Fprintln(color.Error)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(color.Error)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}

		fmt.Fprintf(color.Error, "%s%s %s %s\n",
			blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)

			fmt.Fprintf(color.Error, "%s%s %s\n",
				yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(color.Error, " ")
		}
	}
	r.Fprintln(color.Error, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(color.Error, Version)
	pad(rightmost - len(Author))
	y.Fprintln(color.Error, Author)
	pad(rightmost - len(Description))
	y.Fprintf(color.Error, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<commit_msg>v3.5.1 release<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package format

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/requests"
	"github.com/fatih/color"
)

// Banner is the ASCII art logo used within help output.
const Banner = `
        .+++:.            :                             .+++.
      +W@@@@@@8        &+W@#               o8W8:      +W@@@@@@#.   oW@@@W#+
     &@#+   .o@##.    .@@@o@W.o@@o       :@@#&W8o    .@#:  .:oW+  .@#+++&#&
    +@&        &@&     #@8 +@W@&8@+     :@W.   +@8   +@:          .@8
    8@          @@     8@o  8@8  WW    .@W      W@+  .@W.          o@#:
    WW          &@o    &@:  o@+  o@+   #@.      8@o   +W@#+.        +W@8:
    #@          :@W    &@+  &@+   @8  :@o       o@o     oW@@W+        oW@8
    o@+          @@&   &@+  &@+   #@  &@.      .W@W       .+#@&         o@W.
     WW         +@W@8. &@+  :&    o@+ #@      :@W&@&         &@:  ..     :@o
     :@W:      o@# +Wo &@+        :W: +@W&o++o@W. &@&  8@#o+&@W.  #@:    o@+
      :W@@WWWW@@8       +              :&W@@@@&    &W  .o#@@W&.   :W@WWW@@&
        +o&&&&+.                                                    +oooo.
`

const (
	// Version is used to display the current version of Amass.
	Version = "v3.5.1"

	// Author is used to display the Amass Project Team.
	Author = "OWASP Amass Project - @owaspamass"

	// Description is the slogan for the Amass Project.
	Description = "In-depth Attack Surface Mapping and Asset Discovery"
)

var (
	// Colors used to ease the reading of program output
	y      = color.New(color.FgHiYellow)
	g      = color.New(color.FgHiGreen)
	r      = color.New(color.FgHiRed)
	b      = color.New(color.FgHiBlue)
	fgR    = color.New(color.FgRed)
	fgY    = color.New(color.FgYellow)
	yellow = color.New(color.FgHiYellow).SprintFunc()
	green  = color.New(color.FgHiGreen).SprintFunc()
	blue   = color.New(color.FgHiBlue).SprintFunc()
)

// ASNSummaryData stores information related to discovered ASs and netblocks.
type ASNSummaryData struct {
	Name      string
	Netblocks map[string]int
}

// UpdateSummaryData updates the summary maps using the provided requests.Output data.
func UpdateSummaryData(output *requests.Output, tags map[string]int, asns map[int]*ASNSummaryData) {
	tags[output.Tag]++

	for _, addr := range output.Addresses {
		data, found := asns[addr.ASN]
		if !found {
			asns[addr.ASN] = &ASNSummaryData{
				Name:      addr.Description,
				Netblocks: make(map[string]int),
			}
			data = asns[addr.ASN]
		}
		// Increment how many IPs were in this netblock
		data.Netblocks[addr.Netblock.String()]++
	}
}

// PrintEnumerationSummary outputs the summary information utilized by the command-line tools.
func PrintEnumerationSummary(total int, tags map[string]int, asns map[int]*ASNSummaryData, demo bool) {
	pad := func(num int, chr string) {
		for i := 0; i < num; i++ {
			b.Fprint(color.Error, chr)
		}
	}

	fmt.Fprintln(color.Error)
	// Print the header information
	title := "OWASP Amass "
	site := "https://github.com/OWASP/Amass"
	b.Fprint(color.Error, title+Version)
	num := 80 - (len(title) + len(Version) + len(site))
	pad(num, " ")
	b.Fprintf(color.Error, "%s\n", site)
	pad(8, "----------")
	fmt.Fprintf(color.Error, "\n%s%s", yellow(strconv.Itoa(total)), green(" names discovered - "))
	// Print the stats using tag information
	num, length := 1, len(tags)
	for k, v := range tags {
		fmt.Fprintf(color.Error, "%s: %s", green(k), yellow(strconv.Itoa(v)))
		if num < length {
			g.Fprint(color.Error, ", ")
		}
		num++
	}
	fmt.Fprintln(color.Error)

	if len(asns) == 0 {
		return
	}
	// Another line gets printed
	pad(8, "----------")
	fmt.Fprintln(color.Error)
	// Print the ASN and netblock information
	for asn, data := range asns {
		asnstr := strconv.Itoa(asn)
		datastr := data.Name

		if demo && asn > 0 {
			asnstr = censorString(asnstr, 0, len(asnstr))
			datastr = censorString(datastr, 0, len(datastr))
		}

		fmt.Fprintf(color.Error, "%s%s %s %s\n",
			blue("ASN: "), yellow(asnstr), green("-"), green(datastr))

		for cidr, ips := range data.Netblocks {
			countstr := strconv.Itoa(ips)
			cidrstr := cidr

			if demo {
				cidrstr = censorNetBlock(cidrstr)
			}

			countstr = fmt.Sprintf("\t%-4s", countstr)
			cidrstr = fmt.Sprintf("\t%-18s", cidrstr)

			fmt.Fprintf(color.Error, "%s%s %s\n",
				yellow(cidrstr), yellow(countstr), blue("Subdomain Name(s)"))
		}
	}
}

// PrintBanner outputs the Amass banner the same for all tools.
func PrintBanner() {
	y := color.New(color.FgHiYellow)
	r := color.New(color.FgHiRed)
	rightmost := 76

	pad := func(num int) {
		for i := 0; i < num; i++ {
			fmt.Fprint(color.Error, " ")
		}
	}
	r.Fprintln(color.Error, Banner)
	pad(rightmost - len(Version))
	y.Fprintln(color.Error, Version)
	pad(rightmost - len(Author))
	y.Fprintln(color.Error, Author)
	pad(rightmost - len(Description))
	y.Fprintf(color.Error, "%s\n\n\n", Description)
}

func censorDomain(input string) string {
	return censorString(input, strings.Index(input, "."), len(input))
}

func censorIP(input string) string {
	return censorString(input, 0, strings.LastIndex(input, "."))
}

func censorNetBlock(input string) string {
	return censorString(input, 0, strings.Index(input, "/"))
}

func censorString(input string, start, end int) string {
	runes := []rune(input)
	for i := start; i < end; i++ {
		if runes[i] == '.' ||
			runes[i] == '/' ||
			runes[i] == '-' ||
			runes[i] == ' ' {
			continue
		}
		runes[i] = 'x'
	}
	return string(runes)
}

// OutputLineParts returns the parts of a line to be printed for a requests.Output.
func OutputLineParts(out *requests.Output, src, addrs, demo bool) (source, name, ips string) {
	if src {
		source = fmt.Sprintf("%-18s", "["+out.Source+"] ")
	}
	if addrs {
		for i, a := range out.Addresses {
			if i != 0 {
				ips += ","
			}
			if demo {
				ips += censorIP(a.Address.String())
			} else {
				ips += a.Address.String()
			}
		}
		if ips == "" {
			ips = "N/A"
		}
	}
	name = out.Name
	if demo {
		name = censorDomain(name)
	}
	return
}

// DesiredAddrTypes removes undesired address types from the AddressInfo slice.
func DesiredAddrTypes(addrs []requests.AddressInfo, ipv4, ipv6 bool) []requests.AddressInfo {
	if !ipv4 && !ipv6 {
		return addrs
	}

	var keep []requests.AddressInfo
	for _, addr := range addrs {
		if net.IsIPv4(addr.Address) && !ipv4 {
			continue
		} else if net.IsIPv6(addr.Address) && !ipv6 {
			continue
		}
		keep = append(keep, addr)
	}
	return keep
}
<|endoftext|>


<commit_before>package proxy

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"
)

func (apiServer *ApiServer) newHandler(api *Api) func(http.ResponseWriter, *http.Request) {
	bindPath := api.Path
	log.Println(apiServer.ServerConf.Port, api.Name, "bind path [", bindPath, "]")
	return func(rw http.ResponseWriter, req *http.Request) {
		id := api.PvInc()
		uniqId := apiServer.GetUniqReqId(id)

		broadData := apiServer.initBroadCastData(req)
		broadData.Id = uniqId
		start := time.Now()
		defer func() {
			used := float64(time.Now().Sub(start).Nanoseconds()) / 1e6
			broadData.SetData("used", used)
			go apiServer.BroadcastApiReq(api, broadData)
		}()

		rw.Header().Set("Api-Proxy-Version", API_PROXY_VERSION)
		log.Println(req.URL.String())

		relPath := req.URL.Path[len(bindPath):]
		req.Header.Set("Connection", "close")

		logData := make(map[string]interface{})

		body, err := ioutil.ReadAll(req.Body)

		logData["body_len"] = len(body)

		if err != nil {
			rw.WriteHeader(http.StatusBadGateway)
			rw.Write([]byte("read body failed"))
			broadData.SetError(err.Error())
			return
		}
		//get body must by before  parse callerPref

		hosts, masterHost, cpf := api.getApiHostsByReq(req)
		broadData.SetData("master", masterHost)
		broadData.SetData("remote", cpf.GetIp())

		_uri := req.URL.Path
		if req.URL.RawQuery != "" {
			_uri += "?" + req.URL.RawQuery
		}
		mainLogStr := fmt.Sprintf("uniqid=%s port=%d remote=%s method=%s uri=%s master=%s hostsTotal=%d refer=%s", uniqId, apiServer.ServerConf.Port, req.RemoteAddr, req.Method, _uri, masterHost, len(hosts), req.Referer())

		var printLog = func(logIndex int) {
			totalUsed := fmt.Sprintf("%.3fms", float64(time.Now().Sub(start).Nanoseconds())/1e6)
			log.Println(fmt.Sprintf("logindex=%d", logIndex), mainLogStr, fmt.Sprintf("totalUsed=%s", totalUsed), logData)
		}
		defer (func() {
			printLog(0)
		})()

		rw.Header().Set("Api-Proxy-Master", masterHost)

		addrInfo := strings.Split(req.RemoteAddr, ":")

		if len(hosts) == 0 {
			logData["hostTotal"] = 0
			rw.WriteHeader(http.StatusBadGateway)
			rw.Write([]byte("no backend hosts"))
			broadData.SetError("no backend hosts")
			return
		}

		bodyLen := int64(len(body))
		reqs := make([]*apiHostRequest, 0)

		//build request
		for _, api_host := range hosts {
			isMaster := api_host.Name == masterHost
			urlNew := ""

			serverUrl := api_host.Url
			if api.HostAsProxy {
				serverUrl = "http://" + req.Host + api.Path
			}
			if strings.HasSuffix(urlNew, "/") {
				urlNew += strings.TrimLeft(relPath, "/")
			} else {
				urlNew += relPath
			}
			if req.URL.RawQuery != "" {
				urlNew += "?" + req.URL.RawQuery
			}

			rawUrl := api_host.Url + urlNew

			if isMaster {
				rw.Header().Set("Api-Proxy-Raw-Url", rawUrl)
			}

			urlNew = serverUrl + urlNew

			broadData.SetData("raw_url", rawUrl)

			reqNew, err := http.NewRequest(req.Method, urlNew, ioutil.NopCloser(bytes.NewReader(body)))
			if err != nil {
				log.Println("build req failed:", err)
				if isMaster {
					rw.WriteHeader(http.StatusBadGateway)
					rw.Write([]byte("error:" + err.Error() + "\nraw_url:" + rawUrl))
				}
				broadData.SetError(err.Error())
				return

			}
			copyHeaders(reqNew.Header, req.Header)

			//only accept gzip encode
			accept_encoding := reqNew.Header.Get("Accept-Encoding")
			if accept_encoding != "" && (In_StringSlice("gzip", reqNew.Header["Accept-Encoding"]) || strings.Contains(accept_encoding, "gzip")) {
				reqNew.Header.Set("Accept-Encoding", "gzip")
			}

			if bodyLen > 0 {
				reqNew.ContentLength = bodyLen
				reqNew.Header.Set("Content-Length", fmt.Sprintf("%d", bodyLen))
			}

			reqNew.Header.Set("HTTP_X_FORWARDED_FOR", addrInfo[0])

			transport := &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				Dial: (&net.Dialer{
					Timeout:   time.Duration(api.TimeoutMs) * time.Millisecond,
					KeepAlive: 30 * time.Second,
				}).Dial,
				TLSHandshakeTimeout: 10 * time.Second,
			}
			if api.HostAsProxy {
				transport.Proxy = func(req *http.Request) (*url.URL, error) {
					return url.Parse(api_host.Url)
				}
			}
			apiReq := &apiHostRequest{
				req:       reqNew,
				transport: transport,
				apiHost:   api_host,
				isMaster:  isMaster,
				urlNew:    urlNew,
				urlRaw:    rawUrl,
			}
			reqs = append(reqs, apiReq)
		}

		//call master at first sync
		for index, apiReq := range reqs {
			if !apiReq.isMaster {
				continue
			}
			backLog := make(map[string]interface{})
			logData[fmt.Sprintf("host_%s_%d", apiReq.apiHost.Name, index)] = backLog
			hostStart := time.Now()
			backLog["isMaster"] = apiReq.isMaster
			backLog["start"] = fmt.Sprintf("%.4f", float64(hostStart.UnixNano())/1e9)
			resp, err := apiReq.transport.RoundTrip(apiReq.req)
			if err != nil {
				log.Println("fetch "+apiReq.urlNew, err)
				rw.WriteHeader(http.StatusBadGateway)
				rw.Write([]byte("fetch_error:" + err.Error() + "\nraw_url:" + apiReq.urlRaw + "\nnew_url:" + apiReq.urlNew))
				return
			}
			defer resp.Body.Close()

			apiServer.addBroadCastDataResponse(broadData, resp)

			for k, vs := range resp.Header {
				for _, v := range vs {
					rw.Header().Add(k, v)
				}
			}
			rw.Header().Set("Connection", "close")
			rw.WriteHeader(resp.StatusCode)
			backLog["status"] = resp.StatusCode
			n, err := io.Copy(rw, resp.Body)
			if err != nil {
				log.Println(apiReq.urlNew, "io.copy:", n, err)
			}
			hostEnd := time.Now()
			used := hostEnd.Sub(hostStart)
			backLog["end"] = fmt.Sprintf("%.4f", float64(hostEnd.UnixNano())/1e9)
			backLog["used"] = fmt.Sprintf("%.3fms", float64(used.Nanoseconds())/1e6)

		}

		//call other hosts async
		go (func(reqs []*apiHostRequest) {
			defer (func() {
				printLog(1)
			})()
			var wgOther sync.WaitGroup
			for index, apiReq := range reqs {
				if apiReq.isMaster {
					continue
				}
				wgOther.Add(1)
				go (func(index int, apiReq *apiHostRequest) {
					defer wgOther.Done()
					backLog := make(map[string]interface{})
					logData[fmt.Sprintf("host_%s_%d", apiReq.apiHost.Name, index)] = backLog
					hostStart := time.Now()
					backLog["isMaster"] = apiReq.isMaster
					backLog["start"] = fmt.Sprintf("%.4f", float64(hostStart.UnixNano())/1e9)
					backLog["isMaster"] = apiReq.isMaster
					resp, err := apiReq.transport.RoundTrip(apiReq.req)
					if err != nil {
						log.Println("fetch "+apiReq.urlNew, err)
						return
					}
					backLog["status"] = resp.StatusCode
					defer resp.Body.Close()

					hostEnd := time.Now()
					used := hostEnd.Sub(hostStart)
					backLog["end"] = fmt.Sprintf("%.4f", float64(hostEnd.UnixNano())/1e9)
					backLog["used"] = fmt.Sprintf("%.3fms", float64(used.Nanoseconds())/1e6)
				})(index, apiReq)
			}
			wgOther.Wait()

		})(reqs)
	}
}

type apiHostRequest struct {
	req       *http.Request
	urlRaw    string
	urlNew    string
	transport *http.Transport
	apiHost   *Host
	isMaster  bool
}

var reqCookieDumpLine = regexp.MustCompile(`Cookie: .+\r\n`)

func (apiServer *ApiServer) initBroadCastData(req *http.Request) *BroadCastData {
	data := NewReqBroadCastData(req)
	dumpBody := IsRequestDumpBody(req)

	dump, _ := httputil.DumpRequest(req, dumpBody)
	req_detail := string(dump)
	if apiServer.ServerConf.HiddenCookie {
		req_detail = reqCookieDumpLine.ReplaceAllStringFunc(req_detail, ReqCookieHidden)
	}

	data.SetData("req_detail", base64.StdEncoding.EncodeToString([]byte(req_detail)))
	return data
}

var resCookieDumpLine = regexp.MustCompile(`Set-Cookie: .+\r\n`)

func (apiServer *ApiServer) addBroadCastDataResponse(broadData *BroadCastData, resp *http.Response) {
	dumpBody := true
	if resp.StatusCode == http.StatusOK {
		dumpBody = IsContentTypeText(resp.Header.Get("Content-Type"))
		if dumpBody && resp.ContentLength > 0 && resp.ContentLength > 1e6 {
			dumpBody = false
		}
	}
	switch resp.StatusCode {
	case http.StatusBadRequest:
	case http.StatusNotFound:
	case http.StatusSeeOther:
	case http.StatusForbidden:
	case http.StatusGone:
	case http.StatusBadGateway:
	case http.StatusInternalServerError:
	case http.StatusServiceUnavailable:
	case http.StatusGatewayTimeout:
		dumpBody = true
		break
	default:
		break
	}

	broadData.SetData("resp_status", resp.StatusCode)
	dump, _ := httputil.DumpResponse(resp, false)

	res_detail := string(dump)
	if apiServer.ServerConf.HiddenCookie {
		res_detail = resCookieDumpLine.ReplaceAllStringFunc(res_detail, ResCookieSetHidden)
		//		res_detail = resCookieDumpLine.ReplaceAllString(res_detail, "Set-Cookie: ------hidden------\r\n")
	}
	if !dumpBody {
		res_detail += "---body skipped---"
	} else {
		bd := forgetRead(&resp.Body)
		if resp.Header.Get("Content-Encoding") == "gzip" {
			res_detail += gzipDocode(bd)
		} else {
			res_detail += bd.String()
		}
	}
	//	fmt.Println(res_detail)
	broadData.SetData("res_detail", base64.StdEncoding.EncodeToString([]byte(res_detail)))
}

func (apiServer *ApiServer) BroadcastApiReq(api *Api, data *BroadCastData) {
	apiServer.web.BroadcastApi(api, "req", data)
}
<commit_msg>fix proxy<commit_after>package proxy

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"
)

func (apiServer *ApiServer) newHandler(api *Api) func(http.ResponseWriter, *http.Request) {
	bindPath := api.Path
	log.Println(apiServer.ServerConf.Port, api.Name, "bind path [", bindPath, "]")
	return func(rw http.ResponseWriter, req *http.Request) {
		id := api.PvInc()
		uniqId := apiServer.GetUniqReqId(id)

		broadData := apiServer.initBroadCastData(req)
		broadData.Id = uniqId
		start := time.Now()
		defer func() {
			used := float64(time.Now().Sub(start).Nanoseconds()) / 1e6
			broadData.SetData("used", used)
			go apiServer.BroadcastApiReq(api, broadData)
		}()

		rw.Header().Set("Api-Proxy-Version", API_PROXY_VERSION)
		log.Println(req.URL.String())

		relPath := req.URL.Path[len(bindPath):]
		req.Header.Set("Connection", "close")

		logData := make(map[string]interface{})

		body, err := ioutil.ReadAll(req.Body)

		logData["body_len"] = len(body)

		if err != nil {
			rw.WriteHeader(http.StatusBadGateway)
			rw.Write([]byte("read body failed"))
			broadData.SetError(err.Error())
			return
		}
		//get body must by before  parse callerPref

		hosts, masterHost, cpf := api.getApiHostsByReq(req)
		broadData.SetData("master", masterHost)
		broadData.SetData("remote", cpf.GetIp())

		_uri := req.URL.Path
		if req.URL.RawQuery != "" {
			_uri += "?" + req.URL.RawQuery
		}
		mainLogStr := fmt.Sprintf("uniqid=%s port=%d remote=%s method=%s uri=%s master=%s hostsTotal=%d refer=%s", uniqId, apiServer.ServerConf.Port, req.RemoteAddr, req.Method, _uri, masterHost, len(hosts), req.Referer())

		var printLog = func(logIndex int) {
			totalUsed := fmt.Sprintf("%.3fms", float64(time.Now().Sub(start).Nanoseconds())/1e6)
			log.Println(fmt.Sprintf("logindex=%d", logIndex), mainLogStr, fmt.Sprintf("totalUsed=%s", totalUsed), logData)
		}
		defer (func() {
			printLog(0)
		})()

		rw.Header().Set("Api-Proxy-Master", masterHost)

		addrInfo := strings.Split(req.RemoteAddr, ":")

		if len(hosts) == 0 {
			logData["hostTotal"] = 0
			rw.WriteHeader(http.StatusBadGateway)
			rw.Write([]byte("no backend hosts"))
			broadData.SetError("no backend hosts")
			return
		}

		bodyLen := int64(len(body))
		reqs := make([]*apiHostRequest, 0)

		//build request
		for _, api_host := range hosts {
			isMaster := api_host.Name == masterHost
			urlNew := ""

			serverUrl := api_host.Url
			if api.HostAsProxy {
				serverUrl = "http://" + req.Host + api.Path
			}
			if strings.HasSuffix(urlNew, "/") {
				urlNew += strings.TrimLeft(relPath, "/")
			} else {
				urlNew += relPath
			}
			if req.URL.RawQuery != "" {
				urlNew += "?" + req.URL.RawQuery
			}

			rawUrl := api_host.Url + urlNew

			if isMaster {
				rw.Header().Set("Api-Proxy-Raw-Url", rawUrl)
			}

			urlNew = serverUrl + urlNew

			broadData.SetData("raw_url", rawUrl)

			reqNew, err := http.NewRequest(req.Method, urlNew, ioutil.NopCloser(bytes.NewReader(body)))
			if err != nil {
				log.Println("build req failed:", err)
				if isMaster {
					rw.WriteHeader(http.StatusBadGateway)
					rw.Write([]byte("error:" + err.Error() + "\nraw_url:" + rawUrl))
				}
				broadData.SetError(err.Error())
				return

			}
			copyHeaders(reqNew.Header, req.Header)

			//only accept gzip encode
			accept_encoding := reqNew.Header.Get("Accept-Encoding")
			if accept_encoding != "" && (In_StringSlice("gzip", reqNew.Header["Accept-Encoding"]) || strings.Contains(accept_encoding, "gzip")) {
				reqNew.Header.Set("Accept-Encoding", "gzip")
			}

			if bodyLen > 0 {
				reqNew.ContentLength = bodyLen
				reqNew.Header.Set("Content-Length", fmt.Sprintf("%d", bodyLen))
			}

			reqNew.Header.Set("HTTP_X_FORWARDED_FOR", addrInfo[0])

			transport := &http.Transport{
				Proxy: http.ProxyFromEnvironment,
				Dial: (&net.Dialer{
					Timeout:   time.Duration(api.TimeoutMs) * time.Millisecond,
					KeepAlive: 30 * time.Second,
				}).Dial,
				TLSHandshakeTimeout: 10 * time.Second,
			}
			if api.HostAsProxy {
				transport.Proxy = (func(u string) func(*http.Request) (*url.URL, error) {
					return func(req *http.Request) (*url.URL, error) {
						return url.Parse(u)
					}
				})(api_host.Url)
			}
			apiReq := &apiHostRequest{
				req:       reqNew,
				transport: transport,
				apiHost:   api_host,
				isMaster:  isMaster,
				urlNew:    urlNew,
				urlRaw:    rawUrl,
			}
			reqs = append(reqs, apiReq)
		}

		//call master at first sync
		for index, apiReq := range reqs {
			if !apiReq.isMaster {
				continue
			}
			backLog := make(map[string]interface{})
			logData[fmt.Sprintf("host_%s_%d", apiReq.apiHost.Name, index)] = backLog
			hostStart := time.Now()
			backLog["isMaster"] = apiReq.isMaster
			backLog["start"] = fmt.Sprintf("%.4f", float64(hostStart.UnixNano())/1e9)
			resp, err := apiReq.transport.RoundTrip(apiReq.req)
			if err != nil {
				log.Println("fetch "+apiReq.urlNew, err)
				rw.WriteHeader(http.StatusBadGateway)
				rw.Write([]byte("fetch_error:" + err.Error() + "\nraw_url:" + apiReq.urlRaw + "\nnew_url:" + apiReq.urlNew))
				return
			}
			defer resp.Body.Close()

			apiServer.addBroadCastDataResponse(broadData, resp)

			for k, vs := range resp.Header {
				for _, v := range vs {
					rw.Header().Add(k, v)
				}
			}
			rw.Header().Set("Connection", "close")
			rw.WriteHeader(resp.StatusCode)
			backLog["status"] = resp.StatusCode
			n, err := io.Copy(rw, resp.Body)
			if err != nil {
				log.Println(apiReq.urlNew, "io.copy:", n, err)
			}
			hostEnd := time.Now()
			used := hostEnd.Sub(hostStart)
			backLog["end"] = fmt.Sprintf("%.4f", float64(hostEnd.UnixNano())/1e9)
			backLog["used"] = fmt.Sprintf("%.3fms", float64(used.Nanoseconds())/1e6)

		}

		//call other hosts async
		go (func(reqs []*apiHostRequest) {
			defer (func() {
				printLog(1)
			})()
			var wgOther sync.WaitGroup
			for index, apiReq := range reqs {
				if apiReq.isMaster {
					continue
				}
				wgOther.Add(1)
				go (func(index int, apiReq *apiHostRequest) {
					defer wgOther.Done()
					backLog := make(map[string]interface{})
					logData[fmt.Sprintf("host_%s_%d", apiReq.apiHost.Name, index)] = backLog
					hostStart := time.Now()
					backLog["isMaster"] = apiReq.isMaster
					backLog["start"] = fmt.Sprintf("%.4f", float64(hostStart.UnixNano())/1e9)
					backLog["isMaster"] = apiReq.isMaster
					resp, err := apiReq.transport.RoundTrip(apiReq.req)
					if err != nil {
						log.Println("fetch "+apiReq.urlNew, err)
						return
					}
					backLog["status"] = resp.StatusCode
					defer resp.Body.Close()

					hostEnd := time.Now()
					used := hostEnd.Sub(hostStart)
					backLog["end"] = fmt.Sprintf("%.4f", float64(hostEnd.UnixNano())/1e9)
					backLog["used"] = fmt.Sprintf("%.3fms", float64(used.Nanoseconds())/1e6)
				})(index, apiReq)
			}
			wgOther.Wait()

		})(reqs)
	}
}

type apiHostRequest struct {
	req       *http.Request
	urlRaw    string
	urlNew    string
	transport *http.Transport
	apiHost   *Host
	isMaster  bool
}

var reqCookieDumpLine = regexp.MustCompile(`Cookie: .+\r\n`)

func (apiServer *ApiServer) initBroadCastData(req *http.Request) *BroadCastData {
	data := NewReqBroadCastData(req)
	dumpBody := IsRequestDumpBody(req)

	dump, _ := httputil.DumpRequest(req, dumpBody)
	req_detail := string(dump)
	if apiServer.ServerConf.HiddenCookie {
		req_detail = reqCookieDumpLine.ReplaceAllStringFunc(req_detail, ReqCookieHidden)
	}

	data.SetData("req_detail", base64.StdEncoding.EncodeToString([]byte(req_detail)))
	return data
}

var resCookieDumpLine = regexp.MustCompile(`Set-Cookie: .+\r\n`)

func (apiServer *ApiServer) addBroadCastDataResponse(broadData *BroadCastData, resp *http.Response) {
	dumpBody := true
	if resp.StatusCode == http.StatusOK {
		dumpBody = IsContentTypeText(resp.Header.Get("Content-Type"))
		if dumpBody && resp.ContentLength > 0 && resp.ContentLength > 1e6 {
			dumpBody = false
		}
	}
	switch resp.StatusCode {
	case http.StatusBadRequest:
	case http.StatusNotFound:
	case http.StatusSeeOther:
	case http.StatusForbidden:
	case http.StatusGone:
	case http.StatusBadGateway:
	case http.StatusInternalServerError:
	case http.StatusServiceUnavailable:
	case http.StatusGatewayTimeout:
		dumpBody = true
		break
	default:
		break
	}

	broadData.SetData("resp_status", resp.StatusCode)
	dump, _ := httputil.DumpResponse(resp, false)

	res_detail := string(dump)
	if apiServer.ServerConf.HiddenCookie {
		res_detail = resCookieDumpLine.ReplaceAllStringFunc(res_detail, ResCookieSetHidden)
		//		res_detail = resCookieDumpLine.ReplaceAllString(res_detail, "Set-Cookie: ------hidden------\r\n")
	}
	if !dumpBody {
		res_detail += "---body skipped---"
	} else {
		bd := forgetRead(&resp.Body)
		if resp.Header.Get("Content-Encoding") == "gzip" {
			res_detail += gzipDocode(bd)
		} else {
			res_detail += bd.String()
		}
	}
	//	fmt.Println(res_detail)
	broadData.SetData("res_detail", base64.StdEncoding.EncodeToString([]byte(res_detail)))
}

func (apiServer *ApiServer) BroadcastApiReq(api *Api, data *BroadCastData) {
	apiServer.web.BroadcastApi(api, "req", data)
}
<|endoftext|>


<commit_before>/*

Classical-inheritance-style service declarations.
Services can be started, then stopped, then optionally restarted.
Users can override the OnStart/OnStop methods.
By default, these methods are guaranteed to be called at most once.
A call to Reset will panic, unless OnReset is overwritten, allowing OnStart/OnStop to be called again.
Caller must ensure that Start() and Stop() are not called concurrently.
It is ok to call Stop() without calling Start() first.
Services cannot be re-started unless OnReset is overwritten to allow it.

Typical usage:

type FooService struct {
	BaseService
	// private fields
}

func NewFooService() *FooService {
	fs := &FooService{
		// init
	}
	fs.BaseService = *NewBaseService(log, "FooService", fs)
	return fs
}

func (fs *FooService) OnStart() error {
	fs.BaseService.OnStart() // Always call the overridden method.
	// initialize private fields
	// start subroutines, etc.
}

func (fs *FooService) OnStop() error {
	fs.BaseService.OnStop() // Always call the overridden method.
	// close/destroy private fields
	// stop subroutines, etc.
}

*/
package common

import (
	"sync/atomic"

	"github.com/tendermint/log15"
)

type Service interface {
	Start() (bool, error)
	OnStart() error

	Stop() bool
	OnStop()

	Reset() (bool, error)
	OnReset() error

	IsRunning() bool

	String() string
}

type BaseService struct {
	log     log15.Logger
	name    string
	started uint32 // atomic
	stopped uint32 // atomic
	Quit    chan struct{}

	// The "subclass" of BaseService
	impl Service
}

func NewBaseService(log log15.Logger, name string, impl Service) *BaseService {
	return &BaseService{
		log:  log,
		name: name,
		Quit: make(chan struct{}),
		impl: impl,
	}
}

// Implements Servce
func (bs *BaseService) Start() (bool, error) {
	if atomic.CompareAndSwapUint32(&bs.started, 0, 1) {
		if atomic.LoadUint32(&bs.stopped) == 1 {
			if bs.log != nil {
				bs.log.Warn(Fmt("Not starting %v -- already stopped", bs.name), "impl", bs.impl)
			}
			return false, nil
		} else {
			if bs.log != nil {
				bs.log.Info(Fmt("Starting %v", bs.name), "impl", bs.impl)
			}
		}
		err := bs.impl.OnStart()
		return true, err
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Not starting %v -- already started", bs.name), "impl", bs.impl)
		}
		return false, nil
	}
}

// Implements Service
// NOTE: Do not put anything in here,
// that way users don't need to call BaseService.OnStart()
func (bs *BaseService) OnStart() error { return nil }

// Implements Service
func (bs *BaseService) Stop() bool {
	if atomic.CompareAndSwapUint32(&bs.stopped, 0, 1) {
		if bs.log != nil {
			bs.log.Info(Fmt("Stopping %v", bs.name), "impl", bs.impl)
		}
		bs.impl.OnStop()
		close(bs.Quit)
		return true
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Stopping %v (ignoring: already stopped)", bs.name), "impl", bs.impl)
		}
		return false
	}
}

// Implements Service
// NOTE: Do not put anything in here,
// that way users don't need to call BaseService.OnStop()
func (bs *BaseService) OnStop() {}

// Implements Service
func (bs *BaseService) Reset() (bool, error) {
	if atomic.CompareAndSwapUint32(&bs.stopped, 1, 0) {
		// whether or not we've started, we can reset
		atomic.CompareAndSwapUint32(&bs.started, 1, 0)

		return true, bs.impl.OnReset()
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Can't reset %v. Not stopped", bs.name), "impl", bs.impl)
		}
		return false, nil
	}
	// never happens
	return false, nil
}

// Implements Service
func (bs *BaseService) OnReset() error {
	PanicSanity("The service cannot be reset")
	return nil
}

// Implements Service
func (bs *BaseService) IsRunning() bool {
	return atomic.LoadUint32(&bs.started) == 1 && atomic.LoadUint32(&bs.stopped) == 0
}

func (bs *BaseService) Wait() {
	<-bs.Quit
}

// Implements Servce
func (bs *BaseService) String() string {
	return bs.name
}

//----------------------------------------

type QuitService struct {
	BaseService
}

func NewQuitService(log log15.Logger, name string, impl Service) *QuitService {
	if log != nil {
		log.Warn("QuitService is deprecated, use BaseService instead")
	}
	return &QuitService{
		BaseService: *NewBaseService(log, name, impl),
	}
}
<commit_msg>[service] recreate Quit channel on reset<commit_after>/*

Classical-inheritance-style service declarations.
Services can be started, then stopped, then optionally restarted.
Users can override the OnStart/OnStop methods.
By default, these methods are guaranteed to be called at most once.
A call to Reset will panic, unless OnReset is overwritten, allowing OnStart/OnStop to be called again.
Caller must ensure that Start() and Stop() are not called concurrently.
It is ok to call Stop() without calling Start() first.
Services cannot be re-started unless OnReset is overwritten to allow it.

Typical usage:

type FooService struct {
	BaseService
	// private fields
}

func NewFooService() *FooService {
	fs := &FooService{
		// init
	}
	fs.BaseService = *NewBaseService(log, "FooService", fs)
	return fs
}

func (fs *FooService) OnStart() error {
	fs.BaseService.OnStart() // Always call the overridden method.
	// initialize private fields
	// start subroutines, etc.
}

func (fs *FooService) OnStop() error {
	fs.BaseService.OnStop() // Always call the overridden method.
	// close/destroy private fields
	// stop subroutines, etc.
}

*/
package common

import (
	"sync/atomic"

	"github.com/tendermint/log15"
)

type Service interface {
	Start() (bool, error)
	OnStart() error

	Stop() bool
	OnStop()

	Reset() (bool, error)
	OnReset() error

	IsRunning() bool

	String() string
}

type BaseService struct {
	log     log15.Logger
	name    string
	started uint32 // atomic
	stopped uint32 // atomic
	Quit    chan struct{}

	// The "subclass" of BaseService
	impl Service
}

func NewBaseService(log log15.Logger, name string, impl Service) *BaseService {
	return &BaseService{
		log:  log,
		name: name,
		Quit: make(chan struct{}),
		impl: impl,
	}
}

// Implements Servce
func (bs *BaseService) Start() (bool, error) {
	if atomic.CompareAndSwapUint32(&bs.started, 0, 1) {
		if atomic.LoadUint32(&bs.stopped) == 1 {
			if bs.log != nil {
				bs.log.Warn(Fmt("Not starting %v -- already stopped", bs.name), "impl", bs.impl)
			}
			return false, nil
		} else {
			if bs.log != nil {
				bs.log.Info(Fmt("Starting %v", bs.name), "impl", bs.impl)
			}
		}
		err := bs.impl.OnStart()
		return true, err
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Not starting %v -- already started", bs.name), "impl", bs.impl)
		}
		return false, nil
	}
}

// Implements Service
// NOTE: Do not put anything in here,
// that way users don't need to call BaseService.OnStart()
func (bs *BaseService) OnStart() error { return nil }

// Implements Service
func (bs *BaseService) Stop() bool {
	if atomic.CompareAndSwapUint32(&bs.stopped, 0, 1) {
		if bs.log != nil {
			bs.log.Info(Fmt("Stopping %v", bs.name), "impl", bs.impl)
		}
		bs.impl.OnStop()
		close(bs.Quit)
		return true
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Stopping %v (ignoring: already stopped)", bs.name), "impl", bs.impl)
		}
		return false
	}
}

// Implements Service
// NOTE: Do not put anything in here,
// that way users don't need to call BaseService.OnStop()
func (bs *BaseService) OnStop() {}

// Implements Service
func (bs *BaseService) Reset() (bool, error) {
	if atomic.CompareAndSwapUint32(&bs.stopped, 1, 0) {
		// whether or not we've started, we can reset
		atomic.CompareAndSwapUint32(&bs.started, 1, 0)

		bs.Quit = make(chan struct{})
		return true, bs.impl.OnReset()
	} else {
		if bs.log != nil {
			bs.log.Debug(Fmt("Can't reset %v. Not stopped", bs.name), "impl", bs.impl)
		}
		return false, nil
	}
	// never happens
	return false, nil
}

// Implements Service
func (bs *BaseService) OnReset() error {
	PanicSanity("The service cannot be reset")
	return nil
}

// Implements Service
func (bs *BaseService) IsRunning() bool {
	return atomic.LoadUint32(&bs.started) == 1 && atomic.LoadUint32(&bs.stopped) == 0
}

func (bs *BaseService) Wait() {
	<-bs.Quit
}

// Implements Servce
func (bs *BaseService) String() string {
	return bs.name
}

//----------------------------------------

type QuitService struct {
	BaseService
}

func NewQuitService(log log15.Logger, name string, impl Service) *QuitService {
	if log != nil {
		log.Warn("QuitService is deprecated, use BaseService instead")
	}
	return &QuitService{
		BaseService: *NewBaseService(log, name, impl),
	}
}
<|endoftext|>


<commit_before>package carbon

import (
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"

	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
)

type Collector struct {
	helper.Stoppable
	app            *App
	graphPrefix    string
	metricInterval time.Duration
	data           chan *points.Points
}

func NewCollector(app *App) *Collector {
	c := &Collector{
		app:            app,
		graphPrefix:    app.Config.Common.GraphPrefix,
		metricInterval: app.Config.Common.MetricInterval.Value(),
		data:           make(chan *points.Points, 4096),
	}

	c.Start()

	// collector worker
	c.Go(func(exit chan bool) {
		ticker := time.NewTicker(c.metricInterval)

		for {
			select {
			case <-exit:
				return
			case <-ticker.C:
				c.collect()
			}
		}
	})

	// sender worker
	out := app.Cache.In()

	c.Go(func(exit chan bool) {
		for {
			select {
			case <-exit:
				return
			case p := <-c.data:
				select {
				case out <- p:
				// pass
				case <-exit:
					return
				}
			}
		}
	})

	return c
}

func (c *Collector) collect() {
	app := c.app

	app.Lock()
	defer app.Unlock()

	statModule := func(module string) func(metric string, value float64) {
		return func(metric string, value float64) {
			key := fmt.Sprintf("%s.%s.%s", c.graphPrefix, module, metric)
			logrus.Infof("[stat] %s=%#v", key, value)
			select {
			case c.data <- points.NowPoint(key, value):
				// pass
			default:
				logrus.WithField("key", key).WithField("value", value).
					Warn("[stat] send queue is full. Metric dropped")
			}
		}
	}

	if app.Cache != nil {
		app.Cache.Stat(statModule("cache"))
	}
	if app.UDP != nil {
		app.UDP.Stat(statModule("udp"))
	}
	if app.TCP != nil {
		app.TCP.Stat(statModule("tcp"))
	}
	if app.Pickle != nil {
		app.Pickle.Stat(statModule("pickle"))
	}
	if app.Persister != nil {
		app.Persister.Stat(statModule("persister"))
	}
}
<commit_msg>Fix ticker leak in collector<commit_after>package carbon

import (
	"fmt"
	"time"

	"github.com/Sirupsen/logrus"

	"github.com/lomik/go-carbon/helper"
	"github.com/lomik/go-carbon/points"
)

type Collector struct {
	helper.Stoppable
	app            *App
	graphPrefix    string
	metricInterval time.Duration
	data           chan *points.Points
}

func NewCollector(app *App) *Collector {
	c := &Collector{
		app:            app,
		graphPrefix:    app.Config.Common.GraphPrefix,
		metricInterval: app.Config.Common.MetricInterval.Value(),
		data:           make(chan *points.Points, 4096),
	}

	c.Start()

	// collector worker
	c.Go(func(exit chan bool) {
		ticker := time.NewTicker(c.metricInterval)
		defer ticker.Stop()

		for {
			select {
			case <-exit:
				return
			case <-ticker.C:
				c.collect()
			}
		}
	})

	// sender worker
	out := app.Cache.In()

	c.Go(func(exit chan bool) {
		for {
			select {
			case <-exit:
				return
			case p := <-c.data:
				select {
				case out <- p:
				// pass
				case <-exit:
					return
				}
			}
		}
	})

	return c
}

func (c *Collector) collect() {
	app := c.app

	app.Lock()
	defer app.Unlock()

	statModule := func(module string) func(metric string, value float64) {
		return func(metric string, value float64) {
			key := fmt.Sprintf("%s.%s.%s", c.graphPrefix, module, metric)
			logrus.Infof("[stat] %s=%#v", key, value)
			select {
			case c.data <- points.NowPoint(key, value):
				// pass
			default:
				logrus.WithField("key", key).WithField("value", value).
					Warn("[stat] send queue is full. Metric dropped")
			}
		}
	}

	if app.Cache != nil {
		app.Cache.Stat(statModule("cache"))
	}
	if app.UDP != nil {
		app.UDP.Stat(statModule("udp"))
	}
	if app.TCP != nil {
		app.TCP.Stat(statModule("tcp"))
	}
	if app.Pickle != nil {
		app.Pickle.Stat(statModule("pickle"))
	}
	if app.Persister != nil {
		app.Persister.Stat(statModule("persister"))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"reflect"
	"sync"
	"sync/atomic"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	Registry = ServiceRegistry{
		svcs: make(map[string]*Service, 0),
	}
)

type Service struct {
	sync.Mutex
	Name          string
	Addr          string
	Backends      []*Backend
	Balance       string
	CheckInterval int
	Fall          int
	Rise          int
	ClientTimeout time.Duration
	ServerTimeout time.Duration
	DialTimeout   time.Duration
	Sent          int64
	Rcvd          int64
	Errors        int64

	// Next returns the backends in priority order.
	next func() []*Backend

	// the last backend we used and the number of times we used it
	lastBackend int
	lastCount   int

	// Each Service owns it's own netowrk listener
	listener net.Listener
}

// Stats returned about a service
type ServiceStat struct {
	Name          string        `json:"name"`
	Addr          string        `json:"address"`
	Backends      []BackendStat `json:"backends"`
	Balance       string        `json:"balance"`
	CheckInterval int           `json:"check_interval"`
	Fall          int           `json:"fall"`
	Rise          int           `json:"rise"`
	ClientTimeout int           `json:"client_timeout"`
	ServerTimeout int           `json:"server_timeout"`
	DialTimeout   int           `json:"connect_timeout"`
	Sent          int64         `json:"sent"`
	Rcvd          int64         `json:"received"`
	Errors        int64         `json:"errors"`
	Conns         int64         `json:"connections"`
	Active        int64         `json:"active"`
}

// Subset of service fields needed for configuration.
type ServiceConfig struct {
	Name          string          `json:"name"`
	Addr          string          `json:"address"`
	Backends      []BackendConfig `json:"backends"`
	Balance       string          `json:"balance"`
	CheckInterval int             `json:"check_interval"`
	Fall          int             `json:"fall"`
	Rise          int             `json:"rise"`
	ClientTimeout int             `json:"client_timeout"`
	ServerTimeout int             `json:"server_timeout"`
	DialTimeout   int             `json:"connect_timeout"`
}

// Compare a service's settings, ignoring individual backends.
func (s ServiceConfig) Equal(other ServiceConfig) bool {
	// just remove the backends and compare the rest
	s.Backends = nil
	other.Backends = nil

	// FIXME: Normalize default in one place!

	if s.Balance != other.Balance {
		if s.Balance == "" && other.Balance == "RR" {
			other.Balance = ""
		} else if s.Balance == "RR" && other.Balance == "" {
			other.Balance = "RR"
		}
	}

	if s.CheckInterval == 0 {
		s.CheckInterval = 2000
	}
	if s.Rise == 0 {
		s.Rise = 2
	}
	if s.Fall == 0 {
		s.Fall = 2
	}

	// We handle backends separately
	s.Backends = nil
	other.Backends = nil

	return reflect.DeepEqual(s, other)
}

// Create a Service from a config struct
func NewService(cfg ServiceConfig) *Service {
	s := &Service{
		Name:          cfg.Name,
		Addr:          cfg.Addr,
		CheckInterval: cfg.CheckInterval,
		Fall:          cfg.Fall,
		Rise:          cfg.Rise,
		ClientTimeout: time.Duration(cfg.ClientTimeout) * time.Millisecond,
		ServerTimeout: time.Duration(cfg.ServerTimeout) * time.Millisecond,
		DialTimeout:   time.Duration(cfg.DialTimeout) * time.Millisecond,
	}

	if s.CheckInterval == 0 {
		s.CheckInterval = 2000
	}
	if s.Rise == 0 {
		s.Rise = 2
	}
	if s.Fall == 0 {
		s.Fall = 2
	}

	for _, b := range cfg.Backends {
		s.add(NewBackend(b))
	}

	switch cfg.Balance {
	case "RR", "":
		s.next = s.roundRobin
	case "LC":
		s.next = s.leastConn
	default:
		log.Printf("invalid balancing algorithm '%s'", cfg.Balance)
	}

	return s
}

func (s *Service) Stats() ServiceStat {
	s.Lock()
	defer s.Unlock()

	stats := ServiceStat{
		Name:          s.Name,
		Addr:          s.Addr,
		Balance:       s.Balance,
		CheckInterval: s.CheckInterval,
		Fall:          s.Fall,
		Rise:          s.Rise,
		ClientTimeout: int(s.ClientTimeout / time.Millisecond),
		ServerTimeout: int(s.ServerTimeout / time.Millisecond),
		DialTimeout:   int(s.DialTimeout / time.Millisecond),
	}

	for _, b := range s.Backends {
		stats.Backends = append(stats.Backends, b.Stats())
		stats.Sent += b.Sent
		stats.Rcvd += b.Rcvd
		stats.Errors += b.Errors
		stats.Conns += b.Conns
		stats.Active += b.Active
	}

	return stats
}

func (s *Service) Config() ServiceConfig {
	s.Lock()
	defer s.Unlock()

	config := ServiceConfig{
		Name:          s.Name,
		Addr:          s.Addr,
		Balance:       s.Balance,
		CheckInterval: s.CheckInterval,
		Fall:          s.Fall,
		Rise:          s.Rise,
		ClientTimeout: int(s.ClientTimeout / time.Millisecond),
		ServerTimeout: int(s.ServerTimeout / time.Millisecond),
		DialTimeout:   int(s.DialTimeout / time.Millisecond),
	}
	for _, b := range s.Backends {
		config.Backends = append(config.Backends, b.Config())
	}

	return config
}

func (s *Service) String() string {
	return string(marshal(s.Config()))
}

func (s *Service) get(name string) *Backend {
	s.Lock()
	defer s.Unlock()

	for _, b := range s.Backends {
		if b.Name == name {
			return b
		}
	}
	return nil
}

// Add or replace a Backend in this service
func (s *Service) add(backend *Backend) {
	s.Lock()
	defer s.Unlock()

	backend.up = true
	backend.rwTimeout = s.ServerTimeout
	backend.dialTimeout = s.DialTimeout
	backend.checkInterval = time.Duration(s.CheckInterval) * time.Millisecond

	// replace an existing backend if we have it.
	for i, b := range s.Backends {
		if b.Name == backend.Name {
			b.Stop()
			s.Backends[i] = backend
			backend.Start()
			return
		}
	}

	s.Backends = append(s.Backends, backend)

	backend.Start()
}

// Remove a Backend by name
func (s *Service) remove(name string) bool {
	s.Lock()
	defer s.Unlock()

	for i, b := range s.Backends {
		if b.Name == name {
			last := len(s.Backends) - 1
			deleted := b
			s.Backends[i], s.Backends[last] = s.Backends[last], nil
			s.Backends = s.Backends[:last]
			deleted.Stop()
			return true
		}
	}
	return false
}

// Fill out and verify service
func (s *Service) start() (err error) {
	s.Lock()
	defer s.Unlock()

	s.listener, err = newTimeoutListener(s.Addr, s.ClientTimeout)
	if err != nil {
		return err
	}

	if s.Backends == nil {
		s.Backends = make([]*Backend, 0)
	}

	s.run()
	return nil
}

// Start the Service's Accept loop
func (s *Service) run() {
	go func() {
		for {
			conn, err := s.listener.Accept()
			if err != nil {
				if err, ok := err.(*net.OpError); ok && err.Temporary() {
					log.Println("warning:", err)
					continue
				}
				// we must be getting shut down
				return
			}

			go s.connect(conn)
		}
	}()
}

func (s *Service) connect(cliConn net.Conn) {
	backends := s.next()

	// Try the first backend given, but if that fails, cycle through them all
	// to make a best effort to connect the client.
	for _, b := range backends {
		srvConn, err := net.DialTimeout("tcp", b.Addr, b.dialTimeout)
		if err != nil {
			log.Printf("Error connecting to Backend %s/%s: %s", s.Name, b.Name, err)
			atomic.AddInt64(&b.Errors, 1)
			continue
		}

		b.Proxy(srvConn, cliConn)
		return
	}

	log.Println("Error: no Backend for", s.Name)
	cliConn.Close()
}

// Stop the Service's Accept loop by closing the Listener,
// and stop all backends for this service.
func (s *Service) stop() {
	s.Lock()
	defer s.Unlock()

	log.Debug("Stopping Service", s.Name)
	for _, backend := range s.Backends {
		backend.Stop()
	}

	// the service may have been bad, and the listener failed
	if s.listener == nil {
		return
	}

	err := s.listener.Close()
	if err != nil {
		log.Println(err)
	}
}

// A net.Listener that provides a read/write timeout
type timeoutListener struct {
	net.Listener
	rwTimeout time.Duration
}

func newTimeoutListener(addr string, timeout time.Duration) (net.Listener, error) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	tl := &timeoutListener{
		Listener:  l,
		rwTimeout: timeout,
	}
	return tl, nil
}
func (l *timeoutListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	c, ok := conn.(*net.TCPConn)
	if ok {
		tc := &timeoutConn{
			TCPConn:   c,
			rwTimeout: l.rwTimeout,
		}
		return tc, nil
	}
	return conn, nil
}
<commit_msg>Add virtual_hosts to service config<commit_after>package main

import (
	"net"
	"reflect"
	"sync"
	"sync/atomic"
	"time"

	"github.com/litl/galaxy/log"
)

var (
	Registry = ServiceRegistry{
		svcs: make(map[string]*Service, 0),
	}
)

type Service struct {
	sync.Mutex
	Name          string
	Addr          string
	VirtualHosts  []string
	Backends      []*Backend
	Balance       string
	CheckInterval int
	Fall          int
	Rise          int
	ClientTimeout time.Duration
	ServerTimeout time.Duration
	DialTimeout   time.Duration
	Sent          int64
	Rcvd          int64
	Errors        int64

	// Next returns the backends in priority order.
	next func() []*Backend

	// the last backend we used and the number of times we used it
	lastBackend int
	lastCount   int

	// Each Service owns it's own netowrk listener
	listener net.Listener
}

// Stats returned about a service
type ServiceStat struct {
	Name          string        `json:"name"`
	Addr          string        `json:"address"`
	VirtualHosts  []string      `json:"virtual_hosts"`
	Backends      []BackendStat `json:"backends"`
	Balance       string        `json:"balance"`
	CheckInterval int           `json:"check_interval"`
	Fall          int           `json:"fall"`
	Rise          int           `json:"rise"`
	ClientTimeout int           `json:"client_timeout"`
	ServerTimeout int           `json:"server_timeout"`
	DialTimeout   int           `json:"connect_timeout"`
	Sent          int64         `json:"sent"`
	Rcvd          int64         `json:"received"`
	Errors        int64         `json:"errors"`
	Conns         int64         `json:"connections"`
	Active        int64         `json:"active"`
}

// Subset of service fields needed for configuration.
type ServiceConfig struct {
	Name          string          `json:"name"`
	Addr          string          `json:"address"`
	VirtualHosts  []string        `json:"virtual_hosts"`
	Backends      []BackendConfig `json:"backends"`
	Balance       string          `json:"balance"`
	CheckInterval int             `json:"check_interval"`
	Fall          int             `json:"fall"`
	Rise          int             `json:"rise"`
	ClientTimeout int             `json:"client_timeout"`
	ServerTimeout int             `json:"server_timeout"`
	DialTimeout   int             `json:"connect_timeout"`
}

// Compare a service's settings, ignoring individual backends.
func (s ServiceConfig) Equal(other ServiceConfig) bool {
	// just remove the backends and compare the rest
	s.Backends = nil
	other.Backends = nil

	// FIXME: Normalize default in one place!

	if s.Balance != other.Balance {
		if s.Balance == "" && other.Balance == "RR" {
			other.Balance = ""
		} else if s.Balance == "RR" && other.Balance == "" {
			other.Balance = "RR"
		}
	}

	if s.CheckInterval == 0 {
		s.CheckInterval = 2000
	}
	if s.Rise == 0 {
		s.Rise = 2
	}
	if s.Fall == 0 {
		s.Fall = 2
	}

	// We handle backends separately
	s.Backends = nil
	other.Backends = nil

	return reflect.DeepEqual(s, other)
}

// Create a Service from a config struct
func NewService(cfg ServiceConfig) *Service {
	s := &Service{
		Name:          cfg.Name,
		Addr:          cfg.Addr,
		CheckInterval: cfg.CheckInterval,
		Fall:          cfg.Fall,
		Rise:          cfg.Rise,
		VirtualHosts:  cfg.VirtualHosts,
		ClientTimeout: time.Duration(cfg.ClientTimeout) * time.Millisecond,
		ServerTimeout: time.Duration(cfg.ServerTimeout) * time.Millisecond,
		DialTimeout:   time.Duration(cfg.DialTimeout) * time.Millisecond,
	}

	if s.CheckInterval == 0 {
		s.CheckInterval = 2000
	}
	if s.Rise == 0 {
		s.Rise = 2
	}
	if s.Fall == 0 {
		s.Fall = 2
	}

	for _, b := range cfg.Backends {
		s.add(NewBackend(b))
	}

	switch cfg.Balance {
	case "RR", "":
		s.next = s.roundRobin
	case "LC":
		s.next = s.leastConn
	default:
		log.Printf("invalid balancing algorithm '%s'", cfg.Balance)
	}

	return s
}

func (s *Service) Stats() ServiceStat {
	s.Lock()
	defer s.Unlock()

	stats := ServiceStat{
		Name:          s.Name,
		Addr:          s.Addr,
		VirtualHosts:  s.VirtualHosts,
		Balance:       s.Balance,
		CheckInterval: s.CheckInterval,
		Fall:          s.Fall,
		Rise:          s.Rise,
		ClientTimeout: int(s.ClientTimeout / time.Millisecond),
		ServerTimeout: int(s.ServerTimeout / time.Millisecond),
		DialTimeout:   int(s.DialTimeout / time.Millisecond),
	}

	for _, b := range s.Backends {
		stats.Backends = append(stats.Backends, b.Stats())
		stats.Sent += b.Sent
		stats.Rcvd += b.Rcvd
		stats.Errors += b.Errors
		stats.Conns += b.Conns
		stats.Active += b.Active
	}

	return stats
}

func (s *Service) Config() ServiceConfig {
	s.Lock()
	defer s.Unlock()

	config := ServiceConfig{
		Name:          s.Name,
		Addr:          s.Addr,
		VirtualHosts:  s.VirtualHosts,
		Balance:       s.Balance,
		CheckInterval: s.CheckInterval,
		Fall:          s.Fall,
		Rise:          s.Rise,
		ClientTimeout: int(s.ClientTimeout / time.Millisecond),
		ServerTimeout: int(s.ServerTimeout / time.Millisecond),
		DialTimeout:   int(s.DialTimeout / time.Millisecond),
	}
	for _, b := range s.Backends {
		config.Backends = append(config.Backends, b.Config())
	}

	return config
}

func (s *Service) String() string {
	return string(marshal(s.Config()))
}

func (s *Service) get(name string) *Backend {
	s.Lock()
	defer s.Unlock()

	for _, b := range s.Backends {
		if b.Name == name {
			return b
		}
	}
	return nil
}

// Add or replace a Backend in this service
func (s *Service) add(backend *Backend) {
	s.Lock()
	defer s.Unlock()

	backend.up = true
	backend.rwTimeout = s.ServerTimeout
	backend.dialTimeout = s.DialTimeout
	backend.checkInterval = time.Duration(s.CheckInterval) * time.Millisecond

	// replace an existing backend if we have it.
	for i, b := range s.Backends {
		if b.Name == backend.Name {
			b.Stop()
			s.Backends[i] = backend
			backend.Start()
			return
		}
	}

	s.Backends = append(s.Backends, backend)

	backend.Start()
}

// Remove a Backend by name
func (s *Service) remove(name string) bool {
	s.Lock()
	defer s.Unlock()

	for i, b := range s.Backends {
		if b.Name == name {
			last := len(s.Backends) - 1
			deleted := b
			s.Backends[i], s.Backends[last] = s.Backends[last], nil
			s.Backends = s.Backends[:last]
			deleted.Stop()
			return true
		}
	}
	return false
}

// Fill out and verify service
func (s *Service) start() (err error) {
	s.Lock()
	defer s.Unlock()

	s.listener, err = newTimeoutListener(s.Addr, s.ClientTimeout)
	if err != nil {
		return err
	}

	if s.Backends == nil {
		s.Backends = make([]*Backend, 0)
	}

	s.run()
	return nil
}

// Start the Service's Accept loop
func (s *Service) run() {
	go func() {
		for {
			conn, err := s.listener.Accept()
			if err != nil {
				if err, ok := err.(*net.OpError); ok && err.Temporary() {
					log.Println("warning:", err)
					continue
				}
				// we must be getting shut down
				return
			}

			go s.connect(conn)
		}
	}()
}

func (s *Service) connect(cliConn net.Conn) {
	backends := s.next()

	// Try the first backend given, but if that fails, cycle through them all
	// to make a best effort to connect the client.
	for _, b := range backends {
		srvConn, err := net.DialTimeout("tcp", b.Addr, b.dialTimeout)
		if err != nil {
			log.Printf("Error connecting to Backend %s/%s: %s", s.Name, b.Name, err)
			atomic.AddInt64(&b.Errors, 1)
			continue
		}

		b.Proxy(srvConn, cliConn)
		return
	}

	log.Println("Error: no Backend for", s.Name)
	cliConn.Close()
}

// Stop the Service's Accept loop by closing the Listener,
// and stop all backends for this service.
func (s *Service) stop() {
	s.Lock()
	defer s.Unlock()

	log.Debug("Stopping Service", s.Name)
	for _, backend := range s.Backends {
		backend.Stop()
	}

	// the service may have been bad, and the listener failed
	if s.listener == nil {
		return
	}

	err := s.listener.Close()
	if err != nil {
		log.Println(err)
	}
}

// A net.Listener that provides a read/write timeout
type timeoutListener struct {
	net.Listener
	rwTimeout time.Duration
}

func newTimeoutListener(addr string, timeout time.Duration) (net.Listener, error) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	tl := &timeoutListener{
		Listener:  l,
		rwTimeout: timeout,
	}
	return tl, nil
}
func (l *timeoutListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	c, ok := conn.(*net.TCPConn)
	if ok {
		tc := &timeoutConn{
			TCPConn:   c,
			rwTimeout: l.rwTimeout,
		}
		return tc, nil
	}
	return conn, nil
}
<|endoftext|>


<commit_before>package main

import (
	"math/rand"
)

const (
	// ValidChars chars to consists of random string
	ValidChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
	// DefaultLen default random string length
	DefaultLen = 4
)

// URLServiceImpl main service of shortpath
type URLServiceImpl struct {
	dbs    *DBService
	length int
}

// NewURLServiceImpl return new UrlServiceImpl instance
func NewURLServiceImpl(dbs *DBService) *URLServiceImpl {
	usi := &URLServiceImpl{dbs: dbs, length: DefaultLen}
	return usi
}

// Shorten generate a short, non-repeat path maps to longurl
func (u *URLServiceImpl) Shorten(longurl string, length int) string {
	if length == 0 {
		u.SetLen(DefaultLen)
	} else {
		u.SetLen(length)
	}

	shortpath := generateRandomStr(u.GetLen())
	for u.Seek(shortpath) {
		shortpath = generateRandomStr(u.GetLen())
	}
	u.Put(longurl, shortpath)

	return shortpath
}

// Put save url and shortpath to db
func (u *URLServiceImpl) Put(longurl, shortpath string) {
	u.dbs.InsertShortpath(longurl, shortpath)
}

// GetShortpath get shortpath and url by shortpath from db
func (u *URLServiceImpl) GetShortpath(shortpath string) string {
	return u.dbs.QueryUrlRecord(shortpath)
}

// SetLen set current random string length
func (u *URLServiceImpl) SetLen(length int) {
	u.length = length
}

// GetLen return current random string length
func (u *URLServiceImpl) GetLen() int {
	return u.length
}

// Seek check if path has existed in db
func (u *URLServiceImpl) Seek(path string) bool {
	return u.dbs.CheckPath(path)
}

// generateRandomStr
func generateRandomStr(length int) string {
	shortURL := make([]byte, length)
	seed := make([]byte, length)
	charLength := len(ValidChars)

	_, err := rand.Read(seed)
	if err != nil {
		logq.Fatal(err)
	}

	for k, v := range seed {
		i := int(v)
		shortURL[k] = ValidChars[i%charLength]
	}

	return string(shortURL)
}
<commit_msg>logq -> logrus<commit_after>package main

import (
	"math/rand"

	"github.com/Sirupsen/logrus"
)

const (
	// ValidChars chars to consists of random string
	ValidChars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
	// DefaultLen default random string length
	DefaultLen = 4
)

// URLServiceImpl main service of shortpath
type URLServiceImpl struct {
	dbs    *DBService
	length int
}

// NewURLServiceImpl return new UrlServiceImpl instance
func NewURLServiceImpl(dbs *DBService) *URLServiceImpl {
	usi := &URLServiceImpl{dbs: dbs, length: DefaultLen}
	return usi
}

// Shorten generate a short, non-repeat path maps to longurl
func (u *URLServiceImpl) Shorten(longurl string, length int) string {
	if length == 0 {
		u.SetLen(DefaultLen)
	} else {
		u.SetLen(length)
	}

	shortpath := generateRandomStr(u.GetLen())
	for u.Seek(shortpath) {
		shortpath = generateRandomStr(u.GetLen())
	}
	u.Put(longurl, shortpath)

	return shortpath
}

// Put save url and shortpath to db
func (u *URLServiceImpl) Put(longurl, shortpath string) {
	u.dbs.InsertShortpath(longurl, shortpath)
}

// GetShortpath get shortpath and url by shortpath from db
func (u *URLServiceImpl) GetShortpath(shortpath string) string {
	return u.dbs.QueryUrlRecord(shortpath)
}

// SetLen set current random string length
func (u *URLServiceImpl) SetLen(length int) {
	u.length = length
}

// GetLen return current random string length
func (u *URLServiceImpl) GetLen() int {
	return u.length
}

// Seek check if path has existed in db
func (u *URLServiceImpl) Seek(path string) bool {
	return u.dbs.CheckPath(path)
}

// generateRandomStr
func generateRandomStr(length int) string {
	shortURL := make([]byte, length)
	seed := make([]byte, length)
	charLength := len(ValidChars)

	_, err := rand.Read(seed)
	if err != nil {
		logrus.Fatal(err)
	}

	for k, v := range seed {
		i := int(v)
		shortURL[k] = ValidChars[i%charLength]
	}

	return string(shortURL)
}
<|endoftext|>


<commit_before><commit_msg>Refactor.<commit_after><|endoftext|>


<commit_before>package siesta

import (
	"errors"
	"log"
	"net/http"
	"regexp"
	"strings"
)

var services map[string]*Service = make(map[string]*Service)

type Service struct {
	baseURI string

	pre  []contextHandler
	post []contextHandler

	handlers map[*regexp.Regexp]contextHandler
}

func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:  strings.TrimRight(baseURI, "/"),
		handlers: make(map[*regexp.Regexp]contextHandler),
	}
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)

	if m == nil {
		panic(errors.New("unsupported middleware type"))
	}

	return append(chain, m)
}

func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	quit := false
	for _, m := range s.pre {
		if quit {
			return
		}

		m(c, w, r, func() {
			quit = true
		})
	}

	if !quit {

		var handler contextHandler

		for re, h := range s.handlers {
			req := r.Method + " " + r.URL.Path
			if matches := re.FindStringSubmatch(req); len(matches) > 0 {
				r.ParseForm()
				for i, match := range matches {
					if i > 0 {
						param := re.SubexpNames()[i]
						r.Form.Set(param, match)
					}
				}

				handler = h
				break
			}
		}

		if handler == nil {
			http.NotFoundHandler().ServeHTTP(w, r)
		} else {
			handler(c, w, r, func() {
				quit = true
			})
		}
	}

	for _, m := range s.post {
		if quit {
			return
		}

		m(c, w, r, func() {
			quit = true
		})
	}
}

func (s *Service) Route(verb, pattern, usage string, f interface{}) {
	handler := toContextHandler(f)

	expr := strings.TrimRight(strings.TrimLeft(pattern, "/"), "/")
	expr = strings.Replace(expr, "<", "(?P<", -1)
	expr = strings.Replace(expr, ">", `>[\w-_.]+`, -1)

	end := "?$"
	if len(expr) == 0 {
		end = "/?$"
	}

	if len(expr) > 0 {
		expr += "/"
		if s.baseURI != "/" {
			expr = "/" + expr
		}
	}

	expr = "^" + verb + " " + s.baseURI + expr + end
	re := regexp.MustCompile(expr)

	if _, ok := s.handlers[re]; ok {
		panic("already registered handler for " + verb + " " + pattern)
	} else {
		log.Println("Handling", expr)
	}

	s.handlers[re] = handler
}

func (s *Service) Register() {
	http.Handle(s.baseURI, s)
	http.Handle(s.baseURI+"/", s)
}
<commit_msg>Fix regexp construction at Route()<commit_after>package siesta

import (
	"errors"
	"log"
	"net/http"
	"regexp"
	"strings"
)

var services map[string]*Service = make(map[string]*Service)

type Service struct {
	baseURI string

	pre  []contextHandler
	post []contextHandler

	handlers map[*regexp.Regexp]contextHandler
}

func NewService(baseURI string) *Service {
	if services[baseURI] != nil {
		panic("service already registered")
	}

	return &Service{
		baseURI:  strings.TrimRight(baseURI, "/"),
		handlers: make(map[*regexp.Regexp]contextHandler),
	}
}

func addToChain(f interface{}, chain []contextHandler) []contextHandler {
	m := toContextHandler(f)

	if m == nil {
		panic(errors.New("unsupported middleware type"))
	}

	return append(chain, m)
}

func (s *Service) AddPre(f interface{}) {
	s.pre = addToChain(f, s.pre)
}

func (s *Service) AddPost(f interface{}) {
	s.post = addToChain(f, s.post)
}

// Service satisfies the http.Handler interface.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.ServeHTTPInContext(NewSiestaContext(), w, r)
}

func (s *Service) ServeHTTPInContext(c Context, w http.ResponseWriter, r *http.Request) {
	quit := false
	for _, m := range s.pre {
		if quit {
			return
		}

		m(c, w, r, func() {
			quit = true
		})
	}

	if !quit {

		var handler contextHandler

		for re, h := range s.handlers {
			req := r.Method + " " + r.URL.Path
			if matches := re.FindStringSubmatch(req); len(matches) > 0 {
				r.ParseForm()
				for i, match := range matches {
					if i > 0 {
						param := re.SubexpNames()[i]
						r.Form.Set(param, match)
					}
				}

				handler = h
				break
			}
		}

		if handler == nil {
			http.NotFoundHandler().ServeHTTP(w, r)
		} else {
			handler(c, w, r, func() {
				quit = true
			})
		}
	}

	for _, m := range s.post {
		if quit {
			return
		}

		m(c, w, r, func() {
			quit = true
		})
	}
}

func (s *Service) Route(verb, pattern, usage string, f interface{}) {
	handler := toContextHandler(f)

	expr := strings.TrimRight(strings.TrimLeft(pattern, "/"), "/")
	expr = strings.Replace(expr, "<", "(?P<", -1)
	expr = strings.Replace(expr, ">", `>[\w-_.]+)`, -1)

	end := "?$"
	if len(expr) == 0 {
		end = "/?$"
	}

	if len(expr) > 0 {
		expr += "/"
		if s.baseURI != "/" {
			expr = "/" + expr
		}
	}

	expr = "^" + verb + " " + s.baseURI + expr + end
	re := regexp.MustCompile(expr)

	if _, ok := s.handlers[re]; ok {
		panic("already registered handler for " + verb + " " + pattern)
	} else {
		log.Println("Handling", expr)
	}

	s.handlers[re] = handler
}

func (s *Service) Register() {
	http.Handle(s.baseURI, s)
	http.Handle(s.baseURI+"/", s)
}
<|endoftext|>


<commit_before>package possum

import (
	"net/http"

	"github.com/mikespook/possum/session"
)

const sessionCookieName = "possum"

// SessionFacotryFunc is the default factory function of session.
var SessionFacotryFunc = session.NewFactory(session.CookieStorage(sessionCookieName, nil))

// Session extracts data from the request and returns session instance.
// It uses SessionFacotryFunc to initialise session if no session has been set yet.
func Session(w http.ResponseWriter, req *http.Request) (sn *session.Session, err error) {
	sn, err = SessionFacotryFunc(w, req)
	if err != nil {
		return nil, err
	}
	return sn, nil
}
<commit_msg>add defer func for session<commit_after>package possum

import (
	"net/http"

	"github.com/mikespook/possum/session"
)

const sessionCookieName = "possum"

// SessionFacotryFunc is the default factory function of session.
var SessionFacotryFunc = session.NewFactory(session.CookieStorage(sessionCookieName, nil))

// Session extracts data from the request and returns session instance.
// It uses SessionFacotryFunc to initialise session if no session has been set yet.
func Session(w http.ResponseWriter, req *http.Request) (sn *session.Session, deferFunc func(), err error) {
	sn, err = SessionFacotryFunc(w, req)
	if err != nil {
		return nil, nil, err
	}
	return sn, func() {
		if err := sn.Flush(); err != nil {
			panic(err)
		}
	}, nil
}
<|endoftext|>