<commit_before>package bot

import (
	"github.com/nlopes/slack"
	"gopkg.in/redis.v3"
	"regexp"
	"strings"
)

var (
	remember_re *regexp.Regexp
	tell_re     *regexp.Regexp
)

type Anzu struct {
	*BaseBot
	rc *redis.Client
}

func NewAnzu(token string, stop *chan struct{}, redisClient *redis.Client) *Anzu {
	remember_re = regexp.MustCompile("^안즈쨩? 기억해? ([^/]+)/(.+)")
	tell_re = regexp.MustCompile("^안즈쨩? 알려줘 (.+)")
	return &Anzu{NewBot(token, stop), redisClient}
}

func (bot *Anzu) onMessageEvent(e *slack.MessageEvent) {
	switch {
	case e.Text == "사람은 일을 하고 살아야한다. 메우":
		bot.SendMessage(bot.NewOutgoingMessage("이거 놔라 이 퇴근도 못하는 놈이", e.Channel))
		break
	case e.Text == "안즈쨩 카와이":
		bot.SendMessage(bot.NewOutgoingMessage("뭐... 뭐라는거야", e.Channel))
		break
	case e.Text == "안즈쨩 뭐해?":
		bot.SendMessage(bot.NewOutgoingMessage("숨셔", e.Channel))
		break
	default:
		if matched, ok := MatchRE(e.Text, remember_re); ok {
			key, val := strings.TrimSpace(matched[0]), strings.TrimSpace(matched[1])
			if key == "" || val == "" {
				bot.SendMessage(bot.NewOutgoingMessage("에...?", e.Channel))
			} else if _, ok := MatchRE(val, tell_re); ok {
				bot.SendMessage(bot.NewOutgoingMessage("에... 귀찮아...", e.Channel))
			} else {
				bot.rc.Set(key, val, 0)
				bot.SendMessage(bot.NewOutgoingMessage("에... 귀찮지만 기억했어", e.Channel))
			}
		} else if matched, ok := MatchRE(e.Text, tell_re); ok {
			key := strings.TrimSpace(matched[0])
			val := bot.rc.Get(key).String()
			if val == "" {
				bot.SendMessage(bot.NewOutgoingMessage("그런거 몰라", e.Channel))
			} else {
				bot.SendMessage(bot.NewOutgoingMessage(val, e.Channel))
			}
		}
		break
	}
}
<commit_msg>fix captured index error<commit_after>package bot

import (
	"github.com/nlopes/slack"
	"gopkg.in/redis.v3"
	"regexp"
	"strings"
)

var (
	remember_re *regexp.Regexp
	tell_re     *regexp.Regexp
)

type Anzu struct {
	*BaseBot
	rc *redis.Client
}

func NewAnzu(token string, stop *chan struct{}, redisClient *redis.Client) *Anzu {
	remember_re = regexp.MustCompile("^안즈쨩? 기억해? ([^/]+)/(.+)")
	tell_re = regexp.MustCompile("^안즈쨩? 알려줘 (.+)")
	return &Anzu{NewBot(token, stop), redisClient}
}

func (bot *Anzu) onMessageEvent(e *slack.MessageEvent) {
	switch {
	case e.Text == "사람은 일을 하고 살아야한다. 메우":
		bot.SendMessage(bot.NewOutgoingMessage("이거 놔라 이 퇴근도 못하는 놈이", e.Channel))
		break
	case e.Text == "안즈쨩 카와이":
		bot.SendMessage(bot.NewOutgoingMessage("뭐... 뭐라는거야", e.Channel))
		break
	case e.Text == "안즈쨩 뭐해?":
		bot.SendMessage(bot.NewOutgoingMessage("숨셔", e.Channel))
		break
	default:
		if matched, ok := MatchRE(e.Text, remember_re); ok {
			key, val := strings.TrimSpace(matched[1]), strings.TrimSpace(matched[2])
			if key == "" || val == "" {
				bot.SendMessage(bot.NewOutgoingMessage("에...?", e.Channel))
			} else if _, ok := MatchRE(val, tell_re); ok {
				bot.SendMessage(bot.NewOutgoingMessage("에... 귀찮아...", e.Channel))
			} else {
				bot.rc.Set(key, val, 0)
				bot.SendMessage(bot.NewOutgoingMessage("에... 귀찮지만 기억했어", e.Channel))
			}
		} else if matched, ok := MatchRE(e.Text, tell_re); ok {
			key := strings.TrimSpace(matched[1])
			val := bot.rc.Get(key).Val()
			if val == "" {
				bot.SendMessage(bot.NewOutgoingMessage("그런거 몰라", e.Channel))
			} else {
				bot.SendMessage(bot.NewOutgoingMessage(val, e.Channel))
			}
		}
		break
	}
}
<|endoftext|>


<commit_before>package mithril

import (
	"regexp"
	"strings"
)

// M returns VirtualElement
func M(selector string, opts ...interface{}) *VirtualElement {
	// declare query string
	query := "div"
	// replace query as div if selector does not existed
	if selector != "" {
		query = selector
	}
	// create virtual element
	element := &VirtualElement{"div", NewAttributes(), nil}
	// map options
	for i, opt := range opts {
		switch obj := opt.(type) {
		case *Attributes:
			if i == 0 {
				element.Attrs = obj
			}
		case string, bool, int, *VirtualElement:
			element.Children = obj
		}
	}
	// match all results
	re := regexp.MustCompile(`(^[\w\-]+|\#[\w\-]+|\.[\w\-]+|\[[\w\-]+\=\"[\w\-]+\"\]|\[[\w\-]+\=\'[\w\-]+\'\])`)
	for _, res := range re.FindAllStringSubmatch(query, -1) {
		data := res[0]
		value := data[1:len(data)]
		switch string(data[0]) {
		case "[":
			data := strings.Split(value, "=")
			switch len(data) {
			case 1:
				element.Attrs.Data[data[0][:len(data[0])-1]] = ""
			case 2:
				element.Attrs.Data[data[0]] = data[1][1 : len(data[1])-2]
			}
		case "#":
			element.Attrs.ID = value
		case ".":
			element.Attrs.Class = append(element.Attrs.Class, value)
		default:
			element.Tag = data
		}
	}
	// return virtual element
	return element
}
<commit_msg>VirtualElement creation improvements<commit_after>package mithril

import "regexp"

// M returns VirtualElement
func M(selector string, opts ...interface{}) *VirtualElement {
	// declare query string
	query := "div"
	// replace query as div if selector does not existed
	if selector != "" {
		query = selector
	}
	// create virtual element
	element := &VirtualElement{"div", NewAttributes(), nil}
	// map options
	for i, opt := range opts {
		switch obj := opt.(type) {
		case *Attributes:
			if i == 0 {
				element.Attrs = obj
			}
		case string, bool, int, *VirtualElement:
			element.Children = obj
		}
	}
	// match all results
	matches := regexp.MustCompile(`(?:(^|#|\.)([^#\.\[\]]+))|(\[.+?\])`)
	for _, res := range matches.FindAllStringSubmatch(query, -1) {
		if res[1] == "" && len(res[2]) > 0 {
			element.Tag = res[2]
		} else if res[1] == "#" {
			element.Attrs.ID = res[2]
		} else if res[1] == "." {
			element.Attrs.Class = append(element.Attrs.Class, res[2])
		} else if string(res[3][0]) == "[" {
			matches := regexp.MustCompile(`\[(.+?)(?:=("|'|)(.*?)("|'|))?\]`)
			for _, pair := range matches.FindAllStringSubmatch(res[3], -1) {
				if len(pair[1]) > 0 {
					element.Attrs.Data[pair[1]] = pair[3]
				}
			}
		}
	}
	// return virtual element
	return element
}
<|endoftext|>


<commit_before>package bots

type Bot interface {
	Name() string
	SetName(string)
	Path() string
	Start() error
	SendMessage(QuestionMessage) (*ReplyMessage, error)
}

type QuestionMessage struct {
	GameID      string      `json:"game-id,omitempty" binding:"required"`
	Action      string      `json:"action,omitempty" binding:"required"`
	Game        string      `json:"game,omitempty" binding:"required"`
	Players     int         `json:"players,omitempty"`
	Board       interface{} `json:"board,omitempty"`
	You         interface{} `json:"you,omitempty"`
	PlayerIndex int         `json:"player-index,omitempty" binding:"required"`
}

type ReplyMessage struct {
	Name        string      `json:"name,omitempty"`
	Play        interface{} `json:"play,omitempty"`
	Error       interface{} `json:"error,omitempty"`
	PlayerIndex int         `json:"player-index,omitempty" binding:"required"`
}

// InitTurnBasedBots is an helper that starts and discovers connected bots
func InitTurnBasedBots(bots []Bot, gameName, gameID string) error {
	// start bots
	for _, bot := range bots {
		if err := bot.Start(); err != nil {
			return err
		}
	}

	// send init message to bots
	for idx, bot := range bots {
		reply, err := bot.SendMessage(QuestionMessage{
			GameID:      gameID,
			Action:      "init",
			Game:        gameName,
			Players:     len(bots),
			PlayerIndex: idx,
		})
		if err != nil {
			return err
		}

		// parse reply
		if reply.Name != "" {
			bot.SetName(reply.Name)
		}
	}
	return nil
}
<commit_msg>permit bot to add a comment <commit_after>package bots

type Bot interface {
	Name() string
	SetName(string)
	Path() string
	Start() error
	SendMessage(QuestionMessage) (*ReplyMessage, error)
}

type QuestionMessage struct {
	GameID      string      `json:"game-id,omitempty" binding:"required"`
	Action      string      `json:"action,omitempty" binding:"required"`
	Game        string      `json:"game,omitempty" binding:"required"`
	Players     int         `json:"players,omitempty"`
	Board       interface{} `json:"board,omitempty"`
	You         interface{} `json:"you,omitempty"`
	PlayerIndex int         `json:"player-index,omitempty" binding:"required"`
}

type ReplyMessage struct {
	Name        string      `json:"name,omitempty"`
	Play        interface{} `json:"play,omitempty"`
	Error       interface{} `json:"error,omitempty"`
	Comment     interface{} `json:"comment,omitempty"`
	PlayerIndex int         `json:"player-index,omitempty" binding:"required"`
}

// InitTurnBasedBots is an helper that starts and discovers connected bots
func InitTurnBasedBots(bots []Bot, gameName, gameID string) error {
	// start bots
	for _, bot := range bots {
		if err := bot.Start(); err != nil {
			return err
		}
	}

	// send init message to bots
	for idx, bot := range bots {
		reply, err := bot.SendMessage(QuestionMessage{
			GameID:      gameID,
			Action:      "init",
			Game:        gameName,
			Players:     len(bots),
			PlayerIndex: idx,
		})
		if err != nil {
			return err
		}

		// parse reply
		if reply.Name != "" {
			bot.SetName(reply.Name)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package box

//import "log"

func (r *Run) ensure(nb int) {
	if nb == r.Nalloc {
		r.Grow(r.delta)
		if r.delta < 32768 {
			r.delta *= 2
		}
	}
}
func min(a,b int) int{
	if a<b{
		return a
	}
	return b
}

func (r *Run) Boxscan(s []byte, ymax int) {
	r.Nbox = 0
	r.Nchars = 0
	r.Nchars += int64(len(s))
	i := 0
	nb := 0
	for nl := 0; nl <= ymax; nb++ {
		if nb == r.Nalloc {
			r.Grow(r.delta)
			if r.delta < 32768 {
				r.delta *= 2
			}
		}
		if i == len(s) {
			break
		}
		i++
		c := s[i-1]
		switch c {
		default:
			for _, c = range s[i:min(len(s),MaxBytes)] {
				if special(c) {
					break
				}
				i++
			}
			b := &r.Box[nb]
			b.Nrune=i
			b.Ptr = b.Ptr[:i]
			copy(b.Ptr, s[:i])
			b.Width = r.MeasureBytes(b.Ptr)
//			r.Box[nb] = Box{
//				Nrune: i,
//				Ptr:   s[:i],
//				Width: r.MeasureBytes(s[:i]),
//			}
		case '\t':
			b := &r.Box[nb]
			b.Nrune=-1
			b.Ptr=b.Ptr[:1]
			b.Ptr[0]='\t'
			b.Width=r.maxDx
			b.Minwidth=r.minDx
//			r.Box[nb] = Box{
//				Nrune:    -1,
//				Ptr:      []byte("\t"),
//				Width:    r.minDx,
//				Minwidth: r.minDx,
//			}
		case '\n':
			b := &r.Box[nb]
			b.Nrune=-1
			b.Ptr=b.Ptr[:1]
			b.Ptr[0]='\n'
			b.Width=r.maxDx
//			r.Box[nb] = Box{
//				Nrune: -1,
//				Ptr:   []byte("\n"),
//				Width: r.maxDx,
//			}
			nl++
		}
		s = s[i:]
		i = 0
	}
	r.Nchars -= int64(len(s))
	r.Nbox += nb
}

func special(c byte) bool {
	return c == '\t' || c == '\n'
}
<commit_msg>box: scan: restore plain box behavior for allocations<commit_after>package box

//import "log"

func (r *Run) ensure(nb int) {
	if nb == r.Nalloc {
		r.Grow(r.delta)
		if r.delta < 32768 {
			r.delta *= 2
		}
	}
}
func min(a,b int) int{
	if a<b{
		return a
	}
	return b
}

func (r *Run) Boxscan(s []byte, ymax int) {
	r.Nbox = 0
	r.Nchars = 0
	r.Nchars += int64(len(s))
	i := 0
	nb := 0
	for nl := 0; nl <= ymax; nb++ {
		if nb == r.Nalloc {
			r.Grow(r.delta)
			if r.delta < 32768 {
				r.delta *= 2
			}
		}
		if i == len(s) {
			break
		}
		i++
		c := s[i-1]
		switch c {
		default:
			for _, c = range s[i:min(len(s),MaxBytes)] {
				if special(c) {
					break
				}
				i++
			}
			r.Box[nb] = Box{
				Nrune: i,
				Ptr:   s[:i],
				Width: r.MeasureBytes(s[:i]),
			}
		case '\t':
//			r.Box[nb] = Box{
//				Nrune:    -1,
//				Ptr:      []byte("\t"),
//				Width:    r.minDx,
//				Minwidth: r.minDx,
//			}
		case '\n':
			b := &r.Box[nb]
			b.Nrune=-1
			b.Ptr=b.Ptr[:1]
			b.Ptr[0]='\n'
			b.Width=r.maxDx
//			r.Box[nb] = Box{
//				Nrune: -1,
//				Ptr:   []byte("\n"),
//				Width: r.maxDx,
//			}
			nl++
		}
		s = s[i:]
		i = 0
	}
	r.Nchars -= int64(len(s))
	r.Nbox += nb
}

func special(c byte) bool {
	return c == '\t' || c == '\n'
}
<|endoftext|>


<commit_before>// Copyright 2016 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github_test

import (
	"context"
	"fmt"
	"log"

	"github.com/google/go-github/v18/github"
)

func ExampleClient_Markdown() {
	client := github.NewClient(nil)

	input := "# heading #\n\nLink to issue #1"
	opt := &github.MarkdownOptions{Mode: "gfm", Context: "google/go-github"}

	output, _, err := client.Markdown(context.Background(), input, opt)
	if err != nil {
		fmt.Println(err)
	}

	fmt.Println(output)
}

func ExampleRepositoriesService_GetReadme() {
	client := github.NewClient(nil)

	readme, _, err := client.Repositories.GetReadme(context.Background(), "google", "go-github", nil)
	if err != nil {
		fmt.Println(err)
		return
	}

	content, err := readme.GetContent()
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("google/go-github README:\n%v\n", content)
}

func ExampleRepositoriesService_List() {
	client := github.NewClient(nil)

	user := "willnorris"
	opt := &github.RepositoryListOptions{Type: "owner", Sort: "updated", Direction: "desc"}

	repos, _, err := client.Repositories.List(context.Background(), user, opt)
	if err != nil {
		fmt.Println(err)
	}

	fmt.Printf("Recently updated repositories by %q: %v", user, github.Stringify(repos))
}

func ExampleRepositoriesService_CreateFile() {
	// In this example we're creating a new file in a repository using the
	// Contents API. Only 1 file per commit can be managed through that API.

	// Note that authentication is needed here as you are performing a modification
	// so you will need to modify the example to provide an oauth client to
	// github.NewClient() instead of nil. See the following documentation for more
	// information on how to authenticate with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	ctx := context.Background()
	fileContent := []byte("This is the content of my file\nand the 2nd line of it")

	// Note: the file needs to be absent from the repository as you are not
	// specifying a SHA reference here.
	opts := &github.RepositoryContentFileOptions{
		Message:   github.String("This is my commit message"),
		Content:   fileContent,
		Branch:    github.String("master"),
		Committer: &github.CommitAuthor{Name: github.String("FirstName LastName"), Email: github.String("user@example.com")},
	}
	_, _, err := client.Repositories.CreateFile(ctx, "myOrganization", "myRepository", "myNewFile.md", opts)
	if err != nil {
		fmt.Println(err)
		return
	}
}

func ExampleUsersService_ListAll() {
	client := github.NewClient(nil)
	opts := &github.UserListOptions{}
	for {
		users, _, err := client.Users.ListAll(context.Background(), opts)
		if err != nil {
			log.Fatalf("error listing users: %v", err)
		}
		if len(users) == 0 {
			break
		}
		opts.Since = *users[len(users)-1].ID
		// Process users...
	}
}

func ExamplePullRequestsService_Create() {
	// In this example we're creating a PR and displaying the HTML url at the end.

	// Note that authentication is needed here as you are performing a modification
	// so you will need to modify the example to provide an oauth client to
	// github.NewClient() instead of nil. See the following documentation for more
	// information on how to authenticate with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	newPR := &github.NewPullRequest{
		Title:               github.String("My awesome pull request"),
		Head:                github.String("branch_to_merge"),
		Base:                github.String("master"),
		Body:                github.String("This is the description of the PR created with the package `github.com/google/go-github/github`"),
		MaintainerCanModify: github.Bool(true),
	}

	pr, _, err := client.PullRequests.Create(context.Background(), "myOrganization", "myRepository", newPR)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("PR created: %s\n", pr.GetHTMLURL())
}

func ExampleTeamsService_ListTeams() {
	// This example shows how to get a team ID corresponding to a given team name.

	// Note that authentication is needed here as you are performing a lookup on
	// an organization's administrative configuration, so you will need to modify
	// the example to provide an oauth client to github.NewClient() instead of nil.
	// See the following documentation for more information on how to authenticate
	// with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	teamName := "Developers team"
	ctx := context.Background()
	opts := &github.ListOptions{}

	for {
		teams, resp, err := client.Teams.ListTeams(ctx, "myOrganization", opts)
		if err != nil {
			fmt.Println(err)
			return
		}
		for _, t := range teams {
			if t.GetName() == teamName {
				fmt.Printf("Team %q has ID %d\n", teamName, t.GetID())
				return
			}
		}
		if resp.NextPage == 0 {
			break
		}
		opts.Page = resp.NextPage
	}

	fmt.Printf("Team %q was not found\n", teamName)
}
<commit_msg>Update examples_test.go (#1024)<commit_after>// Copyright 2016 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// These examples are inlined in godoc.

package github_test

import (
	"context"
	"fmt"
	"log"

	"github.com/google/go-github/v18/github"
)

func ExampleClient_Markdown() {
	client := github.NewClient(nil)

	input := "# heading #\n\nLink to issue #1"
	opt := &github.MarkdownOptions{Mode: "gfm", Context: "google/go-github"}

	output, _, err := client.Markdown(context.Background(), input, opt)
	if err != nil {
		fmt.Println(err)
	}

	fmt.Println(output)
}

func ExampleRepositoriesService_GetReadme() {
	client := github.NewClient(nil)

	readme, _, err := client.Repositories.GetReadme(context.Background(), "google", "go-github", nil)
	if err != nil {
		fmt.Println(err)
		return
	}

	content, err := readme.GetContent()
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("google/go-github README:\n%v\n", content)
}

func ExampleRepositoriesService_List() {
	client := github.NewClient(nil)

	user := "willnorris"
	opt := &github.RepositoryListOptions{Type: "owner", Sort: "updated", Direction: "desc"}

	repos, _, err := client.Repositories.List(context.Background(), user, opt)
	if err != nil {
		fmt.Println(err)
	}

	fmt.Printf("Recently updated repositories by %q: %v", user, github.Stringify(repos))
}

func ExampleRepositoriesService_CreateFile() {
	// In this example we're creating a new file in a repository using the
	// Contents API. Only 1 file per commit can be managed through that API.

	// Note that authentication is needed here as you are performing a modification
	// so you will need to modify the example to provide an oauth client to
	// github.NewClient() instead of nil. See the following documentation for more
	// information on how to authenticate with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	ctx := context.Background()
	fileContent := []byte("This is the content of my file\nand the 2nd line of it")

	// Note: the file needs to be absent from the repository as you are not
	// specifying a SHA reference here.
	opts := &github.RepositoryContentFileOptions{
		Message:   github.String("This is my commit message"),
		Content:   fileContent,
		Branch:    github.String("master"),
		Committer: &github.CommitAuthor{Name: github.String("FirstName LastName"), Email: github.String("user@example.com")},
	}
	_, _, err := client.Repositories.CreateFile(ctx, "myOrganization", "myRepository", "myNewFile.md", opts)
	if err != nil {
		fmt.Println(err)
		return
	}
}

func ExampleUsersService_ListAll() {
	client := github.NewClient(nil)
	opts := &github.UserListOptions{}
	for {
		users, _, err := client.Users.ListAll(context.Background(), opts)
		if err != nil {
			log.Fatalf("error listing users: %v", err)
		}
		if len(users) == 0 {
			break
		}
		opts.Since = *users[len(users)-1].ID
		// Process users...
	}
}

func ExamplePullRequestsService_Create() {
	// In this example we're creating a PR and displaying the HTML url at the end.

	// Note that authentication is needed here as you are performing a modification
	// so you will need to modify the example to provide an oauth client to
	// github.NewClient() instead of nil. See the following documentation for more
	// information on how to authenticate with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	newPR := &github.NewPullRequest{
		Title:               github.String("My awesome pull request"),
		Head:                github.String("branch_to_merge"),
		Base:                github.String("master"),
		Body:                github.String("This is the description of the PR created with the package `github.com/google/go-github/github`"),
		MaintainerCanModify: github.Bool(true),
	}

	pr, _, err := client.PullRequests.Create(context.Background(), "myOrganization", "myRepository", newPR)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("PR created: %s\n", pr.GetHTMLURL())
}

func ExampleTeamsService_ListTeams() {
	// This example shows how to get a team ID corresponding to a given team name.

	// Note that authentication is needed here as you are performing a lookup on
	// an organization's administrative configuration, so you will need to modify
	// the example to provide an oauth client to github.NewClient() instead of nil.
	// See the following documentation for more information on how to authenticate
	// with the client:
	// https://godoc.org/github.com/google/go-github/github#hdr-Authentication
	client := github.NewClient(nil)

	teamName := "Developers team"
	ctx := context.Background()
	opts := &github.ListOptions{}

	for {
		teams, resp, err := client.Teams.ListTeams(ctx, "myOrganization", opts)
		if err != nil {
			fmt.Println(err)
			return
		}
		for _, t := range teams {
			if t.GetName() == teamName {
				fmt.Printf("Team %q has ID %d\n", teamName, t.GetID())
				return
			}
		}
		if resp.NextPage == 0 {
			break
		}
		opts.Page = resp.NextPage
	}

	fmt.Printf("Team %q was not found\n", teamName)
}
<|endoftext|>


<commit_before>package orcraft

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	golog "log"
	"os"
	"time"

	"github.com/github/orchestrator/go/db"

	"github.com/hashicorp/raft"
	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
)

type dummyReadCloser struct {
}

func (c *dummyReadCloser) Close() error {
	return nil
}

func (c *dummyReadCloser) Read(p []byte) (n int, err error) {
	return 0, io.EOF
}

// RelSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type RelSnapshotStore struct {
	retain int
	logger *golog.Logger
}

// RelSnapshotSink implements SnapshotSink with a file.
type RelSnapshotSink struct {
	store  *RelSnapshotStore
	logger *golog.Logger
	meta   relSnapshotMeta

	closed bool
}

// relSnapshotMeta is stored on disk. We also put a CRC
// on disk so that we can verify the snapshot.
type relSnapshotMeta struct {
	raft.SnapshotMeta
	CRC []byte
}

// bufferedFile is returned when we open a snapshot. This way
// reads are buffered and the file still gets closed.
type bufferedFile struct {
	bh *bufio.Reader
	fh *os.File
}

func (b *bufferedFile) Read(p []byte) (n int, err error) {
	log.Debugf("===== bufferedFile.Read")
	return b.bh.Read(p)
}

func (b *bufferedFile) Close() error {
	return b.fh.Close()
}

// NewRelSnapshotStoreWithLogger creates a new RelSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewRelSnapshotStoreWithLogger(retain int, logger *golog.Logger) (*RelSnapshotStore, error) {
	if retain < 1 {
		return nil, log.Errorf("must retain at least one snapshot")
	}
	if logger == nil {
		logger = golog.New(os.Stderr, "", golog.LstdFlags)
	}

	// Setup the store
	store := &RelSnapshotStore{
		retain: retain,
		logger: logger,
	}

	return store, nil
}

// NewRelSnapshotStore creates a new RelSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewRelSnapshotStore(retain int, logOutput io.Writer) (*RelSnapshotStore, error) {
	if logOutput == nil {
		logOutput = os.Stderr
	}
	return NewRelSnapshotStoreWithLogger(retain, golog.New(logOutput, "", golog.LstdFlags))
}

// snapshotName generates a name for the snapshot.
func snapshotName(term, index uint64) string {
	now := time.Now()
	msec := now.UnixNano() / int64(time.Millisecond)
	return fmt.Sprintf("%d-%d-%d", term, index, msec)
}

// Create is used to start a new snapshot
func (f *RelSnapshotStore) Create(index, term uint64, peers []byte) (raft.SnapshotSink, error) {
	lastIndex := getRaft().LastIndex()
	log.Debugf("==== RelSnapshotStore create: %+v, %+v,", index, term)

	log.Debugf("==== RelSnapshotStore create: lastIndex=%+v", lastIndex)
	// Create a new path
	name := snapshotName(term, index)

	// Create the sink
	sink := &RelSnapshotSink{
		store:  f,
		logger: f.logger,
		meta: relSnapshotMeta{
			SnapshotMeta: raft.SnapshotMeta{
				ID:    name,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return nil, log.Errore(err)
	}

	// Done
	return sink, nil
}

// List returns available snapshots in the store.
func (f *RelSnapshotStore) List() ([]*raft.SnapshotMeta, error) {
	log.Debugf("===== RelSnapshotStore.List")
	// Get the eligible snapshots
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return nil, log.Errore(err)
	}

	var snapMeta []*raft.SnapshotMeta
	for _, meta := range snapshots {
		snapMeta = append(snapMeta, &meta.SnapshotMeta)
		if len(snapMeta) == f.retain {
			break
		}
	}
	return snapMeta, nil
}

// readSnapshots reads snapshots by query
func (f *RelSnapshotStore) readSnapshots(query string, args []interface{}) (snapMeta []*relSnapshotMeta, err error) {
	log.Debugf("===== RelSnapshotStore.readSnapshots; query=%+v", query)
	err = db.QueryOrchestrator(query, args, func(m sqlutils.RowMap) error {
		snapshotMetaText := m.GetString("snapshot_meta")

		meta := &relSnapshotMeta{}
		if err := json.Unmarshal([]byte(snapshotMetaText), &meta); err != nil {
			return err
		}

		snapMeta = append(snapMeta, meta)
		return nil
	})
	return snapMeta, log.Errore(err)
}

func (f *RelSnapshotStore) readMeta(name string) (*relSnapshotMeta, error) {
	log.Debugf("===== RelSnapshotStore.readMeta; name=%+v", name)
	query := `select snapshot_meta from raft_snapshot where snapshot_name=?`
	snapshots, err := f.readSnapshots(query, sqlutils.Args(name))
	if err != nil {
		return nil, log.Errore(err)
	}
	if len(snapshots) == 1 {
		return snapshots[0], nil
	}
	return nil, log.Errorf("Found %+v snapshots for %s", len(snapshots), name)
}

// getSnapshots returns all the known snapshots.
func (f *RelSnapshotStore) getSnapshots() (snapMeta []*relSnapshotMeta, err error) {
	query := `select snapshot_meta from raft_snapshot order by snapshot_id desc`
	return f.readSnapshots(query, sqlutils.Args())
}

// Open takes a snapshot ID and returns a ReadCloser for that snapshot.
func (f *RelSnapshotStore) Open(id string) (*raft.SnapshotMeta, io.ReadCloser, error) {
	log.Debugf("===== RelSnapshotStore.Open; id=%+v", id)
	// Get the metadata
	meta, err := f.readMeta(id)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get meta data to open snapshot: %v", err)
		return nil, nil, log.Errore(err)
	}
	return &meta.SnapshotMeta, &dummyReadCloser{}, nil
}

// ReapSnapshots reaps any snapshots beyond the retain count.
func (f *RelSnapshotStore) ReapSnapshots() error {
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return log.Errore(err)
	}

	for i := f.retain; i < len(snapshots); i++ {
		if _, err := db.ExecOrchestrator(`delete from raft_snapshot where snapshot_name=?`, snapshots[i].ID); err != nil {
			f.logger.Printf("[ERR] snapshot: Failed to reap snapshot %v: %v", snapshots[i].ID, err)
			return log.Errore(err)
		}
	}
	return nil
}

// ID returns the ID of the snapshot, can be used with Open()
// after the snapshot is finalized.
func (s *RelSnapshotSink) ID() string {
	return s.meta.ID
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches.
func (s *RelSnapshotSink) Write(b []byte) (int, error) {
	return 0, nil
}

// Close is used to indicate a successful end.
func (s *RelSnapshotSink) Close() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true
	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return log.Errore(err)
	}

	// Reap any old snapshots
	if err := s.store.ReapSnapshots(); err != nil {
		return log.Errore(err)
	}

	return nil
}

// Cancel is used to indicate an unsuccessful end.
func (s *RelSnapshotSink) Cancel() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	return nil
}

// writeMeta is used to write out the metadata we have.
func (s *RelSnapshotSink) writeMeta() error {
	b, err := json.Marshal(&s.meta)
	if err != nil {
		return log.Errore(err)
	}
	snapshotMetaText := string(b)
	_, err = db.ExecOrchestrator(`
		replace into raft_snapshot
			(snapshot_id, snapshot_name, snapshot_meta, created_at)
		values
			(null, ?, ?, now())
			`, s.meta.ID, snapshotMetaText)
	return log.Errore(err)
}
<commit_msg>investigating installSnapshot<commit_after>package orcraft

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	golog "log"
	"os"
	"time"

	"github.com/github/orchestrator/go/db"

	"github.com/hashicorp/raft"
	"github.com/openark/golib/log"
	"github.com/openark/golib/sqlutils"
)

type dummyReadCloser struct {
}

func (c *dummyReadCloser) Close() error {
	return nil
}

func (c *dummyReadCloser) Read(p []byte) (n int, err error) {
	return 0, io.EOF
}

// RelSnapshotStore implements the SnapshotStore interface and allows
// snapshots to be made on the local disk.
type RelSnapshotStore struct {
	retain int
	logger *golog.Logger
}

// RelSnapshotSink implements SnapshotSink with a file.
type RelSnapshotSink struct {
	store  *RelSnapshotStore
	logger *golog.Logger
	meta   relSnapshotMeta

	closed bool
}

// relSnapshotMeta is stored on disk. We also put a CRC
// on disk so that we can verify the snapshot.
type relSnapshotMeta struct {
	raft.SnapshotMeta
	CRC []byte
}

// bufferedFile is returned when we open a snapshot. This way
// reads are buffered and the file still gets closed.
type bufferedFile struct {
	bh *bufio.Reader
	fh *os.File
}

func (b *bufferedFile) Read(p []byte) (n int, err error) {
	log.Debugf("===== bufferedFile.Read")
	return b.bh.Read(p)
}

func (b *bufferedFile) Close() error {
	return b.fh.Close()
}

// NewRelSnapshotStoreWithLogger creates a new RelSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewRelSnapshotStoreWithLogger(retain int, logger *golog.Logger) (*RelSnapshotStore, error) {
	if retain < 1 {
		return nil, log.Errorf("must retain at least one snapshot")
	}
	if logger == nil {
		logger = golog.New(os.Stderr, "", golog.LstdFlags)
	}

	// Setup the store
	store := &RelSnapshotStore{
		retain: retain,
		logger: logger,
	}

	return store, nil
}

// NewRelSnapshotStore creates a new RelSnapshotStore based
// on a base directory. The `retain` parameter controls how many
// snapshots are retained. Must be at least 1.
func NewRelSnapshotStore(retain int, logOutput io.Writer) (*RelSnapshotStore, error) {
	if logOutput == nil {
		logOutput = os.Stderr
	}
	return NewRelSnapshotStoreWithLogger(retain, golog.New(logOutput, "", golog.LstdFlags))
}

// snapshotName generates a name for the snapshot.
func snapshotName(term, index uint64) string {
	now := time.Now()
	msec := now.UnixNano() / int64(time.Millisecond)
	return fmt.Sprintf("%d-%d-%d", term, index, msec)
}

// Create is used to start a new snapshot
func (f *RelSnapshotStore) Create(index, term uint64, peers []byte) (raft.SnapshotSink, error) {
	lastIndex := getRaft().LastIndex() // our index
	if lastIndex < index {
		return nil, fmt.Errorf("RelSnapshotStore does not support remot e snaoshot stores")
	}
	log.Debugf("==== RelSnapshotStore create: %+v, %+v,", index, term)

	log.Debugf("==== RelSnapshotStore create: lastIndex=%+v", lastIndex)
	// Create a new path
	name := snapshotName(term, index)

	// Create the sink
	sink := &RelSnapshotSink{
		store:  f,
		logger: f.logger,
		meta: relSnapshotMeta{
			SnapshotMeta: raft.SnapshotMeta{
				ID:    name,
				Index: index,
				Term:  term,
				Peers: peers,
			},
			CRC: nil,
		},
	}

	// Write out the meta data
	if err := sink.writeMeta(); err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return nil, log.Errore(err)
	}

	// Done
	return sink, nil
}

// List returns available snapshots in the store.
func (f *RelSnapshotStore) List() ([]*raft.SnapshotMeta, error) {
	log.Debugf("===== RelSnapshotStore.List")
	// Get the eligible snapshots
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return nil, log.Errore(err)
	}

	var snapMeta []*raft.SnapshotMeta
	for _, meta := range snapshots {
		snapMeta = append(snapMeta, &meta.SnapshotMeta)
		if len(snapMeta) == f.retain {
			break
		}
	}
	return snapMeta, nil
}

// readSnapshots reads snapshots by query
func (f *RelSnapshotStore) readSnapshots(query string, args []interface{}) (snapMeta []*relSnapshotMeta, err error) {
	log.Debugf("===== RelSnapshotStore.readSnapshots; query=%+v", query)
	err = db.QueryOrchestrator(query, args, func(m sqlutils.RowMap) error {
		snapshotMetaText := m.GetString("snapshot_meta")

		meta := &relSnapshotMeta{}
		if err := json.Unmarshal([]byte(snapshotMetaText), &meta); err != nil {
			return err
		}

		snapMeta = append(snapMeta, meta)
		return nil
	})
	return snapMeta, log.Errore(err)
}

func (f *RelSnapshotStore) readMeta(name string) (*relSnapshotMeta, error) {
	log.Debugf("===== RelSnapshotStore.readMeta; name=%+v", name)
	query := `select snapshot_meta from raft_snapshot where snapshot_name=?`
	snapshots, err := f.readSnapshots(query, sqlutils.Args(name))
	if err != nil {
		return nil, log.Errore(err)
	}
	if len(snapshots) == 1 {
		return snapshots[0], nil
	}
	return nil, log.Errorf("Found %+v snapshots for %s", len(snapshots), name)
}

// getSnapshots returns all the known snapshots.
func (f *RelSnapshotStore) getSnapshots() (snapMeta []*relSnapshotMeta, err error) {
	query := `select snapshot_meta from raft_snapshot order by snapshot_id desc`
	return f.readSnapshots(query, sqlutils.Args())
}

// Open takes a snapshot ID and returns a ReadCloser for that snapshot.
func (f *RelSnapshotStore) Open(id string) (*raft.SnapshotMeta, io.ReadCloser, error) {
	log.Debugf("===== RelSnapshotStore.Open; id=%+v", id)
	// Get the metadata
	meta, err := f.readMeta(id)
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get meta data to open snapshot: %v", err)
		return nil, nil, log.Errore(err)
	}
	return &meta.SnapshotMeta, &dummyReadCloser{}, nil
}

// ReapSnapshots reaps any snapshots beyond the retain count.
func (f *RelSnapshotStore) ReapSnapshots() error {
	snapshots, err := f.getSnapshots()
	if err != nil {
		f.logger.Printf("[ERR] snapshot: Failed to get snapshots: %v", err)
		return log.Errore(err)
	}

	for i := f.retain; i < len(snapshots); i++ {
		if _, err := db.ExecOrchestrator(`delete from raft_snapshot where snapshot_name=?`, snapshots[i].ID); err != nil {
			f.logger.Printf("[ERR] snapshot: Failed to reap snapshot %v: %v", snapshots[i].ID, err)
			return log.Errore(err)
		}
	}
	return nil
}

// ID returns the ID of the snapshot, can be used with Open()
// after the snapshot is finalized.
func (s *RelSnapshotSink) ID() string {
	return s.meta.ID
}

// Write is used to append to the state file. We write to the
// buffered IO object to reduce the amount of context switches.
func (s *RelSnapshotSink) Write(b []byte) (int, error) {
	return 0, nil
}

// Close is used to indicate a successful end.
func (s *RelSnapshotSink) Close() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true
	// Write out the meta data
	if err := s.writeMeta(); err != nil {
		s.logger.Printf("[ERR] snapshot: Failed to write metadata: %v", err)
		return log.Errore(err)
	}

	// Reap any old snapshots
	if err := s.store.ReapSnapshots(); err != nil {
		return log.Errore(err)
	}

	return nil
}

// Cancel is used to indicate an unsuccessful end.
func (s *RelSnapshotSink) Cancel() error {
	// Make sure close is idempotent
	if s.closed {
		return nil
	}
	s.closed = true

	return nil
}

// writeMeta is used to write out the metadata we have.
func (s *RelSnapshotSink) writeMeta() error {
	b, err := json.Marshal(&s.meta)
	if err != nil {
		return log.Errore(err)
	}
	snapshotMetaText := string(b)
	_, err = db.ExecOrchestrator(`
		replace into raft_snapshot
			(snapshot_id, snapshot_name, snapshot_meta, created_at)
		values
			(null, ?, ?, now())
			`, s.meta.ID, snapshotMetaText)
	return log.Errore(err)
}
<|endoftext|>


<commit_before>package teams

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

func memberSetup(t *testing.T) (libkb.TestContext, *kbtest.FakeUser, string) {
	tc := libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	u, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name := createTeam(tc)

	return tc, u, name
}

func memberSetupMultiple(t *testing.T) (tc libkb.TestContext, owner, other *kbtest.FakeUser, name string) {
	tc = libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	other, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}
	tc.G.Logout()

	owner, err = kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name = createTeam(tc)

	return tc, owner, other, name
}

func TestMemberOwner(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	ctx := context.Background()
	team, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, team, u.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	aliceRole := usernameRole(ctx, tc, team, "t_alice")
	if aliceRole != keybase1.TeamRole_NONE {
		t.Errorf("role: %s, expected NONE", aliceRole)
	}
}

func TestMemberAddOwner(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleOwner(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", otherRole)
	}
}

func TestMemberAddAdmin(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleAdmin(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_ADMIN {
		t.Errorf("role: %s, expected ADMIN", otherRole)
	}
}

func TestMemberAddWriter(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleWriter(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_WRITER {
		t.Errorf("role: %s, expected WRITER", otherRole)
	}
}

func TestMemberAddReader(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleReader(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_READER {
		t.Errorf("role: %s, expected READER", otherRole)
	}
}

func TestMemberRemove(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleWriter(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	t.Logf("before remove, seqno: %d", s.Chain.GetLatestSeqno())

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_WRITER {
		t.Errorf("role: %s, expected WRITER", otherRole)
	}

	if err := RemoveMember(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx = context.Background()
	s, err = Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	t.Logf("after remove, seqno: %d", s.Chain.GetLatestSeqno())

	role = uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole = usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_NONE {
		t.Errorf("role: %s, expected NONE", otherRole)
	}
}

func TestMemberChangeRole(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleWriter(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_WRITER {
		t.Errorf("role: %s, expected WRITER", otherRole)
	}

	if err := SetRoleReader(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx = context.Background()
	s, err = Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role = uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole = usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_READER {
		t.Errorf("role: %s, expected READER", otherRole)
	}
}

func uidRole(ctx context.Context, tc libkb.TestContext, team *Team, uid keybase1.UID) keybase1.TeamRole {
	uv, err := loadUserVersionByUID(ctx, tc.G, uid)
	if err != nil {
		tc.T.Fatal(err)
	}
	return uvRole(tc, team, uv)
}

func usernameRole(ctx context.Context, tc libkb.TestContext, team *Team, username string) keybase1.TeamRole {
	uv, err := loadUserVersionByUsername(ctx, tc.G, username)
	if err != nil {
		tc.T.Fatal(err)
	}
	return uvRole(tc, team, uv)
}

func uvRole(tc libkb.TestContext, team *Team, uv UserVersion) keybase1.TeamRole {
	role, err := team.Chain.GetUserRole(uv)
	if err != nil {
		tc.T.Fatal(err)
	}
	return role
}
<commit_msg>Cleanup repeat code<commit_after>package teams

import (
	"testing"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/kbtest"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
)

func memberSetup(t *testing.T) (libkb.TestContext, *kbtest.FakeUser, string) {
	tc := libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	u, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name := createTeam(tc)

	return tc, u, name
}

func memberSetupMultiple(t *testing.T) (tc libkb.TestContext, owner, other *kbtest.FakeUser, name string) {
	tc = libkb.SetupTest(t, "team", 1)
	tc.Tp.UpgradePerUserKey = true

	other, err := kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}
	tc.G.Logout()

	owner, err = kbtest.CreateAndSignupFakeUser("team", tc.G)
	if err != nil {
		t.Fatal(err)
	}

	name = createTeam(tc)

	return tc, owner, other, name
}

func TestMemberOwner(t *testing.T) {
	tc, u, name := memberSetup(t)
	defer tc.Cleanup()

	ctx := context.Background()
	team, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, team, u.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	aliceRole := usernameRole(ctx, tc, team, "t_alice")
	if aliceRole != keybase1.TeamRole_NONE {
		t.Errorf("role: %s, expected NONE", aliceRole)
	}
}

type addTest struct {
	name        string
	setRoleFunc func(ctx context.Context, g *libkb.GlobalContext, teamname, username string) error
	afterRole   keybase1.TeamRole
}

var addTests = []addTest{
	addTest{name: "owner", setRoleFunc: SetRoleOwner, afterRole: keybase1.TeamRole_OWNER},
	addTest{name: "admin", setRoleFunc: SetRoleAdmin, afterRole: keybase1.TeamRole_ADMIN},
	addTest{name: "writer", setRoleFunc: SetRoleWriter, afterRole: keybase1.TeamRole_WRITER},
	addTest{name: "reader", setRoleFunc: SetRoleReader, afterRole: keybase1.TeamRole_READER},
}

func TestMemberAddX(t *testing.T) {
	for _, test := range addTests {
		testMemberAdd(t, test)
	}
}

func testMemberAdd(t *testing.T, test addTest) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := test.setRoleFunc(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != test.afterRole {
		t.Errorf("role: %s, expected %s", otherRole, test.afterRole)
	}
}

func TestMemberRemove(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleWriter(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	t.Logf("before remove, seqno: %d", s.Chain.GetLatestSeqno())

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_WRITER {
		t.Errorf("role: %s, expected WRITER", otherRole)
	}

	if err := RemoveMember(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx = context.Background()
	s, err = Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	t.Logf("after remove, seqno: %d", s.Chain.GetLatestSeqno())

	role = uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole = usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_NONE {
		t.Errorf("role: %s, expected NONE", otherRole)
	}
}

func TestMemberChangeRole(t *testing.T) {
	tc, owner, other, name := memberSetupMultiple(t)
	defer tc.Cleanup()

	if err := SetRoleWriter(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx := context.Background()
	s, err := Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role := uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole := usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_WRITER {
		t.Errorf("role: %s, expected WRITER", otherRole)
	}

	if err := SetRoleReader(context.TODO(), tc.G, name, other.Username); err != nil {
		t.Fatal(err)
	}

	ctx = context.Background()
	s, err = Get(ctx, tc.G, name)
	if err != nil {
		t.Fatal(err)
	}

	role = uidRole(ctx, tc, s, owner.User.GetUID())
	if role != keybase1.TeamRole_OWNER {
		t.Errorf("role: %s, expected OWNER", role)
	}

	otherRole = usernameRole(ctx, tc, s, other.Username)
	if otherRole != keybase1.TeamRole_READER {
		t.Errorf("role: %s, expected READER", otherRole)
	}
}

func uidRole(ctx context.Context, tc libkb.TestContext, team *Team, uid keybase1.UID) keybase1.TeamRole {
	uv, err := loadUserVersionByUID(ctx, tc.G, uid)
	if err != nil {
		tc.T.Fatal(err)
	}
	return uvRole(tc, team, uv)
}

func usernameRole(ctx context.Context, tc libkb.TestContext, team *Team, username string) keybase1.TeamRole {
	uv, err := loadUserVersionByUsername(ctx, tc.G, username)
	if err != nil {
		tc.T.Fatal(err)
	}
	return uvRole(tc, team, uv)
}

func uvRole(tc libkb.TestContext, team *Team, uv UserVersion) keybase1.TeamRole {
	role, err := team.Chain.GetUserRole(uv)
	if err != nil {
		tc.T.Fatal(err)
	}
	return role
}
<|endoftext|>


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

package main

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"
)

var (
	confFilename = flag.String("config", "/etc/testimony.conf", "Testimony config")
	logToSyslog  = flag.Bool("syslog", true, "log messages to syslog")
)

type Testimony []SocketConfig

const protocolVersion = 1

type SocketConfig struct {
	SocketName         string
	Interface          string
	BlockSize          int
	NumBlocks          int
	BlockTimeoutMillis int
	FanoutType         int
	FanoutSize         int
	User               string
	Filter             string
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true
		var socks []*Socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}
		os.Remove(sc.SocketName) // ignore errors
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		}
		if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {}
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	v(1, "chowning %q to %d", sc.SocketName, uid)
	if err := syscall.Chown(sc.SocketName, uid, 0); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*Socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*Socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	log.Printf("Received new connection %v", c.RemoteAddr())
	var buf [13]byte
	buf[0] = protocolVersion
	binary.BigEndian.PutUint32(buf[1:], uint32(len(socks)))
	binary.BigEndian.PutUint32(buf[5:], uint32(socks[0].conf.BlockSize))
	binary.BigEndian.PutUint32(buf[9:], uint32(socks[0].conf.NumBlocks))
	if _, err := c.Write(buf[:]); err != nil {
		log.Printf("new conn failed to write version: %v", err)
		return
	}
	var fanoutMsg [4]byte
	if n, err := c.Read(fanoutMsg[:]); n != len(fanoutMsg) || err != nil {
		log.Printf("new conn failed to read conf: %v", err)
		return
	}
	idx := int(binary.BigEndian.Uint32(fanoutMsg[:]))
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn invalid index %v", idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn failed to send file descriptor: %v", err)
		return
	}
	v(2, "new conn spun up, passing off to socket")
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}

func main() {
	flag.Parse()
	if *logToSyslog {
		s, err := syslog.New(syslog.LOG_USER|syslog.LOG_INFO, "testimonyd")
		if err != nil {
			log.Fatalf("could not set up syslog logging: %v", err)
		}
		log.SetOutput(s)
	}
	log.Printf("Starting testimonyd...")
	confdata, err := ioutil.ReadFile(*confFilename)
	if err != nil {
		log.Fatalf("could not read configuration %q: %v", *confFilename, err)
	}
	// Set umask which will affect all of the sockets we create:
	syscall.Umask(0177)
	var t Testimony
	if err := json.NewDecoder(bytes.NewBuffer(confdata)).Decode(&t); err != nil {
		log.Fatalf("could not parse configuration %q: %v", *confFilename, err)
	}
	RunTestimony(t)
}
<commit_msg>Fix syslog import.<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"log/syslog"
	"net"
	"os"
	"os/user"
	"strconv"
	"syscall"
)

var (
	confFilename = flag.String("config", "/etc/testimony.conf", "Testimony config")
	logToSyslog  = flag.Bool("syslog", true, "log messages to syslog")
)

type Testimony []SocketConfig

const protocolVersion = 1

type SocketConfig struct {
	SocketName         string
	Interface          string
	BlockSize          int
	NumBlocks          int
	BlockTimeoutMillis int
	FanoutType         int
	FanoutSize         int
	User               string
	Filter             string
}

func (s SocketConfig) uid() (int, error) {
	var u *user.User
	var err error
	if s.User == "" {
		u, err = user.Current()
	} else {
		u, err = user.Lookup(s.User)
	}
	if err != nil {
		return 0, fmt.Errorf("could not get user: %v", err)
	}
	return strconv.Atoi(u.Uid)
}

func RunTestimony(t Testimony) {
	fanoutID := 0
	names := map[string]bool{}
	for _, sc := range t {
		if names[sc.SocketName] {
			log.Fatalf("invalid config: duplicate socket name %q", sc.SocketName)
		}
		names[sc.SocketName] = true
		var socks []*Socket
		fanoutID++
		for i := 0; i < sc.FanoutSize; i++ {
			sock, err := newSocket(sc, fanoutID, i)
			if err != nil {
				log.Fatalf("invalid config %+v: %v", sc, err)
			}
			socks = append(socks, sock)
			go sock.run()
		}
		os.Remove(sc.SocketName) // ignore errors
		list, err := net.ListenUnix("unix", &net.UnixAddr{Net: "unix", Name: sc.SocketName})
		if err != nil {
			log.Fatalf("failed to listen on socket: %v", err)
		}
		if err := setPermissions(sc); err != nil {
			log.Fatalf("failed to set socket permissions: %v", err)
		}
		go t.run(list, sc, socks)
	}
	// We'd love to drop privs here, but thanks to
	// https://github.com/golang/go/issues/1435 we can't :(
	select {}
}

func setPermissions(sc SocketConfig) error {
	uid, err := sc.uid()
	if err != nil {
		return fmt.Errorf("could not get uid to change to: %v", err)
	}
	v(1, "chowning %q to %d", sc.SocketName, uid)
	if err := syscall.Chown(sc.SocketName, uid, 0); err != nil {
		return fmt.Errorf("unable to chown to (%d, 0): %v", uid, err)
	}
	return nil
}

func (t Testimony) run(list *net.UnixListener, sc SocketConfig, socks []*Socket) {
	for {
		c, err := list.AcceptUnix()
		if err != nil {
			log.Fatalf("failed to accept connection: %v", err)
		}
		go t.handle(socks, c)
	}
}

func (t Testimony) handle(socks []*Socket, c *net.UnixConn) {
	defer func() {
		if c != nil {
			c.Close()
		}
	}()
	log.Printf("Received new connection %v", c.RemoteAddr())
	var buf [13]byte
	buf[0] = protocolVersion
	binary.BigEndian.PutUint32(buf[1:], uint32(len(socks)))
	binary.BigEndian.PutUint32(buf[5:], uint32(socks[0].conf.BlockSize))
	binary.BigEndian.PutUint32(buf[9:], uint32(socks[0].conf.NumBlocks))
	if _, err := c.Write(buf[:]); err != nil {
		log.Printf("new conn failed to write version: %v", err)
		return
	}
	var fanoutMsg [4]byte
	if n, err := c.Read(fanoutMsg[:]); n != len(fanoutMsg) || err != nil {
		log.Printf("new conn failed to read conf: %v", err)
		return
	}
	idx := int(binary.BigEndian.Uint32(fanoutMsg[:]))
	if idx < 0 || idx >= len(socks) {
		log.Printf("new conn invalid index %v", idx)
		return
	}
	sock := socks[idx]
	fdMsg := syscall.UnixRights(sock.fd)
	var msg [1]byte // dummy byte
	n, n2, err := c.WriteMsgUnix(
		msg[:], fdMsg, nil)
	if err != nil || n != len(msg) || n2 != len(fdMsg) {
		log.Printf("new conn failed to send file descriptor: %v", err)
		return
	}
	v(2, "new conn spun up, passing off to socket")
	sock.newConns <- c
	c = nil // so it doesn't get closed by deferred func.
}

func main() {
	flag.Parse()
	if *logToSyslog {
		s, err := syslog.New(syslog.LOG_USER|syslog.LOG_INFO, "testimonyd")
		if err != nil {
			log.Fatalf("could not set up syslog logging: %v", err)
		}
		log.SetOutput(s)
	}
	log.Printf("Starting testimonyd...")
	confdata, err := ioutil.ReadFile(*confFilename)
	if err != nil {
		log.Fatalf("could not read configuration %q: %v", *confFilename, err)
	}
	// Set umask which will affect all of the sockets we create:
	syscall.Umask(0177)
	var t Testimony
	if err := json.NewDecoder(bytes.NewBuffer(confdata)).Decode(&t); err != nil {
		log.Fatalf("could not parse configuration %q: %v", *confFilename, err)
	}
	RunTestimony(t)
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/topo"
)

// Snapshot takes a tablet snapshot.
//
// forceMasterSnapshot: Normally a master is not a viable tablet to snapshot.
// However, there are degenerate cases where you need to override this, for
// instance the initial clone of a new master.
//
// serverMode: if specified, the server will stop its mysqld, and be
// ready to serve the data files directly. Slaves can just download
// these and use them directly. Call SnapshotSourceEnd to return into
// serving mode. If not specified, the server will create an archive
// of the files, store them locally, and restart.
func (wr *Wrangler) Snapshot(tabletAlias topo.TabletAlias, forceMasterSnapshot bool, snapshotConcurrency int, serverMode bool) (manifest string, parent topo.TabletAlias, slaveStartRequired, readOnly bool, originalType topo.TabletType, err error) {
	// read the tablet to be able to RPC to it, and also to get its
	// original type
	var ti *topo.TabletInfo
	ti, err = wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}
	originalType = ti.Tablet.Type

	// execute the remote action, log the results, save the error
	args := &actionnode.SnapshotArgs{
		Concurrency:         snapshotConcurrency,
		ServerMode:          serverMode,
		ForceMasterSnapshot: forceMasterSnapshot,
	}
	logStream, errFunc := wr.ai.Snapshot(ti, args, wr.ActionTimeout())
	for e := range logStream {
		log.Infof("Snapshot: %v", e)
	}
	reply, err := errFunc()

	return reply.ManifestPath, reply.ParentAlias, reply.SlaveStartRequired, reply.ReadOnly, originalType, err
}

// SnapshotSourceEnd will change the tablet back to its original type
// once it's done serving backups.
func (wr *Wrangler) SnapshotSourceEnd(tabletAlias topo.TabletAlias, slaveStartRequired, readWrite bool, originalType topo.TabletType) (err error) {
	var ti *topo.TabletInfo
	ti, err = wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}

	args := &actionnode.SnapshotSourceEndArgs{
		SlaveStartRequired: slaveStartRequired,
		ReadOnly:           !readWrite,
		OriginalType:       originalType,
	}
	return wr.ai.SnapshotSourceEnd(ti, args, wr.ActionTimeout())
}

// ReserveForRestore will make sure a tablet is ready to be used as a restore
// target.
func (wr *Wrangler) ReserveForRestore(srcTabletAlias, dstTabletAlias topo.TabletAlias) (err error) {
	// read our current tablet, verify its state before sending it
	// to the tablet itself
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	if tablet.Type != topo.TYPE_IDLE {
		return fmt.Errorf("expected idle type, not %v: %v", tablet.Type, dstTabletAlias)
	}

	args := &actionnode.ReserveForRestoreArgs{
		SrcTabletAlias: srcTabletAlias,
	}
	return wr.ai.ReserveForRestore(tablet, args, wr.ActionTimeout())
}

// UnreserveForRestore switches the tablet back to its original state,
// the restore won't happen.
func (wr *Wrangler) UnreserveForRestore(dstTabletAlias topo.TabletAlias) (err error) {
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	err = topo.DeleteTabletReplicationData(wr.ts, tablet.Tablet)
	if err != nil {
		return err
	}

	return wr.ChangeType(tablet.Alias, topo.TYPE_IDLE, false)
}

// Restore actually performs the restore action on a tablet.
func (wr *Wrangler) Restore(srcTabletAlias topo.TabletAlias, srcFilePath string, dstTabletAlias, parentAlias topo.TabletAlias, fetchConcurrency, fetchRetryCount int, wasReserved, dontWaitForSlaveStart bool) error {
	// read our current tablet, verify its state before sending it
	// to the tablet itself
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	if wasReserved {
		if tablet.Type != topo.TYPE_RESTORE {
			return fmt.Errorf("expected restore type, not %v: %v", tablet.Type, dstTabletAlias)
		}
	} else {
		if tablet.Type != topo.TYPE_IDLE {
			return fmt.Errorf("expected idle type, not %v: %v", tablet.Type, dstTabletAlias)
		}
	}

	// update the shard record if we need to, to update Cells
	srcTablet, err := wr.ts.GetTablet(srcTabletAlias)
	if err != nil {
		return err
	}
	si, err := wr.ts.GetShard(srcTablet.Keyspace, srcTablet.Shard)
	if err != nil {
		return fmt.Errorf("Cannot read shard: %v", err)
	}
	if err := wr.updateShardCellsAndMaster(si, tablet.Alias, topo.TYPE_SPARE, false); err != nil {
		return err
	}

	// do the work
	args := &actionnode.RestoreArgs{
		SrcTabletAlias:        srcTabletAlias,
		SrcFilePath:           srcFilePath,
		ParentAlias:           parentAlias,
		FetchConcurrency:      fetchConcurrency,
		FetchRetryCount:       fetchRetryCount,
		WasReserved:           wasReserved,
		DontWaitForSlaveStart: dontWaitForSlaveStart,
	}
	logStream, errFunc := wr.ai.Restore(tablet, args, wr.ActionTimeout())
	for e := range logStream {
		log.Infof("Restore: %v", e)
	}
	if err := errFunc(); err != nil {
		return err
	}

	// Restore moves us into the replication graph as a
	// spare. There are no consequences to the replication or
	// serving graphs, so no rebuild required.

	return nil
}

// UnreserveForRestoreMulti calls UnreserveForRestore on all targets.
func (wr *Wrangler) UnreserveForRestoreMulti(dstTabletAliases []topo.TabletAlias) {
	for _, dstTabletAlias := range dstTabletAliases {
		ufrErr := wr.UnreserveForRestore(dstTabletAlias)
		if ufrErr != nil {
			log.Errorf("Failed to UnreserveForRestore destination tablet after failed source snapshot: %v", ufrErr)
		} else {
			log.Infof("Un-reserved %v", dstTabletAlias)
		}
	}
}

// Clone will do all the necessary actions to copy all the data from a
// source to a set of destinations.
func (wr *Wrangler) Clone(srcTabletAlias topo.TabletAlias, dstTabletAliases []topo.TabletAlias, forceMasterSnapshot bool, snapshotConcurrency, fetchConcurrency, fetchRetryCount int, serverMode bool) error {
	// make sure the destination can be restored into (otherwise
	// there is no point in taking the snapshot in the first place),
	// and reserve it.
	reserved := make([]topo.TabletAlias, 0, len(dstTabletAliases))
	for _, dstTabletAlias := range dstTabletAliases {
		err := wr.ReserveForRestore(srcTabletAlias, dstTabletAlias)
		if err != nil {
			wr.UnreserveForRestoreMulti(reserved)
			return err
		}
		reserved = append(reserved, dstTabletAlias)
		log.Infof("Successfully reserved %v for restore", dstTabletAlias)
	}

	// take the snapshot, or put the server in SnapshotSource mode
	srcFilePath, parentAlias, slaveStartRequired, readWrite, originalType, err := wr.Snapshot(srcTabletAlias, forceMasterSnapshot, snapshotConcurrency, serverMode)
	if err != nil {
		// The snapshot failed so un-reserve the destinations
		wr.UnreserveForRestoreMulti(reserved)
	} else {
		// try to restore the snapshot
		// In serverMode, and in the case where we're replicating from
		// the master, we can't wait for replication, as the master is down.
		wg := sync.WaitGroup{}
		er := concurrency.FirstErrorRecorder{}
		for _, dstTabletAlias := range dstTabletAliases {
			wg.Add(1)
			go func(dstTabletAlias topo.TabletAlias) {
				e := wr.Restore(srcTabletAlias, srcFilePath, dstTabletAlias, parentAlias, fetchConcurrency, fetchRetryCount, true, serverMode && originalType == topo.TYPE_MASTER)
				er.RecordError(e)
				wg.Done()
			}(dstTabletAlias)
		}
		wg.Wait()
		err = er.Error()
	}

	// in any case, fix the server
	if serverMode {
		resetErr := wr.SnapshotSourceEnd(srcTabletAlias, slaveStartRequired, readWrite, originalType)
		if resetErr != nil {
			if err == nil {
				// If there is no other error, this matters.
				err = resetErr
			} else {
				// In the context of a larger failure, just log a note to cleanup.
				log.Errorf("Failed to reset snapshot source: %v - vtctl SnapshotSourceEnd is required", resetErr)
			}
		}
	}

	return err
}
<commit_msg>Handling an error case better, some RPC systems may return nil.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/topo"
)

// Snapshot takes a tablet snapshot.
//
// forceMasterSnapshot: Normally a master is not a viable tablet to snapshot.
// However, there are degenerate cases where you need to override this, for
// instance the initial clone of a new master.
//
// serverMode: if specified, the server will stop its mysqld, and be
// ready to serve the data files directly. Slaves can just download
// these and use them directly. Call SnapshotSourceEnd to return into
// serving mode. If not specified, the server will create an archive
// of the files, store them locally, and restart.
func (wr *Wrangler) Snapshot(tabletAlias topo.TabletAlias, forceMasterSnapshot bool, snapshotConcurrency int, serverMode bool) (manifest string, parent topo.TabletAlias, slaveStartRequired, readOnly bool, originalType topo.TabletType, err error) {
	// read the tablet to be able to RPC to it, and also to get its
	// original type
	var ti *topo.TabletInfo
	ti, err = wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}
	originalType = ti.Tablet.Type

	// execute the remote action, log the results, save the error
	args := &actionnode.SnapshotArgs{
		Concurrency:         snapshotConcurrency,
		ServerMode:          serverMode,
		ForceMasterSnapshot: forceMasterSnapshot,
	}
	logStream, errFunc := wr.ai.Snapshot(ti, args, wr.ActionTimeout())
	for e := range logStream {
		log.Infof("Snapshot: %v", e)
	}
	reply, err := errFunc()
	if err != nil {
		return "", topo.TabletAlias{}, false, false, "", err
	}

	return reply.ManifestPath, reply.ParentAlias, reply.SlaveStartRequired, reply.ReadOnly, originalType, nil
}

// SnapshotSourceEnd will change the tablet back to its original type
// once it's done serving backups.
func (wr *Wrangler) SnapshotSourceEnd(tabletAlias topo.TabletAlias, slaveStartRequired, readWrite bool, originalType topo.TabletType) (err error) {
	var ti *topo.TabletInfo
	ti, err = wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}

	args := &actionnode.SnapshotSourceEndArgs{
		SlaveStartRequired: slaveStartRequired,
		ReadOnly:           !readWrite,
		OriginalType:       originalType,
	}
	return wr.ai.SnapshotSourceEnd(ti, args, wr.ActionTimeout())
}

// ReserveForRestore will make sure a tablet is ready to be used as a restore
// target.
func (wr *Wrangler) ReserveForRestore(srcTabletAlias, dstTabletAlias topo.TabletAlias) (err error) {
	// read our current tablet, verify its state before sending it
	// to the tablet itself
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	if tablet.Type != topo.TYPE_IDLE {
		return fmt.Errorf("expected idle type, not %v: %v", tablet.Type, dstTabletAlias)
	}

	args := &actionnode.ReserveForRestoreArgs{
		SrcTabletAlias: srcTabletAlias,
	}
	return wr.ai.ReserveForRestore(tablet, args, wr.ActionTimeout())
}

// UnreserveForRestore switches the tablet back to its original state,
// the restore won't happen.
func (wr *Wrangler) UnreserveForRestore(dstTabletAlias topo.TabletAlias) (err error) {
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	err = topo.DeleteTabletReplicationData(wr.ts, tablet.Tablet)
	if err != nil {
		return err
	}

	return wr.ChangeType(tablet.Alias, topo.TYPE_IDLE, false)
}

// Restore actually performs the restore action on a tablet.
func (wr *Wrangler) Restore(srcTabletAlias topo.TabletAlias, srcFilePath string, dstTabletAlias, parentAlias topo.TabletAlias, fetchConcurrency, fetchRetryCount int, wasReserved, dontWaitForSlaveStart bool) error {
	// read our current tablet, verify its state before sending it
	// to the tablet itself
	tablet, err := wr.ts.GetTablet(dstTabletAlias)
	if err != nil {
		return err
	}
	if wasReserved {
		if tablet.Type != topo.TYPE_RESTORE {
			return fmt.Errorf("expected restore type, not %v: %v", tablet.Type, dstTabletAlias)
		}
	} else {
		if tablet.Type != topo.TYPE_IDLE {
			return fmt.Errorf("expected idle type, not %v: %v", tablet.Type, dstTabletAlias)
		}
	}

	// update the shard record if we need to, to update Cells
	srcTablet, err := wr.ts.GetTablet(srcTabletAlias)
	if err != nil {
		return err
	}
	si, err := wr.ts.GetShard(srcTablet.Keyspace, srcTablet.Shard)
	if err != nil {
		return fmt.Errorf("Cannot read shard: %v", err)
	}
	if err := wr.updateShardCellsAndMaster(si, tablet.Alias, topo.TYPE_SPARE, false); err != nil {
		return err
	}

	// do the work
	args := &actionnode.RestoreArgs{
		SrcTabletAlias:        srcTabletAlias,
		SrcFilePath:           srcFilePath,
		ParentAlias:           parentAlias,
		FetchConcurrency:      fetchConcurrency,
		FetchRetryCount:       fetchRetryCount,
		WasReserved:           wasReserved,
		DontWaitForSlaveStart: dontWaitForSlaveStart,
	}
	logStream, errFunc := wr.ai.Restore(tablet, args, wr.ActionTimeout())
	for e := range logStream {
		log.Infof("Restore: %v", e)
	}
	if err := errFunc(); err != nil {
		return err
	}

	// Restore moves us into the replication graph as a
	// spare. There are no consequences to the replication or
	// serving graphs, so no rebuild required.

	return nil
}

// UnreserveForRestoreMulti calls UnreserveForRestore on all targets.
func (wr *Wrangler) UnreserveForRestoreMulti(dstTabletAliases []topo.TabletAlias) {
	for _, dstTabletAlias := range dstTabletAliases {
		ufrErr := wr.UnreserveForRestore(dstTabletAlias)
		if ufrErr != nil {
			log.Errorf("Failed to UnreserveForRestore destination tablet after failed source snapshot: %v", ufrErr)
		} else {
			log.Infof("Un-reserved %v", dstTabletAlias)
		}
	}
}

// Clone will do all the necessary actions to copy all the data from a
// source to a set of destinations.
func (wr *Wrangler) Clone(srcTabletAlias topo.TabletAlias, dstTabletAliases []topo.TabletAlias, forceMasterSnapshot bool, snapshotConcurrency, fetchConcurrency, fetchRetryCount int, serverMode bool) error {
	// make sure the destination can be restored into (otherwise
	// there is no point in taking the snapshot in the first place),
	// and reserve it.
	reserved := make([]topo.TabletAlias, 0, len(dstTabletAliases))
	for _, dstTabletAlias := range dstTabletAliases {
		err := wr.ReserveForRestore(srcTabletAlias, dstTabletAlias)
		if err != nil {
			wr.UnreserveForRestoreMulti(reserved)
			return err
		}
		reserved = append(reserved, dstTabletAlias)
		log.Infof("Successfully reserved %v for restore", dstTabletAlias)
	}

	// take the snapshot, or put the server in SnapshotSource mode
	srcFilePath, parentAlias, slaveStartRequired, readWrite, originalType, err := wr.Snapshot(srcTabletAlias, forceMasterSnapshot, snapshotConcurrency, serverMode)
	if err != nil {
		// The snapshot failed so un-reserve the destinations
		wr.UnreserveForRestoreMulti(reserved)
	} else {
		// try to restore the snapshot
		// In serverMode, and in the case where we're replicating from
		// the master, we can't wait for replication, as the master is down.
		wg := sync.WaitGroup{}
		er := concurrency.FirstErrorRecorder{}
		for _, dstTabletAlias := range dstTabletAliases {
			wg.Add(1)
			go func(dstTabletAlias topo.TabletAlias) {
				e := wr.Restore(srcTabletAlias, srcFilePath, dstTabletAlias, parentAlias, fetchConcurrency, fetchRetryCount, true, serverMode && originalType == topo.TYPE_MASTER)
				er.RecordError(e)
				wg.Done()
			}(dstTabletAlias)
		}
		wg.Wait()
		err = er.Error()
	}

	// in any case, fix the server
	if serverMode {
		resetErr := wr.SnapshotSourceEnd(srcTabletAlias, slaveStartRequired, readWrite, originalType)
		if resetErr != nil {
			if err == nil {
				// If there is no other error, this matters.
				err = resetErr
			} else {
				// In the context of a larger failure, just log a note to cleanup.
				log.Errorf("Failed to reset snapshot source: %v - vtctl SnapshotSourceEnd is required", resetErr)
			}
		}
	}

	return err
}
<|endoftext|>


<commit_before>package goldprice

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const debug = true

const urlBankYear = "http://rate.bot.com.tw/Pages/UIP005/UIP005INQ3.aspx?view=1&lang=zh-TW"

const viewStatKey = "__VIEWSTATE"
const viewStat = "/wEPDwUKLTc3OTg4NTEyM2QYAgUeX19Db250cm9sc1JlcXVpcmVQb3N0QmFja0tleV9fFgoFBlJhZGlvNQUGUmFkaW8xBQZSYWRpbzIFBlJhZGlvMwUGUmFkaW80BQVjdGwwMgUFY3RsMDMFBWN0bDA0BQVjdGwwNQUFY3RsMDYFCW11bHRpVGFicw8PZAIBZPZJVCxGUh1sKGSIy7aqPTTqBGsH"
const validationKey = "__EVENTVALIDATION"
const validation = "/wEWEQKouumRBAKOkoCCCwKMhc76CQLWssLjAgKf5L6bDQLL3LrjCgLW3JbjCgLS3JbjCgLR3JbjCgLM3JbjCgLl18nSCwLm1/nSCwL9/oqMCQKLts3CAQKUts3CAQKM54rGBgLWlM+bAvA8FJroJ9FZZI52UscKHadSwtMt"

// form input parameter for date
const (
	dateParam   = "term"
	recentDay   = 99
	threeMonth  = 6
	halfYear    = 2
	year        = 3
	specifyDate = 0
	monthParam  = "month"
	yearParam   = "year"
)

// form input parameter for current type
const (
	currentParam = "curcd"
	currentTWN   = "TWD"
	currentUSD   = "USD"
	currentCNY   = "CNY"
)

// form input parameter for when in a day
const (
	whenParam = "afterOrNot"
	before    = 0
	after     = 1
)

// Date 1234/5/6
type Date struct {
	Year, Month, Day int
}

// Time 23:59
type Time struct {
	hour, minute int
}

// Price of buy and sell
type Price struct {
	buy, sell int
}

// GetTaiwanBankGoldPriceYear get whole year gold price from taiwan bank
func GetYearFromTaiwanBank() (dateArray []Date, ret map[Date]Price) {
	var term int //
	var y int
	var m int
	var curcd string // should be TWD
	var when int
	ret = make(map[Date]Price)
	dateArray = make([]Date, 0)

	now := time.Now()
	term = year
	y = now.Year()
	m = int(now.Month())
	curcd = currentTWN
	when = before

	resp, err := http.PostForm(urlBankYear,
		url.Values{
			"Button1":     {"查詢"}, // I don't know what the fuck is this
			viewStatKey:   {viewStat},
			validationKey: {validation},
			dateParam:     {strconv.Itoa(term)},
			yearParam:     {strconv.Itoa(y)},
			monthParam:    {strconv.Itoa(m)},
			currentParam:  {curcd},
			whenParam:     {strconv.Itoa(when)}})
	if err != nil {
		fmt.Printf("%v\n", err)
		return
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	html := strings.Split(string(body), "\n")

	for _, line := range html {
		// if strings.Contains(line, "class=\"color0\"") || strings.Contains(line, "class=\"color1\"") {
		r := regexp.MustCompile(`date=(\d{8}).+class="decimal">(\d+)</td><td class="decimal">(\d+)`)
		res := r.FindStringSubmatch(line)
		if res == nil {
			continue
		}
		tmpDate, _ := strconv.Atoi(res[1])
		date := Date{tmpDate / 10000, (tmpDate % 10000) / 100, tmpDate % 100}
		buy, _ := strconv.Atoi(res[2])
		sell, _ := strconv.Atoi(res[3])
		price := Price{buy, sell}
		ret[date] = price
		dateArray = append(dateArray, date)
	}

	return
	// TODO return map of a year
}

const urlBankDay = "http://rate.bot.com.tw/Pages/UIP005/UIP00511.aspx"

// GetTaiwanBankGoldPriceDay get specifiy date gold price from taiwan bank
func GetDayFromTaiwanBank(date Date) (timeArray []Time, ret map[Time]Price) {
	ret = make(map[Time]Price)
	timeArray = make([]Time, 0)
	dateString := fmt.Sprintf("%d%02d%02d", date.Year, date.Month, date.Day)

	url := urlBankDay +
		"?" +
		"&lang=zh-TW" +
		"&whom=GB0030001000" + // don't know what
		"&date=" + dateString +
		"&afterOrNot=" + strconv.Itoa(before) +
		"&curcd=" + currentTWN

	resp, err := http.Get(url)
	if err != nil {
		fmt.Printf("%v\n", err)
		return
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	html := strings.Split(string(body), "\n")
	for _, line := range html {
		r := regexp.MustCompile(`">(\d{2}):(\d{2})<.+class="decimal">(\d+)</td><td class="decimal">(\d+)`)
		res := r.FindStringSubmatch(line)
		if res == nil {
			continue
		}
		hour, _ := strconv.Atoi(res[1])
		min, _ := strconv.Atoi(res[2])
		buy, _ := strconv.Atoi(res[3])
		sell, _ := strconv.Atoi(res[4])
		eventTime := Time{hour, min}
		price := Price{buy, sell}
		ret[eventTime] = price
		timeArray = append(timeArray, eventTime)
	}

	return
}
<commit_msg>fix comment warning<commit_after>package goldprice

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const debug = true

const urlBankYear = "http://rate.bot.com.tw/Pages/UIP005/UIP005INQ3.aspx?view=1&lang=zh-TW"

const viewStatKey = "__VIEWSTATE"
const viewStat = "/wEPDwUKLTc3OTg4NTEyM2QYAgUeX19Db250cm9sc1JlcXVpcmVQb3N0QmFja0tleV9fFgoFBlJhZGlvNQUGUmFkaW8xBQZSYWRpbzIFBlJhZGlvMwUGUmFkaW80BQVjdGwwMgUFY3RsMDMFBWN0bDA0BQVjdGwwNQUFY3RsMDYFCW11bHRpVGFicw8PZAIBZPZJVCxGUh1sKGSIy7aqPTTqBGsH"
const validationKey = "__EVENTVALIDATION"
const validation = "/wEWEQKouumRBAKOkoCCCwKMhc76CQLWssLjAgKf5L6bDQLL3LrjCgLW3JbjCgLS3JbjCgLR3JbjCgLM3JbjCgLl18nSCwLm1/nSCwL9/oqMCQKLts3CAQKUts3CAQKM54rGBgLWlM+bAvA8FJroJ9FZZI52UscKHadSwtMt"

// form input parameter for date
const (
	dateParam   = "term"
	recentDay   = 99
	threeMonth  = 6
	halfYear    = 2
	year        = 3
	specifyDate = 0
	monthParam  = "month"
	yearParam   = "year"
)

// form input parameter for current type
const (
	currentParam = "curcd"
	currentTWN   = "TWD"
	currentUSD   = "USD"
	currentCNY   = "CNY"
)

// form input parameter for when in a day
const (
	whenParam = "afterOrNot"
	before    = 0
	after     = 1
)

// Date 1234/5/6
type Date struct {
	Year, Month, Day int
}

// Time 23:59
type Time struct {
	hour, minute int
}

// Price of buy and sell
type Price struct {
	buy, sell int
}

// GetYearFromTaiwanBank get whole year gold price from taiwan bank
func GetYearFromTaiwanBank() (dateArray []Date, ret map[Date]Price) {
	var term int //
	var y int
	var m int
	var curcd string // should be TWD
	var when int
	ret = make(map[Date]Price)
	dateArray = make([]Date, 0)

	now := time.Now()
	term = year
	y = now.Year()
	m = int(now.Month())
	curcd = currentTWN
	when = before

	resp, err := http.PostForm(urlBankYear,
		url.Values{
			"Button1":     {"查詢"}, // I don't know what the fuck is this
			viewStatKey:   {viewStat},
			validationKey: {validation},
			dateParam:     {strconv.Itoa(term)},
			yearParam:     {strconv.Itoa(y)},
			monthParam:    {strconv.Itoa(m)},
			currentParam:  {curcd},
			whenParam:     {strconv.Itoa(when)}})
	if err != nil {
		fmt.Printf("%v\n", err)
		return
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	html := strings.Split(string(body), "\n")

	for _, line := range html {
		// if strings.Contains(line, "class=\"color0\"") || strings.Contains(line, "class=\"color1\"") {
		r := regexp.MustCompile(`date=(\d{8}).+class="decimal">(\d+)</td><td class="decimal">(\d+)`)
		res := r.FindStringSubmatch(line)
		if res == nil {
			continue
		}
		tmpDate, _ := strconv.Atoi(res[1])
		date := Date{tmpDate / 10000, (tmpDate % 10000) / 100, tmpDate % 100}
		buy, _ := strconv.Atoi(res[2])
		sell, _ := strconv.Atoi(res[3])
		price := Price{buy, sell}
		ret[date] = price
		dateArray = append(dateArray, date)
	}

	return
	// TODO return map of a year
}

const urlBankDay = "http://rate.bot.com.tw/Pages/UIP005/UIP00511.aspx"

// GetDayFromTaiwanBank get specifiy date gold price from taiwan bank
func GetDayFromTaiwanBank(date Date) (timeArray []Time, ret map[Time]Price) {
	ret = make(map[Time]Price)
	timeArray = make([]Time, 0)
	dateString := fmt.Sprintf("%d%02d%02d", date.Year, date.Month, date.Day)

	url := urlBankDay +
		"?" +
		"&lang=zh-TW" +
		"&whom=GB0030001000" + // don't know what
		"&date=" + dateString +
		"&afterOrNot=" + strconv.Itoa(before) +
		"&curcd=" + currentTWN

	resp, err := http.Get(url)
	if err != nil {
		fmt.Printf("%v\n", err)
		return
	}
	defer resp.Body.Close()

	body, _ := ioutil.ReadAll(resp.Body)

	html := strings.Split(string(body), "\n")
	for _, line := range html {
		r := regexp.MustCompile(`">(\d{2}):(\d{2})<.+class="decimal">(\d+)</td><td class="decimal">(\d+)`)
		res := r.FindStringSubmatch(line)
		if res == nil {
			continue
		}
		hour, _ := strconv.Atoi(res[1])
		min, _ := strconv.Atoi(res[2])
		buy, _ := strconv.Atoi(res[3])
		sell, _ := strconv.Atoi(res[4])
		eventTime := Time{hour, min}
		price := Price{buy, sell}
		ret[eventTime] = price
		timeArray = append(timeArray, eventTime)
	}

	return
}
<|endoftext|>


<commit_before>package ircclient

import (
	"strings"

	"github.com/goshuirc/irc-go/ircmsg"
)

/**
 * ClientCaps holds the capabilities between the client and server
 */
type ClientCaps struct {
	Wanted    []string
	Enabled   map[string]string
	Available map[string]string
}

// CommonCaps returns a slice of caps that both the client and server support
func (caps *ClientCaps) CommonCaps() []string {
	var common []string

	for _, wantedCap := range caps.Wanted {
		_, exists := caps.Available[wantedCap]
		if exists {
			common = append(common, wantedCap)
		}
	}

	return common
}

/**
 * Client is the IRC client
 */
type Client struct {
	Socket
	Nick             string
	Username         string
	Realname         string
	Password         string
	BindHost         string
	Caps             *ClientCaps
	Supported        map[string]string
	HasRegistered    bool
	CommandListeners map[string][]func(*ircmsg.IrcMessage)
}

func NewClient() *Client {
	client := &Client{
		Socket:           *NewSocket(),
		Supported:        make(map[string]string),
		CommandListeners: make(map[string][]func(*ircmsg.IrcMessage)),
	}

	client.Caps = &ClientCaps{
		Enabled:   make(map[string]string),
		Available: make(map[string]string),
	}

	return client
}

func (client *Client) Connect() error {
	err := client.Socket.Connect()
	if err != nil {
		return err
	}

	go client.messageDispatcher()

	client.WriteLine("CAP LS 302")
	client.WriteLine("NICK %s", client.Nick)
	client.WriteLine("USER %s 0 * :%s", client.Username, client.Realname)

	return nil
}

func (client *Client) HandleCommand(command string, fn func(*ircmsg.IrcMessage)) {
	command = strings.ToUpper(command)

	ar, _ := client.CommandListeners[command]
	if ar == nil {
		ar = make([]func(*ircmsg.IrcMessage), 0)
		client.CommandListeners[command] = ar
	}

	client.CommandListeners[command] = append(client.CommandListeners[command], fn)
}

func (client *Client) messageDispatcher() {
	var handlers []func(*ircmsg.IrcMessage)

	for {
		message, isOK := <-client.MessagesIn
		if !isOK {
			break
		}

		// Run our internal command handlers first
		command, commandExists := ServerCommands[message.Command]
		if commandExists {
			command.Run(client, &message)
		}

		// Dispatch any command handler
		handlers, _ = client.CommandListeners[strings.ToUpper(message.Command)]
		if handlers != nil {
			for _, handler := range handlers {
				handler(&message)
			}
		}

		// Dispatch any ALL handlers
		handlers, _ = client.CommandListeners["ALL"]
		if handlers != nil {
			for _, handler := range handlers {
				handler(&message)
			}
		}
	}

	handlers, _ = client.CommandListeners["CLOSED"]
	if handlers != nil {
		for _, handler := range handlers {
			handler(nil)
		}
	}

	client.HasRegistered = false
}

func (client *Client) JoinChannel(channel string, key string) {
	if client.Connected {
		client.WriteLine("JOIN %s %s", channel, key)
	}
}
<commit_msg>Lets request some CAPS. Probably too much right now, just taken from the old irc lib<commit_after>package ircclient

import (
	"strings"

	"github.com/goshuirc/irc-go/ircmsg"
)

/**
 * ClientCaps holds the capabilities between the client and server
 */
type ClientCaps struct {
	Wanted    []string
	Enabled   map[string]string
	Available map[string]string
}

// CommonCaps returns a slice of caps that both the client and server support
func (caps *ClientCaps) CommonCaps() []string {
	var common []string

	for _, wantedCap := range caps.Wanted {
		_, exists := caps.Available[wantedCap]
		if exists {
			common = append(common, wantedCap)
		}
	}

	return common
}

// IsEnabled checks if a particular cap is enabled for this connection
func (caps *ClientCaps) IsEnabled(cap string) bool {
	_, exists := caps.Enabled[cap]
	return exists
}

/**
 * Client is the IRC client
 */
type Client struct {
	Socket
	Nick             string
	Username         string
	Realname         string
	Password         string
	BindHost         string
	Caps             *ClientCaps
	Supported        map[string]string
	HasRegistered    bool
	CommandListeners map[string][]func(*ircmsg.IrcMessage)
}

func NewClient() *Client {
	client := &Client{
		Socket:           *NewSocket(),
		Supported:        make(map[string]string),
		CommandListeners: make(map[string][]func(*ircmsg.IrcMessage)),
	}

	client.Caps = &ClientCaps{
		Enabled:   make(map[string]string),
		Available: make(map[string]string),
	}

	client.Caps.Wanted = append(
		client.Caps.Wanted,
		"account-notify",
		"away-notify",
		"extended-join",
		"multi-prefix",
		"sasl",
		"account-tag",
		"cap-notify",
		"chghost",
		"invite-notify",
		"server-time",
		"userhost-in-names",
	)

	return client
}

func (client *Client) Connect() error {
	err := client.Socket.Connect()
	if err != nil {
		return err
	}

	go client.messageDispatcher()

	client.WriteLine("CAP LS 302")
	client.WriteLine("NICK %s", client.Nick)
	client.WriteLine("USER %s 0 * :%s", client.Username, client.Realname)

	return nil
}

func (client *Client) HandleCommand(command string, fn func(*ircmsg.IrcMessage)) {
	command = strings.ToUpper(command)

	ar, _ := client.CommandListeners[command]
	if ar == nil {
		ar = make([]func(*ircmsg.IrcMessage), 0)
		client.CommandListeners[command] = ar
	}

	client.CommandListeners[command] = append(client.CommandListeners[command], fn)
}

func (client *Client) messageDispatcher() {
	var handlers []func(*ircmsg.IrcMessage)

	for {
		message, isOK := <-client.MessagesIn
		if !isOK {
			break
		}

		// Run our internal command handlers first
		command, commandExists := ServerCommands[message.Command]
		if commandExists {
			command.Run(client, &message)
		}

		// Dispatch any command handler
		handlers, _ = client.CommandListeners[strings.ToUpper(message.Command)]
		if handlers != nil {
			for _, handler := range handlers {
				handler(&message)
			}
		}

		// Dispatch any ALL handlers
		handlers, _ = client.CommandListeners["ALL"]
		if handlers != nil {
			for _, handler := range handlers {
				handler(&message)
			}
		}
	}

	handlers, _ = client.CommandListeners["CLOSED"]
	if handlers != nil {
		for _, handler := range handlers {
			handler(nil)
		}
	}

	client.HasRegistered = false
}

func (client *Client) JoinChannel(channel string, key string) {
	if client.Connected {
		client.WriteLine("JOIN %s %s", channel, key)
	}
}
<|endoftext|>