<commit_before>package service

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/utils/featureflag"

	"github.com/juju/juju/feature"
	"github.com/juju/juju/service/common"
	"github.com/juju/juju/version"
)

// This exists to allow patching during tests.
var getVersion = func() version.Binary {
	return version.Current
}

// DiscoverService returns an interface to a service apropriate
// for the current system
func DiscoverService(name string, conf common.Conf) (Service, error) {
	initName, err := discoverInitSystem()
	if err != nil {
		return nil, errors.Trace(err)
	}

	service, err := NewService(name, conf, initName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return service, nil
}

func discoverInitSystem() (string, error) {
	initName, err := discoverLocalInitSystem()
	if errors.IsNotFound(err) {
		// Fall back to checking the juju version.
		jujuVersion := getVersion()
		versionInitName, ok := VersionInitSystem(jujuVersion)
		if !ok {
			// The key error is the one from discoverLocalInitSystem so
			// that is what we return.
			return "", errors.Trace(err)
		}
		initName = versionInitName
	} else if err != nil {
		return "", errors.Trace(err)
	}
	return initName, nil
}

// VersionInitSystem returns an init system name based on the provided
// version info. If one cannot be identified then false if returned
// for the second return value.
func VersionInitSystem(vers version.Binary) (string, bool) {
	initName, ok := versionInitSystem(vers)
	if !ok {
		logger.Errorf("could not identify init system from juju version info (%#v)", vers)
		return "", false
	}
	logger.Debugf("discovered init system %q from juju version info (%#v)", initName, vers)
	return initName, true
}

func versionInitSystem(vers version.Binary) (string, bool) {
	switch vers.OS {
	case version.Windows:
		return InitSystemWindows, true
	case version.Ubuntu:
		switch vers.Series {
		case "precise", "quantal", "raring", "saucy", "trusty", "utopic":
			return InitSystemUpstart, true
		case "":
			return "", false
		default:
			// Check for pre-precise releases.
			os, _ := version.GetOSFromSeries(vers.Series)
			if os == version.Unknown {
				return "", false
			}
			// vivid and later
			if featureflag.Enabled(feature.LegacyUpstart) {
				return InitSystemUpstart, true
			}
			return InitSystemSystemd, true
		}
		// TODO(ericsnow) Support other OSes, like version.CentOS.
	default:
		return "", false
	}
}

// pid1 is the path to the "file" that contains the path to the init
// system executable on linux.
const pid1 = "/proc/1/cmdline"

// These exist to allow patching during tests.
var (
	runtimeOS    = func() string { return runtime.GOOS }
	pid1Filename = func() string { return pid1 }
	evalSymlinks = filepath.EvalSymlinks

	initExecutable = func() (string, error) {
		pid1File := pid1Filename()
		data, err := ioutil.ReadFile(pid1File)
		if os.IsNotExist(err) {
			return "", errors.NotFoundf("init system (via %q)", pid1File)
		}
		if err != nil {
			return "", errors.Annotatef(err, "failed to identify init system (via %q)", pid1File)
		}
		executable := strings.Split(string(data), "\x00")[0]
		return executable, nil
	}
)

func discoverLocalInitSystem() (string, error) {
	if runtimeOS() == "windows" {
		return InitSystemWindows, nil
	}

	executable, err := initExecutable()
	if err != nil {
		return "", errors.Trace(err)
	}

	initName, ok := identifyInitSystem(executable)
	if !ok {
		return "", errors.NotFoundf("init system (based on %q)", executable)
	}
	logger.Debugf("discovered init system %q from executable %q", initName, executable)
	return initName, nil
}

func identifyInitSystem(executable string) (string, bool) {
	initSystem, ok := identifyExecutable(executable)
	if ok {
		return initSystem, true
	}

	// First fall back to following symlinks (if any).
	executable, err := evalSymlinks(executable)
	if err != nil {
		logger.Errorf("failed to find %q: %v", executable, err)
		return "", false
	}
	initSystem, ok = identifyExecutable(executable)
	if ok {
		return initSystem, true
	}

	// Fall back to checking the "version" text.
	cmd := exec.Command(executable, "--version")
	out, err := cmd.CombinedOutput()
	if err != nil {
		logger.Errorf(`"%s --version" failed (%v): %s`, executable, err, out)
		return "", false
	}

	verText := string(out)
	switch {
	case strings.Contains(verText, "upstart"):
		return InitSystemUpstart, true
	case strings.Contains(verText, "systemd"):
		return InitSystemSystemd, true
	}

	// uh-oh
	return "", false
}

func identifyExecutable(executable string) (string, bool) {
	switch {
	case strings.Contains(executable, "upstart"):
		return InitSystemUpstart, true
	case strings.Contains(executable, "systemd"):
		return InitSystemSystemd, true
	default:
		return "", false
	}
}

// TODO(ericsnow) Build this script more dynamically (using shell.Renderer).
// TODO(ericsnow) Use a case statement in the script?

// DiscoverInitSystemScript is the shell script to use when
// discovering the local init system.
const DiscoverInitSystemScript = `#!/usr/bin/env bash

function checkInitSystem() {
    if [[ $1 == *"systemd"* ]]; then
        echo -n systemd
        exit $?
    elif [[ $1 == *"upstart"* ]]; then
        echo -n upstart
        exit $?
    fi
}

# Find the executable.
executable=$(cat /proc/1/cmdline | awk -F"\0" '{print $1}')
if [[ $? -ne 0 ]]; then
    exit 1
fi

# Check the executable.
checkInitSystem "$executable"

# First fall back to following symlinks.
if [[ -L $executable ]]; then
    linked=$(readlink -f "$executable")
    if [[ $? -eq 0 ]]; then
        executable=$linked

        # Check the linked executable.
        checkInitSystem "$linked"
    fi
fi

# Fall back to checking the "version" text.
verText=$("${executable}" --version)
if [[ $? -eq 0 ]]; then
    checkInitSystem "$verText"
fi

# uh-oh
exit 1
`

func writeDiscoverInitSystemScript(filename string) []string {
	// TODO(ericsnow) Use utils.shell.Renderer.WriteScript.
	return []string{
		fmt.Sprintf(`
cat > %s << 'EOF'
%s
EOF`[1:], filename, DiscoverInitSystemScript),
		"chmod 0755 " + filename,
	}
}

const caseLine = "%sif [[ $%s == \"%s\" ]]; then %s\n"

// newShellSelectCommand creates a bash if statement with an if
// (or elif) clause for each of the executables in linuxExecutables.
// The body of each clause comes from calling the provided handler with
// the init system name. If the handler does not support the args then
// it returns a false "ok" value.
func newShellSelectCommand(envVarName string, handler func(string) (string, bool)) string {
	// TODO(ericsnow) Build the command in a better way?
	// TODO(ericsnow) Use a case statement?

	prefix := ""
	lines := ""
	for _, initSystem := range linuxInitSystems {
		cmd, ok := handler(initSystem)
		if !ok {
			continue
		}
		lines += fmt.Sprintf(caseLine, prefix, envVarName, initSystem, cmd)

		if prefix != "el" {
			prefix = "el"
		}
	}
	if lines != "" {
		lines += "" +
			"else exit 1\n" +
			"fi"
	}
	return lines
}
<commit_msg>Save the original executable for the error message.<commit_after>package service

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/juju/errors"
	"github.com/juju/utils/featureflag"

	"github.com/juju/juju/feature"
	"github.com/juju/juju/service/common"
	"github.com/juju/juju/version"
)

// This exists to allow patching during tests.
var getVersion = func() version.Binary {
	return version.Current
}

// DiscoverService returns an interface to a service apropriate
// for the current system
func DiscoverService(name string, conf common.Conf) (Service, error) {
	initName, err := discoverInitSystem()
	if err != nil {
		return nil, errors.Trace(err)
	}

	service, err := NewService(name, conf, initName)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return service, nil
}

func discoverInitSystem() (string, error) {
	initName, err := discoverLocalInitSystem()
	if errors.IsNotFound(err) {
		// Fall back to checking the juju version.
		jujuVersion := getVersion()
		versionInitName, ok := VersionInitSystem(jujuVersion)
		if !ok {
			// The key error is the one from discoverLocalInitSystem so
			// that is what we return.
			return "", errors.Trace(err)
		}
		initName = versionInitName
	} else if err != nil {
		return "", errors.Trace(err)
	}
	return initName, nil
}

// VersionInitSystem returns an init system name based on the provided
// version info. If one cannot be identified then false if returned
// for the second return value.
func VersionInitSystem(vers version.Binary) (string, bool) {
	initName, ok := versionInitSystem(vers)
	if !ok {
		logger.Errorf("could not identify init system from juju version info (%#v)", vers)
		return "", false
	}
	logger.Debugf("discovered init system %q from juju version info (%#v)", initName, vers)
	return initName, true
}

func versionInitSystem(vers version.Binary) (string, bool) {
	switch vers.OS {
	case version.Windows:
		return InitSystemWindows, true
	case version.Ubuntu:
		switch vers.Series {
		case "precise", "quantal", "raring", "saucy", "trusty", "utopic":
			return InitSystemUpstart, true
		case "":
			return "", false
		default:
			// Check for pre-precise releases.
			os, _ := version.GetOSFromSeries(vers.Series)
			if os == version.Unknown {
				return "", false
			}
			// vivid and later
			if featureflag.Enabled(feature.LegacyUpstart) {
				return InitSystemUpstart, true
			}
			return InitSystemSystemd, true
		}
		// TODO(ericsnow) Support other OSes, like version.CentOS.
	default:
		return "", false
	}
}

// pid1 is the path to the "file" that contains the path to the init
// system executable on linux.
const pid1 = "/proc/1/cmdline"

// These exist to allow patching during tests.
var (
	runtimeOS    = func() string { return runtime.GOOS }
	pid1Filename = func() string { return pid1 }
	evalSymlinks = filepath.EvalSymlinks

	initExecutable = func() (string, error) {
		pid1File := pid1Filename()
		data, err := ioutil.ReadFile(pid1File)
		if os.IsNotExist(err) {
			return "", errors.NotFoundf("init system (via %q)", pid1File)
		}
		if err != nil {
			return "", errors.Annotatef(err, "failed to identify init system (via %q)", pid1File)
		}
		executable := strings.Split(string(data), "\x00")[0]
		return executable, nil
	}
)

func discoverLocalInitSystem() (string, error) {
	if runtimeOS() == "windows" {
		return InitSystemWindows, nil
	}

	executable, err := initExecutable()
	if err != nil {
		return "", errors.Trace(err)
	}

	initName, ok := identifyInitSystem(executable)
	if !ok {
		return "", errors.NotFoundf("init system (based on %q)", executable)
	}
	logger.Debugf("discovered init system %q from executable %q", initName, executable)
	return initName, nil
}

func identifyInitSystem(executable string) (string, bool) {
	initSystem, ok := identifyExecutable(executable)
	if ok {
		return initSystem, true
	}

	// First fall back to following symlinks (if any).
	resolved, err := evalSymlinks(executable)
	if err != nil {
		logger.Errorf("failed to find %q: %v", executable, err)
		return "", false
	}
	executable = resolved
	initSystem, ok = identifyExecutable(executable)
	if ok {
		return initSystem, true
	}

	// Fall back to checking the "version" text.
	cmd := exec.Command(executable, "--version")
	out, err := cmd.CombinedOutput()
	if err != nil {
		logger.Errorf(`"%s --version" failed (%v): %s`, executable, err, out)
		return "", false
	}

	verText := string(out)
	switch {
	case strings.Contains(verText, "upstart"):
		return InitSystemUpstart, true
	case strings.Contains(verText, "systemd"):
		return InitSystemSystemd, true
	}

	// uh-oh
	return "", false
}

func identifyExecutable(executable string) (string, bool) {
	switch {
	case strings.Contains(executable, "upstart"):
		return InitSystemUpstart, true
	case strings.Contains(executable, "systemd"):
		return InitSystemSystemd, true
	default:
		return "", false
	}
}

// TODO(ericsnow) Build this script more dynamically (using shell.Renderer).
// TODO(ericsnow) Use a case statement in the script?

// DiscoverInitSystemScript is the shell script to use when
// discovering the local init system.
const DiscoverInitSystemScript = `#!/usr/bin/env bash

function checkInitSystem() {
    if [[ $1 == *"systemd"* ]]; then
        echo -n systemd
        exit $?
    elif [[ $1 == *"upstart"* ]]; then
        echo -n upstart
        exit $?
    fi
}

# Find the executable.
executable=$(cat /proc/1/cmdline | awk -F"\0" '{print $1}')
if [[ $? -ne 0 ]]; then
    exit 1
fi

# Check the executable.
checkInitSystem "$executable"

# First fall back to following symlinks.
if [[ -L $executable ]]; then
    linked=$(readlink -f "$executable")
    if [[ $? -eq 0 ]]; then
        executable=$linked

        # Check the linked executable.
        checkInitSystem "$linked"
    fi
fi

# Fall back to checking the "version" text.
verText=$("${executable}" --version)
if [[ $? -eq 0 ]]; then
    checkInitSystem "$verText"
fi

# uh-oh
exit 1
`

func writeDiscoverInitSystemScript(filename string) []string {
	// TODO(ericsnow) Use utils.shell.Renderer.WriteScript.
	return []string{
		fmt.Sprintf(`
cat > %s << 'EOF'
%s
EOF`[1:], filename, DiscoverInitSystemScript),
		"chmod 0755 " + filename,
	}
}

const caseLine = "%sif [[ $%s == \"%s\" ]]; then %s\n"

// newShellSelectCommand creates a bash if statement with an if
// (or elif) clause for each of the executables in linuxExecutables.
// The body of each clause comes from calling the provided handler with
// the init system name. If the handler does not support the args then
// it returns a false "ok" value.
func newShellSelectCommand(envVarName string, handler func(string) (string, bool)) string {
	// TODO(ericsnow) Build the command in a better way?
	// TODO(ericsnow) Use a case statement?

	prefix := ""
	lines := ""
	for _, initSystem := range linuxInitSystems {
		cmd, ok := handler(initSystem)
		if !ok {
			continue
		}
		lines += fmt.Sprintf(caseLine, prefix, envVarName, initSystem, cmd)

		if prefix != "el" {
			prefix = "el"
		}
	}
	if lines != "" {
		lines += "" +
			"else exit 1\n" +
			"fi"
	}
	return lines
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. 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 gcp

import (
	"cups-connector/lib"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/golang/oauth2"
)

const baseURL = "https://www.google.com/cloudprint/"

// Interface between Go and the Google Cloud Print API.
type GoogleCloudPrint struct {
	xmppJID        string
	xmppClient     *gcpXMPP
	robotTransport *oauth2.Transport
	userTransport  *oauth2.Transport
	proxyName      string
}

func NewGoogleCloudPrint(xmppJID, robotRefreshToken, userRefreshToken, proxyName string) (*GoogleCloudPrint, error) {
	robotTransport, err := newTransport(robotRefreshToken, lib.ScopeCloudPrint, lib.ScopeGoogleTalk)
	if err != nil {
		return nil, err
	}

	var userTransport *oauth2.Transport
	if userRefreshToken != "" {
		userTransport, err = newTransport(userRefreshToken, lib.ScopeCloudPrint)
		if err != nil {
			return nil, err
		}
	}

	gcp := &GoogleCloudPrint{xmppJID, nil, robotTransport, userTransport, proxyName}
	gcp.restartXMPP()
	return gcp, nil
}

func newTransport(refreshToken string, scopes ...string) (*oauth2.Transport, error) {
	options := &oauth2.Options{
		ClientID:     lib.ClientID,
		ClientSecret: lib.ClientSecret,
		RedirectURL:  lib.RedirectURL,
		Scopes:       scopes,
	}
	oauthConfig, err := oauth2.NewConfig(options, lib.AuthURL, lib.TokenURL)
	if err != nil {
		return nil, err
	}

	transport := oauthConfig.NewTransport()
	transport.SetToken(&oauth2.Token{RefreshToken: refreshToken})
	// Get first access token to be sure we can.
	if err = transport.RefreshToken(); err != nil {
		return nil, err
	}

	return transport, nil
}

func (gcp *GoogleCloudPrint) Quit() {
	gcp.xmppClient.quit()
}

func (gcp *GoogleCloudPrint) CanShare() bool {
	return gcp.userTransport != nil
}

// Tries to start an XMPP conversation multiple times, then panics.
func (gcp *GoogleCloudPrint) restartXMPP() {
	if gcp.xmppClient != nil {
		gcp.xmppClient.quit()
	}

	var err error
	for i := 0; i < 4; i++ {
		var xmpp *gcpXMPP
		xmpp, err = newXMPP(gcp.xmppJID, gcp.robotTransport.Token().AccessToken, gcp.proxyName)
		if err == nil {
			gcp.xmppClient = xmpp
			glog.Warning("Started XMPP successfully")
		}
		// Sleep for 1, 2, 4, 8 seconds.
		time.Sleep(time.Duration((i+1)*2) * time.Second)
	}
	glog.Fatalf("Failed to start XMPP conversation: %s", err)
	panic("unreachable")
}

// Waits for the next batch of jobs from GCP. Blocks until batch arrives.
//
// Calls google.com/cloudprint/fetch.
func (gcp *GoogleCloudPrint) NextJobBatch() ([]lib.Job, error) {
	printerIDb64, err := gcp.xmppClient.nextWaitingPrinter()
	if err != nil {
		if err == Closed {
			return nil, err
		}

		glog.Warningf("Restarting XMPP conversation because: %s", err)
		gcp.restartXMPP()

		// Now try again.
		printerIDb64, err = gcp.xmppClient.nextWaitingPrinter()
		if err != nil {
			if err == Closed {
				return nil, err
			}
			glog.Fatalf("Failed to wait for next printer twice: %s", err)
		}
	}

	printerIDbyte, err := base64.StdEncoding.DecodeString(printerIDb64)
	if err != nil {
		return nil, err
	}

	return gcp.Fetch(string(printerIDbyte))
}

// Calls google.com/cloudprint/control.
func (gcp *GoogleCloudPrint) Control(jobID string, status lib.GCPJobStatus, code, message string) error {
	form := url.Values{}
	form.Set("jobid", jobID)
	form.Set("status", string(status))
	form.Set("code", code)
	form.Set("message", message)

	if _, _, err := post(gcp.robotTransport, "control", form); err != nil {
		return err
	}

	return nil
}

// Calls google.com/cloudprint/delete.
func (gcp *GoogleCloudPrint) Delete(gcpID string) error {
	form := url.Values{}
	form.Set("printerid", gcpID)

	if _, _, err := post(gcp.robotTransport, "delete", form); err != nil {
		return err
	}

	return nil
}

// Gets the outstanding print jobs for a printer.
//
// Calls google.com/cloudprint/fetch.
func (gcp *GoogleCloudPrint) Fetch(gcpID string) ([]lib.Job, error) {
	form := url.Values{}
	form.Set("printerid", gcpID)

	responseBody, errorCode, err := post(gcp.robotTransport, "fetch", form)
	if err != nil {
		if errorCode == 413 {
			// 413 means "Zero print jobs returned", which isn't really an error.
			return []lib.Job{}, nil
		}
		return nil, err
	}

	var jobsData struct {
		Jobs []struct {
			ID        string
			FileURL   string
			TicketURL string
			OwnerID   string
		}
	}
	if err = json.Unmarshal(responseBody, &jobsData); err != nil {
		return nil, err
	}

	jobs := make([]lib.Job, 0, len(jobsData.Jobs))

	for _, jobData := range jobsData.Jobs {
		job := lib.Job{
			GCPPrinterID: gcpID,
			GCPJobID:     jobData.ID,
			FileURL:      jobData.FileURL,
			TicketURL:    jobData.TicketURL,
			OwnerID:      jobData.OwnerID,
		}
		jobs = append(jobs, job)
	}

	return jobs, nil
}

// Gets all GCP printers assigned to the configured proxy.
//
// Calls google.com/cloudprint/list.
func (gcp *GoogleCloudPrint) List() ([]lib.Printer, error) {
	form := url.Values{}
	form.Set("proxy", gcp.proxyName)

	responseBody, _, err := post(gcp.robotTransport, "list", form)
	if err != nil {
		return nil, err
	}

	var listData struct {
		Printers []struct {
			Id                 string
			Name               string
			DefaultDisplayName string
			Description        string
			Status             string
			CapsHash           string
			Tags               []string
		}
	}
	if err = json.Unmarshal(responseBody, &listData); err != nil {
		return nil, err
	}

	printers := make([]lib.Printer, 0, len(listData.Printers))
	for _, p := range listData.Printers {
		tags := make(map[string]string)
		for _, tag := range p.Tags {
			if !strings.HasPrefix(tag, "cups-") {
				continue
			}
			s := strings.SplitN(tag, "=", 2)
			key := s[0][5:]
			var value string
			if len(s) > 1 {
				value = s[1]
			}
			tags[key] = value
		}

		printer := lib.Printer{
			GCPID:              p.Id,
			Name:               p.Name,
			DefaultDisplayName: p.DefaultDisplayName,
			Description:        p.Description,
			Status:             lib.PrinterStatusFromString(p.Status),
			CapsHash:           p.CapsHash,
			Tags:               tags,
		}
		printers = append(printers, printer)
	}

	return printers, nil
}

// Registers a Google Cloud Print Printer. Sets the GCPID field in the printer arg.
//
// Calls google.com/cloudprint/register.
func (gcp *GoogleCloudPrint) Register(printer *lib.Printer, ppd string) error {
	if len(ppd) <= 0 {
		return errors.New("GCP requires a non-empty PPD")
	}

	form := url.Values{}
	form.Set("name", printer.Name)
	form.Set("default_display_name", printer.DefaultDisplayName)
	form.Set("proxy", gcp.proxyName)
	form.Set("capabilities", string(ppd))
	form.Set("description", printer.Description)
	form.Set("status", string(printer.Status))
	form.Set("capsHash", printer.CapsHash)
	form.Set("content_types", "application/pdf")
	for key, value := range printer.Tags {
		form.Add("tag", fmt.Sprintf("cups-%s=%s", key, value))
	}

	responseBody, _, err := post(gcp.robotTransport, "register", form)
	if err != nil {
		return err
	}

	var registerData struct {
		Printers []struct {
			Id string
		}
	}
	if err = json.Unmarshal(responseBody, &registerData); err != nil {
		return err
	}

	printer.GCPID = registerData.Printers[0].Id

	return nil
}

// Updates a Google Cloud Print Printer.
//
// Calls google.com/cloudprint/update.
func (gcp *GoogleCloudPrint) Update(diff *lib.PrinterDiff, ppd string) error {
	form := url.Values{}
	form.Set("printerid", diff.Printer.GCPID)
	form.Set("proxy", gcp.proxyName)

	// Ignore Name field because it never changes.
	if diff.DefaultDisplayNameChanged {
		form.Set("default_display_name", diff.Printer.DefaultDisplayName)
	}

	if diff.DescriptionChanged {
		form.Set("description", diff.Printer.Description)
	}

	if diff.StatusChanged {
		form.Set("status", string(diff.Printer.Status))
	}

	if diff.CapsHashChanged {
		form.Set("capsHash", diff.Printer.CapsHash)
		form.Set("capabilities", ppd)
	}

	if diff.TagsChanged {
		for key, value := range diff.Printer.Tags {
			form.Add("tag", fmt.Sprintf("cups-%s=%s", key, value))
		}
		form.Set("remove_tag", "^cups-.*")
	}

	if _, _, err := post(gcp.robotTransport, "update", form); err != nil {
		return err
	}

	return nil
}

// Shares a GCP printer.
//
// Calls google.com/cloudprint/share.
func (gcp *GoogleCloudPrint) Share(gcpID, shareScope string) error {
	if gcp.userTransport == nil {
		return errors.New("Cannot share because user OAuth credentials not provided.")
	}

	form := url.Values{}
	form.Set("printerid", gcpID)
	form.Set("scope", shareScope)
	form.Set("role", "USER")
	form.Set("skip_notification", "true")

	if _, _, err := post(gcp.userTransport, "share", form); err != nil {
		return err
	}

	return nil
}

// Downloads a url (print job) to a Writer.
func (gcp *GoogleCloudPrint) Download(dst io.Writer, url string) error {
	response, err := get(gcp.robotTransport, url)
	if err != nil {
		return err
	}

	_, err = io.Copy(dst, response.Body)
	if err != nil {
		return err
	}

	return nil
}

// Gets a ticket (job options), returns it as a map.
func (gcp *GoogleCloudPrint) Ticket(ticketURL string) (map[string]string, error) {
	response, err := get(gcp.robotTransport, ticketURL)
	if err != nil {
		return nil, err
	}

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	var m map[string]string
	err = json.Unmarshal(responseBody, &m)
	if err != nil {
		return nil, err
	}

	return m, nil
}

// GETs to a URL. Returns the response object, in case the body is very large.
func get(t *oauth2.Transport, url string) (*http.Response, error) {
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("X-CloudPrint-Proxy", "cups-cloudprint-"+runtime.GOOS)

	response, err := t.RoundTrip(request)
	if err != nil {
		return nil, err
	}
	if response.StatusCode != 200 {
		return nil, fmt.Errorf("GET failed: %s %s", url, response.Status)
	}

	return response, nil
}

// POSTs to a GCP method. Returns the body of the response.
//
// On error, the last two return values are non-zero values.
func post(t *oauth2.Transport, method string, form url.Values) ([]byte, uint, error) {
	requestBody := strings.NewReader(form.Encode())
	request, err := http.NewRequest("POST", baseURL+method, requestBody)
	if err != nil {
		return nil, 0, err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Set("X-CloudPrint-Proxy", "cups-cloudprint-"+runtime.GOOS)

	response, err := t.RoundTrip(request)
	if err != nil {
		return nil, 0, err
	}
	if response.StatusCode != 200 {
		return nil, 0, fmt.Errorf("/%s call failed: %s", method, response.Status)
	}

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, 0, err
	}

	var responseStatus struct {
		Success   bool
		Message   string
		ErrorCode uint
	}
	if err = json.Unmarshal(responseBody, &responseStatus); err != nil {
		return nil, 0, err
	}
	if !responseStatus.Success {
		return nil, responseStatus.ErrorCode, fmt.Errorf(
			"/%s call failed: %s", method, responseStatus.Message)
	}

	return responseBody, 0, nil
}
<commit_msg>One-liner to fix reset function<commit_after>/*
Copyright 2014 Google Inc. 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 gcp

import (
	"cups-connector/lib"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"

	"github.com/golang/glog"
	"github.com/golang/oauth2"
)

const baseURL = "https://www.google.com/cloudprint/"

// Interface between Go and the Google Cloud Print API.
type GoogleCloudPrint struct {
	xmppJID        string
	xmppClient     *gcpXMPP
	robotTransport *oauth2.Transport
	userTransport  *oauth2.Transport
	proxyName      string
}

func NewGoogleCloudPrint(xmppJID, robotRefreshToken, userRefreshToken, proxyName string) (*GoogleCloudPrint, error) {
	robotTransport, err := newTransport(robotRefreshToken, lib.ScopeCloudPrint, lib.ScopeGoogleTalk)
	if err != nil {
		return nil, err
	}

	var userTransport *oauth2.Transport
	if userRefreshToken != "" {
		userTransport, err = newTransport(userRefreshToken, lib.ScopeCloudPrint)
		if err != nil {
			return nil, err
		}
	}

	gcp := &GoogleCloudPrint{xmppJID, nil, robotTransport, userTransport, proxyName}
	gcp.restartXMPP()
	return gcp, nil
}

func newTransport(refreshToken string, scopes ...string) (*oauth2.Transport, error) {
	options := &oauth2.Options{
		ClientID:     lib.ClientID,
		ClientSecret: lib.ClientSecret,
		RedirectURL:  lib.RedirectURL,
		Scopes:       scopes,
	}
	oauthConfig, err := oauth2.NewConfig(options, lib.AuthURL, lib.TokenURL)
	if err != nil {
		return nil, err
	}

	transport := oauthConfig.NewTransport()
	transport.SetToken(&oauth2.Token{RefreshToken: refreshToken})
	// Get first access token to be sure we can.
	if err = transport.RefreshToken(); err != nil {
		return nil, err
	}

	return transport, nil
}

func (gcp *GoogleCloudPrint) Quit() {
	gcp.xmppClient.quit()
}

func (gcp *GoogleCloudPrint) CanShare() bool {
	return gcp.userTransport != nil
}

// Tries to start an XMPP conversation multiple times, then panics.
func (gcp *GoogleCloudPrint) restartXMPP() {
	if gcp.xmppClient != nil {
		gcp.xmppClient.quit()
	}

	var err error
	for i := 0; i < 4; i++ {
		var xmpp *gcpXMPP
		xmpp, err = newXMPP(gcp.xmppJID, gcp.robotTransport.Token().AccessToken, gcp.proxyName)
		if err == nil {
			gcp.xmppClient = xmpp
			glog.Warning("Started XMPP successfully")
			return
		}
		// Sleep for 1, 2, 4, 8 seconds.
		time.Sleep(time.Duration((i+1)*2) * time.Second)
	}
	glog.Fatalf("Failed to start XMPP conversation: %s", err)
	panic("unreachable")
}

// Waits for the next batch of jobs from GCP. Blocks until batch arrives.
//
// Calls google.com/cloudprint/fetch.
func (gcp *GoogleCloudPrint) NextJobBatch() ([]lib.Job, error) {
	printerIDb64, err := gcp.xmppClient.nextWaitingPrinter()
	if err != nil {
		if err == Closed {
			return nil, err
		}

		glog.Warningf("Restarting XMPP conversation because: %s", err)
		gcp.restartXMPP()

		// Now try again.
		printerIDb64, err = gcp.xmppClient.nextWaitingPrinter()
		if err != nil {
			if err == Closed {
				return nil, err
			}
			glog.Fatalf("Failed to wait for next printer twice: %s", err)
		}
	}

	printerIDbyte, err := base64.StdEncoding.DecodeString(printerIDb64)
	if err != nil {
		return nil, err
	}

	return gcp.Fetch(string(printerIDbyte))
}

// Calls google.com/cloudprint/control.
func (gcp *GoogleCloudPrint) Control(jobID string, status lib.GCPJobStatus, code, message string) error {
	form := url.Values{}
	form.Set("jobid", jobID)
	form.Set("status", string(status))
	form.Set("code", code)
	form.Set("message", message)

	if _, _, err := post(gcp.robotTransport, "control", form); err != nil {
		return err
	}

	return nil
}

// Calls google.com/cloudprint/delete.
func (gcp *GoogleCloudPrint) Delete(gcpID string) error {
	form := url.Values{}
	form.Set("printerid", gcpID)

	if _, _, err := post(gcp.robotTransport, "delete", form); err != nil {
		return err
	}

	return nil
}

// Gets the outstanding print jobs for a printer.
//
// Calls google.com/cloudprint/fetch.
func (gcp *GoogleCloudPrint) Fetch(gcpID string) ([]lib.Job, error) {
	form := url.Values{}
	form.Set("printerid", gcpID)

	responseBody, errorCode, err := post(gcp.robotTransport, "fetch", form)
	if err != nil {
		if errorCode == 413 {
			// 413 means "Zero print jobs returned", which isn't really an error.
			return []lib.Job{}, nil
		}
		return nil, err
	}

	var jobsData struct {
		Jobs []struct {
			ID        string
			FileURL   string
			TicketURL string
			OwnerID   string
		}
	}
	if err = json.Unmarshal(responseBody, &jobsData); err != nil {
		return nil, err
	}

	jobs := make([]lib.Job, 0, len(jobsData.Jobs))

	for _, jobData := range jobsData.Jobs {
		job := lib.Job{
			GCPPrinterID: gcpID,
			GCPJobID:     jobData.ID,
			FileURL:      jobData.FileURL,
			TicketURL:    jobData.TicketURL,
			OwnerID:      jobData.OwnerID,
		}
		jobs = append(jobs, job)
	}

	return jobs, nil
}

// Gets all GCP printers assigned to the configured proxy.
//
// Calls google.com/cloudprint/list.
func (gcp *GoogleCloudPrint) List() ([]lib.Printer, error) {
	form := url.Values{}
	form.Set("proxy", gcp.proxyName)

	responseBody, _, err := post(gcp.robotTransport, "list", form)
	if err != nil {
		return nil, err
	}

	var listData struct {
		Printers []struct {
			Id                 string
			Name               string
			DefaultDisplayName string
			Description        string
			Status             string
			CapsHash           string
			Tags               []string
		}
	}
	if err = json.Unmarshal(responseBody, &listData); err != nil {
		return nil, err
	}

	printers := make([]lib.Printer, 0, len(listData.Printers))
	for _, p := range listData.Printers {
		tags := make(map[string]string)
		for _, tag := range p.Tags {
			if !strings.HasPrefix(tag, "cups-") {
				continue
			}
			s := strings.SplitN(tag, "=", 2)
			key := s[0][5:]
			var value string
			if len(s) > 1 {
				value = s[1]
			}
			tags[key] = value
		}

		printer := lib.Printer{
			GCPID:              p.Id,
			Name:               p.Name,
			DefaultDisplayName: p.DefaultDisplayName,
			Description:        p.Description,
			Status:             lib.PrinterStatusFromString(p.Status),
			CapsHash:           p.CapsHash,
			Tags:               tags,
		}
		printers = append(printers, printer)
	}

	return printers, nil
}

// Registers a Google Cloud Print Printer. Sets the GCPID field in the printer arg.
//
// Calls google.com/cloudprint/register.
func (gcp *GoogleCloudPrint) Register(printer *lib.Printer, ppd string) error {
	if len(ppd) <= 0 {
		return errors.New("GCP requires a non-empty PPD")
	}

	form := url.Values{}
	form.Set("name", printer.Name)
	form.Set("default_display_name", printer.DefaultDisplayName)
	form.Set("proxy", gcp.proxyName)
	form.Set("capabilities", string(ppd))
	form.Set("description", printer.Description)
	form.Set("status", string(printer.Status))
	form.Set("capsHash", printer.CapsHash)
	form.Set("content_types", "application/pdf")
	for key, value := range printer.Tags {
		form.Add("tag", fmt.Sprintf("cups-%s=%s", key, value))
	}

	responseBody, _, err := post(gcp.robotTransport, "register", form)
	if err != nil {
		return err
	}

	var registerData struct {
		Printers []struct {
			Id string
		}
	}
	if err = json.Unmarshal(responseBody, &registerData); err != nil {
		return err
	}

	printer.GCPID = registerData.Printers[0].Id

	return nil
}

// Updates a Google Cloud Print Printer.
//
// Calls google.com/cloudprint/update.
func (gcp *GoogleCloudPrint) Update(diff *lib.PrinterDiff, ppd string) error {
	form := url.Values{}
	form.Set("printerid", diff.Printer.GCPID)
	form.Set("proxy", gcp.proxyName)

	// Ignore Name field because it never changes.
	if diff.DefaultDisplayNameChanged {
		form.Set("default_display_name", diff.Printer.DefaultDisplayName)
	}

	if diff.DescriptionChanged {
		form.Set("description", diff.Printer.Description)
	}

	if diff.StatusChanged {
		form.Set("status", string(diff.Printer.Status))
	}

	if diff.CapsHashChanged {
		form.Set("capsHash", diff.Printer.CapsHash)
		form.Set("capabilities", ppd)
	}

	if diff.TagsChanged {
		for key, value := range diff.Printer.Tags {
			form.Add("tag", fmt.Sprintf("cups-%s=%s", key, value))
		}
		form.Set("remove_tag", "^cups-.*")
	}

	if _, _, err := post(gcp.robotTransport, "update", form); err != nil {
		return err
	}

	return nil
}

// Shares a GCP printer.
//
// Calls google.com/cloudprint/share.
func (gcp *GoogleCloudPrint) Share(gcpID, shareScope string) error {
	if gcp.userTransport == nil {
		return errors.New("Cannot share because user OAuth credentials not provided.")
	}

	form := url.Values{}
	form.Set("printerid", gcpID)
	form.Set("scope", shareScope)
	form.Set("role", "USER")
	form.Set("skip_notification", "true")

	if _, _, err := post(gcp.userTransport, "share", form); err != nil {
		return err
	}

	return nil
}

// Downloads a url (print job) to a Writer.
func (gcp *GoogleCloudPrint) Download(dst io.Writer, url string) error {
	response, err := get(gcp.robotTransport, url)
	if err != nil {
		return err
	}

	_, err = io.Copy(dst, response.Body)
	if err != nil {
		return err
	}

	return nil
}

// Gets a ticket (job options), returns it as a map.
func (gcp *GoogleCloudPrint) Ticket(ticketURL string) (map[string]string, error) {
	response, err := get(gcp.robotTransport, ticketURL)
	if err != nil {
		return nil, err
	}

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	var m map[string]string
	err = json.Unmarshal(responseBody, &m)
	if err != nil {
		return nil, err
	}

	return m, nil
}

// GETs to a URL. Returns the response object, in case the body is very large.
func get(t *oauth2.Transport, url string) (*http.Response, error) {
	request, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	request.Header.Set("X-CloudPrint-Proxy", "cups-cloudprint-"+runtime.GOOS)

	response, err := t.RoundTrip(request)
	if err != nil {
		return nil, err
	}
	if response.StatusCode != 200 {
		return nil, fmt.Errorf("GET failed: %s %s", url, response.Status)
	}

	return response, nil
}

// POSTs to a GCP method. Returns the body of the response.
//
// On error, the last two return values are non-zero values.
func post(t *oauth2.Transport, method string, form url.Values) ([]byte, uint, error) {
	requestBody := strings.NewReader(form.Encode())
	request, err := http.NewRequest("POST", baseURL+method, requestBody)
	if err != nil {
		return nil, 0, err
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	request.Header.Set("X-CloudPrint-Proxy", "cups-cloudprint-"+runtime.GOOS)

	response, err := t.RoundTrip(request)
	if err != nil {
		return nil, 0, err
	}
	if response.StatusCode != 200 {
		return nil, 0, fmt.Errorf("/%s call failed: %s", method, response.Status)
	}

	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, 0, err
	}

	var responseStatus struct {
		Success   bool
		Message   string
		ErrorCode uint
	}
	if err = json.Unmarshal(responseBody, &responseStatus); err != nil {
		return nil, 0, err
	}
	if !responseStatus.Success {
		return nil, responseStatus.ErrorCode, fmt.Errorf(
			"/%s call failed: %s", method, responseStatus.Message)
	}

	return responseBody, 0, nil
}
<|endoftext|>


<commit_before>package thumbnailer

import (
	"crypto/sha1"
	"fmt"
	"github.com/abustany/gollery/utils"
	"github.com/gographics/imagick/imagick"
	"github.com/howeyc/fsnotify"
	"github.com/robfig/revel"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	THUMBNAILER_SIZE_PX             = 200
	THUMBNAILER_COMPRESSION_QUALITY = 75
)

type Thumbnailer struct {
	RootDir           string
	CacheDir          string
	monitor           *fsnotify.Watcher
	thumbnailingQueue chan string
	queuedItems       map[string]bool
	queuedItemsMutex  sync.Mutex
}

func init() {
	imagick.Initialize()
}

func makeCacheKey(path string) string {
	h := sha1.New()
	io.WriteString(h, path)
	key := fmt.Sprintf("%.0x", h.Sum(nil))
	return key
}

// returns the list of thumb keys from this directory
func (t *Thumbnailer) checkCacheDir(dirPath string) ([]string, error) {
	dirFd, err := os.Open(dirPath)

	revel.INFO.Printf("Cleaning cache for directory '%s'", dirPath)

	if err != nil {
		return nil, utils.WrapError(err, "Cannot open directory '%s'", dirPath)
	}

	defer dirFd.Close()

	fis, err := dirFd.Readdir(-1)

	if err == io.EOF {
		return nil, nil
	}

	if err != nil {
		return nil, utils.WrapError(err, "Cannot read directory '%s'", dirPath)
	}

	thumbsToCreate := []string{}
	allThumbKeys := []string{}

	for _, f := range fis {
		fPath := path.Join(dirPath, f.Name())

		if f.IsDir() {
			childThumbKeys, err := t.checkCacheDir(fPath)

			if err != nil {
				revel.WARN.Printf("Cannot clean cache directory '%s': %s (skipping)", fPath, err)
			}

			allThumbKeys = append(allThumbKeys, childThumbKeys...)

			continue
		}

		fId := fPath[1+len(t.RootDir):]

		revel.TRACE.Printf("Checking thumbnail for %s", fId)

		cacheKey := makeCacheKey(fId)

		allThumbKeys = append(allThumbKeys, cacheKey)

		cacheFilePath := path.Join(t.CacheDir, cacheKey)

		_, err := os.Stat(cacheFilePath)

		if os.IsNotExist(err) {
			thumbsToCreate = append(thumbsToCreate, fPath)
			continue
		}

		if err != nil {
			revel.WARN.Printf("Error while checking thumbnail for '%s': %s (skipping)", fPath, err)
			continue
		}
	}

	for _, x := range thumbsToCreate {
		t.ScheduleThumbnail(x)
	}

	return allThumbKeys, nil
}

func (t *Thumbnailer) fileMonitorRoutine() {
	for {
		select {
		case ev := <-t.monitor.Event:
			revel.TRACE.Printf("Thumbnailer: file event: %s", ev)

			basename := path.Base(ev.Name)

			if len(basename) > 0 && basename[0] == '.' {
				revel.TRACE.Printf("Skipping event for hidden file %s", ev.Name)
				continue
			}

			if ev.IsDelete() {
				// We can't know if it was a directory or a file... Try to remove anyway
				// I assume if the directory gets deleted, the monitor goes away?
				t.DeleteThumbnail(ev.Name)
				continue
			}

			info, err := os.Stat(ev.Name)

			if err != nil {
				revel.ERROR.Printf("Cannot get file information for %s: %s", ev.Name, err)
				continue
			}

			if info.Mode().IsRegular() {
				t.ScheduleThumbnail(ev.Name)
			} else if info.IsDir() {
				if !ev.IsCreate() {
					continue
				}

				err := t.setupDirMonitor(ev.Name, info, nil)

				if err != nil {
					revel.ERROR.Printf("Cannot setup a file monitor on %s: %s", ev.Name, err)
				}
			} else {
				revel.INFO.Printf("Skipping file event for file of unknown type %s", ev.Name)
			}
		case err := <-t.monitor.Error:
			revel.ERROR.Printf("Thumbnailer: file monitoring error: %s", err)
		}
	}
}

func (t *Thumbnailer) thumbnailQueueRoutine() {
	for filePath := range t.thumbnailingQueue {
		t.queuedItemsMutex.Lock()
		delete(t.queuedItems, filePath)
		t.queuedItemsMutex.Unlock()

		err := t.CreateThumbnail(filePath)

		if err != nil {
			revel.ERROR.Printf("Couldn't create thumbnail for file '%s': %s", filePath, err)
		}

		revel.INFO.Printf("The thumbnailing queue now has %d items", len(t.thumbnailingQueue))
	}
}

func NewThumbnailer(rootDir string, cacheDir string) (*Thumbnailer, error) {
	t := &Thumbnailer{
		RootDir:           rootDir,
		CacheDir:          cacheDir,
		thumbnailingQueue: make(chan string, 256),
		queuedItems:       make(map[string]bool, 256),
	}

	var err error

	t.monitor, err = fsnotify.NewWatcher()

	if err != nil {
		return nil, utils.WrapError(err, "Cannot initialize file monitoring")
	}

	go t.fileMonitorRoutine()
	go t.thumbnailQueueRoutine()

	return t, nil
}

// Creates missing thumbnails
func (t *Thumbnailer) CheckCache() error {
	allThumbKeys, err := t.checkCacheDir(t.RootDir)

	if err != nil {
		return utils.WrapError(err, "Cannot check thumbnail cache")
	}

	keyHash := make(map[string]bool, len(allThumbKeys))

	for _, key := range allThumbKeys {
		keyHash[key] = true
	}

	fd, err := os.Open(t.CacheDir)

	if err != nil {
		return utils.WrapError(err, "Cannot check thumbnail cache for stall thumbnails")
	}

	defer fd.Close()

	for {
		// Don't read all files at all, it might be a lot
		fis, err := fd.Readdir(1024)

		if err == io.EOF {
			break
		}

		if err != nil {
			return utils.WrapError(err, "Cannot list thumbnails while cleaning cache")
		}

		for _, fi := range fis {
			// Thumbnail corresponds to a known picture, leave it alone
			if _, exists := keyHash[fi.Name()]; exists {
				continue
			}

			thumbPath := path.Join(t.CacheDir, fi.Name())

			revel.INFO.Printf("Removing stale thumbnail with key %s", fi.Name())
			err = os.Remove(thumbPath)

			if err != nil && !os.IsNotExist(err) {
				return utils.WrapError(err, "Cannot delete thumbnail '%s' while cleaning up cache", thumbPath)
			}
		}
	}

	return nil
}

func (t *Thumbnailer) setupDirMonitor(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}

	if strings.HasPrefix(path, ".") {
		return filepath.SkipDir
	}

	if !info.IsDir() {
		return nil
	}

	err = t.monitor.Watch(path)

	if err == nil {
		revel.INFO.Printf("Setup a directory monitor on '%s'", path)
	}

	return err
}

func (t *Thumbnailer) SetupMonitors() error {
	return utils.WrapError(filepath.Walk(t.RootDir, t.setupDirMonitor), "Cannot setup monitoring")
}

func (t *Thumbnailer) ScheduleThumbnail(filePath string) {
	revel.INFO.Printf("Scheduling thumbnailing of file %s", filePath)

	t.queuedItemsMutex.Lock()
	defer t.queuedItemsMutex.Unlock()

	if _, alreadyQueued := t.queuedItems[filePath]; alreadyQueued {
		revel.INFO.Printf("Thumbnailing already scheduled for file %s", filePath)
		return
	}

	t.thumbnailingQueue <- filePath
	t.queuedItems[filePath] = true
}

// For absolute paths, check that they are in the root dir
// For relative paths, prepend the root dir path
func (t *Thumbnailer) normalizePath(filePath string) (string, error) {
	if len(filePath) > 0 && filePath[0] == '/' {
		if !strings.HasPrefix(filePath, t.RootDir) {
			return "", fmt.Errorf("Not creating a thumbnail for a file outside the root directory: %s", filePath)
		}

		return filePath, nil
	}

	return path.Join(t.RootDir, filePath), nil
}

func (t *Thumbnailer) CreateThumbnail(filePath string) error {
	normalizedPath, err := t.normalizePath(filePath)

	if err != nil {
		return utils.WrapError(err, "Invalid path '%s'", filePath)
	}

	fileId := normalizedPath[1+len(t.RootDir):]

	startTime := time.Now()

	mw := imagick.NewMagickWand()
	defer mw.Destroy()

	err = mw.ReadImage(normalizedPath)

	if err != nil {
		return utils.WrapError(err, "Cannot read file '%s'", normalizedPath)
	}

	thumbKey := makeCacheKey(fileId)
	thumbPath := path.Join(t.CacheDir, thumbKey)

	width := mw.GetImageWidth()
	height := mw.GetImageHeight()
	var scale float32

	if width > height {
		scale = float32(THUMBNAILER_SIZE_PX) / float32(width)
		width = THUMBNAILER_SIZE_PX
		height = uint(float32(height) * scale)
	} else {
		scale = float32(THUMBNAILER_SIZE_PX) / float32(height)
		height = THUMBNAILER_SIZE_PX
		width = uint(float32(width) * scale)
	}

	// TRIANGLE is a simple linear interpolation, should be fast enough
	err = mw.ResizeImage(width, height, imagick.FILTER_TRIANGLE, 1)

	if err != nil {
		return utils.WrapError(err, "Cannot generate thumbnail for file '%s'", normalizedPath)
	}

	err = mw.SetCompressionQuality(THUMBNAILER_COMPRESSION_QUALITY)

	if err != nil {
		return utils.WrapError(err, "Cannot set compression quality for file '%s'", normalizedPath)
	}

	err = mw.WriteImage(thumbPath)

	if err != nil {
		return utils.WrapError(err, "Cannot write thumbnail '%s' for file '%s'", thumbPath, normalizedPath)
	}

	revel.INFO.Printf("Thumbnailed image '%s' as '%s' in %.2f seconds", normalizedPath, thumbPath, time.Now().Sub(startTime).Seconds())

	return nil
}

func (t *Thumbnailer) DeleteThumbnail(filePath string) error {
	normalizedPath, err := t.normalizePath(filePath)

	if err != nil {
		return utils.WrapError(err, "Invalid path '%s'", normalizedPath)
	}

	fileId := normalizedPath[1+len(t.RootDir):]
	thumbKey := makeCacheKey(fileId)
	thumbPath := path.Join(t.CacheDir, thumbKey)

	err = os.Remove(thumbPath)

	if os.IsNotExist(err) {
		return nil
	}

	if err != nil {
		return utils.WrapError(err, "Cannot remove thumbnail")
	}

	revel.INFO.Printf("Deleted thumbnail for image '%s'", normalizedPath)

	return nil
}

func (t *Thumbnailer) ThumbnailQueueSize() int {
	return len(t.queuedItems)
}
<commit_msg>thumbnailer: Consider renames as deletes<commit_after>package thumbnailer

import (
	"crypto/sha1"
	"fmt"
	"github.com/abustany/gollery/utils"
	"github.com/gographics/imagick/imagick"
	"github.com/howeyc/fsnotify"
	"github.com/robfig/revel"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	THUMBNAILER_SIZE_PX             = 200
	THUMBNAILER_COMPRESSION_QUALITY = 75
)

type Thumbnailer struct {
	RootDir           string
	CacheDir          string
	monitor           *fsnotify.Watcher
	thumbnailingQueue chan string
	queuedItems       map[string]bool
	queuedItemsMutex  sync.Mutex
}

func init() {
	imagick.Initialize()
}

func makeCacheKey(path string) string {
	h := sha1.New()
	io.WriteString(h, path)
	key := fmt.Sprintf("%.0x", h.Sum(nil))
	return key
}

// returns the list of thumb keys from this directory
func (t *Thumbnailer) checkCacheDir(dirPath string) ([]string, error) {
	dirFd, err := os.Open(dirPath)

	revel.INFO.Printf("Cleaning cache for directory '%s'", dirPath)

	if err != nil {
		return nil, utils.WrapError(err, "Cannot open directory '%s'", dirPath)
	}

	defer dirFd.Close()

	fis, err := dirFd.Readdir(-1)

	if err == io.EOF {
		return nil, nil
	}

	if err != nil {
		return nil, utils.WrapError(err, "Cannot read directory '%s'", dirPath)
	}

	thumbsToCreate := []string{}
	allThumbKeys := []string{}

	for _, f := range fis {
		fPath := path.Join(dirPath, f.Name())

		if f.IsDir() {
			childThumbKeys, err := t.checkCacheDir(fPath)

			if err != nil {
				revel.WARN.Printf("Cannot clean cache directory '%s': %s (skipping)", fPath, err)
			}

			allThumbKeys = append(allThumbKeys, childThumbKeys...)

			continue
		}

		fId := fPath[1+len(t.RootDir):]

		revel.TRACE.Printf("Checking thumbnail for %s", fId)

		cacheKey := makeCacheKey(fId)

		allThumbKeys = append(allThumbKeys, cacheKey)

		cacheFilePath := path.Join(t.CacheDir, cacheKey)

		_, err := os.Stat(cacheFilePath)

		if os.IsNotExist(err) {
			thumbsToCreate = append(thumbsToCreate, fPath)
			continue
		}

		if err != nil {
			revel.WARN.Printf("Error while checking thumbnail for '%s': %s (skipping)", fPath, err)
			continue
		}
	}

	for _, x := range thumbsToCreate {
		t.ScheduleThumbnail(x)
	}

	return allThumbKeys, nil
}

func (t *Thumbnailer) fileMonitorRoutine() {
	for {
		select {
		case ev := <-t.monitor.Event:
			revel.TRACE.Printf("Thumbnailer: file event: %s", ev)

			basename := path.Base(ev.Name)

			if len(basename) > 0 && basename[0] == '.' {
				revel.TRACE.Printf("Skipping event for hidden file %s", ev.Name)
				continue
			}

			if ev.IsDelete() || ev.IsRename() {
				// We can't know if it was a directory or a file... Try to remove anyway
				// I assume if the directory gets deleted, the monitor goes away?
				t.DeleteThumbnail(ev.Name)
				continue
			}

			info, err := os.Stat(ev.Name)

			if err != nil {
				revel.ERROR.Printf("Cannot get file information for %s: %s", ev.Name, err)
				continue
			}

			if info.Mode().IsRegular() {
				t.ScheduleThumbnail(ev.Name)
			} else if info.IsDir() {
				if !ev.IsCreate() {
					continue
				}

				err := t.setupDirMonitor(ev.Name, info, nil)

				if err != nil {
					revel.ERROR.Printf("Cannot setup a file monitor on %s: %s", ev.Name, err)
				}
			} else {
				revel.INFO.Printf("Skipping file event for file of unknown type %s", ev.Name)
			}
		case err := <-t.monitor.Error:
			revel.ERROR.Printf("Thumbnailer: file monitoring error: %s", err)
		}
	}
}

func (t *Thumbnailer) thumbnailQueueRoutine() {
	for filePath := range t.thumbnailingQueue {
		t.queuedItemsMutex.Lock()
		delete(t.queuedItems, filePath)
		t.queuedItemsMutex.Unlock()

		err := t.CreateThumbnail(filePath)

		if err != nil {
			revel.ERROR.Printf("Couldn't create thumbnail for file '%s': %s", filePath, err)
		}

		revel.INFO.Printf("The thumbnailing queue now has %d items", len(t.thumbnailingQueue))
	}
}

func NewThumbnailer(rootDir string, cacheDir string) (*Thumbnailer, error) {
	t := &Thumbnailer{
		RootDir:           rootDir,
		CacheDir:          cacheDir,
		thumbnailingQueue: make(chan string, 256),
		queuedItems:       make(map[string]bool, 256),
	}

	var err error

	t.monitor, err = fsnotify.NewWatcher()

	if err != nil {
		return nil, utils.WrapError(err, "Cannot initialize file monitoring")
	}

	go t.fileMonitorRoutine()
	go t.thumbnailQueueRoutine()

	return t, nil
}

// Creates missing thumbnails
func (t *Thumbnailer) CheckCache() error {
	allThumbKeys, err := t.checkCacheDir(t.RootDir)

	if err != nil {
		return utils.WrapError(err, "Cannot check thumbnail cache")
	}

	keyHash := make(map[string]bool, len(allThumbKeys))

	for _, key := range allThumbKeys {
		keyHash[key] = true
	}

	fd, err := os.Open(t.CacheDir)

	if err != nil {
		return utils.WrapError(err, "Cannot check thumbnail cache for stall thumbnails")
	}

	defer fd.Close()

	for {
		// Don't read all files at all, it might be a lot
		fis, err := fd.Readdir(1024)

		if err == io.EOF {
			break
		}

		if err != nil {
			return utils.WrapError(err, "Cannot list thumbnails while cleaning cache")
		}

		for _, fi := range fis {
			// Thumbnail corresponds to a known picture, leave it alone
			if _, exists := keyHash[fi.Name()]; exists {
				continue
			}

			thumbPath := path.Join(t.CacheDir, fi.Name())

			revel.INFO.Printf("Removing stale thumbnail with key %s", fi.Name())
			err = os.Remove(thumbPath)

			if err != nil && !os.IsNotExist(err) {
				return utils.WrapError(err, "Cannot delete thumbnail '%s' while cleaning up cache", thumbPath)
			}
		}
	}

	return nil
}

func (t *Thumbnailer) setupDirMonitor(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}

	if strings.HasPrefix(path, ".") {
		return filepath.SkipDir
	}

	if !info.IsDir() {
		return nil
	}

	err = t.monitor.Watch(path)

	if err == nil {
		revel.INFO.Printf("Setup a directory monitor on '%s'", path)
	}

	return err
}

func (t *Thumbnailer) SetupMonitors() error {
	return utils.WrapError(filepath.Walk(t.RootDir, t.setupDirMonitor), "Cannot setup monitoring")
}

func (t *Thumbnailer) ScheduleThumbnail(filePath string) {
	revel.INFO.Printf("Scheduling thumbnailing of file %s", filePath)

	t.queuedItemsMutex.Lock()
	defer t.queuedItemsMutex.Unlock()

	if _, alreadyQueued := t.queuedItems[filePath]; alreadyQueued {
		revel.INFO.Printf("Thumbnailing already scheduled for file %s", filePath)
		return
	}

	t.thumbnailingQueue <- filePath
	t.queuedItems[filePath] = true
}

// For absolute paths, check that they are in the root dir
// For relative paths, prepend the root dir path
func (t *Thumbnailer) normalizePath(filePath string) (string, error) {
	if len(filePath) > 0 && filePath[0] == '/' {
		if !strings.HasPrefix(filePath, t.RootDir) {
			return "", fmt.Errorf("Not creating a thumbnail for a file outside the root directory: %s", filePath)
		}

		return filePath, nil
	}

	return path.Join(t.RootDir, filePath), nil
}

func (t *Thumbnailer) CreateThumbnail(filePath string) error {
	normalizedPath, err := t.normalizePath(filePath)

	if err != nil {
		return utils.WrapError(err, "Invalid path '%s'", filePath)
	}

	fileId := normalizedPath[1+len(t.RootDir):]

	startTime := time.Now()

	mw := imagick.NewMagickWand()
	defer mw.Destroy()

	err = mw.ReadImage(normalizedPath)

	if err != nil {
		return utils.WrapError(err, "Cannot read file '%s'", normalizedPath)
	}

	thumbKey := makeCacheKey(fileId)
	thumbPath := path.Join(t.CacheDir, thumbKey)

	width := mw.GetImageWidth()
	height := mw.GetImageHeight()
	var scale float32

	if width > height {
		scale = float32(THUMBNAILER_SIZE_PX) / float32(width)
		width = THUMBNAILER_SIZE_PX
		height = uint(float32(height) * scale)
	} else {
		scale = float32(THUMBNAILER_SIZE_PX) / float32(height)
		height = THUMBNAILER_SIZE_PX
		width = uint(float32(width) * scale)
	}

	// TRIANGLE is a simple linear interpolation, should be fast enough
	err = mw.ResizeImage(width, height, imagick.FILTER_TRIANGLE, 1)

	if err != nil {
		return utils.WrapError(err, "Cannot generate thumbnail for file '%s'", normalizedPath)
	}

	err = mw.SetCompressionQuality(THUMBNAILER_COMPRESSION_QUALITY)

	if err != nil {
		return utils.WrapError(err, "Cannot set compression quality for file '%s'", normalizedPath)
	}

	err = mw.WriteImage(thumbPath)

	if err != nil {
		return utils.WrapError(err, "Cannot write thumbnail '%s' for file '%s'", thumbPath, normalizedPath)
	}

	revel.INFO.Printf("Thumbnailed image '%s' as '%s' in %.2f seconds", normalizedPath, thumbPath, time.Now().Sub(startTime).Seconds())

	return nil
}

func (t *Thumbnailer) DeleteThumbnail(filePath string) error {
	normalizedPath, err := t.normalizePath(filePath)

	if err != nil {
		return utils.WrapError(err, "Invalid path '%s'", normalizedPath)
	}

	fileId := normalizedPath[1+len(t.RootDir):]
	thumbKey := makeCacheKey(fileId)
	thumbPath := path.Join(t.CacheDir, thumbKey)

	err = os.Remove(thumbPath)

	if os.IsNotExist(err) {
		return nil
	}

	if err != nil {
		return utils.WrapError(err, "Cannot remove thumbnail")
	}

	revel.INFO.Printf("Deleted thumbnail for image '%s'", normalizedPath)

	return nil
}

func (t *Thumbnailer) ThumbnailQueueSize() int {
	return len(t.queuedItems)
}
<|endoftext|>


<commit_before>package bamstats

import (
	"bytes"
	"fmt"
	"math"
	"sort"

	"github.com/biogo/hts/sam"
)

// TagMap represents a map of sam tags with integer keys
type TagMap map[int]int

// MappedReadsStats represents statistics for mapped reads
type MappedReadsStats struct {
	Total    int    `json:"total,omitempty"`
	Unmapped int    `json:"unmapped,omitempty"`
	Mapped   TagMap `json:"mapped,omitempty"`
}

// MappingsStats represents statistics for mappings
type MappingsStats struct {
	MappedReadsStats
	Continuous int           `json:"continuous"`
	Split      int           `json:"split"`
	Mappings   MultimapStats `json:"mappings"`
}

// MappedPairsStats represents statistcs for mapped read-pairs
type MappedPairsStats struct {
	MappedReadsStats
	InsertSizes TagMap `json:"insert_sizes,omitempty"`
}

// MultimapStats represents statistics for multi-maps
type MultimapStats struct {
	Ratio float64 `json:"ratio"`
	Count int     `json:"count"`
}

// GeneralStats represents general mapping statistics
type GeneralStats struct {
	Reads    MappingsStats    `json:"reads,omitempty"`
	Pairs    MappedPairsStats `json:"pairs,omitempty"`
	Coverage *CoverageStats   `json:"coverage,omitempty"`
}

// Merge updates counts from a channel of Stats instances.
func (s *GeneralStats) Merge(others chan Stats) {
	for other := range others {
		if other, ok := other.(*GeneralStats); ok {
			s.Update(other)
		}
	}
}

// Update updates all counts from a Stats instance.
func (s *GeneralStats) Update(other Stats) {
	if other, ok := other.(*GeneralStats); ok {
		s.Reads.Update(other.Reads)
		s.Pairs.Update(other.Pairs)
		if s.Coverage != nil {
			s.Coverage.Update(other.Coverage)
		}
		if len(s.Pairs.Mapped) > 0 {
			s.Pairs.Total = s.Reads.Total / 2
			s.Pairs.Unmapped = s.Pairs.Total - s.Pairs.Mapped.Total()
		}
	}
}

// Update updates all counts from another MappedReadStats instance.
func (s *MappedReadsStats) Update(other MappedReadsStats) {
	s.Total += other.Total
	s.Unmapped += other.Unmapped
	s.Mapped.Update(other.Mapped)
}

// Update updates all counts from another MappingsStats instance.
func (s *MappingsStats) Update(other MappingsStats) {
	s.MappedReadsStats.Update(other.MappedReadsStats)
	s.Continuous += other.Continuous
	s.Split += other.Split
	s.Mappings.Count += other.Mappings.Count
	s.UpdateMappingsRatio()
}

// Update updates all counts from another MappedPairsStats instance.
func (s *MappedPairsStats) Update(other MappedPairsStats) {
	s.MappedReadsStats.Update(other.MappedReadsStats)
	s.InsertSizes.Update(other.InsertSizes)
}

// FilterInsertSizes filters out insert size lengths having support below the given percentage of total read-pairs.
func (s *MappedPairsStats) FilterInsertSizes(percent float64) {
	for k, v := range s.InsertSizes {
		if float64(v) < float64(s.Total)*(percent/100) {
			delete(s.InsertSizes, k)
		}
	}
}

// UpdateMappingsRatio updates ration of mappings vs total mapped reads.
func (s *MappingsStats) UpdateMappingsRatio() {
	s.Mappings.Ratio = float64(s.Mappings.Count) / float64(s.Mapped.Total())
}

// Unique returns the number of uniquely mapped reads.
func (s *MappedReadsStats) Unique() int {
	return s.Mapped[1]
}

// Update updates all counts from another TagMap instance.
func (tm TagMap) Update(other TagMap) {
	for k := range tm {
		tm[k] += other[k]
	}
	for k := range other {
		if _, ok := tm[k]; !ok {
			tm[k] += other[k]
		}
	}
}

// Total returns the total number of reads in the TagMap
func (tm TagMap) Total() (sum int) {
	for _, v := range tm {
		sum += v
	}
	return
}

// NewGeneralStats creates a new instance of GeneralStats
func NewGeneralStats() *GeneralStats {
	ms := GeneralStats{}
	ms.Pairs = *NewMappedPairsStats()
	ms.Reads.MappedReadsStats = *NewMappedReadsStats()
	return &ms
}

// NewMappedReadsStats creates a new instance of MappedReadsStats
func NewMappedReadsStats() *MappedReadsStats {
	s := MappedReadsStats{}
	s.Mapped = make(TagMap)
	return &s
}

// NewMappedPairsStats creates a new instance of MappedPairsStats
func NewMappedPairsStats() *MappedPairsStats {
	s := MappedPairsStats{}
	s.MappedReadsStats = *NewMappedReadsStats()
	s.InsertSizes = make(TagMap)
	return &s
}

// MarshalJSON returns a JSON representation of a TagMap, numerically sorting the keys.
func (tm TagMap) MarshalJSON() ([]byte, error) {
	buf := &bytes.Buffer{}
	buf.Write([]byte{'{', '\n'})
	var keys []int
	for k := range tm {
		keys = append(keys, int(k))
	}
	sort.Ints(keys)
	l := len(keys)
	for i, k := range keys {
		fmt.Fprintf(buf, "\t\"%d\": \"%v\"", k, tm[k])
		if i < l-1 {
			buf.WriteByte(',')
		}
		buf.WriteByte('\n')
	}
	buf.Write([]byte{'}', '\n'})
	return buf.Bytes(), nil
}

// Collect collects general mapping statistics from a sam.Record.
func (s *GeneralStats) Collect(r *sam.Record, index *RtreeMap) {
	NH, hasNH := r.Tag([]byte("NH"))
	if !hasNH {
		NH, _ = sam.ParseAux([]byte("NH:i:0"))
	}
	NHKey := int(NH.Value().(uint8))
	if isUnmapped(r) {
		s.Reads.Total++
		s.Reads.Unmapped++
		return
	}
	s.Reads.Mappings.Count++
	if isPrimary(r) {
		s.Reads.Total++
		s.Reads.Mapped[NHKey]++
		if isSplit(r) {
			s.Reads.Split++
		} else {
			s.Reads.Continuous++
		}
		if isFirstOfValidPair(r) {
			s.Pairs.Mapped[NHKey]++
			isLen := int(math.Abs(float64(r.TempLen)))
			s.Pairs.InsertSizes[isLen]++
		}
	}
}
<commit_msg>Remove split and continuos reads from general stats<commit_after>package bamstats

import (
	"bytes"
	"fmt"
	"math"
	"sort"

	"github.com/biogo/hts/sam"
)

// TagMap represents a map of sam tags with integer keys
type TagMap map[int]int

// MappedReadsStats represents statistics for mapped reads
type MappedReadsStats struct {
	Total    int    `json:"total,omitempty"`
	Unmapped int    `json:"unmapped,omitempty"`
	Mapped   TagMap `json:"mapped,omitempty"`
}

// MappingsStats represents statistics for mappings
type MappingsStats struct {
	MappedReadsStats
	Mappings MultimapStats `json:"mappings"`
}

// MappedPairsStats represents statistcs for mapped read-pairs
type MappedPairsStats struct {
	MappedReadsStats
	InsertSizes TagMap `json:"insert_sizes,omitempty"`
}

// MultimapStats represents statistics for multi-maps
type MultimapStats struct {
	Ratio float64 `json:"ratio"`
	Count int     `json:"count"`
}

// GeneralStats represents general mapping statistics
type GeneralStats struct {
	Reads    MappingsStats    `json:"reads,omitempty"`
	Pairs    MappedPairsStats `json:"pairs,omitempty"`
	Coverage *CoverageStats   `json:"coverage,omitempty"`
}

// Merge updates counts from a channel of Stats instances.
func (s *GeneralStats) Merge(others chan Stats) {
	for other := range others {
		if other, ok := other.(*GeneralStats); ok {
			s.Update(other)
		}
	}
}

// Update updates all counts from a Stats instance.
func (s *GeneralStats) Update(other Stats) {
	if other, ok := other.(*GeneralStats); ok {
		s.Reads.Update(other.Reads)
		s.Pairs.Update(other.Pairs)
		if s.Coverage != nil {
			s.Coverage.Update(other.Coverage)
		}
		if len(s.Pairs.Mapped) > 0 {
			s.Pairs.Total = s.Reads.Total / 2
			s.Pairs.Unmapped = s.Pairs.Total - s.Pairs.Mapped.Total()
		}
	}
}

// Update updates all counts from another MappedReadStats instance.
func (s *MappedReadsStats) Update(other MappedReadsStats) {
	s.Total += other.Total
	s.Unmapped += other.Unmapped
	s.Mapped.Update(other.Mapped)
}

// Update updates all counts from another MappingsStats instance.
func (s *MappingsStats) Update(other MappingsStats) {
	s.MappedReadsStats.Update(other.MappedReadsStats)
	s.Mappings.Count += other.Mappings.Count
	s.UpdateMappingsRatio()
}

// Update updates all counts from another MappedPairsStats instance.
func (s *MappedPairsStats) Update(other MappedPairsStats) {
	s.MappedReadsStats.Update(other.MappedReadsStats)
	s.InsertSizes.Update(other.InsertSizes)
}

// FilterInsertSizes filters out insert size lengths having support below the given percentage of total read-pairs.
func (s *MappedPairsStats) FilterInsertSizes(percent float64) {
	for k, v := range s.InsertSizes {
		if float64(v) < float64(s.Total)*(percent/100) {
			delete(s.InsertSizes, k)
		}
	}
}

// UpdateMappingsRatio updates ration of mappings vs total mapped reads.
func (s *MappingsStats) UpdateMappingsRatio() {
	s.Mappings.Ratio = float64(s.Mappings.Count) / float64(s.Mapped.Total())
}

// Unique returns the number of uniquely mapped reads.
func (s *MappedReadsStats) Unique() int {
	return s.Mapped[1]
}

// Update updates all counts from another TagMap instance.
func (tm TagMap) Update(other TagMap) {
	for k := range tm {
		tm[k] += other[k]
	}
	for k := range other {
		if _, ok := tm[k]; !ok {
			tm[k] += other[k]
		}
	}
}

// Total returns the total number of reads in the TagMap
func (tm TagMap) Total() (sum int) {
	for _, v := range tm {
		sum += v
	}
	return
}

// NewGeneralStats creates a new instance of GeneralStats
func NewGeneralStats() *GeneralStats {
	ms := GeneralStats{}
	ms.Pairs = *NewMappedPairsStats()
	ms.Reads.MappedReadsStats = *NewMappedReadsStats()
	return &ms
}

// NewMappedReadsStats creates a new instance of MappedReadsStats
func NewMappedReadsStats() *MappedReadsStats {
	s := MappedReadsStats{}
	s.Mapped = make(TagMap)
	return &s
}

// NewMappedPairsStats creates a new instance of MappedPairsStats
func NewMappedPairsStats() *MappedPairsStats {
	s := MappedPairsStats{}
	s.MappedReadsStats = *NewMappedReadsStats()
	s.InsertSizes = make(TagMap)
	return &s
}

// MarshalJSON returns a JSON representation of a TagMap, numerically sorting the keys.
func (tm TagMap) MarshalJSON() ([]byte, error) {
	buf := &bytes.Buffer{}
	buf.Write([]byte{'{', '\n'})
	var keys []int
	for k := range tm {
		keys = append(keys, int(k))
	}
	sort.Ints(keys)
	l := len(keys)
	for i, k := range keys {
		fmt.Fprintf(buf, "\t\"%d\": \"%v\"", k, tm[k])
		if i < l-1 {
			buf.WriteByte(',')
		}
		buf.WriteByte('\n')
	}
	buf.Write([]byte{'}', '\n'})
	return buf.Bytes(), nil
}

// Collect collects general mapping statistics from a sam.Record.
func (s *GeneralStats) Collect(r *sam.Record, index *RtreeMap) {
	NH, hasNH := r.Tag([]byte("NH"))
	if !hasNH {
		NH, _ = sam.ParseAux([]byte("NH:i:0"))
	}
	NHKey := int(NH.Value().(uint8))
	if isUnmapped(r) {
		s.Reads.Total++
		s.Reads.Unmapped++
		return
	}
	s.Reads.Mappings.Count++
	if isPrimary(r) {
		s.Reads.Total++
		s.Reads.Mapped[NHKey]++
		if isFirstOfValidPair(r) {
			s.Pairs.Mapped[NHKey]++
			isLen := int(math.Abs(float64(r.TempLen)))
			s.Pairs.InsertSizes[isLen]++
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2020-2021 Tigera, Inc. 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.

// +build fvtests

package fv_test

import (
	"context"
	"fmt"
	"time"

	"github.com/projectcalico/felix/fv/connectivity"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	log "github.com/sirupsen/logrus"

	api "github.com/projectcalico/api/pkg/apis/projectcalico/v3"
	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/infrastructure"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/felix/fv/workload"
	"github.com/projectcalico/libcalico-go/lib/apiconfig"
	client "github.com/projectcalico/libcalico-go/lib/clientv3"
	"github.com/projectcalico/libcalico-go/lib/options"
)

var _ = infrastructure.DatastoreDescribe("do-not-track policy tests; with 2 nodes", []apiconfig.DatastoreType{apiconfig.EtcdV3, apiconfig.Kubernetes}, func(getInfra infrastructure.InfraFactory) {

	var (
		infra          infrastructure.DatastoreInfra
		felixes        []*infrastructure.Felix
		hostW          [2]*workload.Workload
		client         client.Interface
		cc             *connectivity.Checker
		dumpedDiags    bool
		externalClient *containers.Container
	)

	BeforeEach(func() {
		var err error
		infra = getInfra()

		dumpedDiags = false
		options := infrastructure.DefaultTopologyOptions()
		felixes, client = infrastructure.StartNNodeTopology(2, options, infra)
		cc = &connectivity.Checker{}

		// Start a host networked workload on each host for connectivity checks.
		for ii := range felixes {
			// We tell each workload to open:
			// - its normal (uninteresting) port, 8055
			// - port 2379, which is both an inbound and an outbound failsafe port
			// - port 22, which is an inbound failsafe port.
			// This allows us to test the interaction between do-not-track policy and failsafe
			// ports.
			const portsToOpen = "8055,2379,22"
			hostW[ii] = workload.Run(
				felixes[ii],
				fmt.Sprintf("host%d", ii),
				"default",
				felixes[ii].IP, // Same IP as felix means "run in the host's namespace"
				portsToOpen,
				"tcp")
		}

		// We will use this container to model an external client trying to connect into
		// workloads on a host.  Create a route in the container for the workload CIDR.
		externalClient = containers.Run("external-client",
			containers.RunOpts{AutoRemove: true},
			"--privileged", // So that we can add routes inside the container.
			utils.Config.BusyboxImage,
			"/bin/sh", "-c", "sleep 1000")

		err = infra.AddDefaultDeny()
		Expect(err).To(BeNil())
	})

	// Utility function to dump diags if the test failed.  Should be called in the inner-most
	// AfterEach() to dump diags before the test is torn down.  Only the first call for a given
	// test has any effect.
	dumpDiags := func() {
		if !CurrentGinkgoTestDescription().Failed || dumpedDiags {
			return
		}
		for ii := range felixes {
			iptSave, err := felixes[ii].ExecOutput("iptables-save", "-c")
			if err == nil {
				log.WithField("felix", ii).Info("iptables-save:\n" + iptSave)
			}
			ipR, err := felixes[ii].ExecOutput("ip", "r")
			if err == nil {
				log.WithField("felix", ii).Info("ip route:\n" + ipR)
			}
		}
		infra.DumpErrorData()

	}

	AfterEach(func() {
		dumpDiags()
		for _, f := range felixes {
			f.Stop()
		}
		infra.Stop()
		externalClient.Stop()
	})

	expectFullConnectivity := func() {
		cc.ResetExpectations()
		cc.ExpectSome(felixes[0], hostW[1].Port(8055))
		cc.ExpectSome(felixes[1], hostW[0].Port(8055))
		cc.ExpectSome(felixes[0], hostW[1].Port(2379))
		cc.ExpectSome(felixes[1], hostW[0].Port(2379))
		cc.ExpectSome(felixes[0], hostW[1].Port(22))
		cc.ExpectSome(felixes[1], hostW[0].Port(22))
		cc.ExpectSome(externalClient, hostW[1].Port(22))
		cc.ExpectSome(externalClient, hostW[0].Port(22))
		cc.CheckConnectivityOffset(1)
	}

	It("before adding policy, should have connectivity between hosts", func() {
		expectFullConnectivity()
	})

	Context("after adding host endpoints", func() {
		var (
			ctx    context.Context
			cancel context.CancelFunc
		)

		BeforeEach(func() {
			// Make sure our new host endpoints don't cut felix off from the datastore.
			err := infra.AddAllowToDatastore("host-endpoint=='true'")
			Expect(err).NotTo(HaveOccurred())

			ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)

			for _, f := range felixes {
				hep := api.NewHostEndpoint()
				hep.Name = "eth0-" + f.Name
				hep.Labels = map[string]string{
					"name":          hep.Name,
					"host-endpoint": "true",
				}
				hep.Spec.Node = f.Hostname
				hep.Spec.ExpectedIPs = []string{f.IP}
				_, err := client.HostEndpoints().Create(ctx, hep, options.SetOptions{})
				Expect(err).NotTo(HaveOccurred())
			}
		})

		AfterEach(func() {
			dumpDiags()
			cancel()
		})

		It("should implement untracked policy correctly", func() {
			// This test covers both normal connectivity and failsafe connectivity.  We combine the
			// tests because we rely on the changes of normal connectivity at each step to make sure
			// that the policy has actually flowed through to the dataplane.

			By("having only failsafe connectivity to start with")
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			// Port 22 is inbound-only so it'll be blocked by the (lack of egress policy).
			cc.ExpectNone(felixes[0], hostW[1].Port(22))
			cc.ExpectNone(felixes[1], hostW[0].Port(22))
			// But external client should still be able to access it...
			cc.ExpectSome(externalClient, hostW[1].Port(22))
			cc.ExpectSome(externalClient, hostW[0].Port(22))
			cc.CheckConnectivity()

			host0Selector := fmt.Sprintf("name == 'eth0-%s'", felixes[0].Name)
			host1Selector := fmt.Sprintf("name == 'eth0-%s'", felixes[1].Name)

			By("Having connectivity after installing bidirectional policies")
			host0Pol := api.NewGlobalNetworkPolicy()
			host0Pol.Name = "host-0-pol"
			host0Pol.Spec.Selector = host0Selector
			host0Pol.Spec.DoNotTrack = true
			host0Pol.Spec.ApplyOnForward = true
			host0Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Allow,
					Source: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol, err := client.GlobalNetworkPolicies().Create(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			host1Pol := api.NewGlobalNetworkPolicy()
			host1Pol.Name = "host-1-pol"
			host1Pol.Spec.Selector = host1Selector
			host1Pol.Spec.DoNotTrack = true
			host1Pol.Spec.ApplyOnForward = true
			host1Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Allow,
					Source: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host1Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host1Pol, err = client.GlobalNetworkPolicies().Create(ctx, host1Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			expectFullConnectivity()

			By("Having only failsafe connectivity after replacing host-0's egress rules with Deny")
			// Since there's no conntrack, removing rules in one direction is enough to prevent
			// connectivity in either direction.
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Deny,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			cc.ResetExpectations()
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			cc.ExpectNone(felixes[0], hostW[1].Port(22))     // Now blocked (lack of egress).
			cc.ExpectSome(felixes[1], hostW[0].Port(22))     // Still open due to failsafe.
			cc.ExpectSome(externalClient, hostW[1].Port(22)) // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[0].Port(22)) // Allowed by failsafe
			cc.CheckConnectivity()

			By("Having full connectivity after putting them back")
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			expectFullConnectivity()

			By("Having only failsafe connectivity after replacing host-0's ingress rules with Deny")
			host0Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Deny,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			cc.ResetExpectations()
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			cc.ExpectNone(felixes[0], hostW[1].Port(22))     // Response traffic blocked by policy
			cc.ExpectSome(felixes[1], hostW[0].Port(22))     // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[1].Port(22)) // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[0].Port(22)) // Allowed by failsafe
			cc.CheckConnectivity()
		})
	})
})
<commit_msg>Mark DoNotTrack FV tests to run in BPF mode<commit_after>// Copyright (c) 2020-2021 Tigera, Inc. 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.

// +build fvtests

package fv_test

import (
	"context"
	"fmt"
	"time"

	"github.com/projectcalico/felix/fv/connectivity"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	log "github.com/sirupsen/logrus"

	api "github.com/projectcalico/api/pkg/apis/projectcalico/v3"
	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/infrastructure"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/felix/fv/workload"
	"github.com/projectcalico/libcalico-go/lib/apiconfig"
	client "github.com/projectcalico/libcalico-go/lib/clientv3"
	"github.com/projectcalico/libcalico-go/lib/options"
)

var _ = infrastructure.DatastoreDescribe("_BPF-SAFE_ do-not-track policy tests; with 2 nodes", []apiconfig.DatastoreType{apiconfig.EtcdV3, apiconfig.Kubernetes}, func(getInfra infrastructure.InfraFactory) {

	var (
		infra          infrastructure.DatastoreInfra
		felixes        []*infrastructure.Felix
		hostW          [2]*workload.Workload
		client         client.Interface
		cc             *connectivity.Checker
		dumpedDiags    bool
		externalClient *containers.Container
	)

	BeforeEach(func() {
		var err error
		infra = getInfra()

		dumpedDiags = false
		options := infrastructure.DefaultTopologyOptions()
		felixes, client = infrastructure.StartNNodeTopology(2, options, infra)
		cc = &connectivity.Checker{}

		// Start a host networked workload on each host for connectivity checks.
		for ii := range felixes {
			// We tell each workload to open:
			// - its normal (uninteresting) port, 8055
			// - port 2379, which is both an inbound and an outbound failsafe port
			// - port 22, which is an inbound failsafe port.
			// This allows us to test the interaction between do-not-track policy and failsafe
			// ports.
			const portsToOpen = "8055,2379,22"
			hostW[ii] = workload.Run(
				felixes[ii],
				fmt.Sprintf("host%d", ii),
				"default",
				felixes[ii].IP, // Same IP as felix means "run in the host's namespace"
				portsToOpen,
				"tcp")
		}

		// We will use this container to model an external client trying to connect into
		// workloads on a host.  Create a route in the container for the workload CIDR.
		externalClient = containers.Run("external-client",
			containers.RunOpts{AutoRemove: true},
			"--privileged", // So that we can add routes inside the container.
			utils.Config.BusyboxImage,
			"/bin/sh", "-c", "sleep 1000")

		err = infra.AddDefaultDeny()
		Expect(err).To(BeNil())
	})

	// Utility function to dump diags if the test failed.  Should be called in the inner-most
	// AfterEach() to dump diags before the test is torn down.  Only the first call for a given
	// test has any effect.
	dumpDiags := func() {
		if !CurrentGinkgoTestDescription().Failed || dumpedDiags {
			return
		}
		for ii := range felixes {
			iptSave, err := felixes[ii].ExecOutput("iptables-save", "-c")
			if err == nil {
				log.WithField("felix", ii).Info("iptables-save:\n" + iptSave)
			}
			ipR, err := felixes[ii].ExecOutput("ip", "r")
			if err == nil {
				log.WithField("felix", ii).Info("ip route:\n" + ipR)
			}
		}
		infra.DumpErrorData()

	}

	AfterEach(func() {
		dumpDiags()
		for _, f := range felixes {
			f.Stop()
		}
		infra.Stop()
		externalClient.Stop()
	})

	expectFullConnectivity := func() {
		cc.ResetExpectations()
		cc.ExpectSome(felixes[0], hostW[1].Port(8055))
		cc.ExpectSome(felixes[1], hostW[0].Port(8055))
		cc.ExpectSome(felixes[0], hostW[1].Port(2379))
		cc.ExpectSome(felixes[1], hostW[0].Port(2379))
		cc.ExpectSome(felixes[0], hostW[1].Port(22))
		cc.ExpectSome(felixes[1], hostW[0].Port(22))
		cc.ExpectSome(externalClient, hostW[1].Port(22))
		cc.ExpectSome(externalClient, hostW[0].Port(22))
		cc.CheckConnectivityOffset(1)
	}

	It("before adding policy, should have connectivity between hosts", func() {
		expectFullConnectivity()
	})

	Context("after adding host endpoints", func() {
		var (
			ctx    context.Context
			cancel context.CancelFunc
		)

		BeforeEach(func() {
			// Make sure our new host endpoints don't cut felix off from the datastore.
			err := infra.AddAllowToDatastore("host-endpoint=='true'")
			Expect(err).NotTo(HaveOccurred())

			ctx, cancel = context.WithTimeout(context.Background(), 30*time.Second)

			for _, f := range felixes {
				hep := api.NewHostEndpoint()
				hep.Name = "eth0-" + f.Name
				hep.Labels = map[string]string{
					"name":          hep.Name,
					"host-endpoint": "true",
				}
				hep.Spec.Node = f.Hostname
				hep.Spec.ExpectedIPs = []string{f.IP}
				_, err := client.HostEndpoints().Create(ctx, hep, options.SetOptions{})
				Expect(err).NotTo(HaveOccurred())
			}
		})

		AfterEach(func() {
			dumpDiags()
			cancel()
		})

		It("should implement untracked policy correctly", func() {
			// This test covers both normal connectivity and failsafe connectivity.  We combine the
			// tests because we rely on the changes of normal connectivity at each step to make sure
			// that the policy has actually flowed through to the dataplane.

			By("having only failsafe connectivity to start with")
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			// Port 22 is inbound-only so it'll be blocked by the (lack of egress policy).
			cc.ExpectNone(felixes[0], hostW[1].Port(22))
			cc.ExpectNone(felixes[1], hostW[0].Port(22))
			// But external client should still be able to access it...
			cc.ExpectSome(externalClient, hostW[1].Port(22))
			cc.ExpectSome(externalClient, hostW[0].Port(22))
			cc.CheckConnectivity()

			host0Selector := fmt.Sprintf("name == 'eth0-%s'", felixes[0].Name)
			host1Selector := fmt.Sprintf("name == 'eth0-%s'", felixes[1].Name)

			By("Having connectivity after installing bidirectional policies")
			host0Pol := api.NewGlobalNetworkPolicy()
			host0Pol.Name = "host-0-pol"
			host0Pol.Spec.Selector = host0Selector
			host0Pol.Spec.DoNotTrack = true
			host0Pol.Spec.ApplyOnForward = true
			host0Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Allow,
					Source: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol, err := client.GlobalNetworkPolicies().Create(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			host1Pol := api.NewGlobalNetworkPolicy()
			host1Pol.Name = "host-1-pol"
			host1Pol.Spec.Selector = host1Selector
			host1Pol.Spec.DoNotTrack = true
			host1Pol.Spec.ApplyOnForward = true
			host1Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Allow,
					Source: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host1Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host1Pol, err = client.GlobalNetworkPolicies().Create(ctx, host1Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			expectFullConnectivity()

			By("Having only failsafe connectivity after replacing host-0's egress rules with Deny")
			// Since there's no conntrack, removing rules in one direction is enough to prevent
			// connectivity in either direction.
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Deny,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			cc.ResetExpectations()
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			cc.ExpectNone(felixes[0], hostW[1].Port(22))     // Now blocked (lack of egress).
			cc.ExpectSome(felixes[1], hostW[0].Port(22))     // Still open due to failsafe.
			cc.ExpectSome(externalClient, hostW[1].Port(22)) // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[0].Port(22)) // Allowed by failsafe
			cc.CheckConnectivity()

			By("Having full connectivity after putting them back")
			host0Pol.Spec.Egress = []api.Rule{
				{
					Action: api.Allow,
					Destination: api.EntityRule{
						Selector: host1Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			expectFullConnectivity()

			By("Having only failsafe connectivity after replacing host-0's ingress rules with Deny")
			host0Pol.Spec.Ingress = []api.Rule{
				{
					Action: api.Deny,
					Destination: api.EntityRule{
						Selector: host0Selector,
					},
				},
			}
			host0Pol, err = client.GlobalNetworkPolicies().Update(ctx, host0Pol, options.SetOptions{})
			Expect(err).NotTo(HaveOccurred())

			cc.ResetExpectations()
			cc.ExpectNone(felixes[0], hostW[1].Port(8055))
			cc.ExpectNone(felixes[1], hostW[0].Port(8055))
			cc.ExpectSome(felixes[0], hostW[1].Port(2379))
			cc.ExpectSome(felixes[1], hostW[0].Port(2379))
			cc.ExpectNone(felixes[0], hostW[1].Port(22))     // Response traffic blocked by policy
			cc.ExpectSome(felixes[1], hostW[0].Port(22))     // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[1].Port(22)) // Allowed by failsafe
			cc.ExpectSome(externalClient, hostW[0].Port(22)) // Allowed by failsafe
			cc.CheckConnectivity()
		})
	})
})
<|endoftext|>


<commit_before>package keymap

import (
	"github.com/Bredgren/game1/game/keymap/button"
	"github.com/hajimehoshi/ebiten"
)

// Action is a name/label for an action.
type Action string

// KeyMap groups all map types.
type KeyMap struct {
	KeyMouse    *KeyMouseMap
	GamepadBtn  *GamepadBtnMap
	GamepadAxis *GamepadAxisMap
	btnHandlers ButtonHandlerMap
	gaHandlers  AxisHandlerMap
}

// New creates and returns a new, empty KeyMap. The btnHandlers map shared between keyboard/mouse
// and gamepad buttons. This means that if an action handler for a keyboard/mouse buttton
// stops propagation then a gamepad button that maps to the same action in a later layer
// will not be handled.
func New(btnHandlers ButtonHandlerMap, gamepadAxisHandlers AxisHandlerMap) *KeyMap {
	return &KeyMap{
		KeyMouse:    NewKeyMouseMap(),
		GamepadBtn:  NewGamepadBtnMap(),
		GamepadAxis: NewGamepadAxisMap(),
		btnHandlers: btnHandlers,
		gaHandlers:  gamepadAxisHandlers,
	}
}

// ButtonHandler is a function that handles a button state. The parameter down is true
// if the button is pressed. It should return true if no later handlers should be called
// for the same button.
type ButtonHandler func(down bool) (stopPropagation bool)

// AxisHandler is a function that handles gamepad axis state. The function will directly
// be given the result of ebiten.GamepadAxis. It should return true if no later handlers
// should be called for the same axis.
type AxisHandler func(val float64) (stopPropagation bool)

// ButtonHandlerMap maps button Actions to their handlers.
type ButtonHandlerMap map[Action]ButtonHandler

// AxisHandlerMap maps axis Actions to their handlers.
type AxisHandlerMap map[Action]AxisHandler

// Layers is a slice of KeyMaps. It enables buttons to be overloaded with multiple actions
// with the option of skipping later handlers for buttons handled at earlier layers.
type Layers []*KeyMap

// Update checks input state and calls handlers for any actions triggered. It handles
// each layer in order. If the handler for a button stops propagation then later following
// layers will not handle any actions the same button triggers.
func (l Layers) Update() {
	const gamepadID = 0
	stoppedKeys := map[button.KeyMouse]bool{}
	stoppedBtns := map[ebiten.GamepadButton]bool{}
	stoppedAxes := map[int]bool{}
	for _, keymap := range l {
		if keymap == nil {
			continue
		}
		actions := map[Action]bool{}

		for _, btn := range keymap.KeyMouse.Buttons() {
			if stoppedKeys[btn] {
				continue
			}
			var down bool
			if k, ok := btn.Key(); ok {
				down = ebiten.IsKeyPressed(k)
			} else if mb, ok := btn.Mouse(); ok {
				down = ebiten.IsMouseButtonPressed(mb)
			}
			if a, ok := keymap.KeyMouse.GetAction(btn); ok {
				actions[a] = down || actions[a]
			}
		}

		for _, btn := range keymap.GamepadBtn.Buttons() {
			if stoppedBtns[btn] {
				continue
			}
			down := ebiten.IsGamepadButtonPressed(gamepadID, btn)
			if a, ok := keymap.GamepadBtn.GetAction(btn); ok {
				actions[a] = down || actions[a]
			}
		}

		for action, down := range actions {
			stop := keymap.btnHandlers[action](down)
			if b, ok := keymap.KeyMouse.GetButton(action); ok {
				stoppedKeys[b] = stop
			}
			if b, ok := keymap.GamepadBtn.GetButton(action); ok {
				stoppedBtns[b] = stop
			}
		}

		numAxis := ebiten.GamepadAxisNum(gamepadID)
		for _, axis := range keymap.GamepadAxis.Axes() {
			if stoppedAxes[axis] || axis >= numAxis {
				continue
			}
			val := ebiten.GamepadAxis(gamepadID, axis)
			if act, ok := keymap.GamepadAxis.GetAction(axis); ok {
				stop := keymap.gaHandlers[act](val)
				stoppedAxes[axis] = stop
			}
		}
	}
}
<commit_msg>Skip actions with no handler functons<commit_after>package keymap

import (
	"github.com/Bredgren/game1/game/keymap/button"
	"github.com/hajimehoshi/ebiten"
)

// Action is a name/label for an action.
type Action string

// KeyMap groups all map types.
type KeyMap struct {
	KeyMouse    *KeyMouseMap
	GamepadBtn  *GamepadBtnMap
	GamepadAxis *GamepadAxisMap
	btnHandlers ButtonHandlerMap
	gaHandlers  AxisHandlerMap
}

// New creates and returns a new, empty KeyMap. The btnHandlers map shared between keyboard/mouse
// and gamepad buttons. This means that if an action handler for a keyboard/mouse buttton
// stops propagation then a gamepad button that maps to the same action in a later layer
// will not be handled.
func New(btnHandlers ButtonHandlerMap, gamepadAxisHandlers AxisHandlerMap) *KeyMap {
	return &KeyMap{
		KeyMouse:    NewKeyMouseMap(),
		GamepadBtn:  NewGamepadBtnMap(),
		GamepadAxis: NewGamepadAxisMap(),
		btnHandlers: btnHandlers,
		gaHandlers:  gamepadAxisHandlers,
	}
}

// ButtonHandler is a function that handles a button state. The parameter down is true
// if the button is pressed. It should return true if no later handlers should be called
// for the same button.
type ButtonHandler func(down bool) (stopPropagation bool)

// AxisHandler is a function that handles gamepad axis state. The function will directly
// be given the result of ebiten.GamepadAxis. It should return true if no later handlers
// should be called for the same axis.
type AxisHandler func(val float64) (stopPropagation bool)

// ButtonHandlerMap maps button Actions to their handlers.
type ButtonHandlerMap map[Action]ButtonHandler

// AxisHandlerMap maps axis Actions to their handlers.
type AxisHandlerMap map[Action]AxisHandler

// Layers is a slice of KeyMaps. It enables buttons to be overloaded with multiple actions
// with the option of skipping later handlers for buttons handled at earlier layers.
type Layers []*KeyMap

// Update checks input state and calls handlers for any actions triggered. It handles
// each layer in order. If the handler for a button stops propagation then later following
// layers will not handle any actions the same button triggers.
func (l Layers) Update() {
	const gamepadID = 0
	stoppedKeys := map[button.KeyMouse]bool{}
	stoppedBtns := map[ebiten.GamepadButton]bool{}
	stoppedAxes := map[int]bool{}
	for _, keymap := range l {
		if keymap == nil {
			continue
		}
		actions := map[Action]bool{}

		for _, btn := range keymap.KeyMouse.Buttons() {
			if stoppedKeys[btn] {
				continue
			}
			var down bool
			if k, ok := btn.Key(); ok {
				down = ebiten.IsKeyPressed(k)
			} else if mb, ok := btn.Mouse(); ok {
				down = ebiten.IsMouseButtonPressed(mb)
			}
			if a, ok := keymap.KeyMouse.GetAction(btn); ok {
				actions[a] = down || actions[a]
			}
		}

		for _, btn := range keymap.GamepadBtn.Buttons() {
			if stoppedBtns[btn] {
				continue
			}
			down := ebiten.IsGamepadButtonPressed(gamepadID, btn)
			if a, ok := keymap.GamepadBtn.GetAction(btn); ok {
				actions[a] = down || actions[a]
			}
		}

		for action, down := range actions {
			fn, ok := keymap.btnHandlers[action]
			if !ok {
				continue
			}
			stop := fn(down)
			if b, ok := keymap.KeyMouse.GetButton(action); ok {
				stoppedKeys[b] = stop
			}
			if b, ok := keymap.GamepadBtn.GetButton(action); ok {
				stoppedBtns[b] = stop
			}
		}

		numAxis := ebiten.GamepadAxisNum(gamepadID)
		for _, axis := range keymap.GamepadAxis.Axes() {
			if stoppedAxes[axis] || axis >= numAxis {
				continue
			}
			val := ebiten.GamepadAxis(gamepadID, axis)
			if act, ok := keymap.GamepadAxis.GetAction(axis); ok {
				fn, ok := keymap.gaHandlers[act]
				if !ok {
					continue
				}
				stop := fn(val)
				stoppedAxes[axis] = stop
			}
		}
	}
}
<|endoftext|>