<commit_before>/*
 *
 * Copyright 2014, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

package grpc

import (
	"bytes"
	"io"
	"math"
	"reflect"
	"testing"

	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	perfpb "google.golang.org/grpc/test/codec_perf"
	"google.golang.org/grpc/transport"
)

func TestSimpleParsing(t *testing.T) {
	bigMsg := bytes.Repeat([]byte{'x'}, 1<<24)
	for _, test := range []struct {
		// input
		p []byte
		// outputs
		err error
		b   []byte
		pt  payloadFormat
	}{
		{nil, io.EOF, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 0}, nil, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 1, 'a'}, nil, []byte{'a'}, compressionNone},
		{[]byte{1, 0}, io.ErrUnexpectedEOF, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 10, 'a'}, io.ErrUnexpectedEOF, nil, compressionNone},
		// Check that messages with length >= 2^24 are parsed.
		{append([]byte{0, 1, 0, 0, 0}, bigMsg...), nil, bigMsg, compressionNone},
	} {
		buf := bytes.NewReader(test.p)
		parser := &parser{r: buf}
		pt, b, err := parser.recvMsg(math.MaxInt32)
		if err != test.err || !bytes.Equal(b, test.b) || pt != test.pt {
			t.Fatalf("parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, %v", test.p, pt, b, err, test.pt, test.b, test.err)
		}
	}
}

func TestMultipleParsing(t *testing.T) {
	// Set a byte stream consists of 3 messages with their headers.
	p := []byte{0, 0, 0, 0, 1, 'a', 0, 0, 0, 0, 2, 'b', 'c', 0, 0, 0, 0, 1, 'd'}
	b := bytes.NewReader(p)
	parser := &parser{r: b}

	wantRecvs := []struct {
		pt   payloadFormat
		data []byte
	}{
		{compressionNone, []byte("a")},
		{compressionNone, []byte("bc")},
		{compressionNone, []byte("d")},
	}
	for i, want := range wantRecvs {
		pt, data, err := parser.recvMsg(math.MaxInt32)
		if err != nil || pt != want.pt || !reflect.DeepEqual(data, want.data) {
			t.Fatalf("after %d calls, parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, <nil>",
				i, p, pt, data, err, want.pt, want.data)
		}
	}

	pt, data, err := parser.recvMsg(math.MaxInt32)
	if err != io.EOF {
		t.Fatalf("after %d recvMsgs calls, parser{%v}.recvMsg(_) = %v, %v, %v\nwant _, _, %v",
			len(wantRecvs), p, pt, data, err, io.EOF)
	}
}

func TestEncode(t *testing.T) {
	for _, test := range []struct {
		// input
		msg proto.Message
		cp  Compressor
		// outputs
		b   []byte
		err error
	}{
		{nil, nil, []byte{0, 0, 0, 0, 0}, nil},
	} {
		b, err := encode(protoCodec{}, test.msg, nil, nil, nil)
		if err != test.err || !bytes.Equal(b, test.b) {
			t.Fatalf("encode(_, _, %v, _) = %v, %v\nwant %v, %v", test.cp, b, err, test.b, test.err)
		}
	}
}

func TestCompress(t *testing.T) {
	for _, test := range []struct {
		// input
		data []byte
		cp   Compressor
		dc   Decompressor
		// outputs
		err error
	}{
		{make([]byte, 1024), NewGZIPCompressor(), NewGZIPDecompressor(), nil},
	} {
		b := new(bytes.Buffer)
		if err := test.cp.Do(b, test.data); err != test.err {
			t.Fatalf("Compressor.Do(_, %v) = %v, want %v", test.data, err, test.err)
		}
		if b.Len() >= len(test.data) {
			t.Fatalf("The compressor fails to compress data.")
		}
		if p, err := test.dc.Do(b); err != nil || !bytes.Equal(test.data, p) {
			t.Fatalf("Decompressor.Do(%v) = %v, %v, want %v, <nil>", b, p, err, test.data)
		}
	}
}

func TestToRPCErr(t *testing.T) {
	for _, test := range []struct {
		// input
		errIn error
		// outputs
		errOut error
	}{
		{transport.StreamError{Code: codes.Unknown, Desc: ""}, status.Error(codes.Unknown, "")},
		{transport.ErrConnClosing, status.Error(codes.Internal, transport.ErrConnClosing.Desc)},
	} {
		err := toRPCErr(test.errIn)
		if _, ok := status.FromError(err); !ok {
			t.Fatalf("toRPCErr{%v} returned type %T, want %T", test.errIn, err, status.Error(codes.Unknown, ""))
		}
		if !reflect.DeepEqual(err, test.errOut) {
			t.Fatalf("toRPCErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
		}
	}
}

// bmEncode benchmarks encoding a Protocol Buffer message containing mSize
// bytes.
func bmEncode(b *testing.B, mSize int) {
	msg := &perfpb.Buffer{Body: make([]byte, mSize)}
	encoded, _ := encode(protoCodec{}, msg, nil, nil, nil)
	encodedSz := int64(len(encoded))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		encode(protoCodec{}, msg, nil, nil, nil)
	}
	b.SetBytes(encodedSz)
}

func BenchmarkEncode1B(b *testing.B) {
	bmEncode(b, 1)
}

func BenchmarkEncode1KiB(b *testing.B) {
	bmEncode(b, 1024)
}

func BenchmarkEncode8KiB(b *testing.B) {
	bmEncode(b, 8*1024)
}

func BenchmarkEncode64KiB(b *testing.B) {
	bmEncode(b, 64*1024)
}

func BenchmarkEncode512KiB(b *testing.B) {
	bmEncode(b, 512*1024)
}

func BenchmarkEncode1MiB(b *testing.B) {
	bmEncode(b, 1024*1024)
}

// bmCompressor benchmarks a compressor of a Protocol Buffer message containing
// mSize bytes.
func bmCompressor(b *testing.B, mSize int, cp Compressor) {
	payload := make([]byte, mSize)
	cBuf := bytes.NewBuffer(make([]byte, mSize))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		cp.Do(cBuf, payload)
		cBuf.Reset()
	}
}

func BenchmarkGZIPCompressor1B(b *testing.B) {
	bmCompressor(b, 1, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1KiB(b *testing.B) {
	bmCompressor(b, 1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor8KiB(b *testing.B) {
	bmCompressor(b, 8*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor64KiB(b *testing.B) {
	bmCompressor(b, 64*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor512KiB(b *testing.B) {
	bmCompressor(b, 512*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1MiB(b *testing.B) {
	bmCompressor(b, 1024*1024, NewGZIPCompressor())
}
<commit_msg>Test for EmptyCallOption<commit_after>/*
 *
 * Copyright 2014, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

package grpc

import (
	"bytes"
	"io"
	"math"
	"reflect"
	"testing"

	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	perfpb "google.golang.org/grpc/test/codec_perf"
	"google.golang.org/grpc/transport"
)

var _ CallOption = EmptyCallOption{} // ensure EmptyCallOption implements the interface

func TestSimpleParsing(t *testing.T) {
	bigMsg := bytes.Repeat([]byte{'x'}, 1<<24)
	for _, test := range []struct {
		// input
		p []byte
		// outputs
		err error
		b   []byte
		pt  payloadFormat
	}{
		{nil, io.EOF, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 0}, nil, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 1, 'a'}, nil, []byte{'a'}, compressionNone},
		{[]byte{1, 0}, io.ErrUnexpectedEOF, nil, compressionNone},
		{[]byte{0, 0, 0, 0, 10, 'a'}, io.ErrUnexpectedEOF, nil, compressionNone},
		// Check that messages with length >= 2^24 are parsed.
		{append([]byte{0, 1, 0, 0, 0}, bigMsg...), nil, bigMsg, compressionNone},
	} {
		buf := bytes.NewReader(test.p)
		parser := &parser{r: buf}
		pt, b, err := parser.recvMsg(math.MaxInt32)
		if err != test.err || !bytes.Equal(b, test.b) || pt != test.pt {
			t.Fatalf("parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, %v", test.p, pt, b, err, test.pt, test.b, test.err)
		}
	}
}

func TestMultipleParsing(t *testing.T) {
	// Set a byte stream consists of 3 messages with their headers.
	p := []byte{0, 0, 0, 0, 1, 'a', 0, 0, 0, 0, 2, 'b', 'c', 0, 0, 0, 0, 1, 'd'}
	b := bytes.NewReader(p)
	parser := &parser{r: b}

	wantRecvs := []struct {
		pt   payloadFormat
		data []byte
	}{
		{compressionNone, []byte("a")},
		{compressionNone, []byte("bc")},
		{compressionNone, []byte("d")},
	}
	for i, want := range wantRecvs {
		pt, data, err := parser.recvMsg(math.MaxInt32)
		if err != nil || pt != want.pt || !reflect.DeepEqual(data, want.data) {
			t.Fatalf("after %d calls, parser{%v}.recvMsg(_) = %v, %v, %v\nwant %v, %v, <nil>",
				i, p, pt, data, err, want.pt, want.data)
		}
	}

	pt, data, err := parser.recvMsg(math.MaxInt32)
	if err != io.EOF {
		t.Fatalf("after %d recvMsgs calls, parser{%v}.recvMsg(_) = %v, %v, %v\nwant _, _, %v",
			len(wantRecvs), p, pt, data, err, io.EOF)
	}
}

func TestEncode(t *testing.T) {
	for _, test := range []struct {
		// input
		msg proto.Message
		cp  Compressor
		// outputs
		b   []byte
		err error
	}{
		{nil, nil, []byte{0, 0, 0, 0, 0}, nil},
	} {
		b, err := encode(protoCodec{}, test.msg, nil, nil, nil)
		if err != test.err || !bytes.Equal(b, test.b) {
			t.Fatalf("encode(_, _, %v, _) = %v, %v\nwant %v, %v", test.cp, b, err, test.b, test.err)
		}
	}
}

func TestCompress(t *testing.T) {
	for _, test := range []struct {
		// input
		data []byte
		cp   Compressor
		dc   Decompressor
		// outputs
		err error
	}{
		{make([]byte, 1024), NewGZIPCompressor(), NewGZIPDecompressor(), nil},
	} {
		b := new(bytes.Buffer)
		if err := test.cp.Do(b, test.data); err != test.err {
			t.Fatalf("Compressor.Do(_, %v) = %v, want %v", test.data, err, test.err)
		}
		if b.Len() >= len(test.data) {
			t.Fatalf("The compressor fails to compress data.")
		}
		if p, err := test.dc.Do(b); err != nil || !bytes.Equal(test.data, p) {
			t.Fatalf("Decompressor.Do(%v) = %v, %v, want %v, <nil>", b, p, err, test.data)
		}
	}
}

func TestToRPCErr(t *testing.T) {
	for _, test := range []struct {
		// input
		errIn error
		// outputs
		errOut error
	}{
		{transport.StreamError{Code: codes.Unknown, Desc: ""}, status.Error(codes.Unknown, "")},
		{transport.ErrConnClosing, status.Error(codes.Internal, transport.ErrConnClosing.Desc)},
	} {
		err := toRPCErr(test.errIn)
		if _, ok := status.FromError(err); !ok {
			t.Fatalf("toRPCErr{%v} returned type %T, want %T", test.errIn, err, status.Error(codes.Unknown, ""))
		}
		if !reflect.DeepEqual(err, test.errOut) {
			t.Fatalf("toRPCErr{%v} = %v \nwant %v", test.errIn, err, test.errOut)
		}
	}
}

// bmEncode benchmarks encoding a Protocol Buffer message containing mSize
// bytes.
func bmEncode(b *testing.B, mSize int) {
	msg := &perfpb.Buffer{Body: make([]byte, mSize)}
	encoded, _ := encode(protoCodec{}, msg, nil, nil, nil)
	encodedSz := int64(len(encoded))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		encode(protoCodec{}, msg, nil, nil, nil)
	}
	b.SetBytes(encodedSz)
}

func BenchmarkEncode1B(b *testing.B) {
	bmEncode(b, 1)
}

func BenchmarkEncode1KiB(b *testing.B) {
	bmEncode(b, 1024)
}

func BenchmarkEncode8KiB(b *testing.B) {
	bmEncode(b, 8*1024)
}

func BenchmarkEncode64KiB(b *testing.B) {
	bmEncode(b, 64*1024)
}

func BenchmarkEncode512KiB(b *testing.B) {
	bmEncode(b, 512*1024)
}

func BenchmarkEncode1MiB(b *testing.B) {
	bmEncode(b, 1024*1024)
}

// bmCompressor benchmarks a compressor of a Protocol Buffer message containing
// mSize bytes.
func bmCompressor(b *testing.B, mSize int, cp Compressor) {
	payload := make([]byte, mSize)
	cBuf := bytes.NewBuffer(make([]byte, mSize))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		cp.Do(cBuf, payload)
		cBuf.Reset()
	}
}

func BenchmarkGZIPCompressor1B(b *testing.B) {
	bmCompressor(b, 1, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1KiB(b *testing.B) {
	bmCompressor(b, 1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor8KiB(b *testing.B) {
	bmCompressor(b, 8*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor64KiB(b *testing.B) {
	bmCompressor(b, 64*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor512KiB(b *testing.B) {
	bmCompressor(b, 512*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1MiB(b *testing.B) {
	bmCompressor(b, 1024*1024, NewGZIPCompressor())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	groupcache "github.com/ijibu/groupcache"
	"log"
	"net/http"
	"os"
	"strings"
)

func main() {
	// Usage: ./test_groupcache port
	me := ":" + os.Args[1]
	peers := groupcache.NewHTTPPool("http://localhost" + me)
	peers.Set("http://localhost:8081", "http://localhost:8082", "http://localhost:8083")
	helloworld := groupcache.NewGroup("helloworld", 1024*1024*1024*16, groupcache.GetterFunc(
		func(ctx groupcache.Context, key string, dest groupcache.Sink) error {
			//log.Println(me)
			//dest.SetString(me)
			fileName := key
			dest.SetBytes(generateThumbnail(fileName))
			return nil
		}))
	fmt.Println("GroupName: ", helloworld.Name())
	http.HandleFunc("/xbox/", func(w http.ResponseWriter, r *http.Request) {
		parts := strings.SplitN(r.URL.Path[len("/xbox/"):], "/", 1)
		if len(parts) != 1 {
			http.Error(w, "Bad Request", http.StatusBadRequest)
			return
		}
		var data []byte
		helloworld.Get(nil, parts[0], groupcache.AllocatingByteSliceSink(&data))
		w.Write(data)
		log.Println("Gets: ", helloworld.Stats.Gets.String())
		log.Println("Load: ", helloworld.Stats.Loads.String())
		log.Println("LocalLoad: ", helloworld.Stats.LocalLoads.String())
		log.Println("PeerError: ", helloworld.Stats.PeerErrors.String())
		log.Println("PeerLoad: ", helloworld.Stats.PeerLoads.String())
	})
	http.ListenAndServe(me, nil)
}

func generateThumbnail(filename string) []byte {
	// ...
	return []byte(filename + "123")
}
<commit_msg>修改分布式使用groupcache的实例代码 Signed-off-by: ijibu <ijibu.com@gmail.com><commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/ijibu/groupcache"
)

func main() {
	// Usage: ./test_groupcache port
	me := ":" + os.Args[1]
	peers := groupcache.NewHTTPPool("http://localhost" + me)
	peers.Set("http://localhost:8081", "http://localhost:8082")
	helloworld := groupcache.NewGroup("helloworld", 1024*1024*1024*16, groupcache.GetterFunc(
		func(ctx groupcache.Context, key string, dest groupcache.Sink) error {
			//log.Println(me)
			//dest.SetString(me)
			fileName := key
			dest.SetBytes(generateThumbnail(fileName))
			return nil
		}))
	fmt.Println("GroupName: ", helloworld.Name())
	http.HandleFunc("/xbox/", func(w http.ResponseWriter, r *http.Request) {
		parts := strings.SplitN(r.URL.Path[len("/xbox/"):], "/", 1)
		if len(parts) != 1 {
			http.Error(w, "Bad Request", http.StatusBadRequest)
			return
		}
		var data []byte
		helloworld.Get(nil, parts[0], groupcache.AllocatingByteSliceSink(&data))
		w.Write(data)
		log.Println("Gets: ", helloworld.Stats.Gets.String())
		log.Println("Load: ", helloworld.Stats.Loads.String())
		log.Println("LocalLoad: ", helloworld.Stats.LocalLoads.String())
		log.Println("PeerError: ", helloworld.Stats.PeerErrors.String())
		log.Println("PeerLoad: ", helloworld.Stats.PeerLoads.String())
	})
	http.ListenAndServe(me, nil)
}

func generateThumbnail(filename string) []byte {
	// ...
	return []byte(filename + "123")
}
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"os"

	"github.com/gogo/protobuf/proto"
	"github.com/kopia/kopia/internal/blockmgrpb"

	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	blockIndexShowCommand = blockIndexCommands.Command("show", "List block indexes").Alias("cat")
	blockIndexShowIDs     = blockIndexShowCommand.Arg("id", "IDs of index blocks to show").Required().Strings()
)

func runShowBlockIndexesAction(context *kingpin.ParseContext) error {
	rep := mustOpenRepository(nil)
	defer rep.Close()

	for _, blockID := range *blockIndexShowIDs {
		data, err := rep.Blocks.GetBlock(blockID)
		if err != nil {
			return fmt.Errorf("can't read block %q: %v", blockID, err)
		}

		var d blockmgrpb.Indexes
		if err := proto.Unmarshal(data, &d); err != nil {
			return err
		}

		e := json.NewEncoder(os.Stdout)
		e.SetIndent("", "  ")
		e.Encode(&d)
	}

	return nil
}

func init() {
	blockIndexShowCommand.Action(runShowBlockIndexesAction)
}
<commit_msg>added pretty printing of block indexes<commit_after>package cli

import (
	"fmt"
	"sort"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/kopia/kopia/internal/blockmgrpb"

	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	blockIndexShowCommand = blockIndexCommands.Command("show", "List block indexes").Alias("cat")
	blockIndexShowSort    = blockIndexShowCommand.Flag("sort", "Sort order").Default("offset").Enum("offset", "blockID", "size")
	blockIndexShowIDs     = blockIndexShowCommand.Arg("id", "IDs of index blocks to show").Required().Strings()
)

func runShowBlockIndexesAction(context *kingpin.ParseContext) error {
	rep := mustOpenRepository(nil)
	defer rep.Close()

	for _, blockID := range *blockIndexShowIDs {
		data, err := rep.Blocks.GetBlock(blockID)
		if err != nil {
			return fmt.Errorf("can't read block %q: %v", blockID, err)
		}

		var d blockmgrpb.Indexes
		if err := proto.Unmarshal(data, &d); err != nil {
			return err
		}

		for _, ndx := range d.Indexes {
			fmt.Printf("pack %v len: %v created %v\n", ndx.PackBlockId, ndx.PackLength, time.Unix(0, int64(ndx.CreateTimeNanos)).Local())
			type unpacked struct {
				blockID string
				offset  uint32
				size    uint32
			}
			var lines []unpacked

			for blk, os := range ndx.Items {
				lines = append(lines, unpacked{blk, uint32(os >> 32), uint32(os)})
			}
			switch *blockIndexShowSort {
			case "offset":
				sort.Slice(lines, func(i, j int) bool {
					return lines[i].offset < lines[j].offset
				})
			case "blockID":
				sort.Slice(lines, func(i, j int) bool {
					return lines[i].blockID < lines[j].blockID
				})
			case "size":
				sort.Slice(lines, func(i, j int) bool {
					return lines[i].size < lines[j].size
				})
			}
			for _, l := range lines {
				fmt.Printf("  added %-40v offset:%-10v size:%v\n", l.blockID, l.offset, l.size)
			}
			for _, del := range ndx.DeletedItems {
				fmt.Printf("  deleted %v\n", del)
			}
		}
	}

	return nil
}

func init() {
	blockIndexShowCommand.Action(runShowBlockIndexesAction)
}
<|endoftext|>


<commit_before>package app

import (
	"errors"
	"fmt"
	"github.com/timeredbull/tsuru/api/auth"
	"github.com/timeredbull/tsuru/api/unit"
	"github.com/timeredbull/tsuru/config"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/log"
	"github.com/timeredbull/tsuru/repository"
	"launchpad.net/goyaml"
	"launchpad.net/mgo/bson"
	"path"
	"strings"
	"time"
)

const confSep = "========"

type App struct {
	Env       map[string]string
	Framework string
	Ip        string
	Machine   int
	Name      string
	State     string
	Teams     []auth.Team
	Logs      []Log
}

type Log struct {
	Date    time.Time
	Message string
}

type conf struct {
	PreRestart string `yaml:"pre-restart"`
	PosRestart string `yaml:"pos-restart"`
}

func AllApps() ([]App, error) {
	var apps []App
	err := db.Session.Apps().Find(nil).All(&apps)
	return apps, err
}

func (a *App) Get() error {
	return db.Session.Apps().Find(bson.M{"name": a.Name}).One(&a)
}

func (a *App) Create() error {
	a.State = "Pending"
	err := db.Session.Apps().Insert(a)
	if err != nil {
		return err
	}
	a.Log(fmt.Sprintf("creating app %s", a.Name))
	u := a.unit()
	err = u.Create()
	if err != nil {
		return err
	}
	a.Log(fmt.Sprintf("app %s successfully created", a.Name))
	return nil
}

func (a *App) Destroy() error {
	err := db.Session.Apps().Remove(a)
	if err != nil {
		return err
	}
	u := a.unit()
	u.Destroy()
	return nil
}

func (a *App) findTeam(team *auth.Team) int {
	for i, t := range a.Teams {
		if t.Name == team.Name {
			return i
		}
	}
	return -1
}

func (a *App) hasTeam(team *auth.Team) bool {
	return a.findTeam(team) > -1
}

func (a *App) GrantAccess(team *auth.Team) error {
	if a.hasTeam(team) {
		return errors.New("This team has already access to this app")
	}
	a.Teams = append(a.Teams, *team)
	return nil
}

func (a *App) RevokeAccess(team *auth.Team) error {
	index := a.findTeam(team)
	if index < 0 {
		return errors.New("This team does not have access to this app")
	}
	last := len(a.Teams) - 1
	a.Teams[index] = a.Teams[last]
	a.Teams = a.Teams[:last]
	return nil
}

func (a *App) CheckUserAccess(user *auth.User) bool {
	for _, team := range a.Teams {
		if team.ContainsUser(user) {
			return true
		}
	}
	return false
}

func (a *App) SetEnv(name, value string) {
	if a.Env == nil {
		a.Env = make(map[string]string)
	}
	a.Env[name] = value
}

func (a *App) GetEnv(name string) (value string, err error) {
	var ok bool
	if value, ok = a.Env[name]; !ok {
		err = errors.New("Environment variable not declared for this app.")
	}
	return
}

func deployHookAbsPath(p string) (string, error) {
	repoPath, err := config.GetString("git:unit-repo")
	if err != nil {
		return "", nil
	}
	return path.Join(repoPath, p), nil
}

/*
* Returns app.conf located at app's git repository
 */
func (a *App) conf() (conf, error) {
	var c conf
	u := a.unit()
	uRepo, err := repository.GetPath()
	if err != nil {
		a.Log(fmt.Sprint("Got error while getting repository path: %s", err.Error()))
		return c, err
	}
	cPath := path.Join(uRepo, "app.conf")
	cmd := fmt.Sprintf(`echo "%s";cat %s`, confSep, cPath)
	o, err := u.Command(cmd)
	if err != nil {
		a.Log(fmt.Sprintf("Got error while executing command: %s... Skipping hooks execution", err.Error()))
		return c, nil
	}
	data := strings.Split(string(o), confSep)[1]
	err = goyaml.Unmarshal([]byte(data), &c)
	if err != nil {
		a.Log(fmt.Sprintf("Got error while parsing yaml: %s", err.Error()))
		return c, err
	}
	return c, nil
}

/*
* preRestart is responsible for running user's pre-restart script.
* The path to this script can be found at the app.conf file, at the root of user's app repository.
 */
func (a *App) preRestart(c conf) error {
	if !a.hasRestartHooks(c) {
		a.Log("app.conf file does not exists or is in the right place. Skipping...")
		return nil
	}
	if c.PreRestart == "" {
		a.Log("pre-restart hook section in app conf does not exists... Skipping...")
		return nil
	}
	u := a.unit()
	p, err := deployHookAbsPath(c.PreRestart)
	if err != nil {
		a.Log(fmt.Sprintf("Error obtaining absolute path to hook: %s. Skipping...", err))
		return nil
	}
	out, err := u.Command("/bin/bash", p)
	a.Log("Executing pre-restart hook...")
	a.Log(fmt.Sprintf("Output of pre-restart hook: %s", string(out)))
	return err
}

/*
* posRestart is responsible for running user's pos-restart script.
* The path to this script can be found at the app.conf file, at the root of user's app repository.
 */
func (a *App) posRestart(c conf) error {
	if !a.hasRestartHooks(c) {
		a.Log("app.conf file does not exists or is in the right place. Skipping...")
		return nil
	}
	if c.PosRestart == "" {
		a.Log("pos-restart hook section in app conf does not exists... Skipping...")
		return nil
	}
	u := a.unit()
	p, err := deployHookAbsPath(c.PosRestart)
	if err != nil {
		a.Log(fmt.Sprintf("Error obtaining absolute path to hook: %s. Skipping...", err))
		return nil
	}
	out, err := u.Command("/bin/bash", p)
	a.Log("Executing pos-restart hook...")
	a.Log(fmt.Sprintf("Output of pos-restart hook: %s", string(out)))
	return err
}

func (a *App) hasRestartHooks(c conf) bool {
	return !(c.PreRestart == "" && c.PosRestart == "")
}

func (a *App) updateHooks() error {
	u := a.unit()
	err := u.ExecuteHook("dependencies")
	if err != nil {
		return err
	}
	err = u.ExecuteHook("reload-gunicorn")
	if err != nil {
		return err
	}
	return nil
}

func (a *App) unit() unit.Unit {
	return unit.Unit{Name: a.Name, Type: a.Framework, Machine: a.Machine}
}

func (a *App) Log(message string) error {
	log.Printf(message)
	l := Log{Date: time.Now(), Message: message}
	a.Logs = append(a.Logs, l)
	return db.Session.Apps().Update(bson.M{"name": a.Name}, a)
}
<commit_msg>s/Sprint/Sprintf<commit_after>package app

import (
	"errors"
	"fmt"
	"github.com/timeredbull/tsuru/api/auth"
	"github.com/timeredbull/tsuru/api/unit"
	"github.com/timeredbull/tsuru/config"
	"github.com/timeredbull/tsuru/db"
	"github.com/timeredbull/tsuru/log"
	"github.com/timeredbull/tsuru/repository"
	"launchpad.net/goyaml"
	"launchpad.net/mgo/bson"
	"path"
	"strings"
	"time"
)

const confSep = "========"

type App struct {
	Env       map[string]string
	Framework string
	Ip        string
	Machine   int
	Name      string
	State     string
	Teams     []auth.Team
	Logs      []Log
}

type Log struct {
	Date    time.Time
	Message string
}

type conf struct {
	PreRestart string `yaml:"pre-restart"`
	PosRestart string `yaml:"pos-restart"`
}

func AllApps() ([]App, error) {
	var apps []App
	err := db.Session.Apps().Find(nil).All(&apps)
	return apps, err
}

func (a *App) Get() error {
	return db.Session.Apps().Find(bson.M{"name": a.Name}).One(&a)
}

func (a *App) Create() error {
	a.State = "Pending"
	err := db.Session.Apps().Insert(a)
	if err != nil {
		return err
	}
	a.Log(fmt.Sprintf("creating app %s", a.Name))
	u := a.unit()
	err = u.Create()
	if err != nil {
		return err
	}
	a.Log(fmt.Sprintf("app %s successfully created", a.Name))
	return nil
}

func (a *App) Destroy() error {
	err := db.Session.Apps().Remove(a)
	if err != nil {
		return err
	}
	u := a.unit()
	u.Destroy()
	return nil
}

func (a *App) findTeam(team *auth.Team) int {
	for i, t := range a.Teams {
		if t.Name == team.Name {
			return i
		}
	}
	return -1
}

func (a *App) hasTeam(team *auth.Team) bool {
	return a.findTeam(team) > -1
}

func (a *App) GrantAccess(team *auth.Team) error {
	if a.hasTeam(team) {
		return errors.New("This team has already access to this app")
	}
	a.Teams = append(a.Teams, *team)
	return nil
}

func (a *App) RevokeAccess(team *auth.Team) error {
	index := a.findTeam(team)
	if index < 0 {
		return errors.New("This team does not have access to this app")
	}
	last := len(a.Teams) - 1
	a.Teams[index] = a.Teams[last]
	a.Teams = a.Teams[:last]
	return nil
}

func (a *App) CheckUserAccess(user *auth.User) bool {
	for _, team := range a.Teams {
		if team.ContainsUser(user) {
			return true
		}
	}
	return false
}

func (a *App) SetEnv(name, value string) {
	if a.Env == nil {
		a.Env = make(map[string]string)
	}
	a.Env[name] = value
}

func (a *App) GetEnv(name string) (value string, err error) {
	var ok bool
	if value, ok = a.Env[name]; !ok {
		err = errors.New("Environment variable not declared for this app.")
	}
	return
}

func deployHookAbsPath(p string) (string, error) {
	repoPath, err := config.GetString("git:unit-repo")
	if err != nil {
		return "", nil
	}
	return path.Join(repoPath, p), nil
}

/*
* Returns app.conf located at app's git repository
 */
func (a *App) conf() (conf, error) {
	var c conf
	u := a.unit()
	uRepo, err := repository.GetPath()
	if err != nil {
		a.Log(fmt.Sprintf("Got error while getting repository path: %s", err.Error()))
		return c, err
	}
	cPath := path.Join(uRepo, "app.conf")
	cmd := fmt.Sprintf(`echo "%s";cat %s`, confSep, cPath)
	o, err := u.Command(cmd)
	if err != nil {
		a.Log(fmt.Sprintf("Got error while executing command: %s... Skipping hooks execution", err.Error()))
		return c, nil
	}
	data := strings.Split(string(o), confSep)[1]
	err = goyaml.Unmarshal([]byte(data), &c)
	if err != nil {
		a.Log(fmt.Sprintf("Got error while parsing yaml: %s", err.Error()))
		return c, err
	}
	return c, nil
}

/*
* preRestart is responsible for running user's pre-restart script.
* The path to this script can be found at the app.conf file, at the root of user's app repository.
 */
func (a *App) preRestart(c conf) error {
	if !a.hasRestartHooks(c) {
		a.Log("app.conf file does not exists or is in the right place. Skipping...")
		return nil
	}
	if c.PreRestart == "" {
		a.Log("pre-restart hook section in app conf does not exists... Skipping...")
		return nil
	}
	u := a.unit()
	p, err := deployHookAbsPath(c.PreRestart)
	if err != nil {
		a.Log(fmt.Sprintf("Error obtaining absolute path to hook: %s. Skipping...", err))
		return nil
	}
	out, err := u.Command("/bin/bash", p)
	a.Log("Executing pre-restart hook...")
	a.Log(fmt.Sprintf("Output of pre-restart hook: %s", string(out)))
	return err
}

/*
* posRestart is responsible for running user's pos-restart script.
* The path to this script can be found at the app.conf file, at the root of user's app repository.
 */
func (a *App) posRestart(c conf) error {
	if !a.hasRestartHooks(c) {
		a.Log("app.conf file does not exists or is in the right place. Skipping...")
		return nil
	}
	if c.PosRestart == "" {
		a.Log("pos-restart hook section in app conf does not exists... Skipping...")
		return nil
	}
	u := a.unit()
	p, err := deployHookAbsPath(c.PosRestart)
	if err != nil {
		a.Log(fmt.Sprintf("Error obtaining absolute path to hook: %s. Skipping...", err))
		return nil
	}
	out, err := u.Command("/bin/bash", p)
	a.Log("Executing pos-restart hook...")
	a.Log(fmt.Sprintf("Output of pos-restart hook: %s", string(out)))
	return err
}

func (a *App) hasRestartHooks(c conf) bool {
	return !(c.PreRestart == "" && c.PosRestart == "")
}

func (a *App) updateHooks() error {
	u := a.unit()
	err := u.ExecuteHook("dependencies")
	if err != nil {
		return err
	}
	err = u.ExecuteHook("reload-gunicorn")
	if err != nil {
		return err
	}
	return nil
}

func (a *App) unit() unit.Unit {
	return unit.Unit{Name: a.Name, Type: a.Framework, Machine: a.Machine}
}

func (a *App) Log(message string) error {
	log.Printf(message)
	l := Log{Date: time.Now(), Message: message}
	a.Logs = append(a.Logs, l)
	return db.Session.Apps().Update(bson.M{"name": a.Name}, a)
}
<|endoftext|>


<commit_before>// Copyright 2016 Afshin Darian. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package sleuth

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/ursiform/logger"
)

func init() {
	// Tests should run using a different group than production.
	group = "SLEUTH-vT"
}

// testCodes compares the error codes in an error with a list of wanted codes.
func testCodes(t *testing.T, err error, want []int) {
	codes := err.(*Error).Codes
	if len(codes) != len(want) {
		t.Errorf("expected codes length %d to be %d", len(codes), len(want))
		return
	}
	for i, code := range codes {
		if code != want[i] {
			t.Errorf("expected code [%d] to be [%d]", code, want[i])
			return
		}
	}
}

// Test client.go

func TestClientAddBadMember(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.add(group, "foo", "bar", "", "")
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errAdd})
}

func TestClientDispatchBadAction(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte(group + "FAIL"))
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errDispatchAction})
}

func TestClientDispatchEmpty(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte{})
	if err == nil {
		t.Errorf("expected client dispatch to fail on empty payload")
		return
	}
	testCodes(t, err, []int{errDispatchHeader})
}

func TestClientDoUnknownScheme(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "foo://bar/baz", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown scheme")
		return
	}
	testCodes(t, err, []int{errScheme})
}

func TestClientDoTimeout(t *testing.T) {
	c, _ := New(nil)
	defer c.Close()
	service := "foo"
	c.add(group, "bar", "baz", service, "")
	req, _ := http.NewRequest("POST", "sleuth://"+service+"/", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail by timing out")
		return
	}
	testCodes(t, err, []int{errTimeout})
}

func TestClientDoUnknownService(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "sleuth://foo/bar", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown service")
		return
	}
	testCodes(t, err, []int{errUnknownService})
}

func TestClientReceiveBadHandle(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	res := &response{Handle: "foo"}
	err := c.receive(marshalResponse(res)[len(group)+len(recv):])
	if err == nil {
		t.Errorf("expected client receive to fail on bad handle")
		return
	}
	testCodes(t, err, []int{errRECV})
}

func TestClientReceiveBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.receive([]byte(""))
	if err == nil {
		t.Errorf("expected client receive to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errResUnmarshal, errRECV})
}

func TestClientReplyBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.reply([]byte(""))
	if err == nil {
		t.Errorf("expected client reply to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errReqUnmarshal, errREPL})
}

// Test error.go

func TestError(t *testing.T) {
	code := 1
	message := "test"
	want := "sleuth: test [1]"
	err := newError(code, message)
	if err.Error() != want {
		t.Errorf("expected error to be formatted as: %s", want)
	}
}

// Test request.go

func TestUnmarshalRequestBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalRequest(payload)
	if err == nil {
		t.Errorf("expected unmarshalRequest to fail on bad json")
		return
	}
	testCodes(t, err, []int{errReqUnmarshalJSON})
}

// Test response.go

func TestUnmarshalResponseBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalResponse(payload)
	if err == nil {
		t.Errorf("expected unmarshalResponse to fail on bad json")
		return
	}
	testCodes(t, err, []int{errResUnmarshalJSON})
}

// Test sleuth.go

func TestNewBadInterface(t *testing.T) {
	_, err := New(&Config{Interface: "foo"})
	if err == nil {
		t.Errorf("expected New to fail on start with bad interface")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadLogLevel(t *testing.T) {
	c, _ := New(&Config{LogLevel: "foo"})
	if c.log.Level() != logger.Debug {
		t.Errorf("expected log level 'foo' to be coerced to 'debug'")
		return
	}
}

func TestNewBadPort(t *testing.T) {
	_, err := New(&Config{Port: 1})
	if err == nil {
		t.Errorf("expected New to fail on start with bad port")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadService(t *testing.T) {
	_, err := New(&Config{Handler: http.FileServer(http.Dir("."))})
	if err == nil {
		t.Errorf("expected New to fail with bad service")
		return
	}
	testCodes(t, err, []int{errService})
}

// Test workers.go

func TestWorkersAddDuplicate(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if n := w.add(p); n != w.add(p) {
		t.Error("expected duplicate addition to be ignored")
	}
}

func TestWorkersAvailable(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if w.add(p); !w.available() {
		t.Error("expected workers to be available")
	}
}

func TestWorkersNext(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	w.add(&peer{name: "qux", node: "quux", service: "corge"})
	if w.next().name != "foo" {
		t.Error("expected next to return first added item")
	}
	if w.next().name != "qux" {
		t.Error("expected next to return second added item")
	}
	if w.next().name != "foo" {
		t.Error("expected next to loop back to first added item")
	}
}

func TestWorkersNextNonexistent(t *testing.T) {
	w := newWorkers()
	if p := w.next(); p != nil {
		t.Error("expected nonexistent worker to be nil")
	}
}

func TestWorkersRemove(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	if _, p := w.remove("foo"); p == nil || p.name != "foo" || w.available() {
		t.Error("expected worker to be removed")
	}
}

func TestWorkersRemoveNonexistent(t *testing.T) {
	w := newWorkers()
	if _, p := w.remove("foo"); p != nil {
		t.Error("expected nonexistent worker removal to return nil")
	}
}

// Test writer.go

type goodWhisperer struct{}

func (g *goodWhisperer) Whisper(addr string, payload []byte) error {
	return nil
}

type badWhisperer struct{}

func (b *badWhisperer) Whisper(addr string, payload []byte) error {
	return errors.New("bad whisperer error")
}

func TestWriterWrite(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(goodWhisperer), &destination{node: "qux", handle: "quux"})
	if n, err := w.Write(data); err != nil {
		t.Errorf("expected write to succeed: %s", err.Error())
	} else if n <= 0 {
		t.Errorf("expected (%d) to be greater than 0", n)
	}
}

func TestWriterErrWhisper(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(badWhisperer), &destination{node: "qux", handle: "quux"})
	if _, err := w.Write(data); err == nil {
		t.Errorf("expected writer to fail")
	}
}

// Test zip.go

func TestZipUnzip(t *testing.T) {
	in := []byte("a value that should be zipped")
	zipped := zip(in)
	if out, err := unzip(zipped); err != nil {
		t.Errorf("unzip failed: %s", err.Error())
	} else if string(out) != string(in) {
		t.Errorf("zip failed")
	}
}

func TestUnzipErr(t *testing.T) {
	in := []byte("a value that cannot be unzipped")
	_, err := unzip(in)
	testCodes(t, err, []int{errUnzip})
}

// Test integrated package.

type echoHandler struct{}

func (h *echoHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	body, _ := ioutil.ReadAll(req.Body)
	res.Write(body)
}

type silentHandler struct{}

func (*silentHandler) ServeHTTP(http.ResponseWriter, *http.Request) {
	// Silent handler does nothing.
}

func TestRequestResponseCycle(t *testing.T) {
	// Create client.
	client, err := New(nil)
	if err != nil {
		t.Errorf("client instantiation failed: %s", err.Error())
		return
	}
	defer func(client *Client, t *testing.T) {
		if err := client.Close(); err != nil {
			t.Errorf("client close failed: %s", err.Error())
		}
	}(client, t)
	// Create server.
	addr := "sleuth-test-server-one"
	server, err := New(&Config{
		Handler: new(echoHandler), Service: addr})
	if err != nil {
		t.Errorf("server instantiation failed: %s", err.Error())
		return
	}
	defer func(server *Client, t *testing.T) {
		if err := server.Close(); err != nil {
			t.Errorf("server close failed: %s", err.Error())
		}
	}(server, t)
	// Wait until the server has been added to the client pool.
	client.WaitFor(addr)
	// Set timeout to 10 seconds to accommodate slow test spin-up.
	client.Timeout = time.Second * 10
	if client.block(addr) {
		t.Errorf("call to block should have returned immediately")
	}
	body := "foo bar baz"
	buffer := bytes.NewBuffer([]byte(body))
	request, err := http.NewRequest("GET", scheme+"://"+addr+"/", buffer)
	if err != nil {
		t.Errorf("request instantiation failed: %s", err.Error())
		return
	}
	response, err := client.Do(request)
	if err != nil {
		t.Errorf("client.Do failed: %s", err.Error())
		return
	}
	if response.StatusCode != http.StatusOK {
		t.Errorf("client.Do expected %d got %d", http.StatusOK, response.StatusCode)
		return
	}
	output, _ := ioutil.ReadAll(response.Body)
	// It's not really necessary to close a sleuth response.
	response.Body.Close()
	if string(output) != body {
		t.Errorf("client.Do expected %s to equal %s", string(output), body)
		return
	}
}
<commit_msg>Add comments to tests.<commit_after>// Copyright 2016 Afshin Darian. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package sleuth

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/ursiform/logger"
)

func init() {
	// Tests should run using a different group than production.
	group = "SLEUTH-vT"
}

// badWhisperer generates an error when whispering.
type badWhisperer struct{}

// Whisper allows goodWhisperer to conform to the whisperer interface. It
// returns and error every time.
func (b *badWhisperer) Whisper(addr string, payload []byte) error {
	return errors.New("bad whisperer error")
}

// goodWhisperer can whisper without errors.
type goodWhisperer struct{}

// Whisper allows goodWhisperer to conform to the whisperer interface. It
// succeeds every time.
func (g *goodWhisperer) Whisper(addr string, payload []byte) error {
	return nil
}

// echoHandler is the handler for the server in the integration test.
type echoHandler struct{}

// ServeHTTP allows echoHandler to conform to the http.Handler interface.
func (*echoHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	body, _ := ioutil.ReadAll(req.Body)
	res.Write(body)
}

// testCodes compares the error codes in an error with a list of wanted codes.
func testCodes(t *testing.T, err error, want []int) {
	codes := err.(*Error).Codes
	if len(codes) != len(want) {
		t.Errorf("expected codes length %d to be %d", len(codes), len(want))
		return
	}
	for i, code := range codes {
		if code != want[i] {
			t.Errorf("expected code [%d] to be [%d]", code, want[i])
			return
		}
	}
}

// Test client.go

func TestClientAddBadMember(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.add(group, "foo", "bar", "", "")
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errAdd})
}

func TestClientDispatchBadAction(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte(group + "FAIL"))
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errDispatchAction})
}

func TestClientDispatchEmpty(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte{})
	if err == nil {
		t.Errorf("expected client dispatch to fail on empty payload")
		return
	}
	testCodes(t, err, []int{errDispatchHeader})
}

func TestClientDoTimeout(t *testing.T) {
	c, _ := New(nil)
	defer c.Close()
	service := "foo"
	c.add(group, "bar", "baz", service, "")
	req, _ := http.NewRequest("POST", "sleuth://"+service+"/", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail by timing out")
		return
	}
	testCodes(t, err, []int{errTimeout})
}

func TestClientDoUnknownScheme(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "foo://bar/baz", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown scheme")
		return
	}
	testCodes(t, err, []int{errScheme})
}

func TestClientDoUnknownService(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "sleuth://foo/bar", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown service")
		return
	}
	testCodes(t, err, []int{errUnknownService})
}

func TestClientReceiveBadHandle(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	res := &response{Handle: "foo"}
	err := c.receive(marshalResponse(res)[len(group)+len(recv):])
	if err == nil {
		t.Errorf("expected client receive to fail on bad handle")
		return
	}
	testCodes(t, err, []int{errRECV})
}

func TestClientReceiveBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.receive([]byte(""))
	if err == nil {
		t.Errorf("expected client receive to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errResUnmarshal, errRECV})
}

func TestClientReplyBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.reply([]byte(""))
	if err == nil {
		t.Errorf("expected client reply to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errReqUnmarshal, errREPL})
}

// Test error.go

func TestError(t *testing.T) {
	code := 1
	message := "test"
	want := "sleuth: test [1]"
	err := newError(code, message)
	if err.Error() != want {
		t.Errorf("expected error to be formatted as: %s", want)
	}
}

// Test request.go

func TestRequestUnmarshalBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalRequest(payload)
	if err == nil {
		t.Errorf("expected unmarshalRequest to fail on bad json")
		return
	}
	testCodes(t, err, []int{errReqUnmarshalJSON})
}

// Test response.go

func TestResponseUnmarshalBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalResponse(payload)
	if err == nil {
		t.Errorf("expected unmarshalResponse to fail on bad json")
		return
	}
	testCodes(t, err, []int{errResUnmarshalJSON})
}

// Test sleuth.go

func TestSleuthNewBadInterface(t *testing.T) {
	_, err := New(&Config{Interface: "foo"})
	if err == nil {
		t.Errorf("expected New to fail on start with bad interface")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestSleuthNewBadLogLevel(t *testing.T) {
	c, _ := New(&Config{LogLevel: "foo"})
	if c.log.Level() != logger.Debug {
		t.Errorf("expected log level 'foo' to be coerced to 'debug'")
		return
	}
}

func TestSleuthNewBadPort(t *testing.T) {
	_, err := New(&Config{Port: 1})
	if err == nil {
		t.Errorf("expected New to fail on start with bad port")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestSleuthNewBadService(t *testing.T) {
	_, err := New(&Config{Handler: http.FileServer(http.Dir("."))})
	if err == nil {
		t.Errorf("expected New to fail with bad service")
		return
	}
	testCodes(t, err, []int{errService})
}

// Test workers.go

func TestWorkersAddDuplicate(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if n := w.add(p); n != w.add(p) {
		t.Error("expected duplicate addition to be ignored")
	}
}

func TestWorkersAvailable(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if w.add(p); !w.available() {
		t.Error("expected workers to be available")
	}
}

func TestWorkersNext(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	w.add(&peer{name: "qux", node: "quux", service: "corge"})
	if w.next().name != "foo" {
		t.Error("expected next to return first added item")
	}
	if w.next().name != "qux" {
		t.Error("expected next to return second added item")
	}
	if w.next().name != "foo" {
		t.Error("expected next to loop back to first added item")
	}
}

func TestWorkersNextNonexistent(t *testing.T) {
	w := newWorkers()
	if p := w.next(); p != nil {
		t.Error("expected nonexistent worker to be nil")
	}
}

func TestWorkersRemove(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	if _, p := w.remove("foo"); p == nil || p.name != "foo" || w.available() {
		t.Error("expected worker to be removed")
	}
}

func TestWorkersRemoveNonexistent(t *testing.T) {
	w := newWorkers()
	if _, p := w.remove("foo"); p != nil {
		t.Error("expected nonexistent worker removal to return nil")
	}
}

// Test writer.go

func TestWriterWrite(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(goodWhisperer), &destination{node: "qux", handle: "quux"})
	if n, err := w.Write(data); err != nil {
		t.Errorf("expected write to succeed: %s", err.Error())
	} else if n <= 0 {
		t.Errorf("expected (%d) to be greater than 0", n)
	}
}

func TestWriterWriteBadWhisperer(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(badWhisperer), &destination{node: "qux", handle: "quux"})
	_, err := w.Write(data)
	if err == nil {
		t.Errorf("expected writer to fail using bad whisperer")
		return
	}
	testCodes(t, err, []int{errResWhisper})
}

// Test zip.go

func TestZipUnzipBadInput(t *testing.T) {
	in := []byte("a value that cannot be unzipped")
	_, err := unzip(in)
	if err == nil {
		t.Errorf("expected unzip to fail with bad input")
		return
	}
	testCodes(t, err, []int{errUnzip})
}

func TestZipUnzip(t *testing.T) {
	in := []byte("a value that should be zipped")
	zipped := zip(in)
	if out, err := unzip(zipped); err != nil {
		t.Errorf("unzip failed: %s", err.Error())
	} else if string(out) != string(in) {
		t.Errorf("zip failed")
	}
}

// Test integrated package.

func TestIntegratedCycle(t *testing.T) {
	// Create client.
	client, err := New(nil)
	if err != nil {
		t.Errorf("client instantiation failed: %s", err.Error())
		return
	}
	defer func(client *Client, t *testing.T) {
		if err := client.Close(); err != nil {
			t.Errorf("client close failed: %s", err.Error())
		}
	}(client, t)
	// Create server.
	addr := "sleuth-test-server-one"
	server, err := New(&Config{
		Handler: new(echoHandler), Service: addr})
	if err != nil {
		t.Errorf("server instantiation failed: %s", err.Error())
		return
	}
	defer func(server *Client, t *testing.T) {
		if err := server.Close(); err != nil {
			t.Errorf("server close failed: %s", err.Error())
		}
	}(server, t)
	// Wait until the server has been added to the client pool.
	client.WaitFor(addr)
	// Set timeout to 10 seconds to accommodate slow test spin-up.
	client.Timeout = time.Second * 10
	if client.block(addr) {
		t.Errorf("call to block should have returned immediately")
	}
	body := "foo bar baz"
	buffer := bytes.NewBuffer([]byte(body))
	request, err := http.NewRequest("GET", scheme+"://"+addr+"/", buffer)
	if err != nil {
		t.Errorf("request instantiation failed: %s", err.Error())
		return
	}
	response, err := client.Do(request)
	if err != nil {
		t.Errorf("client.Do failed: %s", err.Error())
		return
	}
	if response.StatusCode != http.StatusOK {
		t.Errorf("client.Do expected %d got %d", http.StatusOK, response.StatusCode)
		return
	}
	output, _ := ioutil.ReadAll(response.Body)
	// It's not really necessary to close a sleuth response.
	response.Body.Close()
	if string(output) != body {
		t.Errorf("client.Do expected %s to equal %s", string(output), body)
		return
	}
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2016-2017 Bitnami

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 function

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/ghodss/yaml"
	cronjobApi "github.com/kubeless/cronjob-trigger/pkg/apis/kubeless/v1beta1"
	cronjobUtils "github.com/kubeless/cronjob-trigger/pkg/utils"
	kubelessApi "github.com/kubeless/kubeless/pkg/apis/kubeless/v1beta1"
	"github.com/kubeless/kubeless/pkg/langruntime"
	kubelessUtils "github.com/kubeless/kubeless/pkg/utils"
	"github.com/robfig/cron"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var deployCmd = &cobra.Command{
	Use:   "deploy <function_name> FLAG",
	Short: "deploy a function to Kubeless",
	Long:  `deploy a function to Kubeless`,
	Run: func(cmd *cobra.Command, args []string) {
		cli := kubelessUtils.GetClientOutOfCluster()
		apiExtensionsClientset := kubelessUtils.GetAPIExtensionsClientOutOfCluster()
		config, err := kubelessUtils.GetKubelessConfig(cli, apiExtensionsClientset)
		if err != nil {
			logrus.Fatalf("Unable to read the configmap: %v", err)
		}

		var lr = langruntime.New(config)
		lr.ReadConfigMap()

		if len(args) != 1 {
			logrus.Fatal("Need exactly one argument - function name")
		}
		funcName := args[0]

		schedule, err := cmd.Flags().GetString("schedule")
		if err != nil {
			logrus.Fatal(err)
		}

		if schedule != "" {
			if _, err := cron.ParseStandard(schedule); err != nil {
				logrus.Fatalf("Invalid value for --schedule. " + err.Error())
			}
		}

		labels, err := cmd.Flags().GetStringSlice("label")
		if err != nil {
			logrus.Fatal(err)
		}

		envs, err := cmd.Flags().GetStringSlice("env")
		if err != nil {
			logrus.Fatal(err)
		}

		runtime, err := cmd.Flags().GetString("runtime")
		if err != nil {
			logrus.Fatal(err)
		}

		if runtime != "" && !lr.IsValidRuntime(runtime) {
			logrus.Fatalf("Invalid runtime: %s. Supported runtimes are: %s",
				runtime, strings.Join(lr.GetRuntimes(), ", "))
		}

		handler, err := cmd.Flags().GetString("handler")
		if err != nil {
			logrus.Fatal(err)
		}

		file, err := cmd.Flags().GetString("from-file")
		if err != nil {
			logrus.Fatal(err)
		}

		ns, err := cmd.Flags().GetString("namespace")
		if err != nil {
			logrus.Fatal(err)
		}
		if ns == "" {
			ns = kubelessUtils.GetDefaultNamespace()
		}

		deps, err := cmd.Flags().GetString("dependencies")
		if err != nil {
			logrus.Fatal(err)
		}

		secrets, err := cmd.Flags().GetStringSlice("secrets")
		if err != nil {
			logrus.Fatal(err)
		}

		runtimeImage, err := cmd.Flags().GetString("runtime-image")
		if err != nil {
			logrus.Fatal(err)
		}

		imagePullPolicy, err := cmd.Flags().GetString("image-pull-policy")
		if err != nil {
			logrus.Fatal(err)
		}

		if imagePullPolicy != "IfNotPresent" && imagePullPolicy != "Always" && imagePullPolicy != "Never" {
			err := fmt.Errorf("image-pull-policy must be {IfNotPresent|Always|Never}")
			logrus.Fatal(err)
		}

		mem, err := cmd.Flags().GetString("memory")
		if err != nil {
			logrus.Fatal(err)
		}

		cpu, err := cmd.Flags().GetString("cpu")
		if err != nil {
			logrus.Fatal(err)
		}

		timeout, err := cmd.Flags().GetString("timeout")
		if err != nil {
			logrus.Fatal(err)
		}

		output, err := cmd.Flags().GetString("output")
		if err != nil {
			logrus.Fatal(err)
		}

		headless, err := cmd.Flags().GetBool("headless")
		if err != nil {
			logrus.Fatal(err)
		}

		dryrun, err := cmd.Flags().GetBool("dryrun")
		if err != nil {
			logrus.Fatal(err)
		}

		port, err := cmd.Flags().GetInt32("port")
		if err != nil {
			logrus.Fatal(err)
		}
		if port <= 0 || port > 65535 {
			logrus.Fatalf("Invalid port number %d specified", port)
		}

		servicePort, err := cmd.Flags().GetInt32("servicePort")
		if err != nil {
			logrus.Fatal(err)
		}
		if servicePort < 0 || servicePort > 65535 {
			logrus.Fatalf("Invalid servicePort number %d specified", servicePort)
		}

		funcDeps := ""
		if deps != "" {
			contentType, err := getContentType(deps)
			if err != nil {
				logrus.Fatal(err)
			}
			funcDeps, _, err = parseContent(deps, contentType)
			if err != nil {
				logrus.Fatal(err)
			}
		}

		if runtime == "" && runtimeImage == "" {
			logrus.Fatal("Either `--runtime` or `--runtime-image` flag must be specified.")
		}

		if runtime != "" && handler == "" {
			logrus.Fatal("You must specify handler for the runtime.")
		}

		defaultFunctionSpec := kubelessApi.Function{}
		defaultFunctionSpec.ObjectMeta.Labels = map[string]string{
			"created-by": "kubeless",
			"function":   funcName,
		}

		f, err := getFunctionDescription(funcName, ns, handler, file, funcDeps, runtime, runtimeImage, mem, cpu, timeout, imagePullPolicy, port, servicePort, headless, envs, labels, secrets, defaultFunctionSpec)
		if err != nil {
			logrus.Fatal(err)
		}

		if dryrun == true {
			if output == "json" {
				j, err := json.MarshalIndent(f, "", "    ")
				if err != nil {
					logrus.Fatal(err)
				}
				fmt.Println(string(j[:]))
				return
			} else if output == "yaml" {
				y, err := yaml.Marshal(f)
				if err != nil {
					logrus.Fatal(err)
				}
				fmt.Println(string(y[:]))
				return
			} else {
				logrus.Infof("Output format needs to be yaml or json")
				return
			}
		}

		kubelessClient, err := kubelessUtils.GetKubelessClientOutCluster()
		if err != nil {
			logrus.Fatal(err)
		}

		logrus.Infof("Deploying function...")
		err = kubelessUtils.CreateFunctionCustomResource(kubelessClient, f)
		if err != nil {
			logrus.Fatalf("Failed to deploy %s. Received:\n%s", funcName, err)
		}
		logrus.Infof("Function %s submitted for deployment", funcName)
		logrus.Infof("Check the deployment status executing 'kubeless function ls %s'", funcName)

		if schedule != "" {
			cronJobTrigger := cronjobApi.CronJobTrigger{}
			cronJobTrigger.TypeMeta = metav1.TypeMeta{
				Kind:       "CronJobTrigger",
				APIVersion: "kubeless.io/v1beta1",
			}
			cronJobTrigger.ObjectMeta = metav1.ObjectMeta{
				Name:      funcName,
				Namespace: ns,
			}
			cronJobTrigger.ObjectMeta.Labels = map[string]string{
				"created-by": "kubeless",
				"function":   funcName,
			}
			cronJobTrigger.Spec.FunctionName = funcName
			cronJobTrigger.Spec.Schedule = schedule
			cronjobClient, err := cronjobUtils.GetKubelessClientOutCluster()
			if err != nil {
				logrus.Fatal(err)
			}
			err = cronjobUtils.CreateCronJobCustomResource(cronjobClient, &cronJobTrigger)
			if err != nil {
				logrus.Fatalf("Failed to deploy cron job trigger %s. Received:\n%s", funcName, err)
			}
		}
	},
}

func init() {
	deployCmd.Flags().StringP("runtime", "r", "", "Specify runtime")
	deployCmd.Flags().StringP("handler", "", "", "Specify handler")
	deployCmd.Flags().StringP("from-file", "f", "", "Specify code file or a URL to the code file")
	deployCmd.Flags().StringSliceP("label", "l", []string{}, "Specify labels of the function. Both separator ':' and '=' are allowed. For example: --label foo1=bar1,foo2:bar2")
	deployCmd.Flags().StringSliceP("secrets", "", []string{}, "Specify Secrets to be mounted to the functions container. For example: --secrets mySecret")
	deployCmd.Flags().StringSliceP("env", "e", []string{}, "Specify environment variable of the function. Both separator ':' and '=' are allowed. For example: --env foo1=bar1,foo2:bar2")
	deployCmd.Flags().StringP("namespace", "n", "", "Specify namespace for the function")
	deployCmd.Flags().StringP("dependencies", "d", "", "Specify a file containing list of dependencies for the function")
	deployCmd.Flags().StringP("schedule", "", "", "Specify schedule in cron format for scheduled function")
	deployCmd.Flags().StringP("memory", "", "", "Request amount of memory, which is measured in bytes, for the function. It is expressed as a plain integer or a fixed-point interger with one of these suffies: E, P, T, G, M, K, Ei, Pi, Ti, Gi, Mi, Ki")
	deployCmd.Flags().StringP("cpu", "", "", "Request amount of cpu for the function, which is measured in units of cores. Please see the following link for more information: https://kubernetes.io/docs/concepts/configuration/manage-compute-resources-container/#meaning-of-cpu")
	deployCmd.Flags().StringP("runtime-image", "", "", "Custom runtime image")
	deployCmd.Flags().StringP("image-pull-policy", "", "Always", "Image pull policy")
	deployCmd.Flags().StringP("timeout", "", "180", "Maximum timeout (in seconds) for the function to complete its execution")
	deployCmd.Flags().StringP("output", "o", "yaml", "Output format")
	deployCmd.Flags().Bool("headless", false, "Deploy http-based function without a single service IP and load balancing support from Kubernetes. See: https://kubernetes.io/docs/concepts/services-networking/service/#headless-services")
	deployCmd.Flags().Bool("dryrun", false, "Output JSON manifest of the function without creating it")
	deployCmd.Flags().Int32("port", 8080, "Deploy http-based function with a custom port")
	deployCmd.Flags().Int32("servicePort", 0, "Deploy http-based function with a custom service port. If not provided the value of 'port' will be used")
}
<commit_msg>Disable runtime check if has no permissions to do so (#992)<commit_after>/*
Copyright (c) 2016-2017 Bitnami

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 function

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/ghodss/yaml"
	cronjobApi "github.com/kubeless/cronjob-trigger/pkg/apis/kubeless/v1beta1"
	cronjobUtils "github.com/kubeless/cronjob-trigger/pkg/utils"
	kubelessApi "github.com/kubeless/kubeless/pkg/apis/kubeless/v1beta1"
	"github.com/kubeless/kubeless/pkg/langruntime"
	kubelessUtils "github.com/kubeless/kubeless/pkg/utils"
	"github.com/robfig/cron"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var deployCmd = &cobra.Command{
	Use:   "deploy <function_name> FLAG",
	Short: "deploy a function to Kubeless",
	Long:  `deploy a function to Kubeless`,
	Run: func(cmd *cobra.Command, args []string) {
		cli := kubelessUtils.GetClientOutOfCluster()
		apiExtensionsClientset := kubelessUtils.GetAPIExtensionsClientOutOfCluster()

		if len(args) != 1 {
			logrus.Fatal("Need exactly one argument - function name")
		}
		funcName := args[0]

		runtime, err := cmd.Flags().GetString("runtime")
		if err != nil {
			logrus.Fatal(err)
		}

		// Checking runtime parameter if allowed by RBAC, otherwide skip the check
		config, err := kubelessUtils.GetKubelessConfig(cli, apiExtensionsClientset)
		if config == nil || err != nil {
			logrus.Warnf("%v. Runtime check is disabled.", err)
		} else {
			lr := langruntime.New(config)
			lr.ReadConfigMap()

			if runtime != "" && !lr.IsValidRuntime(runtime) {
				logrus.Fatalf("Invalid runtime: %s. Supported runtimes are: %s",
					runtime, strings.Join(lr.GetRuntimes(), ", "))
			}
		}

		schedule, err := cmd.Flags().GetString("schedule")
		if err != nil {
			logrus.Fatal(err)
		}

		if schedule != "" {
			if _, err := cron.ParseStandard(schedule); err != nil {
				logrus.Fatalf("Invalid value for --schedule. " + err.Error())
			}
		}

		labels, err := cmd.Flags().GetStringSlice("label")
		if err != nil {
			logrus.Fatal(err)
		}

		envs, err := cmd.Flags().GetStringSlice("env")
		if err != nil {
			logrus.Fatal(err)
		}

		handler, err := cmd.Flags().GetString("handler")
		if err != nil {
			logrus.Fatal(err)
		}

		file, err := cmd.Flags().GetString("from-file")
		if err != nil {
			logrus.Fatal(err)
		}

		ns, err := cmd.Flags().GetString("namespace")
		if err != nil {
			logrus.Fatal(err)
		}
		if ns == "" {
			ns = kubelessUtils.GetDefaultNamespace()
		}

		deps, err := cmd.Flags().GetString("dependencies")
		if err != nil {
			logrus.Fatal(err)
		}

		secrets, err := cmd.Flags().GetStringSlice("secrets")
		if err != nil {
			logrus.Fatal(err)
		}

		runtimeImage, err := cmd.Flags().GetString("runtime-image")
		if err != nil {
			logrus.Fatal(err)
		}

		imagePullPolicy, err := cmd.Flags().GetString("image-pull-policy")
		if err != nil {
			logrus.Fatal(err)
		}

		if imagePullPolicy != "IfNotPresent" && imagePullPolicy != "Always" && imagePullPolicy != "Never" {
			err := fmt.Errorf("image-pull-policy must be {IfNotPresent|Always|Never}")
			logrus.Fatal(err)
		}

		mem, err := cmd.Flags().GetString("memory")
		if err != nil {
			logrus.Fatal(err)
		}

		cpu, err := cmd.Flags().GetString("cpu")
		if err != nil {
			logrus.Fatal(err)
		}

		timeout, err := cmd.Flags().GetString("timeout")
		if err != nil {
			logrus.Fatal(err)
		}

		output, err := cmd.Flags().GetString("output")
		if err != nil {
			logrus.Fatal(err)
		}

		headless, err := cmd.Flags().GetBool("headless")
		if err != nil {
			logrus.Fatal(err)
		}

		dryrun, err := cmd.Flags().GetBool("dryrun")
		if err != nil {
			logrus.Fatal(err)
		}

		port, err := cmd.Flags().GetInt32("port")
		if err != nil {
			logrus.Fatal(err)
		}
		if port <= 0 || port > 65535 {
			logrus.Fatalf("Invalid port number %d specified", port)
		}

		servicePort, err := cmd.Flags().GetInt32("servicePort")
		if err != nil {
			logrus.Fatal(err)
		}
		if servicePort < 0 || servicePort > 65535 {
			logrus.Fatalf("Invalid servicePort number %d specified", servicePort)
		}

		funcDeps := ""
		if deps != "" {
			contentType, err := getContentType(deps)
			if err != nil {
				logrus.Fatal(err)
			}
			funcDeps, _, err = parseContent(deps, contentType)
			if err != nil {
				logrus.Fatal(err)
			}
		}

		if runtime == "" && runtimeImage == "" {
			logrus.Fatal("Either `--runtime` or `--runtime-image` flag must be specified.")
		}

		if runtime != "" && handler == "" {
			logrus.Fatal("You must specify handler for the runtime.")
		}

		defaultFunctionSpec := kubelessApi.Function{}
		defaultFunctionSpec.ObjectMeta.Labels = map[string]string{
			"created-by": "kubeless",
			"function":   funcName,
		}

		f, err := getFunctionDescription(funcName, ns, handler, file, funcDeps, runtime, runtimeImage, mem, cpu, timeout, imagePullPolicy, port, servicePort, headless, envs, labels, secrets, defaultFunctionSpec)
		if err != nil {
			logrus.Fatal(err)
		}

		if dryrun == true {
			if output == "json" {
				j, err := json.MarshalIndent(f, "", "    ")
				if err != nil {
					logrus.Fatal(err)
				}
				fmt.Println(string(j[:]))
				return
			} else if output == "yaml" {
				y, err := yaml.Marshal(f)
				if err != nil {
					logrus.Fatal(err)
				}
				fmt.Println(string(y[:]))
				return
			} else {
				logrus.Infof("Output format needs to be yaml or json")
				return
			}
		}

		kubelessClient, err := kubelessUtils.GetKubelessClientOutCluster()
		if err != nil {
			logrus.Fatal(err)
		}

		logrus.Infof("Deploying function...")
		err = kubelessUtils.CreateFunctionCustomResource(kubelessClient, f)
		if err != nil {
			logrus.Fatalf("Failed to deploy %s. Received:\n%s", funcName, err)
		}
		logrus.Infof("Function %s submitted for deployment", funcName)
		logrus.Infof("Check the deployment status executing 'kubeless function ls %s'", funcName)

		if schedule != "" {
			cronJobTrigger := cronjobApi.CronJobTrigger{}
			cronJobTrigger.TypeMeta = metav1.TypeMeta{
				Kind:       "CronJobTrigger",
				APIVersion: "kubeless.io/v1beta1",
			}
			cronJobTrigger.ObjectMeta = metav1.ObjectMeta{
				Name:      funcName,
				Namespace: ns,
			}
			cronJobTrigger.ObjectMeta.Labels = map[string]string{
				"created-by": "kubeless",
				"function":   funcName,
			}
			cronJobTrigger.Spec.FunctionName = funcName
			cronJobTrigger.Spec.Schedule = schedule
			cronjobClient, err := cronjobUtils.GetKubelessClientOutCluster()
			if err != nil {
				logrus.Fatal(err)
			}
			err = cronjobUtils.CreateCronJobCustomResource(cronjobClient, &cronJobTrigger)
			if err != nil {
				logrus.Fatalf("Failed to deploy cron job trigger %s. Received:\n%s", funcName, err)
			}
		}
	},
}

func init() {
	deployCmd.Flags().StringP("runtime", "r", "", "Specify runtime")
	deployCmd.Flags().StringP("handler", "", "", "Specify handler")
	deployCmd.Flags().StringP("from-file", "f", "", "Specify code file or a URL to the code file")
	deployCmd.Flags().StringSliceP("label", "l", []string{}, "Specify labels of the function. Both separator ':' and '=' are allowed. For example: --label foo1=bar1,foo2:bar2")
	deployCmd.Flags().StringSliceP("secrets", "", []string{}, "Specify Secrets to be mounted to the functions container. For example: --secrets mySecret")
	deployCmd.Flags().StringSliceP("env", "e", []string{}, "Specify environment variable of the function. Both separator ':' and '=' are allowed. For example: --env foo1=bar1,foo2:bar2")
	deployCmd.Flags().StringP("namespace", "n", "", "Specify namespace for the function")
	deployCmd.Flags().StringP("dependencies", "d", "", "Specify a file containing list of dependencies for the function")
	deployCmd.Flags().StringP("schedule", "", "", "Specify schedule in cron format for scheduled function")
	deployCmd.Flags().StringP("memory", "", "", "Request amount of memory, which is measured in bytes, for the function. It is expressed as a plain integer or a fixed-point interger with one of these suffies: E, P, T, G, M, K, Ei, Pi, Ti, Gi, Mi, Ki")
	deployCmd.Flags().StringP("cpu", "", "", "Request amount of cpu for the function, which is measured in units of cores. Please see the following link for more information: https://kubernetes.io/docs/concepts/configuration/manage-compute-resources-container/#meaning-of-cpu")
	deployCmd.Flags().StringP("runtime-image", "", "", "Custom runtime image")
	deployCmd.Flags().StringP("image-pull-policy", "", "Always", "Image pull policy")
	deployCmd.Flags().StringP("timeout", "", "180", "Maximum timeout (in seconds) for the function to complete its execution")
	deployCmd.Flags().StringP("output", "o", "yaml", "Output format")
	deployCmd.Flags().Bool("headless", false, "Deploy http-based function without a single service IP and load balancing support from Kubernetes. See: https://kubernetes.io/docs/concepts/services-networking/service/#headless-services")
	deployCmd.Flags().Bool("dryrun", false, "Output JSON manifest of the function without creating it")
	deployCmd.Flags().Int32("port", 8080, "Deploy http-based function with a custom port")
	deployCmd.Flags().Int32("servicePort", 0, "Deploy http-based function with a custom service port. If not provided the value of 'port' will be used")
}
<|endoftext|>


<commit_before>package app

import (
	"net"
	"os"
	"strings"
	"time"

	"github.com/coredns/coredns/coremain"
	clog "github.com/coredns/coredns/plugin/pkg/log"

	"k8s.io/dns/cmd/kube-dns/app/options"
	"k8s.io/dns/pkg/dns/config"
	"k8s.io/dns/pkg/netif"
	"k8s.io/kubernetes/pkg/util/dbus"
	utiliptables "k8s.io/kubernetes/pkg/util/iptables"
	utilexec "k8s.io/utils/exec"
	utilnet "k8s.io/utils/net"
	utilebtables "k8s.io/utils/net/ebtables"
)

// ConfigParams lists the configuration options that can be provided to node-cache
type ConfigParams struct {
	LocalIPStr           string        // comma separated listen ips for the local cache agent
	LocalIPs             []net.IP      // parsed ip addresses for the local cache agent to listen for dns requests
	LocalPort            string        // port to listen for dns requests
	MetricsListenAddress string        // address to serve metrics on
	SetupInterface       bool          // Indicates whether to setup network interface
	InterfaceName        string        // Name of the interface to be created
	Interval             time.Duration // specifies how often to run iptables rules check
	BaseCoreFile         string        // Path to the template config file for node-cache
	CoreFile             string        // Path to config file used by node-cache
	KubednsCMPath        string        // Directory where kube-dns configmap will be mounted
	UpstreamSvcName      string        // Name of the service whose clusterIP is the upstream for node-cache for cluster domain
	HealthPort           string        // port for the healthcheck
	SetupIptables        bool
	SetupEbtables        bool
	SkipTeardown         bool // Indicates whether the iptables rules and interface should be torn down
}

type iptablesRule struct {
	table utiliptables.Table
	chain utiliptables.Chain
	args  []string
}

type ebtablesRule struct {
	table utilebtables.Table
	chain utilebtables.Chain
	args  []string
}

// CacheApp contains all the config required to run node-cache.
type CacheApp struct {
	iptables      utiliptables.Interface
	iptablesRules []iptablesRule
	ebtables      utilebtables.Interface
	ebtablesRules []ebtablesRule
	params        *ConfigParams
	netifHandle   *netif.NetifManager
	kubednsConfig *options.KubeDNSConfig
	exitChan      chan struct{} // Channel to terminate background goroutines
	clusterDNSIP  net.IP
}

func isLockedErr(err error) bool {
	return strings.Contains(err.Error(), "holding the xtables lock")
}

// Init initializes the parameters and networking setup necessary to run node-cache
func (c *CacheApp) Init() {
	if c.params.SetupInterface {
		c.netifHandle = netif.NewNetifManager(c.params.LocalIPs)
	}
	if c.params.SetupIptables {
		c.initIptables()
	}
	if c.params.SetupEbtables {
		c.initEbtables()
	}
	initMetrics(c.params.MetricsListenAddress)
	// Write the config file from template.
	// this is required in case there is no or erroneous kube-dns configpath specified.
	c.updateCorefile(&config.Config{})
	// Initialize periodic sync for node-local-dns, kube-dns configmap.
	c.initDNSConfigSync()
	// Setup only the network interface first.
	// This is to ensure that iptables rules don't get setup if the cache(coreDNS) is unable to startup due to config
	// error, port conflicts or other reasons.
	c.setupInterface()
	// Initial setup of iptables/ebtables
	c.setupNetworking()
}

// isIPv6 return if the node-cache is working in IPv6 mode
// LocalIPs are guaranteed to have the same family
func (c *CacheApp) isIPv6() bool {
	if len(c.params.LocalIPs) > 0 {
		return utilnet.IsIPv6(c.params.LocalIPs[0])
	}
	return false
}

func (c *CacheApp) initIptables() {
	// using the localIPStr param since we need ip strings here
	for _, localIP := range strings.Split(c.params.LocalIPStr, ",") {
		c.iptablesRules = append(c.iptablesRules, []iptablesRule{
			// Match traffic destined for localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
			{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			// There are rules in filter table to allow tracked connections to be accepted. Since we skipped connection tracking,
			// need these additional filter table rules.
			{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "ACCEPT"}},
			{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "ACCEPT"}},
			// Match traffic from localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "NOTRACK"}},
			// Additional filter table rules for traffic frpm localIp:localPort
			{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "ACCEPT"}},
			{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "udp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "ACCEPT"}},
			// Skip connection tracking for requests to nodelocalDNS that are locally generated, example - by hostNetwork pods
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			// skip connection tracking for healthcheck requests generated by liveness probe to health plugin
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.HealthPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.HealthPort, "-j", "NOTRACK"}},
		}...)
	}
	c.iptables = newIPTables(c.isIPv6())
}

func newIPTables(isIPv6 bool) utiliptables.Interface {
	execer := utilexec.New()
	dbus := dbus.New()
	protocol := utiliptables.ProtocolIpv4
	if isIPv6 {
		protocol = utiliptables.ProtocolIpv6
	}
	return utiliptables.New(execer, dbus, protocol)
}

func (c *CacheApp) initEbtables() {
	protocol := "IPv4"
	if c.isIPv6() {
		protocol = "IPv6"
	}
	// using the localIPStr param since we need ip strings here
	for _, localIP := range strings.Split(c.params.LocalIPStr, ",") {
		c.ebtablesRules = append(c.ebtablesRules, []ebtablesRule{
			// Match traffic destined for localIp and use the MAC address of the bridge port as destination address
			{utilebtables.TableBroute, utilebtables.ChainBrouting, []string{"-p", protocol, "--ip-dst", localIP,
				"-j", "redirect"}},
		}...)
	}
	c.ebtables = newEBTables()
}

func newEBTables() utilebtables.Interface {
	execer := utilexec.New()
	return utilebtables.New(execer)
}

// TeardownNetworking removes all custom iptables rules and network interface added by node-cache
func (c *CacheApp) TeardownNetworking() error {
	clog.Infof("Tearing down")
	if c.exitChan != nil {
		// Stop the goroutine that periodically checks for iptables rules/dummy interface
		// exitChan is a buffered channel of size 1, so this will not block
		c.exitChan <- struct{}{}
	}
	var err error
	if c.params.SetupInterface {
		err = c.netifHandle.RemoveDummyDevice(c.params.InterfaceName)
	}
	if c.params.SetupIptables {
		for _, rule := range c.iptablesRules {
			exists := true
			for exists == true {
				c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
				exists, _ = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			}
			// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
			c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
		}
		if c.params.SetupEbtables {
			for _, rule := range c.ebtablesRules {
				exists := true
				for exists == true {
					c.ebtables.DeleteRule(rule.table, rule.chain, rule.args...)
					exists, _ = c.ebtables.EnsureRule(utilebtables.Append, rule.table, rule.chain, rule.args...)
				}
				// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
				c.ebtables.DeleteRule(rule.table, rule.chain, rule.args...)
			}
		}
	}
	return err
}

func (c *CacheApp) setupNetworking() {
	if c.params.SetupIptables {
		for _, rule := range c.iptablesRules {
			exists, err := c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			switch {
			case exists:
				// debug messages can be printed by including "debug" plugin in coreFile.
				clog.Debugf("iptables rule %v for nodelocaldns already exists", rule)
				continue
			case err == nil:
				clog.Infof("Added back nodelocaldns rule - %v", rule)
				continue
				// if we got here, either iptables check failed or adding rule back failed.
			case isLockedErr(err):
				clog.Infof("Error checking/adding iptables rule %v, due to xtables lock in use, retrying in %v", rule, c.params.Interval)
				setupErrCount.WithLabelValues("iptables_lock").Inc()
			default:
				clog.Errorf("Error adding iptables rule %v - %s", rule, err)
				setupErrCount.WithLabelValues("iptables").Inc()
			}
		}
	}

	if c.params.SetupEbtables {
		for _, rule := range c.ebtablesRules {
			exists, err := c.ebtables.EnsureRule(utilebtables.Append, rule.table, rule.chain, rule.args...)
			switch {
			case exists:
				// debug messages can be printed by including "debug" plugin in coreFile.
				clog.Debugf("ebtables rule %v for nodelocaldns already exists", rule)
				continue
			case err == nil:
				clog.Infof("Added back ebtables rule - %v", rule)
				continue
				// if we got here, either ebtables check failed or adding rule back failed.
			default:
				clog.Errorf("Error adding ebtables rule %v - %s", rule, err)
				setupErrCount.WithLabelValues("ebtables").Inc()
			}
		}
	}

	c.setupInterface()
}

func (c *CacheApp) setupInterface() {
	if c.params.SetupInterface {
		exists, err := c.netifHandle.EnsureDummyDevice(c.params.InterfaceName)
		if !exists {
			if err != nil {
				clog.Errorf("Failed to add non-existent interface %s: %s", c.params.InterfaceName, err)
				setupErrCount.WithLabelValues("interface_add").Inc()
			}
			clog.Infof("Added interface - %s", c.params.InterfaceName)
		}
		if err != nil {
			clog.Errorf("Error checking dummy device %s - %s", c.params.InterfaceName, err)
			setupErrCount.WithLabelValues("interface_check").Inc()
		}
	}
}

func (c *CacheApp) runPeriodic() {
	c.exitChan = make(chan struct{}, 1)
	tick := time.NewTicker(c.params.Interval * time.Second)
	for {
		select {
		case <-tick.C:
			c.setupNetworking()
		case <-c.exitChan:
			clog.Warningf("Exiting iptables/interface check goroutine")
			return
		}
	}
}

// RunApp invokes the background checks and runs coreDNS as a cache
func (c *CacheApp) RunApp() {
	go c.runPeriodic()
	coremain.Run()
	// Unlikely to reach here, if we did it is because coremain exited and the signal was not trapped.
	clog.Errorf("Untrapped signal, tearing down")
	c.TeardownNetworking()
}

// NewCacheApp returns a new instance of CacheApp by applying the specified config params.
func NewCacheApp(params *ConfigParams) (*CacheApp, error) {
	c := &CacheApp{params: params, kubednsConfig: options.NewKubeDNSConfig()}
	c.clusterDNSIP = net.ParseIP(os.ExpandEnv(toSvcEnv(params.UpstreamSvcName)))
	if c.clusterDNSIP == nil {
		clog.Warningf("Unable to lookup IP address of Upstream service %s, env %s `%s`", params.UpstreamSvcName, toSvcEnv(params.UpstreamSvcName), os.ExpandEnv(toSvcEnv(params.UpstreamSvcName)))
	}
	return c, nil
}

// toSvcEnv converts service name to the corresponding ENV variable. This is exposed in every pod and its value is the clusterIP.
// https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables
func toSvcEnv(svcName string) string {
	envName := strings.Replace(svcName, "-", "_", -1)
	return "$" + strings.ToUpper(envName) + "_SERVICE_HOST"
}
<commit_msg>call setupNetworking before loop<commit_after>package app

import (
	"net"
	"os"
	"strings"
	"time"

	"github.com/coredns/coredns/coremain"
	clog "github.com/coredns/coredns/plugin/pkg/log"

	"k8s.io/dns/cmd/kube-dns/app/options"
	"k8s.io/dns/pkg/dns/config"
	"k8s.io/dns/pkg/netif"
	"k8s.io/kubernetes/pkg/util/dbus"
	utiliptables "k8s.io/kubernetes/pkg/util/iptables"
	utilexec "k8s.io/utils/exec"
	utilnet "k8s.io/utils/net"
	utilebtables "k8s.io/utils/net/ebtables"
)

// ConfigParams lists the configuration options that can be provided to node-cache
type ConfigParams struct {
	LocalIPStr           string        // comma separated listen ips for the local cache agent
	LocalIPs             []net.IP      // parsed ip addresses for the local cache agent to listen for dns requests
	LocalPort            string        // port to listen for dns requests
	MetricsListenAddress string        // address to serve metrics on
	SetupInterface       bool          // Indicates whether to setup network interface
	InterfaceName        string        // Name of the interface to be created
	Interval             time.Duration // specifies how often to run iptables rules check
	BaseCoreFile         string        // Path to the template config file for node-cache
	CoreFile             string        // Path to config file used by node-cache
	KubednsCMPath        string        // Directory where kube-dns configmap will be mounted
	UpstreamSvcName      string        // Name of the service whose clusterIP is the upstream for node-cache for cluster domain
	HealthPort           string        // port for the healthcheck
	SetupIptables        bool
	SetupEbtables        bool
	SkipTeardown         bool // Indicates whether the iptables rules and interface should be torn down
}

type iptablesRule struct {
	table utiliptables.Table
	chain utiliptables.Chain
	args  []string
}

type ebtablesRule struct {
	table utilebtables.Table
	chain utilebtables.Chain
	args  []string
}

// CacheApp contains all the config required to run node-cache.
type CacheApp struct {
	iptables      utiliptables.Interface
	iptablesRules []iptablesRule
	ebtables      utilebtables.Interface
	ebtablesRules []ebtablesRule
	params        *ConfigParams
	netifHandle   *netif.NetifManager
	kubednsConfig *options.KubeDNSConfig
	exitChan      chan struct{} // Channel to terminate background goroutines
	clusterDNSIP  net.IP
}

func isLockedErr(err error) bool {
	return strings.Contains(err.Error(), "holding the xtables lock")
}

// Init initializes the parameters and networking setup necessary to run node-cache
func (c *CacheApp) Init() {
	if c.params.SetupInterface {
		c.netifHandle = netif.NewNetifManager(c.params.LocalIPs)
	}
	if c.params.SetupIptables {
		c.initIptables()
	}
	if c.params.SetupEbtables {
		c.initEbtables()
	}
	initMetrics(c.params.MetricsListenAddress)
	// Write the config file from template.
	// this is required in case there is no or erroneous kube-dns configpath specified.
	c.updateCorefile(&config.Config{})
	// Initialize periodic sync for node-local-dns, kube-dns configmap.
	c.initDNSConfigSync()
	// Setup only the network interface during this init. IPTables will be setup via runPeriodic.
	// This is to ensure that iptables rules don't get setup if the cache(coreDNS) is unable to startup due to config
	// error, port conflicts or other reasons.
	setupIptables := c.params.SetupIptables
	c.params.SetupIptables = false
	c.setupNetworking()
	c.params.SetupIptables = setupIptables
}

// isIPv6 return if the node-cache is working in IPv6 mode
// LocalIPs are guaranteed to have the same family
func (c *CacheApp) isIPv6() bool {
	if len(c.params.LocalIPs) > 0 {
		return utilnet.IsIPv6(c.params.LocalIPs[0])
	}
	return false
}

func (c *CacheApp) initIptables() {
	// using the localIPStr param since we need ip strings here
	for _, localIP := range strings.Split(c.params.LocalIPStr, ",") {
		c.iptablesRules = append(c.iptablesRules, []iptablesRule{
			// Match traffic destined for localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
			{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainPrerouting, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			// There are rules in filter table to allow tracked connections to be accepted. Since we skipped connection tracking,
			// need these additional filter table rules.
			{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "ACCEPT"}},
			{utiliptables.TableFilter, utiliptables.ChainInput, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "ACCEPT"}},
			// Match traffic from localIp:localPort and set the flows to be NOTRACKED, this skips connection tracking
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "NOTRACK"}},
			// Additional filter table rules for traffic frpm localIp:localPort
			{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "ACCEPT"}},
			{utiliptables.TableFilter, utiliptables.ChainOutput, []string{"-p", "udp", "-s", localIP,
				"--sport", c.params.LocalPort, "-j", "ACCEPT"}},
			// Skip connection tracking for requests to nodelocalDNS that are locally generated, example - by hostNetwork pods
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "udp", "-d", localIP,
				"--dport", c.params.LocalPort, "-j", "NOTRACK"}},
			// skip connection tracking for healthcheck requests generated by liveness probe to health plugin
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-d", localIP,
				"--dport", c.params.HealthPort, "-j", "NOTRACK"}},
			{utiliptables.Table("raw"), utiliptables.ChainOutput, []string{"-p", "tcp", "-s", localIP,
				"--sport", c.params.HealthPort, "-j", "NOTRACK"}},
		}...)
	}
	c.iptables = newIPTables(c.isIPv6())
}

func newIPTables(isIPv6 bool) utiliptables.Interface {
	execer := utilexec.New()
	dbus := dbus.New()
	protocol := utiliptables.ProtocolIpv4
	if isIPv6 {
		protocol = utiliptables.ProtocolIpv6
	}
	return utiliptables.New(execer, dbus, protocol)
}

func (c *CacheApp) initEbtables() {
	protocol := "IPv4"
	if c.isIPv6() {
		protocol = "IPv6"
	}
	// using the localIPStr param since we need ip strings here
	for _, localIP := range strings.Split(c.params.LocalIPStr, ",") {
		c.ebtablesRules = append(c.ebtablesRules, []ebtablesRule{
			// Match traffic destined for localIp and use the MAC address of the bridge port as destination address
			{utilebtables.TableBroute, utilebtables.ChainBrouting, []string{"-p", protocol, "--ip-dst", localIP,
				"-j", "redirect"}},
		}...)
	}
	c.ebtables = newEBTables()
}

func newEBTables() utilebtables.Interface {
	execer := utilexec.New()
	return utilebtables.New(execer)
}

// TeardownNetworking removes all custom iptables rules and network interface added by node-cache
func (c *CacheApp) TeardownNetworking() error {
	clog.Infof("Tearing down")
	if c.exitChan != nil {
		// Stop the goroutine that periodically checks for iptables rules/dummy interface
		// exitChan is a buffered channel of size 1, so this will not block
		c.exitChan <- struct{}{}
	}
	var err error
	if c.params.SetupInterface {
		err = c.netifHandle.RemoveDummyDevice(c.params.InterfaceName)
	}
	if c.params.SetupIptables {
		for _, rule := range c.iptablesRules {
			exists := true
			for exists == true {
				c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
				exists, _ = c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			}
			// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
			c.iptables.DeleteRule(rule.table, rule.chain, rule.args...)
		}
		if c.params.SetupEbtables {
			for _, rule := range c.ebtablesRules {
				exists := true
				for exists == true {
					c.ebtables.DeleteRule(rule.table, rule.chain, rule.args...)
					exists, _ = c.ebtables.EnsureRule(utilebtables.Append, rule.table, rule.chain, rule.args...)
				}
				// Delete the rule one last time since EnsureRule creates the rule if it doesn't exist
				c.ebtables.DeleteRule(rule.table, rule.chain, rule.args...)
			}
		}
	}
	return err
}

func (c *CacheApp) setupNetworking() {
	if c.params.SetupIptables {
		for _, rule := range c.iptablesRules {
			exists, err := c.iptables.EnsureRule(utiliptables.Prepend, rule.table, rule.chain, rule.args...)
			switch {
			case exists:
				// debug messages can be printed by including "debug" plugin in coreFile.
				clog.Debugf("iptables rule %v for nodelocaldns already exists", rule)
				continue
			case err == nil:
				clog.Infof("Added back nodelocaldns rule - %v", rule)
				continue
				// if we got here, either iptables check failed or adding rule back failed.
			case isLockedErr(err):
				clog.Infof("Error checking/adding iptables rule %v, due to xtables lock in use, retrying in %v", rule, c.params.Interval)
				setupErrCount.WithLabelValues("iptables_lock").Inc()
			default:
				clog.Errorf("Error adding iptables rule %v - %s", rule, err)
				setupErrCount.WithLabelValues("iptables").Inc()
			}
		}
	}

	if c.params.SetupEbtables {
		for _, rule := range c.ebtablesRules {
			exists, err := c.ebtables.EnsureRule(utilebtables.Append, rule.table, rule.chain, rule.args...)
			switch {
			case exists:
				// debug messages can be printed by including "debug" plugin in coreFile.
				clog.Debugf("ebtables rule %v for nodelocaldns already exists", rule)
				continue
			case err == nil:
				clog.Infof("Added back ebtables rule - %v", rule)
				continue
				// if we got here, either ebtables check failed or adding rule back failed.
			default:
				clog.Errorf("Error adding ebtables rule %v - %s", rule, err)
				setupErrCount.WithLabelValues("ebtables").Inc()
			}
		}
	}

	if c.params.SetupInterface {
		exists, err := c.netifHandle.EnsureDummyDevice(c.params.InterfaceName)
		if !exists {
			if err != nil {
				clog.Errorf("Failed to add non-existent interface %s: %s", c.params.InterfaceName, err)
				setupErrCount.WithLabelValues("interface_add").Inc()
			}
			clog.Infof("Added interface - %s", c.params.InterfaceName)
		}
		if err != nil {
			clog.Errorf("Error checking dummy device %s - %s", c.params.InterfaceName, err)
			setupErrCount.WithLabelValues("interface_check").Inc()
		}
	}
}

func (c *CacheApp) runPeriodic() {
	c.exitChan = make(chan struct{}, 1)
	tick := time.NewTicker(c.params.Interval * time.Second)
	c.setupNetworking()
	for {
		select {
		case <-tick.C:
			c.setupNetworking()
		case <-c.exitChan:
			clog.Warningf("Exiting iptables/interface check goroutine")
			return
		}
	}
}

// RunApp invokes the background checks and runs coreDNS as a cache
func (c *CacheApp) RunApp() {
	go c.runPeriodic()
	coremain.Run()
	// Unlikely to reach here, if we did it is because coremain exited and the signal was not trapped.
	clog.Errorf("Untrapped signal, tearing down")
	c.TeardownNetworking()
}

// NewCacheApp returns a new instance of CacheApp by applying the specified config params.
func NewCacheApp(params *ConfigParams) (*CacheApp, error) {
	c := &CacheApp{params: params, kubednsConfig: options.NewKubeDNSConfig()}
	c.clusterDNSIP = net.ParseIP(os.ExpandEnv(toSvcEnv(params.UpstreamSvcName)))
	if c.clusterDNSIP == nil {
		clog.Warningf("Unable to lookup IP address of Upstream service %s, env %s `%s`", params.UpstreamSvcName, toSvcEnv(params.UpstreamSvcName), os.ExpandEnv(toSvcEnv(params.UpstreamSvcName)))
	}
	return c, nil
}

// toSvcEnv converts service name to the corresponding ENV variable. This is exposed in every pod and its value is the clusterIP.
// https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables
func toSvcEnv(svcName string) string {
	envName := strings.Replace(svcName, "-", "_", -1)
	return "$" + strings.ToUpper(envName) + "_SERVICE_HOST"
}
<|endoftext|>