<commit_before>package jsontaggedbson

import (
	"encoding/json"
	"fmt"
	"reflect"

	"gopkg.in/mgo.v2/bson"
)

type Intermediate interface{}

//TaggedIntermediate creates BSON-tagged struct ready for serializing from a JSON struct.
func TaggedIntermediate(obj interface{}) (bdoc Intermediate, err error) {
	JSON, err := json.Marshal(obj)
	if err != nil {
		return bdoc, fmt.Errorf("JSONTaggedBSON: json.Marshal: %v", err)
	}
	err = bson.UnmarshalJSON(JSON, &bdoc)
	if err != nil {
		return bdoc, fmt.Errorf("JSONTaggedBSON: bson.MarshalJSON: %v", err)
	}
	return bdoc, nil
}

//Marshal a JSON-tagged object to BSON.
func Marshal(v interface{}) (BSON []byte, err error) {
	var bdoc Intermediate
	if bdoc, err = TaggedIntermediate(v); err != nil {
		return BSON, fmt.Errorf("TaggedInterface: %v", err)
	}
	return bson.Marshal(bdoc)
}

func isPtr(i interface{}) bool {
	return reflect.ValueOf(i).Kind() == reflect.Ptr
}

//FromTaggedIntermediate assigns the elements of bdoc to the JSON-tagged struct pointed to by v.
func FromTaggedIntermediate(bdoc Intermediate, v interface{}) (err error) {
	if !isPtr(v) {
		return fmt.Errorf("v must be a pointer to a struct")
	}
	var JSON []byte
	if JSON, err = bson.MarshalJSON(bdoc); err != nil {
		return fmt.Errorf("bson.MarshalJSON: %v", err)
	}
	if err = json.Unmarshal(JSON, v); err != nil {
		return fmt.Errorf("json.Unmarshal: %v", err)
	}
	return nil
}

//Unmarshal a BSON-serialized object to the json-tagged object pointed to by v.
func Unmarshal(data []byte, v interface{}) (err error) {

	var JSON []byte
	if err = bson.Unmarshal(data, &intermediate); err != nil {
		return fmt.Errorf("bson.Unmarshal: %v", err)
	}
	if JSON, err = bson.MarshalJSON(intermediate); err != nil {
		return fmt.Errorf("bson.MarshalJSON: %v", err)
	}
	if err = json.Unmarshal(JSON, v); err != nil {
		return fmt.Errorf("json.Unmarshal: %v", err)
	}
	return nil
}
<commit_msg>fixed compile error<commit_after>package jsontaggedbson

import (
	"encoding/json"
	"fmt"
	"reflect"

	"gopkg.in/mgo.v2/bson"
)

type Intermediate interface{}

//TaggedIntermediate creates BSON-tagged struct ready for serializing from a JSON struct.
func TaggedIntermediate(obj interface{}) (bdoc Intermediate, err error) {
	JSON, err := json.Marshal(obj)
	if err != nil {
		return bdoc, fmt.Errorf("JSONTaggedBSON: json.Marshal: %v", err)
	}
	err = bson.UnmarshalJSON(JSON, &bdoc)
	if err != nil {
		return bdoc, fmt.Errorf("JSONTaggedBSON: bson.MarshalJSON: %v", err)
	}
	return bdoc, nil
}

//Marshal a JSON-tagged object to BSON.
func Marshal(v interface{}) (BSON []byte, err error) {
	var bdoc Intermediate
	if bdoc, err = TaggedIntermediate(v); err != nil {
		return BSON, fmt.Errorf("TaggedInterface: %v", err)
	}
	return bson.Marshal(bdoc)
}

func isPtr(i interface{}) bool {
	return reflect.ValueOf(i).Kind() == reflect.Ptr
}

//FromTaggedIntermediate assigns the elements of bdoc to the JSON-tagged struct pointed to by v.
func FromTaggedIntermediate(bdoc Intermediate, v interface{}) (err error) {
	if !isPtr(v) {
		return fmt.Errorf("v must be a pointer to a struct")
	}
	var JSON []byte
	if JSON, err = bson.MarshalJSON(bdoc); err != nil {
		return fmt.Errorf("bson.MarshalJSON: %v", err)
	}
	if err = json.Unmarshal(JSON, v); err != nil {
		return fmt.Errorf("json.Unmarshal: %v", err)
	}
	return nil
}

//Unmarshal a BSON-serialized object to the json-tagged object pointed to by v.
func Unmarshal(data []byte, v interface{}) (err error) {
	var bdoc interface{}
	var JSON []byte
	if err = bson.Unmarshal(data, &bdoc); err != nil {
		return fmt.Errorf("bson.Unmarshal: %v", err)
	}
	if JSON, err = bson.MarshalJSON(bdoc); err != nil {
		return fmt.Errorf("bson.MarshalJSON: %v", err)
	}
	if err = json.Unmarshal(JSON, v); err != nil {
		return fmt.Errorf("json.Unmarshal: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
MIT License

Copyright (c) 2017 Curtis La Graff

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// pad provides basic functionallity for padding sides of a string.
package pad

import "bytes"

type direction bool

// LeftDir is used to specify text will be padded on the left side.
const LeftDir direction = false

// RightDir is used to specify text will be padded on the right side.
const RightDir direction = true

// Pad is used to pad a given string using the specified character to satisfy the difference between the text length and the length specified.
func Pad(text string, dir direction, length int, char byte) string {
	startLength := len(text)
	diff := length - startLength

	// If no difference exists, return the text as it satisfies the length.
	if diff <= 0 {
		return text
	}

	// Writing to a buffer to improve performance.
	buff := bytes.NewBuffer(nil)
	if dir == RightDir {
		buff.WriteString(text)
	}

	for i := 0; i < diff; i++ {
		buff.WriteByte(char)
	}

	if dir == LeftDir {
		buff.WriteString(text)
	}

	return buff.String()
}

// Left will pad the left of the string with spaces to meet the specified length.
func Left(text string, length int) string {
	return Pad(text, LeftDir, length, ' ')
}

// LeftChar will pad the left of the string with the specified character to meet
// the specified length.
func LeftChar(text string, length int, char byte) string {
	return Pad(text, LeftDir, length, char)
}

// Right will pad the right of the string with spaces to meet the
// specified length.
func Right(text string, length int) string {
	return Pad(text, RightDir, length, ' ')
}

// RightChar will pad the right of the string with the specified character to
// meet the specified length.
func RightChar(text string, length int, char byte) string {
	return Pad(text, RightDir, length, char)
}
<commit_msg>fixed misspell and linting error<commit_after>/*
MIT License

Copyright (c) 2017 Curtis La Graff

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

// Package pad provides basic functionality for padding sides of a string.
package pad

import "bytes"

type direction bool

// LeftDir is used to specify text will be padded on the left side.
const LeftDir direction = false

// RightDir is used to specify text will be padded on the right side.
const RightDir direction = true

// Pad is used to pad a given string using the specified character to satisfy the difference between the text length and the length specified.
func Pad(text string, dir direction, length int, char byte) string {
	startLength := len(text)
	diff := length - startLength

	// If no difference exists, return the text as it satisfies the length.
	if diff <= 0 {
		return text
	}

	// Writing to a buffer to improve performance.
	buff := bytes.NewBuffer(nil)
	if dir == RightDir {
		buff.WriteString(text)
	}

	for i := 0; i < diff; i++ {
		buff.WriteByte(char)
	}

	if dir == LeftDir {
		buff.WriteString(text)
	}

	return buff.String()
}

// Left will pad the left of the string with spaces to meet the specified length.
func Left(text string, length int) string {
	return Pad(text, LeftDir, length, ' ')
}

// LeftChar will pad the left of the string with the specified character to meet
// the specified length.
func LeftChar(text string, length int, char byte) string {
	return Pad(text, LeftDir, length, char)
}

// Right will pad the right of the string with spaces to meet the
// specified length.
func Right(text string, length int) string {
	return Pad(text, RightDir, length, ' ')
}

// RightChar will pad the right of the string with the specified character to
// meet the specified length.
func RightChar(text string, length int, char byte) string {
	return Pad(text, RightDir, length, char)
}
<|endoftext|>


<commit_before>package jk

import (
	"errors"
	"strconv"
	"strings"
)

//Argument represent a single argument
type Argument struct {
	Sid string
	Tid int
	Rep string
}

// Arguments is a slice of Argument
type Arguments []Argument

//ArgMap is a map from string to Arguments
type ArgMap map[string]Arguments

// Pas represents a predicate argument structure
type Pas struct {
	Cfid string
	Args ArgMap
}

// NewPas returns a Pas
func NewPas(line string, knpstyle bool) (*Pas, error) {
	if !knpstyle {
		return nil, errors.New("Not Implemented")
	}
	return parseKnpStylePas(line)
}

func parseKnpStylePas(val string) (*Pas, error) {
	pas := new(Pas)
	pas.Args = ArgMap{}
	c0 := strings.Index(val, ":")
	c1 := c0 + 1 + strings.Index(val[c0+1:], ":")
	pas.Cfid = val[:c0] + ":" + val[c0+1:c1]

	if strings.Count(val, ":") < 2 {
		return nil, nil
	}

	ks := strings.Split(val[c1+1:], ";")
	for _, k := range ks {
		items := strings.Split(k, "/")
		casetype := items[1]
		if casetype == "U" || casetype == "-" {
			continue
		}
		if len(items) != 6 {
			return nil, errors.New("Invalid number of values")
		}

		mycase := items[0]
		rep := items[2]
		tid, err := strconv.Atoi(items[3])
		if err != nil {
			return nil, err
		}
		sid := items[5]
		arg := Argument{Sid: sid, Tid: tid, Rep: rep}

		as, ok := pas.Args[mycase]
		if !ok {
			as = Arguments{}
			pas.Args[mycase] = as
		}
		pas.Args[mycase] = append(pas.Args[mycase], arg)
	}
	return pas, nil
}
<commit_msg>Fix a bug<commit_after>package jk

import (
	"errors"
	"strconv"
	"strings"
)

//Argument represent a single argument
type Argument struct {
	Sid string
	Tid int
	Rep string
}

// Arguments is a slice of Argument
type Arguments []Argument

//ArgMap is a map from string to Arguments
type ArgMap map[string]Arguments

// Pas represents a predicate argument structure
type Pas struct {
	Cfid string
	Args ArgMap
}

// NewPas returns a Pas
func NewPas(line string, knpstyle bool) (*Pas, error) {
	if !knpstyle {
		return nil, errors.New("Not Implemented")
	}
	return parseKnpStylePas(line)
}

func parseKnpStylePas(val string) (*Pas, error) {
	if strings.Count(val, ":") < 2 {
		return nil, nil
	}

	pas := new(Pas)
	pas.Args = ArgMap{}
	c0 := strings.Index(val, ":")
	c1 := c0 + 1 + strings.Index(val[c0+1:], ":")
	pas.Cfid = val[:c0] + ":" + val[c0+1:c1]

	ks := strings.Split(val[c1+1:], ";")
	for _, k := range ks {
		items := strings.Split(k, "/")
		casetype := items[1]
		if casetype == "U" || casetype == "-" {
			continue
		}
		if len(items) != 6 {
			return nil, errors.New("Invalid number of values")
		}

		mycase := items[0]
		rep := items[2]
		tid, err := strconv.Atoi(items[3])
		if err != nil {
			return nil, err
		}
		sid := items[5]
		arg := Argument{Sid: sid, Tid: tid, Rep: rep}

		as, ok := pas.Args[mycase]
		if !ok {
			as = Arguments{}
			pas.Args[mycase] = as
		}
		pas.Args[mycase] = append(pas.Args[mycase], arg)
	}
	return pas, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ghmetrics

import (
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

// ghTokenUntilResetGaugeVec provides the 'github_token_reset' gauge that
// enables keeping track of GitHub reset times.
var ghTokenUntilResetGaugeVec = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "github_token_reset",
		Help: "Last reported GitHub token reset time.",
	},
	[]string{"token_hash", "api_version"},
)

// ghTokenUsageGaugeVec provides the 'github_token_usage' gauge that
// enables keeping track of GitHub calls and quotas.
var ghTokenUsageGaugeVec = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "github_token_usage",
		Help: "How many GitHub token requets are remaining for the current hour.",
	},
	[]string{"token_hash", "api_version"},
)

// ghRequestDurationHistVec provides the 'github_request_duration' histogram that keeps track
// of the duration of GitHub requests by API path.
var ghRequestDurationHistVec = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "github_request_duration",
		Help:    "GitHub request duration by API path.",
		Buckets: []float64{0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10},
	},
	[]string{"token_hash", "path", "status", "user_agent"},
)

// cacheCounter provides the 'ghcache_responses' counter vec that is indexed
// by the cache response mode.
var cacheCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "ghcache_responses",
		Help: "How many cache responses of each cache response mode there are.",
	},
	[]string{"mode", "path"},
)

// timeoutDuration provides the 'github_request_timeouts' histogram that keeps
// track of the timeouts of GitHub requests by API path.
var timeoutDuration = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "github_request_timeouts",
		Help:    "GitHub request timeout by API path.",
		Buckets: []float64{45, 60, 90, 120, 300},
	},
	[]string{"token_hash", "path", "user_agent"},
)

var muxTokenUsage, muxRequestMetrics sync.Mutex
var lastGitHubResponse time.Time

func init() {
	prometheus.MustRegister(ghTokenUntilResetGaugeVec)
	prometheus.MustRegister(ghTokenUsageGaugeVec)
	prometheus.MustRegister(ghRequestDurationHistVec)
	prometheus.MustRegister(cacheCounter)
	prometheus.MustRegister(timeoutDuration)
}

// CollectGitHubTokenMetrics publishes the rate limits of the github api to
// `github_token_usage` as well as `github_token_reset` on prometheus.
func CollectGitHubTokenMetrics(tokenHash, apiVersion string, headers http.Header, reqStartTime, responseTime time.Time) {
	remaining := headers.Get("X-RateLimit-Remaining")
	timeUntilReset := timestampStringToTime(headers.Get("X-RateLimit-Reset"))
	durationUntilReset := timeUntilReset.Sub(reqStartTime)

	remainingFloat, err := strconv.ParseFloat(remaining, 64)
	if err != nil {
		logrus.WithError(err).Infof("Couldn't convert number of remaining token requests into gauge value (float)")
	}
	if remainingFloat == 0 {
		logrus.WithFields(logrus.Fields{
			"header":     remaining,
			"user-agent": headers.Get("User-Agent"),
		}).Debug("Parsed GitHub header as indicating no remaining rate-limit.")
	}

	muxTokenUsage.Lock()
	isAfter := lastGitHubResponse.After(responseTime)
	if !isAfter {
		lastGitHubResponse = responseTime
	}
	muxTokenUsage.Unlock()
	if isAfter {
		logrus.WithField("last-github-response", lastGitHubResponse).WithField("response-time", responseTime).Debug("Previously pushed metrics of a newer response, skipping old metrics")
	} else {
		ghTokenUntilResetGaugeVec.With(prometheus.Labels{"token_hash": tokenHash, "api_version": apiVersion}).Set(float64(durationUntilReset.Nanoseconds()))
		ghTokenUsageGaugeVec.With(prometheus.Labels{"token_hash": tokenHash, "api_version": apiVersion}).Set(remainingFloat)
	}
}

// CollectGitHubRequestMetrics publishes the number of requests by API path to
// `github_requests` on prometheus.
func CollectGitHubRequestMetrics(tokenHash, path, statusCode, userAgent string, roundTripTime float64) {
	ghRequestDurationHistVec.With(prometheus.Labels{"token_hash": tokenHash, "path": simplifier.Simplify(path), "status": statusCode, "user_agent": userAgent}).Observe(roundTripTime)
}

// timestampStringToTime takes a unix timestamp and returns a `time.Time`
// from the given time.
func timestampStringToTime(tstamp string) time.Time {
	timestamp, err := strconv.ParseInt(tstamp, 10, 64)
	if err != nil {
		logrus.WithField("timestamp", tstamp).Info("Couldn't convert unix timestamp")
	}
	return time.Unix(timestamp, 0)
}

// CollectCacheRequestMetrics records a cache outcome for a specific path
func CollectCacheRequestMetrics(mode, path string) {
	cacheCounter.With(prometheus.Labels{"mode": mode, "path": simplifier.Simplify(path)}).Inc()
}

// CollectRequestTimeoutMetrics publishes the duration of timed-out requests by
// API path to 'github_request_timeouts' on prometheus.
func CollectRequestTimeoutMetrics(tokenHash, path, userAgent string, reqStartTime, responseTime time.Time) {
	timeoutDuration.With(prometheus.Labels{"token_hash": tokenHash, "path": simplifier.Simplify(path), "user_agent": userAgent}).Observe(float64(responseTime.Sub(reqStartTime).Seconds()))
}
<commit_msg>ghproxy: do not record zero-value rate-limits when empty<commit_after>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package ghmetrics

import (
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

// ghTokenUntilResetGaugeVec provides the 'github_token_reset' gauge that
// enables keeping track of GitHub reset times.
var ghTokenUntilResetGaugeVec = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "github_token_reset",
		Help: "Last reported GitHub token reset time.",
	},
	[]string{"token_hash", "api_version"},
)

// ghTokenUsageGaugeVec provides the 'github_token_usage' gauge that
// enables keeping track of GitHub calls and quotas.
var ghTokenUsageGaugeVec = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "github_token_usage",
		Help: "How many GitHub token requets are remaining for the current hour.",
	},
	[]string{"token_hash", "api_version"},
)

// ghRequestDurationHistVec provides the 'github_request_duration' histogram that keeps track
// of the duration of GitHub requests by API path.
var ghRequestDurationHistVec = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "github_request_duration",
		Help:    "GitHub request duration by API path.",
		Buckets: []float64{0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10},
	},
	[]string{"token_hash", "path", "status", "user_agent"},
)

// cacheCounter provides the 'ghcache_responses' counter vec that is indexed
// by the cache response mode.
var cacheCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "ghcache_responses",
		Help: "How many cache responses of each cache response mode there are.",
	},
	[]string{"mode", "path"},
)

// timeoutDuration provides the 'github_request_timeouts' histogram that keeps
// track of the timeouts of GitHub requests by API path.
var timeoutDuration = prometheus.NewHistogramVec(
	prometheus.HistogramOpts{
		Name:    "github_request_timeouts",
		Help:    "GitHub request timeout by API path.",
		Buckets: []float64{45, 60, 90, 120, 300},
	},
	[]string{"token_hash", "path", "user_agent"},
)

var muxTokenUsage, muxRequestMetrics sync.Mutex
var lastGitHubResponse time.Time

func init() {
	prometheus.MustRegister(ghTokenUntilResetGaugeVec)
	prometheus.MustRegister(ghTokenUsageGaugeVec)
	prometheus.MustRegister(ghRequestDurationHistVec)
	prometheus.MustRegister(cacheCounter)
	prometheus.MustRegister(timeoutDuration)
}

// CollectGitHubTokenMetrics publishes the rate limits of the github api to
// `github_token_usage` as well as `github_token_reset` on prometheus.
func CollectGitHubTokenMetrics(tokenHash, apiVersion string, headers http.Header, reqStartTime, responseTime time.Time) {
	remaining := headers.Get("X-RateLimit-Remaining")
	if remaining == "" {
		return
	}
	timeUntilReset := timestampStringToTime(headers.Get("X-RateLimit-Reset"))
	durationUntilReset := timeUntilReset.Sub(reqStartTime)

	remainingFloat, err := strconv.ParseFloat(remaining, 64)
	if err != nil {
		logrus.WithError(err).Infof("Couldn't convert number of remaining token requests into gauge value (float)")
	}
	if remainingFloat == 0 {
		logrus.WithFields(logrus.Fields{
			"header":     remaining,
			"user-agent": headers.Get("User-Agent"),
		}).Debug("Parsed GitHub header as indicating no remaining rate-limit.")
	}

	muxTokenUsage.Lock()
	isAfter := lastGitHubResponse.After(responseTime)
	if !isAfter {
		lastGitHubResponse = responseTime
	}
	muxTokenUsage.Unlock()
	if isAfter {
		logrus.WithField("last-github-response", lastGitHubResponse).WithField("response-time", responseTime).Debug("Previously pushed metrics of a newer response, skipping old metrics")
	} else {
		ghTokenUntilResetGaugeVec.With(prometheus.Labels{"token_hash": tokenHash, "api_version": apiVersion}).Set(float64(durationUntilReset.Nanoseconds()))
		ghTokenUsageGaugeVec.With(prometheus.Labels{"token_hash": tokenHash, "api_version": apiVersion}).Set(remainingFloat)
	}
}

// CollectGitHubRequestMetrics publishes the number of requests by API path to
// `github_requests` on prometheus.
func CollectGitHubRequestMetrics(tokenHash, path, statusCode, userAgent string, roundTripTime float64) {
	ghRequestDurationHistVec.With(prometheus.Labels{"token_hash": tokenHash, "path": simplifier.Simplify(path), "status": statusCode, "user_agent": userAgent}).Observe(roundTripTime)
}

// timestampStringToTime takes a unix timestamp and returns a `time.Time`
// from the given time.
func timestampStringToTime(tstamp string) time.Time {
	timestamp, err := strconv.ParseInt(tstamp, 10, 64)
	if err != nil {
		logrus.WithField("timestamp", tstamp).Info("Couldn't convert unix timestamp")
	}
	return time.Unix(timestamp, 0)
}

// CollectCacheRequestMetrics records a cache outcome for a specific path
func CollectCacheRequestMetrics(mode, path string) {
	cacheCounter.With(prometheus.Labels{"mode": mode, "path": simplifier.Simplify(path)}).Inc()
}

// CollectRequestTimeoutMetrics publishes the duration of timed-out requests by
// API path to 'github_request_timeouts' on prometheus.
func CollectRequestTimeoutMetrics(tokenHash, path, userAgent string, reqStartTime, responseTime time.Time) {
	timeoutDuration.With(prometheus.Labels{"token_hash": tokenHash, "path": simplifier.Simplify(path), "user_agent": userAgent}).Observe(float64(responseTime.Sub(reqStartTime).Seconds()))
}
<|endoftext|>


<commit_before>package pdfhandler

import (
	"archive/zip"
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

type PDF struct {
	FileName string            `json:"filename"`
	Fields   map[string]string `json:"fields"`
}

func (p PDF) render(rootPath string) ([]byte, error) {
	path := filepath.Join(rootPath, p.FileName)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return nil, err
	}
	tmpfile, err := ioutil.TempFile("", "example")
	if err != nil {
		return nil, err
	}
	defer os.Remove(tmpfile.Name()) // clean up
	if _, err := tmpfile.Write(mapToXFDF(p.Fields)); err != nil {
		return nil, err
	}
	if err := tmpfile.Close(); err != nil {
		return nil, err
	}
	cmd := exec.Command("pdftk", path, "fill_form", tmpfile.Name(), "output", "-")
	var out bytes.Buffer
	cmd.Stdout = &out
	var t bytes.Buffer
	cmd.Stderr = &t
	err = cmd.Run()
	if err != nil {
		return nil, errors.New(t.String())
	}
	return out.Bytes(), nil
}

type PDFHandler struct {
	FilePath string
}

func (ph PDFHandler) multi(mimetype string, pdfs []PDF, w http.ResponseWriter) error {
	dir, err := ioutil.TempDir("", "workpath")
	if err != nil {
		return err
	}
	defer os.RemoveAll(dir)
	var wg sync.WaitGroup
	for idx, p := range pdfs {
		wg.Add(1)
		go func(idx int, p PDF) {
			defer wg.Done()
			tmpfn := filepath.Join(dir, fmt.Sprintf("%d.pdf", idx))
			b, err := p.render(ph.FilePath)
			if err != nil {
				return
			}
			ioutil.WriteFile(tmpfn, b, 0777)
		}(idx, p)
	}
	wg.Wait()

	switch mimetype {
	case "application/pdf":
		cmd := exec.Command("pdftk")
		for idx := range pdfs {
			cmd.Args = append(cmd.Args, filepath.Join(dir, fmt.Sprintf("%d.pdf", idx)))
		}
		cmd.Args = append(cmd.Args, "cat")
		cmd.Args = append(cmd.Args, "output")
		cmd.Args = append(cmd.Args, "-")
		var out bytes.Buffer
		cmd.Stdout = &out
		var t bytes.Buffer
		cmd.Stderr = &t
		err = cmd.Run()
		if err != nil {
			return errors.New(t.String())
		}
		w.Write(out.Bytes())
		return nil
		break
	case "application/zip":
		zw := zip.NewWriter(w)
		for idx := range pdfs {
			f, err := zw.Create(fmt.Sprintf("%d.pdf", idx))
			if err != nil {
				return err
			}
			xx, err := os.Open(filepath.Join(dir, fmt.Sprintf("%d.pdf", idx)))
			if err != nil {
				return err
			}
			_, err = io.Copy(f, xx)
			if err != nil {
				return err
			}
		}
		err := zw.Close()
		if err != nil {
			return err
		}
		break
	}

	return nil
}

func (p PDFHandler) get(w http.ResponseWriter, req *http.Request) {
	files, err := ioutil.ReadDir(p.FilePath)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	ch := make(chan PDF)

	go func() {
		var wg sync.WaitGroup
		for _, file := range files {
			if !strings.HasSuffix(file.Name(), ".pdf") {
				continue
			}
			wg.Add(1)
			go func(fp string) {
				p, err := readFields(p.FilePath, fp)
				if err == nil {
					ch <- *p
				}
				defer wg.Done()
			}(file.Name())
		}
		wg.Wait()
		close(ch)
	}()

	pdfs := []PDF{}
	for p := range ch {
		pdfs = append(pdfs, p)
	}

	enc := json.NewEncoder(w)
	err = enc.Encode(pdfs)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

}

func (p PDFHandler) post(w http.ResponseWriter, req *http.Request) {
	ct := req.Header.Get("Content-Type")
	if ct != "application/json" {
		http.Error(w, "Invalid Content-Type", http.StatusNotAcceptable)
		return
	}

	ac := req.Header.Get("Accept")
	// validate ac

	r := bufio.NewReader(req.Body)
	dec := json.NewDecoder(r)
	ch, _ := r.Peek(1)
	switch string(ch) {
	case "{":
		var x PDF
		err := dec.Decode(&x)
		if err != nil {
			http.Error(w, err.Error(), http.StatusNotAcceptable)
			return
		}
		out, err := x.render(p.FilePath)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.Write(out)
		break
	case "[":
		var pdfs []PDF
		err := dec.Decode(&pdfs)
		if err != nil {
			http.Error(w, err.Error(), http.StatusNotAcceptable)
			return
		}
		err = p.multi(ac, pdfs, w)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		break
	default:
		http.Error(w, "Invalid input", http.StatusNotAcceptable)
		return
	}
	w.Header().Set("Content-Disposition", "attachment; filename=file.pdf")
	w.Header().Set("Content-Type", ac)
}

func (p PDFHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	switch req.Method {
	case "GET":
		p.get(w, req)
	case "POST":
		p.post(w, req)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
}
<commit_msg>Use channel for getting more info out of the job<commit_after>package pdfhandler

import (
	"archive/zip"
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

type PDF struct {
	FileName string            `json:"filename"`
	Fields   map[string]string `json:"fields"`
}

func (p PDF) render(rootPath string) ([]byte, error) {
	path := filepath.Join(rootPath, p.FileName)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return nil, err
	}
	tmpfile, err := ioutil.TempFile("", "example")
	if err != nil {
		return nil, err
	}
	defer os.Remove(tmpfile.Name()) // clean up
	if _, err := tmpfile.Write(mapToXFDF(p.Fields)); err != nil {
		return nil, err
	}
	if err := tmpfile.Close(); err != nil {
		return nil, err
	}
	cmd := exec.Command("pdftk", path, "fill_form", tmpfile.Name(), "output", "-")
	var out bytes.Buffer
	cmd.Stdout = &out
	var t bytes.Buffer
	cmd.Stderr = &t
	err = cmd.Run()
	if err != nil {
		return nil, errors.New(t.String())
	}
	return out.Bytes(), nil
}

type PDFHandler struct {
	FilePath string
}

func (ph PDFHandler) multi(mimetype string, pdfs []PDF, w http.ResponseWriter) error {
	dir, err := ioutil.TempDir("", "workpath")
	if err != nil {
		return err
	}
	defer os.RemoveAll(dir)

	type job struct {
		Pdf  PDF
		File string
		Body []byte
	}
	ch := make(chan job)

	go func() {
		var wg sync.WaitGroup
		for idx, p := range pdfs {
			wg.Add(1)
			go func(idx int, p PDF) {
				defer wg.Done()
				tmpfn := filepath.Join(dir, fmt.Sprintf("%d.pdf", idx))
				b, err := p.render(ph.FilePath)
				if err != nil {
					return
				}
				ioutil.WriteFile(tmpfn, b, 0777)
				ch <- job{p, tmpfn, b}
			}(idx, p)
		}
		wg.Wait()
		close(ch)
	}()

	switch mimetype {
	case "application/pdf":
		cmd := exec.Command("pdftk")
		for j := range ch {
			cmd.Args = append(cmd.Args, j.File)
		}
		cmd.Args = append(cmd.Args, "cat")
		cmd.Args = append(cmd.Args, "output")
		cmd.Args = append(cmd.Args, "-")
		var out bytes.Buffer
		cmd.Stdout = &out
		var t bytes.Buffer
		cmd.Stderr = &t
		err = cmd.Run()
		if err != nil {
			return errors.New(t.String())
		}
		w.Write(out.Bytes())
		return nil
		break
	case "application/zip":
		zw := zip.NewWriter(w)
		for j := range ch {
			f, err := zw.Create(j.Pdf.FileName)
			if err != nil {
				return err
			}
			_, err = f.Write(j.Body)
			if err != nil {
				return err
			}
		}
		err := zw.Close()
		if err != nil {
			return err
		}
		break
	}

	return nil
}

func (p PDFHandler) get(w http.ResponseWriter, req *http.Request) {
	files, err := ioutil.ReadDir(p.FilePath)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	ch := make(chan PDF)

	go func() {
		var wg sync.WaitGroup
		for _, file := range files {
			if !strings.HasSuffix(file.Name(), ".pdf") {
				continue
			}
			wg.Add(1)
			go func(fp string) {
				p, err := readFields(p.FilePath, fp)
				if err == nil {
					ch <- *p
				}
				defer wg.Done()
			}(file.Name())
		}
		wg.Wait()
		close(ch)
	}()

	pdfs := []PDF{}
	for p := range ch {
		pdfs = append(pdfs, p)
	}

	enc := json.NewEncoder(w)
	err = enc.Encode(pdfs)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

}

func (p PDFHandler) post(w http.ResponseWriter, req *http.Request) {
	ct := req.Header.Get("Content-Type")
	if ct != "application/json" {
		http.Error(w, "Invalid Content-Type", http.StatusNotAcceptable)
		return
	}

	ac := req.Header.Get("Accept")
	// validate ac

	r := bufio.NewReader(req.Body)
	dec := json.NewDecoder(r)
	ch, _ := r.Peek(1)
	switch string(ch) {
	case "{":
		var x PDF
		err := dec.Decode(&x)
		if err != nil {
			http.Error(w, err.Error(), http.StatusNotAcceptable)
			return
		}
		out, err := x.render(p.FilePath)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.Write(out)
		break
	case "[":
		var pdfs []PDF
		err := dec.Decode(&pdfs)
		if err != nil {
			http.Error(w, err.Error(), http.StatusNotAcceptable)
			return
		}
		err = p.multi(ac, pdfs, w)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		break
	default:
		http.Error(w, "Invalid input", http.StatusNotAcceptable)
		return
	}
	w.Header().Set("Content-Disposition", "attachment; filename=file.pdf")
	w.Header().Set("Content-Type", ac)
}

func (p PDFHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	switch req.Method {
	case "GET":
		p.get(w, req)
	case "POST":
		p.post(w, req)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

// sessionGaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// sessionGaugeStart and calls the returned function in a defer statement.
func sessionGaugeStart(session *sockjs.Session) (sessionGaugeEnd func()) {
	log.Debug("Client connected: %v", session.Tag)
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge := logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", session.Tag)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	socketId := randomString()
	session.Tag = socketId

	sessionGaugeEnd := sessionGaugeStart(session)
	defer sessionGaugeEnd()

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<commit_msg>broker: gauges should be global vars<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

// randomString() returns a new 16 char length random string
func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

// sessionGaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// sessionGaugeStart and calls the returned function in a defer statement.
func sessionGaugeStart(session *sockjs.Session) (sessionGaugeEnd func()) {
	log.Debug("Client connected: %v", session.Tag)
	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", session.Tag)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	socketId := randomString()
	session.Tag = socketId

	sessionGaugeEnd := sessionGaugeStart(session)
	defer sessionGaugeEnd()

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/letsencrypt/boulder/cmd"
)

type Config struct {
	// Execution plan parameters
	Plan struct {
		Actions   []string // things to do
		Rate      int64    // requests / s
		RateDelta string   // requests / s^2
		MaxRate   int64    // XXX: is this needed?
		Runtime   string   // how long to run for
	}
	ExternalState string   // path to file to load/save registrations etc to/from
	DontSaveState bool     // don't save changes to external state
	APIBase       string   // ACME API address to send requests to
	DomainBase    string   // base domain name to create authorizations for
	ChallTypes    []string // which challenges to complete, empty means use all
	HTTPOneAddr   string   // address to listen for http-01 validation requests on
	TLSOneAddr    string   // address to listen for tls-sni-01 validation requests on
	RealIP        string   // value of the Real-IP header to use when bypassing CDN
	CertKeySize   int      // size of the key to use when creating CSRs
	RegEmail      string   // email to use in registrations
	Results       string   // path to save metrics to
}

func main() {
	configPath := flag.String("config", "", "Path to configuration file for WFE load-generator")
	resultsPath := flag.String("results", "", "Path to latency results file")
	rateArg := flag.Int("rate", 0, "")
	runtimeArg := flag.String("runtime", "", "")
	deltaArg := flag.String("delta", "", "")
	flag.Parse()

	configBytes, err := ioutil.ReadFile(*configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read wfe config file %q: %s\n", &configPath, err)
		os.Exit(1)
	}
	var config Config
	err = json.Unmarshal(configBytes, &config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to parse wfe config file: %s\n", err)
		os.Exit(1)
	}

	if *resultsPath != "" {
		config.Results = *resultsPath
	}
	if *rateArg != 0 {
		config.Plan.Rate = int64(*rateArg)
	}
	if *runtimeArg != "" {
		config.Plan.Runtime = *runtimeArg
	}
	if *deltaArg != "" {
		config.Plan.RateDelta = *deltaArg
	}

	s, err := New(
		config.APIBase,
		config.CertKeySize,
		config.DomainBase,
		config.RealIP,
		config.MaxRegs,
		config.Results,
		config.RegEmail,
		config.Plan.Actions,
	)
	cmd.FailOnError(err, "Failed to create WFE generator")

	if config.ExternalState != "" {
		err = s.Restore(config.ExternalState)
		cmd.FailOnError(err, "Failed to load registration snapshot")
	}

	runtime, err := time.ParseDuration(config.Plan.Runtime)
	cmd.FailOnError(err, "Failed to parse plan runtime")

	var delta *RateDelta
	if config.Plan.RateDelta != "" {
		parts := strings.Split(config.Plan.RateDelta, "/")
		if len(parts) != 2 {
			fmt.Fprintf(os.Stderr, "RateDelta is malformed")
			os.Exit(1)
		}
		rate, err := strconv.Atoi(parts[0])
		cmd.FailOnError(err, "Failed to parse increase portion of RateDelta")
		period, err := time.ParseDuration(parts[1])
		cmd.FailOnError(err, "Failed to parse period portion of RateDelta")
		delta = &RateDelta{Inc: int64(rate), Period: period}
	}

	err = s.Run(config.HTTPOneAddr, config.TLSOneAddr, Plan{
		Runtime: runtime,
		Rate:    config.Plan.Rate,
		Delta:   delta,
	})
	cmd.FailOnError(err, "Failed to run WFE load generator")

	if config.ExternalState != "" && !config.DontSaveState {
		err = s.Snapshot(config.ExternalState)
		cmd.FailOnError(err, "Failed to save registration snapshot")
	}

	fmt.Println("[+] All done, bye bye ^_^")
}
<commit_msg>Whoops, delete correct field<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/letsencrypt/boulder/cmd"
)

type Config struct {
	// Execution plan parameters
	Plan struct {
		Actions   []string // things to do
		Rate      int64    // requests / s
		RateDelta string   // requests / s^2
		Runtime   string   // how long to run for
	}
	ExternalState string   // path to file to load/save registrations etc to/from
	DontSaveState bool     // don't save changes to external state
	APIBase       string   // ACME API address to send requests to
	DomainBase    string   // base domain name to create authorizations for
	ChallTypes    []string // which challenges to complete, empty means use all
	HTTPOneAddr   string   // address to listen for http-01 validation requests on
	TLSOneAddr    string   // address to listen for tls-sni-01 validation requests on
	RealIP        string   // value of the Real-IP header to use when bypassing CDN
	CertKeySize   int      // size of the key to use when creating CSRs
	RegEmail      string   // email to use in registrations
	Results       string   // path to save metrics to
	MaxRegs       int      // maximum number of registrations to create
}

func main() {
	configPath := flag.String("config", "", "Path to configuration file for WFE load-generator")
	resultsPath := flag.String("results", "", "Path to latency results file")
	rateArg := flag.Int("rate", 0, "")
	runtimeArg := flag.String("runtime", "", "")
	deltaArg := flag.String("delta", "", "")
	flag.Parse()

	configBytes, err := ioutil.ReadFile(*configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read wfe config file %q: %s\n", &configPath, err)
		os.Exit(1)
	}
	var config Config
	err = json.Unmarshal(configBytes, &config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to parse wfe config file: %s\n", err)
		os.Exit(1)
	}

	if *resultsPath != "" {
		config.Results = *resultsPath
	}
	if *rateArg != 0 {
		config.Plan.Rate = int64(*rateArg)
	}
	if *runtimeArg != "" {
		config.Plan.Runtime = *runtimeArg
	}
	if *deltaArg != "" {
		config.Plan.RateDelta = *deltaArg
	}

	s, err := New(
		config.APIBase,
		config.CertKeySize,
		config.DomainBase,
		config.RealIP,
		config.MaxRegs,
		config.Results,
		config.RegEmail,
		config.Plan.Actions,
	)
	cmd.FailOnError(err, "Failed to create WFE generator")

	if config.ExternalState != "" {
		err = s.Restore(config.ExternalState)
		cmd.FailOnError(err, "Failed to load registration snapshot")
	}

	runtime, err := time.ParseDuration(config.Plan.Runtime)
	cmd.FailOnError(err, "Failed to parse plan runtime")

	var delta *RateDelta
	if config.Plan.RateDelta != "" {
		parts := strings.Split(config.Plan.RateDelta, "/")
		if len(parts) != 2 {
			fmt.Fprintf(os.Stderr, "RateDelta is malformed")
			os.Exit(1)
		}
		rate, err := strconv.Atoi(parts[0])
		cmd.FailOnError(err, "Failed to parse increase portion of RateDelta")
		period, err := time.ParseDuration(parts[1])
		cmd.FailOnError(err, "Failed to parse period portion of RateDelta")
		delta = &RateDelta{Inc: int64(rate), Period: period}
	}

	err = s.Run(config.HTTPOneAddr, config.TLSOneAddr, Plan{
		Runtime: runtime,
		Rate:    config.Plan.Rate,
		Delta:   delta,
	})
	cmd.FailOnError(err, "Failed to run WFE load generator")

	if config.ExternalState != "" && !config.DontSaveState {
		err = s.Snapshot(config.ExternalState)
		cmd.FailOnError(err, "Failed to save registration snapshot")
	}

	fmt.Println("[+] All done, bye bye ^_^")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"runtime"
	"strings"

	cnc "github.com/noironetworks/cilium-net/common/cilium-net-client"
	ciliumtype "github.com/noironetworks/cilium-net/common/types"

	log "github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ip"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ipam"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ns"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/skel"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/types"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/vishvananda/netlink"
)

const (
	HostInterfacePrefix      = "lxc"
	TemporaryInterfacePrefix = "tmp"
)

type NetConf struct {
	types.NetConf
	NodeIP net.IP `json:"NodeIP"`
	MTU    int    `json:"mtu"`
}

func init() {
	log.SetLevel(log.DebugLevel)
	log.SetOutput(os.Stderr)
	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel)
}

func loadNetConf(bytes []byte) (*NetConf, error) {
	n := &NetConf{}
	if err := json.Unmarshal(bytes, n); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %s", err)
	}
	return n, nil
}

func endpoint2ifname(endpointID string) string {
	return HostInterfacePrefix + endpointID[:5]
}

func removeIfFromNSIfExists(netns *os.File, ifName string) error {
	return ns.WithNetNS(netns, false, func(_ *os.File) error {
		l, err := netlink.LinkByName(ifName)
		if err != nil {
			if strings.Contains(err.Error(), "Link not found") {
				return nil
			} else {
				return err
			}
		}
		return netlink.LinkDel(l)
	})
}

func setupVeth(netNsFile *os.File, ifName, containerID string, mtu int, ep *ciliumtype.Endpoint) (*netlink.Veth, error) {

	lxcIfname := endpoint2ifname(containerID)
	tmpIfname := TemporaryInterfacePrefix + containerID[:5]

	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: lxcIfname},
		PeerName:  tmpIfname,
	}

	if err := netlink.LinkAdd(veth); err != nil {
		return nil, fmt.Errorf("unable to create veth pair: %s", err)
	}
	var err error
	defer func() {
		if err != nil {
			if err = netlink.LinkDel(veth); err != nil {
				log.Warnf("failed to clean up veth %q: %s", veth.Name, err)
			}
		}
	}()

	peer, err := netlink.LinkByName(tmpIfname)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup veth peer just created: %s", err)
	}

	if err = netlink.LinkSetMTU(peer, mtu); err != nil {
		return nil, fmt.Errorf("unable to set MTU to %q: %s", tmpIfname, err)
	}

	hostVeth, err := netlink.LinkByName(lxcIfname)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup veth just created: %s", err)
	}

	if err = netlink.LinkSetMTU(hostVeth, mtu); err != nil {
		return nil, fmt.Errorf("unable to set MTU to %q: %s", lxcIfname, err)
	}

	if err = netlink.LinkSetUp(veth); err != nil {
		return nil, fmt.Errorf("unable to bring up veth pair: %s", err)
	}

	ep.LxcMAC = peer.Attrs().HardwareAddr
	ep.NodeMAC = hostVeth.Attrs().HardwareAddr
	ep.Ifname = lxcIfname

	if err := netlink.LinkSetNsFd(peer, int(netNsFile.Fd())); err != nil {
		return nil, fmt.Errorf("unable to move veth pair %q to netns: %s", peer, err)
	}

	err = ns.WithNetNS(netNsFile, false, func(_ *os.File) error {
		err := renameLink(tmpIfname, ifName)
		if err != nil {
			return fmt.Errorf("failed to rename %q to %q: %s", tmpIfname, ifName, err)
		}
		return nil
	})
	return veth, err
}

func renameLink(curName, newName string) error {
	link, err := netlink.LinkByName(curName)
	if err != nil {
		return err
	}

	return netlink.LinkSetName(link, newName)
}

func cmdAdd(args *skel.CmdArgs) error {
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	c, err := cnc.NewDefaultClient()
	if err != nil {
		return fmt.Errorf("error while starting cilium-client: %s", err)
	}

	netNsFile, err := os.Open(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %s", args.Netns, err)
	}
	defer netNsFile.Close()

	if err := removeIfFromNSIfExists(netNsFile, args.IfName); err != nil {
		return fmt.Errorf("failed removing interface %q from namespace %q: %s",
			args.IfName, args.Netns, err)
	}

	var ep ciliumtype.Endpoint
	veth, err := setupVeth(netNsFile, args.IfName, args.ContainerID, n.MTU, &ep)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			if err = netlink.LinkDel(veth); err != nil {
				log.Warnf("failed to clean up veth %q: %s", veth.Name, err)
			}
		}
	}()

	result, err := ipam.ExecAdd(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			os.Setenv("CNI_COMMAND", "DEL")
			if err = ipam.ExecDel(n.IPAM.Type, args.StdinData); err != nil {
				log.Warnf("failed to release allocated IP %q: %s", result, err)
			}
			os.Setenv("CNI_COMMAND", "ADD")
		}
	}()

	err = ns.WithNetNSPath(args.Netns, false, func(_ *os.File) error {
		return ipam.ConfigureIface(args.IfName, result)
	})
	if err != nil {
		return err
	}

	if result.IP6 == nil {
		err = fmt.Errorf("result.IP6 is nil and it shouldn't be")
		return err
	}
	ep.LxcIP = result.IP6.IP.IP
	ep.NodeIP = n.NodeIP
	ep.SetID()
	if err := c.EndpointJoin(ep); err != nil {
		return fmt.Errorf("unable to create eBPF map: %s", err)
	}

	return result.Print()
}

func cmdDel(args *skel.CmdArgs) error {
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	c, err := cnc.NewDefaultClient()
	if err != nil {
		return fmt.Errorf("error while starting cilium-client: %s", err)
	}

	var containerIP net.IP
	// We need to retrieve the IPv6 address somehow...
	ns.WithNetNSPath(args.Netns, false, func(hostNS *os.File) error {
		l, err := netlink.LinkByName(args.IfName)
		if err != nil {
			return err
		}
		addrs, err := netlink.AddrList(l, netlink.FAMILY_V6)
		if err != nil {
			return err
		}
		log.Debugf("IPv6 addresses found %+v\n", addrs)

		// As long the nodeIP is address 0...
		for _, addr := range addrs {
			if bytes.Compare(n.NodeIP, addr.IPNet.IP.Mask(addr.IPNet.Mask)) == 0 {
				containerIP = addr.IP
				log.Debug("Container IP found ", containerIP)
				break
			}
		}
		return nil
	})

	err = ipam.ExecDel(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}

	var ep ciliumtype.Endpoint
	ep.LxcIP = containerIP
	ep.SetID()
	if err := c.EndpointLeave(ep.ID); err != nil {
		log.Warnf("leaving the endpoint failed: %s\n", err)
	}

	return ns.WithNetNSPath(args.Netns, false, func(hostNS *os.File) error {
		return ip.DelLinkByName(args.IfName)
	})
}
<commit_msg>Added a better IPv6 mask to the k8s-cni DEL<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"runtime"
	"strings"

	cnc "github.com/noironetworks/cilium-net/common/cilium-net-client"
	ciliumtype "github.com/noironetworks/cilium-net/common/types"
	"github.com/noironetworks/cilium-net/common"

	log "github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/Sirupsen/logrus"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ip"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ipam"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/ns"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/skel"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/appc/cni/pkg/types"
	"github.com/noironetworks/cilium-net/Godeps/_workspace/src/github.com/vishvananda/netlink"
)

const (
	HostInterfacePrefix      = "lxc"
	TemporaryInterfacePrefix = "tmp"
)

type NetConf struct {
	types.NetConf
	NodeIP net.IP `json:"NodeIP"`
	MTU    int    `json:"mtu"`
}

func init() {
	log.SetLevel(log.DebugLevel)
	log.SetOutput(os.Stderr)
	// this ensures that main runs only on main thread (thread group leader).
	// since namespace ops (unshare, setns) are done for a single thread, we
	// must ensure that the goroutine does not jump from OS thread to thread
	runtime.LockOSThread()
}

func main() {
	skel.PluginMain(cmdAdd, cmdDel)
}

func loadNetConf(bytes []byte) (*NetConf, error) {
	n := &NetConf{}
	if err := json.Unmarshal(bytes, n); err != nil {
		return nil, fmt.Errorf("failed to load netconf: %s", err)
	}
	return n, nil
}

func endpoint2ifname(endpointID string) string {
	return HostInterfacePrefix + endpointID[:5]
}

func removeIfFromNSIfExists(netns *os.File, ifName string) error {
	return ns.WithNetNS(netns, false, func(_ *os.File) error {
		l, err := netlink.LinkByName(ifName)
		if err != nil {
			if strings.Contains(err.Error(), "Link not found") {
				return nil
			} else {
				return err
			}
		}
		return netlink.LinkDel(l)
	})
}

func setupVeth(netNsFile *os.File, ifName, containerID string, mtu int, ep *ciliumtype.Endpoint) (*netlink.Veth, error) {

	lxcIfname := endpoint2ifname(containerID)
	tmpIfname := TemporaryInterfacePrefix + containerID[:5]

	veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: lxcIfname},
		PeerName:  tmpIfname,
	}

	if err := netlink.LinkAdd(veth); err != nil {
		return nil, fmt.Errorf("unable to create veth pair: %s", err)
	}
	var err error
	defer func() {
		if err != nil {
			if err = netlink.LinkDel(veth); err != nil {
				log.Warnf("failed to clean up veth %q: %s", veth.Name, err)
			}
		}
	}()

	peer, err := netlink.LinkByName(tmpIfname)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup veth peer just created: %s", err)
	}

	if err = netlink.LinkSetMTU(peer, mtu); err != nil {
		return nil, fmt.Errorf("unable to set MTU to %q: %s", tmpIfname, err)
	}

	hostVeth, err := netlink.LinkByName(lxcIfname)
	if err != nil {
		return nil, fmt.Errorf("unable to lookup veth just created: %s", err)
	}

	if err = netlink.LinkSetMTU(hostVeth, mtu); err != nil {
		return nil, fmt.Errorf("unable to set MTU to %q: %s", lxcIfname, err)
	}

	if err = netlink.LinkSetUp(veth); err != nil {
		return nil, fmt.Errorf("unable to bring up veth pair: %s", err)
	}

	ep.LxcMAC = peer.Attrs().HardwareAddr
	ep.NodeMAC = hostVeth.Attrs().HardwareAddr
	ep.Ifname = lxcIfname

	if err := netlink.LinkSetNsFd(peer, int(netNsFile.Fd())); err != nil {
		return nil, fmt.Errorf("unable to move veth pair %q to netns: %s", peer, err)
	}

	err = ns.WithNetNS(netNsFile, false, func(_ *os.File) error {
		err := renameLink(tmpIfname, ifName)
		if err != nil {
			return fmt.Errorf("failed to rename %q to %q: %s", tmpIfname, ifName, err)
		}
		return nil
	})
	return veth, err
}

func renameLink(curName, newName string) error {
	link, err := netlink.LinkByName(curName)
	if err != nil {
		return err
	}

	return netlink.LinkSetName(link, newName)
}

func cmdAdd(args *skel.CmdArgs) error {
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	c, err := cnc.NewDefaultClient()
	if err != nil {
		return fmt.Errorf("error while starting cilium-client: %s", err)
	}

	netNsFile, err := os.Open(args.Netns)
	if err != nil {
		return fmt.Errorf("failed to open netns %q: %s", args.Netns, err)
	}
	defer netNsFile.Close()

	if err := removeIfFromNSIfExists(netNsFile, args.IfName); err != nil {
		return fmt.Errorf("failed removing interface %q from namespace %q: %s",
			args.IfName, args.Netns, err)
	}

	var ep ciliumtype.Endpoint
	veth, err := setupVeth(netNsFile, args.IfName, args.ContainerID, n.MTU, &ep)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			if err = netlink.LinkDel(veth); err != nil {
				log.Warnf("failed to clean up veth %q: %s", veth.Name, err)
			}
		}
	}()

	result, err := ipam.ExecAdd(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			os.Setenv("CNI_COMMAND", "DEL")
			if err = ipam.ExecDel(n.IPAM.Type, args.StdinData); err != nil {
				log.Warnf("failed to release allocated IP %q: %s", result, err)
			}
			os.Setenv("CNI_COMMAND", "ADD")
		}
	}()

	err = ns.WithNetNSPath(args.Netns, false, func(_ *os.File) error {
		return ipam.ConfigureIface(args.IfName, result)
	})
	if err != nil {
		return err
	}

	if result.IP6 == nil {
		err = fmt.Errorf("result.IP6 is nil and it shouldn't be")
		return err
	}
	ep.LxcIP = result.IP6.IP.IP
	ep.NodeIP = n.NodeIP
	ep.SetID()
	if err := c.EndpointJoin(ep); err != nil {
		return fmt.Errorf("unable to create eBPF map: %s", err)
	}

	return result.Print()
}

func cmdDel(args *skel.CmdArgs) error {
	n, err := loadNetConf(args.StdinData)
	if err != nil {
		return err
	}

	c, err := cnc.NewDefaultClient()
	if err != nil {
		return fmt.Errorf("error while starting cilium-client: %s", err)
	}

	var containerIP net.IP
	// We need to retrieve the IPv6 address somehow...
	ns.WithNetNSPath(args.Netns, false, func(hostNS *os.File) error {
		l, err := netlink.LinkByName(args.IfName)
		if err != nil {
			return err
		}
		addrs, err := netlink.AddrList(l, netlink.FAMILY_V6)
		if err != nil {
			return err
		}
		log.Debugf("IPv6 addresses found %+v\n", addrs)

		// As long the nodeIP is address 0...
		for _, addr := range addrs {
			if bytes.Compare(n.NodeIP, addr.IP.Mask(common.NodeIPv6Mask)) == 0 {
				containerIP = addr.IP
				log.Debug("Container IP found ", containerIP)
				break
			}
		}
		return nil
	})

	err = ipam.ExecDel(n.IPAM.Type, args.StdinData)
	if err != nil {
		return err
	}

	var ep ciliumtype.Endpoint
	ep.LxcIP = containerIP
	ep.SetID()
	if err := c.EndpointLeave(ep.ID); err != nil {
		log.Warnf("leaving the endpoint failed: %s\n", err)
	}

	return ns.WithNetNSPath(args.Netns, false, func(hostNS *os.File) error {
		return ip.DelLinkByName(args.IfName)
	})
}
<|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 log

// These tests are too simple.

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strings"
	"testing"
	"time"
)

const (
	Rdate         = `[0-9][0-9][0-9][0-9]/[0-9][0-9]/[0-9][0-9]`
	Rtime         = `[0-9][0-9]:[0-9][0-9]:[0-9][0-9]`
	Rmicroseconds = `\.[0-9][0-9][0-9][0-9][0-9][0-9]`
	Rline         = `(57|59):` // must update if the calls to l.Printf / l.Print below move
	Rlongfile     = `.*/[A-Za-z0-9_\-]+\.go:` + Rline
	Rshortfile    = `[A-Za-z0-9_\-]+\.go:` + Rline
)

type tester struct {
	flag    int
	prefix  string
	pattern string // regexp that log output must match; we add ^ and expected_text$ always
}

var tests = []tester{
	// individual pieces:
	{0, "", ""},
	{0, "XXX", "XXX"},
	{Ldate, "", Rdate + " "},
	{Ltime, "", Rtime + " "},
	{Ltime | Lmicroseconds, "", Rtime + Rmicroseconds + " "},
	{Lmicroseconds, "", Rtime + Rmicroseconds + " "}, // microsec implies time
	{Llongfile, "", Rlongfile + " "},
	{Lshortfile, "", Rshortfile + " "},
	{Llongfile | Lshortfile, "", Rshortfile + " "}, // shortfile overrides longfile
	// everything at once:
	{Ldate | Ltime | Lmicroseconds | Llongfile, "XXX", "XXX" + Rdate + " " + Rtime + Rmicroseconds + " " + Rlongfile + " "},
	{Ldate | Ltime | Lmicroseconds | Lshortfile, "XXX", "XXX" + Rdate + " " + Rtime + Rmicroseconds + " " + Rshortfile + " "},
}

// Test using Println("hello", 23, "world") or using Printf("hello %d world", 23)
func testPrint(t *testing.T, flag int, prefix string, pattern string, useFormat bool) {
	buf := new(bytes.Buffer)
	SetOutput(buf)
	SetFlags(flag)
	SetPrefix(prefix)
	if useFormat {
		Printf("hello %d world", 23)
	} else {
		Println("hello", 23, "world")
	}
	line := buf.String()
	line = line[0 : len(line)-1]
	pattern = "^" + pattern + "hello 23 world$"
	matched, err4 := regexp.MatchString(pattern, line)
	if err4 != nil {
		t.Fatal("pattern did not compile:", err4)
	}
	if !matched {
		t.Errorf("log output should match %q is %q", pattern, line)
	}
	SetOutput(os.Stderr)
}

func TestAll(t *testing.T) {
	for _, testcase := range tests {
		testPrint(t, testcase.flag, testcase.prefix, testcase.pattern, false)
		testPrint(t, testcase.flag, testcase.prefix, testcase.pattern, true)
	}
}

func TestOutput(t *testing.T) {
	const testString = "test"
	var b bytes.Buffer
	l := New(&b, "", 0)
	l.Println(testString)
	if expect := testString + "\n"; b.String() != expect {
		t.Errorf("log output should match %q is %q", expect, b.String())
	}
}

func TestOutputRace(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "", 0)
	for i := 0; i < 100; i++ {
		go func() {
			l.SetFlags(0)
		}()
		l.Output(0, "")
	}
}

func TestFlagAndPrefixSetting(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Test:", LstdFlags)
	f := l.Flags()
	if f != LstdFlags {
		t.Errorf("Flags 1: expected %x got %x", LstdFlags, f)
	}
	l.SetFlags(f | Lmicroseconds)
	f = l.Flags()
	if f != LstdFlags|Lmicroseconds {
		t.Errorf("Flags 2: expected %x got %x", LstdFlags|Lmicroseconds, f)
	}
	p := l.Prefix()
	if p != "Test:" {
		t.Errorf(`Prefix: expected "Test:" got %q`, p)
	}
	l.SetPrefix("Reality:")
	p = l.Prefix()
	if p != "Reality:" {
		t.Errorf(`Prefix: expected "Reality:" got %q`, p)
	}
	// Verify a log message looks right, with our prefix and microseconds present.
	l.Print("hello")
	pattern := "^Reality:" + Rdate + " " + Rtime + Rmicroseconds + " hello\n"
	matched, err := regexp.Match(pattern, b.Bytes())
	if err != nil {
		t.Fatalf("pattern %q did not compile: %s", pattern, err)
	}
	if !matched {
		t.Error("message did not match pattern")
	}
}

func TestUTCFlag(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Test:", LstdFlags)
	l.SetFlags(Ldate | Ltime | LUTC)
	// Verify a log message looks right in the right time zone. Quantize to the second only.
	now := time.Now().UTC()
	l.Print("hello")
	want := fmt.Sprintf("Test:%d/%.2d/%.2d %.2d:%.2d:%.2d hello\n",
		now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
	got := b.String()
	if got == want {
		return
	}
	// It's possible we crossed a second boundary between getting now and logging,
	// so add a second and try again. This should very nearly always work.
	now = now.Add(time.Second)
	want = fmt.Sprintf("Test:%d/%.2d/%.2d %.2d:%.2d:%.2d hello\n",
		now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
	if got == want {
		return
	}
	t.Errorf("got %q; want %q", got, want)
}

func TestEmptyPrintCreatesLine(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Header:", LstdFlags)
	l.Print()
	l.Println("non-empty")
	output := b.String()
	if n := strings.Count(output, "Header"); n != 2 {
		t.Errorf("expected 2 headers, got %d", n)
	}
	if n := strings.Count(output, "\n"); n != 2 {
		t.Errorf("expected 2 lines, got %d", n)
	}
}

func BenchmarkItoa(b *testing.B) {
	dst := make([]byte, 0, 64)
	for i := 0; i < b.N; i++ {
		dst = dst[0:0]
		itoa(&dst, 2015, 4)   // year
		itoa(&dst, 1, 2)      // month
		itoa(&dst, 30, 2)     // day
		itoa(&dst, 12, 2)     // hour
		itoa(&dst, 56, 2)     // minute
		itoa(&dst, 0, 2)      // second
		itoa(&dst, 987654, 6) // microsecond
	}
}

func BenchmarkPrintln(b *testing.B) {
	const testString = "test"
	var buf bytes.Buffer
	l := New(&buf, "", LstdFlags)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		l.Println(testString)
	}
}

func BenchmarkPrintlnNoFlags(b *testing.B) {
	const testString = "test"
	var buf bytes.Buffer
	l := New(&buf, "", 0)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		l.Println(testString)
	}
}
<commit_msg>log: make the name of error clearer<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 log

// These tests are too simple.

import (
	"bytes"
	"fmt"
	"os"
	"regexp"
	"strings"
	"testing"
	"time"
)

const (
	Rdate         = `[0-9][0-9][0-9][0-9]/[0-9][0-9]/[0-9][0-9]`
	Rtime         = `[0-9][0-9]:[0-9][0-9]:[0-9][0-9]`
	Rmicroseconds = `\.[0-9][0-9][0-9][0-9][0-9][0-9]`
	Rline         = `(57|59):` // must update if the calls to l.Printf / l.Print below move
	Rlongfile     = `.*/[A-Za-z0-9_\-]+\.go:` + Rline
	Rshortfile    = `[A-Za-z0-9_\-]+\.go:` + Rline
)

type tester struct {
	flag    int
	prefix  string
	pattern string // regexp that log output must match; we add ^ and expected_text$ always
}

var tests = []tester{
	// individual pieces:
	{0, "", ""},
	{0, "XXX", "XXX"},
	{Ldate, "", Rdate + " "},
	{Ltime, "", Rtime + " "},
	{Ltime | Lmicroseconds, "", Rtime + Rmicroseconds + " "},
	{Lmicroseconds, "", Rtime + Rmicroseconds + " "}, // microsec implies time
	{Llongfile, "", Rlongfile + " "},
	{Lshortfile, "", Rshortfile + " "},
	{Llongfile | Lshortfile, "", Rshortfile + " "}, // shortfile overrides longfile
	// everything at once:
	{Ldate | Ltime | Lmicroseconds | Llongfile, "XXX", "XXX" + Rdate + " " + Rtime + Rmicroseconds + " " + Rlongfile + " "},
	{Ldate | Ltime | Lmicroseconds | Lshortfile, "XXX", "XXX" + Rdate + " " + Rtime + Rmicroseconds + " " + Rshortfile + " "},
}

// Test using Println("hello", 23, "world") or using Printf("hello %d world", 23)
func testPrint(t *testing.T, flag int, prefix string, pattern string, useFormat bool) {
	buf := new(bytes.Buffer)
	SetOutput(buf)
	SetFlags(flag)
	SetPrefix(prefix)
	if useFormat {
		Printf("hello %d world", 23)
	} else {
		Println("hello", 23, "world")
	}
	line := buf.String()
	line = line[0 : len(line)-1]
	pattern = "^" + pattern + "hello 23 world$"
	matched, err := regexp.MatchString(pattern, line)
	if err != nil {
		t.Fatal("pattern did not compile:", err)
	}
	if !matched {
		t.Errorf("log output should match %q is %q", pattern, line)
	}
	SetOutput(os.Stderr)
}

func TestAll(t *testing.T) {
	for _, testcase := range tests {
		testPrint(t, testcase.flag, testcase.prefix, testcase.pattern, false)
		testPrint(t, testcase.flag, testcase.prefix, testcase.pattern, true)
	}
}

func TestOutput(t *testing.T) {
	const testString = "test"
	var b bytes.Buffer
	l := New(&b, "", 0)
	l.Println(testString)
	if expect := testString + "\n"; b.String() != expect {
		t.Errorf("log output should match %q is %q", expect, b.String())
	}
}

func TestOutputRace(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "", 0)
	for i := 0; i < 100; i++ {
		go func() {
			l.SetFlags(0)
		}()
		l.Output(0, "")
	}
}

func TestFlagAndPrefixSetting(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Test:", LstdFlags)
	f := l.Flags()
	if f != LstdFlags {
		t.Errorf("Flags 1: expected %x got %x", LstdFlags, f)
	}
	l.SetFlags(f | Lmicroseconds)
	f = l.Flags()
	if f != LstdFlags|Lmicroseconds {
		t.Errorf("Flags 2: expected %x got %x", LstdFlags|Lmicroseconds, f)
	}
	p := l.Prefix()
	if p != "Test:" {
		t.Errorf(`Prefix: expected "Test:" got %q`, p)
	}
	l.SetPrefix("Reality:")
	p = l.Prefix()
	if p != "Reality:" {
		t.Errorf(`Prefix: expected "Reality:" got %q`, p)
	}
	// Verify a log message looks right, with our prefix and microseconds present.
	l.Print("hello")
	pattern := "^Reality:" + Rdate + " " + Rtime + Rmicroseconds + " hello\n"
	matched, err := regexp.Match(pattern, b.Bytes())
	if err != nil {
		t.Fatalf("pattern %q did not compile: %s", pattern, err)
	}
	if !matched {
		t.Error("message did not match pattern")
	}
}

func TestUTCFlag(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Test:", LstdFlags)
	l.SetFlags(Ldate | Ltime | LUTC)
	// Verify a log message looks right in the right time zone. Quantize to the second only.
	now := time.Now().UTC()
	l.Print("hello")
	want := fmt.Sprintf("Test:%d/%.2d/%.2d %.2d:%.2d:%.2d hello\n",
		now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
	got := b.String()
	if got == want {
		return
	}
	// It's possible we crossed a second boundary between getting now and logging,
	// so add a second and try again. This should very nearly always work.
	now = now.Add(time.Second)
	want = fmt.Sprintf("Test:%d/%.2d/%.2d %.2d:%.2d:%.2d hello\n",
		now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
	if got == want {
		return
	}
	t.Errorf("got %q; want %q", got, want)
}

func TestEmptyPrintCreatesLine(t *testing.T) {
	var b bytes.Buffer
	l := New(&b, "Header:", LstdFlags)
	l.Print()
	l.Println("non-empty")
	output := b.String()
	if n := strings.Count(output, "Header"); n != 2 {
		t.Errorf("expected 2 headers, got %d", n)
	}
	if n := strings.Count(output, "\n"); n != 2 {
		t.Errorf("expected 2 lines, got %d", n)
	}
}

func BenchmarkItoa(b *testing.B) {
	dst := make([]byte, 0, 64)
	for i := 0; i < b.N; i++ {
		dst = dst[0:0]
		itoa(&dst, 2015, 4)   // year
		itoa(&dst, 1, 2)      // month
		itoa(&dst, 30, 2)     // day
		itoa(&dst, 12, 2)     // hour
		itoa(&dst, 56, 2)     // minute
		itoa(&dst, 0, 2)      // second
		itoa(&dst, 987654, 6) // microsecond
	}
}

func BenchmarkPrintln(b *testing.B) {
	const testString = "test"
	var buf bytes.Buffer
	l := New(&buf, "", LstdFlags)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		l.Println(testString)
	}
}

func BenchmarkPrintlnNoFlags(b *testing.B) {
	const testString = "test"
	var buf bytes.Buffer
	l := New(&buf, "", 0)
	for i := 0; i < b.N; i++ {
		buf.Reset()
		l.Println(testString)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom

import (
	"encoding/hex"
	"encoding/json"
)

type RawData struct {
	Data string `json:"data"`
}

func (r *RawData) GetDataBytes() ([]byte, error) {
	return hex.DecodeString(r.Data)
}

func GetRaw(keymr string) ([]byte, error) {
	params := hashRequest{Hash: keymr}
	req := NewJSON2Request("raw-data", APICounter(), params)
	resp, err := factomdRequest(req)
	if err != nil {
		return nil, err
	}
	if resp.Error != nil {
		return nil, resp.Error
	}

	raw := new(RawData)
	if err := json.Unmarshal(resp.JSONResult(), raw); err != nil {
		return nil, err
	}

	return raw.GetDataBytes()
}

func SendRawMsg(message string) (string, error) {
	param := messageRequest{Message: message}
	req := NewJSON2Request("send-raw-message", APICounter(), param)
	resp, err := factomdRequest(req)
	if err != nil {
		return "", err
	}
	if resp.Error != nil {
		return "", resp.Error
	}

	status := new(struct {
		Message string `json:"message"`
	})
	if err := json.Unmarshal(resp.JSONResult(), status); err != nil {
		return "", err
	}

	return status.Message, nil
}
<commit_msg>documentation for raw.go<commit_after>// Copyright 2016 Factom Foundation
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package factom

import (
	"encoding/hex"
	"encoding/json"
)

// RawData is a simple hex encoded byte string
type RawData struct {
	Data string `json:"data"`
}

func (r *RawData) GetDataBytes() ([]byte, error) {
	return hex.DecodeString(r.Data)
}

// GetRaw requests the raw data for any binary block kept in the factomd
// database.
func GetRaw(keymr string) ([]byte, error) {
	params := hashRequest{Hash: keymr}
	req := NewJSON2Request("raw-data", APICounter(), params)
	resp, err := factomdRequest(req)
	if err != nil {
		return nil, err
	}
	if resp.Error != nil {
		return nil, resp.Error
	}

	raw := new(RawData)
	if err := json.Unmarshal(resp.JSONResult(), raw); err != nil {
		return nil, err
	}

	return raw.GetDataBytes()
}

// SendRawMsg sends a raw hex encoded byte string for factomd to send as a
// binary message on the Factom Netwrork.
func SendRawMsg(message string) (string, error) {
	param := messageRequest{Message: message}
	req := NewJSON2Request("send-raw-message", APICounter(), param)
	resp, err := factomdRequest(req)
	if err != nil {
		return "", err
	}
	if resp.Error != nil {
		return "", resp.Error
	}

	status := new(struct {
		Message string `json:"message"`
	})
	if err := json.Unmarshal(resp.JSONResult(), status); err != nil {
		return "", err
	}

	return status.Message, nil
}
<|endoftext|>