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

package cni

import (
	"testing"

	"istio.io/istio/pkg/test/framework"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/environment"
	"istio.io/istio/pkg/test/framework/components/galley"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/pilot"
	"istio.io/istio/tests/integration/security/util/reachability"
)

func TestMain(m *testing.M) {
	framework.
		NewSuite("cni", m).
		SetupOnEnv(environment.Kube, istio.Setup(nil, func(cfg *istio.Config) {
			cfg.ControlPlaneValues = `
components:
  cni:
     enabled: true
values:
  cni:
     hub: gcr.io/istio-testing
     tag: latest
`
		})).
		Run()
}

// This test verifies reachability under different authN scenario:
// - app A to app B using mTLS.
// In each test, the steps are:
// - Configure authn policy.
// - Wait for config propagation.
// - Send HTTP/gRPC requests between apps.
func TestCNIReachability(t *testing.T) {
	framework.NewTest(t).
		Run(func(ctx framework.TestContext) {
			g, err := galley.New(ctx, galley.Config{})
			if err != nil {
				ctx.Fatal(err)
			}
			p, err := pilot.New(ctx, pilot.Config{
				Galley: g,
			})
			if err != nil {
				ctx.Fatal(err)
			}
			rctx := reachability.CreateContext(ctx, g, p)
			systemNM := namespace.ClaimSystemNamespaceOrFail(ctx, ctx)

			testCases := []reachability.TestCase{
				{
					ConfigFile:          "global-mtls-on.yaml",
					Namespace:           systemNM,
					RequiredEnvironment: environment.Kube,
					Include: func(src echo.Instance, opts echo.CallOptions) bool {
						// Exclude calls to the headless TCP port.
						if opts.Target == rctx.Headless && opts.PortName == "tcp" {
							return false
						}
						return true
					},
					ExpectSuccess: func(src echo.Instance, opts echo.CallOptions) bool {
						if src == rctx.Naked && opts.Target == rctx.Naked {
							// naked->naked should always succeed.
							return true
						}

						// If one of the two endpoints is naked, expect failure.
						return src != rctx.Naked && opts.Target != rctx.Naked
					},
				},
			}
			rctx.Run(testCases)
		})
}
<commit_msg>Update the CNI integ test with component level hub and tag. (#20516)<commit_after>//  Copyright 2020 Istio Authors
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

package cni

import (
	"testing"

	"istio.io/istio/pkg/test/framework"
	"istio.io/istio/pkg/test/framework/components/echo"
	"istio.io/istio/pkg/test/framework/components/environment"
	"istio.io/istio/pkg/test/framework/components/galley"
	"istio.io/istio/pkg/test/framework/components/istio"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/pilot"
	"istio.io/istio/tests/integration/security/util/reachability"
)

func TestMain(m *testing.M) {
	framework.
		NewSuite("cni", m).
		SetupOnEnv(environment.Kube, istio.Setup(nil, func(cfg *istio.Config) {
			cfg.ControlPlaneValues = `
components:
  cni:
     enabled: true
     hub: gcr.io/istio-testing
     tag: latest
`
		})).
		Run()
}

// This test verifies reachability under different authN scenario:
// - app A to app B using mTLS.
// In each test, the steps are:
// - Configure authn policy.
// - Wait for config propagation.
// - Send HTTP/gRPC requests between apps.
func TestCNIReachability(t *testing.T) {
	framework.NewTest(t).
		Run(func(ctx framework.TestContext) {
			g, err := galley.New(ctx, galley.Config{})
			if err != nil {
				ctx.Fatal(err)
			}
			p, err := pilot.New(ctx, pilot.Config{
				Galley: g,
			})
			if err != nil {
				ctx.Fatal(err)
			}
			rctx := reachability.CreateContext(ctx, g, p)
			systemNM := namespace.ClaimSystemNamespaceOrFail(ctx, ctx)

			testCases := []reachability.TestCase{
				{
					ConfigFile:          "global-mtls-on.yaml",
					Namespace:           systemNM,
					RequiredEnvironment: environment.Kube,
					Include: func(src echo.Instance, opts echo.CallOptions) bool {
						// Exclude calls to the headless TCP port.
						if opts.Target == rctx.Headless && opts.PortName == "tcp" {
							return false
						}
						return true
					},
					ExpectSuccess: func(src echo.Instance, opts echo.CallOptions) bool {
						if src == rctx.Naked && opts.Target == rctx.Naked {
							// naked->naked should always succeed.
							return true
						}

						// If one of the two endpoints is naked, expect failure.
						return src != rctx.Naked && opts.Target != rctx.Naked
					},
				},
			}
			rctx.Run(testCases)
		})
}
<|endoftext|>


<commit_before>package venom

import (
	"context"
	"fmt"
	"time"
)

//RunTestStep executes a venom testcase is a venom context
func (v *Venom) RunTestStep(tcc TestCaseContext, e *ExecutorWrap, ts *TestSuite, tc *TestCase, step TestStep, l Logger) ExecutorResult {
	var assertRes assertionsApplied

	var retry int
	var result ExecutorResult

	for retry = 0; retry <= e.retry && !assertRes.ok; retry++ {
		if retry > 1 && !assertRes.ok {
			l.Debugf("Sleep %d, it's %d attempt", e.delay, retry)
			time.Sleep(time.Duration(e.delay) * time.Second)
		}

		var err error
		result, err = runTestStepExecutor(tcc, e, ts, step, l)

		if err != nil {
			tc.Failures = append(tc.Failures, Failure{Value: RemoveNotPrintableChar(err.Error())})
			continue
		}

		// add result in templater
		ts.Templater.Add(tc.Name, stringifyExecutorResult(result))

		l.Debugf("Apply assertions")

		if h, ok := e.executor.(executorWithDefaultAssertions); ok {
			assertRes = applyChecks(&result, step, h.GetDefaultAssertions(), l)
		} else {
			assertRes = applyChecks(&result, step, nil, l)
		}
		// add result again for extracts values
		ts.Templater.Add(tc.Name, stringifyExecutorResult(result))

		l.Debugf("result step:%+v", result)

		if assertRes.ok {
			break
		}
	}
	tc.Errors = append(tc.Errors, assertRes.errors...)
	tc.Failures = append(tc.Failures, assertRes.failures...)
	if retry > 1 && (len(assertRes.failures) > 0 || len(assertRes.errors) > 0) {
		tc.Failures = append(tc.Failures, Failure{Value: fmt.Sprintf("It's a failure after %d attempts", retry)})
	}
	tc.Systemout.Value += assertRes.systemout
	tc.Systemerr.Value += assertRes.systemerr

	return result
}

func stringifyExecutorResult(e ExecutorResult) map[string]string {
	out := make(map[string]string)
	for k, v := range e {
		out[k] = fmt.Sprintf("%v", v)
	}
	return out
}

func runTestStepExecutor(tcc TestCaseContext, e *ExecutorWrap, ts *TestSuite, step TestStep, l Logger) (ExecutorResult, error) {
	if e.timeout == 0 {
		return e.executor.Run(tcc, l, step, ts.WorkDir)
	}

	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Duration(e.timeout)*time.Second)
	defer cancel()

	ch := make(chan ExecutorResult)
	cherr := make(chan error)
	go func(tcc TestCaseContext, e *ExecutorWrap, step TestStep, l Logger) {
		result, err := e.executor.Run(tcc, l, step, ts.WorkDir)
		if err != nil {
			cherr <- err
		} else {
			ch <- result
		}
	}(tcc, e, step, l)

	select {
	case err := <-cherr:
		return nil, err
	case result := <-ch:
		return result, nil
	case <-ctxTimeout.Done():
		return nil, fmt.Errorf("Timeout after %d second(s)", e.timeout)
	}
}
<commit_msg>fix: do not count retry failures (#154) (#155)<commit_after>package venom

import (
	"context"
	"fmt"
	"time"
)

//RunTestStep executes a venom testcase is a venom context
func (v *Venom) RunTestStep(tcc TestCaseContext, e *ExecutorWrap, ts *TestSuite, tc *TestCase, step TestStep, l Logger) ExecutorResult {
	var assertRes assertionsApplied

	var retry int
	var result ExecutorResult

	for retry = 0; retry <= e.retry && !assertRes.ok; retry++ {
		if retry > 1 && !assertRes.ok {
			l.Debugf("Sleep %d, it's %d attempt", e.delay, retry)
			time.Sleep(time.Duration(e.delay) * time.Second)
		}

		var err error
		result, err = runTestStepExecutor(tcc, e, ts, step, l)

		if err != nil {
			continue
		}

		// add result in templater
		ts.Templater.Add(tc.Name, stringifyExecutorResult(result))

		l.Debugf("Apply assertions")

		if h, ok := e.executor.(executorWithDefaultAssertions); ok {
			assertRes = applyChecks(&result, step, h.GetDefaultAssertions(), l)
		} else {
			assertRes = applyChecks(&result, step, nil, l)
		}
		// add result again for extracts values
		ts.Templater.Add(tc.Name, stringifyExecutorResult(result))

		l.Debugf("result step:%+v", result)

		if assertRes.ok {
			break
		}
	}
	tc.Errors = append(tc.Errors, assertRes.errors...)
	tc.Failures = append(tc.Failures, assertRes.failures...)
	if retry > 1 && (len(assertRes.failures) > 0 || len(assertRes.errors) > 0) {
		tc.Failures = append(tc.Failures, Failure{Value: fmt.Sprintf("It's a failure after %d attempts", retry)})
	}
	tc.Systemout.Value += assertRes.systemout
	tc.Systemerr.Value += assertRes.systemerr

	return result
}

func stringifyExecutorResult(e ExecutorResult) map[string]string {
	out := make(map[string]string)
	for k, v := range e {
		out[k] = fmt.Sprintf("%v", v)
	}
	return out
}

func runTestStepExecutor(tcc TestCaseContext, e *ExecutorWrap, ts *TestSuite, step TestStep, l Logger) (ExecutorResult, error) {
	if e.timeout == 0 {
		return e.executor.Run(tcc, l, step, ts.WorkDir)
	}

	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Duration(e.timeout)*time.Second)
	defer cancel()

	ch := make(chan ExecutorResult)
	cherr := make(chan error)
	go func(tcc TestCaseContext, e *ExecutorWrap, step TestStep, l Logger) {
		result, err := e.executor.Run(tcc, l, step, ts.WorkDir)
		if err != nil {
			cherr <- err
		} else {
			ch <- result
		}
	}(tcc, e, step, l)

	select {
	case err := <-cherr:
		return nil, err
	case result := <-ch:
		return result, nil
	case <-ctxTimeout.Done():
		return nil, fmt.Errorf("Timeout after %d second(s)", e.timeout)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"net/http"
	"net/url"
	"encoding/json"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"github.com/markllama/hexgame/types/db"
	"github.com/markllama/hexgame/types/api"
)

// MatchHandleFunc processes and responds to HTTP queries
// GET/name - return one
// GET/     - return references
// POST     - create a new one
func MatchHandleFunc(s *mgo.Session) (func(w http.ResponseWriter, r *http.Request)) {

	f := func(w http.ResponseWriter, r *http.Request) {

		sc := s.Copy()
		defer sc.Close()

		switch r.Method {
		// return an existing match or list of matches
		case http.MethodGet:
			_, id := path.Split(r.URL.Path)
			if (id == "") {
				GetMatchList(s, w, r)
			} else {
				GetMatch(s, id, w, r)
			}
		// create a new match
		case http.MethodPost:
			CreateMatch(s, w, r)

		// update/replace an existing match
		//case http.MethodPut:
		//	UpdateMatch(s, w, r)

		// delete a match
		//case http.MethodDelete:
		//	DeleteMatch(s, w, r)
		}
	}

	return f
}

func GetMatchList(s *mgo.Session, w http.ResponseWriter, r *http.Request) {

	var m []db.Match

	c := s.DB("hexgame").C("matches")
	c.Find(nil).All(&m)

	matchrefs := make([]api.MatchRef, len(m))

	murl := url.URL{Scheme: "http", Host: r.Host}

	for index, match := range m {
		murl.Path = path.Join(r.URL.Path, match.Id.Hex())
		matchrefs[index].Id = match.Id
		matchrefs[index].GameId = match.GameId
		matchrefs[index].URL = murl.String()
	}

	jmatch, _ := json.Marshal(matchrefs)

	w.Write([]byte(jmatch))
}

func GetMatch(s *mgo.Session, id string, w http.ResponseWriter, r *http.Request) {
	var m db.Match

	c := s.DB("hexgame").C("matches")
	//c.Find(nil).All(&m)

	q := c.Find(bson.M{"_id": bson.ObjectIdHex(id)})
	// check for errors
	err := q.One(&m)
	if (err != nil) {
		http.Error(w, fmt.Sprintf("match %s not found", id), http.StatusNotFound)
		return
	}

	murl := url.URL{Scheme: "http", Host: r.Host, Path: r.URL.Path}
	m.URL = murl.String()
	p, _ := json.Marshal(m)
	w.Write(p)
}

func CreateMatch(s *mgo.Session, w http.ResponseWriter, r *http.Request) {

	//
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Printf("Error reading body: %v", err)
		http.Error(w, "can't read body", http.StatusBadRequest)
		return
	}

	var m db.Match

	// marshal the POST data into a Match
	err = json.Unmarshal(body, &m)
	if err != nil {
		log.Printf("Error reading body: %v", err)
		http.Error(w, "can't read body", http.StatusBadRequest)
		return
	}

	// Required: owner ID
	// Required: game ID
}

//func DeleteMatch(s *mgo.Session, w, http.ResponseWriter, r *http.Response) {
//
//}
<commit_msg>retrieve game name for match<commit_after>package server

import (
	"fmt"
	"io/ioutil"
	"log"
	"path"
	"net/http"
	"net/url"
	"encoding/json"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"github.com/markllama/hexgame/types/db"
	"github.com/markllama/hexgame/types/api"
)

// MatchHandleFunc processes and responds to HTTP queries
// GET/name - return one
// GET/     - return references
// POST     - create a new one
func MatchHandleFunc(s *mgo.Session) (func(w http.ResponseWriter, r *http.Request)) {

	f := func(w http.ResponseWriter, r *http.Request) {

		sc := s.Copy()
		defer sc.Close()

		switch r.Method {
		// return an existing match or list of matches
		case http.MethodGet:
			_, id := path.Split(r.URL.Path)
			if (id == "") {
				GetMatchList(s, w, r)
			} else {
				GetMatch(s, id, w, r)
			}
		// create a new match
		case http.MethodPost:
			CreateMatch(s, w, r)

		// update/replace an existing match
		//case http.MethodPut:
		//	UpdateMatch(s, w, r)

		// delete a match
		//case http.MethodDelete:
		//	DeleteMatch(s, w, r)
		}
	}

	return f
}

func GetMatchList(s *mgo.Session, w http.ResponseWriter, r *http.Request) {

	var m []db.Match

	c := s.DB("hexgame").C("matches")
	c.Find(nil).All(&m)

	matchrefs := make([]api.MatchRef, len(m))

	murl := url.URL{Scheme: "http", Host: r.Host}

	for index, match := range m {
		murl.Path = path.Join(r.URL.Path, match.Id.Hex())
		matchrefs[index].Id = match.Id
		matchrefs[index].GameId = match.GameId
		matchrefs[index].URL = murl.String()
	}

	jmatch, _ := json.Marshal(matchrefs)

	w.Write([]byte(jmatch))
}

func GetMatch(s *mgo.Session, id string, w http.ResponseWriter, r *http.Request) {
	var m db.Match

	c_matches := s.DB("hexgame").C("matches")
	//c.Find(nil).All(&m)

	q := c_matches.Find(bson.M{"_id": bson.ObjectIdHex(id)})
	// check for errors
	err := q.One(&m)
	if (err != nil) {
		http.Error(w, fmt.Sprintf("match %s not found", id), http.StatusNotFound)
		return
	}

	c_games := s.DB("hexgame").C("games")

	var g db.Game

	q_games := c_games.Find(bson.M{"_id": m.GameId})
	err = q_games.One(&g)
	if (err != nil) {
		http.Error(w, fmt.Sprintf("game not found"), http.StatusInternalServerError)
		return
	}

	murl := url.URL{Scheme: "http", Host: r.Host, Path: r.URL.Path}
	m.URL = murl.String()
	p, _ := json.Marshal(m)
	w.Write(p)
}

func CreateMatch(s *mgo.Session, w http.ResponseWriter, r *http.Request) {

	//
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Printf("Error reading body: %v", err)
		http.Error(w, "can't read body", http.StatusBadRequest)
		return
	}

	var m db.Match

	// marshal the POST data into a Match
	err = json.Unmarshal(body, &m)
	if err != nil {
		log.Printf("Error reading body: %v", err)
		http.Error(w, "can't read body", http.StatusBadRequest)
		return
	}

	// Required: owner ID
	// Required: game ID
}

//func DeleteMatch(s *mgo.Session, w, http.ResponseWriter, r *http.Response) {
//
//}
<|endoftext|>


<commit_before>package transport

import (
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"

	"github.com/scrapli/scrapligo/logging"

	"golang.org/x/crypto/ssh"
)

// Standard the "standard" (standard library) transport option for scrapligo.
type Standard struct {
	BaseTransportArgs     *BaseTransportArgs
	StandardTransportArgs *StandardTransportArgs
	client                *ssh.Client
	session               *ssh.Session
	writer                io.WriteCloser
	reader                io.Reader
}

// StandardTransportArgs struct representing attributes required for the Standard transport.
type StandardTransportArgs struct {
	AuthPassword      string
	AuthPrivateKey    string
	AuthStrictKey     bool
	SSHConfigFile     string
	SSHKnownHostsFile string
}

func keyString(k ssh.PublicKey) string {
	return k.Type() + " " + base64.StdEncoding.EncodeToString(
		k.Marshal(),
	) // e.g. "ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTY...."
}

// https://stackoverflow.com/questions/44269142/ \
// golang-ssh-getting-must-specify-hoskeycallback-error-despite-setting-it-to-n
// basically need to parse ssh config like scrapli does... at some point.
func trustedHostKeyCallback(trustedKey string) ssh.HostKeyCallback {
	if trustedKey == "" {
		return func(_ string, _ net.Addr, k ssh.PublicKey) error {
			log.Printf(
				"ssh key verification is *NOT* in effect: to fix, add this trustedKey: %q",
				keyString(k),
			)

			return nil
		}
	}

	return func(_ string, _ net.Addr, k ssh.PublicKey) error {
		ks := keyString(k)
		if trustedKey != ks {
			return ErrKeyVerificationFailed
		}

		return nil
	}
}

func (t *Standard) open(cfg *ssh.ClientConfig) error {
	var err error
	t.client, err = ssh.Dial(
		"tcp",
		fmt.Sprintf("%s:%d", t.BaseTransportArgs.Host, t.BaseTransportArgs.Port),
		cfg,
	)
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error connecting to host: %v", err)),
		)

		return err
	}

	t.session, err = t.client.NewSession()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating session: %v", err)),
		)

		return err
	}

	t.writer, err = t.session.StdinPipe()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating writer: %v", err)),
		)

		return err
	}

	t.reader, err = t.session.StdoutPipe()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating reader: %v", err)),
		)

		return err
	}

	return nil
}

func (t *Standard) openBase() error {
	/* #nosec G106 */
	hostKeyCallback := ssh.InsecureIgnoreHostKey()
	if t.StandardTransportArgs.AuthStrictKey {
		// trustedKey will need to be gleaned from known hosts how scrapli does at some point
		hostKeyCallback = trustedHostKeyCallback("")
	}

	authMethods := make([]ssh.AuthMethod, 0)

	if t.StandardTransportArgs.AuthPrivateKey != "" {
		key, err := ioutil.ReadFile(t.StandardTransportArgs.AuthPrivateKey)
		if err != nil {
			return err
		}

		signer, err := ssh.ParsePrivateKey(key)

		if err != nil {
			logging.LogError(
				t.FormatLogMessage("error", fmt.Sprintf("unable to parse private key: %v", err)),
			)

			return err
		}

		authMethods = append(authMethods, ssh.PublicKeys(signer))
	}

	if t.StandardTransportArgs.AuthPassword != "" {
		authMethods = append(authMethods, ssh.Password(t.StandardTransportArgs.AuthPassword),
			ssh.KeyboardInteractive(
				func(user, instruction string, questions []string, echos []bool) ([]string, error) {
					answers := make([]string, len(questions))
					for i := range answers {
						answers[i] = t.StandardTransportArgs.AuthPassword
					}

					return answers, nil
				},
			))
	}

	cfg := &ssh.ClientConfig{
		User:            t.BaseTransportArgs.AuthUsername,
		Auth:            authMethods,
		Timeout:         *t.BaseTransportArgs.TimeoutSocket,
		HostKeyCallback: hostKeyCallback,
	}

	err := t.open(cfg)
	if err != nil {
		return err
	}

	// not sure what to do about the tty speeds... figured lets just go fast?
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 115200,
		ssh.TTY_OP_OSPEED: 115200,
	}

	err = t.session.RequestPty(
		"xterm",
		t.BaseTransportArgs.PtyHeight,
		t.BaseTransportArgs.PtyWidth,
		modes,
	)
	if err != nil {
		return err
	}

	return nil
}

// Open open a standard ssh connection.
func (t *Standard) Open() error {
	err := t.openBase()
	if err != nil {
		return err
	}

	err = t.session.Shell()
	if err != nil {
		return err
	}

	return nil
}

// OpenNetconf open a netconf connection.
func (t *Standard) OpenNetconf() error {
	err := t.openBase()
	if err != nil {
		logging.LogError(
			fmt.Sprintf(
				"failed opening base connection, cant attempt to open netconf connection; error: %v",
				err,
			),
		)

		return err
	}

	err = t.session.RequestSubsystem("netconf")
	if err != nil {
		logging.LogError(fmt.Sprintf("failed opening netconf subsystem; error: %v", err))
		return err
	}

	return nil
}

// Close close the transport connection to the device.
func (t *Standard) Close() error {
	err := t.session.Close()
	t.session = nil

	logging.LogDebug(t.FormatLogMessage("debug", "transport connection to host closed"))

	return err
}

func (t *Standard) read(n int) *transportResult {
	b := make([]byte, n)
	_, err := t.reader.Read(b)

	if err != nil {
		return &transportResult{
			result: nil,
			error:  ErrTransportFailure,
		}
	}

	return &transportResult{
		result: b,
		error:  nil,
	}
}

// Read read bytes from the transport.
func (t *Standard) Read() ([]byte, error) {
	b, err := transportTimeout(
		*t.BaseTransportArgs.TimeoutTransport,
		t.read,
		ReadSize,
	)

	if err != nil {
		logging.LogError(t.FormatLogMessage("error", "timed out reading from transport"))
		return b, err
	}

	return b, nil
}

// ReadN read N bytes from the transport.
func (t *Standard) ReadN(n int) ([]byte, error) {
	b, err := transportTimeout(
		*t.BaseTransportArgs.TimeoutTransport,
		t.read,
		n,
	)

	if err != nil {
		logging.LogError(t.FormatLogMessage("error", "timed out reading from transport"))
		return b, err
	}

	return b, nil
}

// Write write bytes to the transport.
func (t *Standard) Write(channelInput []byte) error {
	_, err := t.writer.Write(channelInput)
	if err != nil {
		return err
	}

	return nil
}

// IsAlive indicate if the transport is alive or not.
func (t *Standard) IsAlive() bool {
	return t.session != nil
}

// FormatLogMessage formats log message payload, adding contextual info about the host.
func (t *Standard) FormatLogMessage(level, msg string) string {
	return logging.FormatLogMessage(level, t.BaseTransportArgs.Host, t.BaseTransportArgs.Port, msg)
}
<commit_msg>newline before if<commit_after>package transport

import (
	"encoding/base64"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"

	"github.com/scrapli/scrapligo/logging"

	"golang.org/x/crypto/ssh"
)

// Standard the "standard" (standard library) transport option for scrapligo.
type Standard struct {
	BaseTransportArgs     *BaseTransportArgs
	StandardTransportArgs *StandardTransportArgs
	client                *ssh.Client
	session               *ssh.Session
	writer                io.WriteCloser
	reader                io.Reader
}

// StandardTransportArgs struct representing attributes required for the Standard transport.
type StandardTransportArgs struct {
	AuthPassword      string
	AuthPrivateKey    string
	AuthStrictKey     bool
	SSHConfigFile     string
	SSHKnownHostsFile string
}

func keyString(k ssh.PublicKey) string {
	return k.Type() + " " + base64.StdEncoding.EncodeToString(
		k.Marshal(),
	) // e.g. "ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTY...."
}

// https://stackoverflow.com/questions/44269142/ \
// golang-ssh-getting-must-specify-hoskeycallback-error-despite-setting-it-to-n
// basically need to parse ssh config like scrapli does... at some point.
func trustedHostKeyCallback(trustedKey string) ssh.HostKeyCallback {
	if trustedKey == "" {
		return func(_ string, _ net.Addr, k ssh.PublicKey) error {
			log.Printf(
				"ssh key verification is *NOT* in effect: to fix, add this trustedKey: %q",
				keyString(k),
			)

			return nil
		}
	}

	return func(_ string, _ net.Addr, k ssh.PublicKey) error {
		ks := keyString(k)
		if trustedKey != ks {
			return ErrKeyVerificationFailed
		}

		return nil
	}
}

func (t *Standard) open(cfg *ssh.ClientConfig) error {
	var err error
	t.client, err = ssh.Dial(
		"tcp",
		fmt.Sprintf("%s:%d", t.BaseTransportArgs.Host, t.BaseTransportArgs.Port),
		cfg,
	)

	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error connecting to host: %v", err)),
		)

		return err
	}

	t.session, err = t.client.NewSession()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating session: %v", err)),
		)

		return err
	}

	t.writer, err = t.session.StdinPipe()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating writer: %v", err)),
		)

		return err
	}

	t.reader, err = t.session.StdoutPipe()
	if err != nil {
		logging.LogError(
			t.FormatLogMessage("error", fmt.Sprintf("error allocating reader: %v", err)),
		)

		return err
	}

	return nil
}

func (t *Standard) openBase() error {
	/* #nosec G106 */
	hostKeyCallback := ssh.InsecureIgnoreHostKey()
	if t.StandardTransportArgs.AuthStrictKey {
		// trustedKey will need to be gleaned from known hosts how scrapli does at some point
		hostKeyCallback = trustedHostKeyCallback("")
	}

	authMethods := make([]ssh.AuthMethod, 0)

	if t.StandardTransportArgs.AuthPrivateKey != "" {
		key, err := ioutil.ReadFile(t.StandardTransportArgs.AuthPrivateKey)
		if err != nil {
			return err
		}

		signer, err := ssh.ParsePrivateKey(key)

		if err != nil {
			logging.LogError(
				t.FormatLogMessage("error", fmt.Sprintf("unable to parse private key: %v", err)),
			)

			return err
		}

		authMethods = append(authMethods, ssh.PublicKeys(signer))
	}

	if t.StandardTransportArgs.AuthPassword != "" {
		authMethods = append(authMethods, ssh.Password(t.StandardTransportArgs.AuthPassword),
			ssh.KeyboardInteractive(
				func(user, instruction string, questions []string, echos []bool) ([]string, error) {
					answers := make([]string, len(questions))
					for i := range answers {
						answers[i] = t.StandardTransportArgs.AuthPassword
					}

					return answers, nil
				},
			))
	}

	cfg := &ssh.ClientConfig{
		User:            t.BaseTransportArgs.AuthUsername,
		Auth:            authMethods,
		Timeout:         *t.BaseTransportArgs.TimeoutSocket,
		HostKeyCallback: hostKeyCallback,
	}

	err := t.open(cfg)
	if err != nil {
		return err
	}

	// not sure what to do about the tty speeds... figured lets just go fast?
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 115200,
		ssh.TTY_OP_OSPEED: 115200,
	}

	err = t.session.RequestPty(
		"xterm",
		t.BaseTransportArgs.PtyHeight,
		t.BaseTransportArgs.PtyWidth,
		modes,
	)
	if err != nil {
		return err
	}

	return nil
}

// Open open a standard ssh connection.
func (t *Standard) Open() error {
	err := t.openBase()
	if err != nil {
		return err
	}

	err = t.session.Shell()
	if err != nil {
		return err
	}

	return nil
}

// OpenNetconf open a netconf connection.
func (t *Standard) OpenNetconf() error {
	err := t.openBase()
	if err != nil {
		logging.LogError(
			fmt.Sprintf(
				"failed opening base connection, cant attempt to open netconf connection; error: %v",
				err,
			),
		)

		return err
	}

	err = t.session.RequestSubsystem("netconf")
	if err != nil {
		logging.LogError(fmt.Sprintf("failed opening netconf subsystem; error: %v", err))
		return err
	}

	return nil
}

// Close close the transport connection to the device.
func (t *Standard) Close() error {
	err := t.session.Close()
	t.session = nil

	logging.LogDebug(t.FormatLogMessage("debug", "transport connection to host closed"))

	return err
}

func (t *Standard) read(n int) *transportResult {
	b := make([]byte, n)
	_, err := t.reader.Read(b)

	if err != nil {
		return &transportResult{
			result: nil,
			error:  ErrTransportFailure,
		}
	}

	return &transportResult{
		result: b,
		error:  nil,
	}
}

// Read read bytes from the transport.
func (t *Standard) Read() ([]byte, error) {
	b, err := transportTimeout(
		*t.BaseTransportArgs.TimeoutTransport,
		t.read,
		ReadSize,
	)

	if err != nil {
		logging.LogError(t.FormatLogMessage("error", "timed out reading from transport"))
		return b, err
	}

	return b, nil
}

// ReadN read N bytes from the transport.
func (t *Standard) ReadN(n int) ([]byte, error) {
	b, err := transportTimeout(
		*t.BaseTransportArgs.TimeoutTransport,
		t.read,
		n,
	)

	if err != nil {
		logging.LogError(t.FormatLogMessage("error", "timed out reading from transport"))
		return b, err
	}

	return b, nil
}

// Write write bytes to the transport.
func (t *Standard) Write(channelInput []byte) error {
	_, err := t.writer.Write(channelInput)
	if err != nil {
		return err
	}

	return nil
}

// IsAlive indicate if the transport is alive or not.
func (t *Standard) IsAlive() bool {
	return t.session != nil
}

// FormatLogMessage formats log message payload, adding contextual info about the host.
func (t *Standard) FormatLogMessage(level, msg string) string {
	return logging.FormatLogMessage(level, t.BaseTransportArgs.Host, t.BaseTransportArgs.Port, msg)
}
<|endoftext|>


<commit_before>package data

import (
	"github.com/tuna-timer/tuna-timer-api/models"
	"github.com/tuna-timer/tuna-timer-api/utils"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"testing"
	"time"
	"gopkg.in/tylerb/is.v1"
	"github.com/pavlo/gosuite"
)

func TestPassRepository(t *testing.T) {
	gosuite.Run(t, &PassRepositoryTestSuite{})
}

func (s *PassRepositoryTestSuite) GSTFindByToken(t *testing.T) {
	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    time.Now(),
		ExpiresAt:    time.Now().Add(5 * time.Minute),
		ClaimedAt:    nil,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.NotNil(p1Test)

	s.Equal(p1.ID, p1Test.ID)
}

func (s *PassRepositoryTestSuite) GSTFindByTokenDoesNotGetExpired(t *testing.T) {
	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    time.Now().Add(-10 * time.Minute),
		ExpiresAt:    time.Now().Add(-5 * time.Minute),
		ClaimedAt:    nil,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.Nil(p1Test)
}

func (s *PassRepositoryTestSuite) GSTFindByTokenDoesNotGetClaimed(t *testing.T) {
	now := time.Now()

	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    now,
		ExpiresAt:    now.Add(5 * time.Minute),
		ClaimedAt:    &now,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.Nil(p1Test)
}

func (s *PassRepositoryTestSuite) GSTFindActiveByUserID(t *testing.T) {

	now := time.Now()

	userID := bson.NewObjectId()
	s.userRepository.save(&models.TeamUser{
		ID: userID,
	})

	p1 := &models.Pass{ // a good one
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: userID.Hex(),
	}

	p2 := &models.Pass{ // already claimed
		ID:        bson.NewObjectId(),
		Token:     "p2token",
		CreatedAt: now,
		ExpiresAt: now.Add(5 * time.Minute),
		ClaimedAt: &now,
	}

	p3 := &models.Pass{ // belongs to another user
		ID:         bson.NewObjectId(),
		Token:      "p3token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  &now,
		TeamUserID: "another-user",
	}

	s.repository.insert(p1)
	s.repository.insert(p2)
	s.repository.insert(p3)

	pass, err := s.repository.FindActiveByUserID(userID.Hex())
	s.Nil(err)
	s.NotNil(pass)
	s.Equal("p1token", pass.Token)
}

func (s *PassRepositoryTestSuite) GSTRemoveExpiredPasses(t *testing.T) {

	now := time.Now()

	p1 := &models.Pass{ //should be removed as its expiresAt is in the past
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now.Add(-5 * time.Minute),
		ExpiresAt:  now.Add(-3 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	p2 := &models.Pass{ //should NOT be removed as its expiresAt is in the future
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	claimedAt := now.Add(2 * time.Minute)
	p3 := &models.Pass{ //should NOT be removed as it is claimed
		ID:         bson.NewObjectId(),
		Token:      "p3token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  &claimedAt,
		TeamUserID: "user-id",
	}

	err := s.repository.insert(p1)
	s.Nil(err)

	err = s.repository.insert(p2)
	s.Nil(err)

	err = s.repository.insert(p3)
	s.Nil(err)

	err = s.repository.removeExpiredPasses()
	s.Nil(err)

	p1, err = s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.Nil(p1)

	p2, err = s.repository.findByID(p2.ID.Hex())
	s.Nil(err)
	s.NotNil(p2)

	p3, err = s.repository.findByID(p3.ID.Hex())
	s.Nil(err)
	s.NotNil(p3)
}

func (s *PassRepositoryTestSuite) GSTFindByID(t *testing.T) {
	p1 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  time.Now(),
		ExpiresAt:  time.Now(),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	p2 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  time.Now(),
		ExpiresAt:  time.Now(),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	err := s.repository.insert(p1)
	s.Nil(err)

	err = s.repository.insert(p2)
	s.Nil(err)

	p, err := s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.NotNil(p)
	s.Equal("p1token", p.Token)

	p, err = s.repository.findByID(bson.NewObjectId().Hex())
	s.Nil(err)
	s.Nil(p)
}

func (s *PassRepositoryTestSuite) GSTRemovePassesClaimedBefore(t *testing.T) {

	now := time.Now()

	fiveMinutesInPast := now.Add(-5 * time.Minute)
	p1 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now,
		ExpiresAt:  now,
		ClaimedAt:  &fiveMinutesInPast,
		TeamUserID: "user-id",
	}

	fiveMinutesInFuture := now.Add(5 * time.Minute)
	p2 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  now,
		ExpiresAt:  now,
		ClaimedAt:  &fiveMinutesInFuture,
		TeamUserID: "user-id",
	}

	s.repository.insert(p1)
	s.repository.insert(p2)

	err := s.repository.removePassesClaimedBefore(time.Now())
	s.Nil(err)

	p1, err = s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.Nil(p1)

	p2, err = s.repository.findByID(p2.ID.Hex())
	s.Nil(err)
	s.NotNil(p2)
}

type PassRepositoryTestSuite struct {
	*is.Is
	env            *utils.Environment
	session        *mgo.Session
	repository     *PassRepository
	userRepository *UserRepository
}

func (s *PassRepositoryTestSuite) SetUpSuite(t *testing.T) {
	e := utils.NewEnvironment(utils.TestEnv, "1.0.0")

	session, err := utils.ConnectToDatabase(e.Config)
	if err != nil {
		log.Fatal("Failed to connect to DB!")
	}

	e.MigrateDatabase(session)

	s.env = e
	s.session = session.Clone()
	s.repository = NewPassRepository(s.session)
	s.userRepository = NewUserRepository(session)
	s.Is = is.New(t)
}

func (s *PassRepositoryTestSuite) TearDownSuite() {
	s.session.Close()
}

func (s *PassRepositoryTestSuite) SetUp() {
	utils.TruncateTables(s.session)
}

func (s *PassRepositoryTestSuite) TearDown() {}
<commit_msg>let's break a test to see if SemaphoreCI figures it out<commit_after>package data

import (
	"github.com/tuna-timer/tuna-timer-api/models"
	"github.com/tuna-timer/tuna-timer-api/utils"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"testing"
	"time"
	"gopkg.in/tylerb/is.v1"
	"github.com/pavlo/gosuite"
)

func TestPassRepository(t *testing.T) {
	gosuite.Run(t, &PassRepositoryTestSuite{})
}

func (s *PassRepositoryTestSuite) GSTFindByToken(t *testing.T) {
	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    time.Now(),
		ExpiresAt:    time.Now().Add(5 * time.Minute),
		ClaimedAt:    nil,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.NotNil(p1Test)

	s.Equal(p1.ID, p1Test.ID)
}

func (s *PassRepositoryTestSuite) GSTFindByTokenDoesNotGetExpired(t *testing.T) {
	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    time.Now().Add(-10 * time.Minute),
		ExpiresAt:    time.Now().Add(-5 * time.Minute),
		ClaimedAt:    nil,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.Nil(p1Test)
}

func (s *PassRepositoryTestSuite) GSTFindByTokenDoesNotGetClaimed(t *testing.T) {
	now := time.Now()

	p1 := &models.Pass{
		ID:           bson.NewObjectId(),
		Token:        "token",
		CreatedAt:    now,
		ExpiresAt:    now.Add(5 * time.Minute),
		ClaimedAt:    &now,
		ModelVersion: models.ModelVersionPass,
	}
	err := s.repository.insert(p1)
	s.Nil(err)

	p1Test, err := s.repository.FindActivePassByToken("token")
	s.Nil(err)
	s.Nil(p1Test)
}

func (s *PassRepositoryTestSuite) GSTFindActiveByUserID(t *testing.T) {

	now := time.Now()

	userID := bson.NewObjectId()
	s.userRepository.save(&models.TeamUser{
		ID: userID,
	})

	p1 := &models.Pass{ // a good one
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: userID.Hex(),
	}

	p2 := &models.Pass{ // already claimed
		ID:        bson.NewObjectId(),
		Token:     "p2token",
		CreatedAt: now,
		ExpiresAt: now.Add(5 * time.Minute),
		ClaimedAt: &now,
	}

	p3 := &models.Pass{ // belongs to another user
		ID:         bson.NewObjectId(),
		Token:      "p3token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  &now,
		TeamUserID: "another-user",
	}

	s.repository.insert(p1)
	s.repository.insert(p2)
	s.repository.insert(p3)

	pass, err := s.repository.FindActiveByUserID(userID.Hex())
	s.Nil(err)
	s.NotNil(pass)
	s.Equal("p1token", pass.Token)
}

func (s *PassRepositoryTestSuite) GSTRemoveExpiredPasses(t *testing.T) {

	now := time.Now()

	p1 := &models.Pass{ //should be removed as its expiresAt is in the past
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now.Add(-5 * time.Minute),
		ExpiresAt:  now.Add(-3 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	p2 := &models.Pass{ //should NOT be removed as its expiresAt is in the future
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	claimedAt := now.Add(2 * time.Minute)
	p3 := &models.Pass{ //should NOT be removed as it is claimed
		ID:         bson.NewObjectId(),
		Token:      "p3token",
		CreatedAt:  now,
		ExpiresAt:  now.Add(5 * time.Minute),
		ClaimedAt:  &claimedAt,
		TeamUserID: "user-id",
	}

	err := s.repository.insert(p1)
	s.Nil(err)

	err = s.repository.insert(p2)
	s.Nil(err)

	err = s.repository.insert(p3)
	s.Nil(err)

	err = s.repository.removeExpiredPasses()
	s.Nil(err)

	p1, err = s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.Nil(p1)

	p2, err = s.repository.findByID(p2.ID.Hex())
	s.Nil(err)
	s.NotNil(p2)

	p3, err = s.repository.findByID(p3.ID.Hex())
	s.Nil(err)
	s.NotNil(p3)
}

func (s *PassRepositoryTestSuite) GSTFindByID(t *testing.T) {
	p1 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  time.Now(),
		ExpiresAt:  time.Now(),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	p2 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  time.Now(),
		ExpiresAt:  time.Now(),
		ClaimedAt:  nil,
		TeamUserID: "user-id",
	}

	err := s.repository.insert(p1)
	s.Nil(err)

	err = s.repository.insert(p2)
	s.Nil(err)

	p, err := s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.NotNil(p)
	s.Equal("--p1token", p.Token)

	p, err = s.repository.findByID(bson.NewObjectId().Hex())
	s.Nil(err)
	s.Nil(p)
}

func (s *PassRepositoryTestSuite) GSTRemovePassesClaimedBefore(t *testing.T) {

	now := time.Now()

	fiveMinutesInPast := now.Add(-5 * time.Minute)
	p1 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p1token",
		CreatedAt:  now,
		ExpiresAt:  now,
		ClaimedAt:  &fiveMinutesInPast,
		TeamUserID: "user-id",
	}

	fiveMinutesInFuture := now.Add(5 * time.Minute)
	p2 := &models.Pass{
		ID:         bson.NewObjectId(),
		Token:      "p2token",
		CreatedAt:  now,
		ExpiresAt:  now,
		ClaimedAt:  &fiveMinutesInFuture,
		TeamUserID: "user-id",
	}

	s.repository.insert(p1)
	s.repository.insert(p2)

	err := s.repository.removePassesClaimedBefore(time.Now())
	s.Nil(err)

	p1, err = s.repository.findByID(p1.ID.Hex())
	s.Nil(err)
	s.Nil(p1)

	p2, err = s.repository.findByID(p2.ID.Hex())
	s.Nil(err)
	s.NotNil(p2)
}

type PassRepositoryTestSuite struct {
	*is.Is
	env            *utils.Environment
	session        *mgo.Session
	repository     *PassRepository
	userRepository *UserRepository
}

func (s *PassRepositoryTestSuite) SetUpSuite(t *testing.T) {
	e := utils.NewEnvironment(utils.TestEnv, "1.0.0")

	session, err := utils.ConnectToDatabase(e.Config)
	if err != nil {
		log.Fatal("Failed to connect to DB!")
	}

	e.MigrateDatabase(session)

	s.env = e
	s.session = session.Clone()
	s.repository = NewPassRepository(s.session)
	s.userRepository = NewUserRepository(session)
	s.Is = is.New(t)
}

func (s *PassRepositoryTestSuite) TearDownSuite() {
	s.session.Close()
}

func (s *PassRepositoryTestSuite) SetUp() {
	utils.TruncateTables(s.session)
}

func (s *PassRepositoryTestSuite) TearDown() {}
<|endoftext|>


<commit_before>package aws

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/mitchellh/go-homedir"

	"errors"

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

func resourceAwsLambdaFunction() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLambdaFunctionCreate,
		Read:   resourceAwsLambdaFunctionRead,
		Update: resourceAwsLambdaFunctionUpdate,
		Delete: resourceAwsLambdaFunctionDelete,

		Schema: map[string]*schema.Schema{
			"filename": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"s3_bucket", "s3_key", "s3_object_version"},
			},
			"s3_bucket": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_key": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_object_version": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"function_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"handler": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"memory_size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  128,
			},
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"runtime": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "nodejs",
			},
			"timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  3,
			},
			"vpc_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"subnet_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
						"security_group_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
					},
				},
			},
			"update_code": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"last_modified": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"source_code_hash": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"remote_code_hash": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

// resourceAwsLambdaFunction maps to:
// CreateFunction in the API / SDK
func resourceAwsLambdaFunctionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	functionName := d.Get("function_name").(string)
	iamRole := d.Get("role").(string)

	log.Printf("[DEBUG] Creating Lambda Function %s with role %s", functionName, iamRole)

	var functionCode *lambda.FunctionCode
	if v, ok := d.GetOk("filename"); ok {
		zipfile, shaSum, err := loadLocalZipFile(v.(string))
		if err != nil {
			return err
		}
		d.Set("source_code_hash", shaSum)
		functionCode = &lambda.FunctionCode{
			ZipFile: zipfile,
		}
	} else {
		s3Bucket, bucketOk := d.GetOk("s3_bucket")
		s3Key, keyOk := d.GetOk("s3_key")
		s3ObjectVersion, versionOk := d.GetOk("s3_object_version")
		if !bucketOk || !keyOk {
			return errors.New("s3_bucket and s3_key must all be set while using S3 code source")
		}
		functionCode = &lambda.FunctionCode{
			S3Bucket: aws.String(s3Bucket.(string)),
			S3Key:    aws.String(s3Key.(string)),
		}
		if versionOk {
			functionCode.S3ObjectVersion = aws.String(s3ObjectVersion.(string))
		}
	}

	params := &lambda.CreateFunctionInput{
		Code:         functionCode,
		Description:  aws.String(d.Get("description").(string)),
		FunctionName: aws.String(functionName),
		Handler:      aws.String(d.Get("handler").(string)),
		MemorySize:   aws.Int64(int64(d.Get("memory_size").(int))),
		Role:         aws.String(iamRole),
		Runtime:      aws.String(d.Get("runtime").(string)),
		Timeout:      aws.Int64(int64(d.Get("timeout").(int))),
	}

	if v, ok := d.GetOk("vpc_config"); ok {
		config, err := validateVPCConfig(v)
		if err != nil {
			return err
		}

		var subnetIds []*string
		for _, id := range config["subnet_ids"].(*schema.Set).List() {
			subnetIds = append(subnetIds, aws.String(id.(string)))
		}

		var securityGroupIds []*string
		for _, id := range config["security_group_ids"].(*schema.Set).List() {
			securityGroupIds = append(securityGroupIds, aws.String(id.(string)))
		}

		params.VpcConfig = &lambda.VpcConfig{
			SubnetIds:        subnetIds,
			SecurityGroupIds: securityGroupIds,
		}
	}

	// IAM profiles can take ~10 seconds to propagate in AWS:
	// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#launch-instance-with-role-console
	// Error creating Lambda function: InvalidParameterValueException: The role defined for the task cannot be assumed by Lambda.
	err := resource.Retry(1*time.Minute, func() *resource.RetryError {
		_, err := conn.CreateFunction(params)
		if err != nil {
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "InvalidParameterValueException" {
					log.Printf("[DEBUG] InvalidParameterValueException creating Lambda Function: %s", awserr)
					return resource.RetryableError(awserr)
				}
			}
			log.Printf("[DEBUG] Error creating Lambda Function: %s", err)
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Error creating Lambda function: %s", err)
	}

	d.SetId(d.Get("function_name").(string))

	return resourceAwsLambdaFunctionRead(d, meta)
}

// resourceAwsLambdaFunctionRead maps to:
// GetFunction in the API / SDK
func resourceAwsLambdaFunctionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[DEBUG] Fetching Lambda Function: %s", d.Id())

	params := &lambda.GetFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	getFunctionOutput, err := conn.GetFunction(params)
	if err != nil {
		return err
	}

	// getFunctionOutput.Code.Location is a pre-signed URL pointing at the zip
	// file that we uploaded when we created the resource. You can use it to
	// download the code from AWS. The other part is
	// getFunctionOutput.Configuration which holds metadata.

	function := getFunctionOutput.Configuration
	// TODO error checking / handling on the Set() calls.
	d.Set("arn", function.FunctionArn)
	d.Set("description", function.Description)
	d.Set("handler", function.Handler)
	d.Set("memory_size", function.MemorySize)
	d.Set("last_modified", function.LastModified)
	d.Set("role", function.Role)
	d.Set("runtime", function.Runtime)
	d.Set("timeout", function.Timeout)
	if config := flattenLambdaVpcConfigResponse(function.VpcConfig); len(config) > 0 {
		d.Set("vpc_config", config)
	}

	// Compare code hashes, and see if an update is required to code. If there
	// is, set the "update_code" attribute.

	remoteSum, err := decodeBase64(*function.CodeSha256)
	if err != nil {
		return err
	}
	_, localSum, err := loadLocalZipFile(d.Get("filename").(string))
	if err != nil {
		return err
	}
	d.Set("remote_code_hash", remoteSum)
	d.Set("source_code_hash", localSum)

	if remoteSum != localSum {
		d.Set("update_code", true)
	} else {
		d.Set("update_code", false)
	}

	return nil
}

// resourceAwsLambdaFunction maps to:
// DeleteFunction in the API / SDK
func resourceAwsLambdaFunctionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[INFO] Deleting Lambda Function: %s", d.Id())

	params := &lambda.DeleteFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	_, err := conn.DeleteFunction(params)
	if err != nil {
		return fmt.Errorf("Error deleting Lambda Function: %s", err)
	}

	d.SetId("")

	return nil
}

// resourceAwsLambdaFunctionUpdate maps to:
// UpdateFunctionCode in the API / SDK
func resourceAwsLambdaFunctionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	d.Partial(true)

	codeReq := &lambda.UpdateFunctionCodeInput{
		FunctionName: aws.String(d.Id()),
	}

	codeUpdate := false
	if sourceHash, ok := d.GetOk("source_code_hash"); ok {
		zipfile, shaSum, err := loadLocalZipFile(d.Get("filename").(string))
		if err != nil {
			return err
		}
		if sourceHash != shaSum {
			d.SetPartial("filename")
			d.SetPartial("source_code_hash")
		}
		codeReq.ZipFile = zipfile
		codeUpdate = true
	}
	if d.HasChange("s3_bucket") || d.HasChange("s3_key") || d.HasChange("s3_object_version") {
		codeReq.S3Bucket = aws.String(d.Get("s3_bucket").(string))
		codeReq.S3Key = aws.String(d.Get("s3_key").(string))
		codeReq.S3ObjectVersion = aws.String(d.Get("s3_object_version").(string))
		codeUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Code request: %#v", codeReq)
	if codeUpdate {
		_, err := conn.UpdateFunctionCode(codeReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Code %s: %s", d.Id(), err)
		}

		d.SetPartial("filename")
		d.SetPartial("source_code_hash")
		d.SetPartial("s3_bucket")
		d.SetPartial("s3_key")
		d.SetPartial("s3_object_version")
	}

	configReq := &lambda.UpdateFunctionConfigurationInput{
		FunctionName: aws.String(d.Id()),
	}

	configUpdate := false
	if d.HasChange("description") {
		configReq.Description = aws.String(d.Get("description").(string))
		configUpdate = true
	}
	if d.HasChange("handler") {
		configReq.Handler = aws.String(d.Get("handler").(string))
		configUpdate = true
	}
	if d.HasChange("memory_size") {
		configReq.MemorySize = aws.Int64(int64(d.Get("memory_size").(int)))
		configUpdate = true
	}
	if d.HasChange("role") {
		configReq.Role = aws.String(d.Get("role").(string))
		configUpdate = true
	}
	if d.HasChange("timeout") {
		configReq.Timeout = aws.Int64(int64(d.Get("timeout").(int)))
		configUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Configuration request: %#v", configReq)
	if configUpdate {
		_, err := conn.UpdateFunctionConfiguration(configReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Configuration %s: %s", d.Id(), err)
		}
		d.SetPartial("description")
		d.SetPartial("handler")
		d.SetPartial("memory_size")
		d.SetPartial("role")
		d.SetPartial("timeout")
	}
	d.Partial(false)

	return resourceAwsLambdaFunctionRead(d, meta)
}

// loads the local ZIP data and the SHA sum of the data.
func loadLocalZipFile(v string) ([]byte, string, error) {
	filename, err := homedir.Expand(v)
	if err != nil {
		return nil, "", err
	}
	zipfile, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, "", err
	}
	sum := sha256.Sum256(zipfile)
	return zipfile, fmt.Sprintf("%x", sum), nil
}

// Decodes a base64 string to a string.
func decodeBase64(s string) (string, error) {
	sum, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%x", sum), nil
}

func validateVPCConfig(v interface{}) (map[string]interface{}, error) {
	configs := v.([]interface{})
	if len(configs) > 1 {
		return nil, errors.New("Only a single vpc_config block is expected")
	}

	config, ok := configs[0].(map[string]interface{})

	if !ok {
		return nil, errors.New("vpc_config is <nil>")
	}

	if config["subnet_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.subnet_ids cannot be empty")
	}

	if config["security_group_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.security_group_ids cannot be empty")
	}

	return config, nil
}
<commit_msg>provider/aws: Simplify update logic for Lambda function<commit_after>package aws

import (
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/mitchellh/go-homedir"

	"errors"

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

func resourceAwsLambdaFunction() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLambdaFunctionCreate,
		Read:   resourceAwsLambdaFunctionRead,
		Update: resourceAwsLambdaFunctionUpdate,
		Delete: resourceAwsLambdaFunctionDelete,

		Schema: map[string]*schema.Schema{
			"filename": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"s3_bucket", "s3_key", "s3_object_version"},
			},
			"s3_bucket": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_key": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"s3_object_version": &schema.Schema{
				Type:          schema.TypeString,
				Optional:      true,
				ConflictsWith: []string{"filename"},
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"function_name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},
			"handler": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"memory_size": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  128,
			},
			"role": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"runtime": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
				Default:  "nodejs",
			},
			"timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  3,
			},
			"vpc_config": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				ForceNew: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"subnet_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
						"security_group_ids": &schema.Schema{
							Type:     schema.TypeSet,
							Required: true,
							ForceNew: true,
							Elem:     &schema.Schema{Type: schema.TypeString},
							Set:      schema.HashString,
						},
					},
				},
			},
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"last_modified": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"source_code_hash": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
		},
	}
}

// resourceAwsLambdaFunction maps to:
// CreateFunction in the API / SDK
func resourceAwsLambdaFunctionCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	functionName := d.Get("function_name").(string)
	iamRole := d.Get("role").(string)

	log.Printf("[DEBUG] Creating Lambda Function %s with role %s", functionName, iamRole)

	var functionCode *lambda.FunctionCode
	if v, ok := d.GetOk("filename"); ok {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		functionCode = &lambda.FunctionCode{
			ZipFile: file,
		}
	} else {
		s3Bucket, bucketOk := d.GetOk("s3_bucket")
		s3Key, keyOk := d.GetOk("s3_key")
		s3ObjectVersion, versionOk := d.GetOk("s3_object_version")
		if !bucketOk || !keyOk {
			return errors.New("s3_bucket and s3_key must all be set while using S3 code source")
		}
		functionCode = &lambda.FunctionCode{
			S3Bucket: aws.String(s3Bucket.(string)),
			S3Key:    aws.String(s3Key.(string)),
		}
		if versionOk {
			functionCode.S3ObjectVersion = aws.String(s3ObjectVersion.(string))
		}
	}

	params := &lambda.CreateFunctionInput{
		Code:         functionCode,
		Description:  aws.String(d.Get("description").(string)),
		FunctionName: aws.String(functionName),
		Handler:      aws.String(d.Get("handler").(string)),
		MemorySize:   aws.Int64(int64(d.Get("memory_size").(int))),
		Role:         aws.String(iamRole),
		Runtime:      aws.String(d.Get("runtime").(string)),
		Timeout:      aws.Int64(int64(d.Get("timeout").(int))),
	}

	if v, ok := d.GetOk("vpc_config"); ok {
		config, err := validateVPCConfig(v)
		if err != nil {
			return err
		}

		var subnetIds []*string
		for _, id := range config["subnet_ids"].(*schema.Set).List() {
			subnetIds = append(subnetIds, aws.String(id.(string)))
		}

		var securityGroupIds []*string
		for _, id := range config["security_group_ids"].(*schema.Set).List() {
			securityGroupIds = append(securityGroupIds, aws.String(id.(string)))
		}

		params.VpcConfig = &lambda.VpcConfig{
			SubnetIds:        subnetIds,
			SecurityGroupIds: securityGroupIds,
		}
	}

	// IAM profiles can take ~10 seconds to propagate in AWS:
	// http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#launch-instance-with-role-console
	// Error creating Lambda function: InvalidParameterValueException: The role defined for the task cannot be assumed by Lambda.
	err := resource.Retry(1*time.Minute, func() *resource.RetryError {
		_, err := conn.CreateFunction(params)
		if err != nil {
			if awserr, ok := err.(awserr.Error); ok {
				if awserr.Code() == "InvalidParameterValueException" {
					log.Printf("[DEBUG] InvalidParameterValueException creating Lambda Function: %s", awserr)
					return resource.RetryableError(awserr)
				}
			}
			log.Printf("[DEBUG] Error creating Lambda Function: %s", err)
			return resource.NonRetryableError(err)
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("Error creating Lambda function: %s", err)
	}

	d.SetId(d.Get("function_name").(string))

	return resourceAwsLambdaFunctionRead(d, meta)
}

// resourceAwsLambdaFunctionRead maps to:
// GetFunction in the API / SDK
func resourceAwsLambdaFunctionRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[DEBUG] Fetching Lambda Function: %s", d.Id())

	params := &lambda.GetFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	getFunctionOutput, err := conn.GetFunction(params)
	if err != nil {
		return err
	}

	// getFunctionOutput.Code.Location is a pre-signed URL pointing at the zip
	// file that we uploaded when we created the resource. You can use it to
	// download the code from AWS. The other part is
	// getFunctionOutput.Configuration which holds metadata.

	function := getFunctionOutput.Configuration
	// TODO error checking / handling on the Set() calls.
	d.Set("arn", function.FunctionArn)
	d.Set("description", function.Description)
	d.Set("handler", function.Handler)
	d.Set("memory_size", function.MemorySize)
	d.Set("last_modified", function.LastModified)
	d.Set("role", function.Role)
	d.Set("runtime", function.Runtime)
	d.Set("timeout", function.Timeout)
	if config := flattenLambdaVpcConfigResponse(function.VpcConfig); len(config) > 0 {
		d.Set("vpc_config", config)
	}
	d.Set("source_code_hash", function.CodeSha256)

	return nil
}

// resourceAwsLambdaFunction maps to:
// DeleteFunction in the API / SDK
func resourceAwsLambdaFunctionDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	log.Printf("[INFO] Deleting Lambda Function: %s", d.Id())

	params := &lambda.DeleteFunctionInput{
		FunctionName: aws.String(d.Get("function_name").(string)),
	}

	_, err := conn.DeleteFunction(params)
	if err != nil {
		return fmt.Errorf("Error deleting Lambda Function: %s", err)
	}

	d.SetId("")

	return nil
}

// resourceAwsLambdaFunctionUpdate maps to:
// UpdateFunctionCode in the API / SDK
func resourceAwsLambdaFunctionUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lambdaconn

	d.Partial(true)

	codeReq := &lambda.UpdateFunctionCodeInput{
		FunctionName: aws.String(d.Id()),
	}

	codeUpdate := false
	if v, ok := d.GetOk("filename"); ok && d.HasChange("source_code_hash") {
		file, err := loadFileContent(v.(string))
		if err != nil {
			return fmt.Errorf("Unable to load %q: %s", v.(string), err)
		}
		codeReq.ZipFile = file
		codeUpdate = true
	}
	if d.HasChange("s3_bucket") || d.HasChange("s3_key") || d.HasChange("s3_object_version") {
		codeReq.S3Bucket = aws.String(d.Get("s3_bucket").(string))
		codeReq.S3Key = aws.String(d.Get("s3_key").(string))
		codeReq.S3ObjectVersion = aws.String(d.Get("s3_object_version").(string))
		codeUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Code request: %#v", codeReq)
	if codeUpdate {
		_, err := conn.UpdateFunctionCode(codeReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Code %s: %s", d.Id(), err)
		}

		d.SetPartial("filename")
		d.SetPartial("source_code_hash")
		d.SetPartial("s3_bucket")
		d.SetPartial("s3_key")
		d.SetPartial("s3_object_version")
	}

	configReq := &lambda.UpdateFunctionConfigurationInput{
		FunctionName: aws.String(d.Id()),
	}

	configUpdate := false
	if d.HasChange("description") {
		configReq.Description = aws.String(d.Get("description").(string))
		configUpdate = true
	}
	if d.HasChange("handler") {
		configReq.Handler = aws.String(d.Get("handler").(string))
		configUpdate = true
	}
	if d.HasChange("memory_size") {
		configReq.MemorySize = aws.Int64(int64(d.Get("memory_size").(int)))
		configUpdate = true
	}
	if d.HasChange("role") {
		configReq.Role = aws.String(d.Get("role").(string))
		configUpdate = true
	}
	if d.HasChange("timeout") {
		configReq.Timeout = aws.Int64(int64(d.Get("timeout").(int)))
		configUpdate = true
	}

	log.Printf("[DEBUG] Send Update Lambda Function Configuration request: %#v", configReq)
	if configUpdate {
		_, err := conn.UpdateFunctionConfiguration(configReq)
		if err != nil {
			return fmt.Errorf("Error modifying Lambda Function Configuration %s: %s", d.Id(), err)
		}
		d.SetPartial("description")
		d.SetPartial("handler")
		d.SetPartial("memory_size")
		d.SetPartial("role")
		d.SetPartial("timeout")
	}
	d.Partial(false)

	return resourceAwsLambdaFunctionRead(d, meta)
}

// loadFileContent returns contents of a file in a given path
func loadFileContent(v string) ([]byte, error) {
	filename, err := homedir.Expand(v)
	if err != nil {
		return nil, err
	}
	fileContent, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return fileContent, nil
}

func validateVPCConfig(v interface{}) (map[string]interface{}, error) {
	configs := v.([]interface{})
	if len(configs) > 1 {
		return nil, errors.New("Only a single vpc_config block is expected")
	}

	config, ok := configs[0].(map[string]interface{})

	if !ok {
		return nil, errors.New("vpc_config is <nil>")
	}

	if config["subnet_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.subnet_ids cannot be empty")
	}

	if config["security_group_ids"].(*schema.Set).Len() == 0 {
		return nil, errors.New("vpc_config.security_group_ids cannot be empty")
	}

	return config, nil
}
<|endoftext|>


<commit_before>package server_test

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
	"github.com/wanelo/image-server/core"
	fetcher "github.com/wanelo/image-server/fetcher/http"
	"github.com/wanelo/image-server/logger"
	"github.com/wanelo/image-server/logger/graphite"
	"github.com/wanelo/image-server/paths"
	"github.com/wanelo/image-server/server"
	"github.com/wanelo/image-server/uploader"

	. "github.com/wanelo/image-server/test"
)

func TestNewImageHandlerWithS3(t *testing.T) {
	deleteS3TestDirectory()

	sc := buildTestServerConfiguration()
	uploader.Initialize(sc)

	router := server.NewRouter(sc)

	request, _ := http.NewRequest("POST", "/test_namespace?outputs=x300.jpg,x300.webp&source=http%3A%2F%2Fcdn-s3-3.wanelo.com%2Fproduct%2Fimage%2F15209365%2Fx354.jpg", nil)
	response := httptest.NewRecorder()
	log.Println(sc)

	router.ServeHTTP(response, request)

	url := "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/info.json"
	resp, err := http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "application/json", resp.Header.Get("Content-Type"))

	url = "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.jpg"
	resp, err = http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "image/jpeg", resp.Header.Get("Content-Type"))

	url = "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.webp"
	resp, err = http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "image/webp", resp.Header.Get("Content-Type"))
}

func deleteS3TestDirectory() {
	auth := aws.Auth{
		AccessKey: os.Getenv("AWS_ACCESS_KEY_ID"),
		SecretKey: os.Getenv("AWS_SECRET_KEY"),
	}
	client := s3.New(auth, aws.USEast)
	bucket := client.Bucket(os.Getenv("AWS_BUCKET"))
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/info.json")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/original")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.jpg")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.webp")
}

func buildTestServerConfiguration() *core.ServerConfiguration {
	sc := &core.ServerConfiguration{
		LocalBasePath:  "../public",
		RemoteBasePath: "test",
		DefaultQuality: 90,
		AWSAccessKeyID: os.Getenv("AWS_ACCESS_KEY_ID"),
		AWSSecretKey:   os.Getenv("AWS_SECRET_KEY"),
		AWSBucket:      os.Getenv("AWS_BUCKET"),
	}

	loggers := []core.Logger{
		graphite.New(sc.GraphiteHost, sc.GraphitePort),
	}

	adapters := &core.Adapters{
		Fetcher: &fetcher.Fetcher{},
		Paths:   &paths.Paths{LocalBasePath: sc.LocalBasePath, RemoteBasePath: sc.RemoteBasePath, RemoteBaseURL: sc.RemoteBaseURL},
		Logger:  &logger.Logger{Loggers: loggers},
	}
	sc.Adapters = adapters
	return sc
}
<commit_msg>use same version of goamz in tests<commit_after>package server_test

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/s3"
	"github.com/wanelo/image-server/core"
	fetcher "github.com/wanelo/image-server/fetcher/http"
	"github.com/wanelo/image-server/logger"
	"github.com/wanelo/image-server/logger/graphite"
	"github.com/wanelo/image-server/paths"
	"github.com/wanelo/image-server/server"
	"github.com/wanelo/image-server/uploader"

	. "github.com/wanelo/image-server/test"
)

func TestNewImageHandlerWithS3(t *testing.T) {
	deleteS3TestDirectory()

	sc := buildTestServerConfiguration()
	uploader.Initialize(sc)

	router := server.NewRouter(sc)

	request, _ := http.NewRequest("POST", "/test_namespace?outputs=x300.jpg,x300.webp&source=http%3A%2F%2Fcdn-s3-3.wanelo.com%2Fproduct%2Fimage%2F15209365%2Fx354.jpg", nil)
	response := httptest.NewRecorder()
	log.Println(sc)

	router.ServeHTTP(response, request)

	url := "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/info.json"
	resp, err := http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "application/json", resp.Header.Get("Content-Type"))

	url = "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.jpg"
	resp, err = http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "image/jpeg", resp.Header.Get("Content-Type"))

	url = "https://s3.amazonaws.com/wanelo-dev/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.webp"
	resp, err = http.Head(url)
	Ok(t, err)
	Equals(t, "200 OK", resp.Status)
	Equals(t, "image/webp", resp.Header.Get("Content-Type"))
}

func deleteS3TestDirectory() {
	auth := aws.Auth{
		AccessKey: os.Getenv("AWS_ACCESS_KEY_ID"),
		SecretKey: os.Getenv("AWS_SECRET_KEY"),
	}
	client := s3.New(auth, aws.USEast)
	bucket := client.Bucket(os.Getenv("AWS_BUCKET"))
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/info.json")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/original")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.jpg")
	bucket.Del("/test/test_namespace/6da/b5f/6d8/d4bddc73fdff34d4f0507f7/x300.webp")
}

func buildTestServerConfiguration() *core.ServerConfiguration {
	sc := &core.ServerConfiguration{
		LocalBasePath:  "../public",
		RemoteBasePath: "test",
		DefaultQuality: 90,
		AWSAccessKeyID: os.Getenv("AWS_ACCESS_KEY_ID"),
		AWSSecretKey:   os.Getenv("AWS_SECRET_KEY"),
		AWSBucket:      os.Getenv("AWS_BUCKET"),
	}

	loggers := []core.Logger{
		graphite.New(sc.GraphiteHost, sc.GraphitePort),
	}

	adapters := &core.Adapters{
		Fetcher: &fetcher.Fetcher{},
		Paths:   &paths.Paths{LocalBasePath: sc.LocalBasePath, RemoteBasePath: sc.RemoteBasePath, RemoteBaseURL: sc.RemoteBaseURL},
		Logger:  &logger.Logger{Loggers: loggers},
	}
	sc.Adapters = adapters
	return sc
}
<|endoftext|>


<commit_before>package tsacmd

import (
	"bytes"
	"context"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"io/ioutil"

	yaml "gopkg.in/yaml.v2"

	"os/signal"
	"syscall"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/skymarshal/token"
	"github.com/concourse/concourse/tsa"
	"github.com/concourse/flag"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"golang.org/x/crypto/ssh"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/clientcredentials"
)

type TSACommand struct {
	Logger flag.Lager

	BindIP      flag.IP `long:"bind-ip"   default:"0.0.0.0" description:"IP address on which to listen for SSH."`
	PeerAddress string  `long:"peer-address" default:"127.0.0.1" description:"Network address of this web node, reachable by other web nodes. Used for forwarded worker addresses."`
	BindPort    uint16  `long:"bind-port" default:"2222"    description:"Port on which to listen for SSH."`

	DebugBindIP   flag.IP `long:"debug-bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the pprof debugger endpoints."`
	DebugBindPort uint16  `long:"debug-bind-port" default:"2221"      description:"Port on which to listen for the pprof debugger endpoints."`

	HostKey                *flag.PrivateKey               `long:"host-key"        required:"true" description:"Path to private key to use for the SSH server."`
	AuthorizedKeys         flag.AuthorizedKeys            `long:"authorized-keys" description:"Path to file containing keys to authorize, in SSH authorized_keys format (one public key per line)."`
	TeamAuthorizedKeys     map[string]flag.AuthorizedKeys `long:"team-authorized-keys" value-name:"NAME:PATH" description:"Path to file containing keys to authorize, in SSH authorized_keys format (one public key per line)."`
	TeamAuthorizedKeysFile flag.File                      `long:"team-authorized-keys-file" description:"Path to file containing a YAML array of teams and their authorized SSH keys, e.g. [{team:foo,ssh_keys:[key1,key2]}]."`

	ATCURLs []flag.URL `long:"atc-url" required:"true" description:"ATC API endpoints to which workers will be registered."`

	ClientID     string   `long:"client-id" default:"concourse-worker" description:"Client used to fetch a token from the auth server. NOTE: if you change this value you will also need to change the --system-claim-value flag so the atc knows to allow requests from this client."`
	ClientSecret string   `long:"client-secret" required:"true" description:"Client used to fetch a token from the auth server"`
	TokenURL     flag.URL `long:"token-url" required:"true" description:"Token endpoint of the auth server"`
	Scopes       []string `long:"scope" description:"Scopes to request from the auth server"`

	HeartbeatInterval time.Duration `long:"heartbeat-interval" default:"30s" description:"interval on which to heartbeat workers to the ATC"`

	ClusterName    string `long:"cluster-name" description:"A name for this Concourse cluster, to be displayed on the dashboard page."`
	LogClusterName bool   `long:"log-cluster-name" description:"Log cluster name."`
}

type TeamAuthKeys struct {
	Team     string
	AuthKeys []ssh.PublicKey
}

type yamlTeamAuthorizedKey struct {
	Team string   `yaml:"team"`
	Keys []string `yaml:"ssh_keys,flow"`
}

func (cmd *TSACommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	tsaServerMember := grouper.Member{
		Name:   "tsa-server",
		Runner: sigmon.New(runner),
	}

	tsaDebugMember := grouper.Member{
		Name: "debug-server",
		Runner: http_server.New(
			cmd.debugBindAddr(),
			http.DefaultServeMux,
		)}

	members := []grouper.Member{
		tsaDebugMember,
		tsaServerMember,
	}

	group := grouper.NewParallel(os.Interrupt, members)
	return <-ifrit.Invoke(group).Wait()
}

func (cmd *TSACommand) Runner(args []string) (ifrit.Runner, error) {
	logger, _ := cmd.constructLogger()

	atcEndpointPicker := tsa.NewRandomATCEndpointPicker(cmd.ATCURLs)

	teamAuthorizedKeys, err := cmd.loadTeamAuthorizedKeys()
	if err != nil {
		return nil, fmt.Errorf("failed to load team authorized keys: %s", err)
	}

	if len(cmd.AuthorizedKeys.Keys)+len(cmd.TeamAuthorizedKeys) == 0 {
		logger.Info("starting-tsa-without-authorized-keys")
	}

	sessionAuthTeam := &sessionTeam{
		sessionTeams: make(map[string]string),
		lock:         &sync.RWMutex{},
	}

	config, err := cmd.configureSSHServer(sessionAuthTeam, cmd.AuthorizedKeys.Keys, teamAuthorizedKeys)
	if err != nil {
		return nil, fmt.Errorf("failed to configure SSH server: %s", err)
	}

	listenAddr := fmt.Sprintf("%s:%d", cmd.BindIP, cmd.BindPort)

	authConfig := clientcredentials.Config{
		ClientID:     cmd.ClientID,
		ClientSecret: cmd.ClientSecret,
		TokenURL:     cmd.TokenURL.URL.String(),
		Scopes:       cmd.Scopes,
	}

	ctx := context.Background()

	tokenSource := authConfig.TokenSource(ctx)
	idTokenSource := token.NewTokenSource(tokenSource)
	httpClient := oauth2.NewClient(ctx, idTokenSource)

	server := &server{
		logger:            logger,
		heartbeatInterval: cmd.HeartbeatInterval,
		cprInterval:       1 * time.Second,
		atcEndpointPicker: atcEndpointPicker,
		forwardHost:       cmd.PeerAddress,
		config:            config,
		httpClient:        httpClient,
		sessionTeam:       sessionAuthTeam,
	}
	// Starts a goroutine that his purpose is to basically listen to the
	// SIGHUP syscall to then reload the config.
	// For now it only reload the TSACommand.AuthorizedKeys but any
	// other configuration could be added
	go func() {
		// Set up channel on which to send signal notifications.
		// We must use a buffered channel or risk missing the signal
		// if we're not ready to receive when the signal is sent.
		c := make(chan os.Signal, 1)
		defer close(c)
		signal.Notify(c, syscall.SIGHUP)
		for {

			// Block until a signal is received.
			_ = <-c

			logger.Info("reloading-config")

			err := cmd.AuthorizedKeys.Reload()
			if err != nil {
				logger.Error("failed to reload authorized keys file : %s", err)
				continue
			}

			teamAuthorizedKeys, err = cmd.loadTeamAuthorizedKeys()
			if err != nil {
				logger.Error("failed to load team authorized keys : %s", err)
				continue
			}

			// compute again the config so it's updated
			config, err := cmd.configureSSHServer(sessionAuthTeam, cmd.AuthorizedKeys.Keys, teamAuthorizedKeys)
			if err != nil {
				logger.Error("failed to configure SSH server: %s", err)
				continue
			}

			server.config = config
		}
	}()

	return serverRunner{logger, server, listenAddr}, nil
}

func (cmd *TSACommand) constructLogger() (lager.Logger, *lager.ReconfigurableSink) {
	logger, reconfigurableSink := cmd.Logger.Logger("tsa")
	if cmd.LogClusterName {
		logger = logger.WithData(lager.Data{
			"cluster": cmd.ClusterName,
		})
	}

	return logger, reconfigurableSink
}

func (cmd *TSACommand) loadTeamAuthorizedKeys() ([]TeamAuthKeys, error) {
	var teamKeys []TeamAuthKeys

	for teamName, keys := range cmd.TeamAuthorizedKeys {
		teamKeys = append(teamKeys, TeamAuthKeys{
			Team:     teamName,
			AuthKeys: keys.Keys,
		})
	}

	// load TeamAuthorizedKeysFile
	if cmd.TeamAuthorizedKeysFile != "" {
		logger, _ := cmd.constructLogger()
		var rawTeamAuthorizedKeys []yamlTeamAuthorizedKey

		authorizedKeysBytes, err := ioutil.ReadFile(cmd.TeamAuthorizedKeysFile.Path())
		if err != nil {
			return nil, fmt.Errorf("failed to read yaml authorized keys file: %s", err)
		}
		err = yaml.Unmarshal([]byte(authorizedKeysBytes), &rawTeamAuthorizedKeys)
		if err != nil {
			return nil, fmt.Errorf("failed to parse yaml authorized keys file: %s", err)
		}

		for _, t := range rawTeamAuthorizedKeys {
			var teamAuthorizedKeys []ssh.PublicKey
			for _, k := range t.Keys {
				key, _, _, _, err := ssh.ParseAuthorizedKey([]byte(k))
				if err != nil {
					logger.Error("load-team-authorized-keys-parse", fmt.Errorf("Invalid format, ignoring (%s): %s", k, err.Error()))
					continue
				}
				logger.Info("load-team-authorized-keys-loaded", lager.Data{"team": t.Team, "key": k})
				teamAuthorizedKeys = append(teamAuthorizedKeys, key)
			}
			teamKeys = append(teamKeys, TeamAuthKeys{Team: t.Team, AuthKeys: teamAuthorizedKeys})
		}
	}

	return teamKeys, nil
}

func (cmd *TSACommand) configureSSHServer(sessionAuthTeam *sessionTeam, authorizedKeys []ssh.PublicKey, teamAuthorizedKeys []TeamAuthKeys) (*ssh.ServerConfig, error) {
	certChecker := &ssh.CertChecker{
		IsUserAuthority: func(key ssh.PublicKey) bool {
			return false
		},

		IsHostAuthority: func(key ssh.PublicKey, address string) bool {
			return false
		},

		UserKeyFallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			for _, k := range authorizedKeys {
				if bytes.Equal(k.Marshal(), key.Marshal()) {
					return nil, nil
				}
			}

			for _, teamKeys := range teamAuthorizedKeys {
				for _, k := range teamKeys.AuthKeys {
					if bytes.Equal(k.Marshal(), key.Marshal()) {
						sessionAuthTeam.AuthorizeTeam(string(conn.SessionID()), teamKeys.Team)
						return nil, nil
					}
				}
			}

			return nil, fmt.Errorf("unknown public key")
		},
	}

	config := &ssh.ServerConfig{
		Config: atc.DefaultSSHConfig(),
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return certChecker.Authenticate(conn, key)
		},
	}

	signer, err := ssh.NewSignerFromKey(cmd.HostKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create signer from host key: %s", err)
	}

	config.AddHostKey(signer)

	return config, nil
}

func (cmd *TSACommand) debugBindAddr() string {
	return fmt.Sprintf("%s:%d", cmd.DebugBindIP, cmd.DebugBindPort)
}
<commit_msg>Update tsa/tsacmd/command.go<commit_after>package tsacmd

import (
	"bytes"
	"context"
	"fmt"
	"net/http"
	"os"
	"sync"
	"time"

	"io/ioutil"

	yaml "gopkg.in/yaml.v2"

	"os/signal"
	"syscall"

	"code.cloudfoundry.org/lager"
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/skymarshal/token"
	"github.com/concourse/concourse/tsa"
	"github.com/concourse/flag"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"golang.org/x/crypto/ssh"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/clientcredentials"
)

type TSACommand struct {
	Logger flag.Lager

	BindIP      flag.IP `long:"bind-ip"   default:"0.0.0.0" description:"IP address on which to listen for SSH."`
	PeerAddress string  `long:"peer-address" default:"127.0.0.1" description:"Network address of this web node, reachable by other web nodes. Used for forwarded worker addresses."`
	BindPort    uint16  `long:"bind-port" default:"2222"    description:"Port on which to listen for SSH."`

	DebugBindIP   flag.IP `long:"debug-bind-ip"   default:"127.0.0.1" description:"IP address on which to listen for the pprof debugger endpoints."`
	DebugBindPort uint16  `long:"debug-bind-port" default:"2221"      description:"Port on which to listen for the pprof debugger endpoints."`

	HostKey                *flag.PrivateKey               `long:"host-key"        required:"true" description:"Path to private key to use for the SSH server."`
	AuthorizedKeys         flag.AuthorizedKeys            `long:"authorized-keys" description:"Path to file containing keys to authorize, in SSH authorized_keys format (one public key per line)."`
	TeamAuthorizedKeys     map[string]flag.AuthorizedKeys `long:"team-authorized-keys" value-name:"NAME:PATH" description:"Path to file containing keys to authorize, in SSH authorized_keys format (one public key per line)."`
	TeamAuthorizedKeysFile flag.File                      `long:"team-authorized-keys-file" description:"Path to file containing a YAML array of teams and their authorized SSH keys, e.g. [{team:foo,ssh_keys:[key1,key2]}]."`

	ATCURLs []flag.URL `long:"atc-url" required:"true" description:"ATC API endpoints to which workers will be registered."`

	ClientID     string   `long:"client-id" default:"concourse-worker" description:"Client used to fetch a token from the auth server. NOTE: if you change this value you will also need to change the --system-claim-value flag so the atc knows to allow requests from this client."`
	ClientSecret string   `long:"client-secret" required:"true" description:"Client used to fetch a token from the auth server"`
	TokenURL     flag.URL `long:"token-url" required:"true" description:"Token endpoint of the auth server"`
	Scopes       []string `long:"scope" description:"Scopes to request from the auth server"`

	HeartbeatInterval time.Duration `long:"heartbeat-interval" default:"30s" description:"interval on which to heartbeat workers to the ATC"`

	ClusterName    string `long:"cluster-name" description:"A name for this Concourse cluster, to be displayed on the dashboard page."`
	LogClusterName bool   `long:"log-cluster-name" description:"Log cluster name."`
}

type TeamAuthKeys struct {
	Team     string
	AuthKeys []ssh.PublicKey
}

type yamlTeamAuthorizedKey struct {
	Team string   `yaml:"team"`
	Keys []string `yaml:"ssh_keys,flow"`
}

func (cmd *TSACommand) Execute(args []string) error {
	runner, err := cmd.Runner(args)
	if err != nil {
		return err
	}

	tsaServerMember := grouper.Member{
		Name:   "tsa-server",
		Runner: sigmon.New(runner),
	}

	tsaDebugMember := grouper.Member{
		Name: "debug-server",
		Runner: http_server.New(
			cmd.debugBindAddr(),
			http.DefaultServeMux,
		)}

	members := []grouper.Member{
		tsaDebugMember,
		tsaServerMember,
	}

	group := grouper.NewParallel(os.Interrupt, members)
	return <-ifrit.Invoke(group).Wait()
}

func (cmd *TSACommand) Runner(args []string) (ifrit.Runner, error) {
	logger, _ := cmd.constructLogger()

	atcEndpointPicker := tsa.NewRandomATCEndpointPicker(cmd.ATCURLs)

	teamAuthorizedKeys, err := cmd.loadTeamAuthorizedKeys()
	if err != nil {
		return nil, fmt.Errorf("failed to load team authorized keys: %s", err)
	}

	if len(cmd.AuthorizedKeys.Keys)+len(cmd.TeamAuthorizedKeys) == 0 {
		logger.Info("starting-tsa-without-authorized-keys")
	}

	sessionAuthTeam := &sessionTeam{
		sessionTeams: make(map[string]string),
		lock:         &sync.RWMutex{},
	}

	config, err := cmd.configureSSHServer(sessionAuthTeam, cmd.AuthorizedKeys.Keys, teamAuthorizedKeys)
	if err != nil {
		return nil, fmt.Errorf("failed to configure SSH server: %s", err)
	}

	listenAddr := fmt.Sprintf("%s:%d", cmd.BindIP, cmd.BindPort)

	authConfig := clientcredentials.Config{
		ClientID:     cmd.ClientID,
		ClientSecret: cmd.ClientSecret,
		TokenURL:     cmd.TokenURL.URL.String(),
		Scopes:       cmd.Scopes,
	}

	ctx := context.Background()

	tokenSource := authConfig.TokenSource(ctx)
	idTokenSource := token.NewTokenSource(tokenSource)
	httpClient := oauth2.NewClient(ctx, idTokenSource)

	server := &server{
		logger:            logger,
		heartbeatInterval: cmd.HeartbeatInterval,
		cprInterval:       1 * time.Second,
		atcEndpointPicker: atcEndpointPicker,
		forwardHost:       cmd.PeerAddress,
		config:            config,
		httpClient:        httpClient,
		sessionTeam:       sessionAuthTeam,
	}
	// Starts a goroutine whose purpose is to listen to the
	// SIGHUP syscall and reload configuration upon receiving the signal.
 	// For now it only reloads the TSACommand.AuthorizedKeys but
	// other configuration can potentially be added.
	go func() {
		reloadWorkerKeys := make(chan os.Signal, 1)
		defer close(c)
		signal.Notify(c, syscall.SIGHUP)
		for {

			// Block until a signal is received.
			<-c

			logger.Info("reloading-config")

			err := cmd.AuthorizedKeys.Reload()
			if err != nil {
				logger.Error("failed to reload authorized keys file : %s", err)
				continue
			}

			teamAuthorizedKeys, err = cmd.loadTeamAuthorizedKeys()
			if err != nil {
				logger.Error("failed to load team authorized keys : %s", err)
				continue
			}

			// Reconfigure the SSH server with the new keys
			config, err := cmd.configureSSHServer(sessionAuthTeam, cmd.AuthorizedKeys.Keys, teamAuthorizedKeys)
			if err != nil {
				logger.Error("failed to configure SSH server: %s", err)
				continue
			}

			server.config = config
		}
	}()

	return serverRunner{logger, server, listenAddr}, nil
}

func (cmd *TSACommand) constructLogger() (lager.Logger, *lager.ReconfigurableSink) {
	logger, reconfigurableSink := cmd.Logger.Logger("tsa")
	if cmd.LogClusterName {
		logger = logger.WithData(lager.Data{
			"cluster": cmd.ClusterName,
		})
	}

	return logger, reconfigurableSink
}

func (cmd *TSACommand) loadTeamAuthorizedKeys() ([]TeamAuthKeys, error) {
	var teamKeys []TeamAuthKeys

	for teamName, keys := range cmd.TeamAuthorizedKeys {
		teamKeys = append(teamKeys, TeamAuthKeys{
			Team:     teamName,
			AuthKeys: keys.Keys,
		})
	}

	// load TeamAuthorizedKeysFile
	if cmd.TeamAuthorizedKeysFile != "" {
		logger, _ := cmd.constructLogger()
		var rawTeamAuthorizedKeys []yamlTeamAuthorizedKey

		authorizedKeysBytes, err := ioutil.ReadFile(cmd.TeamAuthorizedKeysFile.Path())
		if err != nil {
			return nil, fmt.Errorf("failed to read yaml authorized keys file: %s", err)
		}
		err = yaml.Unmarshal([]byte(authorizedKeysBytes), &rawTeamAuthorizedKeys)
		if err != nil {
			return nil, fmt.Errorf("failed to parse yaml authorized keys file: %s", err)
		}

		for _, t := range rawTeamAuthorizedKeys {
			var teamAuthorizedKeys []ssh.PublicKey
			for _, k := range t.Keys {
				key, _, _, _, err := ssh.ParseAuthorizedKey([]byte(k))
				if err != nil {
					logger.Error("load-team-authorized-keys-parse", fmt.Errorf("Invalid format, ignoring (%s): %s", k, err.Error()))
					continue
				}
				logger.Info("load-team-authorized-keys-loaded", lager.Data{"team": t.Team, "key": k})
				teamAuthorizedKeys = append(teamAuthorizedKeys, key)
			}
			teamKeys = append(teamKeys, TeamAuthKeys{Team: t.Team, AuthKeys: teamAuthorizedKeys})
		}
	}

	return teamKeys, nil
}

func (cmd *TSACommand) configureSSHServer(sessionAuthTeam *sessionTeam, authorizedKeys []ssh.PublicKey, teamAuthorizedKeys []TeamAuthKeys) (*ssh.ServerConfig, error) {
	certChecker := &ssh.CertChecker{
		IsUserAuthority: func(key ssh.PublicKey) bool {
			return false
		},

		IsHostAuthority: func(key ssh.PublicKey, address string) bool {
			return false
		},

		UserKeyFallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			for _, k := range authorizedKeys {
				if bytes.Equal(k.Marshal(), key.Marshal()) {
					return nil, nil
				}
			}

			for _, teamKeys := range teamAuthorizedKeys {
				for _, k := range teamKeys.AuthKeys {
					if bytes.Equal(k.Marshal(), key.Marshal()) {
						sessionAuthTeam.AuthorizeTeam(string(conn.SessionID()), teamKeys.Team)
						return nil, nil
					}
				}
			}

			return nil, fmt.Errorf("unknown public key")
		},
	}

	config := &ssh.ServerConfig{
		Config: atc.DefaultSSHConfig(),
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return certChecker.Authenticate(conn, key)
		},
	}

	signer, err := ssh.NewSignerFromKey(cmd.HostKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create signer from host key: %s", err)
	}

	config.AddHostKey(signer)

	return config, nil
}

func (cmd *TSACommand) debugBindAddr() string {
	return fmt.Sprintf("%s:%d", cmd.DebugBindIP, cmd.DebugBindPort)
}
<|endoftext|>