<commit_before>package handlers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/memcache"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/user"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type ArchiveData struct {
	Years   *map[string]Year
	Months  *[]string
	Posts   *[]models.Entry
	IsAdmin bool
}

type Year map[string]Month
type Month []Day
type Day int64

var Months = [12]time.Month{
	time.January,
	time.February,
	time.March,
	time.April,
	time.May,
	time.June,
	time.July,
	time.August,
	time.September,
	time.October,
	time.November,
	time.December,
}

func ArchiveQueueHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	t := taskqueue.NewPOSTTask("/archive/work", url.Values{})
	_, err := taskqueue.Add(c, t, "")

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		fmt.Fprint(w, "success.\n")
	}
}

func ArchiveTaskHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)

	entries, err := models.AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	log.Infof(c, "Retrieved data: %d.", len(*entries))

	years := make(map[string]Year)

	oldest := (*entries)[len(*entries)-1].Datetime
	newest := (*entries)[0].Datetime

	log.Infof(c, "Oldest: %v, Newest: %v", oldest, newest)

	for year := oldest.Year(); year <= newest.Year(); year += 1 {
		ystr := strconv.Itoa(year)
		years[ystr] = make(Year)
		log.Infof(c, "Adding %d.", year)
		for _, month := range Months {
			notFirstOrLastYear := year < newest.Year() || year > oldest.Year()
			notMonthsInFuture := year == newest.Year() && month <= newest.Month()
			notMonthsBeforeOldest := year == oldest.Year() && month >= oldest.Month()
			if notFirstOrLastYear || (notMonthsInFuture && notMonthsBeforeOldest) {
				mstr := month.String()
				years[ystr][mstr] = make([]Day, daysIn(month, year)+1)
				log.Debugf(c, "Adding %d/%d - %d days.", year, month, len(years[ystr][mstr]))

				for day := range years[ystr][mstr] {
					if day > 0 {
						e, err := models.PostsForDay(c, int64(year), int64(month), int64(day))
						if err != nil {
							log.Errorf(c, err.Error())
							http.Error(w, err.Error(), 500)
							return
						}
						years[ystr][mstr][day] = Day(len(*e))
					}
				}
			}
		}
	}
	log.Infof(c, "Added posts.")

	// https://blog.golang.org/json-and-go
	b, err := json.Marshal(years)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
	}

	item := &memcache.Item{
		Key:   "archive_data",
		Value: b,
	}

	// Set the item, unconditionally
	if err := memcache.Set(c, item); err != nil {
		log.Errorf(c, "error setting item: %v", err)
	}
}

func ArchiveHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)

	entries, err := models.AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	log.Infof(c, "Retrieved data: %d.", len(*entries))

	// Get the item from the memcache
	var years map[string]Year
	if year_data, err := memcache.Get(c, "archive_data"); err == memcache.ErrCacheMiss {
		log.Infof(c, "item not in the cache")
	} else if err != nil {
		log.Errorf(c, "error getting item: %v", err)
	} else {
		err := json.Unmarshal(year_data.Value, &years)
		if err != nil {
			log.Errorf(c, err.Error())
			http.Error(w, err.Error(), 500)
			return
		}
	}

	data := &ArchiveData{
		Years: &years,
		Months: &[]string{
			time.January.String(),
			time.February.String(),
			time.March.String(),
			time.April.String(),
			time.May.String(),
			time.June.String(),
			time.July.String(),
			time.August.String(),
			time.September.String(),
			time.October.String(),
			time.November.String(),
			time.December.String(),
		},
		IsAdmin: user.IsAdmin(c),
		Posts:   entries,
	}
	w.Render("archive", data)
}

// daysIn returns the number of days in a month for a given year.
func daysIn(m time.Month, year int) int {
	// This is equivalent to time.daysIn(m, year).
	return time.Date(year, m+1, 0, 0, 0, 0, 0, time.UTC).Day()
}
<commit_msg>ken's logic<commit_after>package handlers

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/memcache"
	"google.golang.org/appengine/taskqueue"
	"google.golang.org/appengine/user"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type ArchiveData struct {
	Years   *map[string]Year
	Months  *[]string
	Posts   *[]models.Entry
	IsAdmin bool
}

type Year map[string]Month
type Month []Day
type Day int64

var Months = [12]time.Month{
	time.January,
	time.February,
	time.March,
	time.April,
	time.May,
	time.June,
	time.July,
	time.August,
	time.September,
	time.October,
	time.November,
	time.December,
}

func ArchiveQueueHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	t := taskqueue.NewPOSTTask("/archive/work", url.Values{})
	_, err := taskqueue.Add(c, t, "")

	if err != nil {
		http.Error(w, err.Error(), 500)
	} else {
		fmt.Fprint(w, "success.\n")
	}
}

func ArchiveTaskHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)

	entries, err := models.AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	log.Infof(c, "Retrieved data: %d.", len(*entries))

	years := make(map[string]Year)

	oldest := (*entries)[len(*entries)-1].Datetime
	newest := (*entries)[0].Datetime

	log.Infof(c, "Oldest: %v, Newest: %v", oldest, newest)

	for year := oldest.Year(); year <= newest.Year(); year += 1 {
		ystr := strconv.Itoa(year)
		years[ystr] = make(Year)
		log.Infof(c, "Adding %d.", year)
		for _, month := range Months {
			notFirstOrLastYear := year < newest.Year() || year > oldest.Year()
			notMonthsInFuture := year != newest.Year() || month <= newest.Month()
			notMonthsBeforeOldest := year != oldest.Year() || month >= oldest.Month()
			if notFirstOrLastYear || (notMonthsInFuture && notMonthsBeforeOldest) {
				mstr := month.String()
				years[ystr][mstr] = make([]Day, daysIn(month, year)+1)
				log.Debugf(c, "Adding %d/%d - %d days.", year, month, len(years[ystr][mstr]))

				for day := range years[ystr][mstr] {
					if day > 0 {
						e, err := models.PostsForDay(c, int64(year), int64(month), int64(day))
						if err != nil {
							log.Errorf(c, err.Error())
							http.Error(w, err.Error(), 500)
							return
						}
						years[ystr][mstr][day] = Day(len(*e))
					}
				}
			}
		}
	}
	log.Infof(c, "Added posts.")

	// https://blog.golang.org/json-and-go
	b, err := json.Marshal(years)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
	}

	item := &memcache.Item{
		Key:   "archive_data",
		Value: b,
	}

	// Set the item, unconditionally
	if err := memcache.Set(c, item); err != nil {
		log.Errorf(c, "error setting item: %v", err)
	}
}

func ArchiveHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)

	entries, err := models.AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	log.Infof(c, "Retrieved data: %d.", len(*entries))

	// Get the item from the memcache
	var years map[string]Year
	if year_data, err := memcache.Get(c, "archive_data"); err == memcache.ErrCacheMiss {
		log.Infof(c, "item not in the cache")
	} else if err != nil {
		log.Errorf(c, "error getting item: %v", err)
	} else {
		err := json.Unmarshal(year_data.Value, &years)
		if err != nil {
			log.Errorf(c, err.Error())
			http.Error(w, err.Error(), 500)
			return
		}
	}

	data := &ArchiveData{
		Years: &years,
		Months: &[]string{
			time.January.String(),
			time.February.String(),
			time.March.String(),
			time.April.String(),
			time.May.String(),
			time.June.String(),
			time.July.String(),
			time.August.String(),
			time.September.String(),
			time.October.String(),
			time.November.String(),
			time.December.String(),
		},
		IsAdmin: user.IsAdmin(c),
		Posts:   entries,
	}
	w.Render("archive", data)
}

// daysIn returns the number of days in a month for a given year.
func daysIn(m time.Month, year int) int {
	// This is equivalent to time.daysIn(m, year).
	return time.Date(year, m+1, 0, 0, 0, 0, 0, time.UTC).Day()
}
<|endoftext|>


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

package mgt

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	"github.com/gocraft/work"
	"github.com/goharbor/harbor/src/jobservice/common/query"
	"github.com/goharbor/harbor/src/jobservice/common/rds"
	"github.com/goharbor/harbor/src/jobservice/common/utils"
	"github.com/goharbor/harbor/src/jobservice/errs"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/jobservice/period"
	"github.com/goharbor/harbor/src/lib/errors"
	"github.com/gomodule/redigo/redis"
)

// Manager defies the related operations to handle the management of job stats.
type Manager interface {
	// Get the stats data of all kinds of jobs.
	// Data returned by pagination.
	//
	// Arguments:
	//   q *query.Parameter : the query parameters
	//
	// Returns:
	//   The matched job stats list
	//   The total number of the jobs
	//   Non nil error if any issues meet
	GetJobs(q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the executions of the specified periodic job by pagination
	//
	// Arguments:
	//   pID: ID of the periodic job
	//   q *query.Parameter: query parameters
	//
	// Returns:
	//   The matched job stats list,
	//   The total number of the executions,
	//   Non nil error if any issues meet.
	GetPeriodicExecution(pID string, q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the scheduled jobs
	//
	// Arguments:
	//   q *query.Parameter: query parameters
	//
	// Returns:
	//   The matched job stats list,
	//   The total number of the executions,
	//   Non nil error if any issues meet.
	GetScheduledJobs(q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the stats of the specified job
	//
	// Arguments:
	//   jobID string: ID of the job
	//
	// Returns:
	//   The job stats
	//   Non nil error if any issues meet
	GetJob(jobID string) (*job.Stats, error)

	// Save the job stats
	//
	// Arguments:
	//   job *job.Stats: the saving job stats
	//
	// Returns:
	//   Non nil error if any issues meet
	SaveJob(job *job.Stats) error
}

// basicManager is the default implementation of @manager,
// based on redis.
type basicManager struct {
	// system context
	ctx context.Context
	// db namespace
	namespace string
	// redis conn pool
	pool *redis.Pool
	// go work client
	client *work.Client
}

// NewManager news a basic manager
func NewManager(ctx context.Context, ns string, pool *redis.Pool) Manager {
	return &basicManager{
		ctx:       ctx,
		namespace: ns,
		pool:      pool,
		client:    work.NewClient(ns, pool),
	}
}

// GetJobs is implementation of Manager.GetJobs
// Because of the hash set used to keep the job stats, we can not support a standard pagination.
// A cursor is used to fetch the jobs with several batches.
func (bm *basicManager) GetJobs(q *query.Parameter) ([]*job.Stats, int64, error) {
	cursor, count := int64(0), query.DefaultPageSize
	if q != nil {
		if q.PageSize > 0 {
			count = q.PageSize
		}

		if cur, ok := q.Extras.Get(query.ExtraParamKeyCursor); ok {
			cursor = cur.(int64)
		}
	}

	pattern := rds.KeyJobStats(bm.namespace, "*")
	args := []interface{}{cursor, "MATCH", pattern, "COUNT", count}

	conn := bm.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	values, err := redis.Values(conn.Do("SCAN", args...))
	if err != nil {
		return nil, 0, err
	}
	if len(values) != 2 {
		return nil, 0, errors.New("malform scan results")
	}

	nextCur, err := strconv.ParseUint(string(values[0].([]byte)), 10, 8)
	if err != nil {
		return nil, 0, err
	}
	list := values[1].([]interface{})

	results := make([]*job.Stats, 0)
	for _, v := range list {
		if bytes, ok := v.([]byte); ok {
			statsKey := string(bytes)
			if i := strings.LastIndex(statsKey, ":"); i != -1 {
				jID := statsKey[i+1:]
				t := job.NewBasicTrackerWithID(bm.ctx, jID, bm.namespace, bm.pool, nil, nil)
				if err := t.Load(); err != nil {
					logger.Errorf("retrieve stats data of job %s error: %s", jID, err)
					continue
				}

				results = append(results, t.Job())
			}
		}
	}

	return results, int64(nextCur), nil
}

// GetPeriodicExecution is implementation of Manager.GetPeriodicExecution
func (bm *basicManager) GetPeriodicExecution(pID string, q *query.Parameter) (results []*job.Stats, total int64, err error) {
	if utils.IsEmptyStr(pID) {
		return nil, 0, errors.New("nil periodic job ID")
	}

	tracker := job.NewBasicTrackerWithID(bm.ctx, pID, bm.namespace, bm.pool, nil, nil)
	err = tracker.Load()
	if err != nil {
		return nil, 0, err
	}

	if tracker.Job().Info.JobKind != job.KindPeriodic {
		return nil, 0, errors.Errorf("only periodic job has executions: %s kind is received", tracker.Job().Info.JobKind)
	}

	conn := bm.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	key := rds.KeyUpstreamJobAndExecutions(bm.namespace, pID)

	executionIDs := make([]string, 0)
	// Query executions by "non stopped"
	if nonStoppedOnly, ok := q.Extras.Get(query.ExtraParamKeyNonStoppedOnly); ok {
		if v, yes := nonStoppedOnly.(bool); yes && v {
			executionIDs, total, err = queryExecutions(conn, key, q)
			if err != nil {
				return nil, 0, err
			}
		}
	} else {
		// Query all
		// Pagination
		var pageNumber, pageSize uint = 1, query.DefaultPageSize
		if q != nil {
			if q.PageNumber > 0 {
				pageNumber = q.PageNumber
			}
			if q.PageSize > 0 {
				pageSize = q.PageSize
			}
		}

		// Get total first
		total, err = redis.Int64(conn.Do("ZCARD", key))
		if err != nil {
			return nil, 0, err
		}

		// No items
		if total == 0 || (int64)((pageNumber-1)*pageSize) >= total {
			return results, total, nil
		}

		min, max := (pageNumber-1)*pageSize, pageNumber*pageSize-1
		args := []interface{}{key, min, max}
		list, err := redis.Values(conn.Do("ZREVRANGE", args...))
		if err != nil {
			return nil, 0, err
		}

		for _, item := range list {
			if eID, ok := item.([]byte); ok {
				executionIDs = append(executionIDs, string(eID))
			}
		}
	}

	for _, eID := range executionIDs {
		t := job.NewBasicTrackerWithID(bm.ctx, eID, bm.namespace, bm.pool, nil, nil)
		if er := t.Load(); er != nil {
			logger.Errorf("track job %s error: %s", eID, err)
			continue
		}

		results = append(results, t.Job())
	}

	return
}

// GetScheduledJobs is implementation of Manager.GetScheduledJobs
func (bm *basicManager) GetScheduledJobs(q *query.Parameter) ([]*job.Stats, int64, error) {
	// PageSize is not supported here
	var page uint = 1
	if q != nil && q.PageNumber > 1 {
		page = q.PageNumber
	}

	sJobs, total, err := bm.client.ScheduledJobs(page)
	if err != nil {
		return nil, 0, err
	}

	res := make([]*job.Stats, 0)
	for _, sJob := range sJobs {
		jID := sJob.ID
		if len(sJob.Args) > 0 {
			if _, ok := sJob.Args[period.PeriodicExecutionMark]; ok {
				// Periodic scheduled job
				jID = fmt.Sprintf("%s@%d", sJob.ID, sJob.RunAt)
			}
		}
		t := job.NewBasicTrackerWithID(bm.ctx, jID, bm.namespace, bm.pool, nil, nil)
		err = t.Load()
		if err != nil {
			// Just log it
			logger.Errorf("mgt.basicManager: query scheduled jobs error: %s", err)
			continue
		}

		res = append(res, t.Job())
	}

	return res, total, nil
}

// GetJob is implementation of Manager.GetJob
func (bm *basicManager) GetJob(jobID string) (*job.Stats, error) {
	if utils.IsEmptyStr(jobID) {
		return nil, errs.BadRequestError("empty job ID")
	}

	t := job.NewBasicTrackerWithID(bm.ctx, jobID, bm.namespace, bm.pool, nil, nil)
	if err := t.Load(); err != nil {
		return nil, err
	}

	return t.Job(), nil
}

// SaveJob is implementation of Manager.SaveJob
func (bm *basicManager) SaveJob(j *job.Stats) error {
	if j == nil {
		return errs.BadRequestError("nil saving job stats")
	}

	t := job.NewBasicTrackerWithStats(bm.ctx, j, bm.namespace, bm.pool, nil, nil)
	return t.Save()
}

// queryExecutions queries periodic executions by status
func queryExecutions(conn redis.Conn, dataKey string, q *query.Parameter) ([]string, int64, error) {
	total, err := redis.Int64(conn.Do("ZCOUNT", dataKey, 0, "+inf"))
	if err != nil {
		return nil, 0, err
	}

	var pageNumber, pageSize uint = 1, query.DefaultPageSize
	if q.PageNumber > 0 {
		pageNumber = q.PageNumber
	}
	if q.PageSize > 0 {
		pageSize = q.PageSize
	}

	results := make([]string, 0)
	if total == 0 || (int64)((pageNumber-1)*pageSize) >= total {
		return results, total, nil
	}

	offset := (pageNumber - 1) * pageSize
	args := []interface{}{dataKey, "+inf", 0, "LIMIT", offset, pageSize}

	eIDs, err := redis.Values(conn.Do("ZREVRANGEBYSCORE", args...))
	if err != nil {
		return nil, 0, err
	}

	for _, eID := range eIDs {
		if eIDBytes, ok := eID.([]byte); ok {
			results = append(results, string(eIDBytes))
		}
	}

	return results, total, nil
}
<commit_msg>fix(jobservice): value out of range (#14634)<commit_after>// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mgt

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	"github.com/gocraft/work"
	"github.com/goharbor/harbor/src/jobservice/common/query"
	"github.com/goharbor/harbor/src/jobservice/common/rds"
	"github.com/goharbor/harbor/src/jobservice/common/utils"
	"github.com/goharbor/harbor/src/jobservice/errs"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/jobservice/period"
	"github.com/goharbor/harbor/src/lib/errors"
	"github.com/gomodule/redigo/redis"
)

// Manager defies the related operations to handle the management of job stats.
type Manager interface {
	// Get the stats data of all kinds of jobs.
	// Data returned by pagination.
	//
	// Arguments:
	//   q *query.Parameter : the query parameters
	//
	// Returns:
	//   The matched job stats list
	//   The total number of the jobs
	//   Non nil error if any issues meet
	GetJobs(q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the executions of the specified periodic job by pagination
	//
	// Arguments:
	//   pID: ID of the periodic job
	//   q *query.Parameter: query parameters
	//
	// Returns:
	//   The matched job stats list,
	//   The total number of the executions,
	//   Non nil error if any issues meet.
	GetPeriodicExecution(pID string, q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the scheduled jobs
	//
	// Arguments:
	//   q *query.Parameter: query parameters
	//
	// Returns:
	//   The matched job stats list,
	//   The total number of the executions,
	//   Non nil error if any issues meet.
	GetScheduledJobs(q *query.Parameter) ([]*job.Stats, int64, error)

	// Get the stats of the specified job
	//
	// Arguments:
	//   jobID string: ID of the job
	//
	// Returns:
	//   The job stats
	//   Non nil error if any issues meet
	GetJob(jobID string) (*job.Stats, error)

	// Save the job stats
	//
	// Arguments:
	//   job *job.Stats: the saving job stats
	//
	// Returns:
	//   Non nil error if any issues meet
	SaveJob(job *job.Stats) error
}

// basicManager is the default implementation of @manager,
// based on redis.
type basicManager struct {
	// system context
	ctx context.Context
	// db namespace
	namespace string
	// redis conn pool
	pool *redis.Pool
	// go work client
	client *work.Client
}

// NewManager news a basic manager
func NewManager(ctx context.Context, ns string, pool *redis.Pool) Manager {
	return &basicManager{
		ctx:       ctx,
		namespace: ns,
		pool:      pool,
		client:    work.NewClient(ns, pool),
	}
}

// GetJobs is implementation of Manager.GetJobs
// Because of the hash set used to keep the job stats, we can not support a standard pagination.
// A cursor is used to fetch the jobs with several batches.
func (bm *basicManager) GetJobs(q *query.Parameter) ([]*job.Stats, int64, error) {
	cursor, count := int64(0), query.DefaultPageSize
	if q != nil {
		if q.PageSize > 0 {
			count = q.PageSize
		}

		if cur, ok := q.Extras.Get(query.ExtraParamKeyCursor); ok {
			cursor = cur.(int64)
		}
	}

	pattern := rds.KeyJobStats(bm.namespace, "*")
	args := []interface{}{cursor, "MATCH", pattern, "COUNT", count}

	conn := bm.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	values, err := redis.Values(conn.Do("SCAN", args...))
	if err != nil {
		return nil, 0, err
	}
	if len(values) != 2 {
		return nil, 0, errors.New("malform scan results")
	}

	nextCur, err := strconv.ParseInt(string(values[0].([]byte)), 10, 64)
	if err != nil {
		return nil, 0, err
	}
	list := values[1].([]interface{})

	results := make([]*job.Stats, 0)
	for _, v := range list {
		if bytes, ok := v.([]byte); ok {
			statsKey := string(bytes)
			if i := strings.LastIndex(statsKey, ":"); i != -1 {
				jID := statsKey[i+1:]
				t := job.NewBasicTrackerWithID(bm.ctx, jID, bm.namespace, bm.pool, nil, nil)
				if err := t.Load(); err != nil {
					logger.Errorf("retrieve stats data of job %s error: %s", jID, err)
					continue
				}

				results = append(results, t.Job())
			}
		}
	}

	return results, nextCur, nil
}

// GetPeriodicExecution is implementation of Manager.GetPeriodicExecution
func (bm *basicManager) GetPeriodicExecution(pID string, q *query.Parameter) (results []*job.Stats, total int64, err error) {
	if utils.IsEmptyStr(pID) {
		return nil, 0, errors.New("nil periodic job ID")
	}

	tracker := job.NewBasicTrackerWithID(bm.ctx, pID, bm.namespace, bm.pool, nil, nil)
	err = tracker.Load()
	if err != nil {
		return nil, 0, err
	}

	if tracker.Job().Info.JobKind != job.KindPeriodic {
		return nil, 0, errors.Errorf("only periodic job has executions: %s kind is received", tracker.Job().Info.JobKind)
	}

	conn := bm.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	key := rds.KeyUpstreamJobAndExecutions(bm.namespace, pID)

	executionIDs := make([]string, 0)
	// Query executions by "non stopped"
	if nonStoppedOnly, ok := q.Extras.Get(query.ExtraParamKeyNonStoppedOnly); ok {
		if v, yes := nonStoppedOnly.(bool); yes && v {
			executionIDs, total, err = queryExecutions(conn, key, q)
			if err != nil {
				return nil, 0, err
			}
		}
	} else {
		// Query all
		// Pagination
		var pageNumber, pageSize uint = 1, query.DefaultPageSize
		if q != nil {
			if q.PageNumber > 0 {
				pageNumber = q.PageNumber
			}
			if q.PageSize > 0 {
				pageSize = q.PageSize
			}
		}

		// Get total first
		total, err = redis.Int64(conn.Do("ZCARD", key))
		if err != nil {
			return nil, 0, err
		}

		// No items
		if total == 0 || (int64)((pageNumber-1)*pageSize) >= total {
			return results, total, nil
		}

		min, max := (pageNumber-1)*pageSize, pageNumber*pageSize-1
		args := []interface{}{key, min, max}
		list, err := redis.Values(conn.Do("ZREVRANGE", args...))
		if err != nil {
			return nil, 0, err
		}

		for _, item := range list {
			if eID, ok := item.([]byte); ok {
				executionIDs = append(executionIDs, string(eID))
			}
		}
	}

	for _, eID := range executionIDs {
		t := job.NewBasicTrackerWithID(bm.ctx, eID, bm.namespace, bm.pool, nil, nil)
		if er := t.Load(); er != nil {
			logger.Errorf("track job %s error: %s", eID, err)
			continue
		}

		results = append(results, t.Job())
	}

	return
}

// GetScheduledJobs is implementation of Manager.GetScheduledJobs
func (bm *basicManager) GetScheduledJobs(q *query.Parameter) ([]*job.Stats, int64, error) {
	// PageSize is not supported here
	var page uint = 1
	if q != nil && q.PageNumber > 1 {
		page = q.PageNumber
	}

	sJobs, total, err := bm.client.ScheduledJobs(page)
	if err != nil {
		return nil, 0, err
	}

	res := make([]*job.Stats, 0)
	for _, sJob := range sJobs {
		jID := sJob.ID
		if len(sJob.Args) > 0 {
			if _, ok := sJob.Args[period.PeriodicExecutionMark]; ok {
				// Periodic scheduled job
				jID = fmt.Sprintf("%s@%d", sJob.ID, sJob.RunAt)
			}
		}
		t := job.NewBasicTrackerWithID(bm.ctx, jID, bm.namespace, bm.pool, nil, nil)
		err = t.Load()
		if err != nil {
			// Just log it
			logger.Errorf("mgt.basicManager: query scheduled jobs error: %s", err)
			continue
		}

		res = append(res, t.Job())
	}

	return res, total, nil
}

// GetJob is implementation of Manager.GetJob
func (bm *basicManager) GetJob(jobID string) (*job.Stats, error) {
	if utils.IsEmptyStr(jobID) {
		return nil, errs.BadRequestError("empty job ID")
	}

	t := job.NewBasicTrackerWithID(bm.ctx, jobID, bm.namespace, bm.pool, nil, nil)
	if err := t.Load(); err != nil {
		return nil, err
	}

	return t.Job(), nil
}

// SaveJob is implementation of Manager.SaveJob
func (bm *basicManager) SaveJob(j *job.Stats) error {
	if j == nil {
		return errs.BadRequestError("nil saving job stats")
	}

	t := job.NewBasicTrackerWithStats(bm.ctx, j, bm.namespace, bm.pool, nil, nil)
	return t.Save()
}

// queryExecutions queries periodic executions by status
func queryExecutions(conn redis.Conn, dataKey string, q *query.Parameter) ([]string, int64, error) {
	total, err := redis.Int64(conn.Do("ZCOUNT", dataKey, 0, "+inf"))
	if err != nil {
		return nil, 0, err
	}

	var pageNumber, pageSize uint = 1, query.DefaultPageSize
	if q.PageNumber > 0 {
		pageNumber = q.PageNumber
	}
	if q.PageSize > 0 {
		pageSize = q.PageSize
	}

	results := make([]string, 0)
	if total == 0 || (int64)((pageNumber-1)*pageSize) >= total {
		return results, total, nil
	}

	offset := (pageNumber - 1) * pageSize
	args := []interface{}{dataKey, "+inf", 0, "LIMIT", offset, pageSize}

	eIDs, err := redis.Values(conn.Do("ZREVRANGEBYSCORE", args...))
	if err != nil {
		return nil, 0, err
	}

	for _, eID := range eIDs {
		if eIDBytes, ok := eID.([]byte); ok {
			results = append(results, string(eIDBytes))
		}
	}

	return results, total, nil
}
<|endoftext|>


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

// The iterable package provides several traversal and searching methods.
// It can be used on anything that satisfies the Iterable interface,
// including vector, though certain functions, such as Map, can also be used on
// something that would produce an infinite amount of data.
package iterable

import "container/vector"

type Iterable interface {
	// Iter should return a fresh channel each time it is called.
	Iter() <-chan interface {}
}

func not(f func(interface {}) bool) (func(interface {}) bool) {
	return func(e interface {}) bool { return !f(e) }
}

// All tests whether f is true for every element of iter.
func All(iter Iterable, f func(interface {}) bool) bool {
	for e := range iter.Iter() {
		if !f(e) {
			return false
		}
	}
	return true
}

// Any tests whether f is true for at least one element of iter.
func Any(iter Iterable, f func(interface {}) bool) bool {
	return !All(iter, not(f))
}

// Data returns a slice containing the elements of iter.
func Data(iter Iterable) []interface {} {
	vec := vector.New(0);
	for e := range iter.Iter() {
		vec.Push(e)
	}
	return vec.Data()
}

// filteredIterable is a struct that implements Iterable with each element
// passed through a filter.
type filteredIterable struct {
	it Iterable;
	f func(interface {}) bool;
}

func (f *filteredIterable) iterate(out chan<- interface {}) {
	for e := range f.it.Iter() {
		if f.f(e) {
			out <- e
		}
	}
	close(out)
}

func (f *filteredIterable) Iter() <-chan interface {} {
	ch := make(chan interface {});
	go f.iterate(ch);
	return ch;
}

// Filter returns an Iterable that returns the elements of iter that satisfy f.
func Filter(iter Iterable, f func(interface {}) bool) Iterable {
	return &filteredIterable{ iter, f }
}

// Find returns the first element of iter that satisfies f.
// Returns nil if no such element is found.
func Find(iter Iterable, f func(interface {}) bool) interface {} {
	for e := range Filter(iter, f).Iter() {
		return e
	}
	return nil
}

// An injector function takes two arguments, an accumulated value and an
// element, and returns the next accumulated value. See the Inject function.
type Injector func(interface {}, interface {}) interface{};

// Inject combines the elements of iter by repeatedly calling f with an
// accumulated value and each element in order. The starting accumulated value
// is initial, and after each call the accumulated value is set to the return
// value of f. For instance, to compute a sum:
//   var arr IntArray = []int{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
//   sum := iterable.Inject(arr, 0,
//                          func(ax interface {}, x interface {}) interface {} {
//                            return ax.(int) + x.(int) }).(int)
func Inject(iter Iterable, initial interface {}, f Injector) interface {} {
	acc := initial;
	for e := range iter.Iter() {
		acc = f(acc, e)
	}
	return acc
}

// mappedIterable is a helper struct that implements Iterable, returned by Map.
type mappedIterable struct {
	it Iterable;
	f func(interface {}) interface {};
}

func (m *mappedIterable) iterate(out chan<- interface {}) {
	for e := range m.it.Iter() {
		out <- m.f(e)
	}
	close(out)
}

func (m *mappedIterable) Iter() <-chan interface {} {
	ch := make(chan interface {});
	go m.iterate(ch);
	return ch
}

// Map returns an Iterable that returns the result of applying f to each
// element of iter.
func Map(iter Iterable, f func(interface {}) interface {}) Iterable {
	return &mappedIterable{ iter, f }
}

// Partition(iter, f) returns Filter(iter, f) and Filter(iter, !f).
func Partition(iter Iterable, f func(interface {}) bool) (Iterable, Iterable)  {
	return Filter(iter, f), Filter(iter, not(f))
}

// TODO:
// - Zip
<commit_msg>Oops, forgot to commit this change.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The iterable package provides several traversal and searching methods.
// It can be used on anything that satisfies the Iterable interface,
// including vector, though certain functions, such as Map, can also be used on
// something that would produce an infinite amount of data.
package iterable

import "container/vector"

type Iterable interface {
	// Iter should return a fresh channel each time it is called.
	Iter() <-chan interface {}
}

func not(f func(interface {}) bool) (func(interface {}) bool) {
	return func(e interface {}) bool { return !f(e) }
}

// All tests whether f is true for every element of iter.
func All(iter Iterable, f func(interface {}) bool) bool {
	for e := range iter.Iter() {
		if !f(e) {
			return false
		}
	}
	return true
}

// Any tests whether f is true for at least one element of iter.
func Any(iter Iterable, f func(interface {}) bool) bool {
	return !All(iter, not(f))
}

// Data returns a slice containing the elements of iter.
func Data(iter Iterable) []interface {} {
	vec := vector.New(0);
	for e := range iter.Iter() {
		vec.Push(e)
	}
	return vec.Data()
}

// filteredIterable is a struct that implements Iterable with each element
// passed through a filter.
type filteredIterable struct {
	it Iterable;
	f func(interface {}) bool;
}

func (f *filteredIterable) iterate(out chan<- interface {}) {
	for e := range f.it.Iter() {
		if f.f(e) {
			out <- e
		}
	}
	close(out)
}

func (f *filteredIterable) Iter() <-chan interface {} {
	ch := make(chan interface {});
	go f.iterate(ch);
	return ch;
}

// Filter returns an Iterable that returns the elements of iter that satisfy f.
func Filter(iter Iterable, f func(interface {}) bool) Iterable {
	return &filteredIterable{ iter, f }
}

// Find returns the first element of iter that satisfies f.
// Returns nil if no such element is found.
func Find(iter Iterable, f func(interface {}) bool) interface {} {
	for e := range Filter(iter, f).Iter() {
		return e
	}
	return nil
}

// Injector is a type representing a function that takes two arguments,
// an accumulated value and an element, and returns the next accumulated value.
// See the Inject function.
type Injector func(interface {}, interface {}) interface{};

// Inject combines the elements of iter by repeatedly calling f with an
// accumulated value and each element in order. The starting accumulated value
// is initial, and after each call the accumulated value is set to the return
// value of f. For instance, to compute a sum:
//   var arr IntArray = []int{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
//   sum := iterable.Inject(arr, 0,
//                          func(ax interface {}, x interface {}) interface {} {
//                            return ax.(int) + x.(int) }).(int)
func Inject(iter Iterable, initial interface {}, f Injector) interface {} {
	acc := initial;
	for e := range iter.Iter() {
		acc = f(acc, e)
	}
	return acc
}

// mappedIterable is a helper struct that implements Iterable, returned by Map.
type mappedIterable struct {
	it Iterable;
	f func(interface {}) interface {};
}

func (m *mappedIterable) iterate(out chan<- interface {}) {
	for e := range m.it.Iter() {
		out <- m.f(e)
	}
	close(out)
}

func (m *mappedIterable) Iter() <-chan interface {} {
	ch := make(chan interface {});
	go m.iterate(ch);
	return ch
}

// Map returns an Iterable that returns the result of applying f to each
// element of iter.
func Map(iter Iterable, f func(interface {}) interface {}) Iterable {
	return &mappedIterable{ iter, f }
}

// Partition(iter, f) returns Filter(iter, f) and Filter(iter, !f).
func Partition(iter Iterable, f func(interface {}) bool) (Iterable, Iterable)  {
	return Filter(iter, f), Filter(iter, not(f))
}

// TODO:
// - Zip
<|endoftext|>


<commit_before>package dnssd

import (
	"context"
	"fmt"
	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"math/rand"
	"net"
	"strings"
	"time"
)

// ProbeService probes for the hostname and service instance name of srv.
// If err == nil, the returned service is verified to be unique on the local network.
func ProbeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()

	// After one minute of probing, if the Multicast DNS responder has been
	// unable to find any unused name, it should log an error (RFC6762 9)
	probeCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	// When ready to send its Multicast DNS probe packet(s) the host should
	// first wait for a short random delay time, uniformly distributed in
	// the range 0-250 ms. (RFC6762 8.1)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	delay := time.Duration(r.Intn(250)) * time.Millisecond
	log.Debug.Println("Probing delay", delay)
	time.Sleep(delay)

	return probeService(probeCtx, conn, srv, 1*time.Millisecond, false)
}

func ReprobeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()
	return probeService(ctx, conn, srv, 1*time.Millisecond, true)
}

func probeService(ctx context.Context, conn MDNSConn, srv Service, delay time.Duration, probeOnce bool) (s Service, e error) {
	candidate := srv.Copy()
	prevConflict := probeConflict{}

	// Keep track of the number of conflicts
	numHostConflicts := 0
	numNameConflicts := 0

	for i := 1; i <= 100; i++ {
		conflict, err := probe(ctx, conn, *candidate)
		if err != nil {
			e = err
			return
		}

		if conflict.hasNone() {
			s = *candidate
			return
		}

		candidate = candidate.Copy()

		if conflict.hostname && (prevConflict.hostname || probeOnce) {
			numHostConflicts++
			candidate.Host = fmt.Sprintf("%s-%d", srv.Host, numHostConflicts+1)
			conflict.hostname = false
		}

		if conflict.serviceName && (prevConflict.serviceName || probeOnce) {
			numNameConflicts++
			candidate.Name = fmt.Sprintf("%s-%d", srv.Name, numNameConflicts+1)
			conflict.serviceName = false
		}

		prevConflict = conflict

		if conflict.hasAny() {
			// If the host finds that its own data is lexicographically earlier,
			// then it defers to the winning host by waiting one second,
			// and then begins probing for this record again. (RFC6762 8.2)
			log.Debug.Println("Increase wait time after receiving conflicting data")
			delay = 1 * time.Second
		} else {
			delay = 250 * time.Millisecond
		}

		log.Debug.Println("Probing wait", delay)
		time.Sleep(delay)
	}

	return
}

func probe(ctx context.Context, conn MDNSConn, service Service) (conflict probeConflict, err error) {
	for _, iface := range service.Interfaces() {
		log.Debug.Printf("Probing at %s\n", iface.Name)
		conflict, err := probeAtInterface(ctx, conn, service, iface)
		if conflict.hasAny() {
			return conflict, err
		}
	}

	return probeConflict{}, nil
}

func probeAtInterface(ctx context.Context, conn MDNSConn, service Service, iface net.Interface) (conflict probeConflict, err error) {

	msg := new(dns.Msg)

	instanceQ := dns.Question{
		Name:   service.ServiceInstanceName(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	hostQ := dns.Question{
		Name:   service.Hostname(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	// TODO Responses to probe should be unicast
	// setQuestionUnicast(&instanceQ)
	// setQuestionUnicast(&hostQ)

	msg.Question = []dns.Question{instanceQ, hostQ}

	srv := SRV(service)
	as := A(service, iface)
	aaaas := AAAA(service, iface)

	var authority = []dns.RR{srv}
	for _, a := range as {
		authority = append(authority, a)
	}
	for _, aaaa := range aaaas {
		authority = append(authority, aaaa)
	}
	msg.Ns = authority

	readCtx, readCancel := context.WithCancel(ctx)
	defer readCancel()

	// Multicast DNS responses received *before* the first probe packet is sent
	// MUST be silently ignored. (RFC6762 8.1)
	conn.Drain(readCtx)
	ch := conn.Read(readCtx)

	queryTime := time.After(1 * time.Millisecond)
	queriesCount := 1

	for {
		select {
		case req := <-ch:
			answers := allRecords(req.msg)
			for _, answer := range answers {
				switch rr := answer.(type) {
				case *dns.A:
					for _, a := range as {
						if isDenyingA(rr, a) {
							log.Debug.Printf("%v:%d@%s denies A\n", req.from.IP, req.from.Port, req.IfaceName())
							conflict.hostname = true
							break
						}
					}

				case *dns.AAAA:
					for _, aaaa := range aaaas {
						if isDenyingAAAA(rr, aaaa) {
							log.Debug.Printf("%v:%d@%s denies AAAA\n", req.from.IP, req.from.Port, req.IfaceName())
							conflict.hostname = true
							break
						}
					}

				case *dns.SRV:
					if isDenyingSRV(rr, srv) {
						conflict.serviceName = true
					}

				default:
					break
				}
			}

		case <-ctx.Done():
			err = ctx.Err()
			return

		case <-queryTime:
			// Stop on conflict
			if conflict.hasAny() {
				return
			}

			// Stop after 3 probe queries
			if queriesCount > 3 {
				return
			}

			queriesCount++
			log.Debug.Println("Sending probe", msg)
			q := &Query{msg: msg, iface: &iface}
			conn.SendQuery(q)

			delay := 250 * time.Millisecond
			log.Debug.Println("Waiting for conflicting data", delay)
			queryTime = time.After(delay)
		}
	}

	return
}

type probeConflict struct {
	hostname    bool
	serviceName bool
}

func (pr probeConflict) hasNone() bool {
	return !pr.hostname && !pr.serviceName
}

func (pr probeConflict) hasAny() bool {
	return pr.hostname || pr.serviceName
}

func isDenyingA(this *dns.A, that *dns.A) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same hosts")

		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.A.To4(), that.A.To4()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

// isDenyingAAAA returns true if this denies that.
func isDenyingAAAA(this *dns.AAAA, that *dns.AAAA) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.AAAA.To16(), that.AAAA.To16()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

func areDenyingAs(this []*dns.A, that []*dns.A) bool {
	if len(this) != len(that) {
		// different number of As are a conflict
		return true
	}

	if equalAs(this, that) {
		return false // equal As are no conlict
	}

	if len(this) == 1 && !isDenyingA(this[0], that[0]) {
		// this doesn't deny that but the other way around
		return false
	}

	return true
}

func areDenyingAAAAs(this []*dns.AAAA, that []*dns.AAAA) bool {
	if len(this) != len(that) {
		// different number of AAAAs are a conflict
		return true
	}

	if equalAAAAs(this, that) {
		return false // equal AAAAs are no conlict
	}

	if len(this) == 1 && !isDenyingAAAA(this[0], that[0]) {
		// this doesn't deny that but the other way around
		return false
	}

	return true
}

func equalAs(this []*dns.A, that []*dns.A) bool {
	var tmp = that
	for _, ti := range this {
		var found = false
		for i, ta := range tmp {
			if compareIP(ti.A.To4(), ta.A.To4()) == 0 {
				tmp = append(tmp[:i], tmp[i+1:]...)
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return len(tmp) == 0
}

func equalAAAAs(this []*dns.AAAA, that []*dns.AAAA) bool {
	var tmp = that
	for _, ti := range this {
		var found = false
		for i, ta := range tmp {
			if compareIP(ti.AAAA.To16(), ta.AAAA.To16()) == 0 {
				tmp = append(tmp[:i], tmp[i+1:]...)
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return len(tmp) == 0
}

func containedInAs(this *dns.A, aaas []*dns.A) bool {
	for _, that := range aaas {
		if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
			if !isValidRR(this) {
				log.Debug.Println("Invalid record produces conflict")
				return false
			}

			if compareIP(this.A.To4(), that.A.To4()) == 0 {
				return true
			}
		}
	}

	return false
}

func containedInAAAAs(this *dns.AAAA, aaas []*dns.AAAA) bool {
	for _, that := range aaas {
		if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
			if !isValidRR(this) {
				log.Debug.Println("Invalid record produces conflict")
				return false
			}

			if compareIP(this.AAAA.To16(), that.AAAA.To16()) == 0 {
				return true
			}
		}
	}

	return false
}

// isDenyingSRV returns true if this denies that.
func isDenyingSRV(this *dns.SRV, that *dns.SRV) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same SRV")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareSRV(this, that) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

func isValidRR(rr dns.RR) bool {
	switch r := rr.(type) {
	case *dns.A:
		return !net.IPv4zero.Equal(r.A)
	case *dns.AAAA:
		return !net.IPv6zero.Equal(r.AAAA)
	case *dns.SRV:
		return len(r.Target) > 0 && r.Port != 0
	default:
		break
	}

	return true
}

func compareIP(this net.IP, that net.IP) int {
	count := len(this)
	if count > len(that) {
		count = len(that)
	}

	for i := 0; i < count; i++ {
		if this[i] < that[i] {
			return -1
		} else if this[i] > that[i] {
			return 1
		}
	}

	if len(this) < len(that) {
		return -1
	} else if len(this) > len(that) {
		return 1
	}
	return 0
}

func compareSRV(this *dns.SRV, that *dns.SRV) int {
	if this.Priority < that.Priority {
		return -1
	} else if this.Priority > that.Priority {
		return 1
	}

	if this.Weight < that.Weight {
		return -1
	} else if this.Weight > that.Weight {
		return 1
	}

	if this.Port < that.Port {
		return -1
	} else if this.Port > that.Port {
		return 1
	}

	return strings.Compare(this.Target, that.Target)
}
<commit_msg>Filter responses<commit_after>package dnssd

import (
	"context"
	"fmt"
	"github.com/brutella/dnssd/log"
	"github.com/miekg/dns"
	"math/rand"
	"net"
	"strings"
	"time"
)

// ProbeService probes for the hostname and service instance name of srv.
// If err == nil, the returned service is verified to be unique on the local network.
func ProbeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()

	// After one minute of probing, if the Multicast DNS responder has been
	// unable to find any unused name, it should log an error (RFC6762 9)
	probeCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	// When ready to send its Multicast DNS probe packet(s) the host should
	// first wait for a short random delay time, uniformly distributed in
	// the range 0-250 ms. (RFC6762 8.1)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	delay := time.Duration(r.Intn(250)) * time.Millisecond
	log.Debug.Println("Probing delay", delay)
	time.Sleep(delay)

	return probeService(probeCtx, conn, srv, 1*time.Millisecond, false)
}

func ReprobeService(ctx context.Context, srv Service) (Service, error) {
	conn, err := newMDNSConn()

	if err != nil {
		return srv, err
	}

	defer conn.close()
	return probeService(ctx, conn, srv, 1*time.Millisecond, true)
}

func probeService(ctx context.Context, conn MDNSConn, srv Service, delay time.Duration, probeOnce bool) (s Service, e error) {
	candidate := srv.Copy()
	prevConflict := probeConflict{}

	// Keep track of the number of conflicts
	numHostConflicts := 0
	numNameConflicts := 0

	for i := 1; i <= 100; i++ {
		conflict, err := probe(ctx, conn, *candidate)
		if err != nil {
			e = err
			return
		}

		if conflict.hasNone() {
			s = *candidate
			return
		}

		candidate = candidate.Copy()

		if conflict.hostname && (prevConflict.hostname || probeOnce) {
			numHostConflicts++
			candidate.Host = fmt.Sprintf("%s-%d", srv.Host, numHostConflicts+1)
			conflict.hostname = false
		}

		if conflict.serviceName && (prevConflict.serviceName || probeOnce) {
			numNameConflicts++
			candidate.Name = fmt.Sprintf("%s-%d", srv.Name, numNameConflicts+1)
			conflict.serviceName = false
		}

		prevConflict = conflict

		if conflict.hasAny() {
			// If the host finds that its own data is lexicographically earlier,
			// then it defers to the winning host by waiting one second,
			// and then begins probing for this record again. (RFC6762 8.2)
			log.Debug.Println("Increase wait time after receiving conflicting data")
			delay = 1 * time.Second
		} else {
			delay = 250 * time.Millisecond
		}

		log.Debug.Println("Probing wait", delay)
		time.Sleep(delay)
	}

	return
}

func probe(ctx context.Context, conn MDNSConn, service Service) (conflict probeConflict, err error) {
	for _, iface := range service.Interfaces() {
		log.Debug.Printf("Probing at %s\n", iface.Name)
		conflict, err := probeAtInterface(ctx, conn, service, iface)
		if conflict.hasAny() {
			return conflict, err
		}
	}

	return probeConflict{}, nil
}

func probeAtInterface(ctx context.Context, conn MDNSConn, service Service, iface net.Interface) (conflict probeConflict, err error) {

	msg := new(dns.Msg)

	instanceQ := dns.Question{
		Name:   service.ServiceInstanceName(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	hostQ := dns.Question{
		Name:   service.Hostname(),
		Qtype:  dns.TypeANY,
		Qclass: dns.ClassINET,
	}

	// TODO Responses to probe should be unicast
	// setQuestionUnicast(&instanceQ)
	// setQuestionUnicast(&hostQ)

	msg.Question = []dns.Question{instanceQ, hostQ}

	srv := SRV(service)
	as := A(service, iface)
	aaaas := AAAA(service, iface)

	var authority = []dns.RR{srv}
	for _, a := range as {
		authority = append(authority, a)
	}
	for _, aaaa := range aaaas {
		authority = append(authority, aaaa)
	}
	msg.Ns = authority

	readCtx, readCancel := context.WithCancel(ctx)
	defer readCancel()

	// Multicast DNS responses received *before* the first probe packet is sent
	// MUST be silently ignored. (RFC6762 8.1)
	conn.Drain(readCtx)
	ch := conn.Read(readCtx)

	queryTime := time.After(1 * time.Millisecond)
	queriesCount := 1

	for {
		select {
		case req := <-ch:
			answers := filterRecords(req.msg, &service)
			for _, answer := range answers {
				switch rr := answer.(type) {
				case *dns.A:
					for _, a := range as {
						if isDenyingA(rr, a) {
							log.Debug.Printf("%v:%d@%s denies A\n", req.from.IP, req.from.Port, req.IfaceName())
							conflict.hostname = true
							break
						}
					}

				case *dns.AAAA:
					for _, aaaa := range aaaas {
						if isDenyingAAAA(rr, aaaa) {
							log.Debug.Printf("%v:%d@%s denies AAAA\n", req.from.IP, req.from.Port, req.IfaceName())
							conflict.hostname = true
							break
						}
					}

				case *dns.SRV:
					if isDenyingSRV(rr, srv) {
						conflict.serviceName = true
					}

				default:
					break
				}
			}

		case <-ctx.Done():
			err = ctx.Err()
			return

		case <-queryTime:
			// Stop on conflict
			if conflict.hasAny() {
				return
			}

			// Stop after 3 probe queries
			if queriesCount > 3 {
				return
			}

			queriesCount++
			log.Debug.Println("Sending probe", msg)
			q := &Query{msg: msg, iface: &iface}
			conn.SendQuery(q)

			delay := 250 * time.Millisecond
			log.Debug.Println("Waiting for conflicting data", delay)
			queryTime = time.After(delay)
		}
	}

	return
}

type probeConflict struct {
	hostname    bool
	serviceName bool
}

func (pr probeConflict) hasNone() bool {
	return !pr.hostname && !pr.serviceName
}

func (pr probeConflict) hasAny() bool {
	return pr.hostname || pr.serviceName
}

func isDenyingA(this *dns.A, that *dns.A) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same hosts")

		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.A.To4(), that.A.To4()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

// isDenyingAAAA returns true if this denies that.
func isDenyingAAAA(this *dns.AAAA, that *dns.AAAA) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same hosts")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareIP(this.AAAA.To16(), that.AAAA.To16()) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

func areDenyingAs(this []*dns.A, that []*dns.A) bool {
	if len(this) != len(that) {
		// different number of As are a conflict
		return true
	}

	if equalAs(this, that) {
		return false // equal As are no conlict
	}

	if len(this) == 1 && !isDenyingA(this[0], that[0]) {
		// this doesn't deny that but the other way around
		return false
	}

	return true
}

func areDenyingAAAAs(this []*dns.AAAA, that []*dns.AAAA) bool {
	if len(this) != len(that) {
		// different number of AAAAs are a conflict
		return true
	}

	if equalAAAAs(this, that) {
		return false // equal AAAAs are no conlict
	}

	if len(this) == 1 && !isDenyingAAAA(this[0], that[0]) {
		// this doesn't deny that but the other way around
		return false
	}

	return true
}

func equalAs(this []*dns.A, that []*dns.A) bool {
	var tmp = that
	for _, ti := range this {
		var found = false
		for i, ta := range tmp {
			if compareIP(ti.A.To4(), ta.A.To4()) == 0 {
				tmp = append(tmp[:i], tmp[i+1:]...)
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return len(tmp) == 0
}

func equalAAAAs(this []*dns.AAAA, that []*dns.AAAA) bool {
	var tmp = that
	for _, ti := range this {
		var found = false
		for i, ta := range tmp {
			if compareIP(ti.AAAA.To16(), ta.AAAA.To16()) == 0 {
				tmp = append(tmp[:i], tmp[i+1:]...)
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}

	return len(tmp) == 0
}

func containedInAs(this *dns.A, aaas []*dns.A) bool {
	for _, that := range aaas {
		if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
			if !isValidRR(this) {
				log.Debug.Println("Invalid record produces conflict")
				return false
			}

			if compareIP(this.A.To4(), that.A.To4()) == 0 {
				return true
			}
		}
	}

	return false
}

func containedInAAAAs(this *dns.AAAA, aaas []*dns.AAAA) bool {
	for _, that := range aaas {
		if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
			if !isValidRR(this) {
				log.Debug.Println("Invalid record produces conflict")
				return false
			}

			if compareIP(this.AAAA.To16(), that.AAAA.To16()) == 0 {
				return true
			}
		}
	}

	return false
}

// isDenyingSRV returns true if this denies that.
func isDenyingSRV(this *dns.SRV, that *dns.SRV) bool {
	if strings.EqualFold(this.Hdr.Name, that.Hdr.Name) {
		log.Debug.Println("Same SRV")
		if !isValidRR(this) {
			log.Debug.Println("Invalid record produces conflict")
			return true
		}

		switch compareSRV(this, that) {
		case -1:
			log.Debug.Println("Lexicographical earlier")
			break
		case 1:
			log.Debug.Println("Lexicographical later")
			return true
		default:
			log.Debug.Println("No conflict")
			break
		}
	}

	return false
}

func isValidRR(rr dns.RR) bool {
	switch r := rr.(type) {
	case *dns.A:
		return !net.IPv4zero.Equal(r.A)
	case *dns.AAAA:
		return !net.IPv6zero.Equal(r.AAAA)
	case *dns.SRV:
		return len(r.Target) > 0 && r.Port != 0
	default:
		break
	}

	return true
}

func compareIP(this net.IP, that net.IP) int {
	count := len(this)
	if count > len(that) {
		count = len(that)
	}

	for i := 0; i < count; i++ {
		if this[i] < that[i] {
			return -1
		} else if this[i] > that[i] {
			return 1
		}
	}

	if len(this) < len(that) {
		return -1
	} else if len(this) > len(that) {
		return 1
	}
	return 0
}

func compareSRV(this *dns.SRV, that *dns.SRV) int {
	if this.Priority < that.Priority {
		return -1
	} else if this.Priority > that.Priority {
		return 1
	}

	if this.Weight < that.Weight {
		return -1
	} else if this.Weight > that.Weight {
		return 1
	}

	if this.Port < that.Port {
		return -1
	} else if this.Port > that.Port {
		return 1
	}

	return strings.Compare(this.Target, that.Target)
}
<|endoftext|>


<commit_before>package bbs

import "fmt"

//This has all the structs for various commands

type BBSCommand struct {
	Command string `json:"cmd"`
}

//From start to end inclusive, starting from 1.
type Range struct {
	Start int `json:"start"`
	End   int `json:"end"`
}

func (r *Range) String() string {
	return fmt.Sprintf("%d-%d", r.Start, r.End)
}

func (r *Range) Validate() bool {
	if r.Start > r.End {
		return false
	}
	return true
}

// "hello" message (server -> client)
type HelloMessage struct {
	Command         string     `json:"cmd"`
	Name            string     `json:"name"`
	ProtocolVersion int        `json:"version"`
	Description     string     `json:"description"`
	SecureURL       string     `json:"secure,omitempty"` //https URL, if any
	Options         []string   `json:"options,omitempty"`
	Access          AccessInfo `json:"access"`
	Formats         []string   `json:"format"` //formats the server accepts, the first one should be the primary one
	ServerVersion   string     `json:"server"`
}

// guest commands are commands you can use without logging on (e.g. "list", "get")
// user commands require being logged in first (usually "post" and "reply")
type AccessInfo struct {
	GuestCommands []string `json:"guest,omitempty"`
	UserCommands  []string `json:"user,omitempty"`
}

// "error" message (server -> client)
type ErrorMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Error   string `json:"error"`
}

// session expired or invalid? use this
var SessionErrorMessage *ErrorMessage = &ErrorMessage{"error", "session", "Invalid session."}

// "ok" message (server -> client)
type OKMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Message string `json:"msg,omitempty"`
}

// "login" command (client -> server)
type LoginCommand struct {
	Command         string `json:"cmd"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	ProtocolVersion int    `json:"version"`
	ClientVersion   string `json:"client"`
}

// "welcome" message (server -> client)
type WelcomeMessage struct {
	Command  string `json:"cmd"`
	Username string `json:"username"`
	Session  string `json:"session"`
}

// "logout" command (client -> server)
type LogoutCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

// "get" command (client -> server)
type GetCommand struct {
	Command  string `json:"cmd"`
	Session  string `json:"session,omitempty"`
	ThreadID string `json:"thread"`
	Range    *Range `json:"range"`
	Filter   string `json:"filter,omitempty"`
	Format   string `json:"format,omitempty"`
}

// "list" command (client -> server)
type ListCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	Type    string `json:"type"`
	Query   string `json:"query"`
}

// "reply" command (client -> server)
type ReplyCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	To      string `json:"to"`
	Text    string `json:"body"`
	Format  string `json:"format,omitempty"`
}

// "post" command (client -> server)
type PostCommand struct {
	Command string   `json:"cmd"`
	Session string   `json:"session,omitempty"`
	Title   string   `json:"title"`
	Text    string   `json:"body"`
	Format  string   `json:"format,omitempty"`
	Tags    []string `json:"tags,omitempty"`
}

// "msg" message (server -> client) [response to "get"]
type ThreadMessage struct {
	Command  string     `json:"cmd"`
	ID       string     `json:"id"`
	Title    string     `json:"title"`
	Range    *Range     `json:"range"`
	Filter   string     `json:"filter,omitempty"`
	Tags     []string   `json:"tags,omitempty"`
	Format   string     `json:"format,omitempty"`
	Messages []*Message `json:"messages"`
}

func (t *ThreadMessage) Size() int {
	if t.Messages != nil {
		return len(t.Messages)
	}
	return 0
}

// format for posts used in "msg"
type Message struct {
	ID                  string `json:"id"`
	Author              string `json:"user"`
	AuthorID            string `json:"user_id"`
	Date                string `json:"date"`
	Text                string `json:"body"`
	AuthorTitle         string `json:"user_title,omitempty"`
	AvatarURL           string `json:"avatar,omitempty"`
	PictureURL          string `json:"img,omitempty"`
	PictureThumbnailURL string `json:"thumb,omitempty"`
}

// "list" message (server -> client)
type ListMessage struct {
	Command string           `json:"cmd"`
	Type    string           `json:"type"`
	Query   string           `json:"query"`
	Threads []*ThreadListing `json:"threads"`
}

// format for threads in "list"
type ThreadListing struct {
	ID        string   `json:"id"`
	Title     string   `json:"title"`
	Author    string   `json:"user"`
	AuthorID  string   `json:"user_id"`
	Date      string   `json:"date"`
	PostCount int      `json:"posts"`
	Tags      []string `json:"tags,omitempty"`
}
<commit_msg>better comments, filling gaps ("boards", "imageboard" options)<commit_after>package bbs

import "fmt"

//This has all the structs for various commands

type BBSCommand struct {
	Command string `json:"cmd"`
}

type UserCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

//From start to end inclusive, starting from 1.
type Range struct {
	Start int `json:"start"`
	End   int `json:"end"`
}

func (r *Range) String() string {
	return fmt.Sprintf("%d-%d", r.Start, r.End)
}

func (r *Range) Validate() bool {
	if r.Start > r.End {
		return false
	}
	return true
}

// "hello" message (server -> client)
type HelloMessage struct {
	Command         string     `json:"cmd"`
	Name            string     `json:"name"`
	ProtocolVersion int        `json:"version"`
	Description     string     `json:"description"`
	SecureURL       string     `json:"secure,omitempty"` //https URL, if any
	Options         []string   `json:"options,omitempty"`
	Access          AccessInfo `json:"access"`
	Formats         []string   `json:"format"` //formats the server accepts, the first one should be the primary one
	ServerVersion   string     `json:"server"`
}

// guest commands are commands you can use without logging on (e.g. "list", "get")
// user commands require being logged in first (usually "post" and "reply")
type AccessInfo struct {
	GuestCommands []string `json:"guest,omitempty"`
	UserCommands  []string `json:"user,omitempty"`
}

// "error" message (server -> client)
type ErrorMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Error   string `json:"error"`
}

// session expired or invalid? use this
var SessionErrorMessage *ErrorMessage = &ErrorMessage{"error", "session", "Invalid session."}

// "ok" message (server -> client)
type OKMessage struct {
	Command string `json:"cmd"`
	ReplyTo string `json:"wrt"`
	Message string `json:"msg,omitempty"`
}

// "login" command (client -> server)
type LoginCommand struct {
	Command         string `json:"cmd"`
	Username        string `json:"username"`
	Password        string `json:"password"`
	ProtocolVersion int    `json:"version"`
	ClientVersion   string `json:"client"`
}

// "welcome" message (server -> client)
type WelcomeMessage struct {
	Command  string `json:"cmd"`
	Username string `json:"username,omitempty"` //omit for option 'anon'
	Session  string `json:"session"`
}

// "logout" command (client -> server)
type LogoutCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session"`
}

// "get" command (client -> server)
type GetCommand struct {
	Command  string `json:"cmd"`
	Session  string `json:"session,omitempty"`
	ThreadID string `json:"thread"`
	Board    string `json:"board,omitempty"` //option: "tags"
	Range    *Range `json:"range"`
	Filter   string `json:"filter,omitempty"` //option: "filter"
	Format   string `json:"format,omitempty"`
}

// "list" command (client -> server)
type ListCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	Type    string `json:"type"`
	Query   string `json:"query"` //board for "boards", tag expression for "tags" (like "Dogs+Pizza-Anime")
}

// "reply" command (client -> server)
type ReplyCommand struct {
	Command string `json:"cmd"`
	Session string `json:"session,omitempty"`
	To      string `json:"to"`
	Board   string `json:"board,omitempty"` //option: "boards"
	Text    string `json:"body"`
	Format  string `json:"format,omitempty"`
}

// "post" command (client -> server)
type PostCommand struct {
	Command string   `json:"cmd"`
	Session string   `json:"session,omitempty"`
	Title   string   `json:"title"`
	Text    string   `json:"body"`
	Format  string   `json:"format,omitempty"`
	Board   string   `json:"board,omitempty"` //option: "boards"
	Tags    []string `json:"tags,omitempty"`  //option: "tags"
}

// "msg" message (server -> client) [response to "get"]
type ThreadMessage struct {
	Command  string     `json:"cmd"`
	ID       string     `json:"id"`
	Title    string     `json:"title"`
	Range    *Range     `json:"range,omitempty"`
	Filter   string     `json:"filter,omitempty"` //option: "filter"
	Board    string     `json:"board,omitempty"`  //option: "boards"
	Tags     []string   `json:"tags,omitempty"`   //option: "tags"
	Format   string     `json:"format,omitempty"`
	Messages []*Message `json:"messages"`
}

func (t *ThreadMessage) Size() int {
	if t.Messages != nil {
		return len(t.Messages)
	}
	return 0
}

// format for posts used in "msg"
type Message struct {
	ID           string `json:"id"`
	Author       string `json:"user"`
	AuthorID     string `json:"user_id"`
	Date         string `json:"date"`
	Text         string `json:"body"`
	AuthorTitle  string `json:"user_title,omitempty"` //option: "usertitles"
	AvatarURL    string `json:"avatar,omitempty"`     //option: "avatars"
	PictureURL   string `json:"img,omitempty"`        //option: "imageboard"
	ThumbnailURL string `json:"thumb,omitempty"`      //option: "imageboard"
}

// "list" message (server -> client)
type ListMessage struct {
	Command string           `json:"cmd"`
	Type    string           `json:"type"`
	Query   string           `json:"query"`
	Threads []*ThreadListing `json:"threads"`
}

// format for threads in "list"
type ThreadListing struct {
	ID           string   `json:"id"`
	Title        string   `json:"title"`
	Author       string   `json:"user"`
	AuthorID     string   `json:"user_id"`
	Date         string   `json:"date"`
	PostCount    int      `json:"posts"`
	Tags         []string `json:"tags,omitempty"`  //option: "tags"
	PictureURL   string   `json:"img,omitempty"`   //option: "imageboard"
	ThumbnailURL string   `json:"thumb,omitempty"` //option: "imageboard"
}
<|endoftext|>


<commit_before>/*
   Copyright 2015 Ian Bishop

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

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

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

package main

import (
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/vishvananda/netlink"
)

const routeCheckInterval = 30

type listener struct {
	sync.RWMutex
	extChan, intChan  chan ndp
	errChan           chan error
	ruleNet           *net.IPNet
	started, finished bool
}

type sessionStatus int

// sessions track clients who've previously sent neighbor solicits
type session struct {
	upstream             *listener
	srcIP, dstIP, target net.IP
	status               sessionStatus
	expiry               time.Time
}

type ndp struct {
	ifname  string
	payload gopacket.Payload
	icmp    layers.ICMPv6
	ip6     layers.IPv6
	eth     layers.Ethernet
}

const (
	waiting sessionStatus = iota
	valid
	invalid

	ttl = time.Duration(500 * time.Millisecond)
)

func Proxy(wg *sync.WaitGroup, ifname string, rules []string) {
	defer wg.Done()

	var err error
	upstreams := make(map[string]*listener)
	// shared channels upstreams send to
	errChan := make(chan error)
	intChan := make(chan ndp)
	mainInChan := make(chan ndp)
	tickChan := time.NewTicker(time.Second * routeCheckInterval).C

	var sessions []session

	// launch handler for main interface 'ifname'
	l := &listener{intChan: intChan, extChan: mainInChan, errChan: errChan}
	go l.Handler(ifname)

	err = refreshRoutes(rules, intChan, errChan, upstreams)
	if err != nil {
		fmt.Printf("%s\n", err)
		return
	}

	for {
		select {
		case err = <-errChan:
			fmt.Printf("%s\n", err)
			return
		case n := <-intChan:
			sessions = proxyPacket(ifname, n, mainInChan, upstreams, sessions)
		case <-tickChan:
			sessions = updateSessions(sessions)
			err := refreshRoutes(rules, intChan, errChan, upstreams)
			if err != nil {
				fmt.Printf("%s\n", err)
				return
			}
		}
	}
}

func proxyPacket(ifname string, n ndp, mainInChan chan ndp, upstreams map[string]*listener, sessions []session) []session {
	var target net.IP
	// IPv6 bounds check
	if len(n.payload) >= 16 {
		target = net.IP(n.payload[:16])
	} else {
		return sessions
	}

	switch n.icmp.TypeCode.Type() {
	case layers.ICMPv6TypeNeighborAdvertisement:
		for i, s := range sessions {
			if s.target.Equal(target) {
				sessions[i].status = valid
				n.ip6.DstIP = s.srcIP
				mainInChan <- n
				return sessions
			}
		}
	case layers.ICMPv6TypeNeighborSolicitation:
		if !n.ip6.DstIP.IsMulticast() {
			return sessions
		}
		for _, s := range sessions {
			if s.target.Equal(target) {

				switch s.status {
				case waiting:
				case invalid:
					break
				case valid:
					s.upstream.extChan <- n
				}
				return sessions
			}
		}

		var s *session
		// if msg arrived from the main interface, then send to matching upstreams
		for _, upstream := range upstreams {
			if upstream.ruleNet.Contains(target) {
				s = &session{
					upstream: upstream,
					srcIP:    n.ip6.SrcIP,
					dstIP:    n.ip6.DstIP,
					target:   target,
					status:   waiting,
					expiry:   time.Now().Add(ttl),
				}
			}
		}

		if s != nil {
			sessions = append(sessions, *s)
			s.upstream.extChan <- n
		}
	}

	return sessions
}

func updateSessions(sessions []session) []session {
	log.Printf("update sessions")
	for i := 0; i < len(sessions); i++ {

		if sessions[i].expiry.After(time.Now()) {
			continue
		}

		switch sessions[i].status {
		case waiting:
			sessions[i].status = invalid
			sessions[i].expiry = time.Now().Add(ttl)
		default:
			log.Printf("remove sess %d, target %s", i, sessions[i].target)
			// remove session
			if i == len(sessions)-1 {
				sessions = sessions[:i]
			} else {
				sessions = append(sessions[:i], sessions[i+1:]...)
			}
		}

	}
	return sessions
}

func refreshRoutes(rules []string, intChan chan ndp, errChan chan error, upstreams map[string]*listener) error {
	log.Printf("Refreshing routes...")
	for _, rule := range rules {
		_, ruleNet, err := net.ParseCIDR(rule)
		if err != nil {
			return fmt.Errorf("invalid rule '%s', %s", rule, err)
		}
		routes, err := netlink.RouteGet(ruleNet.IP)
		if err != nil || len(routes) == 0 {
			// cancel any proxies for removed routes
			for _, upstream := range upstreams {
				if upstream.ruleNet.IP.Equal(ruleNet.IP) {
					close(upstream.extChan)
				}
			}
			// route not found, skip
			continue
		}
		links, err := netlink.LinkList()
		if err != nil {
			return fmt.Errorf("error enumerating links, %s", err)
		}
		for _, link := range links {
			for _, route := range routes {
				if link.Attrs().Index == route.LinkIndex && route.Gw == nil {
					if _, ok := upstreams[link.Attrs().Name]; !ok {
						log.Printf("New upstream for link '%s', rule '%s', route '%s'\n", link.Attrs().Name, rule, route.Dst)
						upstreams[link.Attrs().Name] = &listener{
							extChan: make(chan ndp),
							intChan: intChan,
							errChan: errChan,
							ruleNet: ruleNet,
						}
					}
				}
			}
		}
	}
	for name, listener := range upstreams {
		listener.RLock()
		if !listener.started {
			// launch handler for an upstream of the main proxy interface
			go listener.Handler(name)
		}
		if listener.finished {
			delete(upstreams, name)
		}
		listener.RUnlock()
	}
	return nil
}

func (l *listener) Handler(ifname string) {
	var err error
	var handle *pcap.Handle
	log.Printf("Spawning listener for if %s\n", ifname)
	l.Lock()
	l.started = true
	l.Unlock()

	defer func() {
		l.Lock()
		l.finished = true
		l.Unlock()
		if err != nil {
			l.errChan <- err
		}
	}()

	handle, err = pcap.OpenLive(ifname, snaplen, true, 0)
	if err != nil {
		err = fmt.Errorf("pcap open error: %s", err)
		return
	}
	defer handle.Close()
	defer func() {
		log.Printf("%s: handler exit %s", ifname, err)
	}()

	var iface *net.Interface
	iface, err = net.InterfaceByName(ifname)
	if err != nil {
		return
	}

	var addrs []net.Addr
	var linklocal net.IP
	addrs, err = iface.Addrs()
	if err != nil {
		return
	}

	for _, addr := range addrs {
		switch v := addr.(type) {
		case *net.IPNet:
			if v.IP.IsLinkLocalUnicast() {
				linklocal = v.IP
			}
		}
	}

	if linklocal.IsUnspecified() {
		err = fmt.Errorf("error finding link local unicast address for interface %s", ifname)
		return
	}

	var eth layers.Ethernet
	var ip6 layers.IPv6
	var ip6extensions layers.IPv6ExtensionSkipper
	var icmp layers.ICMPv6
	var payload gopacket.Payload
	decoded := []gopacket.LayerType{}
	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeEthernet, &eth, &ip6, &ip6extensions, &icmp, &payload)
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetsChan := packetSource.Packets()
	for {
		select {
		case packet := <-packetsChan:
			parser.DecodeLayers(packet.Data(), &decoded)
			for _, layerType := range decoded {
				switch layerType {
				case layers.LayerTypeICMPv6:

					var target net.IP
					// IPv6 bounds check
					if len(payload) >= 16 {
						target = net.IP(payload[:16])
					} else {
						continue
					}

					switch icmp.TypeCode.Type() {
					case layers.ICMPv6TypeNeighborSolicitation:
					case layers.ICMPv6TypeNeighborAdvertisement:
						n := ndp{eth: eth, ip6: ip6, icmp: icmp, ifname: ifname, payload: payload}
						log.Printf("%s: read ndp %s, ip6 src %s, dst %s, target %s\n", ifname, icmp.TypeCode, ip6.SrcIP, ip6.DstIP, target)
						l.intChan <- n
					}
				}
			}

		case n, ok := <-l.extChan:
			if !ok {
				// channel was closed
				return
			}
			eth := n.eth

			eth.DstMAC = nil
			if n.ip6.DstIP.IsLinkLocalMulticast() {
				// Ethernet MAC is derived by the four low-order octets of IPv6 address
				eth.DstMAC = net.HardwareAddr{0x33, 0x33}
				eth.DstMAC = append(eth.DstMAC, n.ip6.DstIP[12:]...)
			} else {
				var neighbors []netlink.Neigh
				neighbors, err = netlink.NeighList(iface.Index, netlink.FAMILY_V6)
				if err != nil {
					return
				}
				for _, neighbor := range neighbors {
					if neighbor.IP.Equal(n.ip6.DstIP) {
						eth.DstMAC = neighbor.HardwareAddr
					}
				}
			}
			if eth.DstMAC == nil {
				err = fmt.Errorf("could not find destination MAC address for IP %s", n.ip6.DstIP)
				return
			}
			eth.SrcMAC = iface.HardwareAddr
			ipv6 := n.ip6
			ipv6.SrcIP = linklocal
			buf := gopacket.NewSerializeBuffer()
			n.icmp.SetNetworkLayerForChecksum(&ipv6)
			opts := gopacket.SerializeOptions{ComputeChecksums: true}
			switch n.icmp.TypeCode.Type() {
			case layers.ICMPv6TypeNeighborSolicitation:
				// ND solicit opt type, opt length
				n.payload = append(n.payload[:16], 0x01, 0x01)
			case layers.ICMPv6TypeNeighborAdvertisement:
				// ND advert opt type, opt length
				n.payload = append(n.payload[:16], 0x02, 0x01)
				n.icmp.TypeBytes[0] = 0xc0 // router,solicit,override flags
			}
			n.payload = append(n.payload[:18], iface.HardwareAddr...)

			err = gopacket.SerializeLayers(buf, opts, &eth, &ipv6, &n.icmp, &n.payload)
			if err != nil {
				err = fmt.Errorf("serialize layers error: %s", err)
				return
			}

			err = handle.WritePacketData(buf.Bytes())
			if err != nil {
				err = fmt.Errorf("pcap write error: %s", err)
				return
			}
			log.Printf("%s: write ndp %s, ip6 src %s, dst %s, target %s\n", ifname, n.icmp.TypeCode, ipv6.SrcIP, ipv6.DstIP, net.IP(n.payload[:16]))
		}
	}
}
<commit_msg>fix route selection<commit_after>/*
   Copyright 2015 Ian Bishop

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

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

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

package main

import (
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/vishvananda/netlink"
)

const routeCheckInterval = 30

type listener struct {
	sync.RWMutex
	extChan, intChan  chan ndp
	errChan           chan error
	ruleNet           *net.IPNet
	started, finished bool
}

type sessionStatus int

// sessions track clients who've previously sent neighbor solicits
type session struct {
	upstream             *listener
	srcIP, dstIP, target net.IP
	status               sessionStatus
	expiry               time.Time
}

type ndp struct {
	ifname  string
	payload gopacket.Payload
	icmp    layers.ICMPv6
	ip6     layers.IPv6
	eth     layers.Ethernet
}

const (
	waiting sessionStatus = iota
	valid
	invalid

	ttl = time.Duration(500 * time.Millisecond)
)

func Proxy(wg *sync.WaitGroup, ifname string, rules []string) {
	defer wg.Done()

	var err error
	upstreams := make(map[string]*listener)
	// shared channels upstreams send to
	errChan := make(chan error)
	intChan := make(chan ndp)
	mainInChan := make(chan ndp)
	tickChan := time.NewTicker(time.Second * routeCheckInterval).C

	var sessions []session

	// launch handler for main interface 'ifname'
	l := &listener{intChan: intChan, extChan: mainInChan, errChan: errChan}
	go l.Handler(ifname)

	err = refreshRoutes(rules, intChan, errChan, upstreams)
	if err != nil {
		fmt.Printf("%s\n", err)
		return
	}

	for {
		select {
		case err = <-errChan:
			fmt.Printf("%s\n", err)
			return
		case n := <-intChan:
			sessions = proxyPacket(ifname, n, mainInChan, upstreams, sessions)
		case <-tickChan:
			sessions = updateSessions(sessions)
			err := refreshRoutes(rules, intChan, errChan, upstreams)
			if err != nil {
				fmt.Printf("%s\n", err)
				return
			}
		}
	}
}

func proxyPacket(ifname string, n ndp, mainInChan chan ndp, upstreams map[string]*listener, sessions []session) []session {
	var target net.IP
	// IPv6 bounds check
	if len(n.payload) >= 16 {
		target = net.IP(n.payload[:16])
	} else {
		return sessions
	}

	switch n.icmp.TypeCode.Type() {
	case layers.ICMPv6TypeNeighborAdvertisement:
		for i, s := range sessions {
			if s.target.Equal(target) {
				sessions[i].status = valid
				n.ip6.DstIP = s.srcIP
				mainInChan <- n
				return sessions
			}
		}
	case layers.ICMPv6TypeNeighborSolicitation:
		if !n.ip6.DstIP.IsMulticast() {
			return sessions
		}
		for _, s := range sessions {
			if s.target.Equal(target) {

				switch s.status {
				case waiting:
				case invalid:
					break
				case valid:
					s.upstream.extChan <- n
				}
				return sessions
			}
		}

		var s *session
		// if msg arrived from the main interface, then send to matching upstreams
		for _, upstream := range upstreams {
			if upstream.ruleNet.Contains(target) {
				s = &session{
					upstream: upstream,
					srcIP:    n.ip6.SrcIP,
					dstIP:    n.ip6.DstIP,
					target:   target,
					status:   waiting,
					expiry:   time.Now().Add(ttl),
				}
			}
		}

		if s != nil {
			sessions = append(sessions, *s)
			s.upstream.extChan <- n
		}
	}

	return sessions
}

func updateSessions(sessions []session) []session {
	log.Printf("update sessions")
	for i := 0; i < len(sessions); i++ {

		if sessions[i].expiry.After(time.Now()) {
			continue
		}

		switch sessions[i].status {
		case waiting:
			sessions[i].status = invalid
			sessions[i].expiry = time.Now().Add(ttl)
		default:
			log.Printf("remove sess %d, target %s", i, sessions[i].target)
			// remove session
			if i == len(sessions)-1 {
				sessions = sessions[:i]
			} else {
				sessions = append(sessions[:i], sessions[i+1:]...)
			}
		}

	}
	return sessions
}

func refreshRoutes(rules []string, intChan chan ndp, errChan chan error, upstreams map[string]*listener) error {
	log.Printf("Refreshing routes...")
	for _, rule := range rules {
		_, ruleNet, err := net.ParseCIDR(rule)
		if err != nil {
			return fmt.Errorf("invalid rule '%s', %s", rule, err)
		}

		routes, err := netlink.RouteList(nil, netlink.FAMILY_V6)
		if err != nil {
			return fmt.Errorf("error enumerating routes, %s", err)
		}
		var route *netlink.Route
		for _, r := range routes {
			if r.Dst.Contains(ruleNet.IP) {
				route = &r
				break
			}
		}

		if route == nil {
			// cancel any proxies for removed routes
			for _, upstream := range upstreams {
				if upstream.ruleNet.IP.Equal(ruleNet.IP) {
					close(upstream.extChan)
				}
			}
			// route not found, skip
			continue
		}
		links, err := netlink.LinkList()
		if err != nil {
			return fmt.Errorf("error enumerating links, %s", err)
		}
		for _, link := range links {
			if link.Attrs().Index == route.LinkIndex {
				if _, ok := upstreams[link.Attrs().Name]; !ok {
					log.Printf("New upstream for link '%s', rule '%s', route '%s'\n", link.Attrs().Name, rule, route.Dst)
					upstreams[link.Attrs().Name] = &listener{
						extChan: make(chan ndp),
						intChan: intChan,
						errChan: errChan,
						ruleNet: ruleNet,
					}
				}
			}
		}
	}
	for name, listener := range upstreams {
		listener.RLock()
		if !listener.started {
			// launch handler for an upstream of the main proxy interface
			go listener.Handler(name)
		}
		if listener.finished {
			delete(upstreams, name)
		}
		listener.RUnlock()
	}
	return nil
}

func (l *listener) Handler(ifname string) {
	var err error
	var handle *pcap.Handle
	log.Printf("Spawning listener for if %s\n", ifname)
	l.Lock()
	l.started = true
	l.Unlock()

	defer func() {
		l.Lock()
		l.finished = true
		l.Unlock()
		if err != nil {
			l.errChan <- err
		}
	}()

	handle, err = pcap.OpenLive(ifname, snaplen, true, 0)
	if err != nil {
		err = fmt.Errorf("pcap open error: %s", err)
		return
	}
	defer handle.Close()
	defer func() {
		log.Printf("%s: handler exit %s", ifname, err)
	}()

	var iface *net.Interface
	iface, err = net.InterfaceByName(ifname)
	if err != nil {
		return
	}

	var addrs []net.Addr
	var linklocal net.IP
	addrs, err = iface.Addrs()
	if err != nil {
		return
	}

	for _, addr := range addrs {
		switch v := addr.(type) {
		case *net.IPNet:
			if v.IP.IsLinkLocalUnicast() {
				linklocal = v.IP
			}
		}
	}

	if linklocal.IsUnspecified() {
		err = fmt.Errorf("error finding link local unicast address for interface %s", ifname)
		return
	}

	var eth layers.Ethernet
	var ip6 layers.IPv6
	var ip6extensions layers.IPv6ExtensionSkipper
	var icmp layers.ICMPv6
	var payload gopacket.Payload
	decoded := []gopacket.LayerType{}
	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeEthernet, &eth, &ip6, &ip6extensions, &icmp, &payload)
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetsChan := packetSource.Packets()
	for {
		select {
		case packet := <-packetsChan:
			parser.DecodeLayers(packet.Data(), &decoded)
			for _, layerType := range decoded {
				switch layerType {
				case layers.LayerTypeICMPv6:

					var target net.IP
					// IPv6 bounds check
					if len(payload) >= 16 {
						target = net.IP(payload[:16])
					} else {
						continue
					}

					switch icmp.TypeCode.Type() {
					case layers.ICMPv6TypeNeighborSolicitation:
					case layers.ICMPv6TypeNeighborAdvertisement:
						n := ndp{eth: eth, ip6: ip6, icmp: icmp, ifname: ifname, payload: payload}
						log.Printf("%s: read ndp %s, ip6 src %s, dst %s, target %s\n", ifname, icmp.TypeCode, ip6.SrcIP, ip6.DstIP, target)
						l.intChan <- n
					}
				}
			}

		case n, ok := <-l.extChan:
			if !ok {
				// channel was closed
				return
			}
			eth := n.eth

			eth.DstMAC = nil
			if n.ip6.DstIP.IsLinkLocalMulticast() {
				// Ethernet MAC is derived by the four low-order octets of IPv6 address
				eth.DstMAC = net.HardwareAddr{0x33, 0x33}
				eth.DstMAC = append(eth.DstMAC, n.ip6.DstIP[12:]...)
			} else {
				var neighbors []netlink.Neigh
				neighbors, err = netlink.NeighList(iface.Index, netlink.FAMILY_V6)
				if err != nil {
					return
				}
				for _, neighbor := range neighbors {
					if neighbor.IP.Equal(n.ip6.DstIP) {
						eth.DstMAC = neighbor.HardwareAddr
					}
				}
			}
			if eth.DstMAC == nil {
				err = fmt.Errorf("could not find destination MAC address for IP %s", n.ip6.DstIP)
				return
			}
			eth.SrcMAC = iface.HardwareAddr
			ipv6 := n.ip6
			ipv6.SrcIP = linklocal
			buf := gopacket.NewSerializeBuffer()
			n.icmp.SetNetworkLayerForChecksum(&ipv6)
			opts := gopacket.SerializeOptions{ComputeChecksums: true}
			switch n.icmp.TypeCode.Type() {
			case layers.ICMPv6TypeNeighborSolicitation:
				// ND solicit opt type, opt length
				n.payload = append(n.payload[:16], 0x01, 0x01)
			case layers.ICMPv6TypeNeighborAdvertisement:
				// ND advert opt type, opt length
				n.payload = append(n.payload[:16], 0x02, 0x01)
				n.icmp.TypeBytes[0] = 0xc0 // router,solicit,override flags
			}
			n.payload = append(n.payload[:18], iface.HardwareAddr...)

			err = gopacket.SerializeLayers(buf, opts, &eth, &ipv6, &n.icmp, &n.payload)
			if err != nil {
				err = fmt.Errorf("serialize layers error: %s", err)
				return
			}

			err = handle.WritePacketData(buf.Bytes())
			if err != nil {
				err = fmt.Errorf("pcap write error: %s", err)
				return
			}
			log.Printf("%s: write ndp %s, ip6 src %s, dst %s, target %s\n", ifname, n.icmp.TypeCode, ipv6.SrcIP, ipv6.DstIP, net.IP(n.payload[:16]))
		}
	}
}
<|endoftext|>


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

package flags

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

// Marshaler is the interface implemented by types that can marshal themselves
// to a string representation of the flag.
type Marshaler interface {
	// MarshalFlag marshals a flag value to its string representation.
	MarshalFlag() (string, error)
}

// Unmarshaler is the interface implemented by types that can unmarshal a flag
// argument to themselves. The provided value is directly passed from the
// command line.
type Unmarshaler interface {
	// UnmarshalFlag unmarshals a string value representation to the flag
	// value (which therefore needs to be a pointer receiver).
	UnmarshalFlag(value string) error
}

func getBase(options multiTag, base int) (int, error) {
	sbase := options.Get("base")

	var err error
	var ivbase int64

	if sbase != "" {
		ivbase, err = strconv.ParseInt(sbase, 10, 32)
		base = int(ivbase)
	}

	return base, err
}

func convertMarshal(val reflect.Value) (bool, string, error) {
	// Check first for the Marshaler interface
	if val.Type().NumMethod() > 0 && val.CanInterface() {
		if marshaler, ok := val.Interface().(Marshaler); ok {
			ret, err := marshaler.MarshalFlag()
			return true, ret, err
		}
	}

	return false, "", nil
}

func convertToString(val reflect.Value, options multiTag) (string, error) {
	if ok, ret, err := convertMarshal(val); ok {
		return ret, err
	}

	tp := val.Type()

	// Support for time.Duration
	if tp == reflect.TypeOf((*time.Duration)(nil)).Elem() {
		stringer := val.Interface().(fmt.Stringer)
		return stringer.String(), nil
	}

	switch tp.Kind() {
	case reflect.String:
		return val.String(), nil
	case reflect.Bool:
		if val.Bool() {
			return "true", nil
		}

		return "false", nil
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		base, err := getBase(options, 10)

		if err != nil {
			return "", err
		}

		return strconv.FormatInt(val.Int(), base), nil
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		base, err := getBase(options, 10)

		if err != nil {
			return "", err
		}

		return strconv.FormatUint(val.Uint(), base), nil
	case reflect.Float32, reflect.Float64:
		return strconv.FormatFloat(val.Float(), 'g', -1, tp.Bits()), nil
	case reflect.Slice:
		if val.Len() == 0 {
			return "", nil
		}

		ret := "["

		for i := 0; i < val.Len(); i++ {
			if i != 0 {
				ret += ", "
			}

			item, err := convertToString(val.Index(i), options)

			if err != nil {
				return "", err
			}

			ret += item
		}

		return ret + "]", nil
	case reflect.Map:
		ret := "{"

		for i, key := range val.MapKeys() {
			if i != 0 {
				ret += ", "
			}

			keyitem, err := convertToString(key, options)

			if err != nil {
				return "", err
			}

			item, err := convertToString(val.MapIndex(key), options)

			if err != nil {
				return "", err
			}

			ret += keyitem + ":" + item
		}

		return ret + "}", nil
	case reflect.Ptr:
		return convertToString(reflect.Indirect(val), options)
	case reflect.Interface:
		if !val.IsNil() {
			return convertToString(val.Elem(), options)
		}
	}

	return "", nil
}

func convertUnmarshal(val string, retval reflect.Value) (bool, error) {
	if retval.Type().NumMethod() > 0 && retval.CanInterface() {
		if unmarshaler, ok := retval.Interface().(Unmarshaler); ok {
			return true, unmarshaler.UnmarshalFlag(val)
		}
	}

	if retval.Type().Kind() != reflect.Ptr && retval.CanAddr() {
		return convertUnmarshal(val, retval.Addr())
	}

	if retval.Type().Kind() == reflect.Interface && !retval.IsNil() {
		return convertUnmarshal(val, retval.Elem())
	}

	return false, nil
}

func convert(val string, retval reflect.Value, options multiTag) error {
	if ok, err := convertUnmarshal(val, retval); ok {
		return err
	}

	tp := retval.Type()

	// Support for time.Duration
	if tp == reflect.TypeOf((*time.Duration)(nil)).Elem() {
		parsed, err := time.ParseDuration(val)

		if err != nil {
			return err
		}

		retval.SetInt(int64(parsed))
		return nil
	}

	switch tp.Kind() {
	case reflect.String:
		retval.SetString(val)
	case reflect.Bool:
		if val == "" {
			retval.SetBool(true)
		} else {
			b, err := strconv.ParseBool(val)

			if err != nil {
				return err
			}

			retval.SetBool(b)
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		base, err := getBase(options, 10)

		if err != nil {
			return err
		}

		parsed, err := strconv.ParseInt(val, base, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetInt(parsed)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		base, err := getBase(options, 10)

		if err != nil {
			return err
		}

		parsed, err := strconv.ParseUint(val, base, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetUint(parsed)
	case reflect.Float32, reflect.Float64:
		parsed, err := strconv.ParseFloat(val, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetFloat(parsed)
	case reflect.Slice:
		elemtp := tp.Elem()

		elemvalptr := reflect.New(elemtp)
		elemval := reflect.Indirect(elemvalptr)

		if err := convert(val, elemval, options); err != nil {
			return err
		}

		retval.Set(reflect.Append(retval, elemval))
	case reflect.Map:
		parts := strings.SplitN(val, ":", 2)

		key := parts[0]
		var value string

		if len(parts) == 2 {
			value = parts[1]
		}

		keytp := tp.Key()
		keyval := reflect.New(keytp)

		if err := convert(key, keyval, options); err != nil {
			return err
		}

		valuetp := tp.Elem()
		valueval := reflect.New(valuetp)

		if err := convert(value, valueval, options); err != nil {
			return err
		}

		if retval.IsNil() {
			retval.Set(reflect.MakeMap(tp))
		}

		retval.SetMapIndex(reflect.Indirect(keyval), reflect.Indirect(valueval))
	case reflect.Ptr:
		if retval.IsNil() {
			retval.Set(reflect.New(retval.Type().Elem()))
		}

		return convert(val, reflect.Indirect(retval), options)
	case reflect.Interface:
		if !retval.IsNil() {
			return convert(val, retval.Elem(), options)
		}
	}

	return nil
}

func isPrint(s string) bool {
	for _, c := range s {
		if !strconv.IsPrint(c) {
			return false
		}
	}

	return true
}

func quoteIfNeeded(s string) string {
	if !isPrint(s) {
		return strconv.Quote(s)
	}

	return s
}

func quoteIfNeededV(s []string) []string {
	ret := make([]string, len(s))

	for i, v := range s {
		ret[i] = quoteIfNeeded(v)
	}

	return ret
}

func quoteV(s []string) []string {
	ret := make([]string, len(s))

	for i, v := range s {
		ret[i] = strconv.Quote(v)
	}

	return ret
}

func unquoteIfPossible(s string) (string, error) {
	if len(s) == 0 || s[0] != '"' {
		return s, nil
	}

	return strconv.Unquote(s)
}
<commit_msg>vendor change: go-flags Factory<commit_after>// Copyright 2012 Jesse van den Kieboom. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package flags

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"
)

// Marshaler is the interface implemented by types that can marshal themselves
// to a string representation of the flag.
type Marshaler interface {
	// MarshalFlag marshals a flag value to its string representation.
	MarshalFlag() (string, error)
}

// Unmarshaler is the interface implemented by types that can unmarshal a flag
// argument to themselves. The provided value is directly passed from the
// command line.
type Unmarshaler interface {
	// UnmarshalFlag unmarshals a string value representation to the flag
	// value (which therefore needs to be a pointer receiver).
	UnmarshalFlag(value string) error
}

func getBase(options multiTag, base int) (int, error) {
	sbase := options.Get("base")

	var err error
	var ivbase int64

	if sbase != "" {
		ivbase, err = strconv.ParseInt(sbase, 10, 32)
		base = int(ivbase)
	}

	return base, err
}

func convertMarshal(val reflect.Value) (bool, string, error) {
	// Check first for the Marshaler interface
	if val.Type().NumMethod() > 0 && val.CanInterface() {
		if marshaler, ok := val.Interface().(Marshaler); ok {
			ret, err := marshaler.MarshalFlag()
			return true, ret, err
		}
	}

	return false, "", nil
}

func convertToString(val reflect.Value, options multiTag) (string, error) {
	if ok, ret, err := convertMarshal(val); ok {
		return ret, err
	}

	tp := val.Type()

	// Support for time.Duration
	if tp == reflect.TypeOf((*time.Duration)(nil)).Elem() {
		stringer := val.Interface().(fmt.Stringer)
		return stringer.String(), nil
	}

	switch tp.Kind() {
	case reflect.String:
		return val.String(), nil
	case reflect.Bool:
		if val.Bool() {
			return "true", nil
		}

		return "false", nil
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		base, err := getBase(options, 10)

		if err != nil {
			return "", err
		}

		return strconv.FormatInt(val.Int(), base), nil
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		base, err := getBase(options, 10)

		if err != nil {
			return "", err
		}

		return strconv.FormatUint(val.Uint(), base), nil
	case reflect.Float32, reflect.Float64:
		return strconv.FormatFloat(val.Float(), 'g', -1, tp.Bits()), nil
	case reflect.Slice:
		if val.Len() == 0 {
			return "", nil
		}

		ret := "["

		for i := 0; i < val.Len(); i++ {
			if i != 0 {
				ret += ", "
			}

			item, err := convertToString(val.Index(i), options)

			if err != nil {
				return "", err
			}

			ret += item
		}

		return ret + "]", nil
	case reflect.Map:
		ret := "{"

		for i, key := range val.MapKeys() {
			if i != 0 {
				ret += ", "
			}

			keyitem, err := convertToString(key, options)

			if err != nil {
				return "", err
			}

			item, err := convertToString(val.MapIndex(key), options)

			if err != nil {
				return "", err
			}

			ret += keyitem + ":" + item
		}

		return ret + "}", nil
	case reflect.Ptr:
		return convertToString(reflect.Indirect(val), options)
	case reflect.Interface:
		if !val.IsNil() {
			return convertToString(val.Elem(), options)
		}
	}

	return "", nil
}

var FactoryFunc func(interface{})

func convertUnmarshal(val string, retval reflect.Value) (bool, error) {
	if retval.Type().NumMethod() > 0 && retval.CanInterface() {
		if unmarshaler, ok := retval.Interface().(Unmarshaler); ok {
			FactoryFunc(unmarshaler)
			return true, unmarshaler.UnmarshalFlag(val)
		}
	}

	if retval.Type().Kind() != reflect.Ptr && retval.CanAddr() {
		return convertUnmarshal(val, retval.Addr())
	}

	if retval.Type().Kind() == reflect.Interface && !retval.IsNil() {
		return convertUnmarshal(val, retval.Elem())
	}

	return false, nil
}

func convert(val string, retval reflect.Value, options multiTag) error {
	if ok, err := convertUnmarshal(val, retval); ok {
		return err
	}

	tp := retval.Type()

	// Support for time.Duration
	if tp == reflect.TypeOf((*time.Duration)(nil)).Elem() {
		parsed, err := time.ParseDuration(val)

		if err != nil {
			return err
		}

		retval.SetInt(int64(parsed))
		return nil
	}

	switch tp.Kind() {
	case reflect.String:
		retval.SetString(val)
	case reflect.Bool:
		if val == "" {
			retval.SetBool(true)
		} else {
			b, err := strconv.ParseBool(val)

			if err != nil {
				return err
			}

			retval.SetBool(b)
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		base, err := getBase(options, 10)

		if err != nil {
			return err
		}

		parsed, err := strconv.ParseInt(val, base, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetInt(parsed)
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		base, err := getBase(options, 10)

		if err != nil {
			return err
		}

		parsed, err := strconv.ParseUint(val, base, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetUint(parsed)
	case reflect.Float32, reflect.Float64:
		parsed, err := strconv.ParseFloat(val, tp.Bits())

		if err != nil {
			return err
		}

		retval.SetFloat(parsed)
	case reflect.Slice:
		elemtp := tp.Elem()

		elemvalptr := reflect.New(elemtp)
		elemval := reflect.Indirect(elemvalptr)

		if err := convert(val, elemval, options); err != nil {
			return err
		}

		retval.Set(reflect.Append(retval, elemval))
	case reflect.Map:
		parts := strings.SplitN(val, ":", 2)

		key := parts[0]
		var value string

		if len(parts) == 2 {
			value = parts[1]
		}

		keytp := tp.Key()
		keyval := reflect.New(keytp)

		if err := convert(key, keyval, options); err != nil {
			return err
		}

		valuetp := tp.Elem()
		valueval := reflect.New(valuetp)

		if err := convert(value, valueval, options); err != nil {
			return err
		}

		if retval.IsNil() {
			retval.Set(reflect.MakeMap(tp))
		}

		retval.SetMapIndex(reflect.Indirect(keyval), reflect.Indirect(valueval))
	case reflect.Ptr:
		if retval.IsNil() {
			retval.Set(reflect.New(retval.Type().Elem()))
		}

		return convert(val, reflect.Indirect(retval), options)
	case reflect.Interface:
		if !retval.IsNil() {
			return convert(val, retval.Elem(), options)
		}
	}

	return nil
}

func isPrint(s string) bool {
	for _, c := range s {
		if !strconv.IsPrint(c) {
			return false
		}
	}

	return true
}

func quoteIfNeeded(s string) string {
	if !isPrint(s) {
		return strconv.Quote(s)
	}

	return s
}

func quoteIfNeededV(s []string) []string {
	ret := make([]string, len(s))

	for i, v := range s {
		ret[i] = quoteIfNeeded(v)
	}

	return ret
}

func quoteV(s []string) []string {
	ret := make([]string, len(s))

	for i, v := range s {
		ret[i] = strconv.Quote(v)
	}

	return ret
}

func unquoteIfPossible(s string) (string, error) {
	if len(s) == 0 || s[0] != '"' {
		return s, nil
	}

	return strconv.Unquote(s)
}
<|endoftext|>