<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 control

import (
	"errors"
	"fmt"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/gomit"

	"github.com/intelsdi-x/pulse/control/plugin"
	"github.com/intelsdi-x/pulse/core"
	"github.com/intelsdi-x/pulse/core/control_event"
)

var (
	runnerLog = log.WithField("_module", "control-runner")
)

type availablePluginState int

const (
	HandlerRegistrationName = "control.runner"

	// availablePlugin States
	PluginRunning availablePluginState = iota - 1 // Default value (0) is Running
	PluginStopped
	PluginDisabled
)

// TBD
type executablePlugin interface {
	Start() error
	Kill() error
	WaitForResponse(time.Duration) (*plugin.Response, error)
}

// Handles events pertaining to plugins and control the runnning state accordingly.
type runner struct {
	delegates        []gomit.Delegator
	emitter          gomit.Emitter
	monitor          *monitor
	availablePlugins *availablePlugins
	metricCatalog    catalogsMetrics
	pluginManager    managesPlugins
	routingStrategy  RoutingStrategy
}

func newRunner(routingStrategy RoutingStrategy) *runner {
	r := &runner{
		monitor:          newMonitor(),
		availablePlugins: newAvailablePlugins(routingStrategy),
		routingStrategy:  routingStrategy,
	}
	return r
}

func (r *runner) SetStrategy(rs RoutingStrategy) {
	r.routingStrategy = rs
}

func (r *runner) Strategy() RoutingStrategy {
	return r.routingStrategy
}

func (r *runner) SetMetricCatalog(c catalogsMetrics) {
	r.metricCatalog = c
}

func (r *runner) SetEmitter(e gomit.Emitter) {
	r.emitter = e
}

func (r *runner) SetPluginManager(m managesPlugins) {
	r.pluginManager = m
}

func (r *runner) AvailablePlugins() *availablePlugins {
	return r.availablePlugins
}

func (r *runner) Monitor() *monitor {
	return r.monitor
}

// Adds Delegates (gomit.Delegator) for adding Runner handlers to on Start and
// unregistration on Stop.
func (r *runner) AddDelegates(delegates ...gomit.Delegator) {
	// Append the variadic collection of gomit.RegisterHanlders to r.delegates
	r.delegates = append(r.delegates, delegates...)
}

// Begin handing events and managing available plugins
func (r *runner) Start() error {
	// Delegates must be added before starting if none exist
	// then this Runner can do nothing and should not start.
	if len(r.delegates) == 0 {
		return errors.New("No delegates added before called Start()")
	}

	// For each delegate register needed handlers
	for _, del := range r.delegates {
		e := del.RegisterHandler(HandlerRegistrationName, r)
		if e != nil {
			return e
		}
	}

	// Start the monitor
	r.monitor.Start(r.availablePlugins)
	runnerLog.WithFields(log.Fields{
		"_block": "start",
	}).Debug("started")
	return nil
}

// Stop handling, gracefully stop all plugins.
func (r *runner) Stop() []error {
	var errs []error

	// Stop the monitor
	r.monitor.Stop()

	// TODO: Actually stop the plugins

	// For each delegate unregister needed handlers
	for _, del := range r.delegates {
		e := del.UnregisterHandler(HandlerRegistrationName)
		if e != nil {
			errs = append(errs, e)
		}
	}
	defer runnerLog.WithFields(log.Fields{
		"_block": "start-plugin",
	}).Debug("stopped")
	return errs
}

func (r *runner) startPlugin(p executablePlugin) (*availablePlugin, error) {
	e := p.Start()
	if e != nil {
		e_ := errors.New("error while starting plugin: " + e.Error())
		defer runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e_
	}

	// Wait for plugin response
	resp, err := p.WaitForResponse(time.Second * 3)
	if err != nil {
		e := errors.New("error while waiting for response: " + err.Error())
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	if resp == nil {
		e := errors.New("no reponse object returned from plugin")
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	if resp.State != plugin.PluginSuccess {
		e := errors.New("plugin could not start error: " + resp.ErrorMessage)
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	// build availablePlugin
	ap, err := newAvailablePlugin(resp, r.emitter, p)
	if err != nil {
		return nil, err
	}

	if resp.Meta.Unsecure {
		err = ap.client.Ping()
	} else {
		err = ap.client.SetKey()
	}
	if err != nil {
		return nil, err
	}

	r.availablePlugins.insert(ap)
	runnerLog.WithFields(log.Fields{
		"_block":                "start-plugin",
		"available-plugin":      ap.String(),
		"available-plugin-type": ap.TypeName(),
	}).Info("available plugin started")

	return ap, nil
}

func (r *runner) stopPlugin(reason string, ap *availablePlugin) error {
	err := ap.Stop(reason)
	if err != nil {
		return err
	}
	pool, err := r.availablePlugins.getPool(ap.key)
	if err != nil {
		return err
	}
	if pool != nil {
		pool.remove(ap.id)
	}
	return nil
}

// Empty handler acting as placeholder until implementation. This helps tests
// pass to ensure registration works.
func (r *runner) HandleGomitEvent(e gomit.Event) {
	switch v := e.Body.(type) {
	case *control_event.DeadAvailablePluginEvent:
		runnerLog.WithFields(log.Fields{
			"_block":  "handle-events",
			"event":   v.Namespace(),
			"aplugin": v.String,
		}).Warning("handling dead available plugin event")
		pool, err := r.availablePlugins.getPool(v.Key)
		if err != nil {
			runnerLog.WithFields(log.Fields{
				"_block":  "handle-events",
				"aplugin": v.String,
			}).Error(err.Error())
			return
		}
		if pool != nil {
			pool.kill(v.Id, "plugin dead")
		}
	case *control_event.PluginUnsubscriptionEvent:
		runnerLog.WithFields(log.Fields{
			"_block":         "subscribe-pool",
			"event":          v.Namespace(),
			"plugin-name":    v.PluginName,
			"plugin-version": v.PluginVersion,
			"plugin-type":    core.PluginType(v.PluginType).String(),
		}).Debug("handling plugin unsubscription event")

		err := r.handleUnsubscription(core.PluginType(v.PluginType).String(), v.PluginName, v.PluginVersion, v.TaskId)
		if err != nil {
			return
		}
	case *control_event.LoadPluginEvent:
		var pool *apPool
		r.availablePlugins.RLock()
		for key, p := range r.availablePlugins.pools() {
			// tuple of type name and version
			// type @ index 0, name @ index 1, version @ index 2
			tnv := strings.Split(key, ":")
			// make sure we don't panic and crash the service if a junk key is retrieved
			if len(tnv) != 3 {
				runnerLog.WithFields(log.Fields{
					"_block":         "subscribe-pool",
					"event":          v.Namespace(),
					"plugin-name":    v.Name,
					"plugin-version": v.Version,
					"plugin-type":    core.PluginType(v.Type).String(),
					"plugin-signed":  v.Signed,
				}).Info("pool has bad key ", key)
				continue
			}

			// attempt to find a pool whose type and name are the same, and whose version is
			// less than newly loaded plugin.  If we find it, break out of loop.
			if core.PluginType(v.Type).String() == tnv[0] && v.Name == tnv[1] && v.Version > p.version {
				pool = p
				break
			}
		}
		r.availablePlugins.RUnlock()
		// now check to see if anything was put where pool points.
		// if not, there are no older pools whose subscriptions need to be
		// moved.
		if pool == nil {
			runnerLog.WithFields(log.Fields{
				"_block":         "subscribe-pool",
				"event":          v.Namespace(),
				"plugin-name":    v.Name,
				"plugin-version": v.Version,
				"plugin-type":    core.PluginType(v.Type).String(),
			}).Info("No previous pool found for loaded plugin")
			return
		}
		plugin, err := r.pluginManager.get(fmt.Sprintf("%s:%s:%d", core.PluginType(v.Type).String(), v.Name, v.Version))
		if err != nil {
			return
		}
		newPool, err := r.availablePlugins.getOrCreatePool(plugin.Key())
		if err != nil {
			return
		}

		subs := pool.moveSubscriptions(newPool)
		if len(subs) != 0 {
			runnerLog.WithFields(log.Fields{
				"_block":         "subscribe-pool",
				"event":          v.Namespace(),
				"plugin-name":    v.Name,
				"plugin-version": v.Version,
				"plugin-type":    core.PluginType(v.Type).String(),
			}).Info("pool with subscriptions to move found")
			for _, sub := range subs {
				r.emitter.Emit(&control_event.PluginSubscriptionEvent{
					PluginName:       v.Name,
					PluginVersion:    v.Version,
					TaskId:           sub.taskId,
					PluginType:       v.Type,
					SubscriptionType: int(unboundSubscriptionType),
				})
				r.emitter.Emit(&control_event.PluginUnsubscriptionEvent{
					PluginName:    v.Name,
					PluginVersion: pool.version,
					TaskId:        sub.taskId,
					PluginType:    v.Type,
				})
				r.emitter.Emit(&control_event.MovePluginSubscriptionEvent{
					PluginName:      v.Name,
					PreviousVersion: pool.version,
					NewVersion:      v.Version,
					TaskId:          sub.taskId,
					PluginType:      v.Type,
				})
			}
		}
	default:
		runnerLog.WithFields(log.Fields{
			"_block": "handle-events",
			"event":  v.Namespace(),
		}).Info("Nothing to do for this event")
	}
}

func (r *runner) runPlugin(path string) error {
	ePlugin, err := plugin.NewExecutablePlugin(r.pluginManager.GenerateArgs(path), path)
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block": "run-plugin",
			"path":   path,
			"error":  err,
		}).Error("error creating executable plugin")
		return err
	}
	_, err = r.startPlugin(ePlugin)
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block": "run-plugin",
			"path":   path,
			"error":  err,
		}).Error("error starting new plugin")
		return err
	}
	return nil
}

func (r *runner) handleUnsubscription(pType, pName string, pVersion int, taskId string) error {
	pool, err := r.availablePlugins.getPool(fmt.Sprintf("%s:%s:%d", pType, pName, pVersion))
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block":         "handle-unsubscription",
			"plugin-name":    pName,
			"plugin-version": pVersion,
			"plugin-type":    pType,
		}).Error("error retrieving pool")
		return errors.New("error retrieving pool")
	}
	if pool == nil {
		runnerLog.WithFields(log.Fields{
			"_block":         "handle-unsubscription",
			"plugin-name":    pName,
			"plugin-version": pVersion,
			"plugin-type":    pType,
		}).Error("pool not found")
		return errors.New("pool not found")
	}
	if pool.subscriptionCount() < pool.count() {
		runnerLog.WithFields(log.Fields{
			"_block":                  "handle-unsubscription",
			"pool-count":              pool.count(),
			"pool-subscription-count": pool.subscriptionCount(),
		}).Debug(fmt.Sprintf("killing an available plugin in pool  %s:%s:%d", pType, pName, pVersion))
		pool.killLeastUsed("unsubscription event")
	}
	return nil
}
<commit_msg>Increases the timeout for waiting for a plugin response from 3 to 5 seconds<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 control

import (
	"errors"
	"fmt"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/gomit"

	"github.com/intelsdi-x/pulse/control/plugin"
	"github.com/intelsdi-x/pulse/core"
	"github.com/intelsdi-x/pulse/core/control_event"
)

var (
	runnerLog = log.WithField("_module", "control-runner")
)

type availablePluginState int

const (
	HandlerRegistrationName = "control.runner"

	// availablePlugin States
	PluginRunning availablePluginState = iota - 1 // Default value (0) is Running
	PluginStopped
	PluginDisabled
)

// TBD
type executablePlugin interface {
	Start() error
	Kill() error
	WaitForResponse(time.Duration) (*plugin.Response, error)
}

// Handles events pertaining to plugins and control the runnning state accordingly.
type runner struct {
	delegates        []gomit.Delegator
	emitter          gomit.Emitter
	monitor          *monitor
	availablePlugins *availablePlugins
	metricCatalog    catalogsMetrics
	pluginManager    managesPlugins
	routingStrategy  RoutingStrategy
}

func newRunner(routingStrategy RoutingStrategy) *runner {
	r := &runner{
		monitor:          newMonitor(),
		availablePlugins: newAvailablePlugins(routingStrategy),
		routingStrategy:  routingStrategy,
	}
	return r
}

func (r *runner) SetStrategy(rs RoutingStrategy) {
	r.routingStrategy = rs
}

func (r *runner) Strategy() RoutingStrategy {
	return r.routingStrategy
}

func (r *runner) SetMetricCatalog(c catalogsMetrics) {
	r.metricCatalog = c
}

func (r *runner) SetEmitter(e gomit.Emitter) {
	r.emitter = e
}

func (r *runner) SetPluginManager(m managesPlugins) {
	r.pluginManager = m
}

func (r *runner) AvailablePlugins() *availablePlugins {
	return r.availablePlugins
}

func (r *runner) Monitor() *monitor {
	return r.monitor
}

// Adds Delegates (gomit.Delegator) for adding Runner handlers to on Start and
// unregistration on Stop.
func (r *runner) AddDelegates(delegates ...gomit.Delegator) {
	// Append the variadic collection of gomit.RegisterHanlders to r.delegates
	r.delegates = append(r.delegates, delegates...)
}

// Begin handing events and managing available plugins
func (r *runner) Start() error {
	// Delegates must be added before starting if none exist
	// then this Runner can do nothing and should not start.
	if len(r.delegates) == 0 {
		return errors.New("No delegates added before called Start()")
	}

	// For each delegate register needed handlers
	for _, del := range r.delegates {
		e := del.RegisterHandler(HandlerRegistrationName, r)
		if e != nil {
			return e
		}
	}

	// Start the monitor
	r.monitor.Start(r.availablePlugins)
	runnerLog.WithFields(log.Fields{
		"_block": "start",
	}).Debug("started")
	return nil
}

// Stop handling, gracefully stop all plugins.
func (r *runner) Stop() []error {
	var errs []error

	// Stop the monitor
	r.monitor.Stop()

	// TODO: Actually stop the plugins

	// For each delegate unregister needed handlers
	for _, del := range r.delegates {
		e := del.UnregisterHandler(HandlerRegistrationName)
		if e != nil {
			errs = append(errs, e)
		}
	}
	defer runnerLog.WithFields(log.Fields{
		"_block": "start-plugin",
	}).Debug("stopped")
	return errs
}

func (r *runner) startPlugin(p executablePlugin) (*availablePlugin, error) {
	e := p.Start()
	if e != nil {
		e_ := errors.New("error while starting plugin: " + e.Error())
		defer runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e_
	}

	// Wait for plugin response
	resp, err := p.WaitForResponse(time.Second * 5)
	if err != nil {
		e := errors.New("error while waiting for response: " + err.Error())
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	if resp == nil {
		e := errors.New("no reponse object returned from plugin")
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	if resp.State != plugin.PluginSuccess {
		e := errors.New("plugin could not start error: " + resp.ErrorMessage)
		runnerLog.WithFields(log.Fields{
			"_block": "start-plugin",
			"error":  e.Error(),
		}).Error("error starting a plugin")
		return nil, e
	}

	// build availablePlugin
	ap, err := newAvailablePlugin(resp, r.emitter, p)
	if err != nil {
		return nil, err
	}

	if resp.Meta.Unsecure {
		err = ap.client.Ping()
	} else {
		err = ap.client.SetKey()
	}
	if err != nil {
		return nil, err
	}

	r.availablePlugins.insert(ap)
	runnerLog.WithFields(log.Fields{
		"_block":                "start-plugin",
		"available-plugin":      ap.String(),
		"available-plugin-type": ap.TypeName(),
	}).Info("available plugin started")

	return ap, nil
}

func (r *runner) stopPlugin(reason string, ap *availablePlugin) error {
	err := ap.Stop(reason)
	if err != nil {
		return err
	}
	pool, err := r.availablePlugins.getPool(ap.key)
	if err != nil {
		return err
	}
	if pool != nil {
		pool.remove(ap.id)
	}
	return nil
}

// Empty handler acting as placeholder until implementation. This helps tests
// pass to ensure registration works.
func (r *runner) HandleGomitEvent(e gomit.Event) {
	switch v := e.Body.(type) {
	case *control_event.DeadAvailablePluginEvent:
		runnerLog.WithFields(log.Fields{
			"_block":  "handle-events",
			"event":   v.Namespace(),
			"aplugin": v.String,
		}).Warning("handling dead available plugin event")
		pool, err := r.availablePlugins.getPool(v.Key)
		if err != nil {
			runnerLog.WithFields(log.Fields{
				"_block":  "handle-events",
				"aplugin": v.String,
			}).Error(err.Error())
			return
		}
		if pool != nil {
			pool.kill(v.Id, "plugin dead")
		}
	case *control_event.PluginUnsubscriptionEvent:
		runnerLog.WithFields(log.Fields{
			"_block":         "subscribe-pool",
			"event":          v.Namespace(),
			"plugin-name":    v.PluginName,
			"plugin-version": v.PluginVersion,
			"plugin-type":    core.PluginType(v.PluginType).String(),
		}).Debug("handling plugin unsubscription event")

		err := r.handleUnsubscription(core.PluginType(v.PluginType).String(), v.PluginName, v.PluginVersion, v.TaskId)
		if err != nil {
			return
		}
	case *control_event.LoadPluginEvent:
		var pool *apPool
		r.availablePlugins.RLock()
		for key, p := range r.availablePlugins.pools() {
			// tuple of type name and version
			// type @ index 0, name @ index 1, version @ index 2
			tnv := strings.Split(key, ":")
			// make sure we don't panic and crash the service if a junk key is retrieved
			if len(tnv) != 3 {
				runnerLog.WithFields(log.Fields{
					"_block":         "subscribe-pool",
					"event":          v.Namespace(),
					"plugin-name":    v.Name,
					"plugin-version": v.Version,
					"plugin-type":    core.PluginType(v.Type).String(),
					"plugin-signed":  v.Signed,
				}).Info("pool has bad key ", key)
				continue
			}

			// attempt to find a pool whose type and name are the same, and whose version is
			// less than newly loaded plugin.  If we find it, break out of loop.
			if core.PluginType(v.Type).String() == tnv[0] && v.Name == tnv[1] && v.Version > p.version {
				pool = p
				break
			}
		}
		r.availablePlugins.RUnlock()
		// now check to see if anything was put where pool points.
		// if not, there are no older pools whose subscriptions need to be
		// moved.
		if pool == nil {
			runnerLog.WithFields(log.Fields{
				"_block":         "subscribe-pool",
				"event":          v.Namespace(),
				"plugin-name":    v.Name,
				"plugin-version": v.Version,
				"plugin-type":    core.PluginType(v.Type).String(),
			}).Info("No previous pool found for loaded plugin")
			return
		}
		plugin, err := r.pluginManager.get(fmt.Sprintf("%s:%s:%d", core.PluginType(v.Type).String(), v.Name, v.Version))
		if err != nil {
			return
		}
		newPool, err := r.availablePlugins.getOrCreatePool(plugin.Key())
		if err != nil {
			return
		}

		subs := pool.moveSubscriptions(newPool)
		if len(subs) != 0 {
			runnerLog.WithFields(log.Fields{
				"_block":         "subscribe-pool",
				"event":          v.Namespace(),
				"plugin-name":    v.Name,
				"plugin-version": v.Version,
				"plugin-type":    core.PluginType(v.Type).String(),
			}).Info("pool with subscriptions to move found")
			for _, sub := range subs {
				r.emitter.Emit(&control_event.PluginSubscriptionEvent{
					PluginName:       v.Name,
					PluginVersion:    v.Version,
					TaskId:           sub.taskId,
					PluginType:       v.Type,
					SubscriptionType: int(unboundSubscriptionType),
				})
				r.emitter.Emit(&control_event.PluginUnsubscriptionEvent{
					PluginName:    v.Name,
					PluginVersion: pool.version,
					TaskId:        sub.taskId,
					PluginType:    v.Type,
				})
				r.emitter.Emit(&control_event.MovePluginSubscriptionEvent{
					PluginName:      v.Name,
					PreviousVersion: pool.version,
					NewVersion:      v.Version,
					TaskId:          sub.taskId,
					PluginType:      v.Type,
				})
			}
		}
	default:
		runnerLog.WithFields(log.Fields{
			"_block": "handle-events",
			"event":  v.Namespace(),
		}).Info("Nothing to do for this event")
	}
}

func (r *runner) runPlugin(path string) error {
	ePlugin, err := plugin.NewExecutablePlugin(r.pluginManager.GenerateArgs(path), path)
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block": "run-plugin",
			"path":   path,
			"error":  err,
		}).Error("error creating executable plugin")
		return err
	}
	_, err = r.startPlugin(ePlugin)
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block": "run-plugin",
			"path":   path,
			"error":  err,
		}).Error("error starting new plugin")
		return err
	}
	return nil
}

func (r *runner) handleUnsubscription(pType, pName string, pVersion int, taskId string) error {
	pool, err := r.availablePlugins.getPool(fmt.Sprintf("%s:%s:%d", pType, pName, pVersion))
	if err != nil {
		runnerLog.WithFields(log.Fields{
			"_block":         "handle-unsubscription",
			"plugin-name":    pName,
			"plugin-version": pVersion,
			"plugin-type":    pType,
		}).Error("error retrieving pool")
		return errors.New("error retrieving pool")
	}
	if pool == nil {
		runnerLog.WithFields(log.Fields{
			"_block":         "handle-unsubscription",
			"plugin-name":    pName,
			"plugin-version": pVersion,
			"plugin-type":    pType,
		}).Error("pool not found")
		return errors.New("pool not found")
	}
	if pool.subscriptionCount() < pool.count() {
		runnerLog.WithFields(log.Fields{
			"_block":                  "handle-unsubscription",
			"pool-count":              pool.count(),
			"pool-subscription-count": pool.subscriptionCount(),
		}).Debug(fmt.Sprintf("killing an available plugin in pool  %s:%s:%d", pType, pName, pVersion))
		pool.killLeastUsed("unsubscription event")
	}
	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/xmetrics"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

const (
	// DefaultPrimaryAddress is the bind address of the primary server (e.g. talaria, petasos, etc)
	DefaultPrimaryAddress = ":8080"

	// DefaultHealthAddress is the bind address of the health check server
	DefaultHealthAddress = ":8081"

	// DefaultMetricsAddress is the bind address of the metrics server
	DefaultMetricsAddress = ":8082"

	// DefaultHealthLogInterval is the interval at which health statistics are emitted
	// when a non-positive log interval is specified
	DefaultHealthLogInterval time.Duration = time.Duration(60 * time.Second)

	// DefaultLogConnectionState is the default setting for logging connection state messages.  This
	// value is primarily used when a *WebPA value is nil.
	DefaultLogConnectionState = false

	// AlternateSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the alternate server.
	AlternateSuffix = "alternate"

	// DefaultProject is used as a metrics namespace when one is not defined.
	DefaultProject = "xmidt"

	// HealthSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the health server.
	HealthSuffix = "health"

	// PprofSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the pprof server.
	PprofSuffix = "pprof"

	// MetricsSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the metrics server.
	MetricsSuffix = "metrics"

	// FileFlagName is the name of the command-line flag for specifying an alternate
	// configuration file for Viper to hunt for.
	FileFlagName = "file"

	// FileFlagShorthand is the command-line shortcut flag for FileFlagName
	FileFlagShorthand = "f"

	// CPUProfileFlagName is the command-line flag for creating a cpuprofile on the server
	CPUProfileFlagName = "cpuprofile"

	// CPUProfileShortHand is the command-line shortcut for creating cpushorthand on the server
	CPUProfileShorthand = "c"

	// MemProfileFlagName is the command-line flag for creating memprofile on the server
	MemProfileFlagName = "memprofile"

	// MemProfileShortHand is the command-line shortcut for creating memprofile on the server
	MemProfileShorthand = "m"
)

// ConfigureFlagSet adds the standard set of WebPA flags to the supplied FlagSet.  Use of this function
// is optional, and necessary only if the standard flags should be supported.  However, this is highly desirable,
// as ConfigureViper can make use of the standard flags to tailor how configuration is loaded or if gathering cpuprofile
// or memprofile data is needed.
func ConfigureFlagSet(applicationName string, f *pflag.FlagSet) {
	f.StringP(FileFlagName, FileFlagShorthand, applicationName, "base name of the configuration file")
}

// create CPUProfileFiles creates a cpu profile of the server, its triggered by the optional flag cpuprofile
//
// the CPU profile is created on the server's start
func CreateCPUProfileFile(v *viper.Viper, fp *pflag.FlagSet, l log.Logger) {
	flag := fp.Lookup("cpuprofile")
	if flag == nil {
		return
	}

	f, err := os.Create(flag.Value.String())
	if err != nil {
		l.Log("could not create CPU profile: ", err)
	}
	defer f.Close()

	if err := pprof.StartCPUProfile(f); err != nil {
		l.Log("could not start CPU profile: ", err)
	}

	defer pprof.StopCPUProfile()
}

// Create CPUProfileFiles creates a memory profile of the server, its triggered by the optional flag memprofile
//
// the memory profile is created on the server's exit.
// this function should be used within the application.
func CreateMemoryProfileFile(v *viper.Viper, fp *pflag.FlagSet, l log.Logger) {
	flag := fp.Lookup("memprofile")
	if flag == nil {
		return
	}

	f, err := os.Create(flag.Value.String())
	if err != nil {
		l.Log("could not create memory profile: ", err)
	}

	defer f.Close()
	runtime.GC()
	if err := pprof.WriteHeapProfile(f); err != nil {
		l.Log("could not write memory profile: ", err)
	}
}

// ConfigureViper configures a Viper instances using the opinionated WebPA settings.  All WebPA servers should
// use this function.
//
// The flagSet is optional.  If supplied, it will be bound to the given Viper instance.  Additionally, if the
// flagSet has a FileFlagName flag, it will be used as the configuration name to hunt for instead of the
// application name.
func ConfigureViper(applicationName string, f *pflag.FlagSet, v *viper.Viper) (err error) {
	v.AddConfigPath(fmt.Sprintf("/etc/%s", applicationName))
	v.AddConfigPath(fmt.Sprintf("$HOME/.%s", applicationName))
	v.AddConfigPath(".")

	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	v.SetEnvPrefix(applicationName)
	v.AutomaticEnv()

	v.SetDefault("primary.name", applicationName)
	v.SetDefault("primary.address", DefaultPrimaryAddress)
	v.SetDefault("primary.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("alternate.name", fmt.Sprintf("%s.%s", applicationName, AlternateSuffix))

	v.SetDefault("health.name", fmt.Sprintf("%s.%s", applicationName, HealthSuffix))
	v.SetDefault("health.address", DefaultHealthAddress)
	v.SetDefault("health.logInterval", DefaultHealthLogInterval)
	v.SetDefault("health.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("pprof.name", fmt.Sprintf("%s.%s", applicationName, PprofSuffix))
	v.SetDefault("pprof.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("metric.name", fmt.Sprintf("%s.%s", applicationName, MetricsSuffix))
	v.SetDefault("metric.address", DefaultMetricsAddress)

	v.SetDefault("project", DefaultProject)

	configName := applicationName
	if f != nil {
		if fileFlag := f.Lookup(FileFlagName); fileFlag != nil {
			// use the command-line to specify the base name of the file to be searched for
			configName = fileFlag.Value.String()
		}

		err = v.BindPFlags(f)
	}

	v.SetConfigName(configName)
	return
}

/*
Configure is a one-stop shopping function for preparing WebPA configuration.  This function
does not itself read in configuration from the Viper environment.  Typical usage is:

    var (
      f = pflag.NewFlagSet()
      v = viper.New()
    )

    if err := server.Configure("petasos", os.Args, f, v); err != nil {
      // deal with the error, possibly just exiting
    }

    // further customizations to the Viper instance can be done here

    if err := v.ReadInConfig(); err != nil {
      // more error handling
    }

Usage of this function is only necessary if custom configuration is needed.  Normally,
using New will suffice.
*/
func Configure(applicationName string, arguments []string, f *pflag.FlagSet, v *viper.Viper) (err error) {
	if f != nil {
		ConfigureFlagSet(applicationName, f)
		err = f.Parse(arguments)
		if err != nil {
			return
		}
	}

	err = ConfigureViper(applicationName, f, v)
	return
}

/*
Initialize handles the bootstrapping of the server code for a WebPA node.  It configures Viper,
reads configuration, and unmarshals the appropriate objects.  This function is typically all that's
needed to fully instantiate a WebPA server.  Typical usage:

    var (
      f = pflag.NewFlagSet()
      v = viper.New()

      // can customize both the FlagSet and the Viper before invoking New
      logger, registry, webPA, err = server.Initialize("petasos", os.Args, f, v)
    )

    if err != nil {
      // deal with the error, possibly just exiting
    }

Note that the FlagSet is optional but highly encouraged.  If not supplied, then no command-line binding
is done for the unmarshalled configuration.

This function always returns a logger, regardless of any errors.  This allows clients to use the returned
logger when reporting errors.  This function falls back to a logger that writes to os.Stdout if it cannot
create a logger from the Viper environment.
*/
func Initialize(applicationName string, arguments []string, f *pflag.FlagSet, v *viper.Viper, modules ...xmetrics.Module) (logger log.Logger, registry xmetrics.Registry, webPA *WebPA, err error) {
	defer func() {
		if err != nil {
			// never return a WebPA in the presence of an error, to
			// avoid an ambiguous API
			webPA = nil

			// Make sure there's at least a default logger for the caller to use
			logger = logging.DefaultLogger()
		}
	}()

	if err = Configure(applicationName, arguments, f, v); err != nil {
		return
	}

	if err = v.ReadInConfig(); err != nil {
		return
	}

	webPA = &WebPA{
		ApplicationName: applicationName,
	}

	err = v.Unmarshal(webPA)
	if err != nil {
		return
	}

	logger = logging.New(webPA.Log)
	logger.Log(level.Key(), level.InfoValue(), logging.MessageKey(), "initialized Viper environment", "configurationFile", v.ConfigFileUsed())

	if len(webPA.Metric.MetricsOptions.Namespace) == 0 {
		webPA.Metric.MetricsOptions.Namespace = applicationName
	}

	if len(webPA.Metric.MetricsOptions.Subsystem) == 0 {
		webPA.Metric.MetricsOptions.Subsystem = applicationName
	}

	webPA.Metric.MetricsOptions.Logger = logger
	registry, err = webPA.Metric.NewRegistry(modules...)
	if err != nil {
		return
	}

	// CreateCPUProfileFile(v, f, logger)
	return
}
<commit_msg>server: include in ConfigureFlagSet<commit_after>package server

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/Comcast/webpa-common/logging"
	"github.com/Comcast/webpa-common/xmetrics"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"
)

const (
	// DefaultPrimaryAddress is the bind address of the primary server (e.g. talaria, petasos, etc)
	DefaultPrimaryAddress = ":8080"

	// DefaultHealthAddress is the bind address of the health check server
	DefaultHealthAddress = ":8081"

	// DefaultMetricsAddress is the bind address of the metrics server
	DefaultMetricsAddress = ":8082"

	// DefaultHealthLogInterval is the interval at which health statistics are emitted
	// when a non-positive log interval is specified
	DefaultHealthLogInterval time.Duration = time.Duration(60 * time.Second)

	// DefaultLogConnectionState is the default setting for logging connection state messages.  This
	// value is primarily used when a *WebPA value is nil.
	DefaultLogConnectionState = false

	// AlternateSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the alternate server.
	AlternateSuffix = "alternate"

	// DefaultProject is used as a metrics namespace when one is not defined.
	DefaultProject = "xmidt"

	// HealthSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the health server.
	HealthSuffix = "health"

	// PprofSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the pprof server.
	PprofSuffix = "pprof"

	// MetricsSuffix is the suffix appended to the server name, along with a period (.), for
	// logging information pertinent to the metrics server.
	MetricsSuffix = "metrics"

	// FileFlagName is the name of the command-line flag for specifying an alternate
	// configuration file for Viper to hunt for.
	FileFlagName = "file"

	// FileFlagShorthand is the command-line shortcut flag for FileFlagName
	FileFlagShorthand = "f"

	// CPUProfileFlagName is the command-line flag for creating a cpuprofile on the server
	CPUProfileFlagName = "cpuprofile"

	// CPUProfileShortHand is the command-line shortcut for creating cpushorthand on the server
	CPUProfileShorthand = "c"

	// MemProfileFlagName is the command-line flag for creating memprofile on the server
	MemProfileFlagName = "memprofile"

	// MemProfileShortHand is the command-line shortcut for creating memprofile on the server
	MemProfileShorthand = "m"
)

// ConfigureFlagSet adds the standard set of WebPA flags to the supplied FlagSet.  Use of this function
// is optional, and necessary only if the standard flags should be supported.  However, this is highly desirable,
// as ConfigureViper can make use of the standard flags to tailor how configuration is loaded or if gathering cpuprofile
// or memprofile data is needed.
func ConfigureFlagSet(applicationName string, f *pflag.FlagSet) {
	f.StringP(FileFlagName, FileFlagShorthand, applicationName, "base name of the configuration file")
	f.StringP(CPUProfileFlagName, CPUProfileShorthand, "cpuprofile", "base name of the cpuprofile file")
	f.StringP(MemProfileFlagName, MemProfileShorthand, "memprofile", "base name of the memprofile file")
}

// create CPUProfileFiles creates a cpu profile of the server, its triggered by the optional flag cpuprofile
//
// the CPU profile is created on the server's start
func CreateCPUProfileFile(v *viper.Viper, fp *pflag.FlagSet, l log.Logger) {
	flag := fp.Lookup("cpuprofile")
	if flag == nil {
		return
	}

	f, err := os.Create(flag.Value.String())
	if err != nil {
		l.Log("could not create CPU profile: ", err)
	}
	defer f.Close()

	if err := pprof.StartCPUProfile(f); err != nil {
		l.Log("could not start CPU profile: ", err)
	}

	defer pprof.StopCPUProfile()
}

// Create CPUProfileFiles creates a memory profile of the server, its triggered by the optional flag memprofile
//
// the memory profile is created on the server's exit.
// this function should be used within the application.
func CreateMemoryProfileFile(v *viper.Viper, fp *pflag.FlagSet, l log.Logger) {
	flag := fp.Lookup("memprofile")
	if flag == nil {
		return
	}

	f, err := os.Create(flag.Value.String())
	if err != nil {
		l.Log("could not create memory profile: ", err)
	}

	defer f.Close()
	runtime.GC()
	if err := pprof.WriteHeapProfile(f); err != nil {
		l.Log("could not write memory profile: ", err)
	}
}

// ConfigureViper configures a Viper instances using the opinionated WebPA settings.  All WebPA servers should
// use this function.
//
// The flagSet is optional.  If supplied, it will be bound to the given Viper instance.  Additionally, if the
// flagSet has a FileFlagName flag, it will be used as the configuration name to hunt for instead of the
// application name.
func ConfigureViper(applicationName string, f *pflag.FlagSet, v *viper.Viper) (err error) {
	v.AddConfigPath(fmt.Sprintf("/etc/%s", applicationName))
	v.AddConfigPath(fmt.Sprintf("$HOME/.%s", applicationName))
	v.AddConfigPath(".")

	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	v.SetEnvPrefix(applicationName)
	v.AutomaticEnv()

	v.SetDefault("primary.name", applicationName)
	v.SetDefault("primary.address", DefaultPrimaryAddress)
	v.SetDefault("primary.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("alternate.name", fmt.Sprintf("%s.%s", applicationName, AlternateSuffix))

	v.SetDefault("health.name", fmt.Sprintf("%s.%s", applicationName, HealthSuffix))
	v.SetDefault("health.address", DefaultHealthAddress)
	v.SetDefault("health.logInterval", DefaultHealthLogInterval)
	v.SetDefault("health.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("pprof.name", fmt.Sprintf("%s.%s", applicationName, PprofSuffix))
	v.SetDefault("pprof.logConnectionState", DefaultLogConnectionState)

	v.SetDefault("metric.name", fmt.Sprintf("%s.%s", applicationName, MetricsSuffix))
	v.SetDefault("metric.address", DefaultMetricsAddress)

	v.SetDefault("project", DefaultProject)

	configName := applicationName
	if f != nil {
		if fileFlag := f.Lookup(FileFlagName); fileFlag != nil {
			// use the command-line to specify the base name of the file to be searched for
			configName = fileFlag.Value.String()
		}

		err = v.BindPFlags(f)
	}

	v.SetConfigName(configName)
	return
}

/*
Configure is a one-stop shopping function for preparing WebPA configuration.  This function
does not itself read in configuration from the Viper environment.  Typical usage is:

    var (
      f = pflag.NewFlagSet()
      v = viper.New()
    )

    if err := server.Configure("petasos", os.Args, f, v); err != nil {
      // deal with the error, possibly just exiting
    }

    // further customizations to the Viper instance can be done here

    if err := v.ReadInConfig(); err != nil {
      // more error handling
    }

Usage of this function is only necessary if custom configuration is needed.  Normally,
using New will suffice.
*/
func Configure(applicationName string, arguments []string, f *pflag.FlagSet, v *viper.Viper) (err error) {
	if f != nil {
		ConfigureFlagSet(applicationName, f)
		err = f.Parse(arguments)
		if err != nil {
			return
		}
	}

	err = ConfigureViper(applicationName, f, v)
	return
}

/*
Initialize handles the bootstrapping of the server code for a WebPA node.  It configures Viper,
reads configuration, and unmarshals the appropriate objects.  This function is typically all that's
needed to fully instantiate a WebPA server.  Typical usage:

    var (
      f = pflag.NewFlagSet()
      v = viper.New()

      // can customize both the FlagSet and the Viper before invoking New
      logger, registry, webPA, err = server.Initialize("petasos", os.Args, f, v)
    )

    if err != nil {
      // deal with the error, possibly just exiting
    }

Note that the FlagSet is optional but highly encouraged.  If not supplied, then no command-line binding
is done for the unmarshalled configuration.

This function always returns a logger, regardless of any errors.  This allows clients to use the returned
logger when reporting errors.  This function falls back to a logger that writes to os.Stdout if it cannot
create a logger from the Viper environment.
*/
func Initialize(applicationName string, arguments []string, f *pflag.FlagSet, v *viper.Viper, modules ...xmetrics.Module) (logger log.Logger, registry xmetrics.Registry, webPA *WebPA, err error) {
	defer func() {
		if err != nil {
			// never return a WebPA in the presence of an error, to
			// avoid an ambiguous API
			webPA = nil

			// Make sure there's at least a default logger for the caller to use
			logger = logging.DefaultLogger()
		}
	}()

	if err = Configure(applicationName, arguments, f, v); err != nil {
		return
	}

	if err = v.ReadInConfig(); err != nil {
		return
	}

	webPA = &WebPA{
		ApplicationName: applicationName,
	}

	err = v.Unmarshal(webPA)
	if err != nil {
		return
	}

	logger = logging.New(webPA.Log)
	logger.Log(level.Key(), level.InfoValue(), logging.MessageKey(), "initialized Viper environment", "configurationFile", v.ConfigFileUsed())

	if len(webPA.Metric.MetricsOptions.Namespace) == 0 {
		webPA.Metric.MetricsOptions.Namespace = applicationName
	}

	if len(webPA.Metric.MetricsOptions.Subsystem) == 0 {
		webPA.Metric.MetricsOptions.Subsystem = applicationName
	}

	webPA.Metric.MetricsOptions.Logger = logger
	registry, err = webPA.Metric.NewRegistry(modules...)
	if err != nil {
		return
	}

	// CreateCPUProfileFile(v, f, logger)
	return
}
<|endoftext|>


<commit_before>package service

import (
	"demo/cache"
	"demo/dao/mysql"
	"strconv"
	"time"

	"github.com/iiinsomnia/yiigo"
)

type Book struct {
	ID          int       `db:"id"`
	Title       string    `db:"title"`
	SubTitle    string    `db:"subtitle"`
	Author      string    `db:"author"`
	Version     string    `db:"version"`
	Price       string    `db:"price"`
	Publisher   string    `db:"publisher"`
	PublishDate string    `db:"publish_date"`
	CreatedAt   time.Time `db:"created_at"`
	UpdatedAt   time.Time `db:"updated_at"`
}

func GetBookById(id int) (yiigo.X, error) {
	book := &Book{}

	cache := cache.NewBookCache()

	cacheField := strconv.Itoa(id)
	ok := cache.GetBook(cacheField, book)

	if !ok {
		bookDao := mysql.NewBookDao()
		err := bookDao.GetById(id, book)

		if err != nil {
			if err.Error() != "not found" {
				return nil, err
			}

			return yiigo.X{}, nil
		}

		cache.SetBook(cacheField, book)
	}

	data := yiigo.X{
		"id":           book.ID,
		"title":        book.Title,
		"subtitle":     book.SubTitle,
		"author":       book.Author,
		"version":      book.Version,
		"price":        book.Price,
		"publisher":    book.Publisher,
		"publish_date": book.PublishDate,
	}

	return data, nil
}

func GetAllBooks() ([]yiigo.X, error) {
	books := []Book{}

	bookDao := mysql.NewBookDao()
	err := bookDao.GetAll(&books)

	if err != nil {
		return []yiigo.X{}, nil
	}

	data := formatBookList(books)

	return data, nil
}

func AddNewBook(data yiigo.X) (int64, error) {
	bookDao := mysql.NewBookDao()
	id, err := bookDao.AddNewRecord(data)

	return id, err
}

func UpdateBookById(id int, data yiigo.X) error {
	bookDao := mysql.NewBookDao()
	err := bookDao.UpdateById(id, data)

	return err
}

func DeleteBookById(id int) error {
	bookDao := mysql.NewBookDao()
	err := bookDao.DeleteById(id)

	return err
}

func formatBookList(books []Book) []yiigo.X {
	data := []yiigo.X{}

	for _, v := range books {
		item := yiigo.X{
			"id":           v.ID,
			"title":        v.Title,
			"subtitle":     v.SubTitle,
			"author":       v.Author,
			"version":      v.Version,
			"price":        v.Price,
			"publisher":    v.Publisher,
			"publish_date": v.PublishDate,
		}

		data = append(data, item)
	}

	return data
}
<commit_msg>update<commit_after>package service

import (
	"demo/cache"
	"demo/dao/mysql"
	"strconv"
	"time"

	"github.com/iiinsomnia/yiigo"
)

type Book struct {
	ID          int       `db:"id"`
	Title       string    `db:"title"`
	SubTitle    string    `db:"subtitle"`
	Author      string    `db:"author"`
	Version     string    `db:"version"`
	Price       string    `db:"price"`
	Publisher   string    `db:"publisher"`
	PublishDate string    `db:"publish_date"`
	CreatedAt   time.Time `db:"created_at"`
	UpdatedAt   time.Time `db:"updated_at"`
}

func GetBookById(id int) (yiigo.X, error) {
	book := &Book{}

	cache := cache.NewBookCache()

	cacheField := strconv.Itoa(id)
	ok := cache.GetBook(cacheField, book)

	if !ok {
		bookDao := mysql.NewBookDao()
		err := bookDao.GetById(id, book)

		if err != nil {
			if err.Error() != "not found" {
				return nil, err
			}

			return yiigo.X{}, nil
		}

		cache.SetBook(cacheField, book)
	}

	data := yiigo.X{
		"id":           book.ID,
		"title":        book.Title,
		"subtitle":     book.SubTitle,
		"author":       book.Author,
		"version":      book.Version,
		"price":        book.Price,
		"publisher":    book.Publisher,
		"publish_date": book.PublishDate,
	}

	return data, nil
}

func GetAllBooks() ([]yiigo.X, error) {
	books := []Book{}

	bookDao := mysql.NewBookDao()
	err := bookDao.GetAll(&books)

	if err != nil {
		return nil, err
	}

	data := formatBookList(books)

	return data, nil
}

func AddNewBook(data yiigo.X) (int64, error) {
	bookDao := mysql.NewBookDao()
	id, err := bookDao.AddNewRecord(data)

	return id, err
}

func UpdateBookById(id int, data yiigo.X) error {
	bookDao := mysql.NewBookDao()
	err := bookDao.UpdateById(id, data)

	return err
}

func DeleteBookById(id int) error {
	bookDao := mysql.NewBookDao()
	err := bookDao.DeleteById(id)

	return err
}

func formatBookList(books []Book) []yiigo.X {
	data := []yiigo.X{}

	for _, v := range books {
		item := yiigo.X{
			"id":           v.ID,
			"title":        v.Title,
			"subtitle":     v.SubTitle,
			"author":       v.Author,
			"version":      v.Version,
			"price":        v.Price,
			"publisher":    v.Publisher,
			"publish_date": v.PublishDate,
		}

		data = append(data, item)
	}

	return data
}
<|endoftext|>


<commit_before>package objx

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestSet(t *testing.T) {

	m := map[string]interface{}{"name": "Mat"}
	o := New(m)
	assert.Equal(t, o, o.Set("name", "Tyler"), "Set should chain")
	assert.Equal(t, "Tyler", m["name"])

}

func TestSetInDeepMap(t *testing.T) {

	m := map[string]interface{}{"one": map[string]interface{}{"two": map[string]interface{}{"three": 3}}}
	o := New(m)
	assert.Equal(t, o, o.Set("one.two.three", "THREE"), "Set should chain")
	assert.Equal(t, "THREE", m["one"].(map[string]interface{})["two"].(map[string]interface{})["three"])

}

func TestSetInArray(t *testing.T) {

	a := []string{"one", "two", "three"}
	o := New(a)

	o.Set(0, "1").Set(1, "2").Set(2, "3")

	assert.Equal(t, "1", a[0])
	assert.Equal(t, "2", a[1])
	assert.Equal(t, "3", a[2])

}
<commit_msg>added test for setting nil<commit_after>package objx

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestSet(t *testing.T) {

	m := map[string]interface{}{"name": "Mat"}
	o := New(m)
	assert.Equal(t, o, o.Set("name", "Tyler"), "Set should chain")
	assert.Equal(t, "Tyler", m["name"])

}

func TestSetInDeepMap(t *testing.T) {

	m := map[string]interface{}{"one": map[string]interface{}{"two": map[string]interface{}{"three": 3}}}
	o := New(m)
	assert.Equal(t, o, o.Set("one.two.three", "THREE"), "Set should chain")
	assert.Equal(t, "THREE", m["one"].(map[string]interface{})["two"].(map[string]interface{})["three"])

}

func TestSetNilInDeepMap(t *testing.T) {

	m := map[string]interface{}{"one": map[string]interface{}{"two": map[string]interface{}{"three": 3}}}
	o := New(m)
	assert.Equal(t, o, o.Set("one.two.three", nil), "Set should chain")
	assert.Equal(t, nil, m["one"].(map[string]interface{})["two"].(map[string]interface{})["three"])

}

func TestSetInArray(t *testing.T) {

	a := []string{"one", "two", "three"}
	o := New(a)

	o.Set(0, "1").Set(1, "2").Set(2, "3")

	assert.Equal(t, "1", a[0])
	assert.Equal(t, "2", a[1])
	assert.Equal(t, "3", a[2])

}
<|endoftext|>


<commit_before>package roaring

// to run just these tests: go test -run TestSetUtil*

import (
	"testing"
)

func TestSetUtilDifference(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1}
	nl := difference(data1, data2, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Difference is broken")
	}
	expectedresult = []uint16{5, 8, 11}
	nl = difference(data2, data1, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Difference is broken")
	}
}

func TestSetUtilUnion(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1, 2, 3, 4, 5, 8, 9, 11}
	nl := union2by2(data1, data2, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Union is broken")
	}
	nl = union2by2(data2, data1, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Union is broken")
	}
}

func TestSetUtilExclusiveUnion(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1, 5, 8, 11}
	nl := exclusiveUnion2by2(data1, data2, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Exclusive Union is broken")
	}
	nl = exclusiveUnion2by2(data2, data1, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Exclusive Union is broken")
	}
}

func TestSetUtilIntersection(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{2, 3, 4, 9}
	nl := intersection2by2(data1, data2, result)
	result = result[:nl]
	result = result[:len(expectedresult)]
	if !equal(result, expectedresult) {
		t.Errorf("Intersection is broken")
	}
	nl = intersection2by2(data2, data1, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Intersection is broken")
	}
	data1 = []uint16{4}

	data2 = make([]uint16, 10000)
	for i := range data2 {
		data2[i] = uint16(i)
	}
	result = make([]uint16, 0, len(data1)+len(data2))
	expectedresult = data1
	nl = intersection2by2(data1, data2, result)
	result = result[:nl]

	if !equal(result, expectedresult) {
		t.Errorf("Long intersection is broken")
	}
	nl = intersection2by2(data2, data1, result)
	result = result[:nl]
	if !equal(result, expectedresult) {
		t.Errorf("Long intersection is broken")
	}

}

func TestSetUtilIntersection2(t *testing.T) {
	data1 := []uint16{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}
	data2 := []uint16{0, 3, 6, 9, 12, 15, 18}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 6, 12, 18}
	nl := intersection2by2(data1, data2, result)
	result = result[:nl]
	result = result[:len(expectedresult)]
	if !equal(result, expectedresult) {
		t.Errorf("Intersection is broken")
	}
}

func TestSetUtilBinarySearch(t *testing.T) {
	data := make([]uint16, 256)
	for i := range data {
		data[i] = uint16(2 * i)
	}
	for i := 0; i < 2*len(data); i++ {
		key := uint16(i)
		loc := binarySearch(data, key)
		if (key & 1) == 0 {
			if loc != int(key)/2 {
				t.Errorf("binary search is broken")
			}
		} else {
			if loc != -int(key)/2-2 {
				t.Errorf("neg binary search is broken")
			}
		}
	}
}
<commit_msg>Replace convey with assert package for setutil_test.go<commit_after>package roaring

// to run just these tests: go test -run TestSetUtil*

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestSetUtilDifference(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1}
	nl := difference(data1, data2, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)

	expectedresult = []uint16{5, 8, 11}
	nl = difference(data2, data1, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)
}

func TestSetUtilUnion(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1, 2, 3, 4, 5, 8, 9, 11}
	nl := union2by2(data1, data2, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)

	nl = union2by2(data2, data1, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)
}

func TestSetUtilExclusiveUnion(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 1, 5, 8, 11}
	nl := exclusiveUnion2by2(data1, data2, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)

	nl = exclusiveUnion2by2(data2, data1, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)
}

func TestSetUtilIntersection(t *testing.T) {
	data1 := []uint16{0, 1, 2, 3, 4, 9}
	data2 := []uint16{2, 3, 4, 5, 8, 9, 11}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{2, 3, 4, 9}
	nl := intersection2by2(data1, data2, result)
	result = result[:nl]
	result = result[:len(expectedresult)]

	assert.Equal(t, expectedresult, result)

	nl = intersection2by2(data2, data1, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)

	data1 = []uint16{4}
	data2 = make([]uint16, 10000)

	for i := range data2 {
		data2[i] = uint16(i)
	}

	result = make([]uint16, 0, len(data1)+len(data2))
	expectedresult = data1
	nl = intersection2by2(data1, data2, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)

	nl = intersection2by2(data2, data1, result)
	result = result[:nl]

	assert.Equal(t, expectedresult, result)
}

func TestSetUtilIntersection2(t *testing.T) {
	data1 := []uint16{0, 2, 4, 6, 8, 10, 12, 14, 16, 18}
	data2 := []uint16{0, 3, 6, 9, 12, 15, 18}
	result := make([]uint16, 0, len(data1)+len(data2))
	expectedresult := []uint16{0, 6, 12, 18}
	nl := intersection2by2(data1, data2, result)
	result = result[:nl]
	result = result[:len(expectedresult)]

	assert.Equal(t, expectedresult, result)
}

func TestSetUtilBinarySearch(t *testing.T) {
	data := make([]uint16, 256)
	for i := range data {
		data[i] = uint16(2 * i)
	}
	for i := 0; i < 2*len(data); i++ {
		key := uint16(i)
		loc := binarySearch(data, key)
		if (key & 1) == 0 {
			assert.Equal(t, int(key)/2, loc)
		} else {
			assert.Equal(t, -int(key)/2-2, loc)
		}
	}
}
<|endoftext|>


<commit_before>// Licensed under GPL-2.0
package main

import (
	"flag"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
)

/*
#cgo LDFLAGS: -lcap
#include <sys/capability.h>
#include <errno.h>

static int dropAllCaps(void) {
    cap_t state;

    state = cap_init();
    if (!state) {
        cap_free(state);
    }

    if (cap_clear(state) < 0) {
        cap_free(state);
        return errno;
    }

    if (cap_set_proc(state) == -1) {
        cap_free(state);
        return errno;
    }

    cap_free(state);
    return 0;
}
*/
import "C"

var RequestLog *log.Logger
var allowUploads *bool

func dropAllCaps() (err error) {
	errno := C.dropAllCaps()
	if errno != 0 {
		return syscall.Errno(errno)
	}
	return
}

func reqHandler(w http.ResponseWriter, r *http.Request) {
	if !*allowUploads && r.Method == "POST" {
		http.Error(w, "Uploads not allowed", http.StatusForbidden)
		log.Println("Uploads not allowed")
		return
	}

	cwd, _ := os.Getwd()
	filePath := filepath.Join(cwd, r.URL.Path)
	if strings.HasPrefix(filePath, cwd) == false {
		log.Println("Trying to access dir outside of cwd")
		return
	}

	statInfo, statErr := os.Stat(filePath)
	if statErr != nil {
		http.NotFound(w, r)
		log.Println(statErr)
		return
	}

	RequestLog.Println(r.RemoteAddr, fmt.Sprintf("\"%s %s %s\"", r.Method, r.URL, r.Proto))
	if r.Method == "GET" && statInfo.IsDir() {
		files, readErr := ioutil.ReadDir(filePath)
		if readErr != nil {
			http.Error(w, "Directory read error", http.StatusInternalServerError)
			log.Println(readErr)
			return
		}

		w.Header().Set("Content-Type", "text/html")
		dirlistTitle := fmt.Sprintf("Directory listing for %s", r.URL.Path)

		fmt.Fprintf(w, "<!DOCTYPE html>\n"+
			"\n"+
			"\n"+
			"<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\">\n"+
			fmt.Sprintf("<title>%s</title>\n", dirlistTitle)+
			"\n\n"+
			fmt.Sprintf("
%s
\n", dirlistTitle))
		if *allowUploads {
			fmt.Fprintf(w, fmt.Sprintf("<form action=\"%s\" method=\"post\" enctype=\"multipart/form-data\">\n", r.URL.Path)+
				"Upload file: \n"+
				"<input type=\"file\" name=\"file\">\n"+
				"<input type=\"submit\" value=\"Submit\" />\n</form>\n")
		}
		fmt.Fprintf(w, "
\n
\n")

		listFmt := "
\n..\n
\n"
		fmt.Fprintf(w, listFmt, filepath.Dir(r.URL.Path))
		for _, f := range files {
			fFull := html.EscapeString(filepath.Join(r.URL.Path, f.Name()))
			fmt.Fprintf(w, "
\n%s\n
\n", fFull, f.Name())
		}

		fmt.Fprintf(w, "
\n
\n\n")

	} else if r.Method == "GET" {
		mimeType := mime.TypeByExtension(filepath.Ext(filePath))
		if mimeType == "" {
			mimeType = "application/octet-stream"
		}
		w.Header().Set("Content-Type", mimeType)
		w.Header().Set("Content-Length", strconv.FormatInt(statInfo.Size(), 10))

		file, err := os.Open(filePath)
		if err != nil {
			log.Println(err)
			return
		}
		defer file.Close()
		io.Copy(w, file)
	} else if r.Method == "POST" {
		if !statInfo.IsDir() {
			http.Error(w, "Cannot upload to non-directory file", http.StatusForbidden)
			log.Println("Cannot upload to non-directory file")
			return
		}

		r.ParseMultipartForm(15485760)
		formFile, handler, err := r.FormFile("file")
		if err != nil {
			fmt.Println(err)
			return
		}
		defer formFile.Close()

		outFilePath := filepath.Join(filePath, handler.Filename)
		if _, err := os.Stat("/path/to/whatever"); err == nil {
			http.Error(w, "File already exists", http.StatusForbidden)
			log.Println("File already exists")
			return
		}
		f, err := os.OpenFile(outFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			http.Error(w, "Cannot write file", http.StatusForbidden)
			fmt.Println(err)
			return
		}
		defer f.Close()

		io.Copy(f, formFile)
		http.Redirect(w, r, r.URL.Path, 302)
	} else {
		http.Error(w, "Unhandled request", http.StatusBadRequest)
	}
}

func main() {
	RequestLog = log.New(os.Stdout,
		"REQ: ",
		log.Ldate|log.Ltime)

	err := syscall.Chroot(".")
	if err != nil {
		panic(err)
	}
	err = dropAllCaps()
	if err != nil {
		panic(err)
	}

	allowUploads = flag.Bool("allow-uploads", false, "Allow uploading of files")
	listenPort := flag.Int("port", 8000, "Listen on port (default 8000)")
	flag.Parse()

	http.HandleFunc("/", reqHandler)
	fmt.Println("Listening on port", *listenPort)
	err = http.ListenAndServe(fmt.Sprintf(":%d", *listenPort), nil)
	if err != nil {
		panic(err)
	}
}
<commit_msg>Fix code style issues, error handling and not checking the right file when uploading<commit_after>// Licensed under GPL-2.0
package main

import (
	"flag"
	"fmt"
	"html"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
)

/*
#cgo LDFLAGS: -lcap
#include <sys/capability.h>
#include <errno.h>

static int dropAllCaps(void) {
    cap_t state;

    state = cap_init();
    if (!state) {
        cap_free(state);
    }

    if (cap_clear(state) < 0) {
        cap_free(state);
        return errno;
    }

    if (cap_set_proc(state) == -1) {
        cap_free(state);
        return errno;
    }

    cap_free(state);
    return 0;
}
*/
import "C"

var requestLog *log.Logger
var allowUploads *bool

func dropAllCaps() (err error) {
	errno := C.dropAllCaps()
	if errno != 0 {
		return syscall.Errno(errno)
	}
	return
}

func reqHandler(w http.ResponseWriter, r *http.Request) {
	if !*allowUploads && r.Method == "POST" {
		http.Error(w, "Uploads not allowed", http.StatusForbidden)
		log.Println("Uploads not allowed")
		return
	}

	cwd, _ := os.Getwd()
	filePath := filepath.Join(cwd, r.URL.Path)
	if strings.HasPrefix(filePath, cwd) == false {
		log.Println("Trying to access dir outside of cwd")
		return
	}

	statInfo, statErr := os.Stat(filePath)
	if statErr != nil {
		http.NotFound(w, r)
		log.Println(statErr)
		return
	}

	requestLog.Println(r.RemoteAddr, fmt.Sprintf("\"%s %s %s\"", r.Method, r.URL, r.Proto))
	if r.Method == "GET" && statInfo.IsDir() {
		files, readErr := ioutil.ReadDir(filePath)
		if readErr != nil {
			http.Error(w, "Directory read error", http.StatusInternalServerError)
			log.Println(readErr)
			return
		}

		w.Header().Set("Content-Type", "text/html")
		dirlistTitle := fmt.Sprintf("Directory listing for %s", r.URL.Path)

		fmt.Fprintf(w, "<!DOCTYPE html>\n"+
			"\n"+
			"\n"+
			"<meta http-equiv=\"Content-Type\" content=\"text/html; charset=utf-8\">\n"+
			fmt.Sprintf("<title>%s</title>\n", dirlistTitle)+
			"\n\n"+
			fmt.Sprintf("
%s
\n", dirlistTitle))
		if *allowUploads {
			fmt.Fprintf(w, fmt.Sprintf("<form action=\"%s\" method=\"post\" enctype=\"multipart/form-data\">\n", r.URL.Path)+
				"Upload file: \n"+
				"<input type=\"file\" name=\"file\">\n"+
				"<input type=\"submit\" value=\"Submit\" />\n</form>\n")
		}
		fmt.Fprintf(w, "
\n
\n")

		listFmt := "
\n..\n
\n"
		fmt.Fprintf(w, listFmt, filepath.Dir(r.URL.Path))
		for _, f := range files {
			fFull := html.EscapeString(filepath.Join(r.URL.Path, f.Name()))
			fmt.Fprintf(w, "
\n%s\n
\n", fFull, f.Name())
		}

		fmt.Fprintf(w, "
\n
\n\n")

	} else if r.Method == "GET" {
		mimeType := mime.TypeByExtension(filepath.Ext(filePath))
		if mimeType == "" {
			mimeType = "application/octet-stream"
		}
		w.Header().Set("Content-Type", mimeType)
		w.Header().Set("Content-Length", strconv.FormatInt(statInfo.Size(), 10))

		file, err := os.Open(filePath)
		if err != nil {
			log.Println(err)
			return
		}
		defer file.Close()

		_, err = io.Copy(w, file)
		if err != nil {
			log.Println(err)
			return
		}

	} else if r.Method == "POST" {
		if !statInfo.IsDir() {
			http.Error(w, "Cannot upload to non-directory file", http.StatusForbidden)
			log.Println("Cannot upload to non-directory file")
			return
		}

		err := r.ParseMultipartForm(15485760)
		if err != nil {
			log.Println(err)
			return
		}
		formFile, handler, err := r.FormFile("file")
		if err != nil {
			fmt.Println(err)
			return
		}
		defer formFile.Close()

		outFilePath := filepath.Join(filePath, handler.Filename)
		if _, err := os.Stat(outFilePath); err == nil {
			http.Error(w, "File already exists", http.StatusForbidden)
			log.Println("File already exists")
			return
		}
		f, err := os.OpenFile(outFilePath, os.O_WRONLY|os.O_CREATE, 0666)
		if err != nil {
			http.Error(w, "Cannot write file", http.StatusForbidden)
			fmt.Println(err)
			return
		}
		defer f.Close()

		_, err = io.Copy(f, formFile)
		if err != nil {
			log.Println(err)
			return
		}
		http.Redirect(w, r, r.URL.Path, 302)
	} else {
		http.Error(w, "Unhandled request", http.StatusBadRequest)
	}
}

func main() {
	requestLog = log.New(os.Stdout,
		"REQ: ",
		log.Ldate|log.Ltime)

	err := syscall.Chroot(".")
	if err != nil {
		panic(err)
	}
	err = dropAllCaps()
	if err != nil {
		panic(err)
	}

	allowUploads = flag.Bool("allow-uploads", false, "Allow uploading of files")
	listenPort := flag.Int("port", 8000, "Listen on port (default 8000)")
	flag.Parse()

	http.HandleFunc("/", reqHandler)
	fmt.Println("Listening on port", *listenPort)
	err = http.ListenAndServe(fmt.Sprintf(":%d", *listenPort), nil)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/mitchellh/colorstring"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

const stateKey string = "spotify_auth_state"
const scope string = "playlist-modify-public playlist-modify-private " +
	"playlist-read-private"

type SpotifyAuth struct {
	ClientId     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	TokenFile    string `json:"token_file"`
	listen       string
}

func (auth *SpotifyAuth) ListenURL() string {
	if strings.HasPrefix(auth.listen, ":") {
		return "http://localhost" + auth.listen
	}
	return "http://" + auth.listen
}

func (auth *SpotifyAuth) CallbackURL() string {
	return auth.ListenURL() + "/callback"
}

func (auth *SpotifyAuth) authHeader() string {
	data := auth.ClientId + ":" + auth.ClientSecret
	return "Basic " + base64.StdEncoding.EncodeToString(
		[]byte(data))
}

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

func (auth *SpotifyAuth) Login(w http.ResponseWriter, r *http.Request) {
	state, err := base64Rand(32)
	if err != nil {
		http.Error(w, "Failed to generate state value", 400)
		return
	}
	http.SetCookie(w, &http.Cookie{
		Name:  stateKey,
		Value: state,
	})
	params := url.Values{
		"response_type": {"code"},
		"client_id":     {auth.ClientId},
		"scope":         {scope},
		"redirect_uri":  {auth.CallbackURL()},
		"state":         {state},
	}
	url := "https://accounts.spotify.com/authorize?" + params.Encode()
	http.Redirect(w, r, url, http.StatusFound)
}

func (auth *SpotifyAuth) getToken(code []string) (*Spotify, error) {
	formData := url.Values{
		"code":          code,
		"redirect_uri":  {auth.CallbackURL()},
		"grant_type":    {"authorization_code"},
		"client_id":     {auth.ClientId},
		"client_secret": {auth.ClientSecret},
	}
	url := "https://accounts.spotify.com/api/token"
	resp, err := http.PostForm(url, formData)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var token Spotify
	if err := json.Unmarshal(body, &token); err != nil {
		return nil, err
	}
	token.Auth = *auth
	return &token, nil
}

func (auth *SpotifyAuth) Callback(w http.ResponseWriter, r *http.Request) {
	params := r.URL.Query()
	state := params["state"]
	cookie, _ := r.Cookie(stateKey)

	if state == nil || cookie == nil || cookie.Value != state[0] {
		http.Error(w, "Could not validate request", 400)
		return
	}

	code, exists := params["code"]
	if !exists {
		http.Error(w, "Missing required query parameter: code", 400)
		return
	}
	token, err := auth.getToken(code)
	if err != nil {
		http.Error(w, "Failed to retrieve token from Spotify", 400)
		return
	}
	err = token.Save(auth.TokenFile)
	if err != nil {
		fmt.Printf("Failed to save token: %s\n", err)
		http.Error(w, "Failed to save token", 400)
		return
	}
	fmt.Printf(colorstring.Color("Wrote file: [green]%s[reset]\n"),
		auth.TokenFile)
	fmt.Fprintf(w, "Success! Wrote token file to %s", auth.TokenFile)
}

func (auth *SpotifyAuth) Serve() {
	http.HandleFunc("/login", auth.Login)
	http.HandleFunc("/callback", auth.Callback)
	fmt.Printf(colorstring.Color(
		"Visit [green]%s/login[reset] to authenticate "+
			"with Spotify.\n"), auth.ListenURL())
	http.ListenAndServe(auth.listen, nil)
}
<commit_msg>Use URL method when building URLs<commit_after>package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"github.com/mitchellh/colorstring"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

const defaultAuthURL string = "https://accounts.spotify.com"
const stateKey string = "spotify_auth_state"
const scope string = "playlist-modify-public playlist-modify-private " +
	"playlist-read-private"

type SpotifyAuth struct {
	ClientId     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	TokenFile    string `json:"token_file"`
	listen       string
	url          string
}

func (auth *SpotifyAuth) URL() string {
	if auth.url == "" {
		return defaultAuthURL
	}
	return auth.url
}

func (auth *SpotifyAuth) ListenURL() string {
	if strings.HasPrefix(auth.listen, ":") {
		return "http://localhost" + auth.listen
	}
	return "http://" + auth.listen
}

func (auth *SpotifyAuth) CallbackURL() string {
	return auth.ListenURL() + "/callback"
}

func (auth *SpotifyAuth) authHeader() string {
	data := auth.ClientId + ":" + auth.ClientSecret
	return "Basic " + base64.StdEncoding.EncodeToString(
		[]byte(data))
}

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

func (auth *SpotifyAuth) Login(w http.ResponseWriter, r *http.Request) {
	state, err := base64Rand(32)
	if err != nil {
		http.Error(w, "Failed to generate state value", 400)
		return
	}
	http.SetCookie(w, &http.Cookie{
		Name:  stateKey,
		Value: state,
	})
	params := url.Values{
		"response_type": {"code"},
		"client_id":     {auth.ClientId},
		"scope":         {scope},
		"redirect_uri":  {auth.CallbackURL()},
		"state":         {state},
	}
	url := auth.URL() + "/authorize?" + params.Encode()
	http.Redirect(w, r, url, http.StatusFound)
}

func (auth *SpotifyAuth) getToken(code []string) (*Spotify, error) {
	formData := url.Values{
		"code":          code,
		"redirect_uri":  {auth.CallbackURL()},
		"grant_type":    {"authorization_code"},
		"client_id":     {auth.ClientId},
		"client_secret": {auth.ClientSecret},
	}
	url := auth.URL() + "/api/token"
	resp, err := http.PostForm(url, formData)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var token Spotify
	if err := json.Unmarshal(body, &token); err != nil {
		return nil, err
	}
	token.Auth = *auth
	return &token, nil
}

func (auth *SpotifyAuth) Callback(w http.ResponseWriter, r *http.Request) {
	params := r.URL.Query()
	state := params["state"]
	cookie, _ := r.Cookie(stateKey)

	if state == nil || cookie == nil || cookie.Value != state[0] {
		http.Error(w, "Could not validate request", 400)
		return
	}

	code, exists := params["code"]
	if !exists {
		http.Error(w, "Missing required query parameter: code", 400)
		return
	}
	token, err := auth.getToken(code)
	if err != nil {
		http.Error(w, "Failed to retrieve token from Spotify", 400)
		return
	}
	err = token.Save(auth.TokenFile)
	if err != nil {
		fmt.Printf("Failed to save token: %s\n", err)
		http.Error(w, "Failed to save token", 400)
		return
	}
	fmt.Printf(colorstring.Color("Wrote file: [green]%s[reset]\n"),
		auth.TokenFile)
	fmt.Fprintf(w, "Success! Wrote token file to %s", auth.TokenFile)
}

func (auth *SpotifyAuth) Serve() {
	http.HandleFunc("/login", auth.Login)
	http.HandleFunc("/callback", auth.Callback)
	fmt.Printf(colorstring.Color(
		"Visit [green]%s/login[reset] to authenticate "+
			"with Spotify.\n"), auth.ListenURL())
	http.ListenAndServe(auth.listen, nil)
}
<|endoftext|>


<commit_before>// Copyright 2011 Bobby Powers. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
package seshcookie

import (
	"http"
	"io"
	"os"
	"log"
	"time"
	"gob"
	"bytes"
	"sync"
	"sync/atomic"
	"hash"
	"crypto/sha1"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/hmac"
	"crypto/subtle"
	"encoding/base64"
)

// we want 16 byte blocks, for AES-128
const blockSize = 16

var (
	// if you don't need multiple independent seshcookie
	// instances, you can use this RequestSessions instance to
	// manage & access your sessions.  Simply use it as the final
	// parameter in your call to seshcookie.NewSessionHandler, and
	// whenever you want to access the current session from an
	// embedded http.Handler you can simply call:
	//
	//     seshcookie.Session.Get(req)
	Session = &RequestSessions{HttpOnly:true}

	// Hash validation of the decrypted cookie failed. Most likely
	// the session was encoded with a different cookie than we're
	// using to decode it, but its possible the client (or someone
	// else) tried to modify the session.
	HashError = os.NewError("Hash validation failed")

	// The cookie is too short, so we must exit decoding early.
	LenError = os.NewError("Bad cookie length")
)

type sessionResponseWriter struct {
	http.ResponseWriter
	h   *SessionHandler
	req *http.Request
	// int32 so we can use the sync/atomic functions on it
	wroteHeader int32
}

type SessionHandler struct {
	http.Handler
	CookieName string // name of the cookie to store our session in
	CookiePath string // resource path the cookie is valid for
	RS         *RequestSessions
	encKey     []byte
	hmacKey    []byte
}

type RequestSessions struct {
	HttpOnly bool // don't allow javascript to access cookie
	Secure   bool // only send session over HTTPS
	lk       sync.Mutex
	m        map[*http.Request]map[string]interface{}
	// stores a hash of the serialized session (the gob) that we
	// received with the start of the request.  Before setting a
	// cookie for the reply, check to see if the session has
	// actually changed.  If it hasn't, then we don't need to send
	// a new cookie.
	hm map[*http.Request][]byte
}

func (rs *RequestSessions) Get(req *http.Request) map[string]interface{} {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		log.Print("seshcookie: warning! trying to get session " +
			"data for unknown request. Perhaps your handler " +
			"isn't wrapped by a SessionHandler?")
		return nil
	}

	return rs.m[req]
}

func (rs *RequestSessions) getHash(req *http.Request) []byte {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.hm == nil {
		return nil
	}

	return rs.hm[req]
}

func (rs *RequestSessions) Set(req *http.Request, val map[string]interface{}, gobHash []byte) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		rs.m = map[*http.Request]map[string]interface{}{}
		rs.hm = map[*http.Request][]byte{}
	}

	rs.m[req] = val
	rs.hm[req] = gobHash
}

func (rs *RequestSessions) Clear(req *http.Request) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	rs.m[req] = nil, false
}

func encodeGob(obj interface{}) ([]byte, os.Error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(obj)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decodeGob(encoded []byte) (map[string]interface{}, os.Error) {
	buf := bytes.NewBuffer(encoded)
	dec := gob.NewDecoder(buf)
	var out map[string]interface{}
	err := dec.Decode(&out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// encode uses the given block cipher (in CTR mode) to encrypt the
// data, along with a hash, returning the iv and the ciphertext. What
// is returned looks like:
//
//   encrypted(salt + sessionData) + iv + hmac
//
func encode(block cipher.Block, hmac hash.Hash, data []byte) ([]byte, os.Error) {

	buf := bytes.NewBuffer(nil)

	salt := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, err
	}
	buf.Write(salt)
	buf.Write(data)

	session := buf.Bytes()

	iv := make([]byte, block.BlockSize())
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	buf.Write(iv)
	hmac.Write(buf.Bytes())
	buf.Write(hmac.Sum())

	return buf.Bytes(), nil
}

func encodeCookie(content interface{}, encKey, hmacKey []byte) (string, []byte, os.Error) {
	encodedGob, err := encodeGob(content)
	if err != nil {
		return "", nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(encodedGob)

	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return "", nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)

	sessionBytes, err := encode(aesCipher, hmacHash, encodedGob)
	if err != nil {
		return "", nil, err
	}

	return base64.StdEncoding.EncodeToString(sessionBytes), gobHash.Sum(), nil
}

// decode uses the given block cipher (in CTR mode) to decrypt the
// data, and validate the hash.  If hash validation fails, an error is
// returned.
func decode(block cipher.Block, hmac hash.Hash, ciphertext []byte) ([]byte, os.Error) {
	if len(ciphertext) < 2*block.BlockSize()+hmac.Size() {
		return nil, LenError
	}

	receivedHmac := ciphertext[len(ciphertext)-hmac.Size():]
	ciphertext = ciphertext[:len(ciphertext)-hmac.Size()]

	hmac.Write(ciphertext)
	if subtle.ConstantTimeCompare(hmac.Sum(), receivedHmac) != 1 {
		return nil, HashError
	}

	// split the iv and session bytes
	iv := ciphertext[len(ciphertext)-block.BlockSize():]
	session := ciphertext[:len(ciphertext)-block.BlockSize()]

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	// skip past the iv
	session = session[block.BlockSize():]

	return session, nil
}

func decodeCookie(encoded string, encKey, hmacKey []byte) (map[string]interface{}, []byte, os.Error) {
	sessionBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return nil, nil, err
	}
	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)
	gobBytes, err := decode(aesCipher, hmacHash, sessionBytes)
	if err != nil {
		return nil, nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(gobBytes)

	session, err := decodeGob(gobBytes)
	if err != nil {
		log.Printf("decodeGob: %s\n", err)
		return nil, nil, err
	}
	return session, gobHash.Sum(), nil
}

func (s sessionResponseWriter) WriteHeader(code int) {
	if atomic.AddInt32(&s.wroteHeader, 1) == 1 {
		origCookie, err := s.req.Cookie(s.h.CookieName)
		var origCookieVal string
		if err != nil {
			origCookieVal = ""
		} else {
			origCookieVal = origCookie.Value
		}

		session := s.h.RS.Get(s.req)
		if len(session) == 0 {
			// if we have an empty session, but the
			// request didn't start out that way, we
			// assume the user wants us to clear the
			// session
			if origCookieVal != "" {
				log.Println("clearing cookie")
				var cookie http.Cookie
				cookie.Name = s.h.CookieName
				cookie.Value = ""
				cookie.Path = "/"
				// a cookie is expired by setting it
				// with an expiration time in the past
				cookie.Expires = *time.SecondsToUTC(0)
				http.SetCookie(s, &cookie)
			} else {
				log.Println("not setting empty cookie")
			}
			goto write
		}
		encoded, gobHash, err := encodeCookie(session, s.h.encKey, s.h.hmacKey)
		if err != nil {
			log.Printf("createCookie: %s\n", err)
			goto write
		}

		if bytes.Equal(gobHash, s.h.RS.getHash(s.req)) {
			//log.Println("not re-setting identical cookie")
			goto write
		}

		var cookie http.Cookie
		cookie.Name = s.h.CookieName
		cookie.Value = encoded
		cookie.Path = s.h.CookiePath
		cookie.HttpOnly = s.h.RS.HttpOnly
		cookie.Secure = s.h.RS.Secure
		http.SetCookie(s, &cookie)
	}
write:
	s.ResponseWriter.WriteHeader(code)
}

func (h *SessionHandler) getCookieSession(req *http.Request) (map[string]interface{}, []byte) {
	cookie, err := req.Cookie(h.CookieName)
	if err != nil {
		log.Printf("getCookieSesh: '%#v' not found\n",
			h.CookieName)
		return map[string]interface{}{}, nil
	}
	session, gobHash, err := decodeCookie(cookie.Value, h.encKey, h.hmacKey)
	if err != nil {
		log.Printf("decodeCookie: %s\n", err)
		return map[string]interface{}{}, nil
	}

	return session, gobHash
}

func (h *SessionHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// get our session a little early, so that we can add our
	// authentication information to it if we get some
	session, gobHash := h.getCookieSession(req)

	h.RS.Set(req, session, gobHash)

	sessionWriter := sessionResponseWriter{rw, h, req, 0}
	h.Handler.ServeHTTP(sessionWriter, req)

	h.RS.Clear(req)
}

func NewSessionHandler(handler http.Handler, key string, rs *RequestSessions) *SessionHandler {
	// sha1 sums are 20 bytes long.  we use the first 16 bytes as
	// the aes key.
	encHash := sha1.New()
	encHash.Write([]byte(key))
	encHash.Write([]byte("-encryption"))
	hmacHash := sha1.New()
	hmacHash.Write([]byte(key))
	hmacHash.Write([]byte("-hmac"))

	// if the user hasn't specified a session handler, use the
	// package's default one
	if rs == nil {
		rs = Session
	}

	return &SessionHandler{
		Handler:    handler,
		CookieName: "session",
		CookiePath: "/",
		RS:         rs,
		encKey:     encHash.Sum()[:blockSize],
		hmacKey:    hmacHash.Sum()[:blockSize],
	}
}
<commit_msg>remove some debug logging<commit_after>// Copyright 2011 Bobby Powers. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.
package seshcookie

import (
	"http"
	"io"
	"os"
	"log"
	"time"
	"gob"
	"bytes"
	"sync"
	"sync/atomic"
	"hash"
	"crypto/sha1"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/hmac"
	"crypto/subtle"
	"encoding/base64"
)

// we want 16 byte blocks, for AES-128
const blockSize = 16

var (
	// if you don't need multiple independent seshcookie
	// instances, you can use this RequestSessions instance to
	// manage & access your sessions.  Simply use it as the final
	// parameter in your call to seshcookie.NewSessionHandler, and
	// whenever you want to access the current session from an
	// embedded http.Handler you can simply call:
	//
	//     seshcookie.Session.Get(req)
	Session = &RequestSessions{HttpOnly:true}

	// Hash validation of the decrypted cookie failed. Most likely
	// the session was encoded with a different cookie than we're
	// using to decode it, but its possible the client (or someone
	// else) tried to modify the session.
	HashError = os.NewError("Hash validation failed")

	// The cookie is too short, so we must exit decoding early.
	LenError = os.NewError("Bad cookie length")
)

type sessionResponseWriter struct {
	http.ResponseWriter
	h   *SessionHandler
	req *http.Request
	// int32 so we can use the sync/atomic functions on it
	wroteHeader int32
}

type SessionHandler struct {
	http.Handler
	CookieName string // name of the cookie to store our session in
	CookiePath string // resource path the cookie is valid for
	RS         *RequestSessions
	encKey     []byte
	hmacKey    []byte
}

type RequestSessions struct {
	HttpOnly bool // don't allow javascript to access cookie
	Secure   bool // only send session over HTTPS
	lk       sync.Mutex
	m        map[*http.Request]map[string]interface{}
	// stores a hash of the serialized session (the gob) that we
	// received with the start of the request.  Before setting a
	// cookie for the reply, check to see if the session has
	// actually changed.  If it hasn't, then we don't need to send
	// a new cookie.
	hm map[*http.Request][]byte
}

func (rs *RequestSessions) Get(req *http.Request) map[string]interface{} {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		log.Print("seshcookie: warning! trying to get session " +
			"data for unknown request. Perhaps your handler " +
			"isn't wrapped by a SessionHandler?")
		return nil
	}

	return rs.m[req]
}

func (rs *RequestSessions) getHash(req *http.Request) []byte {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.hm == nil {
		return nil
	}

	return rs.hm[req]
}

func (rs *RequestSessions) Set(req *http.Request, val map[string]interface{}, gobHash []byte) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	if rs.m == nil {
		rs.m = map[*http.Request]map[string]interface{}{}
		rs.hm = map[*http.Request][]byte{}
	}

	rs.m[req] = val
	rs.hm[req] = gobHash
}

func (rs *RequestSessions) Clear(req *http.Request) {
	rs.lk.Lock()
	defer rs.lk.Unlock()

	rs.m[req] = nil, false
}

func encodeGob(obj interface{}) ([]byte, os.Error) {
	buf := bytes.NewBuffer(nil)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(obj)
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func decodeGob(encoded []byte) (map[string]interface{}, os.Error) {
	buf := bytes.NewBuffer(encoded)
	dec := gob.NewDecoder(buf)
	var out map[string]interface{}
	err := dec.Decode(&out)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// encode uses the given block cipher (in CTR mode) to encrypt the
// data, along with a hash, returning the iv and the ciphertext. What
// is returned looks like:
//
//   encrypted(salt + sessionData) + iv + hmac
//
func encode(block cipher.Block, hmac hash.Hash, data []byte) ([]byte, os.Error) {

	buf := bytes.NewBuffer(nil)

	salt := make([]byte, block.BlockSize())
	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
		return nil, err
	}
	buf.Write(salt)
	buf.Write(data)

	session := buf.Bytes()

	iv := make([]byte, block.BlockSize())
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	buf.Write(iv)
	hmac.Write(buf.Bytes())
	buf.Write(hmac.Sum())

	return buf.Bytes(), nil
}

func encodeCookie(content interface{}, encKey, hmacKey []byte) (string, []byte, os.Error) {
	encodedGob, err := encodeGob(content)
	if err != nil {
		return "", nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(encodedGob)

	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return "", nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)

	sessionBytes, err := encode(aesCipher, hmacHash, encodedGob)
	if err != nil {
		return "", nil, err
	}

	return base64.StdEncoding.EncodeToString(sessionBytes), gobHash.Sum(), nil
}

// decode uses the given block cipher (in CTR mode) to decrypt the
// data, and validate the hash.  If hash validation fails, an error is
// returned.
func decode(block cipher.Block, hmac hash.Hash, ciphertext []byte) ([]byte, os.Error) {
	if len(ciphertext) < 2*block.BlockSize()+hmac.Size() {
		return nil, LenError
	}

	receivedHmac := ciphertext[len(ciphertext)-hmac.Size():]
	ciphertext = ciphertext[:len(ciphertext)-hmac.Size()]

	hmac.Write(ciphertext)
	if subtle.ConstantTimeCompare(hmac.Sum(), receivedHmac) != 1 {
		return nil, HashError
	}

	// split the iv and session bytes
	iv := ciphertext[len(ciphertext)-block.BlockSize():]
	session := ciphertext[:len(ciphertext)-block.BlockSize()]

	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(session, session)

	// skip past the iv
	session = session[block.BlockSize():]

	return session, nil
}

func decodeCookie(encoded string, encKey, hmacKey []byte) (map[string]interface{}, []byte, os.Error) {
	sessionBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return nil, nil, err
	}
	aesCipher, err := aes.NewCipher(encKey)
	if err != nil {
		return nil, nil, err
	}

	hmacHash := hmac.NewSHA256(hmacKey)
	gobBytes, err := decode(aesCipher, hmacHash, sessionBytes)
	if err != nil {
		return nil, nil, err
	}

	gobHash := sha1.New()
	gobHash.Write(gobBytes)

	session, err := decodeGob(gobBytes)
	if err != nil {
		log.Printf("decodeGob: %s\n", err)
		return nil, nil, err
	}
	return session, gobHash.Sum(), nil
}

func (s sessionResponseWriter) WriteHeader(code int) {
	if atomic.AddInt32(&s.wroteHeader, 1) == 1 {
		origCookie, err := s.req.Cookie(s.h.CookieName)
		var origCookieVal string
		if err != nil {
			origCookieVal = ""
		} else {
			origCookieVal = origCookie.Value
		}

		session := s.h.RS.Get(s.req)
		if len(session) == 0 {
			// if we have an empty session, but the
			// request didn't start out that way, we
			// assume the user wants us to clear the
			// session
			if origCookieVal != "" {
				log.Println("clearing cookie")
				var cookie http.Cookie
				cookie.Name = s.h.CookieName
				cookie.Value = ""
				cookie.Path = "/"
				// a cookie is expired by setting it
				// with an expiration time in the past
				cookie.Expires = *time.SecondsToUTC(0)
				http.SetCookie(s, &cookie)
			} else {
				log.Println("not setting empty cookie")
			}
			goto write
		}
		encoded, gobHash, err := encodeCookie(session, s.h.encKey, s.h.hmacKey)
		if err != nil {
			log.Printf("createCookie: %s\n", err)
			goto write
		}

		if bytes.Equal(gobHash, s.h.RS.getHash(s.req)) {
			//log.Println("not re-setting identical cookie")
			goto write
		}

		var cookie http.Cookie
		cookie.Name = s.h.CookieName
		cookie.Value = encoded
		cookie.Path = s.h.CookiePath
		cookie.HttpOnly = s.h.RS.HttpOnly
		cookie.Secure = s.h.RS.Secure
		http.SetCookie(s, &cookie)
	}
write:
	s.ResponseWriter.WriteHeader(code)
}

func (h *SessionHandler) getCookieSession(req *http.Request) (map[string]interface{}, []byte) {
	cookie, err := req.Cookie(h.CookieName)
	if err != nil {
		//log.Printf("getCookieSesh: '%#v' not found\n",
		//	h.CookieName)
		return map[string]interface{}{}, nil
	}
	session, gobHash, err := decodeCookie(cookie.Value, h.encKey, h.hmacKey)
	if err != nil {
		log.Printf("decodeCookie: %s\n", err)
		return map[string]interface{}{}, nil
	}

	return session, gobHash
}

func (h *SessionHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// get our session a little early, so that we can add our
	// authentication information to it if we get some
	session, gobHash := h.getCookieSession(req)

	h.RS.Set(req, session, gobHash)

	sessionWriter := sessionResponseWriter{rw, h, req, 0}
	h.Handler.ServeHTTP(sessionWriter, req)

	h.RS.Clear(req)
}

func NewSessionHandler(handler http.Handler, key string, rs *RequestSessions) *SessionHandler {
	// sha1 sums are 20 bytes long.  we use the first 16 bytes as
	// the aes key.
	encHash := sha1.New()
	encHash.Write([]byte(key))
	encHash.Write([]byte("-encryption"))
	hmacHash := sha1.New()
	hmacHash.Write([]byte(key))
	hmacHash.Write([]byte("-hmac"))

	// if the user hasn't specified a session handler, use the
	// package's default one
	if rs == nil {
		rs = Session
	}

	return &SessionHandler{
		Handler:    handler,
		CookieName: "session",
		CookiePath: "/",
		RS:         rs,
		encKey:     encHash.Sum()[:blockSize],
		hmacKey:    hmacHash.Sum()[:blockSize],
	}
}
<|endoftext|>


<commit_before>package session

import (
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"runtime"
	"testing"

	"github.com/gin-gonic/gin"
)

var (
	listenAddr    = ":8080"
	connectAddr   = "http://127.0.0.1:8080/"
	connectUrl, _ = url.Parse(connectAddr)
)

var (
	session Session
)

func TestMain(m *testing.M) {
	// Чтобы тестирования не началось случайно до того, как запустится сервер Gin.
	runtime.GOMAXPROCS(1)
	// Тесты запускаются в отдельной от сервера горутине.
	// Эта горутина стартует, только когда заблокируется текущая.
	// А именно, когда заблокируется сервер Gin в ожидании подключений.
	go func() {
		os.Exit(m.Run())
	}()
	// Довольно костыльное решение, но Gin не предоставляет нормального
	// функционала для тестирования.

	hf := func(c *gin.Context) {
		session = FromContext(c)
	}
	r := gin.Default()
	r.GET("/", hf)
	r.Run(listenAddr)
}

// Отправляет с локального клиента запрос на сервер и возвращает
// полученную для него сессию.
func retrieveSession(client *http.Client) Session {
	if client.Jar == nil {
		jar, _ := cookiejar.New(nil)
		client.Jar = jar
	}

	resp, err := client.Get(connectAddr)
	if err != nil {
		panic(err)
	}
	resp.Body.Close()
	client.Jar.SetCookies(connectUrl, resp.Cookies())

	return session
}

func TestCookieStoring(t *testing.T) {

	client := &http.Client{}

	sess1 := retrieveSession(client)
	sess2 := retrieveSession(client)

	client.Jar = nil
	sess3 := retrieveSession(client)

	// Проверки
	if sess1 == nil || sess3 == nil {
		t.Fatal("Сессия должна быть создана, если не найден ключ сессии в куки.")
	}
	if sess2 == nil {
		t.Fatal("Сессия должна быть получена, если ключ найден в куки.")
	}

	if sess1.ID() != sess2.ID() {
		t.Error("Сессии от одинаковых куки должны быть одинаковыми.")
	}
	if sess1.ID() == sess3.ID() || sess2.ID() == sess3.ID() {
		t.Error("Сессии от разных куки должны быть разными.")
	}

}

// Проверяет равенство ключа key в сессии s значению expected.
// Иначе сообщает о нарушении условия cond.
func assertIntKey(t *testing.T, cond string, s Session, key string, expected int) {
	actual := s.GetInt(key)
	if actual != expected {
		t.Errorf("%s Ключ %q: ожидалось %v, получено %v.", cond, key, expected, actual)
	}
}

func TestDataStoring(t *testing.T) {
	client := &http.Client{}

	ints := map[string]int{
		"int1":       12,
		"second int": 34,
	}
	missingInt := "missInt"

	// TODO: Тесты для всех поддерживаемых типов.
	// floats := map[string]float64{
	// 	"float1":       12.04,
	// 	"second float": 3.14,
	// }
	// missingFloat := "missFloat"
	//
	// bytes := map[string][]byte{
	// 	"slice1":       {1, 2, 3, 12},
	// 	"second slice": {},
	// }
	// missingBytes := "missBytes"
	//
	// strings := map[string]string{
	// 	"hello":    "world",
	// 	"username": "onotole",
	// }
	// missingString := "missString"

	sess := retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, val := range ints {
		assertIntKey(t, "Пустая сессия не должна содержать значений int.", sess, key, 0)
		sess.SetInt(key, val)
	}

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, val := range ints {
		assertIntKey(t, "Сессия должна сохранять значения int.", sess, key, val)
	}
	assertIntKey(t, "Сессия должна возвращать 0 для незаданных значений.", sess, missingInt, 0)

	// Обнуление куки.
	client.Jar = nil

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	sess.SetInt(missingInt, 12)

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, _ := range ints {
		assertIntKey(t, "Сессия не должна возвращать значения чужой сессии.", sess, key, 0)
	}
	assertIntKey(t, "Сессия должна сохранять значения int.", sess, missingInt, 12)

}
<commit_msg>Test nil case<commit_after>package session

import (
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"runtime"
	"testing"

	"github.com/gin-gonic/gin"
)

var (
	listenAddr    = ":8080"
	connectAddr   = "http://127.0.0.1:8080/"
	connectUrl, _ = url.Parse(connectAddr)
)

var (
	session Session
)

func TestMain(m *testing.M) {
	// Чтобы тестирования не началось случайно до того, как запустится сервер Gin.
	runtime.GOMAXPROCS(1)
	// Тесты запускаются в отдельной от сервера горутине.
	// Эта горутина стартует, только когда заблокируется текущая.
	// А именно, когда заблокируется сервер Gin в ожидании подключений.
	go func() {
		os.Exit(m.Run())
	}()
	// Довольно костыльное решение, но Gin не предоставляет нормального
	// функционала для тестирования.

	hf := func(c *gin.Context) {
		session = FromContext(c)
	}
	r := gin.Default()
	r.GET("/", hf)
	r.Run(listenAddr)
}

// Отправляет с локального клиента запрос на сервер и возвращает
// полученную для него сессию.
func retrieveSession(client *http.Client) Session {
	if client.Jar == nil {
		jar, _ := cookiejar.New(nil)
		client.Jar = jar
	}

	resp, err := client.Get(connectAddr)
	if err != nil {
		panic(err)
	}
	resp.Body.Close()
	client.Jar.SetCookies(connectUrl, resp.Cookies())

	return session
}

func TestCookieStoring(t *testing.T) {

	sess := FromContext(nil)
	if sess != nil {
		t.Error("FromContext должна возвращать нулевую сессию для нулевого контекста.")
	}

	client := &http.Client{}

	sess1 := retrieveSession(client)
	sess2 := retrieveSession(client)

	client.Jar = nil
	sess3 := retrieveSession(client)

	// Проверки
	if sess1 == nil || sess3 == nil {
		t.Fatal("Сессия должна быть создана, если не найден ключ сессии в куки.")
	}
	if sess2 == nil {
		t.Fatal("Сессия должна быть получена, если ключ найден в куки.")
	}

	if sess1.ID() != sess2.ID() {
		t.Error("Сессии от одинаковых куки должны быть одинаковыми.")
	}
	if sess1.ID() == sess3.ID() || sess2.ID() == sess3.ID() {
		t.Error("Сессии от разных куки должны быть разными.")
	}

}

// Проверяет равенство ключа key в сессии s значению expected.
// Иначе сообщает о нарушении условия cond.
func assertIntKey(t *testing.T, cond string, s Session, key string, expected int) {
	actual := s.GetInt(key)
	if actual != expected {
		t.Errorf("%s Ключ %q: ожидалось %v, получено %v.", cond, key, expected, actual)
	}
}

func TestDataStoring(t *testing.T) {
	client := &http.Client{}

	ints := map[string]int{
		"int1":       12,
		"second int": 34,
	}
	missingInt := "missInt"

	// TODO: Тесты для всех поддерживаемых типов.
	// floats := map[string]float64{
	// 	"float1":       12.04,
	// 	"second float": 3.14,
	// }
	// missingFloat := "missFloat"
	//
	// bytes := map[string][]byte{
	// 	"slice1":       {1, 2, 3, 12},
	// 	"second slice": {},
	// }
	// missingBytes := "missBytes"
	//
	// strings := map[string]string{
	// 	"hello":    "world",
	// 	"username": "onotole",
	// }
	// missingString := "missString"

	sess := retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, val := range ints {
		assertIntKey(t, "Пустая сессия не должна содержать значений int.", sess, key, 0)
		sess.SetInt(key, val)
	}

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, val := range ints {
		assertIntKey(t, "Сессия должна сохранять значения int.", sess, key, val)
	}
	assertIntKey(t, "Сессия должна возвращать 0 для незаданных значений.", sess, missingInt, 0)

	// Обнуление куки.
	client.Jar = nil

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	sess.SetInt(missingInt, 12)

	sess = retrieveSession(client)
	if sess == nil {
		t.Fatal("Сессия не должна быть nil.")
	}
	for key, _ := range ints {
		assertIntKey(t, "Сессия не должна возвращать значения чужой сессии.", sess, key, 0)
	}
	assertIntKey(t, "Сессия должна сохранять значения int.", sess, missingInt, 12)

}
<|endoftext|>