<commit_before>package consul

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/hashicorp/consul/api"
)

var upstreamNameRegex = regexp.MustCompile("[^\\w]+")

func NginxTemplateContextCommand() cli.Command {
	return cli.Command{
		Name:  "nginx-template-context",
		Usage: "Collect and return data for consul-template",
		Action: func(c *cli.Context) error {
			consul, _ := api.NewClient(api.DefaultConfig())

			upstreams := make(map[string][]map[string]interface{})
			servers := make(map[string]map[string]map[string]string)

			allRoutes, _, err := consul.KV().List("services/routes/", nil)
			if err != nil {
				return err
			}

			for _, kv := range allRoutes {
				name := strings.TrimPrefix(kv.Key, "services/routes/")
				upstream := upstreamNameRegex.ReplaceAllString("serve_"+name, "_")

				routes := make([]map[string]string, 0)
				if err := json.Unmarshal(kv.Value, &routes); err != nil {
					return err
				}

				instances, _, err := consul.Health().Service(name, "", true, nil)
				if err != nil {
					panic(err)
				}

				if len(instances) == 0 {
					break
				}

				if _, ok := upstreams[upstream]; !ok {
					upstreams[upstream] = make([]map[string]interface{}, 0)
				}

				for _, inst := range instances {
					address := inst.Node.Address
					if inst.Service.Address != "" {
						address = inst.Service.Address
					}

					upstreams[upstream] = append(upstreams[upstream], map[string]interface{}{
						"address": address,
						"port":    inst.Service.Port,
					})
				}

				for _, route := range routes {
					location, ok := route["location"]
					if !ok {
						location = "/"
					}

					stage, ok := route["stage"]
					if !ok {
						stage = "live"
					}

					if _, ok := servers[route["host"]]; !ok {
						servers[route["host"]] = make(map[string]map[string]string, 0)
					}

					if _, ok := servers[route["host"]][location]; !ok {
						servers[route["host"]][location] = make(map[string]string, 0)
						servers[route["host"]][location]["service"] = name // todo: backward compatibility for inn-ci-tools, remove after
					}

					if _, ok := servers[route["host"]][location][stage]; !ok {
						servers[route["host"]][location][stage] = upstream
					}
				}
			}

			out, _ := json.MarshalIndent(map[string]interface{}{
				"upstreams": upstreams,
				"servers":   servers,
			}, "", "  ")

			fmt.Fprintln(os.Stdout, string(out))
			return nil
		},
	}
}
<commit_msg>= nginx-template-context: support non-daemon services<commit_after>package consul

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/hashicorp/consul/api"
)

var upstreamNameRegex = regexp.MustCompile("[^\\w]+")

func NginxTemplateContextCommand() cli.Command {
	return cli.Command{
		Name:  "nginx-template-context",
		Usage: "Collect and return data for consul-template",
		Action: func(c *cli.Context) error {
			conf := api.DefaultConfig()
			conf.Address = "mesos1-q.qa.inn.ru:8500"
			consul, _ := api.NewClient(conf)

			upstreams := make(map[string][]map[string]interface{})
			services := make(map[string]map[string]map[string]map[string]string)

			allRoutes, _, err := consul.KV().List("services/routes/", nil)
			if err != nil {
				return fmt.Errorf("Error on load routes: %s", err)
			}

			for _, kv := range allRoutes {
				name := strings.TrimPrefix(kv.Key, "services/routes/")
				upstream := ""

				routes := make([]map[string]string, 0)
				if err := json.Unmarshal(kv.Value, &routes); err != nil {
					return fmt.Errorf("Error on parse route json: %s. Serive `%s`, json: %s", err, name, string(kv.Value))
				}

				instances, _, err := consul.Health().Service(name, "", true, nil)
				if err != nil {
					return fmt.Errorf("Error on get service `%s` health: %s", name, err)
				}

				if len(instances) == 0 {
					break
				}

				for _, inst := range instances {
					if inst.Service.Port != 0 {
						upstream = upstreamNameRegex.ReplaceAllString("serve_"+name, "_")

						if _, ok := upstreams[upstream]; !ok {
							upstreams[upstream] = make([]map[string]interface{}, 0)
						}

						address := inst.Node.Address
						if inst.Service.Address != "" {
							address = inst.Service.Address
						}

						upstreams[upstream] = append(upstreams[upstream], map[string]interface{}{
							"address": address,
							"port":    inst.Service.Port,
						})
					}
				}

				for _, route := range routes {
					location, ok := route["location"]
					if !ok {
						location = "/"
					}

					stage, ok := route["stage"]
					if !ok {
						stage = "live"
					}

					if _, ok := services[route["host"]]; !ok {
						services[route["host"]] = make(map[string]map[string]map[string]string, 0)
					}

					if _, ok := services[route["host"]][location]; !ok {
						services[route["host"]][location] = make(map[string]map[string]string, 0)
					}

					if _, ok := services[route["host"]][location][stage]; !ok {
						services[route["host"]][location][stage] = map[string]string{
							"upstream": upstream,
							"service": name,
						}
					}
				}
			}

			out, _ := json.MarshalIndent(map[string]interface{}{
				"upstreams": upstreams,
				"services":  services,
			}, "", "  ")

			fmt.Fprintln(os.Stdout, string(out))
			return nil
		},
	}
}
<|endoftext|>


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

// DNS client: see RFC 1035.
// Has to be linked into package net for Dial.

// TODO(rsc):
//	Check periodically whether /etc/resolv.conf has changed.
//	Could potentially handle many outstanding lookups faster.
//	Could have a small cache.
//	Random UDP source port (net.Dial should do that for us).
//	Random request IDs.

package net

import (
	"bytes"
	"fmt"
	"os"
	"rand"
	"sync"
	"time"
)

// DNSError represents a DNS lookup error.
type DNSError struct {
	Error     string // description of the error
	Name      string // name looked for
	Server    string // server used
	IsTimeout bool
}

func (e *DNSError) String() string {
	if e == nil {
		return "<nil>"
	}
	s := "lookup " + e.Name
	if e.Server != "" {
		s += " on " + e.Server
	}
	s += ": " + e.Error
	return s
}

func (e *DNSError) Timeout() bool   { return e.IsTimeout }
func (e *DNSError) Temporary() bool { return e.IsTimeout }

const noSuchHost = "no such host"

// Send a request on the connection and hope for a reply.
// Up to cfg.attempts attempts.
func exchange(cfg *dnsConfig, c Conn, name string, qtype uint16) (*dnsMsg, os.Error) {
	if len(name) >= 256 {
		return nil, &DNSError{Error: "name too long", Name: name}
	}
	out := new(dnsMsg)
	out.id = uint16(rand.Int()) ^ uint16(time.Nanoseconds())
	out.question = []dnsQuestion{
		{name, qtype, dnsClassINET},
	}
	out.recursion_desired = true
	msg, ok := out.Pack()
	if !ok {
		return nil, &DNSError{Error: "internal error - cannot pack message", Name: name}
	}

	for attempt := 0; attempt < cfg.attempts; attempt++ {
		n, err := c.Write(msg)
		if err != nil {
			return nil, err
		}

		c.SetReadTimeout(int64(cfg.timeout) * 1e9) // nanoseconds

		buf := make([]byte, 2000) // More than enough.
		n, err = c.Read(buf)
		if err != nil {
			if e, ok := err.(Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}
		buf = buf[0:n]
		in := new(dnsMsg)
		if !in.Unpack(buf) || in.id != out.id {
			continue
		}
		return in, nil
	}
	var server string
	if a := c.RemoteAddr(); a != nil {
		server = a.String()
	}
	return nil, &DNSError{Error: "no answer from server", Name: name, Server: server, IsTimeout: true}
}

// Find answer for name in dns message.
// On return, if err == nil, addrs != nil.
func answer(name, server string, dns *dnsMsg, qtype uint16) (addrs []dnsRR, err os.Error) {
	addrs = make([]dnsRR, 0, len(dns.answer))

	if dns.rcode == dnsRcodeNameError && dns.recursion_available {
		return nil, &DNSError{Error: noSuchHost, Name: name}
	}
	if dns.rcode != dnsRcodeSuccess {
		// None of the error codes make sense
		// for the query we sent.  If we didn't get
		// a name error and we didn't get success,
		// the server is behaving incorrectly.
		return nil, &DNSError{Error: "server misbehaving", Name: name, Server: server}
	}

	// Look for the name.
	// Presotto says it's okay to assume that servers listed in
	// /etc/resolv.conf are recursive resolvers.
	// We asked for recursion, so it should have included
	// all the answers we need in this one packet.
Cname:
	for cnameloop := 0; cnameloop < 10; cnameloop++ {
		addrs = addrs[0:0]
		for i := 0; i < len(dns.answer); i++ {
			rr := dns.answer[i]
			h := rr.Header()
			if h.Class == dnsClassINET && h.Name == name {
				switch h.Rrtype {
				case qtype:
					n := len(addrs)
					addrs = addrs[0 : n+1]
					addrs[n] = rr
				case dnsTypeCNAME:
					// redirect to cname
					name = rr.(*dnsRR_CNAME).Cname
					continue Cname
				}
			}
		}
		if len(addrs) == 0 {
			return nil, &DNSError{Error: noSuchHost, Name: name, Server: server}
		}
		return addrs, nil
	}

	return nil, &DNSError{Error: "too many redirects", Name: name, Server: server}
}

// Do a lookup for a single name, which must be rooted
// (otherwise answer will not find the answers).
func tryOneName(cfg *dnsConfig, name string, qtype uint16) (addrs []dnsRR, err os.Error) {
	if len(cfg.servers) == 0 {
		return nil, &DNSError{Error: "no DNS servers", Name: name}
	}
	for i := 0; i < len(cfg.servers); i++ {
		// Calling Dial here is scary -- we have to be sure
		// not to dial a name that will require a DNS lookup,
		// or Dial will call back here to translate it.
		// The DNS config parser has already checked that
		// all the cfg.servers[i] are IP addresses, which
		// Dial will use without a DNS lookup.
		server := cfg.servers[i] + ":53"
		c, cerr := Dial("udp", "", server)
		if cerr != nil {
			err = cerr
			continue
		}
		msg, merr := exchange(cfg, c, name, qtype)
		c.Close()
		if merr != nil {
			err = merr
			continue
		}
		addrs, err = answer(name, server, msg, qtype)
		if err == nil || err.(*DNSError).Error == noSuchHost {
			break
		}
	}
	return
}

func convertRR_A(records []dnsRR) []string {
	addrs := make([]string, len(records))
	for i := 0; i < len(records); i++ {
		rr := records[i]
		a := rr.(*dnsRR_A).A
		addrs[i] = IPv4(byte(a>>24), byte(a>>16), byte(a>>8), byte(a)).String()
	}
	return addrs
}

var cfg *dnsConfig
var dnserr os.Error

func loadConfig() { cfg, dnserr = dnsReadConfig() }

func isDomainName(s string) bool {
	// See RFC 1035, RFC 3696.
	if len(s) == 0 {
		return false
	}
	if len(s) > 255 {
		return false
	}
	if s[len(s)-1] != '.' { // simplify checking loop: make name end in dot
		s += "."
	}

	last := byte('.')
	ok := false // ok once we've seen a letter
	partlen := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		default:
			return false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_':
			ok = true
			partlen++
		case '0' <= c && c <= '9':
			// fine
			partlen++
		case c == '-':
			// byte before dash cannot be dot
			if last == '.' {
				return false
			}
			partlen++
		case c == '.':
			// byte before dot cannot be dot, dash
			if last == '.' || last == '-' {
				return false
			}
			if partlen > 63 || partlen == 0 {
				return false
			}
			partlen = 0
		}
		last = c
	}

	return ok
}

var onceLoadConfig sync.Once

func lookup(name string, qtype uint16) (cname string, addrs []dnsRR, err os.Error) {
	if !isDomainName(name) {
		return name, nil, &DNSError{Error: "invalid domain name", Name: name}
	}
	onceLoadConfig.Do(loadConfig)
	if dnserr != nil || cfg == nil {
		err = dnserr
		return
	}
	// If name is rooted (trailing dot) or has enough dots,
	// try it by itself first.
	rooted := len(name) > 0 && name[len(name)-1] == '.'
	if rooted || count(name, '.') >= cfg.ndots {
		rname := name
		if !rooted {
			rname += "."
		}
		// Can try as ordinary name.
		addrs, err = tryOneName(cfg, rname, qtype)
		if err == nil {
			cname = rname
			return
		}
	}
	if rooted {
		return
	}

	// Otherwise, try suffixes.
	for i := 0; i < len(cfg.search); i++ {
		rname := name + "." + cfg.search[i]
		if rname[len(rname)-1] != '.' {
			rname += "."
		}
		addrs, err = tryOneName(cfg, rname, qtype)
		if err == nil {
			cname = rname
			return
		}
	}

	// Last ditch effort: try unsuffixed.
	rname := name
	if !rooted {
		rname += "."
	}
	addrs, err = tryOneName(cfg, rname, qtype)
	if err == nil {
		cname = rname
		return
	}
	return
}

// LookupHost looks for name using the local hosts file and DNS resolver.
// It returns the canonical name for the host and an array of that
// host's addresses.
func LookupHost(name string) (cname string, addrs []string, err os.Error) {
	onceLoadConfig.Do(loadConfig)
	if dnserr != nil || cfg == nil {
		err = dnserr
		return
	}
	// Use entries from /etc/hosts if they match.
	addrs = lookupStaticHost(name)
	if len(addrs) > 0 {
		cname = name
		return
	}
	var records []dnsRR
	cname, records, err = lookup(name, dnsTypeA)
	if err != nil {
		return
	}
	addrs = convertRR_A(records)
	return
}

type SRV struct {
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}

// LookupSRV tries to resolve an SRV query of the given service,
// protocol, and domain name, as specified in RFC 2782. In most cases
// the proto argument can be the same as the corresponding
// Addr.Network().
func LookupSRV(service, proto, name string) (cname string, addrs []*SRV, err os.Error) {
	target := "_" + service + "._" + proto + "." + name
	var records []dnsRR
	cname, records, err = lookup(target, dnsTypeSRV)
	if err != nil {
		return
	}
	addrs = make([]*SRV, len(records))
	for i := 0; i < len(records); i++ {
		r := records[i].(*dnsRR_SRV)
		addrs[i] = &SRV{r.Target, r.Port, r.Priority, r.Weight}
	}
	return
}

type MX struct {
	Host string
	Pref uint16
}

func LookupMX(name string) (entries []*MX, err os.Error) {
	var records []dnsRR
	_, records, err = lookup(name, dnsTypeMX)
	if err != nil {
		return
	}
	entries = make([]*MX, len(records))
	for i := range records {
		r := records[i].(*dnsRR_MX)
		entries[i] = &MX{r.Mx, r.Pref}
	}
	return
}

// reverseaddr returns the in-addr.arpa. or ip6.arpa. hostname of the IP
// address addr suitable for rDNS (PTR) record lookup or an error if it fails
// to parse the IP address.
func reverseaddr(addr string) (arpa string, err os.Error) {
	ip := ParseIP(addr)
	if ip == nil {
		return "", &DNSError{Error: "unrecognized address", Name: addr}
	}
	if ip.To4() != nil {
		return fmt.Sprintf("%d.%d.%d.%d.in-addr.arpa.", ip[15], ip[14], ip[13], ip[12]), nil
	}
	// Must be IPv6
	var buf bytes.Buffer
	// Add it, in reverse, to the buffer
	for i := len(ip) - 1; i >= 0; i-- {
		s := fmt.Sprintf("%02x", ip[i])
		buf.WriteByte(s[1])
		buf.WriteByte('.')
		buf.WriteByte(s[0])
		buf.WriteByte('.')
	}
	// Append "ip6.arpa." and return (buf already has the final .)
	return buf.String() + "ip6.arpa.", nil
}

// LookupAddr performs a reverse lookup for the given address, returning a list
// of names mapping to that address.
func LookupAddr(addr string) (name []string, err os.Error) {
	name = lookupStaticAddr(addr)
	if len(name) > 0 {
		return
	}
	var arpa string
	arpa, err = reverseaddr(addr)
	if err != nil {
		return
	}
	var records []dnsRR
	_, records, err = lookup(arpa, dnsTypePTR)
	if err != nil {
		return
	}
	name = make([]string, len(records))
	for i := range records {
		r := records[i].(*dnsRR_PTR)
		name[i] = r.Ptr
	}
	return
}
<commit_msg>net: return cname in LookupHost<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS client: see RFC 1035.
// Has to be linked into package net for Dial.

// TODO(rsc):
//	Check periodically whether /etc/resolv.conf has changed.
//	Could potentially handle many outstanding lookups faster.
//	Could have a small cache.
//	Random UDP source port (net.Dial should do that for us).
//	Random request IDs.

package net

import (
	"bytes"
	"fmt"
	"os"
	"rand"
	"sync"
	"time"
)

// DNSError represents a DNS lookup error.
type DNSError struct {
	Error     string // description of the error
	Name      string // name looked for
	Server    string // server used
	IsTimeout bool
}

func (e *DNSError) String() string {
	if e == nil {
		return "<nil>"
	}
	s := "lookup " + e.Name
	if e.Server != "" {
		s += " on " + e.Server
	}
	s += ": " + e.Error
	return s
}

func (e *DNSError) Timeout() bool   { return e.IsTimeout }
func (e *DNSError) Temporary() bool { return e.IsTimeout }

const noSuchHost = "no such host"

// Send a request on the connection and hope for a reply.
// Up to cfg.attempts attempts.
func exchange(cfg *dnsConfig, c Conn, name string, qtype uint16) (*dnsMsg, os.Error) {
	if len(name) >= 256 {
		return nil, &DNSError{Error: "name too long", Name: name}
	}
	out := new(dnsMsg)
	out.id = uint16(rand.Int()) ^ uint16(time.Nanoseconds())
	out.question = []dnsQuestion{
		{name, qtype, dnsClassINET},
	}
	out.recursion_desired = true
	msg, ok := out.Pack()
	if !ok {
		return nil, &DNSError{Error: "internal error - cannot pack message", Name: name}
	}

	for attempt := 0; attempt < cfg.attempts; attempt++ {
		n, err := c.Write(msg)
		if err != nil {
			return nil, err
		}

		c.SetReadTimeout(int64(cfg.timeout) * 1e9) // nanoseconds

		buf := make([]byte, 2000) // More than enough.
		n, err = c.Read(buf)
		if err != nil {
			if e, ok := err.(Error); ok && e.Timeout() {
				continue
			}
			return nil, err
		}
		buf = buf[0:n]
		in := new(dnsMsg)
		if !in.Unpack(buf) || in.id != out.id {
			continue
		}
		return in, nil
	}
	var server string
	if a := c.RemoteAddr(); a != nil {
		server = a.String()
	}
	return nil, &DNSError{Error: "no answer from server", Name: name, Server: server, IsTimeout: true}
}

// Find answer for name in dns message.
// On return, if err == nil, addrs != nil.
func answer(name, server string, dns *dnsMsg, qtype uint16) (cname string, addrs []dnsRR, err os.Error) {
	addrs = make([]dnsRR, 0, len(dns.answer))

	if dns.rcode == dnsRcodeNameError && dns.recursion_available {
		return "", nil, &DNSError{Error: noSuchHost, Name: name}
	}
	if dns.rcode != dnsRcodeSuccess {
		// None of the error codes make sense
		// for the query we sent.  If we didn't get
		// a name error and we didn't get success,
		// the server is behaving incorrectly.
		return "", nil, &DNSError{Error: "server misbehaving", Name: name, Server: server}
	}

	// Look for the name.
	// Presotto says it's okay to assume that servers listed in
	// /etc/resolv.conf are recursive resolvers.
	// We asked for recursion, so it should have included
	// all the answers we need in this one packet.
Cname:
	for cnameloop := 0; cnameloop < 10; cnameloop++ {
		addrs = addrs[0:0]
		for i := 0; i < len(dns.answer); i++ {
			rr := dns.answer[i]
			h := rr.Header()
			if h.Class == dnsClassINET && h.Name == name {
				switch h.Rrtype {
				case qtype:
					n := len(addrs)
					addrs = addrs[0 : n+1]
					addrs[n] = rr
				case dnsTypeCNAME:
					// redirect to cname
					name = rr.(*dnsRR_CNAME).Cname
					continue Cname
				}
			}
		}
		if len(addrs) == 0 {
			return "", nil, &DNSError{Error: noSuchHost, Name: name, Server: server}
		}
		return name, addrs, nil
	}

	return "", nil, &DNSError{Error: "too many redirects", Name: name, Server: server}
}

// Do a lookup for a single name, which must be rooted
// (otherwise answer will not find the answers).
func tryOneName(cfg *dnsConfig, name string, qtype uint16) (cname string, addrs []dnsRR, err os.Error) {
	if len(cfg.servers) == 0 {
		return "", nil, &DNSError{Error: "no DNS servers", Name: name}
	}
	for i := 0; i < len(cfg.servers); i++ {
		// Calling Dial here is scary -- we have to be sure
		// not to dial a name that will require a DNS lookup,
		// or Dial will call back here to translate it.
		// The DNS config parser has already checked that
		// all the cfg.servers[i] are IP addresses, which
		// Dial will use without a DNS lookup.
		server := cfg.servers[i] + ":53"
		c, cerr := Dial("udp", "", server)
		if cerr != nil {
			err = cerr
			continue
		}
		msg, merr := exchange(cfg, c, name, qtype)
		c.Close()
		if merr != nil {
			err = merr
			continue
		}
		cname, addrs, err = answer(name, server, msg, qtype)
		if err == nil || err.(*DNSError).Error == noSuchHost {
			break
		}
	}
	return
}

func convertRR_A(records []dnsRR) []string {
	addrs := make([]string, len(records))
	for i := 0; i < len(records); i++ {
		rr := records[i]
		a := rr.(*dnsRR_A).A
		addrs[i] = IPv4(byte(a>>24), byte(a>>16), byte(a>>8), byte(a)).String()
	}
	return addrs
}

var cfg *dnsConfig
var dnserr os.Error

func loadConfig() { cfg, dnserr = dnsReadConfig() }

func isDomainName(s string) bool {
	// See RFC 1035, RFC 3696.
	if len(s) == 0 {
		return false
	}
	if len(s) > 255 {
		return false
	}
	if s[len(s)-1] != '.' { // simplify checking loop: make name end in dot
		s += "."
	}

	last := byte('.')
	ok := false // ok once we've seen a letter
	partlen := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		switch {
		default:
			return false
		case 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' || c == '_':
			ok = true
			partlen++
		case '0' <= c && c <= '9':
			// fine
			partlen++
		case c == '-':
			// byte before dash cannot be dot
			if last == '.' {
				return false
			}
			partlen++
		case c == '.':
			// byte before dot cannot be dot, dash
			if last == '.' || last == '-' {
				return false
			}
			if partlen > 63 || partlen == 0 {
				return false
			}
			partlen = 0
		}
		last = c
	}

	return ok
}

var onceLoadConfig sync.Once

func lookup(name string, qtype uint16) (cname string, addrs []dnsRR, err os.Error) {
	if !isDomainName(name) {
		return name, nil, &DNSError{Error: "invalid domain name", Name: name}
	}
	onceLoadConfig.Do(loadConfig)
	if dnserr != nil || cfg == nil {
		err = dnserr
		return
	}
	// If name is rooted (trailing dot) or has enough dots,
	// try it by itself first.
	rooted := len(name) > 0 && name[len(name)-1] == '.'
	if rooted || count(name, '.') >= cfg.ndots {
		rname := name
		if !rooted {
			rname += "."
		}
		// Can try as ordinary name.
		cname, addrs, err = tryOneName(cfg, rname, qtype)
		if err == nil {
			return
		}
	}
	if rooted {
		return
	}

	// Otherwise, try suffixes.
	for i := 0; i < len(cfg.search); i++ {
		rname := name + "." + cfg.search[i]
		if rname[len(rname)-1] != '.' {
			rname += "."
		}
		cname, addrs, err = tryOneName(cfg, rname, qtype)
		if err == nil {
			return
		}
	}

	// Last ditch effort: try unsuffixed.
	rname := name
	if !rooted {
		rname += "."
	}
	cname, addrs, err = tryOneName(cfg, rname, qtype)
	if err == nil {
		return
	}
	return
}

// LookupHost looks for name using the local hosts file and DNS resolver.
// It returns the canonical name for the host and an array of that
// host's addresses.
func LookupHost(name string) (cname string, addrs []string, err os.Error) {
	onceLoadConfig.Do(loadConfig)
	if dnserr != nil || cfg == nil {
		err = dnserr
		return
	}
	// Use entries from /etc/hosts if they match.
	addrs = lookupStaticHost(name)
	if len(addrs) > 0 {
		cname = name
		return
	}
	var records []dnsRR
	cname, records, err = lookup(name, dnsTypeA)
	if err != nil {
		return
	}
	addrs = convertRR_A(records)
	return
}

type SRV struct {
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}

// LookupSRV tries to resolve an SRV query of the given service,
// protocol, and domain name, as specified in RFC 2782. In most cases
// the proto argument can be the same as the corresponding
// Addr.Network().
func LookupSRV(service, proto, name string) (cname string, addrs []*SRV, err os.Error) {
	target := "_" + service + "._" + proto + "." + name
	var records []dnsRR
	cname, records, err = lookup(target, dnsTypeSRV)
	if err != nil {
		return
	}
	addrs = make([]*SRV, len(records))
	for i := 0; i < len(records); i++ {
		r := records[i].(*dnsRR_SRV)
		addrs[i] = &SRV{r.Target, r.Port, r.Priority, r.Weight}
	}
	return
}

type MX struct {
	Host string
	Pref uint16
}

func LookupMX(name string) (entries []*MX, err os.Error) {
	var records []dnsRR
	_, records, err = lookup(name, dnsTypeMX)
	if err != nil {
		return
	}
	entries = make([]*MX, len(records))
	for i := range records {
		r := records[i].(*dnsRR_MX)
		entries[i] = &MX{r.Mx, r.Pref}
	}
	return
}

// reverseaddr returns the in-addr.arpa. or ip6.arpa. hostname of the IP
// address addr suitable for rDNS (PTR) record lookup or an error if it fails
// to parse the IP address.
func reverseaddr(addr string) (arpa string, err os.Error) {
	ip := ParseIP(addr)
	if ip == nil {
		return "", &DNSError{Error: "unrecognized address", Name: addr}
	}
	if ip.To4() != nil {
		return fmt.Sprintf("%d.%d.%d.%d.in-addr.arpa.", ip[15], ip[14], ip[13], ip[12]), nil
	}
	// Must be IPv6
	var buf bytes.Buffer
	// Add it, in reverse, to the buffer
	for i := len(ip) - 1; i >= 0; i-- {
		s := fmt.Sprintf("%02x", ip[i])
		buf.WriteByte(s[1])
		buf.WriteByte('.')
		buf.WriteByte(s[0])
		buf.WriteByte('.')
	}
	// Append "ip6.arpa." and return (buf already has the final .)
	return buf.String() + "ip6.arpa.", nil
}

// LookupAddr performs a reverse lookup for the given address, returning a list
// of names mapping to that address.
func LookupAddr(addr string) (name []string, err os.Error) {
	name = lookupStaticAddr(addr)
	if len(name) > 0 {
		return
	}
	var arpa string
	arpa, err = reverseaddr(addr)
	if err != nil {
		return
	}
	var records []dnsRR
	_, records, err = lookup(arpa, dnsTypePTR)
	if err != nil {
		return
	}
	name = make([]string, len(records))
	for i := range records {
		r := records[i].(*dnsRR_PTR)
		name[i] = r.Ptr
	}
	return
}
<|endoftext|>


<commit_before>package divaextractor

import "testing"

func TestDocument(t *testing.T) {
	c := NewWikipedia()

	err := c.Do("明日花キララ")
	if err != nil {
		t.Fatal(err)
	}

	if c.Birthday().Unix() != 591753600 {
		t.Errorf("Unexpected Wikipedia.Birthday: %s", c.Birthday())
	}

	if c.Blood() != "A" {
		t.Errorf("Unexpected Wikipedia.Blood: %s", c.Blood())
	}

	if c.HW() != "162  / ―" {
		t.Errorf("Unexpected Wikipedia.H/Weight: %s", c.HW())
	}
	if c.Height() != 162 {
		t.Errorf("Unexpected Wikipedia.Height: %d", c.Height())
	}
	if c.Weight() != 0 {
		t.Errorf("Unexpected Wikipedia.Weight: %d", c.Weight())
	}

	if c.BWH() != "90 - 58 - 85" {
		t.Errorf("Unexpected Wikipedia.BWH: %s", c.BWH())
	}
	if c.Bust() != 90 {
		t.Errorf("Unexpected Wikipedia.Bust: %d", c.Bust())
	}
	if c.Waist() != 58 {
		t.Errorf("Unexpected Wikipedia.Waist: %d", c.Waist())
	}
	if c.Hip() != 85 {
		t.Errorf("Unexpected Wikipedia.Hip: %d", c.Hip())
	}

	if c.Bracup() != "G" {
		t.Errorf("Unexpected Wikipedia.Bracup: %s", c.Bracup())
	}
}
<commit_msg>fixed testing<commit_after>package divaextractor

import "testing"

func TestDocument(t *testing.T) {
	c := NewWikipedia()

	err := c.Do("明日花キララ")
	if err != nil {
		t.Fatal(err)
	}

	if c.Birthday().Unix() != 591753600 {
		t.Errorf("Unexpected Wikipedia.Birthday: %s(%d)", c.Birthday(), c.Birthday().Unix())
	}

	if c.Blood() != "A" {
		t.Errorf("Unexpected Wikipedia.Blood: %s", c.Blood())
	}

	if c.HW() != "162  / ―" {
		t.Errorf("Unexpected Wikipedia.H/Weight: %s", c.HW())
	}
	if c.Height() != 162 {
		t.Errorf("Unexpected Wikipedia.Height: %d", c.Height())
	}
	if c.Weight() != 0 {
		t.Errorf("Unexpected Wikipedia.Weight: %d", c.Weight())
	}

	if c.BWH() != "90 - 58 - 85" {
		t.Errorf("Unexpected Wikipedia.BWH: %s", c.BWH())
	}
	if c.Bust() != 90 {
		t.Errorf("Unexpected Wikipedia.Bust: %d", c.Bust())
	}
	if c.Waist() != 58 {
		t.Errorf("Unexpected Wikipedia.Waist: %d", c.Waist())
	}
	if c.Hip() != 85 {
		t.Errorf("Unexpected Wikipedia.Hip: %d", c.Hip())
	}

	if c.Bracup() != "G" {
		t.Errorf("Unexpected Wikipedia.Bracup: %s", c.Bracup())
	}
}
<|endoftext|>


<commit_before>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2021-12-06 04:39:32.070052 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "ebfe5f0b981aa536a7f07d0d73a70f2ba874c1e2"

// SpartaGitShortHash is the short version of SpartaGitHash
const SpartaGitShortHash = "ebfe5f0"
<commit_msg>"Autogenerated build info"<commit_after>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2021-12-07 06:24:53.806908 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "459c70dbb2e85b6b476d070c1b450b587843719c"

// SpartaGitShortHash is the short version of SpartaGitHash
const SpartaGitShortHash = "459c70d"
<|endoftext|>


<commit_before>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2019-01-20 00:37:06.531216 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "2a7dae8c93556fe4d6f32a32ecc64e953c566b33"
<commit_msg>"Autogenerated build info"<commit_after>package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: 2019-01-20 00:46:35.341786 +0000 UTC

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "35a4e13b44a6ce0752ee3bd8f68f7d7b4b254ddf"
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

package newbase60

import "testing"

func TestEncodeInt(t *testing.T) {
	tests := []struct {
		n int
		s string
	}{
		{0, ""},
		{60, "10"},
		{1152, "KC"},
		{-1, ""},
	}

	for _, tt := range tests {
		if got := EncodeInt(tt.n); got != tt.s {
			t.Errorf("EncodeInt(%d) = %v, want %v", tt.n, got, tt.s)
		}
	}
}

func TestDecodeToInt(t *testing.T) {
	tests := []struct {
		n int
		s string
	}{
		{0, ""},
		{60, "10"},
		{1152, "KC"},
	}

	for _, tt := range tests {
		if got := DecodeToInt(tt.s); got != tt.n {
			t.Errorf("DecodeToInt(%s) = %v, want %v", tt.s, got, tt.n)
		}
	}
}
<commit_msg>add test for decoding invalid chars<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

package newbase60

import "testing"

func TestEncodeInt(t *testing.T) {
	tests := []struct {
		n int
		s string
	}{
		{0, ""},
		{60, "10"},
		{1152, "KC"},
		{-1, ""}, // negative numbers are ignored
	}

	for _, tt := range tests {
		if got := EncodeInt(tt.n); got != tt.s {
			t.Errorf("EncodeInt(%d) = %v, want %v", tt.n, got, tt.s)
		}
	}
}

func TestDecodeToInt(t *testing.T) {
	tests := []struct {
		s string
		n int
	}{
		{"", 0},
		{"10", 60},
		{"KC", 1152},
		{"!/*#", 0},    // invalid chars are ignored
		{".*KC", 1152}, // invalid chars at beginning have no effect
	}

	for _, tt := range tests {
		if got := DecodeToInt(tt.s); got != tt.n {
			t.Errorf("DecodeToInt(%s) = %v, want %v", tt.s, got, tt.n)
		}
	}
}
<|endoftext|>


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

package test

import (
	"fmt"
	"testing"
	"time"

	. "github.com/prometheus/alertmanager/test"
)

// This file contains acceptance tests around the basic sending logic
// for notifications, which includes batching and ensuring that each
// notification is eventually sent at least once and ideally exactly
// once.

func TestMergeAlerts(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 1s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co := at.Collector("webhook")
	wh := NewWebhook(co)

	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	am.Push(At(1), Alert("alertname", "test").Active(1))
	am.Push(At(1.2), Alert("alertname", "test").Active(1.1))

	co.Want(Between(2, 2.5), Alert("alertname", "test").Active(1))

	am.Push(At(2.1), Alert("alertname", "test").Annotate("ann", "v1").Active(2))

	co.Want(Between(3, 3.5), Alert("alertname", "test").Annotate("ann", "v1").Active(1))

	am.Push(At(3.6), Alert("alertname", "test").Annotate("ann", "v2").Active(1.5))

	co.Want(Between(4, 4.5), Alert("alertname", "test").Annotate("ann", "v2").Active(1))

	at.Run()
}

func TestRepeat(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 1s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
`

	// Create a new acceptance test that instantiates new Alertmanagers
	// with the given configuration and verifies times with the given
	// tollerance.
	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	// Create a collector to which alerts can be written and verified
	// against a set of expected alert notifications.
	co := at.Collector("webhook")
	// Run something that satisfies the webhook interface to which the
	// Alertmanager pushes as defined by its configuration.
	wh := NewWebhook(co)

	// Create a new Alertmanager process listening to a random port
	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	// Declare pushes to be made to the Alertmanager at the given time.
	// Times are provided in fractions of seconds.
	am.Push(At(1), Alert("alertname", "test").Active(1))

	at.Do(At(1.2), func() {
		am.Terminate()
		am.Start()
	})
	am.Push(At(3.5), Alert("alertname", "test").Active(1, 3))

	// Declare which alerts are expected to arrive at the collector within
	// the defined time intervals.
	co.Want(Between(2, 2.5), Alert("alertname", "test").Active(1))
	co.Want(Between(3, 3.5), Alert("alertname", "test").Active(1))
	co.Want(Between(4, 4.5), Alert("alertname", "test").Active(1, 3))

	// Start the flow as defined above and run the checks afterwards.
	at.Run()
}

func TestRetry(t *testing.T) {
	t.Parallel()

	// We create a notification config that fans out into two different
	// webhooks.
	// The succeeding one must still only receive the first successful
	// notifications. Sending to the succeeding one must eventually succeed.
	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 3s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co1 := at.Collector("webhook")
	wh1 := NewWebhook(co1)

	co2 := at.Collector("webhook_failing")
	wh2 := NewWebhook(co2)

	wh2.Func = func(ts float64) bool {
		// Fail the first two interval periods but eventually
		// succeed in the third interval after a few failed attempts.
		return ts < 4.5
	}

	am := at.Alertmanager(fmt.Sprintf(conf, wh1.Address(), wh2.Address()))

	am.Push(At(1), Alert("alertname", "test1"))

	co1.Want(Between(2, 2.5), Alert("alertname", "test1").Active(1))
	co1.Want(Between(5, 5.5), Alert("alertname", "test1").Active(1))

	co2.Want(Between(4.5, 5), Alert("alertname", "test1").Active(1))
}

func TestBatching(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 5s

notification_configs:
- name:            "default"
  webhook_configs:
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co := at.Collector("webhook")
	wh := NewWebhook(co)

	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	am.Push(At(1.1), Alert("alertname", "test1").Active(1))
	am.Push(At(1.9), Alert("alertname", "test5").Active(1))
	am.Push(At(2.3),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

	co.Want(Between(2.0, 2.5),
		Alert("alertname", "test1").Active(1),
		Alert("alertname", "test5").Active(1),
	)
	// Only expect the new ones with the next group interval.
	co.Want(Between(3, 3.5),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

	// While no changes happen expect no additional notifications
	// until the 5s repeat interval has ended.

	// The last three notifications should sent with the first two even
	// though their repeat interval has not yet passed. This way fragmented
	// batches are unified and notification noise reduced.
	co.Want(Between(7, 7.5),
		Alert("alertname", "test1").Active(1),
		Alert("alertname", "test5").Active(1),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

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

package test

import (
	"fmt"
	"testing"
	"time"

	. "github.com/prometheus/alertmanager/test"
)

// This file contains acceptance tests around the basic sending logic
// for notifications, which includes batching and ensuring that each
// notification is eventually sent at least once and ideally exactly
// once.

func TestMergeAlerts(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 1s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co := at.Collector("webhook")
	wh := NewWebhook(co)

	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	// Refresh an alert several times. The starting time must remain at the earliest
	// point in time.
	am.Push(At(1), Alert("alertname", "test").Active(1.1))
	// Another Prometheus server might be sending later but with an earlier start time.
	am.Push(At(1.2), Alert("alertname", "test").Active(1))

	co.Want(Between(2, 2.5), Alert("alertname", "test").Active(1))

	am.Push(At(2.1), Alert("alertname", "test").Annotate("ann", "v1").Active(2))

	co.Want(Between(3, 3.5), Alert("alertname", "test").Annotate("ann", "v1").Active(1))

	// Annotations are always overwritten by the alert that arrived most recently.
	am.Push(At(3.6), Alert("alertname", "test").Annotate("ann", "v2").Active(1.5))

	co.Want(Between(4, 4.5), Alert("alertname", "test").Annotate("ann", "v2").Active(1))

	// If an alert is marked resolved twice, the latest point in time must be
	// set as the eventual resolve time.
	am.Push(At(4.6), Alert("alertname", "test").Annotate("ann", "v2").Active(3, 4.5))
	am.Push(At(4.8), Alert("alertname", "test").Annotate("ann", "v3").Active(2.9, 4.8))
	am.Push(At(4.8), Alert("alertname", "test").Annotate("ann", "v3").Active(2.9, 4.1))

	co.Want(Between(5, 5.5), Alert("alertname", "test").Annotate("ann", "v3").Active(1, 4.8))

	// Reactivate an alert after a previous occurrence has been resolved.
	// No overlap, no merge must occur.
	am.Push(At(5.3), Alert("alertname", "test"))

	co.Want(Between(6, 6.5), Alert("alertname", "test").Active(5.3))

	// Test against a bug which ocurrec after a restart. The previous occurrence of
	// the alert was sent rather than the most recent one.
	at.Do(At(6.7), func() {
		am.Terminate()
		am.Start()
	})

	// On restart the alert is flushed right away as the group_wait has already passed.
	// However, it must be caught in the deduplication stage.
	// The next attempt will be 1s later and won't be filtered in deduping.
	co.Want(Between(7.7, 8), Alert("alertname", "test").Active(5.3))

	at.Run()
}

func TestRepeat(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 1s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
`

	// Create a new acceptance test that instantiates new Alertmanagers
	// with the given configuration and verifies times with the given
	// tollerance.
	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	// Create a collector to which alerts can be written and verified
	// against a set of expected alert notifications.
	co := at.Collector("webhook")
	// Run something that satisfies the webhook interface to which the
	// Alertmanager pushes as defined by its configuration.
	wh := NewWebhook(co)

	// Create a new Alertmanager process listening to a random port
	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	// Declare pushes to be made to the Alertmanager at the given time.
	// Times are provided in fractions of seconds.
	am.Push(At(1), Alert("alertname", "test").Active(1))

	at.Do(At(1.2), func() {
		am.Terminate()
		am.Start()
	})
	am.Push(At(3.5), Alert("alertname", "test").Active(1, 3))

	// Declare which alerts are expected to arrive at the collector within
	// the defined time intervals.
	co.Want(Between(2, 2.5), Alert("alertname", "test").Active(1))
	co.Want(Between(3, 3.5), Alert("alertname", "test").Active(1))
	co.Want(Between(4, 4.5), Alert("alertname", "test").Active(1, 3))

	// Start the flow as defined above and run the checks afterwards.
	at.Run()
}

func TestRetry(t *testing.T) {
	t.Parallel()

	// We create a notification config that fans out into two different
	// webhooks.
	// The succeeding one must still only receive the first successful
	// notifications. Sending to the succeeding one must eventually succeed.
	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 3s

notification_configs:
- name: "default"
  webhook_configs:
  - url: 'http://%s'
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co1 := at.Collector("webhook")
	wh1 := NewWebhook(co1)

	co2 := at.Collector("webhook_failing")
	wh2 := NewWebhook(co2)

	wh2.Func = func(ts float64) bool {
		// Fail the first two interval periods but eventually
		// succeed in the third interval after a few failed attempts.
		return ts < 4.5
	}

	am := at.Alertmanager(fmt.Sprintf(conf, wh1.Address(), wh2.Address()))

	am.Push(At(1), Alert("alertname", "test1"))

	co1.Want(Between(2, 2.5), Alert("alertname", "test1").Active(1))
	co1.Want(Between(5, 5.5), Alert("alertname", "test1").Active(1))

	co2.Want(Between(4.5, 5), Alert("alertname", "test1").Active(1))
}

func TestBatching(t *testing.T) {
	t.Parallel()

	conf := `
route:
  send_to: "default"
  group_by: []
  group_wait:      1s
  group_interval:  1s
  repeat_interval: 5s

notification_configs:
- name:            "default"
  webhook_configs:
  - url: 'http://%s'
`

	at := NewAcceptanceTest(t, &AcceptanceOpts{
		Tolerance: 150 * time.Millisecond,
	})

	co := at.Collector("webhook")
	wh := NewWebhook(co)

	am := at.Alertmanager(fmt.Sprintf(conf, wh.Address()))

	am.Push(At(1.1), Alert("alertname", "test1").Active(1))
	am.Push(At(1.9), Alert("alertname", "test5").Active(1))
	am.Push(At(2.3),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

	co.Want(Between(2.0, 2.5),
		Alert("alertname", "test1").Active(1),
		Alert("alertname", "test5").Active(1),
	)
	// Only expect the new ones with the next group interval.
	co.Want(Between(3, 3.5),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

	// While no changes happen expect no additional notifications
	// until the 5s repeat interval has ended.

	// The last three notifications should sent with the first two even
	// though their repeat interval has not yet passed. This way fragmented
	// batches are unified and notification noise reduced.
	co.Want(Between(7, 7.5),
		Alert("alertname", "test1").Active(1),
		Alert("alertname", "test5").Active(1),
		Alert("alertname", "test2").Active(1.5),
		Alert("alertname", "test3").Active(1.5),
		Alert("alertname", "test4").Active(1.6),
	)

	at.Run()
}
<|endoftext|>


<commit_before>package lib

import "testing"
import "github.com/stretchr/testify/assert"

func TestDependencyLinks(t *testing.T) {
	a := newApplicationMetadata("app-a", "a", &Spec{Name: "app-a", Dependencies: []string{"app-b"}})
	b := newApplicationMetadata("app-b", "b", &Spec{Name: "app-b", Dependencies: []string{"app-c"}})
	c := newApplicationMetadata("app-c", "c", &Spec{Name: "app-c"})

	s := applicationMetadataSet{a, b, c}
	apps, err := s.toApplications(true)
	check(t, err)
	m := apps.indexByName()

	assert.Len(t, apps, 3)
	assert.Equal(t, m["app-b"], m["app-a"].Requires()[0])
	assert.Equal(t, m["app-c"], m["app-b"].Requires()[0])
	assert.Len(t, m["app-c"].Requires(), 0)
	assert.Equal(t, m["app-b"], m["app-c"].RequiredBy()[0])
	assert.Equal(t, m["app-a"], m["app-b"].RequiredBy()[0])
}
<commit_msg>Tests<commit_after>package lib

import "testing"
import "github.com/stretchr/testify/assert"

func TestDependencyLinks(t *testing.T) {
	a := newApplicationMetadata("app-a", "a", &Spec{Name: "app-a", Dependencies: []string{"app-b"}})
	b := newApplicationMetadata("app-b", "b", &Spec{Name: "app-b", Dependencies: []string{"app-c"}})
	c := newApplicationMetadata("app-c", "c", &Spec{Name: "app-c"})

	s := applicationMetadataSet{a, b, c}
	apps, err := s.toApplications(true)
	check(t, err)
	m := apps.indexByName()

	assert.Len(t, apps, 3)
	assert.Equal(t, m["app-b"], m["app-a"].Requires()[0])
	assert.Equal(t, m["app-c"], m["app-b"].Requires()[0])
	assert.Len(t, m["app-c"].Requires(), 0)
	assert.Equal(t, m["app-b"], m["app-c"].RequiredBy()[0])
	assert.Equal(t, m["app-a"], m["app-b"].RequiredBy()[0])
}

func TestVersionCalculation(t *testing.T) {
	a := newApplicationMetadata("app-a", "a", &Spec{Name: "app-a", Dependencies: []string{"app-b"}})
	b := newApplicationMetadata("app-b", "b", &Spec{Name: "app-b"})

	s := applicationMetadataSet{a, b}
	apps, err := s.toApplications(true)
	check(t, err)
	m := apps.indexByName()

	assert.Equal(t, "b", m["app-b"].Version())
	assert.Equal(t, "da23614e02469a0d7c7bd1bdab5c9c474b1904dc", m["app-a"].Version())
}
<|endoftext|>


<commit_before>// +build linux

package vdlog

import (
	"fmt"
	"os/exec"
)

// Mapping between journalctl and vdlog levels

//    journald            vdlog
//    0: emerg            0:LevelError
//    1: alert            0:LevelError
//    2: crit             0:LevelError
//    3: err              0:LevelError
//    4: warning          1:LevelWarn
//    5: notice           2:LevelInfo
//    6: info             3:LevelVerbose
//    7: debug            4:LevelDebug

func createJournaldSink(host string, port int, tcp, local bool) func(e Event) error {
	return func(e Event) error {
		var journaldPriority int
		switch e.Level {
		case LevelError:
			journaldPriority = 3
			break
		case LevelWarn:
			journaldPriority = 4
			break
		case LevelInfo:
			journaldPriority = 5
			break
		case LevelVerbose:
			journaldPriority = 6
			break
		case LevelDebug:
			journaldPriority = 7
			break
		default:
			journaldPriority = 5
		}
		var cmd *exec.Cmd
		if local {
			cmd = exec.Command("/bin/logger", "--journald")
		} else {
			var protoParams string
			if tcp {
				protoParams = "--tcp"
			} else {
				protoParams = "--udp"
			}
			cmd = exec.Command("/bin/logger",
				"--journald",
				fmt.Sprintf("--server %s", host),
				fmt.Sprintf("--port %v", port),
				protoParams,
			)
		}
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return err
		}
		//http://man7.org/linux/man-pages/man7/systemd.journal-fields.7.html
		fmt.Fprintf(stdin, "MESSAGE_ID=%s\n", e.UUID)
		fmt.Fprintf(stdin, "CODE_FILE=%s\n", e.Filename)
		fmt.Fprintf(stdin, "CODE_LINE=%v\n", e.Line)
		fmt.Fprintf(stdin, "SYSLOG_IDENTIFIER=%s\n", e.Type)
		fmt.Fprintf(stdin, "MESSAGE=%q\n", fmt.Sprint(e.Metadata))
		fmt.Fprintf(stdin, "PRIORITY=%d\n", journaldPriority)
		fmt.Fprintf(stdin, "SYSLOG_PID=%d\n", e.Pid)
		fmt.Fprintf(stdin, "OBJECT_PID=%d\n", e.Pid)
		fmt.Fprintf(stdin, "VDLOG_LEVEL=%s\n", e.LevelString)
		fmt.Fprintf(stdin, "VDLOG_CALLED=%s\n", e.Called)
		for k,v := range e.Metadata {
			fmt.Fprintf(stdin, "METADATA_%q=%q\n", k, fmt.Sprint(v))
		}
		err = stdin.Close()
		if err != nil {
			return err
		}
		err = cmd.Run()
		if err != nil {
			return err
		}
		return nil
	}
}

//LogToLocalJournald allows to send messages directly to journald daemon via logger command
func LogToLocalJournald() {
	AddSink("journald local", createJournaldSink("localhost", 514, true, true))
}

//LogToRemoteJournaldViaTCP allows to send messages directly to remote journald daemon via logger command and TCP
func LogToRemoteJournaldViaTCP(server string, port int) {
	AddSink(fmt.Sprintf("journald tcp(%s:%v)", server, port), createJournaldSink(server, port, true, false))
}

//LogToRemoteJournaldViaUDP allows to send messages directly to remote journald daemon via logger command and UDP
func LogToRemoteJournaldViaUDP(server string, port int) {
	AddSink(fmt.Sprintf("journald udp(%s:%v)", server, port), createJournaldSink(server, port, false, false))
}
<commit_msg>fix metadata not shown<commit_after>// +build linux

package vdlog

import (
	"fmt"
	"os/exec"
)

// Mapping between journalctl and vdlog levels

//    journald            vdlog
//    0: emerg            0:LevelError
//    1: alert            0:LevelError
//    2: crit             0:LevelError
//    3: err              0:LevelError
//    4: warning          1:LevelWarn
//    5: notice           2:LevelInfo
//    6: info             3:LevelVerbose
//    7: debug            4:LevelDebug

func createJournaldSink(host string, port int, tcp, local bool) func(e Event) error {
	return func(e Event) error {
		var journaldPriority int
		switch e.Level {
		case LevelError:
			journaldPriority = 3
			break
		case LevelWarn:
			journaldPriority = 4
			break
		case LevelInfo:
			journaldPriority = 5
			break
		case LevelVerbose:
			journaldPriority = 6
			break
		case LevelDebug:
			journaldPriority = 7
			break
		default:
			journaldPriority = 5
		}
		var cmd *exec.Cmd
		if local {
			cmd = exec.Command("/bin/logger", "--journald")
		} else {
			var protoParams string
			if tcp {
				protoParams = "--tcp"
			} else {
				protoParams = "--udp"
			}
			cmd = exec.Command("/bin/logger",
				"--journald",
				fmt.Sprintf("--server %s", host),
				fmt.Sprintf("--port %v", port),
				protoParams,
			)
		}
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return err
		}
		//http://man7.org/linux/man-pages/man7/systemd.journal-fields.7.html
		fmt.Fprintf(stdin, "MESSAGE_ID=%s\n", e.UUID)
		fmt.Fprintf(stdin, "CODE_FILE=%s\n", e.Filename)
		fmt.Fprintf(stdin, "CODE_LINE=%v\n", e.Line)
		fmt.Fprintf(stdin, "SYSLOG_IDENTIFIER=%s\n", e.Type)
		fmt.Fprintf(stdin, "MESSAGE=%q\n", fmt.Sprint(e.Metadata))
		fmt.Fprintf(stdin, "PRIORITY=%d\n", journaldPriority)
		fmt.Fprintf(stdin, "SYSLOG_PID=%d\n", e.Pid)
		fmt.Fprintf(stdin, "OBJECT_PID=%d\n", e.Pid)
		fmt.Fprintf(stdin, "VDLOG_LEVEL=%s\n", e.LevelString)
		fmt.Fprintf(stdin, "VDLOG_CALLED=%s\n", e.Called)
		for k,v := range e.Metadata {
			fmt.Fprintf(stdin, "VDLOG_META_%s=%q\n", k, fmt.Sprint(v))
		}
		err = stdin.Close()
		if err != nil {
			return err
		}
		err = cmd.Run()
		if err != nil {
			return err
		}
		return nil
	}
}

//LogToLocalJournald allows to send messages directly to journald daemon via logger command
func LogToLocalJournald() {
	AddSink("journald local", createJournaldSink("localhost", 514, true, true))
}

//LogToRemoteJournaldViaTCP allows to send messages directly to remote journald daemon via logger command and TCP
func LogToRemoteJournaldViaTCP(server string, port int) {
	AddSink(fmt.Sprintf("journald tcp(%s:%v)", server, port), createJournaldSink(server, port, true, false))
}

//LogToRemoteJournaldViaUDP allows to send messages directly to remote journald daemon via logger command and UDP
func LogToRemoteJournaldViaUDP(server string, port int) {
	AddSink(fmt.Sprintf("journald udp(%s:%v)", server, port), createJournaldSink(server, port, false, false))
}
<|endoftext|>


<commit_before>package md

import (
	"flag"
	"fmt"
	"log"

	"github.com/mrwill84/CtpSender/ctp/async"
	"github.com/mrwill84/goctp"
)

var (
	broker_id    = flag.String("BrokerID", "9999", "经纪公司编号,SimNow BrokerID统一为：9999")
	investor_id  = flag.String("InvestorID", "<InvestorID>", "交易用户代码")
	pass_word    = flag.String("Password", "<Password>", "交易用户密码")
	market_front = flag.String("MarketFront", "tcp://180.168.146.187:10031", "行情前置,SimNow的测试环境: tcp://180.168.146.187:10031")
	trade_front  = flag.String("TradeFront", "tcp://180.168.146.187:10030", "交易前置,SimNow的测试环境: tcp://180.168.146.187:10030")
)

var CTP GoCTPClient

type GoCTPClient struct {
	BrokerID   string
	InvestorID string
	Password   string

	MdFront string
	MdApi   goctp.CThostFtdcMdApi

	TraderFront string
	TraderApi   goctp.CThostFtdcTraderApi

	MdRequestID     int
	TraderRequestID int
	FrontendConnent chan bool
}

func (g *GoCTPClient) GetMdRequestID() int {
	g.MdRequestID += 1
	return g.MdRequestID
}

func NewDirectorCThostFtdcMdSpi(v interface{}) goctp.CThostFtdcMdSpi {

	return goctp.NewDirectorCThostFtdcMdSpi(v)

}

type GoCThostFtdcMdSpi struct {
	Client GoCTPClient
}

func (p *GoCThostFtdcMdSpi) OnRspError(pRspInfo goctp.CThostFtdcRspInfoField, nRequestID int, bIsLast bool) {
	log.Println("GoCThostFtdcMdSpi.OnRspError.")
	p.IsErrorRspInfo(pRspInfo)
}

func (p *GoCThostFtdcMdSpi) OnFrontDisconnected(nReason int) {
	log.Printf("GoCThostFtdcMdSpi.OnFrontDisconnected: %#v\n", nReason)
}

func (p *GoCThostFtdcMdSpi) OnHeartBeatWarning(nTimeLapse int) {
	log.Printf("GoCThostFtdcMdSpi.OnHeartBeatWarning: %v", nTimeLapse)

}

func (p *GoCThostFtdcMdSpi) OnFrontConnected() {
	log.Println("GoCThostFtdcMdSpi.OnFrontConnected.")
	if CTP.FrontendConnent != nil {
		CTP.FrontendConnent <- true
		return
	}
	panic("never")
	//p.ReqUserLogin()
}

func (p *GoCThostFtdcMdSpi) IsErrorRspInfo(pRspInfo goctp.CThostFtdcRspInfoField) bool {
	// 如果ErrorID != 0, 说明收到了错误的响应
	bResult := (pRspInfo.GetErrorID() != 0)
	if bResult {
		log.Printf("ErrorID=%v ErrorMsg=%v\n", pRspInfo.GetErrorID(), pRspInfo.GetErrorMsg())
	}
	return bResult
}

func (p *GoCThostFtdcMdSpi) ReqUserLogin() (interface{}, error) {
	id := p.Client.GetMdRequestID()
	obj, err := async.Alloc(id, func() {
		log.Println("GoCThostFtdcMdSpi.ReqUserLogin.")
		req := goctp.NewCThostFtdcReqUserLoginField()
		req.SetBrokerID(p.Client.BrokerID)
		req.SetUserID(p.Client.InvestorID)
		req.SetPassword(p.Client.Password)
		requestID := id
		iResult := p.Client.MdApi.ReqUserLogin(req, requestID)

		if iResult == 0 {
			log.Println("发送用户登录请求: 成功.")
			return
		}
		log.Println("发送用户登录请求: 失败.")
	})
	return obj, err
}

func (p *GoCThostFtdcMdSpi) OnRspUserLogin(pRspUserLogin goctp.CThostFtdcRspUserLoginField, pRspInfo goctp.CThostFtdcRspInfoField, nRequestID int, bIsLast bool) {

	if bIsLast && !p.IsErrorRspInfo(pRspInfo) {
		async.Put(nRequestID, pRspUserLogin, nil)
		return
	}
	async.Put(nRequestID, nil, fmt.Errorf("pRspUserLogin error "))
}
<commit_msg>debug<commit_after>package md

import (
	"flag"
	"fmt"
	"log"

	"github.com/mrwill84/CtpSender/ctp/async"
	"github.com/mrwill84/goctp"
)

var (
	broker_id    = flag.String("BrokerID", "9999", "经纪公司编号,SimNow BrokerID统一为：9999")
	investor_id  = flag.String("InvestorID", "<InvestorID>", "交易用户代码")
	pass_word    = flag.String("Password", "<Password>", "交易用户密码")
	market_front = flag.String("MarketFront", "tcp://180.168.146.187:10031", "行情前置,SimNow的测试环境: tcp://180.168.146.187:10031")
	trade_front  = flag.String("TradeFront", "tcp://180.168.146.187:10030", "交易前置,SimNow的测试环境: tcp://180.168.146.187:10030")
)

var CTP GoCTPClient

type GoCTPClient struct {
	BrokerID   string
	InvestorID string
	Password   string

	MdFront string
	MdApi   goctp.CThostFtdcMdApi

	TraderFront string
	TraderApi   goctp.CThostFtdcTraderApi

	MdRequestID     int
	TraderRequestID int
	FrontendConnent chan bool
}

func (g *GoCTPClient) GetMdRequestID() int {
	g.MdRequestID += 1
	return g.MdRequestID
}

func NewDirectorCThostFtdcMdSpi(v interface{}) goctp.CThostFtdcMdSpi {

	return goctp.NewDirectorCThostFtdcMdSpi(v)

}

type GoCThostFtdcMdSpi struct {
	Client GoCTPClient
}

func (p *GoCThostFtdcMdSpi) OnRspError(pRspInfo goctp.CThostFtdcRspInfoField, nRequestID int, bIsLast bool) {
	log.Println("GoCThostFtdcMdSpi.OnRspError.")
	p.IsErrorRspInfo(pRspInfo)
}

func (p *GoCThostFtdcMdSpi) OnFrontDisconnected(nReason int) {
	log.Printf("GoCThostFtdcMdSpi.OnFrontDisconnected: %#v\n", nReason)
}

func (p *GoCThostFtdcMdSpi) OnHeartBeatWarning(nTimeLapse int) {
	log.Printf("GoCThostFtdcMdSpi.OnHeartBeatWarning: %v", nTimeLapse)

}

func (p *GoCThostFtdcMdSpi) OnFrontConnected() {
	log.Println("GoCThostFtdcMdSpi.OnFrontConnected.")
	if CTP.FrontendConnent != nil {
		CTP.FrontendConnent <- true
		return
	}
	panic("never")
	//p.ReqUserLogin()
}

func (p *GoCThostFtdcMdSpi) IsErrorRspInfo(pRspInfo goctp.CThostFtdcRspInfoField) bool {
	// 如果ErrorID != 0, 说明收到了错误的响应
	bResult := (pRspInfo.GetErrorID() != 0)
	if bResult {
		log.Printf("ErrorID=%v ErrorMsg=%v\n", pRspInfo.GetErrorID(), pRspInfo.GetErrorMsg())
	}
	return bResult
}

func (p *GoCThostFtdcMdSpi) ReqUserLogin() (interface{}, error) {
	id := p.Client.GetMdRequestID()
	obj, err := async.Alloc(id, func() {
		log.Println("GoCThostFtdcMdSpi.ReqUserLogin.")
		req := goctp.NewCThostFtdcReqUserLoginField()
		req.SetBrokerID(p.Client.BrokerID)
		req.SetUserID(p.Client.InvestorID)
		req.SetPassword(p.Client.Password)
		requestID := id
		fmt.Println("ReqUserLogin(req, requestID)", requestID)
		iResult := p.Client.MdApi.ReqUserLogin(req, requestID)

		if iResult == 0 {
			log.Println("发送用户登录请求: 成功.")
			return
		}
		log.Println("发送用户登录请求: 失败.")
	})
	return obj, err
}

func (p *GoCThostFtdcMdSpi) OnRspUserLogin(pRspUserLogin goctp.CThostFtdcRspUserLoginField, pRspInfo goctp.CThostFtdcRspInfoField, nRequestID int, bIsLast bool) {

	if bIsLast && !p.IsErrorRspInfo(pRspInfo) {
		async.Put(nRequestID, pRspUserLogin, nil)
		return
	}
	async.Put(nRequestID, nil, fmt.Errorf("pRspUserLogin error "))
}
<|endoftext|>


<commit_before>package rabbitsmpp

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

type Consumer interface {
	Consume() (<-chan Job, <-chan error, error)
	Stop() error
	ID() string
}

type consumer struct {
	*client
	channel Channel
	ctx     context.Context
	cancel  context.CancelFunc
}

func NewConsumer(conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(context.Background())

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func NewConsumerWithContext(ctx context.Context, conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(ctx)

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func (c *consumer) ID() string {
	return c.client.QueueName()
}

func (c *consumer) bindWithRetry() chan *amqp.Error {
	closeChan, err := c.Bind()
	for err != nil {
		log.Println("Failed to bind consumer:", err)
		time.Sleep(5 * time.Second)
		closeChan, err = c.Bind()
	}
	return closeChan
}

func (c *consumer) getConsumeChannel() (<-chan amqp.Delivery, error) {
	ch, err := c.Channel()
	if err != nil {
		return nil, err
	}
	c.channel = ch

	q, err := c.channel.QueueDeclare(
		c.QueueName(), // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		nil,           // arguments
	)
	if err != nil {
		return nil, err
	}

	return c.channel.Consume(
		q.Name, // queue
		"",     // consumer
		false,  // auto-ackey
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
}

func (c *consumer) Consume() (<-chan Job, <-chan error, error) {
	if c.channel != nil {
		return nil, nil, errors.New("consumer already active")
	}
	closeChan, err := c.Bind()
	if err != nil {
		return nil, nil, err
	}
	dlvChan, err := c.getConsumeChannel()
	if err != nil {
		return nil, nil, err
	}
	jobChan := make(chan Job)
	errChan := make(chan error)

	go func() {
		defer func() {
			_ = c.channel.Close()
			c.channel = nil
		}()
		defer close(jobChan)
		defer close(errChan)
		for {
			err = c.consume(dlvChan, closeChan, jobChan, errChan)
			// if consume returns without an error, means that it was terminated
			// properly, otherwise something went wrong and it needs to restart
			if err == nil {
				log.Printf("EOF consuming for: %s", c.ID())
				return
			}
			log.Println("stopped consuming jobs:", err)
			closeChan = c.bindWithRetry()
			dlvChan, err = c.getConsumeChannel()
			for err != nil {
				time.Sleep(5 * time.Second)
				dlvChan, err = c.getConsumeChannel()
			}
		}
	}()

	return jobChan, errChan, nil
}

func (c *consumer) consume(dlvChan <-chan amqp.Delivery, closeChan <-chan *amqp.Error, jobChan chan<- Job, errChan chan<- error) error {
	for {
		select {
		case d := <-dlvChan:
			j := Job{}
			err := json.Unmarshal(d.Body, &j)
			if err != nil {
				errChan <- fmt.Errorf("failed to unmarshal PDU: %v", err)
				return err
			}
			j.delivery = &d
			jobChan <- j
		case err := <-closeChan:
			return err
		case <-c.ctx.Done():
			return nil
		}
	}
}

func (c *consumer) Stop() error {
	if c.channel == nil {
		return nil
	}
	// Sends the stop signal
	c.cancel()
	return nil
}
<commit_msg>Return the error in the errorChan outside the consume()<commit_after>package rabbitsmpp

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/streadway/amqp"
	"golang.org/x/net/context"
)

type Consumer interface {
	Consume() (<-chan Job, <-chan error, error)
	Stop() error
	ID() string
}

type consumer struct {
	*client
	channel Channel
	ctx     context.Context
	cancel  context.CancelFunc
}

func NewConsumer(conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(context.Background())

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func NewConsumerWithContext(ctx context.Context, conf Config) (Consumer, error) {
	c := NewClient(conf).(*client)
	ctx, cancel := context.WithCancel(ctx)

	return &consumer{
		client:  c,
		channel: nil,
		ctx:     ctx,
		cancel:  cancel,
	}, nil
}

func (c *consumer) ID() string {
	return c.client.QueueName()
}

func (c *consumer) bindWithRetry() chan *amqp.Error {
	closeChan, err := c.Bind()
	for err != nil {
		log.Println("Failed to bind consumer:", err)
		time.Sleep(5 * time.Second)
		closeChan, err = c.Bind()
	}
	return closeChan
}

func (c *consumer) getConsumeChannel() (<-chan amqp.Delivery, error) {
	ch, err := c.Channel()
	if err != nil {
		return nil, err
	}
	c.channel = ch

	q, err := c.channel.QueueDeclare(
		c.QueueName(), // name
		true,          // durable
		false,         // delete when unused
		false,         // exclusive
		false,         // no-wait
		nil,           // arguments
	)
	if err != nil {
		return nil, err
	}

	return c.channel.Consume(
		q.Name, // queue
		"",     // consumer
		false,  // auto-ackey
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
}

func (c *consumer) Consume() (<-chan Job, <-chan error, error) {
	if c.channel != nil {
		return nil, nil, errors.New("consumer already active")
	}
	closeChan, err := c.Bind()
	if err != nil {
		return nil, nil, err
	}
	dlvChan, err := c.getConsumeChannel()
	if err != nil {
		return nil, nil, err
	}
	jobChan := make(chan Job)
	errChan := make(chan error)

	go func() {
		defer func() {
			_ = c.channel.Close()
			c.channel = nil
		}()
		defer close(jobChan)
		defer close(errChan)
		for {
			err = c.consume(dlvChan, closeChan, jobChan)
			// if consume returns without an error, means that it was terminated
			// properly, otherwise something went wrong and it needs to restart
			if err == nil {
				log.Printf("EOF consuming for: %s", c.ID())
				return
			}
			errChan <- err
			log.Println("stopped consuming jobs:", err)
			closeChan = c.bindWithRetry()
			dlvChan, err = c.getConsumeChannel()
			for err != nil {
				time.Sleep(5 * time.Second)
				dlvChan, err = c.getConsumeChannel()
			}
		}
	}()

	return jobChan, errChan, nil
}

func (c *consumer) consume(dlvChan <-chan amqp.Delivery, closeChan <-chan *amqp.Error, jobChan chan<- Job) error {
	for {
		select {
		case d := <-dlvChan:
			j := Job{}
			err := json.Unmarshal(d.Body, &j)
			if err != nil {
				return fmt.Errorf("failed to unmarshal PDU: %v", err)
			}
			j.delivery = &d
			jobChan <- j
		case err := <-closeChan:
			return err
		case <-c.ctx.Done():
			return nil
		}
	}
}

func (c *consumer) Stop() error {
	if c.channel == nil {
		return nil
	}
	// Sends the stop signal
	c.cancel()
	return nil
}
<|endoftext|>


<commit_before>package main

import "strings"
import "strconv"
import "fmt"
import "errors"
import "net"
import "net/http"
import "net/url"
import "context"
import "time"
import "html"
import "log"
import "io"
import "bufio"
import "bytes"
import "flag"
import "encoding/xml"
import "golang.org/x/tools/blog/atom"

const versionMajor = 0
const versionMinor = 1
const versionBugfix = 0
const versionGit = true

const bugzillaDateFormat = "2006-01-02 15:04:05 -0700"
const userAgentName = "bugzillatoatom"

// Maximum size a request for a bug-xml is read in byte.
var maxBugRequestRead int64

// Maximum number of requests per second. Set to something negative to disable
var maxRequestsPerSecond int

// Channel to block on during too many requests in a second
var tooManyRequestsBlocker chan bool = make(chan bool)

// returns the minimum of the given values
func min(a int, b int) int {
    if a <= b {
        return a
    } else {
        return b
    }
}

// Returns the current version as string in the form major.minor.bugfix(-git)
func getVersion() string {
    if versionGit {
        return strconv.Itoa(versionMajor) + "." + strconv.Itoa(versionMinor) + "." + strconv.Itoa(versionBugfix) + "-git"
    } else {
        return strconv.Itoa(versionMajor) + "." + strconv.Itoa(versionMinor) + "." + strconv.Itoa(versionBugfix)
    }
}

// Read from r until given string is found. Appends toAppend afterwards if string
// is found and in any case returns the result
func readUntilString(r io.Reader, until string, toAppend string) (string, error) {
    var buffer bytes.Buffer
    eofReached := false
    rs := bufio.NewReader(io.LimitReader(r, maxBugRequestRead))

    for !eofReached {
        str, err := rs.ReadString('\n')

        if err == io.EOF {
            eofReached = true
        } else if err != nil {
            return "", errors.New(fmt.Sprintf("Error during reading from url: %s", err))
        }

        index := strings.Index(str, until)

        if index == -1 {
            buffer.WriteString(str)
        } else {
            buffer.WriteString(str[:index])
            buffer.WriteString(toAppend)
            break
        }
    }

    return buffer.String(), nil
}

// Requests the bug from given url
func doRequest(target *url.URL) (string, error) {
    request := http.Request{
        Method: http.MethodGet,
        URL: target,
        Header: http.Header{
            "User-Agent": { userAgentName + "/" + getVersion() },
        },
    }

    resp, err := http.DefaultClient.Do(&request)

    if err != nil {
        return "", err
    }

    defer resp.Body.Close()

    if resp.StatusCode != http.StatusOK {
        return "", errors.New(fmt.Sprintf("Request returned status code %d (%s).", resp.StatusCode, http.StatusText(resp.StatusCode)))
    }

    // TODO: maybe we should search for something more clever to abort. <attachment could be given by a user in a report
    return readUntilString(resp.Body, "<attachment", "</bug></bugzilla>")
}

// converts the given xml string into an atom feed
func convertXmlToAtom(inXml string) (string, error) {
    type Who struct {
        Name string `xml:",chardata"`
        RealName string `xml:"name,attr"`
    }

    getFormatedName := func(w Who) string {
        if w.RealName == "" {
            return w.Name
        } else {
            return w.RealName + " (" + w.Name + ")"
        }
    }

    type Comment struct {
        CommentId int `xml:"commentid"`
        CommentCount int `xml:"comment_count"`
        AttachmentID int `xml:"attachid"`
        Who Who `xml:"who"`
        When string `xml:"bug_when"`
        Text string `xml:"thetext"`
    }

    type InResult struct {
        Urlbase string `xml:"urlbase,attr"`
        BugId int `xml:"bug>bug_id"`
        Description string `xml:"bug>short_desc"`
        Comments []Comment `xml:"bug>long_desc"`
    }

    inResult := InResult{}
    err := xml.Unmarshal([]byte(inXml), &inResult)
    if err != nil {
        return "", errors.New(fmt.Sprintf("Error during unmarshalling the xml: %s", err))
    } else if len(inResult.Comments) == 0 {
        // One comment, the initial one, should always be available
        err := errors.New("Zero comments in bug. There should be at least the initial one.")
        return "", err
    }

    updateTime, err := time.Parse(bugzillaDateFormat, inResult.Comments[len(inResult.Comments)-1].When)
    if err != nil {
        return "", errors.New(fmt.Sprintf("Couldn't parse updateTime in initial comment: %s", err))
    }

    inUrl := fmt.Sprintf("%s/show_bug.cgi?id=%d", inResult.Urlbase, inResult.BugId)
    attachmentUrl := fmt.Sprintf("%s/attachment.cgi?id=", inResult.Urlbase)

    feed := &atom.Feed{
        Title: inResult.Description,
        ID: inUrl,
        Link: []atom.Link{atom.Link{Href: inUrl, Rel: "alternate"}},
        Updated: atom.Time(updateTime),
        Author: &atom.Person{Name: getFormatedName(inResult.Comments[0].Who)},
        Entry: make([]*atom.Entry, 0, len(inResult.Comments)),
    }

    for i, comment := range inResult.Comments {
        creationTime, err := time.Parse(bugzillaDateFormat, comment.When)
        if err != nil {
            return "", errors.New(fmt.Sprintf("Couldn't parse updateTime in comment %d: %s", i, err))
        }

        links := []atom.Link{atom.Link{Href: inUrl + "#c" + strconv.Itoa(comment.CommentCount), Rel: "alternate"}}
        if comment.AttachmentID != 0 {
            links = append(links, atom.Link{Href: attachmentUrl + strconv.Itoa(comment.AttachmentID), Rel: "enclosure"})
        }

        entry := &atom.Entry{
            Title: getFormatedName(comment.Who) + ": " + comment.Text[:min(100, len(comment.Text))],
            ID: inUrl + "#c" + strconv.Itoa(comment.CommentCount),
            Link: links,
            Published: atom.Time(creationTime),
            Author: &atom.Person{Name: getFormatedName(comment.Who)},
            Content: &atom.Text{Type: "html", Body: strings.Replace(html.EscapeString(comment.Text), "\n", "
", -1)},
        }

        feed.Entry = append(feed.Entry, entry)
    }

    atom, err := xml.MarshalIndent(feed, "", "\t")
    if err != nil {
        return "", errors.New(fmt.Sprintf("Error during creating the atom feed: %s", err))
    }

    return xml.Header + string(atom), nil
}

// Checks for not allowed targets, defined by the given networks.
// Returns a list of IPs if target is allowed, otherwise
// throws an error
func targetAllowedIps(target string, forbiddenNetworks []*net.IPNet) ([]net.IP, error) {
    ips, err := net.LookupIP(target)

    if forbiddenNetworks == nil {
        return ips, nil
    }

    if err != nil {
        return []net.IP{}, err
    }

    for _, ip := range ips {
        for _, ipnet := range forbiddenNetworks {
            if ipnet.Contains(ip) {
                log.Printf("Blocked target \"%s\" since it's IP %s is contained in blocked network %s.\n", target, ip, ipnet)
                errStr := fmt.Sprintf("Target \"%s\" blocked.", target)
                return []net.IP{}, errors.New(errStr)
            }
        }
    }

    return ips, nil
}

func handleConvert(w http.ResponseWriter, r *http.Request) {
    // Block during too many requests in the last second
    if maxRequestsPerSecond >= 0 {
        <-tooManyRequestsBlocker
    }

    // Check for a possible recursive call
    if r.Header != nil {
        for _, agent := range r.Header["User-Agent"] {
            if strings.Contains(agent, userAgentName) {
                log.Printf("Blocked request by %s due to User-Agent \"%s\".\n", r.RemoteAddr, agent)
                errStr := fmt.Sprintf("User-Agent \"%s\" blocked.", r.Header["User-Agent"])
                http.Error(w, errStr, http.StatusForbidden)
                return
            }
        }
    }

    formValueUrl := r.FormValue("url")

    // if the user didn't give a protocol simply assume http
    if !(strings.HasPrefix(formValueUrl, "http://") || strings.HasPrefix(formValueUrl, "https://")) {
        formValueUrl = "http://" + formValueUrl
    }

    target, err := url.Parse(formValueUrl)

    if err != nil {
        log.Printf("Error occurred during parsing the url \"%s\": %s.\n", r.FormValue("url"), err.Error())
        errStr := fmt.Sprintf("Error occurred during parsing the url \"%s\": %s\nAre you sure the url is correct?", r.FormValue("url"), err.Error())
        http.Error(w, errStr, http.StatusInternalServerError)
        return
    }

    parsedQuery := target.Query()
    parsedQuery.Set("ctype", "xml")
    target.RawQuery = parsedQuery.Encode()
    target.Fragment = ""

    inXml, err := doRequest(target)
    if err != nil {
        log.Printf("Error occurred during fetching the url \"%s\": %s\n", target.String(), err.Error())
        errStr := fmt.Sprintf("Error occurred during fetching the url \"%s\": %s\nAre you sure the url is correct?", target.String(), err.Error())
        http.Error(w, errStr, http.StatusInternalServerError)
        return
    }

    atom, err := convertXmlToAtom(inXml)
    if err != nil {
        log.Printf("Error occurred during conversion of the url \"%s\" to atom: %s\n", target.String(), err.Error())
        errStr := fmt.Sprintf("Error occurred during conversion of the url \"%s\" to atom: %s\nAre you sure the url is correct?", target.String(), err.Error())
        http.Error(w, errStr, http.StatusInternalServerError)
        return
    }

    fmt.Fprintf(w, "%s", atom)
}

func handleMain(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "%s", `

<title>bugzillatoatom</title>

<form action=convert>
    Convert a Bugzilla bug entry into an Atom feed. Enter an url:
    <input type="text" name="url">

`)
}

// Parses a given IP or CIDR into a CIDR. IPs are treated as CIDRs with full bitmask
func parseIPOrCIDR(str string) (*net.IPNet, error) {

    if !strings.Contains(str, "/") {
        if strings.Contains(str, ":") {
            str = str + "/128"
        } else {
            str = str + "/32"
        }
    }

    _, ipnet, err := net.ParseCIDR(str)
    return ipnet, err
}

// To allow forbiddenNetworks to be parsed as argument
type CIDRList []*net.IPNet

func (forbiddenNetworks *CIDRList) String() string {
    strs := []string{}

    for _, ipnet := range *forbiddenNetworks {
        strs = append(strs, ipnet.String())
    }

    return strings.Join(strs, ", ")
}

func (forbiddenNetworks *CIDRList) Set(value string) error {
    ipnet, err := parseIPOrCIDR(value)

    if err == nil {
        *forbiddenNetworks = append(*forbiddenNetworks, ipnet)
    }

    return err
}

// Sets the http.DefaultClient to a client with timeout which
// blocks connections to a host with an IP in forbiddenNetworks
func setHttpDefaultClient(forbiddenNetworks []*net.IPNet) {
    // Copy the original DefaultTransport and add our dialer
    httpDialer := &net.Dialer{
       Timeout:   30 * time.Second,
       KeepAlive: 30 * time.Second,
    }

    httpTransport := *(http.DefaultTransport.(*http.Transport))
    httpTransport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
        port := ""
        portPosition := strings.Index(addr, ":")

        if portPosition != -1 {
            port = addr[portPosition:]
            addr = addr[:portPosition]
        }

        ips, err := targetAllowedIps(addr, forbiddenNetworks)

        if err != nil {
            return nil, err
        }

        for _, ip := range ips {
            var ipstr string

            if len(ip.To4()) == net.IPv4len {
                ipstr = ip.String() + port
            } else {
                ipstr = "[" + ip.String() + "]" + port
            }

            conn, err := httpDialer.DialContext(ctx, network, ipstr)

            if err == nil {
                return conn, nil
            }
        }

        return nil, errors.New("Unable to connect to " + addr)
    }

    // Add a timeout for the default http.Get() in case something goes wrong
    // on the other side.
    http.DefaultClient = &http.Client{
        Timeout: time.Second * 30,
        Transport: &httpTransport,
    }
}

func main() {
    version := flag.Bool("version", false, "Print the current version and exit")
    port := flag.Uint64("p", 9080, "Port to bind to")
    maxBugRequestReadFlag := flag.Uint64("requestsize", 1 * 1024 * 1024, "Maximum number of bytes to read during a request to another server.") // 1MiB per default
    flag.IntVar(&maxRequestsPerSecond, "persecond", 5, "Maximum number of requests to another server per second. Set to -1 to disable.")
    forbiddenNetworks := CIDRList{}
    flag.Var(&forbiddenNetworks, "b", "IP or Network in CIDR format to block. If a host is available under any blocked IP it will be blocked. Can be given multiple times.\n\tYou probably want to exclude localhost or local networks both on IPv4 and IPv6.")
    flag.Parse()

    if *version {
        log.Fatalln(getVersion())
    }

    if *maxBugRequestReadFlag & (1 << 63) != 0 {
        log.Fatalln("Too large requestsize")
    } else {
        maxBugRequestRead = int64(*maxBugRequestReadFlag)
    }

    setHttpDefaultClient(forbiddenNetworks)

    for i := 0; i < maxRequestsPerSecond; i++ {
        go func() {
            for {
                tooManyRequestsBlocker <- true
                time.Sleep(time.Second)
            }
        }()
    }

    http.HandleFunc("/convert", handleConvert)
    http.HandleFunc("/", handleMain)

    log.Fatal(http.ListenAndServe(":" + strconv.FormatUint(*port, 10), nil))
}<commit_msg>run go fmt<commit_after>package main

import "strings"
import "strconv"
import "fmt"
import "errors"
import "net"
import "net/http"
import "net/url"
import "context"
import "time"
import "html"
import "log"
import "io"
import "bufio"
import "bytes"
import "flag"
import "encoding/xml"
import "golang.org/x/tools/blog/atom"

const versionMajor = 0
const versionMinor = 1
const versionBugfix = 0
const versionGit = true

const bugzillaDateFormat = "2006-01-02 15:04:05 -0700"
const userAgentName = "bugzillatoatom"

// Maximum size a request for a bug-xml is read in byte.
var maxBugRequestRead int64

// Maximum number of requests per second. Set to something negative to disable
var maxRequestsPerSecond int

// Channel to block on during too many requests in a second
var tooManyRequestsBlocker chan bool = make(chan bool)

// returns the minimum of the given values
func min(a int, b int) int {
	if a <= b {
		return a
	} else {
		return b
	}
}

// Returns the current version as string in the form major.minor.bugfix(-git)
func getVersion() string {
	if versionGit {
		return strconv.Itoa(versionMajor) + "." + strconv.Itoa(versionMinor) + "." + strconv.Itoa(versionBugfix) + "-git"
	} else {
		return strconv.Itoa(versionMajor) + "." + strconv.Itoa(versionMinor) + "." + strconv.Itoa(versionBugfix)
	}
}

// Read from r until given string is found. Appends toAppend afterwards if string
// is found and in any case returns the result
func readUntilString(r io.Reader, until string, toAppend string) (string, error) {
	var buffer bytes.Buffer
	eofReached := false
	rs := bufio.NewReader(io.LimitReader(r, maxBugRequestRead))

	for !eofReached {
		str, err := rs.ReadString('\n')

		if err == io.EOF {
			eofReached = true
		} else if err != nil {
			return "", errors.New(fmt.Sprintf("Error during reading from url: %s", err))
		}

		index := strings.Index(str, until)

		if index == -1 {
			buffer.WriteString(str)
		} else {
			buffer.WriteString(str[:index])
			buffer.WriteString(toAppend)
			break
		}
	}

	return buffer.String(), nil
}

// Requests the bug from given url
func doRequest(target *url.URL) (string, error) {
	request := http.Request{
		Method: http.MethodGet,
		URL:    target,
		Header: http.Header{
			"User-Agent": {userAgentName + "/" + getVersion()},
		},
	}

	resp, err := http.DefaultClient.Do(&request)

	if err != nil {
		return "", err
	}

	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", errors.New(fmt.Sprintf("Request returned status code %d (%s).", resp.StatusCode, http.StatusText(resp.StatusCode)))
	}

	// TODO: maybe we should search for something more clever to abort. <attachment could be given by a user in a report
	return readUntilString(resp.Body, "<attachment", "</bug></bugzilla>")
}

// converts the given xml string into an atom feed
func convertXmlToAtom(inXml string) (string, error) {
	type Who struct {
		Name     string `xml:",chardata"`
		RealName string `xml:"name,attr"`
	}

	getFormatedName := func(w Who) string {
		if w.RealName == "" {
			return w.Name
		} else {
			return w.RealName + " (" + w.Name + ")"
		}
	}

	type Comment struct {
		CommentId    int    `xml:"commentid"`
		CommentCount int    `xml:"comment_count"`
		AttachmentID int    `xml:"attachid"`
		Who          Who    `xml:"who"`
		When         string `xml:"bug_when"`
		Text         string `xml:"thetext"`
	}

	type InResult struct {
		Urlbase     string    `xml:"urlbase,attr"`
		BugId       int       `xml:"bug>bug_id"`
		Description string    `xml:"bug>short_desc"`
		Comments    []Comment `xml:"bug>long_desc"`
	}

	inResult := InResult{}
	err := xml.Unmarshal([]byte(inXml), &inResult)
	if err != nil {
		return "", errors.New(fmt.Sprintf("Error during unmarshalling the xml: %s", err))
	} else if len(inResult.Comments) == 0 {
		// One comment, the initial one, should always be available
		err := errors.New("Zero comments in bug. There should be at least the initial one.")
		return "", err
	}

	updateTime, err := time.Parse(bugzillaDateFormat, inResult.Comments[len(inResult.Comments)-1].When)
	if err != nil {
		return "", errors.New(fmt.Sprintf("Couldn't parse updateTime in initial comment: %s", err))
	}

	inUrl := fmt.Sprintf("%s/show_bug.cgi?id=%d", inResult.Urlbase, inResult.BugId)
	attachmentUrl := fmt.Sprintf("%s/attachment.cgi?id=", inResult.Urlbase)

	feed := &atom.Feed{
		Title:   inResult.Description,
		ID:      inUrl,
		Link:    []atom.Link{atom.Link{Href: inUrl, Rel: "alternate"}},
		Updated: atom.Time(updateTime),
		Author:  &atom.Person{Name: getFormatedName(inResult.Comments[0].Who)},
		Entry:   make([]*atom.Entry, 0, len(inResult.Comments)),
	}

	for i, comment := range inResult.Comments {
		creationTime, err := time.Parse(bugzillaDateFormat, comment.When)
		if err != nil {
			return "", errors.New(fmt.Sprintf("Couldn't parse updateTime in comment %d: %s", i, err))
		}

		links := []atom.Link{atom.Link{Href: inUrl + "#c" + strconv.Itoa(comment.CommentCount), Rel: "alternate"}}
		if comment.AttachmentID != 0 {
			links = append(links, atom.Link{Href: attachmentUrl + strconv.Itoa(comment.AttachmentID), Rel: "enclosure"})
		}

		entry := &atom.Entry{
			Title:     getFormatedName(comment.Who) + ": " + comment.Text[:min(100, len(comment.Text))],
			ID:        inUrl + "#c" + strconv.Itoa(comment.CommentCount),
			Link:      links,
			Published: atom.Time(creationTime),
			Author:    &atom.Person{Name: getFormatedName(comment.Who)},
			Content:   &atom.Text{Type: "html", Body: strings.Replace(html.EscapeString(comment.Text), "\n", "
", -1)},
		}

		feed.Entry = append(feed.Entry, entry)
	}

	atom, err := xml.MarshalIndent(feed, "", "\t")
	if err != nil {
		return "", errors.New(fmt.Sprintf("Error during creating the atom feed: %s", err))
	}

	return xml.Header + string(atom), nil
}

// Checks for not allowed targets, defined by the given networks.
// Returns a list of IPs if target is allowed, otherwise
// throws an error
func targetAllowedIps(target string, forbiddenNetworks []*net.IPNet) ([]net.IP, error) {
	ips, err := net.LookupIP(target)

	if forbiddenNetworks == nil {
		return ips, nil
	}

	if err != nil {
		return []net.IP{}, err
	}

	for _, ip := range ips {
		for _, ipnet := range forbiddenNetworks {
			if ipnet.Contains(ip) {
				log.Printf("Blocked target \"%s\" since it's IP %s is contained in blocked network %s.\n", target, ip, ipnet)
				errStr := fmt.Sprintf("Target \"%s\" blocked.", target)
				return []net.IP{}, errors.New(errStr)
			}
		}
	}

	return ips, nil
}

func handleConvert(w http.ResponseWriter, r *http.Request) {
	// Block during too many requests in the last second
	if maxRequestsPerSecond >= 0 {
		<-tooManyRequestsBlocker
	}

	// Check for a possible recursive call
	if r.Header != nil {
		for _, agent := range r.Header["User-Agent"] {
			if strings.Contains(agent, userAgentName) {
				log.Printf("Blocked request by %s due to User-Agent \"%s\".\n", r.RemoteAddr, agent)
				errStr := fmt.Sprintf("User-Agent \"%s\" blocked.", r.Header["User-Agent"])
				http.Error(w, errStr, http.StatusForbidden)
				return
			}
		}
	}

	formValueUrl := r.FormValue("url")

	// if the user didn't give a protocol simply assume http
	if !(strings.HasPrefix(formValueUrl, "http://") || strings.HasPrefix(formValueUrl, "https://")) {
		formValueUrl = "http://" + formValueUrl
	}

	target, err := url.Parse(formValueUrl)

	if err != nil {
		log.Printf("Error occurred during parsing the url \"%s\": %s.\n", r.FormValue("url"), err.Error())
		errStr := fmt.Sprintf("Error occurred during parsing the url \"%s\": %s\nAre you sure the url is correct?", r.FormValue("url"), err.Error())
		http.Error(w, errStr, http.StatusInternalServerError)
		return
	}

	parsedQuery := target.Query()
	parsedQuery.Set("ctype", "xml")
	target.RawQuery = parsedQuery.Encode()
	target.Fragment = ""

	inXml, err := doRequest(target)
	if err != nil {
		log.Printf("Error occurred during fetching the url \"%s\": %s\n", target.String(), err.Error())
		errStr := fmt.Sprintf("Error occurred during fetching the url \"%s\": %s\nAre you sure the url is correct?", target.String(), err.Error())
		http.Error(w, errStr, http.StatusInternalServerError)
		return
	}

	atom, err := convertXmlToAtom(inXml)
	if err != nil {
		log.Printf("Error occurred during conversion of the url \"%s\" to atom: %s\n", target.String(), err.Error())
		errStr := fmt.Sprintf("Error occurred during conversion of the url \"%s\" to atom: %s\nAre you sure the url is correct?", target.String(), err.Error())
		http.Error(w, errStr, http.StatusInternalServerError)
		return
	}

	fmt.Fprintf(w, "%s", atom)
}

func handleMain(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "%s", `

<title>bugzillatoatom</title>

<form action=convert>
    Convert a Bugzilla bug entry into an Atom feed. Enter an url:
    <input type="text" name="url">

`)
}

// Parses a given IP or CIDR into a CIDR. IPs are treated as CIDRs with full bitmask
func parseIPOrCIDR(str string) (*net.IPNet, error) {

	if !strings.Contains(str, "/") {
		if strings.Contains(str, ":") {
			str = str + "/128"
		} else {
			str = str + "/32"
		}
	}

	_, ipnet, err := net.ParseCIDR(str)
	return ipnet, err
}

// To allow forbiddenNetworks to be parsed as argument
type CIDRList []*net.IPNet

func (forbiddenNetworks *CIDRList) String() string {
	strs := []string{}

	for _, ipnet := range *forbiddenNetworks {
		strs = append(strs, ipnet.String())
	}

	return strings.Join(strs, ", ")
}

func (forbiddenNetworks *CIDRList) Set(value string) error {
	ipnet, err := parseIPOrCIDR(value)

	if err == nil {
		*forbiddenNetworks = append(*forbiddenNetworks, ipnet)
	}

	return err
}

// Sets the http.DefaultClient to a client with timeout which
// blocks connections to a host with an IP in forbiddenNetworks
func setHttpDefaultClient(forbiddenNetworks []*net.IPNet) {
	// Copy the original DefaultTransport and add our dialer
	httpDialer := &net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}

	httpTransport := *(http.DefaultTransport.(*http.Transport))
	httpTransport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
		port := ""
		portPosition := strings.Index(addr, ":")

		if portPosition != -1 {
			port = addr[portPosition:]
			addr = addr[:portPosition]
		}

		ips, err := targetAllowedIps(addr, forbiddenNetworks)

		if err != nil {
			return nil, err
		}

		for _, ip := range ips {
			var ipstr string

			if len(ip.To4()) == net.IPv4len {
				ipstr = ip.String() + port
			} else {
				ipstr = "[" + ip.String() + "]" + port
			}

			conn, err := httpDialer.DialContext(ctx, network, ipstr)

			if err == nil {
				return conn, nil
			}
		}

		return nil, errors.New("Unable to connect to " + addr)
	}

	// Add a timeout for the default http.Get() in case something goes wrong
	// on the other side.
	http.DefaultClient = &http.Client{
		Timeout:   time.Second * 30,
		Transport: &httpTransport,
	}
}

func main() {
	version := flag.Bool("version", false, "Print the current version and exit")
	port := flag.Uint64("p", 9080, "Port to bind to")
	maxBugRequestReadFlag := flag.Uint64("requestsize", 1*1024*1024, "Maximum number of bytes to read during a request to another server.") // 1MiB per default
	flag.IntVar(&maxRequestsPerSecond, "persecond", 5, "Maximum number of requests to another server per second. Set to -1 to disable.")
	forbiddenNetworks := CIDRList{}
	flag.Var(&forbiddenNetworks, "b", "IP or Network in CIDR format to block. If a host is available under any blocked IP it will be blocked. Can be given multiple times.\n\tYou probably want to exclude localhost or local networks both on IPv4 and IPv6.")
	flag.Parse()

	if *version {
		log.Fatalln(getVersion())
	}

	if *maxBugRequestReadFlag&(1<<63) != 0 {
		log.Fatalln("Too large requestsize")
	} else {
		maxBugRequestRead = int64(*maxBugRequestReadFlag)
	}

	setHttpDefaultClient(forbiddenNetworks)

	for i := 0; i < maxRequestsPerSecond; i++ {
		go func() {
			for {
				tooManyRequestsBlocker <- true
				time.Sleep(time.Second)
			}
		}()
	}

	http.HandleFunc("/convert", handleConvert)
	http.HandleFunc("/", handleMain)

	log.Fatal(http.ListenAndServe(":"+strconv.FormatUint(*port, 10), nil))
}
<|endoftext|>