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

// DNS server implementation.

package dns

import (
	"io"
	"os"
	"net"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// TODO(mg): fit axfr responses in here too
// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the address of the client that sent the current request
	RemoteAddr() string

	Write([]byte) (int, os.Error)

	// IP based ACL mapping. The contains the string representation
	// of the IP address and a boolean saying it may connect (true) or not.
	Acl() map[string]bool

	// Tsig secrets. Its a mapping of key names to secrets.
	Tsig() map[string]string
}

type conn struct {
	remoteAddr net.Addr     // address of remote side (sans port)
	port       int          // port of the remote side, needed TODO(mg)
	handler    Handler      // request handler
	request    []byte       // bytes read
	_UDP       *net.UDPConn // i/o connection if UDP was used
	_TCP       *net.TCPConn // i/o connection if TCP was used
	hijacked   bool         // connection has been hijacked by hander TODO(mg)
        tsig       map[string]string    // tsig secrets
        acl        map[string]bool      // ip acl list
}

type response struct {
	conn *conn
	req  *Msg
	xfr  bool // {i/a}xfr was requested
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type ServeMux struct {
	m map[string]Handler
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{make(map[string]Handler)} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, reg)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Helper handlers

func Refused(w ResponseWriter, r *Msg) {
        m := new(Msg)
        m.SetReply(r)
        m.MsgHdr.Rcode = RcodeRefused
        buf, _ := m.Pack()
        w.Write(buf)
}

// RefusedHandler return a REFUSED answer
func RefusedHandler() Handler { return HandlerFunc(Refused) }

func ListenAndServe(addr string, network string, handler Handler) os.Error {
	server := &Server{Addr: addr, Network: network, Handler: handler}
	return server.ListenAndServe()
}

func zoneMatch(pattern, zone string) bool {
	if len(pattern) == 0 {
		return false
	}
	n := len(pattern)
        var _ = n
        // better matching from the right TODO(mg)
	return true
}

func (mux *ServeMux) match(zone string) Handler {
	var h Handler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m[pattern] = handler
}

func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message.
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name)
	if h == nil {
		h = RefusedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle register the handler the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patters are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// Serve accepts incoming DNS request on the TCP listener l,
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func ServeTCP(l *net.TCPListener, handler Handler) os.Error {
	srv := &Server{Handler: handler, Network: "tcp"}
	return srv.ServeTCP(l)
}

// Serve accepts incoming DNS request on the UDP Conn l,
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func ServeUDP(l *net.UDPConn, handler Handler) os.Error {
	srv := &Server{Handler: handler, Network: "udp"}
	return srv.ServeUDP(l)
}

// A Server defines parameters for running an HTTP server.
type Server struct {
	Addr         string  // address to listen on, ":dns" if empty
	Network      string  // If "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler // handler to invoke, http.DefaultServeMux if nil
	ReadTimeout  int64   // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64   // the net.Conn.SetWriteTimeout value for new connections
}

// Fixes for udp/tcp
func (srv *Server) ListenAndServe() os.Error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Network {
	case "tcp":
		a, e := net.ResolveTCPAddr(addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP("tcp", a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp":
		a, e := net.ResolveUDPAddr(addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP("udp", a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return nil // os.Error with wrong network
}

func (srv *Server) ServeTCP(l *net.TCPListener) os.Error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadTimeout(srv.ReadTimeout)
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteTimeout(srv.WriteTimeout)
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func (srv *Server) ServeUDP(l *net.UDPConn) os.Error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	for {
		m := make([]byte, DefaultMsgSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadTimeout(srv.ReadTimeout)
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteTimeout(srv.WriteTimeout)
		}
		d, err := newConn(nil, l, a, m, handler)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler) (*conn, os.Error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	if t != nil {
		c.port = a.(*net.TCPAddr).Port
	}
	if u != nil {
		c.port = a.(*net.UDPAddr).Port
	}
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
        for {
                // Request has been read in ServeUDP or ServeTCP
                w := new(response)
                w.conn = c
                w.xfr = false
                req := new(Msg)
                if !req.Unpack(c.request) {
                        break
                }
                w.req = req
                c.handler.ServeDNS(w, w.req) // this does the writing back to the client
                if c.hijacked {
                        return
                }
                break           // TODO(mg) Why is this a loop anyway
        }
        if c._TCP != nil {
                c.close() // Listen and Serve is closed then
        }
}

func (w *response) Write(data []byte) (n int, err os.Error) {
        switch {
        case w.conn._UDP != nil:
                n, err = w.conn._UDP.WriteTo(data, w.conn.remoteAddr)
                if err != nil {
                        return 0, err
                }
        case w.conn._TCP != nil:
                // TODO(mg) len(data) > 64K
                l := make([]byte, 2)
                l[0], l[1] = packUint16(uint16(len(data)))
                n, err = w.conn._TCP.Write(l)
                if err != nil {
                        return n, err
                }
                if n != 2 {
                        return n, io.ErrShortWrite
                }
                n, err = w.conn._TCP.Write(data)
                if err != nil {
                        return n, err
                }
                i := n
                if i < len(data) {
                        j, err := w.conn._TCP.Write(data[i:len(data)])
                        if err != nil {
                                return i, err
                        }
                        i += j
                }
                n = i
        }
        return n, nil
}

// Acl implements the ResponseWriter.Acl
func (w *response) Acl() map[string]bool { return w.conn.acl }

// Tsig implements the ResponseWriter.Tsig
func (w *response) Tsig() map[string]string { return w.conn.tsig }

// RemoteAddr implements the ResponseWriter.RemoteAddr method
func (w *response) RemoteAddr() string { return w.conn.remoteAddr.String() }
<commit_msg>dont set auth bit in refused responses<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS server implementation.

package dns

import (
	"io"
	"os"
	"net"
)

type Handler interface {
	ServeDNS(w ResponseWriter, r *Msg)
}

// TODO(mg): fit axfr responses in here too
// A ResponseWriter interface is used by an DNS handler to
// construct an DNS response.
type ResponseWriter interface {
	// RemoteAddr returns the address of the client that sent the current request
	RemoteAddr() string

	Write([]byte) (int, os.Error)

	// IP based ACL mapping. The contains the string representation
	// of the IP address and a boolean saying it may connect (true) or not.
	Acl() map[string]bool

	// Tsig secrets. Its a mapping of key names to secrets.
	Tsig() map[string]string
}

type conn struct {
	remoteAddr net.Addr     // address of remote side (sans port)
	port       int          // port of the remote side, needed TODO(mg)
	handler    Handler      // request handler
	request    []byte       // bytes read
	_UDP       *net.UDPConn // i/o connection if UDP was used
	_TCP       *net.TCPConn // i/o connection if TCP was used
	hijacked   bool         // connection has been hijacked by hander TODO(mg)
        tsig       map[string]string    // tsig secrets
        acl        map[string]bool      // ip acl list
}

type response struct {
	conn *conn
	req  *Msg
	xfr  bool // {i/a}xfr was requested
}

// ServeMux is an DNS request multiplexer. It matches the
// zone name of each incoming request against a list of
// registered patterns add calls the handler for the pattern
// that most closely matches the zone name.
type ServeMux struct {
	m map[string]Handler
}

// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux { return &ServeMux{make(map[string]Handler)} }

// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()

// The HandlerFunc type is an adapter to allow the use of
// ordinary functions as DNS handlers.  If f is a function
// with the appropriate signature, HandlerFunc(f) is a
// Handler object that calls f.
type HandlerFunc func(ResponseWriter, *Msg)

// ServerDNS calls f(w, reg)
func (f HandlerFunc) ServeDNS(w ResponseWriter, r *Msg) {
	f(w, r)
}

// Helper handlers

func Refused(w ResponseWriter, r *Msg) {
        m := new(Msg)
        m.SetReply(r)
        m.MsgHdr.Rcode = RcodeRefused
        m.MsgHdr.Authoritative = false
        buf, _ := m.Pack()
        w.Write(buf)
}

// RefusedHandler return a REFUSED answer
func RefusedHandler() Handler { return HandlerFunc(Refused) }

func ListenAndServe(addr string, network string, handler Handler) os.Error {
	server := &Server{Addr: addr, Network: network, Handler: handler}
	return server.ListenAndServe()
}

func zoneMatch(pattern, zone string) bool {
	if len(pattern) == 0 {
		return false
	}
	n := len(pattern)
        var _ = n
        // better matching from the right TODO(mg)
	return true
}

func (mux *ServeMux) match(zone string) Handler {
	var h Handler
	var n = 0
	for k, v := range mux.m {
		if !zoneMatch(k, zone) {
			continue
		}
		if h == nil || len(k) > n {
			n = len(k)
			h = v
		}
	}
	return h
}

func (mux *ServeMux) Handle(pattern string, handler Handler) {
	if pattern == "" {
		panic("dns: invalid pattern " + pattern)
	}
	mux.m[pattern] = handler
}

func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	mux.Handle(pattern, HandlerFunc(handler))
}

// ServeDNS dispatches the request to the handler whose
// pattern most closely matches the request message.
func (mux *ServeMux) ServeDNS(w ResponseWriter, request *Msg) {
	h := mux.match(request.Question[0].Name)
	if h == nil {
		h = RefusedHandler()
	}
	h.ServeDNS(w, request)
}

// Handle register the handler the given pattern
// in the DefaultServeMux. The documentation for
// ServeMux explains how patters are matched.
func Handle(pattern string, handler Handler) { DefaultServeMux.Handle(pattern, handler) }

func HandleFunc(pattern string, handler func(ResponseWriter, *Msg)) {
	DefaultServeMux.HandleFunc(pattern, handler)
}

// Serve accepts incoming DNS request on the TCP listener l,
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func ServeTCP(l *net.TCPListener, handler Handler) os.Error {
	srv := &Server{Handler: handler, Network: "tcp"}
	return srv.ServeTCP(l)
}

// Serve accepts incoming DNS request on the UDP Conn l,
// creating a new service thread for each.  The service threads
// read requests and then call handler to reply to them.
// Handler is typically nil, in which case the DefaultServeMux is used.
func ServeUDP(l *net.UDPConn, handler Handler) os.Error {
	srv := &Server{Handler: handler, Network: "udp"}
	return srv.ServeUDP(l)
}

// A Server defines parameters for running an HTTP server.
type Server struct {
	Addr         string  // address to listen on, ":dns" if empty
	Network      string  // If "tcp" it will invoke a TCP listener, otherwise an UDP one
	Handler      Handler // handler to invoke, http.DefaultServeMux if nil
	ReadTimeout  int64   // the net.Conn.SetReadTimeout value for new connections
	WriteTimeout int64   // the net.Conn.SetWriteTimeout value for new connections
}

// Fixes for udp/tcp
func (srv *Server) ListenAndServe() os.Error {
	addr := srv.Addr
	if addr == "" {
		addr = ":domain"
	}
	switch srv.Network {
	case "tcp":
		a, e := net.ResolveTCPAddr(addr)
		if e != nil {
			return e
		}
		l, e := net.ListenTCP("tcp", a)
		if e != nil {
			return e
		}
		return srv.ServeTCP(l)
	case "udp":
		a, e := net.ResolveUDPAddr(addr)
		if e != nil {
			return e
		}
		l, e := net.ListenUDP("udp", a)
		if e != nil {
			return e
		}
		return srv.ServeUDP(l)
	}
	return nil // os.Error with wrong network
}

func (srv *Server) ServeTCP(l *net.TCPListener) os.Error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
forever:
	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			return e
		}
		if srv.ReadTimeout != 0 {
			rw.SetReadTimeout(srv.ReadTimeout)
		}
		if srv.WriteTimeout != 0 {
			rw.SetWriteTimeout(srv.WriteTimeout)
		}
		l := make([]byte, 2)
		n, err := rw.Read(l)
		if err != nil || n != 2 {
			continue
		}
		length, _ := unpackUint16(l, 0)
		if length == 0 {
			continue
		}
		m := make([]byte, int(length))
		n, err = rw.Read(m[:int(length)])
		if err != nil {
			continue
		}
		i := n
		for i < int(length) {
			j, err := rw.Read(m[i:int(length)])
			if err != nil {
				continue forever
			}
			i += j
		}
		n = i
		d, err := newConn(rw, nil, rw.RemoteAddr(), m, handler)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func (srv *Server) ServeUDP(l *net.UDPConn) os.Error {
	defer l.Close()
	handler := srv.Handler
	if handler == nil {
		handler = DefaultServeMux
	}
	for {
		m := make([]byte, DefaultMsgSize)
		n, a, e := l.ReadFromUDP(m)
		if e != nil {
			return e
		}
		m = m[:n]

		if srv.ReadTimeout != 0 {
			l.SetReadTimeout(srv.ReadTimeout)
		}
		if srv.WriteTimeout != 0 {
			l.SetWriteTimeout(srv.WriteTimeout)
		}
		d, err := newConn(nil, l, a, m, handler)
		if err != nil {
			continue
		}
		go d.serve()
	}
	panic("not reached")
}

func newConn(t *net.TCPConn, u *net.UDPConn, a net.Addr, buf []byte, handler Handler) (*conn, os.Error) {
	c := new(conn)
	c.handler = handler
	c._TCP = t
	c._UDP = u
	c.remoteAddr = a
	c.request = buf
	if t != nil {
		c.port = a.(*net.TCPAddr).Port
	}
	if u != nil {
		c.port = a.(*net.UDPAddr).Port
	}
	return c, nil
}

// Close the connection.
func (c *conn) close() {
	switch {
	case c._UDP != nil:
		c._UDP.Close()
		c._UDP = nil
	case c._TCP != nil:
		c._TCP.Close()
		c._TCP = nil
	}
}

// Serve a new connection.
func (c *conn) serve() {
        for {
                // Request has been read in ServeUDP or ServeTCP
                w := new(response)
                w.conn = c
                w.xfr = false
                req := new(Msg)
                if !req.Unpack(c.request) {
                        break
                }
                w.req = req
                c.handler.ServeDNS(w, w.req) // this does the writing back to the client
                if c.hijacked {
                        return
                }
                break           // TODO(mg) Why is this a loop anyway
        }
        if c._TCP != nil {
                c.close() // Listen and Serve is closed then
        }
}

func (w *response) Write(data []byte) (n int, err os.Error) {
        switch {
        case w.conn._UDP != nil:
                n, err = w.conn._UDP.WriteTo(data, w.conn.remoteAddr)
                if err != nil {
                        return 0, err
                }
        case w.conn._TCP != nil:
                // TODO(mg) len(data) > 64K
                l := make([]byte, 2)
                l[0], l[1] = packUint16(uint16(len(data)))
                n, err = w.conn._TCP.Write(l)
                if err != nil {
                        return n, err
                }
                if n != 2 {
                        return n, io.ErrShortWrite
                }
                n, err = w.conn._TCP.Write(data)
                if err != nil {
                        return n, err
                }
                i := n
                if i < len(data) {
                        j, err := w.conn._TCP.Write(data[i:len(data)])
                        if err != nil {
                                return i, err
                        }
                        i += j
                }
                n = i
        }
        return n, nil
}

// Acl implements the ResponseWriter.Acl
func (w *response) Acl() map[string]bool { return w.conn.acl }

// Tsig implements the ResponseWriter.Tsig
func (w *response) Tsig() map[string]string { return w.conn.tsig }

// RemoteAddr implements the ResponseWriter.RemoteAddr method
func (w *response) RemoteAddr() string { return w.conn.remoteAddr.String() }
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/fortuna/ss-example/metrics"
	"github.com/shadowsocks/go-shadowsocks2/core"
	ssnet "github.com/shadowsocks/go-shadowsocks2/net"
	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

var config struct {
	UDPTimeout time.Duration
}

func shadowConn(conn ssnet.DuplexConn, cipherList []shadowaead.Cipher) (ssnet.DuplexConn, int, error) {
	cipher, index, shadowReader, err := findCipher(conn, cipherList)
	if err != nil {
		return nil, -1, err
	}
	shadowWriter := shadowaead.NewShadowsocksWriter(conn, cipher)
	return ssnet.WrapDuplexConn(conn, shadowReader, shadowWriter), index, nil
}

func findCipher(clientReader io.Reader, cipherList []shadowaead.Cipher) (shadowaead.Cipher, int, io.Reader, error) {
	if len(cipherList) == 0 {
		return nil, -1, nil, errors.New("Empty cipher list")
	} else if len(cipherList) == 1 {
		return cipherList[0], 0, shadowaead.NewShadowsocksReader(clientReader, cipherList[0]), nil
	}
	// buffer saves the bytes read from shadowConn, in order to allow for replays.
	var buffer bytes.Buffer
	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	for i, cipher := range cipherList {
		log.Printf("Trying cipher %v", i)
		// tmpReader reuses the bytes read so far, falling back to shadowConn if it needs more
		// bytes. All bytes read from shadowConn are saved in buffer.
		tmpReader := io.MultiReader(bytes.NewReader(buffer.Bytes()), io.TeeReader(clientReader, &buffer))
		// Override the Reader of shadowConn so we can reset it for each cipher test.
		cipherReader := shadowaead.NewShadowsocksReader(tmpReader, cipher)
		// Read should read just enough data to authenticate the payload size.
		_, err := cipherReader.Read(make([]byte, 0))
		if err != nil {
			log.Printf("Failed cipher %v", i)
			continue
		}
		log.Printf("Selected cipher %v", i)
		// We don't need to replay the bytes anymore, but we don't want to drop those
		// read so far.
		return cipher, i, shadowaead.NewShadowsocksReader(io.MultiReader(&buffer, clientReader), cipher), nil
	}
	return nil, -1, nil, fmt.Errorf("could not find valid cipher")
}

func getNetKey(addr net.Addr) (string, error) {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return "", err
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return "", errors.New("Failed to parse ip")
	}
	ipNet := net.IPNet{IP: ip}
	if ip.To4() != nil {
		ipNet.Mask = net.CIDRMask(24, 32)
	} else {
		ipNet.Mask = net.CIDRMask(32, 128)
	}
	return ipNet.String(), nil
}

// Listen on addr for incoming connections.
func tcpRemote(addr string, cipherList []shadowaead.Cipher) {
	accessKeyMetrics := metrics.NewMetricsMap()
	netMetrics := metrics.NewMetricsMap()
	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Printf("failed to listen on %s: %v", addr, err)
		return
	}

	log.Printf("listening TCP on %s", addr)
	for {
		var clientConn ssnet.DuplexConn
		clientConn, err := l.(*net.TCPListener).AcceptTCP()
		if err != nil {
			log.Printf("failed to accept: %v", err)
			return
		}

		go func() {
			defer clientConn.Close()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			accessKey := "INVALID"
			netKey, err := getNetKey(clientConn.RemoteAddr())
			if err != nil {
				netKey = "INVALID"
			}
			var proxyMetrics metrics.ProxyMetrics
			defer func() {
				log.Printf("Done")
				accessKeyMetrics.Add(accessKey, proxyMetrics)
				log.Printf("Key %v: %s", accessKey, metrics.SPrintMetrics(accessKeyMetrics.Get(accessKey)))
				netMetrics.Add(netKey, proxyMetrics)
				log.Printf("Net %v: %s", netKey, metrics.SPrintMetrics(netMetrics.Get(netKey)))
			}()

			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ClientProxy, &proxyMetrics.ProxyClient)
			clientConn, index, err := shadowConn(clientConn, cipherList)
			if err != nil {
				log.Printf("Failed to find a valid cipher: %v", err)
				return
			}
			accessKey = strconv.Itoa(index)

			tgt, err := socks.ReadAddr(clientConn)
			if err != nil {
				log.Printf("failed to get target address: %v", err)
				return
			}

			c, err := net.Dial("tcp", tgt.String())
			if err != nil {
				log.Printf("failed to connect to target: %v", err)
				return
			}
			var tgtConn ssnet.DuplexConn = c.(*net.TCPConn)
			defer tgtConn.Close()
			tgtConn.(*net.TCPConn).SetKeepAlive(true)
			tgtConn = metrics.MeasureConn(tgtConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			log.Printf("proxy %s <-> %s", clientConn.RemoteAddr(), tgt)
			_, _, err = ssnet.Relay(clientConn, tgtConn)
			if err != nil {
				log.Printf("relay error: %v", err)
			}
		}()
	}
}

type cipherList []shadowaead.Cipher

func main() {

	var flags struct {
		Server  string
		Ciphers cipherList
	}

	flag.StringVar(&flags.Server, "s", "", "server listen address")
	flag.Var(&flags.Ciphers, "u", "available ciphers: "+strings.Join(core.ListCipher(), " "))
	flag.DurationVar(&config.UDPTimeout, "udptimeout", 5*time.Minute, "UDP tunnel timeout")
	flag.Parse()

	if flags.Server == "" || len(flags.Ciphers) == 0 {
		flag.Usage()
		return
	}

	go udpRemote(flags.Server, flags.Ciphers)
	go tcpRemote(flags.Server, flags.Ciphers)

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh
}

func (sl *cipherList) Set(flagValue string) error {
	e := strings.SplitN(flagValue, ":", 2)
	if len(e) != 2 {
		return fmt.Errorf("Missing colon")
	}
	cipher, err := core.PickCipher(e[0], nil, e[1])
	if err != nil {
		return err
	}
	aead, ok := cipher.(shadowaead.Cipher)
	if !ok {
		log.Fatal("Only AEAD ciphers are supported")
	}
	*sl = append(*sl, aead)
	return nil
}

func (sl *cipherList) String() string {
	return fmt.Sprint(*sl)
}
<commit_msg>Add more debugging messages<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/fortuna/ss-example/metrics"
	"github.com/shadowsocks/go-shadowsocks2/core"
	ssnet "github.com/shadowsocks/go-shadowsocks2/net"
	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
	"github.com/shadowsocks/go-shadowsocks2/socks"
)

var config struct {
	UDPTimeout time.Duration
}

func shadowConn(conn ssnet.DuplexConn, cipherList []shadowaead.Cipher) (ssnet.DuplexConn, int, error) {
	cipher, index, shadowReader, err := findCipher(conn, cipherList)
	if err != nil {
		return nil, -1, err
	}
	shadowWriter := shadowaead.NewShadowsocksWriter(conn, cipher)
	return ssnet.WrapDuplexConn(conn, shadowReader, shadowWriter), index, nil
}

func findCipher(clientReader io.Reader, cipherList []shadowaead.Cipher) (shadowaead.Cipher, int, io.Reader, error) {
	if len(cipherList) == 0 {
		return nil, -1, nil, errors.New("Empty cipher list")
	} else if len(cipherList) == 1 {
		return cipherList[0], 0, shadowaead.NewShadowsocksReader(clientReader, cipherList[0]), nil
	}
	// buffer saves the bytes read from shadowConn, in order to allow for replays.
	var buffer bytes.Buffer
	// Try each cipher until we find one that authenticates successfully.
	// This assumes that all ciphers are AEAD.
	for i, cipher := range cipherList {
		log.Printf("Trying cipher %v", i)
		// tmpReader reuses the bytes read so far, falling back to shadowConn if it needs more
		// bytes. All bytes read from shadowConn are saved in buffer.
		tmpReader := io.MultiReader(bytes.NewReader(buffer.Bytes()), io.TeeReader(clientReader, &buffer))
		// Override the Reader of shadowConn so we can reset it for each cipher test.
		cipherReader := shadowaead.NewShadowsocksReader(tmpReader, cipher)
		// Read should read just enough data to authenticate the payload size.
		_, err := cipherReader.Read(make([]byte, 0))
		if err != nil {
			log.Printf("Failed cipher %v: %v", i, err)
			continue
		}
		log.Printf("Selected cipher %v", i)
		// We don't need to replay the bytes anymore, but we don't want to drop those
		// read so far.
		return cipher, i, shadowaead.NewShadowsocksReader(io.MultiReader(&buffer, clientReader), cipher), nil
	}
	return nil, -1, nil, fmt.Errorf("could not find valid cipher")
}

func getNetKey(addr net.Addr) (string, error) {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return "", err
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return "", errors.New("Failed to parse ip")
	}
	ipNet := net.IPNet{IP: ip}
	if ip.To4() != nil {
		ipNet.Mask = net.CIDRMask(24, 32)
	} else {
		ipNet.Mask = net.CIDRMask(32, 128)
	}
	return ipNet.String(), nil
}

// Listen on addr for incoming connections.
func tcpRemote(addr string, cipherList []shadowaead.Cipher) {
	accessKeyMetrics := metrics.NewMetricsMap()
	netMetrics := metrics.NewMetricsMap()
	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Printf("failed to listen on %s: %v", addr, err)
		return
	}

	log.Printf("listening TCP on %s", addr)
	for {
		var clientConn ssnet.DuplexConn
		clientConn, err := l.(*net.TCPListener).AcceptTCP()
		if err != nil {
			log.Printf("failed to accept: %v", err)
			return
		}

		go func() {
			defer clientConn.Close()
			clientConn.(*net.TCPConn).SetKeepAlive(true)
			accessKey := "INVALID"
			netKey, err := getNetKey(clientConn.RemoteAddr())
			if err != nil {
				netKey = "INVALID"
			}
			var proxyMetrics metrics.ProxyMetrics
			defer func() {
				log.Printf("Done")
				accessKeyMetrics.Add(accessKey, proxyMetrics)
				log.Printf("Key %v: %s", accessKey, metrics.SPrintMetrics(accessKeyMetrics.Get(accessKey)))
				netMetrics.Add(netKey, proxyMetrics)
				log.Printf("Net %v: %s", netKey, metrics.SPrintMetrics(netMetrics.Get(netKey)))
			}()

			clientConn = metrics.MeasureConn(clientConn, &proxyMetrics.ClientProxy, &proxyMetrics.ProxyClient)
			clientConn, index, err := shadowConn(clientConn, cipherList)
			if err != nil {
				log.Printf("Failed to find a valid cipher: %v", err)
				return
			}
			accessKey = strconv.Itoa(index)

			tgt, err := socks.ReadAddr(clientConn)
			if err != nil {
				log.Printf("failed to get target address: %v", err)
				return
			}

			c, err := net.Dial("tcp", tgt.String())
			if err != nil {
				log.Printf("failed to connect to target: %v", err)
				return
			}
			var tgtConn ssnet.DuplexConn = c.(*net.TCPConn)
			defer tgtConn.Close()
			tgtConn.(*net.TCPConn).SetKeepAlive(true)
			tgtConn = metrics.MeasureConn(tgtConn, &proxyMetrics.ProxyTarget, &proxyMetrics.TargetProxy)

			log.Printf("proxy %s <-> %s", clientConn.RemoteAddr(), tgt)
			_, _, err = ssnet.Relay(clientConn, tgtConn)
			if err != nil {
				log.Printf("relay error: %v", err)
			}
		}()
	}
}

type cipherList []shadowaead.Cipher

func main() {

	var flags struct {
		Server  string
		Ciphers cipherList
	}

	flag.StringVar(&flags.Server, "s", "", "server listen address")
	flag.Var(&flags.Ciphers, "u", "available ciphers: "+strings.Join(core.ListCipher(), " "))
	flag.DurationVar(&config.UDPTimeout, "udptimeout", 5*time.Minute, "UDP tunnel timeout")
	flag.Parse()

	if flags.Server == "" || len(flags.Ciphers) == 0 {
		flag.Usage()
		return
	}

	go udpRemote(flags.Server, flags.Ciphers)
	go tcpRemote(flags.Server, flags.Ciphers)

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	<-sigCh
}

func (sl *cipherList) Set(flagValue string) error {
	e := strings.SplitN(flagValue, ":", 2)
	if len(e) != 2 {
		return fmt.Errorf("Missing colon")
	}
	cipher, err := core.PickCipher(e[0], nil, e[1])
	if err != nil {
		return err
	}
	aead, ok := cipher.(shadowaead.Cipher)
	if !ok {
		log.Fatal("Only AEAD ciphers are supported")
	}
	*sl = append(*sl, aead)
	return nil
}

func (sl *cipherList) String() string {
	return fmt.Sprint(*sl)
}
<|endoftext|>


<commit_before>package mithril

import (
	"fmt"
	"net/http"

	log "github.com/Sirupsen/logrus"

	"github.com/modcloth/mithril/message"
	"github.com/modcloth/mithril/store"
)

type Server struct {
	amqp    *AMQPPublisher
	storage *store.Storage
}

func NewServer(storer *store.Storage, amqp *AMQPPublisher) *Server {
	return &Server{
		storage: storer,
		amqp:    amqp,
	}
}

func (me *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" || r.Method == "PUT" {
		me.processMessage(w, r)
	} else {
		me.respondErr(
			fmt.Errorf(`Only "POST" and "PUT" are accepted, not %q`, r.Method),
			http.StatusMethodNotAllowed,
			w)
	}
}

func (me *Server) processMessage(w http.ResponseWriter, r *http.Request) {
	var (
		msg *message.Message
		err error
	)
	if msg, err = message.NewMessage(r); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	log.Infof("Processing message: ", msg.MessageId)
	if err = me.storage.Store(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	if err = me.amqp.Publish(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	w.WriteHeader(http.StatusNoContent)
	w.Write([]byte(""))
}

func (me *Server) respondErr(err error, status int, w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(status)
	fmt.Fprintf(w, "WOMP WOMP: %v\n", err)
}
<commit_msg>Adding healthcheck url of /heartbeat<commit_after>package mithril

import (
	"fmt"
	"net/http"

	log "github.com/Sirupsen/logrus"

	"github.com/modcloth/mithril/message"
	"github.com/modcloth/mithril/store"
)

type Server struct {
	amqp    *AMQPPublisher
	storage *store.Storage
}

func NewServer(storer *store.Storage, amqp *AMQPPublisher) *Server {
	return &Server{
		storage: storer,
		amqp:    amqp,
	}
}

func (me *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" || r.Method == "PUT" {
		me.processMessage(w, r)
	} else if r.URL.Path == "/heartbeat" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "pong\n")
	} else {
		me.respondErr(
			fmt.Errorf(`Only "POST" and "PUT" are accepted, not %q`, r.Method),
			http.StatusMethodNotAllowed,
			w)
	}
}

func (me *Server) processMessage(w http.ResponseWriter, r *http.Request) {
	var (
		msg *message.Message
		err error
	)
	if msg, err = message.NewMessage(r); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	log.Infof("Processing message: ", msg.MessageId)
	if err = me.storage.Store(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	if err = me.amqp.Publish(msg); err != nil {
		me.respondErr(err, http.StatusBadRequest, w)
		return
	}

	w.WriteHeader(http.StatusNoContent)
	w.Write([]byte(""))
}

func (me *Server) respondErr(err error, status int, w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(status)
	fmt.Fprintf(w, "WOMP WOMP: %v\n", err)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"strconv"

	"runtime"

	"github.com/codegangsta/negroni"
	"github.com/julienschmidt/httprouter"
)

const (
	// StopAfterFailedInserts is the number of failed inserts after which we will stop fetching from Instagram.
	// Inserts fail when the image already exists, 50 was chosen to allow for images that are fetched out of order.
	StopAfterFailedInserts = 50
)

func ping(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
	err := p.Ping()
	if err != nil {
		http.Error(w, "Ping failed", 500)
	}
	fmt.Fprint(w, "pong")
}

func stats(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	cnt, err := notificationCount()
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
	fmt.Fprint(w, "Notifications: ", string(cnt))
}

func fetchMediaForSubscription(sid string, stopAfter int) int {
	sub := getSubscription(sid)
	failed := 0
	counter := 0
	fetchQueue := make(chan *Media)
	stop := make(chan struct{})
	go fetchMedia(sub, fetchQueue, stop)
	for m := range fetchQueue {
		new, err := insert(*m)
		if err != nil {
			fmt.Printf("Error encountered %v when inserting\n", err)
		}
		if !new {
			failed++
			if failed == stopAfter {
				fmt.Printf("Stopping after %v\n", stopAfter)
				stop <- struct{}{}
			}
		} else {
			failed = 0
			counter++
		}
	}
	return counter
}

func fetch(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	sid := p.ByName("sid")
	r.ParseForm()
	stopAfter := StopAfterFailedInserts
	v := r.FormValue("stopAfter")
	if v != "" {
		sa, err := strconv.Atoi(v)
		if err == nil {
			stopAfter = sa
		}
	}
	counter := fetchMediaForSubscription(sid, stopAfter)
	fmt.Fprintf(w, "Fetch completed, fetched %d\n", counter)
}

func initAPI() *negroni.Negroni {
	n := negroni.New(
		negroni.NewRecovery(),
		negroni.NewLogger(), negroni.NewStatic(http.Dir("ui")))
	router := httprouter.New()
	router.GET("/ping", ping)
	router.GET("/stats", stats)
	router.GET("/fetch/:sid", fetch)
	n.UseHandler(router)
	return n
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	db := initDb()
	defer db.Close()

	initClient()

	n := initAPI()

	host := getEnvOrDefault("SPOTO_HOST", "localhost")
	port := getEnvOrDefault("SPOTO_PORT", "3000")
	bindTo := fmt.Sprintf("%s:%s", host, port)

	n.Run(bindTo)
}

func getEnvOrDefault(key, def string) string {
	env := os.Getenv(key)
	if len(env) == 0 {
		env = def
	}
	return env
}

func checkErr(err error) {
	if err != nil {
		panic(err)
	}
}
<commit_msg>Remove printf<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"strconv"

	"runtime"

	"github.com/codegangsta/negroni"
	"github.com/julienschmidt/httprouter"
)

const (
	// StopAfterFailedInserts is the number of failed inserts after which we will stop fetching from Instagram.
	// Inserts fail when the image already exists, 50 was chosen to allow for images that are fetched out of order.
	StopAfterFailedInserts = 50
)

func ping(w http.ResponseWriter, _ *http.Request, _ httprouter.Params) {
	err := p.Ping()
	if err != nil {
		http.Error(w, "Ping failed", 500)
	}
	fmt.Fprint(w, "pong")
}

func stats(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	cnt, err := notificationCount()
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
	fmt.Fprint(w, "Notifications: ", string(cnt))
}

func fetchMediaForSubscription(sid string, stopAfter int) int {
	sub := getSubscription(sid)
	failed := 0
	counter := 0
	fetchQueue := make(chan *Media)
	stop := make(chan struct{})
	go fetchMedia(sub, fetchQueue, stop)
	for m := range fetchQueue {
		new, err := insert(*m)
		if err != nil {
			fmt.Printf("Error encountered %v when inserting\n", err)
		}
		if !new {
			failed++
			if failed == stopAfter {
				stop <- struct{}{}
			}
		} else {
			failed = 0
			counter++
		}
	}
	return counter
}

func fetch(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	sid := p.ByName("sid")
	r.ParseForm()
	stopAfter := StopAfterFailedInserts
	v := r.FormValue("stopAfter")
	if v != "" {
		sa, err := strconv.Atoi(v)
		if err == nil {
			stopAfter = sa
		}
	}
	counter := fetchMediaForSubscription(sid, stopAfter)
	fmt.Fprintf(w, "Fetch completed, fetched %d\n", counter)
}

func initAPI() *negroni.Negroni {
	n := negroni.New(
		negroni.NewRecovery(),
		negroni.NewLogger(), negroni.NewStatic(http.Dir("ui")))
	router := httprouter.New()
	router.GET("/ping", ping)
	router.GET("/stats", stats)
	router.GET("/fetch/:sid", fetch)
	n.UseHandler(router)
	return n
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	db := initDb()
	defer db.Close()

	initClient()

	n := initAPI()

	host := getEnvOrDefault("SPOTO_HOST", "localhost")
	port := getEnvOrDefault("SPOTO_PORT", "3000")
	bindTo := fmt.Sprintf("%s:%s", host, port)

	n.Run(bindTo)
}

func getEnvOrDefault(key, def string) string {
	env := os.Getenv(key)
	if len(env) == 0 {
		env = def
	}
	return env
}

func checkErr(err error) {
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package event

import (
	"encoding/json"
	"testing"

	"github.com/weaveworks/flux/update"
)

var (
	spec = update.ReleaseImageSpec{
		ImageSpec: update.ImageSpecLatest,
	}
	cause = update.Cause{
		User:    "test user",
		Message: "test message",
	}
)

func TestEvent_ParseReleaseMetaData(t *testing.T) {
	origEvent := Event{
		Type: EventRelease,
		Metadata: &ReleaseEventMetadata{
			Cause: cause,
			Spec: ReleaseSpec{
				Type:             ReleaseImageSpecType,
				ReleaseImageSpec: &spec,
			},
		},
	}

	bytes, _ := json.Marshal(origEvent)

	e := Event{}
	err := e.UnmarshalJSON(bytes)
	if err != nil {
		t.Fatal(err)
	}
	switch r := e.Metadata.(type) {
	case *ReleaseEventMetadata:
		if r.Spec.ReleaseImageSpec.ImageSpec != spec.ImageSpec ||
			r.Cause != cause {
			t.Fatal("Release event wasn't marshalled/unmarshalled")
		}
	default:
		t.Fatal("Wrong event type unmarshalled")
	}
}

func TestEvent_ParseNoMetadata(t *testing.T) {
	origEvent := Event{
		Type: EventLock,
	}

	bytes, _ := json.Marshal(origEvent)

	e := Event{}
	err := e.UnmarshalJSON(bytes)
	if err != nil {
		t.Fatal(err)
	}
	if e.Metadata != nil {
		t.Fatal("Hasn't been unmarshalled properly")
	}
}
<commit_msg>Add test for parsing old event format<commit_after>package event

import (
	"encoding/json"
	"testing"

	"github.com/weaveworks/flux/update"
)

var (
	spec = update.ReleaseImageSpec{
		ImageSpec: update.ImageSpecLatest,
	}
	cause = update.Cause{
		User:    "test user",
		Message: "test message",
	}
)

func TestEvent_ParseReleaseMetaData(t *testing.T) {
	origEvent := Event{
		Type: EventRelease,
		Metadata: &ReleaseEventMetadata{
			Cause: cause,
			Spec: ReleaseSpec{
				Type:             ReleaseImageSpecType,
				ReleaseImageSpec: &spec,
			},
		},
	}

	bytes, _ := json.Marshal(origEvent)

	e := Event{}
	err := e.UnmarshalJSON(bytes)
	if err != nil {
		t.Fatal(err)
	}
	switch r := e.Metadata.(type) {
	case *ReleaseEventMetadata:
		if r.Spec.ReleaseImageSpec.ImageSpec != spec.ImageSpec ||
			r.Cause != cause {
			t.Fatal("Release event wasn't marshalled/unmarshalled")
		}
	default:
		t.Fatal("Wrong event type unmarshalled")
	}
}

func TestEvent_ParseNoMetadata(t *testing.T) {
	origEvent := Event{
		Type: EventLock,
	}

	bytes, _ := json.Marshal(origEvent)

	e := Event{}
	err := e.UnmarshalJSON(bytes)
	if err != nil {
		t.Fatal(err)
	}
	if e.Metadata != nil {
		t.Fatal("Hasn't been unmarshalled properly")
	}
}

// TestEvent_ParseOldReleaseMetaData makes sure the parsing code can
// handle the older format events recorded against commits.
func TestEvent_ParseOldReleaseMetaData(t *testing.T) {
	// A minimal example of an old-style ReleaseEventMetadata. NB it
	// must have at least an entry for "spec", since otherwise the
	// JSON unmarshaller will not attempt to unparse the spec and
	// thereby invoke the specialised UnmarshalJSON.
	oldData := `
{
  "spec": {
    "serviceSpecs": ["<all>"]
  }
}
`
	var eventData ReleaseEventMetadata
	if err := json.Unmarshal([]byte(oldData), &eventData); err != nil {
		t.Fatal(err)
	}
	if eventData.Spec.Type != ReleaseImageSpecType {
		t.Error("did not set spec type to ReleaseImageSpecType")
	}
	if eventData.Spec.ReleaseImageSpec == nil {
		t.Error("did not set .ReleaseImageSpec as expected")
	}
	if eventData.Spec.ReleaseContainersSpec != nil {
		t.Error("unexpectedly set .ReleaseContainersSpec")
	}
	if len(eventData.Spec.ReleaseImageSpec.ServiceSpecs) != 1 {
		t.Error("expected service specs of len 1")
	}
}
<|endoftext|>


<commit_before>package assert

func ThatString(t TestingT, actual string) *StringAssert {
	return &StringAssert{t, actual}
}

type StringAssert struct {
	t      TestingT
	actual string
}

func (assert *StringAssert) isTrue(condition bool, format string, args ...interface{}) *StringAssert {
	if !condition {
		assert.t.Errorf(format, args...)
	}
	return assert
}

func (assert *StringAssert) IsEqualTo(expected string) *StringAssert {
	return assert.isTrue(assert.actual == expected,
		"Expected <%s>, but was <%s>.", expected, assert.actual)
}
<commit_msg>Change StringAssert struct name.<commit_after>package assert

func ThatString(t TestingT, actual string) *StringAssertImpl {
	return &StringAssertImpl{t, actual}
}

type StringAssertImpl struct {
	t      TestingT
	actual string
}

func (assert *StringAssertImpl) isTrue(condition bool, format string, args ...interface{}) *StringAssertImpl {
	if !condition {
		assert.t.Errorf(format, args...)
	}
	return assert
}

func (assert *StringAssertImpl) IsEqualTo(expected string) *StringAssertImpl {
	return assert.isTrue(assert.actual == expected,
		"Expected <%s>, but was <%s>.", expected, assert.actual)
}
<|endoftext|>


<commit_before>package bench

import (
	"fmt"
	"os"
	"time"

	"github.com/codahale/hdrhistogram"
)

// Summary contains the results of a Benchmark run.
type Summary struct {
	Connections                 uint64
	RequestRate                 uint64
	SuccessTotal                uint64
	ErrorTotal                  uint64
	TimeElapsed                 time.Duration
	SuccessHistogram            *hdrhistogram.Histogram
	UncorrectedSuccessHistogram *hdrhistogram.Histogram
	ErrorHistogram              *hdrhistogram.Histogram
	UncorrectedErrorHistogram   *hdrhistogram.Histogram
	Throughput                  float64
}

// String returns a stringified version of the Summary.
func (s *Summary) String() string {
	return fmt.Sprintf(
		"\n{Connections: %d, RequestRate: %d, RequestTotal: %d, SuccessTotal: %d, ErrorTotal: %d, TimeElapsed: %s, Throughput: %.2f/s}",
		s.Connections, s.RequestRate, (s.SuccessTotal + s.ErrorTotal), s.SuccessTotal, s.ErrorTotal, s.TimeElapsed, s.Throughput)
}

// GenerateLatencyDistribution generates a text file containing the specified
// latency distribution in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. Percentiles is a
// list of percentiles to include, e.g. 10.0, 50.0, 99.0, 99.99, etc. If
// percentiles is nil, it defaults to a logarithmic percentile scale. If a
// request rate was specified for the benchmark, this will also generate an
// uncorrected distribution file which does not account for coordinated
// omission.
func (s *Summary) GenerateLatencyDistribution(percentiles Percentiles, file string) error {
	return generateLatencyDistribution(s.SuccessHistogram, s.UncorrectedSuccessHistogram, s.RequestRate, percentiles, file)
}

// GenerateErrorLatencyDistribution generates a text file containing the specified
// latency distribution (of requests that resulted in errors) in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. Percentiles is a
// list of percentiles to include, e.g. 10.0, 50.0, 99.0, 99.99, etc. If
// percentiles is nil, it defaults to a logarithmic percentile scale. If a
// request rate was specified for the benchmark, this will also generate an
// uncorrected distribution file which does not account for coordinated
// omission.
func (s *Summary) GenerateErrorLatencyDistribution(percentiles Percentiles, file string) error {
	return generateLatencyDistribution(s.ErrorHistogram, s.UncorrectedErrorHistogram, s.RequestRate, percentiles, file)
}

func getOneByPercentile(percentile float64) float64 {
	if percentile < 1 {
		return 1 / (1 - (percentile / 100))
	}
	return float64(1000000)
}

func generateLatencyDistribution(histogram, unHistogram *hdrhistogram.Histogram, requestRate uint64, percentiles Percentiles, file string) error {
	if percentiles == nil {
		percentiles = Logarithmic
	}
	f, err := os.Create(file)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
	for _, percentile := range percentiles {
		value := float64(histogram.ValueAtQuantile(percentile)) / 1000000
		oneByPercentile := getOneByPercentile(percentile)
		_, err := f.WriteString(fmt.Sprintf("%f    %f        %d            %f\n",
			value, percentile/100, 0, oneByPercentile))
		if err != nil {
			return err
		}
	}

	// Generate uncorrected distribution.
	if requestRate > 0 {
		f, err := os.Create("uncorrected_" + file)
		if err != nil {
			return err
		}
		defer f.Close()

		f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
		for _, percentile := range percentiles {
			value := float64(unHistogram.ValueAtQuantile(percentile)) / 1000000
			oneByPercentile := getOneByPercentile(percentile)
			_, err := f.WriteString(fmt.Sprintf("%f    %f        %d            %f\n",
				value, percentile/100, 0, oneByPercentile))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// merge the other Summary into this one.
func (s *Summary) merge(o *Summary) {
	if o.TimeElapsed > s.TimeElapsed {
		s.TimeElapsed = o.TimeElapsed
	}
	s.SuccessHistogram.Merge(o.SuccessHistogram)
	s.UncorrectedSuccessHistogram.Merge(o.UncorrectedSuccessHistogram)
	s.ErrorHistogram.Merge(o.ErrorHistogram)
	s.UncorrectedErrorHistogram.Merge(o.UncorrectedErrorHistogram)
	s.SuccessTotal += o.SuccessTotal
	s.ErrorTotal += o.ErrorTotal
	s.Throughput += o.Throughput
	s.RequestRate += o.RequestRate
}
<commit_msg>Fix getOneByPercentile<commit_after>package bench

import (
	"fmt"
	"os"
	"time"

	"github.com/codahale/hdrhistogram"
)

// Summary contains the results of a Benchmark run.
type Summary struct {
	Connections                 uint64
	RequestRate                 uint64
	SuccessTotal                uint64
	ErrorTotal                  uint64
	TimeElapsed                 time.Duration
	SuccessHistogram            *hdrhistogram.Histogram
	UncorrectedSuccessHistogram *hdrhistogram.Histogram
	ErrorHistogram              *hdrhistogram.Histogram
	UncorrectedErrorHistogram   *hdrhistogram.Histogram
	Throughput                  float64
}

// String returns a stringified version of the Summary.
func (s *Summary) String() string {
	return fmt.Sprintf(
		"\n{Connections: %d, RequestRate: %d, RequestTotal: %d, SuccessTotal: %d, ErrorTotal: %d, TimeElapsed: %s, Throughput: %.2f/s}",
		s.Connections, s.RequestRate, (s.SuccessTotal + s.ErrorTotal), s.SuccessTotal, s.ErrorTotal, s.TimeElapsed, s.Throughput)
}

// GenerateLatencyDistribution generates a text file containing the specified
// latency distribution in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. Percentiles is a
// list of percentiles to include, e.g. 10.0, 50.0, 99.0, 99.99, etc. If
// percentiles is nil, it defaults to a logarithmic percentile scale. If a
// request rate was specified for the benchmark, this will also generate an
// uncorrected distribution file which does not account for coordinated
// omission.
func (s *Summary) GenerateLatencyDistribution(percentiles Percentiles, file string) error {
	return generateLatencyDistribution(s.SuccessHistogram, s.UncorrectedSuccessHistogram, s.RequestRate, percentiles, file)
}

// GenerateErrorLatencyDistribution generates a text file containing the specified
// latency distribution (of requests that resulted in errors) in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. Percentiles is a
// list of percentiles to include, e.g. 10.0, 50.0, 99.0, 99.99, etc. If
// percentiles is nil, it defaults to a logarithmic percentile scale. If a
// request rate was specified for the benchmark, this will also generate an
// uncorrected distribution file which does not account for coordinated
// omission.
func (s *Summary) GenerateErrorLatencyDistribution(percentiles Percentiles, file string) error {
	return generateLatencyDistribution(s.ErrorHistogram, s.UncorrectedErrorHistogram, s.RequestRate, percentiles, file)
}

func getOneByPercentile(percentile float64) float64 {
	if percentile == 100 {
		return 1 / (1 - (percentile / 100))
	}
	return float64(10000000)
}

func generateLatencyDistribution(histogram, unHistogram *hdrhistogram.Histogram, requestRate uint64, percentiles Percentiles, file string) error {
	if percentiles == nil {
		percentiles = Logarithmic
	}
	f, err := os.Create(file)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
	for _, percentile := range percentiles {
		value := float64(histogram.ValueAtQuantile(percentile)) / 1000000
		oneByPercentile := getOneByPercentile(percentile)
		_, err := f.WriteString(fmt.Sprintf("%f    %f        %d            %f\n",
			value, percentile/100, 0, oneByPercentile))
		if err != nil {
			return err
		}
	}

	// Generate uncorrected distribution.
	if requestRate > 0 {
		f, err := os.Create("uncorrected_" + file)
		if err != nil {
			return err
		}
		defer f.Close()

		f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
		for _, percentile := range percentiles {
			value := float64(unHistogram.ValueAtQuantile(percentile)) / 1000000
			oneByPercentile := getOneByPercentile(percentile)
			_, err := f.WriteString(fmt.Sprintf("%f    %f        %d            %f\n",
				value, percentile/100, 0, oneByPercentile))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// merge the other Summary into this one.
func (s *Summary) merge(o *Summary) {
	if o.TimeElapsed > s.TimeElapsed {
		s.TimeElapsed = o.TimeElapsed
	}
	s.SuccessHistogram.Merge(o.SuccessHistogram)
	s.UncorrectedSuccessHistogram.Merge(o.UncorrectedSuccessHistogram)
	s.ErrorHistogram.Merge(o.ErrorHistogram)
	s.UncorrectedErrorHistogram.Merge(o.UncorrectedErrorHistogram)
	s.SuccessTotal += o.SuccessTotal
	s.ErrorTotal += o.ErrorTotal
	s.Throughput += o.Throughput
	s.RequestRate += o.RequestRate
}
<|endoftext|>


<commit_before>package sup

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"gopkg.in/yaml.v2"
)

// EnvVar represents an environment variable
type EnvVar struct {
	Key   string
	Value string
}

func (e EnvVar) String() string {
	return e.Key + `=` + e.Value
}

// AsExport returns the environment variable as a bash export statement
func (e EnvVar) AsExport() string {
	return `export ` + e.Key + `="` + e.Value + `";`
}

// EnvList is a list of environment variables that maps to a YAML map,
// but maintains order, enabling late variables to reference early variables.
type EnvList []EnvVar

func (e *EnvList) UnmarshalYAML(unmarshal func(interface{}) error) error {
	items := []yaml.MapItem{}

	err := unmarshal(&items)
	if err != nil {
		return err
	}

	*e = make(EnvList, 0, len(items))

	for _, v := range items {
		e.Set(fmt.Sprintf("%v", v.Key), fmt.Sprintf("%v", v.Value))
	}

	return nil
}

// Set key to be equal value in this list.
func (e *EnvList) Set(key, value string) {

	for _, v := range *e {
		if v.Key == key {
			v.Value = value
			return
		}
	}

	*e = append(*e, EnvVar{
		Key:   key,
		Value: value,
	})
}

// Supfile represents the Stack Up configuration YAML file.
type Supfile struct {
	Networks map[string]Network  `yaml:"networks"`
	Commands map[string]Command  `yaml:"commands"`
	Targets  map[string][]string `yaml:"targets"`
	Env      EnvList             `yaml:"env"`
	Version  string              `yaml:"version"`
}

// Network is group of hosts with extra custom env vars.
type Network struct {
	Env       EnvList  `yaml:"env"`
	Inventory string   `yaml:"inventory"`
	Hosts     []string `yaml:"hosts"`
	Bastion   string   `yaml:"bastion"` // Jump host for the environment
}

// Command represents command(s) to be run remotely.
type Command struct {
	Name   string   `yaml:"-"`      // Command name.
	Desc   string   `yaml:"desc"`   // Command description.
	Local  string   `yaml:"local"`  // Command(s) to be run locally.
	Run    string   `yaml:"run"`    // Command(s) to be run remotelly.
	Script string   `yaml:"script"` // Load command(s) from script and run it remotelly.
	Upload []Upload `yaml:"upload"` // See Upload struct.
	Stdin  bool     `yaml:"stdin"`  // Attach localhost STDOUT to remote commands' STDIN?
	Once   bool     `yaml:"once"`   // The command should be run "once" (on one host only).
	Serial int      `yaml:"serial"` // Max number of clients processing a task in parallel.

	// API backward compatibility. Will be deprecated in v1.0.
	RunOnce bool `yaml:"run_once"` // The command should be run once only.
}

// Upload represents file copy operation from localhost Src path to Dst
// path of every host in a given Network.
type Upload struct {
	Src string `yaml:"src"`
	Dst string `yaml:"dst"`
	Exc string `yaml:"exclude"`
}

// NewSupfile parses configuration file and returns Supfile or error.
func NewSupfile(file string) (*Supfile, error) {
	var conf Supfile
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(data, &conf)
	if err != nil {
		return nil, err
	}

	// API backward compatibility. Will be deprecated in v1.0.
	switch conf.Version {
	case "":
		conf.Version = "0.1"
		fallthrough
	case "0.1":
		for _, cmd := range conf.Commands {
			if cmd.RunOnce {
				return nil, errors.New("command.run_once is not supported in Supfile v" + conf.Version)
			}
		}
		fallthrough
	case "0.2":
		for _, cmd := range conf.Commands {
			if cmd.Once {
				return nil, errors.New("command.once is not supported in Supfile v" + conf.Version)
			}
			if cmd.Local != "" {
				return nil, errors.New("command.local is not supported in Supfile v" + conf.Version)
			}
			if cmd.Serial != 0 {
				return nil, errors.New("command.serial is not supported in Supfile v" + conf.Version)
			}
		}
		for _, network := range conf.Networks {
			if network.Inventory != "" {
				return nil, errors.New("network.inventory is not supported in Supfile v" + conf.Version)
			}
		}
		fallthrough
	case "0.3":
		var warning string
		for key, cmd := range conf.Commands {
			if cmd.RunOnce {
				warning = "Warning: command.run_once was deprecated by command.once in Supfile v" + conf.Version + "\n"
				cmd.Once = true
				conf.Commands[key] = cmd
			}
		}
		if warning != "" {
			fmt.Fprintf(os.Stderr, warning)
		}
	default:
		return nil, errors.New("unsupported version, please update sup by `go get -u github.com/pressly/sup`")
	}

	for i, network := range conf.Networks {
		hosts, err := network.ParseInventory()
		if err != nil {
			return nil, err
		}
		network.Hosts = append(network.Hosts, hosts...)
		conf.Networks[i] = network
	}

	return &conf, nil
}

// ParseInventory runs the inventory command, if provided, and appends
// the command's output lines to the manually defined list of hosts.
func (n Network) ParseInventory() ([]string, error) {
	if n.Inventory == "" {
		return nil, nil
	}

	cmd := exec.Command("/bin/sh", "-c", n.Inventory)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var hosts []string
	buf := bytes.NewBuffer(output)
	for {
		host, err := buf.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}

		host = strings.TrimSpace(host)
		// skip empty lines and comments
		if host == "" || host[:1] == "#" {
			continue
		}

		hosts = append(hosts, host)
	}
	return hosts, nil
}
<commit_msg>re-order file putting the most important concepts on top<commit_after>package sup

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"gopkg.in/yaml.v2"
)

// Supfile represents the Stack Up configuration YAML file.
type Supfile struct {
	Networks map[string]Network  `yaml:"networks"`
	Commands map[string]Command  `yaml:"commands"`
	Targets  map[string][]string `yaml:"targets"`
	Env      EnvList             `yaml:"env"`
	Version  string              `yaml:"version"`
}

// Network is group of hosts with extra custom env vars.
type Network struct {
	Env       EnvList  `yaml:"env"`
	Inventory string   `yaml:"inventory"`
	Hosts     []string `yaml:"hosts"`
	Bastion   string   `yaml:"bastion"` // Jump host for the environment
}

// Command represents command(s) to be run remotely.
type Command struct {
	Name   string   `yaml:"-"`      // Command name.
	Desc   string   `yaml:"desc"`   // Command description.
	Local  string   `yaml:"local"`  // Command(s) to be run locally.
	Run    string   `yaml:"run"`    // Command(s) to be run remotelly.
	Script string   `yaml:"script"` // Load command(s) from script and run it remotelly.
	Upload []Upload `yaml:"upload"` // See Upload struct.
	Stdin  bool     `yaml:"stdin"`  // Attach localhost STDOUT to remote commands' STDIN?
	Once   bool     `yaml:"once"`   // The command should be run "once" (on one host only).
	Serial int      `yaml:"serial"` // Max number of clients processing a task in parallel.

	// API backward compatibility. Will be deprecated in v1.0.
	RunOnce bool `yaml:"run_once"` // The command should be run once only.
}

// Upload represents file copy operation from localhost Src path to Dst
// path of every host in a given Network.
type Upload struct {
	Src string `yaml:"src"`
	Dst string `yaml:"dst"`
	Exc string `yaml:"exclude"`
}

// EnvVar represents an environment variable
type EnvVar struct {
	Key   string
	Value string
}

func (e EnvVar) String() string {
	return e.Key + `=` + e.Value
}

// AsExport returns the environment variable as a bash export statement
func (e EnvVar) AsExport() string {
	return `export ` + e.Key + `="` + e.Value + `";`
}

// EnvList is a list of environment variables that maps to a YAML map,
// but maintains order, enabling late variables to reference early variables.
type EnvList []EnvVar

func (e *EnvList) UnmarshalYAML(unmarshal func(interface{}) error) error {
	items := []yaml.MapItem{}

	err := unmarshal(&items)
	if err != nil {
		return err
	}

	*e = make(EnvList, 0, len(items))

	for _, v := range items {
		e.Set(fmt.Sprintf("%v", v.Key), fmt.Sprintf("%v", v.Value))
	}

	return nil
}

// Set key to be equal value in this list.
func (e *EnvList) Set(key, value string) {

	for _, v := range *e {
		if v.Key == key {
			v.Value = value
			return
		}
	}

	*e = append(*e, EnvVar{
		Key:   key,
		Value: value,
	})
}

// NewSupfile parses configuration file and returns Supfile or error.
func NewSupfile(file string) (*Supfile, error) {
	var conf Supfile
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(data, &conf)
	if err != nil {
		return nil, err
	}

	// API backward compatibility. Will be deprecated in v1.0.
	switch conf.Version {
	case "":
		conf.Version = "0.1"
		fallthrough
	case "0.1":
		for _, cmd := range conf.Commands {
			if cmd.RunOnce {
				return nil, errors.New("command.run_once is not supported in Supfile v" + conf.Version)
			}
		}
		fallthrough
	case "0.2":
		for _, cmd := range conf.Commands {
			if cmd.Once {
				return nil, errors.New("command.once is not supported in Supfile v" + conf.Version)
			}
			if cmd.Local != "" {
				return nil, errors.New("command.local is not supported in Supfile v" + conf.Version)
			}
			if cmd.Serial != 0 {
				return nil, errors.New("command.serial is not supported in Supfile v" + conf.Version)
			}
		}
		for _, network := range conf.Networks {
			if network.Inventory != "" {
				return nil, errors.New("network.inventory is not supported in Supfile v" + conf.Version)
			}
		}
		fallthrough
	case "0.3":
		var warning string
		for key, cmd := range conf.Commands {
			if cmd.RunOnce {
				warning = "Warning: command.run_once was deprecated by command.once in Supfile v" + conf.Version + "\n"
				cmd.Once = true
				conf.Commands[key] = cmd
			}
		}
		if warning != "" {
			fmt.Fprintf(os.Stderr, warning)
		}
	default:
		return nil, errors.New("unsupported version, please update sup by `go get -u github.com/pressly/sup`")
	}

	for i, network := range conf.Networks {
		hosts, err := network.ParseInventory()
		if err != nil {
			return nil, err
		}
		network.Hosts = append(network.Hosts, hosts...)
		conf.Networks[i] = network
	}

	return &conf, nil
}

// ParseInventory runs the inventory command, if provided, and appends
// the command's output lines to the manually defined list of hosts.
func (n Network) ParseInventory() ([]string, error) {
	if n.Inventory == "" {
		return nil, nil
	}

	cmd := exec.Command("/bin/sh", "-c", n.Inventory)
	cmd.Stderr = os.Stderr
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var hosts []string
	buf := bytes.NewBuffer(output)
	for {
		host, err := buf.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}

		host = strings.TrimSpace(host)
		// skip empty lines and comments
		if host == "" || host[:1] == "#" {
			continue
		}

		hosts = append(hosts, host)
	}
	return hosts, nil
}
<|endoftext|>


<commit_before>package circonus

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestNewCirconusSink(t *testing.T) {

	// test with invalid config (nil)
	expectedError := errors.New("Invalid check manager configuration (no API token AND no submission url).")
	_, err := NewCirconusSink(nil)
	if err == nil || err.Error() != expectedError.Error() {
		t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
	}

	// test w/submission url and w/o token
	cfg := &Config{}
	cfg.CheckSubmissionURL = "http://127.0.0.1:43191/"
	_, err = NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	// note: a test with a valid token is *not* done as it *will* create a
	//       check resulting in testing the api more than the circonus sink
	// see circonus-gometrics/checkmgr/checkmgr_test.go for testing of api token
}

func TestFlattenKey(t *testing.T) {
	var testKeys = []struct {
		input    []string
		expected string
	}{
		{[]string{"a", "b", "c"}, "a`b`c"},
		{[]string{"a-a", "b_b", "c/c"}, "a-a`b_b`c/c"},
		{[]string{"spaces must", "flatten", "to", "underscores"}, "spaces_must`flatten`to`underscores"},
	}

	c := &CirconusSink{}

	for _, test := range testKeys {
		if actual := c.flattenKey(test.input); actual != test.expected {
			t.Fatalf("Flattening %v failed, expected '%s' got '%s'", test.input, test.expected, actual)
		}
	}
}

func fakeBroker(q chan string) *httptest.Server {
	handler := func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		defer r.Body.Close()
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			q <- err.Error()
			fmt.Fprintln(w, err.Error())
		} else {
			q <- string(body)
			fmt.Fprintln(w, `{"stats":1}`)
		}
	}

	return httptest.NewServer(http.HandlerFunc(handler))
}

func TestSetGauge(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.CheckSubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.SetGauge([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":1}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}

func TestIncrCounter(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.CheckSubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.IncrCounter([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":1}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}

func TestAddSample(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.CheckSubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.AddSample([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":[\"H[1.0e+00]=1\"]}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}
<commit_msg>fix typo Check.SubmissionURL not CheckSubmissionURL<commit_after>package circonus

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestNewCirconusSink(t *testing.T) {

	// test with invalid config (nil)
	expectedError := errors.New("Invalid check manager configuration (no API token AND no submission url).")
	_, err := NewCirconusSink(nil)
	if err == nil || err.Error() != expectedError.Error() {
		t.Errorf("Expected an '%#v' error, got '%#v'", expectedError, err)
	}

	// test w/submission url and w/o token
	cfg := &Config{}
	cfg.Check.SubmissionURL = "http://127.0.0.1:43191/"
	_, err = NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	// note: a test with a valid token is *not* done as it *will* create a
	//       check resulting in testing the api more than the circonus sink
	// see circonus-gometrics/checkmgr/checkmgr_test.go for testing of api token
}

func TestFlattenKey(t *testing.T) {
	var testKeys = []struct {
		input    []string
		expected string
	}{
		{[]string{"a", "b", "c"}, "a`b`c"},
		{[]string{"a-a", "b_b", "c/c"}, "a-a`b_b`c/c"},
		{[]string{"spaces must", "flatten", "to", "underscores"}, "spaces_must`flatten`to`underscores"},
	}

	c := &CirconusSink{}

	for _, test := range testKeys {
		if actual := c.flattenKey(test.input); actual != test.expected {
			t.Fatalf("Flattening %v failed, expected '%s' got '%s'", test.input, test.expected, actual)
		}
	}
}

func fakeBroker(q chan string) *httptest.Server {
	handler := func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Header().Set("Content-Type", "application/json")
		defer r.Body.Close()
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			q <- err.Error()
			fmt.Fprintln(w, err.Error())
		} else {
			q <- string(body)
			fmt.Fprintln(w, `{"stats":1}`)
		}
	}

	return httptest.NewServer(http.HandlerFunc(handler))
}

func TestSetGauge(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.Check.SubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.SetGauge([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":1}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}

func TestIncrCounter(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.Check.SubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.IncrCounter([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":1}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}

func TestAddSample(t *testing.T) {
	q := make(chan string)

	server := fakeBroker(q)
	defer server.Close()

	cfg := &Config{}
	cfg.Check.SubmissionURL = server.URL

	cs, err := NewCirconusSink(cfg)
	if err != nil {
		t.Errorf("Expected no error, got '%v'", err)
	}

	go func() {
		cs.AddSample([]string{"foo", "bar"}, 1)
		cs.Flush()
	}()

	expect := "{\"foo`bar\":{\"_type\":\"n\",\"_value\":[\"H[1.0e+00]=1\"]}}"
	actual := <-q

	if actual != expect {
		t.Errorf("Expected '%s', got '%s'", expect, actual)

	}
}
<|endoftext|>


<commit_before>package revelswagger

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"path"
	"runtime"
	"strings"
	"text/template"

	"github.com/howeyc/fsnotify"
	"github.com/revel/revel"
)

var spec Specification

func init() {
	revel.OnAppStart(func() {
		fmt.Println("[SWAGGER]: Loading schema...")

		loadSpecFile()

		go watchSpecFile()
	})
}

func loadSpecFile() {
	spec = Specification{}

	content, err := ioutil.ReadFile(revel.BasePath + "\\conf\\spec.json")

	if err != nil {
		fmt.Println("[SWAGGER]: Couldn't load spec.json.", err)
		return
	}

	err = json.Unmarshal(content, &spec)

	if err != nil {
		fmt.Println("[SWAGGER]: Error parsing schema file.", err)
	}
}

func watchSpecFile() {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	done := make(chan bool)

	// Process events
	go func() {
		for {
			select {
			case <-watcher.Event:
				loadSpecFile()
			case err := <-watcher.Error:
				fmt.Println("[SWAGGER]: Watcher error:", err)
			}
		}
	}()

	err = watcher.Watch(revel.BasePath + "\\conf\\spec.json")

	if err != nil {
		fmt.Println("[SWAGGER]: Error watching spec file:", err)
	} else {
		fmt.Println("[SWAGGER]: Spec watcher initialized")
	}

	<-done

	/* ... do stuff ... */
	watcher.Close()
}

func Filter(c *revel.Controller, fc []revel.Filter) {
	c.Request.URL.Path = strings.ToLower(c.Request.URL.Path)

	var route *revel.RouteMatch = revel.MainRouter.Route(c.Request.Request)

	if route == nil {
		c.Result = c.NotFound("No matching route found: " + c.Request.RequestURI)
		return
	}

	if len(route.Params) == 0 {
		c.Params.Route = map[string][]string{}
	} else {
		c.Params.Route = route.Params
	}

	if err := c.SetAction(route.ControllerName, route.MethodName); err != nil {
		c.Result = c.NotFound(err.Error())
		return
	}

	// Add the fixed parameters mapped by name.
	// TODO: Pre-calculate this mapping.
	for i, value := range route.FixedParams {
		if c.Params.Fixed == nil {
			c.Params.Fixed = make(url.Values)
		}
		if i < len(c.MethodType.Args) {
			arg := c.MethodType.Args[i]
			c.Params.Fixed.Set(arg.Name, value)
		} else {
			fmt.Println("Too many parameters to", route.Action, "trying to add", value)
			break
		}
	}

	leaf, _ := revel.MainRouter.Tree.Find(treePath(c.Request.Method, c.Request.URL.Path))

	r := leaf.Value.(*revel.Route)

	method := spec.Paths[r.Path].Get

	if method == nil {
		// Check if strict mode is enabled and throw an error, otherwise
		// just move onto the next filter like revel normally would
		if revel.Config.BoolDefault("swagger.strict", true) {
			_, filename, _, _ := runtime.Caller(0)

			t, err := template.ParseFiles(path.Dir(filename) + "/views/notfound.html")

			if err != nil {
				panic(err)
			}

			t.Execute(c.Response.Out, map[string]interface{}{
				"routes": spec.Paths,
				"path":   c.Request.RequestURI,
			})
			return
		} else {
			// Move onto the next filter
			fc[0](c, fc[1:])
			return
		}
	}

	// Action has been found & set, let's validate the parameters
	validateParameters(method.Parameters, c)

	if c.Validation.HasErrors() {
		var errors []string

		for _, e := range c.Validation.Errors {
			errors = append(errors, e.Message)
		}

		c.Result = c.RenderJson(map[string]interface{}{"errors": errors})
		return
	}

	// Move onto the next filter
	fc[0](c, fc[1:])
}

func treePath(method, path string) string {
	if method == "*" {
		method = ":METHOD"
	}
	return "/" + method + path
}
<commit_msg>Fix spec.json warning/error. Proper OS path handling.<commit_after>package revelswagger

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"path"
	"runtime"
	"strings"
	"text/template"

	"github.com/howeyc/fsnotify"
	"github.com/revel/revel"
)

var spec Specification

func init() {
	revel.OnAppStart(func() {
		fmt.Println("[SWAGGER]: Loading schema...")

		loadSpecFile()

		go watchSpecFile()
	})
}

func loadSpecFile() {
	spec = Specification{}
    specPath := "/conf/spec.json"

	if runtime.GOOS == "windows" {
        specPath = "\\conf\\spec.json"
    }

	content, err := ioutil.ReadFile(revel.BasePath + specPath)

	if err != nil {
		fmt.Println("[SWAGGER]: Couldn't load spec.json.", err)
		return
	}

	err = json.Unmarshal(content, &spec)

	if err != nil {
		fmt.Println("[SWAGGER]: Error parsing schema file.", err)
	}
}

func watchSpecFile() {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}

	done := make(chan bool)

	// Process events
	go func() {
		for {
			select {
			case <-watcher.Event:
				loadSpecFile()
			case err := <-watcher.Error:
				fmt.Println("[SWAGGER]: Watcher error:", err)
			}
		}
	}()

    specPath := "/conf/spec.json"
    if runtime.GOOS == "windows" {
        specPath = "\\conf\\spec.json"
    }

	err = watcher.Watch(revel.BasePath + specPath)

	if err != nil {
		fmt.Println("[SWAGGER]: Error watching spec file:", err)
	} else {
		fmt.Println("[SWAGGER]: Spec watcher initialized")
	}

	<-done

	/* ... do stuff ... */
	watcher.Close()
}

func Filter(c *revel.Controller, fc []revel.Filter) {
	c.Request.URL.Path = strings.ToLower(c.Request.URL.Path)

	var route *revel.RouteMatch = revel.MainRouter.Route(c.Request.Request)

	if route == nil {
		c.Result = c.NotFound("No matching route found: " + c.Request.RequestURI)
		return
	}

	if len(route.Params) == 0 {
		c.Params.Route = map[string][]string{}
	} else {
		c.Params.Route = route.Params
	}

	if err := c.SetAction(route.ControllerName, route.MethodName); err != nil {
		c.Result = c.NotFound(err.Error())
		return
	}

	// Add the fixed parameters mapped by name.
	// TODO: Pre-calculate this mapping.
	for i, value := range route.FixedParams {
		if c.Params.Fixed == nil {
			c.Params.Fixed = make(url.Values)
		}
		if i < len(c.MethodType.Args) {
			arg := c.MethodType.Args[i]
			c.Params.Fixed.Set(arg.Name, value)
		} else {
			fmt.Println("Too many parameters to", route.Action, "trying to add", value)
			break
		}
	}

	leaf, _ := revel.MainRouter.Tree.Find(treePath(c.Request.Method, c.Request.URL.Path))

	r := leaf.Value.(*revel.Route)

	method := spec.Paths[r.Path].Get

	if method == nil {
		// Check if strict mode is enabled and throw an error, otherwise
		// just move onto the next filter like revel normally would
		if revel.Config.BoolDefault("swagger.strict", true) {
			_, filename, _, _ := runtime.Caller(0)

			t, err := template.ParseFiles(path.Dir(filename) + "/views/notfound.html")

			if err != nil {
				panic(err)
			}

			t.Execute(c.Response.Out, map[string]interface{}{
				"routes": spec.Paths,
				"path":   c.Request.RequestURI,
			})
			return
		} else {
			// Move onto the next filter
			fc[0](c, fc[1:])
			return
		}
	}

	// Action has been found & set, let's validate the parameters
	validateParameters(method.Parameters, c)

	if c.Validation.HasErrors() {
		var errors []string

		for _, e := range c.Validation.Errors {
			errors = append(errors, e.Message)
		}

		c.Result = c.RenderJson(map[string]interface{}{"errors": errors})
		return
	}

	// Move onto the next filter
	fc[0](c, fc[1:])
}

func treePath(method, path string) string {
	if method == "*" {
		method = ":METHOD"
	}
	return "/" + method + path
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"

	"github.com/brotherlogic/goserver"
	"github.com/brotherlogic/keystore/client"
	"google.golang.org/grpc"

	pb "github.com/brotherlogic/discogssyncer/server"
	pbdi "github.com/brotherlogic/discovery/proto"
	pbd "github.com/brotherlogic/godiscogs"
)

// Syncer the configuration for the syncer
type Syncer struct {
	*goserver.GoServer
	token      string
	retr       saver
	collection *pb.RecordCollection
}

var (
	syncTime int64
)

const (
	//KEY under which we store the collection
	KEY = "/github.com/brotherlogic/discogssyncer/collection"

	//TOKEN for discogs
	TOKEN = "/github.com/brotherlogic/discogssyncer/token"
)

// This is the only method that interacts with disk
func (s *Syncer) readRecordCollection() error {
	log.Printf("Reading collection")
	collection := &pb.RecordCollection{}
	data, err := s.KSclient.Read(KEY, collection)

	log.Printf("READ: %v", data)

	if err != nil {
		log.Printf("Unable to read collection: %v", err)
		return err
	}

	s.collection = data.(*pb.RecordCollection)
	log.Printf("FOLDERS = %v", len(s.collection.Folders))
	return nil
}

func (s *Syncer) saveCollection() {
	log.Printf("Writing collection")
	s.KSclient.Save(KEY, s.collection)
}

func (s *Syncer) deleteRelease(rel *pbd.Release, folder int32) {
	index := -1
	for _, f := range s.collection.Folders {
		if f.Folder.Id == folder {
			for i, r := range f.Releases.Releases {
				if r.Id == rel.Id {
					index = i
				}

			}
			if index >= 0 {
				f.Releases.Releases = append(f.Releases.Releases[:index], f.Releases.Releases[index+1:]...)
			}
		}
	}

}

// DoRegister does RPC registration
func (s Syncer) DoRegister(server *grpc.Server) {
	pb.RegisterDiscogsServiceServer(server, &s)
}

func findServer(name string) (string, int) {
	conn, err := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	if err != nil {
		log.Fatalf("Cannot reach discover server: %v (trying to discover %v)", err, name)
	}
	defer conn.Close()

	registry := pbdi.NewDiscoveryServiceClient(conn)
	rs, err := registry.ListAllServices(context.Background(), &pbdi.Empty{})

	if err != nil {
		log.Fatalf("Failure to list: %v", err)
	}

	for _, r := range rs.Services {
		if r.Name == name {
			log.Printf("%v -> %v", name, r)
			return r.Ip, int(r.Port)
		}
	}

	log.Printf("No %v running", name)

	return "", -1
}

// InitServer builds an initial server
func InitServer() Syncer {
	syncer := Syncer{GoServer: &goserver.GoServer{}, collection: &pb.RecordCollection{Wantlist: &pb.Wantlist{}}}
	syncer.GoServer.KSclient = *keystoreclient.GetClient(findServer)
	err := syncer.readRecordCollection()
	if err != nil {
		log.Fatalf("Unable to read record collection")
	}

	return syncer
}

// Mote promotes/demotes this server
func (s Syncer) Mote(master bool) error {
	return nil
}

// ReportHealth alerts if we're not healthy
func (s Syncer) ReportHealth() bool {
	return true
}

func main() {
	var quiet = flag.Bool("quiet", true, "Show all output")
	var token = flag.String("token", "", "Discogs token")
	flag.Parse()

	syncer := InitServer()

	if len(*token) > 0 {
		syncer.KSclient.Save(TOKEN, &pb.Token{Token: *token})
	}

	tType := &pb.Token{}
	tResp, err := syncer.KSclient.Read(TOKEN, tType)

	if err != nil {
		log.Fatalf("Unable to read token: %v", err)
	}

	sToken := tResp.(*pb.Token).Token
	syncer.retr = pbd.NewDiscogsRetriever(sToken)
	syncer.token = sToken

	//Turn off logging
	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	syncer.Register = syncer
	syncer.PrepServer()
	syncer.RegisterServer("discogssyncer", false)

	log.Printf("PRESERVER %v", len(syncer.collection.Folders))

	syncer.Serve()
}
<commit_msg>Made this more quiet<commit_after>package main

import (
	"context"
	"flag"
	"io/ioutil"
	"log"

	"github.com/brotherlogic/goserver"
	"github.com/brotherlogic/keystore/client"
	"google.golang.org/grpc"

	pb "github.com/brotherlogic/discogssyncer/server"
	pbdi "github.com/brotherlogic/discovery/proto"
	pbd "github.com/brotherlogic/godiscogs"
)

// Syncer the configuration for the syncer
type Syncer struct {
	*goserver.GoServer
	token      string
	retr       saver
	collection *pb.RecordCollection
}

var (
	syncTime int64
)

const (
	//KEY under which we store the collection
	KEY = "/github.com/brotherlogic/discogssyncer/collection"

	//TOKEN for discogs
	TOKEN = "/github.com/brotherlogic/discogssyncer/token"
)

// This is the only method that interacts with disk
func (s *Syncer) readRecordCollection() error {
	log.Printf("Reading collection")
	collection := &pb.RecordCollection{}
	data, err := s.KSclient.Read(KEY, collection)

	log.Printf("READ: %v", data)

	if err != nil {
		log.Printf("Unable to read collection: %v", err)
		return err
	}

	s.collection = data.(*pb.RecordCollection)
	log.Printf("FOLDERS = %v", len(s.collection.Folders))
	return nil
}

func (s *Syncer) saveCollection() {
	log.Printf("Writing collection")
	s.KSclient.Save(KEY, s.collection)
}

func (s *Syncer) deleteRelease(rel *pbd.Release, folder int32) {
	index := -1
	for _, f := range s.collection.Folders {
		if f.Folder.Id == folder {
			for i, r := range f.Releases.Releases {
				if r.Id == rel.Id {
					index = i
				}

			}
			if index >= 0 {
				f.Releases.Releases = append(f.Releases.Releases[:index], f.Releases.Releases[index+1:]...)
			}
		}
	}

}

// DoRegister does RPC registration
func (s Syncer) DoRegister(server *grpc.Server) {
	pb.RegisterDiscogsServiceServer(server, &s)
}

func findServer(name string) (string, int) {
	conn, err := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	if err != nil {
		log.Fatalf("Cannot reach discover server: %v (trying to discover %v)", err, name)
	}
	defer conn.Close()

	registry := pbdi.NewDiscoveryServiceClient(conn)
	rs, err := registry.ListAllServices(context.Background(), &pbdi.Empty{})

	if err != nil {
		log.Fatalf("Failure to list: %v", err)
	}

	for _, r := range rs.Services {
		if r.Name == name {
			log.Printf("%v -> %v", name, r)
			return r.Ip, int(r.Port)
		}
	}

	log.Printf("No %v running", name)

	return "", -1
}

// InitServer builds an initial server
func InitServer() Syncer {
	syncer := Syncer{GoServer: &goserver.GoServer{}, collection: &pb.RecordCollection{Wantlist: &pb.Wantlist{}}}
	syncer.GoServer.KSclient = *keystoreclient.GetClient(findServer)
	err := syncer.readRecordCollection()
	if err != nil {
		log.Fatalf("Unable to read record collection")
	}

	return syncer
}

// Mote promotes/demotes this server
func (s Syncer) Mote(master bool) error {
	return nil
}

// ReportHealth alerts if we're not healthy
func (s Syncer) ReportHealth() bool {
	return true
}

func main() {
	var quiet = flag.Bool("quiet", true, "Show all output")
	var token = flag.String("token", "", "Discogs token")
	flag.Parse()

	//Turn off logging
	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	syncer := InitServer()

	if len(*token) > 0 {
		syncer.KSclient.Save(TOKEN, &pb.Token{Token: *token})
	}

	tType := &pb.Token{}
	tResp, err := syncer.KSclient.Read(TOKEN, tType)

	if err != nil {
		log.Fatalf("Unable to read token: %v", err)
	}

	sToken := tResp.(*pb.Token).Token
	syncer.retr = pbd.NewDiscogsRetriever(sToken)
	syncer.token = sToken

	syncer.Register = syncer
	syncer.PrepServer()
	syncer.RegisterServer("discogssyncer", false)

	log.Printf("PRESERVER %v", len(syncer.collection.Folders))

	syncer.Serve()
}
<|endoftext|>


<commit_before>package gorillamux

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"github.com/ungerik/go-command"
	"github.com/ungerik/go-httpx/httperr"
)

func CommandHandler(commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

func CommandHandlerWithQueryParams(commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)

		// Add query params as arguments by joining them together per key (query
		// param names are not unique).
		for k := range request.URL.Query() {
			if len(request.URL.Query()[k]) > 0 && len(request.URL.Query()[k][0]) > 0 {
				vars[k] = strings.Join(request.URL.Query()[k][:], ";")
			}
		}

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

type RequestBodyArgConverter interface {
	RequestBodyToArg(request *http.Request) (name, value string, err error)
}

type RequestBodyArgConverterFunc func(request *http.Request) (name, value string, err error)

func (f RequestBodyArgConverterFunc) RequestBodyToArg(request *http.Request) (name, value string, err error) {
	return f(request)
}

func RequestBodyAsArg(name string) RequestBodyArgConverterFunc {
	return func(request *http.Request) (string, string, error) {
		defer request.Body.Close()
		b, err := ioutil.ReadAll(request.Body)
		if err != nil {
			return "", "", err
		}
		return name, string(b), nil
	}
}

func CommandHandlerRequestBodyArg(bodyConverter RequestBodyArgConverter, commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)
		name, value, err := bodyConverter.RequestBodyToArg(request)
		if err != nil {
			handleErr(err, writer, request, errHandlers)
			return
		}
		if _, exists := vars[name]; exists {
			err = fmt.Errorf("argument '%s' already set by request URL path", name)
			handleErr(err, writer, request, errHandlers)
			return
		}
		vars[name] = value

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

func handleErr(err error, writer http.ResponseWriter, request *http.Request, errHandlers []httperr.Handler) {
	if err == nil {
		return
	}
	if len(errHandlers) == 0 {
		httperr.Handle(err, writer, request)
	} else {
		for _, errHandler := range errHandlers {
			errHandler.HandleError(err, writer, request)
		}
	}
}

func MapJSONBodyFieldsAsVars(mapping map[string]string, wrappedHandler http.Handler) http.HandlerFunc {
	return func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		body, err := ioutil.ReadAll(request.Body)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		vars := mux.Vars(request)
		err = jsonBodyFieldsAsVars(body, mapping, vars)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		wrappedHandler.ServeHTTP(writer, request)
	}
}

func JSONBodyFieldsAsVars(wrappedHandler http.Handler) http.HandlerFunc {
	return func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		body, err := ioutil.ReadAll(request.Body)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		vars := mux.Vars(request)
		err = jsonBodyFieldsAsVars(body, nil, vars)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		wrappedHandler.ServeHTTP(writer, request)
	}
}

func jsonBodyFieldsAsVars(body []byte, mapping map[string]string, vars map[string]string) error {
	fields := make(map[string]json.RawMessage)
	err := json.Unmarshal(body, &fields)
	if err != nil {
		return err
	}

	if mapping != nil {
		mappedFields := make(map[string]json.RawMessage, len(fields))
		for fieldName, mappedName := range mapping {
			if value, ok := fields[fieldName]; ok {
				mappedFields[mappedName] = value
			}
		}
		fields = mappedFields
	}

	for name, value := range fields {
		if len(value) == 0 {
			// should never happen with well formed JSON
			return fmt.Errorf("JSON body field %q is empty", name)
		}
		valueStr := string(value)
		switch {
		case valueStr == "null":
			// JSON null is handled as empty string command arg
			vars[name] = ""

		case valueStr[0] == '"':
			// Unescape JSON string
			err = json.Unmarshal(value, &valueStr)
			if err != nil {
				return fmt.Errorf("can't unmarshal JSON body field %q as string because of: %w", name, err)
			}
			vars[name] = valueStr

		default:
			// All other JSON types are mapped directly to string
			vars[name] = valueStr
		}
	}
	return nil
}
<commit_msg>Revert "Revert "fix: leave JSON nulls alone""<commit_after>package gorillamux

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"github.com/ungerik/go-command"
	"github.com/ungerik/go-httpx/httperr"
)

func CommandHandler(commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

func CommandHandlerWithQueryParams(commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)

		// Add query params as arguments by joining them together per key (query
		// param names are not unique).
		for k := range request.URL.Query() {
			if len(request.URL.Query()[k]) > 0 && len(request.URL.Query()[k][0]) > 0 {
				vars[k] = strings.Join(request.URL.Query()[k][:], ";")
			}
		}

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

type RequestBodyArgConverter interface {
	RequestBodyToArg(request *http.Request) (name, value string, err error)
}

type RequestBodyArgConverterFunc func(request *http.Request) (name, value string, err error)

func (f RequestBodyArgConverterFunc) RequestBodyToArg(request *http.Request) (name, value string, err error) {
	return f(request)
}

func RequestBodyAsArg(name string) RequestBodyArgConverterFunc {
	return func(request *http.Request) (string, string, error) {
		defer request.Body.Close()
		b, err := ioutil.ReadAll(request.Body)
		if err != nil {
			return "", "", err
		}
		return name, string(b), nil
	}
}

func CommandHandlerRequestBodyArg(bodyConverter RequestBodyArgConverter, commandFunc interface{}, args command.Args, resultsWriter ResultsWriter, errHandlers ...httperr.Handler) http.HandlerFunc {
	cmdFunc := command.MustGetStringMapArgsResultValuesFunc(commandFunc, args)

	return func(writer http.ResponseWriter, request *http.Request) {
		if CatchPanics {
			defer func() {
				handleErr(httperr.AsError(recover()), writer, request, errHandlers)
			}()
		}

		vars := mux.Vars(request)
		name, value, err := bodyConverter.RequestBodyToArg(request)
		if err != nil {
			handleErr(err, writer, request, errHandlers)
			return
		}
		if _, exists := vars[name]; exists {
			err = fmt.Errorf("argument '%s' already set by request URL path", name)
			handleErr(err, writer, request, errHandlers)
			return
		}
		vars[name] = value

		resultVals, err := cmdFunc(request.Context(), vars)

		if resultsWriter != nil {
			err = resultsWriter.WriteResults(args, vars, resultVals, err, writer, request)
		}
		handleErr(err, writer, request, errHandlers)
	}
}

func handleErr(err error, writer http.ResponseWriter, request *http.Request, errHandlers []httperr.Handler) {
	if err == nil {
		return
	}
	if len(errHandlers) == 0 {
		httperr.Handle(err, writer, request)
	} else {
		for _, errHandler := range errHandlers {
			errHandler.HandleError(err, writer, request)
		}
	}
}

func MapJSONBodyFieldsAsVars(mapping map[string]string, wrappedHandler http.Handler) http.HandlerFunc {
	return func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		body, err := ioutil.ReadAll(request.Body)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		vars := mux.Vars(request)
		err = jsonBodyFieldsAsVars(body, mapping, vars)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		wrappedHandler.ServeHTTP(writer, request)
	}
}

func JSONBodyFieldsAsVars(wrappedHandler http.Handler) http.HandlerFunc {
	return func(writer http.ResponseWriter, request *http.Request) {
		defer request.Body.Close()
		body, err := ioutil.ReadAll(request.Body)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		vars := mux.Vars(request)
		err = jsonBodyFieldsAsVars(body, nil, vars)
		if err != nil {
			httperr.BadRequest.ServeHTTP(writer, request)
			return
		}
		wrappedHandler.ServeHTTP(writer, request)
	}
}

func jsonBodyFieldsAsVars(body []byte, mapping map[string]string, vars map[string]string) error {
	fields := make(map[string]json.RawMessage)
	err := json.Unmarshal(body, &fields)
	if err != nil {
		return err
	}

	if mapping != nil {
		mappedFields := make(map[string]json.RawMessage, len(fields))
		for fieldName, mappedName := range mapping {
			if value, ok := fields[fieldName]; ok {
				mappedFields[mappedName] = value
			}
		}
		fields = mappedFields
	}

	for name, value := range fields {
		if len(value) == 0 {
			// should never happen with well formed JSON
			return fmt.Errorf("JSON body field %q is empty", name)
		}
		valueStr := string(value)
		switch {
		case valueStr == "null":
			// JSON nulls are left alone

		case valueStr[0] == '"':
			// Unescape JSON string
			err = json.Unmarshal(value, &valueStr)
			if err != nil {
				return fmt.Errorf("can't unmarshal JSON body field %q as string because of: %w", name, err)
			}
			vars[name] = valueStr

		default:
			// All other JSON types are mapped directly to string
			vars[name] = valueStr
		}
	}
	return nil
}
<|endoftext|>