<commit_before>package markdown

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/mholt/caddy/middleware"
	"github.com/russross/blackfriday"
)

const (
	DefaultTemplate  = "defaultTemplate"
	DefaultStaticDir = "generated_site"
)

type MarkdownData struct {
	middleware.Context
	Var      map[string]interface{}
	Title    string
	Markdown string
}

// Process processes the contents of a page in b. It parses the metadata
// (if any) and uses the template (if found).
func (md Markdown) Process(c Config, requestPath string, b []byte, ctx middleware.Context) ([]byte, error) {
	var metadata = Metadata{Variables: make(map[string]interface{})}
	var markdown []byte
	var err error

	// find parser compatible with page contents
	parser := findParser(b)

	if parser == nil {
		// if not found, assume whole file is markdown (no front matter)
		markdown = b
	} else {
		// if found, assume metadata present and parse.
		markdown, err = parser.Parse(b)
		if err != nil {
			return nil, err
		}
		metadata = parser.Metadata()
	}

	// if template is not specified, check if Default template is set
	if metadata.Template == "" {
		if _, ok := c.Templates[DefaultTemplate]; ok {
			metadata.Template = DefaultTemplate
		}
	}

	// if template is set, load it
	var tmpl []byte
	if metadata.Template != "" {
		if t, ok := c.Templates[metadata.Template]; ok {
			tmpl, err = ioutil.ReadFile(t)
		}
		if err != nil {
			return nil, err
		}
	}

	// process markdown
	markdown = blackfriday.Markdown(markdown, c.Renderer, 0)

	// set it as body for template
	metadata.Variables["markdown"] = string(markdown)

	return md.processTemplate(c, requestPath, tmpl, metadata, ctx)
}

// processTemplate processes a template given a requestPath,
// template (tmpl) and metadata
func (md Markdown) processTemplate(c Config, requestPath string, tmpl []byte, metadata Metadata, ctx middleware.Context) ([]byte, error) {
	// if template is not specified,
	// use the default template
	if tmpl == nil {
		tmpl = defaultTemplate(c, metadata, requestPath)
	}

	// process the template
	b := new(bytes.Buffer)
	t, err := template.New("").Parse(string(tmpl))
	if err != nil {
		return nil, err
	}
	mdData := MarkdownData{
		Context:  ctx,
		Var:      metadata.Variables,
		Title:    metadata.Title,
		Markdown: metadata.Variables["markdown"].(string),
	}

	if err = t.Execute(b, mdData); err != nil {
		return nil, err
	}

	// generate static page
	if err = md.generatePage(c, requestPath, b.Bytes()); err != nil {
		// if static page generation fails,
		// nothing fatal, only log the error.
		// TODO: Report this non-fatal error, but don't log it here
		log.Println(err)
		fmt.Printf("Error: %v", err)
	}

	return b.Bytes(), nil

}

// generatePage generates a static html page from the markdown in content if c.StaticDir
// is a non-empty value, meaning that the user enabled static site generation.
func (md Markdown) generatePage(c Config, requestPath string, content []byte) error {
	// Only generate the page if static site generation is enabled
	if c.StaticDir != "" {
		// if static directory is not existing, create it
		if _, err := os.Stat(c.StaticDir); err != nil {
			err := os.MkdirAll(c.StaticDir, os.FileMode(0755))
			if err != nil {
				return err
			}
		}

		filePath := filepath.Join(c.StaticDir, requestPath)

		// If it is index file, use the directory instead
		if md.IsIndexFile(filepath.Base(requestPath)) {
			filePath, _ = filepath.Split(filePath)
		}

		// Create the directory in case it is not existing
		if err := os.MkdirAll(filePath, os.FileMode(0744)); err != nil {
			return err
		}

		// generate index.html file in the directory
		filePath = filepath.Join(filePath, "index.html")
		err := ioutil.WriteFile(filePath, content, os.FileMode(0664))
		if err != nil {
			return err
		}

		c.StaticFiles[requestPath] = filePath
	}

	return nil
}

// defaultTemplate constructs a default template.
func defaultTemplate(c Config, metadata Metadata, requestPath string) []byte {
	var scripts, styles bytes.Buffer
	for _, style := range c.Styles {
		styles.WriteString(strings.Replace(cssTemplate, "{{url}}", style, 1))
		styles.WriteString("\r\n")
	}
	for _, script := range c.Scripts {
		scripts.WriteString(strings.Replace(jsTemplate, "{{url}}", script, 1))
		scripts.WriteString("\r\n")
	}

	// Title is first line (length-limited), otherwise filename
	title := metadata.Title
	if title == "" {
		title = filepath.Base(requestPath)
		if body, _ := metadata.Variables["markdown"].([]byte); len(body) > 128 {
			title = string(body[:128])
		} else if len(body) > 0 {
			title = string(body)
		}
	}

	html := []byte(htmlTemplate)
	html = bytes.Replace(html, []byte("{{title}}"), []byte(title), 1)
	html = bytes.Replace(html, []byte("{{css}}"), styles.Bytes(), 1)
	html = bytes.Replace(html, []byte("{{js}}"), scripts.Bytes(), 1)

	return html
}

const (
	htmlTemplate = `<!DOCTYPE html>
<html>
	<head>
		<title>{{title}}</title>
		<meta charset="utf-8">
		{{css}}
		{{js}}
	</head>
	<body>
		{{.Markdown}}
	</body>
</html>`
	cssTemplate = `<link rel="stylesheet" href="{{url}}">`
	jsTemplate  = `<script src="{{url}}"></script>`
)
<commit_msg>Remove undesired committed debug logs<commit_after>package markdown

import (
	"bytes"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/mholt/caddy/middleware"
	"github.com/russross/blackfriday"
)

const (
	DefaultTemplate  = "defaultTemplate"
	DefaultStaticDir = "generated_site"
)

type MarkdownData struct {
	middleware.Context
	Var      map[string]interface{}
	Title    string
	Markdown string
}

// Process processes the contents of a page in b. It parses the metadata
// (if any) and uses the template (if found).
func (md Markdown) Process(c Config, requestPath string, b []byte, ctx middleware.Context) ([]byte, error) {
	var metadata = Metadata{Variables: make(map[string]interface{})}
	var markdown []byte
	var err error

	// find parser compatible with page contents
	parser := findParser(b)

	if parser == nil {
		// if not found, assume whole file is markdown (no front matter)
		markdown = b
	} else {
		// if found, assume metadata present and parse.
		markdown, err = parser.Parse(b)
		if err != nil {
			return nil, err
		}
		metadata = parser.Metadata()
	}

	// if template is not specified, check if Default template is set
	if metadata.Template == "" {
		if _, ok := c.Templates[DefaultTemplate]; ok {
			metadata.Template = DefaultTemplate
		}
	}

	// if template is set, load it
	var tmpl []byte
	if metadata.Template != "" {
		if t, ok := c.Templates[metadata.Template]; ok {
			tmpl, err = ioutil.ReadFile(t)
		}
		if err != nil {
			return nil, err
		}
	}

	// process markdown
	markdown = blackfriday.Markdown(markdown, c.Renderer, 0)

	// set it as body for template
	metadata.Variables["markdown"] = string(markdown)

	return md.processTemplate(c, requestPath, tmpl, metadata, ctx)
}

// processTemplate processes a template given a requestPath,
// template (tmpl) and metadata
func (md Markdown) processTemplate(c Config, requestPath string, tmpl []byte, metadata Metadata, ctx middleware.Context) ([]byte, error) {
	// if template is not specified,
	// use the default template
	if tmpl == nil {
		tmpl = defaultTemplate(c, metadata, requestPath)
	}

	// process the template
	b := new(bytes.Buffer)
	t, err := template.New("").Parse(string(tmpl))
	if err != nil {
		return nil, err
	}
	mdData := MarkdownData{
		Context:  ctx,
		Var:      metadata.Variables,
		Title:    metadata.Title,
		Markdown: metadata.Variables["markdown"].(string),
	}

	if err = t.Execute(b, mdData); err != nil {
		return nil, err
	}

	// generate static page
	if err = md.generatePage(c, requestPath, b.Bytes()); err != nil {
		// if static page generation fails,
		// nothing fatal, only log the error.
		// TODO: Report this non-fatal error, but don't log it here
		log.Println(err)
	}

	return b.Bytes(), nil

}

// generatePage generates a static html page from the markdown in content if c.StaticDir
// is a non-empty value, meaning that the user enabled static site generation.
func (md Markdown) generatePage(c Config, requestPath string, content []byte) error {
	// Only generate the page if static site generation is enabled
	if c.StaticDir != "" {
		// if static directory is not existing, create it
		if _, err := os.Stat(c.StaticDir); err != nil {
			err := os.MkdirAll(c.StaticDir, os.FileMode(0755))
			if err != nil {
				return err
			}
		}

		filePath := filepath.Join(c.StaticDir, requestPath)

		// If it is index file, use the directory instead
		if md.IsIndexFile(filepath.Base(requestPath)) {
			filePath, _ = filepath.Split(filePath)
		}

		// Create the directory in case it is not existing
		if err := os.MkdirAll(filePath, os.FileMode(0744)); err != nil {
			return err
		}

		// generate index.html file in the directory
		filePath = filepath.Join(filePath, "index.html")
		err := ioutil.WriteFile(filePath, content, os.FileMode(0664))
		if err != nil {
			return err
		}

		c.StaticFiles[requestPath] = filePath
	}

	return nil
}

// defaultTemplate constructs a default template.
func defaultTemplate(c Config, metadata Metadata, requestPath string) []byte {
	var scripts, styles bytes.Buffer
	for _, style := range c.Styles {
		styles.WriteString(strings.Replace(cssTemplate, "{{url}}", style, 1))
		styles.WriteString("\r\n")
	}
	for _, script := range c.Scripts {
		scripts.WriteString(strings.Replace(jsTemplate, "{{url}}", script, 1))
		scripts.WriteString("\r\n")
	}

	// Title is first line (length-limited), otherwise filename
	title := metadata.Title
	if title == "" {
		title = filepath.Base(requestPath)
		if body, _ := metadata.Variables["markdown"].([]byte); len(body) > 128 {
			title = string(body[:128])
		} else if len(body) > 0 {
			title = string(body)
		}
	}

	html := []byte(htmlTemplate)
	html = bytes.Replace(html, []byte("{{title}}"), []byte(title), 1)
	html = bytes.Replace(html, []byte("{{css}}"), styles.Bytes(), 1)
	html = bytes.Replace(html, []byte("{{js}}"), scripts.Bytes(), 1)

	return html
}

const (
	htmlTemplate = `<!DOCTYPE html>
<html>
	<head>
		<title>{{title}}</title>
		<meta charset="utf-8">
		{{css}}
		{{js}}
	</head>
	<body>
		{{.Markdown}}
	</body>
</html>`
	cssTemplate = `<link rel="stylesheet" href="{{url}}">`
	jsTemplate  = `<script src="{{url}}"></script>`
)
<|endoftext|>


<commit_before>// Copyright 2011 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 main

import (
	"bytes"
	"io"
	"log"
	"os"
	"os/exec"
)

// run is a simple wrapper for exec.Run/Close
func run(envv []string, dir string, argv ...string) error {
	if *verbose {
		log.Println("run", argv)
	}
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = dir
	cmd.Env = envv
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// runLog runs a process and returns the combined stdout/stderr,
// as well as writing it to logfile (if specified). It returns
// process combined stdout and stderr output, exit status and error.
// The error returned is nil, if process is started successfully,
// even if exit status is not 0.
func runLog(envv []string, logfile, dir string, argv ...string) (string, int, error) {
	if *verbose {
		log.Println("runLog", argv)
	}

	b := new(bytes.Buffer)
	var w io.Writer = b
	if logfile != "" {
		f, err := os.OpenFile(logfile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			return "", 0, err
		}
		defer f.Close()
		w = io.MultiWriter(f, b)
	}

	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = dir
	cmd.Env = envv
	cmd.Stdout = w
	cmd.Stderr = w

	err := cmd.Run()
	if err != nil {
		if ws, ok := err.(*exec.ExitError); ok {
			return b.String(), ws.ExitStatus(), nil
		}
	}
	return b.String(), 0, err
}
<commit_msg>builder: update for os.Wait changes.<commit_after>// Copyright 2011 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 main

import (
	"bytes"
	"io"
	"log"
	"os"
	"os/exec"
)

// run is a simple wrapper for exec.Run/Close
func run(envv []string, dir string, argv ...string) error {
	if *verbose {
		log.Println("run", argv)
	}
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = dir
	cmd.Env = envv
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// runLog runs a process and returns the combined stdout/stderr,
// as well as writing it to logfile (if specified). It returns
// process combined stdout and stderr output, exit status and error.
// The error returned is nil, if process is started successfully,
// even if exit status is not successful.
func runLog(envv []string, logfile, dir string, argv ...string) (string, int, error) {
	if *verbose {
		log.Println("runLog", argv)
	}

	b := new(bytes.Buffer)
	var w io.Writer = b
	if logfile != "" {
		f, err := os.OpenFile(logfile, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
		if err != nil {
			return "", 0, err
		}
		defer f.Close()
		w = io.MultiWriter(f, b)
	}

	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = dir
	cmd.Env = envv
	cmd.Stdout = w
	cmd.Stderr = w

	startErr := cmd.Start()
	if startErr != nil {
		return "", 1, startErr
	}
	exitStatus := 0
	if err := cmd.Wait(); err != nil {
		exitStatus = 1 // TODO(bradfitz): this is fake. no callers care, so just return a bool instead.
	}
	return b.String(), exitStatus, nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The CUE 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 cmd

import (
	"github.com/spf13/cobra"

	"cuelang.org/go/internal/encoding"
	"cuelang.org/go/internal/filetypes"
)

// newExportCmd creates and export command
func newExportCmd(c *Command) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "export",
		Short: "output data in a standard format",
		Long: `export evaluates the configuration found in the current
directory and prints the emit value to stdout.

Examples:
Evaluated and emit

	# a single file
	cue export config.cue

	# multiple files: these are combined at the top-level. Order doesn't matter.
	cue export file1.cue foo/file2.cue

	# all files within the "mypkg" package: this includes all files in the
	# current directory and its ancestor directories that are marked with the
	# same package.
	cue export -p cloud

	# the -p flag can be omitted if the directory only contains files for
	# the "mypkg" package.
	cue export

Emit value:
For CUE files, the generated configuration is derived from the top-level
single expression, the emit value. For example, the file

	// config.cue
	arg1: 1
	arg2: "my string"

	{
		a: arg1
		b: arg2
	}

yields the following JSON:

	{
		"a": 1,
		"b", "my string"
	}

In absence of arguments, the current directory is loaded as a package instance.
A package instance for a directory contains all files in the directory and its
ancestor directories, up to the module root, belonging to the same package.
If the package is not explicitly defined by the '-p' flag, it must be uniquely
defined by the files in the current directory.

Formats
The following formats are recognized:

json    output as JSON
               Outputs any CUE value.

text    output as raw text
                The evaluated value must be of type string.

yaml    output as YAML
                Outputs any CUE value.
`,

		RunE: mkRunE(c, runExport),
	}

	addOutFlags(cmd.Flags(), true)
	addOrphanFlags(cmd.Flags())

	cmd.Flags().Bool(string(flagEscape), false, "use HTML escaping")

	cmd.Flags().StringArrayP(string(flagExpression), "e", nil, "export this expression only")

	cmd.Flags().StringArrayP(string(flagInject), "t", nil,
		"set the value of a tagged field")

	return cmd
}

func runExport(cmd *Command, args []string) error {
	b, err := parseArgs(cmd, args, &config{outMode: filetypes.Export})
	exitOnErr(cmd, err, true)

	enc, err := encoding.NewEncoder(b.outFile, b.encConfig)
	exitOnErr(cmd, err, true)
	defer enc.Close()

	iter := b.instances()
	defer iter.close()
	for iter.scan() {
		v := iter.value()
		err = enc.Encode(v)
		exitOnErr(cmd, err, true)
	}
	exitOnErr(cmd, iter.err(), true)
	return nil
}
<commit_msg>cmd/export: fix help<commit_after>// Copyright 2018 The CUE 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 cmd

import (
	"github.com/spf13/cobra"

	"cuelang.org/go/internal/encoding"
	"cuelang.org/go/internal/filetypes"
)

// newExportCmd creates and export command
func newExportCmd(c *Command) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "export",
		Short: "output data in a standard format",
		Long: `export evaluates the configuration found in the current
directory and prints the emit value to stdout.

Examples:
Evaluated and emit

	# a single file
	cue export config.cue

	# multiple files: these are combined at the top-level. Order doesn't matter.
	cue export file1.cue foo/file2.cue

	# all files within the "mypkg" package: this includes all files in the
	# current directory and its ancestor directories that are marked with the
	# same package.
	cue export -p cloud

	# the -p flag can be omitted if the directory only contains files for
	# the "mypkg" package.
	cue export

Emit value:
For CUE files, the generated configuration is derived from the top-level
single expression, the emit value. For example, the file

	// config.cue
	arg1: 1
	arg2: "my string"

	{
		a: arg1
		b: arg2
	}

yields the following JSON:

	{
		"arg1": 1,
		"a": 1,
		"arg2": "my string",
		"b": "my string"
	}

In absence of arguments, the current directory is loaded as a package instance.
A package instance for a directory contains all files in the directory and its
ancestor directories, up to the module root, belonging to the same package.
If the package is not explicitly defined by the '-p' flag, it must be uniquely
defined by the files in the current directory.

Formats
The following formats are recognized:

json    output as JSON
               Outputs any CUE value.

text    output as raw text
                The evaluated value must be of type string.

yaml    output as YAML
                Outputs any CUE value.
`,

		RunE: mkRunE(c, runExport),
	}

	addOutFlags(cmd.Flags(), true)
	addOrphanFlags(cmd.Flags())

	cmd.Flags().Bool(string(flagEscape), false, "use HTML escaping")

	cmd.Flags().StringArrayP(string(flagExpression), "e", nil, "export this expression only")

	cmd.Flags().StringArrayP(string(flagInject), "t", nil,
		"set the value of a tagged field")

	return cmd
}

func runExport(cmd *Command, args []string) error {
	b, err := parseArgs(cmd, args, &config{outMode: filetypes.Export})
	exitOnErr(cmd, err, true)

	enc, err := encoding.NewEncoder(b.outFile, b.encConfig)
	exitOnErr(cmd, err, true)
	defer enc.Close()

	iter := b.instances()
	defer iter.close()
	for iter.scan() {
		v := iter.value()
		err = enc.Encode(v)
		exitOnErr(cmd, err, true)
	}
	exitOnErr(cmd, iter.err(), true)
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Andrew O'Neill

// 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 main

import (
	"log"
	"net"

	"google.golang.org/grpc"

	"github.com/foolusion/choices/elwinstorage"
	"github.com/foolusion/choices/storage/mongo"
)

func main() {
	log.Println("Starting exp-store...")
	server, err := mongo.NewServer("localhost", "elwin")
	server.LoadExampleData()

	lis, err := net.Listen("tcp", ":8080")
	if err != nil {
		log.Fatal(err)
	}

	s := grpc.NewServer()
	storage.RegisterElwinStorageServer(s, server)
	log.Fatal(s.Serve(lis))
}
<commit_msg>cmd/exp-store: add grpc metrics<commit_after>// Copyright 2016 Andrew O'Neill

// 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 main

import (
	"log"
	"net"
	"net/http"

	"google.golang.org/grpc"

	"github.com/foolusion/choices/elwinstorage"
	"github.com/foolusion/choices/storage/mongo"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/prometheus/client_golang/prometheus"
)

func main() {
	log.Println("Starting exp-store...")
	server, err := mongo.NewServer("localhost", "elwin")
	server.LoadExampleData()

	lis, err := net.Listen("tcp", ":8080")
	if err != nil {
		log.Fatal(err)
	}

	s := grpc.NewServer(
		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
	)
	storage.RegisterElwinStorageServer(s, server)
	grpc_prometheus.Register(s)
	go func() {
		http.Handle("/metrics", prometheus.Handler())
		log.Fatal(http.ListenAndServe(":8081", nil))
	}()

	log.Fatal(s.Serve(lis))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/go-redis/redis"
	metrics "github.com/rcrowley/go-metrics"
	log "github.com/sirupsen/logrus"
)

// credentials are used by basic auth and include the hash of a valid password, plus
// a "stage" string which is used to emit metrics that are useful when managing credrolls, so that
// we can track whether or not deprecated passwords are still in use.
type credential struct {
	Stage string `json:stage`
	Hmac  string `json:hmac`
}

func newAuth(config AuthConfig, registry metrics.Registry) (*BasicAuth, error) {
	if config.RedisUrl != "" && config.RedisKey == "" {
		return nil, errors.New("RedisKey must be set if RedisUrl is set")
	}

	if config.RedisUrl == "" && config.Tokens == "" {
		return nil, errors.New("At least one of RedisUrl or Tokens must be set.")
	}

	result, err := NewBasicAuthFromString(config.Tokens, config.HmacKey, registry)

	if err != nil {
		return result, err
	}

	if config.RedisUrl == "" {
		return result, err
	}

	// Parse redis db out of url
	opt, err := redis.ParseURL(config.RedisUrl)
	if err != nil {
		return result, err
	}
	client := redis.NewClient(opt)

	// Refresh once at the start.
	// Ignore errors here - if redis is down, we don't want the process to fail to start.
	_, _ = result.refresh(client, config.HmacKey, config.RedisKey, config.Tokens)

	// Refresh forever.
	go result.startRefresh(client, config, registry)

	return result, err
}

func (auth *BasicAuth) startRefresh(client *redis.Client, config AuthConfig, registry metrics.Registry) {
	pChanges := metrics.GetOrRegisterCounter("log-iss.auth_refresh.changes", registry)
	pFailures := metrics.GetOrRegisterCounter("log-iss.auth_refresh.failures", registry)
	pSuccesses := metrics.GetOrRegisterCounter("log-iss.auth_refresh.successes", registry)
	ticker := time.NewTicker(config.RefreshInterval)

	for ; true; <-ticker.C {
		changed, err := auth.refresh(client, config.HmacKey, config.RedisKey, config.Tokens)
		if err == nil {
			pSuccesses.Inc(1)
			if changed {
				pChanges.Inc(1)
			}
			continue
		}
		log.WithFields(log.Fields{"ns": "auth", "at": "error", "refresh": true, "message": err.Error()}).Info()
		fmt.Printf("Unable to refresh credentials: %s", err)
		pFailures.Inc(1)
	}
}

// Refresh auth credentials.
// Return true if credentials changed, false otherwise.
func (ba *BasicAuth) refresh(client redis.Cmdable, hmacKey string, redisKey string, config string) (bool, error) {
	// Start out using the strings from config
	nba, err := NewBasicAuthFromString(config, hmacKey, ba.registry)
	if err != nil {
		return false, err
	}

	// Retrieve all secrets and construct a string in the format expected by BasicAuth
	val := client.HGetAll(redisKey)
	// Return error if the key does not exist.
	if val == nil {
		return false, fmt.Errorf("Key %s not found", redisKey)
	}

	r, err := val.Result()
	if err != nil {
		return false, err
	}

	for k, v := range r {
		var arr []credential
		err = json.Unmarshal([]byte(v), &arr)
		if err != nil {
			return false, err
		}
		nba.creds[k] = arr
	}

	// Swap the secrets if there are changes
	if !reflect.DeepEqual(ba.creds, nba.creds) {
		ba.Lock()
		defer ba.Unlock()
		ba.creds = nba.creds
		return true, nil
	}
	return false, nil
}

// BasicAuth handles normal user/password Basic Auth requests, multiple
// password for the same user and is safe for concurrent use.
type BasicAuth struct {
	sync.RWMutex
	creds    map[string][]credential
	hmacKey  string
	registry metrics.Registry
}

// NewBasicAuthFromString creates and populates a BasicAuth from the provided
// credentials, encoded as a string, in the following format:
// user:password|user:password|...
func NewBasicAuthFromString(creds string, hmacKey string, registry metrics.Registry) (*BasicAuth, error) {
	ba := NewBasicAuth(registry, hmacKey)

	// If the string is empty, that's allowed.
	if creds == "" {
		return ba, nil
	}

	for _, u := range strings.Split(creds, "|") {
		uparts := strings.SplitN(u, ":", 2)
		if len(uparts) != 2 || len(uparts[0]) == 0 || len(uparts[1]) == 0 {
			return ba, fmt.Errorf("Unable to create credentials from '%s'", u)
		}

		ba.AddPrincipal(uparts[0], hmacEncode(hmacKey, uparts[1]), "env")
	}
	return ba, nil
}

func hmacEncode(key string, message string) string {
	hash := hmac.New(sha512.New, []byte(key))
	hash.Write([]byte(message))
	return hex.EncodeToString(hash.Sum(nil))
}

func NewBasicAuth(registry metrics.Registry, hmacKey string) *BasicAuth {
	return &BasicAuth{
		creds:    make(map[string][]credential),
		hmacKey:  hmacKey,
		registry: registry,
	}
}

// AddPrincipal add's a user/password combo to the list of valid combinations
func (ba *BasicAuth) AddPrincipal(user string, hmac string, stage string) {
	ba.Lock()
	defer ba.Unlock()
	u, exists := ba.creds[user]
	if !exists {
		u = make([]credential, 0, 1)
	}
	ba.creds[user] = append(u, credential{Stage: stage, Hmac: hmac})
}

// Authenticate is true if the Request has a valid BasicAuth signature and
// that signature encodes a known username/password combo.
func (ba *BasicAuth) Authenticate(r *http.Request) bool {
	user, pass, ok := r.BasicAuth()
	if !ok {
		log.WithFields(log.Fields{"ns": "auth", "at": "failure", "no_basic_auth": true}).Info()
		return false
	}

	ba.RLock()
	defer ba.RUnlock()

	credentials, exists := ba.creds[user]
	if !exists {
		log.WithFields(log.Fields{"ns": "auth", "at": "failure", "user": user}).Info()
		return false
	}

	for _, c := range credentials {
		if c.Hmac == hmacEncode(ba.hmacKey, pass) {
			countName := fmt.Sprintf("log-iss.auth.%s.%s.successes", user, c.Stage)
			counter := metrics.GetOrRegisterCounter(countName, ba.registry)
			counter.Inc(1)
			return true
		}
	}
	countName := fmt.Sprintf("log-iss.auth.%s.failures", user)
	counter := metrics.GetOrRegisterCounter(countName, ba.registry)
	counter.Inc(1)
	return false
}
<commit_msg>Remove initial refresh of creds - it's no longer relevant<commit_after>package main

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/go-redis/redis"
	metrics "github.com/rcrowley/go-metrics"
	log "github.com/sirupsen/logrus"
)

// credentials are used by basic auth and include the hash of a valid password, plus
// a "stage" string which is used to emit metrics that are useful when managing credrolls, so that
// we can track whether or not deprecated passwords are still in use.
type credential struct {
	Stage string `json:stage`
	Hmac  string `json:hmac`
}

func newAuth(config AuthConfig, registry metrics.Registry) (*BasicAuth, error) {
	if config.RedisUrl != "" && config.RedisKey == "" {
		return nil, errors.New("RedisKey must be set if RedisUrl is set")
	}

	if config.RedisUrl == "" && config.Tokens == "" {
		return nil, errors.New("At least one of RedisUrl or Tokens must be set.")
	}

	result, err := NewBasicAuthFromString(config.Tokens, config.HmacKey, registry)

	if err != nil {
		return result, err
	}

	if config.RedisUrl == "" {
		return result, err
	}

	// Parse redis db out of url
	opt, err := redis.ParseURL(config.RedisUrl)
	if err != nil {
		return result, err
	}
	client := redis.NewClient(opt)

	// Refresh forever.
	go result.startRefresh(client, config, registry)

	return result, err
}

func (auth *BasicAuth) startRefresh(client *redis.Client, config AuthConfig, registry metrics.Registry) {
	pChanges := metrics.GetOrRegisterCounter("log-iss.auth_refresh.changes", registry)
	pFailures := metrics.GetOrRegisterCounter("log-iss.auth_refresh.failures", registry)
	pSuccesses := metrics.GetOrRegisterCounter("log-iss.auth_refresh.successes", registry)
	ticker := time.NewTicker(config.RefreshInterval)

	for ; true; <-ticker.C {
		changed, err := auth.refresh(client, config.HmacKey, config.RedisKey, config.Tokens)
		if err == nil {
			pSuccesses.Inc(1)
			if changed {
				pChanges.Inc(1)
			}
			continue
		}
		log.WithFields(log.Fields{"ns": "auth", "at": "error", "refresh": true, "message": err.Error()}).Info()
		pFailures.Inc(1)
	}
}

// Refresh auth credentials.
// Return true if credentials changed, false otherwise.
func (ba *BasicAuth) refresh(client redis.Cmdable, hmacKey string, redisKey string, config string) (bool, error) {
	// Start out using the strings from config
	nba, err := NewBasicAuthFromString(config, hmacKey, ba.registry)
	if err != nil {
		return false, err
	}

	// Retrieve all secrets and construct a string in the format expected by BasicAuth
	val := client.HGetAll(redisKey)
	// Return error if the key does not exist.
	if val == nil {
		return false, fmt.Errorf("Key %s not found", redisKey)
	}

	r, err := val.Result()
	if err != nil {
		return false, err
	}

	for k, v := range r {
		var arr []credential
		err = json.Unmarshal([]byte(v), &arr)
		if err != nil {
			return false, err
		}
		nba.creds[k] = arr
	}

	// Swap the secrets if there are changes
	if !reflect.DeepEqual(ba.creds, nba.creds) {
		ba.Lock()
		defer ba.Unlock()
		ba.creds = nba.creds
		return true, nil
	}
	return false, nil
}

// BasicAuth handles normal user/password Basic Auth requests, multiple
// password for the same user and is safe for concurrent use.
type BasicAuth struct {
	sync.RWMutex
	creds    map[string][]credential
	hmacKey  string
	registry metrics.Registry
}

// NewBasicAuthFromString creates and populates a BasicAuth from the provided
// credentials, encoded as a string, in the following format:
// user:password|user:password|...
func NewBasicAuthFromString(creds string, hmacKey string, registry metrics.Registry) (*BasicAuth, error) {
	ba := NewBasicAuth(registry, hmacKey)

	// If the string is empty, that's allowed.
	if creds == "" {
		return ba, nil
	}

	for _, u := range strings.Split(creds, "|") {
		uparts := strings.SplitN(u, ":", 2)
		if len(uparts) != 2 || len(uparts[0]) == 0 || len(uparts[1]) == 0 {
			return ba, fmt.Errorf("Unable to create credentials from '%s'", u)
		}

		ba.AddPrincipal(uparts[0], hmacEncode(hmacKey, uparts[1]), "env")
	}
	return ba, nil
}

func hmacEncode(key string, message string) string {
	hash := hmac.New(sha512.New, []byte(key))
	hash.Write([]byte(message))
	return hex.EncodeToString(hash.Sum(nil))
}

func NewBasicAuth(registry metrics.Registry, hmacKey string) *BasicAuth {
	return &BasicAuth{
		creds:    make(map[string][]credential),
		hmacKey:  hmacKey,
		registry: registry,
	}
}

// AddPrincipal add's a user/password combo to the list of valid combinations
func (ba *BasicAuth) AddPrincipal(user string, hmac string, stage string) {
	ba.Lock()
	defer ba.Unlock()
	u, exists := ba.creds[user]
	if !exists {
		u = make([]credential, 0, 1)
	}
	ba.creds[user] = append(u, credential{Stage: stage, Hmac: hmac})
}

// Authenticate is true if the Request has a valid BasicAuth signature and
// that signature encodes a known username/password combo.
func (ba *BasicAuth) Authenticate(r *http.Request) bool {
	user, pass, ok := r.BasicAuth()
	if !ok {
		log.WithFields(log.Fields{"ns": "auth", "at": "failure", "no_basic_auth": true}).Info()
		return false
	}

	ba.RLock()
	defer ba.RUnlock()

	credentials, exists := ba.creds[user]
	if !exists {
		log.WithFields(log.Fields{"ns": "auth", "at": "failure", "user": user}).Info()
		return false
	}

	for _, c := range credentials {
		if c.Hmac == hmacEncode(ba.hmacKey, pass) {
			countName := fmt.Sprintf("log-iss.auth.%s.%s.successes", user, c.Stage)
			counter := metrics.GetOrRegisterCounter(countName, ba.registry)
			counter.Inc(1)
			return true
		}
	}
	countName := fmt.Sprintf("log-iss.auth.%s.failures", user)
	counter := metrics.GetOrRegisterCounter(countName, ba.registry)
	counter.Inc(1)
	return false
}
<|endoftext|>


<commit_before><commit_msg>reformat<commit_after><|endoftext|>


<commit_before>// Copyright 2016 Red Hat, Inc, and individual contributors.
//
// 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 main

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/kubectl"
)

func setContextCommand(sh *kubesh, args []string) error {
	switch l := len(args); {
	case l > 3:
		fmt.Println("Usage: " + args[0] + "[TYPE [NAME]]")

		//TODO: return an error?
		return nil
	case l == 1:
		sh.context = []string{}
		sh.rl.SetPrompt(prompt(sh.context))

		return nil
	}

	resources, err := sh.finder.Lookup(args[1:])
	if err != nil {

		return err
	}

	typeOnly := len(args) == 2

	if len(resources) > 0 {
		res := resources[0]
		if typeOnly {
			sh.context = []string{res.typeName}
		} else {
			sh.context = []string{res.typeName, res.name}
		}
	}

	sh.rl.SetPrompt(prompt(sh.context))

	return nil
}

func prompt(context []string) string {
	return strings.Join(context, ":") + "> "
}

func applyContext(context []string, args []string, rootCommand *cobra.Command) ([]string, error) {
	if len(context) > 0 {
		subcmd, _, err := rootCommand.Find(args[:1])
		if err != nil {

			return args, err
		}

		// poor man's set
		resourceTypes := map[string]struct{}{}
		for _, t := range kubectl.ResourceAliases(ResourceTypes(subcmd)) {
			resourceTypes[t] = struct{}{}
		}
		newArgs := []string{}
		newArgs = append(newArgs, args[0])

		if _, ok := resourceTypes[context[0]]; ok {
			err := subcmd.ParseFlags(args)
			if err != nil {

				return args, err
			}

			nonFlagArgs := subcmd.Flags().Args()
			// the subcommand is an arg, so we'll always have at least one
			switch l := len(nonFlagArgs); {
			case l == 1:
				newArgs = append(newArgs, context...)
			case l == 2:
				// could be a resource type or a resource name
				arg := nonFlagArgs[1]
				if _, ok := resourceTypes[arg]; !ok {
					// treat as a resource name, use the resource type from the context
					newArgs = append(newArgs, context[0])
				}
			}

			return append(newArgs, args[1:]...), nil
		}
	}

	return args, nil
}
<commit_msg>Add clear subcommand to pin along with useful -h<commit_after>// Copyright 2016 Red Hat, Inc, and individual contributors.
//
// 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 main

import (
	"fmt"
	"strings"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	"k8s.io/kubernetes/pkg/kubectl"
)

var help = dedent.Dedent(`
      Pin to a resource or resource type.

      Pinning causes the shell to remember the given resource and/or
      resource type, and apply it to commands as appropriate, allowing
      you to leave the resource type and/or name out of other command
      invocations.

      # Pin to pods
      pin pods

      # Pin to a particular pod
      pin pod nginx-1234-asdf

      # Clear the pin
      pin clear

      The current pin will be shown in the prompt.
`)

func setContextCommand(sh *kubesh, args []string) error {
	if len(args) == 1 || len(args) > 3 {
		fmt.Println("Usage: " + args[0] + " (-h | clear | [TYPE [NAME]])")

		//TODO: return an error?
		return nil
	}

	switch arg := args[1]; {
	case arg == "clear":
		sh.context = []string{}
		sh.rl.SetPrompt(prompt(sh.context))

		return nil

	case arg == "-h":
		fmt.Println(help)

		return nil
	}

	resources, err := sh.finder.Lookup(args[1:])
	if err != nil {

		return err
	}

	typeOnly := len(args) == 2

	if len(resources) > 0 {
		res := resources[0]
		if typeOnly {
			sh.context = []string{res.typeName}
		} else {
			sh.context = []string{res.typeName, res.name}
		}
	}

	sh.rl.SetPrompt(prompt(sh.context))

	return nil
}

func prompt(context []string) string {
	return strings.Join(context, ":") + "> "
}

func applyContext(context []string, args []string, rootCommand *cobra.Command) ([]string, error) {
	if len(context) > 0 {
		subcmd, _, err := rootCommand.Find(args[:1])
		if err != nil {

			return args, err
		}

		// poor man's set
		resourceTypes := map[string]struct{}{}
		for _, t := range kubectl.ResourceAliases(ResourceTypes(subcmd)) {
			resourceTypes[t] = struct{}{}
		}
		newArgs := []string{}
		newArgs = append(newArgs, args[0])

		if _, ok := resourceTypes[context[0]]; ok {
			err := subcmd.ParseFlags(args)
			if err != nil {

				return args, err
			}

			nonFlagArgs := subcmd.Flags().Args()
			// the subcommand is an arg, so we'll always have at least one
			switch l := len(nonFlagArgs); {
			case l == 1:
				newArgs = append(newArgs, context...)
			case l == 2:
				// could be a resource type or a resource name
				arg := nonFlagArgs[1]
				if _, ok := resourceTypes[arg]; !ok {
					// treat as a resource name, use the resource type from the context
					newArgs = append(newArgs, context[0])
				}
			}

			return append(newArgs, args[1:]...), nil
		}
	}

	return args, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/neurosnap/sentences/english"
	"github.com/spf13/cobra"
)

var VERSION string
var COMMITHASH string

var ver bool
var fname string
var delim string

var sentencesCmd = &cobra.Command{
	Use:   "sentences",
	Short: "Sentence tokenizer",
	Long:  "A utility that will break up a blob of text into sentences.",
	Run: func(cmd *cobra.Command, args []string) {
		var text []byte

		if ver {
			fmt.Println(VERSION)
			fmt.Println(COMMITHASH)
			return
		}

		if fname != "" {
			text, _ = ioutil.ReadFile(fname)
		} else {
			reader := bufio.NewReader(os.Stdin)
			text, _ = ioutil.ReadAll(reader)
		}

		tokenizer, err := english.NewSentenceTokenizer(nil)
		if err != nil {
			panic(err)
		}

		sentences := tokenizer.Tokenize(string(text))
		for _, s := range sentences {
			text := strings.Join(strings.Fields(s.Text), " ")

			text = strings.Join([]string{text, delim}, "")
			fmt.Printf("%s", text)
		}
	},
}

func main() {
	sentencesCmd.Flags().BoolVarP(&ver, "version", "v", false, "Get current version of sentences")
	sentencesCmd.Flags().StringVarP(&fname, "file", "f", "", "Read file as source input instead of stdin")
	sentencesCmd.Flags().StringVarP(&delim, "delimiter", "d", "\n", "Delimiter used to demarcate sentence boundaries")

	if err := sentencesCmd.Execute(); err != nil {
		fmt.Print(err)
	}
}
<commit_msg>more cmd line goodness<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/neurosnap/sentences/english"
	"github.com/spf13/cobra"
)

var VERSION string
var COMMITHASH string

var ver bool
var fname string
var delim string

var sentencesCmd = &cobra.Command{
	Use:   "sentences",
	Short: "Sentence tokenizer",
	Long:  "A utility that will break up a blob of text into sentences.",
	Run: func(cmd *cobra.Command, args []string) {
		var text []byte
		var err error

		if ver {
			fmt.Println(VERSION)
			fmt.Println(COMMITHASH)
			return
		}

		if fname != "" {
			text, err = ioutil.ReadFile(fname)
			if err != nil {
				panic(err)
			}
		} else {
			stat, err := os.Stdin.Stat()
			if err != nil {
				panic(err)
			}

			if (stat.Mode() & os.ModeCharDevice) != 0 {
				return
			}

			reader := bufio.NewReader(os.Stdin)
			text, err = ioutil.ReadAll(reader)
			if err != nil {
				panic(err)
			}
		}

		tokenizer, err := english.NewSentenceTokenizer(nil)
		if err != nil {
			panic(err)
		}

		sentences := tokenizer.Tokenize(string(text))
		for _, s := range sentences {
			text := strings.Join(strings.Fields(s.Text), " ")

			text = strings.Join([]string{text, delim}, "")
			fmt.Printf("%s", text)
		}
	},
}

func main() {
	sentencesCmd.Flags().BoolVarP(&ver, "version", "v", false, "Get current version of sentences")
	sentencesCmd.Flags().StringVarP(&fname, "file", "f", "", "Read file as source input instead of stdin")
	sentencesCmd.Flags().StringVarP(&delim, "delimiter", "d", "\n", "Delimiter used to demarcate sentence boundaries")

	if err := sentencesCmd.Execute(); err != nil {
		fmt.Print(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/simple"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	stdlog "log"
	"os"
	"strings"
	"sync"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}

	if c.GlobalBool("json") {
		log.SetFormatter(&log.JSONFormatter{})
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	for _, output := range c.GlobalStringSlice("output") {
		parts := strings.SplitN(output, "+", 2)
		switch parts[0] {
		case "influxdb":
			out, err := influxdb.NewFromURL(parts[1])
			if err != nil {
				log.WithError(err).Fatal("Couldn't create InfluxDB client")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		case "stdout":
			out := &LogMetricsOutput{Writer: os.Stdout}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		default:
			log.WithField("type", parts[0]).Fatal("Unrecognized metric output")
		}
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func headlessController(c context.Context, t *speedboat.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(t.Script, string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Context that expires at the end of the test
	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	// Configure the VU logger
	logger := &log.Logger{
		Out:       os.Stderr,
		Level:     log.DebugLevel,
		Formatter: &log.TextFormatter{},
	}
	if cc.GlobalBool("json") {
		logger.Formatter = &log.JSONFormatter{}
	}
	ctx = speedboat.WithLogger(ctx, logger)

	// Output metrics appropriately; use a mutex to prevent garbled output
	logMetrics := cc.Bool("log")
	metricsLogger := stdlog.New(os.Stdout, "metrics: ", stdlog.Lmicroseconds)
	metricsMutex := sync.Mutex{}
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if logMetrics {
					metricsMutex.Lock()
					printMetrics(metricsLogger)
					metricsMutex.Unlock()
				}
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	mVUs := sampler.Gauge("vus")
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, vuCancel := context.WithCancel(ctx)
			vus = append(vus, vuCancel)
			go func() {
				defer func() {
					if v := recover(); v != nil {
						switch err := v.(type) {
						case speedboat.FlowControl:
							switch err {
							case speedboat.AbortTest:
								log.Error("Test aborted")
								cancel()
							}
						default:
							log.WithFields(log.Fields{
								"id":    i,
								"error": err,
							}).Error("VU crashed!")
						}
					}
				}()
				runner.RunVU(vuCtx, t, len(vus))
			}()
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print final metrics
	metricsMutex.Lock()
	printMetrics(metricsLogger)
	metricsMutex.Unlock()
	commitMetrics()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringSliceFlag{
			Name:  "output, o",
			Usage: "Output metrics to a file or database",
		},
		cli.BoolFlag{
			Name:  "log, l",
			Usage: "Log metrics to stdout",
		},
		cli.BoolFlag{
			Name:  "json",
			Usage: "Format log messages as JSON",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<commit_msg>[fix] Don't log final metrics without -l<commit_after>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/js"
	"github.com/loadimpact/speedboat/sampler"
	"github.com/loadimpact/speedboat/sampler/influxdb"
	"github.com/loadimpact/speedboat/simple"
	"github.com/urfave/cli"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	stdlog "log"
	"os"
	"strings"
	"sync"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}

	if c.GlobalBool("json") {
		log.SetFormatter(&log.JSONFormatter{})
	}
}

// Configure the global sampler.
func configureSampler(c *cli.Context) {
	sampler.DefaultSampler.OnError = func(err error) {
		log.WithError(err).Error("[Sampler error]")
	}

	for _, output := range c.GlobalStringSlice("output") {
		parts := strings.SplitN(output, "+", 2)
		switch parts[0] {
		case "influxdb":
			out, err := influxdb.NewFromURL(parts[1])
			if err != nil {
				log.WithError(err).Fatal("Couldn't create InfluxDB client")
			}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		case "stdout":
			out := &LogMetricsOutput{Writer: os.Stdout}
			sampler.DefaultSampler.Outputs = append(sampler.DefaultSampler.Outputs, out)
		default:
			log.WithField("type", parts[0]).Fatal("Unrecognized metric output")
		}
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func headlessController(c context.Context, t *speedboat.Test) <-chan int {
	ch := make(chan int)

	go func() {
		defer close(ch)

		select {
		case ch <- t.VUsAt(0):
		case <-c.Done():
			return
		}

		startTime := time.Now()
		ticker := time.NewTicker(100 * time.Millisecond)
		for {
			select {
			case <-ticker.C:
				ch <- t.VUsAt(time.Since(startTime))
			case <-c.Done():
				return
			}
		}
	}()

	return ch
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	case strings.HasSuffix(t.Script, ".js"):
		src, err := ioutil.ReadFile(t.Script)
		if err != nil {
			log.WithError(err).Fatal("Couldn't read script")
		}
		runner = js.New(t.Script, string(src))
	default:
		log.Fatal("No suitable runner found!")
	}

	// Context that expires at the end of the test
	ctx, cancel := context.WithTimeout(context.Background(), t.TotalDuration())

	// Configure the VU logger
	logger := &log.Logger{
		Out:       os.Stderr,
		Level:     log.DebugLevel,
		Formatter: &log.TextFormatter{},
	}
	if cc.GlobalBool("json") {
		logger.Formatter = &log.JSONFormatter{}
	}
	ctx = speedboat.WithLogger(ctx, logger)

	// Output metrics appropriately; use a mutex to prevent garbled output
	logMetrics := cc.Bool("log")
	metricsLogger := stdlog.New(os.Stdout, "metrics: ", stdlog.Lmicroseconds)
	metricsMutex := sync.Mutex{}
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				if logMetrics {
					metricsMutex.Lock()
					printMetrics(metricsLogger)
					metricsMutex.Unlock()
				}
				commitMetrics()
			case <-ctx.Done():
				return
			}
		}
	}()

	// Use a "headless controller" to scale VUs by polling the test ramp
	mVUs := sampler.Gauge("vus")
	vus := []context.CancelFunc{}
	for scale := range headlessController(ctx, &t) {
		for i := len(vus); i < scale; i++ {
			log.WithField("id", i).Debug("Spawning VU")
			vuCtx, vuCancel := context.WithCancel(ctx)
			vus = append(vus, vuCancel)
			go func() {
				defer func() {
					if v := recover(); v != nil {
						switch err := v.(type) {
						case speedboat.FlowControl:
							switch err {
							case speedboat.AbortTest:
								log.Error("Test aborted")
								cancel()
							}
						default:
							log.WithFields(log.Fields{
								"id":    i,
								"error": err,
							}).Error("VU crashed!")
						}
					}
				}()
				runner.RunVU(vuCtx, t, len(vus))
			}()
		}
		for i := len(vus); i > scale; i-- {
			log.WithField("id", i-1).Debug("Dropping VU")
			vus[i-1]()
			vus = vus[:i-1]
		}
		mVUs.Int(len(vus))
	}

	// Wait until the end of the test
	<-ctx.Done()

	// Print final metrics
	if logMetrics {
		metricsMutex.Lock()
		printMetrics(metricsLogger)
		metricsMutex.Unlock()
	}
	commitMetrics()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringSliceFlag{
			Name:  "output, o",
			Usage: "Output metrics to a file or database",
		},
		cli.BoolFlag{
			Name:  "log, l",
			Usage: "Log metrics to stdout",
		},
		cli.BoolFlag{
			Name:  "json",
			Usage: "Format log messages as JSON",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		configureSampler(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	_ "expvar"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/facebookgo/httpcontrol"

	version "bosun.org/_version"

	"bosun.org/cmd/tsdbrelay/denormalize"
	"bosun.org/collect"
	"bosun.org/metadata"
	"bosun.org/opentsdb"
	"bosun.org/slog"
	"bosun.org/util"
)

var (
	listenAddr      = flag.String("l", ":4242", "Listen address.")
	bosunServer     = flag.String("b", "bosun", "Target Bosun server. Can specify port with host:port.")
	secondaryRelays = flag.String("r", "", "Additional relays to send data to. Intended for secondary data center replication. Only response from primary tsdb server wil be relayed to clients.")
	tsdbServer      = flag.String("t", "", "Target OpenTSDB server. Can specify port with host:port.")
	logVerbose      = flag.Bool("v", false, "enable verbose logging")
	toDenormalize   = flag.String("denormalize", "", "List of metrics to denormalize. Comma seperated list of `metric__tagname__tagname` rules. Will be translated to `__tagvalue.tagvalue.metric`")
	flagVersion     = flag.Bool("version", false, "Prints the version and exits.")

	redisHost = flag.String("redis", "", "redis host for aggregating external counters")
	redisDb   = flag.Int("db", 0, "redis db to use for counters")
)

var (
	tsdbPutURL    string
	bosunIndexURL string

	denormalizationRules map[string]*denormalize.DenormalizationRule

	relayPutUrls []string

	tags = opentsdb.TagSet{}
)

type tsdbrelayHTTPTransport struct {
	UserAgent string
	http.RoundTripper
}

func (t *tsdbrelayHTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.Header.Get("User-Agent") == "" {
		req.Header.Add("User-Agent", t.UserAgent)
	}
	if req.Header.Get("X-Access-Token") == "" {
		req.Header.Add("X-Access-Token", t.UserAgent)
	}
	return t.RoundTripper.RoundTrip(req)
}

func init() {
	client := &http.Client{
		Transport: &tsdbrelayHTTPTransport{
			"Tsdbrelay/" + version.ShortVersion(),
			&httpcontrol.Transport{
				RequestTimeout: time.Minute,
			},
		},
	}
	http.DefaultClient = client
	collect.DefaultClient = client
}

func main() {
	var err error
	myHost, err = os.Hostname()
	if err != nil || myHost == "" {
		myHost = "tsdbrelay"
	}

	flag.Parse()
	if *flagVersion {
		fmt.Println(version.GetVersionInfo("tsdbrelay"))
		os.Exit(0)
	}
	if *bosunServer == "" || *tsdbServer == "" {
		slog.Fatal("must specify both bosun and tsdb server")
	}
	slog.Infoln(version.GetVersionInfo("tsdbrelay"))
	slog.Infoln("listen on", *listenAddr)
	slog.Infoln("relay to bosun at", *bosunServer)
	slog.Infoln("relay to tsdb at", *tsdbServer)
	if *toDenormalize != "" {
		var err error
		denormalizationRules, err = denormalize.ParseDenormalizationRules(*toDenormalize)
		if err != nil {
			slog.Fatal(err)
		}
	}

	tsdbURL := &url.URL{
		Scheme: "http",
		Host:   *tsdbServer,
	}

	u := url.URL{
		Scheme: "http",
		Host:   *tsdbServer,
		Path:   "/api/put",
	}
	tsdbPutURL = u.String()
	bosunURL := &url.URL{
		Scheme: "http",
		Host:   *bosunServer,
	}
	u = url.URL{
		Scheme: "http",
		Host:   *bosunServer,
		Path:   "/api/index",
	}
	bosunIndexURL = u.String()
	if *secondaryRelays != "" {
		for _, rUrl := range strings.Split(*secondaryRelays, ",") {
			u = url.URL{
				Scheme: "http",
				Host:   rUrl,
				Path:   "/api/put",
			}
			relayPutUrls = append(relayPutUrls, u.String())
		}
	}

	tsdbProxy := util.NewSingleHostProxy(tsdbURL)
	bosunProxy := util.NewSingleHostProxy(bosunURL)
	rp := &relayProxy{
		TSDBProxy:  tsdbProxy,
		BosunProxy: bosunProxy,
	}
	http.HandleFunc("/api/put", func(w http.ResponseWriter, r *http.Request) {
		rp.relayPut(w, r, true)
	})
	if *redisHost != "" {
		http.HandleFunc("/api/count", collect.HandleCounterPut(*redisHost, *redisDb))
	}
	http.HandleFunc("/api/metadata/put", func(w http.ResponseWriter, r *http.Request) {
		rp.relayMetadata(w, r)
	})
	http.Handle("/", tsdbProxy)

	collectUrl := &url.URL{
		Scheme: "http",
		Host:   *listenAddr,
		Path:   "/api/put",
	}
	if err = collect.Init(collectUrl, "tsdbrelay"); err != nil {
		slog.Fatal(err)
	}
	if err := metadata.Init(collectUrl, false); err != nil {
		slog.Fatal(err)
	}
	// Make sure these get zeroed out instead of going unknown on restart
	collect.Add("puts.relayed", tags, 0)
	collect.Add("puts.error", tags, 0)
	collect.Add("metadata.relayed", tags, 0)
	collect.Add("metadata.error", tags, 0)
	collect.Add("additional.puts.relayed", tags, 0)
	collect.Add("additional.puts.error", tags, 0)
	metadata.AddMetricMeta("tsdbrelay.puts.relayed", metadata.Counter, metadata.Count, "Number of successful puts relayed to opentsdb target")
	metadata.AddMetricMeta("tsdbrelay.puts.error", metadata.Counter, metadata.Count, "Number of puts that could not be relayed to opentsdb target")
	metadata.AddMetricMeta("tsdbrelay.metadata.relayed", metadata.Counter, metadata.Count, "Number of successful metadata puts relayed to bosun target")
	metadata.AddMetricMeta("tsdbrelay.metadata.error", metadata.Counter, metadata.Count, "Number of metadata puts that could not be relayed to bosun target")
	metadata.AddMetricMeta("tsdbrelay.additional.puts.relayed", metadata.Counter, metadata.Count, "Number of successful puts relayed to additional targets")
	metadata.AddMetricMeta("tsdbrelay.additional.puts.error", metadata.Counter, metadata.Count, "Number of puts that could not be relayed to additional targets")
	slog.Fatal(http.ListenAndServe(*listenAddr, nil))
}

func verbose(format string, a ...interface{}) {
	if *logVerbose {
		slog.Infof(format, a...)
	}
}

type relayProxy struct {
	TSDBProxy  *httputil.ReverseProxy
	BosunProxy *httputil.ReverseProxy
}

type passthru struct {
	io.ReadCloser
	buf bytes.Buffer
}

func (p *passthru) Read(b []byte) (int, error) {
	n, err := p.ReadCloser.Read(b)
	p.buf.Write(b[:n])
	return n, err
}

type relayWriter struct {
	http.ResponseWriter
	code int
}

func (rw *relayWriter) WriteHeader(code int) {
	rw.code = code
	rw.ResponseWriter.WriteHeader(code)
}

var (
	relayHeader = "X-Relayed-From"
	encHeader   = "Content-Encoding"
	typeHeader  = "Content-Type"
	myHost      string
)

func (rp *relayProxy) relayPut(responseWriter http.ResponseWriter, r *http.Request, parse bool) {
	isRelayed := r.Header.Get(relayHeader) != ""
	reader := &passthru{ReadCloser: r.Body}
	r.Body = reader
	w := &relayWriter{ResponseWriter: responseWriter}
	rp.TSDBProxy.ServeHTTP(w, r)
	if w.code/100 != 2 {
		verbose("relayPut got status %d", w.code)
		collect.Add("puts.error", tags, 1)
		return
	}
	verbose("relayed to tsdb")
	collect.Add("puts.relayed", tags, 1)
	// Send to bosun in a separate go routine so we can end the source's request.
	go func() {
		body := bytes.NewBuffer(reader.buf.Bytes())
		req, err := http.NewRequest(r.Method, bosunIndexURL, body)
		if err != nil {
			verbose("bosun connect error: %v", err)
			return
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			verbose("bosun relay error: %v", err)
			return
		}
		resp.Body.Close()
		verbose("bosun relay success")
	}()
	// Parse and denormalize datapoints
	if !isRelayed && parse && denormalizationRules != nil {
		go rp.denormalize(bytes.NewReader(reader.buf.Bytes()))
	}

	if !isRelayed && len(relayPutUrls) > 0 {
		go func() {
			for _, relayUrl := range relayPutUrls {
				body := bytes.NewBuffer(reader.buf.Bytes())
				req, err := http.NewRequest(r.Method, relayUrl, body)
				if err != nil {
					verbose("relayPutUrls connect error: %v", err)
					collect.Add("additional.puts.error", tags, 1)
					continue
				}
				if contenttype := r.Header.Get(typeHeader); contenttype != "" {
					req.Header.Set(typeHeader, contenttype)
				}
				if encoding := r.Header.Get(encHeader); encoding != "" {
					req.Header.Set(encHeader, encoding)
				}
				req.Header.Add(relayHeader, myHost)
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					verbose("secondary relay error: %v", err)
					collect.Add("additional.puts.error", tags, 1)
					continue
				}
				resp.Body.Close()
				verbose("secondary relay success")
				collect.Add("additional.puts.relayed", tags, 1)
			}
		}()
	}
}

func (rp *relayProxy) denormalize(body io.Reader) {
	gReader, err := gzip.NewReader(body)
	if err != nil {
		verbose("error making gzip reader: %v", err)
		return
	}
	decoder := json.NewDecoder(gReader)
	dps := []*opentsdb.DataPoint{}
	err = decoder.Decode(&dps)
	if err != nil {
		verbose("error decoding data points: %v", err)
		return
	}
	relayDps := []*opentsdb.DataPoint{}
	for _, dp := range dps {
		if rule, ok := denormalizationRules[dp.Metric]; ok {
			if err = rule.Translate(dp); err == nil {
				relayDps = append(relayDps, dp)
			} else {
				verbose("error translating points: %v", err.Error())
			}
		}
	}
	if len(relayDps) == 0 {
		return
	}
	buf := &bytes.Buffer{}
	gWriter := gzip.NewWriter(buf)
	encoder := json.NewEncoder(gWriter)
	err = encoder.Encode(relayDps)
	if err != nil {
		verbose("error encoding denormalized data points: %v", err)
		return
	}
	if err = gWriter.Close(); err != nil {
		verbose("error zipping denormalized data points: %v", err)
		return
	}
	req, err := http.NewRequest("POST", tsdbPutURL, buf)
	if err != nil {
		verbose("error posting denormalized data points: %v", err)
		return
	}
	req.Header.Set(typeHeader, "application/json")
	req.Header.Set(encHeader, "gzip")

	responseWriter := httptest.NewRecorder()
	rp.relayPut(responseWriter, req, false)

	verbose("relayed %d denormalized data points. Tsdb response: %d", len(relayDps), responseWriter.Code)
}

func (rp *relayProxy) relayMetadata(responseWriter http.ResponseWriter, r *http.Request) {
	reader := &passthru{ReadCloser: r.Body}
	r.Body = reader
	w := &relayWriter{ResponseWriter: responseWriter}
	rp.BosunProxy.ServeHTTP(w, r)
	if w.code != 204 {
		verbose("relayMetadata got status %d", w.code)
		collect.Add("metadata.error", tags, 1)
		return
	}
	verbose("relayed metadata to bosun")
	collect.Add("metadata.relayed", tags, 1)
	if r.Header.Get(relayHeader) != "" {
		return
	}
	if len(relayPutUrls) != 0 {
		go func() {
			for _, relayUrl := range relayPutUrls {
				relayUrl = strings.Replace(relayUrl, "/put", "/metadata/put", 1)
				body := bytes.NewBuffer(reader.buf.Bytes())
				req, err := http.NewRequest(r.Method, relayUrl, body)
				if err != nil {
					verbose("metadata relayPutUrls error %v", err)
					continue
				}
				if contenttype := r.Header.Get(typeHeader); contenttype != "" {
					req.Header.Set(typeHeader, contenttype)
				}
				if encoding := r.Header.Get(encHeader); encoding != "" {
					req.Header.Set(encHeader, encoding)
				}
				req.Header.Add(relayHeader, myHost)
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					verbose("secondary relay metadata error: %v", err)
					continue
				}
				resp.Body.Close()
				verbose("secondary relay metadata success")
			}
		}()
	}
}
<commit_msg>cmd/tsdbrelay: copy X-Access-Token in correct place<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	_ "expvar"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/facebookgo/httpcontrol"

	version "bosun.org/_version"

	"bosun.org/cmd/tsdbrelay/denormalize"
	"bosun.org/collect"
	"bosun.org/metadata"
	"bosun.org/opentsdb"
	"bosun.org/slog"
	"bosun.org/util"
)

var (
	listenAddr      = flag.String("l", ":4242", "Listen address.")
	bosunServer     = flag.String("b", "bosun", "Target Bosun server. Can specify port with host:port.")
	secondaryRelays = flag.String("r", "", "Additional relays to send data to. Intended for secondary data center replication. Only response from primary tsdb server wil be relayed to clients.")
	tsdbServer      = flag.String("t", "", "Target OpenTSDB server. Can specify port with host:port.")
	logVerbose      = flag.Bool("v", false, "enable verbose logging")
	toDenormalize   = flag.String("denormalize", "", "List of metrics to denormalize. Comma seperated list of `metric__tagname__tagname` rules. Will be translated to `__tagvalue.tagvalue.metric`")
	flagVersion     = flag.Bool("version", false, "Prints the version and exits.")

	redisHost = flag.String("redis", "", "redis host for aggregating external counters")
	redisDb   = flag.Int("db", 0, "redis db to use for counters")
)

var (
	tsdbPutURL    string
	bosunIndexURL string

	denormalizationRules map[string]*denormalize.DenormalizationRule

	relayPutUrls []string

	tags = opentsdb.TagSet{}
)

type tsdbrelayHTTPTransport struct {
	UserAgent string
	http.RoundTripper
}

func (t *tsdbrelayHTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.Header.Get("User-Agent") == "" {
		req.Header.Add("User-Agent", t.UserAgent)
	}
	return t.RoundTripper.RoundTrip(req)
}

func init() {
	client := &http.Client{
		Transport: &tsdbrelayHTTPTransport{
			"Tsdbrelay/" + version.ShortVersion(),
			&httpcontrol.Transport{
				RequestTimeout: time.Minute,
			},
		},
	}
	http.DefaultClient = client
	collect.DefaultClient = client
}

func main() {
	var err error
	myHost, err = os.Hostname()
	if err != nil || myHost == "" {
		myHost = "tsdbrelay"
	}

	flag.Parse()
	if *flagVersion {
		fmt.Println(version.GetVersionInfo("tsdbrelay"))
		os.Exit(0)
	}
	if *bosunServer == "" || *tsdbServer == "" {
		slog.Fatal("must specify both bosun and tsdb server")
	}
	slog.Infoln(version.GetVersionInfo("tsdbrelay"))
	slog.Infoln("listen on", *listenAddr)
	slog.Infoln("relay to bosun at", *bosunServer)
	slog.Infoln("relay to tsdb at", *tsdbServer)
	if *toDenormalize != "" {
		var err error
		denormalizationRules, err = denormalize.ParseDenormalizationRules(*toDenormalize)
		if err != nil {
			slog.Fatal(err)
		}
	}

	tsdbURL := &url.URL{
		Scheme: "http",
		Host:   *tsdbServer,
	}

	u := url.URL{
		Scheme: "http",
		Host:   *tsdbServer,
		Path:   "/api/put",
	}
	tsdbPutURL = u.String()
	bosunURL := &url.URL{
		Scheme: "http",
		Host:   *bosunServer,
	}
	u = url.URL{
		Scheme: "http",
		Host:   *bosunServer,
		Path:   "/api/index",
	}
	bosunIndexURL = u.String()
	if *secondaryRelays != "" {
		for _, rUrl := range strings.Split(*secondaryRelays, ",") {
			u = url.URL{
				Scheme: "http",
				Host:   rUrl,
				Path:   "/api/put",
			}
			relayPutUrls = append(relayPutUrls, u.String())
		}
	}

	tsdbProxy := util.NewSingleHostProxy(tsdbURL)
	bosunProxy := util.NewSingleHostProxy(bosunURL)
	rp := &relayProxy{
		TSDBProxy:  tsdbProxy,
		BosunProxy: bosunProxy,
	}
	http.HandleFunc("/api/put", func(w http.ResponseWriter, r *http.Request) {
		rp.relayPut(w, r, true)
	})
	if *redisHost != "" {
		http.HandleFunc("/api/count", collect.HandleCounterPut(*redisHost, *redisDb))
	}
	http.HandleFunc("/api/metadata/put", func(w http.ResponseWriter, r *http.Request) {
		rp.relayMetadata(w, r)
	})
	http.Handle("/", tsdbProxy)

	collectUrl := &url.URL{
		Scheme: "http",
		Host:   *listenAddr,
		Path:   "/api/put",
	}
	if err = collect.Init(collectUrl, "tsdbrelay"); err != nil {
		slog.Fatal(err)
	}
	if err := metadata.Init(collectUrl, false); err != nil {
		slog.Fatal(err)
	}
	// Make sure these get zeroed out instead of going unknown on restart
	collect.Add("puts.relayed", tags, 0)
	collect.Add("puts.error", tags, 0)
	collect.Add("metadata.relayed", tags, 0)
	collect.Add("metadata.error", tags, 0)
	collect.Add("additional.puts.relayed", tags, 0)
	collect.Add("additional.puts.error", tags, 0)
	metadata.AddMetricMeta("tsdbrelay.puts.relayed", metadata.Counter, metadata.Count, "Number of successful puts relayed to opentsdb target")
	metadata.AddMetricMeta("tsdbrelay.puts.error", metadata.Counter, metadata.Count, "Number of puts that could not be relayed to opentsdb target")
	metadata.AddMetricMeta("tsdbrelay.metadata.relayed", metadata.Counter, metadata.Count, "Number of successful metadata puts relayed to bosun target")
	metadata.AddMetricMeta("tsdbrelay.metadata.error", metadata.Counter, metadata.Count, "Number of metadata puts that could not be relayed to bosun target")
	metadata.AddMetricMeta("tsdbrelay.additional.puts.relayed", metadata.Counter, metadata.Count, "Number of successful puts relayed to additional targets")
	metadata.AddMetricMeta("tsdbrelay.additional.puts.error", metadata.Counter, metadata.Count, "Number of puts that could not be relayed to additional targets")
	slog.Fatal(http.ListenAndServe(*listenAddr, nil))
}

func verbose(format string, a ...interface{}) {
	if *logVerbose {
		slog.Infof(format, a...)
	}
}

type relayProxy struct {
	TSDBProxy  *httputil.ReverseProxy
	BosunProxy *httputil.ReverseProxy
}

type passthru struct {
	io.ReadCloser
	buf bytes.Buffer
}

func (p *passthru) Read(b []byte) (int, error) {
	n, err := p.ReadCloser.Read(b)
	p.buf.Write(b[:n])
	return n, err
}

type relayWriter struct {
	http.ResponseWriter
	code int
}

func (rw *relayWriter) WriteHeader(code int) {
	rw.code = code
	rw.ResponseWriter.WriteHeader(code)
}

var (
	relayHeader  = "X-Relayed-From"
	encHeader    = "Content-Encoding"
	typeHeader   = "Content-Type"
	accessHeader = "X-Access-Header"
	myHost       string
)

func (rp *relayProxy) relayPut(responseWriter http.ResponseWriter, r *http.Request, parse bool) {
	isRelayed := r.Header.Get(relayHeader) != ""
	reader := &passthru{ReadCloser: r.Body}
	r.Body = reader
	w := &relayWriter{ResponseWriter: responseWriter}
	rp.TSDBProxy.ServeHTTP(w, r)
	if w.code/100 != 2 {
		verbose("relayPut got status %d", w.code)
		collect.Add("puts.error", tags, 1)
		return
	}
	verbose("relayed to tsdb")
	collect.Add("puts.relayed", tags, 1)
	// Send to bosun in a separate go routine so we can end the source's request.
	go func() {
		body := bytes.NewBuffer(reader.buf.Bytes())
		req, err := http.NewRequest(r.Method, bosunIndexURL, body)
		if err != nil {
			verbose("bosun connect error: %v", err)
			return
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			verbose("bosun relay error: %v", err)
			return
		}
		resp.Body.Close()
		verbose("bosun relay success")
	}()
	// Parse and denormalize datapoints
	if !isRelayed && parse && denormalizationRules != nil {
		go rp.denormalize(bytes.NewReader(reader.buf.Bytes()))
	}

	if !isRelayed && len(relayPutUrls) > 0 {
		go func() {
			for _, relayUrl := range relayPutUrls {
				body := bytes.NewBuffer(reader.buf.Bytes())
				req, err := http.NewRequest(r.Method, relayUrl, body)
				if err != nil {
					verbose("relayPutUrls connect error: %v", err)
					collect.Add("additional.puts.error", tags, 1)
					continue
				}
				if contenttype := r.Header.Get(typeHeader); contenttype != "" {
					req.Header.Set(typeHeader, contenttype)
				}
				if access := r.Header.Get(accessHeader); access != "" {
					req.Header.Set(accessHeader, access)
				}
				if encoding := r.Header.Get(encHeader); encoding != "" {
					req.Header.Set(encHeader, encoding)
				}
				req.Header.Add(relayHeader, myHost)
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					verbose("secondary relay error: %v", err)
					collect.Add("additional.puts.error", tags, 1)
					continue
				}
				resp.Body.Close()
				verbose("secondary relay success")
				collect.Add("additional.puts.relayed", tags, 1)
			}
		}()
	}
}

func (rp *relayProxy) denormalize(body io.Reader) {
	gReader, err := gzip.NewReader(body)
	if err != nil {
		verbose("error making gzip reader: %v", err)
		return
	}
	decoder := json.NewDecoder(gReader)
	dps := []*opentsdb.DataPoint{}
	err = decoder.Decode(&dps)
	if err != nil {
		verbose("error decoding data points: %v", err)
		return
	}
	relayDps := []*opentsdb.DataPoint{}
	for _, dp := range dps {
		if rule, ok := denormalizationRules[dp.Metric]; ok {
			if err = rule.Translate(dp); err == nil {
				relayDps = append(relayDps, dp)
			} else {
				verbose("error translating points: %v", err.Error())
			}
		}
	}
	if len(relayDps) == 0 {
		return
	}
	buf := &bytes.Buffer{}
	gWriter := gzip.NewWriter(buf)
	encoder := json.NewEncoder(gWriter)
	err = encoder.Encode(relayDps)
	if err != nil {
		verbose("error encoding denormalized data points: %v", err)
		return
	}
	if err = gWriter.Close(); err != nil {
		verbose("error zipping denormalized data points: %v", err)
		return
	}
	req, err := http.NewRequest("POST", tsdbPutURL, buf)
	if err != nil {
		verbose("error posting denormalized data points: %v", err)
		return
	}
	req.Header.Set(typeHeader, "application/json")
	req.Header.Set(encHeader, "gzip")

	responseWriter := httptest.NewRecorder()
	rp.relayPut(responseWriter, req, false)

	verbose("relayed %d denormalized data points. Tsdb response: %d", len(relayDps), responseWriter.Code)
}

func (rp *relayProxy) relayMetadata(responseWriter http.ResponseWriter, r *http.Request) {
	reader := &passthru{ReadCloser: r.Body}
	r.Body = reader
	w := &relayWriter{ResponseWriter: responseWriter}
	rp.BosunProxy.ServeHTTP(w, r)
	if w.code != 204 {
		verbose("relayMetadata got status %d", w.code)
		collect.Add("metadata.error", tags, 1)
		return
	}
	verbose("relayed metadata to bosun")
	collect.Add("metadata.relayed", tags, 1)
	if r.Header.Get(relayHeader) != "" {
		return
	}
	if len(relayPutUrls) != 0 {
		go func() {
			for _, relayUrl := range relayPutUrls {
				relayUrl = strings.Replace(relayUrl, "/put", "/metadata/put", 1)
				body := bytes.NewBuffer(reader.buf.Bytes())
				req, err := http.NewRequest(r.Method, relayUrl, body)
				if err != nil {
					verbose("metadata relayPutUrls error %v", err)
					continue
				}
				if contenttype := r.Header.Get(typeHeader); contenttype != "" {
					req.Header.Set(typeHeader, contenttype)
				}
				if access := r.Header.Get(accessHeader); access != "" {
					req.Header.Set(accessHeader, access)
				}
				if encoding := r.Header.Get(encHeader); encoding != "" {
					req.Header.Set(encHeader, encoding)
				}
				req.Header.Add(relayHeader, myHost)
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					verbose("secondary relay metadata error: %v", err)
					continue
				}
				resp.Body.Close()
				verbose("secondary relay metadata success")
			}
		}()
	}
}
<|endoftext|>


<commit_before>/**
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * 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 cmds

import (
	"fmt"

	"github.com/fabric8io/gofabric8/client"
	"github.com/fabric8io/gofabric8/util"
	"github.com/spf13/cobra"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

type cleanUpJenkinsFlags struct {
	confirm bool
}

// NewCmdCleanUpJenkins delete files in the tenants content repository
func NewCmdCleanUpJenkins(f *cmdutil.Factory) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "jenkins",
		Short:   "Deletes all the jenkins jobs in your tenant Jenkins service",
		Long:    `Deletes all the jenkins jobs in your tenant Jenkins service`,
		Aliases: []string{"content-repository"},

		Run: func(cmd *cobra.Command, args []string) {
			p := cleanUpJenkinsFlags{}
			if cmd.Flags().Lookup(yesFlag).Value.String() == "true" {
				p.confirm = true
			}
			err := p.cleanUpJenkins(f)
			if err != nil {
				util.Fatalf("%s\n", err)
			}
			return
		},
	}
	return cmd
}

func (p *cleanUpJenkinsFlags) cleanUpJenkins(f *cmdutil.Factory) error {
	c, cfg := client.NewClient(f)
	ns, _, _ := f.DefaultNamespace()
	oc, _ := client.NewOpenShiftClient(cfg)
	initSchema()

	userNS, err := detectCurrentUserNamespace(ns, c, oc)
	if err != nil {
		return err
	}
	jenkinsNS := fmt.Sprintf("%s-jenkins", userNS)

	if !p.confirm {
		confirm := ""
		util.Warn("WARNING this is destructive and will remove ALL of the jenkins jobs\n")
		util.Info("for your tenant: ")
		util.Successf("%s", userNS)
		util.Info(" running in namespace: ")
		util.Successf("%s\n", jenkinsNS)
		util.Warn("\nContinue [y/N]: ")
		fmt.Scanln(&confirm)
		if confirm != "y" {
			util.Warn("Aborted\n")
			return nil
		}
	}
	util.Info("Cleaning jenkins for tenant: ")
	util.Successf("%s", userNS)
	util.Info(" running in namespace: ")
	util.Successf("%s\n", jenkinsNS)

	err = ensureDeploymentOrDCHasReplicas(c, oc, jenkinsNS, "jenkins", 1)
	if err != nil {
		return err
	}
	pod, err := waitForReadyPodForDeploymentOrDC(c, oc, jenkinsNS, "jenkins")
	if err != nil {
		return err
	}
	util.Infof("Found running jenkins pod %s\n", pod)

	kubeCLI := "kubectl"
	err = runCommand(kubeCLI, "exec", "-it", pod, "-n", jenkinsNS, "--", "bash", "-c", "rm -rf /var/lib/jenkins/jobs/*")
	if err != nil {
		return err
	}
	err = runCommand(kubeCLI, "delete", "pod", pod, "-n", jenkinsNS)
	if err != nil {
		return err
	}

	// now lets remove the Jenkins Job ConfigMaps
	err = runCommand("oc", "delete", "configmap", "-l", "openshift.io/jenkins=job", "-n", userNS)
	if err != nil {
		return err
	}

	// now lets remove the BuildConfigs and Builds
	err = runCommand("oc", "delete", "buildconfig", "--all", "-n", userNS)
	if err != nil {
		return err
	}
	err = runCommand("oc", "delete", "build", "--all", "-n", userNS)

	if err == nil {
		util.Info("Completed!\n")
	}
	return err
}
<commit_msg>avoid kubectl for clean jenkins<commit_after>/**
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * 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 cmds

import (
	"fmt"

	"github.com/fabric8io/gofabric8/client"
	"github.com/fabric8io/gofabric8/util"
	"github.com/spf13/cobra"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

type cleanUpJenkinsFlags struct {
	confirm bool
}

// NewCmdCleanUpJenkins delete files in the tenants content repository
func NewCmdCleanUpJenkins(f *cmdutil.Factory) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "jenkins",
		Short:   "Deletes all the jenkins jobs in your tenant Jenkins service",
		Long:    `Deletes all the jenkins jobs in your tenant Jenkins service`,
		Aliases: []string{"content-repository"},

		Run: func(cmd *cobra.Command, args []string) {
			p := cleanUpJenkinsFlags{}
			if cmd.Flags().Lookup(yesFlag).Value.String() == "true" {
				p.confirm = true
			}
			err := p.cleanUpJenkins(f)
			if err != nil {
				util.Fatalf("%s\n", err)
			}
			return
		},
	}
	return cmd
}

func (p *cleanUpJenkinsFlags) cleanUpJenkins(f *cmdutil.Factory) error {
	c, cfg := client.NewClient(f)
	ns, _, _ := f.DefaultNamespace()
	oc, _ := client.NewOpenShiftClient(cfg)
	initSchema()

	userNS, err := detectCurrentUserNamespace(ns, c, oc)
	if err != nil {
		return err
	}
	jenkinsNS := fmt.Sprintf("%s-jenkins", userNS)

	if !p.confirm {
		confirm := ""
		util.Warn("WARNING this is destructive and will remove ALL of the jenkins jobs\n")
		util.Info("for your tenant: ")
		util.Successf("%s", userNS)
		util.Info(" running in namespace: ")
		util.Successf("%s\n", jenkinsNS)
		util.Warn("\nContinue [y/N]: ")
		fmt.Scanln(&confirm)
		if confirm != "y" {
			util.Warn("Aborted\n")
			return nil
		}
	}
	util.Info("Cleaning jenkins for tenant: ")
	util.Successf("%s", userNS)
	util.Info(" running in namespace: ")
	util.Successf("%s\n", jenkinsNS)

	err = ensureDeploymentOrDCHasReplicas(c, oc, jenkinsNS, "jenkins", 1)
	if err != nil {
		return err
	}
	pod, err := waitForReadyPodForDeploymentOrDC(c, oc, jenkinsNS, "jenkins")
	if err != nil {
		return err
	}
	util.Infof("Found running jenkins pod %s\n", pod)

	kubeCLI := "oc"
	err = runCommand(kubeCLI, "exec", "-it", pod, "-n", jenkinsNS, "--", "bash", "-c", "rm -rf /var/lib/jenkins/jobs/*")
	if err != nil {
		return err
	}
	err = runCommand(kubeCLI, "delete", "pod", pod, "-n", jenkinsNS)
	if err != nil {
		return err
	}

	// now lets remove the Jenkins Job ConfigMaps
	err = runCommand("oc", "delete", "configmap", "-l", "openshift.io/jenkins=job", "-n", userNS)
	if err != nil {
		return err
	}

	// now lets remove the BuildConfigs and Builds
	err = runCommand("oc", "delete", "buildconfig", "--all", "-n", userNS)
	if err != nil {
		return err
	}
	err = runCommand("oc", "delete", "build", "--all", "-n", userNS)

	if err == nil {
		util.Info("Completed!\n")
	}
	return err
}
<|endoftext|>


<commit_before>package html

import (
	"github.com/tdewolff/parse"
	"github.com/tdewolff/parse/html"
)

type tokenBuffer struct {
	tokenizer *html.Tokenizer

	buf []token
	pos int
}

func newTokenBuffer(tokenizer *html.Tokenizer) *tokenBuffer {
	return &tokenBuffer{
		tokenizer: tokenizer,
		buf:       make([]token, 0, 8),
	}
}

func (z *tokenBuffer) Read(p []token) int {
	for i := 0; i < len(p); i++ {
		tt, data := z.tokenizer.Next()
		if !z.tokenizer.IsEOF() {
			data = parse.Copy(data)
		}

		var attrVal []byte
		var hash html.Hash
		if tt == html.AttributeToken {
			attrVal = z.tokenizer.AttrVal()
			if !z.tokenizer.IsEOF() {
				attrVal = parse.Copy(attrVal)
			}
			hash = html.ToHash(data)
		} else if tt == html.StartTagToken || tt == html.EndTagToken {
			hash = z.tokenizer.RawTag()
			if hash == 0 {
				hash = html.ToHash(data)
			}
		}
		p[i] = token{tt, data, attrVal, hash}
		if tt == html.ErrorToken {
			return i + 1
		}
	}
	return len(p)
}

// Peek returns the ith element and possibly does an allocation.
// Peeking past an error will panic.
func (z *tokenBuffer) Peek(i int) *token {
	end := z.pos + i
	if end >= len(z.buf) {
		c := cap(z.buf)
		d := len(z.buf) - z.pos
		var buf []token
		if 2*d > c {
			buf = make([]token, d, 2*c)
		} else {
			buf = z.buf[:d]
		}
		copy(buf, z.buf[z.pos:])

		n := z.Read(buf[d:cap(buf)])
		end -= z.pos
		z.pos, z.buf = 0, buf[:d+n]
	}
	return &z.buf[end]
}

// Shift returns the first element and advances position.
func (z *tokenBuffer) Shift() *token {
	t := z.Peek(0)
	z.pos++
	return t
}
<commit_msg>Commentary change<commit_after>package html // import "github.com/tdewolff/minify/html"

import (
	"github.com/tdewolff/parse"
	"github.com/tdewolff/parse/html"
)

type tokenBuffer struct {
	tokenizer *html.Tokenizer

	buf []token
	pos int
}

func newTokenBuffer(tokenizer *html.Tokenizer) *tokenBuffer {
	return &tokenBuffer{
		tokenizer: tokenizer,
		buf:       make([]token, 0, 8),
	}
}

func (z *tokenBuffer) Read(p []token) int {
	for i := 0; i < len(p); i++ {
		tt, data := z.tokenizer.Next()
		if !z.tokenizer.IsEOF() {
			data = parse.Copy(data)
		}

		var attrVal []byte
		var hash html.Hash
		if tt == html.AttributeToken {
			attrVal = z.tokenizer.AttrVal()
			if !z.tokenizer.IsEOF() {
				attrVal = parse.Copy(attrVal)
			}
			hash = html.ToHash(data)
		} else if tt == html.StartTagToken || tt == html.EndTagToken {
			hash = z.tokenizer.RawTag()
			if hash == 0 {
				hash = html.ToHash(data)
			}
		}
		p[i] = token{tt, data, attrVal, hash}
		if tt == html.ErrorToken {
			return i + 1
		}
	}
	return len(p)
}

// Peek returns the ith element and possibly does an allocation.
// Peeking past an error will panic.
func (z *tokenBuffer) Peek(i int) *token {
	end := z.pos + i
	if end >= len(z.buf) {
		c := cap(z.buf)
		d := len(z.buf) - z.pos
		var buf []token
		if 2*d > c {
			buf = make([]token, d, 2*c)
		} else {
			buf = z.buf[:d]
		}
		copy(buf, z.buf[z.pos:])

		n := z.Read(buf[d:cap(buf)])
		end -= z.pos
		z.pos, z.buf = 0, buf[:d+n]
	}
	return &z.buf[end]
}

// Shift returns the first element and advances position.
func (z *tokenBuffer) Shift() *token {
	t := z.Peek(0)
	z.pos++
	return t
}
<|endoftext|>