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

package fs_test

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/user"
	"strconv"
	"strings"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/perms"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

const (
	filePerms os.FileMode = 0740
	dirPerms              = 0754
)

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

// A struct that can be embedded to inherit common file system test behaviors.
type fsTest struct {
	ctx context.Context

	// Configuration
	serverCfg fs.ServerConfig
	mountCfg  fuse.MountConfig

	// Dependencies. If bucket is set before SetUp is called, it will be used
	// rather than creating a default one.
	clock  timeutil.SimulatedClock
	bucket gcs.Bucket

	// Mount information
	mfs *fuse.MountedFileSystem
	Dir string

	// Files to close when tearing down. Nil entries are skipped.
	f1 *os.File
	f2 *os.File
}

var _ SetUpInterface = &fsTest{}
var _ TearDownInterface = &fsTest{}

func (t *fsTest) SetUp(ti *TestInfo) {
	var err error
	t.ctx = ti.Ctx

	// Set up the clock.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// And the bucket.
	if t.bucket == nil {
		t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")
	}

	// Set up ownership.
	t.serverCfg.Uid, t.serverCfg.Gid, err = perms.MyUserAndGroup()
	AssertEq(nil, err)

	// Set up permissions.
	t.serverCfg.FilePerms = filePerms
	t.serverCfg.DirPerms = dirPerms

	// Use temporary space to speed tests.
	t.serverCfg.TempDirLimit = 1 << 27 // 128 MiB

	// Tests assume SupportNlink is enabled.
	t.serverCfg.SupportNlink = true

	// Set up a temporary directory for mounting.
	t.Dir, err = ioutil.TempDir("", "fs_test")
	AssertEq(nil, err)

	// Create a file system server.
	server, err := fs.NewServer(&t.serverCfg)
	AssertEq(nil, err)

	// Mount the file system.
	mountCfg := t.MountConfig
	mountCfg.OpContext = t.ctx

	t.mfs, err = fuse.Mount(t.Dir, server, &mountCfg)
	AssertEq(nil, err)
}

func (t *fsTest) TearDown() {
	var err error

	// Close any files we opened.
	if t.f1 != nil {
		ExpectEq(nil, t.f1.Close())
	}

	if t.f2 != nil {
		ExpectEq(nil, t.f2.Close())
	}

	// Unmount the file system. Try again on "resource busy" errors.
	delay := 10 * time.Millisecond
	for {
		err := fuse.Unmount(t.mfs.Dir())
		if err == nil {
			break
		}

		if strings.Contains(err.Error(), "resource busy") {
			log.Println("Resource busy error while unmounting; trying again")
			time.Sleep(delay)
			delay = time.Duration(1.3 * float64(delay))
			continue
		}

		AddFailure("MountedFileSystem.Unmount: %v", err)
		AbortTest()
	}

	if err := t.mfs.Join(t.ctx); err != nil {
		AssertEq(nil, err)
	}

	// Unlink the mount point.
	if err = os.Remove(t.Dir); err != nil {
		err = fmt.Errorf("Unlinking mount point: %v", err)
		return
	}
}

func (t *fsTest) createWithContents(name string, contents string) error {
	return t.createObjects(map[string]string{name: contents})
}

func (t *fsTest) createObjects(in map[string]string) error {
	err := gcsutil.CreateObjects(t.ctx, t.bucket, in)
	return err
}

func (t *fsTest) createEmptyObjects(names []string) error {
	err := gcsutil.CreateEmptyObjects(t.ctx, t.bucket, names)
	return err
}

////////////////////////////////////////////////////////////////////////
// Common helpers
////////////////////////////////////////////////////////////////////////

func getFileNames(entries []os.FileInfo) (names []string) {
	for _, e := range entries {
		names = append(names, e.Name())
	}

	return
}

// REQUIRES: n % 4 == 0
func randString(n int) string {
	bytes := make([]byte, n)
	for i := 0; i < n; i += 4 {
		u32 := rand.Uint32()
		bytes[i] = byte(u32 >> 0)
		bytes[i+1] = byte(u32 >> 8)
		bytes[i+2] = byte(u32 >> 16)
		bytes[i+3] = byte(u32 >> 24)
	}

	return string(bytes)
}

func readRange(r io.ReadSeeker, offset int64, n int) (s string, err error) {
	if _, err = r.Seek(offset, 0); err != nil {
		return
	}

	bytes := make([]byte, n)
	if _, err = io.ReadFull(r, bytes); err != nil {
		return
	}

	s = string(bytes)
	return
}

func currentUid() uint32 {
	user, err := user.Current()
	AssertEq(nil, err)

	uid, err := strconv.ParseUint(user.Uid, 10, 32)
	AssertEq(nil, err)

	return uint32(uid)
}

func currentGid() uint32 {
	user, err := user.Current()
	AssertEq(nil, err)

	gid, err := strconv.ParseUint(user.Gid, 10, 32)
	AssertEq(nil, err)

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

package fs_test

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/user"
	"strconv"
	"strings"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/perms"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	"github.com/jacobsa/gcloud/gcs/gcsutil"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

const (
	filePerms os.FileMode = 0740
	dirPerms              = 0754
)

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

// A struct that can be embedded to inherit common file system test behaviors.
type fsTest struct {
	ctx context.Context

	// Configuration
	serverCfg fs.ServerConfig
	mountCfg  fuse.MountConfig

	// Dependencies. If bucket is set before SetUp is called, it will be used
	// rather than creating a default one.
	clock  timeutil.SimulatedClock
	bucket gcs.Bucket

	// Mount information
	mfs *fuse.MountedFileSystem
	Dir string

	// Files to close when tearing down. Nil entries are skipped.
	f1 *os.File
	f2 *os.File
}

var _ SetUpInterface = &fsTest{}
var _ TearDownInterface = &fsTest{}

func (t *fsTest) SetUp(ti *TestInfo) {
	var err error
	t.ctx = ti.Ctx

	// Set up the clock.
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))

	// And the bucket.
	if t.bucket == nil {
		t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")
	}

	// Set up ownership.
	t.serverCfg.Uid, t.serverCfg.Gid, err = perms.MyUserAndGroup()
	AssertEq(nil, err)

	// Set up permissions.
	t.serverCfg.FilePerms = filePerms
	t.serverCfg.DirPerms = dirPerms

	// Use temporary space to speed tests.
	t.serverCfg.TempDirLimit = 1 << 27 // 128 MiB

	// Tests assume SupportNlink is enabled.
	t.serverCfg.SupportNlink = true

	// Set up a temporary directory for mounting.
	t.Dir, err = ioutil.TempDir("", "fs_test")
	AssertEq(nil, err)

	// Create a file system server.
	server, err := fs.NewServer(&t.serverCfg)
	AssertEq(nil, err)

	// Mount the file system.
	mountCfg := t.mountCfg
	mountCfg.OpContext = t.ctx

	t.mfs, err = fuse.Mount(t.Dir, server, &mountCfg)
	AssertEq(nil, err)
}

func (t *fsTest) TearDown() {
	var err error

	// Close any files we opened.
	if t.f1 != nil {
		ExpectEq(nil, t.f1.Close())
	}

	if t.f2 != nil {
		ExpectEq(nil, t.f2.Close())
	}

	// Unmount the file system. Try again on "resource busy" errors.
	delay := 10 * time.Millisecond
	for {
		err := fuse.Unmount(t.mfs.Dir())
		if err == nil {
			break
		}

		if strings.Contains(err.Error(), "resource busy") {
			log.Println("Resource busy error while unmounting; trying again")
			time.Sleep(delay)
			delay = time.Duration(1.3 * float64(delay))
			continue
		}

		AddFailure("MountedFileSystem.Unmount: %v", err)
		AbortTest()
	}

	if err := t.mfs.Join(t.ctx); err != nil {
		AssertEq(nil, err)
	}

	// Unlink the mount point.
	if err = os.Remove(t.Dir); err != nil {
		err = fmt.Errorf("Unlinking mount point: %v", err)
		return
	}
}

func (t *fsTest) createWithContents(name string, contents string) error {
	return t.createObjects(map[string]string{name: contents})
}

func (t *fsTest) createObjects(in map[string]string) error {
	err := gcsutil.CreateObjects(t.ctx, t.bucket, in)
	return err
}

func (t *fsTest) createEmptyObjects(names []string) error {
	err := gcsutil.CreateEmptyObjects(t.ctx, t.bucket, names)
	return err
}

////////////////////////////////////////////////////////////////////////
// Common helpers
////////////////////////////////////////////////////////////////////////

func getFileNames(entries []os.FileInfo) (names []string) {
	for _, e := range entries {
		names = append(names, e.Name())
	}

	return
}

// REQUIRES: n % 4 == 0
func randString(n int) string {
	bytes := make([]byte, n)
	for i := 0; i < n; i += 4 {
		u32 := rand.Uint32()
		bytes[i] = byte(u32 >> 0)
		bytes[i+1] = byte(u32 >> 8)
		bytes[i+2] = byte(u32 >> 16)
		bytes[i+3] = byte(u32 >> 24)
	}

	return string(bytes)
}

func readRange(r io.ReadSeeker, offset int64, n int) (s string, err error) {
	if _, err = r.Seek(offset, 0); err != nil {
		return
	}

	bytes := make([]byte, n)
	if _, err = io.ReadFull(r, bytes); err != nil {
		return
	}

	s = string(bytes)
	return
}

func currentUid() uint32 {
	user, err := user.Current()
	AssertEq(nil, err)

	uid, err := strconv.ParseUint(user.Uid, 10, 32)
	AssertEq(nil, err)

	return uint32(uid)
}

func currentGid() uint32 {
	user, err := user.Current()
	AssertEq(nil, err)

	gid, err := strconv.ParseUint(user.Gid, 10, 32)
	AssertEq(nil, err)

	return uint32(gid)
}
<|endoftext|>


<commit_before>package stdio

import (
	"io"
	"os"
	"strings"
	"sync"

	"github.com/Microsoft/opengcs/service/gcs/transport"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ConnectionSet is a structure defining the readers and writers the Core
// implementation should forward a process's stdio through.
type ConnectionSet struct {
	In, Out, Err transport.Connection
}

// Close closes each stdio connection.
func (s *ConnectionSet) Close() error {
	var err error
	if s.In != nil {
		if cerr := s.In.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdin")
		}
		s.In = nil
	}
	if s.Out != nil {
		if cerr := s.Out.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdout")
		}
		s.Out = nil
	}
	if s.Err != nil {
		if cerr := s.Err.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stderr")
		}
		s.Err = nil
	}
	return err
}

// FileSet contains os.File fields for stdio.
type FileSet struct {
	In, Out, Err *os.File
}

// Close closes all the FileSet handles.
func (fs *FileSet) Close() error {
	var err error
	if fs.In != nil {
		if cerr := fs.In.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdin")
		}
		fs.In = nil
	}
	if fs.Out != nil {
		if cerr := fs.Out.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdout")
		}
		fs.Out = nil
	}
	if fs.Err != nil {
		if cerr := fs.Err.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stderr")
		}
		fs.Err = nil
	}
	return err
}

// Files returns a FileSet with an os.File for each connection
// in the connection set.
func (s *ConnectionSet) Files() (_ *FileSet, err error) {
	fs := &FileSet{}
	defer func() {
		if err != nil {
			fs.Close()
		}
	}()
	if s.In != nil {
		fs.In, err = s.In.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stdin socket for command")
		}
	}
	if s.Out != nil {
		fs.Out, err = s.Out.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stdout socket for command")
		}
	}
	if s.Err != nil {
		fs.Err, err = s.Err.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stderr socket for command")
		}
	}
	return fs, nil
}

// NewPipeRelay returns a new pipe relay wrapping the given connection stdin,
// stdout, stderr set. If s is nil will assume al stdin, stdout, stderr pipes.
func NewPipeRelay(s *ConnectionSet) (_ *PipeRelay, err error) {
	pr := &PipeRelay{s: s}
	defer func() {
		if err != nil {
			pr.closePipes()
		}
	}()

	if s == nil || s.In != nil {
		pr.pipes[0], pr.pipes[1], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stdin pipe relay")
		}
	}
	if s == nil || s.Out != nil {
		pr.pipes[2], pr.pipes[3], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stdout pipe relay")
		}
	}
	if s == nil || s.Err != nil {
		pr.pipes[4], pr.pipes[5], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stderr pipe relay")
		}
	}
	return pr, nil
}

// PipeRelay is a relay built to expose a pipe interface
// for stdin, stdout, stderr on top of a ConnectionSet.
type PipeRelay struct {
	wg sync.WaitGroup
	s  *ConnectionSet
	// pipes format is stdin [0 read, 1 write], stdout [2 read, 3 write], stderr [4 read, 5 write].
	pipes [6]*os.File
}

// ReplaceConnectionSet allows the caller to add a new destination set after
// creating the relay. This can only be called previous to the call to Start.
func (pr *PipeRelay) ReplaceConnectionSet(s *ConnectionSet) {
	pr.s = s
}

// Files returns a FileSet with an os.File for each connection
// in the connection set.
func (pr *PipeRelay) Files() (*FileSet, error) {
	fs := new(FileSet)
	if pr.s == nil || pr.s.In != nil {
		fs.In = pr.pipes[0]
	}
	if pr.s == nil || pr.s.Out != nil {
		fs.Out = pr.pipes[3]
	}
	if pr.s == nil || pr.s.Err != nil {
		fs.Err = pr.pipes[5]
	}
	return fs, nil
}

func copyAndCleanClose(c transport.Connection, r io.Reader, name string) {
	if n, err := io.Copy(c, r); err != nil {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"bytes":         n,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error copying from pipe")
	}
	// Shut down the write end of the socket, then read a byte (which should
	// yield EOF) to wait for the other endpoint to finish reading and close
	// the connection.
	if err := c.CloseWrite(); err == nil {
		var b [1]byte
		_, err = c.Read(b[:])
		if err == nil {
			err = errors.New("unexpected data in socket")
		}
		if err != io.EOF {
			logrus.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"file":          name,
			}).Error("opengcs::PipeRelay::copyAndCleanClose - error reading for clean close")
		}
	} else {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error shutting down socket")
	}
	if err := c.Close(); err != nil {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error closing socket")
	}
}

// Start starts the relay operation. The caller must call Wait to wait
// for the relay to finish and release the associated resources.
func (pr *PipeRelay) Start() {
	if pr.s.In != nil {
		pr.wg.Add(1)
		go func() {
			if n, err := io.Copy(pr.pipes[1], pr.s.In); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
					"bytes":         n,
				}).Error("opengcs::PipeRelay::Start - error copying stdin to pipe")
			}
			if err := pr.pipes[1].Close(); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::PipeRelay::Start - error closing stdin write pipe")
			}
			pr.pipes[1] = nil
			pr.wg.Done()
		}()
	}
	if pr.s.Out != nil {
		pr.wg.Add(1)
		go func() {
			copyAndCleanClose(pr.s.Out, pr.pipes[2], "stdout")
			pr.wg.Done()
		}()
	}
	if pr.s.Err != nil {
		pr.wg.Add(1)
		go func() {
			copyAndCleanClose(pr.s.Err, pr.pipes[4], "stderr")
			pr.wg.Done()
		}()
	}
}

// Wait waits for the relaying to finish and closes the associated
// pipes and connections.
func (pr *PipeRelay) Wait() {
	// Close stdin so that the copying goroutine is safely unblocked; this is necessary
	// because the host expects stdin to be closed before it will report process
	// exit back to the client, and the client expects the process notification before
	// it will close its side of stdin (which io.Copy is waiting on in the copying goroutine).
	if pr.s.In != nil {
		pr.s.In.CloseRead()
	}

	pr.wg.Wait()
	pr.closePipes()
	pr.s.Close()
}

// CloseUnusedPipes gives the caller the ability to close any pipes that do not
// have a cooresponding entry on the ConnectionSet. This is to be used in
// conjunction with NewPipeRelay where s is nil which wil open all pipes and
// later calling ReplaceConnectionSet with the actual connections.
func (pr *PipeRelay) CloseUnusedPipes() {
	if pr.s == nil {
		pr.closePipes()
	} else {
		if pr.s.In == nil {
			// Write end of stdin
			pr.pipes[1].Close()
		}
		if pr.s.Out == nil {
			// Read end of stdout
			pr.pipes[2].Close()
		}
		if pr.s.Err == nil {
			// Read end of stderr
			pr.pipes[4].Close()
		}
	}
}

func (pr *PipeRelay) closePipes() {
	for i := 0; i < len(pr.pipes); i++ {
		if pr.pipes[i] != nil {
			if err := pr.pipes[i].Close(); err != nil {
				if !strings.Contains(err.Error(), "file already closed") {
					logrus.WithFields(logrus.Fields{
						logrus.ErrorKey: err,
					}).Error("opengcs::PipeRelay::closePipes - error closing relay pipe")
				}
			}
			pr.pipes[i] = nil
		}
	}
}

// NewTtyRelay returns a new TTY relay for a given master PTY file.
func NewTtyRelay(s *ConnectionSet, pty *os.File) *TtyRelay {
	return &TtyRelay{s: s, pty: pty}
}

// TtyRelay relays IO between a set of stdio connections and a master PTY file.
type TtyRelay struct {
	m      sync.Mutex
	closed bool
	wg     sync.WaitGroup
	s      *ConnectionSet
	pty    *os.File
}

// ReplaceConnectionSet allows the caller to add a new destination set after
// creating the relay. This can only be called previous to the call to Start.
func (r *TtyRelay) ReplaceConnectionSet(s *ConnectionSet) {
	r.s = s
}

// ResizeConsole sends the appropriate resize to a pTTY FD
func (r *TtyRelay) ResizeConsole(height, width uint16) error {
	r.m.Lock()
	defer r.m.Unlock()

	if r.closed {
		return nil
	}
	return ResizeConsole(r.pty, height, width)
}

// Start starts the relay operation. The caller must call Wait to wait
// for the relay to finish and release the associated resources.
func (r *TtyRelay) Start() {
	if r.s.In != nil {
		r.wg.Add(1)
		go func() {
			if _, err := io.Copy(r.pty, r.s.In); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::TtyRelay::Start - error copying stdin to pty")
			}
			r.wg.Done()
		}()
	}
	if r.s.Out != nil {
		r.wg.Add(1)
		go func() {
			if _, err := io.Copy(r.s.Out, r.pty); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::TtyRelay::Start - error copying pty to stdout")
			}
			r.wg.Done()
		}()
	}
}

// Wait waits for the relaying to finish and closes the associated
// files and connections.
func (r *TtyRelay) Wait() {
	// Close stdin so that the copying goroutine is safely unblocked; this is necessary
	// because the host expects stdin to be closed before it will report process
	// exit back to the client, and the client expects the process notification before
	// it will close its side of stdin (which io.Copy is waiting on in the copying goroutine).
	if r.s.In != nil {
		r.s.In.CloseRead()
	}

	// Wait for all users of stdioSet and master to finish before closing them.
	r.wg.Wait()

	r.m.Lock()
	defer r.m.Unlock()

	r.pty.Close()
	r.closed = true
	r.s.Close()
}
<commit_msg>Fix nil dereference race in Pipe/TtyRelay Wait<commit_after>package stdio

import (
	"io"
	"os"
	"strings"
	"sync"

	"github.com/Microsoft/opengcs/service/gcs/transport"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ConnectionSet is a structure defining the readers and writers the Core
// implementation should forward a process's stdio through.
type ConnectionSet struct {
	In, Out, Err transport.Connection
}

// Close closes each stdio connection.
func (s *ConnectionSet) Close() error {
	var err error
	if s.In != nil {
		if cerr := s.In.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdin")
		}
		s.In = nil
	}
	if s.Out != nil {
		if cerr := s.Out.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdout")
		}
		s.Out = nil
	}
	if s.Err != nil {
		if cerr := s.Err.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stderr")
		}
		s.Err = nil
	}
	return err
}

// FileSet contains os.File fields for stdio.
type FileSet struct {
	In, Out, Err *os.File
}

// Close closes all the FileSet handles.
func (fs *FileSet) Close() error {
	var err error
	if fs.In != nil {
		if cerr := fs.In.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdin")
		}
		fs.In = nil
	}
	if fs.Out != nil {
		if cerr := fs.Out.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stdout")
		}
		fs.Out = nil
	}
	if fs.Err != nil {
		if cerr := fs.Err.Close(); cerr != nil && err == nil {
			err = errors.Wrap(cerr, "failed Close on stderr")
		}
		fs.Err = nil
	}
	return err
}

// Files returns a FileSet with an os.File for each connection
// in the connection set.
func (s *ConnectionSet) Files() (_ *FileSet, err error) {
	fs := &FileSet{}
	defer func() {
		if err != nil {
			fs.Close()
		}
	}()
	if s.In != nil {
		fs.In, err = s.In.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stdin socket for command")
		}
	}
	if s.Out != nil {
		fs.Out, err = s.Out.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stdout socket for command")
		}
	}
	if s.Err != nil {
		fs.Err, err = s.Err.File()
		if err != nil {
			return nil, errors.Wrap(err, "failed to dup stderr socket for command")
		}
	}
	return fs, nil
}

// NewPipeRelay returns a new pipe relay wrapping the given connection stdin,
// stdout, stderr set. If s is nil will assume al stdin, stdout, stderr pipes.
func NewPipeRelay(s *ConnectionSet) (_ *PipeRelay, err error) {
	pr := &PipeRelay{s: s}
	defer func() {
		if err != nil {
			pr.closePipes()
		}
	}()

	if s == nil || s.In != nil {
		pr.pipes[0], pr.pipes[1], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stdin pipe relay")
		}
	}
	if s == nil || s.Out != nil {
		pr.pipes[2], pr.pipes[3], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stdout pipe relay")
		}
	}
	if s == nil || s.Err != nil {
		pr.pipes[4], pr.pipes[5], err = os.Pipe()
		if err != nil {
			return nil, errors.Wrap(err, "failed to create stderr pipe relay")
		}
	}
	return pr, nil
}

// PipeRelay is a relay built to expose a pipe interface
// for stdin, stdout, stderr on top of a ConnectionSet.
type PipeRelay struct {
	wg sync.WaitGroup
	s  *ConnectionSet
	// pipes format is stdin [0 read, 1 write], stdout [2 read, 3 write], stderr [4 read, 5 write].
	pipes [6]*os.File
}

// ReplaceConnectionSet allows the caller to add a new destination set after
// creating the relay. This can only be called previous to the call to Start.
func (pr *PipeRelay) ReplaceConnectionSet(s *ConnectionSet) {
	pr.s = s
}

// Files returns a FileSet with an os.File for each connection
// in the connection set.
func (pr *PipeRelay) Files() (*FileSet, error) {
	fs := new(FileSet)
	if pr.s == nil || pr.s.In != nil {
		fs.In = pr.pipes[0]
	}
	if pr.s == nil || pr.s.Out != nil {
		fs.Out = pr.pipes[3]
	}
	if pr.s == nil || pr.s.Err != nil {
		fs.Err = pr.pipes[5]
	}
	return fs, nil
}

func copyAndCleanClose(c transport.Connection, r io.Reader, name string) {
	if n, err := io.Copy(c, r); err != nil {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"bytes":         n,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error copying from pipe")
	}
	// Shut down the write end of the socket, then read a byte (which should
	// yield EOF) to wait for the other endpoint to finish reading and close
	// the connection.
	if err := c.CloseWrite(); err == nil {
		var b [1]byte
		_, err = c.Read(b[:])
		if err == nil {
			err = errors.New("unexpected data in socket")
		}
		if err != io.EOF {
			logrus.WithFields(logrus.Fields{
				logrus.ErrorKey: err,
				"file":          name,
			}).Error("opengcs::PipeRelay::copyAndCleanClose - error reading for clean close")
		}
	} else {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error shutting down socket")
	}
	if err := c.Close(); err != nil {
		logrus.WithFields(logrus.Fields{
			logrus.ErrorKey: err,
			"file":          name,
		}).Error("opengcs::PipeRelay::copyAndCleanClose - error closing socket")
	}
}

// Start starts the relay operation. The caller must call Wait to wait
// for the relay to finish and release the associated resources.
func (pr *PipeRelay) Start() {
	if pr.s.In != nil {
		pr.wg.Add(1)
		go func() {
			if n, err := io.Copy(pr.pipes[1], pr.s.In); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
					"bytes":         n,
				}).Error("opengcs::PipeRelay::Start - error copying stdin to pipe")
			}
			if err := pr.pipes[1].Close(); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::PipeRelay::Start - error closing stdin write pipe")
			}
			pr.pipes[1] = nil
			pr.wg.Done()
		}()
	}
	if pr.s.Out != nil {
		pr.wg.Add(1)
		go func() {
			copyAndCleanClose(pr.s.Out, pr.pipes[2], "stdout")
			pr.wg.Done()
		}()
	}
	if pr.s.Err != nil {
		pr.wg.Add(1)
		go func() {
			copyAndCleanClose(pr.s.Err, pr.pipes[4], "stderr")
			pr.wg.Done()
		}()
	}
}

// Wait waits for the relaying to finish and closes the associated
// pipes and connections.
func (pr *PipeRelay) Wait() {
	// Close stdin so that the copying goroutine is safely unblocked; this is necessary
	// because the host expects stdin to be closed before it will report process
	// exit back to the client, and the client expects the process notification before
	// it will close its side of stdin (which io.Copy is waiting on in the copying goroutine).
	if pr.s != nil && pr.s.In != nil {
		pr.s.In.CloseRead()
	}

	pr.wg.Wait()
	pr.closePipes()
	if pr.s != nil {
		pr.s.Close()
	}
}

// CloseUnusedPipes gives the caller the ability to close any pipes that do not
// have a cooresponding entry on the ConnectionSet. This is to be used in
// conjunction with NewPipeRelay where s is nil which wil open all pipes and
// later calling ReplaceConnectionSet with the actual connections.
func (pr *PipeRelay) CloseUnusedPipes() {
	if pr.s == nil {
		pr.closePipes()
	} else {
		if pr.s.In == nil {
			// Write end of stdin
			pr.pipes[1].Close()
		}
		if pr.s.Out == nil {
			// Read end of stdout
			pr.pipes[2].Close()
		}
		if pr.s.Err == nil {
			// Read end of stderr
			pr.pipes[4].Close()
		}
	}
}

func (pr *PipeRelay) closePipes() {
	for i := 0; i < len(pr.pipes); i++ {
		if pr.pipes[i] != nil {
			if err := pr.pipes[i].Close(); err != nil {
				if !strings.Contains(err.Error(), "file already closed") {
					logrus.WithFields(logrus.Fields{
						logrus.ErrorKey: err,
					}).Error("opengcs::PipeRelay::closePipes - error closing relay pipe")
				}
			}
			pr.pipes[i] = nil
		}
	}
}

// NewTtyRelay returns a new TTY relay for a given master PTY file.
func NewTtyRelay(s *ConnectionSet, pty *os.File) *TtyRelay {
	return &TtyRelay{s: s, pty: pty}
}

// TtyRelay relays IO between a set of stdio connections and a master PTY file.
type TtyRelay struct {
	m      sync.Mutex
	closed bool
	wg     sync.WaitGroup
	s      *ConnectionSet
	pty    *os.File
}

// ReplaceConnectionSet allows the caller to add a new destination set after
// creating the relay. This can only be called previous to the call to Start.
func (r *TtyRelay) ReplaceConnectionSet(s *ConnectionSet) {
	r.s = s
}

// ResizeConsole sends the appropriate resize to a pTTY FD
func (r *TtyRelay) ResizeConsole(height, width uint16) error {
	r.m.Lock()
	defer r.m.Unlock()

	if r.closed {
		return nil
	}
	return ResizeConsole(r.pty, height, width)
}

// Start starts the relay operation. The caller must call Wait to wait
// for the relay to finish and release the associated resources.
func (r *TtyRelay) Start() {
	if r.s.In != nil {
		r.wg.Add(1)
		go func() {
			if _, err := io.Copy(r.pty, r.s.In); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::TtyRelay::Start - error copying stdin to pty")
			}
			r.wg.Done()
		}()
	}
	if r.s.Out != nil {
		r.wg.Add(1)
		go func() {
			if _, err := io.Copy(r.s.Out, r.pty); err != nil {
				logrus.WithFields(logrus.Fields{
					logrus.ErrorKey: err,
				}).Error("opengcs::TtyRelay::Start - error copying pty to stdout")
			}
			r.wg.Done()
		}()
	}
}

// Wait waits for the relaying to finish and closes the associated
// files and connections.
func (r *TtyRelay) Wait() {
	// Close stdin so that the copying goroutine is safely unblocked; this is necessary
	// because the host expects stdin to be closed before it will report process
	// exit back to the client, and the client expects the process notification before
	// it will close its side of stdin (which io.Copy is waiting on in the copying goroutine).
	if r.s != nil && r.s.In != nil {
		r.s.In.CloseRead()
	}

	// Wait for all users of stdioSet and master to finish before closing them.
	r.wg.Wait()

	r.m.Lock()
	defer r.m.Unlock()

	r.pty.Close()
	r.closed = true
	if r.s != nil {
		r.s.Close()
	}
}
<|endoftext|>


<commit_before>package sakuracloud

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/yamamoto-febc/libsacloud/api"
	"github.com/yamamoto-febc/libsacloud/sacloud"
)

func resourceSakuraCloudPacketFilter() *schema.Resource {
	return &schema.Resource{
		Create: resourceSakuraCloudPacketFilterCreate,
		Read:   resourceSakuraCloudPacketFilterRead,
		Update: resourceSakuraCloudPacketFilterUpdate,
		Delete: resourceSakuraCloudPacketFilterDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"expressions": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"protocol": &schema.Schema{
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validateStringInWord(sacloud.AllowPacketFilterProtocol()),
						},

						"source_nw": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},

						"source_port": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
						"dest_port": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
						"allow": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
						},
						"description": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
					},
				},
			},
			"zone": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				Description:  "target SakuraCloud zone",
				ValidateFunc: validateStringInWord([]string{"is1a", "is1b", "tk1a", "tk1v"}),
			},
		},
	}
}

func resourceSakuraCloudPacketFilterCreate(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	opts := client.PacketFilter.New()

	opts.Name = d.Get("name").(string)
	if description, ok := d.GetOk("description"); ok {
		opts.Description = description.(string)
	}

	if rawExpressions, ok := d.GetOk("expressions"); ok {
		expressions := rawExpressions.([]interface{})
		for _, e := range expressions {
			exp := e.(map[string]interface{})
			protocol := exp["protocol"].(string)
			sourceNW := exp["source_nw"].(string)
			sourcePort := exp["source_port"].(string)
			destPort := exp["dest_port"].(string)
			allow := exp["allow"].(bool)
			desc := exp["description"].(string)

			var err error
			switch protocol {
			case "tcp":
				err = opts.AddTCPRule(sourceNW, sourcePort, destPort, desc, allow)
			case "udp":
				err = opts.AddUDPRule(sourceNW, sourcePort, destPort, desc, allow)
			case "icmp":
				err = opts.AddICMPRule(sourceNW, desc, allow)
			case "fragment":
				err = opts.AddFragmentRule(sourceNW, desc, allow)
			case "ip":
				err = opts.AddIPRule(sourceNW, desc, allow)
			}

			if err != nil {
				return err
			}

		}
	}

	filter, err := client.PacketFilter.Create(opts)
	if err != nil {
		return fmt.Errorf("Failed to create SakuraCloud PacketFilter resource: %s", err)
	}

	d.SetId(filter.ID)
	return resourceSakuraCloudPacketFilterRead(d, meta)
}

func resourceSakuraCloudPacketFilterRead(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	filter, err := client.PacketFilter.Read(d.Id())
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud PacketFilter resource: %s", err)
	}

	d.Set("name", filter.Name)
	d.Set("description", filter.Description)

	if filter.Expression != nil && len(filter.Expression) > 0 {
		expressions := []interface{}{}
		for _, exp := range filter.Expression {
			expression := map[string]interface{}{}
			protocol := exp.Protocol
			switch protocol {
			case "tcp", "udp":
				expression["source_nw"] = exp.SourceNetwork
				expression["source_port"] = exp.SourcePort
				expression["dest_port"] = exp.DestinationPort
			case "icmp", "fragment", "ip":
				expression["source_nw"] = exp.SourceNetwork
			}

			expression["protocol"] = exp.Protocol
			expression["allow"] = (exp.Action == "allow")
			expression["description"] = exp.Description

			expressions = append(expressions, expression)
		}
		d.Set("expressions", expressions)
	} else {
		d.Set("expressions", []interface{}{})
	}

	d.Set("zone", client.Zone)

	return nil
}

func resourceSakuraCloudPacketFilterUpdate(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	filter, err := client.PacketFilter.Read(d.Id())
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud PacketFilter resource: %s", err)
	}

	if d.HasChange("name") {
		filter.Name = d.Get("name").(string)
	}
	if d.HasChange("description") {
		if description, ok := d.GetOk("description"); ok {
			filter.Description = description.(string)
		} else {
			filter.Description = ""
		}
	}

	if d.HasChange("expressions") {
		filter.ClearRules()
		if rawExpressions, ok := d.GetOk("expressions"); ok {
			expressions := rawExpressions.([]interface{})
			for _, e := range expressions {
				exp := e.(map[string]interface{})
				protocol := exp["protocol"].(string)
				sourceNW := exp["source_nw"].(string)
				sourcePort := exp["source_port"].(string)
				destPort := exp["dest_port"].(string)
				allow := exp["allow"].(bool)
				desc := exp["description"].(string)

				var err error
				switch protocol {
				case "tcp":
					err = filter.AddTCPRule(sourceNW, sourcePort, destPort, desc, allow)
				case "udp":
					err = filter.AddUDPRule(sourceNW, sourcePort, destPort, desc, allow)
				case "icmp":
					err = filter.AddICMPRule(sourceNW, desc, allow)
				case "fragment":
					err = filter.AddFragmentRule(sourceNW, desc, allow)
				case "ip":
					err = filter.AddIPRule(sourceNW, desc, allow)
				}

				if err != nil {
					return err
				}

			}
		}

	}

	filter, err = client.PacketFilter.Update(filter.ID, filter)
	if err != nil {
		return fmt.Errorf("Error updating SakuraCloud PacketFilter resource: %s", err)
	}

	d.SetId(filter.ID)
	return resourceSakuraCloudPacketFilterRead(d, meta)
}

func resourceSakuraCloudPacketFilterDelete(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	servers, err := client.Server.Find()
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud Server resource: %s", err)
	}
	for _, server := range servers.Servers {
		for _, i := range server.Interfaces {
			if i.PacketFilter != nil && i.PacketFilter.ID == d.Id() {
				_, err := client.Interface.DisconnectFromPacketFilter(i.ID)
				if err != nil {
					return fmt.Errorf("Error disconnecting SakuraCloud PacketFilter : %s", err)
				}
			}
		}
	}

	_, err = client.PacketFilter.Delete(d.Id())
	if err != nil {
		return fmt.Errorf("Error deleting SakuraCloud PacketFilter resource: %s", err)
	}
	return nil
}
<commit_msg>Add default value to PacketFilter(allow)<commit_after>package sakuracloud

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/yamamoto-febc/libsacloud/api"
	"github.com/yamamoto-febc/libsacloud/sacloud"
)

func resourceSakuraCloudPacketFilter() *schema.Resource {
	return &schema.Resource{
		Create: resourceSakuraCloudPacketFilterCreate,
		Read:   resourceSakuraCloudPacketFilterRead,
		Update: resourceSakuraCloudPacketFilterUpdate,
		Delete: resourceSakuraCloudPacketFilterDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"expressions": &schema.Schema{
				Type:     schema.TypeList,
				Required: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"protocol": &schema.Schema{
							Type:         schema.TypeString,
							Required:     true,
							ForceNew:     true,
							ValidateFunc: validateStringInWord(sacloud.AllowPacketFilterProtocol()),
						},

						"source_nw": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},

						"source_port": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
						"dest_port": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
						"allow": &schema.Schema{
							Type:     schema.TypeBool,
							Optional: true,
							Default:  true,
						},
						"description": &schema.Schema{
							Type:     schema.TypeString,
							Optional: true,
							Default:  "",
						},
					},
				},
			},
			"zone": &schema.Schema{
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				Description:  "target SakuraCloud zone",
				ValidateFunc: validateStringInWord([]string{"is1a", "is1b", "tk1a", "tk1v"}),
			},
		},
	}
}

func resourceSakuraCloudPacketFilterCreate(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	opts := client.PacketFilter.New()

	opts.Name = d.Get("name").(string)
	if description, ok := d.GetOk("description"); ok {
		opts.Description = description.(string)
	}

	if rawExpressions, ok := d.GetOk("expressions"); ok {
		expressions := rawExpressions.([]interface{})
		for _, e := range expressions {
			exp := e.(map[string]interface{})
			protocol := exp["protocol"].(string)
			sourceNW := exp["source_nw"].(string)
			sourcePort := exp["source_port"].(string)
			destPort := exp["dest_port"].(string)
			allow := exp["allow"].(bool)
			desc := exp["description"].(string)

			var err error
			switch protocol {
			case "tcp":
				err = opts.AddTCPRule(sourceNW, sourcePort, destPort, desc, allow)
			case "udp":
				err = opts.AddUDPRule(sourceNW, sourcePort, destPort, desc, allow)
			case "icmp":
				err = opts.AddICMPRule(sourceNW, desc, allow)
			case "fragment":
				err = opts.AddFragmentRule(sourceNW, desc, allow)
			case "ip":
				err = opts.AddIPRule(sourceNW, desc, allow)
			}

			if err != nil {
				return err
			}

		}
	}

	filter, err := client.PacketFilter.Create(opts)
	if err != nil {
		return fmt.Errorf("Failed to create SakuraCloud PacketFilter resource: %s", err)
	}

	d.SetId(filter.ID)
	return resourceSakuraCloudPacketFilterRead(d, meta)
}

func resourceSakuraCloudPacketFilterRead(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	filter, err := client.PacketFilter.Read(d.Id())
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud PacketFilter resource: %s", err)
	}

	d.Set("name", filter.Name)
	d.Set("description", filter.Description)

	if filter.Expression != nil && len(filter.Expression) > 0 {
		expressions := []interface{}{}
		for _, exp := range filter.Expression {
			expression := map[string]interface{}{}
			protocol := exp.Protocol
			switch protocol {
			case "tcp", "udp":
				expression["source_nw"] = exp.SourceNetwork
				expression["source_port"] = exp.SourcePort
				expression["dest_port"] = exp.DestinationPort
			case "icmp", "fragment", "ip":
				expression["source_nw"] = exp.SourceNetwork
			}

			expression["protocol"] = exp.Protocol
			expression["allow"] = (exp.Action == "allow")
			expression["description"] = exp.Description

			expressions = append(expressions, expression)
		}
		d.Set("expressions", expressions)
	} else {
		d.Set("expressions", []interface{}{})
	}

	d.Set("zone", client.Zone)

	return nil
}

func resourceSakuraCloudPacketFilterUpdate(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	filter, err := client.PacketFilter.Read(d.Id())
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud PacketFilter resource: %s", err)
	}

	if d.HasChange("name") {
		filter.Name = d.Get("name").(string)
	}
	if d.HasChange("description") {
		if description, ok := d.GetOk("description"); ok {
			filter.Description = description.(string)
		} else {
			filter.Description = ""
		}
	}

	if d.HasChange("expressions") {
		filter.ClearRules()
		if rawExpressions, ok := d.GetOk("expressions"); ok {
			expressions := rawExpressions.([]interface{})
			for _, e := range expressions {
				exp := e.(map[string]interface{})
				protocol := exp["protocol"].(string)
				sourceNW := exp["source_nw"].(string)
				sourcePort := exp["source_port"].(string)
				destPort := exp["dest_port"].(string)
				allow := exp["allow"].(bool)
				desc := exp["description"].(string)

				var err error
				switch protocol {
				case "tcp":
					err = filter.AddTCPRule(sourceNW, sourcePort, destPort, desc, allow)
				case "udp":
					err = filter.AddUDPRule(sourceNW, sourcePort, destPort, desc, allow)
				case "icmp":
					err = filter.AddICMPRule(sourceNW, desc, allow)
				case "fragment":
					err = filter.AddFragmentRule(sourceNW, desc, allow)
				case "ip":
					err = filter.AddIPRule(sourceNW, desc, allow)
				}

				if err != nil {
					return err
				}

			}
		}

	}

	filter, err = client.PacketFilter.Update(filter.ID, filter)
	if err != nil {
		return fmt.Errorf("Error updating SakuraCloud PacketFilter resource: %s", err)
	}

	d.SetId(filter.ID)
	return resourceSakuraCloudPacketFilterRead(d, meta)
}

func resourceSakuraCloudPacketFilterDelete(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*api.Client)
	client := c.Clone()
	zone, ok := d.GetOk("zone")
	if ok {
		client.Zone = zone.(string)
	}

	servers, err := client.Server.Find()
	if err != nil {
		return fmt.Errorf("Couldn't find SakuraCloud Server resource: %s", err)
	}
	for _, server := range servers.Servers {
		for _, i := range server.Interfaces {
			if i.PacketFilter != nil && i.PacketFilter.ID == d.Id() {
				_, err := client.Interface.DisconnectFromPacketFilter(i.ID)
				if err != nil {
					return fmt.Errorf("Error disconnecting SakuraCloud PacketFilter : %s", err)
				}
			}
		}
	}

	_, err = client.PacketFilter.Delete(d.Id())
	if err != nil {
		return fmt.Errorf("Error deleting SakuraCloud PacketFilter resource: %s", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"

	log "github.com/Sirupsen/logrus"
	"github.com/elazarl/goproxy"
	"github.com/garyburd/redigo/redis"
)

type DBClient struct {
	cache Cache
	http  *http.Client
}

// request holds structure for request
type request struct {
	details requestDetails
}

var emptyResp = &http.Response{}

// requestDetails stores information about request, it's used for creating unique hash and also as a payload structure
type requestDetails struct {
	Path        string              `json:"path"`
	Method      string              `json:"method"`
	Destination string              `json:"destination"`
	Query       string              `json:"query"`
	Body        string              `json:"body"`
	RemoteAddr  string              `json:"remoteAddr"`
	Headers     map[string][]string `json:"headers"`
}

// hash returns unique hash key for request
func (r *request) hash() string {
	h := md5.New()
	io.WriteString(h, fmt.Sprintf("%s%s%s%s", r.details.Destination, r.details.Path, r.details.Method, r.details.Query))
	return fmt.Sprintf("%x", h.Sum(nil))
}

// res structure hold response body from external service, body is not decoded and is supposed
// to be bytes, however headers should provide all required information for later decoding
// by the client.
type response struct {
	Status  int                 `json:"status"`
	Body    string              `json:"body"`
	Headers map[string][]string `json:"headers"`
}

// Payload structure holds request and response structure
type Payload struct {
	Response response       `json:"response"`
	Request  requestDetails `json:"request"`
	ID       string         `json:"id"`
}

// recordRequest saves request for later playback
func (d *DBClient) captureRequest(req *http.Request) (*http.Response, error) {

	// forwarding request
	resp, err := d.doRequest(req)

	if err == nil {

		// getting response body
		body, err := httputil.DumpResponse(resp, true)
		if err != nil {
			// copying the response body did not work
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to copy response body.")
			}
		}

		// saving response body with request/response meta to cache
		go d.save(req, resp, body)
	}

	// return new response or error here
	return resp, err
}

// doRequest performs original request and returns response that should be returned to client and error (if there is one)
func (d *DBClient) doRequest(request *http.Request) (*http.Response, error) {
	// We can't have this set. And it only contains "/pkg/net/http/" anyway
	request.RequestURI = ""

	resp, err := d.http.Do(request)

	if err != nil {
		log.WithFields(log.Fields{
			"error":  err.Error(),
			"host":   request.Host,
			"method": request.Method,
			"path":   request.URL.Path,
		}).Error("Could not forward request.")
		return nil, err
	}

	log.WithFields(log.Fields{
		"host":   request.Host,
		"method": request.Method,
		"path":   request.URL.Path,
	}).Info("Response got successfuly!")

	resp.Header.Set("Gen-proxy", "Was-Here")
	return resp, nil

}

// save gets request fingerprint, extracts request body, status code and headers, then saves it to cache
func (d *DBClient) save(req *http.Request, resp *http.Response, respBody []byte) {
	// record request here
	key := getRequestFingerprint(req)

	if resp == nil {
		resp = emptyResp
	} else {
		responseObj := response{
			Status:  resp.StatusCode,
			Body:    string(respBody),
			Headers: resp.Header,
		}

		log.WithFields(log.Fields{
			"path":          req.URL.Path,
			"rawQuery":      req.URL.RawQuery,
			"requestMethod": req.Method,
			"destination":   req.Host,
			"hashKey":       key,
		}).Info("Recording")

		requestObj := requestDetails{
			Path:        req.URL.Path,
			Method:      req.Method,
			Destination: req.Host,
			Query:       req.URL.RawQuery,
		}

		payload := Payload{
			Response: responseObj,
			Request:  requestObj,
			ID:       key,
		}
		// converting it to json bytes
		bts, err := json.Marshal(payload)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
			}).Error("Failed to marshal json")
		} else {
			d.cache.set(key, bts)
		}
	}

}

// getAllRecordsRaw returns raw (json string) for all records
func (d *DBClient) getAllRecordsRaw() ([]string, error) {
	keys, err := d.cache.getAllKeys()

	if err == nil {

		// checking if there are any records
		if len(keys) == 0 {
			return nil, nil
		}

		jsonStrs, err := d.cache.getAllValues(keys)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
			}).Error("Failed to get all values (raw)")
			return nil, err
		} else {
			return jsonStrs, nil
		}

	} else {
		return nil, err
	}
}

// getAllRecords returns all stored
func (d *DBClient) getAllRecords() ([]Payload, error) {
	var payloads []Payload

	jsonStrs, err := d.getAllRecordsRaw()

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Failed to get all values")
	} else {

		if jsonStrs != nil {
			for _, v := range jsonStrs {
				var pl Payload
				err = json.Unmarshal([]byte(v), &pl)
				if err != nil {
					log.WithFields(log.Fields{
						"error": err.Error(),
						"json":  v,
					}).Warning("Failed to deserialize json")
				} else {
					payloads = append(payloads, pl)
				}
			}
		}
	}

	return payloads, err

}

// deleteAllRecords deletes all recorded requests
func (d *DBClient) deleteAllRecords() error {
	keys, err := d.cache.getAllKeys()

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Warning("Failed to get keys, cannot delete all records")
		return err
	} else {
		for _, v := range keys {
			d.cache.delete(v)
		}
		return nil
	}
}

// getRequestFingerprint returns request hash
func getRequestFingerprint(req *http.Request) string {
	details := requestDetails{Path: req.URL.Path, Method: req.Method, Destination: req.Host, Query: req.URL.RawQuery}
	r := request{details: details}
	return r.hash()
}

// getResponse returns stored response from cache
func (d *DBClient) getResponse(req *http.Request) *http.Response {

	key := getRequestFingerprint(req)
	var payload Payload

	payloadBts, err := redis.Bytes(d.cache.get(key))

	if err == nil {
		// getting cache response
		err = json.Unmarshal(payloadBts, &payload)
		if err != nil {
			log.Error(err)
			// what now?
		}

		c := NewConstructor(req, payload)

		if AppConfig.middleware != "" {
			_ = c.ApplyMiddleware(AppConfig.middleware)
		}

		response := c.reconstructResponse()

		log.WithFields(log.Fields{
			"key":        key,
			"status":     payload.Response.Status,
			"bodyLength": response.ContentLength,
			"mode":       AppConfig.mode,
		}).Info("Response found, returning")

		return response

	} else {
		log.WithFields(log.Fields{
			"error": err.Error(),
			"mode":  AppConfig.mode,
		}).Error("Failed to retrieve response from cache")
		// return error? if we return nil - proxy forwards request to original destination
		return goproxy.NewResponse(req,
			goproxy.ContentTypeText, http.StatusPreconditionFailed,
			"Coudldn't find recorded request, please record it first!")
	}

}
<commit_msg>including more data for each request<commit_after>package main

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httputil"

	log "github.com/Sirupsen/logrus"
	"github.com/elazarl/goproxy"
	"github.com/garyburd/redigo/redis"
)

type DBClient struct {
	cache Cache
	http  *http.Client
}

// request holds structure for request
type request struct {
	details requestDetails
}

var emptyResp = &http.Response{}

// requestDetails stores information about request, it's used for creating unique hash and also as a payload structure
type requestDetails struct {
	Path        string              `json:"path"`
	Method      string              `json:"method"`
	Destination string              `json:"destination"`
	Query       string              `json:"query"`
	Body        string              `json:"body"`
	RemoteAddr  string              `json:"remoteAddr"`
	Headers     map[string][]string `json:"headers"`
}

// hash returns unique hash key for request
func (r *request) hash() string {
	h := md5.New()
	io.WriteString(h, fmt.Sprintf("%s%s%s%s", r.details.Destination, r.details.Path, r.details.Method, r.details.Query))
	return fmt.Sprintf("%x", h.Sum(nil))
}

// res structure hold response body from external service, body is not decoded and is supposed
// to be bytes, however headers should provide all required information for later decoding
// by the client.
type response struct {
	Status  int                 `json:"status"`
	Body    string              `json:"body"`
	Headers map[string][]string `json:"headers"`
}

// Payload structure holds request and response structure
type Payload struct {
	Response response       `json:"response"`
	Request  requestDetails `json:"request"`
	ID       string         `json:"id"`
}

// recordRequest saves request for later playback
func (d *DBClient) captureRequest(req *http.Request) (*http.Response, error) {

	// forwarding request
	resp, err := d.doRequest(req)

	if err == nil {

		// getting response body
		body, err := httputil.DumpResponse(resp, true)
		if err != nil {
			// copying the response body did not work
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to copy response body.")
			}
		}

		// saving response body with request/response meta to cache
		go d.save(req, resp, body)
	}

	// return new response or error here
	return resp, err
}

// doRequest performs original request and returns response that should be returned to client and error (if there is one)
func (d *DBClient) doRequest(request *http.Request) (*http.Response, error) {
	// We can't have this set. And it only contains "/pkg/net/http/" anyway
	request.RequestURI = ""

	resp, err := d.http.Do(request)

	if err != nil {
		log.WithFields(log.Fields{
			"error":  err.Error(),
			"host":   request.Host,
			"method": request.Method,
			"path":   request.URL.Path,
		}).Error("Could not forward request.")
		return nil, err
	}

	log.WithFields(log.Fields{
		"host":   request.Host,
		"method": request.Method,
		"path":   request.URL.Path,
	}).Info("Response got successfuly!")

	resp.Header.Set("Gen-proxy", "Was-Here")
	return resp, nil

}

// save gets request fingerprint, extracts request body, status code and headers, then saves it to cache
func (d *DBClient) save(req *http.Request, resp *http.Response, respBody []byte) {
	// record request here
	key := getRequestFingerprint(req)

	if resp == nil {
		resp = emptyResp
	} else {
		responseObj := response{
			Status:  resp.StatusCode,
			Body:    string(respBody),
			Headers: resp.Header,
		}

		log.WithFields(log.Fields{
			"path":          req.URL.Path,
			"rawQuery":      req.URL.RawQuery,
			"requestMethod": req.Method,
			"destination":   req.Host,
			"hashKey":       key,
		}).Info("Recording")

		b := bufio.NewScanner(req.Body)

		bodyStr := b.Text()

		requestObj := requestDetails{
			Path:        req.URL.Path,
			Method:      req.Method,
			Destination: req.Host,
			Query:       req.URL.RawQuery,
			Body:        bodyStr,
			RemoteAddr:  req.RemoteAddr,
			Headers:     req.Header,
		}

		payload := Payload{
			Response: responseObj,
			Request:  requestObj,
			ID:       key,
		}
		// converting it to json bytes
		bts, err := json.Marshal(payload)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
			}).Error("Failed to marshal json")
		} else {
			d.cache.set(key, bts)
		}
	}

}

// getAllRecordsRaw returns raw (json string) for all records
func (d *DBClient) getAllRecordsRaw() ([]string, error) {
	keys, err := d.cache.getAllKeys()

	if err == nil {

		// checking if there are any records
		if len(keys) == 0 {
			return nil, nil
		}

		jsonStrs, err := d.cache.getAllValues(keys)

		if err != nil {
			log.WithFields(log.Fields{
				"error": err.Error(),
			}).Error("Failed to get all values (raw)")
			return nil, err
		} else {
			return jsonStrs, nil
		}

	} else {
		return nil, err
	}
}

// getAllRecords returns all stored
func (d *DBClient) getAllRecords() ([]Payload, error) {
	var payloads []Payload

	jsonStrs, err := d.getAllRecordsRaw()

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Failed to get all values")
	} else {

		if jsonStrs != nil {
			for _, v := range jsonStrs {
				var pl Payload
				err = json.Unmarshal([]byte(v), &pl)
				if err != nil {
					log.WithFields(log.Fields{
						"error": err.Error(),
						"json":  v,
					}).Warning("Failed to deserialize json")
				} else {
					payloads = append(payloads, pl)
				}
			}
		}
	}

	return payloads, err

}

// deleteAllRecords deletes all recorded requests
func (d *DBClient) deleteAllRecords() error {
	keys, err := d.cache.getAllKeys()

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Warning("Failed to get keys, cannot delete all records")
		return err
	} else {
		for _, v := range keys {
			d.cache.delete(v)
		}
		return nil
	}
}

// getRequestFingerprint returns request hash
func getRequestFingerprint(req *http.Request) string {
	details := requestDetails{Path: req.URL.Path, Method: req.Method, Destination: req.Host, Query: req.URL.RawQuery}
	r := request{details: details}
	return r.hash()
}

// getResponse returns stored response from cache
func (d *DBClient) getResponse(req *http.Request) *http.Response {

	key := getRequestFingerprint(req)
	var payload Payload

	payloadBts, err := redis.Bytes(d.cache.get(key))

	if err == nil {
		// getting cache response
		err = json.Unmarshal(payloadBts, &payload)
		if err != nil {
			log.Error(err)
			// what now?
		}

		c := NewConstructor(req, payload)

		if AppConfig.middleware != "" {
			_ = c.ApplyMiddleware(AppConfig.middleware)
		}

		response := c.reconstructResponse()

		log.WithFields(log.Fields{
			"key":        key,
			"status":     payload.Response.Status,
			"bodyLength": response.ContentLength,
			"mode":       AppConfig.mode,
		}).Info("Response found, returning")

		return response

	} else {
		log.WithFields(log.Fields{
			"error": err.Error(),
			"mode":  AppConfig.mode,
		}).Error("Failed to retrieve response from cache")
		// return error? if we return nil - proxy forwards request to original destination
		return goproxy.NewResponse(req,
			goproxy.ContentTypeText, http.StatusPreconditionFailed,
			"Coudldn't find recorded request, please record it first!")
	}

}
<|endoftext|>


<commit_before>package watch

import (
	"bytes"
	"fmt"
	"io"
	"time"

	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/watch"
)

// Every watcher expects infinite byte stream
// E.g. http.Body can provide continuous byte stream,
// each chunk sent asynchronously (invocation of Read is blocked until new chunk)

// Implementation of io.ReadCloser
type WatchBuffer struct {
	buf   *bytes.Buffer
	read  chan []byte
	write chan []byte
	retc  chan retc

	closed bool
}

type retc struct {
	n int
	e error
}

var _ io.ReadCloser = &WatchBuffer{}

// Read watch events as byte stream
func (w *WatchBuffer) Read(p []byte) (n int, err error) {
	if w.closed {
		return 0, io.EOF
	}
	w.read <- p
	ret := <-w.retc
	return ret.n, ret.e
}

// Close all channels
func (w *WatchBuffer) Close() error {
	if !w.closed {
		close(w.read)
		close(w.write)
		close(w.retc)
		w.closed = true
	}
	return nil
}

// Write
func (w *WatchBuffer) Write(data []byte) (nr int, err error) {
	w.write <- data
	return len(data), nil
}

func (c *WatchBuffer) EmitWatchEvent(eType watch.EventType, object runtime.Object) error {
	//event := watch.Event{
	//	Type: eType,
	//	Object: object,
	//}
	var buffer bytes.Buffer
	buffer.WriteString("{\"type\":\"")
	buffer.WriteString(string(eType))
	buffer.WriteString("\",\"object\":")

	payload := []byte(buffer.String())
	payload = append(payload, ([]byte)(runtime.EncodeOrDie(testapi.Default.Codec(), object))...)
	payload = append(payload, []byte("}")...)

	_, err := c.Write(payload)
	return err
}

func (w *WatchBuffer) loop() {
	var dataIn, dataOut []byte
	var ok bool
	for {
		select {
		case dataIn = <-w.write:
			_, err := w.buf.Write(dataIn)
			if err != nil {
				// TODO(jchaloup): add log message
				fmt.Println("Write error")
				break
			}
		case dataOut = <-w.read:
			if w.buf.Len() == 0 {
				dataIn, ok = <-w.write
				if !ok {
					break
				}
				_, err := w.buf.Write(dataIn)
				if err != nil {
					// TODO(jchaloup): add log message
					fmt.Println("Write error")
					break
				}
			}
			nr, err := w.buf.Read(dataOut)
			if w.closed {
				break
			}
			w.retc <- retc{nr, err}
		}
	}
}

func NewWatchBuffer() *WatchBuffer {
	wb := &WatchBuffer{
		buf:    bytes.NewBuffer(nil),
		read:   make(chan []byte),
		write:  make(chan []byte),
		retc:   make(chan retc),
		closed: false,
	}

	go wb.loop()
	return wb
}

func main() {
	buffer := NewWatchBuffer()

	go func() {
		buffer.Write([]byte("Ahoj"))
		time.Sleep(5 * time.Second)
		buffer.Write([]byte(" Svete"))
		//time.Sleep(time.Second)
	}()

	data := []byte{0, 0, 0, 0, 0, 0}
	buffer.Read(data)
	fmt.Printf("\tdata: %s\n", data)
	buffer.Read(data)
	fmt.Printf("\tdata: %s\n", data)

	time.Sleep(10 * time.Second)

	buffer.Close()
	fmt.Println("Ahoj")
}
<commit_msg>create named watched buffer, fix some panics<commit_after>package watch

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

	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/watch"
)

// Every watcher expects infinite byte stream
// E.g. http.Body can provide continuous byte stream,
// each chunk sent asynchronously (invocation of Read is blocked until new chunk)

// Implementation of io.ReadCloser
type WatchBuffer struct {
	buf      *bytes.Buffer
	read     chan []byte
	write    chan []byte
	retc     chan retc
	Resource string

	closed   bool
	closeMux sync.RWMutex
}

type retc struct {
	n int
	e error
}

var _ io.ReadCloser = &WatchBuffer{}

// Read watch events as byte stream
func (w *WatchBuffer) Read(p []byte) (n int, err error) {
	if w.closed {
		return 0, io.EOF
	}
	w.read <- p
	ret := <-w.retc
	return ret.n, ret.e
}

// Close all channels
func (w *WatchBuffer) Close() error {
	w.closeMux.Lock()
	defer w.closeMux.Unlock()

	if !w.closed {
		w.closed = true
		close(w.read)
		close(w.write)
		close(w.retc)
	}
	return nil
}

// Write
func (w *WatchBuffer) Write(data []byte) (nr int, err error) {
	if w.closed {
		return 0, io.EOF
	}
	w.write <- data
	return len(data), nil
}

func (c *WatchBuffer) EmitWatchEvent(eType watch.EventType, object runtime.Object) error {
	//event := watch.Event{
	//	Type: eType,
	//	Object: object,
	//}

	obj_str := runtime.EncodeOrDie(testapi.Default.Codec(), object)
	obj_str = strings.Replace(obj_str, "\n", "", -1)

	var buffer bytes.Buffer
	buffer.WriteString("{\"type\":\"")
	buffer.WriteString(string(eType))
	buffer.WriteString("\",\"object\":")
	buffer.WriteString(obj_str)
	buffer.WriteString("}")

	_, err := c.Write(buffer.Bytes())
	return err
}

func (w *WatchBuffer) loop() {
	var dataIn, dataOut []byte
	var ok bool
	for {
		select {
		case dataIn = <-w.write:
			// channel closed
			if len(dataIn) == 0 {
				if w.closed {
					return
				}
			}
			_, err := w.buf.Write(dataIn)
			if err != nil {
				// TODO(jchaloup): add log message
				fmt.Println("Write error")
				break
			}
		case dataOut = <-w.read:
			if w.buf.Len() == 0 {
				dataIn, ok = <-w.write
				if !ok {
					break
				}
				_, err := w.buf.Write(dataIn)
				if err != nil {
					// TODO(jchaloup): add log message
					fmt.Println("Write error")
					break
				}
			}
			nr, err := w.buf.Read(dataOut)
			if w.closed {
				break
			}
			w.retc <- retc{nr, err}
		}
	}
}

func NewWatchBuffer(resource string) *WatchBuffer {
	wb := &WatchBuffer{
		buf:      bytes.NewBuffer(nil),
		read:     make(chan []byte),
		write:    make(chan []byte),
		retc:     make(chan retc),
		Resource: resource,
		closed:   false,
	}

	go wb.loop()
	return wb
}

func main() {
	buffer := NewWatchBuffer("pods")

	go func() {
		buffer.Write([]byte("Ahoj"))
		time.Sleep(5 * time.Second)
		buffer.Write([]byte(" Svete"))
		//time.Sleep(time.Second)
	}()

	data := []byte{0, 0, 0, 0, 0, 0}
	buffer.Read(data)
	fmt.Printf("\tdata: %s\n", data)
	buffer.Read(data)
	fmt.Printf("\tdata: %s\n", data)

	time.Sleep(10 * time.Second)

	buffer.Close()
	fmt.Println("Ahoj")
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 cainjector

import (
	"context"
	"fmt"
	"io/ioutil"

	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"

	admissionreg "k8s.io/api/admissionregistration/v1beta1"
	apiext "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/runtime"
	apireg "k8s.io/kube-aggregator/pkg/apis/apiregistration/v1beta1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

// injectorSet describes a particular setup of the injector controller
type injectorSetup struct {
	resourceName string
	injector     CertInjector
	listType     runtime.Object
}

var (
	MutatingWebhookSetup = injectorSetup{
		resourceName: "mutatingwebhookconfiguration",
		injector:     mutatingWebhookInjector{},
		listType:     &admissionreg.MutatingWebhookConfigurationList{},
	}

	ValidatingWebhookSetup = injectorSetup{
		resourceName: "validatingwebhookconfiguration",
		injector:     validatingWebhookInjector{},
		listType:     &admissionreg.ValidatingWebhookConfigurationList{},
	}

	APIServiceSetup = injectorSetup{
		resourceName: "apiservice",
		injector:     apiServiceInjector{},
		listType:     &apireg.APIServiceList{},
	}

	CRDSetup = injectorSetup{
		resourceName: "customresourcedefinition",
		injector:     crdConversionInjector{},
		listType:     &apiext.CustomResourceDefinitionList{},
	}

	injectorSetups  = []injectorSetup{MutatingWebhookSetup, ValidatingWebhookSetup, APIServiceSetup, CRDSetup}
	ControllerNames []string
)

// registerAllInjectors registers all injectors and based on the
// graduation state of the injector decides how to log no kind/resource match errors
func registerAllInjectors(ctx context.Context, groupName string, mgr ctrl.Manager, sources []caDataSource, client client.Client, ca cache.Cache) error {
	controllers := make([]controller.Controller, len(injectorSetups))
	for i, setup := range injectorSetups {
		controller, err := newGenericInjectionController(groupName, mgr, setup, sources, ca, client)
		if err != nil {
			if !meta.IsNoMatchError(err) || !setup.injector.IsAlpha() {
				return err
			}
			ctrl.Log.V(logf.WarnLevel).Info("unable to register injector which is still in an alpha phase."+
				" Enable the feature on the API server in order to use this injector",
				"injector", setup.resourceName)
		}
		controllers[i] = controller
	}
	g, gctx := errgroup.WithContext(ctx)

	g.Go(func() (err error) {
		if err = ca.Start(gctx.Done()); err != nil {
			return errors.WithStack(err)
		}
		return nil
	})
	if ca.WaitForCacheSync(gctx.Done()) {
		for _, controller := range controllers {
			if gctx.Err() != nil {
				continue
			}
			controller := controller
			g.Go(func() (err error) {
				return controller.Start(gctx.Done())
			})
		}
	}
	return g.Wait()
}

// newGenericInjectionController creates a controller and adds relevant watches
// and indexers to the supplied cache.
func newGenericInjectionController(groupName string, mgr ctrl.Manager, setup injectorSetup, sources []caDataSource, ca cache.Cache, client client.Client) (controller.Controller, error) {
	log := ctrl.Log.WithName(groupName).WithName(setup.resourceName)
	typ := setup.injector.NewTarget().AsObject()

	c, err := controller.NewUnmanaged(
		fmt.Sprintf("controller-for-%s-%s", groupName, setup.resourceName),
		mgr,
		controller.Options{
			Reconciler: &genericInjectReconciler{
				Client:       client,
				sources:      sources,
				log:          log.WithName("generic-inject-reconciler"),
				resourceName: setup.resourceName,
				injector:     setup.injector,
			},
			Log:                     log,
			MaxConcurrentReconciles: 5,
		})
	if err != nil {
		return nil, errors.WithStack(err)
	}
	if err := c.Watch(source.NewKindWithCache(typ, ca), &handler.EnqueueRequestForObject{}); err != nil {
		return nil, errors.WithStack(err)
	}

	for _, s := range sources {
		if err := s.ApplyTo(mgr, setup, c, ca); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// dataFromSliceOrFile returns data from the slice (if non-empty), or from the file,
// or an error if an error occurred reading the file
func dataFromSliceOrFile(data []byte, file string) ([]byte, error) {
	if len(data) > 0 {
		return data, nil
	}
	if len(file) > 0 {
		fileData, err := ioutil.ReadFile(file)
		if err != nil {
			return []byte{}, err
		}
		return fileData, nil
	}
	return nil, nil
}

// RegisterCertificateBased registers all known injection controllers that
// target Certificate resources with the  given manager, and adds relevant
// indices.
// The registered controllers require the cert-manager API to be available
// in order to run.
func RegisterCertificateBased(ctx context.Context, mgr ctrl.Manager) error {
	cache, client, err := newIndependentCacheAndDelegatingClient(mgr)
	if err != nil {
		return err
	}
	return registerAllInjectors(
		ctx,
		"certificate",
		mgr,
		[]caDataSource{
			&certificateDataSource{client: cache},
		},
		client,
		cache,
	)
}

// RegisterSecretBased registers all known injection controllers that
// target Secret resources with the  given manager, and adds relevant
// indices.
// The registered controllers only require the corev1 APi to be available in
// order to run.
func RegisterSecretBased(ctx context.Context, mgr ctrl.Manager) error {
	cache, client, err := newIndependentCacheAndDelegatingClient(mgr)
	if err != nil {
		return err
	}
	return registerAllInjectors(
		ctx,
		"secret",
		mgr,
		[]caDataSource{
			&secretDataSource{client: cache},
			&kubeconfigDataSource{},
		},
		client,
		cache,
	)
}

func newIndependentCacheAndDelegatingClient(mgr ctrl.Manager) (cache.Cache, client.Client, error) {
	cacheOptions := cache.Options{
		Scheme: mgr.GetScheme(),
		Mapper: mgr.GetRESTMapper(),
	}
	ca, err := cache.New(mgr.GetConfig(), cacheOptions)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}

	clientOptions := client.Options{
		Scheme: mgr.GetScheme(),
		Mapper: mgr.GetRESTMapper(),
	}
	client, err := manager.DefaultNewClient(ca, mgr.GetConfig(), clientOptions)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	return ca, client, nil
}
<commit_msg>Revert the increased number of reconcile threads<commit_after>/*
Copyright 2019 The Jetstack cert-manager contributors.

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 cainjector

import (
	"context"
	"fmt"
	"io/ioutil"

	logf "github.com/jetstack/cert-manager/pkg/logs"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"

	admissionreg "k8s.io/api/admissionregistration/v1beta1"
	apiext "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/runtime"
	apireg "k8s.io/kube-aggregator/pkg/apis/apiregistration/v1beta1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

// injectorSet describes a particular setup of the injector controller
type injectorSetup struct {
	resourceName string
	injector     CertInjector
	listType     runtime.Object
}

var (
	MutatingWebhookSetup = injectorSetup{
		resourceName: "mutatingwebhookconfiguration",
		injector:     mutatingWebhookInjector{},
		listType:     &admissionreg.MutatingWebhookConfigurationList{},
	}

	ValidatingWebhookSetup = injectorSetup{
		resourceName: "validatingwebhookconfiguration",
		injector:     validatingWebhookInjector{},
		listType:     &admissionreg.ValidatingWebhookConfigurationList{},
	}

	APIServiceSetup = injectorSetup{
		resourceName: "apiservice",
		injector:     apiServiceInjector{},
		listType:     &apireg.APIServiceList{},
	}

	CRDSetup = injectorSetup{
		resourceName: "customresourcedefinition",
		injector:     crdConversionInjector{},
		listType:     &apiext.CustomResourceDefinitionList{},
	}

	injectorSetups  = []injectorSetup{MutatingWebhookSetup, ValidatingWebhookSetup, APIServiceSetup, CRDSetup}
	ControllerNames []string
)

// registerAllInjectors registers all injectors and based on the
// graduation state of the injector decides how to log no kind/resource match errors
func registerAllInjectors(ctx context.Context, groupName string, mgr ctrl.Manager, sources []caDataSource, client client.Client, ca cache.Cache) error {
	controllers := make([]controller.Controller, len(injectorSetups))
	for i, setup := range injectorSetups {
		controller, err := newGenericInjectionController(groupName, mgr, setup, sources, ca, client)
		if err != nil {
			if !meta.IsNoMatchError(err) || !setup.injector.IsAlpha() {
				return err
			}
			ctrl.Log.V(logf.WarnLevel).Info("unable to register injector which is still in an alpha phase."+
				" Enable the feature on the API server in order to use this injector",
				"injector", setup.resourceName)
		}
		controllers[i] = controller
	}
	g, gctx := errgroup.WithContext(ctx)

	g.Go(func() (err error) {
		if err = ca.Start(gctx.Done()); err != nil {
			return errors.WithStack(err)
		}
		return nil
	})
	if ca.WaitForCacheSync(gctx.Done()) {
		for _, controller := range controllers {
			if gctx.Err() != nil {
				continue
			}
			controller := controller
			g.Go(func() (err error) {
				return controller.Start(gctx.Done())
			})
		}
	}
	return g.Wait()
}

// newGenericInjectionController creates a controller and adds relevant watches
// and indexers to the supplied cache.
func newGenericInjectionController(groupName string, mgr ctrl.Manager, setup injectorSetup, sources []caDataSource, ca cache.Cache, client client.Client) (controller.Controller, error) {
	log := ctrl.Log.WithName(groupName).WithName(setup.resourceName)
	typ := setup.injector.NewTarget().AsObject()

	c, err := controller.NewUnmanaged(
		fmt.Sprintf("controller-for-%s-%s", groupName, setup.resourceName),
		mgr,
		controller.Options{
			Reconciler: &genericInjectReconciler{
				Client:       client,
				sources:      sources,
				log:          log.WithName("generic-inject-reconciler"),
				resourceName: setup.resourceName,
				injector:     setup.injector,
			},
			Log: log,
		})
	if err != nil {
		return nil, errors.WithStack(err)
	}
	if err := c.Watch(source.NewKindWithCache(typ, ca), &handler.EnqueueRequestForObject{}); err != nil {
		return nil, errors.WithStack(err)
	}

	for _, s := range sources {
		if err := s.ApplyTo(mgr, setup, c, ca); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// dataFromSliceOrFile returns data from the slice (if non-empty), or from the file,
// or an error if an error occurred reading the file
func dataFromSliceOrFile(data []byte, file string) ([]byte, error) {
	if len(data) > 0 {
		return data, nil
	}
	if len(file) > 0 {
		fileData, err := ioutil.ReadFile(file)
		if err != nil {
			return []byte{}, err
		}
		return fileData, nil
	}
	return nil, nil
}

// RegisterCertificateBased registers all known injection controllers that
// target Certificate resources with the  given manager, and adds relevant
// indices.
// The registered controllers require the cert-manager API to be available
// in order to run.
func RegisterCertificateBased(ctx context.Context, mgr ctrl.Manager) error {
	cache, client, err := newIndependentCacheAndDelegatingClient(mgr)
	if err != nil {
		return err
	}
	return registerAllInjectors(
		ctx,
		"certificate",
		mgr,
		[]caDataSource{
			&certificateDataSource{client: cache},
		},
		client,
		cache,
	)
}

// RegisterSecretBased registers all known injection controllers that
// target Secret resources with the  given manager, and adds relevant
// indices.
// The registered controllers only require the corev1 APi to be available in
// order to run.
func RegisterSecretBased(ctx context.Context, mgr ctrl.Manager) error {
	cache, client, err := newIndependentCacheAndDelegatingClient(mgr)
	if err != nil {
		return err
	}
	return registerAllInjectors(
		ctx,
		"secret",
		mgr,
		[]caDataSource{
			&secretDataSource{client: cache},
			&kubeconfigDataSource{},
		},
		client,
		cache,
	)
}

func newIndependentCacheAndDelegatingClient(mgr ctrl.Manager) (cache.Cache, client.Client, error) {
	cacheOptions := cache.Options{
		Scheme: mgr.GetScheme(),
		Mapper: mgr.GetRESTMapper(),
	}
	ca, err := cache.New(mgr.GetConfig(), cacheOptions)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}

	clientOptions := client.Options{
		Scheme: mgr.GetScheme(),
		Mapper: mgr.GetRESTMapper(),
	}
	client, err := manager.DefaultNewClient(ca, mgr.GetConfig(), clientOptions)
	if err != nil {
		return nil, nil, errors.WithStack(err)
	}
	return ca, client, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	log "github.com/Sirupsen/logrus"
	docker_types "github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/docker/infrakit/pkg/plugin/group/types"
	"github.com/docker/infrakit/pkg/plugin/group/util"
	"github.com/docker/infrakit/pkg/spi/flavor"
	"github.com/docker/infrakit/pkg/spi/instance"
	"golang.org/x/net/context"
	"text/template"
)

type nodeType string

const (
	worker        nodeType = "worker"
	manager       nodeType = "manager"
	ebsAttachment string   = "ebs"
)

// NewSwarmFlavor creates a flavor.Plugin that creates manager and worker nodes connected in a swarm.
func NewSwarmFlavor(dockerClient client.APIClient) flavor.Plugin {
	return &swarmFlavor{client: dockerClient}
}

type swarmFlavor struct {
	client client.APIClient
}

type schema struct {
	Type                 nodeType
	Attachments          map[instance.LogicalID][]instance.Attachment
	DockerRestartCommand string
}

func parseProperties(flavorProperties json.RawMessage) (schema, error) {
	s := schema{}
	err := json.Unmarshal(flavorProperties, &s)
	return s, err
}

func validateIDsAndAttachments(logicalIDs []instance.LogicalID, attachments map[instance.LogicalID][]instance.Attachment) error {
	// Each attachment association must be represented by a logical ID.
	idsMap := map[instance.LogicalID]bool{}
	for _, id := range logicalIDs {
		if _, exists := idsMap[id]; exists {
			return fmt.Errorf("LogicalID %v specified more than once", id)
		}

		idsMap[id] = true
	}
	for id := range attachments {
		if _, exists := idsMap[id]; !exists {
			return fmt.Errorf("LogicalID %v used for an attachment but is not in group LogicalIDs", id)
		}
	}

	// Only EBS attachments are supported.
	for _, atts := range attachments {
		for _, attachment := range atts {
			if attachment.Type == "" {
				return fmt.Errorf(
					"Attachment Type %s must be specified for '%s'",
					ebsAttachment,
					attachment.ID)
			}

			if attachment.Type != ebsAttachment {
				return fmt.Errorf(
					"Invalid attachment Type '%s', only %s is supported",
					attachment.Type,
					ebsAttachment)
			}
		}
	}

	// Each attachment may only be used once.
	allAttachmentIDs := map[string]bool{}
	for _, atts := range attachments {
		for _, attachment := range atts {
			if _, exists := allAttachmentIDs[attachment.ID]; exists {
				return fmt.Errorf("Attachment %v specified more than once", attachment.ID)
			}
			allAttachmentIDs[attachment.ID] = true
		}
	}

	return nil
}

func (s swarmFlavor) Validate(flavorProperties json.RawMessage, allocation types.AllocationMethod) error {
	properties, err := parseProperties(flavorProperties)
	if err != nil {
		return err
	}

	if properties.DockerRestartCommand == "" {
		return errors.New("DockerRestartCommand must be specified")
	}

	switch properties.Type {
	case manager:
		numIDs := len(allocation.LogicalIDs)
		if numIDs != 1 && numIDs != 3 && numIDs != 5 {
			return errors.New("Must have 1, 3, or 5 manager logical IDs")
		}
	case worker:

	default:
		return errors.New("Unrecognized node Type")
	}

	if properties.Type == manager {
		for _, id := range allocation.LogicalIDs {
			if att, exists := properties.Attachments[id]; !exists || len(att) == 0 {
				log.Warnf("LogicalID %s has no attachments, which is needed for durability", id)
			}
		}
	}

	if err := validateIDsAndAttachments(allocation.LogicalIDs, properties.Attachments); err != nil {
		return err
	}

	return nil
}

const (
	// associationTag is a machine tag added to associate machines with Swarm nodes.
	associationTag = "swarm-association-id"

	// bootScript is used to generate node boot scripts.
	bootScript = `#!/bin/sh
set -o errexit
set -o nounset
set -o xtrace

mkdir -p /etc/docker
cat << EOF > /etc/docker/daemon.json
{
  "labels": ["swarm-association-id={{.ASSOCIATION_ID}}"]
}
EOF

{{.RESTART_DOCKER}}

docker swarm join {{.MY_IP}} --token {{.JOIN_TOKEN}}
`
)

func generateInitScript(joinIP, joinToken, associationID, restartCommand string) string {
	buffer := bytes.Buffer{}
	templ := template.Must(template.New("").Parse(bootScript))
	err := templ.Execute(&buffer, map[string]string{
		"MY_IP":          joinIP,
		"JOIN_TOKEN":     joinToken,
		"ASSOCIATION_ID": associationID,
		"RESTART_DOCKER": restartCommand,
	})
	if err != nil {
		panic(err)
	}
	return buffer.String()
}

// Healthy determines whether an instance is healthy.  This is determined by whether it has successfully joined the
// Swarm.
func (s swarmFlavor) Healthy(flavorProperties json.RawMessage, inst instance.Description) (flavor.Health, error) {
	associationID, exists := inst.Tags[associationTag]
	if !exists {
		log.Info("Reporting unhealthy for instance without an association tag", inst.ID)
		return flavor.Unhealthy, nil
	}

	filter := filters.NewArgs()
	filter.Add("label", fmt.Sprintf("%s=%s", associationTag, associationID))

	nodes, err := s.client.NodeList(context.Background(), docker_types.NodeListOptions{Filters: filter})
	if err != nil {
		return flavor.Unknown, err
	}

	switch {
	case len(nodes) == 0:
		// The instance may not yet be joined, so we consider the health unknown.
		return flavor.Unknown, nil

	case len(nodes) == 1:
		return flavor.Healthy, nil

	default:
		log.Warnf("Expected at most one node with label %s, but found %s", associationID, nodes)
		return flavor.Healthy, nil
	}
}

func (s swarmFlavor) Drain(flavorProperties json.RawMessage, inst instance.Description) error {
	properties, err := parseProperties(flavorProperties)
	if err != nil {
		return err
	}

	// Only explicitly remove worker nodes, not manager nodes.  Manager nodes are assumed to have an
	// attached volume for state, and fixed IP addresses.  This allows them to rejoin as the same node.
	if properties.Type != worker {
		return nil
	}

	associationID, exists := inst.Tags[associationTag]
	if !exists {
		return fmt.Errorf("Unable to drain %s without an association tag", inst.ID)
	}

	filter := filters.NewArgs()
	filter.Add("label", fmt.Sprintf("%s=%s", associationTag, associationID))

	nodes, err := s.client.NodeList(context.Background(), docker_types.NodeListOptions{Filters: filter})
	if err != nil {
		return err
	}

	switch {
	case len(nodes) == 0:
		return fmt.Errorf("Unable to drain %s, not found in swarm", inst.ID)

	case len(nodes) == 1:
		err := s.client.NodeRemove(
			context.Background(),
			nodes[0].ID,
			docker_types.NodeRemoveOptions{Force: true})
		if err != nil {
			return err
		}

		return nil

	default:
		return fmt.Errorf("Expected at most one node with label %s, but found %s", associationID, nodes)
	}
}

func (s swarmFlavor) Prepare(
	flavorProperties json.RawMessage,
	spec instance.Spec,
	allocation types.AllocationMethod) (instance.Spec, error) {

	properties, err := parseProperties(flavorProperties)
	if err != nil {
		return spec, err
	}

	swarmStatus, err := s.client.SwarmInspect(context.Background())
	if err != nil {
		return spec, fmt.Errorf("Failed to fetch Swarm join tokens: %s", err)
	}

	nodeInfo, err := s.client.Info(context.Background())
	if err != nil {
		return spec, fmt.Errorf("Failed to fetch node self info: %s", err)
	}

	self, _, err := s.client.NodeInspectWithRaw(context.Background(), nodeInfo.Swarm.NodeID)
	if err != nil {
		return spec, fmt.Errorf("Failed to fetch Swarm node status: %s", err)
	}

	if self.ManagerStatus == nil {
		return spec, errors.New(
			"Swarm node status did not include manager status.  Need to run 'docker swarm init`?")
	}

	associationID := util.RandomAlphaNumericString(8)
	spec.Tags[associationTag] = associationID

	switch properties.Type {
	case worker:
		spec.Init = generateInitScript(
			self.ManagerStatus.Addr,
			swarmStatus.JoinTokens.Worker,
			associationID,
			properties.DockerRestartCommand)

	case manager:
		if spec.LogicalID == nil {
			return spec, errors.New("Manager nodes require a LogicalID, " +
				"which will be used as an assigned private IP address")
		}

		spec.Init = generateInitScript(
			self.ManagerStatus.Addr,
			swarmStatus.JoinTokens.Manager,
			associationID,
			properties.DockerRestartCommand)

	default:
		return spec, errors.New("Unsupported node type")
	}

	if spec.LogicalID != nil {
		if attachments, exists := properties.Attachments[*spec.LogicalID]; exists {
			spec.Attachments = append(spec.Attachments, attachments...)
		}
	}

	// TODO(wfarner): Use the cluster UUID to scope instances for this swarm separately from instances in another
	// swarm.  This will require plumbing back to Scaled (membership tags).
	spec.Tags["swarm-id"] = swarmStatus.ID

	return spec, nil
}
<commit_msg>adding debug statement and ensure we are using the correct port<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"text/template"

	log "github.com/Sirupsen/logrus"
	docker_types "github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/client"
	"github.com/docker/infrakit/pkg/plugin/group/types"
	"github.com/docker/infrakit/pkg/plugin/group/util"
	"github.com/docker/infrakit/pkg/spi/flavor"
	"github.com/docker/infrakit/pkg/spi/instance"
	"golang.org/x/net/context"
)

type nodeType string

const (
	worker        nodeType = "worker"
	manager       nodeType = "manager"
	ebsAttachment string   = "ebs"
)

// NewSwarmFlavor creates a flavor.Plugin that creates manager and worker nodes connected in a swarm.
func NewSwarmFlavor(dockerClient client.APIClient) flavor.Plugin {
	return &swarmFlavor{client: dockerClient}
}

type swarmFlavor struct {
	client client.APIClient
}

type schema struct {
	Type                 nodeType
	Attachments          map[instance.LogicalID][]instance.Attachment
	DockerRestartCommand string
}

func parseProperties(flavorProperties json.RawMessage) (schema, error) {
	s := schema{}
	err := json.Unmarshal(flavorProperties, &s)
	return s, err
}

func validateIDsAndAttachments(logicalIDs []instance.LogicalID, attachments map[instance.LogicalID][]instance.Attachment) error {
	// Each attachment association must be represented by a logical ID.
	idsMap := map[instance.LogicalID]bool{}
	for _, id := range logicalIDs {
		if _, exists := idsMap[id]; exists {
			return fmt.Errorf("LogicalID %v specified more than once", id)
		}

		idsMap[id] = true
	}
	for id := range attachments {
		if _, exists := idsMap[id]; !exists {
			return fmt.Errorf("LogicalID %v used for an attachment but is not in group LogicalIDs", id)
		}
	}

	// Only EBS attachments are supported.
	for _, atts := range attachments {
		for _, attachment := range atts {
			if attachment.Type == "" {
				return fmt.Errorf(
					"Attachment Type %s must be specified for '%s'",
					ebsAttachment,
					attachment.ID)
			}

			if attachment.Type != ebsAttachment {
				return fmt.Errorf(
					"Invalid attachment Type '%s', only %s is supported",
					attachment.Type,
					ebsAttachment)
			}
		}
	}

	// Each attachment may only be used once.
	allAttachmentIDs := map[string]bool{}
	for _, atts := range attachments {
		for _, attachment := range atts {
			if _, exists := allAttachmentIDs[attachment.ID]; exists {
				return fmt.Errorf("Attachment %v specified more than once", attachment.ID)
			}
			allAttachmentIDs[attachment.ID] = true
		}
	}

	return nil
}

func (s swarmFlavor) Validate(flavorProperties json.RawMessage, allocation types.AllocationMethod) error {
	properties, err := parseProperties(flavorProperties)
	if err != nil {
		return err
	}

	if properties.DockerRestartCommand == "" {
		return errors.New("DockerRestartCommand must be specified")
	}

	switch properties.Type {
	case manager:
		numIDs := len(allocation.LogicalIDs)
		if numIDs != 1 && numIDs != 3 && numIDs != 5 {
			return errors.New("Must have 1, 3, or 5 manager logical IDs")
		}
	case worker:

	default:
		return errors.New("Unrecognized node Type")
	}

	if properties.Type == manager {
		for _, id := range allocation.LogicalIDs {
			if att, exists := properties.Attachments[id]; !exists || len(att) == 0 {
				log.Warnf("LogicalID %s has no attachments, which is needed for durability", id)
			}
		}
	}

	if err := validateIDsAndAttachments(allocation.LogicalIDs, properties.Attachments); err != nil {
		return err
	}

	return nil
}

const (
	// associationTag is a machine tag added to associate machines with Swarm nodes.
	associationTag = "swarm-association-id"

	// bootScript is used to generate node boot scripts.
	bootScript = `#!/bin/sh
set -o errexit
set -o nounset
set -o xtrace

mkdir -p /etc/docker
cat << EOF > /etc/docker/daemon.json
{
  "labels": ["swarm-association-id={{.ASSOCIATION_ID}}"]
}
EOF

{{.RESTART_DOCKER}}

docker swarm join {{.MY_IP}} --token {{.JOIN_TOKEN}}
`
)

func generateInitScript(joinIP, joinToken, associationID, restartCommand string) string {
	buffer := bytes.Buffer{}
	templ := template.Must(template.New("").Parse(bootScript))
	err := templ.Execute(&buffer, map[string]string{
		"MY_IP":          joinIP,
		"JOIN_TOKEN":     joinToken,
		"ASSOCIATION_ID": associationID,
		"RESTART_DOCKER": restartCommand,
	})
	if err != nil {
		panic(err)
	}
	return buffer.String()
}

// Healthy determines whether an instance is healthy.  This is determined by whether it has successfully joined the
// Swarm.
func (s swarmFlavor) Healthy(flavorProperties json.RawMessage, inst instance.Description) (flavor.Health, error) {
	associationID, exists := inst.Tags[associationTag]
	if !exists {
		log.Info("Reporting unhealthy for instance without an association tag", inst.ID)
		return flavor.Unhealthy, nil
	}

	filter := filters.NewArgs()
	filter.Add("label", fmt.Sprintf("%s=%s", associationTag, associationID))

	nodes, err := s.client.NodeList(context.Background(), docker_types.NodeListOptions{Filters: filter})
	if err != nil {
		return flavor.Unknown, err
	}

	switch {
	case len(nodes) == 0:
		// The instance may not yet be joined, so we consider the health unknown.
		return flavor.Unknown, nil

	case len(nodes) == 1:
		return flavor.Healthy, nil

	default:
		log.Warnf("Expected at most one node with label %s, but found %s", associationID, nodes)
		return flavor.Healthy, nil
	}
}

func (s swarmFlavor) Drain(flavorProperties json.RawMessage, inst instance.Description) error {
	properties, err := parseProperties(flavorProperties)
	if err != nil {
		return err
	}

	// Only explicitly remove worker nodes, not manager nodes.  Manager nodes are assumed to have an
	// attached volume for state, and fixed IP addresses.  This allows them to rejoin as the same node.
	if properties.Type != worker {
		return nil
	}

	associationID, exists := inst.Tags[associationTag]
	if !exists {
		return fmt.Errorf("Unable to drain %s without an association tag", inst.ID)
	}

	filter := filters.NewArgs()
	filter.Add("label", fmt.Sprintf("%s=%s", associationTag, associationID))

	nodes, err := s.client.NodeList(context.Background(), docker_types.NodeListOptions{Filters: filter})
	if err != nil {
		return err
	}

	switch {
	case len(nodes) == 0:
		return fmt.Errorf("Unable to drain %s, not found in swarm", inst.ID)

	case len(nodes) == 1:
		err := s.client.NodeRemove(
			context.Background(),
			nodes[0].ID,
			docker_types.NodeRemoveOptions{Force: true})
		if err != nil {
			return err
		}

		return nil

	default:
		return fmt.Errorf("Expected at most one node with label %s, but found %s", associationID, nodes)
	}
}

func (s swarmFlavor) Prepare(
	flavorProperties json.RawMessage,
	spec instance.Spec,
	allocation types.AllocationMethod) (instance.Spec, error) {

	log.Info("Prepare phase flavorProperties=", flavorProperties)
	log.WithFields(log.Fields{
		"flavorProperties": flavorProperties,
		"spec":             spec,
		"allocation":       allocation,
	}).Info("arguments into Prepare")

	properties, err := parseProperties(flavorProperties)
	if err != nil {
		log.Error("Prepare phase parse properties error=", err)
		return spec, err
	}

	swarmStatus, err := s.client.SwarmInspect(context.Background())
	if err != nil {
		log.Error("Prepare phase failed to fetch Swarm join tokens error=", err)
		return spec, fmt.Errorf("Failed to fetch Swarm join tokens: %s", err)
	}

	nodeInfo, err := s.client.Info(context.Background())
	if err != nil {
		log.Error("Prepare phase failed to fetch node self info error=", err)
		return spec, fmt.Errorf("Failed to fetch node self info: %s", err)
	}

	self, _, err := s.client.NodeInspectWithRaw(context.Background(), nodeInfo.Swarm.NodeID)
	if err != nil {
		log.Error("Prepare phase failed to fetch Swarm node status error=", err)
		return spec, fmt.Errorf("Failed to fetch Swarm node status: %s", err)
	}

	if self.ManagerStatus == nil {
		return spec, errors.New(
			"Swarm node status did not include manager status.  Need to run 'docker swarm init`?")
	}

	associationID := util.RandomAlphaNumericString(8)
	spec.Tags[associationTag] = associationID

	switch properties.Type {
	case worker:
		spec.Init = generateInitScript(
			self.ManagerStatus.Addr,
			swarmStatus.JoinTokens.Worker,
			associationID,
			properties.DockerRestartCommand)

	case manager:
		if spec.LogicalID == nil {
			return spec, errors.New("Manager nodes require a LogicalID, " +
				"which will be used as an assigned private IP address")
		}

		// check if mgr status addr uses port 2377, if not, force to use port 2377
		addr := self.ManagerStatus.Addr
		if strings.HasSuffix(self.ManagerStatus.Addr, ":2377") == false {
			addr = strings.TrimSuffix(self.ManagerStatus.Addr, ":2377") + ":2377"
		}
		log.Info("Prepare phase roleWorker type detected mgr address is ", addr)

		spec.Init = generateInitScript(
			self.ManagerStatus.Addr,
			swarmStatus.JoinTokens.Manager,
			associationID,
			properties.DockerRestartCommand)

	default:
		return spec, errors.New("Unsupported node type")
	}

	if spec.LogicalID != nil {
		if attachments, exists := properties.Attachments[*spec.LogicalID]; exists {
			spec.Attachments = append(spec.Attachments, attachments...)
		}
	}

	// TODO(wfarner): Use the cluster UUID to scope instances for this swarm separately from instances in another
	// swarm.  This will require plumbing back to Scaled (membership tags).
	spec.Tags["swarm-id"] = swarmStatus.ID

	log.Info("Prepare phase finished with no error, spec=", spec)

	return spec, nil
}
<|endoftext|>