<commit_before>package test

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/anacrolix/missinggo/v2/filecache"
	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/internal/testutil"
	"github.com/anacrolix/torrent/storage"
	"golang.org/x/time/rate"

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

type testClientTransferParams struct {
	Responsive                 bool
	Readahead                  int64
	SetReadahead               bool
	ExportClientStatus         bool
	LeecherStorage             func(string) storage.ClientImplCloser
	SeederStorage              func(string) storage.ClientImplCloser
	SeederUploadRateLimiter    *rate.Limiter
	LeecherDownloadRateLimiter *rate.Limiter
	ConfigureSeeder            ConfigureClient
	ConfigureLeecher           ConfigureClient
}

func assertReadAllGreeting(t *testing.T, r io.ReadSeeker) {
	pos, err := r.Seek(0, io.SeekStart)
	assert.NoError(t, err)
	assert.EqualValues(t, 0, pos)
	_greeting, err := ioutil.ReadAll(r)
	assert.NoError(t, err)
	assert.EqualValues(t, testutil.GreetingFileContents, string(_greeting))
}

// Creates a seeder and a leecher, and ensures the data transfers when a read
// is attempted on the leecher.
func testClientTransfer(t *testing.T, ps testClientTransferParams) {
	greetingTempDir, mi := testutil.GreetingTestTorrent()
	defer os.RemoveAll(greetingTempDir)
	// Create seeder and a Torrent.
	cfg := torrent.TestingConfig()
	cfg.Seed = true
	if ps.SeederUploadRateLimiter != nil {
		cfg.UploadRateLimiter = ps.SeederUploadRateLimiter
	}
	// cfg.ListenAddr = "localhost:4000"
	if ps.SeederStorage != nil {
		storage := ps.SeederStorage(greetingTempDir)
		defer storage.Close()
		cfg.DefaultStorage = storage
	} else {
		cfg.DataDir = greetingTempDir
	}
	if ps.ConfigureSeeder.Config != nil {
		ps.ConfigureSeeder.Config(cfg)
	}
	seeder, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	if ps.ConfigureSeeder.Client != nil {
		ps.ConfigureSeeder.Client(seeder)
	}
	if ps.ExportClientStatus {
		defer testutil.ExportStatusWriter(seeder, "s")()
	}
	seederTorrent, _, _ := seeder.AddTorrentSpec(torrent.TorrentSpecFromMetaInfo(mi))
	// Run a Stats right after Closing the Client. This will trigger the Stats
	// panic in #214 caused by RemoteAddr on Closed uTP sockets.
	defer seederTorrent.Stats()
	defer seeder.Close()
	seederTorrent.VerifyData()
	// Create leecher and a Torrent.
	leecherDataDir, err := ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(leecherDataDir)
	cfg = torrent.TestingConfig()
	if ps.LeecherStorage == nil {
		cfg.DataDir = leecherDataDir
	} else {
		storage := ps.LeecherStorage(leecherDataDir)
		defer storage.Close()
		cfg.DefaultStorage = storage
	}
	if ps.LeecherDownloadRateLimiter != nil {
		cfg.DownloadRateLimiter = ps.LeecherDownloadRateLimiter
	}
	cfg.Seed = false
	if ps.ConfigureLeecher.Config != nil {
		ps.ConfigureLeecher.Config(cfg)
	}
	leecher, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecher.Close()
	if ps.ConfigureLeecher.Client != nil {
		ps.ConfigureLeecher.Client(leecher)
	}
	if ps.ExportClientStatus {
		defer testutil.ExportStatusWriter(leecher, "l")()
	}
	leecherTorrent, new, err := leecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 2
		return
	}())
	require.NoError(t, err)
	assert.True(t, new)

	//// This was used when observing coalescing of piece state changes.
	//logPieceStateChanges(leecherTorrent)

	// Now do some things with leecher and seeder.
	added := leecherTorrent.AddClientPeer(seeder)
	// The Torrent should not be interested in obtaining peers, so the one we
	// just added should be the only one.
	assert.False(t, leecherTorrent.Seeding())
	assert.EqualValues(t, added, leecherTorrent.Stats().PendingPeers)
	r := leecherTorrent.NewReader()
	defer r.Close()
	if ps.Responsive {
		r.SetResponsive()
	}
	if ps.SetReadahead {
		r.SetReadahead(ps.Readahead)
	}
	assertReadAllGreeting(t, r)
	assert.NotEmpty(t, seederTorrent.PeerConns())
	leecherPeerConns := leecherTorrent.PeerConns()
	assert.NotEmpty(t, leecherPeerConns)
	foundSeeder := false
	for _, pc := range leecherPeerConns {
		completed := pc.PeerPieces().Len()
		t.Logf("peer conn %v has %v completed pieces", pc, completed)
		if completed == leecherTorrent.Info().NumPieces() {
			foundSeeder = true
		}
	}
	if !foundSeeder {
		t.Errorf("didn't find seeder amongst leecher peer conns")
	}

	seederStats := seederTorrent.Stats()
	assert.True(t, 13 <= seederStats.BytesWrittenData.Int64())
	assert.True(t, 8 <= seederStats.ChunksWritten.Int64())

	leecherStats := leecherTorrent.Stats()
	assert.True(t, 13 <= leecherStats.BytesReadData.Int64())
	assert.True(t, 8 <= leecherStats.ChunksRead.Int64())

	// Try reading through again for the cases where the torrent data size
	// exceeds the size of the cache.
	assertReadAllGreeting(t, r)
}

type fileCacheClientStorageFactoryParams struct {
	Capacity    int64
	SetCapacity bool
	Wrapper     func(*filecache.Cache) storage.ClientImplCloser
}

func newFileCacheClientStorageFactory(ps fileCacheClientStorageFactoryParams) storageFactory {
	return func(dataDir string) storage.ClientImplCloser {
		fc, err := filecache.NewCache(dataDir)
		if err != nil {
			panic(err)
		}
		if ps.SetCapacity {
			fc.SetCapacity(ps.Capacity)
		}
		return ps.Wrapper(fc)
	}
}

type storageFactory func(string) storage.ClientImplCloser

func TestClientTransferDefault(t *testing.T) {
	testClientTransfer(t, testClientTransferParams{
		ExportClientStatus: true,
		LeecherStorage: newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			Wrapper: fileCachePieceResourceStorage,
		}),
	})
}

func TestClientTransferRateLimitedUpload(t *testing.T) {
	started := time.Now()
	testClientTransfer(t, testClientTransferParams{
		// We are uploading 13 bytes (the length of the greeting torrent). The
		// chunks are 2 bytes in length. Then the smallest burst we can run
		// with is 2. Time taken is (13-burst)/rate.
		SeederUploadRateLimiter: rate.NewLimiter(11, 2),
		ExportClientStatus:      true,
	})
	require.True(t, time.Since(started) > time.Second)
}

func TestClientTransferRateLimitedDownload(t *testing.T) {
	testClientTransfer(t, testClientTransferParams{
		LeecherDownloadRateLimiter: rate.NewLimiter(512, 512),
	})
}

func fileCachePieceResourceStorage(fc *filecache.Cache) storage.ClientImplCloser {
	return storage.NewResourcePieces(fc.AsResourceProvider())
}

func testClientTransferSmallCache(t *testing.T, setReadahead bool, readahead int64) {
	testClientTransfer(t, testClientTransferParams{
		LeecherStorage: newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			SetCapacity: true,
			// Going below the piece length means it can't complete a piece so
			// that it can be hashed.
			Capacity: 5,
			Wrapper:  fileCachePieceResourceStorage,
		}),
		SetReadahead: setReadahead,
		// Can't readahead too far or the cache will thrash and drop data we
		// thought we had.
		Readahead:          readahead,
		ExportClientStatus: true,

		// These tests don't work well with more than 1 connection to the seeder.
		ConfigureLeecher: ConfigureClient{
			Config: func(cfg *torrent.ClientConfig) {
				cfg.DropDuplicatePeerIds = true
				//cfg.DisableIPv6 = true
				//cfg.DisableUTP = true
			},
		},
	})
}

func TestClientTransferSmallCachePieceSizedReadahead(t *testing.T) {
	testClientTransferSmallCache(t, true, 5)
}

func TestClientTransferSmallCacheLargeReadahead(t *testing.T) {
	testClientTransferSmallCache(t, true, 15)
}

func TestClientTransferSmallCacheDefaultReadahead(t *testing.T) {
	testClientTransferSmallCache(t, false, -1)
}

func TestClientTransferVarious(t *testing.T) {
	// Leecher storage
	for _, ls := range []struct {
		name string
		f    storageFactory
	}{
		{"Filecache", newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			Wrapper: fileCachePieceResourceStorage,
		})},
		{"Boltdb", storage.NewBoltDB},
	} {
		t.Run(fmt.Sprintf("LeecherStorage=%s", ls.name), func(t *testing.T) {
			// Seeder storage
			for _, ss := range []struct {
				name string
				f    func(string) storage.ClientImplCloser
			}{
				{"File", storage.NewFile},
				{"Mmap", storage.NewMMap},
			} {
				t.Run(fmt.Sprintf("%sSeederStorage", ss.name), func(t *testing.T) {
					for _, responsive := range []bool{false, true} {
						t.Run(fmt.Sprintf("Responsive=%v", responsive), func(t *testing.T) {
							t.Run("NoReadahead", func(t *testing.T) {
								testClientTransfer(t, testClientTransferParams{
									Responsive:     responsive,
									SeederStorage:  ss.f,
									LeecherStorage: ls.f,
								})
							})
							for _, readahead := range []int64{-1, 0, 1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15, 20} {
								t.Run(fmt.Sprintf("readahead=%v", readahead), func(t *testing.T) {
									testClientTransfer(t, testClientTransferParams{
										SeederStorage:  ss.f,
										Responsive:     responsive,
										SetReadahead:   true,
										Readahead:      readahead,
										LeecherStorage: ls.f,
									})
								})
							}
						})
					}
				})
			}
		})
	}
}

// Check that after completing leeching, a leecher transitions to a seeding
// correctly. Connected in a chain like so: Seeder <-> Leecher <-> LeecherLeecher.
func TestSeedAfterDownloading(t *testing.T) {
	greetingTempDir, mi := testutil.GreetingTestTorrent()
	defer os.RemoveAll(greetingTempDir)

	cfg := torrent.TestingConfig()
	cfg.Seed = true
	cfg.DataDir = greetingTempDir
	seeder, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer seeder.Close()
	defer testutil.ExportStatusWriter(seeder, "s")()
	seederTorrent, ok, err := seeder.AddTorrentSpec(torrent.TorrentSpecFromMetaInfo(mi))
	require.NoError(t, err)
	assert.True(t, ok)
	seederTorrent.VerifyData()

	cfg = torrent.TestingConfig()
	cfg.Seed = true
	cfg.DataDir, err = ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(cfg.DataDir)
	leecher, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecher.Close()
	defer testutil.ExportStatusWriter(leecher, "l")()

	cfg = torrent.TestingConfig()
	cfg.Seed = false
	cfg.DataDir, err = ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(cfg.DataDir)
	leecherLeecher, _ := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecherLeecher.Close()
	defer testutil.ExportStatusWriter(leecherLeecher, "ll")()
	leecherGreeting, ok, err := leecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 2
		return
	}())
	require.NoError(t, err)
	assert.True(t, ok)
	llg, ok, err := leecherLeecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 3
		return
	}())
	require.NoError(t, err)
	assert.True(t, ok)
	// Simultaneously DownloadAll in Leecher, and read the contents
	// consecutively in LeecherLeecher. This non-deterministically triggered a
	// case where the leecher wouldn't unchoke the LeecherLeecher.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		r := llg.NewReader()
		defer r.Close()
		b, err := ioutil.ReadAll(r)
		require.NoError(t, err)
		assert.EqualValues(t, testutil.GreetingFileContents, b)
	}()
	done := make(chan struct{})
	defer close(done)
	go leecherGreeting.AddClientPeer(seeder)
	go leecherGreeting.AddClientPeer(leecherLeecher)
	wg.Add(1)
	go func() {
		defer wg.Done()
		leecherGreeting.DownloadAll()
		leecher.WaitAll()
	}()
	wg.Wait()
}

type ConfigureClient struct {
	Config func(*torrent.ClientConfig)
	Client func(*torrent.Client)
}
<commit_msg>Add tests for https://github.com/anacrolix/torrent/issues/388<commit_after>package test

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/anacrolix/missinggo/v2/filecache"
	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/internal/testutil"
	"github.com/anacrolix/torrent/storage"
	"golang.org/x/time/rate"

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

type testClientTransferParams struct {
	Responsive                 bool
	Readahead                  int64
	SetReadahead               bool
	ExportClientStatus         bool
	LeecherStorage             func(string) storage.ClientImplCloser
	SeederStorage              func(string) storage.ClientImplCloser
	SeederUploadRateLimiter    *rate.Limiter
	LeecherDownloadRateLimiter *rate.Limiter
	ConfigureSeeder            ConfigureClient
	ConfigureLeecher           ConfigureClient

	LeecherStartsWithoutMetadata bool
}

func assertReadAllGreeting(t *testing.T, r io.ReadSeeker) {
	pos, err := r.Seek(0, io.SeekStart)
	assert.NoError(t, err)
	assert.EqualValues(t, 0, pos)
	_greeting, err := ioutil.ReadAll(r)
	assert.NoError(t, err)
	assert.EqualValues(t, testutil.GreetingFileContents, string(_greeting))
}

// Creates a seeder and a leecher, and ensures the data transfers when a read
// is attempted on the leecher.
func testClientTransfer(t *testing.T, ps testClientTransferParams) {
	greetingTempDir, mi := testutil.GreetingTestTorrent()
	defer os.RemoveAll(greetingTempDir)
	// Create seeder and a Torrent.
	cfg := torrent.TestingConfig()
	cfg.Seed = true
	if ps.SeederUploadRateLimiter != nil {
		cfg.UploadRateLimiter = ps.SeederUploadRateLimiter
	}
	// cfg.ListenAddr = "localhost:4000"
	if ps.SeederStorage != nil {
		storage := ps.SeederStorage(greetingTempDir)
		defer storage.Close()
		cfg.DefaultStorage = storage
	} else {
		cfg.DataDir = greetingTempDir
	}
	if ps.ConfigureSeeder.Config != nil {
		ps.ConfigureSeeder.Config(cfg)
	}
	seeder, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	if ps.ConfigureSeeder.Client != nil {
		ps.ConfigureSeeder.Client(seeder)
	}
	if ps.ExportClientStatus {
		defer testutil.ExportStatusWriter(seeder, "s")()
	}
	seederTorrent, _, _ := seeder.AddTorrentSpec(torrent.TorrentSpecFromMetaInfo(mi))
	// Run a Stats right after Closing the Client. This will trigger the Stats
	// panic in #214 caused by RemoteAddr on Closed uTP sockets.
	defer seederTorrent.Stats()
	defer seeder.Close()
	seederTorrent.VerifyData()
	// Create leecher and a Torrent.
	leecherDataDir, err := ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(leecherDataDir)
	cfg = torrent.TestingConfig()
	if ps.LeecherStorage == nil {
		cfg.DataDir = leecherDataDir
	} else {
		storage := ps.LeecherStorage(leecherDataDir)
		defer storage.Close()
		cfg.DefaultStorage = storage
	}
	if ps.LeecherDownloadRateLimiter != nil {
		cfg.DownloadRateLimiter = ps.LeecherDownloadRateLimiter
	}
	cfg.Seed = false
	if ps.ConfigureLeecher.Config != nil {
		ps.ConfigureLeecher.Config(cfg)
	}
	leecher, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecher.Close()
	if ps.ConfigureLeecher.Client != nil {
		ps.ConfigureLeecher.Client(leecher)
	}
	if ps.ExportClientStatus {
		defer testutil.ExportStatusWriter(leecher, "l")()
	}
	leecherTorrent, new, err := leecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 2
		if ps.LeecherStartsWithoutMetadata {
			ret.InfoBytes = nil
		}
		return
	}())
	require.NoError(t, err)
	assert.True(t, new)

	//// This was used when observing coalescing of piece state changes.
	//logPieceStateChanges(leecherTorrent)

	// Now do some things with leecher and seeder.
	added := leecherTorrent.AddClientPeer(seeder)
	assert.False(t, leecherTorrent.Seeding())
	// The leecher will use peers immediately if it doesn't have the metadata. Otherwise, they
	// should be sitting idle until we demand data.
	if !ps.LeecherStartsWithoutMetadata {
		assert.EqualValues(t, added, leecherTorrent.Stats().PendingPeers)
	}
	if ps.LeecherStartsWithoutMetadata {
		<-leecherTorrent.GotInfo()
	}
	r := leecherTorrent.NewReader()
	defer r.Close()
	go leecherTorrent.SetInfoBytes(mi.InfoBytes)
	if ps.Responsive {
		r.SetResponsive()
	}
	if ps.SetReadahead {
		r.SetReadahead(ps.Readahead)
	}
	assertReadAllGreeting(t, r)
	assert.NotEmpty(t, seederTorrent.PeerConns())
	leecherPeerConns := leecherTorrent.PeerConns()
	assert.NotEmpty(t, leecherPeerConns)
	foundSeeder := false
	for _, pc := range leecherPeerConns {
		completed := pc.PeerPieces().Len()
		t.Logf("peer conn %v has %v completed pieces", pc, completed)
		if completed == leecherTorrent.Info().NumPieces() {
			foundSeeder = true
		}
	}
	if !foundSeeder {
		t.Errorf("didn't find seeder amongst leecher peer conns")
	}

	seederStats := seederTorrent.Stats()
	assert.True(t, 13 <= seederStats.BytesWrittenData.Int64())
	assert.True(t, 8 <= seederStats.ChunksWritten.Int64())

	leecherStats := leecherTorrent.Stats()
	assert.True(t, 13 <= leecherStats.BytesReadData.Int64())
	assert.True(t, 8 <= leecherStats.ChunksRead.Int64())

	// Try reading through again for the cases where the torrent data size
	// exceeds the size of the cache.
	assertReadAllGreeting(t, r)
}

type fileCacheClientStorageFactoryParams struct {
	Capacity    int64
	SetCapacity bool
	Wrapper     func(*filecache.Cache) storage.ClientImplCloser
}

func newFileCacheClientStorageFactory(ps fileCacheClientStorageFactoryParams) storageFactory {
	return func(dataDir string) storage.ClientImplCloser {
		fc, err := filecache.NewCache(dataDir)
		if err != nil {
			panic(err)
		}
		if ps.SetCapacity {
			fc.SetCapacity(ps.Capacity)
		}
		return ps.Wrapper(fc)
	}
}

type storageFactory func(string) storage.ClientImplCloser

func TestClientTransferDefault(t *testing.T) {
	testClientTransfer(t, testClientTransferParams{
		ExportClientStatus: true,
		LeecherStorage: newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			Wrapper: fileCachePieceResourceStorage,
		}),
	})
}

func TestClientTransferDefaultNoMetadata(t *testing.T) {
	testClientTransfer(t, testClientTransferParams{
		ExportClientStatus: true,
		LeecherStorage: newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			Wrapper: fileCachePieceResourceStorage,
		}),
		LeecherStartsWithoutMetadata: true,
	})
}

func TestClientTransferRateLimitedUpload(t *testing.T) {
	started := time.Now()
	testClientTransfer(t, testClientTransferParams{
		// We are uploading 13 bytes (the length of the greeting torrent). The
		// chunks are 2 bytes in length. Then the smallest burst we can run
		// with is 2. Time taken is (13-burst)/rate.
		SeederUploadRateLimiter: rate.NewLimiter(11, 2),
		ExportClientStatus:      true,
	})
	require.True(t, time.Since(started) > time.Second)
}

func TestClientTransferRateLimitedDownload(t *testing.T) {
	testClientTransfer(t, testClientTransferParams{
		LeecherDownloadRateLimiter: rate.NewLimiter(512, 512),
	})
}

func fileCachePieceResourceStorage(fc *filecache.Cache) storage.ClientImplCloser {
	return storage.NewResourcePieces(fc.AsResourceProvider())
}

func testClientTransferSmallCache(t *testing.T, setReadahead bool, readahead int64) {
	testClientTransfer(t, testClientTransferParams{
		LeecherStorage: newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			SetCapacity: true,
			// Going below the piece length means it can't complete a piece so
			// that it can be hashed.
			Capacity: 5,
			Wrapper:  fileCachePieceResourceStorage,
		}),
		SetReadahead: setReadahead,
		// Can't readahead too far or the cache will thrash and drop data we
		// thought we had.
		Readahead:          readahead,
		ExportClientStatus: true,

		// These tests don't work well with more than 1 connection to the seeder.
		ConfigureLeecher: ConfigureClient{
			Config: func(cfg *torrent.ClientConfig) {
				cfg.DropDuplicatePeerIds = true
				//cfg.DisableIPv6 = true
				//cfg.DisableUTP = true
			},
		},
	})
}

func TestClientTransferSmallCachePieceSizedReadahead(t *testing.T) {
	testClientTransferSmallCache(t, true, 5)
}

func TestClientTransferSmallCacheLargeReadahead(t *testing.T) {
	testClientTransferSmallCache(t, true, 15)
}

func TestClientTransferSmallCacheDefaultReadahead(t *testing.T) {
	testClientTransferSmallCache(t, false, -1)
}

func TestClientTransferVarious(t *testing.T) {
	// Leecher storage
	for _, ls := range []struct {
		name string
		f    storageFactory
	}{
		{"Filecache", newFileCacheClientStorageFactory(fileCacheClientStorageFactoryParams{
			Wrapper: fileCachePieceResourceStorage,
		})},
		{"Boltdb", storage.NewBoltDB},
	} {
		t.Run(fmt.Sprintf("LeecherStorage=%s", ls.name), func(t *testing.T) {
			// Seeder storage
			for _, ss := range []struct {
				name string
				f    func(string) storage.ClientImplCloser
			}{
				{"File", storage.NewFile},
				{"Mmap", storage.NewMMap},
			} {
				t.Run(fmt.Sprintf("%sSeederStorage", ss.name), func(t *testing.T) {
					for _, responsive := range []bool{false, true} {
						t.Run(fmt.Sprintf("Responsive=%v", responsive), func(t *testing.T) {
							t.Run("NoReadahead", func(t *testing.T) {
								testClientTransfer(t, testClientTransferParams{
									Responsive:     responsive,
									SeederStorage:  ss.f,
									LeecherStorage: ls.f,
								})
							})
							for _, readahead := range []int64{-1, 0, 1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15, 20} {
								t.Run(fmt.Sprintf("readahead=%v", readahead), func(t *testing.T) {
									testClientTransfer(t, testClientTransferParams{
										SeederStorage:  ss.f,
										Responsive:     responsive,
										SetReadahead:   true,
										Readahead:      readahead,
										LeecherStorage: ls.f,
									})
								})
							}
						})
					}
				})
			}
		})
	}
}

// Check that after completing leeching, a leecher transitions to a seeding
// correctly. Connected in a chain like so: Seeder <-> Leecher <-> LeecherLeecher.
func TestSeedAfterDownloading(t *testing.T) {
	greetingTempDir, mi := testutil.GreetingTestTorrent()
	defer os.RemoveAll(greetingTempDir)

	cfg := torrent.TestingConfig()
	cfg.Seed = true
	cfg.DataDir = greetingTempDir
	seeder, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer seeder.Close()
	defer testutil.ExportStatusWriter(seeder, "s")()
	seederTorrent, ok, err := seeder.AddTorrentSpec(torrent.TorrentSpecFromMetaInfo(mi))
	require.NoError(t, err)
	assert.True(t, ok)
	seederTorrent.VerifyData()

	cfg = torrent.TestingConfig()
	cfg.Seed = true
	cfg.DataDir, err = ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(cfg.DataDir)
	leecher, err := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecher.Close()
	defer testutil.ExportStatusWriter(leecher, "l")()

	cfg = torrent.TestingConfig()
	cfg.Seed = false
	cfg.DataDir, err = ioutil.TempDir("", "")
	require.NoError(t, err)
	defer os.RemoveAll(cfg.DataDir)
	leecherLeecher, _ := torrent.NewClient(cfg)
	require.NoError(t, err)
	defer leecherLeecher.Close()
	defer testutil.ExportStatusWriter(leecherLeecher, "ll")()
	leecherGreeting, ok, err := leecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 2
		return
	}())
	require.NoError(t, err)
	assert.True(t, ok)
	llg, ok, err := leecherLeecher.AddTorrentSpec(func() (ret *torrent.TorrentSpec) {
		ret = torrent.TorrentSpecFromMetaInfo(mi)
		ret.ChunkSize = 3
		return
	}())
	require.NoError(t, err)
	assert.True(t, ok)
	// Simultaneously DownloadAll in Leecher, and read the contents
	// consecutively in LeecherLeecher. This non-deterministically triggered a
	// case where the leecher wouldn't unchoke the LeecherLeecher.
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		r := llg.NewReader()
		defer r.Close()
		b, err := ioutil.ReadAll(r)
		require.NoError(t, err)
		assert.EqualValues(t, testutil.GreetingFileContents, b)
	}()
	done := make(chan struct{})
	defer close(done)
	go leecherGreeting.AddClientPeer(seeder)
	go leecherGreeting.AddClientPeer(leecherLeecher)
	wg.Add(1)
	go func() {
		defer wg.Done()
		leecherGreeting.DownloadAll()
		leecher.WaitAll()
	}()
	wg.Wait()
}

type ConfigureClient struct {
	Config func(*torrent.ClientConfig)
	Client func(*torrent.Client)
}
<|endoftext|>


<commit_before>package http2

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/dropbox/godropbox/errors"
)

const (
	// This is a hax to propagate DialError with the Do call
	dialErrorMsgPrefix = "dial "
)

// Pool of persistent HTTP connections. The only limit is on the max # of idle connections we
// cache. Like Python's dropbox.curllib.CurlConnectionPool.
type SimplePool struct {
	// no locking needed, because http.Client has its own locking
	client    *http.Client
	transport *http.Transport

	addr   string
	params ConnectionParams
}

// get from http://golang.org/src/pkg/net/http/transport.go
func getenvEitherCase(k string) string {
	if v := os.Getenv(strings.ToUpper(k)); v != "" {
		return v
	}
	return os.Getenv(strings.ToLower(k))
}

// Creates a new HTTP connection pool using the given address and pool parameters.
//
// 'addr' is a net.Dial()-style 'host:port' destination for making the TCP connection for
// HTTP/HTTPS traffic.  It will be used as the hostname by default for virtual hosting
// and SSL certificate validation; if you'd like to use a different hostname,
// set params.HostHeader.
func NewSimplePool(addr string, params ConnectionParams) *SimplePool {
	pool := &SimplePool{
		addr:   addr,
		params: params,
		client: new(http.Client),
	}

	// It's desirable to enforce the timeout at the client-level since it
	// includes the connection time, redirects and the time to finish reading
	// the full response. Unlike ResponseHeaderTimeout supported by
	// `http.Transport` which merely accounts for the timeout to receive the
	// first response header byte. It ignores the time to send the request or
	// the time to read the full response.
	pool.client.Timeout = params.ResponseTimeout

	// setup HTTP transport
	transport := new(http.Transport)
	transport.ResponseHeaderTimeout = params.ResponseTimeout
	transport.MaxIdleConnsPerHost = params.MaxIdle

	if params.Proxy != nil {
		transport.Proxy = params.Proxy
	} else {
		transport.Proxy = http.ProxyFromEnvironment
	}

	if params.Dial == nil {
		// dialTimeout could only be used in none proxy requests since it talks directly
		// to pool.addr
		if getenvEitherCase("HTTP_PROXY") == "" && params.Proxy == nil {
			transport.Dial = pool.dialTimeout
		}
	} else {
		transport.Dial = params.Dial
	}
	pool.transport = transport
	pool.client.Transport = transport

	if params.UseSSL && params.SkipVerifySSL {
		transport.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	return pool
}

// Adds connection timeout for HTTP client
//
// Note - we do not use the addr passed in, which golang's http library
// has parsed from the URL, so that we can connect to whatever specific host
// was passed in originally as addr to NewSimplePool().
func (pool *SimplePool) dialTimeout(network, addr string) (net.Conn, error) {
	c, err := net.DialTimeout(network, pool.addr, pool.params.ConnectTimeout)
	if err == nil {
		tcp := c.(*net.TCPConn)
		tcp.SetKeepAlive(true)
		tcp.SetKeepAlivePeriod(10 * time.Second)
	}
	return c, err
}

// Performs the HTTP request using our HTTP client
func (pool *SimplePool) Do(req *http.Request) (resp *http.Response, err error) {
	conn, err := pool.Get()
	if err != nil {
		return nil, errors.Wrap(err, err.Error())
	}

	if pool.params.UseSSL {
		req.URL.Scheme = "https"
	} else {
		req.URL.Scheme = "http"
	}
	if pool.params.HostHeader != nil {
		req.URL.Host = *pool.params.HostHeader
	} else {
		req.URL.Host = pool.addr
	}

	// transport.ResponseHeaderTimeout doesn't encompass the time it takes to write the full
	// request, thus timeout for reading the body should be handled by the caller who is
	// consuming the response.
	resp, err = conn.Do(req)
	if err != nil {
		var isDialError bool
		if urlErr, ok := err.(*url.Error); ok {
			if strings.HasPrefix(urlErr.Err.Error(), dialErrorMsgPrefix) {
				isDialError = true
			}
		}
		if isDialError {
			err = DialError{errors.Wrap(err, "SimplePool: Dial Error")}
		} else {
			err = errors.Wrap(err, err.Error())
		}
	}
	return
}

// Returns the HTTP client, which is thread-safe.
//
// Note that we use http.Client, rather than httputil.ClientConn, despite http.Client being higher-
// level. This is normally a liability for backend code, but it has more robust error handling and
// provides functionality that's more comparable to pycurl/curllib.
func (pool *SimplePool) Get() (*http.Client, error) {
	return pool.client, nil
}

// Closes all idle connections in this pool
func (pool *SimplePool) Close() {
	pool.transport.CloseIdleConnections()
}
<commit_msg>Simplify some redundant code<commit_after>package http2

import (
	"crypto/tls"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/dropbox/godropbox/errors"
)

const (
	// This is a hax to propagate DialError with the Do call
	dialErrorMsgPrefix = "dial "
)

// Pool of persistent HTTP connections. The only limit is on the max # of idle connections we
// cache. Like Python's dropbox.curllib.CurlConnectionPool.
type SimplePool struct {
	// no locking needed, because http.Client has its own locking
	client    *http.Client
	transport *http.Transport

	addr   string
	params ConnectionParams
}

// get from http://golang.org/src/pkg/net/http/transport.go
func getenvEitherCase(k string) string {
	if v := os.Getenv(strings.ToUpper(k)); v != "" {
		return v
	}
	return os.Getenv(strings.ToLower(k))
}

// Creates a new HTTP connection pool using the given address and pool parameters.
//
// 'addr' is a net.Dial()-style 'host:port' destination for making the TCP connection for
// HTTP/HTTPS traffic.  It will be used as the hostname by default for virtual hosting
// and SSL certificate validation; if you'd like to use a different hostname,
// set params.HostHeader.
func NewSimplePool(addr string, params ConnectionParams) *SimplePool {
	pool := &SimplePool{
		addr:   addr,
		params: params,
		client: new(http.Client),
	}

	// It's desirable to enforce the timeout at the client-level since it
	// includes the connection time, redirects and the time to finish reading
	// the full response. Unlike ResponseHeaderTimeout supported by
	// `http.Transport` which merely accounts for the timeout to receive the
	// first response header byte. It ignores the time to send the request or
	// the time to read the full response.
	pool.client.Timeout = params.ResponseTimeout

	// setup HTTP transport
	transport := new(http.Transport)
	transport.ResponseHeaderTimeout = params.ResponseTimeout
	transport.MaxIdleConnsPerHost = params.MaxIdle

	if params.Proxy != nil {
		transport.Proxy = params.Proxy
	} else {
		transport.Proxy = http.ProxyFromEnvironment
	}

	if params.Dial == nil {
		// dialTimeout could only be used in none proxy requests since it talks directly
		// to pool.addr
		if getenvEitherCase("HTTP_PROXY") == "" && params.Proxy == nil {
			transport.Dial = pool.dialTimeout
		}
	} else {
		transport.Dial = params.Dial
	}
	pool.transport = transport
	pool.client.Transport = transport

	if params.UseSSL && params.SkipVerifySSL {
		transport.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	return pool
}

// Adds connection timeout for HTTP client
//
// Note - we do not use the addr passed in, which golang's http library
// has parsed from the URL, so that we can connect to whatever specific host
// was passed in originally as addr to NewSimplePool().
func (pool *SimplePool) dialTimeout(network, addr string) (net.Conn, error) {
	c, err := net.DialTimeout(network, pool.addr, pool.params.ConnectTimeout)
	if err == nil {
		tcp := c.(*net.TCPConn)
		tcp.SetKeepAlive(true)
		tcp.SetKeepAlivePeriod(10 * time.Second)
	}
	return c, err
}

// Performs the HTTP request using our HTTP client
func (pool *SimplePool) Do(req *http.Request) (resp *http.Response, err error) {
	conn, err := pool.Get()
	if err != nil {
		return nil, errors.Wrap(err, err.Error())
	}

	if pool.params.UseSSL {
		req.URL.Scheme = "https"
	} else {
		req.URL.Scheme = "http"
	}
	if pool.params.HostHeader != nil {
		req.URL.Host = *pool.params.HostHeader
	} else {
		req.URL.Host = pool.addr
	}

	// transport.ResponseHeaderTimeout doesn't encompass the time it takes to write the full
	// request, thus timeout for reading the body should be handled by the caller who is
	// consuming the response.
	resp, err = conn.Do(req)
	if err != nil {
		if urlErr, ok := err.(*url.Error); ok &&
			strings.HasPrefix(urlErr.Err.Error(), dialErrorMsgPrefix) {
			err = DialError{errors.Wrap(err, "SimplePool: Dial Error")}
		} else {
			err = errors.Wrap(err, err.Error())
		}
	}
	return
}

// Returns the HTTP client, which is thread-safe.
//
// Note that we use http.Client, rather than httputil.ClientConn, despite http.Client being higher-
// level. This is normally a liability for backend code, but it has more robust error handling and
// provides functionality that's more comparable to pycurl/curllib.
func (pool *SimplePool) Get() (*http.Client, error) {
	return pool.client, nil
}

// Closes all idle connections in this pool
func (pool *SimplePool) Close() {
	pool.transport.CloseIdleConnections()
}
<|endoftext|>


<commit_before>// nmon2influxdb
// import nmon data in InfluxDB
// author: adejoux@djouxtech.net

package nmon2influxdblib

import (
	"bufio"
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"regexp"
	"sort"

	"github.com/pkg/sftp"

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

var remoteFileRegexp = regexp.MustCompile(`(\S+):(\S+)`)
var remoteUserRegexp = regexp.MustCompile(`(\S+)@(\S+)`)

const gzipfile = ".gz"
const size = 64000

// File structure used to select nmon files to import
type File struct {
	Name     string
	FileType string
	Host     string
	SSHUser  string
	SSHKey   string
	checksum string
	lines    []string
}

// Files array of File
type Files []File

//Add a file in the NmonFIles structure
func (nmonFiles *Files) Add(file string, fileType string) {
	*nmonFiles = append(*nmonFiles, File{Name: file, FileType: fileType})
}

//AddRemote a remote file in the NmonFIles structure
func (nmonFiles *Files) AddRemote(file string, fileType string, host string, user string, key string) {
	*nmonFiles = append(*nmonFiles, File{Name: file, FileType: fileType, Host: host, SSHUser: user, SSHKey: key})
}

//Valid returns only valid fiels for nmon import
func (nmonFiles *Files) Valid() (validFiles Files) {
	for _, v := range *nmonFiles {
		if v.FileType == ".nmon" || v.FileType == gzipfile {
			validFiles = append(validFiles, v)
		}
	}
	return validFiles
}

// FileScanner struct to manage
type FileScanner struct {
	*os.File
	*bufio.Scanner
}

// RemoteFileScanner struct for remote files
type RemoteFileScanner struct {
	*sftp.File
	*bufio.Scanner
}

// GetRemoteScanner open an nmon file based on file extension and provides a bufio Scanner
func (nmonFile *File) GetRemoteScanner() (*RemoteFileScanner, error) {

	sftpConn := InitSFTP(nmonFile.SSHUser, nmonFile.Host, nmonFile.SSHKey)
	file, err := sftpConn.Open(nmonFile.Name)
	if err != nil {
		return nil, err
	}

	if nmonFile.FileType == gzipfile {
		gr, err := gzip.NewReader(file)
		if err != nil {
			return nil, err
		}
		reader := bufio.NewReader(gr)
		return &RemoteFileScanner{file, bufio.NewScanner(reader)}, nil
	}

	reader := bufio.NewReader(file)
	return &RemoteFileScanner{file, bufio.NewScanner(reader)}, nil
}

//Checksum generates SHA1 file checksum
func (nmonFile *File) Checksum() (fileHash string) {
	if len(nmonFile.checksum) > 0 {
		return nmonFile.checksum
	}
	var result []byte
	if len(nmonFile.Host) > 0 {
		scanner, err := nmonFile.GetRemoteScanner()
		CheckError(err)
		scanner.Seek(-1024, 2)
		hash := sha1.New()
		if _, err = io.Copy(hash, scanner); err != nil {
			return
		}
		fileHash = hex.EncodeToString(hash.Sum(result))
	} else {
		scanner, err := nmonFile.GetScanner()
		CheckError(err)
		scanner.Seek(-1024, 2)
		hash := sha1.New()
		if _, err = io.Copy(hash, scanner); err != nil {
			return
		}
		fileHash = hex.EncodeToString(hash.Sum(result))
	}
	nmonFile.checksum = fileHash
	return
}

// GetScanner open an nmon file based on file extension and provides a bufio Scanner
func (nmonFile *File) GetScanner() (*FileScanner, error) {

	file, err := os.Open(nmonFile.Name)
	if err != nil {
		return nil, err
	}

	if nmonFile.FileType == gzipfile {
		gr, err := gzip.NewReader(file)
		if err != nil {
			return nil, err
		}
		reader := bufio.NewReader(gr)
		return &FileScanner{file, bufio.NewScanner(reader)}, nil
	}

	reader := bufio.NewReader(file)
	return &FileScanner{file, bufio.NewScanner(reader)}, nil
}

// Parse parameters
func (nmonFiles *Files) Parse(args []string, sshUser string, key string) {
	for _, param := range args {
		if remoteFileRegexp.MatchString(param) {
			matched := remoteFileRegexp.FindStringSubmatch(param)
			host := matched[1]

			if remoteUserRegexp.MatchString(host) {
				hostMatched := remoteUserRegexp.FindStringSubmatch(host)
				sshUser = hostMatched[1]
				host = hostMatched[2]
			}
			matchedParam := matched[2]

			sftpConn := InitSFTP(sshUser, host, key)
			paraminfo, err := sftpConn.Stat(matchedParam)
			CheckError(err)
			if err != nil {
				if os.IsNotExist(err) {
					fmt.Printf("%s doesn't exist ! skipped.\n", param)
				}
				continue
			}
			if paraminfo.IsDir() {
				entries, err := sftpConn.ReadDir(matchedParam)
				CheckError(err)
				for _, entry := range entries {
					if !entry.IsDir() {
						file := path.Join(matchedParam, entry.Name())
						nmonFiles.AddRemote(file, path.Ext(file), host, sshUser, key)
					}
				}
				sftpConn.Close()
				continue
			}
			nmonFiles.AddRemote(matchedParam, path.Ext(matchedParam), host, sshUser, key)
			sftpConn.Close()
			continue
		}

		paraminfo, err := os.Stat(param)
		if err != nil {
			if os.IsNotExist(err) {
				fmt.Printf("%s doesn't exist ! skipped.\n", param)
			}
			continue
		}

		if paraminfo.IsDir() {
			entries, err := ioutil.ReadDir(param)
			CheckError(err)
			for _, entry := range entries {
				if !entry.IsDir() {
					file := path.Join(param, entry.Name())
					nmonFiles.Add(file, path.Ext(file))
				}
			}
			continue
		}
		nmonFiles.Add(param, path.Ext(param))
	}
}

//SSHConfig contains SSH parameters
type SSHConfig struct {
	User string
	Key  string
}

//InitSFTP init sftp session
func InitSFTP(sshUser string, host string, key string) *sftp.Client {
	var auths []ssh.AuthMethod

	if IsFile(key) {
		pemBytes, err := ioutil.ReadFile(key)
		if err != nil {
			log.Fatal(err)
		}
		signer, err := ssh.ParsePrivateKey(pemBytes)
		if err != nil {
			log.Fatalf("parse key failed:%v", err)
		}

		auths = append(auths, ssh.PublicKeys(signer))
	}

	// ssh agent support
	if aconn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		auths = append(auths, ssh.PublicKeysCallback(agent.NewClient(aconn).Signers))
	}

	config := &ssh.ClientConfig{
		User: sshUser,
		Auth: auths,
	}
	sshhost := fmt.Sprintf("%s:22", host)
	conn, err := ssh.Dial("tcp", sshhost, config)
	if err != nil {
		log.Fatalf("dial failed:%v", err)
	}

	c, err := sftp.NewClient(conn, sftp.MaxPacket(size))
	if err != nil {
		log.Fatalf("unable to start sftp subsytem: %v", err)
	}
	return c
}

//Content returns the nmon files content sorted in an slice of string format
func (nmonFile *File) Content() []string {
	if len(nmonFile.lines) > 0 {
		return nmonFile.lines
	}
	if len(nmonFile.Host) > 0 {
		scanner, err := nmonFile.GetRemoteScanner()
		CheckError(err)
		scanner.Split(bufio.ScanLines)
		for scanner.Scan() {
			nmonFile.lines = append(nmonFile.lines, scanner.Text())
		}
		scanner.Close()
	} else {
		scanner, err := nmonFile.GetScanner()
		CheckError(err)
		scanner.Split(bufio.ScanLines)
		for scanner.Scan() {
			nmonFile.lines = append(nmonFile.lines, scanner.Text())
		}
		scanner.Close()
	}

	sort.Strings(nmonFile.lines)

	return nmonFile.lines
}
<commit_msg>fixed new ssh host key check<commit_after>// nmon2influxdb
// import nmon data in InfluxDB
// author: adejoux@djouxtech.net

package nmon2influxdblib

import (
	"bufio"
	"compress/gzip"
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"regexp"
	"sort"

	"github.com/pkg/sftp"

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

var remoteFileRegexp = regexp.MustCompile(`(\S+):(\S+)`)
var remoteUserRegexp = regexp.MustCompile(`(\S+)@(\S+)`)

const gzipfile = ".gz"
const size = 64000

// File structure used to select nmon files to import
type File struct {
	Name     string
	FileType string
	Host     string
	SSHUser  string
	SSHKey   string
	checksum string
	lines    []string
}

// Files array of File
type Files []File

//Add a file in the NmonFIles structure
func (nmonFiles *Files) Add(file string, fileType string) {
	*nmonFiles = append(*nmonFiles, File{Name: file, FileType: fileType})
}

//AddRemote a remote file in the NmonFIles structure
func (nmonFiles *Files) AddRemote(file string, fileType string, host string, user string, key string) {
	*nmonFiles = append(*nmonFiles, File{Name: file, FileType: fileType, Host: host, SSHUser: user, SSHKey: key})
}

//Valid returns only valid fiels for nmon import
func (nmonFiles *Files) Valid() (validFiles Files) {
	for _, v := range *nmonFiles {
		if v.FileType == ".nmon" || v.FileType == gzipfile {
			validFiles = append(validFiles, v)
		}
	}
	return validFiles
}

// FileScanner struct to manage
type FileScanner struct {
	*os.File
	*bufio.Scanner
}

// RemoteFileScanner struct for remote files
type RemoteFileScanner struct {
	*sftp.File
	*bufio.Scanner
}

// GetRemoteScanner open an nmon file based on file extension and provides a bufio Scanner
func (nmonFile *File) GetRemoteScanner() (*RemoteFileScanner, error) {

	sftpConn := InitSFTP(nmonFile.SSHUser, nmonFile.Host, nmonFile.SSHKey)
	file, err := sftpConn.Open(nmonFile.Name)
	if err != nil {
		return nil, err
	}

	if nmonFile.FileType == gzipfile {
		gr, err := gzip.NewReader(file)
		if err != nil {
			return nil, err
		}
		reader := bufio.NewReader(gr)
		return &RemoteFileScanner{file, bufio.NewScanner(reader)}, nil
	}

	reader := bufio.NewReader(file)
	return &RemoteFileScanner{file, bufio.NewScanner(reader)}, nil
}

//Checksum generates SHA1 file checksum
func (nmonFile *File) Checksum() (fileHash string) {
	if len(nmonFile.checksum) > 0 {
		return nmonFile.checksum
	}
	var result []byte
	if len(nmonFile.Host) > 0 {
		scanner, err := nmonFile.GetRemoteScanner()
		CheckError(err)
		scanner.Seek(-1024, 2)
		hash := sha1.New()
		if _, err = io.Copy(hash, scanner); err != nil {
			return
		}
		fileHash = hex.EncodeToString(hash.Sum(result))
	} else {
		scanner, err := nmonFile.GetScanner()
		CheckError(err)
		scanner.Seek(-1024, 2)
		hash := sha1.New()
		if _, err = io.Copy(hash, scanner); err != nil {
			return
		}
		fileHash = hex.EncodeToString(hash.Sum(result))
	}
	nmonFile.checksum = fileHash
	return
}

// GetScanner open an nmon file based on file extension and provides a bufio Scanner
func (nmonFile *File) GetScanner() (*FileScanner, error) {

	file, err := os.Open(nmonFile.Name)
	if err != nil {
		return nil, err
	}

	if nmonFile.FileType == gzipfile {
		gr, err := gzip.NewReader(file)
		if err != nil {
			return nil, err
		}
		reader := bufio.NewReader(gr)
		return &FileScanner{file, bufio.NewScanner(reader)}, nil
	}

	reader := bufio.NewReader(file)
	return &FileScanner{file, bufio.NewScanner(reader)}, nil
}

// Parse parameters
func (nmonFiles *Files) Parse(args []string, sshUser string, key string) {
	for _, param := range args {
		if remoteFileRegexp.MatchString(param) {
			matched := remoteFileRegexp.FindStringSubmatch(param)
			host := matched[1]

			if remoteUserRegexp.MatchString(host) {
				hostMatched := remoteUserRegexp.FindStringSubmatch(host)
				sshUser = hostMatched[1]
				host = hostMatched[2]
			}
			matchedParam := matched[2]

			sftpConn := InitSFTP(sshUser, host, key)
			paraminfo, err := sftpConn.Stat(matchedParam)
			CheckError(err)
			if err != nil {
				if os.IsNotExist(err) {
					fmt.Printf("%s doesn't exist ! skipped.\n", param)
				}
				continue
			}
			if paraminfo.IsDir() {
				entries, err := sftpConn.ReadDir(matchedParam)
				CheckError(err)
				for _, entry := range entries {
					if !entry.IsDir() {
						file := path.Join(matchedParam, entry.Name())
						nmonFiles.AddRemote(file, path.Ext(file), host, sshUser, key)
					}
				}
				sftpConn.Close()
				continue
			}
			nmonFiles.AddRemote(matchedParam, path.Ext(matchedParam), host, sshUser, key)
			sftpConn.Close()
			continue
		}

		paraminfo, err := os.Stat(param)
		if err != nil {
			if os.IsNotExist(err) {
				fmt.Printf("%s doesn't exist ! skipped.\n", param)
			}
			continue
		}

		if paraminfo.IsDir() {
			entries, err := ioutil.ReadDir(param)
			CheckError(err)
			for _, entry := range entries {
				if !entry.IsDir() {
					file := path.Join(param, entry.Name())
					nmonFiles.Add(file, path.Ext(file))
				}
			}
			continue
		}
		nmonFiles.Add(param, path.Ext(param))
	}
}

//SSHConfig contains SSH parameters
type SSHConfig struct {
	User string
	Key  string
}

//InitSFTP init sftp session
func InitSFTP(sshUser string, host string, key string) *sftp.Client {
	var auths []ssh.AuthMethod

	if IsFile(key) {
		pemBytes, err := ioutil.ReadFile(key)
		if err != nil {
			log.Fatal(err)
		}
		signer, err := ssh.ParsePrivateKey(pemBytes)
		if err != nil {
			log.Fatalf("parse key failed:%v", err)
		}

		auths = append(auths, ssh.PublicKeys(signer))
	}

	// ssh agent support
	if aconn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		auths = append(auths, ssh.PublicKeysCallback(agent.NewClient(aconn).Signers))
	}

	config := &ssh.ClientConfig{
		User: sshUser,
		Auth: auths,
                HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}
	sshhost := fmt.Sprintf("%s:22", host)
	conn, err := ssh.Dial("tcp", sshhost, config)
	if err != nil {
		log.Fatalf("dial failed:%v", err)
	}

	c, err := sftp.NewClient(conn, sftp.MaxPacket(size))
	if err != nil {
		log.Fatalf("unable to start sftp subsytem: %v", err)
	}
	return c
}

//Content returns the nmon files content sorted in an slice of string format
func (nmonFile *File) Content() []string {
	if len(nmonFile.lines) > 0 {
		return nmonFile.lines
	}
	if len(nmonFile.Host) > 0 {
		scanner, err := nmonFile.GetRemoteScanner()
		CheckError(err)
		scanner.Split(bufio.ScanLines)
		for scanner.Scan() {
			nmonFile.lines = append(nmonFile.lines, scanner.Text())
		}
		scanner.Close()
	} else {
		scanner, err := nmonFile.GetScanner()
		CheckError(err)
		scanner.Split(bufio.ScanLines)
		for scanner.Scan() {
			nmonFile.lines = append(nmonFile.lines, scanner.Text())
		}
		scanner.Close()
	}

	sort.Strings(nmonFile.lines)

	return nmonFile.lines
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/boltdb/bolt"
	"net/http"
	"net/url"
)

func init() {
	RegisterComponent("forward-request-action", ForwardRequestAction{})
}

// ForwardRequestAction is a component that forwards an incoming request to a
// user defined URL.
type ForwardRequestAction struct {
}

// Name returns the name of this component.
func (ForwardRequestAction) Name() string { return "Forward request" }

// Template returns the HTML template name of this component.
func (ForwardRequestAction) Template() string { return "request-forward-action" }

// Init initializes this component. It requires a valid url parameter to be
// present.
func (ForwardRequestAction) Init(h Hook, params map[string]string, b *bolt.Bucket) error {
	uri, ok := params["url"]
	if !ok {
		return errors.New("url is required")
	}

	if _, err := url.Parse(uri); err != nil {
		return fmt.Errorf("url is not valid: %s", err)
	}

	return b.Put([]byte(fmt.Sprintf("%s-url", h.ID)), []byte(uri))
}

// Process forwards the incoming request to the configured URL.
func (ForwardRequestAction) Process(h Hook, r Request, b *bolt.Bucket) error {
	uri := b.Get([]byte(fmt.Sprintf("%s-url", h.ID)))
	if uri == nil {
		return errors.New("forward request action not initialized")
	}

	req, err := http.NewRequest(r.Method, string(uri), bytes.NewReader(r.Body))
	if err != nil {
		return fmt.Errorf("could not create new request: %s", err)
	}

	for k, v := range r.Headers {
		// special handling for some headers
		switch k {
		case "Connection":
			// skip
			continue
		case "User-Agent":
			// rename header, we will set our own user agent
			k = "X-Forwarded-User-Agent"
		}
		req.Header.Set(k, v)
	}

	req.Header.Set("User-Agent", "rehook/0.0.1 (https://github.com/gophergala/rehook)")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("request forward error: %s", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		return fmt.Errorf("request forward unexpected status code received: %d", resp.StatusCode)
	}
	return nil
}
<commit_msg>Fix forward_action imports<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"net/http"
	"net/url"

	"github.com/boltdb/bolt"
)

func init() {
	RegisterComponent("forward-request-action", ForwardRequestAction{})
}

// ForwardRequestAction is a component that forwards an incoming request to a
// user defined URL.
type ForwardRequestAction struct {
}

// Name returns the name of this component.
func (ForwardRequestAction) Name() string { return "Forward request" }

// Template returns the HTML template name of this component.
func (ForwardRequestAction) Template() string { return "request-forward-action" }

// Init initializes this component. It requires a valid url parameter to be
// present.
func (ForwardRequestAction) Init(h Hook, params map[string]string, b *bolt.Bucket) error {
	uri, ok := params["url"]
	if !ok {
		return errors.New("url is required")
	}

	if _, err := url.Parse(uri); err != nil {
		return fmt.Errorf("url is not valid: %s", err)
	}

	return b.Put([]byte(fmt.Sprintf("%s-url", h.ID)), []byte(uri))
}

// Process forwards the incoming request to the configured URL.
func (ForwardRequestAction) Process(h Hook, r Request, b *bolt.Bucket) error {
	uri := b.Get([]byte(fmt.Sprintf("%s-url", h.ID)))
	if uri == nil {
		return errors.New("forward request action not initialized")
	}

	req, err := http.NewRequest(r.Method, string(uri), bytes.NewReader(r.Body))
	if err != nil {
		return fmt.Errorf("could not create new request: %s", err)
	}

	for k, v := range r.Headers {
		// special handling for some headers
		switch k {
		case "Connection":
			// skip
			continue
		case "User-Agent":
			// rename header, we will set our own user agent
			k = "X-Forwarded-User-Agent"
		}
		req.Header.Set(k, v)
	}

	req.Header.Set("User-Agent", "rehook/0.0.1 (https://github.com/gophergala/rehook)")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("request forward error: %s", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 300 {
		return fmt.Errorf("request forward unexpected status code received: %d", resp.StatusCode)
	}
	return nil
}
<|endoftext|>


<commit_before>package filereader

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"github.com/Workiva/go-datastructures/trie/yfast"
	"github.com/dchest/blake2b"
	"github.com/hashicorp/golang-lru"

	"github.com/tsileo/blobstash/ext/filetree/filetreeutil/meta"
)

type BlobStore interface {
	Get(hash string) ([]byte, error)
}

// Download a file by its hash to path
func GetFile(bs BlobStore, hash, path string) error {
	// readResult := &ReadResult{}
	buf, err := os.Create(path)
	defer buf.Close()
	if err != nil {
		return err
	}
	h := blake2b.New256()
	js, err := bs.Get(hash)
	if err != nil {
		return err
	}
	meta, err := meta.NewMetaFromBlob(hash, js)
	if err != nil {
		return fmt.Errorf("failed to get meta %v: %v", hash, err)
	}
	meta.Hash = hash
	ffile := NewFile(bs, meta)
	defer ffile.Close()
	fileReader := io.TeeReader(ffile, h)
	io.Copy(buf, fileReader)
	// readResult.Hash = fmt.Sprintf("%x", h.Sum(nil))
	// readResult.FilesCount++
	// readResult.FilesDownloaded++
	fstat, err := buf.Stat()
	if err != nil {
		return err
	}
	if int(fstat.Size()) != meta.Size {
		return fmt.Errorf("file %+v not successfully restored, size:%d/expected size:%d", fstat.Size, meta.Size)
	}
	return nil
	// readResult.Size = int(fstat.Size())
	// readResult.SizeDownloaded = readResult.Size
	// if readResult.Size != meta.Size {
	// 	return readResult, fmt.Errorf("file %+v not successfully restored, size:%v/expected size:%v",
	// 		meta, readResult.Size, meta.Size)
	// }
	// return readResult, nil
}

type IndexValue struct {
	Index int
	Value string
	I     int
}

// Key is needed for yfast
func (iv *IndexValue) Key() uint64 {
	return uint64(iv.Index)
}

// FakeFile implements io.Reader, and io.ReaderAt.
// It fetch blobs on the fly.
type File struct {
	name    string
	bs      BlobStore
	meta    *meta.Meta
	offset  int
	size    int
	llen    int
	lmrange []*IndexValue
	trie    *yfast.YFastTrie
	lru     *lru.Cache
}

// NewFakeFile creates a new FakeFile instance.
func NewFile(bs BlobStore, meta *meta.Meta) (f *File) {
	// Needed for the blob routing
	cache, err := lru.New(2)
	if err != nil {
		panic(err)
	}
	f = &File{
		bs:      bs,
		meta:    meta,
		size:    meta.Size,
		lmrange: []*IndexValue{},
		trie:    yfast.New(uint64(0)),
		lru:     cache,
	}
	if meta.Size > 0 {
		for idx, m := range meta.Refs {
			data := m.([]interface{})
			var index int
			switch i := data[0].(type) {
			case float64:
				index = int(i)
			case int:
				index = i
			default:
				panic("unexpected index")
			}
			iv := &IndexValue{Index: index, Value: data[1].(string), I: idx}
			f.lmrange = append(f.lmrange, iv)
			f.trie.Insert(iv)
		}
	}
	return
}

func (f *File) Close() error {
	f.lru.Purge()
	return nil
}

// ReadAt implements the io.ReaderAt interface
func (f *File) ReadAt(p []byte, offset int64) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}
	if f.size == 0 || f.offset >= f.size {
		return 0, io.EOF
	}
	buf, err := f.read(int(offset), len(p))
	if err != nil {
		return
	}
	n = copy(p, buf)
	return
}

// Low level read function, read a size from an offset
// Iterate only the needed blobs
func (f *File) read(offset, cnt int) ([]byte, error) {
	//log.Printf("FakeFile %v read(%v, %v)", f.ref, offset, cnt)
	if cnt < 0 || cnt > f.size {
		cnt = f.size
	}
	var buf bytes.Buffer
	var cbuf []byte
	var err error
	written := 0

	if len(f.lmrange) == 0 {
		panic(fmt.Errorf("FakeFile %+v lmrange empty", f))
	}

	tiv := f.trie.Successor(uint64(offset)).(*IndexValue)
	if tiv.Index == offset {
		tiv = f.trie.Successor(uint64(offset + 1)).(*IndexValue)
	}
	for _, iv := range f.lmrange[tiv.I:] {
		if offset > iv.Index {
			continue
		}
		//bbuf, _, _ := f.client.Blobs.Get(iv.Value)
		if cached, ok := f.lru.Get(iv.Value); ok {
			cbuf = cached.([]byte)
		} else {
			bbuf, err := f.bs.Get(iv.Value)
			if err != nil {
				return nil, fmt.Errorf("failed to fetch blob %v: %v", iv.Value, err)
			}
			f.lru.Add(iv.Value, bbuf)
			cbuf = bbuf
		}
		bbuf := cbuf
		foffset := 0
		if offset != 0 {
			// Compute the starting offset of the blob
			blobStart := iv.Index - len(bbuf)
			// and subtract it to get the correct offset
			foffset = offset - blobStart
			offset = 0
		}
		// If the remaining cnt (cnt - written)
		// is greater than the blob slice
		if cnt-written > len(bbuf)-foffset {
			fwritten, err := buf.Write(bbuf[foffset:])
			if err != nil {
				return nil, err
			}
			written += fwritten

		} else {
			// What we need fit in this blob
			// it should return after this
			if foffset+cnt-written > len(bbuf) {
				panic(fmt.Errorf("failed to read from FakeFile %+v [%v:%v]", f, foffset, foffset+cnt-written))
			}
			fwritten, err := buf.Write(bbuf[foffset : foffset+cnt-written])
			if err != nil {
				return nil, err
			}

			written += fwritten
			// Check that the total written bytes equals the requested size
			if written != cnt {
				panic("error reading FakeFile")
			}
		}
		if written == cnt {
			return buf.Bytes(), nil
		}
		cbuf = nil
	}
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), err
}

// Reset resets the offset to 0
func (f *File) Reset() {
	f.offset = 0
}

// Read implements io.Reader
func (f *File) Read(p []byte) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}
	if f.size == 0 || f.offset >= f.size {
		return 0, io.EOF
	}
	n = 0
	limit := len(p)
	if limit > (f.size - f.offset) {
		limit = f.size - f.offset
	}
	b, err := f.read(f.offset, limit)
	if err == io.EOF {
		return 0, io.EOF
	}
	if err != nil {
		return 0, fmt.Errorf("failed to read %+v at range %v-%v", f, f.offset, limit)
	}
	n = copy(p, b)
	f.offset += n
	return
}
<commit_msg>ext/filetree: make `File` implements `io.Seeker`<commit_after>package filereader

import (
	"bytes"
	"fmt"
	"io"
	"os"

	"github.com/Workiva/go-datastructures/trie/yfast"
	"github.com/dchest/blake2b"
	"github.com/hashicorp/golang-lru"

	"github.com/tsileo/blobstash/ext/filetree/filetreeutil/meta"
)

const (
	SEEK_SET int = 0 // seek relative to the origin of the file
	SEEK_CUR int = 1 // seek relative to the current offset
	SEEK_END int = 2 // seek relative to the end
)

type BlobStore interface {
	Get(hash string) ([]byte, error)
}

// Download a file by its hash to path
func GetFile(bs BlobStore, hash, path string) error {
	// readResult := &ReadResult{}
	buf, err := os.Create(path)
	defer buf.Close()
	if err != nil {
		return err
	}
	h := blake2b.New256()
	js, err := bs.Get(hash)
	if err != nil {
		return err
	}
	meta, err := meta.NewMetaFromBlob(hash, js)
	if err != nil {
		return fmt.Errorf("failed to get meta %v: %v", hash, err)
	}
	meta.Hash = hash
	ffile := NewFile(bs, meta)
	defer ffile.Close()
	fileReader := io.TeeReader(ffile, h)
	io.Copy(buf, fileReader)
	// readResult.Hash = fmt.Sprintf("%x", h.Sum(nil))
	// readResult.FilesCount++
	// readResult.FilesDownloaded++
	fstat, err := buf.Stat()
	if err != nil {
		return err
	}
	if int(fstat.Size()) != meta.Size {
		return fmt.Errorf("file %+v not successfully restored, size:%d/expected size:%d", fstat.Size, meta.Size)
	}
	return nil
	// readResult.Size = int(fstat.Size())
	// readResult.SizeDownloaded = readResult.Size
	// if readResult.Size != meta.Size {
	// 	return readResult, fmt.Errorf("file %+v not successfully restored, size:%v/expected size:%v",
	// 		meta, readResult.Size, meta.Size)
	// }
	// return readResult, nil
}

type IndexValue struct {
	Index int64
	Value string
	I     int
}

// Key is needed for yfast
func (iv *IndexValue) Key() uint64 {
	return uint64(iv.Index)
}

// FakeFile implements io.Reader, and io.ReaderAt.
// It fetch blobs on the fly.
type File struct {
	name    string
	bs      BlobStore
	meta    *meta.Meta
	offset  int64
	size    int64
	llen    int
	lmrange []*IndexValue
	trie    *yfast.YFastTrie
	lru     *lru.Cache
}

// NewFakeFile creates a new FakeFile instance.
func NewFile(bs BlobStore, meta *meta.Meta) (f *File) {
	// Needed for the blob routing
	cache, err := lru.New(2)
	if err != nil {
		panic(err)
	}
	f = &File{
		bs:      bs,
		meta:    meta,
		size:    int64(meta.Size),
		lmrange: []*IndexValue{},
		trie:    yfast.New(uint64(0)),
		lru:     cache,
	}
	if meta.Size > 0 {
		for idx, m := range meta.Refs {
			data := m.([]interface{})
			var index int64
			switch i := data[0].(type) {
			case float64:
				index = int64(i)
			case int:
				index = int64(i)
			case int64:
				index = i
			default:
				panic("unexpected index")
			}
			iv := &IndexValue{Index: index, Value: data[1].(string), I: idx}
			f.lmrange = append(f.lmrange, iv)
			f.trie.Insert(iv)
		}
	}
	return
}

func (f *File) Close() error {
	f.lru.Purge()
	return nil
}

// ReadAt implements the io.ReaderAt interface
func (f *File) ReadAt(p []byte, offset int64) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}
	if f.size == 0 || f.offset >= f.size {
		return 0, io.EOF
	}
	buf, err := f.read(offset, len(p))
	if err != nil {
		return
	}
	n = copy(p, buf)
	return
}

// Low level read function, read a size from an offset
// Iterate only the needed blobs
func (f *File) read(offset int64, cnt int) ([]byte, error) {
	//log.Printf("FakeFile %v read(%v, %v)", f.ref, offset, cnt)
	if cnt < 0 || int64(cnt) > f.size {
		cnt = int(f.size)
	}
	var buf bytes.Buffer
	var cbuf []byte
	var err error
	written := 0

	if len(f.lmrange) == 0 {
		panic(fmt.Errorf("FakeFile %+v lmrange empty", f))
	}

	tiv := f.trie.Successor(uint64(offset)).(*IndexValue)
	if tiv.Index == offset {
		tiv = f.trie.Successor(uint64(offset + 1)).(*IndexValue)
	}
	for _, iv := range f.lmrange[tiv.I:] {
		if offset > iv.Index {
			continue
		}
		//bbuf, _, _ := f.client.Blobs.Get(iv.Value)
		if cached, ok := f.lru.Get(iv.Value); ok {
			cbuf = cached.([]byte)
		} else {
			bbuf, err := f.bs.Get(iv.Value)
			if err != nil {
				return nil, fmt.Errorf("failed to fetch blob %v: %v", iv.Value, err)
			}
			f.lru.Add(iv.Value, bbuf)
			cbuf = bbuf
		}
		bbuf := cbuf
		foffset := 0
		if offset != 0 {
			// Compute the starting offset of the blob
			blobStart := iv.Index - int64(len(bbuf))
			// and subtract it to get the correct offset
			foffset = int(offset - int64(blobStart))
			offset = 0
		}
		// If the remaining cnt (cnt - written)
		// is greater than the blob slice
		if cnt-written > len(bbuf)-foffset {
			fwritten, err := buf.Write(bbuf[foffset:])
			if err != nil {
				return nil, err
			}
			written += fwritten

		} else {
			// What we need fit in this blob
			// it should return after this
			if foffset+cnt-written > len(bbuf) {
				panic(fmt.Errorf("failed to read from FakeFile %+v [%v:%v]", f, foffset, foffset+cnt-written))
			}
			fwritten, err := buf.Write(bbuf[foffset : foffset+cnt-written])
			if err != nil {
				return nil, err
			}

			written += fwritten
			// Check that the total written bytes equals the requested size
			if written != cnt {
				panic("error reading FakeFile")
			}
		}
		if written == cnt {
			return buf.Bytes(), nil
		}
		cbuf = nil
	}
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), err
}

// Reset resets the offset to 0
func (f *File) Reset() {
	f.offset = 0
}

// Read implements io.Reader
func (f *File) Read(p []byte) (n int, err error) {
	if len(p) == 0 {
		return 0, nil
	}
	if f.size == 0 || f.offset >= f.size {
		return 0, io.EOF
	}
	n = 0
	limit := len(p)
	if limit > int(f.size-f.offset) {
		limit = int(f.size - f.offset)
	}
	b, err := f.read(f.offset, limit)
	if err == io.EOF {
		return 0, io.EOF
	}
	if err != nil {
		return 0, fmt.Errorf("failed to read %+v at range %v-%v", f, f.offset, limit)
	}
	n = copy(p, b)
	f.offset += int64(n)
	return
}

func (f *File) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	case SEEK_SET:
		f.offset = offset
	case SEEK_CUR:
		f.offset += offset
	case SEEK_END:
		f.offset = f.size - offset
	default:
		return 0, fmt.Errorf("invalid whence: %d", whence)
	}
	return f.offset, nil
}
<|endoftext|>


<commit_before>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import (
	"net/url"
	"testing"

	. "github.com/apache/incubator-trafficcontrol/traffic_ops/traffic_ops_golang/tcstructs"
	"github.com/jmoiron/sqlx"

	sqlmock "gopkg.in/DATA-DOG/go-sqlmock.v1"
)

func getTestServers() []Server {
	servers := []Server{}
	testServer := Server{
		Cachegroup:     "Cachegroup",
		CachegroupId:   1,
		CdnId:          1,
		CdnName:        "cdnName",
		DomainName:     "domainName",
		Guid:           "guid",
		HostName:       "server1",
		HttpsPort:      443,
		Id:             1,
		IloIpAddress:   "iloIpAddress",
		IloIpGateway:   "iloIpGateway",
		IloIpNetmask:   "iloIpNetmask",
		IloPassword:    "iloPassword",
		IloUsername:    "iloUsername",
		InterfaceMtu:   "interfaceMtu",
		InterfaceName:  "interfaceName",
		Ip6Address:     "ip6Address",
		Ip6Gateway:     "ip6Gateway",
		IpAddress:      "ipAddress",
		IpGateway:      "ipGateway",
		IpNetmask:      "ipNetmask",
		LastUpdated:    "lastUpdated",
		MgmtIpAddress:  "mgmtIpAddress",
		MgmtIpGateway:  "mgmtIpGateway",
		MgmtIpNetmask:  "mgmtIpNetmask",
		OfflineReason:  "offlineReason",
		PhysLocation:   "physLocation",
		PhysLocationId: 1,
		Profile:        "profile",
		ProfileDesc:    "profileDesc",
		ProfileId:      1,
		Rack:           "rack",
		RevalPending:   true,
		RouterHostName: "routerHostName",
		RouterPortName: "routerPortName",
		Status:         "status",
		StatusId:       1,
		TcpPort:        80,
		ServerType:     "EDGE",
		ServerTypeId:   1,
		UpdPending:     true,
		XmppId:         "xmppId",
		XmppPasswd:     "xmppPasswd",
	}
	servers = append(servers, testServer)

	testServer2 := testServer
	testServer2.Cachegroup = "cachegroup2"
	testServer2.HostName = "server2"
	servers = append(servers, testServer2)

	return servers
}

func TestGetServersByDsId(t *testing.T) {
	mockDB, mock, err := sqlmock.New()
	defer mockDB.Close()
	db := sqlx.NewDb(mockDB, "sqlmock")
	if err != nil {
		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
	}
	defer db.Close()

	testServers := getTestServers()
	cols := ColsFromStructByTag("db", Server{})
	rows := sqlmock.NewRows(cols)

	//TODO: drichardson - build helper to add these Rows from the struct values
	//                    or by CSV if types get in the way
	for _, ts := range testServers {
		rows = rows.AddRow(
			ts.Cachegroup,
			ts.CachegroupId,
			ts.CdnId,
			ts.CdnName,
			ts.DomainName,
			ts.Guid,
			ts.HostName,
			ts.HttpsPort,
			ts.Id,
			ts.IloIpAddress,
			ts.IloIpGateway,
			ts.IloIpNetmask,
			ts.IloPassword,
			ts.IloUsername,
			ts.InterfaceMtu,
			ts.InterfaceName,
			ts.Ip6Address,
			ts.Ip6Gateway,
			ts.IpAddress,
			ts.IpNetmask,
			ts.IpGateway,
			ts.LastUpdated,
			ts.MgmtIpAddress,
			ts.MgmtIpGateway,
			ts.MgmtIpNetmask,
			ts.OfflineReason,
			ts.PhysLocation,
			ts.PhysLocationId,
			ts.Profile,
			ts.ProfileDesc,
			ts.ProfileId,
			ts.Rack,
			ts.RevalPending,
			ts.RouterHostName,
			ts.RouterPortName,
			ts.Status,
			ts.StatusId,
			ts.TcpPort,
			ts.ServerType,
			ts.ServerTypeId,
			ts.UpdPending,
			ts.XmppId,
			ts.XmppPasswd,
		)
	}
	mock.ExpectQuery("SELECT").WillReturnRows(rows)
	v := url.Values{}
	v.Set("dsId", "1")

	servers, err := getServers(v, db, PrivLevelAdmin)
	if err != nil {
		t.Errorf("getServers expected: nil error, actual: %v", err)
	}

	if len(servers) != 2 {
		t.Errorf("getServers expected: len(servers) == 1, actual: %v", len(servers))
	}

}

type SortableServers []Server

func (s SortableServers) Len() int {
	return len(s)
}
func (s SortableServers) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s SortableServers) Less(i, j int) bool {
	return s[i].HostName < s[j].HostName
}
<commit_msg>fixed test case<commit_after>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

import (
	"net/url"
	"testing"

	. "github.com/apache/incubator-trafficcontrol/traffic_ops/traffic_ops_golang/tcstructs"
	"github.com/jmoiron/sqlx"

	sqlmock "gopkg.in/DATA-DOG/go-sqlmock.v1"
)

func getTestServers() []Server {
	servers := []Server{}
	testServer := Server{
		Cachegroup:     "Cachegroup",
		CachegroupId:   1,
		CdnId:          1,
		CdnName:        "cdnName",
		DomainName:     "domainName",
		Guid:           "guid",
		HostName:       "server1",
		HttpsPort:      443,
		Id:             1,
		IloIpAddress:   "iloIpAddress",
		IloIpGateway:   "iloIpGateway",
		IloIpNetmask:   "iloIpNetmask",
		IloPassword:    "iloPassword",
		IloUsername:    "iloUsername",
		InterfaceMtu:   9500,
		InterfaceName:  "interfaceName",
		Ip6Address:     "ip6Address",
		Ip6Gateway:     "ip6Gateway",
		IpAddress:      "ipAddress",
		IpGateway:      "ipGateway",
		IpNetmask:      "ipNetmask",
		LastUpdated:    "lastUpdated",
		MgmtIpAddress:  "mgmtIpAddress",
		MgmtIpGateway:  "mgmtIpGateway",
		MgmtIpNetmask:  "mgmtIpNetmask",
		OfflineReason:  "offlineReason",
		PhysLocation:   "physLocation",
		PhysLocationId: 1,
		Profile:        "profile",
		ProfileDesc:    "profileDesc",
		ProfileId:      1,
		Rack:           "rack",
		RevalPending:   true,
		RouterHostName: "routerHostName",
		RouterPortName: "routerPortName",
		Status:         "status",
		StatusId:       1,
		TcpPort:        80,
		ServerType:     "EDGE",
		ServerTypeId:   1,
		UpdPending:     true,
		XmppId:         "xmppId",
		XmppPasswd:     "xmppPasswd",
	}
	servers = append(servers, testServer)

	testServer2 := testServer
	testServer2.Cachegroup = "cachegroup2"
	testServer2.HostName = "server2"
	servers = append(servers, testServer2)

	return servers
}

func TestGetServersByDsId(t *testing.T) {
	mockDB, mock, err := sqlmock.New()
	defer mockDB.Close()
	db := sqlx.NewDb(mockDB, "sqlmock")
	if err != nil {
		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
	}
	defer db.Close()

	testServers := getTestServers()
	cols := ColsFromStructByTag("db", Server{})
	rows := sqlmock.NewRows(cols)

	//TODO: drichardson - build helper to add these Rows from the struct values
	//                    or by CSV if types get in the way
	for _, ts := range testServers {
		rows = rows.AddRow(
			ts.Cachegroup,
			ts.CachegroupId,
			ts.CdnId,
			ts.CdnName,
			ts.DomainName,
			ts.Guid,
			ts.HostName,
			ts.HttpsPort,
			ts.Id,
			ts.IloIpAddress,
			ts.IloIpGateway,
			ts.IloIpNetmask,
			ts.IloPassword,
			ts.IloUsername,
			ts.InterfaceMtu,
			ts.InterfaceName,
			ts.Ip6Address,
			ts.Ip6Gateway,
			ts.IpAddress,
			ts.IpNetmask,
			ts.IpGateway,
			ts.LastUpdated,
			ts.MgmtIpAddress,
			ts.MgmtIpGateway,
			ts.MgmtIpNetmask,
			ts.OfflineReason,
			ts.PhysLocation,
			ts.PhysLocationId,
			ts.Profile,
			ts.ProfileDesc,
			ts.ProfileId,
			ts.Rack,
			ts.RevalPending,
			ts.RouterHostName,
			ts.RouterPortName,
			ts.Status,
			ts.StatusId,
			ts.TcpPort,
			ts.ServerType,
			ts.ServerTypeId,
			ts.UpdPending,
			ts.XmppId,
			ts.XmppPasswd,
		)
	}
	mock.ExpectQuery("SELECT").WillReturnRows(rows)
	v := url.Values{}
	v.Set("dsId", "1")

	servers, err := getServers(v, db, PrivLevelAdmin)
	if err != nil {
		t.Errorf("getServers expected: nil error, actual: %v", err)
	}

	if len(servers) != 2 {
		t.Errorf("getServers expected: len(servers) == 1, actual: %v", len(servers))
	}

}

type SortableServers []Server

func (s SortableServers) Len() int {
	return len(s)
}
func (s SortableServers) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s SortableServers) Less(i, j int) bool {
	return s[i].HostName < s[j].HostName
}
<|endoftext|>


<commit_before>package app

import (
	"cf/terminal"
	"github.com/codegangsta/cli"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"
)

var appHelpTemplate = `{{.Title "NAME:"}}
   {{.Name}} - {{.Usage}}

{{.Title "USAGE:"}}
   [environment variables] {{.Name}} [global options] command [arguments...] [command options]

{{.Title "VERSION:"}}
   {{.Version}}
   {{range .Commands}}
{{.SubTitle .Name}}{{range .CommandSubGroups}}
{{range .}}   {{.Name}} {{.Description}}
{{end}}{{end}}{{end}}
{{.Title "GLOBAL OPTIONS:"}}
   {{range .Flags}}{{.}}
   {{end}}
{{.Title "ENVIRONMENT VARIABLES:"}}
   CF_STAGING_TIMEOUT=15 max wait time for buildpack staging, in minutes
   CF_STARTUP_TIMEOUT=5 max wait time for app instance startup, in minutes
   CF_TRACE=true - print API request diagnostics to stdout
   CF_TRACE=path/to/trace.log - append API request diagnostics to a log file
   HTTP_PROXY=http://proxy.example.com:8080 - enable http proxying for API requests
`

type groupedCommands struct {
	Name             string
	CommandSubGroups [][]cmdPresenter
}

func (c groupedCommands) SubTitle(name string) string {
	return terminal.HeaderColor(name + ":")
}

type cmdPresenter struct {
	Name        string
	Description string
}

func newCmdPresenter(app *cli.App, maxNameLen int, cmdName string) (presenter cmdPresenter) {
	cmd := app.Command(cmdName)

	presenter.Name = presentCmdName(*cmd)
	padding := strings.Repeat(" ", maxNameLen-len(presenter.Name))
	presenter.Name = presenter.Name + padding

	presenter.Description = cmd.Description

	return
}

func presentCmdName(cmd cli.Command) (name string) {
	name = cmd.Name
	if cmd.ShortName != "" {
		name = name + ", " + cmd.ShortName
	}
	return
}

type appPresenter struct {
	cli.App
	Commands []groupedCommands
}

func (p appPresenter) Title(name string) string {
	return terminal.HeaderColor(name)
}

func getMaxCmdNameLength(app *cli.App) (length int) {
	for _, cmd := range app.Commands {
		name := presentCmdName(cmd)
		if len(name) > length {
			length = len(name)
		}
	}
	return
}

func newAppPresenter(app *cli.App) (presenter appPresenter) {
	maxNameLen := getMaxCmdNameLength(app)

	presenter.Name = app.Name
	presenter.Usage = app.Usage
	presenter.Version = app.Version
	presenter.Name = app.Name
	presenter.Flags = app.Flags

	presenter.Commands = []groupedCommands{
		{
			Name: "GETTING STARTED",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "login"),
					newCmdPresenter(app, maxNameLen, "logout"),
					newCmdPresenter(app, maxNameLen, "passwd"),
					newCmdPresenter(app, maxNameLen, "target"),
				}, {
					newCmdPresenter(app, maxNameLen, "api"),
					newCmdPresenter(app, maxNameLen, "auth"),
				},
			},
		}, {
			Name: "APPS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "apps"),
					newCmdPresenter(app, maxNameLen, "app"),
				}, {
					newCmdPresenter(app, maxNameLen, "push"),
					newCmdPresenter(app, maxNameLen, "scale"),
					newCmdPresenter(app, maxNameLen, "delete"),
					newCmdPresenter(app, maxNameLen, "rename"),
				}, {
					newCmdPresenter(app, maxNameLen, "start"),
					newCmdPresenter(app, maxNameLen, "stop"),
					newCmdPresenter(app, maxNameLen, "restart"),
				}, {
					newCmdPresenter(app, maxNameLen, "events"),
					newCmdPresenter(app, maxNameLen, "files"),
					newCmdPresenter(app, maxNameLen, "logs"),
				}, {
					newCmdPresenter(app, maxNameLen, "env"),
					newCmdPresenter(app, maxNameLen, "set-env"),
					newCmdPresenter(app, maxNameLen, "unset-env"),
				}, {
					newCmdPresenter(app, maxNameLen, "stacks"),
				},
			},
		}, {
			Name: "SERVICES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "marketplace"),
					newCmdPresenter(app, maxNameLen, "services"),
					newCmdPresenter(app, maxNameLen, "service"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-service"),
					newCmdPresenter(app, maxNameLen, "delete-service"),
					newCmdPresenter(app, maxNameLen, "rename-service"),
				}, {
					newCmdPresenter(app, maxNameLen, "bind-service"),
					newCmdPresenter(app, maxNameLen, "unbind-service"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-user-provided-service"),
					newCmdPresenter(app, maxNameLen, "update-user-provided-service"),
				},
			},
		}, {
			Name: "ORGS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "orgs"),
					newCmdPresenter(app, maxNameLen, "org"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-org"),
					newCmdPresenter(app, maxNameLen, "delete-org"),
					newCmdPresenter(app, maxNameLen, "rename-org"),
				},
			},
		}, {
			Name: "SPACES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "spaces"),
					newCmdPresenter(app, maxNameLen, "space"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-space"),
					newCmdPresenter(app, maxNameLen, "delete-space"),
					newCmdPresenter(app, maxNameLen, "rename-space"),
				},
			},
		}, {
			Name: "DOMAINS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "domains"),
					newCmdPresenter(app, maxNameLen, "create-domain"),
					newCmdPresenter(app, maxNameLen, "share-domain"),
					newCmdPresenter(app, maxNameLen, "delete-domain"),
				},
			},
		}, {
			Name: "ROUTES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "routes"),
					newCmdPresenter(app, maxNameLen, "create-route"),
					newCmdPresenter(app, maxNameLen, "map-route"),
					newCmdPresenter(app, maxNameLen, "unmap-route"),
					newCmdPresenter(app, maxNameLen, "delete-route"),
				},
			},
		}, {
			Name: "BUILDPACKS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "buildpacks"),
					newCmdPresenter(app, maxNameLen, "create-buildpack"),
					newCmdPresenter(app, maxNameLen, "update-buildpack"),
					newCmdPresenter(app, maxNameLen, "delete-buildpack"),
				},
			},
		}, {
			Name: "USER ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "create-user"),
					newCmdPresenter(app, maxNameLen, "delete-user"),
				}, {
					newCmdPresenter(app, maxNameLen, "org-users"),
					newCmdPresenter(app, maxNameLen, "set-org-role"),
					newCmdPresenter(app, maxNameLen, "unset-org-role"),
				}, {
					newCmdPresenter(app, maxNameLen, "space-users"),
					newCmdPresenter(app, maxNameLen, "set-space-role"),
					newCmdPresenter(app, maxNameLen, "unset-space-role"),
				},
			},
		}, {
			Name: "ORG ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "quotas"),
					newCmdPresenter(app, maxNameLen, "set-quota"),
				},
			},
		}, {
			Name: "SERVICE ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "service-auth-tokens"),
					newCmdPresenter(app, maxNameLen, "create-service-auth-token"),
					newCmdPresenter(app, maxNameLen, "update-service-auth-token"),
					newCmdPresenter(app, maxNameLen, "delete-service-auth-token"),
				}, {
					newCmdPresenter(app, maxNameLen, "service-brokers"),
					newCmdPresenter(app, maxNameLen, "create-service-broker"),
					newCmdPresenter(app, maxNameLen, "update-service-broker"),
					newCmdPresenter(app, maxNameLen, "delete-service-broker"),
					newCmdPresenter(app, maxNameLen, "rename-service-broker"),
				},
			},
		}, {
			Name: "ADVANCED",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "curl"),
				},
			},
		},
	}
	return
}

func showAppHelp(app *cli.App) {
	presenter := newAppPresenter(app)

	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
	t := template.Must(template.New("help").Parse(appHelpTemplate))
	t.Execute(w, presenter)
	w.Flush()
}
<commit_msg>list CF_COLOR in the right file, fixing #56<commit_after>package app

import (
	"cf/terminal"
	"github.com/codegangsta/cli"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"
)

var appHelpTemplate = `{{.Title "NAME:"}}
   {{.Name}} - {{.Usage}}

{{.Title "USAGE:"}}
   [environment variables] {{.Name}} [global options] command [arguments...] [command options]

{{.Title "VERSION:"}}
   {{.Version}}
   {{range .Commands}}
{{.SubTitle .Name}}{{range .CommandSubGroups}}
{{range .}}   {{.Name}} {{.Description}}
{{end}}{{end}}{{end}}
{{.Title "GLOBAL OPTIONS:"}}
   {{range .Flags}}{{.}}
   {{end}}
{{.Title "ENVIRONMENT VARIABLES:"}}
   CF_STAGING_TIMEOUT=15 max wait time for buildpack staging, in minutes
   CF_STARTUP_TIMEOUT=5 max wait time for app instance startup, in minutes
   CF_COLOR=false - will not colorize output
   CF_TRACE=true - print API request diagnostics to stdout
   CF_TRACE=path/to/trace.log - append API request diagnostics to a log file
   HTTP_PROXY=http://proxy.example.com:8080 - enable HTTP proxying for API requests
`

type groupedCommands struct {
	Name             string
	CommandSubGroups [][]cmdPresenter
}

func (c groupedCommands) SubTitle(name string) string {
	return terminal.HeaderColor(name + ":")
}

type cmdPresenter struct {
	Name        string
	Description string
}

func newCmdPresenter(app *cli.App, maxNameLen int, cmdName string) (presenter cmdPresenter) {
	cmd := app.Command(cmdName)

	presenter.Name = presentCmdName(*cmd)
	padding := strings.Repeat(" ", maxNameLen-len(presenter.Name))
	presenter.Name = presenter.Name + padding

	presenter.Description = cmd.Description

	return
}

func presentCmdName(cmd cli.Command) (name string) {
	name = cmd.Name
	if cmd.ShortName != "" {
		name = name + ", " + cmd.ShortName
	}
	return
}

type appPresenter struct {
	cli.App
	Commands []groupedCommands
}

func (p appPresenter) Title(name string) string {
	return terminal.HeaderColor(name)
}

func getMaxCmdNameLength(app *cli.App) (length int) {
	for _, cmd := range app.Commands {
		name := presentCmdName(cmd)
		if len(name) > length {
			length = len(name)
		}
	}
	return
}

func newAppPresenter(app *cli.App) (presenter appPresenter) {
	maxNameLen := getMaxCmdNameLength(app)

	presenter.Name = app.Name
	presenter.Usage = app.Usage
	presenter.Version = app.Version
	presenter.Name = app.Name
	presenter.Flags = app.Flags

	presenter.Commands = []groupedCommands{
		{
			Name: "GETTING STARTED",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "login"),
					newCmdPresenter(app, maxNameLen, "logout"),
					newCmdPresenter(app, maxNameLen, "passwd"),
					newCmdPresenter(app, maxNameLen, "target"),
				}, {
					newCmdPresenter(app, maxNameLen, "api"),
					newCmdPresenter(app, maxNameLen, "auth"),
				},
			},
		}, {
			Name: "APPS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "apps"),
					newCmdPresenter(app, maxNameLen, "app"),
				}, {
					newCmdPresenter(app, maxNameLen, "push"),
					newCmdPresenter(app, maxNameLen, "scale"),
					newCmdPresenter(app, maxNameLen, "delete"),
					newCmdPresenter(app, maxNameLen, "rename"),
				}, {
					newCmdPresenter(app, maxNameLen, "start"),
					newCmdPresenter(app, maxNameLen, "stop"),
					newCmdPresenter(app, maxNameLen, "restart"),
				}, {
					newCmdPresenter(app, maxNameLen, "events"),
					newCmdPresenter(app, maxNameLen, "files"),
					newCmdPresenter(app, maxNameLen, "logs"),
				}, {
					newCmdPresenter(app, maxNameLen, "env"),
					newCmdPresenter(app, maxNameLen, "set-env"),
					newCmdPresenter(app, maxNameLen, "unset-env"),
				}, {
					newCmdPresenter(app, maxNameLen, "stacks"),
				},
			},
		}, {
			Name: "SERVICES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "marketplace"),
					newCmdPresenter(app, maxNameLen, "services"),
					newCmdPresenter(app, maxNameLen, "service"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-service"),
					newCmdPresenter(app, maxNameLen, "delete-service"),
					newCmdPresenter(app, maxNameLen, "rename-service"),
				}, {
					newCmdPresenter(app, maxNameLen, "bind-service"),
					newCmdPresenter(app, maxNameLen, "unbind-service"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-user-provided-service"),
					newCmdPresenter(app, maxNameLen, "update-user-provided-service"),
				},
			},
		}, {
			Name: "ORGS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "orgs"),
					newCmdPresenter(app, maxNameLen, "org"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-org"),
					newCmdPresenter(app, maxNameLen, "delete-org"),
					newCmdPresenter(app, maxNameLen, "rename-org"),
				},
			},
		}, {
			Name: "SPACES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "spaces"),
					newCmdPresenter(app, maxNameLen, "space"),
				}, {
					newCmdPresenter(app, maxNameLen, "create-space"),
					newCmdPresenter(app, maxNameLen, "delete-space"),
					newCmdPresenter(app, maxNameLen, "rename-space"),
				},
			},
		}, {
			Name: "DOMAINS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "domains"),
					newCmdPresenter(app, maxNameLen, "create-domain"),
					newCmdPresenter(app, maxNameLen, "share-domain"),
					newCmdPresenter(app, maxNameLen, "delete-domain"),
				},
			},
		}, {
			Name: "ROUTES",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "routes"),
					newCmdPresenter(app, maxNameLen, "create-route"),
					newCmdPresenter(app, maxNameLen, "map-route"),
					newCmdPresenter(app, maxNameLen, "unmap-route"),
					newCmdPresenter(app, maxNameLen, "delete-route"),
				},
			},
		}, {
			Name: "BUILDPACKS",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "buildpacks"),
					newCmdPresenter(app, maxNameLen, "create-buildpack"),
					newCmdPresenter(app, maxNameLen, "update-buildpack"),
					newCmdPresenter(app, maxNameLen, "delete-buildpack"),
				},
			},
		}, {
			Name: "USER ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "create-user"),
					newCmdPresenter(app, maxNameLen, "delete-user"),
				}, {
					newCmdPresenter(app, maxNameLen, "org-users"),
					newCmdPresenter(app, maxNameLen, "set-org-role"),
					newCmdPresenter(app, maxNameLen, "unset-org-role"),
				}, {
					newCmdPresenter(app, maxNameLen, "space-users"),
					newCmdPresenter(app, maxNameLen, "set-space-role"),
					newCmdPresenter(app, maxNameLen, "unset-space-role"),
				},
			},
		}, {
			Name: "ORG ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "quotas"),
					newCmdPresenter(app, maxNameLen, "set-quota"),
				},
			},
		}, {
			Name: "SERVICE ADMIN",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "service-auth-tokens"),
					newCmdPresenter(app, maxNameLen, "create-service-auth-token"),
					newCmdPresenter(app, maxNameLen, "update-service-auth-token"),
					newCmdPresenter(app, maxNameLen, "delete-service-auth-token"),
				}, {
					newCmdPresenter(app, maxNameLen, "service-brokers"),
					newCmdPresenter(app, maxNameLen, "create-service-broker"),
					newCmdPresenter(app, maxNameLen, "update-service-broker"),
					newCmdPresenter(app, maxNameLen, "delete-service-broker"),
					newCmdPresenter(app, maxNameLen, "rename-service-broker"),
				},
			},
		}, {
			Name: "ADVANCED",
			CommandSubGroups: [][]cmdPresenter{
				{
					newCmdPresenter(app, maxNameLen, "curl"),
				},
			},
		},
	}
	return
}

func showAppHelp(app *cli.App) {
	presenter := newAppPresenter(app)

	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
	t := template.Must(template.New("help").Parse(appHelpTemplate))
	t.Execute(w, presenter)
	w.Flush()
}
<|endoftext|>


<commit_before>package termite

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"http"
	"io/ioutil"
	"log"
	"os"
	"rand"
	"rpc"
	"testing"
	"time"
)

type testCase struct {
	worker *WorkerDaemon
	master *Master
	coordinator *Coordinator
	secret []byte
	tmp    string
	socket string
	coordinatorPort int
	workerPort int
}

func NewTestCase() *testCase {
	me := new(testCase)
	me.secret = RandomBytes(20)
	me.tmp, _ = ioutil.TempDir("", "")

	workerTmp := me.tmp + "/worker-tmp"
	os.Mkdir(workerTmp, 0700)
	me.worker = NewWorkerDaemon(me.secret, workerTmp,
		me.tmp+"/worker-cache", 1)

	// TODO - pick unused port
	me.coordinatorPort = int(rand.Int31n(60000) + 1024)
	c := NewCoordinator()
	rpc.Register(c)
	rpc.HandleHTTP()
	go c.PeriodicCheck()

	coordinatorAddr := fmt.Sprintf(":%d", me.coordinatorPort)
	go http.ListenAndServe(coordinatorAddr, nil)
	// TODO - can we do without the sleeps?
	time.Sleep(0.1e9) // wait for daemon to start up

	me.workerPort = int(rand.Int31n(60000) + 1024)
	go me.worker.RunWorkerServer(me.workerPort, coordinatorAddr)

	// wait worker to be registered on coordinator.
	time.Sleep(0.1e9)

	masterCache := NewContentCache(me.tmp + "/master-cache")
	me.master = NewMaster(
		masterCache, coordinatorAddr,
		[]string{},
		me.secret, []string{}, 1)

	me.master.SetKeepAlive(1.0)
	me.socket = me.tmp + "/master-socket"
	go me.master.Start(me.socket)

	wd := me.tmp + "/wd"
	os.MkdirAll(wd, 0755)
	time.Sleep(0.1e9) // wait for all daemons to start up
	return me
}

func (me *testCase) Clean() {
	me.master.mirrors.dropConnections()
	// TODO - should have explicit worker shutdown routine.
	time.Sleep(0.1e9)
	log.Println("removing.")
	os.RemoveAll(me.tmp)
}

// Simple end-to-end test.  It skips the chroot, but should give a
// basic assurance that things work as expected.
func TestEndToEndBasic(t *testing.T) {
	if os.Geteuid() == 0 {
		log.Println("This test should not run as root")
		return
	}

	tc := NewTestCase()
	defer tc.Clean()

	req := WorkRequest{
		StdinId: ConnectionId(),
		Binary:  "/usr/bin/tee",
		Argv:    []string{"/usr/bin/tee", "output.txt"},
		Env:     os.Environ(),

		// Will not be filtered, since /tmp/foo is more
		// specific than /tmp
		Dir: tc.tmp + "/wd",
		Debug: true,
	}

	// TODO - should separate dial/listen in the daemons?
	stdinConn := OpenSocketConnection(tc.socket, req.StdinId)
	go func() {
		stdinConn.Write([]byte("hello"))
		stdinConn.Close()
	}()

	rpcConn := OpenSocketConnection(tc.socket, RPC_CHANNEL)
	client := rpc.NewClient(rpcConn)

	rep := WorkReply{}
	err := client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		log.Fatal("LocalMaster.Run: ", err)
	}

	content, err := ioutil.ReadFile(tc.tmp + "/wd/output.txt")
	if err != nil {
		t.Error(err)
	}
	if string(content) != "hello" {
		t.Error("content:", content)
	}

	req = WorkRequest{
		Binary:  "/bin/rm",
		Argv:    []string{"/bin/rm", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: true,
	}

	rep = WorkReply{}
	err = client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}
	if fi, _ := os.Lstat(tc.tmp + "/wd/output.txt"); fi != nil {
		t.Error("file should have been deleted", fi)
	}

	// Test keepalive.
	time.Sleep(2e9)

	statusReq := &StatusRequest{}
	statusRep := &StatusReply{}
	tc.worker.Status(statusReq, statusRep)
	if statusRep.Processes != 0 {
		t.Error("Processes still alive.")
	}
}

// This shows a case that is not handled correctly yet: we have no way
// to flush the cache on negative entries.
func TestEndToEndNegativeNotify(t *testing.T) {
	if os.Geteuid() == 0 {
		log.Println("This test should not run as root")
		return
	}

	tc := NewTestCase()
	defer tc.Clean()

	rpcConn := OpenSocketConnection(tc.socket, RPC_CHANNEL)
	client := rpc.NewClient(rpcConn)

	req := WorkRequest{
		Binary:  "/bin/cat",
		Argv:    []string{"/bin/cat", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: true,
	}

	rep := WorkReply{}
	err := client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}

	if rep.Exit.ExitStatus() == 0 {
		t.Fatal("expect exit status != 0")
	}

	newContent := []byte("new content")
	hash := tc.master.cache.Save(newContent)
	updated := []FileAttr{
		FileAttr{
			Path:  tc.tmp + "/wd/output.txt",
			FileInfo: &os.FileInfo{Mode: fuse.S_IFREG | 0644, Size: int64(len(newContent))},
			Hash: hash,
			Content: newContent,
		},
	}
	tc.master.mirrors.queueFiles(nil, updated)
	req = WorkRequest{
		Binary:  "/bin/cat",
		Argv:    []string{"/bin/cat", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: 	 true,
	}

	err = client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}
	if rep.Exit.ExitStatus() != 0 {
		t.Fatal("expect exit status == 0")
	}
	log.Println("new content:", rep.Stdout)
	if string(rep.Stdout) != string(newContent) {
		t.Error("Mismatch", string(rep.Stdout), string(newContent))
	}
}
<commit_msg>Oops; reinit WorkReply.<commit_after>package termite

import (
	"fmt"
	"github.com/hanwen/go-fuse/fuse"
	"http"
	"io/ioutil"
	"log"
	"os"
	"rand"
	"rpc"
	"testing"
	"time"
)

type testCase struct {
	worker *WorkerDaemon
	master *Master
	coordinator *Coordinator
	secret []byte
	tmp    string
	socket string
	coordinatorPort int
	workerPort int
}

func NewTestCase() *testCase {
	me := new(testCase)
	me.secret = RandomBytes(20)
	me.tmp, _ = ioutil.TempDir("", "")

	workerTmp := me.tmp + "/worker-tmp"
	os.Mkdir(workerTmp, 0700)
	me.worker = NewWorkerDaemon(me.secret, workerTmp,
		me.tmp+"/worker-cache", 1)

	// TODO - pick unused port
	me.coordinatorPort = int(rand.Int31n(60000) + 1024)
	c := NewCoordinator()
	rpc.Register(c)
	rpc.HandleHTTP()
	go c.PeriodicCheck()

	coordinatorAddr := fmt.Sprintf(":%d", me.coordinatorPort)
	go http.ListenAndServe(coordinatorAddr, nil)
	// TODO - can we do without the sleeps?
	time.Sleep(0.1e9) // wait for daemon to start up

	me.workerPort = int(rand.Int31n(60000) + 1024)
	go me.worker.RunWorkerServer(me.workerPort, coordinatorAddr)

	// wait worker to be registered on coordinator.
	time.Sleep(0.1e9)

	masterCache := NewContentCache(me.tmp + "/master-cache")
	me.master = NewMaster(
		masterCache, coordinatorAddr,
		[]string{},
		me.secret, []string{}, 1)

	me.master.SetKeepAlive(1.0)
	me.socket = me.tmp + "/master-socket"
	go me.master.Start(me.socket)

	wd := me.tmp + "/wd"
	os.MkdirAll(wd, 0755)
	time.Sleep(0.1e9) // wait for all daemons to start up
	return me
}

func (me *testCase) Clean() {
	me.master.mirrors.dropConnections()
	// TODO - should have explicit worker shutdown routine.
	time.Sleep(0.1e9)
	log.Println("removing.")
	os.RemoveAll(me.tmp)
}

// Simple end-to-end test.  It skips the chroot, but should give a
// basic assurance that things work as expected.
func TestEndToEndBasic(t *testing.T) {
	if os.Geteuid() == 0 {
		log.Println("This test should not run as root")
		return
	}

	tc := NewTestCase()
	defer tc.Clean()

	req := WorkRequest{
		StdinId: ConnectionId(),
		Binary:  "/usr/bin/tee",
		Argv:    []string{"/usr/bin/tee", "output.txt"},
		Env:     os.Environ(),

		// Will not be filtered, since /tmp/foo is more
		// specific than /tmp
		Dir: tc.tmp + "/wd",
		Debug: true,
	}

	// TODO - should separate dial/listen in the daemons?
	stdinConn := OpenSocketConnection(tc.socket, req.StdinId)
	go func() {
		stdinConn.Write([]byte("hello"))
		stdinConn.Close()
	}()

	rpcConn := OpenSocketConnection(tc.socket, RPC_CHANNEL)
	client := rpc.NewClient(rpcConn)

	rep := WorkReply{}
	err := client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		log.Fatal("LocalMaster.Run: ", err)
	}

	content, err := ioutil.ReadFile(tc.tmp + "/wd/output.txt")
	if err != nil {
		t.Error(err)
	}
	if string(content) != "hello" {
		t.Error("content:", content)
	}

	req = WorkRequest{
		Binary:  "/bin/rm",
		Argv:    []string{"/bin/rm", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: true,
	}

	rep = WorkReply{}
	err = client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}
	if fi, _ := os.Lstat(tc.tmp + "/wd/output.txt"); fi != nil {
		t.Error("file should have been deleted", fi)
	}

	// Test keepalive.
	time.Sleep(2e9)

	statusReq := &StatusRequest{}
	statusRep := &StatusReply{}
	tc.worker.Status(statusReq, statusRep)
	if statusRep.Processes != 0 {
		t.Error("Processes still alive.")
	}
}

// This shows a case that is not handled correctly yet: we have no way
// to flush the cache on negative entries.
func TestEndToEndNegativeNotify(t *testing.T) {
	if os.Geteuid() == 0 {
		log.Println("This test should not run as root")
		return
	}

	tc := NewTestCase()
	defer tc.Clean()

	rpcConn := OpenSocketConnection(tc.socket, RPC_CHANNEL)
	client := rpc.NewClient(rpcConn)

	req := WorkRequest{
		Binary:  "/bin/cat",
		Argv:    []string{"/bin/cat", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: true,
	}

	rep := WorkReply{}
	err := client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}

	if rep.Exit.ExitStatus() == 0 {
		t.Fatal("expect exit status != 0")
	}

	newContent := []byte("new content")
	hash := tc.master.cache.Save(newContent)
	updated := []FileAttr{
		FileAttr{
			Path:  tc.tmp + "/wd/output.txt",
			FileInfo: &os.FileInfo{Mode: fuse.S_IFREG | 0644, Size: int64(len(newContent))},
			Hash: hash,
			Content: newContent,
		},
	}
	tc.master.mirrors.queueFiles(nil, updated)
	req = WorkRequest{
		Binary:  "/bin/cat",
		Argv:    []string{"/bin/cat", "output.txt"},
		Env:     os.Environ(),
		Dir:     tc.tmp + "/wd",
		Debug: 	 true,
	}

	rep = WorkReply{}
	err = client.Call("LocalMaster.Run", &req, &rep)
	if err != nil {
		t.Fatal("LocalMaster.Run: ", err)
	}
	if rep.Exit.ExitStatus() != 0 {
		t.Fatal("expect exit status == 0", rep.Exit.ExitStatus())
	}
	log.Println("new content:", rep.Stdout)
	if string(rep.Stdout) != string(newContent) {
		t.Error("Mismatch", string(rep.Stdout), string(newContent))
	}
}
<|endoftext|>