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

package client

import (
	"fmt"
	"os"
	"path"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/client/go/qrcode"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type ProvisionUI struct {
	parent *UI
	role   libkb.KexRole
}

func NewProvisionUIProtocol(g *libkb.GlobalContext, role libkb.KexRole) rpc.Protocol {
	return keybase1.ProvisionUiProtocol(g.UI.GetProvisionUI(role))
}

func (p ProvisionUI) ChooseProvisioningMethod(ctx context.Context, arg keybase1.ChooseProvisioningMethodArg) (keybase1.ProvisionMethod, error) {
	p.parent.Output("How would you like to sign this install of Keybase?\n\n")
	p.parent.Output("(1) Use an existing device\n")
	p.parent.Output("(2) Use a paper key\n")
	p.parent.Output("(3) Use my Keybase passphrase\n")
	max := 3
	if arg.GpgOption {
		p.parent.Printf("(4) Use GPG\n")
		max = 4
	}

	var res keybase1.ProvisionMethod
	ret, err := PromptSelectionOrCancel(PromptDescriptorChooseProvisioningMethod, p.parent, "Choose a signing option", 1, max)
	if err != nil {
		if err == ErrInputCanceled {
			return res, libkb.CanceledError{M: "user canceled input"}
		}
		return res, err
	}
	switch ret {
	case 1:
		return keybase1.ProvisionMethod_DEVICE, nil
	case 2:
		return keybase1.ProvisionMethod_PAPER_KEY, nil
	case 3:
		return keybase1.ProvisionMethod_PASSPHRASE, nil
	case 4:
		return keybase1.ProvisionMethod_GPG, nil
	}
	return res, fmt.Errorf("invalid provision option: %d", ret)
}

func (p ProvisionUI) ChooseDeviceType(ctx context.Context, sessionID int) (keybase1.DeviceType, error) {
	p.parent.Output("What type of device would you like to connect this computer with?\n\n")
	p.parent.Output("(1) Desktop or laptop\n")
	p.parent.Output("(2) Mobile phone\n\n")

	var res keybase1.DeviceType
	ret, err := PromptSelectionOrCancel(PromptDescriptorChooseDeviceType, p.parent, "Choose a device type", 1, 2)
	if err != nil {
		if err == ErrInputCanceled {
			return res, libkb.CanceledError{M: "user canceled input"}
		}
		return res, err
	}
	switch ret {
	case 1:
		return keybase1.DeviceType_DESKTOP, nil
	case 2:
		return keybase1.DeviceType_MOBILE, nil
	}
	return res, fmt.Errorf("invalid device type option: %d", ret)

}

func (p ProvisionUI) DisplayAndPromptSecret(ctx context.Context, arg keybase1.DisplayAndPromptSecretArg) (keybase1.SecretResponse, error) {
	var resp keybase1.SecretResponse
	if p.role == libkb.KexRoleProvisioner {
		// This is the provisioner device (device X)
		// For command line app, all secrets are entered on the provisioner only:
		p.parent.Output("\nEnter the verification code from your other device here.  To get\n")
		p.parent.Output("a verification code, run 'keybase login' on your other device.\n\n")

		ret, err := PromptWithChecker(PromptDescriptorProvisionPhrase, p.parent, "Verification code", false, libkb.CheckNotEmpty)
		if err != nil {
			return resp, err
		}
		resp.Phrase = ret
		return resp, nil

	}

	if p.role == libkb.KexRoleProvisionee {
		// this is the provisionee device (device Y)
		// For command line app, the provisionee displays secrets only

		p.parent.Output("Type this verification code into your other device:\n\n")
		p.parent.Output("\t" + arg.Phrase + "\n\n")
		p.parent.Output("If you are using the command line client on your other device, run this command:\n\n")
		p.parent.Output("\tkeybase device add\n\n")
		p.parent.Output("It will then prompt you for the verification code above.\n\n")

		if arg.OtherDeviceType == keybase1.DeviceType_MOBILE {
			encodings, err := qrcode.Encode([]byte(arg.Phrase))
			// ignoring any of these errors...phrase above will suffice.
			if err == nil {
				p.parent.Output("Or, scan this QR Code with the keybase app on your mobile phone:\n\n")
				p.parent.Output(encodings.Terminal)
				fname := path.Join(os.TempDir(), "keybase_qr.png")
				f, ferr := os.Create(fname)
				if ferr == nil {
					f.Write(encodings.PNG)
					f.Close()
					p.parent.Printf("\nThere's also a PNG version in %s that might work better.\n\n", fname)
				}
			}
		}
		return resp, nil
	}

	return resp, libkb.InvalidArgumentError{Msg: fmt.Sprintf("invalid ProvisionUI role: %d", p.role)}
}

func (p ProvisionUI) PromptNewDeviceName(ctx context.Context, arg keybase1.PromptNewDeviceNameArg) (string, error) {
	for i := 0; i < 10; i++ {

		name, err := PromptWithChecker(PromptDescriptorProvisionDeviceName, p.parent, "Enter a public name for this device", false, libkb.CheckDeviceName)
		if err != nil {
			return "", err
		}
		var match bool
		for _, existing := range arg.ExistingDevices {
			if libkb.NameCmp(name, existing) {
				match = true
				p.parent.Printf("Device name %q already in use.  Please try again.\n", name)
				break
			}
		}
		if !match {
			return name, nil
		}
	}
	return "", libkb.RetryExhaustedError{}
}

func (p ProvisionUI) DisplaySecretExchanged(ctx context.Context, sessionID int) error {
	p.parent.Output("Secret successfully exchanged.  On your new device, choose and save a public name for it.\n\n")
	return nil
}

func (p ProvisionUI) ProvisioneeSuccess(ctx context.Context, arg keybase1.ProvisioneeSuccessArg) error {
	p.parent.Printf(CHECK + " Success! You provisioned your device " + ColorString("bold", arg.DeviceName) + ".\n\n")
	p.parent.Printf("You are logged in as " + ColorString("bold", arg.Username) + "\n")
	// turn on when kbfs active:
	if false {
		p.parent.Printf("  - your keybase public directory is available at /keybase/public/%s\n", arg.Username)
		p.parent.Printf("  - your keybase encrypted directory is available at /keybase/private/%s\n", arg.Username)
	}

	p.parent.Printf("  - type `keybase help` for more info.\n")
	return nil
}

func (p ProvisionUI) ProvisionerSuccess(ctx context.Context, arg keybase1.ProvisionerSuccessArg) error {
	p.parent.Printf(CHECK + " Success! You added a new device named " + ColorString("bold", arg.DeviceName) + " to your account.\n\n")
	return nil
}
<commit_msg>Additional prompts during GPG provisioning<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"fmt"
	"os"
	"path"

	"golang.org/x/net/context"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	"github.com/keybase/client/go/qrcode"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
)

type ProvisionUI struct {
	parent *UI
	role   libkb.KexRole
}

func NewProvisionUIProtocol(g *libkb.GlobalContext, role libkb.KexRole) rpc.Protocol {
	return keybase1.ProvisionUiProtocol(g.UI.GetProvisionUI(role))
}

func (p ProvisionUI) ChooseProvisioningMethod(ctx context.Context, arg keybase1.ChooseProvisioningMethodArg) (keybase1.ProvisionMethod, error) {
	p.parent.Output("How would you like to sign this install of Keybase?\n\n")
	p.parent.Output("(1) Use an existing device\n")
	p.parent.Output("(2) Use a paper key\n")
	p.parent.Output("(3) Use my Keybase passphrase\n")
	max := 3
	if arg.GpgOption {
		p.parent.Printf("(4) Use GPG\n")
		max = 4
	}

	var res keybase1.ProvisionMethod
	ret, err := PromptSelectionOrCancel(PromptDescriptorChooseProvisioningMethod, p.parent, "Choose a signing option", 1, max)
	if err != nil {
		if err == ErrInputCanceled {
			return res, libkb.CanceledError{M: "user canceled input"}
		}
		return res, err
	}
	switch ret {
	case 1:
		return keybase1.ProvisionMethod_DEVICE, nil
	case 2:
		return keybase1.ProvisionMethod_PAPER_KEY, nil
	case 3:
		return keybase1.ProvisionMethod_PASSPHRASE, nil
	case 4:
		p.parent.Output("\nIn order to use GPG to sign this install of Keybase, in the next step\n")
		p.parent.Output("you will select one of your GPG keys.  The gpg client will be used to\n")
		p.parent.Output("export the secret key and import it into keybase's local encrypted\n")
		p.parent.Output("key store.\n\n")
		err = p.parent.PromptForConfirmation("Would you like to continue?")
		if err != nil {
			if _, ok := err.(NotConfirmedError); ok {
				p.parent.Output("\n\nWe have an issue for using gpg to sign the install without\nrequiring importing the secret key.  You can view it here:\n\n    https://github.com/keybase/client/issues/1308\n\nThanks!\n")
				return res, libkb.CanceledError{M: "user canceled gpg provisioning"}
			}
			p.parent.Printf("error type: %T\n", err)
			return res, err
		}
		return keybase1.ProvisionMethod_GPG, nil
	}
	return res, fmt.Errorf("invalid provision option: %d", ret)
}

func (p ProvisionUI) ChooseDeviceType(ctx context.Context, sessionID int) (keybase1.DeviceType, error) {
	p.parent.Output("What type of device would you like to connect this computer with?\n\n")
	p.parent.Output("(1) Desktop or laptop\n")
	p.parent.Output("(2) Mobile phone\n\n")

	var res keybase1.DeviceType
	ret, err := PromptSelectionOrCancel(PromptDescriptorChooseDeviceType, p.parent, "Choose a device type", 1, 2)
	if err != nil {
		if err == ErrInputCanceled {
			return res, libkb.CanceledError{M: "user canceled input"}
		}
		return res, err
	}
	switch ret {
	case 1:
		return keybase1.DeviceType_DESKTOP, nil
	case 2:
		return keybase1.DeviceType_MOBILE, nil
	}
	return res, fmt.Errorf("invalid device type option: %d", ret)

}

func (p ProvisionUI) DisplayAndPromptSecret(ctx context.Context, arg keybase1.DisplayAndPromptSecretArg) (keybase1.SecretResponse, error) {
	var resp keybase1.SecretResponse
	if p.role == libkb.KexRoleProvisioner {
		// This is the provisioner device (device X)
		// For command line app, all secrets are entered on the provisioner only:
		p.parent.Output("\nEnter the verification code from your other device here.  To get\n")
		p.parent.Output("a verification code, run 'keybase login' on your other device.\n\n")

		ret, err := PromptWithChecker(PromptDescriptorProvisionPhrase, p.parent, "Verification code", false, libkb.CheckNotEmpty)
		if err != nil {
			return resp, err
		}
		resp.Phrase = ret
		return resp, nil

	}

	if p.role == libkb.KexRoleProvisionee {
		// this is the provisionee device (device Y)
		// For command line app, the provisionee displays secrets only

		p.parent.Output("Type this verification code into your other device:\n\n")
		p.parent.Output("\t" + arg.Phrase + "\n\n")
		p.parent.Output("If you are using the command line client on your other device, run this command:\n\n")
		p.parent.Output("\tkeybase device add\n\n")
		p.parent.Output("It will then prompt you for the verification code above.\n\n")

		if arg.OtherDeviceType == keybase1.DeviceType_MOBILE {
			encodings, err := qrcode.Encode([]byte(arg.Phrase))
			// ignoring any of these errors...phrase above will suffice.
			if err == nil {
				p.parent.Output("Or, scan this QR Code with the keybase app on your mobile phone:\n\n")
				p.parent.Output(encodings.Terminal)
				fname := path.Join(os.TempDir(), "keybase_qr.png")
				f, ferr := os.Create(fname)
				if ferr == nil {
					f.Write(encodings.PNG)
					f.Close()
					p.parent.Printf("\nThere's also a PNG version in %s that might work better.\n\n", fname)
				}
			}
		}
		return resp, nil
	}

	return resp, libkb.InvalidArgumentError{Msg: fmt.Sprintf("invalid ProvisionUI role: %d", p.role)}
}

func (p ProvisionUI) PromptNewDeviceName(ctx context.Context, arg keybase1.PromptNewDeviceNameArg) (string, error) {
	for i := 0; i < 10; i++ {

		name, err := PromptWithChecker(PromptDescriptorProvisionDeviceName, p.parent, "Enter a public name for this device", false, libkb.CheckDeviceName)
		if err != nil {
			return "", err
		}
		var match bool
		for _, existing := range arg.ExistingDevices {
			if libkb.NameCmp(name, existing) {
				match = true
				p.parent.Printf("Device name %q already in use.  Please try again.\n", name)
				break
			}
		}
		if !match {
			return name, nil
		}
	}
	return "", libkb.RetryExhaustedError{}
}

func (p ProvisionUI) DisplaySecretExchanged(ctx context.Context, sessionID int) error {
	p.parent.Output("Secret successfully exchanged.  On your new device, choose and save a public name for it.\n\n")
	return nil
}

func (p ProvisionUI) ProvisioneeSuccess(ctx context.Context, arg keybase1.ProvisioneeSuccessArg) error {
	p.parent.Printf(CHECK + " Success! You provisioned your device " + ColorString("bold", arg.DeviceName) + ".\n\n")
	p.parent.Printf("You are logged in as " + ColorString("bold", arg.Username) + "\n")
	// turn on when kbfs active:
	if false {
		p.parent.Printf("  - your keybase public directory is available at /keybase/public/%s\n", arg.Username)
		p.parent.Printf("  - your keybase encrypted directory is available at /keybase/private/%s\n", arg.Username)
	}

	p.parent.Printf("  - type `keybase help` for more info.\n")
	return nil
}

func (p ProvisionUI) ProvisionerSuccess(ctx context.Context, arg keybase1.ProvisionerSuccessArg) error {
	p.parent.Printf(CHECK + " Success! You added a new device named " + ColorString("bold", arg.DeviceName) + " to your account.\n\n")
	return nil
}
<|endoftext|>


<commit_before>package powermux

import (
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

const (
	MAX_WIDTH  = 500
	MAX_DEPTH  = 100
	FAN_DEPTH  = 4
	FAN_SPREAD = 8
)

type noopHandler struct{}

func (h *noopHandler) ServeHTTP(_ http.ResponseWriter, _ *http.Request) {
	//noop
}

var emptyHandle = &noopHandler{}

func BenchmarkSingleRoute(b *testing.B) {
	r := NewServeMux()
	r.Route("/").Any(emptyHandle)
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, req)
	}
}

func BenchmarkShallowAndWide(b *testing.B) {
	r := NewServeMux()
	requests := make([]*http.Request, 0, MAX_WIDTH)
	for i := 0; i < MAX_WIDTH; i++ {
		route := "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
		r.Handle(route, emptyHandle)
		requests = append(requests, httptest.NewRequest(http.MethodGet, route, nil))
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, requests[i%MAX_WIDTH])
	}
}

// BenchmarkNarrowAndDeep is powermux's worst-case scenario. One route at the end
// of a really long path
func BenchmarkNarrowAndDeep(b *testing.B) {
	r := NewServeMux()
	var route string
	for i := 0; i < MAX_DEPTH; i++ {
		route += "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
	}
	r.Handle(route, emptyHandle)
	req := httptest.NewRequest(http.MethodGet, route, nil)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, req)
	}
}

func addFanRoutes(n int, r *Route) (routes []string) {
	for i := 0; i < FAN_SPREAD; i++ {
		route := "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
		r2 := r.Route(route).Any(emptyHandle)
		routes = append(routes, route)
		if n == 0 {
			break
		}
		addedRoutes := addFanRoutes(n-1, r2)
		routes = append(routes, addedRoutes...)
	}

	return routes
}

func BenchmarkFan(b *testing.B) {
	r := NewServeMux()
	routes := addFanRoutes(FAN_DEPTH, r.Route("/"))
	requests := make([]*http.Request, 0, len(routes))
	for _, route := range routes {
		req := httptest.NewRequest(http.MethodGet, route, nil)
		requests = append(requests, req)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, requests[i%len(requests)])
	}
}
<commit_msg>Add parallel benchmarks<commit_after>package powermux

import (
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
)

const (
	MAX_WIDTH  = 500
	MAX_DEPTH  = 100
	FAN_DEPTH  = 4
	FAN_SPREAD = 8
)

type noopHandler struct{}

func (h *noopHandler) ServeHTTP(_ http.ResponseWriter, _ *http.Request) {
	//noop
}

var emptyHandle = &noopHandler{}

func BenchmarkSingleRoute(b *testing.B) {
	r := NewServeMux()
	r.Route("/").Any(emptyHandle)
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, req)
	}
}

func BenchmarkShallowAndWide(b *testing.B) {
	r := NewServeMux()
	requests := make([]*http.Request, 0, MAX_WIDTH)
	for i := 0; i < MAX_WIDTH; i++ {
		route := "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
		r.Handle(route, emptyHandle)
		requests = append(requests, httptest.NewRequest(http.MethodGet, route, nil))
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, requests[i%MAX_WIDTH])
	}
}

// BenchmarkNarrowAndDeep is powermux's worst-case scenario. One route at the end
// of a really long path
func BenchmarkNarrowAndDeep(b *testing.B) {
	r := NewServeMux()
	var route string
	for i := 0; i < MAX_DEPTH; i++ {
		route += "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
	}
	r.Handle(route, emptyHandle)
	req := httptest.NewRequest(http.MethodGet, route, nil)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, req)
	}
}

func addFanRoutes(n int, r *Route) (routes []string) {
	for i := 0; i < FAN_SPREAD; i++ {
		route := "/" + hex.EncodeToString([]byte(fmt.Sprint(i)))
		r2 := r.Route(route).Any(emptyHandle)
		routes = append(routes, route)
		if n == 0 {
			break
		}
		addedRoutes := addFanRoutes(n-1, r2)
		routes = append(routes, addedRoutes...)
	}

	return routes
}

func BenchmarkFan(b *testing.B) {
	r := NewServeMux()
	routes := addFanRoutes(FAN_DEPTH, r.Route("/"))
	requests := make([]*http.Request, 0, len(routes))
	for _, route := range routes {
		req := httptest.NewRequest(http.MethodGet, route, nil)
		requests = append(requests, req)
	}
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		r.ServeHTTP(nil, requests[i%len(requests)])
	}
}

func BenchmarkSingleRouteParallel(b *testing.B) {
	r := NewServeMux()
	r.Route("/").Any(emptyHandle)
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	b.ResetTimer()

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			r.ServeHTTP(nil, req)
		}
	})
}
<|endoftext|>


<commit_before>/*
   Copyright 2016 Simon J Mudd <sjmudd@pobox.com>

   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 discovery manages a queue of discovery requests: an ordered
queue with no duplicates.

The configured processor function is called on an entry that's
received and this happens in parallel until maxCapacity go-routines
are running in parallel.

Any new requests get queued and processed as a go routine becomes
free.  The queue is processed in FIFO order.  If a request is
received for an instance that is already in the queue or already
being processed then it will be silently ignored.

*/

package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

// Queue is a container for processing the orchestrator discovery requests.
type Queue struct {
	// The number of active go routines processing discovery requests.
	concurrency uint
	// Channel used by the active go routines to say they have finished processing.
	done chan inst.InstanceKey
	// Input channel we are reading the discovery requests from.
	inputChan <-chan inst.InstanceKey
	// This has 2 uses: to indicate if there is a request in the queue
	// (so not being processed) or to indicate that the request is actively
	// being dealth with.  That state is not explicitly stored as it is not
	// really needed.
	knownKeys map[inst.InstanceKey]bool
	// lock while making changes
	lock sync.Mutex
	// The maximum number of go routines allowed to handle the queue at once.
	// This is a configuration parameter provided when creating the Queue.
	maxConcurrency uint
	// process to run on each received key
	processor func(i inst.InstanceKey)
	// This holds the discover requests (the instance name) which need to be
	// processed. but which are not currently being processed. All requests
	// are initially added to the end of this queue, and then the top element
	// will be popped off if the number of active go routines (defined by
	// concurrency) is less than the maximum specified value at which point
	// it will be processed by a new go routine.
	queue []inst.InstanceKey
}

var emptyKey = inst.InstanceKey{}

// NewQueue creates a new Queue entry.
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}
	// log.Debugf("Queue.push(%+v)", key)

	if _, found := q.knownKeys[key]; !found {
		// log.Debugf("Queue.push() adding %+v to knownKeys", key)
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
		// log.Debugf("Queue.push() ignoring knownKey %+v", key)
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	// log.Debugf("Queue.pop() returns %+v", key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// log.Debugf("Queue.dispatch() key: %q, concurrency: %d", key, q.concurrency)

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	// log.Debugf("Queue.acknowledgeJob(%+v) q.concurrency: %d", key, q.concurrency)
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	// log.Debugf("Queue.maybeDispatch() q.concurrency: %d, q.maxConcurrency: %d, len(q.queue): %d", q.concurrency, q.maxConcurrency, len(q.queue))
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	// log.Debugf("Queue.queueAndMaybeDispatch(%+v) concurency: %d", key, q.concurrency)
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<commit_msg>align license with the one used in other files<commit_after>/*
   Copyright 2014 Outbrain 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 discovery manages a queue of discovery requests: an ordered
queue with no duplicates.

The configured processor function is called on an entry that's
received and this happens in parallel until maxCapacity go-routines
are running in parallel.

Any new requests get queued and processed as a go routine becomes
free.  The queue is processed in FIFO order.  If a request is
received for an instance that is already in the queue or already
being processed then it will be silently ignored.

*/

package discovery

import (
	"errors"
	"sync"

	"github.com/outbrain/golib/log"
	"github.com/outbrain/orchestrator/go/inst"
)

// Queue is a container for processing the orchestrator discovery requests.
type Queue struct {
	// The number of active go routines processing discovery requests.
	concurrency uint
	// Channel used by the active go routines to say they have finished processing.
	done chan inst.InstanceKey
	// Input channel we are reading the discovery requests from.
	inputChan <-chan inst.InstanceKey
	// This has 2 uses: to indicate if there is a request in the queue
	// (so not being processed) or to indicate that the request is actively
	// being dealth with.  That state is not explicitly stored as it is not
	// really needed.
	knownKeys map[inst.InstanceKey]bool
	// lock while making changes
	lock sync.Mutex
	// The maximum number of go routines allowed to handle the queue at once.
	// This is a configuration parameter provided when creating the Queue.
	maxConcurrency uint
	// process to run on each received key
	processor func(i inst.InstanceKey)
	// This holds the discover requests (the instance name) which need to be
	// processed. but which are not currently being processed. All requests
	// are initially added to the end of this queue, and then the top element
	// will be popped off if the number of active go routines (defined by
	// concurrency) is less than the maximum specified value at which point
	// it will be processed by a new go routine.
	queue []inst.InstanceKey
}

var emptyKey = inst.InstanceKey{}

// NewQueue creates a new Queue entry.
func NewQueue(maxConcurrency uint, inputChan chan inst.InstanceKey, processor func(i inst.InstanceKey)) *Queue {
	log.Infof("Queue.NewQueue()")
	q := new(Queue)

	q.concurrency = 0                    // explicitly
	q.done = make(chan inst.InstanceKey) // Do I need this to be larger?
	q.inputChan = inputChan
	q.knownKeys = make(map[inst.InstanceKey]bool)
	q.maxConcurrency = maxConcurrency
	q.processor = processor
	q.queue = make([]inst.InstanceKey, 0)

	return q
}

// add the key to the slice if it does not exist in known keys
// - goroutine safe as only called inside the mutex
func (q *Queue) push(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.push(%v) is empty", key)
	}
	// log.Debugf("Queue.push(%+v)", key)

	if _, found := q.knownKeys[key]; !found {
		// log.Debugf("Queue.push() adding %+v to knownKeys", key)
		// add to the items that are being processed
		q.knownKeys[key] = true
		q.queue = append(q.queue, key)
	} else {
		// If key already there we just ignore it as the request is in the queue.
		// the known key also records stuff in the queue, so pending + active jobs.
		// log.Debugf("Queue.push() ignoring knownKey %+v", key)
	}
}

// remove the entry and remove it from known keys
func (q *Queue) pop() (inst.InstanceKey, error) {
	if len(q.queue) == 0 {
		return inst.InstanceKey{}, errors.New("q.pop() on empty queue")
	}
	key := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.knownKeys, key)
	// log.Debugf("Queue.pop() returns %+v", key)
	return key, nil
}

// dispatch a job from the queue (assumes we are in a locked state)
func (q *Queue) dispatch() {
	key, err := q.pop() // should never give an error but let's check anyway
	if err != nil {
		log.Fatal("Queue.dispatch() q.pop() returns: %+v", err)
		return
	}
	if key == emptyKey {
		log.Fatal("Queue.dispatch() key is empty")
	}

	q.concurrency++
	q.knownKeys[key] = true

	// log.Debugf("Queue.dispatch() key: %q, concurrency: %d", key, q.concurrency)

	// dispatch a discoverInstance() but tell us when we're done (to limit concurrency)
	go func() { // discover asynchronously
		q.processor(key)
		q.done <- key
	}()
}

// acknowledge a job has finished
// - we deal with the locking inside
func (q *Queue) acknowledgeJob(key inst.InstanceKey) {
	q.lock.Lock()
	delete(q.knownKeys, key)
	q.concurrency--
	// log.Debugf("Queue.acknowledgeJob(%+v) q.concurrency: %d", key, q.concurrency)
	q.lock.Unlock()
}

// drain queue by dispatching any jobs we have still
func (q *Queue) maybeDispatch() {
	q.lock.Lock()
	// log.Debugf("Queue.maybeDispatch() q.concurrency: %d, q.maxConcurrency: %d, len(q.queue): %d", q.concurrency, q.maxConcurrency, len(q.queue))
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// add an entry to the queue and dispatch something if concurrency is low enough
// - we deal with locking inside
func (q *Queue) queueAndMaybeDispatch(key inst.InstanceKey) {
	if key == emptyKey {
		log.Fatal("Queue.queueAndMaybeDispatch(%v) is empty", key)
	}
	q.lock.Lock()
	// log.Debugf("Queue.queueAndMaybeDispatch(%+v) concurency: %d", key, q.concurrency)
	q.push(key)
	if q.concurrency < q.maxConcurrency && len(q.queue) > 0 {
		q.dispatch()
	}
	q.lock.Unlock()
}

// cleanup is called when the input channel closes.
// we can not sit in the loop so we have to wait for running go-routines to finish
// but also to dispatch anything left in the queue until finally everything is done.
func (q *Queue) cleanup() {
	log.Infof("Queue.cleanup()")
	for q.concurrency > 0 && len(q.queue) > 0 {
		q.maybeDispatch()
		if key, closed := <-q.done; closed {
			return
		} else {
			q.acknowledgeJob(key)
		}
	}
}

// Ends when all elements in the queue have been handled.
// we read from inputChan and call processor up to maxConcurrency times in parallel
func (q *Queue) HandleRequests() {
	if q == nil {
		log.Infof("Queue.HandleRequests() q == nil ??. Should not happen")

		// no queue, nothing to do
		return
	}
	log.Infof("Queue.NewQueue() processing requests")
	for {
		select {
		case key, ok := <-q.inputChan:
			if ok {
				if key != emptyKey {
					q.queueAndMaybeDispatch(key)
				} else {
					log.Warningf("Queue.HandleRequests() q.inputChan received empty key %+v, ignoring (fix the upstream code to prevent this)", key)
				}
			} else {
				q.cleanup()
				log.Infof("Queue.HandleRequests() q.inputChan is closed. returning")
				return
			}
		case key, ok := <-q.done:
			if ok {
				q.acknowledgeJob(key)
			} else {
				log.Infof("Queue.HandleRequests() q.done is closed. returning (shouldn't get here)")
				return // we shouldn't get here as the return above should get triggered first
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "encoding/hex"
    //"runtime"
    "time"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
    "./helpers"
)

// Use structs to pass data around so I can refactor
type AppInfo struct {
    aid freefare.DESFireAid
    aidbytes []byte
    sysid []byte
    acl_read_base []byte
    acl_write_base []byte
    acl_file_id byte
}

type KeyChain struct {
    uid_read_key_id byte
    acl_read_key_id byte
    acl_write_key_id byte

    uid_read_key *freefare.DESFireKey
    acl_read_key *freefare.DESFireKey
    acl_write_key *freefare.DESFireKey
}

var (
    keychain = KeyChain{}
    appinfo = AppInfo{}
)

func init_appinfo() {
    keymap, err := helpers.LoadYAMLFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    appmap, err := helpers.LoadYAMLFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    // Application-id
    appinfo.aid, err = helpers.String2aid(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }

    // Needed for diversification
    appinfo.aidbytes = helpers.Aid2bytes(appinfo.aid)
    appinfo.sysid, err = hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    appinfo.acl_file_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_file_id"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    keychain.uid_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_write_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // The static app key to read UID
    keychain.uid_read_key, err = helpers.String2aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }

    // Bases for the diversified keys
    appinfo.acl_read_base, err = hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    appinfo.acl_write_base, err = hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

}

func recalculate_diversified_keys(realuid []byte) error {
    acl_read_bytes, err := keydiversification.AES128(appinfo.acl_read_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    acl_write_bytes, err := keydiversification.AES128(appinfo.acl_write_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    keychain.acl_read_key = helpers.Bytes2aeskey(acl_read_bytes)
    keychain.acl_write_key = helpers.Bytes2aeskey(acl_write_bytes)
    return nil
}

func handle_tag(desfiretag *freefare.DESFireTag) {
    desfiretag.Connect()

    uid_str := desfiretag.UID()

    fmt.Printf("Found tag %s\n", uid_str)

    fmt.Println("Selecting application");
    err := desfiretag.SelectApplication(appinfo.aid);
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    /*
    uid_str, err := desfiretag.CardUID()
    if err != nil {
        panic(err)
    }
    */

    realuid, err := hex.DecodeString(uid_str)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to parse real UID (%s), skipping tag", err))
        return
    }
    fmt.Println("Got UID: ", hex.EncodeToString(realuid));

    // Calculate the diversified keys
    err = recalculate_diversified_keys(realuid[:])
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to get diversified ACL keys (%s), skipping tag", err))
        return
    }

    fmt.Println("Got real ACL read key: ", keychain.acl_read_key)

    desfiretag.Disconnect()
}

func main() {

    init_appinfo();

    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    for {
        var tags []freefare.Tag
        for {
            tags, err = freefare.GetTags(d);
            if err != nil {
                continue
            }
            if len(tags) > 0 {
                break
            }
            time.Sleep(100 * time.Millisecond)
            //fmt.Println("...polling")
        }

        for i := 0; i < len(tags); i++ {
            tag := tags[i]
            desfiretag := tag.(freefare.DESFireTag)
            handle_tag(&desfiretag)
        }
        //runtime.GC()
    }

}

<commit_msg>still no crash<commit_after>package main

import (
    "fmt"
    "encoding/hex"
    //"runtime"
    "time"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
    "./helpers"
)

// Use structs to pass data around so I can refactor
type AppInfo struct {
    aid freefare.DESFireAid
    aidbytes []byte
    sysid []byte
    acl_read_base []byte
    acl_write_base []byte
    acl_file_id byte
}

type KeyChain struct {
    uid_read_key_id byte
    acl_read_key_id byte
    acl_write_key_id byte

    uid_read_key *freefare.DESFireKey
    acl_read_key *freefare.DESFireKey
    acl_write_key *freefare.DESFireKey
}

var (
    keychain = KeyChain{}
    appinfo = AppInfo{}
)

func init_appinfo() {
    keymap, err := helpers.LoadYAMLFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    appmap, err := helpers.LoadYAMLFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    // Application-id
    appinfo.aid, err = helpers.String2aid(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }

    // Needed for diversification
    appinfo.aidbytes = helpers.Aid2bytes(appinfo.aid)
    appinfo.sysid, err = hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    appinfo.acl_file_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_file_id"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    keychain.uid_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_write_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // The static app key to read UID
    keychain.uid_read_key, err = helpers.String2aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }

    // Bases for the diversified keys
    appinfo.acl_read_base, err = hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    appinfo.acl_write_base, err = hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

}

func recalculate_diversified_keys(realuid []byte) error {
    acl_read_bytes, err := keydiversification.AES128(appinfo.acl_read_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    acl_write_bytes, err := keydiversification.AES128(appinfo.acl_write_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    keychain.acl_read_key = helpers.Bytes2aeskey(acl_read_bytes)
    keychain.acl_write_key = helpers.Bytes2aeskey(acl_write_bytes)
    return nil
}

func handle_tag(desfiretag *freefare.DESFireTag) {
    desfiretag.Connect()

    uid_str := desfiretag.UID()

    fmt.Printf("Found tag %s\n", uid_str)

    fmt.Println("Selecting application");
    err := desfiretag.SelectApplication(appinfo.aid);
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    fmt.Println("Authenticating");
    err = desfiretag.Authenticate(keychain.uid_read_key_id,*keychain.uid_read_key)
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    /*
    realuid_str, err := desfiretag.CardUID()
    if err != nil {
        panic(err)
    }
    */

    realuid, err := hex.DecodeString(uid_str)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to parse real UID (%s), skipping tag", err))
        return
    }
    fmt.Println("Got UID: ", hex.EncodeToString(realuid));

    // Calculate the diversified keys
    err = recalculate_diversified_keys(realuid[:])
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to get diversified ACL keys (%s), skipping tag", err))
        return
    }

    fmt.Println("Got real ACL read key: ", keychain.acl_read_key)

    desfiretag.Disconnect()
}

func main() {

    init_appinfo();

    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    for {
        var tags []freefare.Tag
        for {
            tags, err = freefare.GetTags(d);
            if err != nil {
                continue
            }
            if len(tags) > 0 {
                break
            }
            time.Sleep(100 * time.Millisecond)
            //fmt.Println("...polling")
        }

        for i := 0; i < len(tags); i++ {
            tag := tags[i]
            desfiretag := tag.(freefare.DESFireTag)
            handle_tag(&desfiretag)
        }
        //runtime.GC()
    }

}

<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "encoding/hex"
    //"runtime"
    "time"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
    "./helpers"
)

// Use structs to pass data around so I can refactor
type AppInfo struct {
    aid freefare.DESFireAid
    aidbytes []byte
    sysid []byte
    acl_read_base []byte
    acl_write_base []byte
    acl_file_id byte
}

type KeyChain struct {
    uid_read_key_id byte
    acl_read_key_id byte
    acl_write_key_id byte

    uid_read_key *freefare.DESFireKey
    acl_read_key *freefare.DESFireKey
    acl_write_key *freefare.DESFireKey
}

var (
    keychain = KeyChain{}
    appinfo = AppInfo{}
)

func init_appinfo() {
    keymap, err := helpers.LoadYAMLFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    appmap, err := helpers.LoadYAMLFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    // Application-id
    appinfo.aid, err = helpers.String2aid(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }

    // Needed for diversification
    appinfo.aidbytes = helpers.Aid2bytes(appinfo.aid)
    appinfo.sysid, err = hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    appinfo.acl_file_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_file_id"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    keychain.uid_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_write_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // The static app key to read UID
    keychain.uid_read_key, err = helpers.String2aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }

    // Bases for the diversified keys
    appinfo.acl_read_base, err = hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    appinfo.acl_write_base, err = hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

}

func recalculate_diversified_keys(realuid []byte) error {
    acl_read_bytes, err := keydiversification.AES128(appinfo.acl_read_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    acl_write_bytes, err := keydiversification.AES128(appinfo.acl_write_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    keychain.acl_read_key = helpers.Bytes2aeskey(acl_read_bytes)
    keychain.acl_write_key = helpers.Bytes2aeskey(acl_write_bytes)
    return nil
}

func handle_tag(desfiretag *freefare.DESFireTag) {
    desfiretag.Connect()

    uid_str := desfiretag.UID()

    fmt.Printf("Found tag %s\n", uid_str)

    fmt.Println("Selecting application");
    err := desfiretag.SelectApplication(appinfo.aid);
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    fmt.Println("Authenticating");
    err = desfiretag.Authenticate(keychain.uid_read_key_id,*keychain.uid_read_key)
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    realuid_str, err := desfiretag.CardUID()
    if err != nil {
        panic(err)
    }

    realuid, err := hex.DecodeString(realuid_str)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to parse real UID (%s), skipping tag", err))
        return
    }
    fmt.Println("Got UID: ", hex.EncodeToString(realuid));

    // Calculate the diversified keys
    err = recalculate_diversified_keys(realuid[:])
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to get diversified ACL keys (%s), skipping tag", err))
        return
    }

    fmt.Println("Got real ACL read key: ", keychain.acl_read_key)

    desfiretag.Disconnect()
}

func main() {

    init_appinfo();

    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    for {
        var tags []freefare.Tag
        for {
            tags, err = freefare.GetTags(d);
            if err != nil {
                continue
            }
            if len(tags) > 0 {
                break
            }
            time.Sleep(100 * time.Millisecond)
            //fmt.Println("...polling")
        }

        for i := 0; i < len(tags); i++ {
            tag := tags[i]
            desfiretag := tag.(freefare.DESFireTag)
            handle_tag(&desfiretag)
        }
        //runtime.GC()
    }

}

<commit_msg>still no crash<commit_after>package main

import (
    "fmt"
    "encoding/hex"
    //"runtime"
    "time"
    "github.com/fuzxxl/nfc/2.0/nfc"
    "github.com/fuzxxl/freefare/0.3/freefare"
    "./keydiversification"
    "./helpers"
)

// Use structs to pass data around so I can refactor
type AppInfo struct {
    aid freefare.DESFireAid
    aidbytes []byte
    sysid []byte
    acl_read_base []byte
    acl_write_base []byte
    acl_file_id byte
}

type KeyChain struct {
    uid_read_key_id byte
    acl_read_key_id byte
    acl_write_key_id byte

    uid_read_key *freefare.DESFireKey
    acl_read_key *freefare.DESFireKey
    acl_write_key *freefare.DESFireKey
}

var (
    keychain = KeyChain{}
    appinfo = AppInfo{}
)

func init_appinfo() {
    keymap, err := helpers.LoadYAMLFile("keys.yaml")
    if err != nil {
        panic(err)
    }

    appmap, err := helpers.LoadYAMLFile("apps.yaml")
    if err != nil {
        panic(err)
    }

    // Application-id
    appinfo.aid, err = helpers.String2aid(appmap["hacklab_acl"].(map[interface{}]interface{})["aid"].(string))
    if err != nil {
        panic(err)
    }

    // Needed for diversification
    appinfo.aidbytes = helpers.Aid2bytes(appinfo.aid)
    appinfo.sysid, err = hex.DecodeString(appmap["hacklab_acl"].(map[interface{}]interface{})["sysid"].(string))
    if err != nil {
        panic(err)
    }

    appinfo.acl_file_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_file_id"].(string))
    if err != nil {
        panic(err)
    }

    // Key id numbers from config
    keychain.uid_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["uid_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_read_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_read_key_id"].(string))
    if err != nil {
        panic(err)
    }
    keychain.acl_write_key_id, err = helpers.String2byte(appmap["hacklab_acl"].(map[interface{}]interface{})["acl_write_key_id"].(string))
    if err != nil {
        panic(err)
    }

    // The static app key to read UID
    keychain.uid_read_key, err = helpers.String2aeskey(keymap["uid_read_key"].(string))
    if err != nil {
        panic(err)
    }

    // Bases for the diversified keys
    appinfo.acl_read_base, err = hex.DecodeString(keymap["acl_read_key"].(string))
    if err != nil {
        panic(err)
    }
    appinfo.acl_write_base, err = hex.DecodeString(keymap["acl_write_key"].(string))
    if err != nil {
        panic(err)
    }

}

func recalculate_diversified_keys(realuid []byte) error {
    acl_read_bytes, err := keydiversification.AES128(appinfo.acl_read_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    acl_write_bytes, err := keydiversification.AES128(appinfo.acl_write_base[:], appinfo.aidbytes[:], realuid[:], appinfo.sysid[:])
    if err != nil {
        return err
    }
    keychain.acl_read_key = helpers.Bytes2aeskey(acl_read_bytes)
    keychain.acl_write_key = helpers.Bytes2aeskey(acl_write_bytes)
    return nil
}

func handle_tag(desfiretag *freefare.DESFireTag) {
    desfiretag.Connect()

    uid_str := desfiretag.UID()

    fmt.Printf("Found tag %s\n", uid_str)

    fmt.Println("Selecting application");
    err := desfiretag.SelectApplication(appinfo.aid);
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    fmt.Println("Authenticating");
    err = desfiretag.Authenticate(keychain.uid_read_key_id,*keychain.uid_read_key)
    if err != nil {
        panic(err)
    }
    fmt.Println("Done");

    realuid_str, err := desfiretag.CardUID()
    if err != nil {
        panic(err)
    }

    realuid, err := hex.DecodeString(realuid_str)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to parse real UID (%s), skipping tag", err))
        return
    }
    fmt.Println("Got UID: ", hex.EncodeToString(realuid));

    // Calculate the diversified keys
    err = recalculate_diversified_keys(realuid)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed to get diversified ACL keys (%s), skipping tag", err))
        return
    }

    fmt.Println("Got real ACL read key: ", keychain.acl_read_key)

    fmt.Print("Re-auth with ACL read key, ")
    err = desfiretag.Authenticate(keychain.acl_read_key_id,*keychain.acl_read_key)
    if err != nil {
        fmt.Println(fmt.Sprintf("ERROR: Failed (%s), skipping tag", err))
        return
    }
    fmt.Println("Done")

    desfiretag.Disconnect()
}

func main() {

    init_appinfo();

    d, err := nfc.Open("");
    if err != nil {
        panic(err)
    }

    for {
        var tags []freefare.Tag
        for {
            tags, err = freefare.GetTags(d);
            if err != nil {
                continue
            }
            if len(tags) > 0 {
                break
            }
            time.Sleep(100 * time.Millisecond)
            //fmt.Println("...polling")
        }

        for i := 0; i < len(tags); i++ {
            tag := tags[i]
            desfiretag := tag.(freefare.DESFireTag)
            handle_tag(&desfiretag)
        }
        //runtime.GC()
    }

}

<|endoftext|>


<commit_before>// Copyright 2016, RadiantBlue Technologies, 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 piazza

import (
	"errors"
	"log"
	"net"
	"net/http"
	"syscall"
	"time"

	"github.com/fvbock/endless"
	"github.com/gin-gonic/gin"
)

const ginHammerTime = 3

type GenericServer struct {
	Sys    *SystemConfig
	pid    int
	router http.Handler
}

type RouteData struct {
	Verb    string
	Path    string
	Handler gin.HandlerFunc
}

func (server *GenericServer) Stop() error {
	sys := server.Sys

	err := syscall.Kill(server.pid, syscall.SIGINT)
	if err != nil {
		return err
	}

	err = sys.WaitForServiceToDieByAddress(sys.Name, sys.BindTo)
	if err != nil {
		return err
	}

	return nil
}

func (server *GenericServer) Start() (chan error, error) {

	sys := server.Sys

	done := make(chan error)

	ready := make(chan bool)

	endless.DefaultHammerTime = ginHammerTime * time.Second

	ginServer := endless.NewServer(server.Sys.BindTo, server.router)

	ginServer.BeforeBegin = func(_ string) {
		server.pid = syscall.Getpid()
		log.Printf("Actual pid is %d", server.pid)

		sys.BindTo = ginServer.EndlessListener.Addr().(*net.TCPAddr).String()

		ready <- true
	}

	go func() {
		err := ginServer.ListenAndServe()
		done <- err
	}()

	<-ready

	err := sys.WaitForServiceByAddress(sys.Name, sys.BindTo)
	if err != nil {
		return nil, err
	}

	log.Printf("Server %s started on %s (%s)", sys.Name, sys.Address, sys.BindTo)

	sys.AddService(sys.Name, sys.BindTo)

	return done, nil
}

func (server *GenericServer) Configure(routeData []RouteData) error {
	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	//router.Use(gin.Logger())
	//router.Use(gin.Recovery())

	for _, data := range routeData {
		switch data.Verb {
		case "GET":
			router.GET(data.Path, data.Handler)
		case "POST":
			router.POST(data.Path, data.Handler)
		case "PUT":
			router.PUT(data.Path, data.Handler)
		case "DELETE":
			router.DELETE(data.Path, data.Handler)
		default:
			return errors.New("Invalid verb: " + data.Verb)
		}
	}

	server.router = router

	return nil
}

/*
// ServerLogHandler adds traditional logging support to the http server handlers.
func ServerLogHandler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}
*/
<commit_msg>tidy up<commit_after>// Copyright 2016, RadiantBlue Technologies, 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 piazza

import (
	"errors"
	"log"
	"net"
	"net/http"
	"syscall"
	"time"

	"github.com/fvbock/endless"
	"github.com/gin-gonic/gin"
)

const ginHammerTime = 3 * time.Second

type GenericServer struct {
	Sys    *SystemConfig
	pid    int
	router http.Handler
}

type RouteData struct {
	Verb    string
	Path    string
	Handler gin.HandlerFunc
}

func (server *GenericServer) Stop() error {
	sys := server.Sys

	err := syscall.Kill(server.pid, syscall.SIGINT)
	if err != nil {
		return err
	}

	err = sys.WaitForServiceToDieByAddress(sys.Name, sys.BindTo)
	if err != nil {
		return err
	}

	return nil
}

func (server *GenericServer) Start() (chan error, error) {

	sys := server.Sys

	done := make(chan error)

	ready := make(chan bool)

	endless.DefaultHammerTime = ginHammerTime

	ginServer := endless.NewServer(server.Sys.BindTo, server.router)

	ginServer.BeforeBegin = func(_ string) {
		server.pid = syscall.Getpid()
		log.Printf("Actual pid is %d", server.pid)

		sys.BindTo = ginServer.EndlessListener.Addr().(*net.TCPAddr).String()

		ready <- true
	}

	go func() {
		err := ginServer.ListenAndServe()
		done <- err
	}()

	<-ready

	err := sys.WaitForServiceByAddress(sys.Name, sys.BindTo)
	if err != nil {
		return nil, err
	}

	log.Printf("Server %s started on %s (%s)", sys.Name, sys.Address, sys.BindTo)

	sys.AddService(sys.Name, sys.BindTo)

	return done, nil
}

func (server *GenericServer) Configure(routeData []RouteData) error {
	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	//router.Use(gin.Logger())
	//router.Use(gin.Recovery())

	for _, data := range routeData {
		switch data.Verb {
		case "GET":
			router.GET(data.Path, data.Handler)
		case "POST":
			router.POST(data.Path, data.Handler)
		case "PUT":
			router.PUT(data.Path, data.Handler)
		case "DELETE":
			router.DELETE(data.Path, data.Handler)
		default:
			return errors.New("Invalid verb: " + data.Verb)
		}
	}

	server.router = router

	return nil
}

/*
// ServerLogHandler adds traditional logging support to the http server handlers.
func ServerLogHandler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}
*/
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/PuerkitoBio/goquery"
	"github.com/SlyMarbo/rss"
)

type Download struct {
	rss.Item

	URL url.URL
}

func main() {
	feed, err := rss.Fetch("https://www.reddit.com/r/listentothis/.rss")
	if err != nil {
		panic(err)
	}

	var available []*Download
	for _, item := range feed.Items {
		download, err := ParseDownload(item)
		if err != nil {
			log.Println("don't know how to download %q: %v", item.ID, err)
		}
		available = append(available, download)
	}

	for _, dl := range available {
		fmt.Println(dl.URL)
	}
}

func ParseDownload(item *rss.Item) (*Download, error) {
	doc, err := goquery.NewDocumentFromReader(bytes.NewBufferString(item.Content))
	if err != nil {
		return nil, err
	}
	a := doc.Find("a:contains('[link]')")
	if a.Length() == 0 {
		return nil, fmt.Errorf("download link not found")
	}
	href, ok := a.Attr("href")
	if !ok {
		return nil, fmt.Errorf("missing expected 'href' attribute in element")
	}
	u, err := url.Parse(href)
	if err != nil {
		return nil, err
	}
	if !isYouTubeURL(u) {
		return nil, fmt.Errorf("unsupported download URL: %q", u)
	}

	return &Download{
		Item: *item,
		URL:  *u,
	}, nil
}
<commit_msg>Adding library to keep track of what we've downloaded.<commit_after>package main

import (
	"bytes"
	"fmt"
	"log"
	"net/url"
	"strings"
	"unicode"

	"github.com/PuerkitoBio/goquery"
	"github.com/SlyMarbo/rss"
)

type Download struct {
	rss.Item

	URL url.URL
}

type Library struct {
	*bolt.DB

	Path string
}

func main() {
	feed, err := rss.Fetch("https://www.reddit.com/r/listentothis/.rss")
	if err != nil {
		panic(err)
	}

	var available []*Download
	for _, item := range feed.Items {
		download, err := ParseDownload(item)
		if err != nil {
			log.Printf("don't know how to download %q: %v", item.ID, err)
		} else {
			available = append(available, download)
		}
	}

	lib, err := NewLibrary(defaultPath())
	if err != nil {
		panic(err)
	}
	defer lib.Close()

	for _, dl := range available {
		if !lib.Contains(dl) {
			err := lib.Archive(dl)
			if err != nil {
				log.Println("failed to archive %q: %v", dl.ID, err)
			}
		}
	}
}

func NewLibrary(path string) (*Library, error) {
	dbpath := filepath.Join(path, ".history")
	db, err := bolt.Open(dbpath, 0600, nil)
	if err != nil {
		return nil, err
	}
	return &Library{
		DB:   db,
		Path: path,
	}, nil
}

func (l *Library) Archive(dl *Download) error {
	return l.Update(func(tx *bolt.Tx) error {
		cmd := exec.Command("youtube-dl", "-x", "--audio-format", "mp3", dl.URL.String())
		cmd.Dir = l.Path
		err := cmd.Run()
		if err != nil {
			return err
		}

		b, err := tx.CreateBucketIfNotExists([]byte("downloaded"))
		if err != nil {
			return err
		}
		if b.Get([]byte(dl.ID)) != nil {
			return fmt.Errorf("already downloaded %q", dl.ID)
		}

		err = b.Put([]byte(dl.ID, []byte(dl.URL.String())))
		if err != nil {
			return err
		}

		return nil
	})
}

func ParseDownload(item *rss.Item) (*Download, error) {
	doc, err := goquery.NewDocumentFromReader(bytes.NewBufferString(item.Content))
	if err != nil {
		return nil, err
	}
	a := doc.Find("a:contains('[link]')")
	if a.Length() == 0 {
		return nil, fmt.Errorf("download link not found")
	}
	href, ok := a.Attr("href")
	if !ok {
		return nil, fmt.Errorf("missing expected 'href' attribute in element")
	}
	u, err := url.Parse(href)
	if err != nil {
		return nil, err
	}
	if !isSupportedURL(u) {
		return nil, fmt.Errorf("unsupported download URL: %q", u)
	}

	return &Download{
		Item: *item,
		Info: songInfo,
		URL:  *u,
	}, nil
}

func isSupportedURL(u *url.URL) bool {
	s := strings.Map(func(r rune) rune {
		if unicode.IsLetter(r) {
			return r
		}
		return rune(-1)
	}, u.Host)
	if strings.Contains(s, "youtube") {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"unicode/utf8"

	"github.com/charlievieth/gocode"
	"github.com/golang/groupcache/lru"
)

type GoCode struct {
	Autoinst      bool
	InstallSuffix string
	Env           map[string]string
	Home          string
	Dir           string
	Builtins      bool
	Fn            string
	Src           string
	Pos           int
	calltip       bool `json:"-"` // ignore
}

// Separate type so we can init Calltip to true
type GoCodeCalltip struct {
	GoCode
}

func (g *GoCodeCalltip) Init() {
	g.calltip = true
}

var calltipCache = NewAstCache(200)

func init() {
	registry.Register("gocode_complete", func(b *Broker) Caller {
		return &GoCode{}
	})

	registry.Register("gocode_calltip", func(b *Broker) Caller {
		return &GoCode{calltip: true}
	})
}

type GoCodeResponse struct {
	Candidates []gocode.Candidate
}

var lastCalltip struct {
	sync.Mutex
	Path string
	Src  string
	Off  int
	Res  GoCodeResponse
}

func useLastCalltip(path, src string, off int) (res GoCodeResponse, ok bool) {
	lastCalltip.Lock()
	if off == lastCalltip.Off && path == lastCalltip.Path && src == lastCalltip.Src {
		res = lastCalltip.Res
		ok = true
	}
	lastCalltip.Unlock()
	return
}

func setLastCalltip(path, src string, off int, res GoCodeResponse) {
	lastCalltip.Lock()
	if off != lastCalltip.Off || path != lastCalltip.Path || src != lastCalltip.Src {
		lastCalltip.Off = off
		lastCalltip.Path = path
		lastCalltip.Src = src
		lastCalltip.Res = res
	}
	lastCalltip.Unlock()
	return
}

func (g *GoCode) Call() (interface{}, string) {
	off, err := g.bytePos()
	if err != nil {
		return g.response(nil, err, false)
	}
	path := g.filepath()
	if g.calltip {
		if res, ok := useLastCalltip(path, g.Src, off); ok {
			return res, ""
		}
		gr, err := g.calltips(path, []byte(g.Src), off)
		res, errStr := g.response(gr, err, false)
		setLastCalltip(path, g.Src, off, res)
		return res, errStr
	}
	return g.response(g.complete(path, []byte(g.Src), off), nil, true)
}

var NoGocodeCandidates = []gocode.Candidate{}

func (g *GoCode) response(res []gocode.Candidate, err error, install bool) (GoCodeResponse, string) {
	if res == nil || len(res) == 0 {
		if install && g.Autoinst {
			autoInstall(AutoInstOptions{
				Src:           g.Src,
				Env:           g.Env,
				InstallSuffix: g.InstallSuffix,
			})
		}
		res = NoGocodeCandidates
	}
	var errStr string
	if err != nil {
		errStr = err.Error()
	}
	return GoCodeResponse{Candidates: res}, errStr
}

// Matching GoSublime's behavior here...
func (g *GoCode) filepath() string {
	if filepath.IsAbs(g.Fn) {
		return filepath.Clean(g.Fn)
	}
	var base string
	if g.Dir != "" {
		base = g.Dir
	} else {
		base = g.Home
	}
	var name string
	if g.Fn != "" {
		name = g.Fn
	} else {
		name = "_.go"
	}
	return filepath.Join(base, name)
}

func (g *GoCode) complete(filename string, src []byte, offset int) []gocode.Candidate {
	conf := gocode.Config{
		GOROOT:        g.GOROOT(),
		GOPATH:        g.GOPATH(),
		InstallSuffix: g.InstallSuffix,
		Builtins:      g.Builtins,
	}
	return conf.Complete(src, filename, offset)
}

func (g *GoCode) GOPATH() string {
	if g.Env != nil {
		if s, ok := g.Env["GOPATH"]; ok && s != "" {
			return s
		}
	}
	return os.Getenv("GOPATH")
}

func (g *GoCode) GOROOT() string {
	if g.Env != nil {
		if s, ok := g.Env["GOROOT"]; ok && s != "" {
			return s
		}
	}
	return runtime.GOROOT()
}

// TODO: cache results and potentially cache AST ad file set.
func (g *GoCode) calltips(filename string, src []byte, offset int) ([]gocode.Candidate, error) {
	fset, af, err := calltipCache.ParseFile(filename, src)
	if af == nil {
		return nil, err
	}
	pos, err := g.convertOffset(af, fset, offset)
	if err != nil {
		return nil, err
	}
	id, err := g.findCaller(af, fset, pos)
	if err != nil {
		return nil, err
	}
	end := fset.Position(id.End())
	if !end.IsValid() {
		return nil, errors.New("calltip: invalid end pos")
	}
	if err := g.validLine(fset, pos, end); err != nil {
		return nil, err
	}
	cl := g.complete(filename, src, end.Offset)
	for i := 0; i < len(cl); i++ {
		if strings.EqualFold(id.Name, cl[i].Name) {
			return cl[i : i+1], nil
		}
	}
	return nil, errors.New("calltip: no candidates")
}

// Matching GoSublime's behavior here, see if this can be removed.
func (g *GoCode) validLine(fset *token.FileSet, off token.Pos, end token.Position) error {
	line := fset.Position(off).Line
	if end.Line == line || line == 0 {
		return nil
	}
	return fmt.Errorf("calltip: line mismatch end (%d) line (%d)", end.Line, line)
}

func (g *GoCode) convertOffset(af *ast.File, fset *token.FileSet, offset int) (token.Pos, error) {
	f := fset.File(af.Pos())
	if f == nil {
		return token.NoPos, errors.New("calltip: ast file not in fileset")
	}
	// Prevent Pos from panicking
	if offset > f.Size() {
		return token.NoPos, errors.New("calltip: illegal file offset")
	}
	return f.Pos(offset), nil
}

func (g *GoCode) findCaller(af *ast.File, fset *token.FileSet, pos token.Pos) (*ast.Ident, error) {
	v := calltipVisitor{
		pos:  pos,
		fset: fset,
	}
	ast.Walk(&v, af)
	if v.expr == nil || v.expr.Fun == nil {
		return nil, errors.New("calltip: nil CallExpr")
	}
	switch x := v.expr.Fun.(type) {
	case *ast.Ident:
		return x, nil
	case *ast.SelectorExpr:
		return x.Sel, nil
	}
	return nil, errors.New("calltip: invalid CallExpr.Fun")
}

// calltipVisitor satisfies the ast.Visitor interface.
type calltipVisitor struct {
	pos  token.Pos
	fset *token.FileSet
	expr *ast.CallExpr
}

func (v *calltipVisitor) Visit(node ast.Node) (w ast.Visitor) {
	if node == nil || v.expr != nil {
		return v
	}
	if x, ok := node.(*ast.CallExpr); ok {
		pos := x.Pos()
		end := x.End()
		if pos.IsValid() && end.IsValid() {
			if pos <= v.pos && v.pos <= end {
				v.expr = x
				return nil
			}
		}
	}
	return v
}

func (g *GoCode) bytePos() (int, error) {
	s := g.Src
	off := g.Pos
	if len(s) == 0 || len(s) <= off || off < 0 {
		return -1, errors.New("gocode: nil source")
	}
	i := 0
	var n int
	for n = 0; i < len(s) && n < off; n++ {
		if s[i] < utf8.RuneSelf {
			i++
		} else {
			_, size := utf8.DecodeRuneInString(s[i:])
			i += size
		}
	}
	if n == off && i < len(s) {
		return i, nil
	}
	return -1, fmt.Errorf("gocode: invalid offset: %d", g.Pos)
}

type AstEntry struct {
	Name    string
	Body    []byte
	File    *ast.File
	FileSet *token.FileSet
}

type AstCache struct {
	cache *lru.Cache
	mu    sync.Mutex
}

func NewAstCache(maxEntries int) *AstCache {
	return &AstCache{
		cache: lru.New(maxEntries),
	}
}

func (c *AstCache) Add(filename string, ent *AstEntry) {
	c.mu.Lock()
	c.cache.Add(filename, ent)
	c.mu.Unlock()
}

func (c *AstCache) Get(filename string) (*AstEntry, bool) {
	c.mu.Lock()
	v, ok := c.cache.Get(filename)
	c.mu.Unlock()
	if ok {
		return v.(*AstEntry), true
	}
	return nil, false
}

func (c *AstCache) Remove(filename string) {
	c.mu.Lock()
	c.cache.Remove(filename)
	c.mu.Unlock()
}

func (c *AstCache) ParseFile(filename string, src interface{}) (*token.FileSet, *ast.File, error) {
	text, err := readSource(filename, src)
	if err != nil {
		return nil, nil, err
	}
	if e, ok := c.Get(filename); ok {
		if bytes.Equal(e.Body, text) {
			return e.FileSet, e.File, nil
		}
		c.Remove(filename)
	}
	fset := token.NewFileSet()
	af, err := parser.ParseFile(fset, filename, text, 0)
	if af == nil {
		return nil, nil, err
	}
	c.Add(filename, &AstEntry{
		Name:    filename,
		Body:    text,
		File:    af,
		FileSet: fset,
	})
	return fset, af, nil
}

func readSource(filename string, src interface{}) ([]byte, error) {
	if src != nil {
		switch s := src.(type) {
		case string:
			return []byte(s), nil
		case []byte:
			return s, nil
		case *bytes.Buffer:
			// is io.Reader, but src is already available in []byte form
			if s != nil {
				return s.Bytes(), nil
			}
		case io.Reader:
			var buf bytes.Buffer
			if _, err := io.Copy(&buf, s); err != nil {
				return nil, err
			}
			return buf.Bytes(), nil
		}
		return nil, errors.New("invalid source")
	}
	return ioutil.ReadFile(filename)
}
<commit_msg>gocode: fix calltips for nested functions<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"unicode/utf8"

	"github.com/charlievieth/gocode"
	"github.com/golang/groupcache/lru"
)

type GoCode struct {
	Autoinst      bool
	InstallSuffix string
	Env           map[string]string
	Home          string
	Dir           string
	Builtins      bool
	Fn            string
	Src           string
	Pos           int
	calltip       bool `json:"-"` // ignore
}

// Separate type so we can init Calltip to true
type GoCodeCalltip struct {
	GoCode
}

func (g *GoCodeCalltip) Init() {
	g.calltip = true
}

var calltipCache = NewAstCache(200)

func init() {
	registry.Register("gocode_complete", func(b *Broker) Caller {
		return &GoCode{}
	})

	registry.Register("gocode_calltip", func(b *Broker) Caller {
		return &GoCode{calltip: true}
	})
}

type GoCodeResponse struct {
	Candidates []gocode.Candidate
}

var lastCalltip struct {
	sync.Mutex
	Path string
	Src  string
	Off  int
	Res  GoCodeResponse
}

func useLastCalltip(path, src string, off int) (res GoCodeResponse, ok bool) {
	lastCalltip.Lock()
	if off == lastCalltip.Off && path == lastCalltip.Path && src == lastCalltip.Src {
		res = lastCalltip.Res
		ok = true
	}
	lastCalltip.Unlock()
	return
}

func setLastCalltip(path, src string, off int, res GoCodeResponse) {
	lastCalltip.Lock()
	if off != lastCalltip.Off || path != lastCalltip.Path || src != lastCalltip.Src {
		lastCalltip.Off = off
		lastCalltip.Path = path
		lastCalltip.Src = src
		lastCalltip.Res = res
	}
	lastCalltip.Unlock()
	return
}

func (g *GoCode) Call() (interface{}, string) {
	off, err := g.bytePos()
	if err != nil {
		return g.response(nil, err, false)
	}
	path := g.filepath()
	if g.calltip {
		if res, ok := useLastCalltip(path, g.Src, off); ok {
			return res, ""
		}
		gr, err := g.calltips(path, []byte(g.Src), off)
		res, errStr := g.response(gr, err, false)
		setLastCalltip(path, g.Src, off, res)
		return res, errStr
	}
	return g.response(g.complete(path, []byte(g.Src), off), nil, true)
}

var NoGocodeCandidates = []gocode.Candidate{}

func (g *GoCode) response(res []gocode.Candidate, err error, install bool) (GoCodeResponse, string) {
	if res == nil || len(res) == 0 {
		if install && g.Autoinst {
			autoInstall(AutoInstOptions{
				Src:           g.Src,
				Env:           g.Env,
				InstallSuffix: g.InstallSuffix,
			})
		}
		res = NoGocodeCandidates
	}
	var errStr string
	if err != nil {
		errStr = err.Error()
	}
	return GoCodeResponse{Candidates: res}, errStr
}

// Matching GoSublime's behavior here...
func (g *GoCode) filepath() string {
	if filepath.IsAbs(g.Fn) {
		return filepath.Clean(g.Fn)
	}
	var base string
	if g.Dir != "" {
		base = g.Dir
	} else {
		base = g.Home
	}
	var name string
	if g.Fn != "" {
		name = g.Fn
	} else {
		name = "_.go"
	}
	return filepath.Join(base, name)
}

func (g *GoCode) complete(filename string, src []byte, offset int) []gocode.Candidate {
	conf := gocode.Config{
		GOROOT:        g.GOROOT(),
		GOPATH:        g.GOPATH(),
		InstallSuffix: g.InstallSuffix,
		Builtins:      g.Builtins,
	}
	return conf.Complete(src, filename, offset)
}

func (g *GoCode) GOPATH() string {
	if g.Env != nil {
		if s, ok := g.Env["GOPATH"]; ok && s != "" {
			return s
		}
	}
	return os.Getenv("GOPATH")
}

func (g *GoCode) GOROOT() string {
	if g.Env != nil {
		if s, ok := g.Env["GOROOT"]; ok && s != "" {
			return s
		}
	}
	return runtime.GOROOT()
}

// TODO: cache results and potentially cache AST ad file set.
func (g *GoCode) calltips(filename string, src []byte, offset int) ([]gocode.Candidate, error) {
	fset, af, err := calltipCache.ParseFile(filename, src)
	if af == nil {
		return nil, err
	}
	pos, err := g.convertOffset(af, fset, offset)
	if err != nil {
		return nil, err
	}
	id, err := g.findCaller(af, fset, pos)
	if err != nil {
		return nil, err
	}
	end := fset.Position(id.End())
	if !end.IsValid() {
		return nil, errors.New("calltip: invalid end pos")
	}
	if err := g.validLine(fset, pos, end); err != nil {
		return nil, err
	}
	cl := g.complete(filename, src, end.Offset)
	for i := 0; i < len(cl); i++ {
		if strings.EqualFold(id.Name, cl[i].Name) {
			return cl[i : i+1], nil
		}
	}
	return nil, errors.New("calltip: no candidates")
}

// Matching GoSublime's behavior here, see if this can be removed.
func (g *GoCode) validLine(fset *token.FileSet, off token.Pos, end token.Position) error {
	line := fset.Position(off).Line
	if end.Line == line || line == 0 {
		return nil
	}
	return fmt.Errorf("calltip: line mismatch end (%d) line (%d)", end.Line, line)
}

func (g *GoCode) convertOffset(af *ast.File, fset *token.FileSet, offset int) (token.Pos, error) {
	f := fset.File(af.Pos())
	if f == nil {
		return token.NoPos, errors.New("calltip: ast file not in fileset")
	}
	// Prevent Pos from panicking
	if offset > f.Size() {
		return token.NoPos, errors.New("calltip: illegal file offset")
	}
	return f.Pos(offset), nil
}

func (g *GoCode) findCaller(af *ast.File, fset *token.FileSet, pos token.Pos) (*ast.Ident, error) {
	v := calltipVisitor{
		pos:  pos,
		fset: fset,
	}
	ast.Walk(&v, af)
	if v.expr == nil || v.expr.Fun == nil {
		return nil, errors.New("calltip: nil CallExpr")
	}
	switch x := v.expr.Fun.(type) {
	case *ast.Ident:
		return x, nil
	case *ast.SelectorExpr:
		return x.Sel, nil
	}
	return nil, errors.New("calltip: invalid CallExpr.Fun")
}

// calltipVisitor satisfies the ast.Visitor interface.
type calltipVisitor struct {
	pos  token.Pos
	fset *token.FileSet
	expr *ast.CallExpr
}

func (v *calltipVisitor) Visit(node ast.Node) (w ast.Visitor) {
	if node == nil {
		return v
	}
	if x, ok := node.(*ast.CallExpr); ok {
		pos := x.Pos()
		end := x.End()
		if pos.IsValid() && end.IsValid() {
			if pos <= v.pos && v.pos <= end {
				v.expr = x
			}
		}
	}
	return v
}

func (g *GoCode) bytePos() (int, error) {
	s := g.Src
	off := g.Pos
	if len(s) == 0 || len(s) <= off || off < 0 {
		return -1, errors.New("gocode: nil source")
	}
	i := 0
	var n int
	for n = 0; i < len(s) && n < off; n++ {
		if s[i] < utf8.RuneSelf {
			i++
		} else {
			_, size := utf8.DecodeRuneInString(s[i:])
			i += size
		}
	}
	if n == off && i < len(s) {
		return i, nil
	}
	return -1, fmt.Errorf("gocode: invalid offset: %d", g.Pos)
}

type AstEntry struct {
	Name    string
	Body    []byte
	File    *ast.File
	FileSet *token.FileSet
}

type AstCache struct {
	cache *lru.Cache
	mu    sync.Mutex
}

func NewAstCache(maxEntries int) *AstCache {
	return &AstCache{
		cache: lru.New(maxEntries),
	}
}

func (c *AstCache) Add(filename string, ent *AstEntry) {
	c.mu.Lock()
	c.cache.Add(filename, ent)
	c.mu.Unlock()
}

func (c *AstCache) Get(filename string) (*AstEntry, bool) {
	c.mu.Lock()
	v, ok := c.cache.Get(filename)
	c.mu.Unlock()
	if ok {
		return v.(*AstEntry), true
	}
	return nil, false
}

func (c *AstCache) Remove(filename string) {
	c.mu.Lock()
	c.cache.Remove(filename)
	c.mu.Unlock()
}

func (c *AstCache) ParseFile(filename string, src interface{}) (*token.FileSet, *ast.File, error) {
	text, err := readSource(filename, src)
	if err != nil {
		return nil, nil, err
	}
	if e, ok := c.Get(filename); ok {
		if bytes.Equal(e.Body, text) {
			return e.FileSet, e.File, nil
		}
		c.Remove(filename)
	}
	fset := token.NewFileSet()
	af, err := parser.ParseFile(fset, filename, text, 0)
	if af == nil {
		return nil, nil, err
	}
	c.Add(filename, &AstEntry{
		Name:    filename,
		Body:    text,
		File:    af,
		FileSet: fset,
	})
	return fset, af, nil
}

func readSource(filename string, src interface{}) ([]byte, error) {
	if src != nil {
		switch s := src.(type) {
		case string:
			return []byte(s), nil
		case []byte:
			return s, nil
		case *bytes.Buffer:
			// is io.Reader, but src is already available in []byte form
			if s != nil {
				return s.Bytes(), nil
			}
		case io.Reader:
			var buf bytes.Buffer
			if _, err := io.Copy(&buf, s); err != nil {
				return nil, err
			}
			return buf.Bytes(), nil
		}
		return nil, errors.New("invalid source")
	}
	return ioutil.ReadFile(filename)
}
<|endoftext|>


<commit_before>// Copyright 2013 Unknwon
//
// 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 goconfig

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"
)

// Read reads an io.Reader and returns a configuration representation.
// This representation can be queried with GetValue.
func (c *ConfigFile) read(reader io.Reader) (err error) {
	buf := bufio.NewReader(reader)

	// Handle BOM-UTF8.
	// http://en.wikipedia.org/wiki/Byte_order_mark#Representations_of_byte_order_marks_by_encoding
	mask, err := buf.Peek(3)
	if err == nil && len(mask) >= 3 &&
		mask[0] == 239 && mask[1] == 187 && mask[2] == 191 {
		buf.Read(mask)
	}

	count := 1 // Counter for auto increment.
	// Current section name.
	section := DEFAULT_SECTION
	var comments string
	// Parse line-by-line
	for {
		line, err := buf.ReadString('\n')
		line = strings.TrimSpace(line)
		lineLengh := len(line) //[SWH|+]
		if err != nil {
			if err != io.EOF {
				return err
			}

			// Reached end of file, if nothing to read then break,
			// otherwise handle the last line.
			if lineLengh == 0 {
				break
			}
		}

		// switch written for readability (not performance)
		switch {
		case lineLengh == 0: // Empty line
			continue
		case line[0] == '#' || line[0] == ';': // Comment
			// Append comments
			if len(comments) == 0 {
				comments = line
			} else {
				comments += LineBreak + line
			}
			continue
		case line[0] == '[' && line[lineLengh-1] == ']': // New sction.
			// Get section name.
			section = strings.TrimSpace(line[1 : lineLengh-1])
			// Set section comments and empty if it has comments.
			if len(comments) > 0 {
				c.SetSectionComments(section, comments)
				comments = ""
			}
			// Make section exist even though it does not have any key.
			c.SetValue(section, " ", " ")
			// Reset counter.
			count = 1
			continue
		case section == "": // No section defined so far
			return readError{ERR_BLANK_SECTION_NAME, line}
		default: // Other alternatives
			var (
				i        int
				keyQuote string
				key      string
				valQuote string
				value    string
			)
			//[SWH|+]:支持引号包围起来的字串
			if line[0] == '"' {
				if lineLengh >= 6 && line[0:3] == `"""` {
					keyQuote = `"""`
				} else {
					keyQuote = `"`
				}
			} else if line[0] == '`' {
				keyQuote = "`"
			}
			if keyQuote != "" {
				qLen := len(keyQuote)
				pos := strings.Index(line[qLen:], keyQuote)
				if pos == -1 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				pos = pos + qLen
				i = strings.IndexAny(line[pos:], "=:")
				if i <= 0 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				i = i + pos
				key = line[qLen:pos] //保留引号内的两端的空格
			} else {
				i = strings.IndexAny(line, "=:")
				if i <= 0 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				key = strings.TrimSpace(line[0:i])
			}
			//[SWH|+];

			// Check if it needs auto increment.
			if key == "-" {
				key = "#" + fmt.Sprint(count)
				count++
			}

			//[SWH|+]:支持引号包围起来的字串
			lineRight := strings.TrimSpace(line[i+1:])
			lineRightLength := len(lineRight)
			firstChar := ""
			if lineRightLength >= 2 {
				firstChar = lineRight[0:1]
			}
			if firstChar == "`" {
				valQuote = "`"
			} else if lineRightLength >= 6 && lineRight[0:3] == `"""` {
				valQuote = `"""`
			}
			if valQuote != "" {
				qLen := len(valQuote)
				pos := strings.LastIndex(lineRight[qLen:], valQuote)
				if pos == -1 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				pos = pos + qLen
				value = lineRight[qLen:pos]
			} else {
				value = strings.TrimSpace(lineRight[0:])
			}
			//[SWH|+];

			c.SetValue(section, key, value)
			// Set key comments and empty if it has comments.
			if len(comments) > 0 {
				c.SetKeyComments(section, key, comments)
				comments = ""
			}
		}

		// Reached end of file.
		if err == io.EOF {
			break
		}
	}
	return nil
}

// LoadFromData accepts raw data directly from memory
// and returns a new configuration representation.
// Note that the configuration is written to the system
// temporary folder, so your file should not contain
// sensitive information.
func LoadFromData(data []byte) (c *ConfigFile, err error) {
	// Save memory data to temporary file to support further operations.
	tmpName := path.Join(os.TempDir(), "goconfig", fmt.Sprintf("%d", time.Now().Nanosecond()))
	if err = os.MkdirAll(path.Dir(tmpName), os.ModePerm); err != nil {
		return nil, err
	}
	if err = ioutil.WriteFile(tmpName, data, 0655); err != nil {
		return nil, err
	}

	c = newConfigFile([]string{tmpName})
	err = c.read(bytes.NewBuffer(data))
	return c, err
}

// LoadFromReader accepts raw data directly from a reader
// and returns a new configuration representation.
// You must use ReloadData to reload.
// You cannot append files a configfile read this way.
func LoadFromReader(in io.Reader) (c *ConfigFile, err error) {
	c = newConfigFile([]string{""})
	err = c.read(in)
	return c, err
}

func (c *ConfigFile) loadFile(fileName string) (err error) {
	f, err := os.Open(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	return c.read(f)
}

// LoadConfigFile reads a file and returns a new configuration representation.
// This representation can be queried with GetValue.
func LoadConfigFile(fileName string, moreFiles ...string) (c *ConfigFile, err error) {
	// Append files' name together.
	fileNames := make([]string, 1, len(moreFiles)+1)
	fileNames[0] = fileName
	if len(moreFiles) > 0 {
		fileNames = append(fileNames, moreFiles...)
	}

	c = newConfigFile(fileNames)

	for _, name := range fileNames {
		if err = c.loadFile(name); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Reload reloads configuration file in case it has changes.
func (c *ConfigFile) Reload() (err error) {
	var cfg *ConfigFile
	if len(c.fileNames) == 1 {
		if c.fileNames[0] == "" {
			return fmt.Errorf("file opened from in-memory data, use ReloadData to reload")
		}
		cfg, err = LoadConfigFile(c.fileNames[0])
	} else {
		cfg, err = LoadConfigFile(c.fileNames[0], c.fileNames[1:]...)
	}

	if err == nil {
		*c = *cfg
	}
	return err
}

// ReloadData reloads configuration file from memory
func (c *ConfigFile) ReloadData(in io.Reader) (err error) {
	var cfg *ConfigFile
	if len(c.fileNames) != 1 {
		return fmt.Errorf("Multiple files loaded, unable to mix in-memory and file data")
	}

	cfg, err = LoadFromReader(in)
	if err == nil {
		*c = *cfg
	}
	return err
}

// AppendFiles appends more files to ConfigFile and reload automatically.
func (c *ConfigFile) AppendFiles(files ...string) error {
	if len(c.fileNames) == 1 && c.fileNames[0] == "" {
		return fmt.Errorf("Cannot append file data to in-memory data")
	}
	c.fileNames = append(c.fileNames, files...)
	return c.Reload()
}

// readError occurs when read configuration file with wrong format.
type readError struct {
	Reason  ParseError
	Content string // Line content
}

// Error implement Error interface.
func (err readError) Error() string {
	switch err.Reason {
	case ERR_BLANK_SECTION_NAME:
		return "empty section name not allowed"
	case ERR_COULD_NOT_PARSE:
		return fmt.Sprintf("could not parse line: %s", string(err.Content))
	}
	return "invalid read error"
}
<commit_msg>fix_comments_word<commit_after>// Copyright 2013 Unknwon
//
// 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 goconfig

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"
)

// Read reads an io.Reader and returns a configuration representation.
// This representation can be queried with GetValue.
func (c *ConfigFile) read(reader io.Reader) (err error) {
	buf := bufio.NewReader(reader)

	// Handle BOM-UTF8.
	// http://en.wikipedia.org/wiki/Byte_order_mark#Representations_of_byte_order_marks_by_encoding
	mask, err := buf.Peek(3)
	if err == nil && len(mask) >= 3 &&
		mask[0] == 239 && mask[1] == 187 && mask[2] == 191 {
		buf.Read(mask)
	}

	count := 1 // Counter for auto increment.
	// Current section name.
	section := DEFAULT_SECTION
	var comments string
	// Parse line-by-line
	for {
		line, err := buf.ReadString('\n')
		line = strings.TrimSpace(line)
		lineLengh := len(line) //[SWH|+]
		if err != nil {
			if err != io.EOF {
				return err
			}

			// Reached end of file, if nothing to read then break,
			// otherwise handle the last line.
			if lineLengh == 0 {
				break
			}
		}

		// switch written for readability (not performance)
		switch {
		case lineLengh == 0: // Empty line
			continue
		case line[0] == '#' || line[0] == ';': // Comment
			// Append comments
			if len(comments) == 0 {
				comments = line
			} else {
				comments += LineBreak + line
			}
			continue
		case line[0] == '[' && line[lineLengh-1] == ']': // New section.
			// Get section name.
			section = strings.TrimSpace(line[1 : lineLengh-1])
			// Set section comments and empty if it has comments.
			if len(comments) > 0 {
				c.SetSectionComments(section, comments)
				comments = ""
			}
			// Make section exist even though it does not have any key.
			c.SetValue(section, " ", " ")
			// Reset counter.
			count = 1
			continue
		case section == "": // No section defined so far
			return readError{ERR_BLANK_SECTION_NAME, line}
		default: // Other alternatives
			var (
				i        int
				keyQuote string
				key      string
				valQuote string
				value    string
			)
			//[SWH|+]:支持引号包围起来的字串
			if line[0] == '"' {
				if lineLengh >= 6 && line[0:3] == `"""` {
					keyQuote = `"""`
				} else {
					keyQuote = `"`
				}
			} else if line[0] == '`' {
				keyQuote = "`"
			}
			if keyQuote != "" {
				qLen := len(keyQuote)
				pos := strings.Index(line[qLen:], keyQuote)
				if pos == -1 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				pos = pos + qLen
				i = strings.IndexAny(line[pos:], "=:")
				if i <= 0 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				i = i + pos
				key = line[qLen:pos] //保留引号内的两端的空格
			} else {
				i = strings.IndexAny(line, "=:")
				if i <= 0 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				key = strings.TrimSpace(line[0:i])
			}
			//[SWH|+];

			// Check if it needs auto increment.
			if key == "-" {
				key = "#" + fmt.Sprint(count)
				count++
			}

			//[SWH|+]:支持引号包围起来的字串
			lineRight := strings.TrimSpace(line[i+1:])
			lineRightLength := len(lineRight)
			firstChar := ""
			if lineRightLength >= 2 {
				firstChar = lineRight[0:1]
			}
			if firstChar == "`" {
				valQuote = "`"
			} else if lineRightLength >= 6 && lineRight[0:3] == `"""` {
				valQuote = `"""`
			}
			if valQuote != "" {
				qLen := len(valQuote)
				pos := strings.LastIndex(lineRight[qLen:], valQuote)
				if pos == -1 {
					return readError{ERR_COULD_NOT_PARSE, line}
				}
				pos = pos + qLen
				value = lineRight[qLen:pos]
			} else {
				value = strings.TrimSpace(lineRight[0:])
			}
			//[SWH|+];

			c.SetValue(section, key, value)
			// Set key comments and empty if it has comments.
			if len(comments) > 0 {
				c.SetKeyComments(section, key, comments)
				comments = ""
			}
		}

		// Reached end of file.
		if err == io.EOF {
			break
		}
	}
	return nil
}

// LoadFromData accepts raw data directly from memory
// and returns a new configuration representation.
// Note that the configuration is written to the system
// temporary folder, so your file should not contain
// sensitive information.
func LoadFromData(data []byte) (c *ConfigFile, err error) {
	// Save memory data to temporary file to support further operations.
	tmpName := path.Join(os.TempDir(), "goconfig", fmt.Sprintf("%d", time.Now().Nanosecond()))
	if err = os.MkdirAll(path.Dir(tmpName), os.ModePerm); err != nil {
		return nil, err
	}
	if err = ioutil.WriteFile(tmpName, data, 0655); err != nil {
		return nil, err
	}

	c = newConfigFile([]string{tmpName})
	err = c.read(bytes.NewBuffer(data))
	return c, err
}

// LoadFromReader accepts raw data directly from a reader
// and returns a new configuration representation.
// You must use ReloadData to reload.
// You cannot append files a configfile read this way.
func LoadFromReader(in io.Reader) (c *ConfigFile, err error) {
	c = newConfigFile([]string{""})
	err = c.read(in)
	return c, err
}

func (c *ConfigFile) loadFile(fileName string) (err error) {
	f, err := os.Open(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	return c.read(f)
}

// LoadConfigFile reads a file and returns a new configuration representation.
// This representation can be queried with GetValue.
func LoadConfigFile(fileName string, moreFiles ...string) (c *ConfigFile, err error) {
	// Append files' name together.
	fileNames := make([]string, 1, len(moreFiles)+1)
	fileNames[0] = fileName
	if len(moreFiles) > 0 {
		fileNames = append(fileNames, moreFiles...)
	}

	c = newConfigFile(fileNames)

	for _, name := range fileNames {
		if err = c.loadFile(name); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Reload reloads configuration file in case it has changes.
func (c *ConfigFile) Reload() (err error) {
	var cfg *ConfigFile
	if len(c.fileNames) == 1 {
		if c.fileNames[0] == "" {
			return fmt.Errorf("file opened from in-memory data, use ReloadData to reload")
		}
		cfg, err = LoadConfigFile(c.fileNames[0])
	} else {
		cfg, err = LoadConfigFile(c.fileNames[0], c.fileNames[1:]...)
	}

	if err == nil {
		*c = *cfg
	}
	return err
}

// ReloadData reloads configuration file from memory
func (c *ConfigFile) ReloadData(in io.Reader) (err error) {
	var cfg *ConfigFile
	if len(c.fileNames) != 1 {
		return fmt.Errorf("Multiple files loaded, unable to mix in-memory and file data")
	}

	cfg, err = LoadFromReader(in)
	if err == nil {
		*c = *cfg
	}
	return err
}

// AppendFiles appends more files to ConfigFile and reload automatically.
func (c *ConfigFile) AppendFiles(files ...string) error {
	if len(c.fileNames) == 1 && c.fileNames[0] == "" {
		return fmt.Errorf("Cannot append file data to in-memory data")
	}
	c.fileNames = append(c.fileNames, files...)
	return c.Reload()
}

// readError occurs when read configuration file with wrong format.
type readError struct {
	Reason  ParseError
	Content string // Line content
}

// Error implement Error interface.
func (err readError) Error() string {
	switch err.Reason {
	case ERR_BLANK_SECTION_NAME:
		return "empty section name not allowed"
	case ERR_COULD_NOT_PARSE:
		return fmt.Sprintf("could not parse line: %s", string(err.Content))
	}
	return "invalid read error"
}
<|endoftext|>