<commit_before>package generator

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"io"
	"reflect"
	"strings"
)

// Used for marking fields as not written. Not actually returned to user.
var unsupportedTypeError = errors.New("Unsupported type")
var ignoreFieldError = errors.New("Ignore field")

// GenerateTypeEncoder generates a struct for a single Type.
func writeTypeEncoder(w io.Writer, typeSpec *ast.TypeSpec) error {
	var b bytes.Buffer
	name := typeSpec.Name.Name

	// Ignore type non-struct specs.
	structType, ok := typeSpec.Type.(*ast.StructType)
	if !ok {
		return nil
	}

	// Generate the encoder.
	fmt.Fprintf(&b, "type %sJSONEncoder struct {", name)
	fmt.Fprintln(&b, "w io.Writer")
	fmt.Fprintln(&b, "}")

	// Generate the constructor.
	fmt.Fprintf(&b, "func New%sJSONEncoder(w io.Writer) *%sJSONEncoder {\n", name, name)
	fmt.Fprintf(&b, "return &%sJSONEncoder{w: w}\n", name)
	fmt.Fprintln(&b, "}")
	fmt.Fprintln(&b, "")

	// Generate the encode function.
	fmt.Fprintf(&b, "func (e *%sJSONEncoder) Encode(v *%s) error {\n", name, name)
	fmt.Fprintf(&b, "if v == nil {\nreturn encoding.WriteBytes(e.w, []byte(`null`))\n}\n\n")

	fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, '{'); err != nil {\nreturn err\n}\n")

	index := 0
	for _, field := range structType.Fields.List {
		for _, name := range field.Names {
			// Write to a temporary buffer to check if anything is written.
			var buf bytes.Buffer
			if err := writeFieldEncoding(&buf, name.Name, field); buf.Len() > 0 {
				// Write separating comma after the first field.
				if index > 0 {
					fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, ','); err != nil {\nreturn err\n}\n")
				}

				// Copy over to main buffer.
				buf.WriteTo(&b)
				index++

			} else if err != nil && err != unsupportedTypeError && err != ignoreFieldError {
				return err
			}
		}
	}

	fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, '}'); err != nil {\nreturn err\n}\n")
	fmt.Fprintf(&b, "return nil\n")
	fmt.Fprintf(&b, "}\n")

	// Copy buffer to writer.
	if _, err := b.WriteTo(w); err != nil {
		return err
	}

	return nil
}

// writeFieldEncoding generates the encoder code for a single field.
func writeFieldEncoding(w io.Writer, name string, f *ast.Field) error {
	var tag string
	if f.Tag != nil {
		tag = f.Tag.Value[1:len(f.Tag.Value)-1]
		tag = reflect.StructTag(tag).Get("json")
	}
	tags := strings.Split(tag, ",")

	// Extract json key name.
	var key = name
	if len(tags) > 0 {
		if tags[0] == "-" {
			return ignoreFieldError
		} else if len(tags[0]) > 0 {
			key = tags[0]
		}
	}

	// Primative and complex types need separate handling.
	switch f.Type.(type) {
	case *ast.Ident:
		return writePrimativeFieldEncoding(w, name, key, f, tag)
	case *ast.StarExpr:
		return writePointerFieldEncoding(w, name, key, f, tag)
	case *ast.ArrayType:
		fmt.Println("ARRAY!")
	}
	return unsupportedTypeError
}

// writeFieldEncoding generates the encoder code for a single primative field.
func writePrimativeFieldEncoding(w io.Writer, name string, key string, f *ast.Field, tag string) error {
	typ := f.Type.(*ast.Ident)

	fmt.Fprintf(w, "if err := encoding.WriteString(e.w, \"%s\"); err != nil {\nreturn err\n}\n", key)
	fmt.Fprintf(w, "if err := encoding.WriteByte(e.w, ':'); err != nil {\nreturn err\n}\n")

	switch typ.Name {
	case "string":
		fmt.Fprintf(w, "if err := encoding.WriteString(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "int":
		fmt.Fprintf(w, "if err := encoding.WriteInt(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "int64":
		fmt.Fprintf(w, "if err := encoding.WriteInt64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "uint":
		fmt.Fprintf(w, "if err := encoding.WriteUint(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "uint64":
		fmt.Fprintf(w, "if err := encoding.WriteUint64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "float32":
		fmt.Fprintf(w, "if err := encoding.WriteFloat32(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "float64":
		fmt.Fprintf(w, "if err := encoding.WriteFloat64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "bool":
		fmt.Fprintf(w, "if err := encoding.WriteBool(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	default:
		return unsupportedTypeError
	}

	return nil
}

// writePointerFieldEncoding generates the encoding code for a single field with a pointer type.
func writePointerFieldEncoding(w io.Writer, name string, key string, f *ast.Field, tag string) error {
	typ := f.Type.(*ast.StarExpr)
	x, ok := typ.X.(*ast.Ident)
	if !ok {
		return unsupportedTypeError
	}

	fmt.Fprintf(w, "if err := encoding.WriteString(e.w, \"%s\"); err != nil {\nreturn err\n}\n", name)
	fmt.Fprintf(w, "if err := encoding.WriteByte(e.w, ':'); err != nil {\nreturn err\n}\n")
	fmt.Fprintf(w, "if err := New%sJSONEncoder(e.w).Encode(v.%s); err != nil {\nreturn err\n}\n", x.Name, name)
	return nil
}
<commit_msg>Fix key quoting.<commit_after>package generator

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"io"
	"reflect"
	"strconv"
	"strings"
)

// Used for marking fields as not written. Not actually returned to user.
var unsupportedTypeError = errors.New("Unsupported type")
var ignoreFieldError = errors.New("Ignore field")

// GenerateTypeEncoder generates a struct for a single Type.
func writeTypeEncoder(w io.Writer, typeSpec *ast.TypeSpec) error {
	var b bytes.Buffer
	name := typeSpec.Name.Name

	// Ignore type non-struct specs.
	structType, ok := typeSpec.Type.(*ast.StructType)
	if !ok {
		return nil
	}

	// Generate the encoder.
	fmt.Fprintf(&b, "type %sJSONEncoder struct {", name)
	fmt.Fprintln(&b, "w io.Writer")
	fmt.Fprintln(&b, "}")

	// Generate the constructor.
	fmt.Fprintf(&b, "func New%sJSONEncoder(w io.Writer) *%sJSONEncoder {\n", name, name)
	fmt.Fprintf(&b, "return &%sJSONEncoder{w: w}\n", name)
	fmt.Fprintln(&b, "}")
	fmt.Fprintln(&b, "")

	// Generate the encode function.
	fmt.Fprintf(&b, "func (e *%sJSONEncoder) Encode(v *%s) error {\n", name, name)
	fmt.Fprintf(&b, "if v == nil {\nreturn encoding.WriteBytes(e.w, []byte(`null`))\n}\n\n")

	fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, '{'); err != nil {\nreturn err\n}\n")

	index := 0
	for _, field := range structType.Fields.List {
		for _, name := range field.Names {
			// Write to a temporary buffer to check if anything is written.
			var buf bytes.Buffer
			if err := writeFieldEncoding(&buf, name.Name, field); buf.Len() > 0 {
				// Write separating comma after the first field.
				if index > 0 {
					fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, ','); err != nil {\nreturn err\n}\n")
				}

				// Copy over to main buffer.
				buf.WriteTo(&b)
				index++

			} else if err != nil && err != unsupportedTypeError && err != ignoreFieldError {
				return err
			}
		}
	}

	fmt.Fprintf(&b, "if err := encoding.WriteByte(e.w, '}'); err != nil {\nreturn err\n}\n")
	fmt.Fprintf(&b, "return nil\n")
	fmt.Fprintf(&b, "}\n")

	// Copy buffer to writer.
	if _, err := b.WriteTo(w); err != nil {
		return err
	}

	return nil
}

// writeFieldEncoding generates the encoder code for a single field.
func writeFieldEncoding(w io.Writer, name string, f *ast.Field) error {
	var tag string
	if f.Tag != nil {
		tag = f.Tag.Value[1:len(f.Tag.Value)-1]
		tag = reflect.StructTag(tag).Get("json")
	}
	tags := strings.Split(tag, ",")

	// Extract json key name.
	var key = name
	if len(tags) > 0 {
		if tags[0] == "-" {
			return ignoreFieldError
		} else if len(tags[0]) > 0 {
			key = tags[0]
		}
	}

	// Primative and complex types need separate handling.
	switch f.Type.(type) {
	case *ast.Ident:
		return writePrimativeFieldEncoding(w, name, key, f, tag)
	case *ast.StarExpr:
		return writePointerFieldEncoding(w, name, key, f, tag)
	case *ast.ArrayType:
		fmt.Println("ARRAY!")
	}
	return unsupportedTypeError
}

// writeFieldEncoding generates the encoder code for a single primative field.
func writePrimativeFieldEncoding(w io.Writer, name string, key string, f *ast.Field, tag string) error {
	typ := f.Type.(*ast.Ident)

	fmt.Fprintf(w, "if err := encoding.WriteString(e.w, %s); err != nil {\nreturn err\n}\n", strconv.Quote(key))
	fmt.Fprintf(w, "if err := encoding.WriteByte(e.w, ':'); err != nil {\nreturn err\n}\n")

	switch typ.Name {
	case "string":
		fmt.Fprintf(w, "if err := encoding.WriteString(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "int":
		fmt.Fprintf(w, "if err := encoding.WriteInt(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "int64":
		fmt.Fprintf(w, "if err := encoding.WriteInt64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "uint":
		fmt.Fprintf(w, "if err := encoding.WriteUint(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "uint64":
		fmt.Fprintf(w, "if err := encoding.WriteUint64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "float32":
		fmt.Fprintf(w, "if err := encoding.WriteFloat32(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "float64":
		fmt.Fprintf(w, "if err := encoding.WriteFloat64(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	case "bool":
		fmt.Fprintf(w, "if err := encoding.WriteBool(e.w, v.%s); err != nil {\nreturn err\n}\n", name)
	default:
		return unsupportedTypeError
	}

	return nil
}

// writePointerFieldEncoding generates the encoding code for a single field with a pointer type.
func writePointerFieldEncoding(w io.Writer, name string, key string, f *ast.Field, tag string) error {
	typ := f.Type.(*ast.StarExpr)
	x, ok := typ.X.(*ast.Ident)
	if !ok {
		return unsupportedTypeError
	}

	fmt.Fprintf(w, "if err := encoding.WriteString(e.w, %s); err != nil {\nreturn err\n}\n", strconv.Quote(key))
	fmt.Fprintf(w, "if err := encoding.WriteByte(e.w, ':'); err != nil {\nreturn err\n}\n")
	fmt.Fprintf(w, "if err := New%sJSONEncoder(e.w).Encode(v.%s); err != nil {\nreturn err\n}\n", x.Name, name)
	return nil
}
<|endoftext|>


<commit_before>package build

import (
	"bufio"
	"bytes"
	"github.com/jingweno/gotask/tasking"
	"io"
	"regexp"
	"strings"
)

type manPage struct {
	Name        string
	Usage       string
	Description string
	Flags       []tasking.Flag
}

type manPageParser struct {
	Doc string
}

func (p *manPageParser) Parse() (mp *manPage, err error) {
	sections, err := p.readSections()
	if err != nil {
		return
	}

	mp = &manPage{}
	if opts, ok := sections["OPTIONS"]; ok {
		mp.Flags, err = p.parseFlags(opts)
		if err != nil {
			return
		}
	}
	if nameAndUsage, ok := sections["NAME"]; ok {
		mp.Name, mp.Usage = p.splitNameAndUsage(nameAndUsage)
	}
	mp.Description = sections["DESCRIPTION"]

	return
}

func (p *manPageParser) readSections() (sections map[string]string, err error) {
	sections = make(map[string]string)
	headingRegexp := regexp.MustCompile(`^([A-Z]+)$`)
	reader := bufio.NewReader(bytes.NewReader([]byte(p.Doc)))

	var (
		line    string
		heading string
		content []string
	)
	for err == nil {
		line, err = readLine(reader)

		if headingRegexp.MatchString(line) {
			if heading != line {
				if heading != "" {
					sections[heading] = concatHeadingContent(content)
				}

				heading = line
				content = []string{}
			}
		} else {
			if line != "" {
				line = strings.TrimSpace(line)
			}
			content = append(content, line)
		}
	}
	// the last one
	if heading != "" {
		sections[heading] = concatHeadingContent(content)
	}

	if err == io.EOF {
		err = nil
	}

	return
}

func (p *manPageParser) splitNameAndUsage(nameAndUsage string) (name, usage string) {
	s := strings.SplitN(nameAndUsage, " - ", 2)
	if len(s) == 1 {
		name = ""
		usage = strings.TrimSpace(s[0])
	} else {
		name = strings.TrimSpace(s[0])
		usage = strings.TrimSpace(s[1])
	}

	return
}

func (p *manPageParser) parseFlags(optsStr string) (flags []tasking.Flag, err error) {
	reader := bufio.NewReader(bytes.NewReader([]byte(optsStr)))
	flagRegexp := regexp.MustCompile(`(\-?\-\w+,?)+`)
	var (
		line, name string
		content    []string
	)

	for err == nil {
		line, err = readLine(reader)
		if flagRegexp.MatchString(line) {
			if name != line {
				if name != "" {
					flags = append(flags, tasking.BoolFlag{Name: name, Usage: concatFlagContent(content)})
				}

				name = line
				content = []string{}
			}
		} else {
			if line != "" {
				line = strings.TrimSpace(line)
			}
			content = append(content, line)
		}
	}
	// the last one
	if name != "" {
		flags = append(flags, tasking.BoolFlag{Name: name, Usage: concatFlagContent(content)})
	}

	if err == io.EOF {
		err = nil
	}

	return
}

func concatHeadingContent(content []string) string {
	return strings.TrimSpace(strings.Join(content, "\n   "))
}

func concatFlagContent(content []string) string {
	return strings.TrimSpace(strings.Join(content, "\n"))
}

func readLine(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)

	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}
<commit_msg>Refactor method name<commit_after>package build

import (
	"bufio"
	"bytes"
	"github.com/jingweno/gotask/tasking"
	"io"
	"regexp"
	"strings"
)

type manPage struct {
	Name        string
	Usage       string
	Description string
	Flags       []tasking.Flag
}

type manPageParser struct {
	Doc string
}

func (p *manPageParser) Parse() (mp *manPage, err error) {
	sections, err := p.readSections()
	if err != nil {
		return
	}

	mp = &manPage{}
	if opts, ok := sections["OPTIONS"]; ok {
		mp.Flags, err = p.parseOptions(opts)
		if err != nil {
			return
		}
	}
	if nameAndUsage, ok := sections["NAME"]; ok {
		mp.Name, mp.Usage = p.splitNameAndUsage(nameAndUsage)
	}
	mp.Description = sections["DESCRIPTION"]

	return
}

func (p *manPageParser) readSections() (sections map[string]string, err error) {
	sections = make(map[string]string)
	headingRegexp := regexp.MustCompile(`^([A-Z]+)$`)
	reader := bufio.NewReader(bytes.NewReader([]byte(p.Doc)))

	var (
		line    string
		heading string
		content []string
	)
	for err == nil {
		line, err = readLine(reader)

		if headingRegexp.MatchString(line) {
			if heading != line {
				if heading != "" {
					sections[heading] = concatSectionContent(content)
				}

				heading = line
				content = []string{}
			}
		} else {
			if line != "" {
				line = strings.TrimSpace(line)
			}
			content = append(content, line)
		}
	}
	// the last one
	if heading != "" {
		sections[heading] = concatSectionContent(content)
	}

	if err == io.EOF {
		err = nil
	}

	return
}

func (p *manPageParser) splitNameAndUsage(nameAndUsage string) (name, usage string) {
	s := strings.SplitN(nameAndUsage, " - ", 2)
	if len(s) == 1 {
		name = ""
		usage = strings.TrimSpace(s[0])
	} else {
		name = strings.TrimSpace(s[0])
		usage = strings.TrimSpace(s[1])
	}

	return
}

func (p *manPageParser) parseOptions(optsStr string) (flags []tasking.Flag, err error) {
	reader := bufio.NewReader(bytes.NewReader([]byte(optsStr)))
	flagRegexp := regexp.MustCompile(`(\-?\-\w+,?)+`)

	var (
		line, name string
		content    []string
	)
	for err == nil {
		line, err = readLine(reader)
		if flagRegexp.MatchString(line) {
			if name != line {
				if name != "" {
					flags = append(flags, tasking.BoolFlag{Name: name, Usage: concatFlagContent(content)})
				}

				name = line
				content = []string{}
			}
		} else {
			if line != "" {
				line = strings.TrimSpace(line)
			}
			content = append(content, line)
		}
	}
	// the last one
	if name != "" {
		flags = append(flags, tasking.BoolFlag{Name: name, Usage: concatFlagContent(content)})
	}

	if err == io.EOF {
		err = nil
	}

	return
}

func concatSectionContent(content []string) string {
	return strings.TrimSpace(strings.Join(content, "\n   "))
}

func concatFlagContent(content []string) string {
	return strings.TrimSpace(strings.Join(content, "\n"))
}

func readLine(r *bufio.Reader) (string, error) {
	var (
		isPrefix = true
		err      error
		line, ln []byte
	)

	for isPrefix && err == nil {
		line, isPrefix, err = r.ReadLine()
		ln = append(ln, line...)
	}

	return string(ln), err
}
<|endoftext|>


<commit_before>// Copyright 2017 The Serulian 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 packageloader

import (
	"strings"

	"github.com/serulian/compiler/compilercommon"

	"fmt"

	cmap "github.com/streamrail/concurrent-map"
)

// SourceTracker is a helper struct for tracking the contents and versioning of all source files
// encountered by the package loader during its loading phase.
type SourceTracker struct {
	// sourceFiles is the map of tracked source files.
	sourceFiles map[compilercommon.InputSource]trackedSourceFile

	// pathLoader is the path loader used when loading the source files.
	pathLoader PathLoader
}

// ModifiedSourcePaths returns the source paths that have been modified since originally read,
// if any.
func (st SourceTracker) ModifiedSourcePaths() ([]compilercommon.InputSource, error) {
	var modifiedPaths = make([]compilercommon.InputSource, 0, len(st.sourceFiles))
	for _, file := range st.sourceFiles {
		current, err := st.pathLoader.GetRevisionID(string(file.sourcePath))
		if err != nil {
			return modifiedPaths, err
		}

		if current != file.revisionID {
			modifiedPaths = append(modifiedPaths, file.sourcePath)
		}
	}
	return modifiedPaths, nil
}

// LoadedContents returns the contents of the given source path when loaded by the packageloader, if any.
func (st SourceTracker) LoadedContents(path compilercommon.InputSource) ([]byte, bool) {
	tracked, exists := st.sourceFiles[path]
	if !exists {
		return []byte{}, false
	}

	return tracked.contents, true
}

func (st SourceTracker) RunePositionToLineAndCol(runePosition int, path compilercommon.InputSource) (int, int, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return 0, 0, fmt.Errorf("Could not find path %s", path)
	}

	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.RunePositionToLineAndCol(runePosition)
}

func (st SourceTracker) LineAndColToRunePosition(lineNumber int, colPosition int, path compilercommon.InputSource) (int, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return 0, fmt.Errorf("Could not find path %s", path)
	}

	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.LineAndColToRunePosition(lineNumber, colPosition)
}

func (st SourceTracker) TextForLine(lineNumber int, path compilercommon.InputSource) (string, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return "", fmt.Errorf("Could not find path %s", path)
	}

	lines := strings.Split(string(tsf.contents), "\n")
	if int(lineNumber) >= len(lines) {
		return "", fmt.Errorf("Line number %v not found in path %v", lineNumber, path)
	}

	return lines[int(lineNumber)], nil
}

// trackedSourceFile defines a struct for tracking the versioning and contents of a source file
// encountered during package loading.
type trackedSourceFile struct {
	// sourcePath is the path of the source file encountered.
	sourcePath compilercommon.InputSource

	// sourceKind is the kind of the source file.
	sourceKind string

	// contents are the contents of the source file when it was loaded.
	contents []byte

	// revisionID is the revision ID of the source file when it was loaded, as returned by
	// GetRevisionID on the PathLoader. Typically, this is an mtime or version number. If the
	// revision ID has changed, the contents of the file are assumed to have been altered since
	// the file was read.
	revisionID int64

	// positionMapper is the source position mapper for converting rune positions <-> line+col
	// positions. Only instantiated if necessary.
	positionMapper *compilercommon.SourcePositionMapper
}

// mutableSourceTracker is a source tracker that is built during the package loading phase.
type mutableSourceTracker struct {
	// sourceFiles is a map from source file path to a trackedSourceFile struct.
	sourceFiles cmap.ConcurrentMap

	// pathLoader is the path loader used when loading the source files.
	pathLoader PathLoader
}

func newMutableSourceTracker(pathLoader PathLoader) *mutableSourceTracker {
	return &mutableSourceTracker{
		sourceFiles: cmap.New(),
		pathLoader:  pathLoader,
	}
}

// AddSourceFile adds a source file to the mutable source tracker.
func (m *mutableSourceTracker) AddSourceFile(path compilercommon.InputSource, sourceKind string, contents []byte, revisionID int64) {
	m.sourceFiles.Set(string(path), trackedSourceFile{
		sourcePath: path,
		sourceKind: sourceKind,
		contents:   contents,
		revisionID: revisionID,
	})
}

func (m *mutableSourceTracker) RunePositionToLineAndCol(runePosition int, path compilercommon.InputSource) (int, int, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return 0, 0, fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)
	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.RunePositionToLineAndCol(runePosition)
}

func (m *mutableSourceTracker) LineAndColToRunePosition(lineNumber int, colPosition int, path compilercommon.InputSource) (int, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return 0, fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)
	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.LineAndColToRunePosition(lineNumber, colPosition)
}

func (m *mutableSourceTracker) TextForLine(lineNumber int, path compilercommon.InputSource) (string, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return "", fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)

	lines := strings.Split(string(tsf.contents), "\n")
	if int(lineNumber) >= len(lines) {
		return "", fmt.Errorf("Line number %v not found in path %v", lineNumber, path)
	}

	return lines[int(lineNumber)], nil
}

// Freeze freezes the mutable source tracker into an immutable SourceTracker.
func (m *mutableSourceTracker) Freeze() SourceTracker {
	sourceFileMap := map[compilercommon.InputSource]trackedSourceFile{}
	for entry := range m.sourceFiles.Iter() {
		sourceFileMap[compilercommon.InputSource(entry.Key)] = entry.Val.(trackedSourceFile)
	}
	return SourceTracker{sourceFileMap, m.pathLoader}
}
<commit_msg>Switch the tracked source file to a pointer<commit_after>// Copyright 2017 The Serulian 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 packageloader

import (
	"strings"

	"github.com/serulian/compiler/compilercommon"

	"fmt"

	cmap "github.com/streamrail/concurrent-map"
)

// SourceTracker is a helper struct for tracking the contents and versioning of all source files
// encountered by the package loader during its loading phase.
type SourceTracker struct {
	// sourceFiles is the map of tracked source files.
	sourceFiles map[compilercommon.InputSource]*trackedSourceFile

	// pathLoader is the path loader used when loading the source files.
	pathLoader PathLoader
}

// ModifiedSourcePaths returns the source paths that have been modified since originally read,
// if any.
func (st SourceTracker) ModifiedSourcePaths() ([]compilercommon.InputSource, error) {
	var modifiedPaths = make([]compilercommon.InputSource, 0, len(st.sourceFiles))
	for _, file := range st.sourceFiles {
		current, err := st.pathLoader.GetRevisionID(string(file.sourcePath))
		if err != nil {
			return modifiedPaths, err
		}

		if current != file.revisionID {
			modifiedPaths = append(modifiedPaths, file.sourcePath)
		}
	}
	return modifiedPaths, nil
}

// LoadedContents returns the contents of the given source path when loaded by the packageloader, if any.
func (st SourceTracker) LoadedContents(path compilercommon.InputSource) ([]byte, bool) {
	tracked, exists := st.sourceFiles[path]
	if !exists {
		return []byte{}, false
	}

	return tracked.contents, true
}

func (st SourceTracker) RunePositionToLineAndCol(runePosition int, path compilercommon.InputSource) (int, int, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return 0, 0, fmt.Errorf("Could not find path %s", path)
	}

	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.RunePositionToLineAndCol(runePosition)
}

func (st SourceTracker) LineAndColToRunePosition(lineNumber int, colPosition int, path compilercommon.InputSource) (int, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return 0, fmt.Errorf("Could not find path %s", path)
	}

	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.LineAndColToRunePosition(lineNumber, colPosition)
}

func (st SourceTracker) TextForLine(lineNumber int, path compilercommon.InputSource) (string, error) {
	tsf, exists := st.sourceFiles[path]
	if !exists {
		return "", fmt.Errorf("Could not find path %s", path)
	}

	lines := strings.Split(string(tsf.contents), "\n")
	if int(lineNumber) >= len(lines) {
		return "", fmt.Errorf("Line number %v not found in path %v", lineNumber, path)
	}

	return lines[int(lineNumber)], nil
}

// trackedSourceFile defines a struct for tracking the versioning and contents of a source file
// encountered during package loading.
type trackedSourceFile struct {
	// sourcePath is the path of the source file encountered.
	sourcePath compilercommon.InputSource

	// sourceKind is the kind of the source file.
	sourceKind string

	// contents are the contents of the source file when it was loaded.
	contents []byte

	// revisionID is the revision ID of the source file when it was loaded, as returned by
	// GetRevisionID on the PathLoader. Typically, this is an mtime or version number. If the
	// revision ID has changed, the contents of the file are assumed to have been altered since
	// the file was read.
	revisionID int64

	// positionMapper is the source position mapper for converting rune positions <-> line+col
	// positions. Only instantiated if necessary.
	positionMapper *compilercommon.SourcePositionMapper
}

// mutableSourceTracker is a source tracker that is built during the package loading phase.
type mutableSourceTracker struct {
	// sourceFiles is a map from source file path to a trackedSourceFile struct.
	sourceFiles cmap.ConcurrentMap

	// pathLoader is the path loader used when loading the source files.
	pathLoader PathLoader
}

func newMutableSourceTracker(pathLoader PathLoader) *mutableSourceTracker {
	return &mutableSourceTracker{
		sourceFiles: cmap.New(),
		pathLoader:  pathLoader,
	}
}

// AddSourceFile adds a source file to the mutable source tracker.
func (m *mutableSourceTracker) AddSourceFile(path compilercommon.InputSource, sourceKind string, contents []byte, revisionID int64) {
	m.sourceFiles.Set(string(path), &trackedSourceFile{
		sourcePath: path,
		sourceKind: sourceKind,
		contents:   contents,
		revisionID: revisionID,
	})
}

func (m *mutableSourceTracker) RunePositionToLineAndCol(runePosition int, path compilercommon.InputSource) (int, int, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return 0, 0, fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)
	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.RunePositionToLineAndCol(runePosition)
}

func (m *mutableSourceTracker) LineAndColToRunePosition(lineNumber int, colPosition int, path compilercommon.InputSource) (int, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return 0, fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)
	if tsf.positionMapper == nil {
		tsf.positionMapper = compilercommon.CreateSourcePositionMapper(tsf.contents)
	}

	return tsf.positionMapper.LineAndColToRunePosition(lineNumber, colPosition)
}

func (m *mutableSourceTracker) TextForLine(lineNumber int, path compilercommon.InputSource) (string, error) {
	tracked, exists := m.sourceFiles.Get(string(path))
	if !exists {
		return "", fmt.Errorf("Could not find path %s", path)
	}

	tsf := tracked.(trackedSourceFile)

	lines := strings.Split(string(tsf.contents), "\n")
	if int(lineNumber) >= len(lines) {
		return "", fmt.Errorf("Line number %v not found in path %v", lineNumber, path)
	}

	return lines[int(lineNumber)], nil
}

// Freeze freezes the mutable source tracker into an immutable SourceTracker.
func (m *mutableSourceTracker) Freeze() SourceTracker {
	sourceFileMap := map[compilercommon.InputSource]*trackedSourceFile{}
	for entry := range m.sourceFiles.Iter() {
		sourceFileMap[compilercommon.InputSource(entry.Key)] = entry.Val.(*trackedSourceFile)
	}
	return SourceTracker{sourceFileMap, m.pathLoader}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"net/http"
	"strconv"
	"text/template"
	"time"

	"github.com/pilu/traffic"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/user"
)

type ResponseData struct {
	Entry   *Entry
	IsAdmin bool
	Next    string
	Prev    string
}

func PostHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		log.Errorf(c, "Error parsing int: %+v", err)
		http.Redirect(w, r.Request, "/", 302)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		log.Errorf(c, "Error loading post: %+v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	} else {
		responseData := &ResponseData{
			Entry:   entry,
			IsAdmin: user.IsAdmin(c),
			Next:    entry.NextPost(c),
			Prev:    entry.PrevPost(c)}
		w.Render("post", responseData)
		return
	}
}

type DayData struct {
	Posts   *[]Entry
	IsAdmin bool
	Date    time.Time
}

func DayHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	day, err := strconv.ParseInt(r.Param("day"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	month, err := strconv.ParseInt(r.Param("month"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	year, err := strconv.ParseInt(r.Param("year"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	entries, err := PostsForDay(c, year, month, day)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	responseData := &DayData{
		Posts:   entries,
		IsAdmin: user.IsAdmin(c),
		Date:    time.Date(int(year), time.Month(month), int(day), 0, 0, 0, 0, time.UTC),
	}
	w.Render("day", responseData)
	return
}

func PostMarkdownHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	}

	// Define a template.
	const md = `---

id: {{.Id}}
datetime: "{{.Datetime}}"
title: "{{if .Title}}{{.Title}}{{else}}#{{.Id}}{{end}}"
draft: {{.Draft}}
permalink: "/posts/{{.Id}}"
{{if .Longform}}longform: {{.Longform}}
{{end}}
---

{{.Content}}
`

	// Create a new template and parse the letter into it.
	t := template.Must(template.New("post_md").Parse(md))

	buf := new(bytes.Buffer)
	err = t.Execute(buf, entry)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	w.WriteText(buf.String())
	return
}

func PostJsonHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		log.Errorf(c, "Error loading post: %+v", err)
		http.Error(w, err.Error(), 500)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteJSON(entry)
}
<commit_msg>Add support for html<commit_after>package main

import (
	"bytes"
	"errors"
	htmp "html/template"
	"net/http"
	"strconv"
	"text/template"
	"time"

	"github.com/pilu/traffic"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/user"
)

type ResponseData struct {
	Entry   *Entry
	IsAdmin bool
	Next    string
	Prev    string
}

func PostHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		log.Errorf(c, "Error parsing int: %+v", err)
		http.Redirect(w, r.Request, "/", 302)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		log.Errorf(c, "Error loading post: %+v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	} else {
		responseData := &ResponseData{
			Entry:   entry,
			IsAdmin: user.IsAdmin(c),
			Next:    entry.NextPost(c),
			Prev:    entry.PrevPost(c)}
		w.Render("post", responseData)
		return
	}
}

type DayData struct {
	Posts   *[]Entry
	IsAdmin bool
	Date    time.Time
}

func DayHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	day, err := strconv.ParseInt(r.Param("day"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	month, err := strconv.ParseInt(r.Param("month"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	year, err := strconv.ParseInt(r.Param("year"), 10, 32)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	entries, err := PostsForDay(c, year, month, day)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	responseData := &DayData{
		Posts:   entries,
		IsAdmin: user.IsAdmin(c),
		Date:    time.Date(int(year), time.Month(month), int(day), 0, 0, 0, 0, time.UTC),
	}
	w.Render("day", responseData)
	return
}

func PostMarkdownHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	}

	// Define a template.
	const md = `---

id: {{.Id}}
datetime: "{{.Datetime}}"
title: "{{if .Title}}{{.Title}}{{else}}#{{.Id}}{{end}}"
draft: {{.Draft}}
permalink: "/posts/{{.Id}}"
{{if .Longform}}longform: {{.Longform}}
{{end}}
---

{{.Content}}
`

	// Create a new template and parse the letter into it.
	t := template.Must(template.New("post_md").Parse(md))

	buf := new(bytes.Buffer)
	err = t.Execute(buf, entry)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	w.WriteText(buf.String())
	return
}

type EntryJson struct {
	Entry
	Html htmp.HTML `json:"html"`
}

func PostJsonHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.ParseInt(r.Param("id"), 10, 64)
	if err != nil {
		log.Errorf(c, "Error loading post: %+v", err)
		http.Error(w, err.Error(), 500)
		return
	}
	entry, err := GetEntry(c, id)
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	if entry.Draft && !user.IsAdmin(c) {
		http.Error(w, errors.New("Post is not public").Error(), 403)
		return
	}

	post := &EntryJson{
		Html: entry.Html(),
	}
	post.Id = entry.Id
	post.Title = entry.Title
	post.Content = entry.Content
	post.Datetime = entry.Datetime
	post.Created = entry.Created
	post.Modified = entry.Modified
	post.Tags = entry.Tags

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteJSON(post)
}
<|endoftext|>


<commit_before>package middlewares

import (
	"net/http"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/controllers/router"
	"github.com/NyaaPantsu/nyaa/templates"
	"github.com/NyaaPantsu/nyaa/utils/log"
	msg "github.com/NyaaPantsu/nyaa/utils/messages"
	"github.com/NyaaPantsu/nyaa/utils/oauth2"
	"github.com/gin-gonic/gin"
	"github.com/ory/fosite"
)

// ErrorMiddleware for managing errors on status
func ErrorMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Next()
		if c.Writer.Status() >= 300 && config.Get().Environment == "DEVELOPMENT" {
			messages := msg.GetMessages(c)
			if messages.HasErrors() {
				log.Errorf("Request has errors: %v", messages.GetAllErrors())
			}
		}
		if c.Writer.Status() != http.StatusOK && c.Writer.Size() <= 0 {
			if c.ContentType() == "application/json" {
				messages := msg.GetMessages(c)
				messages.AddErrorT("errors", "404_not_found")
				c.JSON(c.Writer.Status(), messages.GetAllErrors())
				return
			}
			templates.HttpError(c, c.Writer.Status())
		}
	}
}

// ModMiddleware Make sure the user is a moderator, otherwise return forbidden
func ModMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		currentUser := router.GetUser(c)
		if !currentUser.HasAdmin() {
			NotFoundHandler(c)
		}
		c.Next()
	}
}

func ScopesRequired(scopes ...string) gin.HandlerFunc {
	return func(c *gin.Context) {
		mySessionData := oauth2.NewSession("", "")
		ctx, err := oauth2.Oauth2.IntrospectToken(c, fosite.AccessTokenFromRequest(c.Request), fosite.AccessToken, mySessionData, scopes...)
		if err != nil {
			c.Error(err)
			c.Abort()
			return
		}
		// All required scopes are found
		c.Set("fosite", ctx)
		c.Next()
	}
}

// CSP set Content Security Policy http header
func CSP() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Header("Content-Security-Policy", "default-src 'self'; img-src *; media-src *; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'")
		c.Next()
	}
}
<commit_msg>fix base64 encoded imgs (#1407)<commit_after>package middlewares

import (
	"net/http"

	"github.com/NyaaPantsu/nyaa/config"
	"github.com/NyaaPantsu/nyaa/controllers/router"
	"github.com/NyaaPantsu/nyaa/templates"
	"github.com/NyaaPantsu/nyaa/utils/log"
	msg "github.com/NyaaPantsu/nyaa/utils/messages"
	"github.com/NyaaPantsu/nyaa/utils/oauth2"
	"github.com/gin-gonic/gin"
	"github.com/ory/fosite"
)

// ErrorMiddleware for managing errors on status
func ErrorMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Next()
		if c.Writer.Status() >= 300 && config.Get().Environment == "DEVELOPMENT" {
			messages := msg.GetMessages(c)
			if messages.HasErrors() {
				log.Errorf("Request has errors: %v", messages.GetAllErrors())
			}
		}
		if c.Writer.Status() != http.StatusOK && c.Writer.Size() <= 0 {
			if c.ContentType() == "application/json" {
				messages := msg.GetMessages(c)
				messages.AddErrorT("errors", "404_not_found")
				c.JSON(c.Writer.Status(), messages.GetAllErrors())
				return
			}
			templates.HttpError(c, c.Writer.Status())
		}
	}
}

// ModMiddleware Make sure the user is a moderator, otherwise return forbidden
func ModMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		currentUser := router.GetUser(c)
		if !currentUser.HasAdmin() {
			NotFoundHandler(c)
		}
		c.Next()
	}
}

func ScopesRequired(scopes ...string) gin.HandlerFunc {
	return func(c *gin.Context) {
		mySessionData := oauth2.NewSession("", "")
		ctx, err := oauth2.Oauth2.IntrospectToken(c, fosite.AccessTokenFromRequest(c.Request), fosite.AccessToken, mySessionData, scopes...)
		if err != nil {
			c.Error(err)
			c.Abort()
			return
		}
		// All required scopes are found
		c.Set("fosite", ctx)
		c.Next()
	}
}

// CSP set Content Security Policy http header
func CSP() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Header("Content-Security-Policy", "default-src 'self'; img-src * data:; media-src *; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'")
		c.Next()
	}
}
<|endoftext|>


<commit_before>package main

import (
    "github.com/caa-dev-apps/cefmdd_v1/pkg/diag"
    "github.com/caa-dev-apps/cefmdd_v1/pkg/utils"
)

func main() {

    diag.Info(diag.Yellow("cefmdd v1.0.7, (08 Feb 2017)"))

	s_args, err := utils.NewCefArgs()
	if err != nil {
		return
	}

	err = ReadCef(&s_args)
	if err != nil {
		diag.Errorf(diag.BoldRed("Error parsing cef file"), "\n%v", err.Error())
		return
	}

    diag.Info(diag.Yellow("CAA Rocks!\n"))
}
<commit_msg>v1.0.8<commit_after>package main

import (
    "github.com/caa-dev-apps/cefmdd_v1/pkg/diag"
    "github.com/caa-dev-apps/cefmdd_v1/pkg/utils"
)

func main() {

    diag.Info(diag.Yellow("cefmdd v1.0.8, (09 Feb 2017)"))

	s_args, err := utils.NewCefArgs()
	if err != nil {
		return
	}

	err = ReadCef(&s_args)
	if err != nil {
		diag.Errorf(diag.BoldRed("Error parsing cef file"), "\n%v", err.Error())
		return
	}

    diag.Info(diag.Yellow("CAA Rocks!\n"))
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 client

import (
	"encoding/json"
	"fmt"
	"os"

	shttp "github.com/skydive-project/skydive/http"
	"github.com/skydive-project/skydive/logging"
	"github.com/spf13/cobra"
)

var (
	authenticationOpts shttp.AuthenticationOpts
	gremlinQuery       string
)

var Client = &cobra.Command{
	Use:          "client",
	Short:        "Skydive client",
	Long:         "Skydive client",
	SilenceUsage: true,
}

func printJSON(obj interface{}) {
	s, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		logging.GetLogger().Errorf(err.Error())
		os.Exit(1)
	}
	fmt.Println(string(s))
}

func setFromFlag(cmd *cobra.Command, flag string, value *string) {
	if flag := cmd.LocalFlags().Lookup(flag); flag.Changed {
		*value = flag.Value.String()
	}
}

func init() {
	Client.PersistentFlags().StringVarP(&authenticationOpts.Username, "username", "", os.Getenv("SKYDIVE_USERNAME"), "username auth parameter")
	Client.PersistentFlags().StringVarP(&authenticationOpts.Password, "password", "", os.Getenv("SKYDIVE_PASSWORD"), "password auth parameter")

	Client.AddCommand(AlertCmd)
	Client.AddCommand(CaptureCmd)
	Client.AddCommand(TopologyCmd)
}
<commit_msg>Set default analyzer address for client<commit_after>/*
 * Copyright (C) 2016 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 client

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/skydive-project/skydive/config"
	shttp "github.com/skydive-project/skydive/http"
	"github.com/skydive-project/skydive/logging"
	"github.com/spf13/cobra"
)

var (
	authenticationOpts shttp.AuthenticationOpts
	gremlinQuery       string
	analyzerAddr       string
)

var Client = &cobra.Command{
	Use:          "client",
	Short:        "Skydive client",
	Long:         "Skydive client",
	SilenceUsage: true,
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		if flag := cmd.Flags().Lookup("analyzer"); !flag.Changed {
			config.GetConfig().Set("agent.analyzers", "localhost:8082")
		} else {
			config.GetConfig().Set("agent.analyzers", flag.Value)
		}
	},
}

func printJSON(obj interface{}) {
	s, err := json.MarshalIndent(obj, "", "  ")
	if err != nil {
		logging.GetLogger().Errorf(err.Error())
		os.Exit(1)
	}
	fmt.Println(string(s))
}

func setFromFlag(cmd *cobra.Command, flag string, value *string) {
	if flag := cmd.LocalFlags().Lookup(flag); flag.Changed {
		*value = flag.Value.String()
	}
}

func init() {
	Client.PersistentFlags().StringVarP(&authenticationOpts.Username, "username", "", os.Getenv("SKYDIVE_USERNAME"), "username auth parameter")
	Client.PersistentFlags().StringVarP(&authenticationOpts.Password, "password", "", os.Getenv("SKYDIVE_PASSWORD"), "password auth parameter")
	Client.PersistentFlags().String("analyzer", "", "analyzer address")

	Client.AddCommand(AlertCmd)
	Client.AddCommand(CaptureCmd)
	Client.AddCommand(TopologyCmd)
}
<|endoftext|>


<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 main

import (
	"bufio"
	"context"
	"flag"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		expg       func(context.Context) internal.ExperimentSource
		fetchQueue queue.Queue
	)
	proxyClient, err := proxy.New(*proxyURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}
	if *directProxy {
		var pds *proxydatasource.DataSource
		if *bypassLicenseCheck {
			pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
		} else {
			pds = proxydatasource.New(proxyClient)
		}
		dsg = func(context.Context) internal.DataSource { return pds }
		expg = func(context.Context) internal.ExperimentSource {
			return internal.NewLocalExperimentSource(readLocalExperiments(ctx))
		}
	} else {
		// Wrap the postgres driver with OpenCensus instrumentation.
		ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
		if err != nil {
			log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
		}
		ddb, err := openDB(ctx, cfg, ocDriver)
		if err != nil {
			log.Fatal(ctx, err)
		}
		var db *postgres.DB
		if *bypassLicenseCheck {
			db = postgres.NewBypassingLicenseCheck(ddb)
		} else {
			db = postgres.New(ddb)
		}
		defer db.Close()
		dsg = func(context.Context) internal.DataSource { return db }
		expg = func(context.Context) internal.ExperimentSource { return db }
		sourceClient := source.NewClient(config.SourceTimeout)
		// queue.New uses the db argument only while it is constructing the queue.Queue.
		// The closure passed to it is only used for testing and local execution, not in production.
		// So it's okay that in neither case do we use a per-request connection.
		fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, db,
			func(ctx context.Context, modulePath, version string) (int, error) {
				return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
			})
		if err != nil {
			log.Fatalf(ctx, "queue.New: %v", err)
		}
	}
	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}
	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	requestLogger := getLogger(ctx, cfg)
	experimenter, err := middleware.NewExperimenter(ctx, 1*time.Minute, expg)
	if err != nil {
		log.Fatal(ctx, err)
	}
	mw := middleware.Chain(
		middleware.RequestLog(requestLogger),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),      // potentially redirects so should be early in chain
		middleware.SecureHeaders(), // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		return ddb, nil
	}
	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	return database.Open(driver, ci, cfg.InstanceID)
}
func getLogger(ctx context.Context, cfg *config.Config) middleware.Logger {
	if cfg.OnGCP() {
		logger, err := log.UseStackdriver(ctx, cfg, "frontend-log")
		if err != nil {
			log.Fatal(ctx, err)
		}
		return logger
	}
	return middleware.LocalLogger{}
}

// Read a file of experiments used to initialize the local experiment source
// for use in direct proxy mode.
// Format of the file: each line is
//     name,rollout
// For each experiment.
func readLocalExperiments(ctx context.Context) []*internal.Experiment {
	filename := config.GetEnv("GO_DISCOVERY_LOCAL_EXPERIMENTS", "")
	if filename == "" {
		return nil
	}
	f, err := os.Open(filename)
	if err != nil {
		log.Fatal(ctx, err)
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	var experiments []*internal.Experiment
	log.Infof(ctx, "reading experiments from %q for local development", filename)
	for scan.Scan() {
		line := strings.TrimSpace(scan.Text())
		parts := strings.SplitN(line, ",", 3)
		if len(parts) != 2 {
			log.Fatalf(ctx, "invalid experiment in file: %q", line)
		}
		name := parts[0]
		if name == "" {
			log.Fatalf(ctx, "invalid experiment in file (name cannot be empty): %q", line)
		}
		rollout, err := strconv.ParseUint(parts[1], 10, 0)
		if err != nil {
			log.Fatalf(ctx, "invalid experiment in file (invalid rollout): %v", err)
		}
		if rollout > 100 {
			log.Fatalf(ctx, "invalid experiment in file (rollout must be between 0 - 100): %q", line)
		}
		experiments = append(experiments, &internal.Experiment{
			Name:    name,
			Rollout: uint(rollout),
		})
		log.Infof(ctx, "experiment %q: rollout = %d", name, rollout)
	}
	if err := scan.Err(); err != nil {
		log.Fatalf(ctx, "scanning %s: %v", filename, err)
	}
	log.Infof(ctx, "found %d experiment(s)", len(experiments))
	return experiments
}
<commit_msg>x/pkgsite: fix db error not being wrapped<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 main

import (
	"bufio"
	"context"
	"flag"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		expg       func(context.Context) internal.ExperimentSource
		fetchQueue queue.Queue
	)
	proxyClient, err := proxy.New(*proxyURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}
	if *directProxy {
		var pds *proxydatasource.DataSource
		if *bypassLicenseCheck {
			pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
		} else {
			pds = proxydatasource.New(proxyClient)
		}
		dsg = func(context.Context) internal.DataSource { return pds }
		expg = func(context.Context) internal.ExperimentSource {
			return internal.NewLocalExperimentSource(readLocalExperiments(ctx))
		}
	} else {
		// Wrap the postgres driver with OpenCensus instrumentation.
		ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
		if err != nil {
			log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
		}
		ddb, err := openDB(ctx, cfg, ocDriver)
		if err != nil {
			log.Fatal(ctx, err)
		}
		var db *postgres.DB
		if *bypassLicenseCheck {
			db = postgres.NewBypassingLicenseCheck(ddb)
		} else {
			db = postgres.New(ddb)
		}
		defer db.Close()
		dsg = func(context.Context) internal.DataSource { return db }
		expg = func(context.Context) internal.ExperimentSource { return db }
		sourceClient := source.NewClient(config.SourceTimeout)
		// queue.New uses the db argument only while it is constructing the queue.Queue.
		// The closure passed to it is only used for testing and local execution, not in production.
		// So it's okay that in neither case do we use a per-request connection.
		fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, db,
			func(ctx context.Context, modulePath, version string) (int, error) {
				return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
			})
		if err != nil {
			log.Fatalf(ctx, "queue.New: %v", err)
		}
	}
	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}
	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	requestLogger := getLogger(ctx, cfg)
	experimenter, err := middleware.NewExperimenter(ctx, 1*time.Minute, expg)
	if err != nil {
		log.Fatal(ctx, err)
	}
	mw := middleware.Chain(
		middleware.RequestLog(requestLogger),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),      // potentially redirects so should be early in chain
		middleware.SecureHeaders(), // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	defer derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		return ddb, nil
	}
	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	return database.Open(driver, ci, cfg.InstanceID)
}
func getLogger(ctx context.Context, cfg *config.Config) middleware.Logger {
	if cfg.OnGCP() {
		logger, err := log.UseStackdriver(ctx, cfg, "frontend-log")
		if err != nil {
			log.Fatal(ctx, err)
		}
		return logger
	}
	return middleware.LocalLogger{}
}

// Read a file of experiments used to initialize the local experiment source
// for use in direct proxy mode.
// Format of the file: each line is
//     name,rollout
// For each experiment.
func readLocalExperiments(ctx context.Context) []*internal.Experiment {
	filename := config.GetEnv("GO_DISCOVERY_LOCAL_EXPERIMENTS", "")
	if filename == "" {
		return nil
	}
	f, err := os.Open(filename)
	if err != nil {
		log.Fatal(ctx, err)
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	var experiments []*internal.Experiment
	log.Infof(ctx, "reading experiments from %q for local development", filename)
	for scan.Scan() {
		line := strings.TrimSpace(scan.Text())
		parts := strings.SplitN(line, ",", 3)
		if len(parts) != 2 {
			log.Fatalf(ctx, "invalid experiment in file: %q", line)
		}
		name := parts[0]
		if name == "" {
			log.Fatalf(ctx, "invalid experiment in file (name cannot be empty): %q", line)
		}
		rollout, err := strconv.ParseUint(parts[1], 10, 0)
		if err != nil {
			log.Fatalf(ctx, "invalid experiment in file (invalid rollout): %v", err)
		}
		if rollout > 100 {
			log.Fatalf(ctx, "invalid experiment in file (rollout must be between 0 - 100): %q", line)
		}
		experiments = append(experiments, &internal.Experiment{
			Name:    name,
			Rollout: uint(rollout),
		})
		log.Infof(ctx, "experiment %q: rollout = %d", name, rollout)
	}
	if err := scan.Err(); err != nil {
		log.Fatalf(ctx, "scanning %s: %v", filename, err)
	}
	log.Infof(ctx, "found %d experiment(s)", len(experiments))
	return experiments
}
<|endoftext|>


<commit_before>// Copyright 2011 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 main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/scanner"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"code.google.com/p/go.tools/go/gccgoimporter"
	_ "code.google.com/p/go.tools/go/gcimporter"
	"code.google.com/p/go.tools/go/types"
)

var (
	// main operation modes
	allFiles  = flag.Bool("a", false, "use all (incl. _test.go) files when processing a directory")
	allErrors = flag.Bool("e", false, "report all errors (not just the first 10)")
	verbose   = flag.Bool("v", false, "verbose mode")
	gccgo     = flag.Bool("gccgo", false, "use gccgoimporter instead of gcimporter")

	// debugging support
	sequential    = flag.Bool("seq", false, "parse sequentially, rather than in parallel")
	printAST      = flag.Bool("ast", false, "print AST (forces -seq)")
	printTrace    = flag.Bool("trace", false, "print parse trace (forces -seq)")
	parseComments = flag.Bool("comments", false, "parse comments (ignored unless -ast or -trace is provided)")
)

var (
	fset       = token.NewFileSet()
	errorCount = 0
	parserMode parser.Mode
	sizes      types.Sizes
)

func initParserMode() {
	if *allErrors {
		parserMode |= parser.AllErrors
	}
	if *printTrace {
		parserMode |= parser.Trace
	}
	if *parseComments && (*printAST || *printTrace) {
		parserMode |= parser.ParseComments
	}
}

func initSizes() {
	wordSize := 8
	maxAlign := 8
	switch build.Default.GOARCH {
	case "386", "arm":
		wordSize = 4
		maxAlign = 4
		// add more cases as needed
	}
	sizes = &types.StdSizes{WordSize: int64(wordSize), MaxAlign: int64(maxAlign)}
}

func usage() {
	fmt.Fprintln(os.Stderr, "usage: gotype [flags] [path ...]")
	flag.PrintDefaults()
	os.Exit(2)
}

func report(err error) {
	scanner.PrintError(os.Stderr, err)
	if list, ok := err.(scanner.ErrorList); ok {
		errorCount += len(list)
		return
	}
	errorCount++
}

// parse may be called concurrently
func parse(filename string, src interface{}) (*ast.File, error) {
	if *verbose {
		fmt.Println(filename)
	}
	file, err := parser.ParseFile(fset, filename, src, parserMode) // ok to access fset concurrently
	if *printAST {
		ast.Print(fset, file)
	}
	return file, err
}

func parseStdin() (*ast.File, error) {
	src, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}
	return parse("<standard input>", src)
}

func parseFiles(filenames []string) ([]*ast.File, error) {
	files := make([]*ast.File, len(filenames))

	if *sequential {
		for i, filename := range filenames {
			var err error
			files[i], err = parse(filename, nil)
			if err != nil {
				return nil, err // leave unfinished goroutines hanging
			}
		}
	} else {
		type parseResult struct {
			file *ast.File
			err  error
		}

		out := make(chan parseResult)
		for _, filename := range filenames {
			go func(filename string) {
				file, err := parse(filename, nil)
				out <- parseResult{file, err}
			}(filename)
		}

		for i := range filenames {
			res := <-out
			if res.err != nil {
				return nil, res.err // leave unfinished goroutines hanging
			}
			files[i] = res.file
		}
	}

	return files, nil
}

func parseDir(dirname string) ([]*ast.File, error) {
	ctxt := build.Default
	pkginfo, err := ctxt.ImportDir(dirname, 0)
	if _, nogo := err.(*build.NoGoError); err != nil && !nogo {
		return nil, err
	}
	filenames := append(pkginfo.GoFiles, pkginfo.CgoFiles...)
	if *allFiles {
		filenames = append(filenames, pkginfo.TestGoFiles...)
	}

	// complete file names
	for i, filename := range filenames {
		filenames[i] = filepath.Join(dirname, filename)
	}

	return parseFiles(filenames)
}

func getPkgFiles(args []string) ([]*ast.File, error) {
	if len(args) == 0 {
		// stdin
		file, err := parseStdin()
		if err != nil {
			return nil, err
		}
		return []*ast.File{file}, nil
	}

	if len(args) == 1 {
		// possibly a directory
		path := args[0]
		info, err := os.Stat(path)
		if err != nil {
			return nil, err
		}
		if info.IsDir() {
			return parseDir(path)
		}
	}

	// list of files
	return parseFiles(args)
}

func checkPkgFiles(files []*ast.File) {
	type bailout struct{}
	conf := types.Config{
		FakeImportC: true,
		Error: func(err error) {
			if !*allErrors && errorCount >= 10 {
				panic(bailout{})
			}
			report(err)
		},
		Sizes: sizes,
	}
	if *gccgo {
		var inst gccgoimporter.GccgoInstallation
		inst.InitFromDriver("gccgo")
		conf.Import = inst.GetImporter(nil)
	}

	defer func() {
		switch p := recover().(type) {
		case nil, bailout:
			// normal return or early exit
		default:
			// re-panic
			panic(p)
		}
	}()

	const path = "pkg" // any non-empty string will do for now
	conf.Check(path, fset, files, nil)
}

func printStats(d time.Duration) {
	fileCount := 0
	lineCount := 0
	fset.Iterate(func(f *token.File) bool {
		fileCount++
		lineCount += f.LineCount()
		return true
	})

	fmt.Printf(
		"%s (%d files, %d lines, %d lines/s)\n",
		d, fileCount, lineCount, int64(float64(lineCount)/d.Seconds()),
	)
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU()) // remove this once runtime is smarter

	flag.Usage = usage
	flag.Parse()
	if *printAST || *printTrace {
		*sequential = true
	}
	initParserMode()
	initSizes()

	start := time.Now()

	files, err := getPkgFiles(flag.Args())
	if err != nil {
		report(err)
		os.Exit(2)
	}

	checkPkgFiles(files)
	if errorCount > 0 {
		os.Exit(2)
	}

	if *verbose {
		printStats(time.Since(start))
	}
}
<commit_msg>go.tools/cmd/gotype: fix call site (fix build)<commit_after>// Copyright 2011 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 main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/scanner"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"code.google.com/p/go.tools/go/gccgoimporter"
	_ "code.google.com/p/go.tools/go/gcimporter"
	"code.google.com/p/go.tools/go/types"
)

var (
	// main operation modes
	allFiles  = flag.Bool("a", false, "use all (incl. _test.go) files when processing a directory")
	allErrors = flag.Bool("e", false, "report all errors (not just the first 10)")
	verbose   = flag.Bool("v", false, "verbose mode")
	gccgo     = flag.Bool("gccgo", false, "use gccgoimporter instead of gcimporter")

	// debugging support
	sequential    = flag.Bool("seq", false, "parse sequentially, rather than in parallel")
	printAST      = flag.Bool("ast", false, "print AST (forces -seq)")
	printTrace    = flag.Bool("trace", false, "print parse trace (forces -seq)")
	parseComments = flag.Bool("comments", false, "parse comments (ignored unless -ast or -trace is provided)")
)

var (
	fset       = token.NewFileSet()
	errorCount = 0
	parserMode parser.Mode
	sizes      types.Sizes
)

func initParserMode() {
	if *allErrors {
		parserMode |= parser.AllErrors
	}
	if *printTrace {
		parserMode |= parser.Trace
	}
	if *parseComments && (*printAST || *printTrace) {
		parserMode |= parser.ParseComments
	}
}

func initSizes() {
	wordSize := 8
	maxAlign := 8
	switch build.Default.GOARCH {
	case "386", "arm":
		wordSize = 4
		maxAlign = 4
		// add more cases as needed
	}
	sizes = &types.StdSizes{WordSize: int64(wordSize), MaxAlign: int64(maxAlign)}
}

func usage() {
	fmt.Fprintln(os.Stderr, "usage: gotype [flags] [path ...]")
	flag.PrintDefaults()
	os.Exit(2)
}

func report(err error) {
	scanner.PrintError(os.Stderr, err)
	if list, ok := err.(scanner.ErrorList); ok {
		errorCount += len(list)
		return
	}
	errorCount++
}

// parse may be called concurrently
func parse(filename string, src interface{}) (*ast.File, error) {
	if *verbose {
		fmt.Println(filename)
	}
	file, err := parser.ParseFile(fset, filename, src, parserMode) // ok to access fset concurrently
	if *printAST {
		ast.Print(fset, file)
	}
	return file, err
}

func parseStdin() (*ast.File, error) {
	src, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}
	return parse("<standard input>", src)
}

func parseFiles(filenames []string) ([]*ast.File, error) {
	files := make([]*ast.File, len(filenames))

	if *sequential {
		for i, filename := range filenames {
			var err error
			files[i], err = parse(filename, nil)
			if err != nil {
				return nil, err // leave unfinished goroutines hanging
			}
		}
	} else {
		type parseResult struct {
			file *ast.File
			err  error
		}

		out := make(chan parseResult)
		for _, filename := range filenames {
			go func(filename string) {
				file, err := parse(filename, nil)
				out <- parseResult{file, err}
			}(filename)
		}

		for i := range filenames {
			res := <-out
			if res.err != nil {
				return nil, res.err // leave unfinished goroutines hanging
			}
			files[i] = res.file
		}
	}

	return files, nil
}

func parseDir(dirname string) ([]*ast.File, error) {
	ctxt := build.Default
	pkginfo, err := ctxt.ImportDir(dirname, 0)
	if _, nogo := err.(*build.NoGoError); err != nil && !nogo {
		return nil, err
	}
	filenames := append(pkginfo.GoFiles, pkginfo.CgoFiles...)
	if *allFiles {
		filenames = append(filenames, pkginfo.TestGoFiles...)
	}

	// complete file names
	for i, filename := range filenames {
		filenames[i] = filepath.Join(dirname, filename)
	}

	return parseFiles(filenames)
}

func getPkgFiles(args []string) ([]*ast.File, error) {
	if len(args) == 0 {
		// stdin
		file, err := parseStdin()
		if err != nil {
			return nil, err
		}
		return []*ast.File{file}, nil
	}

	if len(args) == 1 {
		// possibly a directory
		path := args[0]
		info, err := os.Stat(path)
		if err != nil {
			return nil, err
		}
		if info.IsDir() {
			return parseDir(path)
		}
	}

	// list of files
	return parseFiles(args)
}

func checkPkgFiles(files []*ast.File) {
	type bailout struct{}
	conf := types.Config{
		FakeImportC: true,
		Error: func(err error) {
			if !*allErrors && errorCount >= 10 {
				panic(bailout{})
			}
			report(err)
		},
		Sizes: sizes,
	}
	if *gccgo {
		var inst gccgoimporter.GccgoInstallation
		inst.InitFromDriver("gccgo")
		conf.Import = inst.GetImporter(nil, nil)
	}

	defer func() {
		switch p := recover().(type) {
		case nil, bailout:
			// normal return or early exit
		default:
			// re-panic
			panic(p)
		}
	}()

	const path = "pkg" // any non-empty string will do for now
	conf.Check(path, fset, files, nil)
}

func printStats(d time.Duration) {
	fileCount := 0
	lineCount := 0
	fset.Iterate(func(f *token.File) bool {
		fileCount++
		lineCount += f.LineCount()
		return true
	})

	fmt.Printf(
		"%s (%d files, %d lines, %d lines/s)\n",
		d, fileCount, lineCount, int64(float64(lineCount)/d.Seconds()),
	)
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU()) // remove this once runtime is smarter

	flag.Usage = usage
	flag.Parse()
	if *printAST || *printTrace {
		*sequential = true
	}
	initParserMode()
	initSizes()

	start := time.Now()

	files, err := getPkgFiles(flag.Args())
	if err != nil {
		report(err)
		os.Exit(2)
	}

	checkPkgFiles(files)
	if errorCount > 0 {
		os.Exit(2)
	}

	if *verbose {
		printStats(time.Since(start))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"io/ioutil"
	"math/rand"
	"os"
	"time"

	"golang.org/x/tools/imports"
	"kego.io/kerr"
	"kego.io/kerr/kerrsource"
)

func main() {
	rand.Seed(time.Now().UTC().UnixNano())
	flag.Parse()
	if err := processFile(flag.Arg(0)); err != nil {
		panic(err.Error())
	}
}

func processFile(filename string) error {

	if filename == "" {
		return kerr.New("CKBKFTNVXI", "File not found")
	}

	f, err := os.Open(filename)
	if err != nil {
		return kerr.Wrap("WCQEVUWMAI", err)
	}
	defer f.Close()

	src, err := ioutil.ReadAll(f)
	if err != nil {
		return kerr.Wrap("DWTHNAVSVY", err)
	}

	var res []byte

	res, err = kerrsource.Process(filename, res)
	if err != nil {
		return kerr.Wrap("UFVVOEXSYE", err)
	}

	res, err = imports.Process(filename, src, nil)
	if err != nil {
		return kerr.Wrap("XFILQRWHSI", err)
	}

	if !bytes.Equal(src, res) {
		// formatting has changed
		err = ioutil.WriteFile(filename, res, 0)
		if err != nil {
			return kerr.Wrap("NFFAOHJQLG", err)
		}
	}

	return nil

}
<commit_msg>Fix kesource bug<commit_after>package main

import (
	"bytes"
	"flag"
	"io/ioutil"
	"math/rand"
	"os"
	"time"

	"golang.org/x/tools/imports"
	"kego.io/kerr"
	"kego.io/kerr/kerrsource"
)

func main() {
	rand.Seed(time.Now().UTC().UnixNano())
	flag.Parse()
	if err := processFile(flag.Arg(0)); err != nil {
		panic(err.Error())
	}
}

func processFile(filename string) error {

	if filename == "" {
		return kerr.New("CKBKFTNVXI", "File not found")
	}

	f, err := os.Open(filename)
	if err != nil {
		return kerr.Wrap("WCQEVUWMAI", err)
	}
	defer f.Close()

	src, err := ioutil.ReadAll(f)
	if err != nil {
		return kerr.Wrap("DWTHNAVSVY", err)
	}

	var res []byte

	res, err = kerrsource.Process(filename, src)
	if err != nil {
		return kerr.Wrap("UFVVOEXSYE", err)
	}

	res, err = imports.Process(filename, src, nil)
	if err != nil {
		return kerr.Wrap("XFILQRWHSI", err)
	}

	if !bytes.Equal(src, res) {
		// formatting has changed
		err = ioutil.WriteFile(filename, res, 0)
		if err != nil {
			return kerr.Wrap("NFFAOHJQLG", err)
		}
	}

	return nil

}
<|endoftext|>


<commit_before>// Copyright (C) 2016  Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package main

type telnetClient struct {
	addr string
}

func newTelnetClient(addr string) OpenTSDBConn {
	return &telnetClient{
		addr: addr,
	}
}

func (c *telnetClient) Put(d *DataPoint) error {
	panic("TODO")
}
<commit_msg>cmd/octsdb: Finally implement the `-tsdb` flag.<commit_after>// Copyright (C) 2016  Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package main

import (
	"bytes"
	"net"

	"github.com/aristanetworks/glog"
)

type telnetClient struct {
	addr string
	conn net.Conn
}

func newTelnetClient(addr string) OpenTSDBConn {
	return &telnetClient{
		addr: addr,
	}
}

func readErrors(conn net.Conn) {
	var buf [4096]byte
	for {
		// TODO: We should add a buffer to read line-by-line properly instead
		// of using a fixed-size buffer and splitting on newlines manually.
		n, err := conn.Read(buf[:])
		if n == 0 {
			return
		} else if n > 0 {
			for line := range bytes.Split(buf[:n], []byte{'\n'}) {
				glog.Info("tsd replied: ", string(line))
			}
		}
		if err != nil {
			return
		}
	}
}

func (c *telnetClient) Put(d *DataPoint) error {
	var err error
	if c.conn == nil {
		c.conn, err = net.Dial("tcp", c.addr)
		if err != nil {
			return err
		}
		go readErrors(c.conn)
	}
	_, err = c.conn.Write([]byte(d.String()))
	if err != nil {
		c.conn.Close()
		c.conn = nil
	}
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"github.com/hanjm/log"
	"io"
	"net"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type Task interface {
	Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error
	IsCompleted() bool
	FileName() string
	ContentLength() int64
}

func NewDownloadTask(sourceURL string) (Task, error) {
	switch {
	case strings.HasPrefix(sourceURL, "http"):
		return NewHTTPTask(sourceURL), nil
	case strings.HasPrefix(sourceURL, "magnet:?xt=urn:btih:") || IsBase64String(sourceURL):
		return NewMagnetTask(sourceURL), nil
	default:
		return nil, fmt.Errorf("sourceURL expect http or magnet, not %s", sourceURL)
	}
}

const (
	DownloadTaskTypeHTTP = iota
	DownloadTaskTypeMagnet
)

type TaskInfo struct {
	SourceURL     string
	StartTime     time.Time
	FileName      string
	ContentLength int64         // B 总大小
	Size          int64         // B 已下载的大小
	Duration      time.Duration // s 耗时
	Speed         int64         // B/s 速度
	IsCompleted   bool          // 是否完成
	IsError       bool          // 是否出错
	Error         string        // 错误消息
}

// download http content
type HTTPTask struct {
	TaskType int
	TaskInfo
}

func NewHTTPTask(sourceUrl string) *HTTPTask {
	return &HTTPTask{
		TaskType: DownloadTaskTypeHTTP,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}

func (t *HTTPTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error {
	var httpClient = &http.Client{
		Timeout: limitTimeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 20 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 20 * time.Second,
		},
	}
	t.StartTime = time.Now()
	resp, err := httpClient.Get(t.SourceURL)
	if err != nil {
		return t.Errorf("http.Client error:%s", err)
	}
	defer resp.Body.Close()
	t.TaskInfo.ContentLength = resp.ContentLength
	contentDisposition := strings.SplitN(resp.Header.Get("Content-Disposition"), "=", 2)
	var attachmentName string
	if len(contentDisposition) > 1 {
		attachmentName = contentDisposition[1]
	}
	if t.TaskInfo.ContentLength <= 0 {
		resp.Body.Close()
		//一些资源是动态生成的,请求第一次是chunked stream,Header不带Content-Length,第二次请求就有Content-length
		resp, err = httpClient.Get(t.SourceURL)
		if err != nil {
			return t.Errorf("http.Client error:%s", err)
		}
		t.TaskInfo.ContentLength = resp.ContentLength
	}
	// if header has attach filename, update it
	if attachmentName != "" {
		attachmentName2 := getSafeFilename(attachmentName)
		if attachmentName2 != t.TaskInfo.FileName {
			t.TaskInfo.FileName = attachmentName2
		}
	}
	log.Infof("create HTTP task: length:%s source:%s filename:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName)
	if t.TaskInfo.ContentLength > limitByteSize {
		return t.Errorf("the content length of sourceUrl is too big:%d, limit:%d", t.TaskInfo.ContentLength, limitByteSize)
	}
	// write file
	filename := downloadDir + "/" + t.TaskInfo.FileName
	os.Remove(filename)
	fp, err := os.Create(filename)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	defer fp.Sync()
	defer fp.Close()
	bufSize := 4096
	bodyReader := bufio.NewReaderSize(resp.Body, bufSize)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	buf := make([]byte, bufSize)
	size := 0
	readSize := 0
	completed := false
	for i := 0; ; i++ {
		readSize, err = bodyReader.Read(buf)
		if err != nil {
			if err == io.EOF {
				// 正常下载完
				completed = true
			} else {
				return t.Errorf("body read error:%s", err)
			}
		}
		_, err = fp.Write(buf[:readSize])
		size += readSize
		t.Size = int64(size)
		if i%1000 == 0 {
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
		}
		if err != nil {
			return t.Errorf("body write error:%s", err)
		}
		if completed {
			break
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Size = int64(size)
	t.TaskInfo.ContentLength = t.Size
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	log.Infof("complete HTTP task: length:%s source:%s filename:%s, duration:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName, t.Duration)
	return nil
}

func (t *HTTPTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *HTTPTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *HTTPTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *HTTPTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

// download magnet content
type MagnetTask struct {
	TaskType int
	TaskInfo
}

func NewMagnetTask(sourceUrl string) *MagnetTask {
	return &MagnetTask{
		TaskType: DownloadTaskTypeMagnet,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}
func (t *MagnetTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) (err error) {
	if !IsAria2cRunning() {
		return t.Errorf("aria2c is not running, cannot download magnet")
	}
	aria2cRPCClient := NewAria2cRPCClient()
	var taskGID string
	// try base64 decode
	data, err := base64.StdEncoding.DecodeString(t.SourceURL)
	if err != nil {
		taskGID, err = aria2cRPCClient.AddURI(t.SourceURL)
		if err != nil {
			return t.Errorf("call aria2c AddURI error:%s", err)
		}
	} else {
		taskGID, err = aria2cRPCClient.AddTorrent(t.SourceURL)
		if err != nil {
			return t.Errorf("call aria2c AddTorrent error:%s", err)
		}
		// save to file and change the sourceURL
		torrentFilename := fmt.Sprintf("%s/%s.torrent", downloadDir, t.SourceURL[:16])
		fp, err := os.Create(torrentFilename)
		if err != nil {
			log.Warnf("os.Create file error:%s", err)
		}
		_, err = fp.Write(data)
		if err != nil {
			log.Warnf("fp.Write error:%s", err)
		}
		t.SourceURL = torrentFilename
	}
	log.Infof("create Magnet task: sourceURL:%s, taskGID:%s", t.SourceURL, taskGID)
	t.StartTime = time.Now()
	timeout := limitTimeout
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
MagnetLoop:
	complete := false
	ticker := time.NewTicker(time.Second * 5)
	for !complete {
		select {
		case <-ticker.C:
			result, err := aria2cRPCClient.TellStatus(taskGID)
			if err != nil {
				return t.Errorf("call aria2c TellStatus error:%s", err)
			}
			// update break condition
			complete = result.Completed()
			// udpate taskInfo
			t.TaskInfo.ContentLength = result.TotalLength
			t.Size = result.CompletedLength
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = result.DownloadSpeed
			// 磁力链接建立任务时无法指定文件名 获得真实文件名后需要重命名
			realFilename := strings.TrimPrefix(result.GetFilePath(), downloadDir+"/")
			pos := strings.Index(realFilename, "/")
			if pos != -1 && pos < len(realFilename)-1 {
				realFilename = realFilename[:pos]
			}
			t.TaskInfo.FileName = realFilename
			// 检查是否有继续下载磁力链接包含的其他文件
			followedBys := result.FollowedBy
			for _, followedTaskGID := range followedBys {
				taskGID = followedTaskGID
				complete = false
				log.Debugf("task status:%+v goto MagnetLoop", result)
				goto MagnetLoop
			}
		case <-ctx.Done():
			if ctx.Err() == context.DeadlineExceeded {
				return t.Errorf("task timeout:%s", timeout)
			}
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	return err
}

func (t *MagnetTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *MagnetTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *MagnetTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *MagnetTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

//utils
var safeFilenameRegexp = regexp.MustCompile(`[\w\d\-.]+`)

func getSafeFilename(str string) string {
	filename := strings.Join(safeFilenameRegexp.FindAllString(str, -1), "")
	if lenOfFilename := len(filename); lenOfFilename > 50 {
		filename = filename[lenOfFilename-50 : lenOfFilename]
	}
	return fmt.Sprintf("%d-%s", time.Now().Unix(), filename)
}

func getHumanSizeString(byteSize int64) string {
	units := []string{"B", "KB", "MB", "GB", "TB", "EB"}
	index := 0
	byteSizeFloat := float64(byteSize)
	for ; byteSizeFloat > 1024; index += 1 {
		byteSizeFloat /= 1024
	}
	var unit string
	if index < len(units) {
		unit = units[index]
	} else {
		unit = "INF"
	}
	return fmt.Sprintf("%.2f %s", byteSizeFloat, unit)
}

func calculateDownloadSpeed(byteSize int64, duration time.Duration) (bytePerSecond int64) {
	if duration <= 0 {
		duration = 1
	}
	return byteSize * 1e9 / int64(duration)
}

func IsBase64String(str string) bool {
	_, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return false
	}
	return true
}
<commit_msg>add more logger<commit_after>package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"github.com/hanjm/log"
	"io"
	"net"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type Task interface {
	Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error
	IsCompleted() bool
	FileName() string
	ContentLength() int64
}

func NewDownloadTask(sourceURL string) (Task, error) {
	switch {
	case strings.HasPrefix(sourceURL, "http"):
		return NewHTTPTask(sourceURL), nil
	case strings.HasPrefix(sourceURL, "magnet:?xt=urn:btih:") || IsBase64String(sourceURL):
		return NewMagnetTask(sourceURL), nil
	default:
		return nil, fmt.Errorf("sourceURL expect http or magnet, not %s", sourceURL)
	}
}

const (
	DownloadTaskTypeHTTP = iota
	DownloadTaskTypeMagnet
)

type TaskInfo struct {
	SourceURL     string
	StartTime     time.Time
	FileName      string
	ContentLength int64         // B 总大小
	Size          int64         // B 已下载的大小
	Duration      time.Duration // s 耗时
	Speed         int64         // B/s 速度
	IsCompleted   bool          // 是否完成
	IsError       bool          // 是否出错
	Error         string        // 错误消息
}

// download http content
type HTTPTask struct {
	TaskType int
	TaskInfo
}

func NewHTTPTask(sourceUrl string) *HTTPTask {
	return &HTTPTask{
		TaskType: DownloadTaskTypeHTTP,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}

func (t *HTTPTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) error {
	var httpClient = &http.Client{
		Timeout: limitTimeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: 20 * time.Second,
			}).Dial,
			TLSHandshakeTimeout: 20 * time.Second,
		},
	}
	t.StartTime = time.Now()
	resp, err := httpClient.Get(t.SourceURL)
	if err != nil {
		return t.Errorf("http.Client error:%s", err)
	}
	defer resp.Body.Close()
	t.TaskInfo.ContentLength = resp.ContentLength
	contentDisposition := strings.SplitN(resp.Header.Get("Content-Disposition"), "=", 2)
	var attachmentName string
	if len(contentDisposition) > 1 {
		attachmentName = contentDisposition[1]
	}
	if t.TaskInfo.ContentLength <= 0 {
		resp.Body.Close()
		//一些资源是动态生成的,请求第一次是chunked stream,Header不带Content-Length,第二次请求就有Content-length
		resp, err = httpClient.Get(t.SourceURL)
		if err != nil {
			return t.Errorf("http.Client error:%s", err)
		}
		t.TaskInfo.ContentLength = resp.ContentLength
	}
	// if header has attach filename, update it
	if attachmentName != "" {
		attachmentName2 := getSafeFilename(attachmentName)
		if attachmentName2 != t.TaskInfo.FileName {
			t.TaskInfo.FileName = attachmentName2
		}
	}
	log.Infof("create HTTP task: length:%s source:%s filename:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName)
	if t.TaskInfo.ContentLength > limitByteSize {
		return t.Errorf("the content length of sourceUrl is too big:%d, limit:%d", t.TaskInfo.ContentLength, limitByteSize)
	}
	// write file
	filename := downloadDir + "/" + t.TaskInfo.FileName
	os.Remove(filename)
	fp, err := os.Create(filename)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	defer fp.Sync()
	defer fp.Close()
	bufSize := 4096
	bodyReader := bufio.NewReaderSize(resp.Body, bufSize)
	if err != nil {
		return t.Errorf("create file error:%s", err)
	}
	buf := make([]byte, bufSize)
	size := 0
	readSize := 0
	completed := false
	for i := 0; ; i++ {
		readSize, err = bodyReader.Read(buf)
		if err != nil {
			if err == io.EOF {
				// 正常下载完
				completed = true
			} else {
				return t.Errorf("body read error:%s", err)
			}
		}
		_, err = fp.Write(buf[:readSize])
		size += readSize
		t.Size = int64(size)
		if i%1000 == 0 {
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
		}
		if err != nil {
			return t.Errorf("body write error:%s", err)
		}
		if completed {
			break
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Size = int64(size)
	t.TaskInfo.ContentLength = t.Size
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	log.Infof("complete HTTP task: length:%s source:%s filename:%s, duration:%s", getHumanSizeString(t.TaskInfo.ContentLength), t.SourceURL, t.TaskInfo.FileName, t.Duration)
	return nil
}

func (t *HTTPTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *HTTPTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *HTTPTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *HTTPTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

// download magnet content
type MagnetTask struct {
	TaskType int
	TaskInfo
}

func NewMagnetTask(sourceUrl string) *MagnetTask {
	return &MagnetTask{
		TaskType: DownloadTaskTypeMagnet,
		TaskInfo: TaskInfo{
			SourceURL: sourceUrl,
			FileName:  getSafeFilename(sourceUrl),
		}}
}
func (t *MagnetTask) Download(downloadDir string, limitByteSize int64, limitTimeout time.Duration) (err error) {
	if !IsAria2cRunning() {
		return t.Errorf("aria2c is not running, cannot download magnet")
	}
	aria2cRPCClient := NewAria2cRPCClient()
	var taskGID string
	// try base64 decode
	data, err := base64.StdEncoding.DecodeString(t.SourceURL)
	if err != nil {
		taskGID, err = aria2cRPCClient.AddURI(t.SourceURL)
		if err != nil {
			return t.Errorf("call aria2c AddURI error:%s", err)
		}
	} else {
		taskGID, err = aria2cRPCClient.AddTorrent(t.SourceURL)
		if err != nil {
			return t.Errorf("call aria2c AddTorrent error:%s", err)
		}
		// save to file and change the sourceURL
		torrentFilename := fmt.Sprintf("%s/%s.torrent", downloadDir, t.SourceURL[:16])
		fp, err := os.Create(torrentFilename)
		if err != nil {
			log.Warnf("os.Create file error:%s", err)
		}
		_, err = fp.Write(data)
		if err != nil {
			log.Warnf("fp.Write error:%s", err)
		}
		t.SourceURL = torrentFilename
	}
	log.Infof("create Magnet task: sourceURL:%s, taskGID:%s", t.SourceURL, taskGID)
	t.StartTime = time.Now()
	timeout := limitTimeout
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
MagnetLoop:
	complete := false
	ticker := time.NewTicker(time.Second * 5)
	for !complete {
		select {
		case <-ticker.C:
			result, err := aria2cRPCClient.TellStatus(taskGID)
			if err != nil {
				return t.Errorf("call aria2c TellStatus error:%s", err)
			}
			// update break condition
			complete = result.Completed()
			// udpate taskInfo
			t.TaskInfo.ContentLength = result.TotalLength
			t.Size = result.CompletedLength
			t.Duration = time.Now().Sub(t.StartTime)
			t.Speed = result.DownloadSpeed
			if result.CompletedLength == result.TotalLength && !result.Completed() {
				// why aria2c wait long time even it seems the download is completed.
				log.Debugf("task status:%+v goto MagnetLoop", result)
			}
			// 磁力链接建立任务时无法指定文件名 获得真实文件名后需要重命名
			realFilename := strings.TrimPrefix(result.GetFilePath(), downloadDir+"/")
			pos := strings.Index(realFilename, "/")
			if pos != -1 && pos < len(realFilename)-1 {
				realFilename = realFilename[:pos]
			}
			t.TaskInfo.FileName = realFilename
			// 检查是否有继续下载磁力链接包含的其他文件
			followedBys := result.FollowedBy
			for _, followedTaskGID := range followedBys {
				taskGID = followedTaskGID
				complete = false
				log.Debugf("task status:%+v goto MagnetLoop", result)
				goto MagnetLoop
			}
		case <-ctx.Done():
			if ctx.Err() == context.DeadlineExceeded {
				return t.Errorf("task timeout:%s", timeout)
			}
		}
	}
	t.TaskInfo.IsCompleted = true
	t.Duration = time.Now().Sub(t.StartTime)
	t.Speed = calculateDownloadSpeed(t.Size, t.Duration)
	return err
}

func (t *MagnetTask) IsCompleted() bool {
	return t.TaskInfo.IsCompleted
}

func (t *MagnetTask) FileName() string {
	return t.TaskInfo.FileName
}

func (t *MagnetTask) ContentLength() int64 {
	return t.TaskInfo.ContentLength
}

func (t *MagnetTask) Errorf(format string, a ...interface{}) (err error) {
	err = fmt.Errorf(format, a...)
	_, file, line, ok := runtime.Caller(1)
	if ok {
		log.Errorf("[%s:%d]%s", file, line, err.Error())
	}
	t.TaskInfo.IsError = true
	t.TaskInfo.IsCompleted = true
	t.TaskInfo.Error = err.Error()
	return err
}

//utils
var safeFilenameRegexp = regexp.MustCompile(`[\w\d\-.]+`)

func getSafeFilename(str string) string {
	filename := strings.Join(safeFilenameRegexp.FindAllString(str, -1), "")
	if lenOfFilename := len(filename); lenOfFilename > 50 {
		filename = filename[lenOfFilename-50 : lenOfFilename]
	}
	return fmt.Sprintf("%d-%s", time.Now().Unix(), filename)
}

func getHumanSizeString(byteSize int64) string {
	units := []string{"B", "KB", "MB", "GB", "TB", "EB"}
	index := 0
	byteSizeFloat := float64(byteSize)
	for ; byteSizeFloat > 1024; index += 1 {
		byteSizeFloat /= 1024
	}
	var unit string
	if index < len(units) {
		unit = units[index]
	} else {
		unit = "INF"
	}
	return fmt.Sprintf("%.2f %s", byteSizeFloat, unit)
}

func calculateDownloadSpeed(byteSize int64, duration time.Duration) (bytePerSecond int64) {
	if duration <= 0 {
		duration = 1
	}
	return byteSize * 1e9 / int64(duration)
}

func IsBase64String(str string) bool {
	_, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return false
	}
	return true
}
<|endoftext|>