<commit_before>/*
Package cors is net/http handler to handle CORS related requests
as defined by http://www.w3.org/TR/cors/

You can configure it by passing an option struct to cors.New:

    c := cors.New(cors.Options{
        AllowedOrigins: []string{"foo.com"},
        AllowedMethods: []string{"GET", "POST", "DELETE"},
        AllowCredentials: true,
    })

Then insert the handler in the chain:

    handler = c.Handler(handler)

See Options documentation for more options.

The resulting handler is a standard net/http handler.
*/
package cors

import (
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/rs/xhandler"
	"golang.org/x/net/context"
)

// Options is a configuration container to setup the CORS middleware.
type Options struct {
	// AllowedOrigins is a list of origins a cross-domain request can be executed from.
	// If the special "*" value is present in the list, all origins will be allowed.
	// An origin may contain a wildcard (*) to replace 0 or more characters
	// (i.e.: http://*.domain.com). Usage of wildcards implies a small performance penality.
	// Only one wildcard can be used per origin.
	// Default value is ["*"]
	AllowedOrigins []string
	// AllowOriginFunc is a custom function to validate the origin. It take the origin
	// as argument and returns true if allowed or false otherwise. If this option is
	// set, the content of AllowedOrigins is ignored.
	AllowOriginFunc func(origin string) bool
	// AllowedMethods is a list of methods the client is allowed to use with
	// cross-domain requests. Default value is simple methods (GET and POST)
	AllowedMethods []string
	// AllowedHeaders is list of non simple headers the client is allowed to use with
	// cross-domain requests.
	// If the special "*" value is present in the list, all headers will be allowed.
	// Default value is [] but "Origin" is always appended to the list.
	AllowedHeaders []string
	// ExposedHeaders indicates which headers are safe to expose to the API of a CORS
	// API specification
	ExposedHeaders []string
	// AllowCredentials indicates whether the request can include user credentials like
	// cookies, HTTP authentication or client side SSL certificates.
	AllowCredentials bool
	// MaxAge indicates how long (in seconds) the results of a preflight request
	// can be cached
	MaxAge int
	// OptionsPassthrough instructs preflight to let other potential next handlers to
	// process the OPTIONS method. Turn this on if your application handles OPTIONS.
	OptionsPassthrough bool
	// Debugging flag adds additional output to debug server side CORS issues
	Debug bool
}

// Cors http handler
type Cors struct {
	// Debug logger
	Log *log.Logger
	// Set to true when allowed origins contains a "*"
	allowedOriginsAll bool
	// Normalized list of plain allowed origins
	allowedOrigins []string
	// List of allowed origins containing wildcards
	allowedWOrigins []wildcard
	// Optional origin validator function
	allowOriginFunc func(origin string) bool
	// Set to true when allowed headers contains a "*"
	allowedHeadersAll bool
	// Normalized list of allowed headers
	allowedHeaders []string
	// Normalized list of allowed methods
	allowedMethods []string
	// Normalized list of exposed headers
	exposedHeaders    []string
	allowCredentials  bool
	maxAge            int
	optionPassthrough bool
}

// New creates a new Cors handler with the provided options.
func New(options Options) *Cors {
	c := &Cors{
		exposedHeaders:    convert(options.ExposedHeaders, http.CanonicalHeaderKey),
		allowOriginFunc:   options.AllowOriginFunc,
		allowCredentials:  options.AllowCredentials,
		maxAge:            options.MaxAge,
		optionPassthrough: options.OptionsPassthrough,
	}
	if options.Debug {
		c.Log = log.New(os.Stdout, "[cors] ", log.LstdFlags)
	}

	// Normalize options
	// Note: for origins and methods matching, the spec requires a case-sensitive matching.
	// As it may error prone, we chose to ignore the spec here.

	// Allowed Origins
	if len(options.AllowedOrigins) == 0 {
		// Default is all origins
		c.allowedOriginsAll = true
	} else {
		c.allowedOrigins = []string{}
		c.allowedWOrigins = []wildcard{}
		for _, origin := range options.AllowedOrigins {
			// Normalize
			origin = strings.ToLower(origin)
			if origin == "*" {
				// If "*" is present in the list, turn the whole list into a match all
				c.allowedOriginsAll = true
				c.allowedOrigins = nil
				c.allowedWOrigins = nil
				break
			} else if i := strings.IndexByte(origin, '*'); i >= 0 {
				// Split the origin in two: start and end string without the *
				w := wildcard{origin[0:i], origin[i+1 : len(origin)]}
				c.allowedWOrigins = append(c.allowedWOrigins, w)
			} else {
				c.allowedOrigins = append(c.allowedOrigins, origin)
			}
		}
	}

	// Allowed Headers
	if len(options.AllowedHeaders) == 0 {
		// Use sensible defaults
		c.allowedHeaders = []string{"Origin", "Accept", "Content-Type"}
	} else {
		// Origin is always appended as some browsers will always request for this header at preflight
		c.allowedHeaders = convert(append(options.AllowedHeaders, "Origin"), http.CanonicalHeaderKey)
		for _, h := range options.AllowedHeaders {
			if h == "*" {
				c.allowedHeadersAll = true
				c.allowedHeaders = nil
				break
			}
		}
	}

	// Allowed Methods
	if len(options.AllowedMethods) == 0 {
		// Default is spec's "simple" methods
		c.allowedMethods = []string{"GET", "POST"}
	} else {
		c.allowedMethods = convert(options.AllowedMethods, strings.ToUpper)
	}

	return c
}

// Default creates a new Cors handler with default options
func Default() *Cors {
	return New(Options{})
}

// Handler apply the CORS specification on the request, and add relevant CORS headers
// as necessary.
func (c *Cors) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "OPTIONS" {
			c.logf("Handler: Preflight request")
			c.handlePreflight(w, r)
			// Preflight requests are standalone and should stop the chain as some other
			// middleware may not handle OPTIONS requests correctly. One typical example
			// is authentication middleware ; OPTIONS requests won't carry authentication
			// headers (see #1)
			if c.optionPassthrough {
				h.ServeHTTP(w, r)
			}
			else {
 +			w.WriteHeader(http.StatusOK)
 +			}
		} else {
			c.logf("Handler: Actual request")
			c.handleActualRequest(w, r)
			h.ServeHTTP(w, r)
		}
	})
}

// HandlerC is net/context aware handler
func (c *Cors) HandlerC(h xhandler.HandlerC) xhandler.HandlerC {
	return xhandler.HandlerFuncC(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		if r.Method == "OPTIONS" {
			c.logf("Handler: Preflight request")
			c.handlePreflight(w, r)
			// Preflight requests are standalone and should stop the chain as some other
			// middleware may not handle OPTIONS requests correctly. One typical example
			// is authentication middleware ; OPTIONS requests won't carry authentication
			// headers (see #1)
			if c.optionPassthrough {
				h.ServeHTTPC(ctx, w, r)
			}
		} else {
			c.logf("Handler: Actual request")
			c.handleActualRequest(w, r)
			h.ServeHTTPC(ctx, w, r)
		}
	})
}

// HandlerFunc provides Martini compatible handler
func (c *Cors) HandlerFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method == "OPTIONS" {
		c.logf("HandlerFunc: Preflight request")
		c.handlePreflight(w, r)
	} else {
		c.logf("HandlerFunc: Actual request")
		c.handleActualRequest(w, r)
	}
}

// Negroni compatible interface
func (c *Cors) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if r.Method == "OPTIONS" {
		c.logf("ServeHTTP: Preflight request")
		c.handlePreflight(w, r)
		// Preflight requests are standalone and should stop the chain as some other
		// middleware may not handle OPTIONS requests correctly. One typical example
		// is authentication middleware ; OPTIONS requests won't carry authentication
		// headers (see #1)
		if c.optionPassthrough {
			next(w, r)
		}
		else {
			w.WriteHeader(http.StatusOK)
		}
	} else {
		c.logf("ServeHTTP: Actual request")
		c.handleActualRequest(w, r)
		next(w, r)
	}
}

// handlePreflight handles pre-flight CORS requests
func (c *Cors) handlePreflight(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	origin := r.Header.Get("Origin")

	if r.Method != "OPTIONS" {
		c.logf("  Preflight aborted: %s!=OPTIONS", r.Method)
		return
	}
	// Always set Vary headers
	// see https://github.com/rs/cors/issues/10,
	//     https://github.com/rs/cors/commit/dbdca4d95feaa7511a46e6f1efb3b3aa505bc43f#commitcomment-12352001
	headers.Add("Vary", "Origin")
	headers.Add("Vary", "Access-Control-Request-Method")
	headers.Add("Vary", "Access-Control-Request-Headers")

	if origin == "" {
		c.logf("  Preflight aborted: empty origin")
		return
	}
	if !c.isOriginAllowed(origin) {
		c.logf("  Preflight aborted: origin '%s' not allowed", origin)
		return
	}

	reqMethod := r.Header.Get("Access-Control-Request-Method")
	if !c.isMethodAllowed(reqMethod) {
		c.logf("  Preflight aborted: method '%s' not allowed", reqMethod)
		return
	}
	reqHeaders := parseHeaderList(r.Header.Get("Access-Control-Request-Headers"))
	if !c.areHeadersAllowed(reqHeaders) {
		c.logf("  Preflight aborted: headers '%v' not allowed", reqHeaders)
		return
	}
	headers.Set("Access-Control-Allow-Origin", origin)
	// Spec says: Since the list of methods can be unbounded, simply returning the method indicated
	// by Access-Control-Request-Method (if supported) can be enough
	headers.Set("Access-Control-Allow-Methods", strings.ToUpper(reqMethod))
	if len(reqHeaders) > 0 {

		// Spec says: Since the list of headers can be unbounded, simply returning supported headers
		// from Access-Control-Request-Headers can be enough
		headers.Set("Access-Control-Allow-Headers", strings.Join(reqHeaders, ", "))
	}
	if c.allowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	if c.maxAge > 0 {
		headers.Set("Access-Control-Max-Age", strconv.Itoa(c.maxAge))
	}
	c.logf("  Preflight response headers: %v", headers)
}

// handleActualRequest handles simple cross-origin requests, actual request or redirects
func (c *Cors) handleActualRequest(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	origin := r.Header.Get("Origin")

	if r.Method == "OPTIONS" {
		c.logf("  Actual request no headers added: method == %s", r.Method)
		return
	}
	// Always set Vary, see https://github.com/rs/cors/issues/10
	headers.Add("Vary", "Origin")
	if origin == "" {
		c.logf("  Actual request no headers added: missing origin")
		return
	}
	if !c.isOriginAllowed(origin) {
		c.logf("  Actual request no headers added: origin '%s' not allowed", origin)
		return
	}

	// Note that spec does define a way to specifically disallow a simple method like GET or
	// POST. Access-Control-Allow-Methods is only used for pre-flight requests and the
	// spec doesn't instruct to check the allowed methods for simple cross-origin requests.
	// We think it's a nice feature to be able to have control on those methods though.
	if !c.isMethodAllowed(r.Method) {
		c.logf("  Actual request no headers added: method '%s' not allowed", r.Method)

		return
	}
	headers.Set("Access-Control-Allow-Origin", origin)
	if len(c.exposedHeaders) > 0 {
		headers.Set("Access-Control-Expose-Headers", strings.Join(c.exposedHeaders, ", "))
	}
	if c.allowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	c.logf("  Actual response added headers: %v", headers)
}

// convenience method. checks if debugging is turned on before printing
func (c *Cors) logf(format string, a ...interface{}) {
	if c.Log != nil {
		c.Log.Printf(format, a...)
	}
}

// isOriginAllowed checks if a given origin is allowed to perform cross-domain requests
// on the endpoint
func (c *Cors) isOriginAllowed(origin string) bool {
	if c.allowOriginFunc != nil {
		return c.allowOriginFunc(origin)
	}
	if c.allowedOriginsAll {
		return true
	}
	origin = strings.ToLower(origin)
	for _, o := range c.allowedOrigins {
		if o == origin {
			return true
		}
	}
	for _, w := range c.allowedWOrigins {
		if w.match(origin) {
			return true
		}
	}
	return false
}

// isMethodAllowed checks if a given method can be used as part of a cross-domain request
// on the endpoing
func (c *Cors) isMethodAllowed(method string) bool {
	if len(c.allowedMethods) == 0 {
		// If no method allowed, always return false, even for preflight request
		return false
	}
	method = strings.ToUpper(method)
	if method == "OPTIONS" {
		// Always allow preflight requests
		return true
	}
	for _, m := range c.allowedMethods {
		if m == method {
			return true
		}
	}
	return false
}

// areHeadersAllowed checks if a given list of headers are allowed to used within
// a cross-domain request.
func (c *Cors) areHeadersAllowed(requestedHeaders []string) bool {
	if c.allowedHeadersAll || len(requestedHeaders) == 0 {
		return true
	}
	for _, header := range requestedHeaders {
		header = http.CanonicalHeaderKey(header)
		found := false
		for _, h := range c.allowedHeaders {
			if h == header {
				found = true
			}
		}
		if !found {
			return false
		}
	}
	return true
}
<commit_msg>fix go if else<commit_after>/*
Package cors is net/http handler to handle CORS related requests
as defined by http://www.w3.org/TR/cors/

You can configure it by passing an option struct to cors.New:

    c := cors.New(cors.Options{
        AllowedOrigins: []string{"foo.com"},
        AllowedMethods: []string{"GET", "POST", "DELETE"},
        AllowCredentials: true,
    })

Then insert the handler in the chain:

    handler = c.Handler(handler)

See Options documentation for more options.

The resulting handler is a standard net/http handler.
*/
package cors

import (
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/rs/xhandler"
	"golang.org/x/net/context"
)

// Options is a configuration container to setup the CORS middleware.
type Options struct {
	// AllowedOrigins is a list of origins a cross-domain request can be executed from.
	// If the special "*" value is present in the list, all origins will be allowed.
	// An origin may contain a wildcard (*) to replace 0 or more characters
	// (i.e.: http://*.domain.com). Usage of wildcards implies a small performance penality.
	// Only one wildcard can be used per origin.
	// Default value is ["*"]
	AllowedOrigins []string
	// AllowOriginFunc is a custom function to validate the origin. It take the origin
	// as argument and returns true if allowed or false otherwise. If this option is
	// set, the content of AllowedOrigins is ignored.
	AllowOriginFunc func(origin string) bool
	// AllowedMethods is a list of methods the client is allowed to use with
	// cross-domain requests. Default value is simple methods (GET and POST)
	AllowedMethods []string
	// AllowedHeaders is list of non simple headers the client is allowed to use with
	// cross-domain requests.
	// If the special "*" value is present in the list, all headers will be allowed.
	// Default value is [] but "Origin" is always appended to the list.
	AllowedHeaders []string
	// ExposedHeaders indicates which headers are safe to expose to the API of a CORS
	// API specification
	ExposedHeaders []string
	// AllowCredentials indicates whether the request can include user credentials like
	// cookies, HTTP authentication or client side SSL certificates.
	AllowCredentials bool
	// MaxAge indicates how long (in seconds) the results of a preflight request
	// can be cached
	MaxAge int
	// OptionsPassthrough instructs preflight to let other potential next handlers to
	// process the OPTIONS method. Turn this on if your application handles OPTIONS.
	OptionsPassthrough bool
	// Debugging flag adds additional output to debug server side CORS issues
	Debug bool
}

// Cors http handler
type Cors struct {
	// Debug logger
	Log *log.Logger
	// Set to true when allowed origins contains a "*"
	allowedOriginsAll bool
	// Normalized list of plain allowed origins
	allowedOrigins []string
	// List of allowed origins containing wildcards
	allowedWOrigins []wildcard
	// Optional origin validator function
	allowOriginFunc func(origin string) bool
	// Set to true when allowed headers contains a "*"
	allowedHeadersAll bool
	// Normalized list of allowed headers
	allowedHeaders []string
	// Normalized list of allowed methods
	allowedMethods []string
	// Normalized list of exposed headers
	exposedHeaders    []string
	allowCredentials  bool
	maxAge            int
	optionPassthrough bool
}

// New creates a new Cors handler with the provided options.
func New(options Options) *Cors {
	c := &Cors{
		exposedHeaders:    convert(options.ExposedHeaders, http.CanonicalHeaderKey),
		allowOriginFunc:   options.AllowOriginFunc,
		allowCredentials:  options.AllowCredentials,
		maxAge:            options.MaxAge,
		optionPassthrough: options.OptionsPassthrough,
	}
	if options.Debug {
		c.Log = log.New(os.Stdout, "[cors] ", log.LstdFlags)
	}

	// Normalize options
	// Note: for origins and methods matching, the spec requires a case-sensitive matching.
	// As it may error prone, we chose to ignore the spec here.

	// Allowed Origins
	if len(options.AllowedOrigins) == 0 {
		// Default is all origins
		c.allowedOriginsAll = true
	} else {
		c.allowedOrigins = []string{}
		c.allowedWOrigins = []wildcard{}
		for _, origin := range options.AllowedOrigins {
			// Normalize
			origin = strings.ToLower(origin)
			if origin == "*" {
				// If "*" is present in the list, turn the whole list into a match all
				c.allowedOriginsAll = true
				c.allowedOrigins = nil
				c.allowedWOrigins = nil
				break
			} else if i := strings.IndexByte(origin, '*'); i >= 0 {
				// Split the origin in two: start and end string without the *
				w := wildcard{origin[0:i], origin[i+1 : len(origin)]}
				c.allowedWOrigins = append(c.allowedWOrigins, w)
			} else {
				c.allowedOrigins = append(c.allowedOrigins, origin)
			}
		}
	}

	// Allowed Headers
	if len(options.AllowedHeaders) == 0 {
		// Use sensible defaults
		c.allowedHeaders = []string{"Origin", "Accept", "Content-Type"}
	} else {
		// Origin is always appended as some browsers will always request for this header at preflight
		c.allowedHeaders = convert(append(options.AllowedHeaders, "Origin"), http.CanonicalHeaderKey)
		for _, h := range options.AllowedHeaders {
			if h == "*" {
				c.allowedHeadersAll = true
				c.allowedHeaders = nil
				break
			}
		}
	}

	// Allowed Methods
	if len(options.AllowedMethods) == 0 {
		// Default is spec's "simple" methods
		c.allowedMethods = []string{"GET", "POST"}
	} else {
		c.allowedMethods = convert(options.AllowedMethods, strings.ToUpper)
	}

	return c
}

// Default creates a new Cors handler with default options
func Default() *Cors {
	return New(Options{})
}

// Handler apply the CORS specification on the request, and add relevant CORS headers
// as necessary.
func (c *Cors) Handler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "OPTIONS" {
			c.logf("Handler: Preflight request")
			c.handlePreflight(w, r)
			// Preflight requests are standalone and should stop the chain as some other
			// middleware may not handle OPTIONS requests correctly. One typical example
			// is authentication middleware ; OPTIONS requests won't carry authentication
			// headers (see #1)
			if c.optionPassthrough {
				h.ServeHTTP(w, r)
			} else {
 +			w.WriteHeader(http.StatusOK)
 +			}
		} else {
			c.logf("Handler: Actual request")
			c.handleActualRequest(w, r)
			h.ServeHTTP(w, r)
		}
	})
}

// HandlerC is net/context aware handler
func (c *Cors) HandlerC(h xhandler.HandlerC) xhandler.HandlerC {
	return xhandler.HandlerFuncC(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
		if r.Method == "OPTIONS" {
			c.logf("Handler: Preflight request")
			c.handlePreflight(w, r)
			// Preflight requests are standalone and should stop the chain as some other
			// middleware may not handle OPTIONS requests correctly. One typical example
			// is authentication middleware ; OPTIONS requests won't carry authentication
			// headers (see #1)
			if c.optionPassthrough {
				h.ServeHTTPC(ctx, w, r)
			}
		} else {
			c.logf("Handler: Actual request")
			c.handleActualRequest(w, r)
			h.ServeHTTPC(ctx, w, r)
		}
	})
}

// HandlerFunc provides Martini compatible handler
func (c *Cors) HandlerFunc(w http.ResponseWriter, r *http.Request) {
	if r.Method == "OPTIONS" {
		c.logf("HandlerFunc: Preflight request")
		c.handlePreflight(w, r)
	} else {
		c.logf("HandlerFunc: Actual request")
		c.handleActualRequest(w, r)
	}
}

// Negroni compatible interface
func (c *Cors) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	if r.Method == "OPTIONS" {
		c.logf("ServeHTTP: Preflight request")
		c.handlePreflight(w, r)
		// Preflight requests are standalone and should stop the chain as some other
		// middleware may not handle OPTIONS requests correctly. One typical example
		// is authentication middleware ; OPTIONS requests won't carry authentication
		// headers (see #1)
		if c.optionPassthrough {
			next(w, r)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	} else {
		c.logf("ServeHTTP: Actual request")
		c.handleActualRequest(w, r)
		next(w, r)
	}
}

// handlePreflight handles pre-flight CORS requests
func (c *Cors) handlePreflight(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	origin := r.Header.Get("Origin")

	if r.Method != "OPTIONS" {
		c.logf("  Preflight aborted: %s!=OPTIONS", r.Method)
		return
	}
	// Always set Vary headers
	// see https://github.com/rs/cors/issues/10,
	//     https://github.com/rs/cors/commit/dbdca4d95feaa7511a46e6f1efb3b3aa505bc43f#commitcomment-12352001
	headers.Add("Vary", "Origin")
	headers.Add("Vary", "Access-Control-Request-Method")
	headers.Add("Vary", "Access-Control-Request-Headers")

	if origin == "" {
		c.logf("  Preflight aborted: empty origin")
		return
	}
	if !c.isOriginAllowed(origin) {
		c.logf("  Preflight aborted: origin '%s' not allowed", origin)
		return
	}

	reqMethod := r.Header.Get("Access-Control-Request-Method")
	if !c.isMethodAllowed(reqMethod) {
		c.logf("  Preflight aborted: method '%s' not allowed", reqMethod)
		return
	}
	reqHeaders := parseHeaderList(r.Header.Get("Access-Control-Request-Headers"))
	if !c.areHeadersAllowed(reqHeaders) {
		c.logf("  Preflight aborted: headers '%v' not allowed", reqHeaders)
		return
	}
	headers.Set("Access-Control-Allow-Origin", origin)
	// Spec says: Since the list of methods can be unbounded, simply returning the method indicated
	// by Access-Control-Request-Method (if supported) can be enough
	headers.Set("Access-Control-Allow-Methods", strings.ToUpper(reqMethod))
	if len(reqHeaders) > 0 {

		// Spec says: Since the list of headers can be unbounded, simply returning supported headers
		// from Access-Control-Request-Headers can be enough
		headers.Set("Access-Control-Allow-Headers", strings.Join(reqHeaders, ", "))
	}
	if c.allowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	if c.maxAge > 0 {
		headers.Set("Access-Control-Max-Age", strconv.Itoa(c.maxAge))
	}
	c.logf("  Preflight response headers: %v", headers)
}

// handleActualRequest handles simple cross-origin requests, actual request or redirects
func (c *Cors) handleActualRequest(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	origin := r.Header.Get("Origin")

	if r.Method == "OPTIONS" {
		c.logf("  Actual request no headers added: method == %s", r.Method)
		return
	}
	// Always set Vary, see https://github.com/rs/cors/issues/10
	headers.Add("Vary", "Origin")
	if origin == "" {
		c.logf("  Actual request no headers added: missing origin")
		return
	}
	if !c.isOriginAllowed(origin) {
		c.logf("  Actual request no headers added: origin '%s' not allowed", origin)
		return
	}

	// Note that spec does define a way to specifically disallow a simple method like GET or
	// POST. Access-Control-Allow-Methods is only used for pre-flight requests and the
	// spec doesn't instruct to check the allowed methods for simple cross-origin requests.
	// We think it's a nice feature to be able to have control on those methods though.
	if !c.isMethodAllowed(r.Method) {
		c.logf("  Actual request no headers added: method '%s' not allowed", r.Method)

		return
	}
	headers.Set("Access-Control-Allow-Origin", origin)
	if len(c.exposedHeaders) > 0 {
		headers.Set("Access-Control-Expose-Headers", strings.Join(c.exposedHeaders, ", "))
	}
	if c.allowCredentials {
		headers.Set("Access-Control-Allow-Credentials", "true")
	}
	c.logf("  Actual response added headers: %v", headers)
}

// convenience method. checks if debugging is turned on before printing
func (c *Cors) logf(format string, a ...interface{}) {
	if c.Log != nil {
		c.Log.Printf(format, a...)
	}
}

// isOriginAllowed checks if a given origin is allowed to perform cross-domain requests
// on the endpoint
func (c *Cors) isOriginAllowed(origin string) bool {
	if c.allowOriginFunc != nil {
		return c.allowOriginFunc(origin)
	}
	if c.allowedOriginsAll {
		return true
	}
	origin = strings.ToLower(origin)
	for _, o := range c.allowedOrigins {
		if o == origin {
			return true
		}
	}
	for _, w := range c.allowedWOrigins {
		if w.match(origin) {
			return true
		}
	}
	return false
}

// isMethodAllowed checks if a given method can be used as part of a cross-domain request
// on the endpoing
func (c *Cors) isMethodAllowed(method string) bool {
	if len(c.allowedMethods) == 0 {
		// If no method allowed, always return false, even for preflight request
		return false
	}
	method = strings.ToUpper(method)
	if method == "OPTIONS" {
		// Always allow preflight requests
		return true
	}
	for _, m := range c.allowedMethods {
		if m == method {
			return true
		}
	}
	return false
}

// areHeadersAllowed checks if a given list of headers are allowed to used within
// a cross-domain request.
func (c *Cors) areHeadersAllowed(requestedHeaders []string) bool {
	if c.allowedHeadersAll || len(requestedHeaders) == 0 {
		return true
	}
	for _, header := range requestedHeaders {
		header = http.CanonicalHeaderKey(header)
		found := false
		for _, h := range c.allowedHeaders {
			if h == header {
				found = true
			}
		}
		if !found {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package handlers

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

// CORSOption represents a functional option for configuring the CORS middleware.
type CORSOption func(*cors) error

type cors struct {
	h                      http.Handler
	allowedHeaders         []string
	allowedMethods         []string
	allowedOrigins         []string
	allowedOriginValidator OriginValidator
	exposedHeaders         []string
	maxAge                 int
	ignoreOptions          bool
	allowCredentials       bool
}

// OriginValidator takes an origin string and returns whether or not that origin is allowed.
type OriginValidator func(string) bool

var (
	defaultCorsMethods = []string{"GET", "HEAD", "POST"}
	defaultCorsHeaders = []string{"Accept", "Accept-Language", "Content-Language", "Origin"}
	// (WebKit/Safari v9 sends the Origin header by default in AJAX requests)
)

const (
	corsOptionMethod           string = "OPTIONS"
	corsAllowOriginHeader      string = "Access-Control-Allow-Origin"
	corsExposeHeadersHeader    string = "Access-Control-Expose-Headers"
	corsMaxAgeHeader           string = "Access-Control-Max-Age"
	corsAllowMethodsHeader     string = "Access-Control-Allow-Methods"
	corsAllowHeadersHeader     string = "Access-Control-Allow-Headers"
	corsAllowCredentialsHeader string = "Access-Control-Allow-Credentials"
	corsRequestMethodHeader    string = "Access-Control-Request-Method"
	corsRequestHeadersHeader   string = "Access-Control-Request-Headers"
	corsOriginHeader           string = "Origin"
	corsVaryHeader             string = "Vary"
	corsOriginMatchAll         string = "*"
)

func (ch *cors) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get(corsOriginHeader)
	if !ch.isOriginAllowed(origin) {
		ch.h.ServeHTTP(w, r)
		return
	}

	if r.Method == corsOptionMethod {
		if ch.ignoreOptions {
			ch.h.ServeHTTP(w, r)
			return
		}

		if _, ok := r.Header[corsRequestMethodHeader]; !ok {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		method := r.Header.Get(corsRequestMethodHeader)
		if !ch.isMatch(method, ch.allowedMethods) {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}

		requestHeaders := strings.Split(r.Header.Get(corsRequestHeadersHeader), ",")
		allowedHeaders := []string{}
		for _, v := range requestHeaders {
			canonicalHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if canonicalHeader == "" || ch.isMatch(canonicalHeader, defaultCorsHeaders) {
				continue
			}

			if !ch.isMatch(canonicalHeader, ch.allowedHeaders) {
				w.WriteHeader(http.StatusForbidden)
				return
			}

			allowedHeaders = append(allowedHeaders, canonicalHeader)
		}

		if len(allowedHeaders) > 0 {
			w.Header().Set(corsAllowHeadersHeader, strings.Join(allowedHeaders, ","))
		}

		if ch.maxAge > 0 {
			w.Header().Set(corsMaxAgeHeader, strconv.Itoa(ch.maxAge))
		}

		if !ch.isMatch(method, defaultCorsMethods) {
			w.Header().Set(corsAllowMethodsHeader, method)
		}
	} else {
		if len(ch.exposedHeaders) > 0 {
			w.Header().Set(corsExposeHeadersHeader, strings.Join(ch.exposedHeaders, ","))
		}
	}

	if ch.allowCredentials {
		w.Header().Set(corsAllowCredentialsHeader, "true")
	}

	if len(ch.allowedOrigins) > 1 {
		w.Header().Set(corsVaryHeader, corsOriginHeader)
	}

	returnOrigin := origin
	if ch.allowedOriginValidator == nil && len(ch.allowedOrigins) == 0 {
		returnOrigin = "*"
	} else {
		for _, o := range ch.allowedOrigins {
			// A configuration of * is different than explicitly setting an allowed
			// origin. Returning arbitrary origin headers an an access control allow
			// origin header is unsafe and is not required by any use case.
			if o == corsOriginMatchAll {
				returnOrigin = "*"
				break
			}
		}
	}
	w.Header().Set(corsAllowOriginHeader, returnOrigin)

	if r.Method == corsOptionMethod {
		return
	}
	ch.h.ServeHTTP(w, r)
}

// CORS provides Cross-Origin Resource Sharing middleware.
// Example:
//
//  import (
//      "net/http"
//
//      "github.com/gorilla/handlers"
//      "github.com/gorilla/mux"
//  )
//
//  func main() {
//      r := mux.NewRouter()
//      r.HandleFunc("/users", UserEndpoint)
//      r.HandleFunc("/projects", ProjectEndpoint)
//
//      // Apply the CORS middleware to our top-level router, with the defaults.
//      http.ListenAndServe(":8000", handlers.CORS()(r))
//  }
//
func CORS(opts ...CORSOption) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		ch := parseCORSOptions(opts...)
		ch.h = h
		return ch
	}
}

func parseCORSOptions(opts ...CORSOption) *cors {
	ch := &cors{
		allowedMethods: defaultCorsMethods,
		allowedHeaders: defaultCorsHeaders,
		allowedOrigins: []string{},
	}

	for _, option := range opts {
		option(ch)
	}

	return ch
}

//
// Functional options for configuring CORS.
//

// AllowedHeaders adds the provided headers to the list of allowed headers in a
// CORS request.
// This is an append operation so the headers Accept, Accept-Language,
// and Content-Language are always allowed.
// Content-Type must be explicitly declared if accepting Content-Types other than
// application/x-www-form-urlencoded, multipart/form-data, or text/plain.
func AllowedHeaders(headers []string) CORSOption {
	return func(ch *cors) error {
		for _, v := range headers {
			normalizedHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if normalizedHeader == "" {
				continue
			}

			if !ch.isMatch(normalizedHeader, ch.allowedHeaders) {
				ch.allowedHeaders = append(ch.allowedHeaders, normalizedHeader)
			}
		}

		return nil
	}
}

// AllowedMethods can be used to explicitly allow methods in the
// Access-Control-Allow-Methods header.
// This is a replacement operation so you must also
// pass GET, HEAD, and POST if you wish to support those methods.
func AllowedMethods(methods []string) CORSOption {
	return func(ch *cors) error {
		ch.allowedMethods = []string{}
		for _, v := range methods {
			normalizedMethod := strings.ToUpper(strings.TrimSpace(v))
			if normalizedMethod == "" {
				continue
			}

			if !ch.isMatch(normalizedMethod, ch.allowedMethods) {
				ch.allowedMethods = append(ch.allowedMethods, normalizedMethod)
			}
		}

		return nil
	}
}

// AllowedOrigins sets the allowed origins for CORS requests, as used in the
// 'Allow-Access-Control-Origin' HTTP header.
// Note: Passing in a []string{"*"} will allow any domain.
func AllowedOrigins(origins []string) CORSOption {
	return func(ch *cors) error {
		for _, v := range origins {
			if v == corsOriginMatchAll {
				ch.allowedOrigins = []string{corsOriginMatchAll}
				return nil
			}
		}

		ch.allowedOrigins = origins
		return nil
	}
}

// AllowedOriginValidator sets a function for evaluating allowed origins in CORS requests, represented by the
// 'Allow-Access-Control-Origin' HTTP header.
func AllowedOriginValidator(fn OriginValidator) CORSOption {
	return func(ch *cors) error {
		ch.allowedOriginValidator = fn
		return nil
	}
}

// ExposeHeaders can be used to specify headers that are available
// and will not be stripped out by the user-agent.
func ExposedHeaders(headers []string) CORSOption {
	return func(ch *cors) error {
		ch.exposedHeaders = []string{}
		for _, v := range headers {
			normalizedHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if normalizedHeader == "" {
				continue
			}

			if !ch.isMatch(normalizedHeader, ch.exposedHeaders) {
				ch.exposedHeaders = append(ch.exposedHeaders, normalizedHeader)
			}
		}

		return nil
	}
}

// MaxAge determines the maximum age (in seconds) between preflight requests. A
// maximum of 10 minutes is allowed. An age above this value will default to 10
// minutes.
func MaxAge(age int) CORSOption {
	return func(ch *cors) error {
		// Maximum of 10 minutes.
		if age > 600 {
			age = 600
		}

		ch.maxAge = age
		return nil
	}
}

// IgnoreOptions causes the CORS middleware to ignore OPTIONS requests, instead
// passing them through to the next handler. This is useful when your application
// or framework has a pre-existing mechanism for responding to OPTIONS requests.
func IgnoreOptions() CORSOption {
	return func(ch *cors) error {
		ch.ignoreOptions = true
		return nil
	}
}

// AllowCredentials can be used to specify that the user agent may pass
// authentication details along with the request.
func AllowCredentials() CORSOption {
	return func(ch *cors) error {
		ch.allowCredentials = true
		return nil
	}
}

func (ch *cors) isOriginAllowed(origin string) bool {
	if origin == "" {
		return false
	}

	if ch.allowedOriginValidator != nil {
		return ch.allowedOriginValidator(origin)
	}

	if len(ch.allowedOrigins) == 0 {
		return true
	}

	for _, allowedOrigin := range ch.allowedOrigins {
		if allowedOrigin == origin || allowedOrigin == corsOriginMatchAll {
			return true
		}
	}

	return false
}

func (ch *cors) isMatch(needle string, haystack []string) bool {
	for _, v := range haystack {
		if v == needle {
			return true
		}
	}

	return false
}
<commit_msg>Add Vary header when allowedOrigins is * (#114)<commit_after>package handlers

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

// CORSOption represents a functional option for configuring the CORS middleware.
type CORSOption func(*cors) error

type cors struct {
	h                      http.Handler
	allowedHeaders         []string
	allowedMethods         []string
	allowedOrigins         []string
	allowedOriginValidator OriginValidator
	exposedHeaders         []string
	maxAge                 int
	ignoreOptions          bool
	allowCredentials       bool
}

// OriginValidator takes an origin string and returns whether or not that origin is allowed.
type OriginValidator func(string) bool

var (
	defaultCorsMethods = []string{"GET", "HEAD", "POST"}
	defaultCorsHeaders = []string{"Accept", "Accept-Language", "Content-Language", "Origin"}
	// (WebKit/Safari v9 sends the Origin header by default in AJAX requests)
)

const (
	corsOptionMethod           string = "OPTIONS"
	corsAllowOriginHeader      string = "Access-Control-Allow-Origin"
	corsExposeHeadersHeader    string = "Access-Control-Expose-Headers"
	corsMaxAgeHeader           string = "Access-Control-Max-Age"
	corsAllowMethodsHeader     string = "Access-Control-Allow-Methods"
	corsAllowHeadersHeader     string = "Access-Control-Allow-Headers"
	corsAllowCredentialsHeader string = "Access-Control-Allow-Credentials"
	corsRequestMethodHeader    string = "Access-Control-Request-Method"
	corsRequestHeadersHeader   string = "Access-Control-Request-Headers"
	corsOriginHeader           string = "Origin"
	corsVaryHeader             string = "Vary"
	corsOriginMatchAll         string = "*"
)

func (ch *cors) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get(corsOriginHeader)
	if !ch.isOriginAllowed(origin) {
		ch.h.ServeHTTP(w, r)
		return
	}

	if r.Method == corsOptionMethod {
		if ch.ignoreOptions {
			ch.h.ServeHTTP(w, r)
			return
		}

		if _, ok := r.Header[corsRequestMethodHeader]; !ok {
			w.WriteHeader(http.StatusBadRequest)
			return
		}

		method := r.Header.Get(corsRequestMethodHeader)
		if !ch.isMatch(method, ch.allowedMethods) {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}

		requestHeaders := strings.Split(r.Header.Get(corsRequestHeadersHeader), ",")
		allowedHeaders := []string{}
		for _, v := range requestHeaders {
			canonicalHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if canonicalHeader == "" || ch.isMatch(canonicalHeader, defaultCorsHeaders) {
				continue
			}

			if !ch.isMatch(canonicalHeader, ch.allowedHeaders) {
				w.WriteHeader(http.StatusForbidden)
				return
			}

			allowedHeaders = append(allowedHeaders, canonicalHeader)
		}

		if len(allowedHeaders) > 0 {
			w.Header().Set(corsAllowHeadersHeader, strings.Join(allowedHeaders, ","))
		}

		if ch.maxAge > 0 {
			w.Header().Set(corsMaxAgeHeader, strconv.Itoa(ch.maxAge))
		}

		if !ch.isMatch(method, defaultCorsMethods) {
			w.Header().Set(corsAllowMethodsHeader, method)
		}
	} else {
		if len(ch.exposedHeaders) > 0 {
			w.Header().Set(corsExposeHeadersHeader, strings.Join(ch.exposedHeaders, ","))
		}
	}

	if ch.allowCredentials {
		w.Header().Set(corsAllowCredentialsHeader, "true")
	}

	if len(ch.allowedOrigins) > 1 || ch.allowedOrigins[0] == corsOriginMatchAll {
		w.Header().Set(corsVaryHeader, corsOriginHeader)
	}

	returnOrigin := origin
	if ch.allowedOriginValidator == nil && len(ch.allowedOrigins) == 0 {
		returnOrigin = "*"
	} else {
		for _, o := range ch.allowedOrigins {
			// A configuration of * is different than explicitly setting an allowed
			// origin. Returning arbitrary origin headers an an access control allow
			// origin header is unsafe and is not required by any use case.
			if o == corsOriginMatchAll {
				returnOrigin = "*"
				break
			}
		}
	}
	w.Header().Set(corsAllowOriginHeader, returnOrigin)

	if r.Method == corsOptionMethod {
		return
	}
	ch.h.ServeHTTP(w, r)
}

// CORS provides Cross-Origin Resource Sharing middleware.
// Example:
//
//  import (
//      "net/http"
//
//      "github.com/gorilla/handlers"
//      "github.com/gorilla/mux"
//  )
//
//  func main() {
//      r := mux.NewRouter()
//      r.HandleFunc("/users", UserEndpoint)
//      r.HandleFunc("/projects", ProjectEndpoint)
//
//      // Apply the CORS middleware to our top-level router, with the defaults.
//      http.ListenAndServe(":8000", handlers.CORS()(r))
//  }
//
func CORS(opts ...CORSOption) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		ch := parseCORSOptions(opts...)
		ch.h = h
		return ch
	}
}

func parseCORSOptions(opts ...CORSOption) *cors {
	ch := &cors{
		allowedMethods: defaultCorsMethods,
		allowedHeaders: defaultCorsHeaders,
		allowedOrigins: []string{},
	}

	for _, option := range opts {
		option(ch)
	}

	return ch
}

//
// Functional options for configuring CORS.
//

// AllowedHeaders adds the provided headers to the list of allowed headers in a
// CORS request.
// This is an append operation so the headers Accept, Accept-Language,
// and Content-Language are always allowed.
// Content-Type must be explicitly declared if accepting Content-Types other than
// application/x-www-form-urlencoded, multipart/form-data, or text/plain.
func AllowedHeaders(headers []string) CORSOption {
	return func(ch *cors) error {
		for _, v := range headers {
			normalizedHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if normalizedHeader == "" {
				continue
			}

			if !ch.isMatch(normalizedHeader, ch.allowedHeaders) {
				ch.allowedHeaders = append(ch.allowedHeaders, normalizedHeader)
			}
		}

		return nil
	}
}

// AllowedMethods can be used to explicitly allow methods in the
// Access-Control-Allow-Methods header.
// This is a replacement operation so you must also
// pass GET, HEAD, and POST if you wish to support those methods.
func AllowedMethods(methods []string) CORSOption {
	return func(ch *cors) error {
		ch.allowedMethods = []string{}
		for _, v := range methods {
			normalizedMethod := strings.ToUpper(strings.TrimSpace(v))
			if normalizedMethod == "" {
				continue
			}

			if !ch.isMatch(normalizedMethod, ch.allowedMethods) {
				ch.allowedMethods = append(ch.allowedMethods, normalizedMethod)
			}
		}

		return nil
	}
}

// AllowedOrigins sets the allowed origins for CORS requests, as used in the
// 'Allow-Access-Control-Origin' HTTP header.
// Note: Passing in a []string{"*"} will allow any domain.
func AllowedOrigins(origins []string) CORSOption {
	return func(ch *cors) error {
		for _, v := range origins {
			if v == corsOriginMatchAll {
				ch.allowedOrigins = []string{corsOriginMatchAll}
				return nil
			}
		}

		ch.allowedOrigins = origins
		return nil
	}
}

// AllowedOriginValidator sets a function for evaluating allowed origins in CORS requests, represented by the
// 'Allow-Access-Control-Origin' HTTP header.
func AllowedOriginValidator(fn OriginValidator) CORSOption {
	return func(ch *cors) error {
		ch.allowedOriginValidator = fn
		return nil
	}
}

// ExposeHeaders can be used to specify headers that are available
// and will not be stripped out by the user-agent.
func ExposedHeaders(headers []string) CORSOption {
	return func(ch *cors) error {
		ch.exposedHeaders = []string{}
		for _, v := range headers {
			normalizedHeader := http.CanonicalHeaderKey(strings.TrimSpace(v))
			if normalizedHeader == "" {
				continue
			}

			if !ch.isMatch(normalizedHeader, ch.exposedHeaders) {
				ch.exposedHeaders = append(ch.exposedHeaders, normalizedHeader)
			}
		}

		return nil
	}
}

// MaxAge determines the maximum age (in seconds) between preflight requests. A
// maximum of 10 minutes is allowed. An age above this value will default to 10
// minutes.
func MaxAge(age int) CORSOption {
	return func(ch *cors) error {
		// Maximum of 10 minutes.
		if age > 600 {
			age = 600
		}

		ch.maxAge = age
		return nil
	}
}

// IgnoreOptions causes the CORS middleware to ignore OPTIONS requests, instead
// passing them through to the next handler. This is useful when your application
// or framework has a pre-existing mechanism for responding to OPTIONS requests.
func IgnoreOptions() CORSOption {
	return func(ch *cors) error {
		ch.ignoreOptions = true
		return nil
	}
}

// AllowCredentials can be used to specify that the user agent may pass
// authentication details along with the request.
func AllowCredentials() CORSOption {
	return func(ch *cors) error {
		ch.allowCredentials = true
		return nil
	}
}

func (ch *cors) isOriginAllowed(origin string) bool {
	if origin == "" {
		return false
	}

	if ch.allowedOriginValidator != nil {
		return ch.allowedOriginValidator(origin)
	}

	if len(ch.allowedOrigins) == 0 {
		return true
	}

	for _, allowedOrigin := range ch.allowedOrigins {
		if allowedOrigin == origin || allowedOrigin == corsOriginMatchAll {
			return true
		}
	}

	return false
}

func (ch *cors) isMatch(needle string, haystack []string) bool {
	for _, v := range haystack {
		if v == needle {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Package cron implements a cron spec parser and runner.
package cron

import (
	"sort"
	"time"
)

// Cron keeps track of any number of entries, invoking the associated func as
// specified by the schedule. It may be started, stopped, and the entries may
// be inspected while running.
type Cron struct {
	entries  []*Entry
	stop     chan struct{}
	add      chan *Entry
	remove   chan EntryID
	snapshot chan []Entry
	running  bool
	nextID   EntryID
}

// Job is an interface for submitted cron jobs.
type Job interface {
	Run()
}

// Schedule describes a job's duty cycle.
type Schedule interface {
	// Next returns the next activation time, later than the given time.
	// Next is invoked initially, and then each time the job is run.
	Next(time.Time) time.Time
}

// EntryID identifies an entry within a Cron instance
type EntryID int

// Entry consists of a schedule and the func to execute on that schedule.
type Entry struct {
	// ID is the cron-assigned ID of this entry, which may be used to look up a
	// snapshot or remove it.
	ID EntryID

	// Schedule on which this job should be run.
	Schedule Schedule

	// Next time the job will run, or the zero time if Cron has not been
	// started or this entry's schedule is unsatisfiable
	Next time.Time

	// Prev is the last time this job was run, or the zero time if never.
	Prev time.Time

	// Job is the thing to run when the Schedule is activated.
	Job Job
}

// Valid returns true if this is not the zero entry.
func (e Entry) Valid() bool { return e.ID != 0 }

// byTime is a wrapper for sorting the entry array by time
// (with zero time at the end).
type byTime []*Entry

func (s byTime) Len() int      { return len(s) }
func (s byTime) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s byTime) Less(i, j int) bool {
	// Two zero times should return false.
	// Otherwise, zero is "greater" than any other time.
	// (To sort it at the end of the list.)
	if s[i].Next.IsZero() {
		return false
	}
	if s[j].Next.IsZero() {
		return true
	}
	return s[i].Next.Before(s[j].Next)
}

// New returns a new Cron job runner.
func New() *Cron {
	return &Cron{
		entries:  nil,
		add:      make(chan *Entry),
		stop:     make(chan struct{}),
		snapshot: make(chan []Entry),
		remove:   make(chan EntryID),
		running:  false,
	}
}

// FuncJob is a wrapper that turns a func() into a cron.Job
type FuncJob func()

func (f FuncJob) Run() { f() }

// AddFunc adds a func to the Cron to be run on the given schedule.
func (c *Cron) AddFunc(spec string, cmd func()) (EntryID, error) {
	return c.AddJob(spec, FuncJob(cmd))
}

// AddJob adds a Job to the Cron to be run on the given schedule.
func (c *Cron) AddJob(spec string, cmd Job) (EntryID, error) {
	schedule, err := Parse(spec)
	if err != nil {
		return 0, err
	}
	return c.Schedule(schedule, cmd), nil
}

// Schedule adds a Job to the Cron to be run on the given schedule.
func (c *Cron) Schedule(schedule Schedule, cmd Job) EntryID {
	c.nextID++
	entry := &Entry{
		ID:       c.nextID,
		Schedule: schedule,
		Job:      cmd,
	}
	if !c.running {
		c.entries = append(c.entries, entry)
	} else {
		c.add <- entry
	}
	return entry.ID
}

// Entries returns a snapshot of the cron entries.
func (c *Cron) Entries() []Entry {
	if c.running {
		c.snapshot <- nil
		return <-c.snapshot
	}
	return c.entrySnapshot()
}

// Entry returns a snapshot of the given entry, or nil if it couldn't be found.
func (c *Cron) Entry(id EntryID) Entry {
	for _, entry := range c.Entries() {
		if id == entry.ID {
			return entry
		}
	}
	return Entry{}
}

// Remove an entry from being run in the future.
func (c *Cron) Remove(id EntryID) {
	if c.running {
		c.remove <- id
	} else {
		c.removeEntry(id)
	}
}

// Start the cron scheduler in its own go-routine.
func (c *Cron) Start() {
	c.running = true
	go c.run()
}

// run the scheduler.. this is private just due to the need to synchronize
// access to the 'running' state variable.
func (c *Cron) run() {
	// Figure out the next activation times for each entry.
	now := time.Now().Local()
	for _, entry := range c.entries {
		entry.Next = entry.Schedule.Next(now)
	}

	for {
		// Determine the next entry to run.
		sort.Sort(byTime(c.entries))

		var effective time.Time
		if len(c.entries) == 0 || c.entries[0].Next.IsZero() {
			// If there are no entries yet, just sleep - it still handles new entries
			// and stop requests.
			effective = now.AddDate(10, 0, 0)
		} else {
			effective = c.entries[0].Next
		}

		select {
		case now = <-time.After(effective.Sub(now)):
			// Run every entry whose next time was this effective time.
			for _, e := range c.entries {
				if e.Next != effective {
					break
				}
				go e.Job.Run()
				e.Prev = e.Next
				e.Next = e.Schedule.Next(effective)
			}
			continue

		case newEntry := <-c.add:
			c.entries = append(c.entries, newEntry)
			newEntry.Next = newEntry.Schedule.Next(now)

		case <-c.snapshot:
			c.snapshot <- c.entrySnapshot()

		case id := <-c.remove:
			c.removeEntry(id)

		case <-c.stop:
			return
		}

		now = time.Now().Local()
	}
}

// Stop the cron scheduler.
func (c *Cron) Stop() {
	c.stop <- struct{}{}
	c.running = false
}

// entrySnapshot returns a copy of the current cron entry list.
func (c *Cron) entrySnapshot() []Entry {
	var entries = make([]Entry, len(c.entries))
	for i, e := range c.entries {
		entries[i] = *e
	}
	return entries
}

func (c *Cron) removeEntry(id EntryID) {
	var entries []*Entry
	for _, e := range c.entries {
		if e.ID != id {
			entries = append(entries, e)
		}
	}
	c.entries = entries
}
<commit_msg>Enforce custom import pat<commit_after>// Package cron implements a cron spec parser and runner.
package cron //import "gopkg.in/robfig/cron.v2"

import (
	"sort"
	"time"
)

// Cron keeps track of any number of entries, invoking the associated func as
// specified by the schedule. It may be started, stopped, and the entries may
// be inspected while running.
type Cron struct {
	entries  []*Entry
	stop     chan struct{}
	add      chan *Entry
	remove   chan EntryID
	snapshot chan []Entry
	running  bool
	nextID   EntryID
}

// Job is an interface for submitted cron jobs.
type Job interface {
	Run()
}

// Schedule describes a job's duty cycle.
type Schedule interface {
	// Next returns the next activation time, later than the given time.
	// Next is invoked initially, and then each time the job is run.
	Next(time.Time) time.Time
}

// EntryID identifies an entry within a Cron instance
type EntryID int

// Entry consists of a schedule and the func to execute on that schedule.
type Entry struct {
	// ID is the cron-assigned ID of this entry, which may be used to look up a
	// snapshot or remove it.
	ID EntryID

	// Schedule on which this job should be run.
	Schedule Schedule

	// Next time the job will run, or the zero time if Cron has not been
	// started or this entry's schedule is unsatisfiable
	Next time.Time

	// Prev is the last time this job was run, or the zero time if never.
	Prev time.Time

	// Job is the thing to run when the Schedule is activated.
	Job Job
}

// Valid returns true if this is not the zero entry.
func (e Entry) Valid() bool { return e.ID != 0 }

// byTime is a wrapper for sorting the entry array by time
// (with zero time at the end).
type byTime []*Entry

func (s byTime) Len() int      { return len(s) }
func (s byTime) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s byTime) Less(i, j int) bool {
	// Two zero times should return false.
	// Otherwise, zero is "greater" than any other time.
	// (To sort it at the end of the list.)
	if s[i].Next.IsZero() {
		return false
	}
	if s[j].Next.IsZero() {
		return true
	}
	return s[i].Next.Before(s[j].Next)
}

// New returns a new Cron job runner.
func New() *Cron {
	return &Cron{
		entries:  nil,
		add:      make(chan *Entry),
		stop:     make(chan struct{}),
		snapshot: make(chan []Entry),
		remove:   make(chan EntryID),
		running:  false,
	}
}

// FuncJob is a wrapper that turns a func() into a cron.Job
type FuncJob func()

func (f FuncJob) Run() { f() }

// AddFunc adds a func to the Cron to be run on the given schedule.
func (c *Cron) AddFunc(spec string, cmd func()) (EntryID, error) {
	return c.AddJob(spec, FuncJob(cmd))
}

// AddJob adds a Job to the Cron to be run on the given schedule.
func (c *Cron) AddJob(spec string, cmd Job) (EntryID, error) {
	schedule, err := Parse(spec)
	if err != nil {
		return 0, err
	}
	return c.Schedule(schedule, cmd), nil
}

// Schedule adds a Job to the Cron to be run on the given schedule.
func (c *Cron) Schedule(schedule Schedule, cmd Job) EntryID {
	c.nextID++
	entry := &Entry{
		ID:       c.nextID,
		Schedule: schedule,
		Job:      cmd,
	}
	if !c.running {
		c.entries = append(c.entries, entry)
	} else {
		c.add <- entry
	}
	return entry.ID
}

// Entries returns a snapshot of the cron entries.
func (c *Cron) Entries() []Entry {
	if c.running {
		c.snapshot <- nil
		return <-c.snapshot
	}
	return c.entrySnapshot()
}

// Entry returns a snapshot of the given entry, or nil if it couldn't be found.
func (c *Cron) Entry(id EntryID) Entry {
	for _, entry := range c.Entries() {
		if id == entry.ID {
			return entry
		}
	}
	return Entry{}
}

// Remove an entry from being run in the future.
func (c *Cron) Remove(id EntryID) {
	if c.running {
		c.remove <- id
	} else {
		c.removeEntry(id)
	}
}

// Start the cron scheduler in its own go-routine.
func (c *Cron) Start() {
	c.running = true
	go c.run()
}

// run the scheduler.. this is private just due to the need to synchronize
// access to the 'running' state variable.
func (c *Cron) run() {
	// Figure out the next activation times for each entry.
	now := time.Now().Local()
	for _, entry := range c.entries {
		entry.Next = entry.Schedule.Next(now)
	}

	for {
		// Determine the next entry to run.
		sort.Sort(byTime(c.entries))

		var effective time.Time
		if len(c.entries) == 0 || c.entries[0].Next.IsZero() {
			// If there are no entries yet, just sleep - it still handles new entries
			// and stop requests.
			effective = now.AddDate(10, 0, 0)
		} else {
			effective = c.entries[0].Next
		}

		select {
		case now = <-time.After(effective.Sub(now)):
			// Run every entry whose next time was this effective time.
			for _, e := range c.entries {
				if e.Next != effective {
					break
				}
				go e.Job.Run()
				e.Prev = e.Next
				e.Next = e.Schedule.Next(effective)
			}
			continue

		case newEntry := <-c.add:
			c.entries = append(c.entries, newEntry)
			newEntry.Next = newEntry.Schedule.Next(now)

		case <-c.snapshot:
			c.snapshot <- c.entrySnapshot()

		case id := <-c.remove:
			c.removeEntry(id)

		case <-c.stop:
			return
		}

		now = time.Now().Local()
	}
}

// Stop the cron scheduler.
func (c *Cron) Stop() {
	c.stop <- struct{}{}
	c.running = false
}

// entrySnapshot returns a copy of the current cron entry list.
func (c *Cron) entrySnapshot() []Entry {
	var entries = make([]Entry, len(c.entries))
	for i, e := range c.entries {
		entries[i] = *e
	}
	return entries
}

func (c *Cron) removeEntry(id EntryID) {
	var entries []*Entry
	for _, e := range c.entries {
		if e.ID != id {
			entries = append(entries, e)
		}
	}
	c.entries = entries
}
<|endoftext|>


<commit_before>package thermal

import (
	"code.google.com/p/go.crypto/nacl/box"
	"code.google.com/p/go.crypto/nacl/secretbox"
	"code.google.com/p/go.crypto/poly1305"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"log"
)

// cs3a is an implementation of the NaCl based cipher set 3a
type cs3a struct {
	id             string
	fingerprintBin []byte
	fingerprintHex string
	publicKey      [32]byte
	privateKey     [32]byte
}

// initialize generates a key pair and sets up the cipher set
func (cs *cs3a) initialize() error {

	// generate the key pair
	publicKey, privateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair in cs3a initialization")
		return err
	}

	// generate the fingerprint hash
	hash256 := sha256.New()
	hash256.Write(publicKey[:])
	fingerprintBin := hash256.Sum(nil)
	fingerprintHex := fmt.Sprintf("%x", fingerprintBin)

	// initialize the struct
	cs.id = "cs3a"
	cs.fingerprintBin = fingerprintBin
	cs.fingerprintHex = fingerprintHex
	cs.publicKey = *publicKey
	cs.privateKey = *privateKey

	return nil

}

func (cs *cs3a) String() string {
	return fmt.Sprintf("%s: %x", cs.id, cs.fingerprint)
}

func (cs *cs3a) csid() string {
	return cs.id
}

// fingerprint returns the csid and fingerprint for use in a 'parts' set
func (cs *cs3a) fingerprint() (string, string) {
	return cs.id, cs.fingerprintHex
}

func (cs *cs3a) encryptOpenPacket(packet []byte, receiverPublicKey *[32]byte) (openPacketBody []byte, err error) {

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// line key pair
	linePublicKey, linePrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair for line")
		return openPacketBody, err
	}

	// todo - store the line key pair?  return it?

	// Encrypt the inner packet
	var nonce [24]byte
	var lineSharedKey [32]byte
	var encInnerPacket []byte

	box.Precompute(&lineSharedKey, receiverPublicKey, linePrivateKey)
	secretbox.Seal(encInnerPacket, packet, &nonce, &lineSharedKey)

	// Generate the mac and assemble the body for the outer packet
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var macKey [32]byte
	var mac [16]byte
	var openPacketData []byte

	box.Precompute(&macKey, receiverPublicKey, &cs.privateKey)
	openPacketData = append(linePublicKey[:], encInnerPacket...)
	poly1305.Sum(&mac, openPacketData, &macKey)
	openPacketBody = append(mac[:], openPacketData...)

	return openPacketBody, nil

}

func (cs *cs3a) decryptOpenPacket(openPacketBody []byte) (packet []byte) {
	return packet
}

func (cs *cs3a) encryptLinePacket(packet []byte) (linePacketBody []byte) {
	return linePacketBody
}

func (cs *cs3a) decryptLinePacket(linePacketBody []byte) (packet []byte) {
	return packet
}
<commit_msg>early shot at cs3a decryptOpenPacket<commit_after>package thermal

import (
	"code.google.com/p/go.crypto/nacl/box"
	"code.google.com/p/go.crypto/nacl/secretbox"
	"code.google.com/p/go.crypto/poly1305"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"log"
)

// cs3a is an implementation of the NaCl based cipher set 3a
type cs3a struct {
	id             string
	fingerprintBin []byte
	fingerprintHex string
	publicKey      [32]byte
	privateKey     [32]byte
}

// initialize generates a key pair and sets up the cipher set
func (cs *cs3a) initialize() error {

	// generate the key pair
	publicKey, privateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair in cs3a initialization")
		return err
	}

	// generate the fingerprint hash
	hash256 := sha256.New()
	hash256.Write(publicKey[:])
	fingerprintBin := hash256.Sum(nil)
	fingerprintHex := fmt.Sprintf("%x", fingerprintBin)

	// initialize the struct
	cs.id = "cs3a"
	cs.fingerprintBin = fingerprintBin
	cs.fingerprintHex = fingerprintHex
	cs.publicKey = *publicKey
	cs.privateKey = *privateKey

	return nil

}

func (cs *cs3a) String() string {
	return fmt.Sprintf("%s: %x", cs.id, cs.fingerprint)
}

func (cs *cs3a) csid() string {
	return cs.id
}

// fingerprint returns the csid and fingerprint for use in a 'parts' set
func (cs *cs3a) fingerprint() (string, string) {
	return cs.id, cs.fingerprintHex
}

func (cs *cs3a) encryptOpenPacket(packet []byte, receiverPublicKey *[32]byte) (openPacketBody []byte, err error) {

	// todo - store or return the line shared key for use in line packet encryption

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// line key pair
	linePublicKey, linePrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair for line")
		return openPacketBody, err
	}

	// Encrypt the inner packet
	var nonce [24]byte
	var lineSharedKey [32]byte
	var encInnerPacket []byte

	box.Precompute(&lineSharedKey, receiverPublicKey, linePrivateKey)
	secretbox.Seal(encInnerPacket, packet, &nonce, &lineSharedKey)

	// Generate the mac and assemble the body for the outer packet
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var macKey [32]byte
	var mac [16]byte
	var openPacketData []byte

	box.Precompute(&macKey, receiverPublicKey, &cs.privateKey)
	openPacketData = append(linePublicKey[:], encInnerPacket...)
	poly1305.Sum(&mac, openPacketData, &macKey)
	openPacketBody = append(mac[:], openPacketData...)

	return openPacketBody, nil

}

func (cs *cs3a) decryptOpenPacket(openPacketBody []byte, senderPublicKey *[32]byte) (packet []byte, err error) {

	// todo - store or return the line shared key for use in line packet encryption

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// Unpack the outer packet body
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var mac [16]byte
	var senderLinePublicKey [32]byte
	var encInnerPacket []byte
	var openPacketData []byte

	copy(mac[:], openPacketBody[:16])
	copy(senderLinePublicKey[:], openPacketBody[16:48])
	copy(encInnerPacket[:], openPacketBody[48:])
	openPacketData = append(senderLinePublicKey[:], encInnerPacket...)

	// Verify the mac
	var authenticated bool
	var macKey [32]byte

	box.Precompute(&macKey, senderPublicKey, &cs.privateKey)
	authenticated = poly1305.Verify(&mac, openPacketData, &macKey)
	if !authenticated {
		msg := "Incoming open packet failed MAC authentication"
		log.Println(msg)
		err = fmt.Errorf(msg)
		return packet, err
	}

	// Decrypt the inner packet
	var nonce [24]byte
	var lineSharedKey [32]byte

	box.Precompute(&lineSharedKey, &senderLinePublicKey, &cs.privateKey)
	secretbox.Open(packet, encInnerPacket, &nonce, &lineSharedKey)

	return packet, nil
}

func (cs *cs3a) encryptLinePacket(packet []byte) (linePacketBody []byte) {
	return linePacketBody
}

func (cs *cs3a) decryptLinePacket(linePacketBody []byte) (packet []byte) {
	return packet
}
<|endoftext|>


<commit_before>package thermal

import (
	"code.google.com/p/go.crypto/nacl/box"
	"code.google.com/p/go.crypto/nacl/secretbox"
	"code.google.com/p/go.crypto/poly1305"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"log"
)

// cs3a is an implementation of the NaCl based cipher set 3a
type cs3a struct {
	id             string
	fingerprintBin []byte
	fingerprintHex string
	publicKey      [32]byte
	privateKey     [32]byte
}

// initialize generates a key pair and sets up the cipher set
func (cs *cs3a) initialize() error {

	// generate the key pair
	publicKey, privateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair in cs3a initialization")
		return err
	}

	// generate the fingerprint hash
	hash256 := sha256.New()
	hash256.Write(publicKey[:])
	fingerprintBin := hash256.Sum(nil)
	fingerprintHex := fmt.Sprintf("%x", fingerprintBin)

	// initialize the struct
	cs.id = "cs3a"
	cs.fingerprintBin = fingerprintBin
	cs.fingerprintHex = fingerprintHex
	cs.publicKey = *publicKey
	cs.privateKey = *privateKey

	return nil

}

func (cs *cs3a) String() string {
	return fmt.Sprintf("%s: %x", cs.id, cs.fingerprint)
}

func (cs *cs3a) csid() string {
	return cs.id
}

// fingerprint returns the csid and fingerprint for use in a 'parts' set
func (cs *cs3a) fingerprint() (string, string) {
	return cs.id, cs.fingerprintHex
}

func (cs *cs3a) encryptOpenPacket(packet []byte, receiverPublicKey *[32]byte) (openPacketBody []byte, err error) {

	// todo - return/store the egress line shared key for use in line packet encryption
	// return lineSharedKey [32]byte or a more generic []byte for a common interface across ciphersets?

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// line key pair
	linePublicKey, linePrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair for line")
		return openPacketBody, err
	}

	// Encrypt the inner packet
	var nonce [24]byte
	var lineSharedKey [32]byte
	var encInnerPacket []byte

	box.Precompute(&lineSharedKey, receiverPublicKey, linePrivateKey)
	secretbox.Seal(encInnerPacket, packet, &nonce, &lineSharedKey)

	// Generate the mac and assemble the body for the outer packet
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var macKey [32]byte
	var mac [16]byte
	var openPacketData []byte

	box.Precompute(&macKey, receiverPublicKey, &cs.privateKey)
	openPacketData = append(linePublicKey[:], encInnerPacket...)
	poly1305.Sum(&mac, openPacketData, &macKey)
	openPacketBody = append(mac[:], openPacketData...)

	return openPacketBody, nil

}

func (cs *cs3a) decryptOpenPacket(openPacketBody []byte, senderPublicKey *[32]byte) (packet []byte, err error) {

	// todo - return/store the ingress line shared key for use in line packet decryption
	// return lineSharedKey [32]byte or a more generic []byte for a common interface across ciphersets?

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// Unpack the outer packet body
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var mac [16]byte
	var senderLinePublicKey [32]byte
	var encInnerPacket []byte
	var openPacketData []byte

	copy(mac[:], openPacketBody[:16])
	copy(senderLinePublicKey[:], openPacketBody[16:48])
	copy(encInnerPacket[:], openPacketBody[48:])
	openPacketData = append(senderLinePublicKey[:], encInnerPacket...)

	// Verify the mac
	var authenticated bool
	var macKey [32]byte

	box.Precompute(&macKey, senderPublicKey, &cs.privateKey)
	authenticated = poly1305.Verify(&mac, openPacketData, &macKey)
	if !authenticated {
		msg := "Incoming open packet failed MAC authentication"
		log.Println(msg)
		err = fmt.Errorf(msg)
		return packet, err
	}

	// Decrypt the inner packet
	var nonce [24]byte
	var lineSharedKey [32]byte

	box.Precompute(&lineSharedKey, &senderLinePublicKey, &cs.privateKey)
	secretbox.Open(packet, encInnerPacket, &nonce, &lineSharedKey)

	return packet, nil
}

func (cs *cs3a) encryptLinePacket(packet []byte, lineSharedKey, senderLineId, receiverLineId []byte) (linePacketBody []byte) {
	var nonce [24]byte

	rand.Reader.Read(nonce[:])

	return linePacketBody
}

func (cs *cs3a) decryptLinePacket(linePacketBody []byte) (packet []byte) {
	return packet
}
<commit_msg>change var name.  update some signatures.<commit_after>package thermal

import (
	"code.google.com/p/go.crypto/nacl/box"
	"code.google.com/p/go.crypto/nacl/secretbox"
	"code.google.com/p/go.crypto/poly1305"
	"crypto/rand"
	"crypto/sha256"
	"fmt"
	"log"
)

// cs3a is an implementation of the NaCl based cipher set 3a
type cs3a struct {
	id             string
	fingerprintBin []byte
	fingerprintHex string
	publicKey      [32]byte
	privateKey     [32]byte
}

// initialize generates a key pair and sets up the cipher set
func (cs *cs3a) initialize() error {

	// generate the key pair
	publicKey, privateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair in cs3a initialization")
		return err
	}

	// generate the fingerprint hash
	hash256 := sha256.New()
	hash256.Write(publicKey[:])
	fingerprintBin := hash256.Sum(nil)
	fingerprintHex := fmt.Sprintf("%x", fingerprintBin)

	// initialize the struct
	cs.id = "cs3a"
	cs.fingerprintBin = fingerprintBin
	cs.fingerprintHex = fingerprintHex
	cs.publicKey = *publicKey
	cs.privateKey = *privateKey

	return nil

}

func (cs *cs3a) String() string {
	return fmt.Sprintf("%s: %x", cs.id, cs.fingerprint)
}

func (cs *cs3a) csid() string {
	return cs.id
}

// fingerprint returns the csid and fingerprint for use in a 'parts' set
func (cs *cs3a) fingerprint() (string, string) {
	return cs.id, cs.fingerprintHex
}

func (cs *cs3a) encryptOpenPacket(packet []byte, receiverPublicKey *[32]byte) (openPacketBody []byte, lineSharedSecret [32]byte, err error) {

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// line key pair
	linePublicKey, linePrivateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Println("Error generating NaCl keypair for line")
		return openPacketBody, lineSharedSecert, err
	}

	// Encrypt the inner packet
	var nonce [24]byte
	var encInnerPacket []byte

	box.Precompute(&lineSharedSecret, receiverPublicKey, linePrivateKey)
	secretbox.Seal(encInnerPacket, packet, &nonce, &lineSharedSecret)

	// Generate the mac and assemble the body for the outer packet
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var macKey [32]byte
	var mac [16]byte
	var openPacketData []byte

	box.Precompute(&macKey, receiverPublicKey, &cs.privateKey)
	openPacketData = append(linePublicKey[:], encInnerPacket...)
	poly1305.Sum(&mac, openPacketData, &macKey)
	openPacketBody = append(mac[:], openPacketData...)

	return openPacketBody, lineSharedSecret, nil

}

func (cs *cs3a) decryptOpenPacket(openPacketBody []byte, senderPublicKey *[32]byte) (packet []byte, lineSharedSecret [32]byte, err error) {

	// switch key pair
	// cs.publicKey and cs.privateKey should already be populated

	// Unpack the outer packet body
	// <mac><sender-line-public-key><encrypted-inner-packet-data>
	var mac [16]byte
	var senderLinePublicKey [32]byte
	var encInnerPacket []byte
	var openPacketData []byte

	copy(mac[:], openPacketBody[:16])
	copy(senderLinePublicKey[:], openPacketBody[16:48])
	copy(encInnerPacket[:], openPacketBody[48:])
	openPacketData = append(senderLinePublicKey[:], encInnerPacket...)

	// Verify the mac
	var authenticated bool
	var macKey [32]byte

	box.Precompute(&macKey, senderPublicKey, &cs.privateKey)
	authenticated = poly1305.Verify(&mac, openPacketData, &macKey)
	if !authenticated {
		msg := "Incoming open packet failed MAC authentication"
		log.Println(msg)
		err = fmt.Errorf(msg)
		return packet, lineSharedSecret, err
	}

	// Decrypt the inner packet
	var nonce [24]byte

	box.Precompute(&lineSharedSecret, &senderLinePublicKey, &cs.privateKey)
	secretbox.Open(packet, encInnerPacket, &nonce, &lineSharedSecret)

	return packet, lineSharedSecret, nil
}

func (cs *cs3a) encryptLinePacket(packet []byte, lineSharedSecret *[32]byte, senderLineId, receiverLineId *[16]byte) (linePacketBody []byte) {
	var nonce [24]byte

	rand.Reader.Read(nonce[:])

	return linePacketBody
}

func (cs *cs3a) decryptLinePacket(linePacketBody []byte, lineSharedSecret *[32]byte, senderLineId, receiverLineId *[16]byte) (packet []byte) {
	return packet
}
<|endoftext|>


<commit_before>package rethink

import (
	r "github.com/dancannon/gorethink"
)

type DB struct {
	Session *r.Session
}

func NewDB(session *r.Session) *DB {
	return &DB{
		Session: session,
	}
}

func (db *DB) Get(table, id string) (map[string]interface{}, error) {
	result, err := r.Table(table).Get(id).RunRow(db.Session)
	if err != nil || result.IsNil() {
		return nil, err
	}

	var response map[string]interface{}
	result.Scan(&response)
	return response, nil
}

func (db *DB) GetAll(query r.RqlTerm) ([]map[string]interface{}, error) {
	var results []map[string]interface{}
	rows, err := query.Run(db.Session)
	if err != nil {
		return results, err
	}

	for rows.Next() {
		var response map[string]interface{}
		rows.Scan(&response)
		results = append(results, response)
	}

	return results, nil
}
<commit_msg>Return error if we don't find a key.<commit_after>package rethink

import (
	r "github.com/dancannon/gorethink"
	"fmt"
)

type DB struct {
	Session *r.Session
}

func NewDB(session *r.Session) *DB {
	return &DB{
		Session: session,
	}
}

var emptyMap map[string]interface{}

func (db *DB) Get(table, id string) (map[string]interface{}, error) {
	result, err := r.Table(table).Get(id).RunRow(db.Session)
	switch {
	case err != nil:
		return emptyMap, err
	case result.IsNil():
		return emptyMap, fmt.Errorf("No such id: %s", id)
	default:
		var response map[string]interface{}
		result.Scan(&response)
		return response, nil
	}
}

func (db *DB) GetAll(query r.RqlTerm) ([]map[string]interface{}, error) {
	var results []map[string]interface{}
	rows, err := query.Run(db.Session)
	if err != nil {
		return results, err
	}

	for rows.Next() {
		var response map[string]interface{}
		rows.Scan(&response)
		results = append(results, response)
	}

	return results, nil
}
<|endoftext|>


<commit_before>package metrix

import (
	"compress/gzip"
	"encoding/json"
	"os"
	"path"
	"time"
)

type Snapshot struct {
	Hostname     string         `json:"hostname,omitempty"`
	Ec2Metadata  *Ec2Metadata   `json:"ec2_metadata,omitempty"`
	Disks        []*Disk        `json:"disks,omitempty"`
	LoadAvg      *LoadAvg       `json:"load_avg,omitempty"`
	Meminfo      *Meminfo       `json:"meminfo,omitempty"`
	ProcStats    []*ProcStat    `json:"proc_stats,omitempty"`
	ProcCmdlines []*ProcCmdline `json:"proc_cmdlines,omitempty"`
	Stat         *Stat          `json:"stat,omitempty"`
}

func (s *Snapshot) Load() error {
	funcs := []func() error{
		s.loadEc2Metadata,
		s.loadDisks,
		s.loadLoadAvg,
		s.loadMeminfo,
		s.loadProcStats,
		s.loadProcCmdlines,
		s.loadStat,
	}
	for _, f := range funcs {
		err := f()
		if err != nil {
			return err
		}
	}
	return nil
}

func (s Snapshot) Store(rootDir string) error {
	var err error
	if s.Hostname == "" {
		s.Hostname, err = os.Hostname()
		if err != nil {
			return err
		}
	}
	localPath := rootDir + "/snapshots/" + time.Now().UTC().Format("2006/01/02/15/2006-01-02T150405") + "-" + s.Hostname + ".json.gz"
	logger.Printf("writing snapshot to %s", localPath)
	tmpPath := localPath + ".tmp"
	err = os.MkdirAll(path.Dir(tmpPath), 0755)
	if err != nil {
		return err
	}
	f, err := os.Create(tmpPath)
	if err != nil {
		return err
	}
	defer f.Close()
	gz := gzip.NewWriter(f)
	defer gz.Close()
	enc := json.NewEncoder(gz)
	err = enc.Encode(s)
	if err != nil {
		return err
	}
	return os.Rename(tmpPath, localPath)
}

func (s *Snapshot) loadEc2Metadata() (err error) {
	defer benchmark("load ec2 metadata")()
	em, err := LoadEc2Metadata()
	if err != nil {
		dbg.Printf("ERROR loading ec2 metadata: %q", err)
	} else {
		s.Ec2Metadata = em
	}
	return nil
}

func (s *Snapshot) loadProcCmdlines() (err error) {
	s.ProcCmdlines, err = LoadProcCmdlines()
	return err
}

func (s *Snapshot) loadProcStats() (err error) {
	s.ProcStats, err = LoadProcStats()
	return err
}

func (s *Snapshot) loadMeminfo() (err error) {
	s.Meminfo, err = LoadMeminfo()
	return err
}

func (s *Snapshot) loadLoadAvg() (err error) {
	s.LoadAvg, err = LoadLoadAvg()
	return err
}

func (s *Snapshot) loadDisks() (err error) {
	s.Disks, err = LoadDisks()
	return err
}

func (s *Snapshot) loadStat() (err error) {
	s.Stat, err = LoadStat()
	return err
}
<commit_msg>add LoadedAt field to Snapshot, also load hostname<commit_after>package metrix

import (
	"compress/gzip"
	"encoding/json"
	"os"
	"path"
	"time"
)

type Snapshot struct {
	Disks        []*Disk        `json:"disks,omitempty"`
	Ec2Metadata  *Ec2Metadata   `json:"ec2_metadata,omitempty"`
	Hostname     string         `json:"hostname,omitempty"`
	LoadAvg      *LoadAvg       `json:"load_avg,omitempty"`
	Meminfo      *Meminfo       `json:"meminfo,omitempty"`
	ProcCmdlines []*ProcCmdline `json:"proc_cmdlines,omitempty"`
	ProcStats    []*ProcStat    `json:"proc_stats,omitempty"`
	Stat         *Stat          `json:"stat,omitempty"`
	TakenAt      time.Time      `json:"taken_at,omitempty"`
}

func (s *Snapshot) Load() error {
	funcs := []func() error{
		s.loadEc2Metadata,
		s.loadDisks,
		s.loadLoadAvg,
		s.loadMeminfo,
		s.loadProcStats,
		s.loadProcCmdlines,
		s.loadStat,
		s.loadTakenAt,
		s.loadHostname,
	}
	for _, f := range funcs {
		if err := f(); err != nil {
			return err
		}
	}
	return nil
}

func (s *Snapshot) loadHostname() (err error) {
	s.Hostname, err = os.Hostname()
	return err
}

func (s Snapshot) loadTakenAt() error {
	s.TakenAt = time.Now().UTC()
	return nil
}

func (s Snapshot) Store(rootDir string) error {
	var err error
	if s.Hostname == "" {
		s.Hostname, err = os.Hostname()
		if err != nil {
			return err
		}
	}
	localPath := rootDir + "/snapshots/" + time.Now().UTC().Format("2006/01/02/15/2006-01-02T150405") + "-" + s.Hostname + ".json.gz"
	logger.Printf("writing snapshot to %s", localPath)
	tmpPath := localPath + ".tmp"
	err = os.MkdirAll(path.Dir(tmpPath), 0755)
	if err != nil {
		return err
	}
	f, err := os.Create(tmpPath)
	if err != nil {
		return err
	}
	defer f.Close()
	gz := gzip.NewWriter(f)
	defer gz.Close()
	enc := json.NewEncoder(gz)
	err = enc.Encode(s)
	if err != nil {
		return err
	}
	return os.Rename(tmpPath, localPath)
}

func (s *Snapshot) loadEc2Metadata() (err error) {
	defer benchmark("load ec2 metadata")()
	em, err := LoadEc2Metadata()
	if err != nil {
		dbg.Printf("ERROR loading ec2 metadata: %q", err)
	} else {
		s.Ec2Metadata = em
	}
	return nil
}

func (s *Snapshot) loadProcCmdlines() (err error) {
	s.ProcCmdlines, err = LoadProcCmdlines()
	return err
}

func (s *Snapshot) loadProcStats() (err error) {
	s.ProcStats, err = LoadProcStats()
	return err
}

func (s *Snapshot) loadMeminfo() (err error) {
	s.Meminfo, err = LoadMeminfo()
	return err
}

func (s *Snapshot) loadLoadAvg() (err error) {
	s.LoadAvg, err = LoadLoadAvg()
	return err
}

func (s *Snapshot) loadDisks() (err error) {
	s.Disks, err = LoadDisks()
	return err
}

func (s *Snapshot) loadStat() (err error) {
	s.Stat, err = LoadStat()
	return err
}
<|endoftext|>


<commit_before>package migrate

import (
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/acsellers/dr/schema"
)

type GenericDB struct {
	DB                *sql.DB
	Convert           Translator
	Log               *log.Logger
	PrimaryKeyDef     string
	LengthableColumns map[string]bool
}

func (g *GenericDB) HasTable(table *schema.Table) (bool, error) {
	return false, fmt.Errorf("Need a RDBMS-specific for schema check functionality")
}

func (g *GenericDB) CreateTable(table *schema.Table) error {
	vals := []interface{}{}
	sql := fmt.Sprintf(
		"CREATE TABLE %s(",
		g.Convert.SQLTable(table.Name),
	)

	defs := []string{}
	for i, column := range table.Columns {
		switch {
		case i == 0 && column.Type == "integer":
			defs = append(
				defs,
				fmt.Sprintf(
					g.PrimaryKeyDef,
					g.Convert.SQLColumn(table.Name, column.Name),
				),
			)
		default:
			coldef := fmt.Sprintf(
				"%s %s",
				g.Convert.SQLColumn(table.Name, column.Name),
				strings.ToUpper(column.Type),
			)
			if column.Length != 0 && g.LengthableColumns[column.Type] {
				coldef += fmt.Sprintf(
					"(%d)", column.Length,
				)
			}
			defs = append(defs, coldef)
		}
	}
	for _, child := range table.ChildOf {
		defs = append(
			defs,
			fmt.Sprintf(
				"FOREIGN KEY(%s) REFERENCES %s(%s)",
				g.Convert.SQLColumn(table.Name, child.ChildColumn.Name),
				g.Convert.SQLTable(child.Parent.Name),
				g.Convert.SQLColumn(child.Parent.Name, child.Parent.PrimaryKeyColumn().Name),
			),
		)
	}

	for _, belonging := range table.BelongsTo {
		defs = append(
			defs,
			fmt.Sprintf(
				"FOREIGN KEY(%s) REFERENCES %s(%s)",
				g.Convert.SQLColumn(table.Name, belonging.ChildColumn.Name),
				g.Convert.SQLTable(belonging.Parent.Name),
				g.Convert.SQLColumn(belonging.Parent.Name, belonging.Parent.PrimaryKeyColumn().Name),
			),
		)
	}

	sql += strings.Join(defs, ", ") + ")"
	fmt.Println(sql)

	_, err := g.DB.Exec(sql, vals...)
	if err != nil {
		return fmt.Errorf("Error Creating Table\nSQL: %s\nError: %s", sql, err.Error())
	}
	return nil
}

func (g *GenericDB) UpdateTable(table *schema.Table) error {
	for _, col := range table.Columns {
		exists, err := g.HasColumn(table, &col)
		if err != nil {
			return err
		}
		if !exists {
			g.Log.Println("Adding New Column", col.Name, "to", table.Name)
			err = g.CreateColumn(table, &col)
			if err != nil {
				g.Log.Println("Error when adding column", err)
				return err
			}
		}
	}

	return nil
}

func (g *GenericDB) RemoveTable(table *schema.Table) error {
	_, err := g.DB.Exec(fmt.Sprint("DROP TABLE ", g.Convert.SQLTable(table.Name)))
	return err
}

func (g *GenericDB) RenameTable(table *schema.Table, oldName string) error {
	_, err := g.DB.Exec(
		fmt.Sprintf("ALTER TABLE %s RENAME TO %s", oldName, g.Convert.SQLTable(table.Name)),
	)
	return err
}

func (g *GenericDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	return false, fmt.Errorf("Need a RDBMS-specific for schema check functionality")
}

func (g *GenericDB) CreateColumn(table *schema.Table, col *schema.Column) error {
	coldef := fmt.Sprintf(
		"%s %s",
		g.Convert.SQLColumn(table.Name, col.Name),
		strings.ToUpper(col.Type),
	)
	if col.Length != 0 && g.LengthableColumns[col.Type] {
		coldef += fmt.Sprintf(
			"(%d)",
			col.Length,
		)
	}
	sql := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s", g.Convert.SQLTable(table.Name), coldef)
	_, err := g.DB.Exec(sql)

	return err
}

func (g *GenericDB) RenameColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

func (g *GenericDB) RemoveColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

func (g *GenericDB) ModifyColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

// SqliteDB is the standard for the GenericDB, so it has no overridden functions
type SqliteDB struct {
	GenericDB
}

func (s *SqliteDB) CreateTable(table *schema.Table) error {
	s.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	s.GenericDB.LengthableColumns = s.LengthableColumns()
	return s.GenericDB.CreateTable(table)
}

func (s *SqliteDB) UpdateTable(table *schema.Table) error {
	s.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	s.GenericDB.LengthableColumns = s.LengthableColumns()
	return s.GenericDB.UpdateTable(table)
}

func (s *SqliteDB) HasTable(table *schema.Table) (bool, error) {
	var cnt int64
	err := s.DB.QueryRow(
		`SELECT COUNT(name) FROM sqlite_master WHERE type='table' AND name=$1`,
		s.Convert.SQLTable(table.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (s *SqliteDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	rows, err := s.DB.Query(fmt.Sprintf("PRAGMA table_info(%s)", s.Convert.SQLTable(table.Name)))
	if err != nil {
		return false, err
	}
	var tinfo struct {
		CID        int
		Name       string
		Type       string
		NotNull    bool
		Default    interface{}
		PrimaryKey bool
	}
	for rows.Next() {
		err = rows.Scan(
			&tinfo.CID,
			&tinfo.Name,
			&tinfo.Type,
			&tinfo.NotNull,
			&tinfo.Default,
			&tinfo.PrimaryKey,
		)
		if err != nil {
			return false, err
		}
		if tinfo.Name == s.Convert.SQLColumn(table.Name, col.Name) {
			return true, nil
		}
	}
	return false, rows.Err()
}

func (*SqliteDB) String() string {
	return "sqlite"
}

func (*SqliteDB) LengthableColumns(table *schema.Table) map[string]bool {
	return map[string]bool{
		"varchar": true,
		"integer": true,
	}
}

// PostgresDB is pretty much the same as sqlite, except it can modify columns
type PostgresDB struct {
	GenericDB
}

func (*PostgresDB) String() string {
	return "Postgres Master Race"
}

func (p PostgresDB) HasTable(table *schema.Table) (bool, error) {
	var cnt int64
	err := p.DB.QueryRow(
		`SELECT COUNT(*) FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema = 'public' AND table_name = $1`,
		p.Convert.SQLTable(table.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (p *PostgresDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	var cnt int64
	err := p.DB.QueryRow(
		`SELECT COUNT(*) FROM information_schema.columns WHERE table_name = $1 AND column_name = $2`,
		p.Convert.SQLTable(table.Name),
		p.Convert.SQLColumn(table.Name, col.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (p *PostgresDB) CreateTable(table *schema.Table) error {
	p.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	p.GenericDB.LengthableColumns = p.LengthableColumns()
	return p.GenericDB.CreateTable(table)
}

func (p *PostgresDB) UpdateTable(table *schema.Table) error {
	p.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	p.GenericDB.LengthableColumns = p.LengthableColumns()
	return p.GenericDB.UpdateTable(table)
}

func (p *PostgresDB) LengthableColumns(table *schema.Table) map[string]bool {
	return map[string]bool{
		"varchar": true,
	}
}

// MysqlDB is ugh
type MysqlDB struct {
	GenericDB
}

func (*MysqlDB) String() string {
	return "myDerpDB"
}

func (*MysqlDB) LengthableColumns(table *schema.Table) map[string]bool {
	return map[string]bool{
		"varchar": true,
		"integer": true,
	}
}

func (m *MysqlDB) CreateTable(table *schema.Table) error {
	m.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	m.GenericDB.LengthableColumns = m.LengthableColumns()
	return m.GenericDB.CreateTable(table)
}

func (m *MysqlDB) UpdateTable(table *schema.Table) error {
	m.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	m.GenericDB.LengthableColumns = m.LengthableColumns()
	return m.GenericDB.UpdateTable(table)
}
<commit_msg>Miswrote LengthableColumns functions<commit_after>package migrate

import (
	"database/sql"
	"fmt"
	"log"
	"strings"

	"github.com/acsellers/dr/schema"
)

type GenericDB struct {
	DB                *sql.DB
	Convert           Translator
	Log               *log.Logger
	PrimaryKeyDef     string
	LengthableColumns map[string]bool
}

func (g *GenericDB) HasTable(table *schema.Table) (bool, error) {
	return false, fmt.Errorf("Need a RDBMS-specific for schema check functionality")
}

func (g *GenericDB) CreateTable(table *schema.Table) error {
	vals := []interface{}{}
	sql := fmt.Sprintf(
		"CREATE TABLE %s(",
		g.Convert.SQLTable(table.Name),
	)

	defs := []string{}
	for i, column := range table.Columns {
		switch {
		case i == 0 && column.Type == "integer":
			defs = append(
				defs,
				fmt.Sprintf(
					g.PrimaryKeyDef,
					g.Convert.SQLColumn(table.Name, column.Name),
				),
			)
		default:
			coldef := fmt.Sprintf(
				"%s %s",
				g.Convert.SQLColumn(table.Name, column.Name),
				strings.ToUpper(column.Type),
			)
			if column.Length != 0 && g.LengthableColumns[column.Type] {
				coldef += fmt.Sprintf(
					"(%d)", column.Length,
				)
			}
			defs = append(defs, coldef)
		}
	}
	for _, child := range table.ChildOf {
		defs = append(
			defs,
			fmt.Sprintf(
				"FOREIGN KEY(%s) REFERENCES %s(%s)",
				g.Convert.SQLColumn(table.Name, child.ChildColumn.Name),
				g.Convert.SQLTable(child.Parent.Name),
				g.Convert.SQLColumn(child.Parent.Name, child.Parent.PrimaryKeyColumn().Name),
			),
		)
	}

	for _, belonging := range table.BelongsTo {
		defs = append(
			defs,
			fmt.Sprintf(
				"FOREIGN KEY(%s) REFERENCES %s(%s)",
				g.Convert.SQLColumn(table.Name, belonging.ChildColumn.Name),
				g.Convert.SQLTable(belonging.Parent.Name),
				g.Convert.SQLColumn(belonging.Parent.Name, belonging.Parent.PrimaryKeyColumn().Name),
			),
		)
	}

	sql += strings.Join(defs, ", ") + ")"
	fmt.Println(sql)

	_, err := g.DB.Exec(sql, vals...)
	if err != nil {
		return fmt.Errorf("Error Creating Table\nSQL: %s\nError: %s", sql, err.Error())
	}
	return nil
}

func (g *GenericDB) UpdateTable(table *schema.Table) error {
	for _, col := range table.Columns {
		exists, err := g.HasColumn(table, &col)
		if err != nil {
			return err
		}
		if !exists {
			g.Log.Println("Adding New Column", col.Name, "to", table.Name)
			err = g.CreateColumn(table, &col)
			if err != nil {
				g.Log.Println("Error when adding column", err)
				return err
			}
		}
	}

	return nil
}

func (g *GenericDB) RemoveTable(table *schema.Table) error {
	_, err := g.DB.Exec(fmt.Sprint("DROP TABLE ", g.Convert.SQLTable(table.Name)))
	return err
}

func (g *GenericDB) RenameTable(table *schema.Table, oldName string) error {
	_, err := g.DB.Exec(
		fmt.Sprintf("ALTER TABLE %s RENAME TO %s", oldName, g.Convert.SQLTable(table.Name)),
	)
	return err
}

func (g *GenericDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	return false, fmt.Errorf("Need a RDBMS-specific for schema check functionality")
}

func (g *GenericDB) CreateColumn(table *schema.Table, col *schema.Column) error {
	coldef := fmt.Sprintf(
		"%s %s",
		g.Convert.SQLColumn(table.Name, col.Name),
		strings.ToUpper(col.Type),
	)
	if col.Length != 0 && g.LengthableColumns[col.Type] {
		coldef += fmt.Sprintf(
			"(%d)",
			col.Length,
		)
	}
	sql := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s", g.Convert.SQLTable(table.Name), coldef)
	_, err := g.DB.Exec(sql)

	return err
}

func (g *GenericDB) RenameColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

func (g *GenericDB) RemoveColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

func (g *GenericDB) ModifyColumn(table *schema.Table, col *schema.Column) error {
	return nil
}

// SqliteDB is the standard for the GenericDB, so it has no overridden functions
type SqliteDB struct {
	GenericDB
}

func (s *SqliteDB) CreateTable(table *schema.Table) error {
	s.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	s.GenericDB.LengthableColumns = s.LengthableColumns()
	return s.GenericDB.CreateTable(table)
}

func (s *SqliteDB) UpdateTable(table *schema.Table) error {
	s.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	s.GenericDB.LengthableColumns = s.LengthableColumns()
	return s.GenericDB.UpdateTable(table)
}

func (s *SqliteDB) HasTable(table *schema.Table) (bool, error) {
	var cnt int64
	err := s.DB.QueryRow(
		`SELECT COUNT(name) FROM sqlite_master WHERE type='table' AND name=$1`,
		s.Convert.SQLTable(table.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (s *SqliteDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	rows, err := s.DB.Query(fmt.Sprintf("PRAGMA table_info(%s)", s.Convert.SQLTable(table.Name)))
	if err != nil {
		return false, err
	}
	var tinfo struct {
		CID        int
		Name       string
		Type       string
		NotNull    bool
		Default    interface{}
		PrimaryKey bool
	}
	for rows.Next() {
		err = rows.Scan(
			&tinfo.CID,
			&tinfo.Name,
			&tinfo.Type,
			&tinfo.NotNull,
			&tinfo.Default,
			&tinfo.PrimaryKey,
		)
		if err != nil {
			return false, err
		}
		if tinfo.Name == s.Convert.SQLColumn(table.Name, col.Name) {
			return true, nil
		}
	}
	return false, rows.Err()
}

func (*SqliteDB) String() string {
	return "sqlite"
}

func (*SqliteDB) LengthableColumns() map[string]bool {
	return map[string]bool{
		"varchar": true,
		"integer": true,
	}
}

// PostgresDB is pretty much the same as sqlite, except it can modify columns
type PostgresDB struct {
	GenericDB
}

func (*PostgresDB) String() string {
	return "Postgres Master Race"
}

func (p PostgresDB) HasTable(table *schema.Table) (bool, error) {
	var cnt int64
	err := p.DB.QueryRow(
		`SELECT COUNT(*) FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema = 'public' AND table_name = $1`,
		p.Convert.SQLTable(table.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (p *PostgresDB) HasColumn(table *schema.Table, col *schema.Column) (bool, error) {
	var cnt int64
	err := p.DB.QueryRow(
		`SELECT COUNT(*) FROM information_schema.columns WHERE table_name = $1 AND column_name = $2`,
		p.Convert.SQLTable(table.Name),
		p.Convert.SQLColumn(table.Name, col.Name),
	).Scan(&cnt)
	if err != nil {
		return false, err
	}
	return cnt == 1, nil
}

func (p *PostgresDB) CreateTable(table *schema.Table) error {
	p.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	p.GenericDB.LengthableColumns = p.LengthableColumns()
	return p.GenericDB.CreateTable(table)
}

func (p *PostgresDB) UpdateTable(table *schema.Table) error {
	p.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	p.GenericDB.LengthableColumns = p.LengthableColumns()
	return p.GenericDB.UpdateTable(table)
}

func (p *PostgresDB) LengthableColumns() map[string]bool {
	return map[string]bool{
		"varchar": true,
	}
}

// MysqlDB is ugh
type MysqlDB struct {
	GenericDB
}

func (*MysqlDB) String() string {
	return "myDerpDB"
}

func (*MysqlDB) LengthableColumns() map[string]bool {
	return map[string]bool{
		"varchar": true,
		"integer": true,
	}
}

func (m *MysqlDB) CreateTable(table *schema.Table) error {
	m.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	m.GenericDB.LengthableColumns = m.LengthableColumns()
	return m.GenericDB.CreateTable(table)
}

func (m *MysqlDB) UpdateTable(table *schema.Table) error {
	m.GenericDB.PrimaryKeyDef = "%s SERIAL PRIMARY KEY"
	m.GenericDB.LengthableColumns = m.LengthableColumns()
	return m.GenericDB.UpdateTable(table)
}
<|endoftext|>