<commit_before>package dns

import (
	"time"
)

// Envelope is used when doing a zone transfer with a remote server.
type Envelope struct {
	RR    []RR  // The set of RRs in the answer section of the xfr reply message.
	Error error // If something went wrong, this contains the error.
}

// A Transfer defines parameters that are used during a zone transfer.
type Transfer struct {
	*Conn
	DialTimeout    time.Duration     // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	TsigSecret     map[string]string // Secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	tsigTimersOnly bool
}

// Think we need to away to stop the transfer

// In performs an incoming transfer with the server in a.
func (t *Transfer) In(q *Msg, a string) (env chan *Envelope, err error) {
	timeout := dnsTimeout
	if t.DialTimeout != 0 {
		timeout = t.DialTimeout
	}
	t.Conn, err = DialTimeout("tcp", a, timeout)
	if err != nil {
		return nil, err
	}
	if err := t.WriteMsg(q); err != nil {
		return nil, err
	}
	env = make(chan *Envelope)
	go func() {
		if q.Question[0].Qtype == TypeAXFR {
			go t.inAxfr(q.Id, env)
			return
		}
		if q.Question[0].Qtype == TypeIXFR {
			go t.inIxfr(q.Id, env)
			return
		}
	}()
	return env, nil
}

func (t *Transfer) inAxfr(id uint16, c chan *Envelope) {
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.Conn.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{nil, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			first = !first
			// only one answer that is SOA, receive more
			if len(in.Answer) == 1 {
				t.tsigTimersOnly = true
				c <- &Envelope{in.Answer, nil}
				continue
			}
		}

		if !first {
			t.tsigTimersOnly = true // Subsequent envelopes use this.
			if isSOALast(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
	panic("dns: not reached")
}

func (t *Transfer) inIxfr(id uint16, c chan *Envelope) {
	serial := uint32(0) // The first serial seen is the current server serial
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{in.Answer, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			// A single SOA RR signals "no changes"
			if len(in.Answer) == 1 && isSOAFirst(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}

			// Check if the returned answer is ok
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			// This serial is important
			serial = in.Answer[0].(*SOA).Serial
			first = !first
		}

		// Now we need to check each message for SOA records, to see what we need to do
		if !first {
			t.tsigTimersOnly = true
			// If the last record in the IXFR contains the servers' SOA,  we should quit
			if v, ok := in.Answer[len(in.Answer)-1].(*SOA); ok {
				if v.Serial == serial {
					c <- &Envelope{in.Answer, nil}
					return
				}
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
}

// Out performs an outgoing transfer with the client connecting in w.
// Basic use pattern:
//
//	ch := make(chan *dns.Envelope)
//	tr := new(dns.Transfer)
//	tr.Out(w, r, ch)
//	c <- &dns.Envelope{RR: []dns.RR{soa, rr1, rr2, rr3, soa}}
//	close(ch)
//	w.Hijack()
//	// w.Close() // Client closes connection
//
// The server is responsible for sending the correct sequence of RRs through the
// channel ch.
func (t *Transfer) Out(w ResponseWriter, q *Msg, ch chan *Envelope) error {
	r := new(Msg)
	// Compress?
	r.SetReply(q)
	r.Authoritative = true

	for x := range ch {
		// assume it fits TODO(miek): fix
		r.Answer = append(r.Answer, x.RR...)
		if err := w.WriteMsg(r); err != nil {
			return err
		}
	}
	w.TsigTimersOnly(true)
	r.Answer = nil
	return nil
}

// ReadMsg reads a message from the transfer connection t.
func (t *Transfer) ReadMsg() (*Msg, error) {
	m := new(Msg)
	p := make([]byte, MaxMsgSize)
	n, err := t.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	if ts := m.IsTsig(); ts != nil && t.TsigSecret != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	}
	return m, err
}

// WriteMsg writes a message through the transfer connection t.
func (t *Transfer) WriteMsg(m *Msg) (err error) {
	var out []byte
	if ts := m.IsTsig(); ts != nil && t.TsigSecret != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, t.tsigRequestMAC, err = TsigGenerate(m, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	if _, err = t.Write(out); err != nil {
		return err
	}
	return nil
}

func isSOAFirst(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[0].Header().Rrtype == TypeSOA
	}
	return false
}

func isSOALast(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
	}
	return false
}
<commit_msg>xfr: Transfer.Out use a new Msg every time<commit_after>package dns

import (
	"time"
)

// Envelope is used when doing a zone transfer with a remote server.
type Envelope struct {
	RR    []RR  // The set of RRs in the answer section of the xfr reply message.
	Error error // If something went wrong, this contains the error.
}

// A Transfer defines parameters that are used during a zone transfer.
type Transfer struct {
	*Conn
	DialTimeout    time.Duration     // net.DialTimeout (ns), defaults to 2 * 1e9
	ReadTimeout    time.Duration     // net.Conn.SetReadTimeout value for connections (ns), defaults to 2 * 1e9
	WriteTimeout   time.Duration     // net.Conn.SetWriteTimeout value for connections (ns), defaults to 2 * 1e9
	TsigSecret     map[string]string // Secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
	tsigTimersOnly bool
}

// Think we need to away to stop the transfer

// In performs an incoming transfer with the server in a.
func (t *Transfer) In(q *Msg, a string) (env chan *Envelope, err error) {
	timeout := dnsTimeout
	if t.DialTimeout != 0 {
		timeout = t.DialTimeout
	}
	t.Conn, err = DialTimeout("tcp", a, timeout)
	if err != nil {
		return nil, err
	}
	if err := t.WriteMsg(q); err != nil {
		return nil, err
	}
	env = make(chan *Envelope)
	go func() {
		if q.Question[0].Qtype == TypeAXFR {
			go t.inAxfr(q.Id, env)
			return
		}
		if q.Question[0].Qtype == TypeIXFR {
			go t.inIxfr(q.Id, env)
			return
		}
	}()
	return env, nil
}

func (t *Transfer) inAxfr(id uint16, c chan *Envelope) {
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.Conn.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{nil, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			first = !first
			// only one answer that is SOA, receive more
			if len(in.Answer) == 1 {
				t.tsigTimersOnly = true
				c <- &Envelope{in.Answer, nil}
				continue
			}
		}

		if !first {
			t.tsigTimersOnly = true // Subsequent envelopes use this.
			if isSOALast(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
	panic("dns: not reached")
}

func (t *Transfer) inIxfr(id uint16, c chan *Envelope) {
	serial := uint32(0) // The first serial seen is the current server serial
	first := true
	defer t.Close()
	defer close(c)
	timeout := dnsTimeout
	if t.ReadTimeout != 0 {
		timeout = t.ReadTimeout
	}
	for {
		t.SetReadDeadline(time.Now().Add(timeout))
		in, err := t.ReadMsg()
		if err != nil {
			c <- &Envelope{in.Answer, err}
			return
		}
		if id != in.Id {
			c <- &Envelope{in.Answer, ErrId}
			return
		}
		if first {
			// A single SOA RR signals "no changes"
			if len(in.Answer) == 1 && isSOAFirst(in) {
				c <- &Envelope{in.Answer, nil}
				return
			}

			// Check if the returned answer is ok
			if !isSOAFirst(in) {
				c <- &Envelope{in.Answer, ErrSoa}
				return
			}
			// This serial is important
			serial = in.Answer[0].(*SOA).Serial
			first = !first
		}

		// Now we need to check each message for SOA records, to see what we need to do
		if !first {
			t.tsigTimersOnly = true
			// If the last record in the IXFR contains the servers' SOA,  we should quit
			if v, ok := in.Answer[len(in.Answer)-1].(*SOA); ok {
				if v.Serial == serial {
					c <- &Envelope{in.Answer, nil}
					return
				}
			}
			c <- &Envelope{in.Answer, nil}
		}
	}
}

// Out performs an outgoing transfer with the client connecting in w.
// Basic use pattern:
//
//	ch := make(chan *dns.Envelope)
//	tr := new(dns.Transfer)
//	tr.Out(w, r, ch)
//	c <- &dns.Envelope{RR: []dns.RR{soa, rr1, rr2, rr3, soa}}
//	close(ch)
//	w.Hijack()
//	// w.Close() // Client closes connection
//
// The server is responsible for sending the correct sequence of RRs through the
// channel ch.
func (t *Transfer) Out(w ResponseWriter, q *Msg, ch chan *Envelope) error {
	for x := range ch {
		r := new(Msg)
		// Compress?
		r.SetReply(q)
		r.Authoritative = true
		// assume it fits TODO(miek): fix
		r.Answer = append(r.Answer, x.RR...)
		if err := w.WriteMsg(r); err != nil {
			return err
		}
	}
	w.TsigTimersOnly(true)
	return nil
}

// ReadMsg reads a message from the transfer connection t.
func (t *Transfer) ReadMsg() (*Msg, error) {
	m := new(Msg)
	p := make([]byte, MaxMsgSize)
	n, err := t.Read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	if ts := m.IsTsig(); ts != nil && t.TsigSecret != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		err = TsigVerify(p, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	}
	return m, err
}

// WriteMsg writes a message through the transfer connection t.
func (t *Transfer) WriteMsg(m *Msg) (err error) {
	var out []byte
	if ts := m.IsTsig(); ts != nil && t.TsigSecret != nil {
		if _, ok := t.TsigSecret[ts.Hdr.Name]; !ok {
			return ErrSecret
		}
		out, t.tsigRequestMAC, err = TsigGenerate(m, t.TsigSecret[ts.Hdr.Name], t.tsigRequestMAC, t.tsigTimersOnly)
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	if _, err = t.Write(out); err != nil {
		return err
	}
	return nil
}

func isSOAFirst(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[0].Header().Rrtype == TypeSOA
	}
	return false
}

func isSOALast(in *Msg) bool {
	if len(in.Answer) > 0 {
		return in.Answer[len(in.Answer)-1].Header().Rrtype == TypeSOA
	}
	return false
}
<|endoftext|>


<commit_before>// Yam (Yet Another Mux)

package main

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
)

type handler func(http.ResponseWriter, *http.Request)

type Yam struct {
	Root *Route
}

func NewYam() *Yam {
	return &Yam{
		Root: &Route{},
	}
}

func route(path string, router *Route) *Route {
	parts := strings.Split(path, "/")[1:]
	routes := router.Routes

	fmt.Println("Start Router:", router.path)
	fmt.Println("Stat Path:", path)
	fullPath := router.path + path

	for i, part := range parts {
		fmt.Println("Part:", part)
		if i == len(parts)-1 {

			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Route Exists")
					fmt.Println("--------------")
					return route
				}
			}

			route := &Route{leaf: part, path: fullPath}
			fmt.Println("Add:", route.path)
			fmt.Println("Router:", router.path)
			router.Routes = append(router.Routes, route)

			fmt.Println("--------------")

			return route

		} else {
			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Leaf:", route.leaf)
					router = route
				} else {
					fmt.Println("Router:", router.path)
					route := &Route{leaf: part, path: router.path + path}
					router.Routes = append(router.Routes, route)
					router = route
				}
			}

		}
	}

	return nil
}

func (y *Yam) Route(path string) *Route {
	return route(path, y.Root)
}

func (y *Yam) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")[1:]
	fmt.Println(parts)
	routes := y.Root.Routes

	for i, part := range parts {
		fmt.Println(part)
		for _, route := range routes {
			fmt.Println("Leaf:", route.leaf)
			match := false
			// Pattern Match
			if strings.HasPrefix(route.leaf, ":") {
				fmt.Println("Pattern Match")
				match = true
				values := url.Values{}
				values.Add(route.leaf, part)
				r.URL.RawQuery = url.Values(values).Encode() + "&" + r.URL.RawQuery
			} else { // Exact match
				fmt.Println("Exact Match")
				if route.leaf == part {
					match = true
				}
			}

			if match {
				fmt.Println("Leaf ==", part)
				if i < len(parts)-1 {
					routes = route.Routes
					break
				} else {
					fmt.Println("Found: ", route.path)

					var handler http.Handler
					switch r.Method {
					case "GET":
						handler = route.getHandler
					case "POST":
						handler = route.postHandler
					case "PUT":
						handler = route.putHandler
					case "PATCH":
						handler = route.patchHandler
					case "DELETE":
						handler = route.deleteHandler
					}

					if handler != nil {
						handler.ServeHTTP(w, r)
						return
					}

					fmt.Println("No handler for method")
					w.WriteHeader(http.StatusMethodNotAllowed)
					return

				}
			}
		}
	}

	// If we get here then we have not found a route
	fmt.Println("Not Found")
	w.WriteHeader(http.StatusNotFound)
}

type Route struct {
	leaf   string   // a part of a URL path, /foo/bar - a leaf would be foo and bar
	path   string   // full url path
	Routes []*Route // Routes that live under this route

	getHandler    http.Handler
	postHandler   http.Handler
	putHandler    http.Handler
	patchHandler  http.Handler
	deleteHandler http.Handler
}

func (r *Route) Route(path string) *Route {
	return route(path, r)
}

func (r *Route) Get(h handler) *Route {
	r.getHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Post(h handler) *Route {
	r.postHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Put(h handler) *Route {
	r.putHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Delete(h handler) *Route {
	r.deleteHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Patch(h handler) *Route {
	r.patchHandler = http.HandlerFunc(h)

	return r
}

func GetRootHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Root Handler"))
}

func GetFooHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Foo Handler"))
}

func GetAHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get A Handler"))
}

func GetBHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get B Handler"))
}

func GetCHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get C Handler"))
}

func GetDHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get D Handler"))
}

func main() {
	y := NewYam()

	y.Route("/").Get(GetRootHandler)
	y.Route("/get").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("GET"))
	})
	y.Route("/post").Post(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("POST"))
	})
	y.Route("/put").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT"))
	})
	y.Route("/patch").Patch(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PATCH"))
	})
	y.Route("/delete").Delete(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("DELETE"))
	})

	a := y.Route("/a").Get(GetAHandler)
	a.Route("/b").Get(GetBHandler)
	a.Route("/b").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT B Handler"))
	})
	c := a.Route("/b/c").Get(GetCHandler)
	c.Route("/d").Get(GetDHandler)
	e := c.Route("/d/e").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("E Handler"))
	})
	e.Route("/f").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("F Handler"))
	})

	// Pattern Matching
	a.Route("/:foo").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("A :foo Handler\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
	})

	fmt.Printf("%+v\n", y)

	http.ListenAndServe(":5000", y)
}
<commit_msg>More test usage w/ pattern matching<commit_after>// Yam (Yet Another Mux)

package main

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
)

type handler func(http.ResponseWriter, *http.Request)

type Yam struct {
	Root *Route
}

func New() *Yam {
	return &Yam{
		Root: &Route{},
	}
}

func route(path string, router *Route) *Route {
	parts := strings.Split(path, "/")[1:]
	routes := router.Routes

	fmt.Println("Start Router:", router.path)
	fmt.Println("Stat Path:", path)
	fullPath := router.path + path

	for i, part := range parts {
		fmt.Println("Part:", part)
		if i == len(parts)-1 {

			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Route Exists")
					fmt.Println("--------------")
					return route
				}
			}

			route := &Route{leaf: part, path: fullPath}
			fmt.Println("Add:", route.path)
			fmt.Println("Router:", router.path)
			router.Routes = append(router.Routes, route)

			fmt.Println("--------------")

			return route

		} else {
			for _, route := range routes {
				if route.leaf == part {
					fmt.Println("Leaf:", route.leaf)
					router = route
				} else {
					fmt.Println("Router:", router.path)
					route := &Route{leaf: part, path: router.path + path}
					router.Routes = append(router.Routes, route)
					router = route
				}
			}

		}
	}

	return nil
}

func (y *Yam) Route(path string) *Route {
	return route(path, y.Root)
}

func (y *Yam) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	parts := strings.Split(r.URL.Path, "/")[1:]
	fmt.Println(parts)
	routes := y.Root.Routes

	for i, part := range parts {
		fmt.Println(part)
		for _, route := range routes {
			fmt.Println("Leaf:", route.leaf)
			match := false
			// Pattern Match
			if strings.HasPrefix(route.leaf, ":") {
				fmt.Println("Pattern Match")
				match = true
				values := url.Values{}
				values.Add(route.leaf, part)
				r.URL.RawQuery = url.Values(values).Encode() + "&" + r.URL.RawQuery
			} else { // Exact match
				fmt.Println("Exact Match")
				if route.leaf == part {
					match = true
				}
			}

			if match {
				fmt.Println("Leaf ==", part)
				if i < len(parts)-1 {
					routes = route.Routes
					break
				} else {
					fmt.Println("Found: ", route.path)

					var handler http.Handler
					switch r.Method {
					case "GET":
						handler = route.getHandler
					case "POST":
						handler = route.postHandler
					case "PUT":
						handler = route.putHandler
					case "PATCH":
						handler = route.patchHandler
					case "DELETE":
						handler = route.deleteHandler
					}

					if handler != nil {
						handler.ServeHTTP(w, r)
						return
					}

					fmt.Println("No handler for method")
					w.WriteHeader(http.StatusMethodNotAllowed)
					return

				}
			}
		}
	}

	// If we get here then we have not found a route
	fmt.Println("Not Found")
	w.WriteHeader(http.StatusNotFound)
}

type Route struct {
	leaf   string   // a part of a URL path, /foo/bar - a leaf would be foo and bar
	path   string   // full url path
	Routes []*Route // Routes that live under this route

	getHandler    http.Handler
	postHandler   http.Handler
	putHandler    http.Handler
	patchHandler  http.Handler
	deleteHandler http.Handler
}

func (r *Route) Route(path string) *Route {
	return route(path, r)
}

func (r *Route) Get(h handler) *Route {
	r.getHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Post(h handler) *Route {
	r.postHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Put(h handler) *Route {
	r.putHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Delete(h handler) *Route {
	r.deleteHandler = http.HandlerFunc(h)

	return r
}

func (r *Route) Patch(h handler) *Route {
	r.patchHandler = http.HandlerFunc(h)

	return r
}

func GetRootHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Root Handler"))
}

func GetFooHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get Foo Handler"))
}

func GetAHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get A Handler"))
}

func GetBHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get B Handler"))
}

func GetCHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get C Handler"))
}

func GetDHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Get D Handler"))
}

func main() {
	y := New()

	y.Route("/").Get(GetRootHandler)
	y.Route("/get").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("GET"))
	})
	y.Route("/post").Post(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("POST"))
	})
	y.Route("/put").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT"))
	})
	y.Route("/patch").Patch(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PATCH"))
	})
	y.Route("/delete").Delete(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("DELETE"))
	})

	a := y.Route("/a").Get(GetAHandler)
	a.Route("/b").Get(GetBHandler)
	a.Route("/b").Put(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("PUT B Handler"))
	})
	c := a.Route("/b/c").Get(GetCHandler)
	c.Route("/d").Get(GetDHandler)
	e := c.Route("/d/e").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("E Handler"))
	})
	e.Route("/f").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("F Handler"))
	})

	// Pattern Matching
	a.Route("/:foo").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("A :foo Handler\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
	})

	bar := a.Route("/:foo/:bar").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("/a/:foo/:bar Handler\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
		w.Write([]byte("\n"))
		w.Write([]byte(r.URL.Query().Get(":bar")))
	})

	bar.Route("/baz").Get(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("baz\n"))
		w.Write([]byte(r.URL.Query().Get(":foo")))
		w.Write([]byte("\n"))
		w.Write([]byte(r.URL.Query().Get(":bar")))
	})

	fmt.Printf("%+v\n", y)

	http.ListenAndServe(":5000", y)
}
<|endoftext|>


<commit_before>package server

import (
	"net/url"
	"testing"

	"github.com/moov-io/ach"
)

func testNextID(tb testing.TB) string {
	id := NextID()
	if len(id) != 16 {
		tb.Errorf("got other length %d for ID %s", len(id), id)
	}
	return id
}

func TestNextID(t *testing.T) {
	id := testNextID(t)

	_, err := url.Parse("https://moov.io/" + id)
	if err != nil {
		t.Fatalf("failed to parse url with id of %q: %v", id, err)
	}
}

func BenchmarkNextID(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testNextID(b)
	}
}

// test mocks are in mock_test.go

// CreateFile tests
func TestCreateFile(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(mockFileHeader())
	if id != "12345" {
		t.Errorf("expected %s received %s w/ error %s", "12345", id, err)
	}
}
func TestCreateFileIDExists(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(ach.FileHeader{ID: "98765"})
	if err != ErrAlreadyExists {
		t.Errorf("expected %s received %s w/ error %s", "ErrAlreadyExists", id, err)
	}
}

func TestCreateFileNoID(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(ach.NewFileHeader())
	if len(id) < 3 {
		t.Errorf("expected %s received %s w/ error %s", "NextID", id, err)
	}
}

// Service.GetFile tests

func TestGetFile(t *testing.T) {
	s := mockServiceInMemory()
	f, err := s.GetFile("98765")
	if err != nil {
		t.Errorf("expected %s received %s w/ error %s", "98765", f.ID, err)
	}
}

func TestGetFileNotFound(t *testing.T) {
	s := mockServiceInMemory()
	f, err := s.GetFile("12345")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", f.ID, err)
	}
}

// Service.GetFiles tests

func TestGetFiles(t *testing.T) {
	s := mockServiceInMemory()
	files := s.GetFiles()
	if len(files) != 1 {
		t.Errorf("expected %s received %v", "1", len(files))
	}
}

// Service.DeleteFile tests

func TestDeleteFile(t *testing.T) {
	s := mockServiceInMemory()
	err := s.DeleteFile("98765")
	if err != nil {
		t.Errorf("expected %s received %s", "nil", err)
	}
	_, err = s.GetFile("98765")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s", "ErrNotFound", err)
	}
}

// Service.CreateBatch tests

// TestCreateBatch tests creating a new batch when file.ID exists and batch.id does not exist
func TestCreateBatch(t *testing.T) {
	s := mockServiceInMemory()
	bh := mockBatchHeaderWeb()
	bh.ID = "11111"
	id, err := s.CreateBatch("98765", *bh)
	if id != "11111" {
		t.Errorf("expected %s received %s w/ error %s", "11111", id, err)
	}
}

// TestCreateBatchIDExists Create a new batch with batch.id already present. Should fail.
func TestCreateBatchIDExists(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateBatch("98765", *mockBatchHeaderWeb())
	if err != ErrAlreadyExists {
		t.Errorf("expected %s received %s w/ error %s", "ErrAlreadyExists", id, err)
	}
}

// TestCreateBatchFileIDExits create a batch when the file.id does not exist. Should fail.
func TestCreateBatchFileIDExits(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateBatch("55555", *mockBatchHeaderWeb())
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", id, err)
	}

}

// TestCreateBatchIDBank create a new batch when the batch.id is nil but file.id is valid. Should generate batch.id and save.
func TestCreateBatchIDBlank(t *testing.T) {
	s := mockServiceInMemory()
	bh := mockBatchHeaderWeb()
	bh.ID = ""
	id, err := s.CreateBatch("98765", *bh)
	if len(id) < 3 {
		t.Errorf("expected %s received %s w/ error %s", "NextID", id, err)
	}
}

// Service.GetBatch

// TestGetBatch return a batch for the existing file.id and batch.id
func TestGetBatch(t *testing.T) {
	s := mockServiceInMemory()
	b, err := s.GetBatch("98765", "54321")
	if b.ID() != "54321" {
		t.Errorf("expected %s received %s w/ error %s", "54321", b.ID(), err)
	}
}

// TestGetBatchNotFound return a failure if the batch.id is not found
func TestGetBatchNotFound(t *testing.T) {
	s := mockServiceInMemory()
	b, err := s.GetBatch("98765", "55555")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", b.ID(), err)
	}
}

// Service.GetBatches

// TestGetBatches return a list of batches for the supplied file.id
func TestGetBatches(t *testing.T) {
	s := mockServiceInMemory()
	batches := s.GetBatches("98765")
	if len(batches) != 1 {
		t.Errorf("expected %s received %v", "1", len(batches))
	}
}

// Service.DeleteBatch

// TestDeleteBatch removes a batch with existing file and batch id.
func TestDeleteBatch(t *testing.T) {
	s := mockServiceInMemory()
	err := s.DeleteBatch("98765", "54321")
	if err != nil {
		t.Errorf("expected %s received error %s", "nil", err)
	}
}
<commit_msg>server: ensure ID's are 16 characters<commit_after>package server

import (
	"net/url"
	"testing"
	"unicode/utf8"

	"github.com/moov-io/ach"
)

func testNextID(tb testing.TB) string {
	id := NextID()
	if utf8.RuneCountInString(id) != 16 {
		tb.Errorf("got other length %d for ID %s", len(id), id)
	}
	return id
}

func TestNextID(t *testing.T) {
	id := testNextID(t)

	_, err := url.Parse("https://moov.io/" + id)
	if err != nil {
		t.Fatalf("failed to parse url with id of %q: %v", id, err)
	}
}

func BenchmarkNextID(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		testNextID(b)
	}
}

// test mocks are in mock_test.go

// CreateFile tests
func TestCreateFile(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(mockFileHeader())
	if id != "12345" {
		t.Errorf("expected %s received %s w/ error %s", "12345", id, err)
	}
}
func TestCreateFileIDExists(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(ach.FileHeader{ID: "98765"})
	if err != ErrAlreadyExists {
		t.Errorf("expected %s received %s w/ error %s", "ErrAlreadyExists", id, err)
	}
}

func TestCreateFileNoID(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateFile(ach.NewFileHeader())
	if len(id) < 3 {
		t.Errorf("expected %s received %s w/ error %s", "NextID", id, err)
	}
}

// Service.GetFile tests

func TestGetFile(t *testing.T) {
	s := mockServiceInMemory()
	f, err := s.GetFile("98765")
	if err != nil {
		t.Errorf("expected %s received %s w/ error %s", "98765", f.ID, err)
	}
}

func TestGetFileNotFound(t *testing.T) {
	s := mockServiceInMemory()
	f, err := s.GetFile("12345")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", f.ID, err)
	}
}

// Service.GetFiles tests

func TestGetFiles(t *testing.T) {
	s := mockServiceInMemory()
	files := s.GetFiles()
	if len(files) != 1 {
		t.Errorf("expected %s received %v", "1", len(files))
	}
}

// Service.DeleteFile tests

func TestDeleteFile(t *testing.T) {
	s := mockServiceInMemory()
	err := s.DeleteFile("98765")
	if err != nil {
		t.Errorf("expected %s received %s", "nil", err)
	}
	_, err = s.GetFile("98765")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s", "ErrNotFound", err)
	}
}

// Service.CreateBatch tests

// TestCreateBatch tests creating a new batch when file.ID exists and batch.id does not exist
func TestCreateBatch(t *testing.T) {
	s := mockServiceInMemory()
	bh := mockBatchHeaderWeb()
	bh.ID = "11111"
	id, err := s.CreateBatch("98765", *bh)
	if id != "11111" {
		t.Errorf("expected %s received %s w/ error %s", "11111", id, err)
	}
}

// TestCreateBatchIDExists Create a new batch with batch.id already present. Should fail.
func TestCreateBatchIDExists(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateBatch("98765", *mockBatchHeaderWeb())
	if err != ErrAlreadyExists {
		t.Errorf("expected %s received %s w/ error %s", "ErrAlreadyExists", id, err)
	}
}

// TestCreateBatchFileIDExits create a batch when the file.id does not exist. Should fail.
func TestCreateBatchFileIDExits(t *testing.T) {
	s := mockServiceInMemory()
	id, err := s.CreateBatch("55555", *mockBatchHeaderWeb())
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", id, err)
	}

}

// TestCreateBatchIDBank create a new batch when the batch.id is nil but file.id is valid. Should generate batch.id and save.
func TestCreateBatchIDBlank(t *testing.T) {
	s := mockServiceInMemory()
	bh := mockBatchHeaderWeb()
	bh.ID = ""
	id, err := s.CreateBatch("98765", *bh)
	if len(id) < 3 {
		t.Errorf("expected %s received %s w/ error %s", "NextID", id, err)
	}
}

// Service.GetBatch

// TestGetBatch return a batch for the existing file.id and batch.id
func TestGetBatch(t *testing.T) {
	s := mockServiceInMemory()
	b, err := s.GetBatch("98765", "54321")
	if b.ID() != "54321" {
		t.Errorf("expected %s received %s w/ error %s", "54321", b.ID(), err)
	}
}

// TestGetBatchNotFound return a failure if the batch.id is not found
func TestGetBatchNotFound(t *testing.T) {
	s := mockServiceInMemory()
	b, err := s.GetBatch("98765", "55555")
	if err != ErrNotFound {
		t.Errorf("expected %s received %s w/ error %s", "ErrNotFound", b.ID(), err)
	}
}

// Service.GetBatches

// TestGetBatches return a list of batches for the supplied file.id
func TestGetBatches(t *testing.T) {
	s := mockServiceInMemory()
	batches := s.GetBatches("98765")
	if len(batches) != 1 {
		t.Errorf("expected %s received %v", "1", len(batches))
	}
}

// Service.DeleteBatch

// TestDeleteBatch removes a batch with existing file and batch id.
func TestDeleteBatch(t *testing.T) {
	s := mockServiceInMemory()
	err := s.DeleteBatch("98765", "54321")
	if err != nil {
		t.Errorf("expected %s received error %s", "nil", err)
	}
}
<|endoftext|>


<commit_before>package lfs

import (
	"encoding/hex"
	"regexp"

	"github.com/git-lfs/git-lfs/git"
)

var z40 = regexp.MustCompile(`\^?0{40}`)

type lockableNameSet struct {
	opt *ScanRefsOptions
	set GitScannerSet
}

// Determines if the given blob sha matches a locked file.
func (s *lockableNameSet) Check(blobSha string) (string, bool) {
	if s == nil || s.opt == nil || s.set == nil {
		return "", false
	}

	name, ok := s.opt.GetName(blobSha)
	if !ok {
		return name, ok
	}

	if s.set.Contains(name) {
		return name, true
	}
	return name, false
}

func noopFoundLockable(name string) {}

// scanRefsToChan takes a ref and returns a channel of WrappedPointer objects
// for all Git LFS pointers it finds for that ref.
// Reports unique oids once only, not multiple times if >1 file uses the same content
func scanRefsToChan(scanner *GitScanner, pointerCb GitScannerFoundPointer, refLeft, refRight string, opt *ScanRefsOptions) error {
	if opt == nil {
		panic("no scan ref options")
	}

	revs, err := revListShas(refLeft, refRight, opt)
	if err != nil {
		return err
	}

	lockableSet := &lockableNameSet{opt: opt, set: scanner.PotentialLockables}
	smallShas, batchLockableCh, err := catFileBatchCheck(revs, lockableSet)
	if err != nil {
		return err
	}

	lockableCb := scanner.FoundLockable
	if lockableCb == nil {
		lockableCb = noopFoundLockable
	}

	go func(cb GitScannerFoundLockable, ch chan string) {
		for name := range ch {
			cb(name)
		}
	}(lockableCb, batchLockableCh)

	pointers, checkLockableCh, err := catFileBatch(smallShas, lockableSet)
	if err != nil {
		return err
	}

	for p := range pointers.Results {
		if name, ok := opt.GetName(p.Sha1); ok {
			p.Name = name
		}
		pointerCb(p, nil)
	}

	for lockableName := range checkLockableCh {
		lockableCb(lockableName)
	}

	if err := pointers.Wait(); err != nil {
		pointerCb(nil, err)
	}

	return nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(refLeft, refRight string, opt *ScanRefsOptions) (*StringChannelWrapper, error) {
	scanner, err := git.NewRevListScanner(refLeft, refRight, &git.ScanRefsOptions{
		Mode:             git.ScanningMode(opt.ScanMode),
		Remote:           opt.RemoteName,
		SkipDeletedBlobs: opt.SkipDeletedBlobs,
		SkippedRefs:      opt.skippedRefs,
		Mutex:            opt.mutex,
		Names:            opt.nameMap,
	})

	if err != nil {
		return nil, err
	}

	revs := make(chan string, chanBufSize)
	errs := make(chan error, 5) // may be multiple errors

	go func() {
		for scanner.Scan() {
			sha := hex.EncodeToString(scanner.OID())
			if name := scanner.Name(); len(name) > 0 {
				opt.SetName(sha, name)
			}
			revs <- sha
		}

		if err = scanner.Err(); err != nil {
			errs <- err
		}

		if err = scanner.Close(); err != nil {
			errs <- err
		}

		close(revs)
		close(errs)
	}()

	return NewStringChannelWrapper(revs, errs), nil
}

// Get additional arguments needed to limit 'git rev-list' to just the changes
// in refTo that are also not on remoteName.
//
// Returns a slice of string command arguments, and a slice of string git
// commits to pass to `git rev-list` via STDIN.
func revListArgsRefVsRemote(refTo, remoteName string, skippedRefs []string) ([]string, []string) {
	if len(skippedRefs) < 1 {
		// Safe to use cached
		return []string{refTo, "--not", "--remotes=" + remoteName}, nil
	}

	// Use only the non-missing refs as 'from' points
	commits := make([]string, 1, len(skippedRefs)+1)
	commits[0] = refTo
	return []string{"--stdin"}, append(commits, skippedRefs...)
}
<commit_msg>lfs: remove unused revListArgsRefVsRemote()<commit_after>package lfs

import (
	"encoding/hex"
	"regexp"

	"github.com/git-lfs/git-lfs/git"
)

var z40 = regexp.MustCompile(`\^?0{40}`)

type lockableNameSet struct {
	opt *ScanRefsOptions
	set GitScannerSet
}

// Determines if the given blob sha matches a locked file.
func (s *lockableNameSet) Check(blobSha string) (string, bool) {
	if s == nil || s.opt == nil || s.set == nil {
		return "", false
	}

	name, ok := s.opt.GetName(blobSha)
	if !ok {
		return name, ok
	}

	if s.set.Contains(name) {
		return name, true
	}
	return name, false
}

func noopFoundLockable(name string) {}

// scanRefsToChan takes a ref and returns a channel of WrappedPointer objects
// for all Git LFS pointers it finds for that ref.
// Reports unique oids once only, not multiple times if >1 file uses the same content
func scanRefsToChan(scanner *GitScanner, pointerCb GitScannerFoundPointer, refLeft, refRight string, opt *ScanRefsOptions) error {
	if opt == nil {
		panic("no scan ref options")
	}

	revs, err := revListShas(refLeft, refRight, opt)
	if err != nil {
		return err
	}

	lockableSet := &lockableNameSet{opt: opt, set: scanner.PotentialLockables}
	smallShas, batchLockableCh, err := catFileBatchCheck(revs, lockableSet)
	if err != nil {
		return err
	}

	lockableCb := scanner.FoundLockable
	if lockableCb == nil {
		lockableCb = noopFoundLockable
	}

	go func(cb GitScannerFoundLockable, ch chan string) {
		for name := range ch {
			cb(name)
		}
	}(lockableCb, batchLockableCh)

	pointers, checkLockableCh, err := catFileBatch(smallShas, lockableSet)
	if err != nil {
		return err
	}

	for p := range pointers.Results {
		if name, ok := opt.GetName(p.Sha1); ok {
			p.Name = name
		}
		pointerCb(p, nil)
	}

	for lockableName := range checkLockableCh {
		lockableCb(lockableName)
	}

	if err := pointers.Wait(); err != nil {
		pointerCb(nil, err)
	}

	return nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(refLeft, refRight string, opt *ScanRefsOptions) (*StringChannelWrapper, error) {
	scanner, err := git.NewRevListScanner(refLeft, refRight, &git.ScanRefsOptions{
		Mode:             git.ScanningMode(opt.ScanMode),
		Remote:           opt.RemoteName,
		SkipDeletedBlobs: opt.SkipDeletedBlobs,
		SkippedRefs:      opt.skippedRefs,
		Mutex:            opt.mutex,
		Names:            opt.nameMap,
	})

	if err != nil {
		return nil, err
	}

	revs := make(chan string, chanBufSize)
	errs := make(chan error, 5) // may be multiple errors

	go func() {
		for scanner.Scan() {
			sha := hex.EncodeToString(scanner.OID())
			if name := scanner.Name(); len(name) > 0 {
				opt.SetName(sha, name)
			}
			revs <- sha
		}

		if err = scanner.Err(); err != nil {
			errs <- err
		}

		if err = scanner.Close(); err != nil {
			errs <- err
		}

		close(revs)
		close(errs)
	}()

	return NewStringChannelWrapper(revs, errs), nil
}
<|endoftext|>


<commit_before>// Copyright 2016-2019 Authors of Cilium
//
// 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.

// +build !privileged_tests

package main

import (
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/datapath"
	fakedatapath "github.com/cilium/cilium/pkg/datapath/fake"
	"github.com/cilium/cilium/pkg/endpoint/regeneration"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/metrics"
	monitorAPI "github.com/cilium/cilium/pkg/monitor/api"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/revert"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type DaemonSuite struct {
	d *Daemon

	// oldPolicyEnabled is the policy enforcement mode that was set before the test,
	// as returned by policy.GetPolicyEnabled().
	oldPolicyEnabled string

	kvstoreInit bool

	// Owners interface mock
	OnTracingEnabled          func() bool
	OnAlwaysAllowLocalhost    func() bool
	OnGetCachedLabelList      func(id identity.NumericIdentity) (labels.LabelArray, error)
	OnGetPolicyRepository     func() *policy.Repository
	OnUpdateProxyRedirect     func(e regeneration.EndpointUpdater, l4 *policy.L4Filter, proxyWaitGroup *completion.WaitGroup) (uint16, error, revert.FinalizeFunc, revert.RevertFunc)
	OnRemoveProxyRedirect     func(e regeneration.EndpointInfoSource, id string, proxyWaitGroup *completion.WaitGroup) (error, revert.FinalizeFunc, revert.RevertFunc)
	OnUpdateNetworkPolicy     func(e regeneration.EndpointUpdater, policy *policy.L4Policy, proxyWaitGroup *completion.WaitGroup) (error, revert.RevertFunc)
	OnRemoveNetworkPolicy     func(e regeneration.EndpointInfoSource)
	OnQueueEndpointBuild      func(ctx context.Context, epID uint64) (func(), error)
	OnRemoveFromEndpointQueue func(epID uint64)
	OnDebugEnabled            func() bool
	OnGetCompilationLock      func() *lock.RWMutex
	OnSendNotification        func(typ monitorAPI.AgentNotification, text string) error
	OnNewProxyLogRecord       func(l *accesslog.LogRecord) error
	OnClearPolicyConsumers    func(id uint16) *sync.WaitGroup
}

func (ds *DaemonSuite) SetUpTest(c *C) {
	option.Config.Populate()
	ds.oldPolicyEnabled = policy.GetPolicyEnabled()
	policy.SetPolicyEnabled(option.DefaultEnforcement)

	// kvstore is initialized before generic SetUpTest so it must have been completed
	ds.kvstoreInit = true

	time.Local = time.UTC
	tempRunDir, err := ioutil.TempDir("", "cilium-test-run")
	c.Assert(err, IsNil)
	err = os.Mkdir(filepath.Join(tempRunDir, "globals"), 0777)
	c.Assert(err, IsNil)

	option.Config.DryMode = true
	option.Config.Opts = option.NewIntOptions(&option.DaemonMutableOptionLibrary)
	option.Config.Device = "undefined"
	option.Config.RunDir = tempRunDir
	option.Config.StateDir = tempRunDir
	option.Config.AccessLog = filepath.Join(tempRunDir, "cilium-access.log")

	// GetConfig the default labels prefix filter
	err = labels.ParseLabelPrefixCfg(nil, "")
	c.Assert(err, IsNil)
	option.Config.Opts.SetBool(option.DropNotify, true)
	option.Config.Opts.SetBool(option.TraceNotify, true)

	// Disable restore of host IPs for unit tests. There can be arbitrary
	// state left on disk.
	option.Config.EnableHostIPRestore = false

	d, _, err := NewDaemon(fakedatapath.NewDatapath())
	c.Assert(err, IsNil)
	ds.d = d

	kvstore.DeletePrefix(common.OperationalPath)
	kvstore.DeletePrefix(kvstore.BaseKeyPrefix)

	ds.OnTracingEnabled = nil
	ds.OnAlwaysAllowLocalhost = nil
	ds.OnGetCachedLabelList = nil
	ds.OnGetPolicyRepository = nil
	ds.OnUpdateProxyRedirect = nil
	ds.OnRemoveProxyRedirect = nil
	ds.OnUpdateNetworkPolicy = nil
	ds.OnRemoveNetworkPolicy = nil
	ds.OnQueueEndpointBuild = nil
	ds.OnRemoveFromEndpointQueue = nil
	ds.OnDebugEnabled = nil
	ds.OnGetCompilationLock = nil
	ds.OnSendNotification = nil
	ds.OnNewProxyLogRecord = nil
	ds.OnClearPolicyConsumers = nil
}

func (ds *DaemonSuite) TearDownTest(c *C) {
	endpointmanager.RemoveAll()

	if ds.d != nil {
		os.RemoveAll(option.Config.RunDir)
	}

	if ds.kvstoreInit {
		kvstore.DeletePrefix(common.OperationalPath)
		kvstore.DeletePrefix(kvstore.BaseKeyPrefix)
	}

	// Restore the policy enforcement mode.
	policy.SetPolicyEnabled(ds.oldPolicyEnabled)

	// Release the identity allocator reference created by NewDaemon. This
	// is done manually here as we have no Close() function daemon
	cache.Close()

	ds.d.Close()

	_, collectors := metrics.CreateConfiguration(common.MapStringStructToSlice(metrics.DefaultMetrics()))
	for _, collector := range collectors {
		metrics.Unregister(collector)
	}
}

type DaemonEtcdSuite struct {
	DaemonSuite
}

var _ = Suite(&DaemonEtcdSuite{})

func (e *DaemonEtcdSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("etcd")
	e.DaemonSuite.SetUpTest(c)
}

func (e *DaemonEtcdSuite) TearDownTest(c *C) {
	e.DaemonSuite.TearDownTest(c)
}

type DaemonConsulSuite struct {
	DaemonSuite
}

var _ = Suite(&DaemonConsulSuite{})

func (e *DaemonConsulSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("consul")
	e.DaemonSuite.SetUpTest(c)
}

func (e *DaemonConsulSuite) TearDownTest(c *C) {
	e.DaemonSuite.TearDownTest(c)
}

func (ds *DaemonSuite) TestMinimumWorkerThreadsIsSet(c *C) {
	c.Assert(numWorkerThreads() >= 2, Equals, true)
	c.Assert(numWorkerThreads() >= runtime.NumCPU(), Equals, true)
}

func (ds *DaemonSuite) AlwaysAllowLocalhost() bool {
	if ds.OnAlwaysAllowLocalhost != nil {
		return ds.OnAlwaysAllowLocalhost()
	}
	panic("AlwaysAllowLocalhost should not have been called")
}

func (ds *DaemonSuite) GetCachedLabelList(id identity.NumericIdentity) (labels.LabelArray, error) {
	if ds.OnGetCachedLabelList != nil {
		return ds.OnGetCachedLabelList(id)
	}
	panic("GetCachedLabelList should not have been called")
}

func (ds *DaemonSuite) GetPolicyRepository() *policy.Repository {
	if ds.OnGetPolicyRepository != nil {
		return ds.OnGetPolicyRepository()
	}
	panic("GetPolicyRepository should not have been called")
}

func (ds *DaemonSuite) UpdateProxyRedirect(e regeneration.EndpointUpdater, l4 *policy.L4Filter, proxyWaitGroup *completion.WaitGroup) (uint16, error, revert.FinalizeFunc, revert.RevertFunc) {
	if ds.OnUpdateProxyRedirect != nil {
		return ds.OnUpdateProxyRedirect(e, l4, proxyWaitGroup)
	}
	panic("UpdateProxyRedirect should not have been called")
}

func (ds *DaemonSuite) RemoveProxyRedirect(e regeneration.EndpointInfoSource, id string, proxyWaitGroup *completion.WaitGroup) (error, revert.FinalizeFunc, revert.RevertFunc) {
	if ds.OnRemoveProxyRedirect != nil {
		return ds.OnRemoveProxyRedirect(e, id, proxyWaitGroup)
	}
	panic("RemoveProxyRedirect should not have been called")
}

func (ds *DaemonSuite) UpdateNetworkPolicy(e regeneration.EndpointUpdater, policy *policy.L4Policy,
	proxyWaitGroup *completion.WaitGroup) (error, revert.RevertFunc) {
	if ds.OnUpdateNetworkPolicy != nil {
		return ds.OnUpdateNetworkPolicy(e, policy, proxyWaitGroup)
	}
	panic("UpdateNetworkPolicy should not have been called")
}

func (ds *DaemonSuite) RemoveNetworkPolicy(e regeneration.EndpointInfoSource) {
	if ds.OnRemoveNetworkPolicy != nil {
		ds.OnRemoveNetworkPolicy(e)
	}
	panic("RemoveNetworkPolicy should not have been called")
}

func (ds *DaemonSuite) QueueEndpointBuild(ctx context.Context, epID uint64) (func(), error) {
	if ds.OnQueueEndpointBuild != nil {
		return ds.OnQueueEndpointBuild(ctx, epID)
	}
	panic("QueueEndpointBuild should not have been called")
}

func (ds *DaemonSuite) RemoveFromEndpointQueue(epID uint64) {
	if ds.OnRemoveFromEndpointQueue != nil {
		ds.OnRemoveFromEndpointQueue(epID)
		return
	}
	panic("RemoveFromEndpointQueue should not have been called")
}

func (ds *DaemonSuite) DebugEnabled() bool {
	if ds.OnDebugEnabled != nil {
		return ds.OnDebugEnabled()
	}
	panic("DebugEnabled should not have been called")
}

func (ds *DaemonSuite) GetCompilationLock() *lock.RWMutex {
	if ds.OnGetCompilationLock != nil {
		return ds.OnGetCompilationLock()
	}
	panic("GetCompilationLock should not have been called")
}

func (ds *DaemonSuite) SendNotification(typ monitorAPI.AgentNotification, text string) error {
	if ds.OnSendNotification != nil {
		return ds.OnSendNotification(typ, text)
	}
	panic("SendNotification should not have been called")
}

func (ds *DaemonSuite) NewProxyLogRecord(l *accesslog.LogRecord) error {
	if ds.OnNewProxyLogRecord != nil {
		return ds.OnNewProxyLogRecord(l)
	}
	panic("NewProxyLogRecord should not have been called")
}

func (ds *DaemonSuite) Datapath() datapath.Datapath {
	return ds.d.datapath
}

func (ds *DaemonSuite) ClearPolicyConsumers(id uint16) *sync.WaitGroup {
	if ds.OnClearPolicyConsumers != nil {
		return ds.OnClearPolicyConsumers(id)
	}
	panic("ClearPolicyConsumers should not have been called")
}

func (ds *DaemonSuite) GetNodeSuffix() string {
	return ds.d.GetNodeSuffix()
}

func (ds *DaemonSuite) UpdateIdentities(added, deleted cache.IdentityCache) {
	ds.d.UpdateIdentities(added, deleted)
}
<commit_msg>daemon: Use TestMain, SetUpSuite, and SetUpTest<commit_after>// Copyright 2016-2019 Authors of Cilium
//
// 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.

// +build !privileged_tests

package main

import (
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"testing"
	"time"

	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/completion"
	"github.com/cilium/cilium/pkg/datapath"
	fakedatapath "github.com/cilium/cilium/pkg/datapath/fake"
	"github.com/cilium/cilium/pkg/endpoint/regeneration"
	"github.com/cilium/cilium/pkg/endpointmanager"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/lock"
	monitorAPI "github.com/cilium/cilium/pkg/monitor/api"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"
	"github.com/cilium/cilium/pkg/proxy/accesslog"
	"github.com/cilium/cilium/pkg/revert"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

type DaemonSuite struct {
	d *Daemon

	// oldPolicyEnabled is the policy enforcement mode that was set before the test,
	// as returned by policy.GetPolicyEnabled().
	oldPolicyEnabled string

	kvstoreInit bool

	// Owners interface mock
	OnTracingEnabled          func() bool
	OnAlwaysAllowLocalhost    func() bool
	OnGetCachedLabelList      func(id identity.NumericIdentity) (labels.LabelArray, error)
	OnGetPolicyRepository     func() *policy.Repository
	OnUpdateProxyRedirect     func(e regeneration.EndpointUpdater, l4 *policy.L4Filter, proxyWaitGroup *completion.WaitGroup) (uint16, error, revert.FinalizeFunc, revert.RevertFunc)
	OnRemoveProxyRedirect     func(e regeneration.EndpointInfoSource, id string, proxyWaitGroup *completion.WaitGroup) (error, revert.FinalizeFunc, revert.RevertFunc)
	OnUpdateNetworkPolicy     func(e regeneration.EndpointUpdater, policy *policy.L4Policy, proxyWaitGroup *completion.WaitGroup) (error, revert.RevertFunc)
	OnRemoveNetworkPolicy     func(e regeneration.EndpointInfoSource)
	OnQueueEndpointBuild      func(ctx context.Context, epID uint64) (func(), error)
	OnRemoveFromEndpointQueue func(epID uint64)
	OnDebugEnabled            func() bool
	OnGetCompilationLock      func() *lock.RWMutex
	OnSendNotification        func(typ monitorAPI.AgentNotification, text string) error
	OnNewProxyLogRecord       func(l *accesslog.LogRecord) error
	OnClearPolicyConsumers    func(id uint16) *sync.WaitGroup
}

func TestMain(m *testing.M) {
	option.Config.Populate()
	time.Local = time.UTC
	tempRunDir, err := ioutil.TempDir("", "cilium-test-run")
	if err != nil {
		panic("TempDir() failed.")
	}
	err = os.Mkdir(filepath.Join(tempRunDir, "globals"), 0777)
	if err != nil {
		panic("Mkdir failed")
	}

	option.Config.DryMode = true
	option.Config.Opts = option.NewIntOptions(&option.DaemonMutableOptionLibrary)
	option.Config.Device = "undefined"
	option.Config.RunDir = tempRunDir
	option.Config.StateDir = tempRunDir
	option.Config.AccessLog = filepath.Join(tempRunDir, "cilium-access.log")

	// GetConfig the default labels prefix filter
	err = labels.ParseLabelPrefixCfg(nil, "")
	if err != nil {
		panic("ParseLabelPrefixCfg() failed")
	}
	option.Config.Opts.SetBool(option.DropNotify, true)
	option.Config.Opts.SetBool(option.TraceNotify, true)

	// Disable restore of host IPs for unit tests. There can be arbitrary
	// state left on disk.
	option.Config.EnableHostIPRestore = false

	os.Exit(m.Run())
}

func (ds *DaemonSuite) SetUpTest(c *C) {
	ds.oldPolicyEnabled = policy.GetPolicyEnabled()
	policy.SetPolicyEnabled(option.DefaultEnforcement)

	d, _, err := NewDaemon(fakedatapath.NewDatapath())
	c.Assert(err, IsNil)
	ds.d = d

	kvstore.DeletePrefix(common.OperationalPath)
	kvstore.DeletePrefix(kvstore.BaseKeyPrefix)

	ds.OnTracingEnabled = nil
	ds.OnAlwaysAllowLocalhost = nil
	ds.OnGetCachedLabelList = nil
	ds.OnGetPolicyRepository = nil
	ds.OnUpdateProxyRedirect = nil
	ds.OnRemoveProxyRedirect = nil
	ds.OnUpdateNetworkPolicy = nil
	ds.OnRemoveNetworkPolicy = nil
	ds.OnQueueEndpointBuild = nil
	ds.OnRemoveFromEndpointQueue = nil
	ds.OnDebugEnabled = nil
	ds.OnGetCompilationLock = nil
	ds.OnSendNotification = nil
	ds.OnNewProxyLogRecord = nil
	ds.OnClearPolicyConsumers = nil
}

func (ds *DaemonSuite) TearDownTest(c *C) {
	endpointmanager.RemoveAll()

	if ds.d != nil {
		os.RemoveAll(option.Config.RunDir)
	}

	if ds.kvstoreInit {
		kvstore.DeletePrefix(common.OperationalPath)
		kvstore.DeletePrefix(kvstore.BaseKeyPrefix)
	}

	// Restore the policy enforcement mode.
	policy.SetPolicyEnabled(ds.oldPolicyEnabled)

	// Release the identity allocator reference created by NewDaemon. This
	// is done manually here as we have no Close() function daemon
	cache.Close()

	ds.d.Close()
}

type DaemonEtcdSuite struct {
	DaemonSuite
}

var _ = Suite(&DaemonEtcdSuite{})

func (e *DaemonEtcdSuite) SetUpSuite(c *C) {
	kvstore.SetupDummy("etcd")
	e.DaemonSuite.kvstoreInit = true
}

func (e *DaemonEtcdSuite) SetUpTest(c *C) {
	e.DaemonSuite.SetUpTest(c)
}

func (e *DaemonEtcdSuite) TearDownTest(c *C) {
	e.DaemonSuite.TearDownTest(c)
}

type DaemonConsulSuite struct {
	DaemonSuite
}

var _ = Suite(&DaemonConsulSuite{})

func (e *DaemonConsulSuite) SetUpSuite(c *C) {
	kvstore.SetupDummy("consul")
	e.DaemonSuite.kvstoreInit = true
}

func (e *DaemonConsulSuite) SetUpTest(c *C) {
	e.DaemonSuite.SetUpTest(c)
}

func (e *DaemonConsulSuite) TearDownTest(c *C) {
	e.DaemonSuite.TearDownTest(c)
}

func (ds *DaemonSuite) TestMinimumWorkerThreadsIsSet(c *C) {
	c.Assert(numWorkerThreads() >= 2, Equals, true)
	c.Assert(numWorkerThreads() >= runtime.NumCPU(), Equals, true)
}

func (ds *DaemonSuite) AlwaysAllowLocalhost() bool {
	if ds.OnAlwaysAllowLocalhost != nil {
		return ds.OnAlwaysAllowLocalhost()
	}
	panic("AlwaysAllowLocalhost should not have been called")
}

func (ds *DaemonSuite) GetCachedLabelList(id identity.NumericIdentity) (labels.LabelArray, error) {
	if ds.OnGetCachedLabelList != nil {
		return ds.OnGetCachedLabelList(id)
	}
	panic("GetCachedLabelList should not have been called")
}

func (ds *DaemonSuite) GetPolicyRepository() *policy.Repository {
	if ds.OnGetPolicyRepository != nil {
		return ds.OnGetPolicyRepository()
	}
	panic("GetPolicyRepository should not have been called")
}

func (ds *DaemonSuite) UpdateProxyRedirect(e regeneration.EndpointUpdater, l4 *policy.L4Filter, proxyWaitGroup *completion.WaitGroup) (uint16, error, revert.FinalizeFunc, revert.RevertFunc) {
	if ds.OnUpdateProxyRedirect != nil {
		return ds.OnUpdateProxyRedirect(e, l4, proxyWaitGroup)
	}
	panic("UpdateProxyRedirect should not have been called")
}

func (ds *DaemonSuite) RemoveProxyRedirect(e regeneration.EndpointInfoSource, id string, proxyWaitGroup *completion.WaitGroup) (error, revert.FinalizeFunc, revert.RevertFunc) {
	if ds.OnRemoveProxyRedirect != nil {
		return ds.OnRemoveProxyRedirect(e, id, proxyWaitGroup)
	}
	panic("RemoveProxyRedirect should not have been called")
}

func (ds *DaemonSuite) UpdateNetworkPolicy(e regeneration.EndpointUpdater, policy *policy.L4Policy,
	proxyWaitGroup *completion.WaitGroup) (error, revert.RevertFunc) {
	if ds.OnUpdateNetworkPolicy != nil {
		return ds.OnUpdateNetworkPolicy(e, policy, proxyWaitGroup)
	}
	panic("UpdateNetworkPolicy should not have been called")
}

func (ds *DaemonSuite) RemoveNetworkPolicy(e regeneration.EndpointInfoSource) {
	if ds.OnRemoveNetworkPolicy != nil {
		ds.OnRemoveNetworkPolicy(e)
	}
	panic("RemoveNetworkPolicy should not have been called")
}

func (ds *DaemonSuite) QueueEndpointBuild(ctx context.Context, epID uint64) (func(), error) {
	if ds.OnQueueEndpointBuild != nil {
		return ds.OnQueueEndpointBuild(ctx, epID)
	}
	panic("QueueEndpointBuild should not have been called")
}

func (ds *DaemonSuite) RemoveFromEndpointQueue(epID uint64) {
	if ds.OnRemoveFromEndpointQueue != nil {
		ds.OnRemoveFromEndpointQueue(epID)
		return
	}
	panic("RemoveFromEndpointQueue should not have been called")
}

func (ds *DaemonSuite) DebugEnabled() bool {
	if ds.OnDebugEnabled != nil {
		return ds.OnDebugEnabled()
	}
	panic("DebugEnabled should not have been called")
}

func (ds *DaemonSuite) GetCompilationLock() *lock.RWMutex {
	if ds.OnGetCompilationLock != nil {
		return ds.OnGetCompilationLock()
	}
	panic("GetCompilationLock should not have been called")
}

func (ds *DaemonSuite) SendNotification(typ monitorAPI.AgentNotification, text string) error {
	if ds.OnSendNotification != nil {
		return ds.OnSendNotification(typ, text)
	}
	panic("SendNotification should not have been called")
}

func (ds *DaemonSuite) NewProxyLogRecord(l *accesslog.LogRecord) error {
	if ds.OnNewProxyLogRecord != nil {
		return ds.OnNewProxyLogRecord(l)
	}
	panic("NewProxyLogRecord should not have been called")
}

func (ds *DaemonSuite) Datapath() datapath.Datapath {
	return ds.d.datapath
}

func (ds *DaemonSuite) ClearPolicyConsumers(id uint16) *sync.WaitGroup {
	if ds.OnClearPolicyConsumers != nil {
		return ds.OnClearPolicyConsumers(id)
	}
	panic("ClearPolicyConsumers should not have been called")
}

func (ds *DaemonSuite) GetNodeSuffix() string {
	return ds.d.GetNodeSuffix()
}

func (ds *DaemonSuite) UpdateIdentities(added, deleted cache.IdentityCache) {
	ds.d.UpdateIdentities(added, deleted)
}
<|endoftext|>


<commit_before>package daemon

import (
	"fmt"

	"github.com/docker/docker/container"
	"github.com/docker/docker/libcontainerd"
)

func (daemon *Daemon) getLibcontainerdCreateOptions(container *container.Container) (*[]libcontainerd.CreateOption, error) {
	createOptions := []libcontainerd.CreateOption{}

	rt := daemon.configStore.GetRuntime(container.HostConfig.Runtime)
	if rt == nil {
		return nil, fmt.Errorf("No such runtime '%s'", container.HostConfig.Runtime)
	}
	createOptions = append(createOptions, libcontainerd.WithRuntime(rt.Path, rt.Args))

	return &createOptions, nil
}
<commit_msg>Fix missing container runtime on upgrade<commit_after>package daemon

import (
	"fmt"

	"github.com/docker/docker/container"
	"github.com/docker/docker/libcontainerd"
	"github.com/docker/engine-api/types"
)

func (daemon *Daemon) getLibcontainerdCreateOptions(container *container.Container) (*[]libcontainerd.CreateOption, error) {
	createOptions := []libcontainerd.CreateOption{}

	// Ensure a runtime has been assigned to this container
	if container.HostConfig.Runtime == "" {
		container.HostConfig.Runtime = types.DefaultRuntimeName
		container.ToDisk()
	}

	rt := daemon.configStore.GetRuntime(container.HostConfig.Runtime)
	if rt == nil {
		return nil, fmt.Errorf("no such runtime '%s'", container.HostConfig.Runtime)
	}
	createOptions = append(createOptions, libcontainerd.WithRuntime(rt.Path, rt.Args))

	return &createOptions, nil
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"bytes"
	"strings"

	"launchpad.net/goamz/aws"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/juju/testing"
	coretesting "launchpad.net/juju-core/testing"
)

type ValidateMetadataSuite struct {
	testing.RepoSuite
	home *coretesting.FakeHome
}

var _ = gc.Suite(&ValidateMetadataSuite{})

func runValidateMetadata(c *gc.C, args ...string) error {
	_, err := coretesting.RunCommand(c, &ValidateMetadataCommand{}, args)
	return err
}

var validateInitErrorTests = []struct {
	args []string
	err  string
}{
	{
		args: []string{"-p", "ec2", "-r", "region", "-d", "dir"},
		err:  `series required if provider type is specified`,
	}, {
		args: []string{"-p", "ec2", "-s", "series", "-d", "dir"},
		err:  `region required if provider type is specified`,
	}, {
		args: []string{"-p", "ec2", "-s", "series", "-r", "region"},
		err:  `metadata directory required if provider type is specified`,
	},
}

func (s *ValidateMetadataSuite) TestInitErrors(c *gc.C) {
	for i, t := range validateInitErrorTests {
		c.Logf("test %d", i)
		err := coretesting.InitCommand(&ValidateMetadataCommand{}, t.args)
		c.Check(err, gc.ErrorMatches, t.err)
	}
}

func (s *ValidateMetadataSuite) TestInvalidProviderError(c *gc.C) {
	err := runValidateMetadata(c, "-p", "foo", "-s", "series", "-r", "region", "-d", "dir")
	c.Check(err, gc.ErrorMatches, `no registered provider for "foo"`)
}

func (s *ValidateMetadataSuite) TestUnsupportedProviderError(c *gc.C) {
	err := runValidateMetadata(c, "-p", "local", "-s", "series", "-r", "region", "-d", "dir")
	c.Check(err, gc.ErrorMatches, `local provider does not support image metadata validation`)
}

func (s *ValidateMetadataSuite) makeLocalMetadata(c *gc.C, id, region, series, endpoint string) error {
	im := imagemetadata.ImageMetadata{
		Id:   id,
		Arch: "amd64",
	}
	cloudSpec := imagemetadata.CloudSpec{
		Region:   region,
		Endpoint: endpoint,
	}
	_, err := imagemetadata.MakeBoilerplate("", series, &im, &cloudSpec, false)
	if err != nil {
		return err
	}
	return nil
}

const metadataTestEnvConfig = `
environments:
    ec2:
        type: ec2
        control-bucket: foo
        access-key: access
        secret-key: secret
        default-series: precise
        region: us-east-1
`

func (s *ValidateMetadataSuite) SetUpTest(c *gc.C) {
	s.RepoSuite.SetUpTest(c)
	s.home = coretesting.MakeFakeHome(c, metadataTestEnvConfig)
}

func (s *ValidateMetadataSuite) TearDownTest(c *gc.C) {
	s.home.Restore()
	s.RepoSuite.TearDownTest(c)
}

func (s *ValidateMetadataSuite) setupEc2LocalMetadata(c *gc.C, region string) {
	ec2Region, ok := aws.Regions[region]
	if !ok {
		c.Fatalf("unknown ec2 region %q")
	}
	endpoint := ec2Region.EC2Endpoint
	s.makeLocalMetadata(c, "1234", region, "precise", endpoint)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataUsingEnvironment(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-east-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{"-e", "ec2", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "us-east-1":.*`)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataWithManualParams(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-west-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "precise", "-r", "us-west-1",
			"-u", "https://ec2.us-west-1.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "us-west-1":.*`)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataNoMatch(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-east-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "raring", "-r", "us-west-1",
			"-u", "https://ec2.us-west-1.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
	code = cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "precise", "-r", "region",
			"-u", "https://ec2.region.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
}

func (s *ValidateMetadataSuite) TestOpenstackLocalMetadataWithManualParams(c *gc.C) {
	s.makeLocalMetadata(c, "1234", "region-2", "raring", "some-auth-url")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "raring", "-r", "region-2",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "region-2":.*`)
}

func (s *ValidateMetadataSuite) TestOpenstackLocalMetadataNoMatch(c *gc.C) {
	s.makeLocalMetadata(c, "1234", "region-2", "raring", "some-auth-url")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "precise", "-r", "region-2",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
	code = cmd.Main(
		&ValidateMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "raring", "-r", "region-3",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
}
<commit_msg>Fix tests<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"bytes"
	"strings"

	"launchpad.net/goamz/aws"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/imagemetadata"
	"launchpad.net/juju-core/juju/testing"
	coretesting "launchpad.net/juju-core/testing"
)

type ValidateMetadataSuite struct {
	testing.RepoSuite
	home *coretesting.FakeHome
}

var _ = gc.Suite(&ValidateMetadataSuite{})

func runValidateMetadata(c *gc.C, args ...string) error {
	_, err := coretesting.RunCommand(c, &ValidateImageMetadataCommand{}, args)
	return err
}

var validateInitErrorTests = []struct {
	args []string
	err  string
}{
	{
		args: []string{"-p", "ec2", "-r", "region", "-d", "dir"},
		err:  `series required if provider type is specified`,
	}, {
		args: []string{"-p", "ec2", "-s", "series", "-d", "dir"},
		err:  `region required if provider type is specified`,
	}, {
		args: []string{"-p", "ec2", "-s", "series", "-r", "region"},
		err:  `metadata directory required if provider type is specified`,
	},
}

func (s *ValidateMetadataSuite) TestInitErrors(c *gc.C) {
	for i, t := range validateInitErrorTests {
		c.Logf("test %d", i)
		err := coretesting.InitCommand(&ValidateImageMetadataCommand{}, t.args)
		c.Check(err, gc.ErrorMatches, t.err)
	}
}

func (s *ValidateMetadataSuite) TestInvalidProviderError(c *gc.C) {
	err := runValidateMetadata(c, "-p", "foo", "-s", "series", "-r", "region", "-d", "dir")
	c.Check(err, gc.ErrorMatches, `no registered provider for "foo"`)
}

func (s *ValidateMetadataSuite) TestUnsupportedProviderError(c *gc.C) {
	err := runValidateMetadata(c, "-p", "local", "-s", "series", "-r", "region", "-d", "dir")
	c.Check(err, gc.ErrorMatches, `local provider does not support image metadata validation`)
}

func (s *ValidateMetadataSuite) makeLocalMetadata(c *gc.C, id, region, series, endpoint string) error {
	im := imagemetadata.ImageMetadata{
		Id:   id,
		Arch: "amd64",
	}
	cloudSpec := imagemetadata.CloudSpec{
		Region:   region,
		Endpoint: endpoint,
	}
	_, err := imagemetadata.MakeBoilerplate("", series, &im, &cloudSpec, false)
	if err != nil {
		return err
	}
	return nil
}

const metadataTestEnvConfig = `
environments:
    ec2:
        type: ec2
        control-bucket: foo
        access-key: access
        secret-key: secret
        default-series: precise
        region: us-east-1
`

func (s *ValidateMetadataSuite) SetUpTest(c *gc.C) {
	s.RepoSuite.SetUpTest(c)
	s.home = coretesting.MakeFakeHome(c, metadataTestEnvConfig)
}

func (s *ValidateMetadataSuite) TearDownTest(c *gc.C) {
	s.home.Restore()
	s.RepoSuite.TearDownTest(c)
}

func (s *ValidateMetadataSuite) setupEc2LocalMetadata(c *gc.C, region string) {
	ec2Region, ok := aws.Regions[region]
	if !ok {
		c.Fatalf("unknown ec2 region %q")
	}
	endpoint := ec2Region.EC2Endpoint
	s.makeLocalMetadata(c, "1234", region, "precise", endpoint)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataUsingEnvironment(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-east-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{"-e", "ec2", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "us-east-1":.*`)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataWithManualParams(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-west-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "precise", "-r", "us-west-1",
			"-u", "https://ec2.us-west-1.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "us-west-1":.*`)
}

func (s *ValidateMetadataSuite) TestEc2LocalMetadataNoMatch(c *gc.C) {
	s.setupEc2LocalMetadata(c, "us-east-1")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "raring", "-r", "us-west-1",
			"-u", "https://ec2.us-west-1.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
	code = cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "ec2", "-s", "precise", "-r", "region",
			"-u", "https://ec2.region.amazonaws.com", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
}

func (s *ValidateMetadataSuite) TestOpenstackLocalMetadataWithManualParams(c *gc.C) {
	s.makeLocalMetadata(c, "1234", "region-2", "raring", "some-auth-url")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "raring", "-r", "region-2",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 0)
	errOut := ctx.Stdout.(*bytes.Buffer).String()
	strippedOut := strings.Replace(errOut, "\n", "", -1)
	c.Check(strippedOut, gc.Matches, `matching image ids for region "region-2":.*`)
}

func (s *ValidateMetadataSuite) TestOpenstackLocalMetadataNoMatch(c *gc.C) {
	s.makeLocalMetadata(c, "1234", "region-2", "raring", "some-auth-url")
	ctx := coretesting.Context(c)
	metadataDir := config.JujuHomePath("")
	code := cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "precise", "-r", "region-2",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
	code = cmd.Main(
		&ValidateImageMetadataCommand{}, ctx, []string{
			"-p", "openstack", "-s", "raring", "-r", "region-3",
			"-u", "some-auth-url", "-d", metadataDir},
	)
	c.Assert(code, gc.Equals, 1)
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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.

// The ct_server binary runs the CT personality.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/etcd/clientv3"
	etcdnaming "github.com/coreos/etcd/clientv3/naming"
	"github.com/golang/glog"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/util"
	"github.com/google/trillian"
	"github.com/google/trillian/monitoring/opencensus"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rs/cors"
	"github.com/tomasen/realip"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/naming"

	// Register PEMKeyFile, PrivateKey and PKCS11Config ProtoHandlers
	_ "github.com/google/trillian/crypto/keys/der/proto"
	_ "github.com/google/trillian/crypto/keys/pem/proto"
	_ "github.com/google/trillian/crypto/keys/pkcs11/proto"
)

// Global flags that affect all log instances.
var (
	httpEndpoint       = flag.String("http_endpoint", "localhost:6962", "Endpoint for HTTP (host:port)")
	metricsEndpoint    = flag.String("metrics_endpoint", "localhost:6963", "Endpoint for serving metrics; if left empty, metrics will be visible on --http_endpoint")
	rpcBackend         = flag.String("log_rpc_server", "localhost:8090", "Backend specification; comma-separated list or etcd service name (if --etcd_servers specified). If unset backends are specified in config (as a LogMultiConfig proto)")
	rpcDeadline        = flag.Duration("rpc_deadline", time.Second*10, "Deadline for backend RPC requests")
	getSTHInterval     = flag.Duration("get_sth_interval", time.Second*180, "Interval between internal get-sth operations (0 to disable)")
	logConfig          = flag.String("log_config", "", "File holding log config in text proto format")
	maxGetEntries      = flag.Int64("max_get_entries", 0, "Max number of entries we allow in a get-entries request (0=>use default 1000)")
	etcdServers        = flag.String("etcd_servers", "", "A comma-separated list of etcd servers")
	etcdHTTPService    = flag.String("etcd_http_service", "trillian-ctfe-http", "Service name to announce our HTTP endpoint under")
	etcdMetricsService = flag.String("etcd_metrics_service", "trillian-ctfe-metrics-http", "Service name to announce our HTTP metrics endpoint under")
	maskInternalErrors = flag.Bool("mask_internal_errors", false, "Don't return error strings with Internal Server Error HTTP responses")
	tracing            = flag.Bool("tracing", false, "If true opencensus Stackdriver tracing will be enabled. See https://opencensus.io/.")
	tracingProjectID   = flag.String("tracing_project_id", "", "project ID to pass to stackdriver. Can be empty for GCP, consult docs for other platforms.")
	tracingPercent     = flag.Int("tracing_percent", 0, "Percent of requests to be traced. Zero is a special case to use the DefaultSampler")
	quotaRemote        = flag.Bool("quota_remote", true, "Enable requesting of quota for IP address sending incoming requests")
	quotaIntermediate  = flag.Bool("quota_intermediate", true, "Enable requesting of quota for intermediate certificates in sumbmitted chains")
	handlerPrefix      = flag.String("handler_prefix", "", "If set e.g. to '/logs' will prefix all handlers that don't define a custom prefix")
)

// nolint:staticcheck
func main() {
	flag.Parse()
	ctx := context.Background()

	if *maxGetEntries > 0 {
		ctfe.MaxGetEntriesAllowed = *maxGetEntries
	}

	var cfg *configpb.LogMultiConfig
	var err error
	// Get log config from file before we start. This is a different proto
	// type if we're using a multi backend configuration (no rpcBackend set
	// in flags). The single-backend config is converted to a multi config so
	// they can be treated the same.
	if len(*rpcBackend) > 0 {
		var cfgs []*configpb.LogConfig
		if cfgs, err = ctfe.LogConfigFromFile(*logConfig); err == nil {
			cfg = ctfe.ToMultiLogConfig(cfgs, *rpcBackend)
		}
	} else {
		cfg, err = ctfe.MultiLogConfigFromFile(*logConfig)
	}

	if err != nil {
		glog.Exitf("Failed to read config: %v", err)
	}

	beMap, err := ctfe.ValidateLogMultiConfig(cfg)
	if err != nil {
		glog.Exitf("Invalid config: %v", err)
	}

	glog.CopyStandardLogTo("WARNING")
	glog.Info("**** CT HTTP Server Starting ****")

	metricsAt := *metricsEndpoint
	if metricsAt == "" {
		metricsAt = *httpEndpoint
	}

	dialOpts := []grpc.DialOption{grpc.WithInsecure()}
	if len(*etcdServers) > 0 {
		// Use etcd to provide endpoint resolution.
		cfg := clientv3.Config{Endpoints: strings.Split(*etcdServers, ","), DialTimeout: 5 * time.Second}
		client, err := clientv3.New(cfg)
		if err != nil {
			glog.Exitf("Failed to connect to etcd at %v: %v", *etcdServers, err)
		}
		etcdRes := &etcdnaming.GRPCResolver{Client: client}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(etcdRes)))

		// Also announce ourselves.
		updateHTTP := naming.Update{Op: naming.Add, Addr: *httpEndpoint}
		updateMetrics := naming.Update{Op: naming.Add, Addr: metricsAt}
		glog.Infof("Announcing our presence in %v with %+v", *etcdHTTPService, updateHTTP)
		etcdRes.Update(ctx, *etcdHTTPService, updateHTTP)
		glog.Infof("Announcing our presence in %v with %+v", *etcdMetricsService, updateMetrics)
		etcdRes.Update(ctx, *etcdMetricsService, updateMetrics)

		byeHTTP := naming.Update{Op: naming.Delete, Addr: *httpEndpoint}
		byeMetrics := naming.Update{Op: naming.Delete, Addr: metricsAt}
		defer func() {
			glog.Infof("Removing our presence in %v with %+v", *etcdHTTPService, byeHTTP)
			etcdRes.Update(ctx, *etcdHTTPService, byeHTTP)
			glog.Infof("Removing our presence in %v with %+v", *etcdMetricsService, byeMetrics)
			etcdRes.Update(ctx, *etcdMetricsService, byeMetrics)
		}()
	} else if strings.Contains(*rpcBackend, ",") {
		glog.Infof("Using FixedBackendResolver")
		// Use a fixed endpoint resolution that just returns the addresses configured on the command line.
		res := util.FixedBackendResolver{}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(res)))
	} else {
		glog.Infof("Using regular DNS resolver")
		dialOpts = append(dialOpts, grpc.WithBalancerName(roundrobin.Name))
	}

	// Dial all our log backends.
	clientMap := make(map[string]trillian.TrillianLogClient)
	for _, be := range beMap {
		glog.Infof("Dialling backend: %v", be)
		if len(beMap) == 1 {
			// If there's only one of them we use the blocking option as we can't
			// serve anything until connected.
			dialOpts = append(dialOpts, grpc.WithBlock())
		}
		conn, err := grpc.Dial(be.BackendSpec, dialOpts...)
		if err != nil {
			glog.Exitf("Could not dial RPC server: %v: %v", be, err)
		}
		defer conn.Close()
		clientMap[be.Name] = trillian.NewTrillianLogClient(conn)
	}

	// Allow cross-origin requests to all handlers registered on corsMux.
	// This is safe for CT log handlers because the log is public and
	// unauthenticated so cross-site scripting attacks are not a concern.
	corsMux := http.NewServeMux()
	corsHandler := cors.AllowAll().Handler(corsMux)
	http.Handle("/", corsHandler)

	// Register handlers for all the configured logs using the correct RPC
	// client.
	for _, c := range cfg.LogConfigs.Config {
		inst, err := setupAndRegister(ctx, clientMap[c.LogBackendName], *rpcDeadline, c, corsMux, *handlerPrefix, *maskInternalErrors)
		if err != nil {
			glog.Exitf("Failed to set up log instance for %+v: %v", cfg, err)
		}
		if *getSTHInterval > 0 {
			go inst.RunUpdateSTH(ctx, *getSTHInterval)
		}
	}

	// Return a 200 on the root, for GCE default health checking :/
	corsMux.HandleFunc("/", func(resp http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/" {
			resp.WriteHeader(http.StatusOK)
		} else {
			resp.WriteHeader(http.StatusNotFound)
		}
	})

	if metricsAt != *httpEndpoint {
		// Run a separate handler for metrics.
		go func() {
			mux := http.NewServeMux()
			mux.Handle("/metrics", promhttp.Handler())
			metricsServer := http.Server{Addr: metricsAt, Handler: mux}
			err := metricsServer.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		// Handle metrics on the DefaultServeMux.
		http.Handle("/metrics", promhttp.Handler())
	}

	// If we're enabling tracing we need to use an instrumented http.Handler.
	var handler http.Handler
	if *tracing {
		handler, err = opencensus.EnableHTTPServerTracing(*tracingProjectID, *tracingPercent)
		if err != nil {
			glog.Exitf("Failed to initialize stackdriver / opencensus tracing: %v", err)
		}
	}

	// Bring up the HTTP server and serve until we get a signal not to.
	srv := http.Server{Addr: *httpEndpoint, Handler: handler}
	shutdownWG := new(sync.WaitGroup)
	go awaitSignal(func() {
		shutdownWG.Add(1)
		defer shutdownWG.Done()
		// Allow 60s for any pending requests to finish then terminate any stragglers
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
		defer cancel()
		glog.Info("Shutting down HTTP server...")
		srv.Shutdown(ctx)
		glog.Info("HTTP server shutdown")
	})

	err = srv.ListenAndServe()
	if err != http.ErrServerClosed {
		glog.Warningf("Server exited: %v", err)
	}
	// Wait will only block if the function passed to awaitSignal was called,
	// in which case it'll block until the HTTP server has gracefully shutdown
	shutdownWG.Wait()
	glog.Flush()
}

// awaitSignal waits for standard termination signals, then runs the given
// function; it should be run as a separate goroutine.
func awaitSignal(doneFn func()) {
	// Arrange notification for the standard set of signals used to terminate a server
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Now block main and wait for a signal
	sig := <-sigs
	glog.Warningf("Signal received: %v", sig)
	glog.Flush()

	doneFn()
}

func setupAndRegister(ctx context.Context, client trillian.TrillianLogClient, deadline time.Duration, cfg *configpb.LogConfig, mux *http.ServeMux, globalHandlerPrefix string, maskInternalErrors bool) (*ctfe.Instance, error) {
	vCfg, err := ctfe.ValidateLogConfig(cfg)
	if err != nil {
		return nil, err
	}

	opts := ctfe.InstanceOptions{
		Validated:          vCfg,
		Client:             client,
		Deadline:           deadline,
		MetricFactory:      prometheus.MetricFactory{},
		RequestLog:         new(ctfe.DefaultRequestLog),
		MaskInternalErrors: maskInternalErrors,
	}
	if *quotaRemote {
		glog.Info("Enabling quota for requesting IP")
		opts.RemoteQuotaUser = realip.FromRequest
	}
	if *quotaIntermediate {
		glog.Info("Enabling quota for intermediate certificates")
		opts.CertificateQuotaUser = ctfe.QuotaUserForCert
	}
	// Full handler pattern will be of the form "/logs/yyz/ct/v1/add-chain", where "/logs" is the
	// HandlerPrefix and "yyz" is the c.Prefix for this particular log. Use the default
	// HandlerPrefix unless the log config overrides it. The custom prefix in
	// the log configuration intended for use in migration scenarios where logs
	// have an existing URL path that differs from the global one. For example
	// if all new logs are served on "/logs/log/..." and a previously existing
	// log is at "/log/..." this is now supported.
	lhp := globalHandlerPrefix
	if ohPrefix := cfg.OverrideHandlerPrefix; len(ohPrefix) > 0 {
		glog.Infof("Log with prefix: %s is using a custom HandlerPrefix: %s", cfg.Prefix, ohPrefix)
		lhp = "/" + strings.Trim(ohPrefix, "/")
	}
	inst, err := ctfe.SetUpInstance(ctx, opts)
	if err != nil {
		return nil, err
	}
	for path, handler := range inst.Handlers {
		mux.Handle(lhp+path, handler)
	}
	return inst, nil
}
<commit_msg>CTFE: Remove default values for --metrics_endpoint and --log_rpc_server<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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.

// The ct_server binary runs the CT personality.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/etcd/clientv3"
	etcdnaming "github.com/coreos/etcd/clientv3/naming"
	"github.com/golang/glog"
	"github.com/google/certificate-transparency-go/trillian/ctfe"
	"github.com/google/certificate-transparency-go/trillian/ctfe/configpb"
	"github.com/google/certificate-transparency-go/trillian/util"
	"github.com/google/trillian"
	"github.com/google/trillian/monitoring/opencensus"
	"github.com/google/trillian/monitoring/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/rs/cors"
	"github.com/tomasen/realip"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer/roundrobin"
	"google.golang.org/grpc/naming"

	// Register PEMKeyFile, PrivateKey and PKCS11Config ProtoHandlers
	_ "github.com/google/trillian/crypto/keys/der/proto"
	_ "github.com/google/trillian/crypto/keys/pem/proto"
	_ "github.com/google/trillian/crypto/keys/pkcs11/proto"
)

// Global flags that affect all log instances.
var (
	httpEndpoint       = flag.String("http_endpoint", "localhost:6962", "Endpoint for HTTP (host:port)")
	metricsEndpoint    = flag.String("metrics_endpoint", "", "Endpoint for serving metrics; if left empty, metrics will be visible on --http_endpoint")
	rpcBackend         = flag.String("log_rpc_server", "", "Backend specification; comma-separated list or etcd service name (if --etcd_servers specified). If unset backends are specified in config (as a LogMultiConfig proto)")
	rpcDeadline        = flag.Duration("rpc_deadline", time.Second*10, "Deadline for backend RPC requests")
	getSTHInterval     = flag.Duration("get_sth_interval", time.Second*180, "Interval between internal get-sth operations (0 to disable)")
	logConfig          = flag.String("log_config", "", "File holding log config in text proto format")
	maxGetEntries      = flag.Int64("max_get_entries", 0, "Max number of entries we allow in a get-entries request (0=>use default 1000)")
	etcdServers        = flag.String("etcd_servers", "", "A comma-separated list of etcd servers")
	etcdHTTPService    = flag.String("etcd_http_service", "trillian-ctfe-http", "Service name to announce our HTTP endpoint under")
	etcdMetricsService = flag.String("etcd_metrics_service", "trillian-ctfe-metrics-http", "Service name to announce our HTTP metrics endpoint under")
	maskInternalErrors = flag.Bool("mask_internal_errors", false, "Don't return error strings with Internal Server Error HTTP responses")
	tracing            = flag.Bool("tracing", false, "If true opencensus Stackdriver tracing will be enabled. See https://opencensus.io/.")
	tracingProjectID   = flag.String("tracing_project_id", "", "project ID to pass to stackdriver. Can be empty for GCP, consult docs for other platforms.")
	tracingPercent     = flag.Int("tracing_percent", 0, "Percent of requests to be traced. Zero is a special case to use the DefaultSampler")
	quotaRemote        = flag.Bool("quota_remote", true, "Enable requesting of quota for IP address sending incoming requests")
	quotaIntermediate  = flag.Bool("quota_intermediate", true, "Enable requesting of quota for intermediate certificates in sumbmitted chains")
	handlerPrefix      = flag.String("handler_prefix", "", "If set e.g. to '/logs' will prefix all handlers that don't define a custom prefix")
)

// nolint:staticcheck
func main() {
	flag.Parse()
	ctx := context.Background()

	if *maxGetEntries > 0 {
		ctfe.MaxGetEntriesAllowed = *maxGetEntries
	}

	var cfg *configpb.LogMultiConfig
	var err error
	// Get log config from file before we start. This is a different proto
	// type if we're using a multi backend configuration (no rpcBackend set
	// in flags). The single-backend config is converted to a multi config so
	// they can be treated the same.
	if len(*rpcBackend) > 0 {
		var cfgs []*configpb.LogConfig
		if cfgs, err = ctfe.LogConfigFromFile(*logConfig); err == nil {
			cfg = ctfe.ToMultiLogConfig(cfgs, *rpcBackend)
		}
	} else {
		cfg, err = ctfe.MultiLogConfigFromFile(*logConfig)
	}

	if err != nil {
		glog.Exitf("Failed to read config: %v", err)
	}

	beMap, err := ctfe.ValidateLogMultiConfig(cfg)
	if err != nil {
		glog.Exitf("Invalid config: %v", err)
	}

	glog.CopyStandardLogTo("WARNING")
	glog.Info("**** CT HTTP Server Starting ****")

	metricsAt := *metricsEndpoint
	if metricsAt == "" {
		metricsAt = *httpEndpoint
	}

	dialOpts := []grpc.DialOption{grpc.WithInsecure()}
	if len(*etcdServers) > 0 {
		// Use etcd to provide endpoint resolution.
		cfg := clientv3.Config{Endpoints: strings.Split(*etcdServers, ","), DialTimeout: 5 * time.Second}
		client, err := clientv3.New(cfg)
		if err != nil {
			glog.Exitf("Failed to connect to etcd at %v: %v", *etcdServers, err)
		}
		etcdRes := &etcdnaming.GRPCResolver{Client: client}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(etcdRes)))

		// Also announce ourselves.
		updateHTTP := naming.Update{Op: naming.Add, Addr: *httpEndpoint}
		updateMetrics := naming.Update{Op: naming.Add, Addr: metricsAt}
		glog.Infof("Announcing our presence in %v with %+v", *etcdHTTPService, updateHTTP)
		etcdRes.Update(ctx, *etcdHTTPService, updateHTTP)
		glog.Infof("Announcing our presence in %v with %+v", *etcdMetricsService, updateMetrics)
		etcdRes.Update(ctx, *etcdMetricsService, updateMetrics)

		byeHTTP := naming.Update{Op: naming.Delete, Addr: *httpEndpoint}
		byeMetrics := naming.Update{Op: naming.Delete, Addr: metricsAt}
		defer func() {
			glog.Infof("Removing our presence in %v with %+v", *etcdHTTPService, byeHTTP)
			etcdRes.Update(ctx, *etcdHTTPService, byeHTTP)
			glog.Infof("Removing our presence in %v with %+v", *etcdMetricsService, byeMetrics)
			etcdRes.Update(ctx, *etcdMetricsService, byeMetrics)
		}()
	} else if strings.Contains(*rpcBackend, ",") {
		glog.Infof("Using FixedBackendResolver")
		// Use a fixed endpoint resolution that just returns the addresses configured on the command line.
		res := util.FixedBackendResolver{}
		dialOpts = append(dialOpts, grpc.WithBalancer(grpc.RoundRobin(res)))
	} else {
		glog.Infof("Using regular DNS resolver")
		dialOpts = append(dialOpts, grpc.WithBalancerName(roundrobin.Name))
	}

	// Dial all our log backends.
	clientMap := make(map[string]trillian.TrillianLogClient)
	for _, be := range beMap {
		glog.Infof("Dialling backend: %v", be)
		if len(beMap) == 1 {
			// If there's only one of them we use the blocking option as we can't
			// serve anything until connected.
			dialOpts = append(dialOpts, grpc.WithBlock())
		}
		conn, err := grpc.Dial(be.BackendSpec, dialOpts...)
		if err != nil {
			glog.Exitf("Could not dial RPC server: %v: %v", be, err)
		}
		defer conn.Close()
		clientMap[be.Name] = trillian.NewTrillianLogClient(conn)
	}

	// Allow cross-origin requests to all handlers registered on corsMux.
	// This is safe for CT log handlers because the log is public and
	// unauthenticated so cross-site scripting attacks are not a concern.
	corsMux := http.NewServeMux()
	corsHandler := cors.AllowAll().Handler(corsMux)
	http.Handle("/", corsHandler)

	// Register handlers for all the configured logs using the correct RPC
	// client.
	for _, c := range cfg.LogConfigs.Config {
		inst, err := setupAndRegister(ctx, clientMap[c.LogBackendName], *rpcDeadline, c, corsMux, *handlerPrefix, *maskInternalErrors)
		if err != nil {
			glog.Exitf("Failed to set up log instance for %+v: %v", cfg, err)
		}
		if *getSTHInterval > 0 {
			go inst.RunUpdateSTH(ctx, *getSTHInterval)
		}
	}

	// Return a 200 on the root, for GCE default health checking :/
	corsMux.HandleFunc("/", func(resp http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/" {
			resp.WriteHeader(http.StatusOK)
		} else {
			resp.WriteHeader(http.StatusNotFound)
		}
	})

	if metricsAt != *httpEndpoint {
		// Run a separate handler for metrics.
		go func() {
			mux := http.NewServeMux()
			mux.Handle("/metrics", promhttp.Handler())
			metricsServer := http.Server{Addr: metricsAt, Handler: mux}
			err := metricsServer.ListenAndServe()
			glog.Warningf("Metrics server exited: %v", err)
		}()
	} else {
		// Handle metrics on the DefaultServeMux.
		http.Handle("/metrics", promhttp.Handler())
	}

	// If we're enabling tracing we need to use an instrumented http.Handler.
	var handler http.Handler
	if *tracing {
		handler, err = opencensus.EnableHTTPServerTracing(*tracingProjectID, *tracingPercent)
		if err != nil {
			glog.Exitf("Failed to initialize stackdriver / opencensus tracing: %v", err)
		}
	}

	// Bring up the HTTP server and serve until we get a signal not to.
	srv := http.Server{Addr: *httpEndpoint, Handler: handler}
	shutdownWG := new(sync.WaitGroup)
	go awaitSignal(func() {
		shutdownWG.Add(1)
		defer shutdownWG.Done()
		// Allow 60s for any pending requests to finish then terminate any stragglers
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
		defer cancel()
		glog.Info("Shutting down HTTP server...")
		srv.Shutdown(ctx)
		glog.Info("HTTP server shutdown")
	})

	err = srv.ListenAndServe()
	if err != http.ErrServerClosed {
		glog.Warningf("Server exited: %v", err)
	}
	// Wait will only block if the function passed to awaitSignal was called,
	// in which case it'll block until the HTTP server has gracefully shutdown
	shutdownWG.Wait()
	glog.Flush()
}

// awaitSignal waits for standard termination signals, then runs the given
// function; it should be run as a separate goroutine.
func awaitSignal(doneFn func()) {
	// Arrange notification for the standard set of signals used to terminate a server
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	// Now block main and wait for a signal
	sig := <-sigs
	glog.Warningf("Signal received: %v", sig)
	glog.Flush()

	doneFn()
}

func setupAndRegister(ctx context.Context, client trillian.TrillianLogClient, deadline time.Duration, cfg *configpb.LogConfig, mux *http.ServeMux, globalHandlerPrefix string, maskInternalErrors bool) (*ctfe.Instance, error) {
	vCfg, err := ctfe.ValidateLogConfig(cfg)
	if err != nil {
		return nil, err
	}

	opts := ctfe.InstanceOptions{
		Validated:          vCfg,
		Client:             client,
		Deadline:           deadline,
		MetricFactory:      prometheus.MetricFactory{},
		RequestLog:         new(ctfe.DefaultRequestLog),
		MaskInternalErrors: maskInternalErrors,
	}
	if *quotaRemote {
		glog.Info("Enabling quota for requesting IP")
		opts.RemoteQuotaUser = realip.FromRequest
	}
	if *quotaIntermediate {
		glog.Info("Enabling quota for intermediate certificates")
		opts.CertificateQuotaUser = ctfe.QuotaUserForCert
	}
	// Full handler pattern will be of the form "/logs/yyz/ct/v1/add-chain", where "/logs" is the
	// HandlerPrefix and "yyz" is the c.Prefix for this particular log. Use the default
	// HandlerPrefix unless the log config overrides it. The custom prefix in
	// the log configuration intended for use in migration scenarios where logs
	// have an existing URL path that differs from the global one. For example
	// if all new logs are served on "/logs/log/..." and a previously existing
	// log is at "/log/..." this is now supported.
	lhp := globalHandlerPrefix
	if ohPrefix := cfg.OverrideHandlerPrefix; len(ohPrefix) > 0 {
		glog.Infof("Log with prefix: %s is using a custom HandlerPrefix: %s", cfg.Prefix, ohPrefix)
		lhp = "/" + strings.Trim(ohPrefix, "/")
	}
	inst, err := ctfe.SetUpInstance(ctx, opts)
	if err != nil {
		return nil, err
	}
	for path, handler := range inst.Handlers {
		mux.Handle(lhp+path, handler)
	}
	return inst, nil
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"
	"time"
)

func Logger(inner http.Handler, name string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		inner.ServeHTTP(w, r)

		log.Printf(
			"%s\t%s\t%s\t%s",
			r.Method,
			r.RequestURI,
			name,
			time.Since(start),
		)
	})
}
<commit_msg>Added recover function for handlers<commit_after>package main

import (
	"log"
	"net/http"
	"time"
	"errors"
)

func Logger(inner http.Handler, name string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		var err error
			defer func() {
					r := recover()
					if r != nil {
							switch t := r.(type) {
							case string:
									err = errors.New(t)
							case error:
									err = t
							default:
									err = errors.New("Unknown error")
							}
							http.Error(w, err.Error(), http.StatusInternalServerError)
					}
			}()
		inner.ServeHTTP(w, r)

		log.Printf(
			"%s\t%s\t%s\t%s",
			r.Method,
			r.RequestURI,
			name,
			time.Since(start),
		)
	})
}
<|endoftext|>