<commit_before>package main

import (
    "compress/bzip2"
    "compress/flate"
    "compress/gzip"
    "compress/lzw"
    "compress/zlib"
    "flag"
    "fmt"
    "io"
    "log"
    "os"
)

var (
    compress   = flag.Bool("compress", false, "Perform compression")
    decompress = flag.Bool("decompress", false, "Perform decompression")
    algorithm  = flag.String("algorithm", "", "The algorithm to use (one of bzip2, flate, gzip, lzw, zlib)")
)

func filename() string {
    return fmt.Sprintf("everything.go.%s", *algorithm)
}

func openOutputFile() *os.File {
    file, err := os.OpenFile(filename(), os.O_WRONLY|os.O_CREATE, 0644)
    if err != nil {
        log.Fatalf("failed opening output file: %s", err)
    }
    return file
}

func openInputFile() *os.File {
    file, err := os.Open(filename())
    if err != nil {
        log.Fatalf("failed opening input file: %s", err)
    }
    return file
}

func openThisFile() *os.File {
    file, err := os.Open("everything.go")
    if err != nil {
        log.Fatalf("failed opening input file: %s", err)
    }
    return file
}

func getCompressor(out io.Writer) io.WriteCloser {
    switch *algorithm {
    case "bzip2":
        log.Fatalf("no compressor for bzip2. Try `bzip2 -c everything.go > everything.go.bzip2`")
    case "flate":
        compressor, err := flate.NewWriter(out, flate.BestCompression)
        if err != nil {
            log.Fatalf("failed making flate compressor: %s", err)
        }
        return compressor
    case "gzip":
        return gzip.NewWriter(out)
    case "lzw":
        // More specific uses of Order and litWidth are in the package docs
        return lzw.NewWriter(out, lzw.MSB, 8)
    case "zlib":
        return zlib.NewWriter(out)
    default:
        log.Fatalf("choose one of bzip2, flate, gzip, lzw, zlib with -algorithm")
    }
    panic("not reached")
}

func getDecompressor(in io.Reader) io.Reader {
    switch *algorithm {
    case "bzip2":
        return bzip2.NewReader(in)
    case "flate":
        return flate.NewReader(in)
    case "gzip":
        decompressor, err := gzip.NewReader(in)
        if err != nil {
            log.Fatalf("failed making gzip decompressor")
        }
        return decompressor
    case "lzw":
        return lzw.NewReader(in, lzw.MSB, 8)
    case "zlib":
        decompressor, err := zlib.NewReader(in)
        if err != nil {
            log.Fatalf("failed making zlib decompressor")
        }
        return decompressor
    }
    panic("not reached")
}

func compression() {
    output := openOutputFile()
    defer output.Close()
    compressor := getCompressor(output)
    defer compressor.Close()
    input := openThisFile()
    defer input.Close()
    io.Copy(compressor, input)
}

func decompression() {
    input := openInputFile()
    defer input.Close()
    decompressor := getDecompressor(input)
    if c, ok := decompressor.(io.Closer); ok {
        defer c.Close()
    }
    io.Copy(os.Stdout, decompressor)
}

func main() {
    flag.Parse()
    switch {
    case *compress:
        compression()
    case *decompress:
        decompression()
    default:
        log.Println("must specify one of -compress or -decompress")
    }
}
<commit_msg>Specify the input<commit_after>package main

import (
    "compress/bzip2"
    "compress/flate"
    "compress/gzip"
    "compress/lzw"
    "compress/zlib"
    "flag"
    "fmt"
    "io"
    "log"
    "os"
)

var (
    compress   = flag.Bool("compress", false, "Perform compression")
    decompress = flag.Bool("decompress", false, "Perform decompression")
    algorithm  = flag.String("algorithm", "", "The algorithm to use (one of bzip2, flate, gzip, lzw, zlib)")
    input      = flag.String("input", "", "The file to compress or decompress")
)

func filename() string {
    return fmt.Sprintf("%s.%s", *input, *algorithm)
}

func openOutputFile() *os.File {
    file, err := os.OpenFile(filename(), os.O_WRONLY|os.O_CREATE, 0644)
    if err != nil {
        log.Fatalf("failed opening output file: %s", err)
    }
    return file
}

func openInputFile() *os.File {
    file, err := os.Open(*input)
    if err != nil {
        log.Fatalf("failed opening input file: %s", err)
    }
    return file
}

func getCompressor(out io.Writer) io.WriteCloser {
    switch *algorithm {
    case "bzip2":
        log.Fatalf("no compressor for bzip2. Try `bzip2 -c everything.go > everything.go.bzip2`")
    case "flate":
        compressor, err := flate.NewWriter(out, flate.BestCompression)
        if err != nil {
            log.Fatalf("failed making flate compressor: %s", err)
        }
        return compressor
    case "gzip":
        return gzip.NewWriter(out)
    case "lzw":
        // More specific uses of Order and litWidth are in the package docs
        return lzw.NewWriter(out, lzw.MSB, 8)
    case "zlib":
        return zlib.NewWriter(out)
    default:
        log.Fatalf("choose one of bzip2, flate, gzip, lzw, zlib with -algorithm")
    }
    panic("not reached")
}

func getDecompressor(in io.Reader) io.Reader {
    switch *algorithm {
    case "bzip2":
        return bzip2.NewReader(in)
    case "flate":
        return flate.NewReader(in)
    case "gzip":
        decompressor, err := gzip.NewReader(in)
        if err != nil {
            log.Fatalf("failed making gzip decompressor")
        }
        return decompressor
    case "lzw":
        return lzw.NewReader(in, lzw.MSB, 8)
    case "zlib":
        decompressor, err := zlib.NewReader(in)
        if err != nil {
            log.Fatalf("failed making zlib decompressor")
        }
        return decompressor
    }
    panic("not reached")
}

func compression() {
    output := openOutputFile()
    defer output.Close()
    compressor := getCompressor(output)
    defer compressor.Close()
    input := openInputFile()
    defer input.Close()
    io.Copy(compressor, input)
}

func decompression() {
    input := openInputFile()
    defer input.Close()
    decompressor := getDecompressor(input)
    if c, ok := decompressor.(io.Closer); ok {
        defer c.Close()
    }
    io.Copy(os.Stdout, decompressor)
}

func main() {
    flag.Parse()
    if *input == "" {
        log.Fatalf("Please specify an input file with -input")
    }
    switch {
    case *compress:
        compression()
    case *decompress:
        decompression()
    default:
        log.Println("must specify one of -compress or -decompress")
    }
}
<|endoftext|>


<commit_before>package ui

import (
	"fmt"
	"regexp"
	"strconv"
	"unicode"

	"github.com/atotto/clipboard"
	"github.com/tcolar/goed/core"
)

func (v *View) ClearSelections() {
	v.selections = []core.Selection{}
}

// Text returns the text contained in the selection of the given view
// Note: **NOT** a rectangle but from pt1 to pt2
func (v *View) SelectionText(s *core.Selection) [][]rune {
	cf := s.ColFrom
	ct := s.ColTo
	lt := s.LineTo
	lf := s.LineFrom
	if lf == lt {
		return *v.backend.Slice(lf, cf, lt, ct).Text()
	}
	// first line
	text := *v.backend.Slice(lf, cf, lf, -1).Text()
	for l := lf + 1; l < lt; l++ {
		// middle
		text = append(text, *v.backend.Slice(l, 0, l, -1).Text()...)
	}
	// last line
	text = append(text, *v.backend.Slice(lt, 0, lt, ct).Text()...)
	return text
}

// Selected returns whether the text at line, col is current selected
// also returns the matching selection, if any.
func (v *View) Selected(col, line int) (bool, *core.Selection) {
	for _, s := range v.selections {
		if line < s.LineFrom || line > s.LineTo {
			continue
		} else if line > s.LineFrom && line < s.LineTo {
			return true, &s
		} else if s.LineFrom == s.LineTo {
			return col >= s.ColFrom && col <= s.ColTo, &s
		} else if line == s.LineFrom && col >= s.ColFrom {
			return true, &s
		} else if line == s.LineTo && col <= s.ColTo {
			return true, &s
		}
	}
	return false, nil
}

func (v *View) Copy() {
	if len(v.selections) == 0 {
		v.SelectLine(v.CurLine())
	}
	v.SelectionCopy(&v.selections[0])
}

func (v *View) Delete() {
	if len(v.selections) == 0 {
		v.SelectLine(v.CurLine())
	}
	v.SelectionDelete(&v.selections[0])
}

func (v *View) SelectionCopy(s *core.Selection) {
	t := v.SelectionText(s)
	core.Ed.SetStatus(fmt.Sprintf("Copied %d lines to clipboard.", len(t)))
	clipboard.WriteAll(core.RunesToString(t))
}

func (v *View) SelectionDelete(s *core.Selection) {
	v.delete(s.LineFrom, s.ColFrom, s.LineTo, s.ColTo)
}

func (v *View) Paste() {
	text, err := clipboard.ReadAll()
	if err != nil {
		core.Ed.SetStatusErr(err.Error())
		return
	}
	if len(v.selections) > 0 {
		v.DeleteCur()
	}
	_, y, x := v.CurChar()
	v.Insert(y, x, text)
}

var locationRegexp = regexp.MustCompile(`([^"\s(){}[\]<>,?|+=&^%#@!;':]+)(:\d+)?(:\d+)?`)

// Try to select a "location" from the given position
// a location is a path with possibly a line number and maybe a column number as well
func (v *View) ExpandSelectionPath(line, col int) *core.Selection {
	l := v.Line(v.slice, line)
	ln := string(l)
	slice := core.NewSlice(0, 0, 0, len(l), [][]rune{l})
	c := v.LineRunesTo(slice, 0, col)
	matches := locationRegexp.FindAllStringIndex(ln, -1)
	var best []int
	// Find the "narrowest" match around the cursor
	for _, s := range matches {
		if s[0] <= c && s[1] >= c {
			if best == nil || s[1]-s[0] < best[1]-best[0] {
				best = s
			}
		}
	}
	if best == nil {
		return nil
	}
	return core.NewSelection(line, best[0], line, best[1]-1)
}

// Try to select the longest "word" from current position.
func (v *View) ExpandSelectionWord(line, col int) *core.Selection {
	l := v.Line(v.slice, line)
	c := v.LineRunesTo(v.slice, line, col)
	if c < 0 || c >= len(l) {
		return nil
	}
	c1, c2 := c, c
	for ; c1 >= 0 && isWordRune(l[c1]); c1-- {
	}
	c1++
	for ; c2 < len(l) && isWordRune(l[c2]); c2++ {
	}
	c2--
	if c1 >= c2 {
		return nil
	}
	return core.NewSelection(line, c1, line, c2)
}

func isWordRune(r rune) bool {
	return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_'
}

// Select the whole given line
func (v *View) SelectLine(line int) {
	s := core.NewSelection(line, 0, line, v.LineLen(v.slice, line))
	v.selections = []core.Selection{
		*s,
	}
}

// Select a word at the given location (if any)
func (v *View) SelectWord(line, col int) {
	s := v.ExpandSelectionWord(line, col)
	if s != nil {
		v.selections = []core.Selection{
			*s,
		}
	}
}

// Parses a selection into a location (file, line, col)
func (v *View) SelectionToLoc(sel *core.Selection) (loc string, line, col int) {
	sub := locationRegexp.FindAllStringSubmatch(core.RunesToString(v.SelectionText(sel)), 1)
	if len(sub) == 0 {
		return
	}
	s := sub[0]
	if len(s) >= 1 {
		loc = s[1]
	}
	if len(s[2]) > 0 {
		line, _ = strconv.Atoi(s[2][1:])
	}
	if len(s[3]) > 0 {
		col, _ = strconv.Atoi(s[3][1:])
	}
	return loc, line, col
}

// Stretch a selection toward a new position
func (v *View) StretchSelection(prevl, prevc, l, c int) {
	if len(v.selections) == 0 {
		s := *core.NewSelection(prevl, prevc, l, c)
		v.selections = []core.Selection{
			s,
		}
	} else {
		s := v.selections[0]
		if s.LineTo == prevl && s.ColTo == prevc {
			s.LineTo, s.ColTo = l, c
		} else {
			s.LineFrom, s.ColFrom = l, c
		}
		s.Normalize()
		v.selections[0] = s
	}
}

// Open what's selected or under the cursor
// if newView is true then open in a new view, otherwise
// replace content of v
func (v *View) OpenSelection(newView bool) {
	ed := core.Ed.(*Editor)
	newView = newView || v.Dirty()
	if len(v.selections) == 0 {
		selection := v.ExpandSelectionPath(v.CurLine(), v.CurCol())
		if selection == nil {
			ed.SetStatusErr("Could not expand location from cursor location.")
			return
		}
		v.selections = []core.Selection{*selection}
	}
	loc, line, col := v.SelectionToLoc(&v.selections[0])
	line-- // we use 0 indexes in views
	col--
	isDir := false
	loc, isDir = core.LookupLocation(v.WorkDir(), loc)
	vv := ed.ViewByLoc(loc)
	if vv != nil {
		// Already open
		ed.ActivateView(vv, line, col)
		return
	}
	v2 := ed.NewView(loc)
	if _, err := ed.Open(loc, v2, v.WorkDir(), false); err != nil {
		ed.SetStatusErr(err.Error())
		return
	}
	if newView {
		if isDir {
			ed.InsertView(v2, v, 0.5)
		} else {
			ed.InsertViewSmart(v2)
		}
	} else {
		ed.ReplaceView(v, v2)
	}
	ed.ActivateView(v2, line, col)
}
<commit_msg>Avoid empty copies to clipboard<commit_after>package ui

import (
	"fmt"
	"regexp"
	"strconv"
	"unicode"

	"github.com/atotto/clipboard"
	"github.com/tcolar/goed/core"
)

func (v *View) ClearSelections() {
	v.selections = []core.Selection{}
}

// Text returns the text contained in the selection of the given view
// Note: **NOT** a rectangle but from pt1 to pt2
func (v *View) SelectionText(s *core.Selection) [][]rune {
	cf := s.ColFrom
	ct := s.ColTo
	lt := s.LineTo
	lf := s.LineFrom
	if lf == lt {
		return *v.backend.Slice(lf, cf, lt, ct).Text()
	}
	// first line
	text := *v.backend.Slice(lf, cf, lf, -1).Text()
	for l := lf + 1; l < lt; l++ {
		// middle
		text = append(text, *v.backend.Slice(l, 0, l, -1).Text()...)
	}
	// last line
	text = append(text, *v.backend.Slice(lt, 0, lt, ct).Text()...)
	return text
}

// Selected returns whether the text at line, col is current selected
// also returns the matching selection, if any.
func (v *View) Selected(col, line int) (bool, *core.Selection) {
	for _, s := range v.selections {
		if line < s.LineFrom || line > s.LineTo {
			continue
		} else if line > s.LineFrom && line < s.LineTo {
			return true, &s
		} else if s.LineFrom == s.LineTo {
			return col >= s.ColFrom && col <= s.ColTo, &s
		} else if line == s.LineFrom && col >= s.ColFrom {
			return true, &s
		} else if line == s.LineTo && col <= s.ColTo {
			return true, &s
		}
	}
	return false, nil
}

func (v *View) Copy() {
	if len(v.selections) == 0 {
		v.SelectLine(v.CurLine())
	}
	v.SelectionCopy(&v.selections[0])
}

func (v *View) Delete() {
	if len(v.selections) == 0 {
		v.SelectLine(v.CurLine())
	}
	v.SelectionDelete(&v.selections[0])
}

func (v *View) SelectionCopy(s *core.Selection) {
	t := v.SelectionText(s)
	text := core.RunesToString(t)
	if len(text) == 0 {
		return
	}
	core.Ed.SetStatus(fmt.Sprintf("Copied %d lines to clipboard.", len(t)))
	clipboard.WriteAll(text)
}

func (v *View) SelectionDelete(s *core.Selection) {
	v.delete(s.LineFrom, s.ColFrom, s.LineTo, s.ColTo)
}

func (v *View) Paste() {
	text, err := clipboard.ReadAll()
	if err != nil {
		core.Ed.SetStatusErr(err.Error())
		return
	}
	if len(v.selections) > 0 {
		v.DeleteCur()
	}
	_, y, x := v.CurChar()
	v.Insert(y, x, text)
}

var locationRegexp = regexp.MustCompile(`([^"\s(){}[\]<>,?|+=&^%#@!;':]+)(:\d+)?(:\d+)?`)

// Try to select a "location" from the given position
// a location is a path with possibly a line number and maybe a column number as well
func (v *View) ExpandSelectionPath(line, col int) *core.Selection {
	l := v.Line(v.slice, line)
	ln := string(l)
	slice := core.NewSlice(0, 0, 0, len(l), [][]rune{l})
	c := v.LineRunesTo(slice, 0, col)
	matches := locationRegexp.FindAllStringIndex(ln, -1)
	var best []int
	// Find the "narrowest" match around the cursor
	for _, s := range matches {
		if s[0] <= c && s[1] >= c {
			if best == nil || s[1]-s[0] < best[1]-best[0] {
				best = s
			}
		}
	}
	if best == nil {
		return nil
	}
	return core.NewSelection(line, best[0], line, best[1]-1)
}

// Try to select the longest "word" from current position.
func (v *View) ExpandSelectionWord(line, col int) *core.Selection {
	l := v.Line(v.slice, line)
	c := v.LineRunesTo(v.slice, line, col)
	if c < 0 || c >= len(l) {
		return nil
	}
	c1, c2 := c, c
	for ; c1 >= 0 && isWordRune(l[c1]); c1-- {
	}
	c1++
	for ; c2 < len(l) && isWordRune(l[c2]); c2++ {
	}
	c2--
	if c1 >= c2 {
		return nil
	}
	return core.NewSelection(line, c1, line, c2)
}

func isWordRune(r rune) bool {
	return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_'
}

// Select the whole given line
func (v *View) SelectLine(line int) {
	s := core.NewSelection(line, 0, line, v.LineLen(v.slice, line))
	v.selections = []core.Selection{
		*s,
	}
}

// Select a word at the given location (if any)
func (v *View) SelectWord(line, col int) {
	s := v.ExpandSelectionWord(line, col)
	if s != nil {
		v.selections = []core.Selection{
			*s,
		}
	}
}

// Parses a selection into a location (file, line, col)
func (v *View) SelectionToLoc(sel *core.Selection) (loc string, line, col int) {
	sub := locationRegexp.FindAllStringSubmatch(core.RunesToString(v.SelectionText(sel)), 1)
	if len(sub) == 0 {
		return
	}
	s := sub[0]
	if len(s) >= 1 {
		loc = s[1]
	}
	if len(s[2]) > 0 {
		line, _ = strconv.Atoi(s[2][1:])
	}
	if len(s[3]) > 0 {
		col, _ = strconv.Atoi(s[3][1:])
	}
	return loc, line, col
}

// Stretch a selection toward a new position
func (v *View) StretchSelection(prevl, prevc, l, c int) {
	if len(v.selections) == 0 {
		s := *core.NewSelection(prevl, prevc, l, c)
		v.selections = []core.Selection{
			s,
		}
	} else {
		s := v.selections[0]
		if s.LineTo == prevl && s.ColTo == prevc {
			s.LineTo, s.ColTo = l, c
		} else {
			s.LineFrom, s.ColFrom = l, c
		}
		s.Normalize()
		v.selections[0] = s
	}
}

// Open what's selected or under the cursor
// if newView is true then open in a new view, otherwise
// replace content of v
func (v *View) OpenSelection(newView bool) {
	ed := core.Ed.(*Editor)
	newView = newView || v.Dirty()
	if len(v.selections) == 0 {
		selection := v.ExpandSelectionPath(v.CurLine(), v.CurCol())
		if selection == nil {
			ed.SetStatusErr("Could not expand location from cursor location.")
			return
		}
		v.selections = []core.Selection{*selection}
	}
	loc, line, col := v.SelectionToLoc(&v.selections[0])
	line-- // we use 0 indexes in views
	col--
	isDir := false
	loc, isDir = core.LookupLocation(v.WorkDir(), loc)
	vv := ed.ViewByLoc(loc)
	if vv != nil {
		// Already open
		ed.ActivateView(vv, line, col)
		return
	}
	v2 := ed.NewView(loc)
	if _, err := ed.Open(loc, v2, v.WorkDir(), false); err != nil {
		ed.SetStatusErr(err.Error())
		return
	}
	if newView {
		if isDir {
			ed.InsertView(v2, v, 0.5)
		} else {
			ed.InsertViewSmart(v2)
		}
	} else {
		ed.ReplaceView(v, v2)
	}
	ed.ActivateView(v2, line, col)
}
<|endoftext|>


<commit_before>package skydns2

import (
	"log"
	"net/url"
	"strconv"
	"strings"

	"github.com/coreos/go-etcd/etcd"
	"github.com/gliderlabs/registrator/bridge"
)

func init() {
	bridge.Register(new(Factory), "skydns2")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	urls := make([]string, 0)
	if uri.Host != "" {
		urls = append(urls, "http://"+uri.Host)
	}
	return &Skydns2Adapter{client: etcd.NewClient(urls), path: domainPath(uri.Path[1:])}
}

type Skydns2Adapter struct {
	client *etcd.Client
	path   string
}

func (r *Skydns2Adapter) Ping() error {
	rr := etcd.NewRawRequest("GET", "version", nil, nil)
	_, err := r.client.SendRequest(rr)
	if err != nil {
		return err
	}
	return nil
}

func (r *Skydns2Adapter) Register(service *bridge.Service) error {
	port := strconv.Itoa(service.Port)
	record := `{"host":"` + service.IP + `","port":` + port + `}`
	_, err := r.client.Set(r.servicePath(service), record, uint64(service.TTL))
	if err != nil {
		log.Println("skydns2: failed to register service:", err)
	}
	return err
}

func (r *Skydns2Adapter) Deregister(service *bridge.Service) error {
	_, err := r.client.Delete(r.servicePath(service), false)
	if err != nil {
		log.Println("skydns2: failed to register service:", err)
	}
	return err
}

func (r *Skydns2Adapter) Refresh(service *bridge.Service) error {
	return r.Register(service)
}

func (r *Skydns2Adapter) servicePath(service *bridge.Service) string {
	return r.path + "/" + service.Name + "/" + service.ID
}

func domainPath(domain string) string {
	components := strings.Split(domain, ".")
	for i, j := 0, len(components)-1; i < j; i, j = i+1, j-1 {
		components[i], components[j] = components[j], components[i]
	}
	return "/skydns/" + strings.Join(components, "/")
}
<commit_msg>Fix issue where registrator would panic from invalid skydns2 URI.<commit_after>package skydns2

import (
	"log"
	"net/url"
	"strconv"
	"strings"

	"github.com/coreos/go-etcd/etcd"
	"github.com/gliderlabs/registrator/bridge"
)

func init() {
	bridge.Register(new(Factory), "skydns2")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	urls := make([]string, 0)
	if uri.Host != "" {
		urls = append(urls, "http://"+uri.Host)
	}

	if len(uri.Path) < 2 {
		log.Fatal("skydns2: dns domain required e.g.: skydns2://<host>/<domain>")
	}

	return &Skydns2Adapter{client: etcd.NewClient(urls), path: domainPath(uri.Path[1:])}
}

type Skydns2Adapter struct {
	client *etcd.Client
	path   string
}

func (r *Skydns2Adapter) Ping() error {
	rr := etcd.NewRawRequest("GET", "version", nil, nil)
	_, err := r.client.SendRequest(rr)
	if err != nil {
		return err
	}
	return nil
}

func (r *Skydns2Adapter) Register(service *bridge.Service) error {
	port := strconv.Itoa(service.Port)
	record := `{"host":"` + service.IP + `","port":` + port + `}`
	_, err := r.client.Set(r.servicePath(service), record, uint64(service.TTL))
	if err != nil {
		log.Println("skydns2: failed to register service:", err)
	}
	return err
}

func (r *Skydns2Adapter) Deregister(service *bridge.Service) error {
	_, err := r.client.Delete(r.servicePath(service), false)
	if err != nil {
		log.Println("skydns2: failed to register service:", err)
	}
	return err
}

func (r *Skydns2Adapter) Refresh(service *bridge.Service) error {
	return r.Register(service)
}

func (r *Skydns2Adapter) servicePath(service *bridge.Service) string {
	return r.path + "/" + service.Name + "/" + service.ID
}

func domainPath(domain string) string {
	components := strings.Split(domain, ".")
	for i, j := 0, len(components)-1; i < j; i, j = i+1, j-1 {
		components[i], components[j] = components[j], components[i]
	}
	return "/skydns/" + strings.Join(components, "/")
}
<|endoftext|>


<commit_before>// ATM - Automatic TempUrl Maker
// A builder of Swift TempURLs
// Copyright (c) 2015 Stuart Glenn
// All rights reserved
// Use of this source code is goverened by a BSD 3-clause license,
// see included LICENSE file for details
package main

import (
	"fmt"
	"log"
	"os/user"

	"github.com/codegangsta/cli"
	"github.com/glennsb/atm"
	"github.com/howeyc/gopass"
)

func main() {
	app := cli.NewApp()
	app.Name = "atm"
	app.Usage = "Automated TempURL Maker"
	app.Version = "0.0.1 - 20151025"
	app.Author = "Stuart Glenn"
	app.Email = "Stuart-Glenn@omrf.org"
	app.Copyright = "2015 Stuart Glenn, All rights reserved"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "V, verbose",
			Usage: "show more output",
		},
	}

	app.Commands = clientCommands()
	app.Commands = append(app.Commands, serverCommand())
	app.RunAndExitOnError()
}

func clientCommands() []cli.Command {
	return []cli.Command{
		cli.Command{
			Name:  "url",
			Usage: "Request a temp url",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "api-key, k",
					Usage:  "account/user atm api-key",
					EnvVar: "ATM_API_KEY",
				},
				cli.StringFlag{
					Name:   "api-secret, s",
					Usage:  "account/user atm api-secret",
					EnvVar: "ATM_API_SECRET",
				},
				cli.StringFlag{
					Name:   "host, h",
					Usage:  "atm server endpoint",
					EnvVar: "ATM_HOST",
				},
				cli.StringFlag{
					Name:  "method, m",
					Usage: "HTTP method requested for temp url",
					Value: "GET",
				},
			},
			Action: func(c *cli.Context) {
				log.Fatal("Not implemented yet")
				return
			},
		},
		cli.Command{
			Name:  "key",
			Usage: "Add/Remove signing key",
			Action: func(c *cli.Context) {
				log.Fatal("Not implemented yet")
				return
			},
		},
	}
}

func serverFlags() []cli.Flag {
	current_user, err := user.Current()
	default_username := ""
	if nil == err {
		default_username = current_user.Username
	}
	return []cli.Flag{
		cli.StringFlag{
			Name:  "database",
			Usage: "name of database",
			Value: "atm",
		},
		cli.StringFlag{
			Name:  "database-host",
			Usage: "hostname of database server",
			Value: "localhost",
		},
		cli.StringFlag{
			Name:  "database-user",
			Usage: "username for database connection",
			Value: default_username,
		},
		cli.IntFlag{
			Name:  "database-port",
			Usage: "port number of database server",
			Value: 3306,
		},
		cli.DurationFlag{
			Name:  "duration",
			Usage: "Default lifetime for generated tempurl",
			Value: atm.DURATION,
		},
		cli.StringFlag{
			Name:  "object-host, host",
			Usage: "Swift service host prefix",
			Value: atm.HOST,
		},
	}
}

func serverCommand() cli.Command {
	return cli.Command{
		Name:  "server",
		Usage: "Run webservice",
		Flags: serverFlags(),
		Action: func(c *cli.Context) {
			db_user := c.String("database-user")
			db_host := c.String("database-host")
			db := c.String("database")

			fmt.Printf("%s@%s/%s password: ", db_user, db_host, db)
			db_pass := string(gopass.GetPasswd())

			ds, err := atm.NewDatastore("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s",
				db_user, db_pass, db_host, c.Int("database-port"), db))
			if nil != err {
				log.Fatal(err)
				return
			}
			db_pass = ""
			defer ds.Close()

			service := &atm.Server{
				Ds:               ds,
				Object_host:      c.String("object-host"),
				Default_duration: int64(c.Duration("duration").Minutes()),
			}
			service.Run()
		},
	}
}
<commit_msg>Expand url command to check for required arguments<commit_after>// ATM - Automatic TempUrl Maker
// A builder of Swift TempURLs
// Copyright (c) 2015 Stuart Glenn
// All rights reserved
// Use of this source code is goverened by a BSD 3-clause license,
// see included LICENSE file for details
package main

import (
	"fmt"
	"log"
	"os"
	"os/user"

	"github.com/codegangsta/cli"
	"github.com/glennsb/atm"
	"github.com/howeyc/gopass"
)

func main() {
	app := cli.NewApp()
	app.Name = "atm"
	app.Usage = "Automated TempURL Maker"
	app.Version = "0.0.1 - 20151025"
	app.Author = "Stuart Glenn"
	app.Email = "Stuart-Glenn@omrf.org"
	app.Copyright = "2015 Stuart Glenn, All rights reserved"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "V, verbose",
			Usage: "show more output",
		},
	}

	app.Commands = clientCommands()
	app.Commands = append(app.Commands, serverCommand())
	app.RunAndExitOnError()
}

func clientCommands() []cli.Command {
	return []cli.Command{
		cli.Command{
			Name:        "url",
			Usage:       "Request a temp url to Account/Container/Object",
			ArgsUsage:   "<Account> <Container> <Object> ",
			Description: "Send a request to the ATM service for a tempurl",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:   "api-key, k",
					Usage:  "account/user atm api-key",
					EnvVar: "ATM_API_KEY",
				},
				cli.StringFlag{
					Name:   "api-secret, s",
					Usage:  "account/user atm api-secret",
					EnvVar: "ATM_API_SECRET",
				},
				cli.StringFlag{
					Name:   "atm-host, a",
					Usage:  "atm server endpoint",
					EnvVar: "ATM_HOST",
				},
				cli.StringFlag{
					Name:  "method, m",
					Usage: "HTTP method requested for temp url",
					Value: "GET",
				},
			},
			Action: func(c *cli.Context) {
				method := c.String("method")
				if "" == method {
					fmt.Fprintf(os.Stderr, "Missing HTTP method option\n")
					cli.ShowSubcommandHelp(c)
					os.Exit(1)
				}
				account := c.Args().Get(0)
				if "" == account {
					fmt.Fprintf(os.Stderr, "Missing Account argument\n")
					cli.ShowSubcommandHelp(c)
					os.Exit(1)
				}
				container := c.Args().Get(1)
				if "" == container {
					fmt.Fprintf(os.Stderr, "Missing Container argument\n")
					cli.ShowSubcommandHelp(c)
					os.Exit(1)
				}
				object := c.Args().Get(2)
				if "" == object {
					fmt.Fprintf(os.Stderr, "Missing Object argument\n")
					cli.ShowSubcommandHelp(c)
					os.Exit(1)
				}
				fmt.Printf("Requesting %s to %s/%s/%s\n", method, account, container, object)
			},
		},

		cli.Command{
			Name:  "key",
			Usage: "Add/Remove signing key",
			Action: func(c *cli.Context) {
				log.Fatal("Not implemented yet")
				return
			},
		},
	}
}

func serverFlags() []cli.Flag {
	current_user, err := user.Current()
	default_username := ""
	if nil == err {
		default_username = current_user.Username
	}
	return []cli.Flag{
		cli.StringFlag{
			Name:  "database",
			Usage: "name of database",
			Value: "atm",
		},
		cli.StringFlag{
			Name:  "database-host",
			Usage: "hostname of database server",
			Value: "localhost",
		},
		cli.StringFlag{
			Name:  "database-user",
			Usage: "username for database connection",
			Value: default_username,
		},
		cli.IntFlag{
			Name:  "database-port",
			Usage: "port number of database server",
			Value: 3306,
		},
		cli.DurationFlag{
			Name:  "duration",
			Usage: "Default lifetime for generated tempurl",
			Value: atm.DURATION,
		},
		cli.StringFlag{
			Name:  "object-host, host",
			Usage: "Swift service host prefix",
			Value: atm.HOST,
		},
	}
}

func serverCommand() cli.Command {
	return cli.Command{
		Name:  "server",
		Usage: "Run webservice",
		Flags: serverFlags(),
		Action: func(c *cli.Context) {
			db_user := c.String("database-user")
			db_host := c.String("database-host")
			db := c.String("database")

			fmt.Printf("%s@%s/%s password: ", db_user, db_host, db)
			db_pass := string(gopass.GetPasswd())

			ds, err := atm.NewDatastore("mysql", fmt.Sprintf("%s:%s@tcp(%s:%d)/%s",
				db_user, db_pass, db_host, c.Int("database-port"), db))
			if nil != err {
				log.Fatal(err)
				return
			}
			db_pass = ""
			defer ds.Close()

			service := &atm.Server{
				Ds:               ds,
				Object_host:      c.String("object-host"),
				Default_duration: int64(c.Duration("duration").Minutes()),
			}
			service.Run()
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"gnat"
	"log"
	"net/http"

	"strings"

	"github.com/ccding/go-stun/stun"
	b58 "github.com/jbenet/go-base58"
)

var maxRouteCacheSize = 500
var addr = flag.String("localhost", ":2222", "http service address")
var dht *gnat.DHT
var hub *Hub
var routes map[string]*gnat.NetworkNode

func main() {
	// keep a route cache
	routes = make(map[string]*gnat.NetworkNode)

	// test the network to discovery what type of NAT (if any)
	// the client is behind.

	fmt.Println("GNAT Node v0.0.1")
	fmt.Println("  *Documentation:   https://gnat.cs.brown.edu/docs")
	fmt.Println("  *Support:         https://gnat.cs.brown.edu/support")
	fmt.Println("  *GitHub:          https://github.com/ogisan/gnat")
	fmt.Println("  For more information, visit: http://gnat.cs.brown.edu")
	fmt.Println("--------------------------------------------------------")
	fmt.Print("1) Testing network...")
	nat, host, err := stun.NewClient().Discover()
	if err != nil {
		fmt.Println("Error:a problem occured while testing your network.")
		fmt.Println("TODO: try again later.")
	}

	fmt.Println("done")
	// acceptable type of NATs
	if nat == stun.NATNone || nat == stun.NATFull {
		fmt.Println("  Network NAT configuration: " + nat.String())
		fmt.Println("  Node address: " + host.String())
		initializeDHT()
		setupServer()
		fmt.Println("GNAT node setup and running!")
	} else {
		fmt.Println("Error: your network configuration does not support running a GNAT node.")
		fmt.Println("TODO: update your router settings to have less restrictive settings and try again.")
	}
}

func onForwardRequestReceived(forwardToIP string, rqst []byte) {
	hub.sendMessageToAddr(forwardToIP, rqst)
}

func onForwardData(fromAddr string, header map[string]string, data []byte) {

	resp := map[string]string{}

	sendTo := header["send_to"]
	if sendTo == "" {
		return
	}

	fmt.Println("Received forwarding request from " + fromAddr)

	resp["from"] = fromAddr
	respHeader, _ := json.Marshal(resp)
	forwardMessage(sendTo, append(respHeader, data...))
}

func handConnectionRequest(w http.ResponseWriter, r *http.Request) {

	// generate digest hash of IP address
	clientIP := strings.Split(r.RemoteAddr, ":")[0]
	ipDigest := sha256.Sum256([]byte(clientIP))
	id := b58.Encode(ipDigest[:])

	// find the node connected to this client ip
	node, err := dht.FindNode(id)

	if err == nil {

		if string(node.ID) == string(dht.GetSelfID()) {
			fmt.Println("New connection from " + r.RemoteAddr)
			//log.Println(r.URL)

			if r.URL.Path != "/" {
				http.Error(w, "Not found", 404)
				return
			}

			if r.Method != "GET" {
				http.Error(w, "Method not allowed", 405)
				return
			}

			http.ServeFile(w, r, "./static/home.html")

		} else {
			fmt.Println("Redirecting " + r.RemoteAddr + " to http://" + node.IP.String() + ":2222")
			http.Redirect(w, r, "http://"+node.IP.String()+":2222", 301)
		}

	} else {
		fmt.Println(err)
	}
}

func setupServer() {

	fmt.Print("4) Setting up HTTP server...")
	flag.Parse()
	hub = newHub()
	go hub.run(onForwardData)
	http.HandleFunc("/", handConnectionRequest)
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		serveWs(hub, w, r)
	})
	fmt.Println("done")
	fmt.Println("Listening on http://127.0.0.1" + *addr)
	err := http.ListenAndServe(*addr, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

func initializeDHT() {
	var ip = flag.String("ip", "0.0.0.0", "IP Address to use")
	var port = flag.String("port", "1443", "Port to use")
	var bIP = flag.String("bip", "45.55.18.163", "IP Address to bootstrap against")
	var bPort = flag.String("bport", "1443", "Port to bootstrap against")
	var stun = flag.Bool("stun", true, "Use STUN")

	flag.Parse()

	var bootstrapNodes []*gnat.NetworkNode
	if *bIP != "" || *bPort != "" {
		bootstrapNode := gnat.NewNetworkNode(*bIP, *bPort)
		bootstrapNodes = append(bootstrapNodes, bootstrapNode)
	}

	var err error
	dht, err = gnat.NewDHT(&gnat.Options{
		BootstrapNodes:   bootstrapNodes,
		IP:               *ip,
		Port:             *port,
		UseStun:          *stun,
		OnForwardRequest: onForwardRequestReceived,
	})

	fmt.Print("2) Opening socket...")

	err = dht.CreateSocket()
	if err != nil {
		panic(err)
	}

	fmt.Println("done")
	go func() {
		err := dht.Listen()
		panic(err)
	}()

	if len(bootstrapNodes) > 0 {
		fmt.Print("3) Bootstrapping into GNAT network...")
		dht.Bootstrap()
		fmt.Println("done")
	}
}

func forwardMessage(ip string, msg []byte) {
	ipDigest := sha256.Sum256([]byte(ip))
	id := b58.Encode(ipDigest[:])
	fmt.Print("Finding forwarding node...")

	var err error
	var foundNode *gnat.NetworkNode
	if node, ok := routes[id]; ok {
		foundNode = node
	} else {
		foundNode, err = dht.FindNode(id)
		routes[id] = foundNode
		if len(routes) > maxRouteCacheSize {
			// clear the list
			routes = make(map[string]*gnat.NetworkNode)
		}
	}

	fmt.Println("done")
	if err != nil {
		fmt.Println(err.Error())
	} else {
		fmt.Println("Forwarding data to ", foundNode.IP.String())
		dht.ForwardData(foundNode, gnat.NewNetworkNode(ip, "0"), msg)
	}
}
<commit_msg>added id printing for requests<commit_after>package main

import (
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"gnat"
	"log"
	"net/http"

	"strings"

	"github.com/ccding/go-stun/stun"
	b58 "github.com/jbenet/go-base58"
)

var maxRouteCacheSize = 500
var addr = flag.String("localhost", ":2222", "http service address")
var dht *gnat.DHT
var hub *Hub
var routes map[string]*gnat.NetworkNode

func main() {
	// keep a route cache
	routes = make(map[string]*gnat.NetworkNode)

	// test the network to discovery what type of NAT (if any)
	// the client is behind.

	fmt.Println("GNAT Node v0.0.1")
	fmt.Println("  *Documentation:   https://gnat.cs.brown.edu/docs")
	fmt.Println("  *Support:         https://gnat.cs.brown.edu/support")
	fmt.Println("  *GitHub:          https://github.com/ogisan/gnat")
	fmt.Println("  For more information, visit: http://gnat.cs.brown.edu")
	fmt.Println("--------------------------------------------------------")
	fmt.Print("1) Testing network...")
	nat, host, err := stun.NewClient().Discover()
	if err != nil {
		fmt.Println("Error:a problem occured while testing your network.")
		fmt.Println("TODO: try again later.")
	}

	fmt.Println("done")
	// acceptable type of NATs
	if nat == stun.NATNone || nat == stun.NATFull {
		fmt.Println("  Network NAT configuration: " + nat.String())
		fmt.Println("  Node address: " + host.String())
		initializeDHT()
		setupServer()
		fmt.Println("GNAT node setup and running!")
	} else {
		fmt.Println("Error: your network configuration does not support running a GNAT node.")
		fmt.Println("TODO: update your router settings to have less restrictive settings and try again.")
	}
}

func onForwardRequestReceived(forwardToIP string, rqst []byte) {
	hub.sendMessageToAddr(forwardToIP, rqst)
}

func onForwardData(fromAddr string, header map[string]string, data []byte) {

	resp := map[string]string{}

	sendTo := header["send_to"]
	if sendTo == "" {
		return
	}

	fmt.Println("Received forwarding request from " + fromAddr)

	resp["from"] = fromAddr
	respHeader, _ := json.Marshal(resp)
	forwardMessage(sendTo, append(respHeader, data...))
}

func handConnectionRequest(w http.ResponseWriter, r *http.Request) {

	// generate digest hash of IP address
	clientIP := strings.Split(r.RemoteAddr, ":")[0]
	ipDigest := sha256.Sum256([]byte(clientIP))
	id := b58.Encode(ipDigest[:])

	// find the node connected to this client ip
	node, err := dht.FindNode(id)

	if err == nil {
		if string(node.ID) == string(dht.GetSelfID()) {
			fmt.Println("New connection from " + r.RemoteAddr + " id: " + id)
			//log.Println(r.URL)

			if r.URL.Path != "/" {
				http.Error(w, "Not found", 404)
				return
			}

			if r.Method != "GET" {
				http.Error(w, "Method not allowed", 405)
				return
			}

			http.ServeFile(w, r, "./static/home.html")

		} else {
			fmt.Println("Redirecting " + r.RemoteAddr + " to http://" + node.IP.String() + ":2222")
			http.Redirect(w, r, "http://"+node.IP.String()+":2222", 301)
		}

	} else {
		fmt.Println(err)
	}
}

func setupServer() {

	fmt.Print("4) Setting up HTTP server...")
	flag.Parse()
	hub = newHub()
	go hub.run(onForwardData)
	http.HandleFunc("/", handConnectionRequest)
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		serveWs(hub, w, r)
	})
	fmt.Println("done")
	fmt.Println("Listening on http://127.0.0.1" + *addr)
	err := http.ListenAndServe(*addr, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

func initializeDHT() {
	var ip = flag.String("ip", "0.0.0.0", "IP Address to use")
	var port = flag.String("port", "1443", "Port to use")
	var bIP = flag.String("bip", "45.55.18.163", "IP Address to bootstrap against")
	var bPort = flag.String("bport", "1443", "Port to bootstrap against")
	var stun = flag.Bool("stun", true, "Use STUN")

	flag.Parse()

	var bootstrapNodes []*gnat.NetworkNode
	if *bIP != "" || *bPort != "" {
		bootstrapNode := gnat.NewNetworkNode(*bIP, *bPort)
		bootstrapNodes = append(bootstrapNodes, bootstrapNode)
	}

	var err error
	dht, err = gnat.NewDHT(&gnat.Options{
		BootstrapNodes:   bootstrapNodes,
		IP:               *ip,
		Port:             *port,
		UseStun:          *stun,
		OnForwardRequest: onForwardRequestReceived,
	})

	fmt.Print("2) Opening socket...")

	err = dht.CreateSocket()
	if err != nil {
		panic(err)
	}

	fmt.Println("done")
	go func() {
		err := dht.Listen()
		panic(err)
	}()

	if len(bootstrapNodes) > 0 {
		fmt.Print("3) Bootstrapping into GNAT network...")
		dht.Bootstrap()
		fmt.Println("done")
	} else {
		fmt.Println("3) Skipping GNAT bootstrap")
	}
}

func forwardMessage(ip string, msg []byte) {
	ipDigest := sha256.Sum256([]byte(ip))
	id := b58.Encode(ipDigest[:])
	fmt.Print("Finding forwarding node for id " + id + "...")

	var err error
	var foundNode *gnat.NetworkNode
	if node, ok := routes[id]; ok {
		foundNode = node
	} else {
		foundNode, err = dht.FindNode(id)
		routes[id] = foundNode
		if len(routes) > maxRouteCacheSize {
			// clear the list
			routes = make(map[string]*gnat.NetworkNode)
		}
	}

	fmt.Println("done")
	if err != nil {
		fmt.Println(err.Error())
	} else {
		fmt.Println("Forwarding data to ", foundNode.IP.String())
		dht.ForwardData(foundNode, gnat.NewNetworkNode(ip, "0"), msg)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/xlplbo/go_protobuf_test/protocol"
)

//Player struct
type Player struct {
	index  uint64
	conn   net.Conn
	s      *Server
	chStop chan error
}

//Play Run
func (p *Player) Play() {
	go func() {
		var data []byte
		buff := make([]byte, protocol.MaxSize)
		for {
			n, err := p.conn.Read(buff)
			if err != nil {
				log.Println(err)
				p.Stop()
				return
			}
			data = append(data, buff[:n]...)
			for {
				offset, serial, buff := protocol.UnPack(data)
				if buff == nil {
					break
				}
				data = data[offset:]
				if f, ok := p.s.handles[serial]; ok {
					f(p, buff)
					continue
				}
				log.Printf("protocol id(%d) not handle\n", serial)
			}
		}
	}()
	err := <-p.chStop
	p.conn.Close()
	p.s.DelPlayer(p.index)
	close(p.chStop)
	log.Println(err)
}

//Stop player
func (p *Player) Stop() {
	p.chStop <- fmt.Errorf("player(%d) stop", p.index)
}

//GetTargetPlayer ...
func (p *Player) GetTargetPlayer(index uint64) *Player {
	player := p.s.GetPlayer(index)
	if player == nil {
		log.Printf("player(%d) non-exsit", index)
		return nil
	}
	return player
}

//SendChat ...
func (p *Player) SendChat(msg string) {
	if err := protocol.Send2Client(p.conn, protocol.S2CCmd_Result, &protocol.S2CResult{
		Context: msg,
	}); err != nil {
		log.Println(err)
	}
}

//GetIndex ...
func (p *Player) GetIndex() uint64 {
	return p.index
}

//Server center
type Server struct {
	index   uint64
	players map[uint64]*Player
	handles map[int32]func(*Player, []byte)
	chStop  chan error
	chConn  chan net.Conn
	chSig   chan os.Signal
}

func (s *Server) getFreeIndex() uint64 {
	var i uint64 = 1
	for i = 1; i <= s.index; i++ {
		if _, ok := s.players[i]; !ok {
			return i
		}
	}
	s.index++
	return s.index
}

func (s *Server) brocastPlayerList() {
	var buf bytes.Buffer
	buf.WriteString("playerlist:")
	var array []string
	for index := range s.players {
		array = append(array, strconv.FormatUint(index, 10))
	}
	buf.WriteString(strings.Join(array, ","))
	for _, p := range s.players {
		p.SendChat(buf.String() + fmt.Sprintf(" your id: %d", p.GetIndex()))
	}
}

//Run start service
func (s *Server) Run() {
	go func() {
		for {
			conn := <-s.chConn
			index := s.getFreeIndex()
			player := &Player{
				index:  index,
				conn:   conn,
				s:      s,
				chStop: make(chan error),
			}
			s.players[index] = player
			go player.Play()
			s.brocastPlayerList()
			log.Printf("player(%d) %s connect.\n", index, conn.RemoteAddr().String())
		}
	}()

	msg := <-s.chStop
	var array []uint64
	for index := range s.players {
		array = append(array, index)
	}
	for _, index := range array {
		s.GetPlayer(index).Stop()
	}
	log.Printf("server stop: %s\n", msg.Error())
}

//ListenTCP only call func use go routine
func (s *Server) ListenTCP(laddr string) {
	l, err := net.Listen("tcp", laddr)
	if err != nil {
		s.chStop <- err
	}
	defer l.Close()
	log.Printf("listen at %s\n", l.Addr().String())
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println(err)
			continue
		}
		s.chConn <- conn
	}
}

//GetPlayer instance
func (s *Server) GetPlayer(key uint64) *Player {
	if player, ok := s.players[key]; ok {
		return player
	}
	return nil
}

//DelPlayer ...
func (s *Server) DelPlayer(key uint64) {
	delete(s.players, key)
}

//RegisterHandle ...
func (s *Server) RegisterHandle(id protocol.C2SCmd, f func(*Player, []byte)) {
	nID := int32(id)
	if _, ok := s.handles[nID]; ok {
		log.Printf("protocol(%d) handle repeat\n", nID)
		return
	}
	s.handles[nID] = f
	log.Printf("register handle protocol(%d)\n", nID)
}

//HandleSignal ...
func (s *Server) HandleSignal() {
	signal.Notify(s.chSig, os.Interrupt)
	sig := <-s.chSig
	s.chStop <- fmt.Errorf(sig.String())
}

//NewServer instance
func NewServer() *Server {
	s := &Server{
		index:   0,
		players: make(map[uint64]*Player),
		handles: make(map[int32]func(*Player, []byte)),
		chStop:  make(chan error),
		chConn:  make(chan net.Conn),
		chSig:   make(chan os.Signal),
	}
	s.RegisterHandle(protocol.C2SCmd_Abnormal, func(p *Player, msg []byte) {
		p.Stop()
	})
	s.RegisterHandle(protocol.C2SCmd_Chat, func(p *Player, msg []byte) {
		var chatMsg protocol.C2SChat
		if err := proto.Unmarshal(msg, &chatMsg); err != nil {
			p.Stop()
		}
		player := p.GetTargetPlayer(chatMsg.Index)
		if player != nil {
			player.SendChat(chatMsg.Context)
		}
	})
	return s
}

func main() {
	app := NewServer()
	go app.HandleSignal()
	go app.ListenTCP(":7788")
	app.Run()
}
<commit_msg>BUG:concurrent map read and map write<commit_after>package main

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"

	"github.com/golang/protobuf/proto"
	"github.com/xlplbo/go_protobuf_test/protocol"
)

//Player struct
type Player struct {
	index  uint64
	conn   net.Conn
	s      *Server
	chStop chan error
}

//Play Run
func (p *Player) Play() {
	go func() {
		var data []byte
		buff := make([]byte, protocol.MaxSize)
		for {
			n, err := p.conn.Read(buff)
			if err != nil {
				log.Println(err)
				p.Stop()
				return
			}
			data = append(data, buff[:n]...)
			for {
				offset, serial, buff := protocol.UnPack(data)
				if buff == nil {
					break
				}
				data = data[offset:]
				if f, ok := p.s.handles[serial]; ok {
					f(p, buff)
					continue
				}
				log.Printf("protocol id(%d) not handle\n", serial)
			}
		}
	}()
	err := <-p.chStop
	p.conn.Close()
	p.s.DelPlayer(p.index)
	log.Println(err)
}

//Stop player
func (p *Player) Stop() {
	p.chStop <- fmt.Errorf("player(%d) stop", p.index)
}

//GetTargetPlayer ...
func (p *Player) GetTargetPlayer(index uint64) *Player {
	player, ok := p.s.GetPlayer(index)
	if !ok {
		log.Printf("player(%d) non-exsit", index)
		return nil
	}
	return player
}

//SendChat ...
func (p *Player) SendChat(msg string) {
	if err := protocol.Send2Client(p.conn, protocol.S2CCmd_Result, &protocol.S2CResult{
		Context: msg,
	}); err != nil {
		log.Println(err)
	}
}

//GetIndex ...
func (p *Player) GetIndex() uint64 {
	return p.index
}

//Server center
type Server struct {
	index   uint64
	players map[uint64]*Player
	mutex   *sync.RWMutex
	handles map[int32]func(*Player, []byte)
	chStop  chan error
	chConn  chan net.Conn
	chSig   chan os.Signal
}

func (s *Server) getFreeIndex() uint64 {
	var i uint64 = 1
	for i = 1; i <= s.index; i++ {
		if _, ok := s.GetPlayer(i); !ok {
			return i
		}
	}
	s.index++
	return s.index
}

func (s *Server) getPlayerList() []*Player {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	var list []*Player
	for _, p := range s.players {
		list = append(list, p)
	}
	return list
}

//GetPlayer ...
func (s *Server) GetPlayer(index uint64) (*Player, bool) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	player, ok := s.players[index]
	return player, ok
}

func (s *Server) setPlayer(index uint64, p *Player) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.players[index] = p
}

func (s *Server) brocastPlayerList() {
	var buf bytes.Buffer
	buf.WriteString("playerlist:")
	var array []string
	for _, p := range s.getPlayerList() {
		array = append(array, strconv.FormatUint(p.GetIndex(), 10))
	}
	buf.WriteString(strings.Join(array, ","))
	for _, p := range s.getPlayerList() {
		p.SendChat(buf.String() + fmt.Sprintf(" your id: %d", p.GetIndex()))
	}
}

//Run start service
func (s *Server) Run() {
	go func() {
		for {
			conn := <-s.chConn
			index := s.getFreeIndex()
			player := &Player{
				index:  index,
				conn:   conn,
				s:      s,
				chStop: make(chan error),
			}
			s.setPlayer(index, player)
			go player.Play()
			s.brocastPlayerList()
			log.Printf("player(%d) %s connect.\n", index, conn.RemoteAddr().String())
		}
	}()

	msg := <-s.chStop
	for _, p := range s.getPlayerList() {
		p.Stop()
	}
	log.Printf("server stop: %s\n", msg.Error())
}

//ListenTCP only call func use go routine
func (s *Server) ListenTCP(laddr string) {
	l, err := net.Listen("tcp", laddr)
	if err != nil {
		s.chStop <- err
	}
	defer l.Close()
	log.Printf("listen at %s\n", l.Addr().String())
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Println(err)
			continue
		}
		s.chConn <- conn
	}
}

//DelPlayer ...
func (s *Server) DelPlayer(key uint64) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	delete(s.players, key)
}

//RegisterHandle ...
func (s *Server) RegisterHandle(id protocol.C2SCmd, f func(*Player, []byte)) {
	nID := int32(id)
	if _, ok := s.handles[nID]; ok {
		log.Printf("protocol(%d) handle repeat\n", nID)
		return
	}
	s.handles[nID] = f
	log.Printf("register handle protocol(%d)\n", nID)
}

//HandleSignal ...
func (s *Server) HandleSignal() {
	signal.Notify(s.chSig, os.Interrupt)
	sig := <-s.chSig
	s.chStop <- fmt.Errorf(sig.String())
}

//NewServer instance
func NewServer() *Server {
	s := &Server{
		index:   0,
		players: make(map[uint64]*Player),
		handles: make(map[int32]func(*Player, []byte)),
		chStop:  make(chan error),
		chConn:  make(chan net.Conn),
		chSig:   make(chan os.Signal),
		mutex:   &sync.RWMutex{},
	}
	s.RegisterHandle(protocol.C2SCmd_Abnormal, func(p *Player, msg []byte) {
		p.Stop()
	})
	s.RegisterHandle(protocol.C2SCmd_Chat, func(p *Player, msg []byte) {
		var chatMsg protocol.C2SChat
		if err := proto.Unmarshal(msg, &chatMsg); err != nil {
			p.Stop()
		}
		player := p.GetTargetPlayer(chatMsg.Index)
		if player != nil {
			player.SendChat(chatMsg.Context)
		}
	})
	return s
}

func main() {
	app := NewServer()
	go app.HandleSignal()
	go app.ListenTCP(":7788")
	app.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"github.com/coreos/etcd/store"
	"github.com/coreos/etcd/web"
	"github.com/coreos/go-raft"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

//------------------------------------------------------------------------------
//
// Initialization
//
//------------------------------------------------------------------------------

var verbose bool

var cluster string

var address string
var clientPort int
var serverPort int
var webPort int

var serverCertFile string
var serverKeyFile string
var serverCAFile string

var clientCertFile string
var clientKeyFile string
var clientCAFile string

var dirPath string

var ignore bool

var maxSize int

func init() {
	flag.BoolVar(&verbose, "v", false, "verbose logging")

	flag.StringVar(&cluster, "C", "", "the ip address and port of a existing cluster")

	flag.StringVar(&address, "a", "0.0.0.0", "the ip address of the local machine")
	flag.IntVar(&clientPort, "c", 4001, "the port to communicate with clients")
	flag.IntVar(&serverPort, "s", 7001, "the port to communicate with servers")
	flag.IntVar(&webPort, "w", -1, "the port of web interface")

	flag.StringVar(&serverCAFile, "serverCAFile", "", "the path of the CAFile")
	flag.StringVar(&serverCertFile, "serverCert", "", "the cert file of the server")
	flag.StringVar(&serverKeyFile, "serverKey", "", "the key file of the server")

	flag.StringVar(&clientCAFile, "clientCAFile", "", "the path of the client CAFile")
	flag.StringVar(&clientCertFile, "clientCert", "", "the cert file of the client")
	flag.StringVar(&clientKeyFile, "clientKey", "", "the key file of the client")

	flag.StringVar(&dirPath, "d", "/tmp/", "the directory to store log and snapshot")

	flag.BoolVar(&ignore, "i", false, "ignore the old configuration, create a new node")

	flag.IntVar(&maxSize, "m", 1024, "the max size of result buffer")
}

// CONSTANTS
const (
	HTTP = iota
	HTTPS
	HTTPSANDVERIFY
)

const (
	SERVER = iota
	CLIENT
)

const (
	ELECTIONTIMTOUT  = 200 * time.Millisecond
	HEARTBEATTIMEOUT = 50 * time.Millisecond
)

//------------------------------------------------------------------------------
//
// Typedefs
//
//------------------------------------------------------------------------------

type Info struct {
	Address    string `json:"address"`
	ServerPort int    `json:"serverPort"`
	ClientPort int    `json:"clientPort"`
	WebPort    int    `json:"webPort"`

	ServerCertFile string `json:"serverCertFile"`
	ServerKeyFile  string `json:"serverKeyFile"`
	ServerCAFile   string `json:"serverCAFile"`

	ClientCertFile string `json:"clientCertFile"`
	ClientKeyFile  string `json:"clientKeyFile"`
	ClientCAFile   string `json:"clientCAFile"`
}

//------------------------------------------------------------------------------
//
// Variables
//
//------------------------------------------------------------------------------

var raftServer *raft.Server
var raftTransporter transporter
var etcdStore *store.Store
var info *Info

//------------------------------------------------------------------------------
//
// Functions
//
//------------------------------------------------------------------------------

//--------------------------------------
// Main
//--------------------------------------

func main() {
	flag.Parse()

	// Setup commands.
	registerCommands()

	// Read server info from file or grab it from user.
	if err := os.MkdirAll(dirPath, 0744); err != nil {
		fatal("Unable to create path: %v", err)
	}

	info = getInfo(dirPath)

	// secrity type
	st := securityType(SERVER)

	clientSt := securityType(CLIENT)

	if st == -1 || clientSt == -1 {
		fatal("Please specify cert and key file or cert and key file and CAFile or none of the three")
	}

	// Create etcd key-value store
	etcdStore = store.CreateStore(maxSize)

	startRaft(st)

	if webPort != -1 {
		// start web
		etcdStore.SetMessager(&storeMsg)
		go webHelper()
		go web.Start(raftServer, webPort)
	}

	startClientTransport(info.ClientPort, clientSt)

}

// Start the raft server
func startRaft(securityType int) {
	var err error

	raftName := fmt.Sprintf("%s:%d", info.Address, info.ServerPort)

	// Create transporter for raft
	raftTransporter = createTransporter(securityType)

	// Create raft server
	raftServer, err = raft.NewServer(raftName, dirPath, raftTransporter, etcdStore, nil)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// LoadSnapshot
	// err = raftServer.LoadSnapshot()

	// if err == nil {
	// 	debug("%s finished load snapshot", raftServer.Name())
	// } else {
	// 	debug(err)
	// }

	raftServer.Initialize()
	raftServer.SetElectionTimeout(ELECTIONTIMTOUT)
	raftServer.SetHeartbeatTimeout(HEARTBEATTIMEOUT)

	if raftServer.IsLogEmpty() {

		// start as a leader in a new cluster
		if cluster == "" {
			raftServer.StartLeader()

			time.Sleep(time.Millisecond * 20)

			// leader need to join self as a peer
			for {
				command := &JoinCommand{}
				command.Name = raftServer.Name()
				_, err := raftServer.Do(command)
				if err == nil {
					break
				}
			}
			debug("%s start as a leader", raftServer.Name())

			// start as a follower in a existing cluster
		} else {
			raftServer.StartFollower()

			err := joinCluster(raftServer, cluster)
			if err != nil {
				fatal(fmt.Sprintln(err))
			}
			debug("%s success join to the cluster", raftServer.Name())
		}

	} else {
		// rejoin the previous cluster
		raftServer.StartFollower()
		debug("%s restart as a follower", raftServer.Name())
	}

	// open the snapshot
	// go server.Snapshot()

	// start to response to raft requests
	go startRaftTransport(info.ServerPort, securityType)

}

// Create transporter using by raft server
// Create http or https transporter based on
// wether the user give the server cert and key
func createTransporter(st int) transporter {
	t := transporter{}

	switch st {
	case HTTP:
		t.client = nil
		return t

	case HTTPS:
		fallthrough
	case HTTPSANDVERIFY:
		tlsCert, err := tls.LoadX509KeyPair(serverCertFile, serverKeyFile)

		if err != nil {
			fatal(fmt.Sprintln(err))
		}

		tr := &http.Transport{
			TLSClientConfig: &tls.Config{
				Certificates:       []tls.Certificate{tlsCert},
				InsecureSkipVerify: true,
			},
			DisableCompression: true,
		}

		t.client = &http.Client{Transport: tr}
		return t
	}

	// for complier
	return transporter{}
}

// Start to listen and response raft command
func startRaftTransport(port int, st int) {

	// internal commands
	http.HandleFunc("/join", JoinHttpHandler)
	http.HandleFunc("/vote", VoteHttpHandler)
	http.HandleFunc("/log", GetLogHttpHandler)
	http.HandleFunc("/log/append", AppendEntriesHttpHandler)
	http.HandleFunc("/snapshot", SnapshotHttpHandler)
	http.HandleFunc("/client", ClientHttpHandler)

	switch st {

	case HTTP:
		fmt.Printf("raft server [%s] listen on http port %v\n", address, port)
		log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))

	case HTTPS:
		fmt.Printf("raft server [%s] listen on https port %v\n", address, port)
		log.Fatal(http.ListenAndServeTLS(fmt.Sprintf(":%d", port), serverCertFile, serverKeyFile, nil))

	case HTTPSANDVERIFY:

		server := &http.Server{
			TLSConfig: &tls.Config{
				ClientAuth: tls.RequireAndVerifyClientCert,
				ClientCAs:  createCertPool(serverCAFile),
			},
			Addr: fmt.Sprintf(":%d", port),
		}
		fmt.Printf("raft server [%s] listen on https port %v\n", address, port)
		err := server.ListenAndServeTLS(serverCertFile, serverKeyFile)

		if err != nil {
			log.Fatal(err)
		}
	}

}

// Start to listen and response client command
func startClientTransport(port int, st int) {
	// external commands
	http.HandleFunc("/"+version+"/keys/", Multiplexer)
	http.HandleFunc("/"+version+"/watch/", WatchHttpHandler)
	http.HandleFunc("/"+version+"/list/", ListHttpHandler)
	http.HandleFunc("/"+version+"/testAndSet/", TestAndSetHttpHandler)
	http.HandleFunc("/leader", LeaderHttpHandler)

	switch st {

	case HTTP:
		fmt.Printf("etcd [%s] listen on http port %v\n", address, clientPort)
		log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))

	case HTTPS:
		fmt.Printf("etcd [%s] listen on https port %v\n", address, clientPort)
		http.ListenAndServeTLS(fmt.Sprintf(":%d", port), clientCertFile, clientKeyFile, nil)

	case HTTPSANDVERIFY:

		server := &http.Server{
			TLSConfig: &tls.Config{
				ClientAuth: tls.RequireAndVerifyClientCert,
				ClientCAs:  createCertPool(clientCAFile),
			},
			Addr: fmt.Sprintf(":%d", port),
		}
		fmt.Printf("etcd [%s] listen on https port %v\n", address, clientPort)
		err := server.ListenAndServeTLS(clientCertFile, clientKeyFile)

		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	}
}

//--------------------------------------
// Config
//--------------------------------------

// Get the security type
func securityType(source int) int {

	var keyFile, certFile, CAFile string

	switch source {

	case SERVER:
		keyFile = info.ServerKeyFile
		certFile = info.ServerCertFile
		CAFile = info.ServerCAFile

	case CLIENT:
		keyFile = info.ClientKeyFile
		certFile = info.ClientCertFile
		CAFile = info.ClientCAFile
	}

	// If the user do not specify key file, cert file and
	// CA file, the type will be HTTP
	if keyFile == "" && certFile == "" && CAFile == "" {

		return HTTP

	}

	if keyFile != "" && certFile != "" {
		if CAFile != "" {
			// If the user specify all the three file, the type
			// will be HTTPS with client cert auth
			return HTTPSANDVERIFY
		}
		// If the user specify key file and cert file but not
		// CA file, the type will be HTTPS without client cert
		// auth
		return HTTPS
	}

	// bad specification
	return -1
}

// Get the server info from previous conf file
// or from the user
func getInfo(path string) *Info {
	info := &Info{}

	// Read in the server info if available.
	infoPath := fmt.Sprintf("%s/info", path)

	// Delete the old configuration if exist
	if ignore {
		logPath := fmt.Sprintf("%s/log", path)
		snapshotPath := fmt.Sprintf("%s/snapshotPath", path)
		os.Remove(infoPath)
		os.Remove(logPath)
		os.RemoveAll(snapshotPath)

	}

	if file, err := os.Open(infoPath); err == nil {
		if content, err := ioutil.ReadAll(file); err != nil {
			fatal("Unable to read info: %v", err)
		} else {
			if err = json.Unmarshal(content, &info); err != nil {
				fatal("Unable to parse info: %v", err)
			}
		}
		file.Close()

	} else {
		// Otherwise ask user for info and write it to file.

		if address == "" {
			fatal("Please give the address of the local machine")
		}

		info.Address = address
		info.Address = strings.TrimSpace(info.Address)
		fmt.Println("address ", info.Address)

		info.ServerPort = serverPort
		info.ClientPort = clientPort
		info.WebPort = webPort

		info.ClientCAFile = clientCAFile
		info.ClientCertFile = clientCertFile
		info.ClientKeyFile = clientKeyFile

		info.ServerCAFile = serverCAFile
		info.ServerKeyFile = serverKeyFile
		info.ServerCertFile = serverCertFile

		// Write to file.
		content, _ := json.Marshal(info)
		content = []byte(string(content) + "\n")
		if err := ioutil.WriteFile(infoPath, content, 0644); err != nil {
			fatal("Unable to write info to file: %v", err)
		}
	}

	return info
}

// Create client auth certpool
func createCertPool(CAFile string) *x509.CertPool {
	pemByte, _ := ioutil.ReadFile(CAFile)

	block, pemByte := pem.Decode(pemByte)

	cert, err := x509.ParseCertificate(block.Bytes)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	certPool := x509.NewCertPool()

	certPool.AddCert(cert)

	return certPool
}

// Send join requests to the leader.
func joinCluster(s *raft.Server, serverName string) error {
	var b bytes.Buffer

	command := &JoinCommand{}
	command.Name = s.Name()

	json.NewEncoder(&b).Encode(command)

	// t must be ok
	t, _ := raftServer.Transporter().(transporter)
	debug("Send Join Request to %s", serverName)
	resp, err := t.Post(fmt.Sprintf("%s/join", serverName), &b)

	for {
		if resp != nil {
			defer resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return nil
			}
			if resp.StatusCode == http.StatusServiceUnavailable {
				address, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					warn("Cannot Read Leader info: %v", err)
				}
				debug("Leader is %s", address)
				debug("Send Join Request to %s", address)
				json.NewEncoder(&b).Encode(command)
				resp, err = t.Post(fmt.Sprintf("%s/join", address), &b)
			}
		}
	}
	return fmt.Errorf("Unable to join: %v", err)
}

// Register commands to raft server
func registerCommands() {
	raft.RegisterCommand(&JoinCommand{})
	raft.RegisterCommand(&SetCommand{})
	raft.RegisterCommand(&GetCommand{})
	raft.RegisterCommand(&DeleteCommand{})
	raft.RegisterCommand(&WatchCommand{})
	raft.RegisterCommand(&ListCommand{})
	raft.RegisterCommand(&TestAndSetCommand{})
}
<commit_msg>return error when cannot join the cluster<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"github.com/coreos/etcd/store"
	"github.com/coreos/etcd/web"
	"github.com/coreos/go-raft"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

//------------------------------------------------------------------------------
//
// Initialization
//
//------------------------------------------------------------------------------

var verbose bool

var cluster string

var address string
var clientPort int
var serverPort int
var webPort int

var serverCertFile string
var serverKeyFile string
var serverCAFile string

var clientCertFile string
var clientKeyFile string
var clientCAFile string

var dirPath string

var ignore bool

var maxSize int

func init() {
	flag.BoolVar(&verbose, "v", false, "verbose logging")

	flag.StringVar(&cluster, "C", "", "the ip address and port of a existing cluster")

	flag.StringVar(&address, "a", "0.0.0.0", "the ip address of the local machine")
	flag.IntVar(&clientPort, "c", 4001, "the port to communicate with clients")
	flag.IntVar(&serverPort, "s", 7001, "the port to communicate with servers")
	flag.IntVar(&webPort, "w", -1, "the port of web interface")

	flag.StringVar(&serverCAFile, "serverCAFile", "", "the path of the CAFile")
	flag.StringVar(&serverCertFile, "serverCert", "", "the cert file of the server")
	flag.StringVar(&serverKeyFile, "serverKey", "", "the key file of the server")

	flag.StringVar(&clientCAFile, "clientCAFile", "", "the path of the client CAFile")
	flag.StringVar(&clientCertFile, "clientCert", "", "the cert file of the client")
	flag.StringVar(&clientKeyFile, "clientKey", "", "the key file of the client")

	flag.StringVar(&dirPath, "d", "/tmp/", "the directory to store log and snapshot")

	flag.BoolVar(&ignore, "i", false, "ignore the old configuration, create a new node")

	flag.IntVar(&maxSize, "m", 1024, "the max size of result buffer")
}

// CONSTANTS
const (
	HTTP = iota
	HTTPS
	HTTPSANDVERIFY
)

const (
	SERVER = iota
	CLIENT
)

const (
	ELECTIONTIMTOUT  = 200 * time.Millisecond
	HEARTBEATTIMEOUT = 50 * time.Millisecond
)

//------------------------------------------------------------------------------
//
// Typedefs
//
//------------------------------------------------------------------------------

type Info struct {
	Address    string `json:"address"`
	ServerPort int    `json:"serverPort"`
	ClientPort int    `json:"clientPort"`
	WebPort    int    `json:"webPort"`

	ServerCertFile string `json:"serverCertFile"`
	ServerKeyFile  string `json:"serverKeyFile"`
	ServerCAFile   string `json:"serverCAFile"`

	ClientCertFile string `json:"clientCertFile"`
	ClientKeyFile  string `json:"clientKeyFile"`
	ClientCAFile   string `json:"clientCAFile"`
}

//------------------------------------------------------------------------------
//
// Variables
//
//------------------------------------------------------------------------------

var raftServer *raft.Server
var raftTransporter transporter
var etcdStore *store.Store
var info *Info

//------------------------------------------------------------------------------
//
// Functions
//
//------------------------------------------------------------------------------

//--------------------------------------
// Main
//--------------------------------------

func main() {
	flag.Parse()

	// Setup commands.
	registerCommands()

	// Read server info from file or grab it from user.
	if err := os.MkdirAll(dirPath, 0744); err != nil {
		fatal("Unable to create path: %v", err)
	}

	info = getInfo(dirPath)

	// secrity type
	st := securityType(SERVER)

	clientSt := securityType(CLIENT)

	if st == -1 || clientSt == -1 {
		fatal("Please specify cert and key file or cert and key file and CAFile or none of the three")
	}

	// Create etcd key-value store
	etcdStore = store.CreateStore(maxSize)

	startRaft(st)

	if webPort != -1 {
		// start web
		etcdStore.SetMessager(&storeMsg)
		go webHelper()
		go web.Start(raftServer, webPort)
	}

	startClientTransport(info.ClientPort, clientSt)

}

// Start the raft server
func startRaft(securityType int) {
	var err error

	raftName := fmt.Sprintf("%s:%d", info.Address, info.ServerPort)

	// Create transporter for raft
	raftTransporter = createTransporter(securityType)

	// Create raft server
	raftServer, err = raft.NewServer(raftName, dirPath, raftTransporter, etcdStore, nil)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// LoadSnapshot
	// err = raftServer.LoadSnapshot()

	// if err == nil {
	// 	debug("%s finished load snapshot", raftServer.Name())
	// } else {
	// 	debug(err)
	// }

	raftServer.Initialize()
	raftServer.SetElectionTimeout(ELECTIONTIMTOUT)
	raftServer.SetHeartbeatTimeout(HEARTBEATTIMEOUT)

	if raftServer.IsLogEmpty() {

		// start as a leader in a new cluster
		if cluster == "" {
			raftServer.StartLeader()

			time.Sleep(time.Millisecond * 20)

			// leader need to join self as a peer
			for {
				command := &JoinCommand{}
				command.Name = raftServer.Name()
				_, err := raftServer.Do(command)
				if err == nil {
					break
				}
			}
			debug("%s start as a leader", raftServer.Name())

			// start as a follower in a existing cluster
		} else {
			raftServer.StartFollower()

			err := joinCluster(raftServer, cluster)
			if err != nil {
				fatal(fmt.Sprintln(err))
			}
			debug("%s success join to the cluster", raftServer.Name())
		}

	} else {
		// rejoin the previous cluster
		raftServer.StartFollower()
		debug("%s restart as a follower", raftServer.Name())
	}

	// open the snapshot
	// go server.Snapshot()

	// start to response to raft requests
	go startRaftTransport(info.ServerPort, securityType)

}

// Create transporter using by raft server
// Create http or https transporter based on
// wether the user give the server cert and key
func createTransporter(st int) transporter {
	t := transporter{}

	switch st {
	case HTTP:
		t.client = nil
		return t

	case HTTPS:
		fallthrough
	case HTTPSANDVERIFY:
		tlsCert, err := tls.LoadX509KeyPair(serverCertFile, serverKeyFile)

		if err != nil {
			fatal(fmt.Sprintln(err))
		}

		tr := &http.Transport{
			TLSClientConfig: &tls.Config{
				Certificates:       []tls.Certificate{tlsCert},
				InsecureSkipVerify: true,
			},
			DisableCompression: true,
		}

		t.client = &http.Client{Transport: tr}
		return t
	}

	// for complier
	return transporter{}
}

// Start to listen and response raft command
func startRaftTransport(port int, st int) {

	// internal commands
	http.HandleFunc("/join", JoinHttpHandler)
	http.HandleFunc("/vote", VoteHttpHandler)
	http.HandleFunc("/log", GetLogHttpHandler)
	http.HandleFunc("/log/append", AppendEntriesHttpHandler)
	http.HandleFunc("/snapshot", SnapshotHttpHandler)
	http.HandleFunc("/client", ClientHttpHandler)

	switch st {

	case HTTP:
		fmt.Printf("raft server [%s] listen on http port %v\n", address, port)
		log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))

	case HTTPS:
		fmt.Printf("raft server [%s] listen on https port %v\n", address, port)
		log.Fatal(http.ListenAndServeTLS(fmt.Sprintf(":%d", port), serverCertFile, serverKeyFile, nil))

	case HTTPSANDVERIFY:

		server := &http.Server{
			TLSConfig: &tls.Config{
				ClientAuth: tls.RequireAndVerifyClientCert,
				ClientCAs:  createCertPool(serverCAFile),
			},
			Addr: fmt.Sprintf(":%d", port),
		}
		fmt.Printf("raft server [%s] listen on https port %v\n", address, port)
		err := server.ListenAndServeTLS(serverCertFile, serverKeyFile)

		if err != nil {
			log.Fatal(err)
		}
	}

}

// Start to listen and response client command
func startClientTransport(port int, st int) {
	// external commands
	http.HandleFunc("/"+version+"/keys/", Multiplexer)
	http.HandleFunc("/"+version+"/watch/", WatchHttpHandler)
	http.HandleFunc("/"+version+"/list/", ListHttpHandler)
	http.HandleFunc("/"+version+"/testAndSet/", TestAndSetHttpHandler)
	http.HandleFunc("/leader", LeaderHttpHandler)

	switch st {

	case HTTP:
		fmt.Printf("etcd [%s] listen on http port %v\n", address, clientPort)
		log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))

	case HTTPS:
		fmt.Printf("etcd [%s] listen on https port %v\n", address, clientPort)
		http.ListenAndServeTLS(fmt.Sprintf(":%d", port), clientCertFile, clientKeyFile, nil)

	case HTTPSANDVERIFY:

		server := &http.Server{
			TLSConfig: &tls.Config{
				ClientAuth: tls.RequireAndVerifyClientCert,
				ClientCAs:  createCertPool(clientCAFile),
			},
			Addr: fmt.Sprintf(":%d", port),
		}
		fmt.Printf("etcd [%s] listen on https port %v\n", address, clientPort)
		err := server.ListenAndServeTLS(clientCertFile, clientKeyFile)

		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	}
}

//--------------------------------------
// Config
//--------------------------------------

// Get the security type
func securityType(source int) int {

	var keyFile, certFile, CAFile string

	switch source {

	case SERVER:
		keyFile = info.ServerKeyFile
		certFile = info.ServerCertFile
		CAFile = info.ServerCAFile

	case CLIENT:
		keyFile = info.ClientKeyFile
		certFile = info.ClientCertFile
		CAFile = info.ClientCAFile
	}

	// If the user do not specify key file, cert file and
	// CA file, the type will be HTTP
	if keyFile == "" && certFile == "" && CAFile == "" {

		return HTTP

	}

	if keyFile != "" && certFile != "" {
		if CAFile != "" {
			// If the user specify all the three file, the type
			// will be HTTPS with client cert auth
			return HTTPSANDVERIFY
		}
		// If the user specify key file and cert file but not
		// CA file, the type will be HTTPS without client cert
		// auth
		return HTTPS
	}

	// bad specification
	return -1
}

// Get the server info from previous conf file
// or from the user
func getInfo(path string) *Info {
	info := &Info{}

	// Read in the server info if available.
	infoPath := fmt.Sprintf("%s/info", path)

	// Delete the old configuration if exist
	if ignore {
		logPath := fmt.Sprintf("%s/log", path)
		snapshotPath := fmt.Sprintf("%s/snapshotPath", path)
		os.Remove(infoPath)
		os.Remove(logPath)
		os.RemoveAll(snapshotPath)

	}

	if file, err := os.Open(infoPath); err == nil {
		if content, err := ioutil.ReadAll(file); err != nil {
			fatal("Unable to read info: %v", err)
		} else {
			if err = json.Unmarshal(content, &info); err != nil {
				fatal("Unable to parse info: %v", err)
			}
		}
		file.Close()

	} else {
		// Otherwise ask user for info and write it to file.

		if address == "" {
			fatal("Please give the address of the local machine")
		}

		info.Address = address
		info.Address = strings.TrimSpace(info.Address)
		fmt.Println("address ", info.Address)

		info.ServerPort = serverPort
		info.ClientPort = clientPort
		info.WebPort = webPort

		info.ClientCAFile = clientCAFile
		info.ClientCertFile = clientCertFile
		info.ClientKeyFile = clientKeyFile

		info.ServerCAFile = serverCAFile
		info.ServerKeyFile = serverKeyFile
		info.ServerCertFile = serverCertFile

		// Write to file.
		content, _ := json.Marshal(info)
		content = []byte(string(content) + "\n")
		if err := ioutil.WriteFile(infoPath, content, 0644); err != nil {
			fatal("Unable to write info to file: %v", err)
		}
	}

	return info
}

// Create client auth certpool
func createCertPool(CAFile string) *x509.CertPool {
	pemByte, _ := ioutil.ReadFile(CAFile)

	block, pemByte := pem.Decode(pemByte)

	cert, err := x509.ParseCertificate(block.Bytes)

	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	certPool := x509.NewCertPool()

	certPool.AddCert(cert)

	return certPool
}

// Send join requests to the leader.
func joinCluster(s *raft.Server, serverName string) error {
	var b bytes.Buffer

	command := &JoinCommand{}
	command.Name = s.Name()

	json.NewEncoder(&b).Encode(command)

	// t must be ok
	t, _ := raftServer.Transporter().(transporter)
	debug("Send Join Request to %s", serverName)
	resp, err := t.Post(fmt.Sprintf("%s/join", serverName), &b)

	for {
		if err != nil {
			return fmt.Errorf("Unable to join: %v", err)
		}
		if resp != nil {
			defer resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return nil
			}
			if resp.StatusCode == http.StatusServiceUnavailable {
				address, err := ioutil.ReadAll(resp.Body)
				if err != nil {
					warn("Cannot Read Leader info: %v", err)
				}
				debug("Leader is %s", address)
				debug("Send Join Request to %s", address)
				json.NewEncoder(&b).Encode(command)
				resp, err = t.Post(fmt.Sprintf("%s/join", address), &b)
			}
		}
	}
	return fmt.Errorf("Unable to join: %v", err)
}

// Register commands to raft server
func registerCommands() {
	raft.RegisterCommand(&JoinCommand{})
	raft.RegisterCommand(&SetCommand{})
	raft.RegisterCommand(&GetCommand{})
	raft.RegisterCommand(&DeleteCommand{})
	raft.RegisterCommand(&WatchCommand{})
	raft.RegisterCommand(&ListCommand{})
	raft.RegisterCommand(&TestAndSetCommand{})
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 volume

import (
	"fmt"

	"github.com/golang/glog"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	"k8s.io/api/core/v1"
)

func (p *flexProvisioner) Delete(volume *v1.PersistentVolume) error {
	glog.Infof("Delete called for volume:", volume.Name)

	provisioned, err := p.provisioned(volume)
	if err != nil {
		return fmt.Errorf("error determining if this provisioner was the one to provision volume %q: %v", volume.Name, err)
	}
	if !provisioned {
		strerr := fmt.Sprintf("this provisioner id %s didn't provision volume %q and so can't delete it; id %s did & can", p.identity, volume.Name, volume.Annotations[annProvisionerID])
		return &controller.IgnoredError{Reason: strerr}
	}

	call := p.NewDriverCall(p.execCommand, deleteCmd)
	call.AppendSpec(volume.Spec.FlexVolume.Options, nil)
	output, err := call.Run()

	if err != nil {
		glog.Errorf("Failed to delete volume %s, output: %s, error: %s", volume, output.Message, err.Error())
		return err
	}
	return nil
}

func (p *flexProvisioner) provisioned(volume *v1.PersistentVolume) (bool, error) {
	provisionerID, ok := volume.Annotations[annProvisionerID]
	if !ok {
		return false, fmt.Errorf("PV doesn't have an annotation %s", annProvisionerID)
	}

	return provisionerID == string(p.identity), nil
}
<commit_msg>flex: Add PV name to delete<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 volume

import (
	"fmt"

	"github.com/golang/glog"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	"k8s.io/api/core/v1"
)

func (p *flexProvisioner) Delete(volume *v1.PersistentVolume) error {
	glog.Infof("Delete called for volume:", volume.Name)

	provisioned, err := p.provisioned(volume)
	if err != nil {
		return fmt.Errorf("error determining if this provisioner was the one to provision volume %q: %v", volume.Name, err)
	}
	if !provisioned {
		strerr := fmt.Sprintf("this provisioner id %s didn't provision volume %q and so can't delete it; id %s did & can", p.identity, volume.Name, volume.Annotations[annProvisionerID])
		return &controller.IgnoredError{Reason: strerr}
	}

	extraOptions := map[string]string{}
	extraOptions[optionPVorVolumeName] = volume.Name

	call := p.NewDriverCall(p.execCommand, deleteCmd)
	call.AppendSpec(volume.Spec.FlexVolume.Options, extraOptions)
	output, err := call.Run()

	if err != nil {
		glog.Errorf("Failed to delete volume %s, output: %s, error: %s", volume, output.Message, err.Error())
		return err
	}
	return nil
}

func (p *flexProvisioner) provisioned(volume *v1.PersistentVolume) (bool, error) {
	provisionerID, ok := volume.Annotations[annProvisionerID]
	if !ok {
		return false, fmt.Errorf("PV doesn't have an annotation %s", annProvisionerID)
	}

	return provisionerID == string(p.identity), nil
}
<|endoftext|>


<commit_before>package gannoy

import (
	"bytes"
	"encoding/binary"
	"io"
	"os"
	"syscall"
)

type File struct {
	tree       int
	dim        int
	K          int
	file       *os.File
	filename   string
	appendFile *os.File
	createChan chan createArgs
	locker     Locker
	nodeSize   int64
}

func newFile(filename string, tree, dim, K int) *File {
	_, err := os.Stat(filename)
	if err != nil {
		f, _ := os.Create(filename)
		f.Close()
	}

	file, _ := os.OpenFile(filename, os.O_RDWR, 0)
	appendFile, _ := os.OpenFile(filename, os.O_RDWR|os.O_APPEND, 0)

	f := &File{
		tree:       tree,
		dim:        dim,
		K:          K,
		file:       file,
		filename:   filename,
		appendFile: appendFile,
		createChan: make(chan createArgs, 1),
		locker:     newLocker(),
		nodeSize: int64(1 + // free
			4 + // nDescendants
			4 + // key
			4*tree + // parents
			4*2 + // children
			8*dim), // v
	}
	go f.creator()
	return f
}

func (f *File) Create(n Node) (int, error) {
	args := createArgs{node: n, result: make(chan createResult)}
	f.createChan <- args
	result := <-args.result
	return result.id, result.err
}

func (f *File) create(n Node) (int, error) {
	buf := &bytes.Buffer{}
	f.nodeToBuf(buf, n)
	id := f.nodeCount()
	_, err := f.appendFile.Write(buf.Bytes())
	return id, err
}

func (f *File) Find(id int) (Node, error) {
	node := Node{}
	node.id = id
	node.storage = f
	offset := f.offset(id)
	err := f.locker.ReadLock(f.file.Fd(), offset, f.nodeSize)
	if err != nil {
		return node, err
	}
	defer f.locker.UnLock(f.file.Fd(), offset, f.nodeSize)

	b := make([]byte, f.nodeSize)
	_, err = syscall.Pread(int(f.file.Fd()), b, offset)
	if err != nil {
		return node, err
	}

	buf := bytes.NewReader(b)

	var free bool
	binary.Read(buf, binary.BigEndian, &free)
	node.free = free

	var nDescendants int32
	binary.Read(buf, binary.BigEndian, &nDescendants)
	node.nDescendants = int(nDescendants)

	var key int32
	binary.Read(buf, binary.BigEndian, &key)
	node.key = int(key)

	parents := make([]int32, f.tree)
	binary.Read(buf, binary.BigEndian, &parents)
	nodeParents := make([]int, f.tree)
	for i, parent := range parents {
		nodeParents[i] = int(parent)
	}
	node.parents = nodeParents

	if node.nDescendants == 1 {
		// leaf node
		buf.Seek(int64(4*2), io.SeekCurrent) // skip children
		node.children = []int{0, 0}

		vec := make([]float64, f.dim)
		binary.Read(buf, binary.BigEndian, &vec)
		node.v = vec
	} else if node.nDescendants <= f.K {
		// bucket node
		children := make([]int32, nDescendants)
		binary.Read(buf, binary.BigEndian, &children)
		nodeChildren := make([]int, nDescendants)
		for i, child := range children {
			nodeChildren[i] = int(child)
		}
		node.children = nodeChildren
	} else {
		// other node
		children := make([]int32, 2)
		binary.Read(buf, binary.BigEndian, &children)
		nodeChildren := make([]int, 2)
		for i, child := range children {
			nodeChildren[i] = int(child)
		}
		node.children = nodeChildren

		vec := make([]float64, f.dim)
		binary.Read(buf, binary.BigEndian, &vec)
		node.v = vec
	}
	return node, nil
}

func (f *File) Update(n Node) error {
	buf := &bytes.Buffer{}
	f.nodeToBuf(buf, n)
	offset := f.offset(n.id)
	file, _ := os.OpenFile(f.filename, os.O_RDWR, 0)
	defer file.Close()

	err := f.locker.WriteLock(file.Fd(), offset, f.nodeSize)
	if err != nil {
		return err
	}
	defer f.locker.UnLock(file.Fd(), offset, f.nodeSize)

	_, err = syscall.Pwrite(int(file.Fd()), buf.Bytes(), offset)
	return err
}

func (f *File) UpdateParent(id, rootIndex, parent int) error {
	offset := f.offset(id) +
		int64(1+ // free
			4+ // nDescendants
			4+ // key
			4*rootIndex) // parents
	buf := &bytes.Buffer{}
	binary.Write(buf, binary.BigEndian, int32(parent))

	file, _ := os.OpenFile(f.filename, os.O_RDWR, 0)
	defer file.Close()

	err := f.locker.WriteLock(file.Fd(), offset, 4)
	if err != nil {
		return err
	}
	defer f.locker.UnLock(file.Fd(), offset, 4)

	_, err = syscall.Pwrite(int(file.Fd()), buf.Bytes(), offset)
	return err
}

func (f *File) Delete(n Node) error {
	n.free = true
	return f.Update(n)
}

func (f *File) Iterate(c chan Node) {
	count := f.nodeCount()
	// TODO: Use goroutine
	for i := 0; i < count; i++ {
		n, err := f.Find(i)
		if err != nil {
			close(c)
		}
		c <- n
	}
	close(c)
}

func (f File) offset(id int) int64 {
	return (int64(id) * f.nodeSize)
}

func (f File) nodeCount() int {
	stat, _ := f.file.Stat()
	size := stat.Size()
	return int(size / f.nodeSize)
}

func (f File) nodeToBuf(buf *bytes.Buffer, node Node) {
	// 1bytes free
	binary.Write(buf, binary.BigEndian, node.free)

	// 4bytes nDescendants
	binary.Write(buf, binary.BigEndian, int32(node.nDescendants))

	// 4bytes key
	binary.Write(buf, binary.BigEndian, int32(node.key))

	// 4bytes parents
	parents := make([]int32, len(node.parents))
	for i, parent := range node.parents {
		parents[i] = int32(parent)
	}
	binary.Write(buf, binary.BigEndian, parents)

	if node.isBucket() {
		// 4bytes children in K
		children := make([]int32, f.K)
		for i, child := range node.children {
			children[i] = int32(child)
		}
		binary.Write(buf, binary.BigEndian, children)

		// padding by zero
		remainingSize := ((2*4 + 8*f.dim) - (4 * f.K))
		binary.Write(buf, binary.BigEndian, make([]int32, remainingSize/4))
	} else {
		// 4bytes children in K
		children := make([]int32, 2)
		for i, child := range node.children {
			children[i] = int32(child)
		}
		binary.Write(buf, binary.BigEndian, children)

		// 8bytes v in f
		vec := make([]float64, f.dim)
		for i, v := range node.v {
			vec[i] = float64(v)
		}
		binary.Write(buf, binary.BigEndian, vec)
	}
}

type createArgs struct {
	node   Node
	result chan createResult
}

type createResult struct {
	id  int
	err error
}

func (f *File) creator() {
	for args := range f.createChan {
		id, err := f.create(args.node)
		args.result <- createResult{
			id:  id,
			err: err,
		}
	}
}

func (f File) size() int64 {
	info, _ := f.file.Stat()
	return info.Size()
}
<commit_msg>Since binary.Read is slow using reflection for []float64, gannoy uses a dedicated method<commit_after>package gannoy

import (
	"bytes"
	"encoding/binary"
	"math"
	"os"
	"syscall"
)

type File struct {
	tree       int
	dim        int
	K          int
	file       *os.File
	filename   string
	appendFile *os.File
	createChan chan createArgs
	locker     Locker
	nodeSize   int64
	offsetOfV  int64
}

func newFile(filename string, tree, dim, K int) *File {
	_, err := os.Stat(filename)
	if err != nil {
		f, _ := os.Create(filename)
		f.Close()
	}

	file, _ := os.OpenFile(filename, os.O_RDWR, 0)
	appendFile, _ := os.OpenFile(filename, os.O_RDWR|os.O_APPEND, 0)

	f := &File{
		tree:       tree,
		dim:        dim,
		K:          K,
		file:       file,
		filename:   filename,
		appendFile: appendFile,
		createChan: make(chan createArgs, 1),
		locker:     newLocker(),
		nodeSize: int64(1 + // free
			4 + // nDescendants
			4 + // key
			4*tree + // parents
			4*2 + // children
			8*dim), // v
		offsetOfV: int64(1 + // free
			4 + // nDescendants
			4 + // key
			4*tree + // parents
			4*2), // children
	}
	go f.creator()
	return f
}

func (f *File) Create(n Node) (int, error) {
	args := createArgs{node: n, result: make(chan createResult)}
	f.createChan <- args
	result := <-args.result
	return result.id, result.err
}

func (f *File) create(n Node) (int, error) {
	buf := &bytes.Buffer{}
	f.nodeToBuf(buf, n)
	id := f.nodeCount()
	_, err := f.appendFile.Write(buf.Bytes())
	return id, err
}

func (f *File) Find(id int) (Node, error) {
	node := Node{}
	node.id = id
	node.storage = f
	offset := f.offset(id)
	err := f.locker.ReadLock(f.file.Fd(), offset, f.nodeSize)
	if err != nil {
		return node, err
	}
	defer f.locker.UnLock(f.file.Fd(), offset, f.nodeSize)

	b := make([]byte, f.nodeSize)
	_, err = syscall.Pread(int(f.file.Fd()), b, offset)
	if err != nil {
		return node, err
	}

	buf := bytes.NewReader(b)

	var free bool
	binary.Read(buf, binary.BigEndian, &free)
	node.free = free

	var nDescendants int32
	binary.Read(buf, binary.BigEndian, &nDescendants)
	node.nDescendants = int(nDescendants)

	var key int32
	binary.Read(buf, binary.BigEndian, &key)
	node.key = int(key)

	parents := make([]int32, f.tree)
	binary.Read(buf, binary.BigEndian, &parents)
	nodeParents := make([]int, f.tree)
	for i, parent := range parents {
		nodeParents[i] = int(parent)
	}
	node.parents = nodeParents

	if node.nDescendants == 1 {
		// leaf node
		node.children = []int{0, 0} // skip children
		node.v = bytesToFloat64s(b[f.offsetOfV:])
	} else if node.nDescendants <= f.K {
		// bucket node
		children := make([]int32, nDescendants)
		binary.Read(buf, binary.BigEndian, &children)
		nodeChildren := make([]int, nDescendants)
		for i, child := range children {
			nodeChildren[i] = int(child)
		}
		node.children = nodeChildren
	} else {
		// other node
		children := make([]int32, 2)
		binary.Read(buf, binary.BigEndian, &children)
		nodeChildren := make([]int, 2)
		for i, child := range children {
			nodeChildren[i] = int(child)
		}
		node.children = nodeChildren
		node.v = bytesToFloat64s(b[f.offsetOfV:])
	}
	return node, nil
}

func (f *File) Update(n Node) error {
	buf := &bytes.Buffer{}
	f.nodeToBuf(buf, n)
	offset := f.offset(n.id)
	file, _ := os.OpenFile(f.filename, os.O_RDWR, 0)
	defer file.Close()

	err := f.locker.WriteLock(file.Fd(), offset, f.nodeSize)
	if err != nil {
		return err
	}
	defer f.locker.UnLock(file.Fd(), offset, f.nodeSize)

	_, err = syscall.Pwrite(int(file.Fd()), buf.Bytes(), offset)
	return err
}

func (f *File) UpdateParent(id, rootIndex, parent int) error {
	offset := f.offset(id) +
		int64(1+ // free
			4+ // nDescendants
			4+ // key
			4*rootIndex) // parents
	buf := &bytes.Buffer{}
	binary.Write(buf, binary.BigEndian, int32(parent))

	file, _ := os.OpenFile(f.filename, os.O_RDWR, 0)
	defer file.Close()

	err := f.locker.WriteLock(file.Fd(), offset, 4)
	if err != nil {
		return err
	}
	defer f.locker.UnLock(file.Fd(), offset, 4)

	_, err = syscall.Pwrite(int(file.Fd()), buf.Bytes(), offset)
	return err
}

func (f *File) Delete(n Node) error {
	n.free = true
	return f.Update(n)
}

func (f *File) Iterate(c chan Node) {
	count := f.nodeCount()
	// TODO: Use goroutine
	for i := 0; i < count; i++ {
		n, err := f.Find(i)
		if err != nil {
			close(c)
		}
		c <- n
	}
	close(c)
}

func (f File) offset(id int) int64 {
	return (int64(id) * f.nodeSize)
}

func (f File) nodeCount() int {
	stat, _ := f.file.Stat()
	size := stat.Size()
	return int(size / f.nodeSize)
}

func (f File) nodeToBuf(buf *bytes.Buffer, node Node) {
	// 1bytes free
	binary.Write(buf, binary.BigEndian, node.free)

	// 4bytes nDescendants
	binary.Write(buf, binary.BigEndian, int32(node.nDescendants))

	// 4bytes key
	binary.Write(buf, binary.BigEndian, int32(node.key))

	// 4bytes parents
	parents := make([]int32, len(node.parents))
	for i, parent := range node.parents {
		parents[i] = int32(parent)
	}
	binary.Write(buf, binary.BigEndian, parents)

	if node.isBucket() {
		// 4bytes children in K
		children := make([]int32, f.K)
		for i, child := range node.children {
			children[i] = int32(child)
		}
		binary.Write(buf, binary.BigEndian, children)

		// padding by zero
		remainingSize := ((2*4 + 8*f.dim) - (4 * f.K))
		binary.Write(buf, binary.BigEndian, make([]int32, remainingSize/4))
	} else {
		// 4bytes children in K
		children := make([]int32, 2)
		for i, child := range node.children {
			children[i] = int32(child)
		}
		binary.Write(buf, binary.BigEndian, children)

		// 8bytes v in f
		vec := make([]float64, f.dim)
		for i, v := range node.v {
			vec[i] = float64(v)
		}
		binary.Write(buf, binary.BigEndian, vec)
	}
}

type createArgs struct {
	node   Node
	result chan createResult
}

type createResult struct {
	id  int
	err error
}

func (f *File) creator() {
	for args := range f.createChan {
		id, err := f.create(args.node)
		args.result <- createResult{
			id:  id,
			err: err,
		}
	}
}

func (f File) size() int64 {
	info, _ := f.file.Stat()
	return info.Size()
}

func bytesToFloat64s(bytes []byte) []float64 {
	size := len(bytes) / 8
	floats := make([]float64, size)
	for i := 0; i < size; i++ {
		floats[i] = math.Float64frombits(binary.BigEndian.Uint64(bytes[0:8]))
		bytes = bytes[8:]
	}
	return floats
}
<|endoftext|>