<commit_before>/*
Copyright The Helm Authors.

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

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

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

package storage // import "helm.sh/helm/v3/pkg/storage"

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"

	rspb "helm.sh/helm/v3/pkg/release"
	relutil "helm.sh/helm/v3/pkg/releaseutil"
	"helm.sh/helm/v3/pkg/storage/driver"
)

// The type field of the Kubernetes storage object which stores the Helm release
// version. It is modified slightly replacing the '/': sh.helm/release.v1
// Note: The version 'v1' is incremented if the release object metadata is
// modified between major releases.
// This constant is used as a prefix for the Kubernetes storage object name.
const HelmStorageType = "sh.helm.release.v1"

// Storage represents a storage engine for a Release.
type Storage struct {
	driver.Driver

	// MaxHistory specifies the maximum number of historical releases that will
	// be retained, including the most recent release. Values of 0 or less are
	// ignored (meaning no limits are imposed).
	MaxHistory int

	Log func(string, ...interface{})
}

// Get retrieves the release from storage. An error is returned
// if the storage driver failed to fetch the release, or the
// release identified by the key, version pair does not exist.
func (s *Storage) Get(name string, version int) (*rspb.Release, error) {
	s.Log("getting release %q", makeKey(name, version))
	return s.Driver.Get(makeKey(name, version))
}

// Create creates a new storage entry holding the release. An
// error is returned if the storage driver failed to store the
// release, or a release with identical an key already exists.
func (s *Storage) Create(rls *rspb.Release) error {
	s.Log("creating release %q", makeKey(rls.Name, rls.Version))
	if s.MaxHistory > 0 {
		// Want to make space for one more release.
		s.removeLeastRecent(rls.Name, s.MaxHistory-1)
	}
	return s.Driver.Create(makeKey(rls.Name, rls.Version), rls)
}

// Update update the release in storage. An error is returned if the
// storage backend fails to update the release or if the release
// does not exist.
func (s *Storage) Update(rls *rspb.Release) error {
	s.Log("updating release %q", makeKey(rls.Name, rls.Version))
	return s.Driver.Update(makeKey(rls.Name, rls.Version), rls)
}

// Delete deletes the release from storage. An error is returned if
// the storage backend fails to delete the release or if the release
// does not exist.
func (s *Storage) Delete(name string, version int) (*rspb.Release, error) {
	s.Log("deleting release %q", makeKey(name, version))
	return s.Driver.Delete(makeKey(name, version))
}

// ListReleases returns all releases from storage. An error is returned if the
// storage backend fails to retrieve the releases.
func (s *Storage) ListReleases() ([]*rspb.Release, error) {
	s.Log("listing all releases in storage")
	return s.Driver.List(func(_ *rspb.Release) bool { return true })
}

// ListUninstalled returns all releases with Status == UNINSTALLED. An error is returned
// if the storage backend fails to retrieve the releases.
func (s *Storage) ListUninstalled() ([]*rspb.Release, error) {
	s.Log("listing uninstalled releases in storage")
	return s.Driver.List(func(rls *rspb.Release) bool {
		return relutil.StatusFilter(rspb.StatusUninstalled).Check(rls)
	})
}

// ListDeployed returns all releases with Status == DEPLOYED. An error is returned
// if the storage backend fails to retrieve the releases.
func (s *Storage) ListDeployed() ([]*rspb.Release, error) {
	s.Log("listing all deployed releases in storage")
	return s.Driver.List(func(rls *rspb.Release) bool {
		return relutil.StatusFilter(rspb.StatusDeployed).Check(rls)
	})
}

// Deployed returns the last deployed release with the provided release name, or
// returns ErrReleaseNotFound if not found.
func (s *Storage) Deployed(name string) (*rspb.Release, error) {
	ls, err := s.DeployedAll(name)
	if err != nil {
		if strings.Contains(err.Error(), "not found") {
			return nil, errors.Errorf("%q has no deployed releases", name)
		}
		return nil, err
	}

	if len(ls) == 0 {
		return nil, errors.Errorf("%q has no deployed releases", name)
	}

	return ls[0], err
}

// DeployedAll returns all deployed releases with the provided name, or
// returns ErrReleaseNotFound if not found.
func (s *Storage) DeployedAll(name string) ([]*rspb.Release, error) {
	s.Log("getting deployed releases from %q history", name)

	ls, err := s.Driver.Query(map[string]string{
		"name":   name,
		"owner":  "helm",
		"status": "deployed",
	})
	if err == nil {
		return ls, nil
	}
	if strings.Contains(err.Error(), "not found") {
		return nil, errors.Errorf("%q has no deployed releases", name)
	}
	return nil, err
}

// History returns the revision history for the release with the provided name, or
// returns ErrReleaseNotFound if no such release name exists.
func (s *Storage) History(name string) ([]*rspb.Release, error) {
	s.Log("getting release history for %q", name)

	return s.Driver.Query(map[string]string{"name": name, "owner": "helm"})
}

// removeLeastRecent removes items from history until the lengh number of releases
// does not exceed max.
//
// We allow max to be set explicitly so that calling functions can "make space"
// for the new records they are going to write.
func (s *Storage) removeLeastRecent(name string, max int) error {
	if max < 0 {
		return nil
	}
	h, err := s.History(name)
	if err != nil {
		return err
	}
	if len(h) <= max {
		return nil
	}
	overage := len(h) - max

	// We want oldest to newest
	relutil.SortByRevision(h)

	// Delete as many as possible. In the case of API throughput limitations,
	// multiple invocations of this function will eventually delete them all.
	toDelete := h[0:overage]
	errs := []error{}
	for _, rel := range toDelete {
		key := makeKey(name, rel.Version)
		_, innerErr := s.Delete(name, rel.Version)
		if innerErr != nil {
			s.Log("error pruning %s from release history: %s", key, innerErr)
			errs = append(errs, innerErr)
		}
	}

	s.Log("Pruned %d record(s) from %s with %d error(s)", len(toDelete), name, len(errs))
	switch c := len(errs); c {
	case 0:
		return nil
	case 1:
		return errs[0]
	default:
		return errors.Errorf("encountered %d deletion errors. First is: %s", c, errs[0])
	}
}

// Last fetches the last revision of the named release.
func (s *Storage) Last(name string) (*rspb.Release, error) {
	s.Log("getting last revision of %q", name)
	h, err := s.History(name)
	if err != nil {
		return nil, err
	}
	if len(h) == 0 {
		return nil, errors.Errorf("no revision for release %q", name)
	}

	relutil.Reverse(h, relutil.SortByRevision)
	return h[0], nil
}

// makeKey concatenates the Kubernetes storage object type, a release name and version
// into a string with format:```<helm_storage_type>.<release_name>.v<release_version>```.
// The storage type is prepended to keep name uniqueness between different
// release storage types. An example of clash when not using the type:
// https://github.com/helm/helm/issues/6435.
// This key is used to uniquely identify storage objects.
func makeKey(rlsname string, version int) string {
	return fmt.Sprintf("%s.%s.v%d", HelmStorageType, rlsname, version)
}

// Init initializes a new storage backend with the driver d.
// If d is nil, the default in-memory driver is used.
func Init(d driver.Driver) *Storage {
	// default driver is in memory
	if d == nil {
		d = driver.NewMemory()
	}
	return &Storage{
		Driver: d,
		Log:    func(_ string, _ ...interface{}) {},
	}
}
<commit_msg>Fix a typo "update" -> "updates" (#7346)<commit_after>/*
Copyright The Helm Authors.

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

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

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

package storage // import "helm.sh/helm/v3/pkg/storage"

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"

	rspb "helm.sh/helm/v3/pkg/release"
	relutil "helm.sh/helm/v3/pkg/releaseutil"
	"helm.sh/helm/v3/pkg/storage/driver"
)

// The type field of the Kubernetes storage object which stores the Helm release
// version. It is modified slightly replacing the '/': sh.helm/release.v1
// Note: The version 'v1' is incremented if the release object metadata is
// modified between major releases.
// This constant is used as a prefix for the Kubernetes storage object name.
const HelmStorageType = "sh.helm.release.v1"

// Storage represents a storage engine for a Release.
type Storage struct {
	driver.Driver

	// MaxHistory specifies the maximum number of historical releases that will
	// be retained, including the most recent release. Values of 0 or less are
	// ignored (meaning no limits are imposed).
	MaxHistory int

	Log func(string, ...interface{})
}

// Get retrieves the release from storage. An error is returned
// if the storage driver failed to fetch the release, or the
// release identified by the key, version pair does not exist.
func (s *Storage) Get(name string, version int) (*rspb.Release, error) {
	s.Log("getting release %q", makeKey(name, version))
	return s.Driver.Get(makeKey(name, version))
}

// Create creates a new storage entry holding the release. An
// error is returned if the storage driver failed to store the
// release, or a release with identical an key already exists.
func (s *Storage) Create(rls *rspb.Release) error {
	s.Log("creating release %q", makeKey(rls.Name, rls.Version))
	if s.MaxHistory > 0 {
		// Want to make space for one more release.
		s.removeLeastRecent(rls.Name, s.MaxHistory-1)
	}
	return s.Driver.Create(makeKey(rls.Name, rls.Version), rls)
}

// Update updates the release in storage. An error is returned if the
// storage backend fails to update the release or if the release
// does not exist.
func (s *Storage) Update(rls *rspb.Release) error {
	s.Log("updating release %q", makeKey(rls.Name, rls.Version))
	return s.Driver.Update(makeKey(rls.Name, rls.Version), rls)
}

// Delete deletes the release from storage. An error is returned if
// the storage backend fails to delete the release or if the release
// does not exist.
func (s *Storage) Delete(name string, version int) (*rspb.Release, error) {
	s.Log("deleting release %q", makeKey(name, version))
	return s.Driver.Delete(makeKey(name, version))
}

// ListReleases returns all releases from storage. An error is returned if the
// storage backend fails to retrieve the releases.
func (s *Storage) ListReleases() ([]*rspb.Release, error) {
	s.Log("listing all releases in storage")
	return s.Driver.List(func(_ *rspb.Release) bool { return true })
}

// ListUninstalled returns all releases with Status == UNINSTALLED. An error is returned
// if the storage backend fails to retrieve the releases.
func (s *Storage) ListUninstalled() ([]*rspb.Release, error) {
	s.Log("listing uninstalled releases in storage")
	return s.Driver.List(func(rls *rspb.Release) bool {
		return relutil.StatusFilter(rspb.StatusUninstalled).Check(rls)
	})
}

// ListDeployed returns all releases with Status == DEPLOYED. An error is returned
// if the storage backend fails to retrieve the releases.
func (s *Storage) ListDeployed() ([]*rspb.Release, error) {
	s.Log("listing all deployed releases in storage")
	return s.Driver.List(func(rls *rspb.Release) bool {
		return relutil.StatusFilter(rspb.StatusDeployed).Check(rls)
	})
}

// Deployed returns the last deployed release with the provided release name, or
// returns ErrReleaseNotFound if not found.
func (s *Storage) Deployed(name string) (*rspb.Release, error) {
	ls, err := s.DeployedAll(name)
	if err != nil {
		if strings.Contains(err.Error(), "not found") {
			return nil, errors.Errorf("%q has no deployed releases", name)
		}
		return nil, err
	}

	if len(ls) == 0 {
		return nil, errors.Errorf("%q has no deployed releases", name)
	}

	return ls[0], err
}

// DeployedAll returns all deployed releases with the provided name, or
// returns ErrReleaseNotFound if not found.
func (s *Storage) DeployedAll(name string) ([]*rspb.Release, error) {
	s.Log("getting deployed releases from %q history", name)

	ls, err := s.Driver.Query(map[string]string{
		"name":   name,
		"owner":  "helm",
		"status": "deployed",
	})
	if err == nil {
		return ls, nil
	}
	if strings.Contains(err.Error(), "not found") {
		return nil, errors.Errorf("%q has no deployed releases", name)
	}
	return nil, err
}

// History returns the revision history for the release with the provided name, or
// returns ErrReleaseNotFound if no such release name exists.
func (s *Storage) History(name string) ([]*rspb.Release, error) {
	s.Log("getting release history for %q", name)

	return s.Driver.Query(map[string]string{"name": name, "owner": "helm"})
}

// removeLeastRecent removes items from history until the lengh number of releases
// does not exceed max.
//
// We allow max to be set explicitly so that calling functions can "make space"
// for the new records they are going to write.
func (s *Storage) removeLeastRecent(name string, max int) error {
	if max < 0 {
		return nil
	}
	h, err := s.History(name)
	if err != nil {
		return err
	}
	if len(h) <= max {
		return nil
	}
	overage := len(h) - max

	// We want oldest to newest
	relutil.SortByRevision(h)

	// Delete as many as possible. In the case of API throughput limitations,
	// multiple invocations of this function will eventually delete them all.
	toDelete := h[0:overage]
	errs := []error{}
	for _, rel := range toDelete {
		key := makeKey(name, rel.Version)
		_, innerErr := s.Delete(name, rel.Version)
		if innerErr != nil {
			s.Log("error pruning %s from release history: %s", key, innerErr)
			errs = append(errs, innerErr)
		}
	}

	s.Log("Pruned %d record(s) from %s with %d error(s)", len(toDelete), name, len(errs))
	switch c := len(errs); c {
	case 0:
		return nil
	case 1:
		return errs[0]
	default:
		return errors.Errorf("encountered %d deletion errors. First is: %s", c, errs[0])
	}
}

// Last fetches the last revision of the named release.
func (s *Storage) Last(name string) (*rspb.Release, error) {
	s.Log("getting last revision of %q", name)
	h, err := s.History(name)
	if err != nil {
		return nil, err
	}
	if len(h) == 0 {
		return nil, errors.Errorf("no revision for release %q", name)
	}

	relutil.Reverse(h, relutil.SortByRevision)
	return h[0], nil
}

// makeKey concatenates the Kubernetes storage object type, a release name and version
// into a string with format:```<helm_storage_type>.<release_name>.v<release_version>```.
// The storage type is prepended to keep name uniqueness between different
// release storage types. An example of clash when not using the type:
// https://github.com/helm/helm/issues/6435.
// This key is used to uniquely identify storage objects.
func makeKey(rlsname string, version int) string {
	return fmt.Sprintf("%s.%s.v%d", HelmStorageType, rlsname, version)
}

// Init initializes a new storage backend with the driver d.
// If d is nil, the default in-memory driver is used.
func Init(d driver.Driver) *Storage {
	// default driver is in memory
	if d == nil {
		d = driver.NewMemory()
	}
	return &Storage{
		Driver: d,
		Log:    func(_ string, _ ...interface{}) {},
	}
}
<|endoftext|>


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

package systemd

import "net"

// DialJournal returns no error if the process can dial journal socket.
// Returns an error if dial failed, whichi indicates journald is not available
// (e.g. run embedded etcd as docker daemon).
// Reference: https://github.com/coreos/go-systemd/blob/master/journal/journal.go.
func DialJournal() error {
	conn, err := net.Dial("unixgram", "/run/systemd/journal/socket")
	if conn != nil {
		defer conn.Close()
	}
	return err
}
<commit_msg>Fix typo in DialJournal<commit_after>// Copyright 2018 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package systemd

import "net"

// DialJournal returns no error if the process can dial journal socket.
// Returns an error if dial failed, which indicates journald is not available
// (e.g. run embedded etcd as docker daemon).
// Reference: https://github.com/coreos/go-systemd/blob/master/journal/journal.go.
func DialJournal() error {
	conn, err := net.Dial("unixgram", "/run/systemd/journal/socket")
	if conn != nil {
		defer conn.Close()
	}
	return err
}
<|endoftext|>


<commit_before>package timelog

import (
	"context"
	"net/http"
	"time"

	"github.com/LassoProject/lasso/pkg/response"

	log "github.com/Sirupsen/logrus"
	// "github.com/mattn/go-isatty"
)

var (
	green   = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white   = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow  = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red     = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue    = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan    = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset   = string([]byte{27, 91, 48, 109})
)

// TimeLog records how long it takes to process the http request and produce the response (latency)
func TimeLog(nextHandler http.Handler) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Debugf("Request received : %v", r)
		start := time.Now()

		// make the call
		v := response.CaptureWriter{w, 0}
		ctx := context.Background()
		nextHandler.ServeHTTP(&v, r.WithContext(ctx))

		// Stop timer
		end := time.Now()

		latency := end.Sub(start)

		log.Debugf("Request handled successfully: %v", v.GetStatusCode())
		var statusCode = v.GetStatusCode()
		statusColor := colorForStatus(statusCode)

		path := r.URL.Path
		host := r.Host
		referer := r.Header.Get("Referer")

		clientIP := r.RemoteAddr
		method := r.Method

		log.Infof("|%s %3d %s| %13v | %s | %s %s %s | %s",
			statusColor, statusCode, reset,
			latency,
			clientIP,
			method, host, path,
			referer)
	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code < 300:
		return green
	case code >= 300 && code < 400:
		return white
	case code >= 400 && code < 500:
		return yellow
	default:
		return red
	}
}
<commit_msg>record request number and average latency<commit_after>package timelog

import (
	"context"
	"net/http"
	"time"

	"github.com/LassoProject/lasso/pkg/response"

	log "github.com/Sirupsen/logrus"
	// "github.com/mattn/go-isatty"
)

var (
	green      = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white      = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow     = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red        = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue       = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta    = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan       = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset      = string([]byte{27, 91, 48, 109})
	req        = int64(0)
	avgLatency = int64(0)
)

// TimeLog records how long it takes to process the http request and produce the response (latency)
func TimeLog(nextHandler http.Handler) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Debugf("Request received : %v", r)
		start := time.Now()

		// make the call
		v := response.CaptureWriter{w, 0}
		ctx := context.Background()
		nextHandler.ServeHTTP(&v, r.WithContext(ctx))

		// Stop timer
		end := time.Now()
		latency := end.Sub(start)
		req++
		avgLatency = avgLatency + ((int64(latency) - avgLatency) / req)
		log.Debugf("Request handled successfully: %v", v.GetStatusCode())
		var statusCode = v.GetStatusCode()
		statusColor := colorForStatus(statusCode)

		path := r.URL.Path
		host := r.Host
		referer := r.Header.Get("Referer")

		clientIP := r.RemoteAddr
		method := r.Method

		log.Infof("|%s %3d %s| %d %10v %10v | %s | %s %s %s | %s",
			statusColor, statusCode, reset,
			req, latency, time.Duration(avgLatency),
			clientIP,
			method, host, path,
			referer)
	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code < 300:
		return green
	case code >= 300 && code < 400:
		return white
	case code >= 400 && code < 500:
		return yellow
	default:
		return red
	}
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017, 2018 Red Hat, Inc.
 *
 */

package vnc

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"k8s.io/client-go/tools/clientcmd"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const LISTEN_TIMEOUT = 60 * time.Second

const FLAG = "vnc"

//#### Tiger VNC ####
//# https://github.com/TigerVNC/tigervnc/releases
const TIGER_VNC = "/Applications/TigerVNC Viewer 1.8.0.app/Contents/MacOS/TigerVNC Viewer"

//#### Chicken VNC ####
//# https://sourceforge.net/projects/chicken/
const CHICKEN_VNC = "/Applications/Chicken.app/Contents/MacOS/Chicken"

//####  Real VNC ####
//# https://www.realvnc.com/en/connect/download/viewer/macos/
const REAL_VNC = "/Applications/VNC Viewer.app/Contents/MacOS/vncviewer"

func NewCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "vnc (vmi)",
		Short:   "Open a vnc connection to a virtual machine.",
		Example: usage(),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := VNC{clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

type VNC struct {
	clientConfig clientcmd.ClientConfig
}

func (o *VNC) Run(cmd *cobra.Command, args []string) error {
	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	vmi := args[0]

	virtCli, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return err
	}

	// setup connection with VM
	vnc, err := virtCli.VirtualMachineInstance(namespace).VNC(vmi)
	if err != nil {
		return fmt.Errorf("Can't access VMI %s: %s", vmi, err.Error())
	}

	lnAddr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("Can't resolve the address: %s", err.Error())
	}

	// The local tcp server is used to proxy the podExec websock connection to remote-viewer
	ln, err := net.ListenTCP("tcp", lnAddr)
	if err != nil {
		return fmt.Errorf("Can't listen on unix socket: %s", err.Error())
	}
	// End of pre-flight checks. Everything looks good, we can start
	// the goroutines and let the data flow

	//                                       -> pipeInWriter  -> pipeInReader
	// remote-viewer -> unix sock connection
	//                                       <- pipeOutReader <- pipeOutWriter
	pipeInReader, pipeInWriter := io.Pipe()
	pipeOutReader, pipeOutWriter := io.Pipe()

	k8ResChan := make(chan error)
	listenResChan := make(chan error)
	viewResChan := make(chan error)
	stopChan := make(chan struct{}, 1)
	doneChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)

	go func() {
		// transfer data from/to the VM
		k8ResChan <- vnc.Stream(kubecli.StreamOptions{
			In:  pipeInReader,
			Out: pipeOutWriter,
		})
	}()

	// wait for remote-viewer to connect to our local proxy server
	go func() {
		start := time.Now()
		glog.Infof("connection timeout: %v", LISTEN_TIMEOUT)
		// exit early if spawning remote-viewer fails
		ln.SetDeadline(time.Now().Add(LISTEN_TIMEOUT))

		fd, err := ln.Accept()
		if err != nil {
			glog.V(2).Infof("Failed to accept unix sock connection. %s", err.Error())
			listenResChan <- err
		}
		defer fd.Close()

		glog.V(2).Infof("remote-viewer connected in %v", time.Now().Sub(start))

		// write to FD <- pipeOutReader
		go func() {
			_, err := io.Copy(fd, pipeOutReader)
			readStop <- err
		}()

		// read from FD -> pipeInWriter
		go func() {
			_, err := io.Copy(pipeInWriter, fd)
			writeStop <- err
		}()

		// don't terminate until remote-viewer is done
		<-doneChan
		listenResChan <- err
	}()

	// execute VNC
	go func() {
		defer close(doneChan)
		port := ln.Addr().(*net.TCPAddr).Port
		args := []string{}

		vncBin := ""
		osType := runtime.GOOS
		switch osType {
		case "darwin":
			if _, err := os.Stat(TIGER_VNC); err == nil {
				vncBin = TIGER_VNC
				args = tigerVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			} else if _, err := os.Stat(CHICKEN_VNC); err == nil {
				vncBin = CHICKEN_VNC
				args = chickenVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			} else if _, err := os.Stat(REAL_VNC); err == nil {
				vncBin = REAL_VNC
				args = realVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			}
		case "linux":
			_, err := exec.LookPath("remote-viewer")
			if exec.ErrNotFound == err {
				viewResChan <- fmt.Errorf("could not find the remote-viewer binary in $PATH")
				return
			} else if err != nil {
				viewResChan <- err
				return
			}
			vncBin = "remote-viewer"
			args = remoteViewerArgs(port)
		default:
			viewResChan <- fmt.Errorf("virtctl does not support VNC on %v", osType)
			return
		}

		if vncBin == "" {
			glog.Errorf("No supported VNC app found in %s", osType)
			err = fmt.Errorf("No supported VNC app found in %s", osType)
		} else {
			if glog.V(4) {
				glog.Infof("Executing commandline: '%s %v'", vncBin, args)
			}
			cmnd := exec.Command(vncBin, args...)
			output, err := cmnd.CombinedOutput()
			if err != nil {
				glog.Errorf("%s execution failed: %v, output: %v", vncBin, err, string(output))
			} else {
				glog.V(2).Infof("remote-viewer output: %v", string(output))
			}
		}
		viewResChan <- err
	}()

	go func() {
		defer close(stopChan)
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-k8ResChan:
	case err = <-viewResChan:
	case err = <-listenResChan:
	}

	if err != nil {
		return fmt.Errorf("Error encountered: %s", err.Error())
	}
	return nil
}

func tigerVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	if glog.V(4) {
		args = append(args, "Log=*:stderr:100")
	}
	return
}

func chickenVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	return
}

func realVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	args = append(args, "-WarnUnencrypted=0")
	args = append(args, "-Shared=0")
	args = append(args, "-ShareFiles=0")
	if glog.V(4) {
		args = append(args, "-log=*:stderr:100")
	}
	return
}

func remoteViewerArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("vnc://127.0.0.1:%d", port))
	if glog.V(4) {
		args = append(args, "--debug")
	}
	return
}

func usage() string {
	usage := "# Connect to testvmi via remote-viewer:\n"
	usage += "./virtctl vnc testvmi"
	return usage
}
<commit_msg>Add support for arbitrary 'remote-viewer' in path, bundle vnc constants<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017, 2018 Red Hat, Inc.
 *
 */

package vnc

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"time"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"k8s.io/client-go/tools/clientcmd"

	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/virtctl/templates"
)

const (
	LISTEN_TIMEOUT = 60 * time.Second
	FLAG           = "vnc"

	//#### Tiger VNC ####
	//# https://github.com/TigerVNC/tigervnc/releases
	TIGER_VNC = "/Applications/TigerVNC Viewer 1.8.0.app/Contents/MacOS/TigerVNC Viewer"

	//#### Chicken VNC ####
	//# https://sourceforge.net/projects/chicken/
	CHICKEN_VNC = "/Applications/Chicken.app/Contents/MacOS/Chicken"

	//####  Real VNC ####
	//# https://www.realvnc.com/en/connect/download/viewer/macos/
	REAL_VNC = "/Applications/VNC Viewer.app/Contents/MacOS/vncviewer"

	REMOTE_VIEWER = "remote-viewer"
)

func NewCommand(clientConfig clientcmd.ClientConfig) *cobra.Command {
	cmd := &cobra.Command{
		Use:     "vnc (vmi)",
		Short:   "Open a vnc connection to a virtual machine.",
		Example: usage(),
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			c := VNC{clientConfig: clientConfig}
			return c.Run(cmd, args)
		},
	}
	cmd.SetUsageTemplate(templates.UsageTemplate())
	return cmd
}

type VNC struct {
	clientConfig clientcmd.ClientConfig
}

func (o *VNC) Run(cmd *cobra.Command, args []string) error {
	namespace, _, err := o.clientConfig.Namespace()
	if err != nil {
		return err
	}

	vmi := args[0]

	virtCli, err := kubecli.GetKubevirtClientFromClientConfig(o.clientConfig)
	if err != nil {
		return err
	}

	// setup connection with VM
	vnc, err := virtCli.VirtualMachineInstance(namespace).VNC(vmi)
	if err != nil {
		return fmt.Errorf("Can't access VMI %s: %s", vmi, err.Error())
	}

	lnAddr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("Can't resolve the address: %s", err.Error())
	}

	// The local tcp server is used to proxy the podExec websock connection to remote-viewer
	ln, err := net.ListenTCP("tcp", lnAddr)
	if err != nil {
		return fmt.Errorf("Can't listen on unix socket: %s", err.Error())
	}
	// End of pre-flight checks. Everything looks good, we can start
	// the goroutines and let the data flow

	//                                       -> pipeInWriter  -> pipeInReader
	// remote-viewer -> unix sock connection
	//                                       <- pipeOutReader <- pipeOutWriter
	pipeInReader, pipeInWriter := io.Pipe()
	pipeOutReader, pipeOutWriter := io.Pipe()

	k8ResChan := make(chan error)
	listenResChan := make(chan error)
	viewResChan := make(chan error)
	stopChan := make(chan struct{}, 1)
	doneChan := make(chan struct{}, 1)
	writeStop := make(chan error)
	readStop := make(chan error)

	go func() {
		// transfer data from/to the VM
		k8ResChan <- vnc.Stream(kubecli.StreamOptions{
			In:  pipeInReader,
			Out: pipeOutWriter,
		})
	}()

	// wait for remote-viewer to connect to our local proxy server
	go func() {
		start := time.Now()
		glog.Infof("connection timeout: %v", LISTEN_TIMEOUT)
		// exit early if spawning remote-viewer fails
		ln.SetDeadline(time.Now().Add(LISTEN_TIMEOUT))

		fd, err := ln.Accept()
		if err != nil {
			glog.V(2).Infof("Failed to accept unix sock connection. %s", err.Error())
			listenResChan <- err
		}
		defer fd.Close()

		glog.V(2).Infof("remote-viewer connected in %v", time.Now().Sub(start))

		// write to FD <- pipeOutReader
		go func() {
			_, err := io.Copy(fd, pipeOutReader)
			readStop <- err
		}()

		// read from FD -> pipeInWriter
		go func() {
			_, err := io.Copy(pipeInWriter, fd)
			writeStop <- err
		}()

		// don't terminate until remote-viewer is done
		<-doneChan
		listenResChan <- err
	}()

	// execute VNC
	go func() {
		defer close(doneChan)
		port := ln.Addr().(*net.TCPAddr).Port
		args := []string{}

		vncBin := ""
		osType := runtime.GOOS
		switch osType {
		case "darwin":
			if _, err := os.Stat(TIGER_VNC); err == nil {
				vncBin = TIGER_VNC
				args = tigerVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			} else if _, err := os.Stat(CHICKEN_VNC); err == nil {
				vncBin = CHICKEN_VNC
				args = chickenVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			} else if _, err := os.Stat(REAL_VNC); err == nil {
				vncBin = REAL_VNC
				args = realVncArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			} else if _, err := exec.LookPath(REMOTE_VIEWER); err == nil {
				// fall back to user supplied script/binary in path
				vncBin = REMOTE_VIEWER
				args = remoteViewerArgs(port)
			} else if !os.IsNotExist(err) {
				viewResChan <- err
				return
			}
		case "linux":
			_, err := exec.LookPath(REMOTE_VIEWER)
			if exec.ErrNotFound == err {
				viewResChan <- fmt.Errorf("could not find the remote-viewer binary in $PATH")
				return
			} else if err != nil {
				viewResChan <- err
				return
			}
			vncBin = REMOTE_VIEWER
			args = remoteViewerArgs(port)
		default:
			viewResChan <- fmt.Errorf("virtctl does not support VNC on %v", osType)
			return
		}

		if vncBin == "" {
			glog.Errorf("No supported VNC app found in %s", osType)
			err = fmt.Errorf("No supported VNC app found in %s", osType)
		} else {
			if glog.V(4) {
				glog.Infof("Executing commandline: '%s %v'", vncBin, args)
			}
			cmnd := exec.Command(vncBin, args...)
			output, err := cmnd.CombinedOutput()
			if err != nil {
				glog.Errorf("%s execution failed: %v, output: %v", vncBin, err, string(output))
			} else {
				glog.V(2).Infof("remote-viewer output: %v", string(output))
			}
		}
		viewResChan <- err
	}()

	go func() {
		defer close(stopChan)
		interrupt := make(chan os.Signal, 1)
		signal.Notify(interrupt, os.Interrupt)
		<-interrupt
	}()

	select {
	case <-stopChan:
	case err = <-readStop:
	case err = <-writeStop:
	case err = <-k8ResChan:
	case err = <-viewResChan:
	case err = <-listenResChan:
	}

	if err != nil {
		return fmt.Errorf("Error encountered: %s", err.Error())
	}
	return nil
}

func tigerVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	if glog.V(4) {
		args = append(args, "Log=*:stderr:100")
	}
	return
}

func chickenVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	return
}

func realVncArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("127.0.0.1:%d", port))
	args = append(args, "-WarnUnencrypted=0")
	args = append(args, "-Shared=0")
	args = append(args, "-ShareFiles=0")
	if glog.V(4) {
		args = append(args, "-log=*:stderr:100")
	}
	return
}

func remoteViewerArgs(port int) (args []string) {
	args = append(args, fmt.Sprintf("vnc://127.0.0.1:%d", port))
	if glog.V(4) {
		args = append(args, "--debug")
	}
	return
}

func usage() string {
	usage := "# Connect to testvmi via remote-viewer:\n"
	usage += "./virtctl vnc testvmi"
	return usage
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc.
//
// 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 yang

import (
	"bytes"
	"testing"
)

func (s1 *Statement) equal(s2 *Statement) bool {
	if s1.Keyword != s2.Keyword ||
		s1.HasArgument != s2.HasArgument ||
		s1.Argument != s2.Argument ||
		len(s1.statements) != len(s2.statements) {
		return false
	}

	for x, ss := range s1.statements {
		if !ss.equal(s2.statements[x]) {
			return false
		}
	}
	return true
}

// SA returns a statement with an argument and optional substatements.
func SA(k, a string, ss ...*Statement) *Statement {
	return &Statement{
		Keyword:     k,
		Argument:    a,
		HasArgument: true,
		statements:  ss,
	}
}

// S returns a statement with no argument and optional substatements.
func S(k string, ss ...*Statement) *Statement {
	return &Statement{
		Keyword:    k,
		statements: ss,
	}
}

func TestParse(t *testing.T) {
	for _, tt := range []struct {
		line int
		in   string
		out  []*Statement
		err  string
	}{
		{line: line()},
		{line: line(), in: `
foo;
`,
			out: []*Statement{
				S("foo"),
			},
		},
		{line: line(), in: `
foo {}
`,
			out: []*Statement{
				S("foo"),
			},
		},
		{line: line(), in: `
foo "";
`,
			out: []*Statement{
				SA("foo", ""),
			},
		},
		{line: line(), in: `
foo bar;
`,
			out: []*Statement{
				SA("foo", "bar"),
			},
		},
		{line: line(), in: `
foo "bar";
`,
			out: []*Statement{
				SA("foo", "bar"),
			},
		},
		{line: line(), in: `
foo "\\ \S \n";
`,
			err: `test.yang:2:5: invalid escape sequence: \S`,
		},
		{line: line(), in: `
pattern "\\ \S \n";
`,
			out: []*Statement{
				SA("pattern", `\ \S
`),
			},
		},
		{line: line(), in: `
foo '\\ \S \n';
`,
			out: []*Statement{
				SA("foo", `\\ \S \n`),
			},
		},
		{line: line(), in: `
pattern '\\ \S \n';
`,
			out: []*Statement{
				SA("pattern", `\\ \S \n`),
			},
		},
		{line: line(), in: `
foo bar;
red black;
`,
			out: []*Statement{
				SA("foo", "bar"),
				SA("red", "black"),
			},
		},
		{line: line(), in: `
foo {
   key value;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
				),
			},
		},
		{line: line(), in: `
foo {
   key value;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
				),
			},
		},
		{line: line(), in: `
foo {
   key value;
   key2;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
					S("key2"),
				),
			},
		},
		{line: line(), in: `
foo1 {
   key value1;
}
foo2 {
   key value2;
}
foo3 value3;
`,
			out: []*Statement{
				S("foo1",
					SA("key", "value1"),
				),
				S("foo2",
					SA("key", "value2"),
				),
				SA("foo3", "value3"),
			},
		},
		{line: line(), in: `
foo1 {
    key value1;
    foo2 {
        key value2;
    }
}
`,
			out: []*Statement{
				S("foo1",
					SA("key", "value1"),
					S("foo2",
						SA("key", "value2"),
					),
				),
			},
		},
		{line: line(), in: `
 }
`,
			err: `test.yang:2:2: unexpected }`,
		},
		{line: line(), in: `
id
`,
			err: `test.yang: unexpected EOF`,
		},
		{line: line(), in: `
   {
`,
			err: `test.yang:2:4: {: not an identifier`,
		},
		{line: line(), in: `
;
`,
			err: `test.yang:2:1: ;: not an identifier`,
		},
		{line: line(), in: `
statement one two { }
`,
			err: `test.yang:2:15: two: syntax error
test.yang:2:19: {: not an identifier
test.yang:2:21: unexpected }`,
		},
		{line: line(), in: `
    }
foo {
	key: "value";
}
`,
			err: `test.yang:2:5: unexpected }`,
		},
		{line: line(), in: `
{
	something: "bad";
}
foo {
	key: "\Value";
	key2: "value2";
	bar {
		key3: "value\3;
	}
}`,
			err: `test.yang:2:1: {: not an identifier
test.yang:4:1: unexpected }
test.yang:6:7: invalid escape sequence: \V
test.yang:9:9: invalid escape sequence: \3
test.yang:9:9: missing closing "
test.yang: unexpected EOF`,
		},
	} {
		s, err := Parse(tt.in, "test.yang")
		if (s == nil) != (tt.out == nil) {
			if s == nil {
				t.Errorf("%d: did not get expected statements: %s", tt.line, tt.out)
			} else {
				t.Errorf("%d: get unexpected statements: %s", tt.line, s)
			}
		}
		switch {
		case err == nil && tt.err == "":
		case tt.err == "":
			t.Errorf("%d: unexpected error %v", tt.line, err)
			continue
		case err == nil:
			t.Errorf("%d: did not get expected error %v", tt.line, tt.err)
			continue
		case err.Error() == tt.err:
			continue
		default:
			t.Errorf("%d: got error:\n%s\nwant:\n%s", tt.line, err, tt.err)
			continue
		}
		s1 := &Statement{statements: s}
		s2 := &Statement{statements: tt.out}
		if !s1.equal(s2) {
			t.Errorf("%d: got:\n%s\nwant:\n%s", tt.line, s1, s2)
		}
	}
}

func TestWrite(t *testing.T) {
Testing:
	for _, tt := range []struct {
		line int
		in   string
		out  string
	}{
		{line: line(),
			in: `key arg { substatement; }`,
			out: `key "arg" {
	substatement;
}
`,
		},
		{line: line(),
			in: `key { substatement { key arg; }}`,
			out: `key {
	substatement {
		key "arg";
	}
}
`,
		},
		{line: line(),
			in: `
module base {
   namespace "urn:mod";
   prefix "base";

   typedef base-type { type int32; }

   grouping base-group {
     description
       "The base-group is used to test the
        'uses' statement below.  This description
        is here to simply include a multi-line
        string as an example of multi-line strings";
     leaf base-group-leaf {
       config false;
       type string;
     }
   }
   uses base-group;
}
`, out: `module "base" {
	namespace "urn:mod";
	prefix "base";
	typedef "base-type" {
		type "int32";
	}
	grouping "base-group" {
		description "The base-group is used to test the
		             'uses' statement below.  This description
		             is here to simply include a multi-line
		             string as an example of multi-line strings";
		leaf "base-group-leaf" {
			config "false";
			type "string";
		}
	}
	uses "base-group";
}
`,
		},
	} {
		in := tt.in
		// Run twice.  The first time we are parsing tt.in, the second
		// time we are parsing the output from the first parsing.
		for i := 0; i < 2; i++ {
			s, err := Parse(in, "test.yang")
			if err != nil {
				t.Errorf("%d: unexpected error %v", tt.line, err)
				continue Testing
			}
			if len(s) != 1 {
				t.Errorf("%d: got %d statements, expected 1", tt.line, len(s))
				continue Testing
			}
			var buf bytes.Buffer
			s[0].Write(&buf, "")
			out := buf.String()
			if out != tt.out {
				t.Errorf("%d: got:\n%swant:\n%s", tt.line, out, tt.out)
				continue Testing
			}
			in = out
		}
	}
}
<commit_msg>Add missing trailing whitespace.<commit_after>// Copyright 2015 Google Inc.
//
// 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 yang

import (
	"bytes"
	"testing"
)

func (s1 *Statement) equal(s2 *Statement) bool {
	if s1.Keyword != s2.Keyword ||
		s1.HasArgument != s2.HasArgument ||
		s1.Argument != s2.Argument ||
		len(s1.statements) != len(s2.statements) {
		return false
	}

	for x, ss := range s1.statements {
		if !ss.equal(s2.statements[x]) {
			return false
		}
	}
	return true
}

// SA returns a statement with an argument and optional substatements.
func SA(k, a string, ss ...*Statement) *Statement {
	return &Statement{
		Keyword:     k,
		Argument:    a,
		HasArgument: true,
		statements:  ss,
	}
}

// S returns a statement with no argument and optional substatements.
func S(k string, ss ...*Statement) *Statement {
	return &Statement{
		Keyword:    k,
		statements: ss,
	}
}

func TestParse(t *testing.T) {
	for _, tt := range []struct {
		line int
		in   string
		out  []*Statement
		err  string
	}{
		{line: line()},
		{line: line(), in: `
foo;
`,
			out: []*Statement{
				S("foo"),
			},
		},
		{line: line(), in: `
foo {}
`,
			out: []*Statement{
				S("foo"),
			},
		},
		{line: line(), in: `
foo "";
`,
			out: []*Statement{
				SA("foo", ""),
			},
		},
		{line: line(), in: `
foo bar;
`,
			out: []*Statement{
				SA("foo", "bar"),
			},
		},
		{line: line(), in: `
foo "bar";
`,
			out: []*Statement{
				SA("foo", "bar"),
			},
		},
		{line: line(), in: `
foo "\\ \S \n";
`,
			err: `test.yang:2:5: invalid escape sequence: \S`,
		},
		{line: line(), in: `
pattern "\\ \S \n";
`,
			out: []*Statement{
				SA("pattern", `\ \S
`),
			},
		},
		{line: line(), in: `
foo '\\ \S \n';
`,
			out: []*Statement{
				SA("foo", `\\ \S \n`),
			},
		},
		{line: line(), in: `
pattern '\\ \S \n';
`,
			out: []*Statement{
				SA("pattern", `\\ \S \n`),
			},
		},
		{line: line(), in: `
foo bar;
red black;
`,
			out: []*Statement{
				SA("foo", "bar"),
				SA("red", "black"),
			},
		},
		{line: line(), in: `
foo {
   key value;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
				),
			},
		},
		{line: line(), in: `
foo {
   key value;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
				),
			},
		},
		{line: line(), in: `
foo {
   key value;
   key2;
}
`,
			out: []*Statement{
				S("foo",
					SA("key", "value"),
					S("key2"),
				),
			},
		},
		{line: line(), in: `
foo1 {
   key value1;
}
foo2 {
   key value2;
}
foo3 value3;
`,
			out: []*Statement{
				S("foo1",
					SA("key", "value1"),
				),
				S("foo2",
					SA("key", "value2"),
				),
				SA("foo3", "value3"),
			},
		},
		{line: line(), in: `
foo1 {
    key value1;
    foo2 {
        key value2;
    }
}
`,
			out: []*Statement{
				S("foo1",
					SA("key", "value1"),
					S("foo2",
						SA("key", "value2"),
					),
				),
			},
		},
		{line: line(), in: `
 }
`,
			err: `test.yang:2:2: unexpected }`,
		},
		{line: line(), in: `
id
`,
			err: `test.yang: unexpected EOF`,
		},
		{line: line(), in: `
   {
`,
			err: `test.yang:2:4: {: not an identifier`,
		},
		{line: line(), in: `
;
`,
			err: `test.yang:2:1: ;: not an identifier`,
		},
		{line: line(), in: `
statement one two { }
`,
			err: `test.yang:2:15: two: syntax error
test.yang:2:19: {: not an identifier
test.yang:2:21: unexpected }`,
		},
		{line: line(), in: `
    }
foo {
	key: "value";
}
`,
			err: `test.yang:2:5: unexpected }`,
		},
		{line: line(), in: `
{
	something: "bad";
}
foo {
	key: "\Value";
	key2: "value2";
	bar {
		key3: "value\3;
	}
}`,
			err: `test.yang:2:1: {: not an identifier
test.yang:4:1: unexpected }
test.yang:6:7: invalid escape sequence: \V
test.yang:9:9: invalid escape sequence: \3
test.yang:9:9: missing closing "
test.yang: unexpected EOF`,
		},
	} {
		s, err := Parse(tt.in, "test.yang")
		if (s == nil) != (tt.out == nil) {
			if s == nil {
				t.Errorf("%d: did not get expected statements: %s", tt.line, tt.out)
			} else {
				t.Errorf("%d: get unexpected statements: %s", tt.line, s)
			}
		}
		switch {
		case err == nil && tt.err == "":
		case tt.err == "":
			t.Errorf("%d: unexpected error %v", tt.line, err)
			continue
		case err == nil:
			t.Errorf("%d: did not get expected error %v", tt.line, tt.err)
			continue
		case err.Error() == tt.err:
			continue
		default:
			t.Errorf("%d: got error:\n%s\nwant:\n%s", tt.line, err, tt.err)
			continue
		}
		s1 := &Statement{statements: s}
		s2 := &Statement{statements: tt.out}
		if !s1.equal(s2) {
			t.Errorf("%d: got:\n%s\nwant:\n%s", tt.line, s1, s2)
		}
	}
}

func TestWrite(t *testing.T) {
Testing:
	for _, tt := range []struct {
		line int
		in   string
		out  string
	}{
		{line: line(),
			in: `key arg { substatement; }`,
			out: `key "arg" {
	substatement;
}
`,
		},
		{line: line(),
			in: `key { substatement { key arg; }}`,
			out: `key {
	substatement {
		key "arg";
	}
}
`,
		},
		{line: line(),
			in: `
module base {
   namespace "urn:mod";
   prefix "base";

   typedef base-type { type int32; }

   grouping base-group {
     description
       "The base-group is used to test the
        'uses' statement below.  This description
        is here to simply include a multi-line
        string as an example of multi-line strings";
     leaf base-group-leaf {
       config false;
       type string;
     }
   }
   uses base-group;
}
`, out: `module "base" {
	namespace "urn:mod";
	prefix "base";
	typedef "base-type" {
		type "int32";
	}
	grouping "base-group" {
		description "The base-group is used to test the
		             'uses' statement below.  This description
		             is here to simply include a multi-line
		             string as an example of multi-line strings";
		leaf "base-group-leaf" {
			config "false";
			type "string";
		}
	}
	uses "base-group";
}
`,
		},
	} {
		in := tt.in
		// Run twice.  The first time we are parsing tt.in, the second
		// time we are parsing the output from the first parsing.
		for i := 0; i < 2; i++ {
			s, err := Parse(in, "test.yang")
			if err != nil {
				t.Errorf("%d: unexpected error %v", tt.line, err)
				continue Testing
			}
			if len(s) != 1 {
				t.Errorf("%d: got %d statements, expected 1", tt.line, len(s))
				continue Testing
			}
			var buf bytes.Buffer
			s[0].Write(&buf, "")
			out := buf.String()
			if out != tt.out {
				t.Errorf("%d: got:\n%swant:\n%s", tt.line, out, tt.out)
				continue Testing
			}
			in = out
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const HOST = "http://localhost:8765/"

func performCall(method string, url string) (string, error) {
	var res *http.Response
	var err error
	if method == "GET" {
		res, err = http.Get(url)
	} else if method == "POST" {
		res, err = http.Post(url, "text/plain", nil) //todo:
	} else if method == "PUT" {
		client := &http.Client{}
		request, err1 := http.NewRequest("PUT", url, nil)
		if err1 != nil {
			return "", err1
		}
		res, err = client.Do(request)
	}

	if err != nil {
		return "", err
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return "", err
	}
	found := string(body[:])
	return found, nil
}

func checkResult(method, url, expected string, t *testing.T) {
	found, err := performCall(method, url)
	if err != nil {
		t.Fatalf("Unexpected error found - %s", err.Error())
	}

	if found != expected {
		t.Errorf("Expected\n---\n%s\n---\nbut found\n---\n%s\n---\n", expected, found)
	}
}

func TestGetAlive(t *testing.T) {
	expected := "I'm alive"
	checkResult("GET", HOST, expected, t)
}

func TestPlay(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds/beep9.mp3")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayDir(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\",\"test_sounds/beep28.mp3\",\"test_sounds/beep36.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayPlaylist(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("sample_playlist")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\",\"test_sounds/beep28.mp3\",\"test_sounds/beep36.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayNonExistingFile(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds/beep1.mp3")
	expected := "{\"Code\":4,\"Message\":\"File cannot be found\"}"
	checkResult("PUT", url, expected, t)
}

func TestPlayInvalidFileFormat(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_broken/abc.txt")
	expected := "{\"Code\":7,\"Message\":\"Format is not supported\"}"
	checkResult("PUT", url, expected, t)
}

func TestPlayBrokenFile(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_broken/abc.txt")
	expected := "{\"Code\":2,\"Message\":\"SoX failed to open input file\"}"
	checkResult("PUT", url, expected, t)
}

func TestPause(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)
	url := HOST + "pause"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Song is paused\",\"Data\":[\"test_sounds/beep28.mp3\"]}"

	checkResult("POST", url, expected, t)
}

func TestPauseNoPlayback(t *testing.T) {
	url := HOST + "pause"
	expected := "{\"Code\":2,\"Message\":\"Cannot pause. No song is playing\"}"
	checkResult("POST", url, expected, t)
}

func TestResume(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)
	pause_url := HOST + "pause"
	performCall("PUT", pause_url)

	url := HOST + "resume"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Song is resumed\",\"Data\":[\"test_sounds/beep28.mp3\"]}"
	checkResult("POST", url, expected, t)

}

func TestResumeNoPlayback(t *testing.T) {
	url := HOST + "resume"
	expected := "{\"Code\":9,\"Message\":\"Cannot resume. No song was paused\"}"
	checkResult("POST", url, expected, t)
}

func TestResumeNoPaused(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)

	url := HOST + "resume"
	expected := "{\"Code\":9,\"Message\":\"Cannot resume. No song was paused\"}"
	checkResult("POST", url, expected, t)
}
<commit_msg>add tests for stop<commit_after>package main

import (
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const HOST = "http://localhost:8765/"

func performCall(method string, url string) (string, error) {
	var res *http.Response
	var err error
	if method == "GET" {
		res, err = http.Get(url)
	} else if method == "POST" {
		res, err = http.Post(url, "text/plain", nil) //todo:
	} else if method == "PUT" {
		client := &http.Client{}
		request, err1 := http.NewRequest("PUT", url, nil)
		if err1 != nil {
			return "", err1
		}
		res, err = client.Do(request)
	}

	if err != nil {
		return "", err
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return "", err
	}
	found := string(body[:])
	return found, nil
}

func checkResult(method, url, expected string, t *testing.T) {
	found, err := performCall(method, url)
	if err != nil {
		t.Fatalf("Unexpected error found - %s", err.Error())
	}

	if found != expected {
		t.Errorf("Expected\n---\n%s\n---\nbut found\n---\n%s\n---\n", expected, found)
	}
}

func TestGetAlive(t *testing.T) {
	expected := "I'm alive"
	checkResult("GET", HOST, expected, t)
}

func TestPlay(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds/beep9.mp3")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayDir(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\",\"test_sounds/beep28.mp3\",\"test_sounds/beep36.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayPlaylist(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("sample_playlist")
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Started playing\",\"Data\":[\"test_sounds/beep9.mp3\",\"test_sounds/beep28.mp3\",\"test_sounds/beep36.mp3\"]}"
	checkResult("PUT", url, expected, t)
}

func TestPlayNonExistingFile(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_sounds/beep1.mp3")
	expected := "{\"Code\":4,\"Message\":\"File cannot be found\"}"
	checkResult("PUT", url, expected, t)
}

func TestPlayInvalidFileFormat(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_broken/abc.txt")
	expected := "{\"Code\":7,\"Message\":\"Format is not supported\"}"
	checkResult("PUT", url, expected, t)
}

func TestPlayBrokenFile(t *testing.T) {
	url := HOST + "play/" + url.QueryEscape("test_broken/abc.txt")
	expected := "{\"Code\":2,\"Message\":\"SoX failed to open input file\"}"
	checkResult("PUT", url, expected, t)
}

func TestPause(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)
	url := HOST + "pause"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Song is paused\",\"Data\":[\"test_sounds/beep28.mp3\"]}"

	checkResult("POST", url, expected, t)
}

func TestPauseNoPlayback(t *testing.T) {
	url := HOST + "pause"
	expected := "{\"Code\":2,\"Message\":\"Cannot pause. No song is playing\"}"
	checkResult("POST", url, expected, t)
}

func TestResume(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)
	pause_url := HOST + "pause"
	performCall("PUT", pause_url)

	url := HOST + "resume"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Song is resumed\",\"Data\":[\"test_sounds/beep28.mp3\"]}"
	checkResult("POST", url, expected, t)

}

func TestResumeNoPlayback(t *testing.T) {
	url := HOST + "resume"
	expected := "{\"Code\":9,\"Message\":\"Cannot resume. No song was paused\"}"
	checkResult("POST", url, expected, t)
}

func TestResumeNoPaused(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)

	url := HOST + "resume"
	expected := "{\"Code\":9,\"Message\":\"Cannot resume. No song was paused\"}"
	checkResult("POST", url, expected, t)
}

func TestStop(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)

	url := HOST + "stop"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Playback is stopped and cleaned\",\"Data\":[]}"
	checkResult("PUT", url, expected, t)
}

func TestStopNoPlayback(t *testing.T) {
	url := HOST + "stop"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Playback is stopped and cleaned\",\"Data\":[]}"
	checkResult("PUT", url, expected, t)
}

func TestStopPaused(t *testing.T) {
	play_url := HOST + "play/" + url.QueryEscape("test_sounds/beep28.mp3")
	performCall("PUT", play_url)
	pause_url := HOST + "pause"
	performCall("PUT", pause_url)

	url := HOST + "stop"
	expected := "{\"Code\":0,\"Message\":\"Success\",\"Info\":\"Playback is stopped and cleaned\",\"Data\":[]}"
	checkResult("PUT", url, expected, t)
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package gofetch

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
)

// ProgressReport represents the current download progress of a given file
type ProgressReport struct {
	URL string
	// Total length in bytes of the file being downloaded
	Total int64
	// Written bytes to disk on a write by write basis. It does not accumulate.
	WrittenBytes int64
}

// goFetch represents an instance of gofetch, holding global configuration options.
type goFetch struct {
	destDir     string
	etag        bool
	concurrency int
	httpClient  *http.Client
}

// Option as explained in http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type Option func(*goFetch)

// DestDir allows you to set the destination directory for the downloaded files.
func DestDir(dir string) Option {
	return func(f *goFetch) {
		f.destDir = dir
	}
}

// Concurrency allows you to set the number of goroutines used to download a specific
// file.
func Concurrency(c int) Option {
	return func(f *goFetch) {
		f.concurrency = c
	}
}

// ETag allows you to disable or enable ETag support, meaning that if an already
// downloaded file is currently on disk and matches the ETag returned by the server,
// it will not be downloaded again.
func ETag(enable bool) Option {
	return func(f *goFetch) {
		f.etag = enable
	}
}

// New creates a new instance of goFetch with the given options.
func New(opts ...Option) *goFetch {
	// Creates instance and assigns defaults.
	gofetch := &goFetch{
		concurrency: 1,
		destDir:     "./",
		etag:        true,
		httpClient:  new(http.Client),
	}

	for _, opt := range opts {
		opt(gofetch)
	}
	return gofetch
}

// Fetch downloads content from the provided URL. It supports resuming and
// parallelizing downloads while being very memory efficient.
func (gf *goFetch) Fetch(url string, progressCh chan<- ProgressReport) (*os.File, error) {
	if url == "" {
		return nil, errors.New("URL is required")
	}

	// We need to make a preflight request to get the size of the content.
	res, err := http.Head(url)
	if err != nil {
		return nil, err
	}

	if !strings.HasPrefix(res.Status, "2") {
		return nil, errors.New("HTTP requests returned a non 2xx status code")
	}

	fileName := path.Base(url)

	var etag string
	if gf.etag {
		etag = res.Header.Get("ETag")
		fileName += strings.Trim(etag, `"`)
	}

	destFilePath := filepath.Join(gf.destDir, fileName)

	fi, err := os.Stat(destFilePath)
	if err == nil && fi.Size() == res.ContentLength {
		if progressCh != nil {
			close(progressCh)
		}
		return os.Open(destFilePath)
	}

	return gf.parallelFetch(url, destFilePath, res.ContentLength, progressCh)
}

// parallelFetch fetches using multiple goroutines, each piece is streamed down
// to disk which makes it very efficient in terms of memory usage.
func (gf *goFetch) parallelFetch(url, destFilePath string, length int64, progressCh chan<- ProgressReport) (*os.File, error) {
	if progressCh != nil {
		defer close(progressCh)
	}

	var wg sync.WaitGroup

	report := ProgressReport{Total: length}
	concurrency := int64(gf.concurrency)
	chunkSize := length / concurrency
	remainingSize := length % concurrency
	chunksDir := filepath.Join(gf.destDir, path.Base(url)+".chunks")

	if err := os.MkdirAll(chunksDir, 0760); err != nil {
		return nil, err
	}

	var errs []error
	for i := int64(0); i < concurrency; i++ {
		min := chunkSize * i
		max := chunkSize * (i + 1)

		if i == (concurrency - 1) {
			// Add the remaining bytes in the last request
			max += remainingSize
		}

		wg.Add(1)
		go func(min, max int64, chunkNumber int) {
			defer wg.Done()
			chunkFile := filepath.Join(chunksDir, strconv.Itoa(chunkNumber))

			err := gf.fetch(url, chunkFile, min, max, report, progressCh)
			if err != nil {
				errs = append(errs, err)
			}
		}(min, max, int(i))
	}
	wg.Wait()

	if len(errs) > 0 {
		return nil, fmt.Errorf("Errors: \n %s", errs)
	}

	file, err := gf.assembleChunks(destFilePath, chunksDir)
	if err != nil {
		return nil, err
	}

	os.RemoveAll(chunksDir)

	// Makes sure to return the file on the correct offset so it can be
	// consumed by users.
	_, err = file.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return file, err
}

// assembleChunks join all the data pieces together
func (gf *goFetch) assembleChunks(destFile, chunksDir string) (*os.File, error) {
	file, err := os.Create(destFile)
	if err != nil {
		return nil, err
	}

	for i := 0; i < gf.concurrency; i++ {
		chunkFile, err := os.Open(filepath.Join(chunksDir, strconv.Itoa(i)))
		if err != nil {
			return nil, err
		}

		if _, err := io.Copy(file, chunkFile); err != nil {
			return nil, err
		}
		chunkFile.Close()
	}
	return file, nil
}

// fetch downloads files using one unbuffered HTTP connection and supports
// resuming downloads if interrupted.
func (gf *goFetch) fetch(url, destFile string, min, max int64,
	report ProgressReport, progressCh chan<- ProgressReport) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	// In order to resume previous interrupted downloads we need to open the file
	// in append mode.
	file, err := os.OpenFile(destFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	if err != nil {
		return err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	currSize := fi.Size()

	// There is nothing to do if file exists and was fully downloaded.
	// We do substraction between max and min to account for the last chunk
	// size, which may be of different size if division between res.ContentLength and config.SizeLimit
	// is not exact.
	if currSize == (max - min) {
		return nil
	}

	// Report bytes written already into the file
	if progressCh != nil {
		report.WrittenBytes = currSize
		progressCh <- report
	}

	// Adjusts min to resume file download from where it was left off.
	if currSize > 0 {
		min = min + currSize
	}

	// Prepares writer to report download progress.
	writer := fetchWriter{
		Writer:         file,
		progressCh:     progressCh,
		progressReport: report,
	}

	brange := fmt.Sprintf("bytes=%d-%d", min, max-1)
	if max == -1 {
		brange = fmt.Sprintf("bytes=%d-", min)
	}

	req.Header.Add("Range", brange)
	res, err := gf.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if !strings.HasPrefix(res.Status, "2") {
		return errors.New("HTTP requests returned a non 2xx status code")
	}

	_, err = io.Copy(&writer, res.Body)
	return err
}

// fetchWriter implements a custom io.Writer so we can send granular
// progress reports when streaming down content.
type fetchWriter struct {
	io.Writer
	//progressCh is the channel sent by the user to get download updates.
	progressCh chan<- ProgressReport
	// report is the structure sent through the progress channel.
	progressReport ProgressReport
}

func (fw *fetchWriter) Write(b []byte) (int, error) {
	n, err := fw.Writer.Write(b)

	if fw.progressCh != nil {
		fw.progressReport.WrittenBytes = int64(n)
		fw.progressCh <- fw.progressReport
	}

	return n, err
}
<commit_msg>Makes Option private so it does not show up in godoc.org<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package gofetch

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
)

// ProgressReport represents the current download progress of a given file
type ProgressReport struct {
	URL string
	// Total length in bytes of the file being downloaded
	Total int64
	// Written bytes to disk on a write by write basis. It does not accumulate.
	WrittenBytes int64
}

// goFetch represents an instance of gofetch, holding global configuration options.
type goFetch struct {
	destDir     string
	etag        bool
	concurrency int
	httpClient  *http.Client
}

// option as explained in http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type option func(*goFetch)

// DestDir allows you to set the destination directory for the downloaded files.
func DestDir(dir string) option {
	return func(f *goFetch) {
		f.destDir = dir
	}
}

// Concurrency allows you to set the number of goroutines used to download a specific
// file.
func Concurrency(c int) option {
	return func(f *goFetch) {
		f.concurrency = c
	}
}

// ETag allows you to disable or enable ETag support, meaning that if an already
// downloaded file is currently on disk and matches the ETag returned by the server,
// it will not be downloaded again.
func ETag(enable bool) option {
	return func(f *goFetch) {
		f.etag = enable
	}
}

// New creates a new instance of goFetch with the given options.
func New(opts ...option) *goFetch {
	// Creates instance and assigns defaults.
	gofetch := &goFetch{
		concurrency: 1,
		destDir:     "./",
		etag:        true,
		httpClient:  new(http.Client),
	}

	for _, opt := range opts {
		opt(gofetch)
	}
	return gofetch
}

// Fetch downloads content from the provided URL. It supports resuming and
// parallelizing downloads while being very memory efficient.
func (gf *goFetch) Fetch(url string, progressCh chan<- ProgressReport) (*os.File, error) {
	if url == "" {
		return nil, errors.New("URL is required")
	}

	// We need to make a preflight request to get the size of the content.
	res, err := http.Head(url)
	if err != nil {
		return nil, err
	}

	if !strings.HasPrefix(res.Status, "2") {
		return nil, errors.New("HTTP requests returned a non 2xx status code")
	}

	fileName := path.Base(url)

	var etag string
	if gf.etag {
		etag = res.Header.Get("ETag")
		fileName += strings.Trim(etag, `"`)
	}

	destFilePath := filepath.Join(gf.destDir, fileName)

	fi, err := os.Stat(destFilePath)
	if err == nil && fi.Size() == res.ContentLength {
		if progressCh != nil {
			close(progressCh)
		}
		return os.Open(destFilePath)
	}

	return gf.parallelFetch(url, destFilePath, res.ContentLength, progressCh)
}

// parallelFetch fetches using multiple goroutines, each piece is streamed down
// to disk which makes it very efficient in terms of memory usage.
func (gf *goFetch) parallelFetch(url, destFilePath string, length int64, progressCh chan<- ProgressReport) (*os.File, error) {
	if progressCh != nil {
		defer close(progressCh)
	}

	var wg sync.WaitGroup

	report := ProgressReport{Total: length}
	concurrency := int64(gf.concurrency)
	chunkSize := length / concurrency
	remainingSize := length % concurrency
	chunksDir := filepath.Join(gf.destDir, path.Base(url)+".chunks")

	if err := os.MkdirAll(chunksDir, 0760); err != nil {
		return nil, err
	}

	var errs []error
	for i := int64(0); i < concurrency; i++ {
		min := chunkSize * i
		max := chunkSize * (i + 1)

		if i == (concurrency - 1) {
			// Add the remaining bytes in the last request
			max += remainingSize
		}

		wg.Add(1)
		go func(min, max int64, chunkNumber int) {
			defer wg.Done()
			chunkFile := filepath.Join(chunksDir, strconv.Itoa(chunkNumber))

			err := gf.fetch(url, chunkFile, min, max, report, progressCh)
			if err != nil {
				errs = append(errs, err)
			}
		}(min, max, int(i))
	}
	wg.Wait()

	if len(errs) > 0 {
		return nil, fmt.Errorf("Errors: \n %s", errs)
	}

	file, err := gf.assembleChunks(destFilePath, chunksDir)
	if err != nil {
		return nil, err
	}

	os.RemoveAll(chunksDir)

	// Makes sure to return the file on the correct offset so it can be
	// consumed by users.
	_, err = file.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return file, err
}

// assembleChunks join all the data pieces together
func (gf *goFetch) assembleChunks(destFile, chunksDir string) (*os.File, error) {
	file, err := os.Create(destFile)
	if err != nil {
		return nil, err
	}

	for i := 0; i < gf.concurrency; i++ {
		chunkFile, err := os.Open(filepath.Join(chunksDir, strconv.Itoa(i)))
		if err != nil {
			return nil, err
		}

		if _, err := io.Copy(file, chunkFile); err != nil {
			return nil, err
		}
		chunkFile.Close()
	}
	return file, nil
}

// fetch downloads files using one unbuffered HTTP connection and supports
// resuming downloads if interrupted.
func (gf *goFetch) fetch(url, destFile string, min, max int64,
	report ProgressReport, progressCh chan<- ProgressReport) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	// In order to resume previous interrupted downloads we need to open the file
	// in append mode.
	file, err := os.OpenFile(destFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	if err != nil {
		return err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	currSize := fi.Size()

	// There is nothing to do if file exists and was fully downloaded.
	// We do substraction between max and min to account for the last chunk
	// size, which may be of different size if division between res.ContentLength and config.SizeLimit
	// is not exact.
	if currSize == (max - min) {
		return nil
	}

	// Report bytes written already into the file
	if progressCh != nil {
		report.WrittenBytes = currSize
		progressCh <- report
	}

	// Adjusts min to resume file download from where it was left off.
	if currSize > 0 {
		min = min + currSize
	}

	// Prepares writer to report download progress.
	writer := fetchWriter{
		Writer:         file,
		progressCh:     progressCh,
		progressReport: report,
	}

	brange := fmt.Sprintf("bytes=%d-%d", min, max-1)
	if max == -1 {
		brange = fmt.Sprintf("bytes=%d-", min)
	}

	req.Header.Add("Range", brange)
	res, err := gf.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if !strings.HasPrefix(res.Status, "2") {
		return errors.New("HTTP requests returned a non 2xx status code")
	}

	_, err = io.Copy(&writer, res.Body)
	return err
}

// fetchWriter implements a custom io.Writer so we can send granular
// progress reports when streaming down content.
type fetchWriter struct {
	io.Writer
	//progressCh is the channel sent by the user to get download updates.
	progressCh chan<- ProgressReport
	// report is the structure sent through the progress channel.
	progressReport ProgressReport
}

func (fw *fetchWriter) Write(b []byte) (int, error) {
	n, err := fw.Writer.Write(b)

	if fw.progressCh != nil {
		fw.progressReport.WrittenBytes = int64(n)
		fw.progressCh <- fw.progressReport
	}

	return n, err
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package gofetch

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/mitchellh/go-homedir"
)

// ProgressReport represents the current download progress of a given file
type ProgressReport struct {
	// Total length in bytes of the file being downloaded
	Total int64
	// Written bytes to disk on a write by write basis. It does not accumulate.
	WrittenBytes int64
}

// goFetch represents an instance of gofetch, holding global configuration options.
type goFetch struct {
	destDir     string
	etag        bool
	concurrency int
	httpClient  *http.Client
}

// option as explained in http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type option func(*goFetch)

// DestDir allows you to set the destination directory for the downloaded files.
// By default it is set to: ./
func DestDir(dir string) option {
	return func(f *goFetch) {
		f.destDir = dir
	}
}

// Concurrency allows you to set the number of goroutines used to download a specific
// file. By default it is set to 1.
func Concurrency(c int) option {
	return func(f *goFetch) {
		f.concurrency = c
	}
}

// ETag allows you to disable or enable ETag support, meaning that if an already
// downloaded file is currently on disk and matches the ETag value returned by the server,
// it will not be downloaded again. By default it is set to true.

// Be aware that different servers, serving the same file, are likely to return
// different ETag values, causing the file to be re-downloaded, even though it
// might already exist on disk.
func ETag(enable bool) option {
	return func(f *goFetch) {
		f.etag = enable
	}
}

var workDir string

func init() {
	dir, err := homedir.Dir()
	if err != nil {
		fmt.Printf(`Unable to get user home directory err=%s\n`, err)
	}
	workDir = filepath.Join(dir, ".gofetch")
}

// New creates a new instance of goFetch with the given options.
func New(opts ...option) *goFetch {
	// Creates instance and assigns defaults.
	gofetch := &goFetch{
		concurrency: 1,
		destDir:     "./",
		etag:        true,
		httpClient:  new(http.Client),
	}

	for _, opt := range opts {
		opt(gofetch)
	}
	return gofetch
}

// Fetch downloads content from the provided URL. It supports resuming and
// parallelizing downloads while being very memory efficient.
func (gf *goFetch) Fetch(url string, progressCh chan<- ProgressReport) (*os.File, error) {
	if url == "" {
		return nil, errors.New("URL is required")
	}

	// We need to make a preflight request to get the size of the content.
	res, err := http.Head(url)
	if err != nil {
		return nil, err
	}

	if !strings.HasPrefix(res.Status, "2") {
		return nil, errors.New("HTTP requests returned a non 2xx status code")
	}

	fileName := path.Base(url)
	destFilePath := filepath.Join(gf.destDir, fileName)

	var etag string
	if gf.etag {
		// Go's stdlib returns header value enclosed in double quotes.
		etag = strings.Trim(res.Header.Get("ETag"), `"`)

		if etag == "" {
			goto FETCH
		}

		// Create directory if it doesn't exist, we ignore error if it already exists.
		cacheDir := filepath.Join(workDir, fileName)
		etagPath := filepath.Join(cacheDir, etag)
		os.MkdirAll(cacheDir, 0700)

		if _, err := os.Stat(etagPath); err == nil {
			// Our file has been already fully downloaded, return a file
			// descriptor to it and skip fetching altogether.
			fi, err := os.Stat(destFilePath)
			if err == nil && fi.Size() == res.ContentLength {
				if progressCh != nil {
					close(progressCh)
				}
				return os.Open(destFilePath)
			}
		} else {
			f, _ := os.Create(etagPath)
			f.Close()
		}
	}

FETCH:
	return gf.parallelFetch(url, destFilePath, res.ContentLength, progressCh)
}

// parallelFetch fetches using multiple goroutines, each piece is streamed down
// to disk which makes it very efficient in terms of memory usage.
func (gf *goFetch) parallelFetch(url, destFilePath string, length int64, progressCh chan<- ProgressReport) (*os.File, error) {
	if progressCh != nil {
		defer close(progressCh)
	}

	var wg sync.WaitGroup

	report := ProgressReport{Total: length}
	concurrency := int64(gf.concurrency)
	chunkSize := length / concurrency
	remainingSize := length % concurrency
	chunksDir := filepath.Join(gf.destDir, path.Base(url)+".chunks")

	if err := os.MkdirAll(chunksDir, 0760); err != nil {
		return nil, err
	}

	var errs []error
	for i := int64(0); i < concurrency; i++ {
		min := chunkSize * i
		max := chunkSize * (i + 1)

		if i == (concurrency - 1) {
			// Add the remaining bytes in the last request
			max += remainingSize
		}

		wg.Add(1)
		go func(min, max int64, chunkNumber int) {
			defer wg.Done()
			chunkFile := filepath.Join(chunksDir, strconv.Itoa(chunkNumber))

			err := gf.fetch(url, chunkFile, min, max, report, progressCh)
			if err != nil {
				errs = append(errs, err)
			}
		}(min, max, int(i))
	}
	wg.Wait()

	if len(errs) > 0 {
		return nil, fmt.Errorf("Errors: \n %s", errs)
	}

	file, err := gf.assembleChunks(destFilePath, chunksDir)
	if err != nil {
		return nil, err
	}

	os.RemoveAll(chunksDir)

	// Makes sure to return the file on the correct offset so it can be
	// consumed by users.
	_, err = file.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return file, err
}

// assembleChunks join all the data pieces together
func (gf *goFetch) assembleChunks(destFile, chunksDir string) (*os.File, error) {
	file, err := os.Create(destFile)
	if err != nil {
		return nil, err
	}

	for i := 0; i < gf.concurrency; i++ {
		chunkFile, err := os.Open(filepath.Join(chunksDir, strconv.Itoa(i)))
		if err != nil {
			return nil, err
		}

		if _, err := io.Copy(file, chunkFile); err != nil {
			return nil, err
		}
		chunkFile.Close()
	}
	return file, nil
}

// fetch downloads files using one unbuffered HTTP connection and supports
// resuming downloads if interrupted.
func (gf *goFetch) fetch(url, destFile string, min, max int64,
	report ProgressReport, progressCh chan<- ProgressReport) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	// In order to resume previous interrupted downloads we need to open the file
	// in append mode.
	file, err := os.OpenFile(destFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	if err != nil {
		return err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	currSize := fi.Size()

	// There is nothing to do if file exists and was fully downloaded.
	// We do substraction between max and min to account for the last chunk
	// size, which may be of different size if division between res.ContentLength and config.SizeLimit
	// is not exact.
	if currSize == (max - min) {
		return nil
	}

	// Report bytes written already into the file
	if progressCh != nil {
		report.WrittenBytes = currSize
		progressCh <- report
	}

	// Adjusts min to resume file download from where it was left off.
	if currSize > 0 {
		min = min + currSize
	}

	// Prepares writer to report download progress.
	writer := fetchWriter{
		Writer:         file,
		progressCh:     progressCh,
		progressReport: report,
	}

	brange := fmt.Sprintf("bytes=%d-%d", min, max-1)
	if max == -1 {
		brange = fmt.Sprintf("bytes=%d-", min)
	}

	req.Header.Add("Range", brange)
	res, err := gf.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if !strings.HasPrefix(res.Status, "2") {
		return errors.New("HTTP requests returned a non 2xx status code")
	}

	_, err = io.Copy(&writer, res.Body)
	return err
}

// fetchWriter implements a custom io.Writer so we can send granular
// progress reports when streaming down content.
type fetchWriter struct {
	io.Writer
	//progressCh is the channel sent by the user to get download updates.
	progressCh chan<- ProgressReport
	// report is the structure sent through the progress channel.
	progressReport ProgressReport
}

func (fw *fetchWriter) Write(b []byte) (int, error) {
	n, err := fw.Writer.Write(b)

	if fw.progressCh != nil {
		fw.progressReport.WrittenBytes = int64(n)
		fw.progressCh <- fw.progressReport
	}

	return n, err
}
<commit_msg>Makes type public so that users can declare gofetch variables with it<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, version 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package gofetch

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/mitchellh/go-homedir"
)

// ProgressReport represents the current download progress of a given file
type ProgressReport struct {
	// Total length in bytes of the file being downloaded
	Total int64
	// Written bytes to disk on a write by write basis. It does not accumulate.
	WrittenBytes int64
}

// Fetcher represents an instance of gofetch, holding global configuration options.
type Fetcher struct {
	destDir     string
	etag        bool
	concurrency int
	httpClient  *http.Client
}

// option as explained in http://commandcenter.blogspot.com/2014/01/self-referential-functions-and-design.html
type option func(*Fetcher)

// DestDir allows you to set the destination directory for the downloaded files.
// By default it is set to: ./
func DestDir(dir string) option {
	return func(f *Fetcher) {
		f.destDir = dir
	}
}

// Concurrency allows you to set the number of goroutines used to download a specific
// file. By default it is set to 1.
func Concurrency(c int) option {
	return func(f *Fetcher) {
		f.concurrency = c
	}
}

// ETag allows you to disable or enable ETag support, meaning that if an already
// downloaded file is currently on disk and matches the ETag value returned by the server,
// it will not be downloaded again. By default it is set to true.

// Be aware that different servers, serving the same file, are likely to return
// different ETag values, causing the file to be re-downloaded, even though it
// might already exist on disk.
func ETag(enable bool) option {
	return func(f *Fetcher) {
		f.etag = enable
	}
}

var workDir string

func init() {
	dir, err := homedir.Dir()
	if err != nil {
		fmt.Printf(`Unable to get user home directory err=%s\n`, err)
	}
	workDir = filepath.Join(dir, ".gofetch")
}

// New creates a new instance of goFetch with the given options.
func New(opts ...option) *Fetcher {
	// Creates instance and assigns defaults.
	gofetch := &Fetcher{
		concurrency: 1,
		destDir:     "./",
		etag:        true,
		httpClient:  new(http.Client),
	}

	for _, opt := range opts {
		opt(gofetch)
	}
	return gofetch
}

// Fetch downloads content from the provided URL. It supports resuming and
// parallelizing downloads while being very memory efficient.
func (gf *Fetcher) Fetch(url string, progressCh chan<- ProgressReport) (*os.File, error) {
	if url == "" {
		return nil, errors.New("URL is required")
	}

	// We need to make a preflight request to get the size of the content.
	res, err := http.Head(url)
	if err != nil {
		return nil, err
	}

	if !strings.HasPrefix(res.Status, "2") {
		return nil, errors.New("HTTP requests returned a non 2xx status code")
	}

	fileName := path.Base(url)
	destFilePath := filepath.Join(gf.destDir, fileName)

	var etag string
	if gf.etag {
		// Go's stdlib returns header value enclosed in double quotes.
		etag = strings.Trim(res.Header.Get("ETag"), `"`)

		if etag == "" {
			goto FETCH
		}

		// Create directory if it doesn't exist, we ignore error if it already exists.
		cacheDir := filepath.Join(workDir, fileName)
		etagPath := filepath.Join(cacheDir, etag)
		os.MkdirAll(cacheDir, 0700)

		if _, err := os.Stat(etagPath); err == nil {
			// Our file has been already fully downloaded, return a file
			// descriptor to it and skip fetching altogether.
			fi, err := os.Stat(destFilePath)
			if err == nil && fi.Size() == res.ContentLength {
				if progressCh != nil {
					close(progressCh)
				}
				return os.Open(destFilePath)
			}
		} else {
			f, _ := os.Create(etagPath)
			f.Close()
		}
	}

FETCH:
	return gf.parallelFetch(url, destFilePath, res.ContentLength, progressCh)
}

// parallelFetch fetches using multiple goroutines, each piece is streamed down
// to disk which makes it very efficient in terms of memory usage.
func (gf *Fetcher) parallelFetch(url, destFilePath string, length int64, progressCh chan<- ProgressReport) (*os.File, error) {
	if progressCh != nil {
		defer close(progressCh)
	}

	var wg sync.WaitGroup

	report := ProgressReport{Total: length}
	concurrency := int64(gf.concurrency)
	chunkSize := length / concurrency
	remainingSize := length % concurrency
	chunksDir := filepath.Join(gf.destDir, path.Base(url)+".chunks")

	if err := os.MkdirAll(chunksDir, 0760); err != nil {
		return nil, err
	}

	var errs []error
	for i := int64(0); i < concurrency; i++ {
		min := chunkSize * i
		max := chunkSize * (i + 1)

		if i == (concurrency - 1) {
			// Add the remaining bytes in the last request
			max += remainingSize
		}

		wg.Add(1)
		go func(min, max int64, chunkNumber int) {
			defer wg.Done()
			chunkFile := filepath.Join(chunksDir, strconv.Itoa(chunkNumber))

			err := gf.fetch(url, chunkFile, min, max, report, progressCh)
			if err != nil {
				errs = append(errs, err)
			}
		}(min, max, int(i))
	}
	wg.Wait()

	if len(errs) > 0 {
		return nil, fmt.Errorf("Errors: \n %s", errs)
	}

	file, err := gf.assembleChunks(destFilePath, chunksDir)
	if err != nil {
		return nil, err
	}

	os.RemoveAll(chunksDir)

	// Makes sure to return the file on the correct offset so it can be
	// consumed by users.
	_, err = file.Seek(0, 0)
	if err != nil {
		return nil, err
	}

	return file, err
}

// assembleChunks join all the data pieces together
func (gf *Fetcher) assembleChunks(destFile, chunksDir string) (*os.File, error) {
	file, err := os.Create(destFile)
	if err != nil {
		return nil, err
	}

	for i := 0; i < gf.concurrency; i++ {
		chunkFile, err := os.Open(filepath.Join(chunksDir, strconv.Itoa(i)))
		if err != nil {
			return nil, err
		}

		if _, err := io.Copy(file, chunkFile); err != nil {
			return nil, err
		}
		chunkFile.Close()
	}
	return file, nil
}

// fetch downloads files using one unbuffered HTTP connection and supports
// resuming downloads if interrupted.
func (gf *Fetcher) fetch(url, destFile string, min, max int64,
	report ProgressReport, progressCh chan<- ProgressReport) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}

	// In order to resume previous interrupted downloads we need to open the file
	// in append mode.
	file, err := os.OpenFile(destFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	if err != nil {
		return err
	}
	defer file.Close()

	fi, err := file.Stat()
	if err != nil {
		return err
	}
	currSize := fi.Size()

	// There is nothing to do if file exists and was fully downloaded.
	// We do substraction between max and min to account for the last chunk
	// size, which may be of different size if division between res.ContentLength and config.SizeLimit
	// is not exact.
	if currSize == (max - min) {
		return nil
	}

	// Report bytes written already into the file
	if progressCh != nil {
		report.WrittenBytes = currSize
		progressCh <- report
	}

	// Adjusts min to resume file download from where it was left off.
	if currSize > 0 {
		min = min + currSize
	}

	// Prepares writer to report download progress.
	writer := fetchWriter{
		Writer:         file,
		progressCh:     progressCh,
		progressReport: report,
	}

	brange := fmt.Sprintf("bytes=%d-%d", min, max-1)
	if max == -1 {
		brange = fmt.Sprintf("bytes=%d-", min)
	}

	req.Header.Add("Range", brange)
	res, err := gf.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if !strings.HasPrefix(res.Status, "2") {
		return errors.New("HTTP requests returned a non 2xx status code")
	}

	_, err = io.Copy(&writer, res.Body)
	return err
}

// fetchWriter implements a custom io.Writer so we can send granular
// progress reports when streaming down content.
type fetchWriter struct {
	io.Writer
	//progressCh is the channel sent by the user to get download updates.
	progressCh chan<- ProgressReport
	// report is the structure sent through the progress channel.
	progressReport ProgressReport
}

func (fw *fetchWriter) Write(b []byte) (int, error) {
	n, err := fw.Writer.Write(b)

	if fw.progressCh != nil {
		fw.progressReport.WrittenBytes = int64(n)
		fw.progressCh <- fw.progressReport
	}

	return n, err
}
<|endoftext|>