<commit_before>package registry

import (
	"encoding/json"
	"errors"
	"sync"

	consul "github.com/hashicorp/consul/api"
)

type consulRegistry struct {
	Address string
	Client  *consul.Client

	mtx      sync.RWMutex
	services map[string]*Service
}

func encodeMetadata(md map[string]string) []string {
	var tags []string
	for k, v := range md {
		if b, err := json.Marshal(map[string]string{
			k: v,
		}); err == nil {
			tags = append(tags, string(b))
		}
	}
	return tags
}

func decodeMetadata(tags []string) map[string]string {
	md := make(map[string]string)
	for _, tag := range tags {
		var kv map[string]string
		if err := json.Unmarshal([]byte(tag), &kv); err == nil {
			for k, v := range kv {
				md[k] = v
			}
		}
	}
	return md
}

func newConsulRegistry(addrs []string, opts ...Option) Registry {
	config := consul.DefaultConfig()
	client, _ := consul.NewClient(config)
	if len(addrs) > 0 {
		config.Address = addrs[0]
	}

	cr := &consulRegistry{
		Address:  config.Address,
		Client:   client,
		services: make(map[string]*Service),
	}

	return cr
}

func (c *consulRegistry) Deregister(s *Service) error {
	if len(s.Nodes) == 0 {
		return errors.New("Require at least one node")
	}

	node := s.Nodes[0]

	_, err := c.Client.Catalog().Deregister(&consul.CatalogDeregistration{
		Node:      node.Id,
		Address:   node.Address,
		ServiceID: node.Id,
	}, nil)

	return err
}

func (c *consulRegistry) Register(s *Service) error {
	if len(s.Nodes) == 0 {
		return errors.New("Require at least one node")
	}

	node := s.Nodes[0]

	tags := encodeMetadata(node.Metadata)

	_, err := c.Client.Catalog().Register(&consul.CatalogRegistration{
		Node:    node.Id,
		Address: node.Address,
		Service: &consul.AgentService{
			ID:      node.Id,
			Service: s.Name,
			Port:    node.Port,
			Tags:    tags,
		},
	}, nil)

	return err
}

func (c *consulRegistry) GetService(name string) (*Service, error) {
	c.mtx.RLock()
	service, ok := c.services[name]
	c.mtx.RUnlock()

	if ok {
		return service, nil
	}

	rsp, _, err := c.Client.Catalog().Service(name, "", nil)
	if err != nil {
		return nil, err
	}

	cs := &Service{}

	for _, s := range rsp {
		if s.ServiceName != name {
			continue
		}

		cs.Name = s.ServiceName
		cs.Nodes = append(cs.Nodes, &Node{
			Id:       s.ServiceID,
			Address:  s.Address,
			Port:     s.ServicePort,
			Metadata: decodeMetadata(s.ServiceTags),
		})
	}

	return cs, nil
}

func (c *consulRegistry) ListServices() ([]*Service, error) {
	c.mtx.RLock()
	serviceMap := c.services
	c.mtx.RUnlock()

	var services []*Service

	if len(serviceMap) > 0 {
		for _, service := range services {
			services = append(services, service)
		}
		return services, nil
	}

	rsp, _, err := c.Client.Catalog().Services(&consul.QueryOptions{})
	if err != nil {
		return nil, err
	}

	for service, _ := range rsp {
		services = append(services, &Service{Name: service})
	}

	return services, nil
}

func (c *consulRegistry) Watch() (Watcher, error) {
	return newConsulWatcher(c)
}
<commit_msg>Remove service id so nodes are also removed from catalog<commit_after>package registry

import (
	"encoding/json"
	"errors"
	"sync"

	consul "github.com/hashicorp/consul/api"
)

type consulRegistry struct {
	Address string
	Client  *consul.Client

	mtx      sync.RWMutex
	services map[string]*Service
}

func encodeMetadata(md map[string]string) []string {
	var tags []string
	for k, v := range md {
		if b, err := json.Marshal(map[string]string{
			k: v,
		}); err == nil {
			tags = append(tags, string(b))
		}
	}
	return tags
}

func decodeMetadata(tags []string) map[string]string {
	md := make(map[string]string)
	for _, tag := range tags {
		var kv map[string]string
		if err := json.Unmarshal([]byte(tag), &kv); err == nil {
			for k, v := range kv {
				md[k] = v
			}
		}
	}
	return md
}

func newConsulRegistry(addrs []string, opts ...Option) Registry {
	config := consul.DefaultConfig()
	client, _ := consul.NewClient(config)
	if len(addrs) > 0 {
		config.Address = addrs[0]
	}

	cr := &consulRegistry{
		Address:  config.Address,
		Client:   client,
		services: make(map[string]*Service),
	}

	return cr
}

func (c *consulRegistry) Deregister(s *Service) error {
	if len(s.Nodes) == 0 {
		return errors.New("Require at least one node")
	}

	node := s.Nodes[0]

	_, err := c.Client.Catalog().Deregister(&consul.CatalogDeregistration{
		Node:    node.Id,
		Address: node.Address,
	}, nil)

	return err
}

func (c *consulRegistry) Register(s *Service) error {
	if len(s.Nodes) == 0 {
		return errors.New("Require at least one node")
	}

	node := s.Nodes[0]

	tags := encodeMetadata(node.Metadata)

	_, err := c.Client.Catalog().Register(&consul.CatalogRegistration{
		Node:    node.Id,
		Address: node.Address,
		Service: &consul.AgentService{
			ID:      node.Id,
			Service: s.Name,
			Port:    node.Port,
			Tags:    tags,
		},
	}, nil)

	return err
}

func (c *consulRegistry) GetService(name string) (*Service, error) {
	c.mtx.RLock()
	service, ok := c.services[name]
	c.mtx.RUnlock()

	if ok {
		return service, nil
	}

	rsp, _, err := c.Client.Catalog().Service(name, "", nil)
	if err != nil {
		return nil, err
	}

	cs := &Service{}

	for _, s := range rsp {
		if s.ServiceName != name {
			continue
		}

		cs.Name = s.ServiceName
		cs.Nodes = append(cs.Nodes, &Node{
			Id:       s.ServiceID,
			Address:  s.Address,
			Port:     s.ServicePort,
			Metadata: decodeMetadata(s.ServiceTags),
		})
	}

	return cs, nil
}

func (c *consulRegistry) ListServices() ([]*Service, error) {
	c.mtx.RLock()
	serviceMap := c.services
	c.mtx.RUnlock()

	var services []*Service

	if len(serviceMap) > 0 {
		for _, service := range services {
			services = append(services, service)
		}
		return services, nil
	}

	rsp, _, err := c.Client.Catalog().Services(&consul.QueryOptions{})
	if err != nil {
		return nil, err
	}

	for service, _ := range rsp {
		services = append(services, &Service{Name: service})
	}

	return services, nil
}

func (c *consulRegistry) Watch() (Watcher, error) {
	return newConsulWatcher(c)
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/aws/sdb/mock"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/crypto/mock"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestRegistry(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type registryTest struct {
	crypter mock_crypto.MockCrypter
	domain  mock_sdb.MockDomain
}

func (t *registryTest) SetUp(i *TestInfo) {
	t.crypter = mock_crypto.NewMockCrypter(i.MockController, "crypter")
	t.domain = mock_sdb.NewMockDomain(i.MockController, "domain")
}

////////////////////////////////////////////////////////////////////////
// NewRegistry
////////////////////////////////////////////////////////////////////////

type NewRegistryTest struct {
	registryTest

	r backup.Registry
	err error
}

func init() { RegisterTestSuite(&NewRegistryTest{}) }

func (t *NewRegistryTest) callConstructor() {
	t.r, t.err = backup.NewRegistry(t.crypter, t.domain)
}

func (t *NewRegistryTest) CallsGetAttributes() {
	// Domain
	ExpectCall(t.domain, "GetAttributes")(
		"comeback_marker",
		false,
		ElementsAre("encrypted_data")).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callConstructor()
}

func (t *NewRegistryTest) GetAttributesReturnsError() {
	// Domain
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callConstructor()

	ExpectThat(t.err, Error(HasSubstr("GetAttributes")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *NewRegistryTest) CallsDecrypt() {
	// Domain
	attr := sdb.Attribute{Name: "encrypted_data", Value: "taco"}
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return([]sdb.Attribute{attr}, nil))

	// Crypter
	ExpectCall(t.crypter, "Decrypt")(DeepEquals([]byte("taco"))).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callConstructor()
}

func (t *NewRegistryTest) DecryptReturnsGenericError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) DecryptReturnsNotAuthenticError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) DecryptSucceeds() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) CallsEncrypt() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) EncryptReturnsError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) CallsPutAttributes() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) PutAttributesReturnsError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) PutAttributesSucceeds() {
	ExpectEq("TODO", "")
}
<commit_msg>NewRegistryTest.DecryptReturnsNotAuthenticError<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/aws/sdb"
	"github.com/jacobsa/aws/sdb/mock"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/crypto"
	"github.com/jacobsa/comeback/crypto/mock"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestRegistry(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type registryTest struct {
	crypter mock_crypto.MockCrypter
	domain  mock_sdb.MockDomain
}

func (t *registryTest) SetUp(i *TestInfo) {
	t.crypter = mock_crypto.NewMockCrypter(i.MockController, "crypter")
	t.domain = mock_sdb.NewMockDomain(i.MockController, "domain")
}

////////////////////////////////////////////////////////////////////////
// NewRegistry
////////////////////////////////////////////////////////////////////////

type NewRegistryTest struct {
	registryTest

	r backup.Registry
	err error
}

func init() { RegisterTestSuite(&NewRegistryTest{}) }

func (t *NewRegistryTest) callConstructor() {
	t.r, t.err = backup.NewRegistry(t.crypter, t.domain)
}

func (t *NewRegistryTest) CallsGetAttributes() {
	// Domain
	ExpectCall(t.domain, "GetAttributes")(
		"comeback_marker",
		false,
		ElementsAre("encrypted_data")).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callConstructor()
}

func (t *NewRegistryTest) GetAttributesReturnsError() {
	// Domain
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callConstructor()

	ExpectThat(t.err, Error(HasSubstr("GetAttributes")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *NewRegistryTest) CallsDecrypt() {
	// Domain
	attr := sdb.Attribute{Name: "encrypted_data", Value: "taco"}
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return([]sdb.Attribute{attr}, nil))

	// Crypter
	ExpectCall(t.crypter, "Decrypt")(DeepEquals([]byte("taco"))).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callConstructor()
}

func (t *NewRegistryTest) DecryptReturnsGenericError() {
	// Domain
	attr := sdb.Attribute{Name: "encrypted_data"}
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return([]sdb.Attribute{attr}, nil))

	// Crypter
	ExpectCall(t.crypter, "Decrypt")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callConstructor()

	ExpectThat(t.err, Error(HasSubstr("Decrypt")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *NewRegistryTest) DecryptReturnsNotAuthenticError() {
	// Domain
	attr := sdb.Attribute{Name: "encrypted_data"}
	ExpectCall(t.domain, "GetAttributes")(Any(), Any(), Any()).
		WillOnce(oglemock.Return([]sdb.Attribute{attr}, nil))

	// Crypter
	ExpectCall(t.crypter, "Decrypt")(Any()).
		WillOnce(oglemock.Return(nil, &crypto.NotAuthenticError{}))

	// Call
	t.callConstructor()

	_, ok := t.err.(*backup.IncompatibleCrypterError)
	AssertTrue(ok, "Error: %v", t.err)

	ExpectThat(t.err, Error(HasSubstr("crypter")))
	ExpectThat(t.err, Error(HasSubstr("incompatible")))
}

func (t *NewRegistryTest) DecryptSucceeds() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) CallsEncrypt() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) EncryptReturnsError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) CallsPutAttributes() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) PutAttributesReturnsError() {
	ExpectEq("TODO", "")
}

func (t *NewRegistryTest) PutAttributesSucceeds() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Workiva, LLC

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 batcher

import (
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestNoCalculateBytes(t *testing.T) {
	_, err := New(0, 0, 100, 5, nil)
	assert.Error(t, err)
}

func TestMaxItems(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	for i := 0; i < 1000; i++ {
		assert.Nil(b.Put("foo bar baz"))
	}

	batch, err := b.Get()
	assert.Len(batch, 100)
	assert.Nil(err)
}

func TestMaxBytes(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10000, 100, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	go func() {
		for i := 0; i < 1000; i++ {
			b.Put("a")
		}
	}()

	batch, err := b.Get()
	assert.Len(batch, 100)
	assert.Nil(err)
}

func TestMaxTime(t *testing.T) {
	assert := assert.New(t)
	b, err := New(time.Millisecond*200, 100000, 100000, 10,
		func(str interface{}) uint {
			return uint(len(str.(string)))
		},
	)
	assert.Nil(err)

	go func() {
		for i := 0; i < 10000; i++ {
			b.Put("a")
			time.Sleep(time.Millisecond)
		}
	}()

	before := time.Now()
	batch, err := b.Get()

	// This delta is normally 1-3 ms but running tests in CI with -race causes
	// this to run much slower. For now, just bump up the threshold.
	assert.InDelta(200, time.Since(before).Seconds()*1000, 50)
	assert.True(len(batch) > 0)
	assert.Nil(err)
}

func TestFlush(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	b.Put("a")
	wait := make(chan bool)
	go func() {
		batch, err := b.Get()
		assert.Equal([]interface{}{"a"}, batch)
		assert.Nil(err)
		wait <- true
	}()

	b.Flush()
	<-wait
}

func TestMultiConsumer(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	var wg sync.WaitGroup
	wg.Add(5)
	for i := 0; i < 5; i++ {
		go func() {
			batch, err := b.Get()
			assert.Len(batch, 100)
			assert.Nil(err)
			wg.Done()
		}()
	}

	go func() {
		for i := 0; i < 500; i++ {
			b.Put("a")
		}
	}()

	wg.Wait()
}

func TestDispose(t *testing.T) {
	assert := assert.New(t)
	b, err := New(1, 2, 100000, 2, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	b.Put("a")
	b.Put("b")
	b.Put("c")

	possibleBatches := [][]interface{}{
		[]interface{}{"a", "b"},
		[]interface{}{"c"},
	}
	batch1, err := b.Get()
	assert.Contains(possibleBatches, batch1)
	assert.Nil(err)
	batch2, err := b.Get()
	assert.Contains(possibleBatches, batch2)
	assert.Nil(err)

	b.Put("d")
	b.Put("e")
	b.Put("f")
	b.Put("g")
	b.Put("h")
	b.Put("i")

	b.Dispose()

	_, err = b.Get()
	assert.Equal(ErrDisposed, err)

	assert.Equal(ErrDisposed, b.Put("j"))
	assert.Equal(ErrDisposed, b.Flush())

}

func TestIsDisposed(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	assert.False(b.IsDisposed())
	b.Dispose()
	assert.True(b.IsDisposed())
}
<commit_msg>Wait for items to make it to batchChan in test<commit_after>/*
Copyright 2015 Workiva, LLC

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 batcher

import (
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestNoCalculateBytes(t *testing.T) {
	_, err := New(0, 0, 100, 5, nil)
	assert.Error(t, err)
}

func TestMaxItems(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	for i := 0; i < 1000; i++ {
		assert.Nil(b.Put("foo bar baz"))
	}

	batch, err := b.Get()
	assert.Len(batch, 100)
	assert.Nil(err)
}

func TestMaxBytes(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10000, 100, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	go func() {
		for i := 0; i < 1000; i++ {
			b.Put("a")
		}
	}()

	batch, err := b.Get()
	assert.Len(batch, 100)
	assert.Nil(err)
}

func TestMaxTime(t *testing.T) {
	assert := assert.New(t)
	b, err := New(time.Millisecond*200, 100000, 100000, 10,
		func(str interface{}) uint {
			return uint(len(str.(string)))
		},
	)
	assert.Nil(err)

	go func() {
		for i := 0; i < 10000; i++ {
			b.Put("a")
			time.Sleep(time.Millisecond)
		}
	}()

	before := time.Now()
	batch, err := b.Get()

	// This delta is normally 1-3 ms but running tests in CI with -race causes
	// this to run much slower. For now, just bump up the threshold.
	assert.InDelta(200, time.Since(before).Seconds()*1000, 50)
	assert.True(len(batch) > 0)
	assert.Nil(err)
}

func TestFlush(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	b.Put("a")
	wait := make(chan bool)
	go func() {
		batch, err := b.Get()
		assert.Equal([]interface{}{"a"}, batch)
		assert.Nil(err)
		wait <- true
	}()

	b.Flush()
	<-wait
}

func TestMultiConsumer(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)

	var wg sync.WaitGroup
	wg.Add(5)
	for i := 0; i < 5; i++ {
		go func() {
			batch, err := b.Get()
			assert.Len(batch, 100)
			assert.Nil(err)
			wg.Done()
		}()
	}

	go func() {
		for i := 0; i < 500; i++ {
			b.Put("a")
		}
	}()

	wg.Wait()
}

func TestDispose(t *testing.T) {
	assert := assert.New(t)
	b, err := New(1, 2, 100000, 2, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	b.Put("a")
	b.Put("b")
	b.Put("c")

	possibleBatches := [][]interface{}{
		[]interface{}{"a", "b"},
		[]interface{}{"c"},
	}

	// Wait for items to get to the channel
	for len(b.(*basicBatcher).batchChan) == 0 {
		time.Sleep(1 * time.Millisecond)
	}
	batch1, err := b.Get()
	assert.Contains(possibleBatches, batch1)
	assert.Nil(err)

	batch2, err := b.Get()
	assert.Contains(possibleBatches, batch2)
	assert.Nil(err)

	b.Put("d")
	b.Put("e")
	b.Put("f")
	b.Put("g")
	b.Put("h")
	b.Put("i")

	b.Dispose()

	_, err = b.Get()
	assert.Equal(ErrDisposed, err)

	assert.Equal(ErrDisposed, b.Put("j"))
	assert.Equal(ErrDisposed, b.Flush())

}

func TestIsDisposed(t *testing.T) {
	assert := assert.New(t)
	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
		return uint(len(str.(string)))
	})
	assert.Nil(err)
	assert.False(b.IsDisposed())
	b.Dispose()
	assert.True(b.IsDisposed())
}
<|endoftext|>


<commit_before>package typescriptify

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"time"
)

type TypeScriptify struct {
	Prefix           string
	Suffix           string
	Indent           string
	CreateFromMethod bool
	BackupExtension  string // If empty no backup

	golangTypes []reflect.Type
	types       map[reflect.Kind]string

	// throwaway, used when converting
	alreadyConverted map[reflect.Type]bool
}

func New() *TypeScriptify {
	result := new(TypeScriptify)
	result.Indent = "\t"
	result.BackupExtension = "backup"

	types := make(map[reflect.Kind]string)

	types[reflect.Bool] = "boolean"

	types[reflect.Int] = "number"
	types[reflect.Int8] = "number"
	types[reflect.Int16] = "number"
	types[reflect.Int32] = "number"
	types[reflect.Int64] = "number"
	types[reflect.Uint] = "number"
	types[reflect.Uint8] = "number"
	types[reflect.Uint16] = "number"
	types[reflect.Uint32] = "number"
	types[reflect.Uint64] = "number"
	types[reflect.Float32] = "number"
	types[reflect.Float64] = "number"

	types[reflect.String] = "string"

	result.types = types

	result.Indent = "    "
	result.CreateFromMethod = true

	return result
}

func deepFields(typeOf reflect.Type) []reflect.StructField {
	fields := make([]reflect.StructField, 0)

	if typeOf.Kind() == reflect.Ptr {
		typeOf = typeOf.Elem()
	}

	if typeOf.Kind() != reflect.Struct {
		return fields
	}

	for i := 0; i < typeOf.NumField(); i++ {
		f := typeOf.Field(i)

		kind := f.Type.Kind()
		if f.Anonymous && kind == reflect.Struct {
			//fmt.Println(v.Interface())
			fields = append(fields, deepFields(f.Type)...)
		} else {
			fields = append(fields, f)
		}
	}

	return fields
}

func (this *TypeScriptify) Add(obj interface{}) {
	this.AddType(reflect.TypeOf(obj))
}

func (this *TypeScriptify) AddType(typeOf reflect.Type) {
	this.golangTypes = append(this.golangTypes, typeOf)
}

func (this *TypeScriptify) Convert(customCode map[string]string) (string, error) {
	this.alreadyConverted = make(map[reflect.Type]bool)

	result := ""
	for _, typeof := range this.golangTypes {
		typeScriptCode, err := this.convertType(typeof, customCode)
		if err != nil {
			return "", err
		}
		result += "\n" + strings.Trim(typeScriptCode, " "+this.Indent+"\r\n")
	}
	return result, nil
}

func loadCustomCode(fileName string) (map[string]string, error) {
	result := make(map[string]string)
	f, err := os.Open(fileName)
	if err != nil {
		if os.IsNotExist(err) {
			return result, nil
		}
		return result, err
	}
	defer f.Close()

	bytes, err := ioutil.ReadAll(f)
	if err != nil {
		return result, err
	}

	var currentName string
	var currentValue string
	lines := strings.Split(string(bytes), "\n")
	for _, line := range lines {
		trimmedLine := strings.TrimSpace(line)
		if strings.HasPrefix(trimmedLine, "//[") && strings.HasSuffix(trimmedLine, ":]") {
			currentName = strings.Replace(strings.Replace(trimmedLine, "//[", "", -1), ":]", "", -1)
			currentValue = ""
		} else if trimmedLine == "//[end]" {
			result[currentName] = strings.TrimRight(currentValue, " \t\r\n")
			currentName = ""
			currentValue = ""
		} else if len(currentName) > 0 {
			currentValue += line + "\n"
		}
	}

	return result, nil
}

func (this TypeScriptify) backup(fileName string) error {
	fileIn, err := os.Open(fileName)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		// No neet to backup, just return:
		return nil
	}
	defer fileIn.Close()

	bytes, err := ioutil.ReadAll(fileIn)
	if err != nil {
		return err
	}

	fileOut, err := os.Create(fmt.Sprintf("%s-%s.%s", fileName, time.Now().Format("2006-01-02T15:04:05.99"), this.BackupExtension))
	if err != nil {
		return err
	}
	defer fileOut.Close()

	_, err = fileOut.Write(bytes)
	if err != nil {
		return err
	}

	return nil
}

func (this TypeScriptify) ConvertToFile(fileName string) error {
	if len(this.BackupExtension) > 0 {
		err := this.backup(fileName)
		if err != nil {
			return err
		}
	}

	customCode, err := loadCustomCode(fileName)
	if err != nil {
		return err
	}

	f, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	converted, err := this.Convert(customCode)
	if err != nil {
		return err
	}

	f.WriteString("/* Do not change, this code is generated from Golang structs */\n\n")
	f.WriteString(converted)
	if err != nil {
		return err
	}

	return nil
}

func (this *TypeScriptify) convertType(typeOf reflect.Type, customCode map[string]string) (string, error) {
	if _, found := this.alreadyConverted[typeOf]; found { // Already converted
		return "", nil
	}

	entityName := fmt.Sprintf("%s%s%s", this.Prefix, this.Suffix, typeOf.Name())
	result := fmt.Sprintf("class %s {\n", entityName)
	builder := typeScriptClassBuilder{
		types:  this.types,
		indent: this.Indent,
	}

	fields := deepFields(typeOf)
	for _, field := range fields {
		jsonTag := field.Tag.Get("json")
		jsonFieldName := ""
		if len(jsonTag) > 0 {
			jsonTagParts := strings.Split(jsonTag, ",")
			if len(jsonTagParts) > 0 {
				jsonFieldName = strings.Trim(jsonTagParts[0], this.Indent)
			}
		}
		if len(jsonFieldName) > 0 && jsonFieldName != "-" {
			var err error
			if field.Type.Kind() == reflect.Struct { // Struct:
				typeScriptChunk, err := this.convertType(field.Type, customCode)
				if err != nil {
					return "", err
				}
				result = typeScriptChunk + "\n" + result
				builder.AddStructField(jsonFieldName, field.Type.Name())
			} else if field.Type.Kind() == reflect.Slice { // Slice:
				if field.Type.Elem().Kind() == reflect.Struct { // Slice of structs:
					typeScriptChunk, err := this.convertType(field.Type.Elem(), customCode)
					if err != nil {
						return "", err
					}
					result = typeScriptChunk + "\n" + result
					builder.AddArrayOfStructsField(jsonFieldName, field.Type.Elem().Name())
				} else { // Slice of simple fields:
					err = builder.AddSimpleArrayField(jsonFieldName, field.Type.Elem().Name(), field.Type.Elem().Kind())
				}
			} else { // Simple field:
				err = builder.AddSimpleField(jsonFieldName, field.Type.Name(), field.Type.Kind())
			}
			if err != nil {
				return "", err
			}
		}
	}

	result += builder.fields
	if this.CreateFromMethod {
		result += fmt.Sprintf("\n%sstatic createFrom(source: any) {\n", this.Indent)
		result += fmt.Sprintf("%s%svar result = new %s();\n", this.Indent, this.Indent, entityName)
		result += builder.createFromMethodBody
		result += fmt.Sprintf("%s%sreturn result;\n", this.Indent, this.Indent)
		result += fmt.Sprintf("%s}\n\n", this.Indent)
	}

	if customCode != nil {
		code := customCode[entityName]
		result += this.Indent + "//[" + entityName + ":]\n" + code + "\n\n" + this.Indent + "//[end]\n"
	}

	result += "}"

	this.alreadyConverted[typeOf] = true

	return result, nil
}

type typeScriptClassBuilder struct {
	types                map[reflect.Kind]string
	indent               string
	fields               string
	createFromMethodBody string
}

func (this *typeScriptClassBuilder) AddSimpleArrayField(fieldName, fieldType string, kind reflect.Kind) error {
	if typeScriptType, ok := this.types[kind]; ok {
		if len(fieldName) > 0 {
			this.fields += fmt.Sprintf("%s%s: %s[];\n", this.indent, fieldName, typeScriptType)
			this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"];\n", this.indent, this.indent, fieldName, fieldName)
			return nil
		}
	}
	return errors.New(fmt.Sprintf("Cannot find type for %s (%s/%s)", kind.String(), fieldName, fieldType))
}

func (this *typeScriptClassBuilder) AddSimpleField(fieldName, fieldType string, kind reflect.Kind) error {
	if typeScriptType, ok := this.types[kind]; ok {
		if len(fieldName) > 0 {
			this.fields += fmt.Sprintf("%s%s: %s;\n", this.indent, fieldName, typeScriptType)
			this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"];\n", this.indent, this.indent, fieldName, fieldName)
			return nil
		}
	}
	return errors.New("Cannot find type for " + fieldType)
}

func (this *typeScriptClassBuilder) AddStructField(fieldName, fieldType string) {
	this.fields += fmt.Sprintf("%s%s: %s;\n", this.indent, fieldName, fieldType)
	this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"] ? %s.createFrom(source[\"%s\"]) : null;\n", this.indent, this.indent, fieldName, fieldName, fieldType, fieldName)
}

func (this *typeScriptClassBuilder) AddArrayOfStructsField(fieldName, fieldType string) {
	this.fields += fmt.Sprintf("%s%s: %s[];\n", this.indent, fieldName, fieldType)
	this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"] ? source[\"%s\"].map(function(element) { return %s.createFrom(element); }) : null;\n", this.indent, this.indent, fieldName, fieldName, fieldName, fieldType)
}
<commit_msg>fix windows backup path<commit_after>package typescriptify

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"time"
)

type TypeScriptify struct {
	Prefix           string
	Suffix           string
	Indent           string
	CreateFromMethod bool
	BackupExtension  string // If empty no backup

	golangTypes []reflect.Type
	types       map[reflect.Kind]string

	// throwaway, used when converting
	alreadyConverted map[reflect.Type]bool
}

func New() *TypeScriptify {
	result := new(TypeScriptify)
	result.Indent = "\t"
	result.BackupExtension = "backup"

	types := make(map[reflect.Kind]string)

	types[reflect.Bool] = "boolean"

	types[reflect.Int] = "number"
	types[reflect.Int8] = "number"
	types[reflect.Int16] = "number"
	types[reflect.Int32] = "number"
	types[reflect.Int64] = "number"
	types[reflect.Uint] = "number"
	types[reflect.Uint8] = "number"
	types[reflect.Uint16] = "number"
	types[reflect.Uint32] = "number"
	types[reflect.Uint64] = "number"
	types[reflect.Float32] = "number"
	types[reflect.Float64] = "number"

	types[reflect.String] = "string"

	result.types = types

	result.Indent = "    "
	result.CreateFromMethod = true

	return result
}

func deepFields(typeOf reflect.Type) []reflect.StructField {
	fields := make([]reflect.StructField, 0)

	if typeOf.Kind() == reflect.Ptr {
		typeOf = typeOf.Elem()
	}

	if typeOf.Kind() != reflect.Struct {
		return fields
	}

	for i := 0; i < typeOf.NumField(); i++ {
		f := typeOf.Field(i)

		kind := f.Type.Kind()
		if f.Anonymous && kind == reflect.Struct {
			//fmt.Println(v.Interface())
			fields = append(fields, deepFields(f.Type)...)
		} else {
			fields = append(fields, f)
		}
	}

	return fields
}

func (this *TypeScriptify) Add(obj interface{}) {
	this.AddType(reflect.TypeOf(obj))
}

func (this *TypeScriptify) AddType(typeOf reflect.Type) {
	this.golangTypes = append(this.golangTypes, typeOf)
}

func (this *TypeScriptify) Convert(customCode map[string]string) (string, error) {
	this.alreadyConverted = make(map[reflect.Type]bool)

	result := ""
	for _, typeof := range this.golangTypes {
		typeScriptCode, err := this.convertType(typeof, customCode)
		if err != nil {
			return "", err
		}
		result += "\n" + strings.Trim(typeScriptCode, " "+this.Indent+"\r\n")
	}
	return result, nil
}

func loadCustomCode(fileName string) (map[string]string, error) {
	result := make(map[string]string)
	f, err := os.Open(fileName)
	if err != nil {
		if os.IsNotExist(err) {
			return result, nil
		}
		return result, err
	}
	defer f.Close()

	bytes, err := ioutil.ReadAll(f)
	if err != nil {
		return result, err
	}

	var currentName string
	var currentValue string
	lines := strings.Split(string(bytes), "\n")
	for _, line := range lines {
		trimmedLine := strings.TrimSpace(line)
		if strings.HasPrefix(trimmedLine, "//[") && strings.HasSuffix(trimmedLine, ":]") {
			currentName = strings.Replace(strings.Replace(trimmedLine, "//[", "", -1), ":]", "", -1)
			currentValue = ""
		} else if trimmedLine == "//[end]" {
			result[currentName] = strings.TrimRight(currentValue, " \t\r\n")
			currentName = ""
			currentValue = ""
		} else if len(currentName) > 0 {
			currentValue += line + "\n"
		}
	}

	return result, nil
}

func (this TypeScriptify) backup(fileName string) error {
	fileIn, err := os.Open(fileName)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		// No neet to backup, just return:
		return nil
	}
	defer fileIn.Close()

	bytes, err := ioutil.ReadAll(fileIn)
	if err != nil {
		return err
	}

	fileOut, err := os.Create(fmt.Sprintf("%s-%s.%s", fileName, time.Now().Format("2006-01-02T15_04_05.99"), this.BackupExtension))
	if err != nil {
		return err
	}
	defer fileOut.Close()

	_, err = fileOut.Write(bytes)
	if err != nil {
		return err
	}

	return nil
}

func (this TypeScriptify) ConvertToFile(fileName string) error {
	if len(this.BackupExtension) > 0 {
		err := this.backup(fileName)
		if err != nil {
			return err
		}
	}

	customCode, err := loadCustomCode(fileName)
	if err != nil {
		return err
	}

	f, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	converted, err := this.Convert(customCode)
	if err != nil {
		return err
	}

	f.WriteString("/* Do not change, this code is generated from Golang structs */\n\n")
	f.WriteString(converted)
	if err != nil {
		return err
	}

	return nil
}

func (this *TypeScriptify) convertType(typeOf reflect.Type, customCode map[string]string) (string, error) {
	if _, found := this.alreadyConverted[typeOf]; found { // Already converted
		return "", nil
	}

	entityName := fmt.Sprintf("%s%s%s", this.Prefix, this.Suffix, typeOf.Name())
	result := fmt.Sprintf("class %s {\n", entityName)
	builder := typeScriptClassBuilder{
		types:  this.types,
		indent: this.Indent,
	}

	fields := deepFields(typeOf)
	for _, field := range fields {
		jsonTag := field.Tag.Get("json")
		jsonFieldName := ""
		if len(jsonTag) > 0 {
			jsonTagParts := strings.Split(jsonTag, ",")
			if len(jsonTagParts) > 0 {
				jsonFieldName = strings.Trim(jsonTagParts[0], this.Indent)
			}
		}
		if len(jsonFieldName) > 0 && jsonFieldName != "-" {
			var err error
			if field.Type.Kind() == reflect.Struct { // Struct:
				typeScriptChunk, err := this.convertType(field.Type, customCode)
				if err != nil {
					return "", err
				}
				result = typeScriptChunk + "\n" + result
				builder.AddStructField(jsonFieldName, field.Type.Name())
			} else if field.Type.Kind() == reflect.Slice { // Slice:
				if field.Type.Elem().Kind() == reflect.Struct { // Slice of structs:
					typeScriptChunk, err := this.convertType(field.Type.Elem(), customCode)
					if err != nil {
						return "", err
					}
					result = typeScriptChunk + "\n" + result
					builder.AddArrayOfStructsField(jsonFieldName, field.Type.Elem().Name())
				} else { // Slice of simple fields:
					err = builder.AddSimpleArrayField(jsonFieldName, field.Type.Elem().Name(), field.Type.Elem().Kind())
				}
			} else { // Simple field:
				err = builder.AddSimpleField(jsonFieldName, field.Type.Name(), field.Type.Kind())
			}
			if err != nil {
				return "", err
			}
		}
	}

	result += builder.fields
	if this.CreateFromMethod {
		result += fmt.Sprintf("\n%sstatic createFrom(source: any) {\n", this.Indent)
		result += fmt.Sprintf("%s%svar result = new %s();\n", this.Indent, this.Indent, entityName)
		result += builder.createFromMethodBody
		result += fmt.Sprintf("%s%sreturn result;\n", this.Indent, this.Indent)
		result += fmt.Sprintf("%s}\n\n", this.Indent)
	}

	if customCode != nil {
		code := customCode[entityName]
		result += this.Indent + "//[" + entityName + ":]\n" + code + "\n\n" + this.Indent + "//[end]\n"
	}

	result += "}"

	this.alreadyConverted[typeOf] = true

	return result, nil
}

type typeScriptClassBuilder struct {
	types                map[reflect.Kind]string
	indent               string
	fields               string
	createFromMethodBody string
}

func (this *typeScriptClassBuilder) AddSimpleArrayField(fieldName, fieldType string, kind reflect.Kind) error {
	if typeScriptType, ok := this.types[kind]; ok {
		if len(fieldName) > 0 {
			this.fields += fmt.Sprintf("%s%s: %s[];\n", this.indent, fieldName, typeScriptType)
			this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"];\n", this.indent, this.indent, fieldName, fieldName)
			return nil
		}
	}
	return errors.New(fmt.Sprintf("Cannot find type for %s (%s/%s)", kind.String(), fieldName, fieldType))
}

func (this *typeScriptClassBuilder) AddSimpleField(fieldName, fieldType string, kind reflect.Kind) error {
	if typeScriptType, ok := this.types[kind]; ok {
		if len(fieldName) > 0 {
			this.fields += fmt.Sprintf("%s%s: %s;\n", this.indent, fieldName, typeScriptType)
			this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"];\n", this.indent, this.indent, fieldName, fieldName)
			return nil
		}
	}
	return errors.New("Cannot find type for " + fieldType)
}

func (this *typeScriptClassBuilder) AddStructField(fieldName, fieldType string) {
	this.fields += fmt.Sprintf("%s%s: %s;\n", this.indent, fieldName, fieldType)
	this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"] ? %s.createFrom(source[\"%s\"]) : null;\n", this.indent, this.indent, fieldName, fieldName, fieldType, fieldName)
}

func (this *typeScriptClassBuilder) AddArrayOfStructsField(fieldName, fieldType string) {
	this.fields += fmt.Sprintf("%s%s: %s[];\n", this.indent, fieldName, fieldType)
	this.createFromMethodBody += fmt.Sprintf("%s%sresult.%s = source[\"%s\"] ? source[\"%s\"].map(function(element) { return %s.createFrom(element); }) : null;\n", this.indent, this.indent, fieldName, fieldName, fieldName, fieldType)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative 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

    https://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 sender

import "net/http"

type requestInterceptor struct {
	before      func(*http.Request)
	transport   http.RoundTripper
	after       func(*http.Request, *http.Response, error)
}

func (r requestInterceptor) RoundTrip(request *http.Request) (*http.Response, error) {
	if r.before != nil {
		r.before(request)
	}
	res, err := r.transport.RoundTrip(request)
	if r.after != nil {
		r.after(request, res, err)
	}
	return res, err
}
<commit_msg>golang format tools (#2268)<commit_after>/*
Copyright 2019 The Knative 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

    https://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 sender

import "net/http"

type requestInterceptor struct {
	before    func(*http.Request)
	transport http.RoundTripper
	after     func(*http.Request, *http.Response, error)
}

func (r requestInterceptor) RoundTrip(request *http.Request) (*http.Response, error) {
	if r.before != nil {
		r.before(request)
	}
	res, err := r.transport.RoundTrip(request)
	if r.after != nil {
		r.after(request, res, err)
	}
	return res, err
}
<|endoftext|>


<commit_before>// pcsclite binding
//
// http://pcsclite.alioth.debian.org/pcsclite.html
// http://pcsclite.alioth.debian.org/api/group__API.html
//
package scard

// #cgo pkg-config: libpcsclite
// #include <stdlib.h>
// #include <winscard.h>
import "C"

import (
	"bytes"
	"unsafe"
)

// Version returns the libpcsclite version string
func Version() string {
	return C.PCSCLITE_VERSION_NUMBER
}

type Context struct {
	ctx C.SCARDCONTEXT
}

type Card struct {
	handle         C.SCARDHANDLE
	activeProtocol Protocol
}

type Protocol uint32

const (
	PROTOCOL_UNDEFINED Protocol = C.SCARD_PROTOCOL_UNDEFINED
	PROTOCOL_T0        Protocol = C.SCARD_PROTOCOL_T0
	PROTOCOL_T1        Protocol = C.SCARD_PROTOCOL_T1
	PROTOCOL_RAW       Protocol = C.SCARD_PROTOCOL_RAW
	PROTOCOL_ANY       Protocol = C.SCARD_PROTOCOL_ANY
)

type ShareMode uint32

const (
	SHARE_EXCLUSIVE ShareMode = C.SCARD_SHARE_EXCLUSIVE
	SHARE_SHARED    ShareMode = C.SCARD_SHARE_SHARED
	SHARE_DIRECT    ShareMode = C.SCARD_SHARE_DIRECT
)

type Disposition uint32

const (
	LEAVE_CARD   Disposition = C.SCARD_LEAVE_CARD
	RESET_CARD   Disposition = C.SCARD_RESET_CARD
	UNPOWER_CARD Disposition = C.SCARD_UNPOWER_CARD
	EJECT_CARD   Disposition = C.SCARD_EJECT_CARD
)

type CardStatus struct {
	Reader         string
	State          uint32
	ActiveProtocol Protocol
	ATR            []byte
}

type ReaderState struct {
	Reader       string
	UserData     interface{}
	CurrentState StateFlag
	EventState   StateFlag
	// TODO: ATR
}

type StateFlag uint32

const (
	STATE_UNAWARE     StateFlag = C.SCARD_STATE_UNAWARE
	STATE_IGNORE      StateFlag = C.SCARD_STATE_IGNORE
	STATE_CHANGED     StateFlag = C.SCARD_STATE_CHANGED
	STATE_UNKNOWN     StateFlag = C.SCARD_STATE_UNKNOWN
	STATE_UNAVAILABLE StateFlag = C.SCARD_STATE_UNAVAILABLE
	STATE_EMPTY       StateFlag = C.SCARD_STATE_EMPTY
	STATE_PRESENT     StateFlag = C.SCARD_STATE_PRESENT
	STATE_ATRMATCH    StateFlag = C.SCARD_STATE_ATRMATCH
	STATE_EXCLUSIVE   StateFlag = C.SCARD_STATE_EXCLUSIVE
	STATE_INUSE       StateFlag = C.SCARD_STATE_INUSE
	STATE_MUTE        StateFlag = C.SCARD_STATE_MUTE
	STATE_UNPOWERED   StateFlag = C.SCARD_STATE_UNPOWERED
)

type Timeout uint32

const (
	INFINITE Timeout = C.INFINITE
)

// wraps SCardEstablishContext
func EstablishContext() (*Context, error) {
	var ctx Context

	r := C.SCardEstablishContext(C.SCARD_SCOPE_SYSTEM, nil, nil, &ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	return &ctx, nil
}

// wraps SCardIsValidContext
func (ctx *Context) IsValid() (bool, error) {
	r := C.SCardIsValidContext(ctx.ctx)
	switch r {
	case C.SCARD_S_SUCCESS:
		return true, nil
	case C.SCARD_E_INVALID_HANDLE:
		return false, nil
	default:
		return false, newError(r)
	}
	panic("unreachable")
}

// wraps SCardCancel
func (ctx *Context) Cancel() error {
	r := C.SCardCancel(ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardReleaseContext
func (ctx *Context) Release() error {
	r := C.SCardReleaseContext(ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardListReaders
func (ctx *Context) ListReaders() ([]string, error) {
	var needed C.DWORD

	r := C.SCardListReaders(ctx.ctx, nil, nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	data := make([]byte, needed)
	cdata := (*C.char)(unsafe.Pointer(&data[0]))

	r = C.SCardListReaders(ctx.ctx, nil, cdata, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var readers []string
	for _, b := range bytes.Split(data, []byte{0}) {
		if len(b) > 0 {
			readers = append(readers, string(b))
		}
	}

	return readers, nil
}

// wraps SCardListReaderGroups
func (ctx *Context) ListReaderGroups() ([]string, error) {
	var needed C.DWORD

	r := C.SCardListReaderGroups(ctx.ctx, nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	data := make([]byte, needed)
	cdata := (*C.char)(unsafe.Pointer(&data[0]))

	r = C.SCardListReaderGroups(ctx.ctx, cdata, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var groups []string
	for _, b := range bytes.Split(data, []byte{0}) {
		if len(b) > 0 {
			groups = append(groups, string(b))
		}
	}

	return groups, nil
}

// wraps SCardGetStatusChange
func (ctx *Context) GetStatusChange(readerStates []ReaderState, timeout Timeout) error {

	crs := make([]C.SCARD_READERSTATE, len(readerStates))

	for i := range readerStates {
		crs[i].szReader = C.CString(readerStates[i].Reader)
		defer C.free(unsafe.Pointer(crs[i].szReader))
		crs[i].dwCurrentState = C.DWORD(readerStates[i].CurrentState)
	}

	r := C.SCardGetStatusChange(ctx.ctx, C.DWORD(timeout),
		(C.LPSCARD_READERSTATE)(unsafe.Pointer(&crs[0])),
		C.DWORD(len(crs)))

	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}

	for i := range readerStates {
		readerStates[i].EventState = StateFlag(crs[i].dwEventState)
	}

	return nil
}

// wraps SCardConnect
func (ctx *Context) Connect(reader string, mode ShareMode, proto Protocol) (*Card, error) {
	var card Card
	var activeProtocol C.DWORD

	creader := C.CString(reader)
	defer C.free(unsafe.Pointer(creader))

	r := C.SCardConnect(ctx.ctx, creader, C.DWORD(mode), C.DWORD(proto), &card.handle, &activeProtocol)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	card.activeProtocol = Protocol(activeProtocol)
	return &card, nil
}

// wraps SCardDisconnect
func (card *Card) Disconnect(d Disposition) error {
	r := C.SCardDisconnect(card.handle, C.DWORD(d))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardReconnect
func (card *Card) Reconnect(mode ShareMode, protocol Protocol, init Disposition) error {
	var activeProtocol C.DWORD

	r := C.SCardReconnect(card.handle, C.DWORD(mode), C.DWORD(protocol), C.DWORD(init), &activeProtocol)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}

	card.activeProtocol = Protocol(activeProtocol)

	return nil
}

// wraps SCardBeginTransaction
func (card *Card) BeginTransaction() error {
	r := C.SCardBeginTransaction(card.handle)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardEndTransaction
func (card *Card) EndTransaction(d Disposition) error {
	r := C.SCardEndTransaction(card.handle, C.DWORD(d))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardStatus
func (card *Card) Status() (*CardStatus, error) {
	var reader [C.MAX_READERNAME + 1]byte
	var readerLen = C.DWORD(len(reader))
	var state, proto C.DWORD
	var atr [C.MAX_ATR_SIZE]byte
	var atrLen = C.DWORD(len(atr))

	r := C.SCardStatus(card.handle, (C.LPSTR)(unsafe.Pointer(&reader[0])), &readerLen, &state, &proto, (*C.BYTE)(&atr[0]), &atrLen)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	status := &CardStatus{
		Reader:         string(reader[0:readerLen]),
		State:          uint32(state),
		ActiveProtocol: Protocol(proto),
		ATR:            atr[0:atrLen],
	}

	return status, nil
}

// wraps SCardTransmit
func (card *Card) Transmit(cmd []byte) ([]byte, error) {
	var sendpci *C.SCARD_IO_REQUEST
	var recvpci C.SCARD_IO_REQUEST

	switch card.activeProtocol {
	case PROTOCOL_T0:
		sendpci = &C.g_rgSCardT0Pci
	case PROTOCOL_T1:
		sendpci = &C.g_rgSCardT1Pci
	case PROTOCOL_RAW:
		sendpci = &C.g_rgSCardRawPci
	default:
		panic("unknown protocol")
	}

	var recv [C.MAX_BUFFER_SIZE_EXTENDED]byte
	var recvlen C.DWORD = C.DWORD(len(recv))

	r := C.SCardTransmit(card.handle, sendpci, (*C.BYTE)(&cmd[0]), C.DWORD(len(cmd)), &recvpci, (*C.BYTE)(&recv[0]), &recvlen)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	rsp := make([]byte, recvlen)
	copy(rsp, recv[0:recvlen])

	return rsp, nil
}

// wraps SCardControl
func (card *Card) Control(ctrl uint32, cmd []byte) ([]byte, error) {
	var recv [C.MAX_BUFFER_SIZE_EXTENDED]byte
	var recvlen C.DWORD

	r := C.SCardControl(card.handle, C.DWORD(ctrl),
		(C.LPCVOID)(unsafe.Pointer(&cmd[0])), C.DWORD(len(cmd)),
		(C.LPVOID)(unsafe.Pointer(&recv[0])), C.DWORD(len(recv)), &recvlen)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	rsp := make([]byte, recvlen)
	copy(rsp, recv[0:recvlen])

	return rsp, nil
}

// wraps SCardGetAttrib
func (card *Card) GetAttrib(id uint32) ([]byte, error) {
	var needed C.DWORD

	r := C.SCardGetAttrib(card.handle, C.DWORD(id), nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var attrib = make([]byte, needed)

	r = C.SCardGetAttrib(card.handle, C.DWORD(id), (*C.BYTE)(&attrib[0]), &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	return attrib[0:needed], nil
}

// wraps SCardSetAttrib
func (card *Card) SetAttrib(id uint32, data []byte) error {
	r := C.SCardSetAttrib(card.handle, C.DWORD(id), (*C.BYTE)(&data[0]), C.DWORD(len(data)))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// SCardFreeMemory is not needed. We (hopefuly) never return buffers allocated by libpcsclite
<commit_msg>fix Card.Control()<commit_after>// pcsclite binding
//
// http://pcsclite.alioth.debian.org/pcsclite.html
// http://pcsclite.alioth.debian.org/api/group__API.html
//
package scard

// #cgo pkg-config: libpcsclite
// #include <stdlib.h>
// #include <winscard.h>
import "C"

import (
	"bytes"
	"unsafe"
)

// Version returns the libpcsclite version string
func Version() string {
	return C.PCSCLITE_VERSION_NUMBER
}

type Context struct {
	ctx C.SCARDCONTEXT
}

type Card struct {
	handle         C.SCARDHANDLE
	activeProtocol Protocol
}

type Protocol uint32

const (
	PROTOCOL_UNDEFINED Protocol = C.SCARD_PROTOCOL_UNDEFINED
	PROTOCOL_T0        Protocol = C.SCARD_PROTOCOL_T0
	PROTOCOL_T1        Protocol = C.SCARD_PROTOCOL_T1
	PROTOCOL_RAW       Protocol = C.SCARD_PROTOCOL_RAW
	PROTOCOL_ANY       Protocol = C.SCARD_PROTOCOL_ANY
)

type ShareMode uint32

const (
	SHARE_EXCLUSIVE ShareMode = C.SCARD_SHARE_EXCLUSIVE
	SHARE_SHARED    ShareMode = C.SCARD_SHARE_SHARED
	SHARE_DIRECT    ShareMode = C.SCARD_SHARE_DIRECT
)

type Disposition uint32

const (
	LEAVE_CARD   Disposition = C.SCARD_LEAVE_CARD
	RESET_CARD   Disposition = C.SCARD_RESET_CARD
	UNPOWER_CARD Disposition = C.SCARD_UNPOWER_CARD
	EJECT_CARD   Disposition = C.SCARD_EJECT_CARD
)

type CardStatus struct {
	Reader         string
	State          uint32
	ActiveProtocol Protocol
	ATR            []byte
}

type ReaderState struct {
	Reader       string
	UserData     interface{}
	CurrentState StateFlag
	EventState   StateFlag
	// TODO: ATR
}

type StateFlag uint32

const (
	STATE_UNAWARE     StateFlag = C.SCARD_STATE_UNAWARE
	STATE_IGNORE      StateFlag = C.SCARD_STATE_IGNORE
	STATE_CHANGED     StateFlag = C.SCARD_STATE_CHANGED
	STATE_UNKNOWN     StateFlag = C.SCARD_STATE_UNKNOWN
	STATE_UNAVAILABLE StateFlag = C.SCARD_STATE_UNAVAILABLE
	STATE_EMPTY       StateFlag = C.SCARD_STATE_EMPTY
	STATE_PRESENT     StateFlag = C.SCARD_STATE_PRESENT
	STATE_ATRMATCH    StateFlag = C.SCARD_STATE_ATRMATCH
	STATE_EXCLUSIVE   StateFlag = C.SCARD_STATE_EXCLUSIVE
	STATE_INUSE       StateFlag = C.SCARD_STATE_INUSE
	STATE_MUTE        StateFlag = C.SCARD_STATE_MUTE
	STATE_UNPOWERED   StateFlag = C.SCARD_STATE_UNPOWERED
)

type Timeout uint32

const (
	INFINITE Timeout = C.INFINITE
)

// wraps SCardEstablishContext
func EstablishContext() (*Context, error) {
	var ctx Context

	r := C.SCardEstablishContext(C.SCARD_SCOPE_SYSTEM, nil, nil, &ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	return &ctx, nil
}

// wraps SCardIsValidContext
func (ctx *Context) IsValid() (bool, error) {
	r := C.SCardIsValidContext(ctx.ctx)
	switch r {
	case C.SCARD_S_SUCCESS:
		return true, nil
	case C.SCARD_E_INVALID_HANDLE:
		return false, nil
	default:
		return false, newError(r)
	}
	panic("unreachable")
}

// wraps SCardCancel
func (ctx *Context) Cancel() error {
	r := C.SCardCancel(ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardReleaseContext
func (ctx *Context) Release() error {
	r := C.SCardReleaseContext(ctx.ctx)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardListReaders
func (ctx *Context) ListReaders() ([]string, error) {
	var needed C.DWORD

	r := C.SCardListReaders(ctx.ctx, nil, nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	data := make([]byte, needed)
	cdata := (*C.char)(unsafe.Pointer(&data[0]))

	r = C.SCardListReaders(ctx.ctx, nil, cdata, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var readers []string
	for _, b := range bytes.Split(data, []byte{0}) {
		if len(b) > 0 {
			readers = append(readers, string(b))
		}
	}

	return readers, nil
}

// wraps SCardListReaderGroups
func (ctx *Context) ListReaderGroups() ([]string, error) {
	var needed C.DWORD

	r := C.SCardListReaderGroups(ctx.ctx, nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	data := make([]byte, needed)
	cdata := (*C.char)(unsafe.Pointer(&data[0]))

	r = C.SCardListReaderGroups(ctx.ctx, cdata, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var groups []string
	for _, b := range bytes.Split(data, []byte{0}) {
		if len(b) > 0 {
			groups = append(groups, string(b))
		}
	}

	return groups, nil
}

// wraps SCardGetStatusChange
func (ctx *Context) GetStatusChange(readerStates []ReaderState, timeout Timeout) error {

	crs := make([]C.SCARD_READERSTATE, len(readerStates))

	for i := range readerStates {
		crs[i].szReader = C.CString(readerStates[i].Reader)
		defer C.free(unsafe.Pointer(crs[i].szReader))
		crs[i].dwCurrentState = C.DWORD(readerStates[i].CurrentState)
	}

	r := C.SCardGetStatusChange(ctx.ctx, C.DWORD(timeout),
		(C.LPSCARD_READERSTATE)(unsafe.Pointer(&crs[0])),
		C.DWORD(len(crs)))

	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}

	for i := range readerStates {
		readerStates[i].EventState = StateFlag(crs[i].dwEventState)
	}

	return nil
}

// wraps SCardConnect
func (ctx *Context) Connect(reader string, mode ShareMode, proto Protocol) (*Card, error) {
	var card Card
	var activeProtocol C.DWORD

	creader := C.CString(reader)
	defer C.free(unsafe.Pointer(creader))

	r := C.SCardConnect(ctx.ctx, creader, C.DWORD(mode), C.DWORD(proto), &card.handle, &activeProtocol)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	card.activeProtocol = Protocol(activeProtocol)
	return &card, nil
}

// wraps SCardDisconnect
func (card *Card) Disconnect(d Disposition) error {
	r := C.SCardDisconnect(card.handle, C.DWORD(d))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardReconnect
func (card *Card) Reconnect(mode ShareMode, protocol Protocol, init Disposition) error {
	var activeProtocol C.DWORD

	r := C.SCardReconnect(card.handle, C.DWORD(mode), C.DWORD(protocol), C.DWORD(init), &activeProtocol)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}

	card.activeProtocol = Protocol(activeProtocol)

	return nil
}

// wraps SCardBeginTransaction
func (card *Card) BeginTransaction() error {
	r := C.SCardBeginTransaction(card.handle)
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardEndTransaction
func (card *Card) EndTransaction(d Disposition) error {
	r := C.SCardEndTransaction(card.handle, C.DWORD(d))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// wraps SCardStatus
func (card *Card) Status() (*CardStatus, error) {
	var reader [C.MAX_READERNAME + 1]byte
	var readerLen = C.DWORD(len(reader))
	var state, proto C.DWORD
	var atr [C.MAX_ATR_SIZE]byte
	var atrLen = C.DWORD(len(atr))

	r := C.SCardStatus(card.handle, (C.LPSTR)(unsafe.Pointer(&reader[0])), &readerLen, &state, &proto, (*C.BYTE)(&atr[0]), &atrLen)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	status := &CardStatus{
		Reader:         string(reader[0:readerLen]),
		State:          uint32(state),
		ActiveProtocol: Protocol(proto),
		ATR:            atr[0:atrLen],
	}

	return status, nil
}

// wraps SCardTransmit
func (card *Card) Transmit(cmd []byte) ([]byte, error) {
	var sendpci *C.SCARD_IO_REQUEST
	var recvpci C.SCARD_IO_REQUEST

	switch card.activeProtocol {
	case PROTOCOL_T0:
		sendpci = &C.g_rgSCardT0Pci
	case PROTOCOL_T1:
		sendpci = &C.g_rgSCardT1Pci
	case PROTOCOL_RAW:
		sendpci = &C.g_rgSCardRawPci
	default:
		panic("unknown protocol")
	}

	var recv [C.MAX_BUFFER_SIZE_EXTENDED]byte
	var recvlen C.DWORD = C.DWORD(len(recv))

	r := C.SCardTransmit(card.handle, sendpci, (*C.BYTE)(&cmd[0]), C.DWORD(len(cmd)), &recvpci, (*C.BYTE)(&recv[0]), &recvlen)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	rsp := make([]byte, recvlen)
	copy(rsp, recv[0:recvlen])

	return rsp, nil
}

// wraps SCardControl
func (card *Card) Control(ctrl uint32, cmd []byte) ([]byte, error) {
	var recv [C.MAX_BUFFER_SIZE_EXTENDED]byte
	var recvlen C.DWORD
	var r C.LONG

	if len(cmd) == 0 {
		r = C.SCardControl(card.handle, C.DWORD(ctrl),
			(C.LPCVOID)(nil), 0,
			(C.LPVOID)(unsafe.Pointer(&recv[0])), C.DWORD(len(recv)), &recvlen)
	} else {
		r = C.SCardControl(card.handle, C.DWORD(ctrl),
			(C.LPCVOID)(unsafe.Pointer(&cmd[0])), C.DWORD(len(cmd)),
			(C.LPVOID)(unsafe.Pointer(&recv[0])), C.DWORD(len(recv)), &recvlen)
	}
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	rsp := make([]byte, recvlen)
	copy(rsp, recv[0:recvlen])

	return rsp, nil
}

// wraps SCardGetAttrib
func (card *Card) GetAttrib(id uint32) ([]byte, error) {
	var needed C.DWORD

	r := C.SCardGetAttrib(card.handle, C.DWORD(id), nil, &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	var attrib = make([]byte, needed)

	r = C.SCardGetAttrib(card.handle, C.DWORD(id), (*C.BYTE)(&attrib[0]), &needed)
	if r != C.SCARD_S_SUCCESS {
		return nil, newError(r)
	}

	return attrib[0:needed], nil
}

// wraps SCardSetAttrib
func (card *Card) SetAttrib(id uint32, data []byte) error {
	r := C.SCardSetAttrib(card.handle, C.DWORD(id), (*C.BYTE)(&data[0]), C.DWORD(len(data)))
	if r != C.SCARD_S_SUCCESS {
		return newError(r)
	}
	return nil
}

// SCardFreeMemory is not needed. We (hopefuly) never return buffers allocated by libpcsclite
<|endoftext|>


<commit_before>// sokoban solver, work in progress
package main

import (
	"fmt"
	"github.com/bertbaron/solve"
	"math"
	"sort"
	"strings"
)

const (
	floor  byte = 0
	wall   byte = 1
	box    byte = 2
	goal   byte = 4
	player byte = 8
)

var chars = map[rune]byte{
	' ': floor,
	'#': wall,
	'$': box,
	'.': goal,
	'@': player,
	'+': player | goal,
	'*': box | goal}

var reverse = map[byte]rune{
	floor:         ' ',
	wall:          '#',
	box:           '$',
	goal:          '.',
	player:        '@',
	player | goal: '+',
	box | goal:    '*'}

// -------- main problem. We only expose the states in which a block is pushed though to limit the search space
//          for the main search.
type sokoban struct {
	// the static world, without player and boxes
	world []byte
	// sorted list of goal positions
	goals  []uint16
	width  int
	height int
}

type mainstate struct {
	// sorted list of box positions
	boxes    []uint16
	position int
	cost     int
}

// returns the index of position in the sorted list of positions. Returns -1 if the position is not found
func binarySearch(positions []uint16, position int) int {
	idx := sort.Search(len(positions), func(i int) bool { return positions[i] >= uint16(position) })
	if idx < len(positions) && positions[idx] == uint16(position) {
		return idx
	}
	return -1
}

func valueOf(s sokoban, m mainstate, position int) byte {
	boxidx := binarySearch(m.boxes, position)
	var additional byte = 0
	if m.position == position {
		additional |= player
	}
	if boxidx >= 0 {
		additional |= box
	}
	return s.world[position] | additional
}

func isEmpty(value byte) bool {
	return value&(wall|box) == 0
}

func isBox(value byte) bool {
	return value&box != 0
}

func isWall(value byte) bool {
	return value&wall != 0
}

func print(s sokoban, m mainstate) {
	for position := range s.world {
		fmt.Print(string(reverse[valueOf(s, m, position)]))
		if position%s.width == s.width-1 {
			fmt.Println()
		}
	}
}

func (s mainstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func abs(value int) int {
	if value < 0 {
		return -value
	}
	return value
}

func (s mainstate) Heuristic(ctx solve.Context) float64 {
	world := ctx.Custom.(sokoban)
	total := 0
	for _, box := range s.boxes {
		min := math.MaxInt32
		bx, by := int(box)%world.width, int(box)/world.width
		for _, goal := range world.goals {
			gx, gy := int(goal)%world.width, int(goal)/world.width
			md := abs(gx-bx) + abs(gy-by)
			if md < min {
				min = md
			}
		}
		total += min
	}
	return float64(total)
}

func (s mainstate) IsGoal(ctx solve.Context) bool {
	for i, value := range ctx.Custom.(sokoban).goals {
		if s.boxes[i] != value {
			return false
		}
	}
	return true
}

func (s mainstate) Expand(ctx solve.Context) []solve.State {
	world := ctx.Custom.(sokoban)
	targets := make([]int, 0)
	for _, box := range s.boxes {
		left := isEmpty(valueOf(world, s, int(box)-1))
		right := isEmpty(valueOf(world, s, int(box)+1))
		up := isEmpty(valueOf(world, s, int(box)-world.width))
		down := isEmpty(valueOf(world, s, int(box)+world.width))
		if left && right {
			targets = append(targets, int(box)-1)
			targets = append(targets, int(box)+1)
		}
		if up && down {
			targets = append(targets, int(box)-world.width)
			targets = append(targets, int(box)+world.width)
		}
	}
	paths := getWalkMoves(world, s, targets)

	var children []solve.State
	for _, path := range paths {
		p := path.position
		for _, dir := range [...]int{-1, 1, -world.width, world.width} {
			if isBox(valueOf(world, s, p+dir)) && isEmpty(valueOf(world, s, p+2*dir)) {
				child := push(world, s, p, dir, path.cost)
				if child != nil {
					//print(world, child)
					children = append(children, *child)
				}
			}
		}
	}
	return children
}

func push(world sokoban, s mainstate, position int, direction int, cost int) *mainstate {
	newposition := position + direction
	newbox := uint16(position + direction*2)
	newboxes := make([]uint16, len(s.boxes))
	copy(newboxes, s.boxes)
	idx := binarySearch(newboxes, newposition)
	newboxes[idx] = newbox

	n := len(newboxes)
	// insertion sort to keep boxes sorted, only needed when moving up or down
	if direction < -1 {
		for idx > 0 && newboxes[idx-1] > newbox {
			newboxes[idx-1], newboxes[idx] = newboxes[idx], newboxes[idx-1]
			idx--
		}
	}
	if direction > 1 {
		for idx < n-1 && newboxes[idx+1] < newbox {
			newboxes[idx+1], newboxes[idx] = newboxes[idx], newboxes[idx+1]
			idx++
		}
	}
	newState := mainstate{newboxes, newposition, s.cost + cost + 1}
	if deadEnd(world, newState, int(newbox)) {
		return nil
	}
	return &newState
}

// looks in a 3x3 pattern around the box position if this is a dead end
func deadEnd(world sokoban, s mainstate, position int) bool {
	if world.world[position]&goal != 0 {
		return false // box is on a goal position
	}

	// corner walls
	lu := world.world[position-1-world.width]&wall != 0
	ru := world.world[position+1-world.width]&wall != 0
	ld := world.world[position-1+world.width]&wall != 0
	rd := world.world[position+1+world.width]&wall != 0

	// orthogonal walls or boxes
	uvalue := valueOf(world, s, position-world.width)
	dvalue := valueOf(world, s, position+world.width)
	lvalue := valueOf(world, s, position-1)
	rvalue := valueOf(world, s, position+1)

	// direction is blocked if it is a wall or a block that is sideways blocked by a wall
	u := isWall(uvalue) || (isBox(uvalue) && (lu || ru))
	d := isWall(dvalue) || (isBox(dvalue) && (ld || rd))
	l := isWall(lvalue) || (isBox(lvalue) && (lu || ld))
	r := isWall(rvalue) || (isBox(rvalue) && (ru || rd))

	return u && r || r && d || d && l || l && u
}

func assertOrdered(positions []uint16) {
	for i, value := range positions[1:] {
		if value <= positions[i] {
			panic(fmt.Sprintf("Ordering invariant violated: %v", positions))
		}
	}
}

// -------------- Sub problem for moving the player to all positions in which a box can be moved -----------

type walkcontext struct {
	// the static world, without player but with boxes because we don't move them here
	world   []byte
	targets []int
	width   int
}

type walkstate struct {
	position int
	cost     int
}

func (s walkstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func (s walkstate) Heuristic(ctx solve.Context) float64 {
	return 0
}

func (s walkstate) IsGoal(ctx solve.Context) bool {
	wc := ctx.Custom.(walkcontext)
	for _, goal := range wc.targets {
		if s.position == goal {
			return true
		}
	}
	return false
}

func (s walkstate) Expand(ctx solve.Context) []solve.State {
	var children []solve.State
	wc := ctx.Custom.(walkcontext)
	children = s.addIfValid(children, s.position-1, wc)
	children = s.addIfValid(children, s.position+1, wc)
	children = s.addIfValid(children, s.position-wc.width, wc)
	children = s.addIfValid(children, s.position+wc.width, wc)
	return children
}

func (s walkstate) addIfValid(children []solve.State, newPosition int, wc walkcontext) []solve.State {
	if wc.world[newPosition]&(wall|box) == 0 {
		return append(children, walkstate{newPosition, s.cost + 1})
	}
	return children
}

type walkstateMap map[int]float64

func (c walkstateMap) Get(state solve.State) (float64, bool) {
	value, ok := c[state.(walkstate).position]
	return value, ok
}

func (c walkstateMap) Put(state solve.State, value float64) {
	c[state.(walkstate).position] = value
}

func (c *walkstateMap) Clear() {
	*c = make(walkstateMap)
}

func getWalkMoves(wc sokoban, s mainstate, targets []int) []walkstate {
	context := walkcontext{wc.world, targets, wc.width}
	context.world = make([]byte, len(wc.world))
	copy(context.world, wc.world)
	for _, boxposition := range s.boxes {
		context.world[boxposition] |= box
	}
	rootstate := walkstate{s.position, 0}
	var wsMap walkstateMap
	solver := solve.NewSolver(rootstate).
		Context(context).
		Constraint(solve.CheapestPathConstraint(&wsMap)).
		Algorithm(solve.BreadthFirst)
	solutions := make([]walkstate, 0)
	for solution := range solver.SolveAll() {
		solutions = append(solutions, solution.GoalState().(walkstate))
	}
	return solutions
}

func parse(level string) (sokoban, mainstate) {
	width := 0
	lines := strings.Split(level, "\n")
	height := len(lines)
	for _, line := range lines {
		if len(line) > width {
			width = len(line)
		}
	}
	var c sokoban
	var s mainstate
	c.width = width
	c.height = height

	c.world = make([]byte, width*height)
	c.goals = make([]uint16, 0)
	s.boxes = make([]uint16, 0)
	for y, row := range lines {
		for x, raw := range row {
			position := y*width + x
			if value, ok := chars[raw]; ok {
				c.world[position] = value &^ player &^ box
				if value&player != 0 {
					s.position = position
				}
				if value&goal != 0 {
					c.goals = append(c.goals, uint16(position))
				}
				if value&box != 0 {
					s.boxes = append(s.boxes, uint16(position))
				}
			} else {
				panic(fmt.Sprintf("Invalid level format, character %v is not valid", value))
			}
		}
	}
	return c, s
}

// For cheapest path constraint
type cpMap map[string]float64

func key(state solve.State) string {
	s := state.(mainstate)
	return fmt.Sprintf("%v:%v", s.position, s.boxes)
}

func (c cpMap) Get(state solve.State) (value float64, ok bool) {
	value, ok = c[key(state)]
	return
}

func (c cpMap) Put(state solve.State, value float64) {
	c[key(state)] = value
}

func (c *cpMap) Clear() {
	*c = make(cpMap)
}

func cheapestPathConstraint() solve.Constraint {
	var m cpMap
	return solve.CheapestPathConstraint(&m)
}

var simpleLevel = `
########
#     @#
#      #
# $ # .#
#  # # #
########`

var level = `
   ####
####  ##
#   $  #
#  *** #
#  . . ##
## * *  #
 ##***  #
  # $ ###
  # @ #
  #####`

func main() {
	world, root := parse(level)
	print(world, root)
	result := solve.NewSolver(root).
		Context(world).
		Algorithm(solve.IDAstar).
		Constraint(cheapestPathConstraint()).
		Solve()
	fmt.Printf("Result:\n ")
	for _, state := range result.Solution {
		print(world, state.(mainstate))
	}
	if result.Solved() {
		fmt.Printf("Solved in %d moves\n", int(result.GoalState().(mainstate).cost))
	}
	fmt.Printf("visited %v main nodes\n", result.Visited)
}
<commit_msg>Working on the sSokoban example<commit_after>// sokoban solver, work in progress
package main

import (
	"fmt"
	"github.com/bertbaron/solve"
	"math"
	"sort"
	"strings"
	"os"
	"log"
	"runtime/pprof"
	"time"
)

const (
	floor  byte = 0
	wall   byte = 1
	box    byte = 2
	goal   byte = 4
	player byte = 8
)

var chars = map[rune]byte{
	' ': floor,
	'#': wall,
	'$': box,
	'.': goal,
	'@': player,
	'+': player | goal,
	'*': box | goal}

var reverse = map[byte]rune{
	floor:         ' ',
	wall:          '#',
	box:           '$',
	goal:          '.',
	player:        '@',
	player | goal: '+',
	box | goal:    '*'}

// -------- main problem. We only expose the states in which a block is pushed though to limit the search space
//          for the main search.
type sokoban struct {
	// the static world, without player and boxes
	world []byte
	// sorted list of goal positions
	goals  []uint16
	width  int
	height int
}

type mainstate struct {
	// sorted list of box positions
	boxes    []uint16
	position int
	cost     int
}

// returns the index of position in the sorted list of positions. Returns -1 if the position is not found
func binarySearch(positions []uint16, position int) int {
	idx := sort.Search(len(positions), func(i int) bool { return positions[i] >= uint16(position) })
	if idx < len(positions) && positions[idx] == uint16(position) {
		return idx
	}
	return -1
}

func valueOf(s sokoban, m mainstate, position int) byte {
	boxidx := binarySearch(m.boxes, position)
	var additional byte = 0
	if m.position == position {
		additional |= player
	}
	if boxidx >= 0 {
		additional |= box
	}
	return s.world[position] | additional
}

func isEmpty(value byte) bool {
	return value&(wall|box) == 0
}

func isBox(value byte) bool {
	return value&box != 0
}

func isWall(value byte) bool {
	return value&wall != 0
}

func print(s sokoban, m mainstate) {
	for position := range s.world {
		fmt.Print(string(reverse[valueOf(s, m, position)]))
		if position%s.width == s.width-1 {
			fmt.Println()
		}
	}
}

func (s mainstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func abs(value int) int {
	if value < 0 {
		return -value
	}
	return value
}

// total manhattan distance for all boxes towards their nearest target - O(boxcount^2)
func minimalManhattan(world sokoban, s mainstate) int {
	total := 0
	for _, box := range s.boxes {
		min := math.MaxInt32
		bx, by := int(box)%world.width, int(box)/world.width
		for _, goal := range world.goals {
			gx, gy := int(goal)%world.width, int(goal)/world.width
			md := abs(gx-bx) + abs(gy-by)
			if md < min {
				min = md
			}
		}
		total += min
	}
	return total
}

// number of boxes not on a goal position - O(boxcount)
func displaced(world sokoban, s mainstate) int {
	displaced := len(s.boxes)
	for _, box := range s.boxes {
		displaced -= int(world.world[box] & goal) >> 2
	}
	return displaced
}

func (s mainstate) Heuristic(ctx solve.Context) float64 {
	world := ctx.Custom.(sokoban)
	//h := displaced(world, s)
	h := minimalManhattan(world, s)
	return float64(h)
}

func (s mainstate) IsGoal(ctx solve.Context) bool {
	for i, value := range ctx.Custom.(sokoban).goals {
		if s.boxes[i] != value {
			return false
		}
	}
	return true
}

func (s mainstate) Expand(ctx solve.Context) []solve.State {
	world := ctx.Custom.(sokoban)
	targets := make([]int, 0)
	for _, box := range s.boxes {
		left := isEmpty(valueOf(world, s, int(box)-1))
		right := isEmpty(valueOf(world, s, int(box)+1))
		up := isEmpty(valueOf(world, s, int(box)-world.width))
		down := isEmpty(valueOf(world, s, int(box)+world.width))
		if left && right {
			targets = append(targets, int(box)-1)
			targets = append(targets, int(box)+1)
		}
		if up && down {
			targets = append(targets, int(box)-world.width)
			targets = append(targets, int(box)+world.width)
		}
	}
	paths := getWalkMoves(world, s, targets)

	var children []solve.State
	for _, path := range paths {
		p := path.position
		for _, dir := range [...]int{-1, 1, -world.width, world.width} {
			if isBox(valueOf(world, s, p+dir)) && isEmpty(valueOf(world, s, p+2*dir)) {
				child := push(world, s, p, dir, path.cost)
				if child != nil {
					//print(world, child)
					children = append(children, *child)
				}
			}
		}
	}
	return children
}

func push(world sokoban, s mainstate, position int, direction int, cost int) *mainstate {
	newposition := position + direction
	newbox := uint16(position + direction*2)
	newboxes := make([]uint16, len(s.boxes))
	copy(newboxes, s.boxes)
	idx := binarySearch(newboxes, newposition)
	newboxes[idx] = newbox

	n := len(newboxes)
	// insertion sort to keep boxes sorted, only needed when moving up or down
	if direction < -1 {
		for idx > 0 && newboxes[idx-1] > newbox {
			newboxes[idx-1], newboxes[idx] = newboxes[idx], newboxes[idx-1]
			idx--
		}
	}
	if direction > 1 {
		for idx < n-1 && newboxes[idx+1] < newbox {
			newboxes[idx+1], newboxes[idx] = newboxes[idx], newboxes[idx+1]
			idx++
		}
	}
	newState := mainstate{newboxes, newposition, s.cost + cost + 1}
	if deadEnd(world, newState, int(newbox)) {
		return nil
	}
	return &newState
}

// looks in a 3x3 pattern around the box position if this is a dead end
func deadEnd(world sokoban, s mainstate, position int) bool {
	if world.world[position]&goal != 0 {
		return false // box is on a goal position
	}

	// corner walls
	lu := world.world[position-1-world.width]&wall != 0
	ru := world.world[position+1-world.width]&wall != 0
	ld := world.world[position-1+world.width]&wall != 0
	rd := world.world[position+1+world.width]&wall != 0

	// orthogonal walls or boxes
	uvalue := valueOf(world, s, position-world.width)
	dvalue := valueOf(world, s, position+world.width)
	lvalue := valueOf(world, s, position-1)
	rvalue := valueOf(world, s, position+1)

	// direction is blocked if it is a wall or a block that is sideways blocked by a wall
	u := isWall(uvalue) || (isBox(uvalue) && (lu || ru))
	d := isWall(dvalue) || (isBox(dvalue) && (ld || rd))
	l := isWall(lvalue) || (isBox(lvalue) && (lu || ld))
	r := isWall(rvalue) || (isBox(rvalue) && (ru || rd))

	return u && r || r && d || d && l || l && u
}

func assertOrdered(positions []uint16) {
	for i, value := range positions[1:] {
		if value <= positions[i] {
			panic(fmt.Sprintf("Ordering invariant violated: %v", positions))
		}
	}
}

// -------------- Sub problem for moving the player to all positions in which a box can be moved -----------

type walkcontext struct {
	// the static world, without player but with boxes because we don't move them here
	world   []byte
	targets []int
	width   int
}

type walkstate struct {
	position int
	cost     int
}

func (s walkstate) Cost(ctx solve.Context) float64 {
	return float64(s.cost)
}

func (s walkstate) Heuristic(ctx solve.Context) float64 {
	return 0
}

func (s walkstate) IsGoal(ctx solve.Context) bool {
	wc := ctx.Custom.(walkcontext)
	for _, goal := range wc.targets {
		if s.position == goal {
			return true
		}
	}
	return false
}

func (s walkstate) Expand(ctx solve.Context) []solve.State {
	var children []solve.State
	wc := ctx.Custom.(walkcontext)
	children = s.addIfValid(children, s.position-1, wc)
	children = s.addIfValid(children, s.position+1, wc)
	children = s.addIfValid(children, s.position-wc.width, wc)
	children = s.addIfValid(children, s.position+wc.width, wc)
	return children
}

func (s walkstate) addIfValid(children []solve.State, newPosition int, wc walkcontext) []solve.State {
	if wc.world[newPosition]&(wall|box) == 0 {
		return append(children, walkstate{newPosition, s.cost + 1})
	}
	return children
}

type walkstateMap map[int]float64

func (c walkstateMap) Get(state solve.State) (float64, bool) {
	value, ok := c[state.(walkstate).position]
	return value, ok
}

func (c walkstateMap) Put(state solve.State, value float64) {
	c[state.(walkstate).position] = value
}

func (c *walkstateMap) Clear() {
	*c = make(walkstateMap)
}

func getWalkMoves(wc sokoban, s mainstate, targets []int) []walkstate {
	context := walkcontext{wc.world, targets, wc.width}
	context.world = make([]byte, len(wc.world))
	copy(context.world, wc.world)
	for _, boxposition := range s.boxes {
		context.world[boxposition] |= box
	}
	rootstate := walkstate{s.position, 0}
	var wsMap walkstateMap
	solver := solve.NewSolver(rootstate).
		Context(context).
		Constraint(solve.CheapestPathConstraint(&wsMap)).
		Algorithm(solve.BreadthFirst)
	solutions := make([]walkstate, 0)
	for solution := range solver.SolveAll() {
		solutions = append(solutions, solution.GoalState().(walkstate))
	}
	return solutions
}

func parse(level string) (sokoban, mainstate) {
	width := 0
	lines := strings.Split(level, "\n")
	height := len(lines)
	for _, line := range lines {
		if len(line) > width {
			width = len(line)
		}
	}
	var c sokoban
	var s mainstate
	c.width = width
	c.height = height

	c.world = make([]byte, width*height)
	c.goals = make([]uint16, 0)
	s.boxes = make([]uint16, 0)
	for y, row := range lines {
		for x, raw := range row {
			position := y*width + x
			if value, ok := chars[raw]; ok {
				c.world[position] = value &^ player &^ box
				if value&player != 0 {
					s.position = position
				}
				if value&goal != 0 {
					c.goals = append(c.goals, uint16(position))
				}
				if value&box != 0 {
					s.boxes = append(s.boxes, uint16(position))
				}
			} else {
				panic(fmt.Sprintf("Invalid level format, character %v is not valid", value))
			}
		}
	}
	return c, s
}

// For cheapest path constraint
type cpMap map[string]float64

func key(state solve.State) string {
	// nasty hack, but string seems to be the only variable-size type
	// supported as map key. Would love to be able to use slices directly
	s := state.(mainstate)
	runes := make([]rune, len(s.boxes) + 1)
	runes[0] = rune(s.position)
	for i, box := range s.boxes {
		runes[i+1] = rune(box)
	}
	return string(runes)
}

func (c cpMap) Get(state solve.State) (value float64, ok bool) {
	value, ok = c[key(state)]
	return
}

func (c cpMap) Put(state solve.State, value float64) {
	c[key(state)] = value
}

func (c *cpMap) Clear() {
	*c = make(cpMap)
}

func cheapestPathConstraint() solve.Constraint {
	var m cpMap
	return solve.CheapestPathConstraint(&m)
}

var simpleLevel = `
########
#     @#
#      #
# $ # .#
#  # # #
########`

var level = `
   ####
####  ##
#   $  #
#  *** #
#  . . ##
## * *  #
 ##***  #
  # $ ###
  # @ #
  #####`

func main() {
	f, err := os.Create("cpu.prof")
	if err != nil {
		log.Fatal(err)
	}
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()

	world, root := parse(level)
	print(world, root)
	start := time.Now()
	result := solve.NewSolver(root).
		Context(world).
		Algorithm(solve.Astar).
		Constraint(cheapestPathConstraint()).
		//Limit(40).
		Solve()
	fmt.Printf("Time: %.1f seconds\n", time.Since(start).Seconds())
	if result.Solved() {
		fmt.Printf("Result:\n ")
		for _, state := range result.Solution {
			print(world, state.(mainstate))
		}
		fmt.Printf("Solved in %d moves\n", int(result.GoalState().(mainstate).cost))
	}
	fmt.Printf("visited %v main nodes\n", result.Visited)
}
<|endoftext|>


<commit_before>package beacon

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"

	"code.cloudfoundry.org/lager"
	"golang.org/x/crypto/ssh"
)

var ErrFailedToReachAnyTSA = errors.New("failed to connect to TSA")

func NewSSHClient(logger lager.Logger, config Config) Client {
	return &sshClient{
		logger: logger,
		config: config,
	}
}

type sshClient struct {
	logger lager.Logger
	client *ssh.Client
	config Config
	conn   ssh.Conn
}

func (c *sshClient) Dial() (Closeable, error) {
	tsaAddr := c.config.TSAConfig.Host[rand.Intn(len(c.config.TSAConfig.Host))]

	conn, err := keepaliveDialer("tcp", tsaAddr, 10*time.Second)
	if err != nil {
		c.logger.Error("failed to connect to TSA", err)
		return nil, ErrFailedToReachAnyTSA
	}

	pk, err := ssh.NewSignerFromKey(c.config.TSAConfig.WorkerPrivateKey.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to construct ssh public key from worker key: %s", err)
	}

	clientConfig := &ssh.ClientConfig{
		User: "beacon", // doesn't matter

		HostKeyCallback: c.config.checkHostKey,

		Auth: []ssh.AuthMethod{ssh.PublicKeys(pk)},
	}

	clientConn, chans, reqs, err := ssh.NewClientConn(conn, tsaAddr, clientConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to construct client connection: %s", err)
	}

	c.client = ssh.NewClient(clientConn, chans, reqs)

	return c.client, nil
}

func (c *sshClient) KeepAlive() (<-chan error, chan<- struct{}) {
	logger := c.logger.Session("keepalive")

	errs := make(chan error, 1)

	kas := time.NewTicker(5 * time.Second)
	cancel := make(chan struct{})

	go func() {
		for {
			// ignore reply; server may just not have handled it, since there's no
			// standard keepalive request name

			_, _, err := c.client.Conn.SendRequest("keepalive", true, []byte("sup"))
			if err != nil {
				logger.Error("failed", err)
				errs <- err
				return
			}

			logger.Debug("ok")

			select {
			case <-kas.C:
			case <-cancel:
				errs <- nil
				return
			}
		}
	}()

	return errs, cancel
}

func (c *sshClient) Listen(n, addr string) (net.Listener, error) {
	return c.client.Listen(n, addr)
}

func (c *sshClient) Close() error {
	return c.client.Close()
}

func (c *sshClient) NewSession(stdin io.Reader, stdout io.Writer, stderr io.Writer) (Session, error) {
	sess, err := c.client.NewSession()
	if err != nil {
		return nil, err
	}

	sess.Stdin = stdin
	sess.Stdout = stdout
	sess.Stderr = stderr

	return sess, nil
}

func (c *sshClient) Proxy(from, to string) error {
	remoteListener, err := c.Listen("tcp", from)
	if err != nil {
		return fmt.Errorf("failed to listen remotely: %s", err)
	}
	go c.proxyListenerTo(remoteListener, to)
	return nil
}

func (c *sshClient) proxyListenerTo(listener net.Listener, addr string) {
	for {
		rConn, err := listener.Accept()
		if err != nil {
			break
		}

		go c.handleForwardedConn(rConn, addr)
	}
}

func (c *sshClient) handleForwardedConn(rConn net.Conn, addr string) {
	defer rConn.Close()

	lConn, err := net.Dial("tcp", addr)
	if err != nil {
		log.Println("failed to forward remote connection:", err)
		return
	}

	wg := new(sync.WaitGroup)

	pipe := func(to io.WriteCloser, from io.ReadCloser) {
		// if either end breaks, close both ends to ensure they're both unblocked,
		// otherwise io.Copy can block forever if e.g. reading after write end has
		// gone away
		defer to.Close()
		defer from.Close()
		defer wg.Done()

		io.Copy(to, from)
	}

	wg.Add(1)
	go pipe(lConn, rConn)

	wg.Add(1)
	go pipe(rConn, lConn)

	wg.Wait()
}
<commit_msg>properly format log msg<commit_after>package beacon

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"

	"code.cloudfoundry.org/lager"
	"golang.org/x/crypto/ssh"
)

var ErrFailedToReachAnyTSA = errors.New("failed to connect to TSA")

func NewSSHClient(logger lager.Logger, config Config) Client {
	return &sshClient{
		logger: logger,
		config: config,
	}
}

type sshClient struct {
	logger lager.Logger
	client *ssh.Client
	config Config
	conn   ssh.Conn
}

func (c *sshClient) Dial() (Closeable, error) {
	tsaAddr := c.config.TSAConfig.Host[rand.Intn(len(c.config.TSAConfig.Host))]

	conn, err := keepaliveDialer("tcp", tsaAddr, 10*time.Second)
	if err != nil {
		c.logger.Error("failed-to-connect-to-tsa", err)
		return nil, ErrFailedToReachAnyTSA
	}

	pk, err := ssh.NewSignerFromKey(c.config.TSAConfig.WorkerPrivateKey.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to construct ssh public key from worker key: %s", err)
	}

	clientConfig := &ssh.ClientConfig{
		User: "beacon", // doesn't matter

		HostKeyCallback: c.config.checkHostKey,

		Auth: []ssh.AuthMethod{ssh.PublicKeys(pk)},
	}

	clientConn, chans, reqs, err := ssh.NewClientConn(conn, tsaAddr, clientConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to construct client connection: %s", err)
	}

	c.client = ssh.NewClient(clientConn, chans, reqs)

	return c.client, nil
}

func (c *sshClient) KeepAlive() (<-chan error, chan<- struct{}) {
	logger := c.logger.Session("keepalive")

	errs := make(chan error, 1)

	kas := time.NewTicker(5 * time.Second)
	cancel := make(chan struct{})

	go func() {
		for {
			// ignore reply; server may just not have handled it, since there's no
			// standard keepalive request name

			_, _, err := c.client.Conn.SendRequest("keepalive", true, []byte("sup"))
			if err != nil {
				logger.Error("failed", err)
				errs <- err
				return
			}

			logger.Debug("ok")

			select {
			case <-kas.C:
			case <-cancel:
				errs <- nil
				return
			}
		}
	}()

	return errs, cancel
}

func (c *sshClient) Listen(n, addr string) (net.Listener, error) {
	return c.client.Listen(n, addr)
}

func (c *sshClient) Close() error {
	return c.client.Close()
}

func (c *sshClient) NewSession(stdin io.Reader, stdout io.Writer, stderr io.Writer) (Session, error) {
	sess, err := c.client.NewSession()
	if err != nil {
		return nil, err
	}

	sess.Stdin = stdin
	sess.Stdout = stdout
	sess.Stderr = stderr

	return sess, nil
}

func (c *sshClient) Proxy(from, to string) error {
	remoteListener, err := c.Listen("tcp", from)
	if err != nil {
		return fmt.Errorf("failed to listen remotely: %s", err)
	}
	go c.proxyListenerTo(remoteListener, to)
	return nil
}

func (c *sshClient) proxyListenerTo(listener net.Listener, addr string) {
	for {
		rConn, err := listener.Accept()
		if err != nil {
			break
		}

		go c.handleForwardedConn(rConn, addr)
	}
}

func (c *sshClient) handleForwardedConn(rConn net.Conn, addr string) {
	defer rConn.Close()

	lConn, err := net.Dial("tcp", addr)
	if err != nil {
		log.Println("failed to forward remote connection:", err)
		return
	}

	wg := new(sync.WaitGroup)

	pipe := func(to io.WriteCloser, from io.ReadCloser) {
		// if either end breaks, close both ends to ensure they're both unblocked,
		// otherwise io.Copy can block forever if e.g. reading after write end has
		// gone away
		defer to.Close()
		defer from.Close()
		defer wg.Done()

		io.Copy(to, from)
	}

	wg.Add(1)
	go pipe(lConn, rConn)

	wg.Add(1)
	go pipe(rConn, lConn)

	wg.Wait()
}
<|endoftext|>