<commit_before>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// BayesianAverageC Constant for the Bayesian average for the RTT
	BayesianAverageC = 10
	// MaxFailures Maximum number of unanswered queries before a server is marked as dead for VacuumPeriod
	MaxFailures = 100
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr     string
	failures uint
	offline  bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

// UpstreamRTT Keep track of the mean RTT
type UpstreamRTT struct {
	lock  sync.Mutex
	RTT   float64
	count float64
}

// QueuedResponse Response to an asynchronous query
type QueuedResponse struct {
	resolved *dns.Msg
	rtt      time.Duration
	err      error
}

// QueuedRequest Asynchronous DNS request
type QueuedRequest struct {
	ts           time.Time
	req          *dns.Msg
	responseChan chan QueuedResponse
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53,8.8.4.4:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
	maxClients         = flag.Uint("maxclients", 10000, "Maximum number of simultaneous clients")
	maxRTT             = flag.Float64("maxrtt", 0.25, "Maximum mean RTT for upstream queries before marking a server as dead")
	upstreamRtt        UpstreamRTT
	resolverRing       chan QueuedRequest
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	resolverRing = make(chan QueuedRequest, *maxClients)
	for i := uint(0); i < *maxClients; i++ {
		go func() {
			resolverThread()
		}()
	}
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	fmt.Println("RPDNS")
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name),
		Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := uint64(len(upstreamServers.live))
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h / (math.MaxUint64 / liveCount)
	if i >= liveCount {
		i = liveCount - 1
	}
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markFailed(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	for i, server := range upstreamServers.servers {
		if server.addr != addr {
			continue
		}
		if server.offline {
			return
		}
		upstreamServers.servers[i].failures++
		if upstreamServers.servers[i].failures < MaxFailures {
			return
		}
		break
	}
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	log.Printf("[%v] is unresponsive", addr)
}

func resetRTT() {
	upstreamRtt.lock.Lock()
	defer upstreamRtt.lock.Unlock()
	upstreamRtt.count = 0.0
	upstreamRtt.RTT = 0.0
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].failures = 0
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	resetRTT()
}

func syncResolve(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, 0, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, rtt, err := client.Exchange(req, *addr)
	if resolved != nil && resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, rtt, err = client.Exchange(req, *addr)
	}
	if err != nil {
		markFailed(*addr)
		return nil, 0, err
	}
	upstreamRtt.lock.Lock()
	upstreamRtt.count++
	upstreamRtt.RTT += rtt.Seconds()
	meanRTT := upstreamRtt.RTT / (upstreamRtt.count + BayesianAverageC)
	upstreamRtt.lock.Unlock()
	if meanRTT > *maxRTT {
		markFailed(*addr)
	}
	return resolved, rtt, nil
}

func resolverThread() {
	for {
		queuedRequest := <-resolverRing
		if time.Since(queuedRequest.ts).Seconds() > *maxRTT {
			response := QueuedResponse{resolved: nil, rtt: 0, err: errors.New("Request too old")}
			queuedRequest.responseChan <- response
			fmt.Println("Request too old")
			continue
		}
		resolved, rtt, err := syncResolve(queuedRequest.req)
		response := QueuedResponse{resolved: resolved, rtt: rtt, err: err}
		queuedRequest.responseChan <- response
	}
}

func resolveViaResolverThreads(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	responseChan := make(chan QueuedResponse)
	queuedRequest := QueuedRequest{ts: time.Now(), req: req, responseChan: responseChan}
	for queued := false; queued == false; {
		select {
		case resolverRing <- queuedRequest:
			queued = true
		default:
			select {
			case old := <-resolverRing:
				evictedResponse := QueuedResponse{resolved: nil, rtt: 0, err: errors.New("Evicted")}
				old.responseChan <- evictedResponse
			}
		}
	}
	response := <-responseChan
	if response.err != nil {
		return nil, response.rtt, errors.New("Stolen")
	}
	return response.resolved, response.rtt, nil
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	resolved, _, err := resolveViaResolverThreads(req)
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func handleSpecialNames(w dns.ResponseWriter, req *dns.Msg) bool {
	question := req.Question[0]
	if question.Qtype != dns.TypeANY {
		return false
	}
	m := new(dns.Msg)
	m.Id = req.Id
	hinfo := new(dns.HINFO)
	hinfo.Hdr = dns.RR_Header{Name: question.Name, Rrtype: dns.TypeHINFO,
		Class: dns.ClassINET, Ttl: 86400}
	hinfo.Cpu = "ANY is not supported any more"
	hinfo.Os = "See draft-jabley-dnsop-refuse-any"
	m.Answer = []dns.RR{hinfo}
	w.WriteMsg(m)
	return true
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	if handleSpecialNames(w, req) {
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<commit_msg>Reset the list of live servers if none seem to be live<commit_after>package main

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/dchest/siphash"
	"github.com/hashicorp/golang-lru"
	"github.com/miekg/dns"
)

const (
	// BayesianAverageC Constant for the Bayesian average for the RTT
	BayesianAverageC = 10
	// MaxFailures Maximum number of unanswered queries before a server is marked as dead for VacuumPeriod
	MaxFailures = 100
	// MinTTL Minimum TTL
	MinTTL = 60
	// MaxTTL Maximum TTL
	MaxTTL = 604800
	// VacuumPeriod Vacuum period in seconds
	VacuumPeriod = 30
)

// SipHashKey SipHash secret key
type SipHashKey struct {
	k1 uint64
	k2 uint64
}

// UpstreamServer Upstream server
type UpstreamServer struct {
	addr     string
	failures uint
	offline  bool
}

// UpstreamServers List of upstream servers
type UpstreamServers struct {
	lock    sync.RWMutex
	servers []UpstreamServer
	live    []string
}

// UpstreamRTT Keep track of the mean RTT
type UpstreamRTT struct {
	lock  sync.Mutex
	RTT   float64
	count float64
}

// QueuedResponse Response to an asynchronous query
type QueuedResponse struct {
	resolved *dns.Msg
	rtt      time.Duration
	err      error
}

// QueuedRequest Asynchronous DNS request
type QueuedRequest struct {
	ts           time.Time
	req          *dns.Msg
	responseChan chan QueuedResponse
}

var (
	address            = flag.String("listen", ":53", "Address to listen to (TCP and UDP)")
	upstreamServersStr = flag.String("upstream", "8.8.8.8:53,8.8.4.4:53", "Comma-delimited list of upstream servers")
	upstreamServers    *UpstreamServers
	cacheSize          = flag.Int("cachesize", 10000000, "Cache size")
	memSize            = flag.Uint64("memsize", 1*1024, "Memory size in MB")
	minLabelsCount     = flag.Int("minlabels", 2, "Minimum number of labels")
	cache              *lru.ARCCache
	sipHashKey         = SipHashKey{k1: 0, k2: 0}
	maxClients         = flag.Uint("maxclients", 10000, "Maximum number of simultaneous clients")
	maxRTT             = flag.Float64("maxrtt", 0.25, "Maximum mean RTT for upstream queries before marking a server as dead")
	upstreamRtt        UpstreamRTT
	resolverRing       chan QueuedRequest
)

func parseUpstreamServers(str string) (*UpstreamServers, error) {
	servers := []UpstreamServer{}
	live := []string{}
	for _, addr := range strings.Split(str, ",") {
		server := UpstreamServer{addr: addr}
		servers = append(servers, server)
		live = append(live, addr)
	}
	res := UpstreamServers{servers: servers, live: live}
	return &res, nil
}

func randUint64() uint64 {
	buf := make([]byte, 8)
	length, err := rand.Read(buf)
	if err != nil || length != len(buf) {
		log.Fatal("RNG failure")
	}
	return binary.LittleEndian.Uint64(buf)
}

func main() {
	flag.Parse()
	*memSize *= 1024 * 1024
	if *cacheSize < 2 {
		log.Fatal("Cache size too small")
	}
	cache, _ = lru.NewARC(*cacheSize)
	upstreamServers, _ = parseUpstreamServers(*upstreamServersStr)
	sipHashKey = SipHashKey{k1: randUint64(), k2: randUint64()}
	resolverRing = make(chan QueuedRequest, *maxClients)
	for i := uint(0); i < *maxClients; i++ {
		go func() {
			resolverThread()
		}()
	}
	dns.HandleFunc(".", route)
	udpServer := &dns.Server{Addr: *address, Net: "udp"}
	tcpServer := &dns.Server{Addr: *address, Net: "tcp"}
	go func() {
		log.Fatal(udpServer.ListenAndServe())
	}()
	go func() {
		log.Fatal(tcpServer.ListenAndServe())
	}()
	fmt.Println("RPDNS")
	vacuumThread()
}

// CacheKey Key for a cache entry
type CacheKey struct {
	Name   string `dns:"cdomain-name"`
	Qtype  uint16
	DNSSEC bool
}

// CacheVal Value for a cache entry
type CacheVal struct {
	ValidUntil time.Time
	Response   *dns.Msg
}

func getKey(req *dns.Msg) (*CacheKey, error) {
	questions := req.Question
	if len(questions) != 1 {
		return nil, errors.New("Invalid number of questions")
	}
	question := questions[0]
	if question.Qclass != dns.ClassINET {
		return nil, errors.New("Unsupported question class")
	}
	if dns.CountLabel(question.Name) < *minLabelsCount {
		return nil, errors.New("Not enough labels")
	}
	dnssec := false
	for _, extra := range req.Extra {
		if extra.Header().Rrtype == dns.TypeOPT {
			dnssec = extra.(*dns.OPT).Do()
		}
	}
	CacheKey := CacheKey{Name: strings.ToLower(question.Name),
		Qtype: question.Qtype, DNSSEC: dnssec}
	return &CacheKey, nil
}

func getMaxPayloadSize(req *dns.Msg) uint16 {
	opt := req.IsEdns0()
	if opt == nil {
		return dns.MinMsgSize
	}
	maxPayloadSize := opt.UDPSize()
	if maxPayloadSize < dns.MinMsgSize {
		maxPayloadSize = dns.MinMsgSize
	}
	return maxPayloadSize
}

func pickUpstream(req *dns.Msg) (*string, error) {
	name := strings.ToLower(req.Question[0].Name)
	h := siphash.Hash(sipHashKey.k1, sipHashKey.k2, []byte(name))
	upstreamServers.lock.RLock()
	liveCount := uint64(len(upstreamServers.live))
	if liveCount <= 0 {
		upstreamServers.lock.RUnlock()
		return nil, errors.New("All upstream servers are down")
	}
	i := h / (math.MaxUint64 / liveCount)
	if i >= liveCount {
		i = liveCount - 1
	}
	res := upstreamServers.live[i]
	upstreamServers.lock.RUnlock()
	return &res, nil
}

func markFailed(addr string) {
	upstreamServers.lock.Lock()
	defer upstreamServers.lock.Unlock()
	for i, server := range upstreamServers.servers {
		if server.addr != addr {
			continue
		}
		if server.offline {
			return
		}
		upstreamServers.servers[i].failures++
		if upstreamServers.servers[i].failures < MaxFailures {
			return
		}
		break
	}
	if len(upstreamServers.live) <= 1 {
		resetUpstreamServersNoLock()
		return
	}
	servers := upstreamServers.servers
	live := []string{}
	for i, server := range upstreamServers.servers {
		if server.addr == addr {
			servers[i].offline = true
		} else if server.offline == false {
			live = append(live, server.addr)
		}
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
}

func resetRTT() {
	upstreamRtt.lock.Lock()
	defer upstreamRtt.lock.Unlock()
	upstreamRtt.count = 0.0
	upstreamRtt.RTT = 0.0
}

func resetUpstreamServersNoLock() {
	servers := upstreamServers.servers
	if len(servers) == len(upstreamServers.live) {
		return
	}
	live := []string{}
	for i, server := range upstreamServers.servers {
		servers[i].failures = 0
		servers[i].offline = false
		live = append(live, server.addr)
	}
	upstreamServers.servers = servers
	upstreamServers.live = live
	resetRTT()
}

func resetUpstreamServers() {
	upstreamServers.lock.Lock()
	resetUpstreamServersNoLock()
	upstreamServers.lock.Unlock()
}

func syncResolve(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	addr, err := pickUpstream(req)
	if err != nil {
		return nil, 0, err
	}
	client := &dns.Client{Net: "udp"}
	client.SingleInflight = true
	resolved, rtt, err := client.Exchange(req, *addr)
	if resolved != nil && resolved.Truncated {
		client = &dns.Client{Net: "tcp"}
		client.SingleInflight = true
		resolved, rtt, err = client.Exchange(req, *addr)
	}
	if err != nil {
		markFailed(*addr)
		return nil, 0, err
	}
	upstreamRtt.lock.Lock()
	upstreamRtt.count++
	upstreamRtt.RTT += rtt.Seconds()
	meanRTT := upstreamRtt.RTT / (upstreamRtt.count + BayesianAverageC)
	upstreamRtt.lock.Unlock()
	if meanRTT > *maxRTT {
		markFailed(*addr)
	}
	return resolved, rtt, nil
}

func resolverThread() {
	for {
		queuedRequest := <-resolverRing
		if time.Since(queuedRequest.ts).Seconds() > *maxRTT {
			response := QueuedResponse{resolved: nil, rtt: 0, err: errors.New("Request too old")}
			queuedRequest.responseChan <- response
			fmt.Println("Request too old")
			continue
		}
		resolved, rtt, err := syncResolve(queuedRequest.req)
		response := QueuedResponse{resolved: resolved, rtt: rtt, err: err}
		queuedRequest.responseChan <- response
	}
}

func resolveViaResolverThreads(req *dns.Msg) (*dns.Msg, time.Duration, error) {
	responseChan := make(chan QueuedResponse)
	queuedRequest := QueuedRequest{ts: time.Now(), req: req, responseChan: responseChan}
	for queued := false; queued == false; {
		select {
		case resolverRing <- queuedRequest:
			queued = true
		default:
			select {
			case old := <-resolverRing:
				evictedResponse := QueuedResponse{resolved: nil, rtt: 0, err: errors.New("Evicted")}
				old.responseChan <- evictedResponse
			}
		}
	}
	response := <-responseChan
	if response.err != nil {
		return nil, response.rtt, errors.New("Stolen")
	}
	return response.resolved, response.rtt, nil
}

func resolve(req *dns.Msg, dnssec bool) (*dns.Msg, error) {
	extra2 := []dns.RR{}
	for _, extra := range req.Extra {
		if extra.Header().Rrtype != dns.TypeOPT {
			extra2 = append(extra2, extra)
		}
	}
	req.Extra = extra2
	req.SetEdns0(dns.MaxMsgSize, dnssec)
	resolved, _, err := resolveViaResolverThreads(req)
	if err != nil {
		return nil, err
	}
	resolved.Compress = true
	return resolved, nil
}

func getMinTTL(resp *dns.Msg) time.Duration {
	ttl := uint32(MaxTTL)
	for _, rr := range resp.Answer {
		if rr.Header().Ttl < ttl {
			ttl = rr.Header().Ttl
		}
	}
	if ttl < MinTTL {
		ttl = MaxTTL
	}
	return time.Duration(ttl) * time.Second
}

func sendTruncated(w dns.ResponseWriter, msgHdr dns.MsgHdr) {
	emptyResp := new(dns.Msg)
	emptyResp.MsgHdr = msgHdr
	emptyResp.Truncated = true
	w.WriteMsg(emptyResp)
}

func vacuumThread() {
	for {
		time.Sleep(VacuumPeriod * time.Second)
		resetUpstreamServers()
		memStats := new(runtime.MemStats)
		runtime.ReadMemStats(memStats)
		if memStats.Alloc > (*memSize)*1024*1024 {
			cache.Purge()
		}
	}
}

func failWithRcode(w dns.ResponseWriter, r *dns.Msg, rCode int) {
	m := new(dns.Msg)
	m.SetRcode(r, rCode)
	w.WriteMsg(m)
}

func handleSpecialNames(w dns.ResponseWriter, req *dns.Msg) bool {
	question := req.Question[0]
	if question.Qtype != dns.TypeANY {
		return false
	}
	m := new(dns.Msg)
	m.Id = req.Id
	hinfo := new(dns.HINFO)
	hinfo.Hdr = dns.RR_Header{Name: question.Name, Rrtype: dns.TypeHINFO,
		Class: dns.ClassINET, Ttl: 86400}
	hinfo.Cpu = "ANY is not supported any more"
	hinfo.Os = "See draft-jabley-dnsop-refuse-any"
	m.Answer = []dns.RR{hinfo}
	w.WriteMsg(m)
	return true
}

func route(w dns.ResponseWriter, req *dns.Msg) {
	keyP, err := getKey(req)
	if err != nil {
		failWithRcode(w, req, dns.RcodeRefused)
		return
	}
	if handleSpecialNames(w, req) {
		return
	}
	maxPayloadSize := getMaxPayloadSize(req)
	var resp *dns.Msg
	cacheValP, _ := cache.Get(*keyP)
	if cacheValP != nil {
		cacheVal := cacheValP.(CacheVal)
		remaining := -time.Since(cacheVal.ValidUntil)
		if remaining > 0 {
			resp = cacheVal.Response.Copy()
			resp.Id = req.Id
			resp.Question = req.Question
		}
	}
	if resp == nil {
		resp, err = resolve(req, keyP.DNSSEC)
		if err != nil {
			dns.HandleFailed(w, req)
			return
		}
		validUntil := time.Now().Add(getMinTTL(resp))
		cache.Add(*keyP, CacheVal{ValidUntil: validUntil, Response: resp})
	}
	packed, _ := resp.Pack()
	packedLen := len(packed)
	if uint16(packedLen) > maxPayloadSize {
		sendTruncated(w, resp.MsgHdr)
	} else {
		w.WriteMsg(resp)
	}
}
<|endoftext|>


<commit_before>package main

// Copyright (c) 2013, Sapphire Cat <https://github.com/sapphirecat>.  All
// rights reserved.  See the accompanying LICENSE file for license terms.

import (
	"fmt"
	"log"
	"net/http"
	"regexp"

	"github.com/elazarl/goproxy"
)

const (
	RuleForHttp = iota
	RuleForTls
)

type Rule struct {
	Matcher *regexp.Regexp
	Actor   Action  // function called when match hits -> string [hostname[:port]]
}

type Mode int;
type Action func(string, Mode) string;
type Ruleset struct {
	items []Rule
}

func (r *Ruleset) Add (a Rule) {
	r.items = append(r.items, a)
}

func (r *Ruleset) Length () int {
	return len(r.items)
}

func NewRuleset(capacity int) Ruleset {
	return Ruleset{
		make([]Rule, 0, capacity),
	}
}

func redirectPort (really bool, host string, port int) string {
	if really == true {
		return fmt.Sprintf("%s:%d", host, port)
	} else {
		return ""
	}
}

func SendHttpTo (host string) Action {
	return func(matched_host string, mode Mode) string {
		if mode == RuleForHttp {
			return host
		} else {
			return ""
		}
	}
}

func SendHttpToPort (host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForHttp, host, port);
	}
}

func SendAllTo (host string) Action {
	return func(matched_host string, mode Mode) string {
		if mode == RuleForHttp {
			return host
		} else {
			return redirectPort(true, host, 443)
		}
	}
}

func SendAllToPort (host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(true, host, port)
	}
}

func SendTlsTo (host string) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForTls, host, 443)
	}
}

func SendTlsToPort (host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForTls, host, port)
	}
}

func getTarget(rules Ruleset, hostname string, mode Mode) string {
	for _, rule := range rules.items {
		if rule.Matcher.MatchString(hostname) == true {
			if result := rule.Actor(hostname, mode) ; result != "" {
				return result
			}
		}
	}

	return ""
}

func NewDefaultHttpsRule(ruleset Ruleset, verbose bool) func(string, *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
	return func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
		target := getTarget(ruleset, host, RuleForTls)
		if target == "" {
			target = host
			if verbose {
				log.Println("!match HTTPS", host)
			}
		} else if verbose {
			log.Println("+HTTPS", host, ctx.Req.URL.Path)
		}

		return goproxy.OkConnect, target
	}
}

func NewDefaultHttpRule(ruleset Ruleset, verbose bool) func(*http.Request, *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	return func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		host := r.URL.Host
		target := getTarget(ruleset, host, RuleForHttp)
		if target != "" {
			r.URL.Host = target
			if verbose {
				log.Println("+plain", host, r.URL.Path)
			}
		} else if verbose {
			log.Println("!match plain", r.URL.Host)
		}

		return r, nil
	}
}

func SetDefaultRules(proxy *goproxy.ProxyHttpServer, rules Ruleset, verbose bool) {
	proxy.OnRequest().HandleConnectFunc(NewDefaultHttpsRule(rules, verbose))
	proxy.OnRequest().DoFunc(NewDefaultHttpRule(rules, verbose))
}
<commit_msg>go fmt new code<commit_after>package main

// Copyright (c) 2013, Sapphire Cat <https://github.com/sapphirecat>.  All
// rights reserved.  See the accompanying LICENSE file for license terms.

import (
	"fmt"
	"log"
	"net/http"
	"regexp"

	"github.com/elazarl/goproxy"
)

const (
	RuleForHttp = iota
	RuleForTls
)

type Rule struct {
	Matcher *regexp.Regexp
	Actor   Action // function called when match hits -> string [hostname[:port]]
}

type Mode int
type Action func(string, Mode) string
type Ruleset struct {
	items []Rule
}

func (r *Ruleset) Add(a Rule) {
	r.items = append(r.items, a)
}

func (r *Ruleset) Length() int {
	return len(r.items)
}

func NewRuleset(capacity int) Ruleset {
	return Ruleset{
		make([]Rule, 0, capacity),
	}
}

func redirectPort(really bool, host string, port int) string {
	if really == true {
		return fmt.Sprintf("%s:%d", host, port)
	} else {
		return ""
	}
}

func SendHttpTo(host string) Action {
	return func(matched_host string, mode Mode) string {
		if mode == RuleForHttp {
			return host
		} else {
			return ""
		}
	}
}

func SendHttpToPort(host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForHttp, host, port)
	}
}

func SendAllTo(host string) Action {
	return func(matched_host string, mode Mode) string {
		if mode == RuleForHttp {
			return host
		} else {
			return redirectPort(true, host, 443)
		}
	}
}

func SendAllToPort(host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(true, host, port)
	}
}

func SendTlsTo(host string) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForTls, host, 443)
	}
}

func SendTlsToPort(host string, port int) Action {
	return func(matched_host string, mode Mode) string {
		return redirectPort(mode == RuleForTls, host, port)
	}
}

func getTarget(rules Ruleset, hostname string, mode Mode) string {
	for _, rule := range rules.items {
		if rule.Matcher.MatchString(hostname) == true {
			if result := rule.Actor(hostname, mode); result != "" {
				return result
			}
		}
	}

	return ""
}

func NewDefaultHttpsRule(ruleset Ruleset, verbose bool) func(string, *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
	return func(host string, ctx *goproxy.ProxyCtx) (*goproxy.ConnectAction, string) {
		target := getTarget(ruleset, host, RuleForTls)
		if target == "" {
			target = host
			if verbose {
				log.Println("!match HTTPS", host)
			}
		} else if verbose {
			log.Println("+HTTPS", host, ctx.Req.URL.Path)
		}

		return goproxy.OkConnect, target
	}
}

func NewDefaultHttpRule(ruleset Ruleset, verbose bool) func(*http.Request, *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	return func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		host := r.URL.Host
		target := getTarget(ruleset, host, RuleForHttp)
		if target != "" {
			r.URL.Host = target
			if verbose {
				log.Println("+plain", host, r.URL.Path)
			}
		} else if verbose {
			log.Println("!match plain", r.URL.Host)
		}

		return r, nil
	}
}

func SetDefaultRules(proxy *goproxy.ProxyHttpServer, rules Ruleset, verbose bool) {
	proxy.OnRequest().HandleConnectFunc(NewDefaultHttpsRule(rules, verbose))
	proxy.OnRequest().DoFunc(NewDefaultHttpRule(rules, verbose))
}
<|endoftext|>


<commit_before>package workflows

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/stelligent/mu/common"
)

var ecsImagePattern = "amzn-ami-*-amazon-ecs-optimized"
var ec2ImagePattern = "amzn-ami-hvm-*-x86_64-gp2"

// NewEnvironmentUpserter create a new workflow for upserting an environment
func NewEnvironmentUpserter(ctx *common.Context, environmentName string) Executor {

	workflow := new(environmentWorkflow)
	ecsStackParams := make(map[string]string)
	elbStackParams := make(map[string]string)
	workflow.codeRevision = ctx.Config.Repo.Revision
	workflow.repoName = ctx.Config.Repo.Slug

	return newPipelineExecutor(
		workflow.environmentFinder(&ctx.Config, environmentName),
		workflow.environmentRolesetUpserter(ctx.RolesetManager, ctx.RolesetManager, ecsStackParams),
		workflow.environmentVpcUpserter(ctx.Config.Namespace, ecsStackParams, elbStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		workflow.environmentElbUpserter(ctx.Config.Namespace, ecsStackParams, elbStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		workflow.environmentUpserter(ctx.Config.Namespace, ecsStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager),
	)
}

// Find an environment in config, by name and set the reference
func (workflow *environmentWorkflow) environmentFinder(config *common.Config, environmentName string) Executor {

	return func() error {
		for _, e := range config.Environments {
			if strings.EqualFold(e.Name, environmentName) {
				if e.Provider == "" {
					e.Provider = common.EnvProviderEcs
				}
				workflow.environment = &e

				if e.Discovery.Provider == "consul" {
					return fmt.Errorf("Consul is no longer supported as a service discovery provider.  Check out the mu-consul extension for an alternative: https://github.com/stelligent/mu-consul")
				}

				return nil
			}
		}
		return common.Warningf("Unable to find environment named '%s' in configuration", environmentName)
	}
}

func (workflow *environmentWorkflow) environmentVpcUpserter(namespace string, ecsStackParams map[string]string, elbStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter, azCounter common.AZCounter) Executor {
	return func() error {
		environment := workflow.environment
		vpcStackParams := make(map[string]string)
		var err error

		var vpcStackName string
		var vpcTemplateName string
		if environment.VpcTarget.Environment != "" {
			targetNamespace := environment.VpcTarget.Namespace
			if targetNamespace == "" {
				targetNamespace = namespace
			}
			log.Debugf("VpcTarget exists for different environment; targeting that VPC")
			vpcStackName = common.CreateStackName(targetNamespace, common.StackTypeVpc, environment.VpcTarget.Environment)
		} else if environment.VpcTarget.VpcID != "" {
			log.Debugf("VpcTarget exists, so we will upsert the VPC stack that references the VPC attributes")
			vpcStackName = common.CreateStackName(namespace, common.StackTypeTarget, environment.Name)
			vpcTemplateName = "vpc-target.yml"

			// target VPC referenced from config
			vpcStackParams["VpcId"] = environment.VpcTarget.VpcID
			vpcStackParams["ElbSubnetIds"] = strings.Join(environment.VpcTarget.ElbSubnetIds, ",")
			vpcStackParams["InstanceSubnetIds"] = strings.Join(environment.VpcTarget.InstanceSubnetIds, ",")
		} else {
			log.Debugf("No VpcTarget, so we will upsert the VPC stack that manages the VPC")
			vpcStackName = common.CreateStackName(namespace, common.StackTypeVpc, environment.Name)
			vpcTemplateName = "vpc.yml"

			if environment.Cluster.InstanceTenancy != "" {
				vpcStackParams["InstanceTenancy"] = string(environment.Cluster.InstanceTenancy)
			}
			if environment.Cluster.SSHAllow != "" {
				vpcStackParams["SshAllow"] = environment.Cluster.SSHAllow
			} else {
				vpcStackParams["SshAllow"] = "0.0.0.0/0"
			}
			if environment.Cluster.KeyName != "" {
				vpcStackParams["BastionKeyName"] = environment.Cluster.KeyName
				vpcStackParams["BastionImageId"], err = imageFinder.FindLatestImageID(ec2ImagePattern)
				if err != nil {
					return err
				}
			}

			vpcStackParams["ElbInternal"] = strconv.FormatBool(environment.Loadbalancer.Internal)
		}

		azCount, err := azCounter.CountAZs()
		if err != nil {
			return err
		}
		if azCount < 2 {
			return fmt.Errorf("Only found %v availability zones...need at least 2", azCount)
		}
		vpcStackParams["AZCount"] = strconv.Itoa(azCount)

		if vpcTemplateName != "" {
			log.Noticef("Upserting VPC environment '%s' ...", environment.Name)

			tags := createTagMap(&EnvironmentTags{
				Environment: environment.Name,
				Type:        string(common.StackTypeVpc),
				Provider:    string(environment.Provider),
				Revision:    workflow.codeRevision,
				Repo:        workflow.repoName,
			})

			err = stackUpserter.UpsertStack(vpcStackName, vpcTemplateName, environment, vpcStackParams, tags, workflow.cloudFormationRoleArn)
			if err != nil {
				return err
			}

			log.Debugf("Waiting for stack '%s' to complete", vpcStackName)
			stack := stackWaiter.AwaitFinalStatus(vpcStackName)

			if stack == nil {
				return fmt.Errorf("Unable to create stack %s", vpcStackName)
			}
			if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
				return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
			}
		}

		ecsStackParams["VpcId"] = fmt.Sprintf("%s-VpcId", vpcStackName)
		ecsStackParams["InstanceSubnetIds"] = fmt.Sprintf("%s-InstanceSubnetIds", vpcStackName)

		elbStackParams["VpcId"] = fmt.Sprintf("%s-VpcId", vpcStackName)
		elbStackParams["ElbSubnetIds"] = fmt.Sprintf("%s-ElbSubnetIds", vpcStackName)

		return nil
	}
}

func (workflow *environmentWorkflow) environmentRolesetUpserter(rolesetUpserter common.RolesetUpserter, rolesetGetter common.RolesetGetter, ecsStackParams map[string]string) Executor {
	return func() error {
		err := rolesetUpserter.UpsertCommonRoleset()
		if err != nil {
			return err
		}

		commonRoleset, err := rolesetGetter.GetCommonRoleset()
		if err != nil {
			return err
		}

		workflow.cloudFormationRoleArn = commonRoleset["CloudFormationRoleArn"]

		err = rolesetUpserter.UpsertEnvironmentRoleset(workflow.environment.Name)
		if err != nil {
			return err
		}

		environmentRoleset, err := rolesetGetter.GetEnvironmentRoleset(workflow.environment.Name)
		if err != nil {
			return err
		}

		ecsStackParams["EC2InstanceProfileArn"] = environmentRoleset["EC2InstanceProfileArn"]

		return nil
	}
}

func (workflow *environmentWorkflow) environmentElbUpserter(namespace string, ecsStackParams map[string]string, elbStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		environment := workflow.environment
		envStackName := common.CreateStackName(namespace, common.StackTypeLoadBalancer, environment.Name)

		log.Noticef("Upserting ELB environment '%s' ...", environment.Name)

		stackParams := elbStackParams

		if environment.Loadbalancer.Certificate != "" {
			stackParams["ElbCert"] = environment.Loadbalancer.Certificate
		}

		if environment.Loadbalancer.HostedZone != "" {
			stackParams["ElbDomainName"] = environment.Loadbalancer.HostedZone

			if environment.Loadbalancer.Name == "" {
				// default to env name
				stackParams["ElbHostName"] = environment.Name
			} else {
				stackParams["ElbHostName"] = environment.Loadbalancer.Name
			}
		}

		if environment.Discovery.Name == "" {
			stackParams["ServiceDiscoveryName"] = fmt.Sprintf("%s.%s.local", environment.Name, namespace)
		} else {
			stackParams["ServiceDiscoveryName"] = environment.Discovery.Name
		}

		stackParams["ElbInternal"] = strconv.FormatBool(environment.Loadbalancer.Internal)

		tags := createTagMap(&EnvironmentTags{
			Environment: environment.Name,
			Type:        string(common.StackTypeLoadBalancer),
			Provider:    string(environment.Provider),
			Revision:    workflow.codeRevision,
			Repo:        workflow.repoName,
		})

		err := stackUpserter.UpsertStack(envStackName, "elb.yml", environment, stackParams, tags, workflow.cloudFormationRoleArn)
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", envStackName)
		stack := stackWaiter.AwaitFinalStatus(envStackName)

		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", envStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		ecsStackParams["ElbSecurityGroup"] = stack.Outputs["ElbInstanceSecurityGroup"]

		return nil
	}
}

func (workflow *environmentWorkflow) environmentUpserter(namespace string, ecsStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Debugf("Using provider '%s' for environment", workflow.environment.Provider)

		environment := workflow.environment
		envStackName := common.CreateStackName(namespace, common.StackTypeEnv, environment.Name)

		stackParams := ecsStackParams

		var templateName string
		var imagePattern string
		if environment.Provider == common.EnvProviderEcs {
			templateName = "env-ecs.yml"
			imagePattern = ecsImagePattern
			stackParams["LaunchType"] = "EC2"
		} else if environment.Provider == common.EnvProviderEcsFargate {
			templateName = "env-ecs.yml"
			imagePattern = ecsImagePattern
			stackParams["LaunchType"] = "FARGATE"
		} else if environment.Provider == common.EnvProviderEc2 {
			templateName = "env-ec2.yml"
			imagePattern = ec2ImagePattern
		}

		log.Noticef("Upserting environment '%s' ...", environment.Name)

		if environment.Cluster.SSHAllow != "" {
			stackParams["SshAllow"] = environment.Cluster.SSHAllow
		} else {
			stackParams["SshAllow"] = "0.0.0.0/0"
		}
		if environment.Cluster.InstanceType != "" {
			stackParams["InstanceType"] = environment.Cluster.InstanceType
		}
		if environment.Cluster.ExtraUserData != "" {
			stackParams["ExtraUserData"] = environment.Cluster.ExtraUserData
		}
		if environment.Cluster.ImageID != "" {
			stackParams["ImageId"] = environment.Cluster.ImageID
		} else {
			var err error
			stackParams["ImageId"], err = imageFinder.FindLatestImageID(imagePattern)
			if err != nil {
				return err
			}

		}
		if environment.Cluster.ImageOsType != "" {
			stackParams["ImageOsType"] = environment.Cluster.ImageOsType
		}
		if environment.Cluster.DesiredCapacity != 0 {
			stackParams["DesiredCapacity"] = strconv.Itoa(environment.Cluster.DesiredCapacity)
		}
		if environment.Cluster.MinSize != 0 {
			stackParams["MinSize"] = strconv.Itoa(environment.Cluster.MinSize)
		}
		if environment.Cluster.MaxSize != 0 {
			stackParams["MaxSize"] = strconv.Itoa(environment.Cluster.MaxSize)
		}
		if environment.Cluster.KeyName != "" {
			stackParams["KeyName"] = environment.Cluster.KeyName
		}
		if environment.Cluster.TargetCPUReservation != 0 {
			stackParams["TargetCPUReservation"] = strconv.Itoa(environment.Cluster.TargetCPUReservation)
		}
		if environment.Cluster.TargetMemoryReservation != 0 {
			stackParams["TargetMemoryReservation"] = strconv.Itoa(environment.Cluster.TargetMemoryReservation)
		}
		if environment.Cluster.HTTPProxy != "" {
			stackParams["HttpProxy"] = environment.Cluster.HTTPProxy
		}

		tags := createTagMap(&EnvironmentTags{
			Environment: environment.Name,
			Type:        string(common.StackTypeEnv),
			Provider:    string(environment.Provider),
			Revision:    workflow.codeRevision,
			Repo:        workflow.repoName,
		})

		err := stackUpserter.UpsertStack(envStackName, templateName, environment, stackParams, tags, workflow.cloudFormationRoleArn)
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", envStackName)
		stack := stackWaiter.AwaitFinalStatus(envStackName)

		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", envStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
<commit_msg>Use a config map instead of conditionals<commit_after>package workflows

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/stelligent/mu/common"
)

var ecsImagePattern = "amzn-ami-*-amazon-ecs-optimized"
var ec2ImagePattern = "amzn-ami-hvm-*-x86_64-gp2"

// NewEnvironmentUpserter create a new workflow for upserting an environment
func NewEnvironmentUpserter(ctx *common.Context, environmentName string) Executor {

	workflow := new(environmentWorkflow)
	ecsStackParams := make(map[string]string)
	elbStackParams := make(map[string]string)
	workflow.codeRevision = ctx.Config.Repo.Revision
	workflow.repoName = ctx.Config.Repo.Slug

	return newPipelineExecutor(
		workflow.environmentFinder(&ctx.Config, environmentName),
		workflow.environmentRolesetUpserter(ctx.RolesetManager, ctx.RolesetManager, ecsStackParams),
		workflow.environmentVpcUpserter(ctx.Config.Namespace, ecsStackParams, elbStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		workflow.environmentElbUpserter(ctx.Config.Namespace, ecsStackParams, elbStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		workflow.environmentUpserter(ctx.Config.Namespace, ecsStackParams, ctx.StackManager, ctx.StackManager, ctx.StackManager),
	)
}

// Find an environment in config, by name and set the reference
func (workflow *environmentWorkflow) environmentFinder(config *common.Config, environmentName string) Executor {

	return func() error {
		for _, e := range config.Environments {
			if strings.EqualFold(e.Name, environmentName) {
				if e.Provider == "" {
					e.Provider = common.EnvProviderEcs
				}
				workflow.environment = &e

				if e.Discovery.Provider == "consul" {
					return fmt.Errorf("Consul is no longer supported as a service discovery provider.  Check out the mu-consul extension for an alternative: https://github.com/stelligent/mu-consul")
				}

				return nil
			}
		}
		return common.Warningf("Unable to find environment named '%s' in configuration", environmentName)
	}
}

func (workflow *environmentWorkflow) environmentVpcUpserter(namespace string, ecsStackParams map[string]string, elbStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter, azCounter common.AZCounter) Executor {
	return func() error {
		environment := workflow.environment
		vpcStackParams := make(map[string]string)
		var err error

		var vpcStackName string
		var vpcTemplateName string
		if environment.VpcTarget.Environment != "" {
			targetNamespace := environment.VpcTarget.Namespace
			if targetNamespace == "" {
				targetNamespace = namespace
			}
			log.Debugf("VpcTarget exists for different environment; targeting that VPC")
			vpcStackName = common.CreateStackName(targetNamespace, common.StackTypeVpc, environment.VpcTarget.Environment)
		} else if environment.VpcTarget.VpcID != "" {
			log.Debugf("VpcTarget exists, so we will upsert the VPC stack that references the VPC attributes")
			vpcStackName = common.CreateStackName(namespace, common.StackTypeTarget, environment.Name)
			vpcTemplateName = "vpc-target.yml"

			// target VPC referenced from config
			vpcStackParams["VpcId"] = environment.VpcTarget.VpcID
			vpcStackParams["ElbSubnetIds"] = strings.Join(environment.VpcTarget.ElbSubnetIds, ",")
			vpcStackParams["InstanceSubnetIds"] = strings.Join(environment.VpcTarget.InstanceSubnetIds, ",")
		} else {
			log.Debugf("No VpcTarget, so we will upsert the VPC stack that manages the VPC")
			vpcStackName = common.CreateStackName(namespace, common.StackTypeVpc, environment.Name)
			vpcTemplateName = "vpc.yml"

			if environment.Cluster.InstanceTenancy != "" {
				vpcStackParams["InstanceTenancy"] = string(environment.Cluster.InstanceTenancy)
			}
			if environment.Cluster.SSHAllow != "" {
				vpcStackParams["SshAllow"] = environment.Cluster.SSHAllow
			} else {
				vpcStackParams["SshAllow"] = "0.0.0.0/0"
			}
			if environment.Cluster.KeyName != "" {
				vpcStackParams["BastionKeyName"] = environment.Cluster.KeyName
				vpcStackParams["BastionImageId"], err = imageFinder.FindLatestImageID(ec2ImagePattern)
				if err != nil {
					return err
				}
			}

			vpcStackParams["ElbInternal"] = strconv.FormatBool(environment.Loadbalancer.Internal)
		}

		azCount, err := azCounter.CountAZs()
		if err != nil {
			return err
		}
		if azCount < 2 {
			return fmt.Errorf("Only found %v availability zones...need at least 2", azCount)
		}
		vpcStackParams["AZCount"] = strconv.Itoa(azCount)

		if vpcTemplateName != "" {
			log.Noticef("Upserting VPC environment '%s' ...", environment.Name)

			tags := createTagMap(&EnvironmentTags{
				Environment: environment.Name,
				Type:        string(common.StackTypeVpc),
				Provider:    string(environment.Provider),
				Revision:    workflow.codeRevision,
				Repo:        workflow.repoName,
			})

			err = stackUpserter.UpsertStack(vpcStackName, vpcTemplateName, environment, vpcStackParams, tags, workflow.cloudFormationRoleArn)
			if err != nil {
				return err
			}

			log.Debugf("Waiting for stack '%s' to complete", vpcStackName)
			stack := stackWaiter.AwaitFinalStatus(vpcStackName)

			if stack == nil {
				return fmt.Errorf("Unable to create stack %s", vpcStackName)
			}
			if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
				return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
			}
		}

		ecsStackParams["VpcId"] = fmt.Sprintf("%s-VpcId", vpcStackName)
		ecsStackParams["InstanceSubnetIds"] = fmt.Sprintf("%s-InstanceSubnetIds", vpcStackName)

		elbStackParams["VpcId"] = fmt.Sprintf("%s-VpcId", vpcStackName)
		elbStackParams["ElbSubnetIds"] = fmt.Sprintf("%s-ElbSubnetIds", vpcStackName)

		return nil
	}
}

func (workflow *environmentWorkflow) environmentRolesetUpserter(rolesetUpserter common.RolesetUpserter, rolesetGetter common.RolesetGetter, ecsStackParams map[string]string) Executor {
	return func() error {
		err := rolesetUpserter.UpsertCommonRoleset()
		if err != nil {
			return err
		}

		commonRoleset, err := rolesetGetter.GetCommonRoleset()
		if err != nil {
			return err
		}

		workflow.cloudFormationRoleArn = commonRoleset["CloudFormationRoleArn"]

		err = rolesetUpserter.UpsertEnvironmentRoleset(workflow.environment.Name)
		if err != nil {
			return err
		}

		environmentRoleset, err := rolesetGetter.GetEnvironmentRoleset(workflow.environment.Name)
		if err != nil {
			return err
		}

		ecsStackParams["EC2InstanceProfileArn"] = environmentRoleset["EC2InstanceProfileArn"]

		return nil
	}
}

func (workflow *environmentWorkflow) environmentElbUpserter(namespace string, ecsStackParams map[string]string, elbStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		environment := workflow.environment
		envStackName := common.CreateStackName(namespace, common.StackTypeLoadBalancer, environment.Name)

		log.Noticef("Upserting ELB environment '%s' ...", environment.Name)

		stackParams := elbStackParams

		if environment.Loadbalancer.Certificate != "" {
			stackParams["ElbCert"] = environment.Loadbalancer.Certificate
		}

		if environment.Loadbalancer.HostedZone != "" {
			stackParams["ElbDomainName"] = environment.Loadbalancer.HostedZone

			if environment.Loadbalancer.Name == "" {
				// default to env name
				stackParams["ElbHostName"] = environment.Name
			} else {
				stackParams["ElbHostName"] = environment.Loadbalancer.Name
			}
		}

		if environment.Discovery.Name == "" {
			stackParams["ServiceDiscoveryName"] = fmt.Sprintf("%s.%s.local", environment.Name, namespace)
		} else {
			stackParams["ServiceDiscoveryName"] = environment.Discovery.Name
		}

		stackParams["ElbInternal"] = strconv.FormatBool(environment.Loadbalancer.Internal)

		tags := createTagMap(&EnvironmentTags{
			Environment: environment.Name,
			Type:        string(common.StackTypeLoadBalancer),
			Provider:    string(environment.Provider),
			Revision:    workflow.codeRevision,
			Repo:        workflow.repoName,
		})

		err := stackUpserter.UpsertStack(envStackName, "elb.yml", environment, stackParams, tags, workflow.cloudFormationRoleArn)
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", envStackName)
		stack := stackWaiter.AwaitFinalStatus(envStackName)

		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", envStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		ecsStackParams["ElbSecurityGroup"] = stack.Outputs["ElbInstanceSecurityGroup"]

		return nil
	}
}

func (workflow *environmentWorkflow) environmentUpserter(namespace string, ecsStackParams map[string]string, imageFinder common.ImageFinder, stackUpserter common.StackUpserter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Debugf("Using provider '%s' for environment", workflow.environment.Provider)

		environment := workflow.environment
		envStackName := common.CreateStackName(namespace, common.StackTypeEnv, environment.Name)

		stackParams := ecsStackParams

		var templateName string
		var imagePattern string
		envMapping := map[common.EnvProvider]map[string]string{
			common.EnvProviderEcs: map[string]string{
				"templateName": "env-ecs.yml",
				"imagePattern": ecsImagePattern,
				"launchType":   "EC2"},
			common.EnvProviderEcsFargate: map[string]string{
				"templateName": "env-ecs.yml",
				"imagePattern": ecsImagePattern,
				"launchType":   "FARGATE"},
			common.EnvProviderEc2: map[string]string{
				"templateName": "env-ec2.yml",
				"imagePattern": ec2ImagePattern,
				"launchType":   ""}}
		templateName = envMapping[environment.Provider]["templateName"]
		imagePattern = envMapping[environment.Provider]["imagePattern"]
		stackParams["LaunchType"] = envMapping[environment.Provider]["launchType"]

		log.Noticef("Upserting environment '%s' ...", environment.Name)

		if environment.Cluster.SSHAllow != "" {
			stackParams["SshAllow"] = environment.Cluster.SSHAllow
		} else {
			stackParams["SshAllow"] = "0.0.0.0/0"
		}
		if environment.Cluster.InstanceType != "" {
			stackParams["InstanceType"] = environment.Cluster.InstanceType
		}
		if environment.Cluster.ExtraUserData != "" {
			stackParams["ExtraUserData"] = environment.Cluster.ExtraUserData
		}
		if environment.Cluster.ImageID != "" {
			stackParams["ImageId"] = environment.Cluster.ImageID
		} else {
			var err error
			stackParams["ImageId"], err = imageFinder.FindLatestImageID(imagePattern)
			if err != nil {
				return err
			}

		}
		if environment.Cluster.ImageOsType != "" {
			stackParams["ImageOsType"] = environment.Cluster.ImageOsType
		}
		if environment.Cluster.DesiredCapacity != 0 {
			stackParams["DesiredCapacity"] = strconv.Itoa(environment.Cluster.DesiredCapacity)
		}
		if environment.Cluster.MinSize != 0 {
			stackParams["MinSize"] = strconv.Itoa(environment.Cluster.MinSize)
		}
		if environment.Cluster.MaxSize != 0 {
			stackParams["MaxSize"] = strconv.Itoa(environment.Cluster.MaxSize)
		}
		if environment.Cluster.KeyName != "" {
			stackParams["KeyName"] = environment.Cluster.KeyName
		}
		if environment.Cluster.TargetCPUReservation != 0 {
			stackParams["TargetCPUReservation"] = strconv.Itoa(environment.Cluster.TargetCPUReservation)
		}
		if environment.Cluster.TargetMemoryReservation != 0 {
			stackParams["TargetMemoryReservation"] = strconv.Itoa(environment.Cluster.TargetMemoryReservation)
		}
		if environment.Cluster.HTTPProxy != "" {
			stackParams["HttpProxy"] = environment.Cluster.HTTPProxy
		}

		tags := createTagMap(&EnvironmentTags{
			Environment: environment.Name,
			Type:        string(common.StackTypeEnv),
			Provider:    string(environment.Provider),
			Revision:    workflow.codeRevision,
			Repo:        workflow.repoName,
		})

		err := stackUpserter.UpsertStack(envStackName, templateName, environment, stackParams, tags, workflow.cloudFormationRoleArn)
		if err != nil {
			return err
		}
		log.Debugf("Waiting for stack '%s' to complete", envStackName)
		stack := stackWaiter.AwaitFinalStatus(envStackName)

		if stack == nil {
			return fmt.Errorf("Unable to create stack %s", envStackName)
		}
		if strings.HasSuffix(stack.Status, "ROLLBACK_COMPLETE") || !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/cloudson/gitql/parser"
	"github.com/cloudson/gitql/runtime"
	"github.com/cloudson/gitql/semantical"
	"os"
	"path/filepath"
)

func main() {

	query := flag.String("q", "", "The Query to search")
	pathString := flag.String("p", ".", "The (optional) path to run gitql")
	version := flag.Bool("v", false, "The version of gitql")
	flag.Parse()

	if *version {
		// @todo refactor to dynamic value
		fmt.Println("Gitql 1.0.0-RC4")
		os.Exit(0)
	}

	if *query == "" {
		flag.Usage()
		os.Exit(1)
	}

	path, errFile := filepath.Abs(*pathString)

	if errFile != nil {
		panic(errFile)
	}

	parser.New(*query)
	ast, errGit := parser.AST()
	if errGit != nil {
		panic(errGit)
	}
	ast.Path = &path
	errGit = semantical.Analysis(ast)
	if errGit != nil {
		panic(errGit)
	}

	runtime.Run(ast)
}
<commit_msg>Switch -q flag to argument #9<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/cloudson/gitql/parser"
	"github.com/cloudson/gitql/runtime"
	"github.com/cloudson/gitql/semantical"
	"os"
	"path/filepath"
)

func main() {

	pathString := flag.String("p", ".", "The (optional) path to run gitql")
	version := flag.Bool("v", false, "The version of gitql")
	flag.Parse()

	if *version {
		// @todo refactor to dynamic value
		fmt.Println("Gitql 1.0.0-RC4")
		os.Exit(0)
	}

	query := flag.Arg(0)
	if query == "" {
		flag.Usage()
		os.Exit(1)
	}

	path, errFile := filepath.Abs(*pathString)

	if errFile != nil {
		panic(errFile)
	}

	parser.New(query)
	ast, errGit := parser.AST()
	if errGit != nil {
		panic(errGit)
	}
	ast.Path = &path
	errGit = semantical.Analysis(ast)
	if errGit != nil {
		panic(errGit)
	}

	runtime.Run(ast)
}
<|endoftext|>


<commit_before>package gleam

import (
	"fmt"
	"log"
	"os"
	"syscall"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/mikespook/golib/signal"
	"github.com/mikespook/schego"
)

type Gleam struct {
	lua        *luaEnv
	config     Config
	mqttClient mqtt.Client
	scheduler  *schego.Scheduler
}

func NewGleam(root string) *Gleam {
	return &Gleam{
		lua: newLuaEnv(root),
	}
}

func (g *Gleam) Init() error {
	if err := g.lua.Init(&g.config); err != nil {
		return err
	}
	if err := g.initMQTT(); err != nil {
		return err
	}
	if err := g.initSchedule(&g.config); err != nil {
		return err
	}
	return g.lua.onEvent("afterInit", g.mqttClient)
}

func (g *Gleam) initSchedule(config *Config) error {
	g.scheduler = schego.New(config.Schedule.Tick * time.Millisecond)
	g.scheduler.ErrorFunc = g.lua.errorFunc
	for name, interval := range config.Schedule.Tasks {
		g.scheduler.Add(name, time.Now(), interval*time.Millisecond, schego.ForEver, g.lua.newExecFunc(name, g.mqttClient))
	}
	return nil
}

func (g *Gleam) initMQTT() error {
	opts := mqtt.NewClientOptions()
	for _, broker := range g.config.MQTT {
		opts.AddBroker(broker.Addr)
		log.Printf("Add Broker: %s@%s", broker.Username, broker.Addr)
		if broker.Username != "" {
			opts.SetUsername(broker.Username).SetPassword(broker.Password)
		}
	}
	opts.SetClientID(g.config.ClientId)
	log.Printf("ClientId: %s", g.config.ClientId)
	opts.SetDefaultPublishHandler(g.lua.defaultMQTTHandler)
	opts.SetAutoReconnect(true)
	g.mqttClient = mqtt.NewClient(opts)
	if token := g.mqttClient.Connect(); token.Wait() && token.Error() != nil {
		return token.Error()
	}
	for name, qos := range g.config.Tasks {
		topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
		if token := g.mqttClient.Subscribe(topic, qos, g.lua.newMQTTHandler(name)); token.Wait() && token.Error() != nil {
			return token.Error()
		}
		log.Printf("Subscribe: %s = %d", topic, qos)
	}
	return nil
}

func (g *Gleam) Serve() error {
	go g.scheduler.Serve()

	sh := signal.New(nil)
	sh.Bind(os.Interrupt, func() uint {
		return signal.BreakExit
	})
	sh.Bind(syscall.SIGHUP, func() uint {
		log.Printf("Reloading scripts")
		for name, qos := range g.config.Tasks {
			topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
			if token := g.mqttClient.Unsubscribe(topic); token.Wait() && token.Error() != nil {
				log.Printf("Unsubscribe error: %s", token.Error())
			}
			log.Printf("Unsubscribe: %s", topic)
			if token := g.mqttClient.Subscribe(topic, qos, g.lua.newMQTTHandler(name)); token.Wait() && token.Error() != nil {
				log.Printf("Subscribe error: %s", token.Error())
			}
			log.Printf("Subscribe: %s = %d", topic, qos)
		}
		return signal.Continue
	})
	sh.Wait()
	return nil
}

func (g *Gleam) Final() error {
	if err := g.scheduler.Close(); err != nil {
		return err
	}

	for name := range g.config.Tasks {
		topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
		if token := g.mqttClient.Unsubscribe(topic); token.Wait() && token.Error() != nil {
			log.Printf("Unsubscribe error: %s", token.Error())
		}
		log.Printf("Unsubscribe: %s", topic)
	}

	if err := g.lua.onEvent("beforeFinalize", g.mqttClient); err != nil {
		log.Printf("BeforeFinalize: %s", err)
	}
	if g.config.FinalTick != 0 {
		time.Sleep(g.config.FinalTick * time.Millisecond)
	}
	g.mqttClient.Disconnect(500)
	return g.lua.Final()
}
<commit_msg>SIGTERM<commit_after>package gleam

import (
	"fmt"
	"log"
	"os"
	"syscall"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
	"github.com/mikespook/golib/signal"
	"github.com/mikespook/schego"
)

type Gleam struct {
	lua        *luaEnv
	config     Config
	mqttClient mqtt.Client
	scheduler  *schego.Scheduler
}

func NewGleam(root string) *Gleam {
	return &Gleam{
		lua: newLuaEnv(root),
	}
}

func (g *Gleam) Init() error {
	if err := g.lua.Init(&g.config); err != nil {
		return err
	}
	if err := g.initMQTT(); err != nil {
		return err
	}
	if err := g.initSchedule(&g.config); err != nil {
		return err
	}
	return g.lua.onEvent("afterInit", g.mqttClient)
}

func (g *Gleam) initSchedule(config *Config) error {
	g.scheduler = schego.New(config.Schedule.Tick * time.Millisecond)
	g.scheduler.ErrorFunc = g.lua.errorFunc
	for name, interval := range config.Schedule.Tasks {
		g.scheduler.Add(name, time.Now(), interval*time.Millisecond, schego.ForEver, g.lua.newExecFunc(name, g.mqttClient))
	}
	return nil
}

func (g *Gleam) initMQTT() error {
	opts := mqtt.NewClientOptions()
	for _, broker := range g.config.MQTT {
		opts.AddBroker(broker.Addr)
		log.Printf("Add Broker: %s@%s", broker.Username, broker.Addr)
		if broker.Username != "" {
			opts.SetUsername(broker.Username).SetPassword(broker.Password)
		}
	}
	opts.SetClientID(g.config.ClientId)
	log.Printf("ClientId: %s", g.config.ClientId)
	opts.SetDefaultPublishHandler(g.lua.defaultMQTTHandler)
	opts.SetAutoReconnect(true)
	g.mqttClient = mqtt.NewClient(opts)
	if token := g.mqttClient.Connect(); token.Wait() && token.Error() != nil {
		return token.Error()
	}
	for name, qos := range g.config.Tasks {
		topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
		if token := g.mqttClient.Subscribe(topic, qos, g.lua.newMQTTHandler(name)); token.Wait() && token.Error() != nil {
			return token.Error()
		}
		log.Printf("Subscribe: %s = %d", topic, qos)
	}
	return nil
}

func (g *Gleam) Serve() error {
	go g.scheduler.Serve()

	sh := signal.New(nil)
	sh.Bind(os.Interrupt, func() uint {
		return signal.BreakExit
	})
	sh.Bind(syscall.SIGTERM, func() uint {
		return signal.BreakExit
	})
	sh.Bind(syscall.SIGHUP, func() uint {
		log.Printf("Reloading scripts")
		for name, qos := range g.config.Tasks {
			topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
			if token := g.mqttClient.Unsubscribe(topic); token.Wait() && token.Error() != nil {
				log.Printf("Unsubscribe error: %s", token.Error())
			}
			log.Printf("Unsubscribe: %s", topic)
			if token := g.mqttClient.Subscribe(topic, qos, g.lua.newMQTTHandler(name)); token.Wait() && token.Error() != nil {
				log.Printf("Subscribe error: %s", token.Error())
			}
			log.Printf("Subscribe: %s = %d", topic, qos)
		}
		return signal.Continue
	})
	sh.Wait()
	return nil
}

func (g *Gleam) Final() error {
	if err := g.scheduler.Close(); err != nil {
		return err
	}

	for name := range g.config.Tasks {
		topic := fmt.Sprintf("%s/%s", g.config.Prefix, name)
		if token := g.mqttClient.Unsubscribe(topic); token.Wait() && token.Error() != nil {
			log.Printf("Unsubscribe error: %s", token.Error())
		}
		log.Printf("Unsubscribe: %s", topic)
	}

	if err := g.lua.onEvent("beforeFinalize", g.mqttClient); err != nil {
		log.Printf("BeforeFinalize: %s", err)
	}
	if g.config.FinalTick != 0 {
		time.Sleep(g.config.FinalTick * time.Millisecond)
	}
	g.mqttClient.Disconnect(500)
	return g.lua.Final()
}
<|endoftext|>


<commit_before>package wikifier

import (
	"strconv"
	"strings"
)

type secBlock struct {
	n           int
	isIntro     bool
	headerLevel int
	*parserBlock
}

func newSecBlock(name string, b *parserBlock) block {
	return &secBlock{parserBlock: b}
}

func (sec *secBlock) parse(page *Page) {
	enable := page.Opt.Page.EnableTitle

	// @page.enable.title causes the first header to be larger than the
	// rest. it also uses @page.title as the first header if no other text
	// is provided.
	sec.n = page.sectionN
	sec.isIntro = sec.n == 0 && enable
	page.sectionN++

	// find the level from the parent section
	level := 1
	var blk block = sec
	for blk != nil {
		if parentSec, ok := blk.parentBlock().(*secBlock); ok {
			level = parentSec.headerLevel + 1
			break
		}
		blk = blk.parentBlock()
	}

	// top-level headers start at h2 when @page.enable.title is true, since the
	// page title is the sole h1. otherwise, h1 is top-level.
	if enable && level == 1 {
		level++
	}

	// intro is always h1
	if sec.isIntro {
		level = 1
	}

	// max is h6
	if level > 6 {
		level = 6
	}

	sec.headerLevel = level

	// this must come last so the section order is correct
	sec.parserBlock.parse(page)
}

func (sec *secBlock) html(page *Page, el element) {
	// HEADING

	// determine if this is the intro section
	typ := "section-title"
	level := sec.headerLevel
	if sec.isIntro {
		typ = "section-page-title"
	}

	// use the page title if no other title is provided and @page.enable.title
	title, fmtTitle := sec.blockName(), HTML("")
	if sec.isIntro && title == "" {
		title = page.Title()
		fmtTitle = page.FmtTitle()
	}

	// we have a title
	if title != "" {

		// format title if we still need to
		if fmtTitle == "" {
			fmtTitle = page.parseFormattedTextOpts(title, &formatterOptions{pos: sec.openPos})
		}

		// TODO: meta section heading ID

		// heading ID
		headingID := PageNameLink(title, false)

		// add -n as needed if this is already used
		n := page.headingIDs[headingID]
		page.headingIDs[headingID]++
		if n != 0 {
			headingID += "-" + strconv.Itoa(n)
		}

		// create the heading
		h := el.createChild("h"+strconv.Itoa(level), typ)
		h.setAttr("id", "qa-"+headingID)
		h.addHTML(fmtTitle)
	}

	// CONTENT

	// iterate over content
	var contentToAdd []posContent
	for _, pc := range sec.posContent() {
		switch item := pc.content.(type) {
		case block:

			// create a section with the text up to this point
			sec.createParagraph(page, el, contentToAdd)
			contentToAdd = nil

			// adopt this block as my own
			item.html(page, item.el())
			el.addChild(item.el())

		case string:

			// if this is an empty line, create a new paragraph
			item = strings.TrimSpace(item)
			if item == "" {
				sec.createParagraph(page, el, contentToAdd)
				contentToAdd = nil
				continue
			}

			// otherwise, add it to the buffer
			contentToAdd = append(contentToAdd, pc)

		default:
			panic("not sure how to handle this content")
		}
	}

	// add whatever's left
	sec.createParagraph(page, el, contentToAdd)
}

func (sec *secBlock) createParagraph(page *Page, el element, pcs []posContent) {

	// this can be passed nothing
	if len(pcs) == 0 {
		return
	}

	// create a paragraph at first text node position
	p := newBlock("p", "", nil, sec, sec, pcs[0].position)
	p.appendContent(pcs, pcs[0].position)

	// parse and generate
	p.parse(page)
	p.html(page, p.el())

	// adopt it as my own
	el.addChild(p.el())
}
<commit_msg>respect page.enable.title<commit_after>package wikifier

import (
	"strconv"
	"strings"
)

type secBlock struct {
	n           int
	isIntro     bool
	headerLevel int
	*parserBlock
}

func newSecBlock(name string, b *parserBlock) block {
	return &secBlock{parserBlock: b}
}

func (sec *secBlock) parse(page *Page) {

	// wiki option
	enable := page.Opt.Page.EnableTitle

	// overwrite with local var if present
	val, _ := page.Get("page.enable.title")
	if boolVal, ok := val.(bool); ok {
		enable = boolVal
	}

	// @page.enable.title causes the first header to be larger than the
	// rest. it also uses @page.title as the first header if no other text
	// is provided.
	sec.n = page.sectionN
	sec.isIntro = sec.n == 0 && enable
	page.sectionN++

	// find the level from the parent section
	level := 1
	var blk block = sec
	for blk != nil {
		if parentSec, ok := blk.parentBlock().(*secBlock); ok {
			level = parentSec.headerLevel + 1
			break
		}
		blk = blk.parentBlock()
	}

	// top-level headers start at h2 when @page.enable.title is true, since the
	// page title is the sole h1. otherwise, h1 is top-level.
	if enable && level == 1 {
		level++
	}

	// intro is always h1
	if sec.isIntro {
		level = 1
	}

	// max is h6
	if level > 6 {
		level = 6
	}

	sec.headerLevel = level

	// this must come last so the section order is correct
	sec.parserBlock.parse(page)
}

func (sec *secBlock) html(page *Page, el element) {
	// HEADING

	// determine if this is the intro section
	typ := "section-title"
	level := sec.headerLevel
	if sec.isIntro {
		typ = "section-page-title"
	}

	// use the page title if no other title is provided and @page.enable.title
	title, fmtTitle := sec.blockName(), HTML("")
	if sec.isIntro && title == "" {
		title = page.Title()
		fmtTitle = page.FmtTitle()
	}

	// we have a title
	if title != "" {

		// format title if we still need to
		if fmtTitle == "" {
			fmtTitle = page.parseFormattedTextOpts(title, &formatterOptions{pos: sec.openPos})
		}

		// TODO: meta section heading ID

		// heading ID
		headingID := PageNameLink(title, false)

		// add -n as needed if this is already used
		n := page.headingIDs[headingID]
		page.headingIDs[headingID]++
		if n != 0 {
			headingID += "-" + strconv.Itoa(n)
		}

		// create the heading
		h := el.createChild("h"+strconv.Itoa(level), typ)
		h.setAttr("id", "qa-"+headingID)
		h.addHTML(fmtTitle)
	}

	// CONTENT

	// iterate over content
	var contentToAdd []posContent
	for _, pc := range sec.posContent() {
		switch item := pc.content.(type) {
		case block:

			// create a section with the text up to this point
			sec.createParagraph(page, el, contentToAdd)
			contentToAdd = nil

			// adopt this block as my own
			item.html(page, item.el())
			el.addChild(item.el())

		case string:

			// if this is an empty line, create a new paragraph
			item = strings.TrimSpace(item)
			if item == "" {
				sec.createParagraph(page, el, contentToAdd)
				contentToAdd = nil
				continue
			}

			// otherwise, add it to the buffer
			contentToAdd = append(contentToAdd, pc)

		default:
			panic("not sure how to handle this content")
		}
	}

	// add whatever's left
	sec.createParagraph(page, el, contentToAdd)
}

func (sec *secBlock) createParagraph(page *Page, el element, pcs []posContent) {

	// this can be passed nothing
	if len(pcs) == 0 {
		return
	}

	// create a paragraph at first text node position
	p := newBlock("p", "", nil, sec, sec, pcs[0].position)
	p.appendContent(pcs, pcs[0].position)

	// parse and generate
	p.parse(page)
	p.html(page, p.el())

	// adopt it as my own
	el.addChild(p.el())
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"
	"os/exec"
	"time"
)

var Char struct {
	class, name, race, acct string
	lvl                     int
	seen                    time.Time
}

func main() {
	f, err := os.OpenFile("bot.log", os.O_RDWR|os.O_APPEND|os.O_CREATE, 0640)
	if err != nil {
		log.Fatal(err)
	}
	log.SetOutput(f)

	// for who.go WhoChar(char, lvl, class, race, acct)
	var char = flag.String("char", "", "Character name for update or import. Ex: Rynshana")
	var lvl = flag.Int("lvl", 0, "Character level for update or import. Ex: 50")
	var class = flag.String("class", "", "Character class for initial import. Ex: \"Cleric\"")
	var race = flag.String("race", "", "Character race for initial import. Ex: \"Moon Elf\"")
	var acct = flag.String("acct", "", "Character account for initial import. Ex: Krimic")
	// for who.go WhoBatch(ppl)
	var who = flag.String("who", "", "Batched who output. Ex: [ 1 Ctr] Rarac  (Orc)|[ 2 War] Xatus  (Troll)")
	// for identify.go Identify(filename)
	var file = flag.String("import", "", "Parse file for identify stats, import to DB. Ex: newstats.txt")
	// for time.go Uptime(curup)
	var time = flag.String("time", "", "Parse uptime for boot tracking. Ex: 58:10:26")
	// for tell.go ReplyTo(char, tell)
	var tell = flag.String("tell", "", "Tell with command and maybe operant. Ex: \"stat a longsword\"")
	// run database backup and restore
	var backup = flag.Bool("bak", false, "Backup the toril.db database.")
	var restore = flag.String("res", "", "Restore the toril.db database from backup file.")

	flag.Parse()

	// only run one command at a time
	switch {
	case *char != "" && 50 >= *lvl && *lvl > 0 && *class != "" && *race != "" && *acct != "":
		WhoChar(*char, *lvl, *class, *race, *acct)
	case *file != "":
		Identify(*file)
	case *time != "":
		Uptime(*time)
	case *char != "" && *tell != "":
		ReplyTo(*char, *tell)
	case *who != "":
		WhoBatch(*who)
	case *backup:
		cmd := exec.Command("sh", "-c", "echo '.dump' | sqlite3 toril.db | gzip -c >toril.db.`date +\"%Y-%m-%d\"`.gz")
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
	case *restore != "":
		cmd := exec.Command("sh", "-c", "zcat "+*restore+" | sqlite3 toril.db")
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
	}
}
<commit_msg>Reformat for line width, add defer close to log file<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"os/exec"
	"time"
)

var Char struct {
	class, name, race, acct string
	lvl                     int
	seen                    time.Time
}

func main() {
	f, err := os.OpenFile("bot.log", os.O_RDWR|os.O_APPEND|os.O_CREATE, 0640)
	defer f.Close()
	if err != nil {
		log.Fatal(err)
	}
	log.SetOutput(f)

	// for who.go WhoChar(char, lvl, class, race, acct)
	var char = flag.String("char", "",
		"Character name for update or import. Ex: Rynshana")
	var lvl = flag.Int("lvl", 0,
		"Character level for update or import. Ex: 50")
	var class = flag.String("class", "",
		"Character class for initial import. Ex: \"Cleric\"")
	var race = flag.String("race", "",
		"Character race for initial import. Ex: \"Moon Elf\"")
	var acct = flag.String("acct", "",
		"Character account for initial import. Ex: Krimic")
	// for who.go WhoBatch(ppl)
	var who = flag.String("who", "",
		"Batched who output. Ex: [ 1 Ctr] Rarac  (Orc)|[ 2 War] Xatus  (Troll)")
	// for identify.go Identify(filename)
	var file = flag.String("import", "",
		"Parse file for identify stats, import to DB. Ex: newstats.txt")
	// for time.go Uptime(curup)
	var time = flag.String("time", "",
		"Parse uptime for boot tracking. Ex: 58:10:26")
	// for tell.go ReplyTo(char, tell)
	var tell = flag.String("tell", "",
		"Tell with command and maybe operant. Ex: \"stat a longsword\"")
	// run database backup and restore
	var backup = flag.Bool("bak", false,
		"Backup the toril.db database.")
	var restore = flag.String("res", "",
		"Restore the toril.db database from backup file.")

	flag.Parse()

	// only run one command at a time
	switch {
	case *char != "" && 50 >= *lvl && *lvl > 0 &&
		*class != "" && *race != "" && *acct != "":
		WhoChar(*char, *lvl, *class, *race, *acct)
	case *file != "":
		Identify(*file)
	case *time != "":
		Uptime(*time)
	case *char != "" && *tell != "":
		ReplyTo(*char, *tell)
	case *who != "":
		WhoBatch(*who)
	case *backup:
		cmd := exec.Command("sh", "-c",
			"echo '.dump' | sqlite3 toril.db | "+
			"gzip -c >toril.db.`date +\"%Y-%m-%d\"`.gz")
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
	case *restore != "":
		cmd := exec.Command("sh", "-c", "zcat "+*restore+" | sqlite3 toril.db")
		err := cmd.Run()
		if err != nil {
			log.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>package gps

import (
	"errors"
	"go/build"
	"go/parser"
	"go/token"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"unicode"
)

var (
	gorootSrc  = filepath.Join(build.Default.GOROOT, "src")
	ignoreTags = []string{} //[]string{"appengine", "ignore"} //TODO: appengine is a special case for now: https://github.com/tools/godep/issues/353
)

// returns the package in dir either from a cache or by importing it and then caching it
func fullPackageInDir(dir string) (pkg *build.Package, err error) {
	pkg, err = build.ImportDir(dir, build.FindOnly)
	if pkg.Goroot {
		pkg, err = build.ImportDir(pkg.Dir, 0)
	} else {
		err = fillPackage(pkg)
	}
	return pkg, err
}

// fillPackage full of info. Assumes p.Dir is set at a minimum
func fillPackage(p *build.Package) error {
	if p.Goroot {
		return nil
	}

	if p.SrcRoot == "" {
		for _, base := range build.Default.SrcDirs() {
			if strings.HasPrefix(p.Dir, base) {
				p.SrcRoot = base
			}
		}
	}

	if p.SrcRoot == "" {
		return errors.New("Unable to find SrcRoot for package " + p.ImportPath)
	}

	if p.Root == "" {
		p.Root = filepath.Dir(p.SrcRoot)
	}

	var buildMatch = "+build "
	var buildFieldSplit = func(r rune) bool {
		return unicode.IsSpace(r) || r == ','
	}

	//debugln("Filling package:", p.ImportPath, "from", p.Dir)
	gofiles, err := filepath.Glob(filepath.Join(p.Dir, "*.go"))
	if err != nil {
		//debugln("Error globbing", err)
		return err
	}

	if len(gofiles) == 0 {
		return &build.NoGoError{Dir: p.Dir}
	}

	var testImports []string
	var imports []string
NextFile:
	for _, file := range gofiles {
		//debugln(file)
		pf, err := parser.ParseFile(token.NewFileSet(), file, nil, parser.ImportsOnly|parser.ParseComments)
		if err != nil {
			return err
		}
		testFile := strings.HasSuffix(file, "_test.go")
		fname := filepath.Base(file)
		for _, c := range pf.Comments {
			if c.Pos() > pf.Package { // +build must come before package
				continue
			}
			ct := c.Text()
			if i := strings.Index(ct, buildMatch); i != -1 {
				for _, t := range strings.FieldsFunc(ct[i+len(buildMatch):], buildFieldSplit) {
					for _, tag := range ignoreTags {
						if t == tag {
							p.IgnoredGoFiles = append(p.IgnoredGoFiles, fname)
							continue NextFile
						}
					}

					//TODO: Needed in GPS?
					/*					if versionMatch.MatchString(t) && !isSameOrNewer(t, majorGoVersion) {
											debugln("Adding", fname, "to ignored list because of version tag", t)
											p.IgnoredGoFiles = append(p.IgnoredGoFiles, fname)
											continue NextFile
										}
										if versionNegativeMatch.MatchString(t) && isSameOrNewer(t[1:], majorGoVersion) {
											debugln("Adding", fname, "to ignored list because of version tag", t)
											p.IgnoredGoFiles = append(p.IgnoredGoFiles, fname)
											continue NextFile
										} */
				}
			}
		}
		if testFile {
			p.TestGoFiles = append(p.TestGoFiles, fname)
			if p.Name == "" {
				p.Name = strings.Split(pf.Name.Name, "_")[0]
			}
		} else {
			if p.Name == "" {
				p.Name = pf.Name.Name
			}
			p.GoFiles = append(p.GoFiles, fname)
		}
		for _, is := range pf.Imports {
			name, err := strconv.Unquote(is.Path.Value)
			if err != nil {
				return err // can't happen?
			}
			if testFile {
				testImports = append(testImports, name)
			} else {
				imports = append(imports, name)
			}
		}
	}
	imports = uniq(imports)
	testImports = uniq(testImports)
	p.Imports = imports
	p.TestImports = testImports
	return nil
}

func uniq(a []string) []string {
	if a == nil {
		return make([]string, 0)
	}
	var s string
	var i int
	if !sort.StringsAreSorted(a) {
		sort.Strings(a)
	}
	for _, t := range a {
		if t != s {
			a[i] = t
			i++
			s = t
		}
	}
	return a[:i]
}
<commit_msg>Remove go version code<commit_after>package gps

import (
	"errors"
	"go/build"
	"go/parser"
	"go/token"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"unicode"
)

var (
	gorootSrc  = filepath.Join(build.Default.GOROOT, "src")
	ignoreTags = []string{} //[]string{"appengine", "ignore"} //TODO: appengine is a special case for now: https://github.com/tools/godep/issues/353
)

// returns the package in dir either from a cache or by importing it and then caching it
func fullPackageInDir(dir string) (pkg *build.Package, err error) {
	pkg, err = build.ImportDir(dir, build.FindOnly)
	if pkg.Goroot {
		pkg, err = build.ImportDir(pkg.Dir, 0)
	} else {
		err = fillPackage(pkg)
	}
	return pkg, err
}

// fillPackage full of info. Assumes p.Dir is set at a minimum
func fillPackage(p *build.Package) error {
	if p.Goroot {
		return nil
	}

	if p.SrcRoot == "" {
		for _, base := range build.Default.SrcDirs() {
			if strings.HasPrefix(p.Dir, base) {
				p.SrcRoot = base
			}
		}
	}

	if p.SrcRoot == "" {
		return errors.New("Unable to find SrcRoot for package " + p.ImportPath)
	}

	if p.Root == "" {
		p.Root = filepath.Dir(p.SrcRoot)
	}

	var buildMatch = "+build "
	var buildFieldSplit = func(r rune) bool {
		return unicode.IsSpace(r) || r == ','
	}

	//debugln("Filling package:", p.ImportPath, "from", p.Dir)
	gofiles, err := filepath.Glob(filepath.Join(p.Dir, "*.go"))
	if err != nil {
		//debugln("Error globbing", err)
		return err
	}

	if len(gofiles) == 0 {
		return &build.NoGoError{Dir: p.Dir}
	}

	var testImports []string
	var imports []string
NextFile:
	for _, file := range gofiles {
		//debugln(file)
		pf, err := parser.ParseFile(token.NewFileSet(), file, nil, parser.ImportsOnly|parser.ParseComments)
		if err != nil {
			return err
		}
		testFile := strings.HasSuffix(file, "_test.go")
		fname := filepath.Base(file)
		for _, c := range pf.Comments {
			if c.Pos() > pf.Package { // +build must come before package
				continue
			}
			ct := c.Text()
			if i := strings.Index(ct, buildMatch); i != -1 {
				for _, t := range strings.FieldsFunc(ct[i+len(buildMatch):], buildFieldSplit) {
					for _, tag := range ignoreTags {
						if t == tag {
							p.IgnoredGoFiles = append(p.IgnoredGoFiles, fname)
							continue NextFile
						}
					}
				}
			}
		}
		if testFile {
			p.TestGoFiles = append(p.TestGoFiles, fname)
			if p.Name == "" {
				p.Name = strings.Split(pf.Name.Name, "_")[0]
			}
		} else {
			if p.Name == "" {
				p.Name = pf.Name.Name
			}
			p.GoFiles = append(p.GoFiles, fname)
		}
		for _, is := range pf.Imports {
			name, err := strconv.Unquote(is.Path.Value)
			if err != nil {
				return err // can't happen?
			}
			if testFile {
				testImports = append(testImports, name)
			} else {
				imports = append(imports, name)
			}
		}
	}
	imports = uniq(imports)
	testImports = uniq(testImports)
	p.Imports = imports
	p.TestImports = testImports
	return nil
}

func uniq(a []string) []string {
	if a == nil {
		return make([]string, 0)
	}
	var s string
	var i int
	if !sort.StringsAreSorted(a) {
		sort.Strings(a)
	}
	for _, t := range a {
		if t != s {
			a[i] = t
			i++
			s = t
		}
	}
	return a[:i]
}
<|endoftext|>


<commit_before>/*
Package goenv provides a virtualenv-like GOPATH and PATH isolation for go.

Usage is as follows:

	Create a goenv:
		goenv [NAME]

	Activate the goenv:
		. [NAME]/bin/activate

	Deactivate the goenv:
		deactivate

Example:
	goenv local
	. local/bin/activate
	deactivate
*/
package goenv

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const template = `
# This file must be used with "source bin/activate" or ". bin/activate"

if [[ -n "${GOENV+1}" ]]; then
	deactivate
fi

export GOENV=__GOENV__
export GOENV_OLDPS1=$PS1
export GOENV_OLDGOPATH=$GOPATH
export GOENV_OLDPATH=$PATH

export GOPATH=$GOENV:$GOPATH
export PATH="$GOENV/bin:$PATH"
export PS1="($(basename $GOENV))$PS1"

deactivate() {
	export PS1=$GOENV_OLDPS1
	export GOPATH=$GOENV_OLDGOPATH
	export PATH=$GOENV_OLDPATH

	unset GOENV GOENV_OLDPS1 GOENV_OLDPATH GOENV_OLDGOPATH
	unset -f deactivate
}
`

// writeScript writes the modifed script template.
func writeScript(goenv, path string) error {

	script := strings.Replace(template, "__GOENV__", goenv, -1)
	err := ioutil.WriteFile(path, []byte(script), 777)

	return err
}

// createSubdirs creates the goenv directory structure.
func createSubdirs(path string, names []string) error {

	for _, name := range names {
		subDir := filepath.Join(path, name)
		if err := os.MkdirAll(subDir, os.ModeDir|0755); err != nil {
			return err
		}
	}

	return nil
}

func init() {

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: goenv DEST_DIR\n")
	}
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	dir, err := os.Getwd()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting current directory: %s", err.Error())
		os.Exit(1)
	}

	path := filepath.Join(dir, flag.Arg(0))
	err = createSubdirs(path, []string{"src", "bin", "pkg"})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating directories: %s", err.Error())
		os.Exit(1)
	}

	scriptPath := filepath.Join(path, "bin/activate")
	err = writeScript(path, scriptPath)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing activate script: %s", err.Error())
		os.Exit(1)
	}
}
<commit_msg>revert package name<commit_after>/*
Package goenv provides a virtualenv-like GOPATH and PATH isolation for go.

Usage is as follows:

	Create a goenv:
		goenv [NAME]

	Activate the goenv:
		. [NAME]/bin/activate

	Deactivate the goenv:
		deactivate

Example:
	goenv local
	. local/bin/activate
	deactivate
*/
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

const template = `
# This file must be used with "source bin/activate" or ". bin/activate"

if [[ -n "${GOENV+1}" ]]; then
	deactivate
fi

export GOENV=__GOENV__
export GOENV_OLDPS1=$PS1
export GOENV_OLDGOPATH=$GOPATH
export GOENV_OLDPATH=$PATH

export GOPATH=$GOENV:$GOPATH
export PATH="$GOENV/bin:$PATH"
export PS1="($(basename $GOENV))$PS1"

deactivate() {
	export PS1=$GOENV_OLDPS1
	export GOPATH=$GOENV_OLDGOPATH
	export PATH=$GOENV_OLDPATH

	unset GOENV GOENV_OLDPS1 GOENV_OLDPATH GOENV_OLDGOPATH
	unset -f deactivate
}
`

// writeScript writes the modifed script template.
func writeScript(goenv, path string) error {

	script := strings.Replace(template, "__GOENV__", goenv, -1)
	err := ioutil.WriteFile(path, []byte(script), 777)

	return err
}

// createSubdirs creates the goenv directory structure.
func createSubdirs(path string, names []string) error {

	for _, name := range names {
		subDir := filepath.Join(path, name)
		if err := os.MkdirAll(subDir, os.ModeDir|0755); err != nil {
			return err
		}
	}

	return nil
}

func init() {

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: goenv DEST_DIR\n")
	}
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	dir, err := os.Getwd()

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting current directory: %s", err.Error())
		os.Exit(1)
	}

	path := filepath.Join(dir, flag.Arg(0))
	err = createSubdirs(path, []string{"src", "bin", "pkg"})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating directories: %s", err.Error())
		os.Exit(1)
	}

	scriptPath := filepath.Join(path, "bin/activate")
	err = writeScript(path, scriptPath)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing activate script: %s", err.Error())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package gorma

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/raphael/goa/design"
	"github.com/raphael/goa/goagen/codegen"
	"gopkg.in/alecthomas/kingpin.v2"
)

// Generator is the application code generator.
type Generator struct {
	genfiles []string
}

// Generate is the generator entry point called by the meta generator.
func Generate(api *design.APIDefinition) ([]string, error) {
	g, err := NewGenerator()
	if err != nil {
		return nil, err
	}
	return g.Generate(api)
}

// NewGenerator returns the application code generator.
func NewGenerator() (*Generator, error) {
	return new(Generator), nil
}

// Generate produces the skeleton main.
func (g *Generator) Generate(api *design.APIDefinition) ([]string, error) {

	os.RemoveAll(ModelDir())
	os.MkdirAll(ModelDir(), 0755)
	app := kingpin.New("Model generator", "model generator")
	codegen.RegisterFlags(app)
	_, err := app.Parse(os.Args[1:])
	if err != nil {
		panic(err)
	}
	imp, err := filepath.Rel(filepath.Join(os.Getenv("GOPATH"), "src"), codegen.OutputDir)
	if err != nil {
		return nil, err
	}
	imp = filepath.Join(imp, "app")
	imports := []*codegen.ImportSpec{
		codegen.SimpleImport(imp),
		codegen.SimpleImport("github.com/jinzhu/gorm"),
		codegen.SimpleImport("github.com/jinzhu/copier"),
		codegen.SimpleImport("database/sql"),
	}

	rbacimports := []*codegen.ImportSpec{
		codegen.SimpleImport(imp),
		codegen.SimpleImport("github.com/mikespook/gorbac"),
	}

	rbactitle := fmt.Sprintf("%s: RBAC", api.Name)
	_, dorbac := api.Metadata["github.com/bketelsen/gorma#rbac"]

	err = api.IterateUserTypes(func(res *design.UserTypeDefinition) error {
		if res.Type.IsObject() {
			title := fmt.Sprintf("%s: Models", api.Name)
			filename := filepath.Join(ModelDir(), res.Type.Name()+"_model.go")
			mtw, err := NewModelWriter(filename)
			if err != nil {
				panic(err)
			}
			mtw.WriteHeader(title, "models", imports)
			if md, ok := res.Metadata["github.com/bketelsen/gorma"]; ok && md == "Model" {
				fmt.Println("Found Gorma Metadata:", md)
				err = mtw.Execute(res)
				if err != nil {
					g.Cleanup()
					return err
				}
			}
			if err := mtw.FormatCode(); err != nil {
				g.Cleanup()

			}
			if err != nil {
				g.genfiles = append(g.genfiles, filename)
			}
			return nil
		}

		return nil

	})
	if dorbac {
		rbacfilename := filepath.Join(ModelDir(), "rbac.go")
		rbacw, err := NewRbacWriter(rbacfilename)
		if err != nil {
			panic(err)
		}
		rbacw.WriteHeader(rbactitle, "models", rbacimports)
		err = rbacw.Execute(api)
		if err != nil {
			g.Cleanup()
			return g.genfiles, err
		}
		if err := rbacw.FormatCode(); err != nil {
			g.Cleanup()
			return nil, err
		}
		if err != nil {
			g.genfiles = append(g.genfiles, rbacfilename)
		}

	}

	return g.genfiles, err
}

// Cleanup removes all the files generated by this generator during the last invokation of Generate.
func (g *Generator) Cleanup() {
	for _, f := range g.genfiles {
		os.Remove(f)
	}
	g.genfiles = nil
}
<commit_msg>trying it<commit_after>package gorma

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/raphael/goa/design"
	"github.com/raphael/goa/goagen/codegen"
	"gopkg.in/alecthomas/kingpin.v2"
)

// Generator is the application code generator.
type Generator struct {
	genfiles []string
}

// Generate is the generator entry point called by the meta generator.
func Generate(api *design.APIDefinition) ([]string, error) {
	g, err := NewGenerator()
	if err != nil {
		return nil, err
	}
	return g.Generate(api)
}

// NewGenerator returns the application code generator.
func NewGenerator() (*Generator, error) {
	return new(Generator), nil
}

// Generate produces the skeleton main.
func (g *Generator) Generate(api *design.APIDefinition) ([]string, error) {

	os.RemoveAll(ModelDir())
	os.MkdirAll(ModelDir(), 0755)
	app := kingpin.New("Model generator", "model generator")
	codegen.RegisterFlags(app)
	_, err := app.Parse(os.Args[1:])
	if err != nil {
		panic(err)
	}
	imp, err := filepath.Rel(filepath.Join(os.Getenv("GOPATH"), "src"), codegen.OutputDir)
	if err != nil {
		return nil, err
	}
	imp = filepath.Join(imp, "app")
	imports := []*codegen.ImportSpec{
		codegen.SimpleImport(imp),
		codegen.SimpleImport("github.com/jinzhu/gorm"),
		codegen.SimpleImport("github.com/jinzhu/copier"),
		codegen.SimpleImport("database/sql"),
	}

	rbacimports := []*codegen.ImportSpec{
		codegen.SimpleImport(imp),
		codegen.SimpleImport("github.com/mikespook/gorbac"),
	}

	rbactitle := fmt.Sprintf("%s: RBAC", api.Name)
	_, dorbac := api.Metadata["github.com/bketelsen/gorma#rbac"]

	err = api.IterateUserTypes(func(res *design.UserTypeDefinition) error {
		if res.Type.IsObject() {
			title := fmt.Sprintf("%s: Models", api.Name)
			filename := filepath.Join(ModelDir(), res.TypeName()+"_model.go")
			mtw, err := NewModelWriter(filename)
			if err != nil {
				panic(err)
			}
			mtw.WriteHeader(title, "models", imports)
			if md, ok := res.Metadata["github.com/bketelsen/gorma"]; ok && md == "Model" {
				fmt.Println("Found Gorma Metadata:", md)
				err = mtw.Execute(res)
				if err != nil {
					g.Cleanup()
					return err
				}
			}
			if err := mtw.FormatCode(); err != nil {
				g.Cleanup()

			}
			if err != nil {
				g.genfiles = append(g.genfiles, filename)
			}
			return nil
		}

		return nil

	})
	if dorbac {
		rbacfilename := filepath.Join(ModelDir(), "rbac.go")
		rbacw, err := NewRbacWriter(rbacfilename)
		if err != nil {
			panic(err)
		}
		rbacw.WriteHeader(rbactitle, "models", rbacimports)
		err = rbacw.Execute(api)
		if err != nil {
			g.Cleanup()
			return g.genfiles, err
		}
		if err := rbacw.FormatCode(); err != nil {
			g.Cleanup()
			return nil, err
		}
		if err != nil {
			g.genfiles = append(g.genfiles, rbacfilename)
		}

	}

	return g.genfiles, err
}

// Cleanup removes all the files generated by this generator during the last invokation of Generate.
func (g *Generator) Cleanup() {
	for _, f := range g.genfiles {
		os.Remove(f)
	}
	g.genfiles = nil
}
<|endoftext|>