<commit_before>package raft

import (
	"fmt"
	"testing"
	"time"
)

const (
	testThisServerId = "s1"
	// Note: value for tests based on Figure 7
	testCurrentTerm = 8

	testTickerDuration     = 15 * time.Millisecond
	testElectionTimeoutLow = 150 * time.Millisecond

	testSleepToLetGoroutineRun = 3 * time.Millisecond
	testSleepJustMoreThanATick = testTickerDuration + testSleepToLetGoroutineRun
)

var testAllServerIds = []ServerId{testThisServerId, "s2", "s3", "s4", "s5"}

func setupManagedConsensusModule(t *testing.T, logTerms []TermNo) *managedConsensusModule {
	mcm, _ := setupManagedConsensusModuleR2(t, logTerms)
	return mcm
}

func setupManagedConsensusModuleR2(
	t *testing.T,
	logTerms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	ps := newIMPSWithCurrentTerm(testCurrentTerm)
	imle := newIMLEWithDummyCommands(logTerms)
	mrs := newMockRpcSender()
	ts := TimeSettings{testTickerDuration, testElectionTimeoutLow}
	ci := NewClusterInfo(testAllServerIds, testThisServerId)
	cm, now := newPassiveConsensusModule(ps, imle, mrs, ci, ts)
	if cm == nil {
		t.Fatal()
	}
	// Bias simulated clock to avoid exact time matches
	now = now.Add(testSleepToLetGoroutineRun)
	mcm := &managedConsensusModule{cm, now}
	return mcm, mrs
}

// #5.2-p1s2: When servers start up, they begin as followers
func TestCM_InitialState(t *testing.T) {
	mcm := setupManagedConsensusModule(t, nil)

	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}

	// Volatile state on all servers
	if mcm.pcm.volatileState != (volatileState{0, 0}) {
		t.Fatal()
	}
}

func test_ExpectPanic(t *testing.T, f func(), expectedRecover interface{}) {
	skipRecover := false
	defer func() {
		if !skipRecover {
			if r := recover(); r != expectedRecover {
				t.Fatal(fmt.Sprintf("Expected panic: %v; got: %v", expectedRecover, r))
			}
		}
	}()

	f()
	skipRecover = true
	t.Fatal(fmt.Sprintf("Expected panic: %v", expectedRecover))
}

func testCM_setupMCMAndExpectPanicFor(
	t *testing.T,
	f func(*managedConsensusModule),
	expectedRecover interface{},
) {
	mcm := setupManagedConsensusModule(t, nil)
	test_ExpectPanic(
		t,
		func() {
			f(mcm)
		},
		expectedRecover,
	)
}

func TestCM_UnknownRpcTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.rpc("s2", &struct{ int }{42})
		},
		"FATAL: unknown rpc type: *struct { int } from: s2",
	)
}

func TestCM_RpcReply_UnknownRpcTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm.rpcReply("s2", &struct{ int }{42}, &struct{ int }{42})
		},
		"FATAL: unknown rpc type: *struct { int } from: s2",
	)
}

func TestCM_RpcReply_UnknownRpcReplyTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			sentRpc := &RpcRequestVote{1, 0, 0}
			mcm.pcm.rpcReply("s2", sentRpc, &RpcAppendEntriesReply{1, false})
		},
		"FATAL: mismatched rpcReply type: *raft.RpcAppendEntriesReply from: s2 - expected *raft.RpcRequestVoteReply",
	)
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			sentRpc := &RpcAppendEntries{1, 0, 0, nil, 0}
			mcm.pcm.rpcReply("s2", sentRpc, &RpcRequestVoteReply{1, false})
		},
		"FATAL: mismatched rpcReply type: *raft.RpcRequestVoteReply from: s2 - expected *raft.RpcAppendEntriesReply",
	)
}

func TestCM_SetServerStateBadServerStatePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._setServerState(42)
		},
		"FATAL: unknown ServerState: 42",
	)
}

func TestCM_BadServerStatePanicsTick(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._unsafe_serverState = 42

			mcm.tick()
		},
		"FATAL: unknown ServerState: 42",
	)
}

func TestCM_BadServerStatePanicsRpc(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._unsafe_serverState = 42

			requestVote := &RpcRequestVote{1, 0, 0}

			mcm.rpc("s2", requestVote)
		},
		"FATAL: unknown ServerState: 42",
	)
}

// #5.2-p1s5: If a follower receives no communication over a period of time
// called the election timeout, then it assumes there is no viable leader
// and begins an election to choose a new leader.
// #5.2-p2s1: To begin an election, a follower increments its current term
// and transitions to candidate state.
// #5.2-p2s2: It then votes for itself and issues RequestVote RPCs in parallel
// to each of the other servers in the cluster.
// #5.2-p6s2: ..., election timeouts are chosen randomly from a fixed
// interval (e.g., 150-300ms)
func testCM_Follower_StartsElectionOnElectionTimeout(
	t *testing.T,
	mcm *managedConsensusModule,
	mrs *mockRpcSender,
) {

	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	if mcm.pcm.persistentState.GetVotedFor() != "" {
		t.Fatal()
	}

	// Test that a tick before election timeout causes no state change.
	mcm.tick()
	if mcm.pcm.persistentState.GetCurrentTerm() != testCurrentTerm {
		t.Fatal()
	}
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}

	testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(t, mcm, mrs, testCurrentTerm+1)
}

func testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(
	t *testing.T,
	mcm *managedConsensusModule,
	mrs *mockRpcSender,
	expectedNewTerm TermNo,
) {
	// Test that election timeout causes a new election
	mcm.tickTilElectionTimeout()
	if mcm.pcm.persistentState.GetCurrentTerm() != expectedNewTerm {
		t.Fatal(expectedNewTerm, mcm.pcm.persistentState.GetCurrentTerm())
	}
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}
	// candidate has voted for itself
	if mcm.pcm.persistentState.GetVotedFor() != testThisServerId {
		t.Fatal()
	}

	// candidate has issued RequestVote RPCs to all other servers.
	lastLogIndex, lastLogTerm := getIndexAndTermOfLastEntry(mcm.pcm.log)
	expectedRpc := &RpcRequestVote{expectedNewTerm, lastLogIndex, lastLogTerm}
	expectedRpcs := []mockSentRpc{
		{"s2", expectedRpc},
		{"s3", expectedRpc},
		{"s4", expectedRpc},
		{"s5", expectedRpc},
	}
	mrs.checkSentRpcs(t, expectedRpcs)
}

func TestCM_Follower_StartsElectionOnElectionTimeout_EmptyLog(t *testing.T) {
	mcm, mrs := setupManagedConsensusModuleR2(t, nil)
	testCM_Follower_StartsElectionOnElectionTimeout(t, mcm, mrs)
}

// #RFS-L1b: repeat during idle periods to prevent election timeout (#5.2)
func TestCM_Leader_SendEmptyAppendEntriesDuringIdlePeriods(t *testing.T) {
	mcm, mrs := testSetupMCM_Leader_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	mrs.checkSentRpcs(t, []mockSentRpc{})

	mcm.tick()

	testIsLeaderWithTermAndSentEmptyAppendEntries(t, mcm, mrs, serverTerm)
}

func testSetupMCM_Follower_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := setupManagedConsensusModuleR2(t, terms)
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	return mcm, mrs
}

func testSetupMCM_Candidate_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := testSetupMCM_Follower_WithTerms(t, terms)
	testCM_Follower_StartsElectionOnElectionTimeout(t, mcm, mrs)
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}
	return mcm, mrs
}

func testSetupMCM_Leader_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := testSetupMCM_Candidate_WithTerms(t, terms)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()
	sentRpc := &RpcRequestVote{serverTerm, 0, 0}
	mcm.pcm.rpcReply("s2", sentRpc, &RpcRequestVoteReply{serverTerm, true})
	mcm.pcm.rpcReply("s3", sentRpc, &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}
	mrs.clearSentRpcs()
	return mcm, mrs
}

func testSetupMCM_Follower_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Follower_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func testSetupMCM_Candidate_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Candidate_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func testSetupMCM_Leader_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Leader_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func TestCM_Follower_StartsElectionOnElectionTimeout_NonEmptyLog(t *testing.T) {
	testSetupMCM_Candidate_Figure7LeaderLine(t)
}

// For most tests, we'll use a passive CM where we control the progress
// of time with helper methods. This simplifies tests and avoids concurrency
// issues with inspecting the internals.
type managedConsensusModule struct {
	pcm *passiveConsensusModule
	now time.Time
}

func (mcm *managedConsensusModule) tick() {
	mcm.pcm.tick(mcm.now)
	mcm.now = mcm.now.Add(testTickerDuration)
}

func (mcm *managedConsensusModule) rpc(
	from ServerId,
	rpc interface{},
) interface{} {
	return mcm.pcm.rpc(from, rpc, mcm.now)
}

func (mcm *managedConsensusModule) tickTilElectionTimeout() {
	electionTimeoutTime := mcm.pcm.electionTimeoutTracker.electionTimeoutTime
	for {
		mcm.tick()
		if mcm.now.After(electionTimeoutTime) {
			break
		}
	}
	if mcm.pcm.electionTimeoutTracker.electionTimeoutTime != electionTimeoutTime {
		panic("electionTimeoutTime changed!")
	}
	// Because tick() increments "now" after calling tick(),
	// we need one more to actually run with a post-timeout "now".
	mcm.tick()
}
<commit_msg>Add test that election also chooses new random timeout<commit_after>package raft

import (
	"fmt"
	"testing"
	"time"
)

const (
	testThisServerId = "s1"
	// Note: value for tests based on Figure 7
	testCurrentTerm = 8

	testTickerDuration     = 15 * time.Millisecond
	testElectionTimeoutLow = 150 * time.Millisecond

	testSleepToLetGoroutineRun = 3 * time.Millisecond
	testSleepJustMoreThanATick = testTickerDuration + testSleepToLetGoroutineRun
)

var testAllServerIds = []ServerId{testThisServerId, "s2", "s3", "s4", "s5"}

func setupManagedConsensusModule(t *testing.T, logTerms []TermNo) *managedConsensusModule {
	mcm, _ := setupManagedConsensusModuleR2(t, logTerms)
	return mcm
}

func setupManagedConsensusModuleR2(
	t *testing.T,
	logTerms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	ps := newIMPSWithCurrentTerm(testCurrentTerm)
	imle := newIMLEWithDummyCommands(logTerms)
	mrs := newMockRpcSender()
	ts := TimeSettings{testTickerDuration, testElectionTimeoutLow}
	ci := NewClusterInfo(testAllServerIds, testThisServerId)
	cm, now := newPassiveConsensusModule(ps, imle, mrs, ci, ts)
	if cm == nil {
		t.Fatal()
	}
	// Bias simulated clock to avoid exact time matches
	now = now.Add(testSleepToLetGoroutineRun)
	mcm := &managedConsensusModule{cm, now}
	return mcm, mrs
}

// #5.2-p1s2: When servers start up, they begin as followers
func TestCM_InitialState(t *testing.T) {
	mcm := setupManagedConsensusModule(t, nil)

	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}

	// Volatile state on all servers
	if mcm.pcm.volatileState != (volatileState{0, 0}) {
		t.Fatal()
	}
}

func test_ExpectPanic(t *testing.T, f func(), expectedRecover interface{}) {
	skipRecover := false
	defer func() {
		if !skipRecover {
			if r := recover(); r != expectedRecover {
				t.Fatal(fmt.Sprintf("Expected panic: %v; got: %v", expectedRecover, r))
			}
		}
	}()

	f()
	skipRecover = true
	t.Fatal(fmt.Sprintf("Expected panic: %v", expectedRecover))
}

func testCM_setupMCMAndExpectPanicFor(
	t *testing.T,
	f func(*managedConsensusModule),
	expectedRecover interface{},
) {
	mcm := setupManagedConsensusModule(t, nil)
	test_ExpectPanic(
		t,
		func() {
			f(mcm)
		},
		expectedRecover,
	)
}

func TestCM_UnknownRpcTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.rpc("s2", &struct{ int }{42})
		},
		"FATAL: unknown rpc type: *struct { int } from: s2",
	)
}

func TestCM_RpcReply_UnknownRpcTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm.rpcReply("s2", &struct{ int }{42}, &struct{ int }{42})
		},
		"FATAL: unknown rpc type: *struct { int } from: s2",
	)
}

func TestCM_RpcReply_UnknownRpcReplyTypePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			sentRpc := &RpcRequestVote{1, 0, 0}
			mcm.pcm.rpcReply("s2", sentRpc, &RpcAppendEntriesReply{1, false})
		},
		"FATAL: mismatched rpcReply type: *raft.RpcAppendEntriesReply from: s2 - expected *raft.RpcRequestVoteReply",
	)
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			sentRpc := &RpcAppendEntries{1, 0, 0, nil, 0}
			mcm.pcm.rpcReply("s2", sentRpc, &RpcRequestVoteReply{1, false})
		},
		"FATAL: mismatched rpcReply type: *raft.RpcRequestVoteReply from: s2 - expected *raft.RpcAppendEntriesReply",
	)
}

func TestCM_SetServerStateBadServerStatePanics(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._setServerState(42)
		},
		"FATAL: unknown ServerState: 42",
	)
}

func TestCM_BadServerStatePanicsTick(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._unsafe_serverState = 42

			mcm.tick()
		},
		"FATAL: unknown ServerState: 42",
	)
}

func TestCM_BadServerStatePanicsRpc(t *testing.T) {
	testCM_setupMCMAndExpectPanicFor(
		t,
		func(mcm *managedConsensusModule) {
			mcm.pcm._unsafe_serverState = 42

			requestVote := &RpcRequestVote{1, 0, 0}

			mcm.rpc("s2", requestVote)
		},
		"FATAL: unknown ServerState: 42",
	)
}

// #5.2-p1s5: If a follower receives no communication over a period of time
// called the election timeout, then it assumes there is no viable leader
// and begins an election to choose a new leader.
// #5.2-p2s1: To begin an election, a follower increments its current term
// and transitions to candidate state.
// #5.2-p2s2: It then votes for itself and issues RequestVote RPCs in parallel
// to each of the other servers in the cluster.
// #5.2-p6s2: ..., election timeouts are chosen randomly from a fixed
// interval (e.g., 150-300ms)
func testCM_Follower_StartsElectionOnElectionTimeout(
	t *testing.T,
	mcm *managedConsensusModule,
	mrs *mockRpcSender,
) {

	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	if mcm.pcm.persistentState.GetVotedFor() != "" {
		t.Fatal()
	}

	// Test that a tick before election timeout causes no state change.
	mcm.tick()
	if mcm.pcm.persistentState.GetCurrentTerm() != testCurrentTerm {
		t.Fatal()
	}
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}

	testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(t, mcm, mrs, testCurrentTerm+1)
}

func testCM_FollowerOrCandidate_StartsElectionOnElectionTimeout_Part2(
	t *testing.T,
	mcm *managedConsensusModule,
	mrs *mockRpcSender,
	expectedNewTerm TermNo,
) {
	timeout1 := mcm.pcm.electionTimeoutTracker.currentElectionTimeout
	// Test that election timeout causes a new election
	mcm.tickTilElectionTimeout()
	if mcm.pcm.persistentState.GetCurrentTerm() != expectedNewTerm {
		t.Fatal(expectedNewTerm, mcm.pcm.persistentState.GetCurrentTerm())
	}
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}
	// candidate has voted for itself
	if mcm.pcm.persistentState.GetVotedFor() != testThisServerId {
		t.Fatal()
	}
	// a new election timeout was chosen
	// Playing the odds here :P
	if mcm.pcm.electionTimeoutTracker.currentElectionTimeout == timeout1 {
		t.Fatal()
	}

	// candidate has issued RequestVote RPCs to all other servers.
	lastLogIndex, lastLogTerm := getIndexAndTermOfLastEntry(mcm.pcm.log)
	expectedRpc := &RpcRequestVote{expectedNewTerm, lastLogIndex, lastLogTerm}
	expectedRpcs := []mockSentRpc{
		{"s2", expectedRpc},
		{"s3", expectedRpc},
		{"s4", expectedRpc},
		{"s5", expectedRpc},
	}
	mrs.checkSentRpcs(t, expectedRpcs)
}

func TestCM_Follower_StartsElectionOnElectionTimeout_EmptyLog(t *testing.T) {
	mcm, mrs := setupManagedConsensusModuleR2(t, nil)
	testCM_Follower_StartsElectionOnElectionTimeout(t, mcm, mrs)
}

// #RFS-L1b: repeat during idle periods to prevent election timeout (#5.2)
func TestCM_Leader_SendEmptyAppendEntriesDuringIdlePeriods(t *testing.T) {
	mcm, mrs := testSetupMCM_Leader_Figure7LeaderLine(t)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()

	mrs.checkSentRpcs(t, []mockSentRpc{})

	mcm.tick()

	testIsLeaderWithTermAndSentEmptyAppendEntries(t, mcm, mrs, serverTerm)
}

func testSetupMCM_Follower_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := setupManagedConsensusModuleR2(t, terms)
	if mcm.pcm.getServerState() != FOLLOWER {
		t.Fatal()
	}
	return mcm, mrs
}

func testSetupMCM_Candidate_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := testSetupMCM_Follower_WithTerms(t, terms)
	testCM_Follower_StartsElectionOnElectionTimeout(t, mcm, mrs)
	if mcm.pcm.getServerState() != CANDIDATE {
		t.Fatal()
	}
	return mcm, mrs
}

func testSetupMCM_Leader_WithTerms(
	t *testing.T,
	terms []TermNo,
) (*managedConsensusModule, *mockRpcSender) {
	mcm, mrs := testSetupMCM_Candidate_WithTerms(t, terms)
	serverTerm := mcm.pcm.persistentState.GetCurrentTerm()
	sentRpc := &RpcRequestVote{serverTerm, 0, 0}
	mcm.pcm.rpcReply("s2", sentRpc, &RpcRequestVoteReply{serverTerm, true})
	mcm.pcm.rpcReply("s3", sentRpc, &RpcRequestVoteReply{serverTerm, true})
	if mcm.pcm.getServerState() != LEADER {
		t.Fatal()
	}
	mrs.clearSentRpcs()
	return mcm, mrs
}

func testSetupMCM_Follower_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Follower_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func testSetupMCM_Candidate_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Candidate_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func testSetupMCM_Leader_Figure7LeaderLine(t *testing.T) (*managedConsensusModule, *mockRpcSender) {
	return testSetupMCM_Leader_WithTerms(t, makeLogTerms_Figure7LeaderLine())
}

func TestCM_Follower_StartsElectionOnElectionTimeout_NonEmptyLog(t *testing.T) {
	testSetupMCM_Candidate_Figure7LeaderLine(t)
}

// For most tests, we'll use a passive CM where we control the progress
// of time with helper methods. This simplifies tests and avoids concurrency
// issues with inspecting the internals.
type managedConsensusModule struct {
	pcm *passiveConsensusModule
	now time.Time
}

func (mcm *managedConsensusModule) tick() {
	mcm.pcm.tick(mcm.now)
	mcm.now = mcm.now.Add(testTickerDuration)
}

func (mcm *managedConsensusModule) rpc(
	from ServerId,
	rpc interface{},
) interface{} {
	return mcm.pcm.rpc(from, rpc, mcm.now)
}

func (mcm *managedConsensusModule) tickTilElectionTimeout() {
	electionTimeoutTime := mcm.pcm.electionTimeoutTracker.electionTimeoutTime
	for {
		mcm.tick()
		if mcm.now.After(electionTimeoutTime) {
			break
		}
	}
	if mcm.pcm.electionTimeoutTracker.electionTimeoutTime != electionTimeoutTime {
		panic("electionTimeoutTime changed!")
	}
	// Because tick() increments "now" after calling tick(),
	// we need one more to actually run with a post-timeout "now".
	mcm.tick()
}
<|endoftext|>


<commit_before>package consumer

import (
	"bufio"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/cloudfoundry/noaa"
	noaa_errors "github.com/cloudfoundry/noaa/errors"
	"github.com/cloudfoundry/sonde-go/events"
	"github.com/gogo/protobuf/proto"
	"github.com/gorilla/websocket"
)

// TailingLogs listens indefinitely for log messages only; other event types
// are dropped.
// Whenever an error is encountered, the error will be sent down the error
// channel and TailingLogs will attempt to reconnect up to 5 times.  After
// five failed reconnection attempts, TailingLogs will give up and close the
// error and LogMessage channels.
//
// If c is closed, the returned channels will both be closed.
//
// Errors must be drained from the returned error channel for it to continue
// retrying; if they are not drained, the connection attempts will hang.
func (c *Consumer) TailingLogs(appGuid, authToken string) (<-chan *events.LogMessage, <-chan error) {
	return c.tailingLogs(appGuid, authToken, maxRetries)
}

// TailingLogsWithoutReconnect functions identically to TailingLogs but without
// any reconnect attempts when errors occur.
func (c *Consumer) TailingLogsWithoutReconnect(appGuid string, authToken string) (<-chan *events.LogMessage, <-chan error) {
	return c.tailingLogs(appGuid, authToken, 0)
}

// Stream listens indefinitely for all log and event messages.
//
// Messages are presented in the order received from the loggregator server.
// Chronological or other ordering is not guaranteed. It is the responsibility
// of the consumer of these channels to provide any desired sorting mechanism.
//
// Whenever an error is encountered, the error will be sent down the error
// channel and Stream will attempt to reconnect up to 5 times.  After five
// failed reconnection attempts, Stream will give up and close the error and
// Envelope channels.
func (c *Consumer) Stream(appGuid string, authToken string) (outputChan <-chan *events.Envelope, errorChan <-chan error) {
	return c.runStream(appGuid, authToken, maxRetries)
}

// StreamWithoutReconnect functions identically to Stream but without any
// reconnect attempts when errors occur.
func (c *Consumer) StreamWithoutReconnect(appGuid string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.runStream(appGuid, authToken, 0)
}

// Firehose streams all data. All clients with the same subscriptionId will
// receive a proportionate share of the message stream.  Each pool of clients
// will receive the entire stream.
//
// Messages are presented in the order received from the loggregator server.
// Chronological or other ordering is not guaranteed. It is the responsibility
// of the consumer of these channels to provide any desired sorting mechanism.
//
// Whenever an error is encountered, the error will be sent down the error
// channel and Firehose will attempt to reconnect up to 5 times.  After five
// failed reconnection attempts, Firehose will give up and close the error and
// Envelope channels.
func (c *Consumer) Firehose(subscriptionId string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.firehose(subscriptionId, authToken, 5)
}

// FirehoseWithoutReconnect functions identically to Firehose but without any
// reconnect attempts when errors occur.
func (c *Consumer) FirehoseWithoutReconnect(subscriptionId string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.firehose(subscriptionId, authToken, 0)
}

// SetDebugPrinter sets the websocket connection to write debug information to
// debugPrinter.
func (c *Consumer) SetDebugPrinter(debugPrinter noaa.DebugPrinter) {
	c.debugPrinter = debugPrinter
}

// SetOnConnectCallback sets a callback function to be called with the
// websocket connection is established.
func (c *Consumer) SetOnConnectCallback(cb func()) {
	c.callbackLock.Lock()
	defer c.callbackLock.Unlock()
	c.callback = cb
}

// Close terminates all previously opened websocket connections to the traffic
// controller.  It will return an error if there are no open connections, or
// if it has problems closing any connection.
//
// The consumer will mark itself closed once it is done, and will reopen the next
// time a websocket connection is made.
func (c *Consumer) Close() error {
	c.connsLock.Lock()
	defer c.connsLock.Unlock()
	if len(c.conns) == 0 {
		return errors.New("connection does not exist")
	}
	for len(c.conns) > 0 {
		if err := c.conns[0].close(); err != nil {
			return err
		}
		c.conns = c.conns[1:]
	}
	return nil
}

// Closed is a legacy method that now returns false.  Consumers cannot be closed;
// they will always be able to make more requests to the TrafficController.
func (c *Consumer) Closed() bool {
	return false
}

func (c *Consumer) SetIdleTimeout(idleTimeout time.Duration) {
	c.idleTimeout = idleTimeout
}

func (c *Consumer) onConnectCallback() func() {
	c.callbackLock.RLock()
	defer c.callbackLock.RUnlock()
	return c.callback
}

func (c *Consumer) tailingLogs(appGuid, authToken string, retries uint) (<-chan *events.LogMessage, <-chan error) {
	outputs := make(chan *events.LogMessage)
	errors := make(chan error, 1)
	callback := func(env *events.Envelope) {
		if env.GetEventType() == events.Envelope_LogMessage {
			outputs <- env.GetLogMessage()
		}
	}

	conn := c.newConn()
	go func() {
		defer close(errors)
		defer close(outputs)
		c.streamAppDataTo(conn, appGuid, authToken, callback, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) runStream(appGuid, authToken string, retries uint) (<-chan *events.Envelope, <-chan error) {
	outputs := make(chan *events.Envelope)
	errors := make(chan error, 1)

	callback := func(env *events.Envelope) {
		outputs <- env
	}

	conn := c.newConn()
	go func() {
		defer close(errors)
		defer close(outputs)
		c.streamAppDataTo(conn, appGuid, authToken, callback, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) streamAppDataTo(conn *connection, appGuid, authToken string, callback func(*events.Envelope), errors chan<- error, retries uint) {
	streamPath := fmt.Sprintf("/apps/%s/stream", appGuid)
	action := func() (error, bool) {
		if conn.Closed() {
			return nil, true
		}
		ws, err := c.establishWebsocketConnection(streamPath, authToken)
		if err != nil {
			return err, false
		}
		conn.setWebsocket(ws)
		return c.listenForMessages(conn, callback), false
	}
	c.retryAction(action, errors, retries)
}

func (c *Consumer) firehose(subID, authToken string, retries uint) (<-chan *events.Envelope, <-chan error) {
	outputs := make(chan *events.Envelope)
	errors := make(chan error, 1)
	callback := func(env *events.Envelope) {
		outputs <- env
	}

	streamPath := "/firehose/" + subID
	conn := c.newConn()
	action := func() (error, bool) {
		if conn.Closed() {
			return nil, true
		}
		ws, err := c.establishWebsocketConnection(streamPath, authToken)
		if err != nil {
			return err, false
		}
		conn.setWebsocket(ws)
		return c.listenForMessages(conn, callback), false
	}
	go func() {
		defer close(errors)
		defer close(outputs)
		c.retryAction(action, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) stream(streamPath string, authToken string, callback func(*events.Envelope)) error {
	ws, err := c.establishWebsocketConnection(streamPath, authToken)
	if err != nil {
		return err
	}
	conn := c.newConn()
	conn.setWebsocket(ws)

	return c.listenForMessages(conn, callback)
}

func (c *Consumer) listenForMessages(conn *connection, callback func(*events.Envelope)) error {
	if conn.Closed() {
		return nil
	}
	ws := conn.websocket()
	for {
		if c.idleTimeout != 0 {
			ws.SetReadDeadline(time.Now().Add(c.idleTimeout))
		}
		_, data, err := ws.ReadMessage()

		// If the connection was closed (i.e. if conn.Close() was called), we
		// will have a non-nil error, but we want to return a nil error.
		if conn.Closed() {
			return nil
		}

		if err != nil {
			return err
		}

		envelope := &events.Envelope{}
		err = proto.Unmarshal(data, envelope)
		if err != nil {
			continue
		}

		callback(envelope)
	}
}

func (c *Consumer) retryAction(action func() (err error, done bool), errors chan<- error, retries uint) {
	reconnectAttempts := uint(0)

	oldConnectCallback := c.onConnectCallback()
	defer c.SetOnConnectCallback(oldConnectCallback)

	c.SetOnConnectCallback(func() {
		reconnectAttempts = 0
		if oldConnectCallback != nil {
			oldConnectCallback()
		}
	})

	for ; reconnectAttempts <= retries; reconnectAttempts++ {
		err, done := action()
		if done {
			return
		}
		errors <- err
		time.Sleep(reconnectTimeout)
	}
}

func (c *Consumer) newConn() *connection {
	conn := &connection{}
	c.connsLock.Lock()
	defer c.connsLock.Unlock()
	c.conns = append(c.conns, conn)
	return conn
}

func (c *Consumer) establishWebsocketConnection(path string, authToken string) (*websocket.Conn, error) {
	header := http.Header{"Origin": []string{"http://localhost"}, "Authorization": []string{authToken}}
	url := c.trafficControllerUrl + path

	c.debugPrinter.Print("WEBSOCKET REQUEST:",
		"GET "+path+" HTTP/1.1\n"+
			"Host: "+c.trafficControllerUrl+"\n"+
			"Upgrade: websocket\nConnection: Upgrade\nSec-WebSocket-Version: 13\nSec-WebSocket-Key: [HIDDEN]\n"+
			headersString(header))

	ws, resp, err := c.dialer.Dial(url, header)
	if resp != nil {
		c.debugPrinter.Print("WEBSOCKET RESPONSE:",
			resp.Proto+" "+resp.Status+"\n"+
				headersString(resp.Header))
	}

	if resp != nil && resp.StatusCode == http.StatusUnauthorized {
		bodyData, _ := ioutil.ReadAll(resp.Body)
		err = noaa_errors.NewUnauthorizedError(string(bodyData))
		return ws, err
	}

	callback := c.onConnectCallback()
	if err == nil && callback != nil {
		callback()
	}

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Error dialing traffic controller server: %s.\nPlease ask your Cloud Foundry Operator to check the platform configuration (traffic controller is %s).", err.Error(), c.trafficControllerUrl))
	}

	return ws, err
}

func (c *Consumer) proxyDial(network, addr string) (net.Conn, error) {
	targetUrl, err := url.Parse("http://" + addr)
	if err != nil {
		return nil, err
	}

	proxy := c.proxy
	if proxy == nil {
		proxy = http.ProxyFromEnvironment
	}

	proxyUrl, err := proxy(&http.Request{URL: targetUrl})
	if err != nil {
		return nil, err
	}
	if proxyUrl == nil {
		return net.Dial(network, addr)
	}

	proxyConn, err := net.Dial(network, proxyUrl.Host)
	if err != nil {
		return nil, err
	}

	connectReq := &http.Request{
		Method: "CONNECT",
		URL:    targetUrl,
		Host:   targetUrl.Host,
		Header: make(http.Header),
	}
	connectReq.Write(proxyConn)

	connectResp, err := http.ReadResponse(bufio.NewReader(proxyConn), connectReq)
	if err != nil {
		proxyConn.Close()
		return nil, err
	}
	if connectResp.StatusCode != http.StatusOK {
		f := strings.SplitN(connectResp.Status, " ", 2)
		proxyConn.Close()
		return nil, errors.New(f[1])
	}

	return proxyConn, nil
}

func headersString(header http.Header) string {
	var result string
	for name, values := range header {
		result += name + ": " + strings.Join(values, ", ") + "\n"
	}
	return result
}

type connection struct {
	ws     *websocket.Conn
	closed bool
	lock   sync.Mutex
}

func (c *connection) websocket() *websocket.Conn {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.ws
}

func (c *connection) setWebsocket(ws *websocket.Conn) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.closed {
		return
	}
	c.ws = ws
}

func (c *connection) close() error {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.closed = true
	if c.ws == nil {
		return nil
	}
	err := c.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
	if err != nil {
		return err
	}
	return c.ws.Close()
}

func (c *connection) Closed() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.closed
}
<commit_msg>Remove dead method<commit_after>package consumer

import (
	"bufio"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/cloudfoundry/noaa"
	noaa_errors "github.com/cloudfoundry/noaa/errors"
	"github.com/cloudfoundry/sonde-go/events"
	"github.com/gogo/protobuf/proto"
	"github.com/gorilla/websocket"
)

// TailingLogs listens indefinitely for log messages only; other event types
// are dropped.
// Whenever an error is encountered, the error will be sent down the error
// channel and TailingLogs will attempt to reconnect up to 5 times.  After
// five failed reconnection attempts, TailingLogs will give up and close the
// error and LogMessage channels.
//
// If c is closed, the returned channels will both be closed.
//
// Errors must be drained from the returned error channel for it to continue
// retrying; if they are not drained, the connection attempts will hang.
func (c *Consumer) TailingLogs(appGuid, authToken string) (<-chan *events.LogMessage, <-chan error) {
	return c.tailingLogs(appGuid, authToken, maxRetries)
}

// TailingLogsWithoutReconnect functions identically to TailingLogs but without
// any reconnect attempts when errors occur.
func (c *Consumer) TailingLogsWithoutReconnect(appGuid string, authToken string) (<-chan *events.LogMessage, <-chan error) {
	return c.tailingLogs(appGuid, authToken, 0)
}

// Stream listens indefinitely for all log and event messages.
//
// Messages are presented in the order received from the loggregator server.
// Chronological or other ordering is not guaranteed. It is the responsibility
// of the consumer of these channels to provide any desired sorting mechanism.
//
// Whenever an error is encountered, the error will be sent down the error
// channel and Stream will attempt to reconnect up to 5 times.  After five
// failed reconnection attempts, Stream will give up and close the error and
// Envelope channels.
func (c *Consumer) Stream(appGuid string, authToken string) (outputChan <-chan *events.Envelope, errorChan <-chan error) {
	return c.runStream(appGuid, authToken, maxRetries)
}

// StreamWithoutReconnect functions identically to Stream but without any
// reconnect attempts when errors occur.
func (c *Consumer) StreamWithoutReconnect(appGuid string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.runStream(appGuid, authToken, 0)
}

// Firehose streams all data. All clients with the same subscriptionId will
// receive a proportionate share of the message stream.  Each pool of clients
// will receive the entire stream.
//
// Messages are presented in the order received from the loggregator server.
// Chronological or other ordering is not guaranteed. It is the responsibility
// of the consumer of these channels to provide any desired sorting mechanism.
//
// Whenever an error is encountered, the error will be sent down the error
// channel and Firehose will attempt to reconnect up to 5 times.  After five
// failed reconnection attempts, Firehose will give up and close the error and
// Envelope channels.
func (c *Consumer) Firehose(subscriptionId string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.firehose(subscriptionId, authToken, 5)
}

// FirehoseWithoutReconnect functions identically to Firehose but without any
// reconnect attempts when errors occur.
func (c *Consumer) FirehoseWithoutReconnect(subscriptionId string, authToken string) (<-chan *events.Envelope, <-chan error) {
	return c.firehose(subscriptionId, authToken, 0)
}

// SetDebugPrinter sets the websocket connection to write debug information to
// debugPrinter.
func (c *Consumer) SetDebugPrinter(debugPrinter noaa.DebugPrinter) {
	c.debugPrinter = debugPrinter
}

// SetOnConnectCallback sets a callback function to be called with the
// websocket connection is established.
func (c *Consumer) SetOnConnectCallback(cb func()) {
	c.callbackLock.Lock()
	defer c.callbackLock.Unlock()
	c.callback = cb
}

// Close terminates all previously opened websocket connections to the traffic
// controller.  It will return an error if there are no open connections, or
// if it has problems closing any connection.
//
// The consumer will mark itself closed once it is done, and will reopen the next
// time a websocket connection is made.
func (c *Consumer) Close() error {
	c.connsLock.Lock()
	defer c.connsLock.Unlock()
	if len(c.conns) == 0 {
		return errors.New("connection does not exist")
	}
	for len(c.conns) > 0 {
		if err := c.conns[0].close(); err != nil {
			return err
		}
		c.conns = c.conns[1:]
	}
	return nil
}

// Closed is a legacy method that now returns false.  Consumers cannot be closed;
// they will always be able to make more requests to the TrafficController.
func (c *Consumer) Closed() bool {
	return false
}

func (c *Consumer) SetIdleTimeout(idleTimeout time.Duration) {
	c.idleTimeout = idleTimeout
}

func (c *Consumer) onConnectCallback() func() {
	c.callbackLock.RLock()
	defer c.callbackLock.RUnlock()
	return c.callback
}

func (c *Consumer) tailingLogs(appGuid, authToken string, retries uint) (<-chan *events.LogMessage, <-chan error) {
	outputs := make(chan *events.LogMessage)
	errors := make(chan error, 1)
	callback := func(env *events.Envelope) {
		if env.GetEventType() == events.Envelope_LogMessage {
			outputs <- env.GetLogMessage()
		}
	}

	conn := c.newConn()
	go func() {
		defer close(errors)
		defer close(outputs)
		c.streamAppDataTo(conn, appGuid, authToken, callback, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) runStream(appGuid, authToken string, retries uint) (<-chan *events.Envelope, <-chan error) {
	outputs := make(chan *events.Envelope)
	errors := make(chan error, 1)

	callback := func(env *events.Envelope) {
		outputs <- env
	}

	conn := c.newConn()
	go func() {
		defer close(errors)
		defer close(outputs)
		c.streamAppDataTo(conn, appGuid, authToken, callback, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) streamAppDataTo(conn *connection, appGuid, authToken string, callback func(*events.Envelope), errors chan<- error, retries uint) {
	streamPath := fmt.Sprintf("/apps/%s/stream", appGuid)
	action := func() (error, bool) {
		if conn.Closed() {
			return nil, true
		}
		ws, err := c.establishWebsocketConnection(streamPath, authToken)
		if err != nil {
			return err, false
		}
		conn.setWebsocket(ws)
		return c.listenForMessages(conn, callback), false
	}
	c.retryAction(action, errors, retries)
}

func (c *Consumer) firehose(subID, authToken string, retries uint) (<-chan *events.Envelope, <-chan error) {
	outputs := make(chan *events.Envelope)
	errors := make(chan error, 1)
	callback := func(env *events.Envelope) {
		outputs <- env
	}

	streamPath := "/firehose/" + subID
	conn := c.newConn()
	action := func() (error, bool) {
		if conn.Closed() {
			return nil, true
		}
		ws, err := c.establishWebsocketConnection(streamPath, authToken)
		if err != nil {
			return err, false
		}
		conn.setWebsocket(ws)
		return c.listenForMessages(conn, callback), false
	}
	go func() {
		defer close(errors)
		defer close(outputs)
		c.retryAction(action, errors, retries)
	}()
	return outputs, errors
}

func (c *Consumer) listenForMessages(conn *connection, callback func(*events.Envelope)) error {
	if conn.Closed() {
		return nil
	}
	ws := conn.websocket()
	for {
		if c.idleTimeout != 0 {
			ws.SetReadDeadline(time.Now().Add(c.idleTimeout))
		}
		_, data, err := ws.ReadMessage()

		// If the connection was closed (i.e. if conn.Close() was called), we
		// will have a non-nil error, but we want to return a nil error.
		if conn.Closed() {
			return nil
		}

		if err != nil {
			return err
		}

		envelope := &events.Envelope{}
		err = proto.Unmarshal(data, envelope)
		if err != nil {
			continue
		}

		callback(envelope)
	}
}

func (c *Consumer) retryAction(action func() (err error, done bool), errors chan<- error, retries uint) {
	reconnectAttempts := uint(0)

	oldConnectCallback := c.onConnectCallback()
	defer c.SetOnConnectCallback(oldConnectCallback)

	c.SetOnConnectCallback(func() {
		reconnectAttempts = 0
		if oldConnectCallback != nil {
			oldConnectCallback()
		}
	})

	for ; reconnectAttempts <= retries; reconnectAttempts++ {
		err, done := action()
		if done {
			return
		}
		errors <- err
		time.Sleep(reconnectTimeout)
	}
}

func (c *Consumer) newConn() *connection {
	conn := &connection{}
	c.connsLock.Lock()
	defer c.connsLock.Unlock()
	c.conns = append(c.conns, conn)
	return conn
}

func (c *Consumer) establishWebsocketConnection(path string, authToken string) (*websocket.Conn, error) {
	header := http.Header{"Origin": []string{"http://localhost"}, "Authorization": []string{authToken}}
	url := c.trafficControllerUrl + path

	c.debugPrinter.Print("WEBSOCKET REQUEST:",
		"GET "+path+" HTTP/1.1\n"+
			"Host: "+c.trafficControllerUrl+"\n"+
			"Upgrade: websocket\nConnection: Upgrade\nSec-WebSocket-Version: 13\nSec-WebSocket-Key: [HIDDEN]\n"+
			headersString(header))

	ws, resp, err := c.dialer.Dial(url, header)
	if resp != nil {
		c.debugPrinter.Print("WEBSOCKET RESPONSE:",
			resp.Proto+" "+resp.Status+"\n"+
				headersString(resp.Header))
	}

	if resp != nil && resp.StatusCode == http.StatusUnauthorized {
		bodyData, _ := ioutil.ReadAll(resp.Body)
		err = noaa_errors.NewUnauthorizedError(string(bodyData))
		return ws, err
	}

	callback := c.onConnectCallback()
	if err == nil && callback != nil {
		callback()
	}

	if err != nil {
		return nil, errors.New(fmt.Sprintf("Error dialing traffic controller server: %s.\nPlease ask your Cloud Foundry Operator to check the platform configuration (traffic controller is %s).", err.Error(), c.trafficControllerUrl))
	}

	return ws, err
}

func (c *Consumer) proxyDial(network, addr string) (net.Conn, error) {
	targetUrl, err := url.Parse("http://" + addr)
	if err != nil {
		return nil, err
	}

	proxy := c.proxy
	if proxy == nil {
		proxy = http.ProxyFromEnvironment
	}

	proxyUrl, err := proxy(&http.Request{URL: targetUrl})
	if err != nil {
		return nil, err
	}
	if proxyUrl == nil {
		return net.Dial(network, addr)
	}

	proxyConn, err := net.Dial(network, proxyUrl.Host)
	if err != nil {
		return nil, err
	}

	connectReq := &http.Request{
		Method: "CONNECT",
		URL:    targetUrl,
		Host:   targetUrl.Host,
		Header: make(http.Header),
	}
	connectReq.Write(proxyConn)

	connectResp, err := http.ReadResponse(bufio.NewReader(proxyConn), connectReq)
	if err != nil {
		proxyConn.Close()
		return nil, err
	}
	if connectResp.StatusCode != http.StatusOK {
		f := strings.SplitN(connectResp.Status, " ", 2)
		proxyConn.Close()
		return nil, errors.New(f[1])
	}

	return proxyConn, nil
}

func headersString(header http.Header) string {
	var result string
	for name, values := range header {
		result += name + ": " + strings.Join(values, ", ") + "\n"
	}
	return result
}

type connection struct {
	ws     *websocket.Conn
	closed bool
	lock   sync.Mutex
}

func (c *connection) websocket() *websocket.Conn {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.ws
}

func (c *connection) setWebsocket(ws *websocket.Conn) {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.closed {
		return
	}
	c.ws = ws
}

func (c *connection) close() error {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.closed = true
	if c.ws == nil {
		return nil
	}
	err := c.ws.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Time{})
	if err != nil {
		return err
	}
	return c.ws.Close()
}

func (c *connection) Closed() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.closed
}
<|endoftext|>


<commit_before>// +build ccm

package ccm

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"strings"
)

func execCmd(args ...string) (*bytes.Buffer, error) {
	cmd := exec.Command("ccm", args...)
	stdout := &bytes.Buffer{}
	cmd.Stdout = stdout
	cmd.Stderr = &bytes.Buffer{}
	if err := cmd.Run(); err != nil {
		return nil, errors.New(cmd.Stderr.(*bytes.Buffer).String())
	}

	return stdout, nil
}

func AllUp() error {
	status, err := Status()
	if err != nil {
		return err
	}

	for _, host := range status {
		if !host.State.IsUp() {
			if err := NodeUp(host.Name); err != nil {
				return err
			}
		}
	}

	return nil
}

func NodeUp(node string) error {
	_, err := execCmd(node, "start", "--wait-for-binary-proto", "--wait-other-notice")
	return err
}

func NodeDown(node string) error {
	_, err := execCmd(node, "stop")
	return err
}

type Host struct {
	State NodeState
	Addr  string
	Name  string
}

type NodeState int

func (n NodeState) String() string {
	if n == NodeStateUp {
		return "UP"
	} else if n == NodeStateDown {
		return "DOWN"
	} else {
		return fmt.Sprintf("UNKNOWN_STATE_%d", n)
	}
}

func (n NodeState) IsUp() bool {
	return n == NodeStateUp
}

const (
	NodeStateUp NodeState = iota
	NodeStateDown
)

func Status() (map[string]Host, error) {
	// TODO: parse into struct o maniuplate
	out, err := execCmd("status", "-v")
	if err != nil {
		return nil, err
	}

	const (
		stateCluster = iota
		stateCommas
		stateNode
		stateOption
	)

	nodes := make(map[string]Host)
	// didnt really want to write a full state machine parser
	state := stateCluster
	sc := bufio.NewScanner(out)

	var host Host

	for sc.Scan() {
		switch state {
		case stateCluster:
			text := sc.Text()
			if !strings.HasPrefix(text, "Cluster:") {
				return nil, fmt.Errorf("expected 'Cluster:' got %q", text)
			}
			state = stateCommas
		case stateCommas:
			text := sc.Text()
			if strings.HasPrefix(text, "-") {
				return nil, fmt.Errorf("expected commas got %q", text)
			}
			state = stateNode
		case stateNode:
			// assume nodes start with node
			text := sc.Text()
			if !strings.HasPrefix(text, "node") {
				return nil, fmt.Errorf("expected 'node' got %q", text)
			}
			line := strings.Split(text, ":")
			host.Name = line[0]

			nodeState := strings.TrimSpace(line[1])
			switch nodeState {
			case "UP":
				host.State = NodeStateUp
			case "DOWN":
				host.State = NodeStateDown
			default:
				return nil, fmt.Errorf("unknown node state from ccm: %q", nodeState)
			}

			state = stateOption
		case stateOption:
			text := sc.Text()
			if text == "" {
				state = stateNode
				nodes[host.Name] = host
				host = Host{}
				continue
			}

			line := strings.Split(strings.TrimSpace(text), "=")
			k, v := line[0], line[1]
			if k == "binary" {
				// could check errors
				// ('127.0.0.1', 9042)
				v = v[2:] // (''
				if i := strings.IndexByte(v, '\''); i < 0 {
					return nil, fmt.Errorf("invalid binary v=%q", v)
				} else {
					host.Addr = v[:i]
					// dont need port
				}
			}
		default:
			return nil, fmt.Errorf("unexpected state: %q", state)
		}
	}

	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("unable to parse ccm status: %v", err)
	}

	return nodes, nil
}
<commit_msg>fix ccm parsing commas<commit_after>// +build ccm

package ccm

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"strings"
)

func execCmd(args ...string) (*bytes.Buffer, error) {
	cmd := exec.Command("ccm", args...)
	stdout := &bytes.Buffer{}
	cmd.Stdout = stdout
	cmd.Stderr = &bytes.Buffer{}
	if err := cmd.Run(); err != nil {
		return nil, errors.New(cmd.Stderr.(*bytes.Buffer).String())
	}

	return stdout, nil
}

func AllUp() error {
	status, err := Status()
	if err != nil {
		return err
	}

	for _, host := range status {
		if !host.State.IsUp() {
			if err := NodeUp(host.Name); err != nil {
				return err
			}
		}
	}

	return nil
}

func NodeUp(node string) error {
	_, err := execCmd(node, "start", "--wait-for-binary-proto", "--wait-other-notice")
	return err
}

func NodeDown(node string) error {
	_, err := execCmd(node, "stop")
	return err
}

type Host struct {
	State NodeState
	Addr  string
	Name  string
}

type NodeState int

func (n NodeState) String() string {
	if n == NodeStateUp {
		return "UP"
	} else if n == NodeStateDown {
		return "DOWN"
	} else {
		return fmt.Sprintf("UNKNOWN_STATE_%d", n)
	}
}

func (n NodeState) IsUp() bool {
	return n == NodeStateUp
}

const (
	NodeStateUp NodeState = iota
	NodeStateDown
)

func Status() (map[string]Host, error) {
	// TODO: parse into struct o maniuplate
	out, err := execCmd("status", "-v")
	if err != nil {
		return nil, err
	}

	const (
		stateCluster = iota
		stateCommas
		stateNode
		stateOption
	)

	nodes := make(map[string]Host)
	// didnt really want to write a full state machine parser
	state := stateCluster
	sc := bufio.NewScanner(out)

	var host Host

	for sc.Scan() {
		switch state {
		case stateCluster:
			text := sc.Text()
			if !strings.HasPrefix(text, "Cluster:") {
				return nil, fmt.Errorf("expected 'Cluster:' got %q", text)
			}
			state = stateCommas
		case stateCommas:
			text := sc.Text()
			if !strings.HasPrefix(text, "-") {
				return nil, fmt.Errorf("expected commas got %q", text)
			}
			state = stateNode
		case stateNode:
			// assume nodes start with node
			text := sc.Text()
			if !strings.HasPrefix(text, "node") {
				return nil, fmt.Errorf("expected 'node' got %q", text)
			}
			line := strings.Split(text, ":")
			host.Name = line[0]

			nodeState := strings.TrimSpace(line[1])
			switch nodeState {
			case "UP":
				host.State = NodeStateUp
			case "DOWN":
				host.State = NodeStateDown
			default:
				return nil, fmt.Errorf("unknown node state from ccm: %q", nodeState)
			}

			state = stateOption
		case stateOption:
			text := sc.Text()
			if text == "" {
				state = stateNode
				nodes[host.Name] = host
				host = Host{}
				continue
			}

			line := strings.Split(strings.TrimSpace(text), "=")
			k, v := line[0], line[1]
			if k == "binary" {
				// could check errors
				// ('127.0.0.1', 9042)
				v = v[2:] // (''
				if i := strings.IndexByte(v, '\''); i < 0 {
					return nil, fmt.Errorf("invalid binary v=%q", v)
				} else {
					host.Addr = v[:i]
					// dont need port
				}
			}
		default:
			return nil, fmt.Errorf("unexpected state: %q", state)
		}
	}

	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("unable to parse ccm status: %v", err)
	}

	return nodes, nil
}
<|endoftext|>


<commit_before>// Package vmess contains protocol definition, io lib for VMess.
package protocol

import (
	"crypto/md5"
	"encoding/binary"
	"hash/fnv"
	"io"

	"github.com/v2ray/v2ray-core/common/alloc"
	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/vmess"
	"github.com/v2ray/v2ray-core/transport"
)

const (
	addrTypeIPv4   = byte(0x01)
	addrTypeIPv6   = byte(0x03)
	addrTypeDomain = byte(0x02)

	CmdTCP = byte(0x01)
	CmdUDP = byte(0x02)

	Version = byte(0x01)

	blockSize = 16
)

// VMessRequest implements the request message of VMess protocol. It only contains the header of a
// request message. The data part will be handled by conection handler directly, in favor of data
// streaming.
type VMessRequest struct {
	Version        byte
	User           vmess.User
	RequestIV      []byte
	RequestKey     []byte
	ResponseHeader []byte
	Command        byte
	Address        v2net.Address
	Port           v2net.Port
}

// Destination is the final destination of this request.
func (this *VMessRequest) Destination() v2net.Destination {
	if this.Command == CmdTCP {
		return v2net.TCPDestination(this.Address, this.Port)
	} else {
		return v2net.UDPDestination(this.Address, this.Port)
	}
}

// VMessRequestReader is a parser to read VMessRequest from a byte stream.
type VMessRequestReader struct {
	vUserSet UserSet
}

// NewVMessRequestReader creates a new VMessRequestReader with a given UserSet
func NewVMessRequestReader(vUserSet UserSet) *VMessRequestReader {
	return &VMessRequestReader{
		vUserSet: vUserSet,
	}
}

// Read reads a VMessRequest from a byte stream.
func (this *VMessRequestReader) Read(reader io.Reader) (*VMessRequest, error) {
	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	nBytes, err := v2net.ReadAllBytes(reader, buffer.Value[:vmess.IDBytesLen])
	if err != nil {
		log.Debug("VMess: Failed to read request ID (%d bytes): %v", nBytes, err)
		return nil, err
	}

	userObj, timeSec, valid := this.vUserSet.GetUser(buffer.Value[:nBytes])
	if !valid {
		return nil, proxy.InvalidAuthentication
	}

	timestampHash := TimestampHash()
	timestampHash.Write(timeSec.HashBytes())
	iv := timestampHash.Sum(nil)
	aesStream, err := v2crypto.NewAesDecryptionStream(userObj.ID().CmdKey(), iv)
	if err != nil {
		log.Debug("VMess: Failed to create AES stream: %v", err)
		return nil, err
	}

	decryptor := v2crypto.NewCryptionReader(aesStream, reader)

	nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[:41])
	if err != nil {
		log.Debug("VMess: Failed to read request header (%d bytes): %v", nBytes, err)
		return nil, err
	}
	bufferLen := nBytes

	request := &VMessRequest{
		User:    userObj,
		Version: buffer.Value[0],
	}

	if request.Version != Version {
		log.Warning("VMess: Invalid protocol version %d", request.Version)
		return nil, proxy.InvalidProtocolVersion
	}

	request.RequestIV = append([]byte(nil), buffer.Value[1:17]...)       // 16 bytes
	request.RequestKey = append([]byte(nil), buffer.Value[17:33]...)     // 16 bytes
	request.ResponseHeader = append([]byte(nil), buffer.Value[33:37]...) // 4 bytes
	request.Command = buffer.Value[37]

	request.Port = v2net.PortFromBytes(buffer.Value[38:40])

	switch buffer.Value[40] {
	case addrTypeIPv4:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:45]) // 4 bytes
		bufferLen += 4
		if err != nil {
			log.Debug("VMess: Failed to read target IPv4 (%d bytes): %v", nBytes, err)
			return nil, err
		}
		request.Address = v2net.IPAddress(buffer.Value[41:45])
	case addrTypeIPv6:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:57]) // 16 bytes
		bufferLen += 16
		if err != nil {
			log.Debug("VMess: Failed to read target IPv6 (%d bytes): %v", nBytes, err)
			return nil, err
		}
		request.Address = v2net.IPAddress(buffer.Value[41:57])
	case addrTypeDomain:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:42])
		if err != nil {
			log.Debug("VMess: Failed to read target domain (%d bytes): %v", nBytes, err)
			return nil, err
		}
		domainLength := int(buffer.Value[41])
		if domainLength == 0 {
			return nil, transport.CorruptedPacket
		}
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[42:42+domainLength])
		if err != nil {
			log.Debug("VMess: Failed to read target domain (%d bytes): %v", nBytes, err)
			return nil, err
		}
		bufferLen += 1 + domainLength
		domainBytes := append([]byte(nil), buffer.Value[42:42+domainLength]...)
		request.Address = v2net.DomainAddress(string(domainBytes))
	}

	nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[bufferLen:bufferLen+4])
	if err != nil {
		log.Debug("VMess: Failed to read checksum (%d bytes): %v", nBytes, err)
		return nil, err
	}

	fnv1a := fnv.New32a()
	fnv1a.Write(buffer.Value[:bufferLen])
	actualHash := fnv1a.Sum32()
	expectedHash := binary.BigEndian.Uint32(buffer.Value[bufferLen : bufferLen+4])

	if actualHash != expectedHash {
		return nil, transport.CorruptedPacket
	}

	return request, nil
}

// ToBytes returns a VMessRequest in the form of byte array.
func (this *VMessRequest) ToBytes(timestampGenerator RandomTimestampGenerator, buffer *alloc.Buffer) (*alloc.Buffer, error) {
	if buffer == nil {
		buffer = alloc.NewSmallBuffer().Clear()
	}

	timestamp := timestampGenerator.Next()
	idHash := IDHash(this.User.AnyValidID().Bytes())
	idHash.Write(timestamp.Bytes())

	buffer.Append(idHash.Sum(nil))

	encryptionBegin := buffer.Len()

	buffer.AppendBytes(this.Version)
	buffer.Append(this.RequestIV)
	buffer.Append(this.RequestKey)
	buffer.Append(this.ResponseHeader)
	buffer.AppendBytes(this.Command)
	buffer.Append(this.Port.Bytes())

	switch {
	case this.Address.IsIPv4():
		buffer.AppendBytes(addrTypeIPv4)
		buffer.Append(this.Address.IP())
	case this.Address.IsIPv6():
		buffer.AppendBytes(addrTypeIPv6)
		buffer.Append(this.Address.IP())
	case this.Address.IsDomain():
		buffer.AppendBytes(addrTypeDomain, byte(len(this.Address.Domain())))
		buffer.Append([]byte(this.Address.Domain()))
	}

	encryptionEnd := buffer.Len()

	fnv1a := fnv.New32a()
	fnv1a.Write(buffer.Value[encryptionBegin:encryptionEnd])

	fnvHash := fnv1a.Sum32()
	buffer.AppendBytes(byte(fnvHash>>24), byte(fnvHash>>16), byte(fnvHash>>8), byte(fnvHash))
	encryptionEnd += 4

	timestampHash := md5.New()
	timestampHash.Write(timestamp.HashBytes())
	iv := timestampHash.Sum(nil)
	aesStream, err := v2crypto.NewAesEncryptionStream(this.User.ID().CmdKey(), iv)
	if err != nil {
		return nil, err
	}
	aesStream.XORKeyStream(buffer.Value[encryptionBegin:encryptionEnd], buffer.Value[encryptionBegin:encryptionEnd])

	return buffer, nil
}
<commit_msg>remove unnecessary memory allocation in vmess<commit_after>// Package vmess contains protocol definition, io lib for VMess.
package protocol

import (
	"crypto/md5"
	"encoding/binary"
	"hash/fnv"
	"io"

	"github.com/v2ray/v2ray-core/common/alloc"
	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/vmess"
	"github.com/v2ray/v2ray-core/transport"
)

const (
	addrTypeIPv4   = byte(0x01)
	addrTypeIPv6   = byte(0x03)
	addrTypeDomain = byte(0x02)

	CmdTCP = byte(0x01)
	CmdUDP = byte(0x02)

	Version = byte(0x01)

	blockSize = 16
)

// VMessRequest implements the request message of VMess protocol. It only contains the header of a
// request message. The data part will be handled by conection handler directly, in favor of data
// streaming.
type VMessRequest struct {
	Version        byte
	User           vmess.User
	RequestIV      []byte
	RequestKey     []byte
	ResponseHeader []byte
	Command        byte
	Address        v2net.Address
	Port           v2net.Port
}

// Destination is the final destination of this request.
func (this *VMessRequest) Destination() v2net.Destination {
	if this.Command == CmdTCP {
		return v2net.TCPDestination(this.Address, this.Port)
	} else {
		return v2net.UDPDestination(this.Address, this.Port)
	}
}

// VMessRequestReader is a parser to read VMessRequest from a byte stream.
type VMessRequestReader struct {
	vUserSet UserSet
}

// NewVMessRequestReader creates a new VMessRequestReader with a given UserSet
func NewVMessRequestReader(vUserSet UserSet) *VMessRequestReader {
	return &VMessRequestReader{
		vUserSet: vUserSet,
	}
}

// Read reads a VMessRequest from a byte stream.
func (this *VMessRequestReader) Read(reader io.Reader) (*VMessRequest, error) {
	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()

	nBytes, err := v2net.ReadAllBytes(reader, buffer.Value[:vmess.IDBytesLen])
	if err != nil {
		log.Debug("VMess: Failed to read request ID (%d bytes): %v", nBytes, err)
		return nil, err
	}

	userObj, timeSec, valid := this.vUserSet.GetUser(buffer.Value[:nBytes])
	if !valid {
		return nil, proxy.InvalidAuthentication
	}

	timestampHash := TimestampHash()
	timestampHash.Write(timeSec.HashBytes())
	iv := timestampHash.Sum(nil)
	aesStream, err := v2crypto.NewAesDecryptionStream(userObj.ID().CmdKey(), iv)
	if err != nil {
		log.Debug("VMess: Failed to create AES stream: %v", err)
		return nil, err
	}

	decryptor := v2crypto.NewCryptionReader(aesStream, reader)

	nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[:41])
	if err != nil {
		log.Debug("VMess: Failed to read request header (%d bytes): %v", nBytes, err)
		return nil, err
	}
	bufferLen := nBytes

	request := &VMessRequest{
		User:    userObj,
		Version: buffer.Value[0],
	}

	if request.Version != Version {
		log.Warning("VMess: Invalid protocol version %d", request.Version)
		return nil, proxy.InvalidProtocolVersion
	}

	request.RequestIV = append([]byte(nil), buffer.Value[1:17]...)       // 16 bytes
	request.RequestKey = append([]byte(nil), buffer.Value[17:33]...)     // 16 bytes
	request.ResponseHeader = append([]byte(nil), buffer.Value[33:37]...) // 4 bytes
	request.Command = buffer.Value[37]

	request.Port = v2net.PortFromBytes(buffer.Value[38:40])

	switch buffer.Value[40] {
	case addrTypeIPv4:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:45]) // 4 bytes
		bufferLen += 4
		if err != nil {
			log.Debug("VMess: Failed to read target IPv4 (%d bytes): %v", nBytes, err)
			return nil, err
		}
		request.Address = v2net.IPAddress(buffer.Value[41:45])
	case addrTypeIPv6:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:57]) // 16 bytes
		bufferLen += 16
		if err != nil {
			log.Debug("VMess: Failed to read target IPv6 (%d bytes): %v", nBytes, err)
			return nil, err
		}
		request.Address = v2net.IPAddress(buffer.Value[41:57])
	case addrTypeDomain:
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[41:42])
		if err != nil {
			log.Debug("VMess: Failed to read target domain (%d bytes): %v", nBytes, err)
			return nil, err
		}
		domainLength := int(buffer.Value[41])
		if domainLength == 0 {
			return nil, transport.CorruptedPacket
		}
		nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[42:42+domainLength])
		if err != nil {
			log.Debug("VMess: Failed to read target domain (%d bytes): %v", nBytes, err)
			return nil, err
		}
		bufferLen += 1 + domainLength
		domainBytes := append([]byte(nil), buffer.Value[42:42+domainLength]...)
		request.Address = v2net.DomainAddress(string(domainBytes))
	}

	nBytes, err = v2net.ReadAllBytes(decryptor, buffer.Value[bufferLen:bufferLen+4])
	if err != nil {
		log.Debug("VMess: Failed to read checksum (%d bytes): %v", nBytes, err)
		return nil, err
	}

	fnv1a := fnv.New32a()
	fnv1a.Write(buffer.Value[:bufferLen])
	actualHash := fnv1a.Sum32()
	expectedHash := binary.BigEndian.Uint32(buffer.Value[bufferLen : bufferLen+4])

	if actualHash != expectedHash {
		return nil, transport.CorruptedPacket
	}

	return request, nil
}

// ToBytes returns a VMessRequest in the form of byte array.
func (this *VMessRequest) ToBytes(timestampGenerator RandomTimestampGenerator, buffer *alloc.Buffer) (*alloc.Buffer, error) {
	if buffer == nil {
		buffer = alloc.NewSmallBuffer().Clear()
	}

	timestamp := timestampGenerator.Next()
	idHash := IDHash(this.User.AnyValidID().Bytes())
	idHash.Write(timestamp.Bytes())

	hashStart := buffer.Len()
	buffer.Slice(0, hashStart+16)
	idHash.Sum(buffer.Value[hashStart:hashStart])

	encryptionBegin := buffer.Len()

	buffer.AppendBytes(this.Version)
	buffer.Append(this.RequestIV)
	buffer.Append(this.RequestKey)
	buffer.Append(this.ResponseHeader)
	buffer.AppendBytes(this.Command)
	buffer.Append(this.Port.Bytes())

	switch {
	case this.Address.IsIPv4():
		buffer.AppendBytes(addrTypeIPv4)
		buffer.Append(this.Address.IP())
	case this.Address.IsIPv6():
		buffer.AppendBytes(addrTypeIPv6)
		buffer.Append(this.Address.IP())
	case this.Address.IsDomain():
		buffer.AppendBytes(addrTypeDomain, byte(len(this.Address.Domain())))
		buffer.Append([]byte(this.Address.Domain()))
	}

	encryptionEnd := buffer.Len()

	fnv1a := fnv.New32a()
	fnv1a.Write(buffer.Value[encryptionBegin:encryptionEnd])

	fnvHash := fnv1a.Sum32()
	buffer.AppendBytes(byte(fnvHash>>24), byte(fnvHash>>16), byte(fnvHash>>8), byte(fnvHash))
	encryptionEnd += 4

	timestampHash := md5.New()
	timestampHash.Write(timestamp.HashBytes())
	iv := timestampHash.Sum(nil)
	aesStream, err := v2crypto.NewAesEncryptionStream(this.User.ID().CmdKey(), iv)
	if err != nil {
		return nil, err
	}
	aesStream.XORKeyStream(buffer.Value[encryptionBegin:encryptionEnd], buffer.Value[encryptionBegin:encryptionEnd])

	return buffer, nil
}
<|endoftext|>


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

package steven

import (
	"math/rand"
	"sync"
	"unsafe"

	"github.com/thinkofdeath/steven/render"
	"github.com/thinkofdeath/steven/render/builder"
	"github.com/thinkofdeath/steven/type/bit"
	"github.com/thinkofdeath/steven/type/direction"
)

type chunkVertex struct {
	X, Y, Z                    float32
	TX, TY, TW, TH             uint16
	TOffsetX, TOffsetY, TAtlas int16
	Pad0                       int16
	R, G, B                    byte
	Pad1                       byte
	BlockLight, SkyLight       uint16
	Pad2, Pad3                 uint16
}

type buildPos struct {
	X, Y, Z int
}

var (
	_, chunkVertexType = builder.Struct(chunkVertex{})
	builderPool        = sync.Pool{
		New: func() interface{} {
			return []chunkVertex(nil)
		},
	}
)

func (cs *chunkSection) build(complete chan<- buildPos) {
	ox, oy, oz := (cs.chunk.X<<4)-2, (cs.Y<<4)-2, (cs.chunk.Z<<4)-2
	bs := getPooledSnapshot(ox, oy, oz)
	// Make relative
	bs.x = -2
	bs.y = -2
	bs.z = -2
	go func() {
		bO := builderPool.Get().([]chunkVertex)[:0]
		bT := builderPool.Get().([]chunkVertex)[:0]
		bOI := new(int)
		bTI := new(int)

		r := rand.New(rand.NewSource(int64(cs.chunk.X) | (int64(cs.chunk.Z) << 32)))

		for y := 0; y < 16; y++ {
			for x := 0; x < 16; x++ {
				for z := 0; z < 16; z++ {

					bl := bs.block(x, y, z)
					if !bl.Renderable() {
						// Use one step of the rng so that
						// if a block is placed in an empty
						// location is variant doesn't change
						r.Int()
						continue
					}
					bI := bOI
					// Translucent models need special handling
					if bl.IsTranslucent() {
						bI = bTI
					}

					// Liquids can't be represented by the model system
					// due to the number of possible states they have
					if l, ok := bl.(*blockLiquid); ok {
						if bl.IsTranslucent() {
							bT = l.renderLiquid(bs, x, y, z, bT, bI)
						} else {
							bO = l.renderLiquid(bs, x, y, z, bO, bI)
						}
						r.Int() // See the comment above for air
						continue
					}

					// The random generator is used to select a 'random' variant
					// which is constant for that position.
					if variant := bl.Models().selectModel(r); variant != nil {
						if bl.IsTranslucent() {
							bT = variant.Render(x, y, z, bs, bT, bI)
						} else {
							bO = variant.Render(x, y, z, bs, bO, bI)
						}
					}
				}
			}
		}

		// Update culling information
		cullBits := buildCullBits(bs)
		snapshotPool.Put(bs)

		// Upload the buffers on the render goroutine
		render.Sync(func() {
			var data, dataT []byte
			if len(bO) > 0 {
				size := len(bO) * int(unsafe.Sizeof(bO[0]))
				data = (*[1 << 28]byte)(unsafe.Pointer(&bO[0]))[:size]
			}
			if len(bT) > 0 {
				size := len(bT) * int(unsafe.Sizeof(bT[0]))
				dataT = (*[1 << 28]byte)(unsafe.Pointer(&bT[0]))[:size]
			}
			cs.Buffer.Upload(data, *bOI, cullBits)
			cs.Buffer.UploadTrans(dataT, *bTI)

			builderPool.Put(bO)
			builderPool.Put(bT)
		})
		// Free up the builder
		complete <- buildPos{cs.chunk.X, cs.Y, cs.chunk.Z}
	}()
}

func buildCullBits(bs *blocksSnapshot) uint64 {
	bits := uint64(0)
	set := func(from, to direction.Type) {
		bits |= 1 << (from*6 + to)
	}

	visited := bit.NewSet(16 * 16 * 16)
	// This tries a flood fill on every block in the chunk
	// section with an optimization of not visiting a block
	// that was visited in a previous fill (as it would already
	// be accounted for).
	for y := 0; y < 16; y++ {
		for z := 0; z < 16; z++ {
			for x := 0; x < 16; x++ {
				if visited.Get(x | (z << 4) | (y << 8)) {
					continue
				}
				touched := floodFill(bs, visited, x, y, z)
				// Minor optimization for a common case
				if touched == 0 {
					continue
				}
				// Mark each face in the set as visible through
				// each other
				for _, d := range direction.Values {
					if touched&(1<<d) != 0 {
						for _, d2 := range direction.Values {
							if touched&(1<<d2) != 0 {
								set(d, d2)
							}
						}
					}
				}
			}
		}
	}

	return bits
}

func floodFill(bs *blocksSnapshot, visited bit.Set, x, y, z int) uint8 {
	i := x | (z << 4) | (y << 8)
	// Make sure we aren't filling the same spot repeatedly or
	// going out of bounds.
	if x < 0 || x > 15 || y < 0 || y > 15 || z < 0 || z > 15 || visited.Get(i) {
		return 0
	}
	visited.Set(i, true)

	// Can't fill into 'solid' spaces (ones that completely fill
	// the block)
	if bs.block(x, y, z).ShouldCullAgainst() {
		return 0
	}

	// bits are used to represent touched faces
	touched := uint8(0)
	if x == 0 {
		touched |= 1 << direction.West
	} else if x == 15 {
		touched |= 1 << direction.East
	}
	if y == 0 {
		touched |= 1 << direction.Down
	} else if y == 15 {
		touched |= 1 << direction.Up
	}
	if z == 0 {
		touched |= 1 << direction.North
	} else if z == 15 {
		touched |= 1 << direction.South
	}

	// Fill around us and add the touched faces to our
	// bits
	for _, d := range direction.Values {
		ox, oy, oz := d.Offset()
		touched |= floodFill(bs, visited, x+ox, y+oy, z+oz)
	}

	return touched
}

// builder.Struct works by reflection which is to slow for this
// as its called so often.
func buildVertex(b *builder.Buffer, v chunkVertex) {
	b.Float(v.X)
	b.Float(v.Y)
	b.Float(v.Z)
	b.UnsignedShort(v.TX)
	b.UnsignedShort(v.TY)
	b.UnsignedShort(v.TW)
	b.UnsignedShort(v.TH)
	b.Short(v.TOffsetX)
	b.Short(v.TOffsetY)
	b.Short(v.TAtlas)
	b.Short(0)
	b.UnsignedByte(v.R)
	b.UnsignedByte(v.G)
	b.UnsignedByte(v.B)
	b.UnsignedByte(255)
	b.UnsignedShort(v.BlockLight)
	b.UnsignedShort(v.SkyLight)
	b.UnsignedShort(0)
	b.UnsignedShort(0)
}
<commit_msg>steven: check to make sure the buffer stilll exists before uploading<commit_after>// Copyright 2015 Matthew Collins
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package steven

import (
	"math/rand"
	"sync"
	"unsafe"

	"github.com/thinkofdeath/steven/render"
	"github.com/thinkofdeath/steven/render/builder"
	"github.com/thinkofdeath/steven/type/bit"
	"github.com/thinkofdeath/steven/type/direction"
)

type chunkVertex struct {
	X, Y, Z                    float32
	TX, TY, TW, TH             uint16
	TOffsetX, TOffsetY, TAtlas int16
	Pad0                       int16
	R, G, B                    byte
	Pad1                       byte
	BlockLight, SkyLight       uint16
	Pad2, Pad3                 uint16
}

type buildPos struct {
	X, Y, Z int
}

var (
	_, chunkVertexType = builder.Struct(chunkVertex{})
	builderPool        = sync.Pool{
		New: func() interface{} {
			return []chunkVertex(nil)
		},
	}
)

func (cs *chunkSection) build(complete chan<- buildPos) {
	ox, oy, oz := (cs.chunk.X<<4)-2, (cs.Y<<4)-2, (cs.chunk.Z<<4)-2
	bs := getPooledSnapshot(ox, oy, oz)
	// Make relative
	bs.x = -2
	bs.y = -2
	bs.z = -2
	go func() {
		bO := builderPool.Get().([]chunkVertex)[:0]
		bT := builderPool.Get().([]chunkVertex)[:0]
		bOI := new(int)
		bTI := new(int)

		r := rand.New(rand.NewSource(int64(cs.chunk.X) | (int64(cs.chunk.Z) << 32)))

		for y := 0; y < 16; y++ {
			for x := 0; x < 16; x++ {
				for z := 0; z < 16; z++ {

					bl := bs.block(x, y, z)
					if !bl.Renderable() {
						// Use one step of the rng so that
						// if a block is placed in an empty
						// location is variant doesn't change
						r.Int()
						continue
					}
					bI := bOI
					// Translucent models need special handling
					if bl.IsTranslucent() {
						bI = bTI
					}

					// Liquids can't be represented by the model system
					// due to the number of possible states they have
					if l, ok := bl.(*blockLiquid); ok {
						if bl.IsTranslucent() {
							bT = l.renderLiquid(bs, x, y, z, bT, bI)
						} else {
							bO = l.renderLiquid(bs, x, y, z, bO, bI)
						}
						r.Int() // See the comment above for air
						continue
					}

					// The random generator is used to select a 'random' variant
					// which is constant for that position.
					if variant := bl.Models().selectModel(r); variant != nil {
						if bl.IsTranslucent() {
							bT = variant.Render(x, y, z, bs, bT, bI)
						} else {
							bO = variant.Render(x, y, z, bs, bO, bI)
						}
					}
				}
			}
		}

		// Update culling information
		cullBits := buildCullBits(bs)
		snapshotPool.Put(bs)

		// Upload the buffers on the render goroutine
		render.Sync(func() {
			if cs.Buffer != nil {
				var data, dataT []byte
				if len(bO) > 0 {
					size := len(bO) * int(unsafe.Sizeof(bO[0]))
					data = (*[1 << 28]byte)(unsafe.Pointer(&bO[0]))[:size]
				}
				if len(bT) > 0 {
					size := len(bT) * int(unsafe.Sizeof(bT[0]))
					dataT = (*[1 << 28]byte)(unsafe.Pointer(&bT[0]))[:size]
				}
				cs.Buffer.Upload(data, *bOI, cullBits)
				cs.Buffer.UploadTrans(dataT, *bTI)
			}

			builderPool.Put(bO)
			builderPool.Put(bT)
		})
		// Free up the builder
		complete <- buildPos{cs.chunk.X, cs.Y, cs.chunk.Z}
	}()
}

func buildCullBits(bs *blocksSnapshot) uint64 {
	bits := uint64(0)
	set := func(from, to direction.Type) {
		bits |= 1 << (from*6 + to)
	}

	visited := bit.NewSet(16 * 16 * 16)
	// This tries a flood fill on every block in the chunk
	// section with an optimization of not visiting a block
	// that was visited in a previous fill (as it would already
	// be accounted for).
	for y := 0; y < 16; y++ {
		for z := 0; z < 16; z++ {
			for x := 0; x < 16; x++ {
				if visited.Get(x | (z << 4) | (y << 8)) {
					continue
				}
				touched := floodFill(bs, visited, x, y, z)
				// Minor optimization for a common case
				if touched == 0 {
					continue
				}
				// Mark each face in the set as visible through
				// each other
				for _, d := range direction.Values {
					if touched&(1<<d) != 0 {
						for _, d2 := range direction.Values {
							if touched&(1<<d2) != 0 {
								set(d, d2)
							}
						}
					}
				}
			}
		}
	}

	return bits
}

func floodFill(bs *blocksSnapshot, visited bit.Set, x, y, z int) uint8 {
	i := x | (z << 4) | (y << 8)
	// Make sure we aren't filling the same spot repeatedly or
	// going out of bounds.
	if x < 0 || x > 15 || y < 0 || y > 15 || z < 0 || z > 15 || visited.Get(i) {
		return 0
	}
	visited.Set(i, true)

	// Can't fill into 'solid' spaces (ones that completely fill
	// the block)
	if bs.block(x, y, z).ShouldCullAgainst() {
		return 0
	}

	// bits are used to represent touched faces
	touched := uint8(0)
	if x == 0 {
		touched |= 1 << direction.West
	} else if x == 15 {
		touched |= 1 << direction.East
	}
	if y == 0 {
		touched |= 1 << direction.Down
	} else if y == 15 {
		touched |= 1 << direction.Up
	}
	if z == 0 {
		touched |= 1 << direction.North
	} else if z == 15 {
		touched |= 1 << direction.South
	}

	// Fill around us and add the touched faces to our
	// bits
	for _, d := range direction.Values {
		ox, oy, oz := d.Offset()
		touched |= floodFill(bs, visited, x+ox, y+oy, z+oz)
	}

	return touched
}

// builder.Struct works by reflection which is to slow for this
// as its called so often.
func buildVertex(b *builder.Buffer, v chunkVertex) {
	b.Float(v.X)
	b.Float(v.Y)
	b.Float(v.Z)
	b.UnsignedShort(v.TX)
	b.UnsignedShort(v.TY)
	b.UnsignedShort(v.TW)
	b.UnsignedShort(v.TH)
	b.Short(v.TOffsetX)
	b.Short(v.TOffsetY)
	b.Short(v.TAtlas)
	b.Short(0)
	b.UnsignedByte(v.R)
	b.UnsignedByte(v.G)
	b.UnsignedByte(v.B)
	b.UnsignedByte(255)
	b.UnsignedShort(v.BlockLight)
	b.UnsignedShort(v.SkyLight)
	b.UnsignedShort(0)
	b.UnsignedShort(0)
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The cert-manager Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package install

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"helm.sh/helm/v3/pkg/action"
	"helm.sh/helm/v3/pkg/chart"
	"helm.sh/helm/v3/pkg/chart/loader"
	"helm.sh/helm/v3/pkg/cli"
	"helm.sh/helm/v3/pkg/cli/values"
	"helm.sh/helm/v3/pkg/getter"
	"helm.sh/helm/v3/pkg/release"
	"k8s.io/cli-runtime/pkg/genericclioptions"

	"github.com/jetstack/cert-manager/cmd/ctl/pkg/build"
	"github.com/jetstack/cert-manager/cmd/ctl/pkg/install/helm"
)

type InstallOptions struct {
	settings  *cli.EnvSettings
	client    *action.Install
	cfg       *action.Configuration
	valueOpts *values.Options

	ChartName string
	DryRun    bool
	Wait      bool

	genericclioptions.IOStreams
}

const (
	installCRDsFlagName         = "installCRDs"
	defaultCertManagerNamespace = "cert-manager"
)

func installDesc() string {
	return build.WithTemplate(`This command installs cert-manager. It uses the Helm libraries to do so.

The latest published cert-manager chart in the "https://charts.jetstack.io" repo is used.
Most of the features supported by 'helm install' are also supported by this command.
In addition, his command will always correctly install the required CRD resources.

Some example uses:
	$ {{.BuildName}} x install
or
	$ {{.BuildName}} x install -n new-cert-manager
or
	$ {{.BuildName}} x install --version v1.4.0
or
	$ {{.BuildName}} x install --set prometheus.enabled=false

To override values in the cert-manager chart, use either the '--values' flag and
pass in a file or use the '--set' flag and pass configuration from the command line.
`)
}

func NewCmdInstall(ctx context.Context, ioStreams genericclioptions.IOStreams) *cobra.Command {
	settings := cli.New()
	cfg := new(action.Configuration)

	options := &InstallOptions{
		settings:  settings,
		cfg:       cfg,
		client:    action.NewInstall(cfg),
		valueOpts: &values.Options{},

		IOStreams: ioStreams,
	}

	cmd := &cobra.Command{
		Use:   "install",
		Short: "Install cert-manager",
		Long:  installDesc(),
		RunE: func(cmd *cobra.Command, args []string) error {
			options.client.Namespace = settings.Namespace()

			rel, err := options.runInstall(ctx)
			if err != nil {
				return err
			}

			if options.DryRun {
				fmt.Fprintf(ioStreams.Out, "%s", rel.Manifest)
				return nil
			}

			printReleaseSummary(ioStreams.Out, rel)
			return nil
		},
		SilenceUsage:  true,
		SilenceErrors: true,
	}

	settings.AddFlags(cmd.Flags())
	cmd.Flag("namespace").DefValue = defaultCertManagerNamespace

	addInstallUninstallFlags(cmd.Flags(), &options.client.Timeout, &options.Wait)

	addInstallFlags(cmd.Flags(), options.client)
	addValueOptionsFlags(cmd.Flags(), options.valueOpts)
	addChartPathOptionsFlags(cmd.Flags(), &options.client.ChartPathOptions)

	cmd.Flags().BoolVar(&options.client.CreateNamespace, "create-namespace", true, "Create the release namespace if not present")
	cmd.Flags().MarkHidden("create-namespace")
	cmd.Flags().StringVar(&options.ChartName, "chart-name", "cert-manager", "Name of the chart to install")
	cmd.Flags().MarkHidden("chart-name")
	cmd.Flags().BoolVar(&options.DryRun, "dry-run", false, "Simulate install and output manifest")

	return cmd
}

// The overall strategy is to install the CRDs first, and not as part of a Helm
// release, and then to install a Helm release without the CRDs.  This is to
// ensure that CRDs are not removed by a subsequent helm uninstall or by a
// future cmctl uninstall. We want the removal of CRDs to only be performed by
// an administrator who understands that the consequences of removing CRDs will
// be the garbage collection of all the related CRs in the cluster.  We first
// do a dry-run install of the chart (effectively helm template
// --validate=false) to render the CRDs from the CRD templates in the Chart.
// The ClientOnly option is required, otherwise Helm will return an error in
// case the CRDs are already installed in the cluster.  We then extract the
// CRDs from the resulting dry-run manifests and install those first.  Finally,
// we perform a helm install to install the remaining non-CRD resources and
// wait for those to be "Ready".
// This creates a Helm "release" artifact in a Secret in the target namespace, which contains
// a record of all the resources installed by Helm (except the CRDs).
func (o *InstallOptions) runInstall(ctx context.Context) (*release.Release, error) {
	log.SetFlags(0)         // Disable prefixing logs with timestamps.
	log.SetOutput(o.ErrOut) // Log everything to stderr so dry-run output does not get corrupted.

	// Find chart
	cp, err := o.client.ChartPathOptions.LocateChart(o.ChartName, o.settings)
	if err != nil {
		return nil, err
	}

	chart, err := loader.Load(cp)
	if err != nil {
		return nil, err
	}

	// Check if chart is installable
	if err := checkIfInstallable(chart); err != nil {
		return nil, err
	}

	// Console print if chart is deprecated
	if chart.Metadata.Deprecated {
		log.Printf("This chart is deprecated")
	}

	// Merge all values flags
	p := getter.All(o.settings)
	chartValues, err := o.valueOpts.MergeValues(p)
	if err != nil {
		return nil, err
	}

	// Dryrun template generation (used for rendering the CRDs in /templates)
	o.client.DryRun = true                  // Do not apply install
	o.client.ClientOnly = true              // Do not validate against cluster (otherwise double CRDs can cause error)
	chartValues[installCRDsFlagName] = true // Make sure to render CRDs
	dryRunResult, err := o.client.Run(chart, chartValues)
	if err != nil {
		return nil, err
	}

	if o.DryRun {
		return dryRunResult, nil
	}

	if err := o.cfg.Init(o.settings.RESTClientGetter(), o.settings.Namespace(), os.Getenv("HELM_DRIVER"), log.Printf); err != nil {
		return nil, err
	}

	// Extract the resource.Info objects from the manifest
	resources, err := helm.ParseMultiDocumentYAML(dryRunResult.Manifest, o.cfg.KubeClient)
	if err != nil {
		return nil, err
	}

	// Filter resource.Info objects and only keep the CRDs
	crds := helm.FilterCrdResources(resources)

	// Abort in case CRDs were not found in chart
	if len(crds) == 0 {
		return nil, fmt.Errorf("Found no CRDs in provided cert-manager chart.")
	}

	// Make sure that no CRDs are currently installed
	originalCRDs, err := helm.FetchResources(crds, o.cfg.KubeClient)
	if err != nil {
		return nil, err
	}

	if len(originalCRDs) > 0 {
		return nil, fmt.Errorf("Found existing installed cert-manager CRDs! Cannot continue with installation.")
	}

	// Install CRDs
	if err := helm.CreateCRDs(crds, o.cfg); err != nil {
		return nil, err
	}

	// Install chart
	o.client.DryRun = false     // Apply DryRun cli flags
	o.client.ClientOnly = false // Perform install against cluster

	o.client.Wait = o.Wait // Wait for resources to be ready
	// If part of the install fails and the Atomic option is set to True,
	// all resource installs are reverted. Atomic cannot be enabled without
	// waiting (if Atomic=True is set, the value for Wait is overwritten with True),
	// so only enable Atomic if we are waiting.
	o.client.Atomic = o.Wait
	// The cert-manager chart currently has only a startupapicheck hook,
	// if waiting is disabled, this hook should be disabled too; otherwise
	// the hook will still wait for the installation to succeed.
	o.client.DisableHooks = !o.Wait

	chartValues[installCRDsFlagName] = false // Do not render CRDs, as this might cause problems when uninstalling using helm

	return o.client.Run(chart, chartValues)
}

func printReleaseSummary(out io.Writer, rel *release.Release) {
	fmt.Fprintf(out, "NAME: %s\n", rel.Name)
	if !rel.Info.LastDeployed.IsZero() {
		fmt.Fprintf(out, "LAST DEPLOYED: %s\n", rel.Info.LastDeployed.Format(time.ANSIC))
	}
	fmt.Fprintf(out, "NAMESPACE: %s\n", rel.Namespace)
	fmt.Fprintf(out, "STATUS: %s\n", rel.Info.Status.String())
	fmt.Fprintf(out, "REVISION: %d\n", rel.Version)
	fmt.Fprintf(out, "DESCRIPTION: %s\n", rel.Info.Description)

	if len(rel.Info.Notes) > 0 {
		fmt.Fprintf(out, "NOTES:\n%s\n", strings.TrimSpace(rel.Info.Notes))
	}
}

// Only Application chart type are installable.
func checkIfInstallable(ch *chart.Chart) error {
	switch ch.Metadata.Type {
	case "", "application":
		return nil
	}
	return fmt.Errorf("%s charts are not installable", ch.Metadata.Type)
}
<commit_msg>Make sure that cmctl x install uses the cert-manager namespace by default<commit_after>/*
Copyright 2021 The cert-manager Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package install

import (
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"helm.sh/helm/v3/pkg/action"
	"helm.sh/helm/v3/pkg/chart"
	"helm.sh/helm/v3/pkg/chart/loader"
	"helm.sh/helm/v3/pkg/cli"
	"helm.sh/helm/v3/pkg/cli/values"
	"helm.sh/helm/v3/pkg/getter"
	"helm.sh/helm/v3/pkg/release"
	"k8s.io/cli-runtime/pkg/genericclioptions"

	"github.com/jetstack/cert-manager/cmd/ctl/pkg/build"
	"github.com/jetstack/cert-manager/cmd/ctl/pkg/install/helm"
)

type InstallOptions struct {
	settings  *cli.EnvSettings
	client    *action.Install
	cfg       *action.Configuration
	valueOpts *values.Options

	ChartName string
	DryRun    bool
	Wait      bool

	genericclioptions.IOStreams
}

const (
	installCRDsFlagName         = "installCRDs"
	defaultCertManagerNamespace = "cert-manager"
)

func installDesc() string {
	return build.WithTemplate(`This command installs cert-manager. It uses the Helm libraries to do so.

The latest published cert-manager chart in the "https://charts.jetstack.io" repo is used.
Most of the features supported by 'helm install' are also supported by this command.
In addition, his command will always correctly install the required CRD resources.

Some example uses:
	$ {{.BuildName}} x install
or
	$ {{.BuildName}} x install -n new-cert-manager
or
	$ {{.BuildName}} x install --version v1.4.0
or
	$ {{.BuildName}} x install --set prometheus.enabled=false

To override values in the cert-manager chart, use either the '--values' flag and
pass in a file or use the '--set' flag and pass configuration from the command line.
`)
}

func NewCmdInstall(ctx context.Context, ioStreams genericclioptions.IOStreams) *cobra.Command {
	settings := cli.New()
	cfg := new(action.Configuration)

	options := &InstallOptions{
		settings:  settings,
		cfg:       cfg,
		client:    action.NewInstall(cfg),
		valueOpts: &values.Options{},

		IOStreams: ioStreams,
	}

	cmd := &cobra.Command{
		Use:   "install",
		Short: "Install cert-manager",
		Long:  installDesc(),
		RunE: func(cmd *cobra.Command, args []string) error {
			// The Helm client flagset and settings are configured to look up
			// the default namespace in the following order:
			// `--namespace` flag value, if the flag is supplied,
			// `--namespace` flag default value, or
			// `HELM_NAMESPACE` environment variable (if set).
			//
			// The Helm cli.New function does not provide an easy way to
			// override the namespace lookup mechanism so here
			// we check whether the user supplied the `--namespace` option,
			// and if not we force the default value to `cert-manager`.
			// The original default value is "" which would result in
			// cert-manager being installed to the "default" namespace.
			// See https://github.com/helm/helm/issues/9790
			//
			// The default value that is shown in the `--help` usage message is
			// overridden elsewhere; when the flagset is configured (below).
			namespaceFlag := cmd.Flags().Lookup("namespace")
			if !namespaceFlag.Changed {
				if err := namespaceFlag.Value.Set(defaultCertManagerNamespace); err != nil {
					return fmt.Errorf("unexpected error while setting the default namespace: %v", err)
				}
			}
			if settings.Debug {
				cmd.DebugFlags()
			}

			options.client.Namespace = settings.Namespace()

			rel, err := options.runInstall(ctx)
			if err != nil {
				return err
			}

			if options.DryRun {
				fmt.Fprintf(ioStreams.Out, "%s", rel.Manifest)
				return nil
			}

			printReleaseSummary(ioStreams.Out, rel)
			return nil
		},
		SilenceUsage:  true,
		SilenceErrors: true,
	}

	settings.AddFlags(cmd.Flags())
	// Here we set the default value shown in the usage message.
	// The actual default value is overridden in the RunE function above.
	cmd.Flag("namespace").DefValue = defaultCertManagerNamespace

	addInstallUninstallFlags(cmd.Flags(), &options.client.Timeout, &options.Wait)

	addInstallFlags(cmd.Flags(), options.client)
	addValueOptionsFlags(cmd.Flags(), options.valueOpts)
	addChartPathOptionsFlags(cmd.Flags(), &options.client.ChartPathOptions)

	cmd.Flags().BoolVar(&options.client.CreateNamespace, "create-namespace", true, "Create the release namespace if not present")
	cmd.Flags().MarkHidden("create-namespace")
	cmd.Flags().StringVar(&options.ChartName, "chart-name", "cert-manager", "Name of the chart to install")
	cmd.Flags().MarkHidden("chart-name")
	cmd.Flags().BoolVar(&options.DryRun, "dry-run", false, "Simulate install and output manifest")

	return cmd
}

// The overall strategy is to install the CRDs first, and not as part of a Helm
// release, and then to install a Helm release without the CRDs.  This is to
// ensure that CRDs are not removed by a subsequent helm uninstall or by a
// future cmctl uninstall. We want the removal of CRDs to only be performed by
// an administrator who understands that the consequences of removing CRDs will
// be the garbage collection of all the related CRs in the cluster.  We first
// do a dry-run install of the chart (effectively helm template
// --validate=false) to render the CRDs from the CRD templates in the Chart.
// The ClientOnly option is required, otherwise Helm will return an error in
// case the CRDs are already installed in the cluster.  We then extract the
// CRDs from the resulting dry-run manifests and install those first.  Finally,
// we perform a helm install to install the remaining non-CRD resources and
// wait for those to be "Ready".
// This creates a Helm "release" artifact in a Secret in the target namespace, which contains
// a record of all the resources installed by Helm (except the CRDs).
func (o *InstallOptions) runInstall(ctx context.Context) (*release.Release, error) {
	log.SetFlags(0)         // Disable prefixing logs with timestamps.
	log.SetOutput(o.ErrOut) // Log everything to stderr so dry-run output does not get corrupted.

	// Find chart
	cp, err := o.client.ChartPathOptions.LocateChart(o.ChartName, o.settings)
	if err != nil {
		return nil, err
	}

	chart, err := loader.Load(cp)
	if err != nil {
		return nil, err
	}

	// Check if chart is installable
	if err := checkIfInstallable(chart); err != nil {
		return nil, err
	}

	// Console print if chart is deprecated
	if chart.Metadata.Deprecated {
		log.Printf("This chart is deprecated")
	}

	// Merge all values flags
	p := getter.All(o.settings)
	chartValues, err := o.valueOpts.MergeValues(p)
	if err != nil {
		return nil, err
	}

	// Dryrun template generation (used for rendering the CRDs in /templates)
	o.client.DryRun = true                  // Do not apply install
	o.client.ClientOnly = true              // Do not validate against cluster (otherwise double CRDs can cause error)
	chartValues[installCRDsFlagName] = true // Make sure to render CRDs
	dryRunResult, err := o.client.Run(chart, chartValues)
	if err != nil {
		return nil, err
	}

	if o.DryRun {
		return dryRunResult, nil
	}

	if err := o.cfg.Init(o.settings.RESTClientGetter(), o.settings.Namespace(), os.Getenv("HELM_DRIVER"), log.Printf); err != nil {
		return nil, err
	}

	// Extract the resource.Info objects from the manifest
	resources, err := helm.ParseMultiDocumentYAML(dryRunResult.Manifest, o.cfg.KubeClient)
	if err != nil {
		return nil, err
	}

	// Filter resource.Info objects and only keep the CRDs
	crds := helm.FilterCrdResources(resources)

	// Abort in case CRDs were not found in chart
	if len(crds) == 0 {
		return nil, fmt.Errorf("Found no CRDs in provided cert-manager chart.")
	}

	// Make sure that no CRDs are currently installed
	originalCRDs, err := helm.FetchResources(crds, o.cfg.KubeClient)
	if err != nil {
		return nil, err
	}

	if len(originalCRDs) > 0 {
		return nil, fmt.Errorf("Found existing installed cert-manager CRDs! Cannot continue with installation.")
	}

	// Install CRDs
	if err := helm.CreateCRDs(crds, o.cfg); err != nil {
		return nil, err
	}

	// Install chart
	o.client.DryRun = false     // Apply DryRun cli flags
	o.client.ClientOnly = false // Perform install against cluster

	o.client.Wait = o.Wait // Wait for resources to be ready
	// If part of the install fails and the Atomic option is set to True,
	// all resource installs are reverted. Atomic cannot be enabled without
	// waiting (if Atomic=True is set, the value for Wait is overwritten with True),
	// so only enable Atomic if we are waiting.
	o.client.Atomic = o.Wait
	// The cert-manager chart currently has only a startupapicheck hook,
	// if waiting is disabled, this hook should be disabled too; otherwise
	// the hook will still wait for the installation to succeed.
	o.client.DisableHooks = !o.Wait

	chartValues[installCRDsFlagName] = false // Do not render CRDs, as this might cause problems when uninstalling using helm

	return o.client.Run(chart, chartValues)
}

func printReleaseSummary(out io.Writer, rel *release.Release) {
	fmt.Fprintf(out, "NAME: %s\n", rel.Name)
	if !rel.Info.LastDeployed.IsZero() {
		fmt.Fprintf(out, "LAST DEPLOYED: %s\n", rel.Info.LastDeployed.Format(time.ANSIC))
	}
	fmt.Fprintf(out, "NAMESPACE: %s\n", rel.Namespace)
	fmt.Fprintf(out, "STATUS: %s\n", rel.Info.Status.String())
	fmt.Fprintf(out, "REVISION: %d\n", rel.Version)
	fmt.Fprintf(out, "DESCRIPTION: %s\n", rel.Info.Description)

	if len(rel.Info.Notes) > 0 {
		fmt.Fprintf(out, "NOTES:\n%s\n", strings.TrimSpace(rel.Info.Notes))
	}
}

// Only Application chart type are installable.
func checkIfInstallable(ch *chart.Chart) error {
	switch ch.Metadata.Type {
	case "", "application":
		return nil
	}
	return fmt.Errorf("%s charts are not installable", ch.Metadata.Type)
}
<|endoftext|>