<commit_before>package mdns

// Advertise network services via multicast DNS

import (
	"log"
	"net"

	"github.com/davecheney/gmx"
	"github.com/miekg/dns"
)

var (
	ipv4mcastaddr = &net.UDPAddr{
		IP:   net.ParseIP("224.0.0.251"),
		Port: 5353,
	}

	ipv6mcastaddr = &net.UDPAddr{
		IP:   net.ParseIP("ff02::fb"),
		Port: 5353,
	}
	local *zone // the local mdns zone
)

func init() {
	local = &zone{
		entries: make(map[string]entries),
		add:     make(chan *entry),
		queries: make(chan *query, 16),
	}
	go local.mainloop()
	if err := local.listen(ipv4mcastaddr); err != nil {
		log.Fatalf("Failed to listen %s: %s", ipv4mcastaddr, err)
	}
	if err := local.listen(ipv6mcastaddr); err != nil {
		log.Printf("Failed to listen %s: %s", ipv6mcastaddr, err)
	}

	// publish gmx stats for the local zone
	gmx.Publish("mdns.zone.local.queries", func() interface{} {
		return local.queryCount
	})
	gmx.Publish("mdns.zone.local.entries", func() interface{} {
		return local.entryCount
	})

}

// Publish adds a record, described in RFC XXX
func Publish(r string) error {
	rr, err := dns.NewRR(r)
	if err != nil {
		return err
	}
	local.add <- &entry{rr}
	return nil
}

type entry struct {
	dns.RR
}

func (e *entry) fqdn() string {
	return e.Header().Name
}

type query struct {
	dns.Question
	result chan *entry
}

type entries []*entry

func (e entries) contains(entry *entry) bool {
	for _, ee := range e {
		if equals(entry, ee) {
			return true
		}
	}
	return false
}

type zone struct {
	entries map[string]entries
	add     chan *entry // add entries to zone
	queries chan *query // query exsting entries in zone

	// gmx statistics
	queryCount, entryCount int
}

func (z *zone) mainloop() {
	for {
		select {
		case entry := <-z.add:
			z.entryCount++
			if !z.entries[entry.fqdn()].contains(entry) {
				z.entries[entry.fqdn()] = append(z.entries[entry.fqdn()], entry)
			}
		case q := <-z.queries:
			z.queryCount++
			for _, entry := range z.entries[q.Question.Name] {
				if q.matches(entry) {
					q.result <- entry
				}
			}
			close(q.result)
		}
	}
}

func (z *zone) query(q dns.Question) (entries []*entry) {
	res := make(chan *entry, 16)
	z.queries <- &query{q, res}
	for e := range res {
		entries = append(entries, e)
	}
	return
}

func (q *query) matches(entry *entry) bool {
	return q.Question.Qtype == dns.TypeANY || q.Question.Qtype == entry.RR.Header().Rrtype
}

func equals(this, that *entry) bool {
	if _, ok := this.RR.(*dns.ANY); ok {
		return true // *ANY matches anything
	}
	if _, ok := that.RR.(*dns.ANY); ok {
		return true // *ANY matches all
	}
	return false
}

type connector struct {
	*net.UDPAddr
	*net.UDPConn
	*zone

	// gmx statistics
	questions, responses int
}

func (z *zone) listen(addr *net.UDPAddr) error {
	conn, err := openSocket(addr)
	if err != nil {
		return err
	}
	c := &connector{
		UDPAddr: addr,
		UDPConn: conn,
		zone:    z,
	}
	go c.mainloop()

	// publish gmx stats for the connector
	gmx.Publish("mdns.connector.questions", func() interface{} {
		return c.questions
	})
	gmx.Publish("mdns.connector.responses", func() interface{} {
		return c.responses
	})

	return nil
}

func openSocket(addr *net.UDPAddr) (*net.UDPConn, error) {
	switch addr.IP.To4() {
	case nil:
		return net.ListenMulticastUDP("udp6", nil, ipv6mcastaddr)
	default:
		return net.ListenMulticastUDP("udp4", nil, ipv4mcastaddr)
	}
	panic("unreachable")
}

type pkt struct {
	*dns.Msg
	*net.UDPAddr
}

func (c *connector) readloop(in chan pkt) {
	for {
		msg, addr, err := c.readMessage()
		if err != nil {
			// log dud packets
			log.Printf("Cound not read from %s: %s", c.UDPConn, err)
		}
		if len(msg.Question) > 0 {
			c.questions++
			in <- pkt{msg, addr}
		}
	}
}

func (c *connector) mainloop() {
	in := make(chan pkt, 32)
	go c.readloop(in)
	for {
		msg := <-in
		msg.MsgHdr.Response = true // convert question to response
		for _, result := range c.query(msg.Question) {
			msg.Answer = append(msg.Answer, result.RR)
		}
		msg.Extra = append(msg.Extra, c.findExtra(msg.Answer...)...)
		if len(msg.Answer) > 0 {
			c.responses++
			// nuke questions
			msg.Question = nil
			if err := c.writeMessage(msg.Msg, msg.UDPAddr); err != nil {
				log.Fatalf("Cannot send: %s", err)
			}
		}
	}
}

func (c *connector) query(qs []dns.Question) (results []*entry) {
	for _, q := range qs {
		results = append(results, c.zone.query(q)...)
	}
	return
}

// recursively probe for related records
func (c *connector) findExtra(r ...dns.RR) (extra []dns.RR) {
	for _, rr := range r {
		var q dns.Question
		switch rr := rr.(type) {
		case *dns.PTR:
			q = dns.Question{
				Name:   rr.Ptr,
				Qtype:  dns.TypeANY,
				Qclass: dns.ClassINET,
			}
		case *dns.SRV:
			q = dns.Question{
				Name:   rr.Target,
				Qtype:  dns.TypeA,
				Qclass: dns.ClassINET,
			}
		default:
			continue
		}
		res := c.zone.query(q)
		if len(res) > 0 {
			for _, entry := range res {
				extra = append(append(extra, entry.RR), c.findExtra(entry.RR)...)
			}
		}
	}
	return
}

// encode an mdns msg and broadcast it on the wire
func (c *connector) writeMessage(msg *dns.Msg, addr *net.UDPAddr) error {
	buf, err := msg.Pack()
	if err != nil {
		return err
	}
	_, err = c.WriteToUDP(buf, addr)
	return err
}

// consume an mdns packet from the wire and decode it
func (c *connector) readMessage() (*dns.Msg, *net.UDPAddr, error) {
	buf := make([]byte, 1500)
	read, addr, err := c.ReadFromUDP(buf)
	if err != nil {
		return nil, nil, err
	}
	var msg dns.Msg
	if err := msg.Unpack(buf[:read]); err != nil {
		return nil, nil, err
	}
	return &msg, addr, nil
}
<commit_msg>Handle nil reference to dns message<commit_after>package mdns

// Advertise network services via multicast DNS

import (
	"log"
	"net"

	"github.com/davecheney/gmx"
	"github.com/miekg/dns"
)

var (
	ipv4mcastaddr = &net.UDPAddr{
		IP:   net.ParseIP("224.0.0.251"),
		Port: 5353,
	}

	ipv6mcastaddr = &net.UDPAddr{
		IP:   net.ParseIP("ff02::fb"),
		Port: 5353,
	}
	local *zone // the local mdns zone
)

func init() {
	local = &zone{
		entries: make(map[string]entries),
		add:     make(chan *entry),
		queries: make(chan *query, 16),
	}
	go local.mainloop()
	if err := local.listen(ipv4mcastaddr); err != nil {
		log.Fatalf("Failed to listen %s: %s", ipv4mcastaddr, err)
	}
	if err := local.listen(ipv6mcastaddr); err != nil {
		log.Printf("Failed to listen %s: %s", ipv6mcastaddr, err)
	}

	// publish gmx stats for the local zone
	gmx.Publish("mdns.zone.local.queries", func() interface{} {
		return local.queryCount
	})
	gmx.Publish("mdns.zone.local.entries", func() interface{} {
		return local.entryCount
	})

}

// Publish adds a record, described in RFC XXX
func Publish(r string) error {
	rr, err := dns.NewRR(r)
	if err != nil {
		return err
	}
	local.add <- &entry{rr}
	return nil
}

type entry struct {
	dns.RR
}

func (e *entry) fqdn() string {
	return e.Header().Name
}

type query struct {
	dns.Question
	result chan *entry
}

type entries []*entry

func (e entries) contains(entry *entry) bool {
	for _, ee := range e {
		if equals(entry, ee) {
			return true
		}
	}
	return false
}

type zone struct {
	entries map[string]entries
	add     chan *entry // add entries to zone
	queries chan *query // query exsting entries in zone

	// gmx statistics
	queryCount, entryCount int
}

func (z *zone) mainloop() {
	for {
		select {
		case entry := <-z.add:
			z.entryCount++
			if !z.entries[entry.fqdn()].contains(entry) {
				z.entries[entry.fqdn()] = append(z.entries[entry.fqdn()], entry)
			}
		case q := <-z.queries:
			z.queryCount++
			for _, entry := range z.entries[q.Question.Name] {
				if q.matches(entry) {
					q.result <- entry
				}
			}
			close(q.result)
		}
	}
}

func (z *zone) query(q dns.Question) (entries []*entry) {
	res := make(chan *entry, 16)
	z.queries <- &query{q, res}
	for e := range res {
		entries = append(entries, e)
	}
	return
}

func (q *query) matches(entry *entry) bool {
	return q.Question.Qtype == dns.TypeANY || q.Question.Qtype == entry.RR.Header().Rrtype
}

func equals(this, that *entry) bool {
	if _, ok := this.RR.(*dns.ANY); ok {
		return true // *ANY matches anything
	}
	if _, ok := that.RR.(*dns.ANY); ok {
		return true // *ANY matches all
	}
	return false
}

type connector struct {
	*net.UDPAddr
	*net.UDPConn
	*zone

	// gmx statistics
	questions, responses int
}

func (z *zone) listen(addr *net.UDPAddr) error {
	conn, err := openSocket(addr)
	if err != nil {
		return err
	}
	c := &connector{
		UDPAddr: addr,
		UDPConn: conn,
		zone:    z,
	}
	go c.mainloop()

	// publish gmx stats for the connector
	gmx.Publish("mdns.connector.questions", func() interface{} {
		return c.questions
	})
	gmx.Publish("mdns.connector.responses", func() interface{} {
		return c.responses
	})

	return nil
}

func openSocket(addr *net.UDPAddr) (*net.UDPConn, error) {
	switch addr.IP.To4() {
	case nil:
		return net.ListenMulticastUDP("udp6", nil, ipv6mcastaddr)
	default:
		return net.ListenMulticastUDP("udp4", nil, ipv4mcastaddr)
	}
	panic("unreachable")
}

type pkt struct {
	*dns.Msg
	*net.UDPAddr
}

func (c *connector) readloop(in chan pkt) {
	for {
		msg, addr, err := c.readMessage()
		if err != nil {
			// log dud packets
			log.Printf("Could not read from %s: %s", c.UDPConn, err)
			continue
		}
		if len(msg.Question) > 0 {
			c.questions++
			in <- pkt{msg, addr}
		}
	}
}

func (c *connector) mainloop() {
	in := make(chan pkt, 32)
	go c.readloop(in)
	for {
		msg := <-in
		msg.MsgHdr.Response = true // convert question to response
		for _, result := range c.query(msg.Question) {
			msg.Answer = append(msg.Answer, result.RR)
		}
		msg.Extra = append(msg.Extra, c.findExtra(msg.Answer...)...)
		if len(msg.Answer) > 0 {
			c.responses++
			// nuke questions
			msg.Question = nil
			if err := c.writeMessage(msg.Msg, msg.UDPAddr); err != nil {
				log.Fatalf("Cannot send: %s", err)
			}
		}
	}
}

func (c *connector) query(qs []dns.Question) (results []*entry) {
	for _, q := range qs {
		results = append(results, c.zone.query(q)...)
	}
	return
}

// recursively probe for related records
func (c *connector) findExtra(r ...dns.RR) (extra []dns.RR) {
	for _, rr := range r {
		var q dns.Question
		switch rr := rr.(type) {
		case *dns.PTR:
			q = dns.Question{
				Name:   rr.Ptr,
				Qtype:  dns.TypeANY,
				Qclass: dns.ClassINET,
			}
		case *dns.SRV:
			q = dns.Question{
				Name:   rr.Target,
				Qtype:  dns.TypeA,
				Qclass: dns.ClassINET,
			}
		default:
			continue
		}
		res := c.zone.query(q)
		if len(res) > 0 {
			for _, entry := range res {
				extra = append(append(extra, entry.RR), c.findExtra(entry.RR)...)
			}
		}
	}
	return
}

// encode an mdns msg and broadcast it on the wire
func (c *connector) writeMessage(msg *dns.Msg, addr *net.UDPAddr) error {
	buf, err := msg.Pack()
	if err != nil {
		return err
	}
	_, err = c.WriteToUDP(buf, addr)
	return err
}

// consume an mdns packet from the wire and decode it
func (c *connector) readMessage() (*dns.Msg, *net.UDPAddr, error) {
	buf := make([]byte, 1500)
	read, addr, err := c.ReadFromUDP(buf)
	if err != nil {
		return nil, nil, err
	}
	var msg dns.Msg
	if err := msg.Unpack(buf[:read]); err != nil {
		return nil, nil, err
	}
	return &msg, addr, nil
}
<|endoftext|>


<commit_before>// Package meta provides signals and slots for the Observer pattern.
//
// Note: This is a proof of concept package, therefore you are not
// supposed to use it in production or any near. Public API of this
// package is likely to change (a lot).
//
// Signals and slots is a language construct introduced in Qt for
// communication between objects which makes it easy to implement
// the Observer pattern while avoiding boilerplate code. The concept
// is that GUI widgets can send signals containing event information
// which can be received by other controls using special functions,
// known as slots.
//
// Once a signal gets emitted, all the slots connected are being
// executed in the distinct goroutines. Value passed to the signal
// is available from the call context object.
package meta

// Signal is an adorable piece of a structure!
type Signal struct {
	cons []con

	// gets incremented on every connect
	conID Connection
}

type con struct {
	slot Slot
	cid  Connection
}

// Slot is a reciver function, usually a wrapper.
type Slot func(*Call)

// Call stands for the call context, handled in the slot.
type Call struct {
	Data interface{}
}

// Connection is a signal-slot connection descriptor, unique
// within the lifetime of the signal.
//
// You may use it to terminate existing connections.
type Connection int

// Connect attaches a new slot to the signal. It also does
// panic if any of the params given is equal to nil.
func Connect(sig *Signal, slot Slot) Connection {
	if sig == nil {
		panic("can't connect a nil signal")
	}

	if slot == nil {
		panic("can't connect a nil slot")
	}

	sig.conID++

	newConnection := con{
		slot: slot,
		cid:  sig.conID,
	}

	sig.cons = append(sig.cons, newConnection)

	return sig.conID
}

// Disconnect does the opposite to connect.
func Disconnect(sig *Signal, cid Connection) {
	if sig == nil {
		panic("can't disconnect a nil signal")
	}

	if cid <= 0 {
		panic("invalid connection id")
	}

	for i := 0; i < len(sig.cons); i++ {
		if sig.cons[i].cid == cid {
			sig.cons = append(sig.cons[:i], sig.cons[i+1:]...)
			return
		}
	}
}

// Emit executes all the connected slots with data given.
func (sig *Signal) Emit(data interface{}) {
	for i := 0; i < len(sig.cons); i++ {
		sig.cons[i].slot(&Call{Data: data})
	}
}
<commit_msg>I'm silly.. and dumb. This resolves #1<commit_after>// Package meta provides signals and slots for the Observer pattern.
//
// Note: This is a proof of concept package, therefore you are not
// supposed to use it in production or any near. Public API of this
// package is likely to change (a lot).
//
// Signals and slots is a language construct introduced in Qt for
// communication between objects which makes it easy to implement
// the Observer pattern while avoiding boilerplate code. The concept
// is that GUI widgets can send signals containing event information
// which can be received by other controls using special functions,
// known as slots.
//
// Once a signal gets emitted, all the slots connected are being
// executed in the distinct goroutines. Value passed to the signal
// is available from the call context object.
package meta

// Signal is an adorable piece of a structure!
type Signal struct {
	cons []con

	// gets incremented on every connect
	conID Connection
}

type con struct {
	slot Slot
	cid  Connection
}

// Slot is a reciver function, usually a wrapper.
type Slot func(*Call)

// Call stands for the call context, handled in the slot.
type Call struct {
	Data interface{}
}

// Connection is a signal-slot connection descriptor, unique
// within the lifetime of the signal.
//
// You may use it to terminate existing connections.
type Connection int

// Connect attaches a new slot to the signal. It also does
// panic if any of the params given is equal to nil.
func Connect(sig *Signal, slot Slot) Connection {
	if sig == nil {
		panic("can't connect a nil signal")
	}

	if slot == nil {
		panic("can't connect a nil slot")
	}

	sig.conID++

	newConnection := con{
		slot: slot,
		cid:  sig.conID,
	}

	sig.cons = append(sig.cons, newConnection)

	return sig.conID
}

// Disconnect does the opposite to connect.
func Disconnect(sig *Signal, cid Connection) {
	if sig == nil {
		panic("can't disconnect a nil signal")
	}

	if cid <= 0 {
		panic("invalid connection id")
	}

	for i := 0; i < len(sig.cons); i++ {
		if sig.cons[i].cid == cid {
			sig.cons = append(sig.cons[:i], sig.cons[i+1:]...)
			return
		}
	}
}

// Emit executes all the connected slots with data given.
func (sig *Signal) Emit(data interface{}) {
	for i := 0; i < len(sig.cons); i++ {
		go sig.cons[i].slot(&Call{Data: data})
	}
}
<|endoftext|>


<commit_before>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap specifies a function for optionally wrapping upstream for
	// inspecting the decrypted HTTP request and response.
	Wrap func(upstream http.Handler) http.Handler

	// CA specifies the root CA for generating leaf certs for each incoming
	// TLS request.
	CA *tls.Certificate

	// TLSServerConfig specifies the tls.Config to use when generating leaf
	// cert using CA.
	TLSServerConfig *tls.Config

	// TLSServerConfig specifies the tls.Config to use when establishing a
	// downstream connection for proxying.
	TLSClientConfig *tls.Config

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		name  = dnsName(r.Host)
	)

	if name == "" {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(name)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert(hello.ServerName)
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.Dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names ...string) (*tls.Certificate, error) {
	return genCert(p.CA, names)
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

// handshake hijacks w's underlying net.Conn, responds to the CONNECT request
// and manually performs the TLS handshake. It returns the net.Conn or and
// error if any.
func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ""
	}
	return host
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

// A oneShotDialer implements net.Dialer whos Dial only returns a
// net.Conn as specified by c followed by an error for each subsequent Dial.
type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) Dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

// A oneShotListener implements net.Listener whos Accept only returns a
// net.Conn as specified by c followed by an error for each subsequent Accept.
type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

// A onCloseConn implements net.Conn and calls its f on Close.
type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<commit_msg>Add caching for certificates, makes things much faster.<commit_after>package mitm

import (
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"sync"
	"time"
)

// Proxy is a forward proxy that substitutes its own certificate
// for incoming TLS connections in place of the upstream server's
// certificate.
type Proxy struct {
	// Wrap specifies a function for optionally wrapping upstream for
	// inspecting the decrypted HTTP request and response.
	Wrap func(upstream http.Handler) http.Handler

	// CA specifies the root CA for generating leaf certs for each incoming
	// TLS request.
	CA *tls.Certificate

	// TLSServerConfig specifies the tls.Config to use when generating leaf
	// cert using CA.
	TLSServerConfig *tls.Config

	// TLSServerConfig specifies the tls.Config to use when establishing a
	// downstream connection for proxying.
	TLSClientConfig *tls.Config

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "CONNECT" {
		p.serveConnect(w, r)
		return
	}
	rp := &httputil.ReverseProxy{
		Director:      httpDirector,
		FlushInterval: p.FlushInterval,
	}
	p.Wrap(rp).ServeHTTP(w, r)
}

func (p *Proxy) serveConnect(w http.ResponseWriter, r *http.Request) {
	var (
		err   error
		sconn *tls.Conn
		name  = dnsName(r.Host)
	)

	if name == "" {
		log.Println("cannot determine cert name for " + r.Host)
		http.Error(w, "no upstream", 503)
		return
	}

	provisionalCert, err := p.cert(name)
	if err != nil {
		log.Println("cert", err)
		http.Error(w, "no upstream", 503)
		return
	}

	sConfig := new(tls.Config)
	if p.TLSServerConfig != nil {
		*sConfig = *p.TLSServerConfig
	}
	sConfig.Certificates = []tls.Certificate{*provisionalCert}
	sConfig.GetCertificate = func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		cConfig := new(tls.Config)
		if p.TLSClientConfig != nil {
			*cConfig = *p.TLSClientConfig
		}
		cConfig.ServerName = hello.ServerName
		sconn, err = tls.Dial("tcp", r.Host, cConfig)
		if err != nil {
			log.Println("dial", r.Host, err)
			return nil, err
		}
		return p.cert(hello.ServerName)
	}

	cconn, err := handshake(w, sConfig)
	if err != nil {
		log.Println("handshake", r.Host, err)
		return
	}
	defer cconn.Close()
	if sconn == nil {
		log.Println("could not determine cert name for " + r.Host)
		return
	}
	defer sconn.Close()

	od := &oneShotDialer{c: sconn}
	rp := &httputil.ReverseProxy{
		Director:      httpsDirector,
		Transport:     &http.Transport{DialTLS: od.Dial},
		FlushInterval: p.FlushInterval,
	}

	ch := make(chan int)
	wc := &onCloseConn{cconn, func() { ch <- 0 }}
	http.Serve(&oneShotListener{wc}, p.Wrap(rp))
	<-ch
}

func (p *Proxy) cert(names ...string) (*tls.Certificate, error) {
	var (
		cache = map[string]tls.Certificate{} // names (joined with ,) -> cert
		m     sync.Mutex
	)

	m.Lock()
	defer m.Unlock()

	key := strings.Join(names, ",")

	if cert, ok := cache[key]; ok {
		return &cert, nil
	} else {
		cert, err := genCert(p.CA, names)

		if err == nil {
			cache[key] = *cert
		}

		return cert, err
	}
}

var okHeader = []byte("HTTP/1.1 200 OK\r\n\r\n")

// handshake hijacks w's underlying net.Conn, responds to the CONNECT request
// and manually performs the TLS handshake. It returns the net.Conn or and
// error if any.
func handshake(w http.ResponseWriter, config *tls.Config) (net.Conn, error) {
	raw, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		http.Error(w, "no upstream", 503)
		return nil, err
	}
	if _, err = raw.Write(okHeader); err != nil {
		raw.Close()
		return nil, err
	}
	conn := tls.Server(raw, config)
	err = conn.Handshake()
	if err != nil {
		conn.Close()
		raw.Close()
		return nil, err
	}
	return conn, nil
}

func httpDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "http"
}

func httpsDirector(r *http.Request) {
	r.URL.Host = r.Host
	r.URL.Scheme = "https"
}

// dnsName returns the DNS name in addr, if any.
func dnsName(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return ""
	}
	return host
}

// namesOnCert returns the dns names
// in the peer's presented cert.
func namesOnCert(conn *tls.Conn) []string {
	// TODO(kr): handle IP addr SANs.
	c := conn.ConnectionState().PeerCertificates[0]
	if len(c.DNSNames) > 0 {
		// If Subject Alt Name is given,
		// we ignore the common name.
		// This matches behavior of crypto/x509.
		return c.DNSNames
	}
	return []string{c.Subject.CommonName}
}

// A oneShotDialer implements net.Dialer whos Dial only returns a
// net.Conn as specified by c followed by an error for each subsequent Dial.
type oneShotDialer struct {
	c  net.Conn
	mu sync.Mutex
}

func (d *oneShotDialer) Dial(network, addr string) (net.Conn, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.c == nil {
		return nil, errors.New("closed")
	}
	c := d.c
	d.c = nil
	return c, nil
}

// A oneShotListener implements net.Listener whos Accept only returns a
// net.Conn as specified by c followed by an error for each subsequent Accept.
type oneShotListener struct {
	c net.Conn
}

func (l *oneShotListener) Accept() (net.Conn, error) {
	if l.c == nil {
		return nil, errors.New("closed")
	}
	c := l.c
	l.c = nil
	return c, nil
}

func (l *oneShotListener) Close() error {
	return nil
}

func (l *oneShotListener) Addr() net.Addr {
	return l.c.LocalAddr()
}

// A onCloseConn implements net.Conn and calls its f on Close.
type onCloseConn struct {
	net.Conn
	f func()
}

func (c *onCloseConn) Close() error {
	if c.f != nil {
		c.f()
		c.f = nil
	}
	return c.Conn.Close()
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package context

import (
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/auth"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"gitea.com/macaron/csrf"
	"gitea.com/macaron/macaron"
)

// ToggleOptions contains required or check options
type ToggleOptions struct {
	SignInRequired  bool
	SignOutRequired bool
	AdminRequired   bool
	DisableCSRF     bool
}

// Toggle returns toggle options as middleware
func Toggle(options *ToggleOptions) macaron.Handler {
	return func(ctx *Context) {
		// Cannot view any page before installation.
		if !setting.InstallLock {
			ctx.Redirect(setting.AppSubURL + "/install")
			return
		}

		// Check prohibit login users.
		if ctx.IsSigned {
			if !ctx.User.IsActive && setting.Service.RegisterEmailConfirm {
				ctx.Data["Title"] = ctx.Tr("auth.active_your_account")
				ctx.HTML(200, "user/auth/activate")
				return
			} else if !ctx.User.IsActive || ctx.User.ProhibitLogin {
				log.Info("Failed authentication attempt for %s from %s", ctx.User.Name, ctx.RemoteAddr())
				ctx.Data["Title"] = ctx.Tr("auth.prohibit_login")
				ctx.HTML(200, "user/auth/prohibit_login")
				return
			}

			if ctx.User.MustChangePassword {
				if ctx.Req.URL.Path != "/user/settings/change_password" {
					ctx.Data["Title"] = ctx.Tr("auth.must_change_password")
					ctx.Data["ChangePasscodeLink"] = setting.AppSubURL + "/user/change_password"
					if ctx.Req.URL.Path != "/user/events" {
						ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
					}
					ctx.Redirect(setting.AppSubURL + "/user/settings/change_password")
					return
				}
			} else if ctx.Req.URL.Path == "/user/settings/change_password" {
				// make sure that the form cannot be accessed by users who don't need this
				ctx.Redirect(setting.AppSubURL + "/")
				return
			}
		}

		// Redirect to dashboard if user tries to visit any non-login page.
		if options.SignOutRequired && ctx.IsSigned && ctx.Req.URL.RequestURI() != "/" {
			ctx.Redirect(setting.AppSubURL + "/")
			return
		}

		if !options.SignOutRequired && !options.DisableCSRF && ctx.Req.Method == "POST" && !auth.IsAPIPath(ctx.Req.URL.Path) {
			csrf.Validate(ctx.Context, ctx.csrf)
			if ctx.Written() {
				return
			}
		}

		if options.SignInRequired {
			if !ctx.IsSigned {
				// Restrict API calls with error message.
				if auth.IsAPIPath(ctx.Req.URL.Path) {
					ctx.JSON(403, map[string]string{
						"message": "Only signed in user is allowed to call APIs.",
					})
					return
				}

				ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
				ctx.Redirect(setting.AppSubURL + "/user/login")
				return
			} else if !ctx.User.IsActive && setting.Service.RegisterEmailConfirm {
				ctx.Data["Title"] = ctx.Tr("auth.active_your_account")
				ctx.HTML(200, "user/auth/activate")
				return
			}
			if ctx.IsSigned && auth.IsAPIPath(ctx.Req.URL.Path) && ctx.IsBasicAuth {
				twofa, err := models.GetTwoFactorByUID(ctx.User.ID)
				if err != nil {
					if models.IsErrTwoFactorNotEnrolled(err) {
						return // No 2FA enrollment for this user
					}
					ctx.Error(500)
					return
				}
				otpHeader := ctx.Req.Header.Get("X-Gitea-OTP")
				ok, err := twofa.ValidateTOTP(otpHeader)
				if err != nil {
					ctx.Error(500)
					return
				}
				if !ok {
					ctx.JSON(403, map[string]string{
						"message": "Only signed in user is allowed to call APIs.",
					})
					return
				}
			}
		}

		// Redirect to log in page if auto-signin info is provided and has not signed in.
		if !options.SignOutRequired && !ctx.IsSigned && !auth.IsAPIPath(ctx.Req.URL.Path) &&
			len(ctx.GetCookie(setting.CookieUserName)) > 0 {
			ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
			ctx.Redirect(setting.AppSubURL + "/user/login")
			return
		}

		if options.AdminRequired {
			if !ctx.User.IsAdmin {
				ctx.Error(403)
				return
			}
			ctx.Data["PageIsAdmin"] = true
		}
	}
}
<commit_msg>Prevent redirect back to /user/events (#12462)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2019 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package context

import (
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/auth"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"gitea.com/macaron/csrf"
	"gitea.com/macaron/macaron"
)

// ToggleOptions contains required or check options
type ToggleOptions struct {
	SignInRequired  bool
	SignOutRequired bool
	AdminRequired   bool
	DisableCSRF     bool
}

// Toggle returns toggle options as middleware
func Toggle(options *ToggleOptions) macaron.Handler {
	return func(ctx *Context) {
		// Cannot view any page before installation.
		if !setting.InstallLock {
			ctx.Redirect(setting.AppSubURL + "/install")
			return
		}

		// Check prohibit login users.
		if ctx.IsSigned {
			if !ctx.User.IsActive && setting.Service.RegisterEmailConfirm {
				ctx.Data["Title"] = ctx.Tr("auth.active_your_account")
				ctx.HTML(200, "user/auth/activate")
				return
			} else if !ctx.User.IsActive || ctx.User.ProhibitLogin {
				log.Info("Failed authentication attempt for %s from %s", ctx.User.Name, ctx.RemoteAddr())
				ctx.Data["Title"] = ctx.Tr("auth.prohibit_login")
				ctx.HTML(200, "user/auth/prohibit_login")
				return
			}

			if ctx.User.MustChangePassword {
				if ctx.Req.URL.Path != "/user/settings/change_password" {
					ctx.Data["Title"] = ctx.Tr("auth.must_change_password")
					ctx.Data["ChangePasscodeLink"] = setting.AppSubURL + "/user/change_password"
					if ctx.Req.URL.Path != "/user/events" {
						ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
					}
					ctx.Redirect(setting.AppSubURL + "/user/settings/change_password")
					return
				}
			} else if ctx.Req.URL.Path == "/user/settings/change_password" {
				// make sure that the form cannot be accessed by users who don't need this
				ctx.Redirect(setting.AppSubURL + "/")
				return
			}
		}

		// Redirect to dashboard if user tries to visit any non-login page.
		if options.SignOutRequired && ctx.IsSigned && ctx.Req.URL.RequestURI() != "/" {
			ctx.Redirect(setting.AppSubURL + "/")
			return
		}

		if !options.SignOutRequired && !options.DisableCSRF && ctx.Req.Method == "POST" && !auth.IsAPIPath(ctx.Req.URL.Path) {
			csrf.Validate(ctx.Context, ctx.csrf)
			if ctx.Written() {
				return
			}
		}

		if options.SignInRequired {
			if !ctx.IsSigned {
				// Restrict API calls with error message.
				if auth.IsAPIPath(ctx.Req.URL.Path) {
					ctx.JSON(403, map[string]string{
						"message": "Only signed in user is allowed to call APIs.",
					})
					return
				}
				if ctx.Req.URL.Path != "/user/events" {
					ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
				}
				ctx.Redirect(setting.AppSubURL + "/user/login")
				return
			} else if !ctx.User.IsActive && setting.Service.RegisterEmailConfirm {
				ctx.Data["Title"] = ctx.Tr("auth.active_your_account")
				ctx.HTML(200, "user/auth/activate")
				return
			}
			if ctx.IsSigned && auth.IsAPIPath(ctx.Req.URL.Path) && ctx.IsBasicAuth {
				twofa, err := models.GetTwoFactorByUID(ctx.User.ID)
				if err != nil {
					if models.IsErrTwoFactorNotEnrolled(err) {
						return // No 2FA enrollment for this user
					}
					ctx.Error(500)
					return
				}
				otpHeader := ctx.Req.Header.Get("X-Gitea-OTP")
				ok, err := twofa.ValidateTOTP(otpHeader)
				if err != nil {
					ctx.Error(500)
					return
				}
				if !ok {
					ctx.JSON(403, map[string]string{
						"message": "Only signed in user is allowed to call APIs.",
					})
					return
				}
			}
		}

		// Redirect to log in page if auto-signin info is provided and has not signed in.
		if !options.SignOutRequired && !ctx.IsSigned && !auth.IsAPIPath(ctx.Req.URL.Path) &&
			len(ctx.GetCookie(setting.CookieUserName)) > 0 {
			if ctx.Req.URL.Path != "/user/events" {
				ctx.SetCookie("redirect_to", setting.AppSubURL+ctx.Req.URL.RequestURI(), 0, setting.AppSubURL)
			}
			ctx.Redirect(setting.AppSubURL + "/user/login")
			return
		}

		if options.AdminRequired {
			if !ctx.User.IsAdmin {
				ctx.Error(403)
				return
			}
			ctx.Data["PageIsAdmin"] = true
		}
	}
}
<|endoftext|>


<commit_before>
package StatUtil

import (
	"math"
)

// referenced from http://tibasicdev.wikidot.com/invnorm
func InverseNormalDist( p float64, mean float64, stdv float64) float64{
	return mean + stdv*InverserNormalCDF(p)
}

func InverserNormalCDF(p float64) float64{
	return math.Sqrt(2)*ErfInv(2*p - 1)
}
//  function reference from Apache commons-math Erf
func ErfInv(x float64) float64 {

	// beware that the logarithm argument must be
    // commputed as (1.0 - x) * (1.0 + x),
    // it must NOT be simplified as 1.0 - x * x as this
    // would induce rounding errors near the boundaries +/-1

	w := - math.Log((1.0 - x) * (1.0 + x))
	var p float64

	if w < 6.25 {
		w -= 3.125
		p =  -3.6444120640178196996e-21
        p =   -1.685059138182016589e-19 + p * w
        p =   1.2858480715256400167e-18 + p * w
        p =    1.115787767802518096e-17 + p * w
        p =   -1.333171662854620906e-16 + p * w
        p =   2.0972767875968561637e-17 + p * w
        p =   6.6376381343583238325e-15 + p * w
        p =  -4.0545662729752068639e-14 + p * w
        p =  -8.1519341976054721522e-14 + p * w
        p =   2.6335093153082322977e-12 + p * w
        p =  -1.2975133253453532498e-11 + p * w
        p =  -5.4154120542946279317e-11 + p * w
        p =    1.051212273321532285e-09 + p * w
        p =  -4.1126339803469836976e-09 + p * w
        p =  -2.9070369957882005086e-08 + p * w
        p =   4.2347877827932403518e-07 + p * w
        p =  -1.3654692000834678645e-06 + p * w
        p =  -1.3882523362786468719e-05 + p * w
        p =    0.0001867342080340571352 + p * w
        p =  -0.00074070253416626697512 + p * w
        p =   -0.0060336708714301490533 + p * w
        p =      0.24015818242558961693 + p * w
        p =       1.6536545626831027356 + p * w
	} else if w < 16.0 {
		w = math.Sqrt(w) - 3.25
		p =   2.2137376921775787049e-09
        p =   9.0756561938885390979e-08 + p * w
        p =  -2.7517406297064545428e-07 + p * w
        p =   1.8239629214389227755e-08 + p * w
        p =   1.5027403968909827627e-06 + p * w
        p =   -4.013867526981545969e-06 + p * w
        p =   2.9234449089955446044e-06 + p * w
        p =   1.2475304481671778723e-05 + p * w
        p =  -4.7318229009055733981e-05 + p * w
        p =   6.8284851459573175448e-05 + p * w
        p =   2.4031110387097893999e-05 + p * w
        p =   -0.0003550375203628474796 + p * w
        p =   0.00095328937973738049703 + p * w
        p =   -0.0016882755560235047313 + p * w
        p =    0.0024914420961078508066 + p * w
        p =   -0.0037512085075692412107 + p * w
        p =     0.005370914553590063617 + p * w
        p =       1.0052589676941592334 + p * w
        p =       3.0838856104922207635 + p * w
	} else if !math.IsInf(w, 0){
		w = math.Sqrt(w) - 5.0
		p =  -2.7109920616438573243e-11
		p =  -2.5556418169965252055e-10 + p * w
        p =   1.5076572693500548083e-09 + p * w
        p =  -3.7894654401267369937e-09 + p * w
        p =   7.6157012080783393804e-09 + p * w
        p =  -1.4960026627149240478e-08 + p * w
        p =   2.9147953450901080826e-08 + p * w
        p =  -6.7711997758452339498e-08 + p * w
        p =   2.2900482228026654717e-07 + p * w
        p =  -9.9298272942317002539e-07 + p * w
        p =   4.5260625972231537039e-06 + p * w
        p =  -1.9681778105531670567e-05 + p * w
        p =   7.5995277030017761139e-05 + p * w
        p =  -0.00021503011930044477347 + p * w
        p =  -0.00013871931833623122026 + p * w
        p =       1.0103004648645343977 + p * w
        p =       4.8499064014085844221 + p * w

	} else {
		p = math.Inf(0)
	}

	return p * x
}
<commit_msg>added proper documentation<commit_after>//Provides Additional common stat functions that is not included in GO standard library.
//GO standard library philosophy is to follow C standard library
package StatUtil

import (
	"math"
)

//Using Inverse Normal CDF function, mean and standard deviation to get the Inverse Norm Dist.
//referenced from http://tibasicdev.wikidot.com/invnorm
func InverseNormalDist( p float64, mean float64, stdv float64) float64{
	return mean + stdv*InverserNormalCDF(p)
}

//Inverse Normal CDF,
//p is probability
//referenced from http://tibasicdev.wikidot.com/invnorm
func InverserNormalCDF(p float64) float64{
	return math.Sqrt(2)*ErfInv(2*p - 1)
}

//Inverse of the Error Function,
//function referenced from Apache commons-math Erf
func ErfInv(x float64) float64 {

	// beware that the logarithm argument must be
    // commputed as (1.0 - x) * (1.0 + x),
    // it must NOT be simplified as 1.0 - x * x as this
    // would induce rounding errors near the boundaries +/-1

	w := - math.Log((1.0 - x) * (1.0 + x))
	var p float64

	if w < 6.25 {
		w -= 3.125
		p =  -3.6444120640178196996e-21
        p =   -1.685059138182016589e-19 + p * w
        p =   1.2858480715256400167e-18 + p * w
        p =    1.115787767802518096e-17 + p * w
        p =   -1.333171662854620906e-16 + p * w
        p =   2.0972767875968561637e-17 + p * w
        p =   6.6376381343583238325e-15 + p * w
        p =  -4.0545662729752068639e-14 + p * w
        p =  -8.1519341976054721522e-14 + p * w
        p =   2.6335093153082322977e-12 + p * w
        p =  -1.2975133253453532498e-11 + p * w
        p =  -5.4154120542946279317e-11 + p * w
        p =    1.051212273321532285e-09 + p * w
        p =  -4.1126339803469836976e-09 + p * w
        p =  -2.9070369957882005086e-08 + p * w
        p =   4.2347877827932403518e-07 + p * w
        p =  -1.3654692000834678645e-06 + p * w
        p =  -1.3882523362786468719e-05 + p * w
        p =    0.0001867342080340571352 + p * w
        p =  -0.00074070253416626697512 + p * w
        p =   -0.0060336708714301490533 + p * w
        p =      0.24015818242558961693 + p * w
        p =       1.6536545626831027356 + p * w
	} else if w < 16.0 {
		w = math.Sqrt(w) - 3.25
		p =   2.2137376921775787049e-09
        p =   9.0756561938885390979e-08 + p * w
        p =  -2.7517406297064545428e-07 + p * w
        p =   1.8239629214389227755e-08 + p * w
        p =   1.5027403968909827627e-06 + p * w
        p =   -4.013867526981545969e-06 + p * w
        p =   2.9234449089955446044e-06 + p * w
        p =   1.2475304481671778723e-05 + p * w
        p =  -4.7318229009055733981e-05 + p * w
        p =   6.8284851459573175448e-05 + p * w
        p =   2.4031110387097893999e-05 + p * w
        p =   -0.0003550375203628474796 + p * w
        p =   0.00095328937973738049703 + p * w
        p =   -0.0016882755560235047313 + p * w
        p =    0.0024914420961078508066 + p * w
        p =   -0.0037512085075692412107 + p * w
        p =     0.005370914553590063617 + p * w
        p =       1.0052589676941592334 + p * w
        p =       3.0838856104922207635 + p * w
	} else if !math.IsInf(w, 0){
		w = math.Sqrt(w) - 5.0
		p =  -2.7109920616438573243e-11
		p =  -2.5556418169965252055e-10 + p * w
        p =   1.5076572693500548083e-09 + p * w
        p =  -3.7894654401267369937e-09 + p * w
        p =   7.6157012080783393804e-09 + p * w
        p =  -1.4960026627149240478e-08 + p * w
        p =   2.9147953450901080826e-08 + p * w
        p =  -6.7711997758452339498e-08 + p * w
        p =   2.2900482228026654717e-07 + p * w
        p =  -9.9298272942317002539e-07 + p * w
        p =   4.5260625972231537039e-06 + p * w
        p =  -1.9681778105531670567e-05 + p * w
        p =   7.5995277030017761139e-05 + p * w
        p =  -0.00021503011930044477347 + p * w
        p =  -0.00013871931833623122026 + p * w
        p =       1.0103004648645343977 + p * w
        p =       4.8499064014085844221 + p * w

	} else {
		p = math.Inf(0)
	}

	return p * x
}
<|endoftext|>


<commit_before>// 5 october 2014
package main

import (
	"os"
	"path/filepath"
	"strings"
	"sort"
	"github.com/andlabs/ohv/mshi"
)

type MSHI struct {
	dir			string
	containers	[][]*mshi.ContainerPath
	assets		[][]*mshi.AssetData
	topics		map[string]*MSHITopic
	books		[]Topic
	orphans		[]Topic
}

func (m *MSHI) Name() string {
	return "All Microsoft Help Viewer-format documentation"
}

func (m *MSHI) Books() []Topic {
	return m.books
}

func (m *MSHI) Orphans() []Topic {
	return m.orphans
}

// TODO adorn errors
func OpenMSHI(dir string) (*MSHI, error) {
	m := new(MSHI)
	m.dir = dir

	// first read all indices
	err := filepath.Walk(m.dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// TODO subdirectories
		if strings.ToLower(filepath.Ext(path)) != ".mshi" {
			return nil
		}
		of, err := os.Open(path)
		if err != nil {
			return err
		}
		f := mshi.Open(of)
		if f.Err() != nil {
			return f.Err()
		}
		c := f.ReadContainerPaths()
		if f.Err() != nil {
			return f.Err()
		}
		a := f.ReadAssetDatas()
		if f.Err() != nil {
			return f.Err()
		}
		of.Close()
		m.containers = append(m.containers, c)
		m.assets = append(m.assets, a)
		return nil
	})
	if err != nil {
		return nil, err
	}

	// then produce a map that goes from ID -> topic
	// TODO deal with versions sensibly
	m.topics = make(map[string]*MSHITopic)
	for container, aa := range m.assets {
		for _, a := range aa {
			if m.topics[a.ID] != nil && m.topics[a.ID].asset.Version > a.Version {
				continue
			}
			m.topics[a.ID] = &MSHITopic{
				containers:	m.containers[container],
				asset:		a,
			}
		}
	}

	// now figure out the hierarchy, books, and orphans
	all := make([]*MSHITopic, 0, len(m.topics))
	for _, v := range m.topics {
		all = append(all, v)
	}
	i := 0
	for i < len(all) {
		if all[i].asset.ParentID == "-1" {		// is top-level
			m.books = append(m.books, all[i])
			all = append(all[:i], all[i + 1:]...)
			continue
		}
		parent, ok := m.topics[all[i].asset.ParentID]
		if ok {			// has parent
			parent.children = append(parent.children, all[i])
			all = append(all[:i], all[i + 1:]...)
			continue
		}
		i++		// parent elsewhere or missing
	}
	for _, v := range all {
		m.orphans = append(m.orphans, v)
	}

	// now sort all children
	for _, t := range m.topics {
		sort.Sort(TopicSorter(t.children))
	}

	return m, nil
}

type MSHITopic struct {
	containers	[]*mshi.ContainerPath
	asset			*mshi.AssetData
	children		[]Topic
}

func (m *MSHITopic) Name() string {
	return m.asset.Title
}

func (m *MSHITopic) Prepare() (string, error) {
	// TODO
	return "", nil
}

func (m *MSHITopic) Children() []Topic {
	return m.children
}

func (m *MSHITopic) Less(t Topic) bool {
	tt := t.(*MSHITopic)
	return m.asset.Order < tt.asset.Order
}

func main() {
	m, err := OpenMSHI(os.Args[1])
	if err != nil { panic(err) }
	println("books:")
	for _, b := range m.Books() {
		println(b.Name())
	}
	println("orphans:", len(m.Orphans()))
}
<commit_msg>Split the MSHI constructor apart to make a CPU profile more reasonable.<commit_after>// 5 october 2014
package main

import (
	"os"
	"path/filepath"
	"strings"
	"sort"
	"github.com/andlabs/ohv/mshi"
)

type MSHI struct {
	dir			string
	containers	[][]*mshi.ContainerPath
	assets		[][]*mshi.AssetData
	topics		map[string]*MSHITopic
	books		[]Topic
	orphans		[]Topic
}

func (m *MSHI) Name() string {
	return "All Microsoft Help Viewer-format documentation"
}

func (m *MSHI) Books() []Topic {
	return m.books
}

func (m *MSHI) Orphans() []Topic {
	return m.orphans
}

// TODO adorn errors
func OpenMSHI(dir string) (*MSHI, error) {
	m := new(MSHI)
	m.dir = dir

	// first read all indices
	err := m.readAllIndices()
	if err != nil {
		return nil, err
	}

	// then produce a map that goes from ID -> topic
	m.collectTopics()

	// now figure out the hierarchy, books, and orphans
	m.buildHierarchy()

	// now sort all children
	m.sortAllChildren()

	return m, nil
}

func (m *MSHI) readAllIndices() error {
	return filepath.Walk(m.dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// TODO subdirectories
		if strings.ToLower(filepath.Ext(path)) != ".mshi" {
			return nil
		}
		of, err := os.Open(path)
		if err != nil {
			return err
		}
		f := mshi.Open(of)
		if f.Err() != nil {
			return f.Err()
		}
		c := f.ReadContainerPaths()
		if f.Err() != nil {
			return f.Err()
		}
		a := f.ReadAssetDatas()
		if f.Err() != nil {
			return f.Err()
		}
		of.Close()
		m.containers = append(m.containers, c)
		m.assets = append(m.assets, a)
		return nil
	})
}

func (m *MSHI) collectTopics() {
	// TODO deal with versions sensibly
	m.topics = make(map[string]*MSHITopic)
	for container, aa := range m.assets {
		for _, a := range aa {
			if m.topics[a.ID] != nil && m.topics[a.ID].asset.Version > a.Version {
				continue
			}
			m.topics[a.ID] = &MSHITopic{
				containers:	m.containers[container],
				asset:		a,
			}
		}
	}
}

func (m *MSHI) buildHierarchy() {
	all := make([]*MSHITopic, 0, len(m.topics))
	for _, v := range m.topics {
		all = append(all, v)
	}
	i := 0
	for i < len(all) {
		if all[i].asset.ParentID == "-1" {		// is top-level
			m.books = append(m.books, all[i])
			all = append(all[:i], all[i + 1:]...)
			continue
		}
		parent, ok := m.topics[all[i].asset.ParentID]
		if ok {			// has parent
			parent.children = append(parent.children, all[i])
			all = append(all[:i], all[i + 1:]...)
			continue
		}
		i++		// parent elsewhere or missing
	}
	for _, v := range all {
		m.orphans = append(m.orphans, v)
	}
}

func (m *MSHI) sortAllChildren() {
	for _, t := range m.topics {
		sort.Sort(TopicSorter(t.children))
	}
}

type MSHITopic struct {
	containers	[]*mshi.ContainerPath
	asset			*mshi.AssetData
	children		[]Topic
}

func (m *MSHITopic) Name() string {
	return m.asset.Title
}

func (m *MSHITopic) Prepare() (string, error) {
	// TODO
	return "", nil
}

func (m *MSHITopic) Children() []Topic {
	return m.children
}

func (m *MSHITopic) Less(t Topic) bool {
	tt := t.(*MSHITopic)
	return m.asset.Order < tt.asset.Order
}

func main() {
	m, err := OpenMSHI(os.Args[1])
	if err != nil { panic(err) }
	println("books:")
	for _, b := range m.Books() {
		println(b.Name())
	}
	println("orphans:", len(m.Orphans()))
}
<|endoftext|>


<commit_before>//
// # Muta Bin
//
// Handle the CLI in/out of a `muta.go` file that was ran.
//
package main

import "github.com/docopt/docopt-go"

func ParseArgs() {
	usage := `Muta(te) Build System

Usage:
	muta [task]
  muta -h | --help
  muta --version

Options:
  -h --help     Show this screen.
  --version     Show version.`

	docopt.Parse(usage, nil, true, "Muta 0.0.0", false)
}

func main() {
}
<commit_msg>Corrected typo<commit_after>//
// # Muta Bin
//
// Handle the CLI in/out of a `Muta` file that was ran.
//
package main

import "github.com/docopt/docopt-go"

func ParseArgs() {
	usage := `Muta(te) Build System

Usage:
	muta [task]
  muta -h | --help
  muta --version

Options:
  -h --help     Show this screen.
  --version     Show version.`

	docopt.Parse(usage, nil, true, "Muta 0.0.0", false)
}

func main() {
}
<|endoftext|>


<commit_before>package lzjson

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"regexp"
)

// Type represents the different type of JSON values
// (string, number, object, array, true, false, null)
// true and false are combined as bool for obvious reason
type Type int

// These constant represents different JSON value types
// as specified in http://www.json.org/
// with some exception:
// 1. true and false are combined as bool for obvious reason; and
// 2. TypeUnknown for empty strings
const (
	TypeError     Type = -1
	TypeUndefined Type = iota
	TypeString
	TypeNumber
	TypeObject
	TypeArray
	TypeBool
	TypeNull
)

// String returns string representations of
// the Type value
func (t Type) String() string {
	switch t {
	case TypeUndefined:
		return "TypeUndefined"
	case TypeString:
		return "TypeString"
	case TypeNumber:
		return "TypeNumber"
	case TypeObject:
		return "TypeObject"
	case TypeArray:
		return "TypeArray"
	case TypeBool:
		return "TypeBool"
	case TypeNull:
		return "TypeNull"
	}
	return "TypeError"
}

// reNumber is the regular expression to match
// any JSON number values
var reNum = regexp.MustCompile(`^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+\-]?\d+)?$`)

// IsNumJSON test a string and see if it match the
// JSON definition of number
func IsNumJSON(b []byte) bool {
	return reNum.Match(b)
}

// Node is an interface for all JSON nodes
type Node interface {

	// Unmarshal parses the JSON node data into variable v
	Unmarshal(v interface{}) error

	// UnmarshalJSON implements json.Unmarshaler
	UnmarshalJSON(b []byte) error

	// Raw returns the raw JSON string in []byte
	Raw() []byte

	// Type returns the Type of the containing JSON value
	Type() Type

	// Get gets object's inner value.
	// Only works with Object value type
	Get(key string) (inner Node)

	// Len gets the length of the value
	// Only works with Array and String value type
	Len() int

	// GetN gets array's inner value.
	// Only works with Array value type.
	// 0 for the first item.
	GetN(nth int) Node

	// String unmarshal the JSON into string then return
	String() (v string)

	// Number unmarshal the JSON into float64 then return
	Number() (v float64)

	// Int unmarshal the JSON into int the return
	Int() (v int)

	// Bool unmarshal the JSON into bool then return
	Bool() (v bool)

	// IsNull tells if the JSON value is null or not
	IsNull() bool

	// Error returns the JSON parse error, if any
	Error() error
}

// NewNode returns an initialized empty Node value
// ready for unmarshaling
func NewNode() Node {
	return &rootNode{}
}

// Decode read and decodes a JSON from io.Reader then
// returns a Node of it
func Decode(reader io.Reader) (n Node, err error) {
	b, err := ioutil.ReadAll(reader)
	n = &rootNode{buf: b}
	return
}

// rootNode is the default implementation of Node
type rootNode struct {
	buf []byte
	err error
}

// Unmarshal implements Node
func (n *rootNode) Unmarshal(v interface{}) error {
	return json.Unmarshal(n.buf, v)
}

// UnmarshalJSON implements Node
func (n *rootNode) UnmarshalJSON(b []byte) error {
	n.buf = b
	return nil
}

// Raw implements Node
func (n *rootNode) Raw() []byte {
	return n.buf
}

// Type implements Node
func (n rootNode) Type() Type {

	switch {
	case n.err != nil:
		// for error, return TypeError
		return TypeError
	case n.buf == nil || len(n.buf) == 0:
		// for nil raw, return TypeUndefined
		return TypeUndefined
	case n.buf[0] == '"':
		// simply examine the first character
		// to determine the value type
		return TypeString
	case n.buf[0] == '{':
		// simply examine the first character
		// to determine the value type
		return TypeObject
	case n.buf[0] == '[':
		// simply examine the first character
		// to determine the value type
		return TypeArray
	case string(n.buf) == "true":
		fallthrough
	case string(n.buf) == "false":
		return TypeBool
	case string(n.buf) == "null":
		return TypeNull
	case IsNumJSON(n.buf):
		return TypeNumber
	}

	// return TypeUnknown for all other cases
	return TypeError
}

// Get implements Node
func (n *rootNode) Get(key string) (inner Node) {
	if n.Type() != TypeObject {
		inner = &rootNode{}
		return
	}

	vmap := map[string]rootNode{}
	if err := n.Unmarshal(&vmap); err != nil {
		inner = &rootNode{err: err} // dump the error
	} else if val, ok := vmap[key]; !ok {
		inner = &rootNode{err: fmt.Errorf("field %#v not found", key)}
	} else {
		inner = &val
	}
	return
}

// Len gets the length of the value
// Only works with Array and String value type
func (n *rootNode) Len() int {
	switch n.Type() {
	case TypeString:
		return len(string(n.buf)) - 2 // subtact the 2 " marks
	case TypeArray:
		vslice := []*rootNode{}
		n.Unmarshal(&vslice)
		return len(vslice)
	}
	// default return -1 (for type mismatch)
	return -1
}

// GetN implements Node
func (n *rootNode) GetN(nth int) Node {
	if n.Type() != TypeArray {
		return nil
	}

	vslice := []rootNode{}
	n.Unmarshal(&vslice)
	if nth < len(vslice) {
		return &vslice[nth]
	}
	return nil
}

// String implements Node
func (n *rootNode) String() (v string) {
	n.Unmarshal(&v)
	return
}

// Number implements Node
func (n *rootNode) Number() (v float64) {
	n.Unmarshal(&v)
	return
}

// Int implements Node
func (n *rootNode) Int() (v int) {
	return int(n.Number())
}

// Bool implements Node
func (n *rootNode) Bool() (v bool) {
	n.Unmarshal(&v)
	return
}

// IsNull implements Node
func (n *rootNode) IsNull() bool {
	return n.Type() == TypeNull
}

// Error implements Node
func (n *rootNode) Error() error {
	return n.err
}
<commit_msg>Type implements GoStringer<commit_after>package lzjson

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"regexp"
)

// Type represents the different type of JSON values
// (string, number, object, array, true, false, null)
// true and false are combined as bool for obvious reason
type Type int

// These constant represents different JSON value types
// as specified in http://www.json.org/
// with some exception:
// 1. true and false are combined as bool for obvious reason; and
// 2. TypeUnknown for empty strings
const (
	TypeError     Type = -1
	TypeUndefined Type = iota
	TypeString
	TypeNumber
	TypeObject
	TypeArray
	TypeBool
	TypeNull
)

// String returns string representations of
// the Type value
func (t Type) String() string {
	switch t {
	case TypeUndefined:
		return "TypeUndefined"
	case TypeString:
		return "TypeString"
	case TypeNumber:
		return "TypeNumber"
	case TypeObject:
		return "TypeObject"
	case TypeArray:
		return "TypeArray"
	case TypeBool:
		return "TypeBool"
	case TypeNull:
		return "TypeNull"
	}
	return "TypeError"
}

// GoString implements fmt.GoStringer
func (t Type) GoString() string {
	return "lzjson." + t.String()
}

// reNumber is the regular expression to match
// any JSON number values
var reNum = regexp.MustCompile(`^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+\-]?\d+)?$`)

// IsNumJSON test a string and see if it match the
// JSON definition of number
func IsNumJSON(b []byte) bool {
	return reNum.Match(b)
}

// Node is an interface for all JSON nodes
type Node interface {

	// Unmarshal parses the JSON node data into variable v
	Unmarshal(v interface{}) error

	// UnmarshalJSON implements json.Unmarshaler
	UnmarshalJSON(b []byte) error

	// Raw returns the raw JSON string in []byte
	Raw() []byte

	// Type returns the Type of the containing JSON value
	Type() Type

	// Get gets object's inner value.
	// Only works with Object value type
	Get(key string) (inner Node)

	// Len gets the length of the value
	// Only works with Array and String value type
	Len() int

	// GetN gets array's inner value.
	// Only works with Array value type.
	// 0 for the first item.
	GetN(nth int) Node

	// String unmarshal the JSON into string then return
	String() (v string)

	// Number unmarshal the JSON into float64 then return
	Number() (v float64)

	// Int unmarshal the JSON into int the return
	Int() (v int)

	// Bool unmarshal the JSON into bool then return
	Bool() (v bool)

	// IsNull tells if the JSON value is null or not
	IsNull() bool

	// Error returns the JSON parse error, if any
	Error() error
}

// NewNode returns an initialized empty Node value
// ready for unmarshaling
func NewNode() Node {
	return &rootNode{}
}

// Decode read and decodes a JSON from io.Reader then
// returns a Node of it
func Decode(reader io.Reader) (n Node, err error) {
	b, err := ioutil.ReadAll(reader)
	n = &rootNode{buf: b}
	return
}

// rootNode is the default implementation of Node
type rootNode struct {
	buf []byte
	err error
}

// Unmarshal implements Node
func (n *rootNode) Unmarshal(v interface{}) error {
	return json.Unmarshal(n.buf, v)
}

// UnmarshalJSON implements Node
func (n *rootNode) UnmarshalJSON(b []byte) error {
	n.buf = b
	return nil
}

// Raw implements Node
func (n *rootNode) Raw() []byte {
	return n.buf
}

// Type implements Node
func (n rootNode) Type() Type {

	switch {
	case n.err != nil:
		// for error, return TypeError
		return TypeError
	case n.buf == nil || len(n.buf) == 0:
		// for nil raw, return TypeUndefined
		return TypeUndefined
	case n.buf[0] == '"':
		// simply examine the first character
		// to determine the value type
		return TypeString
	case n.buf[0] == '{':
		// simply examine the first character
		// to determine the value type
		return TypeObject
	case n.buf[0] == '[':
		// simply examine the first character
		// to determine the value type
		return TypeArray
	case string(n.buf) == "true":
		fallthrough
	case string(n.buf) == "false":
		return TypeBool
	case string(n.buf) == "null":
		return TypeNull
	case IsNumJSON(n.buf):
		return TypeNumber
	}

	// return TypeUnknown for all other cases
	return TypeError
}

// Get implements Node
func (n *rootNode) Get(key string) (inner Node) {
	if n.Type() != TypeObject {
		inner = &rootNode{}
		return
	}

	vmap := map[string]rootNode{}
	if err := n.Unmarshal(&vmap); err != nil {
		inner = &rootNode{err: err} // dump the error
	} else if val, ok := vmap[key]; !ok {
		inner = &rootNode{err: fmt.Errorf("field %#v not found", key)}
	} else {
		inner = &val
	}
	return
}

// Len gets the length of the value
// Only works with Array and String value type
func (n *rootNode) Len() int {
	switch n.Type() {
	case TypeString:
		return len(string(n.buf)) - 2 // subtact the 2 " marks
	case TypeArray:
		vslice := []*rootNode{}
		n.Unmarshal(&vslice)
		return len(vslice)
	}
	// default return -1 (for type mismatch)
	return -1
}

// GetN implements Node
func (n *rootNode) GetN(nth int) Node {
	if n.Type() != TypeArray {
		return nil
	}

	vslice := []rootNode{}
	n.Unmarshal(&vslice)
	if nth < len(vslice) {
		return &vslice[nth]
	}
	return nil
}

// String implements Node
func (n *rootNode) String() (v string) {
	n.Unmarshal(&v)
	return
}

// Number implements Node
func (n *rootNode) Number() (v float64) {
	n.Unmarshal(&v)
	return
}

// Int implements Node
func (n *rootNode) Int() (v int) {
	return int(n.Number())
}

// Bool implements Node
func (n *rootNode) Bool() (v bool) {
	n.Unmarshal(&v)
	return
}

// IsNull implements Node
func (n *rootNode) IsNull() bool {
	return n.Type() == TypeNull
}

// Error implements Node
func (n *rootNode) Error() error {
	return n.err
}
<|endoftext|>


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

// Package stats collects serviced metrics and posts them to the TSDB.

package stats

import (
	"github.com/control-center/go-procfs/linux"
	coordclient "github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/stats/cgroup"
	"github.com/control-center/serviced/utils"
	zkservice "github.com/control-center/serviced/zzk/service"
	"github.com/rcrowley/go-metrics"
	"github.com/zenoss/glog"

	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// StatsReporter collects and posts serviced stats to the TSDB.
type StatsReporter struct {
	destination         string
	closeChannel        chan bool
	conn                coordclient.Connection
	containerRegistries map[registryKey]metrics.Registry
	hostID              string
	hostRegistry        metrics.Registry
	sync.Mutex
}

// Sample is a single metric measurement
type Sample struct {
	Metric    string            `json:"metric"`
	Value     string            `json:"value"`
	Timestamp int64             `json:"timestamp"`
	Tags      map[string]string `json:"tags"`
}

type registryKey struct {
	serviceID  string
	instanceID int
}

// NewStatsReporter creates a new StatsReporter and kicks off the reporting goroutine.
func NewStatsReporter(destination string, interval time.Duration, conn coordclient.Connection) (*StatsReporter, error) {
	hostID, err := utils.HostID()
	if err != nil {
		glog.Errorf("Could not determine host ID.")
		return nil, err
	}
	if conn == nil {
		glog.Errorf("conn can not be nil")
		return nil, fmt.Errorf("conn can not be nil")
	}
	sr := StatsReporter{
		destination:         destination,
		closeChannel:        make(chan bool),
		conn:                conn,
		containerRegistries: make(map[registryKey]metrics.Registry),
		hostID:              hostID,
	}

	sr.hostRegistry = metrics.NewRegistry()
	go sr.report(interval)
	return &sr, nil
}

// getOrCreateContainerRegistry returns a registry for a given service id or creates it
// if it doesn't exist.
func (sr StatsReporter) getOrCreateContainerRegistry(serviceID string, instanceID int) metrics.Registry {
	key := registryKey{serviceID, instanceID}
	if registry, ok := sr.containerRegistries[key]; ok {
		return registry
	}
	sr.Lock()
	defer sr.Unlock()
	sr.containerRegistries[key] = metrics.NewRegistry()
	return sr.containerRegistries[key]
}

func (sr StatsReporter) removeStaleRegistries(running *[]dao.RunningService) {
	// First build a list of what's actually running
	keys := make(map[string][]int)
	for _, rs := range *running {
		if instances, ok := keys[rs.ServiceID]; !ok {
			instances = []int{rs.InstanceID}
			keys[rs.ServiceID] = instances
		} else {
			keys[rs.ServiceID] = append(keys[rs.ServiceID], rs.InstanceID)
		}
	}
	// Now remove any keys that are in the registry but are no longer running
	// on this host
	sr.Lock()
	defer sr.Unlock()
	for key, _ := range sr.containerRegistries {
		if instances, ok := keys[key.serviceID]; !ok {
			delete(sr.containerRegistries, key)
		} else {
			var seen bool
			for _, instanceid := range instances {
				if instanceid == key.instanceID {
					seen = true
					break
				}
			}
			if !seen {
				delete(sr.containerRegistries, key)
			}
		}
	}
}

// Close shuts down the reporting goroutine. Blocks waiting for the goroutine to signal that it
// is indeed shutting down.
func (sr StatsReporter) Close() {
	sr.closeChannel <- true
	_ = <-sr.closeChannel
}

// Updates the default registry, fills out the metric consumer format, and posts
// the data to the TSDB. Stops when close signal is received on closeChannel.
func (sr StatsReporter) report(d time.Duration) {
	tc := time.Tick(d)
	glog.Infof("collecting internal metrics at %s intervals", d)
	for {
		select {
		case _ = <-sr.closeChannel:
			glog.V(3).Info("Ceasing stat reporting.")
			sr.closeChannel <- true
			return
		case t := <-tc:
			glog.V(1).Info("Reporting container stats at:", t)
			sr.updateStats()
			stats := sr.gatherStats(t)
			err := Post(sr.destination, stats)
			if err != nil {
				glog.Errorf("Error reporting container stats: %v", err)
			}
		}
	}
}

func (sr StatsReporter) updateHostStats() {

	loadavg, err := linux.ReadLoadavg()
	if err != nil {
		glog.Errorf("could not read loadavg: %s", err)
		return
	}
	metrics.GetOrRegisterGaugeFloat64("load.avg1m", sr.hostRegistry).Update(float64(loadavg.Avg1m))
	metrics.GetOrRegisterGaugeFloat64("load.avg5m", sr.hostRegistry).Update(float64(loadavg.Avg5m))
	metrics.GetOrRegisterGaugeFloat64("load.avg10m", sr.hostRegistry).Update(float64(loadavg.Avg10m))
	metrics.GetOrRegisterGauge("load.runningprocesses", sr.hostRegistry).Update(int64(loadavg.RunningProcesses))
	metrics.GetOrRegisterGauge("load.totalprocesses", sr.hostRegistry).Update(int64(loadavg.TotalProcesses))

	stat, err := linux.ReadStat()
	if err != nil {
		glog.Errorf("could not read stat: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("cpu.user", sr.hostRegistry).Update(int64(stat.Cpu.User()))
	metrics.GetOrRegisterGauge("cpu.nice", sr.hostRegistry).Update(int64(stat.Cpu.Nice()))
	metrics.GetOrRegisterGauge("cpu.system", sr.hostRegistry).Update(int64(stat.Cpu.System()))
	metrics.GetOrRegisterGauge("cpu.idle", sr.hostRegistry).Update(int64(stat.Cpu.Idle()))
	metrics.GetOrRegisterGauge("cpu.iowait", sr.hostRegistry).Update(int64(stat.Cpu.Iowait()))
	metrics.GetOrRegisterGauge("cpu.irq", sr.hostRegistry).Update(int64(stat.Cpu.Irq()))
	metrics.GetOrRegisterGauge("cpu.softirq", sr.hostRegistry).Update(int64(stat.Cpu.Softirq()))
	var steal int64
	if stat.Cpu.StealSupported() {
		steal = int64(stat.Cpu.Steal())
	}
	metrics.GetOrRegisterGauge("cpu.steal", sr.hostRegistry).Update(steal)

	meminfo, err := linux.ReadMeminfo()
	if err != nil {
		glog.Errorf("could not read meminfo: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("memory.total", sr.hostRegistry).Update(int64(meminfo.MemTotal))
	metrics.GetOrRegisterGauge("memory.free", sr.hostRegistry).Update(int64(meminfo.MemFree))
	metrics.GetOrRegisterGauge("memory.buffers", sr.hostRegistry).Update(int64(meminfo.Buffers))
	metrics.GetOrRegisterGauge("memory.cached", sr.hostRegistry).Update(int64(meminfo.Cached))
	actualFree := int64(meminfo.MemFree) + int64(meminfo.Buffers) + int64(meminfo.Cached)
	metrics.GetOrRegisterGauge("memory.actualfree", sr.hostRegistry).Update(actualFree)
	metrics.GetOrRegisterGauge("memory.actualused", sr.hostRegistry).Update(int64(meminfo.MemTotal) - actualFree)
	metrics.GetOrRegisterGauge("swap.total", sr.hostRegistry).Update(int64(meminfo.SwapTotal))
	metrics.GetOrRegisterGauge("swap.free", sr.hostRegistry).Update(int64(meminfo.SwapFree))

	vmstat, err := linux.ReadVmstat()
	if err != nil {
		glog.Errorf("could not read vmstat: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("vmstat.pgfault", sr.hostRegistry).Update(int64(vmstat.Pgfault))
	metrics.GetOrRegisterGauge("vmstat.pgmajfault", sr.hostRegistry).Update(int64(vmstat.Pgmajfault))
	metrics.GetOrRegisterGauge("vmstat.pgpgout", sr.hostRegistry).Update(int64(vmstat.Pgpgout) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pgpgin", sr.hostRegistry).Update(int64(vmstat.Pgpgin) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pswpout", sr.hostRegistry).Update(int64(vmstat.Pswpout) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pswpin", sr.hostRegistry).Update(int64(vmstat.Pswpin) * 1024)

	if openFileDescriptorCount, err := GetOpenFileDescriptorCount(); err != nil {
		glog.Warningf("Couldn't get open file descriptor count", err)
	} else {
		metrics.GetOrRegisterGauge("Serviced.OpenFileDescriptors", sr.hostRegistry).Update(openFileDescriptorCount)
	}
}

// Updates the default registry.
func (sr StatsReporter) updateStats() {
	// Stats for host.
	sr.updateHostStats()
	// Stats for the containers.
	var running []dao.RunningService
	running, err := zkservice.LoadRunningServicesByHost(sr.conn, sr.hostID)
	if err != nil {
		glog.Errorf("updateStats: zkservice.LoadRunningServicesByHost (conn: %+v hostID: %v) failed: %v", sr.conn, sr.hostID, err)
	}

	for _, rs := range running {
		if rs.DockerID != "" {
			containerRegistry := sr.getOrCreateContainerRegistry(rs.ServiceID, rs.InstanceID)
			if cpuacctStat, err := cgroup.ReadCpuacctStat(cgroup.GetCgroupDockerStatsFilePath(rs.DockerID, cgroup.Cpuacct)); err != nil {
				glog.V(4).Infof("Couldn't read CpuacctStat:", err)
			} else {
				metrics.GetOrRegisterGauge("cgroup.cpuacct.system", containerRegistry).Update(cpuacctStat.System)
				metrics.GetOrRegisterGauge("cgroup.cpuacct.user", containerRegistry).Update(cpuacctStat.User)
			}
			if memoryStat, err := cgroup.ReadMemoryStat(cgroup.GetCgroupDockerStatsFilePath(rs.DockerID, cgroup.Memory)); err != nil {
				glog.V(4).Infof("Couldn't read MemoryStat:", err)
			} else {
				metrics.GetOrRegisterGauge("cgroup.memory.pgmajfault", containerRegistry).Update(memoryStat.Pgfault)
				metrics.GetOrRegisterGauge("cgroup.memory.totalrss", containerRegistry).Update(memoryStat.TotalRss)
				metrics.GetOrRegisterGauge("cgroup.memory.cache", containerRegistry).Update(memoryStat.Cache)
			}
		} else {
			glog.V(4).Infof("Skipping stats update for %s (%s), no container ID exists yet", rs.Name, rs.ServiceID)
		}
	}
	// Clean out old container registries
	sr.removeStaleRegistries(&running)
}

// Fills out the metric consumer format.
func (sr StatsReporter) gatherStats(t time.Time) []Sample {
	stats := []Sample{}
	// Handle the host metrics.
	reg, _ := sr.hostRegistry.(*metrics.StandardRegistry)
	reg.Each(func(name string, i interface{}) {
		if metric, ok := i.(metrics.Gauge); ok {
			tagmap := make(map[string]string)
			tagmap["controlplane_host_id"] = sr.hostID
			stats = append(stats, Sample{name, strconv.FormatInt(metric.Value(), 10), t.Unix(), tagmap})
		}
		if metricf64, ok := i.(metrics.GaugeFloat64); ok {
			tagmap := make(map[string]string)
			tagmap["controlplane_host_id"] = sr.hostID
			stats = append(stats, Sample{name, strconv.FormatFloat(metricf64.Value(), 'f', -1, 32), t.Unix(), tagmap})
		}
	})
	// Handle each container's metrics.
	for key, registry := range sr.containerRegistries {
		reg, _ := registry.(*metrics.StandardRegistry)
		reg.Each(func(name string, i interface{}) {
			if metric, ok := i.(metrics.Gauge); ok {
				tagmap := make(map[string]string)
				tagmap["controlplane_service_id"] = key.serviceID
				tagmap["controlplane_instance_id"] = strconv.FormatInt(int64(key.instanceID), 10)
				tagmap["controlplane_host_id"] = sr.hostID
				stats = append(stats, Sample{name, strconv.FormatInt(metric.Value(), 10), t.Unix(), tagmap})
			}
		})
	}
	return stats
}

// Send the list of stats to the TSDB.
func Post(destination string, stats []Sample) error {
	payload := map[string][]Sample{"metrics": stats}
	data, err := json.Marshal(payload)
	if err != nil {
		glog.Warningf("Couldn't marshal stats: ", err)
		return err
	}
	statsreq, err := http.NewRequest("POST", destination, bytes.NewBuffer(data))
	if err != nil {
		glog.Warningf("Couldn't create stats request: ", err)
		return err
	}
	statsreq.Header["User-Agent"] = []string{"Zenoss Metric Publisher"}
	statsreq.Header["Content-Type"] = []string{"application/json"}
	resp, reqerr := http.DefaultClient.Do(statsreq)
	if reqerr != nil {
		glog.Warningf("Couldn't post stats: ", reqerr)
		return reqerr
	}
	if !strings.Contains(resp.Status, "200 OK") {
		glog.Warningf("couldn't post stats: ", resp.Status)
		return nil
	}
	resp.Body.Close()
	return nil
}
<commit_msg>Fix govet issue - don't pass mutex by value<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package stats collects serviced metrics and posts them to the TSDB.

package stats

import (
	"github.com/control-center/go-procfs/linux"
	coordclient "github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/stats/cgroup"
	"github.com/control-center/serviced/utils"
	zkservice "github.com/control-center/serviced/zzk/service"
	"github.com/rcrowley/go-metrics"
	"github.com/zenoss/glog"

	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

// StatsReporter collects and posts serviced stats to the TSDB.
type StatsReporter struct {
	destination         string
	closeChannel        chan bool
	conn                coordclient.Connection
	containerRegistries map[registryKey]metrics.Registry
	hostID              string
	hostRegistry        metrics.Registry
	sync.Mutex
}

// Sample is a single metric measurement
type Sample struct {
	Metric    string            `json:"metric"`
	Value     string            `json:"value"`
	Timestamp int64             `json:"timestamp"`
	Tags      map[string]string `json:"tags"`
}

type registryKey struct {
	serviceID  string
	instanceID int
}

// NewStatsReporter creates a new StatsReporter and kicks off the reporting goroutine.
func NewStatsReporter(destination string, interval time.Duration, conn coordclient.Connection) (*StatsReporter, error) {
	hostID, err := utils.HostID()
	if err != nil {
		glog.Errorf("Could not determine host ID.")
		return nil, err
	}
	if conn == nil {
		glog.Errorf("conn can not be nil")
		return nil, fmt.Errorf("conn can not be nil")
	}
	sr := StatsReporter{
		destination:         destination,
		closeChannel:        make(chan bool),
		conn:                conn,
		containerRegistries: make(map[registryKey]metrics.Registry),
		hostID:              hostID,
	}

	sr.hostRegistry = metrics.NewRegistry()
	go sr.report(interval)
	return &sr, nil
}

// getOrCreateContainerRegistry returns a registry for a given service id or creates it
// if it doesn't exist.
func (sr *StatsReporter) getOrCreateContainerRegistry(serviceID string, instanceID int) metrics.Registry {
	key := registryKey{serviceID, instanceID}
	if registry, ok := sr.containerRegistries[key]; ok {
		return registry
	}
	sr.Lock()
	defer sr.Unlock()
	sr.containerRegistries[key] = metrics.NewRegistry()
	return sr.containerRegistries[key]
}

func (sr *StatsReporter) removeStaleRegistries(running *[]dao.RunningService) {
	// First build a list of what's actually running
	keys := make(map[string][]int)
	for _, rs := range *running {
		if instances, ok := keys[rs.ServiceID]; !ok {
			instances = []int{rs.InstanceID}
			keys[rs.ServiceID] = instances
		} else {
			keys[rs.ServiceID] = append(keys[rs.ServiceID], rs.InstanceID)
		}
	}
	// Now remove any keys that are in the registry but are no longer running
	// on this host
	sr.Lock()
	defer sr.Unlock()
	for key, _ := range sr.containerRegistries {
		if instances, ok := keys[key.serviceID]; !ok {
			delete(sr.containerRegistries, key)
		} else {
			var seen bool
			for _, instanceid := range instances {
				if instanceid == key.instanceID {
					seen = true
					break
				}
			}
			if !seen {
				delete(sr.containerRegistries, key)
			}
		}
	}
}

// Close shuts down the reporting goroutine. Blocks waiting for the goroutine to signal that it
// is indeed shutting down.
func (sr *StatsReporter) Close() {
	sr.closeChannel <- true
	_ = <-sr.closeChannel
}

// Updates the default registry, fills out the metric consumer format, and posts
// the data to the TSDB. Stops when close signal is received on closeChannel.
func (sr *StatsReporter) report(d time.Duration) {
	tc := time.Tick(d)
	glog.Infof("collecting internal metrics at %s intervals", d)
	for {
		select {
		case _ = <-sr.closeChannel:
			glog.V(3).Info("Ceasing stat reporting.")
			sr.closeChannel <- true
			return
		case t := <-tc:
			glog.V(1).Info("Reporting container stats at:", t)
			sr.updateStats()
			stats := sr.gatherStats(t)
			err := Post(sr.destination, stats)
			if err != nil {
				glog.Errorf("Error reporting container stats: %v", err)
			}
		}
	}
}

func (sr *StatsReporter) updateHostStats() {

	loadavg, err := linux.ReadLoadavg()
	if err != nil {
		glog.Errorf("could not read loadavg: %s", err)
		return
	}
	metrics.GetOrRegisterGaugeFloat64("load.avg1m", sr.hostRegistry).Update(float64(loadavg.Avg1m))
	metrics.GetOrRegisterGaugeFloat64("load.avg5m", sr.hostRegistry).Update(float64(loadavg.Avg5m))
	metrics.GetOrRegisterGaugeFloat64("load.avg10m", sr.hostRegistry).Update(float64(loadavg.Avg10m))
	metrics.GetOrRegisterGauge("load.runningprocesses", sr.hostRegistry).Update(int64(loadavg.RunningProcesses))
	metrics.GetOrRegisterGauge("load.totalprocesses", sr.hostRegistry).Update(int64(loadavg.TotalProcesses))

	stat, err := linux.ReadStat()
	if err != nil {
		glog.Errorf("could not read stat: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("cpu.user", sr.hostRegistry).Update(int64(stat.Cpu.User()))
	metrics.GetOrRegisterGauge("cpu.nice", sr.hostRegistry).Update(int64(stat.Cpu.Nice()))
	metrics.GetOrRegisterGauge("cpu.system", sr.hostRegistry).Update(int64(stat.Cpu.System()))
	metrics.GetOrRegisterGauge("cpu.idle", sr.hostRegistry).Update(int64(stat.Cpu.Idle()))
	metrics.GetOrRegisterGauge("cpu.iowait", sr.hostRegistry).Update(int64(stat.Cpu.Iowait()))
	metrics.GetOrRegisterGauge("cpu.irq", sr.hostRegistry).Update(int64(stat.Cpu.Irq()))
	metrics.GetOrRegisterGauge("cpu.softirq", sr.hostRegistry).Update(int64(stat.Cpu.Softirq()))
	var steal int64
	if stat.Cpu.StealSupported() {
		steal = int64(stat.Cpu.Steal())
	}
	metrics.GetOrRegisterGauge("cpu.steal", sr.hostRegistry).Update(steal)

	meminfo, err := linux.ReadMeminfo()
	if err != nil {
		glog.Errorf("could not read meminfo: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("memory.total", sr.hostRegistry).Update(int64(meminfo.MemTotal))
	metrics.GetOrRegisterGauge("memory.free", sr.hostRegistry).Update(int64(meminfo.MemFree))
	metrics.GetOrRegisterGauge("memory.buffers", sr.hostRegistry).Update(int64(meminfo.Buffers))
	metrics.GetOrRegisterGauge("memory.cached", sr.hostRegistry).Update(int64(meminfo.Cached))
	actualFree := int64(meminfo.MemFree) + int64(meminfo.Buffers) + int64(meminfo.Cached)
	metrics.GetOrRegisterGauge("memory.actualfree", sr.hostRegistry).Update(actualFree)
	metrics.GetOrRegisterGauge("memory.actualused", sr.hostRegistry).Update(int64(meminfo.MemTotal) - actualFree)
	metrics.GetOrRegisterGauge("swap.total", sr.hostRegistry).Update(int64(meminfo.SwapTotal))
	metrics.GetOrRegisterGauge("swap.free", sr.hostRegistry).Update(int64(meminfo.SwapFree))

	vmstat, err := linux.ReadVmstat()
	if err != nil {
		glog.Errorf("could not read vmstat: %s", err)
		return
	}
	metrics.GetOrRegisterGauge("vmstat.pgfault", sr.hostRegistry).Update(int64(vmstat.Pgfault))
	metrics.GetOrRegisterGauge("vmstat.pgmajfault", sr.hostRegistry).Update(int64(vmstat.Pgmajfault))
	metrics.GetOrRegisterGauge("vmstat.pgpgout", sr.hostRegistry).Update(int64(vmstat.Pgpgout) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pgpgin", sr.hostRegistry).Update(int64(vmstat.Pgpgin) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pswpout", sr.hostRegistry).Update(int64(vmstat.Pswpout) * 1024)
	metrics.GetOrRegisterGauge("vmstat.pswpin", sr.hostRegistry).Update(int64(vmstat.Pswpin) * 1024)

	if openFileDescriptorCount, err := GetOpenFileDescriptorCount(); err != nil {
		glog.Warningf("Couldn't get open file descriptor count", err)
	} else {
		metrics.GetOrRegisterGauge("Serviced.OpenFileDescriptors", sr.hostRegistry).Update(openFileDescriptorCount)
	}
}

// Updates the default registry.
func (sr *StatsReporter) updateStats() {
	// Stats for host.
	sr.updateHostStats()
	// Stats for the containers.
	var running []dao.RunningService
	running, err := zkservice.LoadRunningServicesByHost(sr.conn, sr.hostID)
	if err != nil {
		glog.Errorf("updateStats: zkservice.LoadRunningServicesByHost (conn: %+v hostID: %v) failed: %v", sr.conn, sr.hostID, err)
	}

	for _, rs := range running {
		if rs.DockerID != "" {
			containerRegistry := sr.getOrCreateContainerRegistry(rs.ServiceID, rs.InstanceID)
			if cpuacctStat, err := cgroup.ReadCpuacctStat(cgroup.GetCgroupDockerStatsFilePath(rs.DockerID, cgroup.Cpuacct)); err != nil {
				glog.V(4).Infof("Couldn't read CpuacctStat:", err)
			} else {
				metrics.GetOrRegisterGauge("cgroup.cpuacct.system", containerRegistry).Update(cpuacctStat.System)
				metrics.GetOrRegisterGauge("cgroup.cpuacct.user", containerRegistry).Update(cpuacctStat.User)
			}
			if memoryStat, err := cgroup.ReadMemoryStat(cgroup.GetCgroupDockerStatsFilePath(rs.DockerID, cgroup.Memory)); err != nil {
				glog.V(4).Infof("Couldn't read MemoryStat:", err)
			} else {
				metrics.GetOrRegisterGauge("cgroup.memory.pgmajfault", containerRegistry).Update(memoryStat.Pgfault)
				metrics.GetOrRegisterGauge("cgroup.memory.totalrss", containerRegistry).Update(memoryStat.TotalRss)
				metrics.GetOrRegisterGauge("cgroup.memory.cache", containerRegistry).Update(memoryStat.Cache)
			}
		} else {
			glog.V(4).Infof("Skipping stats update for %s (%s), no container ID exists yet", rs.Name, rs.ServiceID)
		}
	}
	// Clean out old container registries
	sr.removeStaleRegistries(&running)
}

// Fills out the metric consumer format.
func (sr *StatsReporter) gatherStats(t time.Time) []Sample {
	stats := []Sample{}
	// Handle the host metrics.
	reg, _ := sr.hostRegistry.(*metrics.StandardRegistry)
	reg.Each(func(name string, i interface{}) {
		if metric, ok := i.(metrics.Gauge); ok {
			tagmap := make(map[string]string)
			tagmap["controlplane_host_id"] = sr.hostID
			stats = append(stats, Sample{name, strconv.FormatInt(metric.Value(), 10), t.Unix(), tagmap})
		}
		if metricf64, ok := i.(metrics.GaugeFloat64); ok {
			tagmap := make(map[string]string)
			tagmap["controlplane_host_id"] = sr.hostID
			stats = append(stats, Sample{name, strconv.FormatFloat(metricf64.Value(), 'f', -1, 32), t.Unix(), tagmap})
		}
	})
	// Handle each container's metrics.
	for key, registry := range sr.containerRegistries {
		reg, _ := registry.(*metrics.StandardRegistry)
		reg.Each(func(name string, i interface{}) {
			if metric, ok := i.(metrics.Gauge); ok {
				tagmap := make(map[string]string)
				tagmap["controlplane_service_id"] = key.serviceID
				tagmap["controlplane_instance_id"] = strconv.FormatInt(int64(key.instanceID), 10)
				tagmap["controlplane_host_id"] = sr.hostID
				stats = append(stats, Sample{name, strconv.FormatInt(metric.Value(), 10), t.Unix(), tagmap})
			}
		})
	}
	return stats
}

// Send the list of stats to the TSDB.
func Post(destination string, stats []Sample) error {
	payload := map[string][]Sample{"metrics": stats}
	data, err := json.Marshal(payload)
	if err != nil {
		glog.Warningf("Couldn't marshal stats: ", err)
		return err
	}
	statsreq, err := http.NewRequest("POST", destination, bytes.NewBuffer(data))
	if err != nil {
		glog.Warningf("Couldn't create stats request: ", err)
		return err
	}
	statsreq.Header["User-Agent"] = []string{"Zenoss Metric Publisher"}
	statsreq.Header["Content-Type"] = []string{"application/json"}
	resp, reqerr := http.DefaultClient.Do(statsreq)
	if reqerr != nil {
		glog.Warningf("Couldn't post stats: ", reqerr)
		return reqerr
	}
	if !strings.Contains(resp.Status, "200 OK") {
		glog.Warningf("couldn't post stats: ", resp.Status)
		return nil
	}
	resp.Body.Close()
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package stats implements a means of tracking processing statistics for a
// BitTorrent tracker.
package stats

import (
	"time"

	"github.com/chihaya/chihaya/config"
	"github.com/pushrax/flatjson"
)

const (
	Announce = iota
	Scrape

	Completed
	NewLeech
	DeletedLeech
	ReapedLeech
	NewSeed
	DeletedSeed
	ReapedSeed

	NewTorrent
	DeletedTorrent
	ReapedTorrent

	AcceptedConnection
	ClosedConnection

	HandledRequest
	ErroredRequest
	ClientError

	ResponseTime
)

// DefaultStats is a default instance of stats tracking that uses an unbuffered
// channel for broadcasting events unless specified otherwise via a command
// line flag.
var DefaultStats *Stats

type PeerClassStats struct {
	Current uint64 // Current peer count.
	Joined  uint64 // Peers that announced.
	Left    uint64 // Peers that paused or stopped.
	Reaped  uint64 // Peers cleaned up after inactivity.
}

type PeerStats struct {
	PeerClassStats `json:"Peers"` // Stats for all peers.

	Seeds     PeerClassStats // Stats for seeds only.
	Completed uint64         // Number of transitions from leech to seed.
}

type PercentileTimes struct {
	P50 *Percentile
	P90 *Percentile
	P95 *Percentile
}

type Stats struct {
	Started time.Time // Time at which Chihaya was booted.

	OpenConnections     uint64 `json:"Connections.Open"`
	ConnectionsAccepted uint64 `json:"Connections.Accepted"`
	BytesTransmitted    uint64 `json:"BytesTransmitted"`

	RequestsHandled uint64 `json:"Requests.Handled"`
	RequestsErrored uint64 `json:"Requests.Errored"`
	ClientErrors    uint64 `json:"Requests.Bad"`
	ResponseTime    PercentileTimes

	Announces uint64 `json:"Tracker.Announces"`
	Scrapes   uint64 `json:"Tracker.Scrapes"`

	TorrentsAdded   uint64 `json:"Torrents.Added"`
	TorrentsRemoved uint64 `json:"Torrents.Removed"`
	TorrentsReaped  uint64 `json:"Torrents.Reaped"`

	IPv4Peers PeerStats `json:"Peers.IPv4"`
	IPv6Peers PeerStats `json:"Peers.IPv6"`

	*MemStatsWrapper `json:",omitempty"`

	events             chan int
	ipv4PeerEvents     chan int
	ipv6PeerEvents     chan int
	responseTimeEvents chan time.Duration
	recordMemStats     <-chan time.Time

	flattened flatjson.FlatMap
}

func New(cfg config.StatsConfig) *Stats {
	s := &Stats{
		Started: time.Now(),
		events:  make(chan int, cfg.BufferSize),

		ipv4PeerEvents:     make(chan int, cfg.BufferSize),
		ipv6PeerEvents:     make(chan int, cfg.BufferSize),
		responseTimeEvents: make(chan time.Duration, cfg.BufferSize),

		ResponseTime: PercentileTimes{
			P50: NewPercentile(0.5),
			P90: NewPercentile(0.9),
			P95: NewPercentile(0.95),
		},
	}

	if cfg.IncludeMem {
		s.MemStatsWrapper = NewMemStatsWrapper(cfg.VerboseMem)
		s.recordMemStats = time.NewTicker(cfg.MemUpdateInterval.Duration).C
	}

	s.flattened = flatjson.Flatten(s)
	go s.handleEvents()
	return s
}

func (s *Stats) Flattened() flatjson.FlatMap {
	return s.flattened
}

func (s *Stats) Close() {
	close(s.events)
}

func (s *Stats) Uptime() time.Duration {
	return time.Since(s.Started)
}

func (s *Stats) RecordEvent(event int) {
	s.events <- event
}

func (s *Stats) RecordPeerEvent(event int, ipv6 bool) {
	if ipv6 {
		s.ipv6PeerEvents <- event
	} else {
		s.ipv4PeerEvents <- event
	}
}

func (s *Stats) RecordTiming(event int, duration time.Duration) {
	switch event {
	case ResponseTime:
		s.responseTimeEvents <- duration
	default:
		panic("stats: RecordTiming called with an unknown event")
	}
}

func (s *Stats) handleEvents() {
	for {
		select {
		case event := <-s.events:
			s.handleEvent(event)

		case event := <-s.ipv4PeerEvents:
			s.handlePeerEvent(&s.IPv4Peers, event)

		case event := <-s.ipv6PeerEvents:
			s.handlePeerEvent(&s.IPv6Peers, event)

		case duration := <-s.responseTimeEvents:
			f := float64(duration) / float64(time.Millisecond)
			s.ResponseTime.P50.AddSample(f)
			s.ResponseTime.P90.AddSample(f)
			s.ResponseTime.P95.AddSample(f)

		case <-s.recordMemStats:
			s.MemStatsWrapper.Update()
		}
	}
}

func (s *Stats) handleEvent(event int) {
	switch event {
	case Announce:
		s.Announces++

	case Scrape:
		s.Scrapes++

	case NewTorrent:
		s.TorrentsAdded++

	case DeletedTorrent:
		s.TorrentsRemoved++

	case ReapedTorrent:
		s.TorrentsReaped++

	case AcceptedConnection:
		s.ConnectionsAccepted++
		s.OpenConnections++

	case ClosedConnection:
		s.OpenConnections--

	case HandledRequest:
		s.RequestsHandled++

	case ClientError:
		s.ClientErrors++

	case ErroredRequest:
		s.RequestsErrored++

	default:
		panic("stats: RecordEvent called with an unknown event")
	}
}

func (s *Stats) handlePeerEvent(ps *PeerStats, event int) {
	switch event {
	case Completed:
		ps.Completed++
		ps.Seeds.Current++

	case NewLeech:
		ps.Joined++
		ps.Current++

	case DeletedLeech:
		ps.Left++
		ps.Current--

	case ReapedLeech:
		ps.Reaped++
		ps.Current--

	case NewSeed:
		ps.Seeds.Joined++
		ps.Seeds.Current++
		ps.Joined++
		ps.Current++

	case DeletedSeed:
		ps.Seeds.Left++
		ps.Seeds.Current--
		ps.Left++
		ps.Current--

	case ReapedSeed:
		ps.Seeds.Reaped++
		ps.Seeds.Current--
		ps.Reaped++
		ps.Current--

	default:
		panic("stats: RecordPeerEvent called with an unknown event")
	}
}

// RecordEvent broadcasts an event to the default stats queue.
func RecordEvent(event int) {
	DefaultStats.RecordEvent(event)
}

// RecordPeerEvent broadcasts a peer event to the default stats queue.
func RecordPeerEvent(event int, ipv6 bool) {
	DefaultStats.RecordPeerEvent(event, ipv6)
}

// RecordTiming broadcasts a timing event to the default stats queue.
func RecordTiming(event int, duration time.Duration) {
	DefaultStats.RecordTiming(event, duration)
}
<commit_msg>Update references to FlatMap<commit_after>// Copyright 2014 The Chihaya Authors. All rights reserved.
// Use of this source code is governed by the BSD 2-Clause license,
// which can be found in the LICENSE file.

// Package stats implements a means of tracking processing statistics for a
// BitTorrent tracker.
package stats

import (
	"time"

	"github.com/chihaya/chihaya/config"
	"github.com/pushrax/flatjson"
)

const (
	Announce = iota
	Scrape

	Completed
	NewLeech
	DeletedLeech
	ReapedLeech
	NewSeed
	DeletedSeed
	ReapedSeed

	NewTorrent
	DeletedTorrent
	ReapedTorrent

	AcceptedConnection
	ClosedConnection

	HandledRequest
	ErroredRequest
	ClientError

	ResponseTime
)

// DefaultStats is a default instance of stats tracking that uses an unbuffered
// channel for broadcasting events unless specified otherwise via a command
// line flag.
var DefaultStats *Stats

type PeerClassStats struct {
	Current uint64 // Current peer count.
	Joined  uint64 // Peers that announced.
	Left    uint64 // Peers that paused or stopped.
	Reaped  uint64 // Peers cleaned up after inactivity.
}

type PeerStats struct {
	PeerClassStats `json:"Peers"` // Stats for all peers.

	Seeds     PeerClassStats // Stats for seeds only.
	Completed uint64         // Number of transitions from leech to seed.
}

type PercentileTimes struct {
	P50 *Percentile
	P90 *Percentile
	P95 *Percentile
}

type Stats struct {
	Started time.Time // Time at which Chihaya was booted.

	OpenConnections     uint64 `json:"Connections.Open"`
	ConnectionsAccepted uint64 `json:"Connections.Accepted"`
	BytesTransmitted    uint64 `json:"BytesTransmitted"`

	RequestsHandled uint64 `json:"Requests.Handled"`
	RequestsErrored uint64 `json:"Requests.Errored"`
	ClientErrors    uint64 `json:"Requests.Bad"`
	ResponseTime    PercentileTimes

	Announces uint64 `json:"Tracker.Announces"`
	Scrapes   uint64 `json:"Tracker.Scrapes"`

	TorrentsAdded   uint64 `json:"Torrents.Added"`
	TorrentsRemoved uint64 `json:"Torrents.Removed"`
	TorrentsReaped  uint64 `json:"Torrents.Reaped"`

	IPv4Peers PeerStats `json:"Peers.IPv4"`
	IPv6Peers PeerStats `json:"Peers.IPv6"`

	*MemStatsWrapper `json:",omitempty"`

	events             chan int
	ipv4PeerEvents     chan int
	ipv6PeerEvents     chan int
	responseTimeEvents chan time.Duration
	recordMemStats     <-chan time.Time

	flattened flatjson.Map
}

func New(cfg config.StatsConfig) *Stats {
	s := &Stats{
		Started: time.Now(),
		events:  make(chan int, cfg.BufferSize),

		ipv4PeerEvents:     make(chan int, cfg.BufferSize),
		ipv6PeerEvents:     make(chan int, cfg.BufferSize),
		responseTimeEvents: make(chan time.Duration, cfg.BufferSize),

		ResponseTime: PercentileTimes{
			P50: NewPercentile(0.5),
			P90: NewPercentile(0.9),
			P95: NewPercentile(0.95),
		},
	}

	if cfg.IncludeMem {
		s.MemStatsWrapper = NewMemStatsWrapper(cfg.VerboseMem)
		s.recordMemStats = time.NewTicker(cfg.MemUpdateInterval.Duration).C
	}

	s.flattened = flatjson.Flatten(s)
	go s.handleEvents()
	return s
}

func (s *Stats) Flattened() flatjson.Map {
	return s.flattened
}

func (s *Stats) Close() {
	close(s.events)
}

func (s *Stats) Uptime() time.Duration {
	return time.Since(s.Started)
}

func (s *Stats) RecordEvent(event int) {
	s.events <- event
}

func (s *Stats) RecordPeerEvent(event int, ipv6 bool) {
	if ipv6 {
		s.ipv6PeerEvents <- event
	} else {
		s.ipv4PeerEvents <- event
	}
}

func (s *Stats) RecordTiming(event int, duration time.Duration) {
	switch event {
	case ResponseTime:
		s.responseTimeEvents <- duration
	default:
		panic("stats: RecordTiming called with an unknown event")
	}
}

func (s *Stats) handleEvents() {
	for {
		select {
		case event := <-s.events:
			s.handleEvent(event)

		case event := <-s.ipv4PeerEvents:
			s.handlePeerEvent(&s.IPv4Peers, event)

		case event := <-s.ipv6PeerEvents:
			s.handlePeerEvent(&s.IPv6Peers, event)

		case duration := <-s.responseTimeEvents:
			f := float64(duration) / float64(time.Millisecond)
			s.ResponseTime.P50.AddSample(f)
			s.ResponseTime.P90.AddSample(f)
			s.ResponseTime.P95.AddSample(f)

		case <-s.recordMemStats:
			s.MemStatsWrapper.Update()
		}
	}
}

func (s *Stats) handleEvent(event int) {
	switch event {
	case Announce:
		s.Announces++

	case Scrape:
		s.Scrapes++

	case NewTorrent:
		s.TorrentsAdded++

	case DeletedTorrent:
		s.TorrentsRemoved++

	case ReapedTorrent:
		s.TorrentsReaped++

	case AcceptedConnection:
		s.ConnectionsAccepted++
		s.OpenConnections++

	case ClosedConnection:
		s.OpenConnections--

	case HandledRequest:
		s.RequestsHandled++

	case ClientError:
		s.ClientErrors++

	case ErroredRequest:
		s.RequestsErrored++

	default:
		panic("stats: RecordEvent called with an unknown event")
	}
}

func (s *Stats) handlePeerEvent(ps *PeerStats, event int) {
	switch event {
	case Completed:
		ps.Completed++
		ps.Seeds.Current++

	case NewLeech:
		ps.Joined++
		ps.Current++

	case DeletedLeech:
		ps.Left++
		ps.Current--

	case ReapedLeech:
		ps.Reaped++
		ps.Current--

	case NewSeed:
		ps.Seeds.Joined++
		ps.Seeds.Current++
		ps.Joined++
		ps.Current++

	case DeletedSeed:
		ps.Seeds.Left++
		ps.Seeds.Current--
		ps.Left++
		ps.Current--

	case ReapedSeed:
		ps.Seeds.Reaped++
		ps.Seeds.Current--
		ps.Reaped++
		ps.Current--

	default:
		panic("stats: RecordPeerEvent called with an unknown event")
	}
}

// RecordEvent broadcasts an event to the default stats queue.
func RecordEvent(event int) {
	DefaultStats.RecordEvent(event)
}

// RecordPeerEvent broadcasts a peer event to the default stats queue.
func RecordPeerEvent(event int, ipv6 bool) {
	DefaultStats.RecordPeerEvent(event, ipv6)
}

// RecordTiming broadcasts a timing event to the default stats queue.
func RecordTiming(event int, duration time.Duration) {
	DefaultStats.RecordTiming(event, duration)
}
<|endoftext|>