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

package main

import (
	"fmt"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/fgrehm/boom-curl/boom/boomer"
)

const (
	headerRegexp = "^([\\w-]+):\\s*(.+)"
	authRegexp   = "^([\\w-\\.]+):(.+)"
)

func main() {
	cli.AppHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}
USAGE:
   {{.Name}} {{if .Flags}}[options] {{end}}<URL>
VERSION:
   {{.Version}}{{if or .Author .Email}}
AUTHOR:{{if .Author}}
  {{.Author}}{{if .Email}} - <{{.Email}}>{{end}}{{else}}
  {{.Email}}{{end}}{{end}}
GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}
`

	app := cli.NewApp()
	app.Name = "bcurl"
	app.Usage = "A cURL like interface for Boom, an HTTP(S) load generator, ApacheBench (ab) replacement"
	app.Author = "Fábio Rehm"
	app.Email = "fgrehm@gmail.com"
	app.Version = VERSION
	app.EnableBashCompletion = true
	app.Flags = []cli.Flag{
		cli.StringSliceFlag{
			Name:  "H, header",
			Value: &cli.StringSlice{},
			Usage: "custom header to pass to server",
		},
		cli.StringFlag{
			Name:  "d, data",
			Value: "",
			Usage: "HTTP POST data",
		},
		cli.IntFlag{
			Name:  "cpus",
			Value: runtime.NumCPU(),
			Usage: fmt.Sprintf("Number of used cpu cores. (default for current machine is %d cores)", runtime.NumCPU()),
		},
		cli.IntFlag{
			Name:  "requests",
			Value: 200,
			Usage: "Number of requests to run",
		},
		cli.IntFlag{
			Name:  "concurrency",
			Value: 50,
			Usage: "Number of requests to run concurrently.",
		},
	}
	app.Action = func(c *cli.Context) {
		if len(c.Args()) != 1 {
			usageAndExit(c, "")
		}

		boom(c)
	}

	// cURL accepts the URL as the first argument but the cli package we are
	// using will blow up since it will recognize it as a subcommand that does
	// not exist, that's why we need to move things around
	if len(os.Args) > 2 && !strings.HasPrefix(os.Args[1], "-") {
		argsWithoutUrl := append([]string{os.Args[0]}, os.Args[2:]...)
		os.Args = append(argsWithoutUrl, os.Args[1])
	}
	app.Run(os.Args)
}

func boom(c *cli.Context) {
	body := c.String("data")
	var m string
	if body != "" {
		m = "POST"
	} else {
		m = "GET"
	}
	cpus := c.Int("cpus")
	num := c.Int("requests")
	conc := c.Int("concurrency")

	q := 0
	t := 0
	contentType := "text/html"
	if m == "POST" {
		contentType = "application/x-www-form-urlencoded"
	}
	output := ""
	insecure := false

	// MAGIC HAPPENS BELOW
	runtime.GOMAXPROCS(cpus)

	if num <= 0 || conc <= 0 {
		usageAndExit(c, "n and c cannot be smaller than 1.")
	}

	var (
		url, method string
		header http.Header = make(http.Header)
	)

	method = strings.ToUpper(m)
	url = c.Args()[0]

	// Prepend http protocol if user did not provide it
	if !regexp.MustCompile(`^https?://`).MatchString(url) {
		url = "http://" + url
	}

	// set content-type
	header.Set("User-Agent", "boom-curl/"+VERSION)
	header.Set("Content-Type", contentType)
	for _, h := range c.StringSlice("H") {
		headerAndValue := strings.SplitAfterN(h, ":", 2)
		h := strings.TrimSuffix(headerAndValue[0], ":")
		value := strings.Trim(strings.Trim(headerAndValue[1], " "), "\"")

		header.Set(h, value)
	}

	if output != "csv" && output != "" {
		usageAndExit(c, "Invalid output type.")
	}

	(&boomer.Boomer{
		Req: &boomer.ReqOpts{
			Method: method,
			URL:    url,
			Body:   body,
			Header: header,
		},
		N:             num,
		C:             conc,
		Qps:           q,
		Timeout:       t,
		AllowInsecure: insecure,
		Output:        output,
	}).Run()
}

func usageAndExit(c *cli.Context, message string) {
	if message != "" {
		fmt.Fprintf(os.Stderr, message)
		fmt.Fprintf(os.Stderr, "\n\n")
	}
	cli.ShowAppHelp(c)
	os.Exit(1)
}
<commit_msg>Reduce concurrency level to 4<commit_after>// Copyright 2014 Fabio Rehm. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/fgrehm/boom-curl/boom/boomer"
)

const (
	headerRegexp = "^([\\w-]+):\\s*(.+)"
	authRegexp   = "^([\\w-\\.]+):(.+)"
)

func main() {
	cli.AppHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}
USAGE:
   {{.Name}} {{if .Flags}}[options] {{end}}<URL>
VERSION:
   {{.Version}}{{if or .Author .Email}}
AUTHOR:{{if .Author}}
  {{.Author}}{{if .Email}} - <{{.Email}}>{{end}}{{else}}
  {{.Email}}{{end}}{{end}}
GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}
`

	app := cli.NewApp()
	app.Name = "bcurl"
	app.Usage = "A cURL like interface for Boom, an HTTP(S) load generator, ApacheBench (ab) replacement"
	app.Author = "Fábio Rehm"
	app.Email = "fgrehm@gmail.com"
	app.Version = VERSION
	app.EnableBashCompletion = true
	app.Flags = []cli.Flag{
		cli.StringSliceFlag{
			Name:  "H, header",
			Value: &cli.StringSlice{},
			Usage: "custom header to pass to server",
		},
		cli.StringFlag{
			Name:  "d, data",
			Value: "",
			Usage: "HTTP POST data",
		},
		cli.IntFlag{
			Name:  "cpus",
			Value: runtime.NumCPU(),
			Usage: fmt.Sprintf("Number of used cpu cores. (default for current machine is %d cores)", runtime.NumCPU()),
		},
		cli.IntFlag{
			Name:  "requests",
			Value: 200,
			Usage: "Number of requests to run",
		},
		cli.IntFlag{
			Name:  "concurrency",
			Value: 4,
			Usage: "Number of requests to run concurrently.",
		},
	}
	app.Action = func(c *cli.Context) {
		if len(c.Args()) != 1 {
			usageAndExit(c, "")
		}

		boom(c)
	}

	// cURL accepts the URL as the first argument but the cli package we are
	// using will blow up since it will recognize it as a subcommand that does
	// not exist, that's why we need to move things around
	if len(os.Args) > 2 && !strings.HasPrefix(os.Args[1], "-") {
		argsWithoutUrl := append([]string{os.Args[0]}, os.Args[2:]...)
		os.Args = append(argsWithoutUrl, os.Args[1])
	}
	app.Run(os.Args)
}

func boom(c *cli.Context) {
	body := c.String("data")
	var m string
	if body != "" {
		m = "POST"
	} else {
		m = "GET"
	}
	cpus := c.Int("cpus")
	num := c.Int("requests")
	conc := c.Int("concurrency")

	q := 0
	t := 0
	contentType := "text/html"
	if m == "POST" {
		contentType = "application/x-www-form-urlencoded"
	}
	output := ""
	insecure := false

	// MAGIC HAPPENS BELOW
	runtime.GOMAXPROCS(cpus)

	if num <= 0 || conc <= 0 {
		usageAndExit(c, "n and c cannot be smaller than 1.")
	}

	var (
		url, method string
		header http.Header = make(http.Header)
	)

	method = strings.ToUpper(m)
	url = c.Args()[0]

	// Prepend http protocol if user did not provide it
	if !regexp.MustCompile(`^https?://`).MatchString(url) {
		url = "http://" + url
	}

	// set content-type
	header.Set("User-Agent", "boom-curl/"+VERSION)
	header.Set("Content-Type", contentType)
	for _, h := range c.StringSlice("H") {
		headerAndValue := strings.SplitAfterN(h, ":", 2)
		h := strings.TrimSuffix(headerAndValue[0], ":")
		value := strings.Trim(strings.Trim(headerAndValue[1], " "), "\"")

		header.Set(h, value)
	}

	if output != "csv" && output != "" {
		usageAndExit(c, "Invalid output type.")
	}

	(&boomer.Boomer{
		Req: &boomer.ReqOpts{
			Method: method,
			URL:    url,
			Body:   body,
			Header: header,
		},
		N:             num,
		C:             conc,
		Qps:           q,
		Timeout:       t,
		AllowInsecure: insecure,
		Output:        output,
	}).Run()
}

func usageAndExit(c *cli.Context, message string) {
	if message != "" {
		fmt.Fprintf(os.Stderr, message)
		fmt.Fprintf(os.Stderr, "\n\n")
	}
	cli.ShowAppHelp(c)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package hive

import (
	"bufio"
	"errors"
	// "fmt"
	"github.com/eaciit/errorlib"
	"io"
	"log"
	"os/exec"
	"strings"
)

const (
	BEE_CLI_STR  = "0: jdbc:hive2:"
	CLOSE_SCRIPT = "!quit"
)

type DuplexTerm struct {
	Writer     *bufio.Writer
	Reader     *bufio.Reader
	Cmd        *exec.Cmd
	CmdStr     string
	Stdin      io.WriteCloser
	Stdout     io.ReadCloser
	FnReceive  FnHiveReceive
	OutputType string
	DateFormat string
}

/*func (d *DuplexTerm) Open() (e error) {
	if d.Stdin, e = d.Cmd.StdinPipe(); e != nil {
		return
	}

	if d.Stdout, e = d.Cmd.StdoutPipe(); e != nil {
		return
	}

	d.Writer = bufio.NewWriter(d.Stdin)
	d.Reader = bufio.NewReader(d.Stdout)

	e = d.Cmd.Start()
	return
}*/

var hr *HiveResult

func (d *DuplexTerm) Open() (e error) {
	if d.CmdStr != "" {
		arg := append([]string{"-c"}, d.CmdStr)
		d.Cmd = exec.Command("sh", arg...)

		if d.Stdin, e = d.Cmd.StdinPipe(); e != nil {
			return
		}

		if d.Stdout, e = d.Cmd.StdoutPipe(); e != nil {
			return
		}

		d.Writer = bufio.NewWriter(d.Stdin)
		d.Reader = bufio.NewReader(d.Stdout)

		if d.FnReceive != nil {
			go func() {
				_, e = d.Wait()
			}()
		}
		e = d.Cmd.Start()
	} else {
		errorlib.Error("", "", "Open", "The Connection Config not Set")
	}

	return
}

func (d *DuplexTerm) Close() {
	result, e := d.SendInput(CLOSE_SCRIPT)

	_ = result
	_ = e

	d.Cmd.Wait()
	d.Stdin.Close()
	d.Stdout.Close()
}

func (d *DuplexTerm) SendInput(input string) (result []string, e error) {
	iwrite, e := d.Writer.WriteString(input + "\n")
	if iwrite == 0 {
		e = errors.New("Writing only 0 byte")
	} else {
		e = d.Writer.Flush()
	}

	if e != nil {
		return
	}

	if d.FnReceive == nil {
		result, e = d.Wait()
	}

	return
}

func (d *DuplexTerm) Wait() (result []string, e error) {
	isHeader := false
	for {
		peekBefore, _ := d.Reader.Peek(14)
		peekBeforeStr := string(peekBefore)
		// log.Printf("peekBefore: %v\n", peekBefore)
		// log.Printf("peekBeforeStr: %v\n", peekBeforeStr)

		bread, e := d.Reader.ReadString('\n')
		bread = strings.TrimRight(bread, "\n")

		peek, _ := d.Reader.Peek(14)
		peekStr := string(peek)
		// log.Printf("peek: %v\n", peek)
		// log.Printf("peekStr: %v\n", peekStr)

		delimiter := "\t"

		if d.OutputType == CSV {
			delimiter = ","
		}

		if BEE_CLI_STR == peekBeforeStr {
			isHeader = true
		}

		if isHeader {
			hr := HiveResult{}
			hr.constructHeader(bread, delimiter)
			log.Printf("model: %v\n", hr)
			log.Printf("headerStr: %v\n", bread)
			for _, val := range hr.Header {
				log.Printf("header: %v\n", val)
			}

			isHeader = false
		}

		if !strings.Contains(bread, BEE_CLI_STR) {
			//result = append(result, bread)
			if d.FnReceive != nil {
				/*Parse(hr.Header, bread, hr.ResultObj, d.OutputType, d.DateFormat)
				log.Printf("model: %v\n", hr.ResultObj)*/
				d.FnReceive(bread)
			} else {
				result = append(result, bread)
			}
		}

		if d.FnReceive != nil {
			if (e != nil && e.Error() == "EOF") || (strings.Contains(peekStr, CLOSE_SCRIPT)) {
				break
			}
		} else {
			if (e != nil && e.Error() == "EOF") || (BEE_CLI_STR == peekStr) {
				break
			}
		}

	}

	return
}
<commit_msg>bug fixing for exec<commit_after>package hive

import (
	"bufio"
	"errors"
	// "fmt"
	"github.com/eaciit/errorlib"
	"io"
	"log"
	"os/exec"
	"strings"
)

const (
	BEE_CLI_STR  = "0: jdbc:hive2:"
	CLOSE_SCRIPT = "!quit"
)

type DuplexTerm struct {
	Writer     *bufio.Writer
	Reader     *bufio.Reader
	Cmd        *exec.Cmd
	CmdStr     string
	Stdin      io.WriteCloser
	Stdout     io.ReadCloser
	FnReceive  FnHiveReceive
	OutputType string
	DateFormat string
}

/*func (d *DuplexTerm) Open() (e error) {
	if d.Stdin, e = d.Cmd.StdinPipe(); e != nil {
		return
	}

	if d.Stdout, e = d.Cmd.StdoutPipe(); e != nil {
		return
	}

	d.Writer = bufio.NewWriter(d.Stdin)
	d.Reader = bufio.NewReader(d.Stdout)

	e = d.Cmd.Start()
	return
}*/

var hr *HiveResult

func (d *DuplexTerm) Open() (e error) {
	if d.CmdStr != "" {
		arg := append([]string{"-c"}, d.CmdStr)
		d.Cmd = exec.Command("sh", arg...)

		if d.Stdin, e = d.Cmd.StdinPipe(); e != nil {
			return
		}

		if d.Stdout, e = d.Cmd.StdoutPipe(); e != nil {
			return
		}

		d.Writer = bufio.NewWriter(d.Stdin)
		d.Reader = bufio.NewReader(d.Stdout)

		if d.FnReceive != nil {
			go func() {
				_, e = d.Wait()
			}()
		}
		e = d.Cmd.Start()
	} else {
		errorlib.Error("", "", "Open", "The Connection Config not Set")
	}

	return
}

func (d *DuplexTerm) Close() {
	result, e := d.SendInput(CLOSE_SCRIPT)

	_ = result
	_ = e

	d.Cmd.Wait()
	d.Stdin.Close()
	d.Stdout.Close()
}

func (d *DuplexTerm) SendInput(input string) (result []string, e error) {
	iwrite, e := d.Writer.WriteString(input + "\n")
	if iwrite == 0 {
		e = errors.New("Writing only 0 byte")
	} else {
		e = d.Writer.Flush()
	}

	if e != nil {
		return
	}

	if d.FnReceive == nil {
		result, e = d.Wait()
	}

	return
}

func (d *DuplexTerm) Wait() (result []string, e error) {
	isHeader := false
	for {
		peekBefore, _ := d.Reader.Peek(14)
		peekBeforeStr := string(peekBefore)
		// log.Printf("peekBefore: %v\n", peekBefore)
		// log.Printf("peekBeforeStr: %v\n", peekBeforeStr)

		bread, e := d.Reader.ReadString('\n')
		bread = strings.TrimRight(bread, "\n")

		peek, _ := d.Reader.Peek(14)
		peekStr := string(peek)
		// log.Printf("peek: %v\n", peek)
		// log.Printf("peekStr: %v\n", peekStr)

		delimiter := "\t"

		if d.OutputType == CSV {
			delimiter = ","
		}

		if isHeader {
			hr := HiveResult{}
			hr.constructHeader(bread, delimiter)
			log.Printf("model: %v\n", hr)
			log.Printf("headerStr: %v\n", bread)
			for _, val := range hr.Header {
				log.Printf("header: %v\n", val)
			}

			isHeader = false
		}

		if BEE_CLI_STR == peekBeforeStr {
			isHeader = true
		}

		if !strings.Contains(bread, BEE_CLI_STR) {
			//result = append(result, bread)
			if d.FnReceive != nil {
				/*Parse(hr.Header, bread, hr.ResultObj, d.OutputType, d.DateFormat)
				log.Printf("model: %v\n", hr.ResultObj)*/
				d.FnReceive(bread)
			} else {
				result = append(result, bread)
			}
		}

		if d.FnReceive != nil {
			if (e != nil && e.Error() == "EOF") || (strings.Contains(peekStr, CLOSE_SCRIPT)) {
				break
			}
		} else {
			if (e != nil && e.Error() == "EOF") || (BEE_CLI_STR == peekStr) {
				break
			}
		}

	}

	return
}
<|endoftext|>


<commit_before>package fmttab

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"unicode/utf8"
)

//A Border of table
type Border int

//A BorderKind type of element on the border of the table
type BorderKind int

//A Align text alignment in column of the table
type Align bool

const (
	//WidthAuto auto sizing of width column
	WidthAuto = 0
	//BorderNone table without borders
	BorderNone = Border(0)
	//BorderThin table with a thin border
	BorderThin = Border(1)
	//BorderDouble table with a double border
	BorderDouble = Border(2)
	//AlignLeft align text along the left edge
	AlignLeft = Align(false)
	//AlignRight align text along the right edge
	AlignRight = Align(true)
)

//The concrete type of the object on the border of the table
const (
	BKLeftTop BorderKind = iota
	BKRighttop
	BKRightBottom
	BKLeftBottom
	BKLeftToRight
	BKRightToLeft
	BKTopToBottom
	BKBottomToTop
	BKBottomCross
	BKHorizontal
	BKVertical
	BKHorizontalBorder
	BKVerticalBorder
)

const (
	trimend = ".."
)

//Borders predefined border types
var Borders = map[Border]map[BorderKind]string{
	BorderNone: map[BorderKind]string{
		BKVertical: "|",
	},
	BorderThin: map[BorderKind]string{
		BKLeftTop:          "\u250c",
		BKRighttop:         "\u2510",
		BKRightBottom:      "\u2518",
		BKLeftBottom:       "\u2514",
		BKLeftToRight:      "\u251c",
		BKRightToLeft:      "\u2524",
		BKTopToBottom:      "\u252c",
		BKBottomToTop:      "\u2534",
		BKBottomCross:      "\u253c",
		BKHorizontal:       "\u2500",
		BKVertical:         "\u2502",
		BKHorizontalBorder: "\u2500",
		BKVerticalBorder:   "\u2502",
	},
	BorderDouble: map[BorderKind]string{
		BKLeftTop:          "\u2554",
		BKRighttop:         "\u2557",
		BKRightBottom:      "\u255d",
		BKLeftBottom:       "\u255a",
		BKLeftToRight:      "\u255f",
		BKRightToLeft:      "\u2562",
		BKTopToBottom:      "\u2564",
		BKBottomToTop:      "\u2567",
		BKBottomCross:      "\u253c",
		BKHorizontal:       "\u2500",
		BKVertical:         "\u2502",
		BKHorizontalBorder: "\u2550",
		BKVerticalBorder:   "\u2551",
	},
}

//A DataGetter functional type for table data
type DataGetter func() (bool, map[string]interface{})

//A Column type of table columns
type Column struct {
	maxLen int
	Name   string
	Width  int
	Aling  Align
}

//A Table is the repository for the columns, the data that are used for printing the table
type Table struct {
	Columns []*Column
	caption string
	border  Border
	Data    []map[string]interface{}
	dataget DataGetter
}

// A trimEnds supplements the text with special characters by limiting the length of the text column width
func trimEnds(val, end string, max int) string {
	l := utf8.RuneCountInString(val)
	if l <= max {
		return val
	}
	lend := utf8.RuneCountInString(end)
	if lend >= max {
		return end[:max]
	}
	return string([]rune(val)[:(max-lend)]) + end
}

//GetMaskFormat returns a pattern string for formatting text in table column alignment
func (c *Column) GetMaskFormat() string {
	if c.Aling == AlignLeft {
		return "%-" + strconv.Itoa(c.getWidth()) + "v"
	}
	return "%" + strconv.Itoa(c.getWidth()) + "v"
}

//must be calculated before call
func (c *Column) getWidth() int {
	if c.Width == WidthAuto {
		return c.maxLen
	}
	return c.Width
}

//AddColumn adds a column to the table
func (t *Table) AddColumn(name string, width int, aling Align) *Table {
	//TODO: check dublicate
	t.Columns = append(t.Columns, &Column{
		Name:  name,
		Width: width,
		Aling: aling,
	})
	return t
}

//AppendData adds the data to the table
func (t *Table) AppendData(rec map[string]interface{}) *Table {
	t.Data = append(t.Data, rec)
	return t
}

//ClearData removes data from a table
func (t *Table) ClearData() *Table {
	t.Data = nil
	return t
}

//CountData the amount of data in the table
func (t *Table) CountData() int {
	return len(t.Data)
}

func (t *Table) writeHeader(w io.Writer) (int, error) {
	var cntwrite int
	dataout := ""
	if t.caption != "" {
		dataout += t.caption + "\n"
	}
	dataout += Borders[t.border][BKLeftTop]
	cntCols := len(t.Columns)
	for num, c := range t.Columns {
		dataout += strings.Repeat(Borders[t.border][BKHorizontalBorder], c.getWidth())
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKTopToBottom]
		} else {
			delim = Borders[t.border][BKRighttop] + "\n"
		}
		dataout += delim
	}
	dataout += Borders[t.border][BKVerticalBorder]
	for num, c := range t.Columns {
		caption := fmt.Sprintf(c.GetMaskFormat(), c.Name)
		dataout += trimEnds(caption, trimend, c.getWidth())
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKVertical]
		} else {
			delim = Borders[t.border][BKVerticalBorder]
		}
		dataout += delim
	}
	dataout += "\n" + Borders[t.border][BKLeftToRight]
	dataout += t.writeBorderTopButtomData(BKHorizontal, BKBottomCross, BKRightToLeft)
	if n, err := w.Write([]byte(dataout)); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeBorderTopButtomData(hr, vbwnCol, vright BorderKind) (data string) {
	cntCols := len(t.Columns)
	empty := true
	for num, c := range t.Columns {
		s := strings.Repeat(Borders[t.border][hr], c.getWidth())
		if len(s) > 0 {
			empty = false
		}
		data += s
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][vbwnCol]
		} else {
			delim = Borders[t.border][vright]
			if !empty {
				delim += "\n"
			}
		}
		if len(delim) > 0 {
			empty = false
		}
		data += delim
	}
	return data
}

func (t *Table) writeBottomBorder(w io.Writer) (int, error) {
	var cntwrite int
	data := Borders[t.border][BKLeftBottom] + t.writeBorderTopButtomData(BKHorizontalBorder, BKBottomToTop, BKRightBottom)
	if n, err := w.Write([]byte(data)); err == nil {
		cntwrite = n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeRecord(data map[string]interface{}, w io.Writer) (int, error) {
	var cntwrite int
	cntCols := len(t.Columns)
	if n, err := w.Write([]byte(Borders[t.border][BKVerticalBorder])); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	for num, c := range t.Columns {
		val, mok := data[c.Name]
		if !mok || val == nil {
			val = ""
		}
		caption := fmt.Sprintf(c.GetMaskFormat(), val)
		if n, err := w.Write([]byte(trimEnds(caption, trimend, c.getWidth()))); err == nil {
			cntwrite += n
		} else {
			return -1, err
		}
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKVertical]
		} else {
			delim = Borders[t.border][BKVerticalBorder]
		}
		if n, err := w.Write([]byte(delim)); err == nil {
			cntwrite += n
		} else {
			return -1, err
		}
	}
	if n, err := w.Write([]byte("\n")); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeData(w io.Writer) (int, error) {
	var cntwrite int
	if t.dataget != nil {
		for {
			ok, data := t.dataget()
			if !ok {
				break
			}
			if n, err := t.writeRecord(data, w); err == nil {
				cntwrite += n
			} else {
				return -1, err
			}
		}
	} else if t.CountData() != 0 {
		for _, data := range t.Data {
			if n, err := t.writeRecord(data, w); err == nil {
				cntwrite += n
			} else {
				return -1, err
			}
		}
	}
	return cntwrite, nil
}

//New creates a Table object. DataGetter can be nil
func New(caption string, border Border, datagetter DataGetter) *Table {
	return &Table{
		caption: caption,
		border:  border,
		dataget: datagetter,
	}
}

// String returns the contents of the table with borders
// as a string.  If error, it returns "".
func (t *Table) String() string {
	var buf bytes.Buffer
	if _, err := t.WriteTo(&buf); err != nil {
		return ""
	}
	return buf.String()
}

func (t *Table) autoWidth() {
	var wa []*Column
	for i := range t.Columns {
		if t.Columns[i].Width == WidthAuto {
			t.Columns[i].maxLen = len(t.Columns[i].Name)
			wa = append(wa, t.Columns[i])
		}
	}
	if len(wa) == 0 {
		return
	}
	for _, data := range t.Data {
		for i := range wa {
			curval := fmt.Sprintf("%v", data[wa[i].Name])
			curlen := len(curval)
			if curlen > wa[i].maxLen {
				wa[i].maxLen = curlen
			}
		}
	}

	fmt.Println(wa)
}

// WriteTo writes data to w until the buffer is drained or an error occurs.
// The return value n is the number of bytes written; it always fits into an
// int, but it is int64 to match the io.WriterTo interface. Any error
// encountered during the write is also returned.
func (t *Table) WriteTo(w io.Writer) (int64, error) {
	if len(t.Columns) == 0 {
		return 0, nil
	}
	t.autoWidth()
	var cntwrite int64
	if n, err := t.writeHeader(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	if n, err := t.writeData(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	if n, err := t.writeBottomBorder(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	return cntwrite, nil
}
<commit_msg>remove trace messages<commit_after>package fmttab

import (
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"unicode/utf8"
)

//A Border of table
type Border int

//A BorderKind type of element on the border of the table
type BorderKind int

//A Align text alignment in column of the table
type Align bool

const (
	//WidthAuto auto sizing of width column
	WidthAuto = 0
	//BorderNone table without borders
	BorderNone = Border(0)
	//BorderThin table with a thin border
	BorderThin = Border(1)
	//BorderDouble table with a double border
	BorderDouble = Border(2)
	//AlignLeft align text along the left edge
	AlignLeft = Align(false)
	//AlignRight align text along the right edge
	AlignRight = Align(true)
)

//The concrete type of the object on the border of the table
const (
	BKLeftTop BorderKind = iota
	BKRighttop
	BKRightBottom
	BKLeftBottom
	BKLeftToRight
	BKRightToLeft
	BKTopToBottom
	BKBottomToTop
	BKBottomCross
	BKHorizontal
	BKVertical
	BKHorizontalBorder
	BKVerticalBorder
)

const (
	trimend = ".."
)

//Borders predefined border types
var Borders = map[Border]map[BorderKind]string{
	BorderNone: map[BorderKind]string{
		BKVertical: "|",
	},
	BorderThin: map[BorderKind]string{
		BKLeftTop:          "\u250c",
		BKRighttop:         "\u2510",
		BKRightBottom:      "\u2518",
		BKLeftBottom:       "\u2514",
		BKLeftToRight:      "\u251c",
		BKRightToLeft:      "\u2524",
		BKTopToBottom:      "\u252c",
		BKBottomToTop:      "\u2534",
		BKBottomCross:      "\u253c",
		BKHorizontal:       "\u2500",
		BKVertical:         "\u2502",
		BKHorizontalBorder: "\u2500",
		BKVerticalBorder:   "\u2502",
	},
	BorderDouble: map[BorderKind]string{
		BKLeftTop:          "\u2554",
		BKRighttop:         "\u2557",
		BKRightBottom:      "\u255d",
		BKLeftBottom:       "\u255a",
		BKLeftToRight:      "\u255f",
		BKRightToLeft:      "\u2562",
		BKTopToBottom:      "\u2564",
		BKBottomToTop:      "\u2567",
		BKBottomCross:      "\u253c",
		BKHorizontal:       "\u2500",
		BKVertical:         "\u2502",
		BKHorizontalBorder: "\u2550",
		BKVerticalBorder:   "\u2551",
	},
}

//A DataGetter functional type for table data
type DataGetter func() (bool, map[string]interface{})

//A Column type of table columns
type Column struct {
	maxLen int
	Name   string
	Width  int
	Aling  Align
}

//A Table is the repository for the columns, the data that are used for printing the table
type Table struct {
	Columns []*Column
	caption string
	border  Border
	Data    []map[string]interface{}
	dataget DataGetter
}

// A trimEnds supplements the text with special characters by limiting the length of the text column width
func trimEnds(val, end string, max int) string {
	l := utf8.RuneCountInString(val)
	if l <= max {
		return val
	}
	lend := utf8.RuneCountInString(end)
	if lend >= max {
		return end[:max]
	}
	return string([]rune(val)[:(max-lend)]) + end
}

//GetMaskFormat returns a pattern string for formatting text in table column alignment
func (c *Column) GetMaskFormat() string {
	if c.Aling == AlignLeft {
		return "%-" + strconv.Itoa(c.getWidth()) + "v"
	}
	return "%" + strconv.Itoa(c.getWidth()) + "v"
}

//must be calculated before call
func (c *Column) getWidth() int {
	if c.Width == WidthAuto {
		return c.maxLen
	}
	return c.Width
}

//AddColumn adds a column to the table
func (t *Table) AddColumn(name string, width int, aling Align) *Table {
	//TODO: check dublicate
	t.Columns = append(t.Columns, &Column{
		Name:  name,
		Width: width,
		Aling: aling,
	})
	return t
}

//AppendData adds the data to the table
func (t *Table) AppendData(rec map[string]interface{}) *Table {
	t.Data = append(t.Data, rec)
	return t
}

//ClearData removes data from a table
func (t *Table) ClearData() *Table {
	t.Data = nil
	return t
}

//CountData the amount of data in the table
func (t *Table) CountData() int {
	return len(t.Data)
}

func (t *Table) writeHeader(w io.Writer) (int, error) {
	var cntwrite int
	dataout := ""
	if t.caption != "" {
		dataout += t.caption + "\n"
	}
	dataout += Borders[t.border][BKLeftTop]
	cntCols := len(t.Columns)
	for num, c := range t.Columns {
		dataout += strings.Repeat(Borders[t.border][BKHorizontalBorder], c.getWidth())
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKTopToBottom]
		} else {
			delim = Borders[t.border][BKRighttop] + "\n"
		}
		dataout += delim
	}
	dataout += Borders[t.border][BKVerticalBorder]
	for num, c := range t.Columns {
		caption := fmt.Sprintf(c.GetMaskFormat(), c.Name)
		dataout += trimEnds(caption, trimend, c.getWidth())
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKVertical]
		} else {
			delim = Borders[t.border][BKVerticalBorder]
		}
		dataout += delim
	}
	dataout += "\n" + Borders[t.border][BKLeftToRight]
	dataout += t.writeBorderTopButtomData(BKHorizontal, BKBottomCross, BKRightToLeft)
	if n, err := w.Write([]byte(dataout)); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeBorderTopButtomData(hr, vbwnCol, vright BorderKind) (data string) {
	cntCols := len(t.Columns)
	empty := true
	for num, c := range t.Columns {
		s := strings.Repeat(Borders[t.border][hr], c.getWidth())
		if len(s) > 0 {
			empty = false
		}
		data += s
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][vbwnCol]
		} else {
			delim = Borders[t.border][vright]
			if !empty {
				delim += "\n"
			}
		}
		if len(delim) > 0 {
			empty = false
		}
		data += delim
	}
	return data
}

func (t *Table) writeBottomBorder(w io.Writer) (int, error) {
	var cntwrite int
	data := Borders[t.border][BKLeftBottom] + t.writeBorderTopButtomData(BKHorizontalBorder, BKBottomToTop, BKRightBottom)
	if n, err := w.Write([]byte(data)); err == nil {
		cntwrite = n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeRecord(data map[string]interface{}, w io.Writer) (int, error) {
	var cntwrite int
	cntCols := len(t.Columns)
	if n, err := w.Write([]byte(Borders[t.border][BKVerticalBorder])); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	for num, c := range t.Columns {
		val, mok := data[c.Name]
		if !mok || val == nil {
			val = ""
		}
		caption := fmt.Sprintf(c.GetMaskFormat(), val)
		if n, err := w.Write([]byte(trimEnds(caption, trimend, c.getWidth()))); err == nil {
			cntwrite += n
		} else {
			return -1, err
		}
		var delim string
		if num < cntCols-1 {
			delim = Borders[t.border][BKVertical]
		} else {
			delim = Borders[t.border][BKVerticalBorder]
		}
		if n, err := w.Write([]byte(delim)); err == nil {
			cntwrite += n
		} else {
			return -1, err
		}
	}
	if n, err := w.Write([]byte("\n")); err == nil {
		cntwrite += n
	} else {
		return -1, err
	}
	return cntwrite, nil
}

func (t *Table) writeData(w io.Writer) (int, error) {
	var cntwrite int
	if t.dataget != nil {
		for {
			ok, data := t.dataget()
			if !ok {
				break
			}
			if n, err := t.writeRecord(data, w); err == nil {
				cntwrite += n
			} else {
				return -1, err
			}
		}
	} else if t.CountData() != 0 {
		for _, data := range t.Data {
			if n, err := t.writeRecord(data, w); err == nil {
				cntwrite += n
			} else {
				return -1, err
			}
		}
	}
	return cntwrite, nil
}

//New creates a Table object. DataGetter can be nil
func New(caption string, border Border, datagetter DataGetter) *Table {
	return &Table{
		caption: caption,
		border:  border,
		dataget: datagetter,
	}
}

// String returns the contents of the table with borders
// as a string.  If error, it returns "".
func (t *Table) String() string {
	var buf bytes.Buffer
	if _, err := t.WriteTo(&buf); err != nil {
		return ""
	}
	return buf.String()
}

func (t *Table) autoWidth() {
	var wa []*Column
	for i := range t.Columns {
		if t.Columns[i].Width == WidthAuto {
			t.Columns[i].maxLen = len(t.Columns[i].Name)
			wa = append(wa, t.Columns[i])
		}
	}
	if len(wa) == 0 {
		return
	}
	for _, data := range t.Data {
		for i := range wa {
			curval := fmt.Sprintf("%v", data[wa[i].Name])
			curlen := len(curval)
			if curlen > wa[i].maxLen {
				wa[i].maxLen = curlen
			}
		}
	}
}

// WriteTo writes data to w until the buffer is drained or an error occurs.
// The return value n is the number of bytes written; it always fits into an
// int, but it is int64 to match the io.WriterTo interface. Any error
// encountered during the write is also returned.
func (t *Table) WriteTo(w io.Writer) (int64, error) {
	if len(t.Columns) == 0 {
		return 0, nil
	}
	t.autoWidth()
	var cntwrite int64
	if n, err := t.writeHeader(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	if n, err := t.writeData(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	if n, err := t.writeBottomBorder(w); err == nil {
		cntwrite += int64(n)
	} else {
		return -1, err
	}
	return cntwrite, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bitmm/bitfinex"
	"fmt"
	"math"
	"os"
	"os/exec"
	"time"
)

// Trade inputs
const (
	SYMBOL    = "ltcusd" // Instrument to trade
	MINCHANGE = 0.0001   // Minumum change required to update prices
	TRADENUM  = 20       // Number of trades to use in calculations
	MAXO      = 100      // Max order size
	MINO      = 0.011    // Min order size
	INEDGE    = 0.04     // Required entry edge
	OUTEDGE   = 0.01     // Required exit edge
)

var (
	api        = bitfinex.New(os.Getenv("BITFINEX_KEY"), os.Getenv("BITFINEX_SECRET"))
	apiErrors  = false
	liveOrders = false
	orderTheo  = 0.0
)

func main() {
	fmt.Println("\nInitializing...")

	// Check for input to break loop
	inputChan := make(chan rune)
	go checkStdin(inputChan)

	// Run loop until user input is received
	runMainLoop(inputChan)
}

// Check for any user input
func checkStdin(inputChan chan<- rune) {
	var ch rune
	fmt.Scanf("%c", &ch)
	inputChan <- ch
}

// Infinite loop
func runMainLoop(inputChan <-chan rune) {
	// Exchange communication channels
	// bookChan := make(chan bitfinex.Book)
	// tradesChan := make(chan bitfinex.Trades)

	var (
		trades bitfinex.Trades
		// book        bitfinex.Book
		orders      bitfinex.Orders
		start       time.Time
		oldPosition float64
		newPosition float64
		theo        float64
		lastTrade   int
	)

	for {
		// Record time for each iteration
		start = time.Now()

		// Exit if anything entered by user
		select {
		case <-inputChan:
			exit()
			return
		default:
		}

		// // Get data in separate goroutines
		// go processTrades(tradesChan)
		// // go processBook(bookChan)
		//
		// // Possibly send orders when trades data returns
		// trades = <-tradesChan

		trades = processTrades()
		if !apiErrors && trades[0].TID != lastTrade { // If new trades
			theo = calculateTheo(trades)
			newPosition = checkPosition()
			if (math.Abs(theo-orderTheo) >= MINCHANGE || math.Abs(oldPosition-
				newPosition) >= MINO || !liveOrders) && !apiErrors {
				orders = sendOrders(theo, newPosition)
			}
		}

		// Print results when book and order data returns
		// book = <-bookChan
		if !apiErrors {
			// printResults(book, trades, orders, theo, newPosition, start)
			printResults(trades, orders, theo, newPosition, start)
			// Reset for next iteration
			oldPosition = newPosition
			lastTrade = trades[0].TID
		}

		// Reset for next iteration
		apiErrors = false
	}
}

// Send orders to the exchange
func sendOrders(theo, position float64) bitfinex.Orders {
	orderTheo = theo

	if liveOrders {
		cancelAll()
	}

	// Send new order request to the exchange
	params := calcOrderParams(position, theo)
	orders, err := api.MultipleNewOrders(params)
	liveOrders = true
	checkErr(err)
	return orders
}

func calcOrderParams(position, theo float64) []bitfinex.OrderParams {
	var params []bitfinex.OrderParams

	if math.Abs(position) < MINO { // No position
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, MAXO, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	} else if position < (-1*MAXO)+MINO { // Max short postion
		params = []bitfinex.OrderParams{
			{SYMBOL, -1 * position, theo - OUTEDGE, "bitfinex", "buy", "limit"},
		}
	} else if position > MAXO-MINO { // Max long postion
		params = []bitfinex.OrderParams{
			{SYMBOL, position, theo + OUTEDGE, "bitfinex", "sell", "limit"},
		}
	} else if (-1*MAXO)+MINO <= position && position <= -1*MINO { // Partial short
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, -1 * position, theo - OUTEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, MAXO + position, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	} else if MINO <= position && position <= MAXO-MINO { // Partial long
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO - position, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, position, theo + OUTEDGE, "bitfinex", "sell", "limit"},
			{SYMBOL, MAXO, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	}

	return params
}

func checkPosition() float64 {
	var position float64
	posSlice, err := api.ActivePositions()
	checkErr(err)
	for _, pos := range posSlice {
		if pos.Symbol == SYMBOL {
			position = pos.Amount
		}
	}

	return position
}

// // Get book data and send to channel
// func processBook(bookChan chan<- bitfinex.Book) {
// 	book, err := api.Orderbook(SYMBOL, 5, 5)
// 	checkErr(err)
//
// 	bookChan <- book
// }

// Get trade data and send to channel
func processTrades() bitfinex.Trades {
	trades, err := api.Trades(SYMBOL, TRADENUM)
	checkErr(err)

	return trades
}

//
// // Get trade data and send to channel
// func processTrades(tradesChan chan<- bitfinex.Trades) {
// 	trades, err := api.Trades(SYMBOL, TRADENUM)
// 	checkErr(err)
//
// 	tradesChan <- trades
// }

// Calculate a volume-weighted moving average of trades
func calculateTheo(trades bitfinex.Trades) float64 {
	var sum1, sum2 float64
	for _, trade := range trades {
		sum1 += trade.Price * trade.Amount
		sum2 += trade.Amount
	}
	return sum1 / sum2
}

// Called on any error
func checkErr(err error) {
	if err != nil {
		cancelAll()
		apiErrors = true
	}
}

// Call on exit
func exit() {
	cancelAll()
	fmt.Println("\nCancelled all orders.")
}

// Cancel all orders
func cancelAll() {
	cancelled := false
	for !cancelled {
		cancelled, _ = api.CancelAll()
	}
	liveOrders = false
}

// Print results
func printResults(trades bitfinex.Trades,
	orders bitfinex.Orders, theo, position float64, start time.Time) {

	clearScreen()

	// fmt.Println("----------------------------")
	// fmt.Printf("%-10s%-10s%8s\n", " Bid", "  Ask", "Size ")
	// fmt.Println("----------------------------")
	// for i := range book.Asks {
	// 	item := book.Asks[len(book.Asks)-1-i]
	// 	fmt.Printf("%-10s%-10.4f%8.2f\n", "", item.Price, item.Amount)
	// }
	// for _, item := range book.Bids {
	// 	fmt.Printf("%-10.4f%-10.2s%8.2f\n", item.Price, "", item.Amount)
	// }
	// fmt.Println("----------------------------")

	fmt.Println("\nLast Trades:")
	for _, trade := range trades {
		fmt.Printf("%-6.4f - size: %6.2f\n", trade.Price, trade.Amount)
	}

	fmt.Printf("\nPosition: %.2f\n", position)
	fmt.Printf("Theo:     %.4f\n", theo)

	fmt.Println("\nActive orders:")
	for _, order := range orders.Orders {
		fmt.Printf("%8.2f %s @ %6.4f\n", order.Amount, SYMBOL, order.Price)
	}

	fmt.Printf("\n%v processing time...", time.Since(start))
}

// // Print results
// func printResults(book bitfinex.Book, trades bitfinex.Trades,
// 	orders bitfinex.Orders, theo, position float64, start time.Time) {
//
// 	clearScreen()
//
// 	fmt.Println("----------------------------")
// 	fmt.Printf("%-10s%-10s%8s\n", " Bid", "  Ask", "Size ")
// 	fmt.Println("----------------------------")
// 	for i := range book.Asks {
// 		item := book.Asks[len(book.Asks)-1-i]
// 		fmt.Printf("%-10s%-10.4f%8.2f\n", "", item.Price, item.Amount)
// 	}
// 	for _, item := range book.Bids {
// 		fmt.Printf("%-10.4f%-10.2s%8.2f\n", item.Price, "", item.Amount)
// 	}
// 	fmt.Println("----------------------------")
//
// 	fmt.Println("\nLast Trades:")
// 	for _, trade := range trades {
// 		fmt.Printf("%-6.4f - size: %6.2f\n", trade.Price, trade.Amount)
// 	}
//
// 	fmt.Printf("\nPosition: %.2f\n", position)
// 	fmt.Printf("Theo:     %.4f\n", theo)
//
// 	fmt.Println("\nActive orders:")
// 	for _, order := range orders.Orders {
// 		fmt.Printf("%8.2f %s @ %6.4f\n", order.Amount, SYMBOL, order.Price)
// 	}
//
// 	fmt.Printf("\n%v processing time...", time.Since(start))
// }
//
// Clear the terminal between prints
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}
<commit_msg>use os args for trade size and edge<commit_after>package main

import (
	"bitmm/bitfinex"
	"fmt"
	"log"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"time"
)

// Trade inputs
const (
	SYMBOL    = "ltcusd" // Instrument to trade
	MINCHANGE = 0.0001   // Minumum change required to update prices
	TRADENUM  = 20       // Number of trades to use in calculations
	MINO      = 0.011    // Min order size
)

var (
	// MAXO maximum order size
	MAXO float64
	// INEDGE edge for position entry orders
	INEDGE float64
	// OUTEDGE edge for position exit orders
	OUTEDGE    float64
	api        = bitfinex.New(os.Getenv("BITFINEX_KEY"), os.Getenv("BITFINEX_SECRET"))
	apiErrors  = false
	liveOrders = false
	orderTheo  = 0.0
)

func main() {
	fmt.Println("\nInitializing...")

	if len(os.Args) < 4 {
		fmt.Printf("usage: %s <size> <entry edge> <exit edge>\n", filepath.Base(os.Args[0]))
		os.Exit(1)
	}

	// Get MAXO, INEDGE, OUTEDGE from user input
	getVars()

	// Check for input to break loop
	inputChan := make(chan rune)
	go checkStdin(inputChan)

	// Run loop until user input is received
	runMainLoop(inputChan)
}

func getVars() {
	var err error
	if MAXO, err = strconv.ParseFloat(os.Args[1], 64); err != nil {
		log.Fatal(err)
	}
	if INEDGE, err = strconv.ParseFloat(os.Args[2], 64); err != nil {
		log.Fatal(err)
	}
	if OUTEDGE, err = strconv.ParseFloat(os.Args[3], 64); err != nil {
		log.Fatal(err)
	}
}

// Check for any user input
func checkStdin(inputChan chan<- rune) {
	var ch rune
	fmt.Scanf("%c", &ch)
	inputChan <- ch
}

// Infinite loop
func runMainLoop(inputChan <-chan rune) {

	var (
		trades bitfinex.Trades
		// book        bitfinex.Book
		orders      bitfinex.Orders
		start       time.Time
		oldPosition float64
		newPosition float64
		theo        float64
		lastTrade   int
	)

	for {
		// Record time for each iteration
		start = time.Now()

		// Exit if anything entered by user
		select {
		case <-inputChan:
			exit()
			return
		default:
		}

		trades = processTrades()
		if !apiErrors && trades[0].TID != lastTrade { // If new trades
			theo = calculateTheo(trades)
			newPosition = checkPosition()
			if (math.Abs(theo-orderTheo) >= MINCHANGE || math.Abs(oldPosition-
				newPosition) >= MINO || !liveOrders) && !apiErrors {
				orders = sendOrders(theo, newPosition)
			}
		}

		if !apiErrors {
			printResults(trades, orders, theo, newPosition, start)
			// Reset for next iteration
			oldPosition = newPosition
			lastTrade = trades[0].TID
		}

		// Reset for next iteration
		apiErrors = false
	}
}

// Send orders to the exchange
func sendOrders(theo, position float64) bitfinex.Orders {
	orderTheo = theo

	if liveOrders {
		cancelAll()
	}

	// Send new order request to the exchange
	params := calcOrderParams(position, theo)
	orders, err := api.MultipleNewOrders(params)
	liveOrders = true
	checkErr(err)
	return orders
}

func calcOrderParams(position, theo float64) []bitfinex.OrderParams {
	var params []bitfinex.OrderParams

	if math.Abs(position) < MINO { // No position
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, MAXO, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	} else if position < (-1*MAXO)+MINO { // Max short postion
		params = []bitfinex.OrderParams{
			{SYMBOL, -1 * position, theo - OUTEDGE, "bitfinex", "buy", "limit"},
		}
	} else if position > MAXO-MINO { // Max long postion
		params = []bitfinex.OrderParams{
			{SYMBOL, position, theo + OUTEDGE, "bitfinex", "sell", "limit"},
		}
	} else if (-1*MAXO)+MINO <= position && position <= -1*MINO { // Partial short
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, -1 * position, theo - OUTEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, MAXO + position, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	} else if MINO <= position && position <= MAXO-MINO { // Partial long
		params = []bitfinex.OrderParams{
			{SYMBOL, MAXO - position, theo - INEDGE, "bitfinex", "buy", "limit"},
			{SYMBOL, position, theo + OUTEDGE, "bitfinex", "sell", "limit"},
			{SYMBOL, MAXO, theo + INEDGE, "bitfinex", "sell", "limit"},
		}
	}

	return params
}

func checkPosition() float64 {
	var position float64
	posSlice, err := api.ActivePositions()
	checkErr(err)
	for _, pos := range posSlice {
		if pos.Symbol == SYMBOL {
			position = pos.Amount
		}
	}

	return position
}

// Get trade data
func processTrades() bitfinex.Trades {
	trades, err := api.Trades(SYMBOL, TRADENUM)
	checkErr(err)

	return trades
}

// Calculate a volume-weighted moving average of trades
func calculateTheo(trades bitfinex.Trades) float64 {
	var sum1, sum2 float64
	for _, trade := range trades {
		sum1 += trade.Price * trade.Amount
		sum2 += trade.Amount
	}
	return sum1 / sum2
}

// Called on any error
func checkErr(err error) {
	if err != nil {
		cancelAll()
		apiErrors = true
	}
}

// Call on exit
func exit() {
	cancelAll()
	fmt.Println("\nCancelled all orders.")
}

// Cancel all orders
func cancelAll() {
	cancelled := false
	for !cancelled {
		cancelled, _ = api.CancelAll()
	}
	liveOrders = false
}

// Print results
func printResults(trades bitfinex.Trades,
	orders bitfinex.Orders, theo, position float64, start time.Time) {

	clearScreen()

	fmt.Println("\nLast Trades:")
	for _, trade := range trades {
		fmt.Printf("%-6.4f - size: %6.2f\n", trade.Price, trade.Amount)
	}

	fmt.Printf("\nPosition: %.2f\n", position)
	fmt.Printf("Theo:     %.4f\n", theo)

	fmt.Println("\nActive orders:")
	for _, order := range orders.Orders {
		fmt.Printf("%8.2f %s @ %6.4f\n", order.Amount, SYMBOL, order.Price)
	}

	fmt.Printf("\n%v processing time...", time.Since(start))
}

// Clear the terminal between prints
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}
<|endoftext|>


<commit_before>package ics

import (
	"bufio"
	"io"
	"unicode/utf8"
)

const lineLength = 75

var (
	delims       = [...]byte{'\r', '\n', ' '}
	continuation = delims[:]
	endLine      = delims[:2]
)

type folder struct {
	bw *bufio.Writer
}

func newFolder(w io.Writer) folder {
	return folder{bufio.NewWriter(w)}
}

func (f folder) writeLine(p []byte) (err error) {
	pos := 0
	for len(p) > 0 {
		r, s := utf8.DecodeRune(p)
		if pos+s > lineLength {
			f.bw.Write(continuation)
			pos = 1
		}
		_, err = f.bw.WriteRune(r)
		if err != nil {
			return err
		}
		pos += s
		p = p[s:]
	}
	_, err = f.bw.Write(endLine)
	if err != nil {
		return err
	}
	return nil
}

func (f folder) flush() error {
	return f.bw.Flush()
}
<commit_msg>removed bufio from folder<commit_after>package ics

import (
	"io"
	"unicode/utf8"
)

const lineLength = 75

var (
	delims       = [...]byte{'\r', '\n', ' '}
	continuation = delims[:]
	endLine      = delims[:2]
)

type folder struct {
	w io.Writer
}

func newFolder(w io.Writer) folder {
	return folder{w}
}

func (f folder) writeLine(p []byte) (err error) {
	pos := 0
	var utf [5]byte
	for len(p) > 0 {
		r, s := utf8.DecodeRune(p)
		if pos+s > lineLength {
			f.w.Write(continuation)
			pos = 1
		}
		l := utf8.EncodeRune(utf[:], r)
		_, err = f.w.Write(utf[:l])
		if err != nil {
			return err
		}
		pos += s
		p = p[s:]
	}
	_, err = f.w.Write(endLine)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// TODO: quit on terminal input
// TODO: make price a volume wieghted moving average

package main

import (
	"./bitfinex"
	"fmt"
	"log"
	"math"
	"os"
	"os/exec"
	"time"
)

// Trade inputs
const (
	SYMBOL    = "ltcusd" // Instrument to trade
	MINCHANGE = 0.0005   // Minumum change required to update prices
	TRADENUM  = 10       // Number of trades to use in calculations
	AMOUNT    = 0.10     // Size to trade
	BIDEDGE   = 0.01     // Required edge for a buy order
	ASKEDGE   = 0.01     // Required edge for a sell order
)

var (
	api = bitfinex.New(os.Getenv("BITFINEX_KEY"), os.Getenv("BITFINEX_SECRET"))
)

func main() {
	fmt.Println("\nStarting up...")

	// Create channels for exchange communication
	bookChan := make(chan bitfinex.Book)
	tradesChan := make(chan bitfinex.Trades)
	bidChan := make(chan bitfinex.Order)
	askChan := make(chan bitfinex.Order)

	// Initial orders
	bid, ask := createOrders()

	var (
		trades bitfinex.Trades
		book   bitfinex.Book
		start  time.Time
		theo   float64
	)

	for {
		start = time.Now()

		// Get data in separate goroutines
		go processBook(bookChan)
		go processTrades(tradesChan)

		// Modify orders in separate goroutines when trade data returns
		trades = <-tradesChan
		theo = calculateTheo(trades)
		go replaceBid(bid, bidChan, theo)
		go replaceAsk(ask, askChan, theo)

		// Print data and current orders when all communication is finished
		bid = <-bidChan
		ask = <-askChan
		book = <-bookChan
		printResults(book, trades, bid, ask)

		// Print processing time
		fmt.Printf("\n%v processing time...", time.Since(start))
	}
}

// Create initial orders
func createOrders() (bitfinex.Order, bitfinex.Order) {
	// Get the current price
	trades, err := api.Trades(SYMBOL, 1)
	checkErr(err)
	price := trades[0].Price

	// Order parameters
	params := []bitfinex.OrderParams{
		{SYMBOL, AMOUNT, price - BIDEDGE, "bitfinex", "buy", "limit"},
		{SYMBOL, AMOUNT, price + ASKEDGE, "bitfinex", "sell", "limit"},
	}

	// Send new order request to the exchange
	orders, err := api.MultipleNewOrders(params)
	checkErr(err)

	return orders.Orders[0], orders.Orders[1]
}

// Get book data and send to channel
func processBook(bookChan chan<- bitfinex.Book) {
	book, err := api.Orderbook(SYMBOL, 5, 5)
	checkErr(err)

	bookChan <- book
}

// Get trade data and send to channel
func processTrades(tradesChan chan<- bitfinex.Trades) {
	trades, err := api.Trades(SYMBOL, TRADENUM)
	checkErr(err)

	tradesChan <- trades
}

// Calculate a volume-weighted moving average of trades
func calculateTheo(trades bitfinex.Trades) float64 {
	var sum1, sum2 float64
	for _, trade := range trades {
		sum1 += trade.Price * trade.Amount
		sum2 += trade.Amount
	}
	return sum1 / sum2
}

// Modify bid order and send to channel
func replaceBid(bid bitfinex.Order, bidChan chan<- bitfinex.Order, theo float64) {
	price := theo - BIDEDGE

	var err error
	if math.Abs(price-bid.Price) >= MINCHANGE {
		bid, err = api.ReplaceOrder(bid.ID, SYMBOL, AMOUNT, price, "bitfinex", "buy", "limit")
		checkErr(err)
		if bid.ID == 0 {
			cancelAll()
			log.Fatal("Failed to replace bid order")
		}
	}

	bidChan <- bid
}

// Modify ask order and send to channel
func replaceAsk(ask bitfinex.Order, askChan chan<- bitfinex.Order, theo float64) {
	price := theo + ASKEDGE

	var err error
	if math.Abs(price-ask.Price) >= MINCHANGE {
		ask, err = api.ReplaceOrder(ask.ID, SYMBOL, AMOUNT, price, "bitfinex", "sell", "limit")
		checkErr(err)
		if ask.ID == 0 {
			cancelAll()
			log.Fatal("Failed to replace ask order")
		}
	}

	askChan <- ask
}

// Print results
func printResults(book bitfinex.Book, trades bitfinex.Trades, bid, ask bitfinex.Order) {
	clearScreen()

	fmt.Println("----------------------------")
	fmt.Printf("%-10s%-10s%8s\n", " Bid", "  Ask", "Size ")
	fmt.Println("----------------------------")
	for i := range book.Asks {
		item := book.Asks[len(book.Asks)-1-i]
		fmt.Printf("%-10s%-10.4f%8.2f\n", "", item.Price, item.Amount)
	}
	for _, item := range book.Bids {
		fmt.Printf("%-10.4f%-10.2s%8.2f\n", item.Price, "", item.Amount)
	}
	fmt.Println("----------------------------")

	fmt.Println("\nLast Trades:")
	for _, trade := range trades {
		fmt.Printf("%-6.4f - size: %6.2f\n", trade.Price, trade.Amount)
	}

	fmt.Printf("\nCurrent Bid: %6.4f\n", bid.Price)
	fmt.Printf("Current Ask: %6.4f\n", ask.Price)
}

// Exit on any errors
func checkErr(err error) {
	if err != nil {
		cancelAll()
		log.Fatal(err)
	}
}

// Clear the terminal between prints
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

// Cancel all orders
func cancelAll() {
	cancelled := false
	for !cancelled {
		cancelled, _ = api.CancelAll()
	}
	fmt.Println("\nALL ORDERS HAVE BEEN CANCELLED")
}
<commit_msg>pre ben meeting<commit_after>package main

import (
	"./bitfinex"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"os/exec"
	"time"
)

// Trade inputs
const (
	SYMBOL    = "ltcusd" // Instrument to trade
	MINCHANGE = 0.0005   // Minumum change required to update prices
	TRADENUM  = 10       // Number of trades to use in calculations
	AMOUNT    = 0.50     // Size to trade
	BIDEDGE   = 0.05     // Required edge for a buy order
	ASKEDGE   = 0.05     // Required edge for a sell order
)

var (
	api = bitfinex.New(os.Getenv("BITFINEX_KEY"), os.Getenv("BITFINEX_SECRET"))
)

func main() {
	fmt.Println("\nStarting up...")

	// Create channels for exchange communication
	bookChan := make(chan bitfinex.Book)
	tradesChan := make(chan bitfinex.Trades)
	bidChan := make(chan bitfinex.Order)
	askChan := make(chan bitfinex.Order)

	// Initial orders
	bid, ask := createOrders()

	var (
		trades bitfinex.Trades
		book   bitfinex.Book
		start  time.Time
		theo   float64
	)

	for {
		start = time.Now()

		// Get data in separate goroutines
		go processBook(bookChan)
		go processTrades(tradesChan)

		// Modify orders in separate goroutines when trade data returns
		trades = <-tradesChan
		theo = calculateTheo(trades)
		go replaceBid(bid, bidChan, theo)
		go replaceAsk(ask, askChan, theo)

		// Print data and current orders when all communication is finished
		bid = <-bidChan
		ask = <-askChan
		book = <-bookChan
		printResults(book, trades, bid, ask)

		// Print processing time
		fmt.Printf("\n%v processing time...", time.Since(start))

		// TODO: Exit if anything entered by user
		// fmt.Scan()
		// io.ReadFull(os.Stdin, buf)
		// bufio.NewReader(os.Stdin) -> Read(), Peek(), Buffered()
		// ioutil.ReadAll(os.Stdin)
		if bytes, err := ioutil.ReadAll(os.Stdin); len(bytes) > 0 || err != nil {
			cancelAll()
			break
		}
	}
}

// Create initial orders
func createOrders() (bitfinex.Order, bitfinex.Order) {
	// Get the current price
	trades, err := api.Trades(SYMBOL, 1)
	checkErr(err)
	price := trades[0].Price

	// Order parameters
	params := []bitfinex.OrderParams{
		{SYMBOL, AMOUNT, price - BIDEDGE, "bitfinex", "buy", "limit"},
		{SYMBOL, AMOUNT, price + ASKEDGE, "bitfinex", "sell", "limit"},
	}

	// Send new order request to the exchange
	orders, err := api.MultipleNewOrders(params)
	checkErr(err)

	return orders.Orders[0], orders.Orders[1]
}

// Get book data and send to channel
func processBook(bookChan chan<- bitfinex.Book) {
	book, err := api.Orderbook(SYMBOL, 5, 5)
	checkErr(err)

	bookChan <- book
}

// Get trade data and send to channel
func processTrades(tradesChan chan<- bitfinex.Trades) {
	trades, err := api.Trades(SYMBOL, TRADENUM)
	checkErr(err)

	tradesChan <- trades
}

// Calculate a volume-weighted moving average of trades
func calculateTheo(trades bitfinex.Trades) float64 {
	var sum1, sum2 float64
	for _, trade := range trades {
		sum1 += trade.Price * trade.Amount
		sum2 += trade.Amount
	}
	return sum1 / sum2
}

// Modify bid order and send to channel
func replaceBid(bid bitfinex.Order, bidChan chan<- bitfinex.Order, theo float64) {
	price := theo - BIDEDGE

	var err error
	if math.Abs(price-bid.Price) >= MINCHANGE {
		bid, err = api.ReplaceOrder(bid.ID, SYMBOL, AMOUNT, price, "bitfinex", "buy", "limit")
		checkErr(err)
		if bid.ID == 0 {
			cancelAll()
			log.Fatal("Failed to replace bid order")
		}
	}

	bidChan <- bid
}

// Modify ask order and send to channel
func replaceAsk(ask bitfinex.Order, askChan chan<- bitfinex.Order, theo float64) {
	price := theo + ASKEDGE

	var err error
	if math.Abs(price-ask.Price) >= MINCHANGE {
		ask, err = api.ReplaceOrder(ask.ID, SYMBOL, AMOUNT, price, "bitfinex", "sell", "limit")
		checkErr(err)
		if ask.ID == 0 {
			cancelAll()
			log.Fatal("Failed to replace ask order")
		}
	}

	askChan <- ask
}

// Print results
func printResults(book bitfinex.Book, trades bitfinex.Trades, bid, ask bitfinex.Order) {
	clearScreen()

	fmt.Println("----------------------------")
	fmt.Printf("%-10s%-10s%8s\n", " Bid", "  Ask", "Size ")
	fmt.Println("----------------------------")
	for i := range book.Asks {
		item := book.Asks[len(book.Asks)-1-i]
		fmt.Printf("%-10s%-10.4f%8.2f\n", "", item.Price, item.Amount)
	}
	for _, item := range book.Bids {
		fmt.Printf("%-10.4f%-10.2s%8.2f\n", item.Price, "", item.Amount)
	}
	fmt.Println("----------------------------")

	fmt.Println("\nLast Trades:")
	for _, trade := range trades {
		fmt.Printf("%-6.4f - size: %6.2f\n", trade.Price, trade.Amount)
	}

	fmt.Printf("\nCurrent Bid: %6.4f\n", bid.Price)
	fmt.Printf("Current Ask: %6.4f\n", ask.Price)
}

// Exit on any errors
func checkErr(err error) {
	if err != nil {
		cancelAll()
		log.Fatal(err)
	}
}

// Clear the terminal between prints
func clearScreen() {
	c := exec.Command("clear")
	c.Stdout = os.Stdout
	c.Run()
}

// Cancel all orders
func cancelAll() {
	cancelled := false
	for !cancelled {
		cancelled, _ = api.CancelAll()
	}
	fmt.Println("\nALL ORDERS HAVE BEEN CANCELLED")
}
<|endoftext|>


<commit_before>package blend

import (
	"image/color"
	"math"
)

const (
	MULTIPLY = iota
	SCREEN
	OVERLAY
	SOFT_LIGHT
	HARD_LIGHT
	COLOR_DODGE
	COLOR_BURN
	LINEAR_COLOR_DODGE
	LINEAR_COLOR_BURN
	DARKEN
	LIGHTEN
	DIFFERENCE
	EXCLUSION
	REFLEX
	LINEAR_LIGHT
	PIN_LIGHT
	VIVID_LIGHT
	HARD_MIX
	// Blending modes in HSL color model.
	HUE
	COLOR
	LUMINOSITY
)

const (
	max = 65535.0
	mid = max / 2.0
)

type blendFunc func(float64, float64) float64

type rgbaf64 struct {
	r, g, b, a float64
}

func (c rgbaf64) RGBA() (uint32, uint32, uint32, uint32) {
	r := float64ToUint16(c.r)
	g := float64ToUint16(c.g)
	b := float64ToUint16(c.b)
	a := float64ToUint16(c.a)
	return uint32(r), uint32(g), uint32(b), uint32(a)
}

func colorTorgbaf64(c color.Color) rgbaf64 {
	r, g, b, a := c.RGBA()
	return rgbaf64{float64(r), float64(g), float64(b), float64(a)}
}

func BlendColor(source, dest color.Color, mode int) color.Color {
	switch mode {
	case MULTIPLY:
		return blend(source, dest, multiply)
	case SCREEN:
		return blend(source, dest, screen)
	case OVERLAY:
		return blend(source, dest, overlay)
	case SOFT_LIGHT:
		return blend(source, dest, soft_light)
	case HARD_LIGHT:
		return blend(source, dest, hard_light)
	}
	return rgbaf64{0.0, 0.0, 0.0, 0.0}
}

func blend(source, dest color.Color, bf blendFunc) (c rgbaf64) {
	s := colorTorgbaf64(source)
	d := colorTorgbaf64(dest)
	c.r = bf(s.r, d.r)
	c.r = bf(s.g, d.g)
	c.b = bf(s.b, d.b)
	c.a = s.a
	return
}

// Blend Modes
func multiply(s, d float64) float64 {
	return s * d / max
}

func screen(s, d float64) float64 {
	return s + d - s*d/max
}

func overlay(s, d float64) float64 {
	if d < mid {
		return 2 * s * d / max
	}
	return max - 2*(max-s)*(max-d)/max
}

func soft_light(s, d float64) float64 {
	if s > mid {
		return d + (max-d)*((s-mid)/mid)*(0.5-math.Abs(d-mid)/max)
	}
	return d - d*((mid-s)/mid)*(0.5-math.Abs(d-mid)/max)
}

func hard_light(s, d float64) float64 {
	if s > mid {
		return d + (max-d)*((s-mid)/mid)
	}
	return d * s / mid
}
<commit_msg>blend func red channel error fixed.<commit_after>package blend

import (
	"image/color"
	"math"
)

const (
	MULTIPLY = iota
	SCREEN
	OVERLAY
	SOFT_LIGHT
	HARD_LIGHT
	COLOR_DODGE
	COLOR_BURN
	LINEAR_COLOR_DODGE
	LINEAR_COLOR_BURN
	DARKEN
	LIGHTEN
	DIFFERENCE
	EXCLUSION
	REFLEX
	LINEAR_LIGHT
	PIN_LIGHT
	VIVID_LIGHT
	HARD_MIX
	// Blending modes in HSL color model.
	HUE
	COLOR
	LUMINOSITY
)

const (
	max = 65535.0
	mid = max / 2.0
)

type blendFunc func(float64, float64) float64

type rgbaf64 struct {
	r, g, b, a float64
}

func (c rgbaf64) RGBA() (uint32, uint32, uint32, uint32) {
	r := float64ToUint16(c.r)
	g := float64ToUint16(c.g)
	b := float64ToUint16(c.b)
	a := float64ToUint16(c.a)
	return uint32(r), uint32(g), uint32(b), uint32(a)
}

func colorTorgbaf64(c color.Color) rgbaf64 {
	r, g, b, a := c.RGBA()
	return rgbaf64{float64(r), float64(g), float64(b), float64(a)}
}

func BlendColor(source, dest color.Color, mode int) color.Color {
	switch mode {
	case MULTIPLY:
		return blend(source, dest, multiply)
	case SCREEN:
		return blend(source, dest, screen)
	case OVERLAY:
		return blend(source, dest, overlay)
	case SOFT_LIGHT:
		return blend(source, dest, soft_light)
	case HARD_LIGHT:
		return blend(source, dest, hard_light)
	}
	return rgbaf64{0.0, 0.0, 0.0, 0.0}
}

func blend(source, dest color.Color, bf blendFunc) (c rgbaf64) {
	s := colorTorgbaf64(source)
	d := colorTorgbaf64(dest)
	c.r = bf(s.r, d.r)
	c.g = bf(s.g, d.g)
	c.b = bf(s.b, d.b)
	c.a = s.a
	return
}

// Blend Modes
func multiply(s, d float64) float64 {
	return s * d / max
}

func screen(s, d float64) float64 {
	return s + d - s*d/max
}

func overlay(s, d float64) float64 {
	if d < mid {
		return 2 * s * d / max
	}
	return max - 2*(max-s)*(max-d)/max
}

func soft_light(s, d float64) float64 {
	if s > mid {
		return d + (max-d)*((s-mid)/mid)*(0.5-math.Abs(d-mid)/max)
	}
	return d - d*((mid-s)/mid)*(0.5-math.Abs(d-mid)/max)
}

func hard_light(s, d float64) float64 {
	if s > mid {
		return d + (max-d)*((s-mid)/mid)
	}
	return d * s / mid
}
<|endoftext|>


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

/*
A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is a local version of FNV,
a non-cryptographic hashing function, seeded with the index number of
the kth hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"io"
	"math"

	"github.com/willf/bitset"
)

type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

// Create a new Bloom filter with _m_ bits and _k_ hashing functions
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{m, k, bitset.New(m)}
}

func baseHashes(data []byte) (h1 uint, h2 uint) {
	h1 = fnvhash(0, data)
	h2 = fnvhash(1, data)
	return
}

func fnvhash(index uint, data []byte) uint {
	hash := uint(index)
	for _, c := range data {
		hash ^= uint(c)
		hash *= 16777619
	}
	return hash
}

func (f *BloomFilter) location(h1 uint, h2 uint, i uint) (location uint) {
	l := (h1 + h2*i) % f.m
	location = uint(l)
	return
}

// estimate parameters. Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func estimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// Create a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := estimateParameters(n, fp)
	return New(m, k)
}

// Return the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// Return the number of hash functions used
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h1, h2, i))
	}
	return f
}

// Tests for the presence of data in the Bloom filter
func (f *BloomFilter) Test(data []byte) bool {
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h1, h2, i)) {
			return false
		}
	}
	return true
}

// Equivalent to calling Test(data) then Add(data).  Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h1, h2, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// Clear all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// Estimate, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// whilst storing n entries; runs 100,000 tests.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds) * float64(100))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}
<commit_msg>rewrite documentation according to Go standards<commit_after>/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is a local version of FNV,
a non-cryptographic hashing function, seeded with the index number of
the kth hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"io"
	"math"

	"github.com/willf/bitset"
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{m, k, bitset.New(m)}
}

// fnvhash returns the FNV hash value of data, using an index as a seed
func fnvhash(index uint, data []byte) uint {
	hash := uint(index)
	for _, c := range data {
		hash ^= uint(c)
		hash *= 16777619
	}
	return hash
}

// baseHashes returns the two hash values of data that are used to create k
// hashes
func baseHashes(data []byte) (h1 uint, h2 uint) {
	h1 = fnvhash(0, data)
	h2 = fnvhash(1, data)
	return
}

// location returns the ith hashed location using the two base hash values
func (f *BloomFilter) location(h1 uint, h2 uint, i uint) (location uint) {
	l := (h1 + h2*i) % f.m
	location = uint(l)
	return
}

// estimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func estimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := estimateParameters(n, fp)
	return New(m, k)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h1, h2, i))
	}
	return f
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) bool {
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h1, h2, i)) {
			return false
		}
	}
	return true
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h1, h2 := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h1, h2, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// ClearAll clears all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// EstimateFalsePositiveRate returns, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// while storing n entries; runs 100,000 tests. This is an empirical
// test using integers as keys. As a side-effect, it clears the BloomFilter.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds) * float64(100))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}
<|endoftext|>


<commit_before>package vkutil

import (
	"encoding/json"
	"log"
	"net/url"

	"github.com/zhuharev/vk"
)

// Topic board item
type Topic struct {
	ID        int       `json:"id"`
	Title     string    `json:"title"`
	Created   EpochTime `json:"created"`
	CreatedBy int       `json:"created_by"`
	Updated   EpochTime `json:"updated"`
	UpdatedBy int       `json:"updated_by"`
	IsClosed  int       `json:"is_closed"`
	Comments  int       `json:"comments"`
}

type RespTopic struct {
	Count int     `json:"count"`
	Items []Topic `json:"items"`
}

// ResponseTopics represent topic items
type ResponseTopics struct {
	Resp RespTopic `json:"response"`
}

// BoardGetCommetns returns board items
func (api *Api) BoardGetCommetns(groupID, topicID int, args ...url.Values) ([]Comment, error) {
	params := setToUrlValues("group_id", groupID, args...)
	params = setToUrlValues("topic_id", topicID, params)
	//params = setToUrlValues("sort", "desc", params)

	bts, err := api.VkApi.Request(vk.METHOD_BOARD_GET_COMMENTS, params)
	if err != nil {
		return nil, err
	}
	log.Println(string(bts))
	var r ResponseComments
	err = json.Unmarshal(bts, &r)
	if err != nil {
		return nil, err
	}
	return r.Response.Items, nil
}
<commit_msg>Add get board topic method<commit_after>package vkutil

import (
	"context"
	"encoding/json"
	"log"
	"net/url"

	"github.com/zhuharev/vk"
)

const (
	MethodBoardGetTopics = "board.getTopics"
)

// Topic board item
type Topic struct {
	ID        int       `json:"id"`
	Title     string    `json:"title"`
	Created   EpochTime `json:"created"`
	CreatedBy int       `json:"created_by"`
	Updated   EpochTime `json:"updated"`
	UpdatedBy int       `json:"updated_by"`
	IsClosed  int       `json:"is_closed"`
	Comments  int       `json:"comments"`
}

type RespTopic struct {
	Count int     `json:"count"`
	Items []Topic `json:"items"`
}

// ResponseTopics represent topic items
type ResponseTopics struct {
	Resp RespTopic `json:"response"`
}

// BoardGetCommetns returns board items
func (api *Api) BoardGetCommetns(groupID, topicID int, args ...url.Values) ([]Comment, error) {
	params := setToUrlValues("group_id", groupID, args...)
	params = setToUrlValues("topic_id", topicID, params)
	//params = setToUrlValues("sort", "desc", params)

	bts, err := api.VkApi.Request(vk.METHOD_BOARD_GET_COMMENTS, params)
	if err != nil {
		return nil, err
	}
	log.Println(string(bts))
	var r ResponseComments
	err = json.Unmarshal(bts, &r)
	if err != nil {
		return nil, err
	}
	return r.Response.Items, nil
}

func (api *Api) BoardGetAllTopics(ctx context.Context, groupID int) ([]Topic, error) {
	var (
		count  = 100
		offset = 0
		result []Topic
	)
	for {
		topics, err := api.BoardGetTopics(ctx, groupID, count, offset)
		if err != nil {
			return nil, err
		}
		offset += len(topics)
		result = append(result, topics...)

		if len(topics) < count {
			break
		}
	}
	return result, nil
}

func (api *Api) BoardGetTopics(ctx context.Context, groupID int, count int, offset int) ([]Topic, error) {
	params := setToUrlValues("group_id", groupID)
	params = setToUrlValues("count", count, params)
	params = setToUrlValues("offset", offset, params)

	bts, err := api.VkApi.Request(MethodBoardGetTopics, params)
	if err != nil {
		return nil, err
	}

	var r ResponseTopics
	err = json.Unmarshal(bts, &r)
	if err != nil {
		return nil, err
	}
	return r.Resp.Items, nil
}
<|endoftext|>


<commit_before>package hunter

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"time"

	"github.com/concourse/atc"
	"github.com/concourse/go-concourse/concourse"
)

const (
	BUILD_ABORTED    = "aborted"
	STATUS_FORBIDDEN = "forbidden"
)

type SearchSpec struct {
	Pattern string
	Limit   int
}

type Searcher struct {
	client Client
}

type Build struct {
	atc.Build

	ConcourseURL string
}

func NewSearcher(client Client) *Searcher {
	return &Searcher{
		client: client,
	}
}

func (s *Searcher) Search(spec SearchSpec) ([]Build, error) {
	builds, _, err := s.client.Builds(concourse.Page{Limit: spec.Limit})
	if err != nil {
		return nil, err
	}

	bs := []Build{}
	count := 1
	for _, build := range builds {
		if build.Status == BUILD_ABORTED {
			continue
		}

		events, err := s.client.BuildEvents(strconv.Itoa(build.ID))
		if err != nil {
			// Not sure why, but concourse.Builds returns builds from other teams
			if err.Error() == STATUS_FORBIDDEN {
				continue
			}
			return []Build{}, err
		}

		ok, err := regexp.Match(spec.Pattern, events)
		if err != nil {
			return []Build{}, errors.New("failed trying to match pattern given")
		}

		if ok {
			concourseURL := fmt.Sprintf("%s%s", s.client.ConcourseURL(), build.URL)
			b := Build{build, concourseURL}
			bs = append(bs, b)
		}

		if count%100 == 0 {
			fmt.Printf("finished count= %+v\n ...", count)
			time.Sleep(10 * time.Second)
		}
		count = count + 1
	}
	return bs, nil
}
<commit_msg>Remove debug message<commit_after>package hunter

import (
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"time"

	"github.com/concourse/atc"
	"github.com/concourse/go-concourse/concourse"
)

const (
	BUILD_ABORTED    = "aborted"
	STATUS_FORBIDDEN = "forbidden"
)

type SearchSpec struct {
	Pattern string
	Limit   int
}

type Searcher struct {
	client Client
}

type Build struct {
	atc.Build

	ConcourseURL string
}

func NewSearcher(client Client) *Searcher {
	return &Searcher{
		client: client,
	}
}

func (s *Searcher) Search(spec SearchSpec) ([]Build, error) {
	builds, _, err := s.client.Builds(concourse.Page{Limit: spec.Limit})
	if err != nil {
		return nil, err
	}

	bs := []Build{}
	count := 1
	for _, build := range builds {
		if build.Status == BUILD_ABORTED {
			continue
		}

		events, err := s.client.BuildEvents(strconv.Itoa(build.ID))
		if err != nil {
			// Not sure why, but concourse.Builds returns builds from other teams
			if err.Error() == STATUS_FORBIDDEN {
				continue
			}
			return []Build{}, err
		}

		ok, err := regexp.Match(spec.Pattern, events)
		if err != nil {
			return []Build{}, errors.New("failed trying to match pattern given")
		}

		if ok {
			concourseURL := fmt.Sprintf("%s%s", s.client.ConcourseURL(), build.URL)
			b := Build{build, concourseURL}
			bs = append(bs, b)
		}

		if count%100 == 0 {
			time.Sleep(10 * time.Second)
		}
		count = count + 1
	}
	return bs, nil
}
<|endoftext|>


<commit_before>/*
Package framed adds basic support for message framing over Streams.

Message are length-prefixed.  The first two bytes are an unsigned 16 bit int
stored in little-endian byte order.  The remaining bytes are the bytes of the
message.

The use of a uint16 means that the maximum possible data length is 65535.

Example:

	package main

	import (
		"github.com/oxtoacart/framed"
		"net"
		"log"
	)

	func main() {
		// Replace host:port with an actual TCP server, for example the echo service
		if conn, err := net.Dial("tcp", "host:port"); err == nil {
			framedConn = Framed{conn}
			if err := framedConn.Write([]byte("Hello World")); err == nil {
				if resp, err := framedConn.Read(); err == nil {
					log.Println("We're done!")
				}
			}
		}
	}
*/
package framed

import (
	"encoding/binary"
	"fmt"
	"io"
)

var endianness = binary.LittleEndian

/*
A Framed enhances an io.ReadWriteCloser to write data in contiguous frames.

It implements the Reader and Writer interfaces, but unlike typical Readers
it only returns whole frames.  Unlike typical Writers, it will not allow
frames to be fragmented.

Although the underlying ReadWriteCloser may be safe to use from multiple
goroutines, a Framed is not.
*/
type Framed struct {
	Stream io.ReadWriteCloser // the raw underlying connection

}

func (framed *Framed) Close() error {
	return framed.Stream.Close()
}

/*
Read implements the function from io.Reader.  Unlike io.Reader.Read,
frame.Read only returns full frames of data.
*/
func (framed *Framed) Read(frame []byte) (n int, err error) {
	var nb uint16
	err = binary.Read(framed.Stream, endianness, &nb)
	if err != nil {
		return
	}

	n = int(nb)
	// Read into buffer
	var totalRead int
	for totalRead = 0; totalRead < n; {
		var nr int
		nr, err = framed.Stream.Read(frame[totalRead:n])
		if err != nil {
			return
		}
		totalRead += nr
	}
	if totalRead != n {
		err = fmt.Errorf("%d bytes read did not match %d bytes expected", totalRead, n)
	}
	return
}

func (framed *Framed) Write(frame []byte) (n int, err error) {
	n = len(frame)

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	if written, err = framed.Stream.Write(frame); err != nil {
		return
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}
<commit_msg>Fixed problem with over-reading into buffers that are too large<commit_after>/*
Package framed adds basic support for message framing over Streams.

Message are length-prefixed.  The first two bytes are an unsigned 16 bit int
stored in little-endian byte order.  The remaining bytes are the bytes of the
message.

The use of a uint16 means that the maximum possible data length is 65535.

Example:

	package main

	import (
		"github.com/oxtoacart/framed"
		"net"
		"log"
	)

	func main() {
		// Replace host:port with an actual TCP server, for example the echo service
		if conn, err := net.Dial("tcp", "host:port"); err == nil {
			framedConn = Framed{conn}
			if err := framedConn.Write([]byte("Hello World")); err == nil {
				if resp, err := framedConn.Read(); err == nil {
					log.Println("We're done!")
				}
			}
		}
	}
*/
package framed

import (
	"encoding/binary"
	"fmt"
	"io"
)

var endianness = binary.LittleEndian

/*
A Framed enhances an io.ReadWriteCloser to write data in contiguous frames.

It implements the Reader and Writer interfaces, but unlike typical Readers
it only returns whole frames.  Unlike typical Writers, it will not allow
frames to be fragmented.

Although the underlying ReadWriteCloser may be safe to use from multiple
goroutines, a Framed is not.
*/
type Framed struct {
	Stream io.ReadWriteCloser // the raw underlying connection

}

func (framed *Framed) Close() error {
	return framed.Stream.Close()
}

/*
Read implements the function from io.Reader.  Unlike io.Reader.Read,
frame.Read only returns full frames of data.
*/
func (framed *Framed) Read(buffer []byte) (n int, err error) {
	var nb uint16
	err = binary.Read(framed.Stream, endianness, &nb)
	if err != nil {
		return
	}

	n = int(nb)

	bufferSize := len(buffer)
	if n > bufferSize {
		return 0, fmt.Errorf("Buffer of size %d is too small to hold frame of size %d", bufferSize, n)
	}
	// Read into buffer
	var totalRead int
	for totalRead = 0; totalRead < n; {
		var nr int
		nr, err = framed.Stream.Read(buffer[totalRead:n])
		if err != nil {
			return
		}
		totalRead += nr
	}
	if totalRead != n {
		err = fmt.Errorf("%d bytes read did not match %d bytes expected", totalRead, n)
	}
	return
}

func (framed *Framed) Write(frame []byte) (n int, err error) {
	n = len(frame)

	// Write the length header
	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
		return
	}

	// Write the data
	var written int
	if written, err = framed.Stream.Write(frame); err != nil {
		return
	}
	if written != n {
		err = fmt.Errorf("%d bytes written, expected to write %d", written, n)
	}
	return
}
<|endoftext|>