<commit_before>/*
Package pipe provides filters that can be chained together in a manner
similar to Unix pipelines.

Each filter is a function that takes as input a sequence of
strings (read from a channel) and produces as output a sequence of
strings (written to a channel).

Filters can be chained together (e.g., via the Run function), the
output of one filter is fed as input to the next filter.  The empty
input is passed to the first filter. The following sequence will
print two lines to standard output:

	err := Run(
		Echo("hello", "world"),
		Reverse(),
		WriteLines(os.Stdout),
	)

An application can implement its own filters easily. For example,
Repeat(n) returns a filter that repeats every input n times:

	func Repeat(n int) Filter {
		return func(arg Arg) {
			for s := range arg.In {
				for i := 0; i < n; i++ {
					arg.Out <- s
				}
			}
		}
	}

	Run(
		Echo("hello"),
		Repeat(10),
	)
*/
package pipe

import (
	"fmt"
	"regexp"
	"sync"
)

// filterErrors records errors accumulated during the execution of a filter.
type filterErrors struct {
	mu     sync.Mutex
	errors []error
}

// Arg contains the data passed to a Filter. Arg.In is a channel that
// produces the input to the filter, and Arg.Out is a channel that
// receives the output from the filter.
type Arg struct {
	In     <-chan string
	Out    chan<- string
	errors *filterErrors
}

// ReportError records an error encountered during an execution of a filter.
// This error will be reported by whatever facility (e.g., ForEach or Run)
// was being used to execute the filters.
func (a *Arg) ReportError(err error) {
	a.errors.mu.Lock()
	defer a.errors.mu.Unlock()
	a.errors.errors = append(a.errors.errors, err)
}

// Filter reads a sequence of strings from a channel and produces a
// sequence on another channel.
type Filter func(Arg)

// Sequence returns a filter that is the concatenation of all filter arguments.
// The output of a filter is fed as input to the next filter.
func Sequence(filters ...Filter) Filter {
	return func(arg Arg) {
		in := arg.In
		for _, f := range filters {
			c := make(chan string, 10000)
			go runAndClose(f, Arg{in, c, arg.errors})
			in = c
		}
		passThrough(Arg{in, arg.Out, arg.errors})
	}
}

// Run executes the sequence of filters and discards all output.
// It returns either nil, an error if any filter reported an error.
func Run(filters ...Filter) error {
	return ForEach(Sequence(filters...), func(s string) {})
}

// ForEach calls fn(s) for every item s in the output of filter and
// returns either nil, or any error reported by the execution of the filter.
func ForEach(filter Filter, fn func(s string)) error {
	in := make(chan string, 0)
	close(in)
	out := make(chan string, 10000)
	e := &filterErrors{}
	go runAndClose(filter, Arg{in, out, e})
	for s := range out {
		fn(s)
	}

	e.mu.Lock()
	defer e.mu.Unlock()
	switch len(e.errors) {
	case 0:
		return nil
	case 1:
		return e.errors[0]
	default:
		return fmt.Errorf("Filter errors: %s", e.errors)
	}
}

func runAndClose(f Filter, arg Arg) {
	f(arg)
	close(arg.Out)
}

// passThrough copies all items read from in to out.
func passThrough(arg Arg) {
	for s := range arg.In {
		arg.Out <- s
	}
}

// Echo emits items.
// Any input items are copied verbatim to the output before items are emitted.
func Echo(items ...string) Filter {
	return func(arg Arg) {
		passThrough(arg)
		for _, s := range items {
			arg.Out <- s
		}
	}
}

// Numbers copies its input and then emits the integers x..y
func Numbers(x, y int) Filter {
	return func(arg Arg) {
		passThrough(arg)
		for i := x; i <= y; i++ {
			arg.Out <- fmt.Sprintf("%d", i)
		}
	}
}

// Map calls fn(x) for every item x and yields the outputs of the fn calls.
func Map(fn func(string) string) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			arg.Out <- fn(s)
		}
	}
}

// If emits every input x for which fn(x) is true.
func If(fn func(string) bool) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			if fn(s) {
				arg.Out <- s
			}
		}
	}
}

func errorFilter(err error) Filter {
	return func(arg Arg) {
		arg.ReportError(err)
		for _ = range arg.In {
			// Drop the input
		}
	}
}

// Grep emits every input x that matches the regular expression r.
func Grep(r string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return If(re.MatchString)
}

// GrepNot emits every input x that does not match the regular expression r.
func GrepNot(r string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return If(func(s string) bool { return !re.MatchString(s) })
}

// Uniq squashes adjacent identical items in arg.In into a single output.
func Uniq() Filter {
	return func(arg Arg) {
		first := true
		last := ""
		for s := range arg.In {
			if first || last != s {
				arg.Out <- s
			}
			last = s
			first = false
		}
	}
}

// UniqWithCount squashes adjacent identical items in arg.In into a single
// output prefixed with the count of identical items.
func UniqWithCount() Filter {
	return func(arg Arg) {
		current := ""
		count := 0
		for s := range arg.In {
			if s != current {
				if count > 0 {
					arg.Out <- fmt.Sprintf("%d %s", count, current)
				}
				count = 0
				current = s
			}
			count++
		}
		if count > 0 {
			arg.Out <- fmt.Sprintf("%d %s", count, current)
		}
	}
}

// Substitute replaces all occurrences of the regular expression r in
// an input item with replacement.  The replacement string can contain
// $1, $2, etc. which represent submatches of r.
func Substitute(r, replacement string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return func(arg Arg) {
		for s := range arg.In {
			arg.Out <- re.ReplaceAllString(s, replacement)
		}
	}
}

// Reverse yields items in the reverse of the order it received them.
func Reverse() Filter {
	return func(arg Arg) {
		var data []string
		for s := range arg.In {
			data = append(data, s)
		}
		for i := len(data) - 1; i >= 0; i-- {
			arg.Out <- data[i]
		}
	}
}

// First yields the first n items that it receives.
func First(n int) Filter {
	return func(arg Arg) {
		emitted := 0
		for s := range arg.In {
			if emitted < n {
				arg.Out <- s
				emitted++
			}
		}
	}
}

// DropFirst yields all items except for the first n items that it receives.
func DropFirst(n int) Filter {
	return func(arg Arg) {
		emitted := 0
		for s := range arg.In {
			if emitted >= n {
				arg.Out <- s
			}
			emitted++
		}
	}
}

// Last yields the last n items that it receives.
func Last(n int) Filter {
	return func(arg Arg) {
		var buf []string
		for s := range arg.In {
			buf = append(buf, s)
			if len(buf) > n {
				buf = buf[1:]
			}
		}
		for _, s := range buf {
			arg.Out <- s
		}
	}
}

// DropLast yields all items except for the last n items that it receives.
func DropLast(n int) Filter {
	return func(arg Arg) {
		var buf []string
		for s := range arg.In {
			buf = append(buf, s)
			if len(buf) > n {
				arg.Out <- buf[0]
				buf = buf[1:]
			}
		}
	}
}

// NumberLines prefixes its item with its index in the input sequence
// (starting at 1).
func NumberLines() Filter {
	return func(arg Arg) {
		line := 1
		for s := range arg.In {
			arg.Out <- fmt.Sprintf("%5d %s", line, s)
			line++
		}
	}
}

// Cut emits just the bytes indexed [start..end] of each input item.
func Cut(start, end int) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			if len(s) > end {
				s = s[:end+1]
			}
			if len(s) < start {
				s = ""
			} else {
				s = s[start:]
			}
			arg.Out <- s
		}
	}
}

// Select splits each item into columns and yields the concatenation
// of the columns numbers passed as arguments to Select.  Columns are
// numbered starting at 1. A column number of 0 is interpreted as the
// full string.
func Select(columns ...int) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			result := ""
			for _, col := range columns {
				if _, c := column(s, col); c != "" {
					if result != "" {
						result = result + " "
					}
					result = result + c
				}
			}
			arg.Out <- result
		}
	}
}
<commit_msg>package doc examples now use fully qualified names<commit_after>/*
Package pipe provides filters that can be chained together in a manner
similar to Unix pipelines.

Each filter is a function that takes as input a sequence of
strings (read from a channel) and produces as output a sequence of
strings (written to a channel).

Filters can be chained together (e.g., via the Run function), the
output of one filter is fed as input to the next filter.  The empty
input is passed to the first filter. The following sequence will
print two lines to standard output:

	err := pipe.Run(
		pipe.Echo("hello", "world"),
		pipe.Reverse(),
		pipe.WriteLines(os.Stdout),
	)

An application can implement its own filters easily. For example,
repeat(n) returns a filter that repeats every input n times:

	func repeat(n int) Filter {
		return func(arg pipe.Arg) {
			for s := range arg.In {
				for i := 0; i < n; i++ {
					arg.Out <- s
				}
			}
		}
	}

	pipe.Run(
		pipe.Echo("hello"),
		repeat(10),
	)
*/
package pipe

import (
	"fmt"
	"regexp"
	"sync"
)

// filterErrors records errors accumulated during the execution of a filter.
type filterErrors struct {
	mu     sync.Mutex
	errors []error
}

// Arg contains the data passed to a Filter. Arg.In is a channel that
// produces the input to the filter, and Arg.Out is a channel that
// receives the output from the filter.
type Arg struct {
	In     <-chan string
	Out    chan<- string
	errors *filterErrors
}

// ReportError records an error encountered during an execution of a filter.
// This error will be reported by whatever facility (e.g., ForEach or Run)
// was being used to execute the filters.
func (a *Arg) ReportError(err error) {
	a.errors.mu.Lock()
	defer a.errors.mu.Unlock()
	a.errors.errors = append(a.errors.errors, err)
}

// Filter reads a sequence of strings from a channel and produces a
// sequence on another channel.
type Filter func(Arg)

// Sequence returns a filter that is the concatenation of all filter arguments.
// The output of a filter is fed as input to the next filter.
func Sequence(filters ...Filter) Filter {
	return func(arg Arg) {
		in := arg.In
		for _, f := range filters {
			c := make(chan string, 10000)
			go runAndClose(f, Arg{in, c, arg.errors})
			in = c
		}
		passThrough(Arg{in, arg.Out, arg.errors})
	}
}

// Run executes the sequence of filters and discards all output.
// It returns either nil, an error if any filter reported an error.
func Run(filters ...Filter) error {
	return ForEach(Sequence(filters...), func(s string) {})
}

// ForEach calls fn(s) for every item s in the output of filter and
// returns either nil, or any error reported by the execution of the filter.
func ForEach(filter Filter, fn func(s string)) error {
	in := make(chan string, 0)
	close(in)
	out := make(chan string, 10000)
	e := &filterErrors{}
	go runAndClose(filter, Arg{in, out, e})
	for s := range out {
		fn(s)
	}

	e.mu.Lock()
	defer e.mu.Unlock()
	switch len(e.errors) {
	case 0:
		return nil
	case 1:
		return e.errors[0]
	default:
		return fmt.Errorf("Filter errors: %s", e.errors)
	}
}

func runAndClose(f Filter, arg Arg) {
	f(arg)
	close(arg.Out)
}

// passThrough copies all items read from in to out.
func passThrough(arg Arg) {
	for s := range arg.In {
		arg.Out <- s
	}
}

// Echo emits items.
// Any input items are copied verbatim to the output before items are emitted.
func Echo(items ...string) Filter {
	return func(arg Arg) {
		passThrough(arg)
		for _, s := range items {
			arg.Out <- s
		}
	}
}

// Numbers copies its input and then emits the integers x..y
func Numbers(x, y int) Filter {
	return func(arg Arg) {
		passThrough(arg)
		for i := x; i <= y; i++ {
			arg.Out <- fmt.Sprintf("%d", i)
		}
	}
}

// Map calls fn(x) for every item x and yields the outputs of the fn calls.
func Map(fn func(string) string) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			arg.Out <- fn(s)
		}
	}
}

// If emits every input x for which fn(x) is true.
func If(fn func(string) bool) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			if fn(s) {
				arg.Out <- s
			}
		}
	}
}

func errorFilter(err error) Filter {
	return func(arg Arg) {
		arg.ReportError(err)
		for _ = range arg.In {
			// Drop the input
		}
	}
}

// Grep emits every input x that matches the regular expression r.
func Grep(r string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return If(re.MatchString)
}

// GrepNot emits every input x that does not match the regular expression r.
func GrepNot(r string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return If(func(s string) bool { return !re.MatchString(s) })
}

// Uniq squashes adjacent identical items in arg.In into a single output.
func Uniq() Filter {
	return func(arg Arg) {
		first := true
		last := ""
		for s := range arg.In {
			if first || last != s {
				arg.Out <- s
			}
			last = s
			first = false
		}
	}
}

// UniqWithCount squashes adjacent identical items in arg.In into a single
// output prefixed with the count of identical items.
func UniqWithCount() Filter {
	return func(arg Arg) {
		current := ""
		count := 0
		for s := range arg.In {
			if s != current {
				if count > 0 {
					arg.Out <- fmt.Sprintf("%d %s", count, current)
				}
				count = 0
				current = s
			}
			count++
		}
		if count > 0 {
			arg.Out <- fmt.Sprintf("%d %s", count, current)
		}
	}
}

// Substitute replaces all occurrences of the regular expression r in
// an input item with replacement.  The replacement string can contain
// $1, $2, etc. which represent submatches of r.
func Substitute(r, replacement string) Filter {
	re, err := regexp.Compile(r)
	if err != nil {
		return errorFilter(err)
	}
	return func(arg Arg) {
		for s := range arg.In {
			arg.Out <- re.ReplaceAllString(s, replacement)
		}
	}
}

// Reverse yields items in the reverse of the order it received them.
func Reverse() Filter {
	return func(arg Arg) {
		var data []string
		for s := range arg.In {
			data = append(data, s)
		}
		for i := len(data) - 1; i >= 0; i-- {
			arg.Out <- data[i]
		}
	}
}

// First yields the first n items that it receives.
func First(n int) Filter {
	return func(arg Arg) {
		emitted := 0
		for s := range arg.In {
			if emitted < n {
				arg.Out <- s
				emitted++
			}
		}
	}
}

// DropFirst yields all items except for the first n items that it receives.
func DropFirst(n int) Filter {
	return func(arg Arg) {
		emitted := 0
		for s := range arg.In {
			if emitted >= n {
				arg.Out <- s
			}
			emitted++
		}
	}
}

// Last yields the last n items that it receives.
func Last(n int) Filter {
	return func(arg Arg) {
		var buf []string
		for s := range arg.In {
			buf = append(buf, s)
			if len(buf) > n {
				buf = buf[1:]
			}
		}
		for _, s := range buf {
			arg.Out <- s
		}
	}
}

// DropLast yields all items except for the last n items that it receives.
func DropLast(n int) Filter {
	return func(arg Arg) {
		var buf []string
		for s := range arg.In {
			buf = append(buf, s)
			if len(buf) > n {
				arg.Out <- buf[0]
				buf = buf[1:]
			}
		}
	}
}

// NumberLines prefixes its item with its index in the input sequence
// (starting at 1).
func NumberLines() Filter {
	return func(arg Arg) {
		line := 1
		for s := range arg.In {
			arg.Out <- fmt.Sprintf("%5d %s", line, s)
			line++
		}
	}
}

// Cut emits just the bytes indexed [start..end] of each input item.
func Cut(start, end int) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			if len(s) > end {
				s = s[:end+1]
			}
			if len(s) < start {
				s = ""
			} else {
				s = s[start:]
			}
			arg.Out <- s
		}
	}
}

// Select splits each item into columns and yields the concatenation
// of the columns numbers passed as arguments to Select.  Columns are
// numbered starting at 1. A column number of 0 is interpreted as the
// full string.
func Select(columns ...int) Filter {
	return func(arg Arg) {
		for s := range arg.In {
			result := ""
			for _, col := range columns {
				if _, c := column(s, col); c != "" {
					if result != "" {
						result = result + " "
					}
					result = result + c
				}
			}
			arg.Out <- result
		}
	}
}
<|endoftext|>


<commit_before>package nio

import (
	"io"
	"lib/channel"
)

func Pipe() (io.ReadCloser, io.WriteCloser) {
	inReader, inWriter := io.Pipe()
	outReader, outWriter := io.Pipe()

	input := make(chan []byte)
	output := make(chan []byte)

	go channel.Chan(input, output)
	go inFeed(inReader, input)
	go outFeed(outWriter, output)

	return outReader, inWriter
}

func inFeed(r io.Reader, in chan<- []byte) {
	for {
		data := make([]byte, 32*1024)
		if n, err := r.Read(data); err == nil {
			in <- data[:n]
		} else {
			close(in)
			return
		}
	}
}

func outFeed(w io.WriteCloser, out <-chan []byte) {
	for output := range out {
		for len(output) > 0 {
			n, _ := w.Write(output)
			output = output[n:]
		}

	}
	w.Close()
}
<commit_msg>Upgraded to fit changes to buffer, channel, and queue<commit_after>package nio

import (
	"io"
	"lib/buffer"
	"lib/channel"
)

func Pipe() (io.ReadCloser, io.WriteCloser) {
	inReader, inWriter := io.Pipe()
	outReader, outWriter := io.Pipe()

	input := make(chan interface{})
	output := make(chan interface{})

	pending := buffer.NewBufferQueue(buffer.NewUnboundedBuffer(32*1024, 100*1024*1024))
	go channel.ChanQueue(input, output, pending)
	go inFeed(inReader, input)
	go outFeed(outWriter, output)

	return outReader, inWriter
}

func inFeed(r io.Reader, in chan<- interface{}) {
	for {
		data := make([]byte, 32*1024)
		if n, err := r.Read(data); err == nil {
			in <- data[:n]
		} else {
			close(in)
			return
		}
	}
}

func outFeed(w io.WriteCloser, out <-chan interface{}) {
	for output := range out {
		data := output.([]byte)
		for len(data) > 0 {
			n, _ := w.Write(data)
			data = data[n:]
		}

	}
	w.Close()
}
<|endoftext|>


<commit_before>// Copyright Istio 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 kube

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/portforward"
	"k8s.io/client-go/transport/spdy"
)

// PortForwarder manages the forwarding of a single port.
type PortForwarder interface {
	// Start runs this forwarder.
	Start() error

	// Address returns the local forwarded address. Only valid while the forwarder is running.
	Address() string

	// Close this forwarder and release an resources.
	Close()

	// WaitForStop blocks until connection closed (e.g. control-C interrupt)
	WaitForStop()
}

var _ PortForwarder = &forwarder{}

type forwarder struct {
	forwarder *portforward.PortForwarder
	stopCh    chan struct{}
	readyCh   <-chan struct{}
	address   string
}

func (f *forwarder) Start() error {
	errCh := make(chan error, 1)
	go func() {
		errCh <- f.forwarder.ForwardPorts()
	}()

	select {
	case err := <-errCh:
		return fmt.Errorf("failure running port forward process: %v", err)
	case <-f.readyCh:
		// The forwarder is now ready.
		return nil
	}
}

func (f *forwarder) Address() string {
	return f.address
}

func (f *forwarder) Close() {
	close(f.stopCh)
	// Closing the stop channel should close anything
	// opened by f.forwarder.ForwardPorts()
}

func (f *forwarder) WaitForStop() {
	<-f.stopCh
}

func newPortForwarder(restConfig *rest.Config, podName, ns, localAddress string, localPort, podPort int) (PortForwarder, error) {
	restClient, err := rest.RESTClientFor(restConfig)
	if err != nil {
		return nil, err
	}

	req := restClient.Post().Resource("pods").Namespace(ns).Name(podName).SubResource("portforward")
	serverURL := req.URL()

	roundTripper, upgrader, err := roundTripperFor(restConfig)
	if err != nil {
		return nil, fmt.Errorf("failure creating roundtripper: %v", err)
	}

	dialer := spdy.NewDialer(upgrader, &http.Client{Transport: roundTripper}, http.MethodPost, serverURL)

	stopCh := make(chan struct{})
	readyCh := make(chan struct{})
	if localAddress == "" {
		localAddress = defaultLocalAddress
	}
	if localPort == 0 {
		localPort, err = availablePort(localAddress)
		if err != nil {
			return nil, fmt.Errorf("failure allocating port: %v", err)
		}
	}
	fw, err := portforward.NewOnAddresses(dialer,
		[]string{localAddress},
		[]string{fmt.Sprintf("%d:%d", localPort, podPort)},
		stopCh,
		readyCh,
		io.Discard,
		os.Stderr)
	if err != nil {
		return nil, fmt.Errorf("failed establishing port-forward: %v", err)
	}

	// Run the same check as k8s.io/kubectl/pkg/cmd/portforward/portforward.go
	// so that we will fail early if there is a problem contacting API server.
	podGet := restClient.Get().Resource("pods").Namespace(ns).Name(podName)
	obj, err := podGet.Do(context.TODO()).Get()
	if err != nil {
		return nil, fmt.Errorf("failed retrieving: %v in the %q namespace", err, ns)
	}
	pod, ok := obj.(*v1.Pod)
	if !ok {
		return nil, fmt.Errorf("failed getting pod, object type is %T", obj)
	}
	if pod.Status.Phase != v1.PodRunning {
		return nil, fmt.Errorf("pod is not running. Status=%v", pod.Status.Phase)
	}

	sLocalPort := fmt.Sprintf("%d", localPort)
	return &forwarder{
		forwarder: fw,
		stopCh:    stopCh,
		readyCh:   readyCh,
		address:   net.JoinHostPort(localAddress, sLocalPort),
	}, nil
}

func availablePort(localAddr string) (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(localAddr, "0"))
	if err != nil {
		return 0, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return 0, err
	}
	port := l.Addr().(*net.TCPAddr).Port
	return port, l.Close()
}
<commit_msg>Make port forwarder resilient to broken connection. (#35658)<commit_after>// Copyright Istio 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 kube

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"

	v1 "k8s.io/api/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/portforward"
	"k8s.io/client-go/transport/spdy"
)

// PortForwarder manages the forwarding of a single port.
type PortForwarder interface {
	// Start runs this forwarder.
	Start() error

	// Address returns the local forwarded address. Only valid while the forwarder is running.
	Address() string

	// Close this forwarder and release an resources.
	Close()

	// WaitForStop blocks until connection closed (e.g. control-C interrupt)
	WaitForStop()
}

var _ PortForwarder = &forwarder{}

type forwarder struct {
	stopCh       chan struct{}
	restConfig   *rest.Config
	podName      string
	ns           string
	localAddress string
	localPort    int
	podPort      int
	address      string
}

func (f *forwarder) Start() error {
	errCh := make(chan error, 1)
	readyCh := make(chan struct{}, 1)
	go func() {
		for {
			select {
			case <-f.stopCh:
				return
			default:
			}
			// Build a new port forwarder.
			fw, err := f.buildK8sPortForwarder(readyCh)
			if err != nil {
				errCh <- err
				break
			}
			if err = fw.ForwardPorts(); err != nil {
				errCh <- err
				break
			}
			// At this point, either the stopCh has been closed, or port forwarder connection is broken.
			// the port forwarder should have already been ready before.
			// No need to notify the ready channel anymore when forwarding again.
			readyCh = nil
		}
	}()

	select {
	case err := <-errCh:
		return fmt.Errorf("failure running port forward process: %v", err)
	case <-readyCh:
		// The forwarder is now ready.
		return nil
	}
}

func (f *forwarder) Address() string {
	return f.address
}

func (f *forwarder) Close() {
	close(f.stopCh)
	// Closing the stop channel should close anything
	// opened by f.forwarder.ForwardPorts()
}

func (f *forwarder) WaitForStop() {
	<-f.stopCh
}

func newPortForwarder(restConfig *rest.Config, podName, ns, localAddress string, localPort, podPort int) (PortForwarder, error) {
	f := &forwarder{
		stopCh:       make(chan struct{}),
		restConfig:   restConfig,
		podName:      podName,
		ns:           ns,
		localAddress: localAddress,
		localPort:    localPort,
		podPort:      podPort,
	}
	if f.localPort == 0 {
		var err error
		f.localPort, err = availablePort(f.localAddress)
		if err != nil {
			return nil, fmt.Errorf("failure allocating port: %v", err)
		}
	}

	sLocalPort := fmt.Sprintf("%d", f.localPort)
	f.address = net.JoinHostPort(localAddress, sLocalPort)
	return f, nil
}

func (f *forwarder) buildK8sPortForwarder(readyCh chan struct{}) (*portforward.PortForwarder, error) {
	restClient, err := rest.RESTClientFor(f.restConfig)
	if err != nil {
		return nil, err
	}

	req := restClient.Post().Resource("pods").Namespace(f.ns).Name(f.podName).SubResource("portforward")
	serverURL := req.URL()

	roundTripper, upgrader, err := roundTripperFor(f.restConfig)
	if err != nil {
		return nil, fmt.Errorf("failure creating roundtripper: %v", err)
	}

	dialer := spdy.NewDialer(upgrader, &http.Client{Transport: roundTripper}, http.MethodPost, serverURL)

	if f.localAddress == "" {
		f.localAddress = defaultLocalAddress
	}
	fw, err := portforward.NewOnAddresses(dialer,
		[]string{f.localAddress},
		[]string{fmt.Sprintf("%d:%d", f.localPort, f.podPort)},
		f.stopCh,
		readyCh,
		io.Discard,
		os.Stderr)
	if err != nil {
		return nil, fmt.Errorf("failed establishing port-forward: %v", err)
	}

	// Run the same check as k8s.io/kubectl/pkg/cmd/portforward/portforward.go
	// so that we will fail early if there is a problem contacting API server.
	podGet := restClient.Get().Resource("pods").Namespace(f.ns).Name(f.podName)
	obj, err := podGet.Do(context.TODO()).Get()
	if err != nil {
		return nil, fmt.Errorf("failed retrieving: %v in the %q namespace", err, f.ns)
	}
	pod, ok := obj.(*v1.Pod)
	if !ok {
		return nil, fmt.Errorf("failed getting pod, object type is %T", obj)
	}
	if pod.Status.Phase != v1.PodRunning {
		return nil, fmt.Errorf("pod is not running. Status=%v", pod.Status.Phase)
	}

	return fw, nil
}

func availablePort(localAddr string) (int, error) {
	addr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(localAddr, "0"))
	if err != nil {
		return 0, err
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return 0, err
	}
	port := l.Addr().(*net.TCPAddr).Port
	return port, l.Close()
}
<|endoftext|>


<commit_before>package offers

import (
	"errors"
	"testing"
	"time"

	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesosphere/kubernetes-mesos/pkg/proc"
)

func TestExpiredOffer(t *testing.T) {
	t.Parallel()

	ttl := 2 * time.Second
	o := Expired("test", "testhost", ttl)

	if o.Id() != "test" {
		t.Error("expiredOffer does not return its Id")
	}
	if o.Host() != "testhost" {
		t.Error("expiredOffer does not return its hostname")
	}
	if o.HasExpired() != true {
		t.Error("expiredOffer is not expired")
	}
	if o.Details() != nil {
		t.Error("expiredOffer does not return nil Details")
	}
	if o.Acquire() != false {
		t.Error("expiredOffer must not be able to be acquired")
	}
	if delay := o.GetDelay(); !(0 < delay && delay <= ttl) {
		t.Error("expiredOffer does not return a valid deadline")
	}
} // TestExpiredOffer

func TestTimedOffer(t *testing.T) {
	t.Parallel()

	ttl := 2 * time.Second
	now := time.Now()
	o := &liveOffer{nil, now.Add(ttl), 0}

	if o.HasExpired() {
		t.Errorf("offer ttl was %v and should not have expired yet", ttl)
	}
	if !o.Acquire() {
		t.Fatal("1st acquisition of offer failed")
	}
	o.Release()
	if !o.Acquire() {
		t.Fatal("2nd acquisition of offer failed")
	}
	if o.Acquire() {
		t.Fatal("3rd acquisition of offer passed but prior claim was not released")
	}
	o.Release()
	if !o.Acquire() {
		t.Fatal("4th acquisition of offer failed")
	}
	o.Release()
	time.Sleep(ttl)
	if !o.HasExpired() {
		t.Fatal("offer not expired after ttl passed")
	}
	if !o.Acquire() {
		t.Fatal("5th acquisition of offer failed; should not be tied to expiration")
	}
	if o.Acquire() {
		t.Fatal("6th acquisition of offer succeeded; should already be acquired")
	}
} // TestTimedOffer

func TestOfferStorage(t *testing.T) {
	ttl := time.Second / 4
	declinedNum := 0
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			declinedNum = declinedNum + 1
			return proc.ErrorChan(nil)
		},
		Compat: func(o *mesos.Offer) bool {
			return o.Hostname == nil || *o.Hostname != "incompatiblehost"
		},
		TTL:           ttl,
		LingerTTL:     2 * ttl,
	}
	storage := CreateRegistry(config)

	done := make(chan struct{})
	storage.Init(done)

	// Add offer
	id := util.NewOfferID("foo")
	o := &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})

	// Added offer should be in the storage
	if obj, ok := storage.Get(id.GetValue()); obj == nil || !ok {
		t.Error("offer not added")
	}
	if obj, _ := storage.Get(id.GetValue()); obj.Details() != o {
		t.Error("added offer differs from returned offer")
	}

	// Not-added offer is not in storage
	if obj, ok := storage.Get("bar"); obj != nil || ok {
		t.Error("offer bar should not exist in storage")
	}

	// Deleted offer lingers in storage, is acquired and declined
	offer, _ := storage.Get(id.GetValue())
	declinedNumBefore := declinedNum
	storage.Delete(id.GetValue(), "deleted for test")
	if obj, _ := storage.Get(id.GetValue()); obj == nil {
		t.Error("deleted offer is not lingering")
	}
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("deleted offer is no expired")
	}
	if ok := offer.Acquire(); ok {
		t.Error("deleted offer can be acquired")
	}
	if declinedNum <= declinedNumBefore {
		t.Error("deleted offer was not declined")
	}

	// Acquired offer is only declined after 2*ttl
	id = util.NewOfferID("foo2")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	offer, _ = storage.Get(id.GetValue())
	declinedNumBefore = declinedNum
	offer.Acquire()
	storage.Delete(id.GetValue(), "deleted for test")
	if declinedNum > declinedNumBefore {
		t.Error("acquired offer is declined")
	}

	offer.Release()
	time.Sleep(3 * ttl)
	if declinedNum <= declinedNumBefore {
		t.Error("released offer is not declined after 2*ttl")
	}

	// Added offer should be expired after ttl, but lingering
	id = util.NewOfferID("foo3")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})

	time.Sleep(2 * ttl)
	obj, ok := storage.Get(id.GetValue())
	if obj == nil || !ok {
		t.Error("offer not lingering after ttl")
	}
	if !obj.HasExpired() {
		t.Error("offer is not expired after ttl")
	}

	// Should be deleted when waiting longer than LingerTTL
	time.Sleep(2 * ttl)
	if obj, ok := storage.Get(id.GetValue()); obj != nil || ok {
		t.Error("offer not deleted after LingerTTL")
	}

	// Incompatible offer is declined
	id = util.NewOfferID("foo4")
	incompatibleHostname := "incompatiblehost"
	o = &mesos.Offer{Id: id, Hostname:&incompatibleHostname}
	declinedNumBefore = declinedNum
	storage.Add([]*mesos.Offer{o})
	if obj, ok := storage.Get(id.GetValue()); obj != nil || ok {
		t.Error("incompatible offer not rejected")
	}
	if declinedNum <= declinedNumBefore {
		t.Error("incompatible offer is not declined")
	}

	// Invalidated offer are not declined, but expired
	id = util.NewOfferID("foo5")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	offer, _ = storage.Get(id.GetValue())
	declinedNumBefore = declinedNum
	storage.Invalidate(id.GetValue())
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer is not expired")
	}
	if declinedNum > declinedNumBefore {
		t.Error("invalidated offer is declined")
	}
	if ok := offer.Acquire(); ok {
		t.Error("invalidated offer can be acquired")
	}

	// Invalidate "" will invalidate all offers
	id = util.NewOfferID("foo6")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	id2 := util.NewOfferID("foo7")
	o2 := &mesos.Offer{Id: id2}
	storage.Add([]*mesos.Offer{o2})
	storage.Invalidate("")
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer is not expired")
	}
	if obj2, _ := storage.Get(id2.GetValue()); !obj2.HasExpired() {
		t.Error("invalidated offer is not expired")
	}

	// InvalidateForSlave invalides all offers for that slave, but only those
	id = util.NewOfferID("foo8")
	slaveId := util.NewSlaveID("test-slave")
	o = &mesos.Offer{Id: id, SlaveId:slaveId}
	storage.Add([]*mesos.Offer{o})
	id2 = util.NewOfferID("foo9")
	o2 = &mesos.Offer{Id: id2}
	storage.Add([]*mesos.Offer{o2})
	storage.InvalidateForSlave(slaveId.GetValue())
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer for test-slave is not expired")
	}
	if obj2, _ := storage.Get(id2.GetValue()); obj2.HasExpired() {
		t.Error("invalidated offer another slave is expired")
	}

	close(done)
} // TestOfferStorage

func TestListen(t *testing.T) {
	ttl := time.Second / 4
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			return proc.ErrorChan(nil)
		},
		Compat: func(o *mesos.Offer) bool {
			return true
		},
		TTL:           ttl,
		ListenerDelay: ttl / 2,
	}
	storage := CreateRegistry(config)

	done := make(chan struct {})
	storage.Init(done)

	// Create two listeners with a hostname filter
	hostname1 := "hostname1"
	hostname2 := "hostname2"
	listener1 := storage.Listen("listener1", func(offer *mesos.Offer) bool {
		return offer.GetHostname() == hostname1
	})
	listener2 := storage.Listen("listener2", func(offer *mesos.Offer) bool {
		return offer.GetHostname() == hostname2
	})

	// Add hostname1 offer
	id := util.NewOfferID("foo")
	o := &mesos.Offer{Id: id, Hostname:&hostname1}
	storage.Add([]*mesos.Offer{o})

	// listener1 is notified by closing channel
	select {
	case _, more := <-listener1:
		if more {
			t.Error("listener1 is not closed")
		}
	}

	// listener2 is not notified within ttl
	select {
		case <-listener2:
			t.Error("listener2 is notified")
		case <-time.After(ttl):
	}

	close(done)
} // TestListen

func TestWalk(t *testing.T) {
	t.Parallel()
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			return proc.ErrorChan(nil)
		},
		TTL:           0 * time.Second,
		LingerTTL:     0 * time.Second,
		ListenerDelay: 0 * time.Second,
	}
	storage := CreateRegistry(config)
	acceptedOfferId := ""
	walked := 0
	walker1 := func(p Perishable) (bool, error) {
		walked++
		if p.Acquire() {
			acceptedOfferId = p.Details().Id.GetValue()
			return true, nil
		}
		return false, nil
	}
	// sanity check
	err := storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 0 {
		t.Fatal("walked empty storage")
	}
	if acceptedOfferId != "" {
		t.Fatal("somehow found an offer when registry was empty")
	}
	impl, ok := storage.(*offerStorage)
	if !ok {
		t.Fatal("unexpected offer storage impl")
	}
	// single offer
	ttl := 2 * time.Second
	now := time.Now()
	o := &liveOffer{&mesos.Offer{Id: util.NewOfferID("foo")}, now.Add(ttl), 0}

	impl.offers.Add(o)
	err = storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 1 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "foo" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	acceptedOfferId = ""
	err = storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 2 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	walker2 := func(p Perishable) (bool, error) {
		walked++
		return true, nil
	}
	err = storage.Walk(walker2)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 3 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	walker3 := func(p Perishable) (bool, error) {
		walked++
		return true, errors.New("baz")
	}
	err = storage.Walk(walker3)
	if err == nil {
		t.Fatal("expected error")
	}
	if walked != 4 {
		t.Fatalf("walk count %d", walked)
	}
}
<commit_msg>format fixes for offer tests<commit_after>package offers

import (
	"errors"
	"testing"
	"time"

	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesosphere/kubernetes-mesos/pkg/proc"
)

func TestExpiredOffer(t *testing.T) {
	t.Parallel()

	ttl := 2 * time.Second
	o := Expired("test", "testhost", ttl)

	if o.Id() != "test" {
		t.Error("expiredOffer does not return its Id")
	}
	if o.Host() != "testhost" {
		t.Error("expiredOffer does not return its hostname")
	}
	if o.HasExpired() != true {
		t.Error("expiredOffer is not expired")
	}
	if o.Details() != nil {
		t.Error("expiredOffer does not return nil Details")
	}
	if o.Acquire() != false {
		t.Error("expiredOffer must not be able to be acquired")
	}
	if delay := o.GetDelay(); !(0 < delay && delay <= ttl) {
		t.Error("expiredOffer does not return a valid deadline")
	}
} // TestExpiredOffer

func TestTimedOffer(t *testing.T) {
	t.Parallel()

	ttl := 2 * time.Second
	now := time.Now()
	o := &liveOffer{nil, now.Add(ttl), 0}

	if o.HasExpired() {
		t.Errorf("offer ttl was %v and should not have expired yet", ttl)
	}
	if !o.Acquire() {
		t.Fatal("1st acquisition of offer failed")
	}
	o.Release()
	if !o.Acquire() {
		t.Fatal("2nd acquisition of offer failed")
	}
	if o.Acquire() {
		t.Fatal("3rd acquisition of offer passed but prior claim was not released")
	}
	o.Release()
	if !o.Acquire() {
		t.Fatal("4th acquisition of offer failed")
	}
	o.Release()
	time.Sleep(ttl)
	if !o.HasExpired() {
		t.Fatal("offer not expired after ttl passed")
	}
	if !o.Acquire() {
		t.Fatal("5th acquisition of offer failed; should not be tied to expiration")
	}
	if o.Acquire() {
		t.Fatal("6th acquisition of offer succeeded; should already be acquired")
	}
} // TestTimedOffer

func TestOfferStorage(t *testing.T) {
	ttl := time.Second / 4
	declinedNum := 0
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			declinedNum = declinedNum + 1
			return proc.ErrorChan(nil)
		},
		Compat: func(o *mesos.Offer) bool {
			return o.Hostname == nil || *o.Hostname != "incompatiblehost"
		},
		TTL:       ttl,
		LingerTTL: 2 * ttl,
	}
	storage := CreateRegistry(config)

	done := make(chan struct{})
	storage.Init(done)

	// Add offer
	id := util.NewOfferID("foo")
	o := &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})

	// Added offer should be in the storage
	if obj, ok := storage.Get(id.GetValue()); obj == nil || !ok {
		t.Error("offer not added")
	}
	if obj, _ := storage.Get(id.GetValue()); obj.Details() != o {
		t.Error("added offer differs from returned offer")
	}

	// Not-added offer is not in storage
	if obj, ok := storage.Get("bar"); obj != nil || ok {
		t.Error("offer bar should not exist in storage")
	}

	// Deleted offer lingers in storage, is acquired and declined
	offer, _ := storage.Get(id.GetValue())
	declinedNumBefore := declinedNum
	storage.Delete(id.GetValue(), "deleted for test")
	if obj, _ := storage.Get(id.GetValue()); obj == nil {
		t.Error("deleted offer is not lingering")
	}
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("deleted offer is no expired")
	}
	if ok := offer.Acquire(); ok {
		t.Error("deleted offer can be acquired")
	}
	if declinedNum <= declinedNumBefore {
		t.Error("deleted offer was not declined")
	}

	// Acquired offer is only declined after 2*ttl
	id = util.NewOfferID("foo2")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	offer, _ = storage.Get(id.GetValue())
	declinedNumBefore = declinedNum
	offer.Acquire()
	storage.Delete(id.GetValue(), "deleted for test")
	if declinedNum > declinedNumBefore {
		t.Error("acquired offer is declined")
	}

	offer.Release()
	time.Sleep(3 * ttl)
	if declinedNum <= declinedNumBefore {
		t.Error("released offer is not declined after 2*ttl")
	}

	// Added offer should be expired after ttl, but lingering
	id = util.NewOfferID("foo3")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})

	time.Sleep(2 * ttl)
	obj, ok := storage.Get(id.GetValue())
	if obj == nil || !ok {
		t.Error("offer not lingering after ttl")
	}
	if !obj.HasExpired() {
		t.Error("offer is not expired after ttl")
	}

	// Should be deleted when waiting longer than LingerTTL
	time.Sleep(2 * ttl)
	if obj, ok := storage.Get(id.GetValue()); obj != nil || ok {
		t.Error("offer not deleted after LingerTTL")
	}

	// Incompatible offer is declined
	id = util.NewOfferID("foo4")
	incompatibleHostname := "incompatiblehost"
	o = &mesos.Offer{Id: id, Hostname: &incompatibleHostname}
	declinedNumBefore = declinedNum
	storage.Add([]*mesos.Offer{o})
	if obj, ok := storage.Get(id.GetValue()); obj != nil || ok {
		t.Error("incompatible offer not rejected")
	}
	if declinedNum <= declinedNumBefore {
		t.Error("incompatible offer is not declined")
	}

	// Invalidated offer are not declined, but expired
	id = util.NewOfferID("foo5")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	offer, _ = storage.Get(id.GetValue())
	declinedNumBefore = declinedNum
	storage.Invalidate(id.GetValue())
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer is not expired")
	}
	if declinedNum > declinedNumBefore {
		t.Error("invalidated offer is declined")
	}
	if ok := offer.Acquire(); ok {
		t.Error("invalidated offer can be acquired")
	}

	// Invalidate "" will invalidate all offers
	id = util.NewOfferID("foo6")
	o = &mesos.Offer{Id: id}
	storage.Add([]*mesos.Offer{o})
	id2 := util.NewOfferID("foo7")
	o2 := &mesos.Offer{Id: id2}
	storage.Add([]*mesos.Offer{o2})
	storage.Invalidate("")
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer is not expired")
	}
	if obj2, _ := storage.Get(id2.GetValue()); !obj2.HasExpired() {
		t.Error("invalidated offer is not expired")
	}

	// InvalidateForSlave invalides all offers for that slave, but only those
	id = util.NewOfferID("foo8")
	slaveId := util.NewSlaveID("test-slave")
	o = &mesos.Offer{Id: id, SlaveId: slaveId}
	storage.Add([]*mesos.Offer{o})
	id2 = util.NewOfferID("foo9")
	o2 = &mesos.Offer{Id: id2}
	storage.Add([]*mesos.Offer{o2})
	storage.InvalidateForSlave(slaveId.GetValue())
	if obj, _ := storage.Get(id.GetValue()); !obj.HasExpired() {
		t.Error("invalidated offer for test-slave is not expired")
	}
	if obj2, _ := storage.Get(id2.GetValue()); obj2.HasExpired() {
		t.Error("invalidated offer another slave is expired")
	}

	close(done)
} // TestOfferStorage

func TestListen(t *testing.T) {
	ttl := time.Second / 4
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			return proc.ErrorChan(nil)
		},
		Compat: func(o *mesos.Offer) bool {
			return true
		},
		TTL:           ttl,
		ListenerDelay: ttl / 2,
	}
	storage := CreateRegistry(config)

	done := make(chan struct{})
	storage.Init(done)

	// Create two listeners with a hostname filter
	hostname1 := "hostname1"
	hostname2 := "hostname2"
	listener1 := storage.Listen("listener1", func(offer *mesos.Offer) bool {
		return offer.GetHostname() == hostname1
	})
	listener2 := storage.Listen("listener2", func(offer *mesos.Offer) bool {
		return offer.GetHostname() == hostname2
	})

	// Add hostname1 offer
	id := util.NewOfferID("foo")
	o := &mesos.Offer{Id: id, Hostname: &hostname1}
	storage.Add([]*mesos.Offer{o})

	// listener1 is notified by closing channel
	select {
	case _, more := <-listener1:
		if more {
			t.Error("listener1 is not closed")
		}
	}

	// listener2 is not notified within ttl
	select {
	case <-listener2:
		t.Error("listener2 is notified")
	case <-time.After(ttl):
	}

	close(done)
} // TestListen

func TestWalk(t *testing.T) {
	t.Parallel()
	config := RegistryConfig{
		DeclineOffer: func(offerId string) <-chan error {
			return proc.ErrorChan(nil)
		},
		TTL:           0 * time.Second,
		LingerTTL:     0 * time.Second,
		ListenerDelay: 0 * time.Second,
	}
	storage := CreateRegistry(config)
	acceptedOfferId := ""
	walked := 0
	walker1 := func(p Perishable) (bool, error) {
		walked++
		if p.Acquire() {
			acceptedOfferId = p.Details().Id.GetValue()
			return true, nil
		}
		return false, nil
	}
	// sanity check
	err := storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 0 {
		t.Fatal("walked empty storage")
	}
	if acceptedOfferId != "" {
		t.Fatal("somehow found an offer when registry was empty")
	}
	impl, ok := storage.(*offerStorage)
	if !ok {
		t.Fatal("unexpected offer storage impl")
	}
	// single offer
	ttl := 2 * time.Second
	now := time.Now()
	o := &liveOffer{&mesos.Offer{Id: util.NewOfferID("foo")}, now.Add(ttl), 0}

	impl.offers.Add(o)
	err = storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 1 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "foo" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	acceptedOfferId = ""
	err = storage.Walk(walker1)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 2 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	walker2 := func(p Perishable) (bool, error) {
		walked++
		return true, nil
	}
	err = storage.Walk(walker2)
	if err != nil {
		t.Fatalf("received impossible error %v", err)
	}
	if walked != 3 {
		t.Fatalf("walk count %d", walked)
	}
	if acceptedOfferId != "" {
		t.Fatalf("found offer %v", acceptedOfferId)
	}

	walker3 := func(p Perishable) (bool, error) {
		walked++
		return true, errors.New("baz")
	}
	err = storage.Walk(walker3)
	if err == nil {
		t.Fatal("expected error")
	}
	if walked != 4 {
		t.Fatalf("walk count %d", walked)
	}
}
<|endoftext|>


<commit_before>package smokescreen

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	log "github.com/sirupsen/logrus"
)

type Config struct {
	Ip                           string
	Port                         int
	CidrBlacklist                []net.IPNet
	CidrBlacklistExemptions      []net.IPNet
	ConnectTimeout               time.Duration
	ExitTimeout                  time.Duration
	MaintenanceFile              string
	StatsdClient                 *statsd.Client
	AllowProxyToLoopback         bool
	EgressAcl                    EgressAcl
	SupportProxyProtocol         bool
	TlsConfig                    *tls.Config
	CrlByAuthorityKeyId          map[string]*pkix.CertificateList
	RoleFromRequest              func(subject *http.Request) (string, error)
	clientCasBySubjectKeyId      map[string]*x509.Certificate
	AdditionalErrorMessageOnDeny string
	Log                          *log.Logger
	DisabledAclPolicyActions     []string

	hostExtractExpr *regexp.Regexp
}

type MissingRoleError error

// RFC 5280,  4.2.1.1
type authKeyId struct {
	Id []byte `asn1:"optional,tag:0"`
}

func (config *Config) Init() error {
	var err error

	if config.CrlByAuthorityKeyId == nil {
		config.CrlByAuthorityKeyId = make(map[string]*pkix.CertificateList)
	}
	if config.clientCasBySubjectKeyId == nil {
		config.clientCasBySubjectKeyId = make(map[string]*x509.Certificate)
	}
	if config.Log == nil {
		config.Log = log.New()
	}

	config.hostExtractExpr, err = regexp.Compile("^([^:]*)(:\\d+)?$")
	if err != nil {
		return err
	}

	// Configure RoleFromRequest for default behavior. It is ultimately meant to be replaced by the user.
	if config.TlsConfig != nil && config.TlsConfig.ClientCAs != nil { // If client certs are set, pick the CN.
		config.RoleFromRequest = func(req *http.Request) (string, error) {
			fail := func(err error) (string, error) { return "", err }
			if len(req.TLS.PeerCertificates) == 0 { // This should be impossible as long as ClientAuth is RequireAndVerifyClientCert
				fail(MissingRoleError(fmt.Errorf("fatal: No PeerCertificates")))
			}
			return req.TLS.PeerCertificates[0].Subject.CommonName, nil
		}
	} else { // Use a custom header
		config.RoleFromRequest = func(req *http.Request) (string, error) {
			fail := func(err error) (string, error) { return "", err }
			idHeader := req.Header["X-Smokescreen-Role"]
			if len(idHeader) != 1 {
				return fail(MissingRoleError(fmt.Errorf("no or multiple 'X-Smokescreen-Role' header provided")))
			}
			return idHeader[0], nil
		}
	}

	return nil
}

func (config *Config) SetupCrls(crlFiles []string) error {
	fail := func(err error) error { fmt.Print(err); return err }

	config.CrlByAuthorityKeyId = make(map[string]*pkix.CertificateList)
	config.clientCasBySubjectKeyId = make(map[string]*x509.Certificate)

	for _, crlFile := range crlFiles {
		crlBytes, err := ioutil.ReadFile(crlFile)
		if err != nil {
			return fail(err)
		}

		certList, err := x509.ParseCRL(crlBytes)
		if err != nil {
			log.Printf("Failed to parse CRL in '%s': %#v\n", crlFile, err)
		}

		// find the X509v3 Authority Key Identifier in the extensions (2.5.29.35)
		crlIssuerId := ""
		extensionOid := []int{2, 5, 29, 35}
		for _, v := range certList.TBSCertList.Extensions {
			if v.Id.Equal(extensionOid) { // Hurray, we found it
				// Boo, it's ASN.1.
				var crlAuthorityKey authKeyId
				_, err := asn1.Unmarshal(v.Value, &crlAuthorityKey)
				if err != nil {
					fmt.Printf("error: Failed to read AuthorityKey: %#v\n", err)
					continue
				}
				crlIssuerId = string(crlAuthorityKey.Id)
				break
			}
		}
		if crlIssuerId == "" {
			log.Print(fmt.Errorf("error: CRL from '%s' has no Authority Key Identifier: ignoring it\n", crlFile))
			continue
		}

		// Make sure we have a CA for this CRL or warn
		caCert, ok := config.clientCasBySubjectKeyId[crlIssuerId]

		if !ok {
			log.Printf("warn: CRL loaded for issuer '%s' but no such CA loaded: ignoring it\n", hex.EncodeToString([]byte(crlIssuerId)))
			fmt.Printf("%#v loaded certs\n", len(config.clientCasBySubjectKeyId))
			continue
		}

		// At this point, we have the CA certificate and the CRL. All that's left before evicting the CRL we currently trust is to verify the new CRL's signature
		err = caCert.CheckCRLSignature(certList)
		if err != nil {
			fmt.Printf("error: Could not trust CRL. Error during signature check: %#v\n", err)
			continue
		}

		// At this point, we have a new CRL which we trust. Let's evict the old one.
		config.CrlByAuthorityKeyId[crlIssuerId] = certList
		fmt.Printf("info: Loaded CRL for Authority ID '%s'\n", hex.EncodeToString([]byte(crlIssuerId)))
	}

	// Verify that all CAs loaded have a CRL
	for k, _ := range config.clientCasBySubjectKeyId {
		_, ok := config.CrlByAuthorityKeyId[k]
		if !ok {
			fmt.Printf("warn: no CRL loaded for Authority ID '%s'\n", hex.EncodeToString([]byte(k)))
		}
	}
	return nil
}

func (config *Config) SetupStatsd(addr, namespace string) error {
	if addr == "" {
		config.StatsdClient = nil
		return nil
	}

	track, err := statsd.New(addr)
	if err != nil {
		return err
	}
	config.StatsdClient = track

	config.StatsdClient.Namespace = namespace

	return nil
}

func (config *Config) SetupEgressAcl(aclFile string) error {
	if aclFile == "" {
		config.EgressAcl = nil
		return nil
	}

	log.Printf("Loading egress ACL from %s", aclFile)
	egressAcl, err := LoadYamlAclFromFilePath(config, aclFile)
	if err != nil {
		log.Print(err)
		return err
	}
	config.EgressAcl = egressAcl

	return nil
}

func (config *Config) SetupTls(tlsServerPemFile string, tlsClientCasFiles []string) error {
	fail := func(err error) error { return err }

	if tlsServerPemFile != "" {

		tlsConfig := tls.Config{}

		fileBytes, err := ioutil.ReadFile(tlsServerPemFile)
		if err != nil {
			return fail(err)
		}

		serverCert, err := ParsePemChain(fileBytes)
		if err != nil {
			return fail(err)
		}

		tlsConfig.Certificates = []tls.Certificate{
			serverCert,
		}

		if len(tlsClientCasFiles) != 0 {
			tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
			tlsConfig.ClientCAs = x509.NewCertPool()
			for _, clientCaFile := range tlsClientCasFiles {
				caBytes, err := ioutil.ReadFile(clientCaFile)
				if err != nil {
					return fail(err)
				}
				success := tlsConfig.ClientCAs.AppendCertsFromPEM(caBytes)
				if !success {
					fail(fmt.Errorf("Problem decoding '%s'", clientCaFile))
				}

				config.populateClientCaMap(caBytes)
			}
		}

		tlsConfig.BuildNameToCertificate()
		config.TlsConfig = &tlsConfig
	} else {
		if len(tlsClientCasFiles) != 0 {
			return fail(fmt.Errorf("It is pointless to set client CAs without setting the server's cert/key."))
		}
		config.TlsConfig = nil
	}
	return nil
}

func ParsePemChain(pemBytes []byte) (tls.Certificate, error) {
	fail := func(err error) (tls.Certificate, error) { return tls.Certificate{}, err }

	var err error
	var cert tls.Certificate

	for {
		var certDerBlock *pem.Block
		certDerBlock, pemBytes = pem.Decode(pemBytes)
		if certDerBlock == nil {
			break
		}

		if strings.HasSuffix(certDerBlock.Type, "PRIVATE KEY") {
			if cert.PrivateKey != nil {
				return fail(fmt.Errorf("Found multiple '*PRIVATE KEY's block in the provided file."))
			}
			cert.PrivateKey, err = parsePrivateKey(certDerBlock.Bytes)
			if err != nil {
				return fail(err)
			}
		} else if certDerBlock.Type == "CERTIFICATE" {
			cert.Certificate = append(cert.Certificate, certDerBlock.Bytes)
		} else {
			log.Printf("warn: Unsupported PEM block '%s'. Resolution: ignoring", certDerBlock.Type)
		}
	}

	if len(cert.Certificate) == 0 {
		return fail(fmt.Errorf("Could not find any 'CERTIFICATE' in the provided file."))
	}

	if cert.PrivateKey == nil {
		return fail(fmt.Errorf("Could not find a '*PRIVATE KEY' in the provided file."))
	}
	// We don't need to parse the IpTypePublic key for TLS, but we so do anyway
	// to check that it looks sane and matches the IpTypePrivate key.
	x509Cert, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		return fail(err)
	}

	switch pub := x509Cert.PublicKey.(type) {
	case *rsa.PublicKey:
		priv, ok := cert.PrivateKey.(*rsa.PrivateKey)
		if !ok {
			return fail(errors.New("tls: IpTypePrivate key type does not match IpTypePublic key type"))
		}
		if pub.N.Cmp(priv.N) != 0 {
			return fail(errors.New("tls: IpTypePrivate key does not match IpTypePublic key"))
		}
	case *ecdsa.PublicKey:
		priv, ok := cert.PrivateKey.(*ecdsa.PrivateKey)
		if !ok {
			return fail(errors.New("tls: IpTypePrivate key type does not match IpTypePublic key type"))
		}
		if pub.X.Cmp(priv.X) != 0 || pub.Y.Cmp(priv.Y) != 0 {
			return fail(errors.New("tls: IpTypePrivate key does not match IpTypePublic key"))
		}
	default:
		return fail(errors.New("tls: unknown IpTypePublic key algorithm"))
	}

	return cert, nil
}

// Cargoculted from pkg/tls/tls.go
// Attempt to parse the given IpTypePrivate key DER block. OpenSSL 0.9.8 generates
// PKCS#1 IpTypePrivate keys by default, while OpenSSL 1.0.0 generates PKCS#8 keys.
// OpenSSL ecparam generates SEC1 EC IpTypePrivate keys for ECDSA. We try all three.
func parsePrivateKey(der []byte) (crypto.PrivateKey, error) {
	if key, err := x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}
	if key, err := x509.ParsePKCS8PrivateKey(der); err == nil {
		switch key := key.(type) {
		case *rsa.PrivateKey, *ecdsa.PrivateKey:
			return key, nil
		default:
			return nil, errors.New("tls: found unknown IpTypePrivate key type in PKCS#8 wrapping")
		}
	}
	if key, err := x509.ParseECPrivateKey(der); err == nil {
		return key, nil
	}

	return nil, errors.New("tls: failed to parse IpTypePrivate key")
}

func (config *Config) populateClientCaMap(pemCerts []byte) (ok bool) {

	for len(pemCerts) > 0 {
		var block *pem.Block
		block, pemCerts = pem.Decode(pemCerts)
		if block == nil {
			break
		}
		if block.Type != "CERTIFICATE" || len(block.Headers) != 0 {
			continue
		}

		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			continue
		}
		fmt.Printf("info: Loaded CA with Authority ID '%s'\n", hex.EncodeToString(cert.SubjectKeyId))
		config.clientCasBySubjectKeyId[string(cert.SubjectKeyId)] = cert
		ok = true
	}
	return
}
<commit_msg>make client cert optional<commit_after>package smokescreen

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	log "github.com/sirupsen/logrus"
)

type Config struct {
	Ip                           string
	Port                         int
	CidrBlacklist                []net.IPNet
	CidrBlacklistExemptions      []net.IPNet
	ConnectTimeout               time.Duration
	ExitTimeout                  time.Duration
	MaintenanceFile              string
	StatsdClient                 *statsd.Client
	AllowProxyToLoopback         bool
	EgressAcl                    EgressAcl
	SupportProxyProtocol         bool
	TlsConfig                    *tls.Config
	CrlByAuthorityKeyId          map[string]*pkix.CertificateList
	RoleFromRequest              func(subject *http.Request) (string, error)
	clientCasBySubjectKeyId      map[string]*x509.Certificate
	AdditionalErrorMessageOnDeny string
	Log                          *log.Logger
	DisabledAclPolicyActions     []string

	hostExtractExpr *regexp.Regexp
}

type MissingRoleError error

// RFC 5280,  4.2.1.1
type authKeyId struct {
	Id []byte `asn1:"optional,tag:0"`
}

func (config *Config) Init() error {
	var err error

	if config.CrlByAuthorityKeyId == nil {
		config.CrlByAuthorityKeyId = make(map[string]*pkix.CertificateList)
	}
	if config.clientCasBySubjectKeyId == nil {
		config.clientCasBySubjectKeyId = make(map[string]*x509.Certificate)
	}
	if config.Log == nil {
		config.Log = log.New()
	}

	config.hostExtractExpr, err = regexp.Compile("^([^:]*)(:\\d+)?$")
	if err != nil {
		return err
	}

	// Configure RoleFromRequest for default behavior. It is ultimately meant to be replaced by the user.
	if config.TlsConfig != nil && config.TlsConfig.ClientCAs != nil { // If client certs are set, pick the CN.
		config.RoleFromRequest = func(req *http.Request) (string, error) {
			fail := func(err error) (string, error) { return "", err }
			if len(req.TLS.PeerCertificates) == 0 { // This should be impossible as long as ClientAuth is RequireAndVerifyClientCert
				fail(MissingRoleError(fmt.Errorf("fatal: No PeerCertificates")))
			}
			return req.TLS.PeerCertificates[0].Subject.CommonName, nil
		}
	} else { // Use a custom header
		config.RoleFromRequest = func(req *http.Request) (string, error) {
			fail := func(err error) (string, error) { return "", err }
			idHeader := req.Header["X-Smokescreen-Role"]
			if len(idHeader) != 1 {
				return fail(MissingRoleError(fmt.Errorf("no or multiple 'X-Smokescreen-Role' header provided")))
			}
			return idHeader[0], nil
		}
	}

	return nil
}

func (config *Config) SetupCrls(crlFiles []string) error {
	fail := func(err error) error { fmt.Print(err); return err }

	config.CrlByAuthorityKeyId = make(map[string]*pkix.CertificateList)
	config.clientCasBySubjectKeyId = make(map[string]*x509.Certificate)

	for _, crlFile := range crlFiles {
		crlBytes, err := ioutil.ReadFile(crlFile)
		if err != nil {
			return fail(err)
		}

		certList, err := x509.ParseCRL(crlBytes)
		if err != nil {
			log.Printf("Failed to parse CRL in '%s': %#v\n", crlFile, err)
		}

		// find the X509v3 Authority Key Identifier in the extensions (2.5.29.35)
		crlIssuerId := ""
		extensionOid := []int{2, 5, 29, 35}
		for _, v := range certList.TBSCertList.Extensions {
			if v.Id.Equal(extensionOid) { // Hurray, we found it
				// Boo, it's ASN.1.
				var crlAuthorityKey authKeyId
				_, err := asn1.Unmarshal(v.Value, &crlAuthorityKey)
				if err != nil {
					fmt.Printf("error: Failed to read AuthorityKey: %#v\n", err)
					continue
				}
				crlIssuerId = string(crlAuthorityKey.Id)
				break
			}
		}
		if crlIssuerId == "" {
			log.Print(fmt.Errorf("error: CRL from '%s' has no Authority Key Identifier: ignoring it\n", crlFile))
			continue
		}

		// Make sure we have a CA for this CRL or warn
		caCert, ok := config.clientCasBySubjectKeyId[crlIssuerId]

		if !ok {
			log.Printf("warn: CRL loaded for issuer '%s' but no such CA loaded: ignoring it\n", hex.EncodeToString([]byte(crlIssuerId)))
			fmt.Printf("%#v loaded certs\n", len(config.clientCasBySubjectKeyId))
			continue
		}

		// At this point, we have the CA certificate and the CRL. All that's left before evicting the CRL we currently trust is to verify the new CRL's signature
		err = caCert.CheckCRLSignature(certList)
		if err != nil {
			fmt.Printf("error: Could not trust CRL. Error during signature check: %#v\n", err)
			continue
		}

		// At this point, we have a new CRL which we trust. Let's evict the old one.
		config.CrlByAuthorityKeyId[crlIssuerId] = certList
		fmt.Printf("info: Loaded CRL for Authority ID '%s'\n", hex.EncodeToString([]byte(crlIssuerId)))
	}

	// Verify that all CAs loaded have a CRL
	for k, _ := range config.clientCasBySubjectKeyId {
		_, ok := config.CrlByAuthorityKeyId[k]
		if !ok {
			fmt.Printf("warn: no CRL loaded for Authority ID '%s'\n", hex.EncodeToString([]byte(k)))
		}
	}
	return nil
}

func (config *Config) SetupStatsd(addr, namespace string) error {
	if addr == "" {
		config.StatsdClient = nil
		return nil
	}

	track, err := statsd.New(addr)
	if err != nil {
		return err
	}
	config.StatsdClient = track

	config.StatsdClient.Namespace = namespace

	return nil
}

func (config *Config) SetupEgressAcl(aclFile string) error {
	if aclFile == "" {
		config.EgressAcl = nil
		return nil
	}

	log.Printf("Loading egress ACL from %s", aclFile)
	egressAcl, err := LoadYamlAclFromFilePath(config, aclFile)
	if err != nil {
		log.Print(err)
		return err
	}
	config.EgressAcl = egressAcl

	return nil
}

func (config *Config) SetupTls(tlsServerPemFile string, tlsClientCasFiles []string) error {
	fail := func(err error) error { return err }

	if tlsServerPemFile != "" {

		tlsConfig := tls.Config{}

		fileBytes, err := ioutil.ReadFile(tlsServerPemFile)
		if err != nil {
			return fail(err)
		}

		serverCert, err := ParsePemChain(fileBytes)
		if err != nil {
			return fail(err)
		}

		tlsConfig.Certificates = []tls.Certificate{
			serverCert,
		}

		if len(tlsClientCasFiles) != 0 {
			tlsConfig.ClientAuth = tls.VerifyClientCertIfGiven
			tlsConfig.ClientCAs = x509.NewCertPool()
			for _, clientCaFile := range tlsClientCasFiles {
				caBytes, err := ioutil.ReadFile(clientCaFile)
				if err != nil {
					return fail(err)
				}
				success := tlsConfig.ClientCAs.AppendCertsFromPEM(caBytes)
				if !success {
					fail(fmt.Errorf("Problem decoding '%s'", clientCaFile))
				}

				config.populateClientCaMap(caBytes)
			}
		}

		tlsConfig.BuildNameToCertificate()
		config.TlsConfig = &tlsConfig
	} else {
		if len(tlsClientCasFiles) != 0 {
			return fail(fmt.Errorf("It is pointless to set client CAs without setting the server's cert/key."))
		}
		config.TlsConfig = nil
	}
	return nil
}

func ParsePemChain(pemBytes []byte) (tls.Certificate, error) {
	fail := func(err error) (tls.Certificate, error) { return tls.Certificate{}, err }

	var err error
	var cert tls.Certificate

	for {
		var certDerBlock *pem.Block
		certDerBlock, pemBytes = pem.Decode(pemBytes)
		if certDerBlock == nil {
			break
		}

		if strings.HasSuffix(certDerBlock.Type, "PRIVATE KEY") {
			if cert.PrivateKey != nil {
				return fail(fmt.Errorf("Found multiple '*PRIVATE KEY's block in the provided file."))
			}
			cert.PrivateKey, err = parsePrivateKey(certDerBlock.Bytes)
			if err != nil {
				return fail(err)
			}
		} else if certDerBlock.Type == "CERTIFICATE" {
			cert.Certificate = append(cert.Certificate, certDerBlock.Bytes)
		} else {
			log.Printf("warn: Unsupported PEM block '%s'. Resolution: ignoring", certDerBlock.Type)
		}
	}

	if len(cert.Certificate) == 0 {
		return fail(fmt.Errorf("Could not find any 'CERTIFICATE' in the provided file."))
	}

	if cert.PrivateKey == nil {
		return fail(fmt.Errorf("Could not find a '*PRIVATE KEY' in the provided file."))
	}
	// We don't need to parse the IpTypePublic key for TLS, but we so do anyway
	// to check that it looks sane and matches the IpTypePrivate key.
	x509Cert, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		return fail(err)
	}

	switch pub := x509Cert.PublicKey.(type) {
	case *rsa.PublicKey:
		priv, ok := cert.PrivateKey.(*rsa.PrivateKey)
		if !ok {
			return fail(errors.New("tls: IpTypePrivate key type does not match IpTypePublic key type"))
		}
		if pub.N.Cmp(priv.N) != 0 {
			return fail(errors.New("tls: IpTypePrivate key does not match IpTypePublic key"))
		}
	case *ecdsa.PublicKey:
		priv, ok := cert.PrivateKey.(*ecdsa.PrivateKey)
		if !ok {
			return fail(errors.New("tls: IpTypePrivate key type does not match IpTypePublic key type"))
		}
		if pub.X.Cmp(priv.X) != 0 || pub.Y.Cmp(priv.Y) != 0 {
			return fail(errors.New("tls: IpTypePrivate key does not match IpTypePublic key"))
		}
	default:
		return fail(errors.New("tls: unknown IpTypePublic key algorithm"))
	}

	return cert, nil
}

// Cargoculted from pkg/tls/tls.go
// Attempt to parse the given IpTypePrivate key DER block. OpenSSL 0.9.8 generates
// PKCS#1 IpTypePrivate keys by default, while OpenSSL 1.0.0 generates PKCS#8 keys.
// OpenSSL ecparam generates SEC1 EC IpTypePrivate keys for ECDSA. We try all three.
func parsePrivateKey(der []byte) (crypto.PrivateKey, error) {
	if key, err := x509.ParsePKCS1PrivateKey(der); err == nil {
		return key, nil
	}
	if key, err := x509.ParsePKCS8PrivateKey(der); err == nil {
		switch key := key.(type) {
		case *rsa.PrivateKey, *ecdsa.PrivateKey:
			return key, nil
		default:
			return nil, errors.New("tls: found unknown IpTypePrivate key type in PKCS#8 wrapping")
		}
	}
	if key, err := x509.ParseECPrivateKey(der); err == nil {
		return key, nil
	}

	return nil, errors.New("tls: failed to parse IpTypePrivate key")
}

func (config *Config) populateClientCaMap(pemCerts []byte) (ok bool) {

	for len(pemCerts) > 0 {
		var block *pem.Block
		block, pemCerts = pem.Decode(pemCerts)
		if block == nil {
			break
		}
		if block.Type != "CERTIFICATE" || len(block.Headers) != 0 {
			continue
		}

		cert, err := x509.ParseCertificate(block.Bytes)
		if err != nil {
			continue
		}
		fmt.Printf("info: Loaded CA with Authority ID '%s'\n", hex.EncodeToString(cert.SubjectKeyId))
		config.clientCasBySubjectKeyId[string(cert.SubjectKeyId)] = cert
		ok = true
	}
	return
}
<|endoftext|>


<commit_before>package appcli

import (
	"io"
	"strings"

	"github.com/urfave/cli"
)

// init sets the help templates for urfave/cli.
// nolint: lll
func init() {

	cli.HelpPrinter = helpPrinter

	cli.AppHelpTemplate = `{{.Name}}{{if .Usage}} - {{.Usage}}{{end}}

Usage:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}} {{if .VisibleFlags}}[global options]{{end}}{{if .Commands}} <task> [task options]{{end}} {{if .ArgsUsage}}{{.ArgsUsage}}{{end}}{{end}}{{if .Version}}{{if not .HideVersion}}

Version:
   {{.Version}}{{end}}{{end}}{{if .Description}}

Description:
{{indent 3 .Description}}{{end}}{{if .VisibleCommands}}

Tasks:{{range .VisibleCategories}}{{if .Name}}
   {{.Name}}:{{end}}{{range .VisibleCommands}}
     {{join .Names ", "}}{{"\t"}}{{.Usage}}{{end}}{{end}}{{end}}{{if .VisibleFlags}}

Global Options:
   {{range $index, $option := .VisibleFlags}}{{if $index}}
   {{end}}{{$option}}{{end}}{{end}}{{if .Copyright}}

Copyright:
   {{.Copyright}}{{end}}
`

	cli.CommandHelpTemplate = `{{.HelpName}}{{if .Usage}} - {{.Usage}}{{end}}

Usage:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}}{{if .VisibleFlags}} [options]{{end}} {{if .ArgsUsage}}{{.ArgsUsage}}{{end}}{{end}}{{if .Category}}

Category:
   {{.Category}}{{end}}{{if .Description}}

Description:
{{indent 3 .Description}}{{end}}{{if .VisibleFlags}}

Options:
   {{range  $index, $option := .VisibleFlags}}{{if $index}}
   {{end}}{{$option}}{{end}}{{end}}
`
}

// ShowAppHelp shows the help for a given app.
func ShowAppHelp(app *cli.App) {
	app.Setup()
	cli.HelpPrinter(app.Writer, cli.AppHelpTemplate, app)
}

// helpPrinter includes the custom indent template function.
func helpPrinter(out io.Writer, templ string, data interface{}) {
	customFunc := map[string]interface{}{
		"indent": func(spaces int, text string) string {
			padding := strings.Repeat(" ", spaces)
			return padding + strings.Replace(text, "\n", "\n"+padding, -1)
		},
	}

	cli.HelpPrinterCustom(out, templ, data, customFunc)
}
<commit_msg>Remove extra spaces in task listing<commit_after>package appcli

import (
	"io"
	"strings"

	"github.com/urfave/cli"
)

// init sets the help templates for urfave/cli.
// nolint: lll
func init() {

	cli.HelpPrinter = helpPrinter

	cli.AppHelpTemplate = `{{.Name}}{{if .Usage}} - {{.Usage}}{{end}}

Usage:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}} {{if .VisibleFlags}}[global options]{{end}}{{if .Commands}} <task> [task options]{{end}} {{if .ArgsUsage}}{{.ArgsUsage}}{{end}}{{end}}{{if .Version}}{{if not .HideVersion}}

Version:
   {{.Version}}{{end}}{{end}}{{if .Description}}

Description:
{{indent 3 .Description}}{{end}}{{if .VisibleCommands}}

Tasks:{{range .VisibleCategories}}{{$categoryName := .Name}}{{if $categoryName}}
   {{$categoryName}}:{{end}}{{range .VisibleCommands}}
   {{if $categoryName}}  {{end}}{{join .Names ", "}}{{"\t"}}{{.Usage}}{{end}}{{end}}{{end}}{{if .VisibleFlags}}

Global Options:
   {{range $index, $option := .VisibleFlags}}{{if $index}}
   {{end}}{{$option}}{{end}}{{end}}{{if .Copyright}}

Copyright:
   {{.Copyright}}{{end}}
`

	cli.CommandHelpTemplate = `{{.HelpName}}{{if .Usage}} - {{.Usage}}{{end}}

Usage:
   {{if .UsageText}}{{.UsageText}}{{else}}{{.HelpName}}{{if .VisibleFlags}} [options]{{end}} {{if .ArgsUsage}}{{.ArgsUsage}}{{end}}{{end}}{{if .Category}}

Category:
   {{.Category}}{{end}}{{if .Description}}

Description:
{{indent 3 .Description}}{{end}}{{if .VisibleFlags}}

Options:
   {{range  $index, $option := .VisibleFlags}}{{if $index}}
   {{end}}{{$option}}{{end}}{{end}}
`
}

// ShowAppHelp shows the help for a given app.
func ShowAppHelp(app *cli.App) {
	app.Setup()
	cli.HelpPrinter(app.Writer, cli.AppHelpTemplate, app)
}

// helpPrinter includes the custom indent template function.
func helpPrinter(out io.Writer, templ string, data interface{}) {
	customFunc := map[string]interface{}{
		"indent": func(spaces int, text string) string {
			padding := strings.Repeat(" ", spaces)
			return padding + strings.Replace(text, "\n", "\n"+padding, -1)
		},
	}

	cli.HelpPrinterCustom(out, templ, data, customFunc)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/user"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"github.com/datawire/teleproxy/internal/pkg/nat"
	"github.com/datawire/teleproxy/internal/pkg/tpu"
	"github.com/miekg/dns"
	"golang.org/x/net/proxy"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/cache"
)

var kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")

func kubeWatch() {
	config, err := clientcmd.BuildConfigFromFlags("", *kubeconfig)
	if err != nil {
		panic(err.Error())
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	watchlist := cache.NewListWatchFromClient(clientset.Core().RESTClient(), "services", v1.NamespaceAll,
		fields.Everything())
	_, controller := cache.NewInformer(
		watchlist,
		&v1.Service{},
		time.Second * 0,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				svc := obj.(*v1.Service)
				log.Printf("ADDED: %s->%s\n", svc.Name, svc.Spec.ClusterIP)
				updateRoute(svc)
			},
			DeleteFunc: func(obj interface{}) {
				svc := obj.(*v1.Service)
				log.Printf("DELETED: %s\n", svc.Name)
				key := svc.Name + "."
				removeRoute(key)
				domainsToAddresses.Delete(key)
			},
			UpdateFunc:func(oldObj, newObj interface{}) {
				svc := newObj.(*v1.Service)
				log.Printf("CHANGED: %s->%s\n", svc.Name, svc.Spec.ClusterIP)
				updateRoute(svc)
			},
		},
	)
	stop := make(chan struct{})
	go controller.Run(stop)
}

var domainsToAddresses sync.Map
// XXX: need to do better than teleproxy
var translator = nat.NewTranslator("teleproxy")

func removeRoute(key string) {
	if old, ok := domainsToAddresses.Load(key); ok {
		translator.ClearTCP(old.(string))
	}
}

func updateRoute(svc *v1.Service) {
	if svc.Spec.ClusterIP == "None" { return }
	domainsToAddresses.Store(strings.ToLower(svc.Name + "."), svc.Spec.ClusterIP)
	translator.ForwardTCP(svc.Spec.ClusterIP, "1234")
	kickDNS()
}

type handler struct{}
func (this *handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	log.Println(r.Question[0].Qtype, "DNS request for", r.Question[0].Name)
	domain := strings.ToLower(r.Question[0].Name)
	switch r.Question[0].Qtype {
	case dns.TypeA:
		log.Println("Looking up", domain)
		address, ok := domainsToAddresses.Load(domain)
		if ok {
			log.Println("Found:", domain)
			msg := dns.Msg{}
			msg.SetReply(r)
			msg.Authoritative = true
			// mac dns seems to fallback if you don't
			// support recursion, if you have more than a
			// single dns server, this will prevent us
			// from intercepting all queries
			msg.RecursionAvailable = true
			// if we don't give back the same domain
			// requested, then mac dns seems to return an
			// nxdomain
			msg.Answer = append(msg.Answer, &dns.A{
				Hdr: dns.RR_Header{ Name: r.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60 },
				A: net.ParseIP(address.(string)),
			})
			w.WriteMsg(&msg)
			return
		}
	default:
		_, ok := domainsToAddresses.Load(domain)
		if ok {
			log.Println("Found:", domain)
			msg := dns.Msg{}
			msg.SetReply(r)
			msg.Authoritative = true
			msg.RecursionAvailable = true
			w.WriteMsg(&msg)
			log.Println("replied with empty")
			return
		}
	}
	in, err := dns.Exchange(r, *fallbackIP + ":53")
	if err != nil {
		log.Println(err)
		return
	}
	w.WriteMsg(in)
}

func dnsMain() {
	// turns out you need to listen on localhost for nat to work
	// properly for udp, otherwise you get an "unexpected source
	// blah thingy" because the dns reply packets look like they
	// are coming from the wrong place
	srv := &dns.Server{Addr: "127.0.0.1:" + strconv.Itoa(1233), Net: "udp"}
	srv.Handler = &handler{}
	if err := srv.ListenAndServe(); err != nil {
		log.Fatalf("Failed to set udp listener %s\n", err.Error())
	}
}

func rlimit() {
	var rLimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error getting rlimit:", err)
	} else {
		log.Println("Initial rlimit:", rLimit)
	}

	rLimit.Max = 999999
	rLimit.Cur = 999999
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error setting rlimit:", err)
	}

	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error getting rlimit:", err)
	} else {
		log.Println("Final rlimit", rLimit)
	}
}

var dnsIP = flag.String("dns", "", "dns ip address")
var fallbackIP = flag.String("fallback", "", "dns fallback")

func main() {
	flag.Parse()

	if *kubeconfig == "" {
		*kubeconfig = os.Getenv("KUBECONFIG")
	}

	if *kubeconfig == "" {
		current, err := user.Current()
		if err != nil { panic(err) }
		home := current.HomeDir
		*kubeconfig = filepath.Join(home, ".kube/config")
	}

	if *dnsIP == "" {
		dat, err := ioutil.ReadFile("/etc/resolv.conf")
		if err != nil { panic(err) }
		for _, line := range strings.Split(string(dat), "\n") {
			if strings.Contains(line, "nameserver") {
				fields := strings.Fields(line)
				*dnsIP = fields[1]
				log.Printf("Automatically set -dns to %v", *dnsIP)
				break
			}
		}
	}

	if *dnsIP == "" {
		panic("couldn't determine dns ip from /etc/resolv.conf")
	}

	if *fallbackIP == "" {
		if *dnsIP == "8.8.8.8" {
			*fallbackIP = "8.8.4.4"
		} else {
			*fallbackIP = "8.8.8.8"
		}
	}

	if *fallbackIP == *dnsIP {
		panic("if your fallbackIP and your dnsIP are the same, you will have a dns loop")
	}

	if runtime.GOOS == "darwin" {
		ifaces, _ := getIfaces()
		for _, iface := range ifaces {
			// setup dns search path
			domain, _ := getSearchDomains(iface)
			setSearchDomains(iface, ".")
			// restore dns search path
			defer setSearchDomains(iface, domain)
		}
	}

	rlimit()

	kubeWatch()
	go dnsMain()

	ln, err := net.Listen("tcp", ":1234")
	if err != nil {
		log.Println(err)
		return
	}

	translator.Enable()
	translator.ForwardUDP(*dnsIP, "1233")
	defer translator.Disable()

	apply := tpu.Keepalive(1, `
---
apiVersion: v1
kind: Pod
metadata:
  name: teleproxy
spec:
  containers:
  - name: proxy
    image: datawire/telepresence-k8s:0.75
    ports:
    - protocol: TCP
      containerPort: 8022
`, "kubectl", "apply", "-f", "-")
	apply.Wait()

	pf := tpu.Keepalive(0, "", "kubectl", "port-forward", "pod/teleproxy", "8022")
        defer pf.Shutdown()
	// XXX: probably need some kind of keepalive check for ssh, first
	// curl after wakeup seems to trigger detection of death
//	ssh := tpu.Keepalive(0, "", "ssh", "-D", "localhost:1080", "-L", "localhost:9050:localhost:9050", "-C", "-N", "-oConnectTimeout=5", "-oExitOnForwardFailure=yes",
	ssh := tpu.Keepalive(0, "", "ssh", "-D", "localhost:1080", "-C", "-N", "-oConnectTimeout=5", "-oExitOnForwardFailure=yes",
		"-oStrictHostKeyChecking=no", "-oUserKnownHostsFile=/dev/null", "telepresence@localhost", "-p", "8022")
	defer ssh.Shutdown()

	limit := 10000
	log.Printf("Listening (limit %v)...", limit)
	go func() {
		sem := tpu.NewSemaphore(limit)
		for {
			conn, err := ln.Accept();
			if err != nil {
				log.Println(err)
			} else {
				switch conn.(type) {
				case *net.TCPConn:
					log.Println("AVAILABLE:", len(sem))
					sem.Acquire()
					go handleConnection(conn.(*net.TCPConn), sem)
				default:
					log.Println("Don't know how to handle conn:", conn)
				}
			}
		}
	}()

	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	defer kickDNS()
}

func handleConnection(conn *net.TCPConn, sem tpu.Semaphore) {
	defer sem.Release()
	// hmm, we may not actually need to get the original destination,
	// we could just forward each ip to a unique port and either
	// listen on that port or run port-forward
	_, host, err := translator.GetOriginalDst(conn)
	if err != nil {
		log.Println("GetOriginalDst:", err)
		return
	}

	log.Println("CONNECT:", conn.RemoteAddr(), host)

	// setting up an ssh tunnel with dynamic socks proxy at this end
	// seems faster than connecting directly to a socks proxy
	dialer, err := proxy.SOCKS5("tcp", "localhost:1080", nil, proxy.Direct)
//	dialer, err := proxy.SOCKS5("tcp", "localhost:9050", nil, proxy.Direct)
	if err != nil {
		log.Println(err)
		conn.Close()
		return
	}

	_proxy, err := dialer.Dial("tcp", host)
	if err != nil {
		log.Println(err)
		conn.Close()
		return
	}
	proxy := _proxy.(*net.TCPConn)

	done := tpu.NewLatch(2)

	go pipe(conn, proxy, done)
	go pipe(proxy, conn, done)

	done.Wait()
}

func pipe(from, to *net.TCPConn, done tpu.Latch) {
	defer func() {
		log.Println("CLOSED WRITE:", to.RemoteAddr())
		to.CloseWrite()
	}()
	defer func() {
		log.Println("CLOSED READ:", from.RemoteAddr())
		from.CloseRead()
	}()
	defer done.Notify()

	const size = 64*1024
	var buf [size]byte
	for {
		n, err := from.Read(buf[0:size])
		if err != nil {
			if err != io.EOF {
				log.Println(err)
			}
			break
		} else {
			_, err := to.Write(buf[0:n])

			if err != nil {
				log.Println(err)
				break
			}
		}
	}
}

func getPIDs() (pids []int, err error) {
	cmd := exec.Command("ps", "-axo", "pid=,command=")
	out, err := cmd.CombinedOutput()
	if err != nil { return }
	if !cmd.ProcessState.Success() {
		err = fmt.Errorf("%s", out)
		return
	}

	for _, line := range strings.Split(string(out), "\n") {
		if strings.Contains(strings.ToLower(line), "mdnsresponder") {
			parts := strings.Fields(line)
			var pid int
			pid, err = strconv.Atoi(parts[0])
			if err != nil { return }
			pids = append(pids, pid)
		}
	}

	return
}

func kickDNS() {
	pids, err := getPIDs()
	if err != nil {
		log.Println(err)
		return
	}

	log.Printf("Kicking DNS: %v\n", pids)
	for _, pid := range pids {
		proc, err := os.FindProcess(pid)
		if err != nil { log.Println(err) }
		err = proc.Signal(syscall.SIGHUP)
		if err != nil { log.Println(err) }
	}
}

func shell(command string) (result string, err error) {
	log.Println(command)
	cmd := exec.Command("sh", "-c", command)
	out, err := cmd.CombinedOutput()
	if err != nil {
		log.Println(err)
		return
	}
	log.Printf("%s", out)
	result = string(out)
	return
}

func getIfaces() (ifaces []string, err error) {
	lines, err := shell("networksetup -listallnetworkservices | fgrep -v '*'")
	if err != nil { return }
	for _, line := range strings.Split(lines, "\n") {
		line = strings.TrimSpace(line)
		if len(line) > 0 {
			ifaces = append(ifaces, line)
		}
	}
	return
}

func getSearchDomains(iface string) (domains string, err error) {
	domains, err = shell(fmt.Sprintf("networksetup -getsearchdomains '%s'", iface))
	domains = strings.TrimSpace(domains)
	return
}

func setSearchDomains(iface, domains string) (err error) {
	_, err = shell(fmt.Sprintf("networksetup -setsearchdomains '%s' '%s'", iface, domains))
	return
}
<commit_msg>Pass kubeconfig flag through to kubectl. Add label to pod.<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/user"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
	"github.com/datawire/teleproxy/internal/pkg/nat"
	"github.com/datawire/teleproxy/internal/pkg/tpu"
	"github.com/miekg/dns"
	"golang.org/x/net/proxy"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/cache"
)

var kubeconfig = flag.String("kubeconfig", "", "absolute path to the kubeconfig file")

func kubeWatch() {
	config, err := clientcmd.BuildConfigFromFlags("", *kubeconfig)
	if err != nil {
		panic(err.Error())
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	watchlist := cache.NewListWatchFromClient(clientset.Core().RESTClient(), "services", v1.NamespaceAll,
		fields.Everything())
	_, controller := cache.NewInformer(
		watchlist,
		&v1.Service{},
		time.Second * 0,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				svc := obj.(*v1.Service)
				log.Printf("ADDED: %s->%s\n", svc.Name, svc.Spec.ClusterIP)
				updateRoute(svc)
			},
			DeleteFunc: func(obj interface{}) {
				svc := obj.(*v1.Service)
				log.Printf("DELETED: %s\n", svc.Name)
				key := svc.Name + "."
				removeRoute(key)
				domainsToAddresses.Delete(key)
			},
			UpdateFunc:func(oldObj, newObj interface{}) {
				svc := newObj.(*v1.Service)
				log.Printf("CHANGED: %s->%s\n", svc.Name, svc.Spec.ClusterIP)
				updateRoute(svc)
			},
		},
	)
	stop := make(chan struct{})
	go controller.Run(stop)
}

var domainsToAddresses sync.Map
// XXX: need to do better than teleproxy
var translator = nat.NewTranslator("teleproxy")

func removeRoute(key string) {
	if old, ok := domainsToAddresses.Load(key); ok {
		translator.ClearTCP(old.(string))
	}
}

func updateRoute(svc *v1.Service) {
	if svc.Spec.ClusterIP == "None" { return }
	domainsToAddresses.Store(strings.ToLower(svc.Name + "."), svc.Spec.ClusterIP)
	translator.ForwardTCP(svc.Spec.ClusterIP, "1234")
	kickDNS()
}

type handler struct{}
func (this *handler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	log.Println(r.Question[0].Qtype, "DNS request for", r.Question[0].Name)
	domain := strings.ToLower(r.Question[0].Name)
	switch r.Question[0].Qtype {
	case dns.TypeA:
		log.Println("Looking up", domain)
		address, ok := domainsToAddresses.Load(domain)
		if ok {
			log.Println("Found:", domain)
			msg := dns.Msg{}
			msg.SetReply(r)
			msg.Authoritative = true
			// mac dns seems to fallback if you don't
			// support recursion, if you have more than a
			// single dns server, this will prevent us
			// from intercepting all queries
			msg.RecursionAvailable = true
			// if we don't give back the same domain
			// requested, then mac dns seems to return an
			// nxdomain
			msg.Answer = append(msg.Answer, &dns.A{
				Hdr: dns.RR_Header{ Name: r.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60 },
				A: net.ParseIP(address.(string)),
			})
			w.WriteMsg(&msg)
			return
		}
	default:
		_, ok := domainsToAddresses.Load(domain)
		if ok {
			log.Println("Found:", domain)
			msg := dns.Msg{}
			msg.SetReply(r)
			msg.Authoritative = true
			msg.RecursionAvailable = true
			w.WriteMsg(&msg)
			log.Println("replied with empty")
			return
		}
	}
	in, err := dns.Exchange(r, *fallbackIP + ":53")
	if err != nil {
		log.Println(err)
		return
	}
	w.WriteMsg(in)
}

func dnsMain() {
	// turns out you need to listen on localhost for nat to work
	// properly for udp, otherwise you get an "unexpected source
	// blah thingy" because the dns reply packets look like they
	// are coming from the wrong place
	srv := &dns.Server{Addr: "127.0.0.1:" + strconv.Itoa(1233), Net: "udp"}
	srv.Handler = &handler{}
	if err := srv.ListenAndServe(); err != nil {
		log.Fatalf("Failed to set udp listener %s\n", err.Error())
	}
}

func rlimit() {
	var rLimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error getting rlimit:", err)
	} else {
		log.Println("Initial rlimit:", rLimit)
	}

	rLimit.Max = 999999
	rLimit.Cur = 999999
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error setting rlimit:", err)
	}

	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		log.Println("Error getting rlimit:", err)
	} else {
		log.Println("Final rlimit", rLimit)
	}
}

var dnsIP = flag.String("dns", "", "dns ip address")
var fallbackIP = flag.String("fallback", "", "dns fallback")

func main() {
	flag.Parse()

	if *kubeconfig == "" {
		*kubeconfig = os.Getenv("KUBECONFIG")
	}

	if *kubeconfig == "" {
		current, err := user.Current()
		if err != nil { panic(err) }
		home := current.HomeDir
		*kubeconfig = filepath.Join(home, ".kube/config")
	}

	if *dnsIP == "" {
		dat, err := ioutil.ReadFile("/etc/resolv.conf")
		if err != nil { panic(err) }
		for _, line := range strings.Split(string(dat), "\n") {
			if strings.Contains(line, "nameserver") {
				fields := strings.Fields(line)
				*dnsIP = fields[1]
				log.Printf("Automatically set -dns to %v", *dnsIP)
				break
			}
		}
	}

	if *dnsIP == "" {
		panic("couldn't determine dns ip from /etc/resolv.conf")
	}

	if *fallbackIP == "" {
		if *dnsIP == "8.8.8.8" {
			*fallbackIP = "8.8.4.4"
		} else {
			*fallbackIP = "8.8.8.8"
		}
	}

	if *fallbackIP == *dnsIP {
		panic("if your fallbackIP and your dnsIP are the same, you will have a dns loop")
	}

	if runtime.GOOS == "darwin" {
		ifaces, _ := getIfaces()
		for _, iface := range ifaces {
			// setup dns search path
			domain, _ := getSearchDomains(iface)
			setSearchDomains(iface, ".")
			// restore dns search path
			defer setSearchDomains(iface, domain)
		}
	}

	rlimit()

	kubeWatch()
	go dnsMain()

	ln, err := net.Listen("tcp", ":1234")
	if err != nil {
		log.Println(err)
		return
	}

	translator.Enable()
	translator.ForwardUDP(*dnsIP, "1233")
	defer translator.Disable()

	apply := tpu.Keepalive(1, `
---
apiVersion: v1
kind: Pod
metadata:
  name: teleproxy
  labels:
    name: teleproxy
spec:
  containers:
  - name: proxy
    image: datawire/telepresence-k8s:0.75
    ports:
    - protocol: TCP
      containerPort: 8022
`, "kubectl", "--kubeconfig", *kubeconfig, "apply", "-f", "-")
	apply.Wait()

	pf := tpu.Keepalive(0, "", "kubectl", "--kubeconfig", *kubeconfig, "port-forward", "pod/teleproxy", "8022")
        defer pf.Shutdown()
	// XXX: probably need some kind of keepalive check for ssh, first
	// curl after wakeup seems to trigger detection of death
//	ssh := tpu.Keepalive(0, "", "ssh", "-D", "localhost:1080", "-L", "localhost:9050:localhost:9050", "-C", "-N", "-oConnectTimeout=5", "-oExitOnForwardFailure=yes",
	ssh := tpu.Keepalive(0, "", "ssh", "-D", "localhost:1080", "-C", "-N", "-oConnectTimeout=5", "-oExitOnForwardFailure=yes",
		"-oStrictHostKeyChecking=no", "-oUserKnownHostsFile=/dev/null", "telepresence@localhost", "-p", "8022")
	defer ssh.Shutdown()

	limit := 10000
	log.Printf("Listening (limit %v)...", limit)
	go func() {
		sem := tpu.NewSemaphore(limit)
		for {
			conn, err := ln.Accept();
			if err != nil {
				log.Println(err)
			} else {
				switch conn.(type) {
				case *net.TCPConn:
					log.Println("AVAILABLE:", len(sem))
					sem.Acquire()
					go handleConnection(conn.(*net.TCPConn), sem)
				default:
					log.Println("Don't know how to handle conn:", conn)
				}
			}
		}
	}()

	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	defer kickDNS()
}

func handleConnection(conn *net.TCPConn, sem tpu.Semaphore) {
	defer sem.Release()
	// hmm, we may not actually need to get the original destination,
	// we could just forward each ip to a unique port and either
	// listen on that port or run port-forward
	_, host, err := translator.GetOriginalDst(conn)
	if err != nil {
		log.Println("GetOriginalDst:", err)
		return
	}

	log.Println("CONNECT:", conn.RemoteAddr(), host)

	// setting up an ssh tunnel with dynamic socks proxy at this end
	// seems faster than connecting directly to a socks proxy
	dialer, err := proxy.SOCKS5("tcp", "localhost:1080", nil, proxy.Direct)
//	dialer, err := proxy.SOCKS5("tcp", "localhost:9050", nil, proxy.Direct)
	if err != nil {
		log.Println(err)
		conn.Close()
		return
	}

	_proxy, err := dialer.Dial("tcp", host)
	if err != nil {
		log.Println(err)
		conn.Close()
		return
	}
	proxy := _proxy.(*net.TCPConn)

	done := tpu.NewLatch(2)

	go pipe(conn, proxy, done)
	go pipe(proxy, conn, done)

	done.Wait()
}

func pipe(from, to *net.TCPConn, done tpu.Latch) {
	defer func() {
		log.Println("CLOSED WRITE:", to.RemoteAddr())
		to.CloseWrite()
	}()
	defer func() {
		log.Println("CLOSED READ:", from.RemoteAddr())
		from.CloseRead()
	}()
	defer done.Notify()

	const size = 64*1024
	var buf [size]byte
	for {
		n, err := from.Read(buf[0:size])
		if err != nil {
			if err != io.EOF {
				log.Println(err)
			}
			break
		} else {
			_, err := to.Write(buf[0:n])

			if err != nil {
				log.Println(err)
				break
			}
		}
	}
}

func getPIDs() (pids []int, err error) {
	cmd := exec.Command("ps", "-axo", "pid=,command=")
	out, err := cmd.CombinedOutput()
	if err != nil { return }
	if !cmd.ProcessState.Success() {
		err = fmt.Errorf("%s", out)
		return
	}

	for _, line := range strings.Split(string(out), "\n") {
		if strings.Contains(strings.ToLower(line), "mdnsresponder") {
			parts := strings.Fields(line)
			var pid int
			pid, err = strconv.Atoi(parts[0])
			if err != nil { return }
			pids = append(pids, pid)
		}
	}

	return
}

func kickDNS() {
	pids, err := getPIDs()
	if err != nil {
		log.Println(err)
		return
	}

	log.Printf("Kicking DNS: %v\n", pids)
	for _, pid := range pids {
		proc, err := os.FindProcess(pid)
		if err != nil { log.Println(err) }
		err = proc.Signal(syscall.SIGHUP)
		if err != nil { log.Println(err) }
	}
}

func shell(command string) (result string, err error) {
	log.Println(command)
	cmd := exec.Command("sh", "-c", command)
	out, err := cmd.CombinedOutput()
	if err != nil {
		log.Println(err)
		return
	}
	log.Printf("%s", out)
	result = string(out)
	return
}

func getIfaces() (ifaces []string, err error) {
	lines, err := shell("networksetup -listallnetworkservices | fgrep -v '*'")
	if err != nil { return }
	for _, line := range strings.Split(lines, "\n") {
		line = strings.TrimSpace(line)
		if len(line) > 0 {
			ifaces = append(ifaces, line)
		}
	}
	return
}

func getSearchDomains(iface string) (domains string, err error) {
	domains, err = shell(fmt.Sprintf("networksetup -getsearchdomains '%s'", iface))
	domains = strings.TrimSpace(domains)
	return
}

func setSearchDomains(iface, domains string) (err error) {
	_, err = shell(fmt.Sprintf("networksetup -setsearchdomains '%s' '%s'", iface, domains))
	return
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"testing"

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

func TestJSONStoreNextID(t *testing.T) {
	t.Parallel()

	j := JSONStore{CurID: 0}

	i := j.NextID()
	assert.Equal(t, i, uint(1))
	assert.Equal(t, j.CurID, uint(1))
}

func TestJSONStoreCommit(t *testing.T) {
	fname := "test.json"
	defer os.Remove(fname)

	j, err := NewJSONStore(fname)
	assert.Nil(t, err)

	err = j.Commit()
	assert.Nil(t, err)

	bytes, err := ioutil.ReadFile(fname)
	assert.Nil(t, err)
	assert.True(t, len(bytes) > 0)
}

func TestJSONStoreSaveTask(t *testing.T) {
	t.Parallel()

	j, err := NewJSONStore("savetask.json")
	assert.Nil(t, err)
	task := Task{}

	err = j.SaveTask(&task)
	assert.Nil(t, err)
	assert.NotEqual(t, task.ID, "")

	task2 := j.Tasks[task.ID]
	assert.Equal(t, task.ID, task2.ID)
}

func TestJSONStoreGetTask(t *testing.T) {
	t.Parallel()

	j, err := NewJSONStore("gettask.json")
	assert.Nil(t, err)

	task := &Task{ID: "foo"}
	j.Tasks[task.ID] = task

	// a task that exists
	task2, err := j.GetTask(task.ID)
	assert.Nil(t, err)
	assert.Equal(t, task.ID, task2.ID)

	// a task that doesn't exist
	task3, err := j.GetTask("bar")
	assert.Equal(t, err, NoSuchTask)
	assert.Nil(t, task3)
}

func TestJSONStoreFilterTasks(t *testing.T) {
	j, err := NewJSONStore("filter.json")
	assert.Nil(t, err)

	foo := &Task{ID: "foo"}
	bar := &Task{ID: "bar"}
	j.SaveTask(foo)
	j.SaveTask(bar)

	// accept everything
	all := func(t *Task) bool { return true }

	result, err := j.FilterTasks(all)
	assert.Nil(t, err)
	in := func(t *testing.T, task *Task, cont []*Task) {
		for _, x := range cont {
			if x.ID == task.ID {
				return
			}
		}
		t.Errorf("Task with id %s not found in %+v", task.ID, cont)
	}
	in(t, foo, result)
	in(t, bar, result)

	// only accept Task with ID "foo"
	some := func(t *Task) bool { return t.ID == "foo" }

	result, err = j.FilterTasks(some)
	assert.Nil(t, err)
	assert.Equal(t, []*Task{foo}, result)
}

func TestJSONStoreSetMeta(t *testing.T) {
	j, err := NewJSONStore("setmeta.json")
	assert.Nil(t, err)

	k, v := "foo", "bar"
	err = j.SetMeta(k, v)
	assert.Nil(t, err)

	// make sure it actually set
	v2, ok := j.Meta[k]
	assert.True(t, ok)
	assert.Equal(t, v, v2)
}

func TestJSONStoreGetMeta(t *testing.T) {
	j, err := NewJSONStore("getmeta.json")
	assert.Nil(t, err)

	k, v := "foo", "bar"
	j.Meta[k] = v

	// get a value that exists
	v2, err := j.GetMeta(k)
	assert.Nil(t, err)
	assert.Equal(t, v, v2)

	// get a value that doesn't exist
	_, err = j.GetMeta("whatever")
	assert.Equal(t, err, NoSuchKey)
}

func TestJSONStoreImplements(t *testing.T) {
	t.Parallel()

	j := new(JSONStore)
	assert.Implements(t, (*Storage)(nil), j)
	assert.Implements(t, (*TaskStore)(nil), j)
	assert.Implements(t, (*MetaStore)(nil), j)
	assert.Implements(t, (*MetaTaskStore)(nil), j)
}
<commit_msg>mark tests as Parallel<commit_after>package main

import (
	"io/ioutil"
	"os"
	"testing"

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

func TestJSONStoreNextID(t *testing.T) {
	t.Parallel()

	j := JSONStore{CurID: 0}

	i := j.NextID()
	assert.Equal(t, i, uint(1))
	assert.Equal(t, j.CurID, uint(1))
}

func TestJSONStoreCommit(t *testing.T) {
	fname := "test.json"
	defer os.Remove(fname)

	j, err := NewJSONStore(fname)
	assert.Nil(t, err)

	err = j.Commit()
	assert.Nil(t, err)

	bytes, err := ioutil.ReadFile(fname)
	assert.Nil(t, err)
	assert.True(t, len(bytes) > 0)
}

func TestJSONStoreSaveTask(t *testing.T) {
	t.Parallel()

	j, err := NewJSONStore("savetask.json")
	assert.Nil(t, err)
	task := Task{}

	err = j.SaveTask(&task)
	assert.Nil(t, err)
	assert.NotEqual(t, task.ID, "")

	task2 := j.Tasks[task.ID]
	assert.Equal(t, task.ID, task2.ID)
}

func TestJSONStoreGetTask(t *testing.T) {
	t.Parallel()

	j, err := NewJSONStore("gettask.json")
	assert.Nil(t, err)

	task := &Task{ID: "foo"}
	j.Tasks[task.ID] = task

	// a task that exists
	task2, err := j.GetTask(task.ID)
	assert.Nil(t, err)
	assert.Equal(t, task.ID, task2.ID)

	// a task that doesn't exist
	task3, err := j.GetTask("bar")
	assert.Equal(t, err, NoSuchTask)
	assert.Nil(t, task3)
}

func TestJSONStoreFilterTasks(t *testing.T) {
	t.Parallel()
	j, err := NewJSONStore("filter.json")
	assert.Nil(t, err)

	foo := &Task{ID: "foo"}
	bar := &Task{ID: "bar"}
	j.SaveTask(foo)
	j.SaveTask(bar)

	// accept everything
	all := func(t *Task) bool { return true }

	result, err := j.FilterTasks(all)
	assert.Nil(t, err)
	in := func(t *testing.T, task *Task, cont []*Task) {
		for _, x := range cont {
			if x.ID == task.ID {
				return
			}
		}
		t.Errorf("Task with id %s not found in %+v", task.ID, cont)
	}
	in(t, foo, result)
	in(t, bar, result)

	// only accept Task with ID "foo"
	some := func(t *Task) bool { return t.ID == "foo" }

	result, err = j.FilterTasks(some)
	assert.Nil(t, err)
	assert.Equal(t, []*Task{foo}, result)
}

func TestJSONStoreSetMeta(t *testing.T) {
	t.Parallel()
	j, err := NewJSONStore("setmeta.json")
	assert.Nil(t, err)

	k, v := "foo", "bar"
	err = j.SetMeta(k, v)
	assert.Nil(t, err)

	// make sure it actually set
	v2, ok := j.Meta[k]
	assert.True(t, ok)
	assert.Equal(t, v, v2)
}

func TestJSONStoreGetMeta(t *testing.T) {
	t.Parallel()
	j, err := NewJSONStore("getmeta.json")
	assert.Nil(t, err)

	k, v := "foo", "bar"
	j.Meta[k] = v

	// get a value that exists
	v2, err := j.GetMeta(k)
	assert.Nil(t, err)
	assert.Equal(t, v, v2)

	// get a value that doesn't exist
	_, err = j.GetMeta("whatever")
	assert.Equal(t, err, NoSuchKey)
}

func TestJSONStoreImplements(t *testing.T) {
	t.Parallel()

	j := new(JSONStore)
	assert.Implements(t, (*Storage)(nil), j)
	assert.Implements(t, (*TaskStore)(nil), j)
	assert.Implements(t, (*MetaStore)(nil), j)
	assert.Implements(t, (*MetaTaskStore)(nil), j)
}
<|endoftext|>