<commit_before>package nethttp

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"io"
	"log"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/lestrrat/go-hsup/internal/genutil"
	"github.com/lestrrat/go-jshschema"
	"github.com/lestrrat/go-jsschema"
	"github.com/lestrrat/go-jsval"
	"github.com/lestrrat/go-pdebug"
)

type Builder struct {
	AppPkg       string
	ValidatorPkg string
}

type genctx struct {
	apppkg             string
	schema             *hschema.HyperSchema
	clientpkg          string
	validatorpkg       string
	methods            map[string]string
	methodPayloadType  map[string]string
	methodNames        []string
	pathToMethods      map[string]string
	requestValidators  map[string]*jsval.JSVal
	responseValidators map[string]*jsval.JSVal
}

func New() *Builder {
	return &Builder{
		AppPkg:       "app",
		ValidatorPkg: "validator",
	}
}

func (b *Builder) ProcessFile(f string) error {
	log.Printf(" ===> Using schema file '%s'", f)
	s, err := hschema.ReadFile(f)
	if err != nil {
		return err
	}
	return b.Process(s)
}

func (b *Builder) Process(s *hschema.HyperSchema) error {
	ctx := genctx{
		schema:             s,
		methodNames:        make([]string, len(s.Links)),
		apppkg:             b.AppPkg,
		clientpkg:          "client",
		validatorpkg:       b.ValidatorPkg,
		methods:            make(map[string]string),
		methodPayloadType:  make(map[string]string),
		pathToMethods:      make(map[string]string),
		requestValidators:  make(map[string]*jsval.JSVal),
		responseValidators: make(map[string]*jsval.JSVal),
	}

	if err := parse(&ctx, s); err != nil {
		return err
	}

	if err := generateFiles(&ctx); err != nil {
		return err
	}

	log.Printf(" <=== All files generated")
	return nil
}

func parse(ctx *genctx, s *hschema.HyperSchema) error {
	for i, link := range s.Links {
		methodName := genutil.TitleToName(link.Title)
		ctx.methodPayloadType[methodName] = "interface{}"

		// Got to do this first, because validators are used in makeMethod()
		if s := link.Schema; s != nil {
			v, err := genutil.MakeValidator(s, ctx.schema)
			if err != nil {
				return err
			}
			v.Name = fmt.Sprintf("HTTP%sRequest", methodName)
			ctx.requestValidators[methodName] = v
		}
		if s := link.TargetSchema; s != nil {
			v, err := genutil.MakeValidator(s, ctx.schema)
			if err != nil {
				return err
			}
			v.Name = fmt.Sprintf("HTTP%sResponse", methodName)
			ctx.responseValidators[methodName] = v
		}
		if ls := link.Schema; ls != nil {
			if !ls.IsResolved() {
				rs, err := ls.Resolve(ctx.schema)
				if err != nil {
					return err
				}
				ls = rs
			}
			if pdebug.Enabled {
				pdebug.Printf("checking extras for %s: %#v", link.Path(), ls.Extras)
			}
			if gt, ok := ls.Extras["gotype"]; ok {
				ctx.methodPayloadType[methodName] = gt.(string)
			}
		}
		ctx.methodNames[i] = methodName
		ctx.pathToMethods[link.Path()] = methodName
		methodBody, err := makeMethod(ctx, methodName, link)
		if err != nil {
			return err
		}
		ctx.methods[methodName] = methodBody
	}

	sort.Strings(ctx.methodNames)
	return nil
}

func makeMethod(ctx *genctx, name string, l *hschema.Link) (string, error) {
	buf := bytes.Buffer{}

	fmt.Fprintf(&buf, `func http%s(w http.ResponseWriter, r *http.Response) {`+"\n", name)
	method := strings.ToLower(l.Method)
	if method == "" {
		method = "get"
	}
	buf.WriteString("if strings.ToLower(r.Method) != `")
	fmt.Fprintf(&buf, "%s", method)
	buf.WriteString("` {\nhttp.Error(w, `Not Found`, http.StatusNotFound)\n}\n")

	if v := ctx.requestValidators[name]; v != nil {
		payloadType := ctx.methodPayloadType[name]
		if method == "get" {
			// If this is a get request, then we'd have to assemble
			// the incoming data from r.Form
			if payloadType == "interface{}" {
				buf.WriteString("\nvar payload map[string]interface{}")
				for k, v := range l.Schema.Properties {
					if !v.IsResolved() {
						rv, err := v.Resolve(ctx.schema)
						if err != nil {
							return "", err
						}
						v = rv
					}

					if len(v.Type) != 1 {
						return "", errors.New("'" + name + "' can't handle input parameters unless the type contains exactly 1 type")
					}

					switch v.Type[0] {
					case schema.IntegerType:
						fmt.Fprintf(&buf, "\nv, err := getInteger(r.Forms, %s)", strconv.Quote(k))
						fmt.Fprintf(&buf, `i
if err != nil {
	http.Error(w, "Invalid parameter " + %s, http.StatusInternalServerError)
	return
}
`, strconv.Quote(k))
					case schema.StringType:
						fmt.Fprintf(&buf, "\nv := r.Forms[%s]", strconv.Quote(k))
					}
					buf.WriteString(`
switch len(v) {
case 0:
	continue
case 1:
	payload[k] = v[0]
default:
	payload[k] = v
}
`)
				}
			}
		} else {
			buf.WriteString("\nvar payload ")
			if genutil.LooksLikeStruct(payloadType) {
				buf.WriteRune('*')
			}
			buf.WriteString(payloadType)
		}
		buf.WriteString("\nif err := json.NewDecoder(r.Body).Decode(payload); err != nil {")
		buf.WriteString("\nhttp.Error(w, `Invalid input`, http.StatusInternalServerError)")
		buf.WriteString("\nreturn")
		buf.WriteString("\n}")
		fmt.Fprintf(&buf, "\n\nif err := %s.%s.Validate(payload); err != nil {", ctx.validatorpkg, v.Name)
		buf.WriteString("\nhttp.Error(w, `Invalid input`, http.StatusInternalServerError)")
		buf.WriteString("\nreturn")
		buf.WriteString("\n}")
	}

	fmt.Fprintf(&buf, "\ndo%s(context.Background(), w, r, payload)", name)
	buf.WriteString("\n}\n")

log.Printf(buf.String())
	return buf.String(), nil
}

func generateFile(ctx *genctx, fn string, cb func(io.Writer, *genctx) error) error {
	log.Printf(" + Generating file '%s'", fn)
	f, err := genutil.CreateFile(fn)
	if err != nil {
		return err
	}
	defer f.Close()
	return cb(f, ctx)
}

func generateFiles(ctxif interface{}) error {
	ctx, ok := ctxif.(*genctx)
	if !ok {
		return errors.New("expected genctx type")
	}

	{
		fn := filepath.Join(ctx.apppkg, fmt.Sprintf("%s.go", ctx.apppkg))
		if err := generateFile(ctx, fn, generateServerCode); err != nil {
			return err
		}
	}

	{
		fn := filepath.Join(ctx.apppkg, ctx.validatorpkg, fmt.Sprintf("%s.go", ctx.validatorpkg))
		if err := generateFile(ctx, fn, generateValidatorCode); err != nil {
			return err
		}
	}

	return nil
}

func generateServerCode(out io.Writer, ctx *genctx) error {
	buf := bytes.Buffer{}

	fmt.Fprintf(&buf, "package %s\n\n", ctx.apppkg)

	genutil.WriteImports(
		&buf,
		[]string{
			"net/http",
			"strings",
		},
		[]string{
			"github.com/gorilla/mux",
			"golang.org/x/context",
		},
	)

	buf.WriteString(`
type Server struct {
	*mux.Router
}

func New() *Server {
	s := &Server{
		Router: mux.NewRouter(),
	}
	s.SetupRoutes()
	return s
}

func getInteger(v url.Values, f string) ([]int64, error) {
	x, ok := v[f]
	if !ok {
		return nil, nil
	}

	ret := make([]int64, len(x))
	for i, e := range x {
		p, err := strconv.ParseInt(e, 10, 64)
		if err != nil {
			return nil, err
		}
		ret[i] = p
	}

	return ret, nil
}

`)

	for _, methodName := range ctx.methodNames {
		buf.WriteString(ctx.methods[methodName])
		buf.WriteString("\n")
	}

	buf.WriteString("func (s *Server) SetupRoutes() {")
	buf.WriteString("\nr := s.Router")
	paths := make([]string, 0, len(ctx.pathToMethods))
	for path := range ctx.pathToMethods {
		paths = append(paths, path)
	}
	sort.Strings(paths)
	for _, path := range paths {
		method := ctx.pathToMethods[path]
		fmt.Fprintf(&buf, "\nr.HandleFunc(`%s`, %s)", path, method)
	}
	buf.WriteString("\n}\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}

	if _, err := out.Write(fsrc); err != nil {
		return err
	}
	return nil
}

func generateValidatorCode(out io.Writer, ctx *genctx) error {
	g := jsval.NewGenerator()
	validators := make([]*jsval.JSVal, 0, len(ctx.requestValidators)+len(ctx.responseValidators))
	for _, v := range ctx.requestValidators {
		validators = append(validators, v)
	}
	for _, v := range ctx.responseValidators {
		validators = append(validators, v)
	}

	buf := bytes.Buffer{}
	buf.WriteString("package " + ctx.validatorpkg + "\n\n")

	genutil.WriteImports(
		&buf,
		nil,
		[]string{
			"github.com/lestrrat/go-jsval",
		},
	)
	if err := g.Process(&buf, validators...); err != nil {
		return err
	}
	buf.WriteString("\n\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}

	if _, err := out.Write(fsrc); err != nil {
		return err
	}

	return nil
}
<commit_msg>remove debug print<commit_after>package nethttp

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"io"
	"log"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/lestrrat/go-hsup/internal/genutil"
	"github.com/lestrrat/go-jshschema"
	"github.com/lestrrat/go-jsschema"
	"github.com/lestrrat/go-jsval"
	"github.com/lestrrat/go-pdebug"
)

type Builder struct {
	AppPkg       string
	ValidatorPkg string
}

type genctx struct {
	apppkg             string
	schema             *hschema.HyperSchema
	clientpkg          string
	validatorpkg       string
	methods            map[string]string
	methodPayloadType  map[string]string
	methodNames        []string
	pathToMethods      map[string]string
	requestValidators  map[string]*jsval.JSVal
	responseValidators map[string]*jsval.JSVal
}

func New() *Builder {
	return &Builder{
		AppPkg:       "app",
		ValidatorPkg: "validator",
	}
}

func (b *Builder) ProcessFile(f string) error {
	log.Printf(" ===> Using schema file '%s'", f)
	s, err := hschema.ReadFile(f)
	if err != nil {
		return err
	}
	return b.Process(s)
}

func (b *Builder) Process(s *hschema.HyperSchema) error {
	ctx := genctx{
		schema:             s,
		methodNames:        make([]string, len(s.Links)),
		apppkg:             b.AppPkg,
		clientpkg:          "client",
		validatorpkg:       b.ValidatorPkg,
		methods:            make(map[string]string),
		methodPayloadType:  make(map[string]string),
		pathToMethods:      make(map[string]string),
		requestValidators:  make(map[string]*jsval.JSVal),
		responseValidators: make(map[string]*jsval.JSVal),
	}

	if err := parse(&ctx, s); err != nil {
		return err
	}

	if err := generateFiles(&ctx); err != nil {
		return err
	}

	log.Printf(" <=== All files generated")
	return nil
}

func parse(ctx *genctx, s *hschema.HyperSchema) error {
	for i, link := range s.Links {
		methodName := genutil.TitleToName(link.Title)
		ctx.methodPayloadType[methodName] = "interface{}"

		// Got to do this first, because validators are used in makeMethod()
		if s := link.Schema; s != nil {
			v, err := genutil.MakeValidator(s, ctx.schema)
			if err != nil {
				return err
			}
			v.Name = fmt.Sprintf("HTTP%sRequest", methodName)
			ctx.requestValidators[methodName] = v
		}
		if s := link.TargetSchema; s != nil {
			v, err := genutil.MakeValidator(s, ctx.schema)
			if err != nil {
				return err
			}
			v.Name = fmt.Sprintf("HTTP%sResponse", methodName)
			ctx.responseValidators[methodName] = v
		}
		if ls := link.Schema; ls != nil {
			if !ls.IsResolved() {
				rs, err := ls.Resolve(ctx.schema)
				if err != nil {
					return err
				}
				ls = rs
			}
			if pdebug.Enabled {
				pdebug.Printf("checking extras for %s: %#v", link.Path(), ls.Extras)
			}
			if gt, ok := ls.Extras["gotype"]; ok {
				ctx.methodPayloadType[methodName] = gt.(string)
			}
		}
		ctx.methodNames[i] = methodName
		ctx.pathToMethods[link.Path()] = methodName
		methodBody, err := makeMethod(ctx, methodName, link)
		if err != nil {
			return err
		}
		ctx.methods[methodName] = methodBody
	}

	sort.Strings(ctx.methodNames)
	return nil
}

func makeMethod(ctx *genctx, name string, l *hschema.Link) (string, error) {
	buf := bytes.Buffer{}

	fmt.Fprintf(&buf, `func http%s(w http.ResponseWriter, r *http.Response) {`+"\n", name)
	method := strings.ToLower(l.Method)
	if method == "" {
		method = "get"
	}
	buf.WriteString("if strings.ToLower(r.Method) != `")
	fmt.Fprintf(&buf, "%s", method)
	buf.WriteString("` {\nhttp.Error(w, `Not Found`, http.StatusNotFound)\n}\n")

	if v := ctx.requestValidators[name]; v != nil {
		payloadType := ctx.methodPayloadType[name]
		if method == "get" {
			// If this is a get request, then we'd have to assemble
			// the incoming data from r.Form
			if payloadType == "interface{}" {
				buf.WriteString("\nvar payload map[string]interface{}")
				for k, v := range l.Schema.Properties {
					if !v.IsResolved() {
						rv, err := v.Resolve(ctx.schema)
						if err != nil {
							return "", err
						}
						v = rv
					}

					if len(v.Type) != 1 {
						return "", errors.New("'" + name + "' can't handle input parameters unless the type contains exactly 1 type")
					}

					switch v.Type[0] {
					case schema.IntegerType:
						fmt.Fprintf(&buf, "\nv, err := getInteger(r.Forms, %s)", strconv.Quote(k))
						fmt.Fprintf(&buf, `i
if err != nil {
	http.Error(w, "Invalid parameter " + %s, http.StatusInternalServerError)
	return
}
`, strconv.Quote(k))
					case schema.StringType:
						fmt.Fprintf(&buf, "\nv := r.Forms[%s]", strconv.Quote(k))
					}
					buf.WriteString(`
switch len(v) {
case 0:
	continue
case 1:
	payload[k] = v[0]
default:
	payload[k] = v
}
`)
				}
			}
		} else {
			buf.WriteString("\nvar payload ")
			if genutil.LooksLikeStruct(payloadType) {
				buf.WriteRune('*')
			}
			buf.WriteString(payloadType)
		}
		buf.WriteString("\nif err := json.NewDecoder(r.Body).Decode(payload); err != nil {")
		buf.WriteString("\nhttp.Error(w, `Invalid input`, http.StatusInternalServerError)")
		buf.WriteString("\nreturn")
		buf.WriteString("\n}")
		fmt.Fprintf(&buf, "\n\nif err := %s.%s.Validate(payload); err != nil {", ctx.validatorpkg, v.Name)
		buf.WriteString("\nhttp.Error(w, `Invalid input`, http.StatusInternalServerError)")
		buf.WriteString("\nreturn")
		buf.WriteString("\n}")
	}

	fmt.Fprintf(&buf, "\ndo%s(context.Background(), w, r, payload)", name)
	buf.WriteString("\n}\n")

	return buf.String(), nil
}

func generateFile(ctx *genctx, fn string, cb func(io.Writer, *genctx) error) error {
	log.Printf(" + Generating file '%s'", fn)
	f, err := genutil.CreateFile(fn)
	if err != nil {
		return err
	}
	defer f.Close()
	return cb(f, ctx)
}

func generateFiles(ctxif interface{}) error {
	ctx, ok := ctxif.(*genctx)
	if !ok {
		return errors.New("expected genctx type")
	}

	{
		fn := filepath.Join(ctx.apppkg, fmt.Sprintf("%s.go", ctx.apppkg))
		if err := generateFile(ctx, fn, generateServerCode); err != nil {
			return err
		}
	}

	{
		fn := filepath.Join(ctx.apppkg, ctx.validatorpkg, fmt.Sprintf("%s.go", ctx.validatorpkg))
		if err := generateFile(ctx, fn, generateValidatorCode); err != nil {
			return err
		}
	}

	return nil
}

func generateServerCode(out io.Writer, ctx *genctx) error {
	buf := bytes.Buffer{}

	fmt.Fprintf(&buf, "package %s\n\n", ctx.apppkg)

	genutil.WriteImports(
		&buf,
		[]string{
			"net/http",
			"strings",
		},
		[]string{
			"github.com/gorilla/mux",
			"golang.org/x/context",
		},
	)

	buf.WriteString(`
type Server struct {
	*mux.Router
}

func New() *Server {
	s := &Server{
		Router: mux.NewRouter(),
	}
	s.SetupRoutes()
	return s
}

func getInteger(v url.Values, f string) ([]int64, error) {
	x, ok := v[f]
	if !ok {
		return nil, nil
	}

	ret := make([]int64, len(x))
	for i, e := range x {
		p, err := strconv.ParseInt(e, 10, 64)
		if err != nil {
			return nil, err
		}
		ret[i] = p
	}

	return ret, nil
}

`)

	for _, methodName := range ctx.methodNames {
		buf.WriteString(ctx.methods[methodName])
		buf.WriteString("\n")
	}

	buf.WriteString("func (s *Server) SetupRoutes() {")
	buf.WriteString("\nr := s.Router")
	paths := make([]string, 0, len(ctx.pathToMethods))
	for path := range ctx.pathToMethods {
		paths = append(paths, path)
	}
	sort.Strings(paths)
	for _, path := range paths {
		method := ctx.pathToMethods[path]
		fmt.Fprintf(&buf, "\nr.HandleFunc(`%s`, %s)", path, method)
	}
	buf.WriteString("\n}\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}

	if _, err := out.Write(fsrc); err != nil {
		return err
	}
	return nil
}

func generateValidatorCode(out io.Writer, ctx *genctx) error {
	g := jsval.NewGenerator()
	validators := make([]*jsval.JSVal, 0, len(ctx.requestValidators)+len(ctx.responseValidators))
	for _, v := range ctx.requestValidators {
		validators = append(validators, v)
	}
	for _, v := range ctx.responseValidators {
		validators = append(validators, v)
	}

	buf := bytes.Buffer{}
	buf.WriteString("package " + ctx.validatorpkg + "\n\n")

	genutil.WriteImports(
		&buf,
		nil,
		[]string{
			"github.com/lestrrat/go-jsval",
		},
	)
	if err := g.Process(&buf, validators...); err != nil {
		return err
	}
	buf.WriteString("\n\n")

	fsrc, err := format.Source(buf.Bytes())
	if err != nil {
		return err
	}

	if _, err := out.Write(fsrc); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

type Activating struct {
	PartialHandler
	ability *Ability
}

// NewActivating returns a Activating handler
func NewActivating(up UnitPair, ability *Ability) *Activating {
	return &Activating{
		PartialHandler: MakePartialHandler(up, ability.ActivationDuration),
		ability:        ability,
	}
}

// Ability returns the ability
func (a *Activating) Ability() *Ability {
	return a.ability
}

// OnAttach checks requirements
func (a *Activating) OnAttach() {
	a.Subject().AddEventHandler(a, EventDead)
	a.Subject().AddEventHandler(a, EventDisableInterrupt)
	a.Subject().AddEventHandler(a, EventGameTick)
	a.Subject().AddEventHandler(a, EventResourceDecreased)
	if a.Object() != nil {
		a.Object().AddEventHandler(a, EventDead)
	}
	ok := a.EverySubjectHandler(func(ha Handler) bool {
		switch ha.(type) {
		case *Activating:
			return false
		}
		return true
	})
	if !ok {
		a.Stop(a)
		return
	}
	if err := a.ability.CheckRequirements(a.Subject(), a.Object()); err != nil {
		a.Stop(a)
		return
	}
	if a.IsExpired() {
		a.perform()
		return
	}
	a.Publish(message{
	// TODO pack message
	})
}

// OnDetach removes the EventHandlers
func (a *Activating) OnDetach() {
	a.Subject().RemoveEventHandler(a, EventDead)
	a.Subject().RemoveEventHandler(a, EventDisableInterrupt)
	a.Subject().RemoveEventHandler(a, EventGameTick)
	a.Subject().RemoveEventHandler(a, EventResourceDecreased)
	if a.Object() != nil {
		a.Object().RemoveEventHandler(a, EventDead)
	}
}

// HandleEvent handles the event
func (a *Activating) HandleEvent(e Event) {
	switch e {
	case EventDead:
		a.Stop(a)
	case EventDisableInterrupt:
		a.perform()
	case EventGameTick:
		a.perform()
	case EventResourceDecreased:
		a.perform()
	}
}

// perform performs the ability
func (a *Activating) perform() {
	if err := a.ability.CheckRequirements(a.Subject(), a.Object()); err != nil {
		a.Publish(message{
		// TODO pack message
		})
		a.Stop(a)
	}
	if !a.IsExpired() {
		return
	}
	a.Publish(message{
	// TODO pack message
	})
	// TODO consume health
	// TODO consume mana
	a.ability.Perform(a.UnitPair)
	a.AttachHandler(NewCooldown(a.Subject(), a.ability))
	a.Stop(a)
}
<commit_msg>Consume unit resource on using abilities<commit_after>package main

type Activating struct {
	PartialHandler
	ability *Ability
}

// NewActivating returns a Activating handler
func NewActivating(up UnitPair, ability *Ability) *Activating {
	return &Activating{
		PartialHandler: MakePartialHandler(up, ability.ActivationDuration),
		ability:        ability,
	}
}

// Ability returns the ability
func (a *Activating) Ability() *Ability {
	return a.ability
}

// OnAttach checks requirements
func (a *Activating) OnAttach() {
	a.Subject().AddEventHandler(a, EventDead)
	a.Subject().AddEventHandler(a, EventDisableInterrupt)
	a.Subject().AddEventHandler(a, EventGameTick)
	a.Subject().AddEventHandler(a, EventResourceDecreased)
	if a.Object() != nil {
		a.Object().AddEventHandler(a, EventDead)
	}
	ok := a.EverySubjectHandler(func(ha Handler) bool {
		switch ha.(type) {
		case *Activating:
			return false
		}
		return true
	})
	if !ok {
		a.Stop(a)
		return
	}
	if err := a.ability.CheckRequirements(a.Subject(), a.Object()); err != nil {
		a.Stop(a)
		return
	}
	if a.IsExpired() {
		a.perform()
		return
	}
	a.Publish(message{
	// TODO pack message
	})
}

// OnDetach removes the EventHandlers
func (a *Activating) OnDetach() {
	a.Subject().RemoveEventHandler(a, EventDead)
	a.Subject().RemoveEventHandler(a, EventDisableInterrupt)
	a.Subject().RemoveEventHandler(a, EventGameTick)
	a.Subject().RemoveEventHandler(a, EventResourceDecreased)
	if a.Object() != nil {
		a.Object().RemoveEventHandler(a, EventDead)
	}
}

// HandleEvent handles the event
func (a *Activating) HandleEvent(e Event) {
	switch e {
	case EventDead:
		a.Stop(a)
	case EventDisableInterrupt:
		a.perform()
	case EventGameTick:
		a.perform()
	case EventResourceDecreased:
		a.perform()
	}
}

// perform performs the ability
func (a *Activating) perform() {
	if err := a.ability.CheckRequirements(a.Subject(), a.Object()); err != nil {
		a.Publish(message{
		// TODO pack message
		})
		a.Stop(a)
	}
	if !a.IsExpired() {
		return
	}
	a.Publish(message{
	// TODO pack message
	})
	a.Subject().ModifyHealth(a.ability.HealthCost)
	a.Subject().ModifyMana(a.ability.ManaCost)
	a.ability.Perform(a.UnitPair)
	a.AttachHandler(NewCooldown(a.Subject(), a.ability))
	a.Stop(a)
}
<|endoftext|>


<commit_before>package saml

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestGetSignedRequest(t *testing.T) {
	assert := assert.New(t)
	sp := ServiceProviderSettings{
		PublicCertPath:              "./default.crt",
		PrivateKeyPath:              "./default.key",
		IDPSSOURL:                   "http://www.onelogin.net",
		IDPSSODescriptorURL:         "http://www.onelogin.net",
		IDPPublicCertPath:           "./default.crt",
		AssertionConsumerServiceURL: "http://localhost:8000/auth/saml/name",
	}
	err := sp.Init()
	assert.NoError(err)

	// Construct an AuthnRequest
	authnRequest := sp.GetAuthnRequest()
	signedXML, err := authnRequest.SignedString(sp.PrivateKeyPath)
	assert.NoError(err)
	assert.NotEmpty(signedXML)

	err = VerifyRequestSignature(signedXML, sp.PublicCertPath)
	assert.NoError(err)
}
<commit_msg>simple test for unsigned requests<commit_after>package saml

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestGetSignedRequest(t *testing.T) {
	assert := assert.New(t)
	sp := ServiceProviderSettings{
		PublicCertPath:              "./default.crt",
		PrivateKeyPath:              "./default.key",
		IDPSSOURL:                   "http://www.onelogin.net",
		IDPSSODescriptorURL:         "http://www.onelogin.net",
		IDPPublicCertPath:           "./default.crt",
		AssertionConsumerServiceURL: "http://localhost:8000/auth/saml/name",
		SPSignRequest:               true,
	}
	err := sp.Init()
	assert.NoError(err)

	// Construct an AuthnRequest
	authnRequest := sp.GetAuthnRequest()
	signedXML, err := authnRequest.SignedString(sp.PrivateKeyPath)
	assert.NoError(err)
	assert.NotEmpty(signedXML)

	err = VerifyRequestSignature(signedXML, sp.PublicCertPath)
	assert.NoError(err)
}

func TestGetUnsignedRequest(t *testing.T) {
	assert := assert.New(t)
	sp := ServiceProviderSettings{
		IDPSSOURL:                   "http://www.onelogin.net",
		IDPSSODescriptorURL:         "http://www.onelogin.net",
		IDPPublicCertPath:           "./default.crt",
		AssertionConsumerServiceURL: "http://localhost:8000/auth/saml/name",
		SPSignRequest:               false,
	}
	err := sp.Init()
	assert.NoError(err)

	// Construct an AuthnRequest
	authnRequest := sp.GetAuthnRequest()
	assert.NoError(err)
	assert.NotEmpty(authnRequest)
}
<|endoftext|>


<commit_before>package authorizer

import (
	"context"

	"github.com/influxdata/influxdb"
)

var _ influxdb.SourceService = (*SourceService)(nil)

// SourceService wraps a influxdb.SourceService and authorizes actions
// against it appropriately.
type SourceService struct {
	s influxdb.SourceService
}

// NewSourceService constructs an instance of an authorizing source service.
func NewSourceService(s influxdb.SourceService) *SourceService {
	return &SourceService{
		s: s,
	}
}

func newSourcePermission(a influxdb.Action, orgID, id influxdb.ID) (*influxdb.Permission, error) {
	return influxdb.NewPermissionAtID(id, a, influxdb.SourcesResourceType, orgID)
}

func authorizeReadSource(ctx context.Context, orgID, id influxdb.ID) error {
	p, err := newSourcePermission(influxdb.ReadAction, orgID, id)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return nil
}

func authorizeWriteSource(ctx context.Context, orgID, id influxdb.ID) error {
	p, err := newSourcePermission(influxdb.WriteAction, orgID, id)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return nil
}

// DefaultSource checks to see if the authorizer on context has read access to the default source.
func (s *SourceService) DefaultSource(ctx context.Context) (*influxdb.Source, error) {
	src, err := s.s.DefaultSource(ctx)
	if err != nil {
		return nil, err
	}

	if err := authorizeReadSource(ctx, src.OrganizationID, src.ID); err != nil {
		return nil, err
	}

	return src, nil
}

// FindSourceByID checks to see if the authorizer on context has read access to the id provided.
func (s *SourceService) FindSourceByID(ctx context.Context, id influxdb.ID) (*influxdb.Source, error) {
	src, err := s.s.FindSourceByID(ctx, id)
	if err != nil {
		return nil, err
	}

	if err := authorizeReadSource(ctx, src.OrganizationID, id); err != nil {
		return nil, err
	}

	return src, nil
}

// FindSources retrieves all sources that match the provided options and then filters the list down to only the resources that are authorized.
func (s *SourceService) FindSources(ctx context.Context, opts influxdb.FindOptions) ([]*influxdb.Source, int, error) {
	// TODO: we'll likely want to push this operation into the database since fetching the whole list of data will likely be expensive.
	ss, _, err := s.s.FindSources(ctx, opts)
	if err != nil {
		return nil, 0, err
	}

	// This filters without allocating
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	sources := ss[:0]
	for _, src := range ss {
		err := authorizeReadSource(ctx, src.OrganizationID, src.ID)
		if err != nil && influxdb.ErrorCode(err) != influxdb.EUnauthorized {
			return nil, 0, err
		}

		if influxdb.ErrorCode(err) == influxdb.EUnauthorized {
			continue
		}

		sources = append(sources, src)
	}

	return sources, len(sources), nil
}

// CreateSource checks to see if the authorizer on context has write access to the global source resource.
func (s *SourceService) CreateSource(ctx context.Context, src *influxdb.Source) error {
	p, err := influxdb.NewPermission(influxdb.WriteAction, influxdb.SourcesResourceType, src.OrganizationID)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return s.s.CreateSource(ctx, src)
}

// UpdateSource checks to see if the authorizer on context has write access to the source provided.
func (s *SourceService) UpdateSource(ctx context.Context, id influxdb.ID, upd influxdb.SourceUpdate) (*influxdb.Source, error) {
	src, err := s.FindSourceByID(ctx, id)
	if err != nil {
		return nil, err
	}

	if err := authorizeWriteSource(ctx, src.OrganizationID, id); err != nil {
		return nil, err
	}

	return s.s.UpdateSource(ctx, id, upd)
}

// DeleteSource checks to see if the authorizer on context has write access to the source provided.
func (s *SourceService) DeleteSource(ctx context.Context, id influxdb.ID) error {
	m, err := s.FindSourceByID(ctx, id)
	if err != nil {
		return err
	}

	if err := authorizeWriteSource(ctx, m.OrganizationID, id); err != nil {
		return err
	}

	return s.s.DeleteSource(ctx, id)
}
<commit_msg>fix(authorizer): do not require read action in order to perform a write<commit_after>package authorizer

import (
	"context"

	"github.com/influxdata/influxdb"
)

var _ influxdb.SourceService = (*SourceService)(nil)

// SourceService wraps a influxdb.SourceService and authorizes actions
// against it appropriately.
type SourceService struct {
	s influxdb.SourceService
}

// NewSourceService constructs an instance of an authorizing source service.
func NewSourceService(s influxdb.SourceService) *SourceService {
	return &SourceService{
		s: s,
	}
}

func newSourcePermission(a influxdb.Action, orgID, id influxdb.ID) (*influxdb.Permission, error) {
	return influxdb.NewPermissionAtID(id, a, influxdb.SourcesResourceType, orgID)
}

func authorizeReadSource(ctx context.Context, orgID, id influxdb.ID) error {
	p, err := newSourcePermission(influxdb.ReadAction, orgID, id)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return nil
}

func authorizeWriteSource(ctx context.Context, orgID, id influxdb.ID) error {
	p, err := newSourcePermission(influxdb.WriteAction, orgID, id)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return nil
}

// DefaultSource checks to see if the authorizer on context has read access to the default source.
func (s *SourceService) DefaultSource(ctx context.Context) (*influxdb.Source, error) {
	src, err := s.s.DefaultSource(ctx)
	if err != nil {
		return nil, err
	}

	if err := authorizeReadSource(ctx, src.OrganizationID, src.ID); err != nil {
		return nil, err
	}

	return src, nil
}

// FindSourceByID checks to see if the authorizer on context has read access to the id provided.
func (s *SourceService) FindSourceByID(ctx context.Context, id influxdb.ID) (*influxdb.Source, error) {
	src, err := s.s.FindSourceByID(ctx, id)
	if err != nil {
		return nil, err
	}

	if err := authorizeReadSource(ctx, src.OrganizationID, id); err != nil {
		return nil, err
	}

	return src, nil
}

// FindSources retrieves all sources that match the provided options and then filters the list down to only the resources that are authorized.
func (s *SourceService) FindSources(ctx context.Context, opts influxdb.FindOptions) ([]*influxdb.Source, int, error) {
	// TODO: we'll likely want to push this operation into the database since fetching the whole list of data will likely be expensive.
	ss, _, err := s.s.FindSources(ctx, opts)
	if err != nil {
		return nil, 0, err
	}

	// This filters without allocating
	// https://github.com/golang/go/wiki/SliceTricks#filtering-without-allocating
	sources := ss[:0]
	for _, src := range ss {
		err := authorizeReadSource(ctx, src.OrganizationID, src.ID)
		if err != nil && influxdb.ErrorCode(err) != influxdb.EUnauthorized {
			return nil, 0, err
		}

		if influxdb.ErrorCode(err) == influxdb.EUnauthorized {
			continue
		}

		sources = append(sources, src)
	}

	return sources, len(sources), nil
}

// CreateSource checks to see if the authorizer on context has write access to the global source resource.
func (s *SourceService) CreateSource(ctx context.Context, src *influxdb.Source) error {
	p, err := influxdb.NewPermission(influxdb.WriteAction, influxdb.SourcesResourceType, src.OrganizationID)
	if err != nil {
		return err
	}

	if err := IsAllowed(ctx, *p); err != nil {
		return err
	}

	return s.s.CreateSource(ctx, src)
}

// UpdateSource checks to see if the authorizer on context has write access to the source provided.
func (s *SourceService) UpdateSource(ctx context.Context, id influxdb.ID, upd influxdb.SourceUpdate) (*influxdb.Source, error) {
	src, err := s.s.FindSourceByID(ctx, id)
	if err != nil {
		return nil, err
	}

	if err := authorizeWriteSource(ctx, src.OrganizationID, id); err != nil {
		return nil, err
	}

	return s.s.UpdateSource(ctx, id, upd)
}

// DeleteSource checks to see if the authorizer on context has write access to the source provided.
func (s *SourceService) DeleteSource(ctx context.Context, id influxdb.ID) error {
	m, err := s.s.FindSourceByID(ctx, id)
	if err != nil {
		return err
	}

	if err := authorizeWriteSource(ctx, m.OrganizationID, id); err != nil {
		return err
	}

	return s.s.DeleteSource(ctx, id)
}
<|endoftext|>


<commit_before>package securesocket

import (
	"encoding/binary"
	"io"
	"net"
)

type Conn struct {
	net.Conn
	*Cipher
	readBuf  []byte
	writeBuf []byte
}

func NewConn(c net.Conn, cipher *Cipher) *Conn {
	return &Conn{
		Conn:     c,
		Cipher:   cipher,
		readBuf:  leakyBuf.Get(),
		writeBuf: leakyBuf.Get(),
	}
}

func (c *Conn) Close() error {
	leakyBuf.Put(c.readBuf)
	leakyBuf.Put(c.writeBuf)
	return c.Conn.Close()
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.dec == nil {
		iv := make([]byte, c.info.ivLen)
		if _, err = io.ReadFull(c.Conn, iv); err != nil {
			return
		}
		if err = c.initDecrypt(iv); err != nil {
			return
		}
	}

	cipherData := c.readBuf
	if len(b) > len(cipherData) {
		cipherData = make([]byte, len(b))
	} else {
		cipherData = cipherData[:len(b)]
	}

	n, err = c.Conn.Read(cipherData)
	if n > 0 {
		c.decrypt(b[0:n], cipherData[0:n])
	}
	return
}

func (c *Conn) Write(b []byte) (n int, err error) {
	var iv []byte
	if c.enc == nil {
		iv, err = c.initEncrypt()
		if err != nil {
			return
		}
	}

	cipherData := c.writeBuf
	dataSize := len(b) + len(iv)
	if dataSize > len(cipherData) {
		cipherData = make([]byte, dataSize)
	} else {
		cipherData = cipherData[:dataSize]
	}

	if iv != nil {
		// Put initialization vector in buffer, do a single write to send both
		// iv and data.
		copy(cipherData, iv)
	}

	c.encrypt(cipherData[len(iv):], b)
	n, err = c.Conn.Write(cipherData)
	return
}

type PacketStreamConn struct {
	net.Conn
}

func NewPacketStreamConn(c net.Conn) *PacketStreamConn {
	return &PacketStreamConn{c}
}

func (c *PacketStreamConn) Write(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	size = uint16(len(b))
	binary.BigEndian.PutUint16(tmp[:], size)
	n, err = c.Conn.Write(append(tmp[:], b...))
	n -= 2
	return
}

func (c *PacketStreamConn) Read(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	n, err = io.ReadFull(c.Conn, tmp[:2])
	if err != nil {
		return 0, err
	}
	size = binary.BigEndian.Uint16(tmp[:])
	n, err = io.ReadFull(c.Conn, b[:int(size)])
	return
}
<commit_msg>Add shor buffer err when the size is too large to buffer.<commit_after>package securesocket

import (
	"encoding/binary"
	"io"
	"net"
)

type Conn struct {
	net.Conn
	*Cipher
	readBuf  []byte
	writeBuf []byte
}

func NewConn(c net.Conn, cipher *Cipher) *Conn {
	return &Conn{
		Conn:     c,
		Cipher:   cipher,
		readBuf:  leakyBuf.Get(),
		writeBuf: leakyBuf.Get(),
	}
}

func (c *Conn) Close() error {
	leakyBuf.Put(c.readBuf)
	leakyBuf.Put(c.writeBuf)
	return c.Conn.Close()
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.dec == nil {
		iv := make([]byte, c.info.ivLen)
		if _, err = io.ReadFull(c.Conn, iv); err != nil {
			return
		}
		if err = c.initDecrypt(iv); err != nil {
			return
		}
	}

	cipherData := c.readBuf
	if len(b) > len(cipherData) {
		cipherData = make([]byte, len(b))
	} else {
		cipherData = cipherData[:len(b)]
	}

	n, err = c.Conn.Read(cipherData)
	if n > 0 {
		c.decrypt(b[0:n], cipherData[0:n])
	}
	return
}

func (c *Conn) Write(b []byte) (n int, err error) {
	var iv []byte
	if c.enc == nil {
		iv, err = c.initEncrypt()
		if err != nil {
			return
		}
	}

	cipherData := c.writeBuf
	dataSize := len(b) + len(iv)
	if dataSize > len(cipherData) {
		cipherData = make([]byte, dataSize)
	} else {
		cipherData = cipherData[:dataSize]
	}

	if iv != nil {
		// Put initialization vector in buffer, do a single write to send both
		// iv and data.
		copy(cipherData, iv)
	}

	c.encrypt(cipherData[len(iv):], b)
	n, err = c.Conn.Write(cipherData)
	return
}

type PacketStreamConn struct {
	net.Conn
}

func NewPacketStreamConn(c net.Conn) *PacketStreamConn {
	return &PacketStreamConn{c}
}

func (c *PacketStreamConn) Write(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	size = uint16(len(b))
	binary.BigEndian.PutUint16(tmp[:], size)
	n, err = c.Conn.Write(append(tmp[:], b...))
	n -= 2
	return
}

func (c *PacketStreamConn) Read(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	n, err = io.ReadFull(c.Conn, tmp[:2])
	if err != nil {
		return 0, err
	}
	size = binary.BigEndian.Uint16(tmp[:])
	if int(size) > len(b) {
		return 0, io.ErrShortBuffer
	}
	n, err = io.ReadFull(c.Conn, b[:int(size)])
	return
}
<|endoftext|>


<commit_before>package securesocket

import (
	"encoding/binary"
	"io"
	"net"
)

type Conn struct {
	net.Conn
	*Cipher
	readBuf  []byte
	writeBuf []byte
}

func NewConn(c net.Conn, cipher *Cipher) *Conn {
	return &Conn{
		Conn:     c,
		Cipher:   cipher,
		readBuf:  leakyBuf.Get(),
		writeBuf: leakyBuf.Get(),
	}
}

func (c *Conn) Close() error {
	leakyBuf.Put(c.readBuf)
	leakyBuf.Put(c.writeBuf)
	return c.Conn.Close()
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.dec == nil {
		iv := make([]byte, c.info.ivLen)
		if _, err = io.ReadFull(c.Conn, iv); err != nil {
			return
		}
		if err = c.initDecrypt(iv); err != nil {
			return
		}
	}

	cipherData := c.readBuf
	if len(b) > len(cipherData) {
		cipherData = make([]byte, len(b))
	} else {
		cipherData = cipherData[:len(b)]
	}

	n, err = c.Conn.Read(cipherData)
	if n > 0 {
		c.decrypt(b[0:n], cipherData[0:n])
	}
	return
}

func (c *Conn) Write(b []byte) (n int, err error) {
	var iv []byte
	if c.enc == nil {
		iv, err = c.initEncrypt()
		if err != nil {
			return
		}
	}

	cipherData := c.writeBuf
	dataSize := len(b) + len(iv)
	if dataSize > len(cipherData) {
		cipherData = make([]byte, dataSize)
	} else {
		cipherData = cipherData[:dataSize]
	}

	if iv != nil {
		// Put initialization vector in buffer, do a single write to send both
		// iv and data.
		copy(cipherData, iv)
	}

	c.encrypt(cipherData[len(iv):], b)
	n, err = c.Conn.Write(cipherData)
	return
}

type PacketStreamConn struct {
	net.Conn
}

func NewPacketStreamConn(c net.Conn) *PacketStreamConn {
	return &PacketStreamConn{c}
}

func (c *PacketStreamConn) Write(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	size = uint16(len(b))
	binary.BigEndian.PutUint16(tmp[:], size)
	n, err = c.Conn.Write(append(tmp[:], b...))
	n -= 2
	return
}

func (c *PacketStreamConn) Read(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	n, err = io.ReadFull(c.Conn, tmp[:2])
	if err != nil {
		return 0, err
	}
	size = binary.BigEndian.Uint16(tmp[:])
	if int(size) > len(b) {
		return 0, io.ErrShortBuffer
	}
	n, err = io.ReadFull(c.Conn, b[:int(size)])
	return
}
<commit_msg>Some thread safety work.<commit_after>package securesocket

import (
	"encoding/binary"
	"io"
	"net"
	"sync"
)

type Conn struct {
	declock sync.Mutex
	enclock sync.Mutex
	net.Conn
	*Cipher
	readBuf  []byte
	writeBuf []byte
}

func NewConn(c net.Conn, cipher *Cipher) *Conn {
	return &Conn{
		Conn:     c,
		Cipher:   cipher,
		readBuf:  leakyBuf.Get(),
		writeBuf: leakyBuf.Get(),
	}
}

func (c *Conn) Close() error {
	leakyBuf.Put(c.readBuf)
	leakyBuf.Put(c.writeBuf)
	return c.Conn.Close()
}

func (c *Conn) Read(b []byte) (n int, err error) {
	if c.dec == nil {
		c.declock.Lock()
		iv := make([]byte, c.info.ivLen)
		if _, err = io.ReadFull(c.Conn, iv); err != nil {
			c.declock.Unlock()
			return
		}
		if err = c.initDecrypt(iv); err != nil {
			c.declock.Unlock()
			return
		}
		c.declock.Unlock()
	}

	cipherData := c.readBuf
	if len(b) > len(cipherData) {
		cipherData = make([]byte, len(b))
	} else {
		cipherData = cipherData[:len(b)]
	}

	n, err = c.Conn.Read(cipherData)
	if n > 0 {
		c.declock.Lock()
		c.decrypt(b[0:n], cipherData[0:n])
		c.declock.Unlock()
	}
	return
}

func (c *Conn) Write(b []byte) (n int, err error) {
	var iv []byte
	if c.enc == nil {
		c.enclock.Lock()
		iv, err = c.initEncrypt()
		if err != nil {
			c.enclock.Unlock()
			return
		}
		c.enclock.Unlock()
	}

	cipherData := c.writeBuf
	dataSize := len(b) + len(iv)
	if dataSize > len(cipherData) {
		cipherData = make([]byte, dataSize)
	} else {
		cipherData = cipherData[:dataSize]
	}

	if iv != nil {
		// Put initialization vector in buffer, do a single write to send both
		// iv and data.
		copy(cipherData, iv)
	}
	c.enclock.Lock()
	c.encrypt(cipherData[len(iv):], b)
	c.enclock.Unlock()
	n, err = c.Conn.Write(cipherData)
	return
}

type PacketStreamConn struct {
	net.Conn
}

func NewPacketStreamConn(c net.Conn) *PacketStreamConn {
	return &PacketStreamConn{c}
}

func (c *PacketStreamConn) Write(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	size = uint16(len(b))
	binary.BigEndian.PutUint16(tmp[:], size)
	n, err = c.Conn.Write(append(tmp[:], b...))
	n -= 2
	return
}

func (c *PacketStreamConn) Read(b []byte) (n int, err error) {
	var size uint16
	var tmp [2]byte
	n, err = io.ReadFull(c.Conn, tmp[:2])
	if err != nil {
		return 0, err
	}
	size = binary.BigEndian.Uint16(tmp[:])
	if int(size) > len(b) {
		return 0, io.ErrShortBuffer
	}
	n, err = io.ReadFull(c.Conn, b[:int(size)])
	return
}
<|endoftext|>


<commit_before>package autopilot

import (
	"sync"
	"sync/atomic"

	"github.com/btcsuite/btcd/btcec"
	"github.com/lightningnetwork/lnd/lnwallet"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/routing"
)

// ManagerCfg houses a set of values and methods that is passed to the Manager
// for it to properly manage its autopilot agent.
type ManagerCfg struct {
	// Self is the public key of the lnd instance. It is used to making
	// sure the autopilot is not opening channels to itself.
	Self *btcec.PublicKey

	// PilotCfg is the config of the autopilot agent managed by the
	// Manager.
	PilotCfg *Config

	// ChannelState is a function closure that returns the current set of
	// channels managed by this node.
	ChannelState func() ([]Channel, error)

	// SubscribeTransactions is used to get a subscription for transactions
	// relevant to this node's wallet.
	SubscribeTransactions func() (lnwallet.TransactionSubscription, error)

	// SubscribeTopology is used to get a subscription for topology changes
	// on the network.
	SubscribeTopology func() (*routing.TopologyClient, error)
}

// Manager is struct that manages an autopilot agent, making it possible to
// enable and disable it at will, and hand it relevant external information.
// It implements the autopilot grpc service, which is used to get data about
// the running autopilot, and give it relevant information.
type Manager struct {
	started uint32 // To be used atomically.
	stopped uint32 // To be used atomically.

	cfg *ManagerCfg

	// pilot is the current autopilot agent. It will be nil if the agent is
	// disabled.
	pilot *Agent

	quit chan struct{}
	wg   sync.WaitGroup
	sync.Mutex
}

// NewManager creates a new instance of the Manager from the passed config.
func NewManager(cfg *ManagerCfg) (*Manager, error) {
	return &Manager{
		cfg:  cfg,
		quit: make(chan struct{}),
	}, nil
}

// Start starts the Manager.
func (m *Manager) Start() error {
	if !atomic.CompareAndSwapUint32(&m.started, 0, 1) {
		return nil
	}

	return nil
}

// Stop stops the Manager. If an autopilot agent is active, it will also be
// stopped.
func (m *Manager) Stop() error {
	if !atomic.CompareAndSwapUint32(&m.stopped, 0, 1) {
		return nil
	}

	if err := m.StopAgent(); err != nil {
		log.Errorf("Unable to stop pilot: %v", err)
	}

	close(m.quit)
	m.wg.Wait()

	return nil
}

// IsActive returns whether the autopilot agent is currently active.
func (m *Manager) IsActive() bool {
	m.Lock()
	defer m.Unlock()

	return m.pilot != nil
}

// StartAgent creates and starts an autopilot agent from the Manager's
// config.
func (m *Manager) StartAgent() error {
	m.Lock()
	defer m.Unlock()

	// Already active.
	if m.pilot != nil {
		return nil
	}

	// Next, we'll fetch the current state of open channels from the
	// database to use as initial state for the auto-pilot agent.
	initialChanState, err := m.cfg.ChannelState()
	if err != nil {
		return err
	}

	// Now that we have all the initial dependencies, we can create the
	// auto-pilot instance itself.
	pilot, err := New(*m.cfg.PilotCfg, initialChanState)
	if err != nil {
		return err
	}

	if err := pilot.Start(); err != nil {
		return err
	}

	// Finally, we'll need to subscribe to two things: incoming
	// transactions that modify the wallet's balance, and also any graph
	// topology updates.
	txnSubscription, err := m.cfg.SubscribeTransactions()
	if err != nil {
		pilot.Stop()
		return err
	}
	graphSubscription, err := m.cfg.SubscribeTopology()
	if err != nil {
		txnSubscription.Cancel()
		pilot.Stop()
		return err
	}

	m.pilot = pilot

	// We'll launch a goroutine to provide the agent with notifications
	// whenever the balance of the wallet changes.
	// TODO(halseth): can lead to panic if in process of shutting down.
	m.wg.Add(1)
	go func() {
		defer txnSubscription.Cancel()
		defer m.wg.Done()

		for {
			select {
			case <-txnSubscription.ConfirmedTransactions():
				pilot.OnBalanceChange()

			// We won't act upon new unconfirmed transaction, as
			// we'll only use confirmed outputs when funding.
			// However, we will still drain this request in order
			// to avoid goroutine leaks, and ensure we promptly
			// read from the channel if available.
			case <-txnSubscription.UnconfirmedTransactions():
			case <-pilot.quit:
				return
			case <-m.quit:
				return
			}
		}

	}()

	// We'll also launch a goroutine to provide the agent with
	// notifications for when the graph topology controlled by the node
	// changes.
	m.wg.Add(1)
	go func() {
		defer graphSubscription.Cancel()
		defer m.wg.Done()

		for {
			select {
			case topChange, ok := <-graphSubscription.TopologyChanges:
				// If the router is shutting down, then we will
				// as well.
				if !ok {
					return
				}

				for _, edgeUpdate := range topChange.ChannelEdgeUpdates {
					// If this isn't an advertisement by
					// the backing lnd node, then we'll
					// continue as we only want to add
					// channels that we've created
					// ourselves.
					if !edgeUpdate.AdvertisingNode.IsEqual(m.cfg.Self) {
						continue
					}

					// If this is indeed a channel we
					// opened, then we'll convert it to the
					// autopilot.Channel format, and notify
					// the pilot of the new channel.
					chanNode := NewNodeID(
						edgeUpdate.ConnectingNode,
					)
					chanID := lnwire.NewShortChanIDFromInt(
						edgeUpdate.ChanID,
					)
					edge := Channel{
						ChanID:   chanID,
						Capacity: edgeUpdate.Capacity,
						Node:     chanNode,
					}
					pilot.OnChannelOpen(edge)
				}

				// For each closed channel, we'll obtain
				// the chanID of the closed channel and send it
				// to the pilot.
				for _, chanClose := range topChange.ClosedChannels {
					chanID := lnwire.NewShortChanIDFromInt(
						chanClose.ChanID,
					)

					pilot.OnChannelClose(chanID)
				}

				// If new nodes were added to the graph, or nod
				// information has changed, we'll poke autopilot
				// to see if it can make use of them.
				if len(topChange.NodeUpdates) > 0 {
					pilot.OnNodeUpdates()
				}

			case <-pilot.quit:
				return
			case <-m.quit:
				return
			}
		}
	}()

	log.Debugf("Manager started autopilot agent")

	return nil
}

// StopAgent stops any active autopilot agent.
func (m *Manager) StopAgent() error {
	m.Lock()
	defer m.Unlock()

	// Not active, so we can return early.
	if m.pilot == nil {
		return nil
	}

	if err := m.pilot.Stop(); err != nil {
		return err
	}

	// Make sure to nil the current agent, indicating it is no longer
	// active.
	m.pilot = nil

	log.Debugf("Manager stopped autopilot agent")

	return nil
}
<commit_msg>autopilot/manager: add QueryHeuristics<commit_after>package autopilot

import (
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/btcsuite/btcd/btcec"
	"github.com/lightningnetwork/lnd/lnwallet"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/routing"
)

// ManagerCfg houses a set of values and methods that is passed to the Manager
// for it to properly manage its autopilot agent.
type ManagerCfg struct {
	// Self is the public key of the lnd instance. It is used to making
	// sure the autopilot is not opening channels to itself.
	Self *btcec.PublicKey

	// PilotCfg is the config of the autopilot agent managed by the
	// Manager.
	PilotCfg *Config

	// ChannelState is a function closure that returns the current set of
	// channels managed by this node.
	ChannelState func() ([]Channel, error)

	// SubscribeTransactions is used to get a subscription for transactions
	// relevant to this node's wallet.
	SubscribeTransactions func() (lnwallet.TransactionSubscription, error)

	// SubscribeTopology is used to get a subscription for topology changes
	// on the network.
	SubscribeTopology func() (*routing.TopologyClient, error)
}

// Manager is struct that manages an autopilot agent, making it possible to
// enable and disable it at will, and hand it relevant external information.
// It implements the autopilot grpc service, which is used to get data about
// the running autopilot, and give it relevant information.
type Manager struct {
	started uint32 // To be used atomically.
	stopped uint32 // To be used atomically.

	cfg *ManagerCfg

	// pilot is the current autopilot agent. It will be nil if the agent is
	// disabled.
	pilot *Agent

	quit chan struct{}
	wg   sync.WaitGroup
	sync.Mutex
}

// NewManager creates a new instance of the Manager from the passed config.
func NewManager(cfg *ManagerCfg) (*Manager, error) {
	return &Manager{
		cfg:  cfg,
		quit: make(chan struct{}),
	}, nil
}

// Start starts the Manager.
func (m *Manager) Start() error {
	if !atomic.CompareAndSwapUint32(&m.started, 0, 1) {
		return nil
	}

	return nil
}

// Stop stops the Manager. If an autopilot agent is active, it will also be
// stopped.
func (m *Manager) Stop() error {
	if !atomic.CompareAndSwapUint32(&m.stopped, 0, 1) {
		return nil
	}

	if err := m.StopAgent(); err != nil {
		log.Errorf("Unable to stop pilot: %v", err)
	}

	close(m.quit)
	m.wg.Wait()

	return nil
}

// IsActive returns whether the autopilot agent is currently active.
func (m *Manager) IsActive() bool {
	m.Lock()
	defer m.Unlock()

	return m.pilot != nil
}

// StartAgent creates and starts an autopilot agent from the Manager's
// config.
func (m *Manager) StartAgent() error {
	m.Lock()
	defer m.Unlock()

	// Already active.
	if m.pilot != nil {
		return nil
	}

	// Next, we'll fetch the current state of open channels from the
	// database to use as initial state for the auto-pilot agent.
	initialChanState, err := m.cfg.ChannelState()
	if err != nil {
		return err
	}

	// Now that we have all the initial dependencies, we can create the
	// auto-pilot instance itself.
	pilot, err := New(*m.cfg.PilotCfg, initialChanState)
	if err != nil {
		return err
	}

	if err := pilot.Start(); err != nil {
		return err
	}

	// Finally, we'll need to subscribe to two things: incoming
	// transactions that modify the wallet's balance, and also any graph
	// topology updates.
	txnSubscription, err := m.cfg.SubscribeTransactions()
	if err != nil {
		pilot.Stop()
		return err
	}
	graphSubscription, err := m.cfg.SubscribeTopology()
	if err != nil {
		txnSubscription.Cancel()
		pilot.Stop()
		return err
	}

	m.pilot = pilot

	// We'll launch a goroutine to provide the agent with notifications
	// whenever the balance of the wallet changes.
	// TODO(halseth): can lead to panic if in process of shutting down.
	m.wg.Add(1)
	go func() {
		defer txnSubscription.Cancel()
		defer m.wg.Done()

		for {
			select {
			case <-txnSubscription.ConfirmedTransactions():
				pilot.OnBalanceChange()

			// We won't act upon new unconfirmed transaction, as
			// we'll only use confirmed outputs when funding.
			// However, we will still drain this request in order
			// to avoid goroutine leaks, and ensure we promptly
			// read from the channel if available.
			case <-txnSubscription.UnconfirmedTransactions():
			case <-pilot.quit:
				return
			case <-m.quit:
				return
			}
		}

	}()

	// We'll also launch a goroutine to provide the agent with
	// notifications for when the graph topology controlled by the node
	// changes.
	m.wg.Add(1)
	go func() {
		defer graphSubscription.Cancel()
		defer m.wg.Done()

		for {
			select {
			case topChange, ok := <-graphSubscription.TopologyChanges:
				// If the router is shutting down, then we will
				// as well.
				if !ok {
					return
				}

				for _, edgeUpdate := range topChange.ChannelEdgeUpdates {
					// If this isn't an advertisement by
					// the backing lnd node, then we'll
					// continue as we only want to add
					// channels that we've created
					// ourselves.
					if !edgeUpdate.AdvertisingNode.IsEqual(m.cfg.Self) {
						continue
					}

					// If this is indeed a channel we
					// opened, then we'll convert it to the
					// autopilot.Channel format, and notify
					// the pilot of the new channel.
					chanNode := NewNodeID(
						edgeUpdate.ConnectingNode,
					)
					chanID := lnwire.NewShortChanIDFromInt(
						edgeUpdate.ChanID,
					)
					edge := Channel{
						ChanID:   chanID,
						Capacity: edgeUpdate.Capacity,
						Node:     chanNode,
					}
					pilot.OnChannelOpen(edge)
				}

				// For each closed channel, we'll obtain
				// the chanID of the closed channel and send it
				// to the pilot.
				for _, chanClose := range topChange.ClosedChannels {
					chanID := lnwire.NewShortChanIDFromInt(
						chanClose.ChanID,
					)

					pilot.OnChannelClose(chanID)
				}

				// If new nodes were added to the graph, or nod
				// information has changed, we'll poke autopilot
				// to see if it can make use of them.
				if len(topChange.NodeUpdates) > 0 {
					pilot.OnNodeUpdates()
				}

			case <-pilot.quit:
				return
			case <-m.quit:
				return
			}
		}
	}()

	log.Debugf("Manager started autopilot agent")

	return nil
}

// StopAgent stops any active autopilot agent.
func (m *Manager) StopAgent() error {
	m.Lock()
	defer m.Unlock()

	// Not active, so we can return early.
	if m.pilot == nil {
		return nil
	}

	if err := m.pilot.Stop(); err != nil {
		return err
	}

	// Make sure to nil the current agent, indicating it is no longer
	// active.
	m.pilot = nil

	log.Debugf("Manager stopped autopilot agent")

	return nil
}

// QueryHeuristics queries the active autopilot agent for node scores.
func (m *Manager) QueryHeuristics(nodes []NodeID) (HeuristicScores, error) {
	m.Lock()
	defer m.Unlock()

	// Not active, so we can return early.
	if m.pilot == nil {
		return nil, fmt.Errorf("autopilot not active")
	}

	n := make(map[NodeID]struct{})
	for _, node := range nodes {
		n[node] = struct{}{}
	}

	log.Debugf("Querying heuristics for %d nodes", len(n))
	return m.pilot.queryHeuristics(n)
}
<|endoftext|>


<commit_before>// Copyright 2016 Ayke van Laethem.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.txt file.

// On Linux, it parses /proc/self/mountinfo. Then it can find the mount point
// based on the st_dev field of the stat result (os.FileInfo.Sys()), and
// determine filesystem capabilities from the filesystem type.
package osfs

import (
	"bufio"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"golang.org/x/sys/unix"
)

// MOUNTINFO_PATH is the common path for the mountinfo file on Linux.
const MOUNTINFO_PATH = "/proc/self/mountinfo"

// From the kernel docs at
// https://www.kernel.org/doc/Documentation/filesystems/proc.txt
//
//   This file contains lines of the form:
//
//   36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue
//   (1)(2)(3)   (4)   (5)      (6)      (7)   (8) (9)   (10)         (11)
//
//   (1) mount ID:  unique identifier of the mount (may be reused after umount)
//   (2) parent ID:  ID of parent (or of self for the top of the mount tree)
//   (3) major:minor:  value of st_dev for files on filesystem
//   (4) root:  root of the mount within the filesystem
//   (5) mount point:  mount point relative to the process's root
//   (6) mount options:  per mount options
//   (7) optional fields:  zero or more fields of the form "tag[:value]"
//   (8) separator:  marks the end of the optional fields
//   (9) filesystem type:  name of filesystem of the form "type[.subtype]"
//   (10) mount source:  filesystem specific information or "none"
//   (11) super options:  per super block options

func defaultFilesystem() Filesystem {
	return Filesystem{
		Permissions: 0777,
		Symlink:     true,
		Hardlink:    true,
		Inode:       true,
	}
}

// Read retuns a list of all mountpoints and their filesystem types.
// It always returns a valid Info object, but may also return an error on
// failures. Errors are worked around as much as possible. Thus, you can safely
// ignore Read() errors while still having reasonable defaults.
func Read() (*Info, error) {
	f, err := os.Open(MOUNTINFO_PATH)
	if err != nil {
		// Maybe an old system that doesn't have the file, or /proc wasn't
		// mounted (yet).
		return &Info{}, err
	}
	return read(f)
}

func read(f io.Reader) (*Info, error) {
	info := &Info{
		mountPaths:   make(map[string]*MountPoint),
		mountNumbers: make(map[uint64]*MountPoint),
	}
	reader := bufio.NewReader(f)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				// We haven't consumed the last 'line', but the file normally
				// ends in a newline anyway.
				return info, nil
			}
			return info, err
		}
		line = line[:len(line)-1]

		fields := strings.Split(line, " ")
		if len(fields) < 10 {
			// There are not enough fields (unexpected).
			continue
		}

		mount := &MountPoint{}

		// Extract major and minor device number.
		stdev := strings.Split(fields[2], ":")
		if len(stdev) != 2 {
			continue
		}
		devMajor, err1 := strconv.ParseUint(stdev[0], 10, 12)
		devMinor, err2 := strconv.ParseUint(stdev[1], 10, 20)
		if err1 != nil || err2 != nil {
			continue
		}

		// We want to know how the major and minor number are encoded in the
		// st_dev field of stat() results, as that's one of the ways we're going
		// to find the right filesystem to a file.
		// I have found two sources for this information: the Linux kernel
		// source and glibc.
		// Linux: include/linux/kdev_t.h in the Linux sources, e.g.
		//    http://lxr.free-electrons.com/source/include/linux/kdev_t.h
		// In particular, new_encode_dev() and new_decode_dev().
		// It looks like the st_dev field is encoded in the following format:
		//    bit 8-20 minor | bit 0-11 major | bit 0-7 minor
		// Additionally, it looks like the Linux kernel decided upon 32-bits
		// numbers (20 minor and 12 major).
		// The glibc source has 64 bits numbers, and takes a slightly different
		// approach (simply giving the last 32 bits to the major number):
		//    bit 12-43 major | bit 8-19 minor | bit 0-11 major | bit 0-7 minor
		// I'll keep it at the current Linux approach, and limit the minor and
		// major number to 20 and 12 bits (see above).
		mount.devNumber = devMinor&0xff | devMajor&0xfff<<8 | devMinor&^0xff<<12

		var ok bool
		mount.FSRoot, _ = mountParse(fields[3])
		mount.Root, ok = mountParse(fields[4])
		if !ok {
			// This is a critical part of the MountPoint struct.
			continue
		}

		pos := 6
		for pos < len(fields) && fields[pos] != "-" {
			pos++
		}
		pos++
		if pos >= len(fields) {
			// Type is another critical field.
			continue
		}
		mount.Type = fields[pos]

		// TODO check for duplicates?
		info.mountPaths[mount.Root] = mount
		info.mountNumbers[mount.devNumber] = mount
	}

	return info, nil
}

// mountParse parses paths like in /etc/fstab, /etc/mtab, and
// /proc/self/mountinfo.
func mountParse(in string) (string, bool) {
	out := make([]byte, len(in))
	i_out := 0
	for i_in := 0; i_in < len(in); i_in++ {
		c := in[i_in]
		if c == '\\' {
			// \xxx escape, where xxx is a three-digit octal number for the
			// character.
			if (i_in + 3) >= len(in) {
				// There aren't at least three more spaces left.
				return "", false
			}
			n, err := strconv.ParseUint(in[i_in+1:i_in+4], 8, 8)
			if err != nil {
				return "", false
			}
			c = byte(n)
			i_in += 3
		}
		out[i_out] = c
		i_out++
	}
	return string(out[:i_out]), true
}

// MountPoint returns the mountpoint associated with the FileInfo, or nil if the
// mount point wasn't found. When nil is returned, you can still safely call
// Filesystem() on the returned pointer (it will return the default filesystem).
// Calling GetReal() on a nil *Info is not allowed.
//
// The filePath must be absolute (pass filepath.IsAbs), otherwise this function
// will panic. If you aren't sure that the path is absolute, use filepath.Abs().
func (info *Info) GetReal(filePath string, fileInfo os.FileInfo) *MountPoint {
	if !filepath.IsAbs(filePath) {
		panic("path must be absolute")
	}

	// Find the st_dev value.
	var deviceId uint64
	switch sys := fileInfo.Sys().(type) {
	case *syscall.Stat_t:
		deviceId = sys.Dev
	case *unix.Stat_t:
		deviceId = sys.Dev
	default:
		return nil
	}

	// See:
	//  * https://criu.org/Filesystems_pecularities
	//  * https://bugzilla.redhat.com/show_bug.cgi?id=711881
	//  * https://www.mail-archive.com/linux-btrfs@vger.kernel.org/msg02875.html

	// First try finding by device number (works with most filesystems).
	if mount, ok := info.mountNumbers[deviceId]; ok {
		return mount
	}
	filePath = filepath.Clean(filePath)

	// Now find the path that is the closest to the requested path.
	for i := len(filePath) - 1; i >= 0; i-- {
		if filePath[i] != '/' {
			continue
		}
		// We've found a directory
		if mount, ok := info.mountPaths[filePath[:i]]; ok {
			return mount
		}
	}

	// We can't find the root: the FileInfo is probably just empty.
	return nil
}
<commit_msg>Fix root directory on Linux<commit_after>// Copyright 2016 Ayke van Laethem.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.txt file.

// On Linux, it parses /proc/self/mountinfo. Then it can find the mount point
// based on the st_dev field of the stat result (os.FileInfo.Sys()), and
// determine filesystem capabilities from the filesystem type.
package osfs

import (
	"bufio"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"golang.org/x/sys/unix"
)

// MOUNTINFO_PATH is the common path for the mountinfo file on Linux.
const MOUNTINFO_PATH = "/proc/self/mountinfo"

// From the kernel docs at
// https://www.kernel.org/doc/Documentation/filesystems/proc.txt
//
//   This file contains lines of the form:
//
//   36 35 98:0 /mnt1 /mnt2 rw,noatime master:1 - ext3 /dev/root rw,errors=continue
//   (1)(2)(3)   (4)   (5)      (6)      (7)   (8) (9)   (10)         (11)
//
//   (1) mount ID:  unique identifier of the mount (may be reused after umount)
//   (2) parent ID:  ID of parent (or of self for the top of the mount tree)
//   (3) major:minor:  value of st_dev for files on filesystem
//   (4) root:  root of the mount within the filesystem
//   (5) mount point:  mount point relative to the process's root
//   (6) mount options:  per mount options
//   (7) optional fields:  zero or more fields of the form "tag[:value]"
//   (8) separator:  marks the end of the optional fields
//   (9) filesystem type:  name of filesystem of the form "type[.subtype]"
//   (10) mount source:  filesystem specific information or "none"
//   (11) super options:  per super block options

func defaultFilesystem() Filesystem {
	return Filesystem{
		Permissions: 0777,
		Symlink:     true,
		Hardlink:    true,
		Inode:       true,
	}
}

// Read retuns a list of all mountpoints and their filesystem types.
// It always returns a valid Info object, but may also return an error on
// failures. Errors are worked around as much as possible. Thus, you can safely
// ignore Read() errors while still having reasonable defaults.
func Read() (*Info, error) {
	f, err := os.Open(MOUNTINFO_PATH)
	if err != nil {
		// Maybe an old system that doesn't have the file, or /proc wasn't
		// mounted (yet).
		return &Info{}, err
	}
	return read(f)
}

func read(f io.Reader) (*Info, error) {
	info := &Info{
		mountPaths:   make(map[string]*MountPoint),
		mountNumbers: make(map[uint64]*MountPoint),
	}
	reader := bufio.NewReader(f)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				// We haven't consumed the last 'line', but the file normally
				// ends in a newline anyway.
				return info, nil
			}
			return info, err
		}
		line = line[:len(line)-1]

		fields := strings.Split(line, " ")
		if len(fields) < 10 {
			// There are not enough fields (unexpected).
			continue
		}

		mount := &MountPoint{}

		// Extract major and minor device number.
		stdev := strings.Split(fields[2], ":")
		if len(stdev) != 2 {
			continue
		}
		devMajor, err1 := strconv.ParseUint(stdev[0], 10, 12)
		devMinor, err2 := strconv.ParseUint(stdev[1], 10, 20)
		if err1 != nil || err2 != nil {
			continue
		}

		// We want to know how the major and minor number are encoded in the
		// st_dev field of stat() results, as that's one of the ways we're going
		// to find the right filesystem to a file.
		// I have found two sources for this information: the Linux kernel
		// source and glibc.
		// Linux: include/linux/kdev_t.h in the Linux sources, e.g.
		//    http://lxr.free-electrons.com/source/include/linux/kdev_t.h
		// In particular, new_encode_dev() and new_decode_dev().
		// It looks like the st_dev field is encoded in the following format:
		//    bit 8-20 minor | bit 0-11 major | bit 0-7 minor
		// Additionally, it looks like the Linux kernel decided upon 32-bits
		// numbers (20 minor and 12 major).
		// The glibc source has 64 bits numbers, and takes a slightly different
		// approach (simply giving the last 32 bits to the major number):
		//    bit 12-43 major | bit 8-19 minor | bit 0-11 major | bit 0-7 minor
		// I'll keep it at the current Linux approach, and limit the minor and
		// major number to 20 and 12 bits (see above).
		mount.devNumber = devMinor&0xff | devMajor&0xfff<<8 | devMinor&^0xff<<12

		var ok bool
		mount.FSRoot, _ = mountParse(fields[3])
		mount.Root, ok = mountParse(fields[4])
		if !ok {
			// This is a critical part of the MountPoint struct.
			continue
		}

		pos := 6
		for pos < len(fields) && fields[pos] != "-" {
			pos++
		}
		pos++
		if pos >= len(fields) {
			// Type is another critical field.
			continue
		}
		mount.Type = fields[pos]

		// TODO check for duplicates?
		info.mountPaths[mount.Root] = mount
		info.mountNumbers[mount.devNumber] = mount
	}

	return info, nil
}

// mountParse parses paths like in /etc/fstab, /etc/mtab, and
// /proc/self/mountinfo.
func mountParse(in string) (string, bool) {
	out := make([]byte, len(in))
	i_out := 0
	for i_in := 0; i_in < len(in); i_in++ {
		c := in[i_in]
		if c == '\\' {
			// \xxx escape, where xxx is a three-digit octal number for the
			// character.
			if (i_in + 3) >= len(in) {
				// There aren't at least three more spaces left.
				return "", false
			}
			n, err := strconv.ParseUint(in[i_in+1:i_in+4], 8, 8)
			if err != nil {
				return "", false
			}
			c = byte(n)
			i_in += 3
		}
		out[i_out] = c
		i_out++
	}
	return string(out[:i_out]), true
}

// MountPoint returns the mountpoint associated with the FileInfo, or nil if the
// mount point wasn't found. When nil is returned, you can still safely call
// Filesystem() on the returned pointer (it will return the default filesystem).
// Calling GetReal() on a nil *Info is not allowed.
//
// The filePath must be absolute (pass filepath.IsAbs), otherwise this function
// will panic. If you aren't sure that the path is absolute, use filepath.Abs().
func (info *Info) GetReal(filePath string, fileInfo os.FileInfo) *MountPoint {
	if !filepath.IsAbs(filePath) {
		panic("path must be absolute")
	}

	// Find the st_dev value.
	var deviceId uint64
	switch sys := fileInfo.Sys().(type) {
	case *syscall.Stat_t:
		deviceId = sys.Dev
	case *unix.Stat_t:
		deviceId = sys.Dev
	default:
		return nil
	}

	// See:
	//  * https://criu.org/Filesystems_pecularities
	//  * https://bugzilla.redhat.com/show_bug.cgi?id=711881
	//  * https://www.mail-archive.com/linux-btrfs@vger.kernel.org/msg02875.html

	// First try finding by device number (works with most filesystems).
	if mount, ok := info.mountNumbers[deviceId]; ok {
		return mount
	}
	filePath = filepath.Clean(filePath)

	// Now find the path that is the closest to the requested path.
	for i := len(filePath) - 1; i >= 0; i-- {
		if filePath[i] != '/' {
			continue
		}
		// We've found a directory

		testPath := filePath[:i]
		if (testPath == "") {
			// Special-case the root directory
			testPath = "/"
		}

		if mount, ok := info.mountPaths[testPath]; ok {
			return mount
		}
	}

	// We can't find the root: the FileInfo is probably just empty.
	return nil
}
<|endoftext|>


<commit_before>package prereg

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/itimofeev/hustledb/util"
	"strings"
)

type PreregComp struct {
	ID            int64              `json:"id" db:"id"`
	PreregId      int                `json:"prereg_id" db:"prereg_id"`
	CompetitionId string             `json:"competition_id" db:"competition_id"`
	Nominations   []PreregNomination `json:"nominations"`
}

type PreregNomination struct {
	ID           int64  `json:"id" db:"id"`
	PreregCompId int64  `json:"prereg_comp_id" db:"prereg_comp_id"`
	Title        string `json:"title" db:"title"`
	Records      string `json:"records"`
}

type PreregRecord struct {
	ID          int64 `json:"id" db:"id"`
	PreregNomId int64 `json:"prereg_nom_id" db:"prereg_nom_id"`
	Index       int   `json:"index" db:"idnex"`
}

type PreregDancer struct {
	ID      int64        `json:"id" db:"id"`
	CodeASH *string      `json:"code_ash" db:"code_ash"`
	Class   string       `json:"class" db:"class"`
	Title   string       `json:"title" db:"title"`
	Clubs   []PreregClub `json:"clubs"`
}

type PreregClub struct {
	ID    int64  `json:"id" db:"id"`
	Title string `json:"title" db:"title"`
}

const mainUrl = "http://www.liveindance.ru/"
const regUrl = "http://www.liveindance.ru/contest/reg.php?id=%d"

func ParseAllPreregLinks() []string {
	data := util.GetUrlContent(mainUrl)
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(data))
	util.CheckErr(err)

	var listLinks []string

	doc.
		Find(".panel-body a").
		FilterFunction(func(i int, s *goquery.Selection) bool {
			link := s.AttrOr("href", "")
			return strings.Contains(link, "http://www.liveindance.ru/contest/registration/list.php?id=")
		}).
		Each(func(i int, s *goquery.Selection) {
			link := s.AttrOr("href", "")
			listLinks = append(listLinks, link)
		})
	return listLinks
}

func ParsePreregId(listLink string) int {
	preregId := strings.Replace(listLink, "http://www.liveindance.ru/contest/registration/list.php?id=", "", 1)

	return util.Atoi(preregId)
}

func GetForumCompetitionId(preregId int) string {
	regUrlFull := fmt.Sprintf(regUrl, preregId)
	data := util.GetUrlContent(regUrlFull)

	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(data))
	util.CheckErr(err)

	forumLinkA := doc.
		Find("table table table tr a").
		FilterFunction(func(i int, s *goquery.Selection) bool {
			return strings.Contains(s.AttrOr("href", ""), "http://hustle-sa.ru/forum/index.php?showtopic=")
		}).First()

	return forumLinkA.AttrOr("href", "")
}

func ParsePreregCompetition(preregId int, fCompUrl string) PreregComp {
 return PreregComp{}
}
<commit_msg>Список предварительной регистрации на турниры #7 * parse prereg id and forum url from main page<commit_after>package prereg

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"github.com/itimofeev/hustledb/util"
	"strings"
)

type PreregComp struct {
	ID            int64              `json:"id" db:"id"`
	PreregId      int                `json:"prereg_id" db:"prereg_id"`
	CompetitionId string             `json:"competition_id" db:"competition_id"`
	Nominations   []PreregNomination `json:"nominations"`
}

type PreregNomination struct {
	ID           int64  `json:"id" db:"id"`
	PreregCompId int64  `json:"prereg_comp_id" db:"prereg_comp_id"`
	Title        string `json:"title" db:"title"`
	Records      string `json:"records"`
}

type PreregRecord struct {
	ID          int64 `json:"id" db:"id"`
	PreregNomId int64 `json:"prereg_nom_id" db:"prereg_nom_id"`
	Index       int   `json:"index" db:"idnex"`
}

type PreregDancer struct {
	ID      int64        `json:"id" db:"id"`
	CodeASH *string      `json:"code_ash" db:"code_ash"`
	Class   string       `json:"class" db:"class"`
	Title   string       `json:"title" db:"title"`
	Clubs   []PreregClub `json:"clubs"`
}

type PreregClub struct {
	ID    int64  `json:"id" db:"id"`
	Title string `json:"title" db:"title"`
}

const mainUrl = "http://www.liveindance.ru/"
const regUrl = "http://www.liveindance.ru/contest/reg.php?id=%d"

func ParseAllPreregLinks() []string {
	data := util.GetUrlContent(mainUrl)
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(data))
	util.CheckErr(err)

	var listLinks []string

	doc.
		Find(".panel-body a").
		FilterFunction(func(i int, s *goquery.Selection) bool {
			link := s.AttrOr("href", "")
			return strings.Contains(link, "http://www.liveindance.ru/contest/registration/list.php?id=")
		}).
		Each(func(i int, s *goquery.Selection) {
			link := s.AttrOr("href", "")
			listLinks = append(listLinks, link)
		})
	return listLinks
}

func ParsePreregId(listLink string) int {
	preregId := strings.Replace(listLink, "http://www.liveindance.ru/contest/registration/list.php?id=", "", 1)

	return util.Atoi(preregId)
}

func GetForumCompetitionId(preregId int) string {
	regUrlFull := fmt.Sprintf(regUrl, preregId)
	data := util.GetUrlContent(regUrlFull)

	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(data))
	util.CheckErr(err)

	forumLinkA := doc.
		Find("table table table tr a").
		FilterFunction(func(i int, s *goquery.Selection) bool {
			return strings.Contains(s.AttrOr("href", ""), "http://hustle-sa.ru/forum/index.php?showtopic=")
		}).First()

	return forumLinkA.AttrOr("href", "")
}

func ParsePreregCompetition(preregId int, fCompUrl string) PreregComp {
	return PreregComp{}
}
<|endoftext|>


<commit_before>package generator

import (
	"bytes"
	"fmt"
	"hash/crc32"

	tl "github.com/xlab/cgogen/translator"
)

func (gen *Generator) getStructHelpers(goStructName []byte, cStructName string, spec tl.CType) (helpers []*Helper) {
	crc := getRefCRC(spec)
	cgoSpec := gen.tr.CGoSpec(spec, true)

	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "func (x *%s) Ref() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			return nil
		}
		return x.ref%2x
	}`, crc)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Ref", goStructName),
		Description: "Ref returns a reference.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Free()", goStructName)
	fmt.Fprintf(buf, `{
		if x != nil && x.allocs%2x != nil {
			x.allocs%2x.(*cgoAllocMap).Free()
			x.ref%2x = nil
		}
	}`, crc, crc, crc)
	helpers = append(helpers, &Helper{
		Name: fmt.Sprintf("%s.Free", goStructName),
		Description: "Free cleanups the memory using the free stdlib function on C side.\n" +
			"Does nothing if object has no pointer.",
		Source: buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func New%sRef(ref unsafe.Pointer) *%s", goStructName, goStructName)
	fmt.Fprintf(buf, `{
		if ref == nil {
			return nil
		}
		obj := new(%s)
		obj.ref%2x = (*%s)(unsafe.Pointer(ref))
		return obj
	}`, goStructName, crc, cgoSpec)

	name := fmt.Sprintf("New%sRef", goStructName)
	helpers = append(helpers, &Helper{
		Name:        name,
		Description: name + " initialises a new struct holding the reference to the originaitng C struct.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) PassRef() (*%s, *cgoAllocMap) {\n", goStructName, cgoSpec)
	buf.Write(gen.getPassRefSource(goStructName, cStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassRef", goStructName),
		Description: "PassRef returns a reference and creates new C object if no refernce yet.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x %s) PassValue() (%s, *cgoAllocMap) {\n", goStructName, cgoSpec)
	buf.Write(gen.getPassValueSource(goStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassValue", goStructName),
		Description: "PassValue creates a new C object if no refernce yet and returns the dereferenced value.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Deref() {\n", goStructName)
	buf.Write(gen.getDerefSource(goStructName, cStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Deref", goStructName),
		Description: "Deref reads the internal fields of struct from its C pointer.",
		Source:      buf.String(),
	})
	return
}

func (gen *Generator) getRawStructHelpers(goStructName []byte, spec tl.CType) (helpers []*Helper) {
	if spec.GetPointers() > 0 {
		return nil // can't addess a pointer receiver
	}
	cgoSpec := gen.tr.CGoSpec(spec, true)

	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "func (x *%s) Ref() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			return nil
		}
		return (*%s)(unsafe.Pointer(x))
	}`, cgoSpec)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Ref", goStructName),
		Description: "Ref returns a reference.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Free()", goStructName)
	fmt.Fprint(buf, `{
		if x != nil  {
			C.free(unsafe.Pointer(x))
		}
	}`)
	helpers = append(helpers, &Helper{
		Name: fmt.Sprintf("%s.Free", goStructName),
		Description: "Free cleanups the memory using the free stdlib function on C side.\n" +
			"Does nothing if object has no pointer.",
		Source: buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func New%sRef(ref *%s) *%s", goStructName, cgoSpec, goStructName)
	fmt.Fprintf(buf, `{
		return (*%s)(unsafe.Pointer(ref))
	}`, goStructName)
	name := fmt.Sprintf("New%sRef", goStructName)
	helpers = append(helpers, &Helper{
		Name:        name,
		Description: name + " initialises a new struct.",
		Source:      buf.String(),
	})

	buf.Reset()
	allocHelper := gen.getAllocMemoryHelper(cgoSpec)
	fmt.Fprintf(buf, "func (x *%s) PassRef() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			x = (*%s)(%s(1))
		}
		return (*%s)(unsafe.Pointer(x))
	}`, goStructName, allocHelper.Name, cgoSpec)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassRef", goStructName),
		Description: "PassRef returns a reference and creates new object if no refernce yet.",
		Source:      buf.String(),
		Requires:    []*Helper{allocHelper},
	})
	return
}

func (gen *Generator) getPassRefSource(goStructName []byte, cStructName string, spec tl.CType) []byte {
	cgoSpec := gen.tr.CGoSpec(spec, false)
	structSpec := spec.(*tl.CStructSpec)
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x == nil {
		return nil, nil
	} else if x.ref%2x != nil {
		return x.ref%2x, nil
	}`, crc, crc)
	writeSpace(buf, 1)

	h := gen.getAllocMemoryHelper(tl.CGoSpec{Base: cgoSpec.Base})
	gen.submitHelper(h)

	fmt.Fprintf(buf, "mem%2x := %s(1)\n", crc, h.Name)
	fmt.Fprintf(buf, "ref%2x := (*%s)(mem%2x)\n", crc, cgoSpec.Base, crc)
	fmt.Fprintf(buf, "allocs%2x := new(cgoAllocMap)", crc)

	writeSpace(buf, 1)

	ptrTipRx, typeTipRx, memTipRx := gen.tr.TipRxsForSpec(tl.TipScopeStruct, cStructName, spec)
	for i, m := range structSpec.Members {
		if len(m.Name) == 0 {
			continue
			// TODO: generate setters
		}

		typeName := m.Spec.GetBase()
		switch m.Spec.Kind() {
		case tl.StructKind, tl.OpaqueStructKind, tl.EnumKind:
			if !gen.tr.IsAcceptableName(tl.TargetType, typeName) {
				continue
			}
		}
		memTip := memTipRx.TipAt(i)
		if !memTip.IsValid() {
			memTip = gen.MemTipOf(m)
		}
		ptrTip := ptrTipRx.TipAt(i)
		if memTip == tl.TipMemRaw {
			ptrTip = tl.TipPtrSRef
		}
		typeTip := typeTipRx.TipAt(i)
		goSpec := gen.tr.TranslateSpec(m.Spec, ptrTip, typeTip)
		cgoSpec := gen.tr.CGoSpec(m.Spec, false)
		const public = true
		goName := "x." + string(gen.tr.TransformName(tl.TargetType, m.Name, public))
		fromProxy, nillable := gen.proxyValueFromGo(memTip, goName, goSpec, cgoSpec)
		if nillable {
			fmt.Fprintf(buf, "if %s != nil {\n", goName)
		}
		fmt.Fprintf(buf, "var c%s_allocs *cgoAllocMap\n", m.Name)
		fmt.Fprintf(buf, "ref%2x.%s, c%s_allocs  = %s\n", crc, m.Name, m.Name, fromProxy)
		fmt.Fprintf(buf, "allocs%2x.Borrow(c%s_allocs)\n", crc, m.Name)
		if nillable {
			fmt.Fprintf(buf, "}\n\n")
		} else {
			fmt.Fprint(buf, "\n")
		}
	}
	fmt.Fprintf(buf, "x.ref%2x = ref%2x\n", crc, crc)
	fmt.Fprintf(buf, "x.allocs%2x = allocs%2x\n", crc, crc)
	fmt.Fprintf(buf, "return ref%2x, allocs%2x\n", crc, crc)
	writeSpace(buf, 1)
	return buf.Bytes()
}

func (gen *Generator) getPassValueSource(goStructName []byte, spec tl.CType) []byte {
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x.ref%2x != nil {
		return *x.ref%2x, nil
	}`, crc, crc)
	writeSpace(buf, 1)
	fmt.Fprintf(buf, "ref, allocs := x.PassRef()\n")
	fmt.Fprintf(buf, "return *ref, allocs\n")
	return buf.Bytes()
}

func getRefCRC(spec tl.CType) uint32 {
	return crc32.ChecksumIEEE([]byte(spec.String()))
}

func (gen *Generator) getDerefSource(goStructName []byte, cStructName string, spec tl.CType) []byte {
	structSpec := spec.(*tl.CStructSpec)
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x.ref%2x == nil {
		return
	}`, crc)
	writeSpace(buf, 1)

	ptrTipRx, typeTipRx, memTipRx := gen.tr.TipRxsForSpec(tl.TipScopeStruct, cStructName, spec)
	for i, m := range structSpec.Members {
		if len(m.Name) == 0 {
			continue
			// TODO: generate getters
		}

		typeName := m.Spec.GetBase()
		switch m.Spec.Kind() {
		case tl.StructKind, tl.OpaqueStructKind, tl.EnumKind:
			if !gen.tr.IsAcceptableName(tl.TargetType, typeName) {
				continue
			}
		}
		memTip := memTipRx.TipAt(i)
		if !memTip.IsValid() {
			memTip = gen.MemTipOf(m)
		}
		ptrTip := ptrTipRx.TipAt(i)
		if memTip == tl.TipMemRaw {
			ptrTip = tl.TipPtrSRef
		}
		typeTip := typeTipRx.TipAt(i)
		goSpec := gen.tr.TranslateSpec(m.Spec, ptrTip, typeTip)
		const public = true
		goName := "x." + string(gen.tr.TransformName(tl.TargetType, m.Name, public))
		cgoName := fmt.Sprintf("x.ref%2x.%s", crc, m.Name)
		cgoSpec := gen.tr.CGoSpec(m.Spec, false)
		toProxy, _ := gen.proxyValueToGo(memTip, goName, cgoName, goSpec, cgoSpec)
		fmt.Fprintln(buf, toProxy)
	}
	return buf.Bytes()
}
<commit_msg>Allocation helpers (NewXXX()) for raw structs.<commit_after>package generator

import (
	"bytes"
	"fmt"
	"hash/crc32"

	tl "github.com/xlab/cgogen/translator"
)

func (gen *Generator) getStructHelpers(goStructName []byte, cStructName string, spec tl.CType) (helpers []*Helper) {
	crc := getRefCRC(spec)
	cgoSpec := gen.tr.CGoSpec(spec, true)

	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "func (x *%s) Ref() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			return nil
		}
		return x.ref%2x
	}`, crc)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Ref", goStructName),
		Description: "Ref returns a reference.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Free()", goStructName)
	fmt.Fprintf(buf, `{
		if x != nil && x.allocs%2x != nil {
			x.allocs%2x.(*cgoAllocMap).Free()
			x.ref%2x = nil
		}
	}`, crc, crc, crc)
	helpers = append(helpers, &Helper{
		Name: fmt.Sprintf("%s.Free", goStructName),
		Description: "Free cleanups the memory using the free stdlib function on C side.\n" +
			"Does nothing if object has no pointer.",
		Source: buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func New%sRef(ref unsafe.Pointer) *%s", goStructName, goStructName)
	fmt.Fprintf(buf, `{
		if ref == nil {
			return nil
		}
		obj := new(%s)
		obj.ref%2x = (*%s)(unsafe.Pointer(ref))
		return obj
	}`, goStructName, crc, cgoSpec)

	name := fmt.Sprintf("New%sRef", goStructName)
	helpers = append(helpers, &Helper{
		Name:        name,
		Description: name + " initialises a new struct holding the reference to the originaitng C struct.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) PassRef() (*%s, *cgoAllocMap) {\n", goStructName, cgoSpec)
	buf.Write(gen.getPassRefSource(goStructName, cStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassRef", goStructName),
		Description: "PassRef returns a reference and creates new C object if no refernce yet.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x %s) PassValue() (%s, *cgoAllocMap) {\n", goStructName, cgoSpec)
	buf.Write(gen.getPassValueSource(goStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassValue", goStructName),
		Description: "PassValue creates a new C object if no refernce yet and returns the dereferenced value.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Deref() {\n", goStructName)
	buf.Write(gen.getDerefSource(goStructName, cStructName, spec))
	buf.WriteRune('}')
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Deref", goStructName),
		Description: "Deref reads the internal fields of struct from its C pointer.",
		Source:      buf.String(),
	})
	return
}

func (gen *Generator) getRawStructHelpers(goStructName []byte, spec tl.CType) (helpers []*Helper) {
	if spec.GetPointers() > 0 {
		return nil // can't addess a pointer receiver
	}
	cgoSpec := gen.tr.CGoSpec(spec, true)

	buf := new(bytes.Buffer)
	fmt.Fprintf(buf, "func (x *%s) Ref() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			return nil
		}
		return (*%s)(unsafe.Pointer(x))
	}`, cgoSpec)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.Ref", goStructName),
		Description: "Ref returns a reference.",
		Source:      buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) Free()", goStructName)
	fmt.Fprint(buf, `{
		if x != nil  {
			C.free(unsafe.Pointer(x))
		}
	}`)
	helpers = append(helpers, &Helper{
		Name: fmt.Sprintf("%s.Free", goStructName),
		Description: "Free cleanups the memory using the free stdlib function on C side.\n" +
			"Does nothing if object has no pointer.",
		Source: buf.String(),
	})

	buf.Reset()
	fmt.Fprintf(buf, "func New%sRef(ref *%s) *%s", goStructName, cgoSpec, goStructName)
	fmt.Fprintf(buf, `{
		return (*%s)(unsafe.Pointer(ref))
	}`, goStructName)
	name := fmt.Sprintf("New%sRef", goStructName)
	helpers = append(helpers, &Helper{
		Name:        name,
		Description: name + " initialises a new struct.",
		Source:      buf.String(),
	})

	buf.Reset()
	allocHelper := gen.getAllocMemoryHelper(cgoSpec)
	fmt.Fprintf(buf, "func New%s() *%s", goStructName, goStructName)
	fmt.Fprintf(buf, `{
		return (*%s)(%s(1))
	}`, goStructName, allocHelper.Name)
	name = fmt.Sprintf("New%s", goStructName)
	helpers = append(helpers, &Helper{
		Name:        name,
		Description: name + " allocates a new object.",
		Source:      buf.String(),
		Requires:    []*Helper{allocHelper},
	})

	buf.Reset()
	fmt.Fprintf(buf, "func (x *%s) PassRef() *%s", goStructName, cgoSpec)
	fmt.Fprintf(buf, `{
		if x == nil {
			x = (*%s)(%s(1))
		}
		return (*%s)(unsafe.Pointer(x))
	}`, goStructName, allocHelper.Name, cgoSpec)
	helpers = append(helpers, &Helper{
		Name:        fmt.Sprintf("%s.PassRef", goStructName),
		Description: "PassRef returns a reference or creates new object if no reference yet.",
		Source:      buf.String(),
		Requires:    []*Helper{allocHelper},
	})
	return
}

func (gen *Generator) getPassRefSource(goStructName []byte, cStructName string, spec tl.CType) []byte {
	cgoSpec := gen.tr.CGoSpec(spec, false)
	structSpec := spec.(*tl.CStructSpec)
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x == nil {
		return nil, nil
	} else if x.ref%2x != nil {
		return x.ref%2x, nil
	}`, crc, crc)
	writeSpace(buf, 1)

	h := gen.getAllocMemoryHelper(tl.CGoSpec{Base: cgoSpec.Base})
	gen.submitHelper(h)

	fmt.Fprintf(buf, "mem%2x := %s(1)\n", crc, h.Name)
	fmt.Fprintf(buf, "ref%2x := (*%s)(mem%2x)\n", crc, cgoSpec.Base, crc)
	fmt.Fprintf(buf, "allocs%2x := new(cgoAllocMap)", crc)

	writeSpace(buf, 1)

	ptrTipRx, typeTipRx, memTipRx := gen.tr.TipRxsForSpec(tl.TipScopeStruct, cStructName, spec)
	for i, m := range structSpec.Members {
		if len(m.Name) == 0 {
			continue
			// TODO: generate setters
		}

		typeName := m.Spec.GetBase()
		switch m.Spec.Kind() {
		case tl.StructKind, tl.OpaqueStructKind, tl.EnumKind:
			if !gen.tr.IsAcceptableName(tl.TargetType, typeName) {
				continue
			}
		}
		memTip := memTipRx.TipAt(i)
		if !memTip.IsValid() {
			memTip = gen.MemTipOf(m)
		}
		ptrTip := ptrTipRx.TipAt(i)
		if memTip == tl.TipMemRaw {
			ptrTip = tl.TipPtrSRef
		}
		typeTip := typeTipRx.TipAt(i)
		goSpec := gen.tr.TranslateSpec(m.Spec, ptrTip, typeTip)
		cgoSpec := gen.tr.CGoSpec(m.Spec, false)
		const public = true
		goName := "x." + string(gen.tr.TransformName(tl.TargetType, m.Name, public))
		fromProxy, nillable := gen.proxyValueFromGo(memTip, goName, goSpec, cgoSpec)
		if nillable {
			fmt.Fprintf(buf, "if %s != nil {\n", goName)
		}
		fmt.Fprintf(buf, "var c%s_allocs *cgoAllocMap\n", m.Name)
		fmt.Fprintf(buf, "ref%2x.%s, c%s_allocs  = %s\n", crc, m.Name, m.Name, fromProxy)
		fmt.Fprintf(buf, "allocs%2x.Borrow(c%s_allocs)\n", crc, m.Name)
		if nillable {
			fmt.Fprintf(buf, "}\n\n")
		} else {
			fmt.Fprint(buf, "\n")
		}
	}
	fmt.Fprintf(buf, "x.ref%2x = ref%2x\n", crc, crc)
	fmt.Fprintf(buf, "x.allocs%2x = allocs%2x\n", crc, crc)
	fmt.Fprintf(buf, "return ref%2x, allocs%2x\n", crc, crc)
	writeSpace(buf, 1)
	return buf.Bytes()
}

func (gen *Generator) getPassValueSource(goStructName []byte, spec tl.CType) []byte {
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x.ref%2x != nil {
		return *x.ref%2x, nil
	}`, crc, crc)
	writeSpace(buf, 1)
	fmt.Fprintf(buf, "ref, allocs := x.PassRef()\n")
	fmt.Fprintf(buf, "return *ref, allocs\n")
	return buf.Bytes()
}

func getRefCRC(spec tl.CType) uint32 {
	return crc32.ChecksumIEEE([]byte(spec.String()))
}

func (gen *Generator) getDerefSource(goStructName []byte, cStructName string, spec tl.CType) []byte {
	structSpec := spec.(*tl.CStructSpec)
	buf := new(bytes.Buffer)
	crc := getRefCRC(spec)
	fmt.Fprintf(buf, `if x.ref%2x == nil {
		return
	}`, crc)
	writeSpace(buf, 1)

	ptrTipRx, typeTipRx, memTipRx := gen.tr.TipRxsForSpec(tl.TipScopeStruct, cStructName, spec)
	for i, m := range structSpec.Members {
		if len(m.Name) == 0 {
			continue
			// TODO: generate getters
		}

		typeName := m.Spec.GetBase()
		switch m.Spec.Kind() {
		case tl.StructKind, tl.OpaqueStructKind, tl.EnumKind:
			if !gen.tr.IsAcceptableName(tl.TargetType, typeName) {
				continue
			}
		}
		memTip := memTipRx.TipAt(i)
		if !memTip.IsValid() {
			memTip = gen.MemTipOf(m)
		}
		ptrTip := ptrTipRx.TipAt(i)
		if memTip == tl.TipMemRaw {
			ptrTip = tl.TipPtrSRef
		}
		typeTip := typeTipRx.TipAt(i)
		goSpec := gen.tr.TranslateSpec(m.Spec, ptrTip, typeTip)
		const public = true
		goName := "x." + string(gen.tr.TransformName(tl.TargetType, m.Name, public))
		cgoName := fmt.Sprintf("x.ref%2x.%s", crc, m.Name)
		cgoSpec := gen.tr.CGoSpec(m.Spec, false)
		toProxy, _ := gen.proxyValueToGo(memTip, goName, cgoName, goSpec, cgoSpec)
		fmt.Fprintln(buf, toProxy)
	}
	return buf.Bytes()
}
<|endoftext|>


<commit_before>package mainfile

import (
	"fmt"
	"text/template"

	"bytes"

	"go/format"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/schema"
)

type generator struct{}

func New() *generator {
	return &generator{}
}

var PathForMain = "%scmd/%s/main.go"

func (g *generator) Name() string {
	return "statements"
}

// GenerateMainFile handles the main file generation for persistent
// connection rpc server
func (g *generator) Generate(context *common.Context, schema *schema.Schema) ([]common.Output, error) {
	moduleName := context.ModuleNameFunc(schema.Title)
	outputs := make([]common.Output, 0)

	for _, def := range schema.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t != "object" {
					continue
				}
			}
		}

		f, err := generateMainFile(schema)
		if err != nil {
			return nil, err
		}

		path := fmt.Sprintf(
			PathForMain,
			context.Config.Target,
			moduleName,
		)

		outputs = append(outputs, common.Output{
			Content: f,
			Path:    path,
		})

	}

	return outputs, nil

}

func generateMainFile(s *schema.Schema) ([]byte, error) {
	const templateName = "mainfile.tmpl"
	temp := template.New(templateName).Funcs(common.TemplateFuncs)

	if _, err := temp.Parse(MainFileTemplate); err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	if err := temp.ExecuteTemplate(&buf, templateName, s); err != nil {
		return nil, err
	}

	return format.Source(buf.Bytes())
}

// MainFileTemplate holds the template for the main file generation
var MainFileTemplate = `
package main

import (
    "fmt"
    "net/http"
    "github.com/youtube/vitess/go/rpcplus"
    "github.com/youtube/vitess/go/rpcplus/jsonrpc"
    "github.com/youtube/vitess/go/rpcwrap"
)

var (
    Name    = "{{.Title}}"
    VERSION string
)

var ContextCreator = func(req *http.Request) context.Context {
    return context.Background()
}

var Mux = http.NewServeMux()

func main() {

    {{$Name := .Title}}
    server := rpcplus.NewServer()
    {{range $key, $value := .Definitions}}
        {{$type := .Type}}
        {{/* export functions if they have any exported function */}}
        {{if len .Functions}}
            {{/* export functions if they are objects */}}
            {{if Equal $type "object"}}
                server.Register(new({{ToLower $Name}}api.{{$key}}))
            {{end}}
        {{end}}
    {{end}}

    rpcwrap.ServeCustomRPC(
        Mux,
        server,
        false,  // use auth
        "json", // codec name
        jsonrpc.NewServerCodec,
    )

    rpcwrap.ServeHTTPRPC(
        Mux,                    // httpmuxer
        server,                 // rpcserver
        "http_json",            // codec name
        jsonrpc.NewServerCodec, // jsoncodec
        ContextCreator,         // contextCreator
    )

    fmt.Println("Server listening on 3000")
    http.ListenAndServe("localhost:3000", Mux)
}`
<commit_msg>Gene: make mainfile cli compatible<commit_after>package mainfile

import (
	"fmt"
	"text/template"

	"bytes"

	"go/format"

	"github.com/cihangir/gene/generators/common"
	"github.com/cihangir/schema"
)

type Generator struct{}

func New() *Generator {
	return &Generator{}
}

func (g *Generator) Name() string {
	return "statements"
}

// GenerateMainFile handles the main file generation for persistent
// connection rpc server
func (g *Generator) Generate(context *common.Context, schema *schema.Schema) ([]common.Output, error) {
	moduleName := context.ModuleNameFunc(schema.Title)
	outputs := make([]common.Output, 0)

	for _, def := range schema.Definitions {
		// create models only for objects
		if def.Type != nil {
			if t, ok := def.Type.(string); ok {
				if t != "object" {
					continue
				}
			}
		}

		f, err := generateMainFile(context, schema)
		if err != nil {
			return nil, err
		}

		path := fmt.Sprintf(
			"%s%s/main.go",
			context.Config.Target,
			moduleName,
		)

		outputs = append(outputs, common.Output{
			Content: f,
			Path:    path,
		})

	}

	return outputs, nil

}

func generateMainFile(context *common.Context, s *schema.Schema) ([]byte, error) {
	const templateName = "mainfile.tmpl"
	temp := template.New(templateName).Funcs(context.TemplateFuncs)

	if _, err := temp.Parse(MainFileTemplate); err != nil {
		return nil, err
	}

	var buf bytes.Buffer

	if err := temp.ExecuteTemplate(&buf, templateName, s); err != nil {
		return nil, err
	}

	return format.Source(buf.Bytes())
}

// MainFileTemplate holds the template for the main file generation
var MainFileTemplate = `
package main

import (
    "fmt"
    "net/http"
    "github.com/youtube/vitess/go/rpcplus"
    "github.com/youtube/vitess/go/rpcplus/jsonrpc"
    "github.com/youtube/vitess/go/rpcwrap"
)

var (
    Name    = "{{.Title}}"
    VERSION string
)

var ContextCreator = func(req *http.Request) context.Context {
    return context.Background()
}

var Mux = http.NewServeMux()

func main() {

    {{$Name := .Title}}
    server := rpcplus.NewServer()
    {{range $key, $value := .Definitions}}
        {{$type := .Type}}
        {{/* export functions if they have any exported function */}}
        {{if len .Functions}}
            {{/* export functions if they are objects */}}
            {{if Equal $type "object"}}
                server.Register(new({{ToLower $Name}}api.{{$key}}))
            {{end}}
        {{end}}
    {{end}}

    rpcwrap.ServeCustomRPC(
        Mux,
        server,
        false,  // use auth
        "json", // codec name
        jsonrpc.NewServerCodec,
    )

    rpcwrap.ServeHTTPRPC(
        Mux,                    // httpmuxer
        server,                 // rpcserver
        "http_json",            // codec name
        jsonrpc.NewServerCodec, // jsoncodec
        ContextCreator,         // contextCreator
    )

    fmt.Println("Server listening on 3000")
    http.ListenAndServe("localhost:3000", Mux)
}`
<|endoftext|>