<commit_before>package gssapi

/*
#cgo LDFLAGS: -lsasl2
#include <sasl/sasl.h>
#include <stdlib.h>
*/
import "C"

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"unsafe"
)

// sendToken send token to service
func sendToken(conn io.Writer, buf []byte) error {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, uint32(len(buf)))
	_, err := conn.Write(b)
	if err != nil {
		return err
	}
	conn.Write(buf)
	if err != nil {
		return err
	}
	return nil
}

// recvToken recive from a token from service
func recvToken(conn io.ReadWriter) ([]byte, error) {
	b := make([]byte, 4)
	_, err := conn.Read(b)
	if err != nil {
		return b, err
	}

	size := binary.BigEndian.Uint32(b)
	buf := make([]byte, size)
	_, err = conn.Read(buf)

	return buf, err
}

func saslErr(ctx *C.sasl_conn_t) error {
	if ctx == nil {
		return nil
	}
	str := C.GoString(C.sasl_errdetail(ctx))
	return errors.New(str)
}

// GssAuth complete handling of establishing SASL/Kerberos Authentication
func GssAuth(rw io.ReadWriter, service string, realm string) error {
	errCode := C.sasl_client_init(nil)
	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_init failed")
	}

	var ctx *C.sasl_conn_t

	cservice := C.CString(service)
	defer C.free(unsafe.Pointer(cservice))

	crealm := C.CString(realm)
	defer C.free(unsafe.Pointer(crealm))

	errCode = C.sasl_client_new(cservice,
		crealm,
		nil,
		nil,
		nil,
		C.uint(0),
		&ctx)

	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_new cannot establish new context, %v", saslErr(ctx))
	}

	defer C.sasl_dispose(&ctx)

	var out *C.char
	var outlen C.uint
	mech := C.CString("GSSAPI")
	defer C.free(unsafe.Pointer(mech))

	errCode = C.sasl_client_start(ctx,
		mech,
		nil,
		&out,
		&outlen,
		nil)

	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_start failed, %v", saslErr(ctx))
	}
	defer C.sasl_client_done()

	// send initial requestToken
	reqtoken := C.GoBytes(unsafe.Pointer(out), C.int(outlen))
	err := sendToken(rw, reqtoken)
	if err != nil {
		return fmt.Errorf("send token to service failed, %v", err)
	}

	for errCode == C.SASL_CONTINUE {
		reptoken, err := recvToken(rw)
		if err != nil {
			return fmt.Errorf("recive token from service failed, %v", err)
		}

		in := unsafe.Pointer(C.CString(string(reptoken)))
		defer C.free(in)

		errCode = C.sasl_client_step(ctx,
			(*C.char)(in),
			C.uint(len(reptoken)),
			nil,
			&out,
			&outlen)

		reqtoken = C.GoBytes(unsafe.Pointer(out), C.int(outlen))
		err = sendToken(rw, reqtoken)
		if err != nil {
			return fmt.Errorf("send token to service failed, %v", err)
		}
	}

	if errCode != C.SASL_OK {
		return fmt.Errorf("Authentication handshake was not completed, %v", saslErr(ctx))
	}

	return nil
}
<commit_msg>move defer C.sasl_client_done()before defer C.sasl_dispose(&ctx)<commit_after>package gssapi

/*
#cgo LDFLAGS: -lsasl2
#include <sasl/sasl.h>
#include <stdlib.h>
*/
import "C"

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"unsafe"
)

// sendToken send token to service
func sendToken(conn io.Writer, buf []byte) error {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, uint32(len(buf)))
	_, err := conn.Write(b)
	if err != nil {
		return err
	}
	conn.Write(buf)
	if err != nil {
		return err
	}
	return nil
}

// recvToken recive from a token from service
func recvToken(conn io.ReadWriter) ([]byte, error) {
	b := make([]byte, 4)
	_, err := conn.Read(b)
	if err != nil {
		return b, err
	}

	size := binary.BigEndian.Uint32(b)
	buf := make([]byte, size)
	_, err = conn.Read(buf)

	return buf, err
}

func saslErr(ctx *C.sasl_conn_t) error {
	if ctx == nil {
		return nil
	}
	str := C.GoString(C.sasl_errdetail(ctx))
	return errors.New(str)
}

// GssAuth complete handling of establishing SASL/Kerberos Authentication
func GssAuth(rw io.ReadWriter, service string, realm string) error {
	errCode := C.sasl_client_init(nil)
	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_init failed")
	}
	defer C.sasl_client_done()

	var ctx *C.sasl_conn_t

	cservice := C.CString(service)
	defer C.free(unsafe.Pointer(cservice))

	crealm := C.CString(realm)
	defer C.free(unsafe.Pointer(crealm))

	errCode = C.sasl_client_new(cservice,
		crealm,
		nil,
		nil,
		nil,
		C.uint(0),
		&ctx)

	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_new cannot establish new context, %v", saslErr(ctx))
	}
	defer C.sasl_dispose(&ctx)

	var out *C.char
	var outlen C.uint
	mech := C.CString("GSSAPI")
	defer C.free(unsafe.Pointer(mech))

	errCode = C.sasl_client_start(ctx,
		mech,
		nil,
		&out,
		&outlen,
		nil)

	if errCode != C.SASL_OK {
		return fmt.Errorf("sasl_client_start failed, %v", saslErr(ctx))
	}

	// send initial requestToken
	reqtoken := C.GoBytes(unsafe.Pointer(out), C.int(outlen))
	err := sendToken(rw, reqtoken)
	if err != nil {
		return fmt.Errorf("send token to service failed, %v", err)
	}

	for errCode == C.SASL_CONTINUE {
		reptoken, err := recvToken(rw)
		if err != nil {
			return fmt.Errorf("recive token from service failed, %v", err)
		}

		in := unsafe.Pointer(C.CString(string(reptoken)))
		defer C.free(in)

		errCode = C.sasl_client_step(ctx,
			(*C.char)(in),
			C.uint(len(reptoken)),
			nil,
			&out,
			&outlen)

		reqtoken = C.GoBytes(unsafe.Pointer(out), C.int(outlen))
		err = sendToken(rw, reqtoken)
		if err != nil {
			return fmt.Errorf("send token to service failed, %v", err)
		}
	}

	if errCode != C.SASL_OK {
		return fmt.Errorf("Authentication handshake was not completed, %v", saslErr(ctx))
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"crypto/sha1"
	"github.com/miekg/dns"
	"log"
	"os"
	"strings"
	"sync"
	"time"
)

const origTTL uint32 = 60

var cache *sigCache = newCache()
var inflight *single = new(single)

// TODO(miek): allow RRSIG queries and NSEC queries?

// ParseKeyFile read a DNSSEC keyfile as generated by dnssec-keygen or other
// utilities. It add ".key" for the public key and ".private" for the private key.
func ParseKeyFile(file string) (*dns.DNSKEY, dns.PrivateKey, error) {
	f, e := os.Open(file + ".key")
	if e != nil {
		return nil, nil, e
	}
	k, e := dns.ReadRR(f, file+".key")
	if e != nil {
		return nil, nil, e
	}
	f, e = os.Open(file + ".private")
	if e != nil {
		return nil, nil, e
	}
	p, e := k.(*dns.DNSKEY).ReadPrivateKey(f, file+".private")
	if e != nil {
		return nil, nil, e
	}
	return k.(*dns.DNSKEY), p, nil
}

// nsec creates (if needed) NSEC records that are included in the reply.
func (s *Server) nsec(m *dns.Msg) {
	if m.Rcode == dns.RcodeNameError {
		// qname nsec
		nsec1 := s.newNSEC(m.Question[0].Name)
		m.Ns = append(m.Ns, nsec1)
		// wildcard nsec
		idx := dns.Split(m.Question[0].Name)
		wildcard := "*." + m.Question[0].Name[idx[0]:]
		nsec2 := s.newNSEC(wildcard)
		if nsec1.Hdr.Name != nsec2.Hdr.Name || nsec1.NextDomain != nsec2.NextDomain {
			// different NSEC, add it
			m.Ns = append(m.Ns, nsec2)
		}
	}
	if m.Rcode == dns.RcodeSuccess && len(m.Ns) == 1 {
		if _, ok := m.Ns[0].(*dns.SOA); ok {
			m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
		}
	}
}

// sign signs a message m, it takes care of negative or nodata responses as
// well by synthesising NSEC records. It will also cache the signatures, using
// a hash of the signed data as a key.
// We also fake the origin TTL in the signature, because we don't want to
// throw away signatures when services decide to have longer TTL. So we just
// set the origTTL to 60.
func (s *Server) sign(m *dns.Msg, bufsize uint16) {
	now := time.Now().UTC()
	incep := uint32(now.Add(-2 * time.Hour).Unix())     // 2 hours, be sure to catch daylight saving time and such
	expir := uint32(now.Add(7 * 24 * time.Hour).Unix()) // sign for a week

	// TODO(miek): repeating this two times?
	for _, r := range rrSets(m.Answer) {
		if r[0].Header().Rrtype == dns.TypeRRSIG {
			continue
		}
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Answer = append(m.Answer, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Answer = append(m.Answer, dns.Copy(sig).(*dns.RRSIG))
	}
	for _, r := range rrSets(m.Ns) {
		if r[0].Header().Rrtype == dns.TypeRRSIG {
			continue
		}
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Ns = append(m.Ns, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Ns = append(m.Ns, dns.Copy(sig).(*dns.RRSIG))
	}
	// TODO(miek): Forget the additional section for now
	if bufsize >= 512 || bufsize <= 4096 {
		m.Truncated = m.Len() > int(bufsize)
	}
	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	o.SetDo()
	o.SetUDPSize(4096)
	m.Extra = append(m.Extra, o)
	return
}

func (s *Server) newRRSIG(incep, expir uint32) *dns.RRSIG {
	sig := new(dns.RRSIG)
	sig.Hdr.Rrtype = dns.TypeRRSIG
	sig.Hdr.Ttl = origTTL
	sig.OrigTtl = origTTL
	sig.Algorithm = s.Dnskey.Algorithm
	sig.KeyTag = s.KeyTag
	sig.Inception = incep
	sig.Expiration = expir
	sig.SignerName = s.Dnskey.Hdr.Name
	return sig
}

// newNSEC returns the NSEC record need to denial qname, or gives back a NODATA NSEC.
func (s *Server) newNSEC(qname string) *dns.NSEC {
	qlabels := dns.SplitDomainName(qname)
	if len(qlabels) < s.domainLabels {
		// TODO(miek): can not happen...?
	}
	// Strip the last s.domainLabels, return up to 4 before
	// that. Four labels is the maximum qname we can handle.
	ls := len(qlabels) - s.domainLabels
	ls4 := ls - 4
	if ls4 < 0 {
		ls4 = 0
	}
	key := qlabels[ls4:ls]
	prev, next := s.registry.GetNSEC(strings.Join(key, "."))
	nsec := &dns.NSEC{Hdr: dns.RR_Header{Name: prev + s.domain + ".", Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60},
		NextDomain: next + s.domain + "."}
	if prev == "" {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeSOA, dns.TypeNS, dns.TypeAAAA, dns.TypeRRSIG, dns.TypeNSEC, dns.TypeDNSKEY}
	} else {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeSRV, dns.TypeRRSIG, dns.TypeNSEC}
	}
	return nsec
}

type rrset struct {
	qname string
	qtype uint16
}

func rrSets(rrs []dns.RR) map[rrset][]dns.RR {
	m := make(map[rrset][]dns.RR)
	for _, r := range rrs {
		if s, ok := m[rrset{r.Header().Name, r.Header().Rrtype}]; ok {
			s = append(s, r)
			m[rrset{r.Header().Name, r.Header().Rrtype}] = s
		} else {
			s := make([]dns.RR, 1, 3)
			s[0] = r
			m[rrset{r.Header().Name, r.Header().Rrtype}] = s
		}
	}
	if len(m) > 0 {
		return m
	}
	return nil
}

type sigCache struct {
	sync.RWMutex
	m map[string]*dns.RRSIG
}

func newCache() *sigCache {
	c := new(sigCache)
	c.m = make(map[string]*dns.RRSIG)
	return c
}

func (c *sigCache) remove(s string) {
	delete(c.m, s)
}

func (c *sigCache) insert(s string, r *dns.RRSIG) {
	c.Lock()
	defer c.Unlock()
	if _, ok := c.m[s]; !ok {
		c.m[s] = r
	}
}

func (c *sigCache) search(s string) *dns.RRSIG {
	c.RLock()
	defer c.RUnlock()
	if s, ok := c.m[s]; ok {
		// we want to return a copy here, because if we didn't the RRSIG
		// could be removed by another goroutine before the packet containing
		// this signature is send out.
		log.Println("DNS Signature retrieved from cache")
		return dns.Copy(s).(*dns.RRSIG)
	}
	return nil
}

// get return a bunch of signatures which have their owner name set to owner name.
// This is an expensive function, because we loop the entire map.
func (c *sigCache) get(ownername string) []*dns.RRSIG {
	return nil
}

// key uses the name, type and rdata, which is serialized and then hashed as the
// key for the lookup
func (c *sigCache) key(rrs []dns.RR) string {
	h := sha1.New()
	i := []byte(rrs[0].Header().Name)
	i = append(i, packUint16(rrs[0].Header().Rrtype)...)
	for _, r := range rrs {
		switch t := r.(type) { // we only do a few type, serialize these manually
		case *dns.SOA:
			i = append(i, packUint32(t.Serial)...)
			// we only fiddle with the serial so store that
		case *dns.SRV:
			i = append(i, packUint16(t.Priority)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, []byte(t.Target)...)
		case *dns.A:
			i = append(i, []byte(t.A)...)
		case *dns.AAAA:
			i = append(i, []byte(t.AAAA)...)
		case *dns.DNSKEY:
			// Need nothing more, the rdata stays the same during a run
		case *dns.NSEC:
			i = append(i, []byte(t.NextDomain)...)
			// bitmap does not differentiate
		default:
			log.Printf("DNS Signature for unhandled type %T seen", t)
		}
	}
	return string(h.Sum(i))
}

// TODO(miek): prolly should use the stdlib ones
func packUint16(i uint16) []byte { return []byte{byte(i >> 8), byte(i)} }
func packUint32(i uint32) []byte { return []byte{byte(i >> 24), byte(i >> 16), byte(i >> 8), byte(i)} }

// Adapted from singleinflight.go from the original Go Code. Copyright 2013 The Go Authors.
type call struct {
	wg   sync.WaitGroup
	val  *dns.RRSIG
	err  error
	dups int
}

type single struct {
	sync.Mutex
	m map[string]*call
}

func (g *single) Do(key string, fn func() (*dns.RRSIG, error)) (*dns.RRSIG, error, bool) {
	g.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	if c, ok := g.m[key]; ok {
		c.dups++
		g.Unlock()
		c.wg.Wait()
		return c.val, c.err, true
	}
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.Unlock()

	c.val, c.err = fn()
	c.wg.Done()

	g.Lock()
	delete(g.m, key)
	g.Unlock()

	return c.val, c.err, c.dups > 0
}
<commit_msg>Remove this, maybe add later<commit_after>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package server

import (
	"crypto/sha1"
	"github.com/miekg/dns"
	"log"
	"os"
	"strings"
	"sync"
	"time"
)

const origTTL uint32 = 60

var cache *sigCache = newCache()
var inflight *single = new(single)

// TODO(miek): allow RRSIG queries and NSEC queries?

// ParseKeyFile read a DNSSEC keyfile as generated by dnssec-keygen or other
// utilities. It add ".key" for the public key and ".private" for the private key.
func ParseKeyFile(file string) (*dns.DNSKEY, dns.PrivateKey, error) {
	f, e := os.Open(file + ".key")
	if e != nil {
		return nil, nil, e
	}
	k, e := dns.ReadRR(f, file+".key")
	if e != nil {
		return nil, nil, e
	}
	f, e = os.Open(file + ".private")
	if e != nil {
		return nil, nil, e
	}
	p, e := k.(*dns.DNSKEY).ReadPrivateKey(f, file+".private")
	if e != nil {
		return nil, nil, e
	}
	return k.(*dns.DNSKEY), p, nil
}

// nsec creates (if needed) NSEC records that are included in the reply.
func (s *Server) nsec(m *dns.Msg) {
	if m.Rcode == dns.RcodeNameError {
		// qname nsec
		nsec1 := s.newNSEC(m.Question[0].Name)
		m.Ns = append(m.Ns, nsec1)
		// wildcard nsec
		idx := dns.Split(m.Question[0].Name)
		wildcard := "*." + m.Question[0].Name[idx[0]:]
		nsec2 := s.newNSEC(wildcard)
		if nsec1.Hdr.Name != nsec2.Hdr.Name || nsec1.NextDomain != nsec2.NextDomain {
			// different NSEC, add it
			m.Ns = append(m.Ns, nsec2)
		}
	}
	if m.Rcode == dns.RcodeSuccess && len(m.Ns) == 1 {
		if _, ok := m.Ns[0].(*dns.SOA); ok {
			m.Ns = append(m.Ns, s.newNSEC(m.Question[0].Name))
		}
	}
}

// sign signs a message m, it takes care of negative or nodata responses as
// well by synthesising NSEC records. It will also cache the signatures, using
// a hash of the signed data as a key.
// We also fake the origin TTL in the signature, because we don't want to
// throw away signatures when services decide to have longer TTL. So we just
// set the origTTL to 60.
func (s *Server) sign(m *dns.Msg, bufsize uint16) {
	now := time.Now().UTC()
	incep := uint32(now.Add(-2 * time.Hour).Unix())     // 2 hours, be sure to catch daylight saving time and such
	expir := uint32(now.Add(7 * 24 * time.Hour).Unix()) // sign for a week

	// TODO(miek): repeating this two times?
	for _, r := range rrSets(m.Answer) {
		if r[0].Header().Rrtype == dns.TypeRRSIG {
			continue
		}
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Answer = append(m.Answer, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Answer = append(m.Answer, dns.Copy(sig).(*dns.RRSIG))
	}
	for _, r := range rrSets(m.Ns) {
		if r[0].Header().Rrtype == dns.TypeRRSIG {
			continue
		}
		key := cache.key(r)
		if s := cache.search(key); s != nil {
			if s.ValidityPeriod(now.Add(-24 * time.Hour)) {
				m.Ns = append(m.Ns, s)
				continue
			}
			cache.remove(key)
		}
		sig, err, shared := inflight.Do(key, func() (*dns.RRSIG, error) {
			sig1 := s.newRRSIG(incep, expir)
			e := sig1.Sign(s.Privkey, r)
			if e != nil {
				log.Printf("Failed to sign: %s\n", e.Error())
			}
			return sig1, e
		})
		if err != nil {
			continue
		}
		if !shared {
			// is it possible to miss this, due the the c.dups > 0 in Do()? TODO(miek)
			cache.insert(key, sig)
		}
		m.Ns = append(m.Ns, dns.Copy(sig).(*dns.RRSIG))
	}
	// TODO(miek): Forget the additional section for now
	if bufsize >= 512 || bufsize <= 4096 {
		m.Truncated = m.Len() > int(bufsize)
	}
	o := new(dns.OPT)
	o.Hdr.Name = "."
	o.Hdr.Rrtype = dns.TypeOPT
	o.SetDo()
	o.SetUDPSize(4096)
	m.Extra = append(m.Extra, o)
	return
}

func (s *Server) newRRSIG(incep, expir uint32) *dns.RRSIG {
	sig := new(dns.RRSIG)
	sig.Hdr.Rrtype = dns.TypeRRSIG
	sig.Hdr.Ttl = origTTL
	sig.OrigTtl = origTTL
	sig.Algorithm = s.Dnskey.Algorithm
	sig.KeyTag = s.KeyTag
	sig.Inception = incep
	sig.Expiration = expir
	sig.SignerName = s.Dnskey.Hdr.Name
	return sig
}

// newNSEC returns the NSEC record need to denial qname, or gives back a NODATA NSEC.
func (s *Server) newNSEC(qname string) *dns.NSEC {
	qlabels := dns.SplitDomainName(qname)
	if len(qlabels) < s.domainLabels {
		// TODO(miek): can not happen...?
	}
	// Strip the last s.domainLabels, return up to 4 before
	// that. Four labels is the maximum qname we can handle.
	ls := len(qlabels) - s.domainLabels
	ls4 := ls - 4
	if ls4 < 0 {
		ls4 = 0
	}
	key := qlabels[ls4:ls]
	prev, next := s.registry.GetNSEC(strings.Join(key, "."))
	nsec := &dns.NSEC{Hdr: dns.RR_Header{Name: prev + s.domain + ".", Rrtype: dns.TypeNSEC, Class: dns.ClassINET, Ttl: 60},
		NextDomain: next + s.domain + "."}
	if prev == "" {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeSOA, dns.TypeNS, dns.TypeAAAA, dns.TypeRRSIG, dns.TypeNSEC, dns.TypeDNSKEY}
	} else {
		nsec.TypeBitMap = []uint16{dns.TypeA, dns.TypeAAAA, dns.TypeSRV, dns.TypeRRSIG, dns.TypeNSEC}
	}
	return nsec
}

type rrset struct {
	qname string
	qtype uint16
}

func rrSets(rrs []dns.RR) map[rrset][]dns.RR {
	m := make(map[rrset][]dns.RR)
	for _, r := range rrs {
		if s, ok := m[rrset{r.Header().Name, r.Header().Rrtype}]; ok {
			s = append(s, r)
			m[rrset{r.Header().Name, r.Header().Rrtype}] = s
		} else {
			s := make([]dns.RR, 1, 3)
			s[0] = r
			m[rrset{r.Header().Name, r.Header().Rrtype}] = s
		}
	}
	if len(m) > 0 {
		return m
	}
	return nil
}

type sigCache struct {
	sync.RWMutex
	m map[string]*dns.RRSIG
}

func newCache() *sigCache {
	c := new(sigCache)
	c.m = make(map[string]*dns.RRSIG)
	return c
}

func (c *sigCache) remove(s string) {
	delete(c.m, s)
}

func (c *sigCache) insert(s string, r *dns.RRSIG) {
	c.Lock()
	defer c.Unlock()
	if _, ok := c.m[s]; !ok {
		c.m[s] = r
	}
}

func (c *sigCache) search(s string) *dns.RRSIG {
	c.RLock()
	defer c.RUnlock()
	if s, ok := c.m[s]; ok {
		// we want to return a copy here, because if we didn't the RRSIG
		// could be removed by another goroutine before the packet containing
		// this signature is send out.
		log.Println("DNS Signature retrieved from cache")
		return dns.Copy(s).(*dns.RRSIG)
	}
	return nil
}

// key uses the name, type and rdata, which is serialized and then hashed as the
// key for the lookup
func (c *sigCache) key(rrs []dns.RR) string {
	h := sha1.New()
	i := []byte(rrs[0].Header().Name)
	i = append(i, packUint16(rrs[0].Header().Rrtype)...)
	for _, r := range rrs {
		switch t := r.(type) { // we only do a few type, serialize these manually
		case *dns.SOA:
			i = append(i, packUint32(t.Serial)...)
			// we only fiddle with the serial so store that
		case *dns.SRV:
			i = append(i, packUint16(t.Priority)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, packUint16(t.Weight)...)
			i = append(i, []byte(t.Target)...)
		case *dns.A:
			i = append(i, []byte(t.A)...)
		case *dns.AAAA:
			i = append(i, []byte(t.AAAA)...)
		case *dns.DNSKEY:
			// Need nothing more, the rdata stays the same during a run
		case *dns.NSEC:
			i = append(i, []byte(t.NextDomain)...)
			// bitmap does not differentiate
		default:
			log.Printf("DNS Signature for unhandled type %T seen", t)
		}
	}
	return string(h.Sum(i))
}

// TODO(miek): prolly should use the stdlib ones
func packUint16(i uint16) []byte { return []byte{byte(i >> 8), byte(i)} }
func packUint32(i uint32) []byte { return []byte{byte(i >> 24), byte(i >> 16), byte(i >> 8), byte(i)} }

// Adapted from singleinflight.go from the original Go Code. Copyright 2013 The Go Authors.
type call struct {
	wg   sync.WaitGroup
	val  *dns.RRSIG
	err  error
	dups int
}

type single struct {
	sync.Mutex
	m map[string]*call
}

func (g *single) Do(key string, fn func() (*dns.RRSIG, error)) (*dns.RRSIG, error, bool) {
	g.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	if c, ok := g.m[key]; ok {
		c.dups++
		g.Unlock()
		c.wg.Wait()
		return c.val, c.err, true
	}
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.Unlock()

	c.val, c.err = fn()
	c.wg.Done()

	g.Lock()
	delete(g.m, key)
	g.Unlock()

	return c.val, c.err, c.dups > 0
}
<|endoftext|>


<commit_before>package jiracmd

import (
	"fmt"
	"strings"

	"github.com/coryb/figtree"
	"github.com/coryb/oreo"

	"gopkg.in/Netflix-Skunkworks/go-jira.v1"
	"gopkg.in/Netflix-Skunkworks/go-jira.v1/jiracli"
	"gopkg.in/Netflix-Skunkworks/go-jira.v1/jiradata"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

type TransitionOptions struct {
	jiracli.CommonOptions `yaml:",inline" json:",inline" figtree:",inline"`
	Overrides             map[string]string `yaml:"overrides,omitempty" json:"overrides,omitempty"`
	Transition            string            `yaml:"transition,omitempty" json:"transition,omitempty"`
	Issue                 string            `yaml:"issue,omitempty" json:"issue,omitempty"`
	Resolution            string            `yaml:"resolution,omitempty" json:"resolution,omitempty"`
}

func CmdTransitionRegistry(transition string) *jiracli.CommandRegistryEntry {
	opts := TransitionOptions{
		CommonOptions: jiracli.CommonOptions{
			Template: figtree.NewStringOption("transition"),
		},
		Overrides: map[string]string{},
	}

	help := "Transition issue to given state"
	if transition != "" {
		help = fmt.Sprintf("Transition issue to %s state", transition)
		opts.SkipEditing = figtree.NewBoolOption(true)
	}

	return &jiracli.CommandRegistryEntry{
		help,
		func(fig *figtree.FigTree, cmd *kingpin.CmdClause) error {
			jiracli.LoadConfigs(cmd, fig, &opts)
			if opts.Transition == "" {
				opts.Transition = transition
			}
			return CmdTransitionUsage(cmd, &opts)
		},
		func(o *oreo.Client, globals *jiracli.GlobalOptions) error {
			return CmdTransition(o, globals, &opts)
		},
	}
}

func CmdTransitionUsage(cmd *kingpin.CmdClause, opts *TransitionOptions) error {
	jiracli.BrowseUsage(cmd, &opts.CommonOptions)
	jiracli.TemplateUsage(cmd, &opts.CommonOptions)
	cmd.Flag("noedit", "Disable opening the editor").SetValue(&opts.SkipEditing)
	cmd.Flag("comment", "Comment message for issue").Short('m').PreAction(func(ctx *kingpin.ParseContext) error {
		opts.Overrides["comment"] = jiracli.FlagValue(ctx, "comment")
		return nil
	}).String()
	cmd.Flag("override", "Set issue property").Short('o').StringMapVar(&opts.Overrides)
	if opts.Transition == "" {
		cmd.Arg("TRANSITION", "State to transition issue to").Required().StringVar(&opts.Transition)
	}
	cmd.Arg("ISSUE", "issue to transition").Required().StringVar(&opts.Issue)
	return nil
}

func defaultResolution(transMeta *jiradata.Transition) string {
	if resField, ok := transMeta.Fields["resolution"]; ok {
		for _, allowedValueRaw := range resField.AllowedValues {
			if allowedValue, ok := allowedValueRaw.(map[string]interface{}); ok {
				if allowedValue["name"] == "Fixed" {
					return "Fixed"
				} else if allowedValue["name"] == "Done" {
					return "Done"
				}
			}
		}
	}
	return ""
}

// CmdTransition will move state of the given issue to the given transtion
func CmdTransition(o *oreo.Client, globals *jiracli.GlobalOptions, opts *TransitionOptions) error {
	issueData, err := jira.GetIssue(o, globals.Endpoint.Value, opts.Issue, nil)
	if err != nil {
		return jiracli.CliError(err)
	}

	meta, err := jira.GetIssueTransitions(o, globals.Endpoint.Value, opts.Issue)
	if err != nil {
		return jiracli.CliError(err)
	}
	transMeta := meta.Transitions.Find(opts.Transition)

	if transMeta == nil {
		possible := []string{}
		for _, trans := range meta.Transitions {
			possible = append(possible, trans.Name)
		}

		if status, ok := issueData.Fields["status"].(map[string]interface{}); ok {
			if name, ok := status["name"].(string); ok {
				return jiracli.CliError(fmt.Errorf("Invalid Transition %q from %q, Available: %s", opts.Transition, name, strings.Join(possible, ", ")))
			}
		}
		return jiracli.CliError(fmt.Errorf("No valid transition found matching %s", opts.Transition))
	}

	// need to default the Resolution, usually Fixed works but sometime need Done
	if opts.Resolution == "" {
		if resField, ok := transMeta.Fields["resolution"]; ok {
			for _, allowedValueRaw := range resField.AllowedValues {
				if allowedValue, ok := allowedValueRaw.(map[string]interface{}); ok {
					if allowedValue["name"] == "Fixed" {
						opts.Resolution = "Fixed"
					} else if allowedValue["name"] == "Done" {
						opts.Resolution = "Done"
					}
				}
			}
		}
	}
	opts.Overrides["resolution"] = opts.Resolution

	type templateInput struct {
		*jiradata.Issue `yaml:",inline"`
		// Yes, Meta and Transition are redundant, but this is for backwards compatibility
		// with old templates
		Meta       *jiradata.Transition `yaml:"meta,omitempty" json:"meta,omitemtpy"`
		Transition *jiradata.Transition `yaml:"transition,omitempty" json:"transition,omitempty"`
		Overrides  map[string]string    `yaml:"overrides,omitempty" json:"overrides,omitempty"`
	}

	if _, ok := transMeta.Fields["comment"]; !ok && opts.Overrides["comment"] != "" {
		comment := jiradata.Comment{
			Body: opts.Overrides["comment"],
		}
		if _, err := jira.IssueAddComment(o, globals.Endpoint.Value, opts.Issue, &comment); err != nil {
			return err
		}
	}

	issueUpdate := jiradata.IssueUpdate{}
	input := templateInput{
		Issue:      issueData,
		Meta:       transMeta,
		Transition: transMeta,
		Overrides:  opts.Overrides,
	}
	err = jiracli.EditLoop(&opts.CommonOptions, &input, &issueUpdate, func() error {
		return jira.TransitionIssue(o, globals.Endpoint.Value, opts.Issue, &issueUpdate)
	})
	if err != nil {
		return jiracli.CliError(err)
	}
	if !globals.Quiet.Value {
		fmt.Printf("OK %s %s/browse/%s\n", issueData.Key, globals.Endpoint.Value, issueData.Key)
	}

	if opts.Browse.Value {
		return CmdBrowse(globals, opts.Issue)
	}
	return nil
}
<commit_msg>Add --resolution option<commit_after>package jiracmd

import (
	"fmt"
	"strings"

	"github.com/coryb/figtree"
	"github.com/coryb/oreo"

	"gopkg.in/Netflix-Skunkworks/go-jira.v1"
	"gopkg.in/Netflix-Skunkworks/go-jira.v1/jiracli"
	"gopkg.in/Netflix-Skunkworks/go-jira.v1/jiradata"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

type TransitionOptions struct {
	jiracli.CommonOptions `yaml:",inline" json:",inline" figtree:",inline"`
	Overrides             map[string]string `yaml:"overrides,omitempty" json:"overrides,omitempty"`
	Transition            string            `yaml:"transition,omitempty" json:"transition,omitempty"`
	Issue                 string            `yaml:"issue,omitempty" json:"issue,omitempty"`
	Resolution            string            `yaml:"resolution,omitempty" json:"resolution,omitempty"`
}

func CmdTransitionRegistry(transition string) *jiracli.CommandRegistryEntry {
	opts := TransitionOptions{
		CommonOptions: jiracli.CommonOptions{
			Template: figtree.NewStringOption("transition"),
		},
		Overrides: map[string]string{},
	}

	help := "Transition issue to given state"
	if transition != "" {
		help = fmt.Sprintf("Transition issue to %s state", transition)
		opts.SkipEditing = figtree.NewBoolOption(true)
	}

	return &jiracli.CommandRegistryEntry{
		help,
		func(fig *figtree.FigTree, cmd *kingpin.CmdClause) error {
			jiracli.LoadConfigs(cmd, fig, &opts)
			if opts.Transition == "" {
				opts.Transition = transition
			}
			return CmdTransitionUsage(cmd, &opts)
		},
		func(o *oreo.Client, globals *jiracli.GlobalOptions) error {
			return CmdTransition(o, globals, &opts)
		},
	}
}

func CmdTransitionUsage(cmd *kingpin.CmdClause, opts *TransitionOptions) error {
	jiracli.BrowseUsage(cmd, &opts.CommonOptions)
	jiracli.TemplateUsage(cmd, &opts.CommonOptions)
	cmd.Flag("noedit", "Disable opening the editor").SetValue(&opts.SkipEditing)
	cmd.Flag("comment", "Comment message for issue").Short('m').PreAction(func(ctx *kingpin.ParseContext) error {
		opts.Overrides["comment"] = jiracli.FlagValue(ctx, "comment")
		return nil
	}).String()
	cmd.Flag("override", "Set issue property").Short('o').StringMapVar(&opts.Overrides)
	if opts.Transition == "" {
		cmd.Arg("TRANSITION", "State to transition issue to").Required().StringVar(&opts.Transition)
	}
	cmd.Arg("ISSUE", "issue to transition").Required().StringVar(&opts.Issue)
	cmd.Flag("resolution", "Set resolution on transition").StringVar(&opts.Resolution)
	return nil
}

func defaultResolution(transMeta *jiradata.Transition) string {
	if resField, ok := transMeta.Fields["resolution"]; ok {
		for _, allowedValueRaw := range resField.AllowedValues {
			if allowedValue, ok := allowedValueRaw.(map[string]interface{}); ok {
				if allowedValue["name"] == "Fixed" {
					return "Fixed"
				} else if allowedValue["name"] == "Done" {
					return "Done"
				}
			}
		}
	}
	return ""
}

// CmdTransition will move state of the given issue to the given transtion
func CmdTransition(o *oreo.Client, globals *jiracli.GlobalOptions, opts *TransitionOptions) error {
	issueData, err := jira.GetIssue(o, globals.Endpoint.Value, opts.Issue, nil)
	if err != nil {
		return jiracli.CliError(err)
	}

	meta, err := jira.GetIssueTransitions(o, globals.Endpoint.Value, opts.Issue)
	if err != nil {
		return jiracli.CliError(err)
	}
	transMeta := meta.Transitions.Find(opts.Transition)

	if transMeta == nil {
		possible := []string{}
		for _, trans := range meta.Transitions {
			possible = append(possible, trans.Name)
		}

		if status, ok := issueData.Fields["status"].(map[string]interface{}); ok {
			if name, ok := status["name"].(string); ok {
				return jiracli.CliError(fmt.Errorf("Invalid Transition %q from %q, Available: %s", opts.Transition, name, strings.Join(possible, ", ")))
			}
		}
		return jiracli.CliError(fmt.Errorf("No valid transition found matching %s", opts.Transition))
	}

	// need to default the Resolution, usually Fixed works but sometime need Done
	if opts.Resolution == "" {
		if resField, ok := transMeta.Fields["resolution"]; ok {
			for _, allowedValueRaw := range resField.AllowedValues {
				if allowedValue, ok := allowedValueRaw.(map[string]interface{}); ok {
					if allowedValue["name"] == "Fixed" {
						opts.Resolution = "Fixed"
					} else if allowedValue["name"] == "Done" {
						opts.Resolution = "Done"
					}
				}
			}
		}
	}
	opts.Overrides["resolution"] = opts.Resolution

	type templateInput struct {
		*jiradata.Issue `yaml:",inline"`
		// Yes, Meta and Transition are redundant, but this is for backwards compatibility
		// with old templates
		Meta       *jiradata.Transition `yaml:"meta,omitempty" json:"meta,omitemtpy"`
		Transition *jiradata.Transition `yaml:"transition,omitempty" json:"transition,omitempty"`
		Overrides  map[string]string    `yaml:"overrides,omitempty" json:"overrides,omitempty"`
	}

	if _, ok := transMeta.Fields["comment"]; !ok && opts.Overrides["comment"] != "" {
		comment := jiradata.Comment{
			Body: opts.Overrides["comment"],
		}
		if _, err := jira.IssueAddComment(o, globals.Endpoint.Value, opts.Issue, &comment); err != nil {
			return err
		}
	}

	issueUpdate := jiradata.IssueUpdate{}
	input := templateInput{
		Issue:      issueData,
		Meta:       transMeta,
		Transition: transMeta,
		Overrides:  opts.Overrides,
	}
	err = jiracli.EditLoop(&opts.CommonOptions, &input, &issueUpdate, func() error {
		return jira.TransitionIssue(o, globals.Endpoint.Value, opts.Issue, &issueUpdate)
	})
	if err != nil {
		return jiracli.CliError(err)
	}
	if !globals.Quiet.Value {
		fmt.Printf("OK %s %s/browse/%s\n", issueData.Key, globals.Endpoint.Value, issueData.Key)
	}

	if opts.Browse.Value {
		return CmdBrowse(globals, opts.Issue)
	}
	return nil
}
<|endoftext|>


<commit_before>package cachemap

import (
    "time"
    "log"
    "errors"
    "sync"
    "runtime"
)

/* NOTE: to use CacheMap, another object must implement CacheMapCallbacks */

type Callbacks interface {
    // persist new key, value somewhere else
    Create (key interface{}, item *Item) error

    // get value from key somewhere else
    Read   (key interface{}) (item *Item, err error)

    // persist updated key, value somewhere else
    Update (key interface{}, item *Item) error

    // delete key somewhere else
    Delete (key interface{}) error
}

// dont create this struct or modify these fields yourself
type Item struct {
    Value interface{}

    // when this item was originally Add()ed
    Created *time.Time

    // last time this value was Get()
    LastUsed *time.Time

    // size of the key, in bytes
    KeyBytes uint64

    // size of the value, in bytes
    ValueBytes uint64
}

type CacheMap struct {
    cache                    map[interface{}]*Item
    callbacks                Callbacks
    callbackMutex            sync.Mutex
    mutex                    sync.Mutex
    CallbacksAreThreadSafe   bool
    CallbackTimeout          time.Duration
    MaxItems                 int
    _bytes                   uint64
    MaxBytes                 uint64
    evictorRunning           bool
    lazyEvictor              bool
    EvictorDelay             time.Duration
}

const NoTimeout = 0 * time.Second
const DefaultEvictorDelay = 3 * time.Second
const DefaultMaxBytes = 128*1024*1024

func NewCacheMap(c Callbacks) CacheMap {
    return CacheMap {
        cache: make(map[interface{}]*Item),
        MaxItems: 10000,
        _bytes: 0,
        EvictorDelay: DefaultEvictorDelay,
        MaxBytes: DefaultMaxBytes,
        callbacks: c,
    }
}

func NewLazyCacheMap(c Callbacks) CacheMap {
    return CacheMap {
        cache: make(map[interface{}]*Item),
        MaxItems: 10000,
        _bytes: 0,
        EvictorDelay: DefaultEvictorDelay,
        MaxBytes: DefaultMaxBytes,
        lazyEvictor: true,
        callbacks: c,
    }
}

func (m CacheMap) CachedItems() int {
    return len(m.cache)
}

func (m CacheMap) Bytes() uint64 {
    return m._bytes
}

func (m CacheMap) wouldNeedEviction(newbytes uint64, items int) bool {
    return newbytes > m.MaxBytes || items > m.MaxItems
}

func (m *CacheMap) callCreate(key interface{}, item *Item) error {
    if ! m.CallbacksAreThreadSafe {
        defer m.callbackMutex.Unlock()
        m.callbackMutex.Lock()
    }
    if m.CallbackTimeout != NoTimeout {
        return m.callbacks.Create(key, item)
    }
    c := make(chan error)
    go func() {
        c <- m.callbacks.Create(key, item)
    }()
    select {
    case err := <-c:
        return err
    case <-time.After(m.CallbackTimeout):
        return errors.New("cachemap callCreate timeout")
    }
}

func (m *CacheMap) callRead(key interface{}) (item *Item, err error) {
    if ! m.CallbacksAreThreadSafe {
        defer m.callbackMutex.Unlock()
        m.callbackMutex.Lock()
    }
    if m.CallbackTimeout != NoTimeout {
        return m.callbacks.Read(key)
    }
    type loadResult struct {
        *Item
        error
    }
    c := make(chan loadResult)
    go func() {
        i, e := m.callbacks.Read(key)
        c <- loadResult{i, e}
    }()
    select {
    case result := <-c:
        err = result.error
        item = result.Item
    case <-time.After(m.CallbackTimeout):
        err = errors.New("cachemap callLoad timeout")
    }
    return
}

func (m *CacheMap) callUpdate(key interface{}, item *Item) error {
    if ! m.CallbacksAreThreadSafe {
        defer m.callbackMutex.Unlock()
        m.callbackMutex.Lock()
    }
    if m.CallbackTimeout != NoTimeout {
        return m.callbacks.Update(key, item)
    }
    c := make(chan error)
    go func() {
        c <- m.callbacks.Update(key, item)
    }()
    select {
    case err := <-c:
        return err
    case <-time.After(m.CallbackTimeout):
        return errors.New("cachemap callUpdate timeout")
    }
}

func (m *CacheMap) callDelete(key interface{}) error {
    if ! m.CallbacksAreThreadSafe {
        defer m.callbackMutex.Unlock()
        m.callbackMutex.Lock()
    }
    if m.CallbackTimeout != NoTimeout {
        return m.callbacks.Delete(key)
    }
    c := make(chan error)
    go func() {
        c <- m.callbacks.Delete(key)
    }()
    select {
    case err := <-c:
        return err
    case <-time.After(m.CallbackTimeout):
        return errors.New("cachemap callDelete timeout")
    }
}

// replace existing item in cache
func (m *CacheMap) updateItem(key interface{}, prev, _new *Item) (err error) {
    err = m.callUpdate(key, _new)
    if err != nil {
        return
    }
    // should an immediate evict happen?
    //
    tempbytes := m._bytes + (_new.KeyBytes + _new.ValueBytes) - (prev.KeyBytes + prev.ValueBytes)
    if !m.lazyEvictor && m.wouldNeedEviction(tempbytes, m.CachedItems()+1) {
        err = m.chooseEvictionVictim()
        if err != nil {
            return
        }
    }
    // update the cache item
    prev.Value = _new.Value
    prev.KeyBytes = _new.KeyBytes
    prev.ValueBytes = _new.ValueBytes
    prev.Created = _new.Created
    m.removeBytes(prev.KeyBytes + prev.ValueBytes)
    m.addBytes(_new.KeyBytes + _new.ValueBytes)
    return
}

var noTime = time.Time{}

// add a new item to cache
func (m *CacheMap) createItem(key interface{}, _new *Item) (err error) {
    err = m.callCreate(key, _new)
    if err != nil {
        return
    }
    // should an immediate evict happen?
    tempbytes := m._bytes + (_new.KeyBytes + _new.ValueBytes)
    if !m.lazyEvictor && m.wouldNeedEviction(tempbytes, m.CachedItems()+1) {
        err = m.chooseEvictionVictim()
        if err != nil {
            return
        }
    }
    // persist the item
    m.cache[key] = _new
    m.addBytes(_new.KeyBytes + _new.ValueBytes)
    return
}

// add key, value to the map
func (m *CacheMap) Add(key, value interface{}, KeyBytes, ValueBytes uint64) error {
    now := time.Now()
    item := &Item{
        Value: value,
        KeyBytes: KeyBytes,
        ValueBytes: ValueBytes,
        Created: &now,
    }
    if prev := m.cache[key]; prev != nil {
        // replace a previous item
        return m.updateItem(key, prev, item)
    }
    // add a new item
    return m.createItem(key, item)
}

// Empty the cache, dont delete anything
func (m *CacheMap) DropCaches() {
    m.cache = make(map[interface{}]*Item)
    m._bytes = 0
}

// get an item from cache.  if not cached, fetch it
func (m *CacheMap) Get(key interface{}) (value interface{}, err error) {
    defer m.mutex.Unlock()
    m.mutex.Lock()
    cached := m.cache[key]
    // is cached?
    if cached != nil {
        now := time.Now()
        cached.LastUsed = &now
        value = cached.Value
        return
    }
    // is persisted?
    item, err := m.callRead(key)
    if err == nil {
        m.cache[key] = item
        value = item.Value
    }
    return
}

func (m *CacheMap) addBytes(bytes uint64) {
    m._bytes += bytes
}

func (m *CacheMap) removeBytes(bytes uint64) {
    m._bytes -= bytes
}

func (m *CacheMap) deleteFromCache(key interface{}) {
    item := m.cache[key]
    if item != nil {
        delete(m.cache, key)
        m.removeBytes(item.KeyBytes + item.ValueBytes)
    }
}

// remove an item from the in-memory cache
// the item was already saved
func (m *CacheMap) Evict(key interface{}) (err error) {
    cached := m.cache[key]
    if cached == nil {
        err = errors.New("cannot persist item not in cache")
        return
    }
    m.deleteFromCache(key)
    return
}

// delete an item from the cache
func (m *CacheMap) Delete(key interface{}) error {
    m.deleteFromCache(key)
    return m.callDelete(key)
}

func (m *CacheMap) chooseEvictionVictim() error {
    // 1. find OLDEST UNUSED item
    var oldestKey *interface{}
    var oldestItem *Item
    for key, item := range m.cache {
        if item.LastUsed == nil {
            continue
        }
        if oldestItem == nil || item.Created.Before(*oldestItem.Created) {
            oldestKey = &key
            oldestItem = item
        }
    }
    // was there an unused item found?
    if oldestKey != nil && oldestItem != nil {
        // dont evict items that are < 10 seconds old
        // if now - 10s < oldestUnused.created
        tenSecondsAgo := time.Now().Add(-10 * time.Second)
        if oldestItem.Created.Before(tenSecondsAgo) {
            return m.Evict(oldestKey)
        }
    }

    // 2. if no unused item(s), find the OLDEST item used
    for key, item := range m.cache {
        if oldestItem == nil || item.LastUsed.Before(*oldestItem.LastUsed) {
            oldestItem = item
            oldestKey = &key
        }
    }
    // ws there an oldest item?
    if oldestKey != nil {
        return m.Evict(oldestKey)
    }
    panic("no item to evict from cachemap")
    return nil
}

func (m CacheMap) IsLazyEvictor() bool {
    return m.lazyEvictor
}

// stop evictor
func (m *CacheMap) StopEvictor() {
    if ! m.lazyEvictor {
        panic("Not a lazy evictor")
    }
    if ! m.evictorRunning {
        panic("evictor not running")
    }
    m.evictorRunning = false
}

// start evictor
func (m *CacheMap) StartEvictor() {
    if ! m.lazyEvictor {
        panic("Not a lazy evictor")
    }
    if m.evictorRunning {
        panic("evictor already started")
    }
    m.evictorRunning = true
    go func() {
        log.Println("evictor started")
        for m.evictorRunning {
            for len(m.cache) > m.MaxItems && m._bytes > m.MaxBytes {
                if err := m.chooseEvictionVictim(); err != nil {
                    log.Println("eviction error", err)
                }
            }
            log.Println("running GC")
            runtime.GC()
            time.Sleep(m.EvictorDelay)
        }
        log.Println("evictor stopped")
    }()
}
<commit_msg>extracted cachemap<commit_after><|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.2.3`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},

	Language{"Perl", mExt(".pl", ".pm"), perlComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Ruby", mExt(".rb"), shComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments   = Commenter{"\000", "\000", "\000", false}
	xmlComments  = Commenter{"\000", `<!--`, `-->`, false}
	cComments    = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments   = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
	hsComments   = Commenter{`--`, `{-`, `-}`, true}
	sqlComments  = Commenter{`--`, `/*`, `*/`, false}
	pyComments   = Commenter{`#`, `"""`, `"""`, false}

	// TODO support POD and __END__
	perlComments = Commenter{`#`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else {
				s.CodeLines++
			}
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<commit_msg>Adding ML support<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime/pprof"
	"sort"
	"text/tabwriter"
)

const VERSION = `0.2.3`

var languages = []Language{
	Language{"Thrift", mExt(".thrift"), cComments},

	Language{"C", mExt(".c", ".h"), cComments},
	Language{"C++", mExt(".cc", ".cpp", ".cxx", ".hh", ".hpp", ".hxx"), cComments},
	Language{"Go", mExt(".go"), cComments},
	Language{"Scala", mExt(".scala"), cComments},
	Language{"Java", mExt(".java"), cComments},

	Language{"YACC", mExt(".y"), cComments},
	Language{"Lex", mExt(".l"), cComments},

	Language{"SQL", mExt(".sql"), sqlComments},

	Language{"Haskell", mExt(".hs", ".lhs"), hsComments},
	Language{"ML", mExt(".ml", ".mli"), mlComments},

	Language{"Perl", mExt(".pl", ".pm"), perlComments},
	Language{"PHP", mExt(".php"), cComments},

	Language{"Shell", mExt(".sh"), shComments},
	Language{"Bash", mExt(".bash"), shComments},

	Language{"Ruby", mExt(".rb"), shComments},
	Language{"Python", mExt(".py"), pyComments},
	Language{"Assembly", mExt(".asm", ".s"), semiComments},
	Language{"Lisp", mExt(".lsp", ".lisp"), semiComments},
	Language{"Scheme", mExt(".scm", ".scheme"), semiComments},

	Language{"Make", mName("makefile", "Makefile", "MAKEFILE"), shComments},
	Language{"CMake", mName("CMakeLists.txt"), shComments},
	Language{"Jam", mName("Jamfile", "Jamrules"), shComments},

	Language{"Markdown", mExt(".md"), noComments},

	Language{"HAML", mExt(".haml"), noComments},
	Language{"SASS", mExt(".sass"), cssComments},
	Language{"SCSS", mExt(".scss"), cssComments},

	Language{"HTML", mExt(".htm", ".html", ".xhtml"), xmlComments},
	Language{"XML", mExt(".xml"), xmlComments},
	Language{"CSS", mExt(".css"), cssComments},
	Language{"JavaScript", mExt(".js"), cComments},
}

type Commenter struct {
	LineComment  string
	StartComment string
	EndComment   string
	Nesting      bool
}

var (
	noComments   = Commenter{"\000", "\000", "\000", false}
	xmlComments  = Commenter{"\000", `<!--`, `-->`, false}
	cComments    = Commenter{`//`, `/*`, `*/`, false}
	cssComments  = Commenter{"\000", `/*`, `*/`, false}
	shComments   = Commenter{`#`, "\000", "\000", false}
	semiComments = Commenter{`;`, "\000", "\000", false}
	hsComments   = Commenter{`--`, `{-`, `-}`, true}
	mlComments   = Commenter{`\000`, `(*`, `*)`, false}
	sqlComments  = Commenter{`--`, `/*`, `*/`, false}
	pyComments   = Commenter{`#`, `"""`, `"""`, false}

	// TODO support POD and __END__
	perlComments = Commenter{`#`, "\000", "\000", false}
)

type Language struct {
	Namer
	Matcher
	Commenter
}

// TODO work properly with unicode
func (l Language) Update(c []byte, s *Stats) {
	s.FileCount++

	inComment := 0 // this is an int for nesting
	inLComment := false
	blank := true
	lc := []byte(l.LineComment)
	sc := []byte(l.StartComment)
	ec := []byte(l.EndComment)
	lp, sp, ep := 0, 0, 0

	for _, b := range c {
		if inComment == 0 && b == lc[lp] {
			lp++
			if lp == len(lc) {
				inLComment = true
				lp = 0
			}
		} else {
			lp = 0
		}
		if !inLComment && b == sc[sp] {
			sp++
			if sp == len(sc) {
				inComment++
				if inComment > 1 && !l.Nesting {
					inComment = 1
				}
				sp = 0
			}
		} else {
			sp = 0
		}
		if !inLComment && inComment > 0 && b == ec[ep] {
			ep++
			if ep == len(ec) {
				if inComment > 0 {
					inComment--
				}
				ep = 0
			}
		} else {
			ep = 0
		}

		if b != byte(' ') && b != byte('\t') && b != byte('\n') {
			blank = false
		}

		// BUG(srl): lines with comment don't count towards code
		// Note that lines with both code and comment count towards
		// each, but are not counted twice in the total.
		if b == byte('\n') {
			s.TotalLines++
			if inComment > 0 || inLComment {
				inLComment = false
				s.CommentLines++
			} else if blank {
				s.BlankLines++
			} else {
				s.CodeLines++
			}
			blank = true
			continue
		}
	}
}

type Namer string

func (l Namer) Name() string { return string(l) }

type Matcher func(string) bool

func (m Matcher) Match(fname string) bool { return m(fname) }

func mExt(exts ...string) Matcher {
	return func(fname string) bool {
		for _, ext := range exts {
			if ext == path.Ext(fname) {
				return true
			}
		}
		return false
	}
}

func mName(names ...string) Matcher {
	return func(fname string) bool {
		for _, name := range names {
			if name == path.Base(fname) {
				return true
			}
		}
		return false
	}
}

type Stats struct {
	FileCount    int
	TotalLines   int
	CodeLines    int
	BlankLines   int
	CommentLines int
}

var info = map[string]*Stats{}

func handleFileLang(fname string, l Language) {
	i, ok := info[l.Name()]
	if !ok {
		i = &Stats{}
		info[l.Name()] = i
	}
	c, err := ioutil.ReadFile(fname)
	if err != nil {
		fmt.Fprintf(os.Stderr, "  ! %s\n", fname)
		return
	}
	l.Update(c, i)
}

func handleFile(fname string) {
	for _, lang := range languages {
		if lang.Match(fname) {
			handleFileLang(fname, lang)
			return
		}
	}
	// TODO No recognized extension - check for hashbang
}

var files []string

func add(n string) {
	fi, err := os.Stat(n)
	if err != nil {
		goto invalid
	}
	if fi.IsDir() {
		fs, err := ioutil.ReadDir(n)
		if err != nil {
			goto invalid
		}
		for _, f := range fs {
			if f.Name()[0] != '.' {
				add(path.Join(n, f.Name()))
			}
		}
		return
	}
	if fi.Mode()&os.ModeType == 0 {
		files = append(files, n)
		return
	}

	println(fi.Mode())

invalid:
	fmt.Fprintf(os.Stderr, "  ! %s\n", n)
}

type LData []LResult

func (d LData) Len() int { return len(d) }

func (d LData) Less(i, j int) bool {
	if d[i].CodeLines == d[j].CodeLines {
		return d[i].Name > d[j].Name
	}
	return d[i].CodeLines > d[j].CodeLines
}

func (d LData) Swap(i, j int) {
	d[i], d[j] = d[j], d[i]
}

type LResult struct {
	Name         string
	FileCount    int
	CodeLines    int
	CommentLines int
	BlankLines   int
	TotalLines   int
}

func (r *LResult) Add(a LResult) {
	r.FileCount += a.FileCount
	r.CodeLines += a.CodeLines
	r.CommentLines += a.CommentLines
	r.BlankLines += a.BlankLines
	r.TotalLines += a.TotalLines
}

func printJSON() {
	bs, err := json.MarshalIndent(info, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(bs))
}

func printInfo() {
	w := tabwriter.NewWriter(os.Stdout, 2, 8, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(w, "Language\tFiles\tCode\tComment\tBlank\tTotal\t")
	d := LData([]LResult{})
	total := &LResult{}
	total.Name = "Total"
	for n, i := range info {
		r := LResult{n, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines}
		d = append(d, r)
		total.Add(r)
	}
	d = append(d, *total)
	sort.Sort(d)
	for _, i := range d {
		fmt.Fprintf(w, "%s\t%d\t%d\t%d\t%d\t%d\t\n", i.Name, i.FileCount, i.CodeLines, i.CommentLines, i.BlankLines, i.TotalLines)
	}

	w.Flush()
}

var (
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	useJson    = flag.Bool("json", false, "JSON-format output")
	version    = flag.Bool("V", false, "display version info and exit")
)

func main() {
	flag.Parse()
	if *version {
		fmt.Printf("sloc %s\n", VERSION)
		return
	}
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: %s\n", err.Error())
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	args := flag.Args()
	if len(args) == 0 {
		args = append(args, `.`)
	}

	for _, n := range args {
		add(n)
	}

	for _, f := range files {
		handleFile(f)
	}

	if *useJson {
		printJSON()
	} else {
		printInfo()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
)

func scrapeSearch(document *goquery.Document, url string) {
	pagesStr := document.Find("a.next_page").Prev().Text()
	pages, _ := strconv.Atoi(pagesStr)
	page := 1
	for page <= pages {
		pageURL := url + "&p=" + strconv.Itoa(page)
		doc := downloadURL(pageURL)
		doc.Find(".user-list-item").Each(func(i int, s *goquery.Selection) {
			email := s.Find("a.email").Text()
			profileURL, _ := s.Find("a").Eq(1).Attr("href")
			username := profileURL[1:len(profileURL)]
			profileURL = "http://github.com" + profileURL
			info := s.Find(".user-list-info")
			_ = info.Find("ul.user-list-meta").Remove()
			_ = info.Find("a").Remove()
			name := strings.TrimSpace(info.Text())
			fmt.Println("Parsed user: " + username)
			user := user{name: name, email: email, url: profileURL, username: username}
			dumpToCSV(user)
		})

		page = page + 1
	}
}

func scrapeProfile(document *goquery.Document) {
	vcard := document.Find(".vcard")

	email := vcard.Find("a.email").Text()
	url := "http://github.com/" + vcard.Find(".vcard-username").Text()
	username := vcard.Find(".vcard-username").Text()
	name := vcard.Find(".vcard-fullname").Text()

	fmt.Println("Parsed user: " + username)
	user := user{name: name, email: email, url: url, username: username}
	dumpToCSV(user)
}

func scrapeOrganization(document *goquery.Document, url string) {
	nextPage := document.Find("a.next_page").Length()
	if nextPage == 0 {
		document.Find(".member-list-item").Each(func(i int, s *goquery.Selection) {
			scrapeOrganizationItem(s)
		})
	} else {
		pagesStr := document.Find("a.next_page").Prev().Text()
		pages, _ := strconv.Atoi(pagesStr)
		page := 1
		for page <= pages {
			pageURL := url + "?page=" + strconv.Itoa(page)
			doc := downloadURL(pageURL)
			doc.Find(".member-list-item").Each(func(i int, s *goquery.Selection) {
				scrapeOrganizationItem(s)
			})

			page = page + 1
		}
	}
}

func scrapeOrganizationItem(element *goquery.Selection) {
	url := "http://github.com/" + element.Find(".member-username").Text()
	username := element.Find(".member-username").Text()
	name := element.Find(".member-fullname").Text()

	doc := downloadURL(url)
	email := doc.Find(".vcard").Find("a.email").Text()
	fmt.Println("Parsed user: " + username)
	user := user{name: name, email: email, url: url, username: username}
	dumpToCSV(user)
}

func scrapeStarGazers(document *goquery.Document, url string) {
	totalStarGazers := 0
	document.Find("ul.pagehead-actions li").Each(func(i int, s *goquery.Selection) {
		if i == 1 {
			count := strings.TrimSpace(s.Find("a.social-count").Text())
			totalStarGazers, _ = strconv.Atoi(count)
		}
	})
	pages := totalStarGazers / 30
	page := 1
	for page <= pages {
		pageURL := url + "?page=" + strconv.Itoa(page)
		doc := downloadURL(pageURL)
		doc.Find(".follow-list-item span.css-truncate a").Each(func(i int, s *goquery.Selection) {
			profileURL, _ := s.Attr("href")
			profileURL = "http://github.com" + profileURL
			profile := downloadURL(profileURL)
			scrapeProfile(profile)
		})

		page = page + 1
	}
}

func downloadURL(url string) *goquery.Document {
	doc, err := goquery.NewDocument(url)
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(1 * time.Second)
	return doc
}
<commit_msg>Set sleep time to 5 seconds.<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
)

func scrapeSearch(document *goquery.Document, url string) {
	pagesStr := document.Find("a.next_page").Prev().Text()
	pages, _ := strconv.Atoi(pagesStr)
	page := 1
	for page <= pages {
		pageURL := url + "&p=" + strconv.Itoa(page)
		doc := downloadURL(pageURL)
		doc.Find(".user-list-item").Each(func(i int, s *goquery.Selection) {
			email := s.Find("a.email").Text()
			profileURL, _ := s.Find("a").Eq(1).Attr("href")
			username := profileURL[1:len(profileURL)]
			profileURL = "http://github.com" + profileURL
			info := s.Find(".user-list-info")
			_ = info.Find("ul.user-list-meta").Remove()
			_ = info.Find("a").Remove()
			name := strings.TrimSpace(info.Text())
			fmt.Println("Parsed user: " + username)
			user := user{name: name, email: email, url: profileURL, username: username}
			dumpToCSV(user)
		})

		page = page + 1
	}
}

func scrapeProfile(document *goquery.Document) {
	vcard := document.Find(".vcard")

	email := vcard.Find("a.email").Text()
	url := "http://github.com/" + vcard.Find(".vcard-username").Text()
	username := vcard.Find(".vcard-username").Text()
	name := vcard.Find(".vcard-fullname").Text()

	fmt.Println("Parsed user: " + username)
	user := user{name: name, email: email, url: url, username: username}
	dumpToCSV(user)
}

func scrapeOrganization(document *goquery.Document, url string) {
	nextPage := document.Find("a.next_page").Length()
	if nextPage == 0 {
		document.Find(".member-list-item").Each(func(i int, s *goquery.Selection) {
			scrapeOrganizationItem(s)
		})
	} else {
		pagesStr := document.Find("a.next_page").Prev().Text()
		pages, _ := strconv.Atoi(pagesStr)
		page := 1
		for page <= pages {
			pageURL := url + "?page=" + strconv.Itoa(page)
			doc := downloadURL(pageURL)
			doc.Find(".member-list-item").Each(func(i int, s *goquery.Selection) {
				scrapeOrganizationItem(s)
			})

			page = page + 1
		}
	}
}

func scrapeOrganizationItem(element *goquery.Selection) {
	url := "http://github.com/" + element.Find(".member-username").Text()
	username := element.Find(".member-username").Text()
	name := element.Find(".member-fullname").Text()

	doc := downloadURL(url)
	email := doc.Find(".vcard").Find("a.email").Text()
	fmt.Println("Parsed user: " + username)
	user := user{name: name, email: email, url: url, username: username}
	dumpToCSV(user)
}

func scrapeStarGazers(document *goquery.Document, url string) {
	totalStarGazers := 0
	document.Find("ul.pagehead-actions li").Each(func(i int, s *goquery.Selection) {
		if i == 1 {
			count := strings.TrimSpace(s.Find("a.social-count").Text())
			totalStarGazers, _ = strconv.Atoi(count)
		}
	})
	pages := totalStarGazers / 30
	page := 1
	for page <= pages {
		pageURL := url + "?page=" + strconv.Itoa(page)
		doc := downloadURL(pageURL)
		doc.Find(".follow-list-item span.css-truncate a").Each(func(i int, s *goquery.Selection) {
			profileURL, _ := s.Attr("href")
			profileURL = "http://github.com" + profileURL
			profile := downloadURL(profileURL)
			scrapeProfile(profile)
		})

		page = page + 1
	}
}

func downloadURL(url string) *goquery.Document {
	doc, err := goquery.NewDocument(url)
	if err != nil {
		log.Fatal(err)
	}
	time.Sleep(5 * time.Second)
	return doc
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"net/http"

	"github.com/gorilla/mux"
)

type CallbackResponse struct {
	UUID    string `json:"uuid"`
	Message string `json:"message"`
	Error   string `json:"error"`
}

func ReceiveCallback(w http.ResponseWriter, r *http.Request) {
	var res CallbackResponse

	vars := mux.Vars(r)
	uuid := vars["uuid"]

	// Stuff will happen

	res = CallbackResponse{
		Message: "Callback received",
		UUID:    uuid,
	}

	json.NewEncoder(w).Encode(res)

}
<commit_msg>TwiML!!!<commit_after>package main

import (
	"encoding/json"
	"net/http"

	"bitbucket.org/ckvist/twilio/twiml"
	"github.com/gorilla/mux"
)

type CallbackResponse struct {
	UUID    string `json:"uuid"`
	Message string `json:"message"`
	Error   string `json:"error"`
}

func ReceiveCallback(w http.ResponseWriter, r *http.Request) {
	var res CallbackResponse

	vars := mux.Vars(r)
	uuid := vars["uuid"]

	// Stuff will happen

	res = CallbackResponse{
		Message: "Callback received",
		UUID:    uuid,
	}

	json.NewEncoder(w).Encode(res)

}

func GenerateTwiML(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	uuid := vars["uuid"]

	if _, exists := NIP.Stoppers[uuid]; !exists {
		http.Error(w, "No active notifications for this UUID", http.StatusNotFound)
		return
	}

	resp := twiml.NewResponse()

	intro := twiml.Say{
		Voice: "woman",
		Text:  "This is Chicken Little with a message for you.",
	}

	gather := twiml.Gather{
		Action:    "http://me/uuid/digits",
		Timeout:   15,
		NumDigits: 1,
	}

	theMessage := twiml.Say{
		Voice: "man",
		Text:  NIP.Messages[uuid],
	}

	pressAny := twiml.Say{
		Voice: "woman",
		Text:  "Press any key to acknowledge receipt of this message",
	}

	resp.Action(intro)
	resp.Gather(gather, theMessage, pressAny)
	resp.Send(w)
}
<|endoftext|>


<commit_before>package gubled

import (
	log "github.com/Sirupsen/logrus"
	"github.com/alexflint/go-arg"
	"github.com/caarlos0/env"

	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/metrics"
	"github.com/smancke/guble/server"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/server/rest"
	"github.com/smancke/guble/server/webserver"
	"github.com/smancke/guble/server/websocket"
	"github.com/smancke/guble/store"

	"expvar"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"path"
	"runtime"
	"syscall"
	"time"
)

var logger = log.WithFields(log.Fields{
	"app":    "guble",
	"module": "gubled",
	"env":    "TBD"})

const (
	healthEndpointPrefix = "/_health"
)

type Args struct {
	Listen      string   `arg:"-l,help: [Host:]Port the address to listen on (:8080)" env:"GUBLE_LISTEN"`
	LogInfo     bool     `arg:"--log-info,help: Log on INFO level (false)" env:"GUBLE_LOG_INFO"`
	LogDebug    bool     `arg:"--log-debug,help: Log on DEBUG level (false)" env:"GUBLE_LOG_DEBUG"`
	StoragePath string   `arg:"--storage-path,help: The path for storing messages and key value data if 'file' is enabled (/var/lib/guble)" env:"GUBLE_STORAGE_PATH"`
	KVBackend   string   `arg:"--kv-backend,help: The storage backend for the key value store to use: file|memory (file)" env:"GUBLE_KV_BACKEND"`
	MSBackend   string   `arg:"--ms-backend,help: The message storage backend : file|memory (file)" env:"GUBLE_MS_BACKEND"`
	GcmEnable   bool     `arg:"--gcm-enable: Enable the Google Cloud Messaging Connector (false)" env:"GUBLE_GCM_ENABLE"`
	GcmApiKey   string   `arg:"--gcm-api-key: The Google API Key for Google Cloud Messaging" env:"GUBLE_GCM_API_KEY"`
	GcmWorkers  int      `arg:"--gcm-workers: The number of workers handling traffic with Google Cloud Messaging (default: GOMAXPROCS)" env:"GUBLE_GCM_WORKERS"`
	Health      string   `arg:"--health: The health endpoint (default: /_health; value for disabling it: \"\" )" env:"GUBLE_HEALTH_ENDPOINT"`
	Metrics     string   `arg:"--metrics: The metrics endpoint (disabled by default; a possible value for enabling it: /_metrics )" env:"GUBLE_METRICS_ENDPOINT"`
	NodeId      string   `arg:"--node-id: The metrics endpoint (node id for guble node in cluster mode)" env:"GUBLE_NODE_ID"`
	NodeUrls    []string `arg:"positional,help: The list of urls in absolute form for other guble nodes in cluster mode"`
}

var ValidateStoragePath = func(args Args) error {
	if args.KVBackend == "file" || args.MSBackend == "file" {
		testfile := path.Join(args.StoragePath, "write-test-file")
		f, err := os.Create(testfile)
		if err != nil {
			logger.WithFields(log.Fields{
				"storagePath": args.StoragePath,
				"err":         err,
			}).Error("Storage path not present/writeable.")

			if args.StoragePath == "/var/lib/guble" {
				logger.WithFields(log.Fields{
					"storagePath": args.StoragePath,
					"err":         err,
				}).Error("Use --storage-path=<path> to override the default location, or create the directory with RW rights.")
			}
			return err
		}
		f.Close()
		os.Remove(testfile)
	}
	return nil
}

var CreateKVStore = func(args Args) store.KVStore {
	switch args.KVBackend {
	case "memory":
		return store.NewMemoryKVStore()
	case "file":
		db := store.NewSqliteKVStore(path.Join(args.StoragePath, "kv-store.db"), true)
		if err := db.Open(); err != nil {
			logger.WithField("err", err).Panic("Could not open db connection")

		}
		return db
	default:
		panic(fmt.Errorf("Unknown key-value backend: %q", args.KVBackend))
	}
}

var CreateMessageStore = func(args Args) store.MessageStore {
	switch args.MSBackend {
	case "none", "":
		return store.NewDummyMessageStore(store.NewMemoryKVStore())
	case "file":
		logger.WithField("storagePath", args.StoragePath).Info("Using FileMessageStore in directory")

		return store.NewFileMessageStore(args.StoragePath)
	default:
		panic(fmt.Errorf("Unknown message-store backend: %q", args.MSBackend))
	}
}

var CreateModules = func(router server.Router, args Args) []interface{} {
	modules := make([]interface{}, 0, 3)

	if wsHandler, err := websocket.NewWSHandler(router, "/stream/"); err != nil {
		logger.WithField("err", err).Error("Error loading WSHandler module:")
	} else {
		modules = append(modules, wsHandler)
	}

	modules = append(modules, rest.NewRestMessageAPI(router, "/api/"))

	if args.GcmEnable {
		if args.GcmApiKey == "" {
			logger.Panic("GCM API Key has to be provided, if GCM is enabled")

		}

		logger.Info("Google cloud messaging: enabled")

		logger.WithField("args.GcmWOrkers", args.GcmWorkers).Debug("Workers")

		if gcm, err := gcm.NewGCMConnector(router, "/gcm/", args.GcmApiKey, args.GcmWorkers); err != nil {
			logger.WithField("err", err).Error("Error loading GCMConnector:")
		} else {
			modules = append(modules, gcm)
		}
	} else {
		logger.Info("Google cloud messaging: disabled")
	}

	return modules
}

func Main() {
	defer func() {
		if p := recover(); p != nil {
			logger.Fatal("Fatal error in gubled after recover")
		}
	}()

	args := loadArgs()
	if args.LogInfo {
		log.SetLevel(log.InfoLevel)
	}
	if args.LogDebug {
		log.SetLevel(log.DebugLevel)
	}

	if err := ValidateStoragePath(args); err != nil {
		logger.Fatal("Fatal error in gubled in validation for storage path")
	}

	BenchmarkCluster(40, 30*time.Second, 17000)

	service := StartService(args)

	waitForTermination(func() {
		err := service.Stop()
		if err != nil {
			logger.WithField("err", err).Error("Error when stopping service")
		}
	})
}

func validateURLs(values []string) []string {
	correctURLs := make([]string, 0)
	for _, posibleUrl := range values {

		u, err := url.Parse(posibleUrl)
		if err == nil && u.IsAbs() {
			correctURLs = append(correctURLs, u.Host)
		}
	}
	return correctURLs
}

func StartService(args Args) *server.Service {
	accessManager := auth.NewAllowAllAccessManager(true)
	messageStore := CreateMessageStore(args)
	kvStore := CreateKVStore(args)

	router := server.NewRouter(accessManager, messageStore, kvStore)
	webserver := webserver.New(args.Listen)

	service := server.NewService(router, webserver).
		HealthEndpointPrefix(args.Health).
		MetricsEndpointPrefix(args.Metrics).
		GubleNodeID(args.NodeId).
		GubleNodesURLs(validateURLs(args.NodeUrls))

	service.RegisterModules(CreateModules(router, args)...)

	if err := service.Start(); err != nil {
		if err := service.Stop(); err != nil {
			logger.WithField("err", err).Error("Error when stopping service after Start() failed")
		}
		logger.WithField("err", err).Fatal("Service could not be started")
	}
	expvar.Publish("guble.args", expvar.Func(func() interface{} {
		return args
	}))

	return service
}

func loadArgs() Args {
	args := Args{
		Listen:      ":8080",
		KVBackend:   "file",
		MSBackend:   "file",
		StoragePath: "/var/lib/guble",
		GcmWorkers:  runtime.GOMAXPROCS(0),
		Health:      healthEndpointPrefix,
	}

	env.Parse(&args)
	arg.MustParse(&args)
	return args
}

func waitForTermination(callback func()) {
	signalC := make(chan os.Signal)
	signal.Notify(signalC, syscall.SIGINT, syscall.SIGTERM)

	logger.Infof("Got signal '%v' .. exiting gracefully now", <-signalC)

	callback()
	metrics.LogOnDebugLevel()
	logger.Info("Exit gracefully now")
	os.Exit(0)
}
<commit_msg>less nodes in benchmark<commit_after>package gubled

import (
	log "github.com/Sirupsen/logrus"
	"github.com/alexflint/go-arg"
	"github.com/caarlos0/env"

	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/metrics"
	"github.com/smancke/guble/server"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/server/rest"
	"github.com/smancke/guble/server/webserver"
	"github.com/smancke/guble/server/websocket"
	"github.com/smancke/guble/store"

	"expvar"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"path"
	"runtime"
	"syscall"
	"time"
)

var logger = log.WithFields(log.Fields{
	"app":    "guble",
	"module": "gubled",
	"env":    "TBD"})

const (
	healthEndpointPrefix = "/_health"
)

type Args struct {
	Listen      string   `arg:"-l,help: [Host:]Port the address to listen on (:8080)" env:"GUBLE_LISTEN"`
	LogInfo     bool     `arg:"--log-info,help: Log on INFO level (false)" env:"GUBLE_LOG_INFO"`
	LogDebug    bool     `arg:"--log-debug,help: Log on DEBUG level (false)" env:"GUBLE_LOG_DEBUG"`
	StoragePath string   `arg:"--storage-path,help: The path for storing messages and key value data if 'file' is enabled (/var/lib/guble)" env:"GUBLE_STORAGE_PATH"`
	KVBackend   string   `arg:"--kv-backend,help: The storage backend for the key value store to use: file|memory (file)" env:"GUBLE_KV_BACKEND"`
	MSBackend   string   `arg:"--ms-backend,help: The message storage backend : file|memory (file)" env:"GUBLE_MS_BACKEND"`
	GcmEnable   bool     `arg:"--gcm-enable: Enable the Google Cloud Messaging Connector (false)" env:"GUBLE_GCM_ENABLE"`
	GcmApiKey   string   `arg:"--gcm-api-key: The Google API Key for Google Cloud Messaging" env:"GUBLE_GCM_API_KEY"`
	GcmWorkers  int      `arg:"--gcm-workers: The number of workers handling traffic with Google Cloud Messaging (default: GOMAXPROCS)" env:"GUBLE_GCM_WORKERS"`
	Health      string   `arg:"--health: The health endpoint (default: /_health; value for disabling it: \"\" )" env:"GUBLE_HEALTH_ENDPOINT"`
	Metrics     string   `arg:"--metrics: The metrics endpoint (disabled by default; a possible value for enabling it: /_metrics )" env:"GUBLE_METRICS_ENDPOINT"`
	NodeId      string   `arg:"--node-id: The metrics endpoint (node id for guble node in cluster mode)" env:"GUBLE_NODE_ID"`
	NodeUrls    []string `arg:"positional,help: The list of urls in absolute form for other guble nodes in cluster mode"`
}

var ValidateStoragePath = func(args Args) error {
	if args.KVBackend == "file" || args.MSBackend == "file" {
		testfile := path.Join(args.StoragePath, "write-test-file")
		f, err := os.Create(testfile)
		if err != nil {
			logger.WithFields(log.Fields{
				"storagePath": args.StoragePath,
				"err":         err,
			}).Error("Storage path not present/writeable.")

			if args.StoragePath == "/var/lib/guble" {
				logger.WithFields(log.Fields{
					"storagePath": args.StoragePath,
					"err":         err,
				}).Error("Use --storage-path=<path> to override the default location, or create the directory with RW rights.")
			}
			return err
		}
		f.Close()
		os.Remove(testfile)
	}
	return nil
}

var CreateKVStore = func(args Args) store.KVStore {
	switch args.KVBackend {
	case "memory":
		return store.NewMemoryKVStore()
	case "file":
		db := store.NewSqliteKVStore(path.Join(args.StoragePath, "kv-store.db"), true)
		if err := db.Open(); err != nil {
			logger.WithField("err", err).Panic("Could not open db connection")

		}
		return db
	default:
		panic(fmt.Errorf("Unknown key-value backend: %q", args.KVBackend))
	}
}

var CreateMessageStore = func(args Args) store.MessageStore {
	switch args.MSBackend {
	case "none", "":
		return store.NewDummyMessageStore(store.NewMemoryKVStore())
	case "file":
		logger.WithField("storagePath", args.StoragePath).Info("Using FileMessageStore in directory")

		return store.NewFileMessageStore(args.StoragePath)
	default:
		panic(fmt.Errorf("Unknown message-store backend: %q", args.MSBackend))
	}
}

var CreateModules = func(router server.Router, args Args) []interface{} {
	modules := make([]interface{}, 0, 3)

	if wsHandler, err := websocket.NewWSHandler(router, "/stream/"); err != nil {
		logger.WithField("err", err).Error("Error loading WSHandler module:")
	} else {
		modules = append(modules, wsHandler)
	}

	modules = append(modules, rest.NewRestMessageAPI(router, "/api/"))

	if args.GcmEnable {
		if args.GcmApiKey == "" {
			logger.Panic("GCM API Key has to be provided, if GCM is enabled")

		}

		logger.Info("Google cloud messaging: enabled")

		logger.WithField("args.GcmWOrkers", args.GcmWorkers).Debug("Workers")

		if gcm, err := gcm.NewGCMConnector(router, "/gcm/", args.GcmApiKey, args.GcmWorkers); err != nil {
			logger.WithField("err", err).Error("Error loading GCMConnector:")
		} else {
			modules = append(modules, gcm)
		}
	} else {
		logger.Info("Google cloud messaging: disabled")
	}

	return modules
}

func Main() {
	defer func() {
		if p := recover(); p != nil {
			logger.Fatal("Fatal error in gubled after recover")
		}
	}()

	args := loadArgs()
	if args.LogInfo {
		log.SetLevel(log.InfoLevel)
	}
	if args.LogDebug {
		log.SetLevel(log.DebugLevel)
	}

	if err := ValidateStoragePath(args); err != nil {
		logger.Fatal("Fatal error in gubled in validation for storage path")
	}

	BenchmarkCluster(3, 10*time.Second, 10000)

	service := StartService(args)

	waitForTermination(func() {
		err := service.Stop()
		if err != nil {
			logger.WithField("err", err).Error("Error when stopping service")
		}
	})
}

func validateURLs(values []string) []string {
	correctURLs := make([]string, 0)
	for _, posibleUrl := range values {

		u, err := url.Parse(posibleUrl)
		if err == nil && u.IsAbs() {
			correctURLs = append(correctURLs, u.Host)
		}
	}
	return correctURLs
}

func StartService(args Args) *server.Service {
	accessManager := auth.NewAllowAllAccessManager(true)
	messageStore := CreateMessageStore(args)
	kvStore := CreateKVStore(args)

	router := server.NewRouter(accessManager, messageStore, kvStore)
	webserver := webserver.New(args.Listen)

	service := server.NewService(router, webserver).
		HealthEndpointPrefix(args.Health).
		MetricsEndpointPrefix(args.Metrics).
		GubleNodeID(args.NodeId).
		GubleNodesURLs(validateURLs(args.NodeUrls))

	service.RegisterModules(CreateModules(router, args)...)

	if err := service.Start(); err != nil {
		if err := service.Stop(); err != nil {
			logger.WithField("err", err).Error("Error when stopping service after Start() failed")
		}
		logger.WithField("err", err).Fatal("Service could not be started")
	}
	expvar.Publish("guble.args", expvar.Func(func() interface{} {
		return args
	}))

	return service
}

func loadArgs() Args {
	args := Args{
		Listen:      ":8080",
		KVBackend:   "file",
		MSBackend:   "file",
		StoragePath: "/var/lib/guble",
		GcmWorkers:  runtime.GOMAXPROCS(0),
		Health:      healthEndpointPrefix,
	}

	env.Parse(&args)
	arg.MustParse(&args)
	return args
}

func waitForTermination(callback func()) {
	signalC := make(chan os.Signal)
	signal.Notify(signalC, syscall.SIGINT, syscall.SIGTERM)

	logger.Infof("Got signal '%v' .. exiting gracefully now", <-signalC)

	callback()
	metrics.LogOnDebugLevel()
	logger.Info("Exit gracefully now")
	os.Exit(0)
}
<|endoftext|>


<commit_before>package gubled

import (
	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server"
	"github.com/smancke/guble/store"

	"fmt"
	"github.com/alexflint/go-arg"
	"github.com/caarlos0/env"
	"github.com/smancke/guble/server/auth"
	"os"
	"os/signal"
	"path"
	"syscall"
)

type Args struct {
	Listen      string `arg:"-l,help: [Host:]Port the address to listen on (:8080)" env:"GUBLE_LISTEN"`
	LogInfo     bool   `arg:"--log-info,help: Log on INFO level (false)" env:"GUBLE_LOG_INFO"`
	LogDebug    bool   `arg:"--log-debug,help: Log on DEBUG level (false)" env:"GUBLE_LOG_DEBUG"`
	StoragePath string `arg:"--storage-path,help: The path for storing messages and key value data if 'file' is enabled (/var/lib/guble)" env:"GUBLE_STORAGE_PATH"`
	KVBackend   string `arg:"--kv-backend,help: The storage backend for the key value store to use: file|memory (file)" env:"GUBLE_KV_BACKEND"`
	MSBackend   string `arg:"--ms-backend,help: The message storage backend : file|memory (file)" env:"GUBLE_MS_BACKEND"`
	GcmEnable   bool   `arg:"--gcm-enable: Enable the Google Cloud Messaging Connector (false)" env:"GUBLE_GCM_ENABLE"`
	GcmApiKey   string `arg:"--gcm-api-key: The Google API Key for Google Cloud Messaging" env:"GUBLE_GCM_API_KEY"`
}

var ValidateStoragePath = func(args Args) error {
	if args.KVBackend == "file" || args.MSBackend == "file" {
		testfile := path.Join(args.StoragePath, "write-test-file")
		f, err := os.Create(testfile)
		if err != nil {
			protocol.ErrWithoutTrace("Storage path not present/writeable %q: %v", args.StoragePath, err)
			if args.StoragePath == "/var/lib/guble" {
				protocol.ErrWithoutTrace("Use --storage-path=<path> to override the default location, or create the directory with RW rights.")
			}
			return err
		}
		f.Close()
		os.Remove(testfile)
	}
	return nil
}

var CreateKVStore = func(args Args) store.KVStore {
	switch args.KVBackend {
	case "memory":
		return store.NewMemoryKVStore()
	case "file":
		db := store.NewSqliteKVStore(path.Join(args.StoragePath, "kv-store.db"), true)
		if err := db.Open(); err != nil {
			panic(err)
		}
		return db
	default:
		panic(fmt.Errorf("unknown key value backend: %q", args.KVBackend))
	}
}

var CreateMessageStore = func(args Args) store.MessageStore {
	switch args.MSBackend {
	case "none", "":
		return store.NewDummyMessageStore()
	case "file":
		protocol.Info("using FileMessageStore in directory: %q", args.StoragePath)
		return store.NewFileMessageStore(args.StoragePath)
	default:
		panic(fmt.Errorf("unknown message store backend: %q", args.MSBackend))
	}
}

var CreateModules = func(
	router server.Router,
	args Args) []interface{} {
	modules := make([]interface{}, 0, 2)

	if wsHandler, err := server.NewWSHandler(router, "/stream/"); err != nil {
		protocol.Err("Error loading WSHandler module: %s", err)
	} else {
		modules = append(modules, wsHandler)
	}

	modules = append(modules, server.NewRestMessageApi(router, "/api/"))

	if args.GcmEnable {
		if args.GcmApiKey == "" {
			panic("gcm api key has to be provided, if gcm is enabled")
		}

		protocol.Info("google cloud messaging: enabled")
		if gcm, err := gcm.NewGCMConnector(router, "/gcm/", args.GcmApiKey); err != nil {
			protocol.Err("Error loading GCMConnector: ", err)
		} else {
			modules = append(modules, gcm)
		}
	} else {
		protocol.Info("google cloud messaging: disabled")
	}

	return modules
}

func Main() {
	defer func() {
		if p := recover(); p != nil {
			protocol.Err("%v", p)
			os.Exit(1)
		}
	}()

	args := loadArgs()
	if args.LogInfo {
		protocol.LogLevel = protocol.LEVEL_INFO
	}
	if args.LogDebug {
		protocol.LogLevel = protocol.LEVEL_DEBUG
	}

	if err := ValidateStoragePath(args); err != nil {
		os.Exit(1)
	}

	service := StartupService(args)

	waitForTermination(func() {
		err := service.Stop()
		if err != nil {
			protocol.Err("Service: ", err)
		}
	})
}

func StartupService(args Args) *server.Service {
	accessManager := auth.NewAllowAllAccessManager(true)
	messageStore := CreateMessageStore(args)
	kvStore := CreateKVStore(args)

	router := server.NewRouter(accessManager, messageStore, kvStore)

	service := server.NewService(args.Listen, router)

	for _, module := range CreateModules(router, args) {
		service.Register(module)
	}

	if err := service.Start(); err != nil {
		protocol.Err(err.Error())
		if err := service.Stop(); err != nil {
			protocol.Err(err.Error())
		}
		os.Exit(1)
	}

	return service
}

func loadArgs() Args {
	args := Args{
		Listen:      ":8080",
		KVBackend:   "file",
		MSBackend:   "file",
		StoragePath: "/var/lib/guble",
	}

	env.Parse(&args)
	arg.MustParse(&args)
	return args
}

func waitForTermination(callback func()) {
	sigc := make(chan os.Signal)
	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
	protocol.Info("Got signal '%v' .. exiting gracefully now", <-sigc)
	callback()
	protocol.Info("exit now")
	os.Exit(0)
}
<commit_msg>DummyMessageStore requires a kvStore in constructor<commit_after>package gubled

import (
	"github.com/smancke/guble/gcm"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server"
	"github.com/smancke/guble/store"

	"fmt"
	"github.com/alexflint/go-arg"
	"github.com/caarlos0/env"
	"github.com/smancke/guble/server/auth"
	"os"
	"os/signal"
	"path"
	"syscall"
)

type Args struct {
	Listen      string `arg:"-l,help: [Host:]Port the address to listen on (:8080)" env:"GUBLE_LISTEN"`
	LogInfo     bool   `arg:"--log-info,help: Log on INFO level (false)" env:"GUBLE_LOG_INFO"`
	LogDebug    bool   `arg:"--log-debug,help: Log on DEBUG level (false)" env:"GUBLE_LOG_DEBUG"`
	StoragePath string `arg:"--storage-path,help: The path for storing messages and key value data if 'file' is enabled (/var/lib/guble)" env:"GUBLE_STORAGE_PATH"`
	KVBackend   string `arg:"--kv-backend,help: The storage backend for the key value store to use: file|memory (file)" env:"GUBLE_KV_BACKEND"`
	MSBackend   string `arg:"--ms-backend,help: The message storage backend : file|memory (file)" env:"GUBLE_MS_BACKEND"`
	GcmEnable   bool   `arg:"--gcm-enable: Enable the Google Cloud Messaging Connector (false)" env:"GUBLE_GCM_ENABLE"`
	GcmApiKey   string `arg:"--gcm-api-key: The Google API Key for Google Cloud Messaging" env:"GUBLE_GCM_API_KEY"`
}

var ValidateStoragePath = func(args Args) error {
	if args.KVBackend == "file" || args.MSBackend == "file" {
		testfile := path.Join(args.StoragePath, "write-test-file")
		f, err := os.Create(testfile)
		if err != nil {
			protocol.ErrWithoutTrace("Storage path not present/writeable %q: %v", args.StoragePath, err)
			if args.StoragePath == "/var/lib/guble" {
				protocol.ErrWithoutTrace("Use --storage-path=<path> to override the default location, or create the directory with RW rights.")
			}
			return err
		}
		f.Close()
		os.Remove(testfile)
	}
	return nil
}

var CreateKVStore = func(args Args) store.KVStore {
	switch args.KVBackend {
	case "memory":
		return store.NewMemoryKVStore()
	case "file":
		db := store.NewSqliteKVStore(path.Join(args.StoragePath, "kv-store.db"), true)
		if err := db.Open(); err != nil {
			panic(err)
		}
		return db
	default:
		panic(fmt.Errorf("unknown key value backend: %q", args.KVBackend))
	}
}

var CreateMessageStore = func(args Args) store.MessageStore {
	switch args.MSBackend {
	case "none", "":
		return store.NewDummyMessageStore(store.NewMemoryKVStore())
	case "file":
		protocol.Info("using FileMessageStore in directory: %q", args.StoragePath)
		return store.NewFileMessageStore(args.StoragePath)
	default:
		panic(fmt.Errorf("unknown message store backend: %q", args.MSBackend))
	}
}

var CreateModules = func(
	router server.Router,
	args Args) []interface{} {
	modules := make([]interface{}, 0, 2)

	if wsHandler, err := server.NewWSHandler(router, "/stream/"); err != nil {
		protocol.Err("Error loading WSHandler module: %s", err)
	} else {
		modules = append(modules, wsHandler)
	}

	modules = append(modules, server.NewRestMessageApi(router, "/api/"))

	if args.GcmEnable {
		if args.GcmApiKey == "" {
			panic("gcm api key has to be provided, if gcm is enabled")
		}

		protocol.Info("google cloud messaging: enabled")
		if gcm, err := gcm.NewGCMConnector(router, "/gcm/", args.GcmApiKey); err != nil {
			protocol.Err("Error loading GCMConnector: ", err)
		} else {
			modules = append(modules, gcm)
		}
	} else {
		protocol.Info("google cloud messaging: disabled")
	}

	return modules
}

func Main() {
	defer func() {
		if p := recover(); p != nil {
			protocol.Err("%v", p)
			os.Exit(1)
		}
	}()

	args := loadArgs()
	if args.LogInfo {
		protocol.LogLevel = protocol.LEVEL_INFO
	}
	if args.LogDebug {
		protocol.LogLevel = protocol.LEVEL_DEBUG
	}

	if err := ValidateStoragePath(args); err != nil {
		os.Exit(1)
	}

	service := StartupService(args)

	waitForTermination(func() {
		err := service.Stop()
		if err != nil {
			protocol.Err("Service: ", err)
		}
	})
}

func StartupService(args Args) *server.Service {
	accessManager := auth.NewAllowAllAccessManager(true)
	messageStore := CreateMessageStore(args)
	kvStore := CreateKVStore(args)

	router := server.NewRouter(accessManager, messageStore, kvStore)

	service := server.NewService(args.Listen, router)

	for _, module := range CreateModules(router, args) {
		service.Register(module)
	}

	if err := service.Start(); err != nil {
		protocol.Err(err.Error())
		if err := service.Stop(); err != nil {
			protocol.Err(err.Error())
		}
		os.Exit(1)
	}

	return service
}

func loadArgs() Args {
	args := Args{
		Listen:      ":8080",
		KVBackend:   "file",
		MSBackend:   "file",
		StoragePath: "/var/lib/guble",
	}

	env.Parse(&args)
	arg.MustParse(&args)
	return args
}

func waitForTermination(callback func()) {
	sigc := make(chan os.Signal)
	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
	protocol.Info("Got signal '%v' .. exiting gracefully now", <-sigc)
	callback()
	protocol.Info("exit now")
	os.Exit(0)
}
<|endoftext|>