<commit_before>/*
Copyright 2017 The Kubernetes 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 ranch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/boskos/common"
)

// Ranch is the place which all of the Resource objects lives.
type Ranch struct {
	Storage       *Storage
	resourcesLock sync.RWMutex
	requestMgr    *RequestManager
	//
	now func() time.Time
}

func updateTime() time.Time {
	return time.Now()
}

// Public errors:

// ResourceNotFound will be returned if requested resource does not exist.
type ResourceNotFound struct {
	name string
}

func (r ResourceNotFound) Error() string {
	return fmt.Sprintf("no available resource %s, try again later.", r.name)
}

// ResourceTypeNotFound will be returned if requested resource type does not exist.
type ResourceTypeNotFound struct {
	rType string
}

func (r ResourceTypeNotFound) Error() string {
	return fmt.Sprintf("resource type %s does not exist", r.rType)
}

// OwnerNotMatch will be returned if request owner does not match current owner for target resource.
type OwnerNotMatch struct {
	request string
	owner   string
}

func (o OwnerNotMatch) Error() string {
	return fmt.Sprintf("owner mismatch request by %s, currently owned by %s", o.request, o.owner)
}

// StateNotMatch will be returned if requested state does not match current state for target resource.
type StateNotMatch struct {
	expect  string
	current string
}

func (s StateNotMatch) Error() string {
	return fmt.Sprintf("state mismatch - expected %v, current %v", s.expect, s.current)
}

// NewRanch creates a new Ranch object.
// In: config - path to resource file
//     storage - path to where to save/restore the state data
// Out: A Ranch object, loaded from config/storage, or error
func NewRanch(config string, s *Storage, ttl time.Duration) (*Ranch, error) {
	newRanch := &Ranch{
		Storage:    s,
		requestMgr: NewRequestManager(ttl),
		now:        time.Now,
	}
	if config != "" {
		if err := newRanch.SyncConfig(config); err != nil {
			return nil, err
		}
		logrus.Infof("Loaded Boskos configuration successfully")
	}
	return newRanch, nil
}

// acquireRequestPriorityKey is used as key for request priority cache.
type acquireRequestPriorityKey struct {
	rType, state string
}

// Acquire checks out a type of resource in certain state without an owner,
// and move the checked out resource to the end of the resource list.
// In: rtype - name of the target resource
//     state - current state of the requested resource
//     dest - destination state of the requested resource
//     owner - requester of the resource
//     requestID - request ID to get a priority in the queue
// Out: A valid Resource object on success, or
//      ResourceNotFound error if target type resource does not exist in target state.
func (r *Ranch) Acquire(rType, state, dest, owner, requestID string) (*common.Resource, error) {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	logger := logrus.WithFields(logrus.Fields{
		"type":       rType,
		"state":      state,
		"dest":       dest,
		"owner":      owner,
		"identifier": requestID,
	})
	logger.Debug("Determining request priority...")
	ts := acquireRequestPriorityKey{rType: rType, state: state}
	rank, new := r.requestMgr.GetRank(ts, requestID)
	logger.WithFields(logrus.Fields{"rank": rank, "new": new}).Debug("Determined request priority.")

	resources, err := r.Storage.GetResources()
	if err != nil {
		logger.WithError(err).Errorf("could not get resources")
		return nil, &ResourceNotFound{rType}
	}
	logger.Debugf("Considering %d resources.", len(resources))

	// For request priority we need to go over all the list until a matching rank
	matchingResoucesCount := 0
	typeCount := 0
	for idx := range resources {
		res := resources[idx]
		if rType == res.Type {
			typeCount++
			if state == res.State && res.Owner == "" {
				matchingResoucesCount++
				if matchingResoucesCount >= rank {
					logger = logger.WithField("resource", res.Name)
					res.Owner = owner
					res.State = dest
					logger.Debug("Updating resource.")
					updatedRes, err := r.Storage.UpdateResource(res)
					if err != nil {
						logger.WithError(err).Errorf("could not update resource %s", res.Name)
						return nil, err
					}
					// Deleting this request since it has been fulfilled
					if requestID != "" {
						logger.Debug("Cleaning up requests.")
						r.requestMgr.Delete(ts, requestID)
					}
					logger.Debug("Successfully acquired resource.")
					return &updatedRes, nil
				}
			}
		}
	}

	if new {
		logger.Debug("Checking for associated dynamic resource type...")
		lifeCycle, err := r.Storage.GetDynamicResourceLifeCycle(rType)
		// Assuming error means no associated dynamic resource
		if err == nil {
			if typeCount < lifeCycle.MaxCount {
				logger.Debug("Adding new dynamic resources...")
				res := common.NewResourceFromNewDynamicResourceLifeCycle(r.Storage.generateName(), &lifeCycle, r.now())
				if err := r.Storage.AddResource(res); err != nil {
					logger.WithError(err).Warningf("unable to add a new resource of type %s", rType)
				}
				logger.Infof("Added dynamic resource %s of type %s", res.Name, res.Type)
			}
		}
	}

	if typeCount > 0 {
		return nil, &ResourceNotFound{rType}
	}
	return nil, &ResourceTypeNotFound{rType}
}

// AcquireByState checks out resources of a given type without an owner,
// that matches a list of resources names.
// In: state - current state of the requested resource
//     dest - destination state of the requested resource
//     owner - requester of the resource
//     names - names of resource to acquire
// Out: A valid list of Resource object on success, or
//      ResourceNotFound error if target type resource does not exist in target state.
func (r *Ranch) AcquireByState(state, dest, owner string, names []string) ([]common.Resource, error) {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	if names == nil {
		return nil, fmt.Errorf("must provide names of expected resources")
	}

	rNames := map[string]bool{}
	for _, t := range names {
		rNames[t] = true
	}

	allResources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Errorf("could not get resources")
		return nil, &ResourceNotFound{state}
	}

	var resources []common.Resource

	for idx := range allResources {
		res := allResources[idx]
		if state == res.State {
			if res.Owner != "" {
				continue
			}
			if rNames[res.Name] {
				res.Owner = owner
				res.State = dest
				updatedRes, err := r.Storage.UpdateResource(res)
				if err != nil {
					logrus.WithError(err).Errorf("could not update resource %s", res.Name)
					return nil, err
				}
				resources = append(resources, updatedRes)
				delete(rNames, res.Name)
			}
		}
	}
	if len(rNames) != 0 {
		var missingResources []string
		for n := range rNames {
			missingResources = append(missingResources, n)
		}
		err := &ResourceNotFound{state}
		logrus.WithError(err).Errorf("could not find required resources %s", strings.Join(missingResources, ", "))
		return resources, err
	}
	return resources, nil
}

// Release unsets owner for target resource and move it to a new state.
// In: name - name of the target resource
//     dest - destination state of the resource
//     owner - owner of the resource
// Out: nil on success, or
//      OwnerNotMatch error if owner does not match current owner of the resource, or
//      ResourceNotFound error if target named resource does not exist.
func (r *Ranch) Release(name, dest, owner string) error {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	res, err := r.Storage.GetResource(name)
	if err != nil {
		logrus.WithError(err).Errorf("unable to release resource %s", name)
		return &ResourceNotFound{name}
	}
	if owner != res.Owner {
		return &OwnerNotMatch{request: owner, owner: res.Owner}
	}

	res.Owner = ""
	res.State = dest

	if lf, err := r.Storage.GetDynamicResourceLifeCycle(res.Type); err == nil {
		// Assuming error means not existing as the only way to differentiate would be to list
		// all resources and find the right one which is more costly.
		if lf.LifeSpan != nil {
			expirationTime := r.now().Add(*lf.LifeSpan)
			res.ExpirationDate = &expirationTime
		}
	} else {
		res.ExpirationDate = nil
	}

	if _, err := r.Storage.UpdateResource(res); err != nil {
		logrus.WithError(err).Errorf("could not update resource %s", res.Name)
		return err
	}
	return nil
}

// Update updates the timestamp of a target resource.
// In: name  - name of the target resource
//     state - current state of the resource
//     owner - current owner of the resource
// 	   info  - information on how to use the resource
// Out: nil on success, or
//      OwnerNotMatch error if owner does not match current owner of the resource, or
//      ResourceNotFound error if target named resource does not exist, or
//      StateNotMatch error if state does not match current state of the resource.
func (r *Ranch) Update(name, owner, state string, ud *common.UserData) error {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	res, err := r.Storage.GetResource(name)
	if err != nil {
		logrus.WithError(err).Errorf("could not find resource %s for update", name)
		return &ResourceNotFound{name}
	}
	if owner != res.Owner {
		return &OwnerNotMatch{request: owner, owner: res.Owner}
	}
	if state != res.State {
		return &StateNotMatch{res.State, state}
	}
	if res.UserData == nil {
		res.UserData = &common.UserData{}
	}
	res.UserData.Update(ud)
	if _, err := r.Storage.UpdateResource(res); err != nil {
		logrus.WithError(err).Errorf("could not update resource %s", res.Name)
		return err
	}
	return nil
}

// Reset unstucks a type of stale resource to a new state.
// In: rtype - type of the resource
//     state - current state of the resource
//     expire - duration before resource's last update
//     dest - destination state of expired resources
// Out: map of resource name - resource owner.
func (r *Ranch) Reset(rtype, state string, expire time.Duration, dest string) (map[string]string, error) {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	ret := make(map[string]string)

	resources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Errorf("cannot find resources")
		return nil, err
	}

	for idx := range resources {
		res := resources[idx]
		if rtype == res.Type && state == res.State && res.Owner != "" {
			if r.now().Sub(res.LastUpdate) > expire {
				ret[res.Name] = res.Owner
				res.Owner = ""
				res.State = dest
				if _, err := r.Storage.UpdateResource(res); err != nil {
					logrus.WithError(err).Errorf("could not update resource %s", res.Name)
					return ret, err
				}
			}
		}
	}
	return ret, nil
}

// SyncConfig updates resource list from a file
func (r *Ranch) SyncConfig(configPath string) error {
	config, err := common.ParseConfig(configPath)
	if err != nil {
		return err
	}
	if err := common.ValidateConfig(config); err != nil {
		return err
	}
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()
	if err := r.Storage.SyncResources(config); err != nil {
		return err
	}
	return nil
}

// StartRequestGC starts the GC of expired requests
func (r *Ranch) StartRequestGC(gcPeriod time.Duration) {
	r.requestMgr.StartGC(gcPeriod)
}

// Metric returns a metric object with metrics filled in
func (r *Ranch) Metric(rtype string) (common.Metric, error) {
	metric := common.Metric{
		Type:    rtype,
		Current: map[string]int{},
		Owners:  map[string]int{},
	}

	resources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Error("cannot find resources")
		return metric, &ResourceNotFound{rtype}
	}

	for _, res := range resources {
		if res.Type != rtype {
			continue
		}

		if _, ok := metric.Current[res.State]; !ok {
			metric.Current[res.State] = 0
		}

		if _, ok := metric.Owners[res.Owner]; !ok {
			metric.Owners[res.Owner] = 0
		}

		metric.Current[res.State]++
		metric.Owners[res.Owner]++
	}

	if len(metric.Current) == 0 && len(metric.Owners) == 0 {
		return metric, &ResourceNotFound{rtype}
	}

	return metric, nil
}
<commit_msg>boskos: ranch: log lock acquire time<commit_after>/*
Copyright 2017 The Kubernetes 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 ranch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/boskos/common"
)

// Ranch is the place which all of the Resource objects lives.
type Ranch struct {
	Storage       *Storage
	resourcesLock sync.RWMutex
	requestMgr    *RequestManager
	//
	now func() time.Time
}

func updateTime() time.Time {
	return time.Now()
}

// Public errors:

// ResourceNotFound will be returned if requested resource does not exist.
type ResourceNotFound struct {
	name string
}

func (r ResourceNotFound) Error() string {
	return fmt.Sprintf("no available resource %s, try again later.", r.name)
}

// ResourceTypeNotFound will be returned if requested resource type does not exist.
type ResourceTypeNotFound struct {
	rType string
}

func (r ResourceTypeNotFound) Error() string {
	return fmt.Sprintf("resource type %s does not exist", r.rType)
}

// OwnerNotMatch will be returned if request owner does not match current owner for target resource.
type OwnerNotMatch struct {
	request string
	owner   string
}

func (o OwnerNotMatch) Error() string {
	return fmt.Sprintf("owner mismatch request by %s, currently owned by %s", o.request, o.owner)
}

// StateNotMatch will be returned if requested state does not match current state for target resource.
type StateNotMatch struct {
	expect  string
	current string
}

func (s StateNotMatch) Error() string {
	return fmt.Sprintf("state mismatch - expected %v, current %v", s.expect, s.current)
}

// NewRanch creates a new Ranch object.
// In: config - path to resource file
//     storage - path to where to save/restore the state data
// Out: A Ranch object, loaded from config/storage, or error
func NewRanch(config string, s *Storage, ttl time.Duration) (*Ranch, error) {
	newRanch := &Ranch{
		Storage:    s,
		requestMgr: NewRequestManager(ttl),
		now:        time.Now,
	}
	if config != "" {
		if err := newRanch.SyncConfig(config); err != nil {
			return nil, err
		}
		logrus.Infof("Loaded Boskos configuration successfully")
	}
	return newRanch, nil
}

// acquireRequestPriorityKey is used as key for request priority cache.
type acquireRequestPriorityKey struct {
	rType, state string
}

// Acquire checks out a type of resource in certain state without an owner,
// and move the checked out resource to the end of the resource list.
// In: rtype - name of the target resource
//     state - current state of the requested resource
//     dest - destination state of the requested resource
//     owner - requester of the resource
//     requestID - request ID to get a priority in the queue
// Out: A valid Resource object on success, or
//      ResourceNotFound error if target type resource does not exist in target state.
func (r *Ranch) Acquire(rType, state, dest, owner, requestID string) (*common.Resource, error) {
	logger := logrus.WithFields(logrus.Fields{
		"type":       rType,
		"state":      state,
		"dest":       dest,
		"owner":      owner,
		"identifier": requestID,
	})
	logger.Debug("Acquiring resource lock...")
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	logger.Debug("Determining request priority...")
	ts := acquireRequestPriorityKey{rType: rType, state: state}
	rank, new := r.requestMgr.GetRank(ts, requestID)
	logger.WithFields(logrus.Fields{"rank": rank, "new": new}).Debug("Determined request priority.")

	resources, err := r.Storage.GetResources()
	if err != nil {
		logger.WithError(err).Errorf("could not get resources")
		return nil, &ResourceNotFound{rType}
	}
	logger.Debugf("Considering %d resources.", len(resources))

	// For request priority we need to go over all the list until a matching rank
	matchingResoucesCount := 0
	typeCount := 0
	for idx := range resources {
		res := resources[idx]
		if rType == res.Type {
			typeCount++
			if state == res.State && res.Owner == "" {
				matchingResoucesCount++
				if matchingResoucesCount >= rank {
					logger = logger.WithField("resource", res.Name)
					res.Owner = owner
					res.State = dest
					logger.Debug("Updating resource.")
					updatedRes, err := r.Storage.UpdateResource(res)
					if err != nil {
						logger.WithError(err).Errorf("could not update resource %s", res.Name)
						return nil, err
					}
					// Deleting this request since it has been fulfilled
					if requestID != "" {
						logger.Debug("Cleaning up requests.")
						r.requestMgr.Delete(ts, requestID)
					}
					logger.Debug("Successfully acquired resource.")
					return &updatedRes, nil
				}
			}
		}
	}

	if new {
		logger.Debug("Checking for associated dynamic resource type...")
		lifeCycle, err := r.Storage.GetDynamicResourceLifeCycle(rType)
		// Assuming error means no associated dynamic resource
		if err == nil {
			if typeCount < lifeCycle.MaxCount {
				logger.Debug("Adding new dynamic resources...")
				res := common.NewResourceFromNewDynamicResourceLifeCycle(r.Storage.generateName(), &lifeCycle, r.now())
				if err := r.Storage.AddResource(res); err != nil {
					logger.WithError(err).Warningf("unable to add a new resource of type %s", rType)
				}
				logger.Infof("Added dynamic resource %s of type %s", res.Name, res.Type)
			}
		}
	}

	if typeCount > 0 {
		return nil, &ResourceNotFound{rType}
	}
	return nil, &ResourceTypeNotFound{rType}
}

// AcquireByState checks out resources of a given type without an owner,
// that matches a list of resources names.
// In: state - current state of the requested resource
//     dest - destination state of the requested resource
//     owner - requester of the resource
//     names - names of resource to acquire
// Out: A valid list of Resource object on success, or
//      ResourceNotFound error if target type resource does not exist in target state.
func (r *Ranch) AcquireByState(state, dest, owner string, names []string) ([]common.Resource, error) {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	if names == nil {
		return nil, fmt.Errorf("must provide names of expected resources")
	}

	rNames := map[string]bool{}
	for _, t := range names {
		rNames[t] = true
	}

	allResources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Errorf("could not get resources")
		return nil, &ResourceNotFound{state}
	}

	var resources []common.Resource

	for idx := range allResources {
		res := allResources[idx]
		if state == res.State {
			if res.Owner != "" {
				continue
			}
			if rNames[res.Name] {
				res.Owner = owner
				res.State = dest
				updatedRes, err := r.Storage.UpdateResource(res)
				if err != nil {
					logrus.WithError(err).Errorf("could not update resource %s", res.Name)
					return nil, err
				}
				resources = append(resources, updatedRes)
				delete(rNames, res.Name)
			}
		}
	}
	if len(rNames) != 0 {
		var missingResources []string
		for n := range rNames {
			missingResources = append(missingResources, n)
		}
		err := &ResourceNotFound{state}
		logrus.WithError(err).Errorf("could not find required resources %s", strings.Join(missingResources, ", "))
		return resources, err
	}
	return resources, nil
}

// Release unsets owner for target resource and move it to a new state.
// In: name - name of the target resource
//     dest - destination state of the resource
//     owner - owner of the resource
// Out: nil on success, or
//      OwnerNotMatch error if owner does not match current owner of the resource, or
//      ResourceNotFound error if target named resource does not exist.
func (r *Ranch) Release(name, dest, owner string) error {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	res, err := r.Storage.GetResource(name)
	if err != nil {
		logrus.WithError(err).Errorf("unable to release resource %s", name)
		return &ResourceNotFound{name}
	}
	if owner != res.Owner {
		return &OwnerNotMatch{request: owner, owner: res.Owner}
	}

	res.Owner = ""
	res.State = dest

	if lf, err := r.Storage.GetDynamicResourceLifeCycle(res.Type); err == nil {
		// Assuming error means not existing as the only way to differentiate would be to list
		// all resources and find the right one which is more costly.
		if lf.LifeSpan != nil {
			expirationTime := r.now().Add(*lf.LifeSpan)
			res.ExpirationDate = &expirationTime
		}
	} else {
		res.ExpirationDate = nil
	}

	if _, err := r.Storage.UpdateResource(res); err != nil {
		logrus.WithError(err).Errorf("could not update resource %s", res.Name)
		return err
	}
	return nil
}

// Update updates the timestamp of a target resource.
// In: name  - name of the target resource
//     state - current state of the resource
//     owner - current owner of the resource
// 	   info  - information on how to use the resource
// Out: nil on success, or
//      OwnerNotMatch error if owner does not match current owner of the resource, or
//      ResourceNotFound error if target named resource does not exist, or
//      StateNotMatch error if state does not match current state of the resource.
func (r *Ranch) Update(name, owner, state string, ud *common.UserData) error {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	res, err := r.Storage.GetResource(name)
	if err != nil {
		logrus.WithError(err).Errorf("could not find resource %s for update", name)
		return &ResourceNotFound{name}
	}
	if owner != res.Owner {
		return &OwnerNotMatch{request: owner, owner: res.Owner}
	}
	if state != res.State {
		return &StateNotMatch{res.State, state}
	}
	if res.UserData == nil {
		res.UserData = &common.UserData{}
	}
	res.UserData.Update(ud)
	if _, err := r.Storage.UpdateResource(res); err != nil {
		logrus.WithError(err).Errorf("could not update resource %s", res.Name)
		return err
	}
	return nil
}

// Reset unstucks a type of stale resource to a new state.
// In: rtype - type of the resource
//     state - current state of the resource
//     expire - duration before resource's last update
//     dest - destination state of expired resources
// Out: map of resource name - resource owner.
func (r *Ranch) Reset(rtype, state string, expire time.Duration, dest string) (map[string]string, error) {
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()

	ret := make(map[string]string)

	resources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Errorf("cannot find resources")
		return nil, err
	}

	for idx := range resources {
		res := resources[idx]
		if rtype == res.Type && state == res.State && res.Owner != "" {
			if r.now().Sub(res.LastUpdate) > expire {
				ret[res.Name] = res.Owner
				res.Owner = ""
				res.State = dest
				if _, err := r.Storage.UpdateResource(res); err != nil {
					logrus.WithError(err).Errorf("could not update resource %s", res.Name)
					return ret, err
				}
			}
		}
	}
	return ret, nil
}

// SyncConfig updates resource list from a file
func (r *Ranch) SyncConfig(configPath string) error {
	config, err := common.ParseConfig(configPath)
	if err != nil {
		return err
	}
	if err := common.ValidateConfig(config); err != nil {
		return err
	}
	r.resourcesLock.Lock()
	defer r.resourcesLock.Unlock()
	if err := r.Storage.SyncResources(config); err != nil {
		return err
	}
	return nil
}

// StartRequestGC starts the GC of expired requests
func (r *Ranch) StartRequestGC(gcPeriod time.Duration) {
	r.requestMgr.StartGC(gcPeriod)
}

// Metric returns a metric object with metrics filled in
func (r *Ranch) Metric(rtype string) (common.Metric, error) {
	metric := common.Metric{
		Type:    rtype,
		Current: map[string]int{},
		Owners:  map[string]int{},
	}

	resources, err := r.Storage.GetResources()
	if err != nil {
		logrus.WithError(err).Error("cannot find resources")
		return metric, &ResourceNotFound{rtype}
	}

	for _, res := range resources {
		if res.Type != rtype {
			continue
		}

		if _, ok := metric.Current[res.State]; !ok {
			metric.Current[res.State] = 0
		}

		if _, ok := metric.Owners[res.Owner]; !ok {
			metric.Owners[res.Owner] = 0
		}

		metric.Current[res.State]++
		metric.Owners[res.Owner]++
	}

	if len(metric.Current) == 0 && len(metric.Owners) == 0 {
		return metric, &ResourceNotFound{rtype}
	}

	return metric, nil
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package queue

import (
	"fmt"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/database/versiondb"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/utils/wrappers"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	// StatusUpdateFrequency is how many containers should be processed between
	// logs
	StatusUpdateFrequency = 2500
)

// Jobs tracks a series of jobs that form a DAG of dependencies.
type Jobs struct {
	// db ensures that database updates are atomically updated.
	db *versiondb.Database
	// state writes the job queue to [db].
	state *state
}

// New attempts to create a new job queue from the provided database.
func New(
	db database.Database,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) (*Jobs, error) {
	vdb := versiondb.New(db)
	state, err := newState(vdb, metricsNamespace, metricsRegisterer)
	if err != nil {
		return nil, fmt.Errorf("couldn't create new jobs state: %s", err)
	}

	return &Jobs{
		db:    vdb,
		state: state,
	}, nil
}

// SetParser tells this job queue how to parse jobs from the database.
func (j *Jobs) SetParser(parser Parser) error { j.state.parser = parser; return nil }

func (j *Jobs) Has(jobID ids.ID) (bool, error) { return j.state.HasJob(jobID) }

// Push adds a new job to the queue. Returns true if [job] was added to the queue and false
// if [job] was already in the queue.
func (j *Jobs) Push(job Job) (bool, error) {
	jobID := job.ID()
	if has, err := j.state.HasJob(jobID); err != nil {
		return false, fmt.Errorf("failed to check for existing job %s due to %w", jobID, err)
	} else if has {
		return false, nil
	}

	deps, err := job.MissingDependencies()
	if err != nil {
		return false, err
	}
	// Store this job into the database.
	if err := j.state.PutJob(job); err != nil {
		return false, fmt.Errorf("failed to write job due to %w", err)
	}

	if deps.Len() != 0 {
		// This job needs to block on a set of dependencies.
		for depID := range deps {
			if err := j.state.AddDependency(depID, jobID); err != nil {
				return false, fmt.Errorf("failed to add blocking for depID %s, jobID %s", depID, jobID)
			}
		}
		return true, nil
	}
	// This job doesn't have any dependencies, so it should be placed onto the
	// executable stack.
	if err := j.state.AddRunnableJob(jobID); err != nil {
		return false, fmt.Errorf("failed to add %s as a runnable job due to %w", jobID, err)
	}
	return true, nil
}

func (j *Jobs) ExecuteAll(ctx *snow.Context, restarted bool, events ...snow.EventDispatcher) (int, error) {
	numExecuted := 0

	// Disable and clear state caches since the hit rate will be near 0 during execution.
	j.state.DisableCaching()
	for {
		job, err := j.state.RemoveRunnableJob()
		if err == database.ErrNotFound {
			break
		}
		if err != nil {
			return 0, fmt.Errorf("failed to removing runnable job with %w", err)
		}

		jobID := job.ID()
		ctx.Log.Debug("Executing: %s", jobID)
		// Note that event.Accept must be called before executing [job]
		// to honor EventDispatcher.Accept's invariant.
		for _, event := range events {
			if err := event.Accept(ctx, job.ID(), job.Bytes()); err != nil {
				return numExecuted, err
			}
		}
		if err := job.Execute(); err != nil {
			return 0, fmt.Errorf("failed to execute job %s due to %w", jobID, err)
		}

		dependentIDs, err := j.state.RemoveDependencies(jobID)
		if err != nil {
			return 0, fmt.Errorf("failed to remove blocking jobs for %s due to %w", jobID, err)
		}

		for _, dependentID := range dependentIDs {
			job, err := j.state.GetJob(dependentID)
			if err != nil {
				return 0, fmt.Errorf("failed to get job %s from blocking jobs due to %w", dependentID, err)
			}
			deps, err := job.MissingDependencies()
			if err != nil {
				return 0, fmt.Errorf("failed to get missing dependencies for %s due to %w", dependentID, err)
			}
			if deps.Len() > 0 {
				continue
			}
			if err := j.state.AddRunnableJob(dependentID); err != nil {
				return 0, fmt.Errorf("failed to add %s as a runnable job due to %w", dependentID, err)
			}
		}
		if err := j.Commit(); err != nil {
			return 0, err
		}

		numExecuted++
		if numExecuted%StatusUpdateFrequency == 0 { // Periodically print progress
			if !restarted {
				ctx.Log.Info("executed %d operations", numExecuted)
			} else {
				ctx.Log.Debug("executed %d operations", numExecuted)
			}
		}
	}

	if !restarted {
		ctx.Log.Info("executed %d operations", numExecuted)
	} else {
		ctx.Log.Debug("executed %d operations", numExecuted)
	}
	return numExecuted, nil
}

// Commit the versionDB to the underlying database.
func (j *Jobs) Commit() error {
	return j.db.Commit()
}

type JobsWithMissing struct {
	*Jobs

	// keep the missing ID set in memory to avoid unnecessary database reads and
	// writes.
	missingIDs                            ids.Set
	removeFromMissingIDs, addToMissingIDs ids.Set
}

func NewWithMissing(
	db database.Database,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) (*JobsWithMissing, error) {
	innerJobs, err := New(db, metricsNamespace, metricsRegisterer)
	if err != nil {
		return nil, err
	}

	jobs := &JobsWithMissing{
		Jobs: innerJobs,
	}

	missingIDs, err := jobs.state.MissingJobIDs()
	jobs.missingIDs.Add(missingIDs...)
	return jobs, err
}

// SetParser tells this job queue how to parse jobs from the database.
func (jm *JobsWithMissing) SetParser(parser Parser) error {
	jm.state.parser = parser
	return jm.cleanRunnableStack()
}

func (jm *JobsWithMissing) Has(jobID ids.ID) (bool, error) {
	if jm.missingIDs.Contains(jobID) {
		return false, nil
	}

	return jm.Jobs.Has(jobID)
}

// Push adds a new job to the queue. Returns true if [job] was added to the queue and false
// if [job] was already in the queue.
func (jm *JobsWithMissing) Push(job Job) (bool, error) {
	jobID := job.ID()
	if has, err := jm.Has(jobID); err != nil {
		return false, fmt.Errorf("failed to check for existing job %s due to %w", jobID, err)
	} else if has {
		return false, nil
	}

	deps, err := job.MissingDependencies()
	if err != nil {
		return false, err
	}
	// Store this job into the database.
	if err := jm.state.PutJob(job); err != nil {
		return false, fmt.Errorf("failed to write job due to %w", err)
	}

	if deps.Len() != 0 {
		// This job needs to block on a set of dependencies.
		for depID := range deps {
			if err := jm.state.AddDependency(depID, jobID); err != nil {
				return false, fmt.Errorf("failed to add blocking for depID %s, jobID %s", depID, jobID)
			}
		}
		return true, nil
	}
	// This job doesn't have any dependencies, so it should be placed onto the
	// executable stack.
	if err := jm.state.AddRunnableJob(jobID); err != nil {
		return false, fmt.Errorf("failed to add %s as a runnable job due to %w", jobID, err)
	}
	return true, nil
}

// AddMissingID adds [jobID] to missingIDs
func (jm *JobsWithMissing) AddMissingID(jobIDs ...ids.ID) {
	for _, jobID := range jobIDs {
		if !jm.missingIDs.Contains(jobID) {
			jm.missingIDs.Add(jobID)
			jm.addToMissingIDs.Add(jobID)
			jm.removeFromMissingIDs.Remove(jobID)
		}
	}
}

// RemoveMissingID removes [jobID] from missingIDs
func (jm *JobsWithMissing) RemoveMissingID(jobIDs ...ids.ID) {
	for _, jobID := range jobIDs {
		if jm.missingIDs.Contains(jobID) {
			jm.missingIDs.Remove(jobID)
			jm.addToMissingIDs.Remove(jobID)
			jm.removeFromMissingIDs.Add(jobID)
		}
	}
}

func (jm *JobsWithMissing) MissingIDs() []ids.ID { return jm.missingIDs.List() }

func (jm *JobsWithMissing) NumMissingIDs() int { return jm.missingIDs.Len() }

// Commit the versionDB to the underlying database.
func (jm *JobsWithMissing) Commit() error {
	if jm.addToMissingIDs.Len() != 0 {
		if err := jm.state.AddMissingJobIDs(jm.addToMissingIDs); err != nil {
			return err
		}
		jm.addToMissingIDs.Clear()
	}
	if jm.removeFromMissingIDs.Len() != 0 {
		if err := jm.state.RemoveMissingJobIDs(jm.removeFromMissingIDs); err != nil {
			return err
		}
		jm.removeFromMissingIDs.Clear()
	}
	return jm.Jobs.Commit()
}

// cleanRunnableStack iterates over the jobs on the runnable stack and resets any job
// that has missing dependencies to block on those dependencies.
// Note: the jobs queue ensures that no job with missing dependencies will be placed
// on the runnable stack in the first place.
// However, for specific VM implementations blocks may be committed via a two stage commit
// (ex. platformvm Proposal and Commit/Abort blocks). This can cause an issue where if the first stage
// is executed immediately before the node dies, it will be removed from the runnable stack
// without writing the state transition to the VM's database. When the node restarts, the
// VM will not have marked the first block (the proposal block as accepted), but it could
// have already been removed from the jobs queue. cleanRunnableStack handles this case.
func (jm *JobsWithMissing) cleanRunnableStack() error {
	runnableJobsIter := jm.state.runnableJobIDs.NewIterator()
	defer runnableJobsIter.Release()

	for runnableJobsIter.Next() {
		jobIDBytes := runnableJobsIter.Key()
		jobID, err := ids.ToID(jobIDBytes)
		if err != nil {
			return fmt.Errorf("failed to convert jobID bytes into ID due to: %w", err)
		}

		job, err := jm.state.GetJob(jobID)
		if err != nil {
			return fmt.Errorf("failed to retrieve job on runnnable stack due to: %w", err)
		}
		deps, err := job.MissingDependencies()
		if err != nil {
			return fmt.Errorf("failed to retrieve missing dependencies of job on runnable stack due to: %w", err)
		}
		if deps.Len() == 0 {
			continue
		}

		// If the job has missing dependencies, remove it from the runnable stack
		if err := jm.state.runnableJobIDs.Delete(jobIDBytes); err != nil {
			return fmt.Errorf("failed to delete jobID from runnable stack due to: %w", err)
		}

		// Add the missing dependencies to the set that needs to be fetched.
		jm.AddMissingID(deps.List()...)
		for depID := range deps {
			if err := jm.state.AddDependency(depID, jobID); err != nil {
				return fmt.Errorf("failed to add blocking for depID %s, jobID %s while cleaning the runnable stack", depID, jobID)
			}
		}
	}

	errs := wrappers.Errs{}
	errs.Add(
		runnableJobsIter.Error(),
		jm.Commit(),
	)
	return errs.Err
}
<commit_msg>Add TODO to remove DisableCaching VM includes free from storage<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package queue

import (
	"fmt"

	"github.com/ava-labs/avalanchego/database"
	"github.com/ava-labs/avalanchego/database/versiondb"
	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/utils/wrappers"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	// StatusUpdateFrequency is how many containers should be processed between
	// logs
	StatusUpdateFrequency = 2500
)

// Jobs tracks a series of jobs that form a DAG of dependencies.
type Jobs struct {
	// db ensures that database updates are atomically updated.
	db *versiondb.Database
	// state writes the job queue to [db].
	state *state
}

// New attempts to create a new job queue from the provided database.
func New(
	db database.Database,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) (*Jobs, error) {
	vdb := versiondb.New(db)
	state, err := newState(vdb, metricsNamespace, metricsRegisterer)
	if err != nil {
		return nil, fmt.Errorf("couldn't create new jobs state: %s", err)
	}

	return &Jobs{
		db:    vdb,
		state: state,
	}, nil
}

// SetParser tells this job queue how to parse jobs from the database.
func (j *Jobs) SetParser(parser Parser) error { j.state.parser = parser; return nil }

func (j *Jobs) Has(jobID ids.ID) (bool, error) { return j.state.HasJob(jobID) }

// Push adds a new job to the queue. Returns true if [job] was added to the queue and false
// if [job] was already in the queue.
func (j *Jobs) Push(job Job) (bool, error) {
	jobID := job.ID()
	if has, err := j.state.HasJob(jobID); err != nil {
		return false, fmt.Errorf("failed to check for existing job %s due to %w", jobID, err)
	} else if has {
		return false, nil
	}

	deps, err := job.MissingDependencies()
	if err != nil {
		return false, err
	}
	// Store this job into the database.
	if err := j.state.PutJob(job); err != nil {
		return false, fmt.Errorf("failed to write job due to %w", err)
	}

	if deps.Len() != 0 {
		// This job needs to block on a set of dependencies.
		for depID := range deps {
			if err := j.state.AddDependency(depID, jobID); err != nil {
				return false, fmt.Errorf("failed to add blocking for depID %s, jobID %s", depID, jobID)
			}
		}
		return true, nil
	}
	// This job doesn't have any dependencies, so it should be placed onto the
	// executable stack.
	if err := j.state.AddRunnableJob(jobID); err != nil {
		return false, fmt.Errorf("failed to add %s as a runnable job due to %w", jobID, err)
	}
	return true, nil
}

func (j *Jobs) ExecuteAll(ctx *snow.Context, restarted bool, events ...snow.EventDispatcher) (int, error) {
	numExecuted := 0

	// Disable and clear state caches to prevent us from attempting to execute
	// a vertex that was previously parsed, but not saved to the VM. Some VMs
	// may only persist containers when they are accepted. This is a stop-gap
	// measure to ensure the job will be re-parsed before executing until the VM
	// provides a more explicit interface for freeing parsed blocks.
	// TODO remove DisableCaching when VM provides better interface for freeing
	// blocks.
	j.state.DisableCaching()
	for {
		job, err := j.state.RemoveRunnableJob()
		if err == database.ErrNotFound {
			break
		}
		if err != nil {
			return 0, fmt.Errorf("failed to removing runnable job with %w", err)
		}

		jobID := job.ID()
		ctx.Log.Debug("Executing: %s", jobID)
		// Note that event.Accept must be called before executing [job]
		// to honor EventDispatcher.Accept's invariant.
		for _, event := range events {
			if err := event.Accept(ctx, job.ID(), job.Bytes()); err != nil {
				return numExecuted, err
			}
		}
		if err := job.Execute(); err != nil {
			return 0, fmt.Errorf("failed to execute job %s due to %w", jobID, err)
		}

		dependentIDs, err := j.state.RemoveDependencies(jobID)
		if err != nil {
			return 0, fmt.Errorf("failed to remove blocking jobs for %s due to %w", jobID, err)
		}

		for _, dependentID := range dependentIDs {
			job, err := j.state.GetJob(dependentID)
			if err != nil {
				return 0, fmt.Errorf("failed to get job %s from blocking jobs due to %w", dependentID, err)
			}
			deps, err := job.MissingDependencies()
			if err != nil {
				return 0, fmt.Errorf("failed to get missing dependencies for %s due to %w", dependentID, err)
			}
			if deps.Len() > 0 {
				continue
			}
			if err := j.state.AddRunnableJob(dependentID); err != nil {
				return 0, fmt.Errorf("failed to add %s as a runnable job due to %w", dependentID, err)
			}
		}
		if err := j.Commit(); err != nil {
			return 0, err
		}

		numExecuted++
		if numExecuted%StatusUpdateFrequency == 0 { // Periodically print progress
			if !restarted {
				ctx.Log.Info("executed %d operations", numExecuted)
			} else {
				ctx.Log.Debug("executed %d operations", numExecuted)
			}
		}
	}

	if !restarted {
		ctx.Log.Info("executed %d operations", numExecuted)
	} else {
		ctx.Log.Debug("executed %d operations", numExecuted)
	}
	return numExecuted, nil
}

// Commit the versionDB to the underlying database.
func (j *Jobs) Commit() error {
	return j.db.Commit()
}

type JobsWithMissing struct {
	*Jobs

	// keep the missing ID set in memory to avoid unnecessary database reads and
	// writes.
	missingIDs                            ids.Set
	removeFromMissingIDs, addToMissingIDs ids.Set
}

func NewWithMissing(
	db database.Database,
	metricsNamespace string,
	metricsRegisterer prometheus.Registerer,
) (*JobsWithMissing, error) {
	innerJobs, err := New(db, metricsNamespace, metricsRegisterer)
	if err != nil {
		return nil, err
	}

	jobs := &JobsWithMissing{
		Jobs: innerJobs,
	}

	missingIDs, err := jobs.state.MissingJobIDs()
	jobs.missingIDs.Add(missingIDs...)
	return jobs, err
}

// SetParser tells this job queue how to parse jobs from the database.
func (jm *JobsWithMissing) SetParser(parser Parser) error {
	jm.state.parser = parser
	return jm.cleanRunnableStack()
}

func (jm *JobsWithMissing) Has(jobID ids.ID) (bool, error) {
	if jm.missingIDs.Contains(jobID) {
		return false, nil
	}

	return jm.Jobs.Has(jobID)
}

// Push adds a new job to the queue. Returns true if [job] was added to the queue and false
// if [job] was already in the queue.
func (jm *JobsWithMissing) Push(job Job) (bool, error) {
	jobID := job.ID()
	if has, err := jm.Has(jobID); err != nil {
		return false, fmt.Errorf("failed to check for existing job %s due to %w", jobID, err)
	} else if has {
		return false, nil
	}

	deps, err := job.MissingDependencies()
	if err != nil {
		return false, err
	}
	// Store this job into the database.
	if err := jm.state.PutJob(job); err != nil {
		return false, fmt.Errorf("failed to write job due to %w", err)
	}

	if deps.Len() != 0 {
		// This job needs to block on a set of dependencies.
		for depID := range deps {
			if err := jm.state.AddDependency(depID, jobID); err != nil {
				return false, fmt.Errorf("failed to add blocking for depID %s, jobID %s", depID, jobID)
			}
		}
		return true, nil
	}
	// This job doesn't have any dependencies, so it should be placed onto the
	// executable stack.
	if err := jm.state.AddRunnableJob(jobID); err != nil {
		return false, fmt.Errorf("failed to add %s as a runnable job due to %w", jobID, err)
	}
	return true, nil
}

// AddMissingID adds [jobID] to missingIDs
func (jm *JobsWithMissing) AddMissingID(jobIDs ...ids.ID) {
	for _, jobID := range jobIDs {
		if !jm.missingIDs.Contains(jobID) {
			jm.missingIDs.Add(jobID)
			jm.addToMissingIDs.Add(jobID)
			jm.removeFromMissingIDs.Remove(jobID)
		}
	}
}

// RemoveMissingID removes [jobID] from missingIDs
func (jm *JobsWithMissing) RemoveMissingID(jobIDs ...ids.ID) {
	for _, jobID := range jobIDs {
		if jm.missingIDs.Contains(jobID) {
			jm.missingIDs.Remove(jobID)
			jm.addToMissingIDs.Remove(jobID)
			jm.removeFromMissingIDs.Add(jobID)
		}
	}
}

func (jm *JobsWithMissing) MissingIDs() []ids.ID { return jm.missingIDs.List() }

func (jm *JobsWithMissing) NumMissingIDs() int { return jm.missingIDs.Len() }

// Commit the versionDB to the underlying database.
func (jm *JobsWithMissing) Commit() error {
	if jm.addToMissingIDs.Len() != 0 {
		if err := jm.state.AddMissingJobIDs(jm.addToMissingIDs); err != nil {
			return err
		}
		jm.addToMissingIDs.Clear()
	}
	if jm.removeFromMissingIDs.Len() != 0 {
		if err := jm.state.RemoveMissingJobIDs(jm.removeFromMissingIDs); err != nil {
			return err
		}
		jm.removeFromMissingIDs.Clear()
	}
	return jm.Jobs.Commit()
}

// cleanRunnableStack iterates over the jobs on the runnable stack and resets any job
// that has missing dependencies to block on those dependencies.
// Note: the jobs queue ensures that no job with missing dependencies will be placed
// on the runnable stack in the first place.
// However, for specific VM implementations blocks may be committed via a two stage commit
// (ex. platformvm Proposal and Commit/Abort blocks). This can cause an issue where if the first stage
// is executed immediately before the node dies, it will be removed from the runnable stack
// without writing the state transition to the VM's database. When the node restarts, the
// VM will not have marked the first block (the proposal block as accepted), but it could
// have already been removed from the jobs queue. cleanRunnableStack handles this case.
func (jm *JobsWithMissing) cleanRunnableStack() error {
	runnableJobsIter := jm.state.runnableJobIDs.NewIterator()
	defer runnableJobsIter.Release()

	for runnableJobsIter.Next() {
		jobIDBytes := runnableJobsIter.Key()
		jobID, err := ids.ToID(jobIDBytes)
		if err != nil {
			return fmt.Errorf("failed to convert jobID bytes into ID due to: %w", err)
		}

		job, err := jm.state.GetJob(jobID)
		if err != nil {
			return fmt.Errorf("failed to retrieve job on runnnable stack due to: %w", err)
		}
		deps, err := job.MissingDependencies()
		if err != nil {
			return fmt.Errorf("failed to retrieve missing dependencies of job on runnable stack due to: %w", err)
		}
		if deps.Len() == 0 {
			continue
		}

		// If the job has missing dependencies, remove it from the runnable stack
		if err := jm.state.runnableJobIDs.Delete(jobIDBytes); err != nil {
			return fmt.Errorf("failed to delete jobID from runnable stack due to: %w", err)
		}

		// Add the missing dependencies to the set that needs to be fetched.
		jm.AddMissingID(deps.List()...)
		for depID := range deps {
			if err := jm.state.AddDependency(depID, jobID); err != nil {
				return fmt.Errorf("failed to add blocking for depID %s, jobID %s while cleaning the runnable stack", depID, jobID)
			}
		}
	}

	errs := wrappers.Errs{}
	errs.Add(
		runnableJobsIter.Error(),
		jm.Commit(),
	)
	return errs.Err
}
<|endoftext|>


<commit_before>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

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

//go:generate sh generate.sh

var (
	Root           *Resolver
	DebugLogger    io.Writer
	Timeout        = 1000 * time.Millisecond
	MaxRecursion   = 10
	MaxNameservers = 2
	MaxIPs         = 2
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{
			DialTimeout:  Timeout,
			ReadTimeout:  Timeout,
			WriteTimeout: Timeout,
		},
		cache: cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a slice of *RR.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will return an empty, non-nil slice.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) []*RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) []*RR {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil
	}
	qname = toLowerFQDN(qname)
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		return rrs
	}
	logResolveStart(qname, qtype, depth)
	defer logResolveEnd(qname, qtype, depth, time.Now())
	return r.resolveNS(qname, qtype, depth)
}

func (r *Resolver) resolveNS(qname string, qtype string, depth int) []*RR {
	success := make(chan bool, 1)
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		if pname == qname && qtype == "NS" { // If we’re looking for [foo.com,NS], then skip to [com,NS]
			continue
		}

		// Query all DNS servers in parallel
		count := 0
		for _, nrr := range r.resolve(pname, "NS", depth) {
			if qtype != "" { // Early out for specific queries
				if rrs := r.cacheGet(qname, qtype); rrs != nil {
					return rrs
				}
			}
			if nrr.Type != "NS" {
				continue
			}
			if count++; count > MaxNameservers {
				break
			}
			go r.exchange(success, nrr.Value, qname, qtype, depth)
		}

		// Wait for first response
		if count > 0 {
			select {
			case <-success:
				return r.resolveCNAMEs(qname, qtype, depth)
			case <-time.After(Timeout):
				continue
			}
		}
	}
	return nil
}

func (r *Resolver) exchange(success chan<- bool, host string, qname string, qtype string, depth int) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	for _, rr := range r.resolve(host, "A", depth) {
		if rr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return
		}

		// Synchronously query this DNS server
		start := time.Now()
		rmsg, _, err := r.client.Exchange(qmsg, rr.Value+":53")
		logExchange(rr.Value, qmsg, depth, start, err)
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			r.cacheAdd(qname, nil)
		}

		// If successful, cache the results
		r.saveDNSRR(rmsg.Answer...)
		r.saveDNSRR(rmsg.Ns...)
		r.saveDNSRR(rmsg.Extra...)

		// Never block
		select {
		case success <- true:
		default:
		}

		// Return after first successful network request
		return
	}
}

func (r *Resolver) resolveCNAMEs(qname string, qtype string, depth int) []*RR {
	rrs := []*RR{} // Return non-nil slice indicating difference between NXDOMAIN and an error
	for _, crr := range r.cacheGet(qname, "") {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" {
			continue
		}
		logCNAME(depth, crr.String())
		for _, rr := range r.resolve(crr.Value, qtype, depth) {
			r.cacheAdd(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

func logMaxRecursion(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s Error: MAX RECURSION @ %s %s %d\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveStart(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveEnd(qname string, qtype string, depth int, start time.Time) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, qname, qtype, depth)
}

func logCNAME(depth int, cname string) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth-1), cname)
}

func logExchange(host string, qmsg *dns.Msg, depth int, start time.Time, err error) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, host, qmsg.Question[0].Name, dns.TypeToString[qmsg.Question[0].Qtype])
	if err != nil {
		fmt.Fprintf(DebugLogger, "%s│    %dms: ERROR: %s\n",
			strings.Repeat("│   ", depth-1), dur/time.Millisecond, err.Error())
	}
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

type key struct {
	Name string
	Type string
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<commit_msg>removed unused key type<commit_after>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

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

//go:generate sh generate.sh

var (
	Root           *Resolver
	DebugLogger    io.Writer
	Timeout        = 1000 * time.Millisecond
	MaxRecursion   = 10
	MaxNameservers = 2
	MaxIPs         = 2
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{
			DialTimeout:  Timeout,
			ReadTimeout:  Timeout,
			WriteTimeout: Timeout,
		},
		cache: cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a slice of *RR.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will return an empty, non-nil slice.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) []*RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) []*RR {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		return nil
	}
	qname = toLowerFQDN(qname)
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		return rrs
	}
	logResolveStart(qname, qtype, depth)
	defer logResolveEnd(qname, qtype, depth, time.Now())
	return r.resolveNS(qname, qtype, depth)
}

func (r *Resolver) resolveNS(qname string, qtype string, depth int) []*RR {
	success := make(chan bool, 1)
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		if pname == qname && qtype == "NS" { // If we’re looking for [foo.com,NS], then skip to [com,NS]
			continue
		}

		// Query all DNS servers in parallel
		count := 0
		for _, nrr := range r.resolve(pname, "NS", depth) {
			if qtype != "" { // Early out for specific queries
				if rrs := r.cacheGet(qname, qtype); rrs != nil {
					return rrs
				}
			}
			if nrr.Type != "NS" {
				continue
			}
			if count++; count > MaxNameservers {
				break
			}
			go r.exchange(success, nrr.Value, qname, qtype, depth)
		}

		// Wait for first response
		if count > 0 {
			select {
			case <-success:
				return r.resolveCNAMEs(qname, qtype, depth)
			case <-time.After(Timeout):
				continue
			}
		}
	}
	return nil
}

func (r *Resolver) exchange(success chan<- bool, host string, qname string, qtype string, depth int) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	for _, rr := range r.resolve(host, "A", depth) {
		if rr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return
		}

		// Synchronously query this DNS server
		start := time.Now()
		rmsg, _, err := r.client.Exchange(qmsg, rr.Value+":53")
		logExchange(rr.Value, qmsg, depth, start, err)
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			r.cacheAdd(qname, nil)
		}

		// If successful, cache the results
		r.saveDNSRR(rmsg.Answer...)
		r.saveDNSRR(rmsg.Ns...)
		r.saveDNSRR(rmsg.Extra...)

		// Never block
		select {
		case success <- true:
		default:
		}

		// Return after first successful network request
		return
	}
}

func (r *Resolver) resolveCNAMEs(qname string, qtype string, depth int) []*RR {
	rrs := []*RR{} // Return non-nil slice indicating difference between NXDOMAIN and an error
	for _, crr := range r.cacheGet(qname, "") {
		rrs = append(rrs, crr)
		if crr.Type != "CNAME" {
			continue
		}
		logCNAME(depth, crr.String())
		for _, rr := range r.resolve(crr.Value, qtype, depth) {
			r.cacheAdd(qname, rr)
			rrs = append(rrs, crr)
		}
	}
	return rrs
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

func logMaxRecursion(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s Error: MAX RECURSION @ %s %s %d\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveStart(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveEnd(qname string, qtype string, depth int, start time.Time) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, qname, qtype, depth)
}

func logCNAME(depth int, cname string) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth-1), cname)
}

func logExchange(host string, qmsg *dns.Msg, depth int, start time.Time, err error) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, host, qmsg.Question[0].Name, dns.TypeToString[qmsg.Question[0].Qtype])
	if err != nil {
		fmt.Fprintf(DebugLogger, "%s│    %dms: ERROR: %s\n",
			strings.Repeat("│   ", depth-1), dur/time.Millisecond, err.Error())
	}
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<|endoftext|>


<commit_before>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

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

//go:generate sh generate.sh

var (
	Root           *Resolver
	DebugLogger    io.Writer
	Timeout        = 1000 * time.Millisecond
	MaxRecursion   = 10
	MaxNameservers = 2
	MaxIPs         = 2
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{
			DialTimeout:  Timeout,
			ReadTimeout:  Timeout,
			WriteTimeout: Timeout,
		},
		cache: cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a channel of *RR.
// The implementation guarantees that the output channel will close, so it is safe to range over.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will simply close the output channel.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) <-chan *RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) <-chan *RR {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		c := make(chan *RR)
		close(c)
		return c
	}
	qname = toLowerFQDN(qname)
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		c := make(chan *RR, len(rrs))
		inject(c, rrs...)
		close(c)
		return c
	}
	c := make(chan *RR, 20)
	go func() {
		defer close(c)
		logResolveStart(qname, qtype, depth)
		defer logResolveEnd(qname, qtype, depth, time.Now())
		r.resolveNS(c, qname, qtype, depth)
	}()
	return c
}

func (r *Resolver) recall(c chan<- *RR, qname string, qtype string) bool {
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		inject(c, rrs...)
		return true
	}
	return false
}

func (r *Resolver) resolveNS(c chan<- *RR, qname string, qtype string, depth int) {
	success := make(chan bool, 1)
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		if pname == qname && qtype == "NS" { // If we’re looking for [foo.com,NS], then skip to [com,NS]
			continue
		}

		// Query all DNS servers in parallel
		count := 0
		for nrr := range r.resolve(pname, "NS", depth) {
			if qtype != "" && r.recall(c, qname, qtype) {
				return
			}
			if nrr.Type != "NS" {
				continue
			}
			if count++; count > MaxNameservers {
				break
			}
			go r.exchange(success, c, nrr.Value, qname, qtype, depth)
		}

		// Wait for first response
		if count > 0 {
			select {
			case <-success:
				r.resolveCNAMEs(c, qname, qtype, depth)
				return
			case <-time.After(Timeout):
				continue
			}
		}
	}
}

func (r *Resolver) exchange(success chan<- bool, c chan<- *RR, host string, qname string, qtype string, depth int) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	for rr := range r.resolve(host, "A", depth) {
		if rr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return
		}

		// Synchronously query this DNS server
		start := time.Now()
		rmsg, _, err := r.client.Exchange(qmsg, rr.Value+":53")
		logExchange(rr.Value, qmsg, depth, start, err)
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			r.cacheAdd(qname, nil)
		}

		// If successful, cache the results and return
		r.saveDNSRR(rmsg.Answer...)
		r.saveDNSRR(rmsg.Ns...)
		r.saveDNSRR(rmsg.Extra...)
		select {
			case success <- true: // Never block
		}
		return
	}
}

func (r *Resolver) resolveCNAMEs(c chan<- *RR, qname string, qtype string, depth int) {
	rrs := r.cacheGet(qname, "")
	if rrs == nil || !inject(c, rrs...) {
		return
	}
	for _, crr := range rrs {
		if crr.Type != "CNAME" {
			continue
		}
		logCNAME(depth, crr.String())
		for rr := range r.resolve(crr.Value, qtype, depth) {
			r.cacheAdd(qname, rr)
			if !inject(c, rr) {
				return
			}
		}
	}
}

func inject(c chan<- *RR, rrs ...*RR) bool {
	for _, rr := range rrs {
		select {
		case c <- rr:
		default:
			// return false
		}
	}
	return true
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

func logMaxRecursion(qname string, qtype string, depth int) {
	fmt.Printf("%s Error: MAX RECURSION @ %s %s %d\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveStart(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveEnd(qname string, qtype string, depth int, start time.Time) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, qname, qtype, depth)
}

func logCNAME(depth int, cname string) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth-1), cname)
}

func logExchange(host string, qmsg *dns.Msg, depth int, start time.Time, err error) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, host, qmsg.Question[0].Name, dns.TypeToString[qmsg.Question[0].Qtype])
	if err != nil {
		fmt.Fprintf(DebugLogger, "%s│    %dms: ERROR: %s\n",
			strings.Repeat("│   ", depth-1), dur/time.Millisecond, err.Error())
	}
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

type key struct {
	Name string
	Type string
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<commit_msg>really never block<commit_after>package dnsr

import (
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

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

//go:generate sh generate.sh

var (
	Root           *Resolver
	DebugLogger    io.Writer
	Timeout        = 1000 * time.Millisecond
	MaxRecursion   = 10
	MaxNameservers = 2
	MaxIPs         = 2
)

func init() {
	Root = New(strings.Count(root, "\n"))
	for t := range dns.ParseZone(strings.NewReader(root), "", "") {
		if t.Error == nil {
			Root.saveDNSRR(t.RR)
		}
	}
}

// Resolver implements a primitive, non-recursive, caching DNS resolver.
type Resolver struct {
	cache  *lru.Cache
	client *dns.Client
}

// New initializes a Resolver with the specified cache size. Cache size defaults to 10,000 if size <= 0.
func New(size int) *Resolver {
	if size <= 0 {
		size = 10000
	}
	cache, _ := lru.New(size)
	r := &Resolver{
		client: &dns.Client{
			DialTimeout:  Timeout,
			ReadTimeout:  Timeout,
			WriteTimeout: Timeout,
		},
		cache: cache,
	}
	return r
}

// Resolve finds DNS records of type qtype for the domain qname. It returns a channel of *RR.
// The implementation guarantees that the output channel will close, so it is safe to range over.
// For nonexistent domains (where a DNS server will return NXDOMAIN), it will simply close the output channel.
// Specify an empty string in qtype to receive any DNS records found (currently A, AAAA, NS, CNAME, and TXT).
func (r *Resolver) Resolve(qname string, qtype string) <-chan *RR {
	return r.resolve(qname, qtype, 0)
}

func (r *Resolver) resolve(qname string, qtype string, depth int) <-chan *RR {
	if depth++; depth > MaxRecursion {
		logMaxRecursion(qname, qtype, depth)
		c := make(chan *RR)
		close(c)
		return c
	}
	qname = toLowerFQDN(qname)
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		c := make(chan *RR, len(rrs))
		inject(c, rrs...)
		close(c)
		return c
	}
	c := make(chan *RR, 20)
	go func() {
		defer close(c)
		logResolveStart(qname, qtype, depth)
		defer logResolveEnd(qname, qtype, depth, time.Now())
		r.resolveNS(c, qname, qtype, depth)
	}()
	return c
}

func (r *Resolver) recall(c chan<- *RR, qname string, qtype string) bool {
	if rrs := r.cacheGet(qname, qtype); rrs != nil {
		inject(c, rrs...)
		return true
	}
	return false
}

func (r *Resolver) resolveNS(c chan<- *RR, qname string, qtype string, depth int) {
	success := make(chan bool, 1)
	for pname, ok := qname, true; ok; pname, ok = parent(pname) {
		if pname == qname && qtype == "NS" { // If we’re looking for [foo.com,NS], then skip to [com,NS]
			continue
		}

		// Query all DNS servers in parallel
		count := 0
		for nrr := range r.resolve(pname, "NS", depth) {
			if qtype != "" && r.recall(c, qname, qtype) {
				return
			}
			if nrr.Type != "NS" {
				continue
			}
			if count++; count > MaxNameservers {
				break
			}
			go r.exchange(success, c, nrr.Value, qname, qtype, depth)
		}

		// Wait for first response
		if count > 0 {
			select {
			case <-success:
				r.resolveCNAMEs(c, qname, qtype, depth)
				return
			case <-time.After(Timeout):
				continue
			}
		}
	}
}

func (r *Resolver) exchange(success chan<- bool, c chan<- *RR, host string, qname string, qtype string, depth int) {
	dtype := dns.StringToType[qtype]
	if dtype == 0 {
		dtype = dns.TypeA
	}
	qmsg := &dns.Msg{}
	qmsg.SetQuestion(qname, dtype)
	qmsg.MsgHdr.RecursionDesired = false

	// Find each A record for the DNS server
	count := 0
	for rr := range r.resolve(host, "A", depth) {
		if rr.Type != "A" { // FIXME: support AAAA records?
			continue
		}

		// Never query more than MaxIPs for any nameserver
		if count++; count > MaxIPs {
			return
		}

		// Synchronously query this DNS server
		start := time.Now()
		rmsg, _, err := r.client.Exchange(qmsg, rr.Value+":53")
		logExchange(rr.Value, qmsg, depth, start, err)
		if err != nil {
			continue
		}

		// FIXME: cache NXDOMAIN responses responsibly
		if rmsg.Rcode == dns.RcodeNameError {
			r.cacheAdd(qname, nil)
		}

		// If successful, cache the results and return
		r.saveDNSRR(rmsg.Answer...)
		r.saveDNSRR(rmsg.Ns...)
		r.saveDNSRR(rmsg.Extra...)
		// Never block
		select {
		case success <- true:
		default:
		}
		return
	}
}

func (r *Resolver) resolveCNAMEs(c chan<- *RR, qname string, qtype string, depth int) {
	rrs := r.cacheGet(qname, "")
	if rrs == nil || !inject(c, rrs...) {
		return
	}
	for _, crr := range rrs {
		if crr.Type != "CNAME" {
			continue
		}
		logCNAME(depth, crr.String())
		for rr := range r.resolve(crr.Value, qtype, depth) {
			r.cacheAdd(qname, rr)
			if !inject(c, rr) {
				return
			}
		}
	}
}

func inject(c chan<- *RR, rrs ...*RR) bool {
	for _, rr := range rrs {
		select {
		case c <- rr:
		default:
			// return false
		}
	}
	return true
}

func parent(name string) (string, bool) {
	labels := dns.SplitDomainName(name)
	if labels == nil {
		return "", false
	}
	return toLowerFQDN(strings.Join(labels[1:], ".")), true
}

func toLowerFQDN(name string) string {
	return dns.Fqdn(strings.ToLower(name))
}

func logMaxRecursion(qname string, qtype string, depth int) {
	fmt.Printf("%s Error: MAX RECURSION @ %s %s %d\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveStart(qname string, qtype string, depth int) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s┌─── resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), qname, qtype, depth)
}

func logResolveEnd(qname string, qtype string, depth int, start time.Time) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s└─── %dms: resolve(\"%s\", \"%s\", %d)\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, qname, qtype, depth)
}

func logCNAME(depth int, cname string) {
	if DebugLogger == nil {
		return
	}
	fmt.Fprintf(DebugLogger, "%s│    CNAME: %s\n", strings.Repeat("│   ", depth-1), cname)
}

func logExchange(host string, qmsg *dns.Msg, depth int, start time.Time, err error) {
	if DebugLogger == nil {
		return
	}
	dur := time.Since(start)
	fmt.Fprintf(DebugLogger, "%s│    %dms: dig @%s %s %s\n",
		strings.Repeat("│   ", depth-1), dur/time.Millisecond, host, qmsg.Question[0].Name, dns.TypeToString[qmsg.Question[0].Qtype])
	if err != nil {
		fmt.Fprintf(DebugLogger, "%s│    %dms: ERROR: %s\n",
			strings.Repeat("│   ", depth-1), dur/time.Millisecond, err.Error())
	}
}

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

func convertRR(drr dns.RR) *RR {
	switch t := drr.(type) {
	case *dns.NS:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Ns}
	case *dns.CNAME:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.Target}
	case *dns.A:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.A.String()}
	case *dns.AAAA:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], t.AAAA.String()}
	case *dns.TXT:
		return &RR{t.Hdr.Name, dns.TypeToString[t.Hdr.Rrtype], strings.Join(t.Txt, "\t")}
	default:
		// fmt.Printf("%s\n", drr.String())
	}
	return nil
}

type key struct {
	Name string
	Type string
}

type entry struct {
	m   sync.RWMutex
	rrs map[RR]struct{}
}

// saveDNSRR saves 1 or more DNS records to the resolver cache.
func (r *Resolver) saveDNSRR(drrs ...dns.RR) {
	for _, drr := range drrs {
		if rr := convertRR(drr); rr != nil {
			r.cacheAdd(rr.Name, rr)
		}
	}
}

// cacheAdd adds 0 or more DNS records to the resolver cache for a specific
// domain name and record type. This ensures the cache entry exists, even
// if empty, for NXDOMAIN responses.
func (r *Resolver) cacheAdd(qname string, rr *RR) {
	qname = toLowerFQDN(qname)
	e := r.getEntry(qname)
	if e == nil {
		e = &entry{rrs: make(map[RR]struct{}, 0)}
		e.m.Lock()
		r.cache.Add(qname, e)
	} else {
		e.m.Lock()
	}
	defer e.m.Unlock()
	if rr != nil {
		e.rrs[*rr] = struct{}{}
	}
}

// cacheGet returns a randomly ordered slice of DNS records.
func (r *Resolver) cacheGet(qname string, qtype string) []*RR {
	e := r.getEntry(qname)
	if e == nil && r != Root {
		e = Root.getEntry(qname)
	}
	if e == nil {
		return nil
	}
	e.m.RLock()
	defer e.m.RUnlock()
	if len(e.rrs) == 0 {
		return []*RR{}
	}
	rrs := make([]*RR, 0, len(e.rrs))
	for rr, _ := range e.rrs {
		// fmt.Printf("%s\n", rr.String())
		if qtype == "" || rr.Type == qtype {
			rrs = append(rrs, &RR{rr.Name, rr.Type, rr.Value})
		}
	}
	if len(rrs) == 0 && (qtype != "" && qtype != "NS") {
		return nil
	}
	return rrs
}

// getEntry returns a single cache entry or nil if an entry does not exist in the cache.
func (r *Resolver) getEntry(qname string) *entry {
	c, ok := r.cache.Get(qname)
	if !ok {
		return nil
	}
	e, ok := c.(*entry)
	if !ok {
		return nil
	}
	return e
}
<|endoftext|>


<commit_before>package goridge

import (
	"bytes"
	"io"
	"net/rpc"
	"sync"

	j "github.com/json-iterator/go"
	"github.com/spiral/errors"
	"go.uber.org/multierr"
)

var json = j.ConfigCompatibleWithStandardLibrary

// Codec represent net/rpc bridge over Goridge socket relay.
type Codec struct {
	relay  Relay
	closed bool
	frame  *Frame
	codec  sync.Map
}

// NewCodec initiates new server rpc codec over socket connection.
func NewCodec(rwc io.ReadWriteCloser) *Codec {
	return &Codec{
		relay: NewSocketRelay(rwc),
		codec: sync.Map{},
	}
}

// NewCodecWithRelay initiates new server rpc codec with a relay of choice.
func NewCodecWithRelay(relay Relay) *Codec {
	return &Codec{relay: relay}
}

// WriteResponse marshals response, byte slice or error to remote party.
func (c *Codec) WriteResponse(r *rpc.Response, body interface{}) error {
	const op = errors.Op("codec: write response")
	frame := NewFrame()

	// SEQ_ID + METHOD_NAME_LEN
	frame.WriteOptions(uint32(r.Seq), uint32(len(r.ServiceMethod)))
	// Write protocol version
	frame.WriteVersion(VERSION_1)

	// load and delete associated codec to not waste memory
	// because we write it to the frame and don't need more information about it
	// as for go.14, Load and Delete are separate methods
	codec, ok := c.codec.Load(r.Seq)
	if !ok {
		// fallback codec
		frame.WriteFlags(CODEC_GOB)
	} else {
		frame.WriteFlags(codec.(FrameFlag))
	}

	// delete the key
	c.codec.Delete(r.Seq)

	// initialize buffer
	buf := new(bytes.Buffer)
	// writeServiceMethod to the buffer
	buf.WriteString(r.ServiceMethod)

	// if error returned, we sending it via relay and return error from WriteResponse
	if r.Error != "" {
		// Append error flag
		return c.handleError(r, frame, buf, errors.Str(r.Error))
	}

	// read flag previously written
	// TODO might be better to save it to local variable
	flags := frame.ReadFlags()

	switch {
	case flags&byte(CODEC_JSON) != 0:
		err := encodeJSON(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_MSGPACK) != 0:
		err := encodeMsgPack(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_RAW) != 0:
		err := encodeRaw(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_GOB) != 0:
		err := encodeGob(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	default:
		return c.handleError(r, frame, buf, errors.E(op, errors.Str("unknown codec")))
	}
}

func (c *Codec) sendBuf(frame *Frame, buf *bytes.Buffer) error {
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	return c.relay.Send(frame)
}

func (c *Codec) handleError(r *rpc.Response, frame *Frame, buf *bytes.Buffer, err error) error {
	// just to be sure, remove all data from buffer and write new
	buf.Truncate(0)
	// write all possible errors
	err = multierr.Append(err, errors.Str(r.Error))
	buf.WriteString(r.ServiceMethod)

	const op = errors.Op("handle codec error")
	frame.WriteFlags(ERROR)
	// error should be here
	if err != nil {
		buf.WriteString(err.Error())
	}
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	_ = c.relay.Send(frame)
	return errors.E(op, errors.Str(r.Error))
}

// ReadRequestHeader receives
func (c *Codec) ReadRequestHeader(r *rpc.Request) error {
	const op = errors.Op("codec: read request header")
	frame := NewFrame()
	err := c.relay.Receive(frame)
	if err != nil {
		return err
	}

	// opts[0] sequence ID
	// opts[1] service method name offset from payload in bytes
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}

	r.Seq = uint64(opts[0])
	r.ServiceMethod = string(frame.Payload()[:opts[1]])
	c.frame = frame
	return c.storeCodec(r, frame.ReadFlags())
}

func (c *Codec) storeCodec(r *rpc.Request, flag byte) error {
	switch {
	case flag&byte(CODEC_JSON) != 0:
		c.codec.Store(r.Seq, CODEC_JSON)
	case flag&byte(CODEC_RAW) != 0:
		c.codec.Store(r.Seq, CODEC_RAW)
	case flag&byte(CODEC_MSGPACK) != 0:
		c.codec.Store(r.Seq, CODEC_MSGPACK)
	case flag&byte(CODEC_GOB) != 0:
		c.codec.Store(r.Seq, CODEC_GOB)
	default:
		c.codec.Store(r.Seq, CODEC_GOB)
	}

	return nil
}

// ReadRequestBody fetches prefixed body data and automatically unmarshal it as json. RawBody flag will populate
// []byte lice argument for rpc method.
func (c *Codec) ReadRequestBody(out interface{}) error {
	const op = errors.Op("codec read request body")
	if out == nil {
		return nil
	}

	defer func() {
		c.frame = nil
	}()

	flags := c.frame.ReadFlags()

	if flags&byte(CODEC_JSON) != byte(0) {
		return decodeJSON(out, c.frame)
	}

	if flags&byte(CODEC_GOB) != byte(0) {
		return decodeGob(out, c.frame)
	}

	if flags&byte(CODEC_RAW) != byte(0) {
		return decodeRaw(out, c.frame)
	}

	if flags&byte(CODEC_MSGPACK) != byte(0) {
		return decodeMsgPack(out, c.frame)
	}

	return errors.E(op, errors.Str("unknown decoder used in frame"))
}

// Close underlying socket.
func (c *Codec) Close() error {
	if c.closed {
		return nil
	}

	c.closed = true
	return c.relay.Close()
}
<commit_msg>simplify code<commit_after>package goridge

import (
	"bytes"
	"io"
	"net/rpc"
	"sync"

	j "github.com/json-iterator/go"
	"github.com/spiral/errors"
	"go.uber.org/multierr"
)

var json = j.ConfigCompatibleWithStandardLibrary

// Codec represent net/rpc bridge over Goridge socket relay.
type Codec struct {
	relay  Relay
	closed bool
	frame  *Frame
	codec  sync.Map
}

// NewCodec initiates new server rpc codec over socket connection.
func NewCodec(rwc io.ReadWriteCloser) *Codec {
	return &Codec{
		relay: NewSocketRelay(rwc),
		codec: sync.Map{},
	}
}

// NewCodecWithRelay initiates new server rpc codec with a relay of choice.
func NewCodecWithRelay(relay Relay) *Codec {
	return &Codec{relay: relay}
}

// WriteResponse marshals response, byte slice or error to remote party.
func (c *Codec) WriteResponse(r *rpc.Response, body interface{}) error {
	const op = errors.Op("codec: write response")
	frame := NewFrame()

	// SEQ_ID + METHOD_NAME_LEN
	frame.WriteOptions(uint32(r.Seq), uint32(len(r.ServiceMethod)))
	// Write protocol version
	frame.WriteVersion(VERSION_1)

	// load and delete associated codec to not waste memory
	// because we write it to the frame and don't need more information about it
	// as for go.14, Load and Delete are separate methods
	codec, ok := c.codec.Load(r.Seq)
	if !ok {
		// fallback codec
		frame.WriteFlags(CODEC_GOB)
	} else {
		frame.WriteFlags(codec.(FrameFlag))
	}

	// delete the key
	c.codec.Delete(r.Seq)

	// initialize buffer
	buf := new(bytes.Buffer)
	// writeServiceMethod to the buffer
	buf.WriteString(r.ServiceMethod)

	// if error returned, we sending it via relay and return error from WriteResponse
	if r.Error != "" {
		// Append error flag
		return c.handleError(r, frame, buf, errors.Str(r.Error))
	}

	// read flag previously written
	// TODO might be better to save it to local variable
	flags := frame.ReadFlags()

	switch {
	case flags&byte(CODEC_JSON) != 0:
		err := encodeJSON(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_MSGPACK) != 0:
		err := encodeMsgPack(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_RAW) != 0:
		err := encodeRaw(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	case flags&byte(CODEC_GOB) != 0:
		err := encodeGob(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf, err)
		}
		// send buffer
		return c.sendBuf(frame, buf)
	default:
		return c.handleError(r, frame, buf, errors.E(op, errors.Str("unknown codec")))
	}
}

func (c *Codec) sendBuf(frame *Frame, buf *bytes.Buffer) error {
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	return c.relay.Send(frame)
}

func (c *Codec) handleError(r *rpc.Response, frame *Frame, buf *bytes.Buffer, err error) error {
	// just to be sure, remove all data from buffer and write new
	buf.Truncate(0)
	// write all possible errors
	err = multierr.Append(err, errors.Str(r.Error))
	buf.WriteString(r.ServiceMethod)

	const op = errors.Op("handle codec error")
	frame.WriteFlags(ERROR)
	// error should be here
	if err != nil {
		buf.WriteString(err.Error())
	}
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	_ = c.relay.Send(frame)
	return errors.E(op, errors.Str(r.Error))
}

// ReadRequestHeader receives
func (c *Codec) ReadRequestHeader(r *rpc.Request) error {
	const op = errors.Op("codec: read request header")
	frame := NewFrame()
	err := c.relay.Receive(frame)
	if err != nil {
		return err
	}

	// opts[0] sequence ID
	// opts[1] service method name offset from payload in bytes
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}

	r.Seq = uint64(opts[0])
	r.ServiceMethod = string(frame.Payload()[:opts[1]])
	c.frame = frame
	return c.storeCodec(r, frame.ReadFlags())
}

func (c *Codec) storeCodec(r *rpc.Request, flag byte) error {
	switch {
	case flag&byte(CODEC_JSON) != 0:
		c.codec.Store(r.Seq, CODEC_JSON)
	case flag&byte(CODEC_RAW) != 0:
		c.codec.Store(r.Seq, CODEC_RAW)
	case flag&byte(CODEC_MSGPACK) != 0:
		c.codec.Store(r.Seq, CODEC_MSGPACK)
	case flag&byte(CODEC_GOB) != 0:
		c.codec.Store(r.Seq, CODEC_GOB)
	default:
		c.codec.Store(r.Seq, CODEC_GOB)
	}

	return nil
}

// ReadRequestBody fetches prefixed body data and automatically unmarshal it as json. RawBody flag will populate
// []byte lice argument for rpc method.
func (c *Codec) ReadRequestBody(out interface{}) error {
	const op = errors.Op("codec read request body")
	if out == nil {
		return nil
	}

	defer func() {
		c.frame = nil
	}()

	flags := c.frame.ReadFlags()

	switch {
	case flags&byte(CODEC_JSON) != 0:
		return decodeJSON(out, c.frame)
	case flags&byte(CODEC_GOB) != 0:
		return decodeGob(out, c.frame)
	case flags&byte(CODEC_RAW) != 0:
		return decodeRaw(out, c.frame)
	case flags&byte(CODEC_MSGPACK) != 0:
		return decodeMsgPack(out, c.frame)
	default:
		return errors.E(op, errors.Str("unknown decoder used in frame"))
	}
}

// Close underlying socket.
func (c *Codec) Close() error {
	if c.closed {
		return nil
	}

	c.closed = true
	return c.relay.Close()
}
<|endoftext|>


<commit_before>package piccolo
<commit_msg>finish color.go<commit_after>package piccolo

import (
	"bytes"
	"fmt"
)

type Color struct {
	before, after string
}

var colorMap = map[string]Color{
	// styles
	"bold":          {"\x1B[1m", "\x1B[22m"},
	"italic":        {"\x1B[3m", "\x1B[23m"},
	"underline":     {"\x1B[4m", "\x1B[24m"},
	"inverse":       {"\x1B[7m", "\x1B[27m"},
	"strikethrough": {"\x1B[9m", "\x1B[29m"},
	// grayscale
	"white": {"\x1B[37m", "\x1B[39m"},
	"grey":  {"\x1B[90m", "\x1B[39m"},
	"black": {"\x1B[30m", "\x1B[39m"},
	// colors
	"blue":    {"\x1B[34m", "\x1B[39m"},
	"cyan":    {"\x1B[36m", "\x1B[39m"},
	"green":   {"\x1B[32m", "\x1B[39m"},
	"magenta": {"\x1B[35m", "\x1B[39m"},
	"red":     {"\x1B[31m", "\x1B[39m"},
	"yellow":  {"\x1B[33m", "\x1B[39m"},
	// background colors
	"whiteBG":   {"\x1B[47m", "\x1B[49m"},
	"greyBG":    {"\x1B[49;5;8m", "\x1B[49m"},
	"blackBG":   {"\x1B[40m", "\x1B[49m"},
	"blueBG":    {"\x1B[44m", "\x1B[49m"},
	"cyanBG":    {"\x1B[46m", "\x1B[49m"},
	"greenBG":   {"\x1B[42m", "\x1B[49m"},
	"magentaBG": {"\x1B[45m", "\x1B[49m"},
	"redBG":     {"\x1B[41m", "\x1B[49m"},
	"yellowBG":  {"\x1B[43m", "\x1B[49m"},
}

func SetTextColor(color string, text string) string {
	c, ok := colorMap[color]
	if !ok {
		panic(fmt.Sprintf("%s not found", color))
	}
	// fastest concat string
	var buffer bytes.Buffer
	buffer.WriteString(c.before)
	buffer.WriteString(text)
	buffer.WriteString(c.after)
	return buffer.String()
}
<|endoftext|>


<commit_before>package notary

import (
	"time"
)

// application wide constants
const (
	// MaxDownloadSize is the maximum size we'll download for metadata if no limit is given
	MaxDownloadSize int64 = 100 << 20
	// MaxTimestampSize is the maximum size of timestamp metadata - 1MiB.
	MaxTimestampSize int64 = 1 << 20
	// MinRSABitSize is the minimum bit size for RSA keys allowed in notary
	MinRSABitSize = 2048
	// MinThreshold requires a minimum of one threshold for roles; currently we do not support a higher threshold
	MinThreshold = 1
	// PrivKeyPerms are the file permissions to use when writing private keys to disk
	PrivKeyPerms = 0700
	// PubCertPerms are the file permissions to use when writing public certificates to disk
	PubCertPerms = 0755
	// Sha256HexSize is how big a Sha256 hex is in number of characters
	Sha256HexSize = 64
	// Sha512HexSize is how big a Sha512 hex is in number of characters
	Sha512HexSize = 128
	// SHA256 is the name of SHA256 hash algorithm
	SHA256 = "sha256"
	// SHA512 is the name of SHA512 hash algorithm
	SHA512 = "sha512"
	// TrustedCertsDir is the directory, under the notary repo base directory, where trusted certs are stored
	TrustedCertsDir = "trusted_certificates"
	// PrivDir is the directory, under the notary repo base directory, where private keys are stored
	PrivDir = "private"
	// RootKeysSubdir is the subdirectory under PrivDir where root private keys are stored
	RootKeysSubdir = "root_keys"
	// NonRootKeysSubdir is the subdirectory under PrivDir where non-root private keys are stored
	NonRootKeysSubdir = "tuf_keys"

	// Day is a duration of one day
	Day  = 24 * time.Hour
	Year = 365 * Day

	// NotaryRootExpiry is the duration representing the expiry time of the Root role
	NotaryRootExpiry      = 10 * Year
	NotaryTargetsExpiry   = 3 * Year
	NotarySnapshotExpiry  = 3 * Year
	NotaryTimestampExpiry = 14 * Day

	ConsistentMetadataCacheMaxAge = 30 * Day
	CurrentMetadataCacheMaxAge    = 5 * time.Minute
	// CacheMaxAgeLimit is the generally recommended maximum age for Cache-Control headers
	// (one year, in seconds, since one year is forever in terms of internet
	// content)
	CacheMaxAgeLimit = 1 * Year

	MySQLBackend     = "mysql"
	MemoryBackend    = "memory"
	SQLiteBackend    = "sqlite3"
	RethinkDBBackend = "rethinkdb"
)

// NotaryDefaultExpiries is the construct used to configure the default expiry times of
// the various role files.
var NotaryDefaultExpiries = map[string]time.Duration{
	"root":      NotaryRootExpiry,
	"targets":   NotaryTargetsExpiry,
	"snapshot":  NotarySnapshotExpiry,
	"timestamp": NotaryTimestampExpiry,
}
<commit_msg>[PATCH 3/8] Add two constants to describe the signals we would like to handle<commit_after>package notary

import (
	"os"
	"syscall"
	"time"
)

// application wide constants
const (
	// MaxDownloadSize is the maximum size we'll download for metadata if no limit is given
	MaxDownloadSize int64 = 100 << 20
	// MaxTimestampSize is the maximum size of timestamp metadata - 1MiB.
	MaxTimestampSize int64 = 1 << 20
	// MinRSABitSize is the minimum bit size for RSA keys allowed in notary
	MinRSABitSize = 2048
	// MinThreshold requires a minimum of one threshold for roles; currently we do not support a higher threshold
	MinThreshold = 1
	// PrivKeyPerms are the file permissions to use when writing private keys to disk
	PrivKeyPerms = 0700
	// PubCertPerms are the file permissions to use when writing public certificates to disk
	PubCertPerms = 0755
	// Sha256HexSize is how big a Sha256 hex is in number of characters
	Sha256HexSize = 64
	// Sha512HexSize is how big a Sha512 hex is in number of characters
	Sha512HexSize = 128
	// SHA256 is the name of SHA256 hash algorithm
	SHA256 = "sha256"
	// SHA512 is the name of SHA512 hash algorithm
	SHA512 = "sha512"
	// TrustedCertsDir is the directory, under the notary repo base directory, where trusted certs are stored
	TrustedCertsDir = "trusted_certificates"
	// PrivDir is the directory, under the notary repo base directory, where private keys are stored
	PrivDir = "private"
	// RootKeysSubdir is the subdirectory under PrivDir where root private keys are stored
	RootKeysSubdir = "root_keys"
	// NonRootKeysSubdir is the subdirectory under PrivDir where non-root private keys are stored
	NonRootKeysSubdir = "tuf_keys"

	// Day is a duration of one day
	Day  = 24 * time.Hour
	Year = 365 * Day

	// NotaryRootExpiry is the duration representing the expiry time of the Root role
	NotaryRootExpiry      = 10 * Year
	NotaryTargetsExpiry   = 3 * Year
	NotarySnapshotExpiry  = 3 * Year
	NotaryTimestampExpiry = 14 * Day

	ConsistentMetadataCacheMaxAge = 30 * Day
	CurrentMetadataCacheMaxAge    = 5 * time.Minute
	// CacheMaxAgeLimit is the generally recommended maximum age for Cache-Control headers
	// (one year, in seconds, since one year is forever in terms of internet
	// content)
	CacheMaxAgeLimit = 1 * Year

	MySQLBackend     = "mysql"
	MemoryBackend    = "memory"
	SQLiteBackend    = "sqlite3"
	RethinkDBBackend = "rethinkdb"
)

// NotaryDefaultExpiries is the construct used to configure the default expiry times of
// the various role files.
var NotaryDefaultExpiries = map[string]time.Duration{
	"root":      NotaryRootExpiry,
	"targets":   NotaryTargetsExpiry,
	"snapshot":  NotarySnapshotExpiry,
	"timestamp": NotaryTimestampExpiry,
}

// NotarySupportedSignals contains the signals we would like to capture:
// - SIGUSR1, indicates a increment of the log level.
// - SIGUSR2, indicates a decrement of the log level.
var NotarySupportedSignals = []os.Signal{syscall.SIGUSR1, syscall.SIGUSR2}
<|endoftext|>