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

package source

import (
	"fmt"
	"os"

	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/telemetry/tag"
	errors "golang.org/x/xerrors"
)

var (
	DefaultOptions = Options{
		Env:                    os.Environ(),
		TextDocumentSyncKind:   protocol.Incremental,
		HoverKind:              SynopsisDocumentation,
		InsertTextFormat:       protocol.PlainTextTextFormat,
		PreferredContentFormat: protocol.PlainText,
		SupportedCodeActions: map[FileKind]map[protocol.CodeActionKind]bool{
			Go: {
				protocol.SourceOrganizeImports: true,
				protocol.QuickFix:              true,
			},
			Mod: {},
			Sum: {},
		},
		Completion: CompletionOptions{
			Documentation: true,
			Deep:          false,
			FuzzyMatching: false,
		},
	}
)

type Options struct {

	// Env is the current set of environment overrides on this view.
	Env []string

	// BuildFlags is used to adjust the build flags applied to the view.
	BuildFlags []string

	HoverKind        HoverKind
	DisabledAnalyses map[string]struct{}

	WatchFileChanges              bool
	InsertTextFormat              protocol.InsertTextFormat
	ConfigurationSupported        bool
	DynamicConfigurationSupported bool
	DynamicWatchedFilesSupported  bool
	PreferredContentFormat        protocol.MarkupKind
	LineFoldingOnly               bool

	SupportedCodeActions map[FileKind]map[protocol.CodeActionKind]bool

	// TODO: Remove the option once we are certain there are no issues here.
	TextDocumentSyncKind protocol.TextDocumentSyncKind

	Completion CompletionOptions
}

type CompletionOptions struct {
	Deep              bool
	FuzzyMatching     bool
	Unimported        bool
	Documentation     bool
	FullDocumentation bool
	Placeholders      bool
}

type HoverKind int

const (
	SingleLine = HoverKind(iota)
	NoDocumentation
	SynopsisDocumentation
	FullDocumentation

	// structured is an experimental setting that returns a structured hover format.
	// This format separates the signature from the documentation, so that the client
	// can do more manipulation of these fields.
	//
	// This should only be used by clients that support this behavior.
	Structured
)

type OptionResults []OptionResult

type OptionResult struct {
	Name  string
	Value interface{}
	Error error

	State       OptionState
	Replacement string
}

type OptionState int

const (
	OptionHandled = OptionState(iota)
	OptionDeprecated
	OptionUnexpected
)

func SetOptions(options *Options, opts interface{}) OptionResults {
	var results OptionResults
	switch opts := opts.(type) {
	case nil:
	case map[string]interface{}:
		for name, value := range opts {
			results = append(results, options.set(name, value))
		}
	default:
		results = append(results, OptionResult{
			Value: opts,
			Error: errors.Errorf("Invalid options type %T", opts),
		})
	}
	return results
}

func (o *Options) ForClientCapabilities(caps protocol.ClientCapabilities) {
	// Check if the client supports snippets in completion items.
	if caps.TextDocument.Completion.CompletionItem != nil &&
		caps.TextDocument.Completion.CompletionItem.SnippetSupport {
		o.InsertTextFormat = protocol.SnippetTextFormat
	}
	// Check if the client supports configuration messages.
	o.ConfigurationSupported = caps.Workspace.Configuration
	o.DynamicConfigurationSupported = caps.Workspace.DidChangeConfiguration.DynamicRegistration
	o.DynamicWatchedFilesSupported = caps.Workspace.DidChangeWatchedFiles.DynamicRegistration

	// Check which types of content format are supported by this client.
	if len(caps.TextDocument.Hover.ContentFormat) > 0 {
		o.PreferredContentFormat = caps.TextDocument.Hover.ContentFormat[0]
	}
	// Check if the client supports only line folding.
	o.LineFoldingOnly = caps.TextDocument.FoldingRange.LineFoldingOnly
}

func (o *Options) set(name string, value interface{}) OptionResult {
	result := OptionResult{Name: name, Value: value}
	switch name {
	case "env":
		menv, ok := value.(map[string]interface{})
		if !ok {
			result.errorf("invalid config gopls.env type %T", value)
			break
		}
		for k, v := range menv {
			o.Env = append(o.Env, fmt.Sprintf("%s=%s", k, v))
		}

	case "buildFlags":
		iflags, ok := value.([]interface{})
		if !ok {
			result.errorf("invalid config gopls.buildFlags type %T", value)
			break
		}
		flags := make([]string, 0, len(iflags))
		for _, flag := range iflags {
			flags = append(flags, fmt.Sprintf("%s", flag))
		}
		o.BuildFlags = flags

	case "noIncrementalSync":
		if v, ok := result.asBool(); ok && v {
			o.TextDocumentSyncKind = protocol.Full
		}
	case "watchFileChanges":
		result.setBool(&o.WatchFileChanges)
	case "completionDocumentation":
		result.setBool(&o.Completion.Documentation)
	case "usePlaceholders":
		result.setBool(&o.Completion.Placeholders)
	case "deepCompletion":
		result.setBool(&o.Completion.Deep)
	case "fuzzyMatching":
		result.setBool(&o.Completion.FuzzyMatching)
	case "completeUnimported":
		result.setBool(&o.Completion.Unimported)

	case "hoverKind":
		hoverKind, ok := value.(string)
		if !ok {
			result.errorf("Invalid type %T for string option %q", value, name)
			break
		}
		switch hoverKind {
		case "NoDocumentation":
			o.HoverKind = NoDocumentation
		case "SingleLine":
			o.HoverKind = SingleLine
		case "SynopsisDocumentation":
			o.HoverKind = SynopsisDocumentation
		case "FullDocumentation":
			o.HoverKind = FullDocumentation
		case "Structured":
			o.HoverKind = Structured
		default:
			result.errorf("Unsupported hover kind", tag.Of("HoverKind", hoverKind))
		}

	case "experimentalDisabledAnalyses":
		disabledAnalyses, ok := value.([]interface{})
		if !ok {
			result.errorf("Invalid type %T for []string option %q", value, name)
			break
		}
		o.DisabledAnalyses = make(map[string]struct{})
		for _, a := range disabledAnalyses {
			o.DisabledAnalyses[fmt.Sprint(a)] = struct{}{}
		}

	// Deprecated settings.
	case "wantSuggestedFixes":
		result.State = OptionDeprecated

	case "disableDeepCompletion":
		result.State = OptionDeprecated
		result.Replacement = "deepCompletion"

	case "disableFuzzyMatching":
		result.State = OptionDeprecated
		result.Replacement = "fuzzyMatching"

	case "wantCompletionDocumentation":
		result.State = OptionDeprecated
		result.Replacement = "completionDocumentation"

	case "wantUnimportedCompletions":
		result.State = OptionDeprecated
		result.Replacement = "completeUnimported"

	default:
		result.State = OptionUnexpected
	}
	return result
}

func (r *OptionResult) errorf(msg string, values ...interface{}) {
	r.Error = errors.Errorf(msg, values...)
}

func (r *OptionResult) asBool() (bool, bool) {
	b, ok := r.Value.(bool)
	if !ok {
		r.errorf("Invalid type %T for bool option %q", r.Value, r.Name)
		return false, false
	}
	return b, true
}

func (r *OptionResult) setBool(b *bool) {
	if v, ok := r.asBool(); ok {
		*b = v
	}
}
<commit_msg>internal/lsp: turn fuzzy matching and deep completions back on<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package source

import (
	"fmt"
	"os"

	"golang.org/x/tools/internal/lsp/protocol"
	"golang.org/x/tools/internal/telemetry/tag"
	errors "golang.org/x/xerrors"
)

var (
	DefaultOptions = Options{
		Env:                    os.Environ(),
		TextDocumentSyncKind:   protocol.Incremental,
		HoverKind:              SynopsisDocumentation,
		InsertTextFormat:       protocol.PlainTextTextFormat,
		PreferredContentFormat: protocol.PlainText,
		SupportedCodeActions: map[FileKind]map[protocol.CodeActionKind]bool{
			Go: {
				protocol.SourceOrganizeImports: true,
				protocol.QuickFix:              true,
			},
			Mod: {},
			Sum: {},
		},
		Completion: CompletionOptions{
			Documentation: true,
			Deep:          true,
			FuzzyMatching: true,
		},
	}
)

type Options struct {

	// Env is the current set of environment overrides on this view.
	Env []string

	// BuildFlags is used to adjust the build flags applied to the view.
	BuildFlags []string

	HoverKind        HoverKind
	DisabledAnalyses map[string]struct{}

	WatchFileChanges              bool
	InsertTextFormat              protocol.InsertTextFormat
	ConfigurationSupported        bool
	DynamicConfigurationSupported bool
	DynamicWatchedFilesSupported  bool
	PreferredContentFormat        protocol.MarkupKind
	LineFoldingOnly               bool

	SupportedCodeActions map[FileKind]map[protocol.CodeActionKind]bool

	// TODO: Remove the option once we are certain there are no issues here.
	TextDocumentSyncKind protocol.TextDocumentSyncKind

	Completion CompletionOptions
}

type CompletionOptions struct {
	Deep              bool
	FuzzyMatching     bool
	Unimported        bool
	Documentation     bool
	FullDocumentation bool
	Placeholders      bool
}

type HoverKind int

const (
	SingleLine = HoverKind(iota)
	NoDocumentation
	SynopsisDocumentation
	FullDocumentation

	// structured is an experimental setting that returns a structured hover format.
	// This format separates the signature from the documentation, so that the client
	// can do more manipulation of these fields.
	//
	// This should only be used by clients that support this behavior.
	Structured
)

type OptionResults []OptionResult

type OptionResult struct {
	Name  string
	Value interface{}
	Error error

	State       OptionState
	Replacement string
}

type OptionState int

const (
	OptionHandled = OptionState(iota)
	OptionDeprecated
	OptionUnexpected
)

func SetOptions(options *Options, opts interface{}) OptionResults {
	var results OptionResults
	switch opts := opts.(type) {
	case nil:
	case map[string]interface{}:
		for name, value := range opts {
			results = append(results, options.set(name, value))
		}
	default:
		results = append(results, OptionResult{
			Value: opts,
			Error: errors.Errorf("Invalid options type %T", opts),
		})
	}
	return results
}

func (o *Options) ForClientCapabilities(caps protocol.ClientCapabilities) {
	// Check if the client supports snippets in completion items.
	if caps.TextDocument.Completion.CompletionItem != nil &&
		caps.TextDocument.Completion.CompletionItem.SnippetSupport {
		o.InsertTextFormat = protocol.SnippetTextFormat
	}
	// Check if the client supports configuration messages.
	o.ConfigurationSupported = caps.Workspace.Configuration
	o.DynamicConfigurationSupported = caps.Workspace.DidChangeConfiguration.DynamicRegistration
	o.DynamicWatchedFilesSupported = caps.Workspace.DidChangeWatchedFiles.DynamicRegistration

	// Check which types of content format are supported by this client.
	if len(caps.TextDocument.Hover.ContentFormat) > 0 {
		o.PreferredContentFormat = caps.TextDocument.Hover.ContentFormat[0]
	}
	// Check if the client supports only line folding.
	o.LineFoldingOnly = caps.TextDocument.FoldingRange.LineFoldingOnly
}

func (o *Options) set(name string, value interface{}) OptionResult {
	result := OptionResult{Name: name, Value: value}
	switch name {
	case "env":
		menv, ok := value.(map[string]interface{})
		if !ok {
			result.errorf("invalid config gopls.env type %T", value)
			break
		}
		for k, v := range menv {
			o.Env = append(o.Env, fmt.Sprintf("%s=%s", k, v))
		}

	case "buildFlags":
		iflags, ok := value.([]interface{})
		if !ok {
			result.errorf("invalid config gopls.buildFlags type %T", value)
			break
		}
		flags := make([]string, 0, len(iflags))
		for _, flag := range iflags {
			flags = append(flags, fmt.Sprintf("%s", flag))
		}
		o.BuildFlags = flags

	case "noIncrementalSync":
		if v, ok := result.asBool(); ok && v {
			o.TextDocumentSyncKind = protocol.Full
		}
	case "watchFileChanges":
		result.setBool(&o.WatchFileChanges)
	case "completionDocumentation":
		result.setBool(&o.Completion.Documentation)
	case "usePlaceholders":
		result.setBool(&o.Completion.Placeholders)
	case "deepCompletion":
		result.setBool(&o.Completion.Deep)
	case "fuzzyMatching":
		result.setBool(&o.Completion.FuzzyMatching)
	case "completeUnimported":
		result.setBool(&o.Completion.Unimported)

	case "hoverKind":
		hoverKind, ok := value.(string)
		if !ok {
			result.errorf("Invalid type %T for string option %q", value, name)
			break
		}
		switch hoverKind {
		case "NoDocumentation":
			o.HoverKind = NoDocumentation
		case "SingleLine":
			o.HoverKind = SingleLine
		case "SynopsisDocumentation":
			o.HoverKind = SynopsisDocumentation
		case "FullDocumentation":
			o.HoverKind = FullDocumentation
		case "Structured":
			o.HoverKind = Structured
		default:
			result.errorf("Unsupported hover kind", tag.Of("HoverKind", hoverKind))
		}

	case "experimentalDisabledAnalyses":
		disabledAnalyses, ok := value.([]interface{})
		if !ok {
			result.errorf("Invalid type %T for []string option %q", value, name)
			break
		}
		o.DisabledAnalyses = make(map[string]struct{})
		for _, a := range disabledAnalyses {
			o.DisabledAnalyses[fmt.Sprint(a)] = struct{}{}
		}

	// Deprecated settings.
	case "wantSuggestedFixes":
		result.State = OptionDeprecated

	case "disableDeepCompletion":
		result.State = OptionDeprecated
		result.Replacement = "deepCompletion"

	case "disableFuzzyMatching":
		result.State = OptionDeprecated
		result.Replacement = "fuzzyMatching"

	case "wantCompletionDocumentation":
		result.State = OptionDeprecated
		result.Replacement = "completionDocumentation"

	case "wantUnimportedCompletions":
		result.State = OptionDeprecated
		result.Replacement = "completeUnimported"

	default:
		result.State = OptionUnexpected
	}
	return result
}

func (r *OptionResult) errorf(msg string, values ...interface{}) {
	r.Error = errors.Errorf(msg, values...)
}

func (r *OptionResult) asBool() (bool, bool) {
	b, ok := r.Value.(bool)
	if !ok {
		r.errorf("Invalid type %T for bool option %q", r.Value, r.Name)
		return false, false
	}
	return b, true
}

func (r *OptionResult) setBool(b *bool) {
	if v, ok := r.asBool(); ok {
		*b = v
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016-2019 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/cespare/xxhash/v2"
	"github.com/shenwei356/bio/seq"
	"github.com/shenwei356/bio/seqio/fastx"
	"github.com/shenwei356/bio/sketches"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
	"github.com/twotwotwo/sorts/sortutil"
)

// sumCmd represents the fq2fa command
var sumCmd = &cobra.Command{
	Use:   "sum",
	Short: "compute message digest for all sequences in FASTA/Q files",
	Long: `compute message digest for all sequences in FASTA/Q files

Attentions:
  1. Sequence headers and qualities are skipped, only sequences matter.
  2. The order of sequences records does not matter.
  3. Circular complete genomes are supported with the flag -c/--circular.
     - The same genomes with different start positions or in reverse
       complement strand will not affect the result.
     - The message digest would change with different values of k-mer size.
  4. Multiple files are processed in parallel (-j/--threads).

Method:
  1. Converting the sequences to low cases, optionally removing gaps (-g).
  2. Computing the hash (xxhash) for all sequences or k-mers of a circular
     complete genome (-c/--circular).
  3. Sorting all hash values, for ignoring the order of sequences.
  4. Computing MD5 digest from the hash values, sequences length, and
     the number of sequences.

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		alphabet := config.Alphabet
		idRegexp := config.IDRegexp
		// lineWidth := config.LineWidth
		outFile := config.OutFile
		seq.AlphabetGuessSeqLengthThreshold = config.AlphabetGuessSeqLength
		seq.ValidateSeq = false
		runtime.GOMAXPROCS(config.Threads)

		basename := getFlagBool(cmd, "basename")
		circular := getFlagBool(cmd, "circular")
		k := getFlagPositiveInt(cmd, "kmer-size")
		removeGaps := getFlagBool(cmd, "remove-gaps")
		gapLetters := getFlagString(cmd, "gap-letters")

		files := getFileListFromArgsAndFile(cmd, args, true, "infile-list", true)

		outfh, err := xopen.Wopen(outFile)
		checkError(err)
		defer outfh.Close()

		tokens := make(chan int, config.Threads)
		done := make(chan int)
		var wg sync.WaitGroup

		type Aresult struct {
			ok     bool
			id     uint64
			result *SumResult
		}
		ch := make(chan *Aresult, config.Threads)

		go func() {
			m := make(map[uint64]*Aresult, config.Threads)
			var id, _id uint64
			var ok bool
			var _r *Aresult

			id = 1
			for r := range ch {
				_id = r.id

				if _id == id { // right there
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", r.result.Digest, r.result.File)
						outfh.Flush()
					}
					id++
					continue
				}

				m[_id] = r // save for later check

				if _r, ok = m[id]; ok { // check buffered
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						outfh.Flush()
					}
					delete(m, id)
					id++
				}
			}

			if len(m) > 0 {
				ids := make([]uint64, len(m))
				i := 0
				for _id = range m {
					ids[i] = _id
					i++
				}
				sortutil.Uint64s(ids)
				for _, _id = range ids {
					_r = m[_id]

					if _r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						outfh.Flush()
					}
				}
			}
			done <- 1
		}()

		var id uint64
		for _, file := range files {
			tokens <- 1
			wg.Add(1)
			id++

			go func(file string, id uint64) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				var n int // number of sequences
				var l int // lengths of all seqs
				var h uint64
				hashes := make([]uint64, 0, 1024)

				var record *fastx.Record
				var fastxReader *fastx.Reader

				fastxReader, err = fastx.NewReader(alphabet, file, idRegexp)
				// checkError(err)
				if err != nil {
					ch <- &Aresult{
						id:     id,
						ok:     false,
						result: nil,
					}
					log.Warningf(fmt.Sprintf("%s: %s", file, err))
					return
				}

				if circular {
					var ok bool
					var iter *sketches.Iterator
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						if k > len(record.Seq.Seq) {
							checkError(fmt.Errorf("k is too big for sequence of %d bp: %s", len(record.Seq.Seq), file))
						}

						if n >= 1 {
							checkError(fmt.Errorf("only one sequence is allowed for circular genome"))
						}

						if removeGaps {
							record.Seq.RemoveGapsInplace(gapLetters)
						}

						record.Seq.Seq = bytes.ToLower(record.Seq.Seq)
						iter, err = sketches.NewHashIterator(record.Seq, k, true, true)

						for {
							h, ok = iter.NextHash()
							if !ok {
								break
							}

							hashes = append(hashes, h)
						}

						n++
						l += len(record.Seq.Seq)
					}
				} else {
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						if removeGaps {
							record.Seq.RemoveGapsInplace(gapLetters)
						}

						h = xxhash.Sum64(bytes.ToLower(record.Seq.Seq))
						hashes = append(hashes, h)

						n++
						l += len(record.Seq.Seq)
					}
				}

				// sequences
				sortutil.Uint64s(hashes)

				var le = binary.LittleEndian
				buf := make([]byte, 8)
				di := xxhash.New()

				for _, h = range hashes {
					le.PutUint64(buf, h)
					di.Write(buf)
				}

				// sequence length
				le.PutUint64(buf, uint64(l))
				di.Write(buf)

				// sequence number
				le.PutUint64(buf, uint64(n))
				di.Write(buf)

				// sum up
				le.PutUint64(buf, di.Sum64())
				digest := md5.Sum(buf)

				// return result
				if basename {
					file = filepath.Base(file)
				}
				ch <- &Aresult{
					id: id,
					ok: true,
					result: &SumResult{
						File:   file,
						SeqNum: n,
						SeqLen: l,
						Digest: hex.EncodeToString(digest[:]),
					},
				}
			}(file, id)
		}

		wg.Wait()
		close(ch)
		<-done
	},
}

func init() {
	RootCmd.AddCommand(sumCmd)

	sumCmd.Flags().BoolP("circular", "c", false, "the file contains a single cicular genome sequence")
	sumCmd.Flags().IntP("kmer-size", "k", 1000, "k-mer size for processing circular genomes")
	sumCmd.Flags().BoolP("basename", "b", false, "only output basename of files")
	sumCmd.Flags().BoolP("remove-gaps", "g", false, "remove gaps")
	sumCmd.Flags().StringP("gap-letters", "G", "- 	.", "gap letters")
}

type SumResult struct {
	File   string
	SeqNum int
	SeqLen int
	Digest string
}
<commit_msg>seqkit sum: fix bug for circular genomes, but become much slower<commit_after>// Copyright © 2016-2019 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"path/filepath"
	"runtime"
	"sync"

	"github.com/cespare/xxhash/v2"
	"github.com/shenwei356/bio/seq"
	"github.com/shenwei356/bio/seqio/fastx"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
	"github.com/twotwotwo/sorts/sortutil"
)

// sumCmd represents the fq2fa command
var sumCmd = &cobra.Command{
	Use:   "sum",
	Short: "compute message digest for all sequences in FASTA/Q files",
	Long: `compute message digest for all sequences in FASTA/Q files

Attentions:
  1. Sequence headers and qualities are skipped, only sequences matter.
  2. The order of sequences records does not matter.
  3. Circular complete genomes are supported with the flag -c/--circular.
     - The same genomes with different start positions or in reverse
       complement strand will not affect the result.
     - The message digest would change with different values of k-mer size.
  4. Multiple files are processed in parallel (-j/--threads).

Method:
  1. Converting the sequences to low cases, optionally removing gaps (-g).
  2. Computing the hash (xxhash) for all sequences or k-mers of a circular
     complete genome (-c/--circular).
  3. Sorting all hash values, for ignoring the order of sequences.
  4. Computing MD5 digest from the hash values, sequences length, and
     the number of sequences.

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)
		alphabet := config.Alphabet
		idRegexp := config.IDRegexp
		// lineWidth := config.LineWidth
		outFile := config.OutFile
		seq.AlphabetGuessSeqLengthThreshold = config.AlphabetGuessSeqLength
		seq.ValidateSeq = false
		runtime.GOMAXPROCS(config.Threads)

		basename := getFlagBool(cmd, "basename")
		circular := getFlagBool(cmd, "circular")
		k := getFlagPositiveInt(cmd, "kmer-size")
		removeGaps := getFlagBool(cmd, "remove-gaps")
		gapLetters := getFlagString(cmd, "gap-letters")

		files := getFileListFromArgsAndFile(cmd, args, true, "infile-list", true)

		outfh, err := xopen.Wopen(outFile)
		checkError(err)
		defer outfh.Close()

		tokens := make(chan int, config.Threads)
		done := make(chan int)
		var wg sync.WaitGroup

		type Aresult struct {
			ok     bool
			id     uint64
			result *SumResult
		}
		ch := make(chan *Aresult, config.Threads)

		go func() {
			m := make(map[uint64]*Aresult, config.Threads)
			var id, _id uint64
			var ok bool
			var _r *Aresult

			id = 1
			for r := range ch {
				_id = r.id

				if _id == id { // right there
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", r.result.Digest, r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", r.result.Digest, r.result.SeqLen, r.result.SeqNum, r.result.File)
						outfh.Flush()
					}
					id++
					continue
				}

				m[_id] = r // save for later check

				if _r, ok = m[id]; ok { // check buffered
					if r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", _r.result.Digest, _r.result.SeqLen, _r.result.SeqNum, _r.result.File)
						outfh.Flush()
					}
					delete(m, id)
					id++
				}
			}

			if len(m) > 0 {
				ids := make([]uint64, len(m))
				i := 0
				for _id = range m {
					ids[i] = _id
					i++
				}
				sortutil.Uint64s(ids)
				for _, _id = range ids {
					_r = m[_id]

					if _r.ok {
						fmt.Fprintf(outfh, "%s\t%s\n", _r.result.Digest, _r.result.File)
						// fmt.Fprintf(outfh, "%s\t%d\t%d\t%s\n", _r.result.Digest, _r.result.SeqLen, _r.result.SeqNum, _r.result.File)
						outfh.Flush()
					}
				}
			}
			done <- 1
		}()

		var id uint64
		for _, file := range files {
			tokens <- 1
			wg.Add(1)
			id++

			go func(file string, id uint64) {
				defer func() {
					<-tokens
					wg.Done()
				}()

				var n int // number of sequences
				var l int // lengths of all seqs
				var h uint64
				hashes := make([]uint64, 0, 1024)

				var record *fastx.Record
				var fastxReader *fastx.Reader

				fastxReader, err = fastx.NewReader(alphabet, file, idRegexp)
				// checkError(err)
				if err != nil {
					ch <- &Aresult{
						id:     id,
						ok:     false,
						result: nil,
					}
					log.Warningf(fmt.Sprintf("%s: %s", file, err))
					return
				}

				if circular {
					var ok bool
					// var iter *sketches.Iterator
					var slider func() (*seq.Seq, bool)
					var s *seq.Seq
					var h2 uint64

					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						if k > len(record.Seq.Seq) {
							checkError(fmt.Errorf("k is too big for sequence of %d bp: %s", len(record.Seq.Seq), file))
						}

						if n >= 1 {
							checkError(fmt.Errorf("only one sequence is allowed for circular genome"))
						}

						if removeGaps {
							record.Seq.RemoveGapsInplace(gapLetters)
						}

						record.Seq.Seq = bytes.ToLower(record.Seq.Seq)

						// ntHash shoud be simpler and faster, but some thing wrong I didn't fingure out.
						// iter, err = sketches.NewHashIterator(record.Seq, k, true, true)
						// for {
						// 	h, ok = iter.NextHash()
						// 	if !ok {
						// 		break
						// 	}

						// 	hashes = append(hashes, h)
						// }

						slider = record.Seq.Slider(k, 1, true, true)

						for {
							s, ok = slider()
							if !ok {
								break
							}

							h = xxhash.Sum64(s.Seq)
							h2 = xxhash.Sum64(s.Clone().RevComInplace().Seq)
							if h < h2 {
								hashes = append(hashes, h)
							} else {
								hashes = append(hashes, h2)
							}
						}

						n++
						l += len(record.Seq.Seq)
					}
				} else {
					for {
						record, err = fastxReader.Read()
						if err != nil {
							if err == io.EOF {
								break
							}
							// checkError(err)
							// break
							ch <- &Aresult{
								id:     id,
								ok:     false,
								result: nil,
							}
							log.Warningf(fmt.Sprintf("%s: %s", file, err))
							return
						}

						if removeGaps {
							record.Seq.RemoveGapsInplace(gapLetters)
						}

						h = xxhash.Sum64(bytes.ToLower(record.Seq.Seq))
						hashes = append(hashes, h)

						n++
						l += len(record.Seq.Seq)
					}
				}

				// sequences
				sortutil.Uint64s(hashes)

				var le = binary.LittleEndian
				buf := make([]byte, 8)
				di := xxhash.New()

				for _, h = range hashes {
					le.PutUint64(buf, h)
					di.Write(buf)
				}

				// sequence length
				le.PutUint64(buf, uint64(l))
				di.Write(buf)

				// sequence number
				le.PutUint64(buf, uint64(n))
				di.Write(buf)

				// sum up
				le.PutUint64(buf, di.Sum64())
				digest := md5.Sum(buf)

				// return result
				if basename {
					file = filepath.Base(file)
				}
				ch <- &Aresult{
					id: id,
					ok: true,
					result: &SumResult{
						File:   file,
						SeqNum: n,
						SeqLen: l,
						Digest: hex.EncodeToString(digest[:]),
					},
				}
			}(file, id)
		}

		wg.Wait()
		close(ch)
		<-done
	},
}

func init() {
	RootCmd.AddCommand(sumCmd)

	sumCmd.Flags().BoolP("circular", "c", false, "the file contains a single cicular genome sequence")
	sumCmd.Flags().IntP("kmer-size", "k", 1000, "k-mer size for processing circular genomes")
	sumCmd.Flags().BoolP("basename", "b", false, "only output basename of files")
	sumCmd.Flags().BoolP("remove-gaps", "g", false, "remove gaps")
	sumCmd.Flags().StringP("gap-letters", "G", "- 	.", "gap letters")
}

type SumResult struct {
	File   string
	SeqNum int
	SeqLen int
	Digest string
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/base64"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/semtech"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/brocaar/lorawan"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// uplinkCmd represents the `uplink` command
var uplinkCmd = &cobra.Command{
	Use:   "uplink [DevAddr] [NwkSKey] [AppSKey] [Payload] [FCnt]",
	Short: "Send uplink messages to the network",
	Long:  `ttnctl uplink sends an uplink message to the network`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 5 {
			ctx.Fatalf("Insufficient arguments")
		}

		// Parse parameters
		devAddrRaw, err := util.Parse32(args[0])
		if err != nil {
			ctx.Fatalf("Invalid DevAddr: %s", err)
		}
		var devAddr lorawan.DevAddr
		copy(devAddr[:], devAddrRaw)

		nwkSKeyRaw, err := util.Parse128(args[1])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}
		var nwkSKey lorawan.AES128Key
		copy(nwkSKey[:], nwkSKeyRaw[:])

		appSKeyRaw, err := util.Parse128(args[2])
		if err != nil {
			ctx.Fatalf("Invalid appSKey: %s", err)
		}
		var appSKey lorawan.AES128Key
		copy(appSKey[:], appSKeyRaw[:])

		fcnt, err := strconv.ParseInt(args[4], 10, 64)
		if err != nil {
			ctx.Fatalf("Invalid FCnt: %s", err)
		}

		// Lorawan Payload
		macPayload := lorawan.NewMACPayload(true)
		macPayload.FHDR = lorawan.FHDR{
			DevAddr: devAddr,
			FCnt:    uint32(fcnt),
		}
		macPayload.FPort = new(uint8)
		*macPayload.FPort = 1
		macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: []byte(args[3])}}
		if err := macPayload.EncryptFRMPayload(appSKey); err != nil {
			ctx.Fatalf("Unable to encrypt frame payload: %s", err)
		}
		phyPayload := lorawan.NewPHYPayload(true)
		phyPayload.MHDR = lorawan.MHDR{
			MType: lorawan.UnconfirmedDataUp,
			Major: lorawan.LoRaWANR1,
		}
		phyPayload.MACPayload = macPayload
		if err := phyPayload.SetMIC(nwkSKey); err != nil {
			ctx.Fatalf("Unable to set MIC: %s", err)
		}

		addr, err := net.ResolveUDPAddr("udp", viper.GetString("router.address"))
		if err != nil {
			ctx.Fatalf("Couldn't resolve UDP address: %s", err)
		}
		conn, err := net.DialUDP("udp", nil, addr)
		if err != nil {
			ctx.Fatalf("Couldn't Dial UDP connection: %s", err)
		}

		// Handle downlink
		chdown := make(chan bool)
		go func() {
			// Get Ack
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt := new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Ack: %s", pkt)

			// Get Downlink, if any
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err = pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Downlink: %s", pkt)
			defer func() { chdown <- true }()

			if pkt.Payload == nil || pkt.Payload.TXPK == nil || pkt.Payload.TXPK.Data == nil {
				ctx.Fatalf("No payload available in downlink response")
			}

			data, err := base64.RawStdEncoding.DecodeString(*pkt.Payload.TXPK.Data)
			if err != nil {
				ctx.Fatalf("Unable to decode data payload: %s", err)
			}

			payload := lorawan.NewPHYPayload(false)
			if err := payload.UnmarshalBinary(data); err != nil {
				ctx.Fatalf("Unable to retrieve LoRaWAN PhyPayload: %s", err)
			}

			macPayload, ok := payload.MACPayload.(*lorawan.MACPayload)
			if !ok || len(macPayload.FRMPayload) != 1 {
				ctx.Fatalf("Unable to retrieve LoRaWAN MACPayload")
			}
			if err := macPayload.DecryptFRMPayload(appSKey); err != nil {
				ctx.Fatalf("Unable to decrypt MACPayload: %s", err)
			}

			ctx.Infof("Frame counter: %d", macPayload.FHDR.FCnt)
			ctx.Infof("Decrypted Payload: %s", string(macPayload.FRMPayload[0].(*lorawan.DataPayload).Bytes))
		}()

		// Router Packet
		data, err := phyPayload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Couldn't construct LoRaWAN physical payload: %s", err)
		}
		encoded := strings.Trim(base64.StdEncoding.EncodeToString(data), "=")
		payload := semtech.Packet{
			Identifier: semtech.PUSH_DATA,
			Token:      util.RandToken(),
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Version:    semtech.VERSION,
			Payload: &semtech.Payload{
				RXPK: []semtech.RXPK{
					{
						Rssi: pointer.Int32(util.RandRssi()),
						Lsnr: pointer.Float32(util.RandLsnr()),
						Freq: pointer.Float32(util.RandFreq()),
						Datr: pointer.String(util.RandDatr()),
						Codr: pointer.String(util.RandCodr()),
						Modu: pointer.String("LoRa"),
						Tmst: pointer.Uint32(1),
						Data: &encoded,
					},
				},
			},
		}

		ctx.Infof("Sending packet: %s", payload.String())

		data, err = payload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Unable to construct framepayload: %v", data)
		}

		_, err = conn.Write(data)
		if err != nil {
			ctx.Fatal("Unable to send payload")
		}

		select {
		case <-chdown:
		case <-time.After(2 * time.Second):
		}
	},
}

func init() {
	RootCmd.AddCommand(uplinkCmd)

	uplinkCmd.Flags().String("ttn-router", "0.0.0.0:1700", "The net address of the TTN Router")
	viper.BindPFlag("router.address", uplinkCmd.Flags().Lookup("ttn-router"))
}
<commit_msg>Add support for confirmed data in ttnctl uplink<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"encoding/base64"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/TheThingsNetwork/ttn/semtech"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/TheThingsNetwork/ttn/utils/pointer"
	"github.com/brocaar/lorawan"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// uplinkCmd represents the `uplink` command
var uplinkCmd = &cobra.Command{
	Use:   "uplink [ShouldConfirm] [DevAddr] [NwkSKey] [AppSKey] [Payload] [FCnt]",
	Short: "Send uplink messages to the network",
	Long:  `ttnctl uplink sends an uplink message to the network`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 6 {
			ctx.Fatalf("Insufficient arguments")
		}

		// Parse parameters
		var mtype lorawan.MType
		switch args[0] {
		case "yes":
			fallthrough
		case "true":
			mtype = lorawan.ConfirmedDataUp
		default:
			mtype = lorawan.UnconfirmedDataUp
		}

		devAddrRaw, err := util.Parse32(args[1])
		if err != nil {
			ctx.Fatalf("Invalid DevAddr: %s", err)
		}
		var devAddr lorawan.DevAddr
		copy(devAddr[:], devAddrRaw)

		nwkSKeyRaw, err := util.Parse128(args[2])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}
		var nwkSKey lorawan.AES128Key
		copy(nwkSKey[:], nwkSKeyRaw[:])

		appSKeyRaw, err := util.Parse128(args[3])
		if err != nil {
			ctx.Fatalf("Invalid appSKey: %s", err)
		}
		var appSKey lorawan.AES128Key
		copy(appSKey[:], appSKeyRaw[:])

		fcnt, err := strconv.ParseInt(args[5], 10, 64)
		if err != nil {
			ctx.Fatalf("Invalid FCnt: %s", err)
		}

		// Lorawan Payload
		macPayload := lorawan.NewMACPayload(true)
		macPayload.FHDR = lorawan.FHDR{
			DevAddr: devAddr,
			FCnt:    uint32(fcnt),
		}
		macPayload.FPort = new(uint8)
		*macPayload.FPort = 1
		macPayload.FRMPayload = []lorawan.Payload{&lorawan.DataPayload{Bytes: []byte(args[4])}}
		if err := macPayload.EncryptFRMPayload(appSKey); err != nil {
			ctx.Fatalf("Unable to encrypt frame payload: %s", err)
		}
		phyPayload := lorawan.NewPHYPayload(true)
		phyPayload.MHDR = lorawan.MHDR{
			MType: mtype,
			Major: lorawan.LoRaWANR1,
		}
		phyPayload.MACPayload = macPayload
		if err := phyPayload.SetMIC(nwkSKey); err != nil {
			ctx.Fatalf("Unable to set MIC: %s", err)
		}

		addr, err := net.ResolveUDPAddr("udp", viper.GetString("router.address"))
		if err != nil {
			ctx.Fatalf("Couldn't resolve UDP address: %s", err)
		}
		conn, err := net.DialUDP("udp", nil, addr)
		if err != nil {
			ctx.Fatalf("Couldn't Dial UDP connection: %s", err)
		}

		// Handle downlink
		chdown := make(chan bool)
		go func() {
			// Get Ack
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt := new(semtech.Packet)
			if err := pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Ack: %s", pkt)

			// Get Downlink, if any
			buf = make([]byte, 1024)
			n, err = conn.Read(buf)
			if err != nil {
				ctx.Fatalf("Error receiving udp datagram: %s", err)
			}
			pkt = new(semtech.Packet)
			if err = pkt.UnmarshalBinary(buf[:n]); err != nil {
				ctx.Fatalf("Invalid udp response: %s", err)
			}
			ctx.Infof("Received Downlink: %s", pkt)
			defer func() { chdown <- true }()

			if pkt.Payload == nil || pkt.Payload.TXPK == nil || pkt.Payload.TXPK.Data == nil {
				ctx.Fatalf("No payload available in downlink response")
			}

			data, err := base64.RawStdEncoding.DecodeString(*pkt.Payload.TXPK.Data)
			if err != nil {
				ctx.Fatalf("Unable to decode data payload: %s", err)
			}

			payload := lorawan.NewPHYPayload(false)
			if err := payload.UnmarshalBinary(data); err != nil {
				ctx.Fatalf("Unable to retrieve LoRaWAN PhyPayload: %s", err)
			}

			macPayload, ok := payload.MACPayload.(*lorawan.MACPayload)
			if !ok || len(macPayload.FRMPayload) > 1 {
				ctx.Fatalf("Unable to retrieve LoRaWAN MACPayload")
			}
			ctx.Infof("Frame counter: %d", macPayload.FHDR.FCnt)
			if len(macPayload.FRMPayload) > 0 {
				if err := macPayload.DecryptFRMPayload(appSKey); err != nil {
					ctx.Fatalf("Unable to decrypt MACPayload: %s", err)
				}
				ctx.Infof("Decrypted Payload: %s", string(macPayload.FRMPayload[0].(*lorawan.DataPayload).Bytes))
			} else {
				ctx.Infof("The frame payload was empty.")
			}
		}()

		// Router Packet
		data, err := phyPayload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Couldn't construct LoRaWAN physical payload: %s", err)
		}
		encoded := strings.Trim(base64.StdEncoding.EncodeToString(data), "=")
		payload := semtech.Packet{
			Identifier: semtech.PUSH_DATA,
			Token:      util.RandToken(),
			GatewayId:  []byte{1, 2, 3, 4, 5, 6, 7, 8},
			Version:    semtech.VERSION,
			Payload: &semtech.Payload{
				RXPK: []semtech.RXPK{
					{
						Rssi: pointer.Int32(util.RandRssi()),
						Lsnr: pointer.Float32(util.RandLsnr()),
						Freq: pointer.Float32(util.RandFreq()),
						Datr: pointer.String(util.RandDatr()),
						Codr: pointer.String(util.RandCodr()),
						Modu: pointer.String("LoRa"),
						Tmst: pointer.Uint32(1),
						Data: &encoded,
					},
				},
			},
		}

		ctx.Infof("Sending packet: %s", payload.String())

		data, err = payload.MarshalBinary()
		if err != nil {
			ctx.Fatalf("Unable to construct framepayload: %v", data)
		}

		_, err = conn.Write(data)
		if err != nil {
			ctx.Fatal("Unable to send payload")
		}

		select {
		case <-chdown:
		case <-time.After(2 * time.Second):
		}
	},
}

func init() {
	RootCmd.AddCommand(uplinkCmd)

	uplinkCmd.Flags().String("ttn-router", "0.0.0.0:1700", "The net address of the TTN Router")
	viper.BindPFlag("router.address", uplinkCmd.Flags().Lookup("ttn-router"))
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/gocraft/web"
	"gopkg.in/pfnet/jasco.v1"
	"gopkg.in/sensorbee/sensorbee.v0/bql"
	"gopkg.in/sensorbee/sensorbee.v0/bql/parser"
	"gopkg.in/sensorbee/sensorbee.v0/bql/udf"
	"gopkg.in/sensorbee/sensorbee.v0/core"
	"gopkg.in/sensorbee/sensorbee.v0/data"
	"gopkg.in/sensorbee/sensorbee.v0/server/config"
	"gopkg.in/sensorbee/sensorbee.v0/server/udsstorage"
	"io"
	"io/ioutil"
)

// Context is a context object for gocraft/web.
type Context struct {
	*jasco.Context

	udsStorage udf.UDSStorage
	topologies TopologyRegistry
	config     *config.Config
	// logger is used by core.Context, not for the server's Context. This logger
	// can be shared with jasco.Context.
	logger *logrus.Logger
}

// SetTopologyRegistry sets the registry of topologies to this context. This
// method must be called in the middleware of Context.
func (c *Context) SetTopologyRegistry(r TopologyRegistry) {
	c.topologies = r
}

// ContextGlobalVariables has fields which are shared through all contexts
// allocated for each request.
type ContextGlobalVariables struct {
	// Logger is used to write log messages.
	Logger *logrus.Logger

	// LogDestination is a writer to which logs are written.
	LogDestination io.WriteCloser

	// Topologies is a registry which manages topologies to support multi
	// tenancy.
	Topologies TopologyRegistry

	// Config has configuration parameters.
	Config *config.Config
}

// SetUpContextGlobalVariables create a new ContextGlobalVariables from a config.
// DO NOT make any change on the config after calling this function. The caller
// can change other members of ContextGlobalVariables.
//
// The caller must Close LogDestination.
func SetUpContextGlobalVariables(conf *config.Config) (*ContextGlobalVariables, error) {
	logger := logrus.New()
	w, err := conf.Logging.CreateWriter()
	if err != nil {
		return nil, err
	}
	closeWriter := true
	defer func() {
		if closeWriter {
			w.Close()
		}
	}()
	logger.Out = w

	closeWriter = false
	return &ContextGlobalVariables{
		Logger:         logger,
		LogDestination: w,
		Topologies:     NewDefaultTopologyRegistry(),
		Config:         conf,
	}, nil
}

// SetUpContextAndRouter creates a router of the API server and its context.
// jascoRoot is a root router returned from jasco.New.
//
// This function returns a new web.Router. Don't use the router returned from
// this function as a handler of HTTP server, but use jascoRoot instead.
func SetUpContextAndRouter(prefix string, jascoRoot *web.Router, gvariables *ContextGlobalVariables) (*web.Router, error) {
	gvars := *gvariables
	udsStorage, err := setUpUDSStorage(&gvars.Config.Storage.UDS)
	if err != nil {
		return nil, err
	}

	// Topologies should be created after setting up everything necessary for it.
	if err := setUpTopologies(gvars.Logger, gvars.Topologies, gvars.Config, udsStorage); err != nil {
		return nil, err
	}

	router := jascoRoot.Subrouter(Context{}, "/")
	router.Middleware(func(c *Context, rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc) {
		c.logger = gvars.Logger
		c.udsStorage = udsStorage
		c.topologies = gvars.Topologies
		c.config = gvars.Config
		next(rw, req)
	})
	return router, nil
}

func setUpUDSStorage(conf *config.UDSStorage) (udf.UDSStorage, error) {
	// Parameters are already validated in conf
	switch conf.Type {
	case "in_memory":
		return udf.NewInMemoryUDSStorage(), nil
	case "fs":
		dir, _ := data.AsString(conf.Params["dir"])
		var tempDir string
		if v, ok := conf.Params["temp_dir"]; ok {
			tempDir, _ = data.AsString(v)
		}
		return udsstorage.NewFS(dir, tempDir)
	default:
		return nil, fmt.Errorf("unsupported uds storage type: %v", conf.Type)
	}
}

func setUpTopologies(logger *logrus.Logger, r TopologyRegistry, conf *config.Config, us udf.UDSStorage) error {
	stopAll := true
	defer func() {
		if stopAll {
			ts, err := r.List()
			if err != nil {
				logger.WithField("err", err).Error("Cannot list topologies for clean up")
				return
			}

			for name, t := range ts {
				if err := t.Topology().Stop(); err != nil {
					logger.WithFields(logrus.Fields{
						"err":      err,
						"topology": name,
					}).Error("Cannot stop the topology")
				}
			}
		}
	}()

	for name := range conf.Topologies {
		logger.WithField("topology", name).Info("Setting up the topology")
		tb, err := setUpTopology(logger, name, conf, us)
		if err != nil {
			return err
		}
		if err := r.Register(name, tb); err != nil {
			logger.WithFields(logrus.Fields{
				"err":      err,
				"topology": name,
			}).Error("Cannot register the topology")
			return err
		}
	}

	stopAll = false
	return nil
}

func setUpTopology(logger *logrus.Logger, name string, conf *config.Config, us udf.UDSStorage) (*bql.TopologyBuilder, error) {
	cc := &core.ContextConfig{
		Logger: logger,
	}
	cc.Flags.DroppedTupleLog.Set(conf.Logging.LogDroppedTuples)
	cc.Flags.DroppedTupleSummarization.Set(conf.Logging.SummarizeDroppedTuples)

	tp, err := core.NewDefaultTopology(core.NewContext(cc), name)
	if err != nil {
		return nil, err
	}
	tb, err := bql.NewTopologyBuilder(tp)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"topology": name,
		}).Error("Cannot create a topology builder")
		return nil, err
	}
	tb.UDSStorage = us

	bqlFilePath := conf.Topologies[name].BQLFile
	if bqlFilePath == "" {
		return tb, nil
	}

	shouldStop := true
	defer func() {
		if shouldStop {
			if err := tp.Stop(); err != nil {
				logger.WithFields(logrus.Fields{
					"err":      err,
					"topology": name,
				}).Error("Cannot stop the topology")
			}
		}
	}()

	queries, err := ioutil.ReadFile(bqlFilePath)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"topology": name,
			// TODO: add stmt.String()
		}).Error("Cannot add a statement to the topology")
		return nil, err
	}

	// TODO: improve error handling
	bp := parser.New()
	stmts, err := bp.ParseStmts(string(queries))
	if err != nil {
		return nil, err
	}

	for _, stmt := range stmts {
		if _, err := tb.AddStmt(stmt); err != nil {
			logger.WithFields(logrus.Fields{
				"err":      err,
				"topology": name,
				// TODO: add stmt.String()
			}).Error("Cannot add a statement to the topology")
			return nil, err
		}
	}

	shouldStop = false
	return tb, nil
}
<commit_msg>Improve log messages on startup<commit_after>package server

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/gocraft/web"
	"gopkg.in/pfnet/jasco.v1"
	"gopkg.in/sensorbee/sensorbee.v0/bql"
	"gopkg.in/sensorbee/sensorbee.v0/bql/parser"
	"gopkg.in/sensorbee/sensorbee.v0/bql/udf"
	"gopkg.in/sensorbee/sensorbee.v0/core"
	"gopkg.in/sensorbee/sensorbee.v0/data"
	"gopkg.in/sensorbee/sensorbee.v0/server/config"
	"gopkg.in/sensorbee/sensorbee.v0/server/udsstorage"
	"io"
	"io/ioutil"
)

// Context is a context object for gocraft/web.
type Context struct {
	*jasco.Context

	udsStorage udf.UDSStorage
	topologies TopologyRegistry
	config     *config.Config
	// logger is used by core.Context, not for the server's Context. This logger
	// can be shared with jasco.Context.
	logger *logrus.Logger
}

// SetTopologyRegistry sets the registry of topologies to this context. This
// method must be called in the middleware of Context.
func (c *Context) SetTopologyRegistry(r TopologyRegistry) {
	c.topologies = r
}

// ContextGlobalVariables has fields which are shared through all contexts
// allocated for each request.
type ContextGlobalVariables struct {
	// Logger is used to write log messages.
	Logger *logrus.Logger

	// LogDestination is a writer to which logs are written.
	LogDestination io.WriteCloser

	// Topologies is a registry which manages topologies to support multi
	// tenancy.
	Topologies TopologyRegistry

	// Config has configuration parameters.
	Config *config.Config
}

// SetUpContextGlobalVariables create a new ContextGlobalVariables from a config.
// DO NOT make any change on the config after calling this function. The caller
// can change other members of ContextGlobalVariables.
//
// The caller must Close LogDestination.
func SetUpContextGlobalVariables(conf *config.Config) (*ContextGlobalVariables, error) {
	logger := logrus.New()
	w, err := conf.Logging.CreateWriter()
	if err != nil {
		return nil, err
	}
	closeWriter := true
	defer func() {
		if closeWriter {
			w.Close()
		}
	}()
	logger.Out = w

	closeWriter = false
	return &ContextGlobalVariables{
		Logger:         logger,
		LogDestination: w,
		Topologies:     NewDefaultTopologyRegistry(),
		Config:         conf,
	}, nil
}

// SetUpContextAndRouter creates a router of the API server and its context.
// jascoRoot is a root router returned from jasco.New.
//
// This function returns a new web.Router. Don't use the router returned from
// this function as a handler of HTTP server, but use jascoRoot instead.
func SetUpContextAndRouter(prefix string, jascoRoot *web.Router, gvariables *ContextGlobalVariables) (*web.Router, error) {
	gvars := *gvariables
	udsStorage, err := setUpUDSStorage(&gvars.Config.Storage.UDS)
	if err != nil {
		return nil, err
	}

	// Topologies should be created after setting up everything necessary for it.
	if err := setUpTopologies(gvars.Logger, gvars.Topologies, gvars.Config, udsStorage); err != nil {
		return nil, err
	}

	router := jascoRoot.Subrouter(Context{}, "/")
	router.Middleware(func(c *Context, rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc) {
		c.logger = gvars.Logger
		c.udsStorage = udsStorage
		c.topologies = gvars.Topologies
		c.config = gvars.Config
		next(rw, req)
	})
	return router, nil
}

func setUpUDSStorage(conf *config.UDSStorage) (udf.UDSStorage, error) {
	// Parameters are already validated in conf
	switch conf.Type {
	case "in_memory":
		return udf.NewInMemoryUDSStorage(), nil
	case "fs":
		dir, _ := data.AsString(conf.Params["dir"])
		var tempDir string
		if v, ok := conf.Params["temp_dir"]; ok {
			tempDir, _ = data.AsString(v)
		}
		return udsstorage.NewFS(dir, tempDir)
	default:
		return nil, fmt.Errorf("unsupported uds storage type: %v", conf.Type)
	}
}

func setUpTopologies(logger *logrus.Logger, r TopologyRegistry, conf *config.Config, us udf.UDSStorage) error {
	stopAll := true
	defer func() {
		if stopAll {
			ts, err := r.List()
			if err != nil {
				logger.WithField("err", err).Error("Cannot list topologies for clean up")
				return
			}

			for name, t := range ts {
				if err := t.Topology().Stop(); err != nil {
					logger.WithFields(logrus.Fields{
						"err":      err,
						"topology": name,
					}).Error("Cannot stop the topology")
				}
			}
		}
	}()

	for name := range conf.Topologies {
		logger.WithField("topology", name).Info("Setting up the topology")
		tb, err := setUpTopology(logger, name, conf, us)
		if err != nil {
			return err
		}
		if err := r.Register(name, tb); err != nil {
			logger.WithFields(logrus.Fields{
				"err":      err,
				"topology": name,
			}).Error("Cannot register the topology")
			return err
		}
	}

	stopAll = false
	return nil
}

func setUpTopology(logger *logrus.Logger, name string, conf *config.Config, us udf.UDSStorage) (*bql.TopologyBuilder, error) {
	cc := &core.ContextConfig{
		Logger: logger,
	}
	cc.Flags.DroppedTupleLog.Set(conf.Logging.LogDroppedTuples)
	cc.Flags.DroppedTupleSummarization.Set(conf.Logging.SummarizeDroppedTuples)

	tp, err := core.NewDefaultTopology(core.NewContext(cc), name)
	if err != nil {
		return nil, err
	}
	tb, err := bql.NewTopologyBuilder(tp)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"topology": name,
		}).Error("Cannot create a topology builder")
		return nil, err
	}
	tb.UDSStorage = us

	bqlFilePath := conf.Topologies[name].BQLFile
	if bqlFilePath == "" {
		return tb, nil
	}

	shouldStop := true
	defer func() {
		if shouldStop {
			if err := tp.Stop(); err != nil {
				logger.WithFields(logrus.Fields{
					"err":      err,
					"topology": name,
				}).Error("Cannot stop the topology")
			}
		}
	}()

	queries, err := ioutil.ReadFile(bqlFilePath)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"err":      err,
			"topology": name,
			"path":     bqlFilePath,
		}).Error("Cannot read a BQL file")
		return nil, err
	}

	// TODO: improve error handling
	bp := parser.New()
	stmts, err := bp.ParseStmts(string(queries))
	if err != nil {
		return nil, err
	}

	for _, stmt := range stmts {
		if _, err := tb.AddStmt(stmt); err != nil {
			logger.WithFields(logrus.Fields{
				"err":      err,
				"topology": name,
				"stmt":     stmt,
			}).Error("Cannot add a statement to the topology")
			return nil, err
		}
	}

	shouldStop = false
	return tb, nil
}
<|endoftext|>


<commit_before>package bucket

import (
	"bufio"
	"github.com/qiniu/qshell/v2/iqshell/common/alert"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/file"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/common/workspace"
	"github.com/qiniu/qshell/v2/iqshell/storage/bucket/internal/list"
	"io"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type ListApiInfo struct {
	Bucket             string    // 空间名	【必选】
	Prefix             string    // 前缀
	Marker             string    // 标记
	Delimiter          string    //
	StartTime          time.Time // list item 的 put time 区间的开始时间 【闭区间】
	EndTime            time.Time // list item 的 put time 区间的终止时间 【闭区间】
	Suffixes           []string  // list item 必须包含后缀
	StorageTypes       []int     // list item 存储类型，多个使用逗号隔开， 0:普通存储 1:低频存储 2:归档存储 3:深度归档存储
	MimeTypes          []string  // list item Mimetype类型，多个使用逗号隔开
	MinFileSize        int64     // 文件最小值，单位: B
	MaxFileSize        int64     // 文件最大值，单位: B
	MaxRetry           int       // -1: 无限重试
	ShowFields         []string  // 需要展示的字段  【必选】
	ApiVersion         string    // list api 版本，v1 / v2【可选】
	V1Limit            int       // 每次请求 size ，list v1 特有
	OutputLimit        int       // 最大输出条数，默认：-1, 无限输出
	OutputFieldsSep    string    // 输出信息，每行的分隔符 【必选】
	OutputFileMaxLines int64     // 输出文件的最大行数，超过则自动创建新的文件，0：不限制输出文件的行数 【可选】
	OutputFileMaxSize  int64     // 输出文件的最大 Size，超过则自动创建新的文件，0：不限制输出文件的大小 【可选】
	EnableRecord       bool      // 是否开启 record 记录，开启后会记录 list 信息，下次 list 会自动指定 Marker 继续 list 【可选】
	CacheDir           string    // 历史数据存储路径 【内部使用】
}

func ListObjectField(field string) string {
	for _, f := range listObjectFields {
		if strings.EqualFold(field, f) {
			return f
		}
	}
	return ""
}

type ListObject = list.Item

// List list 某个 bucket 所有的文件
func List(info ListApiInfo,
	objectHandler func(marker string, object ListObject) (shouldContinue bool, err *data.CodeError),
	errorHandler func(marker string, err *data.CodeError)) {
	if objectHandler == nil {
		log.Error(alert.CannotEmpty("list bucket: object handler", ""))
		return
	}

	if errorHandler == nil {
		errorHandler = func(marker string, err *data.CodeError) {
			log.ErrorF("marker: %s", info.Marker)
			log.ErrorF("list bucket Error: %v", err)
		}
		log.Warning("list bucket: not set error handler")
	}

	bucketManager, err := GetBucketManager()
	if err != nil {
		errorHandler("", err)
		return
	}

	log.Debug("will list bucket")
	log.DebugF("Suffixes:%s", info.Suffixes)
	shouldCheckPutTime := !info.StartTime.IsZero() || !info.EndTime.IsZero()
	shouldCheckSuffixes := len(info.Suffixes) > 0
	shouldCheckStorageTypes := len(info.StorageTypes) > 0
	shouldCheckMimeTypes := len(info.MimeTypes) > 0
	shouldCheckFileSize := info.MinFileSize > 0 || info.MaxFileSize > 0
	isItemExcepted := func(listItem list.Item) (isExcepted bool) {
		if shouldCheckPutTime {
			putTime := time.Unix(listItem.PutTime/1e7, 0)
			if !filterByPutTime(putTime, info.StartTime, info.EndTime) {
				log.DebugF("filter %s: putTime not match, %s out of range [start:%s ~ end:%s]", listItem.Key, putTime, info.StartTime, info.EndTime)
				return false
			}
		}

		if shouldCheckSuffixes && !filterBySuffixes(listItem.Key, info.Suffixes) {
			log.DebugF("filter %s: key not match, key:%s suffixes:%s ", listItem.Key, listItem.Key, info.Suffixes)
			return false
		}

		if shouldCheckStorageTypes && !filterByStorageType(listItem.Type, info.StorageTypes) {
			log.DebugF("filter %s: key not match, storageType:%d StorageTypes:%s ", listItem.Key, listItem.Type, info.Suffixes)
			return false
		}

		if shouldCheckMimeTypes && !filterByMimeType(listItem.MimeType, info.MimeTypes) {
			log.DebugF("filter %s: key not match, mimeType:%s mimeTypes:%s ", listItem.Key, listItem.MimeType, info.MimeTypes)
			return false
		}

		if shouldCheckFileSize && !filterByFileSize(listItem.Fsize, info.MinFileSize, info.MaxFileSize) {
			log.DebugF("filter %s: key not match, fileSize:%d minSize:%d maxSize:%d", listItem.Key, listItem.Fsize, info.MinFileSize, info.MaxFileSize)
			return false
		}

		return true
	}

	listWaiter := sync.WaitGroup{}
	listWaiter.Add(1)
	workspace.AddCancelObserver(func(s os.Signal) {
		listWaiter.Wait()
	})

	cache := &listCache{
		enableRecord: info.EnableRecord,
		cachePath:    filepath.Join(info.CacheDir, "info.json"),
	}
	cacheInfoP, err := cache.loadCache()
	if err != nil {
		log.Debug(err)
	}

	if cacheInfoP != nil && len(cacheInfoP.Marker) > 0 {
		if len(info.Marker) == 0 {
			info.Marker = cacheInfoP.Marker
		}
	}

	if cacheInfoP == nil {
		cacheInfoP = &cacheInfo{}
	} else if len(cacheInfoP.Marker) > 0 {
		log.InfoF("use marker:%s", cacheInfoP.Marker)
	}

	retryCount := 0
	outputCount := 0
	complete := false
	for !complete && (info.MaxRetry < 0 || retryCount <= info.MaxRetry) {
		var hasMore = false
		var lErr *data.CodeError = nil

		if !workspace.IsCmdInterrupt() {
			hasMore, lErr = list.ListBucket(workspace.GetContext(), list.ApiInfo{
				Manager:    bucketManager,
				ApiVersion: list.ApiVersion(info.ApiVersion),
				Bucket:     info.Bucket,
				Prefix:     info.Prefix,
				Delimiter:  info.Delimiter,
				Marker:     info.Marker,
				V1Limit:    info.V1Limit,
			}, func(marker string, dir string, listItem list.Item) (stop bool) {
				if marker != info.Marker {
					info.Marker = marker
				}

				if !isItemExcepted(listItem) {
					return false
				}

				shouldContinue, hErr := objectHandler(marker, listItem)
				if hErr != nil {
					errorHandler(marker, hErr)
				}
				if !shouldContinue {
					complete = true
					return true
				}

				outputCount++
				if info.OutputLimit > 0 && outputCount >= info.OutputLimit {
					complete = true
					return true
				}

				return false
			})
		}

		// 保存信息
		cacheInfoP.Bucket = info.Bucket
		cacheInfoP.Prefix = info.Prefix
		cacheInfoP.Marker = info.Marker
		_ = cache.saveCache(cacheInfoP)

		if workspace.IsCmdInterrupt() && lErr == nil {
			lErr = data.NewError(0, "list is interrupted")
		}

		if lErr != nil || workspace.IsCmdInterrupt() {
			errorHandler(info.Marker, lErr)

			if workspace.IsCmdInterrupt() || // 取消
				strings.Contains(lErr.Error(), "no such bucket") || // 空间不存在，直接结束
				strings.Contains(lErr.Error(), "incorrect zone") || // 空间不正确
				strings.Contains(lErr.Error(), "context canceled") { // 取消
				break
			}

			retryCount++
			time.Sleep(time.Millisecond * 100)
			continue
		}

		if !hasMore || complete || workspace.IsCmdInterrupt() {
			break
		}

		retryCount = 0
	}

	if len(info.Marker) == 0 {
		if rErr := cache.removeCache(); rErr != nil {
			log.ErrorF("list remove cache status error: %v", rErr)
		} else {
			log.InfoF("list success, remove cache status: %s", cache.cachePath)
		}
	} else {
		log.InfoF("Marker: %s", info.Marker)
	}

	log.Debug("list bucket end")

	listWaiter.Done()
}

type ListToFileApiInfo struct {
	ListApiInfo
	FilePath   string // file 不存在则输出到 stdout
	AppendMode bool
	Readable   bool
}

func ListToFile(info ListToFileApiInfo, errorHandler func(marker string, err *data.CodeError)) {
	if errorHandler == nil {
		errorHandler = func(marker string, err *data.CodeError) {
			log.ErrorF("marker: %s", marker)
			log.ErrorF("list bucket Error: %v", err)
		}
		log.Warning("list bucket to file: not set error handler")
	}

	if len(info.ShowFields) == 0 {
		info.ShowFields = listObjectFields
	}

	if len(info.OutputFieldsSep) == 0 {
		info.OutputFieldsSep = data.DefaultLineSeparate
	}

	// 文件头
	title := strings.Join(info.ShowFields, info.OutputFieldsSep)

	var output io.WriteCloser
	if info.FilePath == "" {
		output = data.Stdout()
		_, _ = output.Write([]byte(title + "\n"))
		log.Debug("prepare list bucket to stdout")
	} else {
		var nErr *data.CodeError
		output, nErr = file.NewRotateFile(info.FilePath,
			file.RotateOptionMaxSize(info.OutputFileMaxSize),
			file.RotateOptionMaxLine(info.OutputFileMaxLines),
			file.RotateOptionAppendMode(info.AppendMode),
			file.RotateOptionFileHeader(title),
			file.RotateOptionOnOpenFile(func(filename string) {
				log.InfoF("open output file and prepare to write:%v", filename)
			}))

		if nErr != nil {
			errorHandler("", data.NewEmptyError().AppendDescF("failed to create rotate file:`%s`, error:%v", info.FilePath, nErr))
			return
		}
		defer output.Close()
		log.Debug("prepare list bucket to file")
	}

	bWriter := bufio.NewWriter(output)
	lineCreator := &ListLineCreator{
		Fields:   info.ShowFields,
		Sep:      info.OutputFieldsSep,
		Readable: info.Readable,
	}
	List(info.ListApiInfo, func(marker string, object ListObject) (bool, *data.CodeError) {
		lineData := lineCreator.Create(&object)
		if _, wErr := bWriter.WriteString(lineData + "\n"); wErr != nil {
			return false, data.NewEmptyError().AppendDesc("write error:" + wErr.Error())
		}

		if fErr := bWriter.Flush(); fErr != nil {
			return false, data.NewEmptyError().AppendDesc("flush error:" + fErr.Error())
		}
		return true, nil
	}, errorHandler)
}

func filterByPutTime(putTime, startDate, endDate time.Time) bool {
	switch {
	case startDate.IsZero() && endDate.IsZero():
		return true
	case !startDate.IsZero() && endDate.IsZero() && putTime.After(startDate):
		return true
	case !endDate.IsZero() && startDate.IsZero() && putTime.Before(endDate):
		return true
	case putTime.After(startDate) && putTime.Before(endDate):
		return true
	default:
		return false
	}
}

func filterBySuffixes(key string, suffixes []string) bool {
	hasSuffix := false
	if len(suffixes) == 0 {
		hasSuffix = true
	}
	for _, s := range suffixes {
		if strings.HasSuffix(key, s) {
			hasSuffix = true
			break
		}
	}
	return hasSuffix
}

func filterByStorageType(storageType int, storageTypes []int) bool {
	hasStorageType := false
	if len(storageTypes) == 0 {
		hasStorageType = true
	}
	for _, s := range storageTypes {
		if storageType == s {
			hasStorageType = true
			break
		}
	}
	return hasStorageType
}

func filterByMimeType(mimeType string, mimeTypes []string) bool {
	hasMimeType := false
	if len(mimeTypes) == 0 {
		hasMimeType = true
	}
	for _, s := range mimeTypes {
		if strings.Contains(s, "*") {
			sp := strings.ReplaceAll(s, "*", "")
			if strings.Contains(mimeType, sp) {
				hasMimeType = true
				break
			}
		} else if mimeType == s {
			hasMimeType = true
			break
		}
	}
	return hasMimeType
}

func filterByFileSize(fileSize, minSize, maxSize int64) bool {
	if maxSize < 0 {
		maxSize = math.MaxInt64
	}
	return fileSize >= minSize && fileSize <= maxSize
}
<commit_msg>list bucket: api v1 invalid list limit not retry<commit_after>package bucket

import (
	"bufio"
	"github.com/qiniu/qshell/v2/iqshell/common/alert"
	"github.com/qiniu/qshell/v2/iqshell/common/data"
	"github.com/qiniu/qshell/v2/iqshell/common/file"
	"github.com/qiniu/qshell/v2/iqshell/common/log"
	"github.com/qiniu/qshell/v2/iqshell/common/workspace"
	"github.com/qiniu/qshell/v2/iqshell/storage/bucket/internal/list"
	"io"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

type ListApiInfo struct {
	Bucket             string    // 空间名	【必选】
	Prefix             string    // 前缀
	Marker             string    // 标记
	Delimiter          string    //
	StartTime          time.Time // list item 的 put time 区间的开始时间 【闭区间】
	EndTime            time.Time // list item 的 put time 区间的终止时间 【闭区间】
	Suffixes           []string  // list item 必须包含后缀
	StorageTypes       []int     // list item 存储类型，多个使用逗号隔开， 0:普通存储 1:低频存储 2:归档存储 3:深度归档存储
	MimeTypes          []string  // list item Mimetype类型，多个使用逗号隔开
	MinFileSize        int64     // 文件最小值，单位: B
	MaxFileSize        int64     // 文件最大值，单位: B
	MaxRetry           int       // -1: 无限重试
	ShowFields         []string  // 需要展示的字段  【必选】
	ApiVersion         string    // list api 版本，v1 / v2【可选】
	V1Limit            int       // 每次请求 size ，list v1 特有
	OutputLimit        int       // 最大输出条数，默认：-1, 无限输出
	OutputFieldsSep    string    // 输出信息，每行的分隔符 【必选】
	OutputFileMaxLines int64     // 输出文件的最大行数，超过则自动创建新的文件，0：不限制输出文件的行数 【可选】
	OutputFileMaxSize  int64     // 输出文件的最大 Size，超过则自动创建新的文件，0：不限制输出文件的大小 【可选】
	EnableRecord       bool      // 是否开启 record 记录，开启后会记录 list 信息，下次 list 会自动指定 Marker 继续 list 【可选】
	CacheDir           string    // 历史数据存储路径 【内部使用】
}

func ListObjectField(field string) string {
	for _, f := range listObjectFields {
		if strings.EqualFold(field, f) {
			return f
		}
	}
	return ""
}

type ListObject = list.Item

// List list 某个 bucket 所有的文件
func List(info ListApiInfo,
	objectHandler func(marker string, object ListObject) (shouldContinue bool, err *data.CodeError),
	errorHandler func(marker string, err *data.CodeError)) {
	if objectHandler == nil {
		log.Error(alert.CannotEmpty("list bucket: object handler", ""))
		return
	}

	if errorHandler == nil {
		errorHandler = func(marker string, err *data.CodeError) {
			log.ErrorF("marker: %s", info.Marker)
			log.ErrorF("list bucket Error: %v", err)
		}
		log.Warning("list bucket: not set error handler")
	}

	bucketManager, err := GetBucketManager()
	if err != nil {
		errorHandler("", err)
		return
	}

	log.Debug("will list bucket")
	log.DebugF("Suffixes:%s", info.Suffixes)
	shouldCheckPutTime := !info.StartTime.IsZero() || !info.EndTime.IsZero()
	shouldCheckSuffixes := len(info.Suffixes) > 0
	shouldCheckStorageTypes := len(info.StorageTypes) > 0
	shouldCheckMimeTypes := len(info.MimeTypes) > 0
	shouldCheckFileSize := info.MinFileSize > 0 || info.MaxFileSize > 0
	isItemExcepted := func(listItem list.Item) (isExcepted bool) {
		if shouldCheckPutTime {
			putTime := time.Unix(listItem.PutTime/1e7, 0)
			if !filterByPutTime(putTime, info.StartTime, info.EndTime) {
				log.DebugF("filter %s: putTime not match, %s out of range [start:%s ~ end:%s]", listItem.Key, putTime, info.StartTime, info.EndTime)
				return false
			}
		}

		if shouldCheckSuffixes && !filterBySuffixes(listItem.Key, info.Suffixes) {
			log.DebugF("filter %s: key not match, key:%s suffixes:%s ", listItem.Key, listItem.Key, info.Suffixes)
			return false
		}

		if shouldCheckStorageTypes && !filterByStorageType(listItem.Type, info.StorageTypes) {
			log.DebugF("filter %s: key not match, storageType:%d StorageTypes:%s ", listItem.Key, listItem.Type, info.Suffixes)
			return false
		}

		if shouldCheckMimeTypes && !filterByMimeType(listItem.MimeType, info.MimeTypes) {
			log.DebugF("filter %s: key not match, mimeType:%s mimeTypes:%s ", listItem.Key, listItem.MimeType, info.MimeTypes)
			return false
		}

		if shouldCheckFileSize && !filterByFileSize(listItem.Fsize, info.MinFileSize, info.MaxFileSize) {
			log.DebugF("filter %s: key not match, fileSize:%d minSize:%d maxSize:%d", listItem.Key, listItem.Fsize, info.MinFileSize, info.MaxFileSize)
			return false
		}

		return true
	}

	listWaiter := sync.WaitGroup{}
	listWaiter.Add(1)
	workspace.AddCancelObserver(func(s os.Signal) {
		listWaiter.Wait()
	})

	cache := &listCache{
		enableRecord: info.EnableRecord,
		cachePath:    filepath.Join(info.CacheDir, "info.json"),
	}
	cacheInfoP, err := cache.loadCache()
	if err != nil {
		log.Debug(err)
	}

	if cacheInfoP != nil && len(cacheInfoP.Marker) > 0 {
		if len(info.Marker) == 0 {
			info.Marker = cacheInfoP.Marker
		}
	}

	if cacheInfoP == nil {
		cacheInfoP = &cacheInfo{}
	} else if len(cacheInfoP.Marker) > 0 {
		log.InfoF("use marker:%s", cacheInfoP.Marker)
	}

	retryCount := 0
	outputCount := 0
	complete := false
	for !complete && (info.MaxRetry < 0 || retryCount <= info.MaxRetry) {
		var hasMore = false
		var lErr *data.CodeError = nil

		if !workspace.IsCmdInterrupt() {
			hasMore, lErr = list.ListBucket(workspace.GetContext(), list.ApiInfo{
				Manager:    bucketManager,
				ApiVersion: list.ApiVersion(info.ApiVersion),
				Bucket:     info.Bucket,
				Prefix:     info.Prefix,
				Delimiter:  info.Delimiter,
				Marker:     info.Marker,
				V1Limit:    info.V1Limit,
			}, func(marker string, dir string, listItem list.Item) (stop bool) {
				if marker != info.Marker {
					info.Marker = marker
				}

				if !isItemExcepted(listItem) {
					return false
				}

				shouldContinue, hErr := objectHandler(marker, listItem)
				if hErr != nil {
					errorHandler(marker, hErr)
				}
				if !shouldContinue {
					complete = true
					return true
				}

				outputCount++
				if info.OutputLimit > 0 && outputCount >= info.OutputLimit {
					complete = true
					return true
				}

				return false
			})
		}

		// 保存信息
		cacheInfoP.Bucket = info.Bucket
		cacheInfoP.Prefix = info.Prefix
		cacheInfoP.Marker = info.Marker
		_ = cache.saveCache(cacheInfoP)

		if workspace.IsCmdInterrupt() && lErr == nil {
			lErr = data.NewError(0, "list is interrupted")
		}

		if lErr != nil || workspace.IsCmdInterrupt() {
			errorHandler(info.Marker, lErr)

			if workspace.IsCmdInterrupt() || // 取消
				strings.Contains(lErr.Error(), "no such bucket") || // 空间不存在，直接结束
				strings.Contains(lErr.Error(), "incorrect zone") || // 空间不正确
				strings.Contains(lErr.Error(), "invalid list limit") || //  api v1 list limit
				strings.Contains(lErr.Error(), "context canceled") { // 取消
				break
			}

			retryCount++
			time.Sleep(time.Millisecond * 100)
			continue
		}

		if !hasMore || complete || workspace.IsCmdInterrupt() {
			break
		}

		retryCount = 0
	}

	if len(info.Marker) == 0 {
		if rErr := cache.removeCache(); rErr != nil {
			log.ErrorF("list remove cache status error: %v", rErr)
		} else {
			log.InfoF("list success, remove cache status: %s", cache.cachePath)
		}
	} else {
		log.InfoF("Marker: %s", info.Marker)
	}

	log.Debug("list bucket end")

	listWaiter.Done()
}

type ListToFileApiInfo struct {
	ListApiInfo
	FilePath   string // file 不存在则输出到 stdout
	AppendMode bool
	Readable   bool
}

func ListToFile(info ListToFileApiInfo, errorHandler func(marker string, err *data.CodeError)) {
	if errorHandler == nil {
		errorHandler = func(marker string, err *data.CodeError) {
			log.ErrorF("marker: %s", marker)
			log.ErrorF("list bucket Error: %v", err)
		}
		log.Warning("list bucket to file: not set error handler")
	}

	if len(info.ShowFields) == 0 {
		info.ShowFields = listObjectFields
	}

	if len(info.OutputFieldsSep) == 0 {
		info.OutputFieldsSep = data.DefaultLineSeparate
	}

	// 文件头
	title := strings.Join(info.ShowFields, info.OutputFieldsSep)

	var output io.WriteCloser
	if info.FilePath == "" {
		output = data.Stdout()
		_, _ = output.Write([]byte(title + "\n"))
		log.Debug("prepare list bucket to stdout")
	} else {
		var nErr *data.CodeError
		output, nErr = file.NewRotateFile(info.FilePath,
			file.RotateOptionMaxSize(info.OutputFileMaxSize),
			file.RotateOptionMaxLine(info.OutputFileMaxLines),
			file.RotateOptionAppendMode(info.AppendMode),
			file.RotateOptionFileHeader(title),
			file.RotateOptionOnOpenFile(func(filename string) {
				log.InfoF("open output file and prepare to write:%v", filename)
			}))

		if nErr != nil {
			errorHandler("", data.NewEmptyError().AppendDescF("failed to create rotate file:`%s`, error:%v", info.FilePath, nErr))
			return
		}
		defer output.Close()
		log.Debug("prepare list bucket to file")
	}

	bWriter := bufio.NewWriter(output)
	lineCreator := &ListLineCreator{
		Fields:   info.ShowFields,
		Sep:      info.OutputFieldsSep,
		Readable: info.Readable,
	}
	List(info.ListApiInfo, func(marker string, object ListObject) (bool, *data.CodeError) {
		lineData := lineCreator.Create(&object)
		if _, wErr := bWriter.WriteString(lineData + "\n"); wErr != nil {
			return false, data.NewEmptyError().AppendDesc("write error:" + wErr.Error())
		}

		if fErr := bWriter.Flush(); fErr != nil {
			return false, data.NewEmptyError().AppendDesc("flush error:" + fErr.Error())
		}
		return true, nil
	}, errorHandler)
}

func filterByPutTime(putTime, startDate, endDate time.Time) bool {
	switch {
	case startDate.IsZero() && endDate.IsZero():
		return true
	case !startDate.IsZero() && endDate.IsZero() && putTime.After(startDate):
		return true
	case !endDate.IsZero() && startDate.IsZero() && putTime.Before(endDate):
		return true
	case putTime.After(startDate) && putTime.Before(endDate):
		return true
	default:
		return false
	}
}

func filterBySuffixes(key string, suffixes []string) bool {
	hasSuffix := false
	if len(suffixes) == 0 {
		hasSuffix = true
	}
	for _, s := range suffixes {
		if strings.HasSuffix(key, s) {
			hasSuffix = true
			break
		}
	}
	return hasSuffix
}

func filterByStorageType(storageType int, storageTypes []int) bool {
	hasStorageType := false
	if len(storageTypes) == 0 {
		hasStorageType = true
	}
	for _, s := range storageTypes {
		if storageType == s {
			hasStorageType = true
			break
		}
	}
	return hasStorageType
}

func filterByMimeType(mimeType string, mimeTypes []string) bool {
	hasMimeType := false
	if len(mimeTypes) == 0 {
		hasMimeType = true
	}
	for _, s := range mimeTypes {
		if strings.Contains(s, "*") {
			sp := strings.ReplaceAll(s, "*", "")
			if strings.Contains(mimeType, sp) {
				hasMimeType = true
				break
			}
		} else if mimeType == s {
			hasMimeType = true
			break
		}
	}
	return hasMimeType
}

func filterByFileSize(fileSize, minSize, maxSize int64) bool {
	if maxSize < 0 {
		maxSize = math.MaxInt64
	}
	return fileSize >= minSize && fileSize <= maxSize
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/unrolled/render"
	"net/http"
	"reflect"
)

const (
	UserIdFieldName      = "UserId__"
	AdminUserIdFieldName = "AdminUserId__"
	ResIdFieldName       = "ResourceId__"
	ContentFieldName     = "Content__"
	RequestFieldName     = "Request__"
)

const (
	_ HandlerType = iota
	HandlerTypeJson
	HandlerTypeHtml
	HandlerTypeFeed
)

var decoder = schema.NewDecoder()

type HandlerType int

type httpFunc struct {
	requireAuth    bool
	requireAdmin   bool
	hasResId       bool
	hasContent     bool
	hasRequest     bool
	dummyInput     bool
	allowAnonymous bool
	inputType      reflect.Type
	contentType    reflect.Type
	funcValue      reflect.Value
}

type handler struct {
	htype    HandlerType
	path     string
	template string
	funcs    map[string]*httpFunc
	supports []string
	ts       TokenStore
	renderer *render.Render
}

func init() {
	decoder.IgnoreUnknownKeys(true)
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	f := h.funcs[r.Method]
	var input reflect.Value
	if f.dummyInput {
		input = reflect.ValueOf((*appgo.DummyInput)(nil))
	} else {
		input = reflect.New(f.inputType)
		if err := decoder.Decode(input.Interface(), r.URL.Query()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
	}
	if f.requireAuth {
		user, _ := h.authByHeader(r)
		s := input.Elem()
		field := s.FieldByName(UserIdFieldName)
		if user == 0 {
			if f.allowAnonymous {
				field.SetInt(appgo.AnonymousId)
			} else {
				h.renderError(w, appgo.NewApiErr(
					appgo.ECodeUnauthorized,
					"either remove UserId__ in your input define, or add allowAnonymous tag",
				))
				return
			}
		} else {
			field.SetInt(int64(user))
		}
	} else if f.requireAdmin {
		user, role := h.authByHeader(r)
		s := input.Elem()
		f := s.FieldByName(AdminUserIdFieldName)
		if user == 0 || role != appgo.RoleWebAdmin {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeUnauthorized,
				"admin role required, you could remove AdminUserId__ in your input define"))
			return
		}
		f.SetInt(int64(user))
	}
	if f.hasResId {
		vars := mux.Vars(r)
		id := appgo.IdFromStr(vars["id"])
		if id == 0 {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeNotFound,
				"ResourceId ('{id}' in url) required, you could remove ResourceId__ in your input define"))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ResIdFieldName)
		f.SetInt(int64(id))
	}
	if f.hasContent {
		content := reflect.New(f.contentType.Elem())
		if err := json.NewDecoder(r.Body).Decode(content.Interface()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ContentFieldName)
		f.Set(content)
	}
	if f.hasRequest {
		s := input.Elem()
		f := s.FieldByName(RequestFieldName)
		f.Set(reflect.ValueOf(r))
	}
	argsIn := []reflect.Value{input}
	returns := f.funcValue.Call(argsIn)
	if len(returns) == 0 || len(returns) > 2 {
		h.renderError(w, appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format"))
		return
	}
	// Either returns (reply, error) or returns (error)
	var retErr reflect.Value
	if len(returns) == 1 {
		retErr = returns[0]
	} else {
		retErr = returns[1]
	}
	// First check if err is nil
	if retErr.IsNil() {
		if len(returns) == 2 {
			h.renderData(w, returns[0].Interface())
		} else { // Empty return
			h.renderData(w, map[string]string{})
		}
	} else {
		if aerr, ok := retErr.Interface().(*appgo.ApiError); !ok {
			aerr = appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format")
		} else {
			h.renderError(w, aerr)
		}
	}
}

func (h *handler) authByHeader(r *http.Request) (appgo.Id, appgo.Role) {
	token := auth.Token(r.Header.Get(appgo.CustomTokenHeaderName))
	user, role := token.Validate()
	if user == 0 {
		return 0, 0
	}
	if !h.ts.Validate(token) {
		return 0, 0
	}
	return user, role
}

func newHandler(funcSet interface{}, htype HandlerType,
	ts TokenStore, renderer *render.Render) *handler {
	funcs := make(map[string]*httpFunc)
	// Let if panic if funSet's type is not right
	path := ""
	template := ""
	t := reflect.TypeOf(funcSet).Elem()
	if field, ok := t.FieldByName("META"); !ok {
		log.Panicln("Bad META setting (path, template)")
	} else {
		if p := field.Tag.Get("path"); p == "" {
			log.Panicln("Empty API path")
		} else {
			path = p
		}
		if htype == HandlerTypeHtml {
			if t := field.Tag.Get("template"); t == "" {
				log.Panicln("Empty HTML template")
			} else {
				template = t
			}
		}
	}
	structVal := reflect.Indirect(reflect.ValueOf(funcSet))
	supports := make([]string, 0, 4)
	if htype == HandlerTypeJson {
		methods := []string{"GET", "POST", "PUT", "DELETE"}
		for _, m := range methods {
			if fun, err := newHttpFunc(structVal, m); err != nil {
				log.Panicln(err)
			} else if fun != nil {
				funcs[m] = fun
				supports = append(supports, m)
			}
		}
		if len(supports) == 0 {
			log.Panicln("API supports no HTTP method")
		}
	} else if htype == HandlerTypeHtml {
		if fun, err := newHttpFunc(structVal, "HTML"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No HTML function for html")
		} else {
			funcs["GET"] = fun
		}
	} else if htype == HandlerTypeFeed {
		if fun, err := newHttpFunc(structVal, "FEED"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No FEED function for feed")
		} else {
			funcs["GET"] = fun
		}
	} else {
		log.Panicln("Bad handler type")
	}
	return &handler{htype, path, template, funcs, supports, ts, renderer}
}

func newHttpFunc(structVal reflect.Value, fieldName string) (*httpFunc, error) {
	fieldVal := structVal.MethodByName(fieldName)
	if !fieldVal.IsValid() {
		return nil, nil
	}
	ftype := fieldVal.Type()
	inNum := ftype.NumIn()
	if inNum != 1 {
		return nil, errors.New("API func needs to have exact 1 parameter")
	}
	inputType := ftype.In(0)
	dummyInput := false
	if inputType.Kind() != reflect.Ptr {
		return nil, errors.New("API func's parameter needs to be a pointer")
	}
	if inputType == reflect.TypeOf((*appgo.DummyInput)(nil)) {
		dummyInput = true
	}
	inputType = inputType.Elem()
	requireAuth := false
	allowAnonymous := false
	if fromIdField, ok := inputType.FieldByName(UserIdFieldName); ok {
		requireAuth = true
		if fromIdField.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
		aa := fromIdField.Tag.Get("allowAnonymous")
		allowAnonymous = (aa == "true")
	}
	requireAdmin := false
	if fromIdType, ok := inputType.FieldByName(AdminUserIdFieldName); ok {
		requireAdmin = true
		if fromIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
	}
	hasResId := false
	if resIdType, ok := inputType.FieldByName(ResIdFieldName); ok {
		hasResId = true
		if resIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("ResId needs to be Int64")
		}
	}
	hasContent := false
	var contentType reflect.Type
	if ctype, ok := inputType.FieldByName(ContentFieldName); ok {
		hasContent = true
		contentType = ctype.Type
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Content needs to be a pointer")
		}
	}
	hasRequest := false
	if ctype, ok := inputType.FieldByName(RequestFieldName); ok {
		hasRequest = true
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
		if ctype.Type.Elem() != reflect.TypeOf((*http.Request)(nil)).Elem() {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
	}
	return &httpFunc{requireAuth, requireAdmin, hasResId, hasContent, hasRequest,
		dummyInput, allowAnonymous, inputType, contentType, fieldVal}, nil
}
<commit_msg>support custom template<commit_after>package server

import (
	"encoding/json"
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	"github.com/gorilla/schema"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/unrolled/render"
	"net/http"
	"reflect"
)

const (
	UserIdFieldName      = "UserId__"
	AdminUserIdFieldName = "AdminUserId__"
	ResIdFieldName       = "ResourceId__"
	ContentFieldName     = "Content__"
	RequestFieldName     = "Request__"
)

const (
	_ HandlerType = iota
	HandlerTypeJson
	HandlerTypeHtml
	HandlerTypeFeed
)

var decoder = schema.NewDecoder()

type HandlerType int

type httpFunc struct {
	requireAuth    bool
	requireAdmin   bool
	hasResId       bool
	hasContent     bool
	hasRequest     bool
	dummyInput     bool
	allowAnonymous bool
	inputType      reflect.Type
	contentType    reflect.Type
	funcValue      reflect.Value
}

type handler struct {
	htype    HandlerType
	path     string
	template string
	funcs    map[string]*httpFunc
	supports []string
	ts       TokenStore
	renderer *render.Render
}

func init() {
	decoder.IgnoreUnknownKeys(true)
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	f := h.funcs[r.Method]
	var input reflect.Value
	if f.dummyInput {
		input = reflect.ValueOf((*appgo.DummyInput)(nil))
	} else {
		input = reflect.New(f.inputType)
		if err := decoder.Decode(input.Interface(), r.URL.Query()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
	}
	if f.requireAuth {
		user, _ := h.authByHeader(r)
		s := input.Elem()
		field := s.FieldByName(UserIdFieldName)
		if user == 0 {
			if f.allowAnonymous {
				field.SetInt(appgo.AnonymousId)
			} else {
				h.renderError(w, appgo.NewApiErr(
					appgo.ECodeUnauthorized,
					"either remove UserId__ in your input define, or add allowAnonymous tag",
				))
				return
			}
		} else {
			field.SetInt(int64(user))
		}
	} else if f.requireAdmin {
		user, role := h.authByHeader(r)
		s := input.Elem()
		f := s.FieldByName(AdminUserIdFieldName)
		if user == 0 || role != appgo.RoleWebAdmin {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeUnauthorized,
				"admin role required, you could remove AdminUserId__ in your input define"))
			return
		}
		f.SetInt(int64(user))
	}
	if f.hasResId {
		vars := mux.Vars(r)
		id := appgo.IdFromStr(vars["id"])
		if id == 0 {
			h.renderError(w, appgo.NewApiErr(
				appgo.ECodeNotFound,
				"ResourceId ('{id}' in url) required, you could remove ResourceId__ in your input define"))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ResIdFieldName)
		f.SetInt(int64(id))
	}
	if f.hasContent {
		content := reflect.New(f.contentType.Elem())
		if err := json.NewDecoder(r.Body).Decode(content.Interface()); err != nil {
			h.renderError(w, appgo.NewApiErr(appgo.ECodeBadRequest, err.Error()))
			return
		}
		s := input.Elem()
		f := s.FieldByName(ContentFieldName)
		f.Set(content)
	}
	if f.hasRequest {
		s := input.Elem()
		f := s.FieldByName(RequestFieldName)
		f.Set(reflect.ValueOf(r))
	}
	argsIn := []reflect.Value{input}
	returns := f.funcValue.Call(argsIn)
	rl := len(returns)
	if !(rl == 1 || rl == 2 || (rl == 3 && h.htype == HandlerTypeHtml)) {
		h.renderError(w, appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format"))
		return
	}
	// returns (reply, template-name, error) or (reply, error) or returns (error)
	retErr := returns[rl-1]
	// First check if err is nil
	if retErr.IsNil() {
		if rl == 3 {
			template := returns[1].Interface().(string)
			h.renderHtml(w, template, returns[0].Interface())
		} else if rl == 2 {
			h.renderData(w, returns[0].Interface())
		} else { // Empty return
			h.renderData(w, map[string]string{})
		}
	} else {
		if aerr, ok := retErr.Interface().(*appgo.ApiError); !ok {
			aerr = appgo.NewApiErr(appgo.ECodeInternal, "Bad api-func format")
		} else {
			h.renderError(w, aerr)
		}
	}
}

func (h *handler) authByHeader(r *http.Request) (appgo.Id, appgo.Role) {
	token := auth.Token(r.Header.Get(appgo.CustomTokenHeaderName))
	user, role := token.Validate()
	if user == 0 {
		return 0, 0
	}
	if !h.ts.Validate(token) {
		return 0, 0
	}
	return user, role
}

func newHandler(funcSet interface{}, htype HandlerType,
	ts TokenStore, renderer *render.Render) *handler {
	funcs := make(map[string]*httpFunc)
	// Let if panic if funSet's type is not right
	path := ""
	template := ""
	t := reflect.TypeOf(funcSet).Elem()
	if field, ok := t.FieldByName("META"); !ok {
		log.Panicln("Bad META setting (path, template)")
	} else {
		if p := field.Tag.Get("path"); p == "" {
			log.Panicln("Empty API path")
		} else {
			path = p
		}
		if htype == HandlerTypeHtml {
			t := field.Tag.Get("template")
			template = t
		}
	}
	structVal := reflect.Indirect(reflect.ValueOf(funcSet))
	supports := make([]string, 0, 4)
	if htype == HandlerTypeJson {
		methods := []string{"GET", "POST", "PUT", "DELETE"}
		for _, m := range methods {
			if fun, err := newHttpFunc(structVal, m); err != nil {
				log.Panicln(err)
			} else if fun != nil {
				funcs[m] = fun
				supports = append(supports, m)
			}
		}
		if len(supports) == 0 {
			log.Panicln("API supports no HTTP method")
		}
	} else if htype == HandlerTypeHtml {
		if fun, err := newHttpFunc(structVal, "HTML"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No HTML function for html")
		} else {
			funcs["GET"] = fun
		}
	} else if htype == HandlerTypeFeed {
		if fun, err := newHttpFunc(structVal, "FEED"); err != nil {
			log.Panicln(err)
		} else if fun == nil {
			log.Panicln("No FEED function for feed")
		} else {
			funcs["GET"] = fun
		}
	} else {
		log.Panicln("Bad handler type")
	}
	return &handler{htype, path, template, funcs, supports, ts, renderer}
}

func newHttpFunc(structVal reflect.Value, fieldName string) (*httpFunc, error) {
	fieldVal := structVal.MethodByName(fieldName)
	if !fieldVal.IsValid() {
		return nil, nil
	}
	ftype := fieldVal.Type()
	inNum := ftype.NumIn()
	if inNum != 1 {
		return nil, errors.New("API func needs to have exact 1 parameter")
	}
	inputType := ftype.In(0)
	dummyInput := false
	if inputType.Kind() != reflect.Ptr {
		return nil, errors.New("API func's parameter needs to be a pointer")
	}
	if inputType == reflect.TypeOf((*appgo.DummyInput)(nil)) {
		dummyInput = true
	}
	inputType = inputType.Elem()
	requireAuth := false
	allowAnonymous := false
	if fromIdField, ok := inputType.FieldByName(UserIdFieldName); ok {
		requireAuth = true
		if fromIdField.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
		aa := fromIdField.Tag.Get("allowAnonymous")
		allowAnonymous = (aa == "true")
	}
	requireAdmin := false
	if fromIdType, ok := inputType.FieldByName(AdminUserIdFieldName); ok {
		requireAdmin = true
		if fromIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("API func's 2nd parameter needs to be Int64")
		}
	}
	hasResId := false
	if resIdType, ok := inputType.FieldByName(ResIdFieldName); ok {
		hasResId = true
		if resIdType.Type.Kind() != reflect.Int64 {
			return nil, errors.New("ResId needs to be Int64")
		}
	}
	hasContent := false
	var contentType reflect.Type
	if ctype, ok := inputType.FieldByName(ContentFieldName); ok {
		hasContent = true
		contentType = ctype.Type
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Content needs to be a pointer")
		}
	}
	hasRequest := false
	if ctype, ok := inputType.FieldByName(RequestFieldName); ok {
		hasRequest = true
		if ctype.Type.Kind() != reflect.Ptr {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
		if ctype.Type.Elem() != reflect.TypeOf((*http.Request)(nil)).Elem() {
			return nil, errors.New("Request needs to be a pointer to http.Request")
		}
	}
	return &httpFunc{requireAuth, requireAdmin, hasResId, hasContent, hasRequest,
		dummyInput, allowAnonymous, inputType, contentType, fieldVal}, nil
}
<|endoftext|>