<commit_before>package packfile

import (
	"bytes"
	"io/ioutil"

	"gopkg.in/src-d/go-git.v4/plumbing"
)

// See https://github.com/jelmer/dulwich/blob/master/dulwich/pack.py and
// https://github.com/tarruda/node-git-core/blob/master/src/js/delta.js
// for more info

const (
	// Standard chunk size used to generate fingerprints
	s = 16

	// https://github.com/git/git/blob/f7466e94375b3be27f229c78873f0acf8301c0a5/diff-delta.c#L428
	// Max size of a copy operation (64KB)
	maxCopySize = 64 * 1024
)

// GetDelta returns an EncodedObject of type OFSDeltaObject. Base and Target object,
// will be loaded into memory to be able to create the delta object.
// To generate target again, you will need the obtained object and "base" one.
// Error will be returned if base or target object cannot be read.
func GetDelta(base, target plumbing.EncodedObject) (plumbing.EncodedObject, error) {
	return getDelta(new(deltaIndex), base, target)
}

func getDelta(index *deltaIndex, base, target plumbing.EncodedObject) (plumbing.EncodedObject, error) {
	br, err := base.Reader()
	if err != nil {
		return nil, err
	}
	tr, err := target.Reader()
	if err != nil {
		return nil, err
	}

	bb, err := ioutil.ReadAll(br)
	if err != nil {
		return nil, err
	}

	tb, err := ioutil.ReadAll(tr)
	if err != nil {
		return nil, err
	}

	db := diffDelta(index, bb, tb)
	delta := &plumbing.MemoryObject{}
	_, err = delta.Write(db)
	if err != nil {
		return nil, err
	}

	delta.SetSize(int64(len(db)))
	delta.SetType(plumbing.OFSDeltaObject)

	return delta, nil
}

// DiffDelta returns the delta that transforms src into tgt.
func DiffDelta(src, tgt []byte) []byte {
	return diffDelta(new(deltaIndex), src, tgt)
}

func diffDelta(index *deltaIndex, src []byte, tgt []byte) []byte {
	buf := bufPool.Get().(*bytes.Buffer)
	buf.Reset()
	buf.Write(deltaEncodeSize(len(src)))
	buf.Write(deltaEncodeSize(len(tgt)))

	if len(index.entries) == 0 {
		index.init(src)
	}

	ibuf := bufPool.Get().(*bytes.Buffer)
	ibuf.Reset()
	for i := 0; i < len(tgt); i++ {
		offset, l := index.findMatch(src, tgt, i)

		if l == 0 {
			// couldn't find a match, just write the current byte and continue
			ibuf.WriteByte(tgt[i])
		} else if l < 0 {
			// src is less than blksz, copy the rest of the target to avoid
			// calls to findMatch
			for ; i < len(tgt); i++ {
				ibuf.WriteByte(tgt[i])
			}
		} else if l < s {
			// remaining target is less than blksz, copy what's left of it
			// and avoid calls to findMatch
			for j := i; j < i+l; j++ {
				ibuf.WriteByte(tgt[j])
			}
			i += l - 1
		} else {
			encodeInsertOperation(ibuf, buf)

			rl := l
			aOffset := offset
			for {
				if rl < maxCopySize {
					buf.Write(encodeCopyOperation(aOffset, rl))
					break
				}

				buf.Write(encodeCopyOperation(aOffset, maxCopySize))
				rl -= maxCopySize
				aOffset += maxCopySize
			}

			i += l - 1
		}
	}

	encodeInsertOperation(ibuf, buf)
	bytes := buf.Bytes()

	bufPool.Put(buf)
	bufPool.Put(ibuf)

	return bytes
}

func encodeInsertOperation(ibuf, buf *bytes.Buffer) {
	if ibuf.Len() == 0 {
		return
	}

	b := ibuf.Bytes()
	s := ibuf.Len()
	o := 0
	for {
		if s <= 127 {
			break
		}
		buf.WriteByte(byte(127))
		buf.Write(b[o : o+127])
		s -= 127
		o += 127
	}
	buf.WriteByte(byte(s))
	buf.Write(b[o : o+s])

	ibuf.Reset()
}

func deltaEncodeSize(size int) []byte {
	var ret []byte
	c := size & 0x7f
	size >>= 7
	for {
		if size == 0 {
			break
		}

		ret = append(ret, byte(c|0x80))
		c = size & 0x7f
		size >>= 7
	}
	ret = append(ret, byte(c))

	return ret
}

func encodeCopyOperation(offset, length int) []byte {
	code := 0x80
	var opcodes []byte

	var i uint
	for i = 0; i < 4; i++ {
		f := 0xff << (i * 8)
		if offset&f != 0 {
			opcodes = append(opcodes, byte(offset&f>>(i*8)))
			code |= 0x01 << i
		}
	}

	for i = 0; i < 3; i++ {
		f := 0xff << (i * 8)
		if length&f != 0 {
			opcodes = append(opcodes, byte(length&f>>(i*8)))
			code |= 0x10 << i
		}
	}

	return append([]byte{byte(code)}, opcodes...)
}
<commit_msg>packfile: use buffer pool for diffs<commit_after>package packfile

import (
	"bytes"

	"gopkg.in/src-d/go-git.v4/plumbing"
)

// See https://github.com/jelmer/dulwich/blob/master/dulwich/pack.py and
// https://github.com/tarruda/node-git-core/blob/master/src/js/delta.js
// for more info

const (
	// Standard chunk size used to generate fingerprints
	s = 16

	// https://github.com/git/git/blob/f7466e94375b3be27f229c78873f0acf8301c0a5/diff-delta.c#L428
	// Max size of a copy operation (64KB)
	maxCopySize = 64 * 1024
)

// GetDelta returns an EncodedObject of type OFSDeltaObject. Base and Target object,
// will be loaded into memory to be able to create the delta object.
// To generate target again, you will need the obtained object and "base" one.
// Error will be returned if base or target object cannot be read.
func GetDelta(base, target plumbing.EncodedObject) (plumbing.EncodedObject, error) {
	return getDelta(new(deltaIndex), base, target)
}

func getDelta(index *deltaIndex, base, target plumbing.EncodedObject) (plumbing.EncodedObject, error) {
	br, err := base.Reader()
	if err != nil {
		return nil, err
	}
	defer br.Close()
	tr, err := target.Reader()
	if err != nil {
		return nil, err
	}
	defer tr.Close()

	bb := bufPool.Get().(*bytes.Buffer)
	bb.Reset()
	defer bufPool.Put(bb)

	_, err = bb.ReadFrom(br)
	if err != nil {
		return nil, err
	}

	tb := bufPool.Get().(*bytes.Buffer)
	tb.Reset()
	defer bufPool.Put(tb)

	_, err = tb.ReadFrom(tr)
	if err != nil {
		return nil, err
	}

	db := diffDelta(index, bb.Bytes(), tb.Bytes())
	delta := &plumbing.MemoryObject{}
	_, err = delta.Write(db)
	if err != nil {
		return nil, err
	}

	delta.SetSize(int64(len(db)))
	delta.SetType(plumbing.OFSDeltaObject)

	return delta, nil
}

// DiffDelta returns the delta that transforms src into tgt.
func DiffDelta(src, tgt []byte) []byte {
	return diffDelta(new(deltaIndex), src, tgt)
}

func diffDelta(index *deltaIndex, src []byte, tgt []byte) []byte {
	buf := bufPool.Get().(*bytes.Buffer)
	buf.Reset()
	buf.Write(deltaEncodeSize(len(src)))
	buf.Write(deltaEncodeSize(len(tgt)))

	if len(index.entries) == 0 {
		index.init(src)
	}

	ibuf := bufPool.Get().(*bytes.Buffer)
	ibuf.Reset()
	for i := 0; i < len(tgt); i++ {
		offset, l := index.findMatch(src, tgt, i)

		if l == 0 {
			// couldn't find a match, just write the current byte and continue
			ibuf.WriteByte(tgt[i])
		} else if l < 0 {
			// src is less than blksz, copy the rest of the target to avoid
			// calls to findMatch
			for ; i < len(tgt); i++ {
				ibuf.WriteByte(tgt[i])
			}
		} else if l < s {
			// remaining target is less than blksz, copy what's left of it
			// and avoid calls to findMatch
			for j := i; j < i+l; j++ {
				ibuf.WriteByte(tgt[j])
			}
			i += l - 1
		} else {
			encodeInsertOperation(ibuf, buf)

			rl := l
			aOffset := offset
			for {
				if rl < maxCopySize {
					buf.Write(encodeCopyOperation(aOffset, rl))
					break
				}

				buf.Write(encodeCopyOperation(aOffset, maxCopySize))
				rl -= maxCopySize
				aOffset += maxCopySize
			}

			i += l - 1
		}
	}

	encodeInsertOperation(ibuf, buf)
	bytes := buf.Bytes()

	bufPool.Put(buf)
	bufPool.Put(ibuf)

	return bytes
}

func encodeInsertOperation(ibuf, buf *bytes.Buffer) {
	if ibuf.Len() == 0 {
		return
	}

	b := ibuf.Bytes()
	s := ibuf.Len()
	o := 0
	for {
		if s <= 127 {
			break
		}
		buf.WriteByte(byte(127))
		buf.Write(b[o : o+127])
		s -= 127
		o += 127
	}
	buf.WriteByte(byte(s))
	buf.Write(b[o : o+s])

	ibuf.Reset()
}

func deltaEncodeSize(size int) []byte {
	var ret []byte
	c := size & 0x7f
	size >>= 7
	for {
		if size == 0 {
			break
		}

		ret = append(ret, byte(c|0x80))
		c = size & 0x7f
		size >>= 7
	}
	ret = append(ret, byte(c))

	return ret
}

func encodeCopyOperation(offset, length int) []byte {
	code := 0x80
	var opcodes []byte

	var i uint
	for i = 0; i < 4; i++ {
		f := 0xff << (i * 8)
		if offset&f != 0 {
			opcodes = append(opcodes, byte(offset&f>>(i*8)))
			code |= 0x01 << i
		}
	}

	for i = 0; i < 3; i++ {
		f := 0xff << (i * 8)
		if length&f != 0 {
			opcodes = append(opcodes, byte(length&f>>(i*8)))
			code |= 0x10 << i
		}
	}

	return append([]byte{byte(code)}, opcodes...)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Portworx

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 manager

import (
	"testing"

	"github.com/libopenstorage/openstorage/cluster"
	"github.com/libopenstorage/openstorage/config"
	"github.com/libopenstorage/openstorage/pkg/auth"
	"github.com/libopenstorage/openstorage/pkg/auth/systemtoken"
	"github.com/portworx/kvdb"
	"github.com/portworx/kvdb/mem"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
)

const (
	testClusterId   = "test-cluster-id"
	testClusterUuid = "test-cluster-uuid"
)

var (
	kv kvdb.Kvdb
)

func init() {
	var err error
	kv, err = kvdb.New(mem.Name, "manager_test/"+testClusterId, []string{}, nil, logrus.Panicf)
	if err != nil {
		logrus.Panicf("Failed to initialize KVDB")
	}
	if err := kvdb.SetInstance(kv); err != nil {
		logrus.Panicf("Failed to set KVDB instance")
	}
}

func TestClusterManagerUuid(t *testing.T) {
	oldInst := inst
	defer func() {
		inst = oldInst
	}()

	uuid := "uuid"
	id := "id"
	err := Init(config.ClusterConfig{
		ClusterId:   id,
		ClusterUuid: uuid,
	})
	assert.NoError(t, err)
	assert.Equal(t, uuid, inst.Uuid())
}

func TestUpdateSchedulerNodeName(t *testing.T) {

	nodeID := "node-alpha"
	Init(config.ClusterConfig{
		ClusterId:         testClusterId,
		ClusterUuid:       testClusterUuid,
		NodeId:            nodeID,
		SchedulerNodeName: "old-sched-name",
	})
	manager, err := systemtoken.NewManager(&systemtoken.Config{
		ClusterId:    testClusterId,
		NodeId:       nodeID,
		SharedSecret: "mysecret",
	})
	assert.NoError(t, err)
	auth.InitSystemTokenManager(manager)

	err = inst.StartWithConfiguration(1, false, "1001", []string{}, "", &cluster.ClusterServerConfiguration{
		ConfigSystemTokenManager: manager,
	})
	assert.NoError(t, err)

	node, err := inst.Inspect(nodeID)
	assert.NoError(t, err)
	assert.Equal(t, "old-sched-name", node.SchedulerNodeName)
	assert.Equal(t, node.GossipPort, "1001", "Expected gossip port to be updated in cluster database")

	err = inst.UpdateSchedulerNodeName("new-sched-name")
	assert.NoError(t, err)

	node, err = inst.Inspect(nodeID)
	assert.NoError(t, err)
	assert.Equal(t, "new-sched-name", node.SchedulerNodeName)

	tokenResp, err := inst.GetPairToken(false)
	assert.NoError(t, err)
	assert.True(t, auth.IsJwtToken(tokenResp.Token))

}

func TestClusterDomainClusterManagerAPIs(t *testing.T) {
	nodeID := "node-alpha"
	Init(config.ClusterConfig{
		ClusterId:         testClusterId,
		ClusterUuid:       testClusterUuid,
		NodeId:            nodeID,
		SchedulerNodeName: "old-sched-name",
	})
	selfClusterDomain := "zone1"
	manager, err := systemtoken.NewManager(&systemtoken.Config{
		ClusterId:    testClusterId,
		NodeId:       nodeID,
		SharedSecret: "mysecret",
	})
	assert.NoError(t, err)
	auth.InitSystemTokenManager(manager)

	err = inst.StartWithConfiguration(1, false, "1002", []string{},
		selfMetroDomain,
		&cluster.ClusterServerConfiguration{
			ConfigSystemTokenManager: manager,
		})
	assert.NoError(t, err)

	c, err := inst.Enumerate()
	assert.NoError(t, err)
	assert.Equal(t, 1, len(c.ClusterDomainsActiveMap), "Unexpected length of cluster domain map")
	assert.True(t, c.ClusterDomainsActiveMap[selfClusterDomain], "Expected zone1 to be activated")
	cleanup()
}
<commit_msg>Fix Cluster Manager UTs: Do a cleanup after every test<commit_after>/*
Copyright 2018 Portworx

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 manager

import (
	"testing"
	"time"

	"github.com/libopenstorage/openstorage/cluster"
	"github.com/libopenstorage/openstorage/config"
	"github.com/libopenstorage/openstorage/pkg/auth"
	"github.com/libopenstorage/openstorage/pkg/auth/systemtoken"
	"github.com/portworx/kvdb"
	"github.com/portworx/kvdb/mem"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
)

const (
	testClusterId   = "test-cluster-id"
	testClusterUuid = "test-cluster-uuid"
)

var (
	kv kvdb.Kvdb
)

func init() {
	var err error
	kv, err = kvdb.New(mem.Name, "manager_test/"+testClusterId, []string{}, nil, logrus.Panicf)
	if err != nil {
		logrus.Panicf("Failed to initialize KVDB")
	}
	if err := kvdb.SetInstance(kv); err != nil {
		logrus.Panicf("Failed to set KVDB instance")
	}
}

func cleanup() {
	inst.Shutdown()
	time.Sleep(100 * time.Millisecond)
	kvdb.Instance().Delete(ClusterDBKey)
	inst = nil
}

func TestClusterManagerUuid(t *testing.T) {
	oldInst := inst
	defer func() {
		inst = oldInst
	}()

	uuid := "uuid"
	id := "id"
	err := Init(config.ClusterConfig{
		ClusterId:   id,
		ClusterUuid: uuid,
	})
	assert.NoError(t, err)
	assert.Equal(t, uuid, inst.Uuid())
	cleanup()
}

func TestUpdateSchedulerNodeName(t *testing.T) {

	nodeID := "node-alpha"
	Init(config.ClusterConfig{
		ClusterId:         testClusterId,
		ClusterUuid:       testClusterUuid,
		NodeId:            nodeID,
		SchedulerNodeName: "old-sched-name",
	})
	manager, err := systemtoken.NewManager(&systemtoken.Config{
		ClusterId:    testClusterId,
		NodeId:       nodeID,
		SharedSecret: "mysecret",
	})
	assert.NoError(t, err)
	auth.InitSystemTokenManager(manager)

	err = inst.StartWithConfiguration(1, false, "1001", []string{}, "", &cluster.ClusterServerConfiguration{
		ConfigSystemTokenManager: manager,
	})
	assert.NoError(t, err)

	node, err := inst.Inspect(nodeID)
	assert.NoError(t, err)
	assert.Equal(t, "old-sched-name", node.SchedulerNodeName)
	assert.Equal(t, node.GossipPort, "1001", "Expected gossip port to be updated in cluster database")

	err = inst.UpdateSchedulerNodeName("new-sched-name")
	assert.NoError(t, err)

	node, err = inst.Inspect(nodeID)
	assert.NoError(t, err)
	assert.Equal(t, "new-sched-name", node.SchedulerNodeName)

	tokenResp, err := inst.GetPairToken(false)
	assert.NoError(t, err)
	assert.True(t, auth.IsJwtToken(tokenResp.Token))

	cleanup()
}

func TestClusterDomainClusterManagerAPIs(t *testing.T) {
	nodeID := "node-alpha"
	Init(config.ClusterConfig{
		ClusterId:         testClusterId,
		ClusterUuid:       testClusterUuid,
		NodeId:            nodeID,
		SchedulerNodeName: "old-sched-name",
	})
	selfClusterDomain := "zone1"
	manager, err := systemtoken.NewManager(&systemtoken.Config{
		ClusterId:    testClusterId,
		NodeId:       nodeID,
		SharedSecret: "mysecret",
	})
	assert.NoError(t, err)
	auth.InitSystemTokenManager(manager)

	err = inst.StartWithConfiguration(1, false, "1002", []string{},
		selfClusterDomain,
		&cluster.ClusterServerConfiguration{
			ConfigSystemTokenManager: manager,
		})
	assert.NoError(t, err)

	c, err := inst.Enumerate()
	assert.NoError(t, err)
	assert.Equal(t, 1, len(c.ClusterDomainsActiveMap), "Unexpected length of cluster domain map")
	assert.True(t, c.ClusterDomainsActiveMap[selfClusterDomain], "Expected zone1 to be activated")
	cleanup()
}
<|endoftext|>


<commit_before>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"golang.org/x/sys/unix"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
)

var (
	help    = flag.Bool("h", false, "Help")
	version = flag.Bool("V", false, "Version")
)

func usage() string {
	return "switch_root [-h] [-V]\nswitch_root newroot init"
}

// get device for file descriptor
func getDev(fd int) (dev uint64, err error) {
	var stat syscall.Stat_t

	if err := syscall.Fstat(fd, &stat); err != nil {
		return 0, err
	}

	return stat.Dev, nil
}

// recursively deletes a directory and everything in it
// works with a file descriptor, can delete files not referenceable
// any more (e.g. after chroot)
// does not descent into mounts
func recursiveDelete(fd int) error {
	parentDev, err := getDev(fd)
	if err != nil {
		return err
	}

  // the file descriptor is already open, but allocating a os.File
	// here for it makes reading the files in the dir so much nicer
	dir := os.NewFile(uintptr(fd), "__ignored__") // filename is completely irrelevant
	defer dir.Close()
	names, err := dir.Readdirnames(-1)
	if err != nil {
		return err
	}

	for _, name := range names {
		// loop here, but handle loop in separate function to make defer work as expected
		if err := recusiveDeleteInner(fd, parentDev, name); err != nil {
			return err
		}
	}
	return nil
}

// don't call this directly, this function is the loop content of recursiveDelete
func recusiveDeleteInner(parentFd int, parentDev uint64, childName string) error {
	// O_DIRECTORY and O_NOFOLLOW make this open fail for all files and all symlinks (even when pointing to a dir)
	// we need to filter out symlinks because getDev later follows them
	childFd, err := syscall.Openat(parentFd, childName, syscall.O_DIRECTORY | syscall.O_NOFOLLOW, syscall.O_RDWR)
	if err != nil {
		// either file or symlink, delete in any case
		if err := unix.Unlinkat(parentFd, childName, 0); err != nil {
			return err
		}
	} else {
		// open succeeded, which means it is a real directory
		defer unix.Close(childFd)

		// don't descent into other file systems
		if childFdDev, err := getDev(childFd); err != nil {
			return err
		} else if childFdDev != parentDev {
			// this means continue in recursiveDelete
			return nil
		}

		if err:= recursiveDelete(childFd); err != nil {
			return err
		}
		// back from recursion, dir is now empty, delete
		if err := unix.Unlinkat(parentFd, childName, unix.AT_REMOVEDIR); err != nil {
			return err
		}
	}
	return nil
}

// do a proper exec to retain PID 1.
// It returns an error if the command exits incorrectly.
func execCommand(path string) error {
	return syscall.Exec(path, []string{path}, []string{})
}

// check if dir is empty
func isEmpty(name string) (bool, error) {
	f, err := os.Open(name)
	if err != nil {
		return false, err
	}
	defer f.Close()

	_, err = f.Readdirnames(1)
	if err == io.EOF { // this means emtpy
		return true, nil
	}
	return false, err
}

// moves a mountpoint
func moveMount(oldPath string, newPath string) error {
	return syscall.Mount(oldPath, newPath, "", syscall.MS_MOVE, "")
}

// moves common 'special' mounts if the exist
func specialFS(newRoot string) error {
	var mounts = []string{"/dev", "/proc", "/sys", "/run"}

	for _, mount := range mounts {
		path := filepath.Join(newRoot, mount)
		// skip all mounting if the dir does not exists
		if _, err := os.Stat(mount); os.IsNotExist(err) {
			fmt.Println("switch_root: Skipping", mount)
			continue
		} else if err != nil {
			return err
		}
		// make sure the target dir exist and is empty
		if _, err := os.Stat(path); os.IsNotExist(err) {
			if err := syscall.Mkdir(path, 0); err != nil {
				return err
			}
		} else if err != nil {
			return err
		}
		if empty, err := isEmpty(path); err != nil {
			return err
		} else if !empty {
			return fmt.Errorf("%v must be empty", path)
		}
		if err := moveMount(mount, path); err != nil {
			return err
		}
	}
	return nil
}

// switchroot moves special mounts (dev, proc, sys, run) to the new dir,
// then does a chroot, moves the root mount to the new dir and finally
// DELETES EVERYTHING in the old dir and execs init
func switchRoot(newRoot string, init string) error {
	log.Printf("switch_root: moving mounts")
	if err := specialFS(newRoot); err != nil {
		return fmt.Errorf("switch_root: moving mounts failed %v", err)
	}

	log.Printf("switch_root: Changing directory")
	if err := syscall.Chdir(newRoot); err != nil {
		return fmt.Errorf("switch_root: failed change directory to new_root %v", err)
	}

	// Open "/" now, we need the file descriptor later
	oldRoot, err := os.Open("/")
	if err != nil {
		return err
	}
	defer oldRoot.Close()

	log.Printf("switch_root: Moving /")
	if err := moveMount(newRoot, "/"); err != nil {
		return err
	}

	log.Printf("switch_root: Changing root!")
	if err := syscall.Chroot("."); err != nil {
		return fmt.Errorf("switch_root: fatal chroot error %v", err)
	}

	log.Printf("switch_root: Deleting old /")
	if err := recursiveDelete(int(oldRoot.Fd())); err != nil {
		panic(err)
	}

	log.Printf("switch_root: executing init")
	if err := execCommand(init); err != nil {
		return fmt.Errorf("switch_root: exec failed %v", err)
	}

	return nil
}

func main() {
	flag.Parse()

	if len(flag.Args()) == 0 {
		fmt.Println(usage())
		os.Exit(0)
	}

	if *help {
		fmt.Println(usage())
		os.Exit(0)
	}

	if *version {
		fmt.Println("Version XX")
		os.Exit(0)
	}

	newRoot := flag.Args()[0]
	init := flag.Args()[1]

	if err := switchRoot(newRoot, init); err != nil {
		log.Fatalf("switch_root failed %v\n", err)
	}
}
<commit_msg>switch_root: Add linux build tag<commit_after>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux

package main

import (
	"flag"
	"fmt"
	"golang.org/x/sys/unix"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
)

var (
	help    = flag.Bool("h", false, "Help")
	version = flag.Bool("V", false, "Version")
)

func usage() string {
	return "switch_root [-h] [-V]\nswitch_root newroot init"
}

// get device for file descriptor
func getDev(fd int) (dev uint64, err error) {
	var stat syscall.Stat_t

	if err := syscall.Fstat(fd, &stat); err != nil {
		return 0, err
	}

	return stat.Dev, nil
}

// recursively deletes a directory and everything in it
// works with a file descriptor, can delete files not referenceable
// any more (e.g. after chroot)
// does not descent into mounts
func recursiveDelete(fd int) error {
	parentDev, err := getDev(fd)
	if err != nil {
		return err
	}

  // the file descriptor is already open, but allocating a os.File
	// here for it makes reading the files in the dir so much nicer
	dir := os.NewFile(uintptr(fd), "__ignored__") // filename is completely irrelevant
	defer dir.Close()
	names, err := dir.Readdirnames(-1)
	if err != nil {
		return err
	}

	for _, name := range names {
		// loop here, but handle loop in separate function to make defer work as expected
		if err := recusiveDeleteInner(fd, parentDev, name); err != nil {
			return err
		}
	}
	return nil
}

// don't call this directly, this function is the loop content of recursiveDelete
func recusiveDeleteInner(parentFd int, parentDev uint64, childName string) error {
	// O_DIRECTORY and O_NOFOLLOW make this open fail for all files and all symlinks (even when pointing to a dir)
	// we need to filter out symlinks because getDev later follows them
	childFd, err := syscall.Openat(parentFd, childName, syscall.O_DIRECTORY | syscall.O_NOFOLLOW, syscall.O_RDWR)
	if err != nil {
		// either file or symlink, delete in any case
		if err := unix.Unlinkat(parentFd, childName, 0); err != nil {
			return err
		}
	} else {
		// open succeeded, which means it is a real directory
		defer unix.Close(childFd)

		// don't descent into other file systems
		if childFdDev, err := getDev(childFd); err != nil {
			return err
		} else if childFdDev != parentDev {
			// this means continue in recursiveDelete
			return nil
		}

		if err:= recursiveDelete(childFd); err != nil {
			return err
		}
		// back from recursion, dir is now empty, delete
		if err := unix.Unlinkat(parentFd, childName, unix.AT_REMOVEDIR); err != nil {
			return err
		}
	}
	return nil
}

// do a proper exec to retain PID 1.
// It returns an error if the command exits incorrectly.
func execCommand(path string) error {
	return syscall.Exec(path, []string{path}, []string{})
}

// check if dir is empty
func isEmpty(name string) (bool, error) {
	f, err := os.Open(name)
	if err != nil {
		return false, err
	}
	defer f.Close()

	_, err = f.Readdirnames(1)
	if err == io.EOF { // this means emtpy
		return true, nil
	}
	return false, err
}

// moves a mountpoint
func moveMount(oldPath string, newPath string) error {
	return syscall.Mount(oldPath, newPath, "", syscall.MS_MOVE, "")
}

// moves common 'special' mounts if the exist
func specialFS(newRoot string) error {
	var mounts = []string{"/dev", "/proc", "/sys", "/run"}

	for _, mount := range mounts {
		path := filepath.Join(newRoot, mount)
		// skip all mounting if the dir does not exists
		if _, err := os.Stat(mount); os.IsNotExist(err) {
			fmt.Println("switch_root: Skipping", mount)
			continue
		} else if err != nil {
			return err
		}
		// make sure the target dir exist and is empty
		if _, err := os.Stat(path); os.IsNotExist(err) {
			if err := syscall.Mkdir(path, 0); err != nil {
				return err
			}
		} else if err != nil {
			return err
		}
		if empty, err := isEmpty(path); err != nil {
			return err
		} else if !empty {
			return fmt.Errorf("%v must be empty", path)
		}
		if err := moveMount(mount, path); err != nil {
			return err
		}
	}
	return nil
}

// switchroot moves special mounts (dev, proc, sys, run) to the new dir,
// then does a chroot, moves the root mount to the new dir and finally
// DELETES EVERYTHING in the old dir and execs init
func switchRoot(newRoot string, init string) error {
	log.Printf("switch_root: moving mounts")
	if err := specialFS(newRoot); err != nil {
		return fmt.Errorf("switch_root: moving mounts failed %v", err)
	}

	log.Printf("switch_root: Changing directory")
	if err := syscall.Chdir(newRoot); err != nil {
		return fmt.Errorf("switch_root: failed change directory to new_root %v", err)
	}

	// Open "/" now, we need the file descriptor later
	oldRoot, err := os.Open("/")
	if err != nil {
		return err
	}
	defer oldRoot.Close()

	log.Printf("switch_root: Moving /")
	if err := moveMount(newRoot, "/"); err != nil {
		return err
	}

	log.Printf("switch_root: Changing root!")
	if err := syscall.Chroot("."); err != nil {
		return fmt.Errorf("switch_root: fatal chroot error %v", err)
	}

	log.Printf("switch_root: Deleting old /")
	if err := recursiveDelete(int(oldRoot.Fd())); err != nil {
		panic(err)
	}

	log.Printf("switch_root: executing init")
	if err := execCommand(init); err != nil {
		return fmt.Errorf("switch_root: exec failed %v", err)
	}

	return nil
}

func main() {
	flag.Parse()

	if len(flag.Args()) == 0 {
		fmt.Println(usage())
		os.Exit(0)
	}

	if *help {
		fmt.Println(usage())
		os.Exit(0)
	}

	if *version {
		fmt.Println("Version XX")
		os.Exit(0)
	}

	newRoot := flag.Args()[0]
	init := flag.Args()[1]

	if err := switchRoot(newRoot, init); err != nil {
		log.Fatalf("switch_root failed %v\n", err)
	}
}
<|endoftext|>


<commit_before>package lib

import (
	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/lib/container"
	"github.com/subutai-io/base/agent/lib/fs"
	"github.com/subutai-io/base/agent/lib/net"
	"github.com/subutai-io/base/agent/log"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// LxcPromote promotes the given container name.
func LxcPromote(name string) {
	checkSanity(name)

	// check: start container if it is not running already
	if container.State(name) != "RUNNING" {
		container.Start(name)
		// log.Info("Container " + name + " is started")
	}
	// check: write package list to packages
	pkgCmmdResult, _ := container.AttachExec(name, []string{"dpkg", "-l"})

	strCmdRes := strings.Join(pkgCmmdResult, "\n")
	log.Check(log.FatalLevel, "Write packages",
		ioutil.WriteFile(config.Agent.LxcPrefix+name+"/packages",
			[]byte(strCmdRes), 0755))
	if container.State(name) == "RUNNING" {
		container.Stop(name)
	}
	net.RemoveDefaultGW(name)

	cleanupFS(config.Agent.LxcPrefix+name+"/rootfs/.git", 0000)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/log/", 0775)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/cache", 0775)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/lib/apt/lists/", 0000)

	makeDiff(name)

	iface := container.GetConfigItem(config.Agent.LxcPrefix+name+"/config", "lxc.network.veth.pair")
	net.ConfigureOVS(iface)
	container.ResetNet(name)
	fs.ReadOnly(name, true)
	log.Info(name + " promoted")
}

func clearFile(path string, f os.FileInfo, err error) error {
	if !f.IsDir() {
		ioutil.WriteFile(path, []byte{}, 0775)
	}
	return nil
}

func cleanupFS(path string, perm os.FileMode) {
	if perm == 0000 {
		os.RemoveAll(path)
	} else {
		filepath.Walk(path, clearFile)
	}
}

func makeDiff(name string) {
	parent := container.GetParent(name)
	if parent == name || len(parent) < 1 {
		return
	}
	os.MkdirAll(config.Agent.LxcPrefix+name+"/diff", 0600)
	execDiff(config.Agent.LxcPrefix+parent+"/rootfs", config.Agent.LxcPrefix+name+"/rootfs", config.Agent.LxcPrefix+name+"/diff/rootfs.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/home", config.Agent.LxcPrefix+name+"/home", config.Agent.LxcPrefix+name+"/diff/home.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/opt", config.Agent.LxcPrefix+name+"/opt", config.Agent.LxcPrefix+name+"/diff/opt.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/var", config.Agent.LxcPrefix+name+"/var", config.Agent.LxcPrefix+name+"/diff/var.diff")
}

func execDiff(dir1, dir2, output string) {
	var out []byte
	out, _ = exec.Command("diff", "-Nur", dir1, dir2).Output()
	err := ioutil.WriteFile(output, out, 0600)
	log.Check(log.FatalLevel, "Writing diff to file"+output, err)
}

func checkSanity(name string) {
	// check: if name exists
	if !container.IsContainer(name) {
		log.Error("Container " + name + " does not exist")
	}

	// check: if name is template
	if container.IsTemplate(name) {
		log.Error("Template " + name + " already exists")
	}
	// check: remove default gateway

	parent := container.GetParent(name)
	if parent == name || len(parent) < 1 {
		return
	}
	if !container.IsTemplate(container.GetParent(name)) {
		log.Error("Parent template " + container.GetParent(name) + " not found")
	}
}
<commit_msg>Fixed Appscale template promoting.<commit_after>package lib

import (
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/subutai-io/base/agent/config"
	"github.com/subutai-io/base/agent/lib/container"
	"github.com/subutai-io/base/agent/lib/fs"
	"github.com/subutai-io/base/agent/lib/net"
	"github.com/subutai-io/base/agent/log"
)

// LxcPromote promotes the given container name.
func LxcPromote(name string) {
	checkSanity(name)

	// check: start container if it is not running already
	if container.State(name) != "RUNNING" {
		LxcStart(name)
		// log.Info("Container " + name + " is started")
	}

	// check: write package list to packages
	pkgCmmdResult, _ := container.AttachExec(name, []string{"timeout", "60", "dpkg", "-l"})
	strCmdRes := strings.Join(pkgCmmdResult, "\n")
	log.Check(log.FatalLevel, "Write packages",
		ioutil.WriteFile(config.Agent.LxcPrefix+name+"/packages",
			[]byte(strCmdRes), 0755))
	if container.State(name) == "RUNNING" {
		container.Stop(name)
	}
	net.RemoveDefaultGW(name)

	cleanupFS(config.Agent.LxcPrefix+name+"/rootfs/.git", 0000)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/log/", 0775)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/cache", 0775)
	cleanupFS(config.Agent.LxcPrefix+name+"/var/lib/apt/lists/", 0000)

	makeDiff(name)

	iface := container.GetConfigItem(config.Agent.LxcPrefix+name+"/config", "lxc.network.veth.pair")
	net.ConfigureOVS(iface)
	container.ResetNet(name)
	fs.ReadOnly(name, true)
	log.Info(name + " promoted")
}

func clearFile(path string, f os.FileInfo, err error) error {
	if !f.IsDir() {
		ioutil.WriteFile(path, []byte{}, 0775)
	}
	return nil
}

func cleanupFS(path string, perm os.FileMode) {
	if perm == 0000 {
		os.RemoveAll(path)
	} else {
		filepath.Walk(path, clearFile)
	}
}

func makeDiff(name string) {
	parent := container.GetParent(name)
	if parent == name || len(parent) < 1 {
		return
	}
	os.MkdirAll(config.Agent.LxcPrefix+name+"/diff", 0600)
	execDiff(config.Agent.LxcPrefix+parent+"/rootfs", config.Agent.LxcPrefix+name+"/rootfs", config.Agent.LxcPrefix+name+"/diff/rootfs.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/home", config.Agent.LxcPrefix+name+"/home", config.Agent.LxcPrefix+name+"/diff/home.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/opt", config.Agent.LxcPrefix+name+"/opt", config.Agent.LxcPrefix+name+"/diff/opt.diff")
	execDiff(config.Agent.LxcPrefix+parent+"/var", config.Agent.LxcPrefix+name+"/var", config.Agent.LxcPrefix+name+"/diff/var.diff")
}

func execDiff(dir1, dir2, output string) {
	var out []byte
	out, _ = exec.Command("diff", "-Nur", dir1, dir2).Output()
	err := ioutil.WriteFile(output, out, 0600)
	log.Check(log.FatalLevel, "Writing diff to file"+output, err)
}

func checkSanity(name string) {
	// check: if name exists
	if !container.IsContainer(name) {
		log.Error("Container " + name + " does not exist")
	}

	// check: if name is template
	if container.IsTemplate(name) {
		log.Error("Template " + name + " already exists")
	}
	// check: remove default gateway

	parent := container.GetParent(name)
	if parent == name || len(parent) < 1 {
		return
	}
	if !container.IsTemplate(container.GetParent(name)) {
		log.Error("Parent template " + container.GetParent(name) + " not found")
	}
}
<|endoftext|>


<commit_before>// Package size implements human representation of bytes size.
package size

import "fmt"

// A Byte represents a number of bytes.
type Byte uint64

const (
	_       = iota // ignore first value by assigning to blank identifier
	KB Byte = 1 << (10 * iota)
	MB
	GB
	TB
	PB
	EB
)

// String represents a Byte with the most appropriate unit.
func (b Byte) String() string {
	switch {
	case b >= EB:
		return fmt.Sprintf("%vEB", b.fmt(EB))
	case b >= PB:
		return fmt.Sprintf("%vPB", b.fmt(PB))
	case b >= TB:
		return fmt.Sprintf("%vTB", b.fmt(TB))
	case b >= GB:
		return fmt.Sprintf("%vGB", b.fmt(GB))
	case b >= MB:
		return fmt.Sprintf("%vMB", b.fmt(MB))
	case b >= KB:
		return fmt.Sprintf("%vKB", b.fmt(KB))
	}
	return fmt.Sprintf("%dB", b)
}

// fmt represents a Byte in a unit with 3 significant digits minimum.
func (b Byte) fmt(unit Byte) string {
	f := "%.0f"
	if b < 10*unit {
		f = "%.2f"
	} else if b < 100*unit {
		f = "%.1f"
	}
	return fmt.Sprintf(f, float64(b)/float64(unit))
}
<commit_msg>lib/size: fix lints<commit_after>// Package size implements human representation of bytes size.
package size

import "fmt"

// A Byte represents a number of bytes.
type Byte uint64

// Suffixes for KB, MB, GB, TB, PB and EB.
const (
	_       = iota // ignore first value by assigning to blank identifier
	KB Byte = 1 << (10 * iota)
	MB
	GB
	TB
	PB
	EB
)

// String represents a Byte with the most appropriate unit.
func (b Byte) String() string {
	switch {
	case b >= EB:
		return fmt.Sprintf("%vEB", b.fmt(EB))
	case b >= PB:
		return fmt.Sprintf("%vPB", b.fmt(PB))
	case b >= TB:
		return fmt.Sprintf("%vTB", b.fmt(TB))
	case b >= GB:
		return fmt.Sprintf("%vGB", b.fmt(GB))
	case b >= MB:
		return fmt.Sprintf("%vMB", b.fmt(MB))
	case b >= KB:
		return fmt.Sprintf("%vKB", b.fmt(KB))
	}
	return fmt.Sprintf("%dB", b)
}

// fmt represents a Byte in a unit with 3 significant digits minimum.
func (b Byte) fmt(unit Byte) string {
	f := "%.0f"
	if b < 10*unit {
		f = "%.2f"
	} else if b < 100*unit {
		f = "%.1f"
	}
	return fmt.Sprintf(f, float64(b)/float64(unit))
}
<|endoftext|>


<commit_before>package oauth2

import (
	"github.com/stretchr/gomniauth/common"
	"github.com/stretchr/gomniauth/test"
	"github.com/stretchr/objx"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"
)

func TestOAuth2Provider_CompleteAuth_URLEncodedResponse(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	creds := new(common.Credentials)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "text/plain")
	testResponse.StatusCode = 200
	testResponse.Body = ioutil.NopCloser(strings.NewReader("expires_in=20&access_token=ACCESSTOKEN&refresh_token=REFRESHTOKEN"))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	creds, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.NoError(t, err) {
		if assert.NotNil(t, creds, "Creds should be returned") {

			assert.Equal(t, creds.Get(OAuth2KeyAccessToken).Str(), "ACCESSTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyRefreshToken).Str(), "REFRESHTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyExpiresIn).Data().(time.Duration), 20000000000)

		}
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}

func TestOAuth2Provider_Non200Response(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "text/plain")
	testResponse.StatusCode = 401
	testResponse.Body = ioutil.NopCloser(strings.NewReader("No mate"))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	_, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.Error(t, err) {
		assert.IsType(t, &common.AuthServerError{}, err)
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}

func TestOAuth2Provider_CompleteAuth_JSON(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	creds := new(common.Credentials)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "application/json")
	testResponse.StatusCode = 200
	testResponse.Body = ioutil.NopCloser(strings.NewReader(`{"expires_in":20,"access_token":"ACCESSTOKEN","refresh_token":"REFRESHTOKEN"}`))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	creds, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.NoError(t, err) {
		if assert.NotNil(t, creds, "Creds should be returned") {

			assert.Equal(t, creds.Get(OAuth2KeyAccessToken).Str(), "ACCESSTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyRefreshToken).Str(), "REFRESHTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyExpiresIn).Data().(time.Duration), 20000000000)

		}
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}
<commit_msg>Fix constant 20000000000 overflows int<commit_after>package oauth2

import (
	"github.com/stretchr/gomniauth/common"
	"github.com/stretchr/gomniauth/test"
	"github.com/stretchr/objx"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"
)

func TestOAuth2Provider_CompleteAuth_URLEncodedResponse(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	creds := new(common.Credentials)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "text/plain")
	testResponse.StatusCode = 200
	testResponse.Body = ioutil.NopCloser(strings.NewReader("expires_in=20&access_token=ACCESSTOKEN&refresh_token=REFRESHTOKEN"))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	creds, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.NoError(t, err) {
		if assert.NotNil(t, creds, "Creds should be returned") {

			assert.Equal(t, creds.Get(OAuth2KeyAccessToken).Str(), "ACCESSTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyRefreshToken).Str(), "REFRESHTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyExpiresIn).Data().(time.Duration), time.Duration(20)*time.Second)

		}
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}

func TestOAuth2Provider_Non200Response(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "text/plain")
	testResponse.StatusCode = 401
	testResponse.Body = ioutil.NopCloser(strings.NewReader("No mate"))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	_, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.Error(t, err) {
		assert.IsType(t, &common.AuthServerError{}, err)
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}

func TestOAuth2Provider_CompleteAuth_JSON(t *testing.T) {

	config := &common.Config{
		Map: objx.MSI(
			OAuth2KeyRedirectUrl, OAuth2KeyRedirectUrl,
			OAuth2KeyScope, OAuth2KeyScope,
			OAuth2KeyClientID, OAuth2KeyClientID,
			OAuth2KeySecret, OAuth2KeySecret,
			OAuth2KeyAuthURL, OAuth2KeyAuthURL,
			OAuth2KeyTokenURL, OAuth2KeyTokenURL)}

	testTripperFactory := new(test.TestTripperFactory)
	testTripper := new(test.TestTripper)
	testProvider := new(test.TestProvider)

	creds := new(common.Credentials)

	testResponse := new(http.Response)
	testResponse.Header = make(http.Header)
	testResponse.Header.Set("Content-Type", "application/json")
	testResponse.StatusCode = 200
	testResponse.Body = ioutil.NopCloser(strings.NewReader(`{"expires_in":20,"access_token":"ACCESSTOKEN","refresh_token":"REFRESHTOKEN"}`))

	testTripperFactory.On("NewTripper", common.EmptyCredentials, mock.Anything).Return(testTripper, nil)
	testTripper.On("RoundTrip", mock.Anything).Return(testResponse, nil)

	data := objx.MSI(OAuth2KeyCode, []string{"123"})
	creds, err := CompleteAuth(testTripperFactory, data, config, testProvider)

	if assert.NoError(t, err) {
		if assert.NotNil(t, creds, "Creds should be returned") {

			assert.Equal(t, creds.Get(OAuth2KeyAccessToken).Str(), "ACCESSTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyRefreshToken).Str(), "REFRESHTOKEN")
			assert.Equal(t, creds.Get(OAuth2KeyExpiresIn).Data().(time.Duration), time.Duration(20)*time.Second)

		}
	}

	mock.AssertExpectationsForObjects(t, testTripperFactory.Mock, testTripper.Mock, testProvider.Mock)

}
<|endoftext|>


<commit_before>package rtsengine

import (
	"container/list"
	"fmt"
	"image"
	"sync"
)

/*
 Maintains the A* Pathing algorithm

*/

// AStarPathing will implement the A* pathing algorithm
// A simple description here: http://www.policyalmanac.org/games/aStarTutorial.htm
// Psuedocode here at the bottom of this: http://web.mit.edu/eranki/www/tutorials/search/
// https://github.com/beefsack/go-astar/blob/master/astar.go
type AStarPathing struct {
	// We need to only path-find one at a time otherwise
	// if we path-find as the world changes it will end badly.
	muPathing sync.Mutex
}

// FindPath will find a path between source and destination Points and
// returns a list of Squares of the proper path.
// All coordinates in world coordinates (absolute coordinates) please.
func (path *AStarPathing) FindPath(pool *Pool, grid *Grid, source *image.Point, destination *image.Point) (*list.List, error) {

	// Check if both source and destination are not colliding
	if !grid.In(source) || grid.Collision(source) {
		return nil, fmt.Errorf("Source not in grid or collision! (%d,%d)", source.X, source.Y)
	}

	if !grid.In(destination) || grid.Collision(destination) {
		return nil, fmt.Errorf("Destination not in grid or collision! (%d,%d)", destination.X, destination.Y)
	}

	closedList := list.New()
	openList := list.New()

	// Starting square. 0 out the cost.
	q := pool.Squares(1)[0]
	q.F = 0
	q.G = 0
	q.H = 0
	q.Locus.X = source.X
	q.Locus.Y = source.Y
	q.Position = 0

	// Push onto the openlist to prime the pathing engine
	openList.PushFront(q)

	// While the open list is not empty
	for openList.Len() > 0 {
		//find the node with the least f on the open list, call it "q"
		//remove q from the open list
		q = path.leastF(openList)

		// generate q's 8 successors and set their parents to q
		successors := path.constructSuccessor(pool, q)
		for i, successor := range successors {

			// ensure it is in the grid and there isn't a collision
			if !grid.In(&successor.Locus) || grid.Collision(&successor.Locus) {
				pool.Free(successor)
				continue
			}

			//if successor is the goal, stop the search
			if destination.Eq(successor.Locus) {
				closedList.PushBack(q)
				closedList.PushBack(successor)
				path.freeList(pool, openList)
				path.freeArray(pool, i+1, successors)
				return path.optimizePath(pool, closedList), nil
			}

			//successor.g = q.g + distance between successor and q
			D := grid.Distance(&q.Locus, &successor.Locus)
			successor.G = q.G + D

			// successor.h = distance from goal to successor
			successor.H = grid.Distance(&successor.Locus, destination)

			// successor.f = successor.g + successor.h
			successor.F = successor.G + successor.H

			//  if a node with the same position as successor is in the OPEN list
			//  which has a lower f than successor, skip this successor
			if path.skipSuccessor(successor, openList) {
				pool.Free(successor)
				continue
			}

			// if a node with the same position as successor is in the CLOSED list \
			// which has a lower f than successor, skip this successor
			if path.skipSuccessor(successor, closedList) {
				pool.Free(successor)
				continue
			}

			// otherwise, add the node to the open list
			openList.PushBack(successor)

		} // for successors

		// push q on the closed list
		closedList.PushBack(q)

	} // openList non empty

	// Free all the remaining successors in the open list.
	path.freeList(pool, openList)

	return path.optimizePath(pool, closedList), nil
}

// freeArray will free all squares in array from i .. len(squares)-1
func (path *AStarPathing) freeArray(pool *Pool, i int, squares []*Square) {
	if i >= len(squares) {
		return
	}

	for ; i < len(squares); i++ {
		pool.Free(squares[i])
	}
}

// freeList will free every Square in the list l
func (path *AStarPathing) freeList(pool *Pool, l *list.List) {
	// Free all the remaining successors in the open list.
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)
		pool.Free(square)
	}
}

// optimizePath will optimize the path list passed as a parameter. Any culled
// nodes are freed from the pool.
//
// A path list will contain duplicates at each _position_. Thus you want to
// iterate over the list and remove duplicates at each _position_ leaving the
// square with the least F in the path list.
// For F ties only one is chosen.
func (path *AStarPathing) optimizePath(pool *Pool, l *list.List) *list.List {
	var m map[int]*Square

	m = make(map[int]*Square)
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)

		p, ok := m[square.Position]

		if !ok {
			m[square.Position] = square
		} else {
			if p.F <= square.F {
				pool.Free(square)
			} else {
				m[square.Position] = square
				pool.Free(p)
			}
		}

	}
	result := list.New()

	length := len(m)
	for i := 0; i < length; i++ {
		result.PushBack(m[i])
	}

	return result
}

//

// skipSuccessor will scan list l and if the list l contains an element with a smaller
// F than the successor at the same position, returns TRUE.
func (path *AStarPathing) skipSuccessor(successor *Square, l *list.List) bool {
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)

		if square.Position == successor.Position && square.F <= successor.F {
			return true
		}
	}

	return false
}

// constructSuccessor will construct 8 successors with parent q from the Pool pool.
func (path *AStarPathing) constructSuccessor(pool *Pool, q *Square) []*Square {
	// The successors are the adjoining squares with the source S
	// in the middle. See below. It moves clockwise. We index
	// from zero so index 0 is 1 below.
	// 1  2  3
	// 8  S  4
	// 7  6  5
	successors := pool.Squares(8)

	// 1
	successors[0].Locus.X = q.Locus.X - 1
	successors[0].Locus.Y = q.Locus.Y - 1
	successors[0].Parent = q
	successors[0].Position = q.Position + 1

	// 2
	successors[1].Locus.X = q.Locus.X - 1
	successors[1].Parent = q
	successors[1].Position = q.Position + 1

	// 3
	successors[2].Locus.X = q.Locus.X - 1
	successors[2].Locus.Y = q.Locus.Y + 1
	successors[2].Parent = q
	successors[2].Position = q.Position + 1

	// 4
	successors[3].Locus.Y = q.Locus.Y + 1
	successors[3].Parent = q
	successors[3].Position = q.Position + 1

	// 5
	successors[4].Locus.X = q.Locus.X + 1
	successors[4].Locus.Y = q.Locus.Y + 1
	successors[4].Parent = q
	successors[4].Position = q.Position + 1

	// 6
	successors[5].Locus.X = q.Locus.X + 1
	successors[5].Parent = q
	successors[5].Position = q.Position + 1

	// 7
	successors[6].Locus.X = q.Locus.X + 1
	successors[6].Locus.Y = q.Locus.Y - 1
	successors[6].Parent = q
	successors[6].Position = q.Position + 1

	// 8
	successors[7].Locus.Y = q.Locus.Y - 1
	successors[7].Parent = q
	successors[7].Position = q.Position + 1

	return successors
}

// leastF returns the Square with the least F within list l
// AND remove that Square from list l.
// Returns nil if no item exists.
func (path *AStarPathing) leastF(l *list.List) *Square {

	var leastSquare *Square
	var leastSquareE *list.Element
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)
		if leastSquare == nil || square.F < leastSquare.F {
			leastSquare = square
			leastSquareE = e
		}
	}

	l.Remove(leastSquareE)
	return leastSquare
}
<commit_msg>Fixed the comments.<commit_after>package rtsengine

import (
	"container/list"
	"fmt"
	"image"
	"sync"
)

/*
 Maintains the A* Pathing algorithm

*/

// AStarPathing will implement the A* pathing algorithm
// A simple description here: http://www.policyalmanac.org/games/aStarTutorial.htm
// Psuedocode here at the bottom of this: http://web.mit.edu/eranki/www/tutorials/search/
// https://github.com/beefsack/go-astar/blob/master/astar.go
type AStarPathing struct {
	// We need to only path-find one at a time otherwise
	// if we path-find as the world changes it will end badly.
	muPathing sync.Mutex
}

// FindPath will find a path between source and destination Points and
// returns a list of Squares of the proper path.
// All coordinates in world coordinates (absolute coordinates) please.
func (path *AStarPathing) FindPath(pool *Pool, grid *Grid, source *image.Point, destination *image.Point) (*list.List, error) {

	// Check if both source and destination are not colliding
	if !grid.In(source) || grid.Collision(source) {
		return nil, fmt.Errorf("Source not in grid or collision! (%d,%d)", source.X, source.Y)
	}

	if !grid.In(destination) || grid.Collision(destination) {
		return nil, fmt.Errorf("Destination not in grid or collision! (%d,%d)", destination.X, destination.Y)
	}

	closedList := list.New()
	openList := list.New()

	// Starting square. 0 out the cost.
	q := pool.Squares(1)[0]
	q.F = 0
	q.G = 0
	q.H = 0
	q.Locus.X = source.X
	q.Locus.Y = source.Y
	q.Position = 0

	// Push onto the openlist to prime the pathing engine
	openList.PushFront(q)

	// While the open list is not empty
	for openList.Len() > 0 {
		//find the square with the least f on the open list, call it "q"
		//remove q from the open list
		q = path.leastF(openList)

		// generate q's 8 successors and set their parents to q
		successors := path.constructSuccessor(pool, q)
		for i, successor := range successors {

			// ensure it is in the grid and there isn't a collision
			if !grid.In(&successor.Locus) || grid.Collision(&successor.Locus) {
				pool.Free(successor)
				continue
			}

			//if successor is the goal, stop the search
			if destination.Eq(successor.Locus) {
				closedList.PushBack(q)
				closedList.PushBack(successor)
				path.freeList(pool, openList)
				path.freeArray(pool, i+1, successors)
				return path.optimizePath(pool, closedList), nil
			}

			//successor.g = q.g + distance between successor and q
			D := grid.Distance(&q.Locus, &successor.Locus)
			successor.G = q.G + D

			// successor.h = distance from goal to successor
			successor.H = grid.Distance(&successor.Locus, destination)

			// successor.f = successor.g + successor.h
			successor.F = successor.G + successor.H

			//  if a square with the same position as successor is in the OPEN list
			//  exists and has a lower f than successor, skip this successor
			if path.skipSuccessor(successor, openList) {
				pool.Free(successor)
				continue
			}

			// if a square with the same position as successor is in the CLOSED list
			// exists has a lower f than successor, skip this successor
			if path.skipSuccessor(successor, closedList) {
				pool.Free(successor)
				continue
			}

			// otherwise, add the square to the open list
			openList.PushBack(successor)

		} // for successors

		// push q on the closed list
		closedList.PushBack(q)

	} // openList non empty

	// Free all the remaining successors in the open list.
	path.freeList(pool, openList)

	return path.optimizePath(pool, closedList), nil
}

// freeArray will free all squares in array from i .. len(squares)-1
func (path *AStarPathing) freeArray(pool *Pool, i int, squares []*Square) {
	if i >= len(squares) {
		return
	}

	for ; i < len(squares); i++ {
		pool.Free(squares[i])
	}
}

// freeList will free every Square in the list l
func (path *AStarPathing) freeList(pool *Pool, l *list.List) {
	// Free all the remaining successors in the open list.
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)
		pool.Free(square)
	}
}

// optimizePath will optimize the path list passed as a parameter. Any culled
// squares are freed from the pool.
//
// A path list will contain duplicates at each _position_. Thus you want to
// iterate over the list and remove duplicates at each _position_ leaving the
// square with the least F in the path list.
// For F ties only one is chosen.
func (path *AStarPathing) optimizePath(pool *Pool, l *list.List) *list.List {
	var m map[int]*Square

	m = make(map[int]*Square)
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)

		p, ok := m[square.Position]

		if !ok {
			m[square.Position] = square
		} else {
			if p.F <= square.F {
				pool.Free(square)
			} else {
				m[square.Position] = square
				pool.Free(p)
			}
		}

	}
	result := list.New()

	length := len(m)
	for i := 0; i < length; i++ {
		result.PushBack(m[i])
	}

	return result
}

//

// skipSuccessor will scan list l and if the list l contains an element with a smaller
// F than the successor at the same position, returns TRUE.
func (path *AStarPathing) skipSuccessor(successor *Square, l *list.List) bool {
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)

		if square.Position == successor.Position && square.F <= successor.F {
			return true
		}
	}

	return false
}

// constructSuccessor will construct 8 successors with parent q from the Pool pool.
func (path *AStarPathing) constructSuccessor(pool *Pool, q *Square) []*Square {
	// The successors are the adjoining squares with the source S
	// in the middle. See below. It moves clockwise. We index
	// from zero so index 0 is 1 below.
	// 1  2  3
	// 8  S  4
	// 7  6  5
	successors := pool.Squares(8)

	// 1
	successors[0].Locus.X = q.Locus.X - 1
	successors[0].Locus.Y = q.Locus.Y - 1
	successors[0].Parent = q
	successors[0].Position = q.Position + 1

	// 2
	successors[1].Locus.X = q.Locus.X - 1
	successors[1].Parent = q
	successors[1].Position = q.Position + 1

	// 3
	successors[2].Locus.X = q.Locus.X - 1
	successors[2].Locus.Y = q.Locus.Y + 1
	successors[2].Parent = q
	successors[2].Position = q.Position + 1

	// 4
	successors[3].Locus.Y = q.Locus.Y + 1
	successors[3].Parent = q
	successors[3].Position = q.Position + 1

	// 5
	successors[4].Locus.X = q.Locus.X + 1
	successors[4].Locus.Y = q.Locus.Y + 1
	successors[4].Parent = q
	successors[4].Position = q.Position + 1

	// 6
	successors[5].Locus.X = q.Locus.X + 1
	successors[5].Parent = q
	successors[5].Position = q.Position + 1

	// 7
	successors[6].Locus.X = q.Locus.X + 1
	successors[6].Locus.Y = q.Locus.Y - 1
	successors[6].Parent = q
	successors[6].Position = q.Position + 1

	// 8
	successors[7].Locus.Y = q.Locus.Y - 1
	successors[7].Parent = q
	successors[7].Position = q.Position + 1

	return successors
}

// leastF returns the Square with the least F within list l
// AND remove that Square from list l.
// Returns nil if no item exists.
func (path *AStarPathing) leastF(l *list.List) *Square {

	var leastSquare *Square
	var leastSquareE *list.Element
	for e := l.Front(); e != nil; e = e.Next() {
		square := e.Value.(*Square)
		if leastSquare == nil || square.F < leastSquare.F {
			leastSquare = square
			leastSquareE = e
		}
	}

	l.Remove(leastSquareE)
	return leastSquare
}
<|endoftext|>


<commit_before>package authenticator

import (
	"net/http"
	"net/url"
	"testing"

	"github.com/abbot/go-http-auth"
	"github.com/fxnn/gone/authenticator/bruteblocker"
)

func TestMiddlewareHandler_noAuthentication(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithRequestAuthAndSessionAuth(requestAuth, sessionAuth)
	var delegate = http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {})
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sut.MiddlewareHandler(delegate).ServeHTTP(rsp, req)

	if requestAuth.IsAuthenticated(req) || sessionAuth.IsAuthenticated(req) {
		t.Fatalf("No authentication expected")
	}
}

func TestMiddlewareHandler_copiesUserId(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithRequestAuthAndSessionAuth(requestAuth, sessionAuth)
	var delegate = http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {})
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sessionAuth.SetUserID(rsp, req, "Aladdin")
	sut.MiddlewareHandler(delegate).ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected Aladdin to be authenticated, but was '%v'", userID)
	}
}

func TestLoginHandler_rightCredentials(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(requestAuth, sessionAuth, aladdinsSecret)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected request user to be Aladdin, but was '%v'", userID)
	}
	if userID := sessionAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected session user to be Aladdin, but was '%v'", userID)
	}
	if rsp.status != http.StatusFound {
		t.Fatalf("Expected status authorized, was '%v'", rsp.status)
	}
}

func TestLoginHandler_wrongCredentials(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthAndBasicAuth(requestAuth, noSecrets)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if rsp.status != http.StatusUnauthorized {
		t.Fatalf("Expected status unauthorized, was '%v'", rsp.status)
	}
	var userID = requestAuth.UserID(req)
	if userID != "" {
		t.Fatalf("Expected no authorization, but got user '%v'", userID)
	}
}

func TestLoginHandler_noLoginGiven(t *testing.T) {
	var sut = sutWithBasicAuth(noSecrets)
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if rsp.status != http.StatusUnauthorized {
		t.Fatalf("Expected status unauthorized, was '%v'", rsp.status)
	}
}

func TestAuthAttemptUser_withoutLogin(t *testing.T) {
	var sut = blankSut()
	var req = blankRequest()

	var actualUser = sut.userAttemptingAuth(req)

	if actualUser != "" {
		t.Fatalf("Expected empty result, was '%v'", actualUser)
	}
}

func TestAuthAttemptUser_successfulLogin(t *testing.T) {
	var sut = blankSut()
	var req = requestWithAladdinUser()

	var actualUser = sut.userAttemptingAuth(req)

	if actualUser != "Aladdin" {
		t.Fatalf("Expected Aladdin, was '%v'", actualUser)
	}
}

func sutWithRequestAuthAndSessionAuth(
	requestAuth Authenticator,
	sessionAuth Authenticator,
) (result HttpBasicAuthenticator) {
	result = blankSut()
	result.requestAuth = requestAuth
	result.sessionAuth = sessionAuth
	return
}

func sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(
	requestAuth Authenticator,
	sessionAuth Authenticator,
	secrets auth.SecretProvider,
) (result HttpBasicAuthenticator) {
	result = sutWithBruteblockerRequestAuthAndBasicAuth(requestAuth, secrets)
	result.sessionAuth = sessionAuth
	return
}

func sutWithBruteblockerRequestAuthAndBasicAuth(
	requestAuth Authenticator,
	secrets auth.SecretProvider,
) (result HttpBasicAuthenticator) {
	result = sutWithBasicAuth(secrets)
	result.requestAuth = requestAuth
	result.bruteBlocker = bruteblocker.New(0, 0, 0, 0, 0)
	return
}

func sutWithBasicAuth(secrets auth.SecretProvider) (result HttpBasicAuthenticator) {
	result = blankSut()
	result.basicAuth = auth.NewBasicAuthenticator("Realm", secrets)
	return
}

func blankSut() HttpBasicAuthenticator {
	return HttpBasicAuthenticator{}
}

func requestWithAladdinUser() (result *http.Request) {
	result = blankRequest()

	// from https://en.wikipedia.org/wiki/Basic_access_authentication
	result.Header.Add("Authorization", "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ==")

	return result
}

func blankRequest() *http.Request {
	return &http.Request{Header: make(http.Header), URL: &url.URL{}}
}
<commit_msg>fix #21: tests for #4 require ssl connection<commit_after>package authenticator

import (
	"net/http"
	"net/url"
	"testing"

	"github.com/abbot/go-http-auth"
	"github.com/fxnn/gone/authenticator/bruteblocker"
)

func TestMiddlewareHandler_noAuthentication(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithRequestAuthAndSessionAuth(requestAuth, sessionAuth)
	var delegate = http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {})
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sut.MiddlewareHandler(delegate).ServeHTTP(rsp, req)

	if requestAuth.IsAuthenticated(req) || sessionAuth.IsAuthenticated(req) {
		t.Fatalf("No authentication expected")
	}
}

func TestMiddlewareHandler_copiesUserId(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithRequestAuthAndSessionAuth(requestAuth, sessionAuth)
	var delegate = http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {})
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sessionAuth.SetUserID(rsp, req, "Aladdin")
	sut.MiddlewareHandler(delegate).ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected Aladdin to be authenticated, but was '%v'", userID)
	}
}

func TestLoginHandler_existingHeader(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(requestAuth, sessionAuth, aladdinsSecret)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.loginRequiresHeader = "Test-Header"
	req.Header.Add("Test-Header", "existing")
	sut.LoginHandler().ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected request user to be Aladdin, but was '%v'", userID)
	}
	if userID := sessionAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected session user to be Aladdin, but was '%v'", userID)
	}
	if rsp.status != http.StatusFound {
		t.Fatalf("Expected status authorized, was '%v'", rsp.status)
	}
}

func TestLoginHandler_missingHeader(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(requestAuth, sessionAuth, aladdinsSecret)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.loginRequiresHeader = "test"
	sut.LoginHandler().ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "" {
		t.Fatalf("Expected request user to be empty, but was '%v'", userID)
	}
	if userID := sessionAuth.UserID(req); userID != "" {
		t.Fatalf("Expected session user to be empty, but was '%v'", userID)
	}
	if rsp.status != http.StatusBadRequest {
		t.Fatalf("Expected status bad request, was '%v'", rsp.status)
	}
}

func TestLoginHandler_rightCredentials(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sessionAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(requestAuth, sessionAuth, aladdinsSecret)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if userID := requestAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected request user to be Aladdin, but was '%v'", userID)
	}
	if userID := sessionAuth.UserID(req); userID != "Aladdin" {
		t.Fatalf("Expected session user to be Aladdin, but was '%v'", userID)
	}
	if rsp.status != http.StatusFound {
		t.Fatalf("Expected status authorized, was '%v'", rsp.status)
	}
}

func TestLoginHandler_wrongCredentials(t *testing.T) {
	var requestAuth = newMockAuthenticator()
	var sut = sutWithBruteblockerRequestAuthAndBasicAuth(requestAuth, noSecrets)
	var req = requestWithAladdinUser()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if rsp.status != http.StatusUnauthorized {
		t.Fatalf("Expected status unauthorized, was '%v'", rsp.status)
	}
	var userID = requestAuth.UserID(req)
	if userID != "" {
		t.Fatalf("Expected no authorization, but got user '%v'", userID)
	}
}

func TestLoginHandler_noLoginGiven(t *testing.T) {
	var sut = sutWithBasicAuth(noSecrets)
	var req = blankRequest()
	var rsp = newMockResponseWriter()

	sut.LoginHandler().ServeHTTP(rsp, req)

	if rsp.status != http.StatusUnauthorized {
		t.Fatalf("Expected status unauthorized, was '%v'", rsp.status)
	}
}

func TestAuthAttemptUser_withoutLogin(t *testing.T) {
	var sut = blankSut()
	var req = blankRequest()

	var actualUser = sut.userAttemptingAuth(req)

	if actualUser != "" {
		t.Fatalf("Expected empty result, was '%v'", actualUser)
	}
}

func TestAuthAttemptUser_successfulLogin(t *testing.T) {
	var sut = blankSut()
	var req = requestWithAladdinUser()

	var actualUser = sut.userAttemptingAuth(req)

	if actualUser != "Aladdin" {
		t.Fatalf("Expected Aladdin, was '%v'", actualUser)
	}
}

func sutWithRequestAuthAndSessionAuth(
	requestAuth Authenticator,
	sessionAuth Authenticator,
) (result HttpBasicAuthenticator) {
	result = blankSut()
	result.requestAuth = requestAuth
	result.sessionAuth = sessionAuth
	return
}

func sutWithBruteblockerRequestAuthSessionAuthAndBasicAuth(
	requestAuth Authenticator,
	sessionAuth Authenticator,
	secrets auth.SecretProvider,
) (result HttpBasicAuthenticator) {
	result = sutWithBruteblockerRequestAuthAndBasicAuth(requestAuth, secrets)
	result.sessionAuth = sessionAuth
	return
}

func sutWithBruteblockerRequestAuthAndBasicAuth(
	requestAuth Authenticator,
	secrets auth.SecretProvider,
) (result HttpBasicAuthenticator) {
	result = sutWithBasicAuth(secrets)
	result.requestAuth = requestAuth
	result.bruteBlocker = bruteblocker.New(0, 0, 0, 0, 0)
	return
}

func sutWithBasicAuth(secrets auth.SecretProvider) (result HttpBasicAuthenticator) {
	result = blankSut()
	result.basicAuth = auth.NewBasicAuthenticator("Realm", secrets)
	return
}

func blankSut() HttpBasicAuthenticator {
	return HttpBasicAuthenticator{}
}

func requestWithAladdinUser() (result *http.Request) {
	result = blankRequest()

	// from https://en.wikipedia.org/wiki/Basic_access_authentication
	result.Header.Add("Authorization", "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ==")

	return result
}

func blankRequest() *http.Request {
	return &http.Request{Header: make(http.Header), URL: &url.URL{}}
}
<|endoftext|>


<commit_before>// +build linux,lxc

package lxc

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/hashicorp/hcl2/hcl"
	ctestutil "github.com/hashicorp/nomad/client/testutil"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
	"github.com/hashicorp/nomad/plugins/shared"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/hashicorp/nomad/testutil"
	"github.com/stretchr/testify/require"
	lxc "gopkg.in/lxc/go-lxc.v2"
)

func TestLXCDriver_Fingerprint(t *testing.T) {
	t.Parallel()
	requireLXC(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	harness := drivers.NewDriverHarness(t, d)

	fingerCh, err := harness.Fingerprint(context.Background())
	require.NoError(err)
	select {
	case finger := <-fingerCh:
		require.Equal(drivers.HealthStateHealthy, finger.Health)
		require.True(finger.Attributes["driver.lxc"].GetBool())
		require.NotEmpty(finger.Attributes["driver.lxc.version"].GetString())
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		require.Fail("timeout receiving fingerprint")
	}
}

func TestLXCDriver_FingerprintNotEnabled(t *testing.T) {
	t.Parallel()
	requireLXC(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = false
	harness := drivers.NewDriverHarness(t, d)

	fingerCh, err := harness.Fingerprint(context.Background())
	require.NoError(err)
	select {
	case finger := <-fingerCh:
		require.Equal(drivers.HealthStateUndetected, finger.Health)
		require.Empty(finger.Attributes["driver.lxc"].GetString())
		require.Empty(finger.Attributes["driver.lxc.version"].GetString())
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		require.Fail("timeout receiving fingerprint")
	}
}

func TestLXCDriver_Start_Wait(t *testing.T) {
	if !testutil.IsTravis() {
		t.Parallel()
	}
	requireLXC(t)
	ctestutil.RequireRoot(t)

	require := require.New(t)

	// prepare test file
	testFileContents := []byte("this should be visible under /mnt/tmp")
	tmpFile, err := ioutil.TempFile("/tmp", "testlxcdriver_start_wait")
	if err != nil {
		t.Fatalf("error writing temp file: %v", err)
	}
	defer os.Remove(tmpFile.Name())
	if _, err := tmpFile.Write(testFileContents); err != nil {
		t.Fatalf("error writing temp file: %v", err)
	}
	if err := tmpFile.Close(); err != nil {
		t.Fatalf("error closing temp file: %v", err)
	}

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	d.config.AllowVolumes = true

	harness := drivers.NewDriverHarness(t, d)
	task := &drivers.TaskConfig{
		ID:   uuid.Generate(),
		Name: "test",
		Resources: &drivers.Resources{
			NomadResources: &structs.Resources{
				CPU:      1,
				MemoryMB: 2,
			},
			LinuxResources: &drivers.LinuxResources{
				CPUShares:        1024,
				MemoryLimitBytes: 2 * 1024,
			},
		},
	}
	taskConfig := map[string]interface{}{
		"template": "/usr/share/lxc/templates/lxc-busybox",
		"volumes":  []string{"/tmp/:mnt/tmp"},
	}
	encodeDriverHelper(require, task, taskConfig)

	cleanup := harness.MkAllocDir(task, false)
	defer cleanup()

	handle, _, err := harness.StartTask(task)
	require.NoError(err)
	require.NotNil(handle)

	lxcHandle, ok := d.tasks.Get(task.ID)
	require.True(ok)

	container := lxcHandle.container

	// Destroy container after test
	defer func() {
		container.Stop()
		container.Destroy()
	}()

	// Test that container is running
	testutil.WaitForResult(func() (bool, error) {
		state := container.State()
		if state == lxc.RUNNING {
			return true, nil
		}
		return false, fmt.Errorf("container in state: %v", state)
	}, func(err error) {
		t.Fatalf("container failed to start: %v", err)
	})

	// Test that directories are mounted in their proper location
	containerName := container.Name()
	for _, mnt := range []string{"alloc", "local", "secrets", "mnt/tmp"} {
		fullpath := filepath.Join(d.lxcPath(), containerName, "rootfs", mnt)
		stat, err := os.Stat(fullpath)
		require.NoError(err)
		require.True(stat.IsDir())
	}

	// Test bind mount volumes exist in container:
	mountedContents, err := exec.Command("lxc-attach",
		"-n", containerName, "--",
		"cat", filepath.Join("/mnt/", tmpFile.Name()),
	).Output()
	require.NoError(err)
	require.Equal(string(testFileContents), string(mountedContents))

	// Test that killing container marks container as stopped
	require.NoError(container.Stop())

	testutil.WaitForResult(func() (bool, error) {
		status, err := d.InspectTask(task.ID)
		if err == nil && status.State == drivers.TaskStateExited {
			return true, nil
		}
		return false, fmt.Errorf("task in state: %v", status.State)
	}, func(err error) {
		t.Fatalf("task was not marked as stopped: %v", err)
	})
}

func TestLXCDriver_Start_Stop(t *testing.T) {
	if !testutil.IsTravis() {
		t.Parallel()
	}
	requireLXC(t)
	ctestutil.RequireRoot(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	d.config.AllowVolumes = true

	harness := drivers.NewDriverHarness(t, d)
	task := &drivers.TaskConfig{
		ID:   uuid.Generate(),
		Name: "test",
		Resources: &drivers.Resources{
			NomadResources: &structs.Resources{
				CPU:      1,
				MemoryMB: 2,
			},
			LinuxResources: &drivers.LinuxResources{
				CPUShares:        1024,
				MemoryLimitBytes: 2 * 1024,
			},
		},
	}
	taskConfig := map[string]interface{}{
		"template": "/usr/share/lxc/templates/lxc-busybox",
	}
	encodeDriverHelper(require, task, taskConfig)

	cleanup := harness.MkAllocDir(task, false)
	defer cleanup()

	handle, _, err := harness.StartTask(task)
	require.NoError(err)
	require.NotNil(handle)

	lxcHandle, ok := d.tasks.Get(task.ID)
	require.True(ok)

	container := lxcHandle.container

	// Destroy container after test
	defer func() {
		container.Stop()
		container.Destroy()
	}()

	// Test that container is running
	testutil.WaitForResult(func() (bool, error) {
		state := container.State()
		if state == lxc.RUNNING {
			return true, nil
		}
		return false, fmt.Errorf("container in state: %v", state)
	}, func(err error) {
		t.Fatalf("container failed to start: %v", err)
	})

	require.NoError(d.StopTask(task.ID, 5*time.Second, "kill"))

	testutil.WaitForResult(func() (bool, error) {
		status, err := d.InspectTask(task.ID)
		if err == nil && status.State == drivers.TaskStateExited {
			return true, nil
		}
		return false, fmt.Errorf("task in state: %v", status.State)
	}, func(err error) {
		t.Fatalf("task was not marked as stopped: %v", err)
	})
}

func requireLXC(t *testing.T) {
	if lxc.Version() == "" {
		t.Skip("skipping, lxc not present")
	}
}

func encodeDriverHelper(require *require.Assertions, task *drivers.TaskConfig, taskConfig map[string]interface{}) {
	evalCtx := &hcl.EvalContext{
		Functions: shared.GetStdlibFuncs(),
	}
	spec, diag := hclspec.Convert(taskConfigSpec)
	require.False(diag.HasErrors())
	taskConfigCtyVal, diag := shared.ParseHclInterface(taskConfig, spec, evalCtx)
	require.False(diag.HasErrors())
	err := task.EncodeDriverConfig(taskConfigCtyVal)
	require.Nil(err)
}
<commit_msg>Fix lxc test panic<commit_after>// +build linux,lxc

package lxc

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"
	"time"

	"github.com/hashicorp/hcl2/hcl"
	ctestutil "github.com/hashicorp/nomad/client/testutil"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
	"github.com/hashicorp/nomad/plugins/shared"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/hashicorp/nomad/testutil"
	"github.com/stretchr/testify/require"
	lxc "gopkg.in/lxc/go-lxc.v2"
)

func TestLXCDriver_Fingerprint(t *testing.T) {
	t.Parallel()
	requireLXC(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	harness := drivers.NewDriverHarness(t, d)

	fingerCh, err := harness.Fingerprint(context.Background())
	require.NoError(err)
	select {
	case finger := <-fingerCh:
		require.Equal(drivers.HealthStateHealthy, finger.Health)
		require.True(finger.Attributes["driver.lxc"].GetBool())
		require.NotEmpty(finger.Attributes["driver.lxc.version"].GetString())
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		require.Fail("timeout receiving fingerprint")
	}
}

func TestLXCDriver_FingerprintNotEnabled(t *testing.T) {
	t.Parallel()
	requireLXC(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = false
	harness := drivers.NewDriverHarness(t, d)

	fingerCh, err := harness.Fingerprint(context.Background())
	require.NoError(err)
	select {
	case finger := <-fingerCh:
		require.Equal(drivers.HealthStateUndetected, finger.Health)
		require.Empty(finger.Attributes["driver.lxc"])
		require.Empty(finger.Attributes["driver.lxc.version"])
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		require.Fail("timeout receiving fingerprint")
	}
}

func TestLXCDriver_Start_Wait(t *testing.T) {
	if !testutil.IsTravis() {
		t.Parallel()
	}
	requireLXC(t)
	ctestutil.RequireRoot(t)

	require := require.New(t)

	// prepare test file
	testFileContents := []byte("this should be visible under /mnt/tmp")
	tmpFile, err := ioutil.TempFile("/tmp", "testlxcdriver_start_wait")
	if err != nil {
		t.Fatalf("error writing temp file: %v", err)
	}
	defer os.Remove(tmpFile.Name())
	if _, err := tmpFile.Write(testFileContents); err != nil {
		t.Fatalf("error writing temp file: %v", err)
	}
	if err := tmpFile.Close(); err != nil {
		t.Fatalf("error closing temp file: %v", err)
	}

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	d.config.AllowVolumes = true

	harness := drivers.NewDriverHarness(t, d)
	task := &drivers.TaskConfig{
		ID:   uuid.Generate(),
		Name: "test",
		Resources: &drivers.Resources{
			NomadResources: &structs.Resources{
				CPU:      1,
				MemoryMB: 2,
			},
			LinuxResources: &drivers.LinuxResources{
				CPUShares:        1024,
				MemoryLimitBytes: 2 * 1024,
			},
		},
	}
	taskConfig := map[string]interface{}{
		"template": "/usr/share/lxc/templates/lxc-busybox",
		"volumes":  []string{"/tmp/:mnt/tmp"},
	}
	encodeDriverHelper(require, task, taskConfig)

	cleanup := harness.MkAllocDir(task, false)
	defer cleanup()

	handle, _, err := harness.StartTask(task)
	require.NoError(err)
	require.NotNil(handle)

	lxcHandle, ok := d.tasks.Get(task.ID)
	require.True(ok)

	container := lxcHandle.container

	// Destroy container after test
	defer func() {
		container.Stop()
		container.Destroy()
	}()

	// Test that container is running
	testutil.WaitForResult(func() (bool, error) {
		state := container.State()
		if state == lxc.RUNNING {
			return true, nil
		}
		return false, fmt.Errorf("container in state: %v", state)
	}, func(err error) {
		t.Fatalf("container failed to start: %v", err)
	})

	// Test that directories are mounted in their proper location
	containerName := container.Name()
	for _, mnt := range []string{"alloc", "local", "secrets", "mnt/tmp"} {
		fullpath := filepath.Join(d.lxcPath(), containerName, "rootfs", mnt)
		stat, err := os.Stat(fullpath)
		require.NoError(err)
		require.True(stat.IsDir())
	}

	// Test bind mount volumes exist in container:
	mountedContents, err := exec.Command("lxc-attach",
		"-n", containerName, "--",
		"cat", filepath.Join("/mnt/", tmpFile.Name()),
	).Output()
	require.NoError(err)
	require.Equal(string(testFileContents), string(mountedContents))

	// Test that killing container marks container as stopped
	require.NoError(container.Stop())

	testutil.WaitForResult(func() (bool, error) {
		status, err := d.InspectTask(task.ID)
		if err == nil && status.State == drivers.TaskStateExited {
			return true, nil
		}
		return false, fmt.Errorf("task in state: %v", status.State)
	}, func(err error) {
		t.Fatalf("task was not marked as stopped: %v", err)
	})
}

func TestLXCDriver_Start_Stop(t *testing.T) {
	if !testutil.IsTravis() {
		t.Parallel()
	}
	requireLXC(t)
	ctestutil.RequireRoot(t)

	require := require.New(t)

	d := NewLXCDriver(testlog.HCLogger(t)).(*Driver)
	d.config.Enabled = true
	d.config.AllowVolumes = true

	harness := drivers.NewDriverHarness(t, d)
	task := &drivers.TaskConfig{
		ID:   uuid.Generate(),
		Name: "test",
		Resources: &drivers.Resources{
			NomadResources: &structs.Resources{
				CPU:      1,
				MemoryMB: 2,
			},
			LinuxResources: &drivers.LinuxResources{
				CPUShares:        1024,
				MemoryLimitBytes: 2 * 1024,
			},
		},
	}
	taskConfig := map[string]interface{}{
		"template": "/usr/share/lxc/templates/lxc-busybox",
	}
	encodeDriverHelper(require, task, taskConfig)

	cleanup := harness.MkAllocDir(task, false)
	defer cleanup()

	handle, _, err := harness.StartTask(task)
	require.NoError(err)
	require.NotNil(handle)

	lxcHandle, ok := d.tasks.Get(task.ID)
	require.True(ok)

	container := lxcHandle.container

	// Destroy container after test
	defer func() {
		container.Stop()
		container.Destroy()
	}()

	// Test that container is running
	testutil.WaitForResult(func() (bool, error) {
		state := container.State()
		if state == lxc.RUNNING {
			return true, nil
		}
		return false, fmt.Errorf("container in state: %v", state)
	}, func(err error) {
		t.Fatalf("container failed to start: %v", err)
	})

	require.NoError(d.StopTask(task.ID, 5*time.Second, "kill"))

	testutil.WaitForResult(func() (bool, error) {
		status, err := d.InspectTask(task.ID)
		if err == nil && status.State == drivers.TaskStateExited {
			return true, nil
		}
		return false, fmt.Errorf("task in state: %v", status.State)
	}, func(err error) {
		t.Fatalf("task was not marked as stopped: %v", err)
	})
}

func requireLXC(t *testing.T) {
	if lxc.Version() == "" {
		t.Skip("skipping, lxc not present")
	}
}

func encodeDriverHelper(require *require.Assertions, task *drivers.TaskConfig, taskConfig map[string]interface{}) {
	evalCtx := &hcl.EvalContext{
		Functions: shared.GetStdlibFuncs(),
	}
	spec, diag := hclspec.Convert(taskConfigSpec)
	require.False(diag.HasErrors())
	taskConfigCtyVal, diag := shared.ParseHclInterface(taskConfig, spec, evalCtx)
	require.False(diag.HasErrors())
	err := task.EncodeDriverConfig(taskConfigCtyVal)
	require.Nil(err)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/pagination"
	"github.com/rackspace/gophercloud/rackspace"
	"github.com/rackspace/gophercloud/rackspace/lb/v1/lbs"
	"github.com/rackspace/gophercloud/rackspace/lb/v1/nodes"
)

func waitForReady(client *gophercloud.ServiceClient, loadBalancerID int) {

	// Ensure the load balancer is ready
	gophercloud.WaitFor(60, func() (bool, error) {
		lb, err := lbs.Get(client, loadBalancerID).Extract()
		if err != nil {
			return false, err
		}

		if lb.Status != lbs.ACTIVE {
			// It has not yet reached ACTIVE
			return false, nil
		}

		// It has reached ACTIVE
		return true, nil
	})
}

func findNodeByIPPort(
	client *gophercloud.ServiceClient,
	loadBalancerID int,
	address string,
	port int,
) *nodes.Node {

	// nil until found
	var found *nodes.Node

	pager := nodes.List(client, loadBalancerID, nil)

	pager.EachPage(func(page pagination.Page) (bool, error) {
		lbNodes, err := nodes.ExtractNodes(page)
		if err != nil {
			log.Fatalf("Error during paging load balancer: %v", err)
		}

		for _, trialNode := range lbNodes {
			if trialNode.Address == address && trialNode.Port == port {
				found = &trialNode
				return false, nil
			}

		}

		return true, nil
	})

	return found
}

func getIP(ipPtr *string) (string, error) {

	// IP overriding order will be:
	//  - The -ip flag takes primary precedence
	//  - Service Net Environment variable
	//  - Public Net Environment variable
	//  - Gleaning a 10 dot out of the network interfaces (typically eth1)
	//  - eth0

	serviceNetIPv4 := os.Getenv("RAX_SERVICENET_IPV4")
	publicNetIPv4 := os.Getenv("RAX_PUBLICNET_IPV4")

	switch {
	case *ipPtr != "":
		return *ipPtr, nil
	case serviceNetIPv4 != "":
		return serviceNetIPv4, nil
	case publicNetIPv4 != "":
		return publicNetIPv4, nil
	}

	addrs, err := net.InterfaceAddrs()

	if err != nil {
		return "", err
	}

	for _, addr := range addrs {
		cidr := addr.String()
		ip := strings.Split(cidr, "/")[0]

		if strings.HasPrefix(ip, "10.") {
			return ip, nil
		}
	}

	// Find eth0
	eth0, err := net.InterfaceByName("eth0")
	if err != nil {
		return "", fmt.Errorf("Trouble finding eth0: %v", err)
	}

	addrs, err = eth0.Addrs()
	if err != nil {
		return "", err
	}

	for _, addr := range addrs {
		cidr := addr.String()
		ip := strings.Split(cidr, "/")[0]

		fmt.Println(ip)
	}

	return "", errors.New("Unable to determine an IP for load balancing")

}

func main() {
	var err error

	disabledPtr := flag.Bool("disable", false, "Disable the node on the load balancer")

	//NOTE: peekaboo allows setting the IP by using
	//        - environment variables: RAX_SERVICENET_IPV4 or RAX_PUBLICNET_IPV4
	//        - finding an ip on the system starting with 10. (service net)
	//        - locating the eth0 interface
	//        - providing an ip as a flag is fine too and will take precedence
	ipPtr := flag.String("ip", "", "IP address to register/deregister on the load balancer")
	flag.Parse()

	username := os.Getenv("OS_USERNAME")
	APIKey := os.Getenv("OS_PASSWORD")
	region := os.Getenv("OS_REGION_NAME")

	// These get converted into integers later
	strLoadBalancerID := os.Getenv("LOAD_BALANCER_ID")
	strAppPort := os.Getenv("APP_PORT")

	/**
	 * Retrieve port for load balancer's node, defaulting to 80
	 */
	var nodePort = 80

	if strAppPort == "" {
		log.Printf("$APP_PORT not set, defaulting to 80")
	} else {
		nodePort, err = strconv.Atoi(strAppPort)
		if err != nil {
			log.Fatalf("Unable to parse integer from $APP_PORT: %v\n", strAppPort)
		}
	}

	/**
	 * Determine the IP Address for the load balancer's node
	 */
	nodeAddress, err := getIP(ipPtr)
	if err != nil {
		log.Fatalln(err)
	}

	log.Printf("Determined IP: %v", nodeAddress)
	log.Printf("Determined Port: %v", nodePort)

	/**
	 * Retrive Load Balancer ID
	 */
	loadBalancerID, err := strconv.Atoi(strLoadBalancerID)
	if err != nil {
		log.Fatalf("$LOAD_BALANCER_ID not an integer: %v\n", loadBalancerID)
	}

	provider, err := rackspace.AuthenticatedClient(gophercloud.AuthOptions{
		Username: username,
		APIKey:   APIKey,
	})

	if err != nil {
		log.Fatalf("Trouble authenticating to Rackspace: %v\n", err)
	}

	client, err := rackspace.NewLBV1(provider, gophercloud.EndpointOpts{
		Region: region,
	})

	if err != nil {
		log.Fatalf("Creating load balancer client in %v failed: %v\n", region, err)
	}

	log.Println("Client ready")

	waitForReady(client, loadBalancerID)

	node := findNodeByIPPort(client, loadBalancerID, nodeAddress, nodePort)

	condition := nodes.ENABLED
	if *disabledPtr {
		//TODO: Watch the interface on the right process/container to determine
		//      when connections have dropped, and set to DISABLED
		condition = nodes.DRAINING
	}

	log.Printf("Telling %v on port %v to be %v\n", nodeAddress, nodePort, condition)

	if node != nil {
		log.Printf("Found existing node %v", *node)

		opts := nodes.UpdateOpts{
			Condition: condition,
		}

		updateResult := nodes.Update(client, loadBalancerID, node.ID, opts)
		log.Printf("Update result: %v\n", updateResult)

	} else {
		log.Printf("Creating new node")
		opts := nodes.CreateOpts{
			nodes.CreateOpt{
				Address:   nodeAddress,
				Port:      nodePort,
				Condition: condition,
			},
		}

		nodeList := nodes.Create(client, loadBalancerID, opts)
		log.Printf("Node made, total list: %v\n", nodeList)
	}

}
<commit_msg>Sanity checks on node pointers, updates<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	"github.com/rackspace/gophercloud"
	"github.com/rackspace/gophercloud/pagination"
	"github.com/rackspace/gophercloud/rackspace"
	"github.com/rackspace/gophercloud/rackspace/lb/v1/lbs"
	"github.com/rackspace/gophercloud/rackspace/lb/v1/nodes"
)

func waitForReady(client *gophercloud.ServiceClient, loadBalancerID int) {

	// Ensure the load balancer is ready
	gophercloud.WaitFor(60, func() (bool, error) {
		lb, err := lbs.Get(client, loadBalancerID).Extract()
		if err != nil {
			return false, err
		}

		if lb.Status != lbs.ACTIVE {
			// It has not yet reached ACTIVE
			return false, nil
		}

		// It has reached ACTIVE
		return true, nil
	})
}

func findNodeByIPPort(
	client *gophercloud.ServiceClient,
	loadBalancerID int,
	address string,
	port int,
) *nodes.Node {

	// nil until found
	var found *nodes.Node

	pager := nodes.List(client, loadBalancerID, nil)

	pager.EachPage(func(page pagination.Page) (bool, error) {
		lbNodes, err := nodes.ExtractNodes(page)
		if err != nil {
			log.Fatalf("Error during paging load balancer: %v", err)
		}

		for _, trialNode := range lbNodes {
			if trialNode.Address == address && trialNode.Port == port {
				found = &trialNode
				return false, nil
			}

		}

		return true, nil
	})

	return found
}

func getIP(ipPtr *string) (string, error) {

	// IP overriding order will be:
	//  - The -ip flag takes primary precedence
	//  - Service Net Environment variable
	//  - Public Net Environment variable
	//  - Gleaning a 10 dot out of the network interfaces (typically eth1)
	//  - eth0

	serviceNetIPv4 := os.Getenv("RAX_SERVICENET_IPV4")
	publicNetIPv4 := os.Getenv("RAX_PUBLICNET_IPV4")

	switch {
	case *ipPtr != "":
		return *ipPtr, nil
	case serviceNetIPv4 != "":
		return serviceNetIPv4, nil
	case publicNetIPv4 != "":
		return publicNetIPv4, nil
	}

	addrs, err := net.InterfaceAddrs()

	if err != nil {
		return "", err
	}

	for _, addr := range addrs {
		cidr := addr.String()
		ip := strings.Split(cidr, "/")[0]

		if strings.HasPrefix(ip, "10.") {
			return ip, nil
		}
	}

	// Find eth0
	eth0, err := net.InterfaceByName("eth0")
	if err != nil {
		return "", fmt.Errorf("Trouble finding eth0: %v", err)
	}

	addrs, err = eth0.Addrs()
	if err != nil {
		return "", err
	}

	for _, addr := range addrs {
		cidr := addr.String()
		ip := strings.Split(cidr, "/")[0]

		fmt.Println(ip)
	}

	return "", errors.New("Unable to determine an IP for load balancing")

}

func main() {
	var err error

	disabledPtr := flag.Bool("disable", false, "Disable the node on the load balancer")

	//NOTE: peekaboo allows setting the IP by using
	//        - environment variables: RAX_SERVICENET_IPV4 or RAX_PUBLICNET_IPV4
	//        - finding an ip on the system starting with 10. (service net)
	//        - locating the eth0 interface
	//        - providing an ip as a flag is fine too and will take precedence
	ipPtr := flag.String("ip", "", "IP address to register/deregister on the load balancer")
	flag.Parse()

	username := os.Getenv("OS_USERNAME")
	APIKey := os.Getenv("OS_PASSWORD")
	region := os.Getenv("OS_REGION_NAME")

	// These get converted into integers later
	strLoadBalancerID := os.Getenv("LOAD_BALANCER_ID")
	strAppPort := os.Getenv("APP_PORT")

	/**
	 * Retrieve port for load balancer's node, defaulting to 80
	 */
	var nodePort = 80

	if strAppPort == "" {
		log.Printf("$APP_PORT not set, defaulting to 80")
	} else {
		nodePort, err = strconv.Atoi(strAppPort)
		if err != nil {
			log.Fatalf("Unable to parse integer from $APP_PORT: %v\n", strAppPort)
		}
	}

	/**
	 * Determine the IP Address for the load balancer's node
	 */
	nodeAddress, err := getIP(ipPtr)
	if err != nil {
		log.Fatalln(err)
	}

	log.Printf("Determined IP: %v", nodeAddress)
	log.Printf("Determined Port: %v", nodePort)

	/**
	 * Retrive Load Balancer ID
	 */
	loadBalancerID, err := strconv.Atoi(strLoadBalancerID)
	if err != nil {
		log.Fatalf("$LOAD_BALANCER_ID not an integer: %v\n", loadBalancerID)
	}

	provider, err := rackspace.AuthenticatedClient(gophercloud.AuthOptions{
		Username: username,
		APIKey:   APIKey,
	})

	if err != nil {
		log.Fatalf("Trouble authenticating to Rackspace: %v\n", err)
	}

	client, err := rackspace.NewLBV1(provider, gophercloud.EndpointOpts{
		Region: region,
	})

	if err != nil {
		log.Fatalf("Creating load balancer client in %v failed: %v\n", region, err)
	}

	log.Println("Client ready")

	waitForReady(client, loadBalancerID)

	nodePtr := findNodeByIPPort(client, loadBalancerID, nodeAddress, nodePort)

	condition := nodes.ENABLED
	if *disabledPtr {
		//TODO: Watch the interface on the right process/container to determine
		//      when connections have dropped, and set to DISABLED
		condition = nodes.DRAINING
	}

	log.Printf("Telling %v on port %v to be %v\n", nodeAddress, nodePort, condition)

	if nodePtr != nil {
		log.Printf("Found existing node %v", *nodePtr)

		opts := nodes.UpdateOpts{
			Condition: condition,
		}

		updateResult := nodes.Update(client, loadBalancerID, nodePtr.ID, opts)
		err = updateResult.ExtractErr()
		if err != nil {
			log.Panicf("Updating node failed: %v", err)
		}

	} else {
		log.Printf("Creating new node")
		opts := nodes.CreateOpts{
			nodes.CreateOpt{
				Address:   nodeAddress,
				Port:      nodePort,
				Condition: condition,
			},
		}

		nodePager := nodes.Create(client, loadBalancerID, opts)
		nodeList, err := nodePager.ExtractNodes()
		if err != nil || len(nodeList) != 1 {
			log.Panicf("Something went terribly wrong on node creation: %v\n", nodeList)
		}
		nodePtr = &nodeList[0]
	}

	waitForReady(client, loadBalancerID)

	// After update, get the version of the node Rackspace has
	result := nodes.Get(client, loadBalancerID, nodePtr.ID)
	nodePtr, err = result.Extract()
	if err != nil {
		log.Panicln(err)
	}

	log.Printf("Updated state for node: %v\n", *nodePtr)

}
<|endoftext|>