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

// Package cors provides handlers to enable CORS support.
package cors

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	headerAllowOrigin      = "Access-Control-Allow-Origin"
	headerAllowCredentials = "Access-Control-Allow-Credentials"
	headerAllowHeaders     = "Access-Control-Allow-Headers"
	headerAllowMethods     = "Access-Control-Allow-Methods"
	headerExposeHeaders    = "Access-Control-Expose-Headers"
	headerMaxAge           = "Access-Control-Max-Age"

	headerOrigin         = "Origin"
	headerRequestMethod  = "Access-Control-Request-Method"
	headerRequestHeaders = "Access-Control-Request-Headers"
)

var (
	defaultAllowHeaders = []string{"Origin", "Accept", "Content-Type", "Authorization"}
	// Regex patterns are generated from AllowOrigins. These are used and generated internally.
	allowOriginPatterns = []string{}
)

// Options represents Access Control options.
type Options struct {
	// If set, all origins are allowed.
	AllowAllOrigins bool
	// A list of allowed origins. Wild cards and FQDNs are supported.
	AllowOrigins []string
	// If set, allows to share auth credentials such as cookies.
	AllowCredentials bool
	// A list of allowed HTTP methods.
	AllowMethods []string
	// A list of allowed HTTP headers.
	AllowHeaders []string
	// A list of exposed HTTP headers.
	ExposeHeaders []string
	// Max age of the CORS headers.
	MaxAge time.Duration
}

// Header converts options into CORS headers.
func (o *Options) Header(origin string) (headers map[string]string) {
	headers = make(map[string]string)
	// if origin is not allowed, don't extend the headers
	// with CORS headers.
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}

	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allow credentials
	headers[headerAllowCredentials] = strconv.FormatBool(o.AllowCredentials)

	// add allow methods
	if len(o.AllowMethods) > 0 {
		headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
	}

	// add allow headers
	if len(o.AllowHeaders) > 0 {
		// TODO: Add default headers
		headers[headerAllowHeaders] = strings.Join(o.AllowHeaders, ",")
	}

	// add exposed header
	if len(o.ExposeHeaders) > 0 {
		headers[headerExposeHeaders] = strings.Join(o.ExposeHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

// PreflightHeader converts options into CORS headers for a preflight response.
func (o *Options) PreflightHeader(origin, rMethod, rHeaders string) (headers map[string]string) {
	headers = make(map[string]string)
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}
	// verify if requested method is allowed
	// TODO: Too many for loops
	for _, method := range o.AllowMethods {
		if method == rMethod {
			headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
			break
		}
	}

	// verify if requested headers are allowed
	var allowed []string
	for _, rHeader := range strings.Split(rHeaders, ",") {
		rHeader = strings.TrimSpace(rHeader)
	lookupLoop:
		for _, allowedHeader := range o.AllowHeaders {
			if strings.ToLower(rHeader) == strings.ToLower(allowedHeader) {
				allowed = append(allowed, rHeader)
				break lookupLoop
			}
		}
	}

	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allowed headers
	if len(allowed) > 0 {
		headers[headerAllowHeaders] = strings.Join(allowed, ",")
	}

	// add exposed headers
	if len(o.ExposeHeaders) > 0 {
		headers[headerExposeHeaders] = strings.Join(o.ExposeHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

// IsOriginAllowed looks up if the origin matches one of the patterns
// generated from Options.AllowOrigins patterns.
func (o *Options) IsOriginAllowed(origin string) (allowed bool) {
	for _, pattern := range allowOriginPatterns {
		allowed, _ = regexp.MatchString(pattern, origin)
		if allowed {
			return
		}
	}
	return
}

// Allow enables CORS for requests those match the provided options.
func Allow(opts *Options) http.HandlerFunc {
	// Allow default headers if nothing is specified.
	if len(opts.AllowHeaders) == 0 {
		opts.AllowHeaders = defaultAllowHeaders
	}

	for _, origin := range opts.AllowOrigins {
		pattern := regexp.QuoteMeta(origin)
		pattern = strings.Replace(pattern, "\\*", ".*", -1)
		pattern = strings.Replace(pattern, "\\?", ".", -1)
		allowOriginPatterns = append(allowOriginPatterns, "^"+pattern+"$")
	}

	return func(res http.ResponseWriter, req *http.Request) {
		var (
			origin           = req.Header.Get(headerOrigin)
			requestedMethod  = req.Header.Get(headerRequestMethod)
			requestedHeaders = req.Header.Get(headerRequestHeaders)
			// additional headers to be added
			// to the response.
			headers map[string]string
		)

		if req.Method == "OPTIONS" &&
			(requestedMethod != "" || requestedHeaders != "") {
			// TODO: if preflight, respond with exact headers if allowed
			headers = opts.PreflightHeader(origin, requestedMethod, requestedHeaders)
			for key, value := range headers {
				res.Header().Set(key, value)
			}
			res.WriteHeader(http.StatusOK)
			return
		}
		headers = opts.Header(origin)

		for key, value := range headers {
			res.Header().Set(key, value)
		}
	}
}
<commit_msg>add credentials headers on OPTIONS request<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cors provides handlers to enable CORS support.
package cors

import (
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	headerAllowOrigin      = "Access-Control-Allow-Origin"
	headerAllowCredentials = "Access-Control-Allow-Credentials"
	headerAllowHeaders     = "Access-Control-Allow-Headers"
	headerAllowMethods     = "Access-Control-Allow-Methods"
	headerExposeHeaders    = "Access-Control-Expose-Headers"
	headerMaxAge           = "Access-Control-Max-Age"

	headerOrigin         = "Origin"
	headerRequestMethod  = "Access-Control-Request-Method"
	headerRequestHeaders = "Access-Control-Request-Headers"
)

var (
	defaultAllowHeaders = []string{"Origin", "Accept", "Content-Type", "Authorization"}
	// Regex patterns are generated from AllowOrigins. These are used and generated internally.
	allowOriginPatterns = []string{}
)

// Options represents Access Control options.
type Options struct {
	// If set, all origins are allowed.
	AllowAllOrigins bool
	// A list of allowed origins. Wild cards and FQDNs are supported.
	AllowOrigins []string
	// If set, allows to share auth credentials such as cookies.
	AllowCredentials bool
	// A list of allowed HTTP methods.
	AllowMethods []string
	// A list of allowed HTTP headers.
	AllowHeaders []string
	// A list of exposed HTTP headers.
	ExposeHeaders []string
	// Max age of the CORS headers.
	MaxAge time.Duration
}

// Header converts options into CORS headers.
func (o *Options) Header(origin string) (headers map[string]string) {
	headers = make(map[string]string)
	// if origin is not allowed, don't extend the headers
	// with CORS headers.
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}

	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allow credentials
	headers[headerAllowCredentials] = strconv.FormatBool(o.AllowCredentials)

	// add allow methods
	if len(o.AllowMethods) > 0 {
		headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
	}

	// add allow headers
	if len(o.AllowHeaders) > 0 {
		// TODO: Add default headers
		headers[headerAllowHeaders] = strings.Join(o.AllowHeaders, ",")
	}

	// add exposed header
	if len(o.ExposeHeaders) > 0 {
		headers[headerExposeHeaders] = strings.Join(o.ExposeHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

// PreflightHeader converts options into CORS headers for a preflight response.
func (o *Options) PreflightHeader(origin, rMethod, rHeaders string) (headers map[string]string) {
	headers = make(map[string]string)
	if !o.AllowAllOrigins && !o.IsOriginAllowed(origin) {
		return
	}
	// verify if requested method is allowed
	// TODO: Too many for loops
	for _, method := range o.AllowMethods {
		if method == rMethod {
			headers[headerAllowMethods] = strings.Join(o.AllowMethods, ",")
			break
		}
	}

	// verify if requested headers are allowed
	var allowed []string
	for _, rHeader := range strings.Split(rHeaders, ",") {
		rHeader = strings.TrimSpace(rHeader)
	lookupLoop:
		for _, allowedHeader := range o.AllowHeaders {
			if strings.ToLower(rHeader) == strings.ToLower(allowedHeader) {
				allowed = append(allowed, rHeader)
				break lookupLoop
			}
		}
	}

	headers[headerAllowCredentials] = strconv.FormatBool(o.AllowCredentials)
	// add allow origin
	if o.AllowAllOrigins {
		headers[headerAllowOrigin] = "*"
	} else {
		headers[headerAllowOrigin] = origin
	}

	// add allowed headers
	if len(allowed) > 0 {
		headers[headerAllowHeaders] = strings.Join(allowed, ",")
	}

	// add exposed headers
	if len(o.ExposeHeaders) > 0 {
		headers[headerExposeHeaders] = strings.Join(o.ExposeHeaders, ",")
	}
	// add a max age header
	if o.MaxAge > time.Duration(0) {
		headers[headerMaxAge] = strconv.FormatInt(int64(o.MaxAge/time.Second), 10)
	}
	return
}

// IsOriginAllowed looks up if the origin matches one of the patterns
// generated from Options.AllowOrigins patterns.
func (o *Options) IsOriginAllowed(origin string) (allowed bool) {
	for _, pattern := range allowOriginPatterns {
		allowed, _ = regexp.MatchString(pattern, origin)
		if allowed {
			return
		}
	}
	return
}

// Allow enables CORS for requests those match the provided options.
func Allow(opts *Options) http.HandlerFunc {
	// Allow default headers if nothing is specified.
	if len(opts.AllowHeaders) == 0 {
		opts.AllowHeaders = defaultAllowHeaders
	}

	for _, origin := range opts.AllowOrigins {
		pattern := regexp.QuoteMeta(origin)
		pattern = strings.Replace(pattern, "\\*", ".*", -1)
		pattern = strings.Replace(pattern, "\\?", ".", -1)
		allowOriginPatterns = append(allowOriginPatterns, "^"+pattern+"$")
	}

	return func(res http.ResponseWriter, req *http.Request) {
		var (
			origin           = req.Header.Get(headerOrigin)
			requestedMethod  = req.Header.Get(headerRequestMethod)
			requestedHeaders = req.Header.Get(headerRequestHeaders)
			// additional headers to be added
			// to the response.
			headers map[string]string
		)

		if req.Method == "OPTIONS" &&
			(requestedMethod != "" || requestedHeaders != "") {
			// TODO: if preflight, respond with exact headers if allowed
			headers = opts.PreflightHeader(origin, requestedMethod, requestedHeaders)
			for key, value := range headers {
				res.Header().Set(key, value)
			}
			res.WriteHeader(http.StatusOK)
			return
		}
		headers = opts.Header(origin)

		for key, value := range headers {
			res.Header().Set(key, value)
		}
	}
}
<|endoftext|>


<commit_before>package gron

import (
	"sort"
	"time"
)

// Entry consists of a schedule and the job to be executed on that schedule.
type Entry struct {
	Schedule Schedule
	Job      Job

	// the next time the job will run. This is zero time if Cron has not been
	// started or invalid schedule.
	Next time.Time

	// the last time the job was run. This is zero time if the job has not been
	// run.
	Prev time.Time
}

// byTime is a handy wrapper to chronologically sort entries.
type byTime []*Entry

func (b byTime) Len() int      { return len(b) }
func (b byTime) Swap(i, j int) { b[i], b[j] = b[j], b[i] }

// Less reports `earliest` time i should sort before j.
// zero time is not `earliest` time.
func (b byTime) Less(i, j int) bool {

	if b[i].Next.IsZero() {
		return false
	}
	if b[j].Next.IsZero() {
		return true
	}

	return b[i].Next.Before(b[j].Next)
}

// Job is the interface that wraps the basic Run method.
//
// Run executes the underlying func.
type Job interface {
	Run()
}

// Cron provides a convenient interface for scheduling job such as to clean-up
// database entry every month.
//
// Cron keeps track of any number of entries, invoking the associated func as
// specified by the schedule. It may also be started, stopped and the entries
// may be inspected.
type Cron struct {
	entries []*Entry
	running bool
	add     chan *Entry
	stop    chan struct{}
}

// New instantiates new Cron instant c.
func New() *Cron {
	return &Cron{
		stop: make(chan struct{}),
		add:  make(chan *Entry),
	}
}

// Start signals cron instant c to get up and running.
func (c *Cron) Start() {
	c.running = true
	go c.run()
}

// Add appends schedule, job to entries.
//
// if cron instant is not running, adding to entries is trivial.
// otherwise, to prevent data-race, adds through channel.
func (c *Cron) Add(s Schedule, j Job) {

	entry := &Entry{
		Schedule: s,
		Job:      j,
	}

	if !c.running {
		c.entries = append(c.entries, entry)
	}
	c.add <- entry
}

// AddFunc registers the Job function for the given Schedule.
func (c *Cron) AddFunc(s Schedule, j func()) {
	c.Add(s, JobFunc(j))
}

// Stop halts cron instant c from running.
func (c *Cron) Stop() {

	if !c.running {
		return
	}
	c.running = false
	c.stop <- struct{}{}
}

// run the scheduler...
//
// It needs to be private as it's responsible of synchronizing a critical
// shared state: `running`.
func (c *Cron) run() {

	var effective time.Time

	for _, e := range c.entries {
		e.Next = e.Schedule.Next(time.Now())
	}

	for {

		sort.Sort(byTime(c.entries))
		if len(c.entries) > 0 {
			effective = c.entries[0].Next
		}

		select {
		case now := <-time.After(effective.Sub(time.Now())):
			go c.entries[0].Job.Run()
			c.entries[0].Prev = now
			c.entries[0].Next = c.entries[0].Schedule.Next(now)

		case e := <-c.add:
			e.Next = e.Schedule.Next(time.Now())
			c.entries = append(c.entries, e)
		case <-c.stop:
			return // terminate go-routine.
		}
	}
}

// JobFunc is an adapter to allow the use of ordinary functions as gron.Job
// If f is a function with the appropriate signature, JobFunc(f) is a handler
// that calls f.
//
// todo: possibly func with params? maybe not needed.
type JobFunc func()

// Run calls j()
func (j JobFunc) Run() {
	j()
}
<commit_msg>return after adding entry, prior running<commit_after>package gron

import (
	"sort"
	"time"
)

// Entry consists of a schedule and the job to be executed on that schedule.
type Entry struct {
	Schedule Schedule
	Job      Job

	// the next time the job will run. This is zero time if Cron has not been
	// started or invalid schedule.
	Next time.Time

	// the last time the job was run. This is zero time if the job has not been
	// run.
	Prev time.Time
}

// byTime is a handy wrapper to chronologically sort entries.
type byTime []*Entry

func (b byTime) Len() int      { return len(b) }
func (b byTime) Swap(i, j int) { b[i], b[j] = b[j], b[i] }

// Less reports `earliest` time i should sort before j.
// zero time is not `earliest` time.
func (b byTime) Less(i, j int) bool {

	if b[i].Next.IsZero() {
		return false
	}
	if b[j].Next.IsZero() {
		return true
	}

	return b[i].Next.Before(b[j].Next)
}

// Job is the interface that wraps the basic Run method.
//
// Run executes the underlying func.
type Job interface {
	Run()
}

// Cron provides a convenient interface for scheduling job such as to clean-up
// database entry every month.
//
// Cron keeps track of any number of entries, invoking the associated func as
// specified by the schedule. It may also be started, stopped and the entries
// may be inspected.
type Cron struct {
	entries []*Entry
	running bool
	add     chan *Entry
	stop    chan struct{}
}

// New instantiates new Cron instant c.
func New() *Cron {
	return &Cron{
		stop: make(chan struct{}),
		add:  make(chan *Entry),
	}
}

// Start signals cron instant c to get up and running.
func (c *Cron) Start() {
	c.running = true
	go c.run()
}

// Add appends schedule, job to entries.
//
// if cron instant is not running, adding to entries is trivial.
// otherwise, to prevent data-race, adds through channel.
func (c *Cron) Add(s Schedule, j Job) {

	entry := &Entry{
		Schedule: s,
		Job:      j,
	}

	if !c.running {
		c.entries = append(c.entries, entry)
		return
	}
	c.add <- entry
}

// AddFunc registers the Job function for the given Schedule.
func (c *Cron) AddFunc(s Schedule, j func()) {
	c.Add(s, JobFunc(j))
}

// Stop halts cron instant c from running.
func (c *Cron) Stop() {

	if !c.running {
		return
	}
	c.running = false
	c.stop <- struct{}{}
}

// run the scheduler...
//
// It needs to be private as it's responsible of synchronizing a critical
// shared state: `running`.
func (c *Cron) run() {

	var effective time.Time

	for _, e := range c.entries {
		e.Next = e.Schedule.Next(time.Now())
	}

	for {

		sort.Sort(byTime(c.entries))
		if len(c.entries) > 0 {
			effective = c.entries[0].Next
		}

		select {
		case now := <-time.After(effective.Sub(time.Now())):
			go c.entries[0].Job.Run()
			c.entries[0].Prev = now
			c.entries[0].Next = c.entries[0].Schedule.Next(now)

		case e := <-c.add:
			e.Next = e.Schedule.Next(time.Now())
			c.entries = append(c.entries, e)
		case <-c.stop:
			return // terminate go-routine.
		}
	}
}

// JobFunc is an adapter to allow the use of ordinary functions as gron.Job
// If f is a function with the appropriate signature, JobFunc(f) is a handler
// that calls f.
//
// todo: possibly func with params? maybe not needed.
type JobFunc func()

// Run calls j()
func (j JobFunc) Run() {
	j()
}
<|endoftext|>


<commit_before>package core

import (
	"sync"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringSubscriptionBroadcaster interface {
	model.Stopable
	model.Startable

	HandleStopMonitoringBroadcastEvent(*model.StopMonitoringBroadcastEvent)
	HandleSubscriptionRequest([]*siri.XMLStopMonitoringSubscriptionRequestEntry)
}

type SIRIStopMonitoringSubscriptionBroadcaster struct {
	model.ClockConsumer
	model.UUIDConsumer

	siriConnector

	stopMonitoringBroadcaster SIRIStopMonitoringBroadcaster
	toBroadcast               map[SubscriptionId][]model.StopVisitId
	mutex                     *sync.Mutex //protect the map
}

type SIRIStopMonitoringSubscriptionBroadcasterFactory struct{}

func (factory *SIRIStopMonitoringSubscriptionBroadcasterFactory) CreateConnector(partner *Partner) Connector {
	if _, ok := partner.Connector(SIRI_SUBSCRIPTION_REQUEST_DISPATCHER); !ok {
		partner.CreateSubscriptionRequestDispatcher()
	}
	return newSIRIStopMonitoringSubscriptionBroadcaster(partner)
}

func (factory *SIRIStopMonitoringSubscriptionBroadcasterFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func newSIRIStopMonitoringSubscriptionBroadcaster(partner *Partner) *SIRIStopMonitoringSubscriptionBroadcaster {
	siriStopMonitoringSubscriptionBroadcaster := &SIRIStopMonitoringSubscriptionBroadcaster{}
	siriStopMonitoringSubscriptionBroadcaster.partner = partner
	siriStopMonitoringSubscriptionBroadcaster.mutex = &sync.Mutex{}
	siriStopMonitoringSubscriptionBroadcaster.toBroadcast = make(map[SubscriptionId][]model.StopVisitId)

	siriStopMonitoringSubscriptionBroadcaster.stopMonitoringBroadcaster = NewSIRIStopMonitoringBroadcaster(siriStopMonitoringSubscriptionBroadcaster)

	return siriStopMonitoringSubscriptionBroadcaster
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) Stop() {
	connector.stopMonitoringBroadcaster.Stop()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) Start() {
	connector.stopMonitoringBroadcaster.Start()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) HandleStopMonitoringBroadcastEvent(event *model.StopMonitoringBroadcastEvent) {
	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	switch event.ModelType {
	case "StopVisit":
		sv, ok := tx.Model().StopVisits().Find(model.StopVisitId(event.ModelId))
		subId, ok := connector.checkEvent(sv, tx)
		if ok {
			connector.addStopVisit(subId, sv.Id())
		}
	case "VehicleJourney":
		for _, sv := range tx.Model().StopVisits().FindFollowingByVehicleJourneyId(model.VehicleJourneyId(event.ModelId)) {
			subId, ok := connector.checkEvent(sv, tx)
			if ok {
				connector.addStopVisit(subId, sv.Id())
			}
		}
	default:
		return
	}
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) addStopVisit(subId SubscriptionId, svId model.StopVisitId) {
	connector.mutex.Lock()
	connector.toBroadcast[SubscriptionId(subId)] = append(connector.toBroadcast[SubscriptionId(subId)], svId)
	connector.mutex.Unlock()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) checkEvent(sv model.StopVisit, tx *model.Transaction) (SubscriptionId, bool) {
	subId := SubscriptionId(0) //just to return a correct type for errors

	stopArea, ok := tx.Model().StopAreas().Find(sv.StopAreaId)
	if !ok {
		return subId, false
	}

	obj, ok := stopArea.ObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER))
	if !ok {
		return subId, false
	}

	_, ok = stopArea.ObjectID(connector.Partner().Setting("remote_objectid_kind"))

	sub, ok := connector.partner.Subscriptions().FindByRessourceId(obj.String())
	if !ok {
		return subId, false
	}

	resources := sub.ResourcesByObjectID()

	resource, ok := resources[obj.String()]

	if !ok || resource.SubscribedUntil.Before(connector.Clock().Now()) {
		return subId, false
	}

	lastState, ok := resource.LastStates[string(sv.Id())]

	if ok && !lastState.(*stopMonitoringLastChange).Haschanged(sv) {
		return subId, false
	}

	if !ok {
		smlc := &stopMonitoringLastChange{}
		smlc.SetSubscription(sub)
		resource.LastStates[string(sv.Id())] = smlc
	}

	return sub.Id(), true
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) HandleSubscriptionRequest(request *siri.XMLSubscriptionRequest) []siri.SIRIResponseStatus {
	sms := request.XMLSubscriptionSMEntries()

	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	resps := []siri.SIRIResponseStatus{}

	for _, sm := range sms {
		rs := siri.SIRIResponseStatus{
			RequestMessageRef: sm.MessageIdentifier(),
			SubscriberRef:     sm.SubscriberRef(),
			SubscriptionRef:   sm.SubscriptionIdentifier(),
			ResponseTimestamp: connector.Clock().Now(),
		}

		objectid := model.NewObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER), sm.MonitoringRef())
		sa, ok := tx.Model().StopAreas().FindByObjectId(objectid)
		if !ok {
			resps = append(resps, rs)
			continue
		}

		sub, ok := connector.Partner().Subscriptions().FindByExternalId(sm.SubscriptionIdentifier())
		if !ok {
			sub = connector.Partner().Subscriptions().New("StopArea")
			sub.SetExternalId(sm.SubscriptionIdentifier())
		}

		ref := model.Reference{
			ObjectId: &objectid,
			Id:       string(sa.Id()),
			Type:     "StopArea",
		}

		r := sub.CreateAddNewResource(ref)
		r.SubscribedUntil = sm.InitialTerminationTime()

		connector.fillOptions(sub, r, request, sm)

		rs.Status = true
		rs.ValidUntil = sm.InitialTerminationTime()
		resps = append(resps, rs)

		connector.AddStopAreaStopVisits(sa, sub, r)
	}
	return resps
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) AddStopAreaStopVisits(sa model.StopArea, sub *Subscription, res *SubscribedResource) {
	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	svs := tx.Model().StopVisits().FindFollowingByStopAreaId(sa.Id())
	for _, sv := range svs {
		_, ok := sv.ObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER))
		if !ok {
			continue
		}
		smlc := &stopMonitoringLastChange{}
		smlc.SetSubscription(sub)
		res.LastStates[string(sv.Id())] = smlc
		connector.addStopVisit(sub.Id(), sv.Id())
	}
}

func (smsb *SIRIStopMonitoringSubscriptionBroadcaster) fillOptions(s *Subscription, r *SubscribedResource, request *siri.XMLSubscriptionRequest, sm *siri.XMLStopMonitoringSubscriptionRequestEntry) {
	ro := r.ResourcesOptions()
	ro["StopVisitTypes"] = sm.StopVisitTypes()

	so := s.SubscriptionOptions()

	so["IncrementalUpdates"] = request.IncrementalUpdates()
	so["MaximumStopVisits"] = request.MaximumStopVisits()
	so["ChangeBeforeUpdates"] = request.ChangeBeforeUpdates()
}

// START TEST

type TestSIRIStopMonitoringSubscriptionBroadcasterFactory struct{}

type TestStopMonitoringSubscriptionBroadcaster struct {
	model.UUIDConsumer

	events                    []*model.StopMonitoringBroadcastEvent
	stopMonitoringBroadcaster SIRIStopMonitoringBroadcaster
}

func NewTestStopMonitoringSubscriptionBroadcaster() *TestStopMonitoringSubscriptionBroadcaster {
	connector := &TestStopMonitoringSubscriptionBroadcaster{}
	return connector
}

func (connector *TestStopMonitoringSubscriptionBroadcaster) HandleStopMonitoringBroadcastEvent(event *model.StopMonitoringBroadcastEvent) {
	connector.events = append(connector.events, event)
}

func (factory *TestSIRIStopMonitoringSubscriptionBroadcasterFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestSIRIStopMonitoringSubscriptionBroadcasterFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringSubscriptionBroadcaster()
}

// END TEST
<commit_msg>Wrong stopArea subscription error message<commit_after>package core

import (
	"sync"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

type StopMonitoringSubscriptionBroadcaster interface {
	model.Stopable
	model.Startable

	HandleStopMonitoringBroadcastEvent(*model.StopMonitoringBroadcastEvent)
	HandleSubscriptionRequest([]*siri.XMLStopMonitoringSubscriptionRequestEntry)
}

type SIRIStopMonitoringSubscriptionBroadcaster struct {
	model.ClockConsumer
	model.UUIDConsumer

	siriConnector

	stopMonitoringBroadcaster SIRIStopMonitoringBroadcaster
	toBroadcast               map[SubscriptionId][]model.StopVisitId
	mutex                     *sync.Mutex //protect the map
}

type SIRIStopMonitoringSubscriptionBroadcasterFactory struct{}

func (factory *SIRIStopMonitoringSubscriptionBroadcasterFactory) CreateConnector(partner *Partner) Connector {
	if _, ok := partner.Connector(SIRI_SUBSCRIPTION_REQUEST_DISPATCHER); !ok {
		partner.CreateSubscriptionRequestDispatcher()
	}
	return newSIRIStopMonitoringSubscriptionBroadcaster(partner)
}

func (factory *SIRIStopMonitoringSubscriptionBroadcasterFactory) Validate(apiPartner *APIPartner) bool {
	ok := apiPartner.ValidatePresenceOfSetting("remote_objectid_kind")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_url")
	ok = ok && apiPartner.ValidatePresenceOfSetting("remote_credential")
	return ok
}

func newSIRIStopMonitoringSubscriptionBroadcaster(partner *Partner) *SIRIStopMonitoringSubscriptionBroadcaster {
	siriStopMonitoringSubscriptionBroadcaster := &SIRIStopMonitoringSubscriptionBroadcaster{}
	siriStopMonitoringSubscriptionBroadcaster.partner = partner
	siriStopMonitoringSubscriptionBroadcaster.mutex = &sync.Mutex{}
	siriStopMonitoringSubscriptionBroadcaster.toBroadcast = make(map[SubscriptionId][]model.StopVisitId)

	siriStopMonitoringSubscriptionBroadcaster.stopMonitoringBroadcaster = NewSIRIStopMonitoringBroadcaster(siriStopMonitoringSubscriptionBroadcaster)

	return siriStopMonitoringSubscriptionBroadcaster
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) Stop() {
	connector.stopMonitoringBroadcaster.Stop()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) Start() {
	connector.stopMonitoringBroadcaster.Start()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) HandleStopMonitoringBroadcastEvent(event *model.StopMonitoringBroadcastEvent) {
	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	switch event.ModelType {
	case "StopVisit":
		sv, ok := tx.Model().StopVisits().Find(model.StopVisitId(event.ModelId))
		subId, ok := connector.checkEvent(sv, tx)
		if ok {
			connector.addStopVisit(subId, sv.Id())
		}
	case "VehicleJourney":
		for _, sv := range tx.Model().StopVisits().FindFollowingByVehicleJourneyId(model.VehicleJourneyId(event.ModelId)) {
			subId, ok := connector.checkEvent(sv, tx)
			if ok {
				connector.addStopVisit(subId, sv.Id())
			}
		}
	default:
		return
	}
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) addStopVisit(subId SubscriptionId, svId model.StopVisitId) {
	connector.mutex.Lock()
	connector.toBroadcast[SubscriptionId(subId)] = append(connector.toBroadcast[SubscriptionId(subId)], svId)
	connector.mutex.Unlock()
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) checkEvent(sv model.StopVisit, tx *model.Transaction) (SubscriptionId, bool) {
	subId := SubscriptionId(0) //just to return a correct type for errors

	stopArea, ok := tx.Model().StopAreas().Find(sv.StopAreaId)
	if !ok {
		return subId, false
	}

	obj, ok := stopArea.ObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER))
	if !ok {
		return subId, false
	}

	_, ok = stopArea.ObjectID(connector.Partner().Setting("remote_objectid_kind"))

	sub, ok := connector.partner.Subscriptions().FindByRessourceId(obj.String())
	if !ok {
		return subId, false
	}

	resources := sub.ResourcesByObjectID()

	resource, ok := resources[obj.String()]

	if !ok || resource.SubscribedUntil.Before(connector.Clock().Now()) {
		return subId, false
	}

	lastState, ok := resource.LastStates[string(sv.Id())]

	if ok && !lastState.(*stopMonitoringLastChange).Haschanged(sv) {
		return subId, false
	}

	if !ok {
		smlc := &stopMonitoringLastChange{}
		smlc.SetSubscription(sub)
		resource.LastStates[string(sv.Id())] = smlc
	}

	return sub.Id(), true
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) HandleSubscriptionRequest(request *siri.XMLSubscriptionRequest) []siri.SIRIResponseStatus {
	sms := request.XMLSubscriptionSMEntries()

	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	resps := []siri.SIRIResponseStatus{}

	for _, sm := range sms {
		rs := siri.SIRIResponseStatus{
			RequestMessageRef: sm.MessageIdentifier(),
			SubscriberRef:     sm.SubscriberRef(),
			SubscriptionRef:   sm.SubscriptionIdentifier(),
			ResponseTimestamp: connector.Clock().Now(),
		}

		objectid := model.NewObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER), sm.MonitoringRef())
		sa, ok := tx.Model().StopAreas().FindByObjectId(objectid)
		if !ok {
			rs.ErrorType = "InvalidDataReferencesError"
			rs.ErrorText = "Stop Area not found"
			resps = append(resps, rs)
			continue
		}

		sub, ok := connector.Partner().Subscriptions().FindByExternalId(sm.SubscriptionIdentifier())
		if !ok {
			sub = connector.Partner().Subscriptions().New("StopArea")
			sub.SetExternalId(sm.SubscriptionIdentifier())
		}

		ref := model.Reference{
			ObjectId: &objectid,
			Id:       string(sa.Id()),
			Type:     "StopArea",
		}

		r := sub.CreateAddNewResource(ref)
		r.SubscribedUntil = sm.InitialTerminationTime()

		connector.fillOptions(sub, r, request, sm)

		rs.Status = true
		rs.ValidUntil = sm.InitialTerminationTime()
		resps = append(resps, rs)

		connector.AddStopAreaStopVisits(sa, sub, r)
	}
	return resps
}

func (connector *SIRIStopMonitoringSubscriptionBroadcaster) AddStopAreaStopVisits(sa model.StopArea, sub *Subscription, res *SubscribedResource) {
	tx := connector.Partner().Referential().NewTransaction()
	defer tx.Close()

	svs := tx.Model().StopVisits().FindFollowingByStopAreaId(sa.Id())
	for _, sv := range svs {
		_, ok := sv.ObjectID(connector.partner.RemoteObjectIDKind(SIRI_STOP_MONITORING_SUBSCRIPTION_BROADCASTER))
		if !ok {
			continue
		}
		smlc := &stopMonitoringLastChange{}
		smlc.SetSubscription(sub)
		res.LastStates[string(sv.Id())] = smlc
		connector.addStopVisit(sub.Id(), sv.Id())
	}
}

func (smsb *SIRIStopMonitoringSubscriptionBroadcaster) fillOptions(s *Subscription, r *SubscribedResource, request *siri.XMLSubscriptionRequest, sm *siri.XMLStopMonitoringSubscriptionRequestEntry) {
	ro := r.ResourcesOptions()
	ro["StopVisitTypes"] = sm.StopVisitTypes()

	so := s.SubscriptionOptions()

	so["IncrementalUpdates"] = request.IncrementalUpdates()
	so["MaximumStopVisits"] = request.MaximumStopVisits()
	so["ChangeBeforeUpdates"] = request.ChangeBeforeUpdates()
}

// START TEST

type TestSIRIStopMonitoringSubscriptionBroadcasterFactory struct{}

type TestStopMonitoringSubscriptionBroadcaster struct {
	model.UUIDConsumer

	events                    []*model.StopMonitoringBroadcastEvent
	stopMonitoringBroadcaster SIRIStopMonitoringBroadcaster
}

func NewTestStopMonitoringSubscriptionBroadcaster() *TestStopMonitoringSubscriptionBroadcaster {
	connector := &TestStopMonitoringSubscriptionBroadcaster{}
	return connector
}

func (connector *TestStopMonitoringSubscriptionBroadcaster) HandleStopMonitoringBroadcastEvent(event *model.StopMonitoringBroadcastEvent) {
	connector.events = append(connector.events, event)
}

func (factory *TestSIRIStopMonitoringSubscriptionBroadcasterFactory) Validate(apiPartner *APIPartner) bool {
	return true
}

func (factory *TestSIRIStopMonitoringSubscriptionBroadcasterFactory) CreateConnector(partner *Partner) Connector {
	return NewTestStopMonitoringSubscriptionBroadcaster()
}

// END TEST
<|endoftext|>


<commit_before>package koding

import (
	"errors"
	"fmt"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/kites/kloud/api/amazon"
	"koding/kites/kloud/contexthelper/request"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/dnsstorage"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/kloudctl/command"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/pkg/dnsclient"
	"koding/kites/kloud/pkg/multiec2"
	"koding/kites/kloud/plans"
	"koding/kites/kloud/userdata"
	"time"

	"github.com/fatih/structs"
	"github.com/koding/kite"
	"github.com/koding/logging"
	"golang.org/x/net/context"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type Provider struct {
	DB         *mongodb.MongoDB
	Log        logging.Logger
	Kite       *kite.Kite
	DNSClient  *dnsclient.Route53
	DNSStorage *dnsstorage.MongodbStorage
	EC2Clients *multiec2.Clients
	Userdata   *userdata.Userdata

	PaymentFetcher plans.PaymentFetcher
	CheckerFetcher plans.CheckerFetcher
}

type Credential struct {
	Id        bson.ObjectId `bson:"_id" json:"-"`
	PublicKey string        `bson:"publicKey"`
	Meta      bson.M        `bson:"meta"`
}

func (p *Provider) Machine(ctx context.Context, id string) (interface{}, error) {
	if !bson.IsObjectIdHex(id) {
		return nil, fmt.Errorf("Invalid machine id: %q", id)
	}

	// let's first check if the id exists, because we are going to use
	// findAndModify() and it would be difficult to distinguish if the id
	// really doesn't exist or if there is an assignee which is a different
	// thing. (Because findAndModify() also returns "not found" for the case
	// where the id exist but someone else is the assignee).
	machine := &Machine{}
	if err := p.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.FindId(bson.ObjectIdHex(id)).One(&machine)
	}); err == mgo.ErrNotFound {
		return nil, kloud.NewError(kloud.ErrMachineNotFound)
	}

	req, ok := request.FromContext(ctx)
	if !ok {
		return nil, errors.New("request context is not available")
	}

	if machine.Meta.Region == "" {
		machine.Meta.Region = "us-east-1"
		p.Log.Critical("[%s] region is not set in. Fallback to us-east-1.", machine.Id.Hex())
	} else {
		p.Log.Debug("[%s] using region: %s", machine.Id.Hex(), machine.Meta.Region)
	}

	if err := p.attachSession(ctx, machine); err != nil {
		return nil, err
	}

	// check for validation and permission
	if err := p.validate(machine, req); err != nil {
		return nil, err
	}

	return machine, nil
}

func (p *Provider) attachSession(ctx context.Context, machine *Machine) error {
	// get user model which contains user ssh keys or the list of users that
	// are allowed to use this machine
	if len(machine.Users) == 0 {
		return errors.New("permitted users list is empty")
	}

	user, err := p.getOwner(machine.Users)
	if err != nil {
		return err
	}

	client, err := p.EC2Clients.Region(machine.Meta.Region)
	if err != nil {
		return err
	}

	amazonClient, err := amazon.New(structs.Map(machine.Meta), client)
	if err != nil {
		return fmt.Errorf("koding-amazon err: %s", err)
	}

	// attach user specific log
	machine.Log = p.Log.New(machine.Id.Hex())

	sess := &session.Session{
		DB:         p.DB,
		Kite:       p.Kite,
		DNSClient:  p.DNSClient,
		DNSStorage: p.DNSStorage,
		Userdata:   p.Userdata,
		AWSClient:  amazonClient,
		AWSClients: p.EC2Clients, // used for fallback if something goes wrong
		Log:        machine.Log,
	}

	// we use session a lot of in Machine owned methods, so that's why we
	// assign it to a field for easy access
	machine.Session = sess

	// we pass it also to the context, so other packages, such as plans checker
	// can make use of it.
	ctx = session.NewContext(ctx, sess)

	payment, err := p.PaymentFetcher.Fetch(ctx, user.Name)
	if err != nil {
		machine.Log.Warning("username: %s could not fetch plan. Fallback to Free plan. err: '%s'",
			user.Name, err)
		payment = &plans.PaymentResponse{Plan: "Free"}
	}

	checker, err := p.CheckerFetcher.Fetch(ctx, payment.Plan)
	if err != nil {
		return err
	}

	machine.Payment = payment
	machine.Username = user.Name
	machine.User = user
	machine.cleanFuncs = make([]func(), 0)
	machine.Checker = checker

	ev, ok := eventer.FromContext(ctx)
	if ok {
		machine.Session.Eventer = ev
	}

	return nil
}

// getOwner returns the owner of the machine, if it's not found it returns an
// error.
func (p *Provider) getOwner(users []models.Permissions) (*models.User, error) {
	var ownerId bson.ObjectId

	for _, user := range users {
		if user.Sudo && user.Owner {
			ownerId = user.Id
		}
	}

	if !ownerId.Valid() {
		return nil, errors.New("owner not found")
	}

	var user *models.User
	err := p.DB.Run("jUsers", func(c *mgo.Collection) error {
		return c.FindId(users[0].Id).One(&user)
	})

	if err == mgo.ErrNotFound {
		return nil, fmt.Errorf("User with Id not found: %s", ownerId.Hex())
	}
	if err != nil {
		return nil, fmt.Errorf("username lookup error: %v", err)
	}

	return user, nil
}

func (p *Provider) validate(m *Machine, r *kite.Request) error {
	m.Log.Debug("validating for method '%s'", r.Method)

	// give access to kloudctl immediately
	if r.Auth != nil {
		if r.Auth.Key == command.KloudSecretKey {
			return nil
		}
	}

	if r.Username != m.User.Name {
		return errors.New("username is not permitted to make any action")
	}

	// check for user permissions
	if err := p.checkUser(m.User.ObjectId, m.Users); err != nil {
		return err
	}

	if m.User.Status != "confirmed" {
		return kloud.NewError(kloud.ErrUserNotConfirmed)
	}

	return nil
}

// checkUser checks whether the given username is available in the users list
// and has permission
func (p *Provider) checkUser(userId bson.ObjectId, users []models.Permissions) error {
	// check if the incoming user is in the list of permitted user list
	for _, u := range users {
		if userId == u.Id && u.Owner {
			return nil // ok he/she is good to go!
		}
	}

	return fmt.Errorf("permission denied. user not in the list of permitted users")
}

func (p *Provider) credential(publicKey string) (*Credential, error) {
	credential := &Credential{}
	// we neglect errors because credential is optional
	err := p.DB.Run("jCredentialDatas", func(c *mgo.Collection) error {
		return c.Find(bson.M{"publicKey": publicKey}).One(credential)
	})
	if err != nil {
		return nil, err
	}

	return credential, nil
}

func (m *Machine) UpdateState(reason string, state machinestate.State) error {
	m.Log.Debug("Updating state to '%v'", state)
	err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.Update(
			bson.M{
				"_id": m.Id,
			},
			bson.M{
				"$set": bson.M{
					"status.state":      state.String(),
					"status.modifiedAt": time.Now().UTC(),
					"status.reason":     reason,
				},
			},
		)
	})

	if err != nil {
		return fmt.Errorf("Couldn't update state to '%s' for document: '%s' err: %s",
			state, m.Id.Hex(), err)
	}

	return nil
}
<commit_msg>koding/provider: fix supporting multiple allowed users<commit_after>package koding

import (
	"errors"
	"fmt"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/kites/kloud/api/amazon"
	"koding/kites/kloud/contexthelper/request"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/dnsstorage"
	"koding/kites/kloud/eventer"
	"koding/kites/kloud/kloud"
	"koding/kites/kloud/kloudctl/command"
	"koding/kites/kloud/machinestate"
	"koding/kites/kloud/pkg/dnsclient"
	"koding/kites/kloud/pkg/multiec2"
	"koding/kites/kloud/plans"
	"koding/kites/kloud/userdata"
	"time"

	"github.com/fatih/structs"
	"github.com/koding/kite"
	"github.com/koding/logging"
	"golang.org/x/net/context"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type Provider struct {
	DB         *mongodb.MongoDB
	Log        logging.Logger
	Kite       *kite.Kite
	DNSClient  *dnsclient.Route53
	DNSStorage *dnsstorage.MongodbStorage
	EC2Clients *multiec2.Clients
	Userdata   *userdata.Userdata

	PaymentFetcher plans.PaymentFetcher
	CheckerFetcher plans.CheckerFetcher
}

type Credential struct {
	Id        bson.ObjectId `bson:"_id" json:"-"`
	PublicKey string        `bson:"publicKey"`
	Meta      bson.M        `bson:"meta"`
}

func (p *Provider) Machine(ctx context.Context, id string) (interface{}, error) {
	if !bson.IsObjectIdHex(id) {
		return nil, fmt.Errorf("Invalid machine id: %q", id)
	}

	// let's first check if the id exists, because we are going to use
	// findAndModify() and it would be difficult to distinguish if the id
	// really doesn't exist or if there is an assignee which is a different
	// thing. (Because findAndModify() also returns "not found" for the case
	// where the id exist but someone else is the assignee).
	machine := &Machine{}
	if err := p.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.FindId(bson.ObjectIdHex(id)).One(&machine)
	}); err == mgo.ErrNotFound {
		return nil, kloud.NewError(kloud.ErrMachineNotFound)
	}

	req, ok := request.FromContext(ctx)
	if !ok {
		return nil, errors.New("request context is not available")
	}

	if machine.Meta.Region == "" {
		machine.Meta.Region = "us-east-1"
		p.Log.Critical("[%s] region is not set in. Fallback to us-east-1.", machine.Id.Hex())
	} else {
		p.Log.Debug("[%s] using region: %s", machine.Id.Hex(), machine.Meta.Region)
	}

	if err := p.attachSession(ctx, machine); err != nil {
		return nil, err
	}

	// check for validation and permission
	if err := p.validate(machine, req); err != nil {
		return nil, err
	}

	return machine, nil
}

func (p *Provider) attachSession(ctx context.Context, machine *Machine) error {
	// get user model which contains user ssh keys or the list of users that
	// are allowed to use this machine
	if len(machine.Users) == 0 {
		return errors.New("permitted users list is empty")
	}

	// check if this is called via Kite call
	var requesterUsername string
	req, ok := request.FromContext(ctx)
	if ok {
		requesterUsername = req.Username
	}

	// get the user from the permitted list. If the list contains more than one
	// allowed person, fetch the one that is the same as requesterUsername, if
	// not pick up the first one.
	user, err := p.getOwner(requesterUsername, machine.Users)
	if err != nil {
		return err
	}

	client, err := p.EC2Clients.Region(machine.Meta.Region)
	if err != nil {
		return err
	}

	amazonClient, err := amazon.New(structs.Map(machine.Meta), client)
	if err != nil {
		return fmt.Errorf("koding-amazon err: %s", err)
	}

	// attach user specific log
	machine.Log = p.Log.New(machine.Id.Hex())

	sess := &session.Session{
		DB:         p.DB,
		Kite:       p.Kite,
		DNSClient:  p.DNSClient,
		DNSStorage: p.DNSStorage,
		Userdata:   p.Userdata,
		AWSClient:  amazonClient,
		AWSClients: p.EC2Clients, // used for fallback if something goes wrong
		Log:        machine.Log,
	}

	// we use session a lot of in Machine owned methods, so that's why we
	// assign it to a field for easy access
	machine.Session = sess

	// we pass it also to the context, so other packages, such as plans checker
	// can make use of it.
	ctx = session.NewContext(ctx, sess)

	payment, err := p.PaymentFetcher.Fetch(ctx, user.Name)
	if err != nil {
		machine.Log.Warning("username: %s could not fetch plan. Fallback to Free plan. err: '%s'",
			user.Name, err)
		payment = &plans.PaymentResponse{Plan: "Free"}
	}

	checker, err := p.CheckerFetcher.Fetch(ctx, payment.Plan)
	if err != nil {
		return err
	}

	machine.Payment = payment
	machine.Username = user.Name
	machine.User = user
	machine.cleanFuncs = make([]func(), 0)
	machine.Checker = checker

	ev, ok := eventer.FromContext(ctx)
	if ok {
		machine.Session.Eventer = ev
	}

	return nil
}

// getOwner returns the owner of the machine, if it's not found it returns an
// error. The requestName is optional, if it's not empty and the the users list
// has more than one valid allowed users, we return the one that matches the
// requesterName.
func (p *Provider) getOwner(requesterName string, users []models.Permissions) (*models.User, error) {
	// if the list contains only one user, do a short lookup.
	if len(users) == 1 || requesterName == "" {
		var ownerId bson.ObjectId

		for _, user := range users {
			if user.Sudo && user.Owner {
				ownerId = user.Id
			}
		}

		if !ownerId.Valid() {
			return nil, errors.New("owner not found")
		}

		var user *models.User
		err := p.DB.Run("jUsers", func(c *mgo.Collection) error {
			return c.FindId(users[0].Id).One(&user)
		})

		if err == mgo.ErrNotFound {
			return nil, fmt.Errorf("User with Id not found: %s", ownerId.Hex())
		}
		if err != nil {
			return nil, fmt.Errorf("username lookup error: %v", err)
		}

		return user, nil
	}

	// get the full list of users and return the one that matches the
	// requesterName, if not we return someone that is allowed. Note that we
	// don't do the validation here, this is only to fetch the user, don't put
	// any validation logic here.
	var allUsers []*models.User
	userIds := make([]bson.ObjectId, len(users))
	for i, perm := range users {
		// we only going to fetch users that are allowed
		if perm.Sudo && perm.Owner {
			userIds[i] = perm.Id
		}
	}

	// nothing found, just return
	if len(userIds) == 0 {
		return nil, errors.New("owner not found")
	}

	if err := p.DB.Run("jUsers", func(c *mgo.Collection) error {
		return c.Find(bson.M{"_id": bson.M{"$in": userIds}}).All(&users)
	}); err != nil {
		return nil, fmt.Errorf("username lookup error: %v", err)
	}

	// now we have all allowed users, if we have someone that is in match with
	// the requesterName just return it.
	for _, u := range allUsers {
		if u.Name == requesterName {
			return u, nil
		}
	}

	// nothing found, just return the first one
	return allUsers[0], nil
}

func (p *Provider) validate(m *Machine, r *kite.Request) error {
	m.Log.Debug("validating for method '%s'", r.Method)

	// give access to kloudctl immediately
	if r.Auth != nil {
		if r.Auth.Key == command.KloudSecretKey {
			return nil
		}
	}

	if r.Username != m.User.Name {
		return errors.New("username is not permitted to make any action")
	}

	// check for user permissions
	if err := p.checkUser(m.User.ObjectId, m.Users); err != nil {
		return err
	}

	if m.User.Status != "confirmed" {
		return kloud.NewError(kloud.ErrUserNotConfirmed)
	}

	return nil
}

// checkUser checks whether the given username is available in the users list
// and has permission
func (p *Provider) checkUser(userId bson.ObjectId, users []models.Permissions) error {
	// check if the incoming user is in the list of permitted user list
	for _, u := range users {
		if userId == u.Id && u.Owner {
			return nil // ok he/she is good to go!
		}
	}

	return fmt.Errorf("permission denied. user not in the list of permitted users")
}

func (p *Provider) credential(publicKey string) (*Credential, error) {
	credential := &Credential{}
	// we neglect errors because credential is optional
	err := p.DB.Run("jCredentialDatas", func(c *mgo.Collection) error {
		return c.Find(bson.M{"publicKey": publicKey}).One(credential)
	})
	if err != nil {
		return nil, err
	}

	return credential, nil
}

func (m *Machine) UpdateState(reason string, state machinestate.State) error {
	m.Log.Debug("Updating state to '%v'", state)
	err := m.Session.DB.Run("jMachines", func(c *mgo.Collection) error {
		return c.Update(
			bson.M{
				"_id": m.Id,
			},
			bson.M{
				"$set": bson.M{
					"status.state":      state.String(),
					"status.modifiedAt": time.Now().UTC(),
					"status.reason":     reason,
				},
			},
		)
	})

	if err != nil {
		return fmt.Errorf("Couldn't update state to '%s' for document: '%s' err: %s",
			state, m.Id.Hex(), err)
	}

	return nil
}
<|endoftext|>


<commit_before>package models

import (
	"fmt"
	"time"

	"github.com/koding/bongo"
)

// TO-DO
//
// Are you sure to add this notification settings into the api schema
// If not , think about more for it about pros and cons
//
//~Mehmet Ali
const NotificationSettingsBongoName = "notification.notification_settings"

func (n NotificationSettings) GetId() int64 {
	return n.Id
}

func (n NotificationSettings) BongoName() string {
	return NotificationSettingsBongoName
}

func (n *NotificationSettings) Create() error {
	if n.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	if n.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	return bongo.B.Create(n)
}

func (n *NotificationSettings) Update() error {
	if n.ChannelId == 0 || n.AccountId == 0 {
		return fmt.Errorf("Update failed ChannelId: %s - AccountId:%s", n.ChannelId, n.AccountId)
	}

	return bongo.B.Update(n)
}

func (n *NotificationSettings) AfterCreate() {
	bongo.B.AfterCreate(n)
}

func (n *NotificationSettings) AfterUpdate() {
	bongo.B.AfterUpdate(n)
}

func (n *NotificationSettings) BeforeCreate() error {
	if err := n.validateBeforeOps(); err != nil {
		return err
	}
	now := time.Now().UTC()
	n.CreatedAt = now
	n.UpdatedAt = now

	return nil
}

// BeforeUpdate runs before updating struct
func (n *NotificationSettings) BeforeUpdate() error {
	return n.validateBeforeOps()
}

func (n *NotificationSettings) One(q *bongo.Query) error {
	return bongo.B.One(n, n, q)
}

func (n *NotificationSettings) ById(id int64) error {
	return bongo.B.ById(n, id)
}

func (n *NotificationSettings) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(n, data, q)
}

func (n *NotificationSettings) validateBeforeOps() error {
	if n.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	if n.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	// a := NewAccount()
	_, err := Cache.Account.ById(n.AccountId)
	if err != nil {
		return err
	}

	_, err = Cache.Channel.ById(n.ChannelId)
	if err != nil {
		return err
	}

	// We should add group && group control ??

	return nil
}
<commit_msg>notificationSettings: delete is added into bongo file<commit_after>package models

import (
	"fmt"
	"time"

	"github.com/koding/bongo"
)

// TO-DO
//
// Are you sure to add this notification settings into the api schema
// If not , think about more for it about pros and cons
//
//~Mehmet Ali
const NotificationSettingsBongoName = "notification.notification_settings"

func (n NotificationSettings) GetId() int64 {
	return n.Id
}

func (n NotificationSettings) BongoName() string {
	return NotificationSettingsBongoName
}

func (n *NotificationSettings) Create() error {
	if n.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	if n.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	return bongo.B.Create(n)
}

func (n *NotificationSettings) Update() error {
	if n.ChannelId == 0 || n.AccountId == 0 {
		return fmt.Errorf("Update failed ChannelId: %s - AccountId:%s", n.ChannelId, n.AccountId)
	}

	return bongo.B.Update(n)
}

func (n *NotificationSettings) Delete() error {
	selector := map[string]interface{}{
		"channel_id": n.ChannelId,
		"account_id": n.AccountId,
	}

	if err := n.One(bongo.NewQS(selector)); err != nil {
		return err
	}

	return bongo.B.Delete(n)
}

func (n *NotificationSettings) AfterCreate() {
	bongo.B.AfterCreate(n)
}

func (n *NotificationSettings) AfterUpdate() {
	bongo.B.AfterUpdate(n)
}

func (n *NotificationSettings) BeforeCreate() error {
	if err := n.validateBeforeOps(); err != nil {
		return err
	}
	now := time.Now().UTC()
	n.CreatedAt = now
	n.UpdatedAt = now

	return nil
}

// BeforeUpdate runs before updating struct
func (n *NotificationSettings) BeforeUpdate() error {
	return n.validateBeforeOps()
}

func (n *NotificationSettings) One(q *bongo.Query) error {
	return bongo.B.One(n, n, q)
}

func (n *NotificationSettings) ById(id int64) error {
	return bongo.B.ById(n, id)
}

func (n *NotificationSettings) Some(data interface{}, q *bongo.Query) error {
	return bongo.B.Some(n, data, q)
}

func (n *NotificationSettings) validateBeforeOps() error {
	if n.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	if n.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	// a := NewAccount()
	_, err := Cache.Account.ById(n.AccountId)
	if err != nil {
		return err
	}

	_, err = Cache.Channel.ById(n.ChannelId)
	if err != nil {
		return err
	}

	// We should add group && group control ??

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package tabletmanager

import (
	"flag"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/trace"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools/events"
	"vitess.io/vitess/go/vt/vttablet/tmclient"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

var (
	finalizeReparentTimeout = flag.Duration("finalize_external_reparent_timeout", 30*time.Second, "Timeout for the finalize stage of a fast external reparent reconciliation.")

	externalReparentStats = stats.NewTimings("ExternalReparents", "External reparenting time", "stage", "NewMasterVisible", "FullRebuild")
)

// SetReparentFlags changes flag values. It should only be used in tests.
func SetReparentFlags(timeout time.Duration) {
	*finalizeReparentTimeout = timeout
}

// TabletExternallyReparented updates all topo records so the current
// tablet is the new master for this shard.
func (agent *ActionAgent) TabletExternallyReparented(ctx context.Context, externalID string) error {
	if err := agent.lock(ctx); err != nil {
		return err
	}
	defer agent.unlock()

	startTime := time.Now()

	// If there is a finalize step running, wait for it to finish or time out
	// before checking the global shard record again.
	if agent.finalizeReparentCtx != nil {
		select {
		case <-agent.finalizeReparentCtx.Done():
			agent.finalizeReparentCtx = nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}

	tablet := agent.Tablet()

	// Check the global shard record.
	si, err := agent.TopoServer.GetShard(ctx, tablet.Keyspace, tablet.Shard)
	if err != nil {
		log.Warningf("fastTabletExternallyReparented: failed to read global shard record for %v/%v: %v", tablet.Keyspace, tablet.Shard, err)
		return err
	}

	// The external failover tool told us that we are still the MASTER. Update the
	// timestamp to the current time.
	agent.setExternallyReparentedTime(startTime)

	// Create a reusable Reparent event with available info.
	ev := &events.Reparent{
		ShardInfo: *si,
		NewMaster: *tablet,
		OldMaster: topodatapb.Tablet{
			Alias: si.MasterAlias,
			Type:  topodatapb.TabletType_MASTER,
		},
		ExternalID: externalID,
	}
	defer func() {
		if err != nil {
			event.DispatchUpdate(ev, "failed: "+err.Error())
		}
	}()
	event.DispatchUpdate(ev, "starting external from tablet (fast)")

	// Execute state change to master by force-updating only the local copy of the
	// tablet record. The actual record in topo will be updated later.
	newTablet := proto.Clone(tablet).(*topodatapb.Tablet)

	// We may get called on the current master multiple times in order to fix incomplete external reparents
	// update tablet only if it is not currently master
	if newTablet.Type != topodatapb.TabletType_MASTER {
		log.Infof("fastTabletExternallyReparented: executing change callback for state change to MASTER")
		newTablet.Type = topodatapb.TabletType_MASTER

		// This is where updateState will block for gracePeriod, while it gives
		// vtgate a chance to stop sending replica queries.
		agent.updateState(ctx, newTablet, "fastTabletExternallyReparented")
	}
	// Start the finalize stage with a background context, but connect the trace.
	bgCtx, cancel := context.WithTimeout(agent.batchCtx, *finalizeReparentTimeout)
	bgCtx = trace.CopySpan(bgCtx, ctx)
	agent.finalizeReparentCtx = bgCtx
	go func() {
		err := agent.finalizeTabletExternallyReparented(bgCtx, si, ev)
		cancel()

		if err != nil {
			log.Warningf("finalizeTabletExternallyReparented error: %v", err)
			event.DispatchUpdate(ev, "failed: "+err.Error())
			return
		}
		externalReparentStats.Record("FullRebuild", startTime)
	}()

	return nil
}

// finalizeTabletExternallyReparented performs slow, synchronized reconciliation
// tasks that ensure topology is self-consistent
// it first updates new and old master tablet records, then updates
// the global shard record, then refreshes the old master
// after that it attempts to detect and clean up any lingering old masters
// note that an up-to-date shard record does not necessarily mean that
// the reparent completed all the actions successfully
func (agent *ActionAgent) finalizeTabletExternallyReparented(ctx context.Context, si *topo.ShardInfo, ev *events.Reparent) (err error) {
	var wg sync.WaitGroup
	var errs concurrency.AllErrorRecorder
	var oldMasterTablet *topodatapb.Tablet
	oldMasterAlias := si.MasterAlias

	// Update the new and old master tablet records concurrently.
	event.DispatchUpdate(ev, "updating old and new master tablet records")
	log.Infof("finalizeTabletExternallyReparented: updating tablet records")
	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Infof("finalizeTabletExternallyReparented: updating tablet record for new master: %v", agent.TabletAlias)
		// Update our own record to master if needed
		_, err := agent.TopoServer.UpdateTabletFields(ctx, agent.TabletAlias,
			func(tablet *topodatapb.Tablet) error {
				if tablet.Type != topodatapb.TabletType_MASTER {
					tablet.Type = topodatapb.TabletType_MASTER
				}
				return nil
			})
		if err != nil {
			errs.RecordError(err)
		}
	}()

	if !topoproto.TabletAliasIsZero(oldMasterAlias) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			log.Infof("finalizeTabletExternallyReparented: updating tablet record for old master: %v", oldMasterAlias)

			// Forcibly demote the old master in topology, since we can't rely on the
			// old master to be up to change its own record.
			var err error
			oldMasterTablet, err = agent.TopoServer.UpdateTabletFields(ctx, oldMasterAlias,
				func(tablet *topodatapb.Tablet) error {
					if tablet.Type == topodatapb.TabletType_MASTER {
						tablet.Type = topodatapb.TabletType_REPLICA
					}
					return nil
				})
			if err != nil {
				errs.RecordError(err)
				return
			}

			// We now know more about the old master, so add it to event data.
			// oldMasterTablet will be nil if no update was needed
			if oldMasterTablet != nil {
				ev.OldMaster = *oldMasterTablet
			}
		}()
	}
	// Wait for the tablet records to be updated. At that point, any rebuild will
	// see the new master, so we're ready to mark the reparent as done in the
	// global shard record.
	wg.Wait()
	if errs.HasErrors() {
		return errs.Error()
	}

	masterTablet := agent.Tablet()

	// Update the master field in the global shard record. We don't use a lock
	// here anymore. The lock was only to ensure that the global shard record
	// didn't get modified between the time when we read it and the time when we
	// write it back. Now we use an update loop pattern to do that instead.

	// We also used to do this in parallel with RefreshState on the old master
	// we don't do that any more. we want to update the shard record first
	// and only then attempt to refresh the old master because it is possible
	// that the old master is unreachable
	event.DispatchUpdate(ev, "updating global shard record")
	log.Infof("finalizeTabletExternallyReparented: updating global shard record if needed")
	_, err = agent.TopoServer.UpdateShardFields(ctx, masterTablet.Keyspace, masterTablet.Shard, func(currentSi *topo.ShardInfo) error {
		if topoproto.TabletAliasEqual(currentSi.MasterAlias, masterTablet.Alias) {
			// It is correct to return this error here, UpdateShardFields will ignore it
			return topo.NewError(topo.NoUpdateNeeded, masterTablet.Alias.String())
		}
		if !topoproto.TabletAliasEqual(currentSi.MasterAlias, oldMasterAlias) {
			log.Warningf("old master alias (%v) not found in the global Shard record i.e. it has changed in the meantime."+
				" We're not overwriting the value with the new master (%v) because the current value is probably newer."+
				" (initial Shard record = %#v, current Shard record = %#v)",
				oldMasterAlias, masterTablet.Alias, si, currentSi)
			// It is correct to return this error here, UpdateShardFields will ignore it
			return topo.NewError(topo.NoUpdateNeeded, oldMasterAlias.String())
		}
		currentSi.MasterAlias = masterTablet.Alias
		return nil
	})
	if err != nil {
		errs.RecordError(err)
	}

	if errs.HasErrors() {
		return errs.Error()
	}

	if !topoproto.TabletAliasIsZero(oldMasterAlias) {
		// Tell the old master to re-read its tablet record and change its state.
		// We don't need to put error into errs if this fails, but we need to wait
		// for it to make sure that old master tablet is not stuck in the MASTER
		// state.
		tmc := tmclient.NewTabletManagerClient()
		if err := tmc.RefreshState(ctx, oldMasterTablet); err != nil {
			log.Warningf("Error calling RefreshState on old master %v: %v", topoproto.TabletAliasString(oldMasterTablet.Alias), err)
		}
	}

	tabletMap, err := agent.TopoServer.GetTabletMapForShard(ctx, masterTablet.Keyspace, masterTablet.Shard)
	if err != nil {
		log.Errorf("ignoring error %v from GetTabletMapForShard so that we can process any partial results", err)
	}

	if len(tabletMap) > 0 {
		// make the channel buffer big enough that it doesn't block senders
		tabletsToRefresh := make(chan topodatapb.Tablet, len(tabletMap))
		// update any other tablets claiming to be MASTER also to REPLICA
		for _, tabletInfo := range tabletMap {
			alias := tabletInfo.Tablet.Alias
			if !topoproto.TabletAliasEqual(alias, agent.TabletAlias) && !topoproto.TabletAliasEqual(alias, oldMasterAlias) && tabletInfo.Tablet.Type == topodatapb.TabletType_MASTER {
				log.Infof("finalizeTabletExternallyReparented: updating tablet record for another old master: %v", alias)
				wg.Add(1)
				go func(alias *topodatapb.TabletAlias) {
					defer wg.Done()
					var err error
					tab, err := agent.TopoServer.UpdateTabletFields(ctx, alias,
						func(tablet *topodatapb.Tablet) error {
							if tablet.Type == topodatapb.TabletType_MASTER {
								tablet.Type = topodatapb.TabletType_REPLICA
							}
							return nil
						})
					if err != nil {
						errs.RecordError(err)
						return
					}
					// tab will be nil if no update was needed
					if tab != nil {
						tabletsToRefresh <- *tab
					}
				}(alias)
			}
		}
		// Wait for the tablet records to be updated. At that point, any rebuild will
		// see the new master, so we're ready to mark the reparent as done in the
		// global shard record.
		wg.Wait()
		// we waited for all goroutines to complete, so now close the channel
		close(tabletsToRefresh)
		if errs.HasErrors() {
			return errs.Error()
		}

		for tab := range tabletsToRefresh {
			log.Infof("finalizeTabletExternallyReparented: Refresh state for tablet: %v", topoproto.TabletAliasString(tab.Alias))
			wg.Add(1)
			go func(tablet topodatapb.Tablet) {
				defer wg.Done()

				// Tell the old master(s) to re-read its tablet record and change its state.
				// We don't need to put error into errs if this fails, but we need to wait
				// for it to make sure that an old master tablet is not stuck in the MASTER
				// state.
				tmc := tmclient.NewTabletManagerClient()
				if err := tmc.RefreshState(ctx, &tablet); err != nil {
					log.Warningf("Error calling RefreshState on old master %v: %v", topoproto.TabletAliasString(tablet.Alias), err)
				}
			}(tab)
		}
		wg.Wait()
	}
	if errs.HasErrors() {
		return errs.Error()
	}

	event.DispatchUpdate(ev, "finished")
	return nil
}

// setExternallyReparentedTime remembers the last time when we were told we're
// the master.
// If another tablet claims to be master and offers a more recent time,
// that tablet will be trusted over us.
func (agent *ActionAgent) setExternallyReparentedTime(t time.Time) {
	agent.mutex.Lock()
	defer agent.mutex.Unlock()

	agent._tabletExternallyReparentedTime = t
	agent._replicationDelay = 0
}
<commit_msg>fix/add descriptive comments, avoid unnecessary topo calls, merge 2 loops into 1 and remove use of channel that is no longer needed<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package tabletmanager

import (
	"flag"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
	"vitess.io/vitess/go/event"
	"vitess.io/vitess/go/stats"
	"vitess.io/vitess/go/trace"
	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools/events"
	"vitess.io/vitess/go/vt/vttablet/tmclient"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

var (
	finalizeReparentTimeout = flag.Duration("finalize_external_reparent_timeout", 30*time.Second, "Timeout for the finalize stage of a fast external reparent reconciliation.")

	externalReparentStats = stats.NewTimings("ExternalReparents", "External reparenting time", "stage", "NewMasterVisible", "FullRebuild")
)

// SetReparentFlags changes flag values. It should only be used in tests.
func SetReparentFlags(timeout time.Duration) {
	*finalizeReparentTimeout = timeout
}

// TabletExternallyReparented updates all topo records so the current
// tablet is the new master for this shard.
func (agent *ActionAgent) TabletExternallyReparented(ctx context.Context, externalID string) error {
	if err := agent.lock(ctx); err != nil {
		return err
	}
	defer agent.unlock()

	startTime := time.Now()

	// If there is a finalize step running, wait for it to finish or time out
	// before checking the global shard record again.
	if agent.finalizeReparentCtx != nil {
		select {
		case <-agent.finalizeReparentCtx.Done():
			agent.finalizeReparentCtx = nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}

	tablet := agent.Tablet()

	// Check the global shard record.
	si, err := agent.TopoServer.GetShard(ctx, tablet.Keyspace, tablet.Shard)
	if err != nil {
		log.Warningf("fastTabletExternallyReparented: failed to read global shard record for %v/%v: %v", tablet.Keyspace, tablet.Shard, err)
		return err
	}

	// The external failover tool told us that we are still the MASTER. Update the
	// timestamp to the current time.
	agent.setExternallyReparentedTime(startTime)

	// Create a reusable Reparent event with available info.
	ev := &events.Reparent{
		ShardInfo: *si,
		NewMaster: *tablet,
		OldMaster: topodatapb.Tablet{
			Alias: si.MasterAlias,
			Type:  topodatapb.TabletType_MASTER,
		},
		ExternalID: externalID,
	}
	event.DispatchUpdate(ev, "starting external from tablet (fast)")

	// Execute state change to master by force-updating only the local copy of the
	// tablet record. The actual record in topo will be updated later.
	newTablet := proto.Clone(tablet).(*topodatapb.Tablet)

	// We may get called on the current master multiple times in order to fix incomplete external reparents.
	// We update the tablet here only if it is not currently master
	if newTablet.Type != topodatapb.TabletType_MASTER {
		log.Infof("fastTabletExternallyReparented: executing change callback for state change to MASTER")
		newTablet.Type = topodatapb.TabletType_MASTER

		// This is where updateState will block for gracePeriod, while it gives
		// vtgate a chance to stop sending replica queries.
		agent.updateState(ctx, newTablet, "fastTabletExternallyReparented")
	}
	// Start the finalize stage with a background context, but connect the trace.
	bgCtx, cancel := context.WithTimeout(agent.batchCtx, *finalizeReparentTimeout)
	bgCtx = trace.CopySpan(bgCtx, ctx)
	agent.finalizeReparentCtx = bgCtx
	go func() {
		err := agent.finalizeTabletExternallyReparented(bgCtx, si, ev)
		cancel()

		if err != nil {
			log.Warningf("finalizeTabletExternallyReparented error: %v", err)
			event.DispatchUpdate(ev, "failed: "+err.Error())
			return
		}
		externalReparentStats.Record("FullRebuild", startTime)
	}()

	return nil
}

// finalizeTabletExternallyReparented performs slow, synchronized reconciliation
// tasks that ensure topology is self-consistent.
// It first updates new and old master tablet records, then updates
// the global shard record, then refreshes the old master.
// After that it attempts to detect and clean up any lingering old masters.
// Note that an up-to-date shard record does not necessarily mean that
// the reparent completed all the actions successfully
func (agent *ActionAgent) finalizeTabletExternallyReparented(ctx context.Context, si *topo.ShardInfo, ev *events.Reparent) (err error) {
	var wg sync.WaitGroup
	var errs concurrency.AllErrorRecorder
	var oldMasterTablet *topodatapb.Tablet
	oldMasterAlias := si.MasterAlias

	// Update the new and old master tablet records concurrently.
	event.DispatchUpdate(ev, "updating old and new master tablet records")
	log.Infof("finalizeTabletExternallyReparented: updating tablet records")
	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Infof("finalizeTabletExternallyReparented: updating tablet record for new master: %v", agent.TabletAlias)
		// Update our own record to master if needed
		_, err := agent.TopoServer.UpdateTabletFields(ctx, agent.TabletAlias,
			func(tablet *topodatapb.Tablet) error {
				if tablet.Type != topodatapb.TabletType_MASTER {
					tablet.Type = topodatapb.TabletType_MASTER
					return nil
				}
				// returning NoUpdateNeeded avoids unnecessary calls to UpdateTablet
				return topo.NewError(topo.NoUpdateNeeded, agent.TabletAlias.String())
			})
		if err != nil {
			errs.RecordError(err)
		}
	}()

	if !topoproto.TabletAliasIsZero(oldMasterAlias) {
		wg.Add(1)
		go func() {
			defer wg.Done()
			log.Infof("finalizeTabletExternallyReparented: updating tablet record for old master: %v", oldMasterAlias)

			// Forcibly demote the old master in topology, since we can't rely on the
			// old master to be up to change its own record.
			var err error
			oldMasterTablet, err = agent.TopoServer.UpdateTabletFields(ctx, oldMasterAlias,
				func(tablet *topodatapb.Tablet) error {
					if tablet.Type == topodatapb.TabletType_MASTER {
						tablet.Type = topodatapb.TabletType_REPLICA
						return nil
					}
					// returning NoUpdateNeeded avoids unnecessary calls to UpdateTablet
					return topo.NewError(topo.NoUpdateNeeded, oldMasterAlias.String())
				})
			if err != nil {
				errs.RecordError(err)
				return
			}

			// We now know more about the old master, so add it to event data.
			// oldMasterTablet will be nil if no update was needed
			if oldMasterTablet != nil {
				ev.OldMaster = *oldMasterTablet
			}
		}()
	}
	// Wait for the tablet records to be updated. At that point, any rebuild will
	// see the new master, so we're ready to mark the reparent as done in the
	// global shard record.
	wg.Wait()
	if errs.HasErrors() {
		return errs.Error()
	}

	masterTablet := agent.Tablet()

	event.DispatchUpdate(ev, "updating global shard record")
	log.Infof("finalizeTabletExternallyReparented: updating global shard record if needed")
	wg.Add(1)
	go func() {
		defer wg.Done()
		// Update the master field in the global shard record. We don't use a lock
		// here anymore. The lock was only to ensure that the global shard record
		// didn't get modified between the time when we read it and the time when we
		// write it back. Now we use an update loop pattern to do that instead.
		_, err = agent.TopoServer.UpdateShardFields(ctx, masterTablet.Keyspace, masterTablet.Shard, func(currentSi *topo.ShardInfo) error {
			if topoproto.TabletAliasEqual(currentSi.MasterAlias, masterTablet.Alias) {
				// returning NoUpdateNeeded avoids unnecessary calls to UpdateTablet
				return topo.NewError(topo.NoUpdateNeeded, masterTablet.Alias.String())
			}
			if !topoproto.TabletAliasEqual(currentSi.MasterAlias, oldMasterAlias) {
				log.Warningf("old master alias (%v) not found in the global Shard record i.e. it has changed in the meantime."+
					" We're not overwriting the value with the new master (%v) because the current value is probably newer."+
					" (initial Shard record = %#v, current Shard record = %#v)",
					oldMasterAlias, masterTablet.Alias, si, currentSi)
				// returning NoUpdateNeeded avoids unnecessary calls to UpdateTablet
				return topo.NewError(topo.NoUpdateNeeded, oldMasterAlias.String())
			}
			currentSi.MasterAlias = masterTablet.Alias
			return nil
		})
		if err != nil {
			errs.RecordError(err)
		}
	}()

	if !topoproto.TabletAliasIsZero(oldMasterAlias) && oldMasterTablet != nil {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// Tell the old master to re-read its tablet record and change its state.
			// We don't need to put error into errs if this fails, but we need to wait
			// for it to make sure that old master tablet is not stuck in the MASTER
			// state.
			tmc := tmclient.NewTabletManagerClient()
			if err := tmc.RefreshState(ctx, oldMasterTablet); err != nil {
				log.Warningf("Error calling RefreshState on old master %v: %v", topoproto.TabletAliasString(oldMasterTablet.Alias), err)
			}
		}()
	}

	wg.Wait()
	if errs.HasErrors() {
		return errs.Error()
	}

	// Look for any other tablets claiming to be master and fix them up on a best-effort basis
	tabletMap, err := agent.TopoServer.GetTabletMapForShard(ctx, masterTablet.Keyspace, masterTablet.Shard)
	if err != nil {
		log.Errorf("ignoring error %v from GetTabletMapForShard so that we can process any partial results", err)
	}

	for _, tabletInfo := range tabletMap {
		alias := tabletInfo.Tablet.Alias
		if !topoproto.TabletAliasEqual(alias, agent.TabletAlias) && !topoproto.TabletAliasEqual(alias, oldMasterAlias) && tabletInfo.Tablet.Type == topodatapb.TabletType_MASTER {
			log.Infof("finalizeTabletExternallyReparented: updating tablet record for another old master: %v", alias)
			wg.Add(1)
			go func(alias *topodatapb.TabletAlias) {
				defer wg.Done()
				var err error
				tab, err := agent.TopoServer.UpdateTabletFields(ctx, alias,
					func(tablet *topodatapb.Tablet) error {
						if tablet.Type == topodatapb.TabletType_MASTER {
							tablet.Type = topodatapb.TabletType_REPLICA
							return nil
						}
						return topo.NewError(topo.NoUpdateNeeded, alias.String())
					})
				if err != nil {
					errs.RecordError(err)
					return
				}
				// tab will be nil if no update was needed
				if tab != nil {
					log.Infof("finalizeTabletExternallyReparented: Refresh state for tablet: %v", topoproto.TabletAliasString(tab.Alias))
					tmc := tmclient.NewTabletManagerClient()
					if err := tmc.RefreshState(ctx, tab); err != nil {
						log.Warningf("Error calling RefreshState on old master %v: %v", topoproto.TabletAliasString(tab.Alias), err)
					}
				}
			}(alias)
		}
	}
	wg.Wait()
	if errs.HasErrors() {
		return errs.Error()
	}

	event.DispatchUpdate(ev, "finished")
	return nil
}

// setExternallyReparentedTime remembers the last time when we were told we're
// the master.
// If another tablet claims to be master and offers a more recent time,
// that tablet will be trusted over us.
func (agent *ActionAgent) setExternallyReparentedTime(t time.Time) {
	agent.mutex.Lock()
	defer agent.mutex.Unlock()

	agent._tabletExternallyReparentedTime = t
	agent._replicationDelay = 0
}
<|endoftext|>


<commit_before>// Copyright 2015 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package orm

import (
	"bytes"
	"database/sql"
	"errors"
	"strconv"
	"strings"

	"github.com/issue9/orm/fetch"
)

const (
	and = iota
	or
)

const (
	asc = iota
	desc
)

var (
	ErrEmptyTableName = errors.New("未指定表名")
)

// 以函数链的方式产生SQL语句。
type SQL struct {
	e        engine
	table    string
	cond     *bytes.Buffer
	condArgs []interface{}
	order    *bytes.Buffer
	limit    []int // 元素0为limit，元素1为offset
}

func newSQL(engine engine) *SQL {
	return &SQL{
		e:        engine,
		cond:     new(bytes.Buffer),
		condArgs: []interface{}{},
		order:    new(bytes.Buffer),
		limit:    make([]int, 0, 2), // 最长2个，同时保存limit和offset
	}
}

func (s *SQL) Reset() *SQL {
	s.table = ""
	s.cond.Reset()
	s.condArgs = s.condArgs[:0]
	s.order.Reset()
	s.limit = s.limit[:0]

	return s
}

// 指定Limit相关的值。
func (s *SQL) Limit(limit int, offset ...int) *SQL {
	if len(offset) > 1 {
		panic("offset参数指定了太多的值")
	}

	s.limit = append(s.limit, limit)

	if len(offset) > 0 {
		s.limit = append(s.limit, offset[0])
	}

	return s
}

func (s *SQL) where(op int, cond string, args ...interface{}) *SQL {
	switch {
	case s.cond.Len() == 0:
		s.cond.WriteString(" WHERE(")
	case op == and:
		s.cond.WriteString(" AND(")
	case op == or:
		s.cond.WriteString(" OR(")
	default:
		panic("where:错误的参数op")
	}
	s.cond.WriteString(cond)
	s.cond.WriteByte(')')
	s.condArgs = append(s.condArgs, args...)

	return s
}

// 将之后的语句以and的形式与当前的语句进行连接
func (s *SQL) And(cond string, args ...interface{}) *SQL {
	return s.where(and, cond, args...)
}

// 将之后的语句以or的形式与当前的语句进行连接
func (s *SQL) Or(cond string, args ...interface{}) *SQL {
	return s.where(or, cond, args...)
}

// order by ... asc
// 当cols参数为空时，不产生任何操作。
func (s *SQL) Asc(cols ...string) *SQL {
	if len(cols) == 0 {
		return s
	}

	for _, col := range cols {
		s.order.WriteString(col)
		s.order.WriteByte(',')
	}
	s.order.Truncate(s.order.Len() - 1)
	s.order.WriteString(" ASC,")
	return s
}

// order by ... desc
// 当cols参数为空时，不产生任何操作。
func (s *SQL) Desc(cols ...string) *SQL {
	if len(cols) == 0 {
		return s
	}

	for _, col := range cols {
		s.order.WriteString(col)
		s.order.WriteByte(',')
	}
	s.order.Truncate(s.order.Len() - 1)
	s.order.WriteString(" DESC,")
	return s
}

// 指定表名。
func (s *SQL) Table(tableName string) *SQL {
	s.table = tableName
	return s
}

// 将符合当前条件的所有记录删除。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Delete(replace bool) (sql.Result, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	sql.Reset()
	sql.WriteString("DELETE FROM ")
	s.e.Dialect().Quote(sql, s.table)
	sql.WriteString(s.cond.String())
	return s.e.Exec(replace, sql.String(), s.condArgs...)
}

// 更新符合当前条件的所有记录。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Update(replace bool, data map[string]interface{}) (sql.Result, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	if len(data) == 0 {
		return nil, errors.New("SQL.Update:未指定需要更新的数据")
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	sql.Reset()
	sql.WriteString("UPDATE ")
	s.e.Dialect().Quote(sql, s.table)
	vals := make([]interface{}, 0, len(data)+len(s.condArgs))
	sql.WriteString(" SET ")
	for k, v := range data {
		s.e.Dialect().Quote(sql, k)
		sql.WriteString("=?,")
		vals = append(vals, v)
	}
	sql.Truncate(sql.Len() - 1) // 去掉最后一个逗号

	sql.WriteString(s.cond.String())
	return s.e.Exec(replace, sql.String(), append(vals, s.condArgs...)...)
}

// 返回符合条件的记录数量。
// 若指定了Limit，则相应的条件也会计算在内。
func (s *SQL) Count(replace bool) (int, error) {
	rows, err := s.query(replace, "COUNT(*) AS cnt")
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	cols, err := fetch.ColumnString(true, "cnt", rows)
	if err != nil {
		return 0, err
	}

	if len(cols) == 0 {
		return 0, nil
	}

	return strconv.Atoi(cols[0])
}

// 将符合当前条件的所有记录依次写入objs中。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Select(replace bool, objs interface{}) (int, error) {
	rows, err := s.query(replace)
	if err != nil {
		return 0, err
	}
	cnt, err := fetch.Obj(objs, rows)
	rows.Close()
	return cnt, err
}

// 返回符合当前条件的所有记录。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) SelectMap(replace bool, cols ...string) ([]map[string]interface{}, error) {
	rows, err := s.query(replace, cols...)
	if err != nil {
		return nil, err
	}
	mapped, err := fetch.Map(false, rows)
	rows.Close()
	return mapped, err
}

func (s *SQL) query(replace bool, cols ...string) (*sql.Rows, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	args := make([]interface{}, 0, len(s.condArgs)) // Query对应的参数
	sql.Reset()
	sql.WriteString("SELECT ")
	if len(cols) == 0 {
		sql.WriteString(" * ")
	} else {
		sql.WriteString(strings.Join(cols, ","))
	}

	// from
	sql.WriteString(" FROM ")
	s.e.Dialect().Quote(sql, s.table)

	// where
	sql.WriteString(s.cond.String())
	args = append(args, s.condArgs...)

	// order
	if s.order.Len() > 0 {
		sql.WriteString(" ORDER BY ")
		s.order.WriteTo(sql)
		sql.Truncate(sql.Len() - 1)
	}

	// limit
	if len(s.limit) > 0 {
		vals, err := s.e.Dialect().LimitSQL(sql, s.limit[0], s.limit[1:]...)
		if err != nil {
			return nil, err
		}

		args = append(args, vals[0])
		if len(vals) > 1 {
			args = append(args, vals[1])
		}
	}

	return s.e.Query(replace, sql.String(), args...)
}
<commit_msg>添加sql.SelectMapString函数<commit_after>// Copyright 2015 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package orm

import (
	"bytes"
	"database/sql"
	"errors"
	"strconv"
	"strings"

	"github.com/issue9/orm/fetch"
)

const (
	and = iota
	or
)

const (
	asc = iota
	desc
)

var (
	ErrEmptyTableName = errors.New("未指定表名")
)

// 以函数链的方式产生SQL语句。
type SQL struct {
	e        engine
	table    string
	cond     *bytes.Buffer
	condArgs []interface{}
	order    *bytes.Buffer
	limit    []int // 元素0为limit，元素1为offset
}

func newSQL(engine engine) *SQL {
	return &SQL{
		e:        engine,
		cond:     new(bytes.Buffer),
		condArgs: []interface{}{},
		order:    new(bytes.Buffer),
		limit:    make([]int, 0, 2), // 最长2个，同时保存limit和offset
	}
}

func (s *SQL) Reset() *SQL {
	s.table = ""
	s.cond.Reset()
	s.condArgs = s.condArgs[:0]
	s.order.Reset()
	s.limit = s.limit[:0]

	return s
}

// 指定Limit相关的值。
func (s *SQL) Limit(limit int, offset ...int) *SQL {
	if len(offset) > 1 {
		panic("offset参数指定了太多的值")
	}

	s.limit = append(s.limit, limit)

	if len(offset) > 0 {
		s.limit = append(s.limit, offset[0])
	}

	return s
}

func (s *SQL) where(op int, cond string, args ...interface{}) *SQL {
	switch {
	case s.cond.Len() == 0:
		s.cond.WriteString(" WHERE(")
	case op == and:
		s.cond.WriteString(" AND(")
	case op == or:
		s.cond.WriteString(" OR(")
	default:
		panic("where:错误的参数op")
	}
	s.cond.WriteString(cond)
	s.cond.WriteByte(')')
	s.condArgs = append(s.condArgs, args...)

	return s
}

// 将之后的语句以and的形式与当前的语句进行连接
func (s *SQL) And(cond string, args ...interface{}) *SQL {
	return s.where(and, cond, args...)
}

// 将之后的语句以or的形式与当前的语句进行连接
func (s *SQL) Or(cond string, args ...interface{}) *SQL {
	return s.where(or, cond, args...)
}

// order by ... asc
// 当cols参数为空时，不产生任何操作。
func (s *SQL) Asc(cols ...string) *SQL {
	if len(cols) == 0 {
		return s
	}

	for _, col := range cols {
		s.order.WriteString(col)
		s.order.WriteByte(',')
	}
	s.order.Truncate(s.order.Len() - 1)
	s.order.WriteString(" ASC,")
	return s
}

// order by ... desc
// 当cols参数为空时，不产生任何操作。
func (s *SQL) Desc(cols ...string) *SQL {
	if len(cols) == 0 {
		return s
	}

	for _, col := range cols {
		s.order.WriteString(col)
		s.order.WriteByte(',')
	}
	s.order.Truncate(s.order.Len() - 1)
	s.order.WriteString(" DESC,")
	return s
}

// 指定表名。
func (s *SQL) Table(tableName string) *SQL {
	s.table = tableName
	return s
}

// 将符合当前条件的所有记录删除。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Delete(replace bool) (sql.Result, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	sql.Reset()
	sql.WriteString("DELETE FROM ")
	s.e.Dialect().Quote(sql, s.table)
	sql.WriteString(s.cond.String())
	return s.e.Exec(replace, sql.String(), s.condArgs...)
}

// 更新符合当前条件的所有记录。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Update(replace bool, data map[string]interface{}) (sql.Result, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	if len(data) == 0 {
		return nil, errors.New("SQL.Update:未指定需要更新的数据")
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	sql.Reset()
	sql.WriteString("UPDATE ")
	s.e.Dialect().Quote(sql, s.table)
	vals := make([]interface{}, 0, len(data)+len(s.condArgs))
	sql.WriteString(" SET ")
	for k, v := range data {
		s.e.Dialect().Quote(sql, k)
		sql.WriteString("=?,")
		vals = append(vals, v)
	}
	sql.Truncate(sql.Len() - 1) // 去掉最后一个逗号

	sql.WriteString(s.cond.String())
	return s.e.Exec(replace, sql.String(), append(vals, s.condArgs...)...)
}

// 返回符合条件的记录数量。
// 若指定了Limit，则相应的条件也会计算在内。
func (s *SQL) Count(replace bool) (int, error) {
	rows, err := s.query(replace, "COUNT(*) AS cnt")
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	cols, err := fetch.ColumnString(true, "cnt", rows)
	if err != nil {
		return 0, err
	}

	if len(cols) == 0 {
		return 0, nil
	}

	return strconv.Atoi(cols[0])
}

// 将符合当前条件的所有记录依次写入objs中。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) Select(replace bool, objs interface{}) (int, error) {
	rows, err := s.query(replace)
	if err != nil {
		return 0, err
	}
	cnt, err := fetch.Obj(objs, rows)
	rows.Close()
	return cnt, err
}

// 返回符合当前条件的所有记录。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) SelectMap(replace bool, cols ...string) ([]map[string]interface{}, error) {
	rows, err := s.query(replace, cols...)
	if err != nil {
		return nil, err
	}
	mapped, err := fetch.Map(false, rows)
	rows.Close()
	return mapped, err
}

// 返回符合当前条件的所有记录。
// replace，是否需将对语句的占位符进行替换。
func (s *SQL) SelectMapString(replace bool, cols ...string) ([]map[string]string, error) {
	rows, err := s.query(replace, cols...)
	if err != nil {
		return nil, err
	}
	mapped, err := fetch.MapString(false, rows)
	rows.Close()
	return mapped, err
}

func (s *SQL) query(replace bool, cols ...string) (*sql.Rows, error) {
	if len(s.table) == 0 {
		return nil, ErrEmptyTableName
	}

	sql := pool.Get().(*bytes.Buffer)
	defer pool.Put(sql)

	args := make([]interface{}, 0, len(s.condArgs)) // Query对应的参数
	sql.Reset()
	sql.WriteString("SELECT ")
	if len(cols) == 0 {
		sql.WriteString(" * ")
	} else {
		sql.WriteString(strings.Join(cols, ","))
	}

	// from
	sql.WriteString(" FROM ")
	s.e.Dialect().Quote(sql, s.table)

	// where
	sql.WriteString(s.cond.String())
	args = append(args, s.condArgs...)

	// order
	if s.order.Len() > 0 {
		sql.WriteString(" ORDER BY ")
		s.order.WriteTo(sql)
		sql.Truncate(sql.Len() - 1)
	}

	// limit
	if len(s.limit) > 0 {
		vals, err := s.e.Dialect().LimitSQL(sql, s.limit[0], s.limit[1:]...)
		if err != nil {
			return nil, err
		}

		args = append(args, vals[0])
		if len(vals) > 1 {
			args = append(args, vals[1])
		}
	}

	return s.e.Query(replace, sql.String(), args...)
}
<|endoftext|>