<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// TODO(thockin): This whole pkg is pretty linux-centric.  As soon as we have
// an alternate platform, we will need to abstract further.

package mount

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	utilexec "k8s.io/utils/exec"
)

const (
	// Default mount command if mounter path is not specified.
	defaultMountCommand = "mount"
	// Log message where sensitive mount options were removed
	sensitiveOptionsRemoved = "<masked>"
)

// Interface defines the set of methods to allow for mount operations on a system.
type Interface interface {
	// Mount mounts source to target as fstype with given options.
	// options MUST not contain sensitive material (like passwords).
	Mount(source string, target string, fstype string, options []string) error
	// MountSensitive is the same as Mount() but this method allows
	// sensitiveOptions to be passed in a separate parameter from the normal
	// mount options and ensures the sensitiveOptions are never logged. This
	// method should be used by callers that pass sensitive material (like
	// passwords) as mount options.
	MountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error
	// MountSensitiveWithoutSystemd is the same as MountSensitive() but this method disable using systemd mount.
	MountSensitiveWithoutSystemd(source string, target string, fstype string, options []string, sensitiveOptions []string) error
	// MountSensitiveWithoutSystemdWithMountFlags is the same as MountSensitiveWithoutSystemd() with additional mount flags
	MountSensitiveWithoutSystemdWithMountFlags(source string, target string, fstype string, options []string, sensitiveOptions []string, mountFlags []string) error
	// Unmount unmounts given target.
	Unmount(target string) error
	// List returns a list of all mounted filesystems.  This can be large.
	// On some platforms, reading mounts directly from the OS is not guaranteed
	// consistent (i.e. it could change between chunked reads). This is guaranteed
	// to be consistent.
	List() ([]MountPoint, error)
	// IsLikelyNotMountPoint uses heuristics to determine if a directory
	// is not a mountpoint.
	// It should return ErrNotExist when the directory does not exist.
	// IsLikelyNotMountPoint does NOT properly detect all mountpoint types
	// most notably linux bind mounts and symbolic link. For callers that do not
	// care about such situations, this is a faster alternative to calling List()
	// and scanning that output.
	IsLikelyNotMountPoint(file string) (bool, error)
	// canSafelySkipMountPointCheck indicates whether this mounter returns errors on
	// operations for targets that are not mount points. If this returns true, no such
	// errors will be returned.
	canSafelySkipMountPointCheck() bool
	// GetMountRefs finds all mount references to pathname, returning a slice of
	// paths. Pathname can be a mountpoint path or a normal	directory
	// (for bind mount). On Linux, pathname is excluded from the slice.
	// For example, if /dev/sdc was mounted at /path/a and /path/b,
	// GetMountRefs("/path/a") would return ["/path/b"]
	// GetMountRefs("/path/b") would return ["/path/a"]
	// On Windows there is no way to query all mount points; as long as pathname is
	// a valid mount, it will be returned.
	GetMountRefs(pathname string) ([]string, error)
}

// Compile-time check to ensure all Mounter implementations satisfy
// the mount interface.
var _ Interface = &Mounter{}

type MounterForceUnmounter interface {
	Interface
	// UnmountWithForce unmounts given target but will retry unmounting with force option
	// after given timeout.
	UnmountWithForce(target string, umountTimeout time.Duration) error
}

// MountPoint represents a single line in /proc/mounts or /etc/fstab.
type MountPoint struct { // nolint: golint
	Device string
	Path   string
	Type   string
	Opts   []string // Opts may contain sensitive mount options (like passwords) and MUST be treated as such (e.g. not logged).
	Freq   int
	Pass   int
}

type MountErrorType string // nolint: golint

const (
	FilesystemMismatch  MountErrorType = "FilesystemMismatch"
	HasFilesystemErrors MountErrorType = "HasFilesystemErrors"
	UnformattedReadOnly MountErrorType = "UnformattedReadOnly"
	FormatFailed        MountErrorType = "FormatFailed"
	GetDiskFormatFailed MountErrorType = "GetDiskFormatFailed"
	UnknownMountError   MountErrorType = "UnknownMountError"
)

type MountError struct { // nolint: golint
	Type    MountErrorType
	Message string
}

func (mountError MountError) String() string {
	return mountError.Message
}

func (mountError MountError) Error() string {
	return mountError.Message
}

func NewMountError(mountErrorValue MountErrorType, format string, args ...interface{}) error {
	mountError := MountError{
		Type:    mountErrorValue,
		Message: fmt.Sprintf(format, args...),
	}
	return mountError
}

// SafeFormatAndMount probes a device to see if it is formatted.
// Namely it checks to see if a file system is present. If so it
// mounts it otherwise the device is formatted first then mounted.
type SafeFormatAndMount struct {
	Interface
	Exec utilexec.Interface
}

// FormatAndMount formats the given disk, if needed, and mounts it.
// That is if the disk is not formatted and it is not being mounted as
// read-only it will format it first then mount it. Otherwise, if the
// disk is already formatted or it is being mounted as read-only, it
// will be mounted without formatting.
// options MUST not contain sensitive material (like passwords).
func (mounter *SafeFormatAndMount) FormatAndMount(source string, target string, fstype string, options []string) error {
	return mounter.FormatAndMountSensitive(source, target, fstype, options, nil /* sensitiveOptions */)
}

// FormatAndMountSensitive is the same as FormatAndMount but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal mount
// options and ensures the sensitiveOptions are never logged. This method should
// be used by callers that pass sensitive material (like passwords) as mount
// options.
func (mounter *SafeFormatAndMount) FormatAndMountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	return mounter.formatAndMountSensitive(source, target, fstype, options, sensitiveOptions)
}

// getMountRefsByDev finds all references to the device provided
// by mountPath; returns a list of paths.
// Note that mountPath should be path after the evaluation of any symblolic links.
func getMountRefsByDev(mounter Interface, mountPath string) ([]string, error) {
	mps, err := mounter.List()
	if err != nil {
		return nil, err
	}

	// Finding the device mounted to mountPath.
	diskDev := ""
	for i := range mps {
		if mountPath == mps[i].Path {
			diskDev = mps[i].Device
			break
		}
	}

	// Find all references to the device.
	var refs []string
	for i := range mps {
		if mps[i].Device == diskDev || mps[i].Device == mountPath {
			if mps[i].Path != mountPath {
				refs = append(refs, mps[i].Path)
			}
		}
	}
	return refs, nil
}

// GetDeviceNameFromMount given a mnt point, find the device from /proc/mounts
// returns the device name, reference count, and error code.
func GetDeviceNameFromMount(mounter Interface, mountPath string) (string, int, error) {
	mps, err := mounter.List()
	if err != nil {
		return "", 0, err
	}

	// Find the device name.
	// FIXME if multiple devices mounted on the same mount path, only the first one is returned.
	device := ""
	// If mountPath is symlink, need get its target path.
	slTarget, err := filepath.EvalSymlinks(mountPath)
	if err != nil {
		slTarget = mountPath
	}
	for i := range mps {
		if mps[i].Path == slTarget {
			device = mps[i].Device
			break
		}
	}

	// Find all references to the device.
	refCount := 0
	for i := range mps {
		if mps[i].Device == device {
			refCount++
		}
	}
	return device, refCount, nil
}

// IsNotMountPoint determines if a directory is a mountpoint.
// It should return ErrNotExist when the directory does not exist.
// IsNotMountPoint is more expensive than IsLikelyNotMountPoint.
// IsNotMountPoint detects bind mounts in linux.
// IsNotMountPoint enumerates all the mountpoints using List() and
// the list of mountpoints may be large, then it uses
// isMountPointMatch to evaluate whether the directory is a mountpoint.
func IsNotMountPoint(mounter Interface, file string) (bool, error) {
	// IsLikelyNotMountPoint provides a quick check
	// to determine whether file IS A mountpoint.
	notMnt, notMntErr := mounter.IsLikelyNotMountPoint(file)
	if notMntErr != nil && os.IsPermission(notMntErr) {
		// We were not allowed to do the simple stat() check, e.g. on NFS with
		// root_squash. Fall back to /proc/mounts check below.
		notMnt = true
		notMntErr = nil
	}
	if notMntErr != nil {
		return notMnt, notMntErr
	}
	// identified as mountpoint, so return this fact.
	if notMnt == false {
		return notMnt, nil
	}

	// Resolve any symlinks in file, kernel would do the same and use the resolved path in /proc/mounts.
	resolvedFile, err := filepath.EvalSymlinks(file)
	if err != nil {
		return true, err
	}

	// check all mountpoints since IsLikelyNotMountPoint
	// is not reliable for some mountpoint types.
	mountPoints, mountPointsErr := mounter.List()
	if mountPointsErr != nil {
		return notMnt, mountPointsErr
	}
	for _, mp := range mountPoints {
		if isMountPointMatch(mp, resolvedFile) {
			notMnt = false
			break
		}
	}
	return notMnt, nil
}

// MakeBindOpts detects whether a bind mount is being requested and makes the remount options to
// use in case of bind mount, due to the fact that bind mount doesn't respect mount options.
// The list equals:
//   options - 'bind' + 'remount' (no duplicate)
func MakeBindOpts(options []string) (bool, []string, []string) {
	bind, bindOpts, bindRemountOpts, _ := MakeBindOptsSensitive(options, nil /* sensitiveOptions */)
	return bind, bindOpts, bindRemountOpts
}

// MakeBindOptsSensitive is the same as MakeBindOpts but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal mount
// options and ensures the sensitiveOptions are never logged. This method should
// be used by callers that pass sensitive material (like passwords) as mount
// options.
func MakeBindOptsSensitive(options []string, sensitiveOptions []string) (bool, []string, []string, []string) {
	// Because we have an FD opened on the subpath bind mount, the "bind" option
	// needs to be included, otherwise the mount target will error as busy if you
	// remount as readonly.
	//
	// As a consequence, all read only bind mounts will no longer change the underlying
	// volume mount to be read only.
	bindRemountOpts := []string{"bind", "remount"}
	bindRemountSensitiveOpts := []string{}
	bind := false
	bindOpts := []string{"bind"}

	// _netdev is a userspace mount option and does not automatically get added when
	// bind mount is created and hence we must carry it over.
	if checkForNetDev(options, sensitiveOptions) {
		bindOpts = append(bindOpts, "_netdev")
	}

	for _, option := range options {
		switch option {
		case "bind":
			bind = true
		case "remount":
		default:
			bindRemountOpts = append(bindRemountOpts, option)
		}
	}

	for _, sensitiveOption := range sensitiveOptions {
		switch sensitiveOption {
		case "bind":
			bind = true
		case "remount":
		default:
			bindRemountSensitiveOpts = append(bindRemountSensitiveOpts, sensitiveOption)
		}
	}

	return bind, bindOpts, bindRemountOpts, bindRemountSensitiveOpts
}

func checkForNetDev(options []string, sensitiveOptions []string) bool {
	for _, option := range options {
		if option == "_netdev" {
			return true
		}
	}
	for _, sensitiveOption := range sensitiveOptions {
		if sensitiveOption == "_netdev" {
			return true
		}
	}
	return false
}

// PathWithinBase checks if give path is within given base directory.
func PathWithinBase(fullPath, basePath string) bool {
	rel, err := filepath.Rel(basePath, fullPath)
	if err != nil {
		return false
	}
	if StartsWithBackstep(rel) {
		// Needed to escape the base path.
		return false
	}
	return true
}

// StartsWithBackstep checks if the given path starts with a backstep segment.
func StartsWithBackstep(rel string) bool {
	// normalize to / and check for ../
	return rel == ".." || strings.HasPrefix(filepath.ToSlash(rel), "../")
}

// sanitizedOptionsForLogging will return a comma separated string containing
// options and sensitiveOptions. Each entry in sensitiveOptions will be
// replaced with the string sensitiveOptionsRemoved
// e.g. o1,o2,<masked>,<masked>
func sanitizedOptionsForLogging(options []string, sensitiveOptions []string) string {
	separator := ""
	if len(options) > 0 && len(sensitiveOptions) > 0 {
		separator = ","
	}

	sensitiveOptionsStart := ""
	sensitiveOptionsEnd := ""
	if len(sensitiveOptions) > 0 {
		sensitiveOptionsStart = strings.Repeat(sensitiveOptionsRemoved+",", len(sensitiveOptions)-1)
		sensitiveOptionsEnd = sensitiveOptionsRemoved
	}

	return strings.Join(options, ",") +
		separator +
		sensitiveOptionsStart +
		sensitiveOptionsEnd
}
<commit_msg>Generate and format files<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// TODO(thockin): This whole pkg is pretty linux-centric.  As soon as we have
// an alternate platform, we will need to abstract further.

package mount

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	utilexec "k8s.io/utils/exec"
)

const (
	// Default mount command if mounter path is not specified.
	defaultMountCommand = "mount"
	// Log message where sensitive mount options were removed
	sensitiveOptionsRemoved = "<masked>"
)

// Interface defines the set of methods to allow for mount operations on a system.
type Interface interface {
	// Mount mounts source to target as fstype with given options.
	// options MUST not contain sensitive material (like passwords).
	Mount(source string, target string, fstype string, options []string) error
	// MountSensitive is the same as Mount() but this method allows
	// sensitiveOptions to be passed in a separate parameter from the normal
	// mount options and ensures the sensitiveOptions are never logged. This
	// method should be used by callers that pass sensitive material (like
	// passwords) as mount options.
	MountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error
	// MountSensitiveWithoutSystemd is the same as MountSensitive() but this method disable using systemd mount.
	MountSensitiveWithoutSystemd(source string, target string, fstype string, options []string, sensitiveOptions []string) error
	// MountSensitiveWithoutSystemdWithMountFlags is the same as MountSensitiveWithoutSystemd() with additional mount flags
	MountSensitiveWithoutSystemdWithMountFlags(source string, target string, fstype string, options []string, sensitiveOptions []string, mountFlags []string) error
	// Unmount unmounts given target.
	Unmount(target string) error
	// List returns a list of all mounted filesystems.  This can be large.
	// On some platforms, reading mounts directly from the OS is not guaranteed
	// consistent (i.e. it could change between chunked reads). This is guaranteed
	// to be consistent.
	List() ([]MountPoint, error)
	// IsLikelyNotMountPoint uses heuristics to determine if a directory
	// is not a mountpoint.
	// It should return ErrNotExist when the directory does not exist.
	// IsLikelyNotMountPoint does NOT properly detect all mountpoint types
	// most notably linux bind mounts and symbolic link. For callers that do not
	// care about such situations, this is a faster alternative to calling List()
	// and scanning that output.
	IsLikelyNotMountPoint(file string) (bool, error)
	// canSafelySkipMountPointCheck indicates whether this mounter returns errors on
	// operations for targets that are not mount points. If this returns true, no such
	// errors will be returned.
	canSafelySkipMountPointCheck() bool
	// GetMountRefs finds all mount references to pathname, returning a slice of
	// paths. Pathname can be a mountpoint path or a normal	directory
	// (for bind mount). On Linux, pathname is excluded from the slice.
	// For example, if /dev/sdc was mounted at /path/a and /path/b,
	// GetMountRefs("/path/a") would return ["/path/b"]
	// GetMountRefs("/path/b") would return ["/path/a"]
	// On Windows there is no way to query all mount points; as long as pathname is
	// a valid mount, it will be returned.
	GetMountRefs(pathname string) ([]string, error)
}

// Compile-time check to ensure all Mounter implementations satisfy
// the mount interface.
var _ Interface = &Mounter{}

type MounterForceUnmounter interface {
	Interface
	// UnmountWithForce unmounts given target but will retry unmounting with force option
	// after given timeout.
	UnmountWithForce(target string, umountTimeout time.Duration) error
}

// MountPoint represents a single line in /proc/mounts or /etc/fstab.
type MountPoint struct { // nolint: golint
	Device string
	Path   string
	Type   string
	Opts   []string // Opts may contain sensitive mount options (like passwords) and MUST be treated as such (e.g. not logged).
	Freq   int
	Pass   int
}

type MountErrorType string // nolint: golint

const (
	FilesystemMismatch  MountErrorType = "FilesystemMismatch"
	HasFilesystemErrors MountErrorType = "HasFilesystemErrors"
	UnformattedReadOnly MountErrorType = "UnformattedReadOnly"
	FormatFailed        MountErrorType = "FormatFailed"
	GetDiskFormatFailed MountErrorType = "GetDiskFormatFailed"
	UnknownMountError   MountErrorType = "UnknownMountError"
)

type MountError struct { // nolint: golint
	Type    MountErrorType
	Message string
}

func (mountError MountError) String() string {
	return mountError.Message
}

func (mountError MountError) Error() string {
	return mountError.Message
}

func NewMountError(mountErrorValue MountErrorType, format string, args ...interface{}) error {
	mountError := MountError{
		Type:    mountErrorValue,
		Message: fmt.Sprintf(format, args...),
	}
	return mountError
}

// SafeFormatAndMount probes a device to see if it is formatted.
// Namely it checks to see if a file system is present. If so it
// mounts it otherwise the device is formatted first then mounted.
type SafeFormatAndMount struct {
	Interface
	Exec utilexec.Interface
}

// FormatAndMount formats the given disk, if needed, and mounts it.
// That is if the disk is not formatted and it is not being mounted as
// read-only it will format it first then mount it. Otherwise, if the
// disk is already formatted or it is being mounted as read-only, it
// will be mounted without formatting.
// options MUST not contain sensitive material (like passwords).
func (mounter *SafeFormatAndMount) FormatAndMount(source string, target string, fstype string, options []string) error {
	return mounter.FormatAndMountSensitive(source, target, fstype, options, nil /* sensitiveOptions */)
}

// FormatAndMountSensitive is the same as FormatAndMount but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal mount
// options and ensures the sensitiveOptions are never logged. This method should
// be used by callers that pass sensitive material (like passwords) as mount
// options.
func (mounter *SafeFormatAndMount) FormatAndMountSensitive(source string, target string, fstype string, options []string, sensitiveOptions []string) error {
	return mounter.formatAndMountSensitive(source, target, fstype, options, sensitiveOptions)
}

// getMountRefsByDev finds all references to the device provided
// by mountPath; returns a list of paths.
// Note that mountPath should be path after the evaluation of any symblolic links.
func getMountRefsByDev(mounter Interface, mountPath string) ([]string, error) {
	mps, err := mounter.List()
	if err != nil {
		return nil, err
	}

	// Finding the device mounted to mountPath.
	diskDev := ""
	for i := range mps {
		if mountPath == mps[i].Path {
			diskDev = mps[i].Device
			break
		}
	}

	// Find all references to the device.
	var refs []string
	for i := range mps {
		if mps[i].Device == diskDev || mps[i].Device == mountPath {
			if mps[i].Path != mountPath {
				refs = append(refs, mps[i].Path)
			}
		}
	}
	return refs, nil
}

// GetDeviceNameFromMount given a mnt point, find the device from /proc/mounts
// returns the device name, reference count, and error code.
func GetDeviceNameFromMount(mounter Interface, mountPath string) (string, int, error) {
	mps, err := mounter.List()
	if err != nil {
		return "", 0, err
	}

	// Find the device name.
	// FIXME if multiple devices mounted on the same mount path, only the first one is returned.
	device := ""
	// If mountPath is symlink, need get its target path.
	slTarget, err := filepath.EvalSymlinks(mountPath)
	if err != nil {
		slTarget = mountPath
	}
	for i := range mps {
		if mps[i].Path == slTarget {
			device = mps[i].Device
			break
		}
	}

	// Find all references to the device.
	refCount := 0
	for i := range mps {
		if mps[i].Device == device {
			refCount++
		}
	}
	return device, refCount, nil
}

// IsNotMountPoint determines if a directory is a mountpoint.
// It should return ErrNotExist when the directory does not exist.
// IsNotMountPoint is more expensive than IsLikelyNotMountPoint.
// IsNotMountPoint detects bind mounts in linux.
// IsNotMountPoint enumerates all the mountpoints using List() and
// the list of mountpoints may be large, then it uses
// isMountPointMatch to evaluate whether the directory is a mountpoint.
func IsNotMountPoint(mounter Interface, file string) (bool, error) {
	// IsLikelyNotMountPoint provides a quick check
	// to determine whether file IS A mountpoint.
	notMnt, notMntErr := mounter.IsLikelyNotMountPoint(file)
	if notMntErr != nil && os.IsPermission(notMntErr) {
		// We were not allowed to do the simple stat() check, e.g. on NFS with
		// root_squash. Fall back to /proc/mounts check below.
		notMnt = true
		notMntErr = nil
	}
	if notMntErr != nil {
		return notMnt, notMntErr
	}
	// identified as mountpoint, so return this fact.
	if notMnt == false {
		return notMnt, nil
	}

	// Resolve any symlinks in file, kernel would do the same and use the resolved path in /proc/mounts.
	resolvedFile, err := filepath.EvalSymlinks(file)
	if err != nil {
		return true, err
	}

	// check all mountpoints since IsLikelyNotMountPoint
	// is not reliable for some mountpoint types.
	mountPoints, mountPointsErr := mounter.List()
	if mountPointsErr != nil {
		return notMnt, mountPointsErr
	}
	for _, mp := range mountPoints {
		if isMountPointMatch(mp, resolvedFile) {
			notMnt = false
			break
		}
	}
	return notMnt, nil
}

// MakeBindOpts detects whether a bind mount is being requested and makes the remount options to
// use in case of bind mount, due to the fact that bind mount doesn't respect mount options.
// The list equals:
//
//	options - 'bind' + 'remount' (no duplicate)
func MakeBindOpts(options []string) (bool, []string, []string) {
	bind, bindOpts, bindRemountOpts, _ := MakeBindOptsSensitive(options, nil /* sensitiveOptions */)
	return bind, bindOpts, bindRemountOpts
}

// MakeBindOptsSensitive is the same as MakeBindOpts but this method allows
// sensitiveOptions to be passed in a separate parameter from the normal mount
// options and ensures the sensitiveOptions are never logged. This method should
// be used by callers that pass sensitive material (like passwords) as mount
// options.
func MakeBindOptsSensitive(options []string, sensitiveOptions []string) (bool, []string, []string, []string) {
	// Because we have an FD opened on the subpath bind mount, the "bind" option
	// needs to be included, otherwise the mount target will error as busy if you
	// remount as readonly.
	//
	// As a consequence, all read only bind mounts will no longer change the underlying
	// volume mount to be read only.
	bindRemountOpts := []string{"bind", "remount"}
	bindRemountSensitiveOpts := []string{}
	bind := false
	bindOpts := []string{"bind"}

	// _netdev is a userspace mount option and does not automatically get added when
	// bind mount is created and hence we must carry it over.
	if checkForNetDev(options, sensitiveOptions) {
		bindOpts = append(bindOpts, "_netdev")
	}

	for _, option := range options {
		switch option {
		case "bind":
			bind = true
		case "remount":
		default:
			bindRemountOpts = append(bindRemountOpts, option)
		}
	}

	for _, sensitiveOption := range sensitiveOptions {
		switch sensitiveOption {
		case "bind":
			bind = true
		case "remount":
		default:
			bindRemountSensitiveOpts = append(bindRemountSensitiveOpts, sensitiveOption)
		}
	}

	return bind, bindOpts, bindRemountOpts, bindRemountSensitiveOpts
}

func checkForNetDev(options []string, sensitiveOptions []string) bool {
	for _, option := range options {
		if option == "_netdev" {
			return true
		}
	}
	for _, sensitiveOption := range sensitiveOptions {
		if sensitiveOption == "_netdev" {
			return true
		}
	}
	return false
}

// PathWithinBase checks if give path is within given base directory.
func PathWithinBase(fullPath, basePath string) bool {
	rel, err := filepath.Rel(basePath, fullPath)
	if err != nil {
		return false
	}
	if StartsWithBackstep(rel) {
		// Needed to escape the base path.
		return false
	}
	return true
}

// StartsWithBackstep checks if the given path starts with a backstep segment.
func StartsWithBackstep(rel string) bool {
	// normalize to / and check for ../
	return rel == ".." || strings.HasPrefix(filepath.ToSlash(rel), "../")
}

// sanitizedOptionsForLogging will return a comma separated string containing
// options and sensitiveOptions. Each entry in sensitiveOptions will be
// replaced with the string sensitiveOptionsRemoved
// e.g. o1,o2,<masked>,<masked>
func sanitizedOptionsForLogging(options []string, sensitiveOptions []string) string {
	separator := ""
	if len(options) > 0 && len(sensitiveOptions) > 0 {
		separator = ","
	}

	sensitiveOptionsStart := ""
	sensitiveOptionsEnd := ""
	if len(sensitiveOptions) > 0 {
		sensitiveOptionsStart = strings.Repeat(sensitiveOptionsRemoved+",", len(sensitiveOptions)-1)
		sensitiveOptionsEnd = sensitiveOptionsRemoved
	}

	return strings.Join(options, ",") +
		separator +
		sensitiveOptionsStart +
		sensitiveOptionsEnd
}
<|endoftext|>


<commit_before>// Copyright 2015 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 main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fsutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"golang.org/x/net/context"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s [flags] <mount-point>\n", os.Args[0])
	flag.PrintDefaults()
}

var fBucketName = flag.String("bucket", "", "Name of GCS bucket to mount.")
var fReadOnly = flag.Bool("read_only", false, "Mount in read-only mode.")

var fTempDir = flag.String(
	"temp_dir", "",
	"The temporary directory in which to store local copies of GCS objects. "+
		"If empty, the system default (probably /tmp) will be used.")

var fTempDirLimit = flag.Int64(
	"temp_dir_bytes", 1<<31,
	"A desired limit on the number of bytes used in --temp_dir. May be exceeded "+
		"for dirty files that have not been flushed or closed.")

var fGCSChunkSize = flag.Uint64(
	"gcs_chunk_size", 1<<24,
	"If set to a non-zero value N, split up GCS objects into multiple chunks of "+
		"size at most N when reading, and do not read or cache unnecessary chunks.")

var fImplicitDirs = flag.Bool(
	"implicit_dirs",
	false,
	"Implicitly define directories based on their content. See "+
		"docs/semantics.md.")

var fSupportNlink = flag.Bool(
	"support_nlink",
	false,
	"Return meaningful values for nlink from fstat(2). See docs/semantics.md.")

var fStatCacheTTL = flag.String(
	"stat_cache_ttl",
	"1m",
	"If non-empty, a duration specifying how long to cache StatObject results "+
		"from GCS, e.g. \"2s\" or \"15ms\". See docs/semantics.md for more.")

var fTypeCacheTTL = flag.String(
	"type_cache_ttl",
	"1m",
	"If non-empty, a duration specifying how long to cache name -> file/dir "+
		"type mappings in directory inodes, e.g. \"2s\" or \"15ms\". "+
		"See docs/semantics.md.")

func getBucketName() string {
	s := *fBucketName
	if s == "" {
		fmt.Println("You must set --bucket.")
		os.Exit(1)
	}

	return s
}

func registerSIGINTHandler(mountPoint string) {
	// Register for SIGINT.
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	// Start a goroutine that will unmount when the signal is received.
	go func() {
		for {
			<-signalChan
			log.Println("Received SIGINT, attempting to unmount...")

			err := fuse.Unmount(mountPoint)
			if err != nil {
				log.Printf("Failed to unmount in response to SIGINT: %v", err)
			} else {
				log.Printf("Successfully unmounted in response to SIGINT.")
				return
			}
		}
	}()
}

func getBucket() (b gcs.Bucket) {
	// Set up a GCS connection.
	log.Println("Initializing GCS connection.")
	conn, err := getConn()
	if err != nil {
		log.Fatal("Couldn't get GCS connection: ", err)
	}

	// Extract the appropriate bucket.
	b = conn.GetBucket(getBucketName())

	// Enable cached StatObject results, if appropriate.
	if *fStatCacheTTL != "" {
		ttl, err := time.ParseDuration(*fStatCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --stat_cache_ttl: %v", err)
			return
		}

		const cacheCapacity = 4096
		b = gcscaching.NewFastStatBucket(
			ttl,
			gcscaching.NewStatCache(cacheCapacity),
			timeutil.RealClock(),
			b)
	}

	return
}

func main() {
	// Make logging output better.
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	// Set up flags.
	flag.Usage = usage
	flag.Parse()

	// Grab the mount point.
	if flag.NArg() != 1 {
		usage()
		os.Exit(1)
	}

	mountPoint := flag.Arg(0)

	// Parse --type_cache_ttl
	var typeCacheTTL time.Duration
	if *fTypeCacheTTL != "" {
		var err error
		typeCacheTTL, err = time.ParseDuration(*fTypeCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --type_cache_ttl: %v", err)
			return
		}
	}

	// Sanity check: make sure the temporary directory exists and is writable
	// currently. This gives a better user experience than harder to debug EIO
	// errors when reading files in the future.
	if *fTempDir != "" {
		f, err := fsutil.AnonymousFile(*fTempDir)
		f.Close()

		if err != nil {
			log.Fatalf(
				"Error writing to temporary directory (%q); are you sure it exists "+
					"with the correct permissions?",
				err.Error())
		}
	}

	// Create a file system server.
	serverCfg := &fs.ServerConfig{
		Clock:               timeutil.RealClock(),
		Bucket:              getBucket(),
		TempDir:             *fTempDir,
		TempDirLimit:        *fTempDirLimit,
		GCSChunkSize:        *fGCSChunkSize,
		ImplicitDirectories: *fImplicitDirs,
		SupportNlink:        *fSupportNlink,
		DirTypeCacheTTL:     typeCacheTTL,
	}

	server, err := fs.NewServer(serverCfg)
	if err != nil {
		log.Fatal("fs.NewServer:", err)
	}

	// Mount the file system.
	mountCfg := &fuse.MountConfig{
		ReadOnly: *fReadOnly,
	}

	mountedFS, err := fuse.Mount(mountPoint, server, mountCfg)
	if err != nil {
		log.Fatal("Mount:", err)
	}

	log.Println("File system has been successfully mounted.")

	// Let the user unmount with Ctrl-C (SIGINT).
	registerSIGINTHandler(mountedFS.Dir())

	// Wait for it to be unmounted.
	if err := mountedFS.Join(context.Background()); err != nil {
		log.Fatal("MountedFileSystem.Join:", err)
	}

	log.Println("Successfully exiting.")
}
<commit_msg>Set a file system name.<commit_after>// Copyright 2015 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 main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fsutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"golang.org/x/net/context"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s [flags] <mount-point>\n", os.Args[0])
	flag.PrintDefaults()
}

var fBucketName = flag.String("bucket", "", "Name of GCS bucket to mount.")
var fReadOnly = flag.Bool("read_only", false, "Mount in read-only mode.")

var fTempDir = flag.String(
	"temp_dir", "",
	"The temporary directory in which to store local copies of GCS objects. "+
		"If empty, the system default (probably /tmp) will be used.")

var fTempDirLimit = flag.Int64(
	"temp_dir_bytes", 1<<31,
	"A desired limit on the number of bytes used in --temp_dir. May be exceeded "+
		"for dirty files that have not been flushed or closed.")

var fGCSChunkSize = flag.Uint64(
	"gcs_chunk_size", 1<<24,
	"If set to a non-zero value N, split up GCS objects into multiple chunks of "+
		"size at most N when reading, and do not read or cache unnecessary chunks.")

var fImplicitDirs = flag.Bool(
	"implicit_dirs",
	false,
	"Implicitly define directories based on their content. See "+
		"docs/semantics.md.")

var fSupportNlink = flag.Bool(
	"support_nlink",
	false,
	"Return meaningful values for nlink from fstat(2). See docs/semantics.md.")

var fStatCacheTTL = flag.String(
	"stat_cache_ttl",
	"1m",
	"If non-empty, a duration specifying how long to cache StatObject results "+
		"from GCS, e.g. \"2s\" or \"15ms\". See docs/semantics.md for more.")

var fTypeCacheTTL = flag.String(
	"type_cache_ttl",
	"1m",
	"If non-empty, a duration specifying how long to cache name -> file/dir "+
		"type mappings in directory inodes, e.g. \"2s\" or \"15ms\". "+
		"See docs/semantics.md.")

func getBucketName() string {
	s := *fBucketName
	if s == "" {
		fmt.Println("You must set --bucket.")
		os.Exit(1)
	}

	return s
}

func registerSIGINTHandler(mountPoint string) {
	// Register for SIGINT.
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	// Start a goroutine that will unmount when the signal is received.
	go func() {
		for {
			<-signalChan
			log.Println("Received SIGINT, attempting to unmount...")

			err := fuse.Unmount(mountPoint)
			if err != nil {
				log.Printf("Failed to unmount in response to SIGINT: %v", err)
			} else {
				log.Printf("Successfully unmounted in response to SIGINT.")
				return
			}
		}
	}()
}

func getBucket() (b gcs.Bucket) {
	// Set up a GCS connection.
	log.Println("Initializing GCS connection.")
	conn, err := getConn()
	if err != nil {
		log.Fatal("Couldn't get GCS connection: ", err)
	}

	// Extract the appropriate bucket.
	b = conn.GetBucket(getBucketName())

	// Enable cached StatObject results, if appropriate.
	if *fStatCacheTTL != "" {
		ttl, err := time.ParseDuration(*fStatCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --stat_cache_ttl: %v", err)
			return
		}

		const cacheCapacity = 4096
		b = gcscaching.NewFastStatBucket(
			ttl,
			gcscaching.NewStatCache(cacheCapacity),
			timeutil.RealClock(),
			b)
	}

	return
}

func main() {
	// Make logging output better.
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	// Set up flags.
	flag.Usage = usage
	flag.Parse()

	// Grab the mount point.
	if flag.NArg() != 1 {
		usage()
		os.Exit(1)
	}

	mountPoint := flag.Arg(0)

	// Parse --type_cache_ttl
	var typeCacheTTL time.Duration
	if *fTypeCacheTTL != "" {
		var err error
		typeCacheTTL, err = time.ParseDuration(*fTypeCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --type_cache_ttl: %v", err)
			return
		}
	}

	// Sanity check: make sure the temporary directory exists and is writable
	// currently. This gives a better user experience than harder to debug EIO
	// errors when reading files in the future.
	if *fTempDir != "" {
		f, err := fsutil.AnonymousFile(*fTempDir)
		f.Close()

		if err != nil {
			log.Fatalf(
				"Error writing to temporary directory (%q); are you sure it exists "+
					"with the correct permissions?",
				err.Error())
		}
	}

	// Create a file system server.
	bucket := getBucket()
	serverCfg := &fs.ServerConfig{
		Clock:               timeutil.RealClock(),
		Bucket:              bucket,
		TempDir:             *fTempDir,
		TempDirLimit:        *fTempDirLimit,
		GCSChunkSize:        *fGCSChunkSize,
		ImplicitDirectories: *fImplicitDirs,
		SupportNlink:        *fSupportNlink,
		DirTypeCacheTTL:     typeCacheTTL,
	}

	server, err := fs.NewServer(serverCfg)
	if err != nil {
		log.Fatal("fs.NewServer:", err)
	}

	// Mount the file system.
	mountCfg := &fuse.MountConfig{
		FSName:   bucket.Name(),
		ReadOnly: *fReadOnly,
	}

	mountedFS, err := fuse.Mount(mountPoint, server, mountCfg)
	if err != nil {
		log.Fatal("Mount:", err)
	}

	log.Println("File system has been successfully mounted.")

	// Let the user unmount with Ctrl-C (SIGINT).
	registerSIGINTHandler(mountedFS.Dir())

	// Wait for it to be unmounted.
	if err := mountedFS.Join(context.Background()); err != nil {
		log.Fatal("MountedFileSystem.Join:", err)
	}

	log.Println("Successfully exiting.")
}
<|endoftext|>


<commit_before>package emulator

import (
    "github.com/kierdavis/avr"
    "github.com/kierdavis/avr/spec"
    "log"
)

// An Emulator encapsulates the state of a processor.
type Emulator struct {
    Spec        *spec.MCUSpec
    regions     []Region
    ports       map[avr.PortRef]Port
    prog        []uint16
    ram         []uint8
    pc          uint32
    sp          uint16
    rampx       uint8
    rampy       uint8
    rampz       uint8
    rampd       uint8
    eind        uint8
    regs        [32]uint8
    flags       [8]uint8
    logging     bool
    excessTicks uint
}

// NewEmulator creates and returns an initialised Emulator for the given MCUSpec.
func NewEmulator(mcuSpec *spec.MCUSpec) (em *Emulator) {
    em = &Emulator{
        Spec:        mcuSpec,
        regions:     make([]Region, len(mcuSpec.Regions)),
        ports:       make(map[avr.PortRef]Port),
        prog:        make([]uint16, 1<<mcuSpec.LogProgMemSize),
        ram:         make([]uint8, 1<<mcuSpec.LogRAMSize),
        pc:          0,
    }

    // register standard ports
    em.RegisterPortByName("SPL", SplPort{em})
    em.RegisterPortByName("SPH", SphPort{em})
    em.RegisterPortByName("SREG", SregPort{em})

    // create memory regions
    for i, regionSpec_ := range mcuSpec.Regions {
        switch regionSpec := regionSpec_.(type) {
        case spec.RegsRegionSpec:
            em.regions[i] = RegsRegion{em, regionSpec}
        case spec.IORegionSpec:
            em.regions[i] = IORegion{em, regionSpec}
        case spec.RAMRegionSpec:
            em.regions[i] = RAMRegion{em, regionSpec}
        }
    }

    return em
}

func (em *Emulator) SetLogging(enabled bool) {
    em.logging = enabled
}

func (em *Emulator) RegisterPort(pref avr.PortRef, port Port) {
    em.ports[pref] = port
}

func (em *Emulator) RegisterPortByName(name string, port Port) (ok bool) {
    pref, ok := em.Spec.Ports[name]
    if !ok {
        return false
    }
    em.RegisterPort(pref, port)
    return true
}

func (em *Emulator) UnregisterPort(pref avr.PortRef) {
    delete(em.ports, pref)
}

func (em *Emulator) UnregisterPortByName(name string) (ok bool) {
    pref, ok := em.Spec.Ports[name]
    if !ok {
        return false
    }
    em.UnregisterPort(pref)
    return true
}

func (em *Emulator) InterruptsEnabled() bool {
    return em.flags[avr.FlagI] != 0
}

func (em *Emulator) Interrupt(num uint) {
    if em.InterruptsEnabled() {
        //log.Printf("int %d (mem = $%02x $%02x $%02x $%02x)", num, em.loadDataByte(0x0106), em.loadDataByte(0x0107), em.loadDataByte(0x0108), em.loadDataByte(0x0109))
        em.flags[avr.FlagI] = 0
        em.pushPC()
        em.pc = uint32(num * em.Spec.InterruptVectorSize)
    }
}

func (em *Emulator) InterruptByName(name string) (ok bool) {
    num, ok := em.Spec.Interrupts[name]
    if !ok {
        return false
    }
    em.Interrupt(num)
    return true
}

func (em *Emulator) Run(ticks uint) {
    // subtract ticks that were executed on the last call to Run
    ticksExecuted := em.excessTicks

    reducedCore := em.Spec.Family == spec.ReducedCore

    for ticksExecuted < ticks {
        word := em.fetchProgWord()
        inst := Decode(word, reducedCore)
        if inst < 0 {
            em.warn(InvalidInstructionWarning{em.pc - 1, word})
            ticksExecuted++
            continue
        }

        if !em.Spec.Available[inst] {
            em.warn(UnavailableInstructionWarning{em.pc - 1, inst, em.Spec})
            ticksExecuted++
            continue
        }

        handler := handlers[inst]
        ticksExecuted += uint(handler(em, word))
    }

    em.excessTicks = ticksExecuted - ticks
}

// Copy program words from buf into program memory starting at the given address.
// The method panics if the address is out of range at any point (the size of the
// program memory is equal to 1 << em.Spec.LogProgMemSize).
func (em *Emulator) WriteProg(address uint16, buf []uint16) {
    for _, word := range buf {
        em.prog[address] = word
        address++
    }
}

func (em *Emulator) fetchProgWord() (word uint16) {
    word = em.prog[em.pc]
    em.pc = (em.pc + 1) & ((1 << em.Spec.LogProgMemSize) - 1)
    return word
}

func (em *Emulator) demap(addr uint16) (r Region) {
    // TODO: optimise
    for _, r := range em.regions {
        if r.Contains(addr) {
            return r
        }
    }

    return nil
}

func (em *Emulator) loadDataByte(addr uint16) uint8 {
    r := em.demap(addr)
    if r != nil {
        return r.Load(addr)
    } else {
        em.warn(UnmappedAddressWarning{em.pc - 1, addr})
        return 0
    }
}

func (em *Emulator) storeDataByte(addr uint16, val uint8) {
    r := em.demap(addr)
    if r != nil {
        r.Store(addr, val)
    } else {
        em.warn(UnmappedAddressWarning{em.pc - 1, addr})
    }
}

func (em *Emulator) push(val uint8) {
    em.storeDataByte(em.sp, val)
    em.sp--
}

func (em *Emulator) pop() uint8 {
    em.sp++
    return em.loadDataByte(em.sp)
}

func (em *Emulator) pushPC() {
    if em.Spec.LogProgMemSize > 16 { // pc is 3 bytes
        em.push(uint8(em.pc >> 16))
    }

    em.push(uint8(em.pc >> 8))
    em.push(uint8(em.pc))
}

func (em *Emulator) popPC() {
    em.pc = uint32(em.pop())
    em.pc |= uint32(em.pop()) << 8

    if em.Spec.LogProgMemSize > 16 { // pc is 3 bytes
        em.pc |= uint32(em.pop()) << 16
    }
}

func (em *Emulator) readPort(bankNum uint, index uint16) uint8 {
    port, ok := em.ports[avr.PortRef{bankNum, index}]
    if !ok {
        em.warn(UnmappedPortWarning{em.pc - 1, bankNum, index})
        return 0
    }

    return port.Read()
}

func (em *Emulator) writePort(bankNum uint, index uint16, val uint8) {
    port, ok := em.ports[avr.PortRef{bankNum, index}]
    if !ok {
        em.warn(UnmappedPortWarning{em.pc - 1, bankNum, index})
        return
    }

    port.Write(val)
}

// Skip the next instruction. Returns 1 if one word was skipped or 2 if two
// words were skipped.
func (em *Emulator) skip() (cycles int) {
    word := em.fetchProgWord()
    inst := Decode(word, em.Spec.Family == spec.ReducedCore)
    if inst.IsTwoWord() {
        em.fetchProgWord()
        return 2
    }
    return 1
}

// Log a warning, if warning logging is enabled. Warnings include events such as
// * invalid instructions
// * instructions not available on this particular MCU
// * accesses to an unmapped data memory address
// which are ignored by a real MCU but are often indicative of software errors.
func (em *Emulator) warn(w Warning) {
    if em.logging {
        log.Printf("[avr/emulator:(*Emulator).warn] %s\n", w.String())
    }
}
<commit_msg>Decrease tick time from 38 to 36 nanoseconds by not computing PC increment mask every time an instruction is fetched<commit_after>package emulator

import (
    "github.com/kierdavis/avr"
    "github.com/kierdavis/avr/spec"
    "log"
)

// An Emulator encapsulates the state of a processor.
type Emulator struct {
    Spec        *spec.MCUSpec
    regions     []Region
    ports       map[avr.PortRef]Port
    prog        []uint16
    ram         []uint8
    pc          uint32
    pcmask      uint32
    sp          uint16
    rampx       uint8
    rampy       uint8
    rampz       uint8
    rampd       uint8
    eind        uint8
    regs        [32]uint8
    flags       [8]uint8
    logging     bool
    excessTicks uint
}

// NewEmulator creates and returns an initialised Emulator for the given MCUSpec.
func NewEmulator(mcuSpec *spec.MCUSpec) (em *Emulator) {
    em = &Emulator{
        Spec:        mcuSpec,
        regions:     make([]Region, len(mcuSpec.Regions)),
        ports:       make(map[avr.PortRef]Port),
        prog:        make([]uint16, 1<<mcuSpec.LogProgMemSize),
        ram:         make([]uint8, 1<<mcuSpec.LogRAMSize),
        pc:          0,
        pcmask:      (1 << mcuSpec.LogProgMemSize) - 1,
    }

    // register standard ports
    em.RegisterPortByName("SPL", SplPort{em})
    em.RegisterPortByName("SPH", SphPort{em})
    em.RegisterPortByName("SREG", SregPort{em})

    // create memory regions
    for i, regionSpec_ := range mcuSpec.Regions {
        switch regionSpec := regionSpec_.(type) {
        case spec.RegsRegionSpec:
            em.regions[i] = RegsRegion{em, regionSpec}
        case spec.IORegionSpec:
            em.regions[i] = IORegion{em, regionSpec}
        case spec.RAMRegionSpec:
            em.regions[i] = RAMRegion{em, regionSpec}
        }
    }

    return em
}

func (em *Emulator) SetLogging(enabled bool) {
    em.logging = enabled
}

func (em *Emulator) RegisterPort(pref avr.PortRef, port Port) {
    em.ports[pref] = port
}

func (em *Emulator) RegisterPortByName(name string, port Port) (ok bool) {
    pref, ok := em.Spec.Ports[name]
    if !ok {
        return false
    }
    em.RegisterPort(pref, port)
    return true
}

func (em *Emulator) UnregisterPort(pref avr.PortRef) {
    delete(em.ports, pref)
}

func (em *Emulator) UnregisterPortByName(name string) (ok bool) {
    pref, ok := em.Spec.Ports[name]
    if !ok {
        return false
    }
    em.UnregisterPort(pref)
    return true
}

func (em *Emulator) InterruptsEnabled() bool {
    return em.flags[avr.FlagI] != 0
}

func (em *Emulator) Interrupt(num uint) {
    if em.InterruptsEnabled() {
        //log.Printf("int %d (mem = $%02x $%02x $%02x $%02x)", num, em.loadDataByte(0x0106), em.loadDataByte(0x0107), em.loadDataByte(0x0108), em.loadDataByte(0x0109))
        em.flags[avr.FlagI] = 0
        em.pushPC()
        em.pc = uint32(num * em.Spec.InterruptVectorSize)
    }
}

func (em *Emulator) InterruptByName(name string) (ok bool) {
    num, ok := em.Spec.Interrupts[name]
    if !ok {
        return false
    }
    em.Interrupt(num)
    return true
}

func (em *Emulator) Run(ticks uint) {
    // subtract ticks that were executed on the last call to Run
    ticksExecuted := em.excessTicks

    reducedCore := em.Spec.Family == spec.ReducedCore

    for ticksExecuted < ticks {
        word := em.fetchProgWord()
        inst := Decode(word, reducedCore)
        if inst < 0 {
            em.warn(InvalidInstructionWarning{em.pc - 1, word})
            ticksExecuted++
            continue
        }

        if !em.Spec.Available[inst] {
            em.warn(UnavailableInstructionWarning{em.pc - 1, inst, em.Spec})
            ticksExecuted++
            continue
        }

        handler := handlers[inst]
        ticksExecuted += uint(handler(em, word))
    }

    em.excessTicks = ticksExecuted - ticks
}

// Copy program words from buf into program memory starting at the given address.
// The method panics if the address is out of range at any point (the size of the
// program memory is equal to 1 << em.Spec.LogProgMemSize).
func (em *Emulator) WriteProg(address uint16, buf []uint16) {
    for _, word := range buf {
        em.prog[address] = word
        address++
    }
}

func (em *Emulator) fetchProgWord() (word uint16) {
    word = em.prog[em.pc]
    em.pc = (em.pc + 1) & em.pcmask
    return word
}

func (em *Emulator) demap(addr uint16) (r Region) {
    // TODO: optimise
    for _, r := range em.regions {
        if r.Contains(addr) {
            return r
        }
    }

    return nil
}

func (em *Emulator) loadDataByte(addr uint16) uint8 {
    r := em.demap(addr)
    if r != nil {
        return r.Load(addr)
    } else {
        em.warn(UnmappedAddressWarning{em.pc - 1, addr})
        return 0
    }
}

func (em *Emulator) storeDataByte(addr uint16, val uint8) {
    r := em.demap(addr)
    if r != nil {
        r.Store(addr, val)
    } else {
        em.warn(UnmappedAddressWarning{em.pc - 1, addr})
    }
}

func (em *Emulator) push(val uint8) {
    em.storeDataByte(em.sp, val)
    em.sp--
}

func (em *Emulator) pop() uint8 {
    em.sp++
    return em.loadDataByte(em.sp)
}

func (em *Emulator) pushPC() {
    if em.Spec.LogProgMemSize > 16 { // pc is 3 bytes
        em.push(uint8(em.pc >> 16))
    }

    em.push(uint8(em.pc >> 8))
    em.push(uint8(em.pc))
}

func (em *Emulator) popPC() {
    em.pc = uint32(em.pop())
    em.pc |= uint32(em.pop()) << 8

    if em.Spec.LogProgMemSize > 16 { // pc is 3 bytes
        em.pc |= uint32(em.pop()) << 16
    }
}

func (em *Emulator) readPort(bankNum uint, index uint16) uint8 {
    port, ok := em.ports[avr.PortRef{bankNum, index}]
    if !ok {
        em.warn(UnmappedPortWarning{em.pc - 1, bankNum, index})
        return 0
    }

    return port.Read()
}

func (em *Emulator) writePort(bankNum uint, index uint16, val uint8) {
    port, ok := em.ports[avr.PortRef{bankNum, index}]
    if !ok {
        em.warn(UnmappedPortWarning{em.pc - 1, bankNum, index})
        return
    }

    port.Write(val)
}

// Skip the next instruction. Returns 1 if one word was skipped or 2 if two
// words were skipped.
func (em *Emulator) skip() (cycles int) {
    word := em.fetchProgWord()
    inst := Decode(word, em.Spec.Family == spec.ReducedCore)
    if inst.IsTwoWord() {
        em.fetchProgWord()
        return 2
    }
    return 1
}

// Log a warning, if warning logging is enabled. Warnings include events such as
// * invalid instructions
// * instructions not available on this particular MCU
// * accesses to an unmapped data memory address
// which are ignored by a real MCU but are often indicative of software errors.
func (em *Emulator) warn(w Warning) {
    if em.logging {
        log.Printf("[avr/emulator:(*Emulator).warn] %s\n", w.String())
    }
}
<|endoftext|>


<commit_before>package redsync

import (
	"crypto/rand"
	"encoding/base64"
	"time"

	"github.com/gomodule/redigo/redis"
)

// A DelayFunc is used to decide the amount of time to wait between retries.
type DelayFunc func(tries int) time.Duration

// A Mutex is a distributed mutual exclusion lock.
type Mutex struct {
	name   string
	expiry time.Duration

	tries     int
	delayFunc DelayFunc

	factor float64

	quorum int

	genValueFunc func() (string, error)
	value        string
	until        time.Time

	pools []Pool
}

// Lock locks m. In case it returns an error on failure, you may retry to acquire the lock by calling this method again.
func (m *Mutex) Lock() error {
	value, err := m.genValueFunc()
	if err != nil {
		return err
	}

	for i := 0; i < m.tries; i++ {
		if i != 0 {
			time.Sleep(m.delayFunc(i))
		}

		start := time.Now()

		n := m.actOnPoolsAsync(func(pool Pool) bool {
			return m.acquire(pool, value)
		})

		until := time.Now().Add(m.expiry - time.Now().Sub(start) - time.Duration(int64(float64(m.expiry)*m.factor)) + 2*time.Millisecond)
		if n >= m.quorum && time.Now().Before(until) {
			m.value = value
			m.until = until
			return nil
		}
		m.actOnPoolsAsync(func(pool Pool) bool {
			return m.release(pool, value)
		})
	}

	return ErrFailed
}

// Unlock unlocks m and returns the status of unlock.
func (m *Mutex) Unlock() bool {
	n := m.actOnPoolsAsync(func(pool Pool) bool {
		return m.release(pool, m.value)
	})
	return n >= m.quorum
}

// Extend resets the mutex's expiry and returns the status of expiry extension.
func (m *Mutex) Extend() bool {
	n := m.actOnPoolsAsync(func(pool Pool) bool {
		return m.touch(pool, m.value, int(m.expiry/time.Millisecond))
	})
	return n >= m.quorum
}

func genValue() (string, error) {
	b := make([]byte, 32)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(b), nil
}

func (m *Mutex) acquire(pool Pool, value string) bool {
	conn := pool.Get()
	defer conn.Close()
	reply, err := redis.String(conn.Do("SET", m.name, value, "NX", "PX", int(m.expiry/time.Millisecond)))
	return err == nil && reply == "OK"
}

var deleteScript = redis.NewScript(1, `
	if redis.call("GET", KEYS[1]) == ARGV[1] then
		return redis.call("DEL", KEYS[1])
	else
		return 0
	end
`)

func (m *Mutex) release(pool Pool, value string) bool {
	conn := pool.Get()
	defer conn.Close()
	status, err := deleteScript.Do(conn, m.name, value)
	return err == nil && status != 0
}

var touchScript = redis.NewScript(1, `
	if redis.call("GET", KEYS[1]) == ARGV[1] then
		return redis.call("pexpire", KEYS[1], ARGV[2])
	else
		return 0
	end
`)

func (m *Mutex) touch(pool Pool, value string, expiry int) bool {
	conn := pool.Get()
	defer conn.Close()
	status, err := touchScript.Do(conn, m.name, value, expiry)
	return err == nil && status != 0
}

func (m *Mutex) actOnPoolsAsync(actFn func(Pool) bool) int {
	ch := make(chan bool)
	for _, pool := range m.pools {
		go func(pool Pool) {
			ch <- actFn(pool)
		}(pool)
	}
	n := 0
	for range m.pools {
		if <-ch {
			n++
		}
	}
	return n
}
<commit_msg>Reduce default size of random value; closes #36<commit_after>package redsync

import (
	"crypto/rand"
	"encoding/base64"
	"time"

	"github.com/gomodule/redigo/redis"
)

// A DelayFunc is used to decide the amount of time to wait between retries.
type DelayFunc func(tries int) time.Duration

// A Mutex is a distributed mutual exclusion lock.
type Mutex struct {
	name   string
	expiry time.Duration

	tries     int
	delayFunc DelayFunc

	factor float64

	quorum int

	genValueFunc func() (string, error)
	value        string
	until        time.Time

	pools []Pool
}

// Lock locks m. In case it returns an error on failure, you may retry to acquire the lock by calling this method again.
func (m *Mutex) Lock() error {
	value, err := m.genValueFunc()
	if err != nil {
		return err
	}

	for i := 0; i < m.tries; i++ {
		if i != 0 {
			time.Sleep(m.delayFunc(i))
		}

		start := time.Now()

		n := m.actOnPoolsAsync(func(pool Pool) bool {
			return m.acquire(pool, value)
		})

		until := time.Now().Add(m.expiry - time.Now().Sub(start) - time.Duration(int64(float64(m.expiry)*m.factor)) + 2*time.Millisecond)
		if n >= m.quorum && time.Now().Before(until) {
			m.value = value
			m.until = until
			return nil
		}
		m.actOnPoolsAsync(func(pool Pool) bool {
			return m.release(pool, value)
		})
	}

	return ErrFailed
}

// Unlock unlocks m and returns the status of unlock.
func (m *Mutex) Unlock() bool {
	n := m.actOnPoolsAsync(func(pool Pool) bool {
		return m.release(pool, m.value)
	})
	return n >= m.quorum
}

// Extend resets the mutex's expiry and returns the status of expiry extension.
func (m *Mutex) Extend() bool {
	n := m.actOnPoolsAsync(func(pool Pool) bool {
		return m.touch(pool, m.value, int(m.expiry/time.Millisecond))
	})
	return n >= m.quorum
}

func genValue() (string, error) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(b), nil
}

func (m *Mutex) acquire(pool Pool, value string) bool {
	conn := pool.Get()
	defer conn.Close()
	reply, err := redis.String(conn.Do("SET", m.name, value, "NX", "PX", int(m.expiry/time.Millisecond)))
	return err == nil && reply == "OK"
}

var deleteScript = redis.NewScript(1, `
	if redis.call("GET", KEYS[1]) == ARGV[1] then
		return redis.call("DEL", KEYS[1])
	else
		return 0
	end
`)

func (m *Mutex) release(pool Pool, value string) bool {
	conn := pool.Get()
	defer conn.Close()
	status, err := deleteScript.Do(conn, m.name, value)
	return err == nil && status != 0
}

var touchScript = redis.NewScript(1, `
	if redis.call("GET", KEYS[1]) == ARGV[1] then
		return redis.call("pexpire", KEYS[1], ARGV[2])
	else
		return 0
	end
`)

func (m *Mutex) touch(pool Pool, value string, expiry int) bool {
	conn := pool.Get()
	defer conn.Close()
	status, err := touchScript.Do(conn, m.name, value, expiry)
	return err == nil && status != 0
}

func (m *Mutex) actOnPoolsAsync(actFn func(Pool) bool) int {
	ch := make(chan bool)
	for _, pool := range m.pools {
		go func(pool Pool) {
			ch <- actFn(pool)
		}(pool)
	}
	n := 0
	for range m.pools {
		if <-ch {
			n++
		}
	}
	return n
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"
	"unicode/utf8"

	"golang.org/x/crypto/ssh/terminal"

	_ "github.com/go-sql-driver/mysql"
)

const (
	// Amount of CSV write data to buffer between flushes.
	flushBufferSize = 26214400 // 25MB

	// Timeout length where ctrl+c is ignored.
	signalTimeout = 3 // Seconds

	// Timeout length to wait for a query string sent via stdin.
	stdinTimeout = 10 // Milliseconds
)

type (
	// dbInfo contains information necessary to connect to a database
	dbInfo struct {
		user    string
		pass    string
		host    string
		port    string
		charset string
	}
)

// ShowUsage prints a help screen
func showUsage() {
	fmt.Println(`
	mycsv usage:
	mycsv DB_COMMANDS [CSV OUTPUT FLAGS] [DEBUG FLAGS] [CSV OUTFILE] query

	EXAMPLES:
	mycsv -user=jprunier -pass= -file=my.csv -charset=utf8 -query="select * from jjp.example_table where filter in ('1', 'test', 'another')"
	echo "select * from mysql.plugin" | mycsv -user=jprunier -pass=mypass -host=remotedb > my.csv
	mycsv -user=jprunier -pass= -file=my.csv -d="|" -q="'" < queryfile

	DATABASE FLAGS
	==============
	-user: Database Username (required)
	-pass: Database Password (interactive prompt if blank)
	-host: Database Host (localhost assumed if blank)
	-port: Database Port (3306 default)
	-charset: Database character set (binary default)

	CSV FLAGS
	=========
	-file: CSV output filename (Write to stdout if not supplied)
	-query: MySQL query (required, can be sent via stdin redirection)
	-header: Print initial column name header line (true default)
	-d: CSV field delimiter ("," default)
	-q: CSV quote character ("\"" default)
	-e: CSV escape character ("\\" default)
	-t: CSV line terminator ("\n" default)
	-v: Print more information (false default)

	DEBUG FLAGS
	===========
	-debug_cpu: CPU debugging filename
	-debug_mem: Memory debugging filename

	`)
}

func main() {
	start := time.Now()

	// Profiling flags
	var cpuprofile = flag.String("debug_cpu", "", "CPU debugging filename")
	var memprofile = flag.String("debug_mem", "", "Memory debugging filename")

	// Database flags
	dbUser := flag.String("user", "", "Database Username (required)")
	dbPass := flag.String("pass", "", "Database Password (interactive prompt if blank)")
	dbHost := flag.String("host", "", "Database Host (localhost assumed if blank)")
	dbPort := flag.String("port", "3306", "Database Port")
	dbCharset := flag.String("charset", "binary", "Database character set")

	// CSV formatting flags
	csvDelimiter := flag.String("d", `,`, "CSV field delimiter")
	csvQuote := flag.String("q", `"`, "CSV quote character")
	csvEscape := flag.String("e", `\`, "CSV escape character")
	csvTerminator := flag.String("t", "\n", "CSV line terminator")

	// Other flags
	csvHeader := flag.Bool("header", true, "Print initial column name header line")
	csvFile := flag.String("file", "", "CSV output filename")
	sqlQuery := flag.String("query", "", "MySQL query")
	verbose := flag.Bool("v", false, "Print more information")

	// Parse flags
	flag.Parse()

	// Print usage
	if flag.NFlag() == 0 {
		showUsage()

		os.Exit(0)
	}

	// If query not provided read from standard in
	var query string
	queryChan := make(chan string)
	defer close(queryChan)
	if *sqlQuery == "" {
		go func() {
			b, err := ioutil.ReadAll(os.Stdin)
			checkErr(err)

			queryChan <- string(b)
		}()

		select {
		case q := <-queryChan:
			query = q
		case <-time.After(time.Millisecond * stdinTimeout):
			fmt.Fprintln(os.Stderr, "You must supply a query")
			os.Exit(1)
		}
	} else {
		query = *sqlQuery
	}

	// Make sure the query is a select
	if strings.ToLower(query[0:6]) != "select" {
		fmt.Fprintln(os.Stderr, "Query must be a select!")
		os.Exit(1)
	}

	// Create CSV output file if supplied, otherwise use standard out
	var writeTo string
	var writerDest io.Writer
	var err error
	if *csvFile == "" {
		writeTo = "standard out"
		writerDest = os.Stdout
	} else {
		f, err := os.Open(*csvFile)
		if err == nil {
			fmt.Fprintln(os.Stderr, *csvFile, "already exists!")
			fmt.Fprintln(os.Stderr, "Please remove it or use a different filename")
			os.Exit(1)
		}
		f.Close()
		writerDest, err = os.Create(*csvFile)
		writeTo = *csvFile
	}

	// Create a new CSV writer
	var i int
	CSVWriter := NewWriter(writerDest)
	CSVWriter.Delimiter, i = utf8.DecodeLastRuneInString(*csvDelimiter)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid delimiter character")
		os.Exit(1)
	}
	CSVWriter.Quote, i = utf8.DecodeLastRuneInString(*csvQuote)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid quote character")
		os.Exit(1)
	}
	CSVWriter.Escape, i = utf8.DecodeLastRuneInString(*csvEscape)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid escape character")
		os.Exit(1)
	}

	// Need literal string check here to see all 4 bytes instead of 2 (ascii 13 & 10)
	// Newline is default but check here in case it is manually passed in
	if *csvTerminator == `\r\n` {
		CSVWriter.Terminator = "\r\n"
	} else if *csvTerminator == `\n` {
		CSVWriter.Terminator = "\n"
	} else {
		CSVWriter.Terminator = *csvTerminator
	}

	if *verbose {
		fmt.Println("CSV output will be written to", writeTo)
	}

	// Check if Stdin has been redirected and reset so the user can be prompted for a password
	checkStdin()

	// Catch signals
	catchNotifications()

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *dbHost == "" {
		*dbHost = "127.0.0.1"
	}

	// Need to provide a target
	if *dbUser == "" {
		fmt.Fprintln(os.Stderr, "You must provide a user name!")
		os.Exit(1)
	}

	// If password is blank prompt user
	if *dbPass == "" {
		fmt.Println("Enter password: ")
		pwd, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			if err != io.EOF {
				checkErr(err)
			}
		}
		*dbPass = string(pwd)
	}

	// Populate dbInfo struct with flag values
	dbi := dbInfo{user: *dbUser, pass: *dbPass, host: *dbHost, port: *dbPort, charset: *dbCharset}

	// Create a *sql.DB connection to the source database
	db, err := dbi.Connect()
	defer db.Close()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Create channels
	dataChan := make(chan []sql.RawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)

	// Start reading & writing
	go readRows(db, query, dataChan, quitChan, goChan, *csvHeader)
	rowCount := writeCSV(CSVWriter, dataChan, goChan, *verbose)

	// Block on quitChan until readRows() completes
	<-quitChan
	close(quitChan)
	close(goChan)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if *verbose {
		fmt.Println()
		fmt.Println(rowCount, "rows written")
		fmt.Println("Total runtime =", time.Since(start))
	}
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

// Catch signals
func catchNotifications() {
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	checkErr(err)

	// Deal with SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			// Prevent exiting on accidental signal send
			if time.Now().Sub(timer) < time.Second*signalTimeout {
				terminal.Restore(int(os.Stdin.Fd()), state)
				os.Exit(0)
			}

			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, sig, "signal caught!")
			fmt.Fprintf(os.Stderr, "Send signal again within %v seconds to exit\n", signalTimeout)
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")

			timer = time.Now()
		}
	}()
}

// Create and return a database handle
func (dbi *dbInfo) Connect() (*sql.DB, error) {
	db, err := sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1&tls=skip-verify&charset="+dbi.charset)
	checkErr(err)

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// readRows executes a query and sends each row over a channel to be consumed
func readRows(db *sql.DB, query string, dataChan chan []sql.RawBytes, quitChan chan bool, goChan chan bool, csvHeader bool) {
	rows, err := db.Query(query)
	defer rows.Close()
	if err != nil {
		log.Print(err)
		os.Exit(1)
	}

	cols, err := rows.Columns()
	checkErr(err)

	// Write columns as a header line
	if csvHeader {
		headers := make([]sql.RawBytes, len(cols))
		for i, col := range cols {
			headers[i] = []byte(col)
		}
		dataChan <- headers
		<-goChan
	}

	// Need to scan into empty interface since we don't know how many columns a query might return
	scanVals := make([]interface{}, len(cols))
	vals := make([]sql.RawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		dataChan <- vals

		// Block and wait for writeRows() to signal back it has consumed the data
		// This is necessary because sql.RawBytes is a memory pointer and when rows.Next()
		// loops and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true
}

// writeCSV reads from a channel and writes CSV output
func writeCSV(w *Writer, dataChan chan []sql.RawBytes, goChan chan bool, verbose bool) uint {
	var rowsWritten uint
	var verboseCount uint

	if verbose {
		fmt.Println("A '.' will be shown for every 10,000 CSV rows written")
	}

	// Range over row results from readRows()
	for data := range dataChan {
		// Format the data to CSV and write
		size, err := w.Write(data)
		checkErr(err)

		// Visual write indicator when verbose is enabled
		rowsWritten++
		if verbose {
			verboseCount++
			if verboseCount == 10000 {
				fmt.Printf(".")
				verboseCount = 0
			}
		}

		// Flush CSV writer contents once it exceeds flushBufferSize
		if size > flushBufferSize {
			w.Flush()
			err = w.Error()
			checkErr(err)
		}

		// Signal back to readRows() it can loop and scan the next row
		goChan <- true
	}

	// Flush remaining CSV writer contents
	w.Flush()
	err := w.Error()
	checkErr(err)

	return rowsWritten
}
<commit_msg>Override default help<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"runtime/pprof"
	"strings"
	"time"
	"unicode/utf8"

	"golang.org/x/crypto/ssh/terminal"

	_ "github.com/go-sql-driver/mysql"
)

const (
	// Amount of CSV write data to buffer between flushes.
	flushBufferSize = 26214400 // 25MB

	// Timeout length where ctrl+c is ignored.
	signalTimeout = 3 // Seconds

	// Timeout length to wait for a query string sent via stdin.
	stdinTimeout = 10 // Milliseconds
)

type (
	// dbInfo contains information necessary to connect to a database
	dbInfo struct {
		user    string
		pass    string
		host    string
		port    string
		charset string
	}
)

// ShowUsage prints a help screen
func showUsage() {
	fmt.Println(`
	mycsv usage:
	mycsv DB_COMMANDS [CSV OUTPUT FLAGS] [DEBUG FLAGS] [CSV OUTFILE] query

	EXAMPLES:
	mycsv -user=jprunier -pass= -file=my.csv -charset=utf8 -query="select * from jjp.example_table where filter in ('1', 'test', 'another')"
	echo "select * from mysql.plugin" | mycsv -user=jprunier -pass=mypass -host=remotedb > my.csv
	mycsv -user=jprunier -pass= -file=my.csv -d="|" -q="'" < queryfile

	DATABASE FLAGS
	==============
	-user: Database Username (required)
	-pass: Database Password (interactive prompt if blank)
	-host: Database Host (localhost assumed if blank)
	-port: Database Port (3306 default)
	-charset: Database character set (binary default)

	CSV FLAGS
	=========
	-file: CSV output filename (Write to stdout if not supplied)
	-query: MySQL query (required, can be sent via stdin redirection)
	-header: Print initial column name header line (true default)
	-d: CSV field delimiter ("," default)
	-q: CSV quote character ("\"" default)
	-e: CSV escape character ("\\" default)
	-t: CSV line terminator ("\n" default)
	-v: Print more information (false default)

	DEBUG FLAGS
	===========
	-debug_cpu: CPU debugging filename
	-debug_mem: Memory debugging filename

	`)
}

func main() {
	start := time.Now()

	// Profiling flags
	var cpuprofile = flag.String("debug_cpu", "", "CPU debugging filename")
	var memprofile = flag.String("debug_mem", "", "Memory debugging filename")

	// Database flags
	dbUser := flag.String("user", "", "Database Username (required)")
	dbPass := flag.String("pass", "", "Database Password (interactive prompt if blank)")
	dbHost := flag.String("host", "", "Database Host (localhost assumed if blank)")
	dbPort := flag.String("port", "3306", "Database Port")
	dbCharset := flag.String("charset", "binary", "Database character set")

	// CSV formatting flags
	csvDelimiter := flag.String("d", `,`, "CSV field delimiter")
	csvQuote := flag.String("q", `"`, "CSV quote character")
	csvEscape := flag.String("e", `\`, "CSV escape character")
	csvTerminator := flag.String("t", "\n", "CSV line terminator")

	// Other flags
	csvHeader := flag.Bool("header", true, "Print initial column name header line")
	csvFile := flag.String("file", "", "CSV output filename")
	sqlQuery := flag.String("query", "", "MySQL query")
	verbose := flag.Bool("v", false, "Print more information")

	// Override default help
	help := flag.Bool("help", false, "Show usage")
	h := flag.Bool("h", false, "Show usage")

	// Parse flags
	flag.Parse()

	// Print usage
	if flag.NFlag() == 0 || *help == true || *h == true {
		showUsage()

		os.Exit(0)
	}

	// If query not provided read from standard in
	var query string
	queryChan := make(chan string)
	defer close(queryChan)
	if *sqlQuery == "" {
		go func() {
			b, err := ioutil.ReadAll(os.Stdin)
			checkErr(err)

			queryChan <- string(b)
		}()

		select {
		case q := <-queryChan:
			query = q
		case <-time.After(time.Millisecond * stdinTimeout):
			fmt.Fprintln(os.Stderr, "You must supply a query")
			os.Exit(1)
		}
	} else {
		query = *sqlQuery
	}

	// Make sure the query is a select
	if strings.ToLower(query[0:6]) != "select" {
		fmt.Fprintln(os.Stderr, "Query must be a select!")
		os.Exit(1)
	}

	// Create CSV output file if supplied, otherwise use standard out
	var writeTo string
	var writerDest io.Writer
	var err error
	if *csvFile == "" {
		writeTo = "standard out"
		writerDest = os.Stdout
	} else {
		f, err := os.Open(*csvFile)
		if err == nil {
			fmt.Fprintln(os.Stderr, *csvFile, "already exists!")
			fmt.Fprintln(os.Stderr, "Please remove it or use a different filename")
			os.Exit(1)
		}
		f.Close()
		writerDest, err = os.Create(*csvFile)
		writeTo = *csvFile
	}

	// Create a new CSV writer
	var i int
	CSVWriter := NewWriter(writerDest)
	CSVWriter.Delimiter, i = utf8.DecodeLastRuneInString(*csvDelimiter)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid delimiter character")
		os.Exit(1)
	}
	CSVWriter.Quote, i = utf8.DecodeLastRuneInString(*csvQuote)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid quote character")
		os.Exit(1)
	}
	CSVWriter.Escape, i = utf8.DecodeLastRuneInString(*csvEscape)
	if i == 0 {
		fmt.Fprintln(os.Stderr, "You must supply a valid escape character")
		os.Exit(1)
	}

	// Need literal string check here to see all 4 bytes instead of 2 (ascii 13 & 10)
	// Newline is default but check here in case it is manually passed in
	if *csvTerminator == `\r\n` {
		CSVWriter.Terminator = "\r\n"
	} else if *csvTerminator == `\n` {
		CSVWriter.Terminator = "\n"
	} else {
		CSVWriter.Terminator = *csvTerminator
	}

	if *verbose {
		fmt.Println("CSV output will be written to", writeTo)
	}

	// Check if Stdin has been redirected and reset so the user can be prompted for a password
	checkStdin()

	// Catch signals
	catchNotifications()

	// CPU Profiling
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		checkErr(err)
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// Default to localhost if no host or socket provided
	if *dbHost == "" {
		*dbHost = "127.0.0.1"
	}

	// Need to provide a target
	if *dbUser == "" {
		fmt.Fprintln(os.Stderr, "You must provide a user name!")
		os.Exit(1)
	}

	// If password is blank prompt user
	if *dbPass == "" {
		fmt.Println("Enter password: ")
		pwd, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			if err != io.EOF {
				checkErr(err)
			}
		}
		*dbPass = string(pwd)
	}

	// Populate dbInfo struct with flag values
	dbi := dbInfo{user: *dbUser, pass: *dbPass, host: *dbHost, port: *dbPort, charset: *dbCharset}

	// Create a *sql.DB connection to the source database
	db, err := dbi.Connect()
	defer db.Close()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Create channels
	dataChan := make(chan []sql.RawBytes)
	quitChan := make(chan bool)
	goChan := make(chan bool)

	// Start reading & writing
	go readRows(db, query, dataChan, quitChan, goChan, *csvHeader)
	rowCount := writeCSV(CSVWriter, dataChan, goChan, *verbose)

	// Block on quitChan until readRows() completes
	<-quitChan
	close(quitChan)
	close(goChan)

	// Memory Profiling
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		checkErr(err)
		pprof.WriteHeapProfile(f)
		defer f.Close()
	}

	if *verbose {
		fmt.Println()
		fmt.Println(rowCount, "rows written")
		fmt.Println("Total runtime =", time.Since(start))
	}
}

// Pass the buck error catching
func checkErr(e error) {
	if e != nil {
		log.Panic(e)
	}
}

// Catch signals
func catchNotifications() {
	state, err := terminal.GetState(int(os.Stdin.Fd()))
	checkErr(err)

	// Deal with SIGINT
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)

	var timer time.Time
	go func() {
		for sig := range sigChan {
			// Prevent exiting on accidental signal send
			if time.Now().Sub(timer) < time.Second*signalTimeout {
				terminal.Restore(int(os.Stdin.Fd()), state)
				os.Exit(0)
			}

			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, sig, "signal caught!")
			fmt.Fprintf(os.Stderr, "Send signal again within %v seconds to exit\n", signalTimeout)
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")
			fmt.Fprintln(os.Stderr, "")

			timer = time.Now()
		}
	}()
}

// Create and return a database handle
func (dbi *dbInfo) Connect() (*sql.DB, error) {
	db, err := sql.Open("mysql", dbi.user+":"+dbi.pass+"@tcp("+dbi.host+":"+dbi.port+")/?allowCleartextPasswords=1&tls=skip-verify&charset="+dbi.charset)
	checkErr(err)

	// Ping database to verify credentials
	err = db.Ping()

	return db, err
}

// readRows executes a query and sends each row over a channel to be consumed
func readRows(db *sql.DB, query string, dataChan chan []sql.RawBytes, quitChan chan bool, goChan chan bool, csvHeader bool) {
	rows, err := db.Query(query)
	defer rows.Close()
	if err != nil {
		log.Print(err)
		os.Exit(1)
	}

	cols, err := rows.Columns()
	checkErr(err)

	// Write columns as a header line
	if csvHeader {
		headers := make([]sql.RawBytes, len(cols))
		for i, col := range cols {
			headers[i] = []byte(col)
		}
		dataChan <- headers
		<-goChan
	}

	// Need to scan into empty interface since we don't know how many columns a query might return
	scanVals := make([]interface{}, len(cols))
	vals := make([]sql.RawBytes, len(cols))
	for i := range vals {
		scanVals[i] = &vals[i]
	}

	for rows.Next() {
		err := rows.Scan(scanVals...)
		checkErr(err)

		dataChan <- vals

		// Block and wait for writeRows() to signal back it has consumed the data
		// This is necessary because sql.RawBytes is a memory pointer and when rows.Next()
		// loops and change the memory address before writeRows can properly process the values
		<-goChan
	}

	err = rows.Err()
	checkErr(err)

	close(dataChan)
	quitChan <- true
}

// writeCSV reads from a channel and writes CSV output
func writeCSV(w *Writer, dataChan chan []sql.RawBytes, goChan chan bool, verbose bool) uint {
	var rowsWritten uint
	var verboseCount uint

	if verbose {
		fmt.Println("A '.' will be shown for every 10,000 CSV rows written")
	}

	// Range over row results from readRows()
	for data := range dataChan {
		// Format the data to CSV and write
		size, err := w.Write(data)
		checkErr(err)

		// Visual write indicator when verbose is enabled
		rowsWritten++
		if verbose {
			verboseCount++
			if verboseCount == 10000 {
				fmt.Printf(".")
				verboseCount = 0
			}
		}

		// Flush CSV writer contents once it exceeds flushBufferSize
		if size > flushBufferSize {
			w.Flush()
			err = w.Error()
			checkErr(err)
		}

		// Signal back to readRows() it can loop and scan the next row
		goChan <- true
	}

	// Flush remaining CSV writer contents
	w.Flush()
	err := w.Error()
	checkErr(err)

	return rowsWritten
}
<|endoftext|>