<commit_before>package cmd

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/exoscale/egoscale"
	"github.com/spf13/cobra"
)

// firewallAddCmd represents the add command
var firewallAddCmd = &cobra.Command{
	Use:   "add <security group name | id>  [ssh | telnet | rdp | ...] (default preset rules)",
	Short: "Add rule to a security group",
}

func firewallAddRun(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		firewallAddCmd.Usage()
		return
	}

	securityGroup, err := getSecuGrpWithNameOrID(cs, args[0])
	if err != nil {
		log.Fatal(err)
	}

	desc, err := cmd.Flags().GetString("description")
	if err != nil {
		log.Fatal(err)
	}

	isIpv6, err := cmd.Flags().GetBool("ipv6")
	if err != nil {
		log.Fatal(err)
	}

	rule := &egoscale.AuthorizeSecurityGroupIngress{}
	if len(args) > 1 {
		rule, err = getDefaultRule(args[1], isIpv6)
		if err != nil {
			log.Fatal(err)
		}
	}

	rule.Description = desc
	rule.SecurityGroupID = securityGroup.ID

	isMyIP, err := cmd.Flags().GetBool("my-ip")
	if err != nil {
		log.Fatal(err)
	}

	ip := ""
	if isMyIP {
		cidr, err := getMyCIDR(isIpv6)
		if err != nil {
			log.Fatal(err)
		}
		ip = cidr.String()
	}

	if ip != "" {
		rule.CidrList = []string{ip}
	}

	protocol, err := cmd.Flags().GetString("protocol")
	if err != nil {
		log.Fatal(err)
	}

	if protocol != "" {
		rule.Protocol = strings.ToLower(protocol)
	}

	cidr, err := cmd.Flags().GetString("cidr")
	if err != nil {
		log.Fatal(err)
	}

	if cidr != "" {
		cidrs := getCommaflag(cidr)

		for _, c := range cidrs {
			rule.CidrList = append(rule.CidrList, c)
		}
	}

	sg, err := cmd.Flags().GetString("security-group")
	if err != nil {
		log.Fatal(err)
	}

	if sg != "" {
		sgs := getCommaflag(sg)

		userSecurityGroups, err := getSGs(cs, sgs)
		if err != nil {
			log.Fatal(err)
		}

		rule.UserSecurityGroupList = userSecurityGroups
	}

	icmptype, err := getUint8CustomFlag(cmd, "icmp-type")
	if err != nil {
		log.Fatal(err)
	}
	icmpcode, err := getUint8CustomFlag(cmd, "icmp-code")
	if err != nil {
		log.Fatal(err)
	}

	if icmptype.uint8 != nil {
		rule.IcmpType = *icmptype.uint8
	}

	if icmpcode.uint8 != nil {
		rule.IcmpCode = *icmpcode.uint8
	}

	isEgress, err := cmd.Flags().GetBool("engress")
	if err != nil {
		log.Fatal(err)
	}

	port, err := cmd.Flags().GetString("port")
	if err != nil {
		log.Fatal(err)
	}

	//Not best practis but waiting to find better solution
	if port != "" && (rule.Protocol == "tcp" || rule.Protocol == "udp") {

		ports := getCommaflag(port)
		portsRange, err := getPortsRange(ports)
		if err != nil {
			log.Fatal(err)
		}

		for _, portRange := range portsRange {

			rule.StartPort = portRange.start
			rule.EndPort = portRange.end
			if err := addRule(rule, isEgress); err != nil {
				log.Fatal(err)
			}
		}
	}
	//Not best practis but waiting to find better solution
	if port == "" || !(rule.Protocol == "tcp" || rule.Protocol == "udp") {
		if err := addRule(rule, isEgress); err != nil {
			log.Fatal(err)
		}
	}

	firewallDetails(rule.SecurityGroupID)
}

type portRange struct {
	start uint16
	end   uint16
}

func getPortsRange(ports []string) ([]portRange, error) {
	portsRange := make([]portRange, len(ports))
	for i, p := range ports {
		pRange := strings.Split(p, "-")
		if len(pRange) > 2 || len(pRange) == 0 {
			return nil, fmt.Errorf("failed to get ports")
		}
		p1, err := strconv.ParseUint(pRange[0], 10, 16)
		if err != nil {
			return nil, err
		}
		portsRange[i].start = uint16(p1)
		portsRange[i].end = uint16(p1)
		if len(pRange) == 2 {
			p2, err := strconv.ParseUint(pRange[0], 10, 16)
			if err != nil {
				return nil, err
			}
			portsRange[i].end = uint16(p2)
		}
	}
	return portsRange, nil
}

func getSGs(cs *egoscale.Client, sgs []string) ([]egoscale.UserSecurityGroup, error) {

	sgRes := []egoscale.UserSecurityGroup{}

	for _, sg := range sgs {

		s, err := getSecuGrpWithNameOrID(cs, sg)
		if err != nil {
			return nil, err
		}
		sgRes = append(sgRes, egoscale.UserSecurityGroup{Group: s.Name, Account: s.Account})
	}
	return sgRes, nil
}

func getDefaultRule(ruleName string, isIpv6 bool) (*egoscale.AuthorizeSecurityGroupIngress, error) {

	cidr := defaultCidr
	if isIpv6 {
		cidr = defaultCidr6
	}

	ruleName = strings.ToLower(ruleName)
	if ruleName == "ping" {
		return &egoscale.AuthorizeSecurityGroupIngress{
			Protocol:    "icmp",
			CidrList:    []string{cidr},
			IcmpType:    8,
			IcmpCode:    0,
			Description: "",
		}, nil
	}

	for d := Daytime; d <= Minecraft; d++ {
		if strings.ToLower(port.String(d)) == ruleName {
			return &egoscale.AuthorizeSecurityGroupIngress{
				Protocol:    "tcp",
				CidrList:    []string{cidr},
				StartPort:   uint16(d),
				EndPort:     uint16(d),
				Description: fmt.Sprintf(""),
			}, nil
		}
	}

	return nil, fmt.Errorf("default rule %q not found", ruleName)
}

func addRule(rule *egoscale.AuthorizeSecurityGroupIngress, isEgress bool) error {
	var err error
	if isEgress {
		_, err = cs.Request((*egoscale.AuthorizeSecurityGroupEgress)(rule))
	} else {
		_, err = cs.Request(rule)
	}

	if err != nil {
		return err
	}

	return nil
}

func init() {
	firewallAddCmd.Run = firewallAddRun

	firewallAddCmd.Flags().BoolP("ipv6", "6", false, "Set ipv6 on default rules or on --my-ip")
	firewallAddCmd.Flags().BoolP("my-ip", "", false, "Set CIDR for my ip")
	firewallAddCmd.Flags().BoolP("engress", "e", false, "By default rule is INGRESS (set --engress to have ENGRESS rule)")
	firewallAddCmd.Flags().StringP("protocol", "p", "", "Rule Protocol available [tcp, udp, icmp, icmpv6, ah, esp, gre]")
	firewallAddCmd.Flags().StringP("cidr", "c", "", "Rule Cidr [CIDR 0.0.0.0/0,::/0,...]")
	firewallAddCmd.Flags().StringP("security-group", "s", "", "Rule security group [name or id ex: sg1,sg2...]")
	firewallAddCmd.Flags().StringP("port", "P", "", "Rule port range [80-80,443,22-22]")

	//Flag for icmp
	icmpTypeVarP := new(uint8PtrValue)
	icmpCodeVarP := new(uint8PtrValue)

	firewallAddCmd.Flags().VarP(icmpTypeVarP, "icmp-type", "", "Set icmp type")
	firewallAddCmd.Flags().VarP(icmpCodeVarP, "icmp-code", "", "Set icmp type code")

	firewallAddCmd.Flags().StringP("description", "d", "", "Rule description")
	firewallCmd.AddCommand(firewallAddCmd)
}

type uint8PtrValue struct {
	*uint8
}

func (v *uint8PtrValue) Set(val string) error {
	r, err := strconv.ParseUint(val, 10, 8)
	if err != nil {
		return err
	}
	res := uint8(r)
	v.uint8 = &res
	return nil
}

func (v *uint8PtrValue) Type() string {
	return "uint8"
}

func (v *uint8PtrValue) String() string {

	if v.uint8 == nil {
		return "nil"
	}
	return strconv.FormatUint(uint64(*v.uint8), 10)
}

func getUint8CustomFlag(cmd *cobra.Command, name string) (uint8PtrValue, error) {
	it := cmd.Flags().Lookup(name)
	if it != nil {
		r := it.Value.(*uint8PtrValue)
		if r != nil {
			return *r, nil
		}
	}
	return uint8PtrValue{}, fmt.Errorf("Unable to get flag %s", name)
}
<commit_msg>Exo: firewall add X http https (#188)<commit_after>package cmd

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/exoscale/egoscale"
	"github.com/spf13/cobra"
)

// firewallAddCmd represents the add command
var firewallAddCmd = &cobra.Command{
	Use:   "add <security group name | id>  [ssh | telnet | rdp | ...] (default preset rules)",
	Short: "Add rule to a security group",
}

func firewallAddRun(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		firewallAddCmd.Usage()
		return
	}

	securityGroup, err := getSecuGrpWithNameOrID(cs, args[0])
	if err != nil {
		log.Fatal(err)
	}

	desc, err := cmd.Flags().GetString("description")
	if err != nil {
		log.Fatal(err)
	}

	isIpv6, err := cmd.Flags().GetBool("ipv6")
	if err != nil {
		log.Fatal(err)
	}

	rule := &egoscale.AuthorizeSecurityGroupIngress{}

	for i := 1; true; i++ {

		if i >= len(args) {
			break
		}

		if len(args) > 1 {
			rule, err = getDefaultRule(args[i], isIpv6)
			if err != nil {
				log.Fatal(err)
			}
		}

		rule.Description = desc
		rule.SecurityGroupID = securityGroup.ID

		isMyIP, err := cmd.Flags().GetBool("my-ip")
		if err != nil {
			log.Fatal(err)
		}

		ip := ""
		if isMyIP {
			cidr, err := getMyCIDR(isIpv6)
			if err != nil {
				log.Fatal(err)
			}
			ip = cidr.String()
		}

		if ip != "" {
			rule.CidrList = []string{ip}
		}

		protocol, err := cmd.Flags().GetString("protocol")
		if err != nil {
			log.Fatal(err)
		}

		if protocol != "" {
			rule.Protocol = strings.ToLower(protocol)
		}

		cidr, err := cmd.Flags().GetString("cidr")
		if err != nil {
			log.Fatal(err)
		}

		if cidr != "" {
			cidrs := getCommaflag(cidr)

			for _, c := range cidrs {
				rule.CidrList = append(rule.CidrList, c)
			}
		}

		sg, err := cmd.Flags().GetString("security-group")
		if err != nil {
			log.Fatal(err)
		}

		if sg != "" {
			sgs := getCommaflag(sg)

			userSecurityGroups, err := getSGs(cs, sgs)
			if err != nil {
				log.Fatal(err)
			}

			rule.UserSecurityGroupList = userSecurityGroups
		}

		icmptype, err := getUint8CustomFlag(cmd, "icmp-type")
		if err != nil {
			log.Fatal(err)
		}
		icmpcode, err := getUint8CustomFlag(cmd, "icmp-code")
		if err != nil {
			log.Fatal(err)
		}

		if icmptype.uint8 != nil {
			rule.IcmpType = *icmptype.uint8
		}

		if icmpcode.uint8 != nil {
			rule.IcmpCode = *icmpcode.uint8
		}

		isEgress, err := cmd.Flags().GetBool("engress")
		if err != nil {
			log.Fatal(err)
		}

		port, err := cmd.Flags().GetString("port")
		if err != nil {
			log.Fatal(err)
		}

		//Not best practis but waiting to find better solution
		if port != "" && (rule.Protocol == "tcp" || rule.Protocol == "udp") {

			ports := getCommaflag(port)
			portsRange, err := getPortsRange(ports)
			if err != nil {
				log.Fatal(err)
			}

			for _, portRange := range portsRange {

				rule.StartPort = portRange.start
				rule.EndPort = portRange.end
				if err := addRule(rule, isEgress); err != nil {
					log.Fatal(err)
				}
			}
		}
		//Not best practis but waiting to find better solution
		if port == "" || !(rule.Protocol == "tcp" || rule.Protocol == "udp") {
			if err := addRule(rule, isEgress); err != nil {
				log.Fatal(err)
			}
		}

		if len(args) == 1 {
			break
		}

	}

	firewallDetails(rule.SecurityGroupID)
}

type portRange struct {
	start uint16
	end   uint16
}

func getPortsRange(ports []string) ([]portRange, error) {
	portsRange := make([]portRange, len(ports))
	for i, p := range ports {
		pRange := strings.Split(p, "-")
		if len(pRange) > 2 || len(pRange) == 0 {
			return nil, fmt.Errorf("failed to get ports")
		}
		p1, err := strconv.ParseUint(pRange[0], 10, 16)
		if err != nil {
			return nil, err
		}
		portsRange[i].start = uint16(p1)
		portsRange[i].end = uint16(p1)
		if len(pRange) == 2 {
			p2, err := strconv.ParseUint(pRange[0], 10, 16)
			if err != nil {
				return nil, err
			}
			portsRange[i].end = uint16(p2)
		}
	}
	return portsRange, nil
}

func getSGs(cs *egoscale.Client, sgs []string) ([]egoscale.UserSecurityGroup, error) {

	sgRes := []egoscale.UserSecurityGroup{}

	for _, sg := range sgs {

		s, err := getSecuGrpWithNameOrID(cs, sg)
		if err != nil {
			return nil, err
		}
		sgRes = append(sgRes, egoscale.UserSecurityGroup{Group: s.Name, Account: s.Account})
	}
	return sgRes, nil
}

func getDefaultRule(ruleName string, isIpv6 bool) (*egoscale.AuthorizeSecurityGroupIngress, error) {

	icmpType := uint8(8)
	cidr := defaultCidr
	if isIpv6 {
		cidr = defaultCidr6
		icmpType = uint8(128)
	}

	ruleName = strings.ToLower(ruleName)
	if ruleName == "ping" {
		return &egoscale.AuthorizeSecurityGroupIngress{
			Protocol:    "icmp",
			CidrList:    []string{cidr},
			IcmpType:    icmpType,
			IcmpCode:    0,
			Description: "",
		}, nil
	}

	for d := Daytime; d <= Minecraft; d++ {
		if strings.ToLower(port.String(d)) == ruleName {
			return &egoscale.AuthorizeSecurityGroupIngress{
				Protocol:    "tcp",
				CidrList:    []string{cidr},
				StartPort:   uint16(d),
				EndPort:     uint16(d),
				Description: fmt.Sprintf(""),
			}, nil
		}
	}

	return nil, fmt.Errorf("default rule %q not found", ruleName)
}

func addRule(rule *egoscale.AuthorizeSecurityGroupIngress, isEgress bool) error {
	var err error
	if isEgress {
		_, err = cs.Request((*egoscale.AuthorizeSecurityGroupEgress)(rule))
	} else {
		_, err = cs.Request(rule)
	}

	if err != nil {
		return err
	}

	return nil
}

func init() {
	firewallAddCmd.Run = firewallAddRun

	firewallAddCmd.Flags().BoolP("ipv6", "6", false, "Set ipv6 on default rules or on --my-ip")
	firewallAddCmd.Flags().BoolP("my-ip", "", false, "Set CIDR for my ip")
	firewallAddCmd.Flags().BoolP("engress", "e", false, "By default rule is INGRESS (set --engress to have ENGRESS rule)")
	firewallAddCmd.Flags().StringP("protocol", "p", "", "Rule Protocol available [tcp, udp, icmp, icmpv6, ah, esp, gre]")
	firewallAddCmd.Flags().StringP("cidr", "c", "", "Rule Cidr [CIDR 0.0.0.0/0,::/0,...]")
	firewallAddCmd.Flags().StringP("security-group", "s", "", "Rule security group [name or id ex: sg1,sg2...]")
	firewallAddCmd.Flags().StringP("port", "P", "", "Rule port range [80-80,443,22-22]")

	//Flag for icmp
	icmpTypeVarP := new(uint8PtrValue)
	icmpCodeVarP := new(uint8PtrValue)

	firewallAddCmd.Flags().VarP(icmpTypeVarP, "icmp-type", "", "Set icmp type")
	firewallAddCmd.Flags().VarP(icmpCodeVarP, "icmp-code", "", "Set icmp type code")

	firewallAddCmd.Flags().StringP("description", "d", "", "Rule description")
	firewallCmd.AddCommand(firewallAddCmd)
}

type uint8PtrValue struct {
	*uint8
}

func (v *uint8PtrValue) Set(val string) error {
	r, err := strconv.ParseUint(val, 10, 8)
	if err != nil {
		return err
	}
	res := uint8(r)
	v.uint8 = &res
	return nil
}

func (v *uint8PtrValue) Type() string {
	return "uint8"
}

func (v *uint8PtrValue) String() string {

	if v.uint8 == nil {
		return "nil"
	}
	return strconv.FormatUint(uint64(*v.uint8), 10)
}

func getUint8CustomFlag(cmd *cobra.Command, name string) (uint8PtrValue, error) {
	it := cmd.Flags().Lookup(name)
	if it != nil {
		r := it.Value.(*uint8PtrValue)
		if r != nil {
			return *r, nil
		}
	}
	return uint8PtrValue{}, fmt.Errorf("Unable to get flag %s", name)
}
<|endoftext|>


<commit_before>package cmd

// DO NOT EDIT manually! Generated by hack/update-version.sh
var version = "0.9.0-dev"
<commit_msg>Bump version to 0.10.0<commit_after>package cmd

// DO NOT EDIT manually! Generated by hack/update-version.sh
var version = "0.10.0"
<|endoftext|>


<commit_before>package gateway

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/funkygao/gafka/cmd/kateway/manager"
	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/gafka/cmd/kateway/store"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
)

func isBrokerError(err error) bool {
	if err != store.ErrTooManyConsumers && err != store.ErrRebalancing {
		return true
	}

	return false
}

func getHttpQueryInt(query *url.Values, key string, defaultVal int) (int, error) {
	valStr := query.Get(key)
	if valStr == "" {
		return defaultVal, nil
	}

	return strconv.Atoi(valStr)
}

// getHttpRemoteIp returns ip only, without remote port.
func getHttpRemoteIp(r *http.Request) string {
	forwardFor := r.Header.Get(HttpHeaderXForwardedFor) // client_ip,proxy_ip,proxy_ip,...
	if forwardFor == "" {
		// not behind haproxy, directly connected
		p := strings.SplitN(r.RemoteAddr, ":", 2)
		return p[0]
	}

	return forwardFor // FIXME forwardFor might be comma seperated ip list, but here for performance ignore it
}

func checkUlimit(min int) {
	ulimitN, err := exec.Command("/bin/sh", "-c", "ulimit -n").Output()
	if err != nil {
		panic(err)
	}

	n, err := strconv.Atoi(strings.TrimSpace(string(ulimitN)))
	if err != nil || n < min {
		log.Panicf("ulimit too small: %d, should be at least %d", n, min)
	}
}

func extractFromMetricsName(name string) (appid, topic, ver, realname string) {
	if name[0] != '{' {
		realname = name
		return
	}

	i := strings.Index(name, "}")
	realname = name[i+1:]
	p := strings.SplitN(name[1:i], ".", 3)
	appid, topic, ver = p[0], p[1], p[2]
	return
}

func updateCounter(appid, topic, ver, name string, n int64,
	mu *sync.RWMutex, m map[string]metrics.Counter) {
	tagBuf := make([]byte, 4+len(appid)+len(topic)+len(ver))
	tagBuf[0] = CharBraceletLeft
	idx := 1
	for ; idx <= len(appid); idx++ {
		tagBuf[idx] = appid[idx-1]
	}
	tagBuf[idx] = CharDot
	idx++
	for j := 0; j < len(topic); j++ {
		tagBuf[idx+j] = topic[j]
	}
	idx += len(topic)
	tagBuf[idx] = CharDot
	idx++
	for j := 0; j < len(ver); j++ {
		tagBuf[idx+j] = ver[j]
	}
	idx += len(ver)
	tagBuf[idx] = CharBraceletRight

	mu.RLock()
	// golang has optimization avoids extra allocations when []byte keys are used to
	// lookup entries in map[string] collections: m[string(key)]
	counter, present := m[string(tagBuf)]
	mu.RUnlock()

	if present {
		counter.Inc(1)
		return
	}

	// seldom goes here, needn't optimize

	tag := string(tagBuf)
	mu.Lock()
	m[tag] = metrics.NewRegisteredCounter(tag+name, metrics.DefaultRegistry)
	mu.Unlock()

	m[tag].Inc(n)
}

type SubStatus struct {
	Appid          string `json:"appid,omitempty"`
	Group          string `json:"group"`
	Partition      string `json:"partition"`
	ProducedOldest int64  `json:"pold"`
	ProducedNewest int64  `json:"pubd"`
	Consumed       int64  `json:"subd"`
}

func topicSubStatus(cluster string, myAppid, hisAppid, topic, ver string,
	group string, onlyMine bool) ([]SubStatus, error) {
	zkcluster := meta.Default.ZkCluster(cluster)
	if group != "" {
		// if group is empty, find all groups
		group = myAppid + "." + group
	}
	rawTopic := manager.Default.KafkaTopic(hisAppid, topic, ver)
	consumersByGroup, err := zkcluster.ConsumerGroupsOfTopic(rawTopic)
	if err != nil {
		return nil, err
	}
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for grp, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, grp)
	}
	sort.Strings(sortedGroups)

	out := make([]SubStatus, 0, len(sortedGroups))
	for _, grp := range sortedGroups {
		if group != "" && grp != group {
			continue
		}

		sortedTopicAndPartitionIds := make([]string, 0, len(consumersByGroup[grp]))
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[grp] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]
			if rawTopic != consumer.Topic {
				continue
			}

			p := strings.SplitN(grp, ".", 2) // grp is like 'appid.groupname'
			if onlyMine && myAppid != p[0] {
				// this group does not belong to me
				continue
			}

			stat := SubStatus{
				Group:          p[1],
				Partition:      consumer.PartitionId,
				ProducedOldest: consumer.OldestOffset,
				ProducedNewest: consumer.ProducerOffset,
				Consumed:       consumer.ConsumerOffset,
			}
			if !onlyMine {
				stat.Appid = p[0]
			}

			out = append(out, stat)
		}
	}

	return out, nil
}
<commit_msg>wait on sync.WaitGroup can have timeout mechanism<commit_after>package gateway

import (
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/manager"
	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/gafka/cmd/kateway/store"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
)

func isBrokerError(err error) bool {
	if err != store.ErrTooManyConsumers && err != store.ErrRebalancing {
		return true
	}

	return false
}

// waitTimeout waits for the waitgroup for the specified max timeout.
// Returns true if waiting timed out.
func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
	c := make(chan struct{})
	go func() {
		defer close(c)
		wg.Wait()
	}()

	select {
	case <-c:
		return false // completed normally

	case <-time.After(timeout):
		return true // timed out
	}
}

func getHttpQueryInt(query *url.Values, key string, defaultVal int) (int, error) {
	valStr := query.Get(key)
	if valStr == "" {
		return defaultVal, nil
	}

	return strconv.Atoi(valStr)
}

// getHttpRemoteIp returns ip only, without remote port.
func getHttpRemoteIp(r *http.Request) string {
	forwardFor := r.Header.Get(HttpHeaderXForwardedFor) // client_ip,proxy_ip,proxy_ip,...
	if forwardFor == "" {
		// not behind haproxy, directly connected
		p := strings.SplitN(r.RemoteAddr, ":", 2)
		return p[0]
	}

	return forwardFor // FIXME forwardFor might be comma seperated ip list, but here for performance ignore it
}

func checkUlimit(min int) {
	ulimitN, err := exec.Command("/bin/sh", "-c", "ulimit -n").Output()
	if err != nil {
		panic(err)
	}

	n, err := strconv.Atoi(strings.TrimSpace(string(ulimitN)))
	if err != nil || n < min {
		log.Panicf("ulimit too small: %d, should be at least %d", n, min)
	}
}

func extractFromMetricsName(name string) (appid, topic, ver, realname string) {
	if name[0] != '{' {
		realname = name
		return
	}

	i := strings.Index(name, "}")
	realname = name[i+1:]
	p := strings.SplitN(name[1:i], ".", 3)
	appid, topic, ver = p[0], p[1], p[2]
	return
}

func updateCounter(appid, topic, ver, name string, n int64,
	mu *sync.RWMutex, m map[string]metrics.Counter) {
	tagBuf := make([]byte, 4+len(appid)+len(topic)+len(ver))
	tagBuf[0] = CharBraceletLeft
	idx := 1
	for ; idx <= len(appid); idx++ {
		tagBuf[idx] = appid[idx-1]
	}
	tagBuf[idx] = CharDot
	idx++
	for j := 0; j < len(topic); j++ {
		tagBuf[idx+j] = topic[j]
	}
	idx += len(topic)
	tagBuf[idx] = CharDot
	idx++
	for j := 0; j < len(ver); j++ {
		tagBuf[idx+j] = ver[j]
	}
	idx += len(ver)
	tagBuf[idx] = CharBraceletRight

	mu.RLock()
	// golang has optimization avoids extra allocations when []byte keys are used to
	// lookup entries in map[string] collections: m[string(key)]
	counter, present := m[string(tagBuf)]
	mu.RUnlock()

	if present {
		counter.Inc(1)
		return
	}

	// seldom goes here, needn't optimize

	tag := string(tagBuf)
	mu.Lock()
	m[tag] = metrics.NewRegisteredCounter(tag+name, metrics.DefaultRegistry)
	mu.Unlock()

	m[tag].Inc(n)
}

type SubStatus struct {
	Appid          string `json:"appid,omitempty"`
	Group          string `json:"group"`
	Partition      string `json:"partition"`
	ProducedOldest int64  `json:"pold"`
	ProducedNewest int64  `json:"pubd"`
	Consumed       int64  `json:"subd"`
}

func topicSubStatus(cluster string, myAppid, hisAppid, topic, ver string,
	group string, onlyMine bool) ([]SubStatus, error) {
	zkcluster := meta.Default.ZkCluster(cluster)
	if group != "" {
		// if group is empty, find all groups
		group = myAppid + "." + group
	}
	rawTopic := manager.Default.KafkaTopic(hisAppid, topic, ver)
	consumersByGroup, err := zkcluster.ConsumerGroupsOfTopic(rawTopic)
	if err != nil {
		return nil, err
	}
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for grp, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, grp)
	}
	sort.Strings(sortedGroups)

	out := make([]SubStatus, 0, len(sortedGroups))
	for _, grp := range sortedGroups {
		if group != "" && grp != group {
			continue
		}

		sortedTopicAndPartitionIds := make([]string, 0, len(consumersByGroup[grp]))
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[grp] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]
			if rawTopic != consumer.Topic {
				continue
			}

			p := strings.SplitN(grp, ".", 2) // grp is like 'appid.groupname'
			if onlyMine && myAppid != p[0] {
				// this group does not belong to me
				continue
			}

			stat := SubStatus{
				Group:          p[1],
				Partition:      consumer.PartitionId,
				ProducedOldest: consumer.OldestOffset,
				ProducedNewest: consumer.ProducerOffset,
				Consumed:       consumer.ConsumerOffset,
			}
			if !onlyMine {
				stat.Appid = p[0]
			}

			out = append(out, stat)
		}
	}

	return out, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2019 by Michael Dvorkin and contributors. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/nsf/termbox-go"
)

// Screen is thin wrapper around Termbox library to provide basic display
// capabilities as required by Mop.
type Screen struct {
	width      int        // Current number of columns.
	height     int        // Current number of rows.
	cleared    bool       // True after the screens gets cleared.
	layout     *Layout    // Pointer to layout (gets created by screen).
	markup     *Markup    // Pointer to markup processor (gets created by screen).
	pausedAt   *time.Time // Timestamp of the pause request or nil if none.
	offset     int        // Offset for scolling
	headerLine int        // Line number of header for scroll feature
	max        int        // highest offset
}

// Initializes Termbox, creates screen along with layout and markup, and
// calculates current screen dimensions. Once initialized the screen is
// ready for display.
func NewScreen(profile *Profile) *Screen {
	if err := termbox.Init(); err != nil {
		panic(err)
	}
	screen := &Screen{}
	screen.layout = NewLayout()
	screen.markup = NewMarkup(profile)
	screen.offset = 0

	return screen.Resize()
}

// Close gets called upon program termination to close the Termbox.
func (screen *Screen) Close() *Screen {
	termbox.Close()

	return screen
}

// Resize gets called when the screen is being resized. It recalculates screen
// dimensions and requests to clear the screen on next update.
func (screen *Screen) Resize() *Screen {
	screen.width, screen.height = termbox.Size()
	screen.cleared = false

	return screen
}

// Pause is a toggle function that either creates a timestamp of the pause
// request or resets it to nil.
func (screen *Screen) Pause(pause bool) *Screen {
	if pause {
		screen.pausedAt = new(time.Time)
		*screen.pausedAt = time.Now()
	} else {
		screen.pausedAt = nil
	}

	return screen
}

// Clear makes the entire screen blank using default background color.
func (screen *Screen) Clear() *Screen {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	screen.cleared = true

	return screen
}

// ClearLine erases the contents of the line starting from (x,y) coordinate
// till the end of the line.
func (screen *Screen) ClearLine(x int, y int) *Screen {
	for i := x; i < screen.width; i++ {
		termbox.SetCell(i, y, ' ', termbox.ColorDefault, termbox.ColorDefault)
	}
	termbox.Flush()

	return screen
}

// Increase the offset for scrolling feature by n
// Takes number of tickers as max, so not scrolling down forever
func (screen *Screen) IncreaseOffset(n int) {
	if screen.offset+n <= screen.max {
		screen.offset += n
	} else if screen.max > screen.height {
		screen.offset = screen.max
	}
}

// Decrease the offset for scrolling feature by n
func (screen *Screen) DecreaseOffset(n int) {
	if screen.offset > n {
		screen.offset -= n
	} else {
		screen.offset = 0
	}
}

func (screen *Screen) ScrollTop() {
	screen.offset = 0
}

func (screen *Screen) ScrollBottom() {
	if screen.max > screen.height {
		screen.offset = screen.max
	}
}

func (screen *Screen) DrawOldQuotes(quotes *Quotes) {
	screen.draw(screen.layout.Quotes(quotes), true)
	termbox.Flush()
}

func (screen *Screen) DrawOldMarket(market *Market) {
	screen.draw(screen.layout.Market(market), false)
	termbox.Flush()
}

// Draw accepts variable number of arguments and knows how to display the
// market data, stock quotes, current time, and an arbitrary string.
func (screen *Screen) Draw(objects ...interface{}) *Screen {
	zonename, _ := time.Now().In(time.Local).Zone()
	if screen.pausedAt != nil {
		defer screen.DrawLine(0, 0, `<right><r>`+screen.pausedAt.Format(`3:04:05pm `+zonename)+`</r></right>`)
	}
	for _, ptr := range objects {
		switch ptr.(type) {
		case *Market:
			object := ptr.(*Market)
			screen.draw(screen.layout.Market(object.Fetch()), false)
		case *Quotes:
			object := ptr.(*Quotes)
			screen.draw(screen.layout.Quotes(object.Fetch()), true)
		case time.Time:
			timestamp := ptr.(time.Time).Format(`3:04:05pm ` + zonename)
			screen.DrawLine(0, 0, `<right><time>`+timestamp+`</></right>`)
		default:
			screen.draw(ptr.(string), false)
		}
	}

	termbox.Flush()

	return screen
}

// DrawLine takes the incoming string, tokenizes it to extract markup
// elements, and displays it all starting at (x,y) location.

// DrawLineFlush gives the option to flush screen after drawing

// wrapper for DrawLineFlush
func (screen *Screen) DrawLine(x int, y int, str string) {
	screen.DrawLineFlush(x, y, str, true)
}

func (screen *Screen) DrawLineFlush(x int, y int, str string, flush bool) {
	start, column := 0, 0

	for _, token := range screen.markup.Tokenize(str) {
		// First check if it's a tag. Tags are eaten up and not displayed.
		if screen.markup.IsTag(token) {
			continue
		}

		// Here comes the actual text: display it one character at a time.
		for i, char := range token {
			if !screen.markup.RightAligned {
				start = x + column
				column++
			} else {
				start = screen.width - len(token) + i
			}
			termbox.SetCell(start, y, char, screen.markup.Foreground, screen.markup.Background)
		}
	}
	if flush {
		termbox.Flush()
	}
}

// Underlying workhorse function that takes multiline string, splits it into
// lines, and displays them row by row.
func (screen *Screen) draw(str string, offset bool) {
	if !screen.cleared {
		screen.Clear()
	}
	var allLines []string
	drewHeading := false

	screen.width, screen.height = termbox.Size()

	tempFormat := "%" + strconv.Itoa(screen.width) + "s"
	blankLine := fmt.Sprintf(tempFormat, "")
	allLines = strings.Split(str, "\n")

	if offset {
		screen.max = len(allLines) - screen.height
	}

	// Write the lines being updated.
	for row := 0; row < len(allLines); row++ {
		if offset {
			// Did we draw the underlined heading row?  This is a crude
			// check, but--see comments below...
			// --- Heading row only appears for quotes, so offset is true
			if !drewHeading {
				if strings.Contains(allLines[row], "Ticker") &&
					strings.Contains(allLines[row], "Last") &&
					strings.Contains(allLines[row], "Change") {
					drewHeading = true
					screen.headerLine = row
					screen.DrawLine(0, row, allLines[row])
					// move on to the point to offset to
					row += screen.offset
				}
			} else {
				// only write the necessary lines
				if row <= len(allLines) &&
					row > screen.headerLine {
					screen.DrawLineFlush(0, row-screen.offset, allLines[row], false)
				} else if row > len(allLines) {
					row = len(allLines)
				}
			}
		} else {
			screen.DrawLineFlush(0, row, allLines[row], false)
		}
	}
	// If the quotes lines in this cycle are shorter than in the previous
	// cycles, e.g., because a filter was just applied, then one or more
	// lines from the previous cycles will not be cleared.  Since the
	// incoming lines don't mark explicitly whether they are part of the
	// market summary or quotes, we can't check whether quotes were updated
	// in a way that is robust for code changes.  This is a simple test: if
	// we drew the heading row ("Ticker Last Change..."), then we are
	// updating the quotes section in this cycle, and we should pad the
	// quotes section with blank lines.  If we didn't draw the heading row,
	// then we probably only updated the market summary at the top in this
	// cycle.  In that case, padding with blank lines would overwrite the
	// stocks list.)
	if drewHeading {
		for i := len(allLines) - 1 - screen.offset; i < screen.height; i++ {
			if i > screen.headerLine {
				screen.DrawLine(0, i, blankLine)
			}
		}
	}
}
<commit_msg>Write screen.height + 1 lines<commit_after>// Copyright (c) 2013-2019 by Michael Dvorkin and contributors. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/nsf/termbox-go"
)

// Screen is thin wrapper around Termbox library to provide basic display
// capabilities as required by Mop.
type Screen struct {
	width      int        // Current number of columns.
	height     int        // Current number of rows.
	cleared    bool       // True after the screens gets cleared.
	layout     *Layout    // Pointer to layout (gets created by screen).
	markup     *Markup    // Pointer to markup processor (gets created by screen).
	pausedAt   *time.Time // Timestamp of the pause request or nil if none.
	offset     int        // Offset for scolling
	headerLine int        // Line number of header for scroll feature
	max        int        // highest offset
}

// Initializes Termbox, creates screen along with layout and markup, and
// calculates current screen dimensions. Once initialized the screen is
// ready for display.
func NewScreen(profile *Profile) *Screen {
	if err := termbox.Init(); err != nil {
		panic(err)
	}
	screen := &Screen{}
	screen.layout = NewLayout()
	screen.markup = NewMarkup(profile)
	screen.offset = 0

	return screen.Resize()
}

// Close gets called upon program termination to close the Termbox.
func (screen *Screen) Close() *Screen {
	termbox.Close()

	return screen
}

// Resize gets called when the screen is being resized. It recalculates screen
// dimensions and requests to clear the screen on next update.
func (screen *Screen) Resize() *Screen {
	screen.width, screen.height = termbox.Size()
	screen.cleared = false

	return screen
}

// Pause is a toggle function that either creates a timestamp of the pause
// request or resets it to nil.
func (screen *Screen) Pause(pause bool) *Screen {
	if pause {
		screen.pausedAt = new(time.Time)
		*screen.pausedAt = time.Now()
	} else {
		screen.pausedAt = nil
	}

	return screen
}

// Clear makes the entire screen blank using default background color.
func (screen *Screen) Clear() *Screen {
	termbox.Clear(termbox.ColorDefault, termbox.ColorDefault)
	screen.cleared = true

	return screen
}

// ClearLine erases the contents of the line starting from (x,y) coordinate
// till the end of the line.
func (screen *Screen) ClearLine(x int, y int) *Screen {
	for i := x; i < screen.width; i++ {
		termbox.SetCell(i, y, ' ', termbox.ColorDefault, termbox.ColorDefault)
	}
	termbox.Flush()

	return screen
}

// Increase the offset for scrolling feature by n
// Takes number of tickers as max, so not scrolling down forever
func (screen *Screen) IncreaseOffset(n int) {
	if screen.offset+n <= screen.max {
		screen.offset += n
	} else if screen.max > screen.height {
		screen.offset = screen.max
	}
}

// Decrease the offset for scrolling feature by n
func (screen *Screen) DecreaseOffset(n int) {
	if screen.offset > n {
		screen.offset -= n
	} else {
		screen.offset = 0
	}
}

func (screen *Screen) ScrollTop() {
	screen.offset = 0
}

func (screen *Screen) ScrollBottom() {
	if screen.max > screen.height {
		screen.offset = screen.max
	}
}

func (screen *Screen) DrawOldQuotes(quotes *Quotes) {
	screen.draw(screen.layout.Quotes(quotes), true)
	termbox.Flush()
}

func (screen *Screen) DrawOldMarket(market *Market) {
	screen.draw(screen.layout.Market(market), false)
	termbox.Flush()
}

// Draw accepts variable number of arguments and knows how to display the
// market data, stock quotes, current time, and an arbitrary string.
func (screen *Screen) Draw(objects ...interface{}) *Screen {
	zonename, _ := time.Now().In(time.Local).Zone()
	if screen.pausedAt != nil {
		defer screen.DrawLine(0, 0, `<right><r>`+screen.pausedAt.Format(`3:04:05pm `+zonename)+`</r></right>`)
	}
	for _, ptr := range objects {
		switch ptr.(type) {
		case *Market:
			object := ptr.(*Market)
			screen.draw(screen.layout.Market(object.Fetch()), false)
		case *Quotes:
			object := ptr.(*Quotes)
			screen.draw(screen.layout.Quotes(object.Fetch()), true)
		case time.Time:
			timestamp := ptr.(time.Time).Format(`3:04:05pm ` + zonename)
			screen.DrawLine(0, 0, `<right><time>`+timestamp+`</></right>`)
		default:
			screen.draw(ptr.(string), false)
		}
	}

	termbox.Flush()

	return screen
}

// DrawLine takes the incoming string, tokenizes it to extract markup
// elements, and displays it all starting at (x,y) location.

// DrawLineFlush gives the option to flush screen after drawing

// wrapper for DrawLineFlush
func (screen *Screen) DrawLine(x int, y int, str string) {
	screen.DrawLineFlush(x, y, str, true)
}

func (screen *Screen) DrawLineFlush(x int, y int, str string, flush bool) {
	start, column := 0, 0

	for _, token := range screen.markup.Tokenize(str) {
		// First check if it's a tag. Tags are eaten up and not displayed.
		if screen.markup.IsTag(token) {
			continue
		}

		// Here comes the actual text: display it one character at a time.
		for i, char := range token {
			if !screen.markup.RightAligned {
				start = x + column
				column++
			} else {
				start = screen.width - len(token) + i
			}
			termbox.SetCell(start, y, char, screen.markup.Foreground, screen.markup.Background)
		}
	}
	if flush {
		termbox.Flush()
	}
}

// Underlying workhorse function that takes multiline string, splits it into
// lines, and displays them row by row.
func (screen *Screen) draw(str string, offset bool) {
	if !screen.cleared {
		screen.Clear()
	}
	var allLines []string
	drewHeading := false

	screen.width, screen.height = termbox.Size()

	tempFormat := "%" + strconv.Itoa(screen.width) + "s"
	blankLine := fmt.Sprintf(tempFormat, "")
	allLines = strings.Split(str, "\n")

	if offset {
		screen.max = len(allLines) - screen.height
	}

	// Write the lines being updated.
	for row := 0; row < len(allLines); row++ {
		if offset {
			// Did we draw the underlined heading row?  This is a crude
			// check, but--see comments below...
			// --- Heading row only appears for quotes, so offset is true
			if !drewHeading {
				if strings.Contains(allLines[row], "Ticker") &&
					strings.Contains(allLines[row], "Last") &&
					strings.Contains(allLines[row], "Change") {
					drewHeading = true
					screen.headerLine = row
					screen.DrawLine(0, row, allLines[row])
					// move on to the point to offset to
					row += screen.offset
				}
			} else {
				// only write the necessary lines
				if row <= len(allLines) &&
					row > screen.headerLine {
					screen.DrawLineFlush(0, row-screen.offset, allLines[row], false)
				} else if row > len(allLines) + 1 {
					row = len(allLines)
				}
			}
		} else {
			screen.DrawLineFlush(0, row, allLines[row], false)
		}
	}
	// If the quotes lines in this cycle are shorter than in the previous
	// cycles, e.g., because a filter was just applied, then one or more
	// lines from the previous cycles will not be cleared.  Since the
	// incoming lines don't mark explicitly whether they are part of the
	// market summary or quotes, we can't check whether quotes were updated
	// in a way that is robust for code changes.  This is a simple test: if
	// we drew the heading row ("Ticker Last Change..."), then we are
	// updating the quotes section in this cycle, and we should pad the
	// quotes section with blank lines.  If we didn't draw the heading row,
	// then we probably only updated the market summary at the top in this
	// cycle.  In that case, padding with blank lines would overwrite the
	// stocks list.)
	if drewHeading {
		for i := len(allLines) - 1 - screen.offset; i < screen.height; i++ {
			if i > screen.headerLine {
				screen.DrawLine(0, i, blankLine)
			}
		}
	}
}
<|endoftext|>


<commit_before>package shp

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// Writer is the type that is used to write a new shapefile.
type Writer struct {
	filename     string
	shp          writeSeekCloser
	shx          writeSeekCloser
	GeometryType ShapeType
	num          int32
	bbox         Box

	dbf             writeSeekCloser
	dbfFields       []Field
	dbfHeaderLength int16
	dbfRecordLength int16
}

type writeSeekCloser interface {
	io.Writer
	io.Seeker
	io.Closer
}

// Create returns a point to new Writer and the first error that was
// encountered. In case an error occurred the returned Writer point will be nil
// This also creates a corresponding SHX file. It is important to use Close()
// when done because that method writes all the headers for each file (SHP, SHX
// and DBF).
// If filename does not end on ".shp" already, it will be treated as the basename
// for the file and the ".shp" extension will be appended to that name.
func Create(filename string, t ShapeType) (*Writer, error) {
	if strings.HasSuffix(strings.ToLower(filename), ".shp") {
		filename = filename[0 : len(filename)-4]
	}
	shp, err := os.Create(filename + ".shp")
	if err != nil {
		return nil, err
	}
	shx, err := os.Create(filename + ".shx")
	if err != nil {
		return nil, err
	}
	shp.Seek(100, io.SeekStart)
	shx.Seek(100, io.SeekStart)
	w := &Writer{
		filename:     filename,
		shp:          shp,
		shx:          shx,
		GeometryType: t,
	}
	return w, nil
}

// Append returns a Writer pointer that will append to the given shapefile and
// the first error that was encounted during creation of that Writer. The
// shapefile must have a valid index file.
func Append(filename string) (*Writer, error) {
	shp, err := os.OpenFile(filename, os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}
	ext := filepath.Ext(filename)
	basename := filename[:len(filename)-len(ext)]
	w := &Writer{
		filename: basename,
		shp:      shp,
	}
	_, err = shp.Seek(32, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP geometry type: %v", err)
	}
	err = binary.Read(shp, binary.LittleEndian, &w.GeometryType)
	if err != nil {
		return nil, fmt.Errorf("cannot read geometry type: %v", err)
	}
	er := &errReader{Reader: shp}
	w.bbox.MinX = readFloat64(er)
	w.bbox.MinY = readFloat64(er)
	w.bbox.MaxX = readFloat64(er)
	w.bbox.MaxY = readFloat64(er)
	if er.e != nil {
		return nil, fmt.Errorf("cannot read bounding box: %v", er.e)
	}

	shx, err := os.OpenFile(basename+".shx", os.O_RDWR, 0666)
	if os.IsNotExist(err) {
		// TODO allow index file to not exist, in that case just
		// read through all the shapes and create it on the fly
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open shapefile index: %v", err)
	}
	_, err = shx.Seek(-8, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape index: %v", err)
	}
	var offset int32
	err = binary.Read(shx, binary.BigEndian, &offset)
	if err != nil {
		return nil, fmt.Errorf("cannot read last shape index: %v", err)
	}
	offset = offset * 2
	_, err = shp.Seek(int64(offset), io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape: %v", err)
	}
	err = binary.Read(shp, binary.BigEndian, &w.num)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of last shape: %v", err)
	}
	_, err = shp.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP end: %v", err)
	}
	_, err = shx.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHX end: %v", err)
	}
	w.shx = shx

	dbf, err := os.Open(basename + ".dbf")
	if os.IsNotExist(err) {
		return w, nil // it's okay if the DBF does not exist
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open DBF: %v", err)
	}

	_, err = dbf.Seek(8, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfHeaderLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read header length from DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfRecordLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read record length from DBF: %v", err)
	}

	_, err = dbf.Seek(20, io.SeekCurrent) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	numFields := int(math.Floor(float64(w.dbfHeaderLength-33) / 32.0))
	w.dbfFields = make([]Field, numFields)
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfFields)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of fields from DBF: %v", err)
	}
	_, err = dbf.Seek(0, io.SeekEnd) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek to DBF end: %v", err)
	}
	w.dbf = dbf

	return w, nil
}

// Write shape to the Shapefile. This also creates
// a record in the SHX file and DBF file (if it is
// initialized). Returns the index of the written object
// which can be used in WriteAttribute.
func (w *Writer) Write(shape Shape) int32 {
	// increate bbox
	if w.num == 0 {
		w.bbox = shape.BBox()
	} else {
		w.bbox.Extend(shape.BBox())
	}

	w.num++
	binary.Write(w.shp, binary.BigEndian, w.num)
	w.shp.Seek(4, io.SeekCurrent)
	start, _ := w.shp.Seek(0, io.SeekCurrent)
	binary.Write(w.shp, binary.LittleEndian, w.GeometryType)
	shape.write(w.shp)
	finish, _ := w.shp.Seek(0, io.SeekCurrent)
	length := int32(math.Floor((float64(finish) - float64(start)) / 2.0))
	w.shp.Seek(start-4, io.SeekStart)
	binary.Write(w.shp, binary.BigEndian, length)
	w.shp.Seek(finish, io.SeekStart)

	// write shx
	binary.Write(w.shx, binary.BigEndian, int32((start-8)/2))
	binary.Write(w.shx, binary.BigEndian, length)

	// write empty record to dbf
	if w.dbf != nil {
		w.writeEmptyRecord()
	}

	return w.num - 1
}

// Close closes the Writer. This must be used at the end of
// the transaction because it writes the correct headers
// to the SHP/SHX and DBF files before closing.
func (w *Writer) Close() {
	w.writeHeader(w.shx)
	w.writeHeader(w.shp)
	w.shp.Close()
	w.shx.Close()

	if w.dbf == nil {
		w.SetFields([]Field{})
	}
	w.writeDbfHeader(w.dbf)
	w.dbf.Close()
}

// writeHeader wrires SHP/SHX headers to ws.
func (w *Writer) writeHeader(ws io.WriteSeeker) {
	filelength, _ := ws.Seek(0, io.SeekEnd)
	if filelength == 0 {
		filelength = 100
	}
	ws.Seek(0, io.SeekStart)
	// file code
	binary.Write(ws, binary.BigEndian, []int32{9994, 0, 0, 0, 0, 0})
	// file length
	binary.Write(ws, binary.BigEndian, int32(filelength/2))
	// version and shape type
	binary.Write(ws, binary.LittleEndian, []int32{1000, int32(w.GeometryType)})
	// bounding box
	binary.Write(ws, binary.LittleEndian, w.bbox)
	// elevation, measure
	binary.Write(ws, binary.LittleEndian, []float64{0.0, 0.0, 0.0, 0.0})
}

// writeDbfHeader writes a DBF header to ws.
func (w *Writer) writeDbfHeader(ws io.WriteSeeker) {
	ws.Seek(0, 0)
	// version, year (YEAR-1990), month, day
	binary.Write(ws, binary.LittleEndian, []byte{3, 24, 5, 3})
	// number of records
	binary.Write(ws, binary.LittleEndian, w.num)
	// header length, record length
	binary.Write(ws, binary.LittleEndian, []int16{w.dbfHeaderLength, w.dbfRecordLength})
	// padding
	binary.Write(ws, binary.LittleEndian, make([]byte, 20))

	for _, field := range w.dbfFields {
		binary.Write(ws, binary.LittleEndian, field)
	}

	// end with return
	ws.Write([]byte("\r"))
}

// SetFields sets field values in the DBF. This initializes the DBF file and
// should be used prior to writing any attributes.
func (w *Writer) SetFields(fields []Field) error {
	if w.dbf != nil {
		return errors.New("Cannot set fields in existing dbf")
	}

	var err error
	w.dbf, err = os.Create(w.filename + "dbf")
	if err != nil {
		return fmt.Errorf("Failed to open %s.dbf: %v", w.filename, err)
	}
	w.dbfFields = fields

	// calculate record length
	w.dbfRecordLength = int16(1)
	for _, field := range w.dbfFields {
		w.dbfRecordLength += int16(field.Size)
	}

	// header lengh
	w.dbfHeaderLength = int16(len(w.dbfFields)*32 + 33)

	// fill header space with empty bytes for now
	buf := make([]byte, w.dbfHeaderLength)
	binary.Write(w.dbf, binary.LittleEndian, buf)

	// write empty records
	for n := int32(0); n < w.num; n++ {
		w.writeEmptyRecord()
	}
	return nil
}

// Writes an empty record to the end of the DBF. This
// works by seeking to the end of the file and writing
// dbfRecordLength number of bytes. The first byte is a
// space that indicates a new record.
func (w *Writer) writeEmptyRecord() {
	w.dbf.Seek(0, io.SeekEnd)
	buf := make([]byte, w.dbfRecordLength)
	buf[0] = ' '
	binary.Write(w.dbf, binary.LittleEndian, buf)
}

// WriteAttribute writes value for field into the given row in the DBF. Row
// number should be the same as the order the Shape was written to the
// Shapefile. The field value corresponds to the field in the slice used in
// SetFields.
func (w *Writer) WriteAttribute(row int, field int, value interface{}) error {
	var buf []byte
	switch v := value.(type) {
	case int:
		buf = []byte(strconv.Itoa(v))
	case float64:
		precision := w.dbfFields[field].Precision
		buf = []byte(strconv.FormatFloat(v, 'f', int(precision), 64))
	case string:
		buf = []byte(v)
	default:
		return fmt.Errorf("Unsupported value type: %T", v)
	}

	if w.dbf == nil {
		return errors.New("Initialize DBF by using SetFields first")
	}
	if sz := int(w.dbfFields[field].Size); len(buf) > sz {
		return fmt.Errorf("Unable to write field %v: %q exceeds field length %v", field, buf, sz)
	}

	seekTo := 1 + int64(w.dbfHeaderLength) + (int64(row) * int64(w.dbfRecordLength))
	for n := 0; n < field; n++ {
		seekTo += int64(w.dbfFields[n].Size)
	}
	w.dbf.Seek(seekTo, io.SeekStart)
	return binary.Write(w.dbf, binary.LittleEndian, buf)
}

// BBox returns the bounding box of the Writer.
func (w *Writer) BBox() Box {
	return w.bbox
}
<commit_msg>Fix Create using wrong dbf filename<commit_after>package shp

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// Writer is the type that is used to write a new shapefile.
type Writer struct {
	filename     string
	shp          writeSeekCloser
	shx          writeSeekCloser
	GeometryType ShapeType
	num          int32
	bbox         Box

	dbf             writeSeekCloser
	dbfFields       []Field
	dbfHeaderLength int16
	dbfRecordLength int16
}

type writeSeekCloser interface {
	io.Writer
	io.Seeker
	io.Closer
}

// Create returns a point to new Writer and the first error that was
// encountered. In case an error occurred the returned Writer point will be nil
// This also creates a corresponding SHX file. It is important to use Close()
// when done because that method writes all the headers for each file (SHP, SHX
// and DBF).
// If filename does not end on ".shp" already, it will be treated as the basename
// for the file and the ".shp" extension will be appended to that name.
func Create(filename string, t ShapeType) (*Writer, error) {
	if strings.HasSuffix(strings.ToLower(filename), ".shp") {
		filename = filename[0 : len(filename)-4]
	}
	shp, err := os.Create(filename + ".shp")
	if err != nil {
		return nil, err
	}
	shx, err := os.Create(filename + ".shx")
	if err != nil {
		return nil, err
	}
	shp.Seek(100, io.SeekStart)
	shx.Seek(100, io.SeekStart)
	w := &Writer{
		filename:     filename,
		shp:          shp,
		shx:          shx,
		GeometryType: t,
	}
	return w, nil
}

// Append returns a Writer pointer that will append to the given shapefile and
// the first error that was encounted during creation of that Writer. The
// shapefile must have a valid index file.
func Append(filename string) (*Writer, error) {
	shp, err := os.OpenFile(filename, os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}
	ext := filepath.Ext(filename)
	basename := filename[:len(filename)-len(ext)]
	w := &Writer{
		filename: basename,
		shp:      shp,
	}
	_, err = shp.Seek(32, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP geometry type: %v", err)
	}
	err = binary.Read(shp, binary.LittleEndian, &w.GeometryType)
	if err != nil {
		return nil, fmt.Errorf("cannot read geometry type: %v", err)
	}
	er := &errReader{Reader: shp}
	w.bbox.MinX = readFloat64(er)
	w.bbox.MinY = readFloat64(er)
	w.bbox.MaxX = readFloat64(er)
	w.bbox.MaxY = readFloat64(er)
	if er.e != nil {
		return nil, fmt.Errorf("cannot read bounding box: %v", er.e)
	}

	shx, err := os.OpenFile(basename+".shx", os.O_RDWR, 0666)
	if os.IsNotExist(err) {
		// TODO allow index file to not exist, in that case just
		// read through all the shapes and create it on the fly
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open shapefile index: %v", err)
	}
	_, err = shx.Seek(-8, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape index: %v", err)
	}
	var offset int32
	err = binary.Read(shx, binary.BigEndian, &offset)
	if err != nil {
		return nil, fmt.Errorf("cannot read last shape index: %v", err)
	}
	offset = offset * 2
	_, err = shp.Seek(int64(offset), io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to last shape: %v", err)
	}
	err = binary.Read(shp, binary.BigEndian, &w.num)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of last shape: %v", err)
	}
	_, err = shp.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHP end: %v", err)
	}
	_, err = shx.Seek(0, io.SeekEnd)
	if err != nil {
		return nil, fmt.Errorf("cannot seek to SHX end: %v", err)
	}
	w.shx = shx

	dbf, err := os.Open(basename + ".dbf")
	if os.IsNotExist(err) {
		return w, nil // it's okay if the DBF does not exist
	}
	if err != nil {
		return nil, fmt.Errorf("cannot open DBF: %v", err)
	}

	_, err = dbf.Seek(8, io.SeekStart)
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfHeaderLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read header length from DBF: %v", err)
	}
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfRecordLength)
	if err != nil {
		return nil, fmt.Errorf("cannot read record length from DBF: %v", err)
	}

	_, err = dbf.Seek(20, io.SeekCurrent) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek in DBF: %v", err)
	}
	numFields := int(math.Floor(float64(w.dbfHeaderLength-33) / 32.0))
	w.dbfFields = make([]Field, numFields)
	err = binary.Read(dbf, binary.LittleEndian, &w.dbfFields)
	if err != nil {
		return nil, fmt.Errorf("cannot read number of fields from DBF: %v", err)
	}
	_, err = dbf.Seek(0, io.SeekEnd) // skip padding
	if err != nil {
		return nil, fmt.Errorf("cannot seek to DBF end: %v", err)
	}
	w.dbf = dbf

	return w, nil
}

// Write shape to the Shapefile. This also creates
// a record in the SHX file and DBF file (if it is
// initialized). Returns the index of the written object
// which can be used in WriteAttribute.
func (w *Writer) Write(shape Shape) int32 {
	// increate bbox
	if w.num == 0 {
		w.bbox = shape.BBox()
	} else {
		w.bbox.Extend(shape.BBox())
	}

	w.num++
	binary.Write(w.shp, binary.BigEndian, w.num)
	w.shp.Seek(4, io.SeekCurrent)
	start, _ := w.shp.Seek(0, io.SeekCurrent)
	binary.Write(w.shp, binary.LittleEndian, w.GeometryType)
	shape.write(w.shp)
	finish, _ := w.shp.Seek(0, io.SeekCurrent)
	length := int32(math.Floor((float64(finish) - float64(start)) / 2.0))
	w.shp.Seek(start-4, io.SeekStart)
	binary.Write(w.shp, binary.BigEndian, length)
	w.shp.Seek(finish, io.SeekStart)

	// write shx
	binary.Write(w.shx, binary.BigEndian, int32((start-8)/2))
	binary.Write(w.shx, binary.BigEndian, length)

	// write empty record to dbf
	if w.dbf != nil {
		w.writeEmptyRecord()
	}

	return w.num - 1
}

// Close closes the Writer. This must be used at the end of
// the transaction because it writes the correct headers
// to the SHP/SHX and DBF files before closing.
func (w *Writer) Close() {
	w.writeHeader(w.shx)
	w.writeHeader(w.shp)
	w.shp.Close()
	w.shx.Close()

	if w.dbf == nil {
		w.SetFields([]Field{})
	}
	w.writeDbfHeader(w.dbf)
	w.dbf.Close()
}

// writeHeader wrires SHP/SHX headers to ws.
func (w *Writer) writeHeader(ws io.WriteSeeker) {
	filelength, _ := ws.Seek(0, io.SeekEnd)
	if filelength == 0 {
		filelength = 100
	}
	ws.Seek(0, io.SeekStart)
	// file code
	binary.Write(ws, binary.BigEndian, []int32{9994, 0, 0, 0, 0, 0})
	// file length
	binary.Write(ws, binary.BigEndian, int32(filelength/2))
	// version and shape type
	binary.Write(ws, binary.LittleEndian, []int32{1000, int32(w.GeometryType)})
	// bounding box
	binary.Write(ws, binary.LittleEndian, w.bbox)
	// elevation, measure
	binary.Write(ws, binary.LittleEndian, []float64{0.0, 0.0, 0.0, 0.0})
}

// writeDbfHeader writes a DBF header to ws.
func (w *Writer) writeDbfHeader(ws io.WriteSeeker) {
	ws.Seek(0, 0)
	// version, year (YEAR-1990), month, day
	binary.Write(ws, binary.LittleEndian, []byte{3, 24, 5, 3})
	// number of records
	binary.Write(ws, binary.LittleEndian, w.num)
	// header length, record length
	binary.Write(ws, binary.LittleEndian, []int16{w.dbfHeaderLength, w.dbfRecordLength})
	// padding
	binary.Write(ws, binary.LittleEndian, make([]byte, 20))

	for _, field := range w.dbfFields {
		binary.Write(ws, binary.LittleEndian, field)
	}

	// end with return
	ws.Write([]byte("\r"))
}

// SetFields sets field values in the DBF. This initializes the DBF file and
// should be used prior to writing any attributes.
func (w *Writer) SetFields(fields []Field) error {
	if w.dbf != nil {
		return errors.New("Cannot set fields in existing dbf")
	}

	var err error
	w.dbf, err = os.Create(w.filename + ".dbf")
	if err != nil {
		return fmt.Errorf("Failed to open %s.dbf: %v", w.filename, err)
	}
	w.dbfFields = fields

	// calculate record length
	w.dbfRecordLength = int16(1)
	for _, field := range w.dbfFields {
		w.dbfRecordLength += int16(field.Size)
	}

	// header lengh
	w.dbfHeaderLength = int16(len(w.dbfFields)*32 + 33)

	// fill header space with empty bytes for now
	buf := make([]byte, w.dbfHeaderLength)
	binary.Write(w.dbf, binary.LittleEndian, buf)

	// write empty records
	for n := int32(0); n < w.num; n++ {
		w.writeEmptyRecord()
	}
	return nil
}

// Writes an empty record to the end of the DBF. This
// works by seeking to the end of the file and writing
// dbfRecordLength number of bytes. The first byte is a
// space that indicates a new record.
func (w *Writer) writeEmptyRecord() {
	w.dbf.Seek(0, io.SeekEnd)
	buf := make([]byte, w.dbfRecordLength)
	buf[0] = ' '
	binary.Write(w.dbf, binary.LittleEndian, buf)
}

// WriteAttribute writes value for field into the given row in the DBF. Row
// number should be the same as the order the Shape was written to the
// Shapefile. The field value corresponds to the field in the slice used in
// SetFields.
func (w *Writer) WriteAttribute(row int, field int, value interface{}) error {
	var buf []byte
	switch v := value.(type) {
	case int:
		buf = []byte(strconv.Itoa(v))
	case float64:
		precision := w.dbfFields[field].Precision
		buf = []byte(strconv.FormatFloat(v, 'f', int(precision), 64))
	case string:
		buf = []byte(v)
	default:
		return fmt.Errorf("Unsupported value type: %T", v)
	}

	if w.dbf == nil {
		return errors.New("Initialize DBF by using SetFields first")
	}
	if sz := int(w.dbfFields[field].Size); len(buf) > sz {
		return fmt.Errorf("Unable to write field %v: %q exceeds field length %v", field, buf, sz)
	}

	seekTo := 1 + int64(w.dbfHeaderLength) + (int64(row) * int64(w.dbfRecordLength))
	for n := 0; n < field; n++ {
		seekTo += int64(w.dbfFields[n].Size)
	}
	w.dbf.Seek(seekTo, io.SeekStart)
	return binary.Write(w.dbf, binary.LittleEndian, buf)
}

// BBox returns the bounding box of the Writer.
func (w *Writer) BBox() Box {
	return w.bbox
}
<|endoftext|>


<commit_before>package wsutil

import (
	"bytes"
	"testing"

	"github.com/gobwas/ws"
)

func TestHandlerHandleClose(t *testing.T) {
	for _, test := range []struct {
		name  string
		state ws.State
		in    ws.Frame
		out   ws.Frame
		err   error
	}{
		{
			in:  ws.NewCloseFrame(nil),
			out: ws.NewCloseFrame(nil),
			err: ClosedError{
				Code: ws.StatusNoStatusRcvd,
			},
		},
		{
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			state: ws.StateServerSide,
			in: ws.MaskFrame(ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			))),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusNormalClosure, string([]byte{0, 200}),
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusProtocolError, ws.ErrProtocolInvalidUTF8.Error(),
			)),
			err: ws.ErrProtocolInvalidUTF8,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			var (
				out = bytes.NewBuffer(nil)
				in  = bytes.NewReader(test.in.Payload)
			)
			c := ControlHandler{
				Src:   in,
				Dst:   out,
				State: test.state,
			}

			err := c.Handle(test.in.Header)
			if err != test.err {
				t.Errorf("unexpected error: %v; want %v", err, test.err)
			}
			{
				act := out.Bytes()
				exp := ws.MustCompileFrame(test.out)
				if !bytes.Equal(act, exp) {
					fa := ws.MustReadFrame(bytes.NewReader(act))
					fe := ws.MustReadFrame(bytes.NewReader(exp))
					t.Errorf(
						"unexpected sent frame:\n\tact: %+v\n\texp: %+v\nbytes:\n\tact: %v\n\texp: %v",
						fa, fe, act, exp,
					)
				}
			}

		})
	}
}
<commit_msg>wsutil: increase test coverage<commit_after>package wsutil

import (
	"bytes"
	"runtime"
	"testing"

	"github.com/gobwas/ws"
)

func TestControlHandler(t *testing.T) {
	for _, test := range []struct {
		name  string
		state ws.State
		in    ws.Frame
		out   ws.Frame
		noOut bool
		err   error
	}{
		{
			name: "ping",
			in:   ws.NewPingFrame(nil),
			out:  ws.NewPongFrame(nil),
		},
		{
			name: "ping",
			in:   ws.NewPingFrame([]byte("catch the ball")),
			out:  ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name:  "ping",
			state: ws.StateServerSide,
			in:    ws.MaskFrame(ws.NewPingFrame([]byte("catch the ball"))),
			out:   ws.NewPongFrame([]byte("catch the ball")),
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame(nil),
			noOut: true,
		},
		{
			name:  "pong",
			in:    ws.NewPongFrame([]byte("catched")),
			noOut: true,
		},
		{
			name: "close",
			in:   ws.NewCloseFrame(nil),
			out:  ws.NewCloseFrame(nil),
			err: ClosedError{
				Code: ws.StatusNoStatusRcvd,
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name:  "close",
			state: ws.StateServerSide,
			in: ws.MaskFrame(ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "goodbye!",
			))),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusGoingAway, "",
			)),
			err: ClosedError{
				Code:   ws.StatusGoingAway,
				Reason: "goodbye!",
			},
		},
		{
			name: "close",
			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusNormalClosure, string([]byte{0, 200}),
			)),
			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
				ws.StatusProtocolError, ws.ErrProtocolInvalidUTF8.Error(),
			)),
			err: ws.ErrProtocolInvalidUTF8,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			defer func() {
				if err := recover(); err != nil {
					stack := make([]byte, 4096)
					n := runtime.Stack(stack, true)
					t.Fatalf(
						"panic recovered: %v\n%s",
						err, stack[:n],
					)
				}
			}()
			var (
				out = bytes.NewBuffer(nil)
				in  = bytes.NewReader(test.in.Payload)
			)
			c := ControlHandler{
				Src:   in,
				Dst:   out,
				State: test.state,
			}

			err := c.Handle(test.in.Header)
			if err != test.err {
				t.Errorf("unexpected error: %v; want %v", err, test.err)
			}

			if in.Len() != 0 {
				t.Errorf("handler did not drained the input")
			}

			act := out.Bytes()
			switch {
			case len(act) == 0 && test.noOut:
				return
			case len(act) == 0 && !test.noOut:
				t.Errorf("unexpected silence")
			case len(act) > 0 && test.noOut:
				t.Errorf("unexpected sent frame")
			default:
				exp := ws.MustCompileFrame(test.out)
				if !bytes.Equal(act, exp) {
					fa := ws.MustReadFrame(bytes.NewReader(act))
					fe := ws.MustReadFrame(bytes.NewReader(exp))
					t.Errorf(
						"unexpected sent frame:\n\tact: %+v\n\texp: %+v\nbytes:\n\tact: %v\n\texp: %v",
						fa, fe, act, exp,
					)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"container/vector"
	"flag"
	"fmt"
	"igo/build"
	"igo/deps"
	"igo/set"
	"os"
	"path"
	"strings"
)

// executeCommand runs the specified tool with the supplied arguments (not
// including the path to the tool itself), chdir'ing to the specified directory
// first. It returns true if and only if the child process returns zero.
func executeCommand(tool string, args []string, dir string) bool {
	fmt.Printf("%s %s\n", tool, strings.Join(args, " "))

	var fullArgs vector.StringVector
	fullArgs.Push(tool)
	fullArgs.AppendVector(&args)

	pid, err := os.ForkExec(
		tool,
		fullArgs.Data(),
		os.Environ(),
		dir,
		[]*os.File{os.Stdin, os.Stdout, os.Stderr})
	if err != nil {
		panic(err)
	}

	waitMsg, err := os.Wait(pid, 0)
	if err != nil {
		panic(err)
	}

	return waitMsg.ExitStatus() == 0
}

// compileFiles invokes 6g with the appropriate arguments for compiling the
// supplied set of .go files, and exits the program if the subprocess fails.
func compileFiles(files *set.StringSet, targetBaseName string) {
	compilerPath := path.Join(os.Getenv("GOBIN"), "6g")
	gopackPath := path.Join(os.Getenv("GOBIN"), "gopack")

	targetDir, _ := path.Split(targetBaseName)
	if targetDir != "" {
		os.MkdirAll(path.Join("igo-out", targetDir), 0700)
	}

	// Compile
	var compilerArgs vector.StringVector
	compilerArgs.Push("-o")
	compilerArgs.Push(targetBaseName + ".6")

	for file := range files.Iter() {
		compilerArgs.Push(path.Join("../", file))
	}

	if !executeCommand(compilerPath, compilerArgs.Data(), "igo-out/") {
		os.Exit(1)
	}

	// Pack
	var gopackArgs vector.StringVector
	gopackArgs.Push("grc")
	gopackArgs.Push(targetBaseName + ".a")
	gopackArgs.Push(targetBaseName + ".6")

	if !executeCommand(gopackPath, gopackArgs.Data(), "igo-out/") {
		os.Exit(1)
	}
}

func printUsageAndExit() {
	fmt.Println("Usage:")
	fmt.Println("  igo build <directory name>")
	os.Exit(1)
}

func main() {
	flag.Parse()

	if flag.NArg() != 2 {
		printUsageAndExit()
	}

	command := flag.Arg(0)
	if command != "build" {
		printUsageAndExit()
	}

	// Grab dependency and file information for every local package, starting
	// with the specified one. We consider a package local if it starts with "./".
	requiredFiles := make(map[string]*set.StringSet)
	packageDeps := make(map[string]*set.StringSet)

	var remainingDirs vector.StringVector
	remainingDirs.Push("./" + flag.Arg(1))

	for remainingDirs.Len() > 0 {
		dir := remainingDirs.Pop()

		// Have we already processed this directory?
		_, alreadyDone := packageDeps[dir]
		if alreadyDone {
			continue
		}

		dirInfo := build.GetDirectoryInfo(dir, false)
		if dirInfo.PackageName == "" {
			fmt.Printf("Couldn't find .go files to build in directory: %s\n", dir)
			os.Exit(1)
		}

		// Stash information about this package, and add its local dependencies to
		// the queue.
		requiredFiles[dir] = dirInfo.Files
		packageDeps[dir] = dirInfo.Deps

		for dep := range dirInfo.Deps.Iter() {
			if strings.HasPrefix(dep, "./") {
				remainingDirs.Push(dep)
			}
		}
	}

	// Order the packages by their dependencies.
	totalOrder := deps.BuildTotalOrder(packageDeps)
	fmt.Println("Found these packages to compile:")
	for _, packageName := range totalOrder {
		fmt.Printf("  %s\n", packageName)
	}

	// Create a directory to hold outputs, deleting the old one first.
	os.RemoveAll("igo-out")
	os.Mkdir("igo-out", 0700)

	// Compile each of the packages in turn.
	for _, currentPackage := range totalOrder {
		fmt.Printf("\nCompiling package: %s\n", currentPackage)
		compileFiles(requiredFiles[currentPackage], currentPackage)
	}
}
<commit_msg>Added support for linking binaries.<commit_after>package main

import (
	"container/vector"
	"flag"
	"fmt"
	"igo/build"
	"igo/deps"
	"igo/set"
	"os"
	"path"
	"strings"
)

// executeCommand runs the specified tool with the supplied arguments (not
// including the path to the tool itself), chdir'ing to the specified directory
// first. It returns true if and only if the child process returns zero.
func executeCommand(tool string, args []string, dir string) bool {
	fmt.Printf("%s %s\n", tool, strings.Join(args, " "))

	var fullArgs vector.StringVector
	fullArgs.Push(tool)
	fullArgs.AppendVector(&args)

	pid, err := os.ForkExec(
		tool,
		fullArgs.Data(),
		os.Environ(),
		dir,
		[]*os.File{os.Stdin, os.Stdout, os.Stderr})
	if err != nil {
		panic(err)
	}

	waitMsg, err := os.Wait(pid, 0)
	if err != nil {
		panic(err)
	}

	return waitMsg.ExitStatus() == 0
}

// compileFiles invokes 6g with the appropriate arguments for compiling the
// supplied set of .go files, and exits the program if the subprocess fails.
func compileFiles(files *set.StringSet, targetBaseName string) {
	compilerPath := path.Join(os.Getenv("GOBIN"), "6g")
	gopackPath := path.Join(os.Getenv("GOBIN"), "gopack")

	targetDir, _ := path.Split(targetBaseName)
	if targetDir != "" {
		os.MkdirAll(path.Join("igo-out", targetDir), 0700)
	}

	// Compile
	var compilerArgs vector.StringVector
	compilerArgs.Push("-o")
	compilerArgs.Push(targetBaseName + ".6")

	for file := range files.Iter() {
		compilerArgs.Push(path.Join("../", file))
	}

	if !executeCommand(compilerPath, compilerArgs.Data(), "igo-out/") {
		os.Exit(1)
	}

	// Pack
	var gopackArgs vector.StringVector
	gopackArgs.Push("grc")
	gopackArgs.Push(targetBaseName + ".a")
	gopackArgs.Push(targetBaseName + ".6")

	if !executeCommand(gopackPath, gopackArgs.Data(), "igo-out/") {
		os.Exit(1)
	}
}

// linkBinary calls 6l to link the binary of the given name, which must have
// already been compiled with compileFiles.
func linkBinary(name string) {
	linkerPath := path.Join(os.Getenv("GOBIN"), "6l")

	var linkerArgs vector.StringVector
	linkerArgs.Push("-o")
	linkerArgs.Push(name)
	linkerArgs.Push(name + ".6")

	if !executeCommand(linkerPath, linkerArgs.Data(), "igo-out/") {
		os.Exit(1)
	}
}

func printUsageAndExit() {
	fmt.Println("Usage:")
	fmt.Println("  igo build <directory name>")
	os.Exit(1)
}

func main() {
	flag.Parse()

	if flag.NArg() != 2 {
		printUsageAndExit()
	}

	command := flag.Arg(0)
	if command != "build" {
		printUsageAndExit()
	}

	// Grab dependency and file information for every local package, starting
	// with the specified one. We consider a package local if it starts with "./".
	requiredFiles := make(map[string]*set.StringSet)
	packageDeps := make(map[string]*set.StringSet)

	specifiedDir := flag.Arg(1)
	var remainingDirs vector.StringVector
	remainingDirs.Push("./" + specifiedDir)

	for remainingDirs.Len() > 0 {
		dir := remainingDirs.Pop()

		// Have we already processed this directory?
		_, alreadyDone := packageDeps[dir]
		if alreadyDone {
			continue
		}

		dirInfo := build.GetDirectoryInfo(dir, false)
		if dirInfo.PackageName == "" {
			fmt.Printf("Couldn't find .go files to build in directory: %s\n", dir)
			os.Exit(1)
		}

		// Stash information about this package, and add its local dependencies to
		// the queue.
		requiredFiles[dir] = dirInfo.Files
		packageDeps[dir] = dirInfo.Deps

		for dep := range dirInfo.Deps.Iter() {
			if strings.HasPrefix(dep, "./") {
				remainingDirs.Push(dep)
			}
		}
	}

	// Order the packages by their dependencies.
	totalOrder := deps.BuildTotalOrder(packageDeps)
	fmt.Println("Found these packages to compile:")
	for _, packageName := range totalOrder {
		fmt.Printf("  %s\n", packageName)
	}

	// Create a directory to hold outputs, deleting the old one first.
	os.RemoveAll("igo-out")
	os.Mkdir("igo-out", 0700)

	// Compile each of the packages in turn.
	for _, currentPackage := range totalOrder {
		fmt.Printf("\nCompiling package: %s\n", currentPackage)
		compileFiles(requiredFiles[currentPackage], currentPackage)
	}

	// If this is a binary, also link it.
	if build.GetDirectoryInfo(specifiedDir, false).PackageName == "main" {
		linkBinary(specifiedDir)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/nsf/termbox-go"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"
	"unicode"
)

const SaveDirName = "saves_5CC9DB70-EEC5-47EA-94B6-398BFC12E4A7"

var mainMenu mainMenuUI

func init() {
	mainMenu.rand = rand.New(rand.NewSource(time.Now().UnixNano()))

	for i := 0; i < 16; i++ {
		const characters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
		mainMenu.seed = append(mainMenu.seed, rune(characters[mainMenu.rand.Intn(len(characters))]))
	}

	saves, err := os.Open(SaveDirName)
	if err != nil {
		if os.IsNotExist(err) {
			return
		}
		panic(err)
	}
	names, err := saves.Readdirnames(0)
	if err != nil {
		panic(err)
	}
	for _, name := range names {
		if strings.HasSuffix(name, ".sav") {
			mainMenu.saveNames = append(mainMenu.saveNames, name[:len(name)-len(".sav")])
		}
	}
}

const (
	menuStateMain = iota
	menuStateLoad
	menuStateNew
)

type mainMenuUI struct {
	rand        *rand.Rand
	saveNames   []string
	choiceIndex int
	state       uint
	saveName    []rune
	seed        []rune
}

func (m *mainMenuUI) render(w, h int) {
	m.drawTextFlicker(w, 2, "5CC9DB70-EEC5-47EA-94B6-398BFC12E4A7", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)

	switch m.state {
	case menuStateMain:
		skip := 0

		if m.choiceIndex > h/2-5 {
			skip = m.choiceIndex - (h/2 - 5)
		}

		for i, name := range m.saveNames[skip:] {
			if m.choiceIndex == i+skip {
				m.drawText(w, 5+i, fmt.Sprintf("Load %q", name), termbox.ColorBlack, termbox.ColorWhite)
			} else {
				m.drawText(w, 5+i, fmt.Sprintf("Load %q", name), termbox.ColorWhite, termbox.ColorBlack)
			}
		}
		if m.choiceIndex == len(m.saveNames) {
			m.drawText(w, len(m.saveNames)+5-skip, "New Game", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, len(m.saveNames)+5-skip, "New Game", termbox.ColorWhite, termbox.ColorBlack)
		}

	case menuStateNew:
		m.drawText(w, 5, "Save Name", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)
		if m.choiceIndex == 0 {
			m.drawText(w, 6, string(m.saveName)+"_", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, 6, string(m.saveName), termbox.ColorWhite, termbox.ColorBlack)
		}

		m.drawText(w, 8, "Seed", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)
		if m.choiceIndex == 1 {
			m.drawText(w, 9, string(m.seed)+"_", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, 9, string(m.seed), termbox.ColorWhite, termbox.ColorBlack)
		}
	}
}

func (m *mainMenuUI) drawText(w, y int, s string, fg, bg termbox.Attribute) {
	r := []rune(s)

	for i := 0; i < w; i++ {
		termbox.SetCell(i, y, ' ', fg, bg)
	}

	for i, ch := range r {
		termbox.SetCell((w-len(r))/2+i, y, ch, fg, bg)
	}
}

func (m *mainMenuUI) drawTextFlicker(w, y int, s string, fg, bg termbox.Attribute) {
	r := []rune(s)

	flicker := m.rand.Intn(len(r))
	if (r[flicker] < '0' || r[flicker] > '9') && (r[flicker] < 'A' || r[flicker] > 'F') {
		flicker = -1
	}

	for i := 0; i < w; i++ {
		termbox.SetCell(i, y, ' ', fg, bg)
	}

	for i, ch := range r {
		if i == flicker {
			const hex = "0123456789ABCDEF"
			termbox.SetCell((w-len(r))/2+i, y, rune(hex[m.rand.Intn(len(hex))]), fg^termbox.AttrBold, bg)
		} else {
			termbox.SetCell((w-len(r))/2+i, y, ch, fg, bg)
		}
	}
}

func (m *mainMenuUI) inputKey(key termbox.Key, ch rune, mod termbox.Modifier) bool {
	switch m.state {
	case menuStateMain:
		switch {
		case key == termbox.KeyArrowDown:
			m.choiceIndex = (m.choiceIndex + 1) % (len(m.saveNames) + 1)
		case key == termbox.KeyArrowUp:
			m.choiceIndex = (m.choiceIndex + (len(m.saveNames) + 1 - 1)) % (len(m.saveNames) + 1)
		case key == termbox.KeyArrowLeft || key == termbox.KeyArrowRight:
			// silently ignore
		case key == termbox.KeyEnter:
			if m.choiceIndex < len(m.saveNames) {
				m.loadGame(m.saveNames[m.choiceIndex])
			} else {
				switch m.choiceIndex - len(m.saveNames) {
				case 0:
					m.newGame()
				}
			}
		case key == termbox.KeyEsc:
			return false
		default:
			panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
		}

	case menuStateNew:
		const fieldCount = 2

		switch {
		case key == termbox.KeyEsc:
			m.state = menuStateMain
			m.choiceIndex = len(m.saveNames)
		case key == termbox.KeyEnter:
			m.choiceIndex++
			if m.choiceIndex >= fieldCount {
				m.choiceIndex = fieldCount - 1
				if len(m.saveName) == 0 {
					m.choiceIndex = 0
					fmt.Print("\a")
					return true
				}
				w := &World{
					Seed:     NewSeed(string(m.seed)),
					saveName: filepath.Join(SaveDirName, string(m.saveName)+".sav"),
				}
				err := w.save(os.O_CREATE | os.O_EXCL | os.O_WRONLY)
				if err == nil {
					err = w.AfterLoad()
				}
				if err != nil {
					panic(err)
				}
				worldLock.Lock()
				world = w
				worldLock.Unlock()
			}
		case key == termbox.KeyArrowDown:
			m.choiceIndex = (m.choiceIndex + 1) % fieldCount
		case key == termbox.KeyArrowUp:
			m.choiceIndex = (m.choiceIndex + (fieldCount - 1)) % fieldCount
		case key == termbox.KeyArrowLeft || key == termbox.KeyArrowRight:
			fmt.Print("\a")
		case m.choiceIndex == 0 && (key == termbox.KeyBackspace || key == termbox.KeyBackspace2):
			if len(m.saveName) == 0 {
				fmt.Print("\a")
			} else {
				m.saveName = m.saveName[:len(m.saveName)-1]
			}
		case m.choiceIndex == 0 && key == termbox.KeySpace:
			if len(m.saveName) == 0 {
				fmt.Print("\a")
			} else {
				m.saveName = append(m.saveName, ' ')
			}
		case m.choiceIndex == 0 && ch != 0:
			// unicode.Punctuation is not included due to characters like /
			if ch != '_' && ch != '-' && !unicode.In(ch, unicode.Letter, unicode.Number, unicode.Symbol) {
				fmt.Print("\a")
			} else {
				m.saveName = append(m.saveName, ch)
			}
		case m.choiceIndex == 1 && (key == termbox.KeyBackspace || key == termbox.KeyBackspace2):
			if len(m.seed) == 0 {
				fmt.Print("\a")
			} else {
				m.seed = m.seed[:len(m.seed)-1]
			}
		case m.choiceIndex == 1 && key == termbox.KeySpace:
			m.seed = append(m.seed, ' ')
		case m.choiceIndex == 1 && ch != 0:
			m.seed = append(m.seed, ch)
		default:
			panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
		}

	default:
		panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
	}

	return true
}

func (m *mainMenuUI) inputMouse(x, y int) {
	switch m.state {
	case menuStateMain:
		switch {
		default:
			panic(fmt.Sprintf("%v, %v", x, y))
		}
	default:
		panic(fmt.Sprintf("%v, %v", x, y))
	}
}

func (m *mainMenuUI) loadGame(name string) {
	m.state = menuStateLoad
}

func (m *mainMenuUI) newGame() {
	m.state = menuStateNew
	m.choiceIndex = 0
}
<commit_msg>UI for errors<commit_after>package main

import (
	"compress/gzip"
	"encoding/gob"
	"fmt"
	"github.com/nsf/termbox-go"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"
	"unicode"
)

const SaveDirName = "saves_5CC9DB70-EEC5-47EA-94B6-398BFC12E4A7"

var mainMenu mainMenuUI

func init() {
	mainMenu.rand = rand.New(rand.NewSource(time.Now().UnixNano()))

	for i := 0; i < 16; i++ {
		const characters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
		mainMenu.seed = append(mainMenu.seed, rune(characters[mainMenu.rand.Intn(len(characters))]))
	}

	saves, err := os.Open(SaveDirName)
	if err != nil {
		if os.IsNotExist(err) {
			return
		}
		panic(err)
	}
	names, err := saves.Readdirnames(0)
	if err != nil {
		panic(err)
	}
	for _, name := range names {
		if strings.HasSuffix(name, ".sav") {
			mainMenu.saveNames = append(mainMenu.saveNames, name[:len(name)-len(".sav")])
		}
	}
}

const (
	menuStateMain = iota
	menuStateError
	menuStateNew
)

type mainMenuUI struct {
	rand        *rand.Rand
	saveNames   []string
	choiceIndex int
	state       uint
	saveName    []rune
	seed        []rune
	err         string
}

func (m *mainMenuUI) render(w, h int) {
	m.drawTextFlicker(w, 2, "5CC9DB70-EEC5-47EA-94B6-398BFC12E4A7", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)

	switch m.state {
	case menuStateMain:
		skip := 0

		if m.choiceIndex > h/2-5 {
			skip = m.choiceIndex - (h/2 - 5)
		}

		for i, name := range m.saveNames[skip:] {
			if m.choiceIndex == i+skip {
				m.drawText(w, 5+i, fmt.Sprintf("Load %q", name), termbox.ColorBlack, termbox.ColorWhite)
			} else {
				m.drawText(w, 5+i, fmt.Sprintf("Load %q", name), termbox.ColorWhite, termbox.ColorBlack)
			}
		}
		if m.choiceIndex == len(m.saveNames) {
			m.drawText(w, len(m.saveNames)+5-skip, "New Game", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, len(m.saveNames)+5-skip, "New Game", termbox.ColorWhite, termbox.ColorBlack)
		}

	case menuStateError:
		m.drawText(w, 5, m.err, termbox.ColorRed, termbox.ColorBlack)

	case menuStateNew:
		m.drawText(w, 5, "Save Name", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)
		if m.choiceIndex == 0 {
			m.drawText(w, 6, string(m.saveName)+"_", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, 6, string(m.saveName), termbox.ColorWhite, termbox.ColorBlack)
		}

		m.drawText(w, 8, "Seed", termbox.ColorWhite|termbox.AttrBold, termbox.ColorBlack)
		if m.choiceIndex == 1 {
			m.drawText(w, 9, string(m.seed)+"_", termbox.ColorBlack, termbox.ColorWhite)
		} else {
			m.drawText(w, 9, string(m.seed), termbox.ColorWhite, termbox.ColorBlack)
		}
	}
}

func (m *mainMenuUI) drawText(w, y int, s string, fg, bg termbox.Attribute) {
	r := []rune(s)

	for i := 0; i < w; i++ {
		termbox.SetCell(i, y, ' ', fg, bg)
	}

	for i, ch := range r {
		termbox.SetCell((w-len(r))/2+i, y, ch, fg, bg)
	}
}

func (m *mainMenuUI) drawTextFlicker(w, y int, s string, fg, bg termbox.Attribute) {
	r := []rune(s)

	flicker := m.rand.Intn(len(r))
	if (r[flicker] < '0' || r[flicker] > '9') && (r[flicker] < 'A' || r[flicker] > 'F') {
		flicker = -1
	}

	for i := 0; i < w; i++ {
		termbox.SetCell(i, y, ' ', fg, bg)
	}

	for i, ch := range r {
		if i == flicker {
			const hex = "0123456789ABCDEF"
			termbox.SetCell((w-len(r))/2+i, y, rune(hex[m.rand.Intn(len(hex))]), fg^termbox.AttrBold, bg)
		} else {
			termbox.SetCell((w-len(r))/2+i, y, ch, fg, bg)
		}
	}
}

func (m *mainMenuUI) inputKey(key termbox.Key, ch rune, mod termbox.Modifier) bool {
	switch m.state {
	case menuStateMain:
		switch {
		case key == termbox.KeyArrowDown:
			m.choiceIndex = (m.choiceIndex + 1) % (len(m.saveNames) + 1)
		case key == termbox.KeyArrowUp:
			m.choiceIndex = (m.choiceIndex + (len(m.saveNames) + 1 - 1)) % (len(m.saveNames) + 1)
		case key == termbox.KeyArrowLeft || key == termbox.KeyArrowRight:
			// silently ignore
		case key == termbox.KeyEnter:
			if m.choiceIndex < len(m.saveNames) {
				m.loadGame(m.saveNames[m.choiceIndex])
			} else {
				switch m.choiceIndex - len(m.saveNames) {
				case 0:
					m.newGame()
				}
			}
		case key == termbox.KeyEsc:
			return false
		default:
			panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
		}

	case menuStateError:
		switch {
		case key == termbox.KeyEsc:
			m.err = ""
			m.state = menuStateMain
		default:
			panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
		}

	case menuStateNew:
		const fieldCount = 2

		switch {
		case key == termbox.KeyEsc:
			m.state = menuStateMain
			m.choiceIndex = len(m.saveNames)
		case key == termbox.KeyEnter:
			m.choiceIndex++
			if m.choiceIndex >= fieldCount {
				m.choiceIndex = fieldCount - 1
				if len(m.saveName) == 0 {
					m.choiceIndex = 0
					fmt.Print("\a")
					return true
				}
				w := &World{
					Seed:     NewSeed(string(m.seed)),
					saveName: filepath.Join(SaveDirName, string(m.saveName)+".sav"),
				}
				err := w.save(os.O_CREATE | os.O_EXCL | os.O_WRONLY)
				if err == nil {
					err = w.AfterLoad()
				}
				if err == nil {
					worldLock.Lock()
					world = w
					worldLock.Unlock()
				} else {
					m.err = err.Error()
					m.state = menuStateError
				}
			}
		case key == termbox.KeyArrowDown:
			m.choiceIndex = (m.choiceIndex + 1) % fieldCount
		case key == termbox.KeyArrowUp:
			m.choiceIndex = (m.choiceIndex + (fieldCount - 1)) % fieldCount
		case key == termbox.KeyArrowLeft || key == termbox.KeyArrowRight:
			fmt.Print("\a")
		case m.choiceIndex == 0 && (key == termbox.KeyBackspace || key == termbox.KeyBackspace2):
			if len(m.saveName) == 0 {
				fmt.Print("\a")
			} else {
				m.saveName = m.saveName[:len(m.saveName)-1]
			}
		case m.choiceIndex == 0 && key == termbox.KeySpace:
			if len(m.saveName) == 0 {
				fmt.Print("\a")
			} else {
				m.saveName = append(m.saveName, ' ')
			}
		case m.choiceIndex == 0 && ch != 0:
			// unicode.Punctuation is not included due to characters like /
			if ch != '_' && ch != '-' && !unicode.In(ch, unicode.Letter, unicode.Number, unicode.Symbol) {
				fmt.Print("\a")
			} else {
				m.saveName = append(m.saveName, ch)
			}
		case m.choiceIndex == 1 && (key == termbox.KeyBackspace || key == termbox.KeyBackspace2):
			if len(m.seed) == 0 {
				fmt.Print("\a")
			} else {
				m.seed = m.seed[:len(m.seed)-1]
			}
		case m.choiceIndex == 1 && key == termbox.KeySpace:
			m.seed = append(m.seed, ' ')
		case m.choiceIndex == 1 && ch != 0:
			m.seed = append(m.seed, ch)
		default:
			panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
		}

	default:
		panic(fmt.Sprintf("%v, %v, %v", key, ch, mod))
	}

	return true
}

func (m *mainMenuUI) inputMouse(x, y int) {
	switch m.state {
	case menuStateMain:
		switch {
		default:
			panic(fmt.Sprintf("%v, %v", x, y))
		}
	default:
		panic(fmt.Sprintf("%v, %v", x, y))
	}
}

func (m *mainMenuUI) loadGame(name string) {
	var w World
	f, err := os.Open(filepath.Join(SaveDirName, name+".sav"))
	if err != nil {
		m.err = err.Error()
		m.state = menuStateError
		return
	}
	defer f.Close()
	g, err := gzip.NewReader(f)
	if err != nil {
		m.err = err.Error()
		m.state = menuStateError
		return
	}
	defer g.Close()
	err = gob.NewDecoder(g).Decode(&w)
	if err != nil {
		m.err = err.Error()
		m.state = menuStateError
		return
	}
	err = w.AfterLoad()
	if err != nil {
		m.err = err.Error()
		m.state = menuStateError
		return
	}
	worldLock.Lock()
	world = &w
	worldLock.Unlock()
}

func (m *mainMenuUI) newGame() {
	m.state = menuStateNew
	m.choiceIndex = 0
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"testing"
)

func TestEnd2End(t *testing.T) {
	config, err := getConfig("test.json")
	if err != nil {
		t.Fatal("Open config error: ", err.Error())
	}
	errors := make(chan *RequestError, config.getSize())
	execute(config, errors)
	if len(errors) != 0 {
		for v := range errors {
			log.Println(v.toString())
		}
		t.Error("Expected 0 errors, actual: ", len(errors))
	}
}
<commit_msg>Temporarily disable end2end test<commit_after>package main

/*func TestEnd2End(t *testing.T) {
	config, err := getConfig("test.json")
	if err != nil {
		t.Fatal("Open config error: ", err.Error())
	}
	errors := make(chan *RequestError, config.getSize())
	execute(config, errors)
	if len(errors) != 0 {
		for v := range errors {
			log.Println(v.toString())
		}
		t.Error("Expected 0 errors, actual: ", len(errors))
	}
}*/
<|endoftext|>


<commit_before>package main_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
)

func containsString(list []string, expected string) bool {
	for _, current := range list {
		if current == expected {
			return true
		}
	}
	return false
}

var _ = Describe("Main", func() {
	It("Makes a http get request to each url in a file", func() {
		port := 8000
		list := `http://127.0.0.1:8000/A
http://127.0.0.1:8000/B
http://127.0.0.1:8000/C
		`

		urlsVisited := []string{}

		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			urlsVisited = append(urlsVisited, r.URL.Path)
		})

		server := httptest.NewUnstartedServer(handler)
		server.Listener, _ = net.Listen("tcp", ":"+strconv.Itoa(port))
		server.Start()
		defer server.Close()
		file, err := ioutil.TempFile(os.TempDir(), "prefix")
		if err != nil {
			panic(err)
		}
		file.WriteString(list)
		defer os.Remove(file.Name())

		exePath, err := filepath.Abs("./code-named-something")
		if err != nil {
			panic(err)
		}
		cmd := exec.Command(exePath, "-f", file.Name())
		_, _ = cmd.CombinedOutput()

		client := &http.Client{}
		req, _ := http.NewRequest("GET", "http://127.0.0.1:8000/A", nil)
		_, err = client.Do(req)
		if err != nil {
			fmt.Printf("%v", err)
			panic(err)
		}

		Expect(err).To(BeNil())
		Expect(containsString(urlsVisited, "/A")).To(Equal(true))
		Expect(containsString(urlsVisited, "/B")).To(Equal(true))
		Expect(containsString(urlsVisited, "/C")).To(Equal(true))
	})
})
<commit_msg>Adding list files again<commit_after>package main_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
)

func containsString(list []string, expected string) bool {
	for _, current := range list {
		if current == expected {
			return true
		}
	}
	return false
}

var _ = Describe("Main", func() {
	It("Makes a http get request to each url in a file", func() {
		port := 8000
		list := `http://127.0.0.1:8000/A
http://127.0.0.1:8000/B
http://127.0.0.1:8000/C
		`
		files, _ := ioutil.ReadDir("./")
		for _, f := range files {
			fmt.Println(f.Name())
		}
		urlsVisited := []string{}

		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			urlsVisited = append(urlsVisited, r.URL.Path)
		})

		server := httptest.NewUnstartedServer(handler)
		server.Listener, _ = net.Listen("tcp", ":"+strconv.Itoa(port))
		server.Start()
		defer server.Close()
		file, err := ioutil.TempFile(os.TempDir(), "prefix")
		if err != nil {
			panic(err)
		}
		file.WriteString(list)
		defer os.Remove(file.Name())

		exePath, err := filepath.Abs("./code-named-something")
		if err != nil {
			panic(err)
		}
		cmd := exec.Command(exePath, "-f", file.Name())
		_, _ = cmd.CombinedOutput()

		client := &http.Client{}
		req, _ := http.NewRequest("GET", "http://127.0.0.1:8000/A", nil)
		_, err = client.Do(req)
		if err != nil {
			fmt.Printf("%v", err)
			panic(err)
		}

		Expect(err).To(BeNil())
		Expect(containsString(urlsVisited, "/A")).To(Equal(true))
		Expect(containsString(urlsVisited, "/B")).To(Equal(true))
		Expect(containsString(urlsVisited, "/C")).To(Equal(true))
	})
})
<|endoftext|>


<commit_before>package main

import "testing"

func TestMain(t *testing.T) {
}

// Generates a TLS certificate for server.
// Based on sample from http://golang.org/src/crypto/tls/generate_cert.go (Jan 30, 2015).
func genCert() {
}
<commit_msg>add final certgen code<commit_after>package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"log"
	"math/big"
	"net"
	"testing"
	"time"
)

func TestMain(t *testing.T) {
}

// Generate a self-signed X.509 certificate for a TLS server.
// Based on sample from http://golang.org/src/crypto/tls/generate_cert.go (Jan 30, 2015).
func genCert() x509.Certificate {
	hosts := []string{"nbusy.com"}
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	notBefore := time.Now()
	notAfter := notBefore.Add(290 * 365 * 24 * time.Hour) //290 years
	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("failed to generate serial number: %s", err)
	}

	cert := x509.Certificate{
		IsCA:         true,
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"NBusy"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			cert.IPAddresses = append(cert.IPAddresses, ip)
		} else {
			cert.DNSNames = append(cert.DNSNames, h)
		}
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &cert, &cert, privKey.PublicKey, privKey)
	if err != nil {
		log.Fatalf("Failed to create certificate: %s", err)
	}

	return cert
}
<|endoftext|>


<commit_before>package w32uiautomation

import "time"

func WaitFindFirstWithBreadthFirstSearch(auto *IUIAutomation, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	for {
		found, err := FindFirstWithBreadthFirstSearch(auto, start, matcher)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, err
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func FindFirstWithBreadthFirstSearch(auto *IUIAutomation, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	walker, err := NewTreeWalker(auto)
	if err != nil {
		return nil, err
	}
	defer walker.Release()

	return findFirstWithBreadthFirstSearchHelper(walker, start, matcher)
}

func findFirstWithBreadthFirstSearchHelper(walker *IUIAutomationTreeWalker, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	//fmt.Printf("findFirstWithBreadthFirstSearchHelper start=%v\n", start)
	element, err := findFirstInSiblings(walker, start, matcher)
	if err != nil {
		return nil, err
	}
	if element != nil {
		//fmt.Printf("findFirstWithBreadthFirstSearchHelper found#1=%v\n", element)
		return element, nil
	}

	for parent := start; parent != nil; {
		child, err := walker.GetFirstChildElement(parent)
		if err != nil {
			return nil, err
		}

		if child != nil {
			//fmt.Printf("findFirstWithBreadthFirstSearchHelper first chid=%v\n", child)
			element, err := findFirstWithBreadthFirstSearchHelper(walker, child, matcher)
			if err != nil {
				return nil, err
			}
			if element != nil {
				//fmt.Printf("findFirstWithBreadthFirstSearchHelper found#2=%v\n", element)
				return element, nil
			}
		}

		parent, err = walker.GetNextSiblingElement(parent)
		if err != nil {
			return nil, err
		}
		//fmt.Printf("findFirstWithBreadthFirstSearchHelper next parent=%v\n", parent)
	}
	//fmt.Printf("findFirstWithBreadthFirstSearchHelper exiting\n")
	return nil, nil
}

func findFirstInSiblings(walker *IUIAutomationTreeWalker, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	element := start
	for element != nil {
		matched, err := matcher(element, walker)
		if err != nil {
			return nil, err
		}
		if matched {
			return element, nil
		}

		element, err = walker.GetNextSiblingElement(element)
		if err != nil {
			return nil, err
		}
		//fmt.Printf("findFirstInSiblings next sibling=%v\n", element)
	}
	return nil, nil
}
<commit_msg>Remove comment-outed debug prints<commit_after>package w32uiautomation

import "time"

func WaitFindFirstWithBreadthFirstSearch(auto *IUIAutomation, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	for {
		found, err := FindFirstWithBreadthFirstSearch(auto, start, matcher)
		if err != nil {
			return nil, err
		}
		if found != nil {
			return found, err
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func FindFirstWithBreadthFirstSearch(auto *IUIAutomation, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	walker, err := NewTreeWalker(auto)
	if err != nil {
		return nil, err
	}
	defer walker.Release()

	return findFirstWithBreadthFirstSearchHelper(walker, start, matcher)
}

func findFirstWithBreadthFirstSearchHelper(walker *IUIAutomationTreeWalker, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	element, err := findFirstInSiblings(walker, start, matcher)
	if err != nil {
		return nil, err
	}
	if element != nil {
		return element, nil
	}

	for parent := start; parent != nil; {
		child, err := walker.GetFirstChildElement(parent)
		if err != nil {
			return nil, err
		}

		if child != nil {
			element, err := findFirstWithBreadthFirstSearchHelper(walker, child, matcher)
			if err != nil {
				return nil, err
			}
			if element != nil {
				return element, nil
			}
		}

		parent, err = walker.GetNextSiblingElement(parent)
		if err != nil {
			return nil, err
		}
	}
	return nil, nil
}

func findFirstInSiblings(walker *IUIAutomationTreeWalker, start *IUIAutomationElement, matcher ElemMatcherFunc) (*IUIAutomationElement, error) {
	element := start
	for element != nil {
		matched, err := matcher(element, walker)
		if err != nil {
			return nil, err
		}
		if matched {
			return element, nil
		}

		element, err = walker.GetNextSiblingElement(element)
		if err != nil {
			return nil, err
		}
	}
	return nil, nil
}
<|endoftext|>