<commit_before>package batch

import (
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_auth"
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_conn"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_path"
	"github.com/watermint/toolbox/domain/dropbox/service/sv_member"
	"github.com/watermint/toolbox/domain/dropbox/service/sv_sharedfolder_member"
	"github.com/watermint/toolbox/domain/dropbox/usecase/uc_sharedfolder"
	"github.com/watermint/toolbox/essentials/lang"
	"github.com/watermint/toolbox/essentials/model/mo_string"
	"github.com/watermint/toolbox/essentials/queue/eq_sequence"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/feed/fd_file"
	"github.com/watermint/toolbox/infra/recipe/rc_exec"
	"github.com/watermint/toolbox/infra/recipe/rc_recipe"
	"github.com/watermint/toolbox/infra/report/rp_model"
	"github.com/watermint/toolbox/quality/infra/qt_file"
	"os"
)

type Add struct {
	Peer         dbx_conn.ConnScopedTeam
	File         fd_file.RowFeed
	OperationLog rp_model.TransactionReport
	Silent       bool
	Message      mo_string.OptionalString
}

func (z *Add) Preset() {
	z.Peer.SetScopes(
		dbx_auth.ScopeFilesContentRead,
		dbx_auth.ScopeGroupsRead,
		dbx_auth.ScopeMembersRead,
		dbx_auth.ScopeSharingRead,
		dbx_auth.ScopeSharingWrite,
		dbx_auth.ScopeTeamDataMember,
	)
	z.File.SetModel(&AddMember{})
	z.OperationLog.SetModel(&AddMember{}, nil)
}

func (z *Add) addMember(m *AddMember, resolver uc_sharedfolder.Resolver, svm sv_member.Member, c app_control.Control) error {
	member, err := svm.ResolveByEmail(m.MemberEmail)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	am, err := resolver.GroupOrEmailForAddMember(m.GroupOrEmail, m.AccessLevel)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	cm := z.Peer.Context().AsMemberId(member.TeamMemberId)
	sf, err := uc_sharedfolder.NewResolver(cm).Resolve(mo_path.NewDropboxPath(m.Path))
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	opts := make([]sv_sharedfolder_member.AddOption, 0)
	if z.Silent {
		opts = append(opts, sv_sharedfolder_member.AddQuiet())
	}
	if z.Message.IsExists() {
		opts = append(opts, sv_sharedfolder_member.AddCustomMessage(z.Message.Value()))
	}
	err = sv_sharedfolder_member.New(cm, sf).Add(am, opts...)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}
	z.OperationLog.Success(m, nil)
	return nil
}

func (z *Add) Exec(c app_control.Control) error {
	if err := z.OperationLog.Open(); err != nil {
		return err
	}

	var lastErr, listErr error
	sfr := uc_sharedfolder.NewResolver(z.Peer.Context())
	svm := sv_member.NewCached(z.Peer.Context())

	c.Sequence().Do(func(s eq_sequence.Stage) {
		s.Define("add_member", z.addMember, sfr, svm, c)
		q := s.Get("add_member")
		listErr = z.File.EachRow(func(m interface{}, rowIndex int) error {
			q.Enqueue(m)
			return nil
		})
	}, eq_sequence.ErrorHandler(func(err error, mouldId, batchId string, p interface{}) {
		lastErr = err
	}))

	return lang.NewMultiErrorOrNull(lastErr, listErr)
}

func (z *Add) Test(c app_control.Control) error {
	f, err := qt_file.MakeTestFile("share", "john@example.com,/shared,editor,emma@examplel.com\nemma@example.com,/project,viewer,john@example.com")
	if err != nil {
		return err
	}
	defer func() {
		_ = os.Remove(f)
	}()

	return rc_exec.ExecMock(c, &Add{}, func(r rc_recipe.Recipe) {
		m := r.(*Add)
		m.File.SetFilePath(f)
	})
}
<commit_msg>fix test<commit_after>package batch

import (
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_auth"
	"github.com/watermint/toolbox/domain/dropbox/api/dbx_conn"
	"github.com/watermint/toolbox/domain/dropbox/model/mo_path"
	"github.com/watermint/toolbox/domain/dropbox/service/sv_member"
	"github.com/watermint/toolbox/domain/dropbox/service/sv_sharedfolder_member"
	"github.com/watermint/toolbox/domain/dropbox/usecase/uc_sharedfolder"
	"github.com/watermint/toolbox/essentials/lang"
	"github.com/watermint/toolbox/essentials/model/mo_string"
	"github.com/watermint/toolbox/essentials/queue/eq_sequence"
	"github.com/watermint/toolbox/infra/control/app_control"
	"github.com/watermint/toolbox/infra/feed/fd_file"
	"github.com/watermint/toolbox/infra/report/rp_model"
	"github.com/watermint/toolbox/quality/infra/qt_errors"
)

type Add struct {
	Peer         dbx_conn.ConnScopedTeam
	File         fd_file.RowFeed
	OperationLog rp_model.TransactionReport
	Silent       bool
	Message      mo_string.OptionalString
}

func (z *Add) Preset() {
	z.Peer.SetScopes(
		dbx_auth.ScopeFilesContentRead,
		dbx_auth.ScopeGroupsRead,
		dbx_auth.ScopeMembersRead,
		dbx_auth.ScopeSharingRead,
		dbx_auth.ScopeSharingWrite,
		dbx_auth.ScopeTeamDataMember,
	)
	z.File.SetModel(&AddMember{})
	z.OperationLog.SetModel(&AddMember{}, nil)
}

func (z *Add) addMember(m *AddMember, resolver uc_sharedfolder.Resolver, svm sv_member.Member, c app_control.Control) error {
	member, err := svm.ResolveByEmail(m.MemberEmail)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	am, err := resolver.GroupOrEmailForAddMember(m.GroupOrEmail, m.AccessLevel)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	cm := z.Peer.Context().AsMemberId(member.TeamMemberId)
	sf, err := uc_sharedfolder.NewResolver(cm).Resolve(mo_path.NewDropboxPath(m.Path))
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}

	opts := make([]sv_sharedfolder_member.AddOption, 0)
	if z.Silent {
		opts = append(opts, sv_sharedfolder_member.AddQuiet())
	}
	if z.Message.IsExists() {
		opts = append(opts, sv_sharedfolder_member.AddCustomMessage(z.Message.Value()))
	}
	err = sv_sharedfolder_member.New(cm, sf).Add(am, opts...)
	if err != nil {
		z.OperationLog.Failure(err, m)
		return err
	}
	z.OperationLog.Success(m, nil)
	return nil
}

func (z *Add) Exec(c app_control.Control) error {
	if err := z.OperationLog.Open(); err != nil {
		return err
	}

	var lastErr, listErr error
	sfr := uc_sharedfolder.NewResolver(z.Peer.Context())
	svm := sv_member.NewCached(z.Peer.Context())

	c.Sequence().Do(func(s eq_sequence.Stage) {
		s.Define("add_member", z.addMember, sfr, svm, c)
		q := s.Get("add_member")
		listErr = z.File.EachRow(func(m interface{}, rowIndex int) error {
			q.Enqueue(m)
			return nil
		})
	}, eq_sequence.ErrorHandler(func(err error, mouldId, batchId string, p interface{}) {
		lastErr = err
	}))

	return lang.NewMultiErrorOrNull(lastErr, listErr)
}

func (z *Add) Test(c app_control.Control) error {
	return qt_errors.ErrorHumanInteractionRequired
	//f, err := qt_file.MakeTestFile("share", "john@example.com,/shared,editor,emma@examplel.com\nemma@example.com,/project,viewer,john@example.com")
	//if err != nil {
	//	return err
	//}
	//defer func() {
	//	_ = os.Remove(f)
	//}()
	//
	//return rc_exec.ExecMock(c, &Add{}, func(r rc_recipe.Recipe) {
	//	m := r.(*Add)
	//	m.File.SetFilePath(f)
	//})
}
<|endoftext|>


<commit_before>package websockets

import (
	"encoding/json"
	"fmt"
	"sync/atomic"

	"github.com/rubblelabs/ripple/data"
)

var counter uint64

type Syncer interface {
	Done()
	Fail(message string)
}

type CommandError struct {
	Name    string `json:"error"`
	Code    int    `json:"error_code"`
	Message string `json:"error_message"`
}

type Command struct {
	*CommandError
	Id     uint64        `json:"id"`
	Name   string        `json:"command"`
	Type   string        `json:"type,omitempty"`
	Status string        `json:"status,omitempty"`
	Ready  chan struct{} `json:"-"`
}

func (c *Command) Done() {
	c.Ready <- struct{}{}
}

func (c *Command) Fail(message string) {
	c.CommandError = &CommandError{
		Name:    "Client Error",
		Code:    -1,
		Message: message,
	}
	c.Ready <- struct{}{}
}

func (c *Command) IncrementId() {
	c.Id = atomic.AddUint64(&counter, 1)
}

func (e *CommandError) Error() string {
	return fmt.Sprintf("%s %d %s", e.Name, e.Code, e.Message)
}

func newCommand(command string) *Command {
	return &Command{
		Id:    atomic.AddUint64(&counter, 1),
		Name:  command,
		Ready: make(chan struct{}),
	}
}

type AccountTxCommand struct {
	*Command
	Account   data.Account           `json:"account"`
	MinLedger int64                  `json:"ledger_index_min"`
	MaxLedger int64                  `json:"ledger_index_max"`
	Binary    bool                   `json:"binary,omitempty"`
	Forward   bool                   `json:"forward,omitempty"`
	Limit     int                    `json:"limit,omitempty"`
	Marker    map[string]interface{} `json:"marker,omitempty"`
	Result    *AccountTxResult       `json:"result,omitempty"`
}

type AccountTxResult struct {
	Marker       map[string]interface{} `json:"marker,omitempty"`
	Transactions data.TransactionSlice  `json:"transactions,omitempty"`
}

func newAccountTxCommand(account data.Account, pageSize int, marker map[string]interface{}) *AccountTxCommand {
	return &AccountTxCommand{
		Command:   newCommand("account_tx"),
		Account:   account,
		MinLedger: -1,
		MaxLedger: -1,
		Limit:     pageSize,
		Marker:    marker,
	}
}

func newBinaryLedgerDataCommand(ledger interface{}, marker *data.Hash256) *BinaryLedgerDataCommand {
	return &BinaryLedgerDataCommand{
		Command: newCommand("ledger_data"),
		Binary:  true,
		Marker:  marker,
	}
}

type TxCommand struct {
	*Command
	Transaction data.Hash256 `json:"transaction"`
	Result      *TxResult    `json:"result,omitempty"`
}

type TxResult struct {
	data.TransactionWithMetaData
	Validated bool `json:"validated"`
}

// A shim to populate the Validated field before passing
// control on to TransactionWithMetaData.UnmarshalJSON
func (txr *TxResult) UnmarshalJSON(b []byte) error {
	var extract map[string]interface{}
	if err := json.Unmarshal(b, &extract); err != nil {
		return err
	}
	txr.Validated = extract["validated"].(bool)
	return json.Unmarshal(b, &txr.TransactionWithMetaData)
}

type SubmitCommand struct {
	*Command
	TxBlob string        `json:"tx_blob"`
	Result *SubmitResult `json:"result,omitempty"`
}

type SubmitResult struct {
	EngineResult        data.TransactionResult `json:"engine_result"`
	EngineResultCode    int                    `json:"engine_result_code"`
	EngineResultMessage string                 `json:"engine_result_message"`
	TxBlob              string                 `json:"tx_blob"`
	Tx                  data.Transaction       `json:"tx_json"`
}

type LedgerCommand struct {
	*Command
	LedgerIndex  interface{}   `json:"ledger_index"`
	Accounts     bool          `json:"accounts"`
	Transactions bool          `json:"transactions"`
	Expand       bool          `json:"expand"`
	Result       *LedgerResult `json:"result,omitempty"`
}

type LedgerResult struct {
	Ledger data.Ledger
}

type LedgerHeaderCommand struct {
	*Command
	Ledger interface{} `json:"ledger"`
	Result *LedgerHeaderResult
}

type LedgerHeaderResult struct {
	Ledger         data.Ledger
	LedgerSequence uint32              `json:"ledger_index"`
	Hash           *data.Hash256       `json:"ledger_hash,omitempty"`
	LedgerData     data.VariableLength `json:"ledger_data"`
}

type LedgerDataCommand struct {
	*Command
	Ledger interface{}       `json:"ledger"`
	Marker *data.Hash256     `json:"marker,omitempty"`
	Result *LedgerDataResult `json:"result,omitempty"`
}

type BinaryLedgerDataCommand struct {
	*Command
	Ledger interface{}             `json:"ledger"`
	Binary bool                    `json:"binary"`
	Marker *data.Hash256           `json:"marker,omitempty"`
	Result *BinaryLedgerDataResult `json:"result,omitempty"`
}

type LedgerDataResult struct {
	LedgerSequence uint32                `json:"ledger_index,string"`
	Hash           data.Hash256          `json:"ledger_hash"`
	Marker         *data.Hash256         `json:"marker"`
	State          data.LedgerEntrySlice `json:"state"`
}

type BinaryLedgerData struct {
	Data  string `json:"data"`
	Index string `json:"index"`
}

type BinaryLedgerDataResult struct {
	LedgerSequence uint32             `json:"ledger_index,string"`
	Hash           data.Hash256       `json:"ledger_hash"`
	Marker         *data.Hash256      `json:"marker"`
	State          []BinaryLedgerData `json:"state"`
}

type RipplePathFindCommand struct {
	*Command
	SrcAccount    data.Account          `json:"source_account"`
	SrcCurrencies *[]data.Currency      `json:"source_currencies,omitempty"`
	DestAccount   data.Account          `json:"destination_account"`
	DestAmount    data.Amount           `json:"destination_amount"`
	Result        *RipplePathFindResult `json:"result,omitempty"`
}

type RipplePathFindResult struct {
	Alternatives []struct {
		SrcAmount      data.Amount  `json:"source_amount"`
		PathsComputed  data.PathSet `json:"paths_computed,omitempty"`
		PathsCanonical data.PathSet `json:"paths_canonical,omitempty"`
	}
	DestAccount    data.Account    `json:"destination_account"`
	DestCurrencies []data.Currency `json:"destination_currencies"`
}

type AccountInfoCommand struct {
	*Command
	Account data.Account       `json:"account"`
	Result  *AccountInfoResult `json:"result,omitempty"`
}

type AccountInfoResult struct {
	LedgerSequence uint32           `json:"ledger_current_index"`
	AccountData    data.AccountRoot `json:"account_data"`
}
<commit_msg>revert tx_json<commit_after>package websockets

import (
	"encoding/json"
	"fmt"
	"sync/atomic"

	"github.com/rubblelabs/ripple/data"
)

var counter uint64

type Syncer interface {
	Done()
	Fail(message string)
}

type CommandError struct {
	Name    string `json:"error"`
	Code    int    `json:"error_code"`
	Message string `json:"error_message"`
}

type Command struct {
	*CommandError
	Id     uint64        `json:"id"`
	Name   string        `json:"command"`
	Type   string        `json:"type,omitempty"`
	Status string        `json:"status,omitempty"`
	Ready  chan struct{} `json:"-"`
}

func (c *Command) Done() {
	c.Ready <- struct{}{}
}

func (c *Command) Fail(message string) {
	c.CommandError = &CommandError{
		Name:    "Client Error",
		Code:    -1,
		Message: message,
	}
	c.Ready <- struct{}{}
}

func (c *Command) IncrementId() {
	c.Id = atomic.AddUint64(&counter, 1)
}

func (e *CommandError) Error() string {
	return fmt.Sprintf("%s %d %s", e.Name, e.Code, e.Message)
}

func newCommand(command string) *Command {
	return &Command{
		Id:    atomic.AddUint64(&counter, 1),
		Name:  command,
		Ready: make(chan struct{}),
	}
}

type AccountTxCommand struct {
	*Command
	Account   data.Account           `json:"account"`
	MinLedger int64                  `json:"ledger_index_min"`
	MaxLedger int64                  `json:"ledger_index_max"`
	Binary    bool                   `json:"binary,omitempty"`
	Forward   bool                   `json:"forward,omitempty"`
	Limit     int                    `json:"limit,omitempty"`
	Marker    map[string]interface{} `json:"marker,omitempty"`
	Result    *AccountTxResult       `json:"result,omitempty"`
}

type AccountTxResult struct {
	Marker       map[string]interface{} `json:"marker,omitempty"`
	Transactions data.TransactionSlice  `json:"transactions,omitempty"`
}

func newAccountTxCommand(account data.Account, pageSize int, marker map[string]interface{}) *AccountTxCommand {
	return &AccountTxCommand{
		Command:   newCommand("account_tx"),
		Account:   account,
		MinLedger: -1,
		MaxLedger: -1,
		Limit:     pageSize,
		Marker:    marker,
	}
}

func newBinaryLedgerDataCommand(ledger interface{}, marker *data.Hash256) *BinaryLedgerDataCommand {
	return &BinaryLedgerDataCommand{
		Command: newCommand("ledger_data"),
		Binary:  true,
		Marker:  marker,
	}
}

type TxCommand struct {
	*Command
	Transaction data.Hash256 `json:"transaction"`
	Result      *TxResult    `json:"result,omitempty"`
}

type TxResult struct {
	data.TransactionWithMetaData
	Validated bool `json:"validated"`
}

// A shim to populate the Validated field before passing
// control on to TransactionWithMetaData.UnmarshalJSON
func (txr *TxResult) UnmarshalJSON(b []byte) error {
	var extract map[string]interface{}
	if err := json.Unmarshal(b, &extract); err != nil {
		return err
	}
	txr.Validated = extract["validated"].(bool)
	return json.Unmarshal(b, &txr.TransactionWithMetaData)
}

type SubmitCommand struct {
	*Command
	TxBlob string        `json:"tx_blob"`
	Result *SubmitResult `json:"result,omitempty"`
}

type SubmitResult struct {
	EngineResult        data.TransactionResult `json:"engine_result"`
	EngineResultCode    int                    `json:"engine_result_code"`
	EngineResultMessage string                 `json:"engine_result_message"`
	TxBlob              string                 `json:"tx_blob"`
	Tx                  interface{}            `json:"tx_json"`
}

type LedgerCommand struct {
	*Command
	LedgerIndex  interface{}   `json:"ledger_index"`
	Accounts     bool          `json:"accounts"`
	Transactions bool          `json:"transactions"`
	Expand       bool          `json:"expand"`
	Result       *LedgerResult `json:"result,omitempty"`
}

type LedgerResult struct {
	Ledger data.Ledger
}

type LedgerHeaderCommand struct {
	*Command
	Ledger interface{} `json:"ledger"`
	Result *LedgerHeaderResult
}

type LedgerHeaderResult struct {
	Ledger         data.Ledger
	LedgerSequence uint32              `json:"ledger_index"`
	Hash           *data.Hash256       `json:"ledger_hash,omitempty"`
	LedgerData     data.VariableLength `json:"ledger_data"`
}

type LedgerDataCommand struct {
	*Command
	Ledger interface{}       `json:"ledger"`
	Marker *data.Hash256     `json:"marker,omitempty"`
	Result *LedgerDataResult `json:"result,omitempty"`
}

type BinaryLedgerDataCommand struct {
	*Command
	Ledger interface{}             `json:"ledger"`
	Binary bool                    `json:"binary"`
	Marker *data.Hash256           `json:"marker,omitempty"`
	Result *BinaryLedgerDataResult `json:"result,omitempty"`
}

type LedgerDataResult struct {
	LedgerSequence uint32                `json:"ledger_index,string"`
	Hash           data.Hash256          `json:"ledger_hash"`
	Marker         *data.Hash256         `json:"marker"`
	State          data.LedgerEntrySlice `json:"state"`
}

type BinaryLedgerData struct {
	Data  string `json:"data"`
	Index string `json:"index"`
}

type BinaryLedgerDataResult struct {
	LedgerSequence uint32             `json:"ledger_index,string"`
	Hash           data.Hash256       `json:"ledger_hash"`
	Marker         *data.Hash256      `json:"marker"`
	State          []BinaryLedgerData `json:"state"`
}

type RipplePathFindCommand struct {
	*Command
	SrcAccount    data.Account          `json:"source_account"`
	SrcCurrencies *[]data.Currency      `json:"source_currencies,omitempty"`
	DestAccount   data.Account          `json:"destination_account"`
	DestAmount    data.Amount           `json:"destination_amount"`
	Result        *RipplePathFindResult `json:"result,omitempty"`
}

type RipplePathFindResult struct {
	Alternatives []struct {
		SrcAmount      data.Amount  `json:"source_amount"`
		PathsComputed  data.PathSet `json:"paths_computed,omitempty"`
		PathsCanonical data.PathSet `json:"paths_canonical,omitempty"`
	}
	DestAccount    data.Account    `json:"destination_account"`
	DestCurrencies []data.Currency `json:"destination_currencies"`
}

type AccountInfoCommand struct {
	*Command
	Account data.Account       `json:"account"`
	Result  *AccountInfoResult `json:"result,omitempty"`
}

type AccountInfoResult struct {
	LedgerSequence uint32           `json:"ledger_current_index"`
	AccountData    data.AccountRoot `json:"account_data"`
}
<|endoftext|>


<commit_before>package adapter

import (
	"fmt"
	"github.com/stitchfix/flotilla-os/state"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"strings"
	"time"
)

type EKSAdapter interface {
	AdaptJobToFlotillaRun(job *batchv1.Job, run state.Run, pod *corev1.Pod) (state.Run, error)
	AdaptFlotillaDefinitionAndRunToJob(td state.Definition, run state.Run, sa string, schedulerName string) (batchv1.Job, error)
}
type eksAdapter struct{}

//
// NewEKSAdapter configures and returns an eks adapter for translating
// from EKS api specific objects to our representation
//
func NewEKSAdapter() (EKSAdapter, error) {
	adapter := eksAdapter{}

	return &adapter, nil
}

func (a *eksAdapter) AdaptJobToFlotillaRun(job *batchv1.Job, run state.Run, pod *corev1.Pod) (state.Run, error) {
	updated := run
	if job.Status.Active == 1 && job.Status.CompletionTime == nil {
		updated.Status = state.StatusRunning
	} else if job.Status.Succeeded == 1 {
		var exitCode int64 = 0
		updated.Status = state.StatusStopped
		updated.ExitCode = &exitCode
	} else if job.Status.Failed == 1 {
		var exitCode int64 = 1
		updated.Status = state.StatusStopped
		if pod != nil {
			if pod.Status.ContainerStatuses != nil && len(pod.Status.ContainerStatuses) > 0 {
				containerStatus := pod.Status.ContainerStatuses[len(pod.Status.ContainerStatuses)-1]
				if containerStatus.State.Terminated != nil {
					updated.ExitReason = &containerStatus.State.Terminated.Reason
					exitCode = int64(containerStatus.State.Terminated.ExitCode)
				}
			}
		}
		updated.ExitCode = &exitCode
	}

	if pod != nil && len(pod.Spec.Containers) > 0 {
		container := pod.Spec.Containers[0]
		//First three lines are injected by Flotilla, strip those out.
		if len(container.Command) > 3 {
			cmd := strings.Join(container.Command[3:], "\n")
			run.Command = &cmd
		}
	}

	if job != nil && job.Status.StartTime != nil {
		updated.StartedAt = &job.Status.StartTime.Time
	}

	if updated.Status == state.StatusStopped {
		if job != nil && job.Status.CompletionTime != nil {
			updated.FinishedAt = &job.Status.CompletionTime.Time
		} else {
			finishedAt := time.Now()
			updated.FinishedAt = &finishedAt
		}
	}
	return updated, nil
}

func (a *eksAdapter) AdaptFlotillaDefinitionAndRunToJob(definition state.Definition, run state.Run, sa string, schedulerName string) (batchv1.Job, error) {
	resourceRequirements := a.constructResourceRequirements(definition, run)

	cmd := definition.Command
	if run.Command != nil {
		cmd = *run.Command
	}

	container := corev1.Container{
		Name:      run.RunID,
		Image:     run.Image,
		Command:   a.constructCmdSlice(cmd),
		Resources: resourceRequirements,
		Env:       a.envOverrides(definition, run),
	}

	affinity := a.constructAffinity(definition, run)
	annotations := map[string]string{"cluster-autoscaler.kubernetes.io/safe-to-evict": "false"}

	activeDeadlineSeconds := state.SpotActiveDeadlineSeconds
	if *run.NodeLifecycle == state.OndemandLifecycle {
		activeDeadlineSeconds = state.OndemandActiveDeadlineSeconds
	}

	jobSpec := batchv1.JobSpec{
		TTLSecondsAfterFinished: &state.TTLSecondsAfterFinished,
		ActiveDeadlineSeconds:   &activeDeadlineSeconds,
		BackoffLimit:            &state.EKSBackoffLimit,

		Template: corev1.PodTemplateSpec{
			ObjectMeta: v1.ObjectMeta{
				Annotations: annotations,
			},
			Spec: corev1.PodSpec{
				SchedulerName:      schedulerName,
				Containers:         []corev1.Container{container},
				RestartPolicy:      corev1.RestartPolicyNever,
				ServiceAccountName: sa,
				Affinity:           affinity,
			},
		},
	}

	eksJob := batchv1.Job{
		Spec: jobSpec,
		ObjectMeta: v1.ObjectMeta{
			Name: run.RunID,
		},
	}

	return eksJob, nil
}

func (a *eksAdapter) constructAffinity(definition state.Definition, run state.Run) *corev1.Affinity {
	affinity := &corev1.Affinity{}
	var matchExpressions []corev1.NodeSelectorRequirement

	gpuNodeTypes := []string{"p3.2xlarge", "p3.8xlarge", "p3.16xlarge"}

	var nodeLifecycle []string
	if *run.NodeLifecycle == state.OndemandLifecycle {
		nodeLifecycle = append(nodeLifecycle, "normal")
	} else {
		nodeLifecycle = append(nodeLifecycle, "spot")
	}

	if definition.Gpu == nil || *definition.Gpu <= 0 {
		matchExpressions = append(matchExpressions, corev1.NodeSelectorRequirement{
			Key:      "beta.kubernetes.io/instance-type",
			Operator: corev1.NodeSelectorOpNotIn,
			Values:   gpuNodeTypes,
		})
	}

	matchExpressions = append(matchExpressions, corev1.NodeSelectorRequirement{
		Key:      "kubernetes.io/lifecycle",
		Operator: corev1.NodeSelectorOpIn,
		Values:   nodeLifecycle,
	})

	affinity = &corev1.Affinity{
		NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{
					{
						MatchExpressions: matchExpressions,
					},
				},
			},
		},
	}

	return affinity
}

func (a *eksAdapter) constructResourceRequirements(definition state.Definition, run state.Run) corev1.ResourceRequirements {
	limits := make(corev1.ResourceList)
	cpu := *definition.Cpu
	if run.Cpu != nil {
		cpu = *run.Cpu
	}
	if cpu < state.MinCPU {
		cpu = state.MinCPU
	}

	mem := *definition.Memory
	if run.Memory != nil {
		mem = *run.Memory
	}
	if mem < state.MinMem {
		mem = state.MinMem
	}

	// CPU Override for legacy jobs that are high memory, remove once migration is complete.
	if mem >= 36864 && mem < 131072 && (definition.Gpu == nil || *definition.Gpu == 0) {
		// using the 8x ratios between cpu and memory ~ r5 class of instances
		cpuOverride := mem / 8
		if cpuOverride > cpu {
			cpu = cpuOverride
		}
	}

	cpuQuantity := resource.MustParse(fmt.Sprintf("%dm", cpu))
	assignedCpu := cpuQuantity.ScaledValue(resource.Milli)
	run.Cpu = &assignedCpu

	memoryQuantity := resource.MustParse(fmt.Sprintf("%dM", mem))
	assignedMem := memoryQuantity.ScaledValue(resource.Mega)
	run.Memory = &assignedMem

	if definition.Gpu != nil && *definition.Gpu > 0 {
		limits["nvidia.com/gpu"] = resource.MustParse(fmt.Sprintf("%d", *definition.Gpu))
	}
	if run.EphemeralStorage != nil {
		limits[corev1.ResourceEphemeralStorage] =
			resource.MustParse(fmt.Sprintf("%dGi", *run.EphemeralStorage))
	}
	limits[corev1.ResourceCPU] = cpuQuantity
	limits[corev1.ResourceMemory] = memoryQuantity
	resourceRequirements := corev1.ResourceRequirements{
		Limits: limits,
	}
	return resourceRequirements
}

func (a *eksAdapter) constructCmdSlice(cmdString string) []string {
	bashCmd := "bash"
	optLogin := "-l"
	optStr := "-cex"
	return []string{bashCmd, optLogin, optStr, cmdString}
}

func (a *eksAdapter) envOverrides(definition state.Definition, run state.Run) []corev1.EnvVar {
	pairs := make(map[string]string)
	for _, ev := range *definition.Env {
		name := a.sanitizeEnvVar(ev.Name)
		value := ev.Value
		pairs[name] = value
	}

	for _, ev := range *run.Env {
		name := a.sanitizeEnvVar(ev.Name)
		value := ev.Value
		pairs[name] = value
	}

	var res []corev1.EnvVar
	for key := range pairs {
		if len(key) > 0 {
			res = append(res, corev1.EnvVar{
				Name:  key,
				Value: pairs[key],
			})
		}
	}
	return res
}

func (a *eksAdapter) sanitizeEnvVar(key string) string {
	// Environment variable can't start with a $
	if strings.HasPrefix(key, "$") {
		key = strings.Replace(key, "$", "", 1)
	}
	// Environment variable names can't contain spaces.
	key = strings.Replace(key, " ", "", -1)
	return key
}
<commit_msg>Adding CPU affinity for high cpu jobs (#239)<commit_after>package adapter

import (
	"fmt"
	"github.com/stitchfix/flotilla-os/state"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"strings"
	"time"
)

type EKSAdapter interface {
	AdaptJobToFlotillaRun(job *batchv1.Job, run state.Run, pod *corev1.Pod) (state.Run, error)
	AdaptFlotillaDefinitionAndRunToJob(td state.Definition, run state.Run, sa string, schedulerName string) (batchv1.Job, error)
}
type eksAdapter struct{}

//
// NewEKSAdapter configures and returns an eks adapter for translating
// from EKS api specific objects to our representation
//
func NewEKSAdapter() (EKSAdapter, error) {
	adapter := eksAdapter{}

	return &adapter, nil
}

func (a *eksAdapter) AdaptJobToFlotillaRun(job *batchv1.Job, run state.Run, pod *corev1.Pod) (state.Run, error) {
	updated := run
	if job.Status.Active == 1 && job.Status.CompletionTime == nil {
		updated.Status = state.StatusRunning
	} else if job.Status.Succeeded == 1 {
		var exitCode int64 = 0
		updated.Status = state.StatusStopped
		updated.ExitCode = &exitCode
	} else if job.Status.Failed == 1 {
		var exitCode int64 = 1
		updated.Status = state.StatusStopped
		if pod != nil {
			if pod.Status.ContainerStatuses != nil && len(pod.Status.ContainerStatuses) > 0 {
				containerStatus := pod.Status.ContainerStatuses[len(pod.Status.ContainerStatuses)-1]
				if containerStatus.State.Terminated != nil {
					updated.ExitReason = &containerStatus.State.Terminated.Reason
					exitCode = int64(containerStatus.State.Terminated.ExitCode)
				}
			}
		}
		updated.ExitCode = &exitCode
	}

	if pod != nil && len(pod.Spec.Containers) > 0 {
		container := pod.Spec.Containers[0]
		//First three lines are injected by Flotilla, strip those out.
		if len(container.Command) > 3 {
			cmd := strings.Join(container.Command[3:], "\n")
			run.Command = &cmd
		}
	}

	if job != nil && job.Status.StartTime != nil {
		updated.StartedAt = &job.Status.StartTime.Time
	}

	if updated.Status == state.StatusStopped {
		if job != nil && job.Status.CompletionTime != nil {
			updated.FinishedAt = &job.Status.CompletionTime.Time
		} else {
			finishedAt := time.Now()
			updated.FinishedAt = &finishedAt
		}
	}
	return updated, nil
}

func (a *eksAdapter) AdaptFlotillaDefinitionAndRunToJob(definition state.Definition, run state.Run, sa string, schedulerName string) (batchv1.Job, error) {
	resourceRequirements, run := a.constructResourceRequirements(definition, run)

	cmd := definition.Command
	if run.Command != nil {
		cmd = *run.Command
	}

	container := corev1.Container{
		Name:      run.RunID,
		Image:     run.Image,
		Command:   a.constructCmdSlice(cmd),
		Resources: resourceRequirements,
		Env:       a.envOverrides(definition, run),
	}

	affinity := a.constructAffinity(definition, run)
	annotations := map[string]string{"cluster-autoscaler.kubernetes.io/safe-to-evict": "false"}

	activeDeadlineSeconds := state.SpotActiveDeadlineSeconds
	if *run.NodeLifecycle == state.OndemandLifecycle {
		activeDeadlineSeconds = state.OndemandActiveDeadlineSeconds
	}

	jobSpec := batchv1.JobSpec{
		TTLSecondsAfterFinished: &state.TTLSecondsAfterFinished,
		ActiveDeadlineSeconds:   &activeDeadlineSeconds,
		BackoffLimit:            &state.EKSBackoffLimit,

		Template: corev1.PodTemplateSpec{
			ObjectMeta: v1.ObjectMeta{
				Annotations: annotations,
			},
			Spec: corev1.PodSpec{
				SchedulerName:      schedulerName,
				Containers:         []corev1.Container{container},
				RestartPolicy:      corev1.RestartPolicyNever,
				ServiceAccountName: sa,
				Affinity:           affinity,
			},
		},
	}

	eksJob := batchv1.Job{
		Spec: jobSpec,
		ObjectMeta: v1.ObjectMeta{
			Name: run.RunID,
		},
	}

	return eksJob, nil
}

func (a *eksAdapter) constructAffinity(definition state.Definition, run state.Run) *corev1.Affinity {
	affinity := &corev1.Affinity{}
	var requiredMatch []corev1.NodeSelectorRequirement

	gpuNodeTypes := []string{"p3.2xlarge", "p3.8xlarge", "p3.16xlarge"}
	cpuNodeTypes := []string{"c5.2xlarge", "c5.4xlarge", "c5.9xlarge"}

	var nodeLifecycle []string
	if *run.NodeLifecycle == state.OndemandLifecycle {
		nodeLifecycle = append(nodeLifecycle, "normal")
	} else {
		nodeLifecycle = append(nodeLifecycle, "spot")
	}

	if definition.Gpu == nil || *definition.Gpu <= 0 {
		requiredMatch = append(requiredMatch, corev1.NodeSelectorRequirement{
			Key:      "beta.kubernetes.io/instance-type",
			Operator: corev1.NodeSelectorOpNotIn,
			Values:   gpuNodeTypes,
		})

		//For high cpu jobs - assign to c5 node types.
		if run.Memory != nil &&
			run.Cpu != nil &&
			*run.Cpu > int64(0) &&
			*run.Memory > int64(0) &&
			float64(*run.Cpu)/float64(*run.Memory) >= 0.5 {
		}

		requiredMatch = append(requiredMatch, corev1.NodeSelectorRequirement{
			Key:      "beta.kubernetes.io/instance-type",
			Operator: corev1.NodeSelectorOpIn,
			Values:   cpuNodeTypes,
		})
	}

	requiredMatch = append(requiredMatch, corev1.NodeSelectorRequirement{
		Key:      "kubernetes.io/lifecycle",
		Operator: corev1.NodeSelectorOpIn,
		Values:   nodeLifecycle,
	})

	affinity = &corev1.Affinity{
		NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{
					{
						MatchExpressions: requiredMatch,
					},
				},
			},
		},
	}

	return affinity
}

func (a *eksAdapter) constructResourceRequirements(definition state.Definition, run state.Run) (corev1.ResourceRequirements, state.Run) {
	limits := make(corev1.ResourceList)
	cpu := *definition.Cpu
	if run.Cpu != nil {
		cpu = *run.Cpu
	}
	if cpu < state.MinCPU {
		cpu = state.MinCPU
	}

	mem := *definition.Memory
	if run.Memory != nil {
		mem = *run.Memory
	}
	if mem < state.MinMem {
		mem = state.MinMem
	}

	// CPU Override for legacy jobs that are high memory, remove once migration is complete.
	if mem >= 36864 && mem < 131072 && (definition.Gpu == nil || *definition.Gpu == 0) {
		// using the 8x ratios between cpu and memory ~ r5 class of instances
		cpuOverride := mem / 8
		if cpuOverride > cpu {
			cpu = cpuOverride
		}
	}

	cpuQuantity := resource.MustParse(fmt.Sprintf("%dm", cpu))
	assignedCpu := cpuQuantity.ScaledValue(resource.Milli)
	run.Cpu = &assignedCpu

	memoryQuantity := resource.MustParse(fmt.Sprintf("%dM", mem))
	assignedMem := memoryQuantity.ScaledValue(resource.Mega)
	run.Memory = &assignedMem

	if definition.Gpu != nil && *definition.Gpu > 0 {
		limits["nvidia.com/gpu"] = resource.MustParse(fmt.Sprintf("%d", *definition.Gpu))
	}
	if run.EphemeralStorage != nil {
		limits[corev1.ResourceEphemeralStorage] =
			resource.MustParse(fmt.Sprintf("%dGi", *run.EphemeralStorage))
	}
	limits[corev1.ResourceCPU] = cpuQuantity
	limits[corev1.ResourceMemory] = memoryQuantity
	resourceRequirements := corev1.ResourceRequirements{
		Limits: limits,
	}
	return resourceRequirements, run
}

func (a *eksAdapter) constructCmdSlice(cmdString string) []string {
	bashCmd := "bash"
	optLogin := "-l"
	optStr := "-cex"
	return []string{bashCmd, optLogin, optStr, cmdString}
}

func (a *eksAdapter) envOverrides(definition state.Definition, run state.Run) []corev1.EnvVar {
	pairs := make(map[string]string)
	for _, ev := range *definition.Env {
		name := a.sanitizeEnvVar(ev.Name)
		value := ev.Value
		pairs[name] = value
	}

	for _, ev := range *run.Env {
		name := a.sanitizeEnvVar(ev.Name)
		value := ev.Value
		pairs[name] = value
	}

	var res []corev1.EnvVar
	for key := range pairs {
		if len(key) > 0 {
			res = append(res, corev1.EnvVar{
				Name:  key,
				Value: pairs[key],
			})
		}
	}
	return res
}

func (a *eksAdapter) sanitizeEnvVar(key string) string {
	// Environment variable can't start with a $
	if strings.HasPrefix(key, "$") {
		key = strings.Replace(key, "$", "", 1)
	}
	// Environment variable names can't contain spaces.
	key = strings.Replace(key, " ", "", -1)
	return key
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/gorilla/mux"
	"github.com/vishvananda/netlink"
	. "github.com/weaveworks/weave/common"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
)

const (
	MethodReceiver = "NetworkDriver"
	WeaveContainer = "weave"
	WeaveBridge    = "weave"
)

var version = "(unreleased version)"

type handshakeResp struct {
	Implements []string
}

type joinInfo struct {
	InterfaceNames []*iface
	Gateway        string
	GatewayIPv6    string
	HostsPath      string
	ResolvConfPath string
}

var (
	network        string
	subnet         *net.IPNet
	resolvConfPath string
	client         *docker.Client
)

func main() {
	var (
		justVersion bool
		address     string
		confdir     string
		debug       bool
	)

	flag.BoolVar(&justVersion, "version", false, "print version and exit")
	flag.BoolVar(&debug, "debug", false, "output debugging info to stderr")
	flag.StringVar(&address, "socket", "/var/run/docker-plugin/plugin.sock", "socket on which to listen")
	flag.StringVar(&confdir, "configdir", "/var/run/weave-plugin", "path in which to store temporary config files, e.g., resolv.conf")

	flag.Parse()

	if justVersion {
		fmt.Printf("weave plugin %s\n", version)
		os.Exit(0)
	}

	InitDefaultLogging(debug)

	peers := flag.Args()

	var err error
	client, err = docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		Error.Fatalf("could not connect to docker: %s", err)
	}

	router := mux.NewRouter()
	router.NotFoundHandler = http.HandlerFunc(notFound)
	router.Methods("GET").Path("/status").HandlerFunc(status)
	router.Methods("POST").Path("/Plugin.Activate").HandlerFunc(handshake)

	handleMethod := func(method string, h http.HandlerFunc) {
		router.Methods("POST").Path(fmt.Sprintf("/%s.%s", MethodReceiver, method)).HandlerFunc(h)
	}

	handleMethod("CreateNetwork", createNetwork)
	handleMethod("DeleteNetwork", deleteNetwork)
	handleMethod("CreateEndpoint", createEndpoint)
	handleMethod("DeleteEndpoint", deleteEndpoint)
	handleMethod("EndpointOperInfo", infoEndpoint)
	handleMethod("Join", joinEndpoint)
	handleMethod("Leave", leaveEndpoint)

	// Put the docker bridge IP into a resolv.conf to be used later.
	nameserver := "nameserver 172.17.42.1\n" // get this from the docker bridge
	resolvConfPath = filepath.Join(confdir, "resolv.conf")
	ioutil.WriteFile(resolvConfPath, []byte(nameserver), os.ModePerm)

	var listener net.Listener

	listener, err = net.Listen("unix", address)
	if err != nil {
		Error.Fatalf("[plugin] Unable to listen on %s: %s", address, err)
	}

	sub := "10.2.0.0/16"

	_, subnet, err = net.ParseCIDR(sub)
	if err != nil {
		Error.Fatalf("Invalid subnet CIDR %s", sub)
	}
	weaveArgs := []string{"launch", "-iprange", subnet.String()}
	if err = runWeaveCmd(append(weaveArgs, peers...)); err != nil {
		Error.Fatal("Problem launching Weave: " + err.Error())
	}

	if err = http.Serve(listener, router); err != nil {
		Error.Fatalf("[plugin] Internal error: %s", err)
	}
}

func notFound(w http.ResponseWriter, r *http.Request) {
	Warning.Printf("[plugin] Not found: %+v", r)
	http.NotFound(w, r)
}

func sendError(w http.ResponseWriter, msg string, code int) {
	Error.Printf("%d %s", code, msg)
	http.Error(w, msg, code)
}

func errorResponsef(w http.ResponseWriter, fmtString string, item ...interface{}) {
	json.NewEncoder(w).Encode(map[string]string{
		"Err": fmt.Sprintf(fmtString, item...),
	})
}

func objectResponse(w http.ResponseWriter, obj interface{}) {
	if err := json.NewEncoder(w).Encode(obj); err != nil {
		sendError(w, "Could not JSON encode response", http.StatusInternalServerError)
		return
	}
}

func emptyResponse(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(map[string]string{})
}

func handshake(w http.ResponseWriter, r *http.Request) {
	err := json.NewEncoder(w).Encode(&handshakeResp{
		[]string{"NetworkDriver"},
	})
	if err != nil {
		Error.Fatal("handshake encode:", err)
		sendError(w, "encode error", http.StatusInternalServerError)
		return
	}
	Info.Printf("Handshake completed")
}

func status(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, fmt.Sprintln("weave plugin", version))
}

type networkCreate struct {
	NetworkID string
	Options   map[string]interface{}
}

func createNetwork(w http.ResponseWriter, r *http.Request) {
	var create networkCreate
	err := json.NewDecoder(r.Body).Decode(&create)
	if err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Create network request %+v", &create)

	if network != "" {
		errorResponsef(w, "You get just one network, and you already made %s", network)
		return
	}

	network = create.NetworkID
	emptyResponse(w)
	Info.Printf("Create network %s", network)
}

type networkDelete struct {
	NetworkID string
}

func deleteNetwork(w http.ResponseWriter, r *http.Request) {
	var delete networkDelete
	if err := json.NewDecoder(r.Body).Decode(&delete); err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Delete network request: %+v", &delete)
	if delete.NetworkID != network {
		errorResponsef(w, "Network %s not found", delete.NetworkID)
		return
	}
	if _, err := getWeaveCmd([]string{"stop"}); err != nil {
		sendError(w, "Unable to stop weave: "+err.Error(), http.StatusInternalServerError)
		return
	}
	network = ""
	emptyResponse(w)
	Info.Printf("Destroy network %s", delete.NetworkID)
}

type endpointCreate struct {
	NetworkID  string
	EndpointID string
	Interfaces []*iface
	Options    map[string]interface{}
}

type iface struct {
	ID         int
	SrcName    string
	DstName    string
	Address    string
	MacAddress string
}

type endpointResponse struct {
	Interfaces []*iface
}

func createEndpoint(w http.ResponseWriter, r *http.Request) {
	var create endpointCreate
	if err := json.NewDecoder(r.Body).Decode(&create); err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Create endpoint request %+v", &create)
	netID := create.NetworkID
	endID := create.EndpointID

	if netID != network {
		errorResponsef(w, "No such network %s", netID)
		return
	}

	ip, err := getIP(endID)
	if err != nil {
		Warning.Printf("Error allocating IP:", err)
		sendError(w, "Unable to allocate IP", http.StatusInternalServerError)
		return
	}
	prefix, _ := subnet.Mask.Size()
	mac := makeMac(ip)

	respIface := &iface{
		Address: (&net.IPNet{
			ip,
			net.CIDRMask(prefix, 32),
		}).String(),
		MacAddress: mac,
	}
	resp := &endpointResponse{
		Interfaces: []*iface{respIface},
	}

	Debug.Printf("Create: %+v", &resp)
	objectResponse(w, resp)
	Info.Printf("Create endpoint %s %+v", endID, resp)
}

type endpointDelete struct {
	NetworkID  string
	EndpointID string
}

func deleteEndpoint(w http.ResponseWriter, r *http.Request) {
	var delete endpointDelete
	if err := json.NewDecoder(r.Body).Decode(&delete); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Delete endpoint request: %+v", &delete)
	// TODO
	emptyResponse(w)
	Info.Printf("Delete endpoint %s", delete.EndpointID)
}

type endpointInfoReq struct {
	NetworkID  string
	EndpointID string
}

type endpointInfo struct {
	Value map[string]interface{}
}

func infoEndpoint(w http.ResponseWriter, r *http.Request) {
	var info endpointInfoReq
	if err := json.NewDecoder(r.Body).Decode(&info); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Endpoint info request: %+v", &info)
	objectResponse(w, &endpointInfo{Value: map[string]interface{}{}})
	Info.Printf("Endpoint info %s", info.EndpointID)
}

type join struct {
	NetworkID  string
	EndpointID string
	SandboxKey string
	Options    map[string]interface{}
}

type joinResponse struct {
	HostsPath      string
	ResolvConfPath string
	Gateway        string
	InterfaceNames []*iface
}

func joinEndpoint(w http.ResponseWriter, r *http.Request) {
	var j join
	if err := json.NewDecoder(r.Body).Decode(&j); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Join request: %+v", &j)

	endID := j.EndpointID

	// use the endpoint ID bytes to make veth names, and cross fingers
	localName := "vethwel" + endID[:5]
	guestName := "vethweg" + endID[:5]
	// create and attach local name to the bridge

	var (
		bridge *netlink.Bridge
		local  *netlink.Veth
	)

	local = &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: localName},
		PeerName:  guestName,
	}
	if err := netlink.LinkAdd(local); err != nil {
		Error.Print(err)
		errorResponsef(w, "could not create veth pair")
		return
	}

	if maybeBridge, err := netlink.LinkByName(WeaveBridge); err != nil {
		Error.Print(err)
		errorResponsef(w, `bridge "%s" not present`, WeaveBridge)
		return
	} else {
		var ok bool
		if bridge, ok = maybeBridge.(*netlink.Bridge); !ok {
			Error.Printf("%s is %+v", WeaveBridge, maybeBridge)
			errorResponsef(w, `device "%s" not a bridge`, WeaveBridge)
			return
		}
	}
	if netlink.LinkSetMaster(local, bridge) != nil || netlink.LinkSetUp(local) != nil {
		errorResponsef(w, `unable to bring veth up`)
		return
	}

	ifname := &iface{
		SrcName: guestName,
		DstName: "ethwe",
		ID:      0,
	}
	res := &joinResponse{
		InterfaceNames: []*iface{ifname},
		ResolvConfPath: resolvConfPath,
	}

	objectResponse(w, res)
	Info.Printf("Join endpoint %s:%s to %s", j.NetworkID, j.EndpointID, j.SandboxKey)
}

type leave struct {
	NetworkID  string
	EndpointID string
	Options    map[string]interface{}
}

func leaveEndpoint(w http.ResponseWriter, r *http.Request) {
	var l leave
	if err := json.NewDecoder(r.Body).Decode(&l); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	// TODO
	Debug.Printf("Leave request: %+v", &l)
	emptyResponse(w)
	Info.Printf("Leave %s:%s", l.NetworkID, l.EndpointID)
}

// ===

func getWeaveCmd(args []string) (string, error) {
	cmd := exec.Command("./weave", args...)
	cmd.Env = []string{"PATH=/usr/bin:/usr/local/bin", "WEAVE_DEBUG=true"}
	var buf bytes.Buffer
	cmd.Stderr = &buf
	out, err := cmd.Output()
	if err != nil {
		Warning.Print(buf.String())
	}
	return string(out), err
}

func runWeaveCmd(args []string) error {
	cmd := exec.Command("./weave", append([]string{"--local"}, args...)...)
	cmd.Env = []string{
		"PROCFS=/hostproc",
		"PATH=/sbin:/usr/sbin:/bin:/usr/bin:/usr/local/bin",
		"WEAVE_DEBUG=true"}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		Warning.Print(err.Error())
	}
	return err
}

func getSwitchIP() (string, error) {
	info, err := client.InspectContainer(WeaveContainer)
	if err != nil {
		return "", err
	}
	return info.NetworkSettings.IPAddress, nil
}

// assumed to be in the subnet
func getIP(ID string) (net.IP, error) {
	weaveip, err := getSwitchIP()
	if err != nil {
		return nil, err
	}
	res, err := http.Post(fmt.Sprintf("http://%s:6784/ip/%s", weaveip, ID), "", nil)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("Received status %d from IPAM", res.StatusCode)
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	ip, _, err := net.ParseCIDR(string(body))
	return ip, err
}

func makeMac(ip net.IP) string {
	hw := make(net.HardwareAddr, 6)
	hw[0] = 0x7a
	hw[1] = 0x42
	copy(hw[2:], ip.To4())
	return hw.String()
}
<commit_msg>Release IPs and delete veths<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/gorilla/mux"
	"github.com/vishvananda/netlink"
	. "github.com/weaveworks/weave/common"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
)

const (
	MethodReceiver = "NetworkDriver"
	WeaveContainer = "weave"
	WeaveBridge    = "weave"
)

var version = "(unreleased version)"

type handshakeResp struct {
	Implements []string
}

type joinInfo struct {
	InterfaceNames []*iface
	Gateway        string
	GatewayIPv6    string
	HostsPath      string
	ResolvConfPath string
}

var (
	network        string
	subnet         *net.IPNet
	resolvConfPath string
	client         *docker.Client
)

func main() {
	var (
		justVersion bool
		address     string
		confdir     string
		debug       bool
	)

	flag.BoolVar(&justVersion, "version", false, "print version and exit")
	flag.BoolVar(&debug, "debug", false, "output debugging info to stderr")
	flag.StringVar(&address, "socket", "/var/run/docker-plugin/plugin.sock", "socket on which to listen")
	flag.StringVar(&confdir, "configdir", "/var/run/weave-plugin", "path in which to store temporary config files, e.g., resolv.conf")

	flag.Parse()

	if justVersion {
		fmt.Printf("weave plugin %s\n", version)
		os.Exit(0)
	}

	InitDefaultLogging(debug)

	peers := flag.Args()

	var err error
	client, err = docker.NewClient("unix:///var/run/docker.sock")
	if err != nil {
		Error.Fatalf("could not connect to docker: %s", err)
	}

	router := mux.NewRouter()
	router.NotFoundHandler = http.HandlerFunc(notFound)
	router.Methods("GET").Path("/status").HandlerFunc(status)
	router.Methods("POST").Path("/Plugin.Activate").HandlerFunc(handshake)

	handleMethod := func(method string, h http.HandlerFunc) {
		router.Methods("POST").Path(fmt.Sprintf("/%s.%s", MethodReceiver, method)).HandlerFunc(h)
	}

	handleMethod("CreateNetwork", createNetwork)
	handleMethod("DeleteNetwork", deleteNetwork)
	handleMethod("CreateEndpoint", createEndpoint)
	handleMethod("DeleteEndpoint", deleteEndpoint)
	handleMethod("EndpointOperInfo", infoEndpoint)
	handleMethod("Join", joinEndpoint)
	handleMethod("Leave", leaveEndpoint)

	// Put the docker bridge IP into a resolv.conf to be used later.
	nameserver := "nameserver 172.17.42.1\n" // get this from the docker bridge
	resolvConfPath = filepath.Join(confdir, "resolv.conf")
	ioutil.WriteFile(resolvConfPath, []byte(nameserver), os.ModePerm)

	var listener net.Listener

	listener, err = net.Listen("unix", address)
	if err != nil {
		Error.Fatalf("[plugin] Unable to listen on %s: %s", address, err)
	}

	sub := "10.2.0.0/16"

	_, subnet, err = net.ParseCIDR(sub)
	if err != nil {
		Error.Fatalf("Invalid subnet CIDR %s", sub)
	}
	weaveArgs := []string{"launch", "-iprange", subnet.String()}
	if err = runWeaveCmd(append(weaveArgs, peers...)); err != nil {
		Error.Fatal("Problem launching Weave: " + err.Error())
	}

	if err = http.Serve(listener, router); err != nil {
		Error.Fatalf("[plugin] Internal error: %s", err)
	}
}

func notFound(w http.ResponseWriter, r *http.Request) {
	Warning.Printf("[plugin] Not found: %+v", r)
	http.NotFound(w, r)
}

func sendError(w http.ResponseWriter, msg string, code int) {
	Error.Printf("%d %s", code, msg)
	http.Error(w, msg, code)
}

func errorResponsef(w http.ResponseWriter, fmtString string, item ...interface{}) {
	json.NewEncoder(w).Encode(map[string]string{
		"Err": fmt.Sprintf(fmtString, item...),
	})
}

func objectResponse(w http.ResponseWriter, obj interface{}) {
	if err := json.NewEncoder(w).Encode(obj); err != nil {
		sendError(w, "Could not JSON encode response", http.StatusInternalServerError)
		return
	}
}

func emptyResponse(w http.ResponseWriter) {
	json.NewEncoder(w).Encode(map[string]string{})
}

func handshake(w http.ResponseWriter, r *http.Request) {
	err := json.NewEncoder(w).Encode(&handshakeResp{
		[]string{"NetworkDriver"},
	})
	if err != nil {
		Error.Fatal("handshake encode:", err)
		sendError(w, "encode error", http.StatusInternalServerError)
		return
	}
	Info.Printf("Handshake completed")
}

func status(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, fmt.Sprintln("weave plugin", version))
}

type networkCreate struct {
	NetworkID string
	Options   map[string]interface{}
}

func createNetwork(w http.ResponseWriter, r *http.Request) {
	var create networkCreate
	err := json.NewDecoder(r.Body).Decode(&create)
	if err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Create network request %+v", &create)

	if network != "" {
		errorResponsef(w, "You get just one network, and you already made %s", network)
		return
	}

	network = create.NetworkID
	emptyResponse(w)
	Info.Printf("Create network %s", network)
}

type networkDelete struct {
	NetworkID string
}

func deleteNetwork(w http.ResponseWriter, r *http.Request) {
	var delete networkDelete
	if err := json.NewDecoder(r.Body).Decode(&delete); err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Delete network request: %+v", &delete)
	if delete.NetworkID != network {
		errorResponsef(w, "Network %s not found", delete.NetworkID)
		return
	}
	if _, err := getWeaveCmd([]string{"stop"}); err != nil {
		sendError(w, "Unable to stop weave: "+err.Error(), http.StatusInternalServerError)
		return
	}
	network = ""
	emptyResponse(w)
	Info.Printf("Destroy network %s", delete.NetworkID)
}

type endpointCreate struct {
	NetworkID  string
	EndpointID string
	Interfaces []*iface
	Options    map[string]interface{}
}

type iface struct {
	ID         int
	SrcName    string
	DstName    string
	Address    string
	MacAddress string
}

type endpointResponse struct {
	Interfaces []*iface
}

func createEndpoint(w http.ResponseWriter, r *http.Request) {
	var create endpointCreate
	if err := json.NewDecoder(r.Body).Decode(&create); err != nil {
		sendError(w, "Unable to decode JSON payload: "+err.Error(), http.StatusBadRequest)
		return
	}
	Debug.Printf("Create endpoint request %+v", &create)
	netID := create.NetworkID
	endID := create.EndpointID

	if netID != network {
		errorResponsef(w, "No such network %s", netID)
		return
	}

	ip, err := getIP(endID)
	if err != nil {
		Warning.Printf("Error allocating IP:", err)
		sendError(w, "Unable to allocate IP", http.StatusInternalServerError)
		return
	}
	prefix, _ := subnet.Mask.Size()
	mac := makeMac(ip)

	respIface := &iface{
		Address: (&net.IPNet{
			ip,
			net.CIDRMask(prefix, 32),
		}).String(),
		MacAddress: mac,
	}
	resp := &endpointResponse{
		Interfaces: []*iface{respIface},
	}

	Debug.Printf("Create: %+v", &resp)
	objectResponse(w, resp)
	Info.Printf("Create endpoint %s %+v", endID, resp)
}

type endpointDelete struct {
	NetworkID  string
	EndpointID string
}

func deleteEndpoint(w http.ResponseWriter, r *http.Request) {
	var delete endpointDelete
	if err := json.NewDecoder(r.Body).Decode(&delete); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Delete endpoint request: %+v", &delete)
	if err := releaseIP(delete.EndpointID); err != nil {
		Warning.Printf("error releasing IP: %s", err)
	}
	emptyResponse(w)
	Info.Printf("Delete endpoint %s", delete.EndpointID)
}

type endpointInfoReq struct {
	NetworkID  string
	EndpointID string
}

type endpointInfo struct {
	Value map[string]interface{}
}

func infoEndpoint(w http.ResponseWriter, r *http.Request) {
	var info endpointInfoReq
	if err := json.NewDecoder(r.Body).Decode(&info); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Endpoint info request: %+v", &info)
	objectResponse(w, &endpointInfo{Value: map[string]interface{}{}})
	Info.Printf("Endpoint info %s", info.EndpointID)
}

type join struct {
	NetworkID  string
	EndpointID string
	SandboxKey string
	Options    map[string]interface{}
}

type joinResponse struct {
	HostsPath      string
	ResolvConfPath string
	Gateway        string
	InterfaceNames []*iface
}

func joinEndpoint(w http.ResponseWriter, r *http.Request) {
	var j join
	if err := json.NewDecoder(r.Body).Decode(&j); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Join request: %+v", &j)

	endID := j.EndpointID

	// create and attach local name to the bridge
	local := vethPair(endID[:5])
	if err := netlink.LinkAdd(local); err != nil {
		Error.Print(err)
		errorResponsef(w, "could not create veth pair")
		return
	}

	var bridge *netlink.Bridge
	if maybeBridge, err := netlink.LinkByName(WeaveBridge); err != nil {
		Error.Print(err)
		errorResponsef(w, `bridge "%s" not present`, WeaveBridge)
		return
	} else {
		var ok bool
		if bridge, ok = maybeBridge.(*netlink.Bridge); !ok {
			Error.Printf("%s is %+v", WeaveBridge, maybeBridge)
			errorResponsef(w, `device "%s" not a bridge`, WeaveBridge)
			return
		}
	}
	if netlink.LinkSetMaster(local, bridge) != nil || netlink.LinkSetUp(local) != nil {
		errorResponsef(w, `unable to bring veth up`)
		return
	}

	ifname := &iface{
		SrcName: local.PeerName,
		DstName: "ethwe",
		ID:      0,
	}
	res := &joinResponse{
		InterfaceNames: []*iface{ifname},
		ResolvConfPath: resolvConfPath,
	}

	objectResponse(w, res)
	Info.Printf("Join endpoint %s:%s to %s", j.NetworkID, j.EndpointID, j.SandboxKey)
}

type leave struct {
	NetworkID  string
	EndpointID string
	Options    map[string]interface{}
}

func leaveEndpoint(w http.ResponseWriter, r *http.Request) {
	var l leave
	if err := json.NewDecoder(r.Body).Decode(&l); err != nil {
		sendError(w, "Could not decode JSON encode payload", http.StatusBadRequest)
		return
	}
	Debug.Printf("Leave request: %+v", &l)

	local := vethPair(l.EndpointID[:5])
	if err := netlink.LinkDel(local); err != nil {
		Warning.Printf("unable to delete veth on leave: %s", err)
	}
	emptyResponse(w)
	Info.Printf("Leave %s:%s", l.NetworkID, l.EndpointID)
}

// ===

func vethPair(suffix string) *netlink.Veth {
	return &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "vethwl" + suffix},
		PeerName:  "vethwg" + suffix,
	}
}

func getWeaveCmd(args []string) (string, error) {
	cmd := exec.Command("./weave", args...)
	cmd.Env = []string{"PATH=/usr/bin:/usr/local/bin", "WEAVE_DEBUG=true"}
	var buf bytes.Buffer
	cmd.Stderr = &buf
	out, err := cmd.Output()
	if err != nil {
		Warning.Print(buf.String())
	}
	return string(out), err
}

func runWeaveCmd(args []string) error {
	cmd := exec.Command("./weave", append([]string{"--local"}, args...)...)
	cmd.Env = []string{
		"PROCFS=/hostproc",
		"PATH=/sbin:/usr/sbin:/bin:/usr/bin:/usr/local/bin",
		"WEAVE_DEBUG=true"}
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		Warning.Print(err.Error())
	}
	return err
}

func getSwitchIP() (string, error) {
	info, err := client.InspectContainer(WeaveContainer)
	if err != nil {
		return "", err
	}
	return info.NetworkSettings.IPAddress, nil
}

// assumed to be in the subnet
func getIP(ID string) (net.IP, error) {
	weaveip, err := getSwitchIP()
	if err != nil {
		return nil, err
	}
	res, err := http.Post(fmt.Sprintf("http://%s:6784/ip/%s", weaveip, ID), "", nil)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("Received status %d from IPAM", res.StatusCode)
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	ip, _, err := net.ParseCIDR(string(body))
	return ip, err
}

func releaseIP(ID string) error {
	weaveip, err := getSwitchIP()
	if err != nil {
		return err
	}
	req, err := http.NewRequest("DELETE", fmt.Sprintf("http://%s:6784/ip/%s", weaveip, ID), nil)
	if err != nil {
		return err
	}
	cl := &http.Client{}
	res, err := cl.Do(req)
	if err != nil {
		return err
	}
	if res.StatusCode != http.StatusOK && res.StatusCode != http.StatusNoContent {
		return fmt.Errorf("unexpected HTTP status code from IP release: %d", res.StatusCode)
	}
	return nil
}

func makeMac(ip net.IP) string {
	hw := make(net.HardwareAddr, 6)
	hw[0] = 0x7a
	hw[1] = 0x42
	copy(hw[2:], ip.To4())
	return hw.String()
}
<|endoftext|>


<commit_before>package printer

import (
	"fmt"
	"text/template"

	"github.com/davyxu/tabtoy/exportorv2/i18n"
	"github.com/davyxu/tabtoy/exportorv2/model"
	"strings"
)

const csharpTemplate = `// Generated by github.com/davyxu/tabtoy
// Version: {{.ToolVersion}}
// DO NOT EDIT!!
using System.Collections.Generic;

namespace {{.Namespace}}{{$globalIndex:=.Indexes}}{{$verticalFields:=.VerticalFields}}
{
	{{range .Enums}}
	// Defined in table: {{.DefinedTable}}
	public enum {{.Name}}
	{
	{{range .Fields}}
		{{.Alias}}
		{{.FieldDescriptor.Name}} = {{.FieldDescriptor.EnumValue}}, {{.Comment}}
	{{end}}
	}
	{{end}}
	{{range .Classes}}

	// Defined in table: {{.DefinedTable}}
	{{.CSClassHeader}}
	public partial class {{.Name}}
	{
	{{if .IsCombine}}
		public tabtoy.Logger TableLogger = new tabtoy.Logger();
	{{end}}
	{{range .Fields}}
		{{.Comment}}
		{{.TypeCode}} {{.Alias}}
	{{end}}
	{{if .IsCombine}}
		#region Index code
	 	{{range $globalIndex}}Dictionary<{{.IndexType}}, {{.RowType}}> _{{.RowName}}By{{.IndexName}} = new Dictionary<{{.IndexType}}, {{.RowType}}>();
        public {{.RowType}} Get{{.RowName}}By{{.IndexName}}({{.IndexType}} {{.IndexName}}, {{.RowType}} def = default({{.RowType}}))
        {
            {{.RowType}} ret;
            if ( _{{.RowName}}By{{.IndexName}}.TryGetValue( {{.IndexName}}, out ret ) )
            {
                return ret;
            }

			if ( def == default({{.RowType}}) )
			{
				TableLogger.ErrorLine("Get{{.RowName}}By{{.IndexName}} failed, {{.IndexName}}: {0}", {{.IndexName}});
			}

            return def;
        }
		{{end}}
	{{range $verticalFields}}
		public {{.StructName}} Get{{.Name}}( )
		{
			return {{.Name}}[0];
		}
	{{end}}
		#endregion
		#region Deserialize code
		{{range $.Classes}}
		static tabtoy.DeserializeHandler<{{.Name}}> {{.Name}}DeserializeHandler = new tabtoy.DeserializeHandler<{{.Name}}>(Deserialize);
		public static void Deserialize( {{.Name}} ins, tabtoy.DataReader reader )
		{
 			int tag = -1;
            while ( -1 != (tag = reader.ReadTag()))
            {
                switch (tag)
                { {{range .Fields}}
                	case {{.Tag}}:
                	{
						{{.ReadCode}}
                	}
                	break; {{end}}
                }
             }

			{{range $a, $row :=.IndexedFields}}
			// Build {{$row.FieldDescriptor.Name}} Index
			for( int i = 0;i< ins.{{$row.FieldDescriptor.Name}}.Count;i++)
			{
				var element = ins.{{$row.FieldDescriptor.Name}}[i];
				{{range $b, $key := .IndexKeys}}
				ins._{{$row.FieldDescriptor.Name}}By{{$key.Name}}.Add(element.{{$key.Name}}, element);
				{{end}}
			}
			{{end}}
		}{{end}}
		#endregion
	{{end}}

	} {{end}}

}
`

type indexField struct {
	TableIndex
}

func (self indexField) IndexName() string {
	return self.Index.Name
}

func (self indexField) RowType() string {
	return self.Row.Complex.Name
}

func (self indexField) RowName() string {
	return self.Row.Name
}

func (self indexField) IndexType() string {

	switch self.Index.Type {
	case model.FieldType_Int32:
		return "int"
	case model.FieldType_UInt32:
		return "uint"
	case model.FieldType_Int64:
		return "long"
	case model.FieldType_UInt64:
		return "ulong"
	case model.FieldType_String:
		return "string"
	case model.FieldType_Float:
		return "float"
	case model.FieldType_Bool:
		return "bool"
	case model.FieldType_Enum:

		return self.Index.Complex.Name
	default:
		log.Errorf("%s can not be index ", self.Index.String())
	}

	return "unknown"
}

type csharpField struct {
	*model.FieldDescriptor

	IndexKeys []*model.FieldDescriptor

	parentStruct *structModel
}

func (self csharpField) Alias() string {

	v := self.FieldDescriptor.Meta.GetString("Alias")
	if v == "" {
		return ""
	}

	return "// " + v
}

func (self csharpField) Comment() string {

	if self.FieldDescriptor.Comment == "" {
		return ""
	}

	// zjwps 建议修改
	return "/// <summary> \n		///" + strings.Replace(self.FieldDescriptor.Comment, "\n", "\n		///", -1) + "\n		/// </summary>"
}

func (self csharpField) ReadCode() string {

	var baseType string

	var descHandlerCode string

	switch self.Type {
	case model.FieldType_Int32:
		baseType = "Int32"
	case model.FieldType_UInt32:
		baseType = "UInt32"
	case model.FieldType_Int64:
		baseType = "Int64"
	case model.FieldType_UInt64:
		baseType = "UInt64"
	case model.FieldType_String:
		baseType = "String"
	case model.FieldType_Float:
		baseType = "Float"
	case model.FieldType_Bool:
		baseType = "Bool"
	case model.FieldType_Enum:

		if self.Complex == nil {
			return "unknown"
		}

		baseType = fmt.Sprintf("Enum<%s>", self.Complex.Name)

	case model.FieldType_Struct:
		if self.Complex == nil {
			return "unknown"
		}

		baseType = fmt.Sprintf("Struct<%s>", self.Complex.Name)

	}

	if self.Type == model.FieldType_Struct {
		descHandlerCode = fmt.Sprintf("%sDeserializeHandler", self.Complex.Name)
	}

	if self.IsRepeated {
		return fmt.Sprintf("ins.%s.Add( reader.Read%s(%s) );", self.Name, baseType, descHandlerCode)
	} else {
		return fmt.Sprintf("ins.%s = reader.Read%s(%s);", self.Name, baseType, descHandlerCode)
	}

}

func (self csharpField) Tag() string {

	if self.parentStruct.IsCombine() {
		tag := model.MakeTag(int32(model.FieldType_Table), self.Order)

		return fmt.Sprintf("0x%x", tag)
	}

	return fmt.Sprintf("0x%x", self.FieldDescriptor.Tag())
}

func (self csharpField) StructName() string {
	if self.Complex == nil {
		return "[NotComplex]"
	}

	return self.Complex.Name
}

func (self csharpField) IsVerticalStruct() bool {
	if self.FieldDescriptor.Complex == nil {
		return false
	}

	return self.FieldDescriptor.Complex.File.Pragma.GetBool("Vertical")
}

func (self csharpField) TypeCode() string {

	var raw string

	switch self.Type {
	case model.FieldType_Int32:
		raw = "int"
	case model.FieldType_UInt32:
		raw = "uint"
	case model.FieldType_Int64:
		raw = "long"
	case model.FieldType_UInt64:
		raw = "ulong"
	case model.FieldType_String:
		raw = "string"
	case model.FieldType_Float:
		raw = "float"
	case model.FieldType_Bool:
		raw = "bool"
	case model.FieldType_Enum:
		if self.Complex == nil {
			log.Errorln("unknown enum type ", self.Type)
			return "unknown"
		}

		raw = self.Complex.Name
	case model.FieldType_Struct:
		if self.Complex == nil {
			log.Errorln("unknown struct type ", self.Type, self.FieldDescriptor.Name, self.FieldDescriptor.Parent.Name)
			return "unknown"
		}

		raw = self.Complex.Name

		// 非repeated的结构体
		if !self.IsRepeated {
			return fmt.Sprintf("public %s %s = new %s();", raw, self.Name, raw)
		}

	default:
		raw = "unknown"
	}

	if self.IsRepeated {
		return fmt.Sprintf("public List<%s> %s = new List<%s>();", raw, self.Name, raw)
	}

	return fmt.Sprintf("public %s %s = %s;", raw, self.Name, wrapCSharpDefaultValue(self.FieldDescriptor))
}

func wrapCSharpDefaultValue(fd *model.FieldDescriptor) string {
	switch fd.Type {
	case model.FieldType_Enum:
		return fmt.Sprintf("%s.%s", fd.Complex.Name, fd.DefaultValue())
	case model.FieldType_String:
		return fmt.Sprintf("\"%s\"", fd.DefaultValue())
	case model.FieldType_Float:
		return fmt.Sprintf("%sf", fd.DefaultValue())
	}

	return fd.DefaultValue()
}

type structModel struct {
	*model.Descriptor
	Fields        []csharpField
	IndexedFields []csharpField // 与csharpField.IndexKeys组成树状的索引层次
}

func (self *structModel) CSClassHeader() string {

	// zjwps 提供需求
	return self.File.Pragma.GetString("CSClassHeader")
}

func (self *structModel) DefinedTable() string {
	return self.File.Name
}

func (self *structModel) Name() string {
	return self.Descriptor.Name
}

func (self *structModel) IsCombine() bool {
	return self.Descriptor.Usage == model.DescriptorUsage_CombineStruct
}

type csharpFileModel struct {
	Namespace   string
	ToolVersion string
	Classes     []*structModel
	Enums       []*structModel
	Indexes     []indexField // 全局的索引

	VerticalFields []csharpField
}

type csharpPrinter struct {
}

func (self *csharpPrinter) Run(g *Globals) *Stream {

	tpl, err := template.New("csharp").Parse(csharpTemplate)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	var m csharpFileModel

	m.Namespace = g.FileDescriptor.Pragma.GetString("Package")
	m.ToolVersion = g.Version

	// combinestruct的全局索引
	for _, ti := range g.GlobalIndexes {

		// 索引也限制
		if !ti.Index.Parent.File.MatchTag(".cs") {
			continue
		}

		m.Indexes = append(m.Indexes, indexField{TableIndex: ti})
	}

	// 遍历所有类型
	for _, d := range g.FileDescriptor.Descriptors {

		// 这给被限制输出
		if !d.File.MatchTag(".cs") {
			log.Infof("%s: %s", i18n.String(i18n.Printer_IgnoredByOutputTag), d.Name)
			continue
		}

		var sm structModel
		sm.Descriptor = d

		switch d.Kind {
		case model.DescriptorKind_Struct:
			m.Classes = append(m.Classes, &sm)
		case model.DescriptorKind_Enum:
			m.Enums = append(m.Enums, &sm)
		}

		// 遍历字段
		for _, fd := range d.Fields {

			// 对CombineStruct的XXDefine对应的字段
			if d.Usage == model.DescriptorUsage_CombineStruct {

				// 这个字段被限制输出
				if fd.Complex != nil && !fd.Complex.File.MatchTag(".cs") {
					continue
				}

				// 这个结构有索引才创建
				if fd.Complex != nil && len(fd.Complex.Indexes) > 0 {

					// 被索引的结构
					indexedField := csharpField{FieldDescriptor: fd, parentStruct: &sm}

					// 索引字段
					for _, key := range fd.Complex.Indexes {
						indexedField.IndexKeys = append(indexedField.IndexKeys, key)
					}

					sm.IndexedFields = append(sm.IndexedFields, indexedField)
				}

				if fd.Complex != nil && fd.Complex.File.Pragma.GetBool("Vertical") {
					m.VerticalFields = append(m.VerticalFields, csharpField{FieldDescriptor: fd, parentStruct: &sm})
				}

			}

			csField := csharpField{FieldDescriptor: fd, parentStruct: &sm}

			sm.Fields = append(sm.Fields, csField)

		}

	}

	bf := NewStream()

	err = tpl.Execute(bf.Buffer(), &m)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	return bf
}

func init() {

	RegisterPrinter("cs", &csharpPrinter{})

}
<commit_msg>对齐生成注释代码<commit_after>package printer

import (
	"fmt"
	"text/template"

	"github.com/davyxu/tabtoy/exportorv2/i18n"
	"github.com/davyxu/tabtoy/exportorv2/model"
	"strings"
)

const csharpTemplate = `// Generated by github.com/davyxu/tabtoy
// Version: {{.ToolVersion}}
// DO NOT EDIT!!
using System.Collections.Generic;

namespace {{.Namespace}}{{$globalIndex:=.Indexes}}{{$verticalFields:=.VerticalFields}}
{
	{{range .Enums}}
	// Defined in table: {{.DefinedTable}}
	public enum {{.Name}}
	{
	{{range .Fields}}
		{{.Comment}}
		{{.FieldDescriptor.Name}} = {{.FieldDescriptor.EnumValue}}, {{.Alias}}
	{{end}}
	}
	{{end}}
	{{range .Classes}}

	// Defined in table: {{.DefinedTable}}
	{{.CSClassHeader}}
	public partial class {{.Name}}
	{
	{{if .IsCombine}}
		public tabtoy.Logger TableLogger = new tabtoy.Logger();
	{{end}}
	{{range .Fields}}
		{{.Comment}}
		{{.TypeCode}} {{.Alias}}
	{{end}}
	{{if .IsCombine}}
		#region Index code
	 	{{range $globalIndex}}Dictionary<{{.IndexType}}, {{.RowType}}> _{{.RowName}}By{{.IndexName}} = new Dictionary<{{.IndexType}}, {{.RowType}}>();
        public {{.RowType}} Get{{.RowName}}By{{.IndexName}}({{.IndexType}} {{.IndexName}}, {{.RowType}} def = default({{.RowType}}))
        {
            {{.RowType}} ret;
            if ( _{{.RowName}}By{{.IndexName}}.TryGetValue( {{.IndexName}}, out ret ) )
            {
                return ret;
            }

			if ( def == default({{.RowType}}) )
			{
				TableLogger.ErrorLine("Get{{.RowName}}By{{.IndexName}} failed, {{.IndexName}}: {0}", {{.IndexName}});
			}

            return def;
        }
		{{end}}
	{{range $verticalFields}}
		public {{.StructName}} Get{{.Name}}( )
		{
			return {{.Name}}[0];
		}
	{{end}}
		#endregion
		#region Deserialize code
		{{range $.Classes}}
		static tabtoy.DeserializeHandler<{{.Name}}> {{.Name}}DeserializeHandler = new tabtoy.DeserializeHandler<{{.Name}}>(Deserialize);
		public static void Deserialize( {{.Name}} ins, tabtoy.DataReader reader )
		{
 			int tag = -1;
            while ( -1 != (tag = reader.ReadTag()))
            {
                switch (tag)
                { {{range .Fields}}
                	case {{.Tag}}:
                	{
						{{.ReadCode}}
                	}
                	break; {{end}}
                }
             }

			{{range $a, $row :=.IndexedFields}}
			// Build {{$row.FieldDescriptor.Name}} Index
			for( int i = 0;i< ins.{{$row.FieldDescriptor.Name}}.Count;i++)
			{
				var element = ins.{{$row.FieldDescriptor.Name}}[i];
				{{range $b, $key := .IndexKeys}}
				ins._{{$row.FieldDescriptor.Name}}By{{$key.Name}}.Add(element.{{$key.Name}}, element);
				{{end}}
			}
			{{end}}
		}{{end}}
		#endregion
	{{end}}

	} {{end}}

}
`

type indexField struct {
	TableIndex
}

func (self indexField) IndexName() string {
	return self.Index.Name
}

func (self indexField) RowType() string {
	return self.Row.Complex.Name
}

func (self indexField) RowName() string {
	return self.Row.Name
}

func (self indexField) IndexType() string {

	switch self.Index.Type {
	case model.FieldType_Int32:
		return "int"
	case model.FieldType_UInt32:
		return "uint"
	case model.FieldType_Int64:
		return "long"
	case model.FieldType_UInt64:
		return "ulong"
	case model.FieldType_String:
		return "string"
	case model.FieldType_Float:
		return "float"
	case model.FieldType_Bool:
		return "bool"
	case model.FieldType_Enum:

		return self.Index.Complex.Name
	default:
		log.Errorf("%s can not be index ", self.Index.String())
	}

	return "unknown"
}

type csharpField struct {
	*model.FieldDescriptor

	IndexKeys []*model.FieldDescriptor

	parentStruct *structModel
}

func (self csharpField) Alias() string {

	v := self.FieldDescriptor.Meta.GetString("Alias")
	if v == "" {
		return ""
	}

	return "// " + v
}

func (self csharpField) Comment() string {

	if self.FieldDescriptor.Comment == "" {
		return ""
	}

	// zjwps 建议修改
	return "/// <summary> \n		/// " + strings.Replace(self.FieldDescriptor.Comment, "\n", "\n		///", -1) + "\n		/// </summary>"
}

func (self csharpField) ReadCode() string {

	var baseType string

	var descHandlerCode string

	switch self.Type {
	case model.FieldType_Int32:
		baseType = "Int32"
	case model.FieldType_UInt32:
		baseType = "UInt32"
	case model.FieldType_Int64:
		baseType = "Int64"
	case model.FieldType_UInt64:
		baseType = "UInt64"
	case model.FieldType_String:
		baseType = "String"
	case model.FieldType_Float:
		baseType = "Float"
	case model.FieldType_Bool:
		baseType = "Bool"
	case model.FieldType_Enum:

		if self.Complex == nil {
			return "unknown"
		}

		baseType = fmt.Sprintf("Enum<%s>", self.Complex.Name)

	case model.FieldType_Struct:
		if self.Complex == nil {
			return "unknown"
		}

		baseType = fmt.Sprintf("Struct<%s>", self.Complex.Name)

	}

	if self.Type == model.FieldType_Struct {
		descHandlerCode = fmt.Sprintf("%sDeserializeHandler", self.Complex.Name)
	}

	if self.IsRepeated {
		return fmt.Sprintf("ins.%s.Add( reader.Read%s(%s) );", self.Name, baseType, descHandlerCode)
	} else {
		return fmt.Sprintf("ins.%s = reader.Read%s(%s);", self.Name, baseType, descHandlerCode)
	}

}

func (self csharpField) Tag() string {

	if self.parentStruct.IsCombine() {
		tag := model.MakeTag(int32(model.FieldType_Table), self.Order)

		return fmt.Sprintf("0x%x", tag)
	}

	return fmt.Sprintf("0x%x", self.FieldDescriptor.Tag())
}

func (self csharpField) StructName() string {
	if self.Complex == nil {
		return "[NotComplex]"
	}

	return self.Complex.Name
}

func (self csharpField) IsVerticalStruct() bool {
	if self.FieldDescriptor.Complex == nil {
		return false
	}

	return self.FieldDescriptor.Complex.File.Pragma.GetBool("Vertical")
}

func (self csharpField) TypeCode() string {

	var raw string

	switch self.Type {
	case model.FieldType_Int32:
		raw = "int"
	case model.FieldType_UInt32:
		raw = "uint"
	case model.FieldType_Int64:
		raw = "long"
	case model.FieldType_UInt64:
		raw = "ulong"
	case model.FieldType_String:
		raw = "string"
	case model.FieldType_Float:
		raw = "float"
	case model.FieldType_Bool:
		raw = "bool"
	case model.FieldType_Enum:
		if self.Complex == nil {
			log.Errorln("unknown enum type ", self.Type)
			return "unknown"
		}

		raw = self.Complex.Name
	case model.FieldType_Struct:
		if self.Complex == nil {
			log.Errorln("unknown struct type ", self.Type, self.FieldDescriptor.Name, self.FieldDescriptor.Parent.Name)
			return "unknown"
		}

		raw = self.Complex.Name

		// 非repeated的结构体
		if !self.IsRepeated {
			return fmt.Sprintf("public %s %s = new %s();", raw, self.Name, raw)
		}

	default:
		raw = "unknown"
	}

	if self.IsRepeated {
		return fmt.Sprintf("public List<%s> %s = new List<%s>();", raw, self.Name, raw)
	}

	return fmt.Sprintf("public %s %s = %s;", raw, self.Name, wrapCSharpDefaultValue(self.FieldDescriptor))
}

func wrapCSharpDefaultValue(fd *model.FieldDescriptor) string {
	switch fd.Type {
	case model.FieldType_Enum:
		return fmt.Sprintf("%s.%s", fd.Complex.Name, fd.DefaultValue())
	case model.FieldType_String:
		return fmt.Sprintf("\"%s\"", fd.DefaultValue())
	case model.FieldType_Float:
		return fmt.Sprintf("%sf", fd.DefaultValue())
	}

	return fd.DefaultValue()
}

type structModel struct {
	*model.Descriptor
	Fields        []csharpField
	IndexedFields []csharpField // 与csharpField.IndexKeys组成树状的索引层次
}

func (self *structModel) CSClassHeader() string {

	// zjwps 提供需求
	return self.File.Pragma.GetString("CSClassHeader")
}

func (self *structModel) DefinedTable() string {
	return self.File.Name
}

func (self *structModel) Name() string {
	return self.Descriptor.Name
}

func (self *structModel) IsCombine() bool {
	return self.Descriptor.Usage == model.DescriptorUsage_CombineStruct
}

type csharpFileModel struct {
	Namespace   string
	ToolVersion string
	Classes     []*structModel
	Enums       []*structModel
	Indexes     []indexField // 全局的索引

	VerticalFields []csharpField
}

type csharpPrinter struct {
}

func (self *csharpPrinter) Run(g *Globals) *Stream {

	tpl, err := template.New("csharp").Parse(csharpTemplate)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	var m csharpFileModel

	m.Namespace = g.FileDescriptor.Pragma.GetString("Package")
	m.ToolVersion = g.Version

	// combinestruct的全局索引
	for _, ti := range g.GlobalIndexes {

		// 索引也限制
		if !ti.Index.Parent.File.MatchTag(".cs") {
			continue
		}

		m.Indexes = append(m.Indexes, indexField{TableIndex: ti})
	}

	// 遍历所有类型
	for _, d := range g.FileDescriptor.Descriptors {

		// 这给被限制输出
		if !d.File.MatchTag(".cs") {
			log.Infof("%s: %s", i18n.String(i18n.Printer_IgnoredByOutputTag), d.Name)
			continue
		}

		var sm structModel
		sm.Descriptor = d

		switch d.Kind {
		case model.DescriptorKind_Struct:
			m.Classes = append(m.Classes, &sm)
		case model.DescriptorKind_Enum:
			m.Enums = append(m.Enums, &sm)
		}

		// 遍历字段
		for _, fd := range d.Fields {

			// 对CombineStruct的XXDefine对应的字段
			if d.Usage == model.DescriptorUsage_CombineStruct {

				// 这个字段被限制输出
				if fd.Complex != nil && !fd.Complex.File.MatchTag(".cs") {
					continue
				}

				// 这个结构有索引才创建
				if fd.Complex != nil && len(fd.Complex.Indexes) > 0 {

					// 被索引的结构
					indexedField := csharpField{FieldDescriptor: fd, parentStruct: &sm}

					// 索引字段
					for _, key := range fd.Complex.Indexes {
						indexedField.IndexKeys = append(indexedField.IndexKeys, key)
					}

					sm.IndexedFields = append(sm.IndexedFields, indexedField)
				}

				if fd.Complex != nil && fd.Complex.File.Pragma.GetBool("Vertical") {
					m.VerticalFields = append(m.VerticalFields, csharpField{FieldDescriptor: fd, parentStruct: &sm})
				}

			}

			csField := csharpField{FieldDescriptor: fd, parentStruct: &sm}

			sm.Fields = append(sm.Fields, csField)

		}

	}

	bf := NewStream()

	err = tpl.Execute(bf.Buffer(), &m)
	if err != nil {
		log.Errorln(err)
		return nil
	}

	return bf
}

func init() {

	RegisterPrinter("cs", &csharpPrinter{})

}
<|endoftext|>


<commit_before>package libkbfs

import (
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

// CryptoClient implements the Crypto interface by sending RPCs to the
// keybase daemon to perform signatures using the device's current
// signing key.
type CryptoClient struct {
	CryptoCommon
	client     keybase1.CryptoClient
	shutdownFn func()
}

var _ Crypto = (*CryptoClient)(nil)

var _ rpc.ConnectionHandler = (*CryptoClient)(nil)

// NewCryptoClient constructs a new CryptoClient.
func NewCryptoClient(config Config, kbCtx *libkb.GlobalContext, log logger.Logger) *CryptoClient {
	c := &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: config.Codec(),
			log:   log,
		},
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, c)
	c.client = keybase1.CryptoClient{Cli: conn.GetClient()}
	c.shutdownFn = conn.Shutdown
	return c
}

// newCryptoClientWithClient should only be used for testing.
func newCryptoClientWithClient(codec Codec, client rpc.GenericClient,
	log logger.Logger) *CryptoClient {
	return &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: codec,
			log:   log,
		},
		client: keybase1.CryptoClient{Cli: client},
	}
}

// HandlerName implements the ConnectionHandler interface.
func (CryptoClient) HandlerName() string {
	return "CryptoClient"
}

// OnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnect(ctx context.Context, conn *rpc.Connection,
	_ rpc.GenericClient, server *rpc.Server) error {
	pushConnectionStatusChange(KeybaseServiceName, nil)
	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnectError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: connection error: %q; retrying in %s",
		err, wait)
	pushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (c *CryptoClient) OnDoCommandError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: docommand error: %q; retrying in %s",
		err, wait)
	pushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDisconnected implements the ConnectionHandler interface.
func (c *CryptoClient) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		c.log.Warning("CryptoClient is disconnected")
		pushConnectionStatusChange(KeybaseServiceName, errDisconnected{})
	}
}

// ShouldRetry implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

// Sign implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Sign(ctx context.Context, msg []byte) (
	sigInfo SignatureInfo, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message with %s: err=%v", len(msg),
			sigInfo, err)
	}()

	ed25519SigInfo, err := c.client.SignED25519(ctx, keybase1.SignED25519Arg{
		Msg:    msg,
		Reason: "to use kbfs",
	})
	if err != nil {
		return
	}

	sigInfo = SignatureInfo{
		Version:      SigED25519,
		Signature:    ed25519SigInfo.Sig[:],
		VerifyingKey: MakeVerifyingKey(libkb.NaclSigningKeyPublic(ed25519SigInfo.PublicKey).GetKID()),
	}
	return
}

// SignToString implements the Crypto interface for CryptoClient.
func (c *CryptoClient) SignToString(ctx context.Context, msg []byte) (
	signature string, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message: err=%v", len(msg), err)
	}()
	signature, err = c.client.SignToString(ctx, keybase1.SignToStringArg{
		Msg:    msg,
		Reason: "KBFS Authentication",
	})
	return
}

func (c *CryptoClient) prepareTLFCryptKeyClientHalf(encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	encryptedData keybase1.EncryptedBytes32, nonce keybase1.BoxNonce, err error) {
	if encryptedClientHalf.Version != EncryptionSecretbox {
		err = UnknownEncryptionVer{encryptedClientHalf.Version}
		return
	}

	if len(encryptedClientHalf.EncryptedData) != len(encryptedData) {
		err = libkb.DecryptionError{}
		return
	}
	copy(encryptedData[:], encryptedClientHalf.EncryptedData)

	if len(encryptedClientHalf.Nonce) != len(nonce) {
		err = InvalidNonceError{encryptedClientHalf.Nonce}
		return
	}
	copy(nonce[:], encryptedClientHalf.Nonce)
	return encryptedData, nonce, err
}

// DecryptTLFCryptKeyClientHalf implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalf(ctx context.Context,
	publicKey TLFEphemeralPublicKey,
	encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	clientHalf TLFCryptKeyClientHalf, err error) {
	encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(encryptedClientHalf)
	if err != nil {
		return
	}

	decryptedClientHalf, err := c.client.UnboxBytes32(ctx, keybase1.UnboxBytes32Arg{
		EncryptedBytes32: encryptedData,
		Nonce:            nonce,
		PeersPublicKey:   keybase1.BoxPublicKey(publicKey.data),
		Reason:           "to use kbfs",
	})
	if err != nil {
		return
	}

	clientHalf = MakeTLFCryptKeyClientHalf(decryptedClientHalf)
	return
}

// DecryptTLFCryptKeyClientHalfAny implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalfAny(ctx context.Context,
	keys []EncryptedTLFCryptKeyClientAndEphemeral, promptPaper bool) (
	clientHalf TLFCryptKeyClientHalf, index int, err error) {
	if len(keys) == 0 {
		return clientHalf, index, NoKeysError{}
	}
	bundles := make([]keybase1.CiphertextBundle, 0, len(keys))
	errors := make([]error, 0, len(keys))
	indexLookup := make([]int, 0, len(keys))
	for i, k := range keys {
		encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(k.ClientHalf)
		if err != nil {
			errors = append(errors, err)
		} else {
			bundles = append(bundles, keybase1.CiphertextBundle{
				Kid:        k.PubKey.kidContainer.kid,
				Ciphertext: encryptedData,
				Nonce:      nonce,
				PublicKey:  keybase1.BoxPublicKey(k.EPubKey.data),
			})
			indexLookup = append(indexLookup, i)
		}
	}
	if len(bundles) == 0 {
		err = errors[0]
		return
	}
	res, err := c.client.UnboxBytes32Any(ctx, keybase1.UnboxBytes32AnyArg{
		Bundles:     bundles,
		Reason:      "to rekey for kbfs",
		PromptPaper: promptPaper,
	})
	if err != nil {
		return
	}
	return MakeTLFCryptKeyClientHalf(res.Plaintext), indexLookup[res.Index], nil
}

// Shutdown implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Shutdown() {
	if c.shutdownFn != nil {
		c.shutdownFn()
	}
}
<commit_msg>crypto_client: set session ID for RPCs (#953)<commit_after>package libkbfs

import (
	"sync"
	"time"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/logger"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

// CryptoClient implements the Crypto interface by sending RPCs to the
// keybase daemon to perform signatures using the device's current
// signing key.
type CryptoClient struct {
	CryptoCommon
	client     keybase1.CryptoClient
	shutdownFn func()

	sessionLock sync.Mutex
	sessionID   int
}

var _ Crypto = (*CryptoClient)(nil)

var _ rpc.ConnectionHandler = (*CryptoClient)(nil)

// NewCryptoClient constructs a new CryptoClient.
func NewCryptoClient(config Config, kbCtx *libkb.GlobalContext, log logger.Logger) *CryptoClient {
	c := &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: config.Codec(),
			log:   log,
		},
	}
	conn := NewSharedKeybaseConnection(kbCtx, config, c)
	c.client = keybase1.CryptoClient{Cli: conn.GetClient()}
	c.shutdownFn = conn.Shutdown
	return c
}

// newCryptoClientWithClient should only be used for testing.
func newCryptoClientWithClient(codec Codec, client rpc.GenericClient,
	log logger.Logger) *CryptoClient {
	return &CryptoClient{
		CryptoCommon: CryptoCommon{
			codec: codec,
			log:   log,
		},
		client: keybase1.CryptoClient{Cli: client},
	}
}

// HandlerName implements the ConnectionHandler interface.
func (c *CryptoClient) HandlerName() string {
	return "CryptoClient"
}

func (c *CryptoClient) newSessionID() int {
	c.sessionLock.Lock()
	defer c.sessionLock.Unlock()
	newID := c.sessionID
	c.sessionID++ // overflow is ok
	return newID
}

// OnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnect(ctx context.Context, conn *rpc.Connection,
	_ rpc.GenericClient, server *rpc.Server) error {
	pushConnectionStatusChange(KeybaseServiceName, nil)
	return nil
}

// OnConnectError implements the ConnectionHandler interface.
func (c *CryptoClient) OnConnectError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: connection error: %q; retrying in %s",
		err, wait)
	pushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDoCommandError implements the ConnectionHandler interface.
func (c *CryptoClient) OnDoCommandError(err error, wait time.Duration) {
	c.log.Warning("CryptoClient: docommand error: %q; retrying in %s",
		err, wait)
	pushConnectionStatusChange(KeybaseServiceName, err)
}

// OnDisconnected implements the ConnectionHandler interface.
func (c *CryptoClient) OnDisconnected(_ context.Context,
	status rpc.DisconnectStatus) {
	if status == rpc.StartingNonFirstConnection {
		c.log.Warning("CryptoClient is disconnected")
		pushConnectionStatusChange(KeybaseServiceName, errDisconnected{})
	}
}

// ShouldRetry implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetry(rpcName string, err error) bool {
	return false
}

// ShouldRetryOnConnect implements the ConnectionHandler interface.
func (c *CryptoClient) ShouldRetryOnConnect(err error) bool {
	_, inputCanceled := err.(libkb.InputCanceledError)
	return !inputCanceled
}

// Sign implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Sign(ctx context.Context, msg []byte) (
	sigInfo SignatureInfo, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message with %s: err=%v", len(msg),
			sigInfo, err)
	}()

	ed25519SigInfo, err := c.client.SignED25519(ctx, keybase1.SignED25519Arg{
		SessionID: c.newSessionID(),
		Msg:       msg,
		Reason:    "to use kbfs",
	})
	if err != nil {
		return
	}

	sigInfo = SignatureInfo{
		Version:      SigED25519,
		Signature:    ed25519SigInfo.Sig[:],
		VerifyingKey: MakeVerifyingKey(libkb.NaclSigningKeyPublic(ed25519SigInfo.PublicKey).GetKID()),
	}
	return
}

// SignToString implements the Crypto interface for CryptoClient.
func (c *CryptoClient) SignToString(ctx context.Context, msg []byte) (
	signature string, err error) {
	defer func() {
		c.log.CDebugf(ctx, "Signed %d-byte message: err=%v", len(msg), err)
	}()
	signature, err = c.client.SignToString(ctx, keybase1.SignToStringArg{
		SessionID: c.newSessionID(),
		Msg:       msg,
		Reason:    "KBFS Authentication",
	})
	return
}

func (c *CryptoClient) prepareTLFCryptKeyClientHalf(encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	encryptedData keybase1.EncryptedBytes32, nonce keybase1.BoxNonce, err error) {
	if encryptedClientHalf.Version != EncryptionSecretbox {
		err = UnknownEncryptionVer{encryptedClientHalf.Version}
		return
	}

	if len(encryptedClientHalf.EncryptedData) != len(encryptedData) {
		err = libkb.DecryptionError{}
		return
	}
	copy(encryptedData[:], encryptedClientHalf.EncryptedData)

	if len(encryptedClientHalf.Nonce) != len(nonce) {
		err = InvalidNonceError{encryptedClientHalf.Nonce}
		return
	}
	copy(nonce[:], encryptedClientHalf.Nonce)
	return encryptedData, nonce, err
}

// DecryptTLFCryptKeyClientHalf implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalf(ctx context.Context,
	publicKey TLFEphemeralPublicKey,
	encryptedClientHalf EncryptedTLFCryptKeyClientHalf) (
	clientHalf TLFCryptKeyClientHalf, err error) {
	encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(encryptedClientHalf)
	if err != nil {
		return
	}

	decryptedClientHalf, err := c.client.UnboxBytes32(ctx, keybase1.UnboxBytes32Arg{
		SessionID:        c.newSessionID(),
		EncryptedBytes32: encryptedData,
		Nonce:            nonce,
		PeersPublicKey:   keybase1.BoxPublicKey(publicKey.data),
		Reason:           "to use kbfs",
	})
	if err != nil {
		return
	}

	clientHalf = MakeTLFCryptKeyClientHalf(decryptedClientHalf)
	return
}

// DecryptTLFCryptKeyClientHalfAny implements the Crypto interface for
// CryptoClient.
func (c *CryptoClient) DecryptTLFCryptKeyClientHalfAny(ctx context.Context,
	keys []EncryptedTLFCryptKeyClientAndEphemeral, promptPaper bool) (
	clientHalf TLFCryptKeyClientHalf, index int, err error) {
	if len(keys) == 0 {
		return clientHalf, index, NoKeysError{}
	}
	bundles := make([]keybase1.CiphertextBundle, 0, len(keys))
	errors := make([]error, 0, len(keys))
	indexLookup := make([]int, 0, len(keys))
	for i, k := range keys {
		encryptedData, nonce, err := c.prepareTLFCryptKeyClientHalf(k.ClientHalf)
		if err != nil {
			errors = append(errors, err)
		} else {
			bundles = append(bundles, keybase1.CiphertextBundle{
				Kid:        k.PubKey.kidContainer.kid,
				Ciphertext: encryptedData,
				Nonce:      nonce,
				PublicKey:  keybase1.BoxPublicKey(k.EPubKey.data),
			})
			indexLookup = append(indexLookup, i)
		}
	}
	if len(bundles) == 0 {
		err = errors[0]
		return
	}
	res, err := c.client.UnboxBytes32Any(ctx, keybase1.UnboxBytes32AnyArg{
		SessionID:   c.newSessionID(),
		Bundles:     bundles,
		Reason:      "to rekey for kbfs",
		PromptPaper: promptPaper,
	})
	if err != nil {
		return
	}
	return MakeTLFCryptKeyClientHalf(res.Plaintext), indexLookup[res.Index], nil
}

// Shutdown implements the Crypto interface for CryptoClient.
func (c *CryptoClient) Shutdown() {
	if c.shutdownFn != nil {
		c.shutdownFn()
	}
}
<|endoftext|>


<commit_before>// Package request implements request entity contains url and other relevant informaion.
package request

import (
	"io/ioutil"
	"net/http"
	"os"
	"github.com/bitly/go-simplejson"

	"github.com/hu17889/go_spider/core/common/mlog"
)

// Request represents object waiting for being crawled.
type Request struct {
    url string

    // Responce type: html json jsonp text
    respType string

    // GET POST
    method string

    // POST data
    postdata string

    // name for marking url and distinguish different urls in PageProcesser and Pipeline
    urltag string

    // http header
    header http.Header

    // http cookies
    cookies []*http.Cookie

    // Redirect function for downloader used in http.Client
    // If CheckRedirect returns an error, the Client's Get
    // method returns both the previous Response.
    // If CheckRedirect returns error.New("normal"), the error process after client.Do will ignore the error.
    checkRedirect func(req *http.Request, via []*http.Request) error

    meta interface{}
}

// NewRequest returns initialized Request object.
// The respType is json, jsonp, html, text
/*
func NewRequestSimple(url string, respType string, urltag string) *Request {
    return &Request{url:url, respType:respType}
}
*/

func NewRequest(url string, respType string, urltag string, method string,
    postdata string, header http.Header, cookies []*http.Cookie,
    checkRedirect func(req *http.Request, via []*http.Request) error,
    meta interface{}) *Request {
    return &Request{url, respType, method, postdata, urltag, header, cookies, checkRedirect, meta}
}

func NewRequestWithHeaderFile(url string, respType string, headerFile string) *Request {
	_, err := os.Stat(headerFile)
	if err != nil {
		//file is not exist , using default mode
		return NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
	}
	//read file , parse the header and cookies
	b, err := ioutil.ReadFile(headerFile)
	if err != nil {
		//make be:  share access error
		mlog.LogInst().LogError(err.Error())
	}
	js, _ := simplejson.NewJson(b)
	//constructed to header

	h := make(http.Header)
	h.Add("User-Agent", js.Get("User-Agent").MustString())
	h.Add("Cookie", js.Get("Cookie").MustString())
	h.Add("Cache-Control", "max-age=0")
	h.Add("Connection", "keep-alive")
	return NewRequest(url, respType, "", "GET", "", h, nil, nil, nil)
}

func (this *Request) GetUrl() string {
    return this.url
}

func (this *Request) GetUrlTag() string {
    return this.urltag
}

func (this *Request) GetMethod() string {
    return this.method
}

func (this *Request) GetPostdata() string {
    return this.postdata
}

func (this *Request) GetHeader() http.Header {
    return this.header
}

func (this *Request) GetCookies() []*http.Cookie {
    return this.cookies
}

func (this *Request) GetResponceType() string {
    return this.respType
}

func (this *Request) GetRedirectFunc() func(req *http.Request, via []*http.Request) error {
    return this.checkRedirect
}

func (this *Request) GetMeta() interface{} {
    return this.meta
}
<commit_msg>add method to handling the proxy<commit_after>// Package request implements request entity contains url and other relevant informaion.
package request

import (
	"io/ioutil"
	"net/http"
	"os"
	"github.com/bitly/go-simplejson"

	"github.com/hu17889/go_spider/core/common/mlog"
)

// Request represents object waiting for being crawled.
type Request struct {
	url string

	// Responce type: html json jsonp text
	respType string

	// GET POST
	method string

	// POST data
	postdata string

	// name for marking url and distinguish different urls in PageProcesser and Pipeline
	urltag string

	// http header
	header http.Header

	// http cookies
	cookies []*http.Cookie

	//proxy host   example='localhost:80'
	proxyHost string

	// Redirect function for downloader used in http.Client
	// If CheckRedirect returns an error, the Client's Get
	// method returns both the previous Response.
	// If CheckRedirect returns error.New("normal"), the error process after client.Do will ignore the error.
	checkRedirect func(req *http.Request, via []*http.Request) error

	meta interface{}
}

// NewRequest returns initialized Request object.
// The respType is json, jsonp, html, text
/*
func NewRequestSimple(url string, respType string, urltag string) *Request {
    return &Request{url:url, respType:respType}
}
*/

func NewRequest(url string, respType string, urltag string, method string,
	postdata string, header http.Header, cookies []*http.Cookie,
	checkRedirect func(req *http.Request, via []*http.Request) error,
	meta interface{}) *Request {
	return &Request{url, respType, method, postdata, urltag, header, cookies, "",checkRedirect, meta}
}

func NewRequestWithProxy(url string, respType string, urltag string, method string,
	postdata string, header http.Header, cookies []*http.Cookie,proxyHost string,
	checkRedirect func(req *http.Request, via []*http.Request ) error,
	meta interface{}) *Request {
	return &Request{url, respType, method, postdata, urltag, header, cookies, proxyHost ,checkRedirect, meta}
}

func NewRequestWithHeaderFile(url string, respType string, headerFile string) *Request {
	_, err := os.Stat(headerFile)
	if err != nil {
		//file is not exist , using default mode
		return NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
	}

	h := readHeaderFromFile(headerFile)

	return NewRequest(url, respType, "", "GET", "", h, nil, nil, nil)
}

func readHeaderFromFile(headerFile string) http.Header {
	//read file , parse the header and cookies
	b, err := ioutil.ReadFile(headerFile)
	if err != nil {
		//make be:  share access error
		mlog.LogInst().LogError(err.Error())
		return nil
	}
	js, _ := simplejson.NewJson(b)
	//constructed to header

	h := make(http.Header)
	h.Add("User-Agent", js.Get("User-Agent").MustString())
	h.Add("Referer", js.Get("Referer").MustString())
	h.Add("Cookie", js.Get("Cookie").MustString())
	h.Add("Cache-Control", "max-age=0")
	h.Add("Connection", "keep-alive")
	return h
}

//point to a json file
/* xxx.json
{
	"User-Agent":"curl/7.19.3 (i386-pc-win32) libcurl/7.19.3 OpenSSL/1.0.0d",
	"Referer":"http://weixin.sogou.com/gzh?openid=oIWsFt6Sb7aZmuI98AU7IXlbjJps",
	"Cookie":""
}
*/
func (this *Request) AddHeaderFile(headerFile string) *Request {
	_, err := os.Stat(headerFile)
	if err != nil {
		return this
	}
	h := readHeaderFromFile(headerFile)
	this.header = h
	return this
}

// @host  http://localhost:8765/
func (this *Request) AddProxyHost(host string) *Request {
	this.proxyHost = host
	return this
}

func (this *Request) GetUrl() string {
	return this.url
}

func (this *Request) GetUrlTag() string {
	return this.urltag
}

func (this *Request) GetMethod() string {
	return this.method
}

func (this *Request) GetPostdata() string {
	return this.postdata
}

func (this *Request) GetHeader() http.Header {
	return this.header
}

func (this *Request) GetCookies() []*http.Cookie {
	return this.cookies
}

func (this *Request) GetProxyHost() string {
	return this.proxyHost
}

func (this *Request) GetResponceType() string {
	return this.respType
}

func (this *Request) GetRedirectFunc() func(req *http.Request, via []*http.Request) error {
	return this.checkRedirect
}

func (this *Request) GetMeta() interface{} {
	return this.meta
}
<|endoftext|>


<commit_before><commit_msg>reflection<commit_after><|endoftext|>