<commit_before>package websession

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"strings"

	"zenhack.net/go/sandstorm/capnp/util"
	"zenhack.net/go/sandstorm/capnp/websession"
	"zenhack.net/go/sandstorm/exp/util/bytestream"
	"zenhack.net/go/sandstorm/exp/util/handle"
	"zenhack.net/go/sandstorm/internal/errors"
)

var specialRequestHeaders = map[string]struct{}{
	"Accept":          {},
	"Accept-Encoding": {},
	"Cookie":          {},
	"If-Match":        {},
	"If-None-Match":   {},
}

// Parameters common to all websession request methods
type commonParams interface {
	Path() (string, error)
	Context() (websession.WebSession_Context, error)
}

// Our UiSession implementation; this implements WebSession, and holds both the
// SessionData from the new*Session call and the http.Handler to invoke.
type handlerWebSession struct {
	sessionData SessionData
	handler     http.Handler
}

//// Helpers for common parts of request handling ////

// Intialize a request with the data common to all webession request methods.
//
// The request will have a Context that is derived from ctx, but includes the
// SessionData in its Values.
func (h *handlerWebSession) initRequest(ctx context.Context, params commonParams) (*http.Request, error) {
	path, err := params.Path()
	if err != nil {
		return nil, err
	}

	wsCtx, err := params.Context()
	if err != nil {
		return nil, err
	}

	// Sandstorm gives us a path no leading slash, but Go's http library
	// expects one:
	parsedUrl, err := url.ParseRequestURI("/" + path)
	if err != nil {
		return nil, err
	}

	ctx = context.WithValue(ctx, sessionDataKey, h.sessionData)
	req := &http.Request{
		Header: http.Header{},
		URL:    parsedUrl,
	}

	req = req.WithContext(ctx)
	err = copyContextInfo(req, wsCtx)
	if err != nil {
		return nil, err
	}

	return req, nil
}

// Copy the information from the context into the request.
func copyContextInfo(req *http.Request, wsCtx websession.WebSession_Context) error {
	// cookies

	cookies, err := wsCtx.Cookies()
	if err != nil {
		return err
	}
	numCookies := cookies.Len()
	for i := 0; i < numCookies; i++ {
		kv := cookies.At(i)
		key, err := kv.Key()
		if err != nil {
			return err
		}
		val, err := kv.Value()
		if err != nil {
			return err
		}
		req.AddCookie(&http.Cookie{
			Name:  key,
			Value: val,
		})
	}

	// accept

	accept, err := wsCtx.Accept()
	if err != nil {
		return err
	}
	acceptHeaders := make([]string, accept.Len())
	for i := range acceptHeaders {
		str, err := formatAccept(accept.At(i))
		if err != nil {
			return err
		}
		acceptHeaders[i] = str
	}
	req.Header["Accept"] = acceptHeaders

	acceptEncoding, err := wsCtx.AcceptEncoding()
	if err != nil {
		return err
	}
	acceptEncodingHeaders := make([]string, acceptEncoding.Len())
	for i := range acceptEncodingHeaders {
		encoding, err := acceptEncoding.At(i).ContentCoding()
		if err != nil {
			return err
		}
		acceptEncodingHeaders[i] = fmt.Sprintf(
			"%s;q=%v",
			encoding,
			acceptEncoding.At(i).QValue())
	}
	req.Header["Accept-Encoding"] = acceptEncodingHeaders

	eTagPrecondition := wsCtx.ETagPrecondition()
	switch eTagPrecondition.Which() {
	case websession.WebSession_Context_eTagPrecondition_Which_none:
	case websession.WebSession_Context_eTagPrecondition_Which_exists:
		req.Header.Set("If-Match", "*")
	case websession.WebSession_Context_eTagPrecondition_Which_doesntExist:
		req.Header.Set("If-None-Match", "*")
	case websession.WebSession_Context_eTagPrecondition_Which_matchesOneOf:
		etags, err := eTagPrecondition.MatchesOneOf()
		if err != nil {
			return err
		}
		etagString, err := formatETags(etags)
		if err != nil {
			return err
		}
		req.Header.Set("If-Match", etagString)
	case websession.WebSession_Context_eTagPrecondition_Which_matchesNoneOf:
		etags, err := eTagPrecondition.MatchesNoneOf()
		if err != nil {
			return err
		}
		etagString, err := formatETags(etags)
		if err != nil {
			return err
		}
		req.Header.Set("If-None-Match", etagString)
	}

	additionalHeaders, err := wsCtx.AdditionalHeaders()
	if err != nil {
		return err
	}
	for i := 0; i < additionalHeaders.Len(); i++ {
		hdr := additionalHeaders.At(i)
		name, err := hdr.Name()
		if err != nil {
			return err
		}
		value, err := hdr.Value()
		if err != nil {
			return err
		}
		_, ok := specialRequestHeaders[name]
		if ok {
			log.Printf("Warning: special request header %q in "+
				"websession additionalHeaders field", name)
		}
		req.Header.Set(name, value)
	}

	return nil
}

// format the argument as expected for the value of the "Accept" header.
func formatAccept(typ websession.WebSession_AcceptedType) (string, error) {
	mimeType, err := typ.MimeType()
	if err != nil {
		return "", err
	}
	param := map[string]string{
		"q": fmt.Sprint(typ.QValue()),
	}
	return mime.FormatMediaType(mimeType, param), nil
}

func formatETags(etags websession.WebSession_ETag_List) (string, error) {
	etagStrings := make([]string, etags.Len())
	for i := range etagStrings {
		etag := etags.At(i)
		value, err := etag.Value()
		if err != nil {
			return "", err
		}

		// Sandstorm strips off the quotes, so we have to add them back:
		etagStrings[i] = `"` + value + `"`

		if etag.Weak() {
			etagStrings[i] = "W/" + etagStrings[i]
		}
	}
	return strings.Join(etagStrings, ", "), nil
}

// Common logic for the request methods that return a websession.WebSession_Request.
//
// `ctx` should be the context from the capnp method argument.
//
// `params` should be the capnp Params for the method.
//
// `response` should be the Response object to store the result in.
//
// `customize` is a function which will be called just before ServeHTTP. It
// is responsible for setting the HTTP Method on the request, as well as any
// other capnp method specific data.
func (h *handlerWebSession) handleCommon(
	ctx context.Context,
	params commonParams,
	response websession.WebSession_Response,
	customize func(*http.Request) error,
) error {
	ctx, cancel := handle.WithCancel(ctx)
	req, err := h.initRequest(ctx, params)
	if err != nil {
		return err
	}

	// Make sure it's not nil:
	req.Body = http.NoBody

	wsCtx, err := params.Context()
	if err != nil {
		return err
	}
	responseStream := wsCtx.ResponseStream()

	w := &basicResponseWriter{
		statusCode:     0,
		header:         http.Header{},
		cancel:         cancel,
		responseStream: responseStream,
		bodyWriter:     bytestream.ToWriteCloser(req.Context(), responseStream),
		response:       response,
	}

	err = customize(req)
	if err != nil {
		return err
	}

	h.handler.ServeHTTP(w, req)
	if w.statusCode == 0 {
		w.WriteHeader(200)
	}
	if w.bodyBuffer == nil {
		responseStream.Done(ctx, func(util.ByteStream_done_Params) error {
			return nil
		})
	} else {
		var errBody websession.WebSession_Response_ErrorBody
		if w.statusCode == 500 {
			errBody, err = w.response.ServerError().NonHtmlBody()
		} else {
			errBody, err = w.response.ClientError().NonHtmlBody()
		}
		if err != nil {
			panic("Error fetching ErrorBody:" + err.Error())
		}
		errBody.SetData(w.bodyBuffer.Bytes())
	}
	return nil
}

// PostContent/PutContent
type pContent interface {
	MimeType() (string, error)
	Content() ([]byte, error)
	HasEncoding() bool
	Encoding() (string, error)
}

func copyPContent(req *http.Request, content pContent) error {
	mimeType, err := content.MimeType()
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", mimeType)

	data, err := content.Content()
	if err != nil {
		return err
	}
	req.Body = ioutil.NopCloser(bytes.NewBuffer(data))

	if content.HasEncoding() {
		encoding, err := content.Encoding()
		if err != nil {
			return err
		}
		req.Header.Set("Content-Encoding", encoding)
	}
	return nil
}

// Common logic for capnp methods which take a `pContent` argument.
//
// `ctx`, `params`, and `response` are the same as in `handleCommon`.
//
// `content` is the `content` parameter.
// `method` is the HTTP method to set.
func (h *handlerWebSession) handlePContent(
	ctx context.Context,
	params commonParams,
	response websession.WebSession_Response,
	content pContent,
	method string,
) error {
	return h.handleCommon(ctx, params, response, func(req *http.Request) error {
		err := copyPContent(req, content)
		if err != nil {
			return err
		}
		req.Method = method
		return nil
	})
}

//// Actual WebSession methods ////

func (h *handlerWebSession) Get(p websession.WebSession_get) error {
	return h.handleCommon(p.Ctx, p.Params, p.Results, func(req *http.Request) error {
		if p.Params.IgnoreBody() {
			req.Method = "HEAD"
		} else {
			req.Method = "GET"
		}
		return nil
	})
}

func (h *handlerWebSession) Post(p websession.WebSession_post) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "POST")
}

func (h *handlerWebSession) Put(p websession.WebSession_put) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "PUT")
}

func (h *handlerWebSession) Delete(p websession.WebSession_delete) error {
	return h.handleCommon(p.Ctx, p.Params, p.Results, func(req *http.Request) error {
		req.Method = "DELETE"
		return nil
	})
}

func (h *handlerWebSession) Patch(p websession.WebSession_patch) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "PATCH")
}

//// Stubs for unimplemented WebSession methods ////

func (*handlerWebSession) PostStreaming(p websession.WebSession_postStreaming) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) PutStreaming(p websession.WebSession_putStreaming) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) OpenWebSocket(p websession.WebSession_openWebSocket) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Propfind(p websession.WebSession_propfind) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Proppatch(p websession.WebSession_proppatch) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Mkcol(p websession.WebSession_mkcol) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Copy(p websession.WebSession_copy) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Move(p websession.WebSession_move) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Lock(p websession.WebSession_lock) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Unlock(p websession.WebSession_unlock) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Acl(p websession.WebSession_acl) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Report(p websession.WebSession_report) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (h *handlerWebSession) Options(p websession.WebSession_options) error {
	return errors.UnImplementedExn(p.Results.Segment())
}
<commit_msg>Move http.NoBody assignment to initRequest<commit_after>package websession

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"strings"

	"zenhack.net/go/sandstorm/capnp/util"
	"zenhack.net/go/sandstorm/capnp/websession"
	"zenhack.net/go/sandstorm/exp/util/bytestream"
	"zenhack.net/go/sandstorm/exp/util/handle"
	"zenhack.net/go/sandstorm/internal/errors"
)

var specialRequestHeaders = map[string]struct{}{
	"Accept":          {},
	"Accept-Encoding": {},
	"Cookie":          {},
	"If-Match":        {},
	"If-None-Match":   {},
}

// Parameters common to all websession request methods
type commonParams interface {
	Path() (string, error)
	Context() (websession.WebSession_Context, error)
}

// Our UiSession implementation; this implements WebSession, and holds both the
// SessionData from the new*Session call and the http.Handler to invoke.
type handlerWebSession struct {
	sessionData SessionData
	handler     http.Handler
}

//// Helpers for common parts of request handling ////

// Intialize a request with the data common to all webession request methods.
//
// The request will have a Context that is derived from ctx, but includes the
// SessionData in its Values.
func (h *handlerWebSession) initRequest(ctx context.Context, params commonParams) (*http.Request, error) {
	path, err := params.Path()
	if err != nil {
		return nil, err
	}

	wsCtx, err := params.Context()
	if err != nil {
		return nil, err
	}

	// Sandstorm gives us a path no leading slash, but Go's http library
	// expects one:
	parsedUrl, err := url.ParseRequestURI("/" + path)
	if err != nil {
		return nil, err
	}

	ctx = context.WithValue(ctx, sessionDataKey, h.sessionData)
	req := &http.Request{
		Header: http.Header{},
		URL:    parsedUrl,
	}

	req = req.WithContext(ctx)
	err = copyContextInfo(req, wsCtx)
	if err != nil {
		return nil, err
	}

	// Set this as a default, so it's never nil.
	req.Body = http.NoBody

	return req, nil
}

// Copy the information from the context into the request.
func copyContextInfo(req *http.Request, wsCtx websession.WebSession_Context) error {
	// cookies

	cookies, err := wsCtx.Cookies()
	if err != nil {
		return err
	}
	numCookies := cookies.Len()
	for i := 0; i < numCookies; i++ {
		kv := cookies.At(i)
		key, err := kv.Key()
		if err != nil {
			return err
		}
		val, err := kv.Value()
		if err != nil {
			return err
		}
		req.AddCookie(&http.Cookie{
			Name:  key,
			Value: val,
		})
	}

	// accept

	accept, err := wsCtx.Accept()
	if err != nil {
		return err
	}
	acceptHeaders := make([]string, accept.Len())
	for i := range acceptHeaders {
		str, err := formatAccept(accept.At(i))
		if err != nil {
			return err
		}
		acceptHeaders[i] = str
	}
	req.Header["Accept"] = acceptHeaders

	acceptEncoding, err := wsCtx.AcceptEncoding()
	if err != nil {
		return err
	}
	acceptEncodingHeaders := make([]string, acceptEncoding.Len())
	for i := range acceptEncodingHeaders {
		encoding, err := acceptEncoding.At(i).ContentCoding()
		if err != nil {
			return err
		}
		acceptEncodingHeaders[i] = fmt.Sprintf(
			"%s;q=%v",
			encoding,
			acceptEncoding.At(i).QValue())
	}
	req.Header["Accept-Encoding"] = acceptEncodingHeaders

	eTagPrecondition := wsCtx.ETagPrecondition()
	switch eTagPrecondition.Which() {
	case websession.WebSession_Context_eTagPrecondition_Which_none:
	case websession.WebSession_Context_eTagPrecondition_Which_exists:
		req.Header.Set("If-Match", "*")
	case websession.WebSession_Context_eTagPrecondition_Which_doesntExist:
		req.Header.Set("If-None-Match", "*")
	case websession.WebSession_Context_eTagPrecondition_Which_matchesOneOf:
		etags, err := eTagPrecondition.MatchesOneOf()
		if err != nil {
			return err
		}
		etagString, err := formatETags(etags)
		if err != nil {
			return err
		}
		req.Header.Set("If-Match", etagString)
	case websession.WebSession_Context_eTagPrecondition_Which_matchesNoneOf:
		etags, err := eTagPrecondition.MatchesNoneOf()
		if err != nil {
			return err
		}
		etagString, err := formatETags(etags)
		if err != nil {
			return err
		}
		req.Header.Set("If-None-Match", etagString)
	}

	additionalHeaders, err := wsCtx.AdditionalHeaders()
	if err != nil {
		return err
	}
	for i := 0; i < additionalHeaders.Len(); i++ {
		hdr := additionalHeaders.At(i)
		name, err := hdr.Name()
		if err != nil {
			return err
		}
		value, err := hdr.Value()
		if err != nil {
			return err
		}
		_, ok := specialRequestHeaders[name]
		if ok {
			log.Printf("Warning: special request header %q in "+
				"websession additionalHeaders field", name)
		}
		req.Header.Set(name, value)
	}

	return nil
}

// format the argument as expected for the value of the "Accept" header.
func formatAccept(typ websession.WebSession_AcceptedType) (string, error) {
	mimeType, err := typ.MimeType()
	if err != nil {
		return "", err
	}
	param := map[string]string{
		"q": fmt.Sprint(typ.QValue()),
	}
	return mime.FormatMediaType(mimeType, param), nil
}

func formatETags(etags websession.WebSession_ETag_List) (string, error) {
	etagStrings := make([]string, etags.Len())
	for i := range etagStrings {
		etag := etags.At(i)
		value, err := etag.Value()
		if err != nil {
			return "", err
		}

		// Sandstorm strips off the quotes, so we have to add them back:
		etagStrings[i] = `"` + value + `"`

		if etag.Weak() {
			etagStrings[i] = "W/" + etagStrings[i]
		}
	}
	return strings.Join(etagStrings, ", "), nil
}

// Common logic for the request methods that return a websession.WebSession_Request.
//
// `ctx` should be the context from the capnp method argument.
//
// `params` should be the capnp Params for the method.
//
// `response` should be the Response object to store the result in.
//
// `customize` is a function which will be called just before ServeHTTP. It
// is responsible for setting the HTTP Method on the request, as well as any
// other capnp method specific data.
func (h *handlerWebSession) handleCommon(
	ctx context.Context,
	params commonParams,
	response websession.WebSession_Response,
	customize func(*http.Request) error,
) error {
	ctx, cancel := handle.WithCancel(ctx)
	req, err := h.initRequest(ctx, params)
	if err != nil {
		return err
	}

	wsCtx, err := params.Context()
	if err != nil {
		return err
	}
	responseStream := wsCtx.ResponseStream()

	w := &basicResponseWriter{
		statusCode:     0,
		header:         http.Header{},
		cancel:         cancel,
		responseStream: responseStream,
		bodyWriter:     bytestream.ToWriteCloser(req.Context(), responseStream),
		response:       response,
	}

	err = customize(req)
	if err != nil {
		return err
	}

	h.handler.ServeHTTP(w, req)
	if w.statusCode == 0 {
		w.WriteHeader(200)
	}
	if w.bodyBuffer == nil {
		responseStream.Done(ctx, func(util.ByteStream_done_Params) error {
			return nil
		})
	} else {
		var errBody websession.WebSession_Response_ErrorBody
		if w.statusCode == 500 {
			errBody, err = w.response.ServerError().NonHtmlBody()
		} else {
			errBody, err = w.response.ClientError().NonHtmlBody()
		}
		if err != nil {
			panic("Error fetching ErrorBody:" + err.Error())
		}
		errBody.SetData(w.bodyBuffer.Bytes())
	}
	return nil
}

// PostContent/PutContent
type pContent interface {
	MimeType() (string, error)
	Content() ([]byte, error)
	HasEncoding() bool
	Encoding() (string, error)
}

func copyPContent(req *http.Request, content pContent) error {
	mimeType, err := content.MimeType()
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", mimeType)

	data, err := content.Content()
	if err != nil {
		return err
	}
	req.Body = ioutil.NopCloser(bytes.NewBuffer(data))

	if content.HasEncoding() {
		encoding, err := content.Encoding()
		if err != nil {
			return err
		}
		req.Header.Set("Content-Encoding", encoding)
	}
	return nil
}

// Common logic for capnp methods which take a `pContent` argument.
//
// `ctx`, `params`, and `response` are the same as in `handleCommon`.
//
// `content` is the `content` parameter.
// `method` is the HTTP method to set.
func (h *handlerWebSession) handlePContent(
	ctx context.Context,
	params commonParams,
	response websession.WebSession_Response,
	content pContent,
	method string,
) error {
	return h.handleCommon(ctx, params, response, func(req *http.Request) error {
		err := copyPContent(req, content)
		if err != nil {
			return err
		}
		req.Method = method
		return nil
	})
}

//// Actual WebSession methods ////

func (h *handlerWebSession) Get(p websession.WebSession_get) error {
	return h.handleCommon(p.Ctx, p.Params, p.Results, func(req *http.Request) error {
		if p.Params.IgnoreBody() {
			req.Method = "HEAD"
		} else {
			req.Method = "GET"
		}
		return nil
	})
}

func (h *handlerWebSession) Post(p websession.WebSession_post) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "POST")
}

func (h *handlerWebSession) Put(p websession.WebSession_put) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "PUT")
}

func (h *handlerWebSession) Delete(p websession.WebSession_delete) error {
	return h.handleCommon(p.Ctx, p.Params, p.Results, func(req *http.Request) error {
		req.Method = "DELETE"
		return nil
	})
}

func (h *handlerWebSession) Patch(p websession.WebSession_patch) error {
	content, err := p.Params.Content()
	if err != nil {
		return err
	}
	return h.handlePContent(p.Ctx, p.Params, p.Results, content, "PATCH")
}

//// Stubs for unimplemented WebSession methods ////

func (*handlerWebSession) PostStreaming(p websession.WebSession_postStreaming) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) PutStreaming(p websession.WebSession_putStreaming) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) OpenWebSocket(p websession.WebSession_openWebSocket) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Propfind(p websession.WebSession_propfind) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Proppatch(p websession.WebSession_proppatch) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Mkcol(p websession.WebSession_mkcol) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Copy(p websession.WebSession_copy) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Move(p websession.WebSession_move) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Lock(p websession.WebSession_lock) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Unlock(p websession.WebSession_unlock) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Acl(p websession.WebSession_acl) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (*handlerWebSession) Report(p websession.WebSession_report) error {
	return errors.UnImplementedExn(p.Results.Segment())
}

func (h *handlerWebSession) Options(p websession.WebSession_options) error {
	return errors.UnImplementedExn(p.Results.Segment())
}
<|endoftext|>


<commit_before>package router

import (
	. "atlantis/manager/constant"
	"atlantis/manager/datamodel"
	"atlantis/manager/dns"
	"atlantis/manager/helper"
)

func Register(zone, ip string) (*datamodel.ZkRouter, error) {
	// create ZkRouter
	zkRouter := datamodel.Router(zone, ip)
	if dns.Provider == nil {
		// if we have no dns provider then just save here
		return zkRouter, zkRouter.Save()
	}
	// first delete all entries we may already have for this IP in DNS
	err := dns.DeleteRecordsForIP(ip)
	if err != nil {
		return nil, err
	}
	// choose cname
	routers, err := datamodel.ListRoutersInZone(zone)
	if err != nil {
		return zkRouter, err
	}
	routerMap := map[string]bool{}
	for _, router := range routers {
		tmpRouter, err := datamodel.GetRouter(zone, router)
		if err != nil {
			return zkRouter, err
		}
		routerMap[tmpRouter.CName] = true
	}
	routerNum := 1
	zkRouter.CName = helper.GetRouterCName(routerNum, zone, dns.Provider.Suffix())
	for ; routerMap[zkRouter.CName]; routerNum++ {
		zkRouter.CName = helper.GetRouterCName(routerNum, zone, dns.Provider.Suffix())
	}
	// create basic health check for router
	zkRouter.HealthCheckId, err = dns.Provider.CreateHealthCheck(zkRouter.IP)
	if err != nil {
		return zkRouter, err
	}
	// create the basic entries for the new router
	// e.g. region is us-east-1 && zone is us-east-1a && there are zones a, b, c, and d && this is the 2nd in a
	//   create router.us-east-1.<suffix>   primary
	//   create router2.us-east-1a.<suffix> primary
	//   create router.us-east-1a.<suffix>  primary
	//   create router.us-east-1b.<suffix>  secondary
	//   create router.us-east-1c.<suffix>  secondary
	//   create router.us-east-1d.<suffix>  secondary
	zkRouter.RecordIds = make([]string, 3)
	cnames := make([]dns.CName, 3)
	// PRIMARY router.<region>.<suffix>
	cnames[0] = dns.CName{
		Primary: true,
		CName:   helper.GetRegionRouterCName(dns.Provider.Suffix()),
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[0] = cnames[0].Id()
	// PRIMARY routerX.<region+zone>.<suffix>
	cnames[1] = dns.CName{
		Primary: true,
		CName:   zkRouter.CName,
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[1] = cnames[1].Id()
	// PRIMARY router.<region+zone>.<suffix>
	cnames[2] = dns.CName{
		Primary: true,
		CName:   helper.GetZoneRouterCName(zkRouter.Zone, dns.Provider.Suffix()),
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[2] = cnames[2].Id()
	// SECONDARY router.<region+zone>.<suffix>
	for _, azone := range AvailableZones {
		if azone == zone {
			continue
		}
		cname := dns.CName{
			Primary: false,
			CName:   helper.GetZoneRouterCName(azone, dns.Provider.Suffix()),
			IP:      zkRouter.IP,
		}
		zkRouter.RecordIds = append(zkRouter.RecordIds, cname.Id())
		cnames = append(cnames, cname)
	}
	err, errChan := dns.Provider.CreateCNames("CREATE_ROUTER "+ip+" in "+zone, cnames)
	if err != nil {
		return zkRouter, err
	}
	err = <-errChan // wait for change to propagate
	if err != nil {
		return zkRouter, err
	}
	return zkRouter, zkRouter.Save()
}

func Unregister(zone, ip string) error {
	zkRouter, err := datamodel.GetRouter(zone, ip)
	if err != nil {
		return err
	}
	if dns.Provider == nil {
		// if we have no dns provider then just save here
		return zkRouter.Delete()
	}
	// delete the basic entries for the new router
	// e.g. region is us-east-1 && zone is us-east-1a && there are zones a, b, c, and d && this is the 2nd in a
	//   delete router2.us-east-1a.<suffix> primary
	//   delete router.us-east-1a.<suffix>  primary
	//   delete router.us-east-1b.<suffix>  secondary
	//   delete router.us-east-1c.<suffix>  secondary
	//   delete router.us-east-1d.<suffix>  secondary
	err, errChan := dns.Provider.DeleteRecords("DELETE_ROUTER "+ip+" in "+zone, zkRouter.RecordIds...)
	if err != nil {
		return err
	}
	err = <-errChan // wait for it to propagate
	if err != nil {
		return err
	}
	// delete basic health check for router
	err = dns.Provider.DeleteHealthCheck(zkRouter.HealthCheckId)
	if err != nil {
		return err
	}
	return zkRouter.Delete()
}
<commit_msg>Move comment to top, maybe simplify.<commit_after>package router

import (
	. "atlantis/manager/constant"
	"atlantis/manager/datamodel"
	"atlantis/manager/dns"
	"atlantis/manager/helper"
)

// Registering the 4th router in us-east-1a, with IP (say) 10.0.0.4 needs to create the following
// entries in DNS (route53 in our case), assuming that atlantis.com is the DNS zone delegated to
// the deployment. Also assume that the deployment spans zones a, b and c in region us-east-1.
//
// - An A record for value router4.us-east-1a.atlantis.com pointing at 10.0.0.4
// - A primary failover A record for value router.us-east-1a.atlantis.com pointing at 10.0.0.4
// - A secondary failover A record for value router.us-east-1b.atlantis.com pointing to 10.0.0.4
// - A secondary failover A record for value router.us-east-1c.atlantis.com pointing to 10.0.0.4
// - A round robin A record for value router.us-east-1.atlantis.com pointing at 10.0.0.4
//
// Deleting the router, simply deletes all the records created when adding it.

func Register(zone, ip string) (*datamodel.ZkRouter, error) {
	// create ZkRouter
	zkRouter := datamodel.Router(zone, ip)
	if dns.Provider == nil {
		// if we have no dns provider then just save here
		return zkRouter, zkRouter.Save()
	}
	// first delete all entries we may already have for this IP in DNS
	err := dns.DeleteRecordsForIP(ip)
	if err != nil {
		return nil, err
	}
	// choose cname
	routers, err := datamodel.ListRoutersInZone(zone)
	if err != nil {
		return zkRouter, err
	}
	routerMap := map[string]bool{}
	for _, router := range routers {
		tmpRouter, err := datamodel.GetRouter(zone, router)
		if err != nil {
			return zkRouter, err
		}
		routerMap[tmpRouter.CName] = true
	}
	routerNum := 1
	zkRouter.CName = helper.GetRouterCName(routerNum, zone, dns.Provider.Suffix())
	for ; routerMap[zkRouter.CName]; routerNum++ {
		zkRouter.CName = helper.GetRouterCName(routerNum, zone, dns.Provider.Suffix())
	}
	// create basic health check for router
	zkRouter.HealthCheckId, err = dns.Provider.CreateHealthCheck(zkRouter.IP)
	if err != nil {
		return zkRouter, err
	}
	zkRouter.RecordIds = make([]string, 3)
	cnames := make([]dns.CName, 3)
	// PRIMARY router.<region>.<suffix>
	cnames[0] = dns.CName{
		Primary: true,
		CName:   helper.GetRegionRouterCName(dns.Provider.Suffix()),
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[0] = cnames[0].Id()
	// PRIMARY routerX.<region+zone>.<suffix>
	cnames[1] = dns.CName{
		Primary: true,
		CName:   zkRouter.CName,
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[1] = cnames[1].Id()
	// PRIMARY router.<region+zone>.<suffix>
	cnames[2] = dns.CName{
		Primary: true,
		CName:   helper.GetZoneRouterCName(zkRouter.Zone, dns.Provider.Suffix()),
		IP:      zkRouter.IP,
	}
	zkRouter.RecordIds[2] = cnames[2].Id()
	// SECONDARY router.<region+zone>.<suffix>
	for _, azone := range AvailableZones {
		if azone == zone {
			continue
		}
		cname := dns.CName{
			Primary: false,
			CName:   helper.GetZoneRouterCName(azone, dns.Provider.Suffix()),
			IP:      zkRouter.IP,
		}
		zkRouter.RecordIds = append(zkRouter.RecordIds, cname.Id())
		cnames = append(cnames, cname)
	}
	err, errChan := dns.Provider.CreateCNames("CREATE_ROUTER "+ip+" in "+zone, cnames)
	if err != nil {
		return zkRouter, err
	}
	err = <-errChan // wait for change to propagate
	if err != nil {
		return zkRouter, err
	}
	return zkRouter, zkRouter.Save()
}

func Unregister(zone, ip string) error {
	zkRouter, err := datamodel.GetRouter(zone, ip)
	if err != nil {
		return err
	}
	if dns.Provider == nil {
		// if we have no dns provider then just save here
		return zkRouter.Delete()
	}
	err, errChan := dns.Provider.DeleteRecords("DELETE_ROUTER "+ip+" in "+zone, zkRouter.RecordIds...)
	if err != nil {
		return err
	}
	err = <-errChan // wait for it to propagate
	if err != nil {
		return err
	}
	// delete basic health check for router
	err = dns.Provider.DeleteHealthCheck(zkRouter.HealthCheckId)
	if err != nil {
		return err
	}
	return zkRouter.Delete()
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"

	"github.com/jessevdk/go-flags"
	"github.com/pkg/errors"

	"github.com/nerdalize/nerd/pkg/transfer"
	"github.com/nerdalize/nerd/svc"

	"github.com/mitchellh/cli"
)

//DatasetUpload command
type DatasetUpload struct {
	KubeOpts
	TransferOpts
	Name string `long:"name" short:"n" description:"assign a name to the dataset"`

	*command
}

//DatasetUploadFactory creates the command
func DatasetUploadFactory(ui cli.Ui) cli.CommandFactory {
	cmd := &DatasetUpload{}
	cmd.command = createCommand(ui, cmd.Execute, cmd.Description, cmd.Usage, cmd, flags.PassAfterNonOption)
	return func() (cli.Command, error) {
		return cmd, nil
	}
}

//Execute runs the command
func (cmd *DatasetUpload) Execute(args []string) (err error) {
	if len(args) < 1 {
		return errShowUsage(MessageNotEnoughArguments)
	}

	deps, err := NewDeps(cmd.Logger(), cmd.KubeOpts)
	if err != nil {
		return renderConfigError(err, "failed to configure")
	}

	kube := svc.NewKube(deps)
	mgr, sto, sta, err := cmd.TransferOpts.TransferManager(kube)
	if err != nil {
		return errors.Wrap(err, "failed to setup transfer manager")
	}

	ctx := context.Background()
	var h transfer.Handle
	if h, err = mgr.Create(
		ctx,
		cmd.Name,
		*sto,
		*sta,
	); err != nil {
		return errors.Wrap(err, "failed to create transfer handle")
	}

	defer h.Close()

	err = h.Push(ctx, args[0], &progressBarReporter{})
	if err != nil {
		return errors.Wrap(err, "failed to upload dataset")
	}

	cmd.out.Infof("Uploaded dataset: '%s'", h.Name())
	cmd.out.Infof("To run a job with a dataset, use: 'nerd job run'")
	return nil
}

// Description returns long-form help text
func (cmd *DatasetUpload) Description() string { return cmd.Synopsis() }

// Synopsis returns a one-line
func (cmd *DatasetUpload) Synopsis() string { return "Upload a dataset to your compute cluster." }

// Usage shows usage
func (cmd *DatasetUpload) Usage() string {
	return "nerd dataset upload [--name=] ~/my-project/my-input-1"
}
<commit_msg>delete dataset k8s object when upload fails, fixes #302<commit_after>package cmd

import (
	"context"

	"github.com/jessevdk/go-flags"
	"github.com/pkg/errors"

	"github.com/nerdalize/nerd/pkg/transfer"
	"github.com/nerdalize/nerd/svc"

	"github.com/mitchellh/cli"
)

//DatasetUpload command
type DatasetUpload struct {
	KubeOpts
	TransferOpts
	Name string `long:"name" short:"n" description:"assign a name to the dataset"`

	*command
}

//DatasetUploadFactory creates the command
func DatasetUploadFactory(ui cli.Ui) cli.CommandFactory {
	cmd := &DatasetUpload{}
	cmd.command = createCommand(ui, cmd.Execute, cmd.Description, cmd.Usage, cmd, flags.PassAfterNonOption)
	return func() (cli.Command, error) {
		return cmd, nil
	}
}

//Execute runs the command
func (cmd *DatasetUpload) Execute(args []string) (err error) {
	if len(args) < 1 {
		return errShowUsage(MessageNotEnoughArguments)
	}

	deps, err := NewDeps(cmd.Logger(), cmd.KubeOpts)
	if err != nil {
		return renderConfigError(err, "failed to configure")
	}

	kube := svc.NewKube(deps)
	mgr, sto, sta, err := cmd.TransferOpts.TransferManager(kube)
	if err != nil {
		return errors.Wrap(err, "failed to setup transfer manager")
	}

	ctx := context.Background()
	var h transfer.Handle
	if h, err = mgr.Create(
		ctx,
		cmd.Name,
		*sto,
		*sta,
	); err != nil {
		return errors.Wrap(err, "failed to create transfer handle")
	}

	defer h.Close()

	err = h.Push(ctx, args[0], &progressBarReporter{})
	if err != nil {
		e := mgr.Remove(ctx, h.Name())
		if e != nil {
			return errors.Wrapf(err, "failed to upload dataset: %v", e)
		}
		return errors.Wrap(err, "failed to upload dataset")
	}

	cmd.out.Infof("Uploaded dataset: '%s'", h.Name())
	cmd.out.Infof("To run a job with a dataset, use: 'nerd job run'")
	return nil
}

// Description returns long-form help text
func (cmd *DatasetUpload) Description() string { return cmd.Synopsis() }

// Synopsis returns a one-line
func (cmd *DatasetUpload) Synopsis() string { return "Upload a dataset to your compute cluster." }

// Usage shows usage
func (cmd *DatasetUpload) Usage() string {
	return "nerd dataset upload [--name=] ~/my-project/my-input-1"
}
<|endoftext|>


<commit_before>// SILVER - Service Wrapper
//
// Copyright (c) 2014 PaperCut Software http://www.papercut.com/
// Use of this source code is governed by an MIT or GPL Version 2 license.
// See the project's LICENSE file for more information.
//
package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/kardianos/service"

	"github.com/robfig/cron"

	"github.com/papercutsoftware/silver/lib/logging"
	"github.com/papercutsoftware/silver/lib/osutils"
	"github.com/papercutsoftware/silver/lib/pathutils"
	"github.com/papercutsoftware/silver/service/cmdutil"
	"github.com/papercutsoftware/silver/service/config"
	"github.com/papercutsoftware/silver/service/svcutil"
)

const (
	defaultRefreshPoll = 10 * time.Second
)

type context struct {
	conf         *config.Config
	terminate    chan struct{}
	logger       *log.Logger
	runningGroup sync.WaitGroup
	cronManager  *cron.Cron
}

func main() {

	ctx := &context{}

	// Parse config (we don't action any errors quite yet)
	var err error
	ctx.conf, err = loadConf()
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: Invalid config - %v\n", err)
		os.Exit(1)
	}

	action, actionArgs, err := parse(os.Args)
	if err != nil {
		printUsage(ctx.conf.ServiceDescription.DisplayName, ctx.conf.ServiceDescription.Description)
		os.Exit(1)
	}

	if err := os.Chdir(exeFolder()); err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: Unable to set working directory: %v\n", err)
		os.Exit(1)
	}

	setupEnvironment(ctx.conf)

	switch action {
	case "command":
		execCommand(ctx, actionArgs)
	case "validate":
		fmt.Println("Config is valid")
		os.Exit(0)
	case "install":
		if err := writeProxyConf(); err != nil {
			fmt.Fprintf(os.Stderr, "WARNING: Unable to store HTTP Proxy settings: %v\n", err)
		}
		fallthrough
	default:
		osServiceControl(ctx)
	}
}

func osServiceControl(ctx *context) {
	// Setup log file out
	logFile := ctx.conf.ServiceConfig.LogFile
	maxSize := int64(ctx.conf.ServiceConfig.LogFileMaxSizeMb) * 1024 * 1024
	if logFile == "" {
		logFile = serviceName() + ".log"
	}
	ctx.logger = logging.NewFileLoggerWithMaxSize(logFile, maxSize)

	// Setup service
	svcConfig := &service.Config{
		Name:        serviceName(),
		DisplayName: ctx.conf.ServiceDescription.DisplayName,
		Description: ctx.conf.ServiceDescription.Description,
	}

	osService := &osService{ctx: ctx}
	svc, err := service.New(osService, svcConfig)
	if err != nil {
		fmt.Printf("ERROR: Invalid service config: %v\n", err)
		os.Exit(1)
	}

	if len(os.Args) > 1 && os.Args[1] != "run" {
		err = service.Control(svc, os.Args[1])
		if err != nil {
			fmt.Printf("ERROR: Invalid service command: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	err = svc.Run()
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
		os.Exit(1)
	}

	pidFile := ctx.conf.ServiceConfig.PidFile
	if pidFile != "" {
		ioutil.WriteFile(pidFile, []byte(fmt.Sprintf("%d\n", os.Getpid())), 0644)
	}
}

func printUsage(svcDisplayName, svcDesc string) {
	fmt.Printf("%s (%s)\n", svcDisplayName,
		serviceName())
	fmt.Printf("%s\n\n", svcDesc)
	fmt.Printf("Usage:\n")
	fmt.Printf("%s [install|uninstall|start|stop|command|validate|run|help] [command-name]\n", exeName())
	fmt.Printf("  install   - Install the service.\n")
	fmt.Printf("  uninstall - Remove/uninstall the service.\n")
	fmt.Printf("  start     - Start an installed service.\n")
	fmt.Printf("  stop      - Stop an installed service.\n")
	fmt.Printf("  validate  - Test the configuration file.\n")
	fmt.Printf("  run       - Run service on in command-line mode.\n")
	fmt.Printf("  command   - Run a command [command-name].\n")
	fmt.Printf("  help      - This usage message.\n")
}

func loadConf() (conf *config.Config, err error) {
	// FIXME: Not Get this function out of utils.
	confPath := getConfigFilePath()
	vars := config.ReplacementVars{
		ServiceName: serviceName(),
		ServiceRoot: exeFolder(),
	}
	conf, err = config.LoadConfig(confPath, vars)
	if err != nil {
		return nil, err
	}

	// Merge in any include files
	for _, include := range conf.Include {
		conf, err = config.MergeInclude(*conf, include, vars)
		if err != nil {
			return nil, err
		}
	}
	return conf, err
}

func setupEnvironment(conf *config.Config) {
	// Load Silver spacific
	os.Setenv("SILVER_SERVICE_NAME", conf.ServiceDescription.Name)
	os.Setenv("SILVER_SERVICE_ROOT", exeFolder())
	os.Setenv("SILVER_SERVICE_PID", string(os.Getpid()))

	// If we have HTTP proxy conf, load this
	if b, err := ioutil.ReadFile(proxyConfFile()); err == nil {
		proxy := strings.TrimSpace(string(b))
		os.Setenv("SILVER_HTTP_PROXY", proxy)
	}

	// Load any configured env
	for k, v := range conf.EnvironmentVars {
		os.Setenv(k, v)
	}
}

func writeProxyConf() error {
	proxy, err := osutils.GetHTTPProxy()
	if err != nil || proxy == "" {
		return nil
	}
	return ioutil.WriteFile(proxyConfFile(), []byte(proxy+"\n"), 0644)
}

func proxyConfFile() string {
	return filepath.Join(exeFolder(), "http-proxy.conf")
}

func execCommand(ctx *context, args []string) {
	/*
	*  IMPORTANT:
	*  Don't write to any log files, etc.  Commands are not system service code.
	*  Commands should be thought of as a "symlink" style, and run under a different
	*  user context.
	*
	*  args format: 1st element is the command. Any extras are appended to the command.
	 */
	if len(ctx.conf.Commands) == 0 {
		fmt.Fprintf(os.Stderr, "There are no commands configured!\n")
		os.Exit(1)
	}

	var cmd *config.Command
	if len(args) > 0 {
		cmdName := args[0]
		for _, c := range ctx.conf.Commands {
			if c.Name == cmdName {
				cmd = &c
				break
			}
		}
	}

	if cmd == nil {
		// Print command usage
		fmt.Fprintf(os.Stderr, "Valid commands are:\n")
		for _, command := range ctx.conf.Commands {
			fmt.Fprintf(os.Stderr, "    %s\n", command.Name)
		}
		os.Exit(1)
	}

	cmdConf := cmdutil.CommandConfig{}
	cmdConf.Path = pathutils.FindLastFile(cmd.Path)
	// Append any extra commands
	cmdConf.Args = append(cmd.Args, args[1:]...)
	// FIXME: Maybe unit conversion should be in the config layer?
	cmdConf.ExecTimeout = (time.Second * time.Duration(cmd.TimeoutSecs))

	exitCode, err := cmdutil.Execute(cmdConf)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
	}
	os.Exit(exitCode)
}

type osService struct {
	ctx *context
}

func (o *osService) Start(s service.Service) error {
	msg := fmt.Sprintf("Service '%s' started.", serviceName())
	o.ctx.logger.Printf(msg)
	sysLogger, err := s.Logger(nil)
	if err != nil {
		sysLogger.Info(msg)
	}

	doStart(o.ctx)
	go watchForReload(o.ctx)

	return nil
}

func doStart(ctx *context) {
	sf := stopFileName(ctx)
	if sf != "" {
		os.Remove(sf)
	}
	ctx.terminate = make(chan struct{})
	execStartupTasks(ctx)
	setupScheduledTasks(ctx)
	startServices(ctx)
}

func (o *osService) Stop(s service.Service) error {
	o.ctx.logger.Printf(fmt.Sprintf("Stopping '%s' service...", serviceName()))

	doStop(o.ctx)

	pidFile := o.ctx.conf.ServiceConfig.PidFile
	if pidFile != "" {
		os.Remove(pidFile)
	}

	msg := fmt.Sprintf("Stopped '%s' service.", serviceName())
	o.ctx.logger.Printf(msg)

	sysLogger, err := s.Logger(nil)
	if err != nil {
		sysLogger.Info(msg)
	}
	return nil
}

func stopFileName(ctx *context) string {
	stopFile := ctx.conf.ServiceConfig.StopFile
	if stopFile != "disabled" {
		return ""
	}
	return stopFile
}

func doStop(ctx *context) {
	// Create stop file... another method to signal services to stop.
	sf := stopFileName(ctx)
	if sf != "" {
		ioutil.WriteFile(sf, nil, 0644)
		defer os.Remove(sf)
	}
	if ctx.cronManager != nil {
		ctx.cronManager.Stop()
		ctx.cronManager = nil
	}
	if ctx.terminate != nil {
		close(ctx.terminate)
	}
	ctx.runningGroup.Wait()
}

func watchForReload(ctx *context) {
	f := ctx.conf.ServiceConfig.ReloadFile
	for {
		// FIXME: File system notification rather than polling?
		time.Sleep(defaultRefreshPoll)
		if _, err := os.Stat(f); err == nil {
			if err := os.Remove(f); err == nil {
				ctx.logger.Printf("Reload requested. Services will now restart.")
				doStop(ctx)
				time.Sleep(time.Second)
				// Reload config
				ctx.conf, _ = loadConf()
				doStart(ctx)
			}
		}
	}
}

func execStartupTasks(ctx *context) {
	ctx.logger.Printf("Starting %d startup tasks.", len(ctx.conf.StartupTasks))
	for _, task := range ctx.conf.StartupTasks {
		runTask := func(task config.StartupTask) {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			taskName := path.Base(task.Path)
			taskConfig := createTaskConfig(ctx, task.Task)
			if exitCode, err := svcutil.ExecuteTask(ctx.terminate, taskConfig); err != nil {
				ctx.logger.Printf("ERROR: Startup task '%s' reported: %v", taskName, err)
			} else {
				ctx.logger.Printf("Startup task '%s' finished with exit code %d", taskName, exitCode)
			}
		}
		if task.Async {
			go runTask(task)
		} else {
			if task.StartupDelaySecs > 0 || task.StartupRandomDelaySecs > 0 {
				ctx.logger.Printf("WARNING: Only Async startup tasks should have startup delays.")
			}
			runTask(task)
		}
	}
}

func startServices(ctx *context) {
	ctx.logger.Printf("Starting %d services.", len(ctx.conf.Services))
	for _, service := range ctx.conf.Services {
		go func(service config.Service) {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			serviceName := path.Base(service.Path)
			svcConfig := svcutil.ServiceConfig{}
			svcConfig.Path = pathutils.FindLastFile(service.Path)
			svcConfig.Args = service.Args
			svcConfig.GracefulShutDown = time.Duration(service.GracefulShutdownTimeoutSecs) * time.Second
			svcConfig.StartupDelay = time.Duration(service.StartupDelaySecs) * time.Second
			svcConfig.Logger = ctx.logger
			svcConfig.CrashConfig = svcutil.CrashConfig{
				MaxCountPerHour: service.MaxCrashCountPerHour,
				RestartDelay:    time.Duration(service.RestartDelaySecs) * time.Second,
			}
			if service.MonitorPing != nil {
				svcConfig.MonitorConfig = svcutil.MonitorConfig{
					URL:                   service.MonitorPing.URL,
					StartupDelay:          time.Duration(service.MonitorPing.StartupDelaySecs) * time.Second,
					Interval:              time.Duration(service.MonitorPing.IntervalSecs) * time.Second,
					Timeout:               time.Duration(service.MonitorPing.TimeoutSecs) * time.Second,
					RestartOnFailureCount: service.MonitorPing.RestartOnFailureCount,
				}
			}
			if err := svcutil.ExecuteService(ctx.terminate, svcConfig); err != nil {
				ctx.logger.Printf("ERROR: Service '%s' reported: %v", serviceName, err)
			}
		}(service)
	}
}

func createTaskConfig(ctx *context, task config.Task) svcutil.TaskConfig {
	taskConfig := svcutil.TaskConfig{}
	taskConfig.Path = pathutils.FindLastFile(task.Path)
	taskConfig.Args = task.Args
	taskConfig.ExecTimeout = time.Duration(task.TimeoutSecs) * time.Second
	taskConfig.StartupDelay = time.Duration(task.StartupDelaySecs) * time.Second
	taskConfig.StartupRandomDelay = time.Duration(task.StartupRandomDelaySecs) * time.Second
	taskConfig.Logger = ctx.logger
	return taskConfig
}

func setupScheduledTasks(ctx *context) {
	ctx.logger.Printf("Setting up %d scheduled tasks.", len(ctx.conf.ScheduledTasks))
	ctx.cronManager = cron.New()
	for _, scheduledTask := range ctx.conf.ScheduledTasks {
		taskConfig := createTaskConfig(ctx, scheduledTask.Task)
		runTask := func() {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			taskName := path.Base(taskConfig.Path)
			ctx.logger.Printf("Running schedule task '%s'", taskName)
			if exitCode, err := svcutil.ExecuteTask(ctx.terminate, taskConfig); err != nil {
				ctx.logger.Printf("ERROR: Scheduled task '%s' reported: %v", taskName, err)
			} else {
				ctx.logger.Printf("The task '%s' finished with exit code %d", taskName, exitCode)
			}
		}
		err := ctx.cronManager.AddFunc(scheduledTask.Schedule, runTask)
		if err != nil {
			ctx.logger.Printf("Unable to schedule task '%s': %v", scheduledTask.Path, err)
		}
	}
	ctx.cronManager.Start()
}
<commit_msg>Support wildcard in include directive<commit_after>// SILVER - Service Wrapper
//
// Copyright (c) 2014 PaperCut Software http://www.papercut.com/
// Use of this source code is governed by an MIT or GPL Version 2 license.
// See the project's LICENSE file for more information.
//
package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/kardianos/service"

	"github.com/robfig/cron"

	"github.com/papercutsoftware/silver/lib/logging"
	"github.com/papercutsoftware/silver/lib/osutils"
	"github.com/papercutsoftware/silver/lib/pathutils"
	"github.com/papercutsoftware/silver/service/cmdutil"
	"github.com/papercutsoftware/silver/service/config"
	"github.com/papercutsoftware/silver/service/svcutil"
)

const (
	defaultRefreshPoll = 10 * time.Second
)

type context struct {
	conf         *config.Config
	terminate    chan struct{}
	logger       *log.Logger
	runningGroup sync.WaitGroup
	cronManager  *cron.Cron
}

func main() {

	if err := os.Chdir(exeFolder()); err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: Unable to set working directory: %v\n", err)
		os.Exit(1)
	}

	ctx := &context{}

	// Parse config (we don't action any errors quite yet)
	var err error
	ctx.conf, err = loadConf()
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: Invalid config - %v\n", err)
		os.Exit(1)
	}

	action, actionArgs, err := parse(os.Args)
	if err != nil {
		printUsage(ctx.conf.ServiceDescription.DisplayName, ctx.conf.ServiceDescription.Description)
		os.Exit(1)
	}

	setupEnvironment(ctx.conf)

	switch action {
	case "command":
		execCommand(ctx, actionArgs)
	case "validate":
		fmt.Println("Config is valid")
		os.Exit(0)
	case "install":
		if err := writeProxyConf(); err != nil {
			fmt.Fprintf(os.Stderr, "WARNING: Unable to store HTTP Proxy settings: %v\n", err)
		}
		fallthrough
	default:
		osServiceControl(ctx)
	}
}

func osServiceControl(ctx *context) {
	// Setup log file out
	logFile := ctx.conf.ServiceConfig.LogFile
	maxSize := int64(ctx.conf.ServiceConfig.LogFileMaxSizeMb) * 1024 * 1024
	if logFile == "" {
		logFile = serviceName() + ".log"
	}
	ctx.logger = logging.NewFileLoggerWithMaxSize(logFile, maxSize)

	// Setup service
	svcConfig := &service.Config{
		Name:        serviceName(),
		DisplayName: ctx.conf.ServiceDescription.DisplayName,
		Description: ctx.conf.ServiceDescription.Description,
	}

	osService := &osService{ctx: ctx}
	svc, err := service.New(osService, svcConfig)
	if err != nil {
		fmt.Printf("ERROR: Invalid service config: %v\n", err)
		os.Exit(1)
	}

	if len(os.Args) > 1 && os.Args[1] != "run" {
		err = service.Control(svc, os.Args[1])
		if err != nil {
			fmt.Printf("ERROR: Invalid service command: %v\n", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	err = svc.Run()
	if err != nil {
		fmt.Printf("ERROR: %v\n", err)
		os.Exit(1)
	}

	pidFile := ctx.conf.ServiceConfig.PidFile
	if pidFile != "" {
		ioutil.WriteFile(pidFile, []byte(fmt.Sprintf("%d\n", os.Getpid())), 0644)
	}
}

func printUsage(svcDisplayName, svcDesc string) {
	fmt.Printf("%s (%s)\n", svcDisplayName,
		serviceName())
	fmt.Printf("%s\n\n", svcDesc)
	fmt.Printf("Usage:\n")
	fmt.Printf("%s [install|uninstall|start|stop|command|validate|run|help] [command-name]\n", exeName())
	fmt.Printf("  install   - Install the service.\n")
	fmt.Printf("  uninstall - Remove/uninstall the service.\n")
	fmt.Printf("  start     - Start an installed service.\n")
	fmt.Printf("  stop      - Stop an installed service.\n")
	fmt.Printf("  validate  - Test the configuration file.\n")
	fmt.Printf("  run       - Run service on in command-line mode.\n")
	fmt.Printf("  command   - Run a command [command-name].\n")
	fmt.Printf("  help      - This usage message.\n")
}

func loadConf() (conf *config.Config, err error) {
	// FIXME: Not Get this function out of utils.
	confPath := getConfigFilePath()
	vars := config.ReplacementVars{
		ServiceName: serviceName(),
		ServiceRoot: exeFolder(),
	}
	conf, err = config.LoadConfig(confPath, vars)
	if err != nil {
		return nil, err
	}

	// Merge in any include files
	for _, include := range conf.Include {
		include = pathutils.FindLastFile(include)
		conf, err = config.MergeInclude(*conf, include, vars)
		if err != nil {
			return nil, err
		}
	}
	return conf, err
}

func setupEnvironment(conf *config.Config) {
	// Load Silver spacific
	os.Setenv("SILVER_SERVICE_NAME", conf.ServiceDescription.Name)
	os.Setenv("SILVER_SERVICE_ROOT", exeFolder())
	os.Setenv("SILVER_SERVICE_PID", string(os.Getpid()))

	// If we have HTTP proxy conf, load this
	if b, err := ioutil.ReadFile(proxyConfFile()); err == nil {
		proxy := strings.TrimSpace(string(b))
		os.Setenv("SILVER_HTTP_PROXY", proxy)
	}

	// Load any configured env
	for k, v := range conf.EnvironmentVars {
		os.Setenv(k, v)
	}
}

func writeProxyConf() error {
	proxy, err := osutils.GetHTTPProxy()
	if err != nil || proxy == "" {
		return nil
	}
	return ioutil.WriteFile(proxyConfFile(), []byte(proxy+"\n"), 0644)
}

func proxyConfFile() string {
	return filepath.Join(exeFolder(), "http-proxy.conf")
}

func execCommand(ctx *context, args []string) {
	/*
	*  IMPORTANT:
	*  Don't write to any log files, etc.  Commands are not system service code.
	*  Commands should be thought of as a "symlink" style, and run under a different
	*  user context.
	*
	*  args format: 1st element is the command. Any extras are appended to the command.
	 */
	if len(ctx.conf.Commands) == 0 {
		fmt.Fprintf(os.Stderr, "There are no commands configured!\n")
		os.Exit(1)
	}

	var cmd *config.Command
	if len(args) > 0 {
		cmdName := args[0]
		for _, c := range ctx.conf.Commands {
			if c.Name == cmdName {
				cmd = &c
				break
			}
		}
	}

	if cmd == nil {
		// Print command usage
		fmt.Fprintf(os.Stderr, "Valid commands are:\n")
		for _, command := range ctx.conf.Commands {
			fmt.Fprintf(os.Stderr, "    %s\n", command.Name)
		}
		os.Exit(1)
	}

	cmdConf := cmdutil.CommandConfig{}
	cmdConf.Path = pathutils.FindLastFile(cmd.Path)
	// Append any extra commands
	cmdConf.Args = append(cmd.Args, args[1:]...)
	// FIXME: Maybe unit conversion should be in the config layer?
	cmdConf.ExecTimeout = (time.Second * time.Duration(cmd.TimeoutSecs))

	exitCode, err := cmdutil.Execute(cmdConf)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
	}
	os.Exit(exitCode)
}

type osService struct {
	ctx *context
}

func (o *osService) Start(s service.Service) error {
	msg := fmt.Sprintf("Service '%s' started.", serviceName())
	o.ctx.logger.Printf(msg)
	sysLogger, err := s.Logger(nil)
	if err != nil {
		sysLogger.Info(msg)
	}

	doStart(o.ctx)
	go watchForReload(o.ctx)

	return nil
}

func doStart(ctx *context) {
	sf := stopFileName(ctx)
	if sf != "" {
		os.Remove(sf)
	}
	ctx.terminate = make(chan struct{})
	execStartupTasks(ctx)
	setupScheduledTasks(ctx)
	startServices(ctx)
}

func (o *osService) Stop(s service.Service) error {
	o.ctx.logger.Printf(fmt.Sprintf("Stopping '%s' service...", serviceName()))

	doStop(o.ctx)

	pidFile := o.ctx.conf.ServiceConfig.PidFile
	if pidFile != "" {
		os.Remove(pidFile)
	}

	msg := fmt.Sprintf("Stopped '%s' service.", serviceName())
	o.ctx.logger.Printf(msg)

	sysLogger, err := s.Logger(nil)
	if err != nil {
		sysLogger.Info(msg)
	}
	return nil
}

func stopFileName(ctx *context) string {
	stopFile := ctx.conf.ServiceConfig.StopFile
	if stopFile != "disabled" {
		return ""
	}
	return stopFile
}

func doStop(ctx *context) {
	// Create stop file... another method to signal services to stop.
	sf := stopFileName(ctx)
	if sf != "" {
		ioutil.WriteFile(sf, nil, 0644)
		defer os.Remove(sf)
	}
	if ctx.cronManager != nil {
		ctx.cronManager.Stop()
		ctx.cronManager = nil
	}
	if ctx.terminate != nil {
		close(ctx.terminate)
	}
	ctx.runningGroup.Wait()
}

func watchForReload(ctx *context) {
	f := ctx.conf.ServiceConfig.ReloadFile
	for {
		// FIXME: File system notification rather than polling?
		time.Sleep(defaultRefreshPoll)
		if _, err := os.Stat(f); err == nil {
			if err := os.Remove(f); err == nil {
				ctx.logger.Printf("Reload requested. Services will now restart.")
				doStop(ctx)
				time.Sleep(time.Second)
				// Reload config
				ctx.conf, _ = loadConf()
				doStart(ctx)
			}
		}
	}
}

func execStartupTasks(ctx *context) {
	ctx.logger.Printf("Starting %d startup tasks.", len(ctx.conf.StartupTasks))
	for _, task := range ctx.conf.StartupTasks {
		runTask := func(task config.StartupTask) {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			taskName := path.Base(task.Path)
			taskConfig := createTaskConfig(ctx, task.Task)
			if exitCode, err := svcutil.ExecuteTask(ctx.terminate, taskConfig); err != nil {
				ctx.logger.Printf("ERROR: Startup task '%s' reported: %v", taskName, err)
			} else {
				ctx.logger.Printf("Startup task '%s' finished with exit code %d", taskName, exitCode)
			}
		}
		if task.Async {
			go runTask(task)
		} else {
			if task.StartupDelaySecs > 0 || task.StartupRandomDelaySecs > 0 {
				ctx.logger.Printf("WARNING: Only Async startup tasks should have startup delays.")
			}
			runTask(task)
		}
	}
}

func startServices(ctx *context) {
	ctx.logger.Printf("Starting %d services.", len(ctx.conf.Services))
	for _, service := range ctx.conf.Services {
		go func(service config.Service) {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			serviceName := path.Base(service.Path)
			svcConfig := svcutil.ServiceConfig{}
			svcConfig.Path = pathutils.FindLastFile(service.Path)
			svcConfig.Args = service.Args
			svcConfig.GracefulShutDown = time.Duration(service.GracefulShutdownTimeoutSecs) * time.Second
			svcConfig.StartupDelay = time.Duration(service.StartupDelaySecs) * time.Second
			svcConfig.Logger = ctx.logger
			svcConfig.CrashConfig = svcutil.CrashConfig{
				MaxCountPerHour: service.MaxCrashCountPerHour,
				RestartDelay:    time.Duration(service.RestartDelaySecs) * time.Second,
			}
			if service.MonitorPing != nil {
				svcConfig.MonitorConfig = svcutil.MonitorConfig{
					URL:                   service.MonitorPing.URL,
					StartupDelay:          time.Duration(service.MonitorPing.StartupDelaySecs) * time.Second,
					Interval:              time.Duration(service.MonitorPing.IntervalSecs) * time.Second,
					Timeout:               time.Duration(service.MonitorPing.TimeoutSecs) * time.Second,
					RestartOnFailureCount: service.MonitorPing.RestartOnFailureCount,
				}
			}
			if err := svcutil.ExecuteService(ctx.terminate, svcConfig); err != nil {
				ctx.logger.Printf("ERROR: Service '%s' reported: %v", serviceName, err)
			}
		}(service)
	}
}

func createTaskConfig(ctx *context, task config.Task) svcutil.TaskConfig {
	taskConfig := svcutil.TaskConfig{}
	taskConfig.Path = pathutils.FindLastFile(task.Path)
	taskConfig.Args = task.Args
	taskConfig.ExecTimeout = time.Duration(task.TimeoutSecs) * time.Second
	taskConfig.StartupDelay = time.Duration(task.StartupDelaySecs) * time.Second
	taskConfig.StartupRandomDelay = time.Duration(task.StartupRandomDelaySecs) * time.Second
	taskConfig.Logger = ctx.logger
	return taskConfig
}

func setupScheduledTasks(ctx *context) {
	ctx.logger.Printf("Setting up %d scheduled tasks.", len(ctx.conf.ScheduledTasks))
	ctx.cronManager = cron.New()
	for _, scheduledTask := range ctx.conf.ScheduledTasks {
		taskConfig := createTaskConfig(ctx, scheduledTask.Task)
		runTask := func() {
			ctx.runningGroup.Add(1)
			defer ctx.runningGroup.Done()
			taskName := path.Base(taskConfig.Path)
			ctx.logger.Printf("Running schedule task '%s'", taskName)
			if exitCode, err := svcutil.ExecuteTask(ctx.terminate, taskConfig); err != nil {
				ctx.logger.Printf("ERROR: Scheduled task '%s' reported: %v", taskName, err)
			} else {
				ctx.logger.Printf("The task '%s' finished with exit code %d", taskName, exitCode)
			}
		}
		err := ctx.cronManager.AddFunc(scheduledTask.Schedule, runTask)
		if err != nil {
			ctx.logger.Printf("Unable to schedule task '%s': %v", scheduledTask.Path, err)
		}
	}
	ctx.cronManager.Start()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strconv"

	pluginmanager "github.com/docker/cli/cli-plugins/manager"
	"github.com/docker/cli/cli/command"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

const (
	builderDefaultPlugin = "buildx"
	buildxMissingWarning = `DEPRECATED: The legacy builder is deprecated and will be removed in a future release.
            Install the buildx component to build images with BuildKit:
            https://docs.docker.com/go/buildx/
`

	buildxMissingError = `ERROR: BuildKit is enabled but the buildx component is missing or broken.
       Install the buildx component to build images with BuildKit:
       https://docs.docker.com/go/buildx/
`
)

func newBuilderError(warn bool, err error) error {
	var errorMsg string
	if warn {
		errorMsg = buildxMissingWarning
	} else {
		errorMsg = buildxMissingError
	}
	if pluginmanager.IsNotFound(err) {
		return errors.New(errorMsg)
	}
	if err != nil {
		return fmt.Errorf("%w\n\n%s", err, errorMsg)
	}
	return fmt.Errorf("%s", errorMsg)
}

func processBuilder(dockerCli command.Cli, cmd *cobra.Command, args, osargs []string) ([]string, []string, error) {
	var useLegacy bool
	var useBuilder bool

	// check DOCKER_BUILDKIT env var is present and
	// if not assume we want to use the builder component
	if v, ok := os.LookupEnv("DOCKER_BUILDKIT"); ok {
		enabled, err := strconv.ParseBool(v)
		if err != nil {
			return args, osargs, errors.Wrap(err, "DOCKER_BUILDKIT environment variable expects boolean value")
		}
		if !enabled {
			useLegacy = true
		} else {
			useBuilder = true
		}
	}

	// if a builder alias is defined, use it instead
	// of the default one
	builderAlias := builderDefaultPlugin
	aliasMap := dockerCli.ConfigFile().Aliases
	if v, ok := aliasMap[keyBuilderAlias]; ok {
		useBuilder = true
		builderAlias = v
	}

	// is this a build that should be forwarded to the builder?
	fwargs, fwosargs, forwarded := forwardBuilder(builderAlias, args, osargs)
	if !forwarded {
		return args, osargs, nil
	}

	if useLegacy {
		// display warning if not wcow and continue
		if dockerCli.ServerInfo().OSType != "windows" {
			_, _ = fmt.Fprintln(dockerCli.Err(), newBuilderError(true, nil))
		}
		return args, osargs, nil
	}

	// check plugin is available if cmd forwarded
	plugin, perr := pluginmanager.GetPlugin(builderAlias, dockerCli, cmd.Root())
	if perr == nil && plugin != nil {
		perr = plugin.Err
	}
	if perr != nil {
		// if builder enforced with DOCKER_BUILDKIT=1, cmd must fail if plugin missing or broken
		if useBuilder {
			return args, osargs, newBuilderError(false, perr)
		}
		// otherwise, display warning and continue
		_, _ = fmt.Fprintln(dockerCli.Err(), newBuilderError(true, perr))
		return args, osargs, nil
	}

	return fwargs, fwosargs, nil
}

func forwardBuilder(alias string, args, osargs []string) ([]string, []string, bool) {
	aliases := [][2][]string{
		{
			{"builder"},
			{alias},
		},
		{
			{"build"},
			{alias, "build"},
		},
		{
			{"image", "build"},
			{alias, "build"},
		},
	}
	for _, al := range aliases {
		if fwargs, changed := command.StringSliceReplaceAt(args, al[0], al[1], 0); changed {
			fwosargs, _ := command.StringSliceReplaceAt(osargs, al[0], al[1], -1)
			return fwargs, fwosargs, true
		}
	}
	return args, osargs, false
}
<commit_msg>build: use legacy builder for wcow if not opt-in with a builder component<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"

	pluginmanager "github.com/docker/cli/cli-plugins/manager"
	"github.com/docker/cli/cli/command"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

const (
	builderDefaultPlugin = "buildx"
	buildxMissingWarning = `DEPRECATED: The legacy builder is deprecated and will be removed in a future release.
            Install the buildx component to build images with BuildKit:
            https://docs.docker.com/go/buildx/
`

	buildxMissingError = `ERROR: BuildKit is enabled but the buildx component is missing or broken.
       Install the buildx component to build images with BuildKit:
       https://docs.docker.com/go/buildx/
`
)

func newBuilderError(warn bool, err error) error {
	var errorMsg string
	if warn {
		errorMsg = buildxMissingWarning
	} else {
		errorMsg = buildxMissingError
	}
	if pluginmanager.IsNotFound(err) {
		return errors.New(errorMsg)
	}
	if err != nil {
		return fmt.Errorf("%w\n\n%s", err, errorMsg)
	}
	return fmt.Errorf("%s", errorMsg)
}

func processBuilder(dockerCli command.Cli, cmd *cobra.Command, args, osargs []string) ([]string, []string, error) {
	var useLegacy, useBuilder bool

	// check DOCKER_BUILDKIT env var is present and
	// if not assume we want to use the builder component
	if v, ok := os.LookupEnv("DOCKER_BUILDKIT"); ok {
		enabled, err := strconv.ParseBool(v)
		if err != nil {
			return args, osargs, errors.Wrap(err, "DOCKER_BUILDKIT environment variable expects boolean value")
		}
		if !enabled {
			useLegacy = true
		} else {
			useBuilder = true
		}
	}

	// if a builder alias is defined, use it instead
	// of the default one
	builderAlias := builderDefaultPlugin
	aliasMap := dockerCli.ConfigFile().Aliases
	if v, ok := aliasMap[keyBuilderAlias]; ok {
		useBuilder = true
		builderAlias = v
	}

	// is this a build that should be forwarded to the builder?
	fwargs, fwosargs, forwarded := forwardBuilder(builderAlias, args, osargs)
	if !forwarded {
		return args, osargs, nil
	}

	// wcow build command must use the legacy builder
	// if not opt-in through a builder component
	if !useBuilder && dockerCli.ServerInfo().OSType == "windows" {
		return args, osargs, nil
	}

	if useLegacy {
		// display warning if not wcow and continue
		if dockerCli.ServerInfo().OSType != "windows" {
			_, _ = fmt.Fprintln(dockerCli.Err(), newBuilderError(true, nil))
		}
		return args, osargs, nil
	}

	// check plugin is available if cmd forwarded
	plugin, perr := pluginmanager.GetPlugin(builderAlias, dockerCli, cmd.Root())
	if perr == nil && plugin != nil {
		perr = plugin.Err
	}
	if perr != nil {
		// if builder enforced with DOCKER_BUILDKIT=1, cmd must fail if plugin missing or broken
		if useBuilder {
			return args, osargs, newBuilderError(false, perr)
		}
		// otherwise, display warning and continue
		_, _ = fmt.Fprintln(dockerCli.Err(), newBuilderError(true, perr))
		return args, osargs, nil
	}

	return fwargs, fwosargs, nil
}

func forwardBuilder(alias string, args, osargs []string) ([]string, []string, bool) {
	aliases := [][2][]string{
		{
			{"builder"},
			{alias},
		},
		{
			{"build"},
			{alias, "build"},
		},
		{
			{"image", "build"},
			{alias, "build"},
		},
	}
	for _, al := range aliases {
		if fwargs, changed := command.StringSliceReplaceAt(args, al[0], al[1], 0); changed {
			fwosargs, _ := command.StringSliceReplaceAt(osargs, al[0], al[1], -1)
			return fwargs, fwosargs, true
		}
	}
	return args, osargs, false
}
<|endoftext|>


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

package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/fiorix/freegeoip"
	"github.com/fiorix/go-redis/redis"
	"github.com/gorilla/context"
)

var maxmindFile = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz"

func main() {
	addr := flag.String("addr", ":8080", "Address in form of ip:port to listen on")
	certFile := flag.String("cert", "", "X.509 certificate file")
	keyFile := flag.String("key", "", "X.509 key file")
	public := flag.String("public", "", "Public directory to serve at the / endpoint")
	ipdb := flag.String("db", maxmindFile, "IP database file or URL")
	updateIntvl := flag.Duration("update", 24*time.Hour, "Database update check interval")
	retryIntvl := flag.Duration("retry", time.Hour, "Max time to wait before retrying update")
	useXFF := flag.Bool("use-x-forwarded-for", false, "Use the X-Forwarded-For header when available")
	silent := flag.Bool("silent", false, "Do not log requests to stderr")
	redisAddr := flag.String("redis", "127.0.0.1:6379", "Redis address in form of ip:port for quota")
	quotaMax := flag.Int("quota-max", 0, "Max requests per source IP per interval; Set 0 to turn off")
	quotaIntvl := flag.Duration("quota-interval", time.Hour, "Quota expiration interval")
	flag.Parse()

	rc, err := redis.Dial(*redisAddr)
	if err != nil {
		log.Fatal(err)
	}

	db, err := openDB(*ipdb, *updateIntvl, *retryIntvl)
	if err != nil {
		log.Fatal(err)
	}

	runtime.GOMAXPROCS(runtime.NumCPU())

	encoders := map[string]http.Handler{
		"/csv/":  freegeoip.NewHandler(db, &freegeoip.CSVEncoder{UseCRLF: true}),
		"/xml/":  freegeoip.NewHandler(db, &freegeoip.XMLEncoder{Indent: true}),
		"/json/": freegeoip.NewHandler(db, &freegeoip.JSONEncoder{}),
	}

	if *quotaMax > 0 {
		seconds := int((*quotaIntvl).Seconds())
		for path, f := range encoders {
			encoders[path] = userQuota(rc, *quotaMax, seconds, f)
		}
	}

	mux := http.NewServeMux()
	for path, handler := range encoders {
		mux.Handle(path, handler)
	}

	if len(*public) > 0 {
		mux.Handle("/", http.FileServer(http.Dir(*public)))
	}

	handler := CORS(mux, "GET", "HEAD")

	if !*silent {
		log.Println("freegeoip server starting on", *addr)
		go logEvents(db)
		handler = logHandler(handler)
	}

	if *useXFF {
		handler = freegeoip.ProxyHandler(handler)
	}

	if len(*certFile) > 0 && len(*keyFile) > 0 {
		err = http.ListenAndServeTLS(*addr, *certFile, *keyFile, handler)
	} else {
		err = http.ListenAndServe(*addr, handler)
	}
	if err != nil {
		log.Fatal(err)
	}
}

// openDB opens and returns the IP database.
func openDB(dsn string, updateIntvl, maxRetryIntvl time.Duration) (db *freegeoip.DB, err error) {
	u, err := url.Parse(dsn)
	if err != nil || len(u.Scheme) == 0 {
		db, err = freegeoip.Open(dsn)
	} else {
		db, err = freegeoip.OpenURL(dsn, updateIntvl, maxRetryIntvl)
	}
	return
}

// CORS is an http handler that checks for allowed request methods (verbs)
// and adds CORS headers to all http responses.
//
// See http://en.wikipedia.org/wiki/Cross-origin_resource_sharing for details.
func CORS(f http.Handler, allow ...string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Method",
			strings.Join(allow, ", ")+", OPTIONS")
		if r.Method == "OPTIONS" {
			w.WriteHeader(200)
			return
		}
		for _, method := range allow {
			if r.Method == method {
				f.ServeHTTP(w, r)
				return
			}
		}
		w.Header().Set("Allow", strings.Join(allow, ", ")+", OPTIONS")
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed),
			http.StatusMethodNotAllowed)
	})
}

// userQuota is a handler that provides a rate limiter to the freegeoip API.
// It allows qmax requests per qintvl, in seconds.
//
// If redis is not available it responds with service unavailable.
func userQuota(rc *redis.Client, qmax int, qintvl int, f http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var ip string
		if idx := strings.LastIndex(r.RemoteAddr, ":"); idx != -1 {
			ip = r.RemoteAddr[:idx]
		} else {
			ip = r.RemoteAddr
		}
		sreq, err := rc.Get(ip)
		if err != nil {
			serviceUnavailable(w, r, err.Error())
			return
		}
		if len(sreq) == 0 {
			err = rc.SetEx(ip, qintvl, "1")
			if err != nil {
				serviceUnavailable(w, r, err.Error())
				return
			}
			f.ServeHTTP(w, r)
			return
		}
		nreq, _ := strconv.Atoi(sreq)
		if nreq >= qmax {
			http.Error(w, "Quota exceeded", http.StatusForbidden)
			return
		}
		_, err = rc.Incr(ip)
		if err != nil {
			context.Set(r, "log", err.Error())
		}
		f.ServeHTTP(w, r)
	})
}

// serviceUnavailable writes an http error 501 to a client.
func serviceUnavailable(w http.ResponseWriter, r *http.Request, log string) {
	context.Set(r, "log", log)
	http.Error(w, "Try again later", http.StatusServiceUnavailable)
}

// logEvents logs database events.
func logEvents(db *freegeoip.DB) {
	for {
		select {
		case file := <-db.NotifyOpen():
			log.Println("database loaded:", file)
		case err := <-db.NotifyError():
			log.Println("database error:", err)
		case <-db.NotifyClose():
			return
		}
	}
}

// logHandler logs http requests.
func logHandler(f http.Handler) http.Handler {
	empty := ""
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp := responseWriter{w, http.StatusOK, 0}
		start := time.Now()
		f.ServeHTTP(&resp, r)
		elapsed := time.Since(start)
		extra := context.Get(r, "log")
		if extra != nil {
			defer context.Clear(r)
		} else {
			extra = empty
		}
		log.Printf("%q %d %q %q %s %q %db in %s %q",
			r.Proto,
			resp.status,
			r.Method,
			r.URL.Path,
			remoteIP(r),
			r.Header.Get("User-Agent"),
			resp.bytes,
			elapsed,
			extra,
		)
	})
}

// remoteIP returns the client's address without the port number.
func remoteIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// responseWriter is an http.ResponseWriter that records the returned
// status and bytes written to the client.
type responseWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

// Write implements the http.ResponseWriter interface.
func (f *responseWriter) Write(b []byte) (int, error) {
	n, err := f.ResponseWriter.Write(b)
	if err != nil {
		return 0, err
	}
	f.bytes += n
	return n, nil
}

// WriteHeader implements the http.ResponseWriter interface.
func (f *responseWriter) WriteHeader(code int) {
	f.status = code
	f.ResponseWriter.WriteHeader(code)
}
<commit_msg>Set v3.0.2 and add -version command line flag<commit_after>// Copyright 2009-2014 The freegeoip authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/fiorix/freegeoip"
	"github.com/fiorix/go-redis/redis"
	"github.com/gorilla/context"
)

var VERSION = "3.0.2"
var maxmindFile = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.mmdb.gz"

func main() {
	addr := flag.String("addr", ":8080", "Address in form of ip:port to listen on")
	certFile := flag.String("cert", "", "X.509 certificate file")
	keyFile := flag.String("key", "", "X.509 key file")
	public := flag.String("public", "", "Public directory to serve at the / endpoint")
	ipdb := flag.String("db", maxmindFile, "IP database file or URL")
	updateIntvl := flag.Duration("update", 24*time.Hour, "Database update check interval")
	retryIntvl := flag.Duration("retry", time.Hour, "Max time to wait before retrying update")
	useXFF := flag.Bool("use-x-forwarded-for", false, "Use the X-Forwarded-For header when available")
	silent := flag.Bool("silent", false, "Do not log requests to stderr")
	redisAddr := flag.String("redis", "127.0.0.1:6379", "Redis address in form of ip:port for quota")
	quotaMax := flag.Int("quota-max", 0, "Max requests per source IP per interval; Set 0 to turn off")
	quotaIntvl := flag.Duration("quota-interval", time.Hour, "Quota expiration interval")
	version := flag.Bool("version", false, "Show version and exit")
	flag.Parse()

	if *version {
		fmt.Printf("freegeoip v%s\n", VERSION)
		return
	}

	rc, err := redis.Dial(*redisAddr)
	if err != nil {
		log.Fatal(err)
	}

	db, err := openDB(*ipdb, *updateIntvl, *retryIntvl)
	if err != nil {
		log.Fatal(err)
	}

	runtime.GOMAXPROCS(runtime.NumCPU())

	encoders := map[string]http.Handler{
		"/csv/":  freegeoip.NewHandler(db, &freegeoip.CSVEncoder{UseCRLF: true}),
		"/xml/":  freegeoip.NewHandler(db, &freegeoip.XMLEncoder{Indent: true}),
		"/json/": freegeoip.NewHandler(db, &freegeoip.JSONEncoder{}),
	}

	if *quotaMax > 0 {
		seconds := int((*quotaIntvl).Seconds())
		for path, f := range encoders {
			encoders[path] = userQuota(rc, *quotaMax, seconds, f)
		}
	}

	mux := http.NewServeMux()
	for path, handler := range encoders {
		mux.Handle(path, handler)
	}

	if len(*public) > 0 {
		mux.Handle("/", http.FileServer(http.Dir(*public)))
	}

	handler := CORS(mux, "GET", "HEAD")

	if !*silent {
		log.Println("freegeoip server starting on", *addr)
		go logEvents(db)
		handler = logHandler(handler)
	}

	if *useXFF {
		handler = freegeoip.ProxyHandler(handler)
	}

	if len(*certFile) > 0 && len(*keyFile) > 0 {
		err = http.ListenAndServeTLS(*addr, *certFile, *keyFile, handler)
	} else {
		err = http.ListenAndServe(*addr, handler)
	}
	if err != nil {
		log.Fatal(err)
	}
}

// openDB opens and returns the IP database.
func openDB(dsn string, updateIntvl, maxRetryIntvl time.Duration) (db *freegeoip.DB, err error) {
	u, err := url.Parse(dsn)
	if err != nil || len(u.Scheme) == 0 {
		db, err = freegeoip.Open(dsn)
	} else {
		db, err = freegeoip.OpenURL(dsn, updateIntvl, maxRetryIntvl)
	}
	return
}

// CORS is an http handler that checks for allowed request methods (verbs)
// and adds CORS headers to all http responses.
//
// See http://en.wikipedia.org/wiki/Cross-origin_resource_sharing for details.
func CORS(f http.Handler, allow ...string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Method",
			strings.Join(allow, ", ")+", OPTIONS")
		if r.Method == "OPTIONS" {
			w.WriteHeader(200)
			return
		}
		for _, method := range allow {
			if r.Method == method {
				f.ServeHTTP(w, r)
				return
			}
		}
		w.Header().Set("Allow", strings.Join(allow, ", ")+", OPTIONS")
		http.Error(w, http.StatusText(http.StatusMethodNotAllowed),
			http.StatusMethodNotAllowed)
	})
}

// userQuota is a handler that provides a rate limiter to the freegeoip API.
// It allows qmax requests per qintvl, in seconds.
//
// If redis is not available it responds with service unavailable.
func userQuota(rc *redis.Client, qmax int, qintvl int, f http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var ip string
		if idx := strings.LastIndex(r.RemoteAddr, ":"); idx != -1 {
			ip = r.RemoteAddr[:idx]
		} else {
			ip = r.RemoteAddr
		}
		sreq, err := rc.Get(ip)
		if err != nil {
			serviceUnavailable(w, r, err.Error())
			return
		}
		if len(sreq) == 0 {
			err = rc.SetEx(ip, qintvl, "1")
			if err != nil {
				serviceUnavailable(w, r, err.Error())
				return
			}
			f.ServeHTTP(w, r)
			return
		}
		nreq, _ := strconv.Atoi(sreq)
		if nreq >= qmax {
			http.Error(w, "Quota exceeded", http.StatusForbidden)
			return
		}
		_, err = rc.Incr(ip)
		if err != nil {
			context.Set(r, "log", err.Error())
		}
		f.ServeHTTP(w, r)
	})
}

// serviceUnavailable writes an http error 501 to a client.
func serviceUnavailable(w http.ResponseWriter, r *http.Request, log string) {
	context.Set(r, "log", log)
	http.Error(w, "Try again later", http.StatusServiceUnavailable)
}

// logEvents logs database events.
func logEvents(db *freegeoip.DB) {
	for {
		select {
		case file := <-db.NotifyOpen():
			log.Println("database loaded:", file)
		case err := <-db.NotifyError():
			log.Println("database error:", err)
		case <-db.NotifyClose():
			return
		}
	}
}

// logHandler logs http requests.
func logHandler(f http.Handler) http.Handler {
	empty := ""
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp := responseWriter{w, http.StatusOK, 0}
		start := time.Now()
		f.ServeHTTP(&resp, r)
		elapsed := time.Since(start)
		extra := context.Get(r, "log")
		if extra != nil {
			defer context.Clear(r)
		} else {
			extra = empty
		}
		log.Printf("%q %d %q %q %s %q %db in %s %q",
			r.Proto,
			resp.status,
			r.Method,
			r.URL.Path,
			remoteIP(r),
			r.Header.Get("User-Agent"),
			resp.bytes,
			elapsed,
			extra,
		)
	})
}

// remoteIP returns the client's address without the port number.
func remoteIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// responseWriter is an http.ResponseWriter that records the returned
// status and bytes written to the client.
type responseWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

// Write implements the http.ResponseWriter interface.
func (f *responseWriter) Write(b []byte) (int, error) {
	n, err := f.ResponseWriter.Write(b)
	if err != nil {
		return 0, err
	}
	f.bytes += n
	return n, nil
}

// WriteHeader implements the http.ResponseWriter interface.
func (f *responseWriter) WriteHeader(code int) {
	f.status = code
	f.ResponseWriter.WriteHeader(code)
}
<|endoftext|>


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

package memory

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"

	ds "go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/datastore/serialize"
	"go.chromium.org/luci/common/data/cmpbin"
	"go.chromium.org/luci/common/data/stringset"
)

// MaxQueryComponents was lifted from a hard-coded constant in dev_appserver.
// No idea if it's a real limit or just a convenience in the current dev
// appserver implementation.
const MaxQueryComponents = 100

// MaxIndexColumns is the maximum number of index columns we're willing to
// support.
const MaxIndexColumns = 64

// A queryCursor is:
//   {#orders} ++ IndexColumn* ++ RawRowData
//   IndexColumn will always contain __key__ as the last column, and so #orders
//     must always be >= 1
type queryCursor []byte

func newCursor(s string) (ds.Cursor, error) {
	d, err := base64.URLEncoding.DecodeString(s)
	if err != nil {
		return nil, fmt.Errorf("Failed to Base64-decode cursor: %s", err)
	}
	c := queryCursor(d)
	if _, _, err := c.decode(); err != nil {
		return nil, err
	}
	return c, nil
}

func (q queryCursor) String() string { return base64.URLEncoding.EncodeToString([]byte(q)) }

// decode returns the encoded IndexColumns, the raw row (cursor) data, or an
// error.
func (q queryCursor) decode() ([]ds.IndexColumn, []byte, error) {
	buf := bytes.NewBuffer([]byte(q))
	count, _, err := cmpbin.ReadUint(buf)
	if err != nil {
		return nil, nil, fmt.Errorf("invalid cursor: bad prefix number")
	}

	if count == 0 || count > MaxIndexColumns {
		return nil, nil, fmt.Errorf("invalid cursor: bad column count %d", count)
	}

	if count == 0 {
		return nil, nil, fmt.Errorf("invalid cursor: zero prefix number")
	}

	cols := make([]ds.IndexColumn, count)
	for i := range cols {
		if cols[i], err = serialize.ReadIndexColumn(buf); err != nil {
			return nil, nil, fmt.Errorf("invalid cursor: unable to decode IndexColumn %d: %s", i, err)
		}
	}

	if cols[len(cols)-1].Property != "__key__" {
		return nil, nil, fmt.Errorf("invalid cursor: last column was not __key__: %v", cols[len(cols)-1])
	}

	return cols, buf.Bytes(), nil
}

func sortOrdersEqual(as, bs []ds.IndexColumn) bool {
	if len(as) != len(bs) {
		return false
	}
	for i, a := range as {
		if a != bs[i] {
			return false
		}
	}
	return true
}

func numComponents(fq *ds.FinalizedQuery) int {
	numComponents := len(fq.Orders())
	if p, _, _ := fq.IneqFilterLow(); p != "" {
		numComponents++
	}
	if p, _, _ := fq.IneqFilterHigh(); p != "" {
		numComponents++
	}
	for _, v := range fq.EqFilters() {
		numComponents += v.Len()
	}
	return numComponents
}

// GetBinaryBounds gets the binary encoding of the upper and lower bounds of
// the inequality filter on fq, if any is defined. If a bound does not exist,
// it is nil.
//
// NOTE: if fq specifies a descending sort order for the inequality, the bounds
// will be inverted, incremented, and flipped.
func GetBinaryBounds(fq *ds.FinalizedQuery) (lower, upper []byte) {
	// Pick up the start/end range from the inequalities, if any.
	//
	// start and end in the reducedQuery are normalized so that `start >=
	// X < end`. Because of that, we need to tweak the inequality filters
	// contained in the query if they use the > or <= operators.
	if ineqProp := fq.IneqFilterProp(); ineqProp != "" {
		_, startOp, startV := fq.IneqFilterLow()
		if startOp != "" {
			lower = serialize.ToBytes(startV)
			if startOp == ">" {
				lower = increment(lower)
			}
		}

		_, endOp, endV := fq.IneqFilterHigh()
		if endOp != "" {
			upper = serialize.ToBytes(endV)
			if endOp == "<=" {
				upper = increment(upper)
			}
		}

		// The inequality is specified in natural (ascending) order in the query's
		// Filter syntax, but the order information may indicate to use a descending
		// index column for it. If that's the case, then we must invert, swap and
		// increment the inequality endpoints.
		//
		// Invert so that the desired numbers are represented correctly in the index.
		// Swap so that our iterators still go from >= start to < end.
		// Increment so that >= and < get correctly bounded (since the iterator is
		// still using natrual bytes ordering)
		if fq.Orders()[0].Descending {
			hi, lo := []byte(nil), []byte(nil)
			if len(lower) > 0 {
				lo = increment(serialize.Invert(lower))
			}
			if len(upper) > 0 {
				hi = increment(serialize.Invert(upper))
			}
			upper, lower = lo, hi
		}
	}
	return
}

func reduce(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool) (*reducedQuery, error) {
	if err := fq.Valid(kc); err != nil {
		return nil, err
	}
	if isTxn && fq.Ancestor() == nil {
		return nil, fmt.Errorf("queries within a transaction must include an Ancestor filter")
	}
	if num := numComponents(fq); num > MaxQueryComponents {
		return nil, fmt.Errorf(
			"gae/memory: query is too large. may not have more than "+
				"%d filters + sort orders + ancestor total: had %d",
			MaxQueryComponents, num)
	}

	ret := &reducedQuery{
		kc:           kc,
		kind:         fq.Kind(),
		suffixFormat: fq.Orders(),
	}

	eqFilts := fq.EqFilters()
	ret.eqFilters = make(map[string]stringset.Set, len(eqFilts))
	for prop, vals := range eqFilts {
		sVals := stringset.New(len(vals))
		for _, v := range vals {
			sVals.Add(string(serialize.ToBytes(v)))
		}
		ret.eqFilters[prop] = sVals
	}

	startD, endD := GetBinaryBounds(fq)

	// Now we check the start and end cursors.
	//
	// Cursors are composed of a list of IndexColumns at the beginning, followed
	// by the raw bytes to use for the suffix. The cursor is only valid if all of
	// its IndexColumns match our proposed suffixFormat, as calculated above.
	//
	// Cursors are mutually exclusive with the start/end we picked up from the
	// inequality. In a well formed query, they indicate a subset of results
	// bounded by the inequality. Technically if the start cursor is not >= the
	// low bound, or the end cursor is < the high bound, it's an error, but for
	// simplicity we just cap to the narrowest intersection of the inequality and
	// cursors.
	ret.start = startD
	ret.end = endD
	if start, end := fq.Bounds(); start != nil || end != nil {
		if start != nil {
			if c, ok := start.(queryCursor); ok {
				startCols, startD, err := c.decode()
				if err != nil {
					return nil, err
				}

				if !sortOrdersEqual(startCols, ret.suffixFormat) {
					return nil, errors.New("gae/memory: start cursor is invalid for this query")
				}
				if ret.start == nil || bytes.Compare(ret.start, startD) < 0 {
					ret.start = startD
				}
			} else {
				return nil, errors.New("gae/memory: bad cursor type")
			}
		}

		if end != nil {
			if c, ok := end.(queryCursor); ok {
				endCols, endD, err := c.decode()
				if err != nil {
					return nil, err
				}

				if !sortOrdersEqual(endCols, ret.suffixFormat) {
					return nil, errors.New("gae/memory: end cursor is invalid for this query")
				}
				if ret.end == nil || bytes.Compare(endD, ret.end) < 0 {
					ret.end = endD
				}
			} else {
				return nil, errors.New("gae/memory: bad cursor type")
			}
		}
	}

	// Finally, verify that we could even /potentially/ do work. If we have
	// overlapping range ends, then we don't have anything to do.
	if ret.end != nil && bytes.Compare(ret.start, ret.end) >= 0 {
		return nil, ds.ErrNullQuery
	}

	ret.numCols = len(ret.suffixFormat)
	for prop, vals := range ret.eqFilters {
		if len(ret.suffixFormat) == 1 && prop == "__ancestor__" {
			continue
		}
		ret.numCols += vals.Len()
	}

	return ret, nil
}

func increment(bstr []byte) []byte {
	ret, overflow := serialize.Increment(bstr)
	if overflow {
		// This byte string was ALL 0xFF's. The only safe incrementation to do here
		// would be to add a new byte to the beginning of bstr with the value 0x01,
		// and a byte to the beginning OF ALL OTHER []byte's which bstr may be
		// compared with. This is obviously impossible to do here, so panic. If we
		// hit this, then we would need to add a spare 0 byte before every index
		// column.
		//
		// Another way to think about this is that we just accumulated a 'carry' bit,
		// and the new value has overflowed this representation.
		//
		// Fortunately, the first byte of a serialized index column entry is a
		// PropertyType byte, and the only valid values that we'll be incrementing
		// are never equal to 0xFF, since they have the high bit set (so either they're
		// 0x8*, or 0x7*, depending on if it's inverted).
		impossible(fmt.Errorf("incrementing %v would require more sigfigs", bstr))
	}
	return ret
}
<commit_msg>[impl/memory] Use raw (unpadded) base64 encoding for cursors.<commit_after>// Copyright 2015 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package memory

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"

	ds "go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/datastore/serialize"
	"go.chromium.org/luci/common/data/cmpbin"
	"go.chromium.org/luci/common/data/stringset"
)

// MaxQueryComponents was lifted from a hard-coded constant in dev_appserver.
// No idea if it's a real limit or just a convenience in the current dev
// appserver implementation.
const MaxQueryComponents = 100

// MaxIndexColumns is the maximum number of index columns we're willing to
// support.
const MaxIndexColumns = 64

// A queryCursor is:
//   {#orders} ++ IndexColumn* ++ RawRowData
//   IndexColumn will always contain __key__ as the last column, and so #orders
//     must always be >= 1
type queryCursor []byte

func newCursor(s string) (ds.Cursor, error) {
	d, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil {
		return nil, fmt.Errorf("failed to Base64-decode cursor: %s", err)
	}
	c := queryCursor(d)
	if _, _, err := c.decode(); err != nil {
		return nil, err
	}
	return c, nil
}

func (q queryCursor) String() string { return base64.RawURLEncoding.EncodeToString([]byte(q)) }

// decode returns the encoded IndexColumns, the raw row (cursor) data, or an
// error.
func (q queryCursor) decode() ([]ds.IndexColumn, []byte, error) {
	buf := bytes.NewBuffer([]byte(q))
	count, _, err := cmpbin.ReadUint(buf)
	if err != nil {
		return nil, nil, fmt.Errorf("invalid cursor: bad prefix number")
	}

	if count == 0 || count > MaxIndexColumns {
		return nil, nil, fmt.Errorf("invalid cursor: bad column count %d", count)
	}

	if count == 0 {
		return nil, nil, fmt.Errorf("invalid cursor: zero prefix number")
	}

	cols := make([]ds.IndexColumn, count)
	for i := range cols {
		if cols[i], err = serialize.ReadIndexColumn(buf); err != nil {
			return nil, nil, fmt.Errorf("invalid cursor: unable to decode IndexColumn %d: %s", i, err)
		}
	}

	if cols[len(cols)-1].Property != "__key__" {
		return nil, nil, fmt.Errorf("invalid cursor: last column was not __key__: %v", cols[len(cols)-1])
	}

	return cols, buf.Bytes(), nil
}

func sortOrdersEqual(as, bs []ds.IndexColumn) bool {
	if len(as) != len(bs) {
		return false
	}
	for i, a := range as {
		if a != bs[i] {
			return false
		}
	}
	return true
}

func numComponents(fq *ds.FinalizedQuery) int {
	numComponents := len(fq.Orders())
	if p, _, _ := fq.IneqFilterLow(); p != "" {
		numComponents++
	}
	if p, _, _ := fq.IneqFilterHigh(); p != "" {
		numComponents++
	}
	for _, v := range fq.EqFilters() {
		numComponents += v.Len()
	}
	return numComponents
}

// GetBinaryBounds gets the binary encoding of the upper and lower bounds of
// the inequality filter on fq, if any is defined. If a bound does not exist,
// it is nil.
//
// NOTE: if fq specifies a descending sort order for the inequality, the bounds
// will be inverted, incremented, and flipped.
func GetBinaryBounds(fq *ds.FinalizedQuery) (lower, upper []byte) {
	// Pick up the start/end range from the inequalities, if any.
	//
	// start and end in the reducedQuery are normalized so that `start >=
	// X < end`. Because of that, we need to tweak the inequality filters
	// contained in the query if they use the > or <= operators.
	if ineqProp := fq.IneqFilterProp(); ineqProp != "" {
		_, startOp, startV := fq.IneqFilterLow()
		if startOp != "" {
			lower = serialize.ToBytes(startV)
			if startOp == ">" {
				lower = increment(lower)
			}
		}

		_, endOp, endV := fq.IneqFilterHigh()
		if endOp != "" {
			upper = serialize.ToBytes(endV)
			if endOp == "<=" {
				upper = increment(upper)
			}
		}

		// The inequality is specified in natural (ascending) order in the query's
		// Filter syntax, but the order information may indicate to use a descending
		// index column for it. If that's the case, then we must invert, swap and
		// increment the inequality endpoints.
		//
		// Invert so that the desired numbers are represented correctly in the index.
		// Swap so that our iterators still go from >= start to < end.
		// Increment so that >= and < get correctly bounded (since the iterator is
		// still using natrual bytes ordering)
		if fq.Orders()[0].Descending {
			hi, lo := []byte(nil), []byte(nil)
			if len(lower) > 0 {
				lo = increment(serialize.Invert(lower))
			}
			if len(upper) > 0 {
				hi = increment(serialize.Invert(upper))
			}
			upper, lower = lo, hi
		}
	}
	return
}

func reduce(fq *ds.FinalizedQuery, kc ds.KeyContext, isTxn bool) (*reducedQuery, error) {
	if err := fq.Valid(kc); err != nil {
		return nil, err
	}
	if isTxn && fq.Ancestor() == nil {
		return nil, fmt.Errorf("queries within a transaction must include an Ancestor filter")
	}
	if num := numComponents(fq); num > MaxQueryComponents {
		return nil, fmt.Errorf(
			"gae/memory: query is too large. may not have more than "+
				"%d filters + sort orders + ancestor total: had %d",
			MaxQueryComponents, num)
	}

	ret := &reducedQuery{
		kc:           kc,
		kind:         fq.Kind(),
		suffixFormat: fq.Orders(),
	}

	eqFilts := fq.EqFilters()
	ret.eqFilters = make(map[string]stringset.Set, len(eqFilts))
	for prop, vals := range eqFilts {
		sVals := stringset.New(len(vals))
		for _, v := range vals {
			sVals.Add(string(serialize.ToBytes(v)))
		}
		ret.eqFilters[prop] = sVals
	}

	startD, endD := GetBinaryBounds(fq)

	// Now we check the start and end cursors.
	//
	// Cursors are composed of a list of IndexColumns at the beginning, followed
	// by the raw bytes to use for the suffix. The cursor is only valid if all of
	// its IndexColumns match our proposed suffixFormat, as calculated above.
	//
	// Cursors are mutually exclusive with the start/end we picked up from the
	// inequality. In a well formed query, they indicate a subset of results
	// bounded by the inequality. Technically if the start cursor is not >= the
	// low bound, or the end cursor is < the high bound, it's an error, but for
	// simplicity we just cap to the narrowest intersection of the inequality and
	// cursors.
	ret.start = startD
	ret.end = endD
	if start, end := fq.Bounds(); start != nil || end != nil {
		if start != nil {
			if c, ok := start.(queryCursor); ok {
				startCols, startD, err := c.decode()
				if err != nil {
					return nil, err
				}

				if !sortOrdersEqual(startCols, ret.suffixFormat) {
					return nil, errors.New("gae/memory: start cursor is invalid for this query")
				}
				if ret.start == nil || bytes.Compare(ret.start, startD) < 0 {
					ret.start = startD
				}
			} else {
				return nil, errors.New("gae/memory: bad cursor type")
			}
		}

		if end != nil {
			if c, ok := end.(queryCursor); ok {
				endCols, endD, err := c.decode()
				if err != nil {
					return nil, err
				}

				if !sortOrdersEqual(endCols, ret.suffixFormat) {
					return nil, errors.New("gae/memory: end cursor is invalid for this query")
				}
				if ret.end == nil || bytes.Compare(endD, ret.end) < 0 {
					ret.end = endD
				}
			} else {
				return nil, errors.New("gae/memory: bad cursor type")
			}
		}
	}

	// Finally, verify that we could even /potentially/ do work. If we have
	// overlapping range ends, then we don't have anything to do.
	if ret.end != nil && bytes.Compare(ret.start, ret.end) >= 0 {
		return nil, ds.ErrNullQuery
	}

	ret.numCols = len(ret.suffixFormat)
	for prop, vals := range ret.eqFilters {
		if len(ret.suffixFormat) == 1 && prop == "__ancestor__" {
			continue
		}
		ret.numCols += vals.Len()
	}

	return ret, nil
}

func increment(bstr []byte) []byte {
	ret, overflow := serialize.Increment(bstr)
	if overflow {
		// This byte string was ALL 0xFF's. The only safe incrementation to do here
		// would be to add a new byte to the beginning of bstr with the value 0x01,
		// and a byte to the beginning OF ALL OTHER []byte's which bstr may be
		// compared with. This is obviously impossible to do here, so panic. If we
		// hit this, then we would need to add a spare 0 byte before every index
		// column.
		//
		// Another way to think about this is that we just accumulated a 'carry' bit,
		// and the new value has overflowed this representation.
		//
		// Fortunately, the first byte of a serialized index column entry is a
		// PropertyType byte, and the only valid values that we'll be incrementing
		// are never equal to 0xFF, since they have the high bit set (so either they're
		// 0x8*, or 0x7*, depending on if it's inverted).
		impossible(fmt.Errorf("incrementing %v would require more sigfigs", bstr))
	}
	return ret
}
<|endoftext|>