<commit_before>package template

// Functions available to gondola templates

import (
	"encoding/json"
	"fmt"
	"gondola/assets"
	"html/template"
	"reflect"
	"strconv"
	"strings"
)

func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func neq(args ...interface{}) bool {
	return !eq(args...)
}

func _json(arg interface{}) string {
	if arg == nil {
		return ""
	}
	b, err := json.Marshal(arg)
	if err == nil {
		return string(b)
	}
	return ""
}

func nz(x interface{}) bool {
	switch x := x.(type) {
	case int, uint, int64, uint64, byte, float32, float64:
		return x != 0
	case string:
		return len(x) > 0
	}
	return false
}

func lower(x string) string {
	return strings.ToLower(x)
}

func join(x []string, sep string) string {
	s := ""
	for _, v := range x {
		s += fmt.Sprintf("%v%s", v, sep)
	}
	if len(s) > 0 {
		return s[:len(s)-len(sep)]
	}
	return ""
}

func _map(args ...interface{}) (map[string]interface{}, error) {
	var key string
	m := make(map[string]interface{})
	for ii, v := range args {
		if ii%2 == 0 {
			if s, ok := v.(string); ok {
				key = s
			} else {
				return nil, fmt.Errorf("Invalid argument to map at index %d, %t instead of string", ii, v)
			}
		} else {
			m[key] = v
		}
	}
	return m, nil
}

func mult(args ...interface{}) (float64, error) {
	val := 1.0
	for ii, v := range args {
		value := reflect.ValueOf(v)
		switch value.Kind() {
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			val *= float64(value.Int())
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			val *= float64(value.Uint())
		case reflect.Float32, reflect.Float64:
			val *= value.Float()
		case reflect.String:
			v, err := strconv.ParseFloat(value.String(), 64)
			if err != nil {
				return 0, fmt.Errorf("Error parsing string passed to mult at index %d: %s", ii, err)
			}
			val *= v
		default:
			return 0, fmt.Errorf("Invalid argument of type %T passed to mult at index %d", v, ii)
		}
	}
	return val, nil

}

var templateFuncs template.FuncMap = template.FuncMap{
	"eq":     eq,
	"neq":    neq,
	"json":   _json,
	"nz":     nz,
	"lower":  lower,
	"join":   join,
	"map":    _map,
	"mult":   mult,
	"render": assets.Render,
}
<commit_msg>Add concat() template func<commit_after>package template

// Functions available to gondola templates

import (
	"encoding/json"
	"fmt"
	"gondola/assets"
	"html/template"
	"reflect"
	"strconv"
	"strings"
)

func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func neq(args ...interface{}) bool {
	return !eq(args...)
}

func _json(arg interface{}) string {
	if arg == nil {
		return ""
	}
	b, err := json.Marshal(arg)
	if err == nil {
		return string(b)
	}
	return ""
}

func nz(x interface{}) bool {
	switch x := x.(type) {
	case int, uint, int64, uint64, byte, float32, float64:
		return x != 0
	case string:
		return len(x) > 0
	}
	return false
}

func lower(x string) string {
	return strings.ToLower(x)
}

func join(x []string, sep string) string {
	s := ""
	for _, v := range x {
		s += fmt.Sprintf("%v%s", v, sep)
	}
	if len(s) > 0 {
		return s[:len(s)-len(sep)]
	}
	return ""
}

func _map(args ...interface{}) (map[string]interface{}, error) {
	var key string
	m := make(map[string]interface{})
	for ii, v := range args {
		if ii%2 == 0 {
			if s, ok := v.(string); ok {
				key = s
			} else {
				return nil, fmt.Errorf("Invalid argument to map at index %d, %t instead of string", ii, v)
			}
		} else {
			m[key] = v
		}
	}
	return m, nil
}

func mult(args ...interface{}) (float64, error) {
	val := 1.0
	for ii, v := range args {
		value := reflect.ValueOf(v)
		switch value.Kind() {
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			val *= float64(value.Int())
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
			val *= float64(value.Uint())
		case reflect.Float32, reflect.Float64:
			val *= value.Float()
		case reflect.String:
			v, err := strconv.ParseFloat(value.String(), 64)
			if err != nil {
				return 0, fmt.Errorf("Error parsing string passed to mult at index %d: %s", ii, err)
			}
			val *= v
		default:
			return 0, fmt.Errorf("Invalid argument of type %T passed to mult at index %d", v, ii)
		}
	}
	return val, nil

}

func concat(args ...interface{}) string {
	var s []string
	for _, v := range args {
		s = append(s, fmt.Sprintf("%v", v))
	}
	return strings.Join(s, "")
}

var templateFuncs template.FuncMap = template.FuncMap{
	"eq":     eq,
	"neq":    neq,
	"json":   _json,
	"nz":     nz,
	"lower":  lower,
	"join":   join,
	"map":    _map,
	"mult":   mult,
	"render": assets.Render,
	"concat": concat,
}
<|endoftext|>


<commit_before>package fastq

import (
	"bytes"
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"os"
	"strings"
	"sync"
)

// Fastq (Name, Seq, Qual)
type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) IsFilter() bool {
	return strings.Contains(fq.Name, ":Y:")
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  string
	seq   []byte
	qual  []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}
	if filename == "-" {
		filename = "STDIN"
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = bytes.TrimSpace(ff.s.Bytes())
		if len(line) == 0 { // ingore empty line
			continue
		}
		switch ff.stage {
		case 0: // get fastq name
			if len(line) > 0 && line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = string(line[1:])
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if len(line) > 0 && line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual (%s) length (%d) != seq (%s) length (%d) at line: %d",
					ff.Name, string(ff.name), ff.qual, len(ff.qual), ff.seq, len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}
	if len(ff.qual) < len(ff.seq) {
		ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
			ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
	}
	ff.setErr(io.EOF)
	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: ff.name, Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
		}
		close(ch)
	}(ch)
	return ch
}

func Opens(filenames ...string) ([]*FastqFile, error) {
	fqfiles := make([]*FastqFile, len(filenames))
	for i, filename := range filenames {
		fqfile, err := Open(filename)
		if err != nil {
			return nil, err
		}
		fqfiles[i] = fqfile
	}
	return fqfiles, nil
}

func Load(filenames ...string) (<-chan *Fastq, <-chan error) {
	fqChan := make(chan *Fastq, 2*len(filenames))
	errChan := make(chan error, 1)

	fqfiles, err := Opens(filenames...)
	if err != nil {
		errChan <- err
		return nil, errChan
	}

	go func(fqChan chan *Fastq, errChan chan error, fqfiles []*FastqFile) {
		for _, fqfile := range fqfiles {
			defer fqfile.Close()
			for fqfile.Next() {
				fqChan <- fqfile.Value()
			}
			if err := fqfile.Err(); err != nil {
				errChan <- err
			}
		}
		close(fqChan)
		close(errChan)
	}(fqChan, errChan, fqfiles)
	return fqChan, errChan
}

func LoadMix(filenames ...string) (<-chan *Fastq, <-chan error) {
	fqChan := make(chan *Fastq, 2*len(filenames))
	errChan := make(chan error, 1)

	fqfiles, err := Opens(filenames...)
	if err != nil {
		errChan <- err
		return nil, errChan
	}

	go func(fqfiles []*FastqFile, fqChan chan *Fastq, errChan chan error) {
		wg := &sync.WaitGroup{}
		for _, fqfile := range fqfiles {
			wg.Add(1)
			go func(wg *sync.WaitGroup, fqChan chan *Fastq, errChan chan error, fqfile *FastqFile) {
				defer fqfile.Close()
				for fqfile.Next() {
					fmt.Fprintln(os.Stderr, fqfile.Value())
					fqChan <- fqfile.Value()
				}
				if err := fqfile.Err(); err != nil {
					errChan <- err
				}
				wg.Done()
			}(wg, fqChan, errChan, fqfile)
		}
		wg.Wait()
		close(fqChan)
	}(fqfiles, fqChan, errChan)
	return fqChan, errChan
}
<commit_msg>update fastq.go<commit_after>package fastq

import (
	"bytes"
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"os"
	"strings"
	"sync"
)

// Fastq (Name, Seq, Qual)
type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) IsFilter() bool {
	return strings.Contains(fq.Name, ":Y:")
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  string
	seq   []byte
	qual  []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}
	if filename == "-" {
		filename = "STDIN"
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = bytes.TrimSpace(ff.s.Bytes())
		if len(line) == 0 { // ingore empty line
			continue
		}
		switch ff.stage {
		case 0: // get fastq name
			if len(line) > 0 && line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = string(line[1:])
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if len(line) > 0 && line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual (%s) length (%d) != seq (%s) length (%d) at line: %d",
					ff.Name, string(ff.name), ff.qual, len(ff.qual), ff.seq, len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}
	if len(ff.qual) < len(ff.seq) {
		ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
			ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
	}
	ff.setErr(io.EOF)
	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: ff.name, Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
		}
		close(ch)
	}(ch)
	return ch
}

func Opens(filenames ...string) ([]*FastqFile, error) {
	if len(filenames) == 0 {
		return nil, ErrEmptyInputFile
	}

	fqfiles := make([]*FastqFile, len(filenames))
	for i, filename := range filenames {
		fqfile, err := Open(filename)
		if err != nil {
			return nil, err
		}
		fqfiles[i] = fqfile
	}
	return fqfiles, nil
}

func Load(filenames ...string) (<-chan *Fastq, <-chan error) {
	fqChan := make(chan *Fastq, 2*len(filenames))
	errChan := make(chan error, 1)

	fqfiles, err := Opens(filenames...)
	if err != nil {
		errChan <- err
		return nil, errChan
	}

	go func(fqChan chan *Fastq, errChan chan error, fqfiles []*FastqFile) {
		for _, fqfile := range fqfiles {
			defer fqfile.Close()
			for fqfile.Next() {
				fqChan <- fqfile.Value()
			}
			if err := fqfile.Err(); err != nil {
				errChan <- err
			}
		}
		close(fqChan)
		close(errChan)
	}(fqChan, errChan, fqfiles)
	return fqChan, errChan
}

func LoadMix(filenames ...string) (<-chan *Fastq, <-chan error) {
	fqChan := make(chan *Fastq, 2*len(filenames))
	errChan := make(chan error, 1)

	fqfiles, err := Opens(filenames...)
	if err != nil {
		errChan <- err
		return nil, errChan
	}

	go func(fqfiles []*FastqFile, fqChan chan *Fastq, errChan chan error) {
		wg := &sync.WaitGroup{}
		for _, fqfile := range fqfiles {
			wg.Add(1)
			go func(wg *sync.WaitGroup, fqChan chan *Fastq, errChan chan error, fqfile *FastqFile) {
				defer fqfile.Close()
				for fqfile.Next() {
					fmt.Fprintln(os.Stderr, fqfile.Value())
					fqChan <- fqfile.Value()
				}
				if err := fqfile.Err(); err != nil {
					errChan <- err
				}
				wg.Done()
			}(wg, fqChan, errChan, fqfile)
		}
		wg.Wait()
		close(fqChan)
	}(fqfiles, fqChan, errChan)
	return fqChan, errChan
}
<|endoftext|>


<commit_before>package fastq

import (
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"strings"
)

type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  []byte
	seq   []byte
	qual  []byte
	buf   []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
		buf:  make([]byte, 1024),
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = ff.s.Bytes()
		switch ff.stage {
		case 0: // get fastq name
			if line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = line[1:]
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
					ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}
	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: string(ff.name), Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
<<<<<<< HEAD
	go func(ch chan *Fastq, fqfiles []*FastqFile) {
		wg := &sync.WaitGroup{}
		for _, fqfile := range fqfiles {
			wg.Add(1)
			go func(ch chan *Fastq, fqfile *FastqFile, wg *sync.WaitGroup) {
				defer fqfile.Close()
				for fq := range fqfile.Load() {
					ch <- fq
				}
				wg.Done()
			}(ch, fqfile, wg)
=======
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
>>>>>>> align
		}
		close(ch)
	}(ch)
	return ch
}
<commit_msg>clean code<commit_after>package fastq

import (
	"fmt"
	"gongs/scan"
	"gongs/xopen"
	"io"
	"strings"
)

type Fastq struct {
	Name string
	Seq  []byte
	Qual []byte
}

func (fq Fastq) String() string {
	return fmt.Sprintf("@%v\n%v\n+\n%v", fq.Name, string(fq.Seq), string(fq.Qual))
}

func (fq Fastq) Id() string {
	if n := strings.IndexByte(fq.Name, ' '); n >= 0 {
		return fq.Name[:n]
	}

	// for old solexa data format
	if n := strings.IndexByte(fq.Name, '#'); n >= 0 {
		return fq.Name[:n]
	}

	return fq.Name
}

type FastqFile struct {
	Name  string
	file  io.ReadCloser
	s     *scan.Scanner
	name  []byte
	seq   []byte
	qual  []byte
	buf   []byte
	err   error
	stage int
}

func Open(filename string) (*FastqFile, error) {
	file, err := xopen.Xopen(filename)
	if err != nil {
		return nil, err
	}

	return &FastqFile{
		Name: filename,
		s:    scan.New(file),
		file: file,
		buf:  make([]byte, 1024),
	}, nil
}

func (ff *FastqFile) Close() error {
	return ff.file.Close()
}

func (ff *FastqFile) Err() error {
	if ff.err == nil || ff.err == io.EOF {
		if err := ff.s.Err(); err != nil {
			return err
		}
		return nil
	}
	return ff.err
}

func (ff *FastqFile) setErr(err error) {
	if ff.err == nil {
		ff.err = err
	}
}

func (ff *FastqFile) Next() bool {
	if ff.err != nil {
		return false
	}

	var line []byte
	for ff.s.Scan() {
		line = ff.s.Bytes()
		switch ff.stage {
		case 0: // get fastq name
			if len(line) > 0 && line[0] != '@' {
				ff.setErr(fmt.Errorf("file: %v Wrong Fastq Record Name %s at line: %d", ff.Name, string(line), ff.s.Lid()))
				return false
			}
			ff.stage++
			ff.name = line[1:]
			ff.seq = ff.seq[:0]   // clear seq
			ff.qual = ff.qual[:0] // clear qual
		case 1: // get fastq seq
			if len(line) > 0 && line[0] == '+' {
				ff.stage += 2
				break
			}
			ff.seq = append(ff.seq, line...)
		case 2: // get + line
		case 3: // get fastq qual
			ff.qual = append(ff.qual, line...)
			if len(ff.qual) == len(ff.seq) {
				ff.stage = 0
				return true
			} else if len(ff.qual) > len(ff.seq) {
				ff.setErr(fmt.Errorf("file: %v Fastq Record (%s) qual length (%d) != seq length (%d) at line: %d",
					ff.Name, string(ff.name), len(ff.qual), len(ff.seq), ff.s.Lid()))
				return false
			}
		}
	}

	return false
}

func (ff *FastqFile) Value() *Fastq {
	return &Fastq{Name: string(ff.name), Seq: ff.seq, Qual: ff.qual}
}

func (ff *FastqFile) Iter() <-chan *Fastq {
	ch := make(chan *Fastq)
	go func(ch chan *Fastq) {
		for ff.Next() {
			ch <- ff.Value()
		}
		close(ch)
	}(ch)
	return ch
}
<|endoftext|>


<commit_before>package boil

import (
	"bytes"
	"fmt"
	"regexp"
	"sort"
	"strings"

	"github.com/vattle/sqlboiler/strmangle"
)

var (
	rgxIdentifier = regexp.MustCompile(`^(?i)"?[a-z_][_a-z0-9]*"?(?:\."?[_a-z][_a-z0-9]*"?)*$`)
)

func buildQuery(q *Query) (string, []interface{}) {
	var buf *bytes.Buffer
	var args []interface{}

	switch {
	case len(q.plainSQL.sql) != 0:
		return q.plainSQL.sql, q.plainSQL.args
	case q.delete:
		buf, args = buildDeleteQuery(q)
	case len(q.update) > 0:
		buf, args = buildUpdateQuery(q)
	default:
		buf, args = buildSelectQuery(q)
	}

	return buf.String(), args
}

func buildSelectQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := &bytes.Buffer{}
	var args []interface{}

	buf.WriteString("SELECT ")

	// Wrap the select in the modifier function
	hasModFunc := len(q.modFunction) != 0
	if hasModFunc {
		fmt.Fprintf(buf, "%s(", q.modFunction)
	}

	hasSelectCols := len(q.selectCols) != 0
	hasJoins := len(q.joins) != 0
	if hasSelectCols && hasJoins && !hasModFunc {
		selectColsWithAs := writeAsStatements(q)
		// Don't identQuoteSlice - writeAsStatements does this
		buf.WriteString(strings.Join(selectColsWithAs, ", "))
	} else if hasSelectCols {
		buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.selectCols), ", "))
	} else if hasJoins {
		selectColsWithStars := writeStars(q)
		buf.WriteString(strings.Join(selectColsWithStars, ", "))
	} else {
		buf.WriteByte('*')
	}

	if hasModFunc {
		buf.WriteByte(')')
	}

	fmt.Fprintf(buf, " FROM %s", strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	if len(q.joins) > 0 {
		argsLen := len(args)
		joinBuf := &bytes.Buffer{}
		for _, j := range q.joins {
			if j.kind != JoinInner {
				panic("only inner joins are supported")
			}
			fmt.Fprintf(joinBuf, " INNER JOIN %s", j.clause)
			args = append(args, j.args...)
		}
		fmt.Fprintf(buf, convertQuestionMarks(joinBuf.String(), argsLen+1))
	}

	where, whereArgs := whereClause(q, len(args)+1)
	buf.WriteString(where)

	args = append(args, whereArgs...)
	writeModifiers(q, buf, &args)

	buf.WriteByte(';')
	return buf, args
}

func buildDeleteQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := &bytes.Buffer{}

	buf.WriteString("DELETE FROM ")
	buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	where, args := whereClause(q, 1)
	buf.WriteString(where)

	writeModifiers(q, buf, &args)

	buf.WriteByte(';')

	return buf, args
}

func buildUpdateQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := &bytes.Buffer{}

	buf.WriteString("UPDATE ")
	buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	cols := make(sort.StringSlice, len(q.update))
	var args []interface{}

	count := 0
	for name := range q.update {
		cols[count] = name
		count++
	}

	cols.Sort()

	for i := 0; i < len(cols); i++ {
		args = append(args, q.update[cols[i]])
		cols[i] = strmangle.IdentQuote(cols[i])
	}

	buf.WriteString(fmt.Sprintf(
		" SET (%s) = (%s)",
		strings.Join(cols, ", "),
		strmangle.Placeholders(len(cols), 1, 1)),
	)

	where, whereArgs := whereClause(q, len(args)+1)
	buf.WriteString(where)

	args = append(args, whereArgs...)
	writeModifiers(q, buf, &args)

	buf.WriteByte(';')

	return buf, args
}

func writeModifiers(q *Query, buf *bytes.Buffer, args *[]interface{}) {
	if len(q.groupBy) != 0 {
		fmt.Fprintf(buf, " GROUP BY %s", strings.Join(q.groupBy, ", "))
	}

	if len(q.having) != 0 {
		argsLen := len(*args)
		havingBuf := &bytes.Buffer{}
		fmt.Fprintf(havingBuf, " HAVING ")
		for i, j := range q.having {
			if i > 0 {
				fmt.Fprintf(havingBuf, ", ")
			}
			fmt.Fprintf(havingBuf, j.clause)
			*args = append(*args, j.args...)
		}
		fmt.Fprintf(buf, convertQuestionMarks(havingBuf.String(), argsLen+1))
	}

	if len(q.orderBy) != 0 {
		buf.WriteString(" ORDER BY ")
		buf.WriteString(strings.Join(q.orderBy, ", "))
	}

	if q.limit != 0 {
		fmt.Fprintf(buf, " LIMIT %d", q.limit)
	}
	if q.offset != 0 {
		fmt.Fprintf(buf, " OFFSET %d", q.offset)
	}
}

func writeStars(q *Query) []string {
	cols := make([]string, 0, len(q.from))
	for _, f := range q.from {
		toks := strings.Split(f, " ")
		if len(toks) == 1 {
			cols = append(cols, fmt.Sprintf(`%s.*`, strmangle.IdentQuote(toks[0])))
			continue
		}

		alias, name, ok := parseFromClause(toks)
		if !ok {
		}

		if len(alias) != 0 {
			name = alias
		}
		cols = append(cols, fmt.Sprintf(`%s.*`, strmangle.IdentQuote(name)))
	}

	return cols
}

func writeAsStatements(q *Query) []string {
	cols := make([]string, len(q.selectCols))
	for i, col := range q.selectCols {
		if !rgxIdentifier.MatchString(col) {
			cols[i] = col
			continue
		}

		toks := strings.Split(col, ".")
		if len(toks) == 1 {
			cols[i] = strmangle.IdentQuote(col)
			continue
		}

		asParts := make([]string, len(toks))
		for j, tok := range toks {
			asParts[j] = strings.Trim(tok, `"`)
		}

		cols[i] = fmt.Sprintf(`%s as "%s"`, strmangle.IdentQuote(col), strings.Join(asParts, "."))
	}

	return cols
}

// whereClause parses a where slice and converts it into a
// single WHERE clause like:
// WHERE (a=$1) AND (b=$2)
//
// startAt specifies what number placeholders start at
func whereClause(q *Query, startAt int) (string, []interface{}) {
	if len(q.where) == 0 {
		return "", nil
	}

	buf := &bytes.Buffer{}
	var args []interface{}

	buf.WriteString(" WHERE ")
	for i := 0; i < len(q.where); i++ {
		buf.WriteString(fmt.Sprintf("(%s)", q.where[i].clause))
		args = append(args, q.where[i].args...)

		// break on the last loop
		if i == len(q.where)-1 {
			break
		}

		if q.where[i].orSeparator {
			buf.WriteString(" OR ")
		} else {
			buf.WriteString(" AND ")
		}
	}

	return convertQuestionMarks(buf.String(), startAt), args
}

func convertQuestionMarks(clause string, startAt int) string {
	if startAt == 0 {
		panic("Not a valid start number.")
	}

	paramBuf := &bytes.Buffer{}
	paramIndex := 0

	for ; ; startAt++ {
		if paramIndex >= len(clause) {
			break
		}

		clause = clause[paramIndex:]
		paramIndex = strings.IndexByte(clause, '?')

		if paramIndex == -1 {
			paramBuf.WriteString(clause)
			break
		}

		paramBuf.WriteString(clause[:paramIndex] + fmt.Sprintf("$%d", startAt))
		paramIndex++
	}

	return paramBuf.String()
}

// identifierMapping creates a map of all identifiers to potential model names
func identifierMapping(q *Query) map[string]string {
	var ids map[string]string
	setID := func(alias, name string) {
		if ids == nil {
			ids = make(map[string]string)
		}
		ids[alias] = name
	}

	for _, from := range q.from {
		tokens := strings.Split(from, " ")
		parseIdentifierClause(tokens, setID)
	}

	for _, join := range q.joins {
		tokens := strings.Split(join.clause, " ")
		parseIdentifierClause(tokens, setID)
	}

	return ids
}

// parseBits takes a set of tokens and looks for something of the form:
// a b
// a as b
// where 'a' and 'b' are valid SQL identifiers
// It only evaluates the first 3 tokens (anything past that is superfluous)
// It stops parsing when it finds "on" or an invalid identifier
func parseIdentifierClause(tokens []string, setID func(string, string)) {
	alias, name, ok := parseFromClause(tokens)
	if !ok {
		panic("could not parse from statement")
	}

	if len(alias) > 0 {
		setID(alias, name)
	} else {
		setID(name, name)
	}
}

func parseFromClause(toks []string) (alias, name string, ok bool) {
	if len(toks) > 3 {
		toks = toks[:3]
	}

	sawIdent, sawAs := false, false
	for _, tok := range toks {
		if t := strings.ToLower(tok); sawIdent && t == "as" {
			sawAs = true
			continue
		} else if sawIdent && t == "on" {
			break
		}

		if !rgxIdentifier.MatchString(tok) {
			break
		}

		if sawIdent || sawAs {
			alias = strings.Trim(tok, `"`)
			break
		}

		name = strings.Trim(tok, `"`)
		sawIdent = true
		ok = true
	}

	return alias, name, ok
}
<commit_msg>Add BufPool to query_builders.go<commit_after>package boil

import (
	"bytes"
	"fmt"
	"regexp"
	"sort"
	"strings"

	"github.com/vattle/sqlboiler/strmangle"
)

var (
	rgxIdentifier = regexp.MustCompile(`^(?i)"?[a-z_][_a-z0-9]*"?(?:\."?[_a-z][_a-z0-9]*"?)*$`)
)

func buildQuery(q *Query) (string, []interface{}) {
	var buf *bytes.Buffer
	var args []interface{}

	switch {
	case len(q.plainSQL.sql) != 0:
		return q.plainSQL.sql, q.plainSQL.args
	case q.delete:
		buf, args = buildDeleteQuery(q)
	case len(q.update) > 0:
		buf, args = buildUpdateQuery(q)
	default:
		buf, args = buildSelectQuery(q)
	}

	defer strmangle.PutBuffer(buf)

	return buf.String(), args
}

func buildSelectQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := strmangle.GetBuffer()
	var args []interface{}

	buf.WriteString("SELECT ")

	// Wrap the select in the modifier function
	hasModFunc := len(q.modFunction) != 0
	if hasModFunc {
		fmt.Fprintf(buf, "%s(", q.modFunction)
	}

	hasSelectCols := len(q.selectCols) != 0
	hasJoins := len(q.joins) != 0
	if hasSelectCols && hasJoins && !hasModFunc {
		selectColsWithAs := writeAsStatements(q)
		// Don't identQuoteSlice - writeAsStatements does this
		buf.WriteString(strings.Join(selectColsWithAs, ", "))
	} else if hasSelectCols {
		buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.selectCols), ", "))
	} else if hasJoins {
		selectColsWithStars := writeStars(q)
		buf.WriteString(strings.Join(selectColsWithStars, ", "))
	} else {
		buf.WriteByte('*')
	}

	if hasModFunc {
		buf.WriteByte(')')
	}

	fmt.Fprintf(buf, " FROM %s", strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	if len(q.joins) > 0 {
		argsLen := len(args)
		joinBuf := strmangle.GetBuffer()
		for _, j := range q.joins {
			if j.kind != JoinInner {
				panic("only inner joins are supported")
			}
			fmt.Fprintf(joinBuf, " INNER JOIN %s", j.clause)
			args = append(args, j.args...)
		}
		fmt.Fprintf(buf, convertQuestionMarks(joinBuf.String(), argsLen+1))
		strmangle.PutBuffer(joinBuf)
	}

	where, whereArgs := whereClause(q, len(args)+1)
	buf.WriteString(where)

	args = append(args, whereArgs...)
	writeModifiers(q, buf, &args)

	buf.WriteByte(';')
	return buf, args
}

func buildDeleteQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := strmangle.GetBuffer()

	buf.WriteString("DELETE FROM ")
	buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	where, args := whereClause(q, 1)
	buf.WriteString(where)

	writeModifiers(q, buf, &args)

	buf.WriteByte(';')

	return buf, args
}

func buildUpdateQuery(q *Query) (*bytes.Buffer, []interface{}) {
	buf := strmangle.GetBuffer()

	buf.WriteString("UPDATE ")
	buf.WriteString(strings.Join(strmangle.IdentQuoteSlice(q.from), ", "))

	cols := make(sort.StringSlice, len(q.update))
	var args []interface{}

	count := 0
	for name := range q.update {
		cols[count] = name
		count++
	}

	cols.Sort()

	for i := 0; i < len(cols); i++ {
		args = append(args, q.update[cols[i]])
		cols[i] = strmangle.IdentQuote(cols[i])
	}

	buf.WriteString(fmt.Sprintf(
		" SET (%s) = (%s)",
		strings.Join(cols, ", "),
		strmangle.Placeholders(len(cols), 1, 1)),
	)

	where, whereArgs := whereClause(q, len(args)+1)
	buf.WriteString(where)

	args = append(args, whereArgs...)
	writeModifiers(q, buf, &args)

	buf.WriteByte(';')

	return buf, args
}

func writeModifiers(q *Query, buf *bytes.Buffer, args *[]interface{}) {
	if len(q.groupBy) != 0 {
		fmt.Fprintf(buf, " GROUP BY %s", strings.Join(q.groupBy, ", "))
	}

	if len(q.having) != 0 {
		argsLen := len(*args)
		havingBuf := strmangle.GetBuffer()
		fmt.Fprintf(havingBuf, " HAVING ")
		for i, j := range q.having {
			if i > 0 {
				fmt.Fprintf(havingBuf, ", ")
			}
			fmt.Fprintf(havingBuf, j.clause)
			*args = append(*args, j.args...)
		}
		fmt.Fprintf(buf, convertQuestionMarks(havingBuf.String(), argsLen+1))
		strmangle.PutBuffer(havingBuf)
	}

	if len(q.orderBy) != 0 {
		buf.WriteString(" ORDER BY ")
		buf.WriteString(strings.Join(q.orderBy, ", "))
	}

	if q.limit != 0 {
		fmt.Fprintf(buf, " LIMIT %d", q.limit)
	}
	if q.offset != 0 {
		fmt.Fprintf(buf, " OFFSET %d", q.offset)
	}
}

func writeStars(q *Query) []string {
	cols := make([]string, 0, len(q.from))
	for _, f := range q.from {
		toks := strings.Split(f, " ")
		if len(toks) == 1 {
			cols = append(cols, fmt.Sprintf(`%s.*`, strmangle.IdentQuote(toks[0])))
			continue
		}

		alias, name, ok := parseFromClause(toks)
		if !ok {
		}

		if len(alias) != 0 {
			name = alias
		}
		cols = append(cols, fmt.Sprintf(`%s.*`, strmangle.IdentQuote(name)))
	}

	return cols
}

func writeAsStatements(q *Query) []string {
	cols := make([]string, len(q.selectCols))
	for i, col := range q.selectCols {
		if !rgxIdentifier.MatchString(col) {
			cols[i] = col
			continue
		}

		toks := strings.Split(col, ".")
		if len(toks) == 1 {
			cols[i] = strmangle.IdentQuote(col)
			continue
		}

		asParts := make([]string, len(toks))
		for j, tok := range toks {
			asParts[j] = strings.Trim(tok, `"`)
		}

		cols[i] = fmt.Sprintf(`%s as "%s"`, strmangle.IdentQuote(col), strings.Join(asParts, "."))
	}

	return cols
}

// whereClause parses a where slice and converts it into a
// single WHERE clause like:
// WHERE (a=$1) AND (b=$2)
//
// startAt specifies what number placeholders start at
func whereClause(q *Query, startAt int) (string, []interface{}) {
	if len(q.where) == 0 {
		return "", nil
	}

	buf := strmangle.GetBuffer()
	defer strmangle.PutBuffer(buf)
	var args []interface{}

	buf.WriteString(" WHERE ")
	for i := 0; i < len(q.where); i++ {
		buf.WriteString(fmt.Sprintf("(%s)", q.where[i].clause))
		args = append(args, q.where[i].args...)

		// break on the last loop
		if i == len(q.where)-1 {
			break
		}

		if q.where[i].orSeparator {
			buf.WriteString(" OR ")
		} else {
			buf.WriteString(" AND ")
		}
	}

	return convertQuestionMarks(buf.String(), startAt), args
}

func convertQuestionMarks(clause string, startAt int) string {
	if startAt == 0 {
		panic("Not a valid start number.")
	}

	paramBuf := strmangle.GetBuffer()
	defer strmangle.PutBuffer(paramBuf)
	paramIndex := 0

	for ; ; startAt++ {
		if paramIndex >= len(clause) {
			break
		}

		clause = clause[paramIndex:]
		paramIndex = strings.IndexByte(clause, '?')

		if paramIndex == -1 {
			paramBuf.WriteString(clause)
			break
		}

		paramBuf.WriteString(clause[:paramIndex] + fmt.Sprintf("$%d", startAt))
		paramIndex++
	}

	return paramBuf.String()
}

// identifierMapping creates a map of all identifiers to potential model names
func identifierMapping(q *Query) map[string]string {
	var ids map[string]string
	setID := func(alias, name string) {
		if ids == nil {
			ids = make(map[string]string)
		}
		ids[alias] = name
	}

	for _, from := range q.from {
		tokens := strings.Split(from, " ")
		parseIdentifierClause(tokens, setID)
	}

	for _, join := range q.joins {
		tokens := strings.Split(join.clause, " ")
		parseIdentifierClause(tokens, setID)
	}

	return ids
}

// parseBits takes a set of tokens and looks for something of the form:
// a b
// a as b
// where 'a' and 'b' are valid SQL identifiers
// It only evaluates the first 3 tokens (anything past that is superfluous)
// It stops parsing when it finds "on" or an invalid identifier
func parseIdentifierClause(tokens []string, setID func(string, string)) {
	alias, name, ok := parseFromClause(tokens)
	if !ok {
		panic("could not parse from statement")
	}

	if len(alias) > 0 {
		setID(alias, name)
	} else {
		setID(name, name)
	}
}

func parseFromClause(toks []string) (alias, name string, ok bool) {
	if len(toks) > 3 {
		toks = toks[:3]
	}

	sawIdent, sawAs := false, false
	for _, tok := range toks {
		if t := strings.ToLower(tok); sawIdent && t == "as" {
			sawAs = true
			continue
		} else if sawIdent && t == "on" {
			break
		}

		if !rgxIdentifier.MatchString(tok) {
			break
		}

		if sawIdent || sawAs {
			alias = strings.Trim(tok, `"`)
			break
		}

		name = strings.Trim(tok, `"`)
		sawIdent = true
		ok = true
	}

	return alias, name, ok
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/AlecAivazis/survey"
	surveyCore "github.com/AlecAivazis/survey/core"
	"github.com/cheggaaa/pb"
	"github.com/tarm/serial"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	port     = kingpin.Flag("port", "Serial port.").Short('p').String()
	filename = kingpin.Flag("image", "Firmware image.").Short('i').ExistingFile()
	force    = kingpin.Flag("force", "Force flash.").Short('f').Bool()
	verbose  = kingpin.Flag("verbose", "Verbose mode.").Short('v').Bool()
)

func make_checksum(payload []byte) []byte {
	var checksum uint16 = 0xFFFF
	for i := 0; i < len(payload); i++ {
		checksum -= uint16(payload[i])
	}
	ret := make([]byte, 2)
	binary.LittleEndian.PutUint16(ret, checksum)
	return ret
}

func WriteAll(s *serial.Port, raw []byte) error {
	if *verbose {
		fmt.Println("Write", raw[:2], raw[2:len(raw)-2], raw[len(raw)-2:])
	}
	n, err := s.Write(raw)
	if err != nil {
		return err
	}
	if n != len(raw) {
		return errors.New("Didn't write all bytes.")
	}
	return nil
}

func WriteFrame(s *serial.Port, payload []byte) error {
	length := make([]byte, 2)
	binary.LittleEndian.PutUint16(length, uint16(len(payload)+4))
	frame := append(length, payload...)
	frame = append(frame, make_checksum(frame)...)
	return WriteAll(s, frame)
}

func ReadAll(s *serial.Port, n int) ([]byte, error) {
	buf := make([]byte, n)
	bytes_read := 0
	for bytes_read < n {
		c, err := s.Read(buf[bytes_read:])
		if err != nil {
			return nil, err
		}
		if c == 0 {
			return nil, errors.New("Read timeout")
		}
		bytes_read += c
	}
	return buf, nil
}

func EmptyRx(s *serial.Port) {
	c := 1
	buf := make([]byte, 1024)
	for c > 0 {
		c, _ = s.Read(buf)
	}
}

func ReadFrame(s *serial.Port) ([]byte, error) {
	head, err := ReadAll(s, 3)
	if err != nil {
		return nil, err
	}
	if head[0] != 0x55 {
		return nil, errors.New("Invalid response")
	}
	size := int(binary.LittleEndian.Uint16(head[1:]))
	body, err := ReadAll(s, size-3)
	if err != nil {
		return nil, err
	}
	payload := body[:len(body)-2]
	checksum := body[len(body)-2:]
	if *verbose {
		fmt.Println("Read", head, payload, checksum)
	}
	checksum_cmp := make_checksum(append(head, payload...))
	if !bytes.Equal(checksum, checksum_cmp) {
		return nil, errors.New("Invalid checksum")
	}
	return payload, nil
}

func ping(s *serial.Port) ([]byte, error) {
	err := WriteFrame(s, []byte{0xC0})
	if err != nil {
		return nil, err
	}
	answer, err := ReadFrame(s)
	if err != nil {
		return nil, err
	}
	if answer[0] != 0xC0 {
		return nil, errors.New("Unexpected answer to ping")
	}
	return answer, nil
}

func communicate(s *serial.Port, request []byte, response []byte) error {
	err := WriteFrame(s, request)
	if err != nil {
		return err
	}
	msg, err := ReadFrame(s)
	if err != nil {
		return err
	}
	if !bytes.Equal(msg, response) {
		errors.New("Unexpected response: " + hex.Dump(response))
	}
	return nil
}

func ask_write(s *serial.Port, address int) error {
	ask_permission := []byte{0xc2, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(ask_permission[1:3], uint16(address))
	get_permission := []byte{0xc2, 0x80, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(get_permission[2:4], uint16(address))

	return communicate(s, ask_permission, get_permission)
}

func write_chunk(s *serial.Port, address int, data []byte) error {
	write_instruction := []byte{0xc3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(write_instruction[1:3], uint16(address))
	binary.LittleEndian.PutUint16(write_instruction[5:7], uint16(len(data)))
	write_instruction = append(write_instruction, data...)
	write_confirmation := []byte{0xc3, 0x00, 0x00, 0x00, 0x00}

	return communicate(s, write_instruction, write_confirmation)
}

func update(s *serial.Port, firmware []byte) error {
	start_address := 0x1800

	bar := pb.New(len(firmware)).SetUnits(pb.U_BYTES)
	bar.Start()

	for bytes_written := 0; bytes_written < len(firmware); bytes_written += 1024 {
		var err error
		tries := -1
	try_write:
		tries++
		if 0 < tries && tries < 3 {
			EmptyRx(s)
		} else if tries == 3 {
			return err
		}

		err = ask_write(s, start_address+bytes_written)
		if err != nil {
			goto try_write
		}
		for chunk := 0; chunk < 1024; chunk += 256 {
			offset := bytes_written + chunk

			bytes_to_write := 256
			// are we at the end of the file?
			if len(firmware)-offset < 256 {
				bytes_to_write = len(firmware) - offset
			}
			if bytes_to_write <= 0 {
				break
			}
			err = write_chunk(s, start_address+offset, firmware[offset:offset+bytes_to_write])
			if err != nil {
				goto try_write
			}
			bar.Add(bytes_to_write)
		}
	}

	bar.FinishPrint("Upload completed.")
	return nil
}

func restart(s *serial.Port) error {
	return WriteFrame(s, []byte{0xC1, 0x00})
}

func choose(item string, options []string) (string, error) {
	if len(options) == 1 {
		fmt.Printf("Using %s as %s.\n", options[0], item)
		return options[0], nil
	}

	if runtime.GOOS == "windows" {
		surveyCore.SelectFocusIcon = ">"
	}

	ret := ""
	message := fmt.Sprintf("%s:", strings.Title(item))
	if len(options) == 0 {
		fmt.Printf("Could not autodetect %s. Please enter manually.\n", item)
		prompt := &survey.Input{Message: message}
		survey.AskOne(prompt, &ret, nil)
	} else {
		prompt := &survey.Select{
			Message: message,
			Options: options,
		}
		survey.AskOne(prompt, &ret, nil)
	}

	if ret == "" {
		return "", errors.New("interrupted")
	}
	return ret, nil
}

func main() {
	kingpin.Parse()
	var err error

	if *filename == "" {
		binFiles, _ := filepath.Glob("*.bin")
		hexFiles, _ := filepath.Glob("*.hex")
		*filename, err = choose("firmware image", append(binFiles, hexFiles...))
		if err != nil {
			fmt.Printf("No firmware image selected: %s.\n", err)
			return
		}
	}

	data, err := ioutil.ReadFile(*filename)
	if err != nil {
		fmt.Printf("Cannot read firmware image %s: %s.\n", *filename, err)
		return
	}

	if (len(data) < 0x9000 || len(data) > 0xe7ff) && !*force {
		fmt.Printf("Unexpected firmare size: %d bytes. Use --force to flash anyway.\n", len(data))
		return
	}

	if *port == "" {
		var serialPortFmt string
		if runtime.GOOS == "windows" {
			serialPortFmt = "COM%d"
		} else {
			serialPortFmt = "/dev/ttyUSB%d"
		}

		var candidates []string
		for i := 0; i < 255; i += 1 {
			candidate := fmt.Sprintf(serialPortFmt, i)
			c := &serial.Config{Name: candidate, Baud: 115200, ReadTimeout: time.Second * 1}
			s, err := serial.OpenPort(c)
			if s != nil {
				s.Close()
			}
			if err == nil {
				candidates = append(candidates, candidate)
			}
		}

		*port, err = choose("serial port", candidates)
		if err != nil {
			fmt.Printf("No serial port selected: %s.\n", err)
			return
		}

		if len(candidates) == 1 {
			start := false
			prompt := &survey.Confirm{
				Message: "Start flashing?",
			}
			survey.AskOne(prompt, &start, nil)
			if start == false {
				fmt.Println("Aborted.")
				return
			}
		}
	}

	c := &serial.Config{Name: *port, Baud: 115200, ReadTimeout: time.Second * 1}
	s, err := serial.OpenPort(c)
	if err != nil {
		fmt.Printf("Cannot open serial port %s: %s.\n", *port, err)
		return
	}

	_, err = ping(s)
	if err != nil {
		fmt.Printf("Cannot ping the remote: %s.\n", err)
		return
	}

	err = update(s, data)
	if err != nil {
		fmt.Printf("Error flashing firmware: %s.\n", err)
		return
	}

	err = restart(s)
	if err != nil {
		fmt.Printf("Error restarting the remote: %s.\n", err)
		return
	}

	fmt.Println("Success!")

	s.Close()
}
<commit_msg>improve ping timeout error message<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/AlecAivazis/survey"
	surveyCore "github.com/AlecAivazis/survey/core"
	"github.com/cheggaaa/pb"
	"github.com/tarm/serial"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	port     = kingpin.Flag("port", "Serial port.").Short('p').String()
	filename = kingpin.Flag("image", "Firmware image.").Short('i').ExistingFile()
	force    = kingpin.Flag("force", "Force flash.").Short('f').Bool()
	verbose  = kingpin.Flag("verbose", "Verbose mode.").Short('v').Bool()
)

func make_checksum(payload []byte) []byte {
	var checksum uint16 = 0xFFFF
	for i := 0; i < len(payload); i++ {
		checksum -= uint16(payload[i])
	}
	ret := make([]byte, 2)
	binary.LittleEndian.PutUint16(ret, checksum)
	return ret
}

func WriteAll(s *serial.Port, raw []byte) error {
	if *verbose {
		fmt.Println("Write", raw[:2], raw[2:len(raw)-2], raw[len(raw)-2:])
	}
	n, err := s.Write(raw)
	if err != nil {
		return err
	}
	if n != len(raw) {
		return errors.New("Didn't write all bytes.")
	}
	return nil
}

func WriteFrame(s *serial.Port, payload []byte) error {
	length := make([]byte, 2)
	binary.LittleEndian.PutUint16(length, uint16(len(payload)+4))
	frame := append(length, payload...)
	frame = append(frame, make_checksum(frame)...)
	return WriteAll(s, frame)
}

func ReadAll(s *serial.Port, n int) ([]byte, error) {
	buf := make([]byte, n)
	bytes_read := 0
	for bytes_read < n {
		c, err := s.Read(buf[bytes_read:])
		if err != nil {
			return nil, err
		}
		if c == 0 {
			return nil, errors.New("Read timeout")
		}
		bytes_read += c
	}
	return buf, nil
}

func EmptyRx(s *serial.Port) {
	c := 1
	buf := make([]byte, 1024)
	for c > 0 {
		c, _ = s.Read(buf)
	}
}

func ReadFrame(s *serial.Port) ([]byte, error) {
	head, err := ReadAll(s, 3)
	if err != nil {
		return nil, err
	}
	if head[0] != 0x55 {
		return nil, errors.New("Invalid response")
	}
	size := int(binary.LittleEndian.Uint16(head[1:]))
	body, err := ReadAll(s, size-3)
	if err != nil {
		return nil, err
	}
	payload := body[:len(body)-2]
	checksum := body[len(body)-2:]
	if *verbose {
		fmt.Println("Read", head, payload, checksum)
	}
	checksum_cmp := make_checksum(append(head, payload...))
	if !bytes.Equal(checksum, checksum_cmp) {
		return nil, errors.New("Invalid checksum")
	}
	return payload, nil
}

func ping(s *serial.Port) ([]byte, error) {
	err := WriteFrame(s, []byte{0xC0})
	if err != nil {
		return nil, err
	}
	answer, err := ReadFrame(s)
	if err != nil {
		return nil, err
	}
	if answer[0] != 0xC0 {
		return nil, errors.New("Unexpected answer to ping")
	}
	return answer, nil
}

func communicate(s *serial.Port, request []byte, response []byte) error {
	err := WriteFrame(s, request)
	if err != nil {
		return err
	}
	msg, err := ReadFrame(s)
	if err != nil {
		return err
	}
	if !bytes.Equal(msg, response) {
		errors.New("Unexpected response: " + hex.Dump(response))
	}
	return nil
}

func ask_write(s *serial.Port, address int) error {
	ask_permission := []byte{0xc2, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(ask_permission[1:3], uint16(address))
	get_permission := []byte{0xc2, 0x80, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(get_permission[2:4], uint16(address))

	return communicate(s, ask_permission, get_permission)
}

func write_chunk(s *serial.Port, address int, data []byte) error {
	write_instruction := []byte{0xc3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
	binary.LittleEndian.PutUint16(write_instruction[1:3], uint16(address))
	binary.LittleEndian.PutUint16(write_instruction[5:7], uint16(len(data)))
	write_instruction = append(write_instruction, data...)
	write_confirmation := []byte{0xc3, 0x00, 0x00, 0x00, 0x00}

	return communicate(s, write_instruction, write_confirmation)
}

func update(s *serial.Port, firmware []byte) error {
	start_address := 0x1800

	bar := pb.New(len(firmware)).SetUnits(pb.U_BYTES)
	bar.Start()

	for bytes_written := 0; bytes_written < len(firmware); bytes_written += 1024 {
		var err error
		tries := -1
	try_write:
		tries++
		if 0 < tries && tries < 3 {
			EmptyRx(s)
		} else if tries == 3 {
			return err
		}

		err = ask_write(s, start_address+bytes_written)
		if err != nil {
			goto try_write
		}
		for chunk := 0; chunk < 1024; chunk += 256 {
			offset := bytes_written + chunk

			bytes_to_write := 256
			// are we at the end of the file?
			if len(firmware)-offset < 256 {
				bytes_to_write = len(firmware) - offset
			}
			if bytes_to_write <= 0 {
				break
			}
			err = write_chunk(s, start_address+offset, firmware[offset:offset+bytes_to_write])
			if err != nil {
				goto try_write
			}
			bar.Add(bytes_to_write)
		}
	}

	bar.FinishPrint("Upload completed.")
	return nil
}

func restart(s *serial.Port) error {
	return WriteFrame(s, []byte{0xC1, 0x00})
}

func choose(item string, options []string) (string, error) {
	if len(options) == 1 {
		fmt.Printf("Using %s as %s.\n", options[0], item)
		return options[0], nil
	}

	if runtime.GOOS == "windows" {
		surveyCore.SelectFocusIcon = ">"
	}

	ret := ""
	message := fmt.Sprintf("%s:", strings.Title(item))
	if len(options) == 0 {
		fmt.Printf("Could not autodetect %s. Please enter manually.\n", item)
		prompt := &survey.Input{Message: message}
		survey.AskOne(prompt, &ret, nil)
	} else {
		prompt := &survey.Select{
			Message: message,
			Options: options,
		}
		survey.AskOne(prompt, &ret, nil)
	}

	if ret == "" {
		return "", errors.New("interrupted")
	}
	return ret, nil
}

func main() {
	kingpin.Parse()
	var err error

	if *filename == "" {
		binFiles, _ := filepath.Glob("*.bin")
		hexFiles, _ := filepath.Glob("*.hex")
		*filename, err = choose("firmware image", append(binFiles, hexFiles...))
		if err != nil {
			fmt.Printf("No firmware image selected: %s.\n", err)
			return
		}
	}

	data, err := ioutil.ReadFile(*filename)
	if err != nil {
		fmt.Printf("Cannot read firmware image %s: %s.\n", *filename, err)
		return
	}

	if (len(data) < 0x9000 || len(data) > 0xe7ff) && !*force {
		fmt.Printf("Unexpected firmare size: %d bytes. Use --force to flash anyway.\n", len(data))
		return
	}

	if *port == "" {
		var serialPortFmt string
		if runtime.GOOS == "windows" {
			serialPortFmt = "COM%d"
		} else {
			serialPortFmt = "/dev/ttyUSB%d"
		}

		var candidates []string
		for i := 0; i < 255; i += 1 {
			candidate := fmt.Sprintf(serialPortFmt, i)
			c := &serial.Config{Name: candidate, Baud: 115200, ReadTimeout: time.Second * 1}
			s, err := serial.OpenPort(c)
			if s != nil {
				s.Close()
			}
			if err == nil {
				candidates = append(candidates, candidate)
			}
		}

		*port, err = choose("serial port", candidates)
		if err != nil {
			fmt.Printf("No serial port selected: %s.\n", err)
			return
		}

		if len(candidates) == 1 {
			start := false
			prompt := &survey.Confirm{
				Message: "Start flashing?",
			}
			survey.AskOne(prompt, &start, nil)
			if start == false {
				fmt.Println("Aborted.")
				return
			}
		}
	}

	c := &serial.Config{Name: *port, Baud: 115200, ReadTimeout: time.Second * 1}
	s, err := serial.OpenPort(c)
	if err != nil {
		fmt.Printf("Cannot open serial port %s: %s.\n", *port, err)
		return
	}

	_, err = ping(s)
	if err != nil {
		fmt.Printf("Cannot ping the remote: %s. Is the remote in firmware update mode?\n", err)
		return
	}

	err = update(s, data)
	if err != nil {
		fmt.Printf("Error flashing firmware: %s.\n", err)
		return
	}

	err = restart(s)
	if err != nil {
		fmt.Printf("Error restarting the remote: %s.\n", err)
		return
	}

	fmt.Println("Success!")

	s.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"strconv"
)

type config struct {
	path                       string
	port                       int
	oortValueSyndicate         string
	oortGroupSyndicate         string
	insecureSkipVerify         bool
	skipMutualTLS              bool
	nodeId                     int
	metricsAddr                string
	metricsCollectors          string
	concurrentRequestsPerStore int
	debug                      bool
}

func resolveConfig(c *config) *config {
	cfg := &config{}
	if c != nil {
		*cfg = *c
	}
	if env := os.Getenv("FORMICD_PATH"); env != "" {
		cfg.path = env
	}
	if cfg.path == "" {
		cfg.path = "/var/lib/formic"
	}
	if env := os.Getenv("FORMICD_PORT"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.port = val
		}
	}
	if cfg.port == 0 {
		cfg.port = 8445
	}
	if env := os.Getenv("FORMICD_OORT_VALUE_SYNDICATE"); env != "" {
		log.Println("Value: ", env)
		cfg.oortValueSyndicate = env
	}
	// cfg.oortValueSyndicate == "" means default SRV resolution.
	if env := os.Getenv("FORMICD_OORT_GROUP_SYNDICATE"); env != "" {
		log.Println("Group: ", env)
		cfg.oortGroupSyndicate = env
	}
	// cfg.oortGroupSyndicate == "" means default SRV resolution.
	if env := os.Getenv("FORMICD_INSECURE_SKIP_VERIFY"); env == "true" {
		cfg.insecureSkipVerify = true
	}
	if env := os.Getenv("FORMICD_SKIP_MUTUAL_TLS"); env == "true" {
		cfg.skipMutualTLS = true
	}
	if env := os.Getenv("FORMICD_NODE_ID"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.nodeId = val
		}
	}
	cfg.metricsAddr = ":9100"
	if env := os.Getenv("FORMICD_METRICS_ADDR"); env != "" {
		cfg.metricsAddr = env
	}
	if env := os.Getenv("FORMICD_METRICS_COLLECTORS"); env != "" {
		cfg.metricsCollectors = env
	}
	if env := os.Getenv("FORMICD_CONCURRENT_REQUESTS_PER_STORE"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.concurrentRequestsPerStore = val
		}
	}
	if env := os.Getenv("FORMICD_DEBUG"); env != "true" {
		cfg.debug = true
	}
	return cfg
}
<commit_msg>Whoops on the debug flag<commit_after>package main

import (
	"log"
	"os"
	"strconv"
)

type config struct {
	path                       string
	port                       int
	oortValueSyndicate         string
	oortGroupSyndicate         string
	insecureSkipVerify         bool
	skipMutualTLS              bool
	nodeId                     int
	metricsAddr                string
	metricsCollectors          string
	concurrentRequestsPerStore int
	debug                      bool
}

func resolveConfig(c *config) *config {
	cfg := &config{}
	if c != nil {
		*cfg = *c
	}
	if env := os.Getenv("FORMICD_PATH"); env != "" {
		cfg.path = env
	}
	if cfg.path == "" {
		cfg.path = "/var/lib/formic"
	}
	if env := os.Getenv("FORMICD_PORT"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.port = val
		}
	}
	if cfg.port == 0 {
		cfg.port = 8445
	}
	if env := os.Getenv("FORMICD_OORT_VALUE_SYNDICATE"); env != "" {
		log.Println("Value: ", env)
		cfg.oortValueSyndicate = env
	}
	// cfg.oortValueSyndicate == "" means default SRV resolution.
	if env := os.Getenv("FORMICD_OORT_GROUP_SYNDICATE"); env != "" {
		log.Println("Group: ", env)
		cfg.oortGroupSyndicate = env
	}
	// cfg.oortGroupSyndicate == "" means default SRV resolution.
	if env := os.Getenv("FORMICD_INSECURE_SKIP_VERIFY"); env == "true" {
		cfg.insecureSkipVerify = true
	}
	if env := os.Getenv("FORMICD_SKIP_MUTUAL_TLS"); env == "true" {
		cfg.skipMutualTLS = true
	}
	if env := os.Getenv("FORMICD_NODE_ID"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.nodeId = val
		}
	}
	cfg.metricsAddr = ":9100"
	if env := os.Getenv("FORMICD_METRICS_ADDR"); env != "" {
		cfg.metricsAddr = env
	}
	if env := os.Getenv("FORMICD_METRICS_COLLECTORS"); env != "" {
		cfg.metricsCollectors = env
	}
	if env := os.Getenv("FORMICD_CONCURRENT_REQUESTS_PER_STORE"); env != "" {
		if val, err := strconv.Atoi(env); err == nil {
			cfg.concurrentRequestsPerStore = val
		}
	}
	if env := os.Getenv("FORMICD_DEBUG"); env == "true" {
		cfg.debug = true
	}
	return cfg
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 fs contains file system related functions and types.
package fs

import (
	"fmt"
	"github.com/jacobsa/comeback/sys"
	"io"
	"io/ioutil"
	"os"
	"path"
	"syscall"
)

const (
	// A permissions bit mask that matches chmod(2)'s notion of permissions.
	permissionBits os.FileMode = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky
)

// FileSystem represents operations performed on a real file system, but is an
// interface for mockability.
type FileSystem interface {
	// Read the contents of the directory named by the supplied path, returning
	// an array of directory entries sorted by name. The entries will contain no
	// scores.
	ReadDir(path string) (entries []*DirectoryEntry, err error)

	// Open the file named by the supplied path for reading.
	OpenForReading(path string) (r io.ReadCloser, err error)
}

// Return a FileSystem that uses the real file system.
func NewFileSystem() (fs FileSystem, err error) {
	fsStruct := &fileSystem{}

	// Create a user registry.
	if fsStruct.userRegistry, err = sys.NewUserRegistry(); err != nil {
		err = fmt.Errorf("Creating user registry: %v", err)
		return
	}

	// Create a group registry.
	if fsStruct.groupRegistry, err = sys.NewGroupRegistry(); err != nil {
		err = fmt.Errorf("Creating group registry: %v", err)
		return
	}

	return fsStruct, nil
}

type fileSystem struct {
	userRegistry sys.UserRegistry
	groupRegistry sys.GroupRegistry
}

func convertFileInfo(fi os.FileInfo) (*DirectoryEntry, error) {
	entry := &DirectoryEntry{
		Permissions: fi.Mode() & permissionBits,
		Name:        fi.Name(),
		MTime:       fi.ModTime(),
	}

	// Convert the type.
	typeBits := fi.Mode() & (os.ModeType | os.ModeCharDevice)
	switch typeBits {
	case 0:
		entry.Type = TypeFile
	case os.ModeDir:
		entry.Type = TypeDirectory
	case os.ModeSymlink:
		entry.Type = TypeSymlink
	case os.ModeDevice:
		entry.Type = TypeBlockDevice
		sys, ok := fi.Sys().(*syscall.Stat_t)
		if !ok {
			return nil, fmt.Errorf("Unexpected sys value: %v", fi.Sys())
		}
		entry.Device = sys.Dev
	case os.ModeDevice | os.ModeCharDevice:
		entry.Type = TypeCharDevice
		sys, ok := fi.Sys().(*syscall.Stat_t)
		if !ok {
			return nil, fmt.Errorf("Unexpected sys value: %v", fi.Sys())
		}
		entry.Device = sys.Dev
	case os.ModeNamedPipe:
		entry.Type = TypeNamedPipe
	default:
		return entry, fmt.Errorf("Unhandled mode: %v %u", fi.Mode(), fi.Mode())
	}

	return entry, nil
}

func (f *fileSystem) ReadDir(dirpath string) (entries []*DirectoryEntry, err error) {
	// Call ioutil.
	fileInfos, err := ioutil.ReadDir(dirpath)
	if err != nil {
		return nil, err
	}

	// Convert each entry.
	entries = []*DirectoryEntry{}
	for _, fileInfo := range fileInfos {
		entry, err := convertFileInfo(fileInfo)
		if err != nil {
			return nil, err
		}

		// Handle symlinks.
		if entry.Type == TypeSymlink {
			linkPath := path.Join(dirpath, entry.Name)
			if entry.Target, err = os.Readlink(linkPath); err != nil {
				return nil, err
			}
		}

		entries = append(entries, entry)
	}

	return entries, nil
}

func (f *fileSystem) OpenForReading(path string) (r io.ReadCloser, err error) {
	return os.Open(path)
}
<commit_msg>Added UIDs and GIDs.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 fs contains file system related functions and types.
package fs

import (
	"fmt"
	"github.com/jacobsa/comeback/sys"
	"io"
	"io/ioutil"
	"os"
	"path"
	"syscall"
)

const (
	// A permissions bit mask that matches chmod(2)'s notion of permissions.
	permissionBits os.FileMode = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky
)

// FileSystem represents operations performed on a real file system, but is an
// interface for mockability.
type FileSystem interface {
	// Read the contents of the directory named by the supplied path, returning
	// an array of directory entries sorted by name. The entries will contain no
	// scores.
	ReadDir(path string) (entries []*DirectoryEntry, err error)

	// Open the file named by the supplied path for reading.
	OpenForReading(path string) (r io.ReadCloser, err error)
}

// Return a FileSystem that uses the real file system.
func NewFileSystem() (fs FileSystem, err error) {
	fsStruct := &fileSystem{}

	// Create a user registry.
	if fsStruct.userRegistry, err = sys.NewUserRegistry(); err != nil {
		err = fmt.Errorf("Creating user registry: %v", err)
		return
	}

	// Create a group registry.
	if fsStruct.groupRegistry, err = sys.NewGroupRegistry(); err != nil {
		err = fmt.Errorf("Creating group registry: %v", err)
		return
	}

	return fsStruct, nil
}

type fileSystem struct {
	userRegistry sys.UserRegistry
	groupRegistry sys.GroupRegistry
}

func convertFileInfo(fi os.FileInfo) (*DirectoryEntry, error) {
	// Grab system-specific info.
	statT, ok := fi.Sys().(*syscall.Stat_t)
	if !ok {
		return nil, fmt.Errorf("Unexpected sys value: %v", fi.Sys())
	}

	// Create the basic entry.
	entry := &DirectoryEntry{
		Permissions: fi.Mode() & permissionBits,
		Name:        fi.Name(),
		MTime:       fi.ModTime(),
		Uid:         sys.UserId(statT.Uid),
		Gid:         sys.GroupId(statT.Gid),
	}

	// Convert the type.
	typeBits := fi.Mode() & (os.ModeType | os.ModeCharDevice)
	switch typeBits {
	case 0:
		entry.Type = TypeFile
	case os.ModeDir:
		entry.Type = TypeDirectory
	case os.ModeSymlink:
		entry.Type = TypeSymlink
	case os.ModeDevice:
		entry.Type = TypeBlockDevice
		entry.Device = statT.Dev
	case os.ModeDevice | os.ModeCharDevice:
		entry.Type = TypeCharDevice
		entry.Device = statT.Dev
	case os.ModeNamedPipe:
		entry.Type = TypeNamedPipe
	default:
		return entry, fmt.Errorf("Unhandled mode: %v %u", fi.Mode(), fi.Mode())
	}

	return entry, nil
}

func (f *fileSystem) ReadDir(dirpath string) (entries []*DirectoryEntry, err error) {
	// Call ioutil.
	fileInfos, err := ioutil.ReadDir(dirpath)
	if err != nil {
		return nil, err
	}

	// Convert each entry.
	entries = []*DirectoryEntry{}
	for _, fileInfo := range fileInfos {
		entry, err := convertFileInfo(fileInfo)
		if err != nil {
			return nil, err
		}

		// Handle symlinks.
		if entry.Type == TypeSymlink {
			linkPath := path.Join(dirpath, entry.Name)
			if entry.Target, err = os.Readlink(linkPath); err != nil {
				return nil, err
			}
		}

		entries = append(entries, entry)
	}

	return entries, nil
}

func (f *fileSystem) OpenForReading(path string) (r io.ReadCloser, err error) {
	return os.Open(path)
}
<|endoftext|>


<commit_before>package commands_test

import (
	"strings"
	"testing"

	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/testutil"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/mocks"
)

func TestGrantPrivilege(t *testing.T) {

	defVM := lib.VirtualMachineName{Group: "default", Account: "default-account"}
	defGroup := lib.GroupName{Group: "default", Account: "default-account"}

	tests := []struct {
		Name      string
		Setup     func(config *mocks.Config, c *mocks.Client)
		ShouldErr bool
		Input     string
	}{
		{
			Name: "MissingArguments",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege",
		}, {
			Name: "BadPrivilege",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege smedly",
		}, {
			Name: "MissingAuth",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege cluster_admin on bucholic to no-one",
		}, {
			Name: "GrantGroupAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetGroup").Return(defGroup)
				group := lib.GroupName{
					Group:   "test-group",
					Account: "test-account",
				}

				c.When("GetGroup", group).Return(brain.Group{
					ID: 303,
				}).Times(1)
				c.When("GrantPrivilege", brain.Privilege{
					Username: "test-user",
					Level:    brain.GroupAdminPrivilege,
					GroupID:  303,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege group_admin test-group.test-account test-user",
		},
		{
			Name: "GrantGroupAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetGroup").Return(defGroup)
				group := lib.GroupName{
					Group:   "test-group",
					Account: "test-account",
				}

				c.When("GetGroup", group).Return(brain.Group{
					ID: 303,
				}).Times(1)
				c.When("GrantPrivilege", brain.Privilege{
					Username: "test-user",
					Level:    brain.GroupAdminPrivilege,
					GroupID:  303,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege group_admin on test-group.test-account to test-user",
		},
		{
			Name: "GrantVMAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetVirtualMachine").Return(defVM)
				vm := lib.VirtualMachineName{
					VirtualMachine: "test-vm",
					Group:          "test-group",
					Account:        "test-account",
				}

				c.When("GetVirtualMachine", vm).Return(brain.VirtualMachine{ID: 333}, nil).Times(1)

				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					Level:            brain.VMAdminPrivilege,
					VirtualMachineID: 333,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege vm_admin on test-vm.test-group.test-account to test-user",
		},
		{
			Name: "AccountVMAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console
				config.When("GetIgnoreErr", "account").Return("default-account")

				c.When("GetAccount", "test-account").Return(lib.Account{
					BrainID: 32310,
				})

				c.When("GrantPrivilege", brain.Privilege{
					Username:  "test-user",
					Level:     brain.AccountAdminPrivilege,
					AccountID: 32310,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege account_admin on test-account to test-user",
		}, {
			Name: "ApiKey",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				config.When("GetIgnoreErr", "account").Return("default-account")
				c.When("GetAccount", "account").Return(lib.Account{BrainID: 32310})
				c.When("GrantPrivilege", brain.Privilege{
					Username:  "user",
					AccountID: 32310,
					Level:     brain.AccountAdminPrivilege,
					APIKeyID:  4,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege --api-key-id 4 account_admin on account to user",
		},
	}
	for i, test := range tests {
		t.Run(test.Name, func(t *testing.T) {
			config, c, app := testutil.BaseTestAuthSetup(t, false, commands.Commands)
			test.Setup(config, c)

			err := app.Run(strings.Split(test.Input, " "))
			if test.ShouldErr && err == nil {
				t.Errorf("TestGrantPrivilege %d should err and didn't", i)
			} else if !test.ShouldErr && err != nil {
				t.Errorf("TestGrantPrivilege %d shouldn't err, but: %s", i, err)
			}
			if ok, err := c.Verify(); !ok {
				t.Fatal(err)
			}
		})
	}
}
<commit_msg>shore up grant tests<commit_after>package commands_test

import (
	"strings"
	"testing"

	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/testutil"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/mocks"
)

func TestGrantPrivilege(t *testing.T) {

	defVM := lib.VirtualMachineName{Group: "default", Account: "default-account"}
	defGroup := lib.GroupName{Group: "default", Account: "default-account"}

	tests := []struct {
		Name      string
		Setup     func(config *mocks.Config, c *mocks.Client)
		ShouldErr bool
		Input     string
	}{
		{
			Name: "MissingArguments",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege",
		}, {
			Name: "BadPrivilege",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege smedly",
		}, {
			Name: "BadPrivilegeYubikeyPlusApiKey",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege --api-key jeff --yubikey-required cluster_admin",
		}, {
			Name: "MissingAuth",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// reset to get rid of the default AuthWithToken expectation
				c.Reset()
			},
			ShouldErr: true,
			Input:     "bytemark grant privilege cluster_admin on bucholic to no-one",
		}, {
			Name: "GrantGroupAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetGroup").Return(defGroup)
				group := lib.GroupName{
					Group:   "test-group",
					Account: "test-account",
				}

				c.When("GetGroup", group).Return(brain.Group{
					ID: 303,
				}).Times(1)
				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					PasswordRequired: true,
					Level:            brain.GroupAdminPrivilege,
					GroupID:          303,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege group_admin test-group.test-account test-user",
		},
		{
			Name: "GrantGroupAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetGroup").Return(defGroup)
				group := lib.GroupName{
					Group:   "test-group",
					Account: "test-account",
				}

				c.When("GetGroup", group).Return(brain.Group{
					ID: 303,
				}).Times(1)
				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					PasswordRequired: true,
					Level:            brain.GroupAdminPrivilege,
					GroupID:          303,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege group_admin on test-group.test-account to test-user",
		},
		{
			Name: "GrantVMAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console

				config.When("GetVirtualMachine").Return(defVM)
				vm := lib.VirtualMachineName{
					VirtualMachine: "test-vm",
					Group:          "test-group",
					Account:        "test-account",
				}

				c.When("GetVirtualMachine", vm).Return(brain.VirtualMachine{ID: 333}, nil).Times(1)

				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					PasswordRequired: true,
					Level:            brain.VMAdminPrivilege,
					VirtualMachineID: 333,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege vm_admin on test-vm.test-group.test-account to test-user",
		},
		{
			Name: "AccountVMAdmin",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console
				config.When("GetIgnoreErr", "account").Return("default-account")

				c.When("GetAccount", "test-account").Return(lib.Account{
					BrainID: 32310,
				})

				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					PasswordRequired: true,
					Level:            brain.AccountAdminPrivilege,
					AccountID:        32310,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege account_admin on test-account to test-user",
		}, {
			Name: "AccountVMAdmin with Yubikey",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				// specific to vm_admin/vm_console
				config.When("GetIgnoreErr", "account").Return("default-account")

				c.When("GetAccount", "test-account").Return(lib.Account{
					BrainID: 32310,
				})

				c.When("GrantPrivilege", brain.Privilege{
					Username:         "test-user",
					PasswordRequired: true,
					YubikeyRequired:  true,
					Level:            brain.AccountAdminPrivilege,
					AccountID:        32310,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege --yubikey-required account_admin on test-account to test-user",
		}, {
			Name: "ApiKey",
			Setup: func(config *mocks.Config, c *mocks.Client) {
				config.When("GetIgnoreErr", "account").Return("default-account")
				c.When("GetAccount", "account").Return(lib.Account{BrainID: 32310})
				c.When("GrantPrivilege", brain.Privilege{
					Username:         "user",
					AccountID:        32310,
					PasswordRequired: false,
					Level:            brain.AccountAdminPrivilege,
					APIKeyID:         4,
				}).Return(nil).Times(1)
			},
			ShouldErr: false,
			Input:     "bytemark grant privilege --api-key-id 4 account_admin on account to user",
		},
	}
	for i, test := range tests {
		t.Run(test.Name, func(t *testing.T) {
			config, c, app := testutil.BaseTestAuthSetup(t, false, commands.Commands)
			test.Setup(config, c)

			err := app.Run(strings.Split(test.Input, " "))
			if test.ShouldErr && err == nil {
				t.Errorf("TestGrantPrivilege %d should err and didn't", i)
			} else if !test.ShouldErr && err != nil {
				t.Errorf("TestGrantPrivilege %d shouldn't err, but: %s", i, err)
			}
			if ok, err := c.Verify(); !ok {
				t.Fatal(err)
			}
		})
	}
}
<|endoftext|>


<commit_before>package runtime

import (
	"fmt"
	"os"
	"strings"

	"github.com/go-clang/gen"
)

// PrepareFunctionName prepares C function naming to Go function name.
func PrepareFunctionName(g *gen.Generation, f *gen.Function) string {
	fname := f.Name
	fname = strings.TrimPrefix(fname, "clang_")

	// Trim some whitelisted prefixes by their function name
	if fn := strings.TrimPrefix(fname, "indexLoc_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "index_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "Location_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "Range_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "remap_"); len(fn) != len(fname) {
		fname = fn
	}

	// Trim some whitelisted prefixes by their types
	if len(f.Parameters) > 0 && g.IsEnumOrStruct(f.Parameters[0].Type.GoName) {
		switch f.Parameters[0].Type.GoName {
		case "CodeCompleteResults":
			fname = strings.TrimPrefix(fname, "codeComplete")
		case "CompletionString":
			if f.CName == "clang_getNumCompletionChunks" {
				fname = "NumChunks"
			} else {
				fname = strings.TrimPrefix(fname, "getCompletion")
			}
		case "SourceRange":
			fname = strings.TrimPrefix(fname, "getRange")
		}
	}

	return fname
}

// PrepareFunction prepares C function to Go function.
func PrepareFunction(f *gen.Function) {
	for i := range f.Parameters {
		p := &f.Parameters[i]

		if f.CName == "clang_getRemappingsFromFileList" {
			switch p.CName {
			case "filePaths":
				p.Type.IsSlice = true
			case "numFiles":
				p.Type.LengthOfSlice = "filePaths"
			}

			continue
		}

		// Whiteflag types that are return arguments
		if p.Type.PointerLevel == 1 && (p.Type.GoName == "File" || p.Type.GoName == "FileUniqueID" || p.Type.GoName == "IdxClientFile" || p.Type.GoName == "cxstring" || p.Type.GoName == gen.GoInt32 || p.Type.GoName == gen.GoUInt32 || p.Type.GoName == "CompilationDatabase_Error" || p.Type.GoName == "PlatformAvailability" || p.Type.GoName == "SourceRange" || p.Type.GoName == "LoadDiag_Error") {
			p.Type.IsReturnArgument = true
		}
		if p.Type.PointerLevel == 2 && (p.Type.GoName == "Token" || p.Type.GoName == "Cursor") {
			p.Type.IsReturnArgument = true
		}

		if f.CName == "clang_disposeOverriddenCursors" && p.CName == "overridden" {
			p.Type.IsSlice = true
		}

		// If this is an array length parameter we need to find its partner
		paCName := gen.ArrayNameFromLength(p.CName)

		if paCName != "" {
			for j := range f.Parameters {
				pa := &f.Parameters[j]

				if strings.ToLower(pa.CName) == strings.ToLower(paCName) {
					if pa.Type.GoName == "struct CXUnsavedFile" || pa.Type.GoName == "UnsavedFile" {
						pa.Type.GoName = "UnsavedFile"
						pa.Type.CGoName = "struct_CXUnsavedFile"
					} else if pa.Type.CGoName == gen.CSChar && pa.Type.PointerLevel == 2 {
					} else if pa.Type.GoName == "CompletionResult" {
					} else if pa.Type.GoName == "Token" {
					} else if pa.Type.GoName == "Cursor" {
					} else {
						break
					}

					p.Type.LengthOfSlice = pa.Name
					pa.Type.IsSlice = true

					if pa.Type.IsReturnArgument && p.Type.PointerLevel > 0 {
						p.Type.IsReturnArgument = true
					}

					break
				}
			}
		}
	}

	for i := range f.Parameters {
		p := &f.Parameters[i]

		if p.Type.CGoName == gen.CSChar && p.Type.PointerLevel == 2 && !p.Type.IsSlice {
			p.Type.IsReturnArgument = true
		}
	}
}

// FilterFunction reports whether the f function filtered to a particular condition.
func FilterFunction(f *gen.Function) bool {
	switch f.CName {
	case "clang_CompileCommand_getMappedSourceContent", "clang_CompileCommand_getMappedSourcePath", "clang_CompileCommand_getNumMappedSources":
		// some functions are not compiled in the library see https://lists.launchpad.net/desktop-packages/msg75835.html for a never resolved bug report
		fmt.Fprintf(os.Stderr, "Ignore function %q because it is not compiled within libClang\n", f.CName)

		return false

	case "clang_executeOnThread", "clang_getInclusions":
		// some functions can not be handled automatically by us
		fmt.Fprintf(os.Stderr, "Ignore function %q because it cannot be handled automatically\n", f.CName)

		return false

	case "clang_annotateTokens", "clang_getCursorPlatformAvailability", "clang_visitChildren":
		// some functions are simply manually implemented
		fmt.Fprintf(os.Stderr, "Ignore function %q because it is manually implemented\n", f.CName)

		return false
	}

	// TODO(go-clang): if this function is from CXString.h we ignore it https://github.com/go-clang/gen/issues/25
	for i := range f.IncludeFiles {
		if strings.HasSuffix(i, "CXString.h") {
			return false
		}
	}

	return true
}

// FilterFunctionParameter reports whether the p function parameter filtered to a particular condition.
func FilterFunctionParameter(p gen.FunctionParameter) bool {
	// these pointers are ok
	if p.Type.PointerLevel == 1 {
		if p.Type.CGoName == gen.CSChar {
			return false
		}

		switch p.Type.GoName {
		case "UnsavedFile",
			"CodeCompleteResults",
			"CursorKind",
			"IdxContainerInfo",
			"IdxDeclInfo",
			"IndexerCallbacks",
			"TranslationUnit",
			"IdxEntityInfo",
			"IdxAttrInfo":

			return false
		}
	}

	return true
}

// FixFunctionName fixes the function name under certain conditions.
func FixFunctionName(f *gen.Function) string {
	// needs to be renamed manually since clang_getTranslationUnitCursor will conflict with clang_getCursor
	if f.CName == "clang_getTranslationUnitCursor" {
		return "TranslationUnitCursor"
	}

	return ""
}

// PrepareStructFields prepares struct fields names.
func PrepareStructFields(s *gen.Struct) {
	for _, m := range s.Fields {
		if (strings.HasPrefix(m.CName, "has") || strings.HasPrefix(m.CName, "is")) && m.Type.GoName == gen.GoInt32 {
			m.Type.GoName = gen.GoBool
		}

		// if this is an array length parameter we need to find its partner
		maCName := gen.ArrayNameFromLength(m.CName)

		if maCName != "" {
			for _, ma := range s.Fields {
				if strings.ToLower(ma.CName) == strings.ToLower(maCName) {
					m.Type.LengthOfSlice = ma.CName
					ma.Type.IsSlice = true
					// TODO(go-clang): wrong usage but needed for the getter generation...
					// maybe refactor this LengthOfSlice alltogether?
					// https://github.com/go-clang/gen/issues/49
					ma.Type.LengthOfSlice = m.CName

					break
				}
			}
		}
	}

	prepareStructFieldsArrayStruct(s)
}

// prepareStructFieldsArrayStruct checks if the struct has two member variables, one is an array and the other a plain
// int/uint with size/length/count/len is its name because then this should be an array struct, and we connect them to handle a slice.
func prepareStructFieldsArrayStruct(s *gen.Struct) {
	if len(s.Fields) != 2 {
		return
	}

	if !arrayLengthCombination(&s.Fields[0].Type, &s.Fields[1].Type) && !arrayLengthCombination(&s.Fields[1].Type, &s.Fields[0].Type) {
		return
	}

	// if one of the members is already marked as array/slice another heuristic has already covered both members.
	switch {
	case s.Fields[0].Type.IsArray,
		s.Fields[1].Type.IsArray,
		s.Fields[0].Type.IsSlice,
		s.Fields[1].Type.IsSlice:

		return
	}

	var a *gen.StructField
	var c *gen.StructField

	if s.Fields[0].Type.PointerLevel == 1 {
		a = s.Fields[0]
		c = s.Fields[1]
	} else {
		c = s.Fields[0]
		a = s.Fields[1]
	}

	lengthName := strings.ToLower(c.CName)
	if lengthName != "count" &&
		lengthName != "len" && lengthName != "length" && lengthName != "size" {

		return
	}

	c.Type.LengthOfSlice = a.CName
	a.Type.IsSlice = true
	// TODO(go-clang): wrong usage but needed for the getter generation...
	// maybe refactor this LengthOfSlice alltogether?
	// https://github.com/go-clang/gen/issues/49
	a.Type.LengthOfSlice = c.CName
}

// arrayLengthCombination reports whether the x and y to correct combination.
func arrayLengthCombination(x *gen.Type, y *gen.Type) bool {
	return x.PointerLevel == 1 && y.PointerLevel == 0 &&
		!gen.IsInteger(x) && gen.IsInteger(y)
}

// FilterStructFieldGetter reports whether the m struct member filtered to a particular condition.
func FilterStructFieldGetter(m *gen.StructField) bool {
	// we do not want getters to *int_data members
	return !strings.HasSuffix(m.CName, "int_data")
}
<commit_msg>cmd/go-clang-gen/runtime: fix documentation and variable name<commit_after>package runtime

import (
	"fmt"
	"os"
	"strings"

	"github.com/go-clang/gen"
)

// PrepareFunctionName prepares C function naming to Go function name.
func PrepareFunctionName(g *gen.Generation, f *gen.Function) string {
	fname := f.Name
	fname = strings.TrimPrefix(fname, "clang_")

	// Trim some whitelisted prefixes by their function name
	if fn := strings.TrimPrefix(fname, "indexLoc_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "index_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "Location_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "Range_"); len(fn) != len(fname) {
		fname = fn
	} else if fn := strings.TrimPrefix(fname, "remap_"); len(fn) != len(fname) {
		fname = fn
	}

	// Trim some whitelisted prefixes by their types
	if len(f.Parameters) > 0 && g.IsEnumOrStruct(f.Parameters[0].Type.GoName) {
		switch f.Parameters[0].Type.GoName {
		case "CodeCompleteResults":
			fname = strings.TrimPrefix(fname, "codeComplete")
		case "CompletionString":
			if f.CName == "clang_getNumCompletionChunks" {
				fname = "NumChunks"
			} else {
				fname = strings.TrimPrefix(fname, "getCompletion")
			}
		case "SourceRange":
			fname = strings.TrimPrefix(fname, "getRange")
		}
	}

	return fname
}

// PrepareFunction prepares C function to Go function.
func PrepareFunction(f *gen.Function) {
	for i := range f.Parameters {
		p := &f.Parameters[i]

		if f.CName == "clang_getRemappingsFromFileList" {
			switch p.CName {
			case "filePaths":
				p.Type.IsSlice = true
			case "numFiles":
				p.Type.LengthOfSlice = "filePaths"
			}

			continue
		}

		// Whiteflag types that are return arguments
		if p.Type.PointerLevel == 1 && (p.Type.GoName == "File" || p.Type.GoName == "FileUniqueID" || p.Type.GoName == "IdxClientFile" || p.Type.GoName == "cxstring" || p.Type.GoName == gen.GoInt32 || p.Type.GoName == gen.GoUInt32 || p.Type.GoName == "CompilationDatabase_Error" || p.Type.GoName == "PlatformAvailability" || p.Type.GoName == "SourceRange" || p.Type.GoName == "LoadDiag_Error") {
			p.Type.IsReturnArgument = true
		}
		if p.Type.PointerLevel == 2 && (p.Type.GoName == "Token" || p.Type.GoName == "Cursor") {
			p.Type.IsReturnArgument = true
		}

		if f.CName == "clang_disposeOverriddenCursors" && p.CName == "overridden" {
			p.Type.IsSlice = true
		}

		// If this is an array length parameter we need to find its partner
		paCName := gen.ArrayNameFromLength(p.CName)

		if paCName != "" {
			for j := range f.Parameters {
				pa := &f.Parameters[j]

				if strings.ToLower(pa.CName) == strings.ToLower(paCName) {
					if pa.Type.GoName == "struct CXUnsavedFile" || pa.Type.GoName == "UnsavedFile" {
						pa.Type.GoName = "UnsavedFile"
						pa.Type.CGoName = "struct_CXUnsavedFile"
					} else if pa.Type.CGoName == gen.CSChar && pa.Type.PointerLevel == 2 {
					} else if pa.Type.GoName == "CompletionResult" {
					} else if pa.Type.GoName == "Token" {
					} else if pa.Type.GoName == "Cursor" {
					} else {
						break
					}

					p.Type.LengthOfSlice = pa.Name
					pa.Type.IsSlice = true

					if pa.Type.IsReturnArgument && p.Type.PointerLevel > 0 {
						p.Type.IsReturnArgument = true
					}

					break
				}
			}
		}
	}

	for i := range f.Parameters {
		p := &f.Parameters[i]

		if p.Type.CGoName == gen.CSChar && p.Type.PointerLevel == 2 && !p.Type.IsSlice {
			p.Type.IsReturnArgument = true
		}
	}
}

// FilterFunction reports whether the f function filtered to a particular condition.
func FilterFunction(f *gen.Function) bool {
	switch f.CName {
	case "clang_CompileCommand_getMappedSourceContent", "clang_CompileCommand_getMappedSourcePath", "clang_CompileCommand_getNumMappedSources":
		// some functions are not compiled in the library see https://lists.launchpad.net/desktop-packages/msg75835.html for a never resolved bug report
		fmt.Fprintf(os.Stderr, "Ignore function %q because it is not compiled within libClang\n", f.CName)

		return false

	case "clang_executeOnThread", "clang_getInclusions":
		// some functions can not be handled automatically by us
		fmt.Fprintf(os.Stderr, "Ignore function %q because it cannot be handled automatically\n", f.CName)

		return false

	case "clang_annotateTokens", "clang_getCursorPlatformAvailability", "clang_visitChildren":
		// some functions are simply manually implemented
		fmt.Fprintf(os.Stderr, "Ignore function %q because it is manually implemented\n", f.CName)

		return false
	}

	// TODO(go-clang): if this function is from CXString.h we ignore it https://github.com/go-clang/gen/issues/25
	for i := range f.IncludeFiles {
		if strings.HasSuffix(i, "CXString.h") {
			return false
		}
	}

	return true
}

// FilterFunctionParameter reports whether the p function parameter filtered to a particular condition.
func FilterFunctionParameter(p gen.FunctionParameter) bool {
	// these pointers are ok
	if p.Type.PointerLevel == 1 {
		if p.Type.CGoName == gen.CSChar {
			return false
		}

		switch p.Type.GoName {
		case "UnsavedFile",
			"CodeCompleteResults",
			"CursorKind",
			"IdxContainerInfo",
			"IdxDeclInfo",
			"IndexerCallbacks",
			"TranslationUnit",
			"IdxEntityInfo",
			"IdxAttrInfo":

			return false
		}
	}

	return true
}

// FixFunctionName fixes the function name under certain conditions.
func FixFunctionName(f *gen.Function) string {
	// needs to be renamed manually since clang_getTranslationUnitCursor will conflict with clang_getCursor
	if f.CName == "clang_getTranslationUnitCursor" {
		return "TranslationUnitCursor"
	}

	return ""
}

// PrepareStructFields prepares struct fields names.
func PrepareStructFields(s *gen.Struct) {
	for _, f := range s.Fields {
		if (strings.HasPrefix(f.CName, "has") || strings.HasPrefix(f.CName, "is")) && f.Type.GoName == gen.GoInt32 {
			f.Type.GoName = gen.GoBool
		}

		// if this is an array length parameter we need to find its partner
		faCName := gen.ArrayNameFromLength(f.CName)

		if faCName != "" {
			for _, ma := range s.Fields {
				if strings.ToLower(ma.CName) == strings.ToLower(faCName) {
					f.Type.LengthOfSlice = ma.CName
					ma.Type.IsSlice = true
					// TODO(go-clang): wrong usage but needed for the getter generation...
					// maybe refactor this LengthOfSlice alltogether?
					// https://github.com/go-clang/gen/issues/49
					ma.Type.LengthOfSlice = f.CName

					break
				}
			}
		}
	}

	prepareStructFieldsArrayStruct(s)
}

// prepareStructFieldsArrayStruct checks if the struct has two field variables, one is an array and the other a plain
// int/uint with size/length/count/len is its name because then this should be an array struct, and we connect them to handle a slice.
func prepareStructFieldsArrayStruct(s *gen.Struct) {
	if len(s.Fields) != 2 {
		return
	}

	if !arrayLengthCombination(&s.Fields[0].Type, &s.Fields[1].Type) && !arrayLengthCombination(&s.Fields[1].Type, &s.Fields[0].Type) {
		return
	}

	// if one of the fields is already marked as array/slice another heuristic has already covered both fields
	switch {
	case s.Fields[0].Type.IsArray,
		s.Fields[1].Type.IsArray,
		s.Fields[0].Type.IsSlice,
		s.Fields[1].Type.IsSlice:

		return
	}

	var a *gen.StructField
	var c *gen.StructField

	if s.Fields[0].Type.PointerLevel == 1 {
		a = s.Fields[0]
		c = s.Fields[1]
	} else {
		c = s.Fields[0]
		a = s.Fields[1]
	}

	lengthName := strings.ToLower(c.CName)
	if lengthName != "count" &&
		lengthName != "len" && lengthName != "length" && lengthName != "size" {

		return
	}

	c.Type.LengthOfSlice = a.CName
	a.Type.IsSlice = true
	// TODO(go-clang): wrong usage but needed for the getter generation...
	// maybe refactor this LengthOfSlice alltogether?
	// https://github.com/go-clang/gen/issues/49
	a.Type.LengthOfSlice = c.CName
}

// arrayLengthCombination reports whether the x and y to correct combination.
func arrayLengthCombination(x *gen.Type, y *gen.Type) bool {
	return x.PointerLevel == 1 && y.PointerLevel == 0 &&
		!gen.IsInteger(x) && gen.IsInteger(y)
}

// FilterStructFieldGetter reports whether the m struct field filtered to a particular condition.
func FilterStructFieldGetter(m *gen.StructField) bool {
	// we do not want getters to *int_data fields
	return !strings.HasSuffix(m.CName, "int_data")
}
<|endoftext|>


<commit_before>package store

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestNSS_certdbDiscovery(t *testing.T) {
	// create a dir, add a 'cert8.db' file and verify we can discover it
	tmp, err := ioutil.TempDir("", "nss-discovery")
	if err != nil {
		t.Error(err)
	}
	defer os.Remove(tmp)

	// Very we don't find anything in an empty dir
	if len(collectNssSuggestions([]string{tmp})) > 0 {
		t.Errorf("shouldn't have found cert8.db files in %s", tmp)
	}
	if containsCert8db(tmp) {
		t.Errorf("%s shouldn't contain a cert8.db", tmp)
	}

	// Create a blank cert8.db file
	where := filepath.Join(tmp, "cert8.db")
	err = ioutil.WriteFile(where, []byte("data"), 0644)
	if err != nil {
		t.Error(err)
	}

	// Now we should find the cert8.db path
	if !containsCert8db(tmp) {
		t.Errorf("should have found cert8.db in %s", tmp)
	}
	sugs := collectNssSuggestions([]string{tmp})
	if len(sugs) != 1 {
		t.Errorf("should have found cert8.db in %s", sugs)
	}
}

func TestNSS_trustedForSSL(t *testing.T) {
	trusted := map[string]bool{
		// trusted attrs
		"c,c,c": true,
		",,":    true,
		"CT,,":  true,
		"u,,":   true,
		"wC,,":  true,
		// not trusted
		"p,p,p": false,
		"p,,":   false,
	}
	for attrs, answer := range trusted {
		item := cert8Item{trustAttrs: attrs}
		if res := item.trustedForSSL(); res != answer {
			t.Errorf("attrs (%s), trustedForSSL()=%v should be %v", attrs, res, answer)
		}
	}
}
<commit_msg>store/nss: clarify test names<commit_after>package store

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestStoreNSS_certdbDiscovery(t *testing.T) {
	// create a dir, add a 'cert8.db' file and verify we can discover it
	tmp, err := ioutil.TempDir("", "nss-discovery")
	if err != nil {
		t.Error(err)
	}
	defer os.Remove(tmp)

	// Very we don't find anything in an empty dir
	if len(collectNssSuggestions([]string{tmp})) > 0 {
		t.Errorf("shouldn't have found cert8.db files in %s", tmp)
	}
	if containsCert8db(tmp) {
		t.Errorf("%s shouldn't contain a cert8.db", tmp)
	}

	// Create a blank cert8.db file
	where := filepath.Join(tmp, "cert8.db")
	err = ioutil.WriteFile(where, []byte("data"), 0644)
	if err != nil {
		t.Error(err)
	}

	// Now we should find the cert8.db path
	if !containsCert8db(tmp) {
		t.Errorf("should have found cert8.db in %s", tmp)
	}
	sugs := collectNssSuggestions([]string{tmp})
	if len(sugs) != 1 {
		t.Errorf("should have found cert8.db in %s", sugs)
	}
}

func TestStoreNSS_trustedForSSL(t *testing.T) {
	trusted := map[string]bool{
		// trusted attrs
		"c,c,c": true,
		",,":    true,
		"CT,,":  true,
		"u,,":   true,
		"wC,,":  true,
		// not trusted
		"p,p,p": false,
		"p,,":   false,
	}
	for attrs, answer := range trusted {
		item := cert8Item{trustAttrs: attrs}
		if res := item.trustedForSSL(); res != answer {
			t.Errorf("attrs (%s), trustedForSSL()=%v should be %v", attrs, res, answer)
		}
	}
}
<|endoftext|>


<commit_before>package streaming

import (
	"encoding/json"
	"github.com/boredomist/mixport/mixpanel"
	"io"
	"log"
	"os"
)

// JSONStreamer writes records to a local file in JSON format line by line,
// simply serializing the JSON directly.
//
// Format is simply: `{"key": "value", ...}`, where all `value`s are scalar
// (i.e. not maps or vectors).
func JSONStreamer(name string, records <-chan mixpanel.EventData) {
	fp, err := os.Create(name)
	if err != nil {
		log.Fatalf("Couldn't create file: %s", err)
	}

	defer func() {
		if err := fp.Close(); err != nil {
			panic(err)
		}
	}()

	encoder := json.NewEncoder(io.Writer(fp))

	for record := range records {
		// FIXME: Does this write newlines?
		encoder.Encode(record)
	}
}
<commit_msg>Update documentation<commit_after>package streaming

import (
	"encoding/json"
	"github.com/boredomist/mixport/mixpanel"
	"io"
	"log"
	"os"
)

// JSONStreamer writes records to a local file in JSON format line by line,
// simply serializing the JSON directly.
//
// Format is simply: `{"key": "value", ...}`. `value` should never be a map,
// but may be a vector. Most values are scalar.
func JSONStreamer(name string, records <-chan mixpanel.EventData) {
	fp, err := os.Create(name)
	if err != nil {
		log.Fatalf("Couldn't create file: %s", err)
	}

	defer func() {
		if err := fp.Close(); err != nil {
			panic(err)
		}
	}()

	encoder := json.NewEncoder(io.Writer(fp))

	for record := range records {
		encoder.Encode(record)
	}
}
<|endoftext|>


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

package scanner

import (
	"testing"

	"github.com/caixw/apidoc/core"
	"github.com/issue9/assert"
)

var _ core.ScanFunc = CStyle

var code1 = `
int x = 5;
/* line1
line2
line3*/`

var comment1 = []byte(` line1
line2
line3`)

var code2 = `
int x = 5;
// line1
// line2
// line3
`

var comment2 = []byte(` line1
 line2
 line3
`)

var code3 = `
  int x=5
  // line1
  // line2
  // line3
`
var comment3 = []byte(` line1
 line2
 line3
`)

func TestCStyle(t *testing.T) {
	a := assert.New(t)

	fn := func(code string, comment []byte) {
		block, pos := CStyle([]byte(code))
		a.Equal(block, comment).Equal(pos, len(code))
	}

	fn(code1, comment1)
	fn(code2, comment2)
	fn(code3, comment3)
}
<commit_msg>删除重复的测试数据<commit_after>// Copyright 2015 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package scanner

import (
	"testing"

	"github.com/caixw/apidoc/core"
	"github.com/issue9/assert"
)

var _ core.ScanFunc = CStyle

var (
	code1 = `
int x = 5;
/* line1
line2
line3*/`

	comment1 = []byte(` line1
line2
line3`)

	code2 = `
int x = 5;
// line1
// line2
// line3
`

	comment2 = []byte(` line1
 line2
 line3
`)
)

func TestCStyle(t *testing.T) {
	a := assert.New(t)

	fn := func(code string, comment []byte) {
		block, pos := CStyle([]byte(code))
		a.Equal(block, comment).Equal(pos, len(code))
	}

	fn(code1, comment1)
	fn(code2, comment2)
}
<|endoftext|>