<commit_before>package cloudflare

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

type TeamsRuleSettings struct {
	// Enable block page on rules with action block
	BlockPageEnabled bool `json:"block_page_enabled"`

	// show this string at block page caused by this rule
	BlockReason string `json:"block_reason"`

	// list of ipv4 or ipv6 ips to override with, when action is set to dns override
	OverrideIPs []string `json:"override_ips"`

	// host name to override with when action is set to dns override. Can not be used with OverrideIPs
	OverrideHost string `json:"override_host"`

	// settings for l4(network) level overrides
	L4Override *TeamsL4OverrideSettings `json:"l4override"`

	// settings for adding headers to http requests
	AddHeaders http.Header `json:"add_headers"`

	// settings for browser isolation actions
	BISOAdminControls *TeamsBISOAdminControlSettings `json:"biso_admin_controls"`

	// settings for session check in allow action
	CheckSession *TeamsCheckSessionSettings `json:"check_session"`

	// whether to disable dnssec validation for allow action
	InsecureDisableDNSSECValidation bool `json:"insecure_disable_dnssec_validation"`
}

// TeamsL4OverrideSettings used in l4 filter type rule with action set to override.
type TeamsL4OverrideSettings struct {
	IP   string `json:"ip,omitempty"`
	Port int    `json:"port,omitempty"`
}

type TeamsBISOAdminControlSettings struct {
	DisablePrinting  bool `json:"dp"`
	DisableCopyPaste bool `json:"dcp"`
	DisableDownload  bool `json:"dd"`
	DisableUpload    bool `json:"du"`
	DisableKeyboard  bool `json:"dk"`
}

type TeamsCheckSessionSettings struct {
	Enforce  bool     `json:"enforce"`
	Duration Duration `json:"duration"`
}

type TeamsFilterType string

type TeamsGatewayAction string

const (
	HttpFilter TeamsFilterType = "http"
	DnsFilter  TeamsFilterType = "dns"
	L4Filter   TeamsFilterType = "l4"
)

const (
	Allow        TeamsGatewayAction = "allow"
	Block        TeamsGatewayAction = "block"
	SafeSearch   TeamsGatewayAction = "safesearch"
	YTRestricted TeamsGatewayAction = "ytrestricted"
	On           TeamsGatewayAction = "on"
	Off          TeamsGatewayAction = "off"
	Scan         TeamsGatewayAction = "scan"
	NoScan       TeamsGatewayAction = "noscan"
	Isolate      TeamsGatewayAction = "isolate"
	NoIsolate    TeamsGatewayAction = "noisolate"
	Override     TeamsGatewayAction = "override"
	L4Override   TeamsGatewayAction = "l4_override"
)

func TeamsRulesActionValues() []string {
	return []string{
		string(Allow),
		string(Block),
		string(SafeSearch),
		string(YTRestricted),
		string(On),
		string(Off),
		string(Scan),
		string(NoScan),
		string(Isolate),
		string(NoIsolate),
		string(Override),
		string(L4Override),
	}
}

// TeamsRule represents an Teams wirefilter rule.
type TeamsRule struct {
	ID            string             `json:"id,omitempty"`
	CreatedAt     *time.Time         `json:"created_at,omitempty"`
	UpdatedAt     *time.Time         `json:"updated_at,omitempty"`
	DeletedAt     *time.Time         `json:"deleted_at,omitempty"`
	Name          string             `json:"name"`
	Description   string             `json:"description"`
	Precedence    uint64             `json:"precedence"`
	Enabled       bool               `json:"enabled"`
	Action        TeamsGatewayAction `json:"action"`
	Filters       []TeamsFilterType  `json:"filters"`
	Traffic       string             `json:"traffic"`
	Identity      string             `json:"identity"`
	DevicePosture string             `json:"device_posture"`
	Version       uint64             `json:"version"`
	RuleSettings  TeamsRuleSettings  `json:"rule_settings,omitempty"`
}

// TeamsRuleResponse is the API response, containing a single rule.
type TeamsRuleResponse struct {
	Response
	Result TeamsRule `json:"result"`
}

// TeamsRuleResponse is the API response, containing an array of rules.
type TeamsRulesResponse struct {
	Response
	Result []TeamsRule `json:"result"`
}

// TeamsRulePatchRequest is used to patch an existing rule.
type TeamsRulePatchRequest struct {
	ID           string             `json:"id"`
	Name         string             `json:"name"`
	Description  string             `json:"description"`
	Precedence   uint64             `json:"precedence"`
	Enabled      bool               `json:"enabled"`
	Action       TeamsGatewayAction `json:"action"`
	RuleSettings TeamsRuleSettings  `json:"rule_settings,omitempty"`
}

// TeamsRules returns all rules within an account.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsRules(ctx context.Context, accountID string) ([]TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules", accountID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return []TeamsRule{}, err
	}

	var teamsRulesResponse TeamsRulesResponse
	err = json.Unmarshal(res, &teamsRulesResponse)
	if err != nil {
		return []TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRulesResponse.Result, nil
}

// TeamsRule returns the rule with rule ID in the URL.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsRule(ctx context.Context, accountID string, ruleId string) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsCreateRule creates a rule with wirefilter expression.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsCreateRule(ctx context.Context, accountID string, rule TeamsRule) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules", accountID)

	res, err := api.makeRequestContext(ctx, http.MethodPost, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsUpdateRule updates a rule with wirefilter expression.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsUpdateRule(ctx context.Context, accountID string, ruleId string, rule TeamsRule) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodPut, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsPatchRule patches a rule associated values.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsPatchRule(ctx context.Context, accountID string, ruleId string, rule TeamsRulePatchRequest) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodPatch, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsDeleteRule deletes a rule.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsDeleteRule(ctx context.Context, accountID string, ruleId string) error {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	_, err := api.makeRequestContext(ctx, http.MethodDelete, uri, nil)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Reorder teams rule setting struct<commit_after>package cloudflare

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/pkg/errors"
)

type TeamsRuleSettings struct {
	// list of ipv4 or ipv6 ips to override with, when action is set to dns override
	OverrideIPs []string `json:"override_ips"`

	// show this string at block page caused by this rule
	BlockReason string `json:"block_reason"`

	// host name to override with when action is set to dns override. Can not be used with OverrideIPs
	OverrideHost string `json:"override_host"`

	// settings for browser isolation actions
	BISOAdminControls *TeamsBISOAdminControlSettings `json:"biso_admin_controls"`

	// settings for l4(network) level overrides
	L4Override *TeamsL4OverrideSettings `json:"l4override"`

	// settings for adding headers to http requests
	AddHeaders http.Header `json:"add_headers"`

	// settings for session check in allow action
	CheckSession *TeamsCheckSessionSettings `json:"check_session"`

	// Enable block page on rules with action block
	BlockPageEnabled bool `json:"block_page_enabled"`

	// whether to disable dnssec validation for allow action
	InsecureDisableDNSSECValidation bool `json:"insecure_disable_dnssec_validation"`
}

// TeamsL4OverrideSettings used in l4 filter type rule with action set to override.
type TeamsL4OverrideSettings struct {
	IP   string `json:"ip,omitempty"`
	Port int    `json:"port,omitempty"`
}

type TeamsBISOAdminControlSettings struct {
	DisablePrinting  bool `json:"dp"`
	DisableCopyPaste bool `json:"dcp"`
	DisableDownload  bool `json:"dd"`
	DisableUpload    bool `json:"du"`
	DisableKeyboard  bool `json:"dk"`
}

type TeamsCheckSessionSettings struct {
	Enforce  bool     `json:"enforce"`
	Duration Duration `json:"duration"`
}

type TeamsFilterType string

type TeamsGatewayAction string

const (
	HttpFilter TeamsFilterType = "http"
	DnsFilter  TeamsFilterType = "dns"
	L4Filter   TeamsFilterType = "l4"
)

const (
	Allow        TeamsGatewayAction = "allow"
	Block        TeamsGatewayAction = "block"
	SafeSearch   TeamsGatewayAction = "safesearch"
	YTRestricted TeamsGatewayAction = "ytrestricted"
	On           TeamsGatewayAction = "on"
	Off          TeamsGatewayAction = "off"
	Scan         TeamsGatewayAction = "scan"
	NoScan       TeamsGatewayAction = "noscan"
	Isolate      TeamsGatewayAction = "isolate"
	NoIsolate    TeamsGatewayAction = "noisolate"
	Override     TeamsGatewayAction = "override"
	L4Override   TeamsGatewayAction = "l4_override"
)

func TeamsRulesActionValues() []string {
	return []string{
		string(Allow),
		string(Block),
		string(SafeSearch),
		string(YTRestricted),
		string(On),
		string(Off),
		string(Scan),
		string(NoScan),
		string(Isolate),
		string(NoIsolate),
		string(Override),
		string(L4Override),
	}
}

// TeamsRule represents an Teams wirefilter rule.
type TeamsRule struct {
	ID            string             `json:"id,omitempty"`
	CreatedAt     *time.Time         `json:"created_at,omitempty"`
	UpdatedAt     *time.Time         `json:"updated_at,omitempty"`
	DeletedAt     *time.Time         `json:"deleted_at,omitempty"`
	Name          string             `json:"name"`
	Description   string             `json:"description"`
	Precedence    uint64             `json:"precedence"`
	Enabled       bool               `json:"enabled"`
	Action        TeamsGatewayAction `json:"action"`
	Filters       []TeamsFilterType  `json:"filters"`
	Traffic       string             `json:"traffic"`
	Identity      string             `json:"identity"`
	DevicePosture string             `json:"device_posture"`
	Version       uint64             `json:"version"`
	RuleSettings  TeamsRuleSettings  `json:"rule_settings,omitempty"`
}

// TeamsRuleResponse is the API response, containing a single rule.
type TeamsRuleResponse struct {
	Response
	Result TeamsRule `json:"result"`
}

// TeamsRuleResponse is the API response, containing an array of rules.
type TeamsRulesResponse struct {
	Response
	Result []TeamsRule `json:"result"`
}

// TeamsRulePatchRequest is used to patch an existing rule.
type TeamsRulePatchRequest struct {
	ID           string             `json:"id"`
	Name         string             `json:"name"`
	Description  string             `json:"description"`
	Precedence   uint64             `json:"precedence"`
	Enabled      bool               `json:"enabled"`
	Action       TeamsGatewayAction `json:"action"`
	RuleSettings TeamsRuleSettings  `json:"rule_settings,omitempty"`
}

// TeamsRules returns all rules within an account.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsRules(ctx context.Context, accountID string) ([]TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules", accountID)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return []TeamsRule{}, err
	}

	var teamsRulesResponse TeamsRulesResponse
	err = json.Unmarshal(res, &teamsRulesResponse)
	if err != nil {
		return []TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRulesResponse.Result, nil
}

// TeamsRule returns the rule with rule ID in the URL.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsRule(ctx context.Context, accountID string, ruleId string) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodGet, uri, nil)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsCreateRule creates a rule with wirefilter expression.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsCreateRule(ctx context.Context, accountID string, rule TeamsRule) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules", accountID)

	res, err := api.makeRequestContext(ctx, http.MethodPost, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsUpdateRule updates a rule with wirefilter expression.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsUpdateRule(ctx context.Context, accountID string, ruleId string, rule TeamsRule) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodPut, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsPatchRule patches a rule associated values.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsPatchRule(ctx context.Context, accountID string, ruleId string, rule TeamsRulePatchRequest) (TeamsRule, error) {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	res, err := api.makeRequestContext(ctx, http.MethodPatch, uri, rule)
	if err != nil {
		return TeamsRule{}, err
	}

	var teamsRuleResponse TeamsRuleResponse
	err = json.Unmarshal(res, &teamsRuleResponse)
	if err != nil {
		return TeamsRule{}, errors.Wrap(err, errUnmarshalError)
	}

	return teamsRuleResponse.Result, nil
}

// TeamsDeleteRule deletes a rule.
//
// API reference: https://api.cloudflare.com/#teams-rules-properties
func (api *API) TeamsDeleteRule(ctx context.Context, accountID string, ruleId string) error {
	uri := fmt.Sprintf("/accounts/%s/gateway/rules/%s", accountID, ruleId)

	_, err := api.makeRequestContext(ctx, http.MethodDelete, uri, nil)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package term // import "github.com/amalog/go/term"

import (
	"io"

	"github.com/amalog/go/scanner"
)

// central definition of the term terminator
const terminator = ","

type Reader struct {
	s *scanner.Scanner
}

func NewReader(r io.RuneScanner) *Reader {
	self := new(Reader)
	self.s = scanner.New(r)
	return self
}

func (r *Reader) Read() (Term, error) {
	x, err := r.s.Scan()
	if err != nil {
		return nil, err
	}

	switch x.Class {
	case scanner.Atom:
		return r.readAtomOrStruct(nil, x)
	case scanner.String:
		if len(x.Text) < 2 {
			return nil, &Err{x, "string token too short"}
		}
		if x.Text[0] != '"' {
			return nil, &Err{x, "string missing opening quote"}
		}
		if x.Text[len(x.Text)-1] != '"' {
			return nil, &Err{x, "string missing closing quote"}
		}
		text := x.Text[1 : len(x.Text)-1]

		y, err := r.s.Scan()
		if err != nil {
			return nil, err
		}

		switch y.Class {
		case scanner.Punct:
			if y.Text == terminator {
				return NewString(text), nil
			}
		}
		return nil, &ErrUnexpectedToken{y}
	case scanner.Var:
		y, err := r.s.Scan()
		if err != nil {
			return nil, err
		}

		switch y.Class {
		case scanner.Punct:
			switch y.Text {
			case terminator:
				term := &Var{
					Name:  x.Text,
					Value: nil,
				}
				return term, nil
			case ".":
				z, err := r.s.Scan()
				if err != nil {
					return nil, err
				}
				if z.Class == scanner.Atom {
					return r.readAtomOrStruct(x, z)
				}
			}
		}
		return nil, &ErrUnexpectedToken{y}
	default:
		return nil, &ErrUnexpectedToken{x}
	}

	return nil, nil
}

func (r *Reader) readAtomOrStruct(context, name *scanner.Token) (Term, error) {
	y, err := r.s.Scan()
	if err != nil {
		return nil, err
	}

	switch y.Class {
	case scanner.Punct:
		switch y.Text {
		case terminator:
			return NewAtom(name.Text), nil
		case "(":
			args, err := r.readSeq(y.Text) // consumes closing ')'
			if err != nil {
				return nil, err
			}

			z, err := r.s.Scan()
			if err != nil {
				return nil, err
			}

			t := &Struct{
				Name: NewAtom(name.Text),
				Args: NewSeq(args),
			}
			if context != nil {
				t.Context = Var{Name: context.Text}
			}

			switch z.Class {
			case scanner.Punct:
				switch z.Text {
				case terminator:
					return t, nil
				case "{":
					data, err := r.readSeq(z.Text) // consumes closing '}'
					if err != nil {
						return nil, err
					}
					t.Data = NewDb(data)

					x, err := r.s.Scan()
					if err == io.EOF {
						return nil, &ErrUnexpectedEof{r.s.Pos()}
					}
					if err != nil {
						return nil, err
					}
					switch x.Class {
					case scanner.Punct:
						switch x.Text {
						case terminator:
							return t, nil
						}
					}
					return nil, &ErrUnexpectedToken{z}
				}
			}
			return nil, &ErrUnexpectedToken{z}
		}
	}
	return nil, &ErrUnexpectedToken{y}
}

// reads a sequence of terms. should be called immediately after the '('
// (or '}') token is consumed.  consumes the closing ')' (or '}') token.
func (r *Reader) readSeq(open string) ([]Term, error) {
	args := make([]Term, 0)

	var close string
	switch open {
	case "(":
		close = ")"
	case "{":
		close = "}"
	}

ARGS:
	for {
		t, err := r.Read()
		switch e := err.(type) {
		case nil:
			args = append(args, t)
		case *ErrUnexpectedToken:
			if e.Token.Class == scanner.Punct && e.Token.Text == close {
				break ARGS
			}
			return nil, err
		default:
			return nil, err
		}
	}

	return args, nil
}
<commit_msg>Factor out Reader.terminate()<commit_after>package term // import "github.com/amalog/go/term"

import (
	"io"

	"github.com/amalog/go/scanner"
)

// central definition of the term terminator
const terminator = ","

type Reader struct {
	s *scanner.Scanner
}

func NewReader(r io.RuneScanner) *Reader {
	self := new(Reader)
	self.s = scanner.New(r)
	return self
}

func (r *Reader) Read() (Term, error) {
	x, err := r.s.Scan()
	if err != nil {
		return nil, err
	}

	switch x.Class {
	case scanner.Atom:
		return r.readAtomOrStruct(nil, x)
	case scanner.String:
		if len(x.Text) < 2 {
			return nil, &Err{x, "string token too short"}
		}
		if x.Text[0] != '"' {
			return nil, &Err{x, "string missing opening quote"}
		}
		if x.Text[len(x.Text)-1] != '"' {
			return nil, &Err{x, "string missing closing quote"}
		}
		text := x.Text[1 : len(x.Text)-1]

		t := NewString(text)
		return r.terminate(t)
	case scanner.Var:
		y, err := r.s.Scan()
		if err != nil {
			return nil, err
		}

		switch y.Class {
		case scanner.Punct:
			switch y.Text {
			case terminator:
				term := &Var{
					Name:  x.Text,
					Value: nil,
				}
				return term, nil
			case ".":
				z, err := r.s.Scan()
				if err != nil {
					return nil, err
				}
				if z.Class == scanner.Atom {
					return r.readAtomOrStruct(x, z)
				}
			}
		}
		return nil, &ErrUnexpectedToken{y}
	default:
		return nil, &ErrUnexpectedToken{x}
	}

	return nil, nil
}

func (r *Reader) readAtomOrStruct(context, name *scanner.Token) (Term, error) {
	y, err := r.s.Scan()
	if err != nil {
		return nil, err
	}

	switch y.Class {
	case scanner.Punct:
		switch y.Text {
		case terminator:
			return NewAtom(name.Text), nil
		case "(":
			args, err := r.readSeq(y.Text) // consumes closing ')'
			if err != nil {
				return nil, err
			}

			z, err := r.s.Scan()
			if err != nil {
				return nil, err
			}

			t := &Struct{
				Name: NewAtom(name.Text),
				Args: NewSeq(args),
			}
			if context != nil {
				t.Context = Var{Name: context.Text}
			}

			switch z.Class {
			case scanner.Punct:
				switch z.Text {
				case terminator:
					return t, nil
				case "{":
					data, err := r.readSeq(z.Text) // consumes closing '}'
					if err != nil {
						return nil, err
					}
					t.Data = NewDb(data)
					return r.terminate(t)
				}
			}
			return nil, &ErrUnexpectedToken{z}
		}
	}
	return nil, &ErrUnexpectedToken{y}
}

// reads a sequence of terms. should be called immediately after the '('
// (or '}') token is consumed.  consumes the closing ')' (or '}') token.
func (r *Reader) readSeq(open string) ([]Term, error) {
	args := make([]Term, 0)

	var close string
	switch open {
	case "(":
		close = ")"
	case "{":
		close = "}"
	}

ARGS:
	for {
		t, err := r.Read()
		switch e := err.(type) {
		case nil:
			args = append(args, t)
		case *ErrUnexpectedToken:
			if e.Token.Class == scanner.Punct && e.Token.Text == close {
				break ARGS
			}
			return nil, err
		default:
			return nil, err
		}
	}

	return args, nil
}

// consume a terminator. if successful return the given term
func (r *Reader) terminate(t Term) (Term, error) {
	x, err := r.s.Scan()
	if err == io.EOF {
		return nil, &ErrUnexpectedEof{r.s.Pos()}
	}
	if err != nil {
		return nil, err
	}

	if x.Class == scanner.Punct && x.Text == terminator {
		return t, nil
	}
	return nil, &ErrUnexpectedToken{x}
}
<|endoftext|>


<commit_before>package main

import "sort"

func (d *dungeon) FreeNeighbors(pos position) []position {
	neighbors := [8]position{pos.E(), pos.W(), pos.N(), pos.S(), pos.NE(), pos.NW(), pos.SE(), pos.SW()}
	nb := make([]position, 0, 8)
	for _, npos := range neighbors {
		if d.Valid(npos) && d.Cell(npos).T != WallCell {
			nb = append(nb, npos)
		}
	}
	return nb
}

func (d *dungeon) CardinalFreeNeighbors(pos position) []position {
	neighbors := [4]position{pos.E(), pos.W(), pos.N(), pos.S()}
	nb := make([]position, 0, 4)
	for _, npos := range neighbors {
		if d.Valid(npos) && d.Cell(npos).T != WallCell {
			nb = append(nb, npos)
		}
	}
	return nb
}

type playerPath struct {
	game *game
}

func (pp *playerPath) Neighbors(pos position) []position {
	m := pp.game.Dungeon
	var neighbors []position
	if pp.game.Player.HasStatus(StatusConfusion) {
		neighbors = m.CardinalFreeNeighbors(pos)
	} else {
		neighbors = m.FreeNeighbors(pos)
	}
	freeNeighbors := []position{}
	for _, npos := range neighbors {
		if m.Cell(npos).Explored && !pp.game.UnknownDig[npos] && !pp.game.ExclusionsMap[npos] {
			freeNeighbors = append(freeNeighbors, npos)
		}
	}
	return freeNeighbors
}

func (pp *playerPath) Cost(from, to position) int {
	return 1
}

func (pp *playerPath) Estimation(from, to position) int {
	return from.Distance(to)
}

type noisePath struct {
	game *game
}

func (fp *noisePath) Neighbors(pos position) []position {
	return fp.game.Dungeon.FreeNeighbors(pos)
}

func (fp *noisePath) Cost(from, to position) int {
	return 1
}

type normalPath struct {
	game *game
}

func (np *normalPath) Neighbors(pos position) []position {
	if np.game.Player.HasStatus(StatusConfusion) {
		return np.game.Dungeon.CardinalFreeNeighbors(pos)
	}
	return np.game.Dungeon.FreeNeighbors(pos)
}

func (np *normalPath) Cost(from, to position) int {
	return 1
}

type autoexplorePath struct {
	game *game
}

func (ap *autoexplorePath) Neighbors(pos position) []position {
	if ap.game.ExclusionsMap[pos] {
		return nil
	}
	var neighbors []position
	if ap.game.Player.HasStatus(StatusConfusion) {
		neighbors = ap.game.Dungeon.CardinalFreeNeighbors(pos)
	} else {
		neighbors = ap.game.Dungeon.FreeNeighbors(pos)
	}
	var suitableNeighbors []position
	for _, pos := range neighbors {
		if !ap.game.ExclusionsMap[pos] {
			suitableNeighbors = append(suitableNeighbors, pos)
		}
	}
	return suitableNeighbors
}

func (ap *autoexplorePath) Cost(from, to position) int {
	return 1
}

type monPath struct {
	game    *game
	monster *monster
	wall    bool
}

func (mp *monPath) Neighbors(pos position) []position {
	if mp.monster.Status(MonsConfused) {
		if mp.wall {
			return mp.game.Dungeon.CardinalNeighbors(pos)
		}
		return mp.game.Dungeon.CardinalFreeNeighbors(pos)
	}
	if mp.wall {
		return mp.game.Dungeon.Neighbors(pos)
	}
	return mp.game.Dungeon.FreeNeighbors(pos)
}

func (mp *monPath) Cost(from, to position) int {
	g := mp.game
	mons, _ := g.MonsterAt(to)
	if !mons.Exists() {
		if mp.wall && g.Dungeon.Cell(to).T == WallCell && mp.monster.State != Hunting {
			return 6
		}
		return 1
	}
	return 4
}

func (mp *monPath) Estimation(from, to position) int {
	return from.Distance(to)
}

func (m *monster) APath(g *game, from, to position) []position {
	mp := &monPath{game: g, monster: m}
	if m.Kind == MonsEarthDragon {
		mp.wall = true
	}
	path, _, found := AstarPath(mp, from, to)
	if !found {
		return nil
	}
	return path
}

func (g *game) PlayerPath(from, to position) []position {
	pp := &playerPath{game: g}
	path, _, found := AstarPath(pp, from, to)
	if !found {
		return nil
	}
	return path
}

func (g *game) SortedNearestTo(cells []position, to position) []position {
	ps := posSlice{}
	for _, pos := range cells {
		pp := &playerPath{game: g}
		_, cost, found := AstarPath(pp, pos, to)
		if found {
			ps = append(ps, posCost{pos, cost})
		}
	}
	sort.Sort(ps)
	sorted := []position{}
	for _, pc := range ps {
		sorted = append(sorted, pc.pos)
	}
	return sorted
}

type posCost struct {
	pos  position
	cost int
}

type posSlice []posCost

func (ps posSlice) Len() int           { return len(ps) }
func (ps posSlice) Swap(i, j int)      { ps[i], ps[j] = ps[j], ps[i] }
func (ps posSlice) Less(i, j int) bool { return ps[i].cost < ps[j].cost }
<commit_msg>we know approximately the length<commit_after>package main

import "sort"

func (d *dungeon) FreeNeighbors(pos position) []position {
	neighbors := [8]position{pos.E(), pos.W(), pos.N(), pos.S(), pos.NE(), pos.NW(), pos.SE(), pos.SW()}
	nb := make([]position, 0, 8)
	for _, npos := range neighbors {
		if d.Valid(npos) && d.Cell(npos).T != WallCell {
			nb = append(nb, npos)
		}
	}
	return nb
}

func (d *dungeon) CardinalFreeNeighbors(pos position) []position {
	neighbors := [4]position{pos.E(), pos.W(), pos.N(), pos.S()}
	nb := make([]position, 0, 4)
	for _, npos := range neighbors {
		if d.Valid(npos) && d.Cell(npos).T != WallCell {
			nb = append(nb, npos)
		}
	}
	return nb
}

type playerPath struct {
	game *game
}

func (pp *playerPath) Neighbors(pos position) []position {
	m := pp.game.Dungeon
	var neighbors []position
	if pp.game.Player.HasStatus(StatusConfusion) {
		neighbors = m.CardinalFreeNeighbors(pos)
	} else {
		neighbors = m.FreeNeighbors(pos)
	}
	freeNeighbors := make([]position, 0, len(neighbors))
	for _, npos := range neighbors {
		if m.Cell(npos).Explored && !pp.game.UnknownDig[npos] && !pp.game.ExclusionsMap[npos] {
			freeNeighbors = append(freeNeighbors, npos)
		}
	}
	return freeNeighbors
}

func (pp *playerPath) Cost(from, to position) int {
	return 1
}

func (pp *playerPath) Estimation(from, to position) int {
	return from.Distance(to)
}

type noisePath struct {
	game *game
}

func (fp *noisePath) Neighbors(pos position) []position {
	return fp.game.Dungeon.FreeNeighbors(pos)
}

func (fp *noisePath) Cost(from, to position) int {
	return 1
}

type normalPath struct {
	game *game
}

func (np *normalPath) Neighbors(pos position) []position {
	if np.game.Player.HasStatus(StatusConfusion) {
		return np.game.Dungeon.CardinalFreeNeighbors(pos)
	}
	return np.game.Dungeon.FreeNeighbors(pos)
}

func (np *normalPath) Cost(from, to position) int {
	return 1
}

type autoexplorePath struct {
	game *game
}

func (ap *autoexplorePath) Neighbors(pos position) []position {
	if ap.game.ExclusionsMap[pos] {
		return nil
	}
	var neighbors []position
	if ap.game.Player.HasStatus(StatusConfusion) {
		neighbors = ap.game.Dungeon.CardinalFreeNeighbors(pos)
	} else {
		neighbors = ap.game.Dungeon.FreeNeighbors(pos)
	}
	var suitableNeighbors []position
	for _, pos := range neighbors {
		if !ap.game.ExclusionsMap[pos] {
			suitableNeighbors = append(suitableNeighbors, pos)
		}
	}
	return suitableNeighbors
}

func (ap *autoexplorePath) Cost(from, to position) int {
	return 1
}

type monPath struct {
	game    *game
	monster *monster
	wall    bool
}

func (mp *monPath) Neighbors(pos position) []position {
	if mp.monster.Status(MonsConfused) {
		if mp.wall {
			return mp.game.Dungeon.CardinalNeighbors(pos)
		}
		return mp.game.Dungeon.CardinalFreeNeighbors(pos)
	}
	if mp.wall {
		return mp.game.Dungeon.Neighbors(pos)
	}
	return mp.game.Dungeon.FreeNeighbors(pos)
}

func (mp *monPath) Cost(from, to position) int {
	g := mp.game
	mons, _ := g.MonsterAt(to)
	if !mons.Exists() {
		if mp.wall && g.Dungeon.Cell(to).T == WallCell && mp.monster.State != Hunting {
			return 6
		}
		return 1
	}
	return 4
}

func (mp *monPath) Estimation(from, to position) int {
	return from.Distance(to)
}

func (m *monster) APath(g *game, from, to position) []position {
	mp := &monPath{game: g, monster: m}
	if m.Kind == MonsEarthDragon {
		mp.wall = true
	}
	path, _, found := AstarPath(mp, from, to)
	if !found {
		return nil
	}
	return path
}

func (g *game) PlayerPath(from, to position) []position {
	pp := &playerPath{game: g}
	path, _, found := AstarPath(pp, from, to)
	if !found {
		return nil
	}
	return path
}

func (g *game) SortedNearestTo(cells []position, to position) []position {
	ps := posSlice{}
	for _, pos := range cells {
		pp := &playerPath{game: g}
		_, cost, found := AstarPath(pp, pos, to)
		if found {
			ps = append(ps, posCost{pos, cost})
		}
	}
	sort.Sort(ps)
	sorted := []position{}
	for _, pc := range ps {
		sorted = append(sorted, pc.pos)
	}
	return sorted
}

type posCost struct {
	pos  position
	cost int
}

type posSlice []posCost

func (ps posSlice) Len() int           { return len(ps) }
func (ps posSlice) Swap(i, j int)      { ps[i], ps[j] = ps[j], ps[i] }
func (ps posSlice) Less(i, j int) bool { return ps[i].cost < ps[j].cost }
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	PieceBlock
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool

	requestPieceChannel   chan uint32
	sendPieceBlockChannel chan BlockMessage
	sendHaveChannel       chan uint32
	done                  chan struct{}

	id string
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

type BlockMessage struct {
	index uint32
	begin uint32
	block []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localChoked = true
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	addr := fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)

	var err error
	peer.connection, err = net.Dial("tcp4", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s", addr)
		return
	} else if len(peer.id) != 0 && !bytes.Equal(handshake[48:68], []byte(peer.id)) {
		log.Printf("peer id mismatch from peer: %s", addr)
		return
	}

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(&blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		peer.sendUnchokeMessage()
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		if !peer.remoteInterested {
			return errors.New("peer sent request without showing interest")
		}

		if peer.localChoked {
			return errors.New("peer sent request while choked")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		peer.torrent.blockRequestChannel <- BlockRequestMessage{peer, index, begin, length}
	case PieceBlock:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	if piece, _ := peer.getPeerPiece(index); piece != nil {
		return
	}

	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(9+len(blockMessage.block)))
	packet.WriteByte(PieceBlock)
	binary.Write(&packet, binary.BigEndian, blockMessage.index)
	binary.Write(&packet, binary.BigEndian, blockMessage.begin)
	packet.Write(blockMessage.block)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := bytes.NewBuffer([]byte{0, 0, 0, 5, Have})
	binary.Write(packet, binary.BigEndian, pieceIndex)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendUnchokeMessage() {
	if !peer.localChoked {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Unchoke})
	peer.localChoked = false
}
<commit_msg>Protocol validation<commit_after>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	PieceBlock
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool

	requestPieceChannel   chan uint32
	sendPieceBlockChannel chan BlockMessage
	sendHaveChannel       chan uint32
	done                  chan struct{}

	id string
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

type BlockMessage struct {
	index uint32
	begin uint32
	block []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localChoked = true
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	addr := fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)

	var err error
	peer.connection, err = net.Dial("tcp4", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[0:20], []byte("\x13BitTorrent protocol")) {
		log.Printf("bad protocol from peer: %s\n", addr)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s\n", addr)
		return
	} else if len(peer.id) != 0 && !bytes.Equal(handshake[48:68], []byte(peer.id)) {
		log.Printf("peer id mismatch from peer: %s\n", addr)
		return
	}

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(&blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		peer.sendUnchokeMessage()
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		if !peer.remoteInterested {
			return errors.New("peer sent request without showing interest")
		}

		if peer.localChoked {
			return errors.New("peer sent request while choked")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		peer.torrent.blockRequestChannel <- BlockRequestMessage{peer, index, begin, length}
	case PieceBlock:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	if piece, _ := peer.getPeerPiece(index); piece != nil {
		return
	}

	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(9+len(blockMessage.block)))
	packet.WriteByte(PieceBlock)
	binary.Write(&packet, binary.BigEndian, blockMessage.index)
	binary.Write(&packet, binary.BigEndian, blockMessage.begin)
	packet.Write(blockMessage.block)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := bytes.NewBuffer([]byte{0, 0, 0, 5, Have})
	binary.Write(packet, binary.BigEndian, pieceIndex)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendUnchokeMessage() {
	if !peer.localChoked {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Unchoke})
	peer.localChoked = false
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	Piece
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localInterested  bool
	bitfield         []byte

	requestPieceChannel chan uint32
	done                chan struct{}
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	addr := fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)

	var err error
	peer.connection, err = net.Dial("tcp4", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s", addr)
		return
	}

	// TODO: Validate peer id if provided by tracker

	peer.requestPieceChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		// TODO: Unchoke peer and send files
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		// TODO: Handle request
	case Piece:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}
<commit_msg>Avoid requesting the same piece more than once<commit_after>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	Piece
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localInterested  bool
	bitfield         []byte

	requestPieceChannel chan uint32
	done                chan struct{}
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	addr := fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)

	var err error
	peer.connection, err = net.Dial("tcp4", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s", addr)
		return
	}

	// TODO: Validate peer id if provided by tracker

	peer.requestPieceChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		// TODO: Unchoke peer and send files
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		// TODO: Handle request
	case Piece:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	if piece, _ := peer.getPeerPiece(index); piece != nil {
		return
	}

	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}
<|endoftext|>


<commit_before><commit_msg>ちゃんとするけどとりあえず＞＜<commit_after><|endoftext|>


<commit_before>package db

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	turtle "github.com/gtfierro/hod/goraptor"
	"github.com/gtfierro/hod/query"

	"github.com/google/btree"
)

func (db *DB) RunQuery(q query.Query) QueryResult {
	// "clean" the query by expanding out the prefixes
	// make sure to first do the Filters, then the Or clauses
	for idx, filter := range q.Where.Filters {
		q.Where.Filters[idx] = db.expandFilter(filter)
	}
	for idx, orclause := range q.Where.Ors {
		q.Where.Ors[idx] = db.expandOrClauseFilters(orclause)
	}

	// we flatten the OR clauses to get the array of queries we are going
	// to run and then merge
	orTerms := query.FlattenOrClauseList(q.Where.Ors)
	oldFilters := q.Where.Filters

	unionedRows := btree.New(3)
	fullQueryStart := time.Now()

	// if we have terms that are part of a set of OR statements, then we run
	// parallel queries for each fully-elaborated "branch" or the OR statement,
	// and then merge the results together at the end
	if len(orTerms) > 0 {
		var rowLock sync.Mutex
		var wg sync.WaitGroup
		wg.Add(len(orTerms))
		for _, orTerm := range orTerms {
			tmpQuery := q.Copy()
			// augment with the filters
			tmpQuery.Where.Filters = make([]query.Filter, len(oldFilters)+len(orTerm))
			copy(tmpQuery.Where.Filters, oldFilters)
			copy(tmpQuery.Where.Filters[len(oldFilters):], orTerm)
			go func(q query.Query) {
				results := db.getQueryResults(q)
				rowLock.Lock()
				for _, row := range results {
					unionedRows.ReplaceOrInsert(ResultRow(row))
				}
				rowLock.Unlock()
				wg.Done()
			}(tmpQuery)
		}
		wg.Wait()
	} else {
		results := db.getQueryResults(q)
		for _, row := range results {
			unionedRows.ReplaceOrInsert(ResultRow(row))
		}
	}
	if db.showQueryLatencies {
		log.Noticef("Full Query took %s", time.Since(fullQueryStart))
	}

	var result QueryResult
	result.selectVars = q.Select.Variables

	if q.Select.Count {
		// return the count of results
		result.Count = unionedRows.Len()
	} else if q.Select.HasLinks {
		// resolve the links
		max := unionedRows.Max()
		iter := func(i btree.Item) bool {
			row := i.(ResultRow)
			var links = make(LinkResultMap)
			var hasContent = false
			for idx, selectvar := range q.Select.Variables {
				if len(selectvar.Links) == 0 && !selectvar.AllLinks {
					continue
				}
				// check for select all
				if selectvar.AllLinks {
					hash, err := db.GetHash(row[idx])
					if err != nil {
						log.Fatal(err)
					}
					keys, values, err := db.linkDB.getAll(hash)
					if err != nil {
						log.Fatal(err)
					}
					if len(keys) > 0 {
						hasContent = true
						links[row[idx]] = make(map[string]string)
					}
					for i := 0; i < len(keys); i++ {
						links[row[idx]][string(keys[i][:4])] = string(values[i])
					}
				} else {
					for _, _link := range selectvar.Links {
						link := &Link{URI: row[idx], Key: []byte(_link.Name)}
						value, err := db.linkDB.get(link)
						if err != nil {
							log.Fatal(err)
						}
						if len(links[row[idx]]) == 0 {
							links[row[idx]] = make(map[string]string)
						}
						if len(value) > 0 {
							links[row[idx]][string(link.Key)] = string(value)
							hasContent = true
						}
					}
				}
			}
			if hasContent {
				result.Links = append(result.Links, links)
			}
			return row.Less(max)
		}
		unionedRows.Ascend(iter)
	} else {
		// return the rows
		max := unionedRows.Max()
		iter := func(i btree.Item) bool {
			row := i.(ResultRow)
			m := make(ResultMap)
			for idx, vname := range q.Select.Variables {
				m[vname.Var.String()] = row[idx]
			}
			result.Rows = append(result.Rows, m)
			return row.Less(max)
		}
		unionedRows.Ascend(iter)
	}
	return result
}

func (db *DB) UpdateLinks(updates *LinkUpdates) error {
	tx, err := db.linkDB.startTx()
	if err != nil {
		return err
	}
	for _, linkAdd := range updates.Adding {
		linkAdd.Key = db.policy.SanitizeBytes(linkAdd.Key)
		linkAdd.Value = db.policy.SanitizeBytes(linkAdd.Value)
		db.linkDB.set(tx, linkAdd)
	}
	for _, linkRm := range updates.Removing {
		linkRm.Key = db.policy.SanitizeBytes(linkRm.Key)
		linkRm.Value = db.policy.SanitizeBytes(linkRm.Value)
		db.linkDB.delete(tx, linkRm)
	}
	return tx.Commit()
}

// TODO: add api call for getting links for entities
// for getting links from entities, we probably want to adopt a more generator-based approach
// to actually getting the rows from the database; as we get each row, we get the associated links,
// pipe that out to our accumulator (probably just appending to a list).

// takes a query and returns a DOT representation to visualize
// the construction of the query
func (db *DB) QueryToDOT(querystring io.Reader) (string, error) {
	q, err := query.Parse(querystring)
	if err != nil {
		return "", err
	}

	dot := ""
	dot += "digraph G {\n"
	dot += "ratio=\"auto\"\n"
	dot += "rankdir=\"LR\"\n"
	dot += "size=\"7.5,10\"\n"

	if len(q.Where.Ors) > 0 {
		orTerms := query.FlattenOrClauseList(q.Where.Ors)
		oldFilters := q.Where.Filters
		for _, orTerm := range orTerms {
			filters := append(oldFilters, orTerm...)
			for _, filter := range filters {
				var parts []string
				for _, p := range filter.Path {
					parts = append(parts, fmt.Sprintf("%s%s", p.Predicate, p.Pattern))
				}
				line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", filter.Subject, filter.Object, strings.Join(parts, "/"))
				if !strings.Contains(dot, line) {
					dot += line
				}

			}
		}
	} else {
		for _, filter := range q.Where.Filters {
			var parts []string
			for _, p := range filter.Path {
				parts = append(parts, fmt.Sprintf("%s%s", p.Predicate, p.Pattern))
			}
			line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", filter.Subject, filter.Object, strings.Join(parts, "/"))
			if !strings.Contains(dot, line) {
				dot += line
			}
		}
	}
	for _, sv := range q.Select.Variables {
		dot += fmt.Sprintf("\"%s\" [fillcolor=#e57373]\n", sv.Var)
	}
	dot += "}"
	return dot, nil
}

// executes a query and returns a DOT string of the classes involved
func (db *DB) QueryToClassDOT(querystring io.Reader) (string, error) {
	q, err := query.Parse(querystring)
	if err != nil {
		return "", err
	}
	// create DOT template string
	dot := ""
	dot += "digraph G {\n"

	// get rdf:type predicate hash as a string
	typeURI := turtle.ParseURI("rdf:type")
	typeURI.Namespace = db.namespaces[typeURI.Namespace]
	typeKey, err := db.GetHash(typeURI)
	if err != nil {
		return "", err
	}
	typeKeyString := typeKey.String()

	getClass := func(ent *Entity) (classes []turtle.URI, err error) {
		_classes := ent.OutEdges[typeKeyString]
		for _, class := range _classes {
			classes = append(classes, db.MustGetURI(class))
		}
		return
	}

	getEdges := func(ent *Entity) (predicates, objects []turtle.URI, reterr error) {
		var predKey Key
		for predKeyString, objectList := range ent.OutEdges {
			predKey.FromSlice([]byte(predKeyString))
			predURI, err := db.GetURI(predKey)
			if err != nil {
				reterr = err
				return
			}
			for _, objectKey := range objectList {
				objectEnt, err := db.GetEntityFromHash(objectKey)
				if err != nil {
					reterr = err
					return
				}
				objectClasses, err := getClass(objectEnt)
				if err != nil {
					reterr = err
					return
				}
				for _, class := range objectClasses {
					if predURI.Value != "type" && class.Value != "Class" {
						predicates = append(predicates, predURI)
						objects = append(objects, class)
					}
				}

			}
		}
		return
	}

	result := db.RunQuery(q)
	for _, row := range result.Rows {
		for _, uri := range row {
			ent, err := db.GetEntity(uri)
			if err != nil {
				return "", err
			}
			classList, err := getClass(ent)
			if err != nil {
				return "", err
			}
			preds, objs, err := getEdges(ent)
			if err != nil {
				return "", err
			}
			// add class as node to graph
			for _, class := range classList {
				line := fmt.Sprintf("\"%s\" [fillcolor=\"#4caf50\"];\n", db.Abbreviate(class))
				if !strings.Contains(dot, line) {
					dot += line
				}
				for i := 0; i < len(preds); i++ {
					line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", db.Abbreviate(class), db.Abbreviate(objs[i]), db.Abbreviate(preds[i]))
					if !strings.Contains(dot, line) {
						dot += line
					}
				}
			}

		}
	}

	dot += "}"

	return dot, nil
}

func (db *DB) Abbreviate(uri turtle.URI) string {
	for abbv, ns := range db.namespaces {
		if abbv != "" && ns == uri.Namespace {
			return abbv + ":" + uri.Value
		}
	}
	return uri.Value
}
<commit_msg>add count to results no matter what<commit_after>package db

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	turtle "github.com/gtfierro/hod/goraptor"
	"github.com/gtfierro/hod/query"

	"github.com/google/btree"
)

func (db *DB) RunQuery(q query.Query) QueryResult {
	// "clean" the query by expanding out the prefixes
	// make sure to first do the Filters, then the Or clauses
	for idx, filter := range q.Where.Filters {
		q.Where.Filters[idx] = db.expandFilter(filter)
	}
	for idx, orclause := range q.Where.Ors {
		q.Where.Ors[idx] = db.expandOrClauseFilters(orclause)
	}

	// we flatten the OR clauses to get the array of queries we are going
	// to run and then merge
	orTerms := query.FlattenOrClauseList(q.Where.Ors)
	oldFilters := q.Where.Filters

	unionedRows := btree.New(3)
	fullQueryStart := time.Now()

	// if we have terms that are part of a set of OR statements, then we run
	// parallel queries for each fully-elaborated "branch" or the OR statement,
	// and then merge the results together at the end
	if len(orTerms) > 0 {
		var rowLock sync.Mutex
		var wg sync.WaitGroup
		wg.Add(len(orTerms))
		for _, orTerm := range orTerms {
			tmpQuery := q.Copy()
			// augment with the filters
			tmpQuery.Where.Filters = make([]query.Filter, len(oldFilters)+len(orTerm))
			copy(tmpQuery.Where.Filters, oldFilters)
			copy(tmpQuery.Where.Filters[len(oldFilters):], orTerm)
			go func(q query.Query) {
				results := db.getQueryResults(q)
				rowLock.Lock()
				for _, row := range results {
					unionedRows.ReplaceOrInsert(ResultRow(row))
				}
				rowLock.Unlock()
				wg.Done()
			}(tmpQuery)
		}
		wg.Wait()
	} else {
		results := db.getQueryResults(q)
		for _, row := range results {
			unionedRows.ReplaceOrInsert(ResultRow(row))
		}
	}
	if db.showQueryLatencies {
		log.Noticef("Full Query took %s", time.Since(fullQueryStart))
	}

	var result QueryResult
	result.selectVars = q.Select.Variables

	if q.Select.Count {
		// return the count of results
		result.Count = unionedRows.Len()
	} else if q.Select.HasLinks {
		// resolve the links
		max := unionedRows.Max()
		iter := func(i btree.Item) bool {
			row := i.(ResultRow)
			var links = make(LinkResultMap)
			var hasContent = false
			for idx, selectvar := range q.Select.Variables {
				if len(selectvar.Links) == 0 && !selectvar.AllLinks {
					continue
				}
				// check for select all
				if selectvar.AllLinks {
					hash, err := db.GetHash(row[idx])
					if err != nil {
						log.Fatal(err)
					}
					keys, values, err := db.linkDB.getAll(hash)
					if err != nil {
						log.Fatal(err)
					}
					if len(keys) > 0 {
						hasContent = true
						links[row[idx]] = make(map[string]string)
					}
					for i := 0; i < len(keys); i++ {
						links[row[idx]][string(keys[i][:4])] = string(values[i])
					}
				} else {
					for _, _link := range selectvar.Links {
						link := &Link{URI: row[idx], Key: []byte(_link.Name)}
						value, err := db.linkDB.get(link)
						if err != nil {
							log.Fatal(err)
						}
						if len(links[row[idx]]) == 0 {
							links[row[idx]] = make(map[string]string)
						}
						if len(value) > 0 {
							links[row[idx]][string(link.Key)] = string(value)
							hasContent = true
						}
					}
				}
			}
			if hasContent {
				result.Links = append(result.Links, links)
			}
			return row.Less(max)
		}
		unionedRows.Ascend(iter)
		result.Count = len(result.Rows)
	} else {
		// return the rows
		max := unionedRows.Max()
		iter := func(i btree.Item) bool {
			row := i.(ResultRow)
			m := make(ResultMap)
			for idx, vname := range q.Select.Variables {
				m[vname.Var.String()] = row[idx]
			}
			result.Rows = append(result.Rows, m)
			return row.Less(max)
		}
		unionedRows.Ascend(iter)
		result.Count = len(result.Rows)
	}
	return result
}

func (db *DB) UpdateLinks(updates *LinkUpdates) error {
	tx, err := db.linkDB.startTx()
	if err != nil {
		return err
	}
	for _, linkAdd := range updates.Adding {
		linkAdd.Key = db.policy.SanitizeBytes(linkAdd.Key)
		linkAdd.Value = db.policy.SanitizeBytes(linkAdd.Value)
		db.linkDB.set(tx, linkAdd)
	}
	for _, linkRm := range updates.Removing {
		linkRm.Key = db.policy.SanitizeBytes(linkRm.Key)
		linkRm.Value = db.policy.SanitizeBytes(linkRm.Value)
		db.linkDB.delete(tx, linkRm)
	}
	return tx.Commit()
}

// TODO: add api call for getting links for entities
// for getting links from entities, we probably want to adopt a more generator-based approach
// to actually getting the rows from the database; as we get each row, we get the associated links,
// pipe that out to our accumulator (probably just appending to a list).

// takes a query and returns a DOT representation to visualize
// the construction of the query
func (db *DB) QueryToDOT(querystring io.Reader) (string, error) {
	q, err := query.Parse(querystring)
	if err != nil {
		return "", err
	}

	dot := ""
	dot += "digraph G {\n"
	dot += "ratio=\"auto\"\n"
	dot += "rankdir=\"LR\"\n"
	dot += "size=\"7.5,10\"\n"

	if len(q.Where.Ors) > 0 {
		orTerms := query.FlattenOrClauseList(q.Where.Ors)
		oldFilters := q.Where.Filters
		for _, orTerm := range orTerms {
			filters := append(oldFilters, orTerm...)
			for _, filter := range filters {
				var parts []string
				for _, p := range filter.Path {
					parts = append(parts, fmt.Sprintf("%s%s", p.Predicate, p.Pattern))
				}
				line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", filter.Subject, filter.Object, strings.Join(parts, "/"))
				if !strings.Contains(dot, line) {
					dot += line
				}

			}
		}
	} else {
		for _, filter := range q.Where.Filters {
			var parts []string
			for _, p := range filter.Path {
				parts = append(parts, fmt.Sprintf("%s%s", p.Predicate, p.Pattern))
			}
			line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", filter.Subject, filter.Object, strings.Join(parts, "/"))
			if !strings.Contains(dot, line) {
				dot += line
			}
		}
	}
	for _, sv := range q.Select.Variables {
		dot += fmt.Sprintf("\"%s\" [fillcolor=#e57373]\n", sv.Var)
	}
	dot += "}"
	return dot, nil
}

// executes a query and returns a DOT string of the classes involved
func (db *DB) QueryToClassDOT(querystring io.Reader) (string, error) {
	q, err := query.Parse(querystring)
	if err != nil {
		return "", err
	}
	// create DOT template string
	dot := ""
	dot += "digraph G {\n"

	// get rdf:type predicate hash as a string
	typeURI := turtle.ParseURI("rdf:type")
	typeURI.Namespace = db.namespaces[typeURI.Namespace]
	typeKey, err := db.GetHash(typeURI)
	if err != nil {
		return "", err
	}
	typeKeyString := typeKey.String()

	getClass := func(ent *Entity) (classes []turtle.URI, err error) {
		_classes := ent.OutEdges[typeKeyString]
		for _, class := range _classes {
			classes = append(classes, db.MustGetURI(class))
		}
		return
	}

	getEdges := func(ent *Entity) (predicates, objects []turtle.URI, reterr error) {
		var predKey Key
		for predKeyString, objectList := range ent.OutEdges {
			predKey.FromSlice([]byte(predKeyString))
			predURI, err := db.GetURI(predKey)
			if err != nil {
				reterr = err
				return
			}
			for _, objectKey := range objectList {
				objectEnt, err := db.GetEntityFromHash(objectKey)
				if err != nil {
					reterr = err
					return
				}
				objectClasses, err := getClass(objectEnt)
				if err != nil {
					reterr = err
					return
				}
				for _, class := range objectClasses {
					if predURI.Value != "type" && class.Value != "Class" {
						predicates = append(predicates, predURI)
						objects = append(objects, class)
					}
				}

			}
		}
		return
	}

	result := db.RunQuery(q)
	for _, row := range result.Rows {
		for _, uri := range row {
			ent, err := db.GetEntity(uri)
			if err != nil {
				return "", err
			}
			classList, err := getClass(ent)
			if err != nil {
				return "", err
			}
			preds, objs, err := getEdges(ent)
			if err != nil {
				return "", err
			}
			// add class as node to graph
			for _, class := range classList {
				line := fmt.Sprintf("\"%s\" [fillcolor=\"#4caf50\"];\n", db.Abbreviate(class))
				if !strings.Contains(dot, line) {
					dot += line
				}
				for i := 0; i < len(preds); i++ {
					line := fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", db.Abbreviate(class), db.Abbreviate(objs[i]), db.Abbreviate(preds[i]))
					if !strings.Contains(dot, line) {
						dot += line
					}
				}
			}

		}
	}

	dot += "}"

	return dot, nil
}

func (db *DB) Abbreviate(uri turtle.URI) string {
	for abbv, ns := range db.namespaces {
		if abbv != "" && ns == uri.Namespace {
			return abbv + ":" + uri.Value
		}
	}
	return uri.Value
}
<|endoftext|>