<commit_before>package runner

import (
	"fmt"
	"math/rand"

	"github.com/BurntSushi/toml"
	"github.com/influxdb/influxdb/client"
)

// tag is a struct that contains data
// about a tag for in a series
type tag struct {
	Key   string `toml:"key"`
	Value string `toml:"value"`
}

// tag is a struct that contains data
// about a field for in a series
type field struct {
	Key  string `toml:"key"`
	Type string `toml:"type"`
}

// series is a struct that contains data
// about the series that will be written
// during a stress test
type series struct {
	PointCount  int     `toml:"point_count"`
	Measurement string  `toml:"measurement"`
	SeriesCount int     `toml:"series_count"`
	Tags        []tag   `toml:"tag"`
	Fields      []field `toml:"field"`
}

// write is a struct that contains the business
// logic for the stress test. e.g. where the
// influxdb instance is running, what database
// should points be written into
type write struct {
	Concurrency   int    `toml:"concurrency"`
	BatchSize     int    `toml:"batch_size"`
	BatchInterval string `toml:"batch_interval"`
	Database      string `toml:"database"`
	ResetDatabase bool   `toml:"reset_database"`
	StartingPoint string `toml:"starting_time"`
	Address       string `toml:"address"`
	Precision     string `toml:"precision"`
}

// query is a struct that contains the logic for
// a query that will be ran on during the stress
// test
type query struct {
	Concurrency int    `toml:"concurrency"`
	Measurement string `toml:"measurement"`
	TagKey      string `toml:"tag_key"`
	TimeFrame   string `toml:"time_frame"`
	Statement   string `toml:"statement"`
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	Write   write    `toml:"write"`
	Series  []series `toml:"series"`
	Queries []query  `toml:"query"`
}

// DecodeFile takes a file path for a toml config file
// and returns a pointer to a Config Struct.
func DecodeFile(s string) (*Config, error) {
	t := &Config{}

	if _, err := toml.DecodeFile(s, t); err != nil {
		return nil, err
	}

	return t, nil
}

// seriesIter is a struct that contains a
// series and a count, where count is the
// number of points that have been written
// for the series `s`
type seriesIter struct {
	s     *series
	count int
}

// Iter returns a pointer to a seriesIter
func (s *series) Iter() *seriesIter {
	return &seriesIter{s: s, count: -1}
}

// newTagMap returns a tagset
func (s *series) newTagMap(i int) map[string]string {
	m := map[string]string{}

	for _, tag := range s.Tags {
		m[tag.Key] = fmt.Sprintf("%s-%d", tag.Value, i)
	}

	return m
}

// newFieldMap returns a new field set for
// a given series
func (s *series) newFieldMap() map[string]interface{} {
	m := map[string]interface{}{}

	for _, field := range s.Fields {
		switch field.Type {
		case "float64":
			m[field.Key] = float64(rand.Intn(1000))
		case "int":
			m[field.Key] = rand.Intn(1000)
		case "bool":
			b := rand.Intn(2) == 1
			m[field.Key] = b
		default:
			m[field.Key] = float64(rand.Intn(1000))
		}

	}

	return m
}

// Next returns a new point for a series.
// Currently, there is an off by one bug here.
// When I try to fix it, I end up creating another off
// by one bug.
// BUG: I can either miss one point from n different series,
//      or miss n points from one series. At the moment,
//      I've chosed to go with the n points from one series.
func (iter *seriesIter) Next() (client.Point, bool) {
	iter.count++
	p := client.Point{
		Measurement: iter.s.Measurement,
		Tags:        iter.s.newTagMap(iter.count),
		Fields:      iter.s.newFieldMap(),
	}
	b := iter.count < iter.s.SeriesCount
	return p, b
}
<commit_msg>Remove comment about off by one bug<commit_after>package runner

import (
	"fmt"
	"math/rand"

	"github.com/BurntSushi/toml"
	"github.com/influxdb/influxdb/client"
)

// tag is a struct that contains data
// about a tag for in a series
type tag struct {
	Key   string `toml:"key"`
	Value string `toml:"value"`
}

// tag is a struct that contains data
// about a field for in a series
type field struct {
	Key  string `toml:"key"`
	Type string `toml:"type"`
}

// series is a struct that contains data
// about the series that will be written
// during a stress test
type series struct {
	PointCount  int     `toml:"point_count"`
	Measurement string  `toml:"measurement"`
	SeriesCount int     `toml:"series_count"`
	Tags        []tag   `toml:"tag"`
	Fields      []field `toml:"field"`
}

// write is a struct that contains the business
// logic for the stress test. e.g. where the
// influxdb instance is running, what database
// should points be written into
type write struct {
	Concurrency   int    `toml:"concurrency"`
	BatchSize     int    `toml:"batch_size"`
	BatchInterval string `toml:"batch_interval"`
	Database      string `toml:"database"`
	ResetDatabase bool   `toml:"reset_database"`
	StartingPoint string `toml:"starting_time"`
	Address       string `toml:"address"`
	Precision     string `toml:"precision"`
}

// query is a struct that contains the logic for
// a query that will be ran on during the stress
// test
type query struct {
	Concurrency int    `toml:"concurrency"`
	Measurement string `toml:"measurement"`
	TagKey      string `toml:"tag_key"`
	TimeFrame   string `toml:"time_frame"`
	Statement   string `toml:"statement"`
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	Write   write    `toml:"write"`
	Series  []series `toml:"series"`
	Queries []query  `toml:"query"`
}

// DecodeFile takes a file path for a toml config file
// and returns a pointer to a Config Struct.
func DecodeFile(s string) (*Config, error) {
	t := &Config{}

	if _, err := toml.DecodeFile(s, t); err != nil {
		return nil, err
	}

	return t, nil
}

// seriesIter is a struct that contains a
// series and a count, where count is the
// number of points that have been written
// for the series `s`
type seriesIter struct {
	s     *series
	count int
}

// Iter returns a pointer to a seriesIter
func (s *series) Iter() *seriesIter {
	return &seriesIter{s: s, count: -1}
}

// newTagMap returns a tagset
func (s *series) newTagMap(i int) map[string]string {
	m := map[string]string{}

	for _, tag := range s.Tags {
		m[tag.Key] = fmt.Sprintf("%s-%d", tag.Value, i)
	}

	return m
}

// newFieldMap returns a new field set for
// a given series
func (s *series) newFieldMap() map[string]interface{} {
	m := map[string]interface{}{}

	for _, field := range s.Fields {
		switch field.Type {
		case "float64":
			m[field.Key] = float64(rand.Intn(1000))
		case "int":
			m[field.Key] = rand.Intn(1000)
		case "bool":
			b := rand.Intn(2) == 1
			m[field.Key] = b
		default:
			m[field.Key] = float64(rand.Intn(1000))
		}

	}

	return m
}

// Next returns a new point for a series.
// Currently, there is an off by one bug here.
func (iter *seriesIter) Next() (client.Point, bool) {
	iter.count++
	p := client.Point{
		Measurement: iter.s.Measurement,
		Tags:        iter.s.newTagMap(iter.count),
		Fields:      iter.s.newFieldMap(),
	}
	b := iter.count < iter.s.SeriesCount
	return p, b
}
<|endoftext|>


<commit_before>package runner

import (
	"bytes"
	"fmt"
	"math/rand"
	"time"

	"github.com/BurntSushi/toml"
)

// tag is a struct that contains data
// about a tag for in a series
type tag struct {
	Key   string `toml:"key"`
	Value string `toml:"value"`
}

// tag is a struct that contains data
// about a field for in a series
type field struct {
	Key  string `toml:"key"`
	Type string `toml:"type"`
}

// series is a struct that contains data
// about the series that will be written
// during a stress test
type series struct {
	PointCount  int     `toml:"point_count"`
	Tick        string  `toml:"tick"`
	Jitter      bool    `toml:"jitter"`
	Measurement string  `toml:"measurement"`
	SeriesCount int     `toml:"series_count"`
	TagCount    int     `toml:"tag_count"`
	Tags        []tag   `toml:"tag"`
	Fields      []field `toml:"field"`
}

// write is a struct that contains the business
// logic for the stress test. e.g. where the
// influxdb instance is running, what database
// should points be written into
type write struct {
	Concurrency   int    `toml:"concurrency"`
	BatchSize     int    `toml:"batch_size"`
	BatchInterval string `toml:"batch_interval"`
	Database      string `toml:"database"`
	ResetDatabase bool   `toml:"reset_database"`
	Address       string `toml:"address"`
	Precision     string `toml:"precision"`
	StartDate     string `toml:"start_date"`
}

// query is a struct that contains the logic for
// a query that will be ran on during the stress
// test
type query struct {
	Enabled     bool     `toml:"enabled"`
	Concurrency int      `toml:"concurrency"`
	Aggregates  []string `toml:"aggregates"`
	Fields      []string `toml:"fields"`
}

// measurementQuery is a struct that contains
// the logic that runs a query against a measurement
// over a time period that is specified by
// `Offset`
type measurementQuery struct {
	query
	Offset string `toml:"offset"`
}

// seriesQuery is a struct that contains
// the logic that runs a query against a single
// series
type seriesQuery struct {
	query
	Interval string `toml:"interval"`
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	Write             write            `toml:"write"`
	Series            []series         `toml:"series"`
	MeasurementQuery  measurementQuery `toml:"measurement_query"`
	SeriesQuery       seriesQuery      `toml:"series_query"`
	ChannelBufferSize int              `toml:"channel_buffer_size"`
}

// NewSeries, takes a measurement, and point count,
// and a series count and returns a series
func NewSeries(m string, p int, sc int) series {
	s := series{
		PointCount:  p,
		SeriesCount: sc,
		Tick:        "1s",
		Measurement: m,
		Tags: []tag{
			tag{
				Key:   "host",
				Value: "server",
			},
		},
		Fields: []field{
			field{
				Key: "value",
			},
		},
	}

	return s
}

// NewConfig returns a pointer to a config
// with some default parameters set
func NewConfig() *Config {

	w := write{
		Concurrency:   10,
		BatchSize:     5000,
		BatchInterval: "0s",
		Database:      "stress",
		ResetDatabase: true,
		Address:       "localhost:8086",
		Precision:     "n",
	}

	c := &Config{
		Write: w,
	}

	return c
}

// DecodeFile takes a file path for a toml config file
// and returns a pointer to a Config Struct.
func DecodeFile(s string) (*Config, error) {
	t := &Config{}

	// Decode the toml file
	if _, err := toml.DecodeFile(s, t); err != nil {
		return nil, err
	}

	// Initialize Config struct
	// NOTE: Not happy with the implementation
	// but it will do for now
	for j, srs := range t.Series {
		for i := 0; i < srs.TagCount; i++ {

			tag := tag{
				Key:   fmt.Sprintf("tag-key-%d", i),
				Value: "tag-value",
			}

			srs.Tags = append(srs.Tags, tag)
			fmt.Println(srs)
		}

		t.Series[j] = srs
	}

	return t, nil
}

// seriesIter is a struct that contains a
// series and a count, where count is the
//number of points that have been written
// for the series `s`
type seriesIter struct {
	s         *series
	count     int
	timestamp time.Time
	precision string
}

// writeInterval returns a timestamp for the current time
// interval
func (s *series) writeInterval(i int, start time.Time) time.Time {
	var tick time.Duration
	var j int
	var err error

	tick, err = time.ParseDuration(s.Tick)
	if err != nil {
		panic(err)
	}

	if true {
		j = rand.Intn(int(tick))
	}

	tick = tick*time.Duration(i) + time.Duration(j)

	return start.Add(tick)
}

// Iter returns a pointer to a seriesIter
func (s *series) Iter(i int, start time.Time, p string) *seriesIter {

	return &seriesIter{s: s, count: -1, timestamp: s.writeInterval(i, start), precision: p}
}

// Next returns a new point for a series.
// Currently, there is an off by one bug here.
func (iter *seriesIter) Next() ([]byte, bool) {
	var buf bytes.Buffer
	iter.count++

	buf.Write([]byte(fmt.Sprintf("%v,", iter.s.Measurement)))
	buf.Write(iter.s.newTagSet(iter.count))
	buf.Write([]byte(" "))
	buf.Write(iter.s.newFieldSet(iter.count))
	buf.Write([]byte(" "))

	switch iter.precision {
	case "s":
		buf.Write([]byte(fmt.Sprintf("%v", iter.timestamp.Unix())))
	default:
		buf.Write([]byte(fmt.Sprintf("%v", iter.timestamp.UnixNano())))
	}

	b := iter.count < iter.s.SeriesCount
	byt := buf.Bytes()

	return byt, b
}

// newTagSet returns a byte array representation
// of the tagset for a series
func (s *series) newTagSet(c int) []byte {
	var buf bytes.Buffer
	for _, tag := range s.Tags {
		buf.Write([]byte(fmt.Sprintf("%v=%v-%v,", tag.Key, tag.Value, c)))
	}

	b := buf.Bytes()
	b = b[0 : len(b)-1]

	return b
}

// newFieldSet returns a byte array representation
// of the field-set for a series
func (s *series) newFieldSet(c int) []byte {
	var buf bytes.Buffer

	for _, field := range s.Fields {
		switch field.Type {
		case "float64-flat":
			if rand.Intn(10) > 2 {
				buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, 100)))
			} else {
				buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, 100+rand.Intn(100))))
			}
		case "float64-inc+":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, c+rand.Intn(2))))
		case "float64-inc":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, c)))
		case "float64":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, rand.Intn(1000))))
		case "int":
			buf.Write([]byte(fmt.Sprintf("%v=%vi,", field.Key, rand.Intn(1000))))
		case "bool":
			b := rand.Intn(2) == 1
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, b)))
		default:
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, rand.Intn(1000))))
		}
	}

	b := buf.Bytes()
	b = b[0 : len(b)-1]

	return b
}
<commit_msg>Only add jitter if config option is set<commit_after>package runner

import (
	"bytes"
	"fmt"
	"math/rand"
	"time"

	"github.com/BurntSushi/toml"
)

// tag is a struct that contains data
// about a tag for in a series
type tag struct {
	Key   string `toml:"key"`
	Value string `toml:"value"`
}

// tag is a struct that contains data
// about a field for in a series
type field struct {
	Key  string `toml:"key"`
	Type string `toml:"type"`
}

// series is a struct that contains data
// about the series that will be written
// during a stress test
type series struct {
	PointCount  int     `toml:"point_count"`
	Tick        string  `toml:"tick"`
	Jitter      bool    `toml:"jitter"`
	Measurement string  `toml:"measurement"`
	SeriesCount int     `toml:"series_count"`
	TagCount    int     `toml:"tag_count"`
	Tags        []tag   `toml:"tag"`
	Fields      []field `toml:"field"`
}

// write is a struct that contains the business
// logic for the stress test. e.g. where the
// influxdb instance is running, what database
// should points be written into
type write struct {
	Concurrency   int    `toml:"concurrency"`
	BatchSize     int    `toml:"batch_size"`
	BatchInterval string `toml:"batch_interval"`
	Database      string `toml:"database"`
	ResetDatabase bool   `toml:"reset_database"`
	Address       string `toml:"address"`
	Precision     string `toml:"precision"`
	StartDate     string `toml:"start_date"`
}

// query is a struct that contains the logic for
// a query that will be ran on during the stress
// test
type query struct {
	Enabled     bool     `toml:"enabled"`
	Concurrency int      `toml:"concurrency"`
	Aggregates  []string `toml:"aggregates"`
	Fields      []string `toml:"fields"`
}

// measurementQuery is a struct that contains
// the logic that runs a query against a measurement
// over a time period that is specified by
// `Offset`
type measurementQuery struct {
	query
	Offset string `toml:"offset"`
}

// seriesQuery is a struct that contains
// the logic that runs a query against a single
// series
type seriesQuery struct {
	query
	Interval string `toml:"interval"`
}

// Config is a struct that is passed into the `Run()` function.
type Config struct {
	Write             write            `toml:"write"`
	Series            []series         `toml:"series"`
	MeasurementQuery  measurementQuery `toml:"measurement_query"`
	SeriesQuery       seriesQuery      `toml:"series_query"`
	ChannelBufferSize int              `toml:"channel_buffer_size"`
}

// NewSeries, takes a measurement, and point count,
// and a series count and returns a series
func NewSeries(m string, p int, sc int) series {
	s := series{
		PointCount:  p,
		SeriesCount: sc,
		Tick:        "1s",
		Measurement: m,
		Tags: []tag{
			tag{
				Key:   "host",
				Value: "server",
			},
		},
		Fields: []field{
			field{
				Key: "value",
			},
		},
	}

	return s
}

// NewConfig returns a pointer to a config
// with some default parameters set
func NewConfig() *Config {

	w := write{
		Concurrency:   10,
		BatchSize:     5000,
		BatchInterval: "0s",
		Database:      "stress",
		ResetDatabase: true,
		Address:       "localhost:8086",
		Precision:     "n",
	}

	c := &Config{
		Write: w,
	}

	return c
}

// DecodeFile takes a file path for a toml config file
// and returns a pointer to a Config Struct.
func DecodeFile(s string) (*Config, error) {
	t := &Config{}

	// Decode the toml file
	if _, err := toml.DecodeFile(s, t); err != nil {
		return nil, err
	}

	// Initialize Config struct
	// NOTE: Not happy with the implementation
	// but it will do for now
	for j, srs := range t.Series {
		for i := 0; i < srs.TagCount; i++ {

			tag := tag{
				Key:   fmt.Sprintf("tag-key-%d", i),
				Value: "tag-value",
			}

			srs.Tags = append(srs.Tags, tag)
			fmt.Println(srs)
		}

		t.Series[j] = srs
	}

	return t, nil
}

// seriesIter is a struct that contains a
// series and a count, where count is the
//number of points that have been written
// for the series `s`
type seriesIter struct {
	s         *series
	count     int
	timestamp time.Time
	precision string
}

// writeInterval returns a timestamp for the current time
// interval
func (s *series) writeInterval(i int, start time.Time) time.Time {
	var tick time.Duration
	var j int
	var err error

	tick, err = time.ParseDuration(s.Tick)
	if err != nil {
		panic(err)
	}

	if s.Jitter {
		j = rand.Intn(int(tick))
	}

	tick = tick*time.Duration(i) + time.Duration(j)

	return start.Add(tick)
}

// Iter returns a pointer to a seriesIter
func (s *series) Iter(i int, start time.Time, p string) *seriesIter {

	return &seriesIter{s: s, count: -1, timestamp: s.writeInterval(i, start), precision: p}
}

// Next returns a new point for a series.
// Currently, there is an off by one bug here.
func (iter *seriesIter) Next() ([]byte, bool) {
	var buf bytes.Buffer
	iter.count++

	buf.Write([]byte(fmt.Sprintf("%v,", iter.s.Measurement)))
	buf.Write(iter.s.newTagSet(iter.count))
	buf.Write([]byte(" "))
	buf.Write(iter.s.newFieldSet(iter.count))
	buf.Write([]byte(" "))

	switch iter.precision {
	case "s":
		buf.Write([]byte(fmt.Sprintf("%v", iter.timestamp.Unix())))
	default:
		buf.Write([]byte(fmt.Sprintf("%v", iter.timestamp.UnixNano())))
	}

	b := iter.count < iter.s.SeriesCount
	byt := buf.Bytes()

	return byt, b
}

// newTagSet returns a byte array representation
// of the tagset for a series
func (s *series) newTagSet(c int) []byte {
	var buf bytes.Buffer
	for _, tag := range s.Tags {
		buf.Write([]byte(fmt.Sprintf("%v=%v-%v,", tag.Key, tag.Value, c)))
	}

	b := buf.Bytes()
	b = b[0 : len(b)-1]

	return b
}

// newFieldSet returns a byte array representation
// of the field-set for a series
func (s *series) newFieldSet(c int) []byte {
	var buf bytes.Buffer

	for _, field := range s.Fields {
		switch field.Type {
		case "float64-flat":
			if rand.Intn(10) > 2 {
				buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, 100)))
			} else {
				buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, 100+rand.Intn(100))))
			}
		case "float64-inc+":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, c+rand.Intn(2))))
		case "float64-inc":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, c)))
		case "float64":
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, rand.Intn(1000))))
		case "int":
			buf.Write([]byte(fmt.Sprintf("%v=%vi,", field.Key, rand.Intn(1000))))
		case "bool":
			b := rand.Intn(2) == 1
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, b)))
		default:
			buf.Write([]byte(fmt.Sprintf("%v=%v,", field.Key, rand.Intn(1000))))
		}
	}

	b := buf.Bytes()
	b = b[0 : len(b)-1]

	return b
}
<|endoftext|>


<commit_before>package runner

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/client"
)

// Timer is struct that can be used to track elaspsed time
type Timer struct {
	start time.Time
	end   time.Time
}

// Start returns a Timers start field
func (t *Timer) Start() time.Time {
	return t.start
}

// End returns a Timers end field
func (t *Timer) End() time.Time {
	return t.end
}

// StartTimer sets a timers `start` field to the current time
func (t *Timer) StartTimer() {
	t.start = time.Now()
}

// StopTimer sets a timers `end` field to the current time
func (t *Timer) StopTimer() {
	t.end = time.Now()
}

// Elapsed returns the total elapsed time between the `start`
// and `end` fields on a timer.
func (t *Timer) Elapsed() time.Duration {
	return t.end.Sub(t.start)
}

// NewTimer returns a pointer to a `Timer` struct where the
// timers `start` field has been set to `time.Now()`
func NewTimer() *Timer {
	t := &Timer{}
	t.StartTimer()
	return t
}

// ResponseTime is a struct that contains `Value`
// `Time` pairing.
type ResponseTime struct {
	Value int
	Time  time.Time
}

// newResponseTime returns a new response time
// with value `v` and time `time.Now()`.
func NewResponseTime(v int) ResponseTime {
	r := ResponseTime{Value: v, Time: time.Now()}
	return r
}

type ResponseTimes []ResponseTime

// Implements the `Len` method for the
// sort.Interface type
func (rs ResponseTimes) Len() int {
	return len(rs)
}

// Implements the `Less` method for the
// sort.Interface type
func (rs ResponseTimes) Less(i, j int) bool {
	return rs[i].Value < rs[j].Value
}

// Implements the `Swap` method for the
// sort.Interface type
func (rs ResponseTimes) Swap(i, j int) {
	rs[i], rs[j] = rs[j], rs[i]
}

type Measurements []string

// String returns a string and implements the `String` method for
// the flag.Value interface.
func (ms *Measurements) String() string {
	return fmt.Sprint(*ms)
}

// Set implements the `Set` method for the flag.Value
// interface. Set splits a string of comma separated values
// into a `Measurement`.
func (ms *Measurements) Set(value string) error {
	values := strings.Split(value, ",")
	for _, m := range values {
		*ms = append(*ms, m)
	}
	return nil
}

// newClient returns a pointer to an InfluxDB client for
// a `Config`'s `Address` field. If an error is encountered
// when creating a new client, the function panics.
func (cfg *Config) NewClient() (*client.Client, error) {
	fmt.Printf("CFG: %#v\n", cfg)
	u, _ := url.Parse(fmt.Sprintf("http://%s", cfg.Write.Address))
	c, err := client.NewClient(client.Config{URL: *u})
	if err != nil {
		return nil, err
	}
	return c, nil
}

func resetDB(c *client.Client, database string) error {
	_, err := c.Query(client.Query{
		Command: fmt.Sprintf("DROP DATABASE %s", database),
	})

	if err != nil && !strings.Contains(err.Error(), "database not found") {
		return err
	}

	return nil
}

// Run runs the stress test that is specified by a `Config`.
// It returns the total number of points that were during the test,
// an slice of all of the stress tests response times,
// and the times that the test started at and ended as a `Timer`
func Run(cfg *Config) (totalPoints int, failedRequests int, responseTimes ResponseTimes, timer *Timer) {

	c, err := cfg.NewClient()
	if err != nil {
		panic(err)
	}

	if cfg.Write.ResetDatabase {
		resetDB(c, cfg.Write.Database)
	}

	_, err = c.Query(client.Query{
		Command: fmt.Sprintf("CREATE DATABASE %s", cfg.Write.Database),
	})

	if err != nil && !strings.Contains(err.Error(), "database already exists") {
		fmt.Println(err)
	}

	counter := NewConcurrencyLimiter(cfg.Write.Concurrency)

	var mu sync.Mutex
	var wg sync.WaitGroup
	responseTimes = make(ResponseTimes, 0)

	failedRequests = 0

	totalPoints = 0

	lastSuccess := true

	car := make(chan []client.Point, 10)

	go func(ch chan []client.Point) {
		points := []client.Point{}

		for _, testSeries := range cfg.Series {
			for i := 0; i < testSeries.PointCount; i++ {
				iter := testSeries.Iter(cfg.Write.StartingPoint, i)
				p, ok := iter.Next()
				for ok {
					points = append(points, p)
					if len(points) == cfg.Write.BatchSize {
						ch <- points
						points = []client.Point{}
					}

					p, ok = iter.Next()
				}
			}
		}

		close(car)

	}(car)

	timer = NewTimer()
	defer timer.StopTimer()

	for pnt := range car {
		batch := &client.BatchPoints{
			Database:         cfg.Write.Database,
			WriteConsistency: "any",
			Time:             time.Now(),
			Precision:        cfg.Write.Precision,
			Points:           pnt,
		}

		wg.Add(1)
		counter.Increment()
		totalPoints += len(batch.Points)

		go func(b *client.BatchPoints, total int) {
			st := time.Now()
			if _, err := c.Write(*b); err != nil { // Should retry write if failed
				mu.Lock()
				if lastSuccess {
					fmt.Println("ERROR: ", err.Error())
				}
				failedRequests += 1
				totalPoints -= len(b.Points)
				lastSuccess = false
				mu.Unlock()
			} else {
				mu.Lock()
				if !lastSuccess {
					fmt.Println("success in ", time.Since(st))
				}
				lastSuccess = true
				responseTimes = append(responseTimes, NewResponseTime(int(time.Since(st).Nanoseconds())))
				mu.Unlock()
			}
			batchInterval, _ := time.ParseDuration(cfg.Write.BatchInterval)
			time.Sleep(batchInterval)
			wg.Done()
			counter.Decrement()
			if total%500000 == 0 {
				fmt.Printf("%d total points. %d in %s\n", total, cfg.Write.BatchSize, time.Since(st))
			}
		}(batch, totalPoints)

	}

	wg.Wait()

	return
}

// ConcurrencyLimiter is a go routine safe struct that can be used to
// ensure that no more than a specifid max number of goroutines are
// executing.
type ConcurrencyLimiter struct {
	inc   chan chan struct{}
	dec   chan struct{}
	max   int
	count int
}

// NewConcurrencyLimiter returns a configured limiter that will
// ensure that calls to Increment will block if the max is hit.
func NewConcurrencyLimiter(max int) *ConcurrencyLimiter {
	c := &ConcurrencyLimiter{
		inc: make(chan chan struct{}),
		dec: make(chan struct{}, max),
		max: max,
	}
	go c.handleLimits()
	return c
}

// Increment will increase the count of running goroutines by 1.
// if the number is currently at the max, the call to Increment
// will block until another goroutine decrements.
func (c *ConcurrencyLimiter) Increment() {
	r := make(chan struct{})
	c.inc <- r
	<-r
}

// Decrement will reduce the count of running goroutines by 1
func (c *ConcurrencyLimiter) Decrement() {
	c.dec <- struct{}{}
}

// handleLimits runs in a goroutine to manage the count of
// running goroutines.
func (c *ConcurrencyLimiter) handleLimits() {
	for {
		r := <-c.inc
		if c.count >= c.max {
			<-c.dec
			c.count--
		}
		c.count++
		r <- struct{}{}
	}
}
<commit_msg>Change variable names<commit_after>package runner

import (
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/influxdb/influxdb/client"
)

// Timer is struct that can be used to track elaspsed time
type Timer struct {
	start time.Time
	end   time.Time
}

// Start returns a Timers start field
func (t *Timer) Start() time.Time {
	return t.start
}

// End returns a Timers end field
func (t *Timer) End() time.Time {
	return t.end
}

// StartTimer sets a timers `start` field to the current time
func (t *Timer) StartTimer() {
	t.start = time.Now()
}

// StopTimer sets a timers `end` field to the current time
func (t *Timer) StopTimer() {
	t.end = time.Now()
}

// Elapsed returns the total elapsed time between the `start`
// and `end` fields on a timer.
func (t *Timer) Elapsed() time.Duration {
	return t.end.Sub(t.start)
}

// NewTimer returns a pointer to a `Timer` struct where the
// timers `start` field has been set to `time.Now()`
func NewTimer() *Timer {
	t := &Timer{}
	t.StartTimer()
	return t
}

// ResponseTime is a struct that contains `Value`
// `Time` pairing.
type ResponseTime struct {
	Value int
	Time  time.Time
}

// newResponseTime returns a new response time
// with value `v` and time `time.Now()`.
func NewResponseTime(v int) ResponseTime {
	r := ResponseTime{Value: v, Time: time.Now()}
	return r
}

type ResponseTimes []ResponseTime

// Implements the `Len` method for the
// sort.Interface type
func (rs ResponseTimes) Len() int {
	return len(rs)
}

// Implements the `Less` method for the
// sort.Interface type
func (rs ResponseTimes) Less(i, j int) bool {
	return rs[i].Value < rs[j].Value
}

// Implements the `Swap` method for the
// sort.Interface type
func (rs ResponseTimes) Swap(i, j int) {
	rs[i], rs[j] = rs[j], rs[i]
}

type Measurements []string

// String returns a string and implements the `String` method for
// the flag.Value interface.
func (ms *Measurements) String() string {
	return fmt.Sprint(*ms)
}

// Set implements the `Set` method for the flag.Value
// interface. Set splits a string of comma separated values
// into a `Measurement`.
func (ms *Measurements) Set(value string) error {
	values := strings.Split(value, ",")
	for _, m := range values {
		*ms = append(*ms, m)
	}
	return nil
}

// newClient returns a pointer to an InfluxDB client for
// a `Config`'s `Address` field. If an error is encountered
// when creating a new client, the function panics.
func (cfg *Config) NewClient() (*client.Client, error) {
	fmt.Printf("CFG: %#v\n", cfg)
	u, _ := url.Parse(fmt.Sprintf("http://%s", cfg.Write.Address))
	c, err := client.NewClient(client.Config{URL: *u})
	if err != nil {
		return nil, err
	}
	return c, nil
}

func resetDB(c *client.Client, database string) error {
	_, err := c.Query(client.Query{
		Command: fmt.Sprintf("DROP DATABASE %s", database),
	})

	if err != nil && !strings.Contains(err.Error(), "database not found") {
		return err
	}

	return nil
}

// Run runs the stress test that is specified by a `Config`.
// It returns the total number of points that were during the test,
// an slice of all of the stress tests response times,
// and the times that the test started at and ended as a `Timer`
func Run(cfg *Config) (totalPoints int, failedRequests int, responseTimes ResponseTimes, timer *Timer) {

	c, err := cfg.NewClient()
	if err != nil {
		panic(err)
	}

	if cfg.Write.ResetDatabase {
		resetDB(c, cfg.Write.Database)
	}

	_, err = c.Query(client.Query{
		Command: fmt.Sprintf("CREATE DATABASE %s", cfg.Write.Database),
	})

	if err != nil && !strings.Contains(err.Error(), "database already exists") {
		fmt.Println(err)
	}

	counter := NewConcurrencyLimiter(cfg.Write.Concurrency)

	var mu sync.Mutex
	var wg sync.WaitGroup
	responseTimes = make(ResponseTimes, 0)

	failedRequests = 0

	totalPoints = 0

	lastSuccess := true

	ch := make(chan []client.Point, 100)

	go func() {
		points := []client.Point{}

		for _, testSeries := range cfg.Series {
			for i := 0; i < testSeries.PointCount; i++ {
				iter := testSeries.Iter(cfg.Write.StartingPoint, i)
				p, ok := iter.Next()
				for ok {
					points = append(points, p)
					if len(points) == cfg.Write.BatchSize {
						ch <- points
						points = []client.Point{}
					}

					p, ok = iter.Next()
				}
			}
		}

		close(ch)

	}()

	timer = NewTimer()
	defer timer.StopTimer()

	for pnt := range ch {
		batch := &client.BatchPoints{
			Database:         cfg.Write.Database,
			WriteConsistency: "any",
			Time:             time.Now(),
			Precision:        cfg.Write.Precision,
			Points:           pnt,
		}

		wg.Add(1)
		counter.Increment()
		totalPoints += len(batch.Points)

		go func(b *client.BatchPoints, total int) {
			st := time.Now()
			if _, err := c.Write(*b); err != nil { // Should retry write if failed
				mu.Lock()
				if lastSuccess {
					fmt.Println("ERROR: ", err.Error())
				}
				failedRequests += 1
				totalPoints -= len(b.Points)
				lastSuccess = false
				mu.Unlock()
			} else {
				mu.Lock()
				if !lastSuccess {
					fmt.Println("success in ", time.Since(st))
				}
				lastSuccess = true
				responseTimes = append(responseTimes, NewResponseTime(int(time.Since(st).Nanoseconds())))
				mu.Unlock()
			}
			batchInterval, _ := time.ParseDuration(cfg.Write.BatchInterval)
			time.Sleep(batchInterval)
			wg.Done()
			counter.Decrement()
			if total%500000 == 0 {
				fmt.Printf("%d total points. %d in %s\n", total, cfg.Write.BatchSize, time.Since(st))
			}
		}(batch, totalPoints)

	}

	wg.Wait()

	return
}

// ConcurrencyLimiter is a go routine safe struct that can be used to
// ensure that no more than a specifid max number of goroutines are
// executing.
type ConcurrencyLimiter struct {
	inc   chan chan struct{}
	dec   chan struct{}
	max   int
	count int
}

// NewConcurrencyLimiter returns a configured limiter that will
// ensure that calls to Increment will block if the max is hit.
func NewConcurrencyLimiter(max int) *ConcurrencyLimiter {
	c := &ConcurrencyLimiter{
		inc: make(chan chan struct{}),
		dec: make(chan struct{}, max),
		max: max,
	}
	go c.handleLimits()
	return c
}

// Increment will increase the count of running goroutines by 1.
// if the number is currently at the max, the call to Increment
// will block until another goroutine decrements.
func (c *ConcurrencyLimiter) Increment() {
	r := make(chan struct{})
	c.inc <- r
	<-r
}

// Decrement will reduce the count of running goroutines by 1
func (c *ConcurrencyLimiter) Decrement() {
	c.dec <- struct{}{}
}

// handleLimits runs in a goroutine to manage the count of
// running goroutines.
func (c *ConcurrencyLimiter) handleLimits() {
	for {
		r := <-c.inc
		if c.count >= c.max {
			<-c.dec
			c.count--
		}
		c.count++
		r <- struct{}{}
	}
}
<|endoftext|>


<commit_before>package zabbix_test

import (
	. "."
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

func CreateHost(group *HostGroup, t *testing.T) *Host {
	iface := HostInterface{IP: "127.0.0.1", Port: "42", Type: Agent, UseIP: true, Main: true}
	name := fmt.Sprintf("%s-%d", getHost(), rand.Int())
	hosts := Hosts{{
		Host:       name,
		Name:       "Name for " + name,
		GroupIds:   HostGroupIds{{group.GroupId}},
		Interfaces: HostInterfaces{iface},
	}}

	err := getAPI(t).HostsCreate(hosts)
	if err != nil {
		t.Fatal(err)
	}
	return &hosts[0]
}

func DeleteHost(host *Host, t *testing.T) {
	err := getAPI(t).HostsDelete(Hosts{*host})
	if err != nil {
		t.Fatal(err)
	}
}

func TestHosts(t *testing.T) {
	api := getAPI(t)

	group := CreateHostGroup(t)
	defer DeleteHostGroup(group, t)

	hosts, err := api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 0 {
		t.Errorf("Bad hosts: %#v", hosts)
	}

	host := CreateHost(group, t)
	if host.HostId == "" || host.Host == "" {
		t.Errorf("Something is empty: %#v", host)
	}
	host.GroupIds = nil
	host.Interfaces = nil

	host2, err := api.HostGetByHost(host.Host)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(host, host2) {
		t.Errorf("Hosts are not equal:\n%#v\n%#v", host, host2)
	}

	host2, err = api.HostGetById(host.HostId)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(host, host2) {
		t.Errorf("Hosts are not equal:\n%#v\n%#v", host, host2)
	}

	hosts, err = api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 1 {
		t.Errorf("Bad hosts: %#v", hosts)
	}

	DeleteHost(host, t)

	hosts, err = api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 0 {
		t.Errorf("Bad hosts: %#v", hosts)
	}
}
<commit_msg>Add failing test for #4.<commit_after>package zabbix_test

import (
	. "."
	"fmt"
	"math/rand"
	"reflect"
	"testing"
)

func CreateHost(group *HostGroup, t *testing.T) *Host {
	name := fmt.Sprintf("%s-%d", getHost(), rand.Int())
	iface := HostInterface{DNS: name, Port: "42", Type: Agent, UseIP: false, Main: true}
	hosts := Hosts{{
		Host:       name,
		Name:       "Name for " + name,
		GroupIds:   HostGroupIds{{group.GroupId}},
		Interfaces: HostInterfaces{iface},
	}}

	err := getAPI(t).HostsCreate(hosts)
	if err != nil {
		t.Fatal(err)
	}
	return &hosts[0]
}

func DeleteHost(host *Host, t *testing.T) {
	err := getAPI(t).HostsDelete(Hosts{*host})
	if err != nil {
		t.Fatal(err)
	}
}

func TestHosts(t *testing.T) {
	api := getAPI(t)

	group := CreateHostGroup(t)
	defer DeleteHostGroup(group, t)

	hosts, err := api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 0 {
		t.Errorf("Bad hosts: %#v", hosts)
	}

	host := CreateHost(group, t)
	if host.HostId == "" || host.Host == "" {
		t.Errorf("Something is empty: %#v", host)
	}
	host.GroupIds = nil
	host.Interfaces = nil

	host2, err := api.HostGetByHost(host.Host)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(host, host2) {
		t.Errorf("Hosts are not equal:\n%#v\n%#v", host, host2)
	}

	host2, err = api.HostGetById(host.HostId)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(host, host2) {
		t.Errorf("Hosts are not equal:\n%#v\n%#v", host, host2)
	}

	hosts, err = api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 1 {
		t.Errorf("Bad hosts: %#v", hosts)
	}

	DeleteHost(host, t)

	hosts, err = api.HostsGetByHostGroups(HostGroups{*group})
	if err != nil {
		t.Fatal(err)
	}
	if len(hosts) != 0 {
		t.Errorf("Bad hosts: %#v", hosts)
	}
}
<|endoftext|>


<commit_before>package hosts

import (
	"fmt"
	"io"
	"text/template"

	mackerel "github.com/mackerelio/mackerel-client-go"

	"github.com/mackerelio/mkr/format"
	"github.com/mackerelio/mkr/mackerelclient"
)

type appLogger interface {
	Log(string, string)
	Error(error)
}

type hostApp struct {
	client    mackerelclient.Client
	logger    appLogger
	outStream io.Writer
}

type findHostsParam struct {
	verbose bool

	name     string
	service  string
	roles    []string
	statuses []string

	format string
}

func (ha *hostApp) findHosts(param findHostsParam) error {
	hosts, err := ha.client.FindHosts(&mackerel.FindHostsParam{
		Name:     param.name,
		Service:  param.service,
		Roles:    param.roles,
		Statuses: param.statuses,
	})
	if err != nil {
		return err
	}

	switch {
	case param.format != "":
		t, err := template.New("format").Parse(param.format)
		if err != nil {
			return err
		}
		return t.Execute(ha.outStream, hosts)
	case param.verbose:
		return format.PrettyPrintJSON(ha.outStream, hosts)
	default:
		var hostsFormat []*format.Host
		for _, host := range hosts {
			hostsFormat = append(hostsFormat, &format.Host{
				ID:            host.ID,
				Name:          host.Name,
				DisplayName:   host.DisplayName,
				Status:        host.Status,
				RoleFullnames: host.GetRoleFullnames(),
				IsRetired:     host.IsRetired,
				CreatedAt:     format.ISO8601Extended(host.DateFromCreatedAt()),
				IPAddresses:   host.IPAddresses(),
			})
		}
		return format.PrettyPrintJSON(ha.outStream, hostsFormat)
	}
}

type createHostParam struct {
	name             string
	roleFullnames    []string
	status           string
	customIdentifier string
}

func (ha *hostApp) createHost(param createHostParam) error {
	hostID, err := ha.client.CreateHost(&mackerel.CreateHostParam{
		Name:             param.name,
		RoleFullnames:    param.roleFullnames,
		CustomIdentifier: param.customIdentifier,
	})
	if err != nil {
		ha.error(err)
		return err
	}

	ha.log("created", hostID)

	if param.status != "" {
		err := ha.client.UpdateHostStatus(hostID, param.status)
		if err != nil {
			ha.error(err)
			return err
		}
		ha.log("updated", fmt.Sprintf("%s %s", hostID, param.status))
	}
	return nil
}

func (ha *hostApp) log(prefix, message string) {
	if ha.logger != nil {
		ha.logger.Log(prefix, message)
	}
}

func (ha *hostApp) error(err error) {
	if ha.logger != nil {
		ha.logger.Error(err)
	}
}
<commit_msg>always not nil<commit_after>package hosts

import (
	"fmt"
	"io"
	"text/template"

	mackerel "github.com/mackerelio/mackerel-client-go"

	"github.com/mackerelio/mkr/format"
	"github.com/mackerelio/mkr/mackerelclient"
)

type appLogger interface {
	Log(string, string)
	Error(error)
}

type hostApp struct {
	client    mackerelclient.Client
	logger    appLogger
	outStream io.Writer
}

type findHostsParam struct {
	verbose bool

	name     string
	service  string
	roles    []string
	statuses []string

	format string
}

func (ha *hostApp) findHosts(param findHostsParam) error {
	hosts, err := ha.client.FindHosts(&mackerel.FindHostsParam{
		Name:     param.name,
		Service:  param.service,
		Roles:    param.roles,
		Statuses: param.statuses,
	})
	if err != nil {
		return err
	}

	switch {
	case param.format != "":
		t, err := template.New("format").Parse(param.format)
		if err != nil {
			return err
		}
		return t.Execute(ha.outStream, hosts)
	case param.verbose:
		return format.PrettyPrintJSON(ha.outStream, hosts)
	default:
		var hostsFormat []*format.Host
		for _, host := range hosts {
			hostsFormat = append(hostsFormat, &format.Host{
				ID:            host.ID,
				Name:          host.Name,
				DisplayName:   host.DisplayName,
				Status:        host.Status,
				RoleFullnames: host.GetRoleFullnames(),
				IsRetired:     host.IsRetired,
				CreatedAt:     format.ISO8601Extended(host.DateFromCreatedAt()),
				IPAddresses:   host.IPAddresses(),
			})
		}
		return format.PrettyPrintJSON(ha.outStream, hostsFormat)
	}
}

type createHostParam struct {
	name             string
	roleFullnames    []string
	status           string
	customIdentifier string
}

func (ha *hostApp) createHost(param createHostParam) error {
	hostID, err := ha.client.CreateHost(&mackerel.CreateHostParam{
		Name:             param.name,
		RoleFullnames:    param.roleFullnames,
		CustomIdentifier: param.customIdentifier,
	})
	if err != nil {
		ha.error(err)
		return err
	}

	ha.log("created", hostID)

	if param.status != "" {
		err := ha.client.UpdateHostStatus(hostID, param.status)
		if err != nil {
			ha.error(err)
			return err
		}
		ha.log("updated", fmt.Sprintf("%s %s", hostID, param.status))
	}
	return nil
}

func (ha *hostApp) log(prefix, message string) {
	ha.logger.Log(prefix, message)
}

func (ha *hostApp) error(err error) {
	ha.logger.Error(err)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/jmhodges/howsmyssl/tls"
	"html/template"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	resp500Format = `HTTP/1.%d 500 Internal Server Error
Content-Length: 26
Connection: close
Content-Type: text/plain; charset="utf-8"
Date: %s

500 Internal Server Error
`
)

var (
	httpsAddr = flag.String("httpsAddr", "localhost:10443", "address to boot the HTTPS server on")
	httpAddr  = flag.String("httpAddr", "localhost:10080", "address to boot the HTTP server on")
	vhost     = flag.String("vhost", "localhost", "public domain to use in redirects and templates")
	certPath  = flag.String("cert", "./config/development.crt", "file path to the TLS certificate to serve with")
	keyPath   = flag.String("key", "./config/development.key", "file path to the TLS key to serve with")
	staticDir = flag.String("staticDir", "./static", "file path to the directory of static files to serve")
	tmplDir   = flag.String("templateDir", "./templates", "file path to the directory of templates")

	index *template.Template
)

func main() {
	flag.Parse()
	index = loadIndex()
	_, httpsPort, err := net.SplitHostPort(*httpsAddr)
	if err != nil {
		log.Fatalf("unable to parse httpsAddr: %s", err)
	}
	cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
	if err != nil {
		log.Fatalf("unable to load TLS key cert pair %s: %s", certPath, err)
	}

	tlsConf := &tls.Config{
		Certificates: []tls.Certificate{cert},
		NextProtos:   []string{"https"},
	}
	tlsConf.BuildNameToCertificate()

	tlsListener, err := tls.Listen("tcp", *httpsAddr, tlsConf)
	if err != nil {
		log.Fatalf("unable to listen for the HTTPS server on %s: %s", *httpsAddr, err)
	}
	plaintextListener, err := net.Listen("tcp", *httpAddr)
	if err != nil {
		log.Fatalf("unable to listen for the HTTP server on %s: %s", *httpAddr, err)
	}
	l := &listener{tlsListener}

	m := tlsMux(
		*vhost,
		httpsPort,
		http.HandlerFunc(handleWeb),
		http.HandlerFunc(handleAPI),
		http.StripPrefix("/s/", http.FileServer(http.Dir(*staticDir))))

	log.Printf("Booting HTTPS on %s and HTTP on %s", *httpsAddr, *httpAddr)
	go func() {
		err := http.Serve(l, m)
		if err != nil {
			log.Fatalf("https server error: %s", err)
		}
	}()
	err = http.Serve(plaintextListener, plaintextRedirect(*vhost, httpsPort))
	if err != nil {
		log.Fatalf("http server error: %s", err)
	}
}

func tlsMux(vhost, port string, webHandler, apiHandler, staticHandler http.Handler) *http.ServeMux {
	m := http.NewServeMux()
	prefix := vhost
	if port != "443" {
		prefix = net.JoinHostPort(vhost, port)
	}
	m.Handle(prefix+"/s/", staticHandler)
	m.Handle(prefix+"/a/check", apiHandler)
	m.Handle(prefix+"/", webHandler)
	m.Handle("/", tlsRedirect(vhost, port, webHandler))
	return m
}

func renderHTML(data *clientInfo) ([]byte, error) {
	b := new(bytes.Buffer)
	err := index.Execute(b, data)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func renderJSON(data *clientInfo) ([]byte, error) {
	return json.Marshal(data)
}

func handleWeb(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "404 Not Found", http.StatusNotFound)
		return
	}

	hijackHandle(w, r, "text/html;charset=utf-8", renderHTML)
}

func handleAPI(w http.ResponseWriter, r *http.Request) {
	hijackHandle(w, r, "application/json", renderJSON)
}

func hijackHandle(w http.ResponseWriter, r *http.Request, contentType string, render func(*clientInfo) ([]byte, error)) {
	hj, ok := w.(http.Hijacker)
	if !ok {
		log.Printf("server not hijackable\n")
		http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
		return
	}

	c, brw, err := hj.Hijack()
	if err != nil {
		log.Printf("server errored during hijack: %s\n", err)
		return
	}
	defer c.Close()
	tc, ok := c.(*conn)
	if !ok {
		log.Printf("Unable to convert net.Conn to *conn: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
	}
	data := ClientInfo(tc)
	bs, err := render(data)
	if err != nil {
		log.Printf("Unable to excute index template: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
		return
	}
	contentLength := int64(len(bs))
	h := make(http.Header)
	h.Set("Date", time.Now().Format(http.TimeFormat))
	h.Set("Content-Type", contentType)
	if r.ProtoMinor == 1 { // Assumes HTTP/1.x
		h.Set("Connection", "close")
	}
	h.Set("Content-Length", strconv.FormatInt(contentLength, 10))
	resp := &http.Response{
		StatusCode:    200,
		ContentLength: contentLength,
		Header:        h,
		Body:          ioutil.NopCloser(bytes.NewBuffer(bs)),
		ProtoMajor:    1, // Assumes HTTP/1.x
		ProtoMinor:    r.ProtoMinor,
	}
	bs, err = httputil.DumpResponse(resp, true)
	if err != nil {
		log.Printf("unable to write response: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
		return
	}
	brw.Write(bs)
	brw.Flush()
}

func hijacked500(brw *bufio.ReadWriter, protoMinor int) {
	// Assumes HTTP/1.x
	s := fmt.Sprintf(resp500Format, protoMinor, time.Now().Format(http.TimeFormat))
	brw.WriteString(s)
	brw.Flush()
}

func commonRedirect(w http.ResponseWriter, r *http.Request, vhostWithPort string) {
	if r.URL.Path == "/healthcheck" {
		w.WriteHeader(200)
		return
	}
	var u url.URL
	u = *r.URL
	u.Scheme = "https"
	u.Host = vhostWithPort
	http.Redirect(w, r, u.String(), http.StatusMovedPermanently)
}

func plaintextRedirect(vhost, port string) http.Handler {
	if port != "443" {
		vhost = net.JoinHostPort(vhost, port)
	}
	var h http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		commonRedirect(w, r, vhost)
	}
	return h
}

func tlsRedirect(vhost, port string, webHandler http.Handler) http.Handler {
	vhostWithPort := vhost
	if port != "443" {
		vhostWithPort = net.JoinHostPort(vhost, port)
	}
	var h http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.Host)
		if err != nil {
			host = r.Host
		}
		if vhost == host {
			webHandler.ServeHTTP(w, r)
			return
		}
		commonRedirect(w, r, vhostWithPort)
	}
	return h
}

func loadIndex() *template.Template {
	return template.Must(template.New("index.html").
		Funcs(template.FuncMap{"sentence": sentence, "ratingSpan": ratingSpan}).
		ParseFiles(*tmplDir + "/index.html"))

}

func ratingSpan(rating Rating) template.HTML {
	class := ""
	switch rating {
	case okay:
		class = "okay"
	case improvable:
		class = "improvable"
	case bad:
		class = "bad"
	}
	return template.HTML(fmt.Sprintf(`<span class="%s">%s</span>`, class, rating))
}

func sentence(parts []string) string {
	if len(parts) == 1 {
		return parts[0] + "."
	}
	commaed := parts[:len(parts)-1]
	return strings.Join(commaed, ", ") + ", and " + parts[len(parts)-1] + "."
}
<commit_msg>log requests simply (avoiding Hijack problems)<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/jmhodges/howsmyssl/tls"
	"html/template"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	resp500Format = `HTTP/1.%d 500 Internal Server Error
Content-Length: 26
Connection: close
Content-Type: text/plain; charset="utf-8"
Date: %s

500 Internal Server Error
`
)

var (
	httpsAddr = flag.String("httpsAddr", "localhost:10443", "address to boot the HTTPS server on")
	httpAddr  = flag.String("httpAddr", "localhost:10080", "address to boot the HTTP server on")
	vhost     = flag.String("vhost", "localhost", "public domain to use in redirects and templates")
	certPath  = flag.String("cert", "./config/development.crt", "file path to the TLS certificate to serve with")
	keyPath   = flag.String("key", "./config/development.key", "file path to the TLS key to serve with")
	staticDir = flag.String("staticDir", "./static", "file path to the directory of static files to serve")
	tmplDir   = flag.String("templateDir", "./templates", "file path to the directory of templates")

	index *template.Template
)

func main() {
	flag.Parse()
	index = loadIndex()
	_, httpsPort, err := net.SplitHostPort(*httpsAddr)
	if err != nil {
		log.Fatalf("unable to parse httpsAddr: %s", err)
	}
	cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
	if err != nil {
		log.Fatalf("unable to load TLS key cert pair %s: %s", certPath, err)
	}

	tlsConf := &tls.Config{
		Certificates: []tls.Certificate{cert},
		NextProtos:   []string{"https"},
	}
	tlsConf.BuildNameToCertificate()

	tlsListener, err := tls.Listen("tcp", *httpsAddr, tlsConf)
	if err != nil {
		log.Fatalf("unable to listen for the HTTPS server on %s: %s", *httpsAddr, err)
	}
	plaintextListener, err := net.Listen("tcp", *httpAddr)
	if err != nil {
		log.Fatalf("unable to listen for the HTTP server on %s: %s", *httpAddr, err)
	}
	l := &listener{tlsListener}

	m := tlsMux(
		*vhost,
		httpsPort,
		http.HandlerFunc(handleWeb),
		http.HandlerFunc(handleAPI),
		http.StripPrefix("/s/", http.FileServer(http.Dir(*staticDir))))

	log.Printf("Booting HTTPS on %s and HTTP on %s", *httpsAddr, *httpAddr)
	go func() {
		err := http.Serve(l, m)
		if err != nil {
			log.Fatalf("https server error: %s", err)
		}
	}()
	err = http.Serve(plaintextListener, plaintextRedirect(*vhost, httpsPort))
	if err != nil {
		log.Fatalf("http server error: %s", err)
	}
}

func tlsMux(vhost, port string, webHandler, apiHandler, staticHandler http.Handler) *http.ServeMux {
	m := http.NewServeMux()
	prefix := vhost
	if port != "443" {
		prefix = net.JoinHostPort(vhost, port)
	}
	m.Handle(prefix+"/s/", logHandler{staticHandler})
	m.Handle(prefix+"/a/check", logHandler{apiHandler})
	m.Handle(prefix+"/", logHandler{webHandler})
	m.Handle("/", logHandler{tlsRedirect(vhost, port, webHandler)})
	return m
}

func renderHTML(data *clientInfo) ([]byte, error) {
	b := new(bytes.Buffer)
	err := index.Execute(b, data)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func renderJSON(data *clientInfo) ([]byte, error) {
	return json.Marshal(data)
}

func handleWeb(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "404 Not Found", http.StatusNotFound)
		return
	}

	hijackHandle(w, r, "text/html;charset=utf-8", renderHTML)
}

func handleAPI(w http.ResponseWriter, r *http.Request) {
	hijackHandle(w, r, "application/json", renderJSON)
}

func hijackHandle(w http.ResponseWriter, r *http.Request, contentType string, render func(*clientInfo) ([]byte, error)) {
	hj, ok := w.(http.Hijacker)
	if !ok {
		log.Printf("server not hijackable\n")
		http.Error(w, "500 Internal Server Error", http.StatusInternalServerError)
		return
	}

	c, brw, err := hj.Hijack()
	if err != nil {
		log.Printf("server errored during hijack: %s\n", err)
		return
	}
	defer c.Close()
	tc, ok := c.(*conn)
	if !ok {
		log.Printf("Unable to convert net.Conn to *conn: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
	}
	data := ClientInfo(tc)
	bs, err := render(data)
	if err != nil {
		log.Printf("Unable to excute index template: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
		return
	}
	contentLength := int64(len(bs))
	h := make(http.Header)
	h.Set("Date", time.Now().Format(http.TimeFormat))
	h.Set("Content-Type", contentType)
	if r.ProtoMinor == 1 { // Assumes HTTP/1.x
		h.Set("Connection", "close")
	}
	h.Set("Content-Length", strconv.FormatInt(contentLength, 10))
	resp := &http.Response{
		StatusCode:    200,
		ContentLength: contentLength,
		Header:        h,
		Body:          ioutil.NopCloser(bytes.NewBuffer(bs)),
		ProtoMajor:    1, // Assumes HTTP/1.x
		ProtoMinor:    r.ProtoMinor,
	}
	bs, err = httputil.DumpResponse(resp, true)
	if err != nil {
		log.Printf("unable to write response: %s\n", err)
		hijacked500(brw, r.ProtoMinor)
		return
	}
	brw.Write(bs)
	brw.Flush()
}

func hijacked500(brw *bufio.ReadWriter, protoMinor int) {
	// Assumes HTTP/1.x
	s := fmt.Sprintf(resp500Format, protoMinor, time.Now().Format(http.TimeFormat))
	brw.WriteString(s)
	brw.Flush()
}

func commonRedirect(w http.ResponseWriter, r *http.Request, vhostWithPort string) {
	if r.URL.Path == "/healthcheck" {
		w.WriteHeader(200)
		return
	}
	var u url.URL
	u = *r.URL
	u.Scheme = "https"
	u.Host = vhostWithPort
	http.Redirect(w, r, u.String(), http.StatusMovedPermanently)
}

func plaintextRedirect(vhost, port string) http.Handler {
	if port != "443" {
		vhost = net.JoinHostPort(vhost, port)
	}
	var h http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		commonRedirect(w, r, vhost)
	}
	return h
}

func tlsRedirect(vhost, port string, webHandler http.Handler) http.Handler {
	vhostWithPort := vhost
	if port != "443" {
		vhostWithPort = net.JoinHostPort(vhost, port)
	}
	var h http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.Host)
		if err != nil {
			host = r.Host
		}
		if vhost == host {
			webHandler.ServeHTTP(w, r)
			return
		}
		commonRedirect(w, r, vhostWithPort)
	}
	return h
}

func loadIndex() *template.Template {
	return template.Must(template.New("index.html").
		Funcs(template.FuncMap{"sentence": sentence, "ratingSpan": ratingSpan}).
		ParseFiles(*tmplDir + "/index.html"))

}

func ratingSpan(rating Rating) template.HTML {
	class := ""
	switch rating {
	case okay:
		class = "okay"
	case improvable:
		class = "improvable"
	case bad:
		class = "bad"
	}
	return template.HTML(fmt.Sprintf(`<span class="%s">%s</span>`, class, rating))
}

func sentence(parts []string) string {
	if len(parts) == 1 {
		return parts[0] + "."
	}
	commaed := parts[:len(parts)-1]
	return strings.Join(commaed, ", ") + ", and " + parts[len(parts)-1] + "."
}

type logHandler struct {
	inner http.Handler
}

// Since we have a Hijack in our code, this simple writer will suffice for
// now.
func (h logHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = "0.0.0.0"
	}
	fmt.Printf("%s %s\n", host, r.URL)
	h.inner.ServeHTTP(w, r)
}
<|endoftext|>


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

// cookie.go provides checks for cookies.

package ht

import (
	"fmt"
	"net/http"
	"strings"
	"time"
)

func init() {
	RegisterCheck(&SetCookie{})
}

// SetCookie checks for cookies being properly set.
// Note that the Path and Domain conditions are checked on the received Path
// and/or Domain and not on the interpreted values according to RFC 6265.
type SetCookie struct {
	Name   string    `json:",omitempty"` // Name is the cookie name.
	Value  Condition `json:",omitempty"` // Value is applied to the cookie value
	Path   Condition `json:",omitempty"` // Path is applied to the path value
	Domain Condition `json:",omitempty"` // Domain is applied to the domain value

	// MinLifetime is the expectetd minimum lifetime of the cookie.
	// A positive value enforces a persistent cookie and a negative
	// value is is equivalent to Delete=true.
	MinLifetime Duration `json:",omitempty"`

	// Absent indicates that the cookie with the given Name must not be received.
	Absent bool `json:",omitempty"`

	// Delete indicates that the Set-Cookie header enforces a deletion
	// of the cookie.
	Delete bool `json:",omitempty"`

	// Type is the type of the cookie. It is a space seperated string of
	// the following (case-insensitive) keywords:
	//   - "session": a session cookie
	//   - "persistent": a persistent cookie
	//   - "secure": a secure cookie, to be sont over https only
	//   - "unsafe", to be sent also over http
	//   - "httpOnly": not accesible from JavaScript
	//   - "exposed": accesible from JavaScript, Flash, etc.
	Type string `json:",omitempty"`
}

// Execute implements Check's Execute method.
func (c SetCookie) Execute(t *Test) error {
	var cookie *http.Cookie
	for _, cp := range t.Response.Response.Cookies() {
		if cp.Name == c.Name {
			cookie = cp
			break
		}
	}

	if cookie == nil {
		if c.Absent {
			return nil
		}
		return fmt.Errorf("Missing cookie %s", c.Name)
	}
	if c.Absent {
		return fmt.Errorf("Found cookie %s=%s", c.Name, cookie.Value)
	}

	if err := c.Value.Fullfilled(cookie.Value); err != nil {
		return err
	}
	if err := c.Path.Fullfilled(cookie.Path); err != nil {
		return err
	}
	if err := c.Value.Fullfilled(cookie.Domain); err != nil {
		return err
	}

	if c.Delete {
		err := properlyDeleted(cookie)
		if err != nil {
			return err
		}
		// No need to check MinLifetime and Type is ignored
		// for delete requests.
		return nil
	}

	if c.MinLifetime > 0 {
		if cookie.MaxAge > 0 {
			if int(time.Duration(c.MinLifetime).Seconds()) > cookie.MaxAge {
				return fmt.Errorf("MaxAge %ds of cookie %s too short, want > %s",
					cookie.MaxAge, c.Name, c.MinLifetime)
			}
		} else if !cookie.Expires.IsZero() {
			min := time.Now().Add(time.Duration(c.MinLifetime))
			if min.Before(cookie.Expires) {
				return fmt.Errorf("Expires %ss of cookie %s too early, want > %s",
					cookie.Expires.Format(time.RFC1123), c.Name,
					min.Format(time.RFC1123))
			}
		} else {
			return fmt.Errorf("Cookie %s is session cookie", c.Name)
		}
	}

	return c.checkType(cookie)

}

// Prepare implements Check's Prepare method.
func (c *SetCookie) Prepare() error {
	if err := c.Value.Compile(); err != nil {
		return err
	}
	if err := c.Path.Compile(); err != nil {
		return err
	}
	if err := c.Domain.Compile(); err != nil {
		return err
	}

	if c.MinLifetime < 0 {
		c.MinLifetime = 0
		c.Delete = true
	}

	c.Type = strings.ToLower(c.Type)
	x := c.Type
	for _, t := range strings.Split("session persistent secure unsafe httponly exposed", " ") {
		x = strings.TrimSpace(strings.Replace(x, t, "", -1))
	}
	if x != "" {
		return fmt.Errorf("ht: unknown stuff in cookie type %q", x)
	}

	if c.MinLifetime > 0 {
		if strings.Index(c.Type, "persistent") == -1 {
			c.Type += " persistent"
		}
	}

	return nil
}

// properlyDeleted ensures the received cookie c is a proper
// delete request
func properlyDeleted(cookie *http.Cookie) error {
	if cookie.MaxAge == 0 && cookie.Expires.IsZero() {
		return fmt.Errorf("cookie not deleted", cookie.MaxAge)
	}

	if cookie.MaxAge > 0 {
		return fmt.Errorf("MaxAge=%d > 0, not deleted", cookie.MaxAge)
	}
	if !cookie.Expires.IsZero() {
		// Deletion via Expires value is problematic due to
		// clock variations. Play it safe and require at least
		// 90 seconds backdatedExpires.
		now := time.Now().Add(-90 * time.Second)
		if cookie.Expires.After(now) {
			return fmt.Errorf("Expires %s to late, not deleted", cookie.Expires)
		}
	}

	// Not both of MaxAge and Expires are unset; and none is positive,
	// so at least one is negative. Thus properly deleted.
	return nil
}

func (c *SetCookie) checkType(cookie *http.Cookie) error {
	t := c.Type
	if strings.Index(t, "session") != -1 {
		if cookie.MaxAge > 0 || !cookie.Expires.IsZero() {
			return fmt.Errorf("persistent cookie")
		}
	} else if strings.Index(t, "persistent") != -1 {
		if cookie.MaxAge == 0 && cookie.Expires.IsZero() {
			return fmt.Errorf("session cookie")
		}
	}

	if strings.Index(t, "secure") != -1 && !cookie.Secure {
		return fmt.Errorf("not a secure cookie")
	} else if strings.Index(t, "unsafe") != -1 && cookie.Secure {
		return fmt.Errorf("secure cookie")
	}

	if strings.Index(t, "httponly") != -1 && !cookie.HttpOnly {
		return fmt.Errorf("not a http-only cookie")
	} else if strings.Index(t, "exposed") != -1 && cookie.HttpOnly {
		return fmt.Errorf("http-only cookie")
	}

	return nil
}
<commit_msg>ht: add explicit DeleteCookie check<commit_after>// Copyright 2015 Volker Dobler.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// cookie.go provides checks for cookies.

package ht

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

func init() {
	RegisterCheck(&SetCookie{})
	RegisterCheck(&DeleteCookie{})
}

func findCookiesByName(t *Test, name string) (cookies []*http.Cookie) {
	for _, cp := range t.Response.Response.Cookies() {
		if cp.Name == name {
			cookies = append(cookies, cp)
		}
	}
	return cookies
}

// isProperCookiePath determins whether path is a sensible Path value
// in a Set-Cookie header received from u. E.g. a cookie of the form
//     Set-Cookie: foo=bar; Path=/abc/123
// which is received while accessing
//     http://example.org/some/other/path
// will be rejected by browser because the ath value is not suitable for
// the path in the URL. See RFC 6265 section 5.1.4 for reference.
func isProperCookiePath(u *url.URL, path string) bool {
	if path == "" || u.Path == path {
		return true // Empyt path defaults to actual path, so it is okay.
	}
	if strings.HasPrefix(u.Path, path) {
		if u.Path[len(path)] == '/' || path[len(path)-1] == '/' {
			// "/abc" and "/abc/" both are proper for "/abc/foo" URL path.
			return true
		}
	}
	return false
}

// ----------------------------------------------------------------------------
// SetCookie

// SetCookie checks for cookies being properly set.
// Note that the Path and Domain conditions are checked on the received Path
// and/or Domain and not on the interpreted values according to RFC 6265.
type SetCookie struct {
	Name   string    `json:",omitempty"` // Name is the cookie name.
	Value  Condition `json:",omitempty"` // Value is applied to the cookie value
	Path   Condition `json:",omitempty"` // Path is applied to the path value
	Domain Condition `json:",omitempty"` // Domain is applied to the domain value

	// MinLifetime is the expectetd minimum lifetime of the cookie.
	// A positive value enforces a persistent cookie.
	// Negative values are illegal (use DelteCookie instead).
	MinLifetime Duration `json:",omitempty"`

	// Absent indicates that the cookie with the given Name must not be received.
	Absent bool `json:",omitempty"`

	// Type is the type of the cookie. It is a space seperated string of
	// the following (case-insensitive) keywords:
	//   - "session": a session cookie
	//   - "persistent": a persistent cookie
	//   - "secure": a secure cookie, to be sont over https only
	//   - "unsafe", aka insecure; to be sent also over http
	//   - "httpOnly": not accesible from JavaScript
	//   - "exposed": accesible from JavaScript, Flash, etc.
	Type string `json:",omitempty"`
}

// Execute implements Check's Execute method.
func (c SetCookie) Execute(t *Test) error {
	var cookie *http.Cookie
	for _, cp := range t.Response.Response.Cookies() {
		if cp.Name == c.Name {
			cookie = cp
			break
		}
	}

	if cookie == nil {
		if c.Absent {
			return nil
		}
		return fmt.Errorf("Missing cookie %s", c.Name)
	}
	if c.Absent {
		return fmt.Errorf("Found cookie %s=%s", c.Name, cookie.Value)
	}

	if err := c.Value.Fullfilled(cookie.Value); err != nil {
		return err
	}
	if err := c.Path.Fullfilled(cookie.Path); err != nil {
		return err
	}
	if err := c.Value.Fullfilled(cookie.Domain); err != nil {
		return err
	}

	if c.MinLifetime > 0 {
		if cookie.MaxAge > 0 {
			if int(time.Duration(c.MinLifetime).Seconds()) > cookie.MaxAge {
				return fmt.Errorf("MaxAge %ds of cookie %s too short, want > %s",
					cookie.MaxAge, c.Name, c.MinLifetime)
			}
		} else if !cookie.Expires.IsZero() {
			min := time.Now().Add(time.Duration(c.MinLifetime))
			if min.Before(cookie.Expires) {
				return fmt.Errorf("Expires %ss of cookie %s too early, want > %s",
					cookie.Expires.Format(time.RFC1123), c.Name,
					min.Format(time.RFC1123))
			}
		} else {
			return fmt.Errorf("Cookie %s is session cookie", c.Name)
		}
	}

	return c.checkType(cookie)

}
func (c *SetCookie) checkType(cookie *http.Cookie) error {
	t := c.Type
	if strings.Index(t, "session") != -1 {
		if cookie.MaxAge > 0 || !cookie.Expires.IsZero() {
			return fmt.Errorf("persistent cookie")
		}
	} else if strings.Index(t, "persistent") != -1 {
		if cookie.MaxAge == 0 && cookie.Expires.IsZero() {
			return fmt.Errorf("session cookie")
		}
	}

	if strings.Index(t, "secure") != -1 && !cookie.Secure {
		return fmt.Errorf("not a secure cookie")
	} else if strings.Index(t, "unsafe") != -1 && cookie.Secure {
		return fmt.Errorf("secure cookie")
	}

	if strings.Index(t, "httponly") != -1 && !cookie.HttpOnly {
		return fmt.Errorf("not a http-only cookie")
	} else if strings.Index(t, "exposed") != -1 && cookie.HttpOnly {
		return fmt.Errorf("http-only cookie")
	}

	return nil
}

// Prepare implements Check's Prepare method.
func (c *SetCookie) Prepare() error {
	if err := c.Value.Compile(); err != nil {
		return err
	}
	if err := c.Path.Compile(); err != nil {
		return err
	}
	if err := c.Domain.Compile(); err != nil {
		return err
	}

	if c.MinLifetime < 0 {
		return fmt.Errorf("ht: illegal negative MinLifetime")
	}

	c.Type = strings.ToLower(c.Type)
	x := strings.Replace(c.Type, ",", " ", -1)
	for _, t := range strings.Split("session persistent secure unsafe httponly exposed", " ") {
		x = strings.TrimSpace(strings.Replace(x, t, "", -1))
	}
	if x != "" {
		return fmt.Errorf("ht: unknown stuff in cookie type %q", x)
	}

	if c.MinLifetime > 0 {
		if strings.Index(c.Type, "persistent") == -1 {
			c.Type += " persistent"
		}
	}

	return nil
}

// ----------------------------------------------------------------------------
// DeleteCookie

// DeleteCookie checks that the HTTP response properly deletes all cookies
// matching Name, Path and Domain. Path and Domain are optional in which case
// all cookies with the given Name are checkd for deletion.
type DeleteCookie struct {
	Name   string
	Path   string `json:",omitempty"`
	Domain string `json:",omitempty"`
}

// Execute implements Check's Execute method.
func (c DeleteCookie) Execute(t *Test) error {
	errors := []string{}
	deleted := false
	u := t.Response.Response.Request.URL
	for _, cookie := range findCookiesByName(t, c.Name) {
		// Report malformed (becuase of path) cookies
		if !isProperCookiePath(u, cookie.Path) {
			em := fmt.Sprintf("invalid cookie path %s for URL %s",
				c.Name, cookie.Path, u.String())
			errors = append(errors, em)
			continue
		}
		if c.Path != "" && !isProperCookiePath(u, c.Path) {
			// Well this should be reported during Prepare, not
			// here. Unfortunately Prepare knows nothing about
			// the URL the request will go to. So do it here.
			em := fmt.Sprintf("bogus path %s in check", c.Path)
			errors = append(errors, em)
			continue
		}

		// Work only on cookies matching the optional Path and Domain.
		if c.Path != "" && cookie.Path != c.Path {
			continue
		}
		if c.Domain != "" && cookie.Domain != c.Domain {
			continue
		}

		// Check for deletion.
		if cookie.MaxAge == 0 && cookie.Expires.IsZero() {
			em := fmt.Sprintf("cookie %s;%s;%s not deleted",
				c.Name, cookie.Path, cookie.Domain)
			errors = append(errors, em)
			continue
		}
		if cookie.MaxAge > 0 {
			em := fmt.Sprintf("cookie (%s;%s;%s) not deleted,MaxAge=%d",
				c.Name, cookie.Path, cookie.Domain, cookie.MaxAge)
			errors = append(errors, em)
			continue
		}
		if !cookie.Expires.IsZero() {
			// Deletion via Expires value is problematic due to
			// clock variations. Play it safe and require at least
			// 90 seconds backdatedExpires.
			now := time.Now().Add(-90 * time.Second)
			if cookie.Expires.After(now) {
				em := fmt.Sprintf("cookie (%s;%s;%s) not deleted, Expires=%s",
					c.Name, cookie.Path, cookie.Domain, cookie.RawExpires)
				errors = append(errors, em)
				continue
			}
		}
		// Here not both of MaxAge and Expires are unset
		// and none is positive, so at least one is negative.
		// Thus the cookie is properly deleted.

		// Sanity check: Do not send values when deleting cookies.
		// It technically works but it is _strange_.
		if cookie.Value != "" {
			em := fmt.Sprintf("cookie (%s;%s;%s) deleted but has value %s",
				c.Name, cookie.Path, cookie.Domain, cookie.Value)
			errors = append(errors, em)
			continue
		}

		deleted = true
	}

	if !deleted {
		errors = append(errors, fmt.Sprintf("No cookie (%s;%s;%s) was deleted",
			c.Name, c.Path, c.Domain))
	}

	if len(errors) != 0 {
		return fmt.Errorf("%s", strings.Join(errors, "\n"))
	}
	return nil
}

// Prepare implements Check's Prepare method.
func (c *DeleteCookie) Prepare() error {
	return nil
}
<|endoftext|>


<commit_before>package leptonica

/*
#cgo LDFLAGS: -llept
#include "leptonica/allheaders.h"
#include <stdlib.h>

l_uint8* uglycast(void* value) { return (l_uint8*)value; }

*/
import "C"
import (
	"errors"
	"strconv"
	"sync"
	"unsafe"
)

type ImageType int32

const (
	UNKNOWN ImageType = iota
	BMP
	JFIF_JPEG
	PNG
	TIFF
	TIFF_PACKBITS
	TIFF_RLE
	TIFF_G3
	TIFF_G4
	TIFF_LZW
	TIFF_ZIP
	PNM
	PS
	GIF
	JP2
	WEBP
	LPDF
	DEFAULT
	SPIX
)

type Pix struct {
	cPix   *C.PIX // exported C.PIX so it can be used with other cgo wrap packages
	closed bool
	lock   sync.Mutex
}

func (p *Pix) CPIX() *C.PIX {
	return p.cPix
}

// GetDimensions returns the dimensions in Width, Height, Depth, Error format
func (p *Pix) GetDimensions() (int32, int32, int32, error) {
	p.lock.Lock()
	defer p.lock.Unlock()
	var w, h, d int32
	cW := C.l_int32(w)
	cH := C.l_int32(h)
	cD := C.l_int32(d)
	if !p.closed {
		code := C.pixGetDimensions(p.cPix, &cW, &cH, &cD)
		if code != 0 {
			return 0, 0, 0, errors.New("could not get dimensions")
		}

	}
	return int32(cW), int32(cH), int32(cD), nil
}

func (p *Pix) Close() {
	p.lock.Lock()
	defer p.lock.Unlock()
	if !p.closed {
		// LEPT_DLL extern void pixDestroy ( PIX **ppix );
		C.pixDestroy(&p.cPix)
		C.free(unsafe.Pointer(p.cPix))
		p.closed = true
	}
}

// WriteFile saves to disk the current pix in the given format
func (p *Pix) WriteFile(filename string, format ImageType) error {
	cFilename := C.CString(filename)
	defer C.free(unsafe.Pointer(cFilename))
	cFormat := C.l_int32(format)

	// create new PIX
	code := C.pixWrite(cFilename, p.cPix, cFormat)
	if code != 0 {
		return errors.New("could not write PIX to given filename: " + filename + " (format: " + strconv.Itoa(int(format)) + ")")
	}

	return nil
}

// WriteBytes will return a pointer to a byte array holding the data from PIX in the given format
func (p *Pix) WriteBytes(format ImageType) (*[]byte, error) {
	var memory []byte
	memPtr := C.uglycast(unsafe.Pointer(&(memory)))
	var i int64
	// TODO: Do I need to free this?
	sizePtr := C.size_t(i)
	cFormat := C.l_int32(format)
	code := C.pixWriteMem(&memPtr, &sizePtr, p.cPix, cFormat)
	if code != 0 {
		return nil, errors.New("Cannot write type to given memory space")
	}
	data := C.GoBytes(unsafe.Pointer(memPtr), C.int(sizePtr))
	C.free(unsafe.Pointer(memPtr))
	return &data, nil
}

// LEPT_DLL extern PIX * pixRead ( const char *filename );
// NewPixFromFile creates a new Pix from given filename
func NewPixFromFile(filename string) (*Pix, error) {
	cFilename := C.CString(filename)
	defer C.free(unsafe.Pointer(cFilename))

	// create new PIX
	CPIX := C.pixRead(cFilename)
	if CPIX == nil {
		return nil, errors.New("could not create PIX from given filename")
	}

	// all done
	pix := &Pix{
		cPix: CPIX,
	}
	return pix, nil
}

// NewPixReadMem creates a new Pix instance from a byte array
func NewPixReadMem(image *[]byte) (*Pix, error) {
	//ptr := (*C.l_uint8)(*C.uchar)(unsafe.Pointer(&(*image)[0]))
	ptr := C.uglycast(unsafe.Pointer(&(*image)[0]))
	CPIX := C.pixReadMem(ptr, C.size_t(len(*image)))
	if CPIX == nil {
		return nil, errors.New("Cannot create PIX from given image data")
	}
	pix := &Pix{
		cPix: CPIX,
	}
	return pix, nil
}
<commit_msg>Now return []byte not *[]byte<commit_after>package leptonica

/*
#cgo LDFLAGS: -llept
#include "leptonica/allheaders.h"
#include <stdlib.h>

l_uint8* uglycast(void* value) { return (l_uint8*)value; }

*/
import "C"
import (
	"errors"
	"strconv"
	"sync"
	"unsafe"
)

type ImageType int32

const (
	UNKNOWN ImageType = iota
	BMP
	JFIF_JPEG
	PNG
	TIFF
	TIFF_PACKBITS
	TIFF_RLE
	TIFF_G3
	TIFF_G4
	TIFF_LZW
	TIFF_ZIP
	PNM
	PS
	GIF
	JP2
	WEBP
	LPDF
	DEFAULT
	SPIX
)

type Pix struct {
	cPix   *C.PIX // exported C.PIX so it can be used with other cgo wrap packages
	closed bool
	lock   sync.Mutex
}

func (p *Pix) CPIX() *C.PIX {
	return p.cPix
}

// GetDimensions returns the dimensions in Width, Height, Depth, Error format
func (p *Pix) GetDimensions() (int32, int32, int32, error) {
	p.lock.Lock()
	defer p.lock.Unlock()
	var w, h, d int32
	cW := C.l_int32(w)
	cH := C.l_int32(h)
	cD := C.l_int32(d)
	if !p.closed {
		code := C.pixGetDimensions(p.cPix, &cW, &cH, &cD)
		if code != 0 {
			return 0, 0, 0, errors.New("could not get dimensions")
		}

	}
	return int32(cW), int32(cH), int32(cD), nil
}

func (p *Pix) Close() {
	p.lock.Lock()
	defer p.lock.Unlock()
	if !p.closed {
		// LEPT_DLL extern void pixDestroy ( PIX **ppix );
		C.pixDestroy(&p.cPix)
		C.free(unsafe.Pointer(p.cPix))
		p.closed = true
	}
}

// WriteFile saves to disk the current pix in the given format
func (p *Pix) WriteFile(filename string, format ImageType) error {
	cFilename := C.CString(filename)
	defer C.free(unsafe.Pointer(cFilename))
	cFormat := C.l_int32(format)

	// create new PIX
	code := C.pixWrite(cFilename, p.cPix, cFormat)
	if code != 0 {
		return errors.New("could not write PIX to given filename: " + filename + " (format: " + strconv.Itoa(int(format)) + ")")
	}

	return nil
}

// WriteBytes will return a pointer to a byte array holding the data from PIX in the given format
func (p *Pix) WriteBytes(format ImageType) ([]byte, error) {
	var memory []byte
	memPtr := C.uglycast(unsafe.Pointer(&(memory)))
	var i int64
	// TODO: Do I need to free this?
	sizePtr := C.size_t(i)
	cFormat := C.l_int32(format)
	code := C.pixWriteMem(&memPtr, &sizePtr, p.cPix, cFormat)
	if code != 0 {
		return nil, errors.New("Cannot write type to given memory space")
	}
	data := C.GoBytes(unsafe.Pointer(memPtr), C.int(sizePtr))
	C.free(unsafe.Pointer(memPtr))
	return data, nil
}

// LEPT_DLL extern PIX * pixRead ( const char *filename );
// NewPixFromFile creates a new Pix from given filename
func NewPixFromFile(filename string) (*Pix, error) {
	cFilename := C.CString(filename)
	defer C.free(unsafe.Pointer(cFilename))

	// create new PIX
	CPIX := C.pixRead(cFilename)
	if CPIX == nil {
		return nil, errors.New("could not create PIX from given filename")
	}

	// all done
	pix := &Pix{
		cPix: CPIX,
	}
	return pix, nil
}

// NewPixReadMem creates a new Pix instance from a byte array
func NewPixReadMem(image *[]byte) (*Pix, error) {
	//ptr := (*C.l_uint8)(*C.uchar)(unsafe.Pointer(&(*image)[0]))
	ptr := C.uglycast(unsafe.Pointer(&(*image)[0]))
	CPIX := C.pixReadMem(ptr, C.size_t(len(*image)))
	if CPIX == nil {
		return nil, errors.New("Cannot create PIX from given image data")
	}
	pix := &Pix{
		cPix: CPIX,
	}
	return pix, nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2016-2017 vLife Systems Ltd <http://vlifesystems.com>
// Licensed under an MIT licence.  Please see LICENSE.md for details.

package html

import (
	"bytes"
	"fmt"
	"github.com/vlifesystems/rulehunter/config"
	"github.com/vlifesystems/rulehunter/html/cmd"
	"github.com/vlifesystems/rulehunter/logger"
	"github.com/vlifesystems/rulehunter/progress"
	"github.com/vlifesystems/rulehunter/quitter"
	"html/template"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

// File mode permission used as standard for the html content:
// No special permission bits
// User: Read, Write Execute
// Group: Read
// Other: None
const modePerm = 0740

// This should be run as a goroutine and each time a command is passed to
// cmds the html will be generated
func Run(
	config *config.Config,
	pm *progress.Monitor,
	l logger.Logger,
	quit *quitter.Quitter,
	cmds <-chan cmd.Cmd,
) {
	quit.Add()
	defer quit.Done()
	if err := generate(cmd.All, config, pm); err != nil {
		l.Error(fmt.Errorf("Couldn't generate report: %s", err))
	}

	for {
		select {
		case c := <-cmds:
			if err := generate(c, config, pm); err != nil {
				l.Error(fmt.Errorf("Couldn't generate report: %s", err))
			}
		case <-quit.C:
			if err := generate(cmd.All, config, pm); err != nil {
				l.Error(fmt.Errorf("Couldn't generate report: %s", err))
			}
			return
		}
	}
}

func generate(
	c cmd.Cmd,
	config *config.Config,
	pm *progress.Monitor,
) error {
	switch c {
	case cmd.Progress:
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
		if err := generateFront(config, pm); err != nil {
			return err
		}
	case cmd.Reports:
		if err := generateFront(config, pm); err != nil {
			return err
		}
		if err := generateReports(config); err != nil {
			return err
		}
		if err := generateTagPages(config); err != nil {
			return err
		}
		if err := generateCategoryPages(config); err != nil {
			return err
		}
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
	case cmd.All:
		if err := generateFront(config, pm); err != nil {
			return err
		}
		if err := generateReports(config); err != nil {
			return err
		}
		if err := generateTagPages(config); err != nil {
			return err
		}
		if err := generateCategoryPages(config); err != nil {
			return err
		}
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
	}
	return nil
}

var nonAlphaNumOrSpaceRegexp = regexp.MustCompile("[^[:alnum:] ]")
var spaceRegexp = regexp.MustCompile(" ")
var multipleHyphenRegexp = regexp.MustCompile("-+")

func escapeString(s string) string {
	newS := nonAlphaNumOrSpaceRegexp.ReplaceAllString(s, "")
	newS = spaceRegexp.ReplaceAllString(newS, "-")
	newS = multipleHyphenRegexp.ReplaceAllString(newS, "-")
	newS = strings.Trim(newS, " -")
	return strings.ToLower(newS)
}

// CreatePageError indicates that an html page can't be created
type CreatePageError struct {
	Filename string
	Op       string
	Err      error
}

func (wpe CreatePageError) Error() string {
	return fmt.Sprintf(
		"can't create html page for filename: %s, error: %s, Op: %s",
		wpe.Filename, wpe.Err, wpe.Op)
}

func writeTemplate(
	config *config.Config,
	filename string,
	tpl string,
	tplData interface{},
) error {
	funcMap := template.FuncMap{
		"ToTitle": strings.Title,
	}
	t, err := template.New("webpage").Funcs(funcMap).Parse(tpl)
	if err != nil {
		return CreatePageError{Filename: filename, Op: "Parse", Err: err}
	}

	fullFilename := filepath.Join(config.WWWDir, filename)
	dir := filepath.Dir(fullFilename)
	if err := os.MkdirAll(dir, modePerm); err != nil {
		return CreatePageError{Filename: filename, Op: "MkdirAll", Err: err}
	}

	f, err := os.Create(fullFilename)
	if err != nil {
		return CreatePageError{Filename: filename, Op: "Create", Err: err}
	}
	defer f.Close()

	if err := t.Execute(f, tplData); err != nil {
		return CreatePageError{Filename: filename, Op: "Execute", Err: err}
	}
	return nil
}

func genReportURLDir(category string, title string) string {
	escapedTitle := escapeString(title)
	escapedCategory := escapeString(category)
	if category != "" {
		return fmt.Sprintf("reports/category/%s/%s/", escapedCategory, escapedTitle)
	}
	return fmt.Sprintf("reports/nocategory/%s/", escapedTitle)
}

func countFiles(files []os.FileInfo) int {
	numFiles := 0
	for _, file := range files {
		if !file.IsDir() {
			numFiles++
		}
	}
	return numFiles
}

func makeHtmlNav(menuItem string) template.HTML {
	var doc bytes.Buffer
	validMenuItems := []string{
		"front",
		"reports",
		"tag",
		"category",
		"activity",
	}

	foundValidItem := false
	for _, validMenuItem := range validMenuItems {
		if validMenuItem == menuItem {
			foundValidItem = true
		}
	}
	if !foundValidItem {
		panic(fmt.Sprintf("menuItem not valid: %s", menuItem))
	}

	t, err := template.New("webpage").Parse(navTpl)
	if err != nil {
		panic(fmt.Sprintf("Couldn't create nav html: %s", err))
	}

	tplData := struct{ MenuItem string }{menuItem}

	if err := t.Execute(&doc, tplData); err != nil {
		panic(fmt.Sprintf("Couldn't create nav html: %s", err))
	}
	return template.HTML(doc.String())
}

func makeHtmlHead(c *config.Config) template.HTML {
	var doc bytes.Buffer
	t, err := template.New("webpage").Parse(headTpl)
	if err != nil {
		panic(fmt.Sprintf("Couldn't create head html: %s", err))
	}

	tplData := struct {
		BaseURL   string
		JSComment template.HTML
	}{BaseURL: c.BaseURL, JSComment: template.HTML(headJSComment)}

	if err := t.Execute(&doc, tplData); err != nil {
		panic(fmt.Sprintf("Couldn't create head html: %s", err))
	}
	return template.HTML(doc.String())
}

func makeHtml(c *config.Config, menuItem string) map[string]template.HTML {
	r := make(map[string]template.HTML)
	r["head"] = makeHtmlHead(c)
	r["nav"] = makeHtmlNav(menuItem)
	r["footer"] = template.HTML(footerHtml)
	r["bootstrapJS"] = template.HTML(bootstrapJSHtml)
	return r
}
<commit_msg>Make CreatePageError Ops' lowercase<commit_after>// Copyright (C) 2016-2017 vLife Systems Ltd <http://vlifesystems.com>
// Licensed under an MIT licence.  Please see LICENSE.md for details.

package html

import (
	"bytes"
	"fmt"
	"github.com/vlifesystems/rulehunter/config"
	"github.com/vlifesystems/rulehunter/html/cmd"
	"github.com/vlifesystems/rulehunter/logger"
	"github.com/vlifesystems/rulehunter/progress"
	"github.com/vlifesystems/rulehunter/quitter"
	"html/template"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

// File mode permission used as standard for the html content:
// No special permission bits
// User: Read, Write Execute
// Group: Read
// Other: None
const modePerm = 0740

// This should be run as a goroutine and each time a command is passed to
// cmds the html will be generated
func Run(
	config *config.Config,
	pm *progress.Monitor,
	l logger.Logger,
	quit *quitter.Quitter,
	cmds <-chan cmd.Cmd,
) {
	quit.Add()
	defer quit.Done()
	if err := generate(cmd.All, config, pm); err != nil {
		l.Error(fmt.Errorf("Couldn't generate report: %s", err))
	}

	for {
		select {
		case c := <-cmds:
			if err := generate(c, config, pm); err != nil {
				l.Error(fmt.Errorf("Couldn't generate report: %s", err))
			}
		case <-quit.C:
			if err := generate(cmd.All, config, pm); err != nil {
				l.Error(fmt.Errorf("Couldn't generate report: %s", err))
			}
			return
		}
	}
}

func generate(
	c cmd.Cmd,
	config *config.Config,
	pm *progress.Monitor,
) error {
	switch c {
	case cmd.Progress:
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
		if err := generateFront(config, pm); err != nil {
			return err
		}
	case cmd.Reports:
		if err := generateFront(config, pm); err != nil {
			return err
		}
		if err := generateReports(config); err != nil {
			return err
		}
		if err := generateTagPages(config); err != nil {
			return err
		}
		if err := generateCategoryPages(config); err != nil {
			return err
		}
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
	case cmd.All:
		if err := generateFront(config, pm); err != nil {
			return err
		}
		if err := generateReports(config); err != nil {
			return err
		}
		if err := generateTagPages(config); err != nil {
			return err
		}
		if err := generateCategoryPages(config); err != nil {
			return err
		}
		if err := generateActivityPage(config, pm); err != nil {
			return err
		}
	}
	return nil
}

var nonAlphaNumOrSpaceRegexp = regexp.MustCompile("[^[:alnum:] ]")
var spaceRegexp = regexp.MustCompile(" ")
var multipleHyphenRegexp = regexp.MustCompile("-+")

func escapeString(s string) string {
	newS := nonAlphaNumOrSpaceRegexp.ReplaceAllString(s, "")
	newS = spaceRegexp.ReplaceAllString(newS, "-")
	newS = multipleHyphenRegexp.ReplaceAllString(newS, "-")
	newS = strings.Trim(newS, " -")
	return strings.ToLower(newS)
}

// CreatePageError indicates that an html page can't be created
type CreatePageError struct {
	Filename string
	Op       string
	Err      error
}

func (wpe CreatePageError) Error() string {
	return fmt.Sprintf(
		"can't create html page for filename: %s, error: %s, Op: %s",
		wpe.Filename, wpe.Err, wpe.Op)
}

func writeTemplate(
	config *config.Config,
	filename string,
	tpl string,
	tplData interface{},
) error {
	funcMap := template.FuncMap{
		"ToTitle": strings.Title,
	}
	t, err := template.New("webpage").Funcs(funcMap).Parse(tpl)
	if err != nil {
		return CreatePageError{Filename: filename, Op: "parse", Err: err}
	}

	fullFilename := filepath.Join(config.WWWDir, filename)
	dir := filepath.Dir(fullFilename)
	if err := os.MkdirAll(dir, modePerm); err != nil {
		return CreatePageError{Filename: filename, Op: "mkdir", Err: err}
	}

	f, err := os.Create(fullFilename)
	if err != nil {
		return CreatePageError{Filename: filename, Op: "create", Err: err}
	}
	defer f.Close()

	if err := t.Execute(f, tplData); err != nil {
		return CreatePageError{Filename: filename, Op: "execute", Err: err}
	}
	return nil
}

func genReportURLDir(category string, title string) string {
	escapedTitle := escapeString(title)
	escapedCategory := escapeString(category)
	if category != "" {
		return fmt.Sprintf("reports/category/%s/%s/", escapedCategory, escapedTitle)
	}
	return fmt.Sprintf("reports/nocategory/%s/", escapedTitle)
}

func countFiles(files []os.FileInfo) int {
	numFiles := 0
	for _, file := range files {
		if !file.IsDir() {
			numFiles++
		}
	}
	return numFiles
}

func makeHtmlNav(menuItem string) template.HTML {
	var doc bytes.Buffer
	validMenuItems := []string{
		"front",
		"reports",
		"tag",
		"category",
		"activity",
	}

	foundValidItem := false
	for _, validMenuItem := range validMenuItems {
		if validMenuItem == menuItem {
			foundValidItem = true
		}
	}
	if !foundValidItem {
		panic(fmt.Sprintf("menuItem not valid: %s", menuItem))
	}

	t, err := template.New("webpage").Parse(navTpl)
	if err != nil {
		panic(fmt.Sprintf("Couldn't create nav html: %s", err))
	}

	tplData := struct{ MenuItem string }{menuItem}

	if err := t.Execute(&doc, tplData); err != nil {
		panic(fmt.Sprintf("Couldn't create nav html: %s", err))
	}
	return template.HTML(doc.String())
}

func makeHtmlHead(c *config.Config) template.HTML {
	var doc bytes.Buffer
	t, err := template.New("webpage").Parse(headTpl)
	if err != nil {
		panic(fmt.Sprintf("Couldn't create head html: %s", err))
	}

	tplData := struct {
		BaseURL   string
		JSComment template.HTML
	}{BaseURL: c.BaseURL, JSComment: template.HTML(headJSComment)}

	if err := t.Execute(&doc, tplData); err != nil {
		panic(fmt.Sprintf("Couldn't create head html: %s", err))
	}
	return template.HTML(doc.String())
}

func makeHtml(c *config.Config, menuItem string) map[string]template.HTML {
	r := make(map[string]template.HTML)
	r["head"] = makeHtmlHead(c)
	r["nav"] = makeHtmlNav(menuItem)
	r["footer"] = template.HTML(footerHtml)
	r["bootstrapJS"] = template.HTML(bootstrapJSHtml)
	return r
}
<|endoftext|>


<commit_before>package http

import (
	"context"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"github.com/aptly-dev/aptly/utils"
	"github.com/cavaliercoder/grab"
	"net/http"
	"path/filepath"
	"time"

	"github.com/aptly-dev/aptly/aptly"
)

type GrabDownloader struct {
	client   *grab.Client
	maxTries int
}

// Check interface
var (
	_ aptly.Downloader = (*GrabDownloader)(nil)
)

// NewGrabDownloader creates new expected downloader
func NewGrabDownloader(downLimit int64, maxTries int, progress aptly.Progress) *GrabDownloader {
	// TODO rate limiting and progress
	client := grab.NewClient()
	return &GrabDownloader{
		client:   client,
		maxTries: maxTries,
	}
}

func (d *GrabDownloader) Download(ctx context.Context, url string, destination string) error {
	return d.DownloadWithChecksum(ctx, url, destination, nil, false)
}

func (d *GrabDownloader) DownloadWithChecksum(ctx context.Context, url string, destination string, expected *utils.ChecksumInfo, ignoreMismatch bool) error {
	maxTries := d.maxTries
	const delayMax = time.Duration(5 * time.Minute)
	delay := time.Duration(1 * time.Second)
	const delayMultiplier = 2
	err := fmt.Errorf("No tries available")
	for maxTries > 0 {
		err = d.download(ctx, url, destination, expected, ignoreMismatch)
		if err == nil {
			// Success
			break
		}
		if retryableError(err) {
			maxTries--
			time.Sleep(delay)
		} else {
			// Can't retry
			break
		}
	}
	return err
}

func (d *GrabDownloader) download(ctx context.Context, url string, destination string, expected *utils.ChecksumInfo, ignoreMismatch bool) error {
	// TODO clean up dest dir on permanent failure
	// TODO maxTries
	destDir := filepath.Dir(destination)
	req, _ := grab.NewRequest(destDir, url)

	// TODO ignoreMismatch
	if expected != nil {
		if expected.MD5 != "" {
			req.SetChecksum(md5.New(), []byte(expected.MD5), true)
		} else if expected.SHA1 != "" {
			req.SetChecksum(sha1.New(), []byte(expected.SHA1), true)
		} else if expected.SHA256 != "" {
			req.SetChecksum(sha256.New(), []byte(expected.SHA256), true)
		} else if expected.SHA512 != "" {
			req.SetChecksum(sha512.New(), []byte(expected.SHA512), true)
		}
	}

	resp := d.client.Do(req)

Loop:
	for {
		select {
		case <-resp.Done:
			// download is complete
			break Loop
		}
	}
	return resp.Err()
}

func (d *GrabDownloader) GetProgress() aptly.Progress {
	// TODO
	return nil
}

func (f *GrabDownloader) GetLength(ctx context.Context, url string) (int64, error) {
	resp, err := http.Head(url)
	if err != nil {
		return -1, err
	}

	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		return -1, &Error{Code: resp.StatusCode, URL: url}
	}

	if resp.ContentLength < 0 {
		return -1, fmt.Errorf("could not determine length of %s", url)
	}

	return resp.ContentLength, nil
}
<commit_msg>Improve logging in grab downloader<commit_after>package http

import (
	"context"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"github.com/aptly-dev/aptly/utils"
	"github.com/cavaliercoder/grab"
	"net/http"
	"path/filepath"
	"time"

	"github.com/aptly-dev/aptly/aptly"
)

type GrabDownloader struct {
	client   *grab.Client
	maxTries int
	progress aptly.Progress
}

// Check interface
var (
	_ aptly.Downloader = (*GrabDownloader)(nil)
)

// NewGrabDownloader creates new expected downloader
func NewGrabDownloader(downLimit int64, maxTries int, progress aptly.Progress) *GrabDownloader {
	// TODO rate limiting and progress
	client := grab.NewClient()
	return &GrabDownloader{
		client:   client,
		maxTries: maxTries,
	}
}

func (d *GrabDownloader) Download(ctx context.Context, url string, destination string) error {
	return d.DownloadWithChecksum(ctx, url, destination, nil, false)
}

func (d *GrabDownloader) DownloadWithChecksum(ctx context.Context, url string, destination string, expected *utils.ChecksumInfo, ignoreMismatch bool) error {
	maxTries := d.maxTries
	const delayMax = time.Duration(5 * time.Minute)
	delay := time.Duration(1 * time.Second)
	const delayMultiplier = 2
	err := fmt.Errorf("No tries available")
	for maxTries > 0 {
		err = d.download(ctx, url, destination, expected, ignoreMismatch)
		if err == nil {
			// Success
			break
		}
		d.log("Error downloading %s: %v\n", url, err)
		if retryableError(err) {
			maxTries--
			d.log("Retrying download %s: %d\n", url, maxTries)
			time.Sleep(delay)
		} else {
			// Can't retry
			d.log("Cannot retry download %s\n", url)
			break
		}
	}
	return err
}

func (d *GrabDownloader) log(msg string, a ...interface{}) {
	if d.progress != nil {
		d.progress.Printf(msg, a...)
	}
}

func (d *GrabDownloader) download(ctx context.Context, url string, destination string, expected *utils.ChecksumInfo, ignoreMismatch bool) error {
	// TODO clean up dest dir on permanent failure
	// TODO maxTries
	d.log("Starting download %s -> %s\n", url, destination)
	destDir := filepath.Dir(destination)
	req, _ := grab.NewRequest(destDir, url)

	// TODO ignoreMismatch
	if expected != nil {
		if expected.MD5 != "" {
			req.SetChecksum(md5.New(), []byte(expected.MD5), true)
		} else if expected.SHA1 != "" {
			req.SetChecksum(sha1.New(), []byte(expected.SHA1), true)
		} else if expected.SHA256 != "" {
			req.SetChecksum(sha256.New(), []byte(expected.SHA256), true)
		} else if expected.SHA512 != "" {
			req.SetChecksum(sha512.New(), []byte(expected.SHA512), true)
		}
	}

	resp := d.client.Do(req)

Loop:
	for {
		select {
		case <-resp.Done:
			// download is complete
			break Loop
		}
	}
	return resp.Err()
}

func (d *GrabDownloader) GetProgress() aptly.Progress {
	return d.progress
}

func (f *GrabDownloader) GetLength(ctx context.Context, url string) (int64, error) {
	resp, err := http.Head(url)
	if err != nil {
		return -1, err
	}

	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		return -1, &Error{Code: resp.StatusCode, URL: url}
	}

	if resp.ContentLength < 0 {
		return -1, fmt.Errorf("could not determine length of %s", url)
	}

	return resp.ContentLength, nil
}
<|endoftext|>