<commit_before>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletmanager

// This file handles the health check. It is enabled by passing a
// target_tablet_type command line parameter. The tablet will then go
// to the target tablet type if healthy, and to 'spare' if not.

import (
	"flag"
	"fmt"
	"html/template"
	"reflect"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/timer"
	"github.com/youtube/vitess/go/vt/servenv"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topotools"

	pb "github.com/youtube/vitess/go/vt/proto/query"
)

const (
	defaultDegradedThreshold  = time.Duration(30 * time.Second)
	defaultUnhealthyThreshold = time.Duration(2 * time.Hour)
)

var (
	healthCheckInterval = flag.Duration("health_check_interval", 20*time.Second, "Interval between health checks")
	targetTabletType    = flag.String("target_tablet_type", "", "The tablet type we are thriving to be when healthy. When not healthy, we'll go to spare.")
	degradedThreshold   = flag.Duration("degraded_threshold", defaultDegradedThreshold, "replication lag after which a replica is considered degraded")
	unhealthyThreshold  = flag.Duration("unhealthy_threshold", defaultUnhealthyThreshold, "replication lag  after which a replica is considered unhealthy")
)

// HealthRecord records one run of the health checker
type HealthRecord struct {
	Error            error
	ReplicationDelay time.Duration
	Time             time.Time
}

// Class returns a human-readable one word version of the health state.
func (r *HealthRecord) Class() string {
	switch {
	case r.Error != nil:
		return "unhealthy"
	case r.ReplicationDelay > *degradedThreshold:
		return "unhappy"
	default:
		return "healthy"
	}
}

// HTML returns a html version to be displayed on UIs
func (r *HealthRecord) HTML() template.HTML {
	switch {
	case r.Error != nil:
		return template.HTML(fmt.Sprintf("unhealthy: %v", r.Error))
	case r.ReplicationDelay > *degradedThreshold:
		return template.HTML(fmt.Sprintf("unhappy: %v behind on replication", r.ReplicationDelay))
	default:
		if r.ReplicationDelay > 0 {
			return template.HTML(fmt.Sprintf("healthy: only %v behind on replication", r.ReplicationDelay))
		}
		return template.HTML("healthy")
	}
}

// IsDuplicate implements history.Deduplicable
func (r *HealthRecord) IsDuplicate(other interface{}) bool {
	rother, ok := other.(*HealthRecord)
	if !ok {
		return false
	}
	if !reflect.DeepEqual(r.Error, rother.Error) {
		return false
	}
	unhealthy := r.ReplicationDelay > *degradedThreshold
	unhealthyOther := rother.ReplicationDelay > *degradedThreshold
	return (unhealthy && unhealthyOther) || (!unhealthy && !unhealthyOther)
}

// ConfigHTML returns a formatted summary of health checking config values.
func ConfigHTML() template.HTML {
	return template.HTML(fmt.Sprintf(
		"healthCheckInterval: %v; degradedThreshold: %v; unhealthyThreshold: %v",
		healthCheckInterval, degradedThreshold, unhealthyThreshold))
}

// IsRunningHealthCheck indicates if the agent is configured to run healthchecks.
func (agent *ActionAgent) IsRunningHealthCheck() bool {
	return *targetTabletType != ""
}

func (agent *ActionAgent) initHeathCheck() {
	if !agent.IsRunningHealthCheck() {
		log.Infof("No target_tablet_type specified, disabling any health check")
		return
	}

	log.Infof("Starting periodic health check every %v with target_tablet_type=%v", *healthCheckInterval, *targetTabletType)
	t := timer.NewTimer(*healthCheckInterval)
	servenv.OnTermSync(func() {
		// When we enter lameduck mode, we want to not call
		// the health check any more. After this returns, we
		// are guaranteed to not call it.
		log.Info("Stopping periodic health check timer")
		t.Stop()

		// Now we can finish up and force ourselves to not healthy.
		agent.terminateHealthChecks(topo.TabletType(*targetTabletType))
	})
	t.Start(func() {
		agent.runHealthCheck(topo.TabletType(*targetTabletType))
	})
	t.Trigger()
}

// runHealthCheck takes the action mutex, runs the health check,
// and if we need to change our state, do it.
// If we are the master, we don't change our type, healthy or not.
// If we are not the master, we change to spare if not healthy,
// or to the passed in targetTabletType if healthy.
//
// Note we only update the topo record if we need to, that is if our type or
// health details changed.
func (agent *ActionAgent) runHealthCheck(targetTabletType topo.TabletType) {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()

	// read the current tablet record and tablet control
	agent.mutex.Lock()
	tablet := agent._tablet
	tabletControl := agent._tabletControl
	agent.mutex.Unlock()

	// figure out if we should be running the query service
	shouldQueryServiceBeRunning := false
	var blacklistedTables []string
	if topo.IsRunningQueryService(targetTabletType) && agent.BinlogPlayerMap.size() == 0 {
		shouldQueryServiceBeRunning = true
		if tabletControl != nil {
			blacklistedTables = tabletControl.BlacklistedTables
			if tabletControl.DisableQueryService {
				shouldQueryServiceBeRunning = false
			}
		}
	}

	// run the health check
	typeForHealthCheck := targetTabletType
	if tablet.Type == topo.TYPE_MASTER {
		typeForHealthCheck = topo.TYPE_MASTER
	}
	replicationDelay, err := agent.HealthReporter.Report(topo.IsSlaveType(typeForHealthCheck), shouldQueryServiceBeRunning)
	health := make(map[string]string)
	if err == nil {
		if replicationDelay > *unhealthyThreshold {
			err = fmt.Errorf("reported replication lag: %v higher than unhealthy threshold: %v", replicationDelay.Seconds(), unhealthyThreshold.Seconds())
		} else if replicationDelay > *degradedThreshold {
			health[topo.ReplicationLag] = topo.ReplicationLagHigh
		}
	}

	// Figure out if we should be running QueryService, see if we are,
	// and reconcile.
	if err != nil {
		if shouldQueryServiceBeRunning && tablet.Type == topo.TYPE_WORKER {
			// Despite being not healthy (e.g. replication may be stopped),
			// we do not shutdown the query service because the vtworker command needs it e.g. as source tablet during a SplitClone.
		} else {
			// we are not healthy, we should not be running QueryService
			shouldQueryServiceBeRunning = false
		}
	}
	isQueryServiceRunning := agent.QueryServiceControl.IsServing()
	if shouldQueryServiceBeRunning {
		if !isQueryServiceRunning {
			// send the type we want to be, not the type we are
			currentType := tablet.Type
			if tablet.Type == topo.TYPE_SPARE {
				tablet.Type = targetTabletType
			}

			// we remember this new possible error
			err = agent.allowQueries(tablet.Tablet, blacklistedTables)

			// restore the current type
			tablet.Type = currentType
		}
	} else {
		if isQueryServiceRunning {
			// we are not healthy or should not be running the
			// query service, shut it down.
			agent.disallowQueries()
		}
	}

	// save the health record
	record := &HealthRecord{
		Error:            err,
		ReplicationDelay: replicationDelay,
		Time:             time.Now(),
	}
	agent.History.Add(record)

	// try to figure out the mysql port if we don't have it yet
	if _, ok := tablet.Portmap["mysql"]; !ok {
		// we don't know the port, try to get it from mysqld
		mysqlPort, err := agent.MysqlDaemon.GetMysqlPort()
		if err != nil {
			// Don't log if we're already in a waiting-for-mysql state.
			agent.mutex.Lock()
			if !agent._waitingForMysql {
				log.Warningf("Can't get mysql port, won't populate Tablet record in topology (will retry silently at healthcheck interval %v): %v", *healthCheckInterval, err)
				agent._waitingForMysql = true
			}
			agent.mutex.Unlock()
		} else {
			log.Infof("Updating tablet mysql port to %v", mysqlPort)
			if err := agent.TopoServer.UpdateTabletFields(agent.batchCtx, tablet.Alias, func(tablet *topo.Tablet) error {
				tablet.Portmap["mysql"] = mysqlPort
				return nil
			}); err != nil {
				log.Infof("Error updating mysql port in tablet record: %v", err)
				return
			}

			// save the port so we don't update it again next time
			// we do the health check.
			agent.mutex.Lock()
			agent._tablet.Portmap["mysql"] = mysqlPort
			agent._waitingForMysql = false
			agent.mutex.Unlock()
		}
	}

	// remember our health status
	agent.mutex.Lock()
	agent._healthy = err
	agent._healthyTime = time.Now()
	agent._replicationDelay = replicationDelay
	terTime := agent._tabletExternallyReparentedTime
	agent.mutex.Unlock()

	// send it to our observers
	// (the Target has already been updated when restarting the
	// query service earlier)
	// FIXME(alainjobart,liguo) add CpuUsage
	stats := &pb.RealtimeStats{
		SecondsBehindMaster: uint32(replicationDelay.Seconds()),
	}
	if err != nil {
		stats.HealthError = err.Error()
	}
	defer func() {
		var ts int64
		if !terTime.IsZero() {
			ts = terTime.Unix()
		}
		agent.QueryServiceControl.BroadcastHealth(ts, stats)
	}()

	// Update our topo.Server state, start with no change
	newTabletType := tablet.Type
	if err != nil {
		// The tablet is not healthy, let's see what we need to do
		if tablet.Type != targetTabletType {
			if tablet.Type != topo.TYPE_SPARE {
				// we only log if we're not in spare,
				// as the spare state is normal for a
				// failed health check.
				log.Infof("Tablet not healthy and in state %v, not changing it: %v", tablet.Type, err)
			}
			return
		}

		// Note that if the query service is running, we may
		// need to stop it. The post-action callback will do
		// it, and it will be done after we change our state,
		// so it's the right order, let it do it.
		log.Infof("Tablet not healthy, converting it from %v to spare: %v", targetTabletType, err)
		newTabletType = topo.TYPE_SPARE
	} else {
		// We are healthy, maybe with health, see if we need
		// to update the record. We only change from spare to
		// our target type.
		if tablet.Type == topo.TYPE_SPARE {
			newTabletType = targetTabletType
		}
		if tablet.Type == newTabletType && tablet.IsHealthEqual(health) {
			// no change in health, not logging anything,
			// and we're done
			return
		}

		// we need to update our state
		log.Infof("Updating tablet record as healthy type %v -> %v with health details %v -> %v", tablet.Type, newTabletType, tablet.Health, health)
		agent.lastHealthMapCount.Set(int64(len(health)))
	}

	// Change the Type, update the health. Note we pass in a map
	// that's not nil, meaning if it's empty, we will clear it.
	if err := topotools.ChangeType(agent.batchCtx, agent.TopoServer, tablet.Alias, newTabletType, health); err != nil {
		log.Infof("Error updating tablet record: %v", err)
		return
	}
	tablet.Health = health
	tablet.Type = newTabletType

	// Rebuild the serving graph in our cell, only if we're dealing with
	// a serving type
	if err := agent.updateServingGraph(tablet, targetTabletType); err != nil {
		log.Warningf("updateServingGraph failed (will still run post action callbacks, serving graph might be out of date): %v", err)
	}

	// run the post action callbacks, not much we can do with returned error
	if err := agent.refreshTablet(agent.batchCtx, "healthcheck"); err != nil {
		log.Warningf("refreshTablet failed: %v", err)
	}
}

// terminateHealthChecks is called when we enter lame duck mode.
// We will clean up our state, and shut down query service.
// We only do something if we are in targetTabletType state, and then
// we just go to spare.
func (agent *ActionAgent) terminateHealthChecks(targetTabletType topo.TabletType) {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()
	log.Info("agent.terminateHealthChecks is starting")

	// read the current tablet record
	tablet := agent.Tablet()
	if tablet.Type != targetTabletType {
		log.Infof("Tablet in state %v, not changing it", tablet.Type)
		return
	}

	// Change the Type to spare, update the health. Note we pass in a map
	// that's not nil, meaning we will clear it.
	if err := topotools.ChangeType(agent.batchCtx, agent.TopoServer, tablet.Alias, topo.TYPE_SPARE, make(map[string]string)); err != nil {
		log.Infof("Error updating tablet record: %v", err)
		return
	}

	// Update the serving graph in our cell, only if we're dealing with
	// a serving type
	if err := agent.updateServingGraph(tablet, targetTabletType); err != nil {
		log.Warningf("updateServingGraph failed (will still run post action callbacks, serving graph might be out of date): %v", err)
	}

	// We've already rebuilt the shard, which is the only reason we registered
	// ourself as OnTermSync (synchronous). The rest can be done asynchronously.
	go func() {
		// Run the post action callbacks (let them shutdown the query service)
		if err := agent.refreshTablet(agent.batchCtx, "terminatehealthcheck"); err != nil {
			log.Warningf("refreshTablet failed: %v", err)
		}
	}()
}

// updateServingGraph will update the serving graph if we need to.
func (agent *ActionAgent) updateServingGraph(tablet *topo.TabletInfo, targetTabletType topo.TabletType) error {
	if topo.IsInServingGraph(targetTabletType) {
		if err := topotools.UpdateTabletEndpoints(agent.batchCtx, agent.TopoServer, tablet.Tablet); err != nil {
			return fmt.Errorf("UpdateTabletEndpoints failed: %v", err)
		}
	}
	return nil
}
<commit_msg>healthcheck: Simplify vtworker exception from queryservice shutdown logic.<commit_after>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tabletmanager

// This file handles the health check. It is enabled by passing a
// target_tablet_type command line parameter. The tablet will then go
// to the target tablet type if healthy, and to 'spare' if not.

import (
	"flag"
	"fmt"
	"html/template"
	"reflect"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/timer"
	"github.com/youtube/vitess/go/vt/servenv"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/topotools"

	pb "github.com/youtube/vitess/go/vt/proto/query"
)

const (
	defaultDegradedThreshold  = time.Duration(30 * time.Second)
	defaultUnhealthyThreshold = time.Duration(2 * time.Hour)
)

var (
	healthCheckInterval = flag.Duration("health_check_interval", 20*time.Second, "Interval between health checks")
	targetTabletType    = flag.String("target_tablet_type", "", "The tablet type we are thriving to be when healthy. When not healthy, we'll go to spare.")
	degradedThreshold   = flag.Duration("degraded_threshold", defaultDegradedThreshold, "replication lag after which a replica is considered degraded")
	unhealthyThreshold  = flag.Duration("unhealthy_threshold", defaultUnhealthyThreshold, "replication lag  after which a replica is considered unhealthy")
)

// HealthRecord records one run of the health checker
type HealthRecord struct {
	Error            error
	ReplicationDelay time.Duration
	Time             time.Time
}

// Class returns a human-readable one word version of the health state.
func (r *HealthRecord) Class() string {
	switch {
	case r.Error != nil:
		return "unhealthy"
	case r.ReplicationDelay > *degradedThreshold:
		return "unhappy"
	default:
		return "healthy"
	}
}

// HTML returns a html version to be displayed on UIs
func (r *HealthRecord) HTML() template.HTML {
	switch {
	case r.Error != nil:
		return template.HTML(fmt.Sprintf("unhealthy: %v", r.Error))
	case r.ReplicationDelay > *degradedThreshold:
		return template.HTML(fmt.Sprintf("unhappy: %v behind on replication", r.ReplicationDelay))
	default:
		if r.ReplicationDelay > 0 {
			return template.HTML(fmt.Sprintf("healthy: only %v behind on replication", r.ReplicationDelay))
		}
		return template.HTML("healthy")
	}
}

// IsDuplicate implements history.Deduplicable
func (r *HealthRecord) IsDuplicate(other interface{}) bool {
	rother, ok := other.(*HealthRecord)
	if !ok {
		return false
	}
	if !reflect.DeepEqual(r.Error, rother.Error) {
		return false
	}
	unhealthy := r.ReplicationDelay > *degradedThreshold
	unhealthyOther := rother.ReplicationDelay > *degradedThreshold
	return (unhealthy && unhealthyOther) || (!unhealthy && !unhealthyOther)
}

// ConfigHTML returns a formatted summary of health checking config values.
func ConfigHTML() template.HTML {
	return template.HTML(fmt.Sprintf(
		"healthCheckInterval: %v; degradedThreshold: %v; unhealthyThreshold: %v",
		healthCheckInterval, degradedThreshold, unhealthyThreshold))
}

// IsRunningHealthCheck indicates if the agent is configured to run healthchecks.
func (agent *ActionAgent) IsRunningHealthCheck() bool {
	return *targetTabletType != ""
}

func (agent *ActionAgent) initHeathCheck() {
	if !agent.IsRunningHealthCheck() {
		log.Infof("No target_tablet_type specified, disabling any health check")
		return
	}

	log.Infof("Starting periodic health check every %v with target_tablet_type=%v", *healthCheckInterval, *targetTabletType)
	t := timer.NewTimer(*healthCheckInterval)
	servenv.OnTermSync(func() {
		// When we enter lameduck mode, we want to not call
		// the health check any more. After this returns, we
		// are guaranteed to not call it.
		log.Info("Stopping periodic health check timer")
		t.Stop()

		// Now we can finish up and force ourselves to not healthy.
		agent.terminateHealthChecks(topo.TabletType(*targetTabletType))
	})
	t.Start(func() {
		agent.runHealthCheck(topo.TabletType(*targetTabletType))
	})
	t.Trigger()
}

// runHealthCheck takes the action mutex, runs the health check,
// and if we need to change our state, do it.
// If we are the master, we don't change our type, healthy or not.
// If we are not the master, we change to spare if not healthy,
// or to the passed in targetTabletType if healthy.
//
// Note we only update the topo record if we need to, that is if our type or
// health details changed.
func (agent *ActionAgent) runHealthCheck(targetTabletType topo.TabletType) {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()

	// read the current tablet record and tablet control
	agent.mutex.Lock()
	tablet := agent._tablet
	tabletControl := agent._tabletControl
	agent.mutex.Unlock()

	// figure out if we should be running the query service
	shouldQueryServiceBeRunning := false
	var blacklistedTables []string
	if topo.IsRunningQueryService(targetTabletType) && agent.BinlogPlayerMap.size() == 0 {
		shouldQueryServiceBeRunning = true
		if tabletControl != nil {
			blacklistedTables = tabletControl.BlacklistedTables
			if tabletControl.DisableQueryService {
				shouldQueryServiceBeRunning = false
			}
		}
	}

	// run the health check
	typeForHealthCheck := targetTabletType
	if tablet.Type == topo.TYPE_MASTER {
		typeForHealthCheck = topo.TYPE_MASTER
	}
	replicationDelay, err := agent.HealthReporter.Report(topo.IsSlaveType(typeForHealthCheck), shouldQueryServiceBeRunning)
	health := make(map[string]string)
	if err == nil {
		if replicationDelay > *unhealthyThreshold {
			err = fmt.Errorf("reported replication lag: %v higher than unhealthy threshold: %v", replicationDelay.Seconds(), unhealthyThreshold.Seconds())
		} else if replicationDelay > *degradedThreshold {
			health[topo.ReplicationLag] = topo.ReplicationLagHigh
		}
	}

	// Figure out if we should be running QueryService, see if we are,
	// and reconcile.
	if err != nil {
		if tablet.Type != topo.TYPE_WORKER {
			// We are not healthy and must shut down QueryService.
			// At the moment, the only exception to this are "worker" tablets which
			// still must serve queries e.g. as source tablet during a "SplitClone".
			shouldQueryServiceBeRunning = false
		}
	}
	isQueryServiceRunning := agent.QueryServiceControl.IsServing()
	if shouldQueryServiceBeRunning {
		if !isQueryServiceRunning {
			// send the type we want to be, not the type we are
			currentType := tablet.Type
			if tablet.Type == topo.TYPE_SPARE {
				tablet.Type = targetTabletType
			}

			// we remember this new possible error
			err = agent.allowQueries(tablet.Tablet, blacklistedTables)

			// restore the current type
			tablet.Type = currentType
		}
	} else {
		if isQueryServiceRunning {
			// we are not healthy or should not be running the
			// query service, shut it down.
			agent.disallowQueries()
		}
	}

	// save the health record
	record := &HealthRecord{
		Error:            err,
		ReplicationDelay: replicationDelay,
		Time:             time.Now(),
	}
	agent.History.Add(record)

	// try to figure out the mysql port if we don't have it yet
	if _, ok := tablet.Portmap["mysql"]; !ok {
		// we don't know the port, try to get it from mysqld
		mysqlPort, err := agent.MysqlDaemon.GetMysqlPort()
		if err != nil {
			// Don't log if we're already in a waiting-for-mysql state.
			agent.mutex.Lock()
			if !agent._waitingForMysql {
				log.Warningf("Can't get mysql port, won't populate Tablet record in topology (will retry silently at healthcheck interval %v): %v", *healthCheckInterval, err)
				agent._waitingForMysql = true
			}
			agent.mutex.Unlock()
		} else {
			log.Infof("Updating tablet mysql port to %v", mysqlPort)
			if err := agent.TopoServer.UpdateTabletFields(agent.batchCtx, tablet.Alias, func(tablet *topo.Tablet) error {
				tablet.Portmap["mysql"] = mysqlPort
				return nil
			}); err != nil {
				log.Infof("Error updating mysql port in tablet record: %v", err)
				return
			}

			// save the port so we don't update it again next time
			// we do the health check.
			agent.mutex.Lock()
			agent._tablet.Portmap["mysql"] = mysqlPort
			agent._waitingForMysql = false
			agent.mutex.Unlock()
		}
	}

	// remember our health status
	agent.mutex.Lock()
	agent._healthy = err
	agent._healthyTime = time.Now()
	agent._replicationDelay = replicationDelay
	terTime := agent._tabletExternallyReparentedTime
	agent.mutex.Unlock()

	// send it to our observers
	// (the Target has already been updated when restarting the
	// query service earlier)
	// FIXME(alainjobart,liguo) add CpuUsage
	stats := &pb.RealtimeStats{
		SecondsBehindMaster: uint32(replicationDelay.Seconds()),
	}
	if err != nil {
		stats.HealthError = err.Error()
	}
	defer func() {
		var ts int64
		if !terTime.IsZero() {
			ts = terTime.Unix()
		}
		agent.QueryServiceControl.BroadcastHealth(ts, stats)
	}()

	// Update our topo.Server state, start with no change
	newTabletType := tablet.Type
	if err != nil {
		// The tablet is not healthy, let's see what we need to do
		if tablet.Type != targetTabletType {
			if tablet.Type != topo.TYPE_SPARE {
				// we only log if we're not in spare,
				// as the spare state is normal for a
				// failed health check.
				log.Infof("Tablet not healthy and in state %v, not changing it: %v", tablet.Type, err)
			}
			return
		}

		// Note that if the query service is running, we may
		// need to stop it. The post-action callback will do
		// it, and it will be done after we change our state,
		// so it's the right order, let it do it.
		log.Infof("Tablet not healthy, converting it from %v to spare: %v", targetTabletType, err)
		newTabletType = topo.TYPE_SPARE
	} else {
		// We are healthy, maybe with health, see if we need
		// to update the record. We only change from spare to
		// our target type.
		if tablet.Type == topo.TYPE_SPARE {
			newTabletType = targetTabletType
		}
		if tablet.Type == newTabletType && tablet.IsHealthEqual(health) {
			// no change in health, not logging anything,
			// and we're done
			return
		}

		// we need to update our state
		log.Infof("Updating tablet record as healthy type %v -> %v with health details %v -> %v", tablet.Type, newTabletType, tablet.Health, health)
		agent.lastHealthMapCount.Set(int64(len(health)))
	}

	// Change the Type, update the health. Note we pass in a map
	// that's not nil, meaning if it's empty, we will clear it.
	if err := topotools.ChangeType(agent.batchCtx, agent.TopoServer, tablet.Alias, newTabletType, health); err != nil {
		log.Infof("Error updating tablet record: %v", err)
		return
	}
	tablet.Health = health
	tablet.Type = newTabletType

	// Rebuild the serving graph in our cell, only if we're dealing with
	// a serving type
	if err := agent.updateServingGraph(tablet, targetTabletType); err != nil {
		log.Warningf("updateServingGraph failed (will still run post action callbacks, serving graph might be out of date): %v", err)
	}

	// run the post action callbacks, not much we can do with returned error
	if err := agent.refreshTablet(agent.batchCtx, "healthcheck"); err != nil {
		log.Warningf("refreshTablet failed: %v", err)
	}
}

// terminateHealthChecks is called when we enter lame duck mode.
// We will clean up our state, and shut down query service.
// We only do something if we are in targetTabletType state, and then
// we just go to spare.
func (agent *ActionAgent) terminateHealthChecks(targetTabletType topo.TabletType) {
	agent.actionMutex.Lock()
	defer agent.actionMutex.Unlock()
	log.Info("agent.terminateHealthChecks is starting")

	// read the current tablet record
	tablet := agent.Tablet()
	if tablet.Type != targetTabletType {
		log.Infof("Tablet in state %v, not changing it", tablet.Type)
		return
	}

	// Change the Type to spare, update the health. Note we pass in a map
	// that's not nil, meaning we will clear it.
	if err := topotools.ChangeType(agent.batchCtx, agent.TopoServer, tablet.Alias, topo.TYPE_SPARE, make(map[string]string)); err != nil {
		log.Infof("Error updating tablet record: %v", err)
		return
	}

	// Update the serving graph in our cell, only if we're dealing with
	// a serving type
	if err := agent.updateServingGraph(tablet, targetTabletType); err != nil {
		log.Warningf("updateServingGraph failed (will still run post action callbacks, serving graph might be out of date): %v", err)
	}

	// We've already rebuilt the shard, which is the only reason we registered
	// ourself as OnTermSync (synchronous). The rest can be done asynchronously.
	go func() {
		// Run the post action callbacks (let them shutdown the query service)
		if err := agent.refreshTablet(agent.batchCtx, "terminatehealthcheck"); err != nil {
			log.Warningf("refreshTablet failed: %v", err)
		}
	}()
}

// updateServingGraph will update the serving graph if we need to.
func (agent *ActionAgent) updateServingGraph(tablet *topo.TabletInfo, targetTabletType topo.TabletType) error {
	if topo.IsInServingGraph(targetTabletType) {
		if err := topotools.UpdateTabletEndpoints(agent.batchCtx, agent.TopoServer, tablet.Tablet); err != nil {
			return fmt.Errorf("UpdateTabletEndpoints failed: %v", err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package repo

import (
	"path"
	"strings"

	"github.com/codegangsta/martini"

	"github.com/gogits/git"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/base"
	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/middleware"
)

func Branches(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsValid {
		return
	}

	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		ctx.Handle(200, "repo.Branches", err)
		return
	} else if len(brs) == 0 {
		ctx.Error(404)
		return
	}

	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]

	ctx.Data["Branchname"] = brs[0]
	ctx.Data["Branches"] = brs
	ctx.Data["IsRepoToolbarBranches"] = true

	ctx.HTML(200, "repo/branches")
}

func Single(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsValid {
		return
	}

	if len(params["branchname"]) == 0 {
		params["branchname"] = "master"
	}

	// Get tree path
	treename := params["_1"]

	if len(treename) > 0 && treename[len(treename)-1] == '/' {
		ctx.Redirect("/"+ctx.Repo.Owner.LowerName+"/"+
			ctx.Repo.Repository.Name+"/src/"+params["branchname"]+"/"+treename[:len(treename)-1], 302)
		return
	}

	ctx.Data["IsRepoToolbarSource"] = true

	// Branches.
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		log.Error("repo.Single(GetBranches): %v", err)
		ctx.Error(404)
		return
	} else if len(brs) == 0 {
		ctx.Data["IsBareRepo"] = true
		ctx.HTML(200, "repo/single")
		return
	}

	ctx.Data["Branches"] = brs

	repoFile, err := models.GetTargetFile(params["username"], params["reponame"],
		params["branchname"], params["commitid"], treename)

	if err != nil && err != models.ErrRepoFileNotExist {
		log.Error("repo.Single(GetTargetFile): %v", err)
		ctx.Error(404)
		return
	}

	branchLink := "/" + ctx.Repo.Owner.LowerName + "/" + ctx.Repo.Repository.Name + "/src/" + params["branchname"]

	if repoFile != nil && repoFile.IsFile() {
		if repoFile.Size > 1024*1024 || repoFile.Filemode != git.FileModeBlob {
			ctx.Data["FileIsLarge"] = true
		} else if blob, err := repoFile.LookupBlob(); err != nil {
			log.Error("repo.Single(repoFile.LookupBlob): %v", err)
			ctx.Error(404)
		} else {
			ctx.Data["IsFile"] = true
			ctx.Data["FileName"] = repoFile.Name
			ext := path.Ext(repoFile.Name)
			if len(ext) > 0 {
				ext = ext[1:]
			}
			ctx.Data["FileExt"] = ext

			readmeExist := base.IsMarkdownFile(repoFile.Name) || base.IsReadmeFile(repoFile.Name)
			ctx.Data["ReadmeExist"] = readmeExist
			if readmeExist {
				ctx.Data["FileContent"] = string(base.RenderMarkdown(blob.Contents(), ""))
			} else {
				ctx.Data["FileContent"] = string(blob.Contents())
			}
		}

	} else {
		// Directory and file list.
		files, err := models.GetReposFiles(params["username"], params["reponame"],
			params["branchname"], params["commitid"], treename)
		if err != nil {
			log.Error("repo.Single(GetReposFiles): %v", err)
			ctx.Error(404)
			return
		}

		ctx.Data["Files"] = files

		var readmeFile *models.RepoFile

		for _, f := range files {
			if !f.IsFile() || !base.IsReadmeFile(f.Name) {
				continue
			} else {
				readmeFile = f
				break
			}
		}

		if readmeFile != nil {
			ctx.Data["ReadmeExist"] = true
			// if file large than 1M not show it
			if readmeFile.Size > 1024*1024 || readmeFile.Filemode != git.FileModeBlob {
				ctx.Data["FileIsLarge"] = true
			} else if blob, err := readmeFile.LookupBlob(); err != nil {
				log.Error("repo.Single(readmeFile.LookupBlob): %v", err)
				ctx.Error(404)
				return
			} else {
				// current repo branch link

				ctx.Data["FileName"] = readmeFile.Name
				ctx.Data["FileContent"] = string(base.RenderMarkdown(blob.Contents(), branchLink))
			}
		}
	}

	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]
	ctx.Data["Branchname"] = params["branchname"]

	var treenames []string
	Paths := make([]string, 0)

	if len(treename) > 0 {
		treenames = strings.Split(treename, "/")
		for i, _ := range treenames {
			Paths = append(Paths, strings.Join(treenames[0:i+1], "/"))
		}

		ctx.Data["HasParentPath"] = true
		if len(Paths)-2 >= 0 {
			ctx.Data["ParentPath"] = "/" + Paths[len(Paths)-2]
		}
	}

	// Get latest commit according username and repo name
	commit, err := models.GetCommit(params["username"], params["reponame"],
		params["branchname"], params["commitid"])
	if err != nil {
		log.Error("repo.Single(GetCommit): %v", err)
		ctx.Error(404)
		return
	}
	ctx.Data["LastCommit"] = commit

	ctx.Data["Paths"] = Paths
	ctx.Data["Treenames"] = treenames
	ctx.Data["BranchLink"] = branchLink
	ctx.HTML(200, "repo/single")
}

func Setting(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsOwner {
		ctx.Error(404)
		return
	}

	ctx.Data["IsRepoToolbarSetting"] = true

	// Branches.
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		log.Error("repo.Setting(GetBranches): %v", err)
		ctx.Error(404)
		return
	} else if len(brs) == 0 {
		ctx.Data["IsBareRepo"] = true
		ctx.HTML(200, "repo/setting")
		return
	}

	var title string
	if t, ok := ctx.Data["Title"].(string); ok {
		title = t
	}

	ctx.Data["Title"] = title + " - settings"
	ctx.HTML(200, "repo/setting")
}

func Commits(ctx *middleware.Context, params martini.Params) {
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		ctx.Handle(200, "repo.Commits", err)
		return
	} else if len(brs) == 0 {
		ctx.Error(404)
		return
	}

	ctx.Data["IsRepoToolbarCommits"] = true
	commits, err := models.GetCommits(params["username"],
		params["reponame"], params["branchname"])
	if err != nil {
		ctx.Error(404)
		return
	}
	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]
	ctx.Data["CommitCount"] = commits.Len()
	ctx.Data["Commits"] = commits
	ctx.HTML(200, "repo/commits")
}

func Issues(ctx *middleware.Context) {
	ctx.Data["IsRepoToolbarIssues"] = true
	ctx.HTML(200, "repo/issues")
}

func Pulls(ctx *middleware.Context) {
	ctx.Data["IsRepoToolbarPulls"] = true
	ctx.HTML(200, "repo/pulls")
}

func Action(ctx *middleware.Context, params martini.Params) {
	var err error
	switch params["action"] {
	case "watch":
		err = models.WatchRepo(ctx.User.Id, ctx.Repo.Repository.Id, true)
	case "unwatch":
		err = models.WatchRepo(ctx.User.Id, ctx.Repo.Repository.Id, false)
	}

	if err != nil {
		log.Error("repo.Action(%s): %v", params["action"], err)
		ctx.JSON(200, map[string]interface{}{
			"ok":  false,
			"err": err.Error(),
		})
		return
	}
	ctx.JSON(200, map[string]interface{}{
		"ok": true,
	})
}
<commit_msg>wrong link should 404<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package repo

import (
	"path"
	"strings"

	"github.com/codegangsta/martini"

	"github.com/gogits/git"

	"github.com/gogits/gogs/models"
	"github.com/gogits/gogs/modules/base"
	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/middleware"
)

func Branches(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsValid {
		return
	}

	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		ctx.Handle(200, "repo.Branches", err)
		return
	} else if len(brs) == 0 {
		ctx.Error(404)
		return
	}

	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]

	ctx.Data["Branchname"] = brs[0]
	ctx.Data["Branches"] = brs
	ctx.Data["IsRepoToolbarBranches"] = true

	ctx.HTML(200, "repo/branches")
}

func Single(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsValid {
		return
	}

	if len(params["branchname"]) == 0 {
		params["branchname"] = "master"
	}

	// Get tree path
	treename := params["_1"]

	if len(treename) > 0 && treename[len(treename)-1] == '/' {
		ctx.Redirect("/"+ctx.Repo.Owner.LowerName+"/"+
			ctx.Repo.Repository.Name+"/src/"+params["branchname"]+"/"+treename[:len(treename)-1], 302)
		return
	}

	ctx.Data["IsRepoToolbarSource"] = true

	// Branches.
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		log.Error("repo.Single(GetBranches): %v", err)
		ctx.Error(404)
		return
	} else if len(brs) == 0 {
		ctx.Data["IsBareRepo"] = true
		ctx.HTML(200, "repo/single")
		return
	}

	ctx.Data["Branches"] = brs

	repoFile, err := models.GetTargetFile(params["username"], params["reponame"],
		params["branchname"], params["commitid"], treename)

	if err != nil && err != models.ErrRepoFileNotExist {
		log.Error("repo.Single(GetTargetFile): %v", err)
		ctx.Error(404)
		return
	}

	branchLink := "/" + ctx.Repo.Owner.LowerName + "/" + ctx.Repo.Repository.Name + "/src/" + params["branchname"]

	if len(treename) != 0 && repoFile == nil {
		ctx.Error(404)
		return
	}

	if repoFile != nil && repoFile.IsFile() {
		if repoFile.Size > 1024*1024 || repoFile.Filemode != git.FileModeBlob {
			ctx.Data["FileIsLarge"] = true
		} else if blob, err := repoFile.LookupBlob(); err != nil {
			log.Error("repo.Single(repoFile.LookupBlob): %v", err)
			ctx.Error(404)
		} else {
			ctx.Data["IsFile"] = true
			ctx.Data["FileName"] = repoFile.Name
			ext := path.Ext(repoFile.Name)
			if len(ext) > 0 {
				ext = ext[1:]
			}
			ctx.Data["FileExt"] = ext

			readmeExist := base.IsMarkdownFile(repoFile.Name) || base.IsReadmeFile(repoFile.Name)
			ctx.Data["ReadmeExist"] = readmeExist
			if readmeExist {
				ctx.Data["FileContent"] = string(base.RenderMarkdown(blob.Contents(), ""))
			} else {
				ctx.Data["FileContent"] = string(blob.Contents())
			}
		}

	} else {
		// Directory and file list.
		files, err := models.GetReposFiles(params["username"], params["reponame"],
			params["branchname"], params["commitid"], treename)
		if err != nil {
			log.Error("repo.Single(GetReposFiles): %v", err)
			ctx.Error(404)
			return
		}

		ctx.Data["Files"] = files

		var readmeFile *models.RepoFile

		for _, f := range files {
			if !f.IsFile() || !base.IsReadmeFile(f.Name) {
				continue
			} else {
				readmeFile = f
				break
			}
		}

		if readmeFile != nil {
			ctx.Data["ReadmeExist"] = true
			// if file large than 1M not show it
			if readmeFile.Size > 1024*1024 || readmeFile.Filemode != git.FileModeBlob {
				ctx.Data["FileIsLarge"] = true
			} else if blob, err := readmeFile.LookupBlob(); err != nil {
				log.Error("repo.Single(readmeFile.LookupBlob): %v", err)
				ctx.Error(404)
				return
			} else {
				// current repo branch link

				ctx.Data["FileName"] = readmeFile.Name
				ctx.Data["FileContent"] = string(base.RenderMarkdown(blob.Contents(), branchLink))
			}
		}
	}

	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]
	ctx.Data["Branchname"] = params["branchname"]

	var treenames []string
	Paths := make([]string, 0)

	if len(treename) > 0 {
		treenames = strings.Split(treename, "/")
		for i, _ := range treenames {
			Paths = append(Paths, strings.Join(treenames[0:i+1], "/"))
		}

		ctx.Data["HasParentPath"] = true
		if len(Paths)-2 >= 0 {
			ctx.Data["ParentPath"] = "/" + Paths[len(Paths)-2]
		}
	}

	// Get latest commit according username and repo name
	commit, err := models.GetCommit(params["username"], params["reponame"],
		params["branchname"], params["commitid"])
	if err != nil {
		log.Error("repo.Single(GetCommit): %v", err)
		ctx.Error(404)
		return
	}
	ctx.Data["LastCommit"] = commit

	ctx.Data["Paths"] = Paths
	ctx.Data["Treenames"] = treenames
	ctx.Data["BranchLink"] = branchLink
	ctx.HTML(200, "repo/single")
}

func Setting(ctx *middleware.Context, params martini.Params) {
	if !ctx.Repo.IsOwner {
		ctx.Error(404)
		return
	}

	ctx.Data["IsRepoToolbarSetting"] = true

	// Branches.
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		log.Error("repo.Setting(GetBranches): %v", err)
		ctx.Error(404)
		return
	} else if len(brs) == 0 {
		ctx.Data["IsBareRepo"] = true
		ctx.HTML(200, "repo/setting")
		return
	}

	var title string
	if t, ok := ctx.Data["Title"].(string); ok {
		title = t
	}

	ctx.Data["Title"] = title + " - settings"
	ctx.HTML(200, "repo/setting")
}

func Commits(ctx *middleware.Context, params martini.Params) {
	brs, err := models.GetBranches(params["username"], params["reponame"])
	if err != nil {
		ctx.Handle(200, "repo.Commits", err)
		return
	} else if len(brs) == 0 {
		ctx.Error(404)
		return
	}

	ctx.Data["IsRepoToolbarCommits"] = true
	commits, err := models.GetCommits(params["username"],
		params["reponame"], params["branchname"])
	if err != nil {
		ctx.Error(404)
		return
	}
	ctx.Data["Username"] = params["username"]
	ctx.Data["Reponame"] = params["reponame"]
	ctx.Data["CommitCount"] = commits.Len()
	ctx.Data["Commits"] = commits
	ctx.HTML(200, "repo/commits")
}

func Issues(ctx *middleware.Context) {
	ctx.Data["IsRepoToolbarIssues"] = true
	ctx.HTML(200, "repo/issues")
}

func Pulls(ctx *middleware.Context) {
	ctx.Data["IsRepoToolbarPulls"] = true
	ctx.HTML(200, "repo/pulls")
}

func Action(ctx *middleware.Context, params martini.Params) {
	var err error
	switch params["action"] {
	case "watch":
		err = models.WatchRepo(ctx.User.Id, ctx.Repo.Repository.Id, true)
	case "unwatch":
		err = models.WatchRepo(ctx.User.Id, ctx.Repo.Repository.Id, false)
	}

	if err != nil {
		log.Error("repo.Action(%s): %v", params["action"], err)
		ctx.JSON(200, map[string]interface{}{
			"ok":  false,
			"err": err.Error(),
		})
		return
	}
	ctx.JSON(200, map[string]interface{}{
		"ok": true,
	})
}
<|endoftext|>


<commit_before>package kite

import (
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	stdlog "log"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("Kite")

// GetLogger returns a new logger which is used within the application itsel
// (in main package).
func GetLogger() *logging.Logger {
	return log
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Registered is true if the Kite is registered to kontrol itself
	Registered bool

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Dnode rpc server
	Server *rpc.Server

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*CallOptions) error
}

// New creates, initialize and then returns a new Kite instance. It accept
// three  arguments. options is a config struct that needs to be filled with
// several informations like Name, Port, IP and so on.
func New(options *protocol.Options) *Kite {
	var err error
	if options == nil {
		options, err = utils.ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	// some simple validations for config
	if options.Kitename == "" {
		log.Fatal("ERROR: options.Kitename field is not set")
	}

	if options.Region == "" {
		log.Fatal("ERROR: options.Region field is not set")
	}

	if options.Environment == "" {
		log.Fatal("ERROR: options.Environment field is not set")
	}

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	kodingKey, err := utils.GetKodingKey()
	if err != nil {
		log.Fatal("Couldn't find koding.key. Please run 'kd register'.")
	}

	port := options.Port
	if options.Port == "" {
		port = "0" // OS binds to an automatic port
	}

	if options.KontrolAddr == "" {
		options.KontrolAddr = "127.0.0.1:4000" // local fallback address
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Port:        port,
			Environment: options.Environment,
			Region:      options.Region,

			// PublicIP will be set by Kontrol after registering if it is not set.
			PublicIP: options.PublicIP,
		},
		KodingKey:         kodingKey,
		Server:            rpc.NewServer(),
		KontrolEnabled:    true,
		RegisterToKontrol: true,
		Authenticators:    make(map[string]func(*CallOptions) error),
		handlers:          make(map[string]HandlerFunc),
	}
	k.Kontrol = k.NewKontrol(options.KontrolAddr)

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("status", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.Server.HandleFunc(method, func(msg *dnode.Message, tr dnode.Transport) {
		request, responseCallback, err := k.parseRequest(msg, tr)
		if err != nil {
			log.Notice("Did not understand request: %s", err)
			return
		}

		result, err := handler(request)
		if responseCallback == nil {
			return
		}

		if err != nil {
			err = responseCallback(err.Error(), result)
		} else {
			err = responseCallback(nil, result)
		}

		if err != nil {
			log.Error(err.Error())
		}
	})
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously. It can be started in a goroutine if you wish to use kite as a
// client too.
func (k *Kite) Run() {
	k.parseVersionFlag()
	k.setupLogging()

	err := k.listenAndServe()
	if err != nil {
		log.Fatal(err)
	}
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args, err := r.Args.Array()
	if err != nil {
		return nil, err
	}

	if len(args) != 2 {
		return nil, fmt.Errorf("Invalid args: %s", string(r.Args.Raw))
	}

	seconds, ok := args[0].(float64)
	if !ok {
		return nil, fmt.Errorf("Invalid interval: %s", args[0])
	}

	ping, ok := args[1].(dnode.Function)
	if !ok {
		return nil, fmt.Errorf("Invalid callback: %s", args[1])
	}

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	s, err := r.Args.String()
	if err != nil {
		return nil, err
	}

	log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, s))
	return nil, nil
}

// setupLogging is used to setup the logging format, destination and level.
func (k *Kite) setupLogging() {
	log.Module = k.Name
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))

	stderrBackend := logging.NewLogBackend(os.Stderr, "", stdlog.LstdFlags)
	stderrBackend.Color = true

	syslogBackend, _ := logging.NewSyslogBackend(k.Name)
	logging.SetBackend(stderrBackend, syslogBackend)

	// Set logging level. Default level is INFO.
	level := logging.INFO
	if k.hasDebugFlag() {
		level = logging.DEBUG
	}
	logging.SetLevel(level, log.Module)
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}
	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() error {
	listener, err := net.Listen("tcp4", ":"+k.Port)
	if err != nil {
		return err
	}

	log.Info("Listening: %s", listener.Addr().String())

	// Port is known here if "0" is used as port number
	_, k.Port, _ = net.SplitHostPort(listener.Addr().String())

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if k.RegisterToKontrol {
			k.Kontrol.RemoteKite.Client.OnConnect(k.registerToKontrol)
		}

		k.Kontrol.DialForever()
	}

	return http.Serve(listener, k.Server)
}

func (k *Kite) registerToKontrol() {
	err := k.Kontrol.Register()
	if err != nil {
		log.Fatalf("Cannot register to Kontrol: %s", err)
	}

	log.Info("Registered to Kontrol successfully")
}

// DISABLED TEMPORARILY
// OnDisconnect adds the given function to the list of the users callback list
// which is called when the user is disconnected. There might be several
// connections from one user to the kite, in that case the functions are
// called only when all connections are closed.
// func (k *Kite) OnDisconnect(username string, f func()) {
// 	if addrs == nil {
// 		return
// 	}

// 	for _, addr := range addrs {
// 		client.onDisconnect = append(client.onDisconnect, f)
// 	}
// }
<commit_msg>kite: remove unused comment, cenk is going to implement the OnConnect and OnDisconnect functions<commit_after>package kite

import (
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	stdlog "log"
	"net"
	"net/http"
	"os"
	"time"
)

var log = logging.MustGetLogger("Kite")

// GetLogger returns a new logger which is used within the application itsel
// (in main package).
func GetLogger() *logging.Logger {
	return log
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Registered is true if the Kite is registered to kontrol itself
	Registered bool

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Dnode rpc server
	Server *rpc.Server

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*CallOptions) error
}

// New creates, initialize and then returns a new Kite instance. It accept
// three  arguments. options is a config struct that needs to be filled with
// several informations like Name, Port, IP and so on.
func New(options *protocol.Options) *Kite {
	var err error
	if options == nil {
		options, err = utils.ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	// some simple validations for config
	if options.Kitename == "" {
		log.Fatal("ERROR: options.Kitename field is not set")
	}

	if options.Region == "" {
		log.Fatal("ERROR: options.Region field is not set")
	}

	if options.Environment == "" {
		log.Fatal("ERROR: options.Environment field is not set")
	}

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	kodingKey, err := utils.GetKodingKey()
	if err != nil {
		log.Fatal("Couldn't find koding.key. Please run 'kd register'.")
	}

	port := options.Port
	if options.Port == "" {
		port = "0" // OS binds to an automatic port
	}

	if options.KontrolAddr == "" {
		options.KontrolAddr = "127.0.0.1:4000" // local fallback address
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Port:        port,
			Environment: options.Environment,
			Region:      options.Region,

			// PublicIP will be set by Kontrol after registering if it is not set.
			PublicIP: options.PublicIP,
		},
		KodingKey:         kodingKey,
		Server:            rpc.NewServer(),
		KontrolEnabled:    true,
		RegisterToKontrol: true,
		Authenticators:    make(map[string]func(*CallOptions) error),
		handlers:          make(map[string]HandlerFunc),
	}
	k.Kontrol = k.NewKontrol(options.KontrolAddr)

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("status", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) HandleFunc(method string, handler HandlerFunc) {
	k.Server.HandleFunc(method, func(msg *dnode.Message, tr dnode.Transport) {
		request, responseCallback, err := k.parseRequest(msg, tr)
		if err != nil {
			log.Notice("Did not understand request: %s", err)
			return
		}

		result, err := handler(request)
		if responseCallback == nil {
			return
		}

		if err != nil {
			err = responseCallback(err.Error(), result)
		} else {
			err = responseCallback(nil, result)
		}

		if err != nil {
			log.Error(err.Error())
		}
	})
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously. It can be started in a goroutine if you wish to use kite as a
// client too.
func (k *Kite) Run() {
	k.parseVersionFlag()
	k.setupLogging()

	err := k.listenAndServe()
	if err != nil {
		log.Fatal(err)
	}
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args, err := r.Args.Array()
	if err != nil {
		return nil, err
	}

	if len(args) != 2 {
		return nil, fmt.Errorf("Invalid args: %s", string(r.Args.Raw))
	}

	seconds, ok := args[0].(float64)
	if !ok {
		return nil, fmt.Errorf("Invalid interval: %s", args[0])
	}

	ping, ok := args[1].(dnode.Function)
	if !ok {
		return nil, fmt.Errorf("Invalid callback: %s", args[1])
	}

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	s, err := r.Args.String()
	if err != nil {
		return nil, err
	}

	log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, s))
	return nil, nil
}

// setupLogging is used to setup the logging format, destination and level.
func (k *Kite) setupLogging() {
	log.Module = k.Name
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))

	stderrBackend := logging.NewLogBackend(os.Stderr, "", stdlog.LstdFlags)
	stderrBackend.Color = true

	syslogBackend, _ := logging.NewSyslogBackend(k.Name)
	logging.SetBackend(stderrBackend, syslogBackend)

	// Set logging level. Default level is INFO.
	level := logging.INFO
	if k.hasDebugFlag() {
		level = logging.DEBUG
	}
	logging.SetLevel(level, log.Module)
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}
	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() error {
	listener, err := net.Listen("tcp4", ":"+k.Port)
	if err != nil {
		return err
	}

	log.Info("Listening: %s", listener.Addr().String())

	// Port is known here if "0" is used as port number
	_, k.Port, _ = net.SplitHostPort(listener.Addr().String())

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if k.RegisterToKontrol {
			k.Kontrol.RemoteKite.Client.OnConnect(k.registerToKontrol)
		}

		k.Kontrol.DialForever()
	}

	return http.Serve(listener, k.Server)
}

func (k *Kite) registerToKontrol() {
	err := k.Kontrol.Register()
	if err != nil {
		log.Fatalf("Cannot register to Kontrol: %s", err)
	}

	log.Info("Registered to Kontrol successfully")
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Spinpunch, Inc. All Rights Reserved.
// See License.txt for license information.

package api

import (
	"bytes"
	l4g "code.google.com/p/log4go"
	"fmt"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/gorilla/mux"
	"github.com/mattermost/platform/model"
	"github.com/mattermost/platform/utils"
	"github.com/nfnt/resize"
	_ "golang.org/x/image/bmp"
	"image"
	_ "image/gif"
	"image/jpeg"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

func InitFile(r *mux.Router) {
	l4g.Debug("Initializing file api routes")

	sr := r.PathPrefix("/files").Subrouter()
	sr.Handle("/upload", ApiUserRequired(uploadFile)).Methods("POST")
	sr.Handle("/get/{channel_id:[A-Za-z0-9]+}/{user_id:[A-Za-z0-9]+}/{filename:([A-Za-z0-9]+/)?.+\\.[A-Za-z0-9]{3,}}", ApiAppHandler(getFile)).Methods("GET")
	sr.Handle("/get_public_link", ApiUserRequired(getPublicLink)).Methods("POST")
}

func uploadFile(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("uploadFile", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	err := r.ParseMultipartForm(model.MAX_FILE_SIZE)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	m := r.MultipartForm

	props := m.Value

	if len(props["channel_id"]) == 0 {
		c.SetInvalidParam("uploadFile", "channel_id")
		return
	}
	channelId := props["channel_id"][0]
	if len(channelId) == 0 {
		c.SetInvalidParam("uploadFile", "channel_id")
		return
	}

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	files := m.File["files"]

	resStruct := &model.FileUploadResponse{
		Filenames: []string{}}

	imageNameList := []string{}
	imageDataList := [][]byte{}

	if !c.HasPermissionsToChannel(cchan, "uploadFile") {
		return
	}

	for i, _ := range files {
		file, err := files[i].Open()
		defer file.Close()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		buf := bytes.NewBuffer(nil)
		io.Copy(buf, file)

		filename := filepath.Base(files[i].Filename)

		uid := model.NewId()

		path := "teams/" + c.Session.TeamId + "/channels/" + channelId + "/users/" + c.Session.UserId + "/" + uid + "/" + filename

		if err := writeFile(buf.Bytes(), path); err != nil {
			c.Err = err
			return
		}

		if model.IsFileExtImage(filepath.Ext(files[i].Filename)) {
			imageNameList = append(imageNameList, uid+"/"+filename)
			imageDataList = append(imageDataList, buf.Bytes())
		}

		encName := utils.UrlEncode(filename)

		fileUrl := "/" + channelId + "/" + c.Session.UserId + "/" + uid + "/" + encName
		resStruct.Filenames = append(resStruct.Filenames, fileUrl)
	}

	fireAndForgetHandleImages(imageNameList, imageDataList, c.Session.TeamId, channelId, c.Session.UserId)

	w.Write([]byte(resStruct.ToJson()))
}

func fireAndForgetHandleImages(filenames []string, fileData [][]byte, teamId, channelId, userId string) {

	go func() {
		dest := "teams/" + teamId + "/channels/" + channelId + "/users/" + userId + "/"

		for i, filename := range filenames {
			name := filename[:strings.LastIndex(filename, ".")]
			go func() {
				// Decode image bytes into Image object
				img, _, err := image.Decode(bytes.NewReader(fileData[i]))
				if err != nil {
					l4g.Error("Unable to decode image channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
					return
				}

				// Decode image config
				imgConfig, _, err := image.DecodeConfig(bytes.NewReader(fileData[i]))
				if err != nil {
					l4g.Error("Unable to decode image config channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
					return
				}

				// Create thumbnail
				go func() {
					var thumbnail image.Image
					if imgConfig.Width > int(utils.Cfg.ImageSettings.ThumbnailWidth) {
						thumbnail = resize.Resize(utils.Cfg.ImageSettings.ThumbnailWidth, utils.Cfg.ImageSettings.ThumbnailHeight, img, resize.Lanczos3)
					} else {
						thumbnail = img
					}

					buf := new(bytes.Buffer)
					err = jpeg.Encode(buf, thumbnail, &jpeg.Options{Quality: 90})
					if err != nil {
						l4g.Error("Unable to encode image as jpeg channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}

					if err := writeFile(buf.Bytes(), dest+name+"_thumb.jpg"); err != nil {
						l4g.Error("Unable to upload thumbnail channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}
				}()

				// Create preview
				go func() {
					var preview image.Image
					if imgConfig.Width > int(utils.Cfg.ImageSettings.PreviewWidth) {
						preview = resize.Resize(utils.Cfg.ImageSettings.PreviewWidth, utils.Cfg.ImageSettings.PreviewHeight, img, resize.Lanczos3)
					} else {
						preview = img
					}

					buf := new(bytes.Buffer)

					err = jpeg.Encode(buf, preview, &jpeg.Options{Quality: 90})
					if err != nil {
						l4g.Error("Unable to encode image as preview jpg channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}

					if err := writeFile(buf.Bytes(), dest+name+"_preview.jpg"); err != nil {
						l4g.Error("Unable to upload preview channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}
				}()
			}()
		}
	}()
}

type ImageGetResult struct {
	Error     error
	ImageData []byte
}

func getFile(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("getFile", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	params := mux.Vars(r)

	channelId := params["channel_id"]
	if len(channelId) != 26 {
		c.SetInvalidParam("getFile", "channel_id")
		return
	}

	userId := params["user_id"]
	if len(userId) != 26 {
		c.SetInvalidParam("getFile", "user_id")
		return
	}

	filename := params["filename"]
	if len(filename) == 0 {
		c.SetInvalidParam("getFile", "filename")
		return
	}

	hash := r.URL.Query().Get("h")
	data := r.URL.Query().Get("d")
	teamId := r.URL.Query().Get("t")

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	path := ""
	if len(teamId) == 26 {
		path = "teams/" + teamId + "/channels/" + channelId + "/users/" + userId + "/" + filename
	} else {
		path = "teams/" + c.Session.TeamId + "/channels/" + channelId + "/users/" + userId + "/" + filename
	}

	fileData := make(chan []byte)
	asyncGetFile(path, fileData)

	if len(hash) > 0 && len(data) > 0 && len(teamId) == 26 {
		if !model.ComparePassword(hash, fmt.Sprintf("%v:%v", data, utils.Cfg.ServiceSettings.PublicLinkSalt)) {
			c.Err = model.NewAppError("getFile", "The public link does not appear to be valid", "")
			return
		}
		props := model.MapFromJson(strings.NewReader(data))

		t, err := strconv.ParseInt(props["time"], 10, 64)
		if err != nil || model.GetMillis()-t > 1000*60*60*24*7 { // one week
			c.Err = model.NewAppError("getFile", "The public link has expired", "")
			return
		}
	} else if !c.HasPermissionsToChannel(cchan, "getFile") {
		return
	}

	f := <-fileData

	if f == nil {
		c.Err = model.NewAppError("getFile", "Could not find file.", "path="+path)
		c.Err.StatusCode = http.StatusNotFound
		return
	}

	w.Header().Set("Cache-Control", "max-age=2592000, public")
	w.Header().Set("Content-Length", strconv.Itoa(len(f)))
	w.Write(f)
}

func asyncGetFile(path string, fileData chan []byte) {
	go func() {
		data, getErr := readFile(path)
		if getErr != nil {
			l4g.Error(getErr)
			fileData <- nil
		} else {
			fileData <- data
		}
	}()
}

func getPublicLink(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.Cfg.TeamSettings.AllowPublicLink {
		c.Err = model.NewAppError("getPublicLink", "Public links have been disabled", "")
		c.Err.StatusCode = http.StatusForbidden
	}

	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("getPublicLink", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	props := model.MapFromJson(r.Body)

	filename := props["filename"]
	if len(filename) == 0 {
		c.SetInvalidParam("getPublicLink", "filename")
		return
	}

	matches := model.PartialUrlRegex.FindAllStringSubmatch(filename, -1)
	if len(matches) == 0 || len(matches[0]) < 4 {
		c.SetInvalidParam("getPublicLink", "filename")
		return
	}

	channelId := matches[0][1]
	userId := matches[0][2]
	filename = matches[0][3]

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	newProps := make(map[string]string)
	newProps["filename"] = filename
	newProps["time"] = fmt.Sprintf("%v", model.GetMillis())

	data := model.MapToJson(newProps)
	hash := model.HashPassword(fmt.Sprintf("%v:%v", data, utils.Cfg.ServiceSettings.PublicLinkSalt))

	url := fmt.Sprintf("%s/api/v1/files/get/%s/%s/%s?d=%s&h=%s&t=%s", c.GetSiteURL(), channelId, userId, filename, url.QueryEscape(data), url.QueryEscape(hash), c.Session.TeamId)

	if !c.HasPermissionsToChannel(cchan, "getPublicLink") {
		return
	}

	rData := make(map[string]string)
	rData["public_link"] = url

	w.Write([]byte(model.MapToJson(rData)))
}

func writeFile(f []byte, path string) *model.AppError {

	if utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		var auth aws.Auth
		auth.AccessKey = utils.Cfg.AWSSettings.S3AccessKeyId
		auth.SecretKey = utils.Cfg.AWSSettings.S3SecretAccessKey

		s := s3.New(auth, aws.Regions[utils.Cfg.AWSSettings.S3Region])
		bucket := s.Bucket(utils.Cfg.AWSSettings.S3Bucket)

		ext := filepath.Ext(path)

		var err error
		if model.IsFileExtImage(ext) {
			options := s3.Options{}
			err = bucket.Put(path, f, model.GetImageMimeType(ext), s3.Private, options)

		} else {
			options := s3.Options{}
			err = bucket.Put(path, f, "binary/octet-stream", s3.Private, options)
		}

		if err != nil {
			return model.NewAppError("writeFile", "Encountered an error writing to S3", err.Error())
		}
	} else if utils.Cfg.ServiceSettings.UseLocalStorage && len(utils.Cfg.ServiceSettings.StorageDirectory) > 0 {
		if err := os.MkdirAll(filepath.Dir(utils.Cfg.ServiceSettings.StorageDirectory+path), 0774); err != nil {
			return model.NewAppError("writeFile", "Encountered an error creating the directory for the new file", err.Error())
		}

		if err := ioutil.WriteFile(utils.Cfg.ServiceSettings.StorageDirectory+path, f, 0644); err != nil {
			return model.NewAppError("writeFile", "Encountered an error writing to local server storage", err.Error())
		}
	} else {
		return model.NewAppError("writeFile", "File storage not configured properly. Please configure for either S3 or local server file storage.", "")
	}

	return nil
}

func readFile(path string) ([]byte, *model.AppError) {

	if utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		var auth aws.Auth
		auth.AccessKey = utils.Cfg.AWSSettings.S3AccessKeyId
		auth.SecretKey = utils.Cfg.AWSSettings.S3SecretAccessKey

		s := s3.New(auth, aws.Regions[utils.Cfg.AWSSettings.S3Region])
		bucket := s.Bucket(utils.Cfg.AWSSettings.S3Bucket)

		// try to get the file from S3 with some basic retry logic
		tries := 0
		for {
			tries++

			f, err := bucket.Get(path)

			if f != nil {
				return f, nil
			} else if tries >= 3 {
				return nil, model.NewAppError("readFile", "Unable to get file from S3", "path="+path+", err="+err.Error())
			}
			time.Sleep(3000 * time.Millisecond)
		}
	} else if utils.Cfg.ServiceSettings.UseLocalStorage && len(utils.Cfg.ServiceSettings.StorageDirectory) > 0 {
		if f, err := ioutil.ReadFile(utils.Cfg.ServiceSettings.StorageDirectory + path); err != nil {
			return nil, model.NewAppError("readFile", "Encountered an error reading from local server storage", err.Error())
		} else {
			return f, nil
		}
	} else {
		return nil, model.NewAppError("readFile", "File storage not configured properly. Please configure for either S3 or local server file storage.", "")
	}
}
<commit_msg>Allow users to download file attachments with names that don't include a file extension<commit_after>// Copyright (c) 2015 Spinpunch, Inc. All Rights Reserved.
// See License.txt for license information.

package api

import (
	"bytes"
	l4g "code.google.com/p/log4go"
	"fmt"
	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/gorilla/mux"
	"github.com/mattermost/platform/model"
	"github.com/mattermost/platform/utils"
	"github.com/nfnt/resize"
	_ "golang.org/x/image/bmp"
	"image"
	_ "image/gif"
	"image/jpeg"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

func InitFile(r *mux.Router) {
	l4g.Debug("Initializing file api routes")

	sr := r.PathPrefix("/files").Subrouter()
	sr.Handle("/upload", ApiUserRequired(uploadFile)).Methods("POST")
	sr.Handle("/get/{channel_id:[A-Za-z0-9]+}/{user_id:[A-Za-z0-9]+}/{filename:([A-Za-z0-9]+/)?.+(\\.[A-Za-z0-9]{3,})?}", ApiAppHandler(getFile)).Methods("GET")
	sr.Handle("/get_public_link", ApiUserRequired(getPublicLink)).Methods("POST")
}

func uploadFile(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("uploadFile", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	err := r.ParseMultipartForm(model.MAX_FILE_SIZE)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	m := r.MultipartForm

	props := m.Value

	if len(props["channel_id"]) == 0 {
		c.SetInvalidParam("uploadFile", "channel_id")
		return
	}
	channelId := props["channel_id"][0]
	if len(channelId) == 0 {
		c.SetInvalidParam("uploadFile", "channel_id")
		return
	}

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	files := m.File["files"]

	resStruct := &model.FileUploadResponse{
		Filenames: []string{}}

	imageNameList := []string{}
	imageDataList := [][]byte{}

	if !c.HasPermissionsToChannel(cchan, "uploadFile") {
		return
	}

	for i, _ := range files {
		file, err := files[i].Open()
		defer file.Close()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		buf := bytes.NewBuffer(nil)
		io.Copy(buf, file)

		filename := filepath.Base(files[i].Filename)

		uid := model.NewId()

		path := "teams/" + c.Session.TeamId + "/channels/" + channelId + "/users/" + c.Session.UserId + "/" + uid + "/" + filename

		if err := writeFile(buf.Bytes(), path); err != nil {
			c.Err = err
			return
		}

		if model.IsFileExtImage(filepath.Ext(files[i].Filename)) {
			imageNameList = append(imageNameList, uid+"/"+filename)
			imageDataList = append(imageDataList, buf.Bytes())
		}

		encName := utils.UrlEncode(filename)

		fileUrl := "/" + channelId + "/" + c.Session.UserId + "/" + uid + "/" + encName
		resStruct.Filenames = append(resStruct.Filenames, fileUrl)
	}

	fireAndForgetHandleImages(imageNameList, imageDataList, c.Session.TeamId, channelId, c.Session.UserId)

	w.Write([]byte(resStruct.ToJson()))
}

func fireAndForgetHandleImages(filenames []string, fileData [][]byte, teamId, channelId, userId string) {

	go func() {
		dest := "teams/" + teamId + "/channels/" + channelId + "/users/" + userId + "/"

		for i, filename := range filenames {
			name := filename[:strings.LastIndex(filename, ".")]
			go func() {
				// Decode image bytes into Image object
				img, _, err := image.Decode(bytes.NewReader(fileData[i]))
				if err != nil {
					l4g.Error("Unable to decode image channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
					return
				}

				// Decode image config
				imgConfig, _, err := image.DecodeConfig(bytes.NewReader(fileData[i]))
				if err != nil {
					l4g.Error("Unable to decode image config channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
					return
				}

				// Create thumbnail
				go func() {
					var thumbnail image.Image
					if imgConfig.Width > int(utils.Cfg.ImageSettings.ThumbnailWidth) {
						thumbnail = resize.Resize(utils.Cfg.ImageSettings.ThumbnailWidth, utils.Cfg.ImageSettings.ThumbnailHeight, img, resize.Lanczos3)
					} else {
						thumbnail = img
					}

					buf := new(bytes.Buffer)
					err = jpeg.Encode(buf, thumbnail, &jpeg.Options{Quality: 90})
					if err != nil {
						l4g.Error("Unable to encode image as jpeg channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}

					if err := writeFile(buf.Bytes(), dest+name+"_thumb.jpg"); err != nil {
						l4g.Error("Unable to upload thumbnail channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}
				}()

				// Create preview
				go func() {
					var preview image.Image
					if imgConfig.Width > int(utils.Cfg.ImageSettings.PreviewWidth) {
						preview = resize.Resize(utils.Cfg.ImageSettings.PreviewWidth, utils.Cfg.ImageSettings.PreviewHeight, img, resize.Lanczos3)
					} else {
						preview = img
					}

					buf := new(bytes.Buffer)

					err = jpeg.Encode(buf, preview, &jpeg.Options{Quality: 90})
					if err != nil {
						l4g.Error("Unable to encode image as preview jpg channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}

					if err := writeFile(buf.Bytes(), dest+name+"_preview.jpg"); err != nil {
						l4g.Error("Unable to upload preview channelId=%v userId=%v filename=%v err=%v", channelId, userId, filename, err)
						return
					}
				}()
			}()
		}
	}()
}

type ImageGetResult struct {
	Error     error
	ImageData []byte
}

func getFile(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("getFile", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	params := mux.Vars(r)

	channelId := params["channel_id"]
	if len(channelId) != 26 {
		c.SetInvalidParam("getFile", "channel_id")
		return
	}

	userId := params["user_id"]
	if len(userId) != 26 {
		c.SetInvalidParam("getFile", "user_id")
		return
	}

	filename := params["filename"]
	if len(filename) == 0 {
		c.SetInvalidParam("getFile", "filename")
		return
	}

	hash := r.URL.Query().Get("h")
	data := r.URL.Query().Get("d")
	teamId := r.URL.Query().Get("t")

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	path := ""
	if len(teamId) == 26 {
		path = "teams/" + teamId + "/channels/" + channelId + "/users/" + userId + "/" + filename
	} else {
		path = "teams/" + c.Session.TeamId + "/channels/" + channelId + "/users/" + userId + "/" + filename
	}

	fileData := make(chan []byte)
	asyncGetFile(path, fileData)

	if len(hash) > 0 && len(data) > 0 && len(teamId) == 26 {
		if !model.ComparePassword(hash, fmt.Sprintf("%v:%v", data, utils.Cfg.ServiceSettings.PublicLinkSalt)) {
			c.Err = model.NewAppError("getFile", "The public link does not appear to be valid", "")
			return
		}
		props := model.MapFromJson(strings.NewReader(data))

		t, err := strconv.ParseInt(props["time"], 10, 64)
		if err != nil || model.GetMillis()-t > 1000*60*60*24*7 { // one week
			c.Err = model.NewAppError("getFile", "The public link has expired", "")
			return
		}
	} else if !c.HasPermissionsToChannel(cchan, "getFile") {
		return
	}

	f := <-fileData

	if f == nil {
		c.Err = model.NewAppError("getFile", "Could not find file.", "path="+path)
		c.Err.StatusCode = http.StatusNotFound
		return
	}

	w.Header().Set("Cache-Control", "max-age=2592000, public")
	w.Header().Set("Content-Length", strconv.Itoa(len(f)))
	w.Write(f)
}

func asyncGetFile(path string, fileData chan []byte) {
	go func() {
		data, getErr := readFile(path)
		if getErr != nil {
			l4g.Error(getErr)
			fileData <- nil
		} else {
			fileData <- data
		}
	}()
}

func getPublicLink(c *Context, w http.ResponseWriter, r *http.Request) {
	if !utils.Cfg.TeamSettings.AllowPublicLink {
		c.Err = model.NewAppError("getPublicLink", "Public links have been disabled", "")
		c.Err.StatusCode = http.StatusForbidden
	}

	if !utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		c.Err = model.NewAppError("getPublicLink", "Unable to upload file. Amazon S3 not configured and local server storage turned off. ", "")
		c.Err.StatusCode = http.StatusNotImplemented
		return
	}

	props := model.MapFromJson(r.Body)

	filename := props["filename"]
	if len(filename) == 0 {
		c.SetInvalidParam("getPublicLink", "filename")
		return
	}

	matches := model.PartialUrlRegex.FindAllStringSubmatch(filename, -1)
	if len(matches) == 0 || len(matches[0]) < 4 {
		c.SetInvalidParam("getPublicLink", "filename")
		return
	}

	channelId := matches[0][1]
	userId := matches[0][2]
	filename = matches[0][3]

	cchan := Srv.Store.Channel().CheckPermissionsTo(c.Session.TeamId, channelId, c.Session.UserId)

	newProps := make(map[string]string)
	newProps["filename"] = filename
	newProps["time"] = fmt.Sprintf("%v", model.GetMillis())

	data := model.MapToJson(newProps)
	hash := model.HashPassword(fmt.Sprintf("%v:%v", data, utils.Cfg.ServiceSettings.PublicLinkSalt))

	url := fmt.Sprintf("%s/api/v1/files/get/%s/%s/%s?d=%s&h=%s&t=%s", c.GetSiteURL(), channelId, userId, filename, url.QueryEscape(data), url.QueryEscape(hash), c.Session.TeamId)

	if !c.HasPermissionsToChannel(cchan, "getPublicLink") {
		return
	}

	rData := make(map[string]string)
	rData["public_link"] = url

	w.Write([]byte(model.MapToJson(rData)))
}

func writeFile(f []byte, path string) *model.AppError {

	if utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		var auth aws.Auth
		auth.AccessKey = utils.Cfg.AWSSettings.S3AccessKeyId
		auth.SecretKey = utils.Cfg.AWSSettings.S3SecretAccessKey

		s := s3.New(auth, aws.Regions[utils.Cfg.AWSSettings.S3Region])
		bucket := s.Bucket(utils.Cfg.AWSSettings.S3Bucket)

		ext := filepath.Ext(path)

		var err error
		if model.IsFileExtImage(ext) {
			options := s3.Options{}
			err = bucket.Put(path, f, model.GetImageMimeType(ext), s3.Private, options)

		} else {
			options := s3.Options{}
			err = bucket.Put(path, f, "binary/octet-stream", s3.Private, options)
		}

		if err != nil {
			return model.NewAppError("writeFile", "Encountered an error writing to S3", err.Error())
		}
	} else if utils.Cfg.ServiceSettings.UseLocalStorage && len(utils.Cfg.ServiceSettings.StorageDirectory) > 0 {
		if err := os.MkdirAll(filepath.Dir(utils.Cfg.ServiceSettings.StorageDirectory+path), 0774); err != nil {
			return model.NewAppError("writeFile", "Encountered an error creating the directory for the new file", err.Error())
		}

		if err := ioutil.WriteFile(utils.Cfg.ServiceSettings.StorageDirectory+path, f, 0644); err != nil {
			return model.NewAppError("writeFile", "Encountered an error writing to local server storage", err.Error())
		}
	} else {
		return model.NewAppError("writeFile", "File storage not configured properly. Please configure for either S3 or local server file storage.", "")
	}

	return nil
}

func readFile(path string) ([]byte, *model.AppError) {

	if utils.IsS3Configured() && !utils.Cfg.ServiceSettings.UseLocalStorage {
		var auth aws.Auth
		auth.AccessKey = utils.Cfg.AWSSettings.S3AccessKeyId
		auth.SecretKey = utils.Cfg.AWSSettings.S3SecretAccessKey

		s := s3.New(auth, aws.Regions[utils.Cfg.AWSSettings.S3Region])
		bucket := s.Bucket(utils.Cfg.AWSSettings.S3Bucket)

		// try to get the file from S3 with some basic retry logic
		tries := 0
		for {
			tries++

			f, err := bucket.Get(path)

			if f != nil {
				return f, nil
			} else if tries >= 3 {
				return nil, model.NewAppError("readFile", "Unable to get file from S3", "path="+path+", err="+err.Error())
			}
			time.Sleep(3000 * time.Millisecond)
		}
	} else if utils.Cfg.ServiceSettings.UseLocalStorage && len(utils.Cfg.ServiceSettings.StorageDirectory) > 0 {
		if f, err := ioutil.ReadFile(utils.Cfg.ServiceSettings.StorageDirectory + path); err != nil {
			return nil, model.NewAppError("readFile", "Encountered an error reading from local server storage", err.Error())
		} else {
			return f, nil
		}
	} else {
		return nil, model.NewAppError("readFile", "File storage not configured properly. Please configure for either S3 or local server file storage.", "")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/siesta/neo4j"
	"github.com/streadway/amqp"
	"io/ioutil"
	"koding/databases/mongo"
	oldNeo "koding/databases/neo4j"
	"koding/tools/amqputil"
	"koding/tools/config"
	"log"
	"net/http"
	"strconv"
	"strings"
)

type strToInf map[string]interface{}

var GRAPH_URL = config.Current.Neo4j.Write + ":" + strconv.Itoa(config.Current.Neo4j.Port)

func main() {
	amqpChannel := connectToRabbitMQ()

	coll := mongo.GetCollection("relationships")
	query := strToInf{
		"targetName": strToInf{"$nin": oldNeo.NotAllowedNames},
		"sourceName": strToInf{"$nin": oldNeo.NotAllowedNames},
	}
	iter := coll.Find(query).Batch(1000).Sort("-timestamp").Iter()

	var result oldNeo.Relationship
	for iter.Next(&result) {
		if relationshipNeedsToBeSynced(result) {
			createRelationship(result, amqpChannel)
		}
	}

	log.Println("Neo4j is now synced with Mongodb.")
}

func connectToRabbitMQ() *amqp.Channel {
	conn := amqputil.CreateConnection("syncWorker")
	amqpChannel, err := conn.Channel()
	if err != nil {
		panic(err)
	}
	return amqpChannel
}

func createRelationship(rel oldNeo.Relationship, amqpChannel *amqp.Channel) {
	data := make([]strToInf, 1)
	data[0] = strToInf{
		"_id":        rel.Id,
		"sourceId":   rel.SourceId,
		"sourceName": rel.SourceName,
		"targetId":   rel.TargetId,
		"targetName": rel.TargetName,
		"as":         rel.As,
	}

	eventData := strToInf{"event": "RelationshipSaved", "payload": data}

	neoMessage, err := json.Marshal(eventData)
	if err != nil {
		log.Println("unmarshall error")
		return
	}

	amqpChannel.Publish(
		"graphFeederExchange", // exchange name
		"",    // key
		false, // mandatory
		false, // immediate
		amqp.Publishing{
			Body: neoMessage,
		},
	)
}

func relationshipNeedsToBeSynced(result oldNeo.Relationship) bool {
	exists, sourceId := checkNodeExists(result.SourceId.Hex())
	if exists != true {
		log.Println("err: No SourceNode:", result.SourceName, result.As)
		return true
	}

	exists, targetId := checkNodeExists(result.TargetId.Hex())
	if exists != true {
		log.Println("err: No TargetNode:", result.TargetName, result.Id, result.As)
		return true
	}

	exists = checkRelationshipExists(sourceId, targetId, result.As)
	if exists != true {
		log.Printf("err: No '%v' relationship exists between %v and %v", result.SourceName, result.TargetName, result.As)
		return true
	}

	// everything is fine
	return false
}

func getAndParse(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func checkRelationshipExists(sourceId, targetId, relType string) bool {
	url := fmt.Sprintf("%v/db/data/node/%v/relationships/all/%v", GRAPH_URL, sourceId, relType)

	body, err := getAndParse(url)
	if err != nil {
		return false
	}

	relResponse := make([]neo4j.RelationshipResponse, 1)
	err = json.Unmarshal(body, &relResponse)
	if err != nil {
		return false
	}

	for _, rl := range relResponse {
		id := strings.SplitAfter(rl.End, GRAPH_URL+"/db/data/node/")[1]
		if targetId == id {
			return true
		}
	}

	return false
}

func checkNodeExists(id string) (bool, string) {
	url := fmt.Sprintf("%v/db/data/index/node/koding/id/%v", GRAPH_URL, id)
	body, err := getAndParse(url)
	if err != nil {
		return false, ""
	}

	nodeResponse := make([]neo4j.NodeResponse, 1)
	err = json.Unmarshal(body, &nodeResponse)
	if err != nil {
		return false, ""
	}

	if len(nodeResponse) < 1 {
		return false, ""
	}

	node := nodeResponse[0]
	idd := strings.SplitAfter(node.Self, GRAPH_URL+"/db/data/node/")

	nodeId := string(idd[1])
	if nodeId == "" {
		return false, ""
	}

	return true, nodeId
}
<commit_msg>sync: better logging<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/siesta/neo4j"
	"github.com/streadway/amqp"
	"io/ioutil"
	"koding/databases/mongo"
	oldNeo "koding/databases/neo4j"
	"koding/tools/amqputil"
	"koding/tools/config"
	"log"
	"net/http"
	"strconv"
	"strings"
)

type strToInf map[string]interface{}

var GRAPH_URL = config.Current.Neo4j.Write + ":" + strconv.Itoa(config.Current.Neo4j.Port)

func main() {
	amqpChannel := connectToRabbitMQ()

	coll := mongo.GetCollection("relationships")
	query := strToInf{
		"targetName": strToInf{"$nin": oldNeo.NotAllowedNames},
		"sourceName": strToInf{"$nin": oldNeo.NotAllowedNames},
	}
	iter := coll.Find(query).Batch(1000).Sort("-timestamp").Iter()

	var result oldNeo.Relationship
	for iter.Next(&result) {
		if relationshipNeedsToBeSynced(result) {
			createRelationship(result, amqpChannel)
		}
	}

	log.Println("Neo4j is now synced with Mongodb.")
}

func connectToRabbitMQ() *amqp.Channel {
	conn := amqputil.CreateConnection("syncWorker")
	amqpChannel, err := conn.Channel()
	if err != nil {
		panic(err)
	}
	return amqpChannel
}

func createRelationship(rel oldNeo.Relationship, amqpChannel *amqp.Channel) {
	data := make([]strToInf, 1)
	data[0] = strToInf{
		"_id":        rel.Id,
		"sourceId":   rel.SourceId,
		"sourceName": rel.SourceName,
		"targetId":   rel.TargetId,
		"targetName": rel.TargetName,
		"as":         rel.As,
	}

	eventData := strToInf{"event": "RelationshipSaved", "payload": data}

	neoMessage, err := json.Marshal(eventData)
	if err != nil {
		log.Println("unmarshall error")
		return
	}

	amqpChannel.Publish(
		"graphFeederExchange", // exchange name
		"",    // key
		false, // mandatory
		false, // immediate
		amqp.Publishing{
			Body: neoMessage,
		},
	)
}

func relationshipNeedsToBeSynced(result oldNeo.Relationship) bool {
	exists, sourceId := checkNodeExists(result.SourceId.Hex())
	if exists != true {
		log.Println("err: Id '%v'. No SourceNode '%v' with Id: '%v' for relationship: '%v'", result.Id, result.SourceName, result.SourceId, result.As)
		return true
	}

	exists, targetId := checkNodeExists(result.TargetId.Hex())
	if exists != true {
		log.Println("err: Id '%v'. No TargetNode '%v' with Id: '%v' for relationship: '%v'", result.Id, result.TargetName, result.TargetId, result.As)
		return true
	}

	exists = checkRelationshipExists(sourceId, targetId, result.As)
	if exists != true {
		log.Printf("err: Id '%v'. No '%v' relationship exists between %v and %v", result.Id, result.As, result.SourceName, result.TargetName)
		return true
	}

	// everything is fine
	return false
}

func getAndParse(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func checkRelationshipExists(sourceId, targetId, relType string) bool {
	url := fmt.Sprintf("%v/db/data/node/%v/relationships/all/%v", GRAPH_URL, sourceId, relType)

	body, err := getAndParse(url)
	if err != nil {
		return false
	}

	relResponse := make([]neo4j.RelationshipResponse, 1)
	err = json.Unmarshal(body, &relResponse)
	if err != nil {
		return false
	}

	for _, rl := range relResponse {
		id := strings.SplitAfter(rl.End, GRAPH_URL+"/db/data/node/")[1]
		if targetId == id {
			return true
		}
	}

	return false
}

func checkNodeExists(id string) (bool, string) {
	url := fmt.Sprintf("%v/db/data/index/node/koding/id/%v", GRAPH_URL, id)
	body, err := getAndParse(url)
	if err != nil {
		return false, ""
	}

	nodeResponse := make([]neo4j.NodeResponse, 1)
	err = json.Unmarshal(body, &nodeResponse)
	if err != nil {
		return false, ""
	}

	if len(nodeResponse) < 1 {
		return false, ""
	}

	node := nodeResponse[0]
	idd := strings.SplitAfter(node.Self, GRAPH_URL+"/db/data/node/")

	nodeId := string(idd[1])
	if nodeId == "" {
		return false, ""
	}

	return true, nodeId
}
<|endoftext|>