<commit_before>package restart

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/masterzen/winrm"
)

var DefaultRestartCommand = "shutdown /r /f /t 0 /c \"packer restart\""
var DefaultRestartCheckCommand = winrm.Powershell(`echo "${env:COMPUTERNAME} restarted."`)
var retryableSleep = 5 * time.Second
var TryCheckReboot = "shutdown.exe -f -r -t 60"
var AbortReboot = "shutdown.exe -a"

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The command used to restart the guest machine
	RestartCommand string `mapstructure:"restart_command"`

	// The command used to check if the guest machine has restarted
	// The output of this command will be displayed to the user
	RestartCheckCommand string `mapstructure:"restart_check_command"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	ctx interpolate.Context
}

type Provisioner struct {
	config     Config
	comm       packer.Communicator
	ui         packer.Ui
	cancel     chan struct{}
	cancelLock sync.Mutex
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartCommand == "" {
		p.config.RestartCommand = DefaultRestartCommand
	}

	if p.config.RestartCheckCommand == "" {
		p.config.RestartCheckCommand = DefaultRestartCheckCommand
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 5 * time.Minute
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	p.cancelLock.Lock()
	p.cancel = make(chan struct{})
	p.cancelLock.Unlock()

	ui.Say("Restarting Machine")
	p.comm = comm
	p.ui = ui

	var cmd *packer.RemoteCmd
	command := p.config.RestartCommand
	err := p.retryable(func() error {
		cmd = &packer.RemoteCmd{Command: command}
		return cmd.StartWithUi(comm, ui)
	})

	if err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Restart script exited with non-zero exit status: %d", cmd.ExitStatus)
	}

	return waitForRestart(p, comm)
}

var waitForRestart = func(p *Provisioner, comm packer.Communicator) error {
	ui := p.ui
	ui.Say("Waiting for machine to restart...")
	waitDone := make(chan bool, 1)
	timeout := time.After(p.config.RestartTimeout)
	var err error

	p.comm = comm
	var cmd *packer.RemoteCmd
	trycommand := TryCheckReboot
	abortcommand := AbortReboot
	// Stolen from Vagrant reboot checker
	for {
		log.Printf("Check if machine is rebooting...")
		cmd = &packer.RemoteCmd{Command: trycommand}
		err = cmd.StartWithUi(comm, ui)
		if err != nil {
			// Couldn't execute, we assume machine is rebooting already
			break
		}

		if cmd.ExitStatus == 1115 || cmd.ExitStatus == 1190 {
			// Reboot already in progress but not completed
			log.Printf("Reboot already in progress, waiting...")
			time.Sleep(10 * time.Second)
		}
		if cmd.ExitStatus == 0 {
			// Cancel reboot we created to test if machine was already rebooting
			cmd = &packer.RemoteCmd{Command: abortcommand}
			cmd.StartWithUi(comm, ui)
			break
		}
	}

	go func() {
		log.Printf("Waiting for machine to become available...")
		err = waitForCommunicator(p)
		waitDone <- true
	}()

	log.Printf("Waiting for machine to reboot with timeout: %s", p.config.RestartTimeout)

WaitLoop:
	for {
		// Wait for either WinRM to become available, a timeout to occur,
		// or an interrupt to come through.
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for machine to restart: %s", err))
				return err
			}

			ui.Say("Machine successfully restarted, moving on")
			close(p.cancel)
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for machine to restart.")
			ui.Error(err.Error())
			close(p.cancel)
			return err
		case <-p.cancel:
			close(waitDone)
			return fmt.Errorf("Interrupt detected, quitting waiting for machine to restart")
		}
	}

	return nil

}

var waitForCommunicator = func(p *Provisioner) error {
	runCustomRestartCheck := true
	if p.config.RestartCheckCommand == DefaultRestartCheckCommand {
		runCustomRestartCheck = false
	}
	// this is the user configurable command
	cmdRestartCheck := &packer.RemoteCmd{Command: p.config.RestartCheckCommand}
	log.Printf("Checking that communicator is connected with: '%s'",
		cmdRestartCheck.Command)
	for {
		select {
		case <-p.cancel:
			log.Println("Communicator wait canceled, exiting loop")
			return fmt.Errorf("Communicator wait canceled")
		case <-time.After(retryableSleep):
		}
		if runCustomRestartCheck == true {
			// run user-configured restart check
			err := cmdRestartCheck.StartWithUi(p.comm, p.ui)
			if err != nil {
				log.Printf("Communication connection err: %s", err)
				continue
			}
			log.Printf("Connected to machine")
			runCustomRestartCheck = false
		}
		// this is the non-user-configurable check that powershell
		// modules have loaded
		cmdModuleLoad := &packer.RemoteCmd{Command: DefaultRestartCheckCommand}
		var buf, buf2 bytes.Buffer
		cmdModuleLoad.Stdout = &buf
		cmdModuleLoad.Stdout = io.MultiWriter(cmdModuleLoad.Stdout, &buf2)

		cmdModuleLoad.StartWithUi(p.comm, p.ui)
		stdoutToRead := buf2.String()
		if !strings.Contains(stdoutToRead, "restarted.") {
			log.Printf("echo didn't succeed; retrying...")
			continue
		}
		break
	}

	return nil
}

func (p *Provisioner) Cancel() {
	log.Printf("Received interrupt Cancel()")

	p.cancelLock.Lock()
	defer p.cancelLock.Unlock()
	if p.cancel != nil {
		close(p.cancel)
	}
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.RestartTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Print(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(retryableSleep)
		}
	}
}
<commit_msg>don't pipe restarted stuff through the ui<commit_after>package restart

import (
	"bytes"
	"fmt"

	"log"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/masterzen/winrm"
)

var DefaultRestartCommand = "shutdown /r /f /t 0 /c \"packer restart\""
var DefaultRestartCheckCommand = winrm.Powershell(`echo "${env:COMPUTERNAME} restarted."`)
var retryableSleep = 5 * time.Second
var TryCheckReboot = "shutdown.exe -f -r -t 60"
var AbortReboot = "shutdown.exe -a"

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The command used to restart the guest machine
	RestartCommand string `mapstructure:"restart_command"`

	// The command used to check if the guest machine has restarted
	// The output of this command will be displayed to the user
	RestartCheckCommand string `mapstructure:"restart_check_command"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	ctx interpolate.Context
}

type Provisioner struct {
	config     Config
	comm       packer.Communicator
	ui         packer.Ui
	cancel     chan struct{}
	cancelLock sync.Mutex
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartCommand == "" {
		p.config.RestartCommand = DefaultRestartCommand
	}

	if p.config.RestartCheckCommand == "" {
		p.config.RestartCheckCommand = DefaultRestartCheckCommand
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 5 * time.Minute
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	p.cancelLock.Lock()
	p.cancel = make(chan struct{})
	p.cancelLock.Unlock()

	ui.Say("Restarting Machine")
	p.comm = comm
	p.ui = ui

	var cmd *packer.RemoteCmd
	command := p.config.RestartCommand
	err := p.retryable(func() error {
		cmd = &packer.RemoteCmd{Command: command}
		return cmd.StartWithUi(comm, ui)
	})

	if err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Restart script exited with non-zero exit status: %d", cmd.ExitStatus)
	}

	return waitForRestart(p, comm)
}

var waitForRestart = func(p *Provisioner, comm packer.Communicator) error {
	ui := p.ui
	ui.Say("Waiting for machine to restart...")
	waitDone := make(chan bool, 1)
	timeout := time.After(p.config.RestartTimeout)
	var err error

	p.comm = comm
	var cmd *packer.RemoteCmd
	trycommand := TryCheckReboot
	abortcommand := AbortReboot
	// Stolen from Vagrant reboot checker
	for {
		log.Printf("Check if machine is rebooting...")
		cmd = &packer.RemoteCmd{Command: trycommand}
		err = cmd.StartWithUi(comm, ui)
		if err != nil {
			// Couldn't execute, we assume machine is rebooting already
			break
		}

		if cmd.ExitStatus == 1115 || cmd.ExitStatus == 1190 {
			// Reboot already in progress but not completed
			log.Printf("Reboot already in progress, waiting...")
			time.Sleep(10 * time.Second)
		}
		if cmd.ExitStatus == 0 {
			// Cancel reboot we created to test if machine was already rebooting
			cmd = &packer.RemoteCmd{Command: abortcommand}
			cmd.StartWithUi(comm, ui)
			break
		}
	}

	go func() {
		log.Printf("Waiting for machine to become available...")
		err = waitForCommunicator(p)
		waitDone <- true
	}()

	log.Printf("Waiting for machine to reboot with timeout: %s", p.config.RestartTimeout)

WaitLoop:
	for {
		// Wait for either WinRM to become available, a timeout to occur,
		// or an interrupt to come through.
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for machine to restart: %s", err))
				return err
			}

			ui.Say("Machine successfully restarted, moving on")
			close(p.cancel)
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for machine to restart.")
			ui.Error(err.Error())
			close(p.cancel)
			return err
		case <-p.cancel:
			close(waitDone)
			return fmt.Errorf("Interrupt detected, quitting waiting for machine to restart")
		}
	}

	return nil

}

var waitForCommunicator = func(p *Provisioner) error {
	runCustomRestartCheck := true
	if p.config.RestartCheckCommand == DefaultRestartCheckCommand {
		runCustomRestartCheck = false
	}
	// this is the user configurable command
	cmdRestartCheck := &packer.RemoteCmd{Command: p.config.RestartCheckCommand}
	log.Printf("Checking that communicator is connected with: '%s'",
		cmdRestartCheck.Command)
	for {
		select {
		case <-p.cancel:
			log.Println("Communicator wait canceled, exiting loop")
			return fmt.Errorf("Communicator wait canceled")
		case <-time.After(retryableSleep):
		}
		if runCustomRestartCheck == true {
			// run user-configured restart check
			err := cmdRestartCheck.StartWithUi(p.comm, p.ui)
			if err != nil {
				log.Printf("Communication connection err: %s", err)
				continue
			}
			log.Printf("Connected to machine")
			runCustomRestartCheck = false
		}
		// this is the non-user-configurable check that powershell
		// modules have loaded
		var buf bytes.Buffer
		cmdModuleLoad := &packer.RemoteCmd{
			Command: DefaultRestartCheckCommand,
			Stdin:   nil,
			Stdout:  &buf,
			Stderr:  &buf}

		// cmdModuleLoad.StartWithUi(p.comm, p.ui)
		p.comm.Start(cmdModuleLoad)
		cmdModuleLoad.Wait()

		stdoutToRead := buf.String()
		if !strings.Contains(stdoutToRead, "restarted.") {
			log.Printf("echo didn't succeed; retrying...")
			continue
		}
		break
	}

	return nil
}

func (p *Provisioner) Cancel() {
	log.Printf("Received interrupt Cancel()")

	p.cancelLock.Lock()
	defer p.cancelLock.Unlock()
	if p.cancel != nil {
		close(p.cancel)
	}
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.RestartTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Print(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(retryableSleep)
		}
	}
}
<|endoftext|>


<commit_before>package set1_test

import (
	"bufio"
	"os"

	. "github.com/dcarley/cryptopals/set1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	"encoding/base64"
	"encoding/hex"
)

var _ = Describe("Set1", func() {
	Describe("Challenge1", func() {
		Describe("HexToBase64", func() {
			It("should convert example", func() {
				b64, err := HexToBase64([]byte("49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d"))
				Expect(err).ToNot(HaveOccurred())
				Expect(b64).To(Equal([]byte("SSdtIGtpbGxpbmcgeW91ciBicmFpbiBsaWtlIGEgcG9pc29ub3VzIG11c2hyb29t")))
			})
		})

		Describe("HexDecode", func() {
			It("should decode hex to decimal byte slice", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, hex.EncodedLen(len(input)))
				hex.Encode(encoded, input)

				decoded, err := HexDecode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should handle uppercase and lowercase alphas", func() {
				decoded, err := HexDecode([]byte("6a6B6c6D"))
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal([]byte("jklm")))
			})

			It("should return an error for invalid alphas", func() {
				decoded, err := HexDecode([]byte("6g"))
				Expect(err).To(MatchError("invalid hex character: g"))
				Expect(decoded).To(Equal([]byte{}))
			})

			It("should return an error on odd input sizes", func() {
				decoded, err := HexDecode([]byte("abc"))
				Expect(err).To(MatchError("input must be an even size"))
				Expect(decoded).To(Equal([]byte{}))
			})
		})

		Describe("HexEncode", func() {
			It("should encode decimal to hex byte slice", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, hex.EncodedLen(len(input)))
				hex.Encode(encoded, input)

				Expect(HexEncode(input)).To(Equal(encoded))
			})
		})

		Describe("Base64Decode", func() {
			It("should decode base64 to text", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should decode base64 to text with one character padding", func() {
				input := []byte("hello gophers")
				Expect(len(input) % 3).To(Equal(1))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should decode base64 to text with two character padding", func() {
				input := []byte("hello gophers!")
				Expect(len(input) % 3).To(Equal(2))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should return an error for invalid base64 characters", func() {
				decoded, err := Base64Decode([]byte("abc!def"))
				Expect(err).To(MatchError("invalid base64 character: !"))
				Expect(decoded).To(Equal([]byte{}))
			})
		})

		Describe("Base64Encode", func() {
			It("should encode text to base64", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})

			It("should encode text to base64 with one character padding", func() {
				input := []byte("hello gophers")
				Expect(len(input) % 3).To(Equal(1))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})

			It("should encode text to base64 with two character padding", func() {
				input := []byte("hello gophers!")
				Expect(len(input) % 3).To(Equal(2))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})
		})
	})

	Describe("Challenge2", func() {
		Describe("FixedKeyXOR", func() {
			It("should convert example", func() {
				xor1, err := HexDecode([]byte("1c0111001f010100061a024b53535009181c"))
				Expect(err).ToNot(HaveOccurred())
				xor2, err := HexDecode([]byte("686974207468652062756c6c277320657965"))
				Expect(err).ToNot(HaveOccurred())

				xor, err := FixedKeyXOR(xor1, xor2)
				Expect(err).ToNot(HaveOccurred())
				Expect(xor).To(Equal([]byte("the kid don't play")))

				Expect(HexEncode(xor)).To(Equal([]byte("746865206b696420646f6e277420706c6179")))
			})

			It("should error on unequal lengths", func() {
				xor, err := FixedKeyXOR(
					[]byte("12345678"),
					[]byte("1234"),
				)
				Expect(err).To(MatchError("text and key must be same size: 8 != 4"))
				Expect(xor).To(Equal([]byte{}))
			})
		})
	})

	Describe("Challenge3", func() {
		Describe("BruteForceSingleByteXOR", func() {
			It("should convert example", func() {
				xor, err := HexDecode([]byte("1b37373331363f78151b7f2b783431333d78397828372d363c78373e783a393b3736"))
				Expect(err).ToNot(HaveOccurred())

				score, err := BruteForceSingleByteXOR(xor)
				Expect(err).ToNot(HaveOccurred())
				Expect(score.Text).To(Equal([]byte("Cooking MC's like a pound of bacon")))
			})
		})

		DescribeTable("ScoreEnglish",
			func(text []byte, score int) {
				Expect(ScoreEnglish(text)).To(Equal(score))
			},
			Entry("repeated character", []byte("xxxxxxxxxxxxxxxxxxxxxxx"), 0),
			Entry("pwgen 23 -y", []byte("qui1Chux(euZae9Ua3pooqu"), 13),
			Entry("keyboard bashing", []byte("dgj lqn0[jr1n3ofe we[of w"), 12),
			Entry("numbers only", []byte("01234567890123456789012"), 0),
			Entry("proper English", []byte("I'm writing proper English"), 19),
			Entry("real sentence", []byte("This is a real sentence"), 22),
		)
	})

	Describe("Challenge4", func() {
		It("should solve example", func() {
			file, err := os.Open("fixtures/s1c4")
			Expect(err).ToNot(HaveOccurred())
			defer file.Close()

			highestScore := KeyScore{}
			scanner := bufio.NewScanner(file)
			for scanner.Scan() {
				text, err := HexDecode(scanner.Bytes())
				Expect(err).ToNot(HaveOccurred())

				score, err := BruteForceSingleByteXOR(text)
				Expect(err).ToNot(HaveOccurred())

				if score.Score > highestScore.Score {
					highestScore = score
				}
			}

			Expect(highestScore.Text).To(Equal([]byte("Now that the party is jumping\n")))
		})
	})
})
<commit_msg>s1c5: Implementation<commit_after>package set1_test

import (
	"bufio"
	"os"

	. "github.com/dcarley/cryptopals/set1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	"encoding/base64"
	"encoding/hex"
)

var _ = Describe("Set1", func() {
	Describe("Challenge1", func() {
		Describe("HexToBase64", func() {
			It("should convert example", func() {
				b64, err := HexToBase64([]byte("49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d"))
				Expect(err).ToNot(HaveOccurred())
				Expect(b64).To(Equal([]byte("SSdtIGtpbGxpbmcgeW91ciBicmFpbiBsaWtlIGEgcG9pc29ub3VzIG11c2hyb29t")))
			})
		})

		Describe("HexDecode", func() {
			It("should decode hex to decimal byte slice", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, hex.EncodedLen(len(input)))
				hex.Encode(encoded, input)

				decoded, err := HexDecode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should handle uppercase and lowercase alphas", func() {
				decoded, err := HexDecode([]byte("6a6B6c6D"))
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal([]byte("jklm")))
			})

			It("should return an error for invalid alphas", func() {
				decoded, err := HexDecode([]byte("6g"))
				Expect(err).To(MatchError("invalid hex character: g"))
				Expect(decoded).To(Equal([]byte{}))
			})

			It("should return an error on odd input sizes", func() {
				decoded, err := HexDecode([]byte("abc"))
				Expect(err).To(MatchError("input must be an even size"))
				Expect(decoded).To(Equal([]byte{}))
			})
		})

		Describe("HexEncode", func() {
			It("should encode decimal to hex byte slice", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, hex.EncodedLen(len(input)))
				hex.Encode(encoded, input)

				Expect(HexEncode(input)).To(Equal(encoded))
			})
		})

		Describe("Base64Decode", func() {
			It("should decode base64 to text", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should decode base64 to text with one character padding", func() {
				input := []byte("hello gophers")
				Expect(len(input) % 3).To(Equal(1))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should decode base64 to text with two character padding", func() {
				input := []byte("hello gophers!")
				Expect(len(input) % 3).To(Equal(2))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				decoded, err := Base64Decode(encoded)
				Expect(err).ToNot(HaveOccurred())
				Expect(decoded).To(Equal(input))
			})

			It("should return an error for invalid base64 characters", func() {
				decoded, err := Base64Decode([]byte("abc!def"))
				Expect(err).To(MatchError("invalid base64 character: !"))
				Expect(decoded).To(Equal([]byte{}))
			})
		})

		Describe("Base64Encode", func() {
			It("should encode text to base64", func() {
				input := []byte("hello gopher")

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})

			It("should encode text to base64 with one character padding", func() {
				input := []byte("hello gophers")
				Expect(len(input) % 3).To(Equal(1))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})

			It("should encode text to base64 with two character padding", func() {
				input := []byte("hello gophers!")
				Expect(len(input) % 3).To(Equal(2))

				encoded := make([]byte, base64.StdEncoding.EncodedLen(len(input)))
				base64.StdEncoding.Encode(encoded, input)

				Expect(Base64Encode(input)).To(Equal(encoded))
			})
		})
	})

	Describe("Challenge2", func() {
		Describe("FixedKeyXOR", func() {
			It("should convert example", func() {
				xor1, err := HexDecode([]byte("1c0111001f010100061a024b53535009181c"))
				Expect(err).ToNot(HaveOccurred())
				xor2, err := HexDecode([]byte("686974207468652062756c6c277320657965"))
				Expect(err).ToNot(HaveOccurred())

				xor, err := FixedKeyXOR(xor1, xor2)
				Expect(err).ToNot(HaveOccurred())
				Expect(xor).To(Equal([]byte("the kid don't play")))

				Expect(HexEncode(xor)).To(Equal([]byte("746865206b696420646f6e277420706c6179")))
			})

			It("should error on unequal lengths", func() {
				xor, err := FixedKeyXOR(
					[]byte("12345678"),
					[]byte("1234"),
				)
				Expect(err).To(MatchError("text and key must be same size: 8 != 4"))
				Expect(xor).To(Equal([]byte{}))
			})
		})
	})

	Describe("Challenge3", func() {
		Describe("BruteForceSingleByteXOR", func() {
			It("should convert example", func() {
				xor, err := HexDecode([]byte("1b37373331363f78151b7f2b783431333d78397828372d363c78373e783a393b3736"))
				Expect(err).ToNot(HaveOccurred())

				score, err := BruteForceSingleByteXOR(xor)
				Expect(err).ToNot(HaveOccurred())
				Expect(score.Text).To(Equal([]byte("Cooking MC's like a pound of bacon")))
			})
		})

		DescribeTable("ScoreEnglish",
			func(text []byte, score int) {
				Expect(ScoreEnglish(text)).To(Equal(score))
			},
			Entry("repeated character", []byte("xxxxxxxxxxxxxxxxxxxxxxx"), 0),
			Entry("pwgen 23 -y", []byte("qui1Chux(euZae9Ua3pooqu"), 13),
			Entry("keyboard bashing", []byte("dgj lqn0[jr1n3ofe we[of w"), 12),
			Entry("numbers only", []byte("01234567890123456789012"), 0),
			Entry("proper English", []byte("I'm writing proper English"), 19),
			Entry("real sentence", []byte("This is a real sentence"), 22),
		)
	})

	Describe("Challenge4", func() {
		It("should solve example", func() {
			file, err := os.Open("fixtures/s1c4")
			Expect(err).ToNot(HaveOccurred())
			defer file.Close()

			highestScore := KeyScore{}
			scanner := bufio.NewScanner(file)
			for scanner.Scan() {
				text, err := HexDecode(scanner.Bytes())
				Expect(err).ToNot(HaveOccurred())

				score, err := BruteForceSingleByteXOR(text)
				Expect(err).ToNot(HaveOccurred())

				if score.Score > highestScore.Score {
					highestScore = score
				}
			}

			Expect(highestScore.Text).To(Equal([]byte("Now that the party is jumping\n")))
		})
	})

	Describe("Challenge5", func() {
		It("should solve example", func() {
			output, err := RepeatingKeyXOR(
				[]byte(`Burning 'em, if you ain't quick and nimble
I go crazy when I hear a cymbal`),
				[]byte("ICE"),
			)
			Expect(err).ToNot(HaveOccurred())
			Expect(HexEncode(output)).To(Equal([]byte("0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f")))
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"github.com/spf13/cobra"
	"github.com/spf13/cobra/doc"
)

var (
	mangenCmd = &cobra.Command{
		Use:   "man-generation [path]",
		Short: "Creates unix style manpages.",
		Long: "Creates a man pages at the specified " +
			"directory.\n\n"

		Run: wrap(mangencmd),
	}
)

func mangencmd(path string) {
	rootCmd.GenManTree(path)
}
<commit_msg>removed source entry<commit_after>package main

import (
	"github.com/spf13/cobra"
	"github.com/spf13/cobra/doc"
)

var (
	mangenCmd = &cobra.Command{
		Use:   "man-generation [path]",
		Short: "Creates unix style manpages.",
		Long: "Creates a man pages at the specified " +
			"directory.\n\n",

		Run: wrap(mangencmd),
	}
)

func mangencmd(path string) {
	header := &doc.GenManHeader{
		Section: "1",
		Manual:  "siac Manual",
		Source:  "",
	}

	doc.GenManTree(rootCmd, header, "/tmp")
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative Authors

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

    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 handler

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"go.uber.org/atomic"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	network "knative.dev/networking/pkg"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/ptr"
	rtesting "knative.dev/pkg/reconciler/testing"
	"knative.dev/pkg/tracing"
	tracingconfig "knative.dev/pkg/tracing/config"
	tracetesting "knative.dev/pkg/tracing/testing"
	"knative.dev/serving/pkg/activator"
	activatorconfig "knative.dev/serving/pkg/activator/config"
	activatortest "knative.dev/serving/pkg/activator/testing"
	"knative.dev/serving/pkg/activator/util"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/queue"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"

	. "knative.dev/pkg/configmap/testing"
	"knative.dev/pkg/logging"
	_ "knative.dev/pkg/system/testing"
)

const (
	wantBody      = "♫ everything is awesome! ♫"
	testNamespace = "real-namespace"
	testRevName   = "real-name"
)

type fakeThrottler struct {
	err error
}

func (ft fakeThrottler) Try(ctx context.Context, f func(string) error) error {
	if ft.err != nil {
		return ft.err
	}
	return f("10.10.10.10:1234")
}

func TestActivationHandler(t *testing.T) {
	tests := []struct {
		name      string
		wantBody  string
		wantCode  int
		wantErr   error
		probeErr  error
		probeCode int
		probeResp []string
		throttler Throttler
	}{{
		name:      "active endpoint",
		wantBody:  wantBody,
		wantCode:  http.StatusOK,
		wantErr:   nil,
		throttler: fakeThrottler{},
	}, {
		name:      "request error",
		wantBody:  "request error\n",
		wantCode:  http.StatusBadGateway,
		wantErr:   errors.New("request error"),
		throttler: fakeThrottler{},
	}, {
		name:      "throttler timeout",
		wantBody:  context.DeadlineExceeded.Error() + "\n",
		wantCode:  http.StatusServiceUnavailable,
		wantErr:   nil,
		throttler: fakeThrottler{err: context.DeadlineExceeded},
	}, {
		name:      "overflow",
		wantBody:  "pending request queue full\n",
		wantCode:  http.StatusServiceUnavailable,
		wantErr:   nil,
		throttler: fakeThrottler{err: queue.ErrRequestQueueFull},
	}}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			probeResponses := make([]activatortest.FakeResponse, len(test.probeResp))
			for i := 0; i < len(test.probeResp); i++ {
				probeResponses[i] = activatortest.FakeResponse{
					Err:  test.probeErr,
					Code: test.probeCode,
					Body: test.probeResp[i],
				}
			}
			fakeRT := activatortest.FakeRoundTripper{
				ExpectHost:     "test-host",
				ProbeResponses: probeResponses,
				RequestResponse: &activatortest.FakeResponse{
					Err:  test.wantErr,
					Code: test.wantCode,
					Body: test.wantBody,
				},
			}
			rt := pkgnet.RoundTripperFunc(fakeRT.RT)

			ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
			defer cancel()
			handler := New(ctx, test.throttler, rt)

			resp := httptest.NewRecorder()
			req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
			req.Host = "test-host"

			// Set up config store to populate context.
			configStore := setupConfigStore(t, logging.FromContext(ctx))
			ctx = configStore.ToContext(ctx)
			ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})

			handler.ServeHTTP(resp, req.WithContext(ctx))

			if resp.Code != test.wantCode {
				t.Fatalf("Unexpected response status. Want %d, got %d", test.wantCode, resp.Code)
			}

			gotBody, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				t.Fatal("Error reading body:", err)
			}
			if string(gotBody) != test.wantBody {
				t.Errorf("Unexpected response body. Response body %q, want %q", gotBody, test.wantBody)
			}
		})
	}
}

func TestActivationHandlerProxyHeader(t *testing.T) {
	interceptCh := make(chan *http.Request, 1)
	rt := pkgnet.RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
		interceptCh <- r
		fake := httptest.NewRecorder()
		return fake.Result(), nil
	})

	ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
	defer cancel()

	handler := New(ctx, fakeThrottler{}, rt)

	writer := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)

	// Set up config store to populate context.
	configStore := setupConfigStore(t, logging.FromContext(ctx))
	ctx = configStore.ToContext(req.Context())
	ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})

	handler.ServeHTTP(writer, req.WithContext(ctx))

	select {
	case httpReq := <-interceptCh:
		if got := httpReq.Header.Get(network.ProxyHeaderName); got != activator.Name {
			t.Errorf("Header %q = %q, want:  %q", network.ProxyHeaderName, got, activator.Name)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("Timed out waiting for a request to be intercepted")
	}
}

func TestActivationHandlerTraceSpans(t *testing.T) {
	testcases := []struct {
		name         string
		wantSpans    int
		traceBackend tracingconfig.BackendType
	}{{
		name:         "zipkin trace enabled",
		wantSpans:    3,
		traceBackend: tracingconfig.Zipkin,
	}, {
		name:         "trace disabled",
		wantSpans:    0,
		traceBackend: tracingconfig.None,
	}}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			// Setup transport
			fakeRT := activatortest.FakeRoundTripper{
				RequestResponse: &activatortest.FakeResponse{
					Err:  nil,
					Code: http.StatusOK,
					Body: wantBody,
				},
			}
			rt := pkgnet.RoundTripperFunc(fakeRT.RT)

			// Create tracer with reporter recorder
			reporter, co := tracetesting.FakeZipkinExporter()
			oct := tracing.NewOpenCensusTracer(co)

			cm := &corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name: tracingconfig.ConfigName,
				},
				Data: map[string]string{
					"zipkin-endpoint": "localhost:1234",
					"backend":         string(tc.traceBackend),
					"debug":           "true",
				},
			}
			cfg, err := tracingconfig.NewTracingConfigFromConfigMap(cm)
			if err != nil {
				t.Fatal("Failed to generate config:", err)
			}
			if err := oct.ApplyConfig(cfg); err != nil {
				t.Error("Failed to apply tracer config:", err)
			}

			ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
			defer func() {
				cancel()
				reporter.Close()
				oct.Finish()
			}()

			handler := New(ctx, fakeThrottler{}, rt)

			// Set up config store to populate context.
			configStore := setupConfigStore(t, logging.FromContext(ctx))
			// Update the store with our "new" config explicitly.
			configStore.OnConfigChanged(cm)
			sendRequest(testNamespace, testRevName, handler, configStore)

			gotSpans := reporter.Flush()
			if len(gotSpans) != tc.wantSpans {
				t.Errorf("Got %d spans, expected %d", len(gotSpans), tc.wantSpans)
			}

			spanNames := []string{"throttler_try", "/", "activator_proxy"}
			for i, spanName := range spanNames[0:tc.wantSpans] {
				if gotSpans[i].Name != spanName {
					t.Errorf("Got span %d named %q, expected %q", i, gotSpans[i].Name, spanName)
				}
			}
		})
	}
}

func sendRequest(namespace, revName string, handler http.Handler, store *activatorconfig.Store) *httptest.ResponseRecorder {
	resp := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
	ctx := store.ToContext(req.Context())
	ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: namespace, Name: revName})
	handler.ServeHTTP(resp, req.WithContext(ctx))
	return resp
}

func revision(namespace, name string) *v1.Revision {
	return &v1.Revision{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      name,
			Labels: map[string]string{
				serving.ConfigurationLabelKey: "config-" + name,
				serving.ServiceLabelKey:       "service-" + name,
			},
		},
		Spec: v1.RevisionSpec{
			ContainerConcurrency: ptr.Int64(1),
		},
	}
}

func setupConfigStore(t *testing.T, logger *zap.SugaredLogger) *activatorconfig.Store {
	configStore := activatorconfig.NewStore(logger)
	tracingConfig := ConfigMapFromTestFile(t, tracingconfig.ConfigName)
	configStore.OnConfigChanged(tracingConfig)
	return configStore
}

func BenchmarkHandler(b *testing.B) {
	ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(&testing.T{})
	defer cancel()
	configStore := setupConfigStore(&testing.T{}, logging.FromContext(ctx))

	// bodyLength is in kilobytes.
	for _, bodyLength := range [5]int{2, 16, 32, 64, 128} {
		body := []byte(randomString(1024 * bodyLength))

		rt := pkgnet.RoundTripperFunc(func(*http.Request) (*http.Response, error) {
			return &http.Response{
				Body:       ioutil.NopCloser(bytes.NewReader(body)),
				StatusCode: http.StatusOK,
			}, nil
		})

		handler := New(ctx, fakeThrottler{}, rt)

		request := func() *http.Request {
			req := httptest.NewRequest(http.MethodGet, "http://example.com", nil)
			req.Host = "test-host"

			reqCtx := configStore.ToContext(context.Background())
			reqCtx = util.WithRevID(reqCtx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})
			return req.WithContext(reqCtx)
		}

		test := func(req *http.Request, b *testing.B) {
			resp := &responseRecorder{}
			handler.ServeHTTP(resp, req)
			if resp.code != http.StatusOK {
				b.Fatalf("resp.Code = %d, want: StatusOK(200)", resp.code)
			}
			if got, want := resp.size.Load(), int32(len(body)); got != want {
				b.Fatalf("|body| = %d, want = %d", got, want)
			}
		}

		b.Run(fmt.Sprintf("%03dk-resp-len-sequential", bodyLength), func(b *testing.B) {
			req := request()
			for j := 0; j < b.N; j++ {
				test(req, b)
			}
		})

		b.Run(fmt.Sprintf("%03dk-resp-len-parallel", bodyLength), func(b *testing.B) {
			b.RunParallel(func(pb *testing.PB) {
				req := request()
				for pb.Next() {
					test(req, b)
				}
			})
		})
	}
}

func randomString(n int) string {
	letter := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")

	b := make([]rune, n)
	for i := range b {
		b[i] = letter[rand.Intn(len(letter))]
	}
	return string(b)
}

// responseRecorder is an implementation of http.ResponseWriter and http.Flusher
// that captures the response code and size.
type responseRecorder struct {
	code int
	size atomic.Int32
}

func (rr *responseRecorder) Flush() {}

func (rr *responseRecorder) Header() http.Header {
	return http.Header{}
}

func (rr *responseRecorder) Write(p []byte) (int, error) {
	rr.size.Add(int32(len(p)))
	return ioutil.Discard.Write(p)
}

func (rr *responseRecorder) WriteHeader(code int) {
	rr.code = code
}
<commit_msg>Some test nit cleanup (#9988)<commit_after>/*
Copyright 2018 The Knative Authors

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

    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 handler

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"go.uber.org/atomic"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	network "knative.dev/networking/pkg"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/ptr"
	rtesting "knative.dev/pkg/reconciler/testing"
	"knative.dev/pkg/tracing"
	tracingconfig "knative.dev/pkg/tracing/config"
	tracetesting "knative.dev/pkg/tracing/testing"
	"knative.dev/serving/pkg/activator"
	activatorconfig "knative.dev/serving/pkg/activator/config"
	activatortest "knative.dev/serving/pkg/activator/testing"
	"knative.dev/serving/pkg/activator/util"
	"knative.dev/serving/pkg/apis/serving"
	v1 "knative.dev/serving/pkg/apis/serving/v1"
	"knative.dev/serving/pkg/queue"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"

	. "knative.dev/pkg/configmap/testing"
	"knative.dev/pkg/logging"
	_ "knative.dev/pkg/system/testing"
)

const (
	wantBody      = "♫ everything is awesome! ♫"
	testNamespace = "real-namespace"
	testRevName   = "real-name"
)

type fakeThrottler struct {
	err error
}

func (ft fakeThrottler) Try(ctx context.Context, f func(string) error) error {
	if ft.err != nil {
		return ft.err
	}
	return f("10.10.10.10:1234")
}

func TestActivationHandler(t *testing.T) {
	tests := []struct {
		name      string
		wantBody  string
		wantCode  int
		wantErr   error
		probeErr  error
		probeCode int
		probeResp []string
		throttler Throttler
	}{{
		name:      "active endpoint",
		wantBody:  wantBody,
		wantCode:  http.StatusOK,
		throttler: fakeThrottler{},
	}, {
		name:      "request error",
		wantBody:  "request error\n",
		wantCode:  http.StatusBadGateway,
		wantErr:   errors.New("request error"),
		throttler: fakeThrottler{},
	}, {
		name:      "throttler timeout",
		wantBody:  context.DeadlineExceeded.Error() + "\n",
		wantCode:  http.StatusServiceUnavailable,
		throttler: fakeThrottler{err: context.DeadlineExceeded},
	}, {
		name:      "overflow",
		wantBody:  "pending request queue full\n",
		wantCode:  http.StatusServiceUnavailable,
		throttler: fakeThrottler{err: queue.ErrRequestQueueFull},
	}}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			probeResponses := make([]activatortest.FakeResponse, len(test.probeResp))
			for i := 0; i < len(test.probeResp); i++ {
				probeResponses[i] = activatortest.FakeResponse{
					Err:  test.probeErr,
					Code: test.probeCode,
					Body: test.probeResp[i],
				}
			}
			fakeRT := activatortest.FakeRoundTripper{
				ExpectHost:     "test-host",
				ProbeResponses: probeResponses,
				RequestResponse: &activatortest.FakeResponse{
					Err:  test.wantErr,
					Code: test.wantCode,
					Body: test.wantBody,
				},
			}
			rt := pkgnet.RoundTripperFunc(fakeRT.RT)

			ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
			defer cancel()
			handler := New(ctx, test.throttler, rt)

			resp := httptest.NewRecorder()
			req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
			req.Host = "test-host"

			// Set up config store to populate context.
			configStore := setupConfigStore(t, logging.FromContext(ctx))
			ctx = configStore.ToContext(ctx)
			ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})

			handler.ServeHTTP(resp, req.WithContext(ctx))

			if resp.Code != test.wantCode {
				t.Fatalf("Unexpected response status. Want %d, got %d", test.wantCode, resp.Code)
			}

			gotBody, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				t.Fatal("Error reading body:", err)
			}
			if string(gotBody) != test.wantBody {
				t.Errorf("Response body = %q, want: %q", gotBody, test.wantBody)
			}
		})
	}
}

func TestActivationHandlerProxyHeader(t *testing.T) {
	interceptCh := make(chan *http.Request, 1)
	rt := pkgnet.RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
		interceptCh <- r
		fake := httptest.NewRecorder()
		return fake.Result(), nil
	})

	ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
	defer cancel()

	handler := New(ctx, fakeThrottler{}, rt)

	writer := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)

	// Set up config store to populate context.
	configStore := setupConfigStore(t, logging.FromContext(ctx))
	ctx = configStore.ToContext(req.Context())
	ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})

	handler.ServeHTTP(writer, req.WithContext(ctx))

	select {
	case httpReq := <-interceptCh:
		if got := httpReq.Header.Get(network.ProxyHeaderName); got != activator.Name {
			t.Errorf("Header %q = %q, want: %q", network.ProxyHeaderName, got, activator.Name)
		}
	case <-time.After(1 * time.Second):
		t.Error("Timed out waiting for a request to be intercepted")
	}
}

func TestActivationHandlerTraceSpans(t *testing.T) {
	testcases := []struct {
		name         string
		wantSpans    int
		traceBackend tracingconfig.BackendType
	}{{
		name:         "zipkin trace enabled",
		wantSpans:    3,
		traceBackend: tracingconfig.Zipkin,
	}, {
		name:         "trace disabled",
		traceBackend: tracingconfig.None,
	}}

	spanNames := []string{"throttler_try", "/", "activator_proxy"}
	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			// Setup transport
			fakeRT := activatortest.FakeRoundTripper{
				RequestResponse: &activatortest.FakeResponse{
					Code: http.StatusOK,
					Body: wantBody,
				},
			}
			rt := pkgnet.RoundTripperFunc(fakeRT.RT)

			// Create tracer with reporter recorder
			reporter, co := tracetesting.FakeZipkinExporter()
			oct := tracing.NewOpenCensusTracer(co)

			cm := &corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name: tracingconfig.ConfigName,
				},
				Data: map[string]string{
					"zipkin-endpoint": "localhost:1234",
					"backend":         string(tc.traceBackend),
					"debug":           "true",
				},
			}
			cfg, err := tracingconfig.NewTracingConfigFromConfigMap(cm)
			if err != nil {
				t.Fatal("Failed to generate config:", err)
			}
			if err := oct.ApplyConfig(cfg); err != nil {
				t.Error("Failed to apply tracer config:", err)
			}

			ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(t)
			defer func() {
				cancel()
				reporter.Close()
				oct.Finish()
			}()

			handler := New(ctx, fakeThrottler{}, rt)

			// Set up config store to populate context.
			configStore := setupConfigStore(t, logging.FromContext(ctx))
			// Update the store with our "new" config explicitly.
			configStore.OnConfigChanged(cm)
			sendRequest(testNamespace, testRevName, handler, configStore)

			gotSpans := reporter.Flush()
			if len(gotSpans) != tc.wantSpans {
				t.Errorf("NumSpans = %d, want: %d", len(gotSpans), tc.wantSpans)
			}

			for i, spanName := range spanNames[0:tc.wantSpans] {
				if gotSpans[i].Name != spanName {
					t.Errorf("Span[%d] = %q, expected %q", i, gotSpans[i].Name, spanName)
				}
			}
		})
	}
}

func sendRequest(namespace, revName string, handler http.Handler, store *activatorconfig.Store) *httptest.ResponseRecorder {
	resp := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodPost, "http://example.com", nil)
	ctx := store.ToContext(req.Context())
	ctx = util.WithRevID(ctx, types.NamespacedName{Namespace: namespace, Name: revName})
	handler.ServeHTTP(resp, req.WithContext(ctx))
	return resp
}

func revision(namespace, name string) *v1.Revision {
	return &v1.Revision{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: namespace,
			Name:      name,
			Labels: map[string]string{
				serving.ConfigurationLabelKey: "config-" + name,
				serving.ServiceLabelKey:       "service-" + name,
			},
		},
		Spec: v1.RevisionSpec{
			ContainerConcurrency: ptr.Int64(1),
		},
	}
}

func setupConfigStore(t *testing.T, logger *zap.SugaredLogger) *activatorconfig.Store {
	configStore := activatorconfig.NewStore(logger)
	tracingConfig := ConfigMapFromTestFile(t, tracingconfig.ConfigName)
	configStore.OnConfigChanged(tracingConfig)
	return configStore
}

func BenchmarkHandler(b *testing.B) {
	ctx, cancel, _ := rtesting.SetupFakeContextWithCancel(&testing.T{})
	b.Cleanup(cancel)
	configStore := setupConfigStore(&testing.T{}, logging.FromContext(ctx))

	// bodyLength is in kilobytes.
	for _, bodyLength := range [5]int{2, 16, 32, 64, 128} {
		body := []byte(randomString(1024 * bodyLength))

		rt := pkgnet.RoundTripperFunc(func(*http.Request) (*http.Response, error) {
			return &http.Response{
				Body:       ioutil.NopCloser(bytes.NewReader(body)),
				StatusCode: http.StatusOK,
			}, nil
		})

		handler := New(ctx, fakeThrottler{}, rt)

		request := func() *http.Request {
			req := httptest.NewRequest(http.MethodGet, "http://example.com", nil)
			req.Host = "test-host"

			reqCtx := configStore.ToContext(context.Background())
			reqCtx = util.WithRevID(reqCtx, types.NamespacedName{Namespace: testNamespace, Name: testRevName})
			return req.WithContext(reqCtx)
		}

		test := func(req *http.Request, b *testing.B) {
			resp := &responseRecorder{}
			handler.ServeHTTP(resp, req)
			if resp.code != http.StatusOK {
				b.Fatalf("resp.Code = %d, want: StatusOK(200)", resp.code)
			}
			if got, want := resp.size.Load(), int32(len(body)); got != want {
				b.Fatalf("|body| = %d, want = %d", got, want)
			}
		}

		b.Run(fmt.Sprintf("%03dk-resp-len-sequential", bodyLength), func(b *testing.B) {
			req := request()
			for j := 0; j < b.N; j++ {
				test(req, b)
			}
		})

		b.Run(fmt.Sprintf("%03dk-resp-len-parallel", bodyLength), func(b *testing.B) {
			b.RunParallel(func(pb *testing.PB) {
				req := request()
				for pb.Next() {
					test(req, b)
				}
			})
		})
	}
}

func randomString(n int) string {
	letter := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")

	b := make([]rune, n)
	for i := range b {
		b[i] = letter[rand.Intn(len(letter))]
	}
	return string(b)
}

// responseRecorder is an implementation of http.ResponseWriter and http.Flusher
// that captures the response code and size.
type responseRecorder struct {
	code int
	size atomic.Int32
}

func (rr *responseRecorder) Header() http.Header {
	return http.Header{}
}

func (rr *responseRecorder) Write(p []byte) (int, error) {
	rr.size.Add(int32(len(p)))
	return len(p), nil
}

func (rr *responseRecorder) WriteHeader(code int) {
	rr.code = code
}
<|endoftext|>


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

///////////////////////////////////////////////////////////////////////////
// Important: Run "make update-yaml" to regenerate code after modifying
// this file.
///////////////////////////////////////////////////////////////////////////

// I'm not sure where a better place to put this is, so I'm putting it here:
//
// # API design guidelines
//
// Ambassador's API has inconsistencies because it has historical
// baggage.  Not all of Ambassador's existing API (or even most of
// it!?) follow these guidelines, but new additions to the API should.
// If/when we advance to getambassador.io/v3 and we can break
// compatibility, these are things that we should apply everywhere.
//
// - Prefer `camelCase` to `snake_case`
//   * Exception: Except for consistency with existing fields in the
//     same resource, or symmetry with identical fields in another
//     resource.
//   * Justification: Kubernetes style is to use camelCase. But
//     historically Ambassador used snake_case for everything.
//
// - Prefer for object references to not support namespacing
//   * Exception: If there's a real use-case for it.
//   * Justification: Most native Kubernetes resources don't support
//     referencing things in a different namespace.  We should be
//     opinionated and not support it either, unless there's a good
//     reason to in a specific case.
//
// - Prefer to use `corev1.LocalObjectReference` or
//   `corev1.SecretReference` references instead of
//   `{name}.{namespace}` strings.
//   * Justification: The `{name}.{namespace}` thing evolved "an
//     opaque DNS name" in the `service` field of Mappings, and that
//     was generalized to other things.  Outside of the context of
//     "this is usable as a DNS name to make a request to", it's just
//     confusing and introduces needless ambiguity.  Nothing other
//     than Ambassador uses that notation.
//   * Notes: For things that don't support cross-namespace references
//     (see above), use LocalObjectReference; if you really must
//     support cross-namespace references, then use SecretReference.
//
// - Prefer to use `metav1.Duration` fields instead of "_s" or "_ms"
//   numeric fields.
//
// - Don't have Ambassador populate anything in the `.spec` or
//   `.metadata` of something a user might edit, only let Ambassador
//   set things in the `.status`.
//   * Exception: If Ambassador 100% owns the resource and a user will
//     never edit it.
//   * Notes: I didn't write "Prefer" on this one.  Don't violate it.
//     Just don't do it.  Ever.  Designing the Host resource in
//     violation of this was a HUGE mistake and one that I regret very
//     much.  Learn from my mistakes.
//   * Justification: Having Ambassador-set things in a subresource
//     from user-set things:
//     1. avoids races between the user updating the spec and us
//        updating the status
//     2. allows watt/whatever to only pay attention to
//        .metadata.generation instead of .metadata.resourceVersion;
//        avoiding pointless reconfigures.
//     3. allows the RBAC to be simpler
//     4. avoids the whole class of bugs where we need to make sure
//        that everything round-trips correctly
//     5. provides clarity on which things a user is expected to know
//        how to fill in

package v2

import (
	"encoding/json"
)

// The old `k8s.io/kube-openapi/cmd/openapi-gen` command had ways to
// specify custom schemas for your types (1: define a "OpenAPIDefinition"
// method, or 2: define a "OpenAPIV3Definition" method, or 3: define
// "OpenAPISchemaType" and "OpenAPISchemaFormat" methods).  But the new
// `sigs.k8s.io/controller-tools/controller-gen` command doesn't; it just
// has a small number of "+kubebuilder:" magic comments ("markers") that we
// can use to influence the schema it generates.
//
// So, for example, we'd like to define the AmbassadorID schema as:
//
//    oneOf:
//    - type: "string"
//    - type: "array"
//    items:             # only matters if type=array
//      type: "string"
//
// but if we're going to use just vanilla controller-gen, we're forced to
// be dumb and say `+kubebuilder:validation:Type=""`, to define its schema
// as
//
//    # no `type:` setting because of the +kubebuilder marker
//    items:
//      type: "string"  # because of the raw type
//
// and then kubectl and/or the apiserver won't be able to validate
// AmbassadorID, because it won't be validated until we actually go to
// UnmarshalJSON it when it makes it to Ambassador.  That's pretty much
// what Kubernetes itself[1] does for the JSON Schema types that are unions
// like that.
//
//  > Aside: Some recent work in controller-gen[2] *strongly* suggests that
//  > setting `+kubebuilder:validation:Type=Any` instead of `:Type=""` is
//  > the proper thing to do.  But, um, it doesn't work... kubectl would
//  > say things like:
//  >
//  >    Invalid value: "array": spec.ambassador_id in body must be of type Any: "array"
//
// But honestly that's dumb, and we can do better than that.
//
// So, option one choice would be to send the controller-tools folks a PR
// to support the openapi-gen methods to allow that customization.  That's
// probably the Right Thing, but that seemed like more work than option
// two.  FIXME(lukeshu): Send the controller-tools folks a PR.
//
// Option two: Say something nonsensical like
// `+kubebuilder:validation:Type="d6e-union"`, and teach the `fix-crds`
// script to notice that and delete that nonsensical `type`, replacing it
// with the appropriate `oneOf: [type: A, type: B]` (note that the version
// of JSONSchema that OpenAPI/Kubernetes uses doesn't support type being an
// array).  And so that's what I did.
//
// FIXME(lukeshu): But all of that is still terrible.  Because the very
// structure of our data inherently means that we must have a
// non-structural[3] schema.  With "apiextensions.k8s.io/v1beta1" CRDs,
// non-structural schemas disable several features; and in v1 CRDs,
// non-structural schemas are entirely forbidden.  I mean it doesn't
// _really_ matter right now, because we give out v1beta1 CRDs anyway
// because v1 only became available in Kubernetes 1.16 and we still support
// down to Kubernetes 1.11; but I don't think that we want to lock
// ourselves out from v1 forever.  So I guess that means when it comes time
// for `getambassador.io/v3` (`ambassadorlabs.com/v1`?), we need to
// strictly avoid union types, in order to avoid violating rule 3 of
// structural schemas.  Or hope that the Kubernetes folks decide to relax
// some of the structural-schema rules.
//
// [1]: https://github.com/kubernetes/apiextensions-apiserver/blob/kubernetes-1.18.4/pkg/apis/apiextensions/v1beta1/types_jsonschema.go#L195-L206
// [2]: https://github.com/kubernetes-sigs/controller-tools/pull/427
// [3]: https://kubernetes.io/docs/tasks/extend-kubernetes/custom-resources/custom-resource-definitions/#specifying-a-structural-schema

type CircuitBreaker struct {
	// +kubebuilder:validation:Enum={"default", "high"}
	Priority           string `json:"priority,omitempty"`
	MaxConnections     *int   `json:"max_connections,omitempty"`
	MaxPendingRequests *int   `json:"max_pending_requests,omitempty"`
	MaxRequests        *int   `json:"max_requests,omitempty"`
	MaxRetries         *int   `json:"max_retries,omitempty"`
}

// ErrorResponseTextFormatSource specifies a source for an error response body
type ErrorResponseTextFormatSource struct {
	// The name of a file on the Ambassador pod that contains a format text string.
	Filename string `json:"filename"`
}

// ErorrResponseOverrideBody specifies the body of an error response
type ErrorResponseOverrideBody struct {
	// A format string representing a text response body.
	// Content-Type can be set using the `content_type` field below.
	ErrorResponseTextFormat string `json:"text_format,omitempty"`

	// A JSON response with content-type: application/json. The values can
	// contain format text like in text_format.
	ErrorResponseJsonFormat map[string]string `json:"json_format,omitempty"`

	// A format string sourced from a file on the Ambassador container.
	// Useful for larger response bodies that should not be placed inline
	// in configuration.
	ErrorResponseTextFormatSource *ErrorResponseTextFormatSource `json:"text_format_source,omitempty"`

	// The content type to set on the error response body when
	// using text_format or text_format_source. Defaults to 'text/plain'.
	ContentType string `json:"content_type,omitempty"`
}

// A response rewrite for an HTTP error response
type ErrorResponseOverride struct {
	// The status code to match on -- not a pointer because it's required.
	// +kubebuilder:validation:Required
	// +kubebuilder:validation:Minimum=400
	// +kubebuilder:validation:Maximum=599
	OnStatusCode int `json:"on_status_code,omitempty"`

	// The new response body
	// +kubebuilder:validation:Required
	Body ErrorResponseOverrideBody `json:"body,omitempty"`
}

// AmbassadorID declares which Ambassador instances should pay
// attention to this resource.  May either be a string or a list of
// strings.  If no value is provided, the default is:
//
//    ambassador_id:
//    - "default"
//
// +kubebuilder:validation:Type="d6e-union:string,array"
type AmbassadorID []string

func (aid *AmbassadorID) UnmarshalJSON(data []byte) error {
	return (*StringOrStringList)(aid).UnmarshalJSON(data)
}

// StringOrStringList is just what it says on the tin, but note that it will always
// marshal as a list of strings right now.
// +kubebuilder:validation:Type="d6e-union:string,array"
type StringOrStringList []string

func (sl *StringOrStringList) UnmarshalJSON(data []byte) error {
	if string(data) == "null" {
		*sl = nil
		return nil
	}

	var err error
	var list []string
	var single string

	if err = json.Unmarshal(data, &single); err == nil {
		*sl = StringOrStringList([]string{single})
		return nil
	}

	if err = json.Unmarshal(data, &list); err == nil {
		*sl = StringOrStringList(list)
		return nil
	}

	return err
}

// BoolOrString is a type that can hold a Boolean or a string.
//
// +kubebuilder:validation:Type="d6e-union:string,boolean"
type BoolOrString struct {
	String *string
	Bool   *bool
}

// MarshalJSON is important both so that we generate the proper
// output, and to trigger controller-gen to not try to generate
// jsonschema for our sub-fields:
// https://github.com/kubernetes-sigs/controller-tools/pull/427
func (o BoolOrString) MarshalJSON() ([]byte, error) {
	switch {
	case o.String == nil && o.Bool == nil:
		return json.Marshal(nil)
	case o.String == nil && o.Bool != nil:
		return json.Marshal(o.Bool)
	case o.String != nil && o.Bool == nil:
		return json.Marshal(o.String)
	case o.String != nil && o.Bool != nil:
		panic("invalid BoolOrString")
	}
	panic("not reached")
}

func (o *BoolOrString) UnmarshalJSON(data []byte) error {
	if string(data) == "null" {
		*o = BoolOrString{}
		return nil
	}

	var err error

	var b bool
	if err = json.Unmarshal(data, &b); err == nil {
		*o = BoolOrString{Bool: &b}
		return nil
	}

	var str string
	if err = json.Unmarshal(data, &str); err == nil {
		*o = BoolOrString{String: &str}
		return nil
	}

	return err
}

// UntypedDict is relatively opaque as a Go type, but it preserves its contents in a roundtrippable
// way.
// +kubebuilder:validation:Type="object"
type UntypedDict struct {
	Values map[string]UntypedValue
}

func (u UntypedDict) MarshalJSON() ([]byte, error) {
	return json.Marshal(u.Values)
}

func (u *UntypedDict) UnmarshalJSON(data []byte) error {
	var values map[string]UntypedValue
	err := json.Unmarshal(data, &values)
	if err != nil {
		return err
	}
	*u = UntypedDict{Values: values}
	return nil
}

type UntypedValue struct {
	raw json.RawMessage
}

func (u UntypedValue) MarshalJSON() ([]byte, error) {
	return json.Marshal(u.raw)
}

func (u *UntypedValue) UnmarshalJSON(data []byte) error {
	*u = UntypedValue{raw: json.RawMessage(data)}
	return nil
}
<commit_msg>(from AES) Update API design guidelines.<commit_after>// Copyright 2020 Datawire.  All rights reserved
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///////////////////////////////////////////////////////////////////////////
// Important: Run "make update-yaml" to regenerate code after modifying
// this file.
///////////////////////////////////////////////////////////////////////////

// I'm not sure where a better place to put this is, so I'm putting it here:
//
// # API design guidelines
//
// Ambassador's API has inconsistencies because it has historical
// baggage.  Not all of Ambassador's existing API (or even most of
// it!?) follow these guidelines, but new additions to the API should.
// If/when we advance to getambassador.io/v3 and we can break
// compatibility, these are things that we should apply everywhere.
//
// - Prefer `camelCase` to `snake_case`
//   * Exception: Except for consistency with existing fields in the
//     same resource, or symmetry with identical fields in another
//     resource.
//   * Justification: Kubernetes style is to use camelCase. But
//     historically Ambassador used snake_case for everything.
//
// - Give _every_ field a `json:""` struct tag.
//   * Justification: Marshaling and unmarshaling are key to what we
//     do, and it's critical to carefully define how it happens.
//   * Notes: This is not optional. Do it for _every field_. (It's OK
//     if the tag is literally `json:""` for fields that must never be
//     exposed during marshaling.)
//
// - Prefer `*int`, `*bool`, and `*BoolOrString`, rather than just
//   `int`, `bool`, and `BoolOrString`.
//   * Justification: The Ambassador API is rooted in Python, where
//     it is always possible to tell if a given element was present in
//     in a CRD, or left unset. This is at odds with Go's `omitempty`
//     specifier, which really means "omit if empty _or if set to the
//     default value". For int in particular, this results in a value
//     of 0 being omitted, and for many Ambassador fields, 0 is not
//     the correct default value.
//
//     This resulted in a lot of bugs in the 1.10 timeframe, so be
//     careful going forward.
//
// - Prefer for object references to not support namespacing
//   * Exception: If there's a real use-case for it.
//   * Justification: Most native Kubernetes resources don't support
//     referencing things in a different namespace.  We should be
//     opinionated and not support it either, unless there's a good
//     reason to in a specific case.
//
// - Prefer to use `corev1.LocalObjectReference` or
//   `corev1.SecretReference` references instead of
//   `{name}.{namespace}` strings.
//   * Justification: The `{name}.{namespace}` thing evolved "an
//     opaque DNS name" in the `service` field of Mappings, and that
//     was generalized to other things.  Outside of the context of
//     "this is usable as a DNS name to make a request to", it's just
//     confusing and introduces needless ambiguity.  Nothing other
//     than Ambassador uses that notation.
//   * Notes: For things that don't support cross-namespace references
//     (see above), use LocalObjectReference; if you really must
//     support cross-namespace references, then use SecretReference.
//
// - Prefer to use `metav1.Duration` fields instead of "_s" or "_ms"
//   numeric fields.
//
// - Don't have Ambassador populate anything in the `.spec` or
//   `.metadata` of something a user might edit, only let Ambassador
//   set things in the `.status`.
//   * Exception: If Ambassador 100% owns the resource and a user will
//     never edit it.
//   * Notes: I didn't write "Prefer" on this one.  Don't violate it.
//     Just don't do it.  Ever.  Designing the Host resource in
//     violation of this was a HUGE mistake and one that I regret very
//     much.  Learn from my mistakes.
//   * Justification: Having Ambassador-set things in a subresource
//     from user-set things:
//     1. avoids races between the user updating the spec and us
//        updating the status
//     2. allows watt/whatever to only pay attention to
//        .metadata.generation instead of .metadata.resourceVersion;
//        avoiding pointless reconfigures.
//     3. allows the RBAC to be simpler
//     4. avoids the whole class of bugs where we need to make sure
//        that everything round-trips correctly
//     5. provides clarity on which things a user is expected to know
//        how to fill in

package v2

import (
	"encoding/json"
)

// The old `k8s.io/kube-openapi/cmd/openapi-gen` command had ways to
// specify custom schemas for your types (1: define a "OpenAPIDefinition"
// method, or 2: define a "OpenAPIV3Definition" method, or 3: define
// "OpenAPISchemaType" and "OpenAPISchemaFormat" methods).  But the new
// `sigs.k8s.io/controller-tools/controller-gen` command doesn't; it just
// has a small number of "+kubebuilder:" magic comments ("markers") that we
// can use to influence the schema it generates.
//
// So, for example, we'd like to define the AmbassadorID schema as:
//
//    oneOf:
//    - type: "string"
//    - type: "array"
//    items:             # only matters if type=array
//      type: "string"
//
// but if we're going to use just vanilla controller-gen, we're forced to
// be dumb and say `+kubebuilder:validation:Type=""`, to define its schema
// as
//
//    # no `type:` setting because of the +kubebuilder marker
//    items:
//      type: "string"  # because of the raw type
//
// and then kubectl and/or the apiserver won't be able to validate
// AmbassadorID, because it won't be validated until we actually go to
// UnmarshalJSON it when it makes it to Ambassador.  That's pretty much
// what Kubernetes itself[1] does for the JSON Schema types that are unions
// like that.
//
//  > Aside: Some recent work in controller-gen[2] *strongly* suggests that
//  > setting `+kubebuilder:validation:Type=Any` instead of `:Type=""` is
//  > the proper thing to do.  But, um, it doesn't work... kubectl would
//  > say things like:
//  >
//  >    Invalid value: "array": spec.ambassador_id in body must be of type Any: "array"
//
// But honestly that's dumb, and we can do better than that.
//
// So, option one choice would be to send the controller-tools folks a PR
// to support the openapi-gen methods to allow that customization.  That's
// probably the Right Thing, but that seemed like more work than option
// two.  FIXME(lukeshu): Send the controller-tools folks a PR.
//
// Option two: Say something nonsensical like
// `+kubebuilder:validation:Type="d6e-union"`, and teach the `fix-crds`
// script to notice that and delete that nonsensical `type`, replacing it
// with the appropriate `oneOf: [type: A, type: B]` (note that the version
// of JSONSchema that OpenAPI/Kubernetes uses doesn't support type being an
// array).  And so that's what I did.
//
// FIXME(lukeshu): But all of that is still terrible.  Because the very
// structure of our data inherently means that we must have a
// non-structural[3] schema.  With "apiextensions.k8s.io/v1beta1" CRDs,
// non-structural schemas disable several features; and in v1 CRDs,
// non-structural schemas are entirely forbidden.  I mean it doesn't
// _really_ matter right now, because we give out v1beta1 CRDs anyway
// because v1 only became available in Kubernetes 1.16 and we still support
// down to Kubernetes 1.11; but I don't think that we want to lock
// ourselves out from v1 forever.  So I guess that means when it comes time
// for `getambassador.io/v3` (`ambassadorlabs.com/v1`?), we need to
// strictly avoid union types, in order to avoid violating rule 3 of
// structural schemas.  Or hope that the Kubernetes folks decide to relax
// some of the structural-schema rules.
//
// [1]: https://github.com/kubernetes/apiextensions-apiserver/blob/kubernetes-1.18.4/pkg/apis/apiextensions/v1beta1/types_jsonschema.go#L195-L206
// [2]: https://github.com/kubernetes-sigs/controller-tools/pull/427
// [3]: https://kubernetes.io/docs/tasks/extend-kubernetes/custom-resources/custom-resource-definitions/#specifying-a-structural-schema

type CircuitBreaker struct {
	// +kubebuilder:validation:Enum={"default", "high"}
	Priority           string `json:"priority,omitempty"`
	MaxConnections     *int   `json:"max_connections,omitempty"`
	MaxPendingRequests *int   `json:"max_pending_requests,omitempty"`
	MaxRequests        *int   `json:"max_requests,omitempty"`
	MaxRetries         *int   `json:"max_retries,omitempty"`
}

// ErrorResponseTextFormatSource specifies a source for an error response body
type ErrorResponseTextFormatSource struct {
	// The name of a file on the Ambassador pod that contains a format text string.
	Filename string `json:"filename"`
}

// ErorrResponseOverrideBody specifies the body of an error response
type ErrorResponseOverrideBody struct {
	// A format string representing a text response body.
	// Content-Type can be set using the `content_type` field below.
	ErrorResponseTextFormat string `json:"text_format,omitempty"`

	// A JSON response with content-type: application/json. The values can
	// contain format text like in text_format.
	ErrorResponseJsonFormat map[string]string `json:"json_format,omitempty"`

	// A format string sourced from a file on the Ambassador container.
	// Useful for larger response bodies that should not be placed inline
	// in configuration.
	ErrorResponseTextFormatSource *ErrorResponseTextFormatSource `json:"text_format_source,omitempty"`

	// The content type to set on the error response body when
	// using text_format or text_format_source. Defaults to 'text/plain'.
	ContentType string `json:"content_type,omitempty"`
}

// A response rewrite for an HTTP error response
type ErrorResponseOverride struct {
	// The status code to match on -- not a pointer because it's required.
	// +kubebuilder:validation:Required
	// +kubebuilder:validation:Minimum=400
	// +kubebuilder:validation:Maximum=599
	OnStatusCode int `json:"on_status_code,omitempty"`

	// The new response body
	// +kubebuilder:validation:Required
	Body ErrorResponseOverrideBody `json:"body,omitempty"`
}

// AmbassadorID declares which Ambassador instances should pay
// attention to this resource.  May either be a string or a list of
// strings.  If no value is provided, the default is:
//
//    ambassador_id:
//    - "default"
//
// +kubebuilder:validation:Type="d6e-union:string,array"
type AmbassadorID []string

func (aid *AmbassadorID) UnmarshalJSON(data []byte) error {
	return (*StringOrStringList)(aid).UnmarshalJSON(data)
}

// StringOrStringList is just what it says on the tin, but note that it will always
// marshal as a list of strings right now.
// +kubebuilder:validation:Type="d6e-union:string,array"
type StringOrStringList []string

func (sl *StringOrStringList) UnmarshalJSON(data []byte) error {
	if string(data) == "null" {
		*sl = nil
		return nil
	}

	var err error
	var list []string
	var single string

	if err = json.Unmarshal(data, &single); err == nil {
		*sl = StringOrStringList([]string{single})
		return nil
	}

	if err = json.Unmarshal(data, &list); err == nil {
		*sl = StringOrStringList(list)
		return nil
	}

	return err
}

// BoolOrString is a type that can hold a Boolean or a string.
//
// +kubebuilder:validation:Type="d6e-union:string,boolean"
type BoolOrString struct {
	String *string
	Bool   *bool
}

// MarshalJSON is important both so that we generate the proper
// output, and to trigger controller-gen to not try to generate
// jsonschema for our sub-fields:
// https://github.com/kubernetes-sigs/controller-tools/pull/427
func (o BoolOrString) MarshalJSON() ([]byte, error) {
	switch {
	case o.String == nil && o.Bool == nil:
		return json.Marshal(nil)
	case o.String == nil && o.Bool != nil:
		return json.Marshal(o.Bool)
	case o.String != nil && o.Bool == nil:
		return json.Marshal(o.String)
	case o.String != nil && o.Bool != nil:
		panic("invalid BoolOrString")
	}
	panic("not reached")
}

func (o *BoolOrString) UnmarshalJSON(data []byte) error {
	if string(data) == "null" {
		*o = BoolOrString{}
		return nil
	}

	var err error

	var b bool
	if err = json.Unmarshal(data, &b); err == nil {
		*o = BoolOrString{Bool: &b}
		return nil
	}

	var str string
	if err = json.Unmarshal(data, &str); err == nil {
		*o = BoolOrString{String: &str}
		return nil
	}

	return err
}

// UntypedDict is relatively opaque as a Go type, but it preserves its contents in a roundtrippable
// way.
// +kubebuilder:validation:Type="object"
type UntypedDict struct {
	Values map[string]UntypedValue
}

func (u UntypedDict) MarshalJSON() ([]byte, error) {
	return json.Marshal(u.Values)
}

func (u *UntypedDict) UnmarshalJSON(data []byte) error {
	var values map[string]UntypedValue
	err := json.Unmarshal(data, &values)
	if err != nil {
		return err
	}
	*u = UntypedDict{Values: values}
	return nil
}

type UntypedValue struct {
	raw json.RawMessage
}

func (u UntypedValue) MarshalJSON() ([]byte, error) {
	return json.Marshal(u.raw)
}

func (u *UntypedValue) UnmarshalJSON(data []byte) error {
	*u = UntypedValue{raw: json.RawMessage(data)}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package gce_cloud

import (
	"encoding/json"
	"net/http"
	"time"

	"code.google.com/p/google-api-go-client/googleapi"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

type altTokenSource struct {
	oauthClient *http.Client
	tokenURL    string
}

func (a *altTokenSource) Token() (*oauth2.Token, error) {
	req, err := http.NewRequest("GET", a.tokenURL, nil)
	if err != nil {
		return nil, err
	}
	res, err := a.oauthClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if err := googleapi.CheckResponse(res); err != nil {
		return nil, err
	}
	var tok struct {
		AccessToken       string `json:"accessToken"`
		ExpiryTimeSeconds int64  `json:"expiryTimeSeconds,string"`
	}
	if err := json.NewDecoder(res.Body).Decode(&tok); err != nil {
		return nil, err
	}
	return &oauth2.Token{
		AccessToken: tok.AccessToken,
		Expiry:      time.Unix(tok.ExpiryTimeSeconds, 0),
	}, nil
}

func newAltTokenSource(tokenURL string) oauth2.TokenSource {
	client := oauth2.NewClient(oauth2.NoContext, google.ComputeTokenSource(""))
	a := &altTokenSource{
		oauthClient: client,
		tokenURL:    tokenURL,
	}
	return oauth2.ReuseTokenSource(nil, a)
}
<commit_msg>Add a RateLimiter for the gce altTokenSource.<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package gce_cloud

import (
	"encoding/json"
	"net/http"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"

	"code.google.com/p/google-api-go-client/googleapi"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

const (
	// Max QPS to allow through to the token URL.
	tokenURLQPS = 1
	// Maximum burst of requests to token URL before limiting.
	tokenURLBurst = 3
)

type altTokenSource struct {
	oauthClient *http.Client
	tokenURL    string
	throttle    util.RateLimiter
}

func (a *altTokenSource) Token() (*oauth2.Token, error) {
	a.throttle.Accept()
	req, err := http.NewRequest("GET", a.tokenURL, nil)
	if err != nil {
		return nil, err
	}
	res, err := a.oauthClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if err := googleapi.CheckResponse(res); err != nil {
		return nil, err
	}
	var tok struct {
		AccessToken       string `json:"accessToken"`
		ExpiryTimeSeconds int64  `json:"expiryTimeSeconds,string"`
	}
	if err := json.NewDecoder(res.Body).Decode(&tok); err != nil {
		return nil, err
	}
	return &oauth2.Token{
		AccessToken: tok.AccessToken,
		Expiry:      time.Unix(tok.ExpiryTimeSeconds, 0),
	}, nil
}

func newAltTokenSource(tokenURL string) oauth2.TokenSource {
	client := oauth2.NewClient(oauth2.NoContext, google.ComputeTokenSource(""))
	a := &altTokenSource{
		oauthClient: client,
		tokenURL:    tokenURL,
		throttle:    util.NewTokenBucketRateLimiter(tokenURLQPS, tokenURLBurst),
	}
	return oauth2.ReuseTokenSource(nil, a)
}
<|endoftext|>


<commit_before>package textsearch // import "a4.io/blobstash/pkg/textsearch"

import (
	"bytes"
	"fmt"
	"strings"
	"text/scanner"

	"github.com/blevesearch/segment"
	lru "github.com/hashicorp/golang-lru"
	porterstemmer "github.com/reiver/go-porterstemmer"
)

var cache, _ = lru.New(128)

type searchTerm struct {
	prefix     string
	term       string
	exactMatch bool
}

type SearchTerms []*searchTerm

type IndexedDoc struct {
	Content string         `msgpack:"c"`
	Stems   map[string]int `msgpack:"s"`
}

func NewIndexedDoc(doc map[string]interface{}, fields []string) (*IndexedDoc, error) {
	parts := []string{}
	stems := map[string]int{}
	for _, field := range fields {
		if dat, ok := doc[field]; ok {
			parts = append(parts, dat.(string))
			segmenter := segment.NewWordSegmenter(bytes.NewReader([]byte(dat.(string))))
			for segmenter.Segment() {
				if segmenter.Type() == segment.Letter {
					stem := porterstemmer.StemString(segmenter.Text())
					if _, ok := stems[stem]; ok {
						stems[stem] += 1
					} else {
						stems[stem] = 1
					}
				}
			}
			if err := segmenter.Err(); err != nil {
				return nil, err
			}
		}
	}
	content := strings.Join(parts, " ")

	return &IndexedDoc{Content: content, Stems: stems}, nil
}

func (terms SearchTerms) Match(d *IndexedDoc) bool {
	match := false

	for _, st := range terms {
		cond := false
		switch {
		case st.exactMatch:
			cond = strings.Contains(d.Content, st.term)
		default:
			_, cond = d.Stems[st.term]
		}

		if st.prefix == "+" {
			if !cond {
				return false
			}
		} else if st.prefix == "-" {
			if cond {
				return false
			} else {
				match = true
			}
		}

		if cond {
			match = true
		}
	}

	return match
}

func ParseTextQuery(q string) SearchTerms {
	if cached, ok := cache.Get(q); ok {
		fmt.Printf("ParseTextQuery form cache")
		return cached.(SearchTerms)
	}
	var s scanner.Scanner
	s.Init(strings.NewReader(q))
	out := SearchTerms{}
	var prefix, term string
	var exactMatch bool
	for tok := s.Scan(); tok != scanner.EOF; tok = s.Scan() {
		term = s.TokenText()

		if term == "+" || term == "-" {
			prefix = term
			continue
		}

		if strings.HasPrefix(term, "\"") && strings.HasSuffix(term, "\"") {
			exactMatch = true
			term = term[1 : len(term)-1]
		}

		if !exactMatch {
			term = porterstemmer.StemString(term)
		}

		out = append(out, &searchTerm{
			prefix:     prefix,
			term:       term,
			exactMatch: exactMatch,
		})

		prefix = ""
		exactMatch = false
	}
	cache.Add(q, out)
	return out
}
<commit_msg>docstore/textsearch: update docs<commit_after>/*
Package textsearch implements basic text search features (for matching text fields of JSON documents).
*/
package textsearch // import "a4.io/blobstash/pkg/textsearch"

import (
	"bytes"
	"fmt"
	"strings"
	"text/scanner"

	"github.com/blevesearch/segment"
	lru "github.com/hashicorp/golang-lru"
	porterstemmer "github.com/reiver/go-porterstemmer"
)

// cache for `string` -> `SearchTerms`
var searchTermsCache, _ = lru.New(128)

// searchTerm holds a single search term
type searchTerm struct {
	prefix     string // `+` (for required match) or `-` (for excluding doc matching the term)
	term       string
	exactMatch bool // true if the search term was quoted for exact match
}

// SearchTerms holds a parsed text search query
type SearchTerms []*searchTerm

// IndexedDoc holds a parsed "document"
type IndexedDoc struct {
	Content string         `msgpack:"c"`
	Stems   map[string]int `msgpack:"s"`
}

// NewIndexedDoc returns a parsed "document"
func NewIndexedDoc(doc map[string]interface{}, fields []string) (*IndexedDoc, error) {
	parts := []string{}
	stems := map[string]int{}
	for _, field := range fields {
		if dat, ok := doc[field]; ok {
			parts = append(parts, dat.(string))
			segmenter := segment.NewWordSegmenter(bytes.NewReader([]byte(dat.(string))))
			for segmenter.Segment() {
				if segmenter.Type() == segment.Letter {
					stem := porterstemmer.StemString(segmenter.Text())
					if _, ok := stems[stem]; ok {
						stems[stem] += 1
					} else {
						stems[stem] = 1
					}
				}
			}
			if err := segmenter.Err(); err != nil {
				return nil, err
			}
		}
	}
	content := strings.Join(parts, " ")

	return &IndexedDoc{Content: content, Stems: stems}, nil
}

// ParseTextQuery returns a parsed text query
func ParseTextQuery(q string) SearchTerms {
	if cached, ok := cache.Get(q); ok {
		fmt.Printf("ParseTextQuery form cache")
		return cached.(SearchTerms)
	}
	var s scanner.Scanner
	s.Init(strings.NewReader(q))
	out := SearchTerms{}
	var prefix, term string
	var exactMatch bool
	for tok := s.Scan(); tok != scanner.EOF; tok = s.Scan() {
		term = s.TokenText()

		if term == "+" || term == "-" {
			prefix = term
			continue
		}

		if strings.HasPrefix(term, "\"") && strings.HasSuffix(term, "\"") {
			exactMatch = true
			term = term[1 : len(term)-1]
		}

		if !exactMatch {
			term = porterstemmer.StemString(term)
		}

		out = append(out, &searchTerm{
			prefix:     prefix,
			term:       term,
			exactMatch: exactMatch,
		})

		prefix = ""
		exactMatch = false
	}
	cache.Add(q, out)
	return out
}

// Match returns true if the query matches the given `IndexedDoc`
func (terms SearchTerms) Match(d *IndexedDoc) bool {
	match := false

	for _, st := range terms {
		cond := false
		switch {
		case st.exactMatch:
			cond = strings.Contains(d.Content, st.term)
		default:
			_, cond = d.Stems[st.term]
		}

		if st.prefix == "+" {
			if !cond {
				return false
			}
		} else if st.prefix == "-" {
			if cond {
				return false
			} else {
				match = true
			}
		}

		if cond {
			match = true
		}
	}

	return match
}
<|endoftext|>


<commit_before>// Copyright 2016-2018 Authors of Cilium
//
// 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.

// +build !privileged_tests

package cache

import (
	"sync"

	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/labels"

	. "gopkg.in/check.v1"
)

func (s *IdentityCacheTestSuite) TestAllocateIdentityReserved(c *C) {
	var (
		lbls  labels.Labels
		i     *identity.Identity
		isNew bool
		err   error
	)

	lbls = labels.Labels{
		labels.IDNameHost: labels.NewLabel(labels.IDNameHost, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityHost)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameWorld: labels.NewLabel(labels.IDNameWorld, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityWorld)
	c.Assert(isNew, Equals, false)

	c.Assert(IdentityAllocationIsLocal(labels.LabelHealth), Equals, true)
	i, isNew, err = AllocateIdentity(labels.LabelHealth)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityHealth)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameInit: labels.NewLabel(labels.IDNameInit, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityInit)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameUnmanaged: labels.NewLabel(labels.IDNameUnmanaged, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityUnmanaged)
	c.Assert(isNew, Equals, false)
}

type IdentityAllocatorSuite struct{}

type IdentityAllocatorEtcdSuite struct {
	IdentityAllocatorSuite
}

var _ = Suite(&IdentityAllocatorEtcdSuite{})

func (e *IdentityAllocatorEtcdSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("etcd")
}

type IdentityAllocatorConsulSuite struct {
	IdentityAllocatorSuite
}

var _ = Suite(&IdentityAllocatorConsulSuite{})

func (e *IdentityAllocatorConsulSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("consul")
}

type dummyOwner struct{}

func (d dummyOwner) TriggerPolicyUpdates(force bool, reason string) *sync.WaitGroup {
	return nil
}

func (d dummyOwner) GetNodeSuffix() string {
	return "foo"
}

func (ias *IdentityAllocatorSuite) TestGetIdentityCache(c *C) {
	InitIdentityAllocator(dummyOwner{})
	defer Close()
	defer IdentityAllocator.DeleteAllKeys()

	cache := GetIdentityCache()
	_, ok := cache[identity.ReservedCiliumKVStore]
	c.Assert(ok, Equals, true)
}

func (ias *IdentityAllocatorSuite) TestAllocator(c *C) {
	lbls1 := labels.NewLabelsFromSortedList("id=foo;user=anna;blah=%%//!!")
	lbls2 := labels.NewLabelsFromSortedList("id=bar;user=anna")
	lbls3 := labels.NewLabelsFromSortedList("id=bar;user=susan")

	InitIdentityAllocator(dummyOwner{})
	defer Close()
	defer IdentityAllocator.DeleteAllKeys()

	id1a, isNew, err := AllocateIdentity(lbls1)
	c.Assert(id1a, Not(IsNil))
	c.Assert(err, IsNil)
	c.Assert(isNew, Equals, true)

	id1b, isNew, err := AllocateIdentity(lbls1)
	c.Assert(id1b, Not(IsNil))
	c.Assert(isNew, Equals, false)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Equals, id1b.ID)

	err = Release(id1a)
	c.Assert(err, IsNil)
	err = Release(id1b)
	c.Assert(err, IsNil)

	id1b, isNew, err = AllocateIdentity(lbls1)
	c.Assert(id1b, Not(IsNil))
	c.Assert(err, IsNil)
	// the value key should not have been removed so the same ID should be
	// assigned again the it should not be marked as new
	c.Assert(isNew, Equals, false)
	c.Assert(id1a.ID, Equals, id1b.ID)

	identity := LookupIdentityByID(id1b.ID)
	c.Assert(identity, Not(IsNil))
	c.Assert(lbls1, DeepEquals, identity.Labels)

	id2, isNew, err := AllocateIdentity(lbls2)
	c.Assert(id2, Not(IsNil))
	c.Assert(isNew, Equals, true)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Not(Equals), id2.ID)

	id3, isNew, err := AllocateIdentity(lbls3)
	c.Assert(id3, Not(IsNil))
	c.Assert(isNew, Equals, true)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Not(Equals), id3.ID)
	c.Assert(id2.ID, Not(Equals), id3.ID)

	err = Release(id1b)
	c.Assert(err, IsNil)
	err = Release(id2)
	c.Assert(err, IsNil)
	err = Release(id3)
	c.Assert(err, IsNil)
}
<commit_msg>identity/cache: Use checker.DeepEquals in unit tests<commit_after>// Copyright 2016-2018 Authors of Cilium
//
// 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.

// +build !privileged_tests

package cache

import (
	"sync"

	"github.com/cilium/cilium/pkg/checker"
	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/labels"

	. "gopkg.in/check.v1"
)

func (s *IdentityCacheTestSuite) TestAllocateIdentityReserved(c *C) {
	var (
		lbls  labels.Labels
		i     *identity.Identity
		isNew bool
		err   error
	)

	lbls = labels.Labels{
		labels.IDNameHost: labels.NewLabel(labels.IDNameHost, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityHost)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameWorld: labels.NewLabel(labels.IDNameWorld, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityWorld)
	c.Assert(isNew, Equals, false)

	c.Assert(IdentityAllocationIsLocal(labels.LabelHealth), Equals, true)
	i, isNew, err = AllocateIdentity(labels.LabelHealth)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityHealth)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameInit: labels.NewLabel(labels.IDNameInit, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityInit)
	c.Assert(isNew, Equals, false)

	lbls = labels.Labels{
		labels.IDNameUnmanaged: labels.NewLabel(labels.IDNameUnmanaged, "", labels.LabelSourceReserved),
	}
	c.Assert(IdentityAllocationIsLocal(lbls), Equals, true)
	i, isNew, err = AllocateIdentity(lbls)
	c.Assert(err, IsNil)
	c.Assert(i.ID, Equals, identity.ReservedIdentityUnmanaged)
	c.Assert(isNew, Equals, false)
}

type IdentityAllocatorSuite struct{}

type IdentityAllocatorEtcdSuite struct {
	IdentityAllocatorSuite
}

var _ = Suite(&IdentityAllocatorEtcdSuite{})

func (e *IdentityAllocatorEtcdSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("etcd")
}

type IdentityAllocatorConsulSuite struct {
	IdentityAllocatorSuite
}

var _ = Suite(&IdentityAllocatorConsulSuite{})

func (e *IdentityAllocatorConsulSuite) SetUpTest(c *C) {
	kvstore.SetupDummy("consul")
}

type dummyOwner struct{}

func (d dummyOwner) TriggerPolicyUpdates(force bool, reason string) *sync.WaitGroup {
	return nil
}

func (d dummyOwner) GetNodeSuffix() string {
	return "foo"
}

func (ias *IdentityAllocatorSuite) TestGetIdentityCache(c *C) {
	InitIdentityAllocator(dummyOwner{})
	defer Close()
	defer IdentityAllocator.DeleteAllKeys()

	cache := GetIdentityCache()
	_, ok := cache[identity.ReservedCiliumKVStore]
	c.Assert(ok, Equals, true)
}

func (ias *IdentityAllocatorSuite) TestAllocator(c *C) {
	lbls1 := labels.NewLabelsFromSortedList("id=foo;user=anna;blah=%%//!!")
	lbls2 := labels.NewLabelsFromSortedList("id=bar;user=anna")
	lbls3 := labels.NewLabelsFromSortedList("id=bar;user=susan")

	InitIdentityAllocator(dummyOwner{})
	defer Close()
	defer IdentityAllocator.DeleteAllKeys()

	id1a, isNew, err := AllocateIdentity(lbls1)
	c.Assert(id1a, Not(IsNil))
	c.Assert(err, IsNil)
	c.Assert(isNew, Equals, true)

	id1b, isNew, err := AllocateIdentity(lbls1)
	c.Assert(id1b, Not(IsNil))
	c.Assert(isNew, Equals, false)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Equals, id1b.ID)

	err = Release(id1a)
	c.Assert(err, IsNil)
	err = Release(id1b)
	c.Assert(err, IsNil)

	id1b, isNew, err = AllocateIdentity(lbls1)
	c.Assert(id1b, Not(IsNil))
	c.Assert(err, IsNil)
	// the value key should not have been removed so the same ID should be
	// assigned again the it should not be marked as new
	c.Assert(isNew, Equals, false)
	c.Assert(id1a.ID, Equals, id1b.ID)

	identity := LookupIdentityByID(id1b.ID)
	c.Assert(identity, Not(IsNil))
	c.Assert(lbls1, checker.DeepEquals, identity.Labels)

	id2, isNew, err := AllocateIdentity(lbls2)
	c.Assert(id2, Not(IsNil))
	c.Assert(isNew, Equals, true)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Not(Equals), id2.ID)

	id3, isNew, err := AllocateIdentity(lbls3)
	c.Assert(id3, Not(IsNil))
	c.Assert(isNew, Equals, true)
	c.Assert(err, IsNil)
	c.Assert(id1a.ID, Not(Equals), id3.ID)
	c.Assert(id2.ID, Not(Equals), id3.ID)

	err = Release(id1b)
	c.Assert(err, IsNil)
	err = Release(id2)
	c.Assert(err, IsNil)
	err = Release(id3)
	c.Assert(err, IsNil)
}
<|endoftext|>