<commit_before>package db

import (
	"database/sql"

	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc/db/lock"
)

//go:generate counterfeiter . ResourceFactory

type ResourceFactory interface {
	VisibleResources([]string) ([]Resource, error)
	AllResources() ([]Resource, error)
}

type resourceFactory struct {
	conn        Conn
	lockFactory lock.LockFactory
}

func NewResourceFactory(conn Conn, lockFactory lock.LockFactory) ResourceFactory {
	return &resourceFactory{
		conn:        conn,
		lockFactory: lockFactory,
	}
}

func (r *resourceFactory) VisibleResources(teamNames []string) ([]Resource, error) {
	rows, err := resourcesQuery.
		Where(
			sq.Or{
				sq.Eq{"t.name": teamNames},
				sq.And{
					sq.NotEq{"t.name": teamNames},
					sq.Eq{"p.public": true},
				},
			}).
		OrderBy("r.id ASC").
		RunWith(r.conn).
		Query()
	if err != nil {
		return nil, err
	}

	return r.parseResources(rows)
}

func (r *resourceFactory) AllResources() ([]Resource, error) {
	rows, err := resourcesQuery.
		OrderBy("r.id ASC").
		RunWith(r.conn).
		Query()
	if err != nil {
		return nil, err
	}

	return r.parseResources(rows)
}

func (r *resourceFactory) parseResources(resourceRows *sql.Rows) ([]Resource, error) {
	var resources []Resource

	for resourceRows.Next() {
		resource := &resource{conn: r.conn, lockFactory: r.lockFactory}

		err := scanResource(resource, resourceRows)
		if err != nil {
			return nil, err
		}

		resources = append(resources, resource)
	}

	return resources, nil
}
<commit_msg>Change the parseResources to scanResources<commit_after>package db

import (
	"database/sql"
	sq "github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc/db/lock"
)

//go:generate counterfeiter . ResourceFactory

type ResourceFactory interface {
	VisibleResources([]string) ([]Resource, error)
	AllResources() ([]Resource, error)
}

type resourceFactory struct {
	conn        Conn
	lockFactory lock.LockFactory
}

func NewResourceFactory(conn Conn, lockFactory lock.LockFactory) ResourceFactory {
	return &resourceFactory{
		conn:        conn,
		lockFactory: lockFactory,
	}
}

func (r *resourceFactory) VisibleResources(teamNames []string) ([]Resource, error) {
	rows, err := resourcesQuery.
		Where(
			sq.Or{
				sq.Eq{"t.name": teamNames},
				sq.And{
					sq.NotEq{"t.name": teamNames},
					sq.Eq{"p.public": true},
				},
			}).
		OrderBy("r.id ASC").
		RunWith(r.conn).
		Query()
	if err != nil {
		return nil, err
	}

	return scanResources(rows, r.conn, r.lockFactory)
}

func (r *resourceFactory) AllResources() ([]Resource, error) {
	rows, err := resourcesQuery.
		OrderBy("r.id ASC").
		RunWith(r.conn).
		Query()
	if err != nil {
		return nil, err
	}

	return scanResources(rows, r.conn, r.lockFactory)
}

func scanResources(resourceRows *sql.Rows, conn Conn, lockFactory lock.LockFactory) ([]Resource, error) {
	var resources []Resource

	for resourceRows.Next() {
		resource := &resource{conn: conn, lockFactory: lockFactory}

		err := scanResource(resource, resourceRows)
		if err != nil {
			return nil, err
		}

		resources = append(resources, resource)
	}

	return resources, nil
}
<|endoftext|>


<commit_before>package roller

import (
	"context"
	"fmt"

	"go.skia.org/infra/autoroll/go/recent_rolls"
	"go.skia.org/infra/autoroll/go/repo_manager"
	"go.skia.org/infra/autoroll/go/state_machine"
	"go.skia.org/infra/go/autoroll"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/sklog"
)

type RollImpl interface {
	state_machine.RollCLImpl

	// Insert the roll into the DB.
	InsertIntoDB() error
}

func retrieveGerritIssue(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, rollIntoAndroid bool, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
	info, err := g.GetIssueProperties(issueNum)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to get issue properties: %s", err)
	}
	a, err := autoroll.FromGerritChangeInfo(info, func(h string) (string, error) {
		return rm.FullChildHash(ctx, h)
	}, rollIntoAndroid)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to convert issue format: %s", err)
	}
	tryResults, err := autoroll.GetTryResultsFromGerrit(g, a)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to retrieve try results: %s", err)
	}
	a.TryResults = tryResults
	return info, a, nil
}

// gerritRoll is an implementation of RollImpl.
type gerritRoll struct {
	ci           *gerrit.ChangeInfo
	issue        *autoroll.AutoRollIssue
	g            *gerrit.Gerrit
	recent       *recent_rolls.RecentRolls
	retrieveRoll func(context.Context, int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error)
	result       string
	rm           repo_manager.RepoManager
}

// newGerritRoll obtains a gerritRoll instance from the given Gerrit issue
// number.
func newGerritRoll(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, recent *recent_rolls.RecentRolls, issueNum int64) (RollImpl, error) {
	ci, issue, err := retrieveGerritIssue(ctx, g, rm, false, issueNum)
	if err != nil {
		return nil, err
	}
	return &gerritRoll{
		ci:     ci,
		issue:  issue,
		g:      g,
		recent: recent,
		retrieveRoll: func(ctx context.Context, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
			return retrieveGerritIssue(ctx, g, rm, false, issueNum)
		},
		rm: rm,
	}, nil
}

// See documentation for RollImpl interface.
func (r *gerritRoll) InsertIntoDB() error {
	return r.recent.Add(r.issue)
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) AddComment(msg string) error {
	return r.g.AddComment(r.ci, msg)
}

// Helper function for modifying a roll CL which might fail due to the CL being
// closed by a human or some other process, in which case we don't want to error
// out.
func (r *gerritRoll) withModify(ctx context.Context, action string, fn func() error) error {
	if err := fn(); err != nil {
		// It's possible that somebody abandoned the CL (or the CL
		// landed) while we were working. If that's the case, log an
		// error and move on.
		if err2 := r.Update(ctx); err2 != nil {
			return fmt.Errorf("Failed to %s with error:\n%s\nAnd failed to update it with error:\n%s", action, err, err2)
		}
		if r.ci.IsClosed() {
			sklog.Errorf("Attempted to %s but it is already closed! Error: %s", action, err)
			return nil
		}
		return err
	}
	return r.Update(ctx)
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) Close(ctx context.Context, result, msg string) error {
	sklog.Infof("Closing issue %d (result %q) with message: %s", r.ci.Issue, result, msg)
	r.result = result
	return r.withModify(ctx, "close the CL", func() error {
		return r.g.Abandon(r.ci, msg)
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsFinished() bool {
	return r.ci.IsClosed() || !r.issue.CommitQueue
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsSuccess() bool {
	return r.ci.Status == gerrit.CHANGE_STATUS_MERGED
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsDryRunFinished() bool {
	// The CQ removes the CQ+1 label when the dry run finishes, regardless
	// of success or failure. Since we uploaded with the dry run label set,
	// we know the roll is in progress if the label is still set, and done
	// otherwise.
	return !r.issue.CommitQueueDryRun
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsDryRunSuccess() bool {
	return r.issue.AllTrybotsSucceeded()
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RollingTo() string {
	return r.issue.RollingTo
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) SwitchToDryRun(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL to dry run", func() error {
		return r.g.SendToDryRun(r.ci, "Mode was changed to dry run")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) SwitchToNormal(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL out of dry run", func() error {
		return r.g.SendToCQ(r.ci, "Mode was changed to normal")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RetryCQ(ctx context.Context) error {
	return r.withModify(ctx, "retry the CQ", func() error {
		return r.g.SendToCQ(r.ci, "CQ failed but there are no new commits. Retrying...")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RetryDryRun(ctx context.Context) error {
	return r.withModify(ctx, "retry the CQ (dry run)", func() error {
		return r.g.SendToDryRun(r.ci, "Dry run failed but there are no new commits. Retrying...")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) Update(ctx context.Context) error {
	ci, issue, err := r.retrieveRoll(ctx, r.issue.Issue)
	if err != nil {
		return err
	}
	r.ci = ci
	if r.result != "" {
		issue.Result = r.result
	}
	r.issue = issue
	return r.recent.Update(r.issue)
}

// Special type for Android rolls.
type gerritAndroidRoll struct {
	*gerritRoll
}

// newGerritAndroidRoll obtains a gerritAndroidRoll instance from the given Gerrit issue number.
func newGerritAndroidRoll(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, recent *recent_rolls.RecentRolls, issueNum int64) (RollImpl, error) {
	ci, issue, err := retrieveGerritIssue(ctx, g, rm, true, issueNum)
	if err != nil {
		return nil, err
	}
	return &gerritAndroidRoll{&gerritRoll{
		ci:     ci,
		issue:  issue,
		g:      g,
		recent: recent,
		retrieveRoll: func(ctx context.Context, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
			return retrieveGerritIssue(ctx, g, rm, true, issueNum)
		},
		rm: rm,
	}}, nil
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) IsDryRunFinished() bool {
	if _, ok := r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL]; ok {
		for _, lb := range r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL].All {
			if lb.Value != gerrit.PRESUBMIT_VERIFIED_LABEL_RUNNING {
				return true
			}
		}
	}
	return false
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) IsDryRunSuccess() bool {
	presubmit, ok := r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL]
	if !ok || len(presubmit.All) == 0 {
		// Not done yet.
		return false
	}
	for _, lb := range presubmit.All {
		if lb.Value == gerrit.PRESUBMIT_VERIFIED_LABEL_ACCEPTED {
			return true
		}
	}
	return false
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) SwitchToDryRun(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL to dry run", func() error {
		return r.g.SetReview(r.ci, "Mode was changed to dry run", map[string]interface{}{gerrit.AUTOSUBMIT_LABEL: gerrit.AUTOSUBMIT_LABEL_NONE})
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) SwitchToNormal(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL out of dry run", func() error {
		return r.g.SetReview(r.ci, "Mode was changed to normal", map[string]interface{}{gerrit.AUTOSUBMIT_LABEL: gerrit.AUTOSUBMIT_LABEL_SUBMIT})
	})
}
<commit_msg>Add RetryCQ and RetryDryRun implementations for GerritAndroidRoll<commit_after>package roller

import (
	"context"
	"fmt"

	"go.skia.org/infra/autoroll/go/recent_rolls"
	"go.skia.org/infra/autoroll/go/repo_manager"
	"go.skia.org/infra/autoroll/go/state_machine"
	"go.skia.org/infra/go/autoroll"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/sklog"
)

type RollImpl interface {
	state_machine.RollCLImpl

	// Insert the roll into the DB.
	InsertIntoDB() error
}

func retrieveGerritIssue(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, rollIntoAndroid bool, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
	info, err := g.GetIssueProperties(issueNum)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to get issue properties: %s", err)
	}
	a, err := autoroll.FromGerritChangeInfo(info, func(h string) (string, error) {
		return rm.FullChildHash(ctx, h)
	}, rollIntoAndroid)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to convert issue format: %s", err)
	}
	tryResults, err := autoroll.GetTryResultsFromGerrit(g, a)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to retrieve try results: %s", err)
	}
	a.TryResults = tryResults
	return info, a, nil
}

// gerritRoll is an implementation of RollImpl.
type gerritRoll struct {
	ci           *gerrit.ChangeInfo
	issue        *autoroll.AutoRollIssue
	g            *gerrit.Gerrit
	recent       *recent_rolls.RecentRolls
	retrieveRoll func(context.Context, int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error)
	result       string
	rm           repo_manager.RepoManager
}

// newGerritRoll obtains a gerritRoll instance from the given Gerrit issue
// number.
func newGerritRoll(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, recent *recent_rolls.RecentRolls, issueNum int64) (RollImpl, error) {
	ci, issue, err := retrieveGerritIssue(ctx, g, rm, false, issueNum)
	if err != nil {
		return nil, err
	}
	return &gerritRoll{
		ci:     ci,
		issue:  issue,
		g:      g,
		recent: recent,
		retrieveRoll: func(ctx context.Context, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
			return retrieveGerritIssue(ctx, g, rm, false, issueNum)
		},
		rm: rm,
	}, nil
}

// See documentation for RollImpl interface.
func (r *gerritRoll) InsertIntoDB() error {
	return r.recent.Add(r.issue)
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) AddComment(msg string) error {
	return r.g.AddComment(r.ci, msg)
}

// Helper function for modifying a roll CL which might fail due to the CL being
// closed by a human or some other process, in which case we don't want to error
// out.
func (r *gerritRoll) withModify(ctx context.Context, action string, fn func() error) error {
	if err := fn(); err != nil {
		// It's possible that somebody abandoned the CL (or the CL
		// landed) while we were working. If that's the case, log an
		// error and move on.
		if err2 := r.Update(ctx); err2 != nil {
			return fmt.Errorf("Failed to %s with error:\n%s\nAnd failed to update it with error:\n%s", action, err, err2)
		}
		if r.ci.IsClosed() {
			sklog.Errorf("Attempted to %s but it is already closed! Error: %s", action, err)
			return nil
		}
		return err
	}
	return r.Update(ctx)
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) Close(ctx context.Context, result, msg string) error {
	sklog.Infof("Closing issue %d (result %q) with message: %s", r.ci.Issue, result, msg)
	r.result = result
	return r.withModify(ctx, "close the CL", func() error {
		return r.g.Abandon(r.ci, msg)
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsFinished() bool {
	return r.ci.IsClosed() || !r.issue.CommitQueue
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsSuccess() bool {
	return r.ci.Status == gerrit.CHANGE_STATUS_MERGED
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsDryRunFinished() bool {
	// The CQ removes the CQ+1 label when the dry run finishes, regardless
	// of success or failure. Since we uploaded with the dry run label set,
	// we know the roll is in progress if the label is still set, and done
	// otherwise.
	return !r.issue.CommitQueueDryRun
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) IsDryRunSuccess() bool {
	return r.issue.AllTrybotsSucceeded()
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RollingTo() string {
	return r.issue.RollingTo
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) SwitchToDryRun(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL to dry run", func() error {
		return r.g.SendToDryRun(r.ci, "Mode was changed to dry run")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) SwitchToNormal(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL out of dry run", func() error {
		return r.g.SendToCQ(r.ci, "Mode was changed to normal")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RetryCQ(ctx context.Context) error {
	return r.withModify(ctx, "retry the CQ", func() error {
		return r.g.SendToCQ(r.ci, "CQ failed but there are no new commits. Retrying...")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) RetryDryRun(ctx context.Context) error {
	return r.withModify(ctx, "retry the CQ (dry run)", func() error {
		return r.g.SendToDryRun(r.ci, "Dry run failed but there are no new commits. Retrying...")
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritRoll) Update(ctx context.Context) error {
	ci, issue, err := r.retrieveRoll(ctx, r.issue.Issue)
	if err != nil {
		return err
	}
	r.ci = ci
	if r.result != "" {
		issue.Result = r.result
	}
	r.issue = issue
	return r.recent.Update(r.issue)
}

// Special type for Android rolls.
type gerritAndroidRoll struct {
	*gerritRoll
}

// newGerritAndroidRoll obtains a gerritAndroidRoll instance from the given Gerrit issue number.
func newGerritAndroidRoll(ctx context.Context, g *gerrit.Gerrit, rm repo_manager.RepoManager, recent *recent_rolls.RecentRolls, issueNum int64) (RollImpl, error) {
	ci, issue, err := retrieveGerritIssue(ctx, g, rm, true, issueNum)
	if err != nil {
		return nil, err
	}
	return &gerritAndroidRoll{&gerritRoll{
		ci:     ci,
		issue:  issue,
		g:      g,
		recent: recent,
		retrieveRoll: func(ctx context.Context, issueNum int64) (*gerrit.ChangeInfo, *autoroll.AutoRollIssue, error) {
			return retrieveGerritIssue(ctx, g, rm, true, issueNum)
		},
		rm: rm,
	}}, nil
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) IsDryRunFinished() bool {
	if _, ok := r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL]; ok {
		for _, lb := range r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL].All {
			if lb.Value != gerrit.PRESUBMIT_VERIFIED_LABEL_RUNNING {
				return true
			}
		}
	}
	return false
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) IsDryRunSuccess() bool {
	presubmit, ok := r.ci.Labels[gerrit.PRESUBMIT_VERIFIED_LABEL]
	if !ok || len(presubmit.All) == 0 {
		// Not done yet.
		return false
	}
	for _, lb := range presubmit.All {
		if lb.Value == gerrit.PRESUBMIT_VERIFIED_LABEL_ACCEPTED {
			return true
		}
	}
	return false
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) SwitchToDryRun(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL to dry run", func() error {
		return r.g.SetReview(r.ci, "Mode was changed to dry run", map[string]interface{}{gerrit.AUTOSUBMIT_LABEL: gerrit.AUTOSUBMIT_LABEL_NONE})
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) SwitchToNormal(ctx context.Context) error {
	return r.withModify(ctx, "switch the CL out of dry run", func() error {
		return r.g.SetReview(r.ci, "Mode was changed to normal", map[string]interface{}{gerrit.AUTOSUBMIT_LABEL: gerrit.AUTOSUBMIT_LABEL_SUBMIT})
	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) RetryCQ(ctx context.Context) error {
	return r.withModify(ctx, "retry TH", func() error {
		return r.g.SetReview(r.ci, "TH failed but there are no new commits. Retrying...", map[string]interface{}{gerrit.PRESUBMIT_READY_LABEL: "1"})

	})
}

// See documentation for state_machine.RollCLImpl interface.
func (r *gerritAndroidRoll) RetryDryRun(ctx context.Context) error {
	return r.withModify(ctx, "retry the TH (dry run)", func() error {
		return r.g.SetReview(r.ci, "Dry run failed but there are no new commits. Retrying...", map[string]interface{}{gerrit.PRESUBMIT_READY_LABEL: "1"})
	})
}
<|endoftext|>


<commit_before>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"context"
	"errors"
	"net"
	"sync/atomic"
	"time"
)

type Socket struct {
	pc            net.PacketConn
	ctx           *C.utp_context
	backlog       chan *Conn
	closed        bool
	conns         map[*C.utp_socket]*Conn
	nonUtpReads   chan packet
	writeDeadline time.Time
	readDeadline  time.Time
}

var (
	_ net.PacketConn = (*Socket)(nil)
)

type packet struct {
	b    []byte
	from net.Addr
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := net.ListenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	mu.Lock()
	defer mu.Unlock()
	ctx := C.utp_init(2)
	if ctx == nil {
		panic(ctx)
	}
	ctx.setCallbacks()
	if utpLogging {
		ctx.setOption(C.UTP_LOG_NORMAL, 1)
		ctx.setOption(C.UTP_LOG_MTU, 1)
		ctx.setOption(C.UTP_LOG_DEBUG, 1)
	}
	s := &Socket{
		pc:          pc,
		ctx:         ctx,
		backlog:     make(chan *Conn, 5),
		conns:       make(map[*C.utp_socket]*Conn),
		nonUtpReads: make(chan packet, 100),
	}
	libContextToSocket[ctx] = s
	go s.timeoutChecker()
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s: us,
	}
	c.cond.L = &mu
	s.conns[us] = c
	c.writeDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	c.readDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	return c
}

var reads int64

func (s *Socket) packetReader() {
	var b [0x1000]byte
	for {
		n, addr, err := s.pc.ReadFrom(b[:])
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				return
			}
			panic(err)
		}
		sa, sal := netAddrToLibSockaddr(addr)
		atomic.AddInt64(&reads, 1)
		// log.Printf("received %d bytes, %d packets", n, reads)
		func() {
			mu.Lock()
			defer mu.Unlock()
			if s.closed {
				return
			}
			if C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(n), sa, sal) != 0 {
				socketUtpPacketsReceived.Add(1)
			} else {
				s.onReadNonUtp(b[:n], addr)
			}
			C.utp_issue_deferred_acks(s.ctx)
			C.utp_check_timeouts(s.ctx)
		}()
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		// C.utp_issue_deferred_acks(s.ctx)
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (me *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	if me.closed {
		return nil
	}
	C.utp_destroy(me.ctx)
	me.ctx = nil
	me.pc.Close()
	close(me.backlog)
	close(me.nonUtpReads)
	me.closed = true
	return nil
}

func (me *Socket) Addr() net.Addr {
	return me.pc.LocalAddr()
}

func (me *Socket) LocalAddr() net.Addr {
	return me.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ctx := context.Background()
	if timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	return s.DialContext(ctx, addr)
}

func (s *Socket) DialContext(ctx context.Context, addr string) (net.Conn, error) {
	ua, err := net.ResolveUDPAddr(s.Addr().Network(), addr)
	if err != nil {
		panic(err)
	}
	sa, sl := netAddrToLibSockaddr(ua)
	mu.Lock()
	c := s.newConn(C.utp_create_socket(s.ctx))
	C.utp_connect(c.s, sa, sl)
	defer mu.Unlock()
	err = c.waitForConnect(ctx)
	if err != nil {
		c.close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p, ok := <-s.nonUtpReads
	if !ok {
		err = errors.New("closed")
		return
	}
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	socketNonUtpPacketsReceived.Add(1)
	select {
	case s.nonUtpReads <- packet{append([]byte(nil), b...), from}:
	default:
		// log.Printf("dropped non utp packet: no room in buffer")
		nonUtpPacketsDropped.Add(1)
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}

func (s *Socket) ReadBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_RCVBUF))
}

func (s *Socket) WriteBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_SNDBUF))
}

func (s *Socket) SetWriteBufferLen(len int) {
	mu.Lock()
	defer mu.Unlock()
	i := C.utp_context_set_option(s.ctx, C.UTP_SNDBUF, C.int(len))
	if i != 0 {
		panic(i)
	}
}
<commit_msg>Add comment about utp_issue_deferred_acks when processing UDP packets<commit_after>package utp

/*
#include "utp.h"
*/
import "C"
import (
	"context"
	"errors"
	"net"
	"sync/atomic"
	"time"
)

type Socket struct {
	pc            net.PacketConn
	ctx           *C.utp_context
	backlog       chan *Conn
	closed        bool
	conns         map[*C.utp_socket]*Conn
	nonUtpReads   chan packet
	writeDeadline time.Time
	readDeadline  time.Time
}

var (
	_ net.PacketConn = (*Socket)(nil)
)

type packet struct {
	b    []byte
	from net.Addr
}

func NewSocket(network, addr string) (*Socket, error) {
	pc, err := net.ListenPacket(network, addr)
	if err != nil {
		return nil, err
	}
	mu.Lock()
	defer mu.Unlock()
	ctx := C.utp_init(2)
	if ctx == nil {
		panic(ctx)
	}
	ctx.setCallbacks()
	if utpLogging {
		ctx.setOption(C.UTP_LOG_NORMAL, 1)
		ctx.setOption(C.UTP_LOG_MTU, 1)
		ctx.setOption(C.UTP_LOG_DEBUG, 1)
	}
	s := &Socket{
		pc:          pc,
		ctx:         ctx,
		backlog:     make(chan *Conn, 5),
		conns:       make(map[*C.utp_socket]*Conn),
		nonUtpReads: make(chan packet, 100),
	}
	libContextToSocket[ctx] = s
	go s.timeoutChecker()
	go s.packetReader()
	return s, nil
}

func (s *Socket) onLibSocketDestroyed(ls *C.utp_socket) {
	delete(s.conns, ls)
}

func (s *Socket) newConn(us *C.utp_socket) *Conn {
	c := &Conn{
		s: us,
	}
	c.cond.L = &mu
	s.conns[us] = c
	c.writeDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	c.readDeadlineTimer = time.AfterFunc(-1, c.cond.Broadcast)
	return c
}

var reads int64

func (s *Socket) packetReader() {
	var b [0x1000]byte
	for {
		// In C, all the reads are processed and when it threatens to block,
		// only then do we call utp_issue_deferred_acks. I don't know how we
		// can do this in Go.
		n, addr, err := s.pc.ReadFrom(b[:])
		if err != nil {
			mu.Lock()
			closed := s.closed
			mu.Unlock()
			if closed {
				return
			}
			panic(err)
		}
		sa, sal := netAddrToLibSockaddr(addr)
		atomic.AddInt64(&reads, 1)
		// log.Printf("received %d bytes, %d packets", n, reads)
		func() {
			mu.Lock()
			defer mu.Unlock()
			if s.closed {
				return
			}
			if C.utp_process_udp(s.ctx, (*C.byte)(&b[0]), C.size_t(n), sa, sal) != 0 {
				socketUtpPacketsReceived.Add(1)
			} else {
				s.onReadNonUtp(b[:n], addr)
			}
			C.utp_issue_deferred_acks(s.ctx)
			C.utp_check_timeouts(s.ctx)
		}()
	}
}

func (s *Socket) timeoutChecker() {
	for {
		mu.Lock()
		if s.closed {
			mu.Unlock()
			return
		}
		// C.utp_issue_deferred_acks(s.ctx)
		C.utp_check_timeouts(s.ctx)
		mu.Unlock()
		time.Sleep(500 * time.Millisecond)
	}
}

func (me *Socket) Close() error {
	mu.Lock()
	defer mu.Unlock()
	if me.closed {
		return nil
	}
	C.utp_destroy(me.ctx)
	me.ctx = nil
	me.pc.Close()
	close(me.backlog)
	close(me.nonUtpReads)
	me.closed = true
	return nil
}

func (me *Socket) Addr() net.Addr {
	return me.pc.LocalAddr()
}

func (me *Socket) LocalAddr() net.Addr {
	return me.pc.LocalAddr()
}

func (s *Socket) Accept() (net.Conn, error) {
	nc, ok := <-s.backlog
	if !ok {
		return nil, errors.New("closed")
	}
	return nc, nil
}

func (s *Socket) Dial(addr string) (net.Conn, error) {
	return s.DialTimeout(addr, 0)
}

func (s *Socket) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) {
	ctx := context.Background()
	if timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, timeout)
		defer cancel()
	}
	return s.DialContext(ctx, addr)
}

func (s *Socket) DialContext(ctx context.Context, addr string) (net.Conn, error) {
	ua, err := net.ResolveUDPAddr(s.Addr().Network(), addr)
	if err != nil {
		panic(err)
	}
	sa, sl := netAddrToLibSockaddr(ua)
	mu.Lock()
	c := s.newConn(C.utp_create_socket(s.ctx))
	C.utp_connect(c.s, sa, sl)
	defer mu.Unlock()
	err = c.waitForConnect(ctx)
	if err != nil {
		c.close()
		return nil, err
	}
	return c, nil
}

func (s *Socket) pushBacklog(c *Conn) {
	select {
	case s.backlog <- c:
	default:
		c.close()
	}
}

func (s *Socket) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	p, ok := <-s.nonUtpReads
	if !ok {
		err = errors.New("closed")
		return
	}
	n = copy(b, p.b)
	addr = p.from
	return
}

func (s *Socket) onReadNonUtp(b []byte, from net.Addr) {
	socketNonUtpPacketsReceived.Add(1)
	select {
	case s.nonUtpReads <- packet{append([]byte(nil), b...), from}:
	default:
		// log.Printf("dropped non utp packet: no room in buffer")
		nonUtpPacketsDropped.Add(1)
	}
}

func (s *Socket) SetReadDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetWriteDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) SetDeadline(t time.Time) error {
	panic("not implemented")
}

func (s *Socket) WriteTo(b []byte, addr net.Addr) (int, error) {
	return s.pc.WriteTo(b, addr)
}

func (s *Socket) ReadBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_RCVBUF))
}

func (s *Socket) WriteBufferLen() int {
	mu.Lock()
	defer mu.Unlock()
	return int(C.utp_context_get_option(s.ctx, C.UTP_SNDBUF))
}

func (s *Socket) SetWriteBufferLen(len int) {
	mu.Lock()
	defer mu.Unlock()
	i := C.utp_context_set_option(s.ctx, C.UTP_SNDBUF, C.int(len))
	if i != 0 {
		panic(i)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
)

func sync(c *cli.Context) {
	repo := c.String("repo")
	gopath := c.String("gopath")

	absRepo, err := filepath.Abs(repo)
	if err != nil {
		println("could not resolve repo: " + err.Error())
		os.Exit(1)
	}

	absGopath, err := filepath.Abs(gopath)
	if err != nil {
		println("could not resolve gopath: " + err.Error())
		os.Exit(1)
	}

	pkgRoots := map[string]*Repo{}

	for _, dep := range c.Args() {
		root, repo, err := getDepRoot(absRepo, absGopath, dep)
		if err != nil {
			println("failed to get dependency repo: " + err.Error())
			os.Exit(1)
		}

		pkgRoots[root] = repo
	}

	existingSubmodules, err := detectExistingGoSubmodules(repo, gopath)
	if err != nil {
		println("failed to detect existing submodules: " + err.Error())
		os.Exit(1)
	}

	gitmodules := filepath.Join(repo, ".gitmodules")

	for _, submodule := range existingSubmodules {
		rm := exec.Command("git", "rm", "--cached", "-f", submodule)
		rm.Dir = repo
		rm.Stderr = os.Stderr

		err := rm.Run()
		if err != nil {
			println("error clearing submodule: " + err.Error())
			os.Exit(1)
		}

		gitConfig := exec.Command("git", "config", "--file", gitmodules, "--remove-section", "submodule."+submodule)
		gitConfig.Dir = repo
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error removing submodule config: " + err.Error())
			os.Exit(1)
		}

		gitAdd := exec.Command("git", "add", gitmodules)
		gitAdd.Dir = repo
		gitAdd.Stderr = os.Stderr

		err = gitAdd.Run()
		if err != nil {
			println("error staging submodule config: " + err.Error())
			os.Exit(1)
		}
	}

	for pkgRoot, pkgRepo := range pkgRoots {
		relRoot, err := filepath.Rel(absRepo, pkgRoot)
		if err != nil {
			println("could not resolve submodule: " + err.Error())
			os.Exit(1)
		}

		fmt.Println(relRoot)

		add := exec.Command("git", "add", pkgRoot)
		add.Dir = repo
		add.Stderr = os.Stderr

		err = add.Run()
		if err != nil {
			println("error clearing submodule: " + err.Error())
			os.Exit(1)
		}

		if pkgRepo == nil {
			// non-git dependency; vendored
			continue
		}

		gitConfig := exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".path", relRoot)
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error configuring submodule: " + err.Error())
			os.Exit(1)
		}

		gitConfig = exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".url", pkgRepo.Origin)
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error configuring submodule: " + err.Error())
			os.Exit(1)
		}

		if pkgRepo.Branch != "" {
			gitConfig = exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".branch", pkgRepo.Branch)
			gitConfig.Stderr = os.Stderr

			err = gitConfig.Run()
			if err != nil {
				println("error configuring submodule: " + err.Error())
				os.Exit(1)
			}
		}

		gitAdd := exec.Command("git", "add", gitmodules)
		gitAdd.Dir = repo
		gitAdd.Stderr = os.Stderr

		err = gitAdd.Run()
		if err != nil {
			println("error staging submodule config: " + err.Error())
			os.Exit(1)
		}
	}
}

func detectExistingGoSubmodules(repo string, gopath string) ([]string, error) {
	srcPath := filepath.Join(gopath, "src")

	submoduleStatus := exec.Command("git", "submodule", "status", srcPath)
	submoduleStatus.Dir = repo

	submoduleStatus.Stderr = os.Stderr

	statusOut, err := submoduleStatus.StdoutPipe()
	if err != nil {
		return nil, err
	}

	lineScanner := bufio.NewScanner(statusOut)

	err = submoduleStatus.Start()
	if err != nil {
		return nil, err
	}

	submodules := []string{}
	for lineScanner.Scan() {
		segments := strings.Split(lineScanner.Text()[1:], " ")

		if len(segments) < 2 {
			return nil, fmt.Errorf("invalid git status output: %q", lineScanner.Text())
		}

		submodules = append(submodules, segments[1])
	}

	err = submoduleStatus.Wait()
	if err != nil {
		return nil, err
	}

	return submodules, nil
}
<commit_msg>remove only submodules we know to be gone<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
)

func sync(c *cli.Context) {
	repo := c.String("repo")
	gopath := c.String("gopath")

	absRepo, err := filepath.Abs(repo)
	if err != nil {
		println("could not resolve repo: " + err.Error())
		os.Exit(1)
	}

	absGopath, err := filepath.Abs(gopath)
	if err != nil {
		println("could not resolve gopath: " + err.Error())
		os.Exit(1)
	}

	pkgRoots := map[string]*Repo{}

	for _, dep := range c.Args() {
		root, repo, err := getDepRoot(absRepo, absGopath, dep)
		if err != nil {
			println("failed to get dependency repo: " + err.Error())
			os.Exit(1)
		}

		pkgRoots[root] = repo
	}

	existingSubmodules, err := detectExistingGoSubmodules(repo, gopath)
	if err != nil {
		println("failed to detect existing submodules: " + err.Error())
		os.Exit(1)
	}

	gitmodules := filepath.Join(repo, ".gitmodules")

	submodulesToRemove := map[string]bool{}
	for _, submodule := range existingSubmodules {
		submodulesToRemove[submodule] = true
	}

	for pkgRoot, pkgRepo := range pkgRoots {
		relRoot, err := filepath.Rel(absRepo, pkgRoot)
		if err != nil {
			println("could not resolve submodule: " + err.Error())
			os.Exit(1)
		}

		fmt.Println("\x1b[32msyncing " + relRoot + "\x1b[0m")

		// keep this submodule
		delete(submodulesToRemove, relRoot)

		add := exec.Command("git", "add", pkgRoot)
		add.Dir = repo
		add.Stderr = os.Stderr

		err = add.Run()
		if err != nil {
			println("error clearing submodule: " + err.Error())
			os.Exit(1)
		}

		if pkgRepo == nil {
			// non-git dependency; vendored
			continue
		}

		gitConfig := exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".path", relRoot)
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error configuring submodule: " + err.Error())
			os.Exit(1)
		}

		gitConfig = exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".url", pkgRepo.Origin)
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error configuring submodule: " + err.Error())
			os.Exit(1)
		}

		if pkgRepo.Branch != "" {
			gitConfig = exec.Command("git", "config", "--file", gitmodules, "submodule."+relRoot+".branch", pkgRepo.Branch)
			gitConfig.Stderr = os.Stderr

			err = gitConfig.Run()
			if err != nil {
				println("error configuring submodule: " + err.Error())
				os.Exit(1)
			}
		}

		gitAdd := exec.Command("git", "add", gitmodules)
		gitAdd.Dir = repo
		gitAdd.Stderr = os.Stderr

		err = gitAdd.Run()
		if err != nil {
			println("error staging submodule config: " + err.Error())
			os.Exit(1)
		}
	}

	for submodule, _ := range submodulesToRemove {
		fmt.Println("\x1b[31mremoving " + submodule + "\x1b[0m")

		rm := exec.Command("git", "rm", "--cached", "-f", submodule)
		rm.Dir = repo
		rm.Stderr = os.Stderr

		err := rm.Run()
		if err != nil {
			println("error clearing submodule: " + err.Error())
			os.Exit(1)
		}

		gitConfig := exec.Command("git", "config", "--file", gitmodules, "--remove-section", "submodule."+submodule)
		gitConfig.Dir = repo
		gitConfig.Stderr = os.Stderr

		err = gitConfig.Run()
		if err != nil {
			println("error removing submodule config: " + err.Error())
			os.Exit(1)
		}

		gitAdd := exec.Command("git", "add", gitmodules)
		gitAdd.Dir = repo
		gitAdd.Stderr = os.Stderr

		err = gitAdd.Run()
		if err != nil {
			println("error staging submodule config: " + err.Error())
			os.Exit(1)
		}
	}
}

func detectExistingGoSubmodules(repo string, gopath string) ([]string, error) {
	srcPath := filepath.Join(gopath, "src")

	submoduleStatus := exec.Command("git", "submodule", "status", srcPath)
	submoduleStatus.Dir = repo

	submoduleStatus.Stderr = os.Stderr

	statusOut, err := submoduleStatus.StdoutPipe()
	if err != nil {
		return nil, err
	}

	lineScanner := bufio.NewScanner(statusOut)

	err = submoduleStatus.Start()
	if err != nil {
		return nil, err
	}

	submodules := []string{}
	for lineScanner.Scan() {
		segments := strings.Split(lineScanner.Text()[1:], " ")

		if len(segments) < 2 {
			return nil, fmt.Errorf("invalid git status output: %q", lineScanner.Text())
		}

		submodules = append(submodules, segments[1])
	}

	err = submoduleStatus.Wait()
	if err != nil {
		return nil, err
	}

	return submodules, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/build"
	"os"
	"path/filepath"
)

// Sync synchronizes the state of the project code (what dependencies it has), the pinlist, and all
// (transitive) dependencies in the cache.
//   * First, analyze the project code to get a list of its dependencies
//   * For each dependency, update the cache/dependency list:
//     - If the dependency is not cached, pull down the rev specifed (or latest if dep not in the pinlist)
//     - If the dependency is cached and the cached repo is dirty, fail
//     - If the dependency is cached and the cached rev is different from the specified rev, update the cached
//			 repo to the specified rev
//		 - Add the dep's transitive non-stdlib deps to the dependency list to be processed (if not done already)
//		* Recreate the pinlist from the set of deps and versions in the updated dependency list (this will also
//			have the effect of removing outdated dependencies)
//		* Write out the new pinlist
func Sync(root, gopath string) error {
	// Make a map out of the current pinlist
	pinnedVersions, err := makePinMap(root)
	if err != nil {
		return err
	}

	// Get the list of immediate dependencies from the code
	context := new(build.Context)
	*context = build.Default
	context.GOPATH = gopath

	immediateDeps, projectPackages, err := findProjectDeps(context, root)
	if err != nil {
		return err
	}

	fmt.Println("Found project packages:")
	for _, pkg := range projectPackages {
		fmt.Printf("\t%s\n", pkg)
	}

	// Now sync each dependency, adding transitive deps to the to-process list as we go.
	var processedDeps smap
	toProcessDeps := immediateDeps
	// Multiple packages may map to the same repo, and thus must necessarily must be at the same rev. We will
	// fill in the following two mappings as we go and fail if there's a contradiction.
	pinnedVersionsByRepo := make(map[string]string)
	packageToRepo := make(map[string]string)
	cacheDir := filepath.Join(root, projectDirName, cacheDirName, "src")
	for len(toProcessDeps) > 0 {
		importPath := toProcessDeps[0]
		toProcessDeps.Remove(importPath)
		repo, err := RepoRootForImportPath(importPath, false)
		if err != nil {
			return err
		}
		switch repo.VCS.Cmd.Cmd {
		case "git", "hg":
		default:
			return fmt.Errorf("%s is not a VCS supported by glp", repo.VCS.Cmd.Name)
		}
		packageToRepo[importPath] = repo.Root
		rev, pinned := pinnedVersions[importPath]
		if pinned {
			rev2, ok := pinnedVersionsByRepo[repo.Root]
			if ok && rev != rev2 {
				return fmt.Errorf("Multiple packages with conflicting revs map to the repo %s", repo.Root)
			}
		}

		// Check if the repo is already cached at all.
		repoDir := filepath.Join(cacheDir, repo.Root)
		if _, err := os.Stat(repoDir); err != nil {
			if !os.IsNotExist(err) {
				return err
			}
			// Repo hasn't been downloaded; fetch latest
			fmt.Printf("Dep repo %s does not exist; downloading...", repo.Root)
			if err := repo.VCS.Create(repoDir, repo.Repo); err != nil {
				return err
			}
			fmt.Println("done.")
		}

		// Get the state of the repo (current rev and whether it's dirty)
		currentRev, dirty, err := repo.VCS.GetRev(repoDir)
		if err != nil {
			return err
		}
		if dirty {
			return fmt.Errorf("dep repo at %s is currently dirty (has untracked changes)", repoDir)
		}
		if pinned && currentRev != rev {
			// If the repo is does not matched the pinned version, update to that version
			fmt.Printf("Updating dep repo at %s to rev %s\n", repoDir, rev)
			if err := repo.VCS.UpdateRev(repoDir, rev); err != nil {
				return err
			}
		}
		if !pinned {
			rev = currentRev
		}

		// Update the pin list -- these are possibly no-ops
		pinnedVersions[importPath] = rev
		pinnedVersionsByRepo[repo.Root] = rev

		// Add the dependencies of this package to the to-process list if they haven't already been inspected
		packageDir := filepath.Join(cacheDir, importPath)
		deps, err := findDeps(context, packageDir)
		if err != nil {
			return err
		}
		for _, dep := range deps {
			if !processedDeps.Contains(dep) {
				toProcessDeps.Add(dep)
			}
		}

		processedDeps.Add(importPath)
	}
	fmt.Println("Dependencies are up-to-date.")

	// Now write out the updated pin list. We reconstruct from pinnedVersionsByRepo and packageToRepo becaused
	// pinnedVersions may contain outdated deps that are not needed by the current project code.
	pinlistFilename := filepath.Join(root, projectDirName, pinlistName)
	newPinlist := new(Pinlist)
	var deps smap
	for dep := range packageToRepo {
		deps.Add(dep)
	}
	for _, dep := range deps {
		newPinlist.Deps = append(newPinlist.Deps, Dep{
			Name: dep,
			Rev:  pinnedVersionsByRepo[packageToRepo[dep]],
		})
	}
	return newPinlist.Save(pinlistFilename)
}

// makePinMap creates a map of name -> rev from the entries in the pinlist in the project located at root.
func makePinMap(root string) (map[string]string, error) {
	pinnedVersions := make(map[string]string)
	pinlistFilename := filepath.Join(root, projectDirName, pinlistName)
	pinlist, err := LoadPinlist(pinlistFilename)
	switch {
	case err == nil:
		for _, dep := range pinlist.Deps {
			pinnedVersions[dep.Name] = dep.Rev
		}
	case os.IsNotExist(err):
	default:
		return nil, err
	}
	return pinnedVersions, nil
}

func findProjectDeps(context *build.Context, root string) (immediateDeps, projectPackages smap, err error) {
	srcDir := filepath.Join(root, "src")
	possiblePackageDirs := FindDirsRecursively(srcDir)
	possiblePackageDirs = append(possiblePackageDirs, ".")
	for _, dir := range possiblePackageDirs {
		pkg, err := context.ImportDir(dir, 0)
		if err != nil {
			if _, ok := err.(*build.NoGoError); ok {
				continue
			}
			return nil, nil, err
		}

		packageName := dir
		if dir != "." {
			packageName, err = filepath.Rel(srcDir, dir)
			if err != nil {
				return nil, nil, err
			}
		}
		projectPackages.Add(packageName)
		for _, dep := range findNonStdDeps(context, pkg) {
			immediateDeps.Add(dep)
		}
	}

	// Remove the intra-project deps from immediateDeps
	for _, dep := range immediateDeps {
		if projectPackages.Contains(dep) {
			immediateDeps.Remove(dep)
		}
	}

	return immediateDeps, projectPackages, nil
}

// Find the first-level deps for a package
func findDeps(context *build.Context, dir string) (immediateDeps smap, err error) {
	pkg, err := context.ImportDir(dir, 0)
	if err != nil {
		return nil, err
	}
	return findNonStdDeps(context, pkg), nil
}

func findNonStdDeps(context *build.Context, pkg *build.Package) (immediateDeps smap) {
	imports := pkg.Imports
	imports = append(imports, pkg.TestImports...)
	imports = append(imports, pkg.XTestImports...)
	for _, imp := range imports {
		if !pkgIsStd(context, imp) {
			immediateDeps.Add(imp)
		}
	}
	return immediateDeps
}

func pkgIsStd(context *build.Context, importPath string) bool {
	pkg, err := context.Import(importPath, ".", 0)
	return err == nil && pkg.Goroot
}
<commit_msg>Bugfix for detecting dependencies within the project<commit_after>package main

import (
	"fmt"
	"go/build"
	"os"
	"path/filepath"
)

// Sync synchronizes the state of the project code (what dependencies it has), the pinlist, and all
// (transitive) dependencies in the cache.
//   * First, analyze the project code to get a list of its dependencies
//   * For each dependency, update the cache/dependency list:
//     - If the dependency is not cached, pull down the rev specifed (or latest if dep not in the pinlist)
//     - If the dependency is cached and the cached repo is dirty, fail
//     - If the dependency is cached and the cached rev is different from the specified rev, update the cached
//			 repo to the specified rev
//		 - Add the dep's transitive non-stdlib deps to the dependency list to be processed (if not done already)
//		* Recreate the pinlist from the set of deps and versions in the updated dependency list (this will also
//			have the effect of removing outdated dependencies)
//		* Write out the new pinlist
func Sync(root, gopath string) error {
	// Make a map out of the current pinlist
	pinnedVersions, err := makePinMap(root)
	if err != nil {
		return err
	}

	// Get the list of immediate dependencies from the code
	context := new(build.Context)
	*context = build.Default
	context.GOPATH = gopath

	immediateDeps, projectPackages, err := findProjectDeps(context, root)
	if err != nil {
		return err
	}

	fmt.Println("Found project packages:")
	for _, pkg := range projectPackages {
		fmt.Printf("\t%s\n", pkg)
	}

	// Now sync each dependency, adding transitive deps to the to-process list as we go.
	var processedDeps smap
	toProcessDeps := immediateDeps
	// Multiple packages may map to the same repo, and thus must necessarily must be at the same rev. We will
	// fill in the following two mappings as we go and fail if there's a contradiction.
	pinnedVersionsByRepo := make(map[string]string)
	packageToRepo := make(map[string]string)
	cacheDir := filepath.Join(root, projectDirName, cacheDirName, "src")
	for len(toProcessDeps) > 0 {
		importPath := toProcessDeps[0]
		toProcessDeps.Remove(importPath)
		repo, err := RepoRootForImportPath(importPath, false)
		if err != nil {
			return err
		}
		switch repo.VCS.Cmd.Cmd {
		case "git", "hg":
		default:
			return fmt.Errorf("%s is not a VCS supported by glp", repo.VCS.Cmd.Name)
		}
		packageToRepo[importPath] = repo.Root
		rev, pinned := pinnedVersions[importPath]
		if pinned {
			rev2, ok := pinnedVersionsByRepo[repo.Root]
			if ok && rev != rev2 {
				return fmt.Errorf("Multiple packages with conflicting revs map to the repo %s", repo.Root)
			}
		}

		// Check if the repo is already cached at all.
		repoDir := filepath.Join(cacheDir, repo.Root)
		if _, err := os.Stat(repoDir); err != nil {
			if !os.IsNotExist(err) {
				return err
			}
			// Repo hasn't been downloaded; fetch latest
			fmt.Printf("Dep repo %s does not exist; downloading...", repo.Root)
			if err := repo.VCS.Create(repoDir, repo.Repo); err != nil {
				return err
			}
			fmt.Println("done.")
		}

		// Get the state of the repo (current rev and whether it's dirty)
		currentRev, dirty, err := repo.VCS.GetRev(repoDir)
		if err != nil {
			return err
		}
		if dirty {
			return fmt.Errorf("dep repo at %s is currently dirty (has untracked changes)", repoDir)
		}
		if pinned && currentRev != rev {
			// If the repo is does not matched the pinned version, update to that version
			fmt.Printf("Updating dep repo at %s to rev %s\n", repoDir, rev)
			if err := repo.VCS.UpdateRev(repoDir, rev); err != nil {
				return err
			}
		}
		if !pinned {
			rev = currentRev
		}

		// Update the pin list -- these are possibly no-ops
		pinnedVersions[importPath] = rev
		pinnedVersionsByRepo[repo.Root] = rev

		// Add the dependencies of this package to the to-process list if they haven't already been inspected
		packageDir := filepath.Join(cacheDir, importPath)
		deps, err := findDeps(context, packageDir)
		if err != nil {
			return err
		}
		for _, dep := range deps {
			if !processedDeps.Contains(dep) {
				toProcessDeps.Add(dep)
			}
		}

		processedDeps.Add(importPath)
	}
	fmt.Println("Dependencies are up-to-date.")

	// Now write out the updated pin list. We reconstruct from pinnedVersionsByRepo and packageToRepo becaused
	// pinnedVersions may contain outdated deps that are not needed by the current project code.
	pinlistFilename := filepath.Join(root, projectDirName, pinlistName)
	newPinlist := new(Pinlist)
	var deps smap
	for dep := range packageToRepo {
		deps.Add(dep)
	}
	for _, dep := range deps {
		newPinlist.Deps = append(newPinlist.Deps, Dep{
			Name: dep,
			Rev:  pinnedVersionsByRepo[packageToRepo[dep]],
		})
	}
	return newPinlist.Save(pinlistFilename)
}

// makePinMap creates a map of name -> rev from the entries in the pinlist in the project located at root.
func makePinMap(root string) (map[string]string, error) {
	pinnedVersions := make(map[string]string)
	pinlistFilename := filepath.Join(root, projectDirName, pinlistName)
	pinlist, err := LoadPinlist(pinlistFilename)
	switch {
	case err == nil:
		for _, dep := range pinlist.Deps {
			pinnedVersions[dep.Name] = dep.Rev
		}
	case os.IsNotExist(err):
	default:
		return nil, err
	}
	return pinnedVersions, nil
}

func findProjectDeps(context *build.Context, root string) (immediateDeps, projectPackages smap, err error) {
	srcDir := filepath.Join(root, "src")
	possiblePackageDirs := FindDirsRecursively(srcDir)
	possiblePackageDirs = append(possiblePackageDirs, ".")
	for _, dir := range possiblePackageDirs {
		pkg, err := context.ImportDir(dir, 0)
		if err != nil {
			if _, ok := err.(*build.NoGoError); ok {
				continue
			}
			return nil, nil, err
		}

		packageName := dir
		if dir != "." {
			packageName, err = filepath.Rel(srcDir, dir)
			if err != nil {
				return nil, nil, err
			}
		}
		projectPackages.Add(packageName)
		for _, dep := range findNonStdDeps(context, pkg) {
			immediateDeps.Add(dep)
		}
	}

	// Remove the intra-project deps from immediateDeps. We need two steps because it's not valid to remove
	// items from an smap while traversing it.
	var toRemove []string
	for _, dep := range immediateDeps {
		if projectPackages.Contains(dep) {
			toRemove = append(toRemove, dep)
		}
	}
	for _, dep := range toRemove {
		immediateDeps.Remove(dep)
	}

	return immediateDeps, projectPackages, nil
}

// Find the first-level deps for a package
func findDeps(context *build.Context, dir string) (immediateDeps smap, err error) {
	pkg, err := context.ImportDir(dir, 0)
	if err != nil {
		return nil, err
	}
	return findNonStdDeps(context, pkg), nil
}

func findNonStdDeps(context *build.Context, pkg *build.Package) (immediateDeps smap) {
	imports := pkg.Imports
	imports = append(imports, pkg.TestImports...)
	imports = append(imports, pkg.XTestImports...)
	for _, imp := range imports {
		if !pkgIsStd(context, imp) {
			immediateDeps.Add(imp)
		}
	}
	return immediateDeps
}

func pkgIsStd(context *build.Context, importPath string) bool {
	pkg, err := context.Import(importPath, ".", 0)
	return err == nil && pkg.Goroot
}
<|endoftext|>


<commit_before>package gps

import (
	"fmt"
	"sync"
)

type source interface {
	syncLocal() error
	checkExistence(sourceExistence) bool
	exportVersionTo(Version, string) error
	getManifestAndLock(ProjectRoot, Version) (Manifest, Lock, error)
	listPackages(ProjectRoot, Version) (PackageTree, error)
	listVersions() ([]Version, error)
	revisionPresentIn(Revision) (bool, error)
}

type sourceMetaCache struct {
	//Version  string                   // TODO(sdboyer) use this
	infos  map[Revision]projectInfo
	ptrees map[Revision]PackageTree
	vMap   map[UnpairedVersion]Revision
	rMap   map[Revision][]UnpairedVersion
	// TODO(sdboyer) mutexes. actually probably just one, b/c complexity
}

// projectInfo holds manifest and lock
type projectInfo struct {
	Manifest
	Lock
}

type existence struct {
	// The existence levels for which a search/check has been performed
	s sourceExistence

	// The existence levels verified to be present through searching
	f sourceExistence
}

func newMetaCache() *sourceMetaCache {
	return &sourceMetaCache{
		infos:  make(map[Revision]projectInfo),
		ptrees: make(map[Revision]PackageTree),
		vMap:   make(map[UnpairedVersion]Revision),
		rMap:   make(map[Revision][]UnpairedVersion),
	}
}

type baseVCSSource struct {
	// Object for the cache repository
	crepo *repo

	// Indicates the extent to which we have searched for, and verified, the
	// existence of the project/repo.
	ex existence

	// ProjectAnalyzer used to fulfill getManifestAndLock
	an ProjectAnalyzer

	// The project metadata cache. This is (or is intended to be) persisted to
	// disk, for reuse across solver runs.
	dc *sourceMetaCache

	// lvfunc allows the other vcs source types that embed this type to inject
	// their listVersions func into the baseSource, for use as needed.
	lvfunc func() (vlist []Version, err error)

	// lock to serialize access to syncLocal
	synclock sync.Mutex

	// Globalish flag indicating whether a "full" sync has been performed. Also
	// used as a one-way gate to ensure that the full syncing routine is never
	// run more than once on a given source instance.
	allsync bool

	// The error, if any, that occurred on syncLocal
	syncerr error

	// Whether the cache has the latest info on versions
	cvsync bool
}

func (bs *baseVCSSource) getManifestAndLock(r ProjectRoot, v Version) (Manifest, Lock, error) {
	if err := bs.ensureCacheExistence(); err != nil {
		return nil, nil, err
	}

	rev, err := bs.toRevOrErr(v)
	if err != nil {
		return nil, nil, err
	}

	// Return the info from the cache, if we already have it
	if pi, exists := bs.dc.infos[rev]; exists {
		return pi.Manifest, pi.Lock, nil
	}

	// Cache didn't help; ensure our local is fully up to date.
	err = bs.syncLocal()
	if err != nil {
		return nil, nil, err
	}

	bs.crepo.mut.Lock()
	// Always prefer a rev, if it's available
	if pv, ok := v.(PairedVersion); ok {
		err = bs.crepo.r.UpdateVersion(pv.Underlying().String())
	} else {
		err = bs.crepo.r.UpdateVersion(v.String())
	}
	bs.crepo.mut.Unlock()

	if err != nil {
		// TODO(sdboyer) More-er proper-er error
		panic(fmt.Sprintf("canary - why is checkout/whatever failing: %s %s %s", bs.crepo.r.LocalPath(), v.String(), unwrapVcsErr(err)))
	}

	bs.crepo.mut.RLock()
	m, l, err := bs.an.DeriveManifestAndLock(bs.crepo.r.LocalPath(), r)
	// TODO(sdboyer) cache results
	bs.crepo.mut.RUnlock()

	if err == nil {
		if l != nil {
			l = prepLock(l)
		}

		// If m is nil, prepManifest will provide an empty one.
		pi := projectInfo{
			Manifest: prepManifest(m),
			Lock:     l,
		}

		bs.dc.infos[rev] = pi

		return pi.Manifest, pi.Lock, nil
	}

	return nil, nil, unwrapVcsErr(err)
}

// toRevision turns a Version into a Revision, if doing so is possible based on
// the information contained in the version itself, or in the cache maps.
func (dc *sourceMetaCache) toRevision(v Version) Revision {
	switch t := v.(type) {
	case Revision:
		return t
	case PairedVersion:
		return t.Underlying()
	case UnpairedVersion:
		// This will return the empty rev (empty string) if we don't have a
		// record of it. It's up to the caller to decide, for example, if
		// it's appropriate to update the cache.
		return dc.vMap[t]
	default:
		panic(fmt.Sprintf("Unknown version type %T", v))
	}
}

// toUnpaired turns a Version into an UnpairedVersion, if doing so is possible
// based on the information contained in the version itself, or in the cache
// maps.
//
// If the input is a revision and multiple UnpairedVersions are associated with
// it, whatever happens to be the first is returned.
func (dc *sourceMetaCache) toUnpaired(v Version) UnpairedVersion {
	switch t := v.(type) {
	case UnpairedVersion:
		return t
	case PairedVersion:
		return t.Unpair()
	case Revision:
		if upv, has := dc.rMap[t]; has && len(upv) > 0 {
			return upv[0]
		}
		return nil
	default:
		panic(fmt.Sprintf("unknown version type %T", v))
	}
}

func (bs *baseVCSSource) revisionPresentIn(r Revision) (bool, error) {
	// First and fastest path is to check the data cache to see if the rev is
	// present. This could give us false positives, but the cases where that can
	// occur would require a type of cache staleness that seems *exceedingly*
	// unlikely to occur.
	if _, has := bs.dc.infos[r]; has {
		return true, nil
	} else if _, has := bs.dc.rMap[r]; has {
		return true, nil
	}

	err := bs.ensureCacheExistence()
	if err != nil {
		return false, err
	}

	bs.crepo.mut.RLock()
	defer bs.crepo.mut.RUnlock()
	return bs.crepo.r.IsReference(string(r)), nil
}

func (bs *baseVCSSource) ensureCacheExistence() error {
	// Technically, methods could could attempt to return straight from the
	// metadata cache even if the repo cache doesn't exist on disk. But that
	// would allow weird state inconsistencies (cache exists, but no repo...how
	// does that even happen?) that it'd be better to just not allow so that we
	// don't have to think about it elsewhere
	if !bs.checkExistence(existsInCache) {
		if bs.checkExistence(existsUpstream) {
			bs.crepo.mut.Lock()
			if bs.crepo.synced {
				// A second ensure call coming in while the first is completing
				// isn't terribly unlikely, especially for a large repo. In that
				// event, the synced flag will have flipped on by the time we
				// acquire the lock. If it has, there's no need to do this work
				// twice.
				bs.crepo.mut.Unlock()
				return nil
			}

			err := bs.crepo.r.Get()

			if err != nil {
				bs.crepo.mut.Unlock()
				return fmt.Errorf("failed to create repository cache for %s with err:\n%s", bs.crepo.r.Remote(), unwrapVcsErr(err))
			}

			bs.crepo.synced = true
			bs.ex.s |= existsInCache
			bs.ex.f |= existsInCache
			bs.crepo.mut.Unlock()
		} else {
			return fmt.Errorf("project %s does not exist upstream", bs.crepo.r.Remote())
		}
	}

	return nil
}

// checkExistence provides a direct method for querying existence levels of the
// source. It will only perform actual searching (local fs or over the network)
// if no previous attempt at that search has been made.
//
// Note that this may perform read-ish operations on the cache repo, and it
// takes a lock accordingly. This makes it unsafe to call from a segment where
// the cache repo mutex is already write-locked, as deadlock will occur.
func (bs *baseVCSSource) checkExistence(ex sourceExistence) bool {
	if bs.ex.s&ex != ex {
		if ex&existsInVendorRoot != 0 && bs.ex.s&existsInVendorRoot == 0 {
			panic("should now be implemented in bridge")
		}
		if ex&existsInCache != 0 && bs.ex.s&existsInCache == 0 {
			bs.crepo.mut.RLock()
			bs.ex.s |= existsInCache
			if bs.crepo.r.CheckLocal() {
				bs.ex.f |= existsInCache
			}
			bs.crepo.mut.RUnlock()
		}
		if ex&existsUpstream != 0 && bs.ex.s&existsUpstream == 0 {
			bs.crepo.mut.RLock()
			bs.ex.s |= existsUpstream
			if bs.crepo.r.Ping() {
				bs.ex.f |= existsUpstream
			}
			bs.crepo.mut.RUnlock()
		}
	}

	return ex&bs.ex.f == ex
}

// syncLocal ensures the local data we have about the source is fully up to date
// with what's out there over the network.
func (bs *baseVCSSource) syncLocal() error {
	// Ensure we only have one goroutine doing this at a time
	bs.synclock.Lock()
	defer bs.synclock.Unlock()

	// ...and that we only ever do it once
	if bs.allsync {
		// Return the stored err, if any
		return bs.syncerr
	}

	bs.allsync = true
	// First, ensure the local instance exists
	bs.syncerr = bs.ensureCacheExistence()
	if bs.syncerr != nil {
		return bs.syncerr
	}

	_, bs.syncerr = bs.lvfunc()
	if bs.syncerr != nil {
		return bs.syncerr
	}

	// This case is really just for git repos, where the lvfunc doesn't
	// guarantee that the local repo is synced
	if !bs.crepo.synced {
		bs.crepo.mut.Lock()
		bs.syncerr = fmt.Errorf("failed fetching latest updates with err: %s", unwrapVcsErr(bs.crepo.r.Update()))
		if bs.syncerr != nil {
			bs.crepo.mut.Unlock()
			return bs.syncerr
		}
		bs.crepo.synced = true
		bs.crepo.mut.Unlock()
	}

	return nil
}

func (bs *baseVCSSource) listPackages(pr ProjectRoot, v Version) (ptree PackageTree, err error) {
	if err = bs.ensureCacheExistence(); err != nil {
		return
	}

	var r Revision
	if r, err = bs.toRevOrErr(v); err != nil {
		return
	}

	// Return the ptree from the cache, if we already have it
	var exists bool
	if ptree, exists = bs.dc.ptrees[r]; exists {
		return
	}

	// Not in the cache; check out the version and do the analysis
	bs.crepo.mut.Lock()
	// Check out the desired version for analysis
	if r != "" {
		// Always prefer a rev, if it's available
		err = bs.crepo.r.UpdateVersion(string(r))
	} else {
		// If we don't have a rev, ensure the repo is up to date, otherwise we
		// could have a desync issue
		if !bs.crepo.synced {
			err = bs.crepo.r.Update()
			if err != nil {
				err = fmt.Errorf("could not fetch latest updates into repository: %s", unwrapVcsErr(err))
				return
			}
			bs.crepo.synced = true
		}
		err = bs.crepo.r.UpdateVersion(v.String())
	}

	if err == nil {
		ptree, err = ListPackages(bs.crepo.r.LocalPath(), string(pr))
		// TODO(sdboyer) cache errs?
		if err == nil {
			bs.dc.ptrees[r] = ptree
		}
	} else {
		err = unwrapVcsErr(err)
	}
	bs.crepo.mut.Unlock()

	return
}

// toRevOrErr makes all efforts to convert a Version into a rev, including
// updating the cache repo (if needed). It does not guarantee that the returned
// Revision actually exists in the repository (as one of the cheaper methods may
// have had bad data).
func (bs *baseVCSSource) toRevOrErr(v Version) (r Revision, err error) {
	r = bs.dc.toRevision(v)
	if r == "" {
		// Rev can be empty if:
		//  - The cache is unsynced
		//  - A version was passed that used to exist, but no longer does
		//  - A garbage version was passed. (Functionally indistinguishable from
		//  the previous)
		if !bs.cvsync {
			// call the lvfunc to sync the meta cache
			_, err = bs.lvfunc()
			if err != nil {
				return
			}
		}

		r = bs.dc.toRevision(v)
		// If we still don't have a rev, then the version's no good
		if r == "" {
			err = fmt.Errorf("version %s does not exist in source %s", v, bs.crepo.r.Remote())
		}
	}

	return
}

func (bs *baseVCSSource) exportVersionTo(v Version, to string) error {
	return bs.crepo.exportVersionTo(v, to)
}
<commit_msg>Only make an error if there's actually an error<commit_after>package gps

import (
	"fmt"
	"sync"
)

type source interface {
	syncLocal() error
	checkExistence(sourceExistence) bool
	exportVersionTo(Version, string) error
	getManifestAndLock(ProjectRoot, Version) (Manifest, Lock, error)
	listPackages(ProjectRoot, Version) (PackageTree, error)
	listVersions() ([]Version, error)
	revisionPresentIn(Revision) (bool, error)
}

type sourceMetaCache struct {
	//Version  string                   // TODO(sdboyer) use this
	infos  map[Revision]projectInfo
	ptrees map[Revision]PackageTree
	vMap   map[UnpairedVersion]Revision
	rMap   map[Revision][]UnpairedVersion
	// TODO(sdboyer) mutexes. actually probably just one, b/c complexity
}

// projectInfo holds manifest and lock
type projectInfo struct {
	Manifest
	Lock
}

type existence struct {
	// The existence levels for which a search/check has been performed
	s sourceExistence

	// The existence levels verified to be present through searching
	f sourceExistence
}

func newMetaCache() *sourceMetaCache {
	return &sourceMetaCache{
		infos:  make(map[Revision]projectInfo),
		ptrees: make(map[Revision]PackageTree),
		vMap:   make(map[UnpairedVersion]Revision),
		rMap:   make(map[Revision][]UnpairedVersion),
	}
}

type baseVCSSource struct {
	// Object for the cache repository
	crepo *repo

	// Indicates the extent to which we have searched for, and verified, the
	// existence of the project/repo.
	ex existence

	// ProjectAnalyzer used to fulfill getManifestAndLock
	an ProjectAnalyzer

	// The project metadata cache. This is (or is intended to be) persisted to
	// disk, for reuse across solver runs.
	dc *sourceMetaCache

	// lvfunc allows the other vcs source types that embed this type to inject
	// their listVersions func into the baseSource, for use as needed.
	lvfunc func() (vlist []Version, err error)

	// lock to serialize access to syncLocal
	synclock sync.Mutex

	// Globalish flag indicating whether a "full" sync has been performed. Also
	// used as a one-way gate to ensure that the full syncing routine is never
	// run more than once on a given source instance.
	allsync bool

	// The error, if any, that occurred on syncLocal
	syncerr error

	// Whether the cache has the latest info on versions
	cvsync bool
}

func (bs *baseVCSSource) getManifestAndLock(r ProjectRoot, v Version) (Manifest, Lock, error) {
	if err := bs.ensureCacheExistence(); err != nil {
		return nil, nil, err
	}

	rev, err := bs.toRevOrErr(v)
	if err != nil {
		return nil, nil, err
	}

	// Return the info from the cache, if we already have it
	if pi, exists := bs.dc.infos[rev]; exists {
		return pi.Manifest, pi.Lock, nil
	}

	// Cache didn't help; ensure our local is fully up to date.
	err = bs.syncLocal()
	if err != nil {
		return nil, nil, err
	}

	bs.crepo.mut.Lock()
	// Always prefer a rev, if it's available
	if pv, ok := v.(PairedVersion); ok {
		err = bs.crepo.r.UpdateVersion(pv.Underlying().String())
	} else {
		err = bs.crepo.r.UpdateVersion(v.String())
	}
	bs.crepo.mut.Unlock()

	if err != nil {
		// TODO(sdboyer) More-er proper-er error
		panic(fmt.Sprintf("canary - why is checkout/whatever failing: %s %s %s", bs.crepo.r.LocalPath(), v.String(), unwrapVcsErr(err)))
	}

	bs.crepo.mut.RLock()
	m, l, err := bs.an.DeriveManifestAndLock(bs.crepo.r.LocalPath(), r)
	// TODO(sdboyer) cache results
	bs.crepo.mut.RUnlock()

	if err == nil {
		if l != nil {
			l = prepLock(l)
		}

		// If m is nil, prepManifest will provide an empty one.
		pi := projectInfo{
			Manifest: prepManifest(m),
			Lock:     l,
		}

		bs.dc.infos[rev] = pi

		return pi.Manifest, pi.Lock, nil
	}

	return nil, nil, unwrapVcsErr(err)
}

// toRevision turns a Version into a Revision, if doing so is possible based on
// the information contained in the version itself, or in the cache maps.
func (dc *sourceMetaCache) toRevision(v Version) Revision {
	switch t := v.(type) {
	case Revision:
		return t
	case PairedVersion:
		return t.Underlying()
	case UnpairedVersion:
		// This will return the empty rev (empty string) if we don't have a
		// record of it. It's up to the caller to decide, for example, if
		// it's appropriate to update the cache.
		return dc.vMap[t]
	default:
		panic(fmt.Sprintf("Unknown version type %T", v))
	}
}

// toUnpaired turns a Version into an UnpairedVersion, if doing so is possible
// based on the information contained in the version itself, or in the cache
// maps.
//
// If the input is a revision and multiple UnpairedVersions are associated with
// it, whatever happens to be the first is returned.
func (dc *sourceMetaCache) toUnpaired(v Version) UnpairedVersion {
	switch t := v.(type) {
	case UnpairedVersion:
		return t
	case PairedVersion:
		return t.Unpair()
	case Revision:
		if upv, has := dc.rMap[t]; has && len(upv) > 0 {
			return upv[0]
		}
		return nil
	default:
		panic(fmt.Sprintf("unknown version type %T", v))
	}
}

func (bs *baseVCSSource) revisionPresentIn(r Revision) (bool, error) {
	// First and fastest path is to check the data cache to see if the rev is
	// present. This could give us false positives, but the cases where that can
	// occur would require a type of cache staleness that seems *exceedingly*
	// unlikely to occur.
	if _, has := bs.dc.infos[r]; has {
		return true, nil
	} else if _, has := bs.dc.rMap[r]; has {
		return true, nil
	}

	err := bs.ensureCacheExistence()
	if err != nil {
		return false, err
	}

	bs.crepo.mut.RLock()
	defer bs.crepo.mut.RUnlock()
	return bs.crepo.r.IsReference(string(r)), nil
}

func (bs *baseVCSSource) ensureCacheExistence() error {
	// Technically, methods could could attempt to return straight from the
	// metadata cache even if the repo cache doesn't exist on disk. But that
	// would allow weird state inconsistencies (cache exists, but no repo...how
	// does that even happen?) that it'd be better to just not allow so that we
	// don't have to think about it elsewhere
	if !bs.checkExistence(existsInCache) {
		if bs.checkExistence(existsUpstream) {
			bs.crepo.mut.Lock()
			if bs.crepo.synced {
				// A second ensure call coming in while the first is completing
				// isn't terribly unlikely, especially for a large repo. In that
				// event, the synced flag will have flipped on by the time we
				// acquire the lock. If it has, there's no need to do this work
				// twice.
				bs.crepo.mut.Unlock()
				return nil
			}

			err := bs.crepo.r.Get()

			if err != nil {
				bs.crepo.mut.Unlock()
				return fmt.Errorf("failed to create repository cache for %s with err:\n%s", bs.crepo.r.Remote(), unwrapVcsErr(err))
			}

			bs.crepo.synced = true
			bs.ex.s |= existsInCache
			bs.ex.f |= existsInCache
			bs.crepo.mut.Unlock()
		} else {
			return fmt.Errorf("project %s does not exist upstream", bs.crepo.r.Remote())
		}
	}

	return nil
}

// checkExistence provides a direct method for querying existence levels of the
// source. It will only perform actual searching (local fs or over the network)
// if no previous attempt at that search has been made.
//
// Note that this may perform read-ish operations on the cache repo, and it
// takes a lock accordingly. This makes it unsafe to call from a segment where
// the cache repo mutex is already write-locked, as deadlock will occur.
func (bs *baseVCSSource) checkExistence(ex sourceExistence) bool {
	if bs.ex.s&ex != ex {
		if ex&existsInVendorRoot != 0 && bs.ex.s&existsInVendorRoot == 0 {
			panic("should now be implemented in bridge")
		}
		if ex&existsInCache != 0 && bs.ex.s&existsInCache == 0 {
			bs.crepo.mut.RLock()
			bs.ex.s |= existsInCache
			if bs.crepo.r.CheckLocal() {
				bs.ex.f |= existsInCache
			}
			bs.crepo.mut.RUnlock()
		}
		if ex&existsUpstream != 0 && bs.ex.s&existsUpstream == 0 {
			bs.crepo.mut.RLock()
			bs.ex.s |= existsUpstream
			if bs.crepo.r.Ping() {
				bs.ex.f |= existsUpstream
			}
			bs.crepo.mut.RUnlock()
		}
	}

	return ex&bs.ex.f == ex
}

// syncLocal ensures the local data we have about the source is fully up to date
// with what's out there over the network.
func (bs *baseVCSSource) syncLocal() error {
	// Ensure we only have one goroutine doing this at a time
	bs.synclock.Lock()
	defer bs.synclock.Unlock()

	// ...and that we only ever do it once
	if bs.allsync {
		// Return the stored err, if any
		return bs.syncerr
	}

	bs.allsync = true
	// First, ensure the local instance exists
	bs.syncerr = bs.ensureCacheExistence()
	if bs.syncerr != nil {
		return bs.syncerr
	}

	_, bs.syncerr = bs.lvfunc()
	if bs.syncerr != nil {
		return bs.syncerr
	}

	// This case is really just for git repos, where the lvfunc doesn't
	// guarantee that the local repo is synced
	if !bs.crepo.synced {
		bs.crepo.mut.Lock()
		err := bs.crepo.r.Update()
		if err != nil {
			bs.syncerr = fmt.Errorf("failed fetching latest updates with err: %s", unwrapVcsErr(err))
			bs.crepo.mut.Unlock()
			return bs.syncerr
		}
		bs.crepo.synced = true
		bs.crepo.mut.Unlock()
	}

	return nil
}

func (bs *baseVCSSource) listPackages(pr ProjectRoot, v Version) (ptree PackageTree, err error) {
	if err = bs.ensureCacheExistence(); err != nil {
		return
	}

	var r Revision
	if r, err = bs.toRevOrErr(v); err != nil {
		return
	}

	// Return the ptree from the cache, if we already have it
	var exists bool
	if ptree, exists = bs.dc.ptrees[r]; exists {
		return
	}

	// Not in the cache; check out the version and do the analysis
	bs.crepo.mut.Lock()
	// Check out the desired version for analysis
	if r != "" {
		// Always prefer a rev, if it's available
		err = bs.crepo.r.UpdateVersion(string(r))
	} else {
		// If we don't have a rev, ensure the repo is up to date, otherwise we
		// could have a desync issue
		if !bs.crepo.synced {
			err = bs.crepo.r.Update()
			if err != nil {
				err = fmt.Errorf("could not fetch latest updates into repository: %s", unwrapVcsErr(err))
				return
			}
			bs.crepo.synced = true
		}
		err = bs.crepo.r.UpdateVersion(v.String())
	}

	if err == nil {
		ptree, err = ListPackages(bs.crepo.r.LocalPath(), string(pr))
		// TODO(sdboyer) cache errs?
		if err == nil {
			bs.dc.ptrees[r] = ptree
		}
	} else {
		err = unwrapVcsErr(err)
	}
	bs.crepo.mut.Unlock()

	return
}

// toRevOrErr makes all efforts to convert a Version into a rev, including
// updating the cache repo (if needed). It does not guarantee that the returned
// Revision actually exists in the repository (as one of the cheaper methods may
// have had bad data).
func (bs *baseVCSSource) toRevOrErr(v Version) (r Revision, err error) {
	r = bs.dc.toRevision(v)
	if r == "" {
		// Rev can be empty if:
		//  - The cache is unsynced
		//  - A version was passed that used to exist, but no longer does
		//  - A garbage version was passed. (Functionally indistinguishable from
		//  the previous)
		if !bs.cvsync {
			// call the lvfunc to sync the meta cache
			_, err = bs.lvfunc()
			if err != nil {
				return
			}
		}

		r = bs.dc.toRevision(v)
		// If we still don't have a rev, then the version's no good
		if r == "" {
			err = fmt.Errorf("version %s does not exist in source %s", v, bs.crepo.r.Remote())
		}
	}

	return
}

func (bs *baseVCSSource) exportVersionTo(v Version, to string) error {
	return bs.crepo.exportVersionTo(v, to)
}
<|endoftext|>


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

// Package base provides some helpers to make constructing bar modules easier.
package base

import (
	"os/exec"
	"time"

	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/outputs"
)

// Base is a simple module that satisfies the bar.Module interface, while adding
// some useful functions to make building modules on top somewhat simpler.
type Base struct {
	channel        chan bar.Output
	clickHandler   func(bar.Event)
	updateFunc     func()
	paused         bool
	updateOnResume bool
	outputOnResume bar.Output
	lastError      error
}

// Module implements bar's Module, Clickable, and Pausable,
// and adds a method to trigger updates on demand.
type Module interface {
	bar.Module
	bar.Clickable
	bar.Pausable
	Update()
}

// WithClickHandler extends Module to add support for a generic click handler
// that only receives the bar.Event. Most custom modules should extend Module
// with their own OnClick method that provides module-specific information.
type WithClickHandler interface {
	Module
	OnClick(func(bar.Event)) Module
}

// Stream starts up the worker goroutine, and channels its output to the bar.
func (b *Base) Stream() <-chan bar.Output {
	b.Resume()
	// Constructed when New is called, but is not directly exposed to extending
	// modules. Use Output or Clear to control the bar output.
	return b.channel
}

// Click handles click events from the bar.
// A middle click will always force an update, but if the module
// is currently in an error state, the configured click handler
// will be replaced by one that shows the error message using
// i3-nagbar on left click and updates the module on right click
func (b *Base) Click(e bar.Event) {
	if b.lastError == nil {
		if e.Button == bar.ButtonMiddle {
			b.Update()
		}
		if b.clickHandler != nil {
			b.clickHandler(e)
		}
		return
	}
	switch e.Button {
	case bar.ButtonRight:
	case bar.ButtonMiddle:
		b.lastError = nil
		b.Clear()
		b.Update()
	case bar.ButtonLeft:
		go exec.Command("i3-nagbar", "-m", b.lastError.Error()).Run()
	}
}

// Pause marks the module as paused, which suspends updates
// and outputs to the bar.
func (b *Base) Pause() {
	b.paused = true
}

// Resume continues normal updating of the module, and performs an
// immediate update if any updates occurred while the module was paused.
func (b *Base) Resume() {
	b.paused = false
	if b.outputOnResume != nil {
		b.Output(b.outputOnResume)
		b.outputOnResume = nil
	}
	if b.updateOnResume {
		b.updateOnResume = false
		b.Update()
	}
}

// Update marks the module as ready for an update.
// The actual update may not happen immediately, e.g. if the bar is hidden.
func (b *Base) Update() {
	if b.paused {
		b.updateOnResume = true
		return
	}
	go b.updateFunc()
}

// OnClick sets the click handler.
// This is a minimal default implementation; derived modules should implement an
// alternative OnClick method that exposes module-specific data to the handler function.
// Returns Module to allow bar.Add/bar.Run on the result.
func (b *Base) OnClick(handler func(bar.Event)) Module {
	b.clickHandler = handler
	return b
}

// New constructs a new base module.
func New() *Base {
	return &Base{
		channel:    make(chan bar.Output),
		updateFunc: func() {},
		// Modules start paused so that any modifications prior to Stream()
		// are not applied before the module has started.
		paused: true,
		// Trigger an initial update when Stream is first called.
		updateOnResume: true,
	}
}

// OnUpdate sets the function that will be called when the module needs
// to be update. That function can choose to call Output/Clear/Error to
// update the output, but is not required if no visual update is necessary.
// This method is only called while the bar is visible, to conserve resources
// when possible. For this reason, it is recommended that heavy update work,
// e.g. http requests, should happen here and not in an independent timer.
func (b *Base) OnUpdate(updateFunc func()) {
	b.updateFunc = updateFunc
}

// Clear hides the module from the bar.
func (b *Base) Clear() {
	b.Output(outputs.Empty())
}

// Output updates the module's output.
func (b *Base) Output(out bar.Output) {
	if b.paused {
		b.outputOnResume = out
		return
	}
	go func() { b.channel <- out }()
}

// Error shows an error on the bar.
// It shows an urgent "Error" on the bar (or the full text if it fits),
// and when clicked shows the full error using i3-nagbar.
func (b *Base) Error(err error) bool {
	if err == nil {
		return false
	}
	b.lastError = err
	b.Output(outputs.Error(err))
	return true
}

// Scheduler holds either a timer or a ticker that schedules
// the next update of the module. Since many modules need to
// schedule updates, adding support into base simplifies module
// code and works well with pause/resume.
type Scheduler struct {
	timer  *time.Timer
	ticker *time.Ticker
}

// Stop cancels the next scheduled update, used when the update
// frequency/schedule changes.
func (s Scheduler) Stop() {
	if s.timer != nil {
		s.timer.Stop()
	}
	if s.ticker != nil {
		s.ticker.Stop()
	}
}

// UpdateAt schedules the module for updating at a specific time.
func (b *Base) UpdateAt(when time.Time) Scheduler {
	return b.UpdateAfter(time.Until(when))
}

// UpdateAfter schedules the module for updating after a delay.
func (b *Base) UpdateAfter(delay time.Duration) Scheduler {
	return Scheduler{
		timer: time.AfterFunc(delay, b.Update),
	}
}

// UpdateEvery schedules the module for repeated updating at an interval.
func (b *Base) UpdateEvery(interval time.Duration) Scheduler {
	ticker := time.NewTicker(interval)
	go func() {
		for range ticker.C {
			b.Update()
		}
	}()
	return Scheduler{ticker: ticker}
}
<commit_msg>Add strong ordering for base module outputs<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package base provides some helpers to make constructing bar modules easier.
package base

import (
	"os/exec"
	"time"

	"github.com/soumya92/barista/bar"
	"github.com/soumya92/barista/outputs"
)

// Base is a simple module that satisfies the bar.Module interface, while adding
// some useful functions to make building modules on top somewhat simpler.
type Base struct {
	channel        chan bar.Output
	clickHandler   func(bar.Event)
	updateFunc     func()
	paused         bool
	updateOnResume bool
	outputOnResume bar.Output
	lastError      error
}

// Module implements bar's Module, Clickable, and Pausable,
// and adds a method to trigger updates on demand.
type Module interface {
	bar.Module
	bar.Clickable
	bar.Pausable
	Update()
}

// WithClickHandler extends Module to add support for a generic click handler
// that only receives the bar.Event. Most custom modules should extend Module
// with their own OnClick method that provides module-specific information.
type WithClickHandler interface {
	Module
	OnClick(func(bar.Event)) Module
}

// Stream starts up the worker goroutine, and channels its output to the bar.
func (b *Base) Stream() <-chan bar.Output {
	b.Resume()
	// Constructed when New is called, but is not directly exposed to extending
	// modules. Use Output or Clear to control the bar output.
	return b.channel
}

// Click handles click events from the bar.
// A middle click will always force an update, but if the module
// is currently in an error state, the configured click handler
// will be replaced by one that shows the error message using
// i3-nagbar on left click and updates the module on right click
func (b *Base) Click(e bar.Event) {
	if b.lastError == nil {
		if e.Button == bar.ButtonMiddle {
			b.Update()
		}
		if b.clickHandler != nil {
			b.clickHandler(e)
		}
		return
	}
	switch e.Button {
	case bar.ButtonRight:
	case bar.ButtonMiddle:
		b.lastError = nil
		b.Clear()
		b.Update()
	case bar.ButtonLeft:
		go exec.Command("i3-nagbar", "-m", b.lastError.Error()).Run()
	}
}

// Pause marks the module as paused, which suspends updates
// and outputs to the bar.
func (b *Base) Pause() {
	b.paused = true
}

// Resume continues normal updating of the module, and performs an
// immediate update if any updates occurred while the module was paused.
func (b *Base) Resume() {
	b.paused = false
	if b.outputOnResume != nil {
		b.Output(b.outputOnResume)
		b.outputOnResume = nil
	}
	if b.updateOnResume {
		b.updateOnResume = false
		b.Update()
	}
}

// Update marks the module as ready for an update.
// The actual update may not happen immediately, e.g. if the bar is hidden.
func (b *Base) Update() {
	if b.paused {
		b.updateOnResume = true
		return
	}
	go b.updateFunc()
}

// OnClick sets the click handler.
// This is a minimal default implementation; derived modules should implement an
// alternative OnClick method that exposes module-specific data to the handler function.
// Returns Module to allow bar.Add/bar.Run on the result.
func (b *Base) OnClick(handler func(bar.Event)) Module {
	b.clickHandler = handler
	return b
}

// New constructs a new base module.
func New() *Base {
	return &Base{
		channel:    make(chan bar.Output, 10),
		updateFunc: func() {},
		// Modules start paused so that any modifications prior to Stream()
		// are not applied before the module has started.
		paused: true,
		// Trigger an initial update when Stream is first called.
		updateOnResume: true,
	}
}

// OnUpdate sets the function that will be called when the module needs
// to be update. That function can choose to call Output/Clear/Error to
// update the output, but is not required if no visual update is necessary.
// This method is only called while the bar is visible, to conserve resources
// when possible. For this reason, it is recommended that heavy update work,
// e.g. http requests, should happen here and not in an independent timer.
func (b *Base) OnUpdate(updateFunc func()) {
	b.updateFunc = updateFunc
}

// Clear hides the module from the bar.
func (b *Base) Clear() {
	b.Output(outputs.Empty())
}

// Output updates the module's output.
func (b *Base) Output(out bar.Output) {
	if b.paused {
		b.outputOnResume = out
		return
	}
	b.channel <- out
}

// Error shows an error on the bar.
// It shows an urgent "Error" on the bar (or the full text if it fits),
// and when clicked shows the full error using i3-nagbar.
func (b *Base) Error(err error) bool {
	if err == nil {
		return false
	}
	b.lastError = err
	b.Output(outputs.Error(err))
	return true
}

// Scheduler holds either a timer or a ticker that schedules
// the next update of the module. Since many modules need to
// schedule updates, adding support into base simplifies module
// code and works well with pause/resume.
type Scheduler struct {
	timer  *time.Timer
	ticker *time.Ticker
}

// Stop cancels the next scheduled update, used when the update
// frequency/schedule changes.
func (s Scheduler) Stop() {
	if s.timer != nil {
		s.timer.Stop()
	}
	if s.ticker != nil {
		s.ticker.Stop()
	}
}

// UpdateAt schedules the module for updating at a specific time.
func (b *Base) UpdateAt(when time.Time) Scheduler {
	return b.UpdateAfter(time.Until(when))
}

// UpdateAfter schedules the module for updating after a delay.
func (b *Base) UpdateAfter(delay time.Duration) Scheduler {
	return Scheduler{
		timer: time.AfterFunc(delay, b.Update),
	}
}

// UpdateEvery schedules the module for repeated updating at an interval.
func (b *Base) UpdateEvery(interval time.Duration) Scheduler {
	ticker := time.NewTicker(interval)
	go func() {
		for range ticker.C {
			b.Update()
		}
	}()
	return Scheduler{ticker: ticker}
}
<|endoftext|>


<commit_before>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

*/
package captcha

import (
	"bytes"
	"strconv"
	"strings"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
)

type IDEncoder func(_ echo.Context, id string, from string) (string, error)
type IDDecoder func(_ echo.Context, id string, from string) (string, error)

var DefaultOptions = &Options{
	EnableImage:    true,
	EnableAudio:    true,
	EnableDownload: true,
	AudioLangs:     []string{`zh`, `ru`, `en`},
	Prefix:         `/captcha`,
	CookieName:     `captchaId`,
	HeaderName:     `X-Captcha-ID`,
	IDEncoder: func(_ echo.Context, id string, _ string) (string, error) {
		return id, nil
	},
	IDDecoder: func(_ echo.Context, id string, _ string) (string, error) {
		return id, nil
	},
}

type Options struct {
	EnableImage    bool
	EnableAudio    bool
	EnableDownload bool
	AudioLangs     []string
	Prefix         string
	CookieName     string
	HeaderName     string
	IDEncoder      IDEncoder
	IDDecoder      IDDecoder
}

func (o Options) Wrapper(e echo.RouteRegister) {
	if o.AudioLangs == nil || len(o.AudioLangs) == 0 {
		o.AudioLangs = DefaultOptions.AudioLangs
	}
	if len(o.Prefix) == 0 {
		o.Prefix = DefaultOptions.Prefix
	}
	o.Prefix = strings.TrimRight(o.Prefix, "/")
	e.Get(o.Prefix+"/*", Captcha(&o))
}

func Captcha(opts ...*Options) func(echo.Context) error {
	var o *Options
	if len(opts) > 0 {
		o = opts[0]
	}
	if o == nil {
		o = DefaultOptions
	}
	if len(o.CookieName) == 0 {
		o.CookieName = DefaultOptions.CookieName
	}
	if len(o.HeaderName) == 0 {
		o.HeaderName = DefaultOptions.HeaderName
	}
	if o.IDDecoder == nil {
		o.IDDecoder = DefaultOptions.IDDecoder
	}
	if o.IDEncoder == nil {
		o.IDEncoder = DefaultOptions.IDEncoder
	}
	return func(ctx echo.Context) (err error) {
		var id, ext string
		param := ctx.P(0)
		if p := strings.LastIndex(param, `.`); p > 0 {
			id = param[0:p]
			ext = param[p:]
		}
		if len(ext) == 0 || len(id) == 0 {
			return echo.ErrNotFound
		}
		id, err = o.IDDecoder(ctx, id, `path`)
		if err != nil {
			return err
		}
		w := ctx.Response()
		header := w.Header()
		ids := []string{id}
		var hasCookieValue, hasHeaderValue bool
		if len(o.CookieName) > 0 {
			idByCookie := ctx.GetCookie(o.CookieName)
			if len(idByCookie) > 0 {
				idByCookie, err = o.IDDecoder(ctx, idByCookie, `cookie`)
				if err != nil {
					header.Add(`X-Captcha-ID-Error`, `cookie: `+err.Error())
					ctx.SetCookie(o.CookieName, ``, -1)
				} else {
					ids = append(ids, idByCookie)
				}
				hasCookieValue = true
			}
		}
		if len(o.HeaderName) > 0 {
			idByHeader := ctx.Header(o.HeaderName)
			if len(idByHeader) > 0 {
				idByHeader, err = o.IDDecoder(ctx, idByHeader, `header`)
				if err != nil {
					header.Add(`X-Captcha-ID-Error`, `header: `+err.Error())
				} else {
					ids = append(ids, idByHeader)
				}
				hasHeaderValue = true
			}
		}
		if len(ctx.Query("reload")) > 0 {
			var ok bool
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				if captcha.Reload(id) {
					ok = true
					ids = []string{id}
					break
				}
			}
			if !ok && (hasCookieValue || hasHeaderValue) { // 旧的已经全部失效了，自动申请新ID
				id = captcha.New()
				ids = []string{id}
				if hasCookieValue {
					ctx.SetCookie(o.CookieName, id)
				}
				if hasHeaderValue {
					header.Set(o.HeaderName, id)
				}
			}
		}
		download := o.EnableDownload && ctx.Queryx("download").Bool()
		b := bytes.NewBuffer(nil)
		switch ext {
		case ".png":
			if !o.EnableImage {
				return echo.ErrNotFound
			}
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				err = captcha.WriteImage(b, id, captcha.StdWidth, captcha.StdHeight)
				if err == nil || err != captcha.ErrNotFound {
					break
				}
			}
			if err != nil {
				if err == captcha.ErrNotFound {
					return echo.ErrNotFound
				}
				return
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "image/png")
			}
		case ".wav":
			if !o.EnableAudio {
				return echo.ErrNotFound
			}
			lang := strings.ToLower(ctx.Query("lang"))
			supported := false
			for _, supportedLang := range o.AudioLangs {
				if supportedLang == lang {
					supported = true
					break
				}
			}
			if !supported && len(o.AudioLangs) > 0 {
				lang = o.AudioLangs[0]
			}
			var au *captcha.Audio
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				au, err = captcha.GetAudio(id, lang)
				if err == nil || err != captcha.ErrNotFound {
					break
				}
			}
			if err != nil {
				if err == captcha.ErrNotFound {
					return echo.ErrNotFound
				}
				return
			}
			length := strconv.Itoa(au.EncodedLen())
			_, err = au.WriteTo(b)
			if err != nil {
				return err
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "audio/x-wav")
			}
			header.Set("Content-Length", length)
		default:
			return nil
		}
		return ctx.Blob(b.Bytes())
	}
}
<commit_msg>update<commit_after>/*

   Copyright 2016 Wenhui Shen <www.webx.top>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

*/
package captcha

import (
	"bytes"
	"strconv"
	"strings"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
)

type IDEncoder func(_ echo.Context, id string, from string) (string, error)
type IDDecoder func(_ echo.Context, id string, from string) (string, error)

var DefaultOptions = &Options{
	EnableImage:    true,
	EnableAudio:    true,
	EnableDownload: true,
	AudioLangs:     []string{`zh`, `ru`, `en`},
	Prefix:         `/captcha`,
	CookieName:     `captchaId`,
	HeaderName:     `X-Captcha-ID`,
	IDEncoder: func(_ echo.Context, id string, _ string) (string, error) {
		return id, nil
	},
	IDDecoder: func(_ echo.Context, id string, _ string) (string, error) {
		return id, nil
	},
}

type Options struct {
	EnableImage    bool
	EnableAudio    bool
	EnableDownload bool
	AudioLangs     []string
	Prefix         string
	CookieName     string
	HeaderName     string
	IDEncoder      IDEncoder
	IDDecoder      IDDecoder
}

func (o Options) Wrapper(e echo.RouteRegister) {
	if o.AudioLangs == nil || len(o.AudioLangs) == 0 {
		o.AudioLangs = DefaultOptions.AudioLangs
	}
	if len(o.Prefix) == 0 {
		o.Prefix = DefaultOptions.Prefix
	}
	o.Prefix = strings.TrimRight(o.Prefix, "/")
	e.Get(o.Prefix+"/*", Captcha(&o))
}

func Captcha(opts ...*Options) func(echo.Context) error {
	var o *Options
	if len(opts) > 0 {
		o = opts[0]
	}
	if o == nil {
		o = DefaultOptions
	}
	if len(o.CookieName) == 0 {
		o.CookieName = DefaultOptions.CookieName
	}
	if len(o.HeaderName) == 0 {
		o.HeaderName = DefaultOptions.HeaderName
	}
	if o.IDDecoder == nil {
		o.IDDecoder = DefaultOptions.IDDecoder
	}
	if o.IDEncoder == nil {
		o.IDEncoder = DefaultOptions.IDEncoder
	}
	return func(ctx echo.Context) (err error) {
		var id, ext string
		param := ctx.P(0)
		if p := strings.LastIndex(param, `.`); p > 0 {
			id = param[0:p]
			ext = param[p:]
		}
		if len(ext) == 0 || len(id) == 0 {
			return echo.ErrNotFound
		}
		id, err = o.IDDecoder(ctx, id, `path`)
		if err != nil {
			return err
		}
		w := ctx.Response()
		header := w.Header()
		ids := []string{id}
		var hasCookieValue, hasHeaderValue bool
		if len(o.CookieName) > 0 {
			idByCookie := ctx.GetCookie(o.CookieName)
			if len(idByCookie) > 0 {
				idByCookie, err = o.IDDecoder(ctx, idByCookie, `cookie`)
				if err != nil {
					header.Add(`X-Captcha-ID-Error`, `cookie: `+err.Error())
					ctx.SetCookie(o.CookieName, ``, -1)
				} else {
					ids = append(ids, idByCookie)
				}
				hasCookieValue = true
			}
		}
		if len(o.HeaderName) > 0 {
			idByHeader := ctx.Header(o.HeaderName)
			if len(idByHeader) > 0 {
				idByHeader, err = o.IDDecoder(ctx, idByHeader, `header`)
				if err != nil {
					header.Add(`X-Captcha-ID-Error`, `header: `+err.Error())
				} else {
					ids = append(ids, idByHeader)
				}
				hasHeaderValue = true
			}
		}
		if ctx.Queryx("reload").Bool() {
			var ok bool
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				if captcha.Reload(id) {
					ok = true
					ids = []string{id}
					break
				}
			}
			if !ok && (hasCookieValue || hasHeaderValue) { // 旧的已经全部失效了，自动申请新ID
				id = captcha.New()
				ids = []string{id}
				if hasCookieValue {
					ctx.SetCookie(o.CookieName, id)
				}
				if hasHeaderValue {
					header.Set(o.HeaderName, id)
				}
			}
		}
		download := o.EnableDownload && ctx.Queryx("download").Bool()
		b := bytes.NewBuffer(nil)
		switch ext {
		case ".png":
			if !o.EnableImage {
				return echo.ErrNotFound
			}
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				err = captcha.WriteImage(b, id, captcha.StdWidth, captcha.StdHeight)
				if err == nil || err != captcha.ErrNotFound {
					break
				}
			}
			if err != nil {
				if err == captcha.ErrNotFound {
					return echo.ErrNotFound
				}
				return
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "image/png")
			}
		case ".wav":
			if !o.EnableAudio {
				return echo.ErrNotFound
			}
			lang := strings.ToLower(ctx.Query("lang"))
			supported := false
			for _, supportedLang := range o.AudioLangs {
				if supportedLang == lang {
					supported = true
					break
				}
			}
			if !supported && len(o.AudioLangs) > 0 {
				lang = o.AudioLangs[0]
			}
			var au *captcha.Audio
			for _, id := range ids {
				if len(id) == 0 {
					continue
				}
				au, err = captcha.GetAudio(id, lang)
				if err == nil || err != captcha.ErrNotFound {
					break
				}
			}
			if err != nil {
				if err == captcha.ErrNotFound {
					return echo.ErrNotFound
				}
				return
			}
			length := strconv.Itoa(au.EncodedLen())
			_, err = au.WriteTo(b)
			if err != nil {
				return err
			}
			if download {
				header.Set(echo.HeaderContentType, "application/octet-stream")
			} else {
				header.Set(echo.HeaderContentType, "audio/x-wav")
			}
			header.Set("Content-Length", length)
		default:
			return nil
		}
		return ctx.Blob(b.Bytes())
	}
}
<|endoftext|>