<commit_before>package graph

import (
	"encoding/json"
	"fmt"
	"io/ioutil"

	"github.com/cayleygraph/cayley"
	"github.com/cayleygraph/cayley/graph"
	_ "github.com/cayleygraph/cayley/graph/kv/bolt" // registers bolt quadstore
	"github.com/cayleygraph/cayley/quad"
	"github.com/cayleygraph/cayley/schema"

	"nimona.io/pkg/object"
)

var (
	predicateDependsOn = quad.IRI("dependsOn")
)

type graphObject struct {
	// rdfType struct{}   `quad:"@type > Object"` // nolint: structcheck
	ID       quad.IRI   `quad:"@id"`
	NodeType string     `quad:"type"`
	Context  string     `quad:"context"`
	Display  string     `quad:"display,optional"`
	Parents  []quad.IRI `quad:"dependsOn < *,optional"`
	Data     string     `quad:"data,optional"`
}

func toGraphObject(v *object.Object) (*graphObject, error) {
	b, err := json.Marshal(v.ToMap())
	if err != nil {
		return nil, err
	}
	nType := "object:root"
	if len(v.GetParents()) > 0 {
		nType = "object"
	}
	o := &graphObject{
		ID:       quad.IRI(v.HashBase58()).Full().Short(),
		NodeType: nType,
		Context:  v.GetType(),
		Parents:  []quad.IRI{},
		Data:     string(b),
	}
	if d, ok := v.GetRaw("@display").(string); ok {
		o.Display = d
	}
	for _, p := range v.GetParents() {
		o.Parents = append(o.Parents, quad.IRI(p).Full().Short())
	}
	return o, nil
}

func fromGraphObject(v *graphObject) (*object.Object, error) {
	o := &object.Object{}
	m := map[string]interface{}{}
	if err := json.Unmarshal([]byte(v.Data), &m); err != nil {
		return nil, err
	}
	if err := o.FromMap(m); err != nil {
		return nil, err
	}
	if o.HashBase58() != string(v.ID) {
		return nil, fmt.Errorf(
			"expected hash does not match actual, %s, %s ",
			string(v.ID),
			o.HashBase58(),
		)
	}
	return o, nil
}

// NewCayley constructs a new object store given a quad store
func NewCayley(store *cayley.Handle) Store {
	sch := schema.NewConfig()
	return &Cayley{
		store:  store,
		schema: sch,
	}
}

// NewCayleyWithTempStore constructs a new object store with a temp bolt store
// NOTE: Cayley's in-memory store is not thread safe, please don't use it to
// replace this.
func NewCayleyWithTempStore() (Store, error) {
	dir, err := ioutil.TempDir("", "cayley-temp")
	if err != nil {
		return nil, err
	}

	err = graph.InitQuadStore("bolt", dir, nil)
	if err != nil {
		return nil, err
	}

	cs, err := cayley.NewGraph("bolt", dir, nil)
	if err != nil {
		return nil, err
	}

	sch := schema.NewConfig()
	return &Cayley{
		store:  cs,
		schema: sch,
	}, nil
}

// Cayley object store
type Cayley struct {
	store  *cayley.Handle
	schema *schema.Config
}

// Put an object in the store
func (s *Cayley) Put(v *object.Object) error {
	qw := graph.NewWriter(s.store)
	defer func() {
		if err := qw.Flush(); err != nil {
			// TODO(geoah) log error
			fmt.Println("error flushing qw", err)
		}
		if err := qw.Close(); err != nil {
			// TODO(geoah) log error
			fmt.Println("error closing qw", err)
		}
	}()
	g, err := toGraphObject(v)
	if err != nil {
		return err
	}
	if _, err := s.schema.WriteAsQuads(qw, g); err != nil {
		return err
	}
	return nil
}

// Graph returns all objects in a graph given the hash of of its objects
func (s *Cayley) Graph(hash string) ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		cayley.
			StartMorphism().
			Or(
				cayley.
					StartMorphism().
					In(predicateDependsOn),
			).
			Or(
				cayley.
					StartMorphism().
					Out(predicateDependsOn),
			),
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, err
}

// Get returns an object given its hash
func (s *Cayley) Get(hash string) (*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	)

	g := graphObject{}
	err := schema.LoadPathTo(nil, s.store, &g, p)
	if err != nil {
		return nil, err
	}

	o, gerr := fromGraphObject(&g)
	if gerr != nil {
		return nil, gerr
	}

	return o, err
}

// Children returns the all the children of an object, sorted
func (s *Cayley) Children(hash string) ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).Or(
		cayley.StartPath(
			s.store,
			quad.IRI(hash),
		).FollowRecursive(
			predicateDependsOn,
			0,
			nil,
		),
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, err
}

// Head returns the root object given another object in the graph.
func (s *Cayley) Head(hash string) (*object.Object, error) {
	// TODO(geoah) Figure out how to use cayley to get the graph's head
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		cayley.
			StartMorphism().
			In(predicateDependsOn),
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	if len(gs) == 0 {
		return nil, ErrNotFound
	}

	return fromGraphObject(&gs[len(gs)-1])
}

// Heads returns all the objects that do not have any parents
func (s *Cayley) Heads() ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
	).Has(
		quad.IRI("type"),
		quad.String("object:root"),
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, nil
}

// Dump returns all objects
func (s *Cayley) Dump() ([]graphObject, error) {
	p := cayley.StartPath(
		s.store,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	fmt.Println(dot(gs))

	return gs, nil
}

// Tails returns all the leaf nodes starting from an object
func (s *Cayley) Tails(hash string) ([]*object.Object, error) {
	// TODO(geoah) Figure out how to use cayley to get the graph's leaves
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		predicateDependsOn,
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	// figure out tail nodes
	ts := map[string]bool{}
	for i := range gs {
		g := gs[i]
		if _, ok := ts[g.ID.String()]; !ok {
			ts[g.ID.String()] = true
		}
		for _, p := range g.Parents {
			ts[p.String()] = false
		}
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		if !ts[g.ID.String()] {
			continue
		}
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, nil
}
<commit_msg>fix(store/graph): return ErrNotFound when Graph() has no results<commit_after>package graph

import (
	"encoding/json"
	"fmt"
	"io/ioutil"

	"github.com/cayleygraph/cayley"
	"github.com/cayleygraph/cayley/graph"
	_ "github.com/cayleygraph/cayley/graph/kv/bolt" // registers bolt quadstore
	"github.com/cayleygraph/cayley/quad"
	"github.com/cayleygraph/cayley/schema"

	"nimona.io/pkg/object"
)

var (
	predicateDependsOn = quad.IRI("dependsOn")
)

type graphObject struct {
	// rdfType struct{}   `quad:"@type > Object"` // nolint: structcheck
	ID       quad.IRI   `quad:"@id"`
	NodeType string     `quad:"type"`
	Context  string     `quad:"context"`
	Display  string     `quad:"display,optional"`
	Parents  []quad.IRI `quad:"dependsOn < *,optional"`
	Data     string     `quad:"data,optional"`
}

func toGraphObject(v *object.Object) (*graphObject, error) {
	b, err := json.Marshal(v.ToMap())
	if err != nil {
		return nil, err
	}
	nType := "object:root"
	if len(v.GetParents()) > 0 {
		nType = "object"
	}
	o := &graphObject{
		ID:       quad.IRI(v.HashBase58()).Full().Short(),
		NodeType: nType,
		Context:  v.GetType(),
		Parents:  []quad.IRI{},
		Data:     string(b),
	}
	if d, ok := v.GetRaw("@display").(string); ok {
		o.Display = d
	}
	for _, p := range v.GetParents() {
		o.Parents = append(o.Parents, quad.IRI(p).Full().Short())
	}
	return o, nil
}

func fromGraphObject(v *graphObject) (*object.Object, error) {
	o := &object.Object{}
	m := map[string]interface{}{}
	if err := json.Unmarshal([]byte(v.Data), &m); err != nil {
		return nil, err
	}
	if err := o.FromMap(m); err != nil {
		return nil, err
	}
	if o.HashBase58() != string(v.ID) {
		return nil, fmt.Errorf(
			"expected hash does not match actual, %s, %s ",
			string(v.ID),
			o.HashBase58(),
		)
	}
	return o, nil
}

// NewCayley constructs a new object store given a quad store
func NewCayley(store *cayley.Handle) Store {
	sch := schema.NewConfig()
	return &Cayley{
		store:  store,
		schema: sch,
	}
}

// NewCayleyWithTempStore constructs a new object store with a temp bolt store
// NOTE: Cayley's in-memory store is not thread safe, please don't use it to
// replace this.
func NewCayleyWithTempStore() (Store, error) {
	dir, err := ioutil.TempDir("", "cayley-temp")
	if err != nil {
		return nil, err
	}

	err = graph.InitQuadStore("bolt", dir, nil)
	if err != nil {
		return nil, err
	}

	cs, err := cayley.NewGraph("bolt", dir, nil)
	if err != nil {
		return nil, err
	}

	sch := schema.NewConfig()
	return &Cayley{
		store:  cs,
		schema: sch,
	}, nil
}

// Cayley object store
type Cayley struct {
	store  *cayley.Handle
	schema *schema.Config
}

// Put an object in the store
func (s *Cayley) Put(v *object.Object) error {
	qw := graph.NewWriter(s.store)
	defer func() {
		if err := qw.Flush(); err != nil {
			// TODO(geoah) log error
			fmt.Println("error flushing qw", err)
		}
		if err := qw.Close(); err != nil {
			// TODO(geoah) log error
			fmt.Println("error closing qw", err)
		}
	}()
	g, err := toGraphObject(v)
	if err != nil {
		return err
	}
	if _, err := s.schema.WriteAsQuads(qw, g); err != nil {
		return err
	}
	return nil
}

// Graph returns all objects in a graph given the hash of of its objects
func (s *Cayley) Graph(hash string) ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		cayley.
			StartMorphism().
			Or(
				cayley.
					StartMorphism().
					In(predicateDependsOn),
			).
			Or(
				cayley.
					StartMorphism().
					Out(predicateDependsOn),
			),
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	if len(os) == 0 {
		return nil, ErrNotFound
	}

	return os, nil
}

// Get returns an object given its hash
func (s *Cayley) Get(hash string) (*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	)

	g := graphObject{}
	err := schema.LoadPathTo(nil, s.store, &g, p)
	if err != nil {
		return nil, err
	}

	o, gerr := fromGraphObject(&g)
	if gerr != nil {
		return nil, gerr
	}

	return o, err
}

// Children returns the all the children of an object, sorted
func (s *Cayley) Children(hash string) ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).Or(
		cayley.StartPath(
			s.store,
			quad.IRI(hash),
		).FollowRecursive(
			predicateDependsOn,
			0,
			nil,
		),
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, err
}

// Head returns the root object given another object in the graph.
func (s *Cayley) Head(hash string) (*object.Object, error) {
	// TODO(geoah) Figure out how to use cayley to get the graph's head
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		cayley.
			StartMorphism().
			In(predicateDependsOn),
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	if len(gs) == 0 {
		return nil, ErrNotFound
	}

	return fromGraphObject(&gs[len(gs)-1])
}

// Heads returns all the objects that do not have any parents
func (s *Cayley) Heads() ([]*object.Object, error) {
	p := cayley.StartPath(
		s.store,
	).Has(
		quad.IRI("type"),
		quad.String("object:root"),
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, nil
}

// Dump returns all objects
func (s *Cayley) Dump() ([]graphObject, error) {
	p := cayley.StartPath(
		s.store,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	fmt.Println(dot(gs))

	return gs, nil
}

// Tails returns all the leaf nodes starting from an object
func (s *Cayley) Tails(hash string) ([]*object.Object, error) {
	// TODO(geoah) Figure out how to use cayley to get the graph's leaves
	p := cayley.StartPath(
		s.store,
		quad.IRI(hash),
	).FollowRecursive(
		predicateDependsOn,
		0,
		nil,
	)

	gs := []graphObject{}
	err := schema.LoadPathTo(nil, s.store, &gs, p)
	if err != nil {
		return nil, err
	}

	// figure out tail nodes
	ts := map[string]bool{}
	for i := range gs {
		g := gs[i]
		if _, ok := ts[g.ID.String()]; !ok {
			ts[g.ID.String()] = true
		}
		for _, p := range g.Parents {
			ts[p.String()] = false
		}
	}

	os := []*object.Object{}
	for i := range gs {
		g := gs[i]
		if !ts[g.ID.String()] {
			continue
		}
		o, gerr := fromGraphObject(&g)
		if gerr != nil {
			return nil, gerr
		}
		os = append(os, o)
	}

	return os, nil
}
<|endoftext|>


<commit_before>package consensus

import (
	"bytes"
	"encoding/binary"
	"math/big"

	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/bolt"
	"github.com/NebulousLabs/errors"
)

// Errors returned by this file.
var (
	// errOakHardforkIncompatibility is the error returned if Oak initialization
	// cannot begin because the consensus database was not upgraded before the
	// hardfork height.
	errOakHardforkIncompatibility = errors.New("difficulty adjustment hardfork incompatibility detected")
)

// difficulty.go defines the Oak difficulty adjustment algorithm. Past the
// hardfork trigger height, it the algorithm that Sia uses to adjust the
// difficulty.
//
// A running tally is maintained which keeps the total difficulty and total time
// passed across all blocks. The total difficulty can be divided by the total
// time to get a hashrate. The total is multiplied by 0.995 each block, to keep
// exponential preference on recent blocks with a half life of about 24 hours.
// This estimated hashrate is assumed to closely match the actual hashrate on
// the network.
//
// There is a target block time. If the difficulty increases or decreases, the
// total amount of time that has passed will be more or less than the target
// amount of time passed for the current height. To counteract this, the target
// block time for each block is adjusted based on how far away from the desired
// total time passed the current total time passed is. If the total time passed
// is too low, blocks are targeted to be slightly longer, which helps to correct
// the network. And if the total time passed is too high, blocks are targeted to
// be slightly shorter, to help correct the network.
//
// High variance in block times means that the corrective action should not be
// very strong if the total time passed has only missed the target time passed
// by a few hours. But if the total time passed is significantly off, the block
// time corrections should be much stronger. The square of the total deviation
// is used to figure out what the adjustment should be. At 10,000 seconds
// variance (about 3 hours), blocks will be adjusted by 10 seconds each. At
// 20,000 seconds, blocks will be adjusted by 40 seconds each, a 4x adjustment
// for 2x the error. And at 40,000 seconds, blocks will be adjusted by 160
// seconds each, and so on.
//
// The total amount of blocktime adjustment is capped to 1/3 and 3x the target
// blocktime, to prevent too much disruption on the network. If blocks are
// actually coming out 3x as fast as intended, there will be a (temporary)
// significant increase on the amount of strain on nodes to process blocks. And
// at 1/3 the target blocktime, the total blockchain throughput will decrease
// dramatically.
//
// Finally, one extra cap is applied to the difficulty adjustment - the
// difficulty of finding a block is not allowed to change more than 0.4% every
// block. This maps to a total possible difficulty change of 55x across 1008
// blocks. This clamp helps to prevent wild swings when the hashrate increases
// or decreases rapidly on the network, and it also limits the amount of damange
// that a malicious attacker can do if performing a difficulty raising attack.

// childTargetOak sets the child target based on the total time delta and total
// hashrate of the parent block. The deltas are known for the child block,
// however we do not use the child block deltas because that would allow the
// child block to influence the target of the following block, which makes abuse
// easier in selfish mining scenarios.
func (cs *ConsensusSet) childTargetOak(parentTotalTime int64, parentTotalTarget, currentTarget types.Target, parentHeight types.BlockHeight) types.Target {
	// Determine the detla of the current total time vs. the desired total time.
	expectedTime := types.BlockFrequency * parentHeight
	delta := int64(expectedTime) - parentTotalTime
	// Convert the delta in to a target block time.
	square := delta * delta
	if delta < 0 {
		// If the delta is negative, restore the negative value.
		square *= -1
	}
	shift := square / 10e6 // 10e3 second delta leads to 10 second shift.
	targetBlockTime := int64(types.BlockFrequency) + shift

	// Clamp the block time to 1/3 and 3x the target block time.
	if targetBlockTime < int64(types.BlockFrequency)/3 {
		targetBlockTime = int64(types.BlockFrequency) / 3
	}
	if targetBlockTime > int64(types.BlockFrequency)*3 {
		targetBlockTime = int64(types.BlockFrequency) * 3
	}

	// Determine the hashrate using the total time and total target. Set a
	// minimum total time of 1 to prevent divide by zero and underflows.
	if parentTotalTime < 1 {
		parentTotalTime = 1
	}
	visibleHashrate := parentTotalTarget.Difficulty().Div64(uint64(parentTotalTime)) // Hashes per second.
	if visibleHashrate.IsZero() {
		visibleHashrate = visibleHashrate.Add(types.NewCurrency64(1))
	}

	// Determine the new target by multiplying the visible hashrate by the
	// target block time. Clamp it to a 0.4% difficulty adjustment.
	maxNewTarget := currentTarget.MulDifficulty(types.OakMaxRise)
	minNewTarget := currentTarget.MulDifficulty(types.OakMaxDrop)
	newTarget := types.RatToTarget(new(big.Rat).SetFrac(types.RootDepth.Int(), visibleHashrate.Mul64(uint64(targetBlockTime)).Big()))
	if newTarget.Cmp(maxNewTarget) < 0 {
		newTarget = maxNewTarget
	}
	if newTarget.Cmp(minNewTarget) > 0 {
		newTarget = minNewTarget
	}
	return newTarget
}

// getBlockTotals returns the block totals values that get stored in
// storeBlockTotals.
func (cs *ConsensusSet) getBlockTotals(tx *bolt.Tx, id types.BlockID) (totalTime int64, totalTarget types.Target) {
	totalsBytes := tx.Bucket(BucketOak).Get(id[:])
	totalTime = int64(binary.LittleEndian.Uint64(totalsBytes[:8]))
	copy(totalTarget[:], totalsBytes[8:])
	return
}

// storeBlockTotals computes the new total time and total target for the current
// block and stores that new time in the database. It also returns the new
// totals.
func (cs *ConsensusSet) storeBlockTotals(tx *bolt.Tx, currentHeight types.BlockHeight, currentBlockID types.BlockID, prevTotalTime int64, parentTimestamp, currentTimestamp types.Timestamp, prevTotalTarget, targetOfCurrentBlock types.Target) (newTotalTime int64, newTotalTarget types.Target, err error) {
	// Reset the prevTotalTime to a delta of zero just before the hardfork.
	if currentHeight == types.OakHardforkBlock-1 {
		prevTotalTime = int64(types.BlockFrequency * currentHeight)
	}

	// For each value, first multiply by the decay, and then add in the new
	// delta.
	newTotalTime = (prevTotalTime * types.OakDecayNum / types.OakDecayDenom) + (int64(currentTimestamp) - int64(parentTimestamp))
	newTotalTarget = prevTotalTarget.MulDifficulty(big.NewRat(types.OakDecayNum, types.OakDecayDenom)).AddDifficulties(targetOfCurrentBlock)

	// Store the new total time and total target in the database at the
	// appropriate id.
	bytes := make([]byte, 40)
	binary.LittleEndian.PutUint64(bytes[:8], uint64(newTotalTime))
	copy(bytes[8:], newTotalTarget[:])
	err = tx.Bucket(BucketOak).Put(currentBlockID[:], bytes)
	if err != nil {
		return 0, types.Target{}, errors.Extend(errors.New("unable to store total time values"), err)
	}
	return newTotalTime, newTotalTarget, nil
}

// initOak will initialize all of the oak difficulty adjustment related fields.
// This is separate from the initialization process for compatibility reasons -
// some databases will not have these fields at start, so it much be checked.
//
// After oak initialization is complete, a specific field in the oak bucket is
// marked so that oak initialization can be skipped in the future.
func (cs *ConsensusSet) initOak(tx *bolt.Tx) error {
	// Prep the oak bucket.
	bucketOak, err := tx.CreateBucketIfNotExists(BucketOak)
	if err != nil {
		return errors.Extend(errors.New("unable to create oak bucket"), err)
	}
	// Check whether the init field is set.
	if bytes.Equal(bucketOak.Get(FieldOakInit), ValueOakInit) {
		// The oak fields have been initialized, nothing to do.
		return nil
	}

	// If the current height is greater than the hardfork trigger date, return
	// an error and refuse to initialize.
	height := blockHeight(tx)
	if height > types.OakHardforkBlock {
		return errOakHardforkIncompatibility
	}

	// Store base values for the genesis block.
	totalTime, totalTarget, err := cs.storeBlockTotals(tx, 0, types.GenesisID, 0, types.GenesisTimestamp, types.GenesisTimestamp, types.RootDepth, types.RootTarget)
	if err != nil {
		return errors.Extend(errors.New("unable to store genesis block totals"), err)
	}

	// The Oak fields have not been initialized, scan through the consensus set
	// and set the fields for each block.
	parentTimestamp := types.GenesisTimestamp
	parentChildTarget := types.RootTarget
	for i := types.BlockHeight(1); i <= height; i++ { // Skip Genesis block
		// Fetch the processed block for the current block.
		id, err := getPath(tx, i)
		if err != nil {
			return errors.Extend(errors.New("unable to find block at height"), err)
		}
		pb, err := getBlockMap(tx, id)
		if err != nil {
			return errors.Extend(errors.New("unable to find block from id"), err)
		}

		// Calculate and store the new block totals.
		totalTime, totalTarget, err = cs.storeBlockTotals(tx, i, id, totalTime, parentTimestamp, pb.Block.Timestamp, totalTarget, parentChildTarget)
		if err != nil {
			return errors.Extend(errors.New("unable to store updated block totals"), err)
		}
		// Update the previous values.
		parentTimestamp = pb.Block.Timestamp
		parentChildTarget = pb.ChildTarget
	}

	// Tag the initialization field in the oak bucket, indicating that
	// initialization has completed.
	err = bucketOak.Put(FieldOakInit, ValueOakInit)
	if err != nil {
		return errors.Extend(errors.New("unable to put oak init confirmation into oak bucket"), err)
	}
	return nil
}
<commit_msg>add comment about half-life being 144 data points<commit_after>package consensus

import (
	"bytes"
	"encoding/binary"
	"math/big"

	"github.com/NebulousLabs/Sia/types"

	"github.com/NebulousLabs/bolt"
	"github.com/NebulousLabs/errors"
)

// Errors returned by this file.
var (
	// errOakHardforkIncompatibility is the error returned if Oak initialization
	// cannot begin because the consensus database was not upgraded before the
	// hardfork height.
	errOakHardforkIncompatibility = errors.New("difficulty adjustment hardfork incompatibility detected")
)

// difficulty.go defines the Oak difficulty adjustment algorithm. Past the
// hardfork trigger height, it the algorithm that Sia uses to adjust the
// difficulty.
//
// A running tally is maintained which keeps the total difficulty and total time
// passed across all blocks. The total difficulty can be divided by the total
// time to get a hashrate. The total is multiplied by 0.995 each block, to keep
// exponential preference on recent blocks with a half life of 144 data points.
// This is about 24 hours. This estimated hashrate is assumed to closely match
// the actual hashrate on the network.
//
// There is a target block time. If the difficulty increases or decreases, the
// total amount of time that has passed will be more or less than the target
// amount of time passed for the current height. To counteract this, the target
// block time for each block is adjusted based on how far away from the desired
// total time passed the current total time passed is. If the total time passed
// is too low, blocks are targeted to be slightly longer, which helps to correct
// the network. And if the total time passed is too high, blocks are targeted to
// be slightly shorter, to help correct the network.
//
// High variance in block times means that the corrective action should not be
// very strong if the total time passed has only missed the target time passed
// by a few hours. But if the total time passed is significantly off, the block
// time corrections should be much stronger. The square of the total deviation
// is used to figure out what the adjustment should be. At 10,000 seconds
// variance (about 3 hours), blocks will be adjusted by 10 seconds each. At
// 20,000 seconds, blocks will be adjusted by 40 seconds each, a 4x adjustment
// for 2x the error. And at 40,000 seconds, blocks will be adjusted by 160
// seconds each, and so on.
//
// The total amount of blocktime adjustment is capped to 1/3 and 3x the target
// blocktime, to prevent too much disruption on the network. If blocks are
// actually coming out 3x as fast as intended, there will be a (temporary)
// significant increase on the amount of strain on nodes to process blocks. And
// at 1/3 the target blocktime, the total blockchain throughput will decrease
// dramatically.
//
// Finally, one extra cap is applied to the difficulty adjustment - the
// difficulty of finding a block is not allowed to change more than 0.4% every
// block. This maps to a total possible difficulty change of 55x across 1008
// blocks. This clamp helps to prevent wild swings when the hashrate increases
// or decreases rapidly on the network, and it also limits the amount of damange
// that a malicious attacker can do if performing a difficulty raising attack.

// childTargetOak sets the child target based on the total time delta and total
// hashrate of the parent block. The deltas are known for the child block,
// however we do not use the child block deltas because that would allow the
// child block to influence the target of the following block, which makes abuse
// easier in selfish mining scenarios.
func (cs *ConsensusSet) childTargetOak(parentTotalTime int64, parentTotalTarget, currentTarget types.Target, parentHeight types.BlockHeight) types.Target {
	// Determine the detla of the current total time vs. the desired total time.
	expectedTime := types.BlockFrequency * parentHeight
	delta := int64(expectedTime) - parentTotalTime
	// Convert the delta in to a target block time.
	square := delta * delta
	if delta < 0 {
		// If the delta is negative, restore the negative value.
		square *= -1
	}
	shift := square / 10e6 // 10e3 second delta leads to 10 second shift.
	targetBlockTime := int64(types.BlockFrequency) + shift

	// Clamp the block time to 1/3 and 3x the target block time.
	if targetBlockTime < int64(types.BlockFrequency)/3 {
		targetBlockTime = int64(types.BlockFrequency) / 3
	}
	if targetBlockTime > int64(types.BlockFrequency)*3 {
		targetBlockTime = int64(types.BlockFrequency) * 3
	}

	// Determine the hashrate using the total time and total target. Set a
	// minimum total time of 1 to prevent divide by zero and underflows.
	if parentTotalTime < 1 {
		parentTotalTime = 1
	}
	visibleHashrate := parentTotalTarget.Difficulty().Div64(uint64(parentTotalTime)) // Hashes per second.
	if visibleHashrate.IsZero() {
		visibleHashrate = visibleHashrate.Add(types.NewCurrency64(1))
	}

	// Determine the new target by multiplying the visible hashrate by the
	// target block time. Clamp it to a 0.4% difficulty adjustment.
	maxNewTarget := currentTarget.MulDifficulty(types.OakMaxRise)
	minNewTarget := currentTarget.MulDifficulty(types.OakMaxDrop)
	newTarget := types.RatToTarget(new(big.Rat).SetFrac(types.RootDepth.Int(), visibleHashrate.Mul64(uint64(targetBlockTime)).Big()))
	if newTarget.Cmp(maxNewTarget) < 0 {
		newTarget = maxNewTarget
	}
	if newTarget.Cmp(minNewTarget) > 0 {
		newTarget = minNewTarget
	}
	return newTarget
}

// getBlockTotals returns the block totals values that get stored in
// storeBlockTotals.
func (cs *ConsensusSet) getBlockTotals(tx *bolt.Tx, id types.BlockID) (totalTime int64, totalTarget types.Target) {
	totalsBytes := tx.Bucket(BucketOak).Get(id[:])
	totalTime = int64(binary.LittleEndian.Uint64(totalsBytes[:8]))
	copy(totalTarget[:], totalsBytes[8:])
	return
}

// storeBlockTotals computes the new total time and total target for the current
// block and stores that new time in the database. It also returns the new
// totals.
func (cs *ConsensusSet) storeBlockTotals(tx *bolt.Tx, currentHeight types.BlockHeight, currentBlockID types.BlockID, prevTotalTime int64, parentTimestamp, currentTimestamp types.Timestamp, prevTotalTarget, targetOfCurrentBlock types.Target) (newTotalTime int64, newTotalTarget types.Target, err error) {
	// Reset the prevTotalTime to a delta of zero just before the hardfork.
	if currentHeight == types.OakHardforkBlock-1 {
		prevTotalTime = int64(types.BlockFrequency * currentHeight)
	}

	// For each value, first multiply by the decay, and then add in the new
	// delta.
	newTotalTime = (prevTotalTime * types.OakDecayNum / types.OakDecayDenom) + (int64(currentTimestamp) - int64(parentTimestamp))
	newTotalTarget = prevTotalTarget.MulDifficulty(big.NewRat(types.OakDecayNum, types.OakDecayDenom)).AddDifficulties(targetOfCurrentBlock)

	// Store the new total time and total target in the database at the
	// appropriate id.
	bytes := make([]byte, 40)
	binary.LittleEndian.PutUint64(bytes[:8], uint64(newTotalTime))
	copy(bytes[8:], newTotalTarget[:])
	err = tx.Bucket(BucketOak).Put(currentBlockID[:], bytes)
	if err != nil {
		return 0, types.Target{}, errors.Extend(errors.New("unable to store total time values"), err)
	}
	return newTotalTime, newTotalTarget, nil
}

// initOak will initialize all of the oak difficulty adjustment related fields.
// This is separate from the initialization process for compatibility reasons -
// some databases will not have these fields at start, so it much be checked.
//
// After oak initialization is complete, a specific field in the oak bucket is
// marked so that oak initialization can be skipped in the future.
func (cs *ConsensusSet) initOak(tx *bolt.Tx) error {
	// Prep the oak bucket.
	bucketOak, err := tx.CreateBucketIfNotExists(BucketOak)
	if err != nil {
		return errors.Extend(errors.New("unable to create oak bucket"), err)
	}
	// Check whether the init field is set.
	if bytes.Equal(bucketOak.Get(FieldOakInit), ValueOakInit) {
		// The oak fields have been initialized, nothing to do.
		return nil
	}

	// If the current height is greater than the hardfork trigger date, return
	// an error and refuse to initialize.
	height := blockHeight(tx)
	if height > types.OakHardforkBlock {
		return errOakHardforkIncompatibility
	}

	// Store base values for the genesis block.
	totalTime, totalTarget, err := cs.storeBlockTotals(tx, 0, types.GenesisID, 0, types.GenesisTimestamp, types.GenesisTimestamp, types.RootDepth, types.RootTarget)
	if err != nil {
		return errors.Extend(errors.New("unable to store genesis block totals"), err)
	}

	// The Oak fields have not been initialized, scan through the consensus set
	// and set the fields for each block.
	parentTimestamp := types.GenesisTimestamp
	parentChildTarget := types.RootTarget
	for i := types.BlockHeight(1); i <= height; i++ { // Skip Genesis block
		// Fetch the processed block for the current block.
		id, err := getPath(tx, i)
		if err != nil {
			return errors.Extend(errors.New("unable to find block at height"), err)
		}
		pb, err := getBlockMap(tx, id)
		if err != nil {
			return errors.Extend(errors.New("unable to find block from id"), err)
		}

		// Calculate and store the new block totals.
		totalTime, totalTarget, err = cs.storeBlockTotals(tx, i, id, totalTime, parentTimestamp, pb.Block.Timestamp, totalTarget, parentChildTarget)
		if err != nil {
			return errors.Extend(errors.New("unable to store updated block totals"), err)
		}
		// Update the previous values.
		parentTimestamp = pb.Block.Timestamp
		parentChildTarget = pb.ChildTarget
	}

	// Tag the initialization field in the oak bucket, indicating that
	// initialization has completed.
	err = bucketOak.Put(FieldOakInit, ValueOakInit)
	if err != nil {
		return errors.Extend(errors.New("unable to put oak init confirmation into oak bucket"), err)
	}
	return nil
}
<|endoftext|>


<commit_before>package deploy

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/cenk/backoff"
	"github.com/fatih/color"
	marathon "github.com/gambol99/go-marathon"

	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils"
)

func init() {
	manifest.PluginRegestry.Add("deploy.marathon", DeployMarathon{})
}

type DeployMarathon struct{}

func (p DeployMarathon) Run(data manifest.Manifest) error {
	if data.GetBool("purge") {
		return p.Uninstall(data)
	} else {
		return p.Install(data)
	}
}

func (p DeployMarathon) Install(data manifest.Manifest) error {
	marathonApi, err := MarathonClient(data.GetString("marathon-address"))
	if err != nil {
		return err
	}

	fullName := data.GetString("app-name")

	backoffSeconds := data.GetFloat("backoff-seconds")
	backoffFactor := data.GetFloat("backoff-factor")
	maxLaunchDelaySeconds := data.GetFloat("max-launch-delay-seconds")
	taskKillGracePeriodSeconds := data.GetFloat("task-kill-grace-period-seconds")

	app := &marathon.Application{
		User:                       data.GetString("user"),
		BackoffSeconds:             &backoffSeconds,
		BackoffFactor:              &backoffFactor,
		MaxLaunchDelaySeconds:      &maxLaunchDelaySeconds,
		TaskKillGracePeriodSeconds: &taskKillGracePeriodSeconds,
		UpgradeStrategy: &marathon.UpgradeStrategy{
			MinimumHealthCapacity: data.GetFloat("min-health-capacity"),
			MaximumOverCapacity:   data.GetFloat("max-over-capacity"),
		},
	}

	portArgs := ""
	if port := data.GetString("listen-port"); port != "" {
		portArgs = "--port " + port
	}

	app.Name(fullName)
	app.Command(fmt.Sprintf("exec serve-tools consul supervisor --service '%s' %s start %s", fullName, portArgs, data.GetString("cmd")))
	app.Count(data.GetInt("instances"))
	app.Memory(float64(data.GetInt("mem")))

	if cpu, err := strconv.ParseFloat(data.GetString("cpu"), 64); err == nil {
		app.CPU(cpu)
	}

	if cluster := data.GetString("cluster"); cluster != "" {
		cs := strings.SplitN(cluster, ":", 2)
		if len(cs) < 2 {
			cs = append(cs, "true")
		}
		app.AddConstraint(cs[0], "CLUSTER", cs[1])
		app.AddLabel(cs[0], cs[1])
	}

	for _, cons := range data.GetArray("constraints") {
		if consArr, ok := cons.Unwrap().([]interface{}); ok {
			consStrings := make([]string, len(consArr))
			for i, c := range consArr {
				consStrings[i] = fmt.Sprintf("%s", c)
			}
			app.AddConstraint(consStrings...)
		}
	}

	for _, port := range data.GetArray("ports") {
		app.AddPortDefinition(marathon.PortDefinition{Name: port.GetStringOr("name", "")}.SetPort(port.GetIntOr("port", 0)))
	}

	app.AddEnv("SERVICE_DEPLOY_TIME", time.Now().Format(time.RFC3339)) // force redeploy app

	for k, v := range data.GetMap("environment") {
		app.AddEnv(k, fmt.Sprintf("%s", v.Unwrap()))
	}

	app.AddUris(data.GetString("package-uri"))

	if data.GetBool("docker.enabled") {
		app.Cmd = nil
		app.EmptyUris()
		app.EmptyPortDefinitions()

		app.AddEnv("SERVICE_NAME", fullName)    // docker-registrator: consul Name
		app.AddEnv("SERVICE_CHECK_TCP", "true") // docker-registrator: consul TCP health check

		doc := marathon.NewDockerContainer()
		doc.Docker.Image = data.GetString("docker.image")
		doc.Docker.Network = data.GetString("docker.network")
		doc.Docker.SetForcePullImage(true)
		doc.EmptyVolumes()

		for _, port := range data.GetArray("docker.ports") {
			doc.Docker.ExposePort(marathon.PortMapping{
				ContainerPort: port.GetInt("containerPort"),
				HostPort:      port.GetIntOr("hostPort", 0),
				Protocol:      "tcp",
			})
		}

		for _, vol := range data.GetArray("docker.volumes") {
			doc.Volume(vol.GetString("hostPath"), vol.GetString("containerPath"), vol.GetString("mode"))
		}

		for key, value := range data.GetMap("docker.parameters") {
			switch key {
			case "env":
				envs := make([]string, 0)
				for k, v := range data.GetMap("docker.parameters.env") {
					envs = append(envs, fmt.Sprintf("%s=%s", k, v))
				}
				doc.Docker.AddParameter(key, fmt.Sprintf("[%s]", strings.Join(envs, ",")))
			default:
				doc.Docker.AddParameter(key, value.Unwrap().(string))
			}
		}

		app.Container = doc
	}

	if data.GetString("listen-port") == "$PORT0" {
		health := marathon.NewDefaultHealthCheck()
		health.Protocol = "TCP"
		app.AddHealthCheck(*health)
	}

	if _, err := marathonApi.UpdateApplication(app, false); err != nil {
		color.Yellow("marathon <- %s", app)
		return err
	}

	color.Green("marathon <- %s", app)

	consulApi, err := utils.ConsulClient(data.GetString("consul-address"))
	if err != nil {
		return err
	}

	if err := utils.RegisterPluginData("deploy.marathon", data.GetString("app-name"), data.String(), data.GetString("consul-address")); err != nil {
		return err
	}

	return backoff.Retry(func() error {
		services, _, err := consulApi.Health().Service(fullName, "", true, nil)

		if err != nil {
			log.Println(color.RedString("Error in check health in consul: %v", err))
			return err
		}

		if len(services) == 0 {
			log.Printf("Service `%s` not started yet! Retry...", fullName)
			return fmt.Errorf("Service `%s` not started!", fullName)
		}

		log.Println(color.GreenString("Service `%s` successfully started!", fullName))
		return nil
	}, backoff.NewExponentialBackOff())
}

func (p DeployMarathon) Uninstall(data manifest.Manifest) error {
	marathonApi, err := MarathonClient(data.GetString("marathon-address"))
	if err != nil {
		return err
	}

	name := data.GetString("app-name")

	if _, err := marathonApi.Application(name); err == nil {
		if _, err := marathonApi.DeleteApplication(name, false); err != nil {
			return err
		}
	} else {
		log.Println(color.YellowString("App `%s` doesnt exists in marathon!", name))
	}

	return utils.DeletePluginData("deploy.marathon", name, data.GetString("consul-address"))
}

func MarathonClient(marathonAddress string) (marathon.Marathon, error) {
	conf := marathon.NewDefaultConfig()
	conf.URL = fmt.Sprintf("http://%s", marathonAddress)
	conf.LogOutput = os.Stderr
	return marathon.NewClient(conf)
}
<commit_msg>add env for docker<commit_after>package deploy

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/cenk/backoff"
	"github.com/fatih/color"
	marathon "github.com/gambol99/go-marathon"

	"github.com/InnovaCo/serve/manifest"
	"github.com/InnovaCo/serve/utils"
)

func init() {
	manifest.PluginRegestry.Add("deploy.marathon", DeployMarathon{})
}

type DeployMarathon struct{}

func (p DeployMarathon) Run(data manifest.Manifest) error {
	if data.GetBool("purge") {
		return p.Uninstall(data)
	} else {
		return p.Install(data)
	}
}

func (p DeployMarathon) Install(data manifest.Manifest) error {
	marathonApi, err := MarathonClient(data.GetString("marathon-address"))
	if err != nil {
		return err
	}

	fullName := data.GetString("app-name")

	backoffSeconds := data.GetFloat("backoff-seconds")
	backoffFactor := data.GetFloat("backoff-factor")
	maxLaunchDelaySeconds := data.GetFloat("max-launch-delay-seconds")
	taskKillGracePeriodSeconds := data.GetFloat("task-kill-grace-period-seconds")

	app := &marathon.Application{
		User:                       data.GetString("user"),
		BackoffSeconds:             &backoffSeconds,
		BackoffFactor:              &backoffFactor,
		MaxLaunchDelaySeconds:      &maxLaunchDelaySeconds,
		TaskKillGracePeriodSeconds: &taskKillGracePeriodSeconds,
		UpgradeStrategy: &marathon.UpgradeStrategy{
			MinimumHealthCapacity: data.GetFloat("min-health-capacity"),
			MaximumOverCapacity:   data.GetFloat("max-over-capacity"),
		},
	}

	portArgs := ""
	if port := data.GetString("listen-port"); port != "" {
		portArgs = "--port " + port
	}

	app.Name(fullName)
	app.Command(fmt.Sprintf("exec serve-tools consul supervisor --service '%s' %s start %s", fullName, portArgs, data.GetString("cmd")))
	app.Count(data.GetInt("instances"))
	app.Memory(float64(data.GetInt("mem")))

	if cpu, err := strconv.ParseFloat(data.GetString("cpu"), 64); err == nil {
		app.CPU(cpu)
	}

	if cluster := data.GetString("cluster"); cluster != "" {
		cs := strings.SplitN(cluster, ":", 2)
		if len(cs) < 2 {
			cs = append(cs, "true")
		}
		app.AddConstraint(cs[0], "CLUSTER", cs[1])
		app.AddLabel(cs[0], cs[1])
	}

	for _, cons := range data.GetArray("constraints") {
		if consArr, ok := cons.Unwrap().([]interface{}); ok {
			consStrings := make([]string, len(consArr))
			for i, c := range consArr {
				consStrings[i] = fmt.Sprintf("%s", c)
			}
			app.AddConstraint(consStrings...)
		}
	}

	for _, port := range data.GetArray("ports") {
		app.AddPortDefinition(marathon.PortDefinition{Name: port.GetStringOr("name", "")}.SetPort(port.GetIntOr("port", 0)))
	}

	app.AddEnv("SERVICE_DEPLOY_TIME", time.Now().Format(time.RFC3339)) // force redeploy app

	for k, v := range data.GetMap("environment") {
		app.AddEnv(k, fmt.Sprintf("%s", v.Unwrap()))
	}

	app.AddUris(data.GetString("package-uri"))

	if data.GetBool("docker.enabled") {
		app.Cmd = nil
		app.EmptyUris()
		app.EmptyPortDefinitions()

		app.AddEnv("SERVICE_NAME", fullName)    // docker-registrator: consul Name
		app.AddEnv("SERVICE_CHECK_TCP", "true") // docker-registrator: consul TCP health check

		doc := marathon.NewDockerContainer()
		doc.Docker.Image = data.GetString("docker.image")
		doc.Docker.Network = data.GetString("docker.network")
		doc.Docker.SetForcePullImage(true)
		doc.EmptyVolumes()

		for _, port := range data.GetArray("docker.ports") {
			doc.Docker.ExposePort(marathon.PortMapping{
				ContainerPort: port.GetInt("containerPort"),
				HostPort:      port.GetIntOr("hostPort", 0),
				Protocol:      "tcp",
			})
		}

		for _, vol := range data.GetArray("docker.volumes") {
			doc.Volume(vol.GetString("hostPath"), vol.GetString("containerPath"), vol.GetString("mode"))
		}

		for key, value := range data.GetMap("docker.envs") {
			app.AddEnv(strings.ToUpper(key), value.Unwrap().(string))
		}

		for key, value := range data.GetMap("docker.parameters") {
			doc.Docker.AddParameter(key, value.Unwrap().(string))
		}

		app.Container = doc
	}

	if data.GetString("listen-port") == "$PORT0" {
		health := marathon.NewDefaultHealthCheck()
		health.Protocol = "TCP"
		app.AddHealthCheck(*health)
	}

	if _, err := marathonApi.UpdateApplication(app, false); err != nil {
		color.Yellow("marathon <- %s", app)
		return err
	}

	color.Green("marathon <- %s", app)

	consulApi, err := utils.ConsulClient(data.GetString("consul-address"))
	if err != nil {
		return err
	}

	if err := utils.RegisterPluginData("deploy.marathon", data.GetString("app-name"), data.String(), data.GetString("consul-address")); err != nil {
		return err
	}

	return backoff.Retry(func() error {
		services, _, err := consulApi.Health().Service(fullName, "", true, nil)

		if err != nil {
			log.Println(color.RedString("Error in check health in consul: %v", err))
			return err
		}

		if len(services) == 0 {
			log.Printf("Service `%s` not started yet! Retry...", fullName)
			return fmt.Errorf("Service `%s` not started!", fullName)
		}

		log.Println(color.GreenString("Service `%s` successfully started!", fullName))
		return nil
	}, backoff.NewExponentialBackOff())
}

func (p DeployMarathon) Uninstall(data manifest.Manifest) error {
	marathonApi, err := MarathonClient(data.GetString("marathon-address"))
	if err != nil {
		return err
	}

	name := data.GetString("app-name")

	if _, err := marathonApi.Application(name); err == nil {
		if _, err := marathonApi.DeleteApplication(name, false); err != nil {
			return err
		}
	} else {
		log.Println(color.YellowString("App `%s` doesnt exists in marathon!", name))
	}

	return utils.DeletePluginData("deploy.marathon", name, data.GetString("consul-address"))
}

func MarathonClient(marathonAddress string) (marathon.Marathon, error) {
	conf := marathon.NewDefaultConfig()
	conf.URL = fmt.Sprintf("http://%s", marathonAddress)
	conf.LogOutput = os.Stderr
	return marathon.NewClient(conf)
}
<|endoftext|>


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

package reverseproxy

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/braintree/manners"
	"github.com/nu7hatch/gouuid"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tsuru/planb/log"
)

var (
	emptyResponseBody   = &fixedReadCloser{}
	noRouteResponseBody = &fixedReadCloser{value: noRouteResponseContent}
	noopDirector        = func(*http.Request) {}

	_ ReverseProxy = &NativeReverseProxy{}

	openConnections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "connections_open",
		Help:      "The current number of open connections.",
	})

	requestDurations = prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "request_duration_seconds",
		Help:      "The total request latencies in seconds.",
	})

	backendDurations = prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "backend_request_duration_seconds",
		Help:      "The backends HTTP request latencies in seconds.",
	})

	backendResponse = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "backends_responses",
		Help:      "The total backends responses by status code.",
	}, []string{"status_code"})
)

func init() {
	prometheus.MustRegister(openConnections)
	prometheus.MustRegister(requestDurations)
	prometheus.MustRegister(backendResponse)
	prometheus.MustRegister(backendDurations)
}

type NativeReverseProxy struct {
	http.Transport
	ReverseProxyConfig
	servers []*manners.GracefulServer
	rp      *httputil.ReverseProxy
	dialer  *net.Dialer
}

type fixedReadCloser struct {
	value []byte
}

func (r *fixedReadCloser) Read(p []byte) (n int, err error) {
	return copy(p, r.value), io.EOF
}

func (r *fixedReadCloser) Close() error {
	return nil
}

type bufferPool struct {
	syncPool sync.Pool
}

func (p *bufferPool) Get() []byte {
	b := p.syncPool.Get()
	if b == nil {
		return make([]byte, 32*1024)
	}
	return b.([]byte)
}

func (p *bufferPool) Put(b []byte) {
	p.syncPool.Put(b)
}

func (rp *NativeReverseProxy) Initialize(rpConfig ReverseProxyConfig) error {
	rp.ReverseProxyConfig = rpConfig
	rp.servers = make([]*manners.GracefulServer, 0)

	rp.dialer = &net.Dialer{
		Timeout:   rp.DialTimeout,
		KeepAlive: 30 * time.Second,
	}
	rp.Transport = http.Transport{
		Dial:                rp.dialer.Dial,
		TLSHandshakeTimeout: rp.DialTimeout,
		MaxIdleConnsPerHost: 100,
	}
	rp.rp = &httputil.ReverseProxy{
		Director:      noopDirector,
		Transport:     rp,
		FlushInterval: rp.FlushInterval,
		BufferPool:    &bufferPool{},
	}
	return nil
}

func (rp *NativeReverseProxy) Listen(listener net.Listener) {
	server := manners.NewWithServer(&http.Server{
		Handler: rp,
		ConnState: func(c net.Conn, s http.ConnState) {
			switch s {
			case http.StateNew:
				openConnections.Inc()
			case http.StateClosed:
				openConnections.Dec()
			}
		},
	})
	rp.servers = append(rp.servers, server)
	server.Serve(listener)
}

func (rp *NativeReverseProxy) Stop() {
	for _, server := range rp.servers {
		server.Close()
	}
}

func (rp *NativeReverseProxy) ridString(req *http.Request) string {
	return rp.RequestIDHeader + ":" + req.Header.Get(rp.RequestIDHeader)
}

func (rp *NativeReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if req.Host == "__ping__" && req.URL.Path == "/" {
		rw.WriteHeader(http.StatusOK)
		rw.Write(okResponse)
		return
	}
	if rp.RequestIDHeader != "" && req.Header.Get(rp.RequestIDHeader) == "" {
		unparsedID, err := uuid.NewV4()
		if err == nil {
			req.Header.Set(rp.RequestIDHeader, unparsedID.String())
		}
	}
	upgrade := req.Header.Get("Upgrade")
	if upgrade != "" && strings.ToLower(upgrade) == "websocket" {
		reqData, err := rp.serveWebsocket(rw, req)
		if err != nil {
			reqData.logError(req.URL.Path, rp.ridString(req), err)
			http.Error(rw, "", http.StatusBadGateway)
		}
		return
	}
	rp.rp.ServeHTTP(rw, req)
}

func (rp *NativeReverseProxy) serveWebsocket(rw http.ResponseWriter, req *http.Request) (*RequestData, error) {
	reqData, err := rp.Router.ChooseBackend(req.Host)
	if err != nil {
		return reqData, err
	}
	url, err := url.Parse(reqData.Backend)
	if err != nil {
		return reqData, err
	}
	req.Host = url.Host
	dstConn, err := rp.dialer.Dial("tcp", url.Host)
	if err != nil {
		return reqData, err
	}
	defer dstConn.Close()
	hj, ok := rw.(http.Hijacker)
	if !ok {
		return reqData, errors.New("not a hijacker")
	}
	conn, _, err := hj.Hijack()
	if err != nil {
		return reqData, err
	}
	defer conn.Close()
	var clientIP string
	if clientIP, _, err = net.SplitHostPort(req.RemoteAddr); err == nil {
		if prior, ok := req.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		req.Header.Set("X-Forwarded-For", clientIP)
	}
	err = req.Write(dstConn)
	if err != nil {
		return reqData, err
	}
	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(dstConn, conn)
	go cp(conn, dstConn)
	<-errc
	return reqData, nil
}

func (rp *NativeReverseProxy) RoundTrip(req *http.Request) (*http.Response, error) {
	req.URL.Scheme = ""
	req.URL.Host = ""
	reqData, err := rp.Router.ChooseBackend(req.Host)
	if err != nil {
		reqData.logError(req.URL.Path, rp.ridString(req), err)
		return rp.roundTripWithData(req, reqData, err), nil
	}
	u, err := url.Parse(reqData.Backend)
	if err == nil {
		req.URL.Host = u.Host
		req.URL.Scheme = u.Scheme
	}
	if req.URL.Host == "" {
		req.URL.Scheme = "http"
		req.URL.Host = reqData.Backend
	}
	return rp.roundTripWithData(req, reqData, nil), nil
}

func (rp *NativeReverseProxy) doResponse(req *http.Request, reqData *RequestData, rsp *http.Response, isDebug bool, isDead bool, backendDuration time.Duration) *http.Response {
	totalDuration := time.Since(reqData.StartTime)
	logEntry := func() *log.LogEntry {
		return &log.LogEntry{
			Now:             time.Now(),
			BackendDuration: backendDuration,
			TotalDuration:   totalDuration,
			BackendKey:      reqData.BackendKey,
			RemoteAddr:      req.RemoteAddr,
			Method:          req.Method,
			Path:            req.URL.Path,
			Proto:           req.Proto,
			Referer:         req.Referer(),
			UserAgent:       req.UserAgent(),
			RequestIDHeader: rp.RequestIDHeader,
			RequestID:       req.Header.Get(rp.RequestIDHeader),
			StatusCode:      rsp.StatusCode,
			ContentLength:   rsp.ContentLength,
		}
	}
	rsp.Request = req
	rsp.ProtoMajor = req.ProtoMajor
	rsp.ProtoMinor = req.ProtoMinor
	if rsp.Header == nil {
		rsp.Header = http.Header{}
	}
	if isDebug {
		rsp.Header.Set("X-Debug-Backend-Url", reqData.Backend)
		rsp.Header.Set("X-Debug-Backend-Id", strconv.FormatUint(uint64(reqData.BackendIdx), 10))
		rsp.Header.Set("X-Debug-Frontend-Key", reqData.Host)
	}
	err := rp.Router.EndRequest(reqData, isDead, logEntry)
	if err != nil {
		reqData.logError(req.URL.Path, rp.ridString(req), err)
	}
	backendResponse.WithLabelValues(strconv.Itoa(rsp.StatusCode)).Inc()
	backendDurations.Observe(backendDuration.Seconds())
	requestDurations.Observe(totalDuration.Seconds())
	return rsp
}

func (rp *NativeReverseProxy) roundTripWithData(req *http.Request, reqData *RequestData, err error) (rsp *http.Response) {
	isDebug := req.Header.Get("X-Debug-Router") != ""
	req.Header.Del("X-Debug-Router")
	if err != nil || req.URL.Scheme == "" || req.URL.Host == "" {
		switch err {
		case nil, ErrNoRegisteredBackends:
			rsp = &http.Response{
				StatusCode:    http.StatusBadRequest,
				ContentLength: int64(len(noRouteResponseBody.value)),
				Body:          noRouteResponseBody,
			}
		default:
			rsp = &http.Response{
				StatusCode: http.StatusServiceUnavailable,
				Body:       emptyResponseBody,
			}
		}
		return rp.doResponse(req, reqData, rsp, isDebug, false, 0)
	}
	var timedout int32
	if rp.RequestTimeout > 0 {
		timer := time.AfterFunc(rp.RequestTimeout, func() {
			atomic.AddInt32(&timedout, 1)
			rp.Transport.CancelRequest(req)
		})
		defer timer.Stop()
	}
	host, _, _ := net.SplitHostPort(req.URL.Host)
	if host == "" {
		host = req.URL.Host
	}
	isIP := net.ParseIP(host) != nil
	if !isIP {
		req.Header.Set("X-Host", req.Host)
		req.Header.Set("X-Forwarded-Host", req.Host)
		req.Host = host
	}
	t0 := time.Now().UTC()
	rsp, err = rp.Transport.RoundTrip(req)
	backendDuration := time.Since(t0)
	markAsDead := false
	if err != nil {
		var dialTimeout, requestTimeout bool
		if netErr, ok := err.(net.Error); ok {
			markAsDead = !netErr.Temporary()
			dialTimeout = netErr.Timeout()
		}
		requestTimeout = atomic.LoadInt32(&timedout) == int32(1)
		if requestTimeout {
			markAsDead = false
			err = fmt.Errorf("request timeout after %v: %s", time.Since(reqData.StartTime), err)
		} else if dialTimeout {
			markAsDead = true
			err = fmt.Errorf("dial timeout after %v: %s", time.Since(reqData.StartTime), err)
		} else {
			err = fmt.Errorf("error in backend request: %s", err)
		}
		if markAsDead {
			err = fmt.Errorf("%s *DEAD*", err)
		}
		reqData.logError(req.URL.Path, rp.ridString(req), err)
		rsp = &http.Response{
			StatusCode: http.StatusServiceUnavailable,
			Body:       emptyResponseBody,
		}
	}
	return rp.doResponse(req, reqData, rsp, isDebug, markAsDead, backendDuration)
}
<commit_msg>reverseproxy: better metric naming<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package reverseproxy

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/braintree/manners"
	"github.com/nu7hatch/gouuid"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tsuru/planb/log"
)

var (
	emptyResponseBody   = &fixedReadCloser{}
	noRouteResponseBody = &fixedReadCloser{value: noRouteResponseContent}
	noopDirector        = func(*http.Request) {}

	_ ReverseProxy = &NativeReverseProxy{}

	openConnections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "connections_current_open",
		Help:      "The current number of open connections.",
	})

	requestDurations = prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "request_duration_seconds",
		Help:      "The total request latencies in seconds.",
	})

	backendDurations = prometheus.NewHistogram(prometheus.HistogramOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "backend_request_duration_seconds",
		Help:      "The backends HTTP request latencies in seconds.",
	})

	backendResponse = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "planb",
		Subsystem: "reverseproxy",
		Name:      "backends_responses_total",
		Help:      "The total backends responses by status code.",
	}, []string{"status_code"})
)

func init() {
	prometheus.MustRegister(openConnections)
	prometheus.MustRegister(requestDurations)
	prometheus.MustRegister(backendResponse)
	prometheus.MustRegister(backendDurations)
}

type NativeReverseProxy struct {
	http.Transport
	ReverseProxyConfig
	servers []*manners.GracefulServer
	rp      *httputil.ReverseProxy
	dialer  *net.Dialer
}

type fixedReadCloser struct {
	value []byte
}

func (r *fixedReadCloser) Read(p []byte) (n int, err error) {
	return copy(p, r.value), io.EOF
}

func (r *fixedReadCloser) Close() error {
	return nil
}

type bufferPool struct {
	syncPool sync.Pool
}

func (p *bufferPool) Get() []byte {
	b := p.syncPool.Get()
	if b == nil {
		return make([]byte, 32*1024)
	}
	return b.([]byte)
}

func (p *bufferPool) Put(b []byte) {
	p.syncPool.Put(b)
}

func (rp *NativeReverseProxy) Initialize(rpConfig ReverseProxyConfig) error {
	rp.ReverseProxyConfig = rpConfig
	rp.servers = make([]*manners.GracefulServer, 0)

	rp.dialer = &net.Dialer{
		Timeout:   rp.DialTimeout,
		KeepAlive: 30 * time.Second,
	}
	rp.Transport = http.Transport{
		Dial:                rp.dialer.Dial,
		TLSHandshakeTimeout: rp.DialTimeout,
		MaxIdleConnsPerHost: 100,
	}
	rp.rp = &httputil.ReverseProxy{
		Director:      noopDirector,
		Transport:     rp,
		FlushInterval: rp.FlushInterval,
		BufferPool:    &bufferPool{},
	}
	return nil
}

func (rp *NativeReverseProxy) Listen(listener net.Listener) {
	server := manners.NewWithServer(&http.Server{
		Handler: rp,
		ConnState: func(c net.Conn, s http.ConnState) {
			switch s {
			case http.StateNew:
				openConnections.Inc()
			case http.StateClosed:
				openConnections.Dec()
			}
		},
	})
	rp.servers = append(rp.servers, server)
	server.Serve(listener)
}

func (rp *NativeReverseProxy) Stop() {
	for _, server := range rp.servers {
		server.Close()
	}
}

func (rp *NativeReverseProxy) ridString(req *http.Request) string {
	return rp.RequestIDHeader + ":" + req.Header.Get(rp.RequestIDHeader)
}

func (rp *NativeReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	if req.Host == "__ping__" && req.URL.Path == "/" {
		rw.WriteHeader(http.StatusOK)
		rw.Write(okResponse)
		return
	}
	if rp.RequestIDHeader != "" && req.Header.Get(rp.RequestIDHeader) == "" {
		unparsedID, err := uuid.NewV4()
		if err == nil {
			req.Header.Set(rp.RequestIDHeader, unparsedID.String())
		}
	}
	upgrade := req.Header.Get("Upgrade")
	if upgrade != "" && strings.ToLower(upgrade) == "websocket" {
		reqData, err := rp.serveWebsocket(rw, req)
		if err != nil {
			reqData.logError(req.URL.Path, rp.ridString(req), err)
			http.Error(rw, "", http.StatusBadGateway)
		}
		return
	}
	rp.rp.ServeHTTP(rw, req)
}

func (rp *NativeReverseProxy) serveWebsocket(rw http.ResponseWriter, req *http.Request) (*RequestData, error) {
	reqData, err := rp.Router.ChooseBackend(req.Host)
	if err != nil {
		return reqData, err
	}
	url, err := url.Parse(reqData.Backend)
	if err != nil {
		return reqData, err
	}
	req.Host = url.Host
	dstConn, err := rp.dialer.Dial("tcp", url.Host)
	if err != nil {
		return reqData, err
	}
	defer dstConn.Close()
	hj, ok := rw.(http.Hijacker)
	if !ok {
		return reqData, errors.New("not a hijacker")
	}
	conn, _, err := hj.Hijack()
	if err != nil {
		return reqData, err
	}
	defer conn.Close()
	var clientIP string
	if clientIP, _, err = net.SplitHostPort(req.RemoteAddr); err == nil {
		if prior, ok := req.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		req.Header.Set("X-Forwarded-For", clientIP)
	}
	err = req.Write(dstConn)
	if err != nil {
		return reqData, err
	}
	errc := make(chan error, 2)
	cp := func(dst io.Writer, src io.Reader) {
		_, err := io.Copy(dst, src)
		errc <- err
	}
	go cp(dstConn, conn)
	go cp(conn, dstConn)
	<-errc
	return reqData, nil
}

func (rp *NativeReverseProxy) RoundTrip(req *http.Request) (*http.Response, error) {
	req.URL.Scheme = ""
	req.URL.Host = ""
	reqData, err := rp.Router.ChooseBackend(req.Host)
	if err != nil {
		reqData.logError(req.URL.Path, rp.ridString(req), err)
		return rp.roundTripWithData(req, reqData, err), nil
	}
	u, err := url.Parse(reqData.Backend)
	if err == nil {
		req.URL.Host = u.Host
		req.URL.Scheme = u.Scheme
	}
	if req.URL.Host == "" {
		req.URL.Scheme = "http"
		req.URL.Host = reqData.Backend
	}
	return rp.roundTripWithData(req, reqData, nil), nil
}

func (rp *NativeReverseProxy) doResponse(req *http.Request, reqData *RequestData, rsp *http.Response, isDebug bool, isDead bool, backendDuration time.Duration) *http.Response {
	totalDuration := time.Since(reqData.StartTime)
	logEntry := func() *log.LogEntry {
		return &log.LogEntry{
			Now:             time.Now(),
			BackendDuration: backendDuration,
			TotalDuration:   totalDuration,
			BackendKey:      reqData.BackendKey,
			RemoteAddr:      req.RemoteAddr,
			Method:          req.Method,
			Path:            req.URL.Path,
			Proto:           req.Proto,
			Referer:         req.Referer(),
			UserAgent:       req.UserAgent(),
			RequestIDHeader: rp.RequestIDHeader,
			RequestID:       req.Header.Get(rp.RequestIDHeader),
			StatusCode:      rsp.StatusCode,
			ContentLength:   rsp.ContentLength,
		}
	}
	rsp.Request = req
	rsp.ProtoMajor = req.ProtoMajor
	rsp.ProtoMinor = req.ProtoMinor
	if rsp.Header == nil {
		rsp.Header = http.Header{}
	}
	if isDebug {
		rsp.Header.Set("X-Debug-Backend-Url", reqData.Backend)
		rsp.Header.Set("X-Debug-Backend-Id", strconv.FormatUint(uint64(reqData.BackendIdx), 10))
		rsp.Header.Set("X-Debug-Frontend-Key", reqData.Host)
	}
	err := rp.Router.EndRequest(reqData, isDead, logEntry)
	if err != nil {
		reqData.logError(req.URL.Path, rp.ridString(req), err)
	}
	backendResponse.WithLabelValues(strconv.Itoa(rsp.StatusCode)).Inc()
	backendDurations.Observe(backendDuration.Seconds())
	requestDurations.Observe(totalDuration.Seconds())
	return rsp
}

func (rp *NativeReverseProxy) roundTripWithData(req *http.Request, reqData *RequestData, err error) (rsp *http.Response) {
	isDebug := req.Header.Get("X-Debug-Router") != ""
	req.Header.Del("X-Debug-Router")
	if err != nil || req.URL.Scheme == "" || req.URL.Host == "" {
		switch err {
		case nil, ErrNoRegisteredBackends:
			rsp = &http.Response{
				StatusCode:    http.StatusBadRequest,
				ContentLength: int64(len(noRouteResponseBody.value)),
				Body:          noRouteResponseBody,
			}
		default:
			rsp = &http.Response{
				StatusCode: http.StatusServiceUnavailable,
				Body:       emptyResponseBody,
			}
		}
		return rp.doResponse(req, reqData, rsp, isDebug, false, 0)
	}
	var timedout int32
	if rp.RequestTimeout > 0 {
		timer := time.AfterFunc(rp.RequestTimeout, func() {
			atomic.AddInt32(&timedout, 1)
			rp.Transport.CancelRequest(req)
		})
		defer timer.Stop()
	}
	host, _, _ := net.SplitHostPort(req.URL.Host)
	if host == "" {
		host = req.URL.Host
	}
	isIP := net.ParseIP(host) != nil
	if !isIP {
		req.Header.Set("X-Host", req.Host)
		req.Header.Set("X-Forwarded-Host", req.Host)
		req.Host = host
	}
	t0 := time.Now().UTC()
	rsp, err = rp.Transport.RoundTrip(req)
	backendDuration := time.Since(t0)
	markAsDead := false
	if err != nil {
		var dialTimeout, requestTimeout bool
		if netErr, ok := err.(net.Error); ok {
			markAsDead = !netErr.Temporary()
			dialTimeout = netErr.Timeout()
		}
		requestTimeout = atomic.LoadInt32(&timedout) == int32(1)
		if requestTimeout {
			markAsDead = false
			err = fmt.Errorf("request timeout after %v: %s", time.Since(reqData.StartTime), err)
		} else if dialTimeout {
			markAsDead = true
			err = fmt.Errorf("dial timeout after %v: %s", time.Since(reqData.StartTime), err)
		} else {
			err = fmt.Errorf("error in backend request: %s", err)
		}
		if markAsDead {
			err = fmt.Errorf("%s *DEAD*", err)
		}
		reqData.logError(req.URL.Path, rp.ridString(req), err)
		rsp = &http.Response{
			StatusCode: http.StatusServiceUnavailable,
			Body:       emptyResponseBody,
		}
	}
	return rp.doResponse(req, reqData, rsp, isDebug, markAsDead, backendDuration)
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd Authors.

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

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

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

package containerd

import (
	"bytes"
	"context"
	"encoding/json"
	"io"

	"io/ioutil"
	"math/rand"
	"reflect"
	"runtime"
	"testing"

	"github.com/containerd/containerd/archive/compression"
	"github.com/containerd/containerd/archive/tartest"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/images/archive"
	digest "github.com/opencontainers/go-digest"
	specs "github.com/opencontainers/image-spec/specs-go"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
)

// TestExportAndImport exports testImage as a tar stream,
// and import the tar stream as a new image.
func TestExportAndImport(t *testing.T) {
	// TODO: support windows
	if testing.Short() || runtime.GOOS == "windows" {
		t.Skip()
	}
	ctx, cancel := testContext(t)
	defer cancel()

	client, err := New(address)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	_, err = client.Fetch(ctx, testImage)
	if err != nil {
		t.Fatal(err)
	}

	wb := bytes.NewBuffer(nil)
	err = client.Export(ctx, wb, archive.WithAllPlatforms(), archive.WithImage(client.ImageService(), testImage))
	if err != nil {
		t.Fatal(err)
	}

	opts := []ImportOpt{
		WithImageRefTranslator(archive.AddRefPrefix("foo/bar")),
	}
	imgrecs, err := client.Import(ctx, bytes.NewReader(wb.Bytes()), opts...)
	if err != nil {
		t.Fatalf("Import failed: %+v", err)
	}

	for _, imgrec := range imgrecs {
		if imgrec.Name == testImage {
			continue
		}
		err = client.ImageService().Delete(ctx, imgrec.Name)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestImport(t *testing.T) {
	ctx, cancel := testContext(t)
	defer cancel()

	client, err := New(address)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	tc := tartest.TarContext{}

	b1, d1 := createContent(256, 1)
	empty := []byte("{}")
	version := []byte("1.0")

	c1, d2 := createConfig()

	m1, d3, expManifest := createManifest(c1, [][]byte{b1})

	provider := client.ContentStore()

	checkManifest := func(ctx context.Context, t *testing.T, d ocispec.Descriptor, expManifest *ocispec.Manifest) {
		m, err := images.Manifest(ctx, provider, d, nil)
		if err != nil {
			t.Fatalf("unable to read target blob: %+v", err)
		}

		if m.Config.Digest != d2 {
			t.Fatalf("unexpected digest hash %s, expected %s", m.Config.Digest, d2)
		}

		if len(m.Layers) != 1 {
			t.Fatalf("expected 1 layer, has %d", len(m.Layers))
		}

		if m.Layers[0].Digest != d1 {
			t.Fatalf("unexpected layer hash %s, expected %s", m.Layers[0].Digest, d1)
		}

		if expManifest != nil {
			if !reflect.DeepEqual(m.Layers, expManifest.Layers) {
				t.Fatalf("DeepEqual on Layers failed: %v vs. %v", m.Layers, expManifest.Layers)
			}
			if !reflect.DeepEqual(m.Config, expManifest.Config) {
				t.Fatalf("DeepEqual on Config failed: %v vs. %v", m.Config, expManifest.Config)
			}
		}
	}

	for _, tc := range []struct {
		Name   string
		Writer tartest.WriterToTar
		Check  func(*testing.T, []images.Image)
		Opts   []ImportOpt
	}{
		{
			Name: "DockerV2.0",
			Writer: tartest.TarAll(
				tc.Dir("bd765cd43e95212f7aa2cab51d0a", 0755),
				tc.File("bd765cd43e95212f7aa2cab51d0a/json", empty, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/layer.tar", b1, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/VERSION", version, 0644),
				tc.File("repositories", []byte(`{"any":{"1":"bd765cd43e95212f7aa2cab51d0a"}}`), 0644),
			),
		},
		{
			Name: "DockerV2.1",
			Writer: tartest.TarAll(
				tc.Dir("bd765cd43e95212f7aa2cab51d0a", 0755),
				tc.File("bd765cd43e95212f7aa2cab51d0a/json", empty, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/layer.tar", b1, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/VERSION", version, 0644),
				tc.File("e95212f7aa2cab51d0abd765cd43.json", c1, 0644),
				tc.File("manifest.json", []byte(`[{"Config":"e95212f7aa2cab51d0abd765cd43.json","RepoTags":["test-import:notlatest", "another/repo:tag"],"Layers":["bd765cd43e95212f7aa2cab51d0a/layer.tar"]}]`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				if len(imgs) == 0 {
					t.Fatalf("no images")
				}

				names := []string{
					"docker.io/library/test-import:notlatest",
					"docker.io/another/repo:tag",
				}

				checkImages(t, imgs[0].Target.Digest, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, nil)
			},
		},
		{
			Name: "OCI-BadFormat",
			Writer: tartest.TarAll(
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"2.0.0"}`), 0644),
			),
		},
		{
			Name: "OCI",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"latest",
					"docker.io/lib/img:ok",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
		},
		{
			Name: "OCIPrefixName",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"localhost:5000/myimage:latest",
					"docker.io/lib/img:ok",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
			Opts: []ImportOpt{
				WithImageRefTranslator(archive.AddRefPrefix("localhost:5000/myimage")),
			},
		},
		{
			Name: "OCIPrefixName2",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "localhost:5000/myimage:old", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"localhost:5000/myimage:latest",
					"localhost:5000/myimage:old",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
			Opts: []ImportOpt{
				WithImageRefTranslator(archive.FilterRefPrefix("localhost:5000/myimage")),
			},
		},
	} {
		t.Run(tc.Name, func(t *testing.T) {
			images, err := client.Import(ctx, tartest.TarFromWriterTo(tc.Writer), tc.Opts...)
			if err != nil {
				if tc.Check != nil {
					t.Errorf("unexpected import error: %+v", err)
				}
				return
			} else if tc.Check == nil {
				t.Fatalf("expected error on import")
			}

			tc.Check(t, images)
		})
	}
}

func checkImages(t *testing.T, target digest.Digest, actual []images.Image, names ...string) {
	if len(names) != len(actual) {
		t.Fatalf("expected %d images, got %d", len(names), len(actual))
	}

	for i, n := range names {
		if actual[i].Target.Digest != target {
			t.Fatalf("image(%d) unexpected target %s, expected %s", i, actual[i].Target.Digest, target)
		}
		if actual[i].Name != n {
			t.Fatalf("image(%d) unexpected name %q, expected %q", i, actual[i].Name, n)
		}

		if actual[i].Target.MediaType != ocispec.MediaTypeImageManifest &&
			actual[i].Target.MediaType != images.MediaTypeDockerSchema2Manifest {
			t.Fatalf("image(%d) unexpected media type: %s", i, actual[i].Target.MediaType)
		}
	}

}

func createContent(size int64, seed int64) ([]byte, digest.Digest) {
	b, err := ioutil.ReadAll(io.LimitReader(rand.New(rand.NewSource(seed)), size))
	if err != nil {
		panic(err)
	}
	wb := bytes.NewBuffer(nil)
	cw, err := compression.CompressStream(wb, compression.Gzip)
	if err != nil {
		panic(err)
	}

	if _, err := cw.Write(b); err != nil {
		panic(err)
	}
	b = wb.Bytes()
	return b, digest.FromBytes(b)
}

func createConfig() ([]byte, digest.Digest) {
	image := ocispec.Image{
		OS:           "any",
		Architecture: "any",
		Author:       "test",
	}
	b, _ := json.Marshal(image)

	return b, digest.FromBytes(b)
}

func createManifest(config []byte, layers [][]byte) ([]byte, digest.Digest, *ocispec.Manifest) {
	manifest := ocispec.Manifest{
		Versioned: specs.Versioned{
			SchemaVersion: 2,
		},
		Config: ocispec.Descriptor{
			MediaType: ocispec.MediaTypeImageConfig,
			Digest:    digest.FromBytes(config),
			Size:      int64(len(config)),
			Annotations: map[string]string{
				"ocispec": "manifest.config.descriptor",
			},
		},
	}
	for _, l := range layers {
		manifest.Layers = append(manifest.Layers, ocispec.Descriptor{
			MediaType: ocispec.MediaTypeImageLayer,
			Digest:    digest.FromBytes(l),
			Size:      int64(len(l)),
			Annotations: map[string]string{
				"ocispec": "manifest.layers.descriptor",
			},
		})
	}

	b, _ := json.Marshal(manifest)

	return b, digest.FromBytes(b), &manifest
}

func createIndex(manifest []byte, tags ...string) []byte {
	idx := ocispec.Index{
		Versioned: specs.Versioned{
			SchemaVersion: 2,
		},
	}
	d := ocispec.Descriptor{
		MediaType: ocispec.MediaTypeImageManifest,
		Digest:    digest.FromBytes(manifest),
		Size:      int64(len(manifest)),
	}

	if len(tags) == 0 {
		idx.Manifests = append(idx.Manifests, d)
	} else {
		for _, t := range tags {
			dt := d
			dt.Annotations = map[string]string{
				ocispec.AnnotationRefName: t,
			}
			idx.Manifests = append(idx.Manifests, dt)
		}
	}

	b, _ := json.Marshal(idx)

	return b
}
<commit_msg>Remove an empty line<commit_after>/*
   Copyright The containerd Authors.

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

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

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

package containerd

import (
	"bytes"
	"context"
	"encoding/json"
	"io"

	"io/ioutil"
	"math/rand"
	"reflect"
	"runtime"
	"testing"

	"github.com/containerd/containerd/archive/compression"
	"github.com/containerd/containerd/archive/tartest"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/images/archive"
	digest "github.com/opencontainers/go-digest"
	specs "github.com/opencontainers/image-spec/specs-go"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
)

// TestExportAndImport exports testImage as a tar stream,
// and import the tar stream as a new image.
func TestExportAndImport(t *testing.T) {
	// TODO: support windows
	if testing.Short() || runtime.GOOS == "windows" {
		t.Skip()
	}
	ctx, cancel := testContext(t)
	defer cancel()

	client, err := New(address)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	_, err = client.Fetch(ctx, testImage)
	if err != nil {
		t.Fatal(err)
	}

	wb := bytes.NewBuffer(nil)
	err = client.Export(ctx, wb, archive.WithAllPlatforms(), archive.WithImage(client.ImageService(), testImage))
	if err != nil {
		t.Fatal(err)
	}

	opts := []ImportOpt{
		WithImageRefTranslator(archive.AddRefPrefix("foo/bar")),
	}
	imgrecs, err := client.Import(ctx, bytes.NewReader(wb.Bytes()), opts...)
	if err != nil {
		t.Fatalf("Import failed: %+v", err)
	}

	for _, imgrec := range imgrecs {
		if imgrec.Name == testImage {
			continue
		}
		err = client.ImageService().Delete(ctx, imgrec.Name)
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestImport(t *testing.T) {
	ctx, cancel := testContext(t)
	defer cancel()

	client, err := New(address)
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	tc := tartest.TarContext{}

	b1, d1 := createContent(256, 1)
	empty := []byte("{}")
	version := []byte("1.0")

	c1, d2 := createConfig()

	m1, d3, expManifest := createManifest(c1, [][]byte{b1})

	provider := client.ContentStore()

	checkManifest := func(ctx context.Context, t *testing.T, d ocispec.Descriptor, expManifest *ocispec.Manifest) {
		m, err := images.Manifest(ctx, provider, d, nil)
		if err != nil {
			t.Fatalf("unable to read target blob: %+v", err)
		}

		if m.Config.Digest != d2 {
			t.Fatalf("unexpected digest hash %s, expected %s", m.Config.Digest, d2)
		}

		if len(m.Layers) != 1 {
			t.Fatalf("expected 1 layer, has %d", len(m.Layers))
		}

		if m.Layers[0].Digest != d1 {
			t.Fatalf("unexpected layer hash %s, expected %s", m.Layers[0].Digest, d1)
		}

		if expManifest != nil {
			if !reflect.DeepEqual(m.Layers, expManifest.Layers) {
				t.Fatalf("DeepEqual on Layers failed: %v vs. %v", m.Layers, expManifest.Layers)
			}
			if !reflect.DeepEqual(m.Config, expManifest.Config) {
				t.Fatalf("DeepEqual on Config failed: %v vs. %v", m.Config, expManifest.Config)
			}
		}
	}

	for _, tc := range []struct {
		Name   string
		Writer tartest.WriterToTar
		Check  func(*testing.T, []images.Image)
		Opts   []ImportOpt
	}{
		{
			Name: "DockerV2.0",
			Writer: tartest.TarAll(
				tc.Dir("bd765cd43e95212f7aa2cab51d0a", 0755),
				tc.File("bd765cd43e95212f7aa2cab51d0a/json", empty, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/layer.tar", b1, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/VERSION", version, 0644),
				tc.File("repositories", []byte(`{"any":{"1":"bd765cd43e95212f7aa2cab51d0a"}}`), 0644),
			),
		},
		{
			Name: "DockerV2.1",
			Writer: tartest.TarAll(
				tc.Dir("bd765cd43e95212f7aa2cab51d0a", 0755),
				tc.File("bd765cd43e95212f7aa2cab51d0a/json", empty, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/layer.tar", b1, 0644),
				tc.File("bd765cd43e95212f7aa2cab51d0a/VERSION", version, 0644),
				tc.File("e95212f7aa2cab51d0abd765cd43.json", c1, 0644),
				tc.File("manifest.json", []byte(`[{"Config":"e95212f7aa2cab51d0abd765cd43.json","RepoTags":["test-import:notlatest", "another/repo:tag"],"Layers":["bd765cd43e95212f7aa2cab51d0a/layer.tar"]}]`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				if len(imgs) == 0 {
					t.Fatalf("no images")
				}

				names := []string{
					"docker.io/library/test-import:notlatest",
					"docker.io/another/repo:tag",
				}

				checkImages(t, imgs[0].Target.Digest, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, nil)
			},
		},
		{
			Name: "OCI-BadFormat",
			Writer: tartest.TarAll(
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"2.0.0"}`), 0644),
			),
		},
		{
			Name: "OCI",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"latest",
					"docker.io/lib/img:ok",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
		},
		{
			Name: "OCIPrefixName",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"localhost:5000/myimage:latest",
					"docker.io/lib/img:ok",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
			Opts: []ImportOpt{
				WithImageRefTranslator(archive.AddRefPrefix("localhost:5000/myimage")),
			},
		},
		{
			Name: "OCIPrefixName2",
			Writer: tartest.TarAll(
				tc.Dir("blobs", 0755),
				tc.Dir("blobs/sha256", 0755),
				tc.File("blobs/sha256/"+d1.Encoded(), b1, 0644),
				tc.File("blobs/sha256/"+d2.Encoded(), c1, 0644),
				tc.File("blobs/sha256/"+d3.Encoded(), m1, 0644),
				tc.File("index.json", createIndex(m1, "latest", "localhost:5000/myimage:old", "docker.io/lib/img:ok"), 0644),
				tc.File("oci-layout", []byte(`{"imageLayoutVersion":"1.0.0"}`), 0644),
			),
			Check: func(t *testing.T, imgs []images.Image) {
				names := []string{
					"localhost:5000/myimage:latest",
					"localhost:5000/myimage:old",
				}

				checkImages(t, d3, imgs, names...)
				checkManifest(ctx, t, imgs[0].Target, expManifest)
			},
			Opts: []ImportOpt{
				WithImageRefTranslator(archive.FilterRefPrefix("localhost:5000/myimage")),
			},
		},
	} {
		t.Run(tc.Name, func(t *testing.T) {
			images, err := client.Import(ctx, tartest.TarFromWriterTo(tc.Writer), tc.Opts...)
			if err != nil {
				if tc.Check != nil {
					t.Errorf("unexpected import error: %+v", err)
				}
				return
			} else if tc.Check == nil {
				t.Fatalf("expected error on import")
			}

			tc.Check(t, images)
		})
	}
}

func checkImages(t *testing.T, target digest.Digest, actual []images.Image, names ...string) {
	if len(names) != len(actual) {
		t.Fatalf("expected %d images, got %d", len(names), len(actual))
	}

	for i, n := range names {
		if actual[i].Target.Digest != target {
			t.Fatalf("image(%d) unexpected target %s, expected %s", i, actual[i].Target.Digest, target)
		}
		if actual[i].Name != n {
			t.Fatalf("image(%d) unexpected name %q, expected %q", i, actual[i].Name, n)
		}

		if actual[i].Target.MediaType != ocispec.MediaTypeImageManifest &&
			actual[i].Target.MediaType != images.MediaTypeDockerSchema2Manifest {
			t.Fatalf("image(%d) unexpected media type: %s", i, actual[i].Target.MediaType)
		}
	}
}

func createContent(size int64, seed int64) ([]byte, digest.Digest) {
	b, err := ioutil.ReadAll(io.LimitReader(rand.New(rand.NewSource(seed)), size))
	if err != nil {
		panic(err)
	}
	wb := bytes.NewBuffer(nil)
	cw, err := compression.CompressStream(wb, compression.Gzip)
	if err != nil {
		panic(err)
	}

	if _, err := cw.Write(b); err != nil {
		panic(err)
	}
	b = wb.Bytes()
	return b, digest.FromBytes(b)
}

func createConfig() ([]byte, digest.Digest) {
	image := ocispec.Image{
		OS:           "any",
		Architecture: "any",
		Author:       "test",
	}
	b, _ := json.Marshal(image)

	return b, digest.FromBytes(b)
}

func createManifest(config []byte, layers [][]byte) ([]byte, digest.Digest, *ocispec.Manifest) {
	manifest := ocispec.Manifest{
		Versioned: specs.Versioned{
			SchemaVersion: 2,
		},
		Config: ocispec.Descriptor{
			MediaType: ocispec.MediaTypeImageConfig,
			Digest:    digest.FromBytes(config),
			Size:      int64(len(config)),
			Annotations: map[string]string{
				"ocispec": "manifest.config.descriptor",
			},
		},
	}
	for _, l := range layers {
		manifest.Layers = append(manifest.Layers, ocispec.Descriptor{
			MediaType: ocispec.MediaTypeImageLayer,
			Digest:    digest.FromBytes(l),
			Size:      int64(len(l)),
			Annotations: map[string]string{
				"ocispec": "manifest.layers.descriptor",
			},
		})
	}

	b, _ := json.Marshal(manifest)

	return b, digest.FromBytes(b), &manifest
}

func createIndex(manifest []byte, tags ...string) []byte {
	idx := ocispec.Index{
		Versioned: specs.Versioned{
			SchemaVersion: 2,
		},
	}
	d := ocispec.Descriptor{
		MediaType: ocispec.MediaTypeImageManifest,
		Digest:    digest.FromBytes(manifest),
		Size:      int64(len(manifest)),
	}

	if len(tags) == 0 {
		idx.Manifests = append(idx.Manifests, d)
	} else {
		for _, t := range tags {
			dt := d
			dt.Annotations = map[string]string{
				ocispec.AnnotationRefName: t,
			}
			idx.Manifests = append(idx.Manifests, dt)
		}
	}

	b, _ := json.Marshal(idx)

	return b
}
<|endoftext|>