<commit_before>package deepstylelib

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
)

const (
	DesignDocName = "unprocessed_jobs"
	ViewName      = "unprocessed_jobs"
)

// Adds a time when we first saw this job in it's current state
// so that we can detect "stuck" jobs that never get processed.
type TrackedDeepStyleJob struct {
	firstSeen time.Time
	jobDoc    JobDocument
}

// Keep a map of jobs we are tracking to see if they are "stuck".
// Key is the job id and the value is the job with timestamp metadata
var trackedJobs map[string]TrackedDeepStyleJob = map[string]TrackedDeepStyleJob{}

func numJobsReadyOrBeingProcessed(syncGwAdminUrl string) (metricValue float64, err error) {

	viewResults, err := getJobsReadyOrBeingProcessed(syncGwAdminUrl)
	if err != nil {
		return 0.0, err
	}
	numRows := viewResults["total_rows"].(float64)
	return float64(numRows), nil

}

func getJobDocsBeingProcessed(syncGwAdminUrl string) (jobs []JobDocument, err error) {

	jobs = []JobDocument{}

	db, err := GetDbConnection(syncGwAdminUrl)
	if err != nil {
		return jobs, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}

	config := configuration{
		Database: db,
	}

	viewResults, err := getJobsReadyOrBeingProcessed(syncGwAdminUrl)
	if err != nil {
		return jobs, err
	}
	rows := viewResults["rows"].([]interface{})
	for _, row := range rows {
		rowMap := row.(map[string]interface{})
		docId := rowMap["id"].(string)
		jobDoc, err := NewJobDocument(docId, config)
		if err != nil {
			log.Printf("Error %v retrieving job doc: %v, skipping", err, docId)
			continue
		}
		if jobDoc.State == StateBeingProcessed {
			jobs = append(jobs, *jobDoc)
		}

	}

	return jobs, nil

}

func getJobsReadyOrBeingProcessed(syncGwAdminUrl string) (viewResults map[string]interface{}, err error) {

	// try to query view
	//    curl localhost:4985/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs
	// if we get a 404, then install the view and then requery

	output := map[string]interface{}{}

	db, err := GetDbConnection(syncGwAdminUrl)
	if err != nil {
		return output, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}

	viewUrl := fmt.Sprintf("_design/%v/_view/%v", DesignDocName, ViewName)
	options := map[string]interface{}{}
	options["stale"] = "false"

	err = db.Query(viewUrl, options, &output)
	if err != nil {
		if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "not_found") {
			// the view doesn't exist yet, attempt to install view
			if errInstallView := installView(syncGwAdminUrl); errInstallView != nil {
				// failed to install view, give up
				return output, errInstallView

			}

			// without this workaround, I'm getting:
			// ERROR: HTTP Error 500 Internal Server Error - {"error":"Internal Server Error","reason":"Internal error: error executing view req at http://127.0.0.1:8092/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs?stale=false: 500 Internal Server Error - {\"error\":\"unknown_error\",\"reason\":\"view_undefined\"}\n"}

			log.Printf("Sleeping 10s to wait for view to be ready")
			<-time.After(time.Duration(10) * time.Second)
			log.Printf("Done sleeping 10s to wait for view to be ready")

			// now retry
			errInner := db.Query(viewUrl, options, &output)
			if errInner != nil {
				// failed again, give up
				return output, errInner
			}
		} else {
			return output, err
		}
	}
	return output, nil

}

type ViewParams struct {
	JobDocType string
	JobState1  string
	JobState2  string
	JobState3  string
}

func installView(syncGwAdminUrl string) error {

	// if url has a trailing slash, remove it
	syncGwAdminUrl = strings.TrimSuffix(syncGwAdminUrl, "/")

	viewJsonTemplate := `
{
    "views":{
        "unprocessed_jobs":{
            "map":"function (doc, meta) { if (doc.type != '{{.JobDocType}}') { return; } if (doc.state == '{{.JobState1}}' || doc.state == '{{.JobState2}}' || doc.state == '{{.JobState3}}') { emit(doc.state, meta.id); }}"
        }
    }
}
`

	viewParams := ViewParams{
		JobDocType: Job,
		JobState1:  StateNotReadyToProcess,
		JobState2:  StateReadyToProcess,
		JobState3:  StateBeingProcessed,
	}
	tmpl, err := template.New("UnprocessedJobsView").Parse(viewJsonTemplate)
	if err != nil {
		return err
	}

	var buffer bytes.Buffer // A Buffer needs no initialization.

	err = tmpl.Execute(&buffer, viewParams)
	if err != nil {
		return err
	}

	log.Printf("installView called")

	// curl -X PUT -H "Content-type: application/json" localhost:4985/todolite/_design/all_lists --data @testview
	viewUrl := fmt.Sprintf("%v/_design/%v", syncGwAdminUrl, DesignDocName)

	bufferBytes := buffer.Bytes()
	log.Printf("view: %v", string(bufferBytes))

	req, err := http.NewRequest("PUT", viewUrl, bytes.NewReader(bufferBytes))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}

	log.Printf("put view resp: %v", resp)

	return nil

}

func resetStuckJobs(jobs []JobDocument) error {

	for _, job := range jobs {

		// have we seen it before?
		trackedJob, ok := trackedJobs[job.Id]
		if !ok {
			log.Printf("Tracking job %v which is currently being processed", job.Id)
			// no -- add to jobTracker map with a first_seen timestamp
			trackedJobInsert := TrackedDeepStyleJob{
				jobDoc:    job,
				firstSeen: time.Now(),
			}
			trackedJobs[job.Id] = trackedJobInsert
			continue
		}

		// yes, we've seen it before.  is first_seen more than an hour old?
		duration := time.Since(trackedJob.firstSeen)
		if duration.Minutes() >= 60 {
			// over an hour old, reset job state

			log.Printf("Job %v has been stuck for over an hour.  Resetting state to %v", job.Id, StateReadyToProcess)

			updated, err := job.UpdateState(StateReadyToProcess)
			if !updated {
				log.Printf("Unable to update job state for job: %v", job.Id)
			}
			if err != nil {
				log.Printf("Unable to update job state for job: %v.  Error: %v", job.Id, err)
				continue
			}
			// remove the job from the job tracker map since state changed
			delete(trackedJobs, job.Id)

		} else {
			log.Printf("Job %v has been processing for %v minutes", job.Id, duration.Minutes())
		}

	}
	return nil
}

func AddCloudWatchMetrics(syncGwAdminUrl string) error {

	for {

		jobs, err := getJobDocsBeingProcessed(syncGwAdminUrl)
		if err != nil {
			log.Printf("Error getting jobs being processed: %v", err)
			return err
		}

		err = resetStuckJobs(jobs)
		if err != nil {
			log.Printf("Error resetting stuck jobs: %v", err)
			return err
		}

		log.Printf("Adding metrics for queue")
		addCloudWatchMetric(syncGwAdminUrl)

		numSecondsToSleep := 60
		log.Printf("Sleeping %v seconds", numSecondsToSleep)
		<-time.After(time.Duration(numSecondsToSleep) * time.Second)

	}

}

func addCloudWatchMetric(syncGwAdminUrl string) error {

	metricValue, err := numJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("Adding metric: numJobsReadyOrBeingProcessed = %v", metricValue)
	if err != nil {
		return err
	}

	cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

	metricName := "NumJobsReadyOrBeingProcessed"
	timestamp := time.Now()

	metricDatum := &cloudwatch.MetricDatum{
		MetricName: &metricName,
		Value:      &metricValue,
		Timestamp:  &timestamp,
	}

	metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
	namespace := "DeepStyleQueue"

	putMetricDataInput := &cloudwatch.PutMetricDataInput{
		MetricData: metricDatumSlice,
		Namespace:  &namespace,
	}

	_, err = cloudwatchSvc.PutMetricData(putMetricDataInput)
	if err != nil {
		log.Printf("ERROR adding metric data  %v", err)
		return err
	}

	return nil

}
<commit_msg>add verbose logging<commit_after>package deepstylelib

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/alecthomas/template"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
)

const (
	DesignDocName = "unprocessed_jobs"
	ViewName      = "unprocessed_jobs"
)

// Adds a time when we first saw this job in it's current state
// so that we can detect "stuck" jobs that never get processed.
type TrackedDeepStyleJob struct {
	firstSeen time.Time
	jobDoc    JobDocument
}

// Keep a map of jobs we are tracking to see if they are "stuck".
// Key is the job id and the value is the job with timestamp metadata
var trackedJobs map[string]TrackedDeepStyleJob = map[string]TrackedDeepStyleJob{}

func numJobsReadyOrBeingProcessed(syncGwAdminUrl string) (metricValue float64, err error) {

	viewResults, err := getJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("Job ready or being processed: %+v", viewResults)
	if err != nil {
		return 0.0, err
	}
	numRows := viewResults["total_rows"].(float64)
	return float64(numRows), nil

}

func getJobDocsBeingProcessed(syncGwAdminUrl string) (jobs []JobDocument, err error) {

	jobs = []JobDocument{}

	db, err := GetDbConnection(syncGwAdminUrl)
	if err != nil {
		return jobs, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}

	config := configuration{
		Database: db,
	}

	viewResults, err := getJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("Job ready or being processed: %+v", viewResults)
	if err != nil {
		return jobs, err
	}
	rows := viewResults["rows"].([]interface{})
	for _, row := range rows {
		rowMap := row.(map[string]interface{})
		docId := rowMap["id"].(string)
		jobDoc, err := NewJobDocument(docId, config)
		if err != nil {
			log.Printf("Error %v retrieving job doc: %v, skipping", err, docId)
			continue
		}
		if jobDoc.State == StateBeingProcessed {
			jobs = append(jobs, *jobDoc)
		}

	}

	return jobs, nil

}

func getJobsReadyOrBeingProcessed(syncGwAdminUrl string) (viewResults map[string]interface{}, err error) {

	// try to query view
	//    curl localhost:4985/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs
	// if we get a 404, then install the view and then requery

	output := map[string]interface{}{}

	db, err := GetDbConnection(syncGwAdminUrl)
	if err != nil {
		return output, fmt.Errorf("Error connecting to db: %v.  Err: %v", syncGwAdminUrl, err)
	}

	viewUrl := fmt.Sprintf("_design/%v/_view/%v", DesignDocName, ViewName)
	options := map[string]interface{}{}
	options["stale"] = "false"

	err = db.Query(viewUrl, options, &output)
	if err != nil {
		if strings.Contains(err.Error(), "404") || strings.Contains(err.Error(), "not_found") {
			// the view doesn't exist yet, attempt to install view
			if errInstallView := installView(syncGwAdminUrl); errInstallView != nil {
				// failed to install view, give up
				return output, errInstallView

			}

			// without this workaround, I'm getting:
			// ERROR: HTTP Error 500 Internal Server Error - {"error":"Internal Server Error","reason":"Internal error: error executing view req at http://127.0.0.1:8092/deepstyle/_design/unprocessed_jobs/_view/unprocessed_jobs?stale=false: 500 Internal Server Error - {\"error\":\"unknown_error\",\"reason\":\"view_undefined\"}\n"}

			log.Printf("Sleeping 10s to wait for view to be ready")
			<-time.After(time.Duration(10) * time.Second)
			log.Printf("Done sleeping 10s to wait for view to be ready")

			// now retry
			errInner := db.Query(viewUrl, options, &output)
			if errInner != nil {
				// failed again, give up
				return output, errInner
			}
		} else {
			return output, err
		}
	}
	return output, nil

}

type ViewParams struct {
	JobDocType string
	JobState1  string
	JobState2  string
	JobState3  string
}

func installView(syncGwAdminUrl string) error {

	// if url has a trailing slash, remove it
	syncGwAdminUrl = strings.TrimSuffix(syncGwAdminUrl, "/")

	viewJsonTemplate := `
{
    "views":{
        "unprocessed_jobs":{
            "map":"function (doc, meta) { if (doc.type != '{{.JobDocType}}') { return; } if (doc.state == '{{.JobState1}}' || doc.state == '{{.JobState2}}' || doc.state == '{{.JobState3}}') { emit(doc.state, meta.id); }}"
        }
    }
}
`

	viewParams := ViewParams{
		JobDocType: Job,
		JobState1:  StateNotReadyToProcess,
		JobState2:  StateReadyToProcess,
		JobState3:  StateBeingProcessed,
	}
	tmpl, err := template.New("UnprocessedJobsView").Parse(viewJsonTemplate)
	if err != nil {
		return err
	}

	var buffer bytes.Buffer // A Buffer needs no initialization.

	err = tmpl.Execute(&buffer, viewParams)
	if err != nil {
		return err
	}

	log.Printf("installView called")

	// curl -X PUT -H "Content-type: application/json" localhost:4985/todolite/_design/all_lists --data @testview
	viewUrl := fmt.Sprintf("%v/_design/%v", syncGwAdminUrl, DesignDocName)

	bufferBytes := buffer.Bytes()
	log.Printf("view: %v", string(bufferBytes))

	req, err := http.NewRequest("PUT", viewUrl, bytes.NewReader(bufferBytes))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		return err
	}

	log.Printf("put view resp: %v", resp)

	return nil

}

func resetStuckJobs(jobs []JobDocument) error {

	for _, job := range jobs {

		// have we seen it before?
		trackedJob, ok := trackedJobs[job.Id]
		if !ok {
			log.Printf("Tracking job %v which is currently being processed", job.Id)
			// no -- add to jobTracker map with a first_seen timestamp
			trackedJobInsert := TrackedDeepStyleJob{
				jobDoc:    job,
				firstSeen: time.Now(),
			}
			trackedJobs[job.Id] = trackedJobInsert
			continue
		}

		// yes, we've seen it before.  is first_seen more than an hour old?
		duration := time.Since(trackedJob.firstSeen)
		if duration.Minutes() >= 60 {
			// over an hour old, reset job state

			log.Printf("Job %v has been stuck for over an hour.  Resetting state to %v", job.Id, StateReadyToProcess)

			updated, err := job.UpdateState(StateReadyToProcess)
			if !updated {
				log.Printf("Unable to update job state for job: %v", job.Id)
			}
			if err != nil {
				log.Printf("Unable to update job state for job: %v.  Error: %v", job.Id, err)
				continue
			}
			// remove the job from the job tracker map since state changed
			delete(trackedJobs, job.Id)

		} else {
			log.Printf("Job %v has been processing for %v minutes", job.Id, duration.Minutes())
		}

	}
	return nil
}

func AddCloudWatchMetrics(syncGwAdminUrl string) error {

	for {

		jobs, err := getJobDocsBeingProcessed(syncGwAdminUrl)
		if err != nil {
			log.Printf("Error getting jobs being processed: %v", err)
			return err
		}

		err = resetStuckJobs(jobs)
		if err != nil {
			log.Printf("Error resetting stuck jobs: %v", err)
			return err
		}

		log.Printf("Adding metrics for queue")
		addCloudWatchMetric(syncGwAdminUrl)

		numSecondsToSleep := 60
		log.Printf("Sleeping %v seconds", numSecondsToSleep)
		<-time.After(time.Duration(numSecondsToSleep) * time.Second)

	}

}

func addCloudWatchMetric(syncGwAdminUrl string) error {

	metricValue, err := numJobsReadyOrBeingProcessed(syncGwAdminUrl)
	log.Printf("Adding metric: numJobsReadyOrBeingProcessed = %v", metricValue)
	if err != nil {
		return err
	}

	cloudwatchSvc := cloudwatch.New(session.New(), &aws.Config{Region: aws.String("us-east-1")})

	metricName := "NumJobsReadyOrBeingProcessed"
	timestamp := time.Now()

	metricDatum := &cloudwatch.MetricDatum{
		MetricName: &metricName,
		Value:      &metricValue,
		Timestamp:  &timestamp,
	}

	metricDatumSlice := []*cloudwatch.MetricDatum{metricDatum}
	namespace := "DeepStyleQueue"

	putMetricDataInput := &cloudwatch.PutMetricDataInput{
		MetricData: metricDatumSlice,
		Namespace:  &namespace,
	}

	_, err = cloudwatchSvc.PutMetricData(putMetricDataInput)
	if err != nil {
		log.Printf("ERROR adding metric data  %v", err)
		return err
	}

	return nil

}
<|endoftext|>


<commit_before>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/strib/gomounts"
)

var kbfusePath = fuse.OSXFUSEPaths{
	DevicePrefix: "/dev/kbfuse",
	Load:         "/Library/Filesystems/kbfuse.fs/Contents/Resources/load_kbfuse",
	Mount:        "/Library/Filesystems/kbfuse.fs/Contents/Resources/mount_kbfuse",
	DaemonVar:    "MOUNT_KBFUSE_DAEMON_PATH",
}

const (
	mountpointTimeout = 5 * time.Second
)

type symlink struct {
	link string
}

func (s symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	a.Mode = os.ModeSymlink | a.Mode | 0555
	a.Valid = 0
	return nil
}

func (s symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (
	link string, err error) {
	return s.link, nil
}

type cacheEntry struct {
	mountpoint string
	time       time.Time
}

type root struct {
	lock            sync.RWMutex
	mountpointCache map[uint32]cacheEntry
}

func newRoot() *root {
	return &root{
		mountpointCache: make(map[uint32]cacheEntry),
	}
}

func (r *root) Root() (fs.Node, error) {
	return r, nil
}

func (r *root) Attr(ctx context.Context, attr *fuse.Attr) error {
	attr.Mode = os.ModeDir | 0555
	return nil
}

func (r *root) getCachedMountpoint(uid uint32) string {
	r.lock.RLock()
	defer r.lock.RUnlock()
	entry, ok := r.mountpointCache[uid]
	if !ok {
		return ""
	}
	now := time.Now()
	if now.Sub(entry.time) > mountpointTimeout {
		// Don't bother deleting the entry, since the caller should
		// just overwrite it.
		return ""
	}
	return entry.mountpoint
}

func (r *root) findKBFSMount(ctx context.Context) (
	mountpoint string, err error) {
	uid := ctx.Value(fs.CtxHeaderUIDKey).(uint32)
	// Don't let the root see anything here; we don't want a symlink
	// loop back to this mount.
	if uid == 0 {
		return "", fuse.ENOENT
	}

	mountpoint = r.getCachedMountpoint(uid)
	if mountpoint != "" {
		return mountpoint, nil
	}

	defer func() {
		if err != nil {
			return
		}
		// Cache the entry if we didn't hit an error.
		r.lock.Lock()
		defer r.lock.Unlock()
		r.mountpointCache[uid] = cacheEntry{
			mountpoint: mountpoint,
			time:       time.Now(),
		}
	}()

	u, err := user.LookupId(strconv.FormatUint(uint64(uid), 10))
	if err != nil {
		return "", err
	}

	vols, err := gomounts.GetMountedVolumes()
	if err != nil {
		return "", err
	}
	fuseType := "fuse"
	if runtime.GOOS == "darwin" {
		fuseType = "kbfuse"
	}
	var fuseMountPoints []string
	for _, v := range vols {
		if v.Type != fuseType {
			continue
		}
		if v.Owner != u.Uid {
			continue
		}
		fuseMountPoints = append(fuseMountPoints, v.Path)
	}

	if len(fuseMountPoints) == 0 {
		return "", fuse.ENOENT
	}
	if len(fuseMountPoints) == 1 {
		return fuseMountPoints[0], nil
	}

	// If there is more than one, pick the first one alphabetically
	// that has "keybase" in the path.
	sort.Strings(fuseMountPoints)
	for _, mp := range fuseMountPoints {
		// TODO: a better regexp that will rule out keybase.staging if
		// we're in prod mode, etc.
		if !strings.Contains(mp, "keybase") {
			continue
		}
		// Double-check that this is a real KBFS mount.
		if _, err := os.Stat(filepath.Join(mp, ".kbfs_error")); err != nil {
			continue
		}
		return mp, nil
	}

	// Give up and return the first one.
	return fuseMountPoints[0], nil
}

func (r *root) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	_, err := r.findKBFSMount(ctx)
	if err != nil {
		if err == fuse.ENOENT {
			return []fuse.Dirent{}, nil
		} else {
			return []fuse.Dirent{}, err
		}
	}

	// TODO: show the `kbfs.error.txt" and "kbfs.nologin.txt" files if
	// they exist?  As root, it is hard to figure out if they're
	// there, though.
	return []fuse.Dirent{
		{
			Type: fuse.DT_Link,
			Name: "private",
		},
		{
			Type: fuse.DT_Link,
			Name: "public",
		},
		fuse.Dirent{
			Type: fuse.DT_Link,
			Name: "team",
		},
	}, nil
}

func (r *root) Lookup(
	ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (
	n fs.Node, err error) {
	mountpoint, err := r.findKBFSMount(ctx)
	if err != nil {
		return nil, err
	}

	resp.EntryValid = 0
	switch req.Name {
	case "private", "public", "team", ".kbfs_error", ".kbfs_metrics",
		".kbfs_profiles", ".kbfs_reset_caches", ".kbfs_status",
		"kbfs.error.txt", "kbfs.nologin.txt", ".kbfs_enable_auto_journals",
		".kbfs_disable_auto_journals", ".kbfs_enable_block_prefetching",
		".kbfs_disable_block_prefetching", ".kbfs_enable_debug_server",
		".kbfs_disable_debug_server":
		return symlink{filepath.Join(mountpoint, req.Name)}, nil
	}
	return nil, fuse.ENOENT
}

func main() {
	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	if currUser.Uid != "0" {
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
	}

	options := []fuse.MountOption{fuse.AllowOther()}
	options = append(options, fuse.FSName("keybase-redirector"))
	options = append(options, fuse.ReadOnly())
	if runtime.GOOS == "darwin" {
		options = append(options, fuse.OSXFUSELocations(kbfusePath))
		options = append(options, fuse.VolumeName("keybase-redirector"))
		options = append(options, fuse.NoBrowse())
	}

	c, err := fuse.Mount(os.Args[1], options...)
	if err != nil {
		fmt.Printf("Mount error, exiting cleanly: %+v\n", err)
		os.Exit(0)
	}

	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)
	go func() {
		_ = <-interruptChan
		err := fuse.Unmount(os.Args[1])
		if err != nil {
			fmt.Printf("Couldn't unmount cleanly: %+v", err)
		}
	}()

	srv := fs.New(c, &fs.Config{
		WithContext: func(ctx context.Context, _ fuse.Request) context.Context {
			return context.Background()
		},
	})
	srv.Serve(newRoot())
}
<commit_msg>redirector: fix lint<commit_after>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/strib/gomounts"
)

var kbfusePath = fuse.OSXFUSEPaths{
	DevicePrefix: "/dev/kbfuse",
	Load:         "/Library/Filesystems/kbfuse.fs/Contents/Resources/load_kbfuse",
	Mount:        "/Library/Filesystems/kbfuse.fs/Contents/Resources/mount_kbfuse",
	DaemonVar:    "MOUNT_KBFUSE_DAEMON_PATH",
}

const (
	mountpointTimeout = 5 * time.Second
)

type symlink struct {
	link string
}

func (s symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	a.Mode = os.ModeSymlink | a.Mode | 0555
	a.Valid = 0
	return nil
}

func (s symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (
	link string, err error) {
	return s.link, nil
}

type cacheEntry struct {
	mountpoint string
	time       time.Time
}

type root struct {
	lock            sync.RWMutex
	mountpointCache map[uint32]cacheEntry
}

func newRoot() *root {
	return &root{
		mountpointCache: make(map[uint32]cacheEntry),
	}
}

func (r *root) Root() (fs.Node, error) {
	return r, nil
}

func (r *root) Attr(ctx context.Context, attr *fuse.Attr) error {
	attr.Mode = os.ModeDir | 0555
	return nil
}

func (r *root) getCachedMountpoint(uid uint32) string {
	r.lock.RLock()
	defer r.lock.RUnlock()
	entry, ok := r.mountpointCache[uid]
	if !ok {
		return ""
	}
	now := time.Now()
	if now.Sub(entry.time) > mountpointTimeout {
		// Don't bother deleting the entry, since the caller should
		// just overwrite it.
		return ""
	}
	return entry.mountpoint
}

func (r *root) findKBFSMount(ctx context.Context) (
	mountpoint string, err error) {
	uid := ctx.Value(fs.CtxHeaderUIDKey).(uint32)
	// Don't let the root see anything here; we don't want a symlink
	// loop back to this mount.
	if uid == 0 {
		return "", fuse.ENOENT
	}

	mountpoint = r.getCachedMountpoint(uid)
	if mountpoint != "" {
		return mountpoint, nil
	}

	defer func() {
		if err != nil {
			return
		}
		// Cache the entry if we didn't hit an error.
		r.lock.Lock()
		defer r.lock.Unlock()
		r.mountpointCache[uid] = cacheEntry{
			mountpoint: mountpoint,
			time:       time.Now(),
		}
	}()

	u, err := user.LookupId(strconv.FormatUint(uint64(uid), 10))
	if err != nil {
		return "", err
	}

	vols, err := gomounts.GetMountedVolumes()
	if err != nil {
		return "", err
	}
	fuseType := "fuse"
	if runtime.GOOS == "darwin" {
		fuseType = "kbfuse"
	}
	var fuseMountPoints []string
	for _, v := range vols {
		if v.Type != fuseType {
			continue
		}
		if v.Owner != u.Uid {
			continue
		}
		fuseMountPoints = append(fuseMountPoints, v.Path)
	}

	if len(fuseMountPoints) == 0 {
		return "", fuse.ENOENT
	}
	if len(fuseMountPoints) == 1 {
		return fuseMountPoints[0], nil
	}

	// If there is more than one, pick the first one alphabetically
	// that has "keybase" in the path.
	sort.Strings(fuseMountPoints)
	for _, mp := range fuseMountPoints {
		// TODO: a better regexp that will rule out keybase.staging if
		// we're in prod mode, etc.
		if !strings.Contains(mp, "keybase") {
			continue
		}
		// Double-check that this is a real KBFS mount.
		if _, err := os.Stat(filepath.Join(mp, ".kbfs_error")); err != nil {
			continue
		}
		return mp, nil
	}

	// Give up and return the first one.
	return fuseMountPoints[0], nil
}

func (r *root) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	_, err := r.findKBFSMount(ctx)
	if err != nil {
		if err == fuse.ENOENT {
			return []fuse.Dirent{}, nil
		}
		return []fuse.Dirent{}, err
	}

	// TODO: show the `kbfs.error.txt" and "kbfs.nologin.txt" files if
	// they exist?  As root, it is hard to figure out if they're
	// there, though.
	return []fuse.Dirent{
		{
			Type: fuse.DT_Link,
			Name: "private",
		},
		{
			Type: fuse.DT_Link,
			Name: "public",
		},
		fuse.Dirent{
			Type: fuse.DT_Link,
			Name: "team",
		},
	}, nil
}

func (r *root) Lookup(
	ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (
	n fs.Node, err error) {
	mountpoint, err := r.findKBFSMount(ctx)
	if err != nil {
		return nil, err
	}

	resp.EntryValid = 0
	switch req.Name {
	case "private", "public", "team", ".kbfs_error", ".kbfs_metrics",
		".kbfs_profiles", ".kbfs_reset_caches", ".kbfs_status",
		"kbfs.error.txt", "kbfs.nologin.txt", ".kbfs_enable_auto_journals",
		".kbfs_disable_auto_journals", ".kbfs_enable_block_prefetching",
		".kbfs_disable_block_prefetching", ".kbfs_enable_debug_server",
		".kbfs_disable_debug_server":
		return symlink{filepath.Join(mountpoint, req.Name)}, nil
	}
	return nil, fuse.ENOENT
}

func main() {
	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	if currUser.Uid != "0" {
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
	}

	options := []fuse.MountOption{fuse.AllowOther()}
	options = append(options, fuse.FSName("keybase-redirector"))
	options = append(options, fuse.ReadOnly())
	if runtime.GOOS == "darwin" {
		options = append(options, fuse.OSXFUSELocations(kbfusePath))
		options = append(options, fuse.VolumeName("keybase-redirector"))
		options = append(options, fuse.NoBrowse())
	}

	c, err := fuse.Mount(os.Args[1], options...)
	if err != nil {
		fmt.Printf("Mount error, exiting cleanly: %+v\n", err)
		os.Exit(0)
	}

	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)
	go func() {
		_ = <-interruptChan
		err := fuse.Unmount(os.Args[1])
		if err != nil {
			fmt.Printf("Couldn't unmount cleanly: %+v", err)
		}
	}()

	srv := fs.New(c, &fs.Config{
		WithContext: func(ctx context.Context, _ fuse.Request) context.Context {
			return context.Background()
		},
	})
	srv.Serve(newRoot())
}
<|endoftext|>


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

package main

import (
	"errors"
	"fmt"
	"reflect"
	"unsafe"
	"strings"
    "encoding/json"
	"github.com/hyperledger/fabric/core/chaincode/shim"
)

// ReferralChaincode implementation stores and updates referral information on the blockchain
type ReferralChaincode struct {
}

type CustomerReferral struct {
	referralId string
    customerName string
	contactNumber string
	customerId string
	employeeId string
	departments []string
    createDate int64
	status string
}

func main() {
	err := shim.Start(new(ReferralChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}

func BytesToString(b []byte) string {
    bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
    sh := reflect.StringHeader{bh.Data, bh.Len}
    return *(*string)(unsafe.Pointer(&sh))
}

// Init resets all the things
func (t *ReferralChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	// There is no initialization to do
	return nil, nil
}

// Invoke is our entry point to invoke a chaincode function
func (t *ReferralChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	fmt.Println("invoke is running " + function)

	// Handle different functions
	return []byte("Function: " + function + "being invoked"), nil
	if function == "init" {
		return t.Init(stub, "init", args)
	} else if function == "createReferral" {

		return t.createReferral(stub, args)
	} else if function == "updateReferralStatus" {
		return t.updateReferralStatus(stub, args)
	}

	return []byte("Function: " + function + "not found"), nil
	fmt.Println("invoke did not find func: " + function)

	return nil, errors.New("Received unknown function invocation")
}

// Query is our entry point for queries
func (t *ReferralChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	fmt.Println("query is running " + function)

	// Handle different functions
	if function == "read" { //read a variable
		return t.read(stub, args)
	} else if function == "searchByStatus" {
		return t.searchByStatus(args[0], stub)
	} else if function == "searchByDepartment" {
		return t.searchByDepartment(args[0], stub)
	}
	fmt.Println("query did not find func: " + function)

	return nil, errors.New("Received unknown function query")
}

// Adds the referral id to a ledger list item for the given department allowing for quick search of referrals in a given department
func (t *ReferralChaincode) indexByDepartment(referralId string, department string, stub *shim.ChaincodeStub) (error) {
	valAsbytes, err := stub.GetState(department)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + department + "\"}"
		return errors.New(jsonResp)
	}

	if valAsbytes == nil {
		err = stub.PutState(department, []byte(referralId))
	} else {
	    commaDelimitedStatuses := BytesToString(valAsbytes)
		err = stub.PutState(department, []byte(commaDelimitedStatuses + "," + referralId))
	}

	return err
}

func (t *ReferralChaincode) removeStatusReferralIndex(referralId string, status string, stub *shim.ChaincodeStub) (error) {
	valAsbytes, err := stub.GetState(status)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + status + "\"}"
		return errors.New(jsonResp)
	}

	if valAsbytes == nil {
		return nil;
	} else {
		// Remove the referral from this status type, if it exists
		commaDelimitedStatuses := BytesToString(valAsbytes)
		referralIdsInCurrentStatus := strings.Split(commaDelimitedStatuses, ",")
		updatedReferralIdList := ""

		appendComma := false
		for i := range referralIdsInCurrentStatus {
			if referralIdsInCurrentStatus[i] != referralId {
			    if appendComma == false {
					updatedReferralIdList += referralIdsInCurrentStatus[i]
					appendComma = true
				} else {
					updatedReferralIdList = updatedReferralIdList + "," + referralIdsInCurrentStatus[i]
				}
			}
		}

		err = stub.PutState(status, []byte(updatedReferralIdList))
	}

	if err != nil {
		jsonResp := "{\"Error\":\"Failed to update state for " + status + "\"}"
		return errors.New(jsonResp)
	}

	return nil
}

// Adds the referral id to a ledger list item for the given department allowing for quick search of referrals in a given department
func (t *ReferralChaincode) indexByStatus(referralId string, status string, stub *shim.ChaincodeStub) (error) {
	valAsbytes, err := stub.GetState(status)
	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + status + "\"}"
		return errors.New(jsonResp)
	}

	if valAsbytes == nil {
		err = stub.PutState(status, []byte(referralId))
	} else {
	    commaDelimitedStatuses := BytesToString(valAsbytes)
		err = stub.PutState(status, []byte(commaDelimitedStatuses + "," + referralId))
	}

	if err != nil {
		jsonResp := "{\"Error\":\"Failed to update state for " + status + "\"}"
		return errors.New(jsonResp)
	}

	return nil
}

func (t *ReferralChaincode) unmarshallBytes(valAsBytes []byte) (error, CustomerReferral) {
	var err error
	var referral CustomerReferral
	fmt.Println("Unmarshalling JSON")
	err = json.Unmarshal(valAsBytes, &referral)

	if err != nil {
		fmt.Println("Unmarshalling JSON failed")
	}

	return err, referral
}

func (t *ReferralChaincode) marshallReferral(referral CustomerReferral) (error, []byte) {
	fmt.Println("Marshalling JSON to bytes")
	valAsbytes, err := json.Marshal(referral)

	if err != nil {
		fmt.Println("Marshalling JSON to bytes failed")
		return err, nil
	}

	return nil, valAsbytes
}

func (t *ReferralChaincode) updateStatus(referral CustomerReferral, status string, stub *shim.ChaincodeStub) (error) {
	fmt.Println("Setting status")

	err := t.removeStatusReferralIndex(referral.referralId, referral.status, stub)
	if err != nil {
		return err
	}
	referral.status = status
	err = t.indexByStatus(referral.referralId, status, stub)

	return err
}

// updateReferral - invoke function to updateReferral key/value pair
func (t *ReferralChaincode) updateReferralStatus(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	var key, status string
	var err error
	var referral CustomerReferral
	fmt.Println("running updateReferral()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
	}

	key = args[0] //rename for funsies
	status = args[1]

	valAsbytes, err := stub.GetState(key)
	if err != nil {
		return nil, err
	}
	err, referral = t.unmarshallBytes(valAsbytes)

	if err != nil {
		return nil, err
	}

	t.updateStatus(referral, status, stub)

	fmt.Println("Marshalling JSON to bytes")
	err, valAsbytes = t.marshallReferral(referral)

	if err != nil {
		return nil, err
	}

	err = stub.PutState(key, valAsbytes) //write the variable into the chaincode state
	if err != nil {
		return nil, err
	}

	return nil, nil
}

// createReferral - invoke function to write key/value pair
func (t *ReferralChaincode) createReferral(stub *shim.ChaincodeStub, args []string) ([]byte, error) {

	var key, value string
	var err error
	var referral CustomerReferral
	fmt.Println("running createReferral()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
	}

	key = args[0] //rename for funsies
	value = args[1]

	err = stub.PutState(key, []byte(value)) //write the variable into the chaincode state
	if err != nil {
		return []byte("Could not put the key: " + key + " and value: " + value + " on the ledger"), err
	}

	// Deserialize the input string into a GO data structure to hold the referral
	err, referral = t.unmarshallBytes([]byte(value))
	if err != nil {
		return []byte("Count not unmarshall the bytes from the value: " + value + " on the ledger"), err
	}

	// Create a ledger record that indexes the referral id by the created status
	err = t.indexByStatus(referral.referralId, referral.status, stub)
	if err != nil {
		return []byte("Count not index the bytes by status from the value: " + value + " on the ledger"), err
	}

	// Create a ledger record that indexes the referral id by the created department
	for i := range referral.departments {
		err = t.indexByDepartment(referral.referralId, referral.departments[i], stub)
		if err != nil {
			return []byte("Count not index the bytes by department from the value: " + value + " on the ledger"), err
		}
	}

	return nil, err
}

func (t *ReferralChaincode) processCommaDelimitedReferrals(delimitedReferrals string, stub *shim.ChaincodeStub) ([]byte, error) {
	commaDelimitedReferrals := strings.Split(delimitedReferrals, ",")

	referralResultSet := ""
	appendComma := false

	for i := range commaDelimitedReferrals {
		valAsbytes, err := stub.GetState(commaDelimitedReferrals[i])

		if err != nil {
			return nil, err
		}

		if appendComma == false {
			referralResultSet += BytesToString(valAsbytes)
		} else {
			referralResultSet = referralResultSet + "," + BytesToString(valAsbytes)
		}
	}

	return []byte(referralResultSet), nil
}

func (t *ReferralChaincode) searchByDepartment(department string, stub *shim.ChaincodeStub) ([]byte, error) {
	valAsbytes, err := stub.GetState(department)

	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + department + "\"}"
		return nil, errors.New(jsonResp)
	}

	valAsbytes, err = t.processCommaDelimitedReferrals(BytesToString(valAsbytes), stub)

	if(err != nil) {
		return nil, err
	}

	return valAsbytes, nil
}

func (t *ReferralChaincode) searchByStatus(status string, stub *shim.ChaincodeStub) ([]byte, error) {
	valAsbytes, err := stub.GetState(status)

	if err != nil {
		jsonResp := "{\"Error\":\"Failed to get state for " + status + "\"}"
		return nil, errors.New(jsonResp)
	}

	valAsbytes, err = t.processCommaDelimitedReferrals(BytesToString(valAsbytes), stub)

	if(err != nil) {
		return nil, err
	}

	return valAsbytes, nil
}

// read - query function to read key/value pair
func (t *ReferralChaincode) read(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	var key, jsonResp string
	var err error

	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the key to query")
	}

	key = args[0]
	valAsbytes, err := stub.GetState(key)

	if err != nil {
		jsonResp = "{\"Error\":\"Failed to get state for " + key + "\"}"
		return []byte(jsonResp), err
	}

	if valAsbytes == nil {
		return []byte("Did not find entry for key: " + key), nil
	}
	return valAsbytes, nil
}<commit_msg>Add files via upload<commit_after>/*
Copyright IBM Corp 2016 All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
		 http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"errors"
	"fmt"

	"github.com/hyperledger/fabric/core/chaincode/shim"
)

// ReferralChaincode implementation stores and updates referral information on the blockchain
type ReferralChaincode struct {
}

func main() {
	err := shim.Start(new(ReferralChaincode))
	if err != nil {
		fmt.Printf("Error starting Simple chaincode: %s", err)
	}
}

// Init resets all the things
func (t *ReferralChaincode) Init(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	// There is no initialization to do
	return nil, nil
}

// Invoke is our entry point to invoke a chaincode function
func (t *ReferralChaincode) Invoke(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	fmt.Println("invoke is running " + function)

	// Handle different functions
	if function == "init" {
		return t.Init(stub, "init", args)
	} else if function == "createReferral" {
		return t.createReferral(stub, args)
	} else if function == "updateReferral" {
		return t.updateReferral(stub, args)
	}
	fmt.Println("invoke did not find func: " + function)

	return nil, errors.New("Received unknown function invocation")
}

// Query is our entry point for queries
func (t *ReferralChaincode) Query(stub *shim.ChaincodeStub, function string, args []string) ([]byte, error) {
	fmt.Println("query is running " + function)

	// Handle different functions
	if function == "read" { //read a variable
		return t.read(stub, args)
	}
	fmt.Println("query did not find func: " + function)

	return nil, errors.New("Received unknown function query")
}

// updateReferral - invoke function to updateReferral key/value pair
func (t *ReferralChaincode) updateReferral(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	var key, value string
	var err error
	fmt.Println("running updateReferral()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
	}

	key = args[0] //rename for funsies
	value = args[1]
	err = stub.PutState(key, []byte(value)) //write the variable into the chaincode state
	if err != nil {
		return nil, err
	}
	return nil, nil
}

// createReferral - invoke function to write key/value pair
func (t *ReferralChaincode) createReferral(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	var key, value string
	var err error
	fmt.Println("running createReferral()")

	if len(args) != 2 {
		return nil, errors.New("Incorrect number of arguments. Expecting 2. name of the key and value to set")
	}

	key = args[0] //rename for funsies
	value = args[1]
	err = stub.PutState(key, []byte(value)) //write the variable into the chaincode state
	if err != nil {
		return nil, err
	}
	return nil, nil
}

// read - query function to read key/value pair
func (t *ReferralChaincode) read(stub *shim.ChaincodeStub, args []string) ([]byte, error) {
	var key, jsonResp string
	var err error

	if len(args) != 1 {
		return nil, errors.New("Incorrect number of arguments. Expecting name of the key to query")
	}

	key = args[0]
	valAsbytes, err := stub.GetState(key)
	if err != nil {
		jsonResp = "{\"Error\":\"Failed to get state for " + key + "\"}"
		return nil, errors.New(jsonResp)
	}

	return valAsbytes, nil
}<|endoftext|>


<commit_before>package registry

import (
	"reflect"
	"testing"

	"github.com/coreos/fleet/unit"
)

func TestMarshaling(t *testing.T) {
	units := []string{
		``,

		`[Service]
		ExecStart=/bin/sleep 1`,

		`[Unit]
		Description=Foo

		[Service]
		ExecStart=echo "foo"`,

		`[Path]
		PathExists=/foo`,
	}

	for _, contents := range units {
		u, _ := unit.NewUnit(contents)
		json, err := marshal(u)
		if err != nil {
			t.Error("Error marshaling unit:", err)
		}
		var um unit.Unit
		err = unmarshal(json, &um)
		if err != nil {
			t.Error("Error unmarshaling unit:", err)
		}
		if !reflect.DeepEqual(*u, um) {
			t.Errorf("Unmarshaled unit does not match original!\nOriginal:\n%s\nUnmarshaled:\n%s", *u, um)
		}
	}

}

func TestLegacyPayload(t *testing.T) {
	unitContents := `
[Service]
ExecStart=/bin/sleep 30000
`[1:]
	legacyPayloadContents := `{"Name":"sleep.service","Unit":{"Contents":{"Service":{"ExecStart":"/bin/sleep 30000"}},"Raw":"[Service]\nExecStart=/bin/sleep 30000\n"}}`
	want, _ := unit.NewUnit(unitContents)
	var ljp LegacyJobPayload
	err := unmarshal(legacyPayloadContents, &ljp)
	if err != nil {
		t.Error("Error unmarshaling legacy payload:", err)
	}
	got := ljp.Unit
	if !reflect.DeepEqual(*want, got) {
		t.Errorf("Unit from legacy payload does not match expected!\nwant:\n%s\ngot:\n%s", *want, got)
	}
}
<commit_msg>chore(registry): handle NewUnit errors in registry tests<commit_after>package registry

import (
	"reflect"
	"testing"

	"github.com/coreos/fleet/unit"
)

func TestMarshaling(t *testing.T) {
	units := []string{
		``,

		`[Service]
		ExecStart=/bin/sleep 1`,

		`[Unit]
		Description=Foo

		[Service]
		ExecStart=echo "foo"`,

		`[Path]
		PathExists=/foo`,
	}

	for _, contents := range units {
		u, err := unit.NewUnit(contents)
		if err != nil {
			t.Fatalf("unexpected error creating unit from %q: %v", contents, err)
		}
		json, err := marshal(u)
		if err != nil {
			t.Error("Error marshaling unit:", err)
		}
		var um unit.Unit
		err = unmarshal(json, &um)
		if err != nil {
			t.Error("Error unmarshaling unit:", err)
		}
		if !reflect.DeepEqual(*u, um) {
			t.Errorf("Unmarshaled unit does not match original!\nOriginal:\n%s\nUnmarshaled:\n%s", *u, um)
		}
	}

}

func TestLegacyPayload(t *testing.T) {
	contents := `
[Service]
ExecStart=/bin/sleep 30000
`[1:]
	legacyPayloadContents := `{"Name":"sleep.service","Unit":{"Contents":{"Service":{"ExecStart":"/bin/sleep 30000"}},"Raw":"[Service]\nExecStart=/bin/sleep 30000\n"}}`
	want, err := unit.NewUnit(contents)
	if err != nil {
		t.Fatalf("unexpected error creating unit from %q: %v", contents, err)
	}
	var ljp LegacyJobPayload
	err = unmarshal(legacyPayloadContents, &ljp)
	if err != nil {
		t.Error("Error unmarshaling legacy payload:", err)
	}
	got := ljp.Unit
	if !reflect.DeepEqual(*want, got) {
		t.Errorf("Unit from legacy payload does not match expected!\nwant:\n%s\ngot:\n%s", *want, got)
	}
}
<|endoftext|>


<commit_before>package render

import (
	"github.com/akavel/polyclip-go"
	"github.com/oakmound/oak/alg/floatgeom"
)

// A DrawPolygon is used to determine whether elements should be drawn, defining
// a polygonal area for what things should be visible.
type DrawPolygon struct {
	usingDrawPolygon bool
	drawPolygon      polyclip.Polygon
}

// SetDrawPolygon sets the draw polygon and flags that draw functions
// should check for containment in the polygon before drawing elements
func (dp *DrawPolygon) SetDrawPolygon(p polyclip.Polygon) {
	dp.usingDrawPolygon = true
	dp.drawPolygon = p
}

// ClearDrawPolygon will stop checking the set draw polygon for whether elements
// should be drawn to screen. If SetDrawPolygon was not called before this was
// called, this does nothing.
// This may in the future be called at the start of new scenes.
func (dp *DrawPolygon) ClearDrawPolygon() {
	dp.usingDrawPolygon = false
}

// DrawPolygonDim returns the dimensions of this draw polygon, or (0,0)->(0,0)
// if there is no draw polygon in use.
func (dp *DrawPolygon) DrawPolygonDim() floatgeom.Rect2 {
	if !dp.usingDrawPolygon {
		return floatgeom.Rect2{}
	}
	mbr := dp.drawPolygon.BoundingBox()
	return floatgeom.NewRect2(mbr.Min.X, mbr.Min.Y, mbr.Max.X, mbr.Max.Y)
}

// InDrawPolygon returns whehter a coordinate and dimension set should be drawn
// given the draw polygon
func (dp *DrawPolygon) InDrawPolygon(xi, yi, x2i, y2i int) bool {
	if dp.usingDrawPolygon {
		x := float64(xi)
		y := float64(yi)
		x2 := float64(x2i)
		y2 := float64(y2i)
		p2 := polyclip.Polygon{{{X: x, Y: y}, {X: x, Y: y2}, {X: x2, Y: y2}, {X: x2, Y: y}}}
		intsct := dp.drawPolygon.Construct(polyclip.INTERSECTION, p2)
		return len(intsct) != 0
	}
	return true
}
<commit_msg>Deprecate polyclip, use our own polygon intersection<commit_after>package render

import (
	"github.com/akavel/polyclip-go"
	"github.com/oakmound/oak/alg"
	"github.com/oakmound/oak/alg/floatgeom"
)

// A DrawPolygon is used to determine whether elements should be drawn, defining
// a polygonal area for what things should be visible.
type DrawPolygon struct {
	usingDrawPolygon bool
	drawPolygon      []floatgeom.Point2
	dims             floatgeom.Rect2
	rectangular      bool
}

// SetDrawPolygon sets the draw polygon and flags that draw functions
// should check for containment in the polygon before drawing elements
// Deprecated: use SetPolygon instead
func (dp *DrawPolygon) SetDrawPolygon(p polyclip.Polygon) {
	// get []floatgeom.Point2
	poly := make([]floatgeom.Point2, 0, len(p))
	for _, c := range p {
		for _, pt := range c {
			poly = append(poly, floatgeom.Point2{pt.X, pt.Y})
		}
	}
	dp.SetPolygon(poly)
}

// SetPolygon sets the draw polygon and flags that draw functions
// should check for containment in the polygon before drawing elements.
func (dp *DrawPolygon) SetPolygon(poly []floatgeom.Point2) {
	dp.usingDrawPolygon = true
	dp.dims = floatgeom.NewBoundingRect2(poly...)
	dp.drawPolygon = poly
	if isRectangular(poly...) {
		dp.rectangular = true
	}
}

func isRectangular(pts ...floatgeom.Point2) bool {
	last := pts[len(pts)-1]
	for _, pt := range pts {
		// The last point needs to share an x or y value with this point
		if !alg.F64eq(pt.X(), last.X()) && !alg.F64eq(pt.Y(), last.Y()) {
			return false
		}
		last = pt
	}
	return true
}

// ClearDrawPolygon will stop checking the set draw polygon for whether elements
// should be drawn to screen. If SetDrawPolygon was not called before this was
// called, this does nothing.
// This may in the future be called at the start of new scenes.
func (dp *DrawPolygon) ClearDrawPolygon() {
	dp.usingDrawPolygon = false
	dp.dims = floatgeom.Rect2{}
	dp.rectangular = false
}

// DrawPolygonDim returns the dimensions of this draw polygon, or (0,0)->(0,0)
// if there is no draw polygon in use.
func (dp *DrawPolygon) DrawPolygonDim() floatgeom.Rect2 {
	return dp.dims
}

// InDrawPolygon returns whehter a coordinate and dimension set should be drawn
// given the draw polygon
func (dp *DrawPolygon) InDrawPolygon(xi, yi, x2i, y2i int) bool {
	if dp.usingDrawPolygon {
		x := float64(xi)
		y := float64(yi)
		x2 := float64(x2i)
		y2 := float64(y2i)

		dx := dp.dims.Min.X()
		dy := dp.dims.Min.Y()
		dx2 := dp.dims.Max.X()
		dy2 := dp.dims.Max.Y()

		dimOverlap := false
		if x > dx {
			if x < dx2 {
				dimOverlap = true
			}
		} else {
			if dx < x2 {
				dimOverlap = true
			}
		}
		if y > dy {
			if y < dy2 {
				dimOverlap = true
			}
		} else {
			if dy < y2 {
				dimOverlap = true
			}
		}
		if !dimOverlap {
			return false
		}
		if dp.rectangular {
			return true
		}
		r := floatgeom.NewRect2(x, y, x2, y2)
		diags := [][2]floatgeom.Point2{
			{
				{r.Min.X(), r.Max.Y()},
				{r.Max.X(), r.Min.Y()},
			}, {
				r.Min,
				r.Max,
			},
		}
		last := dp.drawPolygon[len(dp.drawPolygon)-1]
		for i := 0; i < len(dp.drawPolygon); i++ {
			next := dp.drawPolygon[i]
			if r.Contains(next) {
				return true
			}
			// Checking line segment from last to next
			for _, diag := range diags {
				if orient(diag[0], diag[1], last) != orient(diag[0], diag[1], next) &&
					orient(next, last, diag[0]) != orient(next, last, diag[1]) {
					return true
				}
			}

			last = next
		}
		return false
	}
	return true
}

func orient(p1, p2, p3 floatgeom.Point2) int8 {
	val := (p2.Y()-p1.Y())*(p3.X()-p2.X()) -
		(p2.X()-p1.X())*(p3.Y()-p2.Y())
	switch {
	case val < 0:
		return 2
	case val > 0:
		return 1
	default:
		return 0
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package hrpc

import (
	"bytes"
	"github.com/tsuna/gohbase/filter"
	"golang.org/x/net/context"
	"reflect"
	"testing"
)

func TestNewGet(t *testing.T) {
	ctx := context.Background()
	table := "test"
	tableb := []byte(table)
	key := "45"
	keyb := []byte(key)
	fam := make(map[string][]string)
	fam["info"] = []string{"c1"}
	filter1 := filter.NewFirstKeyOnlyFilter()
	get, err := NewGet(ctx, tableb, keyb)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, nil) {
		t.Errorf("Get1 didn't set attributes correctly.")
	}
	get, err = NewGetStr(ctx, table, key)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, nil) {
		t.Errorf("Get2 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Families(fam))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, nil) {
		t.Errorf("Get3 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, filter1) {
		t.Errorf("Get4 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1), Families(fam))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, filter1) {
		t.Errorf("Get5 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1))
	err = Families(fam)(get)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, filter1) {
		t.Errorf("Get6 didn't set attributes correctly.")
	}

}

func confirmGetAttributes(g *Get, ctx context.Context, table, key []byte, fam map[string][]string, filter1 filter.Filter) bool {
	if g.GetContext() != ctx ||
		bytes.Compare(g.Table(), table) != 0 ||
		bytes.Compare(g.Key(), key) != 0 ||
		!reflect.DeepEqual(g.GetFamilies(), fam) ||
		reflect.TypeOf(g.GetFilter()) != reflect.TypeOf(filter1) {
		return false
	}
	return true
}

func TestNewScan(t *testing.T) {
	ctx := context.Background()
	table := "test"
	tableb := []byte(table)
	fam := make(map[string][]string)
	fam["info"] = []string{"c1"}
	filter1 := filter.NewFirstKeyOnlyFilter()
	start := "0"
	stop := "100"
	startb := []byte("0")
	stopb := []byte("100")
	scan, err := NewScan(ctx, tableb)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, nil, nil) {
		t.Errorf("Scan1 didn't set attributes correctly.")
	}
	scan, err = NewScanRange(ctx, tableb, startb, stopb)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, nil, nil) {
		t.Errorf("Scan2 didn't set attributes correctly.")
	}
	scan, err = NewScanStr(ctx, table)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, nil, nil) {
		t.Errorf("Scan3 didn't set attributes correctly.")
	}
	scan, err = NewScanRangeStr(ctx, table, start, stop)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, nil, nil) {
		t.Errorf("Scan4 didn't set attributes correctly.")
	}
	scan, err = NewScanRange(ctx, tableb, startb, stopb, Families(fam), Filters(filter1))
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, fam, filter1) {
		t.Errorf("Scan5 didn't set attributes correctly.")
	}
	scan, err = NewScan(ctx, tableb, Filters(filter1), Families(fam))
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, fam, filter1) {
		t.Errorf("Scan6 didn't set attributes correctly.")
	}
}

func confirmScanAttributes(s *Scan, ctx context.Context, table, start, stop []byte, fam map[string][]string, filter1 filter.Filter) bool {
	if s.GetContext() != ctx ||
		bytes.Compare(s.Table(), table) != 0 ||
		bytes.Compare(s.GetStartRow(), start) != 0 ||
		bytes.Compare(s.GetStopRow(), stop) != 0 ||
		!reflect.DeepEqual(s.GetFamilies(), fam) ||
		reflect.TypeOf(s.GetFilter()) != reflect.TypeOf(filter1) {
		return false
	}
	return true
}
<commit_msg>Added benchmarks for the two serialization options in hrpc/mutate.go<commit_after>// Copyright (C) 2015  The GoHBase Authors.  All rights reserved.
// This file is part of GoHBase.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package hrpc

import (
	"bytes"
	"reflect"
	"testing"

	"github.com/tsuna/gohbase/filter"
	"github.com/tsuna/gohbase/regioninfo"
	"golang.org/x/net/context"
)

func TestNewGet(t *testing.T) {
	ctx := context.Background()
	table := "test"
	tableb := []byte(table)
	key := "45"
	keyb := []byte(key)
	fam := make(map[string][]string)
	fam["info"] = []string{"c1"}
	filter1 := filter.NewFirstKeyOnlyFilter()
	get, err := NewGet(ctx, tableb, keyb)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, nil) {
		t.Errorf("Get1 didn't set attributes correctly.")
	}
	get, err = NewGetStr(ctx, table, key)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, nil) {
		t.Errorf("Get2 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Families(fam))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, nil) {
		t.Errorf("Get3 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, nil, filter1) {
		t.Errorf("Get4 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1), Families(fam))
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, filter1) {
		t.Errorf("Get5 didn't set attributes correctly.")
	}
	get, err = NewGet(ctx, tableb, keyb, Filters(filter1))
	err = Families(fam)(get)
	if err != nil || !confirmGetAttributes(get, ctx, tableb, keyb, fam, filter1) {
		t.Errorf("Get6 didn't set attributes correctly.")
	}

}

func confirmGetAttributes(g *Get, ctx context.Context, table, key []byte, fam map[string][]string, filter1 filter.Filter) bool {
	if g.GetContext() != ctx ||
		bytes.Compare(g.Table(), table) != 0 ||
		bytes.Compare(g.Key(), key) != 0 ||
		!reflect.DeepEqual(g.GetFamilies(), fam) ||
		reflect.TypeOf(g.GetFilter()) != reflect.TypeOf(filter1) {
		return false
	}
	return true
}

func TestNewScan(t *testing.T) {
	ctx := context.Background()
	table := "test"
	tableb := []byte(table)
	fam := make(map[string][]string)
	fam["info"] = []string{"c1"}
	filter1 := filter.NewFirstKeyOnlyFilter()
	start := "0"
	stop := "100"
	startb := []byte("0")
	stopb := []byte("100")
	scan, err := NewScan(ctx, tableb)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, nil, nil) {
		t.Errorf("Scan1 didn't set attributes correctly.")
	}
	scan, err = NewScanRange(ctx, tableb, startb, stopb)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, nil, nil) {
		t.Errorf("Scan2 didn't set attributes correctly.")
	}
	scan, err = NewScanStr(ctx, table)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, nil, nil) {
		t.Errorf("Scan3 didn't set attributes correctly.")
	}
	scan, err = NewScanRangeStr(ctx, table, start, stop)
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, nil, nil) {
		t.Errorf("Scan4 didn't set attributes correctly.")
	}
	scan, err = NewScanRange(ctx, tableb, startb, stopb, Families(fam), Filters(filter1))
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, startb, stopb, fam, filter1) {
		t.Errorf("Scan5 didn't set attributes correctly.")
	}
	scan, err = NewScan(ctx, tableb, Filters(filter1), Families(fam))
	if err != nil || !confirmScanAttributes(scan, ctx, tableb, nil, nil, fam, filter1) {
		t.Errorf("Scan6 didn't set attributes correctly.")
	}
}

func confirmScanAttributes(s *Scan, ctx context.Context, table, start, stop []byte, fam map[string][]string, filter1 filter.Filter) bool {
	if s.GetContext() != ctx ||
		bytes.Compare(s.Table(), table) != 0 ||
		bytes.Compare(s.GetStartRow(), start) != 0 ||
		bytes.Compare(s.GetStopRow(), stop) != 0 ||
		!reflect.DeepEqual(s.GetFamilies(), fam) ||
		reflect.TypeOf(s.GetFilter()) != reflect.TypeOf(filter1) {
		return false
	}
	return true
}

func BenchmarkMutateSerializeWithNestedMaps(b *testing.B) {
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		data := map[string]map[string][]byte{
			"cf": map[string][]byte{
				"a": []byte{10},
				"b": []byte{20},
				"c": []byte{30, 0},
				"d": []byte{40, 0, 0, 0},
				"e": []byte{50, 0, 0, 0, 0, 0, 0, 0},
				"f": []byte{60},
				"g": []byte{70},
				"h": []byte{80, 0},
				"i": []byte{90, 0, 0, 0},
				"j": []byte{100, 0, 0, 0, 0, 0, 0, 0},
				"k": []byte{0, 0, 220, 66},
				"l": []byte{0, 0, 0, 0, 0, 0, 94, 64},
				"m": []byte{0, 0, 2, 67, 0, 0, 0, 0},
				"n": []byte{0, 0, 0, 0, 0, 128, 97, 64, 0, 0, 0, 0, 0, 0, 0, 0},
				"o": []byte{150},
				"p": []byte{4, 8, 15, 26, 23, 42},
				"q": []byte{1, 1, 3, 5, 8, 13, 21, 34, 55},
				"r": []byte("One Ring to rule them all, One Ring to find them, One Ring to bring them all and in the darkness bind them"),
			},
		}
		mutate, err := NewPutStr(context.Background(), "", "", data)
		if err != nil {
			b.Errorf("Error creating mutate: %v", err)
		}
		mutate.SetRegion(&regioninfo.Info{})
		mutate.Serialize()
	}
}

func BenchmarkMutateSerializeWithReflection(b *testing.B) {
	b.ReportAllocs()

	type teststr struct {
		AnInt       int        `hbase:"cf:a"`
		AnInt8      int8       `hbase:"cf:b"`
		AnInt16     int16      `hbase:"cf:c"`
		AnInt32     int32      `hbase:"cf:d"`
		AnInt64     int64      `hbase:"cf:e"`
		AnUInt      uint       `hbase:"cf:f"`
		AnUInt8     uint8      `hbase:"cf:g"`
		AnUInt16    uint16     `hbase:"cf:h"`
		AnUInt32    uint32     `hbase:"cf:i"`
		AnUInt64    uint64     `hbase:"cf:j"`
		AFloat32    float32    `hbase:"cf:k"`
		AFloat64    float64    `hbase:"cf:l"`
		AComplex64  complex64  `hbase:"cf:m"`
		AComplex128 complex128 `hbase:"cf:n"`
		APointer    *int       `hbase:"cf:o"`
		AnArray     [6]uint8   `hbase:"cf:p"`
		ASlice      []uint8    `hbase:"cf:q"`
		AString     string     `hbase:"cf:r"`
	}

	number := 150
	for i := 0; i < b.N; i++ {
		str := teststr{
			AnInt:       10,
			AnInt8:      20,
			AnInt16:     30,
			AnInt32:     40,
			AnInt64:     50,
			AnUInt:      60,
			AnUInt8:     70,
			AnUInt16:    80,
			AnUInt32:    90,
			AnUInt64:    100,
			AFloat32:    110,
			AFloat64:    120,
			AComplex64:  130,
			AComplex128: 140,
			APointer:    &number,
			AnArray:     [6]uint8{4, 8, 15, 26, 23, 42},
			ASlice:      []uint8{1, 1, 3, 5, 8, 13, 21, 34, 55},
			AString:     "One Ring to rule them all, One Ring to find them, One Ring to bring them all and in the darkness bind them",
		}
		mutate, err := NewPutStrRef(context.Background(), "", "", str)
		if err != nil {
			b.Errorf("Error creating mutate: %v", err)
		}
		mutate.SetRegion(&regioninfo.Info{})
		mutate.Serialize()
	}
}
<|endoftext|>


<commit_before>// Copyright (c) Microsoft Open Technologies, Inc.
// All Rights Reserved.
// Licensed under the Apache License, Version 2.0.
// See License.txt in the project root for license information.
package targets

import (
	"fmt"

	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/constants"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/request"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/response"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/response/model"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"time"
)

const (
	powerState_Started                string = "Started"
	powerState_Stopping               string = "Stopping"
	powerState_Stopped                string = "Stopped"
	powerState_Unknown                string = "Unknown"
	instanceStatus_ReadyRole                 = "ReadyRole"
	instanceStatus_FailedStartingRole        = "FailedStartingRole"
	instanceStatus_FailedStartingVM          = "FailedStartingVM"
	instanceStatus_ProvisioningFailed        = "ProvisioningFailed"
	instanceStatus_UnresponsiveRole          = "UnresponsiveRole"
)

type StepPollStatus struct {
	TmpServiceName string
	TmpVmName      string
	OsType         string
}

func (s *StepPollStatus) Run(state multistep.StateBag) multistep.StepAction {
	reqManager := state.Get(constants.RequestManager).(*request.Manager)
	ui := state.Get(constants.Ui).(packer.Ui)

	errorMsg := "Error Polling Temporary Azure VM is ready: %s"

	ui.Say("Polling Temporary Azure VM is ready...")

	if len(s.OsType) == 0 {
		err := fmt.Errorf(errorMsg, "'OsType' param is empty")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	firstSleepMin := time.Duration(2)
	firstSleepTime := time.Minute * firstSleepMin
	log.Printf("Sleeping for %v min to make the VM to start", uint(firstSleepMin))
	time.Sleep(firstSleepTime)

	var count uint = 60
	var duration time.Duration = 40
	sleepTime := time.Second * duration
	total := count * uint(duration)

	//	var err error
	var deployment *model.Deployment

	requestData := reqManager.GetDeployment(s.TmpServiceName, s.TmpVmName)

	for count != 0 {
		resp, err := reqManager.Execute(requestData)
		if err != nil {
			err := fmt.Errorf(errorMsg, err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}

		deployment, err = response.ParseDeployment(resp.Body)
		//		log.Printf("deployment:\n%v", deployment)

		if len(deployment.RoleInstanceList) > 0 {
			powerState := deployment.RoleInstanceList[0].PowerState
			instanceStatus := deployment.RoleInstanceList[0].InstanceStatus

			if powerState == powerState_Started && instanceStatus == instanceStatus_ReadyRole {
				break
			}

			if instanceStatus == instanceStatus_FailedStartingRole ||
				instanceStatus == instanceStatus_FailedStartingVM ||
				instanceStatus == instanceStatus_ProvisioningFailed ||
				instanceStatus == instanceStatus_UnresponsiveRole {
				err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].instanceStatus is "+instanceStatus)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}
			if powerState == powerState_Stopping ||
				powerState == powerState_Stopped ||
				powerState == powerState_Unknown {
				err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].PowerState is "+powerState)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}
		}

		// powerState_Starting or deployment.RoleInstanceList[0] == 0
		log.Println(fmt.Sprintf("Waiting for another %v seconds...", uint(duration)))
		time.Sleep(sleepTime)
		count--
	}

	if count == 0 {
		err := fmt.Errorf(errorMsg, fmt.Sprintf("time is up (%d seconds)", total))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	state.Put(constants.VmRunning, 1)

	log.Println("s.OsType = " + s.OsType)

	if s.OsType == Linux {
		endpoints := deployment.RoleInstanceList[0].InstanceEndpoints
		if len(endpoints) == 0 {
			err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].InstanceEndpoints list is empty")
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}

		vip := endpoints[0].Vip
		port := endpoints[0].PublicPort
		endpoint := fmt.Sprintf("%s:%d", vip, port)

		ui.Message("VM Endpoint: " + endpoint)
		state.Put(constants.AzureVmAddr, endpoint)
	}

	roleList := deployment.RoleList
	if len(roleList) == 0 {
		err := fmt.Errorf(errorMsg, "deployment.RoleList is empty")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	diskName := roleList[0].OSVirtualHardDisk.DiskName
	ui.Message("VM DiskName: " + diskName)
	state.Put(constants.HardDiskName, diskName)

	mediaLink := roleList[0].OSVirtualHardDisk.MediaLink
	ui.Message("VM MediaLink: " + mediaLink)
	state.Put(constants.MediaLink, mediaLink)

	return multistep.ActionContinue
}

func (s *StepPollStatus) Cleanup(state multistep.StateBag) {
	// nothing to do
}
<commit_msg>Change ip/port format back to %s:%s<commit_after>// Copyright (c) Microsoft Open Technologies, Inc.
// All Rights Reserved.
// Licensed under the Apache License, Version 2.0.
// See License.txt in the project root for license information.
package targets

import (
	"fmt"

	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/constants"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/request"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/response"
	"github.com/MSOpenTech/packer-azure/packer/builder/azure/driver_restapi/response/model"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"time"
)

const (
	powerState_Started                string = "Started"
	powerState_Stopping               string = "Stopping"
	powerState_Stopped                string = "Stopped"
	powerState_Unknown                string = "Unknown"
	instanceStatus_ReadyRole                 = "ReadyRole"
	instanceStatus_FailedStartingRole        = "FailedStartingRole"
	instanceStatus_FailedStartingVM          = "FailedStartingVM"
	instanceStatus_ProvisioningFailed        = "ProvisioningFailed"
	instanceStatus_UnresponsiveRole          = "UnresponsiveRole"
)

type StepPollStatus struct {
	TmpServiceName string
	TmpVmName      string
	OsType         string
}

func (s *StepPollStatus) Run(state multistep.StateBag) multistep.StepAction {
	reqManager := state.Get(constants.RequestManager).(*request.Manager)
	ui := state.Get(constants.Ui).(packer.Ui)

	errorMsg := "Error Polling Temporary Azure VM is ready: %s"

	ui.Say("Polling Temporary Azure VM is ready...")

	if len(s.OsType) == 0 {
		err := fmt.Errorf(errorMsg, "'OsType' param is empty")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	firstSleepMin := time.Duration(2)
	firstSleepTime := time.Minute * firstSleepMin
	log.Printf("Sleeping for %v min to make the VM to start", uint(firstSleepMin))
	time.Sleep(firstSleepTime)

	var count uint = 60
	var duration time.Duration = 40
	sleepTime := time.Second * duration
	total := count * uint(duration)

	//	var err error
	var deployment *model.Deployment

	requestData := reqManager.GetDeployment(s.TmpServiceName, s.TmpVmName)

	for count != 0 {
		resp, err := reqManager.Execute(requestData)
		if err != nil {
			err := fmt.Errorf(errorMsg, err)
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}

		deployment, err = response.ParseDeployment(resp.Body)
		//		log.Printf("deployment:\n%v", deployment)

		if len(deployment.RoleInstanceList) > 0 {
			powerState := deployment.RoleInstanceList[0].PowerState
			instanceStatus := deployment.RoleInstanceList[0].InstanceStatus

			if powerState == powerState_Started && instanceStatus == instanceStatus_ReadyRole {
				break
			}

			if instanceStatus == instanceStatus_FailedStartingRole ||
				instanceStatus == instanceStatus_FailedStartingVM ||
				instanceStatus == instanceStatus_ProvisioningFailed ||
				instanceStatus == instanceStatus_UnresponsiveRole {
				err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].instanceStatus is "+instanceStatus)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}
			if powerState == powerState_Stopping ||
				powerState == powerState_Stopped ||
				powerState == powerState_Unknown {
				err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].PowerState is "+powerState)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}
		}

		// powerState_Starting or deployment.RoleInstanceList[0] == 0
		log.Println(fmt.Sprintf("Waiting for another %v seconds...", uint(duration)))
		time.Sleep(sleepTime)
		count--
	}

	if count == 0 {
		err := fmt.Errorf(errorMsg, fmt.Sprintf("time is up (%d seconds)", total))
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	state.Put(constants.VmRunning, 1)

	log.Println("s.OsType = " + s.OsType)

	if s.OsType == Linux {
		endpoints := deployment.RoleInstanceList[0].InstanceEndpoints
		if len(endpoints) == 0 {
			err := fmt.Errorf(errorMsg, "deployment.RoleInstanceList[0].InstanceEndpoints list is empty")
			state.Put("error", err)
			ui.Error(err.Error())
			return multistep.ActionHalt
		}

		vip := endpoints[0].Vip
		port := endpoints[0].PublicPort
		endpoint := fmt.Sprintf("%s:%s", vip, port)

		ui.Message("VM Endpoint: " + endpoint)
		state.Put(constants.AzureVmAddr, endpoint)
	}

	roleList := deployment.RoleList
	if len(roleList) == 0 {
		err := fmt.Errorf(errorMsg, "deployment.RoleList is empty")
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	diskName := roleList[0].OSVirtualHardDisk.DiskName
	ui.Message("VM DiskName: " + diskName)
	state.Put(constants.HardDiskName, diskName)

	mediaLink := roleList[0].OSVirtualHardDisk.MediaLink
	ui.Message("VM MediaLink: " + mediaLink)
	state.Put(constants.MediaLink, mediaLink)

	return multistep.ActionContinue
}

func (s *StepPollStatus) Cleanup(state multistep.StateBag) {
	// nothing to do
}
<|endoftext|>


<commit_before>package auth

import (
	"encoding/gob"
	"errors"
	"net/http"
	"time"

	"github.com/gorilla/sessions"
)

// custom sessionKey type to prevent collision
type sessionKey uint

func init() {
	// need to register our Key with gob so gorilla/sessions can (de)serialize it
	gob.Register(userKey)
	gob.Register(time.Time{})
}

const (
	sessionName = "AuthSession"

	userKey sessionKey = iota
	userTimeout
)

var (
	ErrBadLogin      = errors.New("Bad Login")
	ErrNotAuthorized = errors.New("Not Authorized")
)

// Auther allows for custom authentication backends
type Auther interface {
	// Check should return a non-nil error for failed requests (like ErrBadLogin)
	// and it can pass custom data that is saved in the cookie through the first return argument
	Check(user, pass string) (interface{}, error)
}

type AuthHandler struct {
	auther Auther
	store  sessions.Store
}

func NewHandler(a Auther, store sessions.Store) (ah AuthHandler) {
	ah.auther = a
	ah.store = store
	return
}

func (ah AuthHandler) Authorize(redir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		session, err := ah.store.Get(r, sessionName)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err := r.ParseForm(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		user := r.Form.Get("user")
		pass := r.Form.Get("pass")
		if user == "" || pass == "" {
			http.Error(w, ErrBadLogin.Error(), http.StatusBadRequest)
			return
		}

		id, err := ah.auther.Check(user, pass)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		session.Values[userKey] = id
		session.Values[userTimeout] = time.Now().Add(5 * time.Minute)
		if err := session.Save(r, w); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.Redirect(w, r, redir, http.StatusTemporaryRedirect)
		return
	}
}

func (ah AuthHandler) Authenticate(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := ah.AuthenticateRequest(r); err != nil {
			http.Error(w, err.Error(), http.StatusUnauthorized)
			return
		}

		h.ServeHTTP(w, r)
	})
}

func (ah AuthHandler) AuthenticateRequest(r *http.Request) error {
	session, err := ah.store.Get(r, sessionName)
	if err != nil {
		return err
	}

	if session.IsNew {
		return ErrNotAuthorized
	}

	if _, ok := session.Values[userKey]; !ok {
		return ErrNotAuthorized
	}

	t, ok := session.Values[userTimeout]
	if !ok {
		return ErrNotAuthorized
	}

	tout, ok := t.(time.Time)
	if !ok {
		return ErrNotAuthorized
	}

	if time.Now().After(tout) {
		return ErrNotAuthorized
	}

	return nil
}

func (ah AuthHandler) Logout(redir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		session, err := ah.store.Get(r, sessionName)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		session.Values[userTimeout] = time.Now().Add(-5 * time.Minute)
		session.Options.MaxAge = -1
		if err := session.Save(r, w); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.Redirect(w, r, redir, http.StatusTemporaryRedirect)
		return
	}
}
<commit_msg>auth: renamed handler/removed stutter<commit_after>package auth

import (
	"encoding/gob"
	"errors"
	"net/http"
	"time"

	"github.com/gorilla/sessions"
)

// custom sessionKey type to prevent collision
type sessionKey uint

func init() {
	// need to register our Key with gob so gorilla/sessions can (de)serialize it
	gob.Register(userKey)
	gob.Register(time.Time{})
}

const (
	sessionName = "AuthSession"

	userKey sessionKey = iota
	userTimeout
)

var (
	ErrBadLogin      = errors.New("Bad Login")
	ErrNotAuthorized = errors.New("Not Authorized")
)

// Auther allows for custom authentication backends
type Auther interface {
	// Check should return a non-nil error for failed requests (like ErrBadLogin)
	// and it can pass custom data that is saved in the cookie through the first return argument
	Check(user, pass string) (interface{}, error)
}

type Handler struct {
	auther Auther
	store  sessions.Store
}

func NewHandler(a Auther, store sessions.Store) *Handler {
	var ah Handler
	ah.auther = a
	ah.store = store
	return &ah
}

func (ah Handler) Authorize(redir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		session, err := ah.store.Get(r, sessionName)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		if err := r.ParseForm(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		user := r.Form.Get("user")
		pass := r.Form.Get("pass")
		if user == "" || pass == "" {
			http.Error(w, ErrBadLogin.Error(), http.StatusBadRequest)
			return
		}

		id, err := ah.auther.Check(user, pass)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		session.Values[userKey] = id
		session.Values[userTimeout] = time.Now().Add(5 * time.Minute)
		if err := session.Save(r, w); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.Redirect(w, r, redir, http.StatusTemporaryRedirect)
		return
	}
}

func (ah Handler) Authenticate(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := ah.AuthenticateRequest(r); err != nil {
			http.Error(w, err.Error(), http.StatusUnauthorized)
			return
		}

		h.ServeHTTP(w, r)
	})
}

func (ah Handler) AuthenticateRequest(r *http.Request) error {
	session, err := ah.store.Get(r, sessionName)
	if err != nil {
		return err
	}

	if session.IsNew {
		return ErrNotAuthorized
	}

	if _, ok := session.Values[userKey]; !ok {
		return ErrNotAuthorized
	}

	t, ok := session.Values[userTimeout]
	if !ok {
		return ErrNotAuthorized
	}

	tout, ok := t.(time.Time)
	if !ok {
		return ErrNotAuthorized
	}

	if time.Now().After(tout) {
		return ErrNotAuthorized
	}

	return nil
}

func (ah Handler) Logout(redir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		session, err := ah.store.Get(r, sessionName)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		session.Values[userTimeout] = time.Now().Add(-5 * time.Minute)
		session.Options.MaxAge = -1
		if err := session.Save(r, w); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		http.Redirect(w, r, redir, http.StatusTemporaryRedirect)
		return
	}
}
<|endoftext|>


<commit_before>package http

import (
	"net/http"

	"github.com/hashicorp/vault/logical"
	"github.com/hashicorp/vault/vault"
)

func handleSysRemount(core *vault.Core) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "PUT", "POST":
		default:
			respondError(w, http.StatusMethodNotAllowed, nil)
			return
		}

		// Parse the request if we can
		var req RemountRequest
		if err := parseRequest(r, &req); err != nil {
			respondError(w, http.StatusBadRequest, err)
			return
		}

		_, err := core.HandleRequest(requestAuth(r, &logical.Request{
			Operation:  logical.WriteOperation,
			Path:       "sys/remount",
			Connection: getConnection(r),
			Data: map[string]interface{}{
				"from": req.From,
				"to":   req.To,
			},
		}))
		if err != nil {
			respondError(w, http.StatusInternalServerError, err)
			return
		}

		respondOk(w, nil)
	})
}

type MountRequest struct {
	Type        string `json:"type"`
	Description string `json:"description"`
}

type RemountRequest struct {
	From string `json:"from"`
	To   string `json:"to"`
}
<commit_msg>sys_mount.go is now unnecessary<commit_after><|endoftext|>


<commit_before>package rcmgr

import (
	"fmt"
	"sync"

	"github.com/libp2p/go-buffer-pool"
	"github.com/libp2p/go-libp2p-core/network"
)

type ResourceScope struct {
	sync.Mutex

	limit   Limit
	memory  int64
	buffers map[interface{}][]byte
}

var _ network.ResourceScope = (*ResourceScope)(nil)

func (rc *ResourceScope) checkMemory(rsvp int) error {
	// overflow check; this also has the side-effect that we cannot reserve negative memory.
	newmem := rc.memory + int64(rsvp)
	if newmem < rc.memory {
		return fmt.Errorf("memory reservation overflow: %w", ErrResourceLimitExceeded)
	}

	// limit check
	if newmem > rc.limit.GetMemoryLimit() {
		return fmt.Errorf("cannot reserve memory: %w", ErrResourceLimitExceeded)
	}

	return nil
}

func (rc *ResourceScope) releaseBuffers() {
	for key, buf := range rc.buffers {
		pool.Put(buf)
		delete(rc.buffers, key)
	}
}

func (rc *ResourceScope) ReserveMemory(size int) error {
	rc.Lock()
	defer rc.Unlock()

	if err := rc.checkMemory(size); err != nil {
		return err
	}

	rc.memory += int64(size)
	return nil
}

func (rc *ResourceScope) ReleaseMemory(size int) {
	rc.Lock()
	defer rc.Unlock()

	rc.memory -= int64(size)

	// sanity check for bugs upstream
	if rc.memory < 0 {
		panic("BUG: too much memory released")
	}
}

func (rc *ResourceScope) GetBuffer(size int) ([]byte, error) {
	rc.Lock()
	defer rc.Unlock()

	if err := rc.checkMemory(size); err != nil {
		return nil, err
	}

	buf := pool.Get(size)

	rc.memory += int64(size)
	rc.buffers[buf] = buf

	return buf, nil
}

func (rc *ResourceScope) GrowBuffer(oldbuf []byte, newsize int) ([]byte, error) {
	rc.Lock()
	defer rc.Unlock()

	grow := newsize - len(oldbuf)
	if err := rc.checkMemory(grow); err != nil {
		return nil, err
	}

	newbuf := pool.Get(newsize)
	copy(newbuf, oldbuf)

	rc.memory += int64(grow)
	rc.buffers[newbuf] = newbuf
	delete(rc.buffers, oldbuf)

	return newbuf, nil
}

func (rc *ResourceScope) ReleaseBuffer(buf []byte) {
	rc.Lock()
	defer rc.Unlock()

	rc.memory -= int64(len(buf))

	// sanity check for bugs upstream
	if rc.memory < 0 {
		panic("BUG: too much memory released")
	}

	delete(rc.buffers, buf)
	pool.Put(buf)
}

func (rc *ResourceScope) Stat() network.ScopeStat {
	rc.Lock()
	defer rc.Unlock()

	return network.ScopeStat{Memory: rc.memory}
}
<commit_msg>add nil receiver safety to basic resource scopes<commit_after>package rcmgr

import (
	"fmt"
	"sync"

	"github.com/libp2p/go-buffer-pool"
	"github.com/libp2p/go-libp2p-core/network"
)

type ResourceScope struct {
	sync.Mutex

	limit   Limit
	memory  int64
	buffers map[interface{}][]byte
}

var _ network.ResourceScope = (*ResourceScope)(nil)

func (rc *ResourceScope) checkMemory(rsvp int) error {
	// overflow check; this also has the side-effect that we cannot reserve negative memory.
	newmem := rc.memory + int64(rsvp)
	if newmem < rc.memory {
		return fmt.Errorf("memory reservation overflow: %w", ErrResourceLimitExceeded)
	}

	// limit check
	if newmem > rc.limit.GetMemoryLimit() {
		return fmt.Errorf("cannot reserve memory: %w", ErrResourceLimitExceeded)
	}

	return nil
}

func (rc *ResourceScope) releaseBuffers() {
	for key, buf := range rc.buffers {
		pool.Put(buf)
		delete(rc.buffers, key)
	}
}

func (rc *ResourceScope) ReserveMemory(size int) error {
	if rc == nil {
		return nil
	}

	rc.Lock()
	defer rc.Unlock()

	if err := rc.checkMemory(size); err != nil {
		return err
	}

	rc.memory += int64(size)
	return nil
}

func (rc *ResourceScope) ReleaseMemory(size int) {
	if rc == nil {
		return
	}

	rc.Lock()
	defer rc.Unlock()

	rc.memory -= int64(size)

	// sanity check for bugs upstream
	if rc.memory < 0 {
		panic("BUG: too much memory released")
	}
}

func (rc *ResourceScope) GetBuffer(size int) ([]byte, error) {
	if rc == nil {
		return make([]byte, size), nil
	}

	rc.Lock()
	defer rc.Unlock()

	if err := rc.checkMemory(size); err != nil {
		return nil, err
	}

	buf := pool.Get(size)

	rc.memory += int64(size)
	rc.buffers[buf] = buf

	return buf, nil
}

func (rc *ResourceScope) GrowBuffer(oldbuf []byte, newsize int) ([]byte, error) {
	if rc == nil {
		newbuf := make([]byte, newsize)
		copy(newbuf, oldbuf)
		return newbuf, nil
	}

	rc.Lock()
	defer rc.Unlock()

	grow := newsize - len(oldbuf)
	if err := rc.checkMemory(grow); err != nil {
		return nil, err
	}

	newbuf := pool.Get(newsize)
	copy(newbuf, oldbuf)

	rc.memory += int64(grow)
	rc.buffers[newbuf] = newbuf
	delete(rc.buffers, oldbuf)

	return newbuf, nil
}

func (rc *ResourceScope) ReleaseBuffer(buf []byte) {
	rc.Lock()
	defer rc.Unlock()

	rc.memory -= int64(len(buf))

	// sanity check for bugs upstream
	if rc.memory < 0 {
		panic("BUG: too much memory released")
	}

	delete(rc.buffers, buf)
	pool.Put(buf)
}

func (rc *ResourceScope) Stat() network.ScopeStat {
	if rc == nil {
		return network.ScopeStat{}
	}

	rc.Lock()
	defer rc.Unlock()

	return network.ScopeStat{Memory: rc.memory}
}
<|endoftext|>


<commit_before>package imagemeta

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

const heifBoxHeaderSize = int64(8)

var heicBrand = []byte("heic")
var avifBrand = []byte("avif")
var heifPict = []byte("pict")

type heifData struct {
	Format        string
	Width, Height int64
}

func (d *heifData) IsFilled() bool {
	return len(d.Format) > 0 && d.Width > 0 && d.Height > 0
}

func heifReadBoxHeader(r io.Reader) (boxType string, boxDataSize int64, err error) {
	b := make([]byte, heifBoxHeaderSize)
	_, err = io.ReadFull(r, b)
	if err != nil {
		return
	}

	boxDataSize = int64(binary.BigEndian.Uint32(b[0:4])) - heifBoxHeaderSize
	boxType = string(b[4:8])

	return
}

func heifReadBoxData(r io.Reader, boxDataSize int64) (b []byte, err error) {
	b = make([]byte, boxDataSize)
	_, err = io.ReadFull(r, b)
	return
}

func heifAssignFormat(d *heifData, brand []byte) bool {
	if bytes.Equal(brand, heicBrand) {
		d.Format = "heic"
		return true
	}

	if bytes.Equal(brand, avifBrand) {
		d.Format = "avif"
		return true
	}

	return false
}

func heifReadFtyp(d *heifData, r io.Reader, boxDataSize int64) error {
	if boxDataSize < 8 {
		return errors.New("Invalid ftyp data")
	}

	data, err := heifReadBoxData(r, boxDataSize)
	if err != nil {
		return err
	}

	if heifAssignFormat(d, data[0:4]) {
		return nil
	}

	if boxDataSize >= 12 {
		for i := int64(8); i < boxDataSize; i += 4 {
			if heifAssignFormat(d, data[i:i+4]) {
				return nil
			}
		}
	}

	return errors.New("Image is not compatible with heic/avif")
}

func heifReadMeta(d *heifData, r io.Reader, boxDataSize int64) error {
	if boxDataSize < 4 {
		return errors.New("Invalid meta data")
	}

	if _, err := io.ReadFull(r, make([]byte, 4)); err != nil {
		return err
	}

	if boxDataSize > 4 {
		if err := heifReadBoxes(d, io.LimitReader(r, boxDataSize-4)); err != nil && err != io.EOF {
			return err
		}
	}

	return nil
}

func heifReadHldr(r io.Reader, boxDataSize int64) error {
	if boxDataSize < 12 {
		return errors.New("Invalid hdlr data")
	}

	data, err := heifReadBoxData(r, boxDataSize)
	if err != nil {
		return err
	}

	if !bytes.Equal(data[8:12], heifPict) {
		return fmt.Errorf("Invalid handler. Expected: pict, actual: %s", data[8:12])
	}

	return nil
}

func heifReadIspe(r io.Reader, boxDataSize int64) (w, h int64, err error) {
	if boxDataSize < 12 {
		return 0, 0, errors.New("Invalid ispe data")
	}

	data, err := heifReadBoxData(r, boxDataSize)
	if err != nil {
		return 0, 0, err
	}

	w = int64(binary.BigEndian.Uint32(data[4:8]))
	h = int64(binary.BigEndian.Uint32(data[8:12]))

	return
}

func heifReadBoxes(d *heifData, r io.Reader) error {
	for {
		boxType, boxDataSize, err := heifReadBoxHeader(r)

		if err != nil {
			return err
		}

		if boxDataSize < 0 {
			return errors.New("Invalid box data")
		}

		// log.Printf("Box type: %s; Box data size: %d", boxType, boxDataSize)

		switch boxType {
		case "ftyp":
			if err := heifReadFtyp(d, r, boxDataSize); err != nil {
				return err
			}
		case "meta":
			if err := heifReadMeta(d, r, boxDataSize); err != nil {
				return err
			}
			if !d.IsFilled() {
				return errors.New("Dimensions data wasn't found in meta box")
			}
			return nil
		case "hdlr":
			if err := heifReadHldr(r, boxDataSize); err != nil {
				return nil
			}
		case "iprp", "ipco":
			if err := heifReadBoxes(d, io.LimitReader(r, boxDataSize)); err != nil && err != io.EOF {
				return err
			}
		case "ispe":
			w, h, err := heifReadIspe(r, boxDataSize)
			if err != nil {
				return err
			}
			if w > d.Width || h > d.Height {
				d.Width, d.Height = w, h
			}
		case "mdat":
			return errors.New("mdat box occurred before meta box")
		default:
			if _, err := heifReadBoxData(r, boxDataSize); err != nil {
				return err
			}
		}
	}
}

func DecodeHeifMeta(r io.Reader) (Meta, error) {
	d := new(heifData)

	if err := heifReadBoxes(d, r); err != nil && !d.IsFilled() {
		return nil, err
	}

	return &meta{
		format: d.Format,
		width:  int(d.Width),
		height: int(d.Height),
	}, nil
}

func init() {
	RegisterFormat("????ftypheic", DecodeHeifMeta)
	RegisterFormat("????ftypheix", DecodeHeifMeta)
	RegisterFormat("????ftyphevc", DecodeHeifMeta)
	RegisterFormat("????ftypheim", DecodeHeifMeta)
	RegisterFormat("????ftypheis", DecodeHeifMeta)
	RegisterFormat("????ftyphevm", DecodeHeifMeta)
	RegisterFormat("????ftyphevs", DecodeHeifMeta)
	RegisterFormat("????ftypmif1", DecodeHeifMeta)
	RegisterFormat("????ftypavif", DecodeHeifMeta)
}
<commit_msg>Optimize HEIF meta parsing<commit_after>package imagemeta

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

const heifBoxHeaderSize = int64(8)

var heicBrand = []byte("heic")
var avifBrand = []byte("avif")
var heifPict = []byte("pict")

type heifData struct {
	Format        string
	Width, Height int64
}

func (d *heifData) IsFilled() bool {
	return len(d.Format) > 0 && d.Width > 0 && d.Height > 0
}

func heifReadN(r io.Reader, n int64) (b []byte, err error) {
	if buf, ok := r.(*bytes.Buffer); ok {
		b = buf.Next(int(n))
		if len(b) == 0 {
			return b, io.EOF
		}
		return b, nil
	}

	b = make([]byte, n)
	_, err = io.ReadFull(r, b)
	return
}

func heifDiscardN(r io.Reader, n int64) error {
	if buf, ok := r.(*bytes.Buffer); ok {
		_ = buf.Next(int(n))
		return nil
	}

	_, err := bufio.NewReader(r).Discard(int(n))
	return err
}

func heifReadBoxHeader(r io.Reader) (boxType string, boxDataSize int64, err error) {
	var b []byte

	b, err = heifReadN(r, heifBoxHeaderSize)
	if err != nil {
		return
	}

	boxDataSize = int64(binary.BigEndian.Uint32(b[0:4])) - heifBoxHeaderSize
	boxType = string(b[4:8])

	return
}

func heifAssignFormat(d *heifData, brand []byte) bool {
	if bytes.Equal(brand, heicBrand) {
		d.Format = "heic"
		return true
	}

	if bytes.Equal(brand, avifBrand) {
		d.Format = "avif"
		return true
	}

	return false
}

func heifReadFtyp(d *heifData, r io.Reader, boxDataSize int64) error {
	if boxDataSize < 8 {
		return errors.New("Invalid ftyp data")
	}

	data, err := heifReadN(r, boxDataSize)
	if err != nil {
		return err
	}

	if heifAssignFormat(d, data[0:4]) {
		return nil
	}

	if boxDataSize >= 12 {
		for i := int64(8); i < boxDataSize; i += 4 {
			if heifAssignFormat(d, data[i:i+4]) {
				return nil
			}
		}
	}

	return errors.New("Image is not compatible with heic/avif")
}

func heifReadMeta(d *heifData, r io.Reader, boxDataSize int64) error {
	if boxDataSize < 4 {
		return errors.New("Invalid meta data")
	}

	data, err := heifReadN(r, boxDataSize)
	if err != nil {
		return err
	}

	if boxDataSize > 4 {
		if err := heifReadBoxes(d, bytes.NewBuffer(data[4:])); err != nil && err != io.EOF {
			return err
		}
	}

	return nil
}

func heifReadHldr(r io.Reader, boxDataSize int64) error {
	if boxDataSize < 12 {
		return errors.New("Invalid hdlr data")
	}

	data, err := heifReadN(r, boxDataSize)
	if err != nil {
		return err
	}

	if !bytes.Equal(data[8:12], heifPict) {
		return fmt.Errorf("Invalid handler. Expected: pict, actual: %s", data[8:12])
	}

	return nil
}

func heifReadIspe(r io.Reader, boxDataSize int64) (w, h int64, err error) {
	if boxDataSize < 12 {
		return 0, 0, errors.New("Invalid ispe data")
	}

	data, err := heifReadN(r, boxDataSize)
	if err != nil {
		return 0, 0, err
	}

	w = int64(binary.BigEndian.Uint32(data[4:8]))
	h = int64(binary.BigEndian.Uint32(data[8:12]))

	return
}

func heifReadBoxes(d *heifData, r io.Reader) error {
	for {
		boxType, boxDataSize, err := heifReadBoxHeader(r)

		if err != nil {
			return err
		}

		if boxDataSize < 0 {
			return errors.New("Invalid box data")
		}

		// log.Printf("Box type: %s; Box data size: %d", boxType, boxDataSize)

		switch boxType {
		case "ftyp":
			if err := heifReadFtyp(d, r, boxDataSize); err != nil {
				return err
			}
		case "meta":
			if err := heifReadMeta(d, r, boxDataSize); err != nil {
				return err
			}
			if !d.IsFilled() {
				return errors.New("Dimensions data wasn't found in meta box")
			}
			return nil
		case "hdlr":
			if err := heifReadHldr(r, boxDataSize); err != nil {
				return nil
			}
		case "iprp", "ipco":
			data, err := heifReadN(r, boxDataSize)
			if err != nil {
				return err
			}

			if err := heifReadBoxes(d, bytes.NewBuffer(data)); err != nil && err != io.EOF {
				return err
			}
		case "ispe":
			w, h, err := heifReadIspe(r, boxDataSize)
			if err != nil {
				return err
			}
			if w > d.Width || h > d.Height {
				d.Width, d.Height = w, h
			}
		case "mdat":
			return errors.New("mdat box occurred before meta box")
		default:
			if err := heifDiscardN(r, boxDataSize); err != nil {
				return err
			}
		}
	}
}

func DecodeHeifMeta(r io.Reader) (Meta, error) {
	d := new(heifData)

	if err := heifReadBoxes(d, r); err != nil && !d.IsFilled() {
		return nil, err
	}

	return &meta{
		format: d.Format,
		width:  int(d.Width),
		height: int(d.Height),
	}, nil
}

func init() {
	RegisterFormat("????ftypheic", DecodeHeifMeta)
	RegisterFormat("????ftypheix", DecodeHeifMeta)
	RegisterFormat("????ftyphevc", DecodeHeifMeta)
	RegisterFormat("????ftypheim", DecodeHeifMeta)
	RegisterFormat("????ftypheis", DecodeHeifMeta)
	RegisterFormat("????ftyphevm", DecodeHeifMeta)
	RegisterFormat("????ftyphevs", DecodeHeifMeta)
	RegisterFormat("????ftypmif1", DecodeHeifMeta)
	RegisterFormat("????ftypavif", DecodeHeifMeta)
}
<|endoftext|>