<commit_before>package restic

import (
	"fmt"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/restic/restic/internal/debug"
)

// ExpirePolicy configures which snapshots should be automatically removed.
type ExpirePolicy struct {
	Last          int       // keep the last n snapshots
	Hourly        int       // keep the last n hourly snapshots
	Daily         int       // keep the last n daily snapshots
	Weekly        int       // keep the last n weekly snapshots
	Monthly       int       // keep the last n monthly snapshots
	Yearly        int       // keep the last n yearly snapshots
	Within        Duration  // keep snapshots made within this duration
	WithinHourly  Duration  // keep hourly snapshots made within this duration
	WithinDaily   Duration  // keep daily snapshots made within this duration
	WithinWeekly  Duration  // keep weekly snapshots made within this duration
	WithinMonthly Duration  // keep monthly snapshots made within this duration
	WithinYearly  Duration  // keep yearly snapshots made within this duration
	Tags          []TagList // keep all snapshots that include at least one of the tag lists.
}

func (e ExpirePolicy) String() (s string) {
	var keeps []string
	var keepw []string

	if e.Last > 0 {
		keeps = append(keeps, fmt.Sprintf("%d latest", e.Last))
	}
	if e.Hourly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d hourly", e.Hourly))
	}
	if e.Daily > 0 {
		keeps = append(keeps, fmt.Sprintf("%d daily", e.Daily))
	}
	if e.Weekly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d weekly", e.Weekly))
	}
	if e.Monthly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d monthly", e.Monthly))
	}
	if e.Yearly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d yearly", e.Yearly))
	}

	if !e.WithinHourly.Zero() {
		keepw = append(keepw, fmt.Sprintf("hourly snapshots within %v", e.WithinHourly))
	}

	if !e.WithinDaily.Zero() {
		keepw = append(keepw, fmt.Sprintf("daily snapshots within %v", e.WithinDaily))
	}

	if !e.WithinWeekly.Zero() {
		keepw = append(keepw, fmt.Sprintf("weekly snapshots within %v", e.WithinWeekly))
	}

	if !e.WithinMonthly.Zero() {
		keepw = append(keepw, fmt.Sprintf("monthly snapshots within %v", e.WithinMonthly))
	}

	if !e.WithinYearly.Zero() {
		keepw = append(keepw, fmt.Sprintf("yearly snapshots within %v", e.WithinYearly))
	}

	if len(keeps) > 0 {
		s = fmt.Sprintf("%s snapshots", strings.Join(keeps, ", "))
	}

	if len(keepw) > 0 {
		if s != "" {
			s += ", "
		}
		s += strings.Join(keepw, ", ")
	}

	if len(e.Tags) > 0 {
		if s != "" {
			s += " and "
		}
		s += fmt.Sprintf("all snapshots with tags %s", e.Tags)
	}

	if !e.Within.Zero() {
		if s != "" {
			s += " and "
		}
		s += fmt.Sprintf("all snapshots within %s of the newest", e.Within)
	}

	s = "keep " + s

	return s
}

// Sum returns the maximum number of snapshots to be kept according to this
// policy.
func (e ExpirePolicy) Sum() int {
	return e.Last + e.Hourly + e.Daily + e.Weekly + e.Monthly + e.Yearly
}

// Empty returns true iff no policy has been configured (all values zero).
func (e ExpirePolicy) Empty() bool {
	if len(e.Tags) != 0 {
		return false
	}

	empty := ExpirePolicy{Tags: e.Tags}
	return reflect.DeepEqual(e, empty)
}

// ymdh returns an integer in the form YYYYMMDDHH.
func ymdh(d time.Time, _ int) int {
	return d.Year()*1000000 + int(d.Month())*10000 + d.Day()*100 + d.Hour()
}

// ymd returns an integer in the form YYYYMMDD.
func ymd(d time.Time, _ int) int {
	return d.Year()*10000 + int(d.Month())*100 + d.Day()
}

// yw returns an integer in the form YYYYWW, where WW is the week number.
func yw(d time.Time, _ int) int {
	year, week := d.ISOWeek()
	return year*100 + week
}

// ym returns an integer in the form YYYYMM.
func ym(d time.Time, _ int) int {
	return d.Year()*100 + int(d.Month())
}

// y returns the year of d.
func y(d time.Time, _ int) int {
	return d.Year()
}

// always returns a unique number for d.
func always(d time.Time, nr int) int {
	return nr
}

// findLatestTimestamp returns the time stamp for the newest snapshot.
func findLatestTimestamp(list Snapshots) time.Time {
	if len(list) == 0 {
		panic("list of snapshots is empty")
	}

	var latest time.Time
	for _, sn := range list {
		if sn.Time.After(latest) {
			latest = sn.Time
		}
	}

	return latest
}

// KeepReason specifies why a particular snapshot was kept, and the counters at
// that point in the policy evaluation.
type KeepReason struct {
	Snapshot *Snapshot `json:"snapshot"`

	// description text which criteria match, e.g. "daily", "monthly"
	Matches []string `json:"matches"`

	// the counters after evaluating the current snapshot
	Counters struct {
		Last    int `json:"last,omitempty"`
		Hourly  int `json:"hourly,omitempty"`
		Daily   int `json:"daily,omitempty"`
		Weekly  int `json:"weekly,omitempty"`
		Monthly int `json:"monthly,omitempty"`
		Yearly  int `json:"yearly,omitempty"`
	} `json:"counters"`
}

// ApplyPolicy returns the snapshots from list that are to be kept and removed
// according to the policy p. list is sorted in the process. reasons contains
// the reasons to keep each snapshot, it is in the same order as keep.
func ApplyPolicy(list Snapshots, p ExpirePolicy) (keep, remove Snapshots, reasons []KeepReason) {
	sort.Sort(list)

	if p.Empty() {
		for _, sn := range list {
			reasons = append(reasons, KeepReason{
				Snapshot: sn,
				Matches:  []string{"policy is empty"},
			})
		}
		return list, remove, reasons
	}

	if len(list) == 0 {
		return list, nil, nil
	}

	// These buckets are for keeping last n snapshots of given type
	var buckets = [6]struct {
		Count  int
		bucker func(d time.Time, nr int) int
		Last   int
		reason string
	}{
		{p.Last, always, -1, "last snapshot"},
		{p.Hourly, ymdh, -1, "hourly snapshot"},
		{p.Daily, ymd, -1, "daily snapshot"},
		{p.Weekly, yw, -1, "weekly snapshot"},
		{p.Monthly, ym, -1, "monthly snapshot"},
		{p.Yearly, y, -1, "yearly snapshot"},
	}

	// These buckets are for keeping snapshots of given type within duration
	var bucketsWithin = [5]struct {
		Within Duration
		bucker func(d time.Time, nr int) int
		Last   int
		reason string
	}{
		{p.WithinHourly, ymdh, -1, "hourly within"},
		{p.WithinDaily, ymd, -1, "daily within"},
		{p.WithinWeekly, yw, -1, "weekly within"},
		{p.WithinMonthly, ym, -1, "monthly within"},
		{p.WithinYearly, y, -1, "yearly within"},
	}

	latest := findLatestTimestamp(list)

	for nr, cur := range list {
		var keepSnap bool
		var keepSnapReasons []string

		// Tags are handled specially as they are not counted.
		for _, l := range p.Tags {
			if cur.HasTags(l) {
				keepSnap = true
				keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("has tags %v", l))
			}
		}

		// If the timestamp of the snapshot is within the range, then keep it.
		if !p.Within.Zero() {
			t := latest.AddDate(-p.Within.Years, -p.Within.Months, -p.Within.Days).Add(time.Hour * time.Duration(-p.Within.Hours))
			if cur.Time.After(t) {
				keepSnap = true
				keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("within %v", p.Within))
			}
		}

		// Now update the other buckets and see if they have some counts left.
		for i, b := range buckets {
			if b.Count > 0 {
				val := b.bucker(cur.Time, nr)
				if val != b.Last {
					debug.Log("keep %v %v, bucker %v, val %v\n", cur.Time, cur.id.Str(), i, val)
					keepSnap = true
					buckets[i].Last = val
					buckets[i].Count--
					keepSnapReasons = append(keepSnapReasons, b.reason)
				}
			}
		}

		// If the timestamp is within range, and the snapshot is an hourly/daily/weekly/monthly/yearly snapshot, then keep it
		for i, b := range bucketsWithin {
			if !b.Within.Zero() {
				t := latest.AddDate(-b.Within.Years, -b.Within.Months, -b.Within.Days).Add(time.Hour * time.Duration(-b.Within.Hours))

				if cur.Time.After(t) {
					val := b.bucker(cur.Time, nr)
					if val != b.Last {
						debug.Log("keep %v, time %v, ID %v, bucker %v, val %v %v\n", b.reason, cur.Time, cur.id.Str(), i, val, b.Last)
						keepSnap = true
						bucketsWithin[i].Last = val
						keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("%v %v", b.reason, b.Within))
					}
				}
			}
		}

		if keepSnap {
			keep = append(keep, cur)
			kr := KeepReason{
				Snapshot: cur,
				Matches:  keepSnapReasons,
			}
			kr.Counters.Last = buckets[0].Count
			kr.Counters.Hourly = buckets[1].Count
			kr.Counters.Daily = buckets[2].Count
			kr.Counters.Weekly = buckets[3].Count
			kr.Counters.Monthly = buckets[4].Count
			kr.Counters.Yearly = buckets[5].Count
			reasons = append(reasons, kr)
		} else {
			remove = append(remove, cur)
		}
	}

	return keep, remove, reasons
}
<commit_msg>forget: Ensure future snapshots do not affect --keep-within-*<commit_after>package restic

import (
	"fmt"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/restic/restic/internal/debug"
)

// ExpirePolicy configures which snapshots should be automatically removed.
type ExpirePolicy struct {
	Last          int       // keep the last n snapshots
	Hourly        int       // keep the last n hourly snapshots
	Daily         int       // keep the last n daily snapshots
	Weekly        int       // keep the last n weekly snapshots
	Monthly       int       // keep the last n monthly snapshots
	Yearly        int       // keep the last n yearly snapshots
	Within        Duration  // keep snapshots made within this duration
	WithinHourly  Duration  // keep hourly snapshots made within this duration
	WithinDaily   Duration  // keep daily snapshots made within this duration
	WithinWeekly  Duration  // keep weekly snapshots made within this duration
	WithinMonthly Duration  // keep monthly snapshots made within this duration
	WithinYearly  Duration  // keep yearly snapshots made within this duration
	Tags          []TagList // keep all snapshots that include at least one of the tag lists.
}

func (e ExpirePolicy) String() (s string) {
	var keeps []string
	var keepw []string

	if e.Last > 0 {
		keeps = append(keeps, fmt.Sprintf("%d latest", e.Last))
	}
	if e.Hourly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d hourly", e.Hourly))
	}
	if e.Daily > 0 {
		keeps = append(keeps, fmt.Sprintf("%d daily", e.Daily))
	}
	if e.Weekly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d weekly", e.Weekly))
	}
	if e.Monthly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d monthly", e.Monthly))
	}
	if e.Yearly > 0 {
		keeps = append(keeps, fmt.Sprintf("%d yearly", e.Yearly))
	}

	if !e.WithinHourly.Zero() {
		keepw = append(keepw, fmt.Sprintf("hourly snapshots within %v", e.WithinHourly))
	}

	if !e.WithinDaily.Zero() {
		keepw = append(keepw, fmt.Sprintf("daily snapshots within %v", e.WithinDaily))
	}

	if !e.WithinWeekly.Zero() {
		keepw = append(keepw, fmt.Sprintf("weekly snapshots within %v", e.WithinWeekly))
	}

	if !e.WithinMonthly.Zero() {
		keepw = append(keepw, fmt.Sprintf("monthly snapshots within %v", e.WithinMonthly))
	}

	if !e.WithinYearly.Zero() {
		keepw = append(keepw, fmt.Sprintf("yearly snapshots within %v", e.WithinYearly))
	}

	if len(keeps) > 0 {
		s = fmt.Sprintf("%s snapshots", strings.Join(keeps, ", "))
	}

	if len(keepw) > 0 {
		if s != "" {
			s += ", "
		}
		s += strings.Join(keepw, ", ")
	}

	if len(e.Tags) > 0 {
		if s != "" {
			s += " and "
		}
		s += fmt.Sprintf("all snapshots with tags %s", e.Tags)
	}

	if !e.Within.Zero() {
		if s != "" {
			s += " and "
		}
		s += fmt.Sprintf("all snapshots within %s of the newest", e.Within)
	}

	s = "keep " + s

	return s
}

// Sum returns the maximum number of snapshots to be kept according to this
// policy.
func (e ExpirePolicy) Sum() int {
	return e.Last + e.Hourly + e.Daily + e.Weekly + e.Monthly + e.Yearly
}

// Empty returns true iff no policy has been configured (all values zero).
func (e ExpirePolicy) Empty() bool {
	if len(e.Tags) != 0 {
		return false
	}

	empty := ExpirePolicy{Tags: e.Tags}
	return reflect.DeepEqual(e, empty)
}

// ymdh returns an integer in the form YYYYMMDDHH.
func ymdh(d time.Time, _ int) int {
	return d.Year()*1000000 + int(d.Month())*10000 + d.Day()*100 + d.Hour()
}

// ymd returns an integer in the form YYYYMMDD.
func ymd(d time.Time, _ int) int {
	return d.Year()*10000 + int(d.Month())*100 + d.Day()
}

// yw returns an integer in the form YYYYWW, where WW is the week number.
func yw(d time.Time, _ int) int {
	year, week := d.ISOWeek()
	return year*100 + week
}

// ym returns an integer in the form YYYYMM.
func ym(d time.Time, _ int) int {
	return d.Year()*100 + int(d.Month())
}

// y returns the year of d.
func y(d time.Time, _ int) int {
	return d.Year()
}

// always returns a unique number for d.
func always(d time.Time, nr int) int {
	return nr
}

// findLatestTimestamp returns the time stamp for the latest (newest) snapshot,
// for use with policies based on time relative to latest.
func findLatestTimestamp(list Snapshots) time.Time {
	if len(list) == 0 {
		panic("list of snapshots is empty")
	}

	var latest time.Time
	now := time.Now()
	for _, sn := range list {
		// Find the latest snapshot in the list
		// The latest snapshot must, however, not be in the future.
		if sn.Time.After(latest) && sn.Time.Before(now) {
			latest = sn.Time
		}
	}

	return latest
}

// KeepReason specifies why a particular snapshot was kept, and the counters at
// that point in the policy evaluation.
type KeepReason struct {
	Snapshot *Snapshot `json:"snapshot"`

	// description text which criteria match, e.g. "daily", "monthly"
	Matches []string `json:"matches"`

	// the counters after evaluating the current snapshot
	Counters struct {
		Last    int `json:"last,omitempty"`
		Hourly  int `json:"hourly,omitempty"`
		Daily   int `json:"daily,omitempty"`
		Weekly  int `json:"weekly,omitempty"`
		Monthly int `json:"monthly,omitempty"`
		Yearly  int `json:"yearly,omitempty"`
	} `json:"counters"`
}

// ApplyPolicy returns the snapshots from list that are to be kept and removed
// according to the policy p. list is sorted in the process. reasons contains
// the reasons to keep each snapshot, it is in the same order as keep.
func ApplyPolicy(list Snapshots, p ExpirePolicy) (keep, remove Snapshots, reasons []KeepReason) {
	sort.Sort(list)

	if p.Empty() {
		for _, sn := range list {
			reasons = append(reasons, KeepReason{
				Snapshot: sn,
				Matches:  []string{"policy is empty"},
			})
		}
		return list, remove, reasons
	}

	if len(list) == 0 {
		return list, nil, nil
	}

	// These buckets are for keeping last n snapshots of given type
	var buckets = [6]struct {
		Count  int
		bucker func(d time.Time, nr int) int
		Last   int
		reason string
	}{
		{p.Last, always, -1, "last snapshot"},
		{p.Hourly, ymdh, -1, "hourly snapshot"},
		{p.Daily, ymd, -1, "daily snapshot"},
		{p.Weekly, yw, -1, "weekly snapshot"},
		{p.Monthly, ym, -1, "monthly snapshot"},
		{p.Yearly, y, -1, "yearly snapshot"},
	}

	// These buckets are for keeping snapshots of given type within duration
	var bucketsWithin = [5]struct {
		Within Duration
		bucker func(d time.Time, nr int) int
		Last   int
		reason string
	}{
		{p.WithinHourly, ymdh, -1, "hourly within"},
		{p.WithinDaily, ymd, -1, "daily within"},
		{p.WithinWeekly, yw, -1, "weekly within"},
		{p.WithinMonthly, ym, -1, "monthly within"},
		{p.WithinYearly, y, -1, "yearly within"},
	}

	latest := findLatestTimestamp(list)

	for nr, cur := range list {
		var keepSnap bool
		var keepSnapReasons []string

		// Tags are handled specially as they are not counted.
		for _, l := range p.Tags {
			if cur.HasTags(l) {
				keepSnap = true
				keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("has tags %v", l))
			}
		}

		// If the timestamp of the snapshot is within the range, then keep it.
		if !p.Within.Zero() {
			t := latest.AddDate(-p.Within.Years, -p.Within.Months, -p.Within.Days).Add(time.Hour * time.Duration(-p.Within.Hours))
			if cur.Time.After(t) {
				keepSnap = true
				keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("within %v", p.Within))
			}
		}

		// Now update the other buckets and see if they have some counts left.
		for i, b := range buckets {
			if b.Count > 0 {
				val := b.bucker(cur.Time, nr)
				if val != b.Last {
					debug.Log("keep %v %v, bucker %v, val %v\n", cur.Time, cur.id.Str(), i, val)
					keepSnap = true
					buckets[i].Last = val
					buckets[i].Count--
					keepSnapReasons = append(keepSnapReasons, b.reason)
				}
			}
		}

		// If the timestamp is within range, and the snapshot is an hourly/daily/weekly/monthly/yearly snapshot, then keep it
		for i, b := range bucketsWithin {
			if !b.Within.Zero() {
				t := latest.AddDate(-b.Within.Years, -b.Within.Months, -b.Within.Days).Add(time.Hour * time.Duration(-b.Within.Hours))

				if cur.Time.After(t) {
					val := b.bucker(cur.Time, nr)
					if val != b.Last {
						debug.Log("keep %v, time %v, ID %v, bucker %v, val %v %v\n", b.reason, cur.Time, cur.id.Str(), i, val, b.Last)
						keepSnap = true
						bucketsWithin[i].Last = val
						keepSnapReasons = append(keepSnapReasons, fmt.Sprintf("%v %v", b.reason, b.Within))
					}
				}
			}
		}

		if keepSnap {
			keep = append(keep, cur)
			kr := KeepReason{
				Snapshot: cur,
				Matches:  keepSnapReasons,
			}
			kr.Counters.Last = buckets[0].Count
			kr.Counters.Hourly = buckets[1].Count
			kr.Counters.Daily = buckets[2].Count
			kr.Counters.Weekly = buckets[3].Count
			kr.Counters.Monthly = buckets[4].Count
			kr.Counters.Yearly = buckets[5].Count
			reasons = append(reasons, kr)
		} else {
			remove = append(remove, cur)
		}
	}

	return keep, remove, reasons
}
<|endoftext|>


<commit_before>// Copyright 2017 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// Package handlers 用于处理节点下请求方法与处理函数的对应关系
package handlers

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/issue9/mux/internal/method"
)

type optionsState int8

const (
	optionsStateDefault      optionsState = iota // 默认情况
	optionsStateFixedString                      // 设置为固定的字符串
	optionsStateFixedHandler                     // 设置为固定的 http.Handler
	optionsStateDisable                          // 禁用
)

// Handlers 用于表示某节点下各个请求方法对应的处理函数。
type Handlers struct {
	handlers     map[method.Type]http.Handler // 请求方法及其对应的 http.Handler
	optionsAllow string                       // 缓存的 OPTIONS 请求头的 allow 报头内容。
	optionsState optionsState                 // OPTIONS 报头的处理方式
}

// New 声明一个新的 Handlers 实例
func New() *Handlers {
	ret := &Handlers{
		handlers:     make(map[method.Type]http.Handler, 4), // 大部分不会超过 4 条数据
		optionsState: optionsStateDefault,
	}

	// 添加默认的 OPTIONS 请求内容
	ret.handlers[method.Options] = http.HandlerFunc(ret.optionsServeHTTP)
	ret.optionsAllow = ret.getOptionsAllow()

	return ret
}

// Add 添加一个处理函数
func (hs *Handlers) Add(h http.Handler, methods ...string) error {
	for _, m := range methods {
		i := method.Int(m)
		if i == method.None {
			return fmt.Errorf("不支持的请求方法 %v", m)
		}

		if err := hs.addSingle(h, i); err != nil {
			return err
		}
	}

	return nil
}

func (hs *Handlers) addSingle(h http.Handler, m method.Type) error {
	if m == method.Options { // 强制修改 OPTIONS 方法的处理方式
		if hs.optionsState == optionsStateFixedHandler { // 被强制修改过，不能再受理。
			return errors.New("该请求方法 OPTIONS 已经存在")
		}

		hs.handlers[method.Options] = h
		hs.optionsState = optionsStateFixedHandler
		return nil
	}

	// 非 OPTIONS 请求
	if _, found := hs.handlers[m]; found {
		return fmt.Errorf("该请求方法 %v 已经存在", m.String())
	}
	hs.handlers[m] = h

	// 重新生成 optionsAllow 字符串
	if hs.optionsState == optionsStateDefault {
		hs.optionsAllow = hs.getOptionsAllow()
	}
	return nil
}

func (hs *Handlers) optionsServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Allow", hs.optionsAllow)
}

func (hs *Handlers) getOptionsAllow() string {
	var index method.Type
	for method := range hs.handlers {
		index += method
	}
	return optionsStrings[index]
}

// Remove 移除某个请求方法对应的处理函数
func (hs *Handlers) Remove(methods ...string) bool {
	for _, m := range methods {
		mm := method.Int(m)
		delete(hs.handlers, mm)
		if mm == method.Options { // 明确指出要删除该路由项的 OPTIONS 时，表示禁止
			hs.optionsState = optionsStateDisable
		}
	}

	// 删完了
	if hs.Len() == 0 {
		hs.optionsAllow = ""
		return true
	}

	// 只有一个 OPTIONS 了，且未经外界强制修改，则将其也一并删除。
	if hs.Len() == 1 &&
		hs.handlers[method.Options] != nil &&
		hs.optionsState == optionsStateDefault {
		delete(hs.handlers, method.Options)
		hs.optionsAllow = ""
		return true
	}

	if hs.optionsState == optionsStateDefault {
		hs.optionsAllow = hs.getOptionsAllow()
	}
	return false
}

// SetAllow 设置 Options 请求头的 Allow 报头。
func (hs *Handlers) SetAllow(optionsAllow string) {
	hs.optionsAllow = optionsAllow
	hs.optionsState = optionsStateFixedString
}

// Handler 获取指定方法对应的处理函数
func (hs *Handlers) Handler(m string) http.Handler {
	return hs.handlers[method.Int(m)]
}

// Len 获取当前支持请求方法数量
func (hs *Handlers) Len() int {
	return len(hs.handlers)
}
<commit_msg>更下文档错误<commit_after>// Copyright 2017 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

// Package handlers 用于处理节点下请求方法与处理函数的对应关系
package handlers

import (
	"errors"
	"fmt"
	"net/http"

	"github.com/issue9/mux/internal/method"
)

type optionsState int8

const (
	optionsStateDefault      optionsState = iota // 默认情况
	optionsStateFixedString                      // 设置为固定的字符串
	optionsStateFixedHandler                     // 设置为固定的 http.Handler
	optionsStateDisable                          // 禁用
)

// Handlers 用于表示某节点下各个请求方法对应的处理函数。
type Handlers struct {
	handlers     map[method.Type]http.Handler // 请求方法及其对应的 http.Handler
	optionsAllow string                       // 缓存的 OPTIONS 请求头的 allow 报头内容。
	optionsState optionsState                 // OPTIONS 报头的处理方式
}

// New 声明一个新的 Handlers 实例
func New() *Handlers {
	ret := &Handlers{
		handlers:     make(map[method.Type]http.Handler, 4), // 大部分不会超过 4 条数据
		optionsState: optionsStateDefault,
	}

	// 添加默认的 OPTIONS 请求内容
	ret.handlers[method.Options] = http.HandlerFunc(ret.optionsServeHTTP)
	ret.optionsAllow = ret.getOptionsAllow()

	return ret
}

// Add 添加一个处理函数
func (hs *Handlers) Add(h http.Handler, methods ...string) error {
	for _, m := range methods {
		i := method.Int(m)
		if i == method.None {
			return fmt.Errorf("不支持的请求方法 %v", m)
		}

		if err := hs.addSingle(h, i); err != nil {
			return err
		}
	}

	return nil
}

func (hs *Handlers) addSingle(h http.Handler, m method.Type) error {
	if m == method.Options { // 强制修改 OPTIONS 方法的处理方式
		if hs.optionsState == optionsStateFixedHandler { // 被强制修改过，不能再受理。
			return errors.New("该请求方法 OPTIONS 已经存在")
		}

		hs.handlers[method.Options] = h
		hs.optionsState = optionsStateFixedHandler
		return nil
	}

	// 非 OPTIONS 请求
	if _, found := hs.handlers[m]; found {
		return fmt.Errorf("该请求方法 %v 已经存在", m.String())
	}
	hs.handlers[m] = h

	// 重新生成 optionsAllow 字符串
	if hs.optionsState == optionsStateDefault {
		hs.optionsAllow = hs.getOptionsAllow()
	}
	return nil
}

func (hs *Handlers) optionsServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Allow", hs.optionsAllow)
}

func (hs *Handlers) getOptionsAllow() string {
	var index method.Type
	for method := range hs.handlers {
		index += method
	}
	return optionsStrings[index]
}

// Remove 移除某个请求方法对应的处理函数。
// 返回值表示是否已经被清空。
func (hs *Handlers) Remove(methods ...string) bool {
	for _, m := range methods {
		mm := method.Int(m)
		delete(hs.handlers, mm)
		if mm == method.Options { // 明确指出要删除该路由项的 OPTIONS 时，表示禁止
			hs.optionsState = optionsStateDisable
		}
	}

	// 删完了
	if hs.Len() == 0 {
		hs.optionsAllow = ""
		return true
	}

	// 只有一个 OPTIONS 了，且未经外界强制修改，则将其也一并删除。
	if hs.Len() == 1 &&
		hs.handlers[method.Options] != nil &&
		hs.optionsState == optionsStateDefault {
		delete(hs.handlers, method.Options)
		hs.optionsAllow = ""
		return true
	}

	if hs.optionsState == optionsStateDefault {
		hs.optionsAllow = hs.getOptionsAllow()
	}
	return false
}

// SetAllow 设置 Options 请求头的 Allow 报头。
func (hs *Handlers) SetAllow(optionsAllow string) {
	hs.optionsAllow = optionsAllow
	hs.optionsState = optionsStateFixedString
}

// Handler 获取指定方法对应的处理函数
func (hs *Handlers) Handler(m string) http.Handler {
	return hs.handlers[method.Int(m)]
}

// Len 获取当前支持请求方法数量
func (hs *Handlers) Len() int {
	return len(hs.handlers)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"regexp"
	"encoding/json"

	"github.com/valyala/fasthttp"
)

type Replace struct {
	regex   *regexp.Regexp
	replace []byte
}

func NewReplace(regex, replace string) (result Replace) {
	result.regex = regexp.MustCompile(regex)
	result.replace = []byte(replace)
	return
}

func (v *Replace) apply(data []byte) []byte {
	return v.regex.ReplaceAll(data, v.replace)
}

// Variables for api proxying # api.vk.com
var apiProxy = &fasthttp.HostClient{
	Addr:  "api.vk.com:443",
	IsTLS: true,
}
var apiReplaces []Replace
var apiLongpollPath = []byte("/method/execute")
var apiLongpollReplace Replace
var apiNewsfeedPath = []byte("/method/execute.getNewsfeedSmart")

// Variables for site proxying # vk.com
var siteProxy = &fasthttp.HostClient{
	Addr:  "vk.com:443",
	IsTLS: true,
}
var siteHlsReplace Replace
var sitePath = []byte("/vk.com")
var siteHost = []byte("vk.com")

func InitReplaces() {
	apiReplaces = []Replace{
		NewReplace(`"https:\\/\\/(pu\.vk\.com|[-a-z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video|audio)\.(?:net|com)))\\/([^"]+)`, `"https:\/\/`+config.domain+`\/_\/$1\/$2`),
		NewReplace(`"https:\\/\\/vk\.com\\/(video_hls\.php[^"]+)`, `"https:\/\/`+config.domain+`\/vk.com\/$1`),
		NewReplace(`"https:\\/\\/vk\.com\\/((images|doc[0-9]+_)[^"]+)`, `"https:\/\/`+config.domain+`\/_\/vk.com\/$1`),
		NewReplace(`"preview_url":"https:\\/\\/m\.vk\.com\\/(article[0-9]+)[^"]+"(,"preview_page":"[^"]+",?)?`, ``),
	}
	apiLongpollReplace = NewReplace(`"server":"api.vk.com\\/newuim`, `"server":"`+config.domain+`\/newuim`)

	siteHlsReplace = NewReplace(`https:\/\/([-a-z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video)\.(?:net|com)))\/`, `https://`+config.domain+`/_/$1/`)
}

func reverseProxyHandler(ctx *fasthttp.RequestCtx) {
	trackRequestStart(ctx)

	req := &ctx.Request
	res := &ctx.Response
	proxyClient := preRequest(req)
	if err := proxyClient.Do(req, res); err != nil {
		ctx.Logger().Printf("error when proxying the request: %s", err)
	} else {
		postResponse(req.URI(), res)
	}
}

func preRequest(req *fasthttp.Request) (client *fasthttp.HostClient) {
	path := req.URI().Path()
	if bytes.HasPrefix(req.URI().Path(), sitePath) {
		client = siteProxy
		req.SetRequestURIBytes([]byte(path[7:]))
		req.SetHost("vk.com")
	} else {
		client = apiProxy
		req.SetHost("api.vk.com")
	}
	req.Header.Del("Accept-Encoding")
	return
}

func postResponse(uri *fasthttp.URI, res *fasthttp.Response) {
	res.Header.Del("Set-Cookie")
	body := res.Body()
	if bytes.Compare(uri.Host(), siteHost) == 0 {
		body = siteHlsReplace.apply(body)
	} else {
		for _, replace := range apiReplaces {
			body = replace.apply(body)
		}

		if bytes.Compare(uri.Path(), apiLongpollPath) == 0 {
			body = apiLongpollReplace.apply(body)
		} else

		// Clear feed from SPAM
		if bytes.Compare(uri.Path(), apiNewsfeedPath) == 0 {
			var parsed map[string]interface{}
			if err := json.Unmarshal(body, &parsed); err == nil {
				removed := 0
				if parsed["response"] != nil {
					response := parsed["response"].(map[string]interface{})
					if response["items"] != nil {
						items := response["items"].([]interface{})
						for i := len(items) - 1; i >= 0; i-- {
							post := items[i].(map[string]interface{})
							if post["type"] == "ads" || (post["type"] == "post" && post["marked_as_ads"] != nil && post["marked_as_ads"].(float64) == 1) {
								items[i] = items[len(items)-1]
								items[len(items)-1] = nil
								items = items[:len(items)-1]
								removed++
							}
						}
						if removed > 0 {
							newItems := make([]interface{}, len(items))
							copy(newItems, items)
							response["items"] = newItems
						}
					}
				}
				if removed > 0 {
					body, err = json.Marshal(parsed)
				}
			}
		}
	}
	res.SetBody(body)

	trackRequestEnd(len(body))
}
<commit_msg>Исправлены видеоролики<commit_after>package main

import (
	"bytes"
	"regexp"
	"encoding/json"

	"github.com/valyala/fasthttp"
)

type Replace struct {
	regex   *regexp.Regexp
	replace []byte
}

func NewReplace(regex, replace string) (result Replace) {
	result.regex = regexp.MustCompile(regex)
	result.replace = []byte(replace)
	return
}

func (v *Replace) apply(data []byte) []byte {
	return v.regex.ReplaceAll(data, v.replace)
}

// Variables for api proxying # api.vk.com
var apiProxy = &fasthttp.HostClient{
	Addr:  "api.vk.com:443",
	IsTLS: true,
}
var apiReplaces []Replace
var apiLongpollPath = []byte("/method/execute")
var apiLongpollReplace Replace
var apiNewsfeedPath = []byte("/method/execute.getNewsfeedSmart")

// Variables for site proxying # vk.com
var siteProxy = &fasthttp.HostClient{
	Addr:  "vk.com:443",
	IsTLS: true,
}
var siteHlsReplace Replace
var sitePath = []byte("/vk.com")
var siteHost = []byte("vk.com")

func InitReplaces() {
	apiReplaces = []Replace{
		NewReplace(`"https:\\/\\/(pu\.vk\.com|[-a-z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video|audio)\.(?:net|com)))\\/([^"]+)`, `"https:\/\/`+config.domain+`\/_\/$1\/$2`),
		NewReplace(`"https:\\/\\/vk\.com\\/(video_hls\.php[^"]+)`, `"https:\/\/`+config.domain+`\/vk.com\/$1`),
		NewReplace(`"https:\\/\\/vk\.com\\/((images|doc[0-9]+_)[^"]+)`, `"https:\/\/`+config.domain+`\/_\/vk.com\/$1`),
		NewReplace(`"preview_url":"https:\\/\\/m\.vk\.com\\/(article[0-9]+)[^"]+"(,"preview_page":"[^"]+",?)?`, ``),
	}
	apiLongpollReplace = NewReplace(`"server":"api.vk.com\\/newuim`, `"server":"`+config.domain+`\/newuim`)

	siteHlsReplace = NewReplace(`https:\/\/([-a-z0-9]+\.(?:userapi\.com|vk-cdn\.net|vk\.me|vkuser(?:live|video)\.(?:net|com)))\/`, `https://`+config.domain+`/_/$1/`)
}

func reverseProxyHandler(ctx *fasthttp.RequestCtx) {
	trackRequestStart(ctx)

	req := &ctx.Request
	res := &ctx.Response
	proxyClient := preRequest(req)
	if err := proxyClient.Do(req, res); err != nil {
		ctx.Logger().Printf("error when proxying the request: %s", err)
	} else {
		postResponse(ctx)
	}
}

func preRequest(req *fasthttp.Request) (client *fasthttp.HostClient) {
	path := req.RequestURI()
	if bytes.HasPrefix(req.URI().Path(), sitePath) {
		client = siteProxy
		req.SetRequestURIBytes([]byte(path[7:]))
		req.SetHost("vk.com")
	} else {
		client = apiProxy
		req.SetHost("api.vk.com")
	}
	req.Header.Del("Accept-Encoding")
	return
}

func postResponse(ctx *fasthttp.RequestCtx) {
	uri := ctx.Request.URI()
	res := &ctx.Response
	res.Header.Del("Set-Cookie")
	body := res.Body()
	if bytes.Compare(uri.Host(), siteHost) == 0 {
		body = siteHlsReplace.apply(body)
	} else {
		for _, replace := range apiReplaces {
			body = replace.apply(body)
		}

		if bytes.Compare(uri.Path(), apiLongpollPath) == 0 {
			body = apiLongpollReplace.apply(body)
		} else

		// Clear feed from SPAM
		if bytes.Compare(uri.Path(), apiNewsfeedPath) == 0 {
			var parsed map[string]interface{}
			if err := json.Unmarshal(body, &parsed); err == nil {
				removed := 0
				if parsed["response"] != nil {
					response := parsed["response"].(map[string]interface{})
					if response["items"] != nil {
						items := response["items"].([]interface{})
						for i := len(items) - 1; i >= 0; i-- {
							post := items[i].(map[string]interface{})
							if post["type"] == "ads" || (post["type"] == "post" && post["marked_as_ads"] != nil && post["marked_as_ads"].(float64) == 1) {
								items[i] = items[len(items)-1]
								items[len(items)-1] = nil
								items = items[:len(items)-1]
								removed++
							}
						}
						if removed > 0 {
							newItems := make([]interface{}, len(items))
							copy(newItems, items)
							response["items"] = newItems
						}
					}
				}
				if removed > 0 {
					body, err = json.Marshal(parsed)
				}
			}
		}
	}
	res.SetBody(body)

	trackRequestEnd(len(body))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/elazarl/goproxy"
	//"github.com/elazarl/goproxy/ext/html"
)

type ContextUserData struct {
	Store bool
	Time  int64
}

type Content struct {
	//Id bson.ObjectId
	Request  Request  "request"
	Response Response "response"
	Date    time.Time   "date"
}

type Request struct {
	Body    string      "body"
	//Date    time.Time   "date"
	Host    string      "host"
	Method  string      "method"
	Path    string      "path"
	Time    float32     "time"
	Headers http.Header "headers"
}

type Response struct {
	Body    string      "body"
	Headers http.Header "headers"
	Status  int         "status"
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	mongourl := flag.String("mongourl", "", "record request/response in mongodb")
	mock := flag.Bool("m", false, "send fake responses")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")

	flag.Parse()

	c := &mgo.Collection{}

	if len(*mongourl) != 0 {
		// Mongo DB connection
		session, err := mgo.Dial(*mongourl)
		if err != nil {
			panic(err)
		}
		defer session.Close()

		// Optional. Switch the session to a monotonic behavior.
		session.SetMode(mgo.Monotonic, true)

		c = session.DB("proxyservice").C("log_logentry")
	}

	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose

	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)

	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		log.Printf("Request: %s %s %s", req.Method, req.Host, req.RequestURI)

		host := req.Host
		// Request to domain--name-co-uk.mocky.dev
		// will be forwarded to domain-name.co.uk
		if strings.Contains(host, ".mocky.dev") {
			host = strings.Replace(host, ".mocky.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
			log.Printf("Target Host: %s", host)
			req.Host = host

		}

		if c != nil && *mock && req.Method != "CONNECT" {
			result := Content{}
			ctx.Logf("Looking for existing request")
			/*fmt.Println("RequestURI:", req.RequestURI)
			  fmt.Println("Path:", req.URL.Path)
			  fmt.Println("Host:", req.Host)
			  fmt.Println("Method:", req.Method)*/
			err := c.Find(bson.M{"request.host": req.Host, "request.method": req.Method, "response.status": 200, "request.path": req.URL.Path}).Sort("-date").One(&result)
			if err == nil {
				ctx.Logf("Found one")
				/*fmt.Println("Path:", result.Request.Path)
				  //fmt.Println("Body:", result.Request.Body)
				  fmt.Println("Method:", result.Request.Method)
				  fmt.Println("Host:", result.Request.Host)
				  fmt.Println("Time:", result.Request.Time)
				  fmt.Println("Date:", result.Request.Date)
				  fmt.Println("Headers:", result.Request.Headers)

				  //fmt.Println("Body:", result.Response.Body)
				  fmt.Println("Status:", result.Response.Status)
				  fmt.Println("Headers:", result.Response.Headers)*/

				resp := goproxy.NewResponse(req, goproxy.ContentTypeHtml, result.Response.Status, result.Response.Body)
				ctx.UserData = ContextUserData{Store: false, Time: 0}
				return req, resp
			}

		}

		ctx.UserData = ContextUserData{Store: true, Time: time.Now().UnixNano()}
		return req, nil
	})

	proxy.OnResponse().Do(goproxy.HandleBytes(
		func(b []byte, ctx *goproxy.ProxyCtx) []byte {
			ctx.Logf("Method: %s - host: %s", ctx.Resp.Request.Method, ctx.Resp.Request.Host)
			if c != nil && ctx.UserData != nil && ctx.UserData.(ContextUserData).Store && ctx.Resp.StatusCode == 200 && ctx.Resp.Request.Method != "CONNECT" {
				// Get Body as string
				s := string(b[:])
				contentType := ctx.Resp.Header.Get("Content-Type")
				ctx.Logf(contentType)
				arr := strings.Split(contentType, ";")
				switch arr[0] {
				case "application/json", "text/html":
					ctx.Logf("Recording both request/response")
					content := Content{
						//Id: bson.NewObjectId(),
						Request:  Request{Path: ctx.Resp.Request.URL.Path, Host: ctx.Resp.Request.Host, Method: ctx.Resp.Request.Method, Time: float32(time.Now().UnixNano()-ctx.UserData.(ContextUserData).Time) / 1.0e9, Headers: ctx.Resp.Request.Header},
						Response: Response{Status: ctx.Resp.StatusCode, Headers: ctx.Resp.Header, Body: s},
						Date:     time.Now(),
					}

					err := c.Insert(content)
					if err != nil {
						ctx.Logf("Can't insert document: %v\n", err)
					}
					ctx.Logf("Body: %s", s)
				}
			}

			return b
		}))

	log.Println("Starting Proxy")
	log.Fatalln(http.ListenAndServe(*addr, proxy))
}
<commit_msg>Use same content-type for response than in request<commit_after>package main

import (
	"flag"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/elazarl/goproxy"
	//"github.com/elazarl/goproxy/ext/html"
)

type ContextUserData struct {
	Store bool
	Time  int64
}

type Content struct {
	//Id bson.ObjectId
	Request  Request  "request"
	Response Response "response"
	Date    time.Time   "date"
}

type Request struct {
	Body    string      "body"
	//Date    time.Time   "date"
	Host    string      "host"
	Method  string      "method"
	Path    string      "path"
	Time    float32     "time"
	Headers http.Header "headers"
}

type Response struct {
	Body    string      "body"
	Headers http.Header "headers"
	Status  int         "status"
}

func main() {
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	mongourl := flag.String("mongourl", "", "record request/response in mongodb")
	mock := flag.Bool("m", false, "send fake responses")
	addr := flag.String("l", ":8080", "on which address should the proxy listen")

	flag.Parse()

	c := &mgo.Collection{}

	if len(*mongourl) != 0 {
		// Mongo DB connection
		session, err := mgo.Dial(*mongourl)
		if err != nil {
			panic(err)
		}
		defer session.Close()

		// Optional. Switch the session to a monotonic behavior.
		session.SetMode(mgo.Monotonic, true)

		c = session.DB("proxyservice").C("log_logentry")
	}

	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = *verbose

	proxy.OnRequest().HandleConnect(goproxy.AlwaysMitm)

	proxy.OnRequest().DoFunc(func(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
		log.Printf("Request: %s %s %s", req.Method, req.Host, req.RequestURI)

		host := req.Host
		// Request to domain--name-co-uk.mocky.dev
		// will be forwarded to domain-name.co.uk
		if strings.Contains(host, ".mocky.dev") {
			host = strings.Replace(host, ".mocky.dev", "", 1)
			host = strings.Replace(host, "-", ".", -1)
			host = strings.Replace(host, "..", "-", -1)
			log.Printf("Target Host: %s", host)
			req.Host = host

		}

		if c != nil && *mock && req.Method != "CONNECT" {
			result := Content{}
			ctx.Logf("Looking for existing request")
			/*fmt.Println("RequestURI:", req.RequestURI)
			  fmt.Println("Path:", req.URL.Path)
			  fmt.Println("Host:", req.Host)
			  fmt.Println("Method:", req.Method)*/
			err := c.Find(bson.M{"request.host": req.Host, "request.method": req.Method, "response.status": 200, "request.path": req.URL.Path}).Sort("-date").One(&result)
			if err == nil {
				ctx.Logf("Found one")
				/*fmt.Println("Path:", result.Request.Path)
				  //fmt.Println("Body:", result.Request.Body)
				  fmt.Println("Method:", result.Request.Method)
				  fmt.Println("Host:", result.Request.Host)
				  fmt.Println("Time:", result.Request.Time)
				  fmt.Println("Date:", result.Request.Date)
				  fmt.Println("Headers:", result.Request.Headers)

				  //fmt.Println("Body:", result.Response.Body)
				  fmt.Println("Status:", result.Response.Status)
				  fmt.Println("Headers:", result.Response.Headers)*/

				resp := goproxy.NewResponse(req, result.Response.Headers.Get("Content-Type"), result.Response.Status, result.Response.Body)
				ctx.UserData = ContextUserData{Store: false, Time: 0}
				return req, resp
			}

		}

		ctx.UserData = ContextUserData{Store: true, Time: time.Now().UnixNano()}
		return req, nil
	})

	proxy.OnResponse().Do(goproxy.HandleBytes(
		func(b []byte, ctx *goproxy.ProxyCtx) []byte {
			ctx.Logf("Method: %s - host: %s", ctx.Resp.Request.Method, ctx.Resp.Request.Host)
			if c != nil && ctx.UserData != nil && ctx.UserData.(ContextUserData).Store && ctx.Resp.StatusCode == 200 && ctx.Resp.Request.Method != "CONNECT" {
				// Get Body as string
				s := string(b[:])
				contentType := ctx.Resp.Header.Get("Content-Type")
				ctx.Logf(contentType)
				arr := strings.Split(contentType, ";")
				switch arr[0] {
				case "application/json", "text/html":
					ctx.Logf("Recording both request/response")
					content := Content{
						//Id: bson.NewObjectId(),
						Request:  Request{Path: ctx.Resp.Request.URL.Path, Host: ctx.Resp.Request.Host, Method: ctx.Resp.Request.Method, Time: float32(time.Now().UnixNano()-ctx.UserData.(ContextUserData).Time) / 1.0e9, Headers: ctx.Resp.Request.Header},
						Response: Response{Status: ctx.Resp.StatusCode, Headers: ctx.Resp.Header, Body: s},
						Date:     time.Now(),
					}

					err := c.Insert(content)
					if err != nil {
						ctx.Logf("Can't insert document: %v\n", err)
					}
					ctx.Logf("Body: %s", s)
				}
			}

			return b
		}))

	log.Println("Starting Proxy")
	log.Fatalln(http.ListenAndServe(*addr, proxy))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

// Structs for XML building
type Package struct {
	Xmlns   string     `xml:"xmlns,attr"`
	Types   []MetaType `xml:"types"`
	Version string     `xml:"version"`
}

type MetaType struct {
	Members []string `xml:"members"`
	Name    string   `xml:"name"`
}

func createPackage() Package {
	return Package{
		Version: strings.TrimPrefix(apiVersion, "v"),
		Xmlns:   "http://soap.sforce.com/2006/04/metadata",
	}
}

type metapath struct {
	path       string
	name       string
	hasFolder  bool
	onlyFolder bool
	extension  string
}

var metapaths = []metapath{
	metapath{path: "actionLinkGroupTemplates", name: "ActionLinkGroupTemplate"},
	metapath{path: "analyticSnapshots", name: "AnalyticSnapshot"},
	metapath{path: "applications", name: "CustomApplication"},
	metapath{path: "appMenus", name: "AppMenu"},
	metapath{path: "approvalProcesses", name: "ApprovalProcess"},
	metapath{path: "assignmentRules", name: "AssignmentRules"},
	metapath{path: "authproviders", name: "AuthProvider"},
	metapath{path: "aura", name: "AuraDefinitionBundle", hasFolder: true, onlyFolder: true},
	metapath{path: "autoResponseRules", name: "AutoResponseRules"},
	metapath{path: "callCenters", name: "CallCenter"},
	metapath{path: "cachePartitions", name: "PlatformCachePartition"},
	metapath{path: "certs", name: "Certificate"},
	metapath{path: "channelLayouts", name: "ChannelLayout"},
	metapath{path: "classes", name: "ApexClass"},
	metapath{path: "communities", name: "Community"},
	metapath{path: "components", name: "ApexComponent"},
	metapath{path: "connectedApps", name: "ConnectedApp"},
	metapath{path: "corsWhitelistOrigins", name: "CorsWhitelistOrigin"},
	metapath{path: "customApplicationComponents", name: "CustomApplicationComponent"},
	metapath{path: "customMetadata", name: "CustomMetadata"},
	metapath{path: "customPermissions", name: "CustomPermission"},
	metapath{path: "dashboards", name: "Dashboard", hasFolder: true},
	metapath{path: "dataSources", name: "ExternalDataSource"},
	metapath{path: "datacategorygroups", name: "DataCategoryGroup"},
	metapath{path: "delegateGroups", name: "DelegateGroup"},
	metapath{path: "documents", name: "Document", hasFolder: true},
	metapath{path: "EmbeddedServiceConfig", name: "EmbeddedServiceConfig"},
	metapath{path: "email", name: "EmailTemplate", hasFolder: true},
	metapath{path: "escalationRules", name: "EscalationRules"},
	metapath{path: "feedFilters", name: "CustomFeedFilter"},
	metapath{path: "flexipages", name: "FlexiPage"},
	metapath{path: "flowDefinitions", name: "FlowDefinition"},
	metapath{path: "flows", name: "Flow"},
	metapath{path: "globalPicklists", name: "GlobalPicklist"},
	metapath{path: "groups", name: "Group"},
	metapath{path: "homePageComponents", name: "HomePageComponent"},
	metapath{path: "homePageLayouts", name: "HomePageLayout"},
	metapath{path: "installedPackages", name: "InstalledPackage"},
	metapath{path: "labels", name: "CustomLabels"},
	metapath{path: "layouts", name: "Layout"},
	metapath{path: "LeadConvertSettings", name: "LeadConvertSettings"},
	metapath{path: "letterhead", name: "Letterhead"},
	metapath{path: "matchingRules", name: "MatchingRules"},
	metapath{path: "namedCredentials", name: "NamedCredential"},
	metapath{path: "objects", name: "CustomObject"},
	metapath{path: "objectTranslations", name: "CustomObjectTranslation"},
	metapath{path: "pages", name: "ApexPage"},
	metapath{path: "pathAssistants", name: "PathAssistant"},
	metapath{path: "permissionsets", name: "PermissionSet"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "profiles", name: "Profile", extension: ".profile"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "profiles", name: "Profile"},
	metapath{path: "queues", name: "Queue"},
	metapath{path: "quickActions", name: "QuickAction"},
	metapath{path: "remoteSiteSettings", name: "RemoteSiteSetting"},
	metapath{path: "reports", name: "Report", hasFolder: true},
	metapath{path: "reportTypes", name: "ReportType"},
	metapath{path: "roles", name: "Role"},
	metapath{path: "scontrols", name: "Scontrol"},
	metapath{path: "settings", name: "Settings"},
	metapath{path: "sharingRules", name: "SharingRules"},
	metapath{path: "siteDotComSites", name: "SiteDotCom"},
	metapath{path: "sites", name: "CustomSite"},
	metapath{path: "staticresources", name: "StaticResource"},
	metapath{path: "synonymDictionaries", name: "SynonymDictionary"},
	metapath{path: "tabs", name: "CustomTab"},
	metapath{path: "triggers", name: "ApexTrigger"},
	metapath{path: "weblinks", name: "CustomPageWebLink"},
	metapath{path: "workflows", name: "Workflow"},
}

type PackageBuilder struct {
	IsPush   bool
	Metadata map[string]MetaType
	Files    ForceMetadataFiles
}

func NewPushBuilder() PackageBuilder {
	pb := PackageBuilder{IsPush: true}
	pb.Metadata = make(map[string]MetaType)
	pb.Files = make(ForceMetadataFiles)

	return pb
}

func NewFetchBuilder() PackageBuilder {
	pb := PackageBuilder{IsPush: false}
	pb.Metadata = make(map[string]MetaType)
	pb.Files = make(ForceMetadataFiles)

	return pb
}

// Build and return package.xml
func (pb PackageBuilder) PackageXml() []byte {
	p := createPackage()

	for _, metaType := range pb.Metadata {
		p.Types = append(p.Types, metaType)
	}

	byteXml, _ := xml.MarshalIndent(p, "", "    ")
	byteXml = append([]byte(xml.Header), byteXml...)
	//if err := ioutil.WriteFile("mypackage.xml", byteXml, 0644); err != nil {
	//ErrorAndExit(err.Error())
	//}
	return byteXml
}

// Returns the full ForceMetadataFiles container
func (pb *PackageBuilder) ForceMetadataFiles() ForceMetadataFiles {
	pb.Files["package.xml"] = pb.PackageXml()
	return pb.Files
}

// Returns the source file path for a given metadata file path.
func MetaPathToSourcePath(mpath string) (spath string) {
	spath = strings.TrimSuffix(mpath, "-meta.xml")
	if spath == mpath {
		return
	}

	_, err := os.Stat(spath)
	if err != nil {
		spath = mpath
	}
	return
}

// Add a file to the builder
func (pb *PackageBuilder) AddFile(fpath string) (fname string, err error) {
	fpath, err = filepath.Abs(fpath)
	if err != nil {
		return
	}
	_, err = os.Stat(fpath)
	if err != nil {
		return
	}

	isDestructiveChanges, err := regexp.MatchString("destructiveChanges(Pre|Post)?"+regexp.QuoteMeta(".")+"xml", fpath)
	if err != nil {
		return
	}

	fpath = MetaPathToSourcePath(fpath)
	metaName, fname := getMetaTypeFromPath(fpath)
	if !isDestructiveChanges && !strings.HasSuffix(fpath, "-meta.xml") {
		pb.AddMetaToPackage(metaName, fname)
	}

	// If it's a push, we want to actually add the files
	if pb.IsPush {
		if isDestructiveChanges {
			err = pb.addDestructiveChanges(fpath)
		} else {
			err = pb.addFileToWorkingDir(metaName, fpath)
		}
	}

	return
}

// Adds the file to a temp directory for deploy
func (pb *PackageBuilder) addFileToWorkingDir(metaName string, fpath string) (err error) {
	// Get relative dir from source
	srcDir := filepath.Dir(filepath.Dir(fpath))
	for _, mp := range metapaths {
		if metaName == mp.name && mp.hasFolder {
			srcDir = filepath.Dir(srcDir)
		}
	}
	frel, _ := filepath.Rel(srcDir, fpath)

	// Try to find meta file
	hasMeta := true
	fmeta := fpath + "-meta.xml"
	fmetarel := ""
	if _, err = os.Stat(fmeta); err != nil {
		if os.IsNotExist(err) {
			hasMeta = false
		} else {
			// Has error
			return
		}
	} else {
		// Should be present since we worked back to srcDir
		fmetarel, _ = filepath.Rel(srcDir, fmeta)
	}

	fdata, err := ioutil.ReadFile(fpath)
	if err != nil {
		return
	}

	pb.Files[frel] = fdata
	if hasMeta {
		fdata, err = ioutil.ReadFile(fmeta)
		pb.Files[fmetarel] = fdata
		return
	}

	return
}

func (pb *PackageBuilder) addDestructiveChanges(fpath string) (err error) {
	fdata, err := ioutil.ReadFile(fpath)
	if err != nil {
		return
	}

	frel, _ := filepath.Rel(filepath.Dir(fpath), fpath)
	pb.Files[frel] = fdata

	return
}

func (pb *PackageBuilder) contains(members []string, name string) bool {
	for _, a := range members {
		if a == name {
			return true
		}
	}
	return false
}

// Adds a metadata name to the pending package
func (pb *PackageBuilder) AddMetaToPackage(metaName string, name string) {
	mt := pb.Metadata[metaName]
	if mt.Name == "" {
		mt.Name = metaName
	}

	if !pb.contains(mt.Members, name) {
		mt.Members = append(mt.Members, name)
		pb.Metadata[metaName] = mt
	}
}

// Gets metadata type name and target name from a file path
func getMetaTypeFromPath(fpath string) (metaName string, name string) {
	fpath, err := filepath.Abs(fpath)
	if err != nil {
		ErrorAndExit("Cound not find " + fpath)
	}
	if _, err := os.Stat(fpath); err != nil {
		ErrorAndExit("Cound not open " + fpath)
	}

	// Get the metadata type and name for the file
	metaName, fileName := getMetaForPath(fpath)
	name = strings.TrimSuffix(fileName, filepath.Ext(fileName))
	//name = strings.TrimSuffix(name, filepath.Ext(name))
	return
}

// Gets partial path based on a meta type name
func getPathForMeta(metaname string) string {
	for _, mp := range metapaths {
		if strings.EqualFold(mp.name, metaname) {
			return mp.path
		}
	}

	// Unknown, so use metaname
	return metaname
}

func findMetapathForFile(file string) (path metapath) {
	parentDir := filepath.Dir(file)
	parentName := filepath.Base(parentDir)
	grandparentName := filepath.Base(filepath.Dir(parentDir))
	fileExtension := filepath.Ext(file)

	for _, mp := range metapaths {
		if mp.hasFolder && grandparentName == mp.path {
			return mp
		}
		if mp.path == parentName {
			return mp
		}
	}

	// Hmm, maybe we can use the extension to determine the type
	for _, mp := range metapaths {
		if mp.extension == fileExtension {
			return mp
		}
	}
	return
}

// Gets meta type and name based on a path
func getMetaForPath(path string) (metaName string, objectName string) {
	parentDir := filepath.Dir(path)
	parentName := filepath.Base(parentDir)
	grandparentName := filepath.Base(filepath.Dir(parentDir))
	fileName := filepath.Base(path)

	for _, mp := range metapaths {
		if mp.hasFolder && grandparentName == mp.path {
			metaName = mp.name
			if mp.onlyFolder {
				objectName = parentName
			} else {
				objectName = parentName + "/" + fileName
			}
			return
		}
		if mp.path == parentName {
			metaName = mp.name
			objectName = fileName
			return
		}
	}

	// Unknown, so use path
	metaName = parentName
	objectName = fileName
	return
}
<commit_msg>added missing cspTrusted Site fix #401<commit_after>package main

import (
	"encoding/xml"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

// Structs for XML building
type Package struct {
	Xmlns   string     `xml:"xmlns,attr"`
	Types   []MetaType `xml:"types"`
	Version string     `xml:"version"`
}

type MetaType struct {
	Members []string `xml:"members"`
	Name    string   `xml:"name"`
}

func createPackage() Package {
	return Package{
		Version: strings.TrimPrefix(apiVersion, "v"),
		Xmlns:   "http://soap.sforce.com/2006/04/metadata",
	}
}

type metapath struct {
	path       string
	name       string
	hasFolder  bool
	onlyFolder bool
	extension  string
}

var metapaths = []metapath{
	metapath{path: "actionLinkGroupTemplates", name: "ActionLinkGroupTemplate"},
	metapath{path: "analyticSnapshots", name: "AnalyticSnapshot"},
	metapath{path: "applications", name: "CustomApplication"},
	metapath{path: "appMenus", name: "AppMenu"},
	metapath{path: "approvalProcesses", name: "ApprovalProcess"},
	metapath{path: "assignmentRules", name: "AssignmentRules"},
	metapath{path: "authproviders", name: "AuthProvider"},
	metapath{path: "aura", name: "AuraDefinitionBundle", hasFolder: true, onlyFolder: true},
	metapath{path: "autoResponseRules", name: "AutoResponseRules"},
	metapath{path: "callCenters", name: "CallCenter"},
	metapath{path: "cachePartitions", name: "PlatformCachePartition"},
	metapath{path: "certs", name: "Certificate"},
	metapath{path: "channelLayouts", name: "ChannelLayout"},
	metapath{path: "classes", name: "ApexClass"},
	metapath{path: "communities", name: "Community"},
	metapath{path: "components", name: "ApexComponent"},
	metapath{path: "connectedApps", name: "ConnectedApp"},
	metapath{path: "corsWhitelistOrigins", name: "CorsWhitelistOrigin"},
	metapath{path: "customApplicationComponents", name: "CustomApplicationComponent"},
	metapath{path: "customMetadata", name: "CustomMetadata"},
	metapath{path: "customPermissions", name: "CustomPermission"},
	metapath{path: "dashboards", name: "Dashboard", hasFolder: true},
	metapath{path: "dataSources", name: "ExternalDataSource"},
	metapath{path: "datacategorygroups", name: "DataCategoryGroup"},
	metapath{path: "delegateGroups", name: "DelegateGroup"},
	metapath{path: "documents", name: "Document", hasFolder: true},
	metapath{path: "EmbeddedServiceConfig", name: "EmbeddedServiceConfig"},
	metapath{path: "email", name: "EmailTemplate", hasFolder: true},
	metapath{path: "escalationRules", name: "EscalationRules"},
	metapath{path: "feedFilters", name: "CustomFeedFilter"},
	metapath{path: "flexipages", name: "FlexiPage"},
	metapath{path: "flowDefinitions", name: "FlowDefinition"},
	metapath{path: "flows", name: "Flow"},
	metapath{path: "globalPicklists", name: "GlobalPicklist"},
	metapath{path: "groups", name: "Group"},
	metapath{path: "homePageComponents", name: "HomePageComponent"},
	metapath{path: "homePageLayouts", name: "HomePageLayout"},
	metapath{path: "installedPackages", name: "InstalledPackage"},
	metapath{path: "labels", name: "CustomLabels"},
	metapath{path: "layouts", name: "Layout"},
	metapath{path: "LeadConvertSettings", name: "LeadConvertSettings"},
	metapath{path: "letterhead", name: "Letterhead"},
	metapath{path: "matchingRules", name: "MatchingRules"},
	metapath{path: "namedCredentials", name: "NamedCredential"},
	metapath{path: "objects", name: "CustomObject"},
	metapath{path: "objectTranslations", name: "CustomObjectTranslation"},
	metapath{path: "pages", name: "ApexPage"},
	metapath{path: "pathAssistants", name: "PathAssistant"},
	metapath{path: "permissionsets", name: "PermissionSet"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "profiles", name: "Profile", extension: ".profile"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "postTemplates", name: "PostTemplate"},
	metapath{path: "profiles", name: "Profile"},
	metapath{path: "queues", name: "Queue"},
	metapath{path: "quickActions", name: "QuickAction"},
	metapath{path: "remoteSiteSettings", name: "RemoteSiteSetting"},
	metapath{path: "reports", name: "Report", hasFolder: true},
	metapath{path: "reportTypes", name: "ReportType"},
	metapath{path: "roles", name: "Role"},
	metapath{path: "scontrols", name: "Scontrol"},
	metapath{path: "settings", name: "Settings"},
	metapath{path: "sharingRules", name: "SharingRules"},
	metapath{path: "siteDotComSites", name: "SiteDotCom"},
	metapath{path: "sites", name: "CustomSite"},
	metapath{path: "staticresources", name: "StaticResource"},
	metapath{path: "synonymDictionaries", name: "SynonymDictionary"},
	metapath{path: "tabs", name: "CustomTab"},
	metapath{path: "triggers", name: "ApexTrigger"},
	metapath{path: "weblinks", name: "CustomPageWebLink"},
	metapath{path: "workflows", name: "Workflow"},
	metapath{path: "cspTrustedSites", name: "CspTrustedSite"},
}

type PackageBuilder struct {
	IsPush   bool
	Metadata map[string]MetaType
	Files    ForceMetadataFiles
}

func NewPushBuilder() PackageBuilder {
	pb := PackageBuilder{IsPush: true}
	pb.Metadata = make(map[string]MetaType)
	pb.Files = make(ForceMetadataFiles)

	return pb
}

func NewFetchBuilder() PackageBuilder {
	pb := PackageBuilder{IsPush: false}
	pb.Metadata = make(map[string]MetaType)
	pb.Files = make(ForceMetadataFiles)

	return pb
}

// Build and return package.xml
func (pb PackageBuilder) PackageXml() []byte {
	p := createPackage()

	for _, metaType := range pb.Metadata {
		p.Types = append(p.Types, metaType)
	}

	byteXml, _ := xml.MarshalIndent(p, "", "    ")
	byteXml = append([]byte(xml.Header), byteXml...)
	//if err := ioutil.WriteFile("mypackage.xml", byteXml, 0644); err != nil {
	//ErrorAndExit(err.Error())
	//}
	return byteXml
}

// Returns the full ForceMetadataFiles container
func (pb *PackageBuilder) ForceMetadataFiles() ForceMetadataFiles {
	pb.Files["package.xml"] = pb.PackageXml()
	return pb.Files
}

// Returns the source file path for a given metadata file path.
func MetaPathToSourcePath(mpath string) (spath string) {
	spath = strings.TrimSuffix(mpath, "-meta.xml")
	if spath == mpath {
		return
	}

	_, err := os.Stat(spath)
	if err != nil {
		spath = mpath
	}
	return
}

// Add a file to the builder
func (pb *PackageBuilder) AddFile(fpath string) (fname string, err error) {
	fpath, err = filepath.Abs(fpath)
	if err != nil {
		return
	}
	_, err = os.Stat(fpath)
	if err != nil {
		return
	}

	isDestructiveChanges, err := regexp.MatchString("destructiveChanges(Pre|Post)?"+regexp.QuoteMeta(".")+"xml", fpath)
	if err != nil {
		return
	}

	fpath = MetaPathToSourcePath(fpath)
	metaName, fname := getMetaTypeFromPath(fpath)
	if !isDestructiveChanges && !strings.HasSuffix(fpath, "-meta.xml") {
		pb.AddMetaToPackage(metaName, fname)
	}

	// If it's a push, we want to actually add the files
	if pb.IsPush {
		if isDestructiveChanges {
			err = pb.addDestructiveChanges(fpath)
		} else {
			err = pb.addFileToWorkingDir(metaName, fpath)
		}
	}

	return
}

// Adds the file to a temp directory for deploy
func (pb *PackageBuilder) addFileToWorkingDir(metaName string, fpath string) (err error) {
	// Get relative dir from source
	srcDir := filepath.Dir(filepath.Dir(fpath))
	for _, mp := range metapaths {
		if metaName == mp.name && mp.hasFolder {
			srcDir = filepath.Dir(srcDir)
		}
	}
	frel, _ := filepath.Rel(srcDir, fpath)

	// Try to find meta file
	hasMeta := true
	fmeta := fpath + "-meta.xml"
	fmetarel := ""
	if _, err = os.Stat(fmeta); err != nil {
		if os.IsNotExist(err) {
			hasMeta = false
		} else {
			// Has error
			return
		}
	} else {
		// Should be present since we worked back to srcDir
		fmetarel, _ = filepath.Rel(srcDir, fmeta)
	}

	fdata, err := ioutil.ReadFile(fpath)
	if err != nil {
		return
	}

	pb.Files[frel] = fdata
	if hasMeta {
		fdata, err = ioutil.ReadFile(fmeta)
		pb.Files[fmetarel] = fdata
		return
	}

	return
}

func (pb *PackageBuilder) addDestructiveChanges(fpath string) (err error) {
	fdata, err := ioutil.ReadFile(fpath)
	if err != nil {
		return
	}

	frel, _ := filepath.Rel(filepath.Dir(fpath), fpath)
	pb.Files[frel] = fdata

	return
}

func (pb *PackageBuilder) contains(members []string, name string) bool {
	for _, a := range members {
		if a == name {
			return true
		}
	}
	return false
}

// Adds a metadata name to the pending package
func (pb *PackageBuilder) AddMetaToPackage(metaName string, name string) {
	mt := pb.Metadata[metaName]
	if mt.Name == "" {
		mt.Name = metaName
	}

	if !pb.contains(mt.Members, name) {
		mt.Members = append(mt.Members, name)
		pb.Metadata[metaName] = mt
	}
}

// Gets metadata type name and target name from a file path
func getMetaTypeFromPath(fpath string) (metaName string, name string) {
	fpath, err := filepath.Abs(fpath)
	if err != nil {
		ErrorAndExit("Cound not find " + fpath)
	}
	if _, err := os.Stat(fpath); err != nil {
		ErrorAndExit("Cound not open " + fpath)
	}

	// Get the metadata type and name for the file
	metaName, fileName := getMetaForPath(fpath)
	name = strings.TrimSuffix(fileName, filepath.Ext(fileName))
	//name = strings.TrimSuffix(name, filepath.Ext(name))
	return
}

// Gets partial path based on a meta type name
func getPathForMeta(metaname string) string {
	for _, mp := range metapaths {
		if strings.EqualFold(mp.name, metaname) {
			return mp.path
		}
	}

	// Unknown, so use metaname
	return metaname
}

func findMetapathForFile(file string) (path metapath) {
	parentDir := filepath.Dir(file)
	parentName := filepath.Base(parentDir)
	grandparentName := filepath.Base(filepath.Dir(parentDir))
	fileExtension := filepath.Ext(file)

	for _, mp := range metapaths {
		if mp.hasFolder && grandparentName == mp.path {
			return mp
		}
		if mp.path == parentName {
			return mp
		}
	}

	// Hmm, maybe we can use the extension to determine the type
	for _, mp := range metapaths {
		if mp.extension == fileExtension {
			return mp
		}
	}
	return
}

// Gets meta type and name based on a path
func getMetaForPath(path string) (metaName string, objectName string) {
	parentDir := filepath.Dir(path)
	parentName := filepath.Base(parentDir)
	grandparentName := filepath.Base(filepath.Dir(parentDir))
	fileName := filepath.Base(path)

	for _, mp := range metapaths {
		if mp.hasFolder && grandparentName == mp.path {
			metaName = mp.name
			if mp.onlyFolder {
				objectName = parentName
			} else {
				objectName = parentName + "/" + fileName
			}
			return
		}
		if mp.path == parentName {
			metaName = mp.name
			objectName = fileName
			return
		}
	}

	// Unknown, so use path
	metaName = parentName
	objectName = fileName
	return
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"math/rand"
	"sync"
	"time"

	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// jitter is the percent of jitter added to restart delays.
	jitter = 0.25

	ReasonNoRestartsAllowed   = "Policy allows no restarts"
	ReasonUnrecoverableErrror = "Error was unrecoverable"
	ReasonWithinPolicy        = "Restart within policy"
	ReasonDelay               = "Exceeded allowed attempts, applying a delay"
)

func newRestartTracker(policy *structs.RestartPolicy, jobType string) *RestartTracker {
	onSuccess := true
	if jobType == structs.JobTypeBatch {
		onSuccess = false
	}
	return &RestartTracker{
		startTime: time.Now(),
		onSuccess: onSuccess,
		policy:    policy,
		rand:      rand.New(rand.NewSource(time.Now().Unix())),
	}
}

type RestartTracker struct {
	waitRes   *cstructs.WaitResult
	startErr  error
	count     int       // Current number of attempts.
	onSuccess bool      // Whether to restart on successful exit code.
	startTime time.Time // When the interval began
	reason    string    // The reason for the last state
	policy    *structs.RestartPolicy
	rand      *rand.Rand
	lock      sync.Mutex
}

// SetPolicy updates the policy used to determine restarts.
func (r *RestartTracker) SetPolicy(policy *structs.RestartPolicy) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.policy = policy
}

// SetStartError is used to mark the most recent start error. If starting was
// successful the error should be nil.
func (r *RestartTracker) SetStartError(err error) *RestartTracker {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.startErr = err
	return r
}

// SetWaitResult is used to mark the most recent wait result.
func (r *RestartTracker) SetWaitResult(res *cstructs.WaitResult) *RestartTracker {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.waitRes = res
	return r
}

// GetReason returns a human-readable description for the last state returned by
// GetState.
func (r *RestartTracker) GetReason() string {
	r.lock.Lock()
	defer r.lock.Unlock()
	return r.reason
}

// GetState returns the tasks next state given the set exit code and start
// error. One of the following states are returned:
// * TaskRestarting - Task should be restarted
// * TaskNotRestarting - Task should not be restarted and has exceeded its
//   restart policy.
// * TaskTerminated - Task has terminated successfully and does not need a
//   restart.
//
// If TaskRestarting is returned, the duration is how long to wait until
// starting the task again.
func (r *RestartTracker) GetState() (string, time.Duration) {
	r.lock.Lock()
	defer r.lock.Unlock()

	// Hot path if no attempts are expected
	if r.policy.Attempts == 0 {
		r.reason = ReasonNoRestartsAllowed
		if r.waitRes != nil && r.waitRes.Successful() {
			return structs.TaskTerminated, 0
		}

		return structs.TaskNotRestarting, 0
	}

	r.count++

	// Check if we have entered a new interval.
	end := r.startTime.Add(r.policy.Interval)
	now := time.Now()
	if now.After(end) {
		r.count = 0
		r.startTime = now
	}

	if r.startErr != nil {
		return r.handleStartError()
	} else if r.waitRes != nil {
		return r.handleWaitResult()
	} else {
		return "", 0
	}
}

// handleStartError returns the new state and potential wait duration for
// restarting the task after it was not successfully started. On start errors,
// the restart policy is always treated as fail mode to ensure we don't
// infinitely try to start a task.
func (r *RestartTracker) handleStartError() (string, time.Duration) {
	// If the error is not recoverable, do not restart.
	if rerr, ok := r.startErr.(*cstructs.RecoverableError); !(ok && rerr.Recoverable) {
		r.reason = ReasonUnrecoverableErrror
		return structs.TaskNotRestarting, 0
	}

	if r.count > r.policy.Attempts {
		if r.policy.Mode == structs.RestartPolicyModeFail {
			r.reason = fmt.Sprintf(
				`Exceeded allowed atttempts %d in interval %v and mode is "fail"`,
				r.policy.Attempts, r.policy.Interval)
			return structs.TaskNotRestarting, 0
		} else {
			r.reason = ReasonDelay
			return structs.TaskRestarting, r.getDelay()
		}
	}

	r.reason = ReasonWithinPolicy
	return structs.TaskRestarting, r.jitter()
}

// handleWaitResult returns the new state and potential wait duration for
// restarting the task after it has exited.
func (r *RestartTracker) handleWaitResult() (string, time.Duration) {
	// If the task started successfully and restart on success isn't specified,
	// don't restart but don't mark as failed.
	if r.waitRes.Successful() && !r.onSuccess {
		r.reason = "Restart unnecessary as task terminated successfully"
		return structs.TaskTerminated, 0
	}

	if r.count > r.policy.Attempts {
		if r.policy.Mode == structs.RestartPolicyModeFail {
			r.reason = fmt.Sprintf(
				`Exceeded allowed atttempts %d in interval %v and mode is "fail"`,
				r.policy.Attempts, r.policy.Interval)
			return structs.TaskNotRestarting, 0
		} else {
			r.reason = ReasonDelay
			return structs.TaskRestarting, r.getDelay()
		}
	}

	r.reason = ReasonWithinPolicy
	return structs.TaskRestarting, r.jitter()
}

// getDelay returns the delay time to enter the next interval.
func (r *RestartTracker) getDelay() time.Duration {
	end := r.startTime.Add(r.policy.Interval)
	now := time.Now()
	return end.Sub(now)
}

// jitter returns the delay time plus a jitter.
func (r *RestartTracker) jitter() time.Duration {
	// Get the delay and ensure it is valid.
	d := r.policy.Delay.Nanoseconds()
	if d == 0 {
		d = 1
	}

	j := float64(r.rand.Int63n(d)) * jitter
	return time.Duration(d + int64(j))
}

// Returns a tracker that never restarts.
func noRestartsTracker() *RestartTracker {
	policy := &structs.RestartPolicy{Attempts: 0, Mode: structs.RestartPolicyModeFail}
	return newRestartTracker(policy, structs.JobTypeBatch)
}
<commit_msg>Fix typo: atttempts<commit_after>package client

import (
	"fmt"
	"math/rand"
	"sync"
	"time"

	cstructs "github.com/hashicorp/nomad/client/driver/structs"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// jitter is the percent of jitter added to restart delays.
	jitter = 0.25

	ReasonNoRestartsAllowed   = "Policy allows no restarts"
	ReasonUnrecoverableErrror = "Error was unrecoverable"
	ReasonWithinPolicy        = "Restart within policy"
	ReasonDelay               = "Exceeded allowed attempts, applying a delay"
)

func newRestartTracker(policy *structs.RestartPolicy, jobType string) *RestartTracker {
	onSuccess := true
	if jobType == structs.JobTypeBatch {
		onSuccess = false
	}
	return &RestartTracker{
		startTime: time.Now(),
		onSuccess: onSuccess,
		policy:    policy,
		rand:      rand.New(rand.NewSource(time.Now().Unix())),
	}
}

type RestartTracker struct {
	waitRes   *cstructs.WaitResult
	startErr  error
	count     int       // Current number of attempts.
	onSuccess bool      // Whether to restart on successful exit code.
	startTime time.Time // When the interval began
	reason    string    // The reason for the last state
	policy    *structs.RestartPolicy
	rand      *rand.Rand
	lock      sync.Mutex
}

// SetPolicy updates the policy used to determine restarts.
func (r *RestartTracker) SetPolicy(policy *structs.RestartPolicy) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.policy = policy
}

// SetStartError is used to mark the most recent start error. If starting was
// successful the error should be nil.
func (r *RestartTracker) SetStartError(err error) *RestartTracker {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.startErr = err
	return r
}

// SetWaitResult is used to mark the most recent wait result.
func (r *RestartTracker) SetWaitResult(res *cstructs.WaitResult) *RestartTracker {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.waitRes = res
	return r
}

// GetReason returns a human-readable description for the last state returned by
// GetState.
func (r *RestartTracker) GetReason() string {
	r.lock.Lock()
	defer r.lock.Unlock()
	return r.reason
}

// GetState returns the tasks next state given the set exit code and start
// error. One of the following states are returned:
// * TaskRestarting - Task should be restarted
// * TaskNotRestarting - Task should not be restarted and has exceeded its
//   restart policy.
// * TaskTerminated - Task has terminated successfully and does not need a
//   restart.
//
// If TaskRestarting is returned, the duration is how long to wait until
// starting the task again.
func (r *RestartTracker) GetState() (string, time.Duration) {
	r.lock.Lock()
	defer r.lock.Unlock()

	// Hot path if no attempts are expected
	if r.policy.Attempts == 0 {
		r.reason = ReasonNoRestartsAllowed
		if r.waitRes != nil && r.waitRes.Successful() {
			return structs.TaskTerminated, 0
		}

		return structs.TaskNotRestarting, 0
	}

	r.count++

	// Check if we have entered a new interval.
	end := r.startTime.Add(r.policy.Interval)
	now := time.Now()
	if now.After(end) {
		r.count = 0
		r.startTime = now
	}

	if r.startErr != nil {
		return r.handleStartError()
	} else if r.waitRes != nil {
		return r.handleWaitResult()
	} else {
		return "", 0
	}
}

// handleStartError returns the new state and potential wait duration for
// restarting the task after it was not successfully started. On start errors,
// the restart policy is always treated as fail mode to ensure we don't
// infinitely try to start a task.
func (r *RestartTracker) handleStartError() (string, time.Duration) {
	// If the error is not recoverable, do not restart.
	if rerr, ok := r.startErr.(*cstructs.RecoverableError); !(ok && rerr.Recoverable) {
		r.reason = ReasonUnrecoverableErrror
		return structs.TaskNotRestarting, 0
	}

	if r.count > r.policy.Attempts {
		if r.policy.Mode == structs.RestartPolicyModeFail {
			r.reason = fmt.Sprintf(
				`Exceeded allowed attempts %d in interval %v and mode is "fail"`,
				r.policy.Attempts, r.policy.Interval)
			return structs.TaskNotRestarting, 0
		} else {
			r.reason = ReasonDelay
			return structs.TaskRestarting, r.getDelay()
		}
	}

	r.reason = ReasonWithinPolicy
	return structs.TaskRestarting, r.jitter()
}

// handleWaitResult returns the new state and potential wait duration for
// restarting the task after it has exited.
func (r *RestartTracker) handleWaitResult() (string, time.Duration) {
	// If the task started successfully and restart on success isn't specified,
	// don't restart but don't mark as failed.
	if r.waitRes.Successful() && !r.onSuccess {
		r.reason = "Restart unnecessary as task terminated successfully"
		return structs.TaskTerminated, 0
	}

	if r.count > r.policy.Attempts {
		if r.policy.Mode == structs.RestartPolicyModeFail {
			r.reason = fmt.Sprintf(
				`Exceeded allowed attempts %d in interval %v and mode is "fail"`,
				r.policy.Attempts, r.policy.Interval)
			return structs.TaskNotRestarting, 0
		} else {
			r.reason = ReasonDelay
			return structs.TaskRestarting, r.getDelay()
		}
	}

	r.reason = ReasonWithinPolicy
	return structs.TaskRestarting, r.jitter()
}

// getDelay returns the delay time to enter the next interval.
func (r *RestartTracker) getDelay() time.Duration {
	end := r.startTime.Add(r.policy.Interval)
	now := time.Now()
	return end.Sub(now)
}

// jitter returns the delay time plus a jitter.
func (r *RestartTracker) jitter() time.Duration {
	// Get the delay and ensure it is valid.
	d := r.policy.Delay.Nanoseconds()
	if d == 0 {
		d = 1
	}

	j := float64(r.rand.Int63n(d)) * jitter
	return time.Duration(d + int64(j))
}

// Returns a tracker that never restarts.
func noRestartsTracker() *RestartTracker {
	policy := &structs.RestartPolicy{Attempts: 0, Mode: structs.RestartPolicyModeFail}
	return newRestartTracker(policy, structs.JobTypeBatch)
}
<|endoftext|>


<commit_before>package physical

import "github.com/hashicorp/golang-lru"

const (
	// DefaultCacheSize is used if no cache size is specified for NewCache
	DefaultCacheSize = 32 * 1024
)

// Cache is used to wrap an underlying physical backend
// and provide an LRU cache layer on top. Most of the reads done by
// Vault are for policy objects so there is a large read reduction
// by using a simple write-through cache.
type Cache struct {
	backend Backend
	lru     *lru.Cache
}

// NewCache returns a physical cache of the given size.
// If no size is provided, the default size is used.
func NewCache(b Backend, size int) *Cache {
	if size <= 0 {
		size = DefaultCacheSize
	}
	cache, _ := lru.New(size)
	c := &Cache{
		backend: b,
		lru:     cache,
	}
	return c
}

// Purge is used to clear the cache
func (c *Cache) Purge() {
	c.lru.Purge()
}

func (c *Cache) Put(entry *Entry) error {
	err := c.backend.Put(entry)
	c.lru.Add(entry.Key, entry)
	return err
}

func (c *Cache) Get(key string) (*Entry, error) {
	// Check the LRU first
	if raw, ok := c.lru.Get(key); ok {
		if raw == nil {
			return nil, nil
		} else {
			return raw.(*Entry), nil
		}
	}

	// Read from the underlying backend
	ent, err := c.backend.Get(key)
	if err != nil {
		return nil, err
	}

	// Cache the result
	c.lru.Add(key, ent)
	return ent, err
}

func (c *Cache) Delete(key string) error {
	err := c.backend.Delete(key)
	c.lru.Remove(key)
	return err
}

func (c *Cache) List(prefix string) ([]string, error) {
	// Always pass-through as this would be difficult to cache.
	return c.backend.List(prefix)
}
<commit_msg>physical: fix negative cache issue for core keys<commit_after>package physical

import (
	"strings"

	"github.com/hashicorp/golang-lru"
)

const (
	// DefaultCacheSize is used if no cache size is specified for NewCache
	DefaultCacheSize = 32 * 1024
)

// Cache is used to wrap an underlying physical backend
// and provide an LRU cache layer on top. Most of the reads done by
// Vault are for policy objects so there is a large read reduction
// by using a simple write-through cache.
type Cache struct {
	backend Backend
	lru     *lru.Cache
}

// NewCache returns a physical cache of the given size.
// If no size is provided, the default size is used.
func NewCache(b Backend, size int) *Cache {
	if size <= 0 {
		size = DefaultCacheSize
	}
	cache, _ := lru.New(size)
	c := &Cache{
		backend: b,
		lru:     cache,
	}
	return c
}

// Purge is used to clear the cache
func (c *Cache) Purge() {
	c.lru.Purge()
}

func (c *Cache) Put(entry *Entry) error {
	err := c.backend.Put(entry)
	c.lru.Add(entry.Key, entry)
	return err
}

func (c *Cache) Get(key string) (*Entry, error) {
	// Check the LRU first
	if raw, ok := c.lru.Get(key); ok {
		if raw == nil {
			return nil, nil
		} else {
			return raw.(*Entry), nil
		}
	}

	// Read from the underlying backend
	ent, err := c.backend.Get(key)
	if err != nil {
		return nil, err
	}

	// Cache the result. We do NOT cache negative results
	// for keys in the 'core/' prefix otherwise we risk certain
	// race conditions upstream. The primary issue is with the HA mode,
	// we could potentially negatively cache the leader entry and cause
	// leader discovery to fail.
	if ent != nil || !strings.HasPrefix(key, "core/") {
		c.lru.Add(key, ent)
	}
	return ent, err
}

func (c *Cache) Delete(key string) error {
	err := c.backend.Delete(key)
	c.lru.Remove(key)
	return err
}

func (c *Cache) List(prefix string) ([]string, error) {
	// Always pass-through as this would be difficult to cache.
	return c.backend.List(prefix)
}
<|endoftext|>


<commit_before>package pila

import (
	"reflect"
	"testing"
)

func TestNewPila(t *testing.T) {
	pila := NewPila()
	if pila == nil {
		t.Fatal("pila is nil")
	}
	if pila.Databases == nil {
		t.Error("pila.Databases is nil")
	}

}

func TestPilaCreateDatabase(t *testing.T) {
	pila := NewPila()
	id := pila.CreateDatabase("test-1")

	db, ok := pila.Databases[id]
	if !ok {
		t.Errorf("db %v not added to pila", id)
	} else if !reflect.DeepEqual(db.Pila, pila) {
		t.Errorf("db %v does not contain pile", id)
	}
}

func TestPilaAddDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	err := pila.AddDatabase(db)
	if err != nil {
		t.Fatal("err is not nil")
	}

	id := db.ID
	db, ok := pila.Databases[id]
	if !ok {
		t.Errorf("db %v not added to pila", id)
	} else if !reflect.DeepEqual(db.Pila, pila) {
		t.Errorf("db %v does not contain pila", id)
	}

}

func TestPilaAddDatabase_Error(t *testing.T) {
	pila := NewPila()
	pila2 := NewPila()
	db := NewDatabase("test")
	db2 := NewDatabase("test")

	if err := pila.AddDatabase(db); err != nil {
		t.Errorf("err is not nil, but %v", err)
	}

	if err := pila2.AddDatabase(db); err == nil {
		t.Error("err is nil")
	}

	if err := pila.AddDatabase(db2); err == nil {
		t.Error("err is nil")
	}
}

func TestPilaRemoveDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	pila.AddDatabase(db)

	if ok := pila.RemoveDatabase(db.ID); !ok {
		t.Errorf("RemoveDatabase did not succeed")
	}

	if db.Pila != nil {
		t.Errorf("a pila is assigned to database %v", db.Name)
	}

	if _, ok := pila.Databases[db.ID]; ok {
		t.Errorf("Removed database does exist on pila")
	}
}

func TestPilaRemoveDatabase_False(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")

	if ok := pila.RemoveDatabase(db.ID); ok {
		t.Errorf("RemoveDatabase did succeed")
	}
}

func TestPilaDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	err := pila.AddDatabase(db)
	if err != nil {
		t.Fatal("err is not nil")
	}

	db2, ok := pila.Database(db.ID)
	if !ok {
		t.Errorf("pila has no Database %v", db.ID)
	} else if !reflect.DeepEqual(db2, db) {
		t.Errorf("Database is %v, expected %v", db2, db)
	}
}

func TestPilaDatabase_False(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	db2, ok := pila.Database(db.ID)
	if ok {
		t.Errorf("pila has Database %v", db.ID)
	} else if db2 != nil {
		t.Errorf("Database %v is not nil", db2)
	}
}

func TestPilaStatus(t *testing.T) {
	pila := NewPila()
	db0 := NewDatabase("db0")
	db1 := NewDatabase("db1")
	db2 := NewDatabase("db2")
	pila.AddDatabase(db0)
	pila.AddDatabase(db1)
	pila.AddDatabase(db2)

	expectedStatus := `{"number_of_databases":3,"databases":[{"id":"714e49277eb730717e413b167b76ef78","name":"db0","number_of_stacks":0},{"id":"93c6f621b761cd88017846beae63f4be","name":"db1","number_of_stacks":0},{"id":"5d02dd2c3917fdd29abe20a2c1b5ea1c","name":"db2","number_of_stacks":0}]}`

	if status := pila.Status(); string(status) != expectedStatus {
		t.Errorf("status is %s, expected %s", string(status), expectedStatus)
	}
}

func TestPilaStatus_Empty(t *testing.T) {
	pila := NewPila()

	expectedStatus := `{"number_of_databases":0,"databases":[]}`

	if status := pila.Status(); string(status) != expectedStatus {
		t.Errorf("status is %s, expected %s", string(status), expectedStatus)
	}
}
<commit_msg>pila: reduce PilaTestStatus as order doesn't matter<commit_after>package pila

import (
	"reflect"
	"testing"
)

func TestNewPila(t *testing.T) {
	pila := NewPila()
	if pila == nil {
		t.Fatal("pila is nil")
	}
	if pila.Databases == nil {
		t.Error("pila.Databases is nil")
	}

}

func TestPilaCreateDatabase(t *testing.T) {
	pila := NewPila()
	id := pila.CreateDatabase("test-1")

	db, ok := pila.Databases[id]
	if !ok {
		t.Errorf("db %v not added to pila", id)
	} else if !reflect.DeepEqual(db.Pila, pila) {
		t.Errorf("db %v does not contain pile", id)
	}
}

func TestPilaAddDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	err := pila.AddDatabase(db)
	if err != nil {
		t.Fatal("err is not nil")
	}

	id := db.ID
	db, ok := pila.Databases[id]
	if !ok {
		t.Errorf("db %v not added to pila", id)
	} else if !reflect.DeepEqual(db.Pila, pila) {
		t.Errorf("db %v does not contain pila", id)
	}

}

func TestPilaAddDatabase_Error(t *testing.T) {
	pila := NewPila()
	pila2 := NewPila()
	db := NewDatabase("test")
	db2 := NewDatabase("test")

	if err := pila.AddDatabase(db); err != nil {
		t.Errorf("err is not nil, but %v", err)
	}

	if err := pila2.AddDatabase(db); err == nil {
		t.Error("err is nil")
	}

	if err := pila.AddDatabase(db2); err == nil {
		t.Error("err is nil")
	}
}

func TestPilaRemoveDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	pila.AddDatabase(db)

	if ok := pila.RemoveDatabase(db.ID); !ok {
		t.Errorf("RemoveDatabase did not succeed")
	}

	if db.Pila != nil {
		t.Errorf("a pila is assigned to database %v", db.Name)
	}

	if _, ok := pila.Databases[db.ID]; ok {
		t.Errorf("Removed database does exist on pila")
	}
}

func TestPilaRemoveDatabase_False(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")

	if ok := pila.RemoveDatabase(db.ID); ok {
		t.Errorf("RemoveDatabase did succeed")
	}
}

func TestPilaDatabase(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	err := pila.AddDatabase(db)
	if err != nil {
		t.Fatal("err is not nil")
	}

	db2, ok := pila.Database(db.ID)
	if !ok {
		t.Errorf("pila has no Database %v", db.ID)
	} else if !reflect.DeepEqual(db2, db) {
		t.Errorf("Database is %v, expected %v", db2, db)
	}
}

func TestPilaDatabase_False(t *testing.T) {
	pila := NewPila()
	db := NewDatabase("test")
	db2, ok := pila.Database(db.ID)
	if ok {
		t.Errorf("pila has Database %v", db.ID)
	} else if db2 != nil {
		t.Errorf("Database %v is not nil", db2)
	}
}

func TestPilaStatus(t *testing.T) {
	pila := NewPila()
	db0 := NewDatabase("db0")
	pila.AddDatabase(db0)

	expectedStatus := `{"number_of_databases":1,"databases":[{"id":"714e49277eb730717e413b167b76ef78","name":"db0","number_of_stacks":0}]}`

	if status := pila.Status(); string(status) != expectedStatus {
		t.Errorf("status is %s, expected %s", string(status), expectedStatus)
	}
}

func TestPilaStatus_Empty(t *testing.T) {
	pila := NewPila()

	expectedStatus := `{"number_of_databases":0,"databases":[]}`

	if status := pila.Status(); string(status) != expectedStatus {
		t.Errorf("status is %s, expected %s", string(status), expectedStatus)
	}
}
<|endoftext|>


<commit_before>package mem

import (
	"fmt"
)

// Memory allows read and write access to memory
type Memory struct {
	hwr         *HardwareRegisters
	mbc         mbc
	videoRAM    [0x2000]byte
	internalRAM [0x2000]byte
	oam         [0xa0]byte
	zeroPage    [0x8f]byte
}

// NewMemory creates the memory and initializes it with ROM contents and default values
func NewMemory(hwr *HardwareRegisters, romFilename string) *Memory {
	return &Memory{
		hwr: hwr,
		mbc: newMBC(romFilename),
	}
}

// Read a byte from the chosen memory location
func (mem *Memory) Read(addr uint16) byte {
	switch {
	case addr < 0x8000:
		return mem.mbc.read(addr)
	case addr < 0xa000:
		return mem.videoRAM[addr-0x8000]
	case addr < 0xc000:
		panic(fmt.Sprintf("Read on cartridge RAM is not implemented: 0x%04x", addr))
	case addr < 0xe000:
		return mem.internalRAM[addr-0xc000]
	case addr < 0xfe00:
		return mem.internalRAM[addr-0xe000]
	case addr < 0xfea0:
		return mem.oam[addr-0xfe00]
	case addr < 0xff00:
		return 0 // Unusable region
	case addr < 0xff80:
		return mem.readHardwareRegisters(addr)
	case addr < 0xffff:
		return mem.zeroPage[addr-0xff80]
	case addr == 0xffff:
		return mem.hwr.IE
	default:
		panic(fmt.Sprintf("Read failed: 0x%04x", addr))
	}
}

// VideoRAM for direct access
func (mem *Memory) VideoRAM() *[0x2000]byte {
	return &mem.videoRAM
}

// OAM memory for direct access
func (mem *Memory) OAM() *[0xa0]byte {
	return &mem.oam
}

// Write a byte to the chosen memory location
func (mem *Memory) Write(addr uint16, value byte) {
	switch {
	case addr < 0x8000:
		mem.mbc.write(addr, value)
	case addr < 0xa000:
		mem.videoRAM[addr-0x8000] = value
	case addr < 0xc000:
		panic(fmt.Sprintf("Write on cartridge RAM is not implemented: 0x%04x", addr))
	case addr < 0xe000:
		mem.internalRAM[addr-0xc000] = value
	case addr < 0xfe00:
		mem.internalRAM[addr-0xe000] = value
	case addr < 0xfea0:
		mem.oam[addr-0xfe00] = value
	case addr < 0xff00:
		// Unusable region
	case addr < 0xff80:
		mem.writeHardwareRegisters(addr, value)
	case addr < 0xffff:
		mem.zeroPage[addr-0xff80] = value
	case addr == IE:
		mem.hwr.IE = value
	default:
		panic(fmt.Sprintf("Write failed: 0x%04x", addr))
	}
}
<commit_msg>Don't panic when accessing cart RAM<commit_after>package mem

import (
	"fmt"
)

// Memory allows read and write access to memory
type Memory struct {
	hwr         *HardwareRegisters
	mbc         mbc
	videoRAM    [0x2000]byte
	internalRAM [0x2000]byte
	oam         [0xa0]byte
	zeroPage    [0x8f]byte
}

// NewMemory creates the memory and initializes it with ROM contents and default values
func NewMemory(hwr *HardwareRegisters, romFilename string) *Memory {
	return &Memory{
		hwr: hwr,
		mbc: newMBC(romFilename),
	}
}

// Read a byte from the chosen memory location
func (mem *Memory) Read(addr uint16) byte {
	switch {
	case addr < 0x8000:
		return mem.mbc.read(addr)
	case addr < 0xa000:
		return mem.videoRAM[addr-0x8000]
	case addr < 0xc000:
		// FIXME
		// panic(fmt.Sprintf("Read on cartridge RAM is not implemented: 0x%04x", addr))
		return 0
	case addr < 0xe000:
		return mem.internalRAM[addr-0xc000]
	case addr < 0xfe00:
		return mem.internalRAM[addr-0xe000]
	case addr < 0xfea0:
		return mem.oam[addr-0xfe00]
	case addr < 0xff00:
		return 0 // Unusable region
	case addr < 0xff80:
		return mem.readHardwareRegisters(addr)
	case addr < 0xffff:
		return mem.zeroPage[addr-0xff80]
	case addr == 0xffff:
		return mem.hwr.IE
	default:
		panic(fmt.Sprintf("Read failed: 0x%04x", addr))
	}
}

// VideoRAM for direct access
func (mem *Memory) VideoRAM() *[0x2000]byte {
	return &mem.videoRAM
}

// OAM memory for direct access
func (mem *Memory) OAM() *[0xa0]byte {
	return &mem.oam
}

// Write a byte to the chosen memory location
func (mem *Memory) Write(addr uint16, value byte) {
	switch {
	case addr < 0x8000:
		mem.mbc.write(addr, value)
	case addr < 0xa000:
		mem.videoRAM[addr-0x8000] = value
	case addr < 0xc000:
		// FIXME maybe write to file
		// panic(fmt.Sprintf("Write on cartridge RAM is not implemented: 0x%04x", addr))
	case addr < 0xe000:
		mem.internalRAM[addr-0xc000] = value
	case addr < 0xfe00:
		mem.internalRAM[addr-0xe000] = value
	case addr < 0xfea0:
		mem.oam[addr-0xfe00] = value
	case addr < 0xff00:
		// Unusable region
	case addr < 0xff80:
		mem.writeHardwareRegisters(addr, value)
	case addr < 0xffff:
		mem.zeroPage[addr-0xff80] = value
	case addr == IE:
		mem.hwr.IE = value
	default:
		panic(fmt.Sprintf("Write failed: 0x%04x", addr))
	}
}
<|endoftext|>


<commit_before>package yago

import (
	"database/sql"
	"fmt"
	"reflect"

	"github.com/aacanakin/qb"
)

// Query helps querying structs from the database
type Query struct {
	db         *DB
	mapper     Mapper
	selectStmt qb.SelectStmt
}

// NewQuery creates a new query
func NewQuery(db *DB, mapper Mapper) Query {
	return Query{
		db:         db,
		mapper:     mapper,
		selectStmt: mapper.Table().Select(mapper.FieldList()...),
	}
}

// Where add filter clauses to the query
func (q Query) Where(clause qb.Clause) Query {
	return Query{
		db:         q.db,
		mapper:     q.mapper,
		selectStmt: q.selectStmt.Where(clause),
	}
}

// InnerJoinMapper joins a mapper table
func (q Query) InnerJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.InnerJoin(*mapper.Table(), fromCol, col)
	return q
}

// LeftJoinMapper joins a mapper table
func (q Query) LeftJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.LeftJoin(*mapper.Table(), fromCol, col)
	return q
}

// RightJoinMapper joins a mapper table
func (q Query) RightJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.RightJoin(*mapper.Table(), fromCol, col)
	return q
}

// InnerJoin joins a table
func (q Query) InnerJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.InnerJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// LeftJoin joins a table
func (q Query) LeftJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.LeftJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// RightJoin joins a table
func (q Query) RightJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.RightJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// SQLQuery runs the query
func (q Query) SQLQuery() (*sql.Rows, error) {
	return q.db.Engine.Query(q.selectStmt)
}

// SQLQueryRow runs the query and expects at most one row in the result
func (q Query) SQLQueryRow() *sql.Row {
	return q.db.Engine.QueryRow(q.selectStmt)
}

// One returns one and only one struct from the query.
// If query has no result or more than one, an error is returned
func (q Query) One(s MappedStruct) error {
	rows, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer rows.Close()

	if !rows.Next() {
		return fmt.Errorf("NoResultError")
	}

	err = q.mapper.Scan(rows, s)
	if err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("TooManyResultsError")
	}
	return nil
}

// All load all the structs matching the query
func (q Query) All(value interface{}) error {
	rows, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer rows.Close()

	resultType := q.mapper.StructType()

	results := reflect.Indirect(reflect.ValueOf(value))

	var (
		isPtr     bool
		wrongType bool
	)

	if results.Kind() != reflect.Slice {
		wrongType = true
	} else {
		elemType := results.Type().Elem()
		if elemType.Kind() == reflect.Ptr {
			isPtr = true
			wrongType = results.Type().Elem().Elem() != resultType
		} else {
			wrongType = results.Type().Elem() != resultType
		}
	}
	if wrongType {
		return fmt.Errorf("yago Query.All(): Expected a []%s, got %v", resultType, results.Type())
	}

	// Empty the slice
	results.Set(reflect.MakeSlice(results.Type(), 0, 0))

	for rows.Next() {
		elem := reflect.New(resultType).Elem()
		if err := q.mapper.Scan(rows, elem.Addr().Interface().(MappedStruct)); err != nil {
			return fmt.Errorf("yago Query.All(): Error while scanning: %s", err)
		}
		if isPtr {
			results.Set(reflect.Append(results, elem.Addr()))
		} else {
			results.Set(reflect.Append(results, elem))
		}
	}
	if err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("TooManyResultsError")
	}
	return nil
}

// Count change the columns to COUNT(*), execute the query and returns
// the result
func (q Query) Count(count interface{}) error {
	// XXX mapper should be able to return a list of pkey fields
	// XXX When qb supports COUNT(*), use it
	q.selectStmt = q.selectStmt.Select(qb.Count(
		q.mapper.Table().PrimaryCols()[0]),
	)
	row, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer row.Close()
	if !row.Next() {
		panic("No result")
	}

	err = row.Scan(count)
	if err != nil {
		return err
	}
	return nil
}

// Exists return true if any record matches the current query
func (q Query) Exists() (bool, error) {
	q.selectStmt = q.selectStmt.Select(qb.SQLText("1")).Limit(0, 1)
	row, err := q.SQLQuery()
	if err != nil {
		return false, err
	}
	defer row.Close()
	if !row.Next() {
		panic("No result")
	}
	var exists bool
	err = row.Scan(&exists)
	if err != nil {
		return false, err
	}
	return exists, nil
}
<commit_msg>Add Query.Select()<commit_after>package yago

import (
	"database/sql"
	"fmt"
	"reflect"

	"github.com/aacanakin/qb"
)

// Query helps querying structs from the database
type Query struct {
	db         *DB
	mapper     Mapper
	selectStmt qb.SelectStmt
}

// NewQuery creates a new query
func NewQuery(db *DB, mapper Mapper) Query {
	return Query{
		db:         db,
		mapper:     mapper,
		selectStmt: mapper.Table().Select(mapper.FieldList()...),
	}
}

// Select redefines the SELECT clauses
func (q Query) Select(clause ...qb.Clause) Query {
	q.selectStmt = q.selectStmt.Select(clause...)
	return q
}

// Where add filter clauses to the query
func (q Query) Where(clause qb.Clause) Query {
	return Query{
		db:         q.db,
		mapper:     q.mapper,
		selectStmt: q.selectStmt.Where(clause),
	}
}

// InnerJoinMapper joins a mapper table
func (q Query) InnerJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.InnerJoin(*mapper.Table(), fromCol, col)
	return q
}

// LeftJoinMapper joins a mapper table
func (q Query) LeftJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.LeftJoin(*mapper.Table(), fromCol, col)
	return q
}

// RightJoinMapper joins a mapper table
func (q Query) RightJoinMapper(mapper Mapper, fromCol qb.ColumnElem, col qb.ColumnElem) Query {
	q.selectStmt = q.selectStmt.RightJoin(*mapper.Table(), fromCol, col)
	return q
}

// InnerJoin joins a table
func (q Query) InnerJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.InnerJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// LeftJoin joins a table
func (q Query) LeftJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.LeftJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// RightJoin joins a table
func (q Query) RightJoin(model Model, fields ...ScalarField) Query {
	// TODO if fields is empty, find the relation based on foreign keys
	return q.RightJoinMapper(model.GetMapper(), fields[0].Column, fields[1].Column)
}

// SQLQuery runs the query
func (q Query) SQLQuery() (*sql.Rows, error) {
	return q.db.Engine.Query(q.selectStmt)
}

// SQLQueryRow runs the query and expects at most one row in the result
func (q Query) SQLQueryRow() *sql.Row {
	return q.db.Engine.QueryRow(q.selectStmt)
}

// One returns one and only one struct from the query.
// If query has no result or more than one, an error is returned
func (q Query) One(s MappedStruct) error {
	rows, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer rows.Close()

	if !rows.Next() {
		return fmt.Errorf("NoResultError")
	}

	err = q.mapper.Scan(rows, s)
	if err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("TooManyResultsError")
	}
	return nil
}

// All load all the structs matching the query
func (q Query) All(value interface{}) error {
	rows, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer rows.Close()

	resultType := q.mapper.StructType()

	results := reflect.Indirect(reflect.ValueOf(value))

	var (
		isPtr     bool
		wrongType bool
	)

	if results.Kind() != reflect.Slice {
		wrongType = true
	} else {
		elemType := results.Type().Elem()
		if elemType.Kind() == reflect.Ptr {
			isPtr = true
			wrongType = results.Type().Elem().Elem() != resultType
		} else {
			wrongType = results.Type().Elem() != resultType
		}
	}
	if wrongType {
		return fmt.Errorf("yago Query.All(): Expected a []%s, got %v", resultType, results.Type())
	}

	// Empty the slice
	results.Set(reflect.MakeSlice(results.Type(), 0, 0))

	for rows.Next() {
		elem := reflect.New(resultType).Elem()
		if err := q.mapper.Scan(rows, elem.Addr().Interface().(MappedStruct)); err != nil {
			return fmt.Errorf("yago Query.All(): Error while scanning: %s", err)
		}
		if isPtr {
			results.Set(reflect.Append(results, elem.Addr()))
		} else {
			results.Set(reflect.Append(results, elem))
		}
	}
	if err != nil {
		return err
	}
	if rows.Next() {
		return fmt.Errorf("TooManyResultsError")
	}
	return nil
}

// Count change the columns to COUNT(*), execute the query and returns
// the result
func (q Query) Count(count interface{}) error {
	// XXX mapper should be able to return a list of pkey fields
	// XXX When qb supports COUNT(*), use it
	q.selectStmt = q.selectStmt.Select(qb.Count(
		q.mapper.Table().PrimaryCols()[0]),
	)
	row, err := q.SQLQuery()
	if err != nil {
		return err
	}
	defer row.Close()
	if !row.Next() {
		panic("No result")
	}

	err = row.Scan(count)
	if err != nil {
		return err
	}
	return nil
}

// Exists return true if any record matches the current query
func (q Query) Exists() (bool, error) {
	q.selectStmt = q.selectStmt.Select(qb.SQLText("1")).Limit(0, 1)
	row, err := q.SQLQuery()
	if err != nil {
		return false, err
	}
	defer row.Close()
	if !row.Next() {
		panic("No result")
	}
	var exists bool
	err = row.Scan(&exists)
	if err != nil {
		return false, err
	}
	return exists, nil
}
<|endoftext|>