<commit_before>package chromedp

import (
	"context"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	"github.com/chromedp/cdproto/target"
)

var (
	execPath    string
	testdataDir string

	browserCtx context.Context

	// allocOpts is filled in TestMain
	allocOpts []ExecAllocatorOption
)

func testAllocate(tb testing.TB, path string) (_ context.Context, cancel func()) {
	// Same browser, new tab; not needing to start new chrome browsers for
	// each test gives a huge speed-up.
	ctx, _ := NewContext(browserCtx)

	// Only navigate if we want a path, otherwise leave the blank page.
	if path != "" {
		if err := Run(ctx, Navigate(testdataDir+"/"+path)); err != nil {
			tb.Fatal(err)
		}
	}

	cancelErr := func() {
		if err := Cancel(ctx); err != nil {
			tb.Error(err)
		}
	}
	return ctx, cancelErr
}

func TestMain(m *testing.M) {
	if task := os.Getenv("CHROMEDP_TEST_TASK"); task != "" {
		// The test binary is re-used to run standalone tasks, such as
		// allocating a browser within a Docker container.
		if err := runTask(task); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	wd, err := os.Getwd()
	if err != nil {
		panic(fmt.Sprintf("could not get working directory: %v", err))
	}
	testdataDir = "file://" + path.Join(wd, "testdata")

	// build on top of the default options
	allocOpts = append(allocOpts, DefaultExecAllocatorOptions...)

	// disabling the GPU helps portability with some systems like Travis,
	// and can slightly speed up the tests on other systems
	allocOpts = append(allocOpts, DisableGPU)

	// find the exec path once at startup
	// it's worth noting that newer versions of chrome (64+) run much faster
	// than older ones -- same for headless_shell ...
	execPath = os.Getenv("CHROMEDP_TEST_RUNNER")
	if execPath == "" {
		execPath = findExecPath()
	}
	allocOpts = append(allocOpts, ExecPath(execPath))

	// not explicitly needed to be set, as this vastly speeds up unit tests
	if noSandbox := os.Getenv("CHROMEDP_NO_SANDBOX"); noSandbox != "false" {
		allocOpts = append(allocOpts, NoSandbox)
	}

	allocCtx, cancel := NewExecAllocator(context.Background(), allocOpts...)

	var browserOpts []ContextOption
	if debug := os.Getenv("CHROMEDP_DEBUG"); debug != "" && debug != "false" {
		browserOpts = append(browserOpts, WithDebugf(log.Printf))
	}

	// start the browser
	browserCtx, _ = NewContext(allocCtx, browserOpts...)
	if err := Run(browserCtx); err != nil {
		panic(err)
	}

	code := m.Run()

	cancel()
	os.Exit(code)
}

func runTask(name string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	switch name {
	case "ExecAllocator_Allocate":
		ctx, cancel := NewContext(ctx)
		defer cancel()
		if err := Run(ctx); err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown test binary task: %q", name)
	}
	return nil
}

func BenchmarkTabNavigate(b *testing.B) {
	b.ReportAllocs()

	allocCtx, cancel := NewExecAllocator(context.Background(), allocOpts...)
	defer cancel()

	// start the browser
	bctx, _ := NewContext(allocCtx)
	if err := Run(bctx); err != nil {
		b.Fatal(err)
	}

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			ctx, _ := NewContext(bctx)
			if err := Run(ctx,
				Navigate(testdataDir+"/form.html"),
				WaitVisible(`#form`, ByID), // for form.html
			); err != nil {
				b.Fatal(err)
			}
			if err := Cancel(ctx); err != nil {
				b.Fatal(err)
			}
		}
	})
}

// checkPages fatals if the browser behind the chromedp context has an
// unexpected number of pages (tabs).
func checkTargets(tb testing.TB, ctx context.Context, want int) {
	tb.Helper()
	infos, err := Targets(ctx)
	if err != nil {
		tb.Fatal(err)
	}
	var pages []*target.Info
	for _, info := range infos {
		if info.Type == "page" {
			pages = append(pages, info)
		}
	}
	if got := len(pages); want != got {
		var summaries []string
		for _, info := range pages {
			summaries = append(summaries, fmt.Sprintf("%v", info))
		}
		tb.Fatalf("want %d targets, got %d:\n%s",
			want, got, strings.Join(summaries, "\n"))
	}
}

func TestTargets(t *testing.T) {
	t.Parallel()

	// Start one browser with one tab.
	ctx1, cancel1 := NewContext(context.Background())
	defer cancel1()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	checkTargets(t, ctx1, 1)

	// Start a second tab on the same browser.
	ctx2, cancel2 := NewContext(ctx1)
	defer cancel2()
	if err := Run(ctx2); err != nil {
		t.Fatal(err)
	}
	checkTargets(t, ctx2, 2)

	// The first context should also see both targets.
	checkTargets(t, ctx1, 2)

	// Cancelling the second context should close the second tab alone.
	cancel2()
	checkTargets(t, ctx1, 1)

	// We used to have a bug where Run would reset the first context as if
	// it weren't the first, breaking its cancellation.
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}
}

func TestCancelError(t *testing.T) {
	t.Parallel()

	ctx1, cancel1 := NewContext(context.Background())
	defer cancel1()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	// Open and close a target normally; no error.
	ctx2, cancel2 := NewContext(ctx1)
	defer cancel2()
	if err := Run(ctx2); err != nil {
		t.Fatal(err)
	}
	if err := Cancel(ctx2); err != nil {
		t.Fatalf("expected a nil error, got %v", err)
	}

	// Make "cancel" close the wrong target; error.
	ctx3, cancel3 := NewContext(ctx1)
	defer cancel3()
	if err := Run(ctx3); err != nil {
		t.Fatal(err)
	}
	FromContext(ctx3).Target.TargetID = "wrong"
	if err := Cancel(ctx3); err == nil {
		t.Fatalf("expected a non-nil error, got %v", err)
	}
}

func TestPrematureCancel(t *testing.T) {
	t.Parallel()

	// Cancel before the browser is allocated.
	ctx, cancel := NewContext(context.Background())
	cancel()
	if err := Run(ctx); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestPrematureCancelTab(t *testing.T) {
	t.Parallel()

	ctx1, cancel := NewContext(context.Background())
	defer cancel()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	ctx2, cancel := NewContext(ctx1)
	// Cancel after the browser is allocated, but before we've created a new
	// tab.
	cancel()
	if err := Run(ctx2); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestPrematureCancelAllocator(t *testing.T) {
	t.Parallel()

	// To ensure we don't actually fire any Chrome processes.
	allocCtx, cancel := NewExecAllocator(context.Background(),
		ExecPath("/do-not-run-chrome"))
	// Cancel before the browser is allocated.
	cancel()

	ctx, cancel := NewContext(allocCtx)
	defer cancel()
	if err := Run(ctx); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestConcurrentCancel(t *testing.T) {
	t.Parallel()

	// To ensure we don't actually fire any Chrome processes.
	allocCtx, cancel := NewExecAllocator(context.Background(),
		ExecPath("/do-not-run-chrome"))
	defer cancel()

	// 50 is enough for 'go test -race' to easily spot issues.
	for i := 0; i < 50; i++ {
		ctx, cancel := NewContext(allocCtx)
		go cancel()
		go Run(ctx)
	}
}
<commit_msg>clean up testAllocate with context.CancelFunc<commit_after>package chromedp

import (
	"context"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"testing"
	"time"

	"github.com/chromedp/cdproto/target"
)

var (
	execPath    string
	testdataDir string

	browserCtx context.Context

	// allocOpts is filled in TestMain
	allocOpts []ExecAllocatorOption
)

func testAllocate(tb testing.TB, path string) (context.Context, context.CancelFunc) {
	// Same browser, new tab; not needing to start new chrome browsers for
	// each test gives a huge speed-up.
	ctx, _ := NewContext(browserCtx)

	// Only navigate if we want a path, otherwise leave the blank page.
	if path != "" {
		if err := Run(ctx, Navigate(testdataDir+"/"+path)); err != nil {
			tb.Fatal(err)
		}
	}

	cancel := func() {
		if err := Cancel(ctx); err != nil {
			tb.Error(err)
		}
	}
	return ctx, cancel
}

func TestMain(m *testing.M) {
	if task := os.Getenv("CHROMEDP_TEST_TASK"); task != "" {
		// The test binary is re-used to run standalone tasks, such as
		// allocating a browser within a Docker container.
		if err := runTask(task); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	wd, err := os.Getwd()
	if err != nil {
		panic(fmt.Sprintf("could not get working directory: %v", err))
	}
	testdataDir = "file://" + path.Join(wd, "testdata")

	// build on top of the default options
	allocOpts = append(allocOpts, DefaultExecAllocatorOptions...)

	// disabling the GPU helps portability with some systems like Travis,
	// and can slightly speed up the tests on other systems
	allocOpts = append(allocOpts, DisableGPU)

	// find the exec path once at startup
	// it's worth noting that newer versions of chrome (64+) run much faster
	// than older ones -- same for headless_shell ...
	execPath = os.Getenv("CHROMEDP_TEST_RUNNER")
	if execPath == "" {
		execPath = findExecPath()
	}
	allocOpts = append(allocOpts, ExecPath(execPath))

	// not explicitly needed to be set, as this vastly speeds up unit tests
	if noSandbox := os.Getenv("CHROMEDP_NO_SANDBOX"); noSandbox != "false" {
		allocOpts = append(allocOpts, NoSandbox)
	}

	allocCtx, cancel := NewExecAllocator(context.Background(), allocOpts...)

	var browserOpts []ContextOption
	if debug := os.Getenv("CHROMEDP_DEBUG"); debug != "" && debug != "false" {
		browserOpts = append(browserOpts, WithDebugf(log.Printf))
	}

	// start the browser
	browserCtx, _ = NewContext(allocCtx, browserOpts...)
	if err := Run(browserCtx); err != nil {
		panic(err)
	}

	code := m.Run()

	cancel()
	os.Exit(code)
}

func runTask(name string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	switch name {
	case "ExecAllocator_Allocate":
		ctx, cancel := NewContext(ctx)
		defer cancel()
		if err := Run(ctx); err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown test binary task: %q", name)
	}
	return nil
}

func BenchmarkTabNavigate(b *testing.B) {
	b.ReportAllocs()

	allocCtx, cancel := NewExecAllocator(context.Background(), allocOpts...)
	defer cancel()

	// start the browser
	bctx, _ := NewContext(allocCtx)
	if err := Run(bctx); err != nil {
		b.Fatal(err)
	}

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			ctx, _ := NewContext(bctx)
			if err := Run(ctx,
				Navigate(testdataDir+"/form.html"),
				WaitVisible(`#form`, ByID), // for form.html
			); err != nil {
				b.Fatal(err)
			}
			if err := Cancel(ctx); err != nil {
				b.Fatal(err)
			}
		}
	})
}

// checkPages fatals if the browser behind the chromedp context has an
// unexpected number of pages (tabs).
func checkTargets(tb testing.TB, ctx context.Context, want int) {
	tb.Helper()
	infos, err := Targets(ctx)
	if err != nil {
		tb.Fatal(err)
	}
	var pages []*target.Info
	for _, info := range infos {
		if info.Type == "page" {
			pages = append(pages, info)
		}
	}
	if got := len(pages); want != got {
		var summaries []string
		for _, info := range pages {
			summaries = append(summaries, fmt.Sprintf("%v", info))
		}
		tb.Fatalf("want %d targets, got %d:\n%s",
			want, got, strings.Join(summaries, "\n"))
	}
}

func TestTargets(t *testing.T) {
	t.Parallel()

	// Start one browser with one tab.
	ctx1, cancel1 := NewContext(context.Background())
	defer cancel1()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	checkTargets(t, ctx1, 1)

	// Start a second tab on the same browser.
	ctx2, cancel2 := NewContext(ctx1)
	defer cancel2()
	if err := Run(ctx2); err != nil {
		t.Fatal(err)
	}
	checkTargets(t, ctx2, 2)

	// The first context should also see both targets.
	checkTargets(t, ctx1, 2)

	// Cancelling the second context should close the second tab alone.
	cancel2()
	checkTargets(t, ctx1, 1)

	// We used to have a bug where Run would reset the first context as if
	// it weren't the first, breaking its cancellation.
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}
}

func TestCancelError(t *testing.T) {
	t.Parallel()

	ctx1, cancel1 := NewContext(context.Background())
	defer cancel1()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	// Open and close a target normally; no error.
	ctx2, cancel2 := NewContext(ctx1)
	defer cancel2()
	if err := Run(ctx2); err != nil {
		t.Fatal(err)
	}
	if err := Cancel(ctx2); err != nil {
		t.Fatalf("expected a nil error, got %v", err)
	}

	// Make "cancel" close the wrong target; error.
	ctx3, cancel3 := NewContext(ctx1)
	defer cancel3()
	if err := Run(ctx3); err != nil {
		t.Fatal(err)
	}
	FromContext(ctx3).Target.TargetID = "wrong"
	if err := Cancel(ctx3); err == nil {
		t.Fatalf("expected a non-nil error, got %v", err)
	}
}

func TestPrematureCancel(t *testing.T) {
	t.Parallel()

	// Cancel before the browser is allocated.
	ctx, cancel := NewContext(context.Background())
	cancel()
	if err := Run(ctx); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestPrematureCancelTab(t *testing.T) {
	t.Parallel()

	ctx1, cancel := NewContext(context.Background())
	defer cancel()
	if err := Run(ctx1); err != nil {
		t.Fatal(err)
	}

	ctx2, cancel := NewContext(ctx1)
	// Cancel after the browser is allocated, but before we've created a new
	// tab.
	cancel()
	if err := Run(ctx2); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestPrematureCancelAllocator(t *testing.T) {
	t.Parallel()

	// To ensure we don't actually fire any Chrome processes.
	allocCtx, cancel := NewExecAllocator(context.Background(),
		ExecPath("/do-not-run-chrome"))
	// Cancel before the browser is allocated.
	cancel()

	ctx, cancel := NewContext(allocCtx)
	defer cancel()
	if err := Run(ctx); err != context.Canceled {
		t.Fatalf("wanted canceled context error, got %v", err)
	}
}

func TestConcurrentCancel(t *testing.T) {
	t.Parallel()

	// To ensure we don't actually fire any Chrome processes.
	allocCtx, cancel := NewExecAllocator(context.Background(),
		ExecPath("/do-not-run-chrome"))
	defer cancel()

	// 50 is enough for 'go test -race' to easily spot issues.
	for i := 0; i < 50; i++ {
		ctx, cancel := NewContext(allocCtx)
		go cancel()
		go Run(ctx)
	}
}
<|endoftext|>


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

package fs

import (
	"fmt"
	"log"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	"github.com/jacobsa/gcloud/syncutil"
)

const tmpObjectPrefix = ".gcsfuse_tmp/"

func garbageCollectOnce(
	ctx context.Context,
	bucket gcs.Bucket) (objectsDeleted uint64, err error) {
	const stalenessThreshold = 30 * time.Minute
	b := syncutil.NewBundle(ctx)

	// List all objects with the temporary prefix.
	objects := make(chan *gcs.Object, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(objects)
		err = gcsutil.ListPrefix(ctx, bucket, tmpObjectPrefix, objects)
		if err != nil {
			err = fmt.Errorf("ListPrefix: %v", err)
			return
		}

		return
	})

	// Filter to the names of objects that are stale.
	now := time.Now()
	staleNames := make(chan string, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(staleNames)
		for o := range objects {
			if now.Sub(o.Updated) < stalenessThreshold {
				continue
			}

			select {
			case <-ctx.Done():
				err = ctx.Err()
				return

			case staleNames <- o.Name:
			}
		}

		return
	})

	// Delete those objects.
	b.Add(func(ctx context.Context) (err error) {
		for name := range staleNames {
			err = bucket.DeleteObject(ctx, name)
			if err != nil {
				err = fmt.Errorf("DeleteObject(%q): %v", name, err)
				return
			}

			atomic.AddUint64(&objectsDeleted, 1)
		}

		return
	})

	err = b.Join()
	return
}

// Periodically delete stale temporary objects from the supplied bucket until
// the context is cancelled.
func garbageCollect(
	ctx context.Context,
	bucket gcs.Bucket) {
	const period = 10 * time.Minute
	ticker := time.NewTicker(period)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return

		case <-ticker.C:
		}

		log.Println("Starting a garbage collection run.")

		startTime := time.Now()
		objectsDeleted, err := garbageCollectOnce(ctx, bucket)

		if err != nil {
			log.Printf(
				"Garbage collection failed after deleting %d objects in %v, "+
					"with error: %v",
				objectsDeleted,
				time.Since(startTime),
				err)
		} else {
			log.Printf(
				"Garbage collection succeeded after deleted %d objects in %v.",
				objectsDeleted,
				time.Since(startTime))
		}
	}
}
<commit_msg>Don't hard-code the temp object prefix.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fs

import (
	"fmt"
	"log"
	"sync/atomic"
	"time"

	"golang.org/x/net/context"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	"github.com/jacobsa/gcloud/syncutil"
)

func garbageCollectOnce(
	ctx context.Context,
	tmpObjectPrefix string,
	bucket gcs.Bucket) (objectsDeleted uint64, err error) {
	const stalenessThreshold = 30 * time.Minute
	b := syncutil.NewBundle(ctx)

	// List all objects with the temporary prefix.
	objects := make(chan *gcs.Object, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(objects)
		err = gcsutil.ListPrefix(ctx, bucket, tmpObjectPrefix, objects)
		if err != nil {
			err = fmt.Errorf("ListPrefix: %v", err)
			return
		}

		return
	})

	// Filter to the names of objects that are stale.
	now := time.Now()
	staleNames := make(chan string, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(staleNames)
		for o := range objects {
			if now.Sub(o.Updated) < stalenessThreshold {
				continue
			}

			select {
			case <-ctx.Done():
				err = ctx.Err()
				return

			case staleNames <- o.Name:
			}
		}

		return
	})

	// Delete those objects.
	b.Add(func(ctx context.Context) (err error) {
		for name := range staleNames {
			err = bucket.DeleteObject(ctx, name)
			if err != nil {
				err = fmt.Errorf("DeleteObject(%q): %v", name, err)
				return
			}

			atomic.AddUint64(&objectsDeleted, 1)
		}

		return
	})

	err = b.Join()
	return
}

// Periodically delete stale temporary objects from the supplied bucket until
// the context is cancelled.
func garbageCollect(
	ctx context.Context,
	tmpObjectPrefix string,
	bucket gcs.Bucket) {
	const period = 10 * time.Minute
	ticker := time.NewTicker(period)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return

		case <-ticker.C:
		}

		log.Println("Starting a garbage collection run.")

		startTime := time.Now()
		objectsDeleted, err := garbageCollectOnce(ctx, tmpObjectPrefix, bucket)

		if err != nil {
			log.Printf(
				"Garbage collection failed after deleting %d objects in %v, "+
					"with error: %v",
				objectsDeleted,
				time.Since(startTime),
				err)
		} else {
			log.Printf(
				"Garbage collection succeeded after deleted %d objects in %v.",
				objectsDeleted,
				time.Since(startTime))
		}
	}
}
<|endoftext|>


<commit_before>package fslock_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	. "launchpad.net/gocheck"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/utils/fslock"
)

func Test(t *testing.T) {
	TestingT(t)
}

type fslockSuite struct {
	coretesting.LoggingSuite
}

var _ = Suite(&fslockSuite{})

func (s *fslockSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	fslock.SetLockWaitDelay(1 * time.Millisecond)
}

func (s *fslockSuite) TearDownSuite(c *C) {
	fslock.SetLockWaitDelay(1 * time.Second)
	s.LoggingSuite.TearDownSuite(c)
}

// This test also happens to test that locks can get created when the parent
// lock directory doesn't exist.
func (s *fslockSuite) TestValidNamesLockDir(c *C) {

	for _, name := range []string{
		"a",
		"longer",
		"longer-with.special-characters",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, IsNil)
	}
}

func (s *fslockSuite) TestInvalidNames(c *C) {

	for _, name := range []string{
		"NoCapitals",
		"no+plus",
		"no/slash",
		"no\\backslash",
		"no$dollar",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, ErrorMatches, "Invalid lock name .*")
	}
}

func (s *fslockSuite) TestNewLockWithExistingDir(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	_, err = fslock.NewLock(dir, "special")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestNewLockWithExistingFileInPlace(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	path := path.Join(dir, "locks")
	err = ioutil.WriteFile(path, []byte("foo"), 0644)
	c.Assert(err, IsNil)

	_, err = fslock.NewLock(path, "special")
	c.Assert(err, ErrorMatches, `.* not a directory`)
}

func (s *fslockSuite) TestIsLockHeldBasics(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)

	err = lock.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, true)

	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestIsLockHeldTwoLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestLockBlocks(c *C) {

	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	acquired := make(chan struct{})
	err = lock1.Lock("")
	c.Assert(err, IsNil)

	go func() {
		lock2.Lock("")
		acquired <- struct{}{}
		close(acquired)
	}()

	// Waiting for something not to happen is inherently hard...
	select {
	case <-acquired:
		c.Fatalf("Unexpected lock acquisition")
	case <-time.After(50 * time.Millisecond):
		// all good
	}

	err = lock1.Unlock()
	c.Assert(err, IsNil)

	select {
	case <-acquired:
		// all good
	case <-time.After(50 * time.Millisecond):
		c.Fatalf("Expected lock acquisition")
	}

	c.Assert(lock2.IsLockHeld(), Equals, true)
}

func (s *fslockSuite) TestLockWithTimeoutUnlocked(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestLockWithTimeoutLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, Equals, fslock.ErrTimeout)
}

func (s *fslockSuite) TestUnlock(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)
}

func (s *fslockSuite) TestIsLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	c.Assert(lock1.IsLocked(), Equals, true)
	c.Assert(lock2.IsLocked(), Equals, true)
}

func (s *fslockSuite) TestBreakLock(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.BreakLock()
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLocked(), Equals, false)

	// Normally locks are broken due to client crashes, not duration.
	err = lock1.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	// Breaking a non-existant isn't an error
	err = lock2.BreakLock()
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestMessage(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")

	err = lock.SetMessage("my message")
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	err = lock.Lock("")
	c.Assert(err, IsNil)

	err = lock.SetMessage("my message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "my message")

	// Messages can be changed while the lock is held.
	err = lock.SetMessage("new message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "new message")

	// Unlocking removes the message.
	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")
}

func (s *fslockSuite) TestMessageAcrossLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	err = lock1.SetMessage("very busy")
	c.Assert(err, IsNil)

	c.Assert(lock2.Message(), Equals, "very busy")
}

func (s *fslockSuite) TestInitialMessageWhenLocking(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Lock("initial message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial message")

	err = lock.Unlock()
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "initial timeout message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial timeout message")
}

func (s *fslockSuite) TestStress(c *C) {
	const lockAttempts = 200
	const concurrentLocks = 10

	var counter = new(int64)
	// Use atomics to update lockState to make sure the lock isn't held by
	// someone else. A value of 1 means locked, 0 means unlocked.
	var lockState = new(int32)
	var done = make(chan struct{})
	defer close(done)

	dir := c.MkDir()

	var stress = func(name string) {
		defer func() { done <- struct{}{} }()
		lock, err := fslock.NewLock(dir, "testing")
		if err != nil {
			return
		}
		for i := 0; i < lockAttempts; i++ {
			err = lock.Lock(name)
			c.Check(err, IsNil)
			state := atomic.AddInt32(lockState, 1)
			c.Check(state, Equals, int32(1))
			// Tell the go routine scheduler to give a slice to someone else
			// while we have this locked.
			runtime.Gosched()
			// need to decrement prior to unlock to avoid the race of someone
			// else grabbing the lock before we decrement the state.
			_ = atomic.AddInt32(lockState, -1)
			err = lock.Unlock()
			c.Check(err, IsNil)
			// increment the general counter
			_ = atomic.AddInt64(counter, 1)
		}
	}

	for i := 0; i < concurrentLocks; i++ {
		go stress(fmt.Sprintf("Lock %d", i))
	}
	for i := 0; i < concurrentLocks; i++ {
		<-done
	}
	c.Assert(*counter, Equals, int64(lockAttempts*concurrentLocks))
}
<commit_msg>Added a few more invalid names.<commit_after>package fslock_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

	. "launchpad.net/gocheck"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/utils/fslock"
)

func Test(t *testing.T) {
	TestingT(t)
}

type fslockSuite struct {
	coretesting.LoggingSuite
}

var _ = Suite(&fslockSuite{})

func (s *fslockSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
	fslock.SetLockWaitDelay(1 * time.Millisecond)
}

func (s *fslockSuite) TearDownSuite(c *C) {
	fslock.SetLockWaitDelay(1 * time.Second)
	s.LoggingSuite.TearDownSuite(c)
}

// This test also happens to test that locks can get created when the parent
// lock directory doesn't exist.
func (s *fslockSuite) TestValidNamesLockDir(c *C) {

	for _, name := range []string{
		"a",
		"longer",
		"longer-with.special-characters",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, IsNil)
	}
}

func (s *fslockSuite) TestInvalidNames(c *C) {

	for _, name := range []string{
		".start",
		"-start",
		"NoCapitals",
		"no+plus",
		"no/slash",
		"no\\backslash",
		"no$dollar",
		"no:colon",
	} {
		dir := c.MkDir()
		_, err := fslock.NewLock(dir, name)
		c.Assert(err, ErrorMatches, "Invalid lock name .*")
	}
}

func (s *fslockSuite) TestNewLockWithExistingDir(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	_, err = fslock.NewLock(dir, "special")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestNewLockWithExistingFileInPlace(c *C) {
	dir := c.MkDir()
	err := os.MkdirAll(dir, 0755)
	c.Assert(err, IsNil)
	path := path.Join(dir, "locks")
	err = ioutil.WriteFile(path, []byte("foo"), 0644)
	c.Assert(err, IsNil)

	_, err = fslock.NewLock(path, "special")
	c.Assert(err, ErrorMatches, `.* not a directory`)
}

func (s *fslockSuite) TestIsLockHeldBasics(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)

	err = lock.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, true)

	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestIsLockHeldTwoLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLockHeld(), Equals, false)
}

func (s *fslockSuite) TestLockBlocks(c *C) {

	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	acquired := make(chan struct{})
	err = lock1.Lock("")
	c.Assert(err, IsNil)

	go func() {
		lock2.Lock("")
		acquired <- struct{}{}
		close(acquired)
	}()

	// Waiting for something not to happen is inherently hard...
	select {
	case <-acquired:
		c.Fatalf("Unexpected lock acquisition")
	case <-time.After(50 * time.Millisecond):
		// all good
	}

	err = lock1.Unlock()
	c.Assert(err, IsNil)

	select {
	case <-acquired:
		// all good
	case <-time.After(50 * time.Millisecond):
		c.Fatalf("Expected lock acquisition")
	}

	c.Assert(lock2.IsLockHeld(), Equals, true)
}

func (s *fslockSuite) TestLockWithTimeoutUnlocked(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestLockWithTimeoutLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.LockWithTimeout(10*time.Millisecond, "")
	c.Assert(err, Equals, fslock.ErrTimeout)
}

func (s *fslockSuite) TestUnlock(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)
}

func (s *fslockSuite) TestIsLocked(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	c.Assert(lock1.IsLocked(), Equals, true)
	c.Assert(lock2.IsLocked(), Equals, true)
}

func (s *fslockSuite) TestBreakLock(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)

	err = lock2.BreakLock()
	c.Assert(err, IsNil)
	c.Assert(lock2.IsLocked(), Equals, false)

	// Normally locks are broken due to client crashes, not duration.
	err = lock1.Unlock()
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	// Breaking a non-existant isn't an error
	err = lock2.BreakLock()
	c.Assert(err, IsNil)
}

func (s *fslockSuite) TestMessage(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")

	err = lock.SetMessage("my message")
	c.Assert(err, Equals, fslock.ErrLockNotHeld)

	err = lock.Lock("")
	c.Assert(err, IsNil)

	err = lock.SetMessage("my message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "my message")

	// Messages can be changed while the lock is held.
	err = lock.SetMessage("new message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "new message")

	// Unlocking removes the message.
	err = lock.Unlock()
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "")
}

func (s *fslockSuite) TestMessageAcrossLocks(c *C) {
	dir := c.MkDir()
	lock1, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)
	lock2, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock1.Lock("")
	c.Assert(err, IsNil)
	err = lock1.SetMessage("very busy")
	c.Assert(err, IsNil)

	c.Assert(lock2.Message(), Equals, "very busy")
}

func (s *fslockSuite) TestInitialMessageWhenLocking(c *C) {
	dir := c.MkDir()
	lock, err := fslock.NewLock(dir, "testing")
	c.Assert(err, IsNil)

	err = lock.Lock("initial message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial message")

	err = lock.Unlock()
	c.Assert(err, IsNil)

	err = lock.LockWithTimeout(10*time.Millisecond, "initial timeout message")
	c.Assert(err, IsNil)
	c.Assert(lock.Message(), Equals, "initial timeout message")
}

func (s *fslockSuite) TestStress(c *C) {
	const lockAttempts = 200
	const concurrentLocks = 10

	var counter = new(int64)
	// Use atomics to update lockState to make sure the lock isn't held by
	// someone else. A value of 1 means locked, 0 means unlocked.
	var lockState = new(int32)
	var done = make(chan struct{})
	defer close(done)

	dir := c.MkDir()

	var stress = func(name string) {
		defer func() { done <- struct{}{} }()
		lock, err := fslock.NewLock(dir, "testing")
		if err != nil {
			return
		}
		for i := 0; i < lockAttempts; i++ {
			err = lock.Lock(name)
			c.Check(err, IsNil)
			state := atomic.AddInt32(lockState, 1)
			c.Check(state, Equals, int32(1))
			// Tell the go routine scheduler to give a slice to someone else
			// while we have this locked.
			runtime.Gosched()
			// need to decrement prior to unlock to avoid the race of someone
			// else grabbing the lock before we decrement the state.
			_ = atomic.AddInt32(lockState, -1)
			err = lock.Unlock()
			c.Check(err, IsNil)
			// increment the general counter
			_ = atomic.AddInt64(counter, 1)
		}
	}

	for i := 0; i < concurrentLocks; i++ {
		go stress(fmt.Sprintf("Lock %d", i))
	}
	for i := 0; i < concurrentLocks; i++ {
		<-done
	}
	c.Assert(*counter, Equals, int64(lockAttempts*concurrentLocks))
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced Authors.
// 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 cmd

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"
)

func remove(index int, list ...interface{}) []interface{} {
	var (
		left  []interface{}
		right []interface{}
	)

	switch {
	case index < 0 || index > len(list):
		panic("index out of bounds")
	case index+1 < len(list):
		right = list[index+1:]
		fallthrough
	default:
		left = list[:index]
	}

	return append(left, right...)
}

var editors = []string{"vim", "vi", "nano"}

func findEditor(editor string) (string, error) {
	if editor != "" {
		path, err := exec.LookPath(editor)
		if err != nil {
			return "", fmt.Errorf("editor (%s) not found: %s", editor, err)
		}
		return path, nil
	}
	for _, e := range editors {
		path, err := exec.LookPath(e)
		if err == nil {
			return path, nil
		}
	}
	return "", fmt.Errorf("no editor found")
}

func openEditor(data []byte, name, editor string) (reader io.Reader, err error) {
	if terminal.IsTerminal(syscall.Stdin) {
		editor, err := findEditor(editor)
		if err != nil {
			return nil, err
		}

		f, err := ioutil.TempFile("", name+"_")
		if err != nil {
			return nil, fmt.Errorf("could not open tempfile: %s", err)
		}
		defer os.Remove(f.Name())
		defer f.Close()

		if _, err := f.Write(data); err != nil {
			return nil, fmt.Errorf("could not write tempfile: %s", err)
		}

		e := exec.Command(editor, f.Name())
		e.Stdin = os.Stdin
		e.Stdout = os.Stdout
		e.Stderr = os.Stderr

		if err := e.Run(); err != nil {
			return nil, fmt.Errorf("received error from editor: %s (%s)", err, editor)
		}
		if _, err := f.Seek(0, 0); err != nil {
			return nil, fmt.Errorf("could not seek file: %s", err)
		}

		data, err := ioutil.ReadFile(f.Name())
		if err != nil {
			return nil, fmt.Errorf("could not read file: %s", err)
		}

		reader = bytes.NewReader(data)
	} else {
		if _, err := os.Stdout.Write(data); err != nil {
			return nil, fmt.Errorf("could not write to stdout: %s", err)
		}
		reader = os.Stdin
	}

	return reader, nil
}
<commit_msg>fixes CC-907<commit_after>// Copyright 2014 The Serviced Authors.
// 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 cmd

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"
)

func remove(index int, list ...interface{}) []interface{} {
	var (
		left  []interface{}
		right []interface{}
	)

	switch {
	case index < 0 || index > len(list):
		panic("index out of bounds")
	case index+1 < len(list):
		right = list[index+1:]
		fallthrough
	default:
		left = list[:index]
	}

	return append(left, right...)
}

var editors = []string{"vim", "vi", "nano"}

func findEditor(editor string) (string, error) {
	if editor != "" {
		path, err := exec.LookPath(editor)
		if err != nil {
			return "", fmt.Errorf("editor (%s) not found: %s", editor, err)
		}
		return path, nil
	}
	for _, e := range editors {
		path, err := exec.LookPath(e)
		if err == nil {
			return path, nil
		}
	}
	return "", fmt.Errorf("no editor found")
}

func openEditor(data []byte, name, editor string) (reader io.Reader, err error) {
	if terminal.IsTerminal(syscall.Stdin) {
		ed := strings.Split(editor, " ")
		editor, err := findEditor(ed[0])
		if err != nil {
			return nil, err
		}

		f, err := ioutil.TempFile("", name+"_")
		if err != nil {
			return nil, fmt.Errorf("could not open tempfile: %s", err)
		}
		defer os.Remove(f.Name())
		defer f.Close()

		if _, err := f.Write(data); err != nil {
			return nil, fmt.Errorf("could not write tempfile: %s", err)
		}

		ed = append(ed, f.Name())
		e := exec.Command(editor, ed[1:]...)
		e.Stdin = os.Stdin
		e.Stdout = os.Stdout
		e.Stderr = os.Stderr

		if err := e.Run(); err != nil {
			return nil, fmt.Errorf("received error from editor: %s (%s)", err, editor)
		}
		if _, err := f.Seek(0, 0); err != nil {
			return nil, fmt.Errorf("could not seek file: %s", err)
		}

		data, err := ioutil.ReadFile(f.Name())
		if err != nil {
			return nil, fmt.Errorf("could not read file: %s", err)
		}

		reader = bytes.NewReader(data)
	} else {
		if _, err := os.Stdout.Write(data); err != nil {
			return nil, fmt.Errorf("could not write to stdout: %s", err)
		}
		reader = os.Stdin
	}

	return reader, nil
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strconv"
)

const (
	apiURL          = "http://api.canlii.org/v1/"
	apiKeysFilename = "apiKey"
)

var (
	apiKeys     []APIKey
	keyRotation int
)

func init() {
	log.Printf("Loading API keys from file `%s`", apiKeysFilename)
	var err error
	apiKeys, err = LoadAPIKeysFromFile(apiKeysFilename)
	if err != nil {
		panic(err)
	}
	log.Printf("Done, found %d keys\n", len(apiKeys))
	for _, key := range apiKeys {
		log.Printf("key `%s`, %d perDay, %d perSec",
			key.Key, key.CallPerDay, key.CallPerSecond)
	}
}

func getAPIKey() url.Values {
	if keyRotation >= len(apiKeys)-1 {
		keyRotation = 0
	}
	if !apiKeys[keyRotation].HasLeft() {
		keyRotation++
	}

	key, err := apiKeys[keyRotation].Use()
	if err != nil {
		return nil
	}

	val := url.Values{}

	val.Set("api_key", key)
	return val
}

func doQuery(qURL *url.URL) ([]byte, error) {
	resp, err := http.Get(qURL.String())
	if err != nil {
		return nil, fmt.Errorf("http.Get(url), %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("query `%v` returned status code %d, %s", qURL.String(), resp.StatusCode, resp.Status)
	}

	dump, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("reading response, %v", err)
	}
	return dump, nil
}

type Database struct {
	ID           string `json: "databaseId"`
	Jurisdiction string `json: "jurisdiction"`
	Name         string `json: "name"`
}

func DatabaseList() ([]Database, error) {
	collection := "caseBrowse/en/?"

	qURL, err := url.Parse(apiURL + collection + getAPIKey().Encode())
	if err != nil {
		return nil, fmt.Errorf("parsing collection url, %v", err)
	}

	dump, err := doQuery(qURL)
	if err != nil {
		return nil, fmt.Errorf("failed querying database list, %v", err)
	}

	dbL := struct {
		DbList []Database `json:"caseDatabases"`
	}{}

	if err := json.Unmarshal(dump, &dbL); err != nil {
		return nil, fmt.Errorf("unmarshalling response, %v, query was `%s`, got `%v`", err, qURL.String(), string(dump))
	}

	return dbL.DbList, nil
}

func (d *Database) CaseList(offset, count int) ([]Case, error) {
	collection := "caseBrowse/en/" + d.ID + "/?"
	val := getAPIKey()
	val.Add("offset", strconv.Itoa(offset))
	val.Add("resultCount", strconv.Itoa(count))

	qURL, err := url.Parse(apiURL + collection + val.Encode())
	if err != nil {
		return nil, fmt.Errorf("parsing query url, %v", err)
	}

	dump, err := doQuery(qURL)
	if err != nil {
		return nil, fmt.Errorf("failed querying case list, %v", err)
	}

	caseL := struct {
		Cases []Case `json:"cases"`
	}{nil}

	if err := json.Unmarshal(dump, &caseL); err != nil {
		return nil, fmt.Errorf("unmarshalling cases, %v, data = %v", err, string(dump))
	}

	return caseL.Cases, nil
}

type Case struct {
	DbID     string `json:"databaseId"`
	ID       string `json:"caseId["en"]"`
	Title    string `json:"title"`
	Citation string `json:"citation"`
}
<commit_msg>Fix obvious bug<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strconv"
)

const (
	apiURL          = "http://api.canlii.org/v1/"
	apiKeysFilename = "apiKey"
)

var (
	apiKeys     []APIKey
	keyRotation int
)

func init() {
	log.Printf("Loading API keys from file `%s`", apiKeysFilename)
	var err error
	apiKeys, err = LoadAPIKeysFromFile(apiKeysFilename)
	if err != nil {
		panic(err)
	}
	log.Printf("Done, found %d keys\n", len(apiKeys))
	for _, key := range apiKeys {
		log.Printf("key `%s`, %d perDay, %d perSec",
			key.Key, key.CallPerDay, key.CallPerSecond)
	}
}

func getAPIKey() url.Values {

	if !apiKeys[keyRotation].HasLeft() {
		keyRotation++
	}

	if keyRotation >= len(apiKeys)-1 {
		keyRotation = 0
	}

	key, err := apiKeys[keyRotation].Use()
	if err != nil {
		return nil
	}

	val := url.Values{}

	val.Set("api_key", key)
	return val
}

func doQuery(qURL *url.URL) ([]byte, error) {
	resp, err := http.Get(qURL.String())
	if err != nil {
		return nil, fmt.Errorf("http.Get(url), %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("query `%v` returned status code %d, %s", qURL.String(), resp.StatusCode, resp.Status)
	}

	dump, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("reading response, %v", err)
	}
	return dump, nil
}

type Database struct {
	ID           string `json: "databaseId"`
	Jurisdiction string `json: "jurisdiction"`
	Name         string `json: "name"`
}

func DatabaseList() ([]Database, error) {
	collection := "caseBrowse/en/?"

	qURL, err := url.Parse(apiURL + collection + getAPIKey().Encode())
	if err != nil {
		return nil, fmt.Errorf("parsing collection url, %v", err)
	}

	dump, err := doQuery(qURL)
	if err != nil {
		return nil, fmt.Errorf("failed querying database list, %v", err)
	}

	dbL := struct {
		DbList []Database `json:"caseDatabases"`
	}{}

	if err := json.Unmarshal(dump, &dbL); err != nil {
		return nil, fmt.Errorf("unmarshalling response, %v, query was `%s`, got `%v`", err, qURL.String(), string(dump))
	}

	return dbL.DbList, nil
}

func (d *Database) CaseList(offset, count int) ([]Case, error) {
	collection := "caseBrowse/en/" + d.ID + "/?"
	val := getAPIKey()
	val.Add("offset", strconv.Itoa(offset))
	val.Add("resultCount", strconv.Itoa(count))

	qURL, err := url.Parse(apiURL + collection + val.Encode())
	if err != nil {
		return nil, fmt.Errorf("parsing query url, %v", err)
	}

	dump, err := doQuery(qURL)
	if err != nil {
		return nil, fmt.Errorf("failed querying case list, %v", err)
	}

	caseL := struct {
		Cases []Case `json:"cases"`
	}{nil}

	if err := json.Unmarshal(dump, &caseL); err != nil {
		return nil, fmt.Errorf("unmarshalling cases, %v, data = %v", err, string(dump))
	}

	return caseL.Cases, nil
}

type Case struct {
	DbID     string `json:"databaseId"`
	ID       string `json:"caseId["en"]"`
	Title    string `json:"title"`
	Citation string `json:"citation"`
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"

	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/go.crypto/ssh/agent"

	"github.com/coreos/shortbread/api"
	"github.com/coreos/shortbread/util"
)

const (
	SHORTBREAD_PUBLIC_KEY = "SHORTBREAD_PUBLIC_KEY"
)

func main() {
	svc, err := util.GetHTTPClientService()
	if err != nil {
		log.Printf("call to util.GetHTTPClientService failed: %s\n" , err.Error())
		return
	}

	crtSvc := api.NewCertService(svc)
	publicKeyPath := util.GetenvWithDefault(SHORTBREAD_PUBLIC_KEY, os.ExpandEnv("$HOME/.ssh/id_rsa.pub"))
	privateKeyPath := strings.Split(publicKeyPath, ".pub")[0]
	pk := util.LoadPublicKey(publicKeyPath)

	time.Sleep(2000 * time.Millisecond)
	certsWithKey, err := crtSvc.GetCerts(pk).Do()
	if err != nil {
		log.Printf("Get request to API failed: %s\n", err.Error())
		return
	}
	err = updateSSHAgent(certsWithKey.List, privateKeyPath)
	if err != nil {
		log.Println("Failed to updateSSHAgent: %s\n", err.Error())
	}
}

// updateSSHAgent takes the list of certificates and path to the private key (corresponding to the signed public key). Adds the cert if it's not present in the agent.
func updateSSHAgent(certsWithKeyList []*api.CertificateAndPrivateKey, privateKeyPath string) error {
	conn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		return err
	}
	defer conn.Close()

	sshAgent := agent.NewClient(conn)
	certsInSSHAgent, err := sshAgent.List()
	if err != nil {
		return err
	}

	privateKeyBytes, err := ioutil.ReadFile(privateKeyPath)
	if err != nil {
		return err
	}

	privateKeyInterface, err := ssh.ParseRawPrivateKey(privateKeyBytes)
	if err != nil {
		return err
	}

	// TODO optimize this. currently O(N^2)
	for _, certAndKey := range certsWithKeyList {
		cert, err := util.ParseSSHCert([]byte(certAndKey.Cert))
		for _, key := range certsInSSHAgent {
			certBlob := key.Blob
			if err != nil {
				return err
			}
			if bytes.Equal(certBlob, cert.Marshal()) {
				break
			}
		}
		err = sshAgent.Add(privateKeyInterface, cert, "certificated added by shortbread")
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>removing time.Sleep()<commit_after>package main

import (
	"bytes"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"

	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/go.crypto/ssh/agent"

	"github.com/coreos/shortbread/api"
	"github.com/coreos/shortbread/util"
)

const (
	SHORTBREAD_PUBLIC_KEY = "SHORTBREAD_PUBLIC_KEY"
)

func main() {
	svc, err := util.GetHTTPClientService()
	if err != nil {
		log.Printf("call to util.GetHTTPClientService failed: %s\n", err.Error())
		return
	}

	crtSvc := api.NewCertService(svc)
	publicKeyPath := util.GetenvWithDefault(SHORTBREAD_PUBLIC_KEY, os.ExpandEnv("$HOME/.ssh/id_rsa.pub"))
	privateKeyPath := strings.Split(publicKeyPath, ".pub")[0]
	pk := util.LoadPublicKey(publicKeyPath)

	certsWithKey, err := crtSvc.GetCerts(pk).Do()
	if err != nil {
		log.Printf("Get request to API failed: %s\n", err.Error())
		return
	}
	err = updateSSHAgent(certsWithKey.List, privateKeyPath)
	if err != nil {
		log.Printf("Failed to updateSSHAgent: %s\n", err.Error())
	}
}

// updateSSHAgent takes the list of certificates and path to the private key (corresponding to the signed public key). Adds the cert if it's not present in the agent.
func updateSSHAgent(certsWithKeyList []*api.CertificateAndPrivateKey, privateKeyPath string) error {
	conn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		return err
	}
	defer conn.Close()

	sshAgent := agent.NewClient(conn)
	certsInSSHAgent, err := sshAgent.List()
	if err != nil {
		return err
	}

	privateKeyBytes, err := ioutil.ReadFile(privateKeyPath)
	if err != nil {
		return err
	}

	privateKeyInterface, err := ssh.ParseRawPrivateKey(privateKeyBytes)
	if err != nil {
		return err
	}

	// TODO optimize this. currently O(N^2)
	for _, certAndKey := range certsWithKeyList {
		cert, err := util.ParseSSHCert([]byte(certAndKey.Cert))
		for _, key := range certsInSSHAgent {
			certBlob := key.Blob
			if err != nil {
				return err
			}
			if bytes.Equal(certBlob, cert.Marshal()) {
				break
			}
		}
		err = sshAgent.Add(privateKeyInterface, cert, "certificated added by shortbread")
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
	Copyright (c) 2020 Docker Inc.

	Permission is hereby granted, free of charge, to any person
	obtaining a copy of this software and associated documentation
	files (the "Software"), to deal in the Software without
	restriction, including without limitation the rights to use, copy,
	modify, merge, publish, distribute, sublicense, and/or sell copies
	of the Software, and to permit persons to whom the Software is
	furnished to do so, subject to the following conditions:

	The above copyright notice and this permission notice shall be
	included in all copies or substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
	EXPRESS OR IMPLIED,
	INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
	IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
	HOLDERS BE LIABLE FOR ANY CLAIM,
	DAMAGES OR OTHER LIABILITY,
	WHETHER IN AN ACTION OF CONTRACT,
	TORT OR OTHERWISE,
	ARISING FROM, OUT OF OR IN CONNECTION WITH
	THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package client

import (
	"context"
	"errors"

	"github.com/docker/api/backend"
	backendv1 "github.com/docker/api/backend/v1"
	cliv1 "github.com/docker/api/cli/v1"
	composev1 "github.com/docker/api/compose/v1"
	"github.com/docker/api/containers"
	containersv1 "github.com/docker/api/containers/v1"
	apicontext "github.com/docker/api/context"
	"github.com/docker/api/context/store"
)

// New returns a GRPC client
func New(ctx context.Context) (*Client, error) {
	currentContext := apicontext.CurrentContext(ctx)
	s := store.ContextStore(ctx)

	cc, err := s.Get(currentContext, nil)
	if err != nil {
		return nil, err
	}
	contextType := s.GetType(cc)

	b, err := backend.Get(ctx, contextType)
	if err != nil {
		return nil, err
	}

	if ba, ok := b.(containers.ContainerService); ok {
		return &Client{
			backendType: contextType,
			cc:          ba,
		}, nil
	}
	return nil, errors.New("backend not found")
}

type Client struct {
	backendv1.BackendClient
	cliv1.CliClient
	containersv1.ContainersClient
	composev1.ComposeClient

	backendType string
	cc          containers.ContainerService
}

func (c *Client) ContainerService() containers.ContainerService {
	return c.cc
}
<commit_msg>Invert the logic to preserve the happy path<commit_after>/*
	Copyright (c) 2020 Docker Inc.

	Permission is hereby granted, free of charge, to any person
	obtaining a copy of this software and associated documentation
	files (the "Software"), to deal in the Software without
	restriction, including without limitation the rights to use, copy,
	modify, merge, publish, distribute, sublicense, and/or sell copies
	of the Software, and to permit persons to whom the Software is
	furnished to do so, subject to the following conditions:

	The above copyright notice and this permission notice shall be
	included in all copies or substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
	EXPRESS OR IMPLIED,
	INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
	IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
	HOLDERS BE LIABLE FOR ANY CLAIM,
	DAMAGES OR OTHER LIABILITY,
	WHETHER IN AN ACTION OF CONTRACT,
	TORT OR OTHERWISE,
	ARISING FROM, OUT OF OR IN CONNECTION WITH
	THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/

package client

import (
	"context"
	"errors"

	"github.com/docker/api/backend"
	backendv1 "github.com/docker/api/backend/v1"
	cliv1 "github.com/docker/api/cli/v1"
	composev1 "github.com/docker/api/compose/v1"
	"github.com/docker/api/containers"
	containersv1 "github.com/docker/api/containers/v1"
	apicontext "github.com/docker/api/context"
	"github.com/docker/api/context/store"
)

// New returns a GRPC client
func New(ctx context.Context) (*Client, error) {
	currentContext := apicontext.CurrentContext(ctx)
	s := store.ContextStore(ctx)

	cc, err := s.Get(currentContext, nil)
	if err != nil {
		return nil, err
	}
	contextType := s.GetType(cc)

	b, err := backend.Get(ctx, contextType)
	if err != nil {
		return nil, err
	}

	ba, ok := b.(containers.ContainerService)
	if !ok {
		return nil, errors.New("backend not found")
	}
	return &Client{
		backendType: contextType,
		cc:          ba,
	}, nil

}

type Client struct {
	backendv1.BackendClient
	cliv1.CliClient
	containersv1.ContainersClient
	composev1.ComposeClient

	backendType string
	cc          containers.ContainerService
}

func (c *Client) ContainerService() containers.ContainerService {
	return c.cc
}
<|endoftext|>


<commit_before>// // (c) 2019-2022, Ava Labs, Inc. All rights reserved.
// // See the file LICENSE for licensing terms.

package client

import (
	"context"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/formatting"
	"github.com/ava-labs/timestampvm/timestampvm"
)

// Client defines timestampvm client operations.
type Client interface {
	// ProposeBlock submits data for a block
	ProposeBlock(ctx context.Context, data [timestampvm.DataLen]byte) (bool, error)

	// GetBlock fetches the contents of a block
	GetBlock(ctx context.Context, blockID *ids.ID) (uint64, [timestampvm.DataLen]byte, uint64, ids.ID, ids.ID, error)
}

// New creates a new client object.
func New(uri string) Client {
	req := NewEndpointRequester(uri, "timestampvm")
	return &client{req: req}
}

type client struct {
	req *EndpointRequester
}

func (cli *client) ProposeBlock(ctx context.Context, data [timestampvm.DataLen]byte) (bool, error) {
	bytes, err := formatting.Encode(formatting.Hex, data[:])
	if err != nil {
		return false, err
	}

	resp := new(timestampvm.ProposeBlockReply)
	err = cli.req.SendRequest(ctx,
		"proposeBlock",
		&timestampvm.ProposeBlockArgs{Data: bytes},
		resp,
	)
	if err != nil {
		return false, err
	}
	return resp.Success, nil
}

func (cli *client) GetBlock(ctx context.Context, blockID *ids.ID) (uint64, [timestampvm.DataLen]byte, uint64, ids.ID, ids.ID, error) {
	resp := new(timestampvm.GetBlockReply)
	err := cli.req.SendRequest(ctx,
		"getBlock",
		&timestampvm.GetBlockArgs{ID: blockID},
		resp,
	)
	if err != nil {
		return 0, [timestampvm.DataLen]byte{}, 0, ids.Empty, ids.Empty, err
	}
	bytes, err := formatting.Decode(formatting.Hex, resp.Data)
	if err != nil {
		return 0, [timestampvm.DataLen]byte{}, 0, ids.Empty, ids.Empty, err
	}
	return uint64(resp.Timestamp), timestampvm.BytesToData(bytes), uint64(resp.Height), resp.ID, resp.ParentID, nil
}
<commit_msg>update license<commit_after>// (c) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package client

import (
	"context"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/utils/formatting"
	"github.com/ava-labs/timestampvm/timestampvm"
)

// Client defines timestampvm client operations.
type Client interface {
	// ProposeBlock submits data for a block
	ProposeBlock(ctx context.Context, data [timestampvm.DataLen]byte) (bool, error)

	// GetBlock fetches the contents of a block
	GetBlock(ctx context.Context, blockID *ids.ID) (uint64, [timestampvm.DataLen]byte, uint64, ids.ID, ids.ID, error)
}

// New creates a new client object.
func New(uri string) Client {
	req := NewEndpointRequester(uri, "timestampvm")
	return &client{req: req}
}

type client struct {
	req *EndpointRequester
}

func (cli *client) ProposeBlock(ctx context.Context, data [timestampvm.DataLen]byte) (bool, error) {
	bytes, err := formatting.Encode(formatting.Hex, data[:])
	if err != nil {
		return false, err
	}

	resp := new(timestampvm.ProposeBlockReply)
	err = cli.req.SendRequest(ctx,
		"proposeBlock",
		&timestampvm.ProposeBlockArgs{Data: bytes},
		resp,
	)
	if err != nil {
		return false, err
	}
	return resp.Success, nil
}

func (cli *client) GetBlock(ctx context.Context, blockID *ids.ID) (uint64, [timestampvm.DataLen]byte, uint64, ids.ID, ids.ID, error) {
	resp := new(timestampvm.GetBlockReply)
	err := cli.req.SendRequest(ctx,
		"getBlock",
		&timestampvm.GetBlockArgs{ID: blockID},
		resp,
	)
	if err != nil {
		return 0, [timestampvm.DataLen]byte{}, 0, ids.Empty, ids.Empty, err
	}
	bytes, err := formatting.Decode(formatting.Hex, resp.Data)
	if err != nil {
		return 0, [timestampvm.DataLen]byte{}, 0, ids.Empty, ids.Empty, err
	}
	return uint64(resp.Timestamp), timestampvm.BytesToData(bytes), uint64(resp.Height), resp.ID, resp.ParentID, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
	"github.com/gin-gonic/gin"
)

type BlogService struct {
}

func (s *BlogService) Run() {
	sqlConnection = Cfg["DB_USER"] + ":" + Cfg["DB_PASSWORD"] + "@tcp(" + Cfg["DB_HOST"] + ":" + Cfg["DB_PORT"] + ")/" + Cfg["DB_NAME"] + "?parseTime=true"

	db, err := gorm.Open("mysql", sqlConnection)
	PanicIf(err)

	//db.LogMode(true)
	db.AutoMigrate(&Post{}, &User{})

	postResource := postResource{db: db}
	userResource := userResource{db: db}

	r := gin.Default()

	apiv1 := r.Group("/api/v1")
	{
		apiv1.GET("/posts", postResource.GetPosts)
		apiv1.GET("/post/:id", postResource.GetPost)
		apiv1.POST("/post", postResource.CreatePost)
		apiv1.PUT("/post/:id", postResource.UpdatePost)
		apiv1.DELETE("/post/:id", postResource.DeletePost)

		apiv1.POST("/user/login", userResource.Login)
		apiv1.GET("/user/:login", userResource.GetUser)
		apiv1.POST("/user", userResource.CreateUser)
		apiv1.PUT("/user/:login", userResource.UpdateUser)
	}

	r.StaticFile("/", "./public/index.html")
	r.Static("/public/", "./public/")

	r.Run(":" + Cfg["PORT"])
	fmt.Println("Testing")
}
<commit_msg>delete fmt<commit_after>package main

import (
	_ "github.com/go-sql-driver/mysql"
	"github.com/jinzhu/gorm"
	"github.com/gin-gonic/gin"
)

type BlogService struct {
}

func (s *BlogService) Run() {
	sqlConnection = Cfg["DB_USER"] + ":" + Cfg["DB_PASSWORD"] + "@tcp(" + Cfg["DB_HOST"] + ":" + Cfg["DB_PORT"] + ")/" + Cfg["DB_NAME"] + "?parseTime=true"

	db, err := gorm.Open("mysql", sqlConnection)
	PanicIf(err)

	//db.LogMode(true)
	db.AutoMigrate(&Post{}, &User{})

	postResource := postResource{db: db}
	userResource := userResource{db: db}

	r := gin.Default()

	apiv1 := r.Group("/api/v1")
	{
		apiv1.GET("/posts", postResource.GetPosts)
		apiv1.GET("/post/:id", postResource.GetPost)
		apiv1.POST("/post", postResource.CreatePost)
		apiv1.PUT("/post/:id", postResource.UpdatePost)
		apiv1.DELETE("/post/:id", postResource.DeletePost)

		apiv1.POST("/user/login", userResource.Login)
		apiv1.GET("/user/:login", userResource.GetUser)
		apiv1.POST("/user", userResource.CreateUser)
		apiv1.PUT("/user/:login", userResource.UpdateUser)
	}

	r.StaticFile("/", "./public/index.html")
	r.Static("/public/", "./public/")

	r.Run(":" + Cfg["PORT"])
}
<|endoftext|>


<commit_before>package kocha

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/base64"
	"errors"
	"fmt"
	"io"

	"github.com/ugorji/go/codec"
)

// SessionStore is the interface that session store.
type SessionStore interface {
	Save(sess Session) (key string, err error)
	Load(key string) (sess Session, err error)
}

// Session represents a session data store.
type Session map[string]string

// Get gets a value associated with the given key.
// If there is the no value associated with the given key, Get returns "".
func (sess Session) Get(key string) string {
	return sess[key]
}

// Set sets the value associated with the key.
// If replaces the existing value associated with the key.
func (sess Session) Set(key, value string) {
	sess[key] = value
}

// Del deletes the value associated with the key.
func (sess Session) Del(key string) {
	delete(sess, key)
}

// Clear clear the all session data.
func (sess Session) Clear() {
	for k, _ := range sess {
		delete(sess, k)
	}
}

type ErrSession struct {
	msg string
}

func (e ErrSession) Error() string {
	return e.msg
}

func NewErrSession(msg string) error {
	return ErrSession{
		msg: msg,
	}
}

// Implementation of cookie store.
//
// This session store will be a session save to client-side cookie.
// Session cookie for save is encoded, encrypted and signed.
type SessionCookieStore struct {
	// key for the encryption.
	SecretKey string

	// Key for the cookie singing.
	SigningKey string
}

var codecHandler = &codec.MsgpackHandle{}

// Save saves and returns the key of session cookie.
// Actually, key is session cookie data itself.
func (store *SessionCookieStore) Save(sess Session) (key string, err error) {
	buf := bufPool.Get().(*bytes.Buffer)
	defer func() {
		buf.Reset()
		bufPool.Put(buf)
	}()
	if err := codec.NewEncoder(buf, codecHandler).Encode(sess); err != nil {
		return "", err
	}
	encrypted, err := store.encrypt(buf.Bytes())
	if err != nil {
		return "", err
	}
	return store.encode(store.sign(encrypted)), nil
}

// Load returns the session data that extract from cookie value.
// The key is stored session cookie value.
func (store *SessionCookieStore) Load(key string) (sess Session, err error) {
	decoded, err := store.decode(key)
	if err != nil {
		return nil, err
	}
	unsigned, err := store.verify(decoded)
	if err != nil {
		return nil, err
	}
	decrypted, err := store.decrypt(unsigned)
	if err != nil {
		return nil, err
	}
	if err := codec.NewDecoderBytes(decrypted, codecHandler).Decode(&sess); err != nil {
		return nil, err
	}
	return sess, nil
}

// Validate validates SecretKey size.
func (store *SessionCookieStore) Validate() error {
	switch len(store.SecretKey) {
	case 16, 24, 32:
		return nil
	}
	return fmt.Errorf("kocha: session: %T.SecretKey size must be 16, 24 or 32, but %v", *store, len(store.SecretKey))
}

// encrypt returns encrypted data by AES-256-CBC.
func (store *SessionCookieStore) encrypt(buf []byte) ([]byte, error) {
	block, err := aes.NewCipher([]byte(store.SecretKey))
	if err != nil {
		return nil, err
	}
	// padding for CBC
	rem := (aes.BlockSize - len(buf)%aes.BlockSize) % aes.BlockSize
	for i := 0; i < rem; i++ {
		buf = append(buf, byte(rem))
	}
	encrypted := make([]byte, aes.BlockSize+len(buf))
	iv := encrypted[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(encrypted[aes.BlockSize:], buf)
	return encrypted, nil
}

// decrypt returns decrypted data from crypted data by AES-256-CBC.
func (store *SessionCookieStore) decrypt(buf []byte) ([]byte, error) {
	block, err := aes.NewCipher([]byte(store.SecretKey))
	if err != nil {
		return nil, err
	}
	iv := buf[:aes.BlockSize]
	decrypted := buf[aes.BlockSize:]
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(decrypted, decrypted)
	return decrypted, nil
}

// encode returns encoded string by Base64 with URLEncoding.
// However, encoded string will stripped the padding character of Base64.
func (store *SessionCookieStore) encode(src []byte) string {
	buf := make([]byte, base64.URLEncoding.EncodedLen(len(src)))
	base64.URLEncoding.Encode(buf, src)
	for {
		if buf[len(buf)-1] != '=' {
			break
		}
		buf = buf[:len(buf)-1]
	}
	return string(buf)
}

// decode returns decoded data from encoded data by Base64 with URLEncoding.
func (store *SessionCookieStore) decode(src string) ([]byte, error) {
	size := len(src)
	rem := (4 - size%4) % 4
	buf := make([]byte, size+rem)
	copy(buf, src)
	for i := 0; i < rem; i++ {
		buf[size+i] = '='
	}
	n, err := base64.URLEncoding.Decode(buf, buf)
	if err != nil {
		return nil, err
	}
	return buf[:n], nil
}

// sign returns signed data.
func (store *SessionCookieStore) sign(src []byte) []byte {
	sign := store.hash(src)
	return append(sign, src...)
}

// verify verify signed data and returns unsigned data if valid.
func (store *SessionCookieStore) verify(src []byte) (unsigned []byte, err error) {
	if len(src) <= sha512.Size256 {
		return nil, errors.New("kocha: session cookie value too short")
	}
	sign := src[:sha512.Size256]
	unsigned = src[sha512.Size256:]
	if !hmac.Equal(store.hash(unsigned), sign) {
		return nil, errors.New("kocha: session cookie verification failed")
	}
	return unsigned, nil
}

// hash returns hashed data by HMAC-SHA512/256.
func (store *SessionCookieStore) hash(src []byte) []byte {
	hash := hmac.New(sha512.New512_256, []byte(store.SigningKey))
	hash.Write(src)
	return hash.Sum(nil)
}
<commit_msg>session: Modify to use AES-256-GCM for encryption<commit_after>package kocha

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/base64"
	"errors"
	"fmt"
	"io"

	"github.com/ugorji/go/codec"
)

// SessionStore is the interface that session store.
type SessionStore interface {
	Save(sess Session) (key string, err error)
	Load(key string) (sess Session, err error)
}

// Session represents a session data store.
type Session map[string]string

// Get gets a value associated with the given key.
// If there is the no value associated with the given key, Get returns "".
func (sess Session) Get(key string) string {
	return sess[key]
}

// Set sets the value associated with the key.
// If replaces the existing value associated with the key.
func (sess Session) Set(key, value string) {
	sess[key] = value
}

// Del deletes the value associated with the key.
func (sess Session) Del(key string) {
	delete(sess, key)
}

// Clear clear the all session data.
func (sess Session) Clear() {
	for k, _ := range sess {
		delete(sess, k)
	}
}

type ErrSession struct {
	msg string
}

func (e ErrSession) Error() string {
	return e.msg
}

func NewErrSession(msg string) error {
	return ErrSession{
		msg: msg,
	}
}

// Implementation of cookie store.
//
// This session store will be a session save to client-side cookie.
// Session cookie for save is encoded, encrypted and signed.
type SessionCookieStore struct {
	// key for the encryption.
	SecretKey string

	// Key for the cookie singing.
	SigningKey string
}

var codecHandler = &codec.MsgpackHandle{}

// Save saves and returns the key of session cookie.
// Actually, key is session cookie data itself.
func (store *SessionCookieStore) Save(sess Session) (key string, err error) {
	buf := bufPool.Get().(*bytes.Buffer)
	defer func() {
		buf.Reset()
		bufPool.Put(buf)
	}()
	if err := codec.NewEncoder(buf, codecHandler).Encode(sess); err != nil {
		return "", err
	}
	encrypted, err := store.encrypt(buf.Bytes())
	if err != nil {
		return "", err
	}
	return store.encode(store.sign(encrypted)), nil
}

// Load returns the session data that extract from cookie value.
// The key is stored session cookie value.
func (store *SessionCookieStore) Load(key string) (sess Session, err error) {
	decoded, err := store.decode(key)
	if err != nil {
		return nil, err
	}
	unsigned, err := store.verify(decoded)
	if err != nil {
		return nil, err
	}
	decrypted, err := store.decrypt(unsigned)
	if err != nil {
		return nil, err
	}
	if err := codec.NewDecoderBytes(decrypted, codecHandler).Decode(&sess); err != nil {
		return nil, err
	}
	return sess, nil
}

// Validate validates SecretKey size.
func (store *SessionCookieStore) Validate() error {
	switch len(store.SecretKey) {
	case 16, 24, 32:
		return nil
	}
	return fmt.Errorf("kocha: session: %T.SecretKey size must be 16, 24 or 32, but %v", *store, len(store.SecretKey))
}

// encrypt returns encrypted data by AES-256-CBC.
func (store *SessionCookieStore) encrypt(buf []byte) ([]byte, error) {
	block, err := aes.NewCipher([]byte(store.SecretKey))
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	iv := make([]byte, aead.NonceSize(), len(buf)+aead.NonceSize())
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}
	encrypted := aead.Seal(nil, iv, buf, nil)
	return append(iv, encrypted...), nil
}

// decrypt returns decrypted data from crypted data by AES-256-CBC.
func (store *SessionCookieStore) decrypt(buf []byte) ([]byte, error) {
	block, err := aes.NewCipher([]byte(store.SecretKey))
	if err != nil {
		return nil, err
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	iv := buf[:aead.NonceSize()]
	decrypted := buf[aead.NonceSize():]
	if _, err := aead.Open(decrypted[:0], iv, decrypted, nil); err != nil {
		return nil, err
	}
	return decrypted, nil
}

// encode returns encoded string by Base64 with URLEncoding.
// However, encoded string will stripped the padding character of Base64.
func (store *SessionCookieStore) encode(src []byte) string {
	buf := make([]byte, base64.URLEncoding.EncodedLen(len(src)))
	base64.URLEncoding.Encode(buf, src)
	for {
		if buf[len(buf)-1] != '=' {
			break
		}
		buf = buf[:len(buf)-1]
	}
	return string(buf)
}

// decode returns decoded data from encoded data by Base64 with URLEncoding.
func (store *SessionCookieStore) decode(src string) ([]byte, error) {
	size := len(src)
	rem := (4 - size%4) % 4
	buf := make([]byte, size+rem)
	copy(buf, src)
	for i := 0; i < rem; i++ {
		buf[size+i] = '='
	}
	n, err := base64.URLEncoding.Decode(buf, buf)
	if err != nil {
		return nil, err
	}
	return buf[:n], nil
}

// sign returns signed data.
func (store *SessionCookieStore) sign(src []byte) []byte {
	sign := store.hash(src)
	return append(sign, src...)
}

// verify verify signed data and returns unsigned data if valid.
func (store *SessionCookieStore) verify(src []byte) (unsigned []byte, err error) {
	if len(src) <= sha512.Size256 {
		return nil, errors.New("kocha: session cookie value too short")
	}
	sign := src[:sha512.Size256]
	unsigned = src[sha512.Size256:]
	if !hmac.Equal(store.hash(unsigned), sign) {
		return nil, errors.New("kocha: session cookie verification failed")
	}
	return unsigned, nil
}

// hash returns hashed data by HMAC-SHA512/256.
func (store *SessionCookieStore) hash(src []byte) []byte {
	hash := hmac.New(sha512.New512_256, []byte(store.SigningKey))
	hash.Write(src)
	return hash.Sum(nil)
}
<|endoftext|>


<commit_before>package riffle

import (
	"fmt"
	"time"

	"github.com/gorilla/websocket"
)

type Session struct {
	Connection
	ReceiveTimeout time.Duration
	Auth           map[string]AuthFunc
	ReceiveDone    chan bool
	listeners      map[uint]chan Message
	events         map[uint]*eventDesc
	procedures     map[uint]*procedureDesc
	requestCount   uint
	pdid           string
}

// Connect to the node with the given URL
func Start(url string, domain string) (*Session, error) {
	dialer := websocket.Dialer{Subprotocols: []string{"wamp.2.json"}}

	conn, _, err := dialer.Dial(url, nil)

	if err != nil {
		return nil, err
	}

	connection := &websocketConnection{
		conn:        conn,
		messages:    make(chan Message, 10),
		serializer:  new(JSONSerializer),
		payloadType: websocket.TextMessage,
	}

	if err != nil {
		return nil, err
	}

	go connection.run()

	client := &Session{
		Connection:     connection,
		ReceiveTimeout: 10 * time.Second,
		listeners:      make(map[uint]chan Message),
		events:         make(map[uint]*eventDesc),
		procedures:     make(map[uint]*procedureDesc),
		requestCount:   0,
	}

	client.JoinRealm(domain, nil)
	return client, nil
}

/////////////////////////////////////////////
// Handler methods
/////////////////////////////////////////////

// EventHandler handles a publish event.
type EventHandler func(args []interface{}, kwargs map[string]interface{})

// Subscribe registers the EventHandler to be called for every message in the provided topic.
func (c *Session) Subscribe(topic string, fn interface{}) error {
	id := NewID()
	c.registerListener(id)

	sub := &Subscribe{
		Request: id,
		Options: make(map[string]interface{}),
		Domain:  topic,
	}

	if err := c.Send(sub); err != nil {
		return err
	}

	// wait to receive SUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error subscribing to topic '%v': %v", topic, e.Error)
	} else if subscribed, ok := msg.(*Subscribed); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, SUBSCRIBED))
	} else {
		// register the event handler with this subscription
		c.events[subscribed.Subscription] = &eventDesc{topic, fn}
	}
	return nil
}

// Unsubscribe removes the registered EventHandler from the topic.
func (c *Session) Unsubscribe(topic string) error {
	var (
		subscriptionID uint
		found          bool
	)
	for id, desc := range c.events {
		if desc.topic == topic {
			subscriptionID = id
			found = true
		}
	}
	if !found {
		return fmt.Errorf("Event %s is not registered with this client.", topic)
	}

	id := NewID()
	c.registerListener(id)
	sub := &Unsubscribe{
		Request:      id,
		Subscription: subscriptionID,
	}
	if err := c.Send(sub); err != nil {
		return err
	}
	// wait to receive UNSUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error unsubscribing to topic '%v': %v", topic, e.Error)
	} else if _, ok := msg.(*Unsubscribed); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, UNSUBSCRIBED))
	}
	delete(c.events, subscriptionID)
	return nil
}

// MethodHandler is an RPC endpoint.
type MethodHandler func(
	args []interface{}, kwargs map[string]interface{}, details map[string]interface{},
) (result *CallResult)

func (c *Session) Register(procedure string, fn interface{}, options map[string]interface{}) error {
	id := NewID()
	c.registerListener(id)

	register := &Register{
		Request: id,
		Options: options,
		Domain:  procedure,
	}

	if err := c.Send(register); err != nil {
		return err
	}

	// wait to receive REGISTERED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error registering procedure '%v': %v", procedure, e.Error)
	} else if registered, ok := msg.(*Registered); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, REGISTERED))
	} else {
		// register the event handler with this registration
		c.procedures[registered.Registration] = &procedureDesc{procedure, fn}
	}
	return nil
}

// Unregister removes a procedure with the Node
func (c *Session) Unregister(procedure string) error {
	var (
		procedureID uint
		found       bool
	)

	for id, p := range c.procedures {
		if p.name == procedure {
			procedureID = id
			found = true
		}
	}

	if !found {
		return fmt.Errorf("Domain %s is not registered with this client.", procedure)
	}

	id := NewID()
	c.registerListener(id)
	unregister := &Unregister{
		Request:      id,
		Registration: procedureID,
	}

	if err := c.Send(unregister); err != nil {
		return err
	}

	// wait to receive UNREGISTERED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error unregister to procedure '%v': %v", procedure, e.Error)
	} else if _, ok := msg.(*Unregistered); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, UNREGISTERED))
	}

	// register the event handler with this unregistration
	delete(c.procedures, procedureID)
	return nil
}

// Publish publishes an EVENT to all subscribed peers.
func (c *Session) Publish(topic string, args ...interface{}) error {
	return c.Send(&Publish{
		Request:   NewID(),
		Options:   make(map[string]interface{}),
		Domain:    topic,
		Arguments: args,
	})
}

// Call calls a procedure given a URI.
func (c *Session) Call(procedure string, args ...interface{}) ([]interface{}, error) {
	id := NewID()
	c.registerListener(id)

	call := &Call{
		Request:   id,
		Domain:    procedure,
		Options:   make(map[string]interface{}),
		Arguments: args,
	}

	if err := c.Send(call); err != nil {
		return nil, err
	}

	// wait to receive RESULT message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return nil, err
	} else if e, ok := msg.(*Error); ok {
		return nil, fmt.Errorf("error calling procedure '%v': %v", procedure, e.Error)
	} else if result, ok := msg.(*Result); !ok {
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, RESULT))
	} else {
		return result.Arguments, nil
	}
}

func (c *Session) Leave() error {
	if err := c.Send(goodbyeSession); err != nil {
		return fmt.Errorf("error leaving realm: %v", err)
	}

	if err := c.Connection.Close(); err != nil {
		return fmt.Errorf("error closing client connection: %v", err)
	}

	return nil
}

/////////////////////////////////////////////
// Misc
/////////////////////////////////////////////

// JoinRealm joins a WAMP realm, but does not handle challenge/response authentication.
func (c *Session) JoinRealm(realm string, details map[string]interface{}) (map[string]interface{}, error) {
	if details == nil {
		details = map[string]interface{}{}
	}

	details["roles"] = map[string]map[string]interface{}{
		"publisher":  make(map[string]interface{}),
		"subscriber": make(map[string]interface{}),
		"callee":     make(map[string]interface{}),
		"caller":     make(map[string]interface{}),
	}

	if c.Auth != nil && len(c.Auth) > 0 {
		return c.joinRealmCRA(realm, details)
	}

	if err := c.Send(&Hello{Realm: realm, Details: details}); err != nil {
		c.Connection.Close()
		return nil, err
	}

	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if welcome, ok := msg.(*Welcome); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, WELCOME))
	} else {
		go c.Receive()
		return welcome.Details, nil
	}
}

// joinRealmCRA joins a WAMP realm and handles challenge/response authentication.
func (c *Session) joinRealmCRA(realm string, details map[string]interface{}) (map[string]interface{}, error) {
	authmethods := []interface{}{}
	for m := range c.Auth {
		authmethods = append(authmethods, m)
	}
	details["authmethods"] = authmethods
	if err := c.Send(&Hello{Realm: realm, Details: details}); err != nil {
		c.Connection.Close()
		return nil, err
	}
	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if challenge, ok := msg.(*Challenge); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, CHALLENGE))
	} else if authFunc, ok := c.Auth[challenge.AuthMethod]; !ok {
		c.Send(abortNoAuthHandler)
		c.Connection.Close()
		return nil, fmt.Errorf("no auth handler for method: %s", challenge.AuthMethod)
	} else if signature, authDetails, err := authFunc(details, challenge.Extra); err != nil {
		c.Send(abortAuthFailure)
		c.Connection.Close()
		return nil, err
	} else if err := c.Send(&Authenticate{Signature: signature, Extra: authDetails}); err != nil {
		c.Connection.Close()
		return nil, err
	}
	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if welcome, ok := msg.(*Welcome); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, WELCOME))
	} else {
		go c.Receive()
		return welcome.Details, nil
	}
}
<commit_msg>that was not worth it<commit_after>package riffle

import (
	"fmt"
	"time"

	"github.com/gorilla/websocket"
)

type Session struct {
	Connection
	ReceiveTimeout time.Duration
	Auth           map[string]AuthFunc
	ReceiveDone    chan bool
	listeners      map[uint]chan Message
	events         map[uint]*eventDesc
	procedures     map[uint]*procedureDesc
	requestCount   uint
	pdid           string
}

// Connect to the node with the given URL
func Start(url string, domain string) (*Session, error) {
	dialer := websocket.Dialer{Subprotocols: []string{"wamp.2.json"}}

	conn, _, err := dialer.Dial(url, nil)

	if err != nil {
		return nil, err
	}

	connection := &websocketConnection{
		conn:        conn,
		messages:    make(chan Message, 10),
		serializer:  new(JSONSerializer),
		payloadType: websocket.TextMessage,
	}

	if err != nil {
		return nil, err
	}

	go connection.run()

	client := &Session{
		Connection:     connection,
		ReceiveTimeout: 10 * time.Second,
		listeners:      make(map[uint]chan Message),
		events:         make(map[uint]*eventDesc),
		procedures:     make(map[uint]*procedureDesc),
		requestCount:   0,
	}

	client.JoinRealm(domain, nil)
	return client, nil
}

/////////////////////////////////////////////
// Handler methods
/////////////////////////////////////////////

// Subscribe registers the EventHandler to be called for every message in the provided topic.
func (c *Session) Subscribe(topic string, fn interface{}) error {
	id := NewID()
	c.registerListener(id)

	sub := &Subscribe{
		Request: id,
		Options: make(map[string]interface{}),
		Domain:  topic,
	}

	if err := c.Send(sub); err != nil {
		return err
	}

	// wait to receive SUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error subscribing to topic '%v': %v", topic, e.Error)
	} else if subscribed, ok := msg.(*Subscribed); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, SUBSCRIBED))
	} else {
		// register the event handler with this subscription
		c.events[subscribed.Subscription] = &eventDesc{topic, fn}
	}
	return nil
}

// Unsubscribe removes the registered EventHandler from the topic.
func (c *Session) Unsubscribe(topic string) error {
	var (
		subscriptionID uint
		found          bool
	)

	for id, desc := range c.events {
		if desc.topic == topic {
			subscriptionID = id
			found = true
		}
	}

	if !found {
		return fmt.Errorf("Event %s is not registered with this client.", topic)
	}

	id := NewID()
	c.registerListener(id)

	sub := &Unsubscribe{
		Request:      id,
		Subscription: subscriptionID,
	}

	if err := c.Send(sub); err != nil {
		return err
	}

	// wait to receive UNSUBSCRIBED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error unsubscribing to topic '%v': %v", topic, e.Error)
	} else if _, ok := msg.(*Unsubscribed); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, UNSUBSCRIBED))
	}

	delete(c.events, subscriptionID)
	return nil
}

func (c *Session) Register(procedure string, fn interface{}, options map[string]interface{}) error {
	id := NewID()
	c.registerListener(id)

	register := &Register{
		Request: id,
		Options: options,
		Domain:  procedure,
	}

	if err := c.Send(register); err != nil {
		return err
	}

	// wait to receive REGISTERED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error registering procedure '%v': %v", procedure, e.Error)
	} else if registered, ok := msg.(*Registered); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, REGISTERED))
	} else {
		// register the event handler with this registration
		c.procedures[registered.Registration] = &procedureDesc{procedure, fn}
	}
	return nil
}

// Unregister removes a procedure with the Node
func (c *Session) Unregister(procedure string) error {
	var (
		procedureID uint
		found       bool
	)

	for id, p := range c.procedures {
		if p.name == procedure {
			procedureID = id
			found = true
		}
	}

	if !found {
		return fmt.Errorf("Domain %s is not registered with this client.", procedure)
	}

	id := NewID()
	c.registerListener(id)
	unregister := &Unregister{
		Request:      id,
		Registration: procedureID,
	}

	if err := c.Send(unregister); err != nil {
		return err
	}

	// wait to receive UNREGISTERED message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return err
	} else if e, ok := msg.(*Error); ok {
		return fmt.Errorf("error unregister to procedure '%v': %v", procedure, e.Error)
	} else if _, ok := msg.(*Unregistered); !ok {
		return fmt.Errorf(formatUnexpectedMessage(msg, UNREGISTERED))
	}

	// register the event handler with this unregistration
	delete(c.procedures, procedureID)
	return nil
}

// Publish publishes an EVENT to all subscribed peers.
func (c *Session) Publish(endpoint string, args ...interface{}) error {
	return c.Send(&Publish{
		Request:   NewID(),
		Options:   make(map[string]interface{}),
		Domain:    endpoint,
		Arguments: args,
	})
}

// Call calls a procedure given a URI.
func (c *Session) Call(procedure string, args ...interface{}) ([]interface{}, error) {
	id := NewID()
	c.registerListener(id)

	call := &Call{
		Request:   id,
		Domain:    procedure,
		Options:   make(map[string]interface{}),
		Arguments: args,
	}

	if err := c.Send(call); err != nil {
		return nil, err
	}

	// wait to receive RESULT message
	msg, err := c.waitOnListener(id)
	if err != nil {
		return nil, err
	} else if e, ok := msg.(*Error); ok {
		return nil, fmt.Errorf("error calling procedure '%v': %v", procedure, e.Error)
	} else if result, ok := msg.(*Result); !ok {
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, RESULT))
	} else {
		return result.Arguments, nil
	}
}

func (c *Session) Leave() error {
	if err := c.Send(goodbyeSession); err != nil {
		return fmt.Errorf("error leaving realm: %v", err)
	}

	if err := c.Connection.Close(); err != nil {
		return fmt.Errorf("error closing client connection: %v", err)
	}

	return nil
}

/////////////////////////////////////////////
// Misc
/////////////////////////////////////////////

// JoinRealm joins a WAMP realm, but does not handle challenge/response authentication.
func (c *Session) JoinRealm(realm string, details map[string]interface{}) (map[string]interface{}, error) {
	if details == nil {
		details = map[string]interface{}{}
	}

	details["roles"] = map[string]map[string]interface{}{
		"publisher":  make(map[string]interface{}),
		"subscriber": make(map[string]interface{}),
		"callee":     make(map[string]interface{}),
		"caller":     make(map[string]interface{}),
	}

	if c.Auth != nil && len(c.Auth) > 0 {
		return c.joinRealmCRA(realm, details)
	}

	if err := c.Send(&Hello{Realm: realm, Details: details}); err != nil {
		c.Connection.Close()
		return nil, err
	}

	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if welcome, ok := msg.(*Welcome); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, WELCOME))
	} else {
		go c.Receive()
		return welcome.Details, nil
	}
}

// joinRealmCRA joins a WAMP realm and handles challenge/response authentication.
func (c *Session) joinRealmCRA(realm string, details map[string]interface{}) (map[string]interface{}, error) {
	authmethods := []interface{}{}
	for m := range c.Auth {
		authmethods = append(authmethods, m)
	}
	details["authmethods"] = authmethods
	if err := c.Send(&Hello{Realm: realm, Details: details}); err != nil {
		c.Connection.Close()
		return nil, err
	}
	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if challenge, ok := msg.(*Challenge); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, CHALLENGE))
	} else if authFunc, ok := c.Auth[challenge.AuthMethod]; !ok {
		c.Send(abortNoAuthHandler)
		c.Connection.Close()
		return nil, fmt.Errorf("no auth handler for method: %s", challenge.AuthMethod)
	} else if signature, authDetails, err := authFunc(details, challenge.Extra); err != nil {
		c.Send(abortAuthFailure)
		c.Connection.Close()
		return nil, err
	} else if err := c.Send(&Authenticate{Signature: signature, Extra: authDetails}); err != nil {
		c.Connection.Close()
		return nil, err
	}
	if msg, err := GetMessageTimeout(c.Connection, c.ReceiveTimeout); err != nil {
		c.Connection.Close()
		return nil, err
	} else if welcome, ok := msg.(*Welcome); !ok {
		c.Send(abortUnexpectedMsg)
		c.Connection.Close()
		return nil, fmt.Errorf(formatUnexpectedMessage(msg, WELCOME))
	} else {
		go c.Receive()
		return welcome.Details, nil
	}
}
<|endoftext|>