<commit_before>package mountlib

import (
	"context"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/rc"
	"github.com/rclone/rclone/vfs"
	"github.com/rclone/rclone/vfs/vfsflags"
)

// MountInfo defines the configuration for a mount
type MountInfo struct {
	unmountFn  UnmountFn
	MountPoint string    `json:"MountPoint"`
	MountedOn  time.Time `json:"MountedOn"`
	Fs         string    `json:"Fs"`
}

var (
	// mutex to protect all the variables in this block
	mountMu sync.Mutex
	// Mount functions available
	mountFns = map[string]MountFn{}
	// Map of mounted path => MountInfo
	liveMounts = map[string]MountInfo{}
)

// AddRc adds mount and unmount functionality to rc
func AddRc(mountUtilName string, mountFunction MountFn) {
	mountMu.Lock()
	defer mountMu.Unlock()
	// rcMount allows the mount command to be run from rc
	mountFns[mountUtilName] = mountFunction
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/mount",
		AuthRequired: true,
		Fn:           mountRc,
		Title:        "Create a new mount point",
		Help: `rclone allows Linux, FreeBSD, macOS and Windows to mount any of
Rclone's cloud storage systems as a file system with FUSE.

If no mountType is provided, the priority is given as follows: 1. mount 2.cmount 3.mount2

This takes the following parameters

- fs - a remote path to be mounted (required)
- mountPoint: valid path on the local machine (required)
- mountType: One of the values (mount, cmount, mount2) specifies the mount implementation to use
- mountOpt: a JSON object with Mount options in.
- vfsOpt: a JSON object with VFS options in.

Eg

    rclone rc mount/mount fs=mydrive: mountPoint=/home/<user>/mountPoint
    rclone rc mount/mount fs=mydrive: mountPoint=/home/<user>/mountPoint mountType=mount
    rclone rc mount/mount fs=TestDrive: mountPoint=/mnt/tmp vfsOpt='{"CacheMode": 2}' mountOpt='{"AllowOther": true}'

The vfsOpt are as described in options/get and can be seen in the the
"vfs" section when running and the mountOpt can be seen in the "mount" section.

    rclone rc options/get
`,
	})
}

// mountRc allows the mount command to be run from rc
func mountRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountPoint, err := in.GetString("mountPoint")
	if err != nil {
		return nil, err
	}

	vfsOpt := vfsflags.Opt
	err = in.GetStructMissingOK("vfsOpt", &vfsOpt)
	if err != nil {
		return nil, err
	}

	opt := Opt
	err = in.GetStructMissingOK("mountOpt", &opt)
	if err != nil {
		return nil, err
	}

	mountType, err := in.GetString("mountType")

	mountMu.Lock()
	defer mountMu.Unlock()

	if err != nil || mountType == "" {
		if mountFns["mount"] != nil {
			mountType = "mount"
		} else if mountFns["cmount"] != nil {
			mountType = "cmount"
		} else if mountFns["mount2"] != nil {
			mountType = "mount2"
		}
	}

	// Get Fs.fs to be mounted from fs parameter in the params
	fdst, err := rc.GetFs(in)
	if err != nil {
		return nil, err
	}

	if mountFns[mountType] != nil {
		VFS := vfs.New(fdst, &vfsOpt)
		_, unmountFn, err := mountFns[mountType](VFS, mountPoint, &opt)

		if err != nil {
			log.Printf("mount FAILED: %v", err)
			return nil, err
		}
		// Add mount to list if mount point was successfully created
		liveMounts[mountPoint] = MountInfo{
			unmountFn:  unmountFn,
			MountedOn:  time.Now(),
			Fs:         fdst.Name(),
			MountPoint: mountPoint,
		}

		fs.Debugf(nil, "Mount for %s created at %s using %s", fdst.String(), mountPoint, mountType)
		return nil, nil
	}
	return nil, errors.New("Mount Option specified is not registered, or is invalid")
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/unmount",
		AuthRequired: true,
		Fn:           unMountRc,
		Title:        "Unmount selected active mount",
		Help: `
rclone allows Linux, FreeBSD, macOS and Windows to
mount any of Rclone's cloud storage systems as a file system with
FUSE.

This takes the following parameters

- mountPoint: valid path on the local machine where the mount was created (required)

Eg

    rclone rc mount/unmount mountPoint=/home/<user>/mountPoint
`,
	})
}

// unMountRc allows the umount command to be run from rc
func unMountRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountPoint, err := in.GetString("mountPoint")
	if err != nil {
		return nil, err
	}
	mountMu.Lock()
	defer mountMu.Unlock()
	err = performUnMount(mountPoint)
	if err != nil {
		return nil, err
	}
	return nil, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/types",
		AuthRequired: true,
		Fn:           mountTypesRc,
		Title:        "Show all possible mount types",
		Help: `This shows all possible mount types and returns them as a list.

This takes no parameters and returns

- mountTypes: list of mount types

The mount types are strings like "mount", "mount2", "cmount" and can
be passed to mount/mount as the mountType parameter.

Eg

    rclone rc mount/types
`,
	})
}

// mountTypesRc returns a list of available mount types.
func mountTypesRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	var mountTypes = []string{}
	mountMu.Lock()
	defer mountMu.Unlock()
	for mountType := range mountFns {
		mountTypes = append(mountTypes, mountType)
	}
	sort.Strings(mountTypes)
	return rc.Params{
		"mountTypes": mountTypes,
	}, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/listmounts",
		AuthRequired: true,
		Fn:           listMountsRc,
		Title:        "Show current mount points",
		Help: `This shows currently mounted points, which can be used for performing an unmount

This takes no parameters and returns

- mountPoints: list of current mount points

Eg

    rclone rc mount/listmounts
`,
	})
}

// listMountsRc returns a list of current mounts
func listMountsRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	var mountTypes = []MountInfo{}
	mountMu.Lock()
	defer mountMu.Unlock()
	for _, a := range liveMounts {
		mountTypes = append(mountTypes, a)
	}
	return rc.Params{
		"mountPoints": mountTypes,
	}, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/unmountall",
		AuthRequired: true,
		Fn:           unmountAll,
		Title:        "Show current mount points",
		Help: `This shows currently mounted points, which can be used for performing an unmount

This takes no parameters and returns error if unmount does not succeed.

Eg

    rclone rc mount/unmountall
`,
	})
}

// unmountAll unmounts all the created mounts
func unmountAll(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountMu.Lock()
	defer mountMu.Unlock()
	for key, mountInfo := range liveMounts {
		err = performUnMount(mountInfo.MountPoint)
		if err != nil {
			fs.Debugf(nil, "Couldn't unmount : %s", key)
			return nil, err
		}
	}
	return nil, nil
}

// performUnMount unmounts the specified mountPoint
func performUnMount(mountPoint string) (err error) {
	mountInfo, ok := liveMounts[mountPoint]
	if ok {
		err := mountInfo.unmountFn()
		if err != nil {
			return err
		}
		delete(liveMounts, mountPoint)
	} else {
		return errors.New("mount not found")
	}
	return nil
}
<commit_msg>mount: add VFS and Mount options to mount/listmounts<commit_after>package mountlib

import (
	"context"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/pkg/errors"
	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/rc"
	"github.com/rclone/rclone/vfs"
	"github.com/rclone/rclone/vfs/vfscommon"
	"github.com/rclone/rclone/vfs/vfsflags"
)

// MountInfo defines the configuration for a mount
type MountInfo struct {
	unmountFn  UnmountFn
	MountPoint string    `json:"MountPoint"`
	MountedOn  time.Time `json:"MountedOn"`
	Fs         string    `json:"Fs"`
	MountOpt   *Options
	VFSOpt     *vfscommon.Options
}

var (
	// mutex to protect all the variables in this block
	mountMu sync.Mutex
	// Mount functions available
	mountFns = map[string]MountFn{}
	// Map of mounted path => MountInfo
	liveMounts = map[string]MountInfo{}
)

// AddRc adds mount and unmount functionality to rc
func AddRc(mountUtilName string, mountFunction MountFn) {
	mountMu.Lock()
	defer mountMu.Unlock()
	// rcMount allows the mount command to be run from rc
	mountFns[mountUtilName] = mountFunction
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/mount",
		AuthRequired: true,
		Fn:           mountRc,
		Title:        "Create a new mount point",
		Help: `rclone allows Linux, FreeBSD, macOS and Windows to mount any of
Rclone's cloud storage systems as a file system with FUSE.

If no mountType is provided, the priority is given as follows: 1. mount 2.cmount 3.mount2

This takes the following parameters

- fs - a remote path to be mounted (required)
- mountPoint: valid path on the local machine (required)
- mountType: One of the values (mount, cmount, mount2) specifies the mount implementation to use
- mountOpt: a JSON object with Mount options in.
- vfsOpt: a JSON object with VFS options in.

Eg

    rclone rc mount/mount fs=mydrive: mountPoint=/home/<user>/mountPoint
    rclone rc mount/mount fs=mydrive: mountPoint=/home/<user>/mountPoint mountType=mount
    rclone rc mount/mount fs=TestDrive: mountPoint=/mnt/tmp vfsOpt='{"CacheMode": 2}' mountOpt='{"AllowOther": true}'

The vfsOpt are as described in options/get and can be seen in the the
"vfs" section when running and the mountOpt can be seen in the "mount" section.

    rclone rc options/get
`,
	})
}

// mountRc allows the mount command to be run from rc
func mountRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountPoint, err := in.GetString("mountPoint")
	if err != nil {
		return nil, err
	}

	vfsOpt := vfsflags.Opt
	err = in.GetStructMissingOK("vfsOpt", &vfsOpt)
	if err != nil {
		return nil, err
	}

	mountOpt := Opt
	err = in.GetStructMissingOK("mountOpt", &mountOpt)
	if err != nil {
		return nil, err
	}

	mountType, err := in.GetString("mountType")

	mountMu.Lock()
	defer mountMu.Unlock()

	if err != nil || mountType == "" {
		if mountFns["mount"] != nil {
			mountType = "mount"
		} else if mountFns["cmount"] != nil {
			mountType = "cmount"
		} else if mountFns["mount2"] != nil {
			mountType = "mount2"
		}
	}

	// Get Fs.fs to be mounted from fs parameter in the params
	fdst, err := rc.GetFs(in)
	if err != nil {
		return nil, err
	}

	if mountFns[mountType] != nil {
		VFS := vfs.New(fdst, &vfsOpt)
		_, unmountFn, err := mountFns[mountType](VFS, mountPoint, &mountOpt)

		if err != nil {
			log.Printf("mount FAILED: %v", err)
			return nil, err
		}
		// Add mount to list if mount point was successfully created
		liveMounts[mountPoint] = MountInfo{
			unmountFn:  unmountFn,
			MountedOn:  time.Now(),
			Fs:         fdst.Name(),
			MountPoint: mountPoint,
			VFSOpt:     &vfsOpt,
			MountOpt:   &mountOpt,
		}

		fs.Debugf(nil, "Mount for %s created at %s using %s", fdst.String(), mountPoint, mountType)
		return nil, nil
	}
	return nil, errors.New("Mount Option specified is not registered, or is invalid")
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/unmount",
		AuthRequired: true,
		Fn:           unMountRc,
		Title:        "Unmount selected active mount",
		Help: `
rclone allows Linux, FreeBSD, macOS and Windows to
mount any of Rclone's cloud storage systems as a file system with
FUSE.

This takes the following parameters

- mountPoint: valid path on the local machine where the mount was created (required)

Eg

    rclone rc mount/unmount mountPoint=/home/<user>/mountPoint
`,
	})
}

// unMountRc allows the umount command to be run from rc
func unMountRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountPoint, err := in.GetString("mountPoint")
	if err != nil {
		return nil, err
	}
	mountMu.Lock()
	defer mountMu.Unlock()
	err = performUnMount(mountPoint)
	if err != nil {
		return nil, err
	}
	return nil, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/types",
		AuthRequired: true,
		Fn:           mountTypesRc,
		Title:        "Show all possible mount types",
		Help: `This shows all possible mount types and returns them as a list.

This takes no parameters and returns

- mountTypes: list of mount types

The mount types are strings like "mount", "mount2", "cmount" and can
be passed to mount/mount as the mountType parameter.

Eg

    rclone rc mount/types
`,
	})
}

// mountTypesRc returns a list of available mount types.
func mountTypesRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	var mountTypes = []string{}
	mountMu.Lock()
	defer mountMu.Unlock()
	for mountType := range mountFns {
		mountTypes = append(mountTypes, mountType)
	}
	sort.Strings(mountTypes)
	return rc.Params{
		"mountTypes": mountTypes,
	}, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/listmounts",
		AuthRequired: true,
		Fn:           listMountsRc,
		Title:        "Show current mount points",
		Help: `This shows currently mounted points, which can be used for performing an unmount

This takes no parameters and returns

- mountPoints: list of current mount points

Eg

    rclone rc mount/listmounts
`,
	})
}

// listMountsRc returns a list of current mounts
func listMountsRc(_ context.Context, in rc.Params) (out rc.Params, err error) {
	var mountTypes = []MountInfo{}
	mountMu.Lock()
	defer mountMu.Unlock()
	for _, a := range liveMounts {
		mountTypes = append(mountTypes, a)
	}
	return rc.Params{
		"mountPoints": mountTypes,
	}, nil
}

func init() {
	rc.Add(rc.Call{
		Path:         "mount/unmountall",
		AuthRequired: true,
		Fn:           unmountAll,
		Title:        "Show current mount points",
		Help: `This shows currently mounted points, which can be used for performing an unmount

This takes no parameters and returns error if unmount does not succeed.

Eg

    rclone rc mount/unmountall
`,
	})
}

// unmountAll unmounts all the created mounts
func unmountAll(_ context.Context, in rc.Params) (out rc.Params, err error) {
	mountMu.Lock()
	defer mountMu.Unlock()
	for key, mountInfo := range liveMounts {
		err = performUnMount(mountInfo.MountPoint)
		if err != nil {
			fs.Debugf(nil, "Couldn't unmount : %s", key)
			return nil, err
		}
	}
	return nil, nil
}

// performUnMount unmounts the specified mountPoint
func performUnMount(mountPoint string) (err error) {
	mountInfo, ok := liveMounts[mountPoint]
	if ok {
		err := mountInfo.unmountFn()
		if err != nil {
			return err
		}
		delete(liveMounts, mountPoint)
	} else {
		return errors.New("mount not found")
	}
	return nil
}
<|endoftext|>


<commit_before>// Ottily executes a javascript snippet on each line of an input file in parallel.
//
// Noop:
//
//     $ ottily datasets/simple.ldj
//     {"name": "ottily", "language": "Golang"}
//
// Inline script with -e:
//
//     $ ottily -e 'output=input.length' datasets/simple.ldj
//	   40
//
//     $ ottily -e 'o=JSON.parse(input); o["language"] = "Go"; output=JSON.stringify(o);' datasets/simple.ldj
//     {"language":"Go","name":"ottily"}
//
// Pass a script file:
//
//     $ ottily -s scripts/classified.js datasets/simple.ldj
//     CLASSIFIED
//
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"

	"github.com/robertkrimen/otto"
)

const NOOP_SCRIPT = "output = input"
const VERSION = "0.1.0"

func Worker(lines, out chan *string, script string, wg *sync.WaitGroup) {
	defer wg.Done()
	vm := otto.New()

	compiled, err := vm.Compile("", script)
	if err != nil {
		log.Fatal(err)
	}

	for line := range lines {
		vm.Set("input", *line)
		_, err := vm.Run(compiled)
		if err != nil {
			log.Fatal(err)
		}
		result, err := vm.Get("output")
		if result == otto.NullValue() {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		r := result.String()
		out <- &r
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan *string, done chan bool) {
	for s := range in {
		writer.Write([]byte(*s))
		writer.Write([]byte("\n"))
	}
	done <- true
}

func main() {
	script := flag.String("s", "", "script to execute on each line of input")
	execute := flag.String("e", "", "execute argument on each line of input")
	numWorkers := flag.Int("w", runtime.NumCPU(), "number of workers")
	version := flag.Bool("v", false, "prints current program version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	if flag.NArg() < 1 {
		log.Fatal("input file required")
	}

	content := NOOP_SCRIPT

	if *script != "" {
		ff, err := os.Open(*script)
		if err != nil {
			log.Fatal(err)
		}
		b, err := ioutil.ReadAll(ff)
		if err != nil {
			log.Fatal(err)
		}
		content = string(b)
	}

	if *execute != "" {
		content = *execute
	}

	ff, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}
	defer ff.Close()
	reader := bufio.NewReader(ff)

	if *numWorkers > 0 {
		runtime.GOMAXPROCS(*numWorkers)
	}

	queue := make(chan *string)
	out := make(chan *string)
	done := make(chan bool)
	var wg sync.WaitGroup

	writer := bufio.NewWriter(os.Stdout)
	defer writer.Flush()
	go FanInWriter(writer, out, done)

	for i := 0; i < *numWorkers; i++ {
		wg.Add(1)
		go Worker(queue, out, content, &wg)
	}

	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		line = strings.TrimSpace(line)
		queue <- &line
	}
	close(queue)
	wg.Wait()
	close(out)
	<-done
}
<commit_msg>support script preloading<commit_after>// Ottily executes a javascript snippet on each line of an input file in parallel.
//
// Noop:
//
//     $ ottily datasets/simple.ldj
//     {"name": "ottily", "language": "Golang"}
//
// Inline script with -e:
//
//     $ ottily -e 'output=input.length' datasets/simple.ldj
//	   40
//
//     $ ottily -e 'o=JSON.parse(input); o["language"] = "Go"; output=JSON.stringify(o);' datasets/simple.ldj
//     {"language":"Go","name":"ottily"}
//
// Pass a script file:
//
//     $ ottily -s scripts/classified.js datasets/simple.ldj
//     CLASSIFIED
//
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"

	"github.com/robertkrimen/otto"
)

const NOOP_SCRIPT = "output = input"
const VERSION = "0.1.0"

func Worker(lines, out chan *string, script, preload string, wg *sync.WaitGroup) {
	defer wg.Done()
	vm := otto.New()

	if preload != "" {
		_, err := vm.Run(preload)
		if err != nil {
			log.Fatal(err)
		}
	}

	compiled, err := vm.Compile("", script)
	if err != nil {
		log.Fatal(err)
	}

	for line := range lines {
		vm.Set("input", *line)
		_, err := vm.Run(compiled)
		if err != nil {
			log.Fatal(err)
		}
		result, err := vm.Get("output")
		if result == otto.NullValue() {
			continue
		}
		if err != nil {
			log.Fatal(err)
		}
		r := result.String()
		out <- &r
	}
}

// FanInWriter writes the channel content to the writer
func FanInWriter(writer io.Writer, in chan *string, done chan bool) {
	for s := range in {
		writer.Write([]byte(*s))
		writer.Write([]byte("\n"))
	}
	done <- true
}

func main() {
	preload := flag.String("p", "", "include this js file before processing")
	script := flag.String("s", "", "script to execute on each line of input")
	execute := flag.String("e", "", "execute argument on each line of input")
	numWorkers := flag.Int("w", runtime.NumCPU(), "number of workers")
	version := flag.Bool("v", false, "prints current program version")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile to file")

	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if *version {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	if flag.NArg() < 1 {
		log.Fatal("input file required")
	}

	preloadContent := ""

	if *preload != "" {
		ff, err := os.Open(*preload)
		if err != nil {
			log.Fatal(err)
		}
		b, err := ioutil.ReadAll(ff)
		if err != nil {
			log.Fatal(err)
		}
		preloadContent = string(b)
	}

	content := NOOP_SCRIPT

	if *script != "" {
		ff, err := os.Open(*script)
		if err != nil {
			log.Fatal(err)
		}
		b, err := ioutil.ReadAll(ff)
		if err != nil {
			log.Fatal(err)
		}
		content = string(b)
	}

	if *execute != "" {
		content = *execute
	}

	ff, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}
	defer ff.Close()
	reader := bufio.NewReader(ff)

	if *numWorkers > 0 {
		runtime.GOMAXPROCS(*numWorkers)
	}

	queue := make(chan *string)
	out := make(chan *string)
	done := make(chan bool)
	var wg sync.WaitGroup

	writer := bufio.NewWriter(os.Stdout)
	defer writer.Flush()
	go FanInWriter(writer, out, done)

	for i := 0; i < *numWorkers; i++ {
		wg.Add(1)
		go Worker(queue, out, content, preloadContent, &wg)
	}

	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		line = strings.TrimSpace(line)
		queue <- &line
	}
	close(queue)
	wg.Wait()
	close(out)
	<-done
}
<|endoftext|>


<commit_before>// Copyright 2015 Sorint.lab
//
// 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 main

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"
	"github.com/sorintlab/stolon/pkg/flagutil"
	slog "github.com/sorintlab/stolon/pkg/log"
	"github.com/sorintlab/stolon/pkg/store"
	"github.com/sorintlab/stolon/pkg/util"

	"github.com/davecgh/go-spew/spew"
	"github.com/sorintlab/pollon"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

var log = slog.S()

var cmdProxy = &cobra.Command{
	Use: "stolon-proxy",
	Run: proxy,
}

type config struct {
	storeBackend       string
	storeEndpoints     string
	storeCertFile      string
	storeKeyFile       string
	storeCAFile        string
	storeSkipTlsVerify bool
	clusterName        string
	listenAddress      string
	port               string
	stopListening      bool
	logLevel           string
	debug              bool
}

var cfg config

func init() {
	cmdProxy.PersistentFlags().StringVar(&cfg.storeBackend, "store-backend", "", "store backend type (etcd or consul)")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeEndpoints, "store-endpoints", "", "a comma-delimited list of store endpoints (use https scheme for tls communication) (defaults: http://127.0.0.1:2379 for etcd, http://127.0.0.1:8500 for consul)")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeCertFile, "store-cert-file", "", "certificate file for client identification to the store")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeKeyFile, "store-key", "", "private key file for client identification to the store")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeCAFile, "store-ca-file", "", "verify certificates of HTTPS-enabled store servers using this CA bundle")
	cmdProxy.PersistentFlags().BoolVar(&cfg.storeSkipTlsVerify, "store-skip-tls-verify", false, "skip store certificate verification (insecure!!!)")
	cmdProxy.PersistentFlags().StringVar(&cfg.clusterName, "cluster-name", "", "cluster name")
	cmdProxy.PersistentFlags().StringVar(&cfg.listenAddress, "listen-address", "127.0.0.1", "proxy listening address")
	cmdProxy.PersistentFlags().StringVar(&cfg.port, "port", "5432", "proxy listening port")
	cmdProxy.PersistentFlags().BoolVar(&cfg.stopListening, "stop-listening", true, "stop listening on store error")
	cmdProxy.PersistentFlags().StringVar(&cfg.logLevel, "log-level", "info", "debug, info (default), warn or error")
	cmdProxy.PersistentFlags().BoolVar(&cfg.debug, "debug", false, "enable debug logging")

	cmdProxy.PersistentFlags().MarkDeprecated("debug", "use --log-level=debug instead")
}

func stderr(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stderr, strings.TrimSuffix(out, "\n"))
}

func stdout(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stdout, strings.TrimSuffix(out, "\n"))
}

func die(format string, a ...interface{}) {
	stderr(format, a...)
	os.Exit(1)
}

type ClusterChecker struct {
	uid           string
	listenAddress string
	port          string

	stopListening bool

	listener         *net.TCPListener
	pp               *pollon.Proxy
	e                *store.StoreManager
	endPollonProxyCh chan error
}

func NewClusterChecker(uid string, cfg config) (*ClusterChecker, error) {
	storePath := filepath.Join(common.StoreBasePath, cfg.clusterName)

	kvstore, err := store.NewStore(store.Config{
		Backend:       store.Backend(cfg.storeBackend),
		Endpoints:     cfg.storeEndpoints,
		CertFile:      cfg.storeCertFile,
		KeyFile:       cfg.storeKeyFile,
		CAFile:        cfg.storeCAFile,
		SkipTLSVerify: cfg.storeSkipTlsVerify,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot create store: %v", err)
	}
	e := store.NewStoreManager(kvstore, storePath)

	return &ClusterChecker{
		uid:              uid,
		listenAddress:    cfg.listenAddress,
		port:             cfg.port,
		stopListening:    cfg.stopListening,
		e:                e,
		endPollonProxyCh: make(chan error),
	}, nil
}

func (c *ClusterChecker) startPollonProxy() error {
	if c.pp != nil {
		return nil
	}

	log.Infow("Starting proxying")
	addr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(cfg.listenAddress, cfg.port))
	if err != nil {
		return fmt.Errorf("error resolving tcp addr %q: %v", addr.String(), err)
	}

	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return fmt.Errorf("error listening on tcp addr %q: %v", addr.String(), err)
	}

	pp, err := pollon.NewProxy(listener)
	if err != nil {
		return fmt.Errorf("error creating pollon proxy: %v", err)
	}
	c.pp = pp
	c.listener = listener

	go func() {
		c.endPollonProxyCh <- c.pp.Start()
	}()

	return nil
}

func (c *ClusterChecker) stopPollonProxy() {
	if c.pp != nil {
		log.Infow("Stopping listening")
		c.pp.Stop()
		c.pp = nil
		c.listener.Close()
		c.listener = nil
	}
}

func (c *ClusterChecker) sendPollonConfData(confData pollon.ConfData) {
	if c.pp != nil {
		c.pp.C <- confData
	}
}

func (c *ClusterChecker) SetProxyInfo(e *store.StoreManager, generation int64, ttl time.Duration) error {
	proxyInfo := &cluster.ProxyInfo{
		UID:        c.uid,
		Generation: generation,
	}
	log.Debugf("proxyInfo dump: %s", spew.Sdump(proxyInfo))

	if err := c.e.SetProxyInfo(proxyInfo, ttl); err != nil {
		return err
	}
	return nil
}

// Check reads the cluster data and applies the right pollon configuration.
func (c *ClusterChecker) Check() error {
	cd, _, err := c.e.GetClusterData()
	if err != nil {
		return fmt.Errorf("cannot get cluster data: %v", err)
	}

	// Start pollon if not active
	if err = c.startPollonProxy(); err != nil {
		return fmt.Errorf("failed to start proxy: %v", err)
	}

	log.Debugf("cd dump: %s", spew.Sdump(cd))
	if cd == nil {
		log.Infow("no clusterdata available, closing connections to master")
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return nil
	}
	if cd.FormatVersion != cluster.CurrentCDFormatVersion {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return fmt.Errorf("unsupported clusterdata format version: %d", cd.FormatVersion)
	}
	if err = cd.Cluster.Spec.Validate(); err != nil {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return fmt.Errorf("clusterdata validation failed: %v", err)
	}

	proxy := cd.Proxy
	if proxy == nil {
		log.Infow("no proxy object available, closing connections to master")
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		// ignore errors on setting proxy info
		if err = c.SetProxyInfo(c.e, cluster.NoGeneration, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
			log.Errorw("failed to update proxyInfo", zap.Error(err))
		}
		return nil
	}

	db, ok := cd.DBs[proxy.Spec.MasterDBUID]
	if !ok {
		log.Infow("no db object available, closing connections to master", "db", proxy.Spec.MasterDBUID)
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		// ignore errors on setting proxy info
		if err = c.SetProxyInfo(c.e, proxy.Generation, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
			log.Errorw("failed to update proxyInfo", zap.Error(err))
		}
		return nil
	}

	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", db.Status.ListenAddress, db.Status.Port))
	if err != nil {
		log.Errorw("cannot resolve db address", zap.Error(err))
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return nil
	}
	log.Infow("master address", "address", addr)
	if err = c.SetProxyInfo(c.e, proxy.Generation, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
		// if we failed to update our proxy info when a master is defined we
		// cannot ignore this error since the sentinel won't know that we exist
		// and are sending connections to a master so, when electing a new
		// master, it'll not wait for us to close connections to the old one.
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		log.Errorw("failed to update proxyInfo", zap.Error(err))
		return nil
	}

	// start proxing only if we are inside enabledProxies, this ensures that the
	// sentinel has read our proxyinfo and knows we are alive
	if util.StringInSlice(proxy.Spec.EnabledProxies, c.uid) {
		log.Infow("proxying to master address", "address", addr)
		c.sendPollonConfData(pollon.ConfData{DestAddr: addr})
	} else {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
	}

	return nil
}

func (c *ClusterChecker) TimeoutChecker(checkOkCh chan struct{}) error {
	timeoutTimer := time.NewTimer(cluster.DefaultProxyTimeoutInterval)

	for true {
		select {
		case <-timeoutTimer.C:
			log.Infow("check timeout timer fired")
			// if the check timeouts close all connections and stop listening
			// (for example to avoid load balancers forward connections to us
			// since we aren't ready or in a bad state)
			c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
			if c.stopListening {
				c.stopPollonProxy()
			}

		case <-checkOkCh:
			log.Debugw("check ok message received")

			// ignore if stop succeeded or not due to timer already expired
			timeoutTimer.Stop()
			timeoutTimer = time.NewTimer(cluster.DefaultProxyTimeoutInterval)
		}
	}
	return nil
}

func (c *ClusterChecker) Start() error {
	checkOkCh := make(chan struct{})
	checkCh := make(chan error)
	timerCh := time.NewTimer(0).C

	go c.TimeoutChecker(checkOkCh)

	for true {
		select {
		case <-timerCh:
			go func() {
				checkCh <- c.Check()
			}()
		case err := <-checkCh:
			if err != nil {
				// don't report check ok since it returned an error
				log.Debugw("check function error", zap.Error(err))
			} else {
				// report that check was ok
				checkOkCh <- struct{}{}
			}
			timerCh = time.NewTimer(cluster.DefaultProxyCheckInterval).C
		case err := <-c.endPollonProxyCh:
			if err != nil {
				return fmt.Errorf("proxy error: %v", err)
			}
		}
	}
	return nil
}

func main() {
	flagutil.SetFlagsFromEnv(cmdProxy.PersistentFlags(), "STPROXY")

	cmdProxy.Execute()
}

func proxy(cmd *cobra.Command, args []string) {
	switch cfg.logLevel {
	case "error":
		slog.SetLevel(zap.ErrorLevel)
	case "warn":
		slog.SetLevel(zap.WarnLevel)
	case "info":
		slog.SetLevel(zap.InfoLevel)
	case "debug":
		slog.SetLevel(zap.DebugLevel)
	default:
		die("invalid log level: %v", cfg.logLevel)
	}
	if cfg.debug {
		slog.SetDebug()
	}
	if slog.IsDebug() {
		stdlog := slog.StdLog()
		pollon.SetLogger(stdlog)
	}

	if cfg.clusterName == "" {
		die("cluster name required")
	}
	if cfg.storeBackend == "" {
		die("store backend type required")
	}

	uid := common.UID()
	log.Infow("proxy uid", "uid", uid)

	clusterChecker, err := NewClusterChecker(uid, cfg)
	if err != nil {
		die("cannot create cluster checker: %v", err)
	}
	if err = clusterChecker.Start(); err != nil {
		die("cluster checker ended with error: %v", err)
	}
}
<commit_msg>proxy: fix data race<commit_after>// Copyright 2015 Sorint.lab
//
// 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 main

import (
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"
	"github.com/sorintlab/stolon/pkg/flagutil"
	slog "github.com/sorintlab/stolon/pkg/log"
	"github.com/sorintlab/stolon/pkg/store"
	"github.com/sorintlab/stolon/pkg/util"

	"github.com/davecgh/go-spew/spew"
	"github.com/sorintlab/pollon"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

var log = slog.S()

var cmdProxy = &cobra.Command{
	Use: "stolon-proxy",
	Run: proxy,
}

type config struct {
	storeBackend       string
	storeEndpoints     string
	storeCertFile      string
	storeKeyFile       string
	storeCAFile        string
	storeSkipTlsVerify bool
	clusterName        string
	listenAddress      string
	port               string
	stopListening      bool
	logLevel           string
	debug              bool
}

var cfg config

func init() {
	cmdProxy.PersistentFlags().StringVar(&cfg.storeBackend, "store-backend", "", "store backend type (etcd or consul)")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeEndpoints, "store-endpoints", "", "a comma-delimited list of store endpoints (use https scheme for tls communication) (defaults: http://127.0.0.1:2379 for etcd, http://127.0.0.1:8500 for consul)")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeCertFile, "store-cert-file", "", "certificate file for client identification to the store")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeKeyFile, "store-key", "", "private key file for client identification to the store")
	cmdProxy.PersistentFlags().StringVar(&cfg.storeCAFile, "store-ca-file", "", "verify certificates of HTTPS-enabled store servers using this CA bundle")
	cmdProxy.PersistentFlags().BoolVar(&cfg.storeSkipTlsVerify, "store-skip-tls-verify", false, "skip store certificate verification (insecure!!!)")
	cmdProxy.PersistentFlags().StringVar(&cfg.clusterName, "cluster-name", "", "cluster name")
	cmdProxy.PersistentFlags().StringVar(&cfg.listenAddress, "listen-address", "127.0.0.1", "proxy listening address")
	cmdProxy.PersistentFlags().StringVar(&cfg.port, "port", "5432", "proxy listening port")
	cmdProxy.PersistentFlags().BoolVar(&cfg.stopListening, "stop-listening", true, "stop listening on store error")
	cmdProxy.PersistentFlags().StringVar(&cfg.logLevel, "log-level", "info", "debug, info (default), warn or error")
	cmdProxy.PersistentFlags().BoolVar(&cfg.debug, "debug", false, "enable debug logging")

	cmdProxy.PersistentFlags().MarkDeprecated("debug", "use --log-level=debug instead")
}

func stderr(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stderr, strings.TrimSuffix(out, "\n"))
}

func stdout(format string, a ...interface{}) {
	out := fmt.Sprintf(format, a...)
	fmt.Fprintln(os.Stdout, strings.TrimSuffix(out, "\n"))
}

func die(format string, a ...interface{}) {
	stderr(format, a...)
	os.Exit(1)
}

type ClusterChecker struct {
	uid           string
	listenAddress string
	port          string

	stopListening bool

	listener         *net.TCPListener
	pp               *pollon.Proxy
	e                *store.StoreManager
	endPollonProxyCh chan error

	pollonMutex sync.Mutex
}

func NewClusterChecker(uid string, cfg config) (*ClusterChecker, error) {
	storePath := filepath.Join(common.StoreBasePath, cfg.clusterName)

	kvstore, err := store.NewStore(store.Config{
		Backend:       store.Backend(cfg.storeBackend),
		Endpoints:     cfg.storeEndpoints,
		CertFile:      cfg.storeCertFile,
		KeyFile:       cfg.storeKeyFile,
		CAFile:        cfg.storeCAFile,
		SkipTLSVerify: cfg.storeSkipTlsVerify,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot create store: %v", err)
	}
	e := store.NewStoreManager(kvstore, storePath)

	return &ClusterChecker{
		uid:              uid,
		listenAddress:    cfg.listenAddress,
		port:             cfg.port,
		stopListening:    cfg.stopListening,
		e:                e,
		endPollonProxyCh: make(chan error),
	}, nil
}

func (c *ClusterChecker) startPollonProxy() error {
	c.pollonMutex.Lock()
	defer c.pollonMutex.Unlock()
	if c.pp != nil {
		return nil
	}

	log.Infow("Starting proxying")
	addr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(cfg.listenAddress, cfg.port))
	if err != nil {
		return fmt.Errorf("error resolving tcp addr %q: %v", addr.String(), err)
	}

	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return fmt.Errorf("error listening on tcp addr %q: %v", addr.String(), err)
	}

	pp, err := pollon.NewProxy(listener)
	if err != nil {
		return fmt.Errorf("error creating pollon proxy: %v", err)
	}
	c.pp = pp
	c.listener = listener

	go func() {
		c.endPollonProxyCh <- c.pp.Start()
	}()

	return nil
}

func (c *ClusterChecker) stopPollonProxy() {
	c.pollonMutex.Lock()
	defer c.pollonMutex.Unlock()
	if c.pp != nil {
		log.Infow("Stopping listening")
		c.pp.Stop()
		c.pp = nil
		c.listener.Close()
		c.listener = nil
	}
}

func (c *ClusterChecker) sendPollonConfData(confData pollon.ConfData) {
	c.pollonMutex.Lock()
	defer c.pollonMutex.Unlock()
	if c.pp != nil {
		c.pp.C <- confData
	}
}

func (c *ClusterChecker) SetProxyInfo(e *store.StoreManager, generation int64, ttl time.Duration) error {
	proxyInfo := &cluster.ProxyInfo{
		UID:        c.uid,
		Generation: generation,
	}
	log.Debugf("proxyInfo dump: %s", spew.Sdump(proxyInfo))

	if err := c.e.SetProxyInfo(proxyInfo, ttl); err != nil {
		return err
	}
	return nil
}

// Check reads the cluster data and applies the right pollon configuration.
func (c *ClusterChecker) Check() error {
	cd, _, err := c.e.GetClusterData()
	if err != nil {
		return fmt.Errorf("cannot get cluster data: %v", err)
	}

	// Start pollon if not active
	if err = c.startPollonProxy(); err != nil {
		return fmt.Errorf("failed to start proxy: %v", err)
	}

	log.Debugf("cd dump: %s", spew.Sdump(cd))
	if cd == nil {
		log.Infow("no clusterdata available, closing connections to master")
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return nil
	}
	if cd.FormatVersion != cluster.CurrentCDFormatVersion {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return fmt.Errorf("unsupported clusterdata format version: %d", cd.FormatVersion)
	}
	if err = cd.Cluster.Spec.Validate(); err != nil {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return fmt.Errorf("clusterdata validation failed: %v", err)
	}

	proxy := cd.Proxy
	if proxy == nil {
		log.Infow("no proxy object available, closing connections to master")
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		// ignore errors on setting proxy info
		if err = c.SetProxyInfo(c.e, cluster.NoGeneration, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
			log.Errorw("failed to update proxyInfo", zap.Error(err))
		}
		return nil
	}

	db, ok := cd.DBs[proxy.Spec.MasterDBUID]
	if !ok {
		log.Infow("no db object available, closing connections to master", "db", proxy.Spec.MasterDBUID)
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		// ignore errors on setting proxy info
		if err = c.SetProxyInfo(c.e, proxy.Generation, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
			log.Errorw("failed to update proxyInfo", zap.Error(err))
		}
		return nil
	}

	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", db.Status.ListenAddress, db.Status.Port))
	if err != nil {
		log.Errorw("cannot resolve db address", zap.Error(err))
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		return nil
	}
	log.Infow("master address", "address", addr)
	if err = c.SetProxyInfo(c.e, proxy.Generation, 2*cluster.DefaultProxyTimeoutInterval); err != nil {
		// if we failed to update our proxy info when a master is defined we
		// cannot ignore this error since the sentinel won't know that we exist
		// and are sending connections to a master so, when electing a new
		// master, it'll not wait for us to close connections to the old one.
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
		log.Errorw("failed to update proxyInfo", zap.Error(err))
		return nil
	}

	// start proxing only if we are inside enabledProxies, this ensures that the
	// sentinel has read our proxyinfo and knows we are alive
	if util.StringInSlice(proxy.Spec.EnabledProxies, c.uid) {
		log.Infow("proxying to master address", "address", addr)
		c.sendPollonConfData(pollon.ConfData{DestAddr: addr})
	} else {
		c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
	}

	return nil
}

func (c *ClusterChecker) TimeoutChecker(checkOkCh chan struct{}) error {
	timeoutTimer := time.NewTimer(cluster.DefaultProxyTimeoutInterval)

	for true {
		select {
		case <-timeoutTimer.C:
			log.Infow("check timeout timer fired")
			// if the check timeouts close all connections and stop listening
			// (for example to avoid load balancers forward connections to us
			// since we aren't ready or in a bad state)
			c.sendPollonConfData(pollon.ConfData{DestAddr: nil})
			if c.stopListening {
				c.stopPollonProxy()
			}

		case <-checkOkCh:
			log.Debugw("check ok message received")

			// ignore if stop succeeded or not due to timer already expired
			timeoutTimer.Stop()
			timeoutTimer = time.NewTimer(cluster.DefaultProxyTimeoutInterval)
		}
	}
	return nil
}

func (c *ClusterChecker) Start() error {
	checkOkCh := make(chan struct{})
	checkCh := make(chan error)
	timerCh := time.NewTimer(0).C

	go c.TimeoutChecker(checkOkCh)

	for true {
		select {
		case <-timerCh:
			go func() {
				checkCh <- c.Check()
			}()
		case err := <-checkCh:
			if err != nil {
				// don't report check ok since it returned an error
				log.Debugw("check function error", zap.Error(err))
			} else {
				// report that check was ok
				checkOkCh <- struct{}{}
			}
			timerCh = time.NewTimer(cluster.DefaultProxyCheckInterval).C
		case err := <-c.endPollonProxyCh:
			if err != nil {
				return fmt.Errorf("proxy error: %v", err)
			}
		}
	}
	return nil
}

func main() {
	flagutil.SetFlagsFromEnv(cmdProxy.PersistentFlags(), "STPROXY")

	cmdProxy.Execute()
}

func proxy(cmd *cobra.Command, args []string) {
	switch cfg.logLevel {
	case "error":
		slog.SetLevel(zap.ErrorLevel)
	case "warn":
		slog.SetLevel(zap.WarnLevel)
	case "info":
		slog.SetLevel(zap.InfoLevel)
	case "debug":
		slog.SetLevel(zap.DebugLevel)
	default:
		die("invalid log level: %v", cfg.logLevel)
	}
	if cfg.debug {
		slog.SetDebug()
	}
	if slog.IsDebug() {
		stdlog := slog.StdLog()
		pollon.SetLogger(stdlog)
	}

	if cfg.clusterName == "" {
		die("cluster name required")
	}
	if cfg.storeBackend == "" {
		die("store backend type required")
	}

	uid := common.UID()
	log.Infow("proxy uid", "uid", uid)

	clusterChecker, err := NewClusterChecker(uid, cfg)
	if err != nil {
		die("cannot create cluster checker: %v", err)
	}
	if err = clusterChecker.Start(); err != nil {
		die("cluster checker ended with error: %v", err)
	}
}
<|endoftext|>


<commit_before>// +build !bare

package e2e

import (
	"context"
	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/net"

	. "github.com/onsi/ginkgo"
	"github.com/operator-framework/api/pkg/operators/v1alpha1"
	"github.com/operator-framework/operator-lifecycle-manager/pkg/lib/operatorclient"
)

var _ = Describe("Metrics", func() {
	It("endpoint", func() {

		// TestMetrics tests the metrics endpoint of the OLM pod.

		c := newKubeClient(GinkgoT())
		crc := newCRClient(GinkgoT())

		failingCSV := v1alpha1.ClusterServiceVersion{
			TypeMeta: metav1.TypeMeta{
				Kind:       v1alpha1.ClusterServiceVersionKind,
				APIVersion: v1alpha1.ClusterServiceVersionAPIVersion,
			},
			ObjectMeta: metav1.ObjectMeta{
				Name: genName("failing-csv-test-"),
			},
			Spec: v1alpha1.ClusterServiceVersionSpec{
				InstallStrategy: v1alpha1.NamedInstallStrategy{
					StrategyName: v1alpha1.InstallStrategyNameDeployment,
					StrategySpec: strategy,
				},
			},
		}

		cleanupCSV, err := createCSV(GinkgoT(), c, crc, failingCSV, testNamespace, false, false)
		require.NoError(GinkgoT(), err)
		defer cleanupCSV()

		_, err = fetchCSV(GinkgoT(), crc, failingCSV.Name, testNamespace, csvFailedChecker)
		require.NoError(GinkgoT(), err)

		rawOutput, err := getMetricsFromPod(GinkgoT(), c, getOLMPodName(GinkgoT(), c), operatorNamespace, "8081")
		if err != nil {
			GinkgoT().Fatalf("Metrics test failed: %v\n", err)
		}

		// Verify metrics have been emitted for packageserver csv
		require.Contains(GinkgoT(), rawOutput, "csv_abnormal")
		require.Contains(GinkgoT(), rawOutput, "name=\""+failingCSV.Name+"\"")
		require.Contains(GinkgoT(), rawOutput, "phase=\"Failed\"")
		require.Contains(GinkgoT(), rawOutput, "reason=\"UnsupportedOperatorGroup\"")
		require.Contains(GinkgoT(), rawOutput, "version=\"0.0.0\"")

		require.Contains(GinkgoT(), rawOutput, "csv_succeeded")
		log.Info(rawOutput)
	})
})

func getOLMPodName(t GinkgoTInterface, client operatorclient.ClientInterface) string {
	listOptions := metav1.ListOptions{LabelSelector: "app=olm-operator"}
	podList, err := client.KubernetesInterface().CoreV1().Pods(operatorNamespace).List(context.TODO(), listOptions)
	if err != nil {
		log.Infof("Error %v\n", err)
		t.Fatalf("Listing pods failed: %v\n", err)
	}
	if len(podList.Items) != 1 {
		t.Fatalf("Expected 1 olm-operator pod, got %v", len(podList.Items))
	}

	podName := podList.Items[0].GetName()
	log.Infof("Looking at pod %v in namespace %v", podName, operatorNamespace)
	return podName

}

func getMetricsFromPod(t GinkgoTInterface, client operatorclient.ClientInterface, podName string, namespace string, port string) (string, error) {
	olmPod, err := client.KubernetesInterface().CoreV1().Pods(namespace).Get(context.TODO(), podName, metav1.GetOptions{})
	if err != nil {
		return "", err
	}
	if len(olmPod.Spec.Containers) != 1 {
		t.Fatalf("Expected only 1 container in olm-operator pod, got %v", len(olmPod.Spec.Containers))
	}

	var foundCert bool
	var foundKey bool
	// assuming -tls-cert and -tls-key aren't used anywhere else as a parameter value
	for _, param := range olmPod.Spec.Containers[0].Args {
		if param == "-tls-cert" {
			foundCert = true
		} else if param == "-tls-key" {
			foundKey = true
		}
	}

	var scheme string
	if foundCert && foundKey {
		scheme = "https"
	} else {
		scheme = "http"
	}
	log.Infof("Retrieving metrics using scheme %v\n", scheme)

	rawOutput, err := client.KubernetesInterface().CoreV1().RESTClient().Get().
		Namespace(namespace).
		Resource("pods").
		SubResource("proxy").
		Name(net.JoinSchemeNamePort(scheme, podName, port)).
		Suffix("metrics").
		Do(context.TODO()).Raw()
	if err != nil {
		return "", err
	}
	return string(rawOutput), nil
}
<commit_msg>test(e2e): correct tls detection in metrics spec<commit_after>// +build !bare

package e2e

import (
	"context"
	"fmt"
	"regexp"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	log "github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/net"

	"github.com/operator-framework/api/pkg/operators/v1alpha1"
	"github.com/operator-framework/operator-lifecycle-manager/pkg/lib/operatorclient"
)

var _ = Describe("Metrics", func() {
	It("endpoint", func() {

		// TestMetrics tests the metrics endpoint of the OLM pod.

		c := newKubeClient(GinkgoT())
		crc := newCRClient(GinkgoT())

		failingCSV := v1alpha1.ClusterServiceVersion{
			TypeMeta: metav1.TypeMeta{
				Kind:       v1alpha1.ClusterServiceVersionKind,
				APIVersion: v1alpha1.ClusterServiceVersionAPIVersion,
			},
			ObjectMeta: metav1.ObjectMeta{
				Name: genName("failing-csv-test-"),
			},
			Spec: v1alpha1.ClusterServiceVersionSpec{
				InstallStrategy: v1alpha1.NamedInstallStrategy{
					StrategyName: v1alpha1.InstallStrategyNameDeployment,
					StrategySpec: strategy,
				},
			},
		}

		cleanupCSV, err := createCSV(GinkgoT(), c, crc, failingCSV, testNamespace, false, false)
		Expect(err).ToNot(HaveOccurred())
		defer cleanupCSV()

		_, err = fetchCSV(GinkgoT(), crc, failingCSV.Name, testNamespace, csvFailedChecker)
		Expect(err).ToNot(HaveOccurred())

		// Verify metrics have been emitted for packageserver csv
		Expect(getMetricsFromPod(c, getOLMPod(c), "8081")).To(And(
			ContainSubstring("csv_abnormal"),
			ContainSubstring(fmt.Sprintf("name=\"%s\"", failingCSV.Name)),
			ContainSubstring("phase=\"Failed\""),
			ContainSubstring("reason=\"UnsupportedOperatorGroup\""),
			ContainSubstring("version=\"0.0.0\""),
			ContainSubstring("csv_succeeded"),
		))
	})
})

func getOLMPod(client operatorclient.ClientInterface) *corev1.Pod {
	listOptions := metav1.ListOptions{LabelSelector: "app=olm-operator"}
	var podList *corev1.PodList
	Eventually(func() (err error) {
		podList, err = client.KubernetesInterface().CoreV1().Pods(operatorNamespace).List(context.TODO(), listOptions)
		return
	}).Should(Succeed(), "Failed to list OLM pods")
	Expect(len(podList.Items)).To(Equal(1))

	return &podList.Items[0]
}

func getMetricsFromPod(client operatorclient.ClientInterface, pod *corev1.Pod, port string) string {
	By(fmt.Sprintf("querying pod %s/%s", pod.GetNamespace(), pod.GetName()))

	// assuming -tls-cert and -tls-key aren't used anywhere else as a parameter value
	var foundCert, foundKey bool
	for _, arg := range pod.Spec.Containers[0].Args {
		matched, err := regexp.MatchString(`^-?-tls-cert`, arg)
		Expect(err).ToNot(HaveOccurred())
		foundCert = foundCert || matched

		matched, err = regexp.MatchString(`^-?-tls-key`, arg)
		Expect(err).ToNot(HaveOccurred())
		foundKey = foundKey || matched
	}

	var scheme string
	if foundCert && foundKey {
		scheme = "https"
	} else {
		scheme = "http"
	}
	log.Infof("Retrieving metrics using scheme %v\n", scheme)

	var raw []byte
	Eventually(func() (err error) {
		raw, err = client.KubernetesInterface().CoreV1().RESTClient().Get().
			Namespace(pod.GetNamespace()).
			Resource("pods").
			SubResource("proxy").
			Name(net.JoinSchemeNamePort(scheme, pod.GetName(), port)).
			Suffix("metrics").
			Do(context.Background()).Raw()
		return
	}).Should(Succeed())

	return string(raw)
}
<|endoftext|>


<commit_before>package cmd

import (
	"testing"

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

func TestSearchCmdHasUse(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Use)
}

func TestSearchCmdHasShort(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Short)
}

func TestSearchCmdHasLong(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Long)
}

func TestSearchCmdHasRun(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Run)
}

<commit_msg>Fix gofmt<commit_after>package cmd

import (
	"testing"

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

func TestSearchCmdHasUse(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Use)
}

func TestSearchCmdHasShort(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Short)
}

func TestSearchCmdHasLong(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Long)
}

func TestSearchCmdHasRun(t *testing.T) {
	assert.NotEmpty(t, SearchCmd.Run)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/albrow/prtty"
	"github.com/albrow/temple"
	"github.com/spf13/cobra"
)

const (
	version = "temple version X.X.X (develop)"
)

func main() {
	// Define build command
	cmdBuild := &cobra.Command{
		Use:   "build <src> <dest>",
		Short: "Compile the templates in the src directory and write generated go code to the dest file.",
		Long:  ``,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 2 {
				prtty.Error.Fatal("temple build requires exactly 2 arguments: the src directory and the dest file.")
			}
			partials := cmd.Flag("partials").Value.String()
			layouts := cmd.Flag("layouts").Value.String()
			if err := temple.Build(args[0], args[1], partials, layouts); err != nil {
				prtty.Error.Fatal(err)
			}
		},
	}
	cmdBuild.Flags().String("partials", "", "(optional) The directory to look for partials. Partials are .tmpl files that are associated with layouts and all other templates.")
	cmdBuild.Flags().String("layouts", "", "(optional) The directory to look for layouts. Layouts are .tmpl files which have access to partials and are associated with all other templates.")

	// Define version command
	cmdVersion := &cobra.Command{
		Use:   "version",
		Short: "Print the current version number.",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println(version)
		},
	}

	// Define the root command
	rootCmd := &cobra.Command{
		Use:   "temple",
		Short: "A command line tool for sharing go templates between a client and server.",
		Long: `
A command line tool for sharing go templates between a client and server.
Visit https://github.com/albrow/temple for source code, example usage, documentation, and more.`,
	}
	rootCmd.AddCommand(cmdBuild, cmdVersion)
	if err := rootCmd.Execute(); err != nil {
		prtty.Error.Fatal(err)
	}
}
<commit_msg>Add verbose flag and set it to false by default<commit_after>package main

import (
	"fmt"
	"github.com/albrow/prtty"
	"github.com/albrow/temple"
	"github.com/spf13/cobra"
	"io/ioutil"
	"os"
)

const (
	version = "temple version X.X.X (develop)"
)

var (
	verbose = false
)

func setQuiet() {
	prtty.AllLoggers.SetOutput(ioutil.Discard)
	prtty.Error.Output = os.Stderr
}

func setVerbose() {
	prtty.AllLoggers.SetOutput(os.Stdout)
	prtty.Error.Output = os.Stderr
}

func main() {
	// Define build command
	cmdBuild := &cobra.Command{
		Use:   "build <src> <dest>",
		Short: "Compile the templates in the src directory and write generated go code to the dest file.",
		Long:  ``,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 2 {
				prtty.Error.Fatal("temple build requires exactly 2 arguments: the src directory and the dest file.")
			}
			if verbose {
				setVerbose()
			} else {
				setQuiet()
			}
			partials := cmd.Flag("partials").Value.String()
			layouts := cmd.Flag("layouts").Value.String()
			if err := temple.Build(args[0], args[1], partials, layouts); err != nil {
				prtty.Error.Fatal(err)
			}
		},
	}
	cmdBuild.Flags().String("partials", "", "(optional) The directory to look for partials. Partials are .tmpl files that are associated with layouts and all other templates.")
	cmdBuild.Flags().String("layouts", "", "(optional) The directory to look for layouts. Layouts are .tmpl files which have access to partials and are associated with all other templates.")
	cmdBuild.Flags().BoolVarP(&verbose, "verbose", "v", false, "If set to true, temple will print out information while building.")

	// Define version command
	cmdVersion := &cobra.Command{
		Use:   "version",
		Short: "Print the current version number.",
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Println(version)
		},
	}

	// Define the root command
	rootCmd := &cobra.Command{
		Use:   "temple",
		Short: "A command line tool for sharing go templates between a client and server.",
		Long: `
A command line tool for sharing go templates between a client and server.
Visit https://github.com/albrow/temple for source code, example usage, documentation, and more.`,
	}
	rootCmd.AddCommand(cmdBuild, cmdVersion)
	if err := rootCmd.Execute(); err != nil {
		prtty.Error.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package test

import (
	"flag"
	"sync"
	"testing"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
)

var ext = flag.Bool("ext", false, "Run external client test case.")

// Helper method for testing client implementations in other languages.
// Flow of events for this function is:
// * Send a {"method":"echo", "params":{"message": "..."}} request to client upon connection,
//   and verify that message body is echoed properly in the response body.
// * Echo any incoming request message body as is within a response message.
// * Repeat ad infinitum, until {"method":"close", "params":"{"message": "..."}"} is received. Close message body is logged.
func TestExternalClient(t *testing.T) {
	sh := NewServerHelper(t).Start()
	defer sh.CloseWait()
	var wg sync.WaitGroup

	m := "Hello!"

	// send 'echo' request to client upon connection (blocks test if no response is received)
	wg.Add(1)
	sh.Server.ConnHandler(func(c *neptulon.Conn) error {
		c.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
			defer wg.Done()
			var msg echoMsg
			if err := ctx.Result(&msg); err != nil {
				t.Fatal(err)
			}
			if msg.Message != m {
				t.Fatalf("server: expected: %v got: %v", m, msg.Message)
			}
			t.Logf("server: client sent 'echo' response message: %v", msg.Message)
			return nil
		})
		return nil
	})

	// handle 'echo' requests via the 'echo middleware'
	srout := middleware.NewRouter()
	sh.Middleware(srout.Middleware)
	srout.Request("echo", middleware.Echo)

	// handle 'close' request (blocks test if no response is received)
	wg.Add(1)
	srout.Request("close", func(ctx *neptulon.ReqCtx) error {
		defer wg.Done()
		if err := ctx.Params(&ctx.Res); err != nil {
			return err
		}
		err := ctx.Next()
		ctx.Conn.Close()
		t.Logf("server: closed connection with message from client: %v\n", ctx.Res)
		return err
	})

	if *ext {
		t.Log("Starter server waiting for external client integration test since.")
		wg.Wait()
		return
	}

	// use internal conn implementation instead to test the test case itself
	t.Log("Skipping external client integration test since -ext flag is not provided.")
	ch := sh.GetConnHelper().Connect()
	defer ch.CloseWait()
	cm := "Thanks for echoing! Over and out."

	// handle 'echo' requests via the 'echo middleware'
	crout := middleware.NewRouter()
	ch.Middleware(crout.Middleware)
	crout.Request("echo", middleware.Echo)

	// handle 'echo' request and send 'close' request upon echo response
	wg.Add(1)
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		defer wg.Done()
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("client: expected: %v got: %v", m, msg.Message)
		}
		t.Log("client: server accepted and echoed 'echo' request message body")

		// send close request after getting our echo message back
		wg.Add(1)
		ch.SendRequest("close", echoMsg{Message: cm}, func(ctx *neptulon.ResCtx) error {
			defer wg.Done()
			var msg echoMsg
			if err := ctx.Result(&msg); err != nil {
				t.Fatal(err)
			}
			if msg.Message != cm {
				t.Fatalf("client: expected: %v got: %v", cm, msg.Message)
			}
			t.Log("client: server accepted and echoed 'close' request message body. bye!")
			return nil
		})

		return nil
	})

	wg.Wait()
}
<commit_msg>add waiter in main thread<commit_after>package test

import (
	"flag"
	"sync"
	"testing"

	"github.com/neptulon/neptulon"
	"github.com/neptulon/neptulon/middleware"
)

var ext = flag.Bool("ext", false, "Run external client test case.")

// Helper method for testing client implementations in other languages.
// Flow of events for this function is:
// * Send a {"method":"echo", "params":{"message": "..."}} request to client upon connection,
//   and verify that message body is echoed properly in the response body.
// * Echo any incoming request message body as is within a response message.
// * Repeat ad infinitum, until {"method":"close", "params":"{"message": "..."}"} is received. Close message body is logged.
func TestExternalClient(t *testing.T) {
	sh := NewServerHelper(t).Start()
	// defer sh.CloseWait()
	var wg sync.WaitGroup

	m := "Hello!"

	// send 'echo' request to client upon connection (blocks test if no response is received)
	wg.Add(1)
	sh.Server.ConnHandler(func(c *neptulon.Conn) error {
		c.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
			defer wg.Done()
			var msg echoMsg
			if err := ctx.Result(&msg); err != nil {
				t.Fatal(err)
			}
			if msg.Message != m {
				t.Fatalf("server: expected: %v got: %v", m, msg.Message)
			}
			t.Logf("server: client sent response to our 'echo' request: %v", msg.Message)
			return nil
		})
		return nil
	})

	// handle 'echo' requests via the 'echo middleware'
	srout := middleware.NewRouter()
	sh.Middleware(srout.Middleware)
	srout.Request("echo", middleware.Echo)

	// handle 'close' request (blocks test if no response is received)
	wg.Add(1)
	srout.Request("close", func(ctx *neptulon.ReqCtx) error {
		defer wg.Done()
		if err := ctx.Params(&ctx.Res); err != nil {
			return err
		}
		err := ctx.Next()
		ctx.Conn.Close()
		t.Logf("server: closed connection with message from client: %v\n", ctx.Res)
		return err
	})

	if *ext {
		t.Log("Starter server waiting for external client integration test since.")
		wg.Wait()
		return
	}

	// use internal conn implementation instead to test the test case itself
	t.Log("Skipping external client integration test since -ext flag is not provided.")
	ch := sh.GetConnHelper().Connect()
	// defer ch.CloseWait()
	cm := "Thanks for echoing! Over and out."

	// handle 'echo' requests via the 'echo middleware'
	crout := middleware.NewRouter()
	ch.Middleware(crout.Middleware)
	crout.Request("echo", middleware.Echo)

	// handle 'echo' request and send 'close' request upon echo response
	wg.Add(2)
	ch.SendRequest("echo", echoMsg{Message: m}, func(ctx *neptulon.ResCtx) error {
		defer wg.Done()
		var msg echoMsg
		if err := ctx.Result(&msg); err != nil {
			t.Fatal(err)
		}
		if msg.Message != m {
			t.Fatalf("client: expected: %v got: %v", m, msg.Message)
		}
		t.Log("client: server accepted and echoed 'echo' request message body")

		// send close request after getting our echo message back
		ch.SendRequest("close", echoMsg{Message: cm}, func(ctx *neptulon.ResCtx) error {
			defer wg.Done()
			var msg echoMsg
			if err := ctx.Result(&msg); err != nil {
				t.Fatal(err)
			}
			if msg.Message != cm {
				t.Fatalf("client: expected: %v got: %v", cm, msg.Message)
			}
			t.Log("client: server accepted and echoed 'close' request message body. bye!")
			return nil
		})

		return nil
	})

	wg.Wait()
}
<|endoftext|>


<commit_before>// +build integration,!no-etcd

package integration

import (
	"net"
	"testing"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"

	"github.com/miekg/dns"
	testutil "github.com/openshift/origin/test/util"
)

func TestDNS(t *testing.T) {
	masterConfig, clientFile, err := testutil.StartTestMaster()
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	localIP := net.ParseIP("127.0.0.1")
	var masterIP net.IP
	// verify service DNS entry is visible
	stop := make(chan struct{})
	util.Until(func() {
		m1 := &dns.Msg{
			MsgHdr:   dns.MsgHdr{Id: dns.Id(), RecursionDesired: true},
			Question: []dns.Question{{"kubernetes.default.svc.cluster.local.", dns.TypeA, dns.ClassINET}},
		}
		in, err := dns.Exchange(m1, masterConfig.DNSConfig.BindAddress)
		if err != nil {
			t.Logf("unexpected error: %v", err)
			return
		}
		if len(in.Answer) != 1 {
			t.Logf("unexpected answer: %#v", in)
			return
		}
		if a, ok := in.Answer[0].(*dns.A); ok {
			if a.A == nil {
				t.Fatalf("expected an A record with an IP: %#v", a)
			}
			masterIP = a.A
		} else {
			t.Fatalf("expected an A record: %#v", in)
		}
		t.Log(in)
		close(stop)
	}, 50*time.Millisecond, stop)

	client, err := testutil.GetClusterAdminKubeClient(clientFile)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	for {
		if _, err := client.Services(kapi.NamespaceDefault).Create(&kapi.Service{
			ObjectMeta: kapi.ObjectMeta{
				Name: "headless",
			},
			Spec: kapi.ServiceSpec{
				PortalIP: kapi.PortalIPNone,
				Ports:    []kapi.ServicePort{{Port: 443}},
			},
		}); err != nil {
			if errors.IsForbidden(err) {
				t.Logf("forbidden, sleeping: %v", err)
				time.Sleep(100 * time.Millisecond)
				continue
			}
			t.Fatalf("unexpected error: %v", err)
		}
		if _, err := client.Endpoints(kapi.NamespaceDefault).Create(&kapi.Endpoints{
			ObjectMeta: kapi.ObjectMeta{
				Name: "headless",
			},
			Subsets: []kapi.EndpointSubset{{
				Addresses: []kapi.EndpointAddress{{IP: "172.0.0.1"}},
				Ports: []kapi.EndpointPort{
					{Port: 2345},
				},
			}},
		}); err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		break
	}
	headlessIP := net.ParseIP("172.0.0.1")

	if _, err := client.Services(kapi.NamespaceDefault).Create(&kapi.Service{
		ObjectMeta: kapi.ObjectMeta{
			Name: "headless2",
		},
		Spec: kapi.ServiceSpec{
			PortalIP: kapi.PortalIPNone,
			Ports:    []kapi.ServicePort{{Port: 443}},
		},
	}); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if _, err := client.Endpoints(kapi.NamespaceDefault).Create(&kapi.Endpoints{
		ObjectMeta: kapi.ObjectMeta{
			Name: "headless2",
		},
		Subsets: []kapi.EndpointSubset{{
			Addresses: []kapi.EndpointAddress{{IP: "172.0.0.2"}},
			Ports: []kapi.EndpointPort{
				{Port: 2345, Name: "other"},
				{Port: 2346, Name: "http"},
			},
		}},
	}); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	headless2IP := net.ParseIP("172.0.0.2")

	// verify recursive DNS lookup is visible when expected
	tests := []struct {
		dnsQuestionName   string
		recursionExpected bool
		retry             bool
		expect            []*net.IP
		srv               []*dns.SRV
	}{
		{ // wildcard resolution of a service works
			dnsQuestionName: "foo.kubernetes.default.svc.cluster.local.",
			expect:          []*net.IP{&masterIP},
		},
		{ // resolving endpoints of a service works
			dnsQuestionName: "kubernetes.default.endpoints.cluster.local.",
			expect:          []*net.IP{&localIP},
		},
		{ // openshift override works
			dnsQuestionName: "openshift.default.svc.cluster.local.",
			expect:          []*net.IP{&masterIP},
		},
		{ // headless service
			dnsQuestionName: "headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // specific port of a headless service
			dnsQuestionName: "unknown-port-2345.e1.headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // SRV record for that service
			dnsQuestionName: "headless.default.svc.cluster.local.",
			srv: []*dns.SRV{
				&dns.SRV{
					Target: "unknown-port-2345.e1.headless.",
					Port:   2345,
				},
			},
		},
		{ // the SRV record resolves to the IP
			dnsQuestionName: "unknown-port-2345.e1.headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // headless 2 service
			dnsQuestionName: "headless2.default.svc.cluster.local.",
			expect:          []*net.IP{&headless2IP},
		},
		{ // SRV records for that service
			dnsQuestionName: "headless2.default.svc.cluster.local.",
			srv: []*dns.SRV{
				&dns.SRV{
					Target: "http.e1.headless2.",
					Port:   2346,
				},
				&dns.SRV{
					Target: "other.e1.headless2.",
					Port:   2345,
				},
			},
		},
		{ // the SRV record resolves to the IP
			dnsQuestionName: "other.e1.headless2.default.svc.cluster.local.",
			expect:          []*net.IP{&headless2IP},
		},
		{
			dnsQuestionName:   "www.google.com.",
			recursionExpected: true,
		},
	}
	for i, tc := range tests {
		qType := dns.TypeA
		if tc.srv != nil {
			qType = dns.TypeSRV
		}
		m1 := &dns.Msg{
			MsgHdr:   dns.MsgHdr{Id: dns.Id(), RecursionDesired: true},
			Question: []dns.Question{{tc.dnsQuestionName, qType, dns.ClassINET}},
		}
		ch := make(chan struct{})
		count := 0
		util.Until(func() {
			count++
			if count > 100 {
				t.Errorf("%d: failed after max iterations", i)
				close(ch)
				return
			}
			in, err := dns.Exchange(m1, masterConfig.DNSConfig.BindAddress)
			if err != nil {
				return
			}
			switch {
			case tc.srv != nil:
				if len(in.Answer) != len(tc.srv) {
					t.Logf("%d: incorrect number of answers: %#v", i, in)
					return
				}
			case tc.recursionExpected:
				if len(in.Answer) == 0 {
					t.Errorf("%d: expected forward resolution: %#v", i, in)
				}
				close(ch)
				return
			default:
				if len(in.Answer) != len(tc.expect) {
					t.Logf("%d: did not resolve or unexpected forward resolution: %#v", i, in)
					return
				}
			}
			for _, answer := range in.Answer {
				switch a := answer.(type) {
				case *dns.A:
					matches := false
					if a.A != nil {
						for _, expect := range tc.expect {
							if a.A.String() == expect.String() {
								matches = true
								break
							}
						}
					}
					if !matches {
						t.Errorf("A record does not match any expected answer: %v", a.A)
					}
				case *dns.SRV:
					matches := false
					for _, expect := range tc.srv {
						if expect.Port == a.Port && expect.Target == a.Target {
							matches = true
							break
						}
					}
					if !matches {
						t.Errorf("SRV record does not match any expected answer: %#v", a)
					}
				default:
					t.Errorf("expected an A or SRV record: %#v", in)
				}
			}
			t.Log(in)
			close(ch)
		}, 50*time.Millisecond, ch)
	}
}
<commit_msg>gofmt fix<commit_after>// +build integration,!no-etcd

package integration

import (
	"net"
	"testing"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"

	"github.com/miekg/dns"
	testutil "github.com/openshift/origin/test/util"
)

func TestDNS(t *testing.T) {
	masterConfig, clientFile, err := testutil.StartTestMaster()
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	localIP := net.ParseIP("127.0.0.1")
	var masterIP net.IP
	// verify service DNS entry is visible
	stop := make(chan struct{})
	util.Until(func() {
		m1 := &dns.Msg{
			MsgHdr:   dns.MsgHdr{Id: dns.Id(), RecursionDesired: true},
			Question: []dns.Question{{"kubernetes.default.svc.cluster.local.", dns.TypeA, dns.ClassINET}},
		}
		in, err := dns.Exchange(m1, masterConfig.DNSConfig.BindAddress)
		if err != nil {
			t.Logf("unexpected error: %v", err)
			return
		}
		if len(in.Answer) != 1 {
			t.Logf("unexpected answer: %#v", in)
			return
		}
		if a, ok := in.Answer[0].(*dns.A); ok {
			if a.A == nil {
				t.Fatalf("expected an A record with an IP: %#v", a)
			}
			masterIP = a.A
		} else {
			t.Fatalf("expected an A record: %#v", in)
		}
		t.Log(in)
		close(stop)
	}, 50*time.Millisecond, stop)

	client, err := testutil.GetClusterAdminKubeClient(clientFile)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	for {
		if _, err := client.Services(kapi.NamespaceDefault).Create(&kapi.Service{
			ObjectMeta: kapi.ObjectMeta{
				Name: "headless",
			},
			Spec: kapi.ServiceSpec{
				PortalIP: kapi.PortalIPNone,
				Ports:    []kapi.ServicePort{{Port: 443}},
			},
		}); err != nil {
			if errors.IsForbidden(err) {
				t.Logf("forbidden, sleeping: %v", err)
				time.Sleep(100 * time.Millisecond)
				continue
			}
			t.Fatalf("unexpected error: %v", err)
		}
		if _, err := client.Endpoints(kapi.NamespaceDefault).Create(&kapi.Endpoints{
			ObjectMeta: kapi.ObjectMeta{
				Name: "headless",
			},
			Subsets: []kapi.EndpointSubset{{
				Addresses: []kapi.EndpointAddress{{IP: "172.0.0.1"}},
				Ports: []kapi.EndpointPort{
					{Port: 2345},
				},
			}},
		}); err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		break
	}
	headlessIP := net.ParseIP("172.0.0.1")

	if _, err := client.Services(kapi.NamespaceDefault).Create(&kapi.Service{
		ObjectMeta: kapi.ObjectMeta{
			Name: "headless2",
		},
		Spec: kapi.ServiceSpec{
			PortalIP: kapi.PortalIPNone,
			Ports:    []kapi.ServicePort{{Port: 443}},
		},
	}); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if _, err := client.Endpoints(kapi.NamespaceDefault).Create(&kapi.Endpoints{
		ObjectMeta: kapi.ObjectMeta{
			Name: "headless2",
		},
		Subsets: []kapi.EndpointSubset{{
			Addresses: []kapi.EndpointAddress{{IP: "172.0.0.2"}},
			Ports: []kapi.EndpointPort{
				{Port: 2345, Name: "other"},
				{Port: 2346, Name: "http"},
			},
		}},
	}); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	headless2IP := net.ParseIP("172.0.0.2")

	// verify recursive DNS lookup is visible when expected
	tests := []struct {
		dnsQuestionName   string
		recursionExpected bool
		retry             bool
		expect            []*net.IP
		srv               []*dns.SRV
	}{
		{ // wildcard resolution of a service works
			dnsQuestionName: "foo.kubernetes.default.svc.cluster.local.",
			expect:          []*net.IP{&masterIP},
		},
		{ // resolving endpoints of a service works
			dnsQuestionName: "kubernetes.default.endpoints.cluster.local.",
			expect:          []*net.IP{&localIP},
		},
		{ // openshift override works
			dnsQuestionName: "openshift.default.svc.cluster.local.",
			expect:          []*net.IP{&masterIP},
		},
		{ // headless service
			dnsQuestionName: "headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // specific port of a headless service
			dnsQuestionName: "unknown-port-2345.e1.headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // SRV record for that service
			dnsQuestionName: "headless.default.svc.cluster.local.",
			srv: []*dns.SRV{
				{
					Target: "unknown-port-2345.e1.headless.",
					Port:   2345,
				},
			},
		},
		{ // the SRV record resolves to the IP
			dnsQuestionName: "unknown-port-2345.e1.headless.default.svc.cluster.local.",
			expect:          []*net.IP{&headlessIP},
		},
		{ // headless 2 service
			dnsQuestionName: "headless2.default.svc.cluster.local.",
			expect:          []*net.IP{&headless2IP},
		},
		{ // SRV records for that service
			dnsQuestionName: "headless2.default.svc.cluster.local.",
			srv: []*dns.SRV{
				{
					Target: "http.e1.headless2.",
					Port:   2346,
				},
				{
					Target: "other.e1.headless2.",
					Port:   2345,
				},
			},
		},
		{ // the SRV record resolves to the IP
			dnsQuestionName: "other.e1.headless2.default.svc.cluster.local.",
			expect:          []*net.IP{&headless2IP},
		},
		{
			dnsQuestionName:   "www.google.com.",
			recursionExpected: true,
		},
	}
	for i, tc := range tests {
		qType := dns.TypeA
		if tc.srv != nil {
			qType = dns.TypeSRV
		}
		m1 := &dns.Msg{
			MsgHdr:   dns.MsgHdr{Id: dns.Id(), RecursionDesired: true},
			Question: []dns.Question{{tc.dnsQuestionName, qType, dns.ClassINET}},
		}
		ch := make(chan struct{})
		count := 0
		util.Until(func() {
			count++
			if count > 100 {
				t.Errorf("%d: failed after max iterations", i)
				close(ch)
				return
			}
			in, err := dns.Exchange(m1, masterConfig.DNSConfig.BindAddress)
			if err != nil {
				return
			}
			switch {
			case tc.srv != nil:
				if len(in.Answer) != len(tc.srv) {
					t.Logf("%d: incorrect number of answers: %#v", i, in)
					return
				}
			case tc.recursionExpected:
				if len(in.Answer) == 0 {
					t.Errorf("%d: expected forward resolution: %#v", i, in)
				}
				close(ch)
				return
			default:
				if len(in.Answer) != len(tc.expect) {
					t.Logf("%d: did not resolve or unexpected forward resolution: %#v", i, in)
					return
				}
			}
			for _, answer := range in.Answer {
				switch a := answer.(type) {
				case *dns.A:
					matches := false
					if a.A != nil {
						for _, expect := range tc.expect {
							if a.A.String() == expect.String() {
								matches = true
								break
							}
						}
					}
					if !matches {
						t.Errorf("A record does not match any expected answer: %v", a.A)
					}
				case *dns.SRV:
					matches := false
					for _, expect := range tc.srv {
						if expect.Port == a.Port && expect.Target == a.Target {
							matches = true
							break
						}
					}
					if !matches {
						t.Errorf("SRV record does not match any expected answer: %#v", a)
					}
				default:
					t.Errorf("expected an A or SRV record: %#v", in)
				}
			}
			t.Log(in)
			close(ch)
		}, 50*time.Millisecond, ch)
	}
}
<|endoftext|>


<commit_before>package azurerm

import (
	"fmt"
	"net/http"
	"testing"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAzureRMNetworkSecurityRule_basic(t *testing.T) {
	rInt := acctest.RandInt()
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_basic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test"),
				),
			},
		},
	})
}

func TestAccAzureRMNetworkSecurityRule_disappears(t *testing.T) {
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_basic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test"),
					testCheckAzureRMNetworkSecurityRuleDisappears("azurerm_network_security_rule.test"),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAzureRMNetworkSecurityRule_addingRules(t *testing.T) {
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_updateBasic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test1"),
				),
			},

			{
				Config: testAccAzureRMNetworkSecurityRule_updateExtraRule(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test2"),
				),
			},
		},
	})
}

func testCheckAzureRMNetworkSecurityRuleExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {

		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup, hasResourceGroup := rs.Primary.Attributes["resource_group_name"]
		if !hasResourceGroup {
			return fmt.Errorf("Bad: no resource group found in state for network security rule: %s", sgName)
		}

		conn := testAccProvider.Meta().(*ArmClient).secRuleClient

		resp, err := conn.Get(resourceGroup, sgName, sgrName)
		if err != nil {
			return fmt.Errorf("Bad: Get on secRuleClient: %s", err)
		}

		if resp.StatusCode == http.StatusNotFound {
			return fmt.Errorf("Bad: Network Security Rule %q (resource group: %q) (network security group: %q) does not exist", sgrName, sgName, resourceGroup)
		}

		return nil
	}
}

func testCheckAzureRMNetworkSecurityRuleDisappears(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {

		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup, hasResourceGroup := rs.Primary.Attributes["resource_group_name"]
		if !hasResourceGroup {
			return fmt.Errorf("Bad: no resource group found in state for network security rule: %s", sgName)
		}

		conn := testAccProvider.Meta().(*ArmClient).secRuleClient

		_, err := conn.Delete(resourceGroup, sgName, sgrName, make(chan struct{}))
		if err != nil {
			return fmt.Errorf("Bad: Delete on secRuleClient: %s", err)
		}

		return nil
	}
}

func testCheckAzureRMNetworkSecurityRuleDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*ArmClient).secRuleClient

	for _, rs := range s.RootModule().Resources {

		if rs.Type != "azurerm_network_security_rule" {
			continue
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup := rs.Primary.Attributes["resource_group_name"]

		resp, err := conn.Get(resourceGroup, sgName, sgrName)

		if err != nil {
			return nil
		}

		if resp.StatusCode != http.StatusNotFound {
			return fmt.Errorf("Network Security Rule still exists:\n%#v", resp.Properties)
		}
	}

	return nil
}

func testAccAzureRMNetworkSecurityRule_basic(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acceptanceTestResourceGroup1"
    location = "West US"
}

resource "azurerm_network_security_group" "test" {
    name = "acceptanceTestSecurityGroup1"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
}

resource "azurerm_network_security_rule" "test" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test.name}"
    	network_security_group_name = "${azurerm_network_security_group.test.name}"
}
`, rInt)
}

func testAccAzureRMNetworkSecurityRule_updateBasic(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test1" {
    name = "acceptanceTestResourceGroup2"
    location = "West US"
}

resource "azurerm_network_security_group" "test1" {
    name = "acceptanceTestSecurityGroup2"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test1.name}"
}

resource "azurerm_network_security_rule" "test1" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}
`, rInt)
}

func testAccAzureRMNetworkSecurityRule_updateExtraRule(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test1" {
    name = "acceptanceTestResourceGroup2"
    location = "West US"
}

resource "azurerm_network_security_group" "test1" {
    name = "acceptanceTestSecurityGroup2"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test1.name}"
}

resource "azurerm_network_security_rule" "test1" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}

resource "azurerm_network_security_rule" "test2" {
	name = "testing456"
    	priority = 101
    	direction = "Inbound"
    	access = "Deny"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}
`, rInt)
}
<commit_msg>provider/azurerm: Randomizing the ARM security rule acceptance tests<commit_after>package azurerm

import (
	"fmt"
	"net/http"
	"testing"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAzureRMNetworkSecurityRule_basic(t *testing.T) {
	rInt := acctest.RandInt()
	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_basic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test"),
				),
			},
		},
	})
}

func TestAccAzureRMNetworkSecurityRule_disappears(t *testing.T) {
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_basic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test"),
					testCheckAzureRMNetworkSecurityRuleDisappears("azurerm_network_security_rule.test"),
				),
				ExpectNonEmptyPlan: true,
			},
		},
	})
}

func TestAccAzureRMNetworkSecurityRule_addingRules(t *testing.T) {
	rInt := acctest.RandInt()

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testCheckAzureRMNetworkSecurityRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAzureRMNetworkSecurityRule_updateBasic(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test1"),
				),
			},

			{
				Config: testAccAzureRMNetworkSecurityRule_updateExtraRule(rInt),
				Check: resource.ComposeTestCheckFunc(
					testCheckAzureRMNetworkSecurityRuleExists("azurerm_network_security_rule.test2"),
				),
			},
		},
	})
}

func testCheckAzureRMNetworkSecurityRuleExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {

		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup, hasResourceGroup := rs.Primary.Attributes["resource_group_name"]
		if !hasResourceGroup {
			return fmt.Errorf("Bad: no resource group found in state for network security rule: %s", sgName)
		}

		conn := testAccProvider.Meta().(*ArmClient).secRuleClient

		resp, err := conn.Get(resourceGroup, sgName, sgrName)
		if err != nil {
			return fmt.Errorf("Bad: Get on secRuleClient: %s", err)
		}

		if resp.StatusCode == http.StatusNotFound {
			return fmt.Errorf("Bad: Network Security Rule %q (resource group: %q) (network security group: %q) does not exist", sgrName, sgName, resourceGroup)
		}

		return nil
	}
}

func testCheckAzureRMNetworkSecurityRuleDisappears(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {

		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup, hasResourceGroup := rs.Primary.Attributes["resource_group_name"]
		if !hasResourceGroup {
			return fmt.Errorf("Bad: no resource group found in state for network security rule: %s", sgName)
		}

		conn := testAccProvider.Meta().(*ArmClient).secRuleClient

		_, err := conn.Delete(resourceGroup, sgName, sgrName, make(chan struct{}))
		if err != nil {
			return fmt.Errorf("Bad: Delete on secRuleClient: %s", err)
		}

		return nil
	}
}

func testCheckAzureRMNetworkSecurityRuleDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*ArmClient).secRuleClient

	for _, rs := range s.RootModule().Resources {

		if rs.Type != "azurerm_network_security_rule" {
			continue
		}

		sgName := rs.Primary.Attributes["network_security_group_name"]
		sgrName := rs.Primary.Attributes["name"]
		resourceGroup := rs.Primary.Attributes["resource_group_name"]

		resp, err := conn.Get(resourceGroup, sgName, sgrName)

		if err != nil {
			return nil
		}

		if resp.StatusCode != http.StatusNotFound {
			return fmt.Errorf("Network Security Rule still exists:\n%#v", resp.Properties)
		}
	}

	return nil
}

func testAccAzureRMNetworkSecurityRule_basic(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test" {
    name = "acctestRG-%d"
    location = "West US"
}

resource "azurerm_network_security_group" "test" {
    name = "acceptanceTestSecurityGroup1"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test.name}"
}

resource "azurerm_network_security_rule" "test" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test.name}"
    	network_security_group_name = "${azurerm_network_security_group.test.name}"
}
`, rInt)
}

func testAccAzureRMNetworkSecurityRule_updateBasic(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test1" {
    name = "acctestRG-%d"
    location = "West US"
}

resource "azurerm_network_security_group" "test1" {
    name = "acceptanceTestSecurityGroup2"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test1.name}"
}

resource "azurerm_network_security_rule" "test1" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}
`, rInt)
}

func testAccAzureRMNetworkSecurityRule_updateExtraRule(rInt int) string {
	return fmt.Sprintf(`
resource "azurerm_resource_group" "test1" {
    name = "acctestRG-%d"
    location = "West US"
}

resource "azurerm_network_security_group" "test1" {
    name = "acceptanceTestSecurityGroup2"
    location = "West US"
    resource_group_name = "${azurerm_resource_group.test1.name}"
}

resource "azurerm_network_security_rule" "test1" {
	name = "test123"
    	priority = 100
    	direction = "Outbound"
    	access = "Allow"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}

resource "azurerm_network_security_rule" "test2" {
	name = "testing456"
    	priority = 101
    	direction = "Inbound"
    	access = "Deny"
    	protocol = "Tcp"
    	source_port_range = "*"
    	destination_port_range = "*"
    	source_address_prefix = "*"
    	destination_address_prefix = "*"
    	resource_group_name = "${azurerm_resource_group.test1.name}"
    	network_security_group_name = "${azurerm_network_security_group.test1.name}"
}
`, rInt)
}
<|endoftext|>