<commit_before>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mackerelio/golib/logging"
	"github.com/mackerelio/mackerel-agent/config"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata
// AzureVM: https://docs.microsoft.com/azure/virtual-machines/virtual-machines-instancemetadataservice-overview

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL, azureVMBaseURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
}

var timeout = 100 * time.Millisecond

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator(conf *config.Config) *CloudGenerator {
	// if CloudPlatform is specified, return corresponding one
	switch conf.CloudPlatform {
	case config.CloudPlatformNone:
		return nil
	case config.CloudPlatformEC2:
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	case config.CloudPlatformGCE:
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	case config.CloudPlatformAzureVM:
		return &CloudGenerator{&AzureVMGenerator{azureVMBaseURL}}
	}

	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}
	if isAzure() {
		return &CloudGenerator{&AzureVMGenerator{azureVMBaseURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

// Note: May want to check without using the API.
func isAzure() bool {
	cl := httpCli()
	// '/vmId` is probably Azure VM specific URL
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return false
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == 200
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}

// AzureVMGenerator meta generator for Azure VM
type AzureVMGenerator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *AzureVMGenerator) Generate() (interface{}, error) {
	metadataComputeKeys := map[string]string{
		"location":  "location",
		"offer":     "imageReferenceOffer",
		"osType":    "osSystemType",
		"publisher": "imageReferencePublisher",
		"sku":       "imageReferenceSku",
		"vmId":      "vmID",
		"vmSize":    "virtualMachineSizeType",
	}

	ipAddressKeys := map[string]string{
		"privateIpAddress": "privateIpAddress",
		"publicIpAddress":  "publicIpAddress",
	}

	metadata := make(map[string]string)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/compute/", metadataComputeKeys)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/network/interface/0/ipv4/ipAddress/0/", ipAddressKeys)

	results := make(map[string]interface{})
	results["provider"] = "AzureVM"
	results["metadata"] = metadata

	return results, nil
}

func retrieveAzureVMMetadata(metadataMap map[string]string, baseURL string, urlSuffix string, keys map[string]string) map[string]string {
	cl := httpCli()

	for key, value := range keys {
		req, err := http.NewRequest("GET", baseURL+urlSuffix+key+"?api-version=2017-04-02&format=text", nil)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return metadataMap
		}

		req.Header.Set("Metadata", "true")

		resp, err := cl.Do(req)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return metadataMap
		}
		defer resp.Body.Close()

		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadataMap[value] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}
	return metadataMap
}

// SuggestCustomIdentifier suggests the identifier of the Azure VM instance
func (g *AzureVMGenerator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request vmId. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting vmId cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".virtual_machine.azure.microsoft.com", nil
}
<commit_msg>extend timeout for retrieving cloud metadata<commit_after>package spec

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mackerelio/golib/logging"
	"github.com/mackerelio/mackerel-agent/config"
)

// This Generator collects metadata about cloud instances.
// Currently EC2 and GCE are supported.
// EC2: http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/AESDG-chapter-instancedata.html
// GCE: https://developers.google.com/compute/docs/metadata
// AzureVM: https://docs.microsoft.com/azure/virtual-machines/virtual-machines-instancemetadataservice-overview

// CloudGenerator definition
type CloudGenerator struct {
	CloudMetaGenerator
}

// CloudMetaGenerator interface of metadata generator for each cloud platform
type CloudMetaGenerator interface {
	Generate() (interface{}, error)
	SuggestCustomIdentifier() (string, error)
}

// Key is a root key for the generator.
func (g *CloudGenerator) Key() string {
	return "cloud"
}

var cloudLogger = logging.GetLogger("spec.cloud")

var ec2BaseURL, gceMetaURL, azureVMBaseURL *url.URL

func init() {
	ec2BaseURL, _ = url.Parse("http://169.254.169.254/latest/meta-data")
	gceMetaURL, _ = url.Parse("http://metadata.google.internal/computeMetadata/v1/?recursive=true")
	azureVMBaseURL, _ = url.Parse("http://169.254.169.254/metadata/instance")
}

var timeout = 3 * time.Second

// SuggestCloudGenerator returns suitable CloudGenerator
func SuggestCloudGenerator(conf *config.Config) *CloudGenerator {
	// if CloudPlatform is specified, return corresponding one
	switch conf.CloudPlatform {
	case config.CloudPlatformNone:
		return nil
	case config.CloudPlatformEC2:
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	case config.CloudPlatformGCE:
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	case config.CloudPlatformAzureVM:
		return &CloudGenerator{&AzureVMGenerator{azureVMBaseURL}}
	}

	if isEC2() {
		return &CloudGenerator{&EC2Generator{ec2BaseURL}}
	}
	if isGCE() {
		return &CloudGenerator{&GCEGenerator{gceMetaURL}}
	}
	if isAzure() {
		return &CloudGenerator{&AzureVMGenerator{azureVMBaseURL}}
	}

	return nil
}

func httpCli() *http.Client {
	return &http.Client{
		Timeout: timeout,
		Transport: &http.Transport{
			// don't use HTTP_PROXY when requesting cloud instance metadata APIs
			Proxy: nil,
		},
	}
}

func isGCE() bool {
	_, err := requestGCEMeta()
	return err == nil
}

// Note: May want to check without using the API.
func isAzure() bool {
	cl := httpCli()
	// '/vmId` is probably Azure VM specific URL
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return false
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == 200
}

func requestGCEMeta() ([]byte, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", gceMetaURL.String(), nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Metadata-Flavor", "Google")

	resp, err := cl.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("failed to request gce meta. response code: %d", resp.StatusCode)
	}
	return ioutil.ReadAll(resp.Body)
}

// EC2Generator meta generator for EC2
type EC2Generator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *EC2Generator) Generate() (interface{}, error) {
	cl := httpCli()

	metadataKeys := []string{
		"instance-id",
		"instance-type",
		"placement/availability-zone",
		"security-groups",
		"ami-id",
		"hostname",
		"local-hostname",
		"public-hostname",
		"local-ipv4",
		"public-ipv4",
		"reservation-id",
	}

	metadata := make(map[string]string)

	for _, key := range metadataKeys {
		resp, err := cl.Get(g.baseURL.String() + "/" + key)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on EC2. Error while reading '%s'", key)
			return nil, nil
		}
		defer resp.Body.Close()
		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadata[key] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}

	results := make(map[string]interface{})
	results["provider"] = "ec2"
	results["metadata"] = metadata

	return results, nil
}

// SuggestCustomIdentifier suggests the identifier of the EC2 instance
func (g *EC2Generator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	key := "instance-id"
	resp, err := cl.Get(g.baseURL.String() + "/" + key)
	if err != nil {
		return "", fmt.Errorf("error while retrieving instance-id")
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request instance-id. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting instance-id cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".ec2.amazonaws.com", nil
}

// GCEGenerator generate for GCE
type GCEGenerator struct {
	metaURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *GCEGenerator) Generate() (interface{}, error) {
	bytes, err := requestGCEMeta()
	if err != nil {
		return nil, err
	}
	var data gceMeta
	json.Unmarshal(bytes, &data)
	return data.toGeneratorResults(), nil
}

type gceInstance struct {
	Zone         string
	InstanceType string `json:"machineType"`
	Hostname     string
	InstanceID   uint64 `json:"id"`
}

type gceProject struct {
	ProjectID        string
	NumericProjectID uint64
}

type gceMeta struct {
	Instance *gceInstance
	Project  *gceProject
}

func (g gceMeta) toGeneratorMeta() map[string]string {
	meta := make(map[string]string)

	lastS := func(s string) string {
		ss := strings.Split(s, "/")
		return ss[len(ss)-1]
	}

	if ins := g.Instance; ins != nil {
		meta["hostname"] = ins.Hostname
		meta["instance-id"] = fmt.Sprint(ins.InstanceID)
		meta["instance-type"] = lastS(ins.InstanceType)
		meta["zone"] = lastS(ins.Zone)
	}

	if proj := g.Project; proj != nil {
		meta["projectId"] = proj.ProjectID
	}

	return meta
}

func (g gceMeta) toGeneratorResults() interface{} {
	results := make(map[string]interface{})
	results["provider"] = "gce"
	results["metadata"] = g.toGeneratorMeta()

	return results
}

// SuggestCustomIdentifier for GCE is not implemented yet
func (g *GCEGenerator) SuggestCustomIdentifier() (string, error) {
	return "", nil
}

// AzureVMGenerator meta generator for Azure VM
type AzureVMGenerator struct {
	baseURL *url.URL
}

// Generate collects metadata from cloud platform.
func (g *AzureVMGenerator) Generate() (interface{}, error) {
	metadataComputeKeys := map[string]string{
		"location":  "location",
		"offer":     "imageReferenceOffer",
		"osType":    "osSystemType",
		"publisher": "imageReferencePublisher",
		"sku":       "imageReferenceSku",
		"vmId":      "vmID",
		"vmSize":    "virtualMachineSizeType",
	}

	ipAddressKeys := map[string]string{
		"privateIpAddress": "privateIpAddress",
		"publicIpAddress":  "publicIpAddress",
	}

	metadata := make(map[string]string)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/compute/", metadataComputeKeys)
	metadata = retrieveAzureVMMetadata(metadata, g.baseURL.String(), "/network/interface/0/ipv4/ipAddress/0/", ipAddressKeys)

	results := make(map[string]interface{})
	results["provider"] = "AzureVM"
	results["metadata"] = metadata

	return results, nil
}

func retrieveAzureVMMetadata(metadataMap map[string]string, baseURL string, urlSuffix string, keys map[string]string) map[string]string {
	cl := httpCli()

	for key, value := range keys {
		req, err := http.NewRequest("GET", baseURL+urlSuffix+key+"?api-version=2017-04-02&format=text", nil)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return metadataMap
		}

		req.Header.Set("Metadata", "true")

		resp, err := cl.Do(req)
		if err != nil {
			cloudLogger.Debugf("This host may not be running on Azure VM. Error while reading '%s'", key)
			return metadataMap
		}
		defer resp.Body.Close()

		if resp.StatusCode == 200 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				cloudLogger.Errorf("Results of requesting metadata cannot be read: '%s'", err)
				break
			}
			metadataMap[value] = string(body)
			cloudLogger.Debugf("results %s:%s", key, string(body))
		} else {
			cloudLogger.Debugf("Status code of the result of requesting metadata '%s' is '%d'", key, resp.StatusCode)
		}
	}
	return metadataMap
}

// SuggestCustomIdentifier suggests the identifier of the Azure VM instance
func (g *AzureVMGenerator) SuggestCustomIdentifier() (string, error) {
	cl := httpCli()
	req, err := http.NewRequest("GET", azureVMBaseURL.String()+"/compute/vmId?api-version=2017-04-02&format=text", nil)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	req.Header.Set("Metadata", "true")

	resp, err := cl.Do(req)
	if err != nil {
		return "", fmt.Errorf("error while retrieving vmId")
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to request vmId. response code: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("results of requesting vmId cannot be read: '%s'", err)
	}
	instanceID := string(body)
	if instanceID == "" {
		return "", fmt.Errorf("invalid instance id")
	}
	return instanceID + ".virtual_machine.azure.microsoft.com", nil
}
<|endoftext|>


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

package tls

import (
	"bytes"
	"crypto"
	"crypto/rsa"
	"crypto/subtle"
	"crypto/x509"
	"errors"
	"io"
	"strconv"
)

func (c *Conn) clientHandshake() error {
	finishedHash := newFinishedHash(versionTLS10)

	if c.config == nil {
		c.config = defaultConfig()
	}

	hello := &clientHelloMsg{
		vers:               maxVersion,
		cipherSuites:       c.config.cipherSuites(),
		compressionMethods: []uint8{compressionNone},
		random:             make([]byte, 32),
		ocspStapling:       true,
		serverName:         c.config.ServerName,
		supportedCurves:    []uint16{curveP256, curveP384, curveP521},
		supportedPoints:    []uint8{pointFormatUncompressed},
		nextProtoNeg:       len(c.config.NextProtos) > 0,
	}

	t := uint32(c.config.time().Unix())
	hello.random[0] = byte(t >> 24)
	hello.random[1] = byte(t >> 16)
	hello.random[2] = byte(t >> 8)
	hello.random[3] = byte(t)
	_, err := io.ReadFull(c.config.rand(), hello.random[4:])
	if err != nil {
		c.sendAlert(alertInternalError)
		return errors.New("short read from Rand")
	}

	finishedHash.Write(hello.marshal())
	c.writeRecord(recordTypeHandshake, hello.marshal())

	msg, err := c.readHandshake()
	if err != nil {
		return err
	}
	serverHello, ok := msg.(*serverHelloMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(serverHello.marshal())

	vers, ok := mutualVersion(serverHello.vers)
	if !ok || vers < versionTLS10 {
		// TLS 1.0 is the minimum version supported as a client.
		return c.sendAlert(alertProtocolVersion)
	}
	c.vers = vers
	c.haveVers = true

	if serverHello.compressionMethod != compressionNone {
		return c.sendAlert(alertUnexpectedMessage)
	}

	if !hello.nextProtoNeg && serverHello.nextProtoNeg {
		c.sendAlert(alertHandshakeFailure)
		return errors.New("server advertised unrequested NPN")
	}

	suite := mutualCipherSuite(c.config.cipherSuites(), serverHello.cipherSuite)
	if suite == nil {
		return c.sendAlert(alertHandshakeFailure)
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}
	certMsg, ok := msg.(*certificateMsg)
	if !ok || len(certMsg.certificates) == 0 {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(certMsg.marshal())

	certs := make([]*x509.Certificate, len(certMsg.certificates))
	for i, asn1Data := range certMsg.certificates {
		cert, err := x509.ParseCertificate(asn1Data)
		if err != nil {
			c.sendAlert(alertBadCertificate)
			return errors.New("failed to parse certificate from server: " + err.Error())
		}
		certs[i] = cert
	}

	if !c.config.InsecureSkipVerify {
		opts := x509.VerifyOptions{
			Roots:         c.config.RootCAs,
			CurrentTime:   c.config.time(),
			DNSName:       c.config.ServerName,
			Intermediates: x509.NewCertPool(),
		}

		for i, cert := range certs {
			if i == 0 {
				continue
			}
			opts.Intermediates.AddCert(cert)
		}
		c.verifiedChains, err = certs[0].Verify(opts)
		if err != nil {
			c.sendAlert(alertBadCertificate)
			return err
		}
	}

	if _, ok := certs[0].PublicKey.(*rsa.PublicKey); !ok {
		return c.sendAlert(alertUnsupportedCertificate)
	}

	c.peerCertificates = certs

	if serverHello.ocspStapling {
		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
		cs, ok := msg.(*certificateStatusMsg)
		if !ok {
			return c.sendAlert(alertUnexpectedMessage)
		}
		finishedHash.Write(cs.marshal())

		if cs.statusType == statusTypeOCSP {
			c.ocspResponse = cs.response
		}
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}

	keyAgreement := suite.ka()

	skx, ok := msg.(*serverKeyExchangeMsg)
	if ok {
		finishedHash.Write(skx.marshal())
		err = keyAgreement.processServerKeyExchange(c.config, hello, serverHello, certs[0], skx)
		if err != nil {
			c.sendAlert(alertUnexpectedMessage)
			return err
		}

		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
	}

	var certToSend *Certificate
	certReq, ok := msg.(*certificateRequestMsg)
	if ok {
		// RFC 4346 on the certificateAuthorities field:
		// A list of the distinguished names of acceptable certificate
		// authorities. These distinguished names may specify a desired
		// distinguished name for a root CA or for a subordinate CA;
		// thus, this message can be used to describe both known roots
		// and a desired authorization space. If the
		// certificate_authorities list is empty then the client MAY
		// send any certificate of the appropriate
		// ClientCertificateType, unless there is some external
		// arrangement to the contrary.

		finishedHash.Write(certReq.marshal())

		// For now, we only know how to sign challenges with RSA
		rsaAvail := false
		for _, certType := range certReq.certificateTypes {
			if certType == certTypeRSASign {
				rsaAvail = true
				break
			}
		}

		// We need to search our list of client certs for one
		// where SignatureAlgorithm is RSA and the Issuer is in
		// certReq.certificateAuthorities
	findCert:
		for i, cert := range c.config.Certificates {
			if !rsaAvail {
				continue
			}

			leaf := cert.Leaf
			if leaf == nil {
				if leaf, err = x509.ParseCertificate(cert.Certificate[0]); err != nil {
					c.sendAlert(alertInternalError)
					return errors.New("tls: failed to parse client certificate #" + strconv.Itoa(i) + ": " + err.Error())
				}
			}

			if leaf.PublicKeyAlgorithm != x509.RSA {
				continue
			}

			if len(certReq.certificateAuthorities) == 0 {
				// they gave us an empty list, so just take the
				// first RSA cert from c.config.Certificates
				certToSend = &cert
				break
			}

			for _, ca := range certReq.certificateAuthorities {
				if bytes.Equal(leaf.RawIssuer, ca) {
					certToSend = &cert
					break findCert
				}
			}
		}

		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
	}

	shd, ok := msg.(*serverHelloDoneMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(shd.marshal())

	if certToSend != nil {
		certMsg = new(certificateMsg)
		certMsg.certificates = certToSend.Certificate
		finishedHash.Write(certMsg.marshal())
		c.writeRecord(recordTypeHandshake, certMsg.marshal())
	}

	preMasterSecret, ckx, err := keyAgreement.generateClientKeyExchange(c.config, hello, certs[0])
	if err != nil {
		c.sendAlert(alertInternalError)
		return err
	}
	if ckx != nil {
		finishedHash.Write(ckx.marshal())
		c.writeRecord(recordTypeHandshake, ckx.marshal())
	}

	if certToSend != nil {
		certVerify := new(certificateVerifyMsg)
		digest := make([]byte, 0, 36)
		digest = finishedHash.serverMD5.Sum(digest)
		digest = finishedHash.serverSHA1.Sum(digest)
		signed, err := rsa.SignPKCS1v15(c.config.rand(), c.config.Certificates[0].PrivateKey.(*rsa.PrivateKey), crypto.MD5SHA1, digest)
		if err != nil {
			return c.sendAlert(alertInternalError)
		}
		certVerify.signature = signed

		finishedHash.Write(certVerify.marshal())
		c.writeRecord(recordTypeHandshake, certVerify.marshal())
	}

	masterSecret, clientMAC, serverMAC, clientKey, serverKey, clientIV, serverIV :=
		keysFromPreMasterSecret(c.vers, preMasterSecret, hello.random, serverHello.random, suite.macLen, suite.keyLen, suite.ivLen)

	clientCipher := suite.cipher(clientKey, clientIV, false /* not for reading */)
	clientHash := suite.mac(c.vers, clientMAC)
	c.out.prepareCipherSpec(c.vers, clientCipher, clientHash)
	c.writeRecord(recordTypeChangeCipherSpec, []byte{1})

	if serverHello.nextProtoNeg {
		nextProto := new(nextProtoMsg)
		proto, fallback := mutualProtocol(c.config.NextProtos, serverHello.nextProtos)
		nextProto.proto = proto
		c.clientProtocol = proto
		c.clientProtocolFallback = fallback

		finishedHash.Write(nextProto.marshal())
		c.writeRecord(recordTypeHandshake, nextProto.marshal())
	}

	finished := new(finishedMsg)
	finished.verifyData = finishedHash.clientSum(masterSecret)
	finishedHash.Write(finished.marshal())
	c.writeRecord(recordTypeHandshake, finished.marshal())

	serverCipher := suite.cipher(serverKey, serverIV, true /* for reading */)
	serverHash := suite.mac(c.vers, serverMAC)
	c.in.prepareCipherSpec(c.vers, serverCipher, serverHash)
	c.readRecord(recordTypeChangeCipherSpec)
	if c.err != nil {
		return c.err
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}
	serverFinished, ok := msg.(*finishedMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}

	verify := finishedHash.serverSum(masterSecret)
	if len(verify) != len(serverFinished.verifyData) ||
		subtle.ConstantTimeCompare(verify, serverFinished.verifyData) != 1 {
		return c.sendAlert(alertHandshakeFailure)
	}

	c.handshakeComplete = true
	c.cipherSuite = suite.id
	return nil
}

// mutualProtocol finds the mutual Next Protocol Negotiation protocol given the
// set of client and server supported protocols. The set of client supported
// protocols must not be empty. It returns the resulting protocol and flag
// indicating if the fallback case was reached.
func mutualProtocol(clientProtos, serverProtos []string) (string, bool) {
	for _, s := range serverProtos {
		for _, c := range clientProtos {
			if s == c {
				return s, false
			}
		}
	}

	return clientProtos[0], true
}
<commit_msg>crypto/tls: always send a Certificate message if one was requested.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tls

import (
	"bytes"
	"crypto"
	"crypto/rsa"
	"crypto/subtle"
	"crypto/x509"
	"errors"
	"io"
	"strconv"
)

func (c *Conn) clientHandshake() error {
	finishedHash := newFinishedHash(versionTLS10)

	if c.config == nil {
		c.config = defaultConfig()
	}

	hello := &clientHelloMsg{
		vers:               maxVersion,
		cipherSuites:       c.config.cipherSuites(),
		compressionMethods: []uint8{compressionNone},
		random:             make([]byte, 32),
		ocspStapling:       true,
		serverName:         c.config.ServerName,
		supportedCurves:    []uint16{curveP256, curveP384, curveP521},
		supportedPoints:    []uint8{pointFormatUncompressed},
		nextProtoNeg:       len(c.config.NextProtos) > 0,
	}

	t := uint32(c.config.time().Unix())
	hello.random[0] = byte(t >> 24)
	hello.random[1] = byte(t >> 16)
	hello.random[2] = byte(t >> 8)
	hello.random[3] = byte(t)
	_, err := io.ReadFull(c.config.rand(), hello.random[4:])
	if err != nil {
		c.sendAlert(alertInternalError)
		return errors.New("short read from Rand")
	}

	finishedHash.Write(hello.marshal())
	c.writeRecord(recordTypeHandshake, hello.marshal())

	msg, err := c.readHandshake()
	if err != nil {
		return err
	}
	serverHello, ok := msg.(*serverHelloMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(serverHello.marshal())

	vers, ok := mutualVersion(serverHello.vers)
	if !ok || vers < versionTLS10 {
		// TLS 1.0 is the minimum version supported as a client.
		return c.sendAlert(alertProtocolVersion)
	}
	c.vers = vers
	c.haveVers = true

	if serverHello.compressionMethod != compressionNone {
		return c.sendAlert(alertUnexpectedMessage)
	}

	if !hello.nextProtoNeg && serverHello.nextProtoNeg {
		c.sendAlert(alertHandshakeFailure)
		return errors.New("server advertised unrequested NPN")
	}

	suite := mutualCipherSuite(c.config.cipherSuites(), serverHello.cipherSuite)
	if suite == nil {
		return c.sendAlert(alertHandshakeFailure)
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}
	certMsg, ok := msg.(*certificateMsg)
	if !ok || len(certMsg.certificates) == 0 {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(certMsg.marshal())

	certs := make([]*x509.Certificate, len(certMsg.certificates))
	for i, asn1Data := range certMsg.certificates {
		cert, err := x509.ParseCertificate(asn1Data)
		if err != nil {
			c.sendAlert(alertBadCertificate)
			return errors.New("failed to parse certificate from server: " + err.Error())
		}
		certs[i] = cert
	}

	if !c.config.InsecureSkipVerify {
		opts := x509.VerifyOptions{
			Roots:         c.config.RootCAs,
			CurrentTime:   c.config.time(),
			DNSName:       c.config.ServerName,
			Intermediates: x509.NewCertPool(),
		}

		for i, cert := range certs {
			if i == 0 {
				continue
			}
			opts.Intermediates.AddCert(cert)
		}
		c.verifiedChains, err = certs[0].Verify(opts)
		if err != nil {
			c.sendAlert(alertBadCertificate)
			return err
		}
	}

	if _, ok := certs[0].PublicKey.(*rsa.PublicKey); !ok {
		return c.sendAlert(alertUnsupportedCertificate)
	}

	c.peerCertificates = certs

	if serverHello.ocspStapling {
		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
		cs, ok := msg.(*certificateStatusMsg)
		if !ok {
			return c.sendAlert(alertUnexpectedMessage)
		}
		finishedHash.Write(cs.marshal())

		if cs.statusType == statusTypeOCSP {
			c.ocspResponse = cs.response
		}
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}

	keyAgreement := suite.ka()

	skx, ok := msg.(*serverKeyExchangeMsg)
	if ok {
		finishedHash.Write(skx.marshal())
		err = keyAgreement.processServerKeyExchange(c.config, hello, serverHello, certs[0], skx)
		if err != nil {
			c.sendAlert(alertUnexpectedMessage)
			return err
		}

		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
	}

	var certToSend *Certificate
	var certRequested bool
	certReq, ok := msg.(*certificateRequestMsg)
	if ok {
		certRequested = true

		// RFC 4346 on the certificateAuthorities field:
		// A list of the distinguished names of acceptable certificate
		// authorities. These distinguished names may specify a desired
		// distinguished name for a root CA or for a subordinate CA;
		// thus, this message can be used to describe both known roots
		// and a desired authorization space. If the
		// certificate_authorities list is empty then the client MAY
		// send any certificate of the appropriate
		// ClientCertificateType, unless there is some external
		// arrangement to the contrary.

		finishedHash.Write(certReq.marshal())

		// For now, we only know how to sign challenges with RSA
		rsaAvail := false
		for _, certType := range certReq.certificateTypes {
			if certType == certTypeRSASign {
				rsaAvail = true
				break
			}
		}

		// We need to search our list of client certs for one
		// where SignatureAlgorithm is RSA and the Issuer is in
		// certReq.certificateAuthorities
	findCert:
		for i, cert := range c.config.Certificates {
			if !rsaAvail {
				continue
			}

			leaf := cert.Leaf
			if leaf == nil {
				if leaf, err = x509.ParseCertificate(cert.Certificate[0]); err != nil {
					c.sendAlert(alertInternalError)
					return errors.New("tls: failed to parse client certificate #" + strconv.Itoa(i) + ": " + err.Error())
				}
			}

			if leaf.PublicKeyAlgorithm != x509.RSA {
				continue
			}

			if len(certReq.certificateAuthorities) == 0 {
				// they gave us an empty list, so just take the
				// first RSA cert from c.config.Certificates
				certToSend = &cert
				break
			}

			for _, ca := range certReq.certificateAuthorities {
				if bytes.Equal(leaf.RawIssuer, ca) {
					certToSend = &cert
					break findCert
				}
			}
		}

		msg, err = c.readHandshake()
		if err != nil {
			return err
		}
	}

	shd, ok := msg.(*serverHelloDoneMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}
	finishedHash.Write(shd.marshal())

	// If the server requested a certificate then we have to send a
	// Certificate message, even if it's empty because we don't have a
	// certificate to send.
	if certRequested {
		certMsg = new(certificateMsg)
		if certToSend != nil {
			certMsg.certificates = certToSend.Certificate
		}
		finishedHash.Write(certMsg.marshal())
		c.writeRecord(recordTypeHandshake, certMsg.marshal())
	}

	preMasterSecret, ckx, err := keyAgreement.generateClientKeyExchange(c.config, hello, certs[0])
	if err != nil {
		c.sendAlert(alertInternalError)
		return err
	}
	if ckx != nil {
		finishedHash.Write(ckx.marshal())
		c.writeRecord(recordTypeHandshake, ckx.marshal())
	}

	if certToSend != nil {
		certVerify := new(certificateVerifyMsg)
		digest := make([]byte, 0, 36)
		digest = finishedHash.serverMD5.Sum(digest)
		digest = finishedHash.serverSHA1.Sum(digest)
		signed, err := rsa.SignPKCS1v15(c.config.rand(), c.config.Certificates[0].PrivateKey.(*rsa.PrivateKey), crypto.MD5SHA1, digest)
		if err != nil {
			return c.sendAlert(alertInternalError)
		}
		certVerify.signature = signed

		finishedHash.Write(certVerify.marshal())
		c.writeRecord(recordTypeHandshake, certVerify.marshal())
	}

	masterSecret, clientMAC, serverMAC, clientKey, serverKey, clientIV, serverIV :=
		keysFromPreMasterSecret(c.vers, preMasterSecret, hello.random, serverHello.random, suite.macLen, suite.keyLen, suite.ivLen)

	clientCipher := suite.cipher(clientKey, clientIV, false /* not for reading */)
	clientHash := suite.mac(c.vers, clientMAC)
	c.out.prepareCipherSpec(c.vers, clientCipher, clientHash)
	c.writeRecord(recordTypeChangeCipherSpec, []byte{1})

	if serverHello.nextProtoNeg {
		nextProto := new(nextProtoMsg)
		proto, fallback := mutualProtocol(c.config.NextProtos, serverHello.nextProtos)
		nextProto.proto = proto
		c.clientProtocol = proto
		c.clientProtocolFallback = fallback

		finishedHash.Write(nextProto.marshal())
		c.writeRecord(recordTypeHandshake, nextProto.marshal())
	}

	finished := new(finishedMsg)
	finished.verifyData = finishedHash.clientSum(masterSecret)
	finishedHash.Write(finished.marshal())
	c.writeRecord(recordTypeHandshake, finished.marshal())

	serverCipher := suite.cipher(serverKey, serverIV, true /* for reading */)
	serverHash := suite.mac(c.vers, serverMAC)
	c.in.prepareCipherSpec(c.vers, serverCipher, serverHash)
	c.readRecord(recordTypeChangeCipherSpec)
	if c.err != nil {
		return c.err
	}

	msg, err = c.readHandshake()
	if err != nil {
		return err
	}
	serverFinished, ok := msg.(*finishedMsg)
	if !ok {
		return c.sendAlert(alertUnexpectedMessage)
	}

	verify := finishedHash.serverSum(masterSecret)
	if len(verify) != len(serverFinished.verifyData) ||
		subtle.ConstantTimeCompare(verify, serverFinished.verifyData) != 1 {
		return c.sendAlert(alertHandshakeFailure)
	}

	c.handshakeComplete = true
	c.cipherSuite = suite.id
	return nil
}

// mutualProtocol finds the mutual Next Protocol Negotiation protocol given the
// set of client and server supported protocols. The set of client supported
// protocols must not be empty. It returns the resulting protocol and flag
// indicating if the fallback case was reached.
func mutualProtocol(clientProtos, serverProtos []string) (string, bool) {
	for _, s := range serverProtos {
		for _, c := range clientProtos {
			if s == c {
				return s, false
			}
		}
	}

	return clientProtos[0], true
}
<|endoftext|>


<commit_before>package gocql

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	flagCluster      = flag.String("cluster", "127.0.0.1", "a comma-separated list of host:port tuples")
	flagProto        = flag.Int("proto", 2, "protcol version")
	flagCQL          = flag.String("cql", "3.0.0", "CQL version")
	flagRF           = flag.Int("rf", 1, "replication factor for test keyspace")
	clusterSize      = flag.Int("clusterSize", 1, "the expected size of the cluster")
	flagRetry        = flag.Int("retries", 5, "number of times to retry queries")
	flagAutoWait     = flag.Duration("autowait", 1000*time.Millisecond, "time to wait for autodiscovery to fill the hosts poll")
	flagRunSslTest   = flag.Bool("runssl", false, "Set to true to run ssl test")
	flagRunAuthTest  = flag.Bool("runauth", false, "Set to true to run authentication test")
	flagCompressTest = flag.String("compressor", "", "compressor to use")
	flagTimeout      = flag.Duration("gocql.timeout", 5*time.Second, "sets the connection `timeout` for all operations")
	clusterHosts     []string
)

func init() {
	flag.Parse()
	clusterHosts = strings.Split(*flagCluster, ",")
	log.SetFlags(log.Lshortfile | log.LstdFlags)
}

func addSslOptions(cluster *ClusterConfig) *ClusterConfig {
	if *flagRunSslTest {
		cluster.SslOpts = &SslOptions{
			CertPath:               "testdata/pki/gocql.crt",
			KeyPath:                "testdata/pki/gocql.key",
			CaPath:                 "testdata/pki/ca.crt",
			EnableHostVerification: false,
		}
	}
	return cluster
}

var initOnce sync.Once

func createTable(s *Session, table string) error {
	if err := s.control.query(table).Close(); err != nil {
		return err
	}

	return nil
}

func createCluster() *ClusterConfig {
	cluster := NewCluster(clusterHosts...)
	cluster.ProtoVersion = *flagProto
	cluster.CQLVersion = *flagCQL
	cluster.Timeout = *flagTimeout
	cluster.Consistency = Quorum
	cluster.MaxWaitSchemaAgreement = 2 * time.Minute // travis might be slow
	if *flagRetry > 0 {
		cluster.RetryPolicy = &SimpleRetryPolicy{NumRetries: *flagRetry}
	}

	switch *flagCompressTest {
	case "snappy":
		cluster.Compressor = &SnappyCompressor{}
	case "":
	default:
		panic("invalid compressor: " + *flagCompressTest)
	}

	cluster = addSslOptions(cluster)
	return cluster
}

func createKeyspace(tb testing.TB, cluster *ClusterConfig, keyspace string) {
	c := *cluster
	c.Keyspace = "system"
	c.Timeout = 20 * time.Second
	session, err := c.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}
	defer session.Close()
	defer log.Println("closing keyspace session")

	err = session.control.query(`DROP KEYSPACE IF EXISTS ` + keyspace).Close()
	if err != nil {
		tb.Fatal(err)
	}

	err = session.control.query(fmt.Sprintf(`CREATE KEYSPACE %s
	WITH replication = {
		'class' : 'SimpleStrategy',
		'replication_factor' : %d
	}`, keyspace, *flagRF)).Close()

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

func createSessionFromCluster(cluster *ClusterConfig, tb testing.TB) *Session {
	// Drop and re-create the keyspace once. Different tests should use their own
	// individual tables, but can assume that the table does not exist before.
	initOnce.Do(func() {
		createKeyspace(tb, cluster, "gocql_test")
	})

	cluster.Keyspace = "gocql_test"
	session, err := cluster.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}

	return session
}

func createSession(tb testing.TB) *Session {
	cluster := createCluster()
	return createSessionFromCluster(cluster, tb)
}
<commit_msg>be more sure about schema consistency in tests<commit_after>package gocql

import (
	"flag"
	"fmt"
	"log"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	flagCluster      = flag.String("cluster", "127.0.0.1", "a comma-separated list of host:port tuples")
	flagProto        = flag.Int("proto", 2, "protcol version")
	flagCQL          = flag.String("cql", "3.0.0", "CQL version")
	flagRF           = flag.Int("rf", 1, "replication factor for test keyspace")
	clusterSize      = flag.Int("clusterSize", 1, "the expected size of the cluster")
	flagRetry        = flag.Int("retries", 5, "number of times to retry queries")
	flagAutoWait     = flag.Duration("autowait", 1000*time.Millisecond, "time to wait for autodiscovery to fill the hosts poll")
	flagRunSslTest   = flag.Bool("runssl", false, "Set to true to run ssl test")
	flagRunAuthTest  = flag.Bool("runauth", false, "Set to true to run authentication test")
	flagCompressTest = flag.String("compressor", "", "compressor to use")
	flagTimeout      = flag.Duration("gocql.timeout", 5*time.Second, "sets the connection `timeout` for all operations")
	clusterHosts     []string
)

func init() {
	flag.Parse()
	clusterHosts = strings.Split(*flagCluster, ",")
	log.SetFlags(log.Lshortfile | log.LstdFlags)
}

func addSslOptions(cluster *ClusterConfig) *ClusterConfig {
	if *flagRunSslTest {
		cluster.SslOpts = &SslOptions{
			CertPath:               "testdata/pki/gocql.crt",
			KeyPath:                "testdata/pki/gocql.key",
			CaPath:                 "testdata/pki/ca.crt",
			EnableHostVerification: false,
		}
	}
	return cluster
}

var initOnce sync.Once

func createTable(s *Session, table string) error {
	if err := s.control.query(table).Close(); err != nil {
		return err
	}

	if err := s.control.awaitSchemaAgreement(); err != nil {
		return err
	}

	return nil
}

func createCluster() *ClusterConfig {
	cluster := NewCluster(clusterHosts...)
	cluster.ProtoVersion = *flagProto
	cluster.CQLVersion = *flagCQL
	cluster.Timeout = *flagTimeout
	cluster.Consistency = Quorum
	cluster.MaxWaitSchemaAgreement = 2 * time.Minute // travis might be slow
	if *flagRetry > 0 {
		cluster.RetryPolicy = &SimpleRetryPolicy{NumRetries: *flagRetry}
	}

	switch *flagCompressTest {
	case "snappy":
		cluster.Compressor = &SnappyCompressor{}
	case "":
	default:
		panic("invalid compressor: " + *flagCompressTest)
	}

	cluster = addSslOptions(cluster)
	return cluster
}

func createKeyspace(tb testing.TB, cluster *ClusterConfig, keyspace string) {
	c := *cluster
	c.Keyspace = "system"
	c.Timeout = 20 * time.Second
	session, err := c.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}
	defer session.Close()
	defer log.Println("closing keyspace session")

	err = session.control.query(`DROP KEYSPACE IF EXISTS ` + keyspace).Close()
	if err != nil {
		tb.Fatal(err)
	}

	err = session.control.query(fmt.Sprintf(`CREATE KEYSPACE %s
	WITH replication = {
		'class' : 'SimpleStrategy',
		'replication_factor' : %d
	}`, keyspace, *flagRF)).Close()

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

	// lets just be sure
	if err := session.control.awaitSchemaAgreement(); err != nil {
		tb.Fatal(err)
	}
}

func createSessionFromCluster(cluster *ClusterConfig, tb testing.TB) *Session {
	// Drop and re-create the keyspace once. Different tests should use their own
	// individual tables, but can assume that the table does not exist before.
	initOnce.Do(func() {
		createKeyspace(tb, cluster, "gocql_test")
	})

	cluster.Keyspace = "gocql_test"
	session, err := cluster.CreateSession()
	if err != nil {
		tb.Fatal("createSession:", err)
	}

	return session
}

func createSession(tb testing.TB) *Session {
	cluster := createCluster()
	return createSessionFromCluster(cluster, tb)
}
<|endoftext|>


<commit_before>package atlas

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	url := "https://foo.example.com/"

	n := getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	url = "https://foo.example.com/?page=42"
	n = getPageNum(url)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	url = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(url)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	url = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(url)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	url = ""
	n = getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err, "should be in error")

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	var jsonErr = `error:{status: 501, code: 500, detail: "test"}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be in error")

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err, "should be in error")
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	new := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(new))
	assert.EqualValues(t, map[string]string{"key": "foo"}, new)
}

func TestPrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestPrepareRequest_2(t *testing.T) {
	c, err := NewClient()
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(apiEndpoint + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}
<commit_msg>Add test for mergeGlobalOptions(), rename other test to match signature.<commit_after>package atlas

import (
	"bytes"
	"fmt"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const testURL = "http://localhost:10000"

func TestGetPageNum(t *testing.T) {
	url := "https://foo.example.com/"

	n := getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
	url = "https://foo.example.com/?page=42"
	n = getPageNum(url)
	if n != "42" {
		t.Errorf("n=%s should be 42", n)
	}
	url = "https://foo.example.com/?country=fr&page=43"
	n = getPageNum(url)
	if n != "43" {
		t.Errorf("n=%s should be 43", n)
	}
	url = "https://foo.example.com/?country=fr&page=666&bar=1"
	n = getPageNum(url)
	if n != "666" {
		t.Errorf("n=%s should be 666", n)
	}

	url = ""
	n = getPageNum(url)
	if n != "" {
		t.Errorf("n=%s should be ''", n)
	}
}

func TestClienthandleAPIResponsese(t *testing.T) {
	var (
		r http.Response
		b bytes.Buffer
	)

	client, err := NewClient()
	err = client.handleAPIResponsese(nil)
	assert.Error(t, err, "should be in error")

	r = http.Response{StatusCode: 0}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	r = http.Response{StatusCode: 200}
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be no error")

	var jsonErr = `error:{status: 501, code: 500, detail: "test"}`

	fmt.Fprintf(&b, "%v", jsonErr)
	r.StatusCode = 300
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.NoError(t, err, "should be in error")

	r.StatusCode = 500
	r.Body = ioutil.NopCloser(&b)
	err = client.handleAPIResponsese(&r)
	assert.Error(t, err, "should be in error")
}

func TestAddQueryParameters(t *testing.T) {
	p := AddQueryParameters("", map[string]string{})
	assert.Equal(t, "", p)
}

func TestAddQueryParameters_1(t *testing.T) {
	p := AddQueryParameters("", map[string]string{"": ""})
	assert.Equal(t, "?=", p)
}

func TestAddQueryParameters_2(t *testing.T) {
	p := AddQueryParameters("foo", map[string]string{"bar": "baz"})
	assert.Equal(t, "foo?bar=baz", p)
}

func TestClient_AddAPIKey(t *testing.T) {
	c, err := NewClient(Config{APIKey: "foo"})
	require.NoError(t, err)
	assert.NotNil(t, c)
	assert.NotEmpty(t, c)

	opts := map[string]string{}

	new := c.addAPIKey(opts)
	assert.NotEmpty(t, c.config.APIKey)
	assert.Equal(t, 1, len(new))
	assert.EqualValues(t, map[string]string{"key": "foo"}, new)
}

func TestClient_PrepareRequest(t *testing.T) {
	c, err := NewClient(Config{endpoint: testURL})
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(testURL + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_PrepareRequest_2(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{}
	req := c.prepareRequest("GET", "foo", opts)

	assert.NotNil(t, req)
	assert.IsType(t, (*http.Request)(nil), req)

	res, _ := url.Parse(apiEndpoint + "/foo")
	assert.Equal(t, "GET", req.Method)
	assert.EqualValues(t, res, req.URL)
}

func TestClient_MergeGlobalOptions(t *testing.T) {
	c, err := NewClient(TesCfg)
	require.NoError(t, err)

	opts := map[string]string{"foo": "bar"}
	c.opts = map[string]string{"baz": "xyz"}
	res := map[string]string{"foo": "bar", "baz": "xyz"}

	c.mergeGlobalOptions(opts)
	assert.EqualValues(t, res, opts)
}
<|endoftext|>


<commit_before>// Echo kite is an example to kite that just echoes the argument back.
// Useful for debugging or testing stuff.

package main

import "kite"

func main() {
	options := &kite.Options{
		Kitename:    "echo",
		Version:     "0.0.1",
		Environment: "development",
		Region:      "localhost",
	}

	k := kite.New(options)

	k.HandleFunc("echo", func(r *kite.Request) (interface{}, error) {
		return r.Args.One().MustString(), nil
	})

	k.Run()
}
<commit_msg>make echo port not random, easier to restart and debug<commit_after>// Echo kite is an example to kite that just echoes the argument back.
// Useful for debugging or testing stuff.

package main

import "kite"

func main() {
	options := &kite.Options{
		Kitename:    "echo",
		Version:     "0.0.1",
		Environment: "development",
		Region:      "localhost",
		Port:        "7777",
	}

	k := kite.New(options)

	k.HandleFunc("echo", func(r *kite.Request) (interface{}, error) {
		return r.Args.One().MustString(), nil
	})

	k.Run()
}
<|endoftext|>


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

package eval

import (
	"container/vector";
	"eval";
)

/*
 * Virtual machine
 */

type vm struct {
	pc uint;
	// The current execution frame.  If execution is within a
	// block, this may be a child of the original function
	// activation frame.
	f *Frame;
	// The original function activation frame.  This is used to
	// access function out args.
	activation *Frame;
}

type code []func(*vm)

func (i code) exec(fr *Frame) {
	v := vm{0, fr, fr};

	l := uint(len(i));
	for v.pc < l {
		pc := v.pc;
		v.pc++;
		i[pc](&v);
	}
}

/*
 * Code buffer
 */

type codeBuf struct {
	instrs code;
}

func newCodeBuf() *codeBuf {
	return &codeBuf{make(code, 0, 16)};
}

func (b codeBuf) push(instr func(*vm)) {
	n := len(b.instrs);
	if n >= cap(b.instrs) {
		a := make(code, n, n*2);
		for i := range b.instrs {
			a[i] = b.instrs[i];
		}
		b.instrs = a;
	}
	b.instrs = b.instrs[0:n+1];
	b.instrs[n] = instr;
}

func (b codeBuf) get() code {
	// Freeze this buffer into an array of exactly the right size
	a := make(code, len(b.instrs));
	for i := range b.instrs {
		a[i] = b.instrs[i];
	}
	return code(a);
}

/*
 * User-defined functions
 */

type evalFunc struct {
	sc *Scope;
	fr *Frame;
	code code;
}

func (f *evalFunc) NewFrame() *Frame {
	return f.sc.NewFrame(f.fr);
}

func (f *evalFunc) Call(fr *Frame) {
	f.code.exec(fr);
}
<commit_msg>Fix bug where nothing could ever be added to a code buffer.<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package eval

import (
	"container/vector";
	"eval";
)

/*
 * Virtual machine
 */

type vm struct {
	pc uint;
	// The current execution frame.  If execution is within a
	// block, this may be a child of the original function
	// activation frame.
	f *Frame;
	// The original function activation frame.  This is used to
	// access function out args.
	activation *Frame;
}

type code []func(*vm)

func (i code) exec(fr *Frame) {
	v := vm{0, fr, fr};

	l := uint(len(i));
	for v.pc < l {
		pc := v.pc;
		v.pc++;
		i[pc](&v);
	}
}

/*
 * Code buffer
 */

type codeBuf struct {
	instrs code;
}

func newCodeBuf() *codeBuf {
	return &codeBuf{make(code, 0, 16)};
}

func (b *codeBuf) push(instr func(*vm)) {
	n := len(b.instrs);
	if n >= cap(b.instrs) {
		a := make(code, n, n*2);
		for i := range b.instrs {
			a[i] = b.instrs[i];
		}
		b.instrs = a;
	}
	b.instrs = b.instrs[0:n+1];
	b.instrs[n] = instr;
}

func (b *codeBuf) get() code {
	// Freeze this buffer into an array of exactly the right size
	a := make(code, len(b.instrs));
	for i := range b.instrs {
		a[i] = b.instrs[i];
	}
	return code(a);
}

/*
 * User-defined functions
 */

type evalFunc struct {
	sc *Scope;
	fr *Frame;
	code code;
}

func (f *evalFunc) NewFrame() *Frame {
	return f.sc.NewFrame(f.fr);
}

func (f *evalFunc) Call(fr *Frame) {
	f.code.exec(fr);
}
<|endoftext|>


<commit_before>// +build !appengine

package geoip

import (
	"gnd.la/app"
	"gnd.la/log"

	"gopkgs.com/geoip.v1"
)

const (
	geoIPKey  = "__gondola_geoip"
	recordKey = "__gondola_record"
)

func getGeoIPRecord(ctx *app.Context) *geoip.Record {
	if rec, _ := ctx.Get(recordKey).(*geoip.Record); rec != nil {
		return rec
	}
	g, _ := ctx.App().Get(geoIPKey).(*geoip.GeoIP)
	if g == nil {
		log.Warning("no GeoIP data loaded - did you call geoip.Load()?")
		return nil
	}
	rec, _ := g.Lookup(ctx.RemoteAddress())
	ctx.Set(recordKey, rec)
	return rec
}

func country(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && rec.Country != nil {
		return rec.Country.Code
	}
	return ""
}

func region(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && len(rec.Subdivisions) > 0 {
		return rec.Subdivisions[0].Name.String()
	}
	return ""
}

func city(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && rec.City != nil {
		return rec.City.Name.String()
	}
	return ""
}

func latLong(ctx *app.Context) (float64, float64, bool) {
	rec := getGeoIPRecord(ctx)
	if rec != nil && (rec.Latitude != 0 || rec.Longitude != 0) {
		return rec.Latitude, rec.Longitude, true
	}
	return 0, 0, false
}

func loadDatabase(a *app.App, filename string) error {
	g, err := geoip.Open(filename)
	if err != nil {
		return err
	}
	a.Set(geoIPKey, g)
	return nil
}
<commit_msg>Import geoip pkg directly from github<commit_after>// +build !appengine

package geoip

import (
	"gnd.la/app"
	"gnd.la/log"

	"github.com/rainycape/geoip"
)

const (
	geoIPKey  = "__gondola_geoip"
	recordKey = "__gondola_record"
)

func getGeoIPRecord(ctx *app.Context) *geoip.Record {
	if rec, _ := ctx.Get(recordKey).(*geoip.Record); rec != nil {
		return rec
	}
	g, _ := ctx.App().Get(geoIPKey).(*geoip.GeoIP)
	if g == nil {
		log.Warning("no GeoIP data loaded - did you call geoip.Load()?")
		return nil
	}
	rec, _ := g.Lookup(ctx.RemoteAddress())
	ctx.Set(recordKey, rec)
	return rec
}

func country(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && rec.Country != nil {
		return rec.Country.Code
	}
	return ""
}

func region(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && len(rec.Subdivisions) > 0 {
		return rec.Subdivisions[0].Name.String()
	}
	return ""
}

func city(ctx *app.Context) string {
	rec := getGeoIPRecord(ctx)
	if rec != nil && rec.City != nil {
		return rec.City.Name.String()
	}
	return ""
}

func latLong(ctx *app.Context) (float64, float64, bool) {
	rec := getGeoIPRecord(ctx)
	if rec != nil && (rec.Latitude != 0 || rec.Longitude != 0) {
		return rec.Latitude, rec.Longitude, true
	}
	return 0, 0, false
}

func loadDatabase(a *app.App, filename string) error {
	g, err := geoip.Open(filename)
	if err != nil {
		return err
	}
	a.Set(geoIPKey, g)
	return nil
}
<|endoftext|>


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

package rocserv

import (
	"encoding/json"
	"fmt"
	"github.com/shawnfeng/hystrix-go/hystrix"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"sync"
	"time"
)

type ItemConf struct {
	*hystrix.CommandConfig
	Name   string `json:"name"`
	Enable bool   `json:"enable"`
	Source int32  `json:"source"`
}

type BreakerConf struct {
	globalConf []*ItemConf
	servConf   []*ItemConf
	funcConf   map[string]*ItemConf
	rwMutex    sync.RWMutex

	rawServConf   string
	rawGlobalConf string
	rawConfMutex  sync.RWMutex
}

func NewBreakerConf() *BreakerConf {
	return &BreakerConf{}
}

func (m *BreakerConf) tryUpdate(servname, rawGlobalConf, rawServConf string) {
	fun := "BreakerConf.tryUpdate -->"

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		slog.Infof("%s servname:%s rawGlobalConf:%s rawServConf:%s dur:%d", fun, servname, rawGlobalConf, rawServConf, dur)
	}()

	if len(rawGlobalConf) > 0 && rawGlobalConf != m.getRawGlobalConf() {

		var globalConf []*ItemConf
		err := json.Unmarshal([]byte(rawGlobalConf), &globalConf)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawGlobalConf:%s", fun, servname, rawGlobalConf)
		} else {
			m.setGlobalConf(globalConf)
			m.setRawGlobalConf(rawGlobalConf)
		}
	}

	if len(rawServConf) > 0 && rawServConf != m.getRawServConf() {
		var items []*ItemConf
		err := json.Unmarshal([]byte(rawServConf), &items)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawServConf:%s", fun, servname, rawServConf)
		} else {

			var servConf []*ItemConf
			funcConf := make(map[string]*ItemConf)
			for _, item := range items {
				if item.Name == "default" {
					servConf = append(servConf, item)
					continue
				}

				key := fmt.Sprintf("%d.%s", item.Source, item.Name)
				funcConf[key] = item
			}
			m.setServConf(servConf)
			m.setFuncConf(funcConf)

			m.setRawServConf(rawServConf)
		}
	}
}

func (m *BreakerConf) getRawServConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawServConf
}

func (m *BreakerConf) getRawGlobalConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawGlobalConf
}

func (m *BreakerConf) setRawGlobalConf(rawGlobalConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawGlobalConf = rawGlobalConf
}

func (m *BreakerConf) setRawServConf(rawServConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawServConf = rawServConf
}

func (m *BreakerConf) getFuncConf(source int32, funcName string) *ItemConf {

	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	key := fmt.Sprintf("%d.%s", source, funcName)
	if c, ok := m.funcConf[key]; ok {
		return c
	}

	if m.servConf != nil {
		for _, item := range m.servConf {
			if item.Source == source {
				return item
			}
		}
	}

	if m.globalConf != nil {
		for _, item := range m.globalConf {
			if item.Source == source {
				return item
			}
		}
	}

	return &ItemConf{
		CommandConfig: &hystrix.CommandConfig{},
		Enable:        false,
	}
}

func (m *BreakerConf) setGlobalConf(globalConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.globalConf = globalConf
}

func (m *BreakerConf) setServConf(servConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.servConf = servConf
}

func (m *BreakerConf) setFuncConf(funcConf map[string]*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.funcConf = funcConf
}

type Breaker struct {
	conf     *BreakerConf
	servName string

	clientLookup ClientLookup
	useFuncConf  map[string]*ItemConf
	rwMutex      sync.RWMutex
}

func NewBreaker(clientLookup ClientLookup) *Breaker {
	hystrix.SetLogger(slog.GetLogger())
	m := &Breaker{
		useFuncConf:  make(map[string]*ItemConf),
		clientLookup: clientLookup,
		servName:     clientLookup.ServKey(),
		conf:         NewBreakerConf(),
	}

	go m.run()

	return m
}

func (m *Breaker) run() {

	for {
		rawServConf := m.clientLookup.GetBreakerServConf()
		rawGlobalConf := m.clientLookup.GetBreakerGlobalConf()
		m.conf.tryUpdate(m.servName, rawGlobalConf, rawServConf)

		time.Sleep(time.Second * 30)
	}
}

func (m *Breaker) generateKey(source int32, servid int, funcName string) string {
	return fmt.Sprintf("%d.%s.%d.%s", source, m.servName, servid, funcName)
}

func (m *Breaker) getUseFuncConf(key string) (*ItemConf, bool) {
	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	conf, ok := m.useFuncConf[key]
	return conf, ok
}

func (m *Breaker) setUseFuncConf(key string, conf *ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.useFuncConf[key] = conf
}

func (m *Breaker) checkOrUpdateConf(source int32, servid int, funcName string) (needBreaker bool) {
	//fun := "Breaker.checkOrUpdateConf -->"

	key := m.generateKey(source, servid, funcName)

	newConf := m.conf.getFuncConf(source, funcName)
	//slog.Infof("%s servid:%d funcName:%s key:%s, newConf:%v", fun, servid, funcName, key, *newConf)

	if newConf.Enable == false {
		return false
	}

	useConf, ok := m.getUseFuncConf(key)
	if ok == false ||
		useConf.Enable || newConf.Enable ||
		useConf.Timeout != newConf.Timeout ||
		useConf.SleepWindow != newConf.SleepWindow ||
		useConf.RequestVolumeThreshold != newConf.RequestVolumeThreshold ||
		useConf.ErrorPercentThreshold != newConf.ErrorPercentThreshold ||
		useConf.MaxConcurrentRequests != newConf.MaxConcurrentRequests {

		hystrix.ConfigureCommand(key, *newConf.CommandConfig)

		m.setUseFuncConf(key, newConf)
	}

	return true
}

func (m *Breaker) Do(source int32, servid int, funcName string, run func() error, fallback func(error) error) error {
	fun := "Breaker.Do -->"

	if m.checkOrUpdateConf(source, servid, funcName) == false {
		return run()
	}

	st := stime.NewTimeStat()
	key := m.generateKey(source, servid, funcName)

	err := hystrix.Do(key, run, fallback)
	if err != nil {
		slog.Errorf("%s key:%s err:%s", fun, key, err.Error())
	}

	dur := st.Duration()
	slog.Infof("%s servid:%d funcName:%s key:%s dur:%d", fun, servid, funcName, key, dur)

	return err
}
<commit_msg>MOD:熔断增加统计报警<commit_after>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"encoding/json"
	"fmt"
	"github.com/shawnfeng/hystrix-go/hystrix"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"sync"
	"time"
)

type ItemConf struct {
	*hystrix.CommandConfig
	Name   string `json:"name"`
	Enable bool   `json:"enable"`
	Source int32  `json:"source"`
}

type BreakerConf struct {
	globalConf []*ItemConf
	servConf   []*ItemConf
	funcConf   map[string]*ItemConf
	rwMutex    sync.RWMutex

	rawServConf   string
	rawGlobalConf string
	rawConfMutex  sync.RWMutex
}

func NewBreakerConf() *BreakerConf {
	return &BreakerConf{}
}

func (m *BreakerConf) tryUpdate(servname, rawGlobalConf, rawServConf string) {
	fun := "BreakerConf.tryUpdate -->"

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		slog.Infof("%s servname:%s rawGlobalConf:%s rawServConf:%s dur:%d", fun, servname, rawGlobalConf, rawServConf, dur)
	}()

	if len(rawGlobalConf) > 0 && rawGlobalConf != m.getRawGlobalConf() {

		var globalConf []*ItemConf
		err := json.Unmarshal([]byte(rawGlobalConf), &globalConf)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawGlobalConf:%s", fun, servname, rawGlobalConf)
		} else {
			m.setGlobalConf(globalConf)
			m.setRawGlobalConf(rawGlobalConf)
		}
	}

	if len(rawServConf) > 0 && rawServConf != m.getRawServConf() {
		var items []*ItemConf
		err := json.Unmarshal([]byte(rawServConf), &items)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawServConf:%s", fun, servname, rawServConf)
		} else {

			var servConf []*ItemConf
			funcConf := make(map[string]*ItemConf)
			for _, item := range items {
				if item.Name == "default" {
					servConf = append(servConf, item)
					continue
				}

				key := fmt.Sprintf("%d.%s", item.Source, item.Name)
				funcConf[key] = item
			}
			m.setServConf(servConf)
			m.setFuncConf(funcConf)

			m.setRawServConf(rawServConf)
		}
	}
}

func (m *BreakerConf) getRawServConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawServConf
}

func (m *BreakerConf) getRawGlobalConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawGlobalConf
}

func (m *BreakerConf) setRawGlobalConf(rawGlobalConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawGlobalConf = rawGlobalConf
}

func (m *BreakerConf) setRawServConf(rawServConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawServConf = rawServConf
}

func (m *BreakerConf) getFuncConf(source int32, funcName string) *ItemConf {

	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	key := fmt.Sprintf("%d.%s", source, funcName)
	if c, ok := m.funcConf[key]; ok {
		return c
	}

	if m.servConf != nil {
		for _, item := range m.servConf {
			if item.Source == source {
				return item
			}
		}
	}

	if m.globalConf != nil {
		for _, item := range m.globalConf {
			if item.Source == source {
				return item
			}
		}
	}

	return &ItemConf{
		CommandConfig: &hystrix.CommandConfig{},
		Enable:        false,
	}
}

func (m *BreakerConf) setGlobalConf(globalConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.globalConf = globalConf
}

func (m *BreakerConf) setServConf(servConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.servConf = servConf
}

func (m *BreakerConf) setFuncConf(funcConf map[string]*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.funcConf = funcConf
}

type BreakerStat struct {
	key   string
	fail  int64
	total int64
}

type Breaker struct {
	conf     *BreakerConf
	servName string

	clientLookup ClientLookup
	useFuncConf  map[string]*ItemConf
	rwMutex      sync.RWMutex

	statCounter map[string]*BreakerStat
	statChan    chan *BreakerStat
}

func NewBreaker(clientLookup ClientLookup) *Breaker {
	hystrix.SetLogger(slog.GetLogger())
	m := &Breaker{
		useFuncConf:  make(map[string]*ItemConf),
		clientLookup: clientLookup,
		servName:     clientLookup.ServKey(),
		conf:         NewBreakerConf(),
		statCounter:  make(map[string]*BreakerStat),
		statChan:     make(chan *BreakerStat, 1024*10),
	}

	go m.run()
	go m.monitor()

	return m
}

func (m *Breaker) run() {

	for {
		rawServConf := m.clientLookup.GetBreakerServConf()
		rawGlobalConf := m.clientLookup.GetBreakerGlobalConf()
		m.conf.tryUpdate(m.servName, rawGlobalConf, rawServConf)

		time.Sleep(time.Second * 30)
	}
}

func (m *Breaker) monitor() {
	fun := "Breaker.monitor -->"

	ticker := time.NewTicker(60 * time.Second)
	for {
		select {
		case <-ticker.C:
			for _, stat := range m.statCounter {
				if stat.fail > 5 && stat.total > 5 &&
					(float64(stat.fail)/float64(stat.total)) > 0.02 {
					slog.Errorf("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				} else {
					slog.Infof("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				}
			}

			m.statCounter = make(map[string]*BreakerStat)

		case stat := <-m.statChan:
			tmp := &BreakerStat{
				key: stat.key,
			}
			if item, ok := m.statCounter[stat.key]; ok {
				tmp = item
			}

			tmp.fail += stat.fail
			tmp.total += stat.total

			m.statCounter[tmp.key] = tmp
		}
	}
}

func (m *Breaker) doStat(key string, total, fail int64) {
	fun := "Breaker.doStat -->"

	stat := &BreakerStat{
		key:   key,
		total: total,
		fail:  fail,
	}

	select {
	case m.statChan <- stat:
	default:
		slog.Errorf("%s drop, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
	}
}

func (m *Breaker) generateKey(source int32, servid int, funcName string) string {
	return fmt.Sprintf("%d.%s.%d.%s", source, m.servName, servid, funcName)
}

func (m *Breaker) getUseFuncConf(key string) (*ItemConf, bool) {
	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	conf, ok := m.useFuncConf[key]
	return conf, ok
}

func (m *Breaker) setUseFuncConf(key string, conf *ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.useFuncConf[key] = conf
}

func (m *Breaker) checkOrUpdateConf(source int32, servid int, funcName string) (needBreaker bool) {
	//fun := "Breaker.checkOrUpdateConf -->"

	key := m.generateKey(source, servid, funcName)

	newConf := m.conf.getFuncConf(source, funcName)
	//slog.Infof("%s servid:%d funcName:%s key:%s, newConf:%v", fun, servid, funcName, key, *newConf)

	if newConf.Enable == false {
		return false
	}

	useConf, ok := m.getUseFuncConf(key)
	if ok == false ||
		useConf.Enable || newConf.Enable ||
		useConf.Timeout != newConf.Timeout ||
		useConf.SleepWindow != newConf.SleepWindow ||
		useConf.RequestVolumeThreshold != newConf.RequestVolumeThreshold ||
		useConf.ErrorPercentThreshold != newConf.ErrorPercentThreshold ||
		useConf.MaxConcurrentRequests != newConf.MaxConcurrentRequests {

		hystrix.ConfigureCommand(key, *newConf.CommandConfig)

		m.setUseFuncConf(key, newConf)
	}

	return true
}

func (m *Breaker) Do(source int32, servid int, funcName string, run func() error, fallback func(error) error) error {
	fun := "Breaker.Do -->"

	if m.checkOrUpdateConf(source, servid, funcName) == false {
		return run()
	}

	st := stime.NewTimeStat()
	key := m.generateKey(source, servid, funcName)

	fail := int64(0)
	err := hystrix.Do(key, run, fallback)
	if err != nil {
		slog.Warnf("%s key:%s err:%s", fun, key, err.Error())
		fail = 1
	}

	m.doStat(key, 1, fail)

	dur := st.Duration()
	slog.Infof("%s servid:%d funcName:%s key:%s dur:%d", fun, servid, funcName, key, dur)

	return err
}
<|endoftext|>


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

package rocserv

import (
	"encoding/json"
	"fmt"
	"github.com/shawnfeng/hystrix-go/hystrix"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"sync"
	"time"
)

type ItemConf struct {
	*hystrix.CommandConfig
	Name   string `json:"name"`
	Enable bool   `json:"enable"`
	Source int32  `json:"source"`
}

type BreakerConf struct {
	globalConf []*ItemConf
	servConf   []*ItemConf
	funcConf   map[string]*ItemConf
	rwMutex    sync.RWMutex

	rawServConf   string
	rawGlobalConf string
	rawConfMutex  sync.RWMutex
}

func NewBreakerConf() *BreakerConf {
	return &BreakerConf{}
}

func (m *BreakerConf) tryUpdate(servname, rawGlobalConf, rawServConf string) {
	fun := "BreakerConf.tryUpdate -->"

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		slog.Infof("%s servname:%s rawGlobalConf:%s rawServConf:%s dur:%d", fun, servname, rawGlobalConf, rawServConf, dur)
	}()

	if len(rawGlobalConf) > 0 && rawGlobalConf != m.getRawGlobalConf() {

		var globalConf []*ItemConf
		err := json.Unmarshal([]byte(rawGlobalConf), &globalConf)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawGlobalConf:%s", fun, servname, rawGlobalConf)
		} else {
			m.setGlobalConf(globalConf)
			m.setRawGlobalConf(rawGlobalConf)
		}
	}

	if len(rawServConf) > 0 && rawServConf != m.getRawServConf() {
		var items []*ItemConf
		err := json.Unmarshal([]byte(rawServConf), &items)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawServConf:%s", fun, servname, rawServConf)
		} else {

			var servConf []*ItemConf
			funcConf := make(map[string]*ItemConf)
			for _, item := range items {
				if item.Name == "default" {
					servConf = append(servConf, item)
					continue
				}

				key := fmt.Sprintf("%d.%s", item.Source, item.Name)
				funcConf[key] = item
			}
			m.setServConf(servConf)
			m.setFuncConf(funcConf)

			m.setRawServConf(rawServConf)
		}
	}
}

func (m *BreakerConf) getRawServConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawServConf
}

func (m *BreakerConf) getRawGlobalConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawGlobalConf
}

func (m *BreakerConf) setRawGlobalConf(rawGlobalConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawGlobalConf = rawGlobalConf
}

func (m *BreakerConf) setRawServConf(rawServConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawServConf = rawServConf
}

func (m *BreakerConf) getFuncConf(source int32, funcName string) *ItemConf {

	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	key := fmt.Sprintf("%d.%s", source, funcName)
	if c, ok := m.funcConf[key]; ok {
		return c
	}

	if m.servConf != nil {
		for _, item := range m.servConf {
			if item.Source == source {
				return item
			}
		}
	}

	if m.globalConf != nil {
		for _, item := range m.globalConf {
			if item.Source == source {
				return item
			}
		}
	}

	return &ItemConf{
		CommandConfig: &hystrix.CommandConfig{},
		Enable:        false,
	}
}

func (m *BreakerConf) setGlobalConf(globalConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.globalConf = globalConf
}

func (m *BreakerConf) setServConf(servConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.servConf = servConf
}

func (m *BreakerConf) setFuncConf(funcConf map[string]*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.funcConf = funcConf
}

type BreakerStat struct {
	key   string
	fail  int64
	total int64
}

type Breaker struct {
	conf     *BreakerConf
	servName string

	clientLookup ClientLookup
	useFuncConf  map[string]*ItemConf
	rwMutex      sync.RWMutex

	statCounter map[string]*BreakerStat
	statChan    chan *BreakerStat
}

func NewBreaker(clientLookup ClientLookup) *Breaker {
	hystrix.SetLogger(slog.GetLogger())
	m := &Breaker{
		useFuncConf:  make(map[string]*ItemConf),
		clientLookup: clientLookup,
		servName:     clientLookup.ServKey(),
		conf:         NewBreakerConf(),
		statCounter:  make(map[string]*BreakerStat),
		statChan:     make(chan *BreakerStat, 1024*10),
	}

	go m.run()
	go m.monitor()

	return m
}

func (m *Breaker) run() {

	for {
		rawServConf := m.clientLookup.GetBreakerServConf()
		rawGlobalConf := m.clientLookup.GetBreakerGlobalConf()
		m.conf.tryUpdate(m.servName, rawGlobalConf, rawServConf)

		time.Sleep(time.Second * 30)
	}
}

func (m *Breaker) monitor() {
	fun := "Breaker.monitor -->"

	ticker := time.NewTicker(60 * time.Second)
	for {
		select {
		case <-ticker.C:
			for _, stat := range m.statCounter {
				if stat.fail > 5 && stat.total > 5 &&
					(float64(stat.fail)/float64(stat.total)) > 0.02 {
					slog.Errorf("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				} else {
					slog.Infof("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				}
			}

			m.statCounter = make(map[string]*BreakerStat)

		case stat := <-m.statChan:
			tmp := &BreakerStat{
				key: stat.key,
			}
			if item, ok := m.statCounter[stat.key]; ok {
				tmp = item
			}

			tmp.fail += stat.fail
			tmp.total += stat.total

			m.statCounter[tmp.key] = tmp
		}
	}
}

func (m *Breaker) doStat(key string, total, fail int64) {
	fun := "Breaker.doStat -->"

	stat := &BreakerStat{
		key:   key,
		total: total,
		fail:  fail,
	}

	select {
	case m.statChan <- stat:
	default:
		slog.Errorf("%s drop, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
	}
}

func (m *Breaker) generateKey(source int32, servid int, funcName string) string {
	return fmt.Sprintf("%d.%s.%d.%s", source, m.servName, servid, funcName)
}

func (m *Breaker) getUseFuncConf(key string) (*ItemConf, bool) {
	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	conf, ok := m.useFuncConf[key]
	return conf, ok
}

func (m *Breaker) setUseFuncConf(key string, conf *ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.useFuncConf[key] = conf
}

func (m *Breaker) checkOrUpdateConf(source int32, servid int, funcName string) (needBreaker bool) {
	//fun := "Breaker.checkOrUpdateConf -->"

	key := m.generateKey(source, servid, funcName)

	newConf := m.conf.getFuncConf(source, funcName)
	//slog.Infof("%s servid:%d funcName:%s key:%s, newConf:%v", fun, servid, funcName, key, *newConf)

	if newConf.Enable == false {
		return false
	}

	useConf, ok := m.getUseFuncConf(key)
	if ok == false ||
		useConf.Enable || newConf.Enable ||
		useConf.Timeout != newConf.Timeout ||
		useConf.SleepWindow != newConf.SleepWindow ||
		useConf.RequestVolumeThreshold != newConf.RequestVolumeThreshold ||
		useConf.ErrorPercentThreshold != newConf.ErrorPercentThreshold ||
		useConf.MaxConcurrentRequests != newConf.MaxConcurrentRequests {

		hystrix.ConfigureCommand(key, *newConf.CommandConfig)

		m.setUseFuncConf(key, newConf)
	}

	return true
}

func (m *Breaker) Do(source int32, servid int, funcName string, run func() error, fallback func(error) error) error {
	fun := "Breaker.Do -->"

	if m.checkOrUpdateConf(source, servid, funcName) == false {
		return run()
	}

	st := stime.NewTimeStat()
	key := m.generateKey(source, servid, funcName)

	fail := int64(0)
	err := hystrix.Do(key, run, fallback)
	if err != nil {
		slog.Warnf("%s key:%s err:%s", fun, key, err.Error())
		fail = 1
	}

	m.doStat(key, 1, fail)

	dur := st.Duration()
	slog.Infof("%s key:%s dur:%d", fun, servid, funcName, key, dur)

	return err
}
<commit_msg>MOD:log<commit_after>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"encoding/json"
	"fmt"
	"github.com/shawnfeng/hystrix-go/hystrix"
	"github.com/shawnfeng/sutil/slog"
	"github.com/shawnfeng/sutil/stime"
	"sync"
	"time"
)

type ItemConf struct {
	*hystrix.CommandConfig
	Name   string `json:"name"`
	Enable bool   `json:"enable"`
	Source int32  `json:"source"`
}

type BreakerConf struct {
	globalConf []*ItemConf
	servConf   []*ItemConf
	funcConf   map[string]*ItemConf
	rwMutex    sync.RWMutex

	rawServConf   string
	rawGlobalConf string
	rawConfMutex  sync.RWMutex
}

func NewBreakerConf() *BreakerConf {
	return &BreakerConf{}
}

func (m *BreakerConf) tryUpdate(servname, rawGlobalConf, rawServConf string) {
	fun := "BreakerConf.tryUpdate -->"

	st := stime.NewTimeStat()
	defer func() {
		dur := st.Duration()
		slog.Infof("%s servname:%s rawGlobalConf:%s rawServConf:%s dur:%d", fun, servname, rawGlobalConf, rawServConf, dur)
	}()

	if len(rawGlobalConf) > 0 && rawGlobalConf != m.getRawGlobalConf() {

		var globalConf []*ItemConf
		err := json.Unmarshal([]byte(rawGlobalConf), &globalConf)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawGlobalConf:%s", fun, servname, rawGlobalConf)
		} else {
			m.setGlobalConf(globalConf)
			m.setRawGlobalConf(rawGlobalConf)
		}
	}

	if len(rawServConf) > 0 && rawServConf != m.getRawServConf() {
		var items []*ItemConf
		err := json.Unmarshal([]byte(rawServConf), &items)
		if err != nil {
			slog.Errorf("%s servname:%s Unmarshal err, rawServConf:%s", fun, servname, rawServConf)
		} else {

			var servConf []*ItemConf
			funcConf := make(map[string]*ItemConf)
			for _, item := range items {
				if item.Name == "default" {
					servConf = append(servConf, item)
					continue
				}

				key := fmt.Sprintf("%d.%s", item.Source, item.Name)
				funcConf[key] = item
			}
			m.setServConf(servConf)
			m.setFuncConf(funcConf)

			m.setRawServConf(rawServConf)
		}
	}
}

func (m *BreakerConf) getRawServConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawServConf
}

func (m *BreakerConf) getRawGlobalConf() string {
	m.rawConfMutex.RLock()
	defer m.rawConfMutex.RUnlock()

	return m.rawGlobalConf
}

func (m *BreakerConf) setRawGlobalConf(rawGlobalConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawGlobalConf = rawGlobalConf
}

func (m *BreakerConf) setRawServConf(rawServConf string) {
	m.rawConfMutex.Lock()
	defer m.rawConfMutex.Unlock()

	m.rawServConf = rawServConf
}

func (m *BreakerConf) getFuncConf(source int32, funcName string) *ItemConf {

	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	key := fmt.Sprintf("%d.%s", source, funcName)
	if c, ok := m.funcConf[key]; ok {
		return c
	}

	if m.servConf != nil {
		for _, item := range m.servConf {
			if item.Source == source {
				return item
			}
		}
	}

	if m.globalConf != nil {
		for _, item := range m.globalConf {
			if item.Source == source {
				return item
			}
		}
	}

	return &ItemConf{
		CommandConfig: &hystrix.CommandConfig{},
		Enable:        false,
	}
}

func (m *BreakerConf) setGlobalConf(globalConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.globalConf = globalConf
}

func (m *BreakerConf) setServConf(servConf []*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.servConf = servConf
}

func (m *BreakerConf) setFuncConf(funcConf map[string]*ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.funcConf = funcConf
}

type BreakerStat struct {
	key   string
	fail  int64
	total int64
}

type Breaker struct {
	conf     *BreakerConf
	servName string

	clientLookup ClientLookup
	useFuncConf  map[string]*ItemConf
	rwMutex      sync.RWMutex

	statCounter map[string]*BreakerStat
	statChan    chan *BreakerStat
}

func NewBreaker(clientLookup ClientLookup) *Breaker {
	hystrix.SetLogger(slog.GetLogger())
	m := &Breaker{
		useFuncConf:  make(map[string]*ItemConf),
		clientLookup: clientLookup,
		servName:     clientLookup.ServKey(),
		conf:         NewBreakerConf(),
		statCounter:  make(map[string]*BreakerStat),
		statChan:     make(chan *BreakerStat, 1024*10),
	}

	go m.run()
	go m.monitor()

	return m
}

func (m *Breaker) run() {

	for {
		rawServConf := m.clientLookup.GetBreakerServConf()
		rawGlobalConf := m.clientLookup.GetBreakerGlobalConf()
		m.conf.tryUpdate(m.servName, rawGlobalConf, rawServConf)

		time.Sleep(time.Second * 30)
	}
}

func (m *Breaker) monitor() {
	fun := "Breaker.monitor -->"

	ticker := time.NewTicker(60 * time.Second)
	for {
		select {
		case <-ticker.C:
			for _, stat := range m.statCounter {
				if stat.fail > 5 && stat.total > 5 &&
					(float64(stat.fail)/float64(stat.total)) > 0.02 {
					slog.Errorf("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				} else {
					slog.Infof("%s breaker stat, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
				}
			}

			m.statCounter = make(map[string]*BreakerStat)

		case stat := <-m.statChan:
			tmp := &BreakerStat{
				key: stat.key,
			}
			if item, ok := m.statCounter[stat.key]; ok {
				tmp = item
			}

			tmp.fail += stat.fail
			tmp.total += stat.total

			m.statCounter[tmp.key] = tmp
		}
	}
}

func (m *Breaker) doStat(key string, total, fail int64) {
	fun := "Breaker.doStat -->"

	stat := &BreakerStat{
		key:   key,
		total: total,
		fail:  fail,
	}

	select {
	case m.statChan <- stat:
	default:
		slog.Errorf("%s drop, key:%s, total:%d, fail:%d", fun, stat.key, stat.total, stat.fail)
	}
}

func (m *Breaker) generateKey(source int32, servid int, funcName string) string {
	return fmt.Sprintf("%d.%s.%d.%s", source, m.servName, servid, funcName)
}

func (m *Breaker) getUseFuncConf(key string) (*ItemConf, bool) {
	m.rwMutex.RLock()
	defer m.rwMutex.RUnlock()

	conf, ok := m.useFuncConf[key]
	return conf, ok
}

func (m *Breaker) setUseFuncConf(key string, conf *ItemConf) {
	m.rwMutex.Lock()
	defer m.rwMutex.Unlock()

	m.useFuncConf[key] = conf
}

func (m *Breaker) checkOrUpdateConf(source int32, servid int, funcName string) (needBreaker bool) {
	//fun := "Breaker.checkOrUpdateConf -->"

	key := m.generateKey(source, servid, funcName)

	newConf := m.conf.getFuncConf(source, funcName)
	//slog.Infof("%s servid:%d funcName:%s key:%s, newConf:%v", fun, servid, funcName, key, *newConf)

	if newConf.Enable == false {
		return false
	}

	useConf, ok := m.getUseFuncConf(key)
	if ok == false ||
		useConf.Enable || newConf.Enable ||
		useConf.Timeout != newConf.Timeout ||
		useConf.SleepWindow != newConf.SleepWindow ||
		useConf.RequestVolumeThreshold != newConf.RequestVolumeThreshold ||
		useConf.ErrorPercentThreshold != newConf.ErrorPercentThreshold ||
		useConf.MaxConcurrentRequests != newConf.MaxConcurrentRequests {

		hystrix.ConfigureCommand(key, *newConf.CommandConfig)

		m.setUseFuncConf(key, newConf)
	}

	return true
}

func (m *Breaker) Do(source int32, servid int, funcName string, run func() error, fallback func(error) error) error {
	fun := "Breaker.Do -->"

	if m.checkOrUpdateConf(source, servid, funcName) == false {
		return run()
	}

	st := stime.NewTimeStat()
	key := m.generateKey(source, servid, funcName)

	fail := int64(0)
	err := hystrix.Do(key, run, fallback)
	if err != nil {
		slog.Warnf("%s key:%s err:%s", fun, key, err.Error())
		fail = 1
	}

	m.doStat(key, 1, fail)

	dur := st.Duration()
	slog.Infof("%s key:%s dur:%d", fun, key, dur)

	return err
}
<|endoftext|>