<commit_before>package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"syscall"
	"time"

	"github.com/weaveworks/procspy"
	"github.com/weaveworks/scope/report"
	"github.com/weaveworks/scope/xfer"
)

func main() {
	var (
		httpListen         = flag.String("http.listen", "", "listen address for HTTP profiling and instrumentation server")
		publishInterval    = flag.Duration("publish.interval", 1*time.Second, "publish (output) interval")
		spyInterval        = flag.Duration("spy.interval", 100*time.Millisecond, "spy (scan) interval")
		listen             = flag.String("listen", ":"+strconv.Itoa(xfer.ProbePort), "listen address")
		prometheusEndpoint = flag.String("prometheus.endpoint", "/metrics", "Prometheus metrics exposition endpoint (requires -http.listen)")
		spyProcs           = flag.Bool("processes", true, "report processes (needs root)")
		cgroupsRoot        = flag.String("cgroups.root", "", "if provided, enrich -processes with cgroup names from this root (e.g. /mnt/cgroups)")
		cgroupsInterval    = flag.Duration("cgroups.interval", 10*time.Second, "how often to update cgroup names")
		dockerMapper       = flag.Bool("docker", true, "collect docker-related attributes for processes.")
		dockerInterval     = flag.Duration("docker.interval", 10*time.Second, "how often to update docker container info")
		procRoot           = flag.String("proc.root", "/proc", "location of the proc filesystem")
	)
	flag.Parse()

	if len(flag.Args()) != 0 {
		flag.Usage()
		os.Exit(1)
	}

	procspy.SetProcRoot(*procRoot)

	if *httpListen != "" {
		log.Printf("profiling data being exported to %s", *httpListen)
		log.Printf("go tool pprof http://%s/debug/pprof/{profile,heap,block}", *httpListen)
		if *prometheusEndpoint != "" {
			log.Printf("exposing Prometheus endpoint at %s%s", *httpListen, *prometheusEndpoint)
			http.Handle(*prometheusEndpoint, makePrometheusHandler())
		}
		go http.ListenAndServe(*httpListen, nil)
	}

	if *spyProcs && os.Getegid() != 0 {
		log.Printf("warning: process reporting enabled, but that requires root to find everything")
	}

	publisher, err := xfer.NewTCPPublisher(*listen)
	if err != nil {
		log.Fatal(err)
	}
	defer publisher.Close()

	pms := []processMapper{identityMapper{}}

	if *cgroupsRoot != "" {
		if fi, err := os.Stat(*cgroupsRoot); err == nil && fi.IsDir() {
			log.Printf("enriching -processes with cgroup names from %s", *cgroupsRoot)
			pms = append(pms, newCgroupMapper(*cgroupsRoot, *cgroupsInterval))
		} else {
			log.Printf("-cgroups.root=%s: %v", *cgroupsRoot, err)
		}
	}

	if *dockerMapper {
		docker := newDockerMapper(*procRoot, *dockerInterval)
		pms = append(pms, &dockerIDMapper{docker}, &dockerNameMapper{docker})
	}

	log.Printf("listening on %s", *listen)

	go func() {
		var (
			hostname = hostname()
			nodeID   = hostname // TODO: we should sanitize the hostname
			pubTick  = time.Tick(*publishInterval)
			spyTick  = time.Tick(*spyInterval)
			r        = report.NewReport()
		)

		for {
			select {
			case <-pubTick:
				publishTicks.WithLabelValues().Add(1)
				r.HostMetadatas[nodeID] = hostMetadata(hostname)
				publisher.Publish(r)
				r = report.NewReport()

			case <-spyTick:
				r.Merge(spy(hostname, hostname, *spyProcs, pms))
				// log.Printf("merged report:\n%#v\n", r)
			}
		}
	}()

	log.Printf("%s", <-interrupt())
}

func interrupt() chan os.Signal {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	return c
}

// hostMetadata produces an instantaneous HostMetadata for this host. No need
// to do this more than once per published report.
func hostMetadata(hostname string) report.HostMetadata {
	loadOne, loadFive, loadFifteen := getLoads()

	host := report.HostMetadata{
		Timestamp:   time.Now().UTC(),
		Hostname:    hostname,
		OS:          runtime.GOOS,
		LoadOne:     loadOne,
		LoadFive:    loadFive,
		LoadFifteen: loadFifteen,
	}

	if localNets, err := net.InterfaceAddrs(); err == nil {
		// Not all networks are IP networks.
		for _, localNet := range localNets {
			if net, ok := localNet.(*net.IPNet); ok {
				host.LocalNets = append(host.LocalNets, net)
			}
		}
	}

	return host
}
<commit_msg>Reduce default spy and publish intervals<commit_after>package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"syscall"
	"time"

	"github.com/weaveworks/procspy"
	"github.com/weaveworks/scope/report"
	"github.com/weaveworks/scope/xfer"
)

func main() {
	var (
		httpListen         = flag.String("http.listen", "", "listen address for HTTP profiling and instrumentation server")
		publishInterval    = flag.Duration("publish.interval", 3*time.Second, "publish (output) interval")
		spyInterval        = flag.Duration("spy.interval", time.Second, "spy (scan) interval")
		listen             = flag.String("listen", ":"+strconv.Itoa(xfer.ProbePort), "listen address")
		prometheusEndpoint = flag.String("prometheus.endpoint", "/metrics", "Prometheus metrics exposition endpoint (requires -http.listen)")
		spyProcs           = flag.Bool("processes", true, "report processes (needs root)")
		cgroupsRoot        = flag.String("cgroups.root", "", "if provided, enrich -processes with cgroup names from this root (e.g. /mnt/cgroups)")
		cgroupsInterval    = flag.Duration("cgroups.interval", 10*time.Second, "how often to update cgroup names")
		dockerMapper       = flag.Bool("docker", true, "collect Docker-related attributes for processes")
		dockerInterval     = flag.Duration("docker.interval", 10*time.Second, "how often to update Docker attributes")
		procRoot           = flag.String("proc.root", "/proc", "location of the proc filesystem")
	)
	flag.Parse()

	if len(flag.Args()) != 0 {
		flag.Usage()
		os.Exit(1)
	}

	procspy.SetProcRoot(*procRoot)

	if *httpListen != "" {
		log.Printf("profiling data being exported to %s", *httpListen)
		log.Printf("go tool pprof http://%s/debug/pprof/{profile,heap,block}", *httpListen)
		if *prometheusEndpoint != "" {
			log.Printf("exposing Prometheus endpoint at %s%s", *httpListen, *prometheusEndpoint)
			http.Handle(*prometheusEndpoint, makePrometheusHandler())
		}
		go http.ListenAndServe(*httpListen, nil)
	}

	if *spyProcs && os.Getegid() != 0 {
		log.Printf("warning: process reporting enabled, but that requires root to find everything")
	}

	publisher, err := xfer.NewTCPPublisher(*listen)
	if err != nil {
		log.Fatal(err)
	}
	defer publisher.Close()

	pms := []processMapper{identityMapper{}}

	if *cgroupsRoot != "" {
		if fi, err := os.Stat(*cgroupsRoot); err == nil && fi.IsDir() {
			log.Printf("enriching -processes with cgroup names from %s", *cgroupsRoot)
			pms = append(pms, newCgroupMapper(*cgroupsRoot, *cgroupsInterval))
		} else {
			log.Printf("-cgroups.root=%s: %v", *cgroupsRoot, err)
		}
	}

	if *dockerMapper {
		docker := newDockerMapper(*procRoot, *dockerInterval)
		pms = append(pms, &dockerIDMapper{docker}, &dockerNameMapper{docker})
	}

	log.Printf("listening on %s", *listen)

	go func() {
		var (
			hostname = hostname()
			nodeID   = hostname // TODO: we should sanitize the hostname
			pubTick  = time.Tick(*publishInterval)
			spyTick  = time.Tick(*spyInterval)
			r        = report.NewReport()
		)

		for {
			select {
			case <-pubTick:
				publishTicks.WithLabelValues().Add(1)
				r.HostMetadatas[nodeID] = hostMetadata(hostname)
				publisher.Publish(r)
				r = report.NewReport()

			case <-spyTick:
				r.Merge(spy(hostname, hostname, *spyProcs, pms))
				// log.Printf("merged report:\n%#v\n", r)
			}
		}
	}()

	log.Printf("%s", <-interrupt())
}

func interrupt() chan os.Signal {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	return c
}

// hostMetadata produces an instantaneous HostMetadata for this host. No need
// to do this more than once per published report.
func hostMetadata(hostname string) report.HostMetadata {
	loadOne, loadFive, loadFifteen := getLoads()

	host := report.HostMetadata{
		Timestamp:   time.Now().UTC(),
		Hostname:    hostname,
		OS:          runtime.GOOS,
		LoadOne:     loadOne,
		LoadFive:    loadFive,
		LoadFifteen: loadFifteen,
	}

	if localNets, err := net.InterfaceAddrs(); err == nil {
		// Not all networks are IP networks.
		for _, localNet := range localNets {
			if net, ok := localNet.(*net.IPNet); ok {
				host.LocalNets = append(host.LocalNets, net)
			}
		}
	}

	return host
}
<|endoftext|>


<commit_before>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 totp

import (
	"encoding/base32"
	"github.com/stretchr/testify/require"

	"testing"
	"time"
)

type tc struct {
	TS     int64
	TOTP   string
	Mode   Algorithm
	Secret string
}

//
// Test vectors from http://tools.ietf.org/html/rfc6238#appendix-B
// NOTE -- the test vectors are documented as having the SAME
// secret -- this is WRONG -- they have a variable secret
// depending upon the hmac algorithm:
// 		http://www.rfc-editor.org/errata_search.php?rfc=6238
// this only took a few hours of head/desk interaction to figure out.
//
func TestValidateRFCMatrix(t *testing.T) {
	secSha1 := base32.StdEncoding.EncodeToString([]byte("12345678901234567890"))
	secSha256 := base32.StdEncoding.EncodeToString([]byte("12345678901234567890123456789012"))
	secSha512 := base32.StdEncoding.EncodeToString([]byte("1234567890123456789012345678901234567890123456789012345678901234"))

	tests := []tc{
		tc{59, "94287082", AlgorithmSHA1, secSha1},
		tc{59, "46119246", AlgorithmSHA256, secSha256},
		tc{59, "90693936", AlgorithmSHA512, secSha512},
		tc{1111111109, "07081804", AlgorithmSHA1, secSha1},
		tc{1111111109, "68084774", AlgorithmSHA256, secSha256},
		tc{1111111109, "25091201", AlgorithmSHA512, secSha512},
		tc{1111111111, "14050471", AlgorithmSHA1, secSha1},
		tc{1111111111, "67062674", AlgorithmSHA256, secSha256},
		tc{1111111111, "99943326", AlgorithmSHA512, secSha512},
		tc{1234567890, "89005924", AlgorithmSHA1, secSha1},
		tc{1234567890, "91819424", AlgorithmSHA256, secSha256},
		tc{1234567890, "93441116", AlgorithmSHA512, secSha512},
		tc{2000000000, "69279037", AlgorithmSHA1, secSha1},
		tc{2000000000, "90698825", AlgorithmSHA256, secSha256},
		tc{2000000000, "38618901", AlgorithmSHA512, secSha512},
		tc{20000000000, "65353130", AlgorithmSHA1, secSha1},
		tc{20000000000, "77737706", AlgorithmSHA256, secSha256},
		tc{20000000000, "47863826", AlgorithmSHA512, secSha512},
	}

	for _, tx := range tests {
		valid, err := ValidateCustom(tx.TOTP, tx.Secret, time.Unix(tx.TS, 0).UTC(),
			ValidateOpts{
				Digits:    DigitsEight,
				Algorithm: tx.Mode,
			})
		require.NoError(t, err,
			"unexpected error totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
		require.True(t, valid,
			"unexpected totp failure totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
	}
}

func TestGenerate(t *testing.T) {
	k, err := Generate(GenerateOpts{
		Issuer:      "SnakeOil",
		AccountName: "alice@example.com",
	})
	require.NoError(t, err, "generate basic TOTP")
	require.Equal(t, "SnakeOil", k.Issuer(), "Extracting Issuer")
	require.Equal(t, "alice@example.com", k.AccountName(), "Extracting Account Name")
	require.Equal(t, 16, len(k.Secret()), "Secret is 16 bytes long as base32.")

	k, err = Generate(GenerateOpts{
		Issuer:      "SnakeOil",
		AccountName: "alice@example.com",
		SecretSize:  20,
	})
	require.NoError(t, err, "generate larger TOTP")
	require.Equal(t, 32, len(k.Secret()), "Secret is 32 bytes long as base32.")
}
<commit_msg>Add test for Skew support in TOTP<commit_after>/**
 *  Copyright 2014 Paul Querna
 *
 *  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 totp

import (
	"encoding/base32"
	"github.com/stretchr/testify/require"

	"testing"
	"time"
)

type tc struct {
	TS     int64
	TOTP   string
	Mode   Algorithm
	Secret string
}

//
// Test vectors from http://tools.ietf.org/html/rfc6238#appendix-B
// NOTE -- the test vectors are documented as having the SAME
// secret -- this is WRONG -- they have a variable secret
// depending upon the hmac algorithm:
// 		http://www.rfc-editor.org/errata_search.php?rfc=6238
// this only took a few hours of head/desk interaction to figure out.
//
func TestValidateRFCMatrix(t *testing.T) {
	secSha1 := base32.StdEncoding.EncodeToString([]byte("12345678901234567890"))
	secSha256 := base32.StdEncoding.EncodeToString([]byte("12345678901234567890123456789012"))
	secSha512 := base32.StdEncoding.EncodeToString([]byte("1234567890123456789012345678901234567890123456789012345678901234"))

	tests := []tc{
		tc{59, "94287082", AlgorithmSHA1, secSha1},
		tc{59, "46119246", AlgorithmSHA256, secSha256},
		tc{59, "90693936", AlgorithmSHA512, secSha512},
		tc{1111111109, "07081804", AlgorithmSHA1, secSha1},
		tc{1111111109, "68084774", AlgorithmSHA256, secSha256},
		tc{1111111109, "25091201", AlgorithmSHA512, secSha512},
		tc{1111111111, "14050471", AlgorithmSHA1, secSha1},
		tc{1111111111, "67062674", AlgorithmSHA256, secSha256},
		tc{1111111111, "99943326", AlgorithmSHA512, secSha512},
		tc{1234567890, "89005924", AlgorithmSHA1, secSha1},
		tc{1234567890, "91819424", AlgorithmSHA256, secSha256},
		tc{1234567890, "93441116", AlgorithmSHA512, secSha512},
		tc{2000000000, "69279037", AlgorithmSHA1, secSha1},
		tc{2000000000, "90698825", AlgorithmSHA256, secSha256},
		tc{2000000000, "38618901", AlgorithmSHA512, secSha512},
		tc{20000000000, "65353130", AlgorithmSHA1, secSha1},
		tc{20000000000, "77737706", AlgorithmSHA256, secSha256},
		tc{20000000000, "47863826", AlgorithmSHA512, secSha512},
	}

	for _, tx := range tests {
		valid, err := ValidateCustom(tx.TOTP, tx.Secret, time.Unix(tx.TS, 0).UTC(),
			ValidateOpts{
				Digits:    DigitsEight,
				Algorithm: tx.Mode,
			})
		require.NoError(t, err,
			"unexpected error totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
		require.True(t, valid,
			"unexpected totp failure totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
	}
}

func TestValidateSkew(t *testing.T) {
	secSha1 := base32.StdEncoding.EncodeToString([]byte("12345678901234567890"))

	tests := []tc{
		tc{29, "94287082", AlgorithmSHA1, secSha1},
		tc{59, "94287082", AlgorithmSHA1, secSha1},
		tc{61, "94287082", AlgorithmSHA1, secSha1},
	}

	for _, tx := range tests {
		valid, err := ValidateCustom(tx.TOTP, tx.Secret, time.Unix(tx.TS, 0).UTC(),
			ValidateOpts{
				Digits:    DigitsEight,
				Algorithm: tx.Mode,
				Skew:      1,
			})
		require.NoError(t, err,
			"unexpected error totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
		require.True(t, valid,
			"unexpected totp failure totp=%s mode=%v ts=%v", tx.TOTP, tx.Mode, tx.TS)
	}
}

func TestGenerate(t *testing.T) {
	k, err := Generate(GenerateOpts{
		Issuer:      "SnakeOil",
		AccountName: "alice@example.com",
	})
	require.NoError(t, err, "generate basic TOTP")
	require.Equal(t, "SnakeOil", k.Issuer(), "Extracting Issuer")
	require.Equal(t, "alice@example.com", k.AccountName(), "Extracting Account Name")
	require.Equal(t, 16, len(k.Secret()), "Secret is 16 bytes long as base32.")

	k, err = Generate(GenerateOpts{
		Issuer:      "SnakeOil",
		AccountName: "alice@example.com",
		SecretSize:  20,
	})
	require.NoError(t, err, "generate larger TOTP")
	require.Equal(t, 32, len(k.Secret()), "Secret is 32 bytes long as base32.")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/color"
	"github.com/jroimartin/gocui"
)

var state = &State{
	ActionIndex: -1,
	Settings:    &Settings{},
	HideHelp:    len(os.Args) > 1,
}

// State is the central function managing the information of claws.
type State struct {
	// important to running the application as a whole
	ActionIndex int
	Mode        int
	Writer      io.Writer
	Conn        *WebSocket

	// important for drawing
	FirstDrawDone     bool
	ShouldQuit        bool
	HideHelp          bool
	KeepAutoscrolling bool

	// functions
	ExecuteFunc func(func(*gocui.Gui) error)

	Settings *Settings
}

// PushAction adds an action to LastActions
func (s *State) PushAction(act string) {
	s.Settings.LastActions = append([]string{act}, s.Settings.LastActions...)
	if len(s.Settings.LastActions) > 100 {
		s.Settings.LastActions = s.Settings.LastActions[:100]
	}
	s.Settings.Save()
}

// BrowseActions changes the ActionIndex and returns the value at the specified index.
// move is the number of elements to move (negatives go into more recent history,
// 0 returns the current element, positives go into older history)
func (s *State) BrowseActions(move int) string {
	s.ActionIndex += move
	if s.ActionIndex >= len(s.Settings.LastActions) {
		s.ActionIndex = len(s.Settings.LastActions) - 1
	} else if s.ActionIndex < -1 {
		s.ActionIndex = -1
	}

	// -1 always indicates the "next" element, thus empty
	if s.ActionIndex == -1 {
		return ""
	}
	return s.Settings.LastActions[s.ActionIndex]
}

// StartConnection begins a WebSocket connection to url.
func (s *State) StartConnection(url string) error {
	if s.Conn != nil {
		return errors.New("state: conn is not nil")
	}
	ws, err := CreateWebSocket(url)
	if err != nil {
		return err
	}
	s.Conn = ws
	connectionStarted = time.Now()
	go s.wsReader()
	return nil
}

func (s *State) wsReader() {
	ch := s.Conn.ReadChannel()
	for msg := range ch {
		s.Server(msg)
	}
}

var (
	printDebug  = color.New(color.FgCyan).Fprint
	printError  = color.New(color.FgRed).Fprint
	printUser   = color.New(color.FgGreen).Fprint
	printServer = color.New(color.FgWhite).Fprint
)

// Debug prints debug information to the Writer, using light blue.
func (s *State) Debug(x string) {
	s.printToOut(printDebug, x)
}

// Error prints an error to the Writer, using red.
func (s *State) Error(x string) {
	s.printToOut(printError, x)
}

// User prints user-provided messages to the Writer, using green.
func (s *State) User(x string) {
	res, err := s.pipe(x, "out", s.Settings.Pipe.Out)
	if err != nil {
		s.Error(err.Error())
		if res == "" || res == "\n" {
			return
		}
	}
	s.printToOut(printUser, res)
}

// Server prints server-returned messages to the Writer, using white.
func (s *State) Server(x string) {
	res, err := s.pipe(x, "in", s.Settings.Pipe.In)
	if err != nil {
		s.Error(err.Error())
		if res == "" || res == "\n" {
			return
		}
	}
	if s.Settings.JSONFormatting {
		res = attemptJSONFormatting(res)
	}
	s.printToOut(printServer, x)
}

var (
	sessionStarted    = time.Now()
	connectionStarted time.Time
)

func (s *State) pipe(data, t string, command []string) (string, error) {
	if len(command) < 1 {
		return data, nil
	}
	// prepare the command: create it, set up env variables
	c := exec.Command(command[0], command[1:]...)
	c.Env = append(
		os.Environ(),
		"CLAWS_PIPE_TYPE="+t,
		"CLAWS_SESSION="+strconv.FormatInt(sessionStarted.UnixNano(), 10),
		"CLAWS_CONNECTION="+strconv.FormatInt(sessionStarted.UnixNano(), 10),
	)
	// set up stdin
	stdin := strings.NewReader(data)
	c.Stdin = stdin

	// run the command
	res, err := c.Output()
	return string(res), err
}

func (s *State) printToOut(f func(io.Writer, ...interface{}) (int, error), str string) {
	if s.Settings.Timestamp != "" {
		str = time.Now().Format(s.Settings.Timestamp) + str
	}

	if str != "" && str[len(str)-1] != '\n' {
		str += "\n"
	}
	s.ExecuteFunc(func(*gocui.Gui) error {
		_, err := f(s.Writer, str)
		return err
	})
}

// Settings contains persistent information about the usage of claws.
type Settings struct {
	JSONFormatting   bool
	Timestamp        string
	LastWebsocketURL string
	LastActions      []string
	Pipe             struct {
		In  []string
		Out []string
	}
}

// Load loads settings from ~/.config/claws.json
func (s *Settings) Load() error {
	folder, err := getConfigFolder()
	if err != nil {
		return err
	}

	f, err := os.Open(folder + "claws.json")
	if err != nil {
		// silently ignore ErrNotExist
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	defer f.Close()

	return json.NewDecoder(f).Decode(s)
}

// Save saves settings to ~/.config/claws.json
func (s Settings) Save() error {
	folder, err := getConfigFolder()
	if err != nil {
		return err
	}

	f, err := os.Create(folder + "claws.json")
	if err != nil {
		return err
	}
	defer f.Close()

	e := json.NewEncoder(f)
	e.SetIndent("", "\t")
	return e.Encode(s)
}

func getConfigFolder() (string, error) {
	u, err := user.Current()
	if err != nil {
		return "", err
	}
	folder := u.HomeDir + "/.config/"

	err = os.MkdirAll(folder, 0755)

	return folder, err
}
<commit_msg>add CLAWS_WS_URL<commit_after>package main

import (
	"encoding/json"
	"errors"
	"io"
	"os"
	"os/exec"
	"os/user"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/color"
	"github.com/jroimartin/gocui"
)

var state = &State{
	ActionIndex: -1,
	Settings:    &Settings{},
	HideHelp:    len(os.Args) > 1,
}

// State is the central function managing the information of claws.
type State struct {
	// important to running the application as a whole
	ActionIndex int
	Mode        int
	Writer      io.Writer
	Conn        *WebSocket

	// important for drawing
	FirstDrawDone     bool
	ShouldQuit        bool
	HideHelp          bool
	KeepAutoscrolling bool

	// functions
	ExecuteFunc func(func(*gocui.Gui) error)

	Settings *Settings
}

// PushAction adds an action to LastActions
func (s *State) PushAction(act string) {
	s.Settings.LastActions = append([]string{act}, s.Settings.LastActions...)
	if len(s.Settings.LastActions) > 100 {
		s.Settings.LastActions = s.Settings.LastActions[:100]
	}
	s.Settings.Save()
}

// BrowseActions changes the ActionIndex and returns the value at the specified index.
// move is the number of elements to move (negatives go into more recent history,
// 0 returns the current element, positives go into older history)
func (s *State) BrowseActions(move int) string {
	s.ActionIndex += move
	if s.ActionIndex >= len(s.Settings.LastActions) {
		s.ActionIndex = len(s.Settings.LastActions) - 1
	} else if s.ActionIndex < -1 {
		s.ActionIndex = -1
	}

	// -1 always indicates the "next" element, thus empty
	if s.ActionIndex == -1 {
		return ""
	}
	return s.Settings.LastActions[s.ActionIndex]
}

// StartConnection begins a WebSocket connection to url.
func (s *State) StartConnection(url string) error {
	if s.Conn != nil {
		return errors.New("state: conn is not nil")
	}
	ws, err := CreateWebSocket(url)
	if err != nil {
		return err
	}
	s.Conn = ws
	connectionStarted = time.Now()
	go s.wsReader()
	return nil
}

func (s *State) wsReader() {
	ch := s.Conn.ReadChannel()
	for msg := range ch {
		s.Server(msg)
	}
}

var (
	printDebug  = color.New(color.FgCyan).Fprint
	printError  = color.New(color.FgRed).Fprint
	printUser   = color.New(color.FgGreen).Fprint
	printServer = color.New(color.FgWhite).Fprint
)

// Debug prints debug information to the Writer, using light blue.
func (s *State) Debug(x string) {
	s.printToOut(printDebug, x)
}

// Error prints an error to the Writer, using red.
func (s *State) Error(x string) {
	s.printToOut(printError, x)
}

// User prints user-provided messages to the Writer, using green.
func (s *State) User(x string) {
	res, err := s.pipe(x, "out", s.Settings.Pipe.Out)
	if err != nil {
		s.Error(err.Error())
		if res == "" || res == "\n" {
			return
		}
	}
	s.printToOut(printUser, res)
}

// Server prints server-returned messages to the Writer, using white.
func (s *State) Server(x string) {
	res, err := s.pipe(x, "in", s.Settings.Pipe.In)
	if err != nil {
		s.Error(err.Error())
		if res == "" || res == "\n" {
			return
		}
	}
	if s.Settings.JSONFormatting {
		res = attemptJSONFormatting(res)
	}
	s.printToOut(printServer, x)
}

var (
	sessionStarted    = time.Now()
	connectionStarted time.Time
)

func (s *State) pipe(data, t string, command []string) (string, error) {
	if len(command) < 1 {
		return data, nil
	}
	// prepare the command: create it, set up env variables
	c := exec.Command(command[0], command[1:]...)
	c.Env = append(
		os.Environ(),
		"CLAWS_PIPE_TYPE="+t,
		"CLAWS_SESSION="+strconv.FormatInt(sessionStarted.UnixNano(), 10),
		"CLAWS_CONNECTION="+strconv.FormatInt(sessionStarted.UnixNano(), 10),
		"CLAWS_WS_URL="+s.Conn.URL(),
	)
	// set up stdin
	stdin := strings.NewReader(data)
	c.Stdin = stdin

	// run the command
	res, err := c.Output()
	return string(res), err
}

func (s *State) printToOut(f func(io.Writer, ...interface{}) (int, error), str string) {
	if s.Settings.Timestamp != "" {
		str = time.Now().Format(s.Settings.Timestamp) + str
	}

	if str != "" && str[len(str)-1] != '\n' {
		str += "\n"
	}
	s.ExecuteFunc(func(*gocui.Gui) error {
		_, err := f(s.Writer, str)
		return err
	})
}

// Settings contains persistent information about the usage of claws.
type Settings struct {
	JSONFormatting   bool
	Timestamp        string
	LastWebsocketURL string
	LastActions      []string
	Pipe             struct {
		In  []string
		Out []string
	}
}

// Load loads settings from ~/.config/claws.json
func (s *Settings) Load() error {
	folder, err := getConfigFolder()
	if err != nil {
		return err
	}

	f, err := os.Open(folder + "claws.json")
	if err != nil {
		// silently ignore ErrNotExist
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	defer f.Close()

	return json.NewDecoder(f).Decode(s)
}

// Save saves settings to ~/.config/claws.json
func (s Settings) Save() error {
	folder, err := getConfigFolder()
	if err != nil {
		return err
	}

	f, err := os.Create(folder + "claws.json")
	if err != nil {
		return err
	}
	defer f.Close()

	e := json.NewEncoder(f)
	e.SetIndent("", "\t")
	return e.Encode(s)
}

func getConfigFolder() (string, error) {
	u, err := user.Current()
	if err != nil {
		return "", err
	}
	folder := u.HomeDir + "/.config/"

	err = os.MkdirAll(folder, 0755)

	return folder, err
}
<|endoftext|>


<commit_before>package fuse

import (
	"fmt"
	"io"
	"os"
	"syscall"
)

var _ = fmt.Println

// DataFile is for implementing read-only filesystems.  This
// assumes we already have the data in memory.
type DataFile struct {
	data []byte

	DefaultFile
}

func (me *DataFile) String() string {
	l := len(me.data)
	if l > 10 {
		l = 10
	}

	return fmt.Sprintf("DataFile(%x)", me.data[:l])
}

func (me *DataFile) GetAttr() (*Attr, Status) {
	return &Attr{Mode: S_IFREG | 0644, Size: uint64(len(me.data))}, OK
}

func NewDataFile(data []byte) *DataFile {
	f := new(DataFile)
	f.data = data
	return f
}

func (me *DataFile) Read(input *ReadIn, bp BufferPool) ([]byte, Status) {
	end := int(input.Offset) + int(input.Size)
	if end > len(me.data) {
		end = len(me.data)
	}

	return me.data[input.Offset:end], OK
}

////////////////

// DevNullFile accepts any write, and always returns EOF.
type DevNullFile struct {
	DefaultFile
}

func NewDevNullFile() *DevNullFile {
	return new(DevNullFile)
}

func (me *DevNullFile) String() string {
	return "DevNullFile"
}

func (me *DevNullFile) Read(input *ReadIn, bp BufferPool) ([]byte, Status) {
	return []byte{}, OK
}

func (me *DevNullFile) Write(input *WriteIn, content []byte) (uint32, Status) {
	return uint32(len(content)), OK
}

func (me *DevNullFile) Flush() Status {
	return OK
}

func (me *DevNullFile) Fsync(flags int) (code Status) {
	return OK
}

func (me *DevNullFile) Truncate(size uint64) (code Status) {
	return OK
}

////////////////

// LoopbackFile delegates all operations back to an underlying os.File.
type LoopbackFile struct {
	File *os.File

	DefaultFile
}

func (me *LoopbackFile) String() string {
	return fmt.Sprintf("LoopbackFile(%s)", me.File.Name())
}

func (me *LoopbackFile) Read(input *ReadIn, buffers BufferPool) ([]byte, Status) {
	slice := buffers.AllocBuffer(input.Size)

	n, err := me.File.ReadAt(slice, int64(input.Offset))
	if err == io.EOF {
		err = nil
	}
	return slice[:n], ToStatus(err)
}

func (me *LoopbackFile) Write(input *WriteIn, data []byte) (uint32, Status) {
	n, err := me.File.WriteAt(data, int64(input.Offset))
	return uint32(n), ToStatus(err)
}

func (me *LoopbackFile) Release() {
	me.File.Close()
}

func (me *LoopbackFile) Fsync(flags int) (code Status) {
	return ToStatus(syscall.Fsync(int(me.File.Fd())))
}

func (me *LoopbackFile) Truncate(size uint64) Status {
	return ToStatus(syscall.Ftruncate(int(me.File.Fd()), int64(size)))
}

// futimens missing from 6g runtime.

func (me *LoopbackFile) Chmod(mode uint32) Status {
	return ToStatus(me.File.Chmod(os.FileMode(mode)))
}

func (me *LoopbackFile) Chown(uid uint32, gid uint32) Status {
	return ToStatus(me.File.Chown(int(uid), int(gid)))
}

func (me *LoopbackFile) GetAttr() (*Attr, Status) {
	st := syscall.Stat_t{}
	err := syscall.Fstat(int(me.File.Fd()), &st)
	if err != nil {
		return nil, ToStatus(err)
	}
	a := &Attr{}
	a.FromStat(&st)
	return a, OK
}

////////////////////////////////////////////////////////////////

// ReadOnlyFile is a wrapper that denies writable operations
type ReadOnlyFile struct {
	File
}

func (me *ReadOnlyFile) String() string {
	return fmt.Sprintf("ReadOnlyFile(%s)", me.File.String())
}

func (me *ReadOnlyFile) Write(input *WriteIn, data []byte) (uint32, Status) {
	return 0, EPERM
}

func (me *ReadOnlyFile) Fsync(flag int) (code Status) {
	return OK
}

func (me *ReadOnlyFile) Truncate(size uint64) Status {
	return EPERM
}

func (me *ReadOnlyFile) Chmod(mode uint32) Status {
	return EPERM
}

func (me *ReadOnlyFile) Chown(uid uint32, gid uint32) Status {
	return EPERM
}
<commit_msg>LoopbackFile.Flush() always returns OK.<commit_after>package fuse

import (
	"fmt"
	"io"
	"os"
	"syscall"
)

var _ = fmt.Println

// DataFile is for implementing read-only filesystems.  This
// assumes we already have the data in memory.
type DataFile struct {
	data []byte

	DefaultFile
}

func (me *DataFile) String() string {
	l := len(me.data)
	if l > 10 {
		l = 10
	}

	return fmt.Sprintf("DataFile(%x)", me.data[:l])
}

func (me *DataFile) GetAttr() (*Attr, Status) {
	return &Attr{Mode: S_IFREG | 0644, Size: uint64(len(me.data))}, OK
}

func NewDataFile(data []byte) *DataFile {
	f := new(DataFile)
	f.data = data
	return f
}

func (me *DataFile) Read(input *ReadIn, bp BufferPool) ([]byte, Status) {
	end := int(input.Offset) + int(input.Size)
	if end > len(me.data) {
		end = len(me.data)
	}

	return me.data[input.Offset:end], OK
}

////////////////

// DevNullFile accepts any write, and always returns EOF.
type DevNullFile struct {
	DefaultFile
}

func NewDevNullFile() *DevNullFile {
	return new(DevNullFile)
}

func (me *DevNullFile) String() string {
	return "DevNullFile"
}

func (me *DevNullFile) Read(input *ReadIn, bp BufferPool) ([]byte, Status) {
	return []byte{}, OK
}

func (me *DevNullFile) Write(input *WriteIn, content []byte) (uint32, Status) {
	return uint32(len(content)), OK
}

func (me *DevNullFile) Flush() Status {
	return OK
}

func (me *DevNullFile) Fsync(flags int) (code Status) {
	return OK
}

func (me *DevNullFile) Truncate(size uint64) (code Status) {
	return OK
}

////////////////

// LoopbackFile delegates all operations back to an underlying os.File.
type LoopbackFile struct {
	File *os.File

	DefaultFile
}

func (me *LoopbackFile) String() string {
	return fmt.Sprintf("LoopbackFile(%s)", me.File.Name())
}

func (me *LoopbackFile) Read(input *ReadIn, buffers BufferPool) ([]byte, Status) {
	slice := buffers.AllocBuffer(input.Size)

	n, err := me.File.ReadAt(slice, int64(input.Offset))
	if err == io.EOF {
		err = nil
	}
	return slice[:n], ToStatus(err)
}

func (me *LoopbackFile) Write(input *WriteIn, data []byte) (uint32, Status) {
	n, err := me.File.WriteAt(data, int64(input.Offset))
	return uint32(n), ToStatus(err)
}

func (me *LoopbackFile) Release() {
	me.File.Close()
}

func (me *LoopbackFile) Flush() Status {
	return OK
}

func (me *LoopbackFile) Fsync(flags int) (code Status) {
	return ToStatus(syscall.Fsync(int(me.File.Fd())))
}

func (me *LoopbackFile) Truncate(size uint64) Status {
	return ToStatus(syscall.Ftruncate(int(me.File.Fd()), int64(size)))
}

// futimens missing from 6g runtime.

func (me *LoopbackFile) Chmod(mode uint32) Status {
	return ToStatus(me.File.Chmod(os.FileMode(mode)))
}

func (me *LoopbackFile) Chown(uid uint32, gid uint32) Status {
	return ToStatus(me.File.Chown(int(uid), int(gid)))
}

func (me *LoopbackFile) GetAttr() (*Attr, Status) {
	st := syscall.Stat_t{}
	err := syscall.Fstat(int(me.File.Fd()), &st)
	if err != nil {
		return nil, ToStatus(err)
	}
	a := &Attr{}
	a.FromStat(&st)
	return a, OK
}

////////////////////////////////////////////////////////////////

// ReadOnlyFile is a wrapper that denies writable operations
type ReadOnlyFile struct {
	File
}

func (me *ReadOnlyFile) String() string {
	return fmt.Sprintf("ReadOnlyFile(%s)", me.File.String())
}

func (me *ReadOnlyFile) Write(input *WriteIn, data []byte) (uint32, Status) {
	return 0, EPERM
}

func (me *ReadOnlyFile) Fsync(flag int) (code Status) {
	return OK
}

func (me *ReadOnlyFile) Truncate(size uint64) Status {
	return EPERM
}

func (me *ReadOnlyFile) Chmod(mode uint32) Status {
	return EPERM
}

func (me *ReadOnlyFile) Chown(uid uint32, gid uint32) Status {
	return EPERM
}
<|endoftext|>


<commit_before>package termite

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

var _ = fmt.Println

type FsServer struct {
	contentServer  *ContentServer
	contentCache   *ContentCache
	Root           string
	excluded       map[string]bool
	excludePrivate bool

	multiplyPaths func(string) []string

	hashCacheMutex sync.RWMutex
	hashCacheCond  *sync.Cond
	hashCache      map[string]string
	hashBusyMap    map[string]bool

	attrCacheMutex sync.RWMutex
	attrCache      map[string]*FileAttr
	attrCacheCond  *sync.Cond
	attrCacheBusy  map[string]bool

	// TODO - add counters and check that the rpcFs.fetchCond is
	// working.
}

var paranoia = false

func (me *FsServer) verify() {
	if !paranoia {
		return
	}
	me.attrCacheMutex.RLock()
	defer me.attrCacheMutex.RUnlock()

	for k, v := range me.attrCache {
		if v.Path != k {
			panic(fmt.Sprintf("attrCache mismatch %q %#v", k, v))
		}
		if _, ok := me.attrCacheBusy[k]; ok {
			panic(fmt.Sprintf("attrCacheBusy and attrCache entry for %q", k))
		}
	}
}

func NewFsServer(root string, cache *ContentCache, excluded []string) *FsServer {
	fs := &FsServer{
		contentCache:  cache,
		contentServer: &ContentServer{Cache: cache},
		Root:          root,
		hashCache:     make(map[string]string),
		hashBusyMap:   map[string]bool{},

		attrCache:      make(map[string]*FileAttr),
		attrCacheBusy:  map[string]bool{},
		excludePrivate: true,
	}

	fs.hashCacheCond = sync.NewCond(&fs.hashCacheMutex)
	fs.attrCacheCond = sync.NewCond(&fs.attrCacheMutex)
	fs.excluded = make(map[string]bool)
	for _, e := range excluded {
		fs.excluded[e] = true
	}
	return fs
}

func (me FileAttr) String() string {
	id := ""
	if me.Hash != "" {
		id = fmt.Sprintf(" sz %d", me.FileInfo.Size)
	}
	if me.Link != "" {
		id = fmt.Sprintf(" -> %s", me.Link)
	}
	if me.Deletion() {
		id = " (del)"
	} else if !me.Status.Ok() {
		id = " " + me.Status.String()
	}
	if me.Status.Ok() {
		id += fmt.Sprintf(" m=%o", me.FileInfo.Mode)
	}

	return fmt.Sprintf("%s%s", me.Path, id)
}

func (me FileAttr) Deletion() bool {
	return me.Status == fuse.ENOENT
}

func (me *FsServer) path(n string) string {
	if me.Root == "" {
		return n
	}
	return filepath.Join(me.Root, strings.TrimLeft(n, "/"))
}

func (me *FsServer) FileContent(req *ContentRequest, rep *ContentResponse) os.Error {
	return me.contentServer.FileContent(req, rep)
}

func (me *FsServer) ReadDir(req *DirRequest, r *DirResponse) os.Error {
	d, e := ioutil.ReadDir(me.path(req.Name))
	log.Println("ReadDir", req)
	r.NameModeMap = make(map[string]uint32)
	for _, v := range d {
		r.NameModeMap[v.Name] = v.Mode
	}
	r.Status = fuse.OsErrorToErrno(e)
	return nil
}

func (me *FsServer) GetAttr(req *AttrRequest, rep *AttrResponse) os.Error {
	log.Println("GetAttr req", req.Name)
	names := []string{}
	if me.multiplyPaths != nil {
		names = me.multiplyPaths(req.Name)
	} else {
		names = append(names, req.Name)
	}
	for _, n := range names {
		a := me.oneGetAttr(n)
		if a.Hash != "" {
			log.Printf("GetAttr %q %v %x", n, a, a.Hash)
		}
		rep.Attrs = append(rep.Attrs, a)
	}
	return nil
}

func (me *FsServer) oneGetAttr(name string) (rep *FileAttr) {
	defer me.verify()

	if me.excluded[name] {
		log.Printf("Denied access to excluded file %q", name)
		return &FileAttr{
			Path:   name,
			Status: fuse.ENOENT,
		}
	}

	me.attrCacheMutex.RLock()
	rep, ok := me.attrCache[name]
	me.attrCacheMutex.RUnlock()
	if ok {
		return rep
	}

	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()
	for me.attrCacheBusy[name] && me.attrCache[name] == nil {
		me.attrCacheCond.Wait()
	}
	rep, ok = me.attrCache[name]
	if ok {
		return rep
	}
	me.attrCacheBusy[name] = true
	me.attrCacheMutex.Unlock()

	p := me.path(name)
	fi, err := os.Lstat(p)
	rep = &FileAttr{
		FileInfo: fi,
		Status:   fuse.OsErrorToErrno(err),
		Path:     name,
	}

	// We don't want to expose the master's private files to the
	// world.
	if me.excludePrivate && fi != nil && fi.Mode&0077 == 0 {
		log.Printf("Denied access to private file %q", name)
		rep.FileInfo = nil
		rep.Status = fuse.EPERM
		fi = nil
	}

	if fi != nil {
		me.fillContent(rep)
	}

	me.attrCacheMutex.Lock()
	me.attrCache[name] = rep
	me.attrCacheCond.Broadcast()
	me.attrCacheBusy[name] = false, false
	return rep
}

func (me *FsServer) fillContent(rep *FileAttr) {
	if rep.FileInfo.IsSymlink() {
		rep.Link, _ = os.Readlink(rep.Path)
	}
	if rep.FileInfo.IsRegular() {
		rep.Hash = me.getHash(rep.Path)
		if rep.Hash == "" {
			// Typically happens if we want to open /etc/shadow as normal user.
			rep.Status = fuse.EPERM
		}
	}
}

func (me *FsServer) updateFiles(infos []*FileAttr) {
	me.updateHashes(infos)
	me.updateAttrs(infos)
}

func (me *FsServer) updateAttrs(infos []*FileAttr) {
	defer me.verify()

	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()

	for _, r := range infos {
		me.attrCache[r.Path] = r
	}
}

func (me *FsServer) updateHashes(infos []*FileAttr) {
	defer me.verify()

	me.hashCacheMutex.Lock()
	defer me.hashCacheMutex.Unlock()

	for _, r := range infos {
		name := r.Path
		if !r.Status.Ok() || r.Link != "" {
			me.hashCache[name] = "", false
			me.hashBusyMap[name] = false, false
		}
		if r.Hash != "" {
			me.hashCache[name] = r.Hash
			me.hashBusyMap[name] = false, false
		}
	}
}

func (me *FsServer) getHash(name string) (hash string) {
	fullPath := me.path(name)

	me.hashCacheMutex.RLock()
	hash = me.hashCache[name]
	me.hashCacheMutex.RUnlock()

	if hash != "" {
		return hash
	}

	me.hashCacheMutex.Lock()
	defer me.hashCacheMutex.Unlock()
	for me.hashBusyMap[name] && me.hashCache[name] == "" {
		me.hashCacheCond.Wait()
	}

	hash = me.hashCache[name]
	if hash != "" {
		return hash
	}
	me.hashBusyMap[name] = true
	me.hashCacheMutex.Unlock()

	// TODO - /usr should be configurable.
	if HasDirPrefix(fullPath, "/usr") && !HasDirPrefix(fullPath, "/usr/local") {
		hash = me.contentCache.SaveImmutablePath(fullPath)
	} else {
		hash = me.contentCache.SavePath(fullPath)
	}

	me.hashCacheMutex.Lock()
	me.hashCache[name] = hash
	me.hashBusyMap[name] = false, false
	me.hashCacheCond.Broadcast()

	return hash
}

func (me *FsServer) refreshAttributeCache(prefix string) []*FileAttr {
	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()

	updated := []*FileAttr{}
	for key, attr := range me.attrCache {
		// TODO -should just do everything?
		if !HasDirPrefix(key, prefix) {
			continue
		}

		fi, _ := os.Lstat(me.path(key))
		if fi == nil && attr.Status.Ok() {
			del := FileAttr{
				Path:   key,
				Status: fuse.ENOENT,
			}
			updated = append(updated, &del)
		}
		if fi != nil && attr.FileInfo != nil && EncodeFileInfo(*attr.FileInfo) != EncodeFileInfo(*fi) {
			newEnt := FileAttr{
				Path:     key,
				Status:   fuse.OK,
				FileInfo: fi,
			}
			me.fillContent(&newEnt)
			updated = append(updated, &newEnt)
		}
	}

	for _, u := range updated {
		copy := *u
		me.attrCache[u.Path] = &copy
	}
	return updated
}

func (me *FsServer) copyCache() []*FileAttr {
	me.attrCacheMutex.RLock()
	defer me.attrCacheMutex.RUnlock()

	dump := []*FileAttr{}
	for _, attr := range me.attrCache {
		copy := *attr
		dump = append(dump, &copy)
	}

	return dump
}
<commit_msg>Drop superfluous string.<commit_after>package termite

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

var _ = fmt.Println

type FsServer struct {
	contentServer  *ContentServer
	contentCache   *ContentCache
	Root           string
	excluded       map[string]bool
	excludePrivate bool

	multiplyPaths func(string) []string

	hashCacheMutex sync.RWMutex
	hashCacheCond  *sync.Cond
	hashCache      map[string]string
	hashBusyMap    map[string]bool

	attrCacheMutex sync.RWMutex
	attrCache      map[string]*FileAttr
	attrCacheCond  *sync.Cond
	attrCacheBusy  map[string]bool

	// TODO - add counters and check that the rpcFs.fetchCond is
	// working.
}

var paranoia = false

func (me *FsServer) verify() {
	if !paranoia {
		return
	}
	me.attrCacheMutex.RLock()
	defer me.attrCacheMutex.RUnlock()

	for k, v := range me.attrCache {
		if v.Path != k {
			panic(fmt.Sprintf("attrCache mismatch %q %#v", k, v))
		}
		if _, ok := me.attrCacheBusy[k]; ok {
			panic(fmt.Sprintf("attrCacheBusy and attrCache entry for %q", k))
		}
	}
}

func NewFsServer(root string, cache *ContentCache, excluded []string) *FsServer {
	fs := &FsServer{
		contentCache:  cache,
		contentServer: &ContentServer{Cache: cache},
		Root:          root,
		hashCache:     make(map[string]string),
		hashBusyMap:   map[string]bool{},

		attrCache:      make(map[string]*FileAttr),
		attrCacheBusy:  map[string]bool{},
		excludePrivate: true,
	}

	fs.hashCacheCond = sync.NewCond(&fs.hashCacheMutex)
	fs.attrCacheCond = sync.NewCond(&fs.attrCacheMutex)
	fs.excluded = make(map[string]bool)
	for _, e := range excluded {
		fs.excluded[e] = true
	}
	return fs
}

func (me FileAttr) String() string {
	id := ""
	if me.Hash != "" {
		id = fmt.Sprintf(" sz %d", me.FileInfo.Size)
	}
	if me.Link != "" {
		id = fmt.Sprintf(" -> %s", me.Link)
	}
	if me.Deletion() {
		id = " (del)"
	} else if !me.Status.Ok() {
		id = " " + me.Status.String()
	}
	if me.Status.Ok() {
		id += fmt.Sprintf(" m=%o", me.FileInfo.Mode)
	}

	return fmt.Sprintf("%s%s", me.Path, id)
}

func (me FileAttr) Deletion() bool {
	return me.Status == fuse.ENOENT
}

func (me *FsServer) path(n string) string {
	if me.Root == "" {
		return n
	}
	return filepath.Join(me.Root, strings.TrimLeft(n, "/"))
}

func (me *FsServer) FileContent(req *ContentRequest, rep *ContentResponse) os.Error {
	return me.contentServer.FileContent(req, rep)
}

func (me *FsServer) ReadDir(req *DirRequest, r *DirResponse) os.Error {
	d, e := ioutil.ReadDir(me.path(req.Name))
	log.Println("ReadDir", req)
	r.NameModeMap = make(map[string]uint32)
	for _, v := range d {
		r.NameModeMap[v.Name] = v.Mode
	}
	r.Status = fuse.OsErrorToErrno(e)
	return nil
}

func (me *FsServer) GetAttr(req *AttrRequest, rep *AttrResponse) os.Error {
	log.Println("GetAttr req", req.Name)
	names := []string{}
	if me.multiplyPaths != nil {
		names = me.multiplyPaths(req.Name)
	} else {
		names = append(names, req.Name)
	}
	for _, n := range names {
		a := me.oneGetAttr(n)
		if a.Hash != "" {
			log.Printf("GetAttr %v %x", n, a, a.Hash)
		}
		rep.Attrs = append(rep.Attrs, a)
	}
	return nil
}

func (me *FsServer) oneGetAttr(name string) (rep *FileAttr) {
	defer me.verify()

	if me.excluded[name] {
		log.Printf("Denied access to excluded file %q", name)
		return &FileAttr{
			Path:   name,
			Status: fuse.ENOENT,
		}
	}

	me.attrCacheMutex.RLock()
	rep, ok := me.attrCache[name]
	me.attrCacheMutex.RUnlock()
	if ok {
		return rep
	}

	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()
	for me.attrCacheBusy[name] && me.attrCache[name] == nil {
		me.attrCacheCond.Wait()
	}
	rep, ok = me.attrCache[name]
	if ok {
		return rep
	}
	me.attrCacheBusy[name] = true
	me.attrCacheMutex.Unlock()

	p := me.path(name)
	fi, err := os.Lstat(p)
	rep = &FileAttr{
		FileInfo: fi,
		Status:   fuse.OsErrorToErrno(err),
		Path:     name,
	}

	// We don't want to expose the master's private files to the
	// world.
	if me.excludePrivate && fi != nil && fi.Mode&0077 == 0 {
		log.Printf("Denied access to private file %q", name)
		rep.FileInfo = nil
		rep.Status = fuse.EPERM
		fi = nil
	}

	if fi != nil {
		me.fillContent(rep)
	}

	me.attrCacheMutex.Lock()
	me.attrCache[name] = rep
	me.attrCacheCond.Broadcast()
	me.attrCacheBusy[name] = false, false
	return rep
}

func (me *FsServer) fillContent(rep *FileAttr) {
	if rep.FileInfo.IsSymlink() {
		rep.Link, _ = os.Readlink(rep.Path)
	}
	if rep.FileInfo.IsRegular() {
		rep.Hash = me.getHash(rep.Path)
		if rep.Hash == "" {
			// Typically happens if we want to open /etc/shadow as normal user.
			rep.Status = fuse.EPERM
		}
	}
}

func (me *FsServer) updateFiles(infos []*FileAttr) {
	me.updateHashes(infos)
	me.updateAttrs(infos)
}

func (me *FsServer) updateAttrs(infos []*FileAttr) {
	defer me.verify()

	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()

	for _, r := range infos {
		me.attrCache[r.Path] = r
	}
}

func (me *FsServer) updateHashes(infos []*FileAttr) {
	defer me.verify()

	me.hashCacheMutex.Lock()
	defer me.hashCacheMutex.Unlock()

	for _, r := range infos {
		name := r.Path
		if !r.Status.Ok() || r.Link != "" {
			me.hashCache[name] = "", false
			me.hashBusyMap[name] = false, false
		}
		if r.Hash != "" {
			me.hashCache[name] = r.Hash
			me.hashBusyMap[name] = false, false
		}
	}
}

func (me *FsServer) getHash(name string) (hash string) {
	fullPath := me.path(name)

	me.hashCacheMutex.RLock()
	hash = me.hashCache[name]
	me.hashCacheMutex.RUnlock()

	if hash != "" {
		return hash
	}

	me.hashCacheMutex.Lock()
	defer me.hashCacheMutex.Unlock()
	for me.hashBusyMap[name] && me.hashCache[name] == "" {
		me.hashCacheCond.Wait()
	}

	hash = me.hashCache[name]
	if hash != "" {
		return hash
	}
	me.hashBusyMap[name] = true
	me.hashCacheMutex.Unlock()

	// TODO - /usr should be configurable.
	if HasDirPrefix(fullPath, "/usr") && !HasDirPrefix(fullPath, "/usr/local") {
		hash = me.contentCache.SaveImmutablePath(fullPath)
	} else {
		hash = me.contentCache.SavePath(fullPath)
	}

	me.hashCacheMutex.Lock()
	me.hashCache[name] = hash
	me.hashBusyMap[name] = false, false
	me.hashCacheCond.Broadcast()

	return hash
}

func (me *FsServer) refreshAttributeCache(prefix string) []*FileAttr {
	me.attrCacheMutex.Lock()
	defer me.attrCacheMutex.Unlock()

	updated := []*FileAttr{}
	for key, attr := range me.attrCache {
		// TODO -should just do everything?
		if !HasDirPrefix(key, prefix) {
			continue
		}

		fi, _ := os.Lstat(me.path(key))
		if fi == nil && attr.Status.Ok() {
			del := FileAttr{
				Path:   key,
				Status: fuse.ENOENT,
			}
			updated = append(updated, &del)
		}
		if fi != nil && attr.FileInfo != nil && EncodeFileInfo(*attr.FileInfo) != EncodeFileInfo(*fi) {
			newEnt := FileAttr{
				Path:     key,
				Status:   fuse.OK,
				FileInfo: fi,
			}
			me.fillContent(&newEnt)
			updated = append(updated, &newEnt)
		}
	}

	for _, u := range updated {
		copy := *u
		me.attrCache[u.Path] = &copy
	}
	return updated
}

func (me *FsServer) copyCache() []*FileAttr {
	me.attrCacheMutex.RLock()
	defer me.attrCacheMutex.RUnlock()

	dump := []*FileAttr{}
	for _, attr := range me.attrCache {
		copy := *attr
		dump = append(dump, &copy)
	}

	return dump
}
<|endoftext|>


<commit_before>package terse

import (
	"strings"
	"testing"

	"github.com/acsellers/multitemplate"
	"github.com/acsellers/multitemplate/helpers"
)

func TestErrors(t *testing.T) {
	helpers.LoadHelpers("all")
	for _, test := range errorTests {
		tmpl := multitemplate.New("terse")
		var e error
		tmpl, e = tmpl.Parse("parse", test.Source, "terse")
		if e == nil {
			t.Error("No Error:", test.Name)
			continue
		} else if test.Contains != "" {
			if !strings.Contains(e.Error(), test.Contains) {
				t.Error("Incorrect error for", test.Name, "Expected:", test.Contains, "Was:", e)
			}
		}
	}
}

var errorTests = []errorTest{
	errorTest{
		Name:     "Unclosed attribute parentheses",
		Source:   "html name=(link_to",
		Contains: "Unclosed parentheses",
	},
	errorTest{
		Name:     "Empty attribute parentheses",
		Source:   "html name=()",
		Contains: "Empty parentheses",
	},
	errorTest{
		Name:     "Unclosed attribute single quotes ",
		Source:   "html name='link_to",
		Contains: "Unclosed single quotes",
	},
	errorTest{
		Name:     "Empty attribute single quotes",
		Source:   "html name=''",
		Contains: "Empty single quotes",
	},
	errorTest{
		Name:     "Unclosed attribute double quotes ",
		Source:   "html name=\"link_to",
		Contains: "Unclosed double quotes",
	},
	errorTest{
		Name:     "Empty attribute double quotes",
		Source:   "html name=\"\"",
		Contains: "Empty double quotes",
	},
	errorTest{
		Name:     "Blank attribute variable",
		Source:   "html name=",
		Contains: "Blank attribute",
	},
	errorTest{
		Name:     "Headless Totem Pole",
		Source:   "html > head >",
		Contains: "Missing tag",
	},
	errorTest{
		Name:   "Malformed code test",
		Source: "= link_to (",
	},
	errorTest{
		Name:     "Multiple errors test",
		Source:   "head >\n= link_to (",
		Contains: "Missing tag",
	},
	errorTest{
		Name:   "Bad function name for if",
		Source: "?asdfasdfsafd\n  no",
	},
	errorTest{
		Name:   "Bad function name for if",
		Source: "&asdfasdfsafd\n  no",
	},
	errorTest{
		Name:   "Bad interpolated code 1",
		Source: "asdf {{ asdfasdf }}",
	},
	errorTest{
		Name:   "Bad interpolated code 2",
		Source: "asdf\n  {{ asdfasdf }}",
	},
	errorTest{
		Name:   "Bad interpolated code 3",
		Source: "asdf\n  {{ {}{}} }}",
	},
	errorTest{
		Name:   "Bad interpolated code 3",
		Source: "asdf}} {{  }}",
	},
	errorTest{
		Name:   "Weird if statement",
		Source: "?!?",
	},
	errorTest{
		Name:     "Weird if statement",
		Source:   "html(",
		Contains: "Missing nested attributes",
	},
}

type errorTest struct {
	Name     string
	Source   string
	Contains string
}
<commit_msg>Add new error test for bad range statements<commit_after>package terse

import (
	"strings"
	"testing"

	"github.com/acsellers/multitemplate"
	"github.com/acsellers/multitemplate/helpers"
)

func TestErrors(t *testing.T) {
	helpers.LoadHelpers("all")
	for _, test := range errorTests {
		tmpl := multitemplate.New("terse")
		var e error
		tmpl, e = tmpl.Parse("parse", test.Source, "terse")
		if e == nil {
			t.Error("No Error:", test.Name)
			continue
		} else if test.Contains != "" {
			if !strings.Contains(e.Error(), test.Contains) {
				t.Error("Incorrect error for", test.Name, "Expected:", test.Contains, "Was:", e)
			}
		}
	}
}

var errorTests = []errorTest{
	errorTest{
		Name:     "Unclosed attribute parentheses",
		Source:   "html name=(link_to",
		Contains: "Unclosed parentheses",
	},
	errorTest{
		Name:     "Empty attribute parentheses",
		Source:   "html name=()",
		Contains: "Empty parentheses",
	},
	errorTest{
		Name:     "Unclosed attribute single quotes ",
		Source:   "html name='link_to",
		Contains: "Unclosed single quotes",
	},
	errorTest{
		Name:     "Empty attribute single quotes",
		Source:   "html name=''",
		Contains: "Empty single quotes",
	},
	errorTest{
		Name:     "Unclosed attribute double quotes ",
		Source:   "html name=\"link_to",
		Contains: "Unclosed double quotes",
	},
	errorTest{
		Name:     "Empty attribute double quotes",
		Source:   "html name=\"\"",
		Contains: "Empty double quotes",
	},
	errorTest{
		Name:     "Blank attribute variable",
		Source:   "html name=",
		Contains: "Blank attribute",
	},
	errorTest{
		Name:     "Headless Totem Pole",
		Source:   "html > head >",
		Contains: "Missing tag",
	},
	errorTest{
		Name:   "Malformed code test",
		Source: "= link_to (",
	},
	errorTest{
		Name:     "Multiple errors test",
		Source:   "head >\n= link_to (",
		Contains: "Missing tag",
	},
	errorTest{
		Name:   "Bad function name for if",
		Source: "?asdfasdfsafd\n  no",
	},
	errorTest{
		Name:   "Bad function name for if",
		Source: "&asdfasdfsafd\n  no",
	},
	errorTest{
		Name:   "Bad interpolated code 1",
		Source: "asdf {{ asdfasdf }}",
	},
	errorTest{
		Name:   "Bad interpolated code 2",
		Source: "asdf\n  {{ asdfasdf }}",
	},
	errorTest{
		Name:   "Bad interpolated code 3",
		Source: "asdf\n  {{ {}{}} }}",
	},
	errorTest{
		Name:   "Bad interpolated code 3",
		Source: "asdf}} {{  }}",
	},
	errorTest{
		Name:   "Weird if statement",
		Source: "?!?",
	},
	errorTest{
		Name:     "No attribute if statement",
		Source:   "html(",
		Contains: "Missing nested attributes",
	},
	errorTest{
		Name:   "Bad range statement",
		Source: "&nothere\n  = .",
	},
}

type errorTest struct {
	Name     string
	Source   string
	Contains string
}
<|endoftext|>


<commit_before>package setup

import (
	"context"
	"fmt"
	"log"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/itchio/savior"

	"github.com/itchio/httpkit/progress"

	"github.com/itchio/wharf/eos"
	"github.com/itchio/wharf/eos/option"
	"github.com/itchio/wharf/pwr"

	"github.com/itchio/wharf/pools/fspool"
	"github.com/itchio/wharf/pwr/bowl"

	"github.com/itchio/go-itchio"
	"github.com/itchio/savior/filesource"
	"github.com/itchio/savior/zipextractor"

	"github.com/itchio/wharf/pwr/patcher"
	"github.com/itchio/wharf/taskgroup"
	"github.com/pkg/errors"
)

type localState struct {
	appDir  string
	version string
}

type remoteState struct {
	version string
}

type patchPlan struct {
	path      *BrothUpgradePath
	totalSize int64
}

type archivePlan struct {
	totalSize int64
}

func (i *Installer) Upgrade(mv Multiverse) error {
	EnableJSON()
	defer DisableJSON()

	var ls *localState
	var rs *remoteState

	ctx := context.Background()
	err := taskgroup.Do(ctx,
		// check latest version
		func() error {
			latestVersion, err := i.brothGetString("/LATEST")
			if err != nil {
				return err
			}
			rs = &remoteState{
				version: latestVersion,
			}
			return nil
		},

		// check local version
		func() error {
			currentBuild := mv.GetCurrentVersion()
			if currentBuild == nil {
				return errors.Errorf("No version currently installed")
			}

			ls = &localState{
				appDir:  currentBuild.Path,
				version: currentBuild.Version,
			}
			return nil
		},
	)
	if err != nil {
		return err
	}

	log.Printf("Installed %s", ls.version)
	log.Printf("Latest    %s", rs.version)

	if ls.version == rs.version {
		log.Printf("We're up-to-date!")
		Emit(NoUpdateAvailable{})
		return nil
	}

	if mv.HasReadyPending() {
		log.Printf("Current is behind, but we have a ready version...")
		if mv.ReadyPendingIs(rs.version) {
			log.Printf("...and it is the latest! (%s)", rs.version)
		}
		Emit(UpdateReady{Version: rs.version})
		return nil
	}

	var pp *patchPlan
	var ap *archivePlan

	err = taskgroup.Do(ctx,
		// try to find patch plan
		func() error {
			upgradePath := &BrothUpgradePath{}
			err = i.brothGetResponse(upgradePath, "/%s/upgrade-paths/%s",
				ls.version,
				rs.version,
			)
			if err != nil {
				log.Printf("While looking for upgrade path: %v", err)
				log.Printf("Giving up patch plan")
				return nil
			}

			vnames := []string{ls.version}
			for _, bp := range upgradePath.Patches {
				vnames = append(vnames, bp.Version)
			}
			log.Printf("Upgrade path: %s",
				strings.Join(vnames, " → "),
			)

			var totalSize int64
			for _, bp := range upgradePath.Patches {
				f := bp.FindSubType(itchio.BuildFileSubTypeDefault)
				if f == nil {
					log.Printf("Missing patch for version %s, giving up patch plan", bp.Version)
					return nil
				}

				of := bp.FindSubType(itchio.BuildFileSubTypeOptimized)
				if of != nil && of.Size < f.Size {
					f = of
				}

				totalSize += f.Size
			}
			pp = &patchPlan{
				path:      upgradePath,
				totalSize: totalSize,
			}

			return nil
		},

		// try to find archive plan
		func() error {
			buildInfo := &BrothBuildInfo{}
			err = i.brothGetResponse(buildInfo, "/%s/info", rs.version)
			if err != nil {
				return errors.WithMessage(err, "While looking for archive plan")
			}

			found := false
			for _, f := range buildInfo.Files {
				if f.Type == itchio.BuildFileTypeArchive &&
					f.SubType == itchio.BuildFileSubTypeDefault {

					ap = &archivePlan{
						totalSize: f.Size,
					}
					found = true
					break
				}
			}

			if !found {
				errMsg := fmt.Sprintf("Default archive not found for version %s", rs.version)
				return errors.WithMessage(err, errMsg)
			}

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

	if pp == nil {
		log.Printf("↑ No patch-based upgrade path found")
	} else {
		log.Printf("↑ Patching cost: %s (in %d patches)",
			progress.FormatBytes(pp.totalSize),
			len(pp.path.Patches),
		)
	}
	log.Printf("↺ Archive  cost: %s",
		progress.FormatBytes(ap.totalSize),
	)

	if pp != nil && pp.totalSize < ap.totalSize {
		err = i.applyPatches(mv, ls, pp)
		if err == nil {
			log.Printf("Patching went fine!")
			return nil
		}

		log.Printf("Patching went wrong, falling back to archive.")
		log.Printf("The patching error was: %+v", err)
	}

	err = i.applyArchive(mv, rs, ap)
	if err != nil {
		Emit(UpdateFailed{Message: fmt.Sprintf("%+v", err)})
		return err
	}

	Emit(UpdateReady{Version: rs.version})
	return nil
}

func (i *Installer) applyPatches(mv Multiverse, ls *localState, pp *patchPlan) error {
	up := pp.path
	log.Printf("Applying %d patches...", len(up.Patches))

	{
		log.Printf("But first, let's check (%s) is a valid build for (%s)", ls.appDir, ls.version)

		signatureURL := fmt.Sprintf("%s/%s/signature/default", i.brothPackageURL(), ls.version)
		log.Printf("☁ %s", signatureURL)

		consumer := newConsumer()
		sigSource, err := filesource.Open(signatureURL, option.WithConsumer(consumer))
		if err != nil {
			return err
		}
		defer sigSource.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		sigInfo, err := pwr.ReadSignature(ctx, sigSource)
		if err != nil {
			return err
		}

		vc := pwr.ValidatorContext{
			Consumer:   consumer,
			FailFast:   true,
			NumWorkers: 2,
		}
		err = vc.Validate(ctx, ls.appDir, sigInfo)
		if err != nil {
			return err
		}
	}

	stagingDir, err := mv.MakeStagingFolder()
	if err != nil {
		return err
	}
	defer mv.CleanStagingFolder()
	log.Printf("Using (%s) as staging directory", stagingDir)

	applyOne := func(bp *BrothPatch, targetDir string, outputDir string) error {
		log.Printf("Upgrading to %s...", bp.Version)
		Emit(InstallingUpdate{Version: bp.Version})

		f := bp.FindSubType(itchio.BuildFileSubTypeDefault)
		if f == nil {
			return errors.Errorf("Could not find default patch file for version %s, giving up", bp.Version)
		}

		of := bp.FindSubType(itchio.BuildFileSubTypeOptimized)
		if of != nil && of.Size < f.Size {
			f = of
		}
		log.Printf("Using (%s) patch (%s)", f.SubType, progress.FormatBytes(f.Size))

		consumer := newConsumer()

		patchURL := fmt.Sprintf("%s/%s/patch/%s", i.brothPackageURL(), bp.Version, f.SubType)
		log.Printf("☁ %s", patchURL)

		patchSource, err := filesource.Open(patchURL, option.WithConsumer(consumer))
		if err != nil {
			return err
		}
		defer patchSource.Close()

		tracker := progress.NewTracker()
		tracker.SetSilent(true)
		tracker.SetTotalBytes(patchSource.Size())
		tracker.Start()
		defer tracker.Finish()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		p, err := patcher.New(patchSource, consumer)
		if err != nil {
			return err
		}
		startPrintingProgress(ctx, tracker)

		targetPool := fspool.New(p.GetTargetContainer(), targetDir)

		bwl, err := bowl.NewFreshBowl(&bowl.FreshBowlParams{
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			SourceContainer: p.GetSourceContainer(),
			OutputFolder:    outputDir,
		})
		if err != nil {
			return err
		}

		err = p.Resume(nil, targetPool, bwl)
		if err != nil {
			return err
		}

		err = bwl.Commit()
		if err != nil {
			return err
		}

		return nil
	}

	targetDir := ls.appDir
	var outputDir string
	var latestVersion string
	for _, p := range up.Patches {
		outputDir = filepath.Join(stagingDir, fmt.Sprintf("app-%s", p.Version))
		err := applyOne(p, targetDir, outputDir)
		if err != nil {
			return err
		}
		targetDir = outputDir
		latestVersion = p.Version
	}

	log.Printf("Fully upgraded into (%s)", outputDir)
	err = mv.QueueReady(&BuildFolder{
		Version: latestVersion,
		Path:    outputDir,
	})
	if err != nil {
		return err
	}

	return nil
}

func (i *Installer) applyArchive(mv Multiverse, rs *remoteState, ap *archivePlan) error {
	log.Printf("Upgrading to (%s) using archive...", rs.version)
	Emit(InstallingUpdate{Version: rs.version})

	archiveURL := fmt.Sprintf("%s/%s/archive/default", i.brothPackageURL(), rs.version)
	log.Printf("☁ %s", archiveURL)

	consumer := newConsumer()

	archiveFile, err := eos.Open(archiveURL, option.WithConsumer(i.consumer))
	if err != nil {
		return err
	}

	archiveStats, err := archiveFile.Stat()
	if err != nil {
		return err
	}

	ex, err := zipextractor.New(archiveFile, archiveStats.Size())
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tracker := progress.NewTracker()
	tracker.SetSilent(true)
	tracker.SetTotalBytes(archiveStats.Size())
	tracker.Start()
	defer tracker.Finish()

	consumer.OnProgress = tracker.SetProgress
	ex.SetConsumer(consumer)
	startPrintingProgress(ctx, tracker)

	stagingFolder, err := mv.MakeStagingFolder()
	if err != nil {
		return err
	}
	defer mv.CleanStagingFolder()

	outputDir := filepath.Join(stagingFolder, fmt.Sprintf("app-%s", rs.version))
	log.Printf("Extracting %s to (%s)", progress.FormatBytes(archiveStats.Size()), outputDir)

	sink := &savior.FolderSink{
		Consumer:  consumer,
		Directory: outputDir,
	}
	var closeSinkOnce sync.Once
	defer closeSinkOnce.Do(func() {
		sink.Close()
	})

	startTime := time.Now()

	res, err := ex.Resume(nil, sink)
	if err != nil {
		return err
	}

	duration := time.Since(startTime)

	log.Printf("Overall extract speed: %s (%s total)",
		progress.FormatBPS(res.Size(), duration),
		progress.FormatDuration(duration),
	)
	closeSinkOnce.Do(func() {
		sink.Close()
	})

	err = mv.QueueReady(&BuildFolder{
		Version: rs.version,
		Path:    outputDir,
	})
	if err != nil {
		return err
	}

	return nil
}

func startPrintingProgress(ctx context.Context, tracker *progress.Tracker) {
	go func() {
		for {
			select {
			case <-time.After(1 * time.Second):
				Emit(Progress{
					Progress: tracker.Progress(),
					ETA:      tracker.ETA().Seconds(),
					BPS:      tracker.BPS(),
				})
				log.Printf("%.2f%% done - %s / s, ETA %s",
					tracker.Progress()*100,
					progress.FormatBytes(int64(tracker.BPS())),
					tracker.ETA(),
				)
			case <-ctx.Done():
				return
			}
		}
	}()
}
<commit_msg>Emit update-ready after patching, closes https://github.com/itchio/itch/issues/1831<commit_after>package setup

import (
	"context"
	"fmt"
	"log"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/itchio/savior"

	"github.com/itchio/httpkit/progress"

	"github.com/itchio/wharf/eos"
	"github.com/itchio/wharf/eos/option"
	"github.com/itchio/wharf/pwr"

	"github.com/itchio/wharf/pools/fspool"
	"github.com/itchio/wharf/pwr/bowl"

	"github.com/itchio/go-itchio"
	"github.com/itchio/savior/filesource"
	"github.com/itchio/savior/zipextractor"

	"github.com/itchio/wharf/pwr/patcher"
	"github.com/itchio/wharf/taskgroup"
	"github.com/pkg/errors"
)

type localState struct {
	appDir  string
	version string
}

type remoteState struct {
	version string
}

type patchPlan struct {
	path      *BrothUpgradePath
	totalSize int64
}

type archivePlan struct {
	totalSize int64
}

func (i *Installer) Upgrade(mv Multiverse) error {
	EnableJSON()
	defer DisableJSON()

	var ls *localState
	var rs *remoteState

	ctx := context.Background()
	err := taskgroup.Do(ctx,
		// check latest version
		func() error {
			latestVersion, err := i.brothGetString("/LATEST")
			if err != nil {
				return err
			}
			rs = &remoteState{
				version: latestVersion,
			}
			return nil
		},

		// check local version
		func() error {
			currentBuild := mv.GetCurrentVersion()
			if currentBuild == nil {
				return errors.Errorf("No version currently installed")
			}

			ls = &localState{
				appDir:  currentBuild.Path,
				version: currentBuild.Version,
			}
			return nil
		},
	)
	if err != nil {
		return err
	}

	log.Printf("Installed %s", ls.version)
	log.Printf("Latest    %s", rs.version)

	if ls.version == rs.version {
		log.Printf("We're up-to-date!")
		Emit(NoUpdateAvailable{})
		return nil
	}

	if mv.HasReadyPending() {
		log.Printf("Current is behind, but we have a ready version...")
		if mv.ReadyPendingIs(rs.version) {
			log.Printf("...and it is the latest! (%s)", rs.version)
		}
		Emit(UpdateReady{Version: rs.version})
		return nil
	}

	var pp *patchPlan
	var ap *archivePlan

	err = taskgroup.Do(ctx,
		// try to find patch plan
		func() error {
			upgradePath := &BrothUpgradePath{}
			err = i.brothGetResponse(upgradePath, "/%s/upgrade-paths/%s",
				ls.version,
				rs.version,
			)
			if err != nil {
				log.Printf("While looking for upgrade path: %v", err)
				log.Printf("Giving up patch plan")
				return nil
			}

			vnames := []string{ls.version}
			for _, bp := range upgradePath.Patches {
				vnames = append(vnames, bp.Version)
			}
			log.Printf("Upgrade path: %s",
				strings.Join(vnames, " → "),
			)

			var totalSize int64
			for _, bp := range upgradePath.Patches {
				f := bp.FindSubType(itchio.BuildFileSubTypeDefault)
				if f == nil {
					log.Printf("Missing patch for version %s, giving up patch plan", bp.Version)
					return nil
				}

				of := bp.FindSubType(itchio.BuildFileSubTypeOptimized)
				if of != nil && of.Size < f.Size {
					f = of
				}

				totalSize += f.Size
			}
			pp = &patchPlan{
				path:      upgradePath,
				totalSize: totalSize,
			}

			return nil
		},

		// try to find archive plan
		func() error {
			buildInfo := &BrothBuildInfo{}
			err = i.brothGetResponse(buildInfo, "/%s/info", rs.version)
			if err != nil {
				return errors.WithMessage(err, "While looking for archive plan")
			}

			found := false
			for _, f := range buildInfo.Files {
				if f.Type == itchio.BuildFileTypeArchive &&
					f.SubType == itchio.BuildFileSubTypeDefault {

					ap = &archivePlan{
						totalSize: f.Size,
					}
					found = true
					break
				}
			}

			if !found {
				errMsg := fmt.Sprintf("Default archive not found for version %s", rs.version)
				return errors.WithMessage(err, errMsg)
			}

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

	if pp == nil {
		log.Printf("↑ No patch-based upgrade path found")
	} else {
		log.Printf("↑ Patching cost: %s (in %d patches)",
			progress.FormatBytes(pp.totalSize),
			len(pp.path.Patches),
		)
	}
	log.Printf("↺ Archive  cost: %s",
		progress.FormatBytes(ap.totalSize),
	)

	if pp != nil && pp.totalSize < ap.totalSize {
		err = i.applyPatches(mv, ls, pp)
		if err == nil {
			log.Printf("Patching went fine!")
			Emit(UpdateReady{Version: rs.version})
			return nil
		}

		log.Printf("Patching went wrong, falling back to archive.")
		log.Printf("The patching error was: %+v", err)
	}

	err = i.applyArchive(mv, rs, ap)
	if err != nil {
		Emit(UpdateFailed{Message: fmt.Sprintf("%+v", err)})
		return err
	}

	Emit(UpdateReady{Version: rs.version})
	return nil
}

func (i *Installer) applyPatches(mv Multiverse, ls *localState, pp *patchPlan) error {
	up := pp.path
	log.Printf("Applying %d patches...", len(up.Patches))

	{
		log.Printf("But first, let's check (%s) is a valid build for (%s)", ls.appDir, ls.version)

		signatureURL := fmt.Sprintf("%s/%s/signature/default", i.brothPackageURL(), ls.version)
		log.Printf("☁ %s", signatureURL)

		consumer := newConsumer()
		sigSource, err := filesource.Open(signatureURL, option.WithConsumer(consumer))
		if err != nil {
			return err
		}
		defer sigSource.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		sigInfo, err := pwr.ReadSignature(ctx, sigSource)
		if err != nil {
			return err
		}

		vc := pwr.ValidatorContext{
			Consumer:   consumer,
			FailFast:   true,
			NumWorkers: 2,
		}
		err = vc.Validate(ctx, ls.appDir, sigInfo)
		if err != nil {
			return err
		}
	}

	stagingDir, err := mv.MakeStagingFolder()
	if err != nil {
		return err
	}
	defer mv.CleanStagingFolder()
	log.Printf("Using (%s) as staging directory", stagingDir)

	applyOne := func(bp *BrothPatch, targetDir string, outputDir string) error {
		log.Printf("Upgrading to %s...", bp.Version)
		Emit(InstallingUpdate{Version: bp.Version})

		f := bp.FindSubType(itchio.BuildFileSubTypeDefault)
		if f == nil {
			return errors.Errorf("Could not find default patch file for version %s, giving up", bp.Version)
		}

		of := bp.FindSubType(itchio.BuildFileSubTypeOptimized)
		if of != nil && of.Size < f.Size {
			f = of
		}
		log.Printf("Using (%s) patch (%s)", f.SubType, progress.FormatBytes(f.Size))

		consumer := newConsumer()

		patchURL := fmt.Sprintf("%s/%s/patch/%s", i.brothPackageURL(), bp.Version, f.SubType)
		log.Printf("☁ %s", patchURL)

		patchSource, err := filesource.Open(patchURL, option.WithConsumer(consumer))
		if err != nil {
			return err
		}
		defer patchSource.Close()

		tracker := progress.NewTracker()
		tracker.SetSilent(true)
		tracker.SetTotalBytes(patchSource.Size())
		tracker.Start()
		defer tracker.Finish()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		p, err := patcher.New(patchSource, consumer)
		if err != nil {
			return err
		}
		startPrintingProgress(ctx, tracker)

		targetPool := fspool.New(p.GetTargetContainer(), targetDir)

		bwl, err := bowl.NewFreshBowl(&bowl.FreshBowlParams{
			TargetContainer: p.GetTargetContainer(),
			TargetPool:      targetPool,
			SourceContainer: p.GetSourceContainer(),
			OutputFolder:    outputDir,
		})
		if err != nil {
			return err
		}

		err = p.Resume(nil, targetPool, bwl)
		if err != nil {
			return err
		}

		err = bwl.Commit()
		if err != nil {
			return err
		}

		return nil
	}

	targetDir := ls.appDir
	var outputDir string
	var latestVersion string
	for _, p := range up.Patches {
		outputDir = filepath.Join(stagingDir, fmt.Sprintf("app-%s", p.Version))
		err := applyOne(p, targetDir, outputDir)
		if err != nil {
			return err
		}
		targetDir = outputDir
		latestVersion = p.Version
	}

	log.Printf("Fully upgraded into (%s)", outputDir)
	err = mv.QueueReady(&BuildFolder{
		Version: latestVersion,
		Path:    outputDir,
	})
	if err != nil {
		return err
	}

	return nil
}

func (i *Installer) applyArchive(mv Multiverse, rs *remoteState, ap *archivePlan) error {
	log.Printf("Upgrading to (%s) using archive...", rs.version)
	Emit(InstallingUpdate{Version: rs.version})

	archiveURL := fmt.Sprintf("%s/%s/archive/default", i.brothPackageURL(), rs.version)
	log.Printf("☁ %s", archiveURL)

	consumer := newConsumer()

	archiveFile, err := eos.Open(archiveURL, option.WithConsumer(i.consumer))
	if err != nil {
		return err
	}

	archiveStats, err := archiveFile.Stat()
	if err != nil {
		return err
	}

	ex, err := zipextractor.New(archiveFile, archiveStats.Size())
	if err != nil {
		return err
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tracker := progress.NewTracker()
	tracker.SetSilent(true)
	tracker.SetTotalBytes(archiveStats.Size())
	tracker.Start()
	defer tracker.Finish()

	consumer.OnProgress = tracker.SetProgress
	ex.SetConsumer(consumer)
	startPrintingProgress(ctx, tracker)

	stagingFolder, err := mv.MakeStagingFolder()
	if err != nil {
		return err
	}
	defer mv.CleanStagingFolder()

	outputDir := filepath.Join(stagingFolder, fmt.Sprintf("app-%s", rs.version))
	log.Printf("Extracting %s to (%s)", progress.FormatBytes(archiveStats.Size()), outputDir)

	sink := &savior.FolderSink{
		Consumer:  consumer,
		Directory: outputDir,
	}
	var closeSinkOnce sync.Once
	defer closeSinkOnce.Do(func() {
		sink.Close()
	})

	startTime := time.Now()

	res, err := ex.Resume(nil, sink)
	if err != nil {
		return err
	}

	duration := time.Since(startTime)

	log.Printf("Overall extract speed: %s (%s total)",
		progress.FormatBPS(res.Size(), duration),
		progress.FormatDuration(duration),
	)
	closeSinkOnce.Do(func() {
		sink.Close()
	})

	err = mv.QueueReady(&BuildFolder{
		Version: rs.version,
		Path:    outputDir,
	})
	if err != nil {
		return err
	}

	return nil
}

func startPrintingProgress(ctx context.Context, tracker *progress.Tracker) {
	go func() {
		for {
			select {
			case <-time.After(1 * time.Second):
				Emit(Progress{
					Progress: tracker.Progress(),
					ETA:      tracker.ETA().Seconds(),
					BPS:      tracker.BPS(),
				})
				log.Printf("%.2f%% done - %s / s, ETA %s",
					tracker.Progress()*100,
					progress.FormatBytes(int64(tracker.BPS())),
					tracker.ETA(),
				)
			case <-ctx.Done():
				return
			}
		}
	}()
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 logrus

import (
	"bytes"
	"fmt"
	"testing"

	lg "github.com/Sirupsen/logrus"

	"github.com/onsi/gomega"
)

func TestEntryWithError(t *testing.T) {

	gomega.RegisterTestingT(t)

	defer func() {
		ErrorKey = "error"
	}()

	err := fmt.Errorf("kaboom at layer %d", 4711)

	gomega.Expect(err).To(gomega.BeEquivalentTo(WithError(err).ent.Data["error"]))

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)

	gomega.Expect(err).To(gomega.BeEquivalentTo(entry.WithError(err).ent.Data["error"]))

	ErrorKey = "err"

	gomega.Expect(err).To(gomega.BeEquivalentTo(entry.WithError(err).ent.Data["err"]))

}

func TestEntryPanicln(t *testing.T) {
	gomega.RegisterTestingT(t)

	errBoom := fmt.Errorf("boom time")

	defer func() {
		p := recover()
		gomega.Expect(p).NotTo(gomega.BeNil())

		switch pVal := p.(type) {
		case *lg.Entry:
			gomega.Expect("kaboom").To(gomega.BeEquivalentTo(pVal.Message))
			gomega.Expect(errBoom).To(gomega.BeEquivalentTo(pVal.Data["err"]))
		default:
			t.Fatalf("want type *Entry, got %T: %#v", pVal, pVal)
		}
	}()

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)
	entry.withField("err", errBoom).Panicln("kaboom")
}

func TestEntryPanicf(t *testing.T) {
	errBoom := fmt.Errorf("boom again")

	defer func() {
		p := recover()
		gomega.Expect(p).NotTo(gomega.BeNil())

		switch pVal := p.(type) {
		case *lg.Entry:
			gomega.Expect("kaboom true").To(gomega.BeEquivalentTo(pVal.Message))
			gomega.Expect(errBoom).To(gomega.BeEquivalentTo(pVal.Data["err"]))
		default:
			t.Fatalf("want type *Entry, got %T: %#v", pVal, pVal)
		}
	}()

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)
	entry.withField("err", errBoom).Panicf("kaboom %v", true)
}
<commit_msg>Fix logger test<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 logrus

import (
	"bytes"
	"fmt"
	"testing"

	lg "github.com/Sirupsen/logrus"

	"github.com/onsi/gomega"
)

func TestEntryWithError(t *testing.T) {

	gomega.RegisterTestingT(t)

	defer func() {
		ErrorKey = "error"
	}()

	err := fmt.Errorf("kaboom at layer %d", 4711)

	gomega.Expect(err).To(gomega.BeEquivalentTo(WithError(err).Entry.Data["error"]))

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)

	gomega.Expect(err).To(gomega.BeEquivalentTo(entry.WithError(err).Entry.Data["error"]))

	ErrorKey = "err"

	gomega.Expect(err).To(gomega.BeEquivalentTo(entry.WithError(err).Entry.Data["err"]))

}

func TestEntryPanicln(t *testing.T) {
	gomega.RegisterTestingT(t)

	errBoom := fmt.Errorf("boom time")

	defer func() {
		p := recover()
		gomega.Expect(p).NotTo(gomega.BeNil())

		switch pVal := p.(type) {
		case *lg.Entry:
			gomega.Expect("kaboom").To(gomega.BeEquivalentTo(pVal.Message))
			gomega.Expect(errBoom).To(gomega.BeEquivalentTo(pVal.Data["err"]))
		default:
			t.Fatalf("want type *Entry, got %T: %#v", pVal, pVal)
		}
	}()

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)
	entry.withField("err", errBoom).Panicln("kaboom")
}

func TestEntryPanicf(t *testing.T) {
	errBoom := fmt.Errorf("boom again")

	defer func() {
		p := recover()
		gomega.Expect(p).NotTo(gomega.BeNil())

		switch pVal := p.(type) {
		case *lg.Entry:
			gomega.Expect("kaboom true").To(gomega.BeEquivalentTo(pVal.Message))
			gomega.Expect(errBoom).To(gomega.BeEquivalentTo(pVal.Data["err"]))
		default:
			t.Fatalf("want type *Entry, got %T: %#v", pVal, pVal)
		}
	}()

	logger := New()
	logger.std.Out = &bytes.Buffer{}
	entry := NewEntry(logger)
	entry.withField("err", errBoom).Panicf("kaboom %v", true)
}
<|endoftext|>


<commit_before>package stats

import (
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"
)

type Stats struct {
	Lock                sync.RWMutex
	Uptime              time.Time
	Pid                 int
	ResponseCounts      map[string]int
	TotalResponseCounts map[string]int
	TotalResponseTime   time.Time
}

func New() *Stats {
	stats := &Stats{
		Uptime:              time.Now(),
		Pid:                 os.Getpid(),
		ResponseCounts:      map[string]int{},
		TotalResponseCounts: map[string]int{},
		TotalResponseTime:   time.Time{},
	}

	go func() {
		for {
			stats.ResetResponseCounts()

			time.Sleep(time.Second * 1)
		}
	}()

	return stats
}

func (mw *Stats) ResetResponseCounts() {
	mw.Lock.Lock()
	defer mw.Lock.Unlock()
	mw.ResponseCounts = map[string]int{}
}

type recorderResponseWriter struct {
	http.ResponseWriter
	StatusCode int
}

func (w *recorderResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.WriteHeader(code)
	w.StatusCode = code
}

// MiddlewareFunc makes Stats implement the Middleware interface.
func (mw *Stats) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		beginning, recorder := mw.Begin(w)

		h.ServeHTTP(recorder, r)

		mw.End(beginning, recorder)
	})
}

// Negroni compatible interface
func (mw *Stats) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	beginning, recorder := mw.Begin(w)

	next(recorder, r)

	mw.End(beginning, recorder)
}

func (mw *Stats) Begin(w http.ResponseWriter) (time.Time, *recorderResponseWriter) {
	start := time.Now()

	writer := &recorderResponseWriter{w, 200}

	return start, writer
}

func (mw *Stats) End(start time.Time, writer *recorderResponseWriter) {
	end := time.Now()

	responseTime := end.Sub(start)

	mw.Lock.Lock()

	defer mw.Lock.Unlock()

	statusCode := fmt.Sprintf("%d", writer.StatusCode)

	mw.ResponseCounts[statusCode]++
	mw.TotalResponseCounts[statusCode]++
	mw.TotalResponseTime = mw.TotalResponseTime.Add(responseTime)
}

type data struct {
	Pid                    int            `json:"pid"`
	UpTime                 string         `json:"uptime"`
	UpTimeSec              float64        `json:"uptime_sec"`
	Time                   string         `json:"time"`
	TimeUnix               int64          `json:"unixtime"`
	StatusCodeCount        map[string]int `json:"status_code_count"`
	TotalStatusCodeCount   map[string]int `json:"total_status_code_count"`
	Count                  int            `json:"count"`
	TotalCount             int            `json:"total_count"`
	TotalResponseTime      string         `json:"total_response_time`
	TotalResponseTimeSec   float64        `json:"total_response_time_sec"`
	AverageResponseTime    string         `json:"average_response_time"`
	AverageResponseTimeSec float64        `json:"average_response_time_sec"`
}

func (mw *Stats) Data() *data {

	mw.Lock.RLock()

	now := time.Now()

	uptime := now.Sub(mw.Uptime)

	count := 0
	for _, current := range mw.ResponseCounts {
		count += current
	}

	totalCount := 0
	for _, count := range mw.TotalResponseCounts {
		totalCount += count
	}

	totalResponseTime := mw.TotalResponseTime.Sub(time.Time{})

	averageResponseTime := time.Duration(0)
	if totalCount > 0 {
		avgNs := int64(totalResponseTime) / int64(totalCount)
		averageResponseTime = time.Duration(avgNs)
	}

	r := &data{
		Pid:                    mw.Pid,
		UpTime:                 uptime.String(),
		UpTimeSec:              uptime.Seconds(),
		Time:                   now.String(),
		TimeUnix:               now.Unix(),
		StatusCodeCount:        mw.ResponseCounts,
		TotalStatusCodeCount:   mw.TotalResponseCounts,
		Count:                  count,
		TotalCount:             totalCount,
		TotalResponseTime:      totalResponseTime.String(),
		TotalResponseTimeSec:   totalResponseTime.Seconds(),
		AverageResponseTime:    averageResponseTime.String(),
		AverageResponseTimeSec: averageResponseTime.Seconds(),
	}

	mw.Lock.RUnlock()

	return r
}
<commit_msg>Do not export Lock<commit_after>package stats

import (
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"
)

type Stats struct {
	mu                sync.RWMutex
	Uptime              time.Time
	Pid                 int
	ResponseCounts      map[string]int
	TotalResponseCounts map[string]int
	TotalResponseTime   time.Time
}

func New() *Stats {
	stats := &Stats{
		Uptime:              time.Now(),
		Pid:                 os.Getpid(),
		ResponseCounts:      map[string]int{},
		TotalResponseCounts: map[string]int{},
		TotalResponseTime:   time.Time{},
	}

	go func() {
		for {
			stats.ResetResponseCounts()

			time.Sleep(time.Second * 1)
		}
	}()

	return stats
}

func (mw *Stats) ResetResponseCounts() {
	mw.mu.Lock()
	defer mw.mu.Unlock()
	mw.ResponseCounts = map[string]int{}
}

type recorderResponseWriter struct {
	http.ResponseWriter
	StatusCode int
}

func (w *recorderResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.WriteHeader(code)
	w.StatusCode = code
}

// MiddlewareFunc makes Stats implement the Middleware interface.
func (mw *Stats) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		beginning, recorder := mw.Begin(w)

		h.ServeHTTP(recorder, r)

		mw.End(beginning, recorder)
	})
}

// Negroni compatible interface
func (mw *Stats) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	beginning, recorder := mw.Begin(w)

	next(recorder, r)

	mw.End(beginning, recorder)
}

func (mw *Stats) Begin(w http.ResponseWriter) (time.Time, *recorderResponseWriter) {
	start := time.Now()

	writer := &recorderResponseWriter{w, 200}

	return start, writer
}

func (mw *Stats) End(start time.Time, writer *recorderResponseWriter) {
	end := time.Now()

	responseTime := end.Sub(start)

	mw.mu.Lock()

	defer mw.mu.Unlock()

	statusCode := fmt.Sprintf("%d", writer.StatusCode)

	mw.ResponseCounts[statusCode]++
	mw.TotalResponseCounts[statusCode]++
	mw.TotalResponseTime = mw.TotalResponseTime.Add(responseTime)
}

type data struct {
	Pid                    int            `json:"pid"`
	UpTime                 string         `json:"uptime"`
	UpTimeSec              float64        `json:"uptime_sec"`
	Time                   string         `json:"time"`
	TimeUnix               int64          `json:"unixtime"`
	StatusCodeCount        map[string]int `json:"status_code_count"`
	TotalStatusCodeCount   map[string]int `json:"total_status_code_count"`
	Count                  int            `json:"count"`
	TotalCount             int            `json:"total_count"`
	TotalResponseTime      string         `json:"total_response_time`
	TotalResponseTimeSec   float64        `json:"total_response_time_sec"`
	AverageResponseTime    string         `json:"average_response_time"`
	AverageResponseTimeSec float64        `json:"average_response_time_sec"`
}

func (mw *Stats) Data() *data {

	mw.mu.RLock()

	now := time.Now()

	uptime := now.Sub(mw.Uptime)

	count := 0
	for _, current := range mw.ResponseCounts {
		count += current
	}

	totalCount := 0
	for _, count := range mw.TotalResponseCounts {
		totalCount += count
	}

	totalResponseTime := mw.TotalResponseTime.Sub(time.Time{})

	averageResponseTime := time.Duration(0)
	if totalCount > 0 {
		avgNs := int64(totalResponseTime) / int64(totalCount)
		averageResponseTime = time.Duration(avgNs)
	}

	r := &data{
		Pid:                    mw.Pid,
		UpTime:                 uptime.String(),
		UpTimeSec:              uptime.Seconds(),
		Time:                   now.String(),
		TimeUnix:               now.Unix(),
		StatusCodeCount:        mw.ResponseCounts,
		TotalStatusCodeCount:   mw.TotalResponseCounts,
		Count:                  count,
		TotalCount:             totalCount,
		TotalResponseTime:      totalResponseTime.String(),
		TotalResponseTimeSec:   totalResponseTime.Seconds(),
		AverageResponseTime:    averageResponseTime.String(),
		AverageResponseTimeSec: averageResponseTime.Seconds(),
	}

	mw.mu.RUnlock()

	return r
}
<|endoftext|>


<commit_before>package msgqueue

import (
	"context"
	"fmt"
	"github.com/labstack/echo/v4"
	"github.com/nats-io/stan.go"
	"net/http"
	"xorkevin.dev/governor"
	"xorkevin.dev/governor/util/uid"
)

const (
	uidSize = 16
)

type (
	// Msgqueue is a service wrapper around a nats streaming queue client instance
	Msgqueue interface {
		Subscribe(queueid, queuegroup string, worker func(msgdata []byte)) (Subscription, error)
		Publish(queueid string, msgdata []byte) error
	}

	Service interface {
		governor.Service
		Msgqueue
	}

	getClientRes struct {
		client stan.Conn
		err    error
	}

	getOp struct {
		res chan<- getClientRes
	}

	service struct {
		client     stan.Conn
		auth       string
		clientid   string
		clusterid  string
		addr       string
		hbinterval int
		hbmaxfail  int
		config     governor.SecretReader
		logger     governor.Logger
		ops        chan getOp
		done       <-chan struct{}
	}

	Subscription interface {
		Close() error
	}

	subscription struct {
		s      *service
		logger governor.Logger
		sub    stan.Subscription
		worker func(data []byte)
	}
)

// New creates a new msgqueue service
func New() Service {
	return &service{}
}

func (s *service) Register(r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	r.SetDefault("auth", "")
	r.SetDefault("host", "localhost")
	r.SetDefault("port", "4222")
	r.SetDefault("cluster", "nss")
	r.SetDefault("hbinterval", 5)
	r.SetDefault("hbmaxfail", 5)
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, g *echo.Group) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	s.config = r

	clientid, err := uid.New(uidSize)
	if err != nil {
		return governor.NewError("Failed to create new uid", http.StatusInternalServerError, err)
	}
	s.clientid = clientid.Base64()

	conf := r.GetStrMap("")
	s.clusterid = conf["cluster"]
	s.addr = fmt.Sprintf("%s:%s", conf["host"], conf["port"])
	s.hbinterval = r.GetInt("hbinterval")
	s.hbmaxfail = r.GetInt("hbmaxfail")

	done := make(chan struct{})
	go s.execute(ctx, done)
	s.done = done

	if _, err := s.getClient(); err != nil {
		return err
	}
	return nil
}

func (s *service) execute(ctx context.Context, done chan<- struct{}) {
	defer close(done)
	for {
		select {
		case <-ctx.Done():
			s.closeClient()
			return
		case op := <-s.ops:
			client, err := s.handleGetClient()
			op.res <- getClientRes{
				client: client,
				err:    err,
			}
			close(op.res)
		}
	}
}

func (s *service) handleGetClient() (stan.Conn, error) {
	authsecret, err := s.config.GetSecret("auth")
	if err != nil {
		return nil, err
	}
	auth := authsecret["password"].(string)
	if auth == s.auth {
		return s.client, nil
	}

	s.closeClient()

	conn, err := stan.Connect(s.clusterid, s.clientid,
		stan.NatsURL(fmt.Sprintf("nats://%s@%s", auth, s.addr)),
		stan.Pings(s.hbinterval, s.hbmaxfail),
		stan.SetConnectionLostHandler(func(_ stan.Conn, err error) {
			s.logger.Error("Lost connection to msgqueue", map[string]string{
				"error": err.Error(),
			})
			s.config.InvalidateSecret("auth")
		}))
	if err != nil {
		return nil, governor.NewError("Failed to connect to msgqueue", http.StatusInternalServerError, err)
	}

	s.client = conn
	s.auth = auth
	s.logger.Info(fmt.Sprintf("established connection to %s %s with client %s", s.addr, s.clusterid, s.clientid), nil)
	return s.client, nil
}

func (s *service) closeClient() {
	if s.client == nil {
		return
	}
	if err := s.client.Close(); err != nil {
		s.logger.Error("failed to close msgqueue connection", map[string]string{
			"error":      err.Error(),
			"actiontype": "closemsgqueueerr",
			"address":    s.addr,
			"clusterid":  s.clusterid,
			"clientid":   s.clientid,
		})
	} else {
		s.logger.Info("closed msgqueue connection", map[string]string{
			"actiontype": "closemsgqueueok",
			"address":    s.addr,
			"clusterid":  s.clusterid,
			"clientid":   s.clientid,
		})
	}
}

func (s *service) getClient() (stan.Conn, error) {
	res := make(chan getClientRes)
	op := getOp{
		res: res,
	}
	select {
	case <-s.done:
		return nil, governor.NewError("Msgqueue service shutdown", http.StatusInternalServerError, nil)
	case s.ops <- op:
		v := <-res
		return v.client, v.err
	}
}

func (s *service) Setup(req governor.ReqSetup) error {
	return nil
}

func (s *service) Start(ctx context.Context) error {
	return nil
}

func (s *service) Stop(ctx context.Context) {
	l := s.logger.WithData(map[string]string{
		"phase": "stop",
	})
	select {
	case <-s.done:
		return
	case <-ctx.Done():
		l.Warn("failed to stop", nil)
	}
}

func (s *service) Health() error {
	return nil
}

// SubscribeQueue subscribes to a queue
func (s *service) Subscribe(queueid, queuegroup string, worker func(msgdata []byte)) (Subscription, error) {
	l := s.logger.WithData(map[string]string{
		"agent":      "subscriber",
		"queueid":    queueid,
		"queuegroup": queuegroup,
	})
	msub := &subscription{
		s:      s,
		logger: l,
		worker: worker,
	}
	client, err := s.getClient()
	if err != nil {
		return nil, err
	}
	sub, err := client.QueueSubscribe(queueid, queuegroup, msub.subscriber, stan.DurableName(queuegroup+"-durable"), stan.SetManualAckMode())
	if err != nil {
		return nil, governor.NewError("Failed to create subscription to queue", http.StatusInternalServerError, err)
	}
	msub.sub = sub
	return msub, nil
}

func (s *service) Publish(queueid string, msgdata []byte) error {
	client, err := s.getClient()
	if err != nil {
		return err
	}
	if err := client.Publish(queueid, msgdata); err != nil {
		return governor.NewError("Failed to publish message: ", http.StatusInternalServerError, err)
	}
	return nil
}

func (s *subscription) subscriber(msg *stan.Msg) {
	s.worker(msg.Data)
	if err := msg.Ack(); err != nil {
		s.logger.Error("Failed to ack message", map[string]string{
			"error":      err.Error(),
			"actiontype": "ackmessage",
		})
	}
}

// Close closes the subscription
func (s *subscription) Close() error {
	if err := s.sub.Close(); err != nil {
		return governor.NewError("Failed to close subscription", http.StatusInternalServerError, err)
	}
	return nil
}
<commit_msg>Improve msgqueue subscription config<commit_after>package msgqueue

import (
	"context"
	"fmt"
	"github.com/labstack/echo/v4"
	"github.com/nats-io/stan.go"
	"net/http"
	"time"
	"xorkevin.dev/governor"
	"xorkevin.dev/governor/util/uid"
)

const (
	uidSize = 16
)

type (
	WorkerFunc func(msgdata []byte) error

	// Msgqueue is a service wrapper around a nats streaming client instance
	Msgqueue interface {
		Subscribe(channel, group string, ackwait time.Duration, inflight int, worker WorkerFunc) (Subscription, error)
		Publish(channel string, msgdata []byte) error
	}

	Service interface {
		governor.Service
		Msgqueue
	}

	getClientRes struct {
		client stan.Conn
		err    error
	}

	getOp struct {
		res chan<- getClientRes
	}

	service struct {
		client     stan.Conn
		auth       string
		clientid   string
		clusterid  string
		addr       string
		hbinterval int
		hbmaxfail  int
		config     governor.SecretReader
		logger     governor.Logger
		ops        chan getOp
		done       <-chan struct{}
	}

	Subscription interface {
		Unsubscribe() error
		Close() error
	}

	subscription struct {
		s      *service
		logger governor.Logger
		sub    stan.Subscription
		worker WorkerFunc
	}
)

// New creates a new msgqueue service
func New() Service {
	return &service{}
}

func (s *service) Register(r governor.ConfigRegistrar, jr governor.JobRegistrar) {
	r.SetDefault("auth", "")
	r.SetDefault("host", "localhost")
	r.SetDefault("port", "4222")
	r.SetDefault("cluster", "nss")
	r.SetDefault("hbinterval", 5)
	r.SetDefault("hbmaxfail", 5)
}

func (s *service) Init(ctx context.Context, c governor.Config, r governor.ConfigReader, l governor.Logger, g *echo.Group) error {
	s.logger = l
	l = s.logger.WithData(map[string]string{
		"phase": "init",
	})

	s.config = r

	clientid, err := uid.New(uidSize)
	if err != nil {
		return governor.NewError("Failed to create new uid", http.StatusInternalServerError, err)
	}
	s.clientid = clientid.Base64()

	conf := r.GetStrMap("")
	s.clusterid = conf["cluster"]
	s.addr = fmt.Sprintf("%s:%s", conf["host"], conf["port"])
	s.hbinterval = r.GetInt("hbinterval")
	s.hbmaxfail = r.GetInt("hbmaxfail")

	done := make(chan struct{})
	go s.execute(ctx, done)
	s.done = done

	if _, err := s.getClient(); err != nil {
		return err
	}
	return nil
}

func (s *service) execute(ctx context.Context, done chan<- struct{}) {
	defer close(done)
	for {
		select {
		case <-ctx.Done():
			s.closeClient()
			return
		case op := <-s.ops:
			client, err := s.handleGetClient()
			op.res <- getClientRes{
				client: client,
				err:    err,
			}
			close(op.res)
		}
	}
}

func (s *service) handleGetClient() (stan.Conn, error) {
	authsecret, err := s.config.GetSecret("auth")
	if err != nil {
		return nil, err
	}
	auth := authsecret["password"].(string)
	if auth == s.auth {
		return s.client, nil
	}

	s.closeClient()

	conn, err := stan.Connect(s.clusterid, s.clientid,
		stan.NatsURL(fmt.Sprintf("nats://%s@%s", auth, s.addr)),
		stan.Pings(s.hbinterval, s.hbmaxfail),
		stan.SetConnectionLostHandler(func(_ stan.Conn, err error) {
			s.logger.Error("Lost connection to msgqueue", map[string]string{
				"error": err.Error(),
			})
			s.config.InvalidateSecret("auth")
		}))
	if err != nil {
		return nil, governor.NewError("Failed to connect to msgqueue", http.StatusInternalServerError, err)
	}

	s.client = conn
	s.auth = auth
	s.logger.Info(fmt.Sprintf("established connection to %s %s with client %s", s.addr, s.clusterid, s.clientid), nil)
	return s.client, nil
}

func (s *service) closeClient() {
	if s.client == nil {
		return
	}
	if err := s.client.Close(); err != nil {
		s.logger.Error("failed to close msgqueue connection", map[string]string{
			"error":      err.Error(),
			"actiontype": "closemsgqueueerr",
			"address":    s.addr,
			"clusterid":  s.clusterid,
			"clientid":   s.clientid,
		})
	} else {
		s.logger.Info("closed msgqueue connection", map[string]string{
			"actiontype": "closemsgqueueok",
			"address":    s.addr,
			"clusterid":  s.clusterid,
			"clientid":   s.clientid,
		})
	}
}

func (s *service) getClient() (stan.Conn, error) {
	res := make(chan getClientRes)
	op := getOp{
		res: res,
	}
	select {
	case <-s.done:
		return nil, governor.NewError("Msgqueue service shutdown", http.StatusInternalServerError, nil)
	case s.ops <- op:
		v := <-res
		return v.client, v.err
	}
}

func (s *service) Setup(req governor.ReqSetup) error {
	return nil
}

func (s *service) Start(ctx context.Context) error {
	return nil
}

func (s *service) Stop(ctx context.Context) {
	l := s.logger.WithData(map[string]string{
		"phase": "stop",
	})
	select {
	case <-s.done:
		return
	case <-ctx.Done():
		l.Warn("failed to stop", nil)
	}
}

func (s *service) Health() error {
	return nil
}

// Subscribe subscribes to a channel
func (s *service) Subscribe(channel, group string, ackwait time.Duration, inflight int, worker WorkerFunc) (Subscription, error) {
	l := s.logger.WithData(map[string]string{
		"agent":   "subscriber",
		"channel": channel,
		"group":   group,
	})
	msub := &subscription{
		s:      s,
		logger: l,
		worker: worker,
	}
	client, err := s.getClient()
	if err != nil {
		return nil, err
	}
	sub, err := client.QueueSubscribe(channel, group,
		msub.subscriber,
		stan.DurableName(group+"-durable"),
		stan.SetManualAckMode(),
		stan.AckWait(ackwait),
		stan.MaxInflight(inflight))
	if err != nil {
		return nil, governor.NewError("Failed to create subscription to channel", http.StatusInternalServerError, err)
	}
	msub.sub = sub
	return msub, nil
}

// Publish publishes to a channel
func (s *service) Publish(channel string, msgdata []byte) error {
	client, err := s.getClient()
	if err != nil {
		return err
	}
	if err := client.Publish(channel, msgdata); err != nil {
		return governor.NewError("Failed to publish message: ", http.StatusInternalServerError, err)
	}
	return nil
}

func (s *subscription) subscriber(msg *stan.Msg) {
	if err := s.worker(msg.Data); err != nil {
		s.logger.Error("Failed to execute message handler", map[string]string{
			"error":      err.Error(),
			"actiontype": "execworker",
		})
		return
	}
	if err := msg.Ack(); err != nil {
		s.logger.Error("Failed to ack message", map[string]string{
			"error":      err.Error(),
			"actiontype": "ackmessage",
		})
	}
}

// Unsubscribe removes the subscription
func (s *subscription) Unsubscribe() error {
	if err := s.sub.Unsubscribe(); err != nil {
		return governor.NewError("Failed to unsubscribe", http.StatusInternalServerError, err)
	}
	return nil
}

// Close closes the subscription
func (s *subscription) Close() error {
	if err := s.sub.Close(); err != nil {
		return governor.NewError("Failed to close subscription", http.StatusInternalServerError, err)
	}
	return nil
}
<|endoftext|>