<commit_before>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"compress/zlib"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	chars     = "abcdefghijklmnopqrstuvwxyz0123456789"

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	idSize                   int
	maxSize                  ByteSize

	validId       *regexp.Regexp
	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	pasteInfos    = make(map[Id]PasteInfo)
	customRand    *rand.Rand
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
	flag.IntVar(&idSize, "i", 8, "Size of the paste ids (between 6 and 256)")
	validId = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.Itoa(idSize) + "}$")
	customRand = rand.New(rand.NewSource(time.Now().UnixNano()))
}

type PasteInfo struct {
	ModTime time.Time
}

type Id string

func IdFromPath(idPath string) (Id, error) {
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 3 {
		return "", errors.New("Found invalid number of directories at " + idPath)
	}
	rawId := parts[0] + parts[1] + parts[2]
	if !validId.MatchString(rawId) {
		return "", errors.New("Found invalid id " + rawId)
	}
	return Id(rawId), nil
}

func RandomId() Id {
	s := make([]byte, idSize)
	var offset int = 0
MainLoop:
	for {
		r := customRand.Int63()
		for i := 0; i < 8; i++ {
			randbyte := int(r&0xff) % len(chars)
			s[offset] = chars[randbyte]
			offset++
			if offset == idSize {
				break MainLoop
			}
			r >>= 8
		}
	}
	return Id(s)
}

func (id Id) String() string {
	return string(id)
}

func (id Id) Path() string {
	return path.Join(string(id[0:2]), string(id[2:4]), string(id[4:]))
}

func (id Id) EndLife() {
	err := os.Remove(id.Path())
	if err == nil {
		delete(pasteInfos, id)
		log.Printf("Removed paste: %s", id)
	} else {
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, siteUrl)
			return
		case "/form":
			formTemplate.Execute(w, siteUrl)
			return
		}
		rawId := r.URL.Path[1:]
		if !validId.MatchString(rawId) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id := Id(strings.ToLower(rawId))
		pasteInfo, e := pasteInfos[id]
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		etag := fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", etag)
		compReader, err := zlib.NewReader(pasteFile)
		if err != nil {
			log.Printf("Could not open a compression reader for %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer compReader.Close()
		io.Copy(w, compReader)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var id Id
		var content string
		found := false
		for i := 0; i < 10; i++ {
			id = RandomId()
			if _, e := pasteInfos[id]; !e {
				found = true
				break
			}
		}
		if !found {
			log.Printf("Gave up trying to find an unused random id")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		compWriter := zlib.NewWriter(pasteFile)
		defer compWriter.Close()
		b, err := io.WriteString(compWriter, content)
		if err != nil {
			log.Printf("Could not write compressed data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		pasteInfos[id] = PasteInfo{ModTime: time.Now()}
		log.Printf("Created a new paste: %s (%s)", id, writtenSize)
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	log.Printf("Recovered paste %s has %s left", id, lifeLeft)
	pasteInfos[id] = PasteInfo{ModTime: modTime}
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	flag.Parse()
	if idSize < 6 || idSize > 256 {
		log.Fatalf("Provided id size %d is not between 6 and 256", idSize)
	}
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<commit_msg>Fill Last-Modified in responses<commit_after>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"compress/zlib"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	indexTmpl = "index.html"
	formTmpl  = "form.html"
	chars     = "abcdefghijklmnopqrstuvwxyz0123456789"

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl, listen, dataDir string
	lifeTime                 time.Duration
	maxSizeStr               string
	idSize                   int
	maxSize                  ByteSize

	validId       *regexp.Regexp
	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
	formTemplate  *template.Template
	pasteInfos    = make(map[Id]PasteInfo)
	customRand    *rand.Rand
)

func init() {
	flag.StringVar(&siteUrl, "u", "http://localhost:8080", "URL of the site")
	flag.StringVar(&listen, "l", "localhost:8080", "Host and port to listen to")
	flag.StringVar(&dataDir, "d", "data", "Directory to store all the pastes in")
	flag.DurationVar(&lifeTime, "t", 12*time.Hour, "Lifetime of the pastes (units: s,m,h)")
	flag.StringVar(&maxSizeStr, "s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")
	flag.IntVar(&idSize, "i", 8, "Size of the paste ids (between 6 and 256)")
	validId = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.Itoa(idSize) + "}$")
	customRand = rand.New(rand.NewSource(time.Now().UnixNano()))
}

type PasteInfo struct {
	ModTime time.Time
}

type Id string

func IdFromPath(idPath string) (Id, error) {
	parts := strings.Split(idPath, string(filepath.Separator))
	if len(parts) != 3 {
		return "", errors.New("Found invalid number of directories at " + idPath)
	}
	rawId := parts[0] + parts[1] + parts[2]
	if !validId.MatchString(rawId) {
		return "", errors.New("Found invalid id " + rawId)
	}
	return Id(rawId), nil
}

func RandomId() Id {
	s := make([]byte, idSize)
	var offset int = 0
MainLoop:
	for {
		r := customRand.Int63()
		for i := 0; i < 8; i++ {
			randbyte := int(r&0xff) % len(chars)
			s[offset] = chars[randbyte]
			offset++
			if offset == idSize {
				break MainLoop
			}
			r >>= 8
		}
	}
	return Id(s)
}

func (id Id) String() string {
	return string(id)
}

func (id Id) Path() string {
	return path.Join(string(id[0:2]), string(id[2:4]), string(id[4:]))
}

func (id Id) EndLife() {
	err := os.Remove(id.Path())
	if err == nil {
		delete(pasteInfos, id)
		log.Printf("Removed paste: %s", id)
	} else {
		log.Printf("Could not end the life of %s: %s", id, err)
		id.EndLifeAfter(2 * time.Minute)
	}
}

func (id Id) EndLifeAfter(duration time.Duration) {
	timer := time.NewTimer(duration)
	go func() {
		<-timer.C
		id.EndLife()
	}()
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		switch r.URL.Path {
		case "/":
			indexTemplate.Execute(w, siteUrl)
			return
		case "/form":
			formTemplate.Execute(w, siteUrl)
			return
		}
		rawId := r.URL.Path[1:]
		if !validId.MatchString(rawId) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id := Id(strings.ToLower(rawId))
		pasteInfo, e := pasteInfos[id]
		if !e {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		etag := fmt.Sprintf("%d-%s", pasteInfo.ModTime.Unix(), id)
		if inm := r.Header.Get("If-None-Match"); inm != "" {
			if etag == inm || inm == "*" {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		pastePath := id.Path()
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		w.Header().Set("Etag", etag)
		w.Header().Set("Last-Modified", pasteInfo.ModTime.Format(http.TimeFormat))
		compReader, err := zlib.NewReader(pasteFile)
		if err != nil {
			log.Printf("Could not open a compression reader for %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer compReader.Close()
		io.Copy(w, compReader)

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var id Id
		var content string
		found := false
		for i := 0; i < 10; i++ {
			id = RandomId()
			if _, e := pasteInfos[id]; !e {
				found = true
				break
			}
		}
		if !found {
			log.Printf("Gave up trying to find an unused random id")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		if vs, found := r.Form["paste"]; found {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		pastePath := id.Path()
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		id.EndLifeAfter(lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		defer pasteFile.Close()
		compWriter := zlib.NewWriter(pasteFile)
		defer compWriter.Close()
		b, err := io.WriteString(compWriter, content)
		if err != nil {
			log.Printf("Could not write compressed data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		pasteInfos[id] = PasteInfo{ModTime: time.Now()}
		log.Printf("Created a new paste: %s (%s)", id, writtenSize)
		fmt.Fprintf(w, "%s/%s\n", siteUrl, id)
	}
}

func walkFunc(filePath string, fileInfo os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if fileInfo.IsDir() {
		return nil
	}
	id, err := IdFromPath(filePath)
	if err != nil {
		return err
	}
	modTime := fileInfo.ModTime()
	deathTime := modTime.Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go id.EndLife()
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	log.Printf("Recovered paste %s has %s left", id, lifeLeft)
	pasteInfos[id] = PasteInfo{ModTime: modTime}
	id.EndLifeAfter(lifeLeft)
	return nil
}

func main() {
	var err error
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", siteUrl)
	log.Printf("listen   = %s", listen)
	log.Printf("dataDir  = %s", dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	flag.Parse()
	if idSize < 6 || idSize > 256 {
		log.Fatalf("Provided id size %d is not between 6 and 256", idSize)
	}
	if maxSize, err = parseByteSize(maxSizeStr); err != nil {
		log.Fatalf("Invalid max size '%s': %s", maxSizeStr, err)
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", indexTmpl, err)
	}
	if formTemplate, err = template.ParseFiles(formTmpl); err != nil {
		log.Fatalf("Could not load template %s: %s", formTmpl, err)
	}
	if err = os.MkdirAll(dataDir, 0700); err != nil {
		log.Fatalf("Could not create data directory %s: %s", dataDir, err)
	}
	if err = os.Chdir(dataDir); err != nil {
		log.Fatalf("Could not enter data directory %s: %s", dataDir, err)
	}
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Fatalf("Could not recover data directory %s: %s", dataDir, err)
	}
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(listen, nil))
}
<|endoftext|>


<commit_before>package librato

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/influxdata/telegraf/plugins/serializers/graphite"
)

// Librato structure for configuration and client
type Librato struct {
	APIUser   string
	APIToken  string
	Debug     bool
	SourceTag string // Deprecated, keeping for backward-compatibility
	Timeout   internal.Duration
	Template  string

	APIUrl string
	client *http.Client
}

// https://www.librato.com/docs/kb/faq/best_practices/naming_convention_metrics_sources.html#naming-limitations-for-sources-and-metrics
var reUnacceptedChar = regexp.MustCompile("[^.a-zA-Z0-9_-]")

var sampleConfig = `
  ## Librator API Docs
  ## http://dev.librato.com/v1/metrics-authentication
  ## Librato API user
  api_user = "telegraf@influxdb.com" # required.
  ## Librato API token
  api_token = "my-secret-token" # required.
  ## Debug
  # debug = false
  ## Connection timeout.
  # timeout = "5s"
  ## Output source Template (same as graphite buckets)
  ## see https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_OUTPUT.md#graphite
  ## This template is used in librato's source (not metric's name)
  template = "host"

`

// LMetrics is the default struct for Librato's API fromat
type LMetrics struct {
	Gauges []*Gauge `json:"gauges"`
}

// Gauge is the gauge format for Librato's API fromat
type Gauge struct {
	Name        string  `json:"name"`
	Value       float64 `json:"value"`
	Source      string  `json:"source"`
	MeasureTime int64   `json:"measure_time"`
}

const libratoAPI = "https://metrics-api.librato.com/v1/metrics"

// NewLibrato is the main constructor for librato output plugins
func NewLibrato(apiURL string) *Librato {
	return &Librato{
		APIUrl:   apiURL,
		Template: "host",
	}
}

// Connect is the default output plugin connection function who make sure it
// can connect to the endpoint
func (l *Librato) Connect() error {
	if l.APIUser == "" || l.APIToken == "" {
		return fmt.Errorf(
			"api_user and api_token are required fields for librato output")
	}
	l.client = &http.Client{
		Timeout: l.Timeout.Duration,
	}
	return nil
}

func (l *Librato) Write(metrics []telegraf.Metric) error {

	if len(metrics) == 0 {
		return nil
	}
	if l.Template == "" {
		l.Template = "host"
	}
	if l.SourceTag != "" {
		l.Template = l.SourceTag
	}

	tempGauges := []*Gauge{}

	for _, m := range metrics {
		if gauges, err := l.buildGauges(m); err == nil {
			for _, gauge := range gauges {
				tempGauges = append(tempGauges, gauge)
				log.Printf("D! Got a gauge: %v\n", gauge)

			}
		} else {
			log.Printf("I! unable to build Gauge for %s, skipping\n", m.Name())
			log.Printf("D! Couldn't build gauge: %v\n", err)

		}
	}

	metricCounter := len(tempGauges)
	// make sur we send a batch of maximum 300
	sizeBatch := 300
	for start := 0; start < metricCounter; start += sizeBatch {
		lmetrics := LMetrics{}
		end := start + sizeBatch
		if end > metricCounter {
			end = metricCounter
			sizeBatch = end - start
		}
		lmetrics.Gauges = make([]*Gauge, sizeBatch)
		copy(lmetrics.Gauges, tempGauges[start:end])
		metricsBytes, err := json.Marshal(lmetrics)
		if err != nil {
			return fmt.Errorf("unable to marshal Metrics, %s\n", err.Error())
		}

		log.Printf("D! Librato request: %v\n", string(metricsBytes))

		req, err := http.NewRequest(
			"POST",
			l.APIUrl,
			bytes.NewBuffer(metricsBytes))
		if err != nil {
			return fmt.Errorf(
				"unable to create http.Request, %s\n",
				err.Error())
		}
		req.Header.Add("Content-Type", "application/json")
		req.SetBasicAuth(l.APIUser, l.APIToken)

		resp, err := l.client.Do(req)
		if err != nil {
			log.Printf("D! Error POSTing metrics: %v\n", err.Error())
			return fmt.Errorf("error POSTing metrics, %s\n", err.Error())
		}
		defer resp.Body.Close()

		if resp.StatusCode != 200 || l.Debug {
			htmlData, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Printf("D! Couldn't get response! (%v)\n", err)
			}
			if resp.StatusCode != 200 {
				return fmt.Errorf(
					"received bad status code, %d\n %s",
					resp.StatusCode,
					string(htmlData))
			}
			log.Printf("D! Librato response: %v\n", string(htmlData))
		}
	}

	return nil
}

// SampleConfig is function who return the default configuration for this
// output
func (l *Librato) SampleConfig() string {
	return sampleConfig
}

// Description is function who return the Description of this output
func (l *Librato) Description() string {
	return "Configuration for Librato API to send metrics to."
}

func (l *Librato) buildGauges(m telegraf.Metric) ([]*Gauge, error) {

	gauges := []*Gauge{}
	if m.Time().Unix() == 0 {
		return gauges, fmt.Errorf(
			"Measure time must not be zero\n <%s> \n",
			m.String())
	}
	metricSource := graphite.InsertField(
		graphite.SerializeBucketName("", m.Tags(), l.Template, ""),
		"value")
	if metricSource == "" {
		return gauges,
			fmt.Errorf("undeterminable Source type from Field, %s\n",
				l.Template)
	}
	for fieldName, value := range m.Fields() {

		metricName := m.Name()
		if fieldName != "value" {
			metricName = fmt.Sprintf("%s.%s", m.Name(), fieldName)
		}

		gauge := &Gauge{
			Source:      reUnacceptedChar.ReplaceAllString(metricSource, "-"),
			Name:        reUnacceptedChar.ReplaceAllString(metricName, "-"),
			MeasureTime: m.Time().Unix(),
		}
		if !verifyValue(value) {
			continue
		}
		if err := gauge.setValue(value); err != nil {
			return gauges, fmt.Errorf(
				"unable to extract value from Fields, %s\n",
				err.Error())
		}
		gauges = append(gauges, gauge)
	}

	log.Printf("D! Built gauges: %v\n", gauges)
	return gauges, nil
}

func verifyValue(v interface{}) bool {
	switch v.(type) {
	case string:
		return false
	}
	return true
}

func (g *Gauge) setValue(v interface{}) error {
	switch d := v.(type) {
	case int:
		g.Value = float64(int(d))
	case int32:
		g.Value = float64(int32(d))
	case int64:
		g.Value = float64(int64(d))
	case float32:
		g.Value = float64(d)
	case float64:
		g.Value = float64(d)
	default:
		return fmt.Errorf("undeterminable type %+v", d)
	}
	return nil
}

//Close is used to close the connection to librato Output
func (l *Librato) Close() error {
	return nil
}

func init() {
	outputs.Add("librato", func() telegraf.Output {
		return NewLibrato(libratoAPI)
	})
}
<commit_msg>Fix for broken librato output (#2225)<commit_after>package librato

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	"github.com/influxdata/telegraf/plugins/outputs"
	"github.com/influxdata/telegraf/plugins/serializers/graphite"
)

// Librato structure for configuration and client
type Librato struct {
	APIUser   string `toml:"api_user"`
	APIToken  string `toml:"api_token"`
	Debug     bool
	SourceTag string // Deprecated, keeping for backward-compatibility
	Timeout   internal.Duration
	Template  string

	APIUrl string
	client *http.Client
}

// https://www.librato.com/docs/kb/faq/best_practices/naming_convention_metrics_sources.html#naming-limitations-for-sources-and-metrics
var reUnacceptedChar = regexp.MustCompile("[^.a-zA-Z0-9_-]")

var sampleConfig = `
  ## Librator API Docs
  ## http://dev.librato.com/v1/metrics-authentication
  ## Librato API user
  api_user = "telegraf@influxdb.com" # required.
  ## Librato API token
  api_token = "my-secret-token" # required.
  ## Debug
  # debug = false
  ## Connection timeout.
  # timeout = "5s"
  ## Output source Template (same as graphite buckets)
  ## see https://github.com/influxdata/telegraf/blob/master/docs/DATA_FORMATS_OUTPUT.md#graphite
  ## This template is used in librato's source (not metric's name)
  template = "host"

`

// LMetrics is the default struct for Librato's API fromat
type LMetrics struct {
	Gauges []*Gauge `json:"gauges"`
}

// Gauge is the gauge format for Librato's API fromat
type Gauge struct {
	Name        string  `json:"name"`
	Value       float64 `json:"value"`
	Source      string  `json:"source"`
	MeasureTime int64   `json:"measure_time"`
}

const libratoAPI = "https://metrics-api.librato.com/v1/metrics"

// NewLibrato is the main constructor for librato output plugins
func NewLibrato(apiURL string) *Librato {
	return &Librato{
		APIUrl:   apiURL,
		Template: "host",
	}
}

// Connect is the default output plugin connection function who make sure it
// can connect to the endpoint
func (l *Librato) Connect() error {
	if l.APIUser == "" || l.APIToken == "" {
		return fmt.Errorf(
			"api_user and api_token are required fields for librato output")
	}
	l.client = &http.Client{
		Timeout: l.Timeout.Duration,
	}
	return nil
}

func (l *Librato) Write(metrics []telegraf.Metric) error {

	if len(metrics) == 0 {
		return nil
	}
	if l.Template == "" {
		l.Template = "host"
	}
	if l.SourceTag != "" {
		l.Template = l.SourceTag
	}

	tempGauges := []*Gauge{}

	for _, m := range metrics {
		if gauges, err := l.buildGauges(m); err == nil {
			for _, gauge := range gauges {
				tempGauges = append(tempGauges, gauge)
				log.Printf("D! Got a gauge: %v\n", gauge)

			}
		} else {
			log.Printf("I! unable to build Gauge for %s, skipping\n", m.Name())
			log.Printf("D! Couldn't build gauge: %v\n", err)

		}
	}

	metricCounter := len(tempGauges)
	// make sur we send a batch of maximum 300
	sizeBatch := 300
	for start := 0; start < metricCounter; start += sizeBatch {
		lmetrics := LMetrics{}
		end := start + sizeBatch
		if end > metricCounter {
			end = metricCounter
			sizeBatch = end - start
		}
		lmetrics.Gauges = make([]*Gauge, sizeBatch)
		copy(lmetrics.Gauges, tempGauges[start:end])
		metricsBytes, err := json.Marshal(lmetrics)
		if err != nil {
			return fmt.Errorf("unable to marshal Metrics, %s\n", err.Error())
		}

		log.Printf("D! Librato request: %v\n", string(metricsBytes))

		req, err := http.NewRequest(
			"POST",
			l.APIUrl,
			bytes.NewBuffer(metricsBytes))
		if err != nil {
			return fmt.Errorf(
				"unable to create http.Request, %s\n",
				err.Error())
		}
		req.Header.Add("Content-Type", "application/json")
		req.SetBasicAuth(l.APIUser, l.APIToken)

		resp, err := l.client.Do(req)
		if err != nil {
			log.Printf("D! Error POSTing metrics: %v\n", err.Error())
			return fmt.Errorf("error POSTing metrics, %s\n", err.Error())
		}
		defer resp.Body.Close()

		if resp.StatusCode != 200 || l.Debug {
			htmlData, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				log.Printf("D! Couldn't get response! (%v)\n", err)
			}
			if resp.StatusCode != 200 {
				return fmt.Errorf(
					"received bad status code, %d\n %s",
					resp.StatusCode,
					string(htmlData))
			}
			log.Printf("D! Librato response: %v\n", string(htmlData))
		}
	}

	return nil
}

// SampleConfig is function who return the default configuration for this
// output
func (l *Librato) SampleConfig() string {
	return sampleConfig
}

// Description is function who return the Description of this output
func (l *Librato) Description() string {
	return "Configuration for Librato API to send metrics to."
}

func (l *Librato) buildGauges(m telegraf.Metric) ([]*Gauge, error) {

	gauges := []*Gauge{}
	if m.Time().Unix() == 0 {
		return gauges, fmt.Errorf(
			"Measure time must not be zero\n <%s> \n",
			m.String())
	}
	metricSource := graphite.InsertField(
		graphite.SerializeBucketName("", m.Tags(), l.Template, ""),
		"value")
	if metricSource == "" {
		return gauges,
			fmt.Errorf("undeterminable Source type from Field, %s\n",
				l.Template)
	}
	for fieldName, value := range m.Fields() {

		metricName := m.Name()
		if fieldName != "value" {
			metricName = fmt.Sprintf("%s.%s", m.Name(), fieldName)
		}

		gauge := &Gauge{
			Source:      reUnacceptedChar.ReplaceAllString(metricSource, "-"),
			Name:        reUnacceptedChar.ReplaceAllString(metricName, "-"),
			MeasureTime: m.Time().Unix(),
		}
		if !verifyValue(value) {
			continue
		}
		if err := gauge.setValue(value); err != nil {
			return gauges, fmt.Errorf(
				"unable to extract value from Fields, %s\n",
				err.Error())
		}
		gauges = append(gauges, gauge)
	}

	log.Printf("D! Built gauges: %v\n", gauges)
	return gauges, nil
}

func verifyValue(v interface{}) bool {
	switch v.(type) {
	case string:
		return false
	}
	return true
}

func (g *Gauge) setValue(v interface{}) error {
	switch d := v.(type) {
	case int:
		g.Value = float64(int(d))
	case int32:
		g.Value = float64(int32(d))
	case int64:
		g.Value = float64(int64(d))
	case float32:
		g.Value = float64(d)
	case float64:
		g.Value = float64(d)
	default:
		return fmt.Errorf("undeterminable type %+v", d)
	}
	return nil
}

//Close is used to close the connection to librato Output
func (l *Librato) Close() error {
	return nil
}

func init() {
	outputs.Add("librato", func() telegraf.Output {
		return NewLibrato(libratoAPI)
	})
}
<|endoftext|>


<commit_before>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"testing"

	"github.com/casbin/casbin/persist"
)

func TestGetAndSetModel(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	model := e.GetModel()
	e.SetModel(model)
}

func TestCreateModelManually(t *testing.T) {
	m := NewModel()
	m.AddDef("r", "r", "sub, obj, act")
	m.AddDef("p", "p", "sub, obj, act")
	m.AddDef("e", "e", "some(where (p.eft == allow))")
	m.AddDef("m", "m", "r.sub == p.sub && keyMatch(r.obj, p.obj) && regexMatch(r.act, p.act)")

	a := persist.NewFileAdapter("examples/keymatch_policy.csv")

	e := NewEnforcer(m, a)

	testEnforce(t, e, "alice", "/alice_data/resource1", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource1", "POST", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "POST", false)

	testEnforce(t, e, "bob", "/alice_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/alice_data/resource1", "POST", false)
	testEnforce(t, e, "bob", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "bob", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "POST", true)
	testEnforce(t, e, "bob", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource2", "POST", true)

	testEnforce(t, e, "cathy", "/cathy_data", "GET", true)
	testEnforce(t, e, "cathy", "/cathy_data", "POST", true)
	testEnforce(t, e, "cathy", "/cathy_data", "DELETE", false)

	e = NewEnforcer(m)
	a.LoadPolicy(e.GetModel())

	testEnforce(t, e, "alice", "/alice_data/resource1", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource1", "POST", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "POST", false)

	testEnforce(t, e, "bob", "/alice_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/alice_data/resource1", "POST", false)
	testEnforce(t, e, "bob", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "bob", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "POST", true)
	testEnforce(t, e, "bob", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource2", "POST", true)

	testEnforce(t, e, "cathy", "/cathy_data", "GET", true)
	testEnforce(t, e, "cathy", "/cathy_data", "POST", true)
	testEnforce(t, e, "cathy", "/cathy_data", "DELETE", false)
}

func TestReloadPolicy(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	e.LoadPolicy()
	testGetPolicy(t, e, [][]string{{"alice", "data1", "read"}, {"bob", "data2", "write"}, {"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}})
}

func TestSavePolicy(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	e.SavePolicy()
}

func TestEnable(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	e.Enable(false)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", true)
	testEnforce(t, e, "alice", "data2", "read", true)
	testEnforce(t, e, "alice", "data2", "write", true)
	testEnforce(t, e, "bob", "data1", "read", true)
	testEnforce(t, e, "bob", "data1", "write", true)
	testEnforce(t, e, "bob", "data2", "read", true)
	testEnforce(t, e, "bob", "data2", "write", true)

	e.Enable(true)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}

func TestEnableLog(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv", true)
	// The log is enabled by default, so the above is the same with:
	// e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)

	// The log can also be enabled or disabled at run-time.
	e.EnableLog(false)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}

func TestInitWithAdapter(t *testing.T) {
	adapter := persist.NewFileAdapter("examples/basic_policy.csv")
	e := NewEnforcer("examples/basic_model.conf", adapter)

	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}
<commit_msg>Add TestCreateModelManually2() example.<commit_after>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"testing"

	"github.com/casbin/casbin/persist"
)

func TestGetAndSetModel(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	model := e.GetModel()
	e.SetModel(model)
}

func TestCreateModelManually(t *testing.T) {
	m := NewModel()
	m.AddDef("r", "r", "sub, obj, act")
	m.AddDef("p", "p", "sub, obj, act")
	m.AddDef("e", "e", "some(where (p.eft == allow))")
	m.AddDef("m", "m", "r.sub == p.sub && keyMatch(r.obj, p.obj) && regexMatch(r.act, p.act)")

	a := persist.NewFileAdapter("examples/keymatch_policy.csv")

	e := NewEnforcer(m, a)

	testEnforce(t, e, "alice", "/alice_data/resource1", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource1", "POST", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "POST", false)

	testEnforce(t, e, "bob", "/alice_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/alice_data/resource1", "POST", false)
	testEnforce(t, e, "bob", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "bob", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "POST", true)
	testEnforce(t, e, "bob", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource2", "POST", true)

	testEnforce(t, e, "cathy", "/cathy_data", "GET", true)
	testEnforce(t, e, "cathy", "/cathy_data", "POST", true)
	testEnforce(t, e, "cathy", "/cathy_data", "DELETE", false)

	e = NewEnforcer(m)
	a.LoadPolicy(e.GetModel())

	testEnforce(t, e, "alice", "/alice_data/resource1", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource1", "POST", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "alice", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource1", "POST", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "alice", "/bob_data/resource2", "POST", false)

	testEnforce(t, e, "bob", "/alice_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/alice_data/resource1", "POST", false)
	testEnforce(t, e, "bob", "/alice_data/resource2", "GET", true)
	testEnforce(t, e, "bob", "/alice_data/resource2", "POST", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource1", "POST", true)
	testEnforce(t, e, "bob", "/bob_data/resource2", "GET", false)
	testEnforce(t, e, "bob", "/bob_data/resource2", "POST", true)

	testEnforce(t, e, "cathy", "/cathy_data", "GET", true)
	testEnforce(t, e, "cathy", "/cathy_data", "POST", true)
	testEnforce(t, e, "cathy", "/cathy_data", "DELETE", false)
}

func TestCreateModelManually2(t *testing.T) {
	m := NewModel()
	m.AddDef("r", "r", "sub, obj, act")
	m.AddDef("p", "p", "sub, obj, act")
	m.AddDef("g", "g", "_, _")
	m.AddDef("e", "e", "some(where (p.eft == allow))")
	m.AddDef("m", "m", "g(r.sub, p.sub) && r.obj == p.obj && r.act == p.act")

	e := NewEnforcer(m)

	e.AddPermissionForUser("alice", "data1", "read")
	e.AddPermissionForUser("bob", "data2", "write")
	e.AddPermissionForUser("data2_admin", "data2", "read")
	e.AddPermissionForUser("data2_admin", "data2", "write")
	e.AddRoleForUser("alice", "data2_admin")

	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", true)
	testEnforce(t, e, "alice", "data2", "write", true)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}

func TestReloadPolicy(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	e.LoadPolicy()
	testGetPolicy(t, e, [][]string{{"alice", "data1", "read"}, {"bob", "data2", "write"}, {"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}})
}

func TestSavePolicy(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	e.SavePolicy()
}

func TestEnable(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	e.Enable(false)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", true)
	testEnforce(t, e, "alice", "data2", "read", true)
	testEnforce(t, e, "alice", "data2", "write", true)
	testEnforce(t, e, "bob", "data1", "read", true)
	testEnforce(t, e, "bob", "data1", "write", true)
	testEnforce(t, e, "bob", "data2", "read", true)
	testEnforce(t, e, "bob", "data2", "write", true)

	e.Enable(true)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}

func TestEnableLog(t *testing.T) {
	e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv", true)
	// The log is enabled by default, so the above is the same with:
	// e := NewEnforcer("examples/basic_model.conf", "examples/basic_policy.csv")

	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)

	// The log can also be enabled or disabled at run-time.
	e.EnableLog(false)
	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}

func TestInitWithAdapter(t *testing.T) {
	adapter := persist.NewFileAdapter("examples/basic_policy.csv")
	e := NewEnforcer("examples/basic_model.conf", adapter)

	testEnforce(t, e, "alice", "data1", "read", true)
	testEnforce(t, e, "alice", "data1", "write", false)
	testEnforce(t, e, "alice", "data2", "read", false)
	testEnforce(t, e, "alice", "data2", "write", false)
	testEnforce(t, e, "bob", "data1", "read", false)
	testEnforce(t, e, "bob", "data1", "write", false)
	testEnforce(t, e, "bob", "data2", "read", false)
	testEnforce(t, e, "bob", "data2", "write", true)
}
<|endoftext|>


<commit_before>package engine

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"

	_ "github.com/dropbox/changes-client/adapter/basic"
	_ "github.com/dropbox/changes-client/adapter/lxc"
	_ "github.com/dropbox/changes-client/reporter/artifactstore"
	_ "github.com/dropbox/changes-client/reporter/jenkins"
	_ "github.com/dropbox/changes-client/reporter/mesos"
	_ "github.com/dropbox/changes-client/reporter/multireporter"
)

const (
	STATUS_QUEUED      = "queued"
	STATUS_IN_PROGRESS = "in_progress"
	STATUS_FINISHED    = "finished"

	RESULT_PASSED  Result = "passed"
	RESULT_FAILED  Result = "failed"
	RESULT_ABORTED Result = "aborted"
	// Test results unreliable or unavailable due to infrastructure
	// issues.
	RESULT_INFRA_FAILED Result = "infra_failed"

	SNAPSHOT_ACTIVE = "active"
	SNAPSHOT_FAILED = "failed"
)

type Result string

func (r Result) String() string {
	return string(r)
}

// Convenience method to check for all types of failure.
func (r Result) IsFailure() bool {
	switch r {
	case RESULT_FAILED, RESULT_INFRA_FAILED:
		return true
	}
	return false
}

var (
	selectedAdapter  string
	selectedReporter string
	outputSnapshot   string
)

type Engine struct {
	config    *client.Config
	clientLog *client.Log
	adapter   adapter.Adapter
	reporter  reporter.Reporter
}

func RunBuildPlan(config *client.Config) (Result, error) {
	var err error

	currentAdapter, err := adapter.Get(selectedAdapter)
	if err != nil {
		// TODO(dcramer): handle this error
		return RESULT_FAILED, err
	}

	currentReporter, err := reporter.Get(selectedReporter)
	if err != nil {
		log.Printf("[engine] failed to initialize reporter: %s", selectedReporter)
		return RESULT_FAILED, err
	}
	log.Printf("[engine] started with reporter %s, adapter %s", selectedReporter, selectedAdapter)

	engine := &Engine{
		config:    config,
		clientLog: client.NewLog(),
		adapter:   currentAdapter,
		reporter:  currentReporter,
	}

	return engine.Run()
}

func (e *Engine) Run() (Result, error) {
	defer e.reporter.Shutdown()

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		reportLogChunks("console", e.clientLog, e.reporter)
		wg.Done()
	}()

	e.reporter.Init(e.config)

	e.reporter.PushJobstepStatus(STATUS_IN_PROGRESS, "")

	result, err := e.runBuildPlan()

	e.clientLog.Writeln(fmt.Sprintf("==> Build finished! Recorded result as %s", result))

	e.reporter.PushJobstepStatus(STATUS_FINISHED, result.String())

	e.clientLog.Close()

	wg.Wait()

	return result, err
}

func (e *Engine) executeCommands() Result {
	result := RESULT_PASSED

	wg := sync.WaitGroup{}

	for _, cmdConfig := range e.config.Cmds {
		e.clientLog.Writeln(fmt.Sprintf("==> Running command %s", cmdConfig.ID))
		e.clientLog.Writeln(fmt.Sprintf("==>     with script %s", cmdConfig.Script))
		cmd, err := client.NewCommand(cmdConfig.ID, cmdConfig.Script)
		if err != nil {
			e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 255)
			result = RESULT_FAILED
			e.clientLog.Writeln(fmt.Sprintf("==> Error: %s", err.Error()))
			break
		}
		e.reporter.PushCommandStatus(cmd.ID, STATUS_IN_PROGRESS, -1)

		cmd.CaptureOutput = cmdConfig.CaptureOutput

		env := os.Environ()
		for k, v := range cmdConfig.Env {
			env = append(env, k+"="+v)
		}
		cmd.Env = env

		if len(cmdConfig.Cwd) > 0 {
			cmd.Cwd = cmdConfig.Cwd
		}

		cmdResult, err := e.adapter.Run(cmd, e.clientLog)

		if err != nil {
			e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 255)
			e.clientLog.Writeln(fmt.Sprintf("==> Error: %s", err.Error()))
			result = RESULT_FAILED
		} else {
			if cmdResult.Success {
				if cmd.CaptureOutput {
					e.reporter.PushCommandOutput(cmd.ID, STATUS_FINISHED, 0, cmdResult.Output)
				} else {
					e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 0)
				}
			} else {
				e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 1)
				result = RESULT_FAILED
			}
		}

		wg.Add(1)
		go func() {
			// publishArtifacts is a synchronous operation and doesnt follow the normal queue flow of
			// other operations
			e.reporter.PublishArtifacts(cmdConfig, e.adapter, e.clientLog)
			wg.Done()
		}()

		if result.IsFailure() {
			break
		}
	}

	wg.Wait()

	return result
}

func (e *Engine) captureSnapshot() error {
	log.Printf("[adapter] Capturing snapshot %s", outputSnapshot)
	err := e.adapter.CaptureSnapshot(outputSnapshot, e.clientLog)
	if err != nil {
		log.Printf("[adapter] Failed to capture snapshot: %s", err.Error())
		return err
	}
	return nil
}

func (e *Engine) runBuildPlan() (Result, error) {
	var (
		result Result
		err    error
	)

	// cancellation signal
	cancel := make(chan struct{})

	// capture ctrl+c and enforce a clean shutdown
	sigchan := make(chan os.Signal, 1)
	signal.Notify(sigchan, os.Interrupt)
	go func() {
		shuttingDown := false
		for _ = range sigchan {
			if shuttingDown {
				log.Printf("Second interrupt received. Terminating!")
				os.Exit(1)
			}

			shuttingDown = true

			log.Printf("Interrupted! Cancelling execution and cleaning up..")
			cancel <- struct{}{}
		}
	}()

	// We need to ensure that we're able to abort the build if upstream suggests
	// that it's been cancelled.
	if !e.config.Debug {
		go func() {
			um := &UpstreamMonitor{
				Config: e.config,
			}
			um.WaitUntilAbort()
			cancel <- struct{}{}
		}()
	}

	err = e.adapter.Init(e.config)
	if err != nil {
		log.Print(fmt.Sprintf("[adapter] %s", err.Error()))
		e.clientLog.Writeln(fmt.Sprintf("==> ERROR: Failed to initialize %s adapter", selectedAdapter))
		return RESULT_FAILED, err
	}

	err = e.adapter.Prepare(e.clientLog)
	defer e.adapter.Shutdown(e.clientLog)
	if err != nil {
		log.Print(fmt.Sprintf("[adapter] %s", err.Error()))
		e.clientLog.Writeln(fmt.Sprintf("==> ERROR: %s adapter failed to prepare: %s", selectedAdapter, err))
		return RESULT_FAILED, err
	}

	// actually begin executing the build plan
	finished := make(chan struct{})
	go func() {
		result = e.executeCommands()
		finished <- struct{}{}
	}()

	select {
	case <-finished:
	case <-cancel:
		e.clientLog.Writeln("==> ERROR: Build was aborted by upstream")
		return RESULT_ABORTED, nil
	}

	if result == RESULT_PASSED && outputSnapshot != "" {
		err = e.captureSnapshot()
		if err != nil {
			e.reporter.PushSnapshotImageStatus(outputSnapshot, SNAPSHOT_FAILED)
		} else {
			e.reporter.PushSnapshotImageStatus(outputSnapshot, SNAPSHOT_ACTIVE)
		}
	}

	return result, nil
}

func reportLogChunks(name string, clientLog *client.Log, r reporter.Reporter) {
	for chunk := range clientLog.Chan {
		r.PushLogChunk(name, chunk)
	}
}

func init() {
	flag.StringVar(&selectedAdapter, "adapter", "basic", "Adapter to run build against")
	flag.StringVar(&selectedReporter, "reporter", "multireporter", "Reporter to send results to")
	flag.StringVar(&outputSnapshot, "save-snapshot", "", "Save the resulting container snapshot")
}
<commit_msg>Make failures not resulting for a failing exit code from a command infra failures.<commit_after>package engine

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"

	_ "github.com/dropbox/changes-client/adapter/basic"
	_ "github.com/dropbox/changes-client/adapter/lxc"
	_ "github.com/dropbox/changes-client/reporter/artifactstore"
	_ "github.com/dropbox/changes-client/reporter/jenkins"
	_ "github.com/dropbox/changes-client/reporter/mesos"
	_ "github.com/dropbox/changes-client/reporter/multireporter"
)

const (
	STATUS_QUEUED      = "queued"
	STATUS_IN_PROGRESS = "in_progress"
	STATUS_FINISHED    = "finished"

	RESULT_PASSED  Result = "passed"
	RESULT_FAILED  Result = "failed"
	RESULT_ABORTED Result = "aborted"
	// Test results unreliable or unavailable due to infrastructure
	// issues.
	RESULT_INFRA_FAILED Result = "infra_failed"

	SNAPSHOT_ACTIVE = "active"
	SNAPSHOT_FAILED = "failed"
)

type Result string

func (r Result) String() string {
	return string(r)
}

// Convenience method to check for all types of failure.
func (r Result) IsFailure() bool {
	switch r {
	case RESULT_FAILED, RESULT_INFRA_FAILED:
		return true
	}
	return false
}

var (
	selectedAdapter  string
	selectedReporter string
	outputSnapshot   string
)

type Engine struct {
	config    *client.Config
	clientLog *client.Log
	adapter   adapter.Adapter
	reporter  reporter.Reporter
}

func RunBuildPlan(config *client.Config) (Result, error) {
	var err error

	currentAdapter, err := adapter.Get(selectedAdapter)
	if err != nil {
		return RESULT_INFRA_FAILED, err
	}

	currentReporter, err := reporter.Get(selectedReporter)
	if err != nil {
		log.Printf("[engine] failed to initialize reporter: %s", selectedReporter)
		return RESULT_INFRA_FAILED, err
	}
	log.Printf("[engine] started with reporter %s, adapter %s", selectedReporter, selectedAdapter)

	engine := &Engine{
		config:    config,
		clientLog: client.NewLog(),
		adapter:   currentAdapter,
		reporter:  currentReporter,
	}

	return engine.Run()
}

func (e *Engine) Run() (Result, error) {
	defer e.reporter.Shutdown()

	wg := sync.WaitGroup{}

	wg.Add(1)
	go func() {
		reportLogChunks("console", e.clientLog, e.reporter)
		wg.Done()
	}()

	e.reporter.Init(e.config)

	e.reporter.PushJobstepStatus(STATUS_IN_PROGRESS, "")

	result, err := e.runBuildPlan()

	e.clientLog.Writeln(fmt.Sprintf("==> Build finished! Recorded result as %s", result))

	e.reporter.PushJobstepStatus(STATUS_FINISHED, result.String())

	e.clientLog.Close()

	wg.Wait()

	return result, err
}

func (e *Engine) executeCommands() Result {
	result := RESULT_PASSED

	wg := sync.WaitGroup{}

	for _, cmdConfig := range e.config.Cmds {
		e.clientLog.Writeln(fmt.Sprintf("==> Running command %s", cmdConfig.ID))
		e.clientLog.Writeln(fmt.Sprintf("==>     with script %s", cmdConfig.Script))
		cmd, err := client.NewCommand(cmdConfig.ID, cmdConfig.Script)
		if err != nil {
			e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 255)
			result = RESULT_INFRA_FAILED
			e.clientLog.Writeln(fmt.Sprintf("==> Error: %s", err.Error()))
			break
		}
		e.reporter.PushCommandStatus(cmd.ID, STATUS_IN_PROGRESS, -1)

		cmd.CaptureOutput = cmdConfig.CaptureOutput

		env := os.Environ()
		for k, v := range cmdConfig.Env {
			env = append(env, k+"="+v)
		}
		cmd.Env = env

		if len(cmdConfig.Cwd) > 0 {
			cmd.Cwd = cmdConfig.Cwd
		}

		cmdResult, err := e.adapter.Run(cmd, e.clientLog)

		if err != nil {
			e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 255)
			e.clientLog.Writeln(fmt.Sprintf("==> Error: %s", err.Error()))
			result = RESULT_INFRA_FAILED
		} else {
			if cmdResult.Success {
				if cmd.CaptureOutput {
					e.reporter.PushCommandOutput(cmd.ID, STATUS_FINISHED, 0, cmdResult.Output)
				} else {
					e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 0)
				}
			} else {
				e.reporter.PushCommandStatus(cmd.ID, STATUS_FINISHED, 1)
				result = RESULT_FAILED
			}
		}

		wg.Add(1)
		go func() {
			// publishArtifacts is a synchronous operation and doesnt follow the normal queue flow of
			// other operations
			e.reporter.PublishArtifacts(cmdConfig, e.adapter, e.clientLog)
			wg.Done()
		}()

		if result.IsFailure() {
			break
		}
	}

	wg.Wait()

	return result
}

func (e *Engine) captureSnapshot() error {
	log.Printf("[adapter] Capturing snapshot %s", outputSnapshot)
	err := e.adapter.CaptureSnapshot(outputSnapshot, e.clientLog)
	if err != nil {
		log.Printf("[adapter] Failed to capture snapshot: %s", err.Error())
		return err
	}
	return nil
}

func (e *Engine) runBuildPlan() (Result, error) {
	var err error

	// cancellation signal
	cancel := make(chan struct{})

	// capture ctrl+c and enforce a clean shutdown
	sigchan := make(chan os.Signal, 1)
	signal.Notify(sigchan, os.Interrupt)
	go func() {
		shuttingDown := false
		for _ = range sigchan {
			if shuttingDown {
				log.Printf("Second interrupt received. Terminating!")
				os.Exit(1)
			}

			shuttingDown = true

			log.Printf("Interrupted! Cancelling execution and cleaning up..")
			cancel <- struct{}{}
		}
	}()

	// We need to ensure that we're able to abort the build if upstream suggests
	// that it's been cancelled.
	if !e.config.Debug {
		go func() {
			um := &UpstreamMonitor{
				Config: e.config,
			}
			um.WaitUntilAbort()
			cancel <- struct{}{}
		}()
	}

	err = e.adapter.Init(e.config)
	if err != nil {
		log.Print(fmt.Sprintf("[adapter] %s", err.Error()))
		e.clientLog.Writeln(fmt.Sprintf("==> ERROR: Failed to initialize %s adapter", selectedAdapter))
		return RESULT_INFRA_FAILED, err
	}

	err = e.adapter.Prepare(e.clientLog)
	defer e.adapter.Shutdown(e.clientLog)
	if err != nil {
		log.Print(fmt.Sprintf("[adapter] %s", err.Error()))
		e.clientLog.Writeln(fmt.Sprintf("==> ERROR: %s adapter failed to prepare: %s", selectedAdapter, err))
		return RESULT_INFRA_FAILED, err
	}

	var result Result
	// actually begin executing the build plan
	finished := make(chan struct{})
	go func() {
		result = e.executeCommands()
		finished <- struct{}{}
	}()

	select {
	case <-finished:
	case <-cancel:
		e.clientLog.Writeln("==> ERROR: Build was aborted by upstream")
		return RESULT_ABORTED, nil
	}

	if result == RESULT_PASSED && outputSnapshot != "" {
		err = e.captureSnapshot()
		if err != nil {
			e.reporter.PushSnapshotImageStatus(outputSnapshot, SNAPSHOT_FAILED)
		} else {
			e.reporter.PushSnapshotImageStatus(outputSnapshot, SNAPSHOT_ACTIVE)
		}
	}

	return result, nil
}

func reportLogChunks(name string, clientLog *client.Log, r reporter.Reporter) {
	for chunk := range clientLog.Chan {
		r.PushLogChunk(name, chunk)
	}
}

func init() {
	flag.StringVar(&selectedAdapter, "adapter", "basic", "Adapter to run build against")
	flag.StringVar(&selectedReporter, "reporter", "multireporter", "Reporter to send results to")
	flag.StringVar(&outputSnapshot, "save-snapshot", "", "Save the resulting container snapshot")
}
<|endoftext|>


<commit_before>package engine

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net"

	"github.com/imdario/mergo"
	"github.com/influx6/relay"
	"gopkg.in/yaml.v2"
)

// DefaultConfig provides a default configuration for the app
var DefaultConfig = Config{
	Addr:   ":8080",
	UseTLS: false,
}

//TLSConfig provides a base config for tls configuration
type TLSConfig struct {
	Cert *tls.Config
}

type tlsconf struct {
	Key  string
	Cert string
}

// UnmarshalYaml unmarshalls the incoming data for use
func (t *TLSConfig) UnmarshalYaml(unmarshal func(interface{}) error) error {
	to := tlsconf{}

	if err := unmarshal(&to); err != nil {
		return err
	}

	co, err := relay.LoadTLS(to.Cert, to.Key)

	if err != nil {
		return err
	}

	t.Cert = co
	return nil
}

// Config provides configuration for Afro
type Config struct {
	Addr   string
	UseTLS bool
	C      TLSConfig
}

// NewConfig returns a new configuration file
func NewConfig() *Config {
	c := DefaultConfig
	return &c
}

// Load loads the configuration from a yaml file
func (c *Config) Load(file string) error {
	data, err := ioutil.ReadFile(file)

	if err != nil {
		return err
	}

	conf := Config{}
	err = yaml.Unmarshal(data, &conf)

	if err != nil {
		return err
	}

	return mergo.MergeWithOverwrite(c, conf)
}

// UnmarshalYAML unmarshals and sets the configuration options
func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {

	return nil
}

//Engine provides a base luncher for a service
type Engine struct {
	*relay.Routes
	li     net.Listener
	config *Config
}

//NewEngine returns a new app configuration
func NewEngine(c *Config) *Engine {
	return &Engine{
		Routes: relay.NewRoutes(),
		config: c,
	}
}

// Serve serves the app and configuration and loads the routes and serivices settings
func (a *Engine) Serve() error {
	var err error
	var li net.Listener

	if a.config.UseTLS && a.config.C.Cert != nil {
		_, li, err = relay.CreateTLS(a.config.Addr, a.config.C.Cert, a)
	} else {
		_, li, err = relay.CreateHTTP(a.config.Addr, a)
	}

	if err != nil {
		log.Fatalf("Server failed to start: %+s", err.Error())
		return err
	}

	a.li = li

	return nil
}

// Addr returns the address of the app
func (a *Engine) Addr() net.Addr {
	return a.li.Addr()
}

// Close closes and returns an error of the internal listener
func (a *Engine) Close() error {
	return a.li.Close()
}
<commit_msg>updating engine<commit_after>package engine

import (
	"crypto/tls"
	"io/ioutil"
	"log"
	"net"

	"github.com/imdario/mergo"
	"github.com/influx6/relay"
	"gopkg.in/yaml.v2"
)

// DefaultConfig provides a default configuration for the app
var DefaultConfig = Config{
	Addr:   ":8080",
	UseTLS: false,
	Folders: Folders{
		Assets: "./assets",
		Models: "./models",
		Views:  "./views",
	},
}

//TLSConfig provides a base config for tls configuration
type TLSConfig struct {
	Cert *tls.Config
}

type tlsconf struct {
	Key  string
	Cert string
}

// UnmarshalYaml unmarshalls the incoming data for use
func (t *TLSConfig) UnmarshalYaml(unmarshal func(interface{}) error) error {
	to := tlsconf{}

	if err := unmarshal(&to); err != nil {
		return err
	}

	co, err := relay.LoadTLS(to.Cert, to.Key)

	if err != nil {
		return err
	}

	t.Cert = co
	return nil
}

// Folders provide a configuration for app-used folders
type Folders struct {
	Assets string
	Models string
	Views  string
}

// Config provides configuration for Afro
type Config struct {
	Addr    string
	UseTLS  bool
	C       TLSConfig
	Folders Folders
}

// NewConfig returns a new configuration file
func NewConfig() *Config {
	c := DefaultConfig
	return &c
}

// Load loads the configuration from a yaml file
func (c *Config) Load(file string) error {
	data, err := ioutil.ReadFile(file)

	if err != nil {
		return err
	}

	conf := Config{}
	err = yaml.Unmarshal(data, &conf)

	if err != nil {
		return err
	}

	return mergo.MergeWithOverwrite(c, conf)
}

// UnmarshalYAML unmarshals and sets the configuration options
func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {

	return nil
}

//Engine provides a base luncher for a service
type Engine struct {
	*relay.Routes
	li     net.Listener
	config *Config
}

//NewEngine returns a new app configuration
func NewEngine(c *Config) *Engine {
	return &Engine{
		Routes: relay.NewRoutes(),
		config: c,
	}
}

// Serve serves the app and configuration and loads the routes and serivices settings
func (a *Engine) Serve() error {
	var err error
	var li net.Listener

	if a.config.UseTLS && a.config.C.Cert != nil {
		_, li, err = relay.CreateTLS(a.config.Addr, a.config.C.Cert, a)
	} else {
		_, li, err = relay.CreateHTTP(a.config.Addr, a)
	}

	if err != nil {
		log.Fatalf("Server failed to start: %+s", err.Error())
		return err
	}

	a.li = li

	return nil
}

// Addr returns the address of the app
func (a *Engine) Addr() net.Addr {
	return a.li.Addr()
}

// Close closes and returns an error of the internal listener
func (a *Engine) Close() error {
	return a.li.Close()
}
<|endoftext|>


<commit_before>package search

import (
	"fmt"
	"regexp"
	"sort"
	"strings"
	"sync"

	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/collect"
	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/opentsdb"
)

// Search is a struct to hold indexed data about OpenTSDB metric and tag data.
// It is suited to answering questions about: available metrics for a tag set,
// available tag keys for a metric, and available tag values for a metric and
// tag key.
type Search struct {
	// tagk + tagv -> metrics
	Metric qmap
	// metric -> tag keys
	Tagk smap
	// metric + tagk -> tag values
	Tagv qmap
	// Each Record
	MetricTags mtsmap

	sync.RWMutex
	ch chan opentsdb.MultiDataPoint
}

type MetricTagSet struct {
	Metric string          `json:"metric"`
	Tags   opentsdb.TagSet `json:"tags"`
}

func (mts *MetricTagSet) key() string {
	s := make([]string, len(mts.Tags))
	i := 0
	for k, v := range mts.Tags {
		s[i] = fmt.Sprintf("%v=%v,", k, v)
		i++
	}
	sort.Strings(s)
	s = append(s, fmt.Sprintf("metric=%v", mts.Metric))
	return strings.Join(s, "")
}

type qmap map[duple]present
type smap map[string]present
type mtsmap map[string]MetricTagSet
type present map[string]struct{}

type duple struct {
	A, B string
}

func NewSearch() *Search {
	s := Search{
		Metric:     make(qmap),
		Tagk:       make(smap),
		Tagv:       make(qmap),
		MetricTags: make(mtsmap),
		ch:         make(chan opentsdb.MultiDataPoint, 1000),
	}
	go s.Process()
	collect.Set("search.ch_len", nil, func() interface{} {
		return len(s.ch)
	})
	return &s
}

func (s *Search) Process() {
	for mdp := range s.ch {
		s.Lock()
		for _, dp := range mdp {
			var mts MetricTagSet
			mts.Metric = dp.Metric
			mts.Tags = dp.Tags
			s.MetricTags[mts.key()] = mts
			var q duple
			for k, v := range dp.Tags {
				q.A, q.B = k, v
				if _, ok := s.Metric[q]; !ok {
					s.Metric[q] = make(present)
				}
				s.Metric[q][dp.Metric] = struct{}{}

				if _, ok := s.Tagk[dp.Metric]; !ok {
					s.Tagk[dp.Metric] = make(present)
				}
				s.Tagk[dp.Metric][k] = struct{}{}

				q.A, q.B = dp.Metric, k
				if _, ok := s.Tagv[q]; !ok {
					s.Tagv[q] = make(present)
				}
				s.Tagv[q][v] = struct{}{}
			}
		}
		s.Unlock()
	}
}

func (s *Search) Index(mdp opentsdb.MultiDataPoint) {
	select {
	case s.ch <- mdp:
	default:
		collect.Add("search.batch_drop", nil, 1)
	}
}

// Match returns all matching values against search. search is a regex, except
// that `.` is literal, `*` can be used for `.*`, and the entire string is
// searched (`^` and `&` added to ends of search).
func Match(search string, values []string) ([]string, error) {
	v := strings.Replace(search, ".", `\.`, -1)
	v = strings.Replace(v, "*", ".*", -1)
	v = "^" + v + "$"
	re, err := regexp.Compile(v)
	if err != nil {
		return nil, err
	}
	var nvs []string
	for _, nv := range values {
		if re.MatchString(nv) {
			nvs = append(nvs, nv)
		}
	}
	return nvs, nil
}

func (s *Search) Expand(q *opentsdb.Query) error {
	for k, ov := range q.Tags {
		v := ov
		if v == "*" || !strings.Contains(v, "*") || strings.Contains(v, "|") {
			continue
		}
		vs := s.TagValuesByMetricTagKey(q.Metric, k)
		nvs, err := Match(v, vs)
		if err != nil {
			return err
		}
		if len(nvs) == 0 {
			return fmt.Errorf("expr: no tags matching %s=%s", k, ov)
		}
		q.Tags[k] = strings.Join(nvs, "|")
	}
	return nil
}

func (s *Search) UniqueMetrics() []string {
	s.RLock()
	defer s.RUnlock()
	metrics := make([]string, len(s.Tagk))
	i := 0
	for k := range s.Tagk {
		metrics[i] = k
		i++
	}
	sort.Strings(metrics)
	return metrics
}

func (s *Search) TagValuesByTagKey(Tagk string) []string {
	um := s.UniqueMetrics()
	s.RLock()
	defer s.RUnlock()
	tagvset := make(map[string]bool)
	for _, Metric := range um {
		for _, Tagv := range s.tagValuesByMetricTagKey(Metric, Tagk) {
			tagvset[Tagv] = true
		}
	}
	tagvs := make([]string, len(tagvset))
	i := 0
	for k := range tagvset {
		tagvs[i] = k
		i++
	}
	sort.Strings(tagvs)
	return tagvs
}

func (s *Search) MetricsByTagPair(Tagk, Tagv string) []string {
	s.RLock()
	defer s.RUnlock()
	r := make([]string, 0)
	for k := range s.Metric[duple{Tagk, Tagv}] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) TagKeysByMetric(Metric string) []string {
	s.RLock()
	defer s.RUnlock()
	r := make([]string, 0)
	for k := range s.Tagk[Metric] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) tagValuesByMetricTagKey(Metric, Tagk string) []string {
	r := make([]string, 0)
	for k := range s.Tagv[duple{Metric, Tagk}] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) TagValuesByMetricTagKey(Metric, Tagk string) []string {
	s.RLock()
	defer s.RUnlock()
	return s.tagValuesByMetricTagKey(Metric, Tagk)
}

func (s *Search) FilteredTagValuesByMetricTagKey(Metric, Tagk string, tsf map[string]string) []string {
	s.RLock()
	defer s.RUnlock()
	tagvset := make(map[string]bool)
	for _, mts := range s.MetricTags {
		if Metric == mts.Metric {
			match := true
			if Tagv, ok := mts.Tags[Tagk]; ok {
				for tpk, tpv := range tsf {
					if v, ok := mts.Tags[tpk]; ok {
						if !(v == tpv) {
							match = false
						}
					} else {
						match = false
					}
				}
				if match {
					tagvset[Tagv] = true
				}
			}
		}
	}
	tagvs := make([]string, len(tagvset))
	i := 0
	for k := range tagvset {
		tagvs[i] = k
		i++
	}
	sort.Strings(tagvs)
	return tagvs
}
<commit_msg>Process search metrics synchronously<commit_after>package search

import (
	"fmt"
	"regexp"
	"sort"
	"strings"
	"sync"

	"github.com/StackExchange/bosun/_third_party/github.com/StackExchange/scollector/opentsdb"
)

// Search is a struct to hold indexed data about OpenTSDB metric and tag data.
// It is suited to answering questions about: available metrics for a tag set,
// available tag keys for a metric, and available tag values for a metric and
// tag key.
type Search struct {
	// tagk + tagv -> metrics
	Metric qmap
	// metric -> tag keys
	Tagk smap
	// metric + tagk -> tag values
	Tagv qmap
	// Each Record
	MetricTags mtsmap

	sync.RWMutex
}

type MetricTagSet struct {
	Metric string          `json:"metric"`
	Tags   opentsdb.TagSet `json:"tags"`
}

func (mts *MetricTagSet) key() string {
	s := make([]string, len(mts.Tags))
	i := 0
	for k, v := range mts.Tags {
		s[i] = fmt.Sprintf("%v=%v,", k, v)
		i++
	}
	sort.Strings(s)
	s = append(s, fmt.Sprintf("metric=%v", mts.Metric))
	return strings.Join(s, "")
}

type qmap map[duple]present
type smap map[string]present
type mtsmap map[string]MetricTagSet
type present map[string]struct{}

type duple struct {
	A, B string
}

func NewSearch() *Search {
	s := Search{
		Metric:     make(qmap),
		Tagk:       make(smap),
		Tagv:       make(qmap),
		MetricTags: make(mtsmap),
	}
	return &s
}

func (s *Search) Index(mdp opentsdb.MultiDataPoint) {
	s.Lock()
	for _, dp := range mdp {
		var mts MetricTagSet
		mts.Metric = dp.Metric
		mts.Tags = dp.Tags
		s.MetricTags[mts.key()] = mts
		var q duple
		for k, v := range dp.Tags {
			q.A, q.B = k, v
			if _, ok := s.Metric[q]; !ok {
				s.Metric[q] = make(present)
			}
			s.Metric[q][dp.Metric] = struct{}{}

			if _, ok := s.Tagk[dp.Metric]; !ok {
				s.Tagk[dp.Metric] = make(present)
			}
			s.Tagk[dp.Metric][k] = struct{}{}

			q.A, q.B = dp.Metric, k
			if _, ok := s.Tagv[q]; !ok {
				s.Tagv[q] = make(present)
			}
			s.Tagv[q][v] = struct{}{}
		}
	}
	s.Unlock()
}

// Match returns all matching values against search. search is a regex, except
// that `.` is literal, `*` can be used for `.*`, and the entire string is
// searched (`^` and `&` added to ends of search).
func Match(search string, values []string) ([]string, error) {
	v := strings.Replace(search, ".", `\.`, -1)
	v = strings.Replace(v, "*", ".*", -1)
	v = "^" + v + "$"
	re, err := regexp.Compile(v)
	if err != nil {
		return nil, err
	}
	var nvs []string
	for _, nv := range values {
		if re.MatchString(nv) {
			nvs = append(nvs, nv)
		}
	}
	return nvs, nil
}

func (s *Search) Expand(q *opentsdb.Query) error {
	for k, ov := range q.Tags {
		v := ov
		if v == "*" || !strings.Contains(v, "*") || strings.Contains(v, "|") {
			continue
		}
		vs := s.TagValuesByMetricTagKey(q.Metric, k)
		nvs, err := Match(v, vs)
		if err != nil {
			return err
		}
		if len(nvs) == 0 {
			return fmt.Errorf("expr: no tags matching %s=%s", k, ov)
		}
		q.Tags[k] = strings.Join(nvs, "|")
	}
	return nil
}

func (s *Search) UniqueMetrics() []string {
	s.RLock()
	defer s.RUnlock()
	metrics := make([]string, len(s.Tagk))
	i := 0
	for k := range s.Tagk {
		metrics[i] = k
		i++
	}
	sort.Strings(metrics)
	return metrics
}

func (s *Search) TagValuesByTagKey(Tagk string) []string {
	um := s.UniqueMetrics()
	s.RLock()
	defer s.RUnlock()
	tagvset := make(map[string]bool)
	for _, Metric := range um {
		for _, Tagv := range s.tagValuesByMetricTagKey(Metric, Tagk) {
			tagvset[Tagv] = true
		}
	}
	tagvs := make([]string, len(tagvset))
	i := 0
	for k := range tagvset {
		tagvs[i] = k
		i++
	}
	sort.Strings(tagvs)
	return tagvs
}

func (s *Search) MetricsByTagPair(Tagk, Tagv string) []string {
	s.RLock()
	defer s.RUnlock()
	r := make([]string, 0)
	for k := range s.Metric[duple{Tagk, Tagv}] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) TagKeysByMetric(Metric string) []string {
	s.RLock()
	defer s.RUnlock()
	r := make([]string, 0)
	for k := range s.Tagk[Metric] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) tagValuesByMetricTagKey(Metric, Tagk string) []string {
	r := make([]string, 0)
	for k := range s.Tagv[duple{Metric, Tagk}] {
		r = append(r, k)
	}
	sort.Strings(r)
	return r
}

func (s *Search) TagValuesByMetricTagKey(Metric, Tagk string) []string {
	s.RLock()
	defer s.RUnlock()
	return s.tagValuesByMetricTagKey(Metric, Tagk)
}

func (s *Search) FilteredTagValuesByMetricTagKey(Metric, Tagk string, tsf map[string]string) []string {
	s.RLock()
	defer s.RUnlock()
	tagvset := make(map[string]bool)
	for _, mts := range s.MetricTags {
		if Metric == mts.Metric {
			match := true
			if Tagv, ok := mts.Tags[Tagk]; ok {
				for tpk, tpv := range tsf {
					if v, ok := mts.Tags[tpk]; ok {
						if !(v == tpv) {
							match = false
						}
					} else {
						match = false
					}
				}
				if match {
					tagvset[Tagv] = true
				}
			}
		}
	}
	tagvs := make([]string, len(tagvset))
	i := 0
	for k := range tagvset {
		tagvs[i] = k
		i++
	}
	sort.Strings(tagvs)
	return tagvs
}
<|endoftext|>


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

package search

import (
	"bufio"
	"os"
	"sort"
	"strings"
)

type Store struct {
	index map[string][]string
}

func NewStore(files ...string) (*Store, error) {
	s := Store{
		index: make(map[string][]string),
	}
	sort.Strings(files)
	for _, name := range files {
		f, err := os.Open(name)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		scanner := bufio.NewScanner(f)
		scanner.Split(bufio.ScanWords)
		for scanner.Scan() {
			t := scanner.Text()
			t = strings.Trim(t, ".?!-,;:")
			l := s.index[t]
			if n := sort.SearchStrings(l, name); n == len(l) {
				l = append(l, name)
				s.index[t] = l
			}
		}
	}
	return &s, nil
}
<commit_msg>search: also ignore "<commit_after>// Copyright 2013 Francisco Souza. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package search

import (
	"bufio"
	"os"
	"sort"
	"strings"
)

type Store struct {
	index map[string][]string
}

func NewStore(files ...string) (*Store, error) {
	s := Store{
		index: make(map[string][]string),
	}
	sort.Strings(files)
	for _, name := range files {
		f, err := os.Open(name)
		if err != nil {
			return nil, err
		}
		defer f.Close()
		scanner := bufio.NewScanner(f)
		scanner.Split(bufio.ScanWords)
		for scanner.Scan() {
			t := scanner.Text()
			t = strings.Trim(t, `.?!-,;:"`)
			l := s.index[t]
			if n := sort.SearchStrings(l, name); n == len(l) {
				l = append(l, name)
				s.index[t] = l
			}
		}
	}
	return &s, nil
}
<|endoftext|>


<commit_before>package tcpoverdns

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"compress/lzw"
	"compress/zlib"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"reflect"
	"testing"
	"time"
)

func TestSegment_Packet(t *testing.T) {
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   []byte{1, 2, 3, 4},
	}

	packet := want.Packet()
	got := SegmentFromPacket(packet)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("recovered: %+#v original: %+#v", got, want)
	}

	want.Flags = FlagHandshakeSyn
	packet = want.Packet()
	got = SegmentFromPacket(packet)
	if !reflect.DeepEqual(got, Segment{Flags: FlagMalformed}) {
		t.Fatal("did not identify malformed segment without initiator config")
	}
}

func TestSegmentFromMalformedPacket(t *testing.T) {
	want := Segment{Flags: FlagMalformed}
	segWithData := Segment{Data: []byte{1, 2}}
	segWithMalformedLen := segWithData.Packet()
	for _, seg := range [][]byte{nil, {1}, segWithMalformedLen[:SegmentHeaderLen+1]} {
		if got := SegmentFromPacket(seg); !reflect.DeepEqual(got, want) {
			t.Fatalf("got: %+#v, want: %+#v", got, want)
		}
	}
}

func TestFlags(t *testing.T) {
	allFlags := FlagHandshakeSyn | FlagHandshakeAck | FlagAckOnly | FlagKeepAlive | FlagReset | FlagMalformed
	for _, flag := range []Flag{FlagHandshakeSyn, FlagHandshakeAck, FlagAckOnly, FlagKeepAlive, FlagReset, FlagMalformed} {
		if !allFlags.Has(flag) {
			t.Fatalf("missing %d", flag)
		}
	}
	if allFlags.Has(1 << 6) {
		t.Fatalf("should not have had flag %d", 1<<4)
	}
}

func TestSegment_Equals(t *testing.T) {
	original := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   []byte{1, 2, 3, 4},
	}
	if !original.Equals(original) {
		t.Errorf("should have been equal")
	}

	tests := []struct {
		a, b Segment
	}{
		{a: Segment{ID: 1}, b: Segment{ID: 2}},
		{a: Segment{Flags: 1}, b: Segment{Flags: 2}},
		{a: Segment{SeqNum: 1}, b: Segment{SeqNum: 2}},
		{a: Segment{AckNum: 1}, b: Segment{AckNum: 2}},
		{a: Segment{Data: []byte{0}}, b: Segment{Data: []byte{1}}},
	}
	for _, test := range tests {
		if test.a.Equals(test.b) {
			t.Errorf("should not have been equal: %+v, %+v", test.a, test.b)
		}
	}
}

func TestCompression(t *testing.T) {
	// original := `The words of the Teacher, son of David, king in Jerusalem: “Meaningless! Meaningless!” says the Teacher. “Utterly meaningless! Everything is meaningless.” What do people gain from all their labors at which they toil under the sun?`
	original := `<!doctype html><html itemscope="" itemtype="http://schema.org/WebPage" lang="en-IE"><head><meta charset="UTF-8"><meta content="dark" name="color-scheme"><meta content="origin" name="referrer"><meta content="/images/branding/googleg/1x/googleg_standard_colo`

	t.Run("zlib", func(t *testing.T) {
		var b bytes.Buffer
		w, err := zlib.NewWriterLevel(&b, zlib.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("zlib", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("lzw", func(t *testing.T) {
		var b bytes.Buffer
		w := lzw.NewWriter(&b, lzw.MSB, 8)
		w.Write([]byte(original))
		w.Close()
		fmt.Println("lzw", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("flate", func(t *testing.T) {
		var b bytes.Buffer
		w, err := flate.NewWriter(&b, flate.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("flate", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("gzip", func(t *testing.T) {
		var b bytes.Buffer
		w, err := gzip.NewWriterLevel(&b, gzip.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("gzip", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("base64", func(t *testing.T) {
		fmt.Println(base64.StdEncoding.DecodeString("TQ=="))
	})
}

func TestCompressDecompressBytes(t *testing.T) {
	tests := [][]byte{
		{},
		{0},
		{0, 1},
		{0, 1, 2},
		[]byte(`<!doctype html><html itemscope="" itemtype="http://schema.org/WebPage" lang="en-IE"><head><meta charset="UTF-8"><meta content="dark" name="color-scheme"><meta content="origin" name="referrer"><meta content="/images/branding/googleg/1x/googleg_standard_colo`),
	}
	for _, original := range tests {
		compressed := CompressBytes(original)
		got, err := DecompressBytes(compressed)
		if err != nil || !reflect.DeepEqual(got, original) {
			t.Fatalf("DecompressBytes(%+v): got %+v, want %+v", compressed, got, original)
		}
	}
}

func TestSegment_DNSQuestion(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}
	dnsQuestion := want.DNSQuestion("prefix-label", "example.com")
	fmt.Println(dnsQuestion.Name.String())
	got := SegmentFromDNSQuery(2, dnsQuestion.Name.String())
	if !reflect.DeepEqual(got, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
	// Try the same conversion, but without the domain name.
	qWithoutDomain := want.DNSQuestion("prefix-label", "")
	fmt.Println(qWithoutDomain.Name.String())
	gotWithoutDomain := SegmentFromDNSQuery(0, qWithoutDomain.Name.String())
	if !reflect.DeepEqual(gotWithoutDomain, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
}

func TestInitiatorConfig(t *testing.T) {
	want := &InitiatorConfig{
		SetConfig:               true,
		MaxSegmentLenExclHeader: 123,
		IOTimeoutSec:            234,
		KeepAliveIntervalSec:    345,
		Debug:                   true,
	}
	serialised := want.Bytes()
	got := DeserialiseInitiatorConfig(serialised)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("got: %+#v want: %+#v", got, want)
	}

	wantTC := &TransmissionControl{
		MaxSegmentLenExclHeader: 123,
		MaxSlidingWindow:        123 * 8,
		ReadTimeout:             234 * time.Second,
		WriteTimeout:            234 * time.Second,
		KeepAliveInterval:       345 * time.Second,
		Debug:                   true,
	}
	gotTC := &TransmissionControl{}
	got.Config(gotTC)
	if !reflect.DeepEqual(gotTC, wantTC) {
		t.Fatalf("got: %+#v want: %+#v", gotTC, wantTC)
	}
}

func TestSegment_DNSNameQuery(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}
	query := want.DNSNameQuery("prefix-label", "example.com")
	fmt.Println(query)
	got := SegmentFromDNSQuery(2, query)
	if !reflect.DeepEqual(got, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
}

func TestSegment_DNSResource(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}

	addrResource := want.DNSResource()
	addrs := make([]net.IP, 0)
	for _, addr := range addrResource {
		ip := make([]byte, 4)
		copy(ip[:], addr.A[:])
		addrs = append(addrs, net.IP(ip))
	}
	got := SegmentFromIPs(addrs)
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}
}
<commit_msg>tcpoverdns: fix failed InitiatorConfig test<commit_after>package tcpoverdns

import (
	"bytes"
	"compress/flate"
	"compress/gzip"
	"compress/lzw"
	"compress/zlib"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net"
	"reflect"
	"testing"
	"time"
)

func TestSegment_Packet(t *testing.T) {
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   []byte{1, 2, 3, 4},
	}

	packet := want.Packet()
	got := SegmentFromPacket(packet)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("recovered: %+#v original: %+#v", got, want)
	}

	want.Flags = FlagHandshakeSyn
	packet = want.Packet()
	got = SegmentFromPacket(packet)
	if !reflect.DeepEqual(got, Segment{Flags: FlagMalformed}) {
		t.Fatal("did not identify malformed segment without initiator config")
	}
}

func TestSegmentFromMalformedPacket(t *testing.T) {
	want := Segment{Flags: FlagMalformed}
	segWithData := Segment{Data: []byte{1, 2}}
	segWithMalformedLen := segWithData.Packet()
	for _, seg := range [][]byte{nil, {1}, segWithMalformedLen[:SegmentHeaderLen+1]} {
		if got := SegmentFromPacket(seg); !reflect.DeepEqual(got, want) {
			t.Fatalf("got: %+#v, want: %+#v", got, want)
		}
	}
}

func TestFlags(t *testing.T) {
	allFlags := FlagHandshakeSyn | FlagHandshakeAck | FlagAckOnly | FlagKeepAlive | FlagReset | FlagMalformed
	for _, flag := range []Flag{FlagHandshakeSyn, FlagHandshakeAck, FlagAckOnly, FlagKeepAlive, FlagReset, FlagMalformed} {
		if !allFlags.Has(flag) {
			t.Fatalf("missing %d", flag)
		}
	}
	if allFlags.Has(1 << 6) {
		t.Fatalf("should not have had flag %d", 1<<4)
	}
}

func TestSegment_Equals(t *testing.T) {
	original := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   []byte{1, 2, 3, 4},
	}
	if !original.Equals(original) {
		t.Errorf("should have been equal")
	}

	tests := []struct {
		a, b Segment
	}{
		{a: Segment{ID: 1}, b: Segment{ID: 2}},
		{a: Segment{Flags: 1}, b: Segment{Flags: 2}},
		{a: Segment{SeqNum: 1}, b: Segment{SeqNum: 2}},
		{a: Segment{AckNum: 1}, b: Segment{AckNum: 2}},
		{a: Segment{Data: []byte{0}}, b: Segment{Data: []byte{1}}},
	}
	for _, test := range tests {
		if test.a.Equals(test.b) {
			t.Errorf("should not have been equal: %+v, %+v", test.a, test.b)
		}
	}
}

func TestCompression(t *testing.T) {
	// original := `The words of the Teacher, son of David, king in Jerusalem: “Meaningless! Meaningless!” says the Teacher. “Utterly meaningless! Everything is meaningless.” What do people gain from all their labors at which they toil under the sun?`
	original := `<!doctype html><html itemscope="" itemtype="http://schema.org/WebPage" lang="en-IE"><head><meta charset="UTF-8"><meta content="dark" name="color-scheme"><meta content="origin" name="referrer"><meta content="/images/branding/googleg/1x/googleg_standard_colo`

	t.Run("zlib", func(t *testing.T) {
		var b bytes.Buffer
		w, err := zlib.NewWriterLevel(&b, zlib.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("zlib", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("lzw", func(t *testing.T) {
		var b bytes.Buffer
		w := lzw.NewWriter(&b, lzw.MSB, 8)
		w.Write([]byte(original))
		w.Close()
		fmt.Println("lzw", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("flate", func(t *testing.T) {
		var b bytes.Buffer
		w, err := flate.NewWriter(&b, flate.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("flate", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("gzip", func(t *testing.T) {
		var b bytes.Buffer
		w, err := gzip.NewWriterLevel(&b, gzip.BestCompression)
		if err != nil {
			t.Fatal(err)
		}
		w.Write([]byte(original))
		w.Close()
		fmt.Println("gzip", len(b.Bytes()), len(base64.StdEncoding.EncodeToString(b.Bytes())))
	})

	t.Run("base64", func(t *testing.T) {
		fmt.Println(base64.StdEncoding.DecodeString("TQ=="))
	})
}

func TestCompressDecompressBytes(t *testing.T) {
	tests := [][]byte{
		{},
		{0},
		{0, 1},
		{0, 1, 2},
		[]byte(`<!doctype html><html itemscope="" itemtype="http://schema.org/WebPage" lang="en-IE"><head><meta charset="UTF-8"><meta content="dark" name="color-scheme"><meta content="origin" name="referrer"><meta content="/images/branding/googleg/1x/googleg_standard_colo`),
	}
	for _, original := range tests {
		compressed := CompressBytes(original)
		got, err := DecompressBytes(compressed)
		if err != nil || !reflect.DeepEqual(got, original) {
			t.Fatalf("DecompressBytes(%+v): got %+v, want %+v", compressed, got, original)
		}
	}
}

func TestSegment_DNSQuestion(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}
	dnsQuestion := want.DNSQuestion("prefix-label", "example.com")
	fmt.Println(dnsQuestion.Name.String())
	got := SegmentFromDNSQuery(2, dnsQuestion.Name.String())
	if !reflect.DeepEqual(got, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
	// Try the same conversion, but without the domain name.
	qWithoutDomain := want.DNSQuestion("prefix-label", "")
	fmt.Println(qWithoutDomain.Name.String())
	gotWithoutDomain := SegmentFromDNSQuery(0, qWithoutDomain.Name.String())
	if !reflect.DeepEqual(gotWithoutDomain, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
}

func TestInitiatorConfig(t *testing.T) {
	want := &InitiatorConfig{
		SetConfig:               true,
		MaxSegmentLenExclHeader: 123,
		IOTimeoutSec:            234,
		KeepAliveIntervalSec:    345,
		Debug:                   true,
	}
	serialised := want.Bytes()
	got := DeserialiseInitiatorConfig(serialised)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("got: %+#v want: %+#v", got, want)
	}

	wantTC := &TransmissionControl{
		MaxSegmentLenExclHeader: 123,
		MaxSlidingWindow:        123 * 4,
		ReadTimeout:             234 * time.Second,
		WriteTimeout:            234 * time.Second,
		KeepAliveInterval:       345 * time.Second,
		Debug:                   true,
	}
	gotTC := &TransmissionControl{}
	got.Config(gotTC)
	if !reflect.DeepEqual(gotTC, wantTC) {
		t.Fatalf("got: %+#v want: %+#v", gotTC, wantTC)
	}
}

func TestSegment_DNSNameQuery(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}
	query := want.DNSNameQuery("prefix-label", "example.com")
	fmt.Println(query)
	got := SegmentFromDNSQuery(2, query)
	if !reflect.DeepEqual(got, want) {
		t.Errorf("recovered: %+#v original: %+#v", got, want)
	}
}

func TestSegment_DNSResource(t *testing.T) {
	randData := make([]byte, 100)
	if _, err := rand.Read(randData); err != nil {
		t.Fatal(err)
	}
	want := Segment{
		ID:     12345,
		Flags:  FlagHandshakeAck & FlagHandshakeSyn,
		SeqNum: 23456,
		AckNum: 34567,
		Data:   randData,
	}

	addrResource := want.DNSResource()
	addrs := make([]net.IP, 0)
	for _, addr := range addrResource {
		ip := make([]byte, 4)
		copy(ip[:], addr.A[:])
		addrs = append(addrs, net.IP(ip))
	}
	got := SegmentFromIPs(addrs)
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got: %v, want: %v", got, want)
	}
}
<|endoftext|>