<commit_before>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 codegen

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"unsafe"

	"github.com/google/gapid/core/app/linker"
	"github.com/google/gapid/core/os/device"

	"llvm/bindings/go/llvm"
)

func init() {
	llvm.InitializeAllTargetInfos()
	llvm.InitializeAllTargets()
	llvm.InitializeAllTargetMCs()
	llvm.InitializeAllAsmPrinters()
}

// Executor executes module functions.
type Executor struct {
	llvm     llvm.ExecutionEngine
	funcPtrs map[string]unsafe.Pointer
}

// Object compiles the module down to an object file.
func (m *Module) Object(optimize bool) ([]byte, error) {
	t, err := llvm.GetTargetFromTriple(m.triple.String())
	if err != nil {
		return nil, fmt.Errorf("Couldn't get target for triple '%v': %v", m.triple, err)
	}
	cpu := ""
	features := ""
	opt := llvm.CodeGenLevelNone
	if optimize {
		opt = llvm.CodeGenLevelDefault
	}
	reloc := llvm.RelocPIC
	model := llvm.CodeModelDefault
	tm := t.CreateTargetMachine(m.triple.String(), cpu, features, opt, reloc, model)
	defer tm.Dispose()

	// Check target data is as expected.
	td := tm.CreateTargetData()
	defer td.Dispose()
	m.validateTargetData(td)

	buf, err := tm.EmitToMemoryBuffer(m.llvm, llvm.ObjectFile)
	if err != nil {
		return nil, err
	}
	defer buf.Dispose()
	return buf.Bytes(), nil
}

func (m *Module) validateTargetData(td llvm.TargetData) {
	abi := m.target
	errs := []string{}
	check := func(llvm, gapid interface{}, name string) bool {
		if reflect.DeepEqual(llvm, gapid) {
			return true
		}
		errs = append(errs, fmt.Sprintf("%v target mismatch for %v: %v (llvm) != %v (gapid)", name, abi.Name, llvm, gapid))
		return false
	}
	checkTD := func(ty Type, dtl *device.DataTypeLayout) {
		check(td.TypeStoreSize(ty.llvmTy()), uint64(dtl.Size), ty.String()+"-size")
		check(td.ABITypeAlignment(ty.llvmTy()), int(dtl.Alignment), ty.String()+"-align")
	}

	layout := abi.MemoryLayout
	isLE := td.ByteOrder() == llvm.LittleEndian
	check(isLE, layout.Endian == device.LittleEndian, "is-little-endian")
	check(td.PointerSize(), int(layout.Pointer.Size), "pointer-size")

	checkTD(m.Types.Pointer(m.Types.Int), layout.Pointer)
	checkTD(m.Types.Int, layout.Integer)
	checkTD(m.Types.Size, layout.Size)
	checkTD(m.Types.Int64, layout.I64)
	checkTD(m.Types.Int32, layout.I32)
	checkTD(m.Types.Int16, layout.I16)
	checkTD(m.Types.Int8, layout.I8)
	checkTD(m.Types.Float32, layout.F32)
	checkTD(m.Types.Float64, layout.F64)

	for _, s := range m.Types.structs {
		if !s.hasBody {
			continue
		}
		if !check(int(td.TypeStoreSize(s.llvm))*8, s.SizeInBits(), fmt.Sprintf("%v-size", s.name)) ||
			!check(int(td.ABITypeAlignment(s.llvm))*8, s.AlignInBits(), fmt.Sprintf("%v-align", s.name)) {
			errs = append(errs, fmt.Sprintf("%v: %v", s.name, s))
		}
		for i := range s.Fields() {
			llvm := int(td.ElementOffset(s.llvm, i)) * 8
			gapid := s.FieldOffsetInBits(i)
			check(llvm, gapid, fmt.Sprintf("%v-field-offset %d", s.name, i))
		}
	}

	for _, s := range m.Types.arrays {
		check(int(td.TypeStoreSize(s.llvm))*8, s.SizeInBits(), fmt.Sprintf("%v-size", s.name))
		check(int(td.ABITypeAlignment(s.llvm))*8, s.AlignInBits(), fmt.Sprintf("%v-align", s.name))
	}

	if len(errs) > 0 {
		panic(fmt.Errorf("%v has ABI mismatches!\n%v", abi.Name, strings.Join(errs, "\n")))
	}
}

// Optimize optimizes the module.
func (m *Module) Optimize() {
	fpm := llvm.NewFunctionPassManagerForModule(m.llvm)
	defer fpm.Dispose()

	mpm := llvm.NewPassManager()
	defer mpm.Dispose()

	pmb := llvm.NewPassManagerBuilder()
	defer pmb.Dispose()

	pmb.SetOptLevel(int(llvm.CodeGenLevelDefault))
	pmb.SetSizeLevel(0)

	mpm.AddVerifierPass()
	fpm.AddVerifierPass()

	pmb.Populate(mpm)
	pmb.PopulateFunc(fpm)

	fpm.InitializeFunc()
	for fn := m.llvm.FirstFunction(); !fn.IsNil(); fn = llvm.NextFunction(fn) {
		fpm.RunFunc(fn)
	}
	fpm.FinalizeFunc()

	mpm.Run(m.llvm)
}

// Executor constructs an executor.
func (m *Module) Executor(optimize bool) (*Executor, error) {
	if err := m.Verify(); err != nil {
		return nil, err
	}

	opts := llvm.NewMCJITCompilerOptions()
	if optimize {
		opts.SetMCJITOptimizationLevel(2)
	} else {
		opts.SetMCJITOptimizationLevel(0)
	}

	engine, err := llvm.NewMCJITCompiler(m.llvm, opts)
	if err != nil {
		return nil, err
	}

	// Check target data is as expected.
	m.validateTargetData(engine.TargetData())

	// Check for unresolved extern symbols.
	var unresolved []string
	for _, f := range m.funcs {
		if f.built || strings.HasPrefix(f.Name, "llvm.") {
			continue
		}
		if linker.ProcAddress(f.Name) == 0 {
			unresolved = append(unresolved, fmt.Sprint(f))
		}
	}
	if len(unresolved) > 0 {
		sort.Strings(unresolved)
		msg := fmt.Sprintf("Unresolved external functions:\n%v", strings.Join(unresolved, "\n"))
		fail(msg)
	}

	engine.RunStaticConstructors()

	return &Executor{
		llvm:     engine,
		funcPtrs: map[string]unsafe.Pointer{},
	}, nil
}

func (e *Executor) FunctionAddress(f *Function) unsafe.Pointer {
	ptr, ok := e.funcPtrs[f.Name]
	if !ok {
		ptr = e.llvm.PointerToGlobal(f.llvm)
		e.funcPtrs[f.Name] = ptr
	}
	return ptr
}

// SizeOf returns the offset in bytes between successive objects of the
// specified type, including alignment padding.
func (e *Executor) SizeOf(t Type) int {
	return int(e.llvm.TargetData().TypeAllocSize(t.llvmTy()))
}

// AlignOf returns the preferred stack/global alignment for the specified type.
func (e *Executor) AlignOf(t Type) int {
	// TODO: Preferred alignment vs ABI alignment. Which one?
	return e.llvm.TargetData().PrefTypeAlignment(t.llvmTy())
}

func (e *Executor) FieldOffsets(s *Struct) []int {
	td := e.llvm.TargetData()
	out := make([]int, len(s.Fields()))
	for i := range s.Fields() {
		out[i] = int(td.ElementOffset(s.llvm, i))
	}
	return out
}

func (e *Executor) StructLayout(s *Struct) string {
	w := bytes.Buffer{}
	w.WriteString(s.TypeName())
	w.WriteString("{\n")
	e.writeStructLayout(s, &w, 0, "")
	w.WriteString("}")
	return w.String()
}

func (e *Executor) writeStructLayout(s *Struct, w *bytes.Buffer, base int, prefix string) {
	fields := s.Fields()
	for i, o := range e.FieldOffsets(s) {
		f := fields[i]
		w.WriteString(fmt.Sprintf(" 0x%.4x: ", base+o))
		w.WriteString(prefix)
		w.WriteString(f.Name)
		w.WriteRune('\n')
		if s, ok := f.Type.(*Struct); ok {
			e.writeStructLayout(s, w, base+o, prefix+f.Name+".")
		}
	}
}
<commit_msg>core/codegen: Add Executor.GlobalAddress()<commit_after>// Copyright (C) 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 codegen

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"unsafe"

	"github.com/google/gapid/core/app/linker"
	"github.com/google/gapid/core/os/device"

	"llvm/bindings/go/llvm"
)

func init() {
	llvm.InitializeAllTargetInfos()
	llvm.InitializeAllTargets()
	llvm.InitializeAllTargetMCs()
	llvm.InitializeAllAsmPrinters()
}

// Executor executes module functions.
type Executor struct {
	llvm     llvm.ExecutionEngine
	funcPtrs map[string]unsafe.Pointer
}

// Object compiles the module down to an object file.
func (m *Module) Object(optimize bool) ([]byte, error) {
	t, err := llvm.GetTargetFromTriple(m.triple.String())
	if err != nil {
		return nil, fmt.Errorf("Couldn't get target for triple '%v': %v", m.triple, err)
	}
	cpu := ""
	features := ""
	opt := llvm.CodeGenLevelNone
	if optimize {
		opt = llvm.CodeGenLevelDefault
	}
	reloc := llvm.RelocPIC
	model := llvm.CodeModelDefault
	tm := t.CreateTargetMachine(m.triple.String(), cpu, features, opt, reloc, model)
	defer tm.Dispose()

	// Check target data is as expected.
	td := tm.CreateTargetData()
	defer td.Dispose()
	m.validateTargetData(td)

	buf, err := tm.EmitToMemoryBuffer(m.llvm, llvm.ObjectFile)
	if err != nil {
		return nil, err
	}
	defer buf.Dispose()
	return buf.Bytes(), nil
}

func (m *Module) validateTargetData(td llvm.TargetData) {
	abi := m.target
	errs := []string{}
	check := func(llvm, gapid interface{}, name string) bool {
		if reflect.DeepEqual(llvm, gapid) {
			return true
		}
		errs = append(errs, fmt.Sprintf("%v target mismatch for %v: %v (llvm) != %v (gapid)", name, abi.Name, llvm, gapid))
		return false
	}
	checkTD := func(ty Type, dtl *device.DataTypeLayout) {
		check(td.TypeStoreSize(ty.llvmTy()), uint64(dtl.Size), ty.String()+"-size")
		check(td.ABITypeAlignment(ty.llvmTy()), int(dtl.Alignment), ty.String()+"-align")
	}

	layout := abi.MemoryLayout
	isLE := td.ByteOrder() == llvm.LittleEndian
	check(isLE, layout.Endian == device.LittleEndian, "is-little-endian")
	check(td.PointerSize(), int(layout.Pointer.Size), "pointer-size")

	checkTD(m.Types.Pointer(m.Types.Int), layout.Pointer)
	checkTD(m.Types.Int, layout.Integer)
	checkTD(m.Types.Size, layout.Size)
	checkTD(m.Types.Int64, layout.I64)
	checkTD(m.Types.Int32, layout.I32)
	checkTD(m.Types.Int16, layout.I16)
	checkTD(m.Types.Int8, layout.I8)
	checkTD(m.Types.Float32, layout.F32)
	checkTD(m.Types.Float64, layout.F64)

	for _, s := range m.Types.structs {
		if !s.hasBody {
			continue
		}
		if !check(int(td.TypeStoreSize(s.llvm))*8, s.SizeInBits(), fmt.Sprintf("%v-size", s.name)) ||
			!check(int(td.ABITypeAlignment(s.llvm))*8, s.AlignInBits(), fmt.Sprintf("%v-align", s.name)) {
			errs = append(errs, fmt.Sprintf("%v: %v", s.name, s))
		}
		for i := range s.Fields() {
			llvm := int(td.ElementOffset(s.llvm, i)) * 8
			gapid := s.FieldOffsetInBits(i)
			check(llvm, gapid, fmt.Sprintf("%v-field-offset %d", s.name, i))
		}
	}

	for _, s := range m.Types.arrays {
		check(int(td.TypeStoreSize(s.llvm))*8, s.SizeInBits(), fmt.Sprintf("%v-size", s.name))
		check(int(td.ABITypeAlignment(s.llvm))*8, s.AlignInBits(), fmt.Sprintf("%v-align", s.name))
	}

	if len(errs) > 0 {
		panic(fmt.Errorf("%v has ABI mismatches!\n%v", abi.Name, strings.Join(errs, "\n")))
	}
}

// Optimize optimizes the module.
func (m *Module) Optimize() {
	fpm := llvm.NewFunctionPassManagerForModule(m.llvm)
	defer fpm.Dispose()

	mpm := llvm.NewPassManager()
	defer mpm.Dispose()

	pmb := llvm.NewPassManagerBuilder()
	defer pmb.Dispose()

	pmb.SetOptLevel(int(llvm.CodeGenLevelDefault))
	pmb.SetSizeLevel(0)

	mpm.AddVerifierPass()
	fpm.AddVerifierPass()

	pmb.Populate(mpm)
	pmb.PopulateFunc(fpm)

	fpm.InitializeFunc()
	for fn := m.llvm.FirstFunction(); !fn.IsNil(); fn = llvm.NextFunction(fn) {
		fpm.RunFunc(fn)
	}
	fpm.FinalizeFunc()

	mpm.Run(m.llvm)
}

// Executor constructs an executor.
func (m *Module) Executor(optimize bool) (*Executor, error) {
	if err := m.Verify(); err != nil {
		return nil, err
	}

	opts := llvm.NewMCJITCompilerOptions()
	if optimize {
		opts.SetMCJITOptimizationLevel(2)
	} else {
		opts.SetMCJITOptimizationLevel(0)
	}

	engine, err := llvm.NewMCJITCompiler(m.llvm, opts)
	if err != nil {
		return nil, err
	}

	// Check target data is as expected.
	m.validateTargetData(engine.TargetData())

	// Check for unresolved extern symbols.
	var unresolved []string
	for _, f := range m.funcs {
		if f.built || strings.HasPrefix(f.Name, "llvm.") {
			continue
		}
		if linker.ProcAddress(f.Name) == 0 {
			unresolved = append(unresolved, fmt.Sprint(f))
		}
	}
	if len(unresolved) > 0 {
		sort.Strings(unresolved)
		msg := fmt.Sprintf("Unresolved external functions:\n%v", strings.Join(unresolved, "\n"))
		fail(msg)
	}

	engine.RunStaticConstructors()

	return &Executor{
		llvm:     engine,
		funcPtrs: map[string]unsafe.Pointer{},
	}, nil
}

// FunctionAddress returns the address of the function f.
func (e *Executor) FunctionAddress(f *Function) unsafe.Pointer {
	ptr, ok := e.funcPtrs[f.Name]
	if !ok {
		ptr = e.llvm.PointerToGlobal(f.llvm)
		e.funcPtrs[f.Name] = ptr
	}
	return ptr
}

// GlobalAddress returns the address of the global g.
func (e *Executor) GlobalAddress(g Global) unsafe.Pointer {
	return e.llvm.PointerToGlobal(g.llvm)
}

// SizeOf returns the offset in bytes between successive objects of the
// specified type, including alignment padding.
func (e *Executor) SizeOf(t Type) int {
	return int(e.llvm.TargetData().TypeAllocSize(t.llvmTy()))
}

// AlignOf returns the preferred stack/global alignment for the specified type.
func (e *Executor) AlignOf(t Type) int {
	// TODO: Preferred alignment vs ABI alignment. Which one?
	return e.llvm.TargetData().PrefTypeAlignment(t.llvmTy())
}

func (e *Executor) FieldOffsets(s *Struct) []int {
	td := e.llvm.TargetData()
	out := make([]int, len(s.Fields()))
	for i := range s.Fields() {
		out[i] = int(td.ElementOffset(s.llvm, i))
	}
	return out
}

func (e *Executor) StructLayout(s *Struct) string {
	w := bytes.Buffer{}
	w.WriteString(s.TypeName())
	w.WriteString("{\n")
	e.writeStructLayout(s, &w, 0, "")
	w.WriteString("}")
	return w.String()
}

func (e *Executor) writeStructLayout(s *Struct, w *bytes.Buffer, base int, prefix string) {
	fields := s.Fields()
	for i, o := range e.FieldOffsets(s) {
		f := fields[i]
		w.WriteString(fmt.Sprintf(" 0x%.4x: ", base+o))
		w.WriteString(prefix)
		w.WriteString(f.Name)
		w.WriteRune('\n')
		if s, ok := f.Type.(*Struct); ok {
			e.writeStructLayout(s, w, base+o, prefix+f.Name+".")
		}
	}
}
<|endoftext|>


<commit_before>package util

import (
	"strings"
	"unicode/utf16"
)

func ReplaceAll(s, old, _new string) string {
	return strings.Replace(s, old, _new, -1)
}

// utf8 -> utf16
func StringToUtf16(str string) []uint16 {
	runes := []rune(str)
	return utf16.Encode(runes) // func Encode(s []rune) []uint16
}

// utf16 -> utf8
func Utf16ToString(s []uint16) string {
	runes := utf16.Decode(s) // func Decode(s []uint16) []rune
	return string(runes)
}
<commit_msg>rename arg<commit_after>package util

import (
	"strings"
	"unicode/utf16"
)

func ReplaceAll(s, old, _new string) string {
	return strings.Replace(s, old, _new, -1)
}

// utf8 -> utf16
func StringToUtf16(s string) []uint16 {
	runes := []rune(s)
	return utf16.Encode(runes) // func Encode(s []rune) []uint16
}

// utf16 -> utf8
func Utf16ToString(s []uint16) string {
	runes := utf16.Decode(s) // func Decode(s []uint16) []rune
	return string(runes)
}
<|endoftext|>


<commit_before>package core

import "sync"

func KeyValueSource() SourceSpec {
	return SourceSpec{
		Name: "key-value",
		Type: KEY_VALUE,
		New:  NewKeyValue,
	}
}

func NewKeyValue() Source {
	return &KeyValue{
		kv:   make(map[string]Message),
		quit: make(chan bool),
	}
}

func (k KeyValue) GetType() SourceType {
	return KEY_VALUE
}

type KeyValue struct {
	kv   map[string]Message
	quit chan bool
	sync.Mutex
}

func (k KeyValue) Serve() {
	<-k.quit
}

func (k KeyValue) Stop() {
	k.quit <- true
}

func (k KeyValue) SetSourceParameter(key, value string) {
}

func (k *KeyValue) Describe() map[string]string {
	return map[string]string{}
}

// retrieves a value from the key value store
func kvGet() Spec {
	return Spec{
		Name: "kvGet",
		Inputs: []Pin{
			Pin{"key"},
		},
		Outputs: []Pin{
			Pin{"value"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if value, ok := kv.kv[key]; !ok {
				out[0] = NewError("Key not found")
			} else {
				out[0] = value
			}
			return nil
		},
	}
}

// sets an entry in a key value store
// if the entry is new, emits true
func kvSet() Spec {
	return Spec{
		Name: "kvSet",
		Inputs: []Pin{
			Pin{"key"},
			Pin{"value"},
		},
		Outputs: []Pin{
			Pin{"new"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if _, ok := kv.kv[key]; !ok {
				out[0] = true
			} else {
				out[0] = false
			}

			kv.kv[in[0].(string)] = in[1]
			return nil
		},
	}
}

// clears the entire map
// TODO: prefer "empty"
// change interface{} to message
func kvClear() Spec {
	return Spec{
		Name: "kvClear",
		Inputs: []Pin{
			Pin{"clear"},
		},
		Outputs: []Pin{
			Pin{"cleared"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			kv.kv = make(map[string]Message)
			out[0] = true
			return nil
		},
	}
}

// dumps the entire map into a message
// should output be named "object" ?
// TODO: convert interface{} to message
// !! should probably double check this to ensure that we don't need a deep copy
func kvDump() Spec {
	return Spec{
		Name: "kvDump",
		Inputs: []Pin{
			Pin{"dump"},
		},
		Outputs: []Pin{
			Pin{"object"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			outMap := make(map[string]Message)
			for k, v := range kv.kv {
				outMap[k] = v
			}
			out[0] = outMap
			return nil
		},
	}
}

// deletes an entry in a key value store
func kvDelete() Spec {
	return Spec{
		Name: "kvDelete",
		Inputs: []Pin{
			Pin{"key"},
		},
		Outputs: []Pin{
			Pin{"deleted"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if _, ok := kv.kv[key]; !ok {
				out[0] = false
			} else {
				delete(kv.kv, key)
				out[0] = true
			}
			return nil
		},
	}
}
<commit_msg>source key value -> store key value<commit_after>package core

import "sync"

func KeyValueStore() SourceSpec {
	return SourceSpec{
		Name: "key-value",
		Type: KEY_VALUE,
		New:  NewKeyValue,
	}
}

func NewKeyValue() Source {
	return &KeyValue{
		kv:   make(map[string]Message),
		quit: make(chan bool),
	}
}

func (k KeyValue) GetType() SourceType {
	return KEY_VALUE
}

type KeyValue struct {
	kv   map[string]Message
	quit chan bool
	sync.Mutex
}

// retrieves a value from the key value store
func kvGet() Spec {
	return Spec{
		Name: "kvGet",
		Inputs: []Pin{
			Pin{"key"},
		},
		Outputs: []Pin{
			Pin{"value"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if value, ok := kv.kv[key]; !ok {
				out[0] = NewError("Key not found")
			} else {
				out[0] = value
			}
			return nil
		},
	}
}

// sets an entry in a key value store
// if the entry is new, emits true
func kvSet() Spec {
	return Spec{
		Name: "kvSet",
		Inputs: []Pin{
			Pin{"key"},
			Pin{"value"},
		},
		Outputs: []Pin{
			Pin{"new"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if _, ok := kv.kv[key]; !ok {
				out[0] = true
			} else {
				out[0] = false
			}

			kv.kv[in[0].(string)] = in[1]
			return nil
		},
	}
}

// clears the entire map
// TODO: prefer "empty"
// change interface{} to message
func kvClear() Spec {
	return Spec{
		Name: "kvClear",
		Inputs: []Pin{
			Pin{"clear"},
		},
		Outputs: []Pin{
			Pin{"cleared"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			kv.kv = make(map[string]Message)
			out[0] = true
			return nil
		},
	}
}

// dumps the entire map into a message
// should output be named "object" ?
// TODO: convert interface{} to message
// !! should probably double check this to ensure that we don't need a deep copy
func kvDump() Spec {
	return Spec{
		Name: "kvDump",
		Inputs: []Pin{
			Pin{"dump"},
		},
		Outputs: []Pin{
			Pin{"object"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			outMap := make(map[string]Message)
			for k, v := range kv.kv {
				outMap[k] = v
			}
			out[0] = outMap
			return nil
		},
	}
}

// deletes an entry in a key value store
func kvDelete() Spec {
	return Spec{
		Name: "kvDelete",
		Inputs: []Pin{
			Pin{"key"},
		},
		Outputs: []Pin{
			Pin{"deleted"},
		},
		Source: KEY_VALUE,
		Kernel: func(in, out, internal MessageMap, s Source, i chan Interrupt) Interrupt {
			kv := s.(*KeyValue)
			key, ok := in[0].(string)
			if !ok {
				out[0] = NewError("Key is not type string")
				return nil
			}

			if _, ok := kv.kv[key]; !ok {
				out[0] = false
			} else {
				delete(kv.kv, key)
				out[0] = true
			}
			return nil
		},
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The Periph Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Package gpioreg defines a registry for the known digital pins.
package gpioreg

import (
	"errors"
	"fmt"
	"sort"
	"strconv"
	"sync"

	"periph.io/x/periph/conn/gpio"
)

// ByNumber returns a GPIO pin from its number.
//
// Returns nil in case the pin is not present.
func ByNumber(number int) gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	return getByNumber(number)
}

// ByName returns a GPIO pin from its name.
//
// This can be strings like GPIO2, PB8, etc.
//
// This function also parses string representation of numbers, so that calling
// with "6" will return the pin registered as number 6.
//
// Returns nil in case the pin is not present.
func ByName(name string) gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	if p, ok := byName[0][name]; ok {
		return p
	}
	if p, ok := byName[1][name]; ok {
		return p
	}
	if p, ok := byAlias[name]; ok {
		if p.PinIO == nil {
			if p.PinIO = getByNumber(p.number); p.PinIO == nil {
				return nil
			}
		}
		return p
	}
	if i, err := strconv.Atoi(name); err == nil {
		return getByNumber(i)
	}
	return nil
}

// All returns all the GPIO pins available on this host.
//
// The list is guaranteed to be in order of number.
//
// This list excludes aliases.
//
// This list excludes non-GPIO pins like GROUND, V3_3, etc.
func All() []gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	out := make(pinList, 0, len(byNumber))
	seen := make(map[int]struct{}, len(byNumber[0]))
	// Memory-mapped pins have highest priority, include all of them.
	for _, p := range byNumber[0] {
		out = append(out, p)
		seen[p.Number()] = struct{}{}
	}
	// Add in OS accessible pins that cannot be accessed via memory-map.
	for _, p := range byNumber[1] {
		if _, ok := seen[p.Number()]; !ok {
			out = append(out, p)
		}
	}
	sort.Sort(out)
	return out
}

// Aliases returns all pin aliases.
//
// The list is guaranteed to be in order of aliase name.
func Aliases() []gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	out := make(pinList, 0, len(byAlias))
	for _, p := range byAlias {
		// Skip aliases that were not resolved.
		if p.PinIO == nil {
			if p.PinIO = getByNumber(p.number); p.PinIO == nil {
				continue
			}
		}
		out = append(out, p)
	}
	sort.Sort(out)
	return out
}

// Register registers a GPIO pin.
//
// Registering the same pin number or name twice is an error.
//
// `preferred` should be true when the pin being registered is exposing as much
// functionality as possible via the underlying hardware. This is normally done
// by accessing the CPU memory mapped registers directly.
//
// `preferred` should be false when the functionality is provided by the OS and
// is limited or slower.
//
// The pin registered cannot implement the interface RealPin.
func Register(p gpio.PinIO, preferred bool) error {
	name := p.Name()
	if len(name) == 0 {
		return errors.New("gpio: can't register a pin with no name")
	}
	if _, err := strconv.Atoi(name); err == nil {
		return fmt.Errorf("gpio: can't register a pin with a name being only a number %q", name)
	}
	number := p.Number()
	if number < 0 {
		return fmt.Errorf("gpio: can't register a pin with a negative number %d", number)
	}
	i := 0
	if !preferred {
		i = 1
	}

	mu.Lock()
	defer mu.Unlock()
	if orig, ok := byNumber[i][number]; ok {
		return fmt.Errorf("gpio: can't register the same pin %d twice; had %q, registering %q", number, orig, p)
	}
	if orig, ok := byName[i][name]; ok {
		return fmt.Errorf("gpio: can't register the same pin %q twice; had %q, registering %q", name, orig, p)
	}
	if r, ok := p.(gpio.RealPin); ok {
		return fmt.Errorf("gpio: can't register %q, which is an aliased for %q, use RegisterAlias() instead", p, r)
	}
	if alias, ok := byAlias[name]; ok {
		return fmt.Errorf("gpio: can't register %q for which an alias %q already exists", p, alias)
	}
	byNumber[i][number] = p
	byName[i][name] = p
	return nil
}

// RegisterAlias registers an alias for a GPIO pin.
//
// It is possible to register an alias for a pin number that itself has not
// been registered yet.
func RegisterAlias(alias string, number int) error {
	if len(alias) == 0 {
		return errors.New("gpio: can't register an alias with no name")
	}
	if _, err := strconv.Atoi(alias); err == nil {
		return fmt.Errorf("gpio: can't register an alias being only a number %q", alias)
	}
	if number < 0 {
		return fmt.Errorf("gpio: can't register an alias to a pin with a negative number %d", number)
	}

	mu.Lock()
	defer mu.Unlock()
	if orig := byAlias[alias]; orig != nil {
		return fmt.Errorf("gpio: can't register alias %q for pin %d: it is already aliased to %q", alias, number, orig)
	}
	byAlias[alias] = &pinAlias{name: alias, number: number}
	return nil
}

//

var (
	mu sync.Mutex
	// The first map is preferred pins, the second is for more limited pins,
	// usually going through OS-provided abstraction layer.
	byNumber = [2]map[int]gpio.PinIO{{}, {}}
	byName   = [2]map[string]gpio.PinIO{{}, {}}
	byAlias  = map[string]*pinAlias{}
)

// pinAlias implements an alias for a PinIO.
//
// pinAlias also implements the RealPin interface, which allows querying for
// the real pin under the alias.
type pinAlias struct {
	gpio.PinIO
	name   string
	number int
}

// String returns the alias name along the real pin's Name() in parenthesis, if
// known, else the real pin's number.
func (a *pinAlias) String() string {
	if a.PinIO == nil {
		return fmt.Sprintf("%s(%d)", a.name, a.number)
	}
	return fmt.Sprintf("%s(%s)", a.name, a.PinIO.Name())
}

// Name returns the pinAlias's name.
func (a *pinAlias) Name() string {
	return a.name
}

// Real returns the real pin behind the alias
func (a *pinAlias) Real() gpio.PinIO {
	return a.PinIO
}

func getByNumber(number int) gpio.PinIO {
	if p, ok := byNumber[0][number]; ok {
		return p
	}
	if p, ok := byNumber[1][number]; ok {
		return p
	}
	return nil
}

type pinList []gpio.PinIO

func (p pinList) Len() int           { return len(p) }
func (p pinList) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p pinList) Less(i, j int) bool { return p[i].Number() < p[j].Number() }
<commit_msg>gpioreg: refactor errors to be coherent with other registries<commit_after>// Copyright 2017 The Periph Authors. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

// Package gpioreg defines a registry for the known digital pins.
package gpioreg

import (
	"fmt"
	"sort"
	"strconv"
	"sync"

	"periph.io/x/periph/conn/gpio"
)

// ByNumber returns a GPIO pin from its number.
//
// Returns nil in case the pin is not present.
func ByNumber(number int) gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	return getByNumber(number)
}

// ByName returns a GPIO pin from its name.
//
// This can be strings like GPIO2, PB8, etc.
//
// This function also parses string representation of numbers, so that calling
// with "6" will return the pin registered as number 6.
//
// Returns nil in case the pin is not present.
func ByName(name string) gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	if p, ok := byName[0][name]; ok {
		return p
	}
	if p, ok := byName[1][name]; ok {
		return p
	}
	if p, ok := byAlias[name]; ok {
		if p.PinIO == nil {
			if p.PinIO = getByNumber(p.number); p.PinIO == nil {
				return nil
			}
		}
		return p
	}
	if i, err := strconv.Atoi(name); err == nil {
		return getByNumber(i)
	}
	return nil
}

// All returns all the GPIO pins available on this host.
//
// The list is guaranteed to be in order of number.
//
// This list excludes aliases.
//
// This list excludes non-GPIO pins like GROUND, V3_3, etc.
func All() []gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	out := make(pinList, 0, len(byNumber))
	seen := make(map[int]struct{}, len(byNumber[0]))
	// Memory-mapped pins have highest priority, include all of them.
	for _, p := range byNumber[0] {
		out = append(out, p)
		seen[p.Number()] = struct{}{}
	}
	// Add in OS accessible pins that cannot be accessed via memory-map.
	for _, p := range byNumber[1] {
		if _, ok := seen[p.Number()]; !ok {
			out = append(out, p)
		}
	}
	sort.Sort(out)
	return out
}

// Aliases returns all pin aliases.
//
// The list is guaranteed to be in order of aliase name.
func Aliases() []gpio.PinIO {
	mu.Lock()
	defer mu.Unlock()
	out := make(pinList, 0, len(byAlias))
	for _, p := range byAlias {
		// Skip aliases that were not resolved.
		if p.PinIO == nil {
			if p.PinIO = getByNumber(p.number); p.PinIO == nil {
				continue
			}
		}
		out = append(out, p)
	}
	sort.Sort(out)
	return out
}

// Register registers a GPIO pin.
//
// Registering the same pin number or name twice is an error.
//
// `preferred` should be true when the pin being registered is exposing as much
// functionality as possible via the underlying hardware. This is normally done
// by accessing the CPU memory mapped registers directly.
//
// `preferred` should be false when the functionality is provided by the OS and
// is limited or slower.
//
// The pin registered cannot implement the interface RealPin.
func Register(p gpio.PinIO, preferred bool) error {
	name := p.Name()
	if len(name) == 0 {
		return wrapf("can't register a pin with no name")
	}
	if _, err := strconv.Atoi(name); err == nil {
		return wrapf("can't register pin %q with name being only a number", name)
	}
	number := p.Number()
	if number < 0 {
		return wrapf("can't register pin %q with invalid pin number %d", name, number)
	}
	i := 0
	if !preferred {
		i = 1
	}

	mu.Lock()
	defer mu.Unlock()
	if orig, ok := byNumber[i][number]; ok {
		return wrapf("can't register pin %q twice with the same number %d; already registered as %s", name, number, orig)
	}
	if orig, ok := byName[i][name]; ok {
		return wrapf("can't register pin %q twice; already registered as %s", name, orig)
	}
	if r, ok := p.(gpio.RealPin); ok {
		return wrapf("can't register pin %q, it is already an alias: %s; use RegisterAlias() instead", name, r)
	}
	if alias, ok := byAlias[name]; ok {
		return wrapf("can't register pin %q; an alias already exist: %s", name, alias)
	}
	byNumber[i][number] = p
	byName[i][name] = p
	return nil
}

// RegisterAlias registers an alias for a GPIO pin.
//
// It is possible to register an alias for a pin number that itself has not
// been registered yet.
func RegisterAlias(alias string, number int) error {
	if len(alias) == 0 {
		return wrapf("can't register an alias with no name")
	}
	if _, err := strconv.Atoi(alias); err == nil {
		return wrapf("can't register alias %q with name being only a number", alias)
	}
	if number < 0 {
		return wrapf("can't register alias %q with invalid pin number %d", alias, number)
	}

	mu.Lock()
	defer mu.Unlock()
	if orig := byAlias[alias]; orig != nil {
		return wrapf("can't register alias %q twice; it is already an alias: %v", alias, orig)
	}
	byAlias[alias] = &pinAlias{name: alias, number: number}
	return nil
}

//

var (
	mu sync.Mutex
	// The first map is preferred pins, the second is for more limited pins,
	// usually going through OS-provided abstraction layer.
	byNumber = [2]map[int]gpio.PinIO{{}, {}}
	byName   = [2]map[string]gpio.PinIO{{}, {}}
	byAlias  = map[string]*pinAlias{}
)

// pinAlias implements an alias for a PinIO.
//
// pinAlias also implements the RealPin interface, which allows querying for
// the real pin under the alias.
type pinAlias struct {
	gpio.PinIO
	name   string
	number int
}

// String returns the alias name along the real pin's Name() in parenthesis, if
// known, else the real pin's number.
func (a *pinAlias) String() string {
	if a.PinIO == nil {
		return fmt.Sprintf("%s(%d)", a.name, a.number)
	}
	return fmt.Sprintf("%s(%s)", a.name, a.PinIO.Name())
}

// Name returns the pinAlias's name.
func (a *pinAlias) Name() string {
	return a.name
}

// Real returns the real pin behind the alias
func (a *pinAlias) Real() gpio.PinIO {
	return a.PinIO
}

func getByNumber(number int) gpio.PinIO {
	if p, ok := byNumber[0][number]; ok {
		return p
	}
	if p, ok := byNumber[1][number]; ok {
		return p
	}
	return nil
}

// wrapf returns an error that is wrapped with the package name.
func wrapf(format string, a ...interface{}) error {
	return fmt.Errorf("gpioreg: "+format, a...)
}

type pinList []gpio.PinIO

func (p pinList) Len() int           { return len(p) }
func (p pinList) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p pinList) Less(i, j int) bool { return p[i].Number() < p[j].Number() }
<|endoftext|>


<commit_before>package types

import (
	"bytes"
	"errors"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/go-graphite/carbonapi/expr/consolidations"
	"github.com/go-graphite/carbonapi/expr/tags"
	pbv2 "github.com/go-graphite/protocol/carbonapi_v2_pb"
	pb "github.com/go-graphite/protocol/carbonapi_v3_pb"
	pickle "github.com/lomik/og-rek"
)

var (
	// ErrWildcardNotAllowed is an eval error returned when a wildcard/glob argument is found where a single series is required.
	ErrWildcardNotAllowed = errors.New("found wildcard where series expected")
	// ErrTooManyArguments is an eval error returned when too many arguments are provided.
	ErrTooManyArguments = errors.New("too many arguments")
)

// MetricData contains necessary data to represent parsed metric (ready to be send out or drawn)
type MetricData struct {
	pb.FetchResponse

	GraphOptions

	ValuesPerPoint    int
	aggregatedValues  []float64
	Tags              map[string]string
	AggregateFunction func([]float64) float64 `json:"-"`
}

// MarshalCSV marshals metric data to CSV
func MarshalCSV(results []*MetricData) []byte {

	var b []byte

	for _, r := range results {

		step := r.StepTime
		t := r.StartTime
		for _, v := range r.Values {
			b = append(b, "\""+r.Name+"\","+time.Unix(t, 0).UTC().Format("2006-01-02 15:04:05")+","...)
			if !math.IsNaN(v) {
				b = strconv.AppendFloat(b, v, 'f', -1, 64)
			}
			b = append(b, '\n')
			t += step
		}
	}
	return b
}

// ConsolidateJSON consolidates values to maxDataPoints size
func ConsolidateJSON(maxDataPoints int64, results []*MetricData) {
	if len(results) == 0 {
		return
	}
	startTime := results[0].StartTime
	endTime := results[0].StopTime
	for _, r := range results {
		t := r.StartTime
		if startTime > t {
			startTime = t
		}
		t = r.StopTime
		if endTime < t {
			endTime = t
		}
	}

	timeRange := endTime - startTime

	if timeRange <= 0 {
		return
	}

	for _, r := range results {
		numberOfDataPoints := math.Floor(float64(timeRange) / float64(r.StepTime))
		if numberOfDataPoints > float64(maxDataPoints) {
			valuesPerPoint := math.Ceil(numberOfDataPoints / float64(maxDataPoints))
			r.SetValuesPerPoint(int(valuesPerPoint))
		}
	}
}

// MarshalJSON marshals metric data to JSON
func MarshalJSON(results []*MetricData, timestampMultiplier int64, noNullPoints bool) []byte {
	var b []byte
	b = append(b, '[')

	var topComma bool
	for _, r := range results {
		if r == nil {
			continue
		}

		if topComma {
			b = append(b, ',')
		}
		topComma = true

		b = append(b, `{"target":`...)
		b = strconv.AppendQuoteToASCII(b, r.Name)
		b = append(b, `,"datapoints":[`...)

		var innerComma bool
		t := r.StartTime * timestampMultiplier
		for _, v := range r.AggregatedValues() {
			if noNullPoints && math.IsNaN(v) {
				t += r.AggregatedTimeStep() * timestampMultiplier
			} else {
				if innerComma {
					b = append(b, ',')
				}
				innerComma = true

				b = append(b, '[')

				if math.IsNaN(v) || math.IsInf(v, 1) || math.IsInf(v, -1) {
					b = append(b, "null"...)
				} else {
					b = strconv.AppendFloat(b, v, 'f', -1, 64)
				}

				b = append(b, ',')

				b = strconv.AppendInt(b, t, 10)

				b = append(b, ']')

				t += r.AggregatedTimeStep() * timestampMultiplier
			}
		}

		b = append(b, `],"tags":{`...)
		notFirstTag := false
		responseTags := make([]string, 0, len(r.Tags))
		for tag := range r.Tags {
			responseTags = append(responseTags, tag)
		}
		sort.Strings(responseTags)
		for _, tag := range responseTags {
			v := r.Tags[tag]
			if notFirstTag {
				b = append(b, ',')
			}
			b = strconv.AppendQuoteToASCII(b, tag)
			b = append(b, ':')
			b = strconv.AppendQuoteToASCII(b, v)
			notFirstTag = true
		}

		b = append(b, `}}`...)
	}

	b = append(b, ']')

	return b
}

// MarshalPickle marshals metric data to pickle format
func MarshalPickle(results []*MetricData) []byte {

	var p []map[string]interface{}

	for _, r := range results {
		values := make([]interface{}, len(r.Values))
		for i, v := range r.Values {
			if math.IsNaN(v) {
				values[i] = pickle.None{}
			} else {
				values[i] = v
			}

		}
		p = append(p, map[string]interface{}{
			"name":              r.Name,
			"pathExpression":    r.PathExpression,
			"consolidationFunc": r.ConsolidationFunc,
			"start":             r.StartTime,
			"end":               r.StopTime,
			"step":              r.StepTime,
			"xFilesFactor":      r.XFilesFactor,
			"values":            values,
		})
	}

	var buf bytes.Buffer

	penc := pickle.NewEncoder(&buf)
	_ = penc.Encode(p)

	return buf.Bytes()
}

// MarshalProtobufV3 marshals metric data to protobuf
func MarshalProtobufV2(results []*MetricData) ([]byte, error) {
	response := pbv2.MultiFetchResponse{}
	for _, metric := range results {
		fmv3 := (*metric).FetchResponse
		v := make([]float64, len(fmv3.Values))
		isAbsent := make([]bool, len(fmv3.Values))
		for i := range fmv3.Values {
			if math.IsNaN(fmv3.Values[i]) {
				v[i] = 0
				isAbsent[i] = true
			} else {
				v[i] = fmv3.Values[i]
			}
		}
		fm := pbv2.FetchResponse{
			Name:      fmv3.Name,
			StartTime: int32(fmv3.StartTime),
			StopTime:  int32(fmv3.StopTime),
			StepTime:  int32(fmv3.StepTime),
			Values:    v,
			IsAbsent:  isAbsent,
		}
		response.Metrics = append(response.Metrics, fm)
	}
	b, err := response.Marshal()
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MarshalProtobufV3 marshals metric data to protobuf
func MarshalProtobufV3(results []*MetricData) ([]byte, error) {
	response := pb.MultiFetchResponse{}
	for _, metric := range results {
		response.Metrics = append(response.Metrics, (*metric).FetchResponse)
	}
	b, err := response.Marshal()
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MarshalRaw marshals metric data to graphite's internal format, called 'raw'
func MarshalRaw(results []*MetricData) []byte {

	var b []byte

	for _, r := range results {

		b = append(b, r.Name...)

		b = append(b, ',')
		b = strconv.AppendInt(b, r.StartTime, 10)
		b = append(b, ',')
		b = strconv.AppendInt(b, r.StopTime, 10)
		b = append(b, ',')
		b = strconv.AppendInt(b, r.StepTime, 10)
		b = append(b, '|')

		var comma bool
		for _, v := range r.Values {
			if comma {
				b = append(b, ',')
			}
			comma = true
			if math.IsNaN(v) {
				b = append(b, "None"...)
			} else {
				b = strconv.AppendFloat(b, v, 'f', -1, 64)
			}
		}

		b = append(b, '\n')
	}
	return b
}

// SetValuesPerPoint sets value per point coefficient.
func (r *MetricData) SetValuesPerPoint(v int) {
	r.ValuesPerPoint = v
	r.aggregatedValues = nil
}

// AggregatedTimeStep aggregates time step
func (r *MetricData) AggregatedTimeStep() int64 {
	if r.ValuesPerPoint == 1 || r.ValuesPerPoint == 0 {
		return r.StepTime
	}

	return r.StepTime * int64(r.ValuesPerPoint)
}

// GetAggregateFunction returns MetricData.AggregateFunction and set it, if it's not yet
func (r *MetricData) GetAggregateFunction() func([]float64) float64 {
	if r.AggregateFunction == nil {
		var ok bool
		if r.AggregateFunction, ok = consolidations.ConsolidationToFunc[strings.ToLower(r.ConsolidationFunc)]; !ok {
			// if consolidation function is not known, we should fall back to average
			r.AggregateFunction = consolidations.AvgValue
		}
	}

	return r.AggregateFunction
}

// AggregatedValues aggregates values (with cache)
func (r *MetricData) AggregatedValues() []float64 {
	if r.aggregatedValues == nil {
		r.AggregateValues()
	}
	return r.aggregatedValues
}

// AggregateValues aggregates values
func (r *MetricData) AggregateValues() {
	if r.ValuesPerPoint == 1 || r.ValuesPerPoint == 0 {
		r.aggregatedValues = make([]float64, len(r.Values))
		copy(r.aggregatedValues, r.Values)
		return
	}
	aggFunc := r.GetAggregateFunction()

	n := len(r.Values)/r.ValuesPerPoint + 1
	aggV := make([]float64, 0, n)

	v := r.Values

	for len(v) >= r.ValuesPerPoint {
		val := aggFunc(v[:r.ValuesPerPoint])
		aggV = append(aggV, val)
		v = v[r.ValuesPerPoint:]
	}

	if len(v) > 0 {
		val := aggFunc(v)
		aggV = append(aggV, val)
	}

	r.aggregatedValues = aggV
}

// Copy returns the copy of r. If includeValues set to true, it copies values as well.
func (r *MetricData) Copy(includeValues bool) *MetricData {
	var values, aggregatedValues []float64
	values = make([]float64, 0)
	appliedFunctions := make([]string, 0)
	aggregatedValues = nil

	if includeValues {
		values = make([]float64, len(r.Values))
		copy(values, r.Values)

		if r.aggregatedValues != nil {
			aggregatedValues = make([]float64, len(r.aggregatedValues))
			copy(aggregatedValues, r.aggregatedValues)
		}

		appliedFunctions = make([]string, len(r.AppliedFunctions))
		copy(appliedFunctions, r.AppliedFunctions)
	}

	tags := make(map[string]string)
	for k, v := range r.Tags {
		tags[k] = v
	}

	return &MetricData{
		FetchResponse: pb.FetchResponse{
			Name:                    r.Name,
			PathExpression:          r.PathExpression,
			ConsolidationFunc:       r.ConsolidationFunc,
			StartTime:               r.StartTime,
			StopTime:                r.StopTime,
			StepTime:                r.StepTime,
			XFilesFactor:            r.XFilesFactor,
			HighPrecisionTimestamps: r.HighPrecisionTimestamps,
			Values:                  values,
			AppliedFunctions:        appliedFunctions,
			RequestStartTime:        r.RequestStartTime,
			RequestStopTime:         r.RequestStopTime,
		},
		GraphOptions:      r.GraphOptions,
		ValuesPerPoint:    r.ValuesPerPoint,
		aggregatedValues:  aggregatedValues,
		Tags:              tags,
		AggregateFunction: r.AggregateFunction,
	}
}

// CopyMetricDataSlice returns the slice of metrics that should be changed later.
// It allows to avoid a changing of source data, e.g. by AlignMetrics
func CopyMetricDataSlice(args []*MetricData) (newData []*MetricData) {
	newData = make([]*MetricData, len(args))
	for i, m := range args {
		newData[i] = m.Copy(true)
	}
	return newData
}

// MakeMetricData creates new metrics data with given metric timeseries
func MakeMetricData(name string, values []float64, step, start int64) *MetricData {
	return makeMetricDataWithTags(name, values, step, start, tags.ExtractTags(name))
}

// MakeMetricDataWithTags creates new metrics data with given metric Time Series (with tags)
func makeMetricDataWithTags(name string, values []float64, step, start int64, tags map[string]string) *MetricData {
	stop := start + int64(len(values))*step

	return &MetricData{
		FetchResponse: pb.FetchResponse{
			Name:      name,
			Values:    values,
			StartTime: start,
			StepTime:  step,
			StopTime:  stop,
		},
		Tags: tags,
	}
}
<commit_msg>MetricData: CopyLink for copy tags and leave origin values slice<commit_after>package types

import (
	"bytes"
	"errors"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/go-graphite/carbonapi/expr/consolidations"
	"github.com/go-graphite/carbonapi/expr/tags"
	pbv2 "github.com/go-graphite/protocol/carbonapi_v2_pb"
	pb "github.com/go-graphite/protocol/carbonapi_v3_pb"
	pickle "github.com/lomik/og-rek"
)

var (
	// ErrWildcardNotAllowed is an eval error returned when a wildcard/glob argument is found where a single series is required.
	ErrWildcardNotAllowed = errors.New("found wildcard where series expected")
	// ErrTooManyArguments is an eval error returned when too many arguments are provided.
	ErrTooManyArguments = errors.New("too many arguments")
)

// MetricData contains necessary data to represent parsed metric (ready to be send out or drawn)
type MetricData struct {
	pb.FetchResponse

	GraphOptions

	ValuesPerPoint    int
	aggregatedValues  []float64
	Tags              map[string]string
	AggregateFunction func([]float64) float64 `json:"-"`
}

// MarshalCSV marshals metric data to CSV
func MarshalCSV(results []*MetricData) []byte {

	var b []byte

	for _, r := range results {

		step := r.StepTime
		t := r.StartTime
		for _, v := range r.Values {
			b = append(b, "\""+r.Name+"\","+time.Unix(t, 0).UTC().Format("2006-01-02 15:04:05")+","...)
			if !math.IsNaN(v) {
				b = strconv.AppendFloat(b, v, 'f', -1, 64)
			}
			b = append(b, '\n')
			t += step
		}
	}
	return b
}

// ConsolidateJSON consolidates values to maxDataPoints size
func ConsolidateJSON(maxDataPoints int64, results []*MetricData) {
	if len(results) == 0 {
		return
	}
	startTime := results[0].StartTime
	endTime := results[0].StopTime
	for _, r := range results {
		t := r.StartTime
		if startTime > t {
			startTime = t
		}
		t = r.StopTime
		if endTime < t {
			endTime = t
		}
	}

	timeRange := endTime - startTime

	if timeRange <= 0 {
		return
	}

	for _, r := range results {
		numberOfDataPoints := math.Floor(float64(timeRange) / float64(r.StepTime))
		if numberOfDataPoints > float64(maxDataPoints) {
			valuesPerPoint := math.Ceil(numberOfDataPoints / float64(maxDataPoints))
			r.SetValuesPerPoint(int(valuesPerPoint))
		}
	}
}

// MarshalJSON marshals metric data to JSON
func MarshalJSON(results []*MetricData, timestampMultiplier int64, noNullPoints bool) []byte {
	var b []byte
	b = append(b, '[')

	var topComma bool
	for _, r := range results {
		if r == nil {
			continue
		}

		if topComma {
			b = append(b, ',')
		}
		topComma = true

		b = append(b, `{"target":`...)
		b = strconv.AppendQuoteToASCII(b, r.Name)
		b = append(b, `,"datapoints":[`...)

		var innerComma bool
		t := r.StartTime * timestampMultiplier
		for _, v := range r.AggregatedValues() {
			if noNullPoints && math.IsNaN(v) {
				t += r.AggregatedTimeStep() * timestampMultiplier
			} else {
				if innerComma {
					b = append(b, ',')
				}
				innerComma = true

				b = append(b, '[')

				if math.IsNaN(v) || math.IsInf(v, 1) || math.IsInf(v, -1) {
					b = append(b, "null"...)
				} else {
					b = strconv.AppendFloat(b, v, 'f', -1, 64)
				}

				b = append(b, ',')

				b = strconv.AppendInt(b, t, 10)

				b = append(b, ']')

				t += r.AggregatedTimeStep() * timestampMultiplier
			}
		}

		b = append(b, `],"tags":{`...)
		notFirstTag := false
		responseTags := make([]string, 0, len(r.Tags))
		for tag := range r.Tags {
			responseTags = append(responseTags, tag)
		}
		sort.Strings(responseTags)
		for _, tag := range responseTags {
			v := r.Tags[tag]
			if notFirstTag {
				b = append(b, ',')
			}
			b = strconv.AppendQuoteToASCII(b, tag)
			b = append(b, ':')
			b = strconv.AppendQuoteToASCII(b, v)
			notFirstTag = true
		}

		b = append(b, `}}`...)
	}

	b = append(b, ']')

	return b
}

// MarshalPickle marshals metric data to pickle format
func MarshalPickle(results []*MetricData) []byte {

	var p []map[string]interface{}

	for _, r := range results {
		values := make([]interface{}, len(r.Values))
		for i, v := range r.Values {
			if math.IsNaN(v) {
				values[i] = pickle.None{}
			} else {
				values[i] = v
			}

		}
		p = append(p, map[string]interface{}{
			"name":              r.Name,
			"pathExpression":    r.PathExpression,
			"consolidationFunc": r.ConsolidationFunc,
			"start":             r.StartTime,
			"end":               r.StopTime,
			"step":              r.StepTime,
			"xFilesFactor":      r.XFilesFactor,
			"values":            values,
		})
	}

	var buf bytes.Buffer

	penc := pickle.NewEncoder(&buf)
	_ = penc.Encode(p)

	return buf.Bytes()
}

// MarshalProtobufV3 marshals metric data to protobuf
func MarshalProtobufV2(results []*MetricData) ([]byte, error) {
	response := pbv2.MultiFetchResponse{}
	for _, metric := range results {
		fmv3 := (*metric).FetchResponse
		v := make([]float64, len(fmv3.Values))
		isAbsent := make([]bool, len(fmv3.Values))
		for i := range fmv3.Values {
			if math.IsNaN(fmv3.Values[i]) {
				v[i] = 0
				isAbsent[i] = true
			} else {
				v[i] = fmv3.Values[i]
			}
		}
		fm := pbv2.FetchResponse{
			Name:      fmv3.Name,
			StartTime: int32(fmv3.StartTime),
			StopTime:  int32(fmv3.StopTime),
			StepTime:  int32(fmv3.StepTime),
			Values:    v,
			IsAbsent:  isAbsent,
		}
		response.Metrics = append(response.Metrics, fm)
	}
	b, err := response.Marshal()
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MarshalProtobufV3 marshals metric data to protobuf
func MarshalProtobufV3(results []*MetricData) ([]byte, error) {
	response := pb.MultiFetchResponse{}
	for _, metric := range results {
		response.Metrics = append(response.Metrics, (*metric).FetchResponse)
	}
	b, err := response.Marshal()
	if err != nil {
		return nil, err
	}

	return b, nil
}

// MarshalRaw marshals metric data to graphite's internal format, called 'raw'
func MarshalRaw(results []*MetricData) []byte {

	var b []byte

	for _, r := range results {

		b = append(b, r.Name...)

		b = append(b, ',')
		b = strconv.AppendInt(b, r.StartTime, 10)
		b = append(b, ',')
		b = strconv.AppendInt(b, r.StopTime, 10)
		b = append(b, ',')
		b = strconv.AppendInt(b, r.StepTime, 10)
		b = append(b, '|')

		var comma bool
		for _, v := range r.Values {
			if comma {
				b = append(b, ',')
			}
			comma = true
			if math.IsNaN(v) {
				b = append(b, "None"...)
			} else {
				b = strconv.AppendFloat(b, v, 'f', -1, 64)
			}
		}

		b = append(b, '\n')
	}
	return b
}

// SetValuesPerPoint sets value per point coefficient.
func (r *MetricData) SetValuesPerPoint(v int) {
	r.ValuesPerPoint = v
	r.aggregatedValues = nil
}

// AggregatedTimeStep aggregates time step
func (r *MetricData) AggregatedTimeStep() int64 {
	if r.ValuesPerPoint == 1 || r.ValuesPerPoint == 0 {
		return r.StepTime
	}

	return r.StepTime * int64(r.ValuesPerPoint)
}

// GetAggregateFunction returns MetricData.AggregateFunction and set it, if it's not yet
func (r *MetricData) GetAggregateFunction() func([]float64) float64 {
	if r.AggregateFunction == nil {
		var ok bool
		if r.AggregateFunction, ok = consolidations.ConsolidationToFunc[strings.ToLower(r.ConsolidationFunc)]; !ok {
			// if consolidation function is not known, we should fall back to average
			r.AggregateFunction = consolidations.AvgValue
		}
	}

	return r.AggregateFunction
}

// AggregatedValues aggregates values (with cache)
func (r *MetricData) AggregatedValues() []float64 {
	if r.aggregatedValues == nil {
		r.AggregateValues()
	}
	return r.aggregatedValues
}

// AggregateValues aggregates values
func (r *MetricData) AggregateValues() {
	if r.ValuesPerPoint == 1 || r.ValuesPerPoint == 0 {
		r.aggregatedValues = make([]float64, len(r.Values))
		copy(r.aggregatedValues, r.Values)
		return
	}
	aggFunc := r.GetAggregateFunction()

	n := len(r.Values)/r.ValuesPerPoint + 1
	aggV := make([]float64, 0, n)

	v := r.Values

	for len(v) >= r.ValuesPerPoint {
		val := aggFunc(v[:r.ValuesPerPoint])
		aggV = append(aggV, val)
		v = v[r.ValuesPerPoint:]
	}

	if len(v) > 0 {
		val := aggFunc(v)
		aggV = append(aggV, val)
	}

	r.aggregatedValues = aggV
}

// Copy returns the copy of r. If includeValues set to true, it copies values as well.
func (r *MetricData) Copy(includeValues bool) *MetricData {
	var values, aggregatedValues []float64
	values = make([]float64, 0)
	appliedFunctions := make([]string, 0)
	aggregatedValues = nil

	if includeValues {
		values = make([]float64, len(r.Values))
		copy(values, r.Values)

		if r.aggregatedValues != nil {
			aggregatedValues = make([]float64, len(r.aggregatedValues))
			copy(aggregatedValues, r.aggregatedValues)
		}

		appliedFunctions = make([]string, len(r.AppliedFunctions))
		copy(appliedFunctions, r.AppliedFunctions)
	}

	tags := make(map[string]string)
	for k, v := range r.Tags {
		tags[k] = v
	}

	return &MetricData{
		FetchResponse: pb.FetchResponse{
			Name:                    r.Name,
			PathExpression:          r.PathExpression,
			ConsolidationFunc:       r.ConsolidationFunc,
			StartTime:               r.StartTime,
			StopTime:                r.StopTime,
			StepTime:                r.StepTime,
			XFilesFactor:            r.XFilesFactor,
			HighPrecisionTimestamps: r.HighPrecisionTimestamps,
			Values:                  values,
			AppliedFunctions:        appliedFunctions,
			RequestStartTime:        r.RequestStartTime,
			RequestStopTime:         r.RequestStopTime,
		},
		GraphOptions:      r.GraphOptions,
		ValuesPerPoint:    r.ValuesPerPoint,
		aggregatedValues:  aggregatedValues,
		Tags:              tags,
		AggregateFunction: r.AggregateFunction,
	}
}

// Copy returns the copy of r. Values not copied and link from parent.
func (r *MetricData) CopyLink() *MetricData {
	tags := make(map[string]string)
	for k, v := range r.Tags {
		tags[k] = v
	}

	return &MetricData{
		FetchResponse: pb.FetchResponse{
			Name:                    r.Name,
			PathExpression:          r.PathExpression,
			ConsolidationFunc:       r.ConsolidationFunc,
			StartTime:               r.StartTime,
			StopTime:                r.StopTime,
			StepTime:                r.StepTime,
			XFilesFactor:            r.XFilesFactor,
			HighPrecisionTimestamps: r.HighPrecisionTimestamps,
			Values:                  r.Values,
			AppliedFunctions:        r.AppliedFunctions,
			RequestStartTime:        r.RequestStartTime,
			RequestStopTime:         r.RequestStopTime,
		},
		GraphOptions:      r.GraphOptions,
		ValuesPerPoint:    r.ValuesPerPoint,
		aggregatedValues:  r.aggregatedValues,
		Tags:              tags,
		AggregateFunction: r.AggregateFunction,
	}
}

// CopyMetricDataSlice returns the slice of metrics that should be changed later.
// It allows to avoid a changing of source data, e.g. by AlignMetrics
func CopyMetricDataSlice(args []*MetricData) (newData []*MetricData) {
	newData = make([]*MetricData, len(args))
	for i, m := range args {
		newData[i] = m.Copy(true)
	}
	return newData
}

// MakeMetricData creates new metrics data with given metric timeseries
func MakeMetricData(name string, values []float64, step, start int64) *MetricData {
	return makeMetricDataWithTags(name, values, step, start, tags.ExtractTags(name))
}

// MakeMetricDataWithTags creates new metrics data with given metric Time Series (with tags)
func makeMetricDataWithTags(name string, values []float64, step, start int64, tags map[string]string) *MetricData {
	stop := start + int64(len(values))*step

	return &MetricData{
		FetchResponse: pb.FetchResponse{
			Name:      name,
			Values:    values,
			StartTime: start,
			StepTime:  step,
			StopTime:  stop,
		},
		Tags: tags,
	}
}
<|endoftext|>


<commit_before>package gopisysfs

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"time"
)

const (
	sys_i2c   = "sys/class/i2c-dev"
	i2c_SLAVE = 0x703
)

func ListI2CDevs() ([]string, error) {
	devdir := file(sys_i2c)
	files, err := ioutil.ReadDir(devdir)
	if err != nil {
		return nil, err
	}

	names := []string{}
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		name := f.Name()
		dev := filepath.Join("/dev", name)
		info("I2C Checking %v\n", dev)
		if _, err := os.Stat(dev); err != nil {
			continue
		}
		names = append(names, dev)
	}
	return names, nil
}

type Recording struct {
	Timestamp time.Time
	Data      []byte
}

func copyBytes(buffer []byte, count int) []byte {
	ret := make([]byte, count)
	copy(ret, buffer)
	return ret
}

func PollI2C(dev string, address int, bytes int, interval time.Duration) (<-chan Recording, func(), error) {

	ctrl, err := os.OpenFile(dev, os.O_RDWR, 0)
	if err != nil {
		return nil, nil, err
	}

	killer := make(chan bool, 1)

	termfn := func() {
		killer <- true
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(ctrl.Fd()), i2c_SLAVE, uintptr(address))
	if errno != 0 {
		return nil, nil, errno
	}

	buffer := make([]byte, bytes)
	n, err := ctrl.Read(buffer)
	if err != nil {
		ctrl.Close()
		return nil, nil, err
	}

	// single slot channel
	data := make(chan Recording, 1)

	data <- Recording{time.Now(), copyBytes(buffer, n)}

	go func() {
		defer close(data)
		defer ctrl.Close()

		tick := time.NewTicker(interval)
		defer tick.Stop()

		var stamp time.Time

		for {
			select {
			case <-killer:
				return
			case stamp = <-tick.C:
			}
			n, err := ctrl.Read(buffer)
			if err != nil {
				info("I2C Unexpected error reading %v: %v\n", dev, err)
				return
			}
			data <- Recording{stamp, copyBytes(buffer, n)}
		}
	}()

	return data, termfn, nil

}
<commit_msg>Better handling of channel selection<commit_after>package gopisysfs

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"syscall"
	"time"
)

const (
	sys_i2c   = "sys/class/i2c-dev"
	i2c_SLAVE = 0x703
)

func I2CListDevices() ([]string, error) {
	devdir := file(sys_i2c)
	files, err := ioutil.ReadDir(devdir)
	if err != nil {
		return nil, err
	}

	names := []string{}
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		name := f.Name()
		dev := filepath.Join("/dev", name)
		info("I2C Checking %v\n", dev)
		if _, err := os.Stat(dev); err != nil {
			continue
		}
		names = append(names, dev)
	}
	return names, nil
}

type I2CRecording struct {
	Timestamp time.Time
	Data      []byte
}

func copyBytes(buffer []byte, count int) []byte {
	ret := make([]byte, count)
	copy(ret, buffer)
	return ret
}

// I2CPoll establishes a connection to a slave I2C device and periodically reads a fixed number of bytes from that device.
// The dev, address, and bytes parameters indicates which device to read and how much to read each time.
// The bufferdepth determines how deep the returned channel's buffer is.
// All samples taken after the buffer is filled will be discarded until space is available.
// An unbuffered return is supported, and guarantees that a receive on that channel gets the most recent sample.
// The interval indicates the period to sample at.
// The returned channel will be closed if there's an error reading the device or the poller is closed using the returned termination function.
// Call the termination function returned when you no longer need to receive polling data.
func I2CPoll(dev string, address int, bytes int, bufferdepth int, interval time.Duration) (<-chan I2CRecording, func(), error) {

	ctrl, err := os.OpenFile(dev, os.O_RDWR, 0)
	if err != nil {
		return nil, nil, err
	}

	killer := make(chan bool, 1)

	termfn := func() {
		killer <- true
	}

	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(ctrl.Fd()), i2c_SLAVE, uintptr(address))
	if errno != 0 {
		return nil, nil, errno
	}

	buffer := make([]byte, bytes)
	n, err := ctrl.Read(buffer)
	if err != nil {
		ctrl.Close()
		return nil, nil, err
	}

	// unbuffered channel - reader only gets data when asking to receive it, and they get the most recently available value.
	data := make(chan Recording, 0)
	record := I2CRecording{time.Now(), copyBytes(buffer, n)}

	go func() {
		defer close(data)
		defer ctrl.Close()

		tick := time.NewTicker(interval)
		defer tick.Stop()

		var stamp time.Time

		// we do some nil channel tricks to manipulate the select statement. dest is part of that.
		dest := data

		for {
			select {
			case <-killer:
				return
			case dest <- record:
				// disable dest until there's a new record.
				dest = nil
			case stamp = <-tick.C:
				n, err := ctrl.Read(buffer)
				if err != nil {
					info("I2C Unexpected error reading %v: %v\n", dev, err)
					return
				}
				record = I2CRecording{stamp, copyBytes(buffer, n)}
				// indicate there's data to send and reenable dest.
				dest = data
			}
		}
	}()

	return data, termfn, nil

}
<|endoftext|>


<commit_before>package jwt

import (
	"testing"
	"time"

	gojwt "github.com/dgrijalva/jwt-go"
)

const (
	testSecret   = "testSecret"
	testIdentity = "testIdentity"
	testDuration = time.Hour * 20
)

var (
	testMethod = gojwt.SigningMethodHS256
)

func TestValidation(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	_, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	if err != nil {
		t.Error(err)
	}
}

func TestRefreshFlagError(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, RefreshToken, testIdentity)
	_, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	if err == nil {
		t.Error("Error refresh flag parsing not detected")
	}
}

func TestExpired(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	if err != nil {
		t.Error(err)
	}

	if _, ok := token.Claims[expiredKey]; ok &&
		getExpiredFromClaims(token.Claims, expiredKey) != time.Now().Add(testDuration).Unix() {
		t.Error("Duration Error")
	}
}

func TestIdentity(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	if err != nil {
		t.Error(err)
	}

	if _, ok := token.Claims[identityKey]; ok && testIdentity != token.Claims[identityKey] {
		t.Error("Read identity Error")
	}
}

func TestRefresh(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	if err != nil {
		t.Error(err)
	}

	if getTokenTypeFromClaims(token.Claims, tokenTypeKey) != AccessToken {
		t.Error("Read refresh flag from token error")
	}
}
<commit_msg>testify<commit_after>package jwt

import (
	"testing"
	"time"

	gojwt "github.com/dgrijalva/jwt-go"
	"github.com/stretchr/testify/assert"
)

const (
	testSecret   = "testSecret"
	testIdentity = "testIdentity"
	testDuration = time.Hour * 20
)

var (
	testMethod = gojwt.SigningMethodHS256
)

func TestValidation(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	_, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	assert.NoError(t, err)
}

func TestRefreshFlagError(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, RefreshToken, testIdentity)
	_, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	assert.Error(t, err, "Error refresh flag parsing not detected")
}

func TestExpired(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	assert.NoError(t, err)

	if _, ok := token.Claims[expiredKey]; ok &&
		getExpiredFromClaims(token.Claims, expiredKey) != time.Now().Add(testDuration).Unix() {
		t.Error("Duration Error")
	}
}

func TestIdentity(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	assert.NoError(t, err)

	if _, ok := token.Claims[identityKey]; ok && testIdentity != token.Claims[identityKey] {
		t.Error("Read identity Error")
	}
}

func TestRefresh(t *testing.T) {
	tokenString, _ := encodeToken(testSecret, testMethod, testDuration, AccessToken, testIdentity)
	token, err := decodeToken(testSecret, testMethod, AccessToken, tokenString)

	assert.NoError(t, err)
	assert.NotEqual(t, getTokenTypeFromClaims(token.Claims, tokenTypeKey),
		AccessToken, "Read refresh flag from token error")
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"fmt"
	"os"
	"time"

	"github.com/keybase/go-updater"
	"github.com/keybase/go-updater/command"
)

// validCodeSigningKIDs are the list of valid code signing IDs for saltpack verify
var validCodeSigningKIDs = map[string]bool{
	"9092ae4e790763dc7343851b977930f35b16cf43ab0ad900a2af3d3ad5cea1a1": true, // keybot (device)
	"d3458bbecdfc0d0ae39fec05722c6e3e897c169223835977a8aa208dfcd902d3": true, // max (device, home)
	"65ae849d1949a8b0021b165b0edaf722e2a7a9036e07817e056e2d721bddcc0e": true, // max (paper key, cry glass)
	"3a5a45c545ef4f661b8b7573711aaecee3fd5717053484a3a3e725cd68abaa5a": true, // chris (device, ccpro)
	"03d86864fb20e310590042ad3d5492c3f5d06728620175b03c717c211bfaccc2": true, // chris (paper key, clay harbor)
}

// Log is the logging interface for the keybase package
type Log interface {
	Debug(...interface{})
	Info(...interface{})
	Debugf(s string, args ...interface{})
	Infof(s string, args ...interface{})
	Warningf(s string, args ...interface{})
	Errorf(s string, args ...interface{})
}

// context is an updater.Context implementation
type context struct {
	// config is updater config
	config Config
	// log is the logger
	log Log
}

// endpoints define all the url locations for reporting, etc
type endpoints struct {
	update  string
	action  string
	success string
	err     string
}

var defaultEndpoints = endpoints{
	update:  "https://api.keybase.io/_/api/1.0/pkg/update.json",
	action:  "https://api.keybase.io/_/api/1.0/pkg/act.json",
	success: "https://api.keybase.io/_/api/1.0/pkg/success.json",
	err:     "https://api.keybase.io/_/api/1.0/pkg/error.json",
}

func newContext(cfg Config, log Log) *context {
	ctx := context{
		config: cfg,
		log:    log,
	}
	return &ctx
}

// NewUpdaterContext returns an updater context for Keybase
func NewUpdaterContext(pathToKeybase string, log Log) (updater.Context, *updater.Updater) {
	cfg, err := newConfig("Keybase", pathToKeybase, log)
	if err != nil {
		log.Warningf("Error loading config for context: %s", err)
	}

	src := NewUpdateSource(cfg, log)
	// For testing
	// (cd /Applications; ditto -c -k --sequesterRsrc --keepParent Keybase.app /tmp/Keybase.zip)
	//src := updater.NewLocalUpdateSource("/tmp/Keybase.zip", log)
	upd := updater.NewUpdater(src, &cfg, log)
	return newContext(&cfg, log), upd
}

// UpdateOptions returns update options
func (c *context) UpdateOptions() updater.UpdateOptions {
	return c.config.updaterOptions()
}

// GetUpdateUI returns Update UI
func (c *context) GetUpdateUI() updater.UpdateUI {
	return c
}

// GetLog returns log
func (c context) GetLog() Log {
	return c.log
}

// Verify verifies the signature
func (c context) Verify(update updater.Update) error {
	return updater.SaltpackVerifyDetachedFileAtPath(update.Asset.LocalPath, update.Asset.Signature, validCodeSigningKIDs, c.log)
}

type checkInUseResult struct {
	InUse bool `json:"in_use"`
}

func (c context) checkInUse() (bool, error) {
	var result checkInUseResult
	if err := command.ExecForJSON(c.config.keybasePath(), []string{"update", "check-in-use"}, &result, time.Minute, c.log); err != nil {
		return false, err
	}
	return result.InUse, nil
}

// BeforeApply is called before an update is applied
func (c context) BeforeApply(update updater.Update) error {
	inUse, err := c.checkInUse()
	if err != nil {
		c.log.Warningf("Error trying to check in use: %s", err)
	}
	if inUse {
		if cancel := c.PausedPrompt(); cancel {
			return fmt.Errorf("Canceled by user from paused prompt")
		}
	}
	return nil
}

// AfterApply is called after an update is applied
func (c context) AfterApply(update updater.Update) error {
	result, err := command.Exec(c.config.keybasePath(), []string{"update", "notify", "after-apply"}, 2*time.Minute, c.log)
	if err != nil {
		c.log.Warningf("Error in after apply: %s (%s)", err, result.CombinedOutput())
	}
	return nil
}

func (c context) AfterUpdateCheck(update *updater.Update) {
	if update != nil {
		// If we received an update from the check let's exit, so the watchdog
		// process (e.g. launchd on darwin) can restart us, no matter what, even if
		// there was an error, and even if the update was or wasn't applied.
		// There is no difference between doing another update check in a loop after
		// delay and restarting the service.
		c.log.Infof("%s", "Exiting for restart")
		os.Exit(0)
	}
}
<commit_msg>Add key id for winbot (#85)<commit_after>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package keybase

import (
	"fmt"
	"os"
	"time"

	"github.com/keybase/go-updater"
	"github.com/keybase/go-updater/command"
)

// validCodeSigningKIDs are the list of valid code signing IDs for saltpack verify
var validCodeSigningKIDs = map[string]bool{
	"9092ae4e790763dc7343851b977930f35b16cf43ab0ad900a2af3d3ad5cea1a1": true, // keybot (device)
	"d3458bbecdfc0d0ae39fec05722c6e3e897c169223835977a8aa208dfcd902d3": true, // max (device, home)
	"65ae849d1949a8b0021b165b0edaf722e2a7a9036e07817e056e2d721bddcc0e": true, // max (paper key, cry glass)
	"3a5a45c545ef4f661b8b7573711aaecee3fd5717053484a3a3e725cd68abaa5a": true, // chris (device, ccpro)
	"03d86864fb20e310590042ad3d5492c3f5d06728620175b03c717c211bfaccc2": true, // chris (paper key, clay harbor)
	"3fab8bd234c902f431b1c1bd0c8e7d1b2290c784d359fd555ec360b1a8102506": true, // winbot (device, Build)
}

// Log is the logging interface for the keybase package
type Log interface {
	Debug(...interface{})
	Info(...interface{})
	Debugf(s string, args ...interface{})
	Infof(s string, args ...interface{})
	Warningf(s string, args ...interface{})
	Errorf(s string, args ...interface{})
}

// context is an updater.Context implementation
type context struct {
	// config is updater config
	config Config
	// log is the logger
	log Log
}

// endpoints define all the url locations for reporting, etc
type endpoints struct {
	update  string
	action  string
	success string
	err     string
}

var defaultEndpoints = endpoints{
	update:  "https://api.keybase.io/_/api/1.0/pkg/update.json",
	action:  "https://api.keybase.io/_/api/1.0/pkg/act.json",
	success: "https://api.keybase.io/_/api/1.0/pkg/success.json",
	err:     "https://api.keybase.io/_/api/1.0/pkg/error.json",
}

func newContext(cfg Config, log Log) *context {
	ctx := context{
		config: cfg,
		log:    log,
	}
	return &ctx
}

// NewUpdaterContext returns an updater context for Keybase
func NewUpdaterContext(pathToKeybase string, log Log) (updater.Context, *updater.Updater) {
	cfg, err := newConfig("Keybase", pathToKeybase, log)
	if err != nil {
		log.Warningf("Error loading config for context: %s", err)
	}

	src := NewUpdateSource(cfg, log)
	// For testing
	// (cd /Applications; ditto -c -k --sequesterRsrc --keepParent Keybase.app /tmp/Keybase.zip)
	//src := updater.NewLocalUpdateSource("/tmp/Keybase.zip", log)
	upd := updater.NewUpdater(src, &cfg, log)
	return newContext(&cfg, log), upd
}

// UpdateOptions returns update options
func (c *context) UpdateOptions() updater.UpdateOptions {
	return c.config.updaterOptions()
}

// GetUpdateUI returns Update UI
func (c *context) GetUpdateUI() updater.UpdateUI {
	return c
}

// GetLog returns log
func (c context) GetLog() Log {
	return c.log
}

// Verify verifies the signature
func (c context) Verify(update updater.Update) error {
	return updater.SaltpackVerifyDetachedFileAtPath(update.Asset.LocalPath, update.Asset.Signature, validCodeSigningKIDs, c.log)
}

type checkInUseResult struct {
	InUse bool `json:"in_use"`
}

func (c context) checkInUse() (bool, error) {
	var result checkInUseResult
	if err := command.ExecForJSON(c.config.keybasePath(), []string{"update", "check-in-use"}, &result, time.Minute, c.log); err != nil {
		return false, err
	}
	return result.InUse, nil
}

// BeforeApply is called before an update is applied
func (c context) BeforeApply(update updater.Update) error {
	inUse, err := c.checkInUse()
	if err != nil {
		c.log.Warningf("Error trying to check in use: %s", err)
	}
	if inUse {
		if cancel := c.PausedPrompt(); cancel {
			return fmt.Errorf("Canceled by user from paused prompt")
		}
	}
	return nil
}

// AfterApply is called after an update is applied
func (c context) AfterApply(update updater.Update) error {
	result, err := command.Exec(c.config.keybasePath(), []string{"update", "notify", "after-apply"}, 2*time.Minute, c.log)
	if err != nil {
		c.log.Warningf("Error in after apply: %s (%s)", err, result.CombinedOutput())
	}
	return nil
}

func (c context) AfterUpdateCheck(update *updater.Update) {
	if update != nil {
		// If we received an update from the check let's exit, so the watchdog
		// process (e.g. launchd on darwin) can restart us, no matter what, even if
		// there was an error, and even if the update was or wasn't applied.
		// There is no difference between doing another update check in a loop after
		// delay and restarting the service.
		c.log.Infof("%s", "Exiting for restart")
		os.Exit(0)
	}
}
<|endoftext|>


<commit_before>// Copyright 2010 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 net

import (
	"once"
	"os"
	"sync"
	"syscall"
	"unsafe"
)

// BUG(brainman): The Windows implementation does not implement SetTimeout.

// IO completion result parameters.
type ioResult struct {
	key   uint32
	qty   uint32
	errno int
}

// Network file descriptor.
type netFD struct {
	// locking/lifetime of sysfd
	sysmu   sync.Mutex
	sysref  int
	closing bool

	// immutable until Close
	sysfd   int
	family  int
	proto   int
	sysfile *os.File
	cr      chan *ioResult
	cw      chan *ioResult
	net     string
	laddr   Addr
	raddr   Addr

	// owned by client
	rdeadline_delta int64
	rdeadline       int64
	rio             sync.Mutex
	wdeadline_delta int64
	wdeadline       int64
	wio             sync.Mutex
}

type InvalidConnError struct{}

func (e *InvalidConnError) String() string  { return "invalid net.Conn" }
func (e *InvalidConnError) Temporary() bool { return false }
func (e *InvalidConnError) Timeout() bool   { return false }

// pollServer will run around waiting for io completion request
// to arrive. Every request received will contain channel to signal
// io owner about the completion.

type pollServer struct {
	iocp int32
}

func newPollServer() (s *pollServer, err os.Error) {
	s = new(pollServer)
	var e int
	if s.iocp, e = syscall.CreateIoCompletionPort(-1, 0, 0, 1); e != 0 {
		return nil, os.NewSyscallError("CreateIoCompletionPort", e)
	}
	go s.Run()
	return s, nil
}

type ioPacket struct {
	// Used by IOCP interface,
	// it must be first field of the struct,
	// as our code rely on it.
	o syscall.Overlapped

	// Link to the io owner.
	c chan *ioResult
}

func (s *pollServer) getCompletedIO() (ov *syscall.Overlapped, result *ioResult, err os.Error) {
	var r ioResult
	var o *syscall.Overlapped
	_, e := syscall.GetQueuedCompletionStatus(s.iocp, &r.qty, &r.key, &o, syscall.INFINITE)
	switch {
	case e == 0:
		// Dequeued successfully completed io packet.
		return o, &r, nil
	case e == syscall.WAIT_TIMEOUT && o == nil:
		// Wait has timed out (should not happen now, but might be used in the future).
		return nil, &r, os.NewSyscallError("GetQueuedCompletionStatus", e)
	case o == nil:
		// Failed to dequeue anything -> report the error.
		return nil, &r, os.NewSyscallError("GetQueuedCompletionStatus", e)
	default:
		// Dequeued failed io packet.
		r.errno = e
		return o, &r, nil
	}
	return
}

func (s *pollServer) Run() {
	for {
		o, r, err := s.getCompletedIO()
		if err != nil {
			panic("Run pollServer: " + err.String() + "\n")
		}
		p := (*ioPacket)(unsafe.Pointer(o))
		p.c <- r
	}
}

// Network FD methods.
// All the network FDs use a single pollServer.

var pollserver *pollServer

func startServer() {
	p, err := newPollServer()
	if err != nil {
		panic("Start pollServer: " + err.String() + "\n")
	}
	pollserver = p
}

var initErr os.Error

func newFD(fd, family, proto int, net string, laddr, raddr Addr) (f *netFD, err os.Error) {
	if initErr != nil {
		return nil, initErr
	}
	once.Do(startServer)
	// Associate our socket with pollserver.iocp.
	if _, e := syscall.CreateIoCompletionPort(int32(fd), pollserver.iocp, 0, 0); e != 0 {
		return nil, &OpError{"CreateIoCompletionPort", net, laddr, os.Errno(e)}
	}
	f = &netFD{
		sysfd:  fd,
		family: family,
		proto:  proto,
		cr:     make(chan *ioResult),
		cw:     make(chan *ioResult),
		net:    net,
		laddr:  laddr,
		raddr:  raddr,
	}
	var ls, rs string
	if laddr != nil {
		ls = laddr.String()
	}
	if raddr != nil {
		rs = raddr.String()
	}
	f.sysfile = os.NewFile(fd, net+":"+ls+"->"+rs)
	return f, nil
}

// Add a reference to this fd.
func (fd *netFD) incref() {
	fd.sysmu.Lock()
	fd.sysref++
	fd.sysmu.Unlock()
}

// Remove a reference to this FD and close if we've been asked to do so (and
// there are no references left.
func (fd *netFD) decref() {
	fd.sysmu.Lock()
	fd.sysref--
	if fd.closing && fd.sysref == 0 && fd.sysfd >= 0 {
		// In case the user has set linger, switch to blocking mode so
		// the close blocks.  As long as this doesn't happen often, we
		// can handle the extra OS processes.  Otherwise we'll need to
		// use the pollserver for Close too.  Sigh.
		syscall.SetNonblock(fd.sysfd, false)
		fd.sysfile.Close()
		fd.sysfile = nil
		fd.sysfd = -1
	}
	fd.sysmu.Unlock()
}

func (fd *netFD) Close() os.Error {
	if fd == nil || fd.sysfile == nil {
		return os.EINVAL
	}

	fd.incref()
	syscall.Shutdown(fd.sysfd, syscall.SHUT_RDWR)
	fd.closing = true
	fd.decref()
	return nil
}

func newWSABuf(p []byte) *syscall.WSABuf {
	return &syscall.WSABuf{uint32(len(p)), (*byte)(unsafe.Pointer(&p[0]))}
}

func (fd *netFD) Read(p []byte) (n int, err os.Error) {
	if fd == nil {
		return 0, os.EINVAL
	}
	fd.rio.Lock()
	defer fd.rio.Unlock()
	fd.incref()
	defer fd.decref()
	if fd.sysfile == nil {
		return 0, os.EINVAL
	}
	// Submit receive request.
	var pckt ioPacket
	pckt.c = fd.cr
	var done uint32
	flags := uint32(0)
	e := syscall.WSARecv(uint32(fd.sysfd), newWSABuf(p), 1, &done, &flags, &pckt.o, nil)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		return 0, &OpError{"WSARecv", fd.net, fd.laddr, os.Errno(e)}
	}
	// Wait for our request to complete.
	r := <-pckt.c
	if r.errno != 0 {
		err = &OpError{"WSARecv", fd.net, fd.laddr, os.Errno(r.errno)}
	}
	n = int(r.qty)
	return
}

func (fd *netFD) ReadFrom(p []byte) (n int, sa syscall.Sockaddr, err os.Error) {
	var r syscall.Sockaddr
	return 0, r, nil
}

func (fd *netFD) Write(p []byte) (n int, err os.Error) {
	if fd == nil {
		return 0, os.EINVAL
	}
	fd.wio.Lock()
	defer fd.wio.Unlock()
	fd.incref()
	defer fd.decref()
	if fd.sysfile == nil {
		return 0, os.EINVAL
	}
	// Submit send request.
	var pckt ioPacket
	pckt.c = fd.cw
	var done uint32
	e := syscall.WSASend(uint32(fd.sysfd), newWSABuf(p), 1, &done, uint32(0), &pckt.o, nil)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		return 0, &OpError{"WSASend", fd.net, fd.laddr, os.Errno(e)}
	}
	// Wait for our request to complete.
	r := <-pckt.c
	if r.errno != 0 {
		err = &OpError{"WSASend", fd.net, fd.laddr, os.Errno(r.errno)}
	}
	n = int(r.qty)
	return
}

func (fd *netFD) WriteTo(p []byte, sa syscall.Sockaddr) (n int, err os.Error) {
	return 0, nil
}

func (fd *netFD) accept(toAddr func(syscall.Sockaddr) Addr) (nfd *netFD, err os.Error) {
	if fd == nil || fd.sysfile == nil {
		return nil, os.EINVAL
	}
	fd.incref()
	defer fd.decref()

	// Get new socket.
	// See ../syscall/exec.go for description of ForkLock.
	syscall.ForkLock.RLock()
	s, e := syscall.Socket(fd.family, fd.proto, 0)
	if e != 0 {
		syscall.ForkLock.RUnlock()
		return nil, os.Errno(e)
	}
	syscall.CloseOnExec(s)
	syscall.ForkLock.RUnlock()

	// Associate our new socket with IOCP.
	once.Do(startServer)
	if _, e = syscall.CreateIoCompletionPort(int32(s), pollserver.iocp, 0, 0); e != 0 {
		return nil, &OpError{"CreateIoCompletionPort", fd.net, fd.laddr, os.Errno(e)}
	}

	// Submit accept request.
	// Will use new unique channel here, because, unlike Read or Write,
	// Accept is expected to be executed by many goroutines simultaniously.
	var pckt ioPacket
	pckt.c = make(chan *ioResult)
	attrs, e := syscall.AcceptIOCP(fd.sysfd, s, &pckt.o)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		syscall.Close(s)
		return nil, &OpError{"AcceptEx", fd.net, fd.laddr, os.Errno(e)}
	}

	// Wait for peer connection.
	r := <-pckt.c
	if r.errno != 0 {
		syscall.Close(s)
		return nil, &OpError{"AcceptEx", fd.net, fd.laddr, os.Errno(r.errno)}
	}

	// Inherit properties of the listening socket.
	e = syscall.SetsockoptInt(s, syscall.SOL_SOCKET, syscall.SO_UPDATE_ACCEPT_CONTEXT, fd.sysfd)
	if e != 0 {
		syscall.Close(s)
		return nil, &OpError{"Setsockopt", fd.net, fd.laddr, os.Errno(r.errno)}
	}

	// Get local and peer addr out of AcceptEx buffer.
	lsa, rsa := syscall.GetAcceptIOCPSockaddrs(attrs)

	// Create our netFD and return it for further use.
	laddr := toAddr(lsa)
	raddr := toAddr(rsa)

	f := &netFD{
		sysfd:  s,
		family: fd.family,
		proto:  fd.proto,
		cr:     make(chan *ioResult),
		cw:     make(chan *ioResult),
		net:    fd.net,
		laddr:  laddr,
		raddr:  raddr,
	}
	var ls, rs string
	if laddr != nil {
		ls = laddr.String()
	}
	if raddr != nil {
		rs = raddr.String()
	}
	f.sysfile = os.NewFile(s, fd.net+":"+ls+"->"+rs)
	return f, nil
}

func init() {
	var d syscall.WSAData
	e := syscall.WSAStartup(uint32(0x101), &d)
	if e != 0 {
		initErr = os.NewSyscallError("WSAStartup", e)
	}
}
<commit_msg>net(windows): properly handle EOF in (*netFD).Read().<commit_after>// Copyright 2010 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 net

import (
	"once"
	"os"
	"sync"
	"syscall"
	"unsafe"
)

// BUG(brainman): The Windows implementation does not implement SetTimeout.

// IO completion result parameters.
type ioResult struct {
	key   uint32
	qty   uint32
	errno int
}

// Network file descriptor.
type netFD struct {
	// locking/lifetime of sysfd
	sysmu   sync.Mutex
	sysref  int
	closing bool

	// immutable until Close
	sysfd   int
	family  int
	proto   int
	sysfile *os.File
	cr      chan *ioResult
	cw      chan *ioResult
	net     string
	laddr   Addr
	raddr   Addr

	// owned by client
	rdeadline_delta int64
	rdeadline       int64
	rio             sync.Mutex
	wdeadline_delta int64
	wdeadline       int64
	wio             sync.Mutex
}

type InvalidConnError struct{}

func (e *InvalidConnError) String() string  { return "invalid net.Conn" }
func (e *InvalidConnError) Temporary() bool { return false }
func (e *InvalidConnError) Timeout() bool   { return false }

// pollServer will run around waiting for io completion request
// to arrive. Every request received will contain channel to signal
// io owner about the completion.

type pollServer struct {
	iocp int32
}

func newPollServer() (s *pollServer, err os.Error) {
	s = new(pollServer)
	var e int
	if s.iocp, e = syscall.CreateIoCompletionPort(-1, 0, 0, 1); e != 0 {
		return nil, os.NewSyscallError("CreateIoCompletionPort", e)
	}
	go s.Run()
	return s, nil
}

type ioPacket struct {
	// Used by IOCP interface,
	// it must be first field of the struct,
	// as our code rely on it.
	o syscall.Overlapped

	// Link to the io owner.
	c chan *ioResult
}

func (s *pollServer) getCompletedIO() (ov *syscall.Overlapped, result *ioResult, err os.Error) {
	var r ioResult
	var o *syscall.Overlapped
	_, e := syscall.GetQueuedCompletionStatus(s.iocp, &r.qty, &r.key, &o, syscall.INFINITE)
	switch {
	case e == 0:
		// Dequeued successfully completed io packet.
		return o, &r, nil
	case e == syscall.WAIT_TIMEOUT && o == nil:
		// Wait has timed out (should not happen now, but might be used in the future).
		return nil, &r, os.NewSyscallError("GetQueuedCompletionStatus", e)
	case o == nil:
		// Failed to dequeue anything -> report the error.
		return nil, &r, os.NewSyscallError("GetQueuedCompletionStatus", e)
	default:
		// Dequeued failed io packet.
		r.errno = e
		return o, &r, nil
	}
	return
}

func (s *pollServer) Run() {
	for {
		o, r, err := s.getCompletedIO()
		if err != nil {
			panic("Run pollServer: " + err.String() + "\n")
		}
		p := (*ioPacket)(unsafe.Pointer(o))
		p.c <- r
	}
}

// Network FD methods.
// All the network FDs use a single pollServer.

var pollserver *pollServer

func startServer() {
	p, err := newPollServer()
	if err != nil {
		panic("Start pollServer: " + err.String() + "\n")
	}
	pollserver = p
}

var initErr os.Error

func newFD(fd, family, proto int, net string, laddr, raddr Addr) (f *netFD, err os.Error) {
	if initErr != nil {
		return nil, initErr
	}
	once.Do(startServer)
	// Associate our socket with pollserver.iocp.
	if _, e := syscall.CreateIoCompletionPort(int32(fd), pollserver.iocp, 0, 0); e != 0 {
		return nil, &OpError{"CreateIoCompletionPort", net, laddr, os.Errno(e)}
	}
	f = &netFD{
		sysfd:  fd,
		family: family,
		proto:  proto,
		cr:     make(chan *ioResult),
		cw:     make(chan *ioResult),
		net:    net,
		laddr:  laddr,
		raddr:  raddr,
	}
	var ls, rs string
	if laddr != nil {
		ls = laddr.String()
	}
	if raddr != nil {
		rs = raddr.String()
	}
	f.sysfile = os.NewFile(fd, net+":"+ls+"->"+rs)
	return f, nil
}

// Add a reference to this fd.
func (fd *netFD) incref() {
	fd.sysmu.Lock()
	fd.sysref++
	fd.sysmu.Unlock()
}

// Remove a reference to this FD and close if we've been asked to do so (and
// there are no references left.
func (fd *netFD) decref() {
	fd.sysmu.Lock()
	fd.sysref--
	if fd.closing && fd.sysref == 0 && fd.sysfd >= 0 {
		// In case the user has set linger, switch to blocking mode so
		// the close blocks.  As long as this doesn't happen often, we
		// can handle the extra OS processes.  Otherwise we'll need to
		// use the pollserver for Close too.  Sigh.
		syscall.SetNonblock(fd.sysfd, false)
		fd.sysfile.Close()
		fd.sysfile = nil
		fd.sysfd = -1
	}
	fd.sysmu.Unlock()
}

func (fd *netFD) Close() os.Error {
	if fd == nil || fd.sysfile == nil {
		return os.EINVAL
	}

	fd.incref()
	syscall.Shutdown(fd.sysfd, syscall.SHUT_RDWR)
	fd.closing = true
	fd.decref()
	return nil
}

func newWSABuf(p []byte) *syscall.WSABuf {
	return &syscall.WSABuf{uint32(len(p)), (*byte)(unsafe.Pointer(&p[0]))}
}

func (fd *netFD) Read(p []byte) (n int, err os.Error) {
	if fd == nil {
		return 0, os.EINVAL
	}
	fd.rio.Lock()
	defer fd.rio.Unlock()
	fd.incref()
	defer fd.decref()
	if fd.sysfile == nil {
		return 0, os.EINVAL
	}
	// Submit receive request.
	var pckt ioPacket
	pckt.c = fd.cr
	var done uint32
	flags := uint32(0)
	e := syscall.WSARecv(uint32(fd.sysfd), newWSABuf(p), 1, &done, &flags, &pckt.o, nil)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		return 0, &OpError{"WSARecv", fd.net, fd.laddr, os.Errno(e)}
	}
	// Wait for our request to complete.
	r := <-pckt.c
	if r.errno != 0 {
		err = &OpError{"WSARecv", fd.net, fd.laddr, os.Errno(r.errno)}
	}
	n = int(r.qty)
	if err == nil && n == 0 {
		err = os.EOF
	}
	return
}

func (fd *netFD) ReadFrom(p []byte) (n int, sa syscall.Sockaddr, err os.Error) {
	var r syscall.Sockaddr
	return 0, r, nil
}

func (fd *netFD) Write(p []byte) (n int, err os.Error) {
	if fd == nil {
		return 0, os.EINVAL
	}
	fd.wio.Lock()
	defer fd.wio.Unlock()
	fd.incref()
	defer fd.decref()
	if fd.sysfile == nil {
		return 0, os.EINVAL
	}
	// Submit send request.
	var pckt ioPacket
	pckt.c = fd.cw
	var done uint32
	e := syscall.WSASend(uint32(fd.sysfd), newWSABuf(p), 1, &done, uint32(0), &pckt.o, nil)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		return 0, &OpError{"WSASend", fd.net, fd.laddr, os.Errno(e)}
	}
	// Wait for our request to complete.
	r := <-pckt.c
	if r.errno != 0 {
		err = &OpError{"WSASend", fd.net, fd.laddr, os.Errno(r.errno)}
	}
	n = int(r.qty)
	return
}

func (fd *netFD) WriteTo(p []byte, sa syscall.Sockaddr) (n int, err os.Error) {
	return 0, nil
}

func (fd *netFD) accept(toAddr func(syscall.Sockaddr) Addr) (nfd *netFD, err os.Error) {
	if fd == nil || fd.sysfile == nil {
		return nil, os.EINVAL
	}
	fd.incref()
	defer fd.decref()

	// Get new socket.
	// See ../syscall/exec.go for description of ForkLock.
	syscall.ForkLock.RLock()
	s, e := syscall.Socket(fd.family, fd.proto, 0)
	if e != 0 {
		syscall.ForkLock.RUnlock()
		return nil, os.Errno(e)
	}
	syscall.CloseOnExec(s)
	syscall.ForkLock.RUnlock()

	// Associate our new socket with IOCP.
	once.Do(startServer)
	if _, e = syscall.CreateIoCompletionPort(int32(s), pollserver.iocp, 0, 0); e != 0 {
		return nil, &OpError{"CreateIoCompletionPort", fd.net, fd.laddr, os.Errno(e)}
	}

	// Submit accept request.
	// Will use new unique channel here, because, unlike Read or Write,
	// Accept is expected to be executed by many goroutines simultaniously.
	var pckt ioPacket
	pckt.c = make(chan *ioResult)
	attrs, e := syscall.AcceptIOCP(fd.sysfd, s, &pckt.o)
	switch e {
	case 0:
		// IO completed immediately, but we need to get our completion message anyway.
	case syscall.ERROR_IO_PENDING:
		// IO started, and we have to wait for it's completion.
	default:
		syscall.Close(s)
		return nil, &OpError{"AcceptEx", fd.net, fd.laddr, os.Errno(e)}
	}

	// Wait for peer connection.
	r := <-pckt.c
	if r.errno != 0 {
		syscall.Close(s)
		return nil, &OpError{"AcceptEx", fd.net, fd.laddr, os.Errno(r.errno)}
	}

	// Inherit properties of the listening socket.
	e = syscall.SetsockoptInt(s, syscall.SOL_SOCKET, syscall.SO_UPDATE_ACCEPT_CONTEXT, fd.sysfd)
	if e != 0 {
		syscall.Close(s)
		return nil, &OpError{"Setsockopt", fd.net, fd.laddr, os.Errno(r.errno)}
	}

	// Get local and peer addr out of AcceptEx buffer.
	lsa, rsa := syscall.GetAcceptIOCPSockaddrs(attrs)

	// Create our netFD and return it for further use.
	laddr := toAddr(lsa)
	raddr := toAddr(rsa)

	f := &netFD{
		sysfd:  s,
		family: fd.family,
		proto:  fd.proto,
		cr:     make(chan *ioResult),
		cw:     make(chan *ioResult),
		net:    fd.net,
		laddr:  laddr,
		raddr:  raddr,
	}
	var ls, rs string
	if laddr != nil {
		ls = laddr.String()
	}
	if raddr != nil {
		rs = raddr.String()
	}
	f.sysfile = os.NewFile(s, fd.net+":"+ls+"->"+rs)
	return f, nil
}

func init() {
	var d syscall.WSAData
	e := syscall.WSAStartup(uint32(0x101), &d)
	if e != 0 {
		initErr = os.NewSyscallError("WSAStartup", e)
	}
}
<|endoftext|>