<commit_before>/*
 * Copyright 2016 Xiaomi Corporation. All rights reserved.
 * Use of this source code is governed by a BSD-style
 * license that can be found in the LICENSE file.
 *
 * Authors:    Yu Bo <yubo@xiaomi.com>
 */
package govs

import (
	"fmt"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"strconv"
	"strings"
)

const (
	URL = "/tmp/dpvs.sock"

	/*
	 *      IPVS Connection Flags
	 */
	VS_CONN_F_FWD_MASK   = 0x000f /* mask for the fwd methods */
	VS_CONN_F_MASQ       = 0x0000 /* masquerading/NAT */
	VS_CONN_F_LOCALNODE  = 0x0001 /* local node */
	VS_CONN_F_TUNNEL     = 0x0002 /* tunneling */
	VS_CONN_F_DROUTE     = 0x0003 /* direct routing */
	VS_CONN_F_BYPASS     = 0x0004 /* cache bypass */
	VS_CONN_F_FULLNAT    = 0x0005 /* full nat */
	VS_CONN_F_DSNAT      = 0x0008 /* dsnat flag */
	VS_CONN_F_SYNC       = 0x0020 /* entry created by sync */
	VS_CONN_F_HASHED     = 0x0040 /* hashed entry */
	VS_CONN_F_NOOUTPUT   = 0x0080 /* no output packets */
	VS_CONN_F_INACTIVE   = 0x0100 /* not established */
	VS_CONN_F_OUT_SEQ    = 0x0200 /* must do output seq adjust */
	VS_CONN_F_IN_SEQ     = 0x0400 /* must do input seq adjust */
	VS_CONN_F_SEQ_MASK   = 0x0600 /* in/out sequence mask */
	VS_CONN_F_NO_CPORT   = 0x0800 /* no client port set yet */
	VS_CONN_F_TEMPLATE   = 0x1000 /* template, not connection */
	VS_CONN_F_ONE_PACKET = 0x2000 /* forward only one packet */
	VS_CONN_F_SYNPROXY   = 0x8000 /* syn proxy flag */

	VS_SVC_F_PERSISTENT = 0x0001             /* persistent port */
	VS_SVC_F_HASHED     = 0x0002             /* hashed entry */
	VS_SVC_F_ONEPACKET  = 0x0004             /* one-packet scheduling */
	VS_SVC_F_SYNPROXY   = VS_CONN_F_SYNPROXY /* synproxy flag */
	VS_SVC_F_DSNAT      = VS_CONN_F_DSNAT    /* dsnat flag */
	VS_SVC_F_MASK       = (VS_SVC_F_PERSISTENT |
		VS_SVC_F_ONEPACKET |
		VS_SVC_F_SYNPROXY |
		VS_SVC_F_DSNAT)
)

var (
	client *rpc.Client
	CmdOpt CmdOptions
)

type CallOptions struct {
	Opt  CmdOptions
	Args []string
}

type CmdOptions struct {
	/* status */
	Typ string
	Id  int
	/* service */
	Addr       Addr4
	Nic        uint
	Protocol   Protocol
	TCP        string
	UDP        string
	Sched_name string
	Flags      uint
	Number     int
	Timeout    uint
	Timeout_s  string
	Netmask    Be32

	/* dest */
	D           bool
	Dnic        uint
	Daddr       Addr4
	Conn_flags  uint
	Weight      int
	U_threshold uint
	L_threshold uint

	/* local addr */
	L    bool
	Lnic uint
	Lip  Be32

	/* timeout */
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

type Be32 uint32

func (p *Be32) Set(value string) error {
	if value == "" {
		return nil
	}
	if ip := net.ParseIP(value).To4(); ip != nil {
		*p = Htonl(ipToU32(ip))
	}
	return nil
}

func (p Be32) String() string {
	return be32_to_addr(p)
}

type Be16 uint16

func (p *Be16) Set(value string) error {
	if value == "" {
		return nil
	}

	port, err := strconv.Atoi(value)
	if err != nil {
		return err
	}

	*p = Htons(uint16(port))
	return nil
}

func (p Be16) String() string {
	return fmt.Sprintf("%d", Ntohs(p))
}

type Addr4 struct {
	Ip   Be32
	Port Be16
}

func (p *Addr4) Set(value string) error {
	if value == "" {
		p.Ip = 0
		p.Port = 0
		return nil
	}

	fields := strings.Split(value, ":")
	if len(fields) == 1 || len(fields) == 2 {
		if ip := net.ParseIP(fields[0]).To4(); ip != nil {
			p.Ip = Htonl(ipToU32(ip))
			if len(fields) == 2 {
				if port, err := strconv.Atoi(fields[1]); err != nil {
					return errIpv4Addr
				} else {
					p.Port = Htons(uint16(port))
				}
			} else {
				p.Port = 0
			}
			return nil
		}
	}
	return errIpv4Addr
}

func (p *Addr4) String() string {
	return fmt.Sprintf("%s:%d", be32_to_addr(p.Ip), Ntohs(p.Port))
}

type Vs_timeout_q struct {
	Cmd             int
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

func (t *Vs_timeout_q) Set(value string) (err error) {
	fields := strings.Split(value, ",")
	if len(fields) != 3 {
		return errTimeout
	}
	if t.Tcp_timeout, err = strconv.Atoi(fields[0]); err != nil {
		return err
	}
	if t.Tcp_fin_timeout, err = strconv.Atoi(fields[1]); err != nil {
		return err
	}
	if t.Udp_timeout, err = strconv.Atoi(fields[2]); err != nil {
		return err
	}
	return nil
}

type Vs_timeout_r struct {
	Code            int
	Msg             string
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

func (r Vs_timeout_r) String() string {
	if r.Code == 0 {
		return fmt.Sprintf("tcp:%d tcp_fin:%d udp:%d",
			r.Tcp_timeout, r.Tcp_fin_timeout, r.Udp_timeout)
	} else {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}
}

const (
	VS_STATS_IO = iota
	VS_STATS_WORKER
	VS_ESTATS_WORKER
	VS_STATS_DEV
	VS_STATS_CTL
	VS_STATS_MEM
	__VS_STATS_MAX
)

const (
	VS_CMD_UNSPEC       = iota
	VS_CMD_NEW_SERVICE  /* add service */
	VS_CMD_SET_SERVICE  /* modify service */
	VS_CMD_DEL_SERVICE  /* delete service */
	VS_CMD_GET_SERVICE  /* get service info */
	VS_CMD_GET_SERVICES /* get services info */
	VS_CMD_NEW_DEST     /* add destination */
	VS_CMD_SET_DEST     /* modify destination */
	VS_CMD_DEL_DEST     /* delete destination */
	VS_CMD_GET_DEST     /* get destination info */
	VS_CMD_NEW_DAEMON   /* start sync daemon */
	VS_CMD_DEL_DAEMON   /* stop sync daemon */
	VS_CMD_GET_DAEMON   /* get sync daemon status */
	VS_CMD_SET_CONFIG   /* set config settings */
	VS_CMD_GET_CONFIG   /* get config settings */
	VS_CMD_SET_INFO     /* only used in GET_INFO reply */
	VS_CMD_GET_INFO     /* get general IPVS info */
	VS_CMD_ZERO         /* zero all counters and stats */
	VS_CMD_FLUSH        /* flush services and dests */
	VS_CMD_NEW_LADDR    /* add local address */
	VS_CMD_DEL_LADDR    /* del local address */
	VS_CMD_GET_LADDR    /* dump local address */
	VS_CMD_GET_STATS    /* dump workers/ctl  stats */
	__VS_CMD_MAX
)

const (
	IPPROTO_IP      = 0   /* Dummy protocol for TCP		*/
	IPPROTO_ICMP    = 1   /* Internet Control Message Protocol	*/
	IPPROTO_IGMP    = 2   /* Internet Group Management Protocol	*/
	IPPROTO_IPIP    = 4   /* IPIP tunnels (older KA9Q tunnels use 94) */
	IPPROTO_TCP     = 6   /* Transmission Control Protocol	*/
	IPPROTO_EGP     = 8   /* Exterior Gateway Protocol		*/
	IPPROTO_PUP     = 12  /* PUP protocol				*/
	IPPROTO_UDP     = 17  /* User Datagram Protocol		*/
	IPPROTO_IDP     = 22  /* XNS IDP protocol			*/
	IPPROTO_DCCP    = 33  /* Datagram Congestion Control Protocol */
	IPPROTO_IPV6    = 41  /* IPv6-in-IPv4 tunnelling		*/
	IPPROTO_RSVP    = 46  /* RSVP protocol			*/
	IPPROTO_GRE     = 47  /* Cisco GRE tunnels (rfc 1701,1702)	*/
	IPPROTO_ESP     = 50  /* Encapsulation Security Payload protocol */
	IPPROTO_AH      = 51  /* Authentication Header protocol       */
	IPPROTO_BEETPH  = 94  /* IP option pseudo header for BEET */
	IPPROTO_PIM     = 103 /* Protocol Independent Multicast	*/
	IPPROTO_COMP    = 108 /* Compression Header protocol */
	IPPROTO_SCTP    = 132 /* Stream Control Transport Protocol	*/
	IPPROTO_UDPLITE = 136 /* UDP-Lite (RFC 3828)			*/
	IPPROTO_RAW     = 255 /* Raw IP packets			*/
)

type Protocol uint8

func (p *Protocol) Set(value string) error {
	if value == "tcp" {
		*p = IPPROTO_TCP
	} else if value == "udp" {
		*p = IPPROTO_UDP
	} else {
		return errProtocol
	}
	return nil
}

func (p *Protocol) String() string {
	return get_protocol_name(uint8(*p))
}

type Vs_timeout_user struct {
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

type Vs_cmd_q struct {
	Cmd int
}

type Vs_cmd_r struct {
	Code int
	Msg  string
}

func (r Vs_cmd_r) String() string {
	if r.Code == 0 {
		return fmt.Sprintf("done")
	} else {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}
}

type Vs_version_r struct {
	Code    int
	Msg     string
	Version int
	Size    int
}

func (r Vs_version_r) String() string {
	if r.Code != 0 {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}

	return fmt.Sprintf("version\t\t%d.%d.%d\n"+
		"conn table size\t%d",
		(r.Version>>16)&0xff,
		(r.Version>>8)&0xff,
		(r.Version)&0xff,
		r.Size)
}
func Vs_dial() error {
	conn, err := net.Dial("unix", URL)
	if err != nil {
		return err
	}
	client = jsonrpc.NewClient(conn)
	return nil
}

func Vs_close() {
	if client != nil {
		client.Close()
	}
}

func Get_version() (*Vs_version_r, error) {
	var reply Vs_version_r
	args := Vs_cmd_q{VS_CMD_GET_INFO}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Get_timeout(o *CmdOptions) (*Vs_timeout_r, error) {
	args := Vs_timeout_q{Cmd: VS_CMD_GET_CONFIG}
	reply := &Vs_timeout_r{}

	err := client.Call("api", args, reply)
	return reply, err
}

func Set_flush(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_cmd_q{VS_CMD_FLUSH}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Set_timeout(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_timeout_q{Cmd: VS_CMD_SET_CONFIG}

	if err := args.Set(o.Timeout_s); err != nil {
		return nil, err
	}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Set_zero(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_service_q{
		Cmd: VS_CMD_ZERO,
		Service: Vs_service_user{
			Addr:     o.Addr.Ip,
			Port:     o.Addr.Port,
			Protocol: uint8(o.Protocol),
		},
	}

	err := client.Call("api", args, &reply)
	return &reply, err
}
<commit_msg>fix: be32 unmarshalJSON<commit_after>/*
 * Copyright 2016 Xiaomi Corporation. All rights reserved.
 * Use of this source code is governed by a BSD-style
 * license that can be found in the LICENSE file.
 *
 * Authors:    Yu Bo <yubo@xiaomi.com>
 */
package govs

import (
	"encoding/json"
	"fmt"
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"strconv"
	"strings"
)

const (
	URL = "/tmp/dpvs.sock"

	/*
	 *      IPVS Connection Flags
	 */
	VS_CONN_F_FWD_MASK   = 0x000f /* mask for the fwd methods */
	VS_CONN_F_MASQ       = 0x0000 /* masquerading/NAT */
	VS_CONN_F_LOCALNODE  = 0x0001 /* local node */
	VS_CONN_F_TUNNEL     = 0x0002 /* tunneling */
	VS_CONN_F_DROUTE     = 0x0003 /* direct routing */
	VS_CONN_F_BYPASS     = 0x0004 /* cache bypass */
	VS_CONN_F_FULLNAT    = 0x0005 /* full nat */
	VS_CONN_F_DSNAT      = 0x0008 /* dsnat flag */
	VS_CONN_F_SYNC       = 0x0020 /* entry created by sync */
	VS_CONN_F_HASHED     = 0x0040 /* hashed entry */
	VS_CONN_F_NOOUTPUT   = 0x0080 /* no output packets */
	VS_CONN_F_INACTIVE   = 0x0100 /* not established */
	VS_CONN_F_OUT_SEQ    = 0x0200 /* must do output seq adjust */
	VS_CONN_F_IN_SEQ     = 0x0400 /* must do input seq adjust */
	VS_CONN_F_SEQ_MASK   = 0x0600 /* in/out sequence mask */
	VS_CONN_F_NO_CPORT   = 0x0800 /* no client port set yet */
	VS_CONN_F_TEMPLATE   = 0x1000 /* template, not connection */
	VS_CONN_F_ONE_PACKET = 0x2000 /* forward only one packet */
	VS_CONN_F_SYNPROXY   = 0x8000 /* syn proxy flag */

	VS_SVC_F_PERSISTENT = 0x0001             /* persistent port */
	VS_SVC_F_HASHED     = 0x0002             /* hashed entry */
	VS_SVC_F_ONEPACKET  = 0x0004             /* one-packet scheduling */
	VS_SVC_F_SYNPROXY   = VS_CONN_F_SYNPROXY /* synproxy flag */
	VS_SVC_F_DSNAT      = VS_CONN_F_DSNAT    /* dsnat flag */
	VS_SVC_F_MASK       = (VS_SVC_F_PERSISTENT |
		VS_SVC_F_ONEPACKET |
		VS_SVC_F_SYNPROXY |
		VS_SVC_F_DSNAT)
)

var (
	client *rpc.Client
	CmdOpt CmdOptions
)

type CallOptions struct {
	Opt  CmdOptions
	Args []string
}

type CmdOptions struct {
	/* status */
	Typ string
	Id  int
	/* service */
	Addr       Addr4
	Nic        uint
	Protocol   Protocol
	TCP        string
	UDP        string
	Sched_name string
	Flags      uint
	Number     int
	Timeout    uint
	Timeout_s  string
	Netmask    Be32

	/* dest */
	D           bool
	Dnic        uint
	Daddr       Addr4
	Conn_flags  uint
	Weight      int
	U_threshold uint
	L_threshold uint

	/* local addr */
	L    bool
	Lnic uint
	Lip  Be32

	/* timeout */
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

type Be32 uint32

func (p *Be32) Set(value string) error {
	if value == "" {
		return nil
	}
	if ip := net.ParseIP(value).To4(); ip != nil {
		*p = Htonl(ipToU32(ip))
	}
	return nil
}

func (p Be32) String() string {
	return be32_to_addr(p)
}

func (p *Be32) UnmarshalJSON(data []byte) error {
	var i int32
	if err := json.Unmarshal(data, &i); err != nil {
		return err
	}
	*p = Be32(uint32(i))
	return nil
}

type Be16 uint16

func (p *Be16) Set(value string) error {
	if value == "" {
		return nil
	}

	port, err := strconv.Atoi(value)
	if err != nil {
		return err
	}

	*p = Htons(uint16(port))
	return nil
}

func (p Be16) String() string {
	return fmt.Sprintf("%d", Ntohs(p))
}

type Addr4 struct {
	Ip   Be32
	Port Be16
}

func (p *Addr4) Set(value string) error {
	if value == "" {
		p.Ip = 0
		p.Port = 0
		return nil
	}

	fields := strings.Split(value, ":")
	if len(fields) == 1 || len(fields) == 2 {
		if ip := net.ParseIP(fields[0]).To4(); ip != nil {
			p.Ip = Htonl(ipToU32(ip))
			if len(fields) == 2 {
				if port, err := strconv.Atoi(fields[1]); err != nil {
					return errIpv4Addr
				} else {
					p.Port = Htons(uint16(port))
				}
			} else {
				p.Port = 0
			}
			return nil
		}
	}
	return errIpv4Addr
}

func (p *Addr4) String() string {
	return fmt.Sprintf("%s:%d", be32_to_addr(p.Ip), Ntohs(p.Port))
}

type Vs_timeout_q struct {
	Cmd             int
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

func (t *Vs_timeout_q) Set(value string) (err error) {
	fields := strings.Split(value, ",")
	if len(fields) != 3 {
		return errTimeout
	}
	if t.Tcp_timeout, err = strconv.Atoi(fields[0]); err != nil {
		return err
	}
	if t.Tcp_fin_timeout, err = strconv.Atoi(fields[1]); err != nil {
		return err
	}
	if t.Udp_timeout, err = strconv.Atoi(fields[2]); err != nil {
		return err
	}
	return nil
}

type Vs_timeout_r struct {
	Code            int
	Msg             string
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

func (r Vs_timeout_r) String() string {
	if r.Code == 0 {
		return fmt.Sprintf("tcp:%d tcp_fin:%d udp:%d",
			r.Tcp_timeout, r.Tcp_fin_timeout, r.Udp_timeout)
	} else {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}
}

const (
	VS_STATS_IO = iota
	VS_STATS_WORKER
	VS_ESTATS_WORKER
	VS_STATS_DEV
	VS_STATS_CTL
	VS_STATS_MEM
	__VS_STATS_MAX
)

const (
	VS_CMD_UNSPEC       = iota
	VS_CMD_NEW_SERVICE  /* add service */
	VS_CMD_SET_SERVICE  /* modify service */
	VS_CMD_DEL_SERVICE  /* delete service */
	VS_CMD_GET_SERVICE  /* get service info */
	VS_CMD_GET_SERVICES /* get services info */
	VS_CMD_NEW_DEST     /* add destination */
	VS_CMD_SET_DEST     /* modify destination */
	VS_CMD_DEL_DEST     /* delete destination */
	VS_CMD_GET_DEST     /* get destination info */
	VS_CMD_NEW_DAEMON   /* start sync daemon */
	VS_CMD_DEL_DAEMON   /* stop sync daemon */
	VS_CMD_GET_DAEMON   /* get sync daemon status */
	VS_CMD_SET_CONFIG   /* set config settings */
	VS_CMD_GET_CONFIG   /* get config settings */
	VS_CMD_SET_INFO     /* only used in GET_INFO reply */
	VS_CMD_GET_INFO     /* get general IPVS info */
	VS_CMD_ZERO         /* zero all counters and stats */
	VS_CMD_FLUSH        /* flush services and dests */
	VS_CMD_NEW_LADDR    /* add local address */
	VS_CMD_DEL_LADDR    /* del local address */
	VS_CMD_GET_LADDR    /* dump local address */
	VS_CMD_GET_STATS    /* dump workers/ctl  stats */
	__VS_CMD_MAX
)

const (
	IPPROTO_IP      = 0   /* Dummy protocol for TCP		*/
	IPPROTO_ICMP    = 1   /* Internet Control Message Protocol	*/
	IPPROTO_IGMP    = 2   /* Internet Group Management Protocol	*/
	IPPROTO_IPIP    = 4   /* IPIP tunnels (older KA9Q tunnels use 94) */
	IPPROTO_TCP     = 6   /* Transmission Control Protocol	*/
	IPPROTO_EGP     = 8   /* Exterior Gateway Protocol		*/
	IPPROTO_PUP     = 12  /* PUP protocol				*/
	IPPROTO_UDP     = 17  /* User Datagram Protocol		*/
	IPPROTO_IDP     = 22  /* XNS IDP protocol			*/
	IPPROTO_DCCP    = 33  /* Datagram Congestion Control Protocol */
	IPPROTO_IPV6    = 41  /* IPv6-in-IPv4 tunnelling		*/
	IPPROTO_RSVP    = 46  /* RSVP protocol			*/
	IPPROTO_GRE     = 47  /* Cisco GRE tunnels (rfc 1701,1702)	*/
	IPPROTO_ESP     = 50  /* Encapsulation Security Payload protocol */
	IPPROTO_AH      = 51  /* Authentication Header protocol       */
	IPPROTO_BEETPH  = 94  /* IP option pseudo header for BEET */
	IPPROTO_PIM     = 103 /* Protocol Independent Multicast	*/
	IPPROTO_COMP    = 108 /* Compression Header protocol */
	IPPROTO_SCTP    = 132 /* Stream Control Transport Protocol	*/
	IPPROTO_UDPLITE = 136 /* UDP-Lite (RFC 3828)			*/
	IPPROTO_RAW     = 255 /* Raw IP packets			*/
)

type Protocol uint8

func (p *Protocol) Set(value string) error {
	if value == "tcp" {
		*p = IPPROTO_TCP
	} else if value == "udp" {
		*p = IPPROTO_UDP
	} else {
		return errProtocol
	}
	return nil
}

func (p *Protocol) String() string {
	return get_protocol_name(uint8(*p))
}

type Vs_timeout_user struct {
	Tcp_timeout     int
	Tcp_fin_timeout int
	Udp_timeout     int
}

type Vs_cmd_q struct {
	Cmd int
}

type Vs_cmd_r struct {
	Code int
	Msg  string
}

func (r Vs_cmd_r) String() string {
	if r.Code == 0 {
		return fmt.Sprintf("done")
	} else {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}
}

type Vs_version_r struct {
	Code    int
	Msg     string
	Version int
	Size    int
}

func (r Vs_version_r) String() string {
	if r.Code != 0 {
		return fmt.Sprintf("%s:%s", Ecode(r.Code), r.Msg)
	}

	return fmt.Sprintf("version\t\t%d.%d.%d\n"+
		"conn table size\t%d",
		(r.Version>>16)&0xff,
		(r.Version>>8)&0xff,
		(r.Version)&0xff,
		r.Size)
}
func Vs_dial() error {
	conn, err := net.Dial("unix", URL)
	if err != nil {
		return err
	}
	client = jsonrpc.NewClient(conn)
	return nil
}

func Vs_close() {
	if client != nil {
		client.Close()
	}
}

func Get_version() (*Vs_version_r, error) {
	var reply Vs_version_r
	args := Vs_cmd_q{VS_CMD_GET_INFO}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Get_timeout(o *CmdOptions) (*Vs_timeout_r, error) {
	args := Vs_timeout_q{Cmd: VS_CMD_GET_CONFIG}
	reply := &Vs_timeout_r{}

	err := client.Call("api", args, reply)
	return reply, err
}

func Set_flush(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_cmd_q{VS_CMD_FLUSH}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Set_timeout(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_timeout_q{Cmd: VS_CMD_SET_CONFIG}

	if err := args.Set(o.Timeout_s); err != nil {
		return nil, err
	}

	err := client.Call("api", args, &reply)
	return &reply, err
}

func Set_zero(o *CmdOptions) (*Vs_cmd_r, error) {
	var reply Vs_cmd_r
	args := Vs_service_q{
		Cmd: VS_CMD_ZERO,
		Service: Vs_service_user{
			Addr:     o.Addr.Ip,
			Port:     o.Addr.Port,
			Protocol: uint8(o.Protocol),
		},
	}

	err := client.Call("api", args, &reply)
	return &reply, err
}
<|endoftext|>


<commit_before>package rst

import (
	"net/http"
	"strconv"
	"strings"
	"time"
)

// DefaultAccessControl defines a limited CORS policy that only allows simple
// cross-origin requests.
var DefaultAccessControl = &AccessControlResponse{
	Origin:      "*",
	Credentials: true,
	Headers:     nil,
	Methods:     nil,
	MaxAge:      24 * time.Hour,
}

// PermissiveAccessControl defines a permissive CORS policy in which all methods
// and all headers are allowed for all origins.
var PermissiveAccessControl = &AccessControlResponse{
	Origin:      "*",
	Credentials: true,
	Headers:     []string{},
	Methods:     []string{},
	MaxAge:      24 * time.Hour,
}

/*
Preflighter is implemented by endpoints wishing to customize the response to
a CORS preflighted request.

	func (e *endpoint) Preflight(req *rst.AccessControlRequest, vars rst.RouteVars, r *http.Request) *rst.AccessControlResponse {
		if time.Now().Hour() < 12 {
			return &rst.AccessControlResponse{
				Origin: "morning.example.com",
				Methods: []string{"GET"},
			}
		}

		return &rst.AccessControlResponse{
			Origin: "afternoon.example.com",
			Methods: []string{"POST"},
		}
	}
*/
type Preflighter interface {
	Preflight(*AccessControlRequest, RouteVars, *http.Request) *AccessControlResponse
}

// AccessControlRequest represents the headers of a CORS access control request.
type AccessControlRequest struct {
	Origin  string
	Method  string
	Headers []string
}

func (ac *AccessControlRequest) isEmpty() bool {
	return ac.Origin == "" && ac.Method == "" && len(ac.Headers) == 0
}

// ParseAccessControlRequest returns a new instance of AccessControlRequest
// filled with CORS headers found in r.
func ParseAccessControlRequest(r *http.Request) *AccessControlRequest {
	var headers []string
	if h := r.Header.Get("Access-Control-Request-Headers"); h != "" {
		headers = strings.Split(strings.Replace(r.Header.Get("Access-Control-Request-Headers"), " ", "", -1), ",")
	}
	return &AccessControlRequest{
		Origin:  r.Header.Get("Origin"),
		Method:  r.Header.Get("Access-Control-Request-Method"),
		Headers: headers,
	}

	// TODO: remove duplicated headers before serving them back.
}

// AccessControlResponse defines the response headers to a CORS access control
// request.
type AccessControlResponse struct {
	Origin      string
	Methods     []string // Empty array means any, nil means none.
	Headers     []string // Empty array means any, nil means none.
	Credentials bool
	MaxAge      time.Duration
}

type accessControlHandler struct {
	endpoint Endpoint
	*AccessControlResponse
}

func newAccessControlHandler(endpoint Endpoint, ac *AccessControlResponse) *accessControlHandler {
	return &accessControlHandler{
		endpoint:              endpoint,
		AccessControlResponse: ac,
	}
}

func (h *accessControlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, exists := r.Header["Origin"]; !exists {
		return
	}

	req := ParseAccessControlRequest(r)

	// If Options and endpoint implements Preflighter, call Preflight.
	var resp *AccessControlResponse
	if preflighter, implemented := h.endpoint.(Preflighter); implemented && strings.ToUpper(r.Method) == Options {
		resp = preflighter.Preflight(req, getVars(r), r)
	} else if h.AccessControlResponse != nil {
		resp = h.AccessControlResponse
	} else {
		return
	}
	w.Header().Add("Vary", "Origin")

	// Writing response headers
	if resp.Origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", resp.Origin)
	}
	w.Header().Set("Access-Control-Allow-Credentials", strconv.FormatBool(resp.Credentials))

	// OPTIONS only
	if strings.ToUpper(r.Method) != Options {
		return
	}

	w.Header().Set("Access-Control-Max-Age", strconv.Itoa(int(resp.MaxAge.Seconds())))

	if req.Method != "" && resp.Methods != nil {
		var methods []string
		if len(resp.Methods) == 0 {
			methods = AllowedMethods(h.endpoint)
		} else {
			methods = resp.Methods
		}
		w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ", "))
	}

	if len(req.Headers) > 0 && resp.Headers != nil {
		var headers []string
		if len(resp.Headers) == 0 {
			headers = req.Headers
		} else {
			headers = resp.Headers
		}
		w.Header().Set("Access-Control-Expose-Headers", strings.Join(headers, ", "))
	}
}
<commit_msg>Vary now altered when a specific origin is allowed<commit_after>package rst

import (
	"net/http"
	"strconv"
	"strings"
	"time"
)

// DefaultAccessControl defines a limited CORS policy that only allows simple
// cross-origin requests.
var DefaultAccessControl = &AccessControlResponse{
	Origin:      "*",
	Credentials: true,
	Headers:     nil,
	Methods:     nil,
	MaxAge:      24 * time.Hour,
}

// PermissiveAccessControl defines a permissive CORS policy in which all methods
// and all headers are allowed for all origins.
var PermissiveAccessControl = &AccessControlResponse{
	Origin:      "*",
	Credentials: true,
	Headers:     []string{},
	Methods:     []string{},
	MaxAge:      24 * time.Hour,
}

/*
Preflighter is implemented by endpoints wishing to customize the response to
a CORS preflighted request.

	func (e *endpoint) Preflight(req *rst.AccessControlRequest, vars rst.RouteVars, r *http.Request) *rst.AccessControlResponse {
		if time.Now().Hour() < 12 {
			return &rst.AccessControlResponse{
				Origin: "morning.example.com",
				Methods: []string{"GET"},
			}
		}

		return &rst.AccessControlResponse{
			Origin: "afternoon.example.com",
			Methods: []string{"POST"},
		}
	}
*/
type Preflighter interface {
	Preflight(*AccessControlRequest, RouteVars, *http.Request) *AccessControlResponse
}

// AccessControlRequest represents the headers of a CORS access control request.
type AccessControlRequest struct {
	Origin  string
	Method  string
	Headers []string
}

func (ac *AccessControlRequest) isEmpty() bool {
	return ac.Origin == "" && ac.Method == "" && len(ac.Headers) == 0
}

// ParseAccessControlRequest returns a new instance of AccessControlRequest
// filled with CORS headers found in r.
func ParseAccessControlRequest(r *http.Request) *AccessControlRequest {
	var headers []string
	if h := r.Header.Get("Access-Control-Request-Headers"); h != "" {
		headers = strings.Split(strings.Replace(r.Header.Get("Access-Control-Request-Headers"), " ", "", -1), ",")
	}
	return &AccessControlRequest{
		Origin:  r.Header.Get("Origin"),
		Method:  r.Header.Get("Access-Control-Request-Method"),
		Headers: headers,
	}

	// TODO: remove duplicated headers before serving them back.
}

// AccessControlResponse defines the response headers to a CORS access control
// request.
type AccessControlResponse struct {
	Origin      string
	Methods     []string // Empty array means any, nil means none.
	Headers     []string // Empty array means any, nil means none.
	Credentials bool
	MaxAge      time.Duration
}

type accessControlHandler struct {
	endpoint Endpoint
	*AccessControlResponse
}

func newAccessControlHandler(endpoint Endpoint, ac *AccessControlResponse) *accessControlHandler {
	return &accessControlHandler{
		endpoint:              endpoint,
		AccessControlResponse: ac,
	}
}

func (h *accessControlHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if _, exists := r.Header["Origin"]; !exists {
		return
	}

	req := ParseAccessControlRequest(r)

	// If Options and endpoint implements Preflighter, call Preflight.
	var resp *AccessControlResponse
	if preflighter, implemented := h.endpoint.(Preflighter); implemented && strings.ToUpper(r.Method) == Options {
		resp = preflighter.Preflight(req, getVars(r), r)
	} else if h.AccessControlResponse != nil {
		resp = h.AccessControlResponse
	} else {
		return
	}

	// Adding a vary if an origin is specified in the response.
	defer func() {
		if allowed := w.Header().Get("Access-Control-Allow-Origin"); allowed != "" && allowed != "*" {
			w.Header().Add("Vary", "Origin")
		}
	}()

	// Writing response headers
	if resp.Origin != "" {
		w.Header().Set("Access-Control-Allow-Origin", resp.Origin)
	}
	w.Header().Set("Access-Control-Allow-Credentials", strconv.FormatBool(resp.Credentials))

	// OPTIONS only
	if strings.ToUpper(r.Method) != Options {
		return
	}

	w.Header().Set("Access-Control-Max-Age", strconv.Itoa(int(resp.MaxAge.Seconds())))

	if req.Method != "" && resp.Methods != nil {
		var methods []string
		if len(resp.Methods) == 0 {
			methods = AllowedMethods(h.endpoint)
		} else {
			methods = resp.Methods
		}
		w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ", "))
	}

	if len(req.Headers) > 0 && resp.Headers != nil {
		var headers []string
		if len(resp.Headers) == 0 {
			headers = req.Headers
		} else {
			headers = resp.Headers
		}
		w.Header().Set("Access-Control-Expose-Headers", strings.Join(headers, ", "))
	}
}
<|endoftext|>


<commit_before>package cuid

import (
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"time"
)

type Counter interface {
	Next() int64
}

const (
	BLOCK_SIZE = 4
	BASE       = 36
)

var (
	counter        Counter = nil
	discreteValues         = int64(math.Pow(BASE, BLOCK_SIZE))
	fingerprint            = ""
	randomSource           = rand.NewSource(time.Now().Unix())
	random                 = rand.New(randomSource)
	padding                = strings.Repeat("0", BLOCK_SIZE)
)

func init() {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "dummy-host"
	}
	acc := int64(len(hostname) + BASE)
	for i := range hostname {
		acc = acc + int64(hostname[i])
	}

	hostID := pad(strconv.FormatInt(int64(os.Getpid()), BASE), 2)
	host := pad(strconv.FormatInt(acc, 10), 2)
	fingerprint = hostID + host
}

func New() string {
	if counter == nil {
		counter = NewDefaultCounter()
	}

	timestampBlock := strconv.FormatInt(time.Now().Unix()*1000, BASE)
	counterBlock := pad(strconv.FormatInt(counter.Next(), BASE), BLOCK_SIZE)
	randomBlock1 := pad(strconv.FormatInt(random.Int63n(discreteValues), BASE), BLOCK_SIZE)
	randomBlock2 := pad(strconv.FormatInt(random.Int63n(discreteValues), BASE), BLOCK_SIZE)

	return "c" + timestampBlock + counterBlock + fingerprint + randomBlock1 + randomBlock2
}

func pad(str string, size int) string {
	s := padding + str
	i := len(s) - size
	return s[i:]
}
<commit_msg>Cleanup and padding optimisation.<commit_after>package cuid

import (
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"time"
)

type Counter interface {
	Next() int64
}

const (
	BLOCK_SIZE = 4
	BASE       = 36
)

var (
	defaultCounter Counter = nil
	defaultRandom          = rand.New(rand.NewSource(time.Now().Unix()))
	discreteValues         = int64(math.Pow(BASE, BLOCK_SIZE))
	padding                = strings.Repeat("0", BLOCK_SIZE)
	fingerprint            = ""
)

func init() {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "dummy-host"
	}
	acc := int64(len(hostname) + BASE)
	for i := range hostname {
		acc = acc + int64(hostname[i])
	}

	hostID := pad(strconv.FormatInt(int64(os.Getpid()), BASE), 2)
	host := pad(strconv.FormatInt(acc, 10), 2)
	fingerprint = hostID + host
}

func New() string {
	if defaultCounter == nil {
		defaultCounter = NewDefaultCounter()
	}

	timestampBlock := strconv.FormatInt(time.Now().Unix()*1000, BASE)
	counterBlock := pad(strconv.FormatInt(defaultCounter.Next(), BASE), BLOCK_SIZE)
	randomBlock1 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)
	randomBlock2 := pad(strconv.FormatInt(defaultRandom.Int63n(discreteValues), BASE), BLOCK_SIZE)

	return "c" + timestampBlock + counterBlock + fingerprint + randomBlock1 + randomBlock2
}

func pad(str string, size int) string {
	if len(str) == size {
		return str
	}

	if len(str) < size {
		str = padding + str
	}

	i := len(str) - size

	return str[i:]
}
<|endoftext|>


<commit_before>// Copyright 2015 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package forms

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

// Data holds data obtained from the request body and url query parameters.
// Because Data is built from multiple sources, sometimes there will be more
// than one value for a given key. You can use Get, Set, Add, and Del to access
// the first element for a given key or access the Values and Files properties directly
// to access additional elements for a given key. You can also use helper methods to convert
// the first value for a given key to a different type (e.g. bool or int).
type Data struct {
	// Values holds any basic key-value string data
	// This includes all fields from a json body or
	// urlencoded form, and the form fields only (not
	// files) from a multipart form
	Values url.Values
	// Files holds files from a multipart form only.
	// For any other type of request, it will always
	// be empty. Files only supports one file per key,
	// since this is by far the most common use. If you
	// need to have more than one file per key, parse the
	// files manually using req.MultipartForm.File.
	Files map[string]*multipart.FileHeader
	// jsonBody holds the original body of the request.
	// Only available for json requests.
	jsonBody []byte
}

func newData() *Data {
	return &Data{
		Values: url.Values{},
		Files:  map[string]*multipart.FileHeader{},
	}
}

// Parse parses the request body and url query parameters into
// Data. The content in the body of the request has a higher priority,
// will be added to Data first, and will be the result of any operation
// which gets the first element for a given key (e.g. Get, GetInt, or GetBool).
func Parse(req *http.Request) (*Data, error) {
	data := newData()
	contentType := req.Header.Get("Content-Type")
	if strings.Contains(contentType, "multipart/form-data") {
		if err := req.ParseMultipartForm(2048); err != nil {
			return nil, err
		}
		for key, vals := range req.MultipartForm.Value {
			for _, val := range vals {
				data.Add(key, val)
			}
		}
		for key, files := range req.MultipartForm.File {
			if len(files) != 0 {
				data.AddFile(key, files[0])
			}
		}
	} else if strings.Contains(contentType, "form-urlencoded") {
		if err := req.ParseForm(); err != nil {
			return nil, err
		}
		for key, vals := range req.PostForm {
			for _, val := range vals {
				data.Add(key, val)
			}
		}
	} else if strings.Contains(contentType, "application/json") {
		body, err := ioutil.ReadAll(req.Body)
		if err != nil {
			return nil, err
		}
		data.jsonBody = body
		if err := parseJSON(data.Values, data.jsonBody); err != nil {
			return nil, err
		}
	}
	for key, vals := range req.URL.Query() {
		for _, val := range vals {
			data.Add(key, val)
		}
	}
	return data, nil
}

// CreateFromMap returns a Data object with keys and values matching
// the map.
func CreateFromMap(m map[string]string) *Data {
	data := newData()
	for key, value := range m {
		data.Add(key, value)
	}
	return data
}

func parseJSON(values url.Values, body []byte) error {
	if len(body) == 0 {
		// don't attempt to parse empty bodies
		return nil
	}
	rawData := map[string]interface{}{}
	if err := json.Unmarshal(body, &rawData); err != nil {
		return err
	}
	// Whatever the underlying type is, we need to convert it to a
	// string. There are only a few possible types, so we can just
	// do a type switch over the possibilities.
	for key, val := range rawData {
		switch val.(type) {
		case string, bool, float64:
			values.Add(key, fmt.Sprint(val))
		case nil:
			values.Add(key, "")
		case map[string]interface{}, []interface{}:
			// for more complicated data structures, convert back to
			// a JSON string and let user decide how to unmarshal
			jsonVal, err := json.Marshal(val)
			if err != nil {
				return err
			}
			values.Add(key, string(jsonVal))
		}
	}
	return nil
}

// Add adds the value to key. It appends to any existing values associated with key.
func (d *Data) Add(key string, value string) {
	d.Values.Add(key, value)
}

// AddFile adds the multipart form file to data with the given key.
func (d *Data) AddFile(key string, file *multipart.FileHeader) {
	d.Files[key] = file
}

// Del deletes the values associated with key.
func (d *Data) Del(key string) {
	d.Values.Del(key)
}

// DelFile deletes the file associated with key (if any).
// If there is no file associated with key, it does nothing.
func (d *Data) DelFile(key string) {
	delete(d.Files, key)
}

// Encode encodes the values into “URL encoded” form ("bar=baz&foo=quux") sorted by key.
// Any files in d will be ignored because there is no direct way to convert a file to a
// URL encoded value.
func (d *Data) Encode() string {
	return d.Values.Encode()
}

// Get gets the first value associated with the given key. If there are no values
// associated with the key, Get returns the empty string. To access multiple values,
// use the map directly.
func (d Data) Get(key string) string {
	return d.Values.Get(key)
}

// GetFile returns the multipart form file associated with key, if any, as a *multipart.FileHeader.
// If there is no file associated with key, it returns nil. If you just want the body of the
// file, use GetFileBytes.
func (d Data) GetFile(key string) *multipart.FileHeader {
	return d.Files[key]
}

// Set sets the key to value. It replaces any existing values.
func (d *Data) Set(key string, value string) {
	d.Values.Set(key, value)
}

// KeyExists returns true iff data.Values[key] exists. When parsing a request body, the key
// is considered to be in existence if it was provided in the request body, even if its value
// is empty.
func (d Data) KeyExists(key string) bool {
	_, found := d.Values[key]
	return found
}

// FileExists returns true iff data.Files[key] exists. When parsing a request body, the key
// is considered to be in existence if it was provided in the request body, even if the file
// is empty.
func (d Data) FileExists(key string) bool {
	_, found := d.Files[key]
	return found
}

// GetInt returns the first element in data[key] converted to an int.
func (d Data) GetInt(key string) int {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return 0
	}
	str := d.Get(key)
	if result, err := strconv.Atoi(str); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetFloat returns the first element in data[key] converted to a float.
func (d Data) GetFloat(key string) float64 {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return 0.0
	}
	str := d.Get(key)
	if result, err := strconv.ParseFloat(str, 64); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetBool returns the first element in data[key] converted to a bool.
func (d Data) GetBool(key string) bool {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return false
	}
	str := d.Get(key)
	if result, err := strconv.ParseBool(str); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetBytes returns the first element in data[key] converted to a slice of bytes.
func (d Data) GetBytes(key string) []byte {
	return []byte(d.Get(key))
}

// GetFileBytes returns the body of the file associated with key. If there is no
// file associated with key, it returns nil (not an error). It may return an error if
// there was a problem reading the file. If you need to know whether or not the file
// exists (i.e. whether it was provided in the request), use the FileExists method.
func (d Data) GetFileBytes(key string) ([]byte, error) {
	fileHeader, found := d.Files[key]
	if !found {
		return nil, nil
	} else {
		file, err := fileHeader.Open()
		if err != nil {
			return nil, err
		}
		return ioutil.ReadAll(file)
	}
}

// GetStringsSplit returns the first element in data[key] split into a slice delimited by delim.
func (d Data) GetStringsSplit(key string, delim string) []string {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil
	}
	return strings.Split(d.Values[key][0], delim)
}

// BindJSON binds v to the json data in the request body. It calls json.Unmarshal and
// sets the value of v.
func (d Data) BindJSON(v interface{}) error {
	if len(d.jsonBody) == 0 {
		return nil
	}
	return json.Unmarshal(d.jsonBody, v)
}

// GetMapFromJSON assumes that the first element in data[key] is a json string, attempts to
// unmarshal it into a map[string]interface{}, and if successful, returns the result. If
// unmarshaling was not successful, returns an error.
func (d Data) GetMapFromJSON(key string) (map[string]interface{}, error) {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil, nil
	}
	result := map[string]interface{}{}
	if err := json.Unmarshal([]byte(d.Get(key)), &result); err != nil {
		return nil, err
	} else {
		return result, nil
	}
}

// GetSliceFromJSON assumes that the first element in data[key] is a json string, attempts to
// unmarshal it into a []interface{}, and if successful, returns the result. If unmarshaling
// was not successful, returns an error.
func (d Data) GetSliceFromJSON(key string) ([]interface{}, error) {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil, nil
	}
	result := []interface{}{}
	if err := json.Unmarshal([]byte(d.Get(key)), &result); err != nil {
		return nil, err
	} else {
		return result, nil
	}
}

// GetAndUnmarshalJSON assumes that the first element in data[key] is a json string and
// attempts to unmarshal it into v. If unmarshaling was not successful, returns an error.
// v should be a pointer to some data structure.
func (d Data) GetAndUnmarshalJSON(key string, v interface{}) error {
	if err := json.Unmarshal([]byte(d.Get(key)), v); err != nil {
		return err
	}
	return nil
}

// Validator returns a Validator which can be used to easily validate data.
func (d *Data) Validator() *Validator {
	return &Validator{
		data: d,
	}
}
<commit_msg>adding ParseMax w/ max form size in signature - forwarded to from Parse<commit_after>// Copyright 2015 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

package forms

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

const DefaultMaxFormSize = 2048

// Data holds data obtained from the request body and url query parameters.
// Because Data is built from multiple sources, sometimes there will be more
// than one value for a given key. You can use Get, Set, Add, and Del to access
// the first element for a given key or access the Values and Files properties directly
// to access additional elements for a given key. You can also use helper methods to convert
// the first value for a given key to a different type (e.g. bool or int).
type Data struct {
	// Values holds any basic key-value string data
	// This includes all fields from a json body or
	// urlencoded form, and the form fields only (not
	// files) from a multipart form
	Values url.Values
	// Files holds files from a multipart form only.
	// For any other type of request, it will always
	// be empty. Files only supports one file per key,
	// since this is by far the most common use. If you
	// need to have more than one file per key, parse the
	// files manually using req.MultipartForm.File.
	Files map[string]*multipart.FileHeader
	// jsonBody holds the original body of the request.
	// Only available for json requests.
	jsonBody []byte
}

func newData() *Data {
	return &Data{
		Values: url.Values{},
		Files:  map[string]*multipart.FileHeader{},
	}
}

// ParseMax parses the request body and url query parameters into
// Data. The content in the body of the request has a higher priority,
// will be added to Data first, and will be the result of any operation
// which gets the first element for a given key (e.g. Get, GetInt, or GetBool).
func ParseMax(req *http.Request, max int64) (*Data, error) {
	data := newData()
	contentType := req.Header.Get("Content-Type")
	if strings.Contains(contentType, "multipart/form-data") {
		if err := req.ParseMultipartForm(max); err != nil {
			return nil, err
		}
		for key, vals := range req.MultipartForm.Value {
			for _, val := range vals {
				data.Add(key, val)
			}
		}
		for key, files := range req.MultipartForm.File {
			if len(files) != 0 {
				data.AddFile(key, files[0])
			}
		}
	} else if strings.Contains(contentType, "form-urlencoded") {
		if err := req.ParseForm(); err != nil {
			return nil, err
		}
		for key, vals := range req.PostForm {
			for _, val := range vals {
				data.Add(key, val)
			}
		}
	} else if strings.Contains(contentType, "application/json") {
		body, err := ioutil.ReadAll(req.Body)
		if err != nil {
			return nil, err
		}
		data.jsonBody = body
		if err := parseJSON(data.Values, data.jsonBody); err != nil {
			return nil, err
		}
	}
	for key, vals := range req.URL.Query() {
		for _, val := range vals {
			data.Add(key, val)
		}
	}
	return data, nil
}

// Parse uses the default max form size defined above and calls ParseMax
func Parse(req *http.Request) (*Data, error) {
	return ParseMax(req, DefaultMaxFormSize)
}

// CreateFromMap returns a Data object with keys and values matching
// the map.
func CreateFromMap(m map[string]string) *Data {
	data := newData()
	for key, value := range m {
		data.Add(key, value)
	}
	return data
}

func parseJSON(values url.Values, body []byte) error {
	if len(body) == 0 {
		// don't attempt to parse empty bodies
		return nil
	}
	rawData := map[string]interface{}{}
	if err := json.Unmarshal(body, &rawData); err != nil {
		return err
	}
	// Whatever the underlying type is, we need to convert it to a
	// string. There are only a few possible types, so we can just
	// do a type switch over the possibilities.
	for key, val := range rawData {
		switch val.(type) {
		case string, bool, float64:
			values.Add(key, fmt.Sprint(val))
		case nil:
			values.Add(key, "")
		case map[string]interface{}, []interface{}:
			// for more complicated data structures, convert back to
			// a JSON string and let user decide how to unmarshal
			jsonVal, err := json.Marshal(val)
			if err != nil {
				return err
			}
			values.Add(key, string(jsonVal))
		}
	}
	return nil
}

// Add adds the value to key. It appends to any existing values associated with key.
func (d *Data) Add(key string, value string) {
	d.Values.Add(key, value)
}

// AddFile adds the multipart form file to data with the given key.
func (d *Data) AddFile(key string, file *multipart.FileHeader) {
	d.Files[key] = file
}

// Del deletes the values associated with key.
func (d *Data) Del(key string) {
	d.Values.Del(key)
}

// DelFile deletes the file associated with key (if any).
// If there is no file associated with key, it does nothing.
func (d *Data) DelFile(key string) {
	delete(d.Files, key)
}

// Encode encodes the values into “URL encoded” form ("bar=baz&foo=quux") sorted by key.
// Any files in d will be ignored because there is no direct way to convert a file to a
// URL encoded value.
func (d *Data) Encode() string {
	return d.Values.Encode()
}

// Get gets the first value associated with the given key. If there are no values
// associated with the key, Get returns the empty string. To access multiple values,
// use the map directly.
func (d Data) Get(key string) string {
	return d.Values.Get(key)
}

// GetFile returns the multipart form file associated with key, if any, as a *multipart.FileHeader.
// If there is no file associated with key, it returns nil. If you just want the body of the
// file, use GetFileBytes.
func (d Data) GetFile(key string) *multipart.FileHeader {
	return d.Files[key]
}

// Set sets the key to value. It replaces any existing values.
func (d *Data) Set(key string, value string) {
	d.Values.Set(key, value)
}

// KeyExists returns true iff data.Values[key] exists. When parsing a request body, the key
// is considered to be in existence if it was provided in the request body, even if its value
// is empty.
func (d Data) KeyExists(key string) bool {
	_, found := d.Values[key]
	return found
}

// FileExists returns true iff data.Files[key] exists. When parsing a request body, the key
// is considered to be in existence if it was provided in the request body, even if the file
// is empty.
func (d Data) FileExists(key string) bool {
	_, found := d.Files[key]
	return found
}

// GetInt returns the first element in data[key] converted to an int.
func (d Data) GetInt(key string) int {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return 0
	}
	str := d.Get(key)
	if result, err := strconv.Atoi(str); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetFloat returns the first element in data[key] converted to a float.
func (d Data) GetFloat(key string) float64 {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return 0.0
	}
	str := d.Get(key)
	if result, err := strconv.ParseFloat(str, 64); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetBool returns the first element in data[key] converted to a bool.
func (d Data) GetBool(key string) bool {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return false
	}
	str := d.Get(key)
	if result, err := strconv.ParseBool(str); err != nil {
		panic(err)
	} else {
		return result
	}
}

// GetBytes returns the first element in data[key] converted to a slice of bytes.
func (d Data) GetBytes(key string) []byte {
	return []byte(d.Get(key))
}

// GetFileBytes returns the body of the file associated with key. If there is no
// file associated with key, it returns nil (not an error). It may return an error if
// there was a problem reading the file. If you need to know whether or not the file
// exists (i.e. whether it was provided in the request), use the FileExists method.
func (d Data) GetFileBytes(key string) ([]byte, error) {
	fileHeader, found := d.Files[key]
	if !found {
		return nil, nil
	} else {
		file, err := fileHeader.Open()
		if err != nil {
			return nil, err
		}
		return ioutil.ReadAll(file)
	}
}

// GetStringsSplit returns the first element in data[key] split into a slice delimited by delim.
func (d Data) GetStringsSplit(key string, delim string) []string {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil
	}
	return strings.Split(d.Values[key][0], delim)
}

// BindJSON binds v to the json data in the request body. It calls json.Unmarshal and
// sets the value of v.
func (d Data) BindJSON(v interface{}) error {
	if len(d.jsonBody) == 0 {
		return nil
	}
	return json.Unmarshal(d.jsonBody, v)
}

// GetMapFromJSON assumes that the first element in data[key] is a json string, attempts to
// unmarshal it into a map[string]interface{}, and if successful, returns the result. If
// unmarshaling was not successful, returns an error.
func (d Data) GetMapFromJSON(key string) (map[string]interface{}, error) {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil, nil
	}
	result := map[string]interface{}{}
	if err := json.Unmarshal([]byte(d.Get(key)), &result); err != nil {
		return nil, err
	} else {
		return result, nil
	}
}

// GetSliceFromJSON assumes that the first element in data[key] is a json string, attempts to
// unmarshal it into a []interface{}, and if successful, returns the result. If unmarshaling
// was not successful, returns an error.
func (d Data) GetSliceFromJSON(key string) ([]interface{}, error) {
	if !d.KeyExists(key) || len(d.Values[key]) == 0 {
		return nil, nil
	}
	result := []interface{}{}
	if err := json.Unmarshal([]byte(d.Get(key)), &result); err != nil {
		return nil, err
	} else {
		return result, nil
	}
}

// GetAndUnmarshalJSON assumes that the first element in data[key] is a json string and
// attempts to unmarshal it into v. If unmarshaling was not successful, returns an error.
// v should be a pointer to some data structure.
func (d Data) GetAndUnmarshalJSON(key string, v interface{}) error {
	if err := json.Unmarshal([]byte(d.Get(key)), v); err != nil {
		return err
	}
	return nil
}

// Validator returns a Validator which can be used to easily validate data.
func (d *Data) Validator() *Validator {
	return &Validator{
		data: d,
	}
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package api

import (
	"net/http"
	"strings"

	"github.com/vmware/harbor/dao"
	"github.com/vmware/harbor/models"
	svc_utils "github.com/vmware/harbor/service/utils"
	"github.com/vmware/harbor/utils/log"
)

// StatisticAPI handles request to /api/statistics/
type StatisticAPI struct {
	BaseAPI
	userID   int
	username string
}

//Prepare validates the URL and the user
func (s *StatisticAPI) Prepare() {
	userID, ok := s.GetSession("userId").(int)
	if !ok {
		s.userID = dao.NonExistUserID
	} else {
		s.userID = userID
	}
	username, ok := s.GetSession("username").(string)
	if !ok {
		log.Warning("failed to get username from session")
		s.username = ""
	} else {
		s.username = username
	}
}

// Get total projects and repos of the user
func (s *StatisticAPI) Get() {
	queryProject := models.Project{UserID: s.userID}
	projectList, err := dao.QueryProject(queryProject)
	if err != nil {
		log.Errorf("Error occured in QueryProject, error: %v", err)
		s.CustomAbort(http.StatusInternalServerError, "Internal error.")
	}
	proMap := map[string]int{}
	isAdmin, err := dao.IsAdminRole(s.userID)
	if err != nil {
		log.Errorf("Error occured in check admin, error: %v", err)
		s.CustomAbort(http.StatusInternalServerError, "Internal error.")
	}
	if isAdmin {
		proMap["total_project_count"] = len(projectList)
		proMap["total_repo_count"] = getTotalRepoCount()
	}
	for i := 0; i < len(projectList); i++ {
		if projectList[i].Role == models.PROJECTADMIN || projectList[i].Role == models.DEVELOPER ||
			projectList[i].Role == models.GUEST {
			proMap["my_project_count"]++
			proMap["my_repo_count"] += getRepoCountByProject(projectList[i].Name)
		}
		if projectList[i].Public == 1 {
			proMap["public_project_count"]++
			proMap["public_repo_count"] += getRepoCountByProject(projectList[i].Name)
		}
	}
	s.Data["json"] = proMap
	s.ServeJSON()
}

//getReposByProject returns repo numbers of specified project
func getRepoCountByProject(projectName string) int {
	repoList, err := svc_utils.GetRepoFromCache()
	if err != nil {
		log.Errorf("Failed to get repo from cache, error: %v", err)
		return 0
	}
	var resp int
	if len(projectName) > 0 {
		for _, r := range repoList {
			if strings.Contains(r, "/") && r[0:strings.LastIndex(r, "/")] == projectName {
				resp += 1
			}
		}
		return resp
	}
	return 0
}

//getTotalRepoCount returns total repo count
func getTotalRepoCount() int {
	repoList, err := svc_utils.GetRepoFromCache()
	if err != nil {
		log.Errorf("Failed to get repo from cache, error: %v", err)
		return 0
	}
	return len(repoList)

}
<commit_msg>modify code where CI unpassed<commit_after>/*
   Copyright (c) 2016 VMware, Inc. All Rights Reserved.
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package api

import (
	"net/http"
	"strings"

	"github.com/vmware/harbor/dao"
	"github.com/vmware/harbor/models"
	svc_utils "github.com/vmware/harbor/service/utils"
	"github.com/vmware/harbor/utils/log"
)

// StatisticAPI handles request to /api/statistics/
type StatisticAPI struct {
	BaseAPI
	userID   int
	username string
}

//Prepare validates the URL and the user
func (s *StatisticAPI) Prepare() {
	userID, ok := s.GetSession("userId").(int)
	if !ok {
		s.userID = dao.NonExistUserID
	} else {
		s.userID = userID
	}
	username, ok := s.GetSession("username").(string)
	if !ok {
		log.Warning("failed to get username from session")
		s.username = ""
	} else {
		s.username = username
	}
}

// Get total projects and repos of the user
func (s *StatisticAPI) Get() {
	queryProject := models.Project{UserID: s.userID}
	projectList, err := dao.QueryProject(queryProject)
	if err != nil {
		log.Errorf("Error occured in QueryProject, error: %v", err)
		s.CustomAbort(http.StatusInternalServerError, "Internal error.")
	}
	proMap := map[string]int{}
	isAdmin, err := dao.IsAdminRole(s.userID)
	if err != nil {
		log.Errorf("Error occured in check admin, error: %v", err)
		s.CustomAbort(http.StatusInternalServerError, "Internal error.")
	}
	if isAdmin {
		proMap["total_project_count"] = len(projectList)
		proMap["total_repo_count"] = getTotalRepoCount()
	}
	for i := 0; i < len(projectList); i++ {
		if projectList[i].Role == models.PROJECTADMIN || projectList[i].Role == models.DEVELOPER ||
			projectList[i].Role == models.GUEST {
			proMap["my_project_count"]++
			proMap["my_repo_count"] += getRepoCountByProject(projectList[i].Name)
		}
		if projectList[i].Public == 1 {
			proMap["public_project_count"]++
			proMap["public_repo_count"] += getRepoCountByProject(projectList[i].Name)
		}
	}
	s.Data["json"] = proMap
	s.ServeJSON()
}

//getReposByProject returns repo numbers of specified project
func getRepoCountByProject(projectName string) int {
	repoList, err := svc_utils.GetRepoFromCache()
	if err != nil {
		log.Errorf("Failed to get repo from cache, error: %v", err)
		return 0
	}
	var resp int
	if len(projectName) > 0 {
		for _, r := range repoList {
			if strings.Contains(r, "/") && r[0:strings.LastIndex(r, "/")] == projectName {
				resp++
			}
		}
		return resp
	}
	return 0
}

//getTotalRepoCount returns total repo count
func getTotalRepoCount() int {
	repoList, err := svc_utils.GetRepoFromCache()
	if err != nil {
		log.Errorf("Failed to get repo from cache, error: %v", err)
		return 0
	}
	return len(repoList)

}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/golang/glog"
	"github.com/oikomi/FishChatServer2/codec"
	"github.com/oikomi/FishChatServer2/libnet"
	"github.com/oikomi/FishChatServer2/server/gateway/conf"
	"github.com/oikomi/FishChatServer2/server/gateway/job"
	"github.com/oikomi/FishChatServer2/server/gateway/server"
)

func init() {
	flag.Set("alsologtostderr", "true")
	flag.Set("log_dir", "false")
}

func main() {
	var err error
	flag.Parse()
	if err = conf.Init(); err != nil { // 载入配置文件gateway.toml进行初始化
		glog.Error("conf.Init() error: ", err)
		panic(err)
	}
	gwServer := server.New()     // gateway的server对象，然后下面的libnetServe负责创建server，监听端口
	protobuf := codec.Protobuf() // protobuf协议解析
	if gwServer.Server, err = libnet.Serve(conf.Conf.Server.Proto, conf.Conf.Server.Addr, protobuf, 0); err != nil {//创建server，监听端口
		glog.Error(err)
		panic(err)
	}
	go job.ConfDiscoveryProc() //启一个协程监控etcd服务，有配置更新的消息推送则更新本地配置
	gwServer.Loop() //服务循环
}
<commit_msg>update github import path<commit_after>package main

import (
	"flag"

	"github.com/golang/glog"
	"github.com/liberalman/im/codec"
	"github.com/liberalman/im/libnet"
	"github.com/liberalman/im/server/gateway/conf"
	"github.com/liberalman/im/server/gateway/job"
	"github.com/liberalman/im/server/gateway/server"
)

func init() {
	flag.Set("alsologtostderr", "true")
	flag.Set("log_dir", "false")
}

func main() {
	var err error
	flag.Parse()
	if err = conf.Init(); err != nil { // 载入配置文件gateway.toml进行初始化
		glog.Error("conf.Init() error: ", err)
		panic(err)
	}
	gwServer := server.New()                                                                                // gateway的server对象，然后下面的libnetServe负责创建server，监听端口
	protobuf := codec.Protobuf()                                                                                // protobuf协议解析
	if gwServer.Server, err = libnet.Serve(conf.Conf.Server.Proto, conf.Conf.Server.Addr, protobuf, 0); err != nil { //创建server，监听端口
		glog.Error(err)
		panic(err)
	}
	go job.ConfDiscoveryProc() //启一个协程监控etcd服务，有配置更新的消息推送则更新本地配置
	gwServer.Loop()            //服务循环
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"code.google.com/p/go-uuid/uuid"
	"github.com/pachyderm/pfs/lib/btrfs"
	"github.com/pachyderm/pfs/lib/mapreduce"
)

var dataRepo, compRepo string
var logFile string
var shard, modulos uint64

func parseArgs() {
	// os.Args[1] looks like 2-16
	dataRepo = "data-" + os.Args[1] + btrfs.RandSeq(4)
	compRepo = "comp-" + os.Args[1] + btrfs.RandSeq(4)
	logFile = "log-" + os.Args[1] + btrfs.RandSeq(4)
	s_m := strings.Split(os.Args[1], "-")
	var err error
	shard, err = strconv.ParseUint(s_m[0], 10, 64)
	if err != nil {
		log.Fatal(err)
	}
	modulos, err = strconv.ParseUint(s_m[1], 10, 64)
	if err != nil {
		log.Fatal(err)
	}
}

var jobDir string = "job"

func commitParam(r *http.Request) string {
	if p := r.URL.Query().Get("commit"); p != "" {
		return p
	}
	return "master"
}

func branchParam(r *http.Request) string {
	if p := r.URL.Query().Get("branch"); p != "" {
		return p
	}
	return "master"
}

func materializeParam(r *http.Request) string {
	if _, ok := r.URL.Query()["run"]; ok {
		return "true"
	}
	return "false"
}

func indexOf(haystack []string, needle string) int {
	for i, s := range haystack {
		if s == needle {
			return i
		}
	}
	return -1
}

func cat(w http.ResponseWriter, name string) {
	f, err := btrfs.Open(name)
	if err != nil {
		http.Error(w, err.Error(), 500)
		log.Print(err)
	}
	defer f.Close()

	if _, err := io.Copy(w, f); err != nil {
		http.Error(w, err.Error(), 500)
		log.Print(err)
	}
}

// genericFileHandler serves files from fs. It's used after branch and commit
// info have already been extracted and ignores those aspects of the URL.
func genericFileHandler(fs string, w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like: /foo/bar/.../file/<file>
	fileStart := indexOf(url, "file") + 1
	// file is the path in the filesystem we're getting
	file := path.Join(append([]string{fs}, url[fileStart:]...)...)

	if r.Method == "GET" {
		if strings.Contains(file, "*") {
			if !strings.HasSuffix(file, "*") {
				http.Error(w, "Illegal path containing internal `*`. `*` is currently only allowed as the last character of a path.", 400)
			} else {
				dir := path.Dir(file)
				files, err := btrfs.ReadDir(dir)
				if err != nil {
					http.Error(w, err.Error(), 500)
					return
				}
				for _, fi := range files {
					if fi.IsDir() {
						continue
					} else {
						cat(w, path.Join(dir, fi.Name()))
					}
				}
			}
		} else {
			cat(w, file)
		}
	} else if r.Method == "POST" {
		btrfs.MkdirAll(path.Dir(file))
		size, err := btrfs.CreateFromReader(file, r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created %s, size: %d.\n", file, size)
	} else if r.Method == "PUT" {
		btrfs.MkdirAll(path.Dir(file))
		size, err := btrfs.WriteFile(file, r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created %s, size: %d.\n", file, size)
	} else if r.Method == "DELETE" {
		if err := btrfs.Remove(file); err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Deleted %s.\n", file)
	}
}

// FileHandler is the core route for modifying the contents of the fileystem.
func FileHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" || r.Method == "DELETE" || r.Method == "PUT" {
		genericFileHandler(path.Join(dataRepo, branchParam(r)), w, r)
	} else if r.Method == "GET" {
		genericFileHandler(path.Join(dataRepo, commitParam(r)), w, r)
	} else {
		http.Error(w, "Invalid method.", 405)
	}
}

// CommitHandler creates a snapshot of outstanding changes.
func CommitHandler(w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, commit, <commit>, file, <file>]
	if len(url) > 3 && url[3] == "file" {
		genericFileHandler(path.Join(dataRepo, url[2]), w, r)
		return
	}
	if r.Method == "GET" {
		commits, err := btrfs.ReadDir(dataRepo)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		for _, ci := range commits {
			if uuid.Parse(ci.Name()) != nil {
				fmt.Fprintf(w, "%s    %s\n", ci.Name(), ci.ModTime().Format("2006-01-02T15:04:05.999999-07:00"))
			}
		}
	} else if r.Method == "POST" {
		err := btrfs.Commit(dataRepo, commitParam(r), branchParam(r))
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		if materializeParam(r) == "true" {
			go func() {
				err := mapreduce.Materialize(dataRepo, branchParam(r), commitParam(r), compRepo, jobDir, shard, modulos)
				if err != nil {
					log.Print(err)
				}
			}()
		}

		fmt.Fprint(w, commitParam(r))
	} else {
		http.Error(w, "Unsupported method.", http.StatusMethodNotAllowed)
		log.Printf("Unsupported method %s in request to %s.", r.Method, r.URL.String())
		return
	}
}

// BranchHandler creates a new branch from commit.
func BranchHandler(w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, commit, <commit>, file, <file>]
	if len(url) > 3 && url[3] == "file" {
		genericFileHandler(path.Join(dataRepo, url[2]), w, r)
		return
	}
	if r.Method == "GET" {
		branches, err := btrfs.ReadDir(dataRepo)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		for _, bi := range branches {
			if uuid.Parse(bi.Name()) == nil {
				fmt.Fprintf(w, "%s    %s\n", bi.Name(), bi.ModTime().Format("2006-01-02T15:04:05.999999-07:00"))
			}
		}
	} else if r.Method == "POST" {
		if err := btrfs.Branch(dataRepo, commitParam(r), branchParam(r)); err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created branch. (%s) -> %s.\n", commitParam(r), branchParam(r))
	} else {
		http.Error(w, "Invalid method.", 405)
		log.Print("Invalid method %s.", r.Method)
		return
	}
}

func JobHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("URL in job handler:\n", r.URL)
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, job, <job>, file, <file>]
	if r.Method == "GET" && len(url) > 3 && url[3] == "file" {
		if commitParam(r) != "master" {
			mapreduce.WaitJob(dataRepo, commitParam(r), url[2])
		}
		genericFileHandler(path.Join(compRepo, commitParam(r), url[2]), w, r)
		return
	} else if r.Method == "POST" {
		r.URL.Path = path.Join("/file", jobDir, url[2])
		log.Print("URL with reset path:\n", r.URL)
		genericFileHandler(path.Join(dataRepo, branchParam(r)), w, r)
	}
}

// MasterMux creates a multiplexer for a Master writing to the passed in FS.
func MasterMux() *http.ServeMux {
	mux := http.NewServeMux()

	mux.HandleFunc("/branch", BranchHandler)
	mux.HandleFunc("/commit", CommitHandler)
	mux.HandleFunc("/file/", FileHandler)
	mux.HandleFunc("/job/", JobHandler)
	mux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, "pong\n") })

	return mux
}

// RunServer runs a master server listening on port 80
func RunServer() {
	http.ListenAndServe(":80", MasterMux())
}

func main() {
	log.SetFlags(log.Lshortfile)
	parseArgs()
	if err := os.MkdirAll("/var/lib/pfs/log", 0777); err != nil {
		log.Fatal(err)
	}
	logF, err := os.Create(path.Join("/var/lib/pfs/log", logFile))
	if err != nil {
		log.Fatal(err)
	}
	defer logF.Close()
	log.Print("Logging to: ", path.Join("/var/lib/pfs/log", logFile))
	log.SetOutput(logF)
	if err := btrfs.Ensure(dataRepo); err != nil {
		log.Fatal(err)
	}
	if err := btrfs.Ensure(compRepo); err != nil {
		log.Fatal(err)
	}
	log.Print("Listening on port 80...")
	log.Printf("dataRepo: %s, compRepo: %s.", dataRepo, compRepo)
	RunServer()
}
<commit_msg>Remove call to WaitJob.<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"code.google.com/p/go-uuid/uuid"
	"github.com/pachyderm/pfs/lib/btrfs"
	"github.com/pachyderm/pfs/lib/mapreduce"
)

var dataRepo, compRepo string
var logFile string
var shard, modulos uint64

func parseArgs() {
	// os.Args[1] looks like 2-16
	dataRepo = "data-" + os.Args[1] + btrfs.RandSeq(4)
	compRepo = "comp-" + os.Args[1] + btrfs.RandSeq(4)
	logFile = "log-" + os.Args[1] + btrfs.RandSeq(4)
	s_m := strings.Split(os.Args[1], "-")
	var err error
	shard, err = strconv.ParseUint(s_m[0], 10, 64)
	if err != nil {
		log.Fatal(err)
	}
	modulos, err = strconv.ParseUint(s_m[1], 10, 64)
	if err != nil {
		log.Fatal(err)
	}
}

var jobDir string = "job"

func commitParam(r *http.Request) string {
	if p := r.URL.Query().Get("commit"); p != "" {
		return p
	}
	return "master"
}

func branchParam(r *http.Request) string {
	if p := r.URL.Query().Get("branch"); p != "" {
		return p
	}
	return "master"
}

func materializeParam(r *http.Request) string {
	if _, ok := r.URL.Query()["run"]; ok {
		return "true"
	}
	return "false"
}

func indexOf(haystack []string, needle string) int {
	for i, s := range haystack {
		if s == needle {
			return i
		}
	}
	return -1
}

func cat(w http.ResponseWriter, name string) {
	f, err := btrfs.Open(name)
	if err != nil {
		http.Error(w, err.Error(), 500)
		log.Print(err)
	}
	defer f.Close()

	if _, err := io.Copy(w, f); err != nil {
		http.Error(w, err.Error(), 500)
		log.Print(err)
	}
}

// genericFileHandler serves files from fs. It's used after branch and commit
// info have already been extracted and ignores those aspects of the URL.
func genericFileHandler(fs string, w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like: /foo/bar/.../file/<file>
	fileStart := indexOf(url, "file") + 1
	// file is the path in the filesystem we're getting
	file := path.Join(append([]string{fs}, url[fileStart:]...)...)

	if r.Method == "GET" {
		if strings.Contains(file, "*") {
			if !strings.HasSuffix(file, "*") {
				http.Error(w, "Illegal path containing internal `*`. `*` is currently only allowed as the last character of a path.", 400)
			} else {
				dir := path.Dir(file)
				files, err := btrfs.ReadDir(dir)
				if err != nil {
					http.Error(w, err.Error(), 500)
					return
				}
				for _, fi := range files {
					if fi.IsDir() {
						continue
					} else {
						cat(w, path.Join(dir, fi.Name()))
					}
				}
			}
		} else {
			cat(w, file)
		}
	} else if r.Method == "POST" {
		btrfs.MkdirAll(path.Dir(file))
		size, err := btrfs.CreateFromReader(file, r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created %s, size: %d.\n", file, size)
	} else if r.Method == "PUT" {
		btrfs.MkdirAll(path.Dir(file))
		size, err := btrfs.WriteFile(file, r.Body)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created %s, size: %d.\n", file, size)
	} else if r.Method == "DELETE" {
		if err := btrfs.Remove(file); err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Deleted %s.\n", file)
	}
}

// FileHandler is the core route for modifying the contents of the fileystem.
func FileHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" || r.Method == "DELETE" || r.Method == "PUT" {
		genericFileHandler(path.Join(dataRepo, branchParam(r)), w, r)
	} else if r.Method == "GET" {
		genericFileHandler(path.Join(dataRepo, commitParam(r)), w, r)
	} else {
		http.Error(w, "Invalid method.", 405)
	}
}

// CommitHandler creates a snapshot of outstanding changes.
func CommitHandler(w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, commit, <commit>, file, <file>]
	if len(url) > 3 && url[3] == "file" {
		genericFileHandler(path.Join(dataRepo, url[2]), w, r)
		return
	}
	if r.Method == "GET" {
		commits, err := btrfs.ReadDir(dataRepo)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		for _, ci := range commits {
			if uuid.Parse(ci.Name()) != nil {
				fmt.Fprintf(w, "%s    %s\n", ci.Name(), ci.ModTime().Format("2006-01-02T15:04:05.999999-07:00"))
			}
		}
	} else if r.Method == "POST" {
		err := btrfs.Commit(dataRepo, commitParam(r), branchParam(r))
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		if materializeParam(r) == "true" {
			go func() {
				err := mapreduce.Materialize(dataRepo, branchParam(r), commitParam(r), compRepo, jobDir, shard, modulos)
				if err != nil {
					log.Print(err)
				}
			}()
		}

		fmt.Fprint(w, commitParam(r))
	} else {
		http.Error(w, "Unsupported method.", http.StatusMethodNotAllowed)
		log.Printf("Unsupported method %s in request to %s.", r.Method, r.URL.String())
		return
	}
}

// BranchHandler creates a new branch from commit.
func BranchHandler(w http.ResponseWriter, r *http.Request) {
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, commit, <commit>, file, <file>]
	if len(url) > 3 && url[3] == "file" {
		genericFileHandler(path.Join(dataRepo, url[2]), w, r)
		return
	}
	if r.Method == "GET" {
		branches, err := btrfs.ReadDir(dataRepo)
		if err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}

		for _, bi := range branches {
			if uuid.Parse(bi.Name()) == nil {
				fmt.Fprintf(w, "%s    %s\n", bi.Name(), bi.ModTime().Format("2006-01-02T15:04:05.999999-07:00"))
			}
		}
	} else if r.Method == "POST" {
		if err := btrfs.Branch(dataRepo, commitParam(r), branchParam(r)); err != nil {
			http.Error(w, err.Error(), 500)
			log.Print(err)
			return
		}
		fmt.Fprintf(w, "Created branch. (%s) -> %s.\n", commitParam(r), branchParam(r))
	} else {
		http.Error(w, "Invalid method.", 405)
		log.Print("Invalid method %s.", r.Method)
		return
	}
}

func JobHandler(w http.ResponseWriter, r *http.Request) {
	log.Print("URL in job handler:\n", r.URL)
	url := strings.Split(r.URL.Path, "/")
	// url looks like [, job, <job>, file, <file>]
	if r.Method == "GET" && len(url) > 3 && url[3] == "file" {
		genericFileHandler(path.Join(compRepo, commitParam(r), url[2]), w, r)
		return
	} else if r.Method == "POST" {
		r.URL.Path = path.Join("/file", jobDir, url[2])
		log.Print("URL with reset path:\n", r.URL)
		genericFileHandler(path.Join(dataRepo, branchParam(r)), w, r)
	}
}

// MasterMux creates a multiplexer for a Master writing to the passed in FS.
func MasterMux() *http.ServeMux {
	mux := http.NewServeMux()

	mux.HandleFunc("/branch", BranchHandler)
	mux.HandleFunc("/commit", CommitHandler)
	mux.HandleFunc("/file/", FileHandler)
	mux.HandleFunc("/job/", JobHandler)
	mux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, "pong\n") })

	return mux
}

// RunServer runs a master server listening on port 80
func RunServer() {
	http.ListenAndServe(":80", MasterMux())
}

func main() {
	log.SetFlags(log.Lshortfile)
	parseArgs()
	if err := os.MkdirAll("/var/lib/pfs/log", 0777); err != nil {
		log.Fatal(err)
	}
	logF, err := os.Create(path.Join("/var/lib/pfs/log", logFile))
	if err != nil {
		log.Fatal(err)
	}
	defer logF.Close()
	log.Print("Logging to: ", path.Join("/var/lib/pfs/log", logFile))
	log.SetOutput(logF)
	if err := btrfs.Ensure(dataRepo); err != nil {
		log.Fatal(err)
	}
	if err := btrfs.Ensure(compRepo); err != nil {
		log.Fatal(err)
	}
	log.Print("Listening on port 80...")
	log.Printf("dataRepo: %s, compRepo: %s.", dataRepo, compRepo)
	RunServer()
}
<|endoftext|>


<commit_before>package test

import (
	"context"
	"errors"
	"fmt"
	"net/http/httptest"
	"net/http/httptrace"
	"os"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Clever/wag/samples/gen-go/client"
	"github.com/Clever/wag/samples/gen-go/models"
	"github.com/Clever/wag/samples/gen-go/server"
	"github.com/afex/hystrix-go/hystrix"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type ClientContextTest struct {
	getCount      int
	getTimes      []time.Time
	getErrorCount int
	postCount     int
}

func (c *ClientContextTest) GetBooks(ctx context.Context, input *models.GetBooksInput) ([]models.Book, error) {
	c.getCount++
	c.getTimes = append(c.getTimes, time.Now())
	if c.getCount <= c.getErrorCount {
		return nil, fmt.Errorf("Error count: %d", c.getCount)
	}
	return []models.Book{}, nil
}
func (c *ClientContextTest) GetBookByID(ctx context.Context, input *models.GetBookByIDInput) (*models.Book, error) {
	return nil, nil
}
func (c *ClientContextTest) GetBookByID2(ctx context.Context, id string) (*models.Book, error) {
	return nil, nil
}
func (c *ClientContextTest) CreateBook(ctx context.Context, input *models.Book) (*models.Book, error) {
	c.postCount++
	if c.postCount == 1 {
		return nil, fmt.Errorf("Error count: %d", c.postCount)
	}
	return &models.Book{}, nil
}

func (c *ClientContextTest) HealthCheck(ctx context.Context) error {
	return nil
}

type ClientCircuitTest struct {
	down bool
}

func (c *ClientCircuitTest) GetBooks(ctx context.Context, input *models.GetBooksInput) ([]models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return []models.Book{}, nil
}
func (c *ClientCircuitTest) GetBookByID(ctx context.Context, input *models.GetBookByIDInput) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return nil, nil
}
func (c *ClientCircuitTest) GetBookByID2(ctx context.Context, id string) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return nil, nil
}
func (c *ClientCircuitTest) CreateBook(ctx context.Context, input *models.Book) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return &models.Book{}, nil
}

func (c *ClientCircuitTest) HealthCheck(ctx context.Context) error {
	if c.down {
		return errors.New("fail")
	}
	return nil
}

func TestExponentialClientRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 2}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	c := client.New(testServer.URL)
	c.SetRetryPolicy(client.ExponentialRetryPolicy{})
	_, err := c.GetBooks(context.Background(), &models.GetBooksInput{})
	require.NoError(t, err)
	require.Equal(t, len(controller.getTimes), 3, "expected three requests")
	assert.WithinDuration(t, controller.getTimes[1], controller.getTimes[0].Add(100*time.Millisecond), 10*time.Millisecond,
		"expected first backoff to be about 100ms")
	assert.WithinDuration(t, controller.getTimes[2], controller.getTimes[1].Add(200*time.Millisecond), 20*time.Millisecond,
		"expected first backoff to be about 200ms")
}

func TestCustomClientRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()

	// Should fail if no retries
	c := client.New(testServer.URL)
	c.SetRetryPolicy(client.NoRetryPolicy{})
	_, err := c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.getCount)
}

func TestCustomContextRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()

	// Should fail if no retries
	c := client.New(testServer.URL)
	_, err := c.GetBooks(client.WithRetryPolicy(context.Background(), client.NoRetryPolicy{}), &models.GetBooksInput{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.getCount)
}

func TestNonGetRetries(t *testing.T) {
	controller := ClientContextTest{}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	c := client.New(testServer.URL)
	_, err := c.CreateBook(context.Background(), &models.Book{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.postCount)
}

func TestNetworkErrorRetries(t *testing.T) {
	c := client.New("https://thisshouldnotresolve1234567890.com/")
	_, err := c.CreateBook(context.Background(), &models.Book{})
	assert.Error(t, err)
}

func TestNewWithDiscovery(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)

	// Should be an err if env vars aren't set
	_, err := client.NewFromDiscovery()
	assert.Error(t, err)

	splitURL := strings.Split(testServer.URL, ":")
	assert.Equal(t, 3, len(splitURL))

	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_PROTO", "http")
	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_PORT", splitURL[2])
	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_HOST", splitURL[1][2:])

	c, err := client.NewFromDiscovery()
	assert.NoError(t, err)
	_, err = c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.NoError(t, err)
	assert.Equal(t, 2, controller.getCount)

	// Testing fallback
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_PROTO")
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_PORT")
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_HOST")
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_PROTO", "http")
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_PORT", splitURL[2])
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_HOST", splitURL[1][2:])

	c, err = client.NewFromDiscovery()
	assert.NoError(t, err)
	_, err = c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.NoError(t, err)
	assert.Equal(t, 3, controller.getCount)
}

func TestCircuitBreaker(t *testing.T) {
	controller := ClientCircuitTest{}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	hystrix.Flush()
	c := client.New(testServer.URL)
	c.SetCircuitBreakerDebug(false)
	c.SetCircuitBreakerSettings(client.CircuitBreakerSettings{
		MaxConcurrentRequests:  client.DefaultCircuitBreakerSettings.MaxConcurrentRequests,
		RequestVolumeThreshold: 1,
		SleepWindow:            2000,
		ErrorPercentThreshold:  client.DefaultCircuitBreakerSettings.ErrorPercentThreshold,
	})

	// the circuit should open after one or two failed attempts, since the volume
	// threshold (set above) is 1.
	controller.down = true
	var connAttempts int64
	ctx := httptrace.WithClientTrace(context.Background(),
		&httptrace.ClientTrace{
			GetConn: func(hostPort string) {
				atomic.AddInt64(&connAttempts, 1)
			},
		})

	_, err := c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	_, err = c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	_, err = c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	assert.Equal(t, true, connAttempts <= int64(2), "circuit should have opened, saw too many connection attempts: %d", connAttempts)

	// we should see an attempt go through after two seconds (this is the
	// sleep window configured above).
	circuitOpened := time.Now()
	for _ = range time.Tick(100 * time.Millisecond) {
		_, err := c.CreateBook(ctx, &models.Book{})
		assert.Error(t, err)
		if connAttempts == 2 {
			assert.WithinDuration(t, time.Now(), circuitOpened,
				2*time.Second+500*time.Millisecond)
			break
		}
		if time.Now().Sub(circuitOpened) > 10*time.Second {
			t.Fatal("circuit should let through a 2nd attempt by now")
		}
	}

	// bring the server back up, and we should see successes after another
	// two seconds, for a total of 4 seconds.
	controller.down = false
	for _ = range time.Tick(100 * time.Millisecond) {
		_, err := c.CreateBook(ctx, &models.Book{})
		if err == nil {
			assert.WithinDuration(t, time.Now(), circuitOpened,
				4*time.Second+500*time.Millisecond)
			break
		}
		if time.Now().Sub(circuitOpened) > 10*time.Second {
			t.Fatal("circuit should have closed by now")
		}
	}
}
<commit_msg>Increase duration error margin in test<commit_after>package test

import (
	"context"
	"errors"
	"fmt"
	"net/http/httptest"
	"net/http/httptrace"
	"os"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Clever/wag/samples/gen-go/client"
	"github.com/Clever/wag/samples/gen-go/models"
	"github.com/Clever/wag/samples/gen-go/server"
	"github.com/afex/hystrix-go/hystrix"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type ClientContextTest struct {
	getCount      int
	getTimes      []time.Time
	getErrorCount int
	postCount     int
}

func (c *ClientContextTest) GetBooks(ctx context.Context, input *models.GetBooksInput) ([]models.Book, error) {
	c.getCount++
	c.getTimes = append(c.getTimes, time.Now())
	if c.getCount <= c.getErrorCount {
		return nil, fmt.Errorf("Error count: %d", c.getCount)
	}
	return []models.Book{}, nil
}
func (c *ClientContextTest) GetBookByID(ctx context.Context, input *models.GetBookByIDInput) (*models.Book, error) {
	return nil, nil
}
func (c *ClientContextTest) GetBookByID2(ctx context.Context, id string) (*models.Book, error) {
	return nil, nil
}
func (c *ClientContextTest) CreateBook(ctx context.Context, input *models.Book) (*models.Book, error) {
	c.postCount++
	if c.postCount == 1 {
		return nil, fmt.Errorf("Error count: %d", c.postCount)
	}
	return &models.Book{}, nil
}

func (c *ClientContextTest) HealthCheck(ctx context.Context) error {
	return nil
}

type ClientCircuitTest struct {
	down bool
}

func (c *ClientCircuitTest) GetBooks(ctx context.Context, input *models.GetBooksInput) ([]models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return []models.Book{}, nil
}
func (c *ClientCircuitTest) GetBookByID(ctx context.Context, input *models.GetBookByIDInput) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return nil, nil
}
func (c *ClientCircuitTest) GetBookByID2(ctx context.Context, id string) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return nil, nil
}
func (c *ClientCircuitTest) CreateBook(ctx context.Context, input *models.Book) (*models.Book, error) {
	if c.down {
		return nil, errors.New("fail")
	}
	return &models.Book{}, nil
}

func (c *ClientCircuitTest) HealthCheck(ctx context.Context) error {
	if c.down {
		return errors.New("fail")
	}
	return nil
}

func TestExponentialClientRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 2}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	c := client.New(testServer.URL)
	c.SetRetryPolicy(client.ExponentialRetryPolicy{})
	_, err := c.GetBooks(context.Background(), &models.GetBooksInput{})
	require.NoError(t, err)
	require.Equal(t, len(controller.getTimes), 3, "expected three requests")
	assert.WithinDuration(t, controller.getTimes[1], controller.getTimes[0].Add(100*time.Millisecond), 20*time.Millisecond,
		"expected first backoff to be about 100ms")
	assert.WithinDuration(t, controller.getTimes[2], controller.getTimes[1].Add(200*time.Millisecond), 20*time.Millisecond,
		"expected first backoff to be about 200ms")
}

func TestCustomClientRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()

	// Should fail if no retries
	c := client.New(testServer.URL)
	c.SetRetryPolicy(client.NoRetryPolicy{})
	_, err := c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.getCount)
}

func TestCustomContextRetries(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()

	// Should fail if no retries
	c := client.New(testServer.URL)
	_, err := c.GetBooks(client.WithRetryPolicy(context.Background(), client.NoRetryPolicy{}), &models.GetBooksInput{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.getCount)
}

func TestNonGetRetries(t *testing.T) {
	controller := ClientContextTest{}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	c := client.New(testServer.URL)
	_, err := c.CreateBook(context.Background(), &models.Book{})
	assert.Error(t, err)
	assert.Equal(t, 1, controller.postCount)
}

func TestNetworkErrorRetries(t *testing.T) {
	c := client.New("https://thisshouldnotresolve1234567890.com/")
	_, err := c.CreateBook(context.Background(), &models.Book{})
	assert.Error(t, err)
}

func TestNewWithDiscovery(t *testing.T) {
	controller := ClientContextTest{getErrorCount: 1}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)

	// Should be an err if env vars aren't set
	_, err := client.NewFromDiscovery()
	assert.Error(t, err)

	splitURL := strings.Split(testServer.URL, ":")
	assert.Equal(t, 3, len(splitURL))

	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_PROTO", "http")
	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_PORT", splitURL[2])
	os.Setenv("SERVICE_SWAGGER_TEST_DEFAULT_HOST", splitURL[1][2:])

	c, err := client.NewFromDiscovery()
	assert.NoError(t, err)
	_, err = c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.NoError(t, err)
	assert.Equal(t, 2, controller.getCount)

	// Testing fallback
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_PROTO")
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_PORT")
	os.Unsetenv("SERVICE_SWAGGER_TEST_DEFAULT_HOST")
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_PROTO", "http")
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_PORT", splitURL[2])
	os.Setenv("SERVICE_SWAGGER_TEST_HTTP_HOST", splitURL[1][2:])

	c, err = client.NewFromDiscovery()
	assert.NoError(t, err)
	_, err = c.GetBooks(context.Background(), &models.GetBooksInput{})
	assert.NoError(t, err)
	assert.Equal(t, 3, controller.getCount)
}

func TestCircuitBreaker(t *testing.T) {
	controller := ClientCircuitTest{}
	s := server.New(&controller, "")
	testServer := httptest.NewServer(s.Handler)
	defer testServer.Close()
	hystrix.Flush()
	c := client.New(testServer.URL)
	c.SetCircuitBreakerDebug(false)
	c.SetCircuitBreakerSettings(client.CircuitBreakerSettings{
		MaxConcurrentRequests:  client.DefaultCircuitBreakerSettings.MaxConcurrentRequests,
		RequestVolumeThreshold: 1,
		SleepWindow:            2000,
		ErrorPercentThreshold:  client.DefaultCircuitBreakerSettings.ErrorPercentThreshold,
	})

	// the circuit should open after one or two failed attempts, since the volume
	// threshold (set above) is 1.
	controller.down = true
	var connAttempts int64
	ctx := httptrace.WithClientTrace(context.Background(),
		&httptrace.ClientTrace{
			GetConn: func(hostPort string) {
				atomic.AddInt64(&connAttempts, 1)
			},
		})

	_, err := c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	_, err = c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	_, err = c.CreateBook(ctx, &models.Book{})
	assert.Error(t, err)
	assert.Equal(t, true, connAttempts <= int64(2), "circuit should have opened, saw too many connection attempts: %d", connAttempts)

	// we should see an attempt go through after two seconds (this is the
	// sleep window configured above).
	circuitOpened := time.Now()
	for _ = range time.Tick(100 * time.Millisecond) {
		_, err := c.CreateBook(ctx, &models.Book{})
		assert.Error(t, err)
		if connAttempts == 2 {
			assert.WithinDuration(t, time.Now(), circuitOpened,
				2*time.Second+500*time.Millisecond)
			break
		}
		if time.Now().Sub(circuitOpened) > 10*time.Second {
			t.Fatal("circuit should let through a 2nd attempt by now")
		}
	}

	// bring the server back up, and we should see successes after another
	// two seconds, for a total of 4 seconds.
	controller.down = false
	for _ = range time.Tick(100 * time.Millisecond) {
		_, err := c.CreateBook(ctx, &models.Book{})
		if err == nil {
			assert.WithinDuration(t, time.Now(), circuitOpened,
				4*time.Second+500*time.Millisecond)
			break
		}
		if time.Now().Sub(circuitOpened) > 10*time.Second {
			t.Fatal("circuit should have closed by now")
		}
	}
}
<|endoftext|>