<commit_before>package vitali

import (
    "os"
    "io"
    "log"
    "fmt"
    "strconv"
    "net/http"
    "net/http/httputil"
    "html/template"
    "time"
    "regexp"
    "strings"
    "reflect"
)

type RouteRule struct {
    Pattern string
    Resource interface{}
}

type PatternMapping struct {
    Re *regexp.Regexp
    Names []string
}

type webApp struct {
    RouteRules []RouteRule
    PatternMappings []PatternMapping
    UserProvider UserProvider
    Settings map[string]string
    DumpRequest bool
    ErrTemplate *template.Template
}

func checkPermission(perm reflect.StructTag, method Method, role string) bool {
    requiredRole := perm.Get(string(method))
    if requiredRole == "" {
        requiredRole = perm.Get("*")
    }
    if requiredRole == "" {
        return true
    }
    return role == requiredRole
}

func checkMediaType(accept Accept, method Method, mediaType MediaType) bool {
    acceptedTypes, exist := accept[method]
    if !exist {
        return true
    }
    for _, acceptedType := range acceptedTypes {
        if mediaType == acceptedType {
            return true
        }
    }
    return false
}

type typeWithPriority struct {
    t string
    q float64
}

func chooseType(provided MediaTypes, acceptHeader string) MediaType {
    if acceptHeader == "" {
        acceptHeader = "*/*"
    }

    typeAndParams := strings.Split(acceptHeader, ",")
    typeWithPriorities := make([]typeWithPriority, len(typeAndParams))
    for i, tpstr := range(typeAndParams) {
        tppair := strings.Split(tpstr, ";")
        var q float64
        if len(tppair) == 1 {
            q = 1.0
        } else {
            q, _ = strconv.ParseFloat(strings.TrimSpace(tppair[1])[2:], 32)
        }
        j := 0
        for ; j<i ; j++ {
            if q > typeWithPriorities[j].q {
                break
            }
        }
        typeWithPriorities = append(typeWithPriorities[:j],
            append([]typeWithPriority{typeWithPriority{strings.TrimSpace(tppair[0]), q}},
                typeWithPriorities[j:]...)...)[:len(typeAndParams)]
    }

    for _, t := range(typeWithPriorities) {
        for _, p := range(provided) {
            matched, _ := regexp.MatchString(fmt.Sprintf("^%s$",
                strings.Replace(t.t, "*", "[^/]+", -1)), string(p))
            if matched {
                return p
            }
        }
    }
    return ""
}

func (c webApp) matchRules(w *wrappedWriter, r *http.Request) (result interface{}, chosenType MediaType) {
    for i, routeRule := range c.RouteRules {
        params := c.PatternMappings[i].Re.FindStringSubmatch(r.URL.Path)
        if params != nil {
            pathParams := make(map[string]string)
            if len(params) > 1 {
                for j, param := range params[1:] {
                    pathParams[c.PatternMappings[i].Names[j]] = param
                }
            }

            user, role := c.UserProvider.GetUserAndRole(r)
            ctx := Ctx {
                pathParams: pathParams,
                Username: user,
                Role: role,
                Request: r,
                ResponseWriter: w,
            }

            vResource := reflect.ValueOf(routeRule.Resource)
            vProvides := vResource.FieldByName("Provides")
            if vProvides.IsValid() {
                provided := vProvides.Interface().(Provides)[Method(r.Method)]
                if len(provided) > 0 {
                    ctx.ChosenType = chooseType(provided, r.Header.Get("Accept"))
                    if ctx.ChosenType == "" {
                        return notAcceptable{vProvides.Interface().(Provides)[Method(r.Method)]}, ""
                    }
                    w.Header().Set("Content-Type", string(ctx.ChosenType))
                }
            }

            vNewResource := reflect.New(reflect.TypeOf(routeRule.Resource)).Elem()
            for i := 0; i < vResource.NumField(); i++ {
                srcField := vResource.Field(i)
                newField := vNewResource.Field(i)

                switch reflect.TypeOf(srcField.Interface()).Name() {
                case "Ctx":
                    newField.Set(reflect.ValueOf(ctx))
                case "Perm":
                    if !checkPermission(vResource.Type().Field(i).Tag, Method(r.Method),
                            ctx.Role) {
                        w.Header()["WWW-Authenticate"] = []string{c.UserProvider.AuthHeader(r)}
                        if c.Settings["401_PAGE"] != "" {
                            w.Header().Set("Content-Type", "text/html")
                            w.WriteHeader(http.StatusUnauthorized)
                            f, err := os.Open(c.Settings["401_PAGE"])
                            if err != nil {
                                panic(err)
                            }
                            io.Copy(w, f)
                        } else {
                            http.Error(w, "unauthorized", http.StatusUnauthorized)
                        }
                        return w, ""
                    }
                case "Accept":
                    if !checkMediaType(srcField.Interface().(Accept), Method(r.Method),
                            MediaType(r.Header.Get("Content-Type"))) {
                        return unsupportedMediaType{}, ""
                    }
                default:
                    newField.Set(srcField)
                }
            }
            resource := vNewResource.Interface()

            result := getResult(r.Method, resource)
            return result, ctx.ChosenType
        }
    }
    return notFound{}, ""
}

func getAllowed(resource interface{}) (allowed []string) {
    _, ok := resource.(Getter)
    if ok {
        allowed = append(allowed, "GET", "HEAD")
    }
    _, ok = resource.(Poster)
    if ok {
        allowed = append(allowed, "POST")
    }
    _, ok = resource.(Putter)
    if ok {
        allowed = append(allowed, "PUT")
    }
    _, ok = resource.(Deleter)
    if ok {
        allowed = append(allowed, "DELETE")
    }
    return
}

func getResult(method string, resource interface{}) (result interface{}) {
    defer func() {
        if r := recover(); r != nil {
            rstr := fmt.Sprintf("%s", r)
            result = internalError {
                where: lineInfo(3),
                why: rstr + fullTrace(5, "\n\t"),
                code: errorCode(rstr),
            }
        }
    }()

    switch method {
    case "HEAD", "GET":
        h, ok := resource.(Getter)
        if ok {
            result = h.Get()
        }
    case "POST":
        h, ok := resource.(Poster)
        if ok {
            result = h.Post()
        }
    case "PUT":
        h, ok := resource.(Putter)
        if ok {
            result = h.Put()
        }
    case "DELETE":
        h, ok := resource.(Deleter)
        if ok {
            result = h.Delete()
        }
    default:
        return notImplemented{}
    }

    if result == nil {
        return methodNotAllowed{getAllowed(resource)}
    }
    return
}

func (c webApp) logRequest(w *wrappedWriter, r *http.Request, elapsedMs float64,
        result interface{}) {
    if w.status == 0 {
        log.Printf("%s %s %s Client Disconnected (%.2f ms)", r.RemoteAddr, r.Method,
            r.URL.Path, elapsedMs)
    } else {
        errMsg := ""
        if w.err.why != "" {
            errMsg = fmt.Sprintf("%s #%d %s ", w.err.where, w.err.code, w.err.why)
        }
        switch result.(type) {
        case unsupportedMediaType:
            errMsg = fmt.Sprintf(": %s ", r.Header.Get("Content-Type"))
        }
        log.Printf("%s %s %s %s %s(%.2f ms, %d bytes)", r.RemoteAddr, r.Method, r.URL.Path,
            http.StatusText(w.status), errMsg, elapsedMs, w.written)

        if c.DumpRequest {
            dump, _ := httputil.DumpRequest(r, false)
            log.Printf("%s", dump)
        }
    }
}

func (c webApp) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    ww := &wrappedWriter{
        status: 0,
        writer: w,
        inTime: time.Now(),
    }
    r.ParseForm()
    result, chosenType := c.matchRules(ww, r)
    c.writeResponse(ww, r, result, chosenType)

    elapsedMs := float64(time.Now().UnixNano() - ww.inTime.UnixNano()) / 1000000
    c.logRequest(ww, r, elapsedMs, result)
}

func CreateWebApp(rules []RouteRule) webApp {
    patternMappings := make([]PatternMapping, len(rules))
    for i, v := range rules {
        re := regexp.MustCompile("/{[^}]*}")
        params := re.FindAllString(v.Pattern, -1)
        names := make([]string, len(params))

        transformedPattern := v.Pattern
        for j, param := range params {
            names[j] = param[2:len(param)-1]
            transformedPattern = strings.Replace(transformedPattern, param, "[/]{0,1}([^/]*)", -1)
        }
        patternMappings[i] = PatternMapping{regexp.MustCompile("^"+transformedPattern+"$"), names}
    }

    return webApp{
        RouteRules: rules,
        PatternMappings: patternMappings,
        UserProvider: EmptyUserProvider{},
        Settings: make(map[string]string),
    }
}

type Method string
type Perm struct{}

type MediaType string
type MediaTypes []MediaType
type Accept map[Method]MediaTypes
type Provides map[Method]MediaTypes
<commit_msg>move vitali.provides to struct field tag<commit_after>package vitali

import (
    "os"
    "io"
    "log"
    "fmt"
    "strconv"
    "net/http"
    "net/http/httputil"
    "html/template"
    "time"
    "regexp"
    "strings"
    "reflect"
)

type RouteRule struct {
    Pattern string
    Resource interface{}
}

type PatternMapping struct {
    Re *regexp.Regexp
    Names []string
}

type webApp struct {
    RouteRules []RouteRule
    PatternMappings []PatternMapping
    UserProvider UserProvider
    Settings map[string]string
    DumpRequest bool
    ErrTemplate *template.Template
}

func checkPermission(perm reflect.StructTag, method Method, role string) bool {
    requiredRole := perm.Get(string(method))
    if requiredRole == "" {
        requiredRole = perm.Get("*")
    }
    if requiredRole == "" {
        return true
    }
    return role == requiredRole
}

func checkMediaType(accept Accept, method Method, mediaType MediaType) bool {
    acceptedTypes, exist := accept[method]
    if !exist {
        return true
    }
    for _, acceptedType := range acceptedTypes {
        if mediaType == acceptedType {
            return true
        }
    }
    return false
}

type typeWithPriority struct {
    t string
    q float64
}

func chooseType(provided MediaTypes, acceptHeader string) MediaType {
    if acceptHeader == "" {
        acceptHeader = "*/*"
    }

    typeAndParams := strings.Split(acceptHeader, ",")
    typeWithPriorities := make([]typeWithPriority, len(typeAndParams))
    for i, tpstr := range(typeAndParams) {
        tppair := strings.Split(tpstr, ";")
        var q float64
        if len(tppair) == 1 {
            q = 1.0
        } else {
            q, _ = strconv.ParseFloat(strings.TrimSpace(tppair[1])[2:], 32)
        }
        j := 0
        for ; j<i ; j++ {
            if q > typeWithPriorities[j].q {
                break
            }
        }
        typeWithPriorities = append(typeWithPriorities[:j],
            append([]typeWithPriority{typeWithPriority{strings.TrimSpace(tppair[0]), q}},
                typeWithPriorities[j:]...)...)[:len(typeAndParams)]
    }

    for _, t := range(typeWithPriorities) {
        for _, p := range(provided) {
            matched, _ := regexp.MatchString(fmt.Sprintf("^%s$",
                strings.Replace(t.t, "*", "[^/]+", -1)), string(p))
            if matched {
                return p
            }
        }
    }
    return ""
}

func (c webApp) matchRules(w *wrappedWriter, r *http.Request) (result interface{}, chosenType MediaType) {
    for i, routeRule := range c.RouteRules {
        params := c.PatternMappings[i].Re.FindStringSubmatch(r.URL.Path)
        if params != nil {
            pathParams := make(map[string]string)
            if len(params) > 1 {
                for j, param := range params[1:] {
                    pathParams[c.PatternMappings[i].Names[j]] = param
                }
            }

            user, role := c.UserProvider.GetUserAndRole(r)
            ctx := Ctx {
                pathParams: pathParams,
                Username: user,
                Role: role,
                Request: r,
                ResponseWriter: w,
            }

            vResource := reflect.ValueOf(routeRule.Resource)
            tProvides, found := reflect.TypeOf(routeRule.Resource).FieldByName("Provides")
            if found {
                providedStr := tProvides.Tag.Get(r.Method)
                if providedStr != "" {
                    providedTmp := strings.Split(providedStr, ",")
                    provided := make(MediaTypes, len(providedTmp))
                    for i, v := range providedTmp {
                        provided[i] = MediaType(v)
                    }

                    ctx.ChosenType = MediaType(chooseType(provided, r.Header.Get("Accept")))
                    if ctx.ChosenType == "" {
                        return notAcceptable{provided}, ""
                    }
                    w.Header().Set("Content-Type", string(ctx.ChosenType))
                }
            }

            vNewResource := reflect.New(reflect.TypeOf(routeRule.Resource)).Elem()
            for i := 0; i < vResource.NumField(); i++ {
                srcField := vResource.Field(i)
                newField := vNewResource.Field(i)

                switch reflect.TypeOf(srcField.Interface()).Name() {
                case "Ctx":
                    newField.Set(reflect.ValueOf(ctx))
                case "Perm":
                    if !checkPermission(vResource.Type().Field(i).Tag, Method(r.Method),
                            ctx.Role) {
                        w.Header()["WWW-Authenticate"] = []string{c.UserProvider.AuthHeader(r)}
                        if c.Settings["401_PAGE"] != "" {
                            w.Header().Set("Content-Type", "text/html")
                            w.WriteHeader(http.StatusUnauthorized)
                            f, err := os.Open(c.Settings["401_PAGE"])
                            if err != nil {
                                panic(err)
                            }
                            io.Copy(w, f)
                        } else {
                            http.Error(w, "unauthorized", http.StatusUnauthorized)
                        }
                        return w, ""
                    }
                case "Accept":
                    if !checkMediaType(srcField.Interface().(Accept), Method(r.Method),
                            MediaType(r.Header.Get("Content-Type"))) {
                        return unsupportedMediaType{}, ""
                    }
                default:
                    newField.Set(srcField)
                }
            }
            resource := vNewResource.Interface()

            result := getResult(r.Method, resource)
            return result, ctx.ChosenType
        }
    }
    return notFound{}, ""
}

func getAllowed(resource interface{}) (allowed []string) {
    _, ok := resource.(Getter)
    if ok {
        allowed = append(allowed, "GET", "HEAD")
    }
    _, ok = resource.(Poster)
    if ok {
        allowed = append(allowed, "POST")
    }
    _, ok = resource.(Putter)
    if ok {
        allowed = append(allowed, "PUT")
    }
    _, ok = resource.(Deleter)
    if ok {
        allowed = append(allowed, "DELETE")
    }
    return
}

func getResult(method string, resource interface{}) (result interface{}) {
    defer func() {
        if r := recover(); r != nil {
            rstr := fmt.Sprintf("%s", r)
            result = internalError {
                where: lineInfo(3),
                why: rstr + fullTrace(5, "\n\t"),
                code: errorCode(rstr),
            }
        }
    }()

    switch method {
    case "HEAD", "GET":
        h, ok := resource.(Getter)
        if ok {
            result = h.Get()
        }
    case "POST":
        h, ok := resource.(Poster)
        if ok {
            result = h.Post()
        }
    case "PUT":
        h, ok := resource.(Putter)
        if ok {
            result = h.Put()
        }
    case "DELETE":
        h, ok := resource.(Deleter)
        if ok {
            result = h.Delete()
        }
    default:
        return notImplemented{}
    }

    if result == nil {
        return methodNotAllowed{getAllowed(resource)}
    }
    return
}

func (c webApp) logRequest(w *wrappedWriter, r *http.Request, elapsedMs float64,
        result interface{}) {
    if w.status == 0 {
        log.Printf("%s %s %s Client Disconnected (%.2f ms)", r.RemoteAddr, r.Method,
            r.URL.Path, elapsedMs)
    } else {
        errMsg := ""
        if w.err.why != "" {
            errMsg = fmt.Sprintf("%s #%d %s ", w.err.where, w.err.code, w.err.why)
        }
        switch result.(type) {
        case unsupportedMediaType:
            errMsg = fmt.Sprintf(": %s ", r.Header.Get("Content-Type"))
        }
        log.Printf("%s %s %s %s %s(%.2f ms, %d bytes)", r.RemoteAddr, r.Method, r.URL.Path,
            http.StatusText(w.status), errMsg, elapsedMs, w.written)

        if c.DumpRequest {
            dump, _ := httputil.DumpRequest(r, false)
            log.Printf("%s", dump)
        }
    }
}

func (c webApp) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    ww := &wrappedWriter{
        status: 0,
        writer: w,
        inTime: time.Now(),
    }
    r.ParseForm()
    result, chosenType := c.matchRules(ww, r)
    c.writeResponse(ww, r, result, chosenType)

    elapsedMs := float64(time.Now().UnixNano() - ww.inTime.UnixNano()) / 1000000
    c.logRequest(ww, r, elapsedMs, result)
}

func CreateWebApp(rules []RouteRule) webApp {
    patternMappings := make([]PatternMapping, len(rules))
    for i, v := range rules {
        re := regexp.MustCompile("/{[^}]*}")
        params := re.FindAllString(v.Pattern, -1)
        names := make([]string, len(params))

        transformedPattern := v.Pattern
        for j, param := range params {
            names[j] = param[2:len(param)-1]
            transformedPattern = strings.Replace(transformedPattern, param, "[/]{0,1}([^/]*)", -1)
        }
        patternMappings[i] = PatternMapping{regexp.MustCompile("^"+transformedPattern+"$"), names}
    }

    return webApp{
        RouteRules: rules,
        PatternMappings: patternMappings,
        UserProvider: EmptyUserProvider{},
        Settings: make(map[string]string),
    }
}

type Method string
type Perm struct{}

type MediaType string
type MediaTypes []MediaType
type Accept map[Method]MediaTypes
type Provides map[Method]MediaTypes
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 certfetcher

import (
	"crypto"
	"crypto/x509"
	"strconv"

	"github.com/WICG/webpackage/go/signedexchange"
	"github.com/go-acme/lego/v3/certcrypto"
	"github.com/go-acme/lego/v3/challenge/http01"
	"github.com/go-acme/lego/v3/challenge/tlsalpn01"
	"github.com/go-acme/lego/v3/lego"
	"github.com/go-acme/lego/v3/providers/http/webroot"
	"github.com/go-acme/lego/v3/registration"
	"github.com/pkg/errors"
)

type CertFetcher struct {
	AcmeDiscoveryURL string
	AcmeUser         AcmeUser
	legoClient       *lego.Client
	CertSignRequest  *x509.CertificateRequest
}

// Implements registration.User
type AcmeUser struct {
	Email        string
	Registration *registration.Resource
	key          crypto.PrivateKey
}

func (u *AcmeUser) GetEmail() string {
	return u.Email
}
func (u AcmeUser) GetRegistration() *registration.Resource {
	return u.Registration
}
func (u *AcmeUser) GetPrivateKey() crypto.PrivateKey {
	return u.key
}

// Initializes the cert fetcher with information it needs to fetch new certificates in the future.
// TODO(banaag): per gregable@ comments:
// Callsite could have some structure like:
//
// fetcher := CertFetcher()
// fetcher.setUser(email, privateKey)
// fetcher.bindToPort(port)
func New(email string, eabKid string, eabHmac string, certSignRequest *x509.CertificateRequest,
	privateKey crypto.PrivateKey, acmeDiscoURL string, httpChallengePort int, httpChallengeWebRoot string,
	tlsChallengePort int, dnsProvider string, shouldRegister bool) (*CertFetcher, error) {

	acmeUser := AcmeUser{
		Email: email,
		key:   privateKey,
	}
	config := lego.NewConfig(&acmeUser)

	config.CADirURL = acmeDiscoURL
	config.Certificate.KeyType = certcrypto.EC256

	// A client facilitates communication with the CA server.
	client, err := lego.NewClient(config)
	if err != nil {
		return nil, errors.Wrap(err, "Obtaining LEGO client.")
	}

	// We specify an http port of `httpChallengePort`
	// because we aren't running as root and can't bind a listener to port 80 and 443
	// (used later when we attempt to pass challenges). Keep in mind that you still
	// need to proxy challenge traffic to port `acmeChallengePort`.
	if httpChallengePort != 0 {
		err := client.Challenge.SetHTTP01Provider(
			http01.NewProviderServer("", strconv.Itoa(httpChallengePort)))
		if err != nil {
			return nil, errors.Wrap(err, "Setting up HTTP01 challenge provider.")
		}
	}
	if httpChallengeWebRoot != "" {
		httpProvider, err := webroot.NewHTTPProvider(httpChallengeWebRoot)
		if err != nil {
			return nil, errors.Wrap(err, "Getting HTTP01 challenge provider.")
		}
		err = client.Challenge.SetHTTP01Provider(httpProvider)
		if err != nil {
			return nil, errors.Wrap(err, "Setting up HTTP01 challenge provider.")
		}
	}

	if tlsChallengePort != 0 {
		err := client.Challenge.SetTLSALPN01Provider(tlsalpn01.NewProviderServer("", strconv.Itoa(tlsChallengePort)))
		if err != nil {
			return nil, errors.Wrap(err, "Setting up TLSALPN01 challenge provider.")
		}
	}

	if dnsProvider != "" {
		provider, err := DNSProvider(dnsProvider)
		if err != nil {
			return nil, errors.Wrap(err, "Getting DNS01 challenge provider.")
		}
		err = client.Challenge.SetDNS01Provider(provider)
		if err != nil {
			return nil, errors.Wrap(err, "Setting up DNS01 challenge provider.")
		}
	}

	// Theoretically, this should always be set to false as users should have pre-registered for access
	// to the ACME CA and agreed to the TOS.
	// TODO(banaag): revisit this when trying the class out with Digicert CA.
	if !shouldRegister {
		acmeUser.Registration = new(registration.Resource)
	} else {
		var reg *registration.Resource
		var err error

		// TODO(banaag) make sure we present the TOS URL to the user and prompt for confirmation.
		// The plan is to move this to some separate setup command outside the server which would be
		// executed one time. Alternatively, we can have a field in the toml file that is documented
		// to indicate agreement with TOS.
		if eabKid == "" && eabHmac == "" {
			reg, err = client.Registration.Register(registration.RegisterOptions{
				TermsOfServiceAgreed: true})
		} else {
			reg, err = client.Registration.RegisterWithExternalAccountBinding(registration.RegisterEABOptions{
				TermsOfServiceAgreed: true,
				Kid:                  eabKid,
				HmacEncoded:          eabHmac})
		}

		if err != nil {
			return nil, errors.Wrap(err, "ACME CA client registration")
		}
		acmeUser.Registration = reg
	}

	return &CertFetcher{
		AcmeDiscoveryURL: acmeDiscoURL,
		AcmeUser:         acmeUser,
		legoClient:       client,
		CertSignRequest:  certSignRequest,
	}, nil
}

func (f *CertFetcher) FetchNewCert() ([]*x509.Certificate, error) {
	// Each resource comes back with the cert bytes, the bytes of the client's
	// private key, and a certificate URL.
	resource, err := f.legoClient.Certificate.ObtainForCSR(*f.CertSignRequest, true)
	if err != nil {
		return nil, err
	}

	if resource == nil {
		return nil, errors.New("No resource returned.")
	}

	if resource.Certificate == nil {
		return nil, errors.New("No certificates were returned.")
	}

	cert, err := signedexchange.ParseCertificates(resource.Certificate)
	if err != nil {
		return nil, err
	}

	return cert, err
}
<commit_msg>Fix EAB account workflow to check for existing accounts using ResolveAccountByKey. (#488)<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//     https://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 certfetcher

import (
	"crypto"
	"crypto/x509"
	"strconv"

	"github.com/WICG/webpackage/go/signedexchange"
	"github.com/go-acme/lego/v3/certcrypto"
	"github.com/go-acme/lego/v3/challenge/http01"
	"github.com/go-acme/lego/v3/challenge/tlsalpn01"
	"github.com/go-acme/lego/v3/lego"
	"github.com/go-acme/lego/v3/providers/http/webroot"
	"github.com/go-acme/lego/v3/registration"
	"github.com/pkg/errors"
)

type CertFetcher struct {
	AcmeDiscoveryURL string
	AcmeUser         AcmeUser
	legoClient       *lego.Client
	CertSignRequest  *x509.CertificateRequest
}

// Implements registration.User
type AcmeUser struct {
	Email        string
	Registration *registration.Resource
	key          crypto.PrivateKey
}

func (u *AcmeUser) GetEmail() string {
	return u.Email
}
func (u AcmeUser) GetRegistration() *registration.Resource {
	return u.Registration
}
func (u *AcmeUser) GetPrivateKey() crypto.PrivateKey {
	return u.key
}

// Initializes the cert fetcher with information it needs to fetch new certificates in the future.
// TODO(banaag): per gregable@ comments:
// Callsite could have some structure like:
//
// fetcher := CertFetcher()
// fetcher.setUser(email, privateKey)
// fetcher.bindToPort(port)
func New(email string, eabKid string, eabHmac string, certSignRequest *x509.CertificateRequest,
	privateKey crypto.PrivateKey, acmeDiscoURL string, httpChallengePort int, httpChallengeWebRoot string,
	tlsChallengePort int, dnsProvider string, shouldRegister bool) (*CertFetcher, error) {

	acmeUser := AcmeUser{
		Email: email,
		key:   privateKey,
	}
	config := lego.NewConfig(&acmeUser)

	config.CADirURL = acmeDiscoURL
	config.Certificate.KeyType = certcrypto.EC256

	// A client facilitates communication with the CA server.
	client, err := lego.NewClient(config)
	if err != nil {
		return nil, errors.Wrap(err, "Obtaining LEGO client.")
	}

	// We specify an http port of `httpChallengePort`
	// because we aren't running as root and can't bind a listener to port 80 and 443
	// (used later when we attempt to pass challenges). Keep in mind that you still
	// need to proxy challenge traffic to port `acmeChallengePort`.
	if httpChallengePort != 0 {
		err := client.Challenge.SetHTTP01Provider(
			http01.NewProviderServer("", strconv.Itoa(httpChallengePort)))
		if err != nil {
			return nil, errors.Wrap(err, "Setting up HTTP01 challenge provider.")
		}
	}
	if httpChallengeWebRoot != "" {
		httpProvider, err := webroot.NewHTTPProvider(httpChallengeWebRoot)
		if err != nil {
			return nil, errors.Wrap(err, "Getting HTTP01 challenge provider.")
		}
		err = client.Challenge.SetHTTP01Provider(httpProvider)
		if err != nil {
			return nil, errors.Wrap(err, "Setting up HTTP01 challenge provider.")
		}
	}

	if tlsChallengePort != 0 {
		err := client.Challenge.SetTLSALPN01Provider(tlsalpn01.NewProviderServer("", strconv.Itoa(tlsChallengePort)))
		if err != nil {
			return nil, errors.Wrap(err, "Setting up TLSALPN01 challenge provider.")
		}
	}

	if dnsProvider != "" {
		provider, err := DNSProvider(dnsProvider)
		if err != nil {
			return nil, errors.Wrap(err, "Getting DNS01 challenge provider.")
		}
		err = client.Challenge.SetDNS01Provider(provider)
		if err != nil {
			return nil, errors.Wrap(err, "Setting up DNS01 challenge provider.")
		}
	}

	var reg *registration.Resource
	if !shouldRegister {
		acmeUser.Registration = new(registration.Resource)
	} else if reg, err = client.Registration.ResolveAccountByKey(); err == nil {
		// Check if we already have an account.
		acmeUser.Registration = reg
	} else {
		// We need to reset the LEGO client after calling Registration.ResolveAccountByKey().
		client, err = lego.NewClient(config)
		if err != nil {
			return nil, errors.Wrap(err, "Obtaining LEGO client.")
		}
		// TODO(banaag) make sure we present the TOS URL to the user and prompt for confirmation.
		// The plan is to move this to some separate setup command outside the server which would be
		// executed one time. Alternatively, we can have a field in the toml file that is documented
		// to indicate agreement with TOS.
		if eabKid == "" && eabHmac == "" {
			reg, err = client.Registration.Register(registration.RegisterOptions{
				TermsOfServiceAgreed: true})
		} else {
			reg, err = client.Registration.RegisterWithExternalAccountBinding(registration.RegisterEABOptions{
				TermsOfServiceAgreed: true,
				Kid:                  eabKid,
				HmacEncoded:          eabHmac})
		}

		if err != nil {
			return nil, errors.Wrap(err, "ACME CA client registration")
		}
		acmeUser.Registration = reg
	}

	return &CertFetcher{
		AcmeDiscoveryURL: acmeDiscoURL,
		AcmeUser:         acmeUser,
		legoClient:       client,
		CertSignRequest:  certSignRequest,
	}, nil
}

func (f *CertFetcher) FetchNewCert() ([]*x509.Certificate, error) {
	// Each resource comes back with the cert bytes, the bytes of the client's
	// private key, and a certificate URL.
	resource, err := f.legoClient.Certificate.ObtainForCSR(*f.CertSignRequest, true)
	if err != nil {
		return nil, err
	}

	if resource == nil {
		return nil, errors.New("No resource returned.")
	}

	if resource.Certificate == nil {
		return nil, errors.New("No certificates were returned.")
	}

	cert, err := signedexchange.ParseCertificates(resource.Certificate)
	if err != nil {
		return nil, err
	}

	return cert, err
}
<|endoftext|>


<commit_before>// Copyright 2015 Prometheus Team
// 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 main

import (
	"reflect"
	"testing"
	"time"

	"github.com/prometheus/common/model"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/alertmanager/config"
)

func TestRouteMatch(t *testing.T) {
	in := `
receiver: 'notify-def'

routes:
- match:
    owner: 'team-A'

  receiver: 'notify-A'

  routes:
  - match:
      env: 'testing'

    receiver: 'notify-testing'
    group_by: []

  - match:
      env: "production"

    receiver: 'notify-productionA'
    group_wait: 1m

    continue: true

  - match_re:
      env: "produ.*"

    receiver: 'notify-productionB'
    group_wait: 30s
    group_interval: 5m
    repeat_interval: 1h
    group_by: ['job']

- match_re:
    owner: 'team-(B|C)'

  group_by: ['foo', 'bar']
  group_wait: 2m
  receiver: 'notify-BC'
`

	var ctree config.Route
	if err := yaml.Unmarshal([]byte(in), &ctree); err != nil {
		t.Fatal(err)
	}
	var (
		def  = DefaultRouteOpts
		tree = NewRoute(&ctree, nil)
	)
	lset := func(labels ...string) map[model.LabelName]struct{} {
		s := map[model.LabelName]struct{}{}
		for _, ls := range labels {
			s[model.LabelName(ls)] = struct{}{}
		}
		return s
	}

	tests := []struct {
		input  model.LabelSet
		result []*RouteOpts
	}{
		{
			input: model.LabelSet{
				"owner": "team-A",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-A",
					GroupBy:        def.GroupBy,
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "unset",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-A",
					GroupBy:        def.GroupBy,
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-C",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-BC",
					GroupBy:        lset("foo", "bar"),
					GroupWait:      2 * time.Minute,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "testing",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-testing",
					GroupBy:        lset(),
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "production",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-productionA",
					GroupBy:        def.GroupBy,
					GroupWait:      1 * time.Minute,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
				{
					Receiver:       "notify-productionB",
					GroupBy:        lset("job"),
					GroupWait:      30 * time.Second,
					GroupInterval:  5 * time.Minute,
					RepeatInterval: 1 * time.Hour,
				},
			},
		},
	}

	for _, test := range tests {
		var matches []*RouteOpts
		for _, r := range tree.Match(test.input) {
			matches = append(matches, &r.RouteOpts)
		}

		if !reflect.DeepEqual(matches, test.result) {
			t.Errorf("\nexpected:\n%v\ngot:\n%v", test.result, matches)
		}
	}
}
<commit_msg>Add route tests for receiver inheritance<commit_after>// Copyright 2015 Prometheus Team
// 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 main

import (
	"reflect"
	"testing"
	"time"

	"github.com/prometheus/common/model"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/alertmanager/config"
)

func TestRouteMatch(t *testing.T) {
	in := `
receiver: 'notify-def'

routes:
- match:
    owner: 'team-A'

  receiver: 'notify-A'

  routes:
  - match:
      env: 'testing'

    receiver: 'notify-testing'
    group_by: []

  - match:
      env: "production"

    receiver: 'notify-productionA'
    group_wait: 1m

    continue: true

  - match_re:
      env: "produ.*"

    receiver: 'notify-productionB'
    group_wait: 30s
    group_interval: 5m
    repeat_interval: 1h
    group_by: ['job']

- match_re:
    owner: 'team-(B|C)'

  group_by: ['foo', 'bar']
  group_wait: 2m
  receiver: 'notify-BC'

- match:
    group_by: 'role'
  group_by: ['role']

  routes:
  - match:
      env: 'testing'
    receiver: 'notify-testing'
    routes:
    - match:
        wait: 'long'
      group_wait: 2m
`

	var ctree config.Route
	if err := yaml.Unmarshal([]byte(in), &ctree); err != nil {
		t.Fatal(err)
	}
	var (
		def  = DefaultRouteOpts
		tree = NewRoute(&ctree, nil)
	)
	lset := func(labels ...string) map[model.LabelName]struct{} {
		s := map[model.LabelName]struct{}{}
		for _, ls := range labels {
			s[model.LabelName(ls)] = struct{}{}
		}
		return s
	}

	tests := []struct {
		input  model.LabelSet
		result []*RouteOpts
	}{
		{
			input: model.LabelSet{
				"owner": "team-A",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-A",
					GroupBy:        def.GroupBy,
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "unset",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-A",
					GroupBy:        def.GroupBy,
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-C",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-BC",
					GroupBy:        lset("foo", "bar"),
					GroupWait:      2 * time.Minute,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "testing",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-testing",
					GroupBy:        lset(),
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"owner": "team-A",
				"env":   "production",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-productionA",
					GroupBy:        def.GroupBy,
					GroupWait:      1 * time.Minute,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
				{
					Receiver:       "notify-productionB",
					GroupBy:        lset("job"),
					GroupWait:      30 * time.Second,
					GroupInterval:  5 * time.Minute,
					RepeatInterval: 1 * time.Hour,
				},
			},
		},
		{
			input: model.LabelSet{
				"group_by": "role",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-def",
					GroupBy:        lset("role"),
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"env":      "testing",
				"group_by": "role",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-testing",
					GroupBy:        lset("role"),
					GroupWait:      def.GroupWait,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
		{
			input: model.LabelSet{
				"env":      "testing",
				"group_by": "role",
				"wait":     "long",
			},
			result: []*RouteOpts{
				{
					Receiver:       "notify-testing",
					GroupBy:        lset("role"),
					GroupWait:      2 * time.Minute,
					GroupInterval:  def.GroupInterval,
					RepeatInterval: def.RepeatInterval,
				},
			},
		},
	}

	for _, test := range tests {
		var matches []*RouteOpts
		for _, r := range tree.Match(test.input) {
			matches = append(matches, &r.RouteOpts)
		}

		if !reflect.DeepEqual(matches, test.result) {
			t.Errorf("\nexpected:\n%v\ngot:\n%v", test.result, matches)
		}
	}
}
<|endoftext|>


<commit_before>package handlers

import (
	"fmt"
	"net/http"
    "net/url"
    "io/ioutil"
    "strings"
    "github.com/nelsonleduc/calmanbot/handlers/models"
    "encoding/json"
    "bytes"
    "regexp"
    "sort"
)

func isValidHTTPURLString(s string)  bool {
    URL, _ := url.Parse(s)
    return (URL.Scheme == "http" || URL.Scheme == "https")
}

func HandleCalman(w http.ResponseWriter, r *http.Request) {

    message := ParseMessageJSON(r.Body)
    bot, _ := models.FetchBot(message.GroupID)

    if !strings.HasPrefix(strings.ToLower(message.Text[1:]), strings.ToLower(bot.BotName)) {
        return
    }

    actions, _ := models.FetchActions(true)
    sort.Sort(models.ByPriority(actions))

    var (
        act models.Action
        sMatch string
    )
    for _, a := range actions {
        r, _ := regexp.Compile("(?i)" + *a.Pattern)
        matched := r.FindStringSubmatch(message.Text)
        if len(matched) > 1 && matched[1] != "" {
            sMatch = matched[1]
            act = a
            break
        }
    }

    updateAction(&act, sMatch)

    postString := ""
    for {
        if act.IsURLType() {
            postString = handleURLAction(act, w, bot)
        } else {
            postString = act.Content
        }

        if postString != "" || act.FallbackAction == nil {
            break
        } else {
            act, _ = models.FetchAction(*act.FallbackAction)
        }
    }

    if postString != "" {
        fmt.Printf("Action: %v\n", act.Content)
        fmt.Printf("Posting: %v\n", postString)
        postText(bot, postString)
    }
}

func handleURLAction(a models.Action, w http.ResponseWriter, b models.Bot) string {

    fmt.Fprintln(w, a)
    resp, err := http.Get(a.Content)

    if err == nil {

        content, _ := ioutil.ReadAll(resp.Body)
        pathString := *a.DataPath

        str := ParseJSON(content, pathString)
        if str == "" {
            return ""
        } else {

            if !validateURL(str) {
                fmt.Printf("Invalid URL: %v\n", str)

                oldStr := str
                for i := 0; i < 3 && oldStr == str; i++ {
                    str = ParseJSON(content, pathString)
                }

                if !validateURL(str) {
                    return ""
                } else {
                    return str
                }
            } else {
                return str
            }
        }
    } else {
        return ""
    }

    resp.Body.Close()
    return ""
}

func postText(b models.Bot, t string) {

    postURL := "https://api.groupme.com/v3/bots/post"
    postBody := map[string]string {
        "bot_id": b.Key,
        "text": t,
    }

    encoded, _ := json.Marshal(postBody)

    http.Post(postURL, "application/json", bytes.NewReader(encoded))
}

func validateURL(u string) bool {

    client := http.Client{}
    if isValidHTTPURLString(u) {
        req, err := http.NewRequest("HEAD", u, nil)
        if err != nil {
            return false
        }

        resp, err := client.Do(req)

        if err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300 {
            return true
        } else {
            return false
        }
    } else {
        return true
    }

    return true
}

func updateAction(a *models.Action, text string) {
    text = url.QueryEscape(text)

    a.Content = strings.Replace(a.Content, "{_text_}", text, -1)
}<commit_msg>Update fallback actions<commit_after>package handlers

import (
	"fmt"
	"net/http"
    "net/url"
    "io/ioutil"
    "strings"
    "github.com/nelsonleduc/calmanbot/handlers/models"
    "encoding/json"
    "bytes"
    "regexp"
    "sort"
)

func isValidHTTPURLString(s string)  bool {
    URL, _ := url.Parse(s)
    return (URL.Scheme == "http" || URL.Scheme == "https")
}

func HandleCalman(w http.ResponseWriter, r *http.Request) {

    message := ParseMessageJSON(r.Body)
    bot, _ := models.FetchBot(message.GroupID)

    if !strings.HasPrefix(strings.ToLower(message.Text[1:]), strings.ToLower(bot.BotName)) {
        return
    }

    actions, _ := models.FetchActions(true)
    sort.Sort(models.ByPriority(actions))

    var (
        act models.Action
        sMatch string
    )
    for _, a := range actions {
        r, _ := regexp.Compile("(?i)" + *a.Pattern)
        matched := r.FindStringSubmatch(message.Text)
        if len(matched) > 1 && matched[1] != "" {
            sMatch = matched[1]
            act = a
            break
        }
    }

    postString := ""
    for {
        updateAction(&act, sMatch)
        if act.IsURLType() {
            postString = handleURLAction(act, w, bot)
        } else {
            postString = act.Content
        }

        if postString != "" || act.FallbackAction == nil {
            break
        } else {
            act, _ = models.FetchAction(*act.FallbackAction)
        }
    }

    if postString != "" {
        fmt.Printf("Action: %v\n", act.Content)
        fmt.Printf("Posting: %v\n", postString)
        postText(bot, postString)
    }
}

func handleURLAction(a models.Action, w http.ResponseWriter, b models.Bot) string {

    fmt.Fprintln(w, a)
    resp, err := http.Get(a.Content)

    if err == nil {

        content, _ := ioutil.ReadAll(resp.Body)
        pathString := *a.DataPath

        str := ParseJSON(content, pathString)
        if str == "" {
            return ""
        } else {

            if !validateURL(str) {
                fmt.Printf("Invalid URL: %v\n", str)

                oldStr := str
                for i := 0; i < 3 && oldStr == str; i++ {
                    str = ParseJSON(content, pathString)
                }

                if !validateURL(str) {
                    return ""
                } else {
                    return str
                }
            } else {
                return str
            }
        }
    } else {
        return ""
    }

    resp.Body.Close()
    return ""
}

func postText(b models.Bot, t string) {

    postURL := "https://api.groupme.com/v3/bots/post"
    postBody := map[string]string {
        "bot_id": b.Key,
        "text": t,
    }

    encoded, _ := json.Marshal(postBody)

    http.Post(postURL, "application/json", bytes.NewReader(encoded))
}

func validateURL(u string) bool {

    client := http.Client{}
    if isValidHTTPURLString(u) {
        req, err := http.NewRequest("HEAD", u, nil)
        if err != nil {
            return false
        }

        resp, err := client.Do(req)

        if err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300 {
            return true
        } else {
            return false
        }
    } else {
        return true
    }

    return true
}

func updateAction(a *models.Action, text string) {
    text = url.QueryEscape(text)

    a.Content = strings.Replace(a.Content, "{_text_}", text, -1)
}<|endoftext|>


<commit_before>package handlers

import (
	"log"
	"net/http"

	"github.com/bulletind/khabar/core"
	"github.com/bulletind/khabar/db"
	"github.com/bulletind/khabar/dbapi/topics"
	"github.com/bulletind/khabar/utils"
	"gopkg.in/mgo.v2"
	"gopkg.in/simversity/gottp.v2"
)

type TopicChannel struct {
	gottp.BaseHandler
}

func (self *TopicChannel) Delete(request *gottp.Request) {
	intopic := new(topics.Topic)

	channelIdent := request.GetArgument("channel").(string)

	if !core.IsChannelAvailable(channelIdent) {
		request.Raise(gottp.HttpError{
			http.StatusBadRequest,
			"Channel is not supported",
		})

		return
	}

	intopic.Ident = request.GetArgument("ident").(string)

	request.ConvertArguments(intopic)

	topic, err := topics.Get(
		intopic.User,
		intopic.Organization,
		intopic.Ident,
	)

	if err != nil && err != mgo.ErrNotFound {
		log.Println(err)
		request.Raise(gottp.HttpError{
			http.StatusInternalServerError,
			"Unable to fetch data, Please try again later.",
		})

		return

	}

	var hasData bool

	if topic == nil {
		log.Println("Creating new document")
		intopic.AddChannel(channelIdent)

		intopic.PrepareSave()
		if !intopic.IsValid(db.INSERT_OPERATION) {
			request.Raise(gottp.HttpError{
				http.StatusBadRequest,
				"Atleast one of the user or org must be present.",
			})

			return
		}

		if !utils.ValidateAndRaiseError(request, intopic) {
			log.Println("Validation Failed")
			return
		}

		topic = intopic

	} else {
		hasData = true

		for _, ident := range topic.Channels {
			if ident == channelIdent {
				request.Raise(gottp.HttpError{
					http.StatusConflict,
					"You have already unsubscribed this channel",
				})

				return
				break
			}
		}

		topic.AddChannel(channelIdent)
	}

	if hasData {
		err = topics.Update(
			topic.User,
			topic.Organization,
			topic.Ident,
			&utils.M{"channels": topic.Channels},
		)

		if err != nil {
			log.Println("Error while inserting document :" + err.Error())
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Internal server error.",
			})

			return
		} else {
			request.Write(utils.R{
				Data:       nil,
				Message:    "true",
				StatusCode: http.StatusNoContent,
			})

			return
		}
	} else {
		log.Println("Successfull call: Inserting document")
		topics.Insert(topic)
		request.Write(utils.R{
			Data:       nil,
			Message:    "true",
			StatusCode: http.StatusNoContent,
		})

		return
	}
}

func (self *TopicChannel) Post(request *gottp.Request) {
	topic := new(topics.Topic)

	channelIdent := request.GetArgument("channel").(string)
	topic.Ident = request.GetArgument("ident").(string)

	if !core.IsChannelAvailable(channelIdent) {
		request.Raise(gottp.HttpError{
			http.StatusBadRequest,
			"Channel is not supported",
		})

		return
	}

	request.ConvertArguments(topic)

	topic, err := topics.Get(
		topic.User,
		topic.Organization,
		topic.Ident,
	)

	if err != nil {
		if err != mgo.ErrNotFound {
			log.Println(err)
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Unable to fetch data, Please try again later.",
			})

		} else {
			request.Write(utils.R{
				Data:       nil,
				Message:    "true",
				StatusCode: http.StatusNoContent,
			})
		}

		return
	}

	topic.RemoveChannel(channelIdent)
	log.Println(topic.Channels)

	if len(topic.Channels) == 0 {
		log.Println("Deleting from database, since channels are now empty.")
		err = topics.Delete(

			&utils.M{
				"org":   topic.Organization,
				"user":  topic.User,
				"ident": topic.Ident,
			},
		)

	} else {
		log.Println("Updating...")

		err = topics.Update(
			topic.User, topic.Organization,
			topic.Ident, &utils.M{"channels": topic.Channels},
		)
	}

	if err != nil {
		request.Raise(gottp.HttpError{
			http.StatusInternalServerError,
			"Unable to delete.",
		})

		return
	}

	request.Write(utils.R{
		Data:       nil,
		Message:    "true",
		StatusCode: http.StatusNoContent,
	})

	return
}

type Topics struct {
	gottp.BaseHandler
}

func (self *Topics) Get(request *gottp.Request) {
	var args struct {
		Organization string `json:"org"`
		AppName      string `json:"app_name"`
		User         string `json:"user"`
	}

	request.ConvertArguments(&args)

	channels := []string{}
	for ident, _ := range core.ChannelMap {
		channels = append(channels, ident)
	}

	iter, err := topics.GetAll(args.User, args.AppName, args.Organization, channels)

	if err != nil {
		if err != mgo.ErrNotFound {
			log.Println(err)
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Unable to fetch data, Please try again later.",
			})

		} else {
			request.Raise(gottp.HttpError{
				http.StatusNotFound,
				"Not Found.",
			})
		}

		return
	}

	ret := []topics.ChotaTopic{}
	for _, singleRet := range iter {
		ret = append(ret, singleRet)
	}

	request.Write(ret)
	return
}

func (self *Topics) Post(request *gottp.Request) {
	var args struct {
		AppName string `json:"app_name"`
		Ident   string `json:"ident" required:"required"`
	}

	request.ConvertArguments(&args)
}

type Topic struct {
	gottp.BaseHandler
}

func (self *Topics) Delete(request *gottp.Request) {
	var args struct {
		Ident string `json:"ident" required:"required"`
	}

	request.ConvertArguments(&args)
	utils.ValidateAndRaiseError(request, args)

	topics.DeleteTopic(args.Ident)
	request.Write(true)
	return
}
<commit_msg>Make app_name compulsary<commit_after>package handlers

import (
	"log"
	"net/http"

	"github.com/bulletind/khabar/core"
	"github.com/bulletind/khabar/db"
	"github.com/bulletind/khabar/dbapi/topics"
	"github.com/bulletind/khabar/utils"
	"gopkg.in/mgo.v2"
	"gopkg.in/simversity/gottp.v2"
)

type TopicChannel struct {
	gottp.BaseHandler
}

func (self *TopicChannel) Delete(request *gottp.Request) {
	intopic := new(topics.Topic)

	channelIdent := request.GetArgument("channel").(string)

	if !core.IsChannelAvailable(channelIdent) {
		request.Raise(gottp.HttpError{
			http.StatusBadRequest,
			"Channel is not supported",
		})

		return
	}

	intopic.Ident = request.GetArgument("ident").(string)

	request.ConvertArguments(intopic)

	topic, err := topics.Get(
		intopic.User,
		intopic.Organization,
		intopic.Ident,
	)

	if err != nil && err != mgo.ErrNotFound {
		log.Println(err)
		request.Raise(gottp.HttpError{
			http.StatusInternalServerError,
			"Unable to fetch data, Please try again later.",
		})

		return

	}

	var hasData bool

	if topic == nil {
		log.Println("Creating new document")
		intopic.AddChannel(channelIdent)

		intopic.PrepareSave()
		if !intopic.IsValid(db.INSERT_OPERATION) {
			request.Raise(gottp.HttpError{
				http.StatusBadRequest,
				"Atleast one of the user or org must be present.",
			})

			return
		}

		if !utils.ValidateAndRaiseError(request, intopic) {
			log.Println("Validation Failed")
			return
		}

		topic = intopic

	} else {
		hasData = true

		for _, ident := range topic.Channels {
			if ident == channelIdent {
				request.Raise(gottp.HttpError{
					http.StatusConflict,
					"You have already unsubscribed this channel",
				})

				return
				break
			}
		}

		topic.AddChannel(channelIdent)
	}

	if hasData {
		err = topics.Update(
			topic.User,
			topic.Organization,
			topic.Ident,
			&utils.M{"channels": topic.Channels},
		)

		if err != nil {
			log.Println("Error while inserting document :" + err.Error())
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Internal server error.",
			})

			return
		} else {
			request.Write(utils.R{
				Data:       nil,
				Message:    "true",
				StatusCode: http.StatusNoContent,
			})

			return
		}
	} else {
		log.Println("Successfull call: Inserting document")
		topics.Insert(topic)
		request.Write(utils.R{
			Data:       nil,
			Message:    "true",
			StatusCode: http.StatusNoContent,
		})

		return
	}
}

func (self *TopicChannel) Post(request *gottp.Request) {
	topic := new(topics.Topic)

	channelIdent := request.GetArgument("channel").(string)
	topic.Ident = request.GetArgument("ident").(string)

	if !core.IsChannelAvailable(channelIdent) {
		request.Raise(gottp.HttpError{
			http.StatusBadRequest,
			"Channel is not supported",
		})

		return
	}

	request.ConvertArguments(topic)

	topic, err := topics.Get(
		topic.User,
		topic.Organization,
		topic.Ident,
	)

	if err != nil {
		if err != mgo.ErrNotFound {
			log.Println(err)
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Unable to fetch data, Please try again later.",
			})

		} else {
			request.Write(utils.R{
				Data:       nil,
				Message:    "true",
				StatusCode: http.StatusNoContent,
			})
		}

		return
	}

	topic.RemoveChannel(channelIdent)
	log.Println(topic.Channels)

	if len(topic.Channels) == 0 {
		log.Println("Deleting from database, since channels are now empty.")
		err = topics.Delete(

			&utils.M{
				"org":   topic.Organization,
				"user":  topic.User,
				"ident": topic.Ident,
			},
		)

	} else {
		log.Println("Updating...")

		err = topics.Update(
			topic.User, topic.Organization,
			topic.Ident, &utils.M{"channels": topic.Channels},
		)
	}

	if err != nil {
		request.Raise(gottp.HttpError{
			http.StatusInternalServerError,
			"Unable to delete.",
		})

		return
	}

	request.Write(utils.R{
		Data:       nil,
		Message:    "true",
		StatusCode: http.StatusNoContent,
	})

	return
}

type Topics struct {
	gottp.BaseHandler
}

func (self *Topics) Get(request *gottp.Request) {
	var args struct {
		Organization string `json:"org"`
		AppName      string `json:"app_name"`
		User         string `json:"user"`
	}

	request.ConvertArguments(&args)

	channels := []string{}
	for ident, _ := range core.ChannelMap {
		channels = append(channels, ident)
	}

	iter, err := topics.GetAll(args.User, args.AppName, args.Organization, channels)

	if err != nil {
		if err != mgo.ErrNotFound {
			log.Println(err)
			request.Raise(gottp.HttpError{
				http.StatusInternalServerError,
				"Unable to fetch data, Please try again later.",
			})

		} else {
			request.Raise(gottp.HttpError{
				http.StatusNotFound,
				"Not Found.",
			})
		}

		return
	}

	ret := []topics.ChotaTopic{}
	for _, singleRet := range iter {
		ret = append(ret, singleRet)
	}

	request.Write(ret)
	return
}

func (self *Topics) Post(request *gottp.Request) {
	var args struct {
		AppName string `json:"app_name" required:"required"`
		Ident   string `json:"ident" required:"required"`
	}

	request.ConvertArguments(&args)
}

type Topic struct {
	gottp.BaseHandler
}

func (self *Topics) Delete(request *gottp.Request) {
	var args struct {
		Ident string `json:"ident" required:"required"`
	}

	request.ConvertArguments(&args)
	utils.ValidateAndRaiseError(request, args)

	topics.DeleteTopic(args.Ident)
	request.Write(true)
	return
}
<|endoftext|>


<commit_before>package vshard

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"log"
	"math/big"
	"time"

	farm "github.com/dgryski/go-farm"
	jump "github.com/dgryski/go-jump"
	"github.com/youtube/vitess/go/memcache"
	"github.com/youtube/vitess/go/pools"
)

var (
	// ErrKeyNotFound defines the error mensage when key is not found on memcached
	ErrKeyNotFound = errors.New("error: key not found")
)

// VitessResource implements the expected interface for vitess internal pool
type VitessResource struct {
	*memcache.Connection
}

// ServerStrategy defines the signature for the sharding function
type ServerStrategy func(key string, numServers int) int

// Close closes connections in a pool
func (r VitessResource) Close() {
	r.Connection.Close()
}

// Pool defines the pool
type Pool struct {
	Servers        []string
	ServerStrategy ServerStrategy
	numServers     int
	pool           []*pools.ResourcePool
}

// PoolStats defines all stats vitess memcached driver exposes
type PoolStats struct {
	Slot        int
	Server      string
	Capacity    int64
	Available   int64
	MaxCap      int64
	WaitCount   int64
	WaitTime    time.Duration
	IdleTimeout time.Duration
}

// NewPool returns a new VitessPool
func NewPool(servers []string, capacity, maxCap int, idleTimeout time.Duration) (*Pool, error) {
	numServers := len(servers)

	pool := &Pool{
		Servers:        servers,
		numServers:     numServers,
		pool:           []*pools.ResourcePool{},
		ServerStrategy: ShardedServerStrategyFarmhash,
	}

	for i, server := range servers {
		func(_pool *[]*pools.ResourcePool, _server string) {
			*_pool = append(*_pool, pools.NewResourcePool(func() (pools.Resource, error) {
				c, err := memcache.Connect(_server, time.Minute)
				return VitessResource{c}, err
			}, capacity, maxCap, idleTimeout))

			conn, err := pool.GetPoolConnection(i)
			defer pool.ReturnConnection(i, conn)
			if err != nil {
				log.Fatalf("Can't connect to memcached: %s", err)
			}
		}(&pool.pool, server)
	}

	return pool, nil
}

// ShardedServerStrategyMD5 uses md5+jump to pick a server
func ShardedServerStrategyMD5(key string, numServers int) int {
	if numServers == 1 {
		return 0
	}

	hash := md5.Sum([]byte(key))
	hashInt := big.NewInt(0)
	hashInt.SetString(hex.EncodeToString(hash[:]), 16)

	return int(jump.Hash(hashInt.Uint64(), numServers))
}

// ShardedServerStrategyFarmhash uses farmhash+jump to pick a server
func ShardedServerStrategyFarmhash(key string, numServers int) int {
	if numServers == 1 {
		return 0
	}

	return int(jump.Hash(farm.Fingerprint64([]byte(key)), numServers))
}

// GetConnection returns a connection from the sharding pool, based on the key
func (v *Pool) GetConnection(key string) (*VitessResource, int, error) {
	poolNum := v.ServerStrategy(key, v.numServers)

	connection, err := v.GetPoolConnection(poolNum)
	if err != nil {
		return nil, -1, err
	}

	return connection, poolNum, nil
}

// GetPoolConnection returns a connection from a specific pool number
func (v *Pool) GetPoolConnection(poolNum int) (*VitessResource, error) {
	ctx := context.Background()

	resource, err := v.pool[poolNum].Get(ctx)
	if err != nil {
		return nil, err
	}

	connection := resource.(VitessResource)

	return &connection, nil
}

// ReturnConnection returns a connection to the pool
func (v *Pool) ReturnConnection(poolNum int, resource *VitessResource) {
	v.pool[poolNum].Put(*resource)
}

// GetKeyMapping returns a mapping of server to a list of keys, useful for Gets()
func (v *Pool) GetKeyMapping(keys ...string) map[int][]string {
	mapping := make(map[int][]string)

	for i := 0; i < v.numServers; i++ {
		mapping[i] = []string{}
	}

	for _, key := range keys {
		poolNum := v.ServerStrategy(key, v.numServers)
		mapping[poolNum] = append(mapping[poolNum], key)
	}

	return mapping
}
<commit_msg>hashing keys (with farmhash) to store on memcached adding locks to access the pool (needs benchmarking)<commit_after>package vshard

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"log"
	"math/big"
	"strconv"
	"sync"
	"time"

	farm "github.com/dgryski/go-farm"
	jump "github.com/dgryski/go-jump"
	"github.com/youtube/vitess/go/memcache"
	"github.com/youtube/vitess/go/pools"
)

var (
	// ErrKeyNotFound defines the error mensage when key is not found on memcached
	ErrKeyNotFound = errors.New("error: key not found")
)

// VitessResource implements the expected interface for vitess internal pool
type VitessResource struct {
	*memcache.Connection
}

// ServerStrategy defines the signature for the sharding function
type ServerStrategy func(key string, numServers int) int

// HashKeyStrategy defines the signature for the key hashing function
type HashKeyStrategy func(key string) string

// Close closes connections in a pool
func (r VitessResource) Close() {
	r.Connection.Close()
}

// Pool defines the pool
type Pool struct {
	Servers         []string
	ServerStrategy  ServerStrategy
	HashKeyStrategy HashKeyStrategy
	numServers      int
	pool            []*pools.ResourcePool
	sync.RWMutex
}

// PoolStats defines all stats vitess memcached driver exposes
type PoolStats struct {
	Slot        int
	Server      string
	Capacity    int64
	Available   int64
	MaxCap      int64
	WaitCount   int64
	WaitTime    time.Duration
	IdleTimeout time.Duration
}

// NewPool returns a new VitessPool
func NewPool(servers []string, capacity, maxCap int, idleTimeout time.Duration) (*Pool, error) {
	numServers := len(servers)

	pool := &Pool{
		Servers:         servers,
		numServers:      numServers,
		pool:            []*pools.ResourcePool{},
		ServerStrategy:  ShardedServerStrategyFarmhash,
		HashKeyStrategy: HashKeyStrategyFarmhash,
	}

	for i, server := range servers {
		func(_pool *Pool, _server string) {
			_pool.Lock()
			_pool.pool = append(_pool.pool, pools.NewResourcePool(func() (pools.Resource, error) {
				c, err := memcache.Connect(_server, time.Minute)
				return VitessResource{c}, err
			}, capacity, maxCap, idleTimeout))
			_pool.Unlock()

			conn, err := pool.GetPoolConnection(i)
			defer pool.ReturnConnection(i, conn)
			if err != nil {
				log.Fatalf("Can't connect to memcached: %s", err)
			}
		}(pool, server)
	}

	return pool, nil
}

// ShardedServerStrategyMD5 uses md5+jump to pick a server
func ShardedServerStrategyMD5(key string, numServers int) int {
	if numServers == 1 {
		return 0
	}

	hash := md5.Sum([]byte(key))
	hashInt := big.NewInt(0)
	hashInt.SetString(hex.EncodeToString(hash[:]), 16)

	return int(jump.Hash(hashInt.Uint64(), numServers))
}

// ShardedServerStrategyFarmhash uses farmhash+jump to pick a server
func ShardedServerStrategyFarmhash(key string, numServers int) int {
	if numServers == 1 {
		return 0
	}

	return int(jump.Hash(farm.Fingerprint64([]byte(key)), numServers))
}

// HashKeyStrategyFarmhash uses farmhash to normalize key names for storage
func HashKeyStrategyFarmhash(key string) string {
	return strconv.FormatUint(farm.Fingerprint64([]byte(key)), 10)
}

// GetConnection returns a connection from the sharding pool, based on the key
func (v *Pool) GetConnection(key string) (*VitessResource, int, error) {
	poolNum := v.ServerStrategy(key, v.numServers)

	connection, err := v.GetPoolConnection(poolNum)
	if err != nil {
		return nil, -1, err
	}

	return connection, poolNum, nil
}

// GetPoolConnection returns a connection from a specific pool number
func (v *Pool) GetPoolConnection(poolNum int) (*VitessResource, error) {
	ctx := context.Background()

	v.RLock()
	resource, err := v.pool[poolNum].Get(ctx)
	v.RUnlock()
	if err != nil {
		return nil, err
	}

	connection := resource.(VitessResource)

	return &connection, nil
}

// ReturnConnection returns a connection to the pool
func (v *Pool) ReturnConnection(poolNum int, resource *VitessResource) {
	if poolNum > v.numServers || poolNum < 0 {
		log.Fatalf("error: invalid server %d (of total %d)", poolNum, v.numServers)
	}

	v.RLock()
	v.pool[poolNum].Put(*resource)
	v.RUnlock()
}

// GetKeyMapping returns a mapping of server to a list of keys, useful for Gets()
func (v *Pool) GetKeyMapping(keys ...string) map[int][]string {
	mapping := make(map[int][]string)

	for i := 0; i < v.numServers; i++ {
		mapping[i] = []string{}
	}

	for _, key := range keys {
		poolNum := v.ServerStrategy(key, v.numServers)
		mapping[poolNum] = append(mapping[poolNum], v.HashKeyStrategy(key))
	}

	return mapping
}
<|endoftext|>


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

package grumpy

import (
	"fmt"
	"os"
	"reflect"
	"runtime/pprof"
	"strings"
	"sync"
)

type moduleState int

const (
	moduleStateNew moduleState = iota
	moduleStateInitializing
	moduleStateReady
)

var (
	importMutex    sync.Mutex
	moduleRegistry = map[string]*Code{}
	// ModuleType is the object representing the Python 'module' type.
	ModuleType = newBasisType("module", reflect.TypeOf(Module{}), toModuleUnsafe, ObjectType)
	// SysModules is the global dict of imported modules, aka sys.modules.
	SysModules = NewDict()
)

// Module represents Python 'module' objects.
type Module struct {
	Object
	mutex recursiveMutex
	state moduleState
	code  *Code
}

// ModuleInit functions are called when importing Grumpy modules to execute the
// top level code for that module.
type ModuleInit func(f *Frame, m *Module) *BaseException

// RegisterModule adds the named module to the registry so that it can be
// subsequently imported.
func RegisterModule(name string, c *Code) {
	err := ""
	importMutex.Lock()
	if moduleRegistry[name] == nil {
		moduleRegistry[name] = c
	} else {
		err = "module already registered: " + name
	}
	importMutex.Unlock()
	if err != "" {
		logFatal(err)
	}
}

// ImportModule takes a fully qualified module name (e.g. a.b.c) and a slice of
// code objects where the name of the i'th module is the prefix of name
// ending in the i'th dot. The number of dot delimited parts of name must be the
// same as the number of code objects. For each successive prefix, ImportModule
// looks in sys.modules for an existing module with that name and if not
// present creates a new module object, adds it to sys.modules and initializes
// it with the corresponding code object. If the module was already present in
// sys.modules, it is not re-initialized. The returned slice contains each
// package and module initialized in this way in order.
//
// For example, ImportModule(f, "a.b", []*Code{a.Code, b.Code})
// causes the initialization and entry into sys.modules of Grumpy module a and
// then Grumpy module b. The two initialized modules are returned.
//
// If ImportModule is called in two threads concurrently to import the same
// module, both invocations will produce the same module object and the module
// is guaranteed to only be initialized once. The second invocation will not
// return the module until it is fully initialized.
func ImportModule(f *Frame, name string) ([]*Object, *BaseException) {
	if strings.Contains(name, "/") {
		o, raised := importOne(f, name)
		if raised != nil {
			return nil, raised
		}
		return []*Object{o}, nil
	}
	parts := strings.Split(name, ".")
	numParts := len(parts)
	result := make([]*Object, numParts)
	var prev *Object
	for i := 0; i < numParts; i++ {
		name := strings.Join(parts[:i+1], ".")
		o, raised := importOne(f, name)
		if raised != nil {
			return nil, raised
		}
		if prev != nil {
			if raised := SetAttr(f, prev, NewStr(parts[i]), o); raised != nil {
				return nil, raised
			}
		}
		result[i] = o
		prev = o
	}
	return result, nil
}

func importOne(f *Frame, name string) (*Object, *BaseException) {
	var c *Code
	// We do very limited locking here resulting in some
	// sys.modules consistency gotchas.
	importMutex.Lock()
	o, raised := SysModules.GetItemString(f, name)
	if raised == nil && o == nil {
		if c = moduleRegistry[name]; c == nil {
			raised = f.RaiseType(ImportErrorType, name)
		} else {
			o = newModule(name, c.filename).ToObject()
			raised = SysModules.SetItemString(f, name, o)
		}
	}
	importMutex.Unlock()
	if raised != nil {
		return nil, raised
	}
	if o.isInstance(ModuleType) {
		var raised *BaseException
		m := toModuleUnsafe(o)
		m.mutex.Lock(f)
		if m.state == moduleStateNew {
			m.state = moduleStateInitializing
			if _, raised = c.Eval(f, m.Dict(), nil, nil); raised == nil {
				m.state = moduleStateReady
			} else {
				// If the module failed to initialize
				// then before we relinquish the module
				// lock, remove it from sys.modules.
				// Threads waiting on this module will
				// fail when they don't find it in
				// sys.modules below.
				e, tb := f.ExcInfo()
				if _, raised := SysModules.DelItemString(f, name); raised != nil {
					f.RestoreExc(e, tb)
				}
			}
		}
		m.mutex.Unlock(f)
		if raised != nil {
			return nil, raised
		}
		// The result should be what's in sys.modules, not
		// necessarily the originally created module since this
		// is CPython's behavior.
		o, raised = SysModules.GetItemString(f, name)
		if raised != nil {
			return nil, raised
		}
		if o == nil {
			// This can happen in the pathological case
			// where the module clears itself from
			// sys.modules during execution and is handled
			// by CPython in PyImport_ExecCodeModuleEx in
			// import.c.
			format := "Loaded module %s not found in sys.modules"
			return nil, f.RaiseType(ImportErrorType, fmt.Sprintf(format, name))
		}
	}
	return o, nil
}

// LoadMembers scans over all the members in module
// and populates globals with them, taking __all__ into
// account.
func LoadMembers(f *Frame, module *Object) *BaseException {
	allAttr, raised := GetAttr(f, module, NewStr("__all__"), nil)
	if raised != nil && !raised.isInstance(AttributeErrorType) {
		return raised
	}
	f.RestoreExc(nil, nil)

	if raised == nil {
		raised = loadMembersFromIterable(f, module, allAttr, nil)
		if raised != nil {
			return raised
		}
		return nil
	}

	// Fall back on __dict__
	dictAttr := module.dict.ToObject()
	raised = loadMembersFromIterable(f, module, dictAttr, func(key *Object) bool {
		return strings.HasPrefix(toStrUnsafe(key).value, "_")
	})
	if raised != nil {
		return raised
	}
	return nil
}

func loadMembersFromIterable(f *Frame, module, iterable *Object, filterF func(*Object) bool) *BaseException {
	globals := f.Globals()
	raised := seqForEach(f, iterable, func(memberName *Object) *BaseException {
		if !memberName.isInstance(StrType) {
			errorMessage := fmt.Sprintf("attribute name must be string, not '%v'", memberName.typ.Name())
			return f.RaiseType(AttributeErrorType, errorMessage)
		}
		member, raised := GetAttr(f, module, toStrUnsafe(memberName), nil)
		if raised != nil {
			return raised
		}
		if filterF != nil && filterF(memberName) {
			return nil
		}
		raised = globals.SetItem(f, memberName, member)
		if raised != nil {
			return raised
		}
		return nil
	})
	return raised
}

// newModule creates a new Module object with the given fully qualified name
// (e.g a.b.c) and its corresponding Python filename and package.
func newModule(name, filename string) *Module {
	pkgName := ""
	if strings.Contains(name, ".") {
		pkgParts := strings.Split(name, ".")
		pkgName = strings.Join(pkgParts[:len(pkgParts)-1], ".")
	}

	d := newStringDict(map[string]*Object{
		"__file__":    NewStr(filename).ToObject(),
		"__name__":    NewStr(name).ToObject(),
		"__package__": NewStr(pkgName).ToObject(),
	})
	return &Module{Object: Object{typ: ModuleType, dict: d}}
}

func toModuleUnsafe(o *Object) *Module {
	return (*Module)(o.toPointer())
}

// GetFilename returns the __file__ attribute of m, raising SystemError if it
// does not exist.
func (m *Module) GetFilename(f *Frame) (*Str, *BaseException) {
	fileAttr, raised := GetAttr(f, m.ToObject(), NewStr("__file__"), None)
	if raised != nil {
		return nil, raised
	}
	if !fileAttr.isInstance(StrType) {
		return nil, f.RaiseType(SystemErrorType, "module filename missing")
	}
	return toStrUnsafe(fileAttr), nil
}

// GetName returns the __name__ attribute of m, raising SystemError if it does
// not exist.
func (m *Module) GetName(f *Frame) (*Str, *BaseException) {
	nameAttr, raised := GetAttr(f, m.ToObject(), internedName, None)
	if raised != nil {
		return nil, raised
	}
	if !nameAttr.isInstance(StrType) {
		return nil, f.RaiseType(SystemErrorType, "nameless module")
	}
	return toStrUnsafe(nameAttr), nil
}

// ToObject upcasts m to an Object.
func (m *Module) ToObject() *Object {
	return &m.Object
}

func moduleInit(f *Frame, o *Object, args Args, _ KWArgs) (*Object, *BaseException) {
	expectedTypes := []*Type{StrType, ObjectType}
	argc := len(args)
	if argc == 1 {
		expectedTypes = expectedTypes[:1]
	}
	if raised := checkFunctionArgs(f, "__init__", args, expectedTypes...); raised != nil {
		return nil, raised
	}
	if raised := SetAttr(f, o, internedName, args[0]); raised != nil {
		return nil, raised
	}
	if argc > 1 {
		if raised := SetAttr(f, o, NewStr("__doc__"), args[1]); raised != nil {
			return nil, raised
		}
	}
	return None, nil
}

func moduleRepr(f *Frame, o *Object) (*Object, *BaseException) {
	m := toModuleUnsafe(o)
	name := "?"
	nameAttr, raised := m.GetName(f)
	if raised == nil {
		name = nameAttr.Value()
	} else {
		f.RestoreExc(nil, nil)
	}
	file := "(built-in)"
	fileAttr, raised := m.GetFilename(f)
	if raised == nil {
		file = fmt.Sprintf("from '%s'", fileAttr.Value())
	} else {
		f.RestoreExc(nil, nil)
	}
	return NewStr(fmt.Sprintf("<module '%s' %s>", name, file)).ToObject(), nil
}

func initModuleType(map[string]*Object) {
	ModuleType.slots.Init = &initSlot{moduleInit}
	ModuleType.slots.Repr = &unaryOpSlot{moduleRepr}
}

// RunMain execs the given code object as a module under the name "__main__".
// It handles any exceptions raised during module execution. If no exceptions
// were raised then the return value is zero. If a SystemExit was raised then
// the return value depends on its code attribute: None -> zero, integer values
// are returned as-is. Other code values and exception types produce a return
// value of 1.
func RunMain(code *Code) int {
	if file := os.Getenv("GRUMPY_PROFILE"); file != "" {
		f, err := os.Create(file)
		if err != nil {
			logFatal(err.Error())
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			logFatal(err.Error())
		}
		defer pprof.StopCPUProfile()
	}
	m := newModule("__main__", code.filename)
	m.state = moduleStateInitializing
	f := NewRootFrame()
	f.code = code
	f.globals = m.Dict()
	if raised := SysModules.SetItemString(f, "__main__", m.ToObject()); raised != nil {
		Stderr.writeString(raised.String())
	}
	_, e := code.fn(f, nil)
	if e == nil {
		return 0
	}
	if !e.isInstance(SystemExitType) {
		Stderr.writeString(FormatExc(f))
		return 1
	}
	f.RestoreExc(nil, nil)
	o, raised := GetAttr(f, e.ToObject(), NewStr("code"), nil)
	if raised != nil {
		return 1
	}
	if o.isInstance(IntType) {
		return toIntUnsafe(o).Value()
	}
	if o == None {
		return 0
	}
	if s, raised := ToStr(f, o); raised == nil {
		Stderr.writeString(s.Value() + "\n")
	}
	return 1
}
<commit_msg>Correct the ImportError message<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grumpy

import (
	"fmt"
	"os"
	"reflect"
	"runtime/pprof"
	"strings"
	"sync"
)

type moduleState int

const (
	moduleStateNew moduleState = iota
	moduleStateInitializing
	moduleStateReady
)

var (
	importMutex    sync.Mutex
	moduleRegistry = map[string]*Code{}
	// ModuleType is the object representing the Python 'module' type.
	ModuleType = newBasisType("module", reflect.TypeOf(Module{}), toModuleUnsafe, ObjectType)
	// SysModules is the global dict of imported modules, aka sys.modules.
	SysModules = NewDict()
)

// Module represents Python 'module' objects.
type Module struct {
	Object
	mutex recursiveMutex
	state moduleState
	code  *Code
}

// ModuleInit functions are called when importing Grumpy modules to execute the
// top level code for that module.
type ModuleInit func(f *Frame, m *Module) *BaseException

// RegisterModule adds the named module to the registry so that it can be
// subsequently imported.
func RegisterModule(name string, c *Code) {
	err := ""
	importMutex.Lock()
	if moduleRegistry[name] == nil {
		moduleRegistry[name] = c
	} else {
		err = "module already registered: " + name
	}
	importMutex.Unlock()
	if err != "" {
		logFatal(err)
	}
}

// ImportModule takes a fully qualified module name (e.g. a.b.c) and a slice of
// code objects where the name of the i'th module is the prefix of name
// ending in the i'th dot. The number of dot delimited parts of name must be the
// same as the number of code objects. For each successive prefix, ImportModule
// looks in sys.modules for an existing module with that name and if not
// present creates a new module object, adds it to sys.modules and initializes
// it with the corresponding code object. If the module was already present in
// sys.modules, it is not re-initialized. The returned slice contains each
// package and module initialized in this way in order.
//
// For example, ImportModule(f, "a.b", []*Code{a.Code, b.Code})
// causes the initialization and entry into sys.modules of Grumpy module a and
// then Grumpy module b. The two initialized modules are returned.
//
// If ImportModule is called in two threads concurrently to import the same
// module, both invocations will produce the same module object and the module
// is guaranteed to only be initialized once. The second invocation will not
// return the module until it is fully initialized.
func ImportModule(f *Frame, name string) ([]*Object, *BaseException) {
	if strings.Contains(name, "/") {
		o, raised := importOne(f, name)
		if raised != nil {
			return nil, raised
		}
		return []*Object{o}, nil
	}
	parts := strings.Split(name, ".")
	numParts := len(parts)
	result := make([]*Object, numParts)
	var prev *Object
	for i := 0; i < numParts; i++ {
		name := strings.Join(parts[:i+1], ".")
		o, raised := importOne(f, name)
		if raised != nil {
			return nil, raised
		}
		if prev != nil {
			if raised := SetAttr(f, prev, NewStr(parts[i]), o); raised != nil {
				return nil, raised
			}
		}
		result[i] = o
		prev = o
	}
	return result, nil
}

func importOne(f *Frame, name string) (*Object, *BaseException) {
	var c *Code
	// We do very limited locking here resulting in some
	// sys.modules consistency gotchas.
	importMutex.Lock()
	o, raised := SysModules.GetItemString(f, name)
	if raised == nil && o == nil {
		if c = moduleRegistry[name]; c == nil {
			msg := fmt.Sprintf("No module named %s", name)
			raised = f.RaiseType(ImportErrorType, msg)
		} else {
			o = newModule(name, c.filename).ToObject()
			raised = SysModules.SetItemString(f, name, o)
		}
	}
	importMutex.Unlock()
	if raised != nil {
		return nil, raised
	}
	if o.isInstance(ModuleType) {
		var raised *BaseException
		m := toModuleUnsafe(o)
		m.mutex.Lock(f)
		if m.state == moduleStateNew {
			m.state = moduleStateInitializing
			if _, raised = c.Eval(f, m.Dict(), nil, nil); raised == nil {
				m.state = moduleStateReady
			} else {
				// If the module failed to initialize
				// then before we relinquish the module
				// lock, remove it from sys.modules.
				// Threads waiting on this module will
				// fail when they don't find it in
				// sys.modules below.
				e, tb := f.ExcInfo()
				if _, raised := SysModules.DelItemString(f, name); raised != nil {
					f.RestoreExc(e, tb)
				}
			}
		}
		m.mutex.Unlock(f)
		if raised != nil {
			return nil, raised
		}
		// The result should be what's in sys.modules, not
		// necessarily the originally created module since this
		// is CPython's behavior.
		o, raised = SysModules.GetItemString(f, name)
		if raised != nil {
			return nil, raised
		}
		if o == nil {
			// This can happen in the pathological case
			// where the module clears itself from
			// sys.modules during execution and is handled
			// by CPython in PyImport_ExecCodeModuleEx in
			// import.c.
			format := "Loaded module %s not found in sys.modules"
			return nil, f.RaiseType(ImportErrorType, fmt.Sprintf(format, name))
		}
	}
	return o, nil
}

// LoadMembers scans over all the members in module
// and populates globals with them, taking __all__ into
// account.
func LoadMembers(f *Frame, module *Object) *BaseException {
	allAttr, raised := GetAttr(f, module, NewStr("__all__"), nil)
	if raised != nil && !raised.isInstance(AttributeErrorType) {
		return raised
	}
	f.RestoreExc(nil, nil)

	if raised == nil {
		raised = loadMembersFromIterable(f, module, allAttr, nil)
		if raised != nil {
			return raised
		}
		return nil
	}

	// Fall back on __dict__
	dictAttr := module.dict.ToObject()
	raised = loadMembersFromIterable(f, module, dictAttr, func(key *Object) bool {
		return strings.HasPrefix(toStrUnsafe(key).value, "_")
	})
	if raised != nil {
		return raised
	}
	return nil
}

func loadMembersFromIterable(f *Frame, module, iterable *Object, filterF func(*Object) bool) *BaseException {
	globals := f.Globals()
	raised := seqForEach(f, iterable, func(memberName *Object) *BaseException {
		if !memberName.isInstance(StrType) {
			errorMessage := fmt.Sprintf("attribute name must be string, not '%v'", memberName.typ.Name())
			return f.RaiseType(AttributeErrorType, errorMessage)
		}
		member, raised := GetAttr(f, module, toStrUnsafe(memberName), nil)
		if raised != nil {
			return raised
		}
		if filterF != nil && filterF(memberName) {
			return nil
		}
		raised = globals.SetItem(f, memberName, member)
		if raised != nil {
			return raised
		}
		return nil
	})
	return raised
}

// newModule creates a new Module object with the given fully qualified name
// (e.g a.b.c) and its corresponding Python filename and package.
func newModule(name, filename string) *Module {
	pkgName := ""
	if strings.Contains(name, ".") {
		pkgParts := strings.Split(name, ".")
		pkgName = strings.Join(pkgParts[:len(pkgParts)-1], ".")
	}

	d := newStringDict(map[string]*Object{
		"__file__":    NewStr(filename).ToObject(),
		"__name__":    NewStr(name).ToObject(),
		"__package__": NewStr(pkgName).ToObject(),
	})
	return &Module{Object: Object{typ: ModuleType, dict: d}}
}

func toModuleUnsafe(o *Object) *Module {
	return (*Module)(o.toPointer())
}

// GetFilename returns the __file__ attribute of m, raising SystemError if it
// does not exist.
func (m *Module) GetFilename(f *Frame) (*Str, *BaseException) {
	fileAttr, raised := GetAttr(f, m.ToObject(), NewStr("__file__"), None)
	if raised != nil {
		return nil, raised
	}
	if !fileAttr.isInstance(StrType) {
		return nil, f.RaiseType(SystemErrorType, "module filename missing")
	}
	return toStrUnsafe(fileAttr), nil
}

// GetName returns the __name__ attribute of m, raising SystemError if it does
// not exist.
func (m *Module) GetName(f *Frame) (*Str, *BaseException) {
	nameAttr, raised := GetAttr(f, m.ToObject(), internedName, None)
	if raised != nil {
		return nil, raised
	}
	if !nameAttr.isInstance(StrType) {
		return nil, f.RaiseType(SystemErrorType, "nameless module")
	}
	return toStrUnsafe(nameAttr), nil
}

// ToObject upcasts m to an Object.
func (m *Module) ToObject() *Object {
	return &m.Object
}

func moduleInit(f *Frame, o *Object, args Args, _ KWArgs) (*Object, *BaseException) {
	expectedTypes := []*Type{StrType, ObjectType}
	argc := len(args)
	if argc == 1 {
		expectedTypes = expectedTypes[:1]
	}
	if raised := checkFunctionArgs(f, "__init__", args, expectedTypes...); raised != nil {
		return nil, raised
	}
	if raised := SetAttr(f, o, internedName, args[0]); raised != nil {
		return nil, raised
	}
	if argc > 1 {
		if raised := SetAttr(f, o, NewStr("__doc__"), args[1]); raised != nil {
			return nil, raised
		}
	}
	return None, nil
}

func moduleRepr(f *Frame, o *Object) (*Object, *BaseException) {
	m := toModuleUnsafe(o)
	name := "?"
	nameAttr, raised := m.GetName(f)
	if raised == nil {
		name = nameAttr.Value()
	} else {
		f.RestoreExc(nil, nil)
	}
	file := "(built-in)"
	fileAttr, raised := m.GetFilename(f)
	if raised == nil {
		file = fmt.Sprintf("from '%s'", fileAttr.Value())
	} else {
		f.RestoreExc(nil, nil)
	}
	return NewStr(fmt.Sprintf("<module '%s' %s>", name, file)).ToObject(), nil
}

func initModuleType(map[string]*Object) {
	ModuleType.slots.Init = &initSlot{moduleInit}
	ModuleType.slots.Repr = &unaryOpSlot{moduleRepr}
}

// RunMain execs the given code object as a module under the name "__main__".
// It handles any exceptions raised during module execution. If no exceptions
// were raised then the return value is zero. If a SystemExit was raised then
// the return value depends on its code attribute: None -> zero, integer values
// are returned as-is. Other code values and exception types produce a return
// value of 1.
func RunMain(code *Code) int {
	if file := os.Getenv("GRUMPY_PROFILE"); file != "" {
		f, err := os.Create(file)
		if err != nil {
			logFatal(err.Error())
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			logFatal(err.Error())
		}
		defer pprof.StopCPUProfile()
	}
	m := newModule("__main__", code.filename)
	m.state = moduleStateInitializing
	f := NewRootFrame()
	f.code = code
	f.globals = m.Dict()
	if raised := SysModules.SetItemString(f, "__main__", m.ToObject()); raised != nil {
		Stderr.writeString(raised.String())
	}
	_, e := code.fn(f, nil)
	if e == nil {
		return 0
	}
	if !e.isInstance(SystemExitType) {
		Stderr.writeString(FormatExc(f))
		return 1
	}
	f.RestoreExc(nil, nil)
	o, raised := GetAttr(f, e.ToObject(), NewStr("code"), nil)
	if raised != nil {
		return 1
	}
	if o.isInstance(IntType) {
		return toIntUnsafe(o).Value()
	}
	if o == None {
		return 0
	}
	if s, raised := ToStr(f, o); raised == nil {
		Stderr.writeString(s.Value() + "\n")
	}
	return 1
}
<|endoftext|>