<commit_before>package mdata

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/grafana/grafana/pkg/log"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	"github.com/raintank/raintank-metric/app"
	"github.com/raintank/raintank-metric/instrumented_nsq"
)

var (
	hostPool            hostpool.HostPool
	producers           map[string]*nsq.Producer
	CluStatus           *ClusterStatus
	persistMessageBatch *PersistMessageBatch
)

//PersistMessage format version
const PersistMessageBatchV1 = 1

// ClusterStatus has Exported fields but don't touch them directly
// it's only for json marshaling. use the accessor methods.
type ClusterStatus struct {
	sync.Mutex
	Instance   string    `json:"instance"`
	Primary    bool      `json:"primary"`
	LastChange time.Time `json:"lastChange"`
}

func NewClusterStatus(instance string, initialState bool) *ClusterStatus {
	return &ClusterStatus{
		Instance:   instance,
		Primary:    initialState,
		LastChange: time.Now(),
	}
}

func (c *ClusterStatus) Marshal() ([]byte, error) {
	c.Lock()
	defer c.Unlock()
	return json.Marshal(c)
}

func (c *ClusterStatus) Set(newState bool) {
	c.Lock()
	c.Primary = newState
	c.LastChange = time.Now()
	c.Unlock()
}

func (c *ClusterStatus) IsPrimary() bool {
	c.Lock()
	defer c.Unlock()
	return c.Primary
}

type PersistMessage struct {
	Instance string `json:"instance"`
	Key      string `json:"key"`
	T0       uint32 `json:"t0"`
}

type PersistMessageBatch struct {
	sync.Mutex  `json:"-"`
	Topic       string       `json:"-"`
	Instance    string       `json:"instance"`
	SavedChunks []savedChunk `json:"saved_chunks"`
}

type savedChunk struct {
	Key string `json:"key"`
	T0  uint32 `json:"t0"`
}

func (p *PersistMessageBatch) AddChunk(key string, t0 uint32) {
	p.Lock()
	defer p.Unlock()
	p.SavedChunks = append(p.SavedChunks, savedChunk{Key: key, T0: t0})
}

func SendPersistMessage(key string, t0 uint32) {
	persistMessageBatch.AddChunk(key, t0)
}

func (p *PersistMessageBatch) flush() {
	ticker := time.NewTicker(time.Second)
	for range ticker.C {
		// get current savedChunks and clear our buffer.
		p.Lock()
		c := make([]savedChunk, len(p.SavedChunks))
		copy(c, p.SavedChunks)
		p.SavedChunks = nil
		msg := PersistMessageBatch{Instance: p.Instance, SavedChunks: c}
		p.Unlock()

		if len(c) == 0 {
			continue
		}

		if p.Topic == "" {
			continue
		}

		log.Debug("CLU sending %d batch metricPersist messages", len(c))

		data, err := json.Marshal(&msg)
		if err != nil {
			log.Fatal(4, "failed to marshal persistMessage to json.")
		}
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.LittleEndian, uint8(PersistMessageBatchV1))
		buf.Write(data)

		sent := false
		for !sent {
			// This will always return a host. If all hosts are currently marked as dead,
			// then all hosts will be reset to alive and we will try them all again. This
			// will result in this loop repeating forever until we successfully publish our msg.
			hostPoolResponse := hostPool.Get()
			prod := producers[hostPoolResponse.Host()]
			err = prod.Publish(p.Topic, buf.Bytes())
			// Hosts that are marked as dead will be retried after 30seconds.  If we published
			// successfully, then sending a nil error will mark the host as alive again.
			hostPoolResponse.Mark(err)
			if err != nil {
				log.Warn("publisher marking host %s as faulty due to %s", hostPoolResponse.Host(), err)
			} else {
				sent = true
			}
		}
	}
}

type MetricPersistHandler struct {
	instance string
	metrics  Metrics
}

func NewMetricPersistHandler(instance string, metrics Metrics) *MetricPersistHandler {
	return &MetricPersistHandler{
		instance: instance,
		metrics:  metrics,
	}
}

func (h *MetricPersistHandler) HandleMessage(m *nsq.Message) error {
	version := uint8(m.Body[0])
	if version == uint8(PersistMessageBatchV1) {
		// new batch format.
		batch := PersistMessageBatch{}
		err := json.Unmarshal(m.Body[1:], &batch)
		if err != nil {
			log.Error(3, "failed to unmarsh batch message. skipping.", err)
			return nil
		}
		if batch.Instance == h.instance {
			log.Debug("CLU skipping batch message we generated.")
			return nil
		}
		for _, c := range batch.SavedChunks {
			if agg, ok := h.metrics.Get(c.Key); ok {
				agg.(*AggMetric).SyncChunkSaveState(c.T0)
			}
		}
	} else {
		// assume the old format.
		ms := PersistMessage{}
		err := json.Unmarshal(m.Body, &ms)
		if err != nil {
			log.Error(3, "skipping message. %s", err)
			return nil
		}
		if ms.Instance == h.instance {
			log.Debug("CLU skipping message we generated. %s - %s:%d", ms.Instance, ms.Key, ms.T0)
			return nil
		}

		// get metric
		if agg, ok := h.metrics.Get(ms.Key); ok {
			agg.(*AggMetric).SyncChunkSaveState(ms.T0)
		}
	}
	return nil
}

func InitCluster(metrics Metrics, stats met.Backend, instance, topic, channel, producerOpts, consumerOpts string, nsqdAdds, lookupdAdds []string, maxInFlight int) {
	persistMessageBatch = &PersistMessageBatch{Instance: instance, SavedChunks: make([]savedChunk, 0), Topic: topic}
	// init producers
	pCfg := nsq.NewConfig()
	pCfg.UserAgent = "metrics_tank"
	err := app.ParseOpts(pCfg, producerOpts)
	if err != nil {
		log.Fatal(4, "failed to parse nsq producer options. %s", err)
	}
	hostPool = hostpool.NewEpsilonGreedy(nsqdAdds, 0, &hostpool.LinearEpsilonValueCalculator{})
	producers = make(map[string]*nsq.Producer)

	for _, addr := range nsqdAdds {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatal(4, "failed creating producer %s", err.Error())
		}
		producers[addr] = producer
	}

	// init consumers
	cfg := nsq.NewConfig()
	cfg.UserAgent = "metrics_tank_cluster"
	err = app.ParseOpts(cfg, consumerOpts)
	if err != nil {
		log.Fatal(4, "failed to parse nsq consumer options. %s", err)
	}
	cfg.MaxInFlight = maxInFlight

	consumer, err := insq.NewConsumer(topic, channel, cfg, "metric_persist.%s", stats)
	if err != nil {
		log.Fatal(4, "Failed to create NSQ consumer. %s", err)
	}
	handler := NewMetricPersistHandler(instance, metrics)
	consumer.AddConcurrentHandlers(handler, 2)

	err = consumer.ConnectToNSQDs(nsqdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQDs. %s", err)
	}
	log.Info("persist consumer connected to nsqd")

	err = consumer.ConnectToNSQLookupds(lookupdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQLookupds. %s", err)
	}
	go persistMessageBatch.flush()
}

// Handle requests for /cluster. POST to set primary flag, GET to get current state.
func (c *ClusterStatus) HttpHandler(w http.ResponseWriter, req *http.Request) {
	if req.Method == "GET" {
		c.getClusterStatus(w, req)
		return
	}
	if req.Method == "POST" {
		c.setClusterStatus(w, req)
		return
	}
	http.Error(w, "not found.", http.StatusNotFound)
}

func (c *ClusterStatus) setClusterStatus(w http.ResponseWriter, req *http.Request) {
	req.ParseForm()
	newState := req.Form.Get("primary")
	if newState == "" {
		http.Error(w, "primary field missing from payload.", http.StatusBadRequest)
		return
	}

	primary, err := strconv.ParseBool(newState)
	if err != nil {
		http.Error(w, "primary field could not be parsed to boolean value.", http.StatusBadRequest)
		return
	}
	c.Set(primary)
	log.Info("primary status is now %t", primary)
	w.Write([]byte("OK"))
}

func (c *ClusterStatus) getClusterStatus(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	resp, err := c.Marshal()
	if err != nil {
		http.Error(w, "could not marshal status to json", http.StatusInternalServerError)
		return
	}
	w.Write(resp)
}
<commit_msg>no need for extra allocation & copying<commit_after>package mdata

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/grafana/grafana/pkg/log"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	"github.com/raintank/raintank-metric/app"
	"github.com/raintank/raintank-metric/instrumented_nsq"
)

var (
	hostPool            hostpool.HostPool
	producers           map[string]*nsq.Producer
	CluStatus           *ClusterStatus
	persistMessageBatch *PersistMessageBatch
)

//PersistMessage format version
const PersistMessageBatchV1 = 1

// ClusterStatus has Exported fields but don't touch them directly
// it's only for json marshaling. use the accessor methods.
type ClusterStatus struct {
	sync.Mutex
	Instance   string    `json:"instance"`
	Primary    bool      `json:"primary"`
	LastChange time.Time `json:"lastChange"`
}

func NewClusterStatus(instance string, initialState bool) *ClusterStatus {
	return &ClusterStatus{
		Instance:   instance,
		Primary:    initialState,
		LastChange: time.Now(),
	}
}

func (c *ClusterStatus) Marshal() ([]byte, error) {
	c.Lock()
	defer c.Unlock()
	return json.Marshal(c)
}

func (c *ClusterStatus) Set(newState bool) {
	c.Lock()
	c.Primary = newState
	c.LastChange = time.Now()
	c.Unlock()
}

func (c *ClusterStatus) IsPrimary() bool {
	c.Lock()
	defer c.Unlock()
	return c.Primary
}

type PersistMessage struct {
	Instance string `json:"instance"`
	Key      string `json:"key"`
	T0       uint32 `json:"t0"`
}

type PersistMessageBatch struct {
	sync.Mutex  `json:"-"`
	Topic       string       `json:"-"`
	Instance    string       `json:"instance"`
	SavedChunks []savedChunk `json:"saved_chunks"`
}

type savedChunk struct {
	Key string `json:"key"`
	T0  uint32 `json:"t0"`
}

func (p *PersistMessageBatch) AddChunk(key string, t0 uint32) {
	p.Lock()
	defer p.Unlock()
	p.SavedChunks = append(p.SavedChunks, savedChunk{Key: key, T0: t0})
}

func SendPersistMessage(key string, t0 uint32) {
	persistMessageBatch.AddChunk(key, t0)
}

func (p *PersistMessageBatch) flush() {
	ticker := time.NewTicker(time.Second)
	for range ticker.C {
		p.Lock()
		msg := PersistMessageBatch{Instance: p.Instance, SavedChunks: p.SavedChunks}
		p.SavedChunks = nil
		p.Unlock()

		if len(msg.SavedChunks) == 0 {
			continue
		}

		if p.Topic == "" {
			continue
		}

		log.Debug("CLU sending %d batch metricPersist messages", len(msg.SavedChunks))

		data, err := json.Marshal(&msg)
		if err != nil {
			log.Fatal(4, "failed to marshal persistMessage to json.")
		}
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.LittleEndian, uint8(PersistMessageBatchV1))
		buf.Write(data)

		sent := false
		for !sent {
			// This will always return a host. If all hosts are currently marked as dead,
			// then all hosts will be reset to alive and we will try them all again. This
			// will result in this loop repeating forever until we successfully publish our msg.
			hostPoolResponse := hostPool.Get()
			prod := producers[hostPoolResponse.Host()]
			err = prod.Publish(p.Topic, buf.Bytes())
			// Hosts that are marked as dead will be retried after 30seconds.  If we published
			// successfully, then sending a nil error will mark the host as alive again.
			hostPoolResponse.Mark(err)
			if err != nil {
				log.Warn("publisher marking host %s as faulty due to %s", hostPoolResponse.Host(), err)
			} else {
				sent = true
			}
		}
	}
}

type MetricPersistHandler struct {
	instance string
	metrics  Metrics
}

func NewMetricPersistHandler(instance string, metrics Metrics) *MetricPersistHandler {
	return &MetricPersistHandler{
		instance: instance,
		metrics:  metrics,
	}
}

func (h *MetricPersistHandler) HandleMessage(m *nsq.Message) error {
	version := uint8(m.Body[0])
	if version == uint8(PersistMessageBatchV1) {
		// new batch format.
		batch := PersistMessageBatch{}
		err := json.Unmarshal(m.Body[1:], &batch)
		if err != nil {
			log.Error(3, "failed to unmarsh batch message. skipping.", err)
			return nil
		}
		if batch.Instance == h.instance {
			log.Debug("CLU skipping batch message we generated.")
			return nil
		}
		for _, c := range batch.SavedChunks {
			if agg, ok := h.metrics.Get(c.Key); ok {
				agg.(*AggMetric).SyncChunkSaveState(c.T0)
			}
		}
	} else {
		// assume the old format.
		ms := PersistMessage{}
		err := json.Unmarshal(m.Body, &ms)
		if err != nil {
			log.Error(3, "skipping message. %s", err)
			return nil
		}
		if ms.Instance == h.instance {
			log.Debug("CLU skipping message we generated. %s - %s:%d", ms.Instance, ms.Key, ms.T0)
			return nil
		}

		// get metric
		if agg, ok := h.metrics.Get(ms.Key); ok {
			agg.(*AggMetric).SyncChunkSaveState(ms.T0)
		}
	}
	return nil
}

func InitCluster(metrics Metrics, stats met.Backend, instance, topic, channel, producerOpts, consumerOpts string, nsqdAdds, lookupdAdds []string, maxInFlight int) {
	persistMessageBatch = &PersistMessageBatch{Instance: instance, SavedChunks: make([]savedChunk, 0), Topic: topic}
	// init producers
	pCfg := nsq.NewConfig()
	pCfg.UserAgent = "metrics_tank"
	err := app.ParseOpts(pCfg, producerOpts)
	if err != nil {
		log.Fatal(4, "failed to parse nsq producer options. %s", err)
	}
	hostPool = hostpool.NewEpsilonGreedy(nsqdAdds, 0, &hostpool.LinearEpsilonValueCalculator{})
	producers = make(map[string]*nsq.Producer)

	for _, addr := range nsqdAdds {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatal(4, "failed creating producer %s", err.Error())
		}
		producers[addr] = producer
	}

	// init consumers
	cfg := nsq.NewConfig()
	cfg.UserAgent = "metrics_tank_cluster"
	err = app.ParseOpts(cfg, consumerOpts)
	if err != nil {
		log.Fatal(4, "failed to parse nsq consumer options. %s", err)
	}
	cfg.MaxInFlight = maxInFlight

	consumer, err := insq.NewConsumer(topic, channel, cfg, "metric_persist.%s", stats)
	if err != nil {
		log.Fatal(4, "Failed to create NSQ consumer. %s", err)
	}
	handler := NewMetricPersistHandler(instance, metrics)
	consumer.AddConcurrentHandlers(handler, 2)

	err = consumer.ConnectToNSQDs(nsqdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQDs. %s", err)
	}
	log.Info("persist consumer connected to nsqd")

	err = consumer.ConnectToNSQLookupds(lookupdAdds)
	if err != nil {
		log.Fatal(4, "failed to connect to NSQLookupds. %s", err)
	}
	go persistMessageBatch.flush()
}

// Handle requests for /cluster. POST to set primary flag, GET to get current state.
func (c *ClusterStatus) HttpHandler(w http.ResponseWriter, req *http.Request) {
	if req.Method == "GET" {
		c.getClusterStatus(w, req)
		return
	}
	if req.Method == "POST" {
		c.setClusterStatus(w, req)
		return
	}
	http.Error(w, "not found.", http.StatusNotFound)
}

func (c *ClusterStatus) setClusterStatus(w http.ResponseWriter, req *http.Request) {
	req.ParseForm()
	newState := req.Form.Get("primary")
	if newState == "" {
		http.Error(w, "primary field missing from payload.", http.StatusBadRequest)
		return
	}

	primary, err := strconv.ParseBool(newState)
	if err != nil {
		http.Error(w, "primary field could not be parsed to boolean value.", http.StatusBadRequest)
		return
	}
	c.Set(primary)
	log.Info("primary status is now %t", primary)
	w.Write([]byte("OK"))
}

func (c *ClusterStatus) getClusterStatus(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	resp, err := c.Marshal()
	if err != nil {
		http.Error(w, "could not marshal status to json", http.StatusInternalServerError)
		return
	}
	w.Write(resp)
}
<|endoftext|>


<commit_before>package terminal

import (
	"fmt"
	"strconv"

	"github.com/opsidian/parsley/ast"
	"github.com/opsidian/parsley/data"
	"github.com/opsidian/parsley/parser"
	"github.com/opsidian/parsley/reader"
	"github.com/opsidian/parsley/text"
	"github.com/opsidian/parsley/text/token"
)

// Float matches a float literal
func Float(bitSize int) parser.Func {
	return parser.Func(func(ctx data.IntMap, r reader.Reader) (cp data.IntSet, rs parser.ResultSet) {
		tr := r.(*text.Reader)
		if matches, pos := tr.ReadMatch("^[-+]?[0-9]*\\.?[0-9]+(?:[eE][-+]?[0-9]+)?"); matches != nil {
			val, err := strconv.ParseFloat(matches[0], bitSize)
			if err != nil {
				panic(fmt.Sprintf("Invalid float value encountered: %s", matches[0]))
			}
			if bitSize == 32 {
				rs = parser.NewResult(ast.NewTerminalNode(token.FLOAT, pos, float32(val)), r).AsSet()
			} else {
				rs = parser.NewResult(ast.NewTerminalNode(token.FLOAT, pos, val), r).AsSet()
			}
		}
		return
	})
}
<commit_msg>Don't use named returns in float terminal parser<commit_after>package terminal

import (
	"fmt"
	"strconv"

	"github.com/opsidian/parsley/ast"
	"github.com/opsidian/parsley/data"
	"github.com/opsidian/parsley/parser"
	"github.com/opsidian/parsley/reader"
	"github.com/opsidian/parsley/text"
	"github.com/opsidian/parsley/text/token"
)

// Float matches a float literal
func Float(bitSize int) parser.Func {
	return parser.Func(func(ctx data.IntMap, r reader.Reader) (data.IntSet, parser.ResultSet) {
		tr := r.(*text.Reader)
		if matches, pos := tr.ReadMatch("^[-+]?[0-9]*\\.?[0-9]+(?:[eE][-+]?[0-9]+)?"); matches != nil {
			val, err := strconv.ParseFloat(matches[0], bitSize)
			if err != nil {
				panic(fmt.Sprintf("Invalid float value encountered: %s", matches[0]))
			}
			var rs parser.ResultSet
			if bitSize == 32 {
				rs = parser.NewResult(ast.NewTerminalNode(token.FLOAT, pos, float32(val)), r).AsSet()
			} else {
				rs = parser.NewResult(ast.NewTerminalNode(token.FLOAT, pos, val), r).AsSet()
			}
			return parser.NoCurtailingParsers(), rs
		}
		return parser.NoCurtailingParsers(), nil
	})
}
<|endoftext|>


<commit_before>package gorpc

import (
	"bufio"
	"encoding/gob"
	"io"
	"net"
	"sync"
	"time"
)

type Client struct {
	Addr  string
	Conns int

	requestsChan chan *clientMessage
}

func (c *Client) Start() {
	c.requestsChan = make(chan *clientMessage, 1024)
	if c.Conns == 0 {
		c.Conns = 1
	}
	for i := 0; i < c.Conns; i++ {
		go clientHandler(c)
	}
}

func (c *Client) Send(request interface{}) interface{} {
	m := clientMessage{
		Request: request,
		Done:    make(chan struct{}, 1),
	}
	select {
	case c.requestsChan <- &m:
		<-m.Done
		return m.Response
	default:
		logError("rpc.Client: [%s]. Requests' queue with size=%d is overflown", cap(c.requestsChan), c.Addr)
		return nil
	}
}

func clientHandler(c *Client) {
	for {
		conn, err := net.Dial("tcp", c.Addr)
		if err != nil {
			logError("rpc.Client: [%s]. Cannot establish rpc connection: [%s]", c.Addr, err)
			time.Sleep(time.Second)
			continue
		}
		clientHandleConnection(c, conn)
	}
}

func clientHandleConnection(c *Client, conn net.Conn) {
	stopChan := make(chan struct{})

	pendingRequests := make(map[uint64]*clientMessage)
	var pendingRequestsLock sync.Mutex

	writerDone := make(chan struct{}, 1)
	go clientWriter(c, conn, pendingRequests, &pendingRequestsLock, stopChan, writerDone)

	readerDone := make(chan struct{}, 1)
	go clientReader(c, conn, pendingRequests, &pendingRequestsLock, readerDone)

	select {
	case <-writerDone:
		close(stopChan)
		conn.Close()
		<-readerDone
	case <-readerDone:
		close(stopChan)
		conn.Close()
		<-writerDone
	}

	for _, m := range pendingRequests {
		m.Done <- struct{}{}
	}
}

type clientMessage struct {
	Request  interface{}
	Response interface{}
	Done     chan struct{}
}

func clientWriter(c *Client, w io.Writer, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, stopChan <-chan struct{}, done chan<- struct{}) {
	defer func() { done <- struct{}{} }()

	var msgID uint64
	bw := bufio.NewWriter(w)
	e := gob.NewEncoder(bw)
	for {
		var rpcM *clientMessage

		msgID++
		select {
		case <-stopChan:
			return
		case rpcM = <-c.requestsChan:
		case <-time.After(10 * time.Millisecond):
			if err := bw.Flush(); err != nil {
				logError("rpc.Client: [%s]. Cannot flush requests to wire: [%s]", c.Addr, err)
				rpcM.Done <- struct{}{}
				return
			}
			select {
			case <-stopChan:
				return
			case rpcM = <-c.requestsChan:
			}
		}

		pendingRequestsLock.Lock()
		pendingRequests[msgID] = rpcM
		pendingRequestsLock.Unlock()

		m := wireMessage{
			ID:   msgID,
			Data: rpcM.Request,
		}
		if err := e.Encode(&m); err != nil {
			logError("rpc.Client: [%s]. Cannot send request to wire: [%s]", c.Addr, err)
			rpcM.Done <- struct{}{}
			pendingRequestsLock.Lock()
			delete(pendingRequests, msgID)
			pendingRequestsLock.Unlock()
			return
		}
	}
}

func clientReader(c *Client, r io.Reader, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, done chan<- struct{}) {
	defer func() { done <- struct{}{} }()

	br := bufio.NewReader(r)
	d := gob.NewDecoder(br)
	for {
		var m wireMessage
		if err := d.Decode(&m); err != nil {
			logError("rpc.Client: [%s]. Cannot read response from wire: [%s]", c.Addr, err)
			return
		}

		pendingRequestsLock.Lock()
		rpcM, ok := pendingRequests[m.ID]
		delete(pendingRequests, m.ID)
		pendingRequestsLock.Unlock()
		if !ok {
			logError("rpc.Client: [%s]. Unexpected msgID=[%d] obtained from server", c.Addr, m.ID)
			return
		}

		rpcM.Response = m.Data
		rpcM.Done <- struct{}{}
	}
}
<commit_msg>fixed a typo<commit_after>package gorpc

import (
	"bufio"
	"encoding/gob"
	"io"
	"net"
	"sync"
	"time"
)

type Client struct {
	Addr  string
	Conns int

	requestsChan chan *clientMessage
}

func (c *Client) Start() {
	c.requestsChan = make(chan *clientMessage, 1024)
	if c.Conns == 0 {
		c.Conns = 1
	}
	for i := 0; i < c.Conns; i++ {
		go clientHandler(c)
	}
}

func (c *Client) Send(request interface{}) interface{} {
	m := clientMessage{
		Request: request,
		Done:    make(chan struct{}, 1),
	}
	select {
	case c.requestsChan <- &m:
		<-m.Done
		return m.Response
	default:
		logError("rpc.Client: [%s]. Requests' queue with size=%d is overflown", c.Addr, cap(c.requestsChan))
		return nil
	}
}

func clientHandler(c *Client) {
	for {
		conn, err := net.Dial("tcp", c.Addr)
		if err != nil {
			logError("rpc.Client: [%s]. Cannot establish rpc connection: [%s]", c.Addr, err)
			time.Sleep(time.Second)
			continue
		}
		clientHandleConnection(c, conn)
	}
}

func clientHandleConnection(c *Client, conn net.Conn) {
	stopChan := make(chan struct{})

	pendingRequests := make(map[uint64]*clientMessage)
	var pendingRequestsLock sync.Mutex

	writerDone := make(chan struct{}, 1)
	go clientWriter(c, conn, pendingRequests, &pendingRequestsLock, stopChan, writerDone)

	readerDone := make(chan struct{}, 1)
	go clientReader(c, conn, pendingRequests, &pendingRequestsLock, readerDone)

	select {
	case <-writerDone:
		close(stopChan)
		conn.Close()
		<-readerDone
	case <-readerDone:
		close(stopChan)
		conn.Close()
		<-writerDone
	}

	for _, m := range pendingRequests {
		m.Done <- struct{}{}
	}
}

type clientMessage struct {
	Request  interface{}
	Response interface{}
	Done     chan struct{}
}

func clientWriter(c *Client, w io.Writer, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, stopChan <-chan struct{}, done chan<- struct{}) {
	defer func() { done <- struct{}{} }()

	var msgID uint64
	bw := bufio.NewWriter(w)
	e := gob.NewEncoder(bw)
	for {
		var rpcM *clientMessage

		msgID++
		select {
		case <-stopChan:
			return
		case rpcM = <-c.requestsChan:
		case <-time.After(10 * time.Millisecond):
			if err := bw.Flush(); err != nil {
				logError("rpc.Client: [%s]. Cannot flush requests to wire: [%s]", c.Addr, err)
				rpcM.Done <- struct{}{}
				return
			}
			select {
			case <-stopChan:
				return
			case rpcM = <-c.requestsChan:
			}
		}

		pendingRequestsLock.Lock()
		pendingRequests[msgID] = rpcM
		pendingRequestsLock.Unlock()

		m := wireMessage{
			ID:   msgID,
			Data: rpcM.Request,
		}
		if err := e.Encode(&m); err != nil {
			logError("rpc.Client: [%s]. Cannot send request to wire: [%s]", c.Addr, err)
			rpcM.Done <- struct{}{}
			pendingRequestsLock.Lock()
			delete(pendingRequests, msgID)
			pendingRequestsLock.Unlock()
			return
		}
	}
}

func clientReader(c *Client, r io.Reader, pendingRequests map[uint64]*clientMessage, pendingRequestsLock *sync.Mutex, done chan<- struct{}) {
	defer func() { done <- struct{}{} }()

	br := bufio.NewReader(r)
	d := gob.NewDecoder(br)
	for {
		var m wireMessage
		if err := d.Decode(&m); err != nil {
			logError("rpc.Client: [%s]. Cannot read response from wire: [%s]", c.Addr, err)
			return
		}

		pendingRequestsLock.Lock()
		rpcM, ok := pendingRequests[m.ID]
		delete(pendingRequests, m.ID)
		pendingRequestsLock.Unlock()
		if !ok {
			logError("rpc.Client: [%s]. Unexpected msgID=[%d] obtained from server", c.Addr, m.ID)
			return
		}

		rpcM.Response = m.Data
		rpcM.Done <- struct{}{}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/sheenobu/go-xco"
)

func main() {
	var config Config
	config = new(StaticConfig)
	_, err := toml.DecodeFile(os.Args[1], &config)
	if err != nil {
		panic(err)
	}

	opts := xco.Options{
		Name:         config.ComponentName(),
		SharedSecret: config.SharedSecret(),
		Address:      fmt.Sprintf("%s:%d", config.XmppHost(), config.XmppPort()),
	}
	c, err := xco.NewComponent(opts)
	if err != nil {
		panic(err)
	}

	// Uppercase Echo Component
	c.MessageHandler = xco.BodyResponseHandler(func(msg *xco.Message) (string, error) {
		return strings.ToUpper(msg.Body), nil
	})

	c.Run()
}
<commit_msg>Split main() in half<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/sheenobu/go-xco"
)

func main() {
	config := new(StaticConfig)
	_, err := toml.DecodeFile(os.Args[1], &config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "can't read config file '%s': %s\n", os.Args[1], err)
		os.Exit(1)
	}
	Main(config)
}

func Main(config Config) {
	opts := xco.Options{
		Name:         config.ComponentName(),
		SharedSecret: config.SharedSecret(),
		Address:      fmt.Sprintf("%s:%d", config.XmppHost(), config.XmppPort()),
	}
	c, err := xco.NewComponent(opts)
	if err != nil {
		panic(err)
	}

	// Uppercase Echo Component
	c.MessageHandler = xco.BodyResponseHandler(func(msg *xco.Message) (string, error) {
		return strings.ToUpper(msg.Body), nil
	})

	c.Run()
}
<|endoftext|>


<commit_before><commit_msg>1419778729<commit_after><|endoftext|>


<commit_before>package cuckoo

import (
	"strconv"
	"testing"
)

// Standardizing our entry formats
func keyAndValue(i int) (key keytype, value valuetype) {
	keystr := strconv.Itoa(i)
	key = keytype(keystr)
	value = valuetype("testvalue-" + keystr)
	return
}

// Helper function to populate the table
func fillTable(t *testing.T, ct *Table, limit int) {
	for i := 0; i < limit; i++ {
		k, v := keyAndValue(i)
		ct.Put(k, v)
	}
}

func validateFound(t *testing.T, ct *Table, start int, limit int, testName string) {
	for i := start; i < limit; i++ {
		k, wantedValue := keyAndValue(i)
		v, found := ct.Get(k)
		if !found {
			t.Fatal(testName, "could not find key", k)
		}
		if v != wantedValue {
			t.Fatalf(testName, "wrong value, expected %v got %v", i, wantedValue, v)
		}
	}
}

func TestNewTable(t *testing.T) {
	ct := NewTable()
	if ct == nil {
		t.Fatal("Could not allocate table")
	}
}

func TestBasic(t *testing.T) {
	ct := NewTable()
	nKeys := 1000
	fillTable(t, ct, nKeys)
	validateFound(t, ct, 0, nKeys, "TestBasic")
}

func TestFill(t *testing.T) {
	ct := NewTablePowerOfTwo(10)
	// Should be able to hold at least 950 elements, but will have to
	// cuckoo a lot to fill those last bits.  Stress test the cuckooing.
	limit := 874 // 875 fails - we're not BFS'ing well enough yet
	fillTable(t, ct, limit)
	validateFound(t, ct, 0, limit, "TestFill")
}

func TestDelete(t *testing.T) {
	ct := NewTable()
	limit := 1000
	fillTable(t, ct, limit)
	validateFound(t, ct, 0, limit, "TestDelete")
	for i := 0; i < limit; i++ {
		k, _ := keyAndValue(i)
		ct.Delete(k)
		foundVal, found := ct.Get(k)
		if found {
			t.Fatalf("TestDelete failed - item %v still present as %v", k, foundVal)
		}
		validateFound(t, ct, i+1, limit, "TestDelete")
	}
}
<commit_msg>Slightly more involved test involving keeping the table full through deletions and more insertions<commit_after>package cuckoo

import (
	"strconv"
	"testing"
)

// Standardizing our entry formats
func keyAndValue(i int) (key keytype, value valuetype) {
	keystr := strconv.Itoa(i)
	key = keytype(keystr)
	value = valuetype("testvalue-" + keystr)
	return
}

// Helper function to populate the table
func fillTable(t *testing.T, ct *Table, limit int) {
	for i := 0; i < limit; i++ {
		k, v := keyAndValue(i)
		ct.Put(k, v)
	}
}

func fatalIfNotFound(t *testing.T, ct *Table, i int) {
	k, v := keyAndValue(i)
	foundVal, found := ct.Get(k)
	if !found {
		t.Fatalf("failed - item %v not present", k)
	}
	if v != foundVal {
		t.Fatalf("failed - item %v value incorrect, wanted %v got %v", k, v, foundVal)
	}
}

func fatalIfFound(t *testing.T, ct *Table, i int) {
	k, _ := keyAndValue(i)
	foundVal, found := ct.Get(k)
	if found {
		t.Fatalf("failed - item %v not deleted, value %v", k, foundVal)
	}
}

func validateFound(t *testing.T, ct *Table, start int, limit int, testName string) {
	for i := start; i < limit; i++ {
		fatalIfNotFound(t, ct, i)
	}
}

func TestNewTable(t *testing.T) {
	ct := NewTable()
	if ct == nil {
		t.Fatal("Could not allocate table")
	}
}

func TestBasic(t *testing.T) {
	ct := NewTable()
	nKeys := 1000
	fillTable(t, ct, nKeys)
	validateFound(t, ct, 0, nKeys, "TestBasic")
}

func TestFill(t *testing.T) {
	ct := NewTablePowerOfTwo(10)
	// Should be able to hold at least 950 elements, but will have to
	// cuckoo a lot to fill those last bits.  Stress test the cuckooing.
	limit := 874 // 875 fails - we're not BFS'ing well enough yet
	fillTable(t, ct, limit)
	validateFound(t, ct, 0, limit, "TestFill")
}

func TestDelete(t *testing.T) {
	ct := NewTable()
	limit := 1000
	fillTable(t, ct, limit)
	validateFound(t, ct, 0, limit, "TestDelete")
	for i := 0; i < limit; i++ {
		k, _ := keyAndValue(i)
		ct.Delete(k)
		fatalIfFound(t, ct, i)
		validateFound(t, ct, i+1, limit, "TestDelete")
	}
}

func TestRunningNearFull(t *testing.T) {
	ct := NewTablePowerOfTwo(10)
	tn := "TestRunningNearFull"
	limit := 800
	fillTable(t, ct, limit)
	validateFound(t, ct, 0, limit, tn)
	for i := 0; i < 10000; i++ {
		fatalIfNotFound(t, ct, i)
		k, _ := keyAndValue(i)
		ct.Delete(k)
		fatalIfFound(t, ct, i)
		k2, v2 := keyAndValue(i+limit)
		ct.Put(k2, v2)
		fatalIfNotFound(t, ct, i+limit)
	}
}
<|endoftext|>


<commit_before>package eval

import (
	"errors"
	"fmt"
	"os"
	"syscall"

	"github.com/elves/elvish/util"
)

// FdNil is a special impossible fd value used for "close fd" in
// syscall.ProcAttr.Files.
const fdNil uintptr = ^uintptr(0)

var ErrCdNoArg = errors.New("implicit cd accepts no arguments")

// ExternalCmd is an external command.
type ExternalCmd struct {
	Name string
}

func (ExternalCmd) Kind() string {
	return "fn"
}

func (e ExternalCmd) Repr(int) string {
	return "<external " + e.Name + " >"
}

// Call calls an external command.
func (e ExternalCmd) Call(ec *EvalCtx, argVals []Value) {
	if util.DontSearch(e.Name) {
		stat, err := os.Stat(e.Name)
		if err == nil && stat.IsDir() {
			// implicit cd
			if len(argVals) > 0 {
				throw(ErrCdNoArg)
			}
			cdInner(e.Name, ec)
			return
		}
	}

	files := make([]uintptr, len(ec.ports))
	for i, port := range ec.ports {
		if port == nil || port.File == nil {
			files[i] = fdNil
		} else {
			files[i] = port.File.Fd()
		}
	}

	args := make([]string, len(argVals)+1)
	for i, a := range argVals {
		// NOTE Maybe we should enfore string arguments instead of coercing all
		// args into string
		args[i+1] = ToString(a)
	}

	sys := syscall.SysProcAttr{}
	if ec.Stub != nil {
		sys.Setpgid = true
		sys.Pgid = ec.Stub.Process().Pid
	}
	attr := syscall.ProcAttr{Env: os.Environ(), Files: files[:], Sys: &sys}

	path, err := ec.Search(e.Name)
	if err != nil {
		throw(errors.New("search: " + err.Error()))
	}

	args[0] = path
	pid, err := syscall.ForkExec(path, args, &attr)
	if err != nil {
		throw(errors.New("forkExec: " + err.Error()))
	}

	var ws syscall.WaitStatus
	_, err = syscall.Wait4(pid, &ws, syscall.WUNTRACED, nil)
	if err != nil {
		throw(fmt.Errorf("wait: %s", err.Error()))
	} else {
		maybeThrow(NewExternalCmdExit(ws, pid))
	}
}
<commit_msg>Fix stuttering of cnf error.<commit_after>package eval

import (
	"errors"
	"fmt"
	"os"
	"syscall"

	"github.com/elves/elvish/util"
)

// FdNil is a special impossible fd value used for "close fd" in
// syscall.ProcAttr.Files.
const fdNil uintptr = ^uintptr(0)

var ErrCdNoArg = errors.New("implicit cd accepts no arguments")

// ExternalCmd is an external command.
type ExternalCmd struct {
	Name string
}

func (ExternalCmd) Kind() string {
	return "fn"
}

func (e ExternalCmd) Repr(int) string {
	return "<external " + e.Name + " >"
}

// Call calls an external command.
func (e ExternalCmd) Call(ec *EvalCtx, argVals []Value) {
	if util.DontSearch(e.Name) {
		stat, err := os.Stat(e.Name)
		if err == nil && stat.IsDir() {
			// implicit cd
			if len(argVals) > 0 {
				throw(ErrCdNoArg)
			}
			cdInner(e.Name, ec)
			return
		}
	}

	files := make([]uintptr, len(ec.ports))
	for i, port := range ec.ports {
		if port == nil || port.File == nil {
			files[i] = fdNil
		} else {
			files[i] = port.File.Fd()
		}
	}

	args := make([]string, len(argVals)+1)
	for i, a := range argVals {
		// NOTE Maybe we should enfore string arguments instead of coercing all
		// args into string
		args[i+1] = ToString(a)
	}

	sys := syscall.SysProcAttr{}
	if ec.Stub != nil {
		sys.Setpgid = true
		sys.Pgid = ec.Stub.Process().Pid
	}
	attr := syscall.ProcAttr{Env: os.Environ(), Files: files[:], Sys: &sys}

	path, err := ec.Search(e.Name)
	if err != nil {
		throw(err)
	}

	args[0] = path
	pid, err := syscall.ForkExec(path, args, &attr)
	if err != nil {
		throw(errors.New("forkExec: " + err.Error()))
	}

	var ws syscall.WaitStatus
	_, err = syscall.Wait4(pid, &ws, syscall.WUNTRACED, nil)
	if err != nil {
		throw(fmt.Errorf("wait: %s", err.Error()))
	} else {
		maybeThrow(NewExternalCmdExit(ws, pid))
	}
}
<|endoftext|>


<commit_before>package maldocker

import (
	"os"
	"os/exec"
	"runtime"

	log "github.com/Sirupsen/logrus"
	docker "github.com/docker/engine-api/client"
	"github.com/maliceio/malice/config"
	er "github.com/maliceio/malice/malice/errors"
)

// NOTE: https://github.com/eris-ltd/eris-cli/blob/master/perform/docker_run.go

// Docker is the Malice docker client
type Docker struct {
	Client *docker.Client
	ip     string
	port   string
}

// NewDockerClient creates a new Docker Client
func NewDockerClient() *Docker {
	var client *docker.Client
	var ip, port string
	var err error

	// create docker client base on OS
	switch runtime.GOOS {
	case "darwin", "windows":
		// Create a new Client from Env
		if _, exists := os.LookupEnv("DOCKER_HOST"); exists {
			ip, port, err = parseDockerEndoint(os.Getenv("DOCKER_HOST"))
			er.CheckError(err)
			client, err = docker.NewEnvClient()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			// Create a new Client from config.Conf.Docker.EndPoint
			ip, port, err = parseDockerEndoint(config.Conf.Docker.EndPoint)
			er.CheckError(err)
			defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
			client, err = docker.NewClient(config.Conf.Docker.EndPoint, "v1.22", nil, defaultHeaders)
			if err != nil {
				log.Fatal(err)
			}
		}
	case "linux":
		defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
		client, err = docker.NewClient("unix:///var/run/docker.sock", "v1.22", nil, defaultHeaders)
		if err != nil {
			log.Fatal(err)
		}
	}
	// Check if client can connect
	_, err = client.Info()
	if err != nil {
		log.Fatal(err)
	}

	return &Docker{
		Client: client,
		ip:     ip,
		port:   port,
	}
}

// GetIP returns IP of docker client
func (client *Docker) GetIP() string {
	return client.ip
}

// TODO: Make this betta MUCHO betta
func handleClientError(dockerError error) {
	if dockerError != nil {
		log.WithFields(log.Fields{"env": config.Conf.Environment.Run}).Error("Unable to connect to docker client")
		switch runtime.GOOS {
		case "darwin":
			if _, err := exec.LookPath("docker-machine"); err != nil {
				log.Infof("Please install docker-machine by running: \n\tbrew install docker-machine\n\tdocker-machine create -d virtualbox %s\n\teval $(docker-machine env %s)\n", config.Conf.Docker.Name, config.Conf.Docker.Name)
			} else {
				log.Infof("Please start and source the docker-machine env by running: \n\tdocker-machine start %s\n\teval $(docker-machine env %s)\n", config.Conf.Docker.Name, config.Conf.Docker.Name)
			}
		case "linux":
			log.Info("Please start the docker daemon.")
		case "windows":
			if _, err := exec.LookPath("docker-machine.exe"); err != nil {
				log.Info("Please install docker-machine - https://www.docker.com/docker-toolbox")
			} else {
				log.Infof("Please start and source the docker-machine env by running: \n\tdocker-machine start %s\n\teval $(docker-machine env %s)\n", config.Conf.Docker.Name, config.Conf.Docker.Name)
			}
		}
		// TODO Decide if I want to make docker machines or rely on use to create their own.
		// log.Info("Trying to create new docker-machine: ", "test")
		// MakeDockerMachine("test")
		os.Exit(2)
	}
}
<commit_msg>re-introduce handleClientError<commit_after>package maldocker

import (
	"os"
	"os/exec"
	"runtime"

	log "github.com/Sirupsen/logrus"
	docker "github.com/docker/engine-api/client"
	"github.com/maliceio/malice/config"
	er "github.com/maliceio/malice/malice/errors"
)

// NOTE: https://github.com/eris-ltd/eris-cli/blob/master/perform/docker_run.go

// Docker is the Malice docker client
type Docker struct {
	Client *docker.Client
	ip     string
	port   string
}

// NewDockerClient creates a new Docker Client
func NewDockerClient() *Docker {
	var client *docker.Client
	var ip, port string
	var err error

	// Create docker client based on OS
	switch runtime.GOOS {
	case "darwin", "windows":
		// Create a new Client from Env
		if _, exists := os.LookupEnv("DOCKER_HOST"); exists {
			ip, port, err = parseDockerEndoint(os.Getenv("DOCKER_HOST"))
			er.CheckError(err)
			client, err = docker.NewEnvClient()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			// Create a new Client from config.Conf.Docker.EndPoint
			ip, port, err = parseDockerEndoint(config.Conf.Docker.EndPoint)
			er.CheckError(err)
			defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
			client, err = docker.NewClient(config.Conf.Docker.EndPoint, "v1.22", nil, defaultHeaders)
			if err != nil {
				log.Fatal(err)
			}
		}
	case "linux":
		defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
		client, err = docker.NewClient("unix:///var/run/docker.sock", "v1.22", nil, defaultHeaders)
		if err != nil {
			log.Fatal(err)
		}
	}
	// Check if client can connect
	_, err = client.Info()
	if err != nil {
		handleClientError(err)
		// log.Fatal(err)
	}

	return &Docker{
		Client: client,
		ip:     ip,
		port:   port,
	}
}

// GetIP returns IP of docker client
func (client *Docker) GetIP() string {
	return client.ip
}

// TODO: Make this betta MUCHO betta
func handleClientError(dockerError error) {
	if dockerError != nil {
		log.WithFields(log.Fields{"env": config.Conf.Environment.Run}).Error("Unable to connect to docker client")
		switch runtime.GOOS {
		case "darwin":
			if _, err := exec.LookPath("docker-machine"); err != nil {
				log.Info("Please install docker-machine by running: ")
				log.Info(" - brew install docker-machine")
				log.Infof(" - brew install docker-machine\n\tdocker-machine create -d virtualbox %s", config.Conf.Docker.Name)
				log.Infof(" - eval $(docker-machine env %s)", config.Conf.Docker.Name)
			} else {
				log.Info("Please start and source the docker-machine env by running: ")
				log.Infof(" - docker-machine start %s", config.Conf.Docker.Name)
				log.Infof(" - eval $(docker-machine env %s)", config.Conf.Docker.Name)
			}
		case "linux":
			log.Info("Please start the docker daemon.")
		case "windows":
			if _, err := exec.LookPath("docker-machine.exe"); err != nil {
				log.Info("Please install docker-machine - https://www.docker.com/docker-toolbox")
			} else {
				log.Info("Please start and source the docker-machine env by running: ")
				log.Infof(" - docker-machine start %", config.Conf.Docker.Name)
				log.Infof(" - eval $(docker-machine env %s)", config.Conf.Docker.Name)
			}
		}
		// TODO Decide if I want to make docker machines or rely on use to create their own.
		// log.Info("Trying to create new docker-machine: ", "test")
		// MakeDockerMachine("test")
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>package requestmatching

import (
	"net/http"
)

//MustMatch match request with given pattern
//Panic if any error raised.
func MustMatch(r *http.Request, p Pattern) bool {
	result, err := p.MatchRequest(r)
	if err != nil {
		panic(err)
	}
	return result
}

//Pattern request matching pattern interafe
type Pattern interface {
	//MatchRequest match request.
	//Return result and any error if raised.
	MatchRequest(r *http.Request) (bool, error)
}

//MatchAll match given request and given partterns.
//Match will fail if any pattern fail.
//Return result and any error if raised.
func MatchAll(r *http.Request, p ...Pattern) (bool, error) {
	var result bool
	var err error
	for k := range p {
		result, err = p[k].MatchRequest(r)
		if err != nil || result == false {
			return false, err
		}
	}
	return true, nil
}

//MustMatchAll match given request and given partterns.
//Match will fail if any pattern fail.
//Return result and painc if any error raised.
func MustMatchAll(r *http.Request, p ...Pattern) bool {
	result, err := MatchAll(r, p...)
	if err != nil {
		panic(err)
	}
	return result
}

//MatchAny match given request and given partterns.
//Match will success if any pattern success or no patterns given.
//Return result and any error if raised.
func MatchAny(r *http.Request, p ...Pattern) (bool, error) {
	var result bool
	var err error
	if len(p) == 0 {
		return true, nil
	}
	for k := range p {
		result, err = p[k].MatchRequest(r)
		if err != nil {
			return false, err
		}
		if result {
			return true, nil
		}
	}
	return false, nil
}

//MustMatchAny match given request and given partterns.
//Match will success if any pattern success.
//Return result and painc if any error raised.
func MustMatchAny(r *http.Request, p ...Pattern) bool {
	result, err := MatchAny(r, p...)
	if err != nil {
		panic(err)
	}
	return result
}

//PlainPattern plain pattern struct
type PlainPattern struct {
	IPNets   *IPNets
	Methods  *Methods
	Exts     *Exts
	Paths    *Paths
	Prefixs  *Prefixs
	Disabled bool
	Not      bool
	And      bool
	Patterns []Pattern
}

//NewPlainPattern create new pattern struct.
func NewPlainPattern() *PlainPattern {
	return &PlainPattern{
		IPNets:  NewIPNets(),
		Methods: NewMethods(),
		Exts:    NewExts(),
		Paths:   NewPaths(),
		Prefixs: NewPrefixs(),
	}
}

//MatchRequest match request.
//Return result and any error if raised.
func (p *PlainPattern) MatchRequest(r *http.Request) (bool, error) {
	if p.Disabled {
		return false, nil
	}
	result, err := p.matchRequest(r)
	if err != nil {
		return false, err
	}
	return result != p.Not, nil
}

func (p *PlainPattern) matchRequest(r *http.Request) (bool, error) {
	result, err := MatchAll(r,
		p.IPNets,
		p.Methods,
		p.Exts,
		p.Paths,
		p.Prefixs,
	)
	if err != nil {
		return false, err
	}
	if p.Patterns == nil {
		return result, nil
	}
	if p.And {
		return MatchAll(r, p.Patterns...)
	}
	if result == true {
		return true, nil
	}
	return MatchAny(r, p.Patterns...)
}

//PatternConfig plainpattern config struct
type PatternConfig struct {
	IPList     []string
	URLList    []string
	PrefixList []string
	ExtList    []string
	MethodList []string
	Disabled   bool
	Not        bool
	And        bool
	Patterns   []*PatternConfig
}

//CreatePattern create plain pattern.
//Retyurn pattern created and any error if raised.
func (c *PatternConfig) CreatePattern() (Pattern, error) {
	p := NewPlainPattern()
	for k := range c.IPList {
		err := p.IPNets.Add(c.IPList[k])
		if err != nil {
			return nil, err
		}
	}
	for k := range c.URLList {
		p.Paths.Add(c.URLList[k])
	}
	for k := range c.PrefixList {
		p.Prefixs.Add(c.PrefixList[k])
	}
	for k := range c.ExtList {
		p.Exts.Add(c.ExtList[k])
	}
	for k := range c.MethodList {
		p.Methods.Add(c.MethodList[k])
	}
	p.Disabled = c.Disabled
	p.Not = c.Not
	p.And = c.And
	for k := range c.Patterns {
		pattern, err := c.Patterns[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		p.Patterns = append(p.Patterns, pattern)
	}
	return p, nil
}

//Filters pattern filters type
type Filters []Pattern

//MatchRequest match request.
//Return result and any error if raised.
func (f *Filters) MatchRequest(r *http.Request) (bool, error) {
	return MatchAny(r, (*f)...)
}

//FiltersConfig filters config struct
type FiltersConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
//Match will success if filters is empty
func (c *FiltersConfig) CreatePattern() (Pattern, error) {
	f := Filters{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		f = append(f, pattern)
	}
	return &f, nil
}

//Whitelist whitelist pattern type
type Whitelist []Pattern

//MatchRequest match request.
//Return result and any error if raised.
//Match will fail if whitelist is empty
func (w *Whitelist) MatchRequest(r *http.Request) (bool, error) {
	if len(*w) == 0 {
		return false, nil
	}
	return MatchAny(r, (*w)...)
}

//WhitelistConfig whitelist config type
type WhitelistConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
func (c *WhitelistConfig) CreatePattern() (Pattern, error) {
	w := Whitelist{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		w = append(w, pattern)
	}
	return &w, nil
}

//PatternAll pattern all type
type PatternAll []Pattern

//MatchRequest match request.
//Return result and any error if raised.
func (p *PatternAll) MatchRequest(r *http.Request) (bool, error) {
	return MatchAll(r, (*p)...)
}

//PatternAllConfig pattern all type
type PatternAllConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
func (c *PatternAllConfig) CreatePattern() (Pattern, error) {
	p := PatternAll{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		p = append(p, pattern)
	}
	return &p, nil
}

//Factory Factory interface
type Factory interface {
	//CreatePattern create plain pattern.
	//Retyurn pattern created and any error if raised.
	CreatePattern() (Pattern, error)
}

//MustCreatePattern create pattern with given factory.
//Panic if any error raised.
func MustCreatePattern(f Factory) Pattern {
	p, err := f.CreatePattern()
	if err != nil {
		panic(err)
	}
	return p
}
<commit_msg>update<commit_after>package requestmatching

import (
	"net/http"
)

//MustMatch match request with given pattern
//Return false if pattern is nil.
//Panic if any error raised.
func MustMatch(r *http.Request, p Pattern) bool {
	if p == nil {
		return false
	}
	result, err := p.MatchRequest(r)
	if err != nil {
		panic(err)
	}
	return result
}

//Pattern request matching pattern interafe
type Pattern interface {
	//MatchRequest match request.
	//Return result and any error if raised.
	MatchRequest(r *http.Request) (bool, error)
}

//MatchAll match given request and given partterns.
//Match will fail if any pattern fail.
//Return result and any error if raised.
func MatchAll(r *http.Request, p ...Pattern) (bool, error) {
	var result bool
	var err error
	for k := range p {
		result, err = p[k].MatchRequest(r)
		if err != nil || result == false {
			return false, err
		}
	}
	return true, nil
}

//MustMatchAll match given request and given partterns.
//Match will fail if any pattern fail.
//Return result and painc if any error raised.
func MustMatchAll(r *http.Request, p ...Pattern) bool {
	result, err := MatchAll(r, p...)
	if err != nil {
		panic(err)
	}
	return result
}

//MatchAny match given request and given partterns.
//Match will success if any pattern success or no patterns given.
//Return result and any error if raised.
func MatchAny(r *http.Request, p ...Pattern) (bool, error) {
	var result bool
	var err error
	if len(p) == 0 {
		return true, nil
	}
	for k := range p {
		result, err = p[k].MatchRequest(r)
		if err != nil {
			return false, err
		}
		if result {
			return true, nil
		}
	}
	return false, nil
}

//MustMatchAny match given request and given partterns.
//Match will success if any pattern success.
//Return result and painc if any error raised.
func MustMatchAny(r *http.Request, p ...Pattern) bool {
	result, err := MatchAny(r, p...)
	if err != nil {
		panic(err)
	}
	return result
}

//PlainPattern plain pattern struct
type PlainPattern struct {
	IPNets   *IPNets
	Methods  *Methods
	Exts     *Exts
	Paths    *Paths
	Prefixs  *Prefixs
	Disabled bool
	Not      bool
	And      bool
	Patterns []Pattern
}

//NewPlainPattern create new pattern struct.
func NewPlainPattern() *PlainPattern {
	return &PlainPattern{
		IPNets:  NewIPNets(),
		Methods: NewMethods(),
		Exts:    NewExts(),
		Paths:   NewPaths(),
		Prefixs: NewPrefixs(),
	}
}

//MatchRequest match request.
//Return result and any error if raised.
func (p *PlainPattern) MatchRequest(r *http.Request) (bool, error) {
	if p.Disabled {
		return false, nil
	}
	result, err := p.matchRequest(r)
	if err != nil {
		return false, err
	}
	return result != p.Not, nil
}

func (p *PlainPattern) matchRequest(r *http.Request) (bool, error) {
	result, err := MatchAll(r,
		p.IPNets,
		p.Methods,
		p.Exts,
		p.Paths,
		p.Prefixs,
	)
	if err != nil {
		return false, err
	}
	if p.Patterns == nil {
		return result, nil
	}
	if p.And {
		return MatchAll(r, p.Patterns...)
	}
	if result == true {
		return true, nil
	}
	return MatchAny(r, p.Patterns...)
}

//PatternConfig plainpattern config struct
type PatternConfig struct {
	IPList     []string
	URLList    []string
	PrefixList []string
	ExtList    []string
	MethodList []string
	Disabled   bool
	Not        bool
	And        bool
	Patterns   []*PatternConfig
}

//CreatePattern create plain pattern.
//Retyurn pattern created and any error if raised.
func (c *PatternConfig) CreatePattern() (Pattern, error) {
	p := NewPlainPattern()
	for k := range c.IPList {
		err := p.IPNets.Add(c.IPList[k])
		if err != nil {
			return nil, err
		}
	}
	for k := range c.URLList {
		p.Paths.Add(c.URLList[k])
	}
	for k := range c.PrefixList {
		p.Prefixs.Add(c.PrefixList[k])
	}
	for k := range c.ExtList {
		p.Exts.Add(c.ExtList[k])
	}
	for k := range c.MethodList {
		p.Methods.Add(c.MethodList[k])
	}
	p.Disabled = c.Disabled
	p.Not = c.Not
	p.And = c.And
	for k := range c.Patterns {
		pattern, err := c.Patterns[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		p.Patterns = append(p.Patterns, pattern)
	}
	return p, nil
}

//Filters pattern filters type
type Filters []Pattern

//MatchRequest match request.
//Return result and any error if raised.
func (f *Filters) MatchRequest(r *http.Request) (bool, error) {
	return MatchAny(r, (*f)...)
}

//FiltersConfig filters config struct
type FiltersConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
//Match will success if filters is empty
func (c *FiltersConfig) CreatePattern() (Pattern, error) {
	f := Filters{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		f = append(f, pattern)
	}
	return &f, nil
}

//Whitelist whitelist pattern type
type Whitelist []Pattern

//MatchRequest match request.
//Return result and any error if raised.
//Match will fail if whitelist is empty
func (w *Whitelist) MatchRequest(r *http.Request) (bool, error) {
	if len(*w) == 0 {
		return false, nil
	}
	return MatchAny(r, (*w)...)
}

//WhitelistConfig whitelist config type
type WhitelistConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
func (c *WhitelistConfig) CreatePattern() (Pattern, error) {
	w := Whitelist{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		w = append(w, pattern)
	}
	return &w, nil
}

//PatternAll pattern all type
type PatternAll []Pattern

//MatchRequest match request.
//Return result and any error if raised.
func (p *PatternAll) MatchRequest(r *http.Request) (bool, error) {
	return MatchAll(r, (*p)...)
}

//PatternAllConfig pattern all type
type PatternAllConfig []*PatternConfig

//CreatePattern create plain pattern.
//Return pattern created and any error if raised.
func (c *PatternAllConfig) CreatePattern() (Pattern, error) {
	p := PatternAll{}
	for k := range *c {
		pattern, err := (*c)[k].CreatePattern()
		if err != nil {
			return nil, err
		}
		p = append(p, pattern)
	}
	return &p, nil
}

//Factory Factory interface
type Factory interface {
	//CreatePattern create plain pattern.
	//Retyurn pattern created and any error if raised.
	CreatePattern() (Pattern, error)
}

//MustCreatePattern create pattern with given factory.
//Panic if any error raised.
func MustCreatePattern(f Factory) Pattern {
	p, err := f.CreatePattern()
	if err != nil {
		panic(err)
	}
	return p
}
<|endoftext|>


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

// HTTP server that uses OAuth to create security.Blessings objects.
package server

import (
	"crypto/rand"
	"fmt"
	mrand "math/rand"
	"net"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/x/ref/lib/security/audit"
	"v.io/x/ref/lib/signals"
	"v.io/x/ref/services/discharger"
	"v.io/x/ref/services/identity/internal/auditor"
	"v.io/x/ref/services/identity/internal/blesser"
	"v.io/x/ref/services/identity/internal/caveats"
	"v.io/x/ref/services/identity/internal/dischargerlib"
	"v.io/x/ref/services/identity/internal/handlers"
	"v.io/x/ref/services/identity/internal/oauth"
	"v.io/x/ref/services/identity/internal/revocation"
	"v.io/x/ref/services/identity/internal/templates"
)

const (
	// TODO(ataly, ashankar, suharshs): The name "google" for the oauthBlesserService does
	// not seem appropriate given our modular construction of the identity server. The
	// oauthBlesserService can use any oauthProvider of its choosing, i.e., it does not
	// always have to be "google". One option would be change the value to "oauth". This
	// would also make the name analogous to that of macaroonService. Note that this option
	// also requires changing the extension.
	oauthBlesserService = "google"
	macaroonService     = "macaroon"
	dischargerService   = "discharger"
)

type IdentityServer struct {
	oauthProvider      oauth.OAuthProvider
	auditor            audit.Auditor
	blessingLogReader  auditor.BlessingLogReader
	revocationManager  revocation.RevocationManager
	oauthBlesserParams blesser.OAuthBlesserParams
	caveatSelector     caveats.CaveatSelector
	rootedObjectAddrs  []naming.Endpoint
	assetsPrefix       string
	mountNamePrefix    string
}

// NewIdentityServer returns a IdentityServer that:
// - uses oauthProvider to authenticate users
// - auditor and blessingLogReader to audit the root principal and read audit logs
// - revocationManager to store revocation data and grant discharges
// - oauthBlesserParams to configure the identity.OAuthBlesser service
func NewIdentityServer(oauthProvider oauth.OAuthProvider, auditor audit.Auditor, blessingLogReader auditor.BlessingLogReader, revocationManager revocation.RevocationManager, oauthBlesserParams blesser.OAuthBlesserParams, caveatSelector caveats.CaveatSelector, assetsPrefix, mountNamePrefix string) *IdentityServer {
	return &IdentityServer{
		oauthProvider:      oauthProvider,
		auditor:            auditor,
		blessingLogReader:  blessingLogReader,
		revocationManager:  revocationManager,
		oauthBlesserParams: oauthBlesserParams,
		caveatSelector:     caveatSelector,
		assetsPrefix:       assetsPrefix,
		mountNamePrefix:    mountNamePrefix,
	}
}

// findUnusedPort finds an unused port and returns it. Of course, no guarantees
// are made that the port will actually be available by the time the caller
// gets around to binding to it. If no port can be found, (0, nil) is returned.
// If an error occurs while creating a socket, that error is returned and the
// other return value is 0.
func findUnusedPort() (int, error) {
	random := mrand.New(mrand.NewSource(time.Now().UnixNano()))
	for i := 0; i < 1000; i++ {
		fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
		if err != nil {
			return 0, err
		}

		port := int(1024 + random.Int31n(64512))
		sa := &syscall.SockaddrInet4{Port: port}
		err = syscall.Bind(fd, sa)
		syscall.Close(fd)
		if err == nil {
			return port, nil
		}
	}

	return 0, nil
}

func (s *IdentityServer) Serve(ctx *context.T, externalHttpAddr, httpAddr, tlsConfig string) {
	ctx, err := v23.WithPrincipal(ctx, audit.NewPrincipal(ctx, s.auditor))
	if err != nil {
		ctx.Panic(err)
	}
	httphost, httpport, err := net.SplitHostPort(httpAddr)
	if err != nil || httpport == "0" {
		httpportNum, err := findUnusedPort()
		if err != nil {
			ctx.Panic(err)
		}
		httpAddr = net.JoinHostPort(httphost, strconv.Itoa(httpportNum))
	}
	rpcServer, _, externalAddr := s.Listen(ctx, externalHttpAddr, httpAddr, tlsConfig)
	fmt.Printf("HTTP_ADDR=%s\n", externalAddr)
	if len(s.rootedObjectAddrs) > 0 {
		fmt.Printf("NAME=%s\n", s.rootedObjectAddrs[0].Name())
	}
	<-signals.ShutdownOnSignals(ctx)
	ctx.Infof("Received shutdown request.")
	if err := rpcServer.Stop(); err != nil {
		ctx.Errorf("Failed to stop rpc server: %v", err)
	}
	ctx.Infof("Successfully stopped the rpc server.")
}

func (s *IdentityServer) Listen(ctx *context.T, externalHttpAddr, httpAddr, tlsConfig string) (rpc.Server, []string, string) {
	// json-encoded public key and blessing names of this server
	principal := v23.GetPrincipal(ctx)
	http.Handle("/auth/blessing-root", handlers.BlessingRoot{principal})

	macaroonKey := make([]byte, 32)
	if _, err := rand.Read(macaroonKey); err != nil {
		ctx.Fatalf("macaroonKey generation failed: %v", err)
	}

	rpcServer, published, err := s.setupBlessingServices(ctx, macaroonKey)
	if err != nil {
		ctx.Fatalf("Failed to setup vanadium services for blessing: %v", err)
	}

	externalHttpAddr = httpAddress(externalHttpAddr, httpAddr)

	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})

	n := "/auth/google/"
	args := oauth.HandlerArgs{
		Principal:          principal,
		MacaroonKey:        macaroonKey,
		Addr:               fmt.Sprintf("%s%s", externalHttpAddr, n),
		BlessingLogReader:  s.blessingLogReader,
		RevocationManager:  s.revocationManager,
		DischargerLocation: naming.JoinAddressName(published[0], dischargerService),
		MacaroonBlessingService: func() []string {
			status := rpcServer.Status()
			names := make([]string, len(status.Endpoints))
			for i, e := range status.Endpoints {
				names[i] = naming.JoinAddressName(e.Name(), macaroonService)
			}
			return names
		},
		OAuthProvider:  s.oauthProvider,
		CaveatSelector: s.caveatSelector,
		AssetsPrefix:   s.assetsPrefix,
	}
	if s.revocationManager != nil {
		args.DischargeServers = appendSuffixTo(published, dischargerService)
	}
	var emptyParams blesser.OAuthBlesserParams
	if !reflect.DeepEqual(s.oauthBlesserParams, emptyParams) {
		args.GoogleServers = appendSuffixTo(published, oauthBlesserService)
	}
	h, err := oauth.NewHandler(ctx, args)
	if err != nil {
		ctx.Fatalf("Failed to create HTTP handler for oauth authentication: %v", err)
	}
	http.Handle(n, h)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		tmplArgs := struct {
			Self                            security.Blessings
			GoogleServers, DischargeServers []string
			ListBlessingsRoute              string
			AssetsPrefix                    string
			Email                           string
		}{
			Self:               principal.BlessingStore().Default(),
			GoogleServers:      args.GoogleServers,
			DischargeServers:   args.DischargeServers,
			ListBlessingsRoute: oauth.ListBlessingsRoute,
			AssetsPrefix:       s.assetsPrefix,
		}
		if err := templates.Home.Execute(w, tmplArgs); err != nil {
			ctx.Info("Failed to render template:", err)
		}
	})
	ctx.Infof("Running HTTP server at: %v", externalHttpAddr)
	go runHTTPSServer(ctx, httpAddr, tlsConfig)
	return rpcServer, published, externalHttpAddr
}

func appendSuffixTo(objectname []string, suffix string) []string {
	names := make([]string, len(objectname))
	for i, o := range objectname {
		names[i] = naming.JoinAddressName(o, suffix)
	}
	return names
}

// Starts the Vanadium and HTTP services for blessing, and the Vanadium service for discharging.
// All Vanadium services are started on the same port.
func (s *IdentityServer) setupBlessingServices(ctx *context.T, macaroonKey []byte) (rpc.Server, []string, error) {
	disp := newDispatcher(macaroonKey, s.oauthBlesserParams)
	blessingNames := security.BlessingNames(v23.GetPrincipal(ctx), v23.GetPrincipal(ctx).BlessingStore().Default())
	if len(blessingNames) == 0 {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("identity server has no blessings?"))
	}
	if len(blessingNames) > 1 {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("cannot configure identity server with >1 (%d = %v) blessings - not quite sure what names to select for the discharger service etc.", len(blessingNames), blessingNames))
	}
	objectAddr := naming.Join(s.mountNamePrefix, naming.EncodeAsNameElement(blessingNames[0]))
	ctx, server, err := v23.WithNewDispatchingServer(ctx, objectAddr, disp)
	if err != nil {
		return nil, nil, err
	}
	// TODO(ashankar): Remove when https://github.com/vanadium/issues/issues/739
	// is resolved - at that point don't need to export the old style names
	// (with /s instead of :s)
	bug739BlessingName := strings.Replace(blessingNames[0], security.ChainSeparator, "/", -1)
	if err := server.AddName(bug739BlessingName); err != nil {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("failed to publish name %q: %v", bug739BlessingName, err))
	}
	s.rootedObjectAddrs = server.Status().Endpoints
	var rootedObjectAddr string
	if naming.Rooted(objectAddr) {
		rootedObjectAddr = objectAddr
	} else if nsroots := v23.GetNamespace(ctx).Roots(); len(nsroots) >= 1 {
		rootedObjectAddr = naming.Join(nsroots[0], objectAddr)
	} else {
		rootedObjectAddr = s.rootedObjectAddrs[0].Name()
	}
	disp.activate(rootedObjectAddr)
	ctx.Infof("Vanadium Blessing and discharger services will be published at %v", rootedObjectAddr)
	// Start the HTTP Handler for the OAuth2 access token based blesser.
	s.oauthBlesserParams.DischargerLocation = naming.Join(rootedObjectAddr, dischargerService)
	http.Handle("/auth/google/bless", handlers.NewOAuthBlessingHandler(ctx, s.oauthBlesserParams))
	return server, []string{rootedObjectAddr}, nil
}

// newDispatcher returns a dispatcher for both the blessing and the
// discharging service.
func newDispatcher(macaroonKey []byte, blesserParams blesser.OAuthBlesserParams) *dispatcher {
	d := &dispatcher{}
	d.macaroonKey = macaroonKey
	d.blesserParams = blesserParams
	d.wg.Add(1) // Will be removed at activate.
	return d
}

type dispatcher struct {
	m             map[string]interface{}
	wg            sync.WaitGroup
	macaroonKey   []byte
	blesserParams blesser.OAuthBlesserParams
}

func (d *dispatcher) Lookup(ctx *context.T, suffix string) (interface{}, security.Authorizer, error) {
	d.wg.Wait() // Wait until activate is called.
	if invoker := d.m[suffix]; invoker != nil {
		return invoker, security.AllowEveryone(), nil
	}
	return nil, nil, verror.New(verror.ErrNoExist, ctx, suffix)
}

func (d *dispatcher) activate(serverName string) {
	d.blesserParams.DischargerLocation = naming.Join(serverName, dischargerService)
	d.m = map[string]interface{}{
		macaroonService:     blesser.NewMacaroonBlesserServer(d.macaroonKey),
		dischargerService:   discharger.DischargerServer(dischargerlib.NewDischarger()),
		oauthBlesserService: blesser.NewOAuthBlesserServer(d.blesserParams),
	}
	// Set up the glob invoker.
	var children []string
	for k, _ := range d.m {
		children = append(children, k)
	}
	d.m[""] = rpc.ChildrenGlobberInvoker(children...)
	d.wg.Done() // Trigger any pending lookups.
}

func runHTTPSServer(ctx *context.T, addr, tlsConfig string) {
	if len(tlsConfig) == 0 {
		ctx.Fatal("Please set the --tls-config flag")
	}
	paths := strings.Split(tlsConfig, ",")
	if len(paths) != 2 {
		ctx.Fatalf("Could not parse --tls-config. Must have exactly two components, separated by a comma")
	}
	ctx.Infof("Starting HTTP server with TLS using certificate [%s] and private key [%s] at https://%s", paths[0], paths[1], addr)
	if err := http.ListenAndServeTLS(addr, paths[0], paths[1], nil); err != nil {
		ctx.Fatalf("http.ListenAndServeTLS failed: %v", err)
	}
}

func httpAddress(externalHttpAddr, httpAddr string) string {
	// If an externalHttpAddr is provided use that.
	if externalHttpAddr != "" {
		httpAddr = externalHttpAddr
	}
	return fmt.Sprintf("https://%v", httpAddr)
}
<commit_msg>Fix typo in fix for vanadium/issues#924<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP server that uses OAuth to create security.Blessings objects.
package server

import (
	"crypto/rand"
	"fmt"
	mrand "math/rand"
	"net"
	"net/http"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/naming"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/v23/verror"
	"v.io/x/ref/lib/security/audit"
	"v.io/x/ref/lib/signals"
	"v.io/x/ref/services/discharger"
	"v.io/x/ref/services/identity/internal/auditor"
	"v.io/x/ref/services/identity/internal/blesser"
	"v.io/x/ref/services/identity/internal/caveats"
	"v.io/x/ref/services/identity/internal/dischargerlib"
	"v.io/x/ref/services/identity/internal/handlers"
	"v.io/x/ref/services/identity/internal/oauth"
	"v.io/x/ref/services/identity/internal/revocation"
	"v.io/x/ref/services/identity/internal/templates"
)

const (
	// TODO(ataly, ashankar, suharshs): The name "google" for the oauthBlesserService does
	// not seem appropriate given our modular construction of the identity server. The
	// oauthBlesserService can use any oauthProvider of its choosing, i.e., it does not
	// always have to be "google". One option would be change the value to "oauth". This
	// would also make the name analogous to that of macaroonService. Note that this option
	// also requires changing the extension.
	oauthBlesserService = "google"
	macaroonService     = "macaroon"
	dischargerService   = "discharger"
)

type IdentityServer struct {
	oauthProvider      oauth.OAuthProvider
	auditor            audit.Auditor
	blessingLogReader  auditor.BlessingLogReader
	revocationManager  revocation.RevocationManager
	oauthBlesserParams blesser.OAuthBlesserParams
	caveatSelector     caveats.CaveatSelector
	rootedObjectAddrs  []naming.Endpoint
	assetsPrefix       string
	mountNamePrefix    string
}

// NewIdentityServer returns a IdentityServer that:
// - uses oauthProvider to authenticate users
// - auditor and blessingLogReader to audit the root principal and read audit logs
// - revocationManager to store revocation data and grant discharges
// - oauthBlesserParams to configure the identity.OAuthBlesser service
func NewIdentityServer(oauthProvider oauth.OAuthProvider, auditor audit.Auditor, blessingLogReader auditor.BlessingLogReader, revocationManager revocation.RevocationManager, oauthBlesserParams blesser.OAuthBlesserParams, caveatSelector caveats.CaveatSelector, assetsPrefix, mountNamePrefix string) *IdentityServer {
	return &IdentityServer{
		oauthProvider:      oauthProvider,
		auditor:            auditor,
		blessingLogReader:  blessingLogReader,
		revocationManager:  revocationManager,
		oauthBlesserParams: oauthBlesserParams,
		caveatSelector:     caveatSelector,
		assetsPrefix:       assetsPrefix,
		mountNamePrefix:    mountNamePrefix,
	}
}

// findUnusedPort finds an unused port and returns it. Of course, no guarantees
// are made that the port will actually be available by the time the caller
// gets around to binding to it. If no port can be found, (0, nil) is returned.
// If an error occurs while creating a socket, that error is returned and the
// other return value is 0.
func findUnusedPort() (int, error) {
	random := mrand.New(mrand.NewSource(time.Now().UnixNano()))
	for i := 0; i < 1000; i++ {
		fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP)
		if err != nil {
			return 0, err
		}

		port := int(1024 + random.Int31n(64512))
		sa := &syscall.SockaddrInet4{Port: port}
		err = syscall.Bind(fd, sa)
		syscall.Close(fd)
		if err == nil {
			return port, nil
		}
	}

	return 0, nil
}

func (s *IdentityServer) Serve(ctx *context.T, externalHttpAddr, httpAddr, tlsConfig string) {
	ctx, err := v23.WithPrincipal(ctx, audit.NewPrincipal(ctx, s.auditor))
	if err != nil {
		ctx.Panic(err)
	}
	httphost, httpport, err := net.SplitHostPort(httpAddr)
	if err != nil || httpport == "0" {
		httpportNum, err := findUnusedPort()
		if err != nil {
			ctx.Panic(err)
		}
		httpAddr = net.JoinHostPort(httphost, strconv.Itoa(httpportNum))
	}
	rpcServer, _, externalAddr := s.Listen(ctx, externalHttpAddr, httpAddr, tlsConfig)
	fmt.Printf("HTTP_ADDR=%s\n", externalAddr)
	if len(s.rootedObjectAddrs) > 0 {
		fmt.Printf("NAME=%s\n", s.rootedObjectAddrs[0].Name())
	}
	<-signals.ShutdownOnSignals(ctx)
	ctx.Infof("Received shutdown request.")
	if err := rpcServer.Stop(); err != nil {
		ctx.Errorf("Failed to stop rpc server: %v", err)
	}
	ctx.Infof("Successfully stopped the rpc server.")
}

func (s *IdentityServer) Listen(ctx *context.T, externalHttpAddr, httpAddr, tlsConfig string) (rpc.Server, []string, string) {
	// json-encoded public key and blessing names of this server
	principal := v23.GetPrincipal(ctx)
	http.Handle("/auth/blessing-root", handlers.BlessingRoot{principal})

	macaroonKey := make([]byte, 32)
	if _, err := rand.Read(macaroonKey); err != nil {
		ctx.Fatalf("macaroonKey generation failed: %v", err)
	}

	rpcServer, published, err := s.setupBlessingServices(ctx, macaroonKey)
	if err != nil {
		ctx.Fatalf("Failed to setup vanadium services for blessing: %v", err)
	}

	externalHttpAddr = httpAddress(externalHttpAddr, httpAddr)

	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})

	n := "/auth/google/"
	args := oauth.HandlerArgs{
		Principal:          principal,
		MacaroonKey:        macaroonKey,
		Addr:               fmt.Sprintf("%s%s", externalHttpAddr, n),
		BlessingLogReader:  s.blessingLogReader,
		RevocationManager:  s.revocationManager,
		DischargerLocation: naming.JoinAddressName(published[0], dischargerService),
		MacaroonBlessingService: func() []string {
			status := rpcServer.Status()
			names := make([]string, len(status.Endpoints))
			for i, e := range status.Endpoints {
				names[i] = naming.JoinAddressName(e.Name(), macaroonService)
			}
			return names
		},
		OAuthProvider:  s.oauthProvider,
		CaveatSelector: s.caveatSelector,
		AssetsPrefix:   s.assetsPrefix,
	}
	if s.revocationManager != nil {
		args.DischargeServers = appendSuffixTo(published, dischargerService)
	}
	var emptyParams blesser.OAuthBlesserParams
	if !reflect.DeepEqual(s.oauthBlesserParams, emptyParams) {
		args.GoogleServers = appendSuffixTo(published, oauthBlesserService)
	}
	h, err := oauth.NewHandler(ctx, args)
	if err != nil {
		ctx.Fatalf("Failed to create HTTP handler for oauth authentication: %v", err)
	}
	http.Handle(n, h)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		tmplArgs := struct {
			Self                            security.Blessings
			GoogleServers, DischargeServers []string
			ListBlessingsRoute              string
			AssetsPrefix                    string
			Email                           string
		}{
			Self:               principal.BlessingStore().Default(),
			GoogleServers:      args.GoogleServers,
			DischargeServers:   args.DischargeServers,
			ListBlessingsRoute: oauth.ListBlessingsRoute,
			AssetsPrefix:       s.assetsPrefix,
		}
		if err := templates.Home.Execute(w, tmplArgs); err != nil {
			ctx.Info("Failed to render template:", err)
		}
	})
	ctx.Infof("Running HTTP server at: %v", externalHttpAddr)
	go runHTTPSServer(ctx, httpAddr, tlsConfig)
	return rpcServer, published, externalHttpAddr
}

func appendSuffixTo(objectname []string, suffix string) []string {
	names := make([]string, len(objectname))
	for i, o := range objectname {
		names[i] = naming.JoinAddressName(o, suffix)
	}
	return names
}

// Starts the Vanadium and HTTP services for blessing, and the Vanadium service for discharging.
// All Vanadium services are started on the same port.
func (s *IdentityServer) setupBlessingServices(ctx *context.T, macaroonKey []byte) (rpc.Server, []string, error) {
	disp := newDispatcher(macaroonKey, s.oauthBlesserParams)
	blessingNames := security.BlessingNames(v23.GetPrincipal(ctx), v23.GetPrincipal(ctx).BlessingStore().Default())
	if len(blessingNames) == 0 {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("identity server has no blessings?"))
	}
	if len(blessingNames) > 1 {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("cannot configure identity server with >1 (%d = %v) blessings - not quite sure what names to select for the discharger service etc.", len(blessingNames), blessingNames))
	}
	objectAddr := naming.Join(s.mountNamePrefix, naming.EncodeAsNameElement(blessingNames[0]))
	ctx, server, err := v23.WithNewDispatchingServer(ctx, objectAddr, disp)
	if err != nil {
		return nil, nil, err
	}
	// TODO(ashankar): Remove when https://github.com/vanadium/issues/issues/739
	// is resolved - at that point don't need to export the old style names
	// (with /s instead of :s)
	bug739BlessingName := naming.Join(s.mountNamePrefix, strings.Replace(blessingNames[0], security.ChainSeparator, "/", -1))
	if err := server.AddName(bug739BlessingName); err != nil {
		return nil, nil, verror.New(verror.ErrInternal, ctx, fmt.Sprintf("failed to publish name %q: %v", bug739BlessingName, err))
	}
	ctx.Infof("Also publishing under %q", bug739BlessingName)
	s.rootedObjectAddrs = server.Status().Endpoints
	var rootedObjectAddr string
	if naming.Rooted(objectAddr) {
		rootedObjectAddr = objectAddr
	} else if nsroots := v23.GetNamespace(ctx).Roots(); len(nsroots) >= 1 {
		rootedObjectAddr = naming.Join(nsroots[0], objectAddr)
	} else {
		rootedObjectAddr = s.rootedObjectAddrs[0].Name()
	}
	disp.activate(rootedObjectAddr)
	ctx.Infof("Vanadium Blessing and discharger services will be published at %v", rootedObjectAddr)
	// Start the HTTP Handler for the OAuth2 access token based blesser.
	s.oauthBlesserParams.DischargerLocation = naming.Join(rootedObjectAddr, dischargerService)
	http.Handle("/auth/google/bless", handlers.NewOAuthBlessingHandler(ctx, s.oauthBlesserParams))
	return server, []string{rootedObjectAddr}, nil
}

// newDispatcher returns a dispatcher for both the blessing and the
// discharging service.
func newDispatcher(macaroonKey []byte, blesserParams blesser.OAuthBlesserParams) *dispatcher {
	d := &dispatcher{}
	d.macaroonKey = macaroonKey
	d.blesserParams = blesserParams
	d.wg.Add(1) // Will be removed at activate.
	return d
}

type dispatcher struct {
	m             map[string]interface{}
	wg            sync.WaitGroup
	macaroonKey   []byte
	blesserParams blesser.OAuthBlesserParams
}

func (d *dispatcher) Lookup(ctx *context.T, suffix string) (interface{}, security.Authorizer, error) {
	d.wg.Wait() // Wait until activate is called.
	if invoker := d.m[suffix]; invoker != nil {
		return invoker, security.AllowEveryone(), nil
	}
	return nil, nil, verror.New(verror.ErrNoExist, ctx, suffix)
}

func (d *dispatcher) activate(serverName string) {
	d.blesserParams.DischargerLocation = naming.Join(serverName, dischargerService)
	d.m = map[string]interface{}{
		macaroonService:     blesser.NewMacaroonBlesserServer(d.macaroonKey),
		dischargerService:   discharger.DischargerServer(dischargerlib.NewDischarger()),
		oauthBlesserService: blesser.NewOAuthBlesserServer(d.blesserParams),
	}
	// Set up the glob invoker.
	var children []string
	for k, _ := range d.m {
		children = append(children, k)
	}
	d.m[""] = rpc.ChildrenGlobberInvoker(children...)
	d.wg.Done() // Trigger any pending lookups.
}

func runHTTPSServer(ctx *context.T, addr, tlsConfig string) {
	if len(tlsConfig) == 0 {
		ctx.Fatal("Please set the --tls-config flag")
	}
	paths := strings.Split(tlsConfig, ",")
	if len(paths) != 2 {
		ctx.Fatalf("Could not parse --tls-config. Must have exactly two components, separated by a comma")
	}
	ctx.Infof("Starting HTTP server with TLS using certificate [%s] and private key [%s] at https://%s", paths[0], paths[1], addr)
	if err := http.ListenAndServeTLS(addr, paths[0], paths[1], nil); err != nil {
		ctx.Fatalf("http.ListenAndServeTLS failed: %v", err)
	}
}

func httpAddress(externalHttpAddr, httpAddr string) string {
	// If an externalHttpAddr is provided use that.
	if externalHttpAddr != "" {
		httpAddr = externalHttpAddr
	}
	return fmt.Sprintf("https://%v", httpAddr)
}
<|endoftext|>


<commit_before>package models

type User struct {
	Id       int64
	Name     string
	Role     string
	Username string
	Password []byte
}
<commit_msg>Delete user.go<commit_after><|endoftext|>


<commit_before>// Package junit defines a JUnit XML report and includes convenience methods
// for working with these reports.
package junit

import (
	"encoding/xml"
	"fmt"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
)

// Testsuites is a collection of JUnit testsuites.
type Testsuites struct {
	XMLName xml.Name `xml:"testsuites"`

	Name     string `xml:"name,attr,omitempty"`
	Time     string `xml:"time,attr,omitempty"` // total duration in seconds
	Tests    int    `xml:"tests,attr,omitempty"`
	Errors   int    `xml:"errors,attr,omitempty"`
	Failures int    `xml:"failures,attr,omitempty"`
	Skipped  int    `xml:"skipped,attr,omitempty"`
	Disabled int    `xml:"disabled,attr,omitempty"`

	Suites []Testsuite `xml:"testsuite,omitempty"`
}

// AddSuite adds a Testsuite and updates this testssuites' totals.
func (t *Testsuites) AddSuite(ts Testsuite) {
	t.Suites = append(t.Suites, ts)
	t.Tests += ts.Tests
	t.Errors += ts.Errors
	t.Failures += ts.Failures
	t.Skipped += ts.Skipped
	t.Disabled += ts.Disabled
}

// Testsuite is a single JUnit testsuite containing testcases.
type Testsuite struct {
	// required attributes
	Name     string `xml:"name,attr"`
	Tests    int    `xml:"tests,attr"`
	Failures int    `xml:"failures,attr"`
	Errors   int    `xml:"errors,attr"`
	ID       int    `xml:"id,attr"`

	// optional attributes
	Disabled  int    `xml:"disabled,attr,omitempty"`
	Hostname  string `xml:"hostname,attr,omitempty"`
	Package   string `xml:"package,attr,omitempty"`
	Skipped   int    `xml:"skipped,attr,omitempty"`
	Time      string `xml:"time,attr"`                // duration in seconds
	Timestamp string `xml:"timestamp,attr,omitempty"` // date and time in ISO8601

	Properties *[]Property `xml:"properties>property,omitempty"`
	Testcases  []Testcase  `xml:"testcase,omitempty"`
	SystemOut  *Output     `xml:"system-out,omitempty"`
	SystemErr  *Output     `xml:"system-err,omitempty"`
}

// AddProperty adds a property with the given name and value to this Testsuite.
func (t *Testsuite) AddProperty(name, value string) {
	prop := Property{Name: name, Value: value}
	if t.Properties == nil {
		t.Properties = &[]Property{prop}
		return
	}
	props := append(*t.Properties, prop)
	t.Properties = &props
}

// AddTestcase adds Testcase tc to this Testsuite.
func (t *Testsuite) AddTestcase(tc Testcase) {
	t.Testcases = append(t.Testcases, tc)
	t.Tests += 1

	if tc.Error != nil {
		t.Errors += 1
	}

	if tc.Failure != nil {
		t.Failures += 1
	}

	if tc.Skipped != nil {
		t.Skipped += 1
	}
}

// SetTimestamp sets the timestamp in this Testsuite.
func (ts *Testsuite) SetTimestamp(t time.Time) {
	ts.Timestamp = t.Format(time.RFC3339)
}

// Testcase represents a single test with its results.
type Testcase struct {
	// required attributes
	Name      string `xml:"name,attr"`
	Classname string `xml:"classname,attr"`

	// optional attributes
	Time   string `xml:"time,attr,omitempty"` // duration in seconds
	Status string `xml:"status,attr,omitempty"`

	Skipped   *Result `xml:"skipped,omitempty"`
	Error     *Result `xml:"error,omitempty"`
	Failure   *Result `xml:"failure,omitempty"`
	SystemOut *Output `xml:"system-out,omitempty"`
	SystemErr *Output `xml:"system-err,omitempty"`
}

// Property represents a key/value pair.
type Property struct {
	Name  string `xml:"name,attr"`
	Value string `xml:"value,attr"`
}

// Result represents the result of a single test.
type Result struct {
	Message string `xml:"message,attr"`
	Type    string `xml:"type,attr,omitempty"`
	Data    string `xml:",cdata"`
}

// Output represents output written to stdout or sderr.
type Output struct {
	Data string `xml:",cdata"`
}

// CreateFromReport creates a JUnit representation of the given gtr.Report.
func CreateFromReport(report gtr.Report, hostname string) Testsuites {
	var suites Testsuites
	for _, pkg := range report.Packages {
		var duration time.Duration
		suite := Testsuite{
			Name:     pkg.Name,
			Hostname: hostname,
			ID:       len(suites.Suites),
		}

		if !pkg.Timestamp.IsZero() {
			suite.SetTimestamp(pkg.Timestamp)
		}

		for k, v := range pkg.Properties {
			suite.AddProperty(k, v)
		}

		if len(pkg.Output) > 0 {
			suite.SystemOut = &Output{Data: formatOutput(pkg.Output, 0)}
		}

		if pkg.Coverage > 0 {
			suite.AddProperty("coverage.statements.pct", fmt.Sprintf("%.2f", pkg.Coverage))
		}

		for _, test := range pkg.Tests {
			duration += test.Duration
			suite.AddTestcase(createTestcaseForTest(pkg.Name, test))
		}

		for _, bm := range groupBenchmarksByName(pkg.Benchmarks) {
			suite.AddTestcase(createTestcaseForBenchmark(pkg.Name, bm))
		}

		// JUnit doesn't have a good way of dealing with build or runtime
		// errors that happen before a test has started, so we create a single
		// failing test that contains the build error details.
		if pkg.BuildError.Name != "" {
			tc := Testcase{
				Classname: pkg.BuildError.Name,
				Name:      pkg.BuildError.Cause,
				Time:      formatDuration(0),
				Error: &Result{
					Message: "Build error",
					Data:    strings.Join(pkg.BuildError.Output, "\n"),
				},
			}
			suite.AddTestcase(tc)
		}

		if pkg.RunError.Name != "" {
			tc := Testcase{
				Classname: pkg.RunError.Name,
				Name:      "Failure",
				Time:      formatDuration(0),
				Error: &Result{
					Message: "Runtime error",
					Data:    strings.Join(pkg.RunError.Output, "\n"),
				},
			}
			suite.AddTestcase(tc)
		}

		if (pkg.Duration) == 0 {
			suite.Time = formatDuration(duration)
		} else {
			suite.Time = formatDuration(pkg.Duration)
		}
		suites.AddSuite(suite)
	}
	return suites
}

func createTestcaseForTest(pkgName string, test gtr.Test) Testcase {
	tc := Testcase{
		Classname: pkgName,
		Name:      test.Name,
		Time:      formatDuration(test.Duration),
	}

	if test.Result == gtr.Fail {
		tc.Failure = &Result{
			Message: "Failed",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if test.Result == gtr.Skip {
		tc.Skipped = &Result{
			Message: "Skipped",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if test.Result == gtr.Unknown {
		tc.Error = &Result{
			Message: "No test result found",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if len(test.Output) > 0 {
		tc.SystemOut = &Output{Data: formatOutput(test.Output, test.Level)}
	}
	return tc
}

func createTestcaseForBenchmark(pkgName string, bm gtr.Benchmark) Testcase {
	tc := Testcase{
		Classname: pkgName,
		Name:      bm.Name,
		Time:      formatBenchmarkTime(time.Duration(bm.NsPerOp) * time.Duration(bm.Iterations)),
	}

	if bm.Result == gtr.Fail {
		tc.Failure = &Result{
			Message: "Failed",
		}
	} else if bm.Result == gtr.Skip {
		tc.Skipped = &Result{
			Message: "Skipped",
		}
	} else if len(bm.Output) > 0 {
		tc.SystemOut = &Output{Data: formatOutput(bm.Output, 0)}
	}
	return tc
}

// formatDuration returns the JUnit string representation of the given
// duration.
func formatDuration(d time.Duration) string {
	return fmt.Sprintf("%.3f", d.Seconds())
}

// formatBenchmarkTime returns the JUnit string representation of the given
// benchmark time.
func formatBenchmarkTime(d time.Duration) string {
	return fmt.Sprintf("%.9f", d.Seconds())
}

// formatOutput trims the test whitespace prefix from each line and joins all
// the lines.
func formatOutput(output []string, indent int) string {
	var lines []string
	for _, line := range output {
		lines = append(lines, line)
	}
	return strings.Join(lines, "\n")
}

func groupBenchmarksByName(benchmarks []gtr.Benchmark) []gtr.Benchmark {
	if len(benchmarks) == 0 {
		return nil
	}

	var grouped []gtr.Benchmark
	benchmap := make(map[string][]gtr.Benchmark)
	for _, bm := range benchmarks {
		if _, ok := benchmap[bm.Name]; !ok {
			grouped = append(grouped, gtr.Benchmark{Name: bm.Name})
		}
		benchmap[bm.Name] = append(benchmap[bm.Name], bm)
	}

	for i, bm := range grouped {
		n := 0
		for _, b := range benchmap[bm.Name] {
			if b.Result != gtr.Pass {
				continue
			}
			bm.Iterations += b.Iterations
			bm.NsPerOp += b.NsPerOp
			bm.MBPerSec += b.MBPerSec
			bm.BytesPerOp += b.BytesPerOp
			bm.AllocsPerOp += b.AllocsPerOp
			n++
		}

		bm.Result = groupResults(benchmap[bm.Name])
		if n > 0 {
			bm.NsPerOp = bm.NsPerOp / float64(n)
			bm.MBPerSec = bm.MBPerSec / float64(n)
			bm.BytesPerOp = bm.BytesPerOp / int64(n)
			bm.AllocsPerOp = bm.AllocsPerOp / int64(n)
		}
		grouped[i] = bm
	}
	return grouped
}

func groupResults(benchmarks []gtr.Benchmark) gtr.Result {
	var result gtr.Result
	for _, bm := range benchmarks {
		if bm.Result == gtr.Fail {
			return gtr.Fail
		}
		if result != gtr.Pass {
			result = bm.Result
		}
	}
	return result
}
<commit_msg>junit: Simplify the formatOutput function<commit_after>// Package junit defines a JUnit XML report and includes convenience methods
// for working with these reports.
package junit

import (
	"encoding/xml"
	"fmt"
	"strings"
	"time"

	"github.com/jstemmer/go-junit-report/v2/gtr"
)

// Testsuites is a collection of JUnit testsuites.
type Testsuites struct {
	XMLName xml.Name `xml:"testsuites"`

	Name     string `xml:"name,attr,omitempty"`
	Time     string `xml:"time,attr,omitempty"` // total duration in seconds
	Tests    int    `xml:"tests,attr,omitempty"`
	Errors   int    `xml:"errors,attr,omitempty"`
	Failures int    `xml:"failures,attr,omitempty"`
	Skipped  int    `xml:"skipped,attr,omitempty"`
	Disabled int    `xml:"disabled,attr,omitempty"`

	Suites []Testsuite `xml:"testsuite,omitempty"`
}

// AddSuite adds a Testsuite and updates this testssuites' totals.
func (t *Testsuites) AddSuite(ts Testsuite) {
	t.Suites = append(t.Suites, ts)
	t.Tests += ts.Tests
	t.Errors += ts.Errors
	t.Failures += ts.Failures
	t.Skipped += ts.Skipped
	t.Disabled += ts.Disabled
}

// Testsuite is a single JUnit testsuite containing testcases.
type Testsuite struct {
	// required attributes
	Name     string `xml:"name,attr"`
	Tests    int    `xml:"tests,attr"`
	Failures int    `xml:"failures,attr"`
	Errors   int    `xml:"errors,attr"`
	ID       int    `xml:"id,attr"`

	// optional attributes
	Disabled  int    `xml:"disabled,attr,omitempty"`
	Hostname  string `xml:"hostname,attr,omitempty"`
	Package   string `xml:"package,attr,omitempty"`
	Skipped   int    `xml:"skipped,attr,omitempty"`
	Time      string `xml:"time,attr"`                // duration in seconds
	Timestamp string `xml:"timestamp,attr,omitempty"` // date and time in ISO8601

	Properties *[]Property `xml:"properties>property,omitempty"`
	Testcases  []Testcase  `xml:"testcase,omitempty"`
	SystemOut  *Output     `xml:"system-out,omitempty"`
	SystemErr  *Output     `xml:"system-err,omitempty"`
}

// AddProperty adds a property with the given name and value to this Testsuite.
func (t *Testsuite) AddProperty(name, value string) {
	prop := Property{Name: name, Value: value}
	if t.Properties == nil {
		t.Properties = &[]Property{prop}
		return
	}
	props := append(*t.Properties, prop)
	t.Properties = &props
}

// AddTestcase adds Testcase tc to this Testsuite.
func (t *Testsuite) AddTestcase(tc Testcase) {
	t.Testcases = append(t.Testcases, tc)
	t.Tests += 1

	if tc.Error != nil {
		t.Errors += 1
	}

	if tc.Failure != nil {
		t.Failures += 1
	}

	if tc.Skipped != nil {
		t.Skipped += 1
	}
}

// SetTimestamp sets the timestamp in this Testsuite.
func (ts *Testsuite) SetTimestamp(t time.Time) {
	ts.Timestamp = t.Format(time.RFC3339)
}

// Testcase represents a single test with its results.
type Testcase struct {
	// required attributes
	Name      string `xml:"name,attr"`
	Classname string `xml:"classname,attr"`

	// optional attributes
	Time   string `xml:"time,attr,omitempty"` // duration in seconds
	Status string `xml:"status,attr,omitempty"`

	Skipped   *Result `xml:"skipped,omitempty"`
	Error     *Result `xml:"error,omitempty"`
	Failure   *Result `xml:"failure,omitempty"`
	SystemOut *Output `xml:"system-out,omitempty"`
	SystemErr *Output `xml:"system-err,omitempty"`
}

// Property represents a key/value pair.
type Property struct {
	Name  string `xml:"name,attr"`
	Value string `xml:"value,attr"`
}

// Result represents the result of a single test.
type Result struct {
	Message string `xml:"message,attr"`
	Type    string `xml:"type,attr,omitempty"`
	Data    string `xml:",cdata"`
}

// Output represents output written to stdout or sderr.
type Output struct {
	Data string `xml:",cdata"`
}

// CreateFromReport creates a JUnit representation of the given gtr.Report.
func CreateFromReport(report gtr.Report, hostname string) Testsuites {
	var suites Testsuites
	for _, pkg := range report.Packages {
		var duration time.Duration
		suite := Testsuite{
			Name:     pkg.Name,
			Hostname: hostname,
			ID:       len(suites.Suites),
		}

		if !pkg.Timestamp.IsZero() {
			suite.SetTimestamp(pkg.Timestamp)
		}

		for k, v := range pkg.Properties {
			suite.AddProperty(k, v)
		}

		if len(pkg.Output) > 0 {
			suite.SystemOut = &Output{Data: formatOutput(pkg.Output, 0)}
		}

		if pkg.Coverage > 0 {
			suite.AddProperty("coverage.statements.pct", fmt.Sprintf("%.2f", pkg.Coverage))
		}

		for _, test := range pkg.Tests {
			duration += test.Duration
			suite.AddTestcase(createTestcaseForTest(pkg.Name, test))
		}

		for _, bm := range groupBenchmarksByName(pkg.Benchmarks) {
			suite.AddTestcase(createTestcaseForBenchmark(pkg.Name, bm))
		}

		// JUnit doesn't have a good way of dealing with build or runtime
		// errors that happen before a test has started, so we create a single
		// failing test that contains the build error details.
		if pkg.BuildError.Name != "" {
			tc := Testcase{
				Classname: pkg.BuildError.Name,
				Name:      pkg.BuildError.Cause,
				Time:      formatDuration(0),
				Error: &Result{
					Message: "Build error",
					Data:    strings.Join(pkg.BuildError.Output, "\n"),
				},
			}
			suite.AddTestcase(tc)
		}

		if pkg.RunError.Name != "" {
			tc := Testcase{
				Classname: pkg.RunError.Name,
				Name:      "Failure",
				Time:      formatDuration(0),
				Error: &Result{
					Message: "Runtime error",
					Data:    strings.Join(pkg.RunError.Output, "\n"),
				},
			}
			suite.AddTestcase(tc)
		}

		if (pkg.Duration) == 0 {
			suite.Time = formatDuration(duration)
		} else {
			suite.Time = formatDuration(pkg.Duration)
		}
		suites.AddSuite(suite)
	}
	return suites
}

func createTestcaseForTest(pkgName string, test gtr.Test) Testcase {
	tc := Testcase{
		Classname: pkgName,
		Name:      test.Name,
		Time:      formatDuration(test.Duration),
	}

	if test.Result == gtr.Fail {
		tc.Failure = &Result{
			Message: "Failed",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if test.Result == gtr.Skip {
		tc.Skipped = &Result{
			Message: "Skipped",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if test.Result == gtr.Unknown {
		tc.Error = &Result{
			Message: "No test result found",
			Data:    formatOutput(test.Output, test.Level),
		}
	} else if len(test.Output) > 0 {
		tc.SystemOut = &Output{Data: formatOutput(test.Output, test.Level)}
	}
	return tc
}

func createTestcaseForBenchmark(pkgName string, bm gtr.Benchmark) Testcase {
	tc := Testcase{
		Classname: pkgName,
		Name:      bm.Name,
		Time:      formatBenchmarkTime(time.Duration(bm.NsPerOp) * time.Duration(bm.Iterations)),
	}

	if bm.Result == gtr.Fail {
		tc.Failure = &Result{
			Message: "Failed",
		}
	} else if bm.Result == gtr.Skip {
		tc.Skipped = &Result{
			Message: "Skipped",
		}
	} else if len(bm.Output) > 0 {
		tc.SystemOut = &Output{Data: formatOutput(bm.Output, 0)}
	}
	return tc
}

// formatDuration returns the JUnit string representation of the given
// duration.
func formatDuration(d time.Duration) string {
	return fmt.Sprintf("%.3f", d.Seconds())
}

// formatBenchmarkTime returns the JUnit string representation of the given
// benchmark time.
func formatBenchmarkTime(d time.Duration) string {
	return fmt.Sprintf("%.9f", d.Seconds())
}

// formatOutput combines the lines from the given output into a single string.
func formatOutput(output []string, indent int) string {
	return strings.Join(output, "\n")
}

func groupBenchmarksByName(benchmarks []gtr.Benchmark) []gtr.Benchmark {
	if len(benchmarks) == 0 {
		return nil
	}

	var grouped []gtr.Benchmark
	benchmap := make(map[string][]gtr.Benchmark)
	for _, bm := range benchmarks {
		if _, ok := benchmap[bm.Name]; !ok {
			grouped = append(grouped, gtr.Benchmark{Name: bm.Name})
		}
		benchmap[bm.Name] = append(benchmap[bm.Name], bm)
	}

	for i, bm := range grouped {
		n := 0
		for _, b := range benchmap[bm.Name] {
			if b.Result != gtr.Pass {
				continue
			}
			bm.Iterations += b.Iterations
			bm.NsPerOp += b.NsPerOp
			bm.MBPerSec += b.MBPerSec
			bm.BytesPerOp += b.BytesPerOp
			bm.AllocsPerOp += b.AllocsPerOp
			n++
		}

		bm.Result = groupResults(benchmap[bm.Name])
		if n > 0 {
			bm.NsPerOp = bm.NsPerOp / float64(n)
			bm.MBPerSec = bm.MBPerSec / float64(n)
			bm.BytesPerOp = bm.BytesPerOp / int64(n)
			bm.AllocsPerOp = bm.AllocsPerOp / int64(n)
		}
		grouped[i] = bm
	}
	return grouped
}

func groupResults(benchmarks []gtr.Benchmark) gtr.Result {
	var result gtr.Result
	for _, bm := range benchmarks {
		if bm.Result == gtr.Fail {
			return gtr.Fail
		}
		if result != gtr.Pass {
			result = bm.Result
		}
	}
	return result
}
<|endoftext|>