<commit_before>package commands

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"

	key "github.com/ipfs/go-ipfs/blocks/key"
	cmds "github.com/ipfs/go-ipfs/commands"
	"github.com/ipfs/go-ipfs/core"
	dag "github.com/ipfs/go-ipfs/merkledag"
	path "github.com/ipfs/go-ipfs/path"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []key.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	buf := new(bytes.Buffer)
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object.",
		ShortDescription: `
Lists the hashes of all the links an IPFS or IPNS object(s) contains,
with the following format:

  <link base58 hash>

NOTE: List all references recursively by using the flag '-r'.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. Available tokens: <src> <dst> <linkname>.").Default("<dst>"),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`.").Default(false),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output.").Default(false),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(ctx, n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			rw := RefWriter{
				out:       out,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					out <- &RefWrapper{Err: err.Error()}
					return
				}
			}
		}()
	},
	Marshalers: refsMarshallerMap,
	Type:       RefWrapper{},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references.",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			for k := range allKeys {
				out <- &RefWrapper{Ref: k.B58String()}
			}
		}()
	},
	Marshalers: refsMarshallerMap,
	Type:       RefWrapper{},
}

var refsMarshallerMap = cmds.MarshalerMap{
	cmds.Text: func(res cmds.Response) (io.Reader, error) {
		outChan, ok := res.Output().(<-chan interface{})
		if !ok {
			return nil, u.ErrCast()
		}

		marshal := func(v interface{}) (io.Reader, error) {
			obj, ok := v.(*RefWrapper)
			if !ok {
				fmt.Println("%#v", v)
				return nil, u.ErrCast()
			}

			if obj.Err != "" {
				return nil, errors.New(obj.Err)
			}

			return strings.NewReader(obj.Ref + "\n"), nil
		}

		return &cmds.ChannelMarshaler{
			Channel:   outChan,
			Marshaler: marshal,
			Res:       res,
		}, nil
	},
}

func objectsForPaths(ctx context.Context, n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := core.Resolve(ctx, n, path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

type RefWrapper struct {
	Ref string
	Err string
}

type RefWriter struct {
	out chan interface{}
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[key.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	if rw.Recursive {
		return rw.writeRefsRecursive(n)
	}
	return rw.writeRefsSingle(n)
}

func (rw *RefWriter) writeRefsRecursive(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	var count int
	for i, ng := range dag.GetDAG(rw.Ctx, rw.DAG, n) {
		lk := key.Key(n.Links[i].Hash)
		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, n.Links[i].Name); err != nil {
			return count, err
		}

		nd, err := ng.Get(rw.Ctx)
		if err != nil {
			return count, err
		}

		c, err := rw.writeRefsRecursive(nd)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

func (rw *RefWriter) writeRefsSingle(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := key.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k key.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[key.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to key.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.B58String(), -1)
		s = strings.Replace(s, "<dst>", to.B58String(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.B58String() + " -> " + to.B58String()
	default:
		s += to.B58String()
	}

	rw.out <- &RefWrapper{Ref: s}
	return nil
}
<commit_msg>Rremove unneeded print<commit_after>package commands

import (
	"bytes"
	"errors"
	"io"
	"strings"

	key "github.com/ipfs/go-ipfs/blocks/key"
	cmds "github.com/ipfs/go-ipfs/commands"
	"github.com/ipfs/go-ipfs/core"
	dag "github.com/ipfs/go-ipfs/merkledag"
	path "github.com/ipfs/go-ipfs/path"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
)

// KeyList is a general type for outputting lists of keys
type KeyList struct {
	Keys []key.Key
}

// KeyListTextMarshaler outputs a KeyList as plaintext, one key per line
func KeyListTextMarshaler(res cmds.Response) (io.Reader, error) {
	output := res.Output().(*KeyList)
	buf := new(bytes.Buffer)
	for _, key := range output.Keys {
		buf.WriteString(key.B58String() + "\n")
	}
	return buf, nil
}

var RefsCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists links (references) from an object.",
		ShortDescription: `
Lists the hashes of all the links an IPFS or IPNS object(s) contains,
with the following format:

  <link base58 hash>

NOTE: List all references recursively by using the flag '-r'.
`,
	},
	Subcommands: map[string]*cmds.Command{
		"local": RefsLocalCmd,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("ipfs-path", true, true, "Path to the object(s) to list refs from.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "Emit edges with given format. Available tokens: <src> <dst> <linkname>.").Default("<dst>"),
		cmds.BoolOption("edges", "e", "Emit edge format: `<from> -> <to>`.").Default(false),
		cmds.BoolOption("unique", "u", "Omit duplicate refs from output.").Default(false),
		cmds.BoolOption("recursive", "r", "Recursively list links of child nodes.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		unique, _, err := req.Option("unique").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		recursive, _, err := req.Option("recursive").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		edges, _, err := req.Option("edges").Bool()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		format, _, err := req.Option("format").String()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		objs, err := objectsForPaths(ctx, n, req.Arguments())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			rw := RefWriter{
				out:       out,
				DAG:       n.DAG,
				Ctx:       ctx,
				Unique:    unique,
				PrintEdge: edges,
				PrintFmt:  format,
				Recursive: recursive,
			}

			for _, o := range objs {
				if _, err := rw.WriteRefs(o); err != nil {
					out <- &RefWrapper{Err: err.Error()}
					return
				}
			}
		}()
	},
	Marshalers: refsMarshallerMap,
	Type:       RefWrapper{},
}

var RefsLocalCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Lists all local references.",
		ShortDescription: `
Displays the hashes of all local objects.
`,
	},

	Run: func(req cmds.Request, res cmds.Response) {
		ctx := req.Context()
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		// todo: make async
		allKeys, err := n.Blockstore.AllKeysChan(ctx)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))

		go func() {
			defer close(out)

			for k := range allKeys {
				out <- &RefWrapper{Ref: k.B58String()}
			}
		}()
	},
	Marshalers: refsMarshallerMap,
	Type:       RefWrapper{},
}

var refsMarshallerMap = cmds.MarshalerMap{
	cmds.Text: func(res cmds.Response) (io.Reader, error) {
		outChan, ok := res.Output().(<-chan interface{})
		if !ok {
			return nil, u.ErrCast()
		}

		marshal := func(v interface{}) (io.Reader, error) {
			obj, ok := v.(*RefWrapper)
			if !ok {
				return nil, u.ErrCast()
			}

			if obj.Err != "" {
				return nil, errors.New(obj.Err)
			}

			return strings.NewReader(obj.Ref + "\n"), nil
		}

		return &cmds.ChannelMarshaler{
			Channel:   outChan,
			Marshaler: marshal,
			Res:       res,
		}, nil
	},
}

func objectsForPaths(ctx context.Context, n *core.IpfsNode, paths []string) ([]*dag.Node, error) {
	objects := make([]*dag.Node, len(paths))
	for i, p := range paths {
		o, err := core.Resolve(ctx, n, path.Path(p))
		if err != nil {
			return nil, err
		}
		objects[i] = o
	}
	return objects, nil
}

type RefWrapper struct {
	Ref string
	Err string
}

type RefWriter struct {
	out chan interface{}
	DAG dag.DAGService
	Ctx context.Context

	Unique    bool
	Recursive bool
	PrintEdge bool
	PrintFmt  string

	seen map[key.Key]struct{}
}

// WriteRefs writes refs of the given object to the underlying writer.
func (rw *RefWriter) WriteRefs(n *dag.Node) (int, error) {
	if rw.Recursive {
		return rw.writeRefsRecursive(n)
	}
	return rw.writeRefsSingle(n)
}

func (rw *RefWriter) writeRefsRecursive(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	var count int
	for i, ng := range dag.GetDAG(rw.Ctx, rw.DAG, n) {
		lk := key.Key(n.Links[i].Hash)
		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, n.Links[i].Name); err != nil {
			return count, err
		}

		nd, err := ng.Get(rw.Ctx)
		if err != nil {
			return count, err
		}

		c, err := rw.writeRefsRecursive(nd)
		count += c
		if err != nil {
			return count, err
		}
	}
	return count, nil
}

func (rw *RefWriter) writeRefsSingle(n *dag.Node) (int, error) {
	nkey, err := n.Key()
	if err != nil {
		return 0, err
	}

	if rw.skip(nkey) {
		return 0, nil
	}

	count := 0
	for _, l := range n.Links {
		lk := key.Key(l.Hash)

		if rw.skip(lk) {
			continue
		}

		if err := rw.WriteEdge(nkey, lk, l.Name); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// skip returns whether to skip a key
func (rw *RefWriter) skip(k key.Key) bool {
	if !rw.Unique {
		return false
	}

	if rw.seen == nil {
		rw.seen = make(map[key.Key]struct{})
	}

	_, found := rw.seen[k]
	if !found {
		rw.seen[k] = struct{}{}
	}
	return found
}

// Write one edge
func (rw *RefWriter) WriteEdge(from, to key.Key, linkname string) error {
	if rw.Ctx != nil {
		select {
		case <-rw.Ctx.Done(): // just in case.
			return rw.Ctx.Err()
		default:
		}
	}

	var s string
	switch {
	case rw.PrintFmt != "":
		s = rw.PrintFmt
		s = strings.Replace(s, "<src>", from.B58String(), -1)
		s = strings.Replace(s, "<dst>", to.B58String(), -1)
		s = strings.Replace(s, "<linkname>", linkname, -1)
	case rw.PrintEdge:
		s = from.B58String() + " -> " + to.B58String()
	default:
		s += to.B58String()
	}

	rw.out <- &RefWrapper{Ref: s}
	return nil
}
<|endoftext|>


<commit_before>package commands

import (
	cmds "github.com/jbenet/go-ipfs/commands"
	"strings"
)

type TestOutput struct {
	Foo string
	Bar int
}

var Root = &cmds.Command{
	Options: []cmds.Option{
		cmds.Option{[]string{"config", "c"}, cmds.String},
		cmds.Option{[]string{"debug", "D"}, cmds.Bool},
	},
	Help: `ipfs - global versioned p2p merkledag file system

Basic commands:

    init          Initialize ipfs local configuration.
    add <path>    Add an object to ipfs.
    cat <ref>     Show ipfs object data.
    ls <ref>      List links from an object.
    refs <ref>    List link hashes from an object.

Tool commands:

    config        Manage configuration.
    version       Show ipfs version information.
    commands      List all available commands.

Advanced Commands:

    mount         Mount an ipfs read-only mountpoint.
    serve         Serve an interface to ipfs.
    net-diag      Print network diagnostic.

Use "ipfs help <command>" for more information about a command.
`,
	Subcommands: map[string]*cmds.Command{
		"beep": &cmds.Command{
			Run: func(req cmds.Request, res cmds.Response) {
				v := TestOutput{"hello, world", 1337}
				res.SetValue(v)
			},
		},
		"boop": &cmds.Command{
			Run: func(req cmds.Request, res cmds.Response) {
				v := strings.NewReader("hello, world")
				res.SetValue(v)
			},
		},
	},
}
<commit_msg>core/commands: Added more advanced test subcommand<commit_after>package commands

import (
	"fmt"
	cmds "github.com/jbenet/go-ipfs/commands"
	"strings"
)

type TestOutput struct {
	Foo string
	Bar int
}

var Root = &cmds.Command{
	Options: []cmds.Option{
		cmds.Option{[]string{"config", "c"}, cmds.String},
		cmds.Option{[]string{"debug", "D"}, cmds.Bool},
	},
	Help: `ipfs - global versioned p2p merkledag file system

Basic commands:

    init          Initialize ipfs local configuration.
    add <path>    Add an object to ipfs.
    cat <ref>     Show ipfs object data.
    ls <ref>      List links from an object.
    refs <ref>    List link hashes from an object.

Tool commands:

    config        Manage configuration.
    version       Show ipfs version information.
    commands      List all available commands.

Advanced Commands:

    mount         Mount an ipfs read-only mountpoint.
    serve         Serve an interface to ipfs.
    net-diag      Print network diagnostic.

Use "ipfs help <command>" for more information about a command.
`,
	Subcommands: map[string]*cmds.Command{
		"beep": &cmds.Command{
			Run: func(req cmds.Request, res cmds.Response) {
				v := TestOutput{"hello, world", 1337}
				res.SetValue(v)
			},
		},
		"boop": &cmds.Command{
			Run: func(req cmds.Request, res cmds.Response) {
				v := strings.NewReader("hello, world")
				res.SetValue(v)
			},
		},
		"warp": &cmds.Command{
			Options: []cmds.Option{
				cmds.Option{[]string{"power", "p"}, cmds.Float},
			},
			Run: func(req cmds.Request, res cmds.Response) {
				threshold := 1.21

				if power, found := req.Option("power"); found && power.(float64) >= threshold {
					res.SetValue(struct {
						Status string
						Power  float64
					}{"Flux capacitor activated!", power.(float64)})

				} else {
					err := fmt.Errorf("Insufficient power (%v jiggawatts required)", threshold)
					res.SetError(err, cmds.ErrClient)
				}
			},
		},
	},
}
<|endoftext|>


<commit_before>// craw master module
package spider

import (
    "github.com/hu17889/go_spider/core/common/mlog"
    "github.com/hu17889/go_spider/core/common/page"
    "github.com/hu17889/go_spider/core/common/page_items"
    "github.com/hu17889/go_spider/core/common/request"
    "github.com/hu17889/go_spider/core/common/resource_manage"
    "github.com/hu17889/go_spider/core/downloader"
    "github.com/hu17889/go_spider/core/page_processer"
    "github.com/hu17889/go_spider/core/pipeline"
    "github.com/hu17889/go_spider/core/scheduler"
    "math/rand"
    //"net/http"
    "time"
    //"fmt"
)

type Spider struct {
    taskname string

    pPageProcesser page_processer.PageProcesser

    pDownloader downloader.Downloader

    pScheduler scheduler.Scheduler

    pPiplelines []pipeline.Pipeline

    mc  resource_manage.ResourceManage

    threadnum uint

    exitWhenComplete bool

    // Sleeptype can be fixed or rand.
    startSleeptime uint
    endSleeptime   uint
    sleeptype      string
}

// Spider is scheduler module for all the other modules, like downloader, pipeline, scheduler and etc.
// The taskname could be empty string too, or it can be used in Pipeline for record the result crawled by which task;
func NewSpider(pageinst page_processer.PageProcesser, taskname string) *Spider {
    mlog.StraceInst().Open()

    ap := &Spider{taskname: taskname, pPageProcesser: pageinst}

    // init filelog.
    ap.CloseFileLog()
    ap.exitWhenComplete = true
    ap.sleeptype = "fixed"
    ap.startSleeptime = 0

    // init spider
    if ap.pScheduler == nil {
        ap.SetScheduler(scheduler.NewQueueScheduler(false))
    }

    if ap.pDownloader == nil {
        ap.SetDownloader(downloader.NewHttpDownloader())
    }

    mlog.StraceInst().Println("** start spider **")
    ap.pPiplelines = make([]pipeline.Pipeline, 0)

    return ap
}

func (this *Spider) Taskname() string {
    return this.taskname
}

// Deal with one url and return the PageItems.
func (this *Spider) Get(url string, respType string) *page_items.PageItems {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    return this.GetByRequest(req)
}

// Deal with several urls and return the PageItems slice.
func (this *Spider) GetAll(urls []string, respType string) []*page_items.PageItems {
    for _, u := range urls {
        req := request.NewRequest(u, respType, "", "GET", "", nil, nil, nil, nil)
        this.AddRequest(req)
    }

    pip := pipeline.NewCollectPipelinePageItems()
    this.AddPipeline(pip)

    this.Run()

    return pip.GetCollected()
}

// Deal with one url and return the PageItems with other setting.
func (this *Spider) GetByRequest(req *request.Request) *page_items.PageItems {
    var reqs []*request.Request
    reqs = append(reqs, req)
    items := this.GetAllByRequest(reqs)
    if len(items) != 0 {
        return items[0]
    }
    return nil
}

// Deal with several urls and return the PageItems slice
func (this *Spider) GetAllByRequest(reqs []*request.Request) []*page_items.PageItems {
    // push url
    for _, req := range reqs {
        //req := request.NewRequest(u, respType, urltag, method, postdata, header, cookies)
        this.AddRequest(req)
    }

    pip := pipeline.NewCollectPipelinePageItems()
    this.AddPipeline(pip)

    this.Run()

    return pip.GetCollected()
}

func (this *Spider) Run() {
    if this.threadnum == 0 {
        this.threadnum = 1
    }
    this.mc = resource_manage.NewResourceManageChan(this.threadnum)

	//init db  by sorawa

    for {
        req := this.pScheduler.Poll()

        // mc is not atomic
        if this.mc.Has() == 0 && req == nil && this.exitWhenComplete {
            mlog.StraceInst().Println("** end spider **")
            break
        } else if req == nil {
            time.Sleep(500 * time.Millisecond)
            //mlog.StraceInst().Println("scheduler is empty")
            continue
        }
        this.mc.GetOne()

        // Asynchronous fetching
        go func(req *request.Request) {
            defer this.mc.FreeOne()
            //time.Sleep( time.Duration(rand.Intn(5)) * time.Second)
            mlog.StraceInst().Println("start crawl : " + req.GetUrl())
            this.pageProcess(req)
        }(req)
    }
    this.close()
}

func (this *Spider) close() {
    this.SetScheduler(scheduler.NewQueueScheduler(false))
    this.SetDownloader(downloader.NewHttpDownloader())
    this.pPiplelines = make([]pipeline.Pipeline, 0)
    this.exitWhenComplete = true
}

func (this *Spider) AddPipeline(p pipeline.Pipeline) *Spider {
    this.pPiplelines = append(this.pPiplelines, p)
    return this
}

func (this *Spider) SetScheduler(s scheduler.Scheduler) *Spider {
    this.pScheduler = s
    return this
}

func (this *Spider) GetScheduler() scheduler.Scheduler {
    return this.pScheduler
}

func (this *Spider) SetDownloader(d downloader.Downloader) *Spider {
    this.pDownloader = d
    return this
}

func (this *Spider) GetDownloader() downloader.Downloader {
    return this.pDownloader
}

func (this *Spider) SetThreadnum(i uint) *Spider {
    this.threadnum = i
    return this
}

func (this *Spider) GetThreadnum() uint {
    return this.threadnum
}

// If exit when each crawl task is done.
// If you want to keep spider in memory all the time and add url from outside, you can set it true.
func (this *Spider) SetExitWhenComplete(e bool) *Spider {
    this.exitWhenComplete = e
    return this
}

func (this *Spider) GetExitWhenComplete() bool {
    return this.exitWhenComplete
}

// The OpenFileLog initialize the log path and open log.
// If log is opened, error info or other useful info in spider will be logged in file of the filepath.
// Log command is mlog.LogInst().LogError("info") or mlog.LogInst().LogInfo("info").
// Spider's default log is closed.
// The filepath is absolute path.
func (this *Spider) OpenFileLog(filePath string) *Spider {
    mlog.InitFilelog(true, filePath)
    return this
}

// OpenFileLogDefault open file log with default file path like "WD/log/log.2014-9-1".
func (this *Spider) OpenFileLogDefault() *Spider {
    mlog.InitFilelog(true, "")
    return this
}

// The CloseFileLog close file log.
func (this *Spider) CloseFileLog() *Spider {
    mlog.InitFilelog(false, "")
    return this
}

// The OpenStrace open strace that output progress info on the screen.
// Spider's default strace is opened.
func (this *Spider) OpenStrace() *Spider {
    mlog.StraceInst().Open()
    return this
}

// The CloseStrace close strace.
func (this *Spider) CloseStrace() *Spider {
    mlog.StraceInst().Close()
    return this
}

// The SetSleepTime set sleep time after each crawl task.
// The unit is millisecond.
// If sleeptype is "fixed", the s is the sleep time and e is useless.
// If sleeptype is "rand", the sleep time is rand between s and e.
func (this *Spider) SetSleepTime(sleeptype string, s uint, e uint) *Spider {
    this.sleeptype = sleeptype
    this.startSleeptime = s
    this.endSleeptime = e
    if this.sleeptype == "rand" && this.startSleeptime >= this.endSleeptime {
        panic("startSleeptime must smaller than endSleeptime")
    }
    return this
}

func (this *Spider) sleep() {
    if this.sleeptype == "fixed" {
        time.Sleep(time.Duration(this.startSleeptime) * time.Millisecond)
    } else if this.sleeptype == "rand" {
        sleeptime := rand.Intn(int(this.endSleeptime-this.startSleeptime)) + int(this.startSleeptime)
        time.Sleep(time.Duration(sleeptime) * time.Millisecond)
    }
}

func (this *Spider) AddUrl(url string, respType string) *Spider {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    this.AddRequest(req)
    return this
}

func (this *Spider) AddUrlEx(url string, respType string,headerFile string,proxyHost string) *Spider {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    this.AddRequest(req.AddHeaderFile(headerFile).AddProxyHost(proxyHost))
    return this
}

func (this *Spider) AddUrlWithHeaderFile(url string, respType string,headerFile string) *Spider {
    req := request.NewRequestWithHeaderFile(url, respType, headerFile)
    this.AddRequest(req)
    return this
}

func (this *Spider) AddUrls(urls []string, respType string) *Spider {
    for _, url := range urls {
        req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
        this.AddRequest(req)
    }
    return this
}

func (this *Spider) AddUrlsWithHeaderFile(urls []string, respType string,headerFile string) *Spider {
	for _, url := range urls {
		req := request.NewRequestWithHeaderFile(url, respType, headerFile)
		this.AddRequest(req)
	}
    return this
}

func (this *Spider) AddUrlsEx(urls []string, respType string,headerFile string,proxyHost string) *Spider {
	for _, url := range urls {
		req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
		this.AddRequest(req.AddHeaderFile(headerFile).AddProxyHost(proxyHost))
	}
    return this
}

// add Request to Schedule
func (this *Spider) AddRequest(req *request.Request) *Spider {
    if req == nil {
        mlog.LogInst().LogError("request is nil")
        return this
    } else if req.GetUrl() == "" {
        mlog.LogInst().LogError("request is empty")
        return this
    }
    this.pScheduler.Push(req)
    return this
}

//
func (this *Spider) AddRequests(reqs []*request.Request) *Spider {
    for _, req := range reqs {
        this.AddRequest(req)
    }
    return this
}

// core processer
func (this *Spider) pageProcess(req *request.Request) {
    var p *page.Page

    defer func() {
        if err := recover(); err != nil { // do not affect other
            if strerr, ok := err.(string); ok {
                mlog.LogInst().LogError(strerr)
            } else {
                mlog.LogInst().LogError("pageProcess error")
            }
        }
    }()

    // download page
    for i := 0; i < 3; i++ {
        this.sleep()
        p = this.pDownloader.Download(req)
        if p.IsSucc() { // if fail retry 3 times
            break
        }

    }

    if !p.IsSucc() { // if fail do not need process
        return
    }

    this.pPageProcesser.Process(p)
    for _, req := range p.GetTargetRequests() {
        this.AddRequest(req)
    }

    // output
    if !p.GetSkip() {
        for _, pip := range this.pPiplelines {
            pip.Process(p.GetPageItems(), this)
        }
    }
}
<commit_msg>add func finish()  before end spider<commit_after>// craw master module
package spider

import (
    "github.com/hu17889/go_spider/core/common/mlog"
    "github.com/hu17889/go_spider/core/common/page"
    "github.com/hu17889/go_spider/core/common/page_items"
    "github.com/hu17889/go_spider/core/common/request"
    "github.com/hu17889/go_spider/core/common/resource_manage"
    "github.com/hu17889/go_spider/core/downloader"
    "github.com/hu17889/go_spider/core/page_processer"
    "github.com/hu17889/go_spider/core/pipeline"
    "github.com/hu17889/go_spider/core/scheduler"
    "math/rand"
    //"net/http"
    "time"
    //"fmt"
)

type Spider struct {
    taskname string

    pPageProcesser page_processer.PageProcesser

    pDownloader downloader.Downloader

    pScheduler scheduler.Scheduler

    pPiplelines []pipeline.Pipeline

    mc  resource_manage.ResourceManage

    threadnum uint

    exitWhenComplete bool

    // Sleeptype can be fixed or rand.
    startSleeptime uint
    endSleeptime   uint
    sleeptype      string
}

// Spider is scheduler module for all the other modules, like downloader, pipeline, scheduler and etc.
// The taskname could be empty string too, or it can be used in Pipeline for record the result crawled by which task;
func NewSpider(pageinst page_processer.PageProcesser, taskname string) *Spider {
    mlog.StraceInst().Open()

    ap := &Spider{taskname: taskname, pPageProcesser: pageinst}

    // init filelog.
    ap.CloseFileLog()
    ap.exitWhenComplete = true
    ap.sleeptype = "fixed"
    ap.startSleeptime = 0

    // init spider
    if ap.pScheduler == nil {
        ap.SetScheduler(scheduler.NewQueueScheduler(false))
    }

    if ap.pDownloader == nil {
        ap.SetDownloader(downloader.NewHttpDownloader())
    }

    mlog.StraceInst().Println("** start spider **")
    ap.pPiplelines = make([]pipeline.Pipeline, 0)

    return ap
}

func (this *Spider) Taskname() string {
    return this.taskname
}

// Deal with one url and return the PageItems.
func (this *Spider) Get(url string, respType string) *page_items.PageItems {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    return this.GetByRequest(req)
}

// Deal with several urls and return the PageItems slice.
func (this *Spider) GetAll(urls []string, respType string) []*page_items.PageItems {
    for _, u := range urls {
        req := request.NewRequest(u, respType, "", "GET", "", nil, nil, nil, nil)
        this.AddRequest(req)
    }

    pip := pipeline.NewCollectPipelinePageItems()
    this.AddPipeline(pip)

    this.Run()

    return pip.GetCollected()
}

// Deal with one url and return the PageItems with other setting.
func (this *Spider) GetByRequest(req *request.Request) *page_items.PageItems {
    var reqs []*request.Request
    reqs = append(reqs, req)
    items := this.GetAllByRequest(reqs)
    if len(items) != 0 {
        return items[0]
    }
    return nil
}

// Deal with several urls and return the PageItems slice
func (this *Spider) GetAllByRequest(reqs []*request.Request) []*page_items.PageItems {
    // push url
    for _, req := range reqs {
        //req := request.NewRequest(u, respType, urltag, method, postdata, header, cookies)
        this.AddRequest(req)
    }

    pip := pipeline.NewCollectPipelinePageItems()
    this.AddPipeline(pip)

    this.Run()

    return pip.GetCollected()
}

func (this *Spider) Run() {
    if this.threadnum == 0 {
        this.threadnum = 1
    }
    this.mc = resource_manage.NewResourceManageChan(this.threadnum)

	//init db  by sorawa

    for {
        req := this.pScheduler.Poll()

        // mc is not atomic
        if this.mc.Has() == 0 && req == nil && this.exitWhenComplete {
	    mlog.StraceInst().Println("** executed callback **")
	    this.pPageProcesser.Finish()
            mlog.StraceInst().Println("** end spider **")
            break
        } else if req == nil {
            time.Sleep(500 * time.Millisecond)
            //mlog.StraceInst().Println("scheduler is empty")
            continue
        }
        this.mc.GetOne()

        // Asynchronous fetching
        go func(req *request.Request) {
            defer this.mc.FreeOne()
            //time.Sleep( time.Duration(rand.Intn(5)) * time.Second)
            mlog.StraceInst().Println("start crawl : " + req.GetUrl())
            this.pageProcess(req)
        }(req)
    }
    this.close()
}

func (this *Spider) close() {
    this.SetScheduler(scheduler.NewQueueScheduler(false))
    this.SetDownloader(downloader.NewHttpDownloader())
    this.pPiplelines = make([]pipeline.Pipeline, 0)
    this.exitWhenComplete = true
}

func (this *Spider) AddPipeline(p pipeline.Pipeline) *Spider {
    this.pPiplelines = append(this.pPiplelines, p)
    return this
}

func (this *Spider) SetScheduler(s scheduler.Scheduler) *Spider {
    this.pScheduler = s
    return this
}

func (this *Spider) GetScheduler() scheduler.Scheduler {
    return this.pScheduler
}

func (this *Spider) SetDownloader(d downloader.Downloader) *Spider {
    this.pDownloader = d
    return this
}

func (this *Spider) GetDownloader() downloader.Downloader {
    return this.pDownloader
}

func (this *Spider) SetThreadnum(i uint) *Spider {
    this.threadnum = i
    return this
}

func (this *Spider) GetThreadnum() uint {
    return this.threadnum
}

// If exit when each crawl task is done.
// If you want to keep spider in memory all the time and add url from outside, you can set it true.
func (this *Spider) SetExitWhenComplete(e bool) *Spider {
    this.exitWhenComplete = e
    return this
}

func (this *Spider) GetExitWhenComplete() bool {
    return this.exitWhenComplete
}

// The OpenFileLog initialize the log path and open log.
// If log is opened, error info or other useful info in spider will be logged in file of the filepath.
// Log command is mlog.LogInst().LogError("info") or mlog.LogInst().LogInfo("info").
// Spider's default log is closed.
// The filepath is absolute path.
func (this *Spider) OpenFileLog(filePath string) *Spider {
    mlog.InitFilelog(true, filePath)
    return this
}

// OpenFileLogDefault open file log with default file path like "WD/log/log.2014-9-1".
func (this *Spider) OpenFileLogDefault() *Spider {
    mlog.InitFilelog(true, "")
    return this
}

// The CloseFileLog close file log.
func (this *Spider) CloseFileLog() *Spider {
    mlog.InitFilelog(false, "")
    return this
}

// The OpenStrace open strace that output progress info on the screen.
// Spider's default strace is opened.
func (this *Spider) OpenStrace() *Spider {
    mlog.StraceInst().Open()
    return this
}

// The CloseStrace close strace.
func (this *Spider) CloseStrace() *Spider {
    mlog.StraceInst().Close()
    return this
}

// The SetSleepTime set sleep time after each crawl task.
// The unit is millisecond.
// If sleeptype is "fixed", the s is the sleep time and e is useless.
// If sleeptype is "rand", the sleep time is rand between s and e.
func (this *Spider) SetSleepTime(sleeptype string, s uint, e uint) *Spider {
    this.sleeptype = sleeptype
    this.startSleeptime = s
    this.endSleeptime = e
    if this.sleeptype == "rand" && this.startSleeptime >= this.endSleeptime {
        panic("startSleeptime must smaller than endSleeptime")
    }
    return this
}

func (this *Spider) sleep() {
    if this.sleeptype == "fixed" {
        time.Sleep(time.Duration(this.startSleeptime) * time.Millisecond)
    } else if this.sleeptype == "rand" {
        sleeptime := rand.Intn(int(this.endSleeptime-this.startSleeptime)) + int(this.startSleeptime)
        time.Sleep(time.Duration(sleeptime) * time.Millisecond)
    }
}

func (this *Spider) AddUrl(url string, respType string) *Spider {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    this.AddRequest(req)
    return this
}

func (this *Spider) AddUrlEx(url string, respType string,headerFile string,proxyHost string) *Spider {
    req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
    this.AddRequest(req.AddHeaderFile(headerFile).AddProxyHost(proxyHost))
    return this
}

func (this *Spider) AddUrlWithHeaderFile(url string, respType string,headerFile string) *Spider {
    req := request.NewRequestWithHeaderFile(url, respType, headerFile)
    this.AddRequest(req)
    return this
}

func (this *Spider) AddUrls(urls []string, respType string) *Spider {
    for _, url := range urls {
        req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
        this.AddRequest(req)
    }
    return this
}

func (this *Spider) AddUrlsWithHeaderFile(urls []string, respType string,headerFile string) *Spider {
	for _, url := range urls {
		req := request.NewRequestWithHeaderFile(url, respType, headerFile)
		this.AddRequest(req)
	}
    return this
}

func (this *Spider) AddUrlsEx(urls []string, respType string,headerFile string,proxyHost string) *Spider {
	for _, url := range urls {
		req := request.NewRequest(url, respType, "", "GET", "", nil, nil, nil, nil)
		this.AddRequest(req.AddHeaderFile(headerFile).AddProxyHost(proxyHost))
	}
    return this
}

// add Request to Schedule
func (this *Spider) AddRequest(req *request.Request) *Spider {
    if req == nil {
        mlog.LogInst().LogError("request is nil")
        return this
    } else if req.GetUrl() == "" {
        mlog.LogInst().LogError("request is empty")
        return this
    }
    this.pScheduler.Push(req)
    return this
}

//
func (this *Spider) AddRequests(reqs []*request.Request) *Spider {
    for _, req := range reqs {
        this.AddRequest(req)
    }
    return this
}

// core processer
func (this *Spider) pageProcess(req *request.Request) {
    var p *page.Page

    defer func() {
        if err := recover(); err != nil { // do not affect other
            if strerr, ok := err.(string); ok {
                mlog.LogInst().LogError(strerr)
            } else {
                mlog.LogInst().LogError("pageProcess error")
            }
        }
    }()

    // download page
    for i := 0; i < 3; i++ {
        this.sleep()
        p = this.pDownloader.Download(req)
        if p.IsSucc() { // if fail retry 3 times
            break
        }

    }

    if !p.IsSucc() { // if fail do not need process
        return
    }

    this.pPageProcesser.Process(p)
    for _, req := range p.GetTargetRequests() {
        this.AddRequest(req)
    }

    // output
    if !p.GetSkip() {
        for _, pip := range this.pPiplelines {
            pip.Process(p.GetPageItems(), this)
        }
    }
}
<|endoftext|>


<commit_before>package main

import "net/http"

import (
	"errors"
	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"
	"github.com/mitchellh/mapstructure"
)

// KeyExists will check if the key being used to access the API is in the request data,
// and then if the key is in the storage engine
type AuthKey struct {
	TykMiddleware
}

type AuthKeyConfiguration struct {
	Auth struct {
		AuthHeaderName string `mapstructure:"auth_header_name" bson:"auth_header_name" json:"auth_header_name"`
	} `mapstructure:"auth" bson:"auth" json:"auth"`
}

func (k AuthKey) New() {}

// GetConfig retrieves the configuration from the API config
func (k *AuthKey) GetConfig() (interface{}, error) {
	var thisModuleConfig AuthKeyConfiguration

	err := mapstructure.Decode(k.TykMiddleware.Spec.APIDefinition.RawData, &thisModuleConfig)
	if err != nil {
		log.Error(err)
		return nil, err
	}

	return thisModuleConfig, nil
}

func (k *AuthKey) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	var thisConfig AuthKeyConfiguration
	thisConfig = configuration.(AuthKeyConfiguration)

	authHeaderValue := r.Header.Get(thisConfig.Auth.AuthHeaderName)
	if authHeaderValue == "" {
		// No header value, fail
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
		}).Info("Attempted access with malformed header, no auth header found.")

		return errors.New("Authorization field missing"), 400
	}

	// Check if API key valid
	thisSessionState, keyExists := k.TykMiddleware.CheckSessionAndIdentityForValidKey(authHeaderValue)
	if !keyExists {
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
			"key":    authHeaderValue,
		}).Info("Attempted access with non-existent key.")

		// Fire Authfailed Event
		AuthFailed(k.TykMiddleware, r, authHeaderValue)

		return errors.New("Key not authorised"), 403
	}

	// Set session state on context, we will need it later
	context.Set(r, SessionData, thisSessionState)
	context.Set(r, AuthHeaderValue, authHeaderValue)

	return nil, 200
}

func AuthFailed (m TykMiddleware, r *http.Request, authHeaderValue string) {
	go m.FireEvent(EVENT_AuthFailure,
		EVENT_AuthFailureMeta{
		EventMetaDefault: EventMetaDefault{Message: "Auth Failure"},
		Path: r.URL.Path,
		Origin: r.RemoteAddr,
		Key: authHeaderValue,
	})
}
<commit_msg>Moved some middleware data back into main<commit_after>package main

import "net/http"

import (
	"errors"
	"github.com/Sirupsen/logrus"
	"github.com/gorilla/context"
)

// KeyExists will check if the key being used to access the API is in the request data,
// and then if the key is in the storage engine
type AuthKey struct {
	TykMiddleware
}

func (k AuthKey) New() {}

// GetConfig retrieves the configuration from the API config
func (k *AuthKey) GetConfig() (interface{}, error) {
	return k.TykMiddleware.Spec.APIDefinition.Auth, nil
}

func (k *AuthKey) ProcessRequest(w http.ResponseWriter, r *http.Request, configuration interface{}) (error, int) {
	thisConfig := k.TykMiddleware.Spec.APIDefinition.Auth

	authHeaderValue := r.Header.Get(thisConfig.AuthHeaderName)
	if authHeaderValue == "" {
		// No header value, fail
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
		}).Info("Attempted access with malformed header, no auth header found.")

		return errors.New("Authorization field missing"), 400
	}

	// Check if API key valid
	thisSessionState, keyExists := k.TykMiddleware.CheckSessionAndIdentityForValidKey(authHeaderValue)
	if !keyExists {
		log.WithFields(logrus.Fields{
			"path":   r.URL.Path,
			"origin": r.RemoteAddr,
			"key":    authHeaderValue,
		}).Info("Attempted access with non-existent key.")

		// Fire Authfailed Event
		AuthFailed(k.TykMiddleware, r, authHeaderValue)

		return errors.New("Key not authorised"), 403
	}

	// Set session state on context, we will need it later
	context.Set(r, SessionData, thisSessionState)
	context.Set(r, AuthHeaderValue, authHeaderValue)

	return nil, 200
}

func AuthFailed (m TykMiddleware, r *http.Request, authHeaderValue string) {
	go m.FireEvent(EVENT_AuthFailure,
		EVENT_AuthFailureMeta{
		EventMetaDefault: EventMetaDefault{Message: "Auth Failure"},
		Path: r.URL.Path,
		Origin: r.RemoteAddr,
		Key: authHeaderValue,
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package resource

import (
	"fmt"

	"k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/watch"
)

// Selector is a Visitor for resources that match a label selector.
type Selector struct {
	Client    RESTClient
	Mapping   *meta.RESTMapping
	Namespace string
	Selector  labels.Selector
	Export    bool
}

// NewSelector creates a resource selector which hides details of getting items by their label selector.
func NewSelector(client RESTClient, mapping *meta.RESTMapping, namespace string, selector labels.Selector, export bool) *Selector {
	return &Selector{
		Client:    client,
		Mapping:   mapping,
		Namespace: namespace,
		Selector:  selector,
		Export:    export,
	}
}

// Visit implements Visitor
func (r *Selector) Visit(fn VisitorFunc) error {
	list, err := NewHelper(r.Client, r.Mapping).List(r.Namespace, r.ResourceMapping().GroupVersionKind.GroupVersion().String(), r.Selector, r.Export)
	if err != nil {
		if errors.IsBadRequest(err) || errors.IsNotFound(err) {
			if r.Selector.Empty() {
				return fmt.Errorf("Unable to list %q: %v", r.Mapping.Resource, err)
			} else {
				return fmt.Errorf("Unable to find %q that match the selector %q: %v", r.Mapping.Resource, r.Selector, err)
			}
		}
		return err
	}
	accessor := r.Mapping.MetadataAccessor
	resourceVersion, _ := accessor.ResourceVersion(list)
	info := &Info{
		Client:    r.Client,
		Mapping:   r.Mapping,
		Namespace: r.Namespace,

		Object:          list,
		ResourceVersion: resourceVersion,
	}
	return fn(info, nil)
}

func (r *Selector) Watch(resourceVersion string) (watch.Interface, error) {
	return NewHelper(r.Client, r.Mapping).Watch(r.Namespace, resourceVersion, r.ResourceMapping().GroupVersionKind.GroupVersion().String(), r.Selector)
}

// ResourceMapping returns the mapping for this resource and implements ResourceMapping
func (r *Selector) ResourceMapping() *meta.RESTMapping {
	return r.Mapping
}
<commit_msg>UPSTREAM: 27243: Don't alter error type from server<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package resource

import (
	"fmt"

	"k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/meta"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/watch"
)

// Selector is a Visitor for resources that match a label selector.
type Selector struct {
	Client    RESTClient
	Mapping   *meta.RESTMapping
	Namespace string
	Selector  labels.Selector
	Export    bool
}

// NewSelector creates a resource selector which hides details of getting items by their label selector.
func NewSelector(client RESTClient, mapping *meta.RESTMapping, namespace string, selector labels.Selector, export bool) *Selector {
	return &Selector{
		Client:    client,
		Mapping:   mapping,
		Namespace: namespace,
		Selector:  selector,
		Export:    export,
	}
}

// Visit implements Visitor
func (r *Selector) Visit(fn VisitorFunc) error {
	list, err := NewHelper(r.Client, r.Mapping).List(r.Namespace, r.ResourceMapping().GroupVersionKind.GroupVersion().String(), r.Selector, r.Export)
	if err != nil {
		if errors.IsBadRequest(err) || errors.IsNotFound(err) {
			if se, ok := err.(*errors.StatusError); ok {
				// modify the message without hiding this is an API error
				if r.Selector.Empty() {
					se.ErrStatus.Message = fmt.Sprintf("Unable to list %q: %v", r.Mapping.Resource, se.ErrStatus.Message)
				} else {
					se.ErrStatus.Message = fmt.Sprintf("Unable to find %q that match the selector %q: %v", r.Mapping.Resource, r.Selector, se.ErrStatus.Message)
				}
				return se
			}
			if r.Selector.Empty() {
				return fmt.Errorf("Unable to list %q: %v", r.Mapping.Resource, err)
			} else {
				return fmt.Errorf("Unable to find %q that match the selector %q: %v", r.Mapping.Resource, r.Selector, err)
			}
		}
		return err
	}
	accessor := r.Mapping.MetadataAccessor
	resourceVersion, _ := accessor.ResourceVersion(list)
	info := &Info{
		Client:    r.Client,
		Mapping:   r.Mapping,
		Namespace: r.Namespace,

		Object:          list,
		ResourceVersion: resourceVersion,
	}
	return fn(info, nil)
}

func (r *Selector) Watch(resourceVersion string) (watch.Interface, error) {
	return NewHelper(r.Client, r.Mapping).Watch(r.Namespace, resourceVersion, r.ResourceMapping().GroupVersionKind.GroupVersion().String(), r.Selector)
}

// ResourceMapping returns the mapping for this resource and implements ResourceMapping
func (r *Selector) ResourceMapping() *meta.RESTMapping {
	return r.Mapping
}
<|endoftext|>


<commit_before>// Package events provides the ability to retrieve Google+ events for a signed-user
// and a trigger to aggregate all event pictures by commmunicating with the Compute Engine.
package events

import (
	"appengine"
	"appengine/urlfetch"
	"appengine/taskqueue"

	"strings"
	"time"

	"code.google.com/p/goauth2/oauth"
	"code.google.com/p/google-api-go-client/calendar/v3"

	"socialvibes/config"
	"socialvibes/model"
)

// ParseEventInfos uses the Google Calendar API to aggregate all event information since a month ago.
// It picks out all Google+ events by looking for the corresponding event link.
// It returns all found events as array and any error encountered.
func ParseEventInfos(transport *oauth.Transport, context appengine.Context) (events []*model.Event, err error) {
	// Create an API service
	service, err := calendar.New(transport.Client())
	if err != nil {
		context.Errorf("events > event.go > ParseEventInfos > calendar.New(): %v", err)
		return events, err
	}

	// Create an event list call for the primary calendar of the given user
	listCall := service.Events.List("primary")

	// Get all calendar events since a month ago
	t := time.Now().Add(time.Hour * 24 * 31 * (-1))
	listCall.TimeMin(t.Format(time.RFC3339))

	// Execute the event list call
	eventList, err := listCall.Do()
	if err != nil {
		context.Errorf("events > event.go > ParseEventInfos > listCall.Do(): %v", err)
		return events, err
	}

	// Extract all events from the list
	items := eventList.Items

	// Iterate over all found events
	for _, item := range items {
		// Find all Google+ events by filtering htmllink
		if strings.HasPrefix(item.HtmlLink, "https://plus.google.com/events/") {
			// Create a new event object for the found calendar event and fill it with all found information
			event := new(model.Event)
			event.Url = item.HtmlLink
			event.Id = "c" + item.Id
			event.Name = item.Summary
			event.Start = item.Start.DateTime
			event.End = item.End.DateTime
			event.Visibility = item.Visibility
			event.Creator = item.Creator.Self
			event.Location = item.Location

			events = append(events, event)
		}
	}
	return events, nil
}

// RefreshEventGallery creates a task (for the App Engine Task Queue),
// which will be executed by the Compute Engine to (re-)aggregate all event pictures for the given event.
// It returns any error encountered.
func RefreshEventGallery(context appengine.Context, eventId string) (err error) {
	// Create a task for a pull queue
	// Tags are used for grouping tasks inside the Compute Engine
	task := &taskqueue.Task{
	    Payload: []byte(eventId),
	    Method:  "PULL",
	    Tag: eventId,
	}
	// Insert the created task into the App Engine Task Queue
	_, err = taskqueue.Add(context, task, "picturerequest")
	if err != nil {
		context.Errorf("events > event.go > RefreshEventGallery > taskqueue.Add(): %v", err)
		return
	}

    // Wait 3 seconds till the task is definitely in the task queue
    time.Sleep(3 * time.Second)

    // Notify the task consumer in the Compute Engine via RPC
    req := `{"method":"EventService.PullTask","params":[{"PullType":"picturerequest", "EventId":"` + eventId + `"}], "id":"1"}`

    // Make a secure POST request via App Engine URL Fetch service
    client := urlfetch.Client(context)
    resp, err := client.Post("http://" + *config.ComputeEngineAddress +"/rpc", "application/json", strings.NewReader(req))
    if err != nil || resp.StatusCode != 200 {
        context.Errorf("events > event.go > RefreshEventGallery > client.Post(): %v", err)
    }

	return nil
}<commit_msg>increased RPC call deadline to 60 seconds due to urlfetch: DEADLINE_EXCEEDED error<commit_after>// Package events provides the ability to retrieve Google+ events for a signed-user
// and a trigger to aggregate all event pictures by commmunicating with the Compute Engine.
package events

import (
	"appengine"
	"appengine/urlfetch"
	"appengine/taskqueue"

	"strings"
	"time"

	"code.google.com/p/goauth2/oauth"
	"code.google.com/p/google-api-go-client/calendar/v3"

	"socialvibes/config"
	"socialvibes/model"
)

// ParseEventInfos uses the Google Calendar API to aggregate all event information since a month ago.
// It picks out all Google+ events by looking for the corresponding event link.
// It returns all found events as array and any error encountered.
func ParseEventInfos(transport *oauth.Transport, context appengine.Context) (events []*model.Event, err error) {
	// Create an API service
	service, err := calendar.New(transport.Client())
	if err != nil {
		context.Errorf("events > event.go > ParseEventInfos > calendar.New(): %v", err)
		return events, err
	}

	// Create an event list call for the primary calendar of the given user
	listCall := service.Events.List("primary")

	// Get all calendar events since a month ago
	t := time.Now().Add(time.Hour * 24 * 31 * (-1))
	listCall.TimeMin(t.Format(time.RFC3339))

	// Execute the event list call
	eventList, err := listCall.Do()
	if err != nil {
		context.Errorf("events > event.go > ParseEventInfos > listCall.Do(): %v", err)
		return events, err
	}

	// Extract all events from the list
	items := eventList.Items

	// Iterate over all found events
	for _, item := range items {
		// Find all Google+ events by filtering htmllink
		if strings.HasPrefix(item.HtmlLink, "https://plus.google.com/events/") {
			// Create a new event object for the found calendar event and fill it with all found information
			event := new(model.Event)
			event.Url = item.HtmlLink
			event.Id = "c" + item.Id
			event.Name = item.Summary
			event.Start = item.Start.DateTime
			event.End = item.End.DateTime
			event.Visibility = item.Visibility
			event.Creator = item.Creator.Self
			event.Location = item.Location

			events = append(events, event)
		}
	}
	return events, nil
}

// RefreshEventGallery creates a task (for the App Engine Task Queue),
// which will be executed by the Compute Engine to (re-)aggregate all event pictures for the given event.
// It returns any error encountered.
func RefreshEventGallery(context appengine.Context, eventId string) (err error) {
	// Create a task for a pull queue
	// Tags are used for grouping tasks inside the Compute Engine
	task := &taskqueue.Task{
	    Payload: []byte(eventId),
	    Method:  "PULL",
	    Tag: eventId,
	}
	// Insert the created task into the App Engine Task Queue
	_, err = taskqueue.Add(context, task, "picturerequest")
	if err != nil {
		context.Errorf("events > event.go > RefreshEventGallery > taskqueue.Add(): %v", err)
		return
	}

    // Wait 3 seconds till the task is definitely in the task queue
    time.Sleep(3 * time.Second)

    // Notify the task consumer in the Compute Engine via RPC
    req := `{"method":"EventService.PullTask","params":[{"PullType":"picturerequest", "EventId":"` + eventId + `"}], "id":"1"}`

    // Make a secure POST request via App Engine URL Fetch service
    client := urlfetch.Client(context)
    // Increase default deadline to 60 seconds
    client.Transport.(*urlfetch.Transport).Deadline = time.Minute
    context.Infof("events > event.go > RefreshEventGallery > client.Transport.Deadline: %v", client.Transport.(*urlfetch.Transport).Deadline)

    resp, err := client.Post("http://" + *config.ComputeEngineAddress +"/rpc", "application/json", strings.NewReader(req))
    context.Infof("events > event.go > RefreshEventGallery > client.Post(); StatusCode: %v; Status: %s", resp.StatusCode, resp.Status)
    if err != nil || resp.StatusCode != 200 {
        context.Errorf("events > event.go > RefreshEventGallery > client.Post(): %v", err)
    }

	return nil
}<|endoftext|>


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

//+build !codes

package test_driver

const panicInfo = "This branch is not implemented. " +
	"This is because you are trying to test something specific to TiDB's MyDecimal implementation. " +
	"It is recommended to do this in TiDB repository."

// constant values.
const (
	maxWordBufLen = 9 // A MyDecimal holds 9 words.
	digitsPerWord = 9 // A word holds 9 digits.
	digMask       = 100_000_000
)

var (
	wordBufLen = 9
)

// fixWordCntError limits word count in wordBufLen, and returns overflow or truncate error.
func fixWordCntError(wordsInt, wordsFrac int) (newWordsInt int, newWordsFrac int, err error) {
	if wordsInt+wordsFrac > wordBufLen {
		panic(panicInfo)
	}
	return wordsInt, wordsFrac, nil
}

/*
  countLeadingZeroes returns the number of leading zeroes that can be removed from fraction.

  @param   i    start index
  @param   word value to compare against list of powers of 10
*/
func countLeadingZeroes(i int, word int32) int {
	leading := 0
	for word < pow10(i) {
		i--
		leading++
	}
	return leading
}

func digitsToWords(digits int) int {
	return (digits + digitsPerWord - 1) / digitsPerWord
}

// MyDecimal represents a decimal value.
type MyDecimal struct {
	digitsInt int8 // the number of *decimal* digits before the point.

	digitsFrac int8 // the number of decimal digits after the point.

	resultFrac int8 // result fraction digits.

	negative bool

	// wordBuf is an array of int32 words.
	// A word is an int32 value can hold 9 digits.(0 <= word < wordBase)
	wordBuf [maxWordBufLen]int32
}

// String returns the decimal string representation rounded to resultFrac.
func (d *MyDecimal) String() string {
	tmp := *d
	return string(tmp.ToString())
}

func (d *MyDecimal) stringSize() int {
	// sign, zero integer and dot.
	return int(d.digitsInt + d.digitsFrac + 3)
}

func (d *MyDecimal) removeLeadingZeros() (wordIdx int, digitsInt int) {
	digitsInt = int(d.digitsInt)
	i := ((digitsInt - 1) % digitsPerWord) + 1
	for digitsInt > 0 && d.wordBuf[wordIdx] == 0 {
		digitsInt -= i
		i = digitsPerWord
		wordIdx++
	}
	if digitsInt > 0 {
		digitsInt -= countLeadingZeroes((digitsInt-1)%digitsPerWord, d.wordBuf[wordIdx])
	} else {
		digitsInt = 0
	}
	return
}

// ToString converts decimal to its printable string representation without rounding.
//
//  RETURN VALUE
//
//      str       - result string
//      errCode   - eDecOK/eDecTruncate/eDecOverflow
//
func (d *MyDecimal) ToString() (str []byte) {
	str = make([]byte, d.stringSize())
	digitsFrac := int(d.digitsFrac)
	wordStartIdx, digitsInt := d.removeLeadingZeros()
	if digitsInt+digitsFrac == 0 {
		digitsInt = 1
		wordStartIdx = 0
	}

	digitsIntLen := digitsInt
	if digitsIntLen == 0 {
		digitsIntLen = 1
	}
	digitsFracLen := digitsFrac
	length := digitsIntLen + digitsFracLen
	if d.negative {
		length++
	}
	if digitsFrac > 0 {
		length++
	}
	str = str[:length]
	strIdx := 0
	if d.negative {
		str[strIdx] = '-'
		strIdx++
	}
	var fill int
	if digitsFrac > 0 {
		fracIdx := strIdx + digitsIntLen
		fill = digitsFracLen - digitsFrac
		wordIdx := wordStartIdx + digitsToWords(digitsInt)
		str[fracIdx] = '.'
		fracIdx++
		for ; digitsFrac > 0; digitsFrac -= digitsPerWord {
			x := d.wordBuf[wordIdx]
			wordIdx++
			for i := myMin(digitsFrac, digitsPerWord); i > 0; i-- {
				y := x / digMask
				str[fracIdx] = byte(y) + '0'
				fracIdx++
				x -= y * digMask
				x *= 10
			}
		}
		for ; fill > 0; fill-- {
			str[fracIdx] = '0'
			fracIdx++
		}
	}
	fill = digitsIntLen - digitsInt
	if digitsInt == 0 {
		fill-- /* symbol 0 before digital point */
	}
	for ; fill > 0; fill-- {
		str[strIdx] = '0'
		strIdx++
	}
	if digitsInt > 0 {
		strIdx += digitsInt
		wordIdx := wordStartIdx + digitsToWords(digitsInt)
		for ; digitsInt > 0; digitsInt -= digitsPerWord {
			wordIdx--
			x := d.wordBuf[wordIdx]
			for i := myMin(digitsInt, digitsPerWord); i > 0; i-- {
				y := x / 10
				strIdx--
				str[strIdx] = '0' + byte(x-y*10)
				x = y
			}
		}
	} else {
		str[strIdx] = '0'
	}
	return
}

// FromString parses decimal from string.
func (d *MyDecimal) FromString(str []byte) error {
	for i := 0; i < len(str); i++ {
		if !isSpace(str[i]) {
			str = str[i:]
			break
		}
	}
	if len(str) == 0 {
		panic(panicInfo)
	}
	switch str[0] {
	case '-':
		d.negative = true
		fallthrough
	case '+':
		str = str[1:]
	}
	var strIdx int
	for strIdx < len(str) && isDigit(str[strIdx]) {
		strIdx++
	}
	digitsInt := strIdx
	var digitsFrac int
	var endIdx int
	if strIdx < len(str) && str[strIdx] == '.' {
		endIdx = strIdx + 1
		for endIdx < len(str) && isDigit(str[endIdx]) {
			endIdx++
		}
		digitsFrac = endIdx - strIdx - 1
	} else {
		digitsFrac = 0
		endIdx = strIdx
	}
	if digitsInt+digitsFrac == 0 {
		panic(panicInfo)
	}
	wordsInt := digitsToWords(digitsInt)
	wordsFrac := digitsToWords(digitsFrac)
	wordsInt, wordsFrac, err := fixWordCntError(wordsInt, wordsFrac)
	if err != nil {
		panic(panicInfo)
	}
	d.digitsInt = int8(digitsInt)
	d.digitsFrac = int8(digitsFrac)
	wordIdx := wordsInt
	strIdxTmp := strIdx
	var word int32
	var innerIdx int
	for digitsInt > 0 {
		digitsInt--
		strIdx--
		word += int32(str[strIdx]-'0') * pow10(innerIdx)
		innerIdx++
		if innerIdx == digitsPerWord {
			wordIdx--
			d.wordBuf[wordIdx] = word
			word = 0
			innerIdx = 0
		}
	}
	if innerIdx != 0 {
		wordIdx--
		d.wordBuf[wordIdx] = word
	}

	wordIdx = wordsInt
	strIdx = strIdxTmp
	word = 0
	innerIdx = 0
	for digitsFrac > 0 {
		digitsFrac--
		strIdx++
		word = int32(str[strIdx]-'0') + word*10
		innerIdx++
		if innerIdx == digitsPerWord {
			d.wordBuf[wordIdx] = word
			wordIdx++
			word = 0
			innerIdx = 0
		}
	}
	if innerIdx != 0 {
		d.wordBuf[wordIdx] = word * pow10(digitsPerWord-innerIdx)
	}
	if endIdx+1 <= len(str) && (str[endIdx] == 'e' || str[endIdx] == 'E') {
		panic(panicInfo)
	}
	allZero := true
	for i := 0; i < wordBufLen; i++ {
		if d.wordBuf[i] != 0 {
			allZero = false
			break
		}
	}
	if allZero {
		d.negative = false
	}
	d.resultFrac = d.digitsFrac
	return err
}
<commit_msg>[parser] remove digit seperator (#683)<commit_after>// Copyright 2019 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

//+build !codes

package test_driver

const panicInfo = "This branch is not implemented. " +
	"This is because you are trying to test something specific to TiDB's MyDecimal implementation. " +
	"It is recommended to do this in TiDB repository."

// constant values.
const (
	maxWordBufLen = 9 // A MyDecimal holds 9 words.
	digitsPerWord = 9 // A word holds 9 digits.
	digMask       = 100000000
)

var (
	wordBufLen = 9
)

// fixWordCntError limits word count in wordBufLen, and returns overflow or truncate error.
func fixWordCntError(wordsInt, wordsFrac int) (newWordsInt int, newWordsFrac int, err error) {
	if wordsInt+wordsFrac > wordBufLen {
		panic(panicInfo)
	}
	return wordsInt, wordsFrac, nil
}

/*
  countLeadingZeroes returns the number of leading zeroes that can be removed from fraction.

  @param   i    start index
  @param   word value to compare against list of powers of 10
*/
func countLeadingZeroes(i int, word int32) int {
	leading := 0
	for word < pow10(i) {
		i--
		leading++
	}
	return leading
}

func digitsToWords(digits int) int {
	return (digits + digitsPerWord - 1) / digitsPerWord
}

// MyDecimal represents a decimal value.
type MyDecimal struct {
	digitsInt int8 // the number of *decimal* digits before the point.

	digitsFrac int8 // the number of decimal digits after the point.

	resultFrac int8 // result fraction digits.

	negative bool

	// wordBuf is an array of int32 words.
	// A word is an int32 value can hold 9 digits.(0 <= word < wordBase)
	wordBuf [maxWordBufLen]int32
}

// String returns the decimal string representation rounded to resultFrac.
func (d *MyDecimal) String() string {
	tmp := *d
	return string(tmp.ToString())
}

func (d *MyDecimal) stringSize() int {
	// sign, zero integer and dot.
	return int(d.digitsInt + d.digitsFrac + 3)
}

func (d *MyDecimal) removeLeadingZeros() (wordIdx int, digitsInt int) {
	digitsInt = int(d.digitsInt)
	i := ((digitsInt - 1) % digitsPerWord) + 1
	for digitsInt > 0 && d.wordBuf[wordIdx] == 0 {
		digitsInt -= i
		i = digitsPerWord
		wordIdx++
	}
	if digitsInt > 0 {
		digitsInt -= countLeadingZeroes((digitsInt-1)%digitsPerWord, d.wordBuf[wordIdx])
	} else {
		digitsInt = 0
	}
	return
}

// ToString converts decimal to its printable string representation without rounding.
//
//  RETURN VALUE
//
//      str       - result string
//      errCode   - eDecOK/eDecTruncate/eDecOverflow
//
func (d *MyDecimal) ToString() (str []byte) {
	str = make([]byte, d.stringSize())
	digitsFrac := int(d.digitsFrac)
	wordStartIdx, digitsInt := d.removeLeadingZeros()
	if digitsInt+digitsFrac == 0 {
		digitsInt = 1
		wordStartIdx = 0
	}

	digitsIntLen := digitsInt
	if digitsIntLen == 0 {
		digitsIntLen = 1
	}
	digitsFracLen := digitsFrac
	length := digitsIntLen + digitsFracLen
	if d.negative {
		length++
	}
	if digitsFrac > 0 {
		length++
	}
	str = str[:length]
	strIdx := 0
	if d.negative {
		str[strIdx] = '-'
		strIdx++
	}
	var fill int
	if digitsFrac > 0 {
		fracIdx := strIdx + digitsIntLen
		fill = digitsFracLen - digitsFrac
		wordIdx := wordStartIdx + digitsToWords(digitsInt)
		str[fracIdx] = '.'
		fracIdx++
		for ; digitsFrac > 0; digitsFrac -= digitsPerWord {
			x := d.wordBuf[wordIdx]
			wordIdx++
			for i := myMin(digitsFrac, digitsPerWord); i > 0; i-- {
				y := x / digMask
				str[fracIdx] = byte(y) + '0'
				fracIdx++
				x -= y * digMask
				x *= 10
			}
		}
		for ; fill > 0; fill-- {
			str[fracIdx] = '0'
			fracIdx++
		}
	}
	fill = digitsIntLen - digitsInt
	if digitsInt == 0 {
		fill-- /* symbol 0 before digital point */
	}
	for ; fill > 0; fill-- {
		str[strIdx] = '0'
		strIdx++
	}
	if digitsInt > 0 {
		strIdx += digitsInt
		wordIdx := wordStartIdx + digitsToWords(digitsInt)
		for ; digitsInt > 0; digitsInt -= digitsPerWord {
			wordIdx--
			x := d.wordBuf[wordIdx]
			for i := myMin(digitsInt, digitsPerWord); i > 0; i-- {
				y := x / 10
				strIdx--
				str[strIdx] = '0' + byte(x-y*10)
				x = y
			}
		}
	} else {
		str[strIdx] = '0'
	}
	return
}

// FromString parses decimal from string.
func (d *MyDecimal) FromString(str []byte) error {
	for i := 0; i < len(str); i++ {
		if !isSpace(str[i]) {
			str = str[i:]
			break
		}
	}
	if len(str) == 0 {
		panic(panicInfo)
	}
	switch str[0] {
	case '-':
		d.negative = true
		fallthrough
	case '+':
		str = str[1:]
	}
	var strIdx int
	for strIdx < len(str) && isDigit(str[strIdx]) {
		strIdx++
	}
	digitsInt := strIdx
	var digitsFrac int
	var endIdx int
	if strIdx < len(str) && str[strIdx] == '.' {
		endIdx = strIdx + 1
		for endIdx < len(str) && isDigit(str[endIdx]) {
			endIdx++
		}
		digitsFrac = endIdx - strIdx - 1
	} else {
		digitsFrac = 0
		endIdx = strIdx
	}
	if digitsInt+digitsFrac == 0 {
		panic(panicInfo)
	}
	wordsInt := digitsToWords(digitsInt)
	wordsFrac := digitsToWords(digitsFrac)
	wordsInt, wordsFrac, err := fixWordCntError(wordsInt, wordsFrac)
	if err != nil {
		panic(panicInfo)
	}
	d.digitsInt = int8(digitsInt)
	d.digitsFrac = int8(digitsFrac)
	wordIdx := wordsInt
	strIdxTmp := strIdx
	var word int32
	var innerIdx int
	for digitsInt > 0 {
		digitsInt--
		strIdx--
		word += int32(str[strIdx]-'0') * pow10(innerIdx)
		innerIdx++
		if innerIdx == digitsPerWord {
			wordIdx--
			d.wordBuf[wordIdx] = word
			word = 0
			innerIdx = 0
		}
	}
	if innerIdx != 0 {
		wordIdx--
		d.wordBuf[wordIdx] = word
	}

	wordIdx = wordsInt
	strIdx = strIdxTmp
	word = 0
	innerIdx = 0
	for digitsFrac > 0 {
		digitsFrac--
		strIdx++
		word = int32(str[strIdx]-'0') + word*10
		innerIdx++
		if innerIdx == digitsPerWord {
			d.wordBuf[wordIdx] = word
			wordIdx++
			word = 0
			innerIdx = 0
		}
	}
	if innerIdx != 0 {
		d.wordBuf[wordIdx] = word * pow10(digitsPerWord-innerIdx)
	}
	if endIdx+1 <= len(str) && (str[endIdx] == 'e' || str[endIdx] == 'E') {
		panic(panicInfo)
	}
	allZero := true
	for i := 0; i < wordBufLen; i++ {
		if d.wordBuf[i] != 0 {
			allZero = false
			break
		}
	}
	if allZero {
		d.negative = false
	}
	d.resultFrac = d.digitsFrac
	return err
}
<|endoftext|>


<commit_before>package xpi // import "go.mozilla.org/autograph/signer/xpi"

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.mozilla.org/autograph/signer"
	"go.mozilla.org/cose"
)

// stringToCOSEAlg returns the cose.Algorithm for a string or nil if
// the algorithm isn't implemented
func stringToCOSEAlg(s string) (v *cose.Algorithm) {
	switch strings.ToUpper(s) {
	case cose.PS256.Name:
		v = cose.PS256
	case cose.ES256.Name:
		v = cose.ES256
	case cose.ES384.Name:
		v = cose.ES384
	case cose.ES512.Name:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// generateIssuerEEKeyPair returns a public and private key pair for
// the provided COSEAlgorithm
func (s *PKCS7Signer) generateCOSEKeyPair(coseAlg *cose.Algorithm) (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	var signer *cose.Signer

	if coseAlg == nil {
		err = fmt.Errorf("Cannot generate private key for nil cose Algorithm")
		return
	} else if coseAlg == cose.PS256 {
		const size = 2048
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	} else {
		signer, err = cose.NewSigner(coseAlg, nil)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate private key")
			return
		}
		eeKey = signer.PrivateKey
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

// isSupportedCOSEAlgValue returns whether the COSE alg value is supported or not
func isSupportedCOSEAlgValue(algValue interface{}) bool {
	return algValue == cose.PS256.Value || algValue == cose.ES256.Value || algValue == cose.ES384.Value || algValue == cose.ES512.Value
}

// isValidCOSESignature checks whether a COSE signature is a valid for XPIs
func isValidCOSESignature(sig cose.Signature) (eeCert *x509.Certificate, resultErr error) {
	if len(sig.Headers.Unprotected) != 0 {
		resultErr = fmt.Errorf("XPI COSE Signature must have an empty Unprotected Header")
		return
	}

	if len(sig.Headers.Protected) != 2 {
		resultErr = fmt.Errorf("XPI COSE Signature must have exactly two Protected Headers")
		return
	}
	algValue, ok := sig.Headers.Protected[1] // 1 is the compressed key for "alg"
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature must have alg in Protected Headers")
		return
	}
	if !isSupportedCOSEAlgValue(algValue) {
		resultErr = fmt.Errorf("XPI COSE Signature must have alg %+v is not supported", algValue)
		return
	}

	kidValue, ok := sig.Headers.Protected[4] // 4 is the compressed key for "kid"
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature must have kid in Protected Headers")
		return
	}
	kidBytes, ok := kidValue.([]byte)
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature kid value is not bytes")
		return
	}

	eeCert, err := x509.ParseCertificate(kidBytes) // eeCert
	if err != nil {
		resultErr = errors.Wrapf(err, "XPI COSE Signature kid must decode to a parseable X509 cert")
		return
	}
	return
}

// isValidCOSEMessage checks whether a COSE SignMessage is a valid for
// XPIs and returns parsed intermediate and end entity certs
func isValidCOSEMessage(msg cose.SignMessage) (intermediateCerts, eeCerts []*x509.Certificate, resultErr error) {
	if msg.Payload != nil {
		resultErr = fmt.Errorf("Expected SignMessage payload to be nil, but got %+v", msg.Payload)
		return
	}
	if len(msg.Headers.Unprotected) != 0 {
		resultErr = fmt.Errorf("Expected SignMessage Unprotected headers to be empty, but got %+v", msg.Headers.Unprotected)
		return
	}

	if len(msg.Headers.Protected) != 1 {
		resultErr = fmt.Errorf("Expected SignMessage Protected headers must contain one value, but got %d", len(msg.Headers.Protected))
		return
	}
	kidValue, ok := msg.Headers.Protected[4] // 4 is the compressed key for "kid"
	if !ok {
		resultErr = fmt.Errorf("Expected SignMessage must have kid in Protected Headers")
		return
	}
	// check that all kid values are bytes and decode into certs
	kidArray, ok := kidValue.([]interface{})
	if !ok {
		resultErr = fmt.Errorf("Expected SignMessage Protected Headers kid value to be an array got %+v with type %T", kidValue, kidValue)
		return
	}
	for i, cert := range kidArray {
		certBytes, ok := cert.([]byte)
		if !ok {
			resultErr = fmt.Errorf("Expected SignMessage Protected Headers kid value %d to be a byte slice got %+v with type %T", i, cert, cert)
			return
		}
		intermediateCert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			resultErr = errors.Wrapf(err, "SignMessage Signature Protected Headers kid value %d does not decode to a parseable X509 cert", i)
			return
		}
		intermediateCerts = append(intermediateCerts, intermediateCert)
	}

	for i, sig := range msg.Signatures {
		eeCert, err := isValidCOSESignature(sig)
		if err != nil {
			resultErr = errors.Wrapf(err, "cose signature %d is invalid", i)
			return
		}
		eeCerts = append(eeCerts, eeCert)
	}

	return
}

// verifyCOSESignatures checks that:
//
// 1) COSE manifest and signature files are present
// 2) the PKCS7 manifest is present
// 3) the COSE and PKCS7 manifests do not include COSE files
// 4) we can decode the COSE signature and it has the right format for an XPI
// 5) the right number of signatures are present and all intermediate and end entity certs parse properly
// TODO: 6) there is a trusted path from the included COSE EE certs to the signer cert using the provided intermediates
//
func verifyCOSESignatures(signedFile signer.SignedFile, signOptions Options) error {
	coseManifestBytes, err := readFileFromZIP(signedFile, "META-INF/cose.manifest")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/cose.manifest from signed zip: %v", err)
	}
	coseMsgBytes, err := readFileFromZIP(signedFile, "META-INF/cose.sig")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/cose.sig from signed zip: %v", err)
	}
	pkcs7ManifestBytes, err := readFileFromZIP(signedFile, "META-INF/manifest.mf")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/manifest.mf from signed zip: %v", err)
	}
	coseManifest := string(coseManifestBytes)
	pkcs7Manifest := string(pkcs7ManifestBytes)

	if !strings.Contains(pkcs7Manifest, "cose") {
		return fmt.Errorf("pkcs7 manifest does not contain cose files: %s", pkcs7Manifest)
	}
	if strings.Contains(coseManifest, "cose") {
		return fmt.Errorf("cose manifest contains cose files: %s", coseManifest)
	}

	xpiSig, err := Unmarshal(string(coseMsgBytes), nil)
	if err != nil {
		return errors.Wrap(err, "error unmarshaling cose.sig")
	}
	coseMsg := *xpiSig.signMessage

	if len(coseMsg.Signatures) != len(signOptions.COSEAlgorithms) {
		return fmt.Errorf("cose.sig contains %d signatures, but expected %d", len(coseMsg.Signatures), len(signOptions.COSEAlgorithms))
	}

	// intermediateCerts, eeCerts, err := isValidCOSEMessage(coseMsg)
	_, _, err = isValidCOSEMessage(coseMsg)
	if err != nil {
		return errors.Wrap(err, "cose.sig is not a valid COSE SignMessage")
	}

	// check that we can verify EE certs with the provided intermediates
	// roots, intermediates := x509.NewCertPool(), x509.NewCertPool()
	// ok = roots.AppendCertsFromPEM([]byte(testcase.Certificate))
	// if !ok {
	// 	return fmt.Errorf("failed to add root cert to pool")
	// }
	// for _, intermediateCert := range intermediateCerts {
	// 	intermediates.AddCert(intermediateCert)
	// }
	// for i, eeCert := range eeCerts {
	// 	opts := x509.VerifyOptions{
	// 		DNSName:       signOptions.ID,
	// 		Roots:         roots,
	// 		Intermediates: intermediates,
	// 	}
	// 	if _, err := eeCert.Verify(opts); err != nil {
	// 		return fmt.Errorf("failed to verify EECert %d %s", i, err)
	// 	}
	// }
	return nil
}

// coseSignature returns a CBOR-marshalled COSE SignMessage
// after generating EE certs and signatures for the COSE algorithms
func coseSignature(cn string, manifest []byte, algs []*cose.Algorithm, s *PKCS7Signer) (coseSig []byte, err error) {
	var (
		coseSigners []cose.Signer
		tmp         = cose.NewSignMessage()
		msg         = &tmp
	)
	msg.Payload = manifest

	// Add list of DER encoded intermediate certificates as message key id
	msg.Headers.Protected["kid"] = [][]byte{s.issuerCert.Raw[:]}

	for _, alg := range algs {
		// create a cert and key
		eeCert, eeKey, err := s.MakeEndEntity(cn, alg)
		if err != nil {
			return nil, err
		}

		// create a COSE.Signer
		signer, err := cose.NewSignerFromKey(alg, eeKey)
		if err != nil {
			return nil, errors.Wrap(err, "xpi: COSE signer creation failed")
		}
		coseSigners = append(coseSigners, *signer)

		// create a COSE Signature holder
		sig := cose.NewSignature()
		sig.Headers.Protected["alg"] = alg.Name
		sig.Headers.Protected["kid"] = eeCert.Raw[:]
		msg.AddSignature(sig)
	}

	// external_aad data must be nil and not byte("")
	err = msg.Sign(rand.Reader, nil, coseSigners)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: COSE signing failed")
	}
	// for addons the signature is detached and the payload is always nil / null
	msg.Payload = nil

	coseSig, err = cose.Marshal(msg)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: error serializing COSE signatures to CBOR")
	}

	return
}
<commit_msg>xpi: rm dead code<commit_after>package xpi // import "go.mozilla.org/autograph/signer/xpi"

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.mozilla.org/autograph/signer"
	"go.mozilla.org/cose"
)

// stringToCOSEAlg returns the cose.Algorithm for a string or nil if
// the algorithm isn't implemented
func stringToCOSEAlg(s string) (v *cose.Algorithm) {
	switch strings.ToUpper(s) {
	case cose.PS256.Name:
		v = cose.PS256
	case cose.ES256.Name:
		v = cose.ES256
	case cose.ES384.Name:
		v = cose.ES384
	case cose.ES512.Name:
		v = cose.ES512
	default:
		v = nil
	}
	return v
}

// generateIssuerEEKeyPair returns a public and private key pair for
// the provided COSEAlgorithm
func (s *PKCS7Signer) generateCOSEKeyPair(coseAlg *cose.Algorithm) (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	var signer *cose.Signer

	if coseAlg == nil {
		err = fmt.Errorf("Cannot generate private key for nil cose Algorithm")
		return
	} else if coseAlg == cose.PS256 {
		const size = 2048
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	} else {
		signer, err = cose.NewSigner(coseAlg, nil)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate private key")
			return
		}
		eeKey = signer.PrivateKey
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

// isSupportedCOSEAlgValue returns whether the COSE alg value is supported or not
func isSupportedCOSEAlgValue(algValue interface{}) bool {
	return algValue == cose.PS256.Value || algValue == cose.ES256.Value || algValue == cose.ES384.Value || algValue == cose.ES512.Value
}

// isValidCOSESignature checks whether a COSE signature is a valid for XPIs
func isValidCOSESignature(sig cose.Signature) (eeCert *x509.Certificate, resultErr error) {
	if len(sig.Headers.Unprotected) != 0 {
		resultErr = fmt.Errorf("XPI COSE Signature must have an empty Unprotected Header")
		return
	}

	if len(sig.Headers.Protected) != 2 {
		resultErr = fmt.Errorf("XPI COSE Signature must have exactly two Protected Headers")
		return
	}
	algValue, ok := sig.Headers.Protected[1] // 1 is the compressed key for "alg"
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature must have alg in Protected Headers")
		return
	}
	if !isSupportedCOSEAlgValue(algValue) {
		resultErr = fmt.Errorf("XPI COSE Signature must have alg %+v is not supported", algValue)
		return
	}

	kidValue, ok := sig.Headers.Protected[4] // 4 is the compressed key for "kid"
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature must have kid in Protected Headers")
		return
	}
	kidBytes, ok := kidValue.([]byte)
	if !ok {
		resultErr = fmt.Errorf("XPI COSE Signature kid value is not bytes")
		return
	}

	eeCert, err := x509.ParseCertificate(kidBytes) // eeCert
	if err != nil {
		resultErr = errors.Wrapf(err, "XPI COSE Signature kid must decode to a parseable X509 cert")
		return
	}
	return
}

// isValidCOSEMessage checks whether a COSE SignMessage is a valid for
// XPIs and returns parsed intermediate and end entity certs
func isValidCOSEMessage(msg cose.SignMessage) (intermediateCerts, eeCerts []*x509.Certificate, resultErr error) {
	if msg.Payload != nil {
		resultErr = fmt.Errorf("Expected SignMessage payload to be nil, but got %+v", msg.Payload)
		return
	}
	if len(msg.Headers.Unprotected) != 0 {
		resultErr = fmt.Errorf("Expected SignMessage Unprotected headers to be empty, but got %+v", msg.Headers.Unprotected)
		return
	}

	if len(msg.Headers.Protected) != 1 {
		resultErr = fmt.Errorf("Expected SignMessage Protected headers must contain one value, but got %d", len(msg.Headers.Protected))
		return
	}
	kidValue, ok := msg.Headers.Protected[4] // 4 is the compressed key for "kid"
	if !ok {
		resultErr = fmt.Errorf("Expected SignMessage must have kid in Protected Headers")
		return
	}
	// check that all kid values are bytes and decode into certs
	kidArray, ok := kidValue.([]interface{})
	if !ok {
		resultErr = fmt.Errorf("Expected SignMessage Protected Headers kid value to be an array got %+v with type %T", kidValue, kidValue)
		return
	}
	for i, cert := range kidArray {
		certBytes, ok := cert.([]byte)
		if !ok {
			resultErr = fmt.Errorf("Expected SignMessage Protected Headers kid value %d to be a byte slice got %+v with type %T", i, cert, cert)
			return
		}
		intermediateCert, err := x509.ParseCertificate(certBytes)
		if err != nil {
			resultErr = errors.Wrapf(err, "SignMessage Signature Protected Headers kid value %d does not decode to a parseable X509 cert", i)
			return
		}
		intermediateCerts = append(intermediateCerts, intermediateCert)
	}

	for i, sig := range msg.Signatures {
		eeCert, err := isValidCOSESignature(sig)
		if err != nil {
			resultErr = errors.Wrapf(err, "cose signature %d is invalid", i)
			return
		}
		eeCerts = append(eeCerts, eeCert)
	}

	return
}

// verifyCOSESignatures checks that:
//
// 1) COSE manifest and signature files are present
// 2) the PKCS7 manifest is present
// 3) the COSE and PKCS7 manifests do not include COSE files
// 4) we can decode the COSE signature and it has the right format for an XPI
// 5) the right number of signatures are present and all intermediate and end entity certs parse properly
// TODO: 6) there is a trusted path from the included COSE EE certs to the signer cert using the provided intermediates
//
func verifyCOSESignatures(signedFile signer.SignedFile, signOptions Options) error {
	coseManifestBytes, err := readFileFromZIP(signedFile, "META-INF/cose.manifest")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/cose.manifest from signed zip: %v", err)
	}
	coseMsgBytes, err := readFileFromZIP(signedFile, "META-INF/cose.sig")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/cose.sig from signed zip: %v", err)
	}
	pkcs7ManifestBytes, err := readFileFromZIP(signedFile, "META-INF/manifest.mf")
	if err != nil {
		return fmt.Errorf("failed to read META-INF/manifest.mf from signed zip: %v", err)
	}
	coseManifest := string(coseManifestBytes)
	pkcs7Manifest := string(pkcs7ManifestBytes)

	if !strings.Contains(pkcs7Manifest, "cose") {
		return fmt.Errorf("pkcs7 manifest does not contain cose files: %s", pkcs7Manifest)
	}
	if strings.Contains(coseManifest, "cose") {
		return fmt.Errorf("cose manifest contains cose files: %s", coseManifest)
	}

	xpiSig, err := Unmarshal(string(coseMsgBytes), nil)
	if err != nil {
		return errors.Wrap(err, "error unmarshaling cose.sig")
	}
	coseMsg := *xpiSig.signMessage

	if len(coseMsg.Signatures) != len(signOptions.COSEAlgorithms) {
		return fmt.Errorf("cose.sig contains %d signatures, but expected %d", len(coseMsg.Signatures), len(signOptions.COSEAlgorithms))
	}

	_, _, err = isValidCOSEMessage(coseMsg)
	if err != nil {
		return errors.Wrap(err, "cose.sig is not a valid COSE SignMessage")
	}

	return nil
}

// coseSignature returns a CBOR-marshalled COSE SignMessage
// after generating EE certs and signatures for the COSE algorithms
func coseSignature(cn string, manifest []byte, algs []*cose.Algorithm, s *PKCS7Signer) (coseSig []byte, err error) {
	var (
		coseSigners []cose.Signer
		tmp         = cose.NewSignMessage()
		msg         = &tmp
	)
	msg.Payload = manifest

	// Add list of DER encoded intermediate certificates as message key id
	msg.Headers.Protected["kid"] = [][]byte{s.issuerCert.Raw[:]}

	for _, alg := range algs {
		// create a cert and key
		eeCert, eeKey, err := s.MakeEndEntity(cn, alg)
		if err != nil {
			return nil, err
		}

		// create a COSE.Signer
		signer, err := cose.NewSignerFromKey(alg, eeKey)
		if err != nil {
			return nil, errors.Wrap(err, "xpi: COSE signer creation failed")
		}
		coseSigners = append(coseSigners, *signer)

		// create a COSE Signature holder
		sig := cose.NewSignature()
		sig.Headers.Protected["alg"] = alg.Name
		sig.Headers.Protected["kid"] = eeCert.Raw[:]
		msg.AddSignature(sig)
	}

	// external_aad data must be nil and not byte("")
	err = msg.Sign(rand.Reader, nil, coseSigners)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: COSE signing failed")
	}
	// for addons the signature is detached and the payload is always nil / null
	msg.Payload = nil

	coseSig, err = cose.Marshal(msg)
	if err != nil {
		return nil, errors.Wrap(err, "xpi: error serializing COSE signatures to CBOR")
	}

	return
}
<|endoftext|>


<commit_before>package cb

import (
	"container/heap"

	"github.com/catorpilor/LeetCode/utils"
)

func assignBikes(workers, bikes [][]int) []int {
	return pq(workers, bikes)
}

func pq(workers, bikes [][]int) []int {
	nw, nb := len(workers), len(bikes)
	if nw == 0 || nb == 0 {
		return []int{}
	}
	tupq := utils.TuplePriorityQueue{}
	// loop through every pairs of bikes and workers
	for i := 0; i < nw; i++ {
		w := workers[i]
		for j := 0; j < nb; j++ {
			b := bikes[j]
			dist := utils.Abs(w[0]-b[0]) + utils.Abs(w[1]-b[1])
			heap.Push(&tupq, [3]int{dist, i, j})
		}
	}
	res := make([]int, nb)
	// init res to -1 labled as none visited
	for i := range res {
		res[i] = -1
	}

	set := make(map[int]bool, nb)
	for len(set) < nb {
		tup := heap.Pop(&tupq).([3]int)
		if res[tup[1]] == -1 && !set[tup[2]] {
			res[tup[1]] = tup[2]
			set[tup[2]] = true
		}
	}
	return res
}

// bucket use bucket sort to solve the problem
func bucket(workers, bikes [][]int) []int {
	nw, nb := len(workers), len(bikes)
	if nw == 0 || nb == 0 {
		return []int{}
	}
	// since distance ranges from [0,2000]
	// we create a bucket to store the (worker, bike) pairs
	// for bucket[i] stores the distance=i's pairs
	type pair struct {
		w, b int
	}
	bkt := [2001][]pair{}
	for i := 0; i < nw; i++ {
		for j := 0; j < nb; j++ {
			dis := utils.Abs(workers[i][0]-bikes[j][0]) + utils.Abs(workers[i][1]-bikes[j][1])
			bkt[dis] = append(bkt[dis], pair{w: i, b: j})
		}
	}
	res := make([]int, nb)
	for i := range res {
		res[i] = -1
	}
	set := make(map[int]bool, nb)
	for d := 0; d <= 2000; d++ {
		for k := 0; k < len(bkt[d]); k++ {
			if res[bkt[d][k].w] == -1 && !set[bkt[d][k].b] {
				set[bkt[d][k].b] = true
				res[bkt[d][k].w] = bkt[d][k].b
			}
		}
	}
	return res
}
<commit_msg>fix: typo<commit_after>package cb

import (
	"container/heap"

	"github.com/catorpilor/LeetCode/utils"
)

func assignBikes(workers, bikes [][]int) []int {
	return pq(workers, bikes)
}

func pq(workers, bikes [][]int) []int {
	nw, nb := len(workers), len(bikes)
	if nw == 0 || nb == 0 {
		return []int{}
	}
	tupq := utils.TuplePriorityQueue{}
	// loop through every pairs of bikes and workers
	for i := 0; i < nw; i++ {
		w := workers[i]
		for j := 0; j < nb; j++ {
			b := bikes[j]
			dist := utils.Abs(w[0]-b[0]) + utils.Abs(w[1]-b[1])
			heap.Push(&tupq, [3]int{dist, i, j})
		}
	}
	res := make([]int, nb)
	// init res to -1 labled as none visited
	for i := range res {
		res[i] = -1
	}

	set := make(map[int]bool, nb)
	for len(set) < nb {
		tup := heap.Pop(&tupq).([3]int)
		if res[tup[1]] == -1 && !set[tup[2]] {
			res[tup[1]] = tup[2]
			set[tup[2]] = true
		}
	}
	return res
}

// bucket use bucket sort to solve the problem
func bucket(workers, bikes [][]int) []int {
	nw, nb := len(workers), len(bikes)
	if nw == 0 || nb == 0 {
		return []int{}
	}
	// since distance ranges from [0,2000]
	// we create a bucket to store the (worker, bike) pairs
	// for bucket[i] stores the distance=i's pairs
	type pair struct {
		w, b int
	}
	bkt := [2001][]pair{}
	for i := 0; i < nw; i++ {
		for j := 0; j < nb; j++ {
			dis := utils.Abs(workers[i][0]-bikes[j][0]) + utils.Abs(workers[i][1]-bikes[j][1])
			bkt[dis] = append(bkt[dis], pair{w: i, b: j})
		}
	}
	res := make([]int, nw)
	for i := range res {
		res[i] = -1
	}
	set := make(map[int]bool, nb)
	for d := 0; d <= 2000; d++ {
		for k := 0; k < len(bkt[d]); k++ {
			if res[bkt[d][k].w] == -1 && !set[bkt[d][k].b] {
				set[bkt[d][k].b] = true
				res[bkt[d][k].w] = bkt[d][k].b
			}
		}
	}
	return res
}
<|endoftext|>