<commit_before>package metric

import (
	"github.com/concourse/concourse/atc/db/lock"
	"strconv"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc/db"
)

var Databases []db.Conn
var DatabaseQueries = Meter(0)

var ContainersCreated = Meter(0)
var VolumesCreated = Meter(0)

var FailedContainers = Meter(0)
var FailedVolumes = Meter(0)

var ContainersDeleted = Meter(0)
var VolumesDeleted = Meter(0)

type SchedulingFullDuration struct {
	PipelineName string
	Duration     time.Duration
}

func (event SchedulingFullDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("full-scheduling-duration"),
		Event{
			Name:  "scheduling: full duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
			},
		},
	)
}

type SchedulingLoadVersionsDuration struct {
	PipelineName string
	Duration     time.Duration
}

func (event SchedulingLoadVersionsDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("loading-versions-duration"),
		Event{
			Name:  "scheduling: loading versions duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
			},
		},
	)
}

type SchedulingJobDuration struct {
	PipelineName string
	JobName      string
	Duration     time.Duration
}

func (event SchedulingJobDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("job-scheduling-duration"),
		Event{
			Name:  "scheduling: job duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
				"job":      event.JobName,
			},
		},
	)
}

type WorkerContainers struct {
	WorkerName string
	Platform   string
	Containers int
}

func (event WorkerContainers) Emit(logger lager.Logger) {
	emit(
		logger.Session("worker-containers"),
		Event{
			Name:  "worker containers",
			Value: event.Containers,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker":   event.WorkerName,
				"platform": event.Platform,
			},
		},
	)
}

type WorkerVolumes struct {
	WorkerName string
	Platform   string
	Volumes    int
}

func (event WorkerVolumes) Emit(logger lager.Logger) {
	emit(
		logger.Session("worker-volumes"),
		Event{
			Name:  "worker volumes",
			Value: event.Volumes,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker":   event.WorkerName,
				"platform": event.Platform,
			},
		},
	)
}

type VolumesToBeGarbageCollected struct {
	Volumes int
}

func (event VolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-orphaned-volumes-for-deletion"),
		Event{
			Name:       "orphaned volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatingContainersToBeGarbageCollected struct {
	Containers int
}

func (event CreatingContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-creating-containers-for-deletion"),
		Event{
			Name:       "creating containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatedContainersToBeGarbageCollected struct {
	Containers int
}

func (event CreatedContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-created-ccontainers-for-deletion"),
		Event{
			Name:       "created containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type DestroyingContainersToBeGarbageCollected struct {
	Containers int
}

func (event DestroyingContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-destroying-containers-for-deletion"),
		Event{
			Name:       "destroying containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type FailedContainersToBeGarbageCollected struct {
	Containers int
}

func (event FailedContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-failed-containers-for-deletion"),
		Event{
			Name:       "failed containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatedVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event CreatedVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-created-volumes-for-deletion"),
		Event{
			Name:       "created volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type DestroyingVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event DestroyingVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-destroying-volumes-for-deletion"),
		Event{
			Name:       "destroying volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type FailedVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event FailedVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-failed-volumes-for-deletion"),
		Event{
			Name:       "failed volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type GarbageCollectionContainerCollectorJobDropped struct {
	WorkerName string
}

func (event GarbageCollectionContainerCollectorJobDropped) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-container-collector-dropped"),
		Event{
			Name:  "GC container collector job dropped",
			Value: 1,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker": event.WorkerName,
			},
		},
	)
}

type BuildStarted struct {
	PipelineName string
	JobName      string
	BuildName    string
	BuildID      int
	TeamName     string
}

func (event BuildStarted) Emit(logger lager.Logger) {
	emit(
		logger.Session("build-started"),
		Event{
			Name:  "build started",
			Value: event.BuildID,
			State: EventStateOK,
			Attributes: map[string]string{
				"pipeline":   event.PipelineName,
				"job":        event.JobName,
				"build_name": event.BuildName,
				"build_id":   strconv.Itoa(event.BuildID),
				"team_name":  event.TeamName,
			},
		},
	)
}

type BuildFinished struct {
	PipelineName  string
	JobName       string
	BuildName     string
	BuildID       int
	BuildStatus   db.BuildStatus
	BuildDuration time.Duration
	TeamName      string
}

func (event BuildFinished) Emit(logger lager.Logger) {
	emit(
		logger.Session("build-finished"),
		Event{
			Name:  "build finished",
			Value: ms(event.BuildDuration),
			State: EventStateOK,
			Attributes: map[string]string{
				"pipeline":     event.PipelineName,
				"job":          event.JobName,
				"build_name":   event.BuildName,
				"build_id":     strconv.Itoa(event.BuildID),
				"build_status": string(event.BuildStatus),
				"team_name":    event.TeamName,
			},
		},
	)
}

func ms(duration time.Duration) float64 {
	return float64(duration) / 1000000
}

type HTTPResponseTime struct {
	Route    string
	Path     string
	Method   string
	Duration time.Duration
}

func (event HTTPResponseTime) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > 100*time.Millisecond {
		state = EventStateWarning
	}

	if event.Duration > 1*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("http-response-time"),
		Event{
			Name:  "http response time",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"route":  event.Route,
				"path":   event.Path,
				"method": event.Method,
			},
		},
	)
}

type ResourceCheck struct {
	PipelineName string
	ResourceName string
	TeamName     string
	Success      bool
}

func (event ResourceCheck) Emit(logger lager.Logger) {
	state := EventStateOK
	if !event.Success {
		state = EventStateWarning
	}
	emit(
		logger.Session("resource-check"),
		Event{
			Name:  "resource checked",
			Value: 1,
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
				"resource": event.ResourceName,
				"team":     event.TeamName,
			},
		},
	)
}

var lockTypeNames = map[int]string{
	lock.LockTypeResourceConfigChecking: "ResourceConfigChecking",
	lock.LockTypeBuildTracking:          "BuildTracking",
	lock.LockTypePipelineScheduling:     "PipelineScheduling",
	lock.LockTypeBatch:                  "Batch",
	lock.LockTypeVolumeCreating:         "VolumeCreating",
	lock.LockTypeContainerCreating:      "ContainerCreating",
	lock.LockTypeDatabaseMigration:      "DatabaseMigration",
}

type LockAcquired struct {
	LockType string
	ObjectID int
}

func (event LockAcquired) Emit(logger lager.Logger) {
	emit(
		logger.Session("lock-acquired"),
		Event{
			Name:  "lock held",
			Value: 1,
			State: EventStateOK,
			Attributes: map[string]string{
				"type":      event.LockType,
				"object_id": strconv.Itoa(event.ObjectID),
			},
		},
	)
}

type LockReleased struct {
	LockType string
	ObjectID int
}

func (event LockReleased) Emit(logger lager.Logger) {
	emit(
		logger.Session("lock-released"),
		Event{
			Name:  "lock held",
			Value: 0,
			State: EventStateOK,
			Attributes: map[string]string{
				"type":      event.LockType,
				"object_id": strconv.Itoa(event.ObjectID),
			},
		},
	)
}

func LogLockAcquired(logger lager.Logger, lockID lock.LockID) {
	logger.Debug("acquired")

	if len(lockID) > 0 {
		if lockType, ok := lockTypeNames[lockID[0]]; ok {
			LockAcquired{
				LockType: lockType,
				ObjectID: lockID[1],
			}.Emit(logger)
		}
	}
}

func LogLockReleased(logger lager.Logger, lockID lock.LockID) {
	logger.Debug("released")

	if len(lockID) > 0 {
		if lockType, ok := lockTypeNames[lockID[0]]; ok {
			LockReleased{
				LockType: lockType,
				ObjectID: lockID[1],
			}.Emit(logger)
		}
	}
}

type WorkersState struct {
	WorkerStateByName map[string]db.WorkerState
}

func (event WorkersState) Emit(logger lager.Logger) {
	var eventState EventState

	for workerName, workerState := range event.WorkerStateByName {
		numericState := 0

		eventState = EventStateOK

		if workerState == db.WorkerStateStalled {
			eventState = EventStateWarning
		}

		switch workerState {
		case db.WorkerStateStalled:
			numericState = 1
		case db.WorkerStateRetiring:
			numericState = 2
		case db.WorkerStateLanded:
			numericState = 3
		case db.WorkerStateLanding:
			numericState = 4
		case db.WorkerStateRunning:
			numericState = 5
		}

		emit(
			logger.Session("worker-state"),
			Event{
				Name:  "worker state",
				Value: numericState,
				State: eventState,
				Attributes: map[string]string{
					"name":         workerName,
					"worker_state": string(workerState),
				},
			},
		)
	}
}<commit_msg>atc : guard against locks with no object id<commit_after>package metric

import (
	"github.com/concourse/concourse/atc/db/lock"
	"strconv"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc/db"
)

var Databases []db.Conn
var DatabaseQueries = Meter(0)

var ContainersCreated = Meter(0)
var VolumesCreated = Meter(0)

var FailedContainers = Meter(0)
var FailedVolumes = Meter(0)

var ContainersDeleted = Meter(0)
var VolumesDeleted = Meter(0)

type SchedulingFullDuration struct {
	PipelineName string
	Duration     time.Duration
}

func (event SchedulingFullDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("full-scheduling-duration"),
		Event{
			Name:  "scheduling: full duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
			},
		},
	)
}

type SchedulingLoadVersionsDuration struct {
	PipelineName string
	Duration     time.Duration
}

func (event SchedulingLoadVersionsDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("loading-versions-duration"),
		Event{
			Name:  "scheduling: loading versions duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
			},
		},
	)
}

type SchedulingJobDuration struct {
	PipelineName string
	JobName      string
	Duration     time.Duration
}

func (event SchedulingJobDuration) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > time.Second {
		state = EventStateWarning
	}

	if event.Duration > 5*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("job-scheduling-duration"),
		Event{
			Name:  "scheduling: job duration (ms)",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
				"job":      event.JobName,
			},
		},
	)
}

type WorkerContainers struct {
	WorkerName string
	Platform   string
	Containers int
}

func (event WorkerContainers) Emit(logger lager.Logger) {
	emit(
		logger.Session("worker-containers"),
		Event{
			Name:  "worker containers",
			Value: event.Containers,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker":   event.WorkerName,
				"platform": event.Platform,
			},
		},
	)
}

type WorkerVolumes struct {
	WorkerName string
	Platform   string
	Volumes    int
}

func (event WorkerVolumes) Emit(logger lager.Logger) {
	emit(
		logger.Session("worker-volumes"),
		Event{
			Name:  "worker volumes",
			Value: event.Volumes,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker":   event.WorkerName,
				"platform": event.Platform,
			},
		},
	)
}

type VolumesToBeGarbageCollected struct {
	Volumes int
}

func (event VolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-orphaned-volumes-for-deletion"),
		Event{
			Name:       "orphaned volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatingContainersToBeGarbageCollected struct {
	Containers int
}

func (event CreatingContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-creating-containers-for-deletion"),
		Event{
			Name:       "creating containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatedContainersToBeGarbageCollected struct {
	Containers int
}

func (event CreatedContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-created-ccontainers-for-deletion"),
		Event{
			Name:       "created containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type DestroyingContainersToBeGarbageCollected struct {
	Containers int
}

func (event DestroyingContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-destroying-containers-for-deletion"),
		Event{
			Name:       "destroying containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type FailedContainersToBeGarbageCollected struct {
	Containers int
}

func (event FailedContainersToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-failed-containers-for-deletion"),
		Event{
			Name:       "failed containers to be garbage collected",
			Value:      event.Containers,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type CreatedVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event CreatedVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-created-volumes-for-deletion"),
		Event{
			Name:       "created volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type DestroyingVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event DestroyingVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-destroying-volumes-for-deletion"),
		Event{
			Name:       "destroying volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type FailedVolumesToBeGarbageCollected struct {
	Volumes int
}

func (event FailedVolumesToBeGarbageCollected) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-found-failed-volumes-for-deletion"),
		Event{
			Name:       "failed volumes to be garbage collected",
			Value:      event.Volumes,
			State:      EventStateOK,
			Attributes: map[string]string{},
		},
	)
}

type GarbageCollectionContainerCollectorJobDropped struct {
	WorkerName string
}

func (event GarbageCollectionContainerCollectorJobDropped) Emit(logger lager.Logger) {
	emit(
		logger.Session("gc-container-collector-dropped"),
		Event{
			Name:  "GC container collector job dropped",
			Value: 1,
			State: EventStateOK,
			Attributes: map[string]string{
				"worker": event.WorkerName,
			},
		},
	)
}

type BuildStarted struct {
	PipelineName string
	JobName      string
	BuildName    string
	BuildID      int
	TeamName     string
}

func (event BuildStarted) Emit(logger lager.Logger) {
	emit(
		logger.Session("build-started"),
		Event{
			Name:  "build started",
			Value: event.BuildID,
			State: EventStateOK,
			Attributes: map[string]string{
				"pipeline":   event.PipelineName,
				"job":        event.JobName,
				"build_name": event.BuildName,
				"build_id":   strconv.Itoa(event.BuildID),
				"team_name":  event.TeamName,
			},
		},
	)
}

type BuildFinished struct {
	PipelineName  string
	JobName       string
	BuildName     string
	BuildID       int
	BuildStatus   db.BuildStatus
	BuildDuration time.Duration
	TeamName      string
}

func (event BuildFinished) Emit(logger lager.Logger) {
	emit(
		logger.Session("build-finished"),
		Event{
			Name:  "build finished",
			Value: ms(event.BuildDuration),
			State: EventStateOK,
			Attributes: map[string]string{
				"pipeline":     event.PipelineName,
				"job":          event.JobName,
				"build_name":   event.BuildName,
				"build_id":     strconv.Itoa(event.BuildID),
				"build_status": string(event.BuildStatus),
				"team_name":    event.TeamName,
			},
		},
	)
}

func ms(duration time.Duration) float64 {
	return float64(duration) / 1000000
}

type HTTPResponseTime struct {
	Route    string
	Path     string
	Method   string
	Duration time.Duration
}

func (event HTTPResponseTime) Emit(logger lager.Logger) {
	state := EventStateOK

	if event.Duration > 100*time.Millisecond {
		state = EventStateWarning
	}

	if event.Duration > 1*time.Second {
		state = EventStateCritical
	}

	emit(
		logger.Session("http-response-time"),
		Event{
			Name:  "http response time",
			Value: ms(event.Duration),
			State: state,
			Attributes: map[string]string{
				"route":  event.Route,
				"path":   event.Path,
				"method": event.Method,
			},
		},
	)
}

type ResourceCheck struct {
	PipelineName string
	ResourceName string
	TeamName     string
	Success      bool
}

func (event ResourceCheck) Emit(logger lager.Logger) {
	state := EventStateOK
	if !event.Success {
		state = EventStateWarning
	}
	emit(
		logger.Session("resource-check"),
		Event{
			Name:  "resource checked",
			Value: 1,
			State: state,
			Attributes: map[string]string{
				"pipeline": event.PipelineName,
				"resource": event.ResourceName,
				"team":     event.TeamName,
			},
		},
	)
}

var lockTypeNames = map[int]string{
	lock.LockTypeResourceConfigChecking: "ResourceConfigChecking",
	lock.LockTypeBuildTracking:          "BuildTracking",
	lock.LockTypePipelineScheduling:     "PipelineScheduling",
	lock.LockTypeBatch:                  "Batch",
	lock.LockTypeVolumeCreating:         "VolumeCreating",
	lock.LockTypeContainerCreating:      "ContainerCreating",
	lock.LockTypeDatabaseMigration:      "DatabaseMigration",
}

type LockAcquired struct {
	LockType string
	ObjectID int
}

func (event LockAcquired) Emit(logger lager.Logger) {
	emit(
		logger.Session("lock-acquired"),
		Event{
			Name:  "lock held",
			Value: 1,
			State: EventStateOK,
			Attributes: map[string]string{
				"type":      event.LockType,
				"object_id": strconv.Itoa(event.ObjectID),
			},
		},
	)
}

type LockReleased struct {
	LockType string
	ObjectID int
}

func (event LockReleased) Emit(logger lager.Logger) {
	emit(
		logger.Session("lock-released"),
		Event{
			Name:  "lock held",
			Value: 0,
			State: EventStateOK,
			Attributes: map[string]string{
				"type":      event.LockType,
				"object_id": strconv.Itoa(event.ObjectID),
			},
		},
	)
}

func LogLockAcquired(logger lager.Logger, lockID lock.LockID) {
	logger.Debug("acquired")

	if len(lockID) > 0 {
		if lockType, ok := lockTypeNames[lockID[0]]; ok {
			var objectID int
			if len(lockID) > 1 {
				objectID = lockID[1];
			}

			LockAcquired{
				LockType:lockType,
				ObjectID: objectID,
			}.Emit(logger)
		}
	}
}

func LogLockReleased(logger lager.Logger, lockID lock.LockID) {
	logger.Debug("released")

	if len(lockID) > 0 {
		if lockType, ok := lockTypeNames[lockID[0]]; ok {
			var objectID int
			if len(lockID) > 1 {
				objectID = lockID[1];
			}
			LockReleased{
				LockType: lockType,
				ObjectID: objectID,
			}.Emit(logger)
		}
	}
}

type WorkersState struct {
	WorkerStateByName map[string]db.WorkerState
}

func (event WorkersState) Emit(logger lager.Logger) {
	var eventState EventState

	for workerName, workerState := range event.WorkerStateByName {
		numericState := 0

		eventState = EventStateOK

		if workerState == db.WorkerStateStalled {
			eventState = EventStateWarning
		}

		switch workerState {
		case db.WorkerStateStalled:
			numericState = 1
		case db.WorkerStateRetiring:
			numericState = 2
		case db.WorkerStateLanded:
			numericState = 3
		case db.WorkerStateLanding:
			numericState = 4
		case db.WorkerStateRunning:
			numericState = 5
		}

		emit(
			logger.Session("worker-state"),
			Event{
				Name:  "worker state",
				Value: numericState,
				State: eventState,
				Attributes: map[string]string{
					"name":         workerName,
					"worker_state": string(workerState),
				},
			},
		)
	}
}<|endoftext|>


<commit_before>package be

import (
	"testing"

	. "github.com/chai2010/assert"
	"github.com/coocood/qbs"
)

func TestUserAPI(t *testing.T) {
	CreateAndInitDB()
	db, err := qbs.GetQbs()
	AssertNil(t, err)
	defer func() {
		WipeDB()
		db.Close()
	}()

	testApi, err := NewTestAPI()
	AssertNil(t, err)
	defer testApi.Stop()

	users, err := FindUsers(0, 100, db)
	AssertNil(t, err)
	AssertEqual(t, 0, len(users), "Need to have 0 users when starting")

	user, err := CreateUser("adrian@monk.example.com", "Adrian", "Monk", false, db)
	AssertNil(t, err)
	_, err = CreatePassword("1234", user.Id, db)
	AssertNil(t, err)
	staff, err := CreateUser("sherona@monk.example.com", "Sherona", "Smith", true, db)
	AssertNil(t, err)
	_, err = CreatePassword("1234", staff.Id, db)
	AssertNil(t, err)

	Assert403(t, "GET", testApi.URL()+"/user/")
	Assert403(t, "GET", testApi.URL()+"/user/"+user.UUID)

	userClient, err := NewClient(testApi.URL())
	AssertNil(t, err)
	err = userClient.Authenticate(user.Email, "")
	AssertNotNil(t, err, "Should have failed with empty password")
	err = userClient.Authenticate("", "1234")
	AssertNotNil(t, err, "Should have failed with empty email")
	err = userClient.Authenticate("", "")
	AssertNotNil(t, err, "Should have failed with empty login info")
	err = userClient.Authenticate(user.Email, "4321")
	AssertNotNil(t, err, "Should have failed with incorrect password")
	err = userClient.Authenticate(user.Email, "1234")
	AssertNil(t, err, "Should have authenticated with proper email and username")

	user2 := new(User)
	err = userClient.GetJSON("/user/current", user2)
	AssertNil(t, err, "Error fetching current user")
	AssertEqual(t, user.Id, user2.Id)
	_, err = userClient.GetList("/user/")
	AssertNotNil(t, err, "Users API should be staff only")
	err = userClient.GetJSON("/user/"+user2.UUID, user2)
	AssertNotNil(t, err, "User API should be staff only")

	staffClient, err := NewClient(testApi.URL())
	AssertNil(t, err)
	err = staffClient.Authenticate(staff.Email, "1234")
	AssertNil(t, err)
	err = staffClient.GetJSON("/user/"+user2.UUID, user2)
	AssertNil(t, err, "API should be readable by staff")
	AssertEqual(t, user.Id, user2.Id)

	list, err := staffClient.GetList("/user/")
	AssertNil(t, err)
	arr := list.Objects.([]interface{})
	AssertEqual(t, 2, len(arr))

	// Test that staff can update a User
	staff2 := new(User)
	err = staffClient.GetJSON("/user/current", staff2)
	AssertNil(t, err)
	staff2.FirstName = "Pickles"
	staff2.LastName = "McGee"
	err = staffClient.UpdateUser(staff2)
	AssertNil(t, err)
	AssertEqual(t, staff2.FirstName, "Pickles")
	AssertEqual(t, staff2.LastName, "McGee")
	staff3 := new(User)
	err = staffClient.GetJSON("/user/current", staff3)
	AssertNil(t, err)
	AssertEqual(t, staff2.FirstName, staff3.FirstName)
	AssertEqual(t, staff2.LastName, staff3.LastName)
}

func TestUser(t *testing.T) {
	CreateAndInitDB()
	db, err := qbs.GetQbs()
	AssertNil(t, err)
	defer func() {
		WipeDB()
		db.Close()
	}()

	user, err := CreateUser("adrian@monk.example.com", "Adrian", "Monk", false, db)
	AssertNil(t, err)
	AssertNotEqual(t, user.UUID, "")

	_, err = FindUser("not-a-uuid", db)
	AssertNotNil(t, err)

	user2, err := FindUser(user.UUID, db)
	AssertNil(t, err)
	AssertEqual(t, user2.UUID, user.UUID)
	AssertEqual(t, user2.Email, user.Email)

	user2.Email = "crosby@bing.example.com"
	err = UpdateUser(user2, db)
	AssertNil(t, err)
	AssertEqual(t, user2.UUID, user.UUID)
	user3, err := FindUser(user2.UUID, db)
	AssertNil(t, err)
	AssertEqual(t, user2.Email, user3.Email)

	// TODO
	/*
		Test schema API
		Test versioning enforcement
	*/
}
<commit_msg>Furthered testing for user and schema resources.<commit_after>package be

import (
	"testing"

	. "github.com/chai2010/assert"
	"github.com/coocood/qbs"
)

func TestUserAPI(t *testing.T) {
	CreateAndInitDB()
	db, err := qbs.GetQbs()
	AssertNil(t, err)
	defer func() {
		WipeDB()
		db.Close()
	}()

	testApi, err := NewTestAPI()
	AssertNil(t, err)
	defer testApi.Stop()

	users, err := FindUsers(0, 100, db)
	AssertNil(t, err)
	AssertEqual(t, 0, len(users), "Need to have 0 users when starting")

	user, err := CreateUser("adrian@monk.example.com", "Adrian", "Monk", false, db)
	AssertNil(t, err)
	_, err = CreatePassword("1234", user.Id, db)
	AssertNil(t, err)
	staff, err := CreateUser("sherona@monk.example.com", "Sherona", "Smith", true, db)
	AssertNil(t, err)
	_, err = CreatePassword("1234", staff.Id, db)
	AssertNil(t, err)

	Assert403(t, "GET", testApi.URL()+"/user/")
	Assert403(t, "GET", testApi.URL()+"/user/"+user.UUID)

	userClient, err := NewClient(testApi.URL())
	AssertNil(t, err)
	err = userClient.Authenticate(user.Email, "")
	AssertNotNil(t, err, "Should have failed with empty password")
	err = userClient.Authenticate("", "1234")
	AssertNotNil(t, err, "Should have failed with empty email")
	err = userClient.Authenticate("", "")
	AssertNotNil(t, err, "Should have failed with empty login info")
	err = userClient.Authenticate(user.Email, "4321")
	AssertNotNil(t, err, "Should have failed with incorrect password")
	err = userClient.Authenticate(user.Email, "1234")
	AssertNil(t, err, "Should have authenticated with proper email and username")

	user2 := new(User)
	err = userClient.GetJSON("/user/current", user2)
	AssertNil(t, err, "Error fetching current user")
	AssertEqual(t, user.Id, user2.Id)
	_, err = userClient.GetList("/user/")
	AssertNotNil(t, err, "Users API should be staff only")
	err = userClient.GetJSON("/user/"+user2.UUID, user2)
	AssertNotNil(t, err, "User API should be staff only")

	staffClient, err := NewClient(testApi.URL())
	AssertNil(t, err)
	err = staffClient.Authenticate(staff.Email, "1234")
	AssertNil(t, err)
	err = staffClient.GetJSON("/user/"+user2.UUID, user2)
	AssertNil(t, err, "API should be readable by staff")
	AssertEqual(t, user.Id, user2.Id)

	list, err := staffClient.GetList("/user/")
	AssertNil(t, err)
	arr := list.Objects.([]interface{})
	AssertEqual(t, 2, len(arr))

	// Test that staff can update a User
	staff2 := new(User)
	err = staffClient.GetJSON("/user/current", staff2)
	AssertNil(t, err)
	staff2.FirstName = "Pickles"
	staff2.LastName = "McGee"
	err = staffClient.UpdateUser(staff2)
	AssertNil(t, err)
	AssertEqual(t, staff2.FirstName, "Pickles")
	AssertEqual(t, staff2.LastName, "McGee")
	staff3 := new(User)
	err = staffClient.GetJSON("/user/current", staff3)
	AssertNil(t, err)
	AssertEqual(t, staff2.FirstName, staff3.FirstName)
	AssertEqual(t, staff2.LastName, staff3.LastName)
}

func TestUser(t *testing.T) {
	CreateAndInitDB()
	db, err := qbs.GetQbs()
	AssertNil(t, err)
	defer func() {
		WipeDB()
		db.Close()
	}()

	user, err := CreateUser("adrian@monk.example.com", "Adrian", "Monk", false, db)
	AssertNil(t, err)
	AssertNotEqual(t, user.UUID, "")

	_, err = FindUser("not-a-uuid", db)
	AssertNotNil(t, err)

	user2, err := FindUser(user.UUID, db)
	AssertNil(t, err)
	AssertEqual(t, user2.UUID, user.UUID)
	AssertEqual(t, user2.Email, user.Email)

	user2.Email = "crosby@bing.example.com"
	err = UpdateUser(user2, db)
	AssertNil(t, err)
	AssertEqual(t, user2.UUID, user.UUID)
	user3, err := FindUser(user2.UUID, db)
	AssertNil(t, err)
	AssertEqual(t, user2.Email, user3.Email)

	// TODO
	/*
		Test versioning enforcement
	*/
}
<|endoftext|>


<commit_before>// Copyright 2021 The Ebiten 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 audio

import (
	"io"
	"runtime"
	"sync"
	"time"
)

// readerDriver represents a driver using io.ReadClosers.
type readerDriver interface {
	NewPlayer(io.Reader) readerDriverPlayer
	io.Closer
}

type readerDriverPlayer interface {
	Pause()
	Play()
	Reset()
	Volume() float64
	SetVolume(volume float64)
	io.Closer
}

type readerPlayerFactory struct {
	driver readerDriver
}

func newReaderPlayerFactory(sampleRate int) *readerPlayerFactory {
	return &readerPlayerFactory{
		driver: newReaderDriverImpl(sampleRate),
	}
	// TODO: Consider the hooks.
}

type readerPlayer struct {
	context *Context
	player  readerDriverPlayer
	src     *timeStream
	playing bool
	m       sync.Mutex
}

func (c *readerPlayerFactory) newPlayerImpl(context *Context, src io.Reader) (playerImpl, error) {
	s, err := newTimeStream(src, context.SampleRate())
	if err != nil {
		return nil, err
	}

	p := &readerPlayer{
		context: context,
		player:  c.driver.NewPlayer(src),
		src:     s,
	}
	runtime.SetFinalizer(p, (*readerPlayer).Close)
	return p, nil
}

func (p *readerPlayer) Play() {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.Play()
	p.playing = true
	p.context.addPlayer(p)
}

func (p *readerPlayer) Pause() {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.Pause()
	p.playing = false
}

func (p *readerPlayer) IsPlaying() bool {
	p.m.Lock()
	defer p.m.Unlock()

	return p.playing
}

func (p *readerPlayer) Volume() float64 {
	p.m.Lock()
	defer p.m.Unlock()

	return p.player.Volume()
}

func (p *readerPlayer) SetVolume(volume float64) {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.SetVolume(volume)
}

func (p *readerPlayer) Close() error {
	p.m.Lock()
	defer p.m.Unlock()

	runtime.SetFinalizer(p, nil)
	p.context.removePlayer(p)
	p.playing = false
	return p.player.Close()
}

func (p *readerPlayer) Current() time.Duration {
	p.m.Lock()
	defer p.m.Unlock()

	// TODO: Add a new function to readerDriverPlayer and use it.
	return p.src.Current()
}

func (p *readerPlayer) Rewind() error {
	return p.Seek(0)
}

func (p *readerPlayer) Seek(offset time.Duration) error {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.Reset()
	return p.src.Seek(offset)
}

func (p *readerPlayer) source() io.Reader {
	return p.src
}

type timeStream struct {
	r          io.Reader
	sampleRate int
	pos        int64
}

func newTimeStream(r io.Reader, sampleRate int) (*timeStream, error) {
	s := &timeStream{
		r:          r,
		sampleRate: sampleRate,
	}
	if seeker, ok := s.r.(io.Seeker); ok {
		// Get the current position of the source.
		pos, err := seeker.Seek(0, io.SeekCurrent)
		if err != nil {
			return nil, err
		}
		s.pos = pos
	}
	return s, nil
}

func (s *timeStream) Read(buf []byte) (int, error) {
	n, err := s.Read(buf)
	s.pos += int64(n)
	return n, err
}

func (s *timeStream) Seek(offset time.Duration) error {
	o := int64(offset) * bytesPerSample * int64(s.sampleRate) / int64(time.Second)

	// Align the byte position with the samples.
	o -= o % bytesPerSample
	o += s.pos % bytesPerSample

	seeker, ok := s.r.(io.Seeker)
	if !ok {
		panic("audio: the source must be io.Seeker when seeking but not")
	}
	pos, err := seeker.Seek(o, io.SeekStart)
	if err != nil {
		return err
	}

	s.pos = pos
	return nil
}

func (s *timeStream) Current() time.Duration {
	sample := s.pos / bytesPerSample
	return time.Duration(sample) * time.Second / time.Duration(s.sampleRate)
}
<commit_msg>audio: Seek the source first at readerPlayer<commit_after>// Copyright 2021 The Ebiten 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 audio

import (
	"io"
	"runtime"
	"sync"
	"time"
)

// readerDriver represents a driver using io.ReadClosers.
type readerDriver interface {
	NewPlayer(io.Reader) readerDriverPlayer
	io.Closer
}

type readerDriverPlayer interface {
	Pause()
	Play()
	Reset()
	Volume() float64
	SetVolume(volume float64)
	io.Closer
}

type readerPlayerFactory struct {
	driver readerDriver
}

func newReaderPlayerFactory(sampleRate int) *readerPlayerFactory {
	return &readerPlayerFactory{
		driver: newReaderDriverImpl(sampleRate),
	}
	// TODO: Consider the hooks.
}

type readerPlayer struct {
	context *Context
	player  readerDriverPlayer
	src     *timeStream
	playing bool
	m       sync.Mutex
}

func (c *readerPlayerFactory) newPlayerImpl(context *Context, src io.Reader) (playerImpl, error) {
	s, err := newTimeStream(src, context.SampleRate())
	if err != nil {
		return nil, err
	}

	p := &readerPlayer{
		context: context,
		player:  c.driver.NewPlayer(src),
		src:     s,
	}
	runtime.SetFinalizer(p, (*readerPlayer).Close)
	return p, nil
}

func (p *readerPlayer) Play() {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.Play()
	p.playing = true
	p.context.addPlayer(p)
}

func (p *readerPlayer) Pause() {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.Pause()
	p.playing = false
}

func (p *readerPlayer) IsPlaying() bool {
	p.m.Lock()
	defer p.m.Unlock()

	return p.playing
}

func (p *readerPlayer) Volume() float64 {
	p.m.Lock()
	defer p.m.Unlock()

	return p.player.Volume()
}

func (p *readerPlayer) SetVolume(volume float64) {
	p.m.Lock()
	defer p.m.Unlock()

	p.player.SetVolume(volume)
}

func (p *readerPlayer) Close() error {
	p.m.Lock()
	defer p.m.Unlock()

	runtime.SetFinalizer(p, nil)
	p.context.removePlayer(p)
	p.playing = false
	return p.player.Close()
}

func (p *readerPlayer) Current() time.Duration {
	p.m.Lock()
	defer p.m.Unlock()

	// TODO: Add a new function to readerDriverPlayer and use it.
	return p.src.Current()
}

func (p *readerPlayer) Rewind() error {
	return p.Seek(0)
}

func (p *readerPlayer) Seek(offset time.Duration) error {
	p.m.Lock()
	defer p.m.Unlock()

	if err := p.src.Seek(offset); err != nil {
		return err
	}
	p.player.Reset()
	return nil
}

func (p *readerPlayer) source() io.Reader {
	return p.src
}

type timeStream struct {
	r          io.Reader
	sampleRate int
	pos        int64
}

func newTimeStream(r io.Reader, sampleRate int) (*timeStream, error) {
	s := &timeStream{
		r:          r,
		sampleRate: sampleRate,
	}
	if seeker, ok := s.r.(io.Seeker); ok {
		// Get the current position of the source.
		pos, err := seeker.Seek(0, io.SeekCurrent)
		if err != nil {
			return nil, err
		}
		s.pos = pos
	}
	return s, nil
}

func (s *timeStream) Read(buf []byte) (int, error) {
	n, err := s.Read(buf)
	s.pos += int64(n)
	return n, err
}

func (s *timeStream) Seek(offset time.Duration) error {
	o := int64(offset) * bytesPerSample * int64(s.sampleRate) / int64(time.Second)

	// Align the byte position with the samples.
	o -= o % bytesPerSample
	o += s.pos % bytesPerSample

	seeker, ok := s.r.(io.Seeker)
	if !ok {
		panic("audio: the source must be io.Seeker when seeking but not")
	}
	pos, err := seeker.Seek(o, io.SeekStart)
	if err != nil {
		return err
	}

	s.pos = pos
	return nil
}

func (s *timeStream) Current() time.Duration {
	sample := s.pos / bytesPerSample
	return time.Duration(sample) * time.Second / time.Duration(s.sampleRate)
}
<|endoftext|>


<commit_before>// The goapp command implements a simple App Engine app to demonstrate how to
// use the Service Control API v2 for admission control. For more information,
// see https://cloud.google.com/service-infrastructure/docs/admission-control.
package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	// WARNING:`go get google.golang.org/api/servicecontrol/v2" may take
	// 30 minutes or longer, depending on your network speed.
	"google.golang.org/api/servicecontrol/v2"
)

// Check calls Service Control API v2 for admission control.
// Name specifies the target resource name. Permission specifies
// the required permission on the target resource. Received
// specifies the timestamp when the request is received.
func check(w http.ResponseWriter, r *http.Request, name string, permission string, received time.Time, client *servicecontrol.Service) (string, error) {
	// Construct CheckRequest from the incoming HTTP request.
	// The code assumes the incoming request processed by App Engine ingress.
	req := &servicecontrol.CheckRequest{
		ServiceConfigId: "latest",
		Attributes: &servicecontrol.AttributeContext{
			Origin: &servicecontrol.Peer{
				Ip: r.Header.Get("x-appengine-user-ip"),
			},
			Api: &servicecontrol.Api{
				Service:   "endpointsapis.appspot.com",
				Operation: "google.example.endpointsapis.v1.Workspaces.GetWorkspace",
				Version:   "v1",
				Protocol:  r.Header.Get("x-forwarded-proto"),
			},
			Request: &servicecontrol.Request{
				Id:     r.Header.Get("x-appengine-request-log-id"),
				Time:   received.UTC().Format(time.RFC3339),
				Method: r.Method,
				Scheme: r.Header.Get("x-forwarded-proto"),
				Host:   r.Host,
				Path:   r.URL.Path,
				Headers: map[string]string{
					"authorization": r.Header.Get("authorization"),
					"user-agent":    r.Header.Get("user-agent"),
					"origin":        r.Header.Get("origin"),
					"referer":       r.Header.Get("referer"),
				},
			},
			Resource: &servicecontrol.Resource{
				Name: name,
			},
		},
		Resources: []*servicecontrol.ResourceInfo{
			{
				Name:       name,
				Type:       "endpointsapis.appspot.com/Workspace",
				Permission: permission,
			},
		},
	}
	resp, err := client.Services.Check("endpointsapis.appspot.com", req).Do()
	if err != nil {
		return "", err
	}
	json, err := resp.MarshalJSON()
	if err != nil {
		return "", err
	}
	return string(json), nil
}

// Report calls Service Control API v2 for telemetry reporting.
// Name specifies the target resource name. ResponseCode specifies
// the response code returned to user. Received specifies the
// timestamp when the request is received.
func report(w http.ResponseWriter, r *http.Request, name string, responseCode int64, received time.Time, client *servicecontrol.Service) (string, error) {
	// Construct ReportRequest from the incoming HTTP request.
	// The code assumes the incoming request processed by App Engine ingress.
	req := &servicecontrol.ReportRequest{
		ServiceConfigId: "latest",
		Operations: []*servicecontrol.AttributeContext{
			{
				Api: &servicecontrol.Api{
					Service:   "endpointsapis.appspot.com",
					Operation: "google.example.endpointsapis.v1.Workspaces.GetWorkspace",
					Version:   "v1",
					Protocol:  r.Header.Get("x-forwarded-proto"),
				},
				Request: &servicecontrol.Request{
					Size: r.ContentLength,
					Time: received.UTC().Format(time.RFC3339),
				},
				Response: &servicecontrol.Response{
					Time:           time.Now().UTC().Format(time.RFC3339),
					Code:           responseCode,
					BackendLatency: time.Duration(70) * time.Millisecond,
				},
				Destination: &servicecontrol.Peer{
					RegionCode: "us-central1",
				},
				Resource: &servicecontrol.Resource{
					Name: name,
				},
			},
		},
	}
	_, err := client.Services.Report("endpointsapis.appspot.com", req).Do()
	if err != nil {
		return "", err
	}
	return "{}", nil
}

// Parse processes the request path and extract the resource name and
// permissions.
func parse(r *http.Request) (string, string, error) {
	// Split the request path.
	segments := strings.Split(r.URL.Path, "/")

	// The request path must match "/v1/projects/*/locations/*/workspaces/*" or
	// "/v1/projects/*/locations/*/workspaces". They correspond to the
	// GetWorkspace() and ListWorkspaces() methods defined in ../v1/workspace.proto.
	if segments[0] != "" || segments[1] != "v1" || segments[2] != "projects" || segments[4] != "locations" || segments[6] != "workspaces" || len(segments) > 8 {
		return "", "", errors.New("Resource '" + r.URL.Path + "' not found.")
	}

	// Skip prefix "/v1/".
	resource := r.URL.Path[4:]
	permission := "endpointsapis.appspot.com/workspaces.list"
	if len(segments) == 8 {
		permission = "endpointsapis.appspot.com/workspaces.get"
	}
	return resource, permission, nil
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	received := time.Now()

	// Create a client for Service Control API v2.
	client, err := servicecontrol.NewService(r.Context())
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		return
	}

	resource, permission, err := parse(r)
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		return
	}

	// Perform admission control.
	result, err := check(w, r, resource, permission, received, client)

	var responseCode int64 = 200
	// Print the admission control result.
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		responseCode = 403
	} else {
		fmt.Fprintln(w, "CheckResponse:")
		fmt.Fprintln(w, result)
	}

	// Print all environment variables.
	fmt.Fprintln(w, "Environments:")
	fmt.Fprintln(w, strings.Join(os.Environ(), "\n"))

	// Print all request headers.
	fmt.Fprintln(w, "Headers:")
	for key, values := range r.Header {
		for _, value := range values {
			fmt.Fprintf(w, "%v: %v\n", key, value)
		}
	}

	// Perform telemetry report.
	report(w, r, resource, responseCode, received, client)
}

func main() {
	http.HandleFunc("/", indexHandler)

	port := os.Getenv("PORT")

	log.Printf("Listen and serve on port %s", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>feat: Update the example for report v2 integration.<commit_after>// The goapp command implements a simple App Engine app to demonstrate how to
// use the Service Control API v2 for admission control. For more information,
// see https://cloud.google.com/service-infrastructure/docs/admission-control.
package main

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	// WARNING:`go get google.golang.org/api/servicecontrol/v2" may take
	// 30 minutes or longer, depending on your network speed.
	"google.golang.org/api/servicecontrol/v2"
)

// Check calls Service Control API v2 for admission control.
// Name specifies the target resource name. Permission specifies
// the required permission on the target resource. Received
// specifies the timestamp when the request is received.
func check(w http.ResponseWriter, r *http.Request, name string, permission string, received time.Time, client *servicecontrol.Service) (string, error) {
	// Construct CheckRequest from the incoming HTTP request.
	// The code assumes the incoming request processed by App Engine ingress.
	req := &servicecontrol.CheckRequest{
		ServiceConfigId: "latest",
		Attributes: &servicecontrol.AttributeContext{
			Origin: &servicecontrol.Peer{
				Ip: r.Header.Get("x-appengine-user-ip"),
			},
			Api: &servicecontrol.Api{
				Service:   "endpointsapis.appspot.com",
				Operation: "google.example.endpointsapis.v1.Workspaces.GetWorkspace",
				Version:   "v1",
				Protocol:  r.Header.Get("x-forwarded-proto"),
			},
			Request: &servicecontrol.Request{
				Id:     r.Header.Get("x-appengine-request-log-id"),
				Time:   received.UTC().Format(time.RFC3339),
				Method: r.Method,
				Scheme: r.Header.Get("x-forwarded-proto"),
				Host:   r.Host,
				Path:   r.URL.Path,
				Headers: map[string]string{
					"authorization": r.Header.Get("authorization"),
					"user-agent":    r.Header.Get("user-agent"),
					"origin":        r.Header.Get("origin"),
					"referer":       r.Header.Get("referer"),
				},
			},
			Resource: &servicecontrol.Resource{
				Name: name,
			},
		},
		Resources: []*servicecontrol.ResourceInfo{
			{
				Name:       name,
				Type:       "endpointsapis.appspot.com/Workspace",
				Permission: permission,
			},
		},
	}
	resp, err := client.Services.Check("endpointsapis.appspot.com", req).Do()
	if err != nil {
		return "", err
	}
	json, err := resp.MarshalJSON()
	if err != nil {
		return "", err
	}
	return string(json), nil
}

// Report calls Service Control API v2 for telemetry reporting.
// Name specifies the target resource name. ResponseCode specifies
// the response code returned to user. Received specifies the
// timestamp when the request is received.
func report(w http.ResponseWriter, r *http.Request, name string, responseCode int64, received time.Time, client *servicecontrol.Service) (string, error) {
	// Construct ReportRequest from the incoming HTTP request.
	// The code assumes the incoming request processed by App Engine ingress.
	req := &servicecontrol.ReportRequest{
		ServiceConfigId: "latest",
		Operations: []*servicecontrol.AttributeContext{
			{
				Api: &servicecontrol.Api{
					Service:   "endpointsapis.appspot.com",
					Operation: "google.example.endpointsapis.v1.Workspaces.GetWorkspace",
					Version:   "v1",
					Protocol:  r.Header.Get("x-forwarded-proto"),
				},
				Request: &servicecontrol.Request{
					Size: r.ContentLength,
					Time: received.UTC().Format(time.RFC3339),
				},
				Response: &servicecontrol.Response{
					Time: time.Now().UTC().Format(time.RFC3339),
					Code: responseCode,
					Headers: map[string]string{
						"x-backend-latency": "0.007",
					},
				},
				Destination: &servicecontrol.Peer{
					RegionCode: "us-central1",
				},
				Resource: &servicecontrol.Resource{
					Name: name,
				},
			},
		},
	}
	_, err := client.Services.Report("endpointsapis.appspot.com", req).Do()
	if err != nil {
		return "", err
	}
	return "{}", nil
}

// Parse processes the request path and extract the resource name and
// permissions.
func parse(r *http.Request) (string, string, error) {
	// Split the request path.
	segments := strings.Split(r.URL.Path, "/")

	// The request path must match "/v1/projects/*/locations/*/workspaces/*" or
	// "/v1/projects/*/locations/*/workspaces". They correspond to the
	// GetWorkspace() and ListWorkspaces() methods defined in ../v1/workspace.proto.
	if segments[0] != "" || segments[1] != "v1" || segments[2] != "projects" || segments[4] != "locations" || segments[6] != "workspaces" || len(segments) > 8 {
		return "", "", errors.New("Resource '" + r.URL.Path + "' not found.")
	}

	// Skip prefix "/v1/".
	resource := r.URL.Path[4:]
	permission := "endpointsapis.appspot.com/workspaces.list"
	if len(segments) == 8 {
		permission = "endpointsapis.appspot.com/workspaces.get"
	}
	return resource, permission, nil
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	received := time.Now()

	// Create a client for Service Control API v2.
	client, err := servicecontrol.NewService(r.Context())
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		return
	}

	resource, permission, err := parse(r)
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		return
	}

	// Perform admission control.
	result, err := check(w, r, resource, permission, received, client)

	var responseCode int64 = 200
	// Print the admission control result.
	if err != nil {
		fmt.Fprintln(w, "Error:")
		fmt.Fprintln(w, err.Error())
		responseCode = 403
	} else {
		fmt.Fprintln(w, "CheckResponse:")
		fmt.Fprintln(w, result)
	}

	// Print all environment variables.
	fmt.Fprintln(w, "Environments:")
	fmt.Fprintln(w, strings.Join(os.Environ(), "\n"))

	// Print all request headers.
	fmt.Fprintln(w, "Headers:")
	for key, values := range r.Header {
		for _, value := range values {
			fmt.Fprintf(w, "%v: %v\n", key, value)
		}
	}

	// Perform telemetry report.
	report(w, r, resource, responseCode, received, client)
}

func main() {
	http.HandleFunc("/", indexHandler)

	port := os.Getenv("PORT")

	log.Printf("Listen and serve on port %s", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package internal_test

import (
	"github.com/gogf/gf/g/encoding/internal"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

var testData = []struct {
	utf8, other, otherEncoding string
}{
	{"Résumé", "Résumé", "utf8"},
	{"Résumé", "R\xe9sum\xe9", "latin-1"},
	{"これは漢字です。", "S0\x8c0o0\"oW[g0Y0\x020", "UTF-16LE"},
	{"これは漢字です。", "0S0\x8c0oo\"[W0g0Y0\x02", "UTF-16BE"},
	{"これは漢字です。", "\xfe\xff0S0\x8c0oo\"[W0g0Y0\x02", "UTF-16"},
	{"𝄢𝄞𝄪𝄫", "\xfe\xff\xd8\x34\xdd\x22\xd8\x34\xdd\x1e\xd8\x34\xdd\x2a\xd8\x34\xdd\x2b", "UTF-16"},
	{"Hello, world", "Hello, world", "ASCII"},
	{"Gdańsk", "Gda\xf1sk", "ISO-8859-2"},
	{"Ââ Čč Đđ Ŋŋ Õõ Šš Žž Åå Ää", "\xc2\xe2 \xc8\xe8 \xa9\xb9 \xaf\xbf \xd5\xf5 \xaa\xba \xac\xbc \xc5\xe5 \xc4\xe4", "ISO-8859-10"},
	{"สำหรับ", "\xca\xd3\xcb\xc3\u047a", "ISO-8859-11"},
	{"latviešu", "latvie\xf0u", "ISO-8859-13"},
	{"Seònaid", "Se\xf2naid", "ISO-8859-14"},
	{"€1 is cheap", "\xa41 is cheap", "ISO-8859-15"},
	{"românește", "rom\xe2ne\xbate", "ISO-8859-16"},
	{"nutraĵo", "nutra\xbco", "ISO-8859-3"},
	{"Kalâdlit", "Kal\xe2dlit", "ISO-8859-4"},
	{"русский", "\xe0\xe3\xe1\xe1\xda\xd8\xd9", "ISO-8859-5"},
	{"ελληνικά", "\xe5\xeb\xeb\xe7\xed\xe9\xea\xdc", "ISO-8859-7"},
	{"Kağan", "Ka\xf0an", "ISO-8859-9"},
	{"Résumé", "R\x8esum\x8e", "macintosh"},
	{"Gdańsk", "Gda\xf1sk", "windows-1250"},
	{"русский", "\xf0\xf3\xf1\xf1\xea\xe8\xe9", "windows-1251"},
	{"Résumé", "R\xe9sum\xe9", "windows-1252"},
	{"ελληνικά", "\xe5\xeb\xeb\xe7\xed\xe9\xea\xdc", "windows-1253"},
	{"Kağan", "Ka\xf0an", "windows-1254"},
	{"עִבְרִית", "\xf2\xc4\xe1\xc0\xf8\xc4\xe9\xfa", "windows-1255"},
	{"العربية", "\xc7\xe1\xda\xd1\xc8\xed\xc9", "windows-1256"},
	{"latviešu", "latvie\xf0u", "windows-1257"},
	{"Việt", "Vi\xea\xf2t", "windows-1258"},
	{"สำหรับ", "\xca\xd3\xcb\xc3\u047a", "windows-874"},
	{"русский", "\xd2\xd5\xd3\xd3\xcb\xc9\xca", "KOI8-R"},
	{"українська", "\xd5\xcb\xd2\xc1\xa7\xce\xd3\xd8\xcb\xc1", "KOI8-U"},
	{"Hello 常用國字標準字體表", "Hello \xb1`\xa5\u03b0\xea\xa6r\xbc\u0437\u01e6r\xc5\xe9\xaa\xed", "big5"},
	{"Hello 常用國字標準字體表", "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed", "gbk"},
	{"Hello 常用國字標準字體表", "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed", "gb18030"},
	{"花间一壶酒，独酌无相亲。", "~{;(<dR;:x>F#,6@WCN^O`GW!#", "GB2312"},
	{"花间一壶酒，独酌无相亲。", "~{;(<dR;:x>F#,6@WCN^O`GW!#", "HZGB2312"},
	{"עִבְרִית", "\x81\x30\xfb\x30\x81\x30\xf6\x34\x81\x30\xf9\x33\x81\x30\xf6\x30\x81\x30\xfb\x36\x81\x30\xf6\x34\x81\x30\xfa\x31\x81\x30\xfb\x38", "gb18030"},
	{"㧯", "\x82\x31\x89\x38", "gb18030"},
	{"㧯", "㧯", "UTF-8"},
	{"これは漢字です。", "\x82\xb1\x82\xea\x82\xcd\x8a\xbf\x8e\x9a\x82\xc5\x82\xb7\x81B", "SJIS"},
	{"これは漢字です。", "\xa4\xb3\xa4\xec\xa4\u03f4\xc1\xbb\xfa\xa4\u01e4\xb9\xa1\xa3", "EUC-JP"},
}

func TestDecode(t *testing.T) {
	for _, data := range testData {
		str := ""
		str, err := internal.Convert("UTF-8", data.otherEncoding, data.other)
		if err != nil {
			t.Errorf("Could not create decoder for %v", err)
			continue
		}

		if str != data.utf8 {
			t.Errorf("Unexpected value: %#v (expected %#v) %v", str, data.utf8, data.otherEncoding)
		}
	}
}

func TestEncode(t *testing.T) {
	for _, data := range testData {
		str := ""
		str, err := internal.Convert(data.otherEncoding, "UTF-8", data.utf8)
		if err != nil {
			t.Errorf("Could not create decoder for %v", err)
			continue
		}

		if str != data.other {
			t.Errorf("Unexpected value: %#v (expected %#v)", str, data.other)
		}
	}
}

func TestConvert(t *testing.T) {
	srcCharset := "big5"
	src := "Hello \xb1`\xa5\u03b0\xea\xa6r\xbc\u0437\u01e6r\xc5\xe9\xaa\xed"
	dstCharset := "gbk"
	dst := "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed"

	str, err := internal.Convert(dstCharset, srcCharset, src)
	if err != nil {
		t.Errorf("convert error. %v", err)
		return
	}

	if str != dst {
		t.Errorf("unexpected value:%#v (expected %#v)", str, dst)
	}
}

func TestGetCharset(t *testing.T) {
	gtest.Case(t, func() {
		if internal.GetCharset("XX") {
			t.Errorf("unexpected value:%v (expected %v)", true, false)
		}
	})

	gtest.Case(t, func() {
		if internal.GetCharset("UTF-8") == false {
			t.Errorf("unexpected value:%v (expected %v)", false, true)
		}
	})

	gtest.Case(t, func() {
		if internal.GetCharset("gbk") == false {
			t.Errorf("unexpected value:%v (expected %v)", false, true)
		}
	})
}
<commit_msg>框架中增加字符集转换的标准库<commit_after>// Copyright 2018 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

package internal_test

import (
	"github.com/gogf/gf/g/encoding/internal"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

var testData = []struct {
	utf8, other, otherEncoding string
}{
	{"Résumé", "Résumé", "utf8"},
	{"Résumé", "R\xe9sum\xe9", "latin-1"},
	{"これは漢字です。", "S0\x8c0o0\"oW[g0Y0\x020", "UTF-16LE"},
	{"これは漢字です。", "0S0\x8c0oo\"[W0g0Y0\x02", "UTF-16BE"},
	{"これは漢字です。", "\xfe\xff0S0\x8c0oo\"[W0g0Y0\x02", "UTF-16"},
	{"𝄢𝄞𝄪𝄫", "\xfe\xff\xd8\x34\xdd\x22\xd8\x34\xdd\x1e\xd8\x34\xdd\x2a\xd8\x34\xdd\x2b", "UTF-16"},
	{"Hello, world", "Hello, world", "ASCII"},
	{"Gdańsk", "Gda\xf1sk", "ISO-8859-2"},
	{"Ââ Čč Đđ Ŋŋ Õõ Šš Žž Åå Ää", "\xc2\xe2 \xc8\xe8 \xa9\xb9 \xaf\xbf \xd5\xf5 \xaa\xba \xac\xbc \xc5\xe5 \xc4\xe4", "ISO-8859-10"},
	{"สำหรับ", "\xca\xd3\xcb\xc3\u047a", "ISO-8859-11"},
	{"latviešu", "latvie\xf0u", "ISO-8859-13"},
	{"Seònaid", "Se\xf2naid", "ISO-8859-14"},
	{"€1 is cheap", "\xa41 is cheap", "ISO-8859-15"},
	{"românește", "rom\xe2ne\xbate", "ISO-8859-16"},
	{"nutraĵo", "nutra\xbco", "ISO-8859-3"},
	{"Kalâdlit", "Kal\xe2dlit", "ISO-8859-4"},
	{"русский", "\xe0\xe3\xe1\xe1\xda\xd8\xd9", "ISO-8859-5"},
	{"ελληνικά", "\xe5\xeb\xeb\xe7\xed\xe9\xea\xdc", "ISO-8859-7"},
	{"Kağan", "Ka\xf0an", "ISO-8859-9"},
	{"Résumé", "R\x8esum\x8e", "macintosh"},
	{"Gdańsk", "Gda\xf1sk", "windows-1250"},
	{"русский", "\xf0\xf3\xf1\xf1\xea\xe8\xe9", "windows-1251"},
	{"Résumé", "R\xe9sum\xe9", "windows-1252"},
	{"ελληνικά", "\xe5\xeb\xeb\xe7\xed\xe9\xea\xdc", "windows-1253"},
	{"Kağan", "Ka\xf0an", "windows-1254"},
	{"עִבְרִית", "\xf2\xc4\xe1\xc0\xf8\xc4\xe9\xfa", "windows-1255"},
	{"العربية", "\xc7\xe1\xda\xd1\xc8\xed\xc9", "windows-1256"},
	{"latviešu", "latvie\xf0u", "windows-1257"},
	{"Việt", "Vi\xea\xf2t", "windows-1258"},
	{"สำหรับ", "\xca\xd3\xcb\xc3\u047a", "windows-874"},
	{"русский", "\xd2\xd5\xd3\xd3\xcb\xc9\xca", "KOI8-R"},
	{"українська", "\xd5\xcb\xd2\xc1\xa7\xce\xd3\xd8\xcb\xc1", "KOI8-U"},
	{"Hello 常用國字標準字體表", "Hello \xb1`\xa5\u03b0\xea\xa6r\xbc\u0437\u01e6r\xc5\xe9\xaa\xed", "big5"},
	{"Hello 常用國字標準字體表", "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed", "gbk"},
	{"Hello 常用國字標準字體表", "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed", "gb18030"},
	{"花间一壶酒，独酌无相亲。", "~{;(<dR;:x>F#,6@WCN^O`GW!#", "GB2312"},
	{"花间一壶酒，独酌无相亲。", "~{;(<dR;:x>F#,6@WCN^O`GW!#", "HZGB2312"},
	{"עִבְרִית", "\x81\x30\xfb\x30\x81\x30\xf6\x34\x81\x30\xf9\x33\x81\x30\xf6\x30\x81\x30\xfb\x36\x81\x30\xf6\x34\x81\x30\xfa\x31\x81\x30\xfb\x38", "gb18030"},
	{"㧯", "\x82\x31\x89\x38", "gb18030"},
	{"㧯", "㧯", "UTF-8"},
	{"これは漢字です。", "\x82\xb1\x82\xea\x82\xcd\x8a\xbf\x8e\x9a\x82\xc5\x82\xb7\x81B", "SJIS"},
	{"これは漢字です。", "\xa4\xb3\xa4\xec\xa4\u03f4\xc1\xbb\xfa\xa4\u01e4\xb9\xa1\xa3", "EUC-JP"},
}

func TestDecode(t *testing.T) {
	for _, data := range testData {
		str := ""
		str, err := internal.Convert("UTF-8", data.otherEncoding, data.other)
		if err != nil {
			t.Errorf("Could not create decoder for %v", err)
			continue
		}

		if str != data.utf8 {
			t.Errorf("Unexpected value: %#v (expected %#v) %v", str, data.utf8, data.otherEncoding)
		}
	}
}

func TestEncode(t *testing.T) {
	for _, data := range testData {
		str := ""
		str, err := internal.Convert(data.otherEncoding, "UTF-8", data.utf8)
		if err != nil {
			t.Errorf("Could not create decoder for %v", err)
			continue
		}

		if str != data.other {
			t.Errorf("Unexpected value: %#v (expected %#v)", str, data.other)
		}
	}
}

func TestConvert(t *testing.T) {
	srcCharset := "big5"
	src := "Hello \xb1`\xa5\u03b0\xea\xa6r\xbc\u0437\u01e6r\xc5\xe9\xaa\xed"
	dstCharset := "gbk"
	dst := "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed"

	str, err := internal.Convert(dstCharset, srcCharset, src)
	if err != nil {
		t.Errorf("convert error. %v", err)
		return
	}

	if str != dst {
		t.Errorf("unexpected value:%#v (expected %#v)", str, dst)
	}
}

func TestGetCharset(t *testing.T) {
	gtest.Case(t, func() {
		if internal.GetCharset("XX") {
			t.Errorf("unexpected value:%v (expected %v)", true, false)
		}
	})

	gtest.Case(t, func() {
		if internal.GetCharset("UTF-8") == false {
			t.Errorf("unexpected value:%v (expected %v)", false, true)
		}
	})

	gtest.Case(t, func() {
		if internal.GetCharset("gbk") == false {
			t.Errorf("unexpected value:%v (expected %v)", false, true)
		}
	})
}

func TestErrConvert(t *testing.T) {
	srcCharset := "XX"
	src := "Hello \xb1`\xa5\u03b0\xea\xa6r\xbc\u0437\u01e6r\xc5\xe9\xaa\xed"
	dstCharset := "gbk"
	//dst := "Hello \xb3\xa3\xd3\xc3\x87\xf8\xd7\xd6\x98\xcb\x9c\xca\xd7\xd6\xf3\x77\xb1\xed"

	_, err := internal.Convert(dstCharset, srcCharset, src)
	if err == nil {
		t.Errorf("convert error. %v", err)
		return
	}

}
<|endoftext|>


<commit_before>package dns

import (
	"context"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	multiQuestionDNS = map[net.Address]bool{
		net.IPAddress([]byte{8, 8, 8, 8}): true,
		net.IPAddress([]byte{8, 8, 4, 4}): true,
		net.IPAddress([]byte{9, 9, 9, 9}): true,
	}
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher core.Dispatcher) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (s *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	s.Lock()
	for id, r := range s.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(s.requests, id)
	}
	s.Unlock()
}

func (s *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	s.Lock()
	if len(s.requests) > CleanupThreshold && s.nextCleanup.Before(time.Now()) {
		s.nextCleanup = time.Now().Add(CleanupInterval)
		go s.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := s.requests[id]; found {
			continue
		}
		newError("add pending request id ", id).AtDebug().WriteToLog()
		s.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	s.Unlock()
	return id
}

func (s *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err == dns.ErrTruncated {
		newError("truncated message received. DNS server should still work. If you see anything abnormal, please submit an issue to v2ray-core.").AtWarning().WriteToLog()
	} else if err != nil {
		newError("failed to parse DNS response").Base(err).AtWarning().WriteToLog()
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	newError("handling response for id ", id, " content: ", msg).AtDebug().WriteToLog()

	s.Lock()
	request, found := s.requests[id]
	if !found {
		s.Unlock()
		return
	}
	delete(s.requests, id)
	s.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (s *UDPNameServer) buildAMsg(domain string, id uint16) *dns.Msg {
	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{
		{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeA,
			Qclass: dns.ClassINET,
		}}
	if multiQuestionDNS[s.address.Address] {
		msg.Question = append(msg.Question, dns.Question{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeAAAA,
			Qclass: dns.ClassINET,
		})
	}

	return msg
}

func msgToBuffer(msg *dns.Msg) (*buf.Buffer, error) {
	buffer := buf.New()
	if err := buffer.Reset(func(b []byte) (int, error) {
		writtenBuffer, err := msg.PackBuffer(b)
		return len(writtenBuffer), err
	}); err != nil {
		return nil, err
	}
	return buffer, nil
}

func (s *UDPNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)
	id := s.AssignUnusedID(response)

	msg := s.buildAMsg(domain, id)
	b, err := msgToBuffer(msg)
	if err != nil {
		newError("failed to build A query for domain ", domain).Base(err).WriteToLog()
		close(response)
		return response
	}

	ctx, cancel := context.WithCancel(context.Background())
	s.udpServer.Dispatch(ctx, s.address, b, s.HandleResponse)

	go func() {
		for i := 0; i < 2; i++ {
			time.Sleep(time.Second)
			s.Lock()
			_, found := s.requests[id]
			s.Unlock()
			if !found {
				break
			}
			b, _ := msgToBuffer(msg)
			s.udpServer.Dispatch(ctx, s.address, b, s.HandleResponse)
		}
		cancel()
	}()

	return response
}

type LocalNameServer struct {
}

func (*LocalNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)

	go func() {
		defer close(response)

		resolver := net.SystemIPResolver()
		ips, err := resolver.LookupIP(domain)
		if err != nil {
			newError("failed to lookup IPs for domain ", domain).Base(err).AtWarning().WriteToLog()
			return
		}

		response <- &ARecord{
			IPs:    ips,
			Expire: time.Now().Add(time.Hour),
		}
	}()

	return response
}
<commit_msg>fix LocalNameServer<commit_after>package dns

import (
	"context"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	multiQuestionDNS = map[net.Address]bool{
		net.IPAddress([]byte{8, 8, 8, 8}): true,
		net.IPAddress([]byte{8, 8, 4, 4}): true,
		net.IPAddress([]byte{9, 9, 9, 9}): true,
	}
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher core.Dispatcher) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (s *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	s.Lock()
	for id, r := range s.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(s.requests, id)
	}
	s.Unlock()
}

func (s *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	s.Lock()
	if len(s.requests) > CleanupThreshold && s.nextCleanup.Before(time.Now()) {
		s.nextCleanup = time.Now().Add(CleanupInterval)
		go s.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := s.requests[id]; found {
			continue
		}
		newError("add pending request id ", id).AtDebug().WriteToLog()
		s.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	s.Unlock()
	return id
}

func (s *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err == dns.ErrTruncated {
		newError("truncated message received. DNS server should still work. If you see anything abnormal, please submit an issue to v2ray-core.").AtWarning().WriteToLog()
	} else if err != nil {
		newError("failed to parse DNS response").Base(err).AtWarning().WriteToLog()
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	newError("handling response for id ", id, " content: ", msg).AtDebug().WriteToLog()

	s.Lock()
	request, found := s.requests[id]
	if !found {
		s.Unlock()
		return
	}
	delete(s.requests, id)
	s.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (s *UDPNameServer) buildAMsg(domain string, id uint16) *dns.Msg {
	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{
		{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeA,
			Qclass: dns.ClassINET,
		}}
	if multiQuestionDNS[s.address.Address] {
		msg.Question = append(msg.Question, dns.Question{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeAAAA,
			Qclass: dns.ClassINET,
		})
	}

	return msg
}

func msgToBuffer(msg *dns.Msg) (*buf.Buffer, error) {
	buffer := buf.New()
	if err := buffer.Reset(func(b []byte) (int, error) {
		writtenBuffer, err := msg.PackBuffer(b)
		return len(writtenBuffer), err
	}); err != nil {
		return nil, err
	}
	return buffer, nil
}

func (s *UDPNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)
	id := s.AssignUnusedID(response)

	msg := s.buildAMsg(domain, id)
	b, err := msgToBuffer(msg)
	if err != nil {
		newError("failed to build A query for domain ", domain).Base(err).WriteToLog()
		close(response)
		return response
	}

	ctx, cancel := context.WithCancel(context.Background())
	s.udpServer.Dispatch(ctx, s.address, b, s.HandleResponse)

	go func() {
		for i := 0; i < 2; i++ {
			time.Sleep(time.Second)
			s.Lock()
			_, found := s.requests[id]
			s.Unlock()
			if !found {
				break
			}
			b, _ := msgToBuffer(msg)
			s.udpServer.Dispatch(ctx, s.address, b, s.HandleResponse)
		}
		cancel()
	}()

	return response
}

type LocalNameServer struct {
}

func (*LocalNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)

	go func() {
		defer close(response)

		ips, err := net.LookupIP(domain)
		if err != nil {
			newError("failed to lookup IPs for domain ", domain).Base(err).AtWarning().WriteToLog()
			return
		}

		response <- &ARecord{
			IPs:    ips,
			Expire: time.Now().Add(time.Hour),
		}
	}()

	return response
}
<|endoftext|>


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

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	os.ErrorString
}

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
// Even if Read returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available rather than block waiting for more.
//
// At the end of the input stream, Read returns 0, os.EOF.
// Read may return a non-zero number of bytes with a non-nil err.
// In particular, a Read that exhausts the input may return n > 0, os.EOF.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying data stream.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// Even if ReadAt returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, ReadAt blocks
// until either all the data is available or an error occurs.
//
// At the end of the input stream, ReadAt returns 0, os.EOF.
// ReadAt may return a non-zero number of bytes with a non-nil err.
// In particular, a ReadAt that exhausts the input may return n > 0, os.EOF.
//
// If ReadAt is reading from an data stream with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<commit_msg>io: add ByteScanner, RuneScanner interfaces<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	os.ErrorString
}

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
// Even if Read returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available rather than block waiting for more.
//
// At the end of the input stream, Read returns 0, os.EOF.
// Read may return a non-zero number of bytes with a non-nil err.
// In particular, a Read that exhausts the input may return n > 0, os.EOF.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying data stream.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// Even if ReadAt returns n < len(p),
// it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, ReadAt blocks
// until either all the data is available or an error occurs.
//
// At the end of the input stream, ReadAt returns 0, os.EOF.
// ReadAt may return a non-zero number of bytes with a non-nil err.
// In particular, a ReadAt that exhausts the input may return n > 0, os.EOF.
//
// If ReadAt is reading from an data stream with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// ByteScanner is the interface that adds the UnreadByte method to the
// basic ReadByte method.
//
// UnreadByte causes the next call to ReadByte to return the same byte
// as the previous call to ReadByte.
// It may be an error to call UnreadByte twice without an intervening
// call to ReadByte.
type ByteScanner interface {
	ByteReader
	UnreadByte() os.Error
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// RuneScanner is the interface that adds the UnreadRune method to the
// basic ReadRune method.
//
// UnreadRune causes the next call to ReadRune to return the same rune
// as the previous call to ReadRune.
// It may be an error to call UnreadRune twice without an intervening
// call to ReadRune.
type RuneScanner interface {
	RuneReader
	UnreadRune() os.Error
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the error, if any.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<|endoftext|>