<commit_before>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"code.google.com/p/go.crypto/ssh/terminal"
	"github.com/tsuru/tsuru/cmd"
	tsuruIo "github.com/tsuru/tsuru/io"
	"launchpad.net/gnuflag"
)

var httpHeaderRegexp = regexp.MustCompile(`HTTP/.*? (\d+)`)

type moveContainersCmd struct{}

type progressFormatter struct{}

func (progressFormatter) Format(out io.Writer, data []byte) error {
	var logEntry progressLog
	err := json.Unmarshal(data, &logEntry)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "%s\n", logEntry.Message)
	return nil
}

func (c *moveContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "containers-move",
		Usage:   "containers-move <from host> <to host>",
		Desc:    "Move all containers from one host to another.\nThis command is especially useful for host maintenance.",
		MinArgs: 2,
	}
}

func (c *moveContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/docker/containers/move")
	if err != nil {
		return err
	}
	params := map[string]string{
		"from": context.Args[0],
		"to":   context.Args[1],
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

type fixContainersCmd struct{}

func (fixContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "fix-containers",
		Usage: "fix-containers",
		Desc:  "Fix containers that are broken in the cluster.",
	}
}

func (fixContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/docker/fix-containers")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("POST", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	return err
}

type moveContainerCmd struct{}

func (c *moveContainerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "container-move",
		Usage:   "container-move <container id> <to host>",
		Desc:    "Move specified container to another host.",
		MinArgs: 2,
	}
}

func (c *moveContainerCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL(fmt.Sprintf("/docker/container/%s/move", context.Args[0]))
	if err != nil {
		return err
	}
	params := map[string]string{
		"to": context.Args[1],
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

type rebalanceContainersCmd struct {
	cmd.ConfirmationCommand
	fs  *gnuflag.FlagSet
	dry bool
}

func (c *rebalanceContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "containers-rebalance",
		Usage:   "containers-rebalance [--dry] [-y/--assume-yes]",
		Desc:    "Move containers creating a more even distribution between docker nodes.",
		MinArgs: 0,
	}
}

func (c *rebalanceContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.dry && !c.Confirm(context, "Are you sure you want to rebalance containers?") {
		return nil
	}
	url, err := cmd.GetURL("/docker/containers/rebalance")
	if err != nil {
		return err
	}
	params := map[string]string{
		"dry": fmt.Sprintf("%t", c.dry),
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

func (c *rebalanceContainersCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = c.ConfirmationCommand.Flags()
		c.fs.BoolVar(&c.dry, "dry", false, "Dry run, only shows what would be done")
	}
	return c.fs
}

type sshToContainerCmd struct {
	cmd.GuessingCommand
}

func (c *sshToContainerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "ssh",
		Usage:   "ssh <[-a/--app <appname>]|[container-id]>",
		Desc:    "Open an SSH shell to the given container, or to one of the containers of the given app.",
		MinArgs: 0,
	}
}

func (c *sshToContainerCmd) Run(context *cmd.Context, client *cmd.Client) error {
	var width, height int
	if stdin, ok := context.Stdin.(*os.File); ok {
		fd := int(stdin.Fd())
		if terminal.IsTerminal(fd) {
			width, height, _ = terminal.GetSize(fd)
			oldState, err := terminal.MakeRaw(fd)
			if err != nil {
				return err
			}
			defer terminal.Restore(fd, oldState)
			sigChan := make(chan os.Signal, 2)
			go func(c <-chan os.Signal) {
				if _, ok := <-c; ok {
					terminal.Restore(fd, oldState)
					os.Exit(1)
				}
			}(sigChan)
			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)
		}
	}
	queryString := make(url.Values)
	queryString.Set("width", strconv.Itoa(width))
	queryString.Set("height", strconv.Itoa(height))
	container, err := c.getContainer(context, client)
	if err != nil {
		return err
	}
	serverURL, err := cmd.GetURL("/docker/ssh/" + container + "?" + queryString.Encode())
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", serverURL, nil)
	if err != nil {
		return err
	}
	request.Close = true
	token, err := cmd.ReadToken()
	if err == nil {
		request.Header.Set("Authorization", "bearer "+token)
	}
	parsedURL, _ := url.Parse(serverURL)
	conn, err := net.Dial("tcp", parsedURL.Host)
	if err != nil {
		return err
	}
	defer conn.Close()
	request.Write(conn)
	bytesLimit := 50
	var readStr string
	byteBuffer := make([]byte, 1)
	for i := 0; i < bytesLimit && byteBuffer[0] != '\n'; i++ {
		_, err := conn.Read(byteBuffer)
		if err != nil {
			break
		}
		readStr += string(byteBuffer)
	}
	matches := httpHeaderRegexp.FindAllStringSubmatch(readStr, -1)
	if len(matches) > 0 && len(matches[0]) > 1 {
		return errors.New(strings.TrimSpace(readStr))
	} else {
		context.Stdout.Write([]byte(readStr))
	}
	errs := make(chan error, 2)
	quit := make(chan bool)
	go io.Copy(conn, context.Stdin)
	go func() {
		defer close(quit)
		_, err := io.Copy(context.Stdout, conn)
		if err != nil && err != io.EOF {
			errs <- err
		}
	}()
	<-quit
	close(errs)
	return <-errs
}

func (c *sshToContainerCmd) getContainer(ctx *cmd.Context, client *cmd.Client) (string, error) {
	if appName, _ := c.Guess(); appName != "" {
		url, err := cmd.GetURL("/apps/" + appName)
		if err != nil {
			return "", err
		}
		request, _ := http.NewRequest("GET", url, nil)
		resp, err := client.Do(request)
		if err != nil {
			return "", err
		}
		defer resp.Body.Close()
		var app apiApp
		err = json.NewDecoder(resp.Body).Decode(&app)
		if err != nil {
			return "", err
		}
		if len(app.Units) < 1 {
			return "", errors.New("app must have at least one container")
		}
		return app.Units[0].Name, nil
	}
	if len(ctx.Args) < 1 {
		return "", errors.New("you need to specify either the container id or the app name")
	}
	return ctx.Args[0], nil
}

type unit struct {
	Name string
}

type apiApp struct {
	Units []unit
}
<commit_msg>provision/docker: the container ID has preference of the app name in ssh<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"code.google.com/p/go.crypto/ssh/terminal"
	"github.com/tsuru/tsuru/cmd"
	tsuruIo "github.com/tsuru/tsuru/io"
	"launchpad.net/gnuflag"
)

var httpHeaderRegexp = regexp.MustCompile(`HTTP/.*? (\d+)`)

type moveContainersCmd struct{}

type progressFormatter struct{}

func (progressFormatter) Format(out io.Writer, data []byte) error {
	var logEntry progressLog
	err := json.Unmarshal(data, &logEntry)
	if err != nil {
		return err
	}
	fmt.Fprintf(out, "%s\n", logEntry.Message)
	return nil
}

func (c *moveContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "containers-move",
		Usage:   "containers-move <from host> <to host>",
		Desc:    "Move all containers from one host to another.\nThis command is especially useful for host maintenance.",
		MinArgs: 2,
	}
}

func (c *moveContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/docker/containers/move")
	if err != nil {
		return err
	}
	params := map[string]string{
		"from": context.Args[0],
		"to":   context.Args[1],
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

type fixContainersCmd struct{}

func (fixContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:  "fix-containers",
		Usage: "fix-containers",
		Desc:  "Fix containers that are broken in the cluster.",
	}
}

func (fixContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL("/docker/fix-containers")
	if err != nil {
		return err
	}
	request, err := http.NewRequest("POST", url, nil)
	if err != nil {
		return err
	}
	_, err = client.Do(request)
	return err
}

type moveContainerCmd struct{}

func (c *moveContainerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "container-move",
		Usage:   "container-move <container id> <to host>",
		Desc:    "Move specified container to another host.",
		MinArgs: 2,
	}
}

func (c *moveContainerCmd) Run(context *cmd.Context, client *cmd.Client) error {
	url, err := cmd.GetURL(fmt.Sprintf("/docker/container/%s/move", context.Args[0]))
	if err != nil {
		return err
	}
	params := map[string]string{
		"to": context.Args[1],
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

type rebalanceContainersCmd struct {
	cmd.ConfirmationCommand
	fs  *gnuflag.FlagSet
	dry bool
}

func (c *rebalanceContainersCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "containers-rebalance",
		Usage:   "containers-rebalance [--dry] [-y/--assume-yes]",
		Desc:    "Move containers creating a more even distribution between docker nodes.",
		MinArgs: 0,
	}
}

func (c *rebalanceContainersCmd) Run(context *cmd.Context, client *cmd.Client) error {
	if !c.dry && !c.Confirm(context, "Are you sure you want to rebalance containers?") {
		return nil
	}
	url, err := cmd.GetURL("/docker/containers/rebalance")
	if err != nil {
		return err
	}
	params := map[string]string{
		"dry": fmt.Sprintf("%t", c.dry),
	}
	b, err := json.Marshal(params)
	if err != nil {
		return err
	}
	buffer := bytes.NewBuffer(b)
	request, err := http.NewRequest("POST", url, buffer)
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := client.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	w := tsuruIo.NewStreamWriter(context.Stdout, progressFormatter{})
	for n := int64(1); n > 0 && err == nil; n, err = io.Copy(w, response.Body) {
	}
	return nil
}

func (c *rebalanceContainersCmd) Flags() *gnuflag.FlagSet {
	if c.fs == nil {
		c.fs = c.ConfirmationCommand.Flags()
		c.fs.BoolVar(&c.dry, "dry", false, "Dry run, only shows what would be done")
	}
	return c.fs
}

type sshToContainerCmd struct {
	cmd.GuessingCommand
}

func (c *sshToContainerCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "ssh",
		Usage:   "ssh <[-a/--app <appname>]|[container-id]>",
		Desc:    "Open an SSH shell to the given container, or to one of the containers of the given app.",
		MinArgs: 0,
	}
}

func (c *sshToContainerCmd) Run(context *cmd.Context, client *cmd.Client) error {
	var width, height int
	if stdin, ok := context.Stdin.(*os.File); ok {
		fd := int(stdin.Fd())
		if terminal.IsTerminal(fd) {
			width, height, _ = terminal.GetSize(fd)
			oldState, err := terminal.MakeRaw(fd)
			if err != nil {
				return err
			}
			defer terminal.Restore(fd, oldState)
			sigChan := make(chan os.Signal, 2)
			go func(c <-chan os.Signal) {
				if _, ok := <-c; ok {
					terminal.Restore(fd, oldState)
					os.Exit(1)
				}
			}(sigChan)
			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)
		}
	}
	queryString := make(url.Values)
	queryString.Set("width", strconv.Itoa(width))
	queryString.Set("height", strconv.Itoa(height))
	container, err := c.getContainer(context, client)
	if err != nil {
		return err
	}
	serverURL, err := cmd.GetURL("/docker/ssh/" + container + "?" + queryString.Encode())
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", serverURL, nil)
	if err != nil {
		return err
	}
	request.Close = true
	token, err := cmd.ReadToken()
	if err == nil {
		request.Header.Set("Authorization", "bearer "+token)
	}
	parsedURL, _ := url.Parse(serverURL)
	conn, err := net.Dial("tcp", parsedURL.Host)
	if err != nil {
		return err
	}
	defer conn.Close()
	request.Write(conn)
	bytesLimit := 50
	var readStr string
	byteBuffer := make([]byte, 1)
	for i := 0; i < bytesLimit && byteBuffer[0] != '\n'; i++ {
		_, err := conn.Read(byteBuffer)
		if err != nil {
			break
		}
		readStr += string(byteBuffer)
	}
	matches := httpHeaderRegexp.FindAllStringSubmatch(readStr, -1)
	if len(matches) > 0 && len(matches[0]) > 1 {
		return errors.New(strings.TrimSpace(readStr))
	} else {
		context.Stdout.Write([]byte(readStr))
	}
	errs := make(chan error, 2)
	quit := make(chan bool)
	go io.Copy(conn, context.Stdin)
	go func() {
		defer close(quit)
		_, err := io.Copy(context.Stdout, conn)
		if err != nil && err != io.EOF {
			errs <- err
		}
	}()
	<-quit
	close(errs)
	return <-errs
}

func (c *sshToContainerCmd) getContainer(ctx *cmd.Context, client *cmd.Client) (string, error) {
	if len(ctx.Args) > 0 {
		return ctx.Args[0], nil
	}
	if appName, _ := c.Guess(); appName != "" {
		url, err := cmd.GetURL("/apps/" + appName)
		if err != nil {
			return "", err
		}
		request, _ := http.NewRequest("GET", url, nil)
		resp, err := client.Do(request)
		if err != nil {
			return "", err
		}
		defer resp.Body.Close()
		var app apiApp
		err = json.NewDecoder(resp.Body).Decode(&app)
		if err != nil {
			return "", err
		}
		if len(app.Units) < 1 {
			return "", errors.New("app must have at least one container")
		}
		return app.Units[0].Name, nil
	}
	return "", errors.New("you need to specify either the container id or the app name")
}

type unit struct {
	Name string
}

type apiApp struct {
	Units []unit
}
<|endoftext|>


<commit_before>// Copyright 2017 Authors of Cilium
//
// 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 helpers

import (
	"bytes"
	"context"
	"fmt"
	log "github.com/sirupsen/logrus"
	"io"
	"os"
)

// SSHMeta contains metadata to SSH into a remote location to run tests
type SSHMeta struct {
	sshClient *SSHClient
	env       []string
}

// CreateSSHMeta returns an SSHMeta with the specified host, port, and user, as
// well as an according SSHClient.
func CreateSSHMeta(host string, port int, user string) *SSHMeta {
	return &SSHMeta{
		sshClient: GetSSHClient(host, port, user),
	}
}

func (s *SSHMeta) String() string {
	return fmt.Sprintf("environment: %s, SSHClient: %s", s.env, s.sshClient.String())

}

// GetVagrantSSHMetadata returns a SSHMeta initialized based on the provided
// SSH-config target.
func GetVagrantSSHMetadata(vmName string) *SSHMeta {
	var vagrant Vagrant
	config, err := vagrant.GetVagrantSSHMetadata(vmName)
	if err != nil {
		return nil
	}

	log.Debugf("generated SSHConfig for node %s", vmName)
	nodes, err := ImportSSHconfig(config)
	if err != nil {
		return nil
	}
	log.Debugf("done importing ssh config")
	node := nodes[vmName]
	if node == nil {
		return nil
	}

	return &SSHMeta{
		sshClient: node.GetSSHClient(),
	}
}

// Execute executes cmd on the provided node and stores the stdout / stderr of
// the command in the provided buffers. Returns false if the command failed
// during its execution.
func (s *SSHMeta) Execute(cmd string, stdout io.Writer, stderr io.Writer) bool {
	if stdout == nil {
		stdout = os.Stdout
	}

	if stderr == nil {
		stderr = os.Stderr
	}

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	err := s.sshClient.RunCommand(command)
	if err != nil {
		return false
	}
	return true
}

//ExecWithSudo executes the provided command using sudo privileges. The stdout
//and stderr of the command are written to the specified stdout / stderr
//buffers accordingly. Returns false if execution of cmd failed.
func (s *SSHMeta) ExecWithSudo(cmd string, stdout io.Writer, stderr io.Writer) bool {
	command := fmt.Sprintf("sudo %s", cmd)
	return s.Execute(command, stdout, stderr)
}

//Exec executes the provided cmd and returns metadata about its result in CmdRes
func (s *SSHMeta) Exec(cmd string) *CmdRes {
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	exit := s.Execute(cmd, stdout, stderr)

	return &CmdRes{
		cmd:    cmd,
		stdout: stdout,
		stderr: stderr,
		exit:   exit,
	}
}

//ExecContext run a command in background and stop when cancel the context
func (s *SSHMeta) ExecContext(ctx context.Context, cmd string) *CmdRes {
	if ctx == nil {
		panic("no context provided")
	}

	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}

	go func() {
		s.sshClient.RunCommandContext(ctx, command)
	}()

	return &CmdRes{
		cmd:    cmd,
		stdout: stdout,
		stderr: stderr,
		exit:   false,
	}
}
<commit_msg>test/helpers: fix import order in node.go<commit_after>// Copyright 2017 Authors of Cilium
//
// 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 helpers

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"

	log "github.com/sirupsen/logrus"
)

// SSHMeta contains metadata to SSH into a remote location to run tests
type SSHMeta struct {
	sshClient *SSHClient
	env       []string
}

// CreateSSHMeta returns an SSHMeta with the specified host, port, and user, as
// well as an according SSHClient.
func CreateSSHMeta(host string, port int, user string) *SSHMeta {
	return &SSHMeta{
		sshClient: GetSSHClient(host, port, user),
	}
}

func (s *SSHMeta) String() string {
	return fmt.Sprintf("environment: %s, SSHClient: %s", s.env, s.sshClient.String())

}

// GetVagrantSSHMetadata returns a SSHMeta initialized based on the provided
// SSH-config target.
func GetVagrantSSHMetadata(vmName string) *SSHMeta {
	var vagrant Vagrant
	config, err := vagrant.GetVagrantSSHMetadata(vmName)
	if err != nil {
		return nil
	}

	log.Debugf("generated SSHConfig for node %s", vmName)
	nodes, err := ImportSSHconfig(config)
	if err != nil {
		return nil
	}
	log.Debugf("done importing ssh config")
	node := nodes[vmName]
	if node == nil {
		return nil
	}

	return &SSHMeta{
		sshClient: node.GetSSHClient(),
	}
}

// Execute executes cmd on the provided node and stores the stdout / stderr of
// the command in the provided buffers. Returns false if the command failed
// during its execution.
func (s *SSHMeta) Execute(cmd string, stdout io.Writer, stderr io.Writer) bool {
	if stdout == nil {
		stdout = os.Stdout
	}

	if stderr == nil {
		stderr = os.Stderr
	}

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}
	err := s.sshClient.RunCommand(command)
	if err != nil {
		return false
	}
	return true
}

//ExecWithSudo executes the provided command using sudo privileges. The stdout
//and stderr of the command are written to the specified stdout / stderr
//buffers accordingly. Returns false if execution of cmd failed.
func (s *SSHMeta) ExecWithSudo(cmd string, stdout io.Writer, stderr io.Writer) bool {
	command := fmt.Sprintf("sudo %s", cmd)
	return s.Execute(command, stdout, stderr)
}

//Exec executes the provided cmd and returns metadata about its result in CmdRes
func (s *SSHMeta) Exec(cmd string) *CmdRes {
	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	exit := s.Execute(cmd, stdout, stderr)

	return &CmdRes{
		cmd:    cmd,
		stdout: stdout,
		stderr: stderr,
		exit:   exit,
	}
}

//ExecContext run a command in background and stop when cancel the context
func (s *SSHMeta) ExecContext(ctx context.Context, cmd string) *CmdRes {
	if ctx == nil {
		panic("no context provided")
	}

	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)

	command := &SSHCommand{
		Path:   cmd,
		Stdin:  os.Stdin,
		Stdout: stdout,
		Stderr: stderr,
	}

	go func() {
		s.sshClient.RunCommandContext(ctx, command)
	}()

	return &CmdRes{
		cmd:    cmd,
		stdout: stdout,
		stderr: stderr,
		exit:   false,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Tekton 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 resources

import (
	"testing"

	tb "github.com/tektoncd/pipeline/internal/builder/v1beta1"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1beta1"
)

func TestValidateParamTypesMatching_Valid(t *testing.T) {
	tcs := []struct {
		name          string
		p             *v1beta1.Pipeline
		pr            *v1beta1.PipelineRun
		errorExpected bool
	}{{
		name: "proper param types",
		p: tb.Pipeline("a-pipeline", tb.PipelineSpec(
			tb.PipelineParamSpec("correct-type-1", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("mismatching-type", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("correct-type-2", v1beta1.ParamTypeArray))),
		pr: tb.PipelineRun("a-pipelinerun", tb.PipelineRunSpec(
			"test-pipeline",
			tb.PipelineRunParam("correct-type-1", "somestring"),
			tb.PipelineRunParam("mismatching-type", "astring"),
			tb.PipelineRunParam("correct-type-2", "another", "array"))),
		errorExpected: false,
	}, {
		name:          "no params to get wrong",
		p:             tb.Pipeline("a-pipeline"),
		pr:            tb.PipelineRun("a-pipelinerun"),
		errorExpected: false,
	}, {
		name: "string-array mismatch",
		p: tb.Pipeline("a-pipeline", tb.PipelineSpec(
			tb.PipelineParamSpec("correct-type-1", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("mismatching-type", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("correct-type-2", v1beta1.ParamTypeArray))),
		pr: tb.PipelineRun("a-pipelinerun",
			tb.PipelineRunSpec("test-pipeline",
				tb.PipelineRunParam("correct-type-1", "somestring"),
				tb.PipelineRunParam("mismatching-type", "an", "array"),
				tb.PipelineRunParam("correct-type-2", "another", "array"))),
		errorExpected: true,
	}, {
		name: "array-string mismatch",
		p: tb.Pipeline("a-pipeline", tb.PipelineSpec(
			tb.PipelineParamSpec("correct-type-1", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("mismatching-type", v1beta1.ParamTypeArray),
			tb.PipelineParamSpec("correct-type-2", v1beta1.ParamTypeArray))),
		pr: tb.PipelineRun("a-pipelinerun",
			tb.PipelineRunSpec("test-pipeline",
				tb.PipelineRunParam("correct-type-1", "somestring"),
				tb.PipelineRunParam("mismatching-type", "astring"),
				tb.PipelineRunParam("correct-type-2", "another", "array"))),
		errorExpected: true,
	}}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			err := ValidateParamTypesMatching(&tc.p.Spec, tc.pr)
			if (!tc.errorExpected) && (err != nil) {
				t.Errorf("Pipeline.Validate() returned error: %v", err)
			}

			if tc.errorExpected && (err == nil) {
				t.Error("Pipeline.Validate() did not return error, wanted error")
			}
		})
	}
}

func TestValidateParamTypesMatching_Invalid(t *testing.T) {
	tcs := []struct {
		name string
		p    *v1beta1.Pipeline
		pr   *v1beta1.PipelineRun
	}{{
		name: "string-array mismatch",
		p: tb.Pipeline("a-pipeline", tb.PipelineSpec(
			tb.PipelineParamSpec("correct-type-1", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("mismatching-type", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("correct-type-2", v1beta1.ParamTypeArray))),
		pr: tb.PipelineRun("a-pipelinerun",
			tb.PipelineRunSpec("test-pipeline",
				tb.PipelineRunParam("correct-type-1", "somestring"),
				tb.PipelineRunParam("mismatching-type", "an", "array"),
				tb.PipelineRunParam("correct-type-2", "another", "array"))),
	}, {
		name: "array-string mismatch",
		p: tb.Pipeline("a-pipeline", tb.PipelineSpec(
			tb.PipelineParamSpec("correct-type-1", v1beta1.ParamTypeString),
			tb.PipelineParamSpec("mismatching-type", v1beta1.ParamTypeArray),
			tb.PipelineParamSpec("correct-type-2", v1beta1.ParamTypeArray))),
		pr: tb.PipelineRun("a-pipelinerun",
			tb.PipelineRunSpec("test-pipeline",
				tb.PipelineRunParam("correct-type-1", "somestring"),
				tb.PipelineRunParam("mismatching-type", "astring"),
				tb.PipelineRunParam("correct-type-2", "another", "array"))),
	}}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			if err := ValidateParamTypesMatching(&tc.p.Spec, tc.pr); err == nil {
				t.Errorf("Expected to see error when validating PipelineRun/Pipeline param types but saw none")
			}
		})
	}
}

func TestValidateRequiredParametersProvided_Valid(t *testing.T) {
	tcs := []struct {
		name string
		pp   []v1beta1.ParamSpec
		prp  []v1beta1.Param
	}{{
		name: "required string params provided",
		pp: []v1beta1.ParamSpec{
			{
				Name: "required-string-param",
				Type: v1beta1.ParamTypeString,
			},
		},
		prp: []v1beta1.Param{
			{
				Name:  "required-string-param",
				Value: *v1beta1.NewArrayOrString("somestring"),
			},
		},
	}, {
		name: "required array params provided",
		pp: []v1beta1.ParamSpec{
			{
				Name: "required-array-param",
				Type: v1beta1.ParamTypeArray,
			},
		},
		prp: []v1beta1.Param{
			{
				Name:  "required-array-param",
				Value: *v1beta1.NewArrayOrString("another", "array"),
			},
		},
	}, {
		name: "string params provided in default",
		pp: []v1beta1.ParamSpec{
			{
				Name:    "string-param",
				Type:    v1beta1.ParamTypeString,
				Default: v1beta1.NewArrayOrString("somedefault"),
			},
		},
		prp: []v1beta1.Param{
			{
				Name:  "another-string-param",
				Value: *v1beta1.NewArrayOrString("somestring"),
			},
		},
	}}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			if err := ValidateRequiredParametersProvided(&tc.pp, &tc.prp); err != nil {
				t.Errorf("Didn't expect to see error when validating valid PipelineRun parameters but got: %v", err)
			}
		})
	}
}

func TestValidateRequiredParametersProvided_Invalid(t *testing.T) {
	tcs := []struct {
		name string
		pp   []v1beta1.ParamSpec
		prp  []v1beta1.Param
	}{{
		name: "required string param missing",
		pp: []v1beta1.ParamSpec{
			{
				Name: "required-string-param",
				Type: v1beta1.ParamTypeString,
			},
		},
		prp: []v1beta1.Param{
			{
				Name:  "another-string-param",
				Value: *v1beta1.NewArrayOrString("anotherstring"),
			},
		},
	}, {
		name: "required array param missing",
		pp: []v1beta1.ParamSpec{
			{
				Name: "required-array-param",
				Type: v1beta1.ParamTypeArray,
			},
		},
		prp: []v1beta1.Param{
			{
				Name:  "another-array-param",
				Value: *v1beta1.NewArrayOrString("anotherstring"),
			},
		},
	}}
	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			if err := ValidateRequiredParametersProvided(&tc.pp, &tc.prp); err == nil {
				t.Errorf("Expected to see error when validating invalid PipelineRun parameters but saw none")
			}
		})
	}
}
<commit_msg>Remove test builders from validate_params_test.go<commit_after>/*
Copyright 2019 The Tekton 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 resources

import (
	"testing"

	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestValidateParamTypesMatching_Valid(t *testing.T) {

	stringValue := *v1beta1.NewArrayOrString("stringValue")
	arrayValue := *v1beta1.NewArrayOrString("arrayValue", "arrayValue")

	for _, tc := range []struct {
		name        string
		description string
		pp          []v1beta1.ParamSpec
		prp         []v1beta1.Param
	}{{
		name: "proper param types",
		pp: []v1beta1.ParamSpec{
			{Name: "correct-type-1", Type: v1beta1.ParamTypeString},
			{Name: "correct-type-2", Type: v1beta1.ParamTypeArray},
		},
		prp: []v1beta1.Param{
			{Name: "correct-type-1", Value: stringValue},
			{Name: "correct-type-2", Value: arrayValue},
		},
	}, {
		name: "no params to get wrong",
		pp:   []v1beta1.ParamSpec{},
		prp:  []v1beta1.Param{},
	}} {
		t.Run(tc.name, func(t *testing.T) {
			ps := &v1beta1.PipelineSpec{Params: tc.pp}
			pr := &v1beta1.PipelineRun{
				ObjectMeta: metav1.ObjectMeta{Name: "pipeline"},
				Spec:       v1beta1.PipelineRunSpec{Params: tc.prp},
			}

			if err := ValidateParamTypesMatching(ps, pr); err != nil {
				t.Errorf("Pipeline.Validate() returned error: %v", err)
			}
		})
	}
}

func TestValidateParamTypesMatching_Invalid(t *testing.T) {

	stringValue := *v1beta1.NewArrayOrString("stringValue")
	arrayValue := *v1beta1.NewArrayOrString("arrayValue", "arrayValue")

	for _, tc := range []struct {
		name        string
		description string
		pp          []v1beta1.ParamSpec
		prp         []v1beta1.Param
	}{{
		name: "string-array mismatch",
		pp: []v1beta1.ParamSpec{
			{Name: "correct-type-1", Type: v1beta1.ParamTypeString},
			{Name: "correct-type-2", Type: v1beta1.ParamTypeArray},
			{Name: "incorrect-type", Type: v1beta1.ParamTypeString},
		},
		prp: []v1beta1.Param{
			{Name: "correct-type-1", Value: stringValue},
			{Name: "correct-type-2", Value: arrayValue},
			{Name: "incorrect-type", Value: arrayValue},
		},
	}, {
		name: "array-string mismatch",
		pp: []v1beta1.ParamSpec{
			{Name: "correct-type-1", Type: v1beta1.ParamTypeString},
			{Name: "correct-type-2", Type: v1beta1.ParamTypeArray},
			{Name: "incorrect-type", Type: v1beta1.ParamTypeArray},
		},
		prp: []v1beta1.Param{
			{Name: "correct-type-1", Value: stringValue},
			{Name: "correct-type-2", Value: arrayValue},
			{Name: "incorrect-type", Value: stringValue},
		},
	}} {
		t.Run(tc.name, func(t *testing.T) {
			ps := &v1beta1.PipelineSpec{Params: tc.pp}
			pr := &v1beta1.PipelineRun{
				ObjectMeta: metav1.ObjectMeta{Name: "pipeline"},
				Spec:       v1beta1.PipelineRunSpec{Params: tc.prp},
			}

			if err := ValidateParamTypesMatching(ps, pr); err == nil {
				t.Errorf("Expected to see error when validating PipelineRun/Pipeline param types but saw none")
			}
		})
	}
}

func TestValidateRequiredParametersProvided_Valid(t *testing.T) {

	stringValue := *v1beta1.NewArrayOrString("stringValue")
	arrayValue := *v1beta1.NewArrayOrString("arrayValue", "arrayValue")

	for _, tc := range []struct {
		name        string
		description string
		pp          []v1beta1.ParamSpec
		prp         []v1beta1.Param
	}{{
		name: "required string params provided",
		pp: []v1beta1.ParamSpec{
			{Name: "required-string-param", Type: v1beta1.ParamTypeString},
		},
		prp: []v1beta1.Param{
			{Name: "required-string-param", Value: stringValue},
		},
	}, {
		name: "required array params provided",
		pp: []v1beta1.ParamSpec{
			{Name: "required-array-param", Type: v1beta1.ParamTypeArray},
		},
		prp: []v1beta1.Param{
			{Name: "required-array-param", Value: arrayValue},
		},
	}, {
		name: "string params provided in default",
		pp: []v1beta1.ParamSpec{
			{Name: "string-param", Type: v1beta1.ParamTypeString, Default: &stringValue},
		},
		prp: []v1beta1.Param{
			{Name: "another-string-param", Value: stringValue},
		},
	}} {
		t.Run(tc.name, func(t *testing.T) {
			if err := ValidateRequiredParametersProvided(&tc.pp, &tc.prp); err != nil {
				t.Errorf("Didn't expect to see error when validating valid PipelineRun parameters but got: %v", err)
			}
		})
	}
}

func TestValidateRequiredParametersProvided_Invalid(t *testing.T) {

	stringValue := *v1beta1.NewArrayOrString("stringValue")
	arrayValue := *v1beta1.NewArrayOrString("arrayValue", "arrayValue")

	for _, tc := range []struct {
		name        string
		description string
		pp          []v1beta1.ParamSpec
		prp         []v1beta1.Param
	}{{
		name: "required string param missing",
		pp: []v1beta1.ParamSpec{
			{Name: "required-string-param", Type: v1beta1.ParamTypeString},
		},
		prp: []v1beta1.Param{
			{Name: "another-string-param", Value: stringValue},
		},
	}, {
		name: "required array param missing",
		pp: []v1beta1.ParamSpec{
			{Name: "required-array-param", Type: v1beta1.ParamTypeArray},
		},
		prp: []v1beta1.Param{
			{Name: "another-array-param", Value: arrayValue},
		},
	}} {
		t.Run(tc.name, func(t *testing.T) {
			if err := ValidateRequiredParametersProvided(&tc.pp, &tc.prp); err == nil {
				t.Errorf("Expected to see error when validating invalid PipelineRun parameters but saw none")
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"html/template"
	"koding/artifact"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/go-webserver/templates"
	"koding/tools/config"
	"net/http"
	"runtime"
	"time"

	"github.com/koding/logging"
)

var (
	Name        = "gowebserver"
	flagConfig  = flag.String("c", "", "Configuration profile from file")
	log         = logging.NewLogger(Name)
	kodingGroup *models.Group
	conf        *config.Config
)

type HomeContent struct {
	Version       string
	Runtime       config.RuntimeOptions
	User          LoggedInUser
	Title         string
	Description   string
	ShareUrl      string
	Impersonating bool
}

type LoggedInUser struct {
	Account       *models.Account
	Machines      []*modelhelper.MachineContainer
	Workspaces    []*models.Workspace
	Group         *models.Group
	Username      string
	SessionId     string
	Impersonating bool
}

func initialize() {
	runtime.GOMAXPROCS(runtime.NumCPU() - 1)

	flag.Parse()
	if *flagConfig == "" {
		log.Critical("Please define config file with -c")
	}

	conf = config.MustConfig(*flagConfig)
	modelhelper.Initialize(conf.Mongo)

	var err error
	kodingGroup, err = modelhelper.GetGroup("koding")
	if err != nil {
		log.Critical("Couldn't fetching `koding` group: %v", err)
		panic(err)
	}
}

func main() {
	initialize()

	url := fmt.Sprintf(":%d", conf.Gowebserver.Port)

	log.Info("Starting gowebserver on %v", url)

	http.HandleFunc("/", HomeHandler)
	http.HandleFunc("/version", artifact.VersionHandler())
	http.HandleFunc("/healthCheck", artifact.HealthCheckHandler(Name))
	http.ListenAndServe(url, nil)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	cookie, err := r.Cookie("clientId")
	if err != nil {
		if err != http.ErrNoCookie {
			log.Error("Couldn't fetch the cookie: %s", err)
		}
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	if cookie.Value == "" {
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	clientId := cookie.Value

	session, err := modelhelper.GetSession(clientId)
	if err != nil {
		log.Error("Couldn't fetch session with clientId %s: %s", clientId, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w) // TODO: clean up session

		return
	}

	username := session.Username
	if username == "" {
		log.Error("Username is empty for session with clientId: %s", clientId)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)

		return
	}

	//----------------------------------------------------------
	// Account
	//----------------------------------------------------------

	account, err := modelhelper.GetAccount(username)
	if err != nil {
		log.Error("Couldn't fetch account with username %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)

		return
	}

	if account.Type != "registered" {
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	//----------------------------------------------------------
	// Machines
	//----------------------------------------------------------

	user, err := modelhelper.GetUser(username)
	if err != nil {
		log.Error("Couldn't get user of %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)
	}

	machines, err := modelhelper.GetMachines(user.ObjectId)
	if err != nil {
		log.Error("Couldn't fetch machines: %s", err)
		machines = []*modelhelper.MachineContainer{}
	}

	//----------------------------------------------------------
	// Workspaces
	//----------------------------------------------------------

	workspaces, err := modelhelper.GetWorkspaces(account.Id)
	if err != nil {
		log.Error("Couldn't fetch workspaces: %s", err)
		workspaces = []*models.Workspace{}
	}

	loggedInUser := LoggedInUser{
		SessionId:     clientId,
		Group:         kodingGroup,
		Workspaces:    workspaces,
		Machines:      machines,
		Account:       account,
		Username:      username,
		Impersonating: session.Impersonating,
	}

	renderLoggedInHome(w, loggedInUser)

	log.Info("loggedin page took: %s", time.Since(start))
}

func renderLoggedInHome(w http.ResponseWriter, u LoggedInUser) {
	homeTmpl := buildHomeTemplate(templates.LoggedInHome)

	hc := buildHomeContent()
	hc.Runtime = conf.Client.RuntimeOptions
	hc.User = u

	if u.Impersonating {
		hc.Impersonating = true
	}

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedin page: %s", err)
		renderLoggedOutHome(w)

		return
	}

	fmt.Fprintf(w, buf.String())
}

func renderLoggedOutHome(w http.ResponseWriter) {
	homeTmpl := buildHomeTemplate(templates.LoggedOutHome)

	hc := buildHomeContent()

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedout page: %s", err)
	}

	fmt.Fprintf(w, buf.String())
}

func buildHomeContent() HomeContent {
	hc := HomeContent{
		Version:  conf.Version,
		ShareUrl: conf.Client.RuntimeOptions.MainUri,
	}
	hc.Title = "Koding | Say goodbye to your localhost and write code in the cloud"
	hc.Description = "Koding is a cloud-based development environment complete with free VMs, IDE & sudo enabled terminal where you can learn Ruby, Go,  Java, NodeJS, PHP, C, C++, Perl, Python, etc."

	return hc
}

func buildHomeTemplate(content string) *template.Template {
	homeTmpl := template.Must(template.New("home").Parse(content))
	headerTmpl := template.Must(template.New("header").Parse(templates.Header))
	analyticsTmpl := template.Must(template.New("analytics").Parse(templates.Analytics))

	homeTmpl.AddParseTree("header", headerTmpl.Tree)
	homeTmpl.AddParseTree("analytics", analyticsTmpl.Tree)

	return homeTmpl
}
<commit_msg>go-webserver: minor fix<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"html/template"
	"koding/artifact"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/go-webserver/templates"
	"koding/tools/config"
	"net/http"
	"runtime"
	"time"

	"github.com/koding/logging"
)

var (
	Name        = "gowebserver"
	flagConfig  = flag.String("c", "", "Configuration profile from file")
	log         = logging.NewLogger(Name)
	kodingGroup *models.Group
	conf        *config.Config
)

type HomeContent struct {
	Version       string
	Runtime       config.RuntimeOptions
	User          LoggedInUser
	Title         string
	Description   string
	ShareUrl      string
	Impersonating bool
}

type LoggedInUser struct {
	Account       *models.Account
	Machines      []*modelhelper.MachineContainer
	Workspaces    []*models.Workspace
	Group         *models.Group
	Username      string
	SessionId     string
	Impersonating bool
}

func initialize() {
	runtime.GOMAXPROCS(runtime.NumCPU() - 1)

	flag.Parse()
	if *flagConfig == "" {
		log.Critical("Please define config file with -c")
	}

	conf = config.MustConfig(*flagConfig)
	modelhelper.Initialize(conf.Mongo)

	var err error
	kodingGroup, err = modelhelper.GetGroup("koding")
	if err != nil {
		log.Critical("Couldn't fetching `koding` group: %v", err)
		panic(err)
	}
}

func main() {
	initialize()

	url := fmt.Sprintf(":%d", conf.Gowebserver.Port)

	log.Info("Starting gowebserver on %v", url)

	http.HandleFunc("/", HomeHandler)
	http.HandleFunc("/version", artifact.VersionHandler())
	http.HandleFunc("/healthCheck", artifact.HealthCheckHandler(Name))
	http.ListenAndServe(url, nil)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	cookie, err := r.Cookie("clientId")
	if err != nil {
		if err != http.ErrNoCookie {
			log.Error("Couldn't fetch the cookie: %s", err)
		}
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	if cookie.Value == "" {
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	clientId := cookie.Value

	session, err := modelhelper.GetSession(clientId)
	if err != nil {
		log.Error("Couldn't fetch session with clientId %s: %s", clientId, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w) // TODO: clean up session

		return
	}

	username := session.Username
	if username == "" {
		log.Error("Username is empty for session with clientId: %s", clientId)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)

		return
	}

	//----------------------------------------------------------
	// Account
	//----------------------------------------------------------

	account, err := modelhelper.GetAccount(username)
	if err != nil {
		log.Error("Couldn't fetch account with username %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)

		return
	}

	if account.Type != "registered" {
		log.Info("loggedout page took: %s", time.Since(start))
		renderLoggedOutHome(w)

		return
	}

	//----------------------------------------------------------
	// Machines
	//----------------------------------------------------------

	user, err := modelhelper.GetUser(username)
	if err != nil {
		log.Error("Couldn't get user of %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		renderLoggedOutHome(w)
	}

	machines, err := modelhelper.GetMachines(user.ObjectId)
	if err != nil {
		log.Error("Couldn't fetch machines: %s", err)
		machines = []*modelhelper.MachineContainer{}
	}

	//----------------------------------------------------------
	// Workspaces
	//----------------------------------------------------------

	workspaces, err := modelhelper.GetWorkspaces(account.Id)
	if err != nil {
		log.Error("Couldn't fetch workspaces: %s", err)
		workspaces = []*models.Workspace{}
	}

	loggedInUser := LoggedInUser{
		SessionId:     clientId,
		Group:         kodingGroup,
		Workspaces:    workspaces,
		Machines:      machines,
		Account:       account,
		Username:      username,
		Impersonating: session.Impersonating,
	}

	renderLoggedInHome(w, loggedInUser)

	log.Info("loggedin page took: %s", time.Since(start))
}

func renderLoggedInHome(w http.ResponseWriter, u LoggedInUser) {
	homeTmpl := buildHomeTemplate(templates.LoggedInHome)

	hc := buildHomeContent()
	hc.Runtime = conf.Client.RuntimeOptions
	hc.User = u
	hc.Impersonating = u.Impersonating

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedin page: %s", err)
		renderLoggedOutHome(w)

		return
	}

	fmt.Fprintf(w, buf.String())
}

func renderLoggedOutHome(w http.ResponseWriter) {
	homeTmpl := buildHomeTemplate(templates.LoggedOutHome)

	hc := buildHomeContent()

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedout page: %s", err)
	}

	fmt.Fprintf(w, buf.String())
}

func buildHomeContent() HomeContent {
	hc := HomeContent{
		Version:  conf.Version,
		ShareUrl: conf.Client.RuntimeOptions.MainUri,
	}
	hc.Title = "Koding | Say goodbye to your localhost and write code in the cloud"
	hc.Description = "Koding is a cloud-based development environment complete with free VMs, IDE & sudo enabled terminal where you can learn Ruby, Go,  Java, NodeJS, PHP, C, C++, Perl, Python, etc."

	return hc
}

func buildHomeTemplate(content string) *template.Template {
	homeTmpl := template.Must(template.New("home").Parse(content))
	headerTmpl := template.Must(template.New("header").Parse(templates.Header))
	analyticsTmpl := template.Must(template.New("analytics").Parse(templates.Analytics))

	homeTmpl.AddParseTree("header", headerTmpl.Tree)
	homeTmpl.AddParseTree("analytics", analyticsTmpl.Tree)

	return homeTmpl
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"html/template"
	"koding/artifact"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/go-webserver/templates"
	"koding/tools/config"
	"net/http"
	"runtime"
	"time"

	"github.com/koding/logging"
)

var (
	Name              = "gowebserver"
	kodingTitle       = "Koding | Say goodbye to your localhost and code in the cloud."
	kodingDescription = "Koding is a cloud-based development environment complete with free VMs, IDE & sudo enabled terminal where you can learn Ruby, Go, Java, NodeJS, PHP, C, C++, Perl, Python, etc."
	kodingShareUrl    = "https://koding.com"
	kodingGpImage     = "koding.com/a/site.landing/images/share.g+.jpg"
	kodingFbImage     = "koding.com/a/site.landing/images/share.fb.jpg"
	kodingTwImage     = "koding.com/a/site.landing/images/share.tw.jpg"
	flagConfig        = flag.String("c", "dev", "Configuration profile from file")
	log               = logging.NewLogger(Name)

	kodingGroup *models.Group
	conf        *config.Config
)

type HomeContent struct {
	Version       string
	Runtime       config.RuntimeOptions
	User          LoggedInUser
	Title         string
	Description   string
	ShareUrl      string
	GpImage       string
	FbImage       string
	TwImage       string
	Impersonating bool
}

type LoggedInUser struct {
	Account       *models.Account
	Machines      []*modelhelper.MachineContainer
	Workspaces    []*models.Workspace
	Group         *models.Group
	Username      string
	SessionId     string
	Impersonating bool
}

func initialize() {
	runtime.GOMAXPROCS(runtime.NumCPU() - 1)

	flag.Parse()
	if *flagConfig == "" {
		log.Critical("Please define config file with -c")
	}

	conf = config.MustConfig(*flagConfig)
	modelhelper.Initialize(conf.Mongo)

	var err error
	kodingGroup, err = modelhelper.GetGroup("koding")
	if err != nil {
		log.Critical("Couldn't fetching `koding` group: %v", err)
		panic(err)
	}
}

func main() {
	initialize()

	http.HandleFunc("/", HomeHandler)
	http.HandleFunc("/version", artifact.VersionHandler())
	http.HandleFunc("/healthCheck", artifact.HealthCheckHandler(Name))

	url := fmt.Sprintf(":%d", conf.Gowebserver.Port)
	log.Info("Starting gowebserver on: %v", url)

	http.ListenAndServe(url, nil)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	cookie, err := r.Cookie("clientId")
	if err != nil {
		if err != http.ErrNoCookie {
			log.Error("Couldn't fetch 'clientId' cookie value: %s", err)
		}

		log.Info("loggedout page took: %s", time.Since(start))
		writeLoggedOutHomeToResp(w)

		return
	}

	if cookie.Value == "" {
		log.Info("loggedout page took: %s", time.Since(start))
		writeLoggedOutHomeToResp(w)

		return
	}

	clientId := cookie.Value

	session, err := modelhelper.GetSession(clientId)
	if err != nil {
		log.Error("Couldn't fetch session with clientId %s: %s", clientId, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	username := session.Username
	if username == "" {
		log.Error("Username is empty for session with clientId: %s", clientId)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	//----------------------------------------------------------
	// Account
	//----------------------------------------------------------

	account, err := modelhelper.GetAccount(username)
	if err != nil {
		log.Error("Couldn't fetch account with username %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	if account.Type != "registered" {
		log.Error(
			"Account type: %s is not 'registered' for %s's session.",
			account.Type, username,
		)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	//----------------------------------------------------------
	// Machines
	//----------------------------------------------------------

	user, err := modelhelper.GetUser(username)
	if err != nil {
		log.Error("Couldn't get user of %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)
	}

	machines, err := modelhelper.GetMachines(user.ObjectId)
	if err != nil {
		log.Error("Couldn't fetch machines: %s", err)
		machines = []*modelhelper.MachineContainer{}
	}

	//----------------------------------------------------------
	// Workspaces
	//----------------------------------------------------------

	workspaces, err := modelhelper.GetWorkspaces(account.Id)
	if err != nil {
		log.Error("Couldn't fetch workspaces: %s", err)
		workspaces = []*models.Workspace{}
	}

	loggedInUser := LoggedInUser{
		SessionId:     clientId,
		Group:         kodingGroup,
		Workspaces:    workspaces,
		Machines:      machines,
		Account:       account,
		Username:      username,
		Impersonating: session.Impersonating,
	}

	writeLoggedInHomeToResp(w, loggedInUser)

	log.Info("loggedin page took: %s", time.Since(start))
}

func writeLoggedInHomeToResp(w http.ResponseWriter, u LoggedInUser) {
	homeTmpl := buildHomeTemplate(templates.LoggedInHome)

	hc := buildHomeContent()
	hc.Runtime = conf.Client.RuntimeOptions
	hc.User = u
	hc.Impersonating = u.Impersonating

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedin page: %s", err)
		writeLoggedOutHomeToResp(w)

		return
	}

	fmt.Fprintf(w, buf.String())
}

func writeLoggedOutHomeToResp(w http.ResponseWriter) {
	homeTmpl := buildHomeTemplate(templates.LoggedOutHome)

	hc := buildHomeContent()

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedout page: %s", err)
	}

	fmt.Fprintf(w, buf.String())
}

func buildHomeContent() HomeContent {
	hc := HomeContent{
		Version:     conf.Version,
		ShareUrl:    kodingShareUrl,
		Title:       kodingTitle,
		Description: kodingDescription,
		GpImage:     kodingGpImage,
		FbImage:     kodingFbImage,
		TwImage:     kodingTwImage,
	}

	return hc
}

func buildHomeTemplate(content string) *template.Template {
	homeTmpl := template.Must(template.New("home").Parse(content))
	headerTmpl := template.Must(template.New("header").Parse(templates.Header))
	analyticsTmpl := template.Must(template.New("analytics").Parse(templates.Analytics))

	homeTmpl.AddParseTree("header", headerTmpl.Tree)
	homeTmpl.AddParseTree("analytics", analyticsTmpl.Tree)

	return homeTmpl
}

func expireCookie(w http.ResponseWriter, cookie *http.Cookie) {
	cookie.Expires = time.Now()
	http.SetCookie(w, cookie)
}
<commit_msg>Go-webserver: HomeHandler: on each request update clientIp in JSession<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"html/template"
	"koding/artifact"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/go-webserver/templates"
	"koding/go-webserver/utils"
	"koding/tools/config"
	"net/http"
	"runtime"
	"time"

	"github.com/koding/logging"
)

var (
	Name              = "gowebserver"
	kodingTitle       = "Koding | Say goodbye to your localhost and code in the cloud."
	kodingDescription = "Koding is a cloud-based development environment complete with free VMs, IDE & sudo enabled terminal where you can learn Ruby, Go, Java, NodeJS, PHP, C, C++, Perl, Python, etc."
	kodingShareUrl    = "https://koding.com"
	kodingGpImage     = "koding.com/a/site.landing/images/share.g+.jpg"
	kodingFbImage     = "koding.com/a/site.landing/images/share.fb.jpg"
	kodingTwImage     = "koding.com/a/site.landing/images/share.tw.jpg"
	flagConfig        = flag.String("c", "dev", "Configuration profile from file")
	log               = logging.NewLogger(Name)

	kodingGroup *models.Group
	conf        *config.Config
)

type HomeContent struct {
	Version       string
	Runtime       config.RuntimeOptions
	User          LoggedInUser
	Title         string
	Description   string
	ShareUrl      string
	GpImage       string
	FbImage       string
	TwImage       string
	Impersonating bool
}

type LoggedInUser struct {
	Account       *models.Account
	Machines      []*modelhelper.MachineContainer
	Workspaces    []*models.Workspace
	Group         *models.Group
	Username      string
	SessionId     string
	Impersonating bool
}

func initialize() {
	runtime.GOMAXPROCS(runtime.NumCPU() - 1)

	flag.Parse()
	if *flagConfig == "" {
		log.Critical("Please define config file with -c")
	}

	conf = config.MustConfig(*flagConfig)
	modelhelper.Initialize(conf.Mongo)

	var err error
	kodingGroup, err = modelhelper.GetGroup("koding")
	if err != nil {
		log.Critical("Couldn't fetching `koding` group: %v", err)
		panic(err)
	}
}

func main() {
	initialize()

	http.HandleFunc("/", HomeHandler)
	http.HandleFunc("/version", artifact.VersionHandler())
	http.HandleFunc("/healthCheck", artifact.HealthCheckHandler(Name))

	url := fmt.Sprintf(":%d", conf.Gowebserver.Port)
	log.Info("Starting gowebserver on: %v", url)

	http.ListenAndServe(url, nil)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	start := time.Now()

	cookie, err := r.Cookie("clientId")
	if err != nil {
		if err != http.ErrNoCookie {
			log.Error("Couldn't fetch 'clientId' cookie value: %s", err)
		}

		log.Info("loggedout page took: %s", time.Since(start))
		writeLoggedOutHomeToResp(w)

		return
	}

	if cookie.Value == "" {
		log.Info("loggedout page took: %s", time.Since(start))
		writeLoggedOutHomeToResp(w)

		return
	}

	clientId := cookie.Value

	session, err := modelhelper.GetSession(clientId)
	if err != nil {
		log.Error("Couldn't fetch session with clientId %s: %s", clientId, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	modelhelper.UpdateSessionIP(clientId, utils.GetIpAddress(r))

	username := session.Username
	if username == "" {
		log.Error("Username is empty for session with clientId: %s", clientId)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	//----------------------------------------------------------
	// Account
	//----------------------------------------------------------

	account, err := modelhelper.GetAccount(username)
	if err != nil {
		log.Error("Couldn't fetch account with username %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	if account.Type != "registered" {
		log.Error(
			"Account type: %s is not 'registered' for %s's session.",
			account.Type, username,
		)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)

		return
	}

	//----------------------------------------------------------
	// Machines
	//----------------------------------------------------------

	user, err := modelhelper.GetUser(username)
	if err != nil {
		log.Error("Couldn't get user of %s: %s", username, err)
		log.Info("loggedout page took: %s", time.Since(start))

		expireCookie(w, cookie)
		writeLoggedOutHomeToResp(w)
	}

	machines, err := modelhelper.GetMachines(user.ObjectId)
	if err != nil {
		log.Error("Couldn't fetch machines: %s", err)
		machines = []*modelhelper.MachineContainer{}
	}

	//----------------------------------------------------------
	// Workspaces
	//----------------------------------------------------------

	workspaces, err := modelhelper.GetWorkspaces(account.Id)
	if err != nil {
		log.Error("Couldn't fetch workspaces: %s", err)
		workspaces = []*models.Workspace{}
	}

	loggedInUser := LoggedInUser{
		SessionId:     clientId,
		Group:         kodingGroup,
		Workspaces:    workspaces,
		Machines:      machines,
		Account:       account,
		Username:      username,
		Impersonating: session.Impersonating,
	}

	writeLoggedInHomeToResp(w, loggedInUser)

	log.Info("loggedin page took: %s", time.Since(start))
}

func writeLoggedInHomeToResp(w http.ResponseWriter, u LoggedInUser) {
	homeTmpl := buildHomeTemplate(templates.LoggedInHome)

	hc := buildHomeContent()
	hc.Runtime = conf.Client.RuntimeOptions
	hc.User = u
	hc.Impersonating = u.Impersonating

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedin page: %s", err)
		writeLoggedOutHomeToResp(w)

		return
	}

	fmt.Fprintf(w, buf.String())
}

func writeLoggedOutHomeToResp(w http.ResponseWriter) {
	homeTmpl := buildHomeTemplate(templates.LoggedOutHome)

	hc := buildHomeContent()

	var buf bytes.Buffer
	if err := homeTmpl.Execute(&buf, hc); err != nil {
		log.Error("Failed to render loggedout page: %s", err)
	}

	fmt.Fprintf(w, buf.String())
}

func buildHomeContent() HomeContent {
	hc := HomeContent{
		Version:     conf.Version,
		ShareUrl:    kodingShareUrl,
		Title:       kodingTitle,
		Description: kodingDescription,
		GpImage:     kodingGpImage,
		FbImage:     kodingFbImage,
		TwImage:     kodingTwImage,
	}

	return hc
}

func buildHomeTemplate(content string) *template.Template {
	homeTmpl := template.Must(template.New("home").Parse(content))
	headerTmpl := template.Must(template.New("header").Parse(templates.Header))
	analyticsTmpl := template.Must(template.New("analytics").Parse(templates.Analytics))

	homeTmpl.AddParseTree("header", headerTmpl.Tree)
	homeTmpl.AddParseTree("analytics", analyticsTmpl.Tree)

	return homeTmpl
}

func expireCookie(w http.ResponseWriter, cookie *http.Cookie) {
	cookie.Expires = time.Now()
	http.SetCookie(w, cookie)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/db"
	"koding/virt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"math"
	"net"
	"net/http"
	"net/url"
	"strings"
)

type UserInfo struct {
	Domain   *proxyconfig.Domain
	IP       string
	Country  string
	Target   *url.URL
	Redirect bool
}

func NewUserInfo(domain *proxyconfig.Domain) *UserInfo {
	return &UserInfo{
		Domain: domain,
	}
}

func populateUser(outreq *http.Request) (*UserInfo, error) {
	user, err := parseDomain(outreq.Host)
	if err != nil {
		return nil, err
	}

	host, err := user.populateIP(outreq.RemoteAddr)
	if err != nil {
		fmt.Println(err)
	} else {
		user.populateCountry(host)
	}

	err = user.populateTarget()
	if err != nil {
		return nil, err
	}

	fmt.Printf("--\nmode '%s'\t: %s %s\n", user.Domain.Proxy.Mode, user.IP, user.Country)
	return user, nil
}

func (u *UserInfo) populateIP(remoteAddr string) (string, error) {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		fmt.Printf("could not split host and port: %s", err.Error())
		return "", err
	}
	u.IP = host
	return host, nil
}

func (u *UserInfo) populateCountry(host string) {
	if geoIP != nil {
		loc := geoIP.GetLocationByIP(host)
		if loc != nil {
			u.Country = loc.CountryName
		}
	}
}

func (u *UserInfo) populateTarget() error {
	var err error
	var hostname string

	username := u.Domain.Proxy.Username
	servicename := u.Domain.Proxy.Servicename
	key := u.Domain.Proxy.Key
	fullurl := u.Domain.Proxy.FullUrl

	switch u.Domain.Proxy.Mode {
	case "redirect":
		u.Target, err = url.Parse("http://" + fullurl)
		if err != nil {
			return err
		}
		u.Redirect = true
		return nil
	case "vm":
		switch u.Domain.LoadBalancer.Mode {
		case "roundrobin":
			N := float64(len(u.Domain.HostnameAlias))
			n := int(math.Mod(float64(u.Domain.LoadBalancer.Index+1), N))
			hostname = u.Domain.HostnameAlias[n]

			u.Domain.LoadBalancer.Index = n
			go proxyDB.UpdateDomain(u.Domain)
		case "sticky":
			hostname = u.Domain.HostnameAlias[u.Domain.LoadBalancer.Index]
		case "default":
			hostname = u.Domain.HostnameAlias[0]
		}

		var vm virt.VM
		if err := db.VMs.Find(bson.M{"hostnameAlias": hostname}).One(&vm); err != nil {
			u.Target, _ = url.Parse("http://www.koding.com/notfound.html")
			u.Redirect = true
			return nil
		}
		if vm.IP == nil {
			u.Target, _ = url.Parse("http://www.koding.com/notactive.html")
			u.Redirect = true
			return nil
		}
		u.Target, err = url.Parse("http://" + vm.IP.String())
		if err != nil {
			return err
		}

		return nil
	case "internal":
		break // internal is done below
	}

	keyData, err := proxyDB.GetKey(username, servicename, key)
	if err != nil {
		return fmt.Errorf("no keyData for username '%s', servicename '%s' and key '%s'", username, servicename, key)
	}

	switch keyData.Mode {
	case "roundrobin":
		N := float64(len(keyData.Host))
		n := int(math.Mod(float64(keyData.CurrentIndex+1), N))
		hostname = keyData.Host[n]

		keyData.CurrentIndex = n
		go proxyDB.UpdateKeyData(username, servicename, keyData)
	case "sticky":
		hostname = keyData.Host[keyData.CurrentIndex]
	}

	u.Target, err = url.Parse("http://" + hostname)
	if err != nil {
		return err
	}

	u.Redirect = false

	return nil
}

func parseDomain(host string) (*UserInfo, error) {
	// Then make a lookup for domains
	domain, err := proxyDB.GetDomain(host)
	if err != nil {
		if err != mgo.ErrNotFound {
			return &UserInfo{}, fmt.Errorf("domain lookup error '%s'", err)
		}

		// lookup didn't found anything, move on to .x.koding.com domains
		if strings.HasSuffix(host, "x.koding.com") {
			// hostsin form {name}-{key}.kd.io or {name}-{key}.x.koding.com is used by koding
			subdomain := strings.TrimSuffix(host, ".x.koding.com")
			servicename := strings.Split(subdomain, "-")[0]
			key := strings.Split(subdomain, "-")[1]

			domain := proxyconfig.NewDomain(host, "internal", "koding", servicename, key, "", []string{})
			return NewUserInfo(domain), nil
		}

		return &UserInfo{}, fmt.Errorf("domain %s is unknown.", host)
	}

	return NewUserInfo(&domain), nil

	// // Handle kd.io domains first
	// if strings.HasSuffix(host, "kd.io") {
	// 	return NewUserInfo("", "", "", "", "vm", host), nil
	// }

	//
	// 	switch counts := strings.Count(host, "-"); {
	// 	case counts == 1:
	// 		// host is in form {name}-{key}.kd.io, used by koding
	// 		subdomain := strings.TrimSuffix(host, ".kd.io")
	// 		servicename := strings.Split(subdomain, "-")[0]
	// 		key := strings.Split(subdomain, "-")[1]
	//
	// 		return NewUserInfo("koding", servicename, key, "", "internal", host), nil
	// 	case counts > 1:
	// 		// host is in form {name}-{key}-{username}.kd.io, used by users
	// 		firstSub := strings.Split(host, ".")[0]
	//
	// 		partsSecond := strings.SplitN(firstSub, "-", 3)
	// 		servicename := partsSecond[0]
	// 		key := partsSecond[1]
	// 		username := partsSecond[2]
	//
	// 		return NewUserInfo(username, servicename, key, "", "internal", host), nil
	// 	}
	// return &UserInfo{}, fmt.Errorf("no data available for proxy. can't parse domain %s", host)
}

func validate(u *UserInfo) (bool, error) {
	ruleId, err := proxyDB.GetDomainRuleId(u.Domain.Id)
	if err != nil {
		return true, nil //don't block if we don't get a rule (pre-caution))
	}

	rule, err := proxyDB.GetRuleByID(ruleId)
	if err != nil {
		return true, nil //don't block if we don't get a rule (pre-caution))
	}

	return validator(rule, u).AddRules().Check()
}

// func lookupRabbitKey(username, servicename, key string) (string, error) {
// 	res, err := proxyDB.GetKey(username, servicename, key)
// 	if err != nil {
// 		return "", fmt.Errorf("no rabbitkey available for user '%s'\n", username)
// 	}
//
// 	if res.Mode == "roundrobin" {
// 		return "", fmt.Errorf("round-robin is disabled for user %s\n", username)
// 	}
//
// 	if res.RabbitKey == "" {
// 		return "", fmt.Errorf("rabbitkey is empty for user %s\n", username)
// 	}
//
// 	return res.RabbitKey, nil
// }

func isWebsocket(req *http.Request) bool {
	conn_hdr := ""
	conn_hdrs := req.Header["Connection"]
	if len(conn_hdrs) > 0 {
		conn_hdr = conn_hdrs[0]
	}

	upgrade_websocket := false
	if strings.ToLower(conn_hdr) == "upgrade" {
		upgrade_hdrs := req.Header["Upgrade"]
		if len(upgrade_hdrs) > 0 {
			upgrade_websocket = (strings.ToLower(upgrade_hdrs[0]) == "websocket")
		}
	}

	return upgrade_websocket
}

func logDomainRequests(domain string) {
	if domain == "" {
		return
	}

	err := proxyDB.AddDomainRequests(domain)
	if err != nil {
		fmt.Printf("could not add domain statistisitcs for %s\n", err.Error())
	}
}

func logProxyStat(name, country string) {
	err := proxyDB.AddProxyStat(name, country)
	if err != nil {
		fmt.Printf("could not add proxy statistisitcs for %s\n", err.Error())
	}
}

func logDomainDenied(domain, ip, country, reason string) {
	if domain == "" {
		return
	}

	err := proxyDB.AddDomainDenied(domain, ip, country, reason)
	if err != nil {
		fmt.Printf("could not add domain statistisitcs for %s\n", err.Error())
	}
}
<commit_msg>kontrolproxy: enable maintancence mode<commit_after>package main

import (
	"fmt"
	"koding/kontrol/kontrolproxy/proxyconfig"
	"koding/tools/db"
	"koding/virt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"math"
	"net"
	"net/http"
	"net/url"
	"strings"
)

type UserInfo struct {
	Domain   *proxyconfig.Domain
	IP       string
	Country  string
	Target   *url.URL
	Redirect bool
}

func NewUserInfo(domain *proxyconfig.Domain) *UserInfo {
	return &UserInfo{
		Domain: domain,
	}
}

func populateUser(outreq *http.Request) (*UserInfo, error) {
	user, err := parseDomain(outreq.Host)
	if err != nil {
		return nil, err
	}

	host, err := user.populateIP(outreq.RemoteAddr)
	if err != nil {
		fmt.Println(err)
	} else {
		user.populateCountry(host)
	}

	err = user.populateTarget()
	if err != nil {
		return nil, err
	}

	fmt.Printf("--\nmode '%s'\t: %s %s\n", user.Domain.Proxy.Mode, user.IP, user.Country)
	return user, nil
}

func (u *UserInfo) populateIP(remoteAddr string) (string, error) {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		fmt.Printf("could not split host and port: %s", err.Error())
		return "", err
	}
	u.IP = host
	return host, nil
}

func (u *UserInfo) populateCountry(host string) {
	if geoIP != nil {
		loc := geoIP.GetLocationByIP(host)
		if loc != nil {
			u.Country = loc.CountryName
		}
	}
}

func (u *UserInfo) populateTarget() error {
	var err error
	var hostname string

	username := u.Domain.Proxy.Username
	servicename := u.Domain.Proxy.Servicename
	key := u.Domain.Proxy.Key
	fullurl := u.Domain.Proxy.FullUrl

	switch u.Domain.Proxy.Mode {
	case "maintenance":
		return
	case "redirect":
		u.Target, err = url.Parse("http://" + fullurl)
		if err != nil {
			return err
		}
		u.Redirect = true
		return nil
	case "vm":
		switch u.Domain.LoadBalancer.Mode {
		case "roundrobin":
			N := float64(len(u.Domain.HostnameAlias))
			n := int(math.Mod(float64(u.Domain.LoadBalancer.Index+1), N))
			hostname = u.Domain.HostnameAlias[n]

			u.Domain.LoadBalancer.Index = n
			go proxyDB.UpdateDomain(u.Domain)
		case "sticky":
			hostname = u.Domain.HostnameAlias[u.Domain.LoadBalancer.Index]
		case "default":
			hostname = u.Domain.HostnameAlias[0]
		}

		var vm virt.VM
		if err := db.VMs.Find(bson.M{"hostnameAlias": hostname}).One(&vm); err != nil {
			u.Target, _ = url.Parse("http://www.koding.com/notfound.html")
			u.Redirect = true
			return nil
		}
		if vm.IP == nil {
			u.Target, _ = url.Parse("http://www.koding.com/notactive.html")
			u.Redirect = true
			return nil
		}
		u.Target, err = url.Parse("http://" + vm.IP.String())
		if err != nil {
			return err
		}

		return nil
	case "internal":
		break // internal is done below
	}

	keyData, err := proxyDB.GetKey(username, servicename, key)
	if err != nil {
		return fmt.Errorf("no keyData for username '%s', servicename '%s' and key '%s'", username, servicename, key)
	}

	switch keyData.Mode {
	case "roundrobin":
		N := float64(len(keyData.Host))
		n := int(math.Mod(float64(keyData.CurrentIndex+1), N))
		hostname = keyData.Host[n]

		keyData.CurrentIndex = n
		go proxyDB.UpdateKeyData(username, servicename, keyData)
	case "sticky":
		hostname = keyData.Host[keyData.CurrentIndex]
	}

	u.Target, err = url.Parse("http://" + hostname)
	if err != nil {
		return err
	}

	u.Redirect = false

	return nil
}

func parseDomain(host string) (*UserInfo, error) {
	// Then make a lookup for domains
	domain, err := proxyDB.GetDomain(host)
	if err != nil {
		if err != mgo.ErrNotFound {
			return &UserInfo{}, fmt.Errorf("domain lookup error '%s'", err)
		}

		// lookup didn't found anything, move on to .x.koding.com domains
		if strings.HasSuffix(host, "x.koding.com") {
			// hostsin form {name}-{key}.kd.io or {name}-{key}.x.koding.com is used by koding
			subdomain := strings.TrimSuffix(host, ".x.koding.com")
			servicename := strings.Split(subdomain, "-")[0]
			key := strings.Split(subdomain, "-")[1]

			domain := proxyconfig.NewDomain(host, "internal", "koding", servicename, key, "", []string{})
			return NewUserInfo(domain), nil
		}

		return &UserInfo{}, fmt.Errorf("domain %s is unknown.", host)
	}

	return NewUserInfo(&domain), nil

	// // Handle kd.io domains first
	// if strings.HasSuffix(host, "kd.io") {
	// 	return NewUserInfo("", "", "", "", "vm", host), nil
	// }

	//
	// 	switch counts := strings.Count(host, "-"); {
	// 	case counts == 1:
	// 		// host is in form {name}-{key}.kd.io, used by koding
	// 		subdomain := strings.TrimSuffix(host, ".kd.io")
	// 		servicename := strings.Split(subdomain, "-")[0]
	// 		key := strings.Split(subdomain, "-")[1]
	//
	// 		return NewUserInfo("koding", servicename, key, "", "internal", host), nil
	// 	case counts > 1:
	// 		// host is in form {name}-{key}-{username}.kd.io, used by users
	// 		firstSub := strings.Split(host, ".")[0]
	//
	// 		partsSecond := strings.SplitN(firstSub, "-", 3)
	// 		servicename := partsSecond[0]
	// 		key := partsSecond[1]
	// 		username := partsSecond[2]
	//
	// 		return NewUserInfo(username, servicename, key, "", "internal", host), nil
	// 	}
	// return &UserInfo{}, fmt.Errorf("no data available for proxy. can't parse domain %s", host)
}

func validate(u *UserInfo) (bool, error) {
	ruleId, err := proxyDB.GetDomainRuleId(u.Domain.Id)
	if err != nil {
		return true, nil //don't block if we don't get a rule (pre-caution))
	}

	rule, err := proxyDB.GetRuleByID(ruleId)
	if err != nil {
		return true, nil //don't block if we don't get a rule (pre-caution))
	}

	return validator(rule, u).AddRules().Check()
}

// func lookupRabbitKey(username, servicename, key string) (string, error) {
// 	res, err := proxyDB.GetKey(username, servicename, key)
// 	if err != nil {
// 		return "", fmt.Errorf("no rabbitkey available for user '%s'\n", username)
// 	}
//
// 	if res.Mode == "roundrobin" {
// 		return "", fmt.Errorf("round-robin is disabled for user %s\n", username)
// 	}
//
// 	if res.RabbitKey == "" {
// 		return "", fmt.Errorf("rabbitkey is empty for user %s\n", username)
// 	}
//
// 	return res.RabbitKey, nil
// }

func isWebsocket(req *http.Request) bool {
	conn_hdr := ""
	conn_hdrs := req.Header["Connection"]
	if len(conn_hdrs) > 0 {
		conn_hdr = conn_hdrs[0]
	}

	upgrade_websocket := false
	if strings.ToLower(conn_hdr) == "upgrade" {
		upgrade_hdrs := req.Header["Upgrade"]
		if len(upgrade_hdrs) > 0 {
			upgrade_websocket = (strings.ToLower(upgrade_hdrs[0]) == "websocket")
		}
	}

	return upgrade_websocket
}

func logDomainRequests(domain string) {
	if domain == "" {
		return
	}

	err := proxyDB.AddDomainRequests(domain)
	if err != nil {
		fmt.Printf("could not add domain statistisitcs for %s\n", err.Error())
	}
}

func logProxyStat(name, country string) {
	err := proxyDB.AddProxyStat(name, country)
	if err != nil {
		fmt.Printf("could not add proxy statistisitcs for %s\n", err.Error())
	}
}

func logDomainDenied(domain, ip, country, reason string) {
	if domain == "" {
		return
	}

	err := proxyDB.AddDomainDenied(domain, ip, country, reason)
	if err != nil {
		fmt.Printf("could not add domain statistisitcs for %s\n", err.Error())
	}
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
)

func dataSourceGoogleProjects() *schema.Resource {
	return &schema.Resource{
		Read: datasourceGoogleProjectsRead,
		Schema: map[string]*schema.Schema{
			"filter": {
				Type:     schema.TypeString,
				Required: true,
			},
			"projects": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"project_id": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
		},
	}
}

func datasourceGoogleProjectsRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	params := make(map[string]string)

	params["filter"] = d.Get("filter").(string)
	url := "https://cloudresourcemanager.googleapis.com/v1/projects"

	url, err := addQueryParams(url, params)
	if err != nil {
		return err
	}

	res, err := sendRequest(config, "GET", url, nil)
	if err != nil {
		return fmt.Errorf("Error retrieving projects: %s", err)
	}

	if err := d.Set("projects", flattenDatasourceGoogleProjectsProjects(res["projects"], d)); err != nil {
		return fmt.Errorf("Error retrieving projects: %s", err)
	}

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

	return nil
}

func flattenDatasourceGoogleProjectsProjects(v interface{}, d *schema.ResourceData) interface{} {
	if v == nil {
		return v
	}

	l := v.([]interface{})
	transformed := make([]interface{}, 0, len(l))
	for _, raw := range l {
		original := raw.(map[string]interface{})
		if len(original) < 1 {
			// Do not include empty json objects coming back from the api
			continue
		}
		transformed = append(transformed, map[string]interface{}{
			"project_id": original["projectId"],
		})
	}

	return transformed
}
<commit_msg>Add pagination to google projects data source (#626)<commit_after>package google

import (
	"fmt"
	"github.com/hashicorp/terraform/helper/schema"
)

func dataSourceGoogleProjects() *schema.Resource {
	return &schema.Resource{
		Read: datasourceGoogleProjectsRead,
		Schema: map[string]*schema.Schema{
			"filter": {
				Type:     schema.TypeString,
				Required: true,
			},
			"projects": {
				Type:     schema.TypeList,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"project_id": {
							Type:     schema.TypeString,
							Computed: true,
						},
					},
				},
			},
		},
	}
}

func datasourceGoogleProjectsRead(d *schema.ResourceData, meta interface{}) error {
	config := meta.(*Config)

	params := make(map[string]string)
	projects := make([]map[string]interface{}, 0)

	for {
		params["filter"] = d.Get("filter").(string)
		url := "https://cloudresourcemanager.googleapis.com/v1/projects"

		url, err := addQueryParams(url, params)
		if err != nil {
			return err
		}

		res, err := sendRequest(config, "GET", url, nil)
		if err != nil {
			return fmt.Errorf("Error retrieving projects: %s", err)
		}

		pageProjects := flattenDatasourceGoogleProjectsList(res["projects"])
		projects = append(projects, pageProjects...)

		pToken, ok := res["nextPageToken"]
		if ok && pToken != nil && pToken.(string) != "" {
			params["pageToken"] = pToken.(string)
		} else {
			break
		}
	}

	if err := d.Set("projects", projects); err != nil {
		return fmt.Errorf("Error retrieving projects: %s", err)
	}

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

	return nil
}

func flattenDatasourceGoogleProjectsList(v interface{}) []map[string]interface{} {
	if v == nil {
		return make([]map[string]interface{}, 0)
	}

	ls := v.([]interface{})
	projects := make([]map[string]interface{}, 0, len(ls))
	for _, raw := range ls {
		p := raw.(map[string]interface{})
		if pId, ok := p["projectId"]; ok {
			projects = append(projects, map[string]interface{}{
				"project_id": pId,
			})
		}
	}

	return projects
}
<|endoftext|>


<commit_before>/*
	The mongo package is a very simple wrapper around the labix.org/v2/mgo
	package. It's purpose is to allow you to do CRUD operations with very
	little code. It's not exhaustive and not meant to do everything for you.
*/
package mongo

import (
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	"errors"
	"fmt"
	"reflect"
	"time"
)

var (
	mgoSession *mgo.Session
	servers    string
	database   string
	NoPtr      = errors.New("You must pass in a pointer")
)

// Set the mongo servers and the database
func SetServers(servers, db string) error {
	var err error

	database = db

	mgoSession, err = mgo.Dial(servers)
	return err
}

// Insert a single record. Must pass in a pointer to a struct. The struct must
// contain an Id field of type bson.ObjectId.
func Insert(records ...interface{}) error {
	for _, rec := range records {
		if !isPtr(rec) {
			return NoPtr
		}

		if err := addNewFields(rec); err != nil {
			return err
		}

		s, err := GetMongoSession()
		if err != nil {
			return err
		}
		defer s.Close()

		coll := getColl(s, typeName(rec))
		err = coll.Insert(rec)
		if err != nil {
			return err
		}
	}

	return nil
}

// Find one or more records. If a single struct is passed in we'll return one record.
// If a slice is passed in all records will be returned. Must pass in a pointer to a
// struct or slice of structs.
func Find(i interface{}, q bson.M) error {
	if !isPtr(i) {
		return NoPtr
	}

	s, err := GetMongoSession()
	if err != nil {
		return err
	}
	defer s.Close()

	coll := getColl(s, typeName(i))

	query := coll.Find(q)

	if isSlice(reflect.TypeOf(i)) {
		err = query.All(i)
	} else {
		err = query.One(i)
	}
	return err
}

// Find a single record by id. Must pass a pointer to a struct.
func FindById(i interface{}, id string) error {
	return Find(i, bson.M{"_id": id})
}

// Updates a record. Uses the Id to identify the record to update. Must pass in a pointer
// to a struct.
func Update(i interface{}) error {
	if !isPtr(i) {
		return NoPtr
	}

	err := addCurrentDateTime(i, "UpdatedAt")
	if err != nil {
		return err
	}

	s, err := GetMongoSession()
	if err != nil {
		return err
	}
	defer s.Close()

	id, err := getObjIdFromStruct(i)
	if err != nil {
		return err
	}

	return getColl(s, typeName(i)).Update(bson.M{"_id": id}, i)
}

// Deletes a record. Uses the Id to identify the record to delete. Must pass in a pointer
// to a struct.
func Delete(i interface{}) error {
	if !isPtr(i) {
		return NoPtr
	}

	s, err := GetMongoSession()
	if err != nil {
		return err
	}
	defer s.Close()

	id, err := getObjIdFromStruct(i)
	if err != nil {
		return err
	}

	return getColl(s, typeName(i)).RemoveId(id)
}

// Returns a Mongo session. You must call Session.Close() when you're done.
func GetMongoSession() (*mgo.Session, error) {
	var err error

	if mgoSession == nil {
		mgoSession, err = mgo.Dial(servers)
		if err != nil {
			return nil, err
		}
	}

	return mgoSession.Clone(), nil
}

// We pass in the session because that is a clone of the original and the
// caller will need to close it when finished.
func getColl(session *mgo.Session, coll string) *mgo.Collection {
	return session.DB(database).C(coll)
}

func getObjIdFromStruct(i interface{}) (bson.ObjectId, error) {
	v := reflect.ValueOf(i)

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return bson.ObjectId(""), errors.New("Can't delete record. Type must be a struct.")
	}

	f := v.FieldByName("Id")
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	return f.Interface().(bson.ObjectId), nil
}

func isPtr(i interface{}) bool {
	return reflect.ValueOf(i).Kind() == reflect.Ptr
}

func typeName(i interface{}) string {
	t := reflect.TypeOf(i)

	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	if isSlice(t) {
		t = t.Elem()

		if t.Kind() == reflect.Ptr {
			t = t.Elem()
		}
	}

	return t.Name()
}

// returns true if the interface is a slice
func isSlice(t reflect.Type) bool {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return t.Kind() == reflect.Slice
}

func addNewFields(i interface{}) error {
	err := addId(i)
	if err != nil {
		return err
	}

	if err := addCurrentDateTime(i, "CreatedAt"); err != nil {
		return err
	}

	return addCurrentDateTime(i, "UpdatedAt")
}

func addCurrentDateTime(i interface{}, name string) error {
	if !hasStructField(i, name) {
		return nil
	}

	now := time.Now()

	v := reflect.ValueOf(i)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	f := v.FieldByName(name)
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	if reflect.TypeOf(now) != f.Type() {
		return fmt.Errorf("%v must be time.Time type.", name)
	}

	if !f.CanSet() {
		return fmt.Errorf("Couldn't set time for field: %v", name)
	}

	f.Set(reflect.ValueOf(now))

	return nil
}

func hasStructField(i interface{}, field string) bool {
	t := reflect.TypeOf(i)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	if t.Kind() != reflect.Struct {
		return false
	}

	_, found := t.FieldByName(field)
	return found
}

func addId(i interface{}) error {
	v := reflect.ValueOf(i)

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return errors.New("Record must be a struct")
	}

	f := v.FieldByName("Id")
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	if f.Kind() == reflect.String {
		if !f.Interface().(bson.ObjectId).Valid() {
			id := reflect.ValueOf(bson.NewObjectId())
			f.Set(id)
		}
	}

	return nil
}
<commit_msg>Clean up function name.<commit_after>/*
	The mongo package is a very simple wrapper around the labix.org/v2/mgo
	package. It's purpose is to allow you to do CRUD operations with very
	little code. It's not exhaustive and not meant to do everything for you.
*/
package mongo

import (
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	"errors"
	"fmt"
	"reflect"
	"time"
)

var (
	mgoSession *mgo.Session
	servers    string
	database   string
	NoPtr      = errors.New("You must pass in a pointer")
)

// Set the mongo servers and the database
func SetServers(servers, db string) error {
	var err error

	database = db

	mgoSession, err = mgo.Dial(servers)
	return err
}

// Insert a single record. Must pass in a pointer to a struct. The struct must
// contain an Id field of type bson.ObjectId.
func Insert(records ...interface{}) error {
	for _, rec := range records {
		if !isPtr(rec) {
			return NoPtr
		}

		if err := addNewFields(rec); err != nil {
			return err
		}

		s, err := GetSession()
		if err != nil {
			return err
		}
		defer s.Close()

		coll := getColl(s, typeName(rec))
		err = coll.Insert(rec)
		if err != nil {
			return err
		}
	}

	return nil
}

// Find one or more records. If a single struct is passed in we'll return one record.
// If a slice is passed in all records will be returned. Must pass in a pointer to a
// struct or slice of structs.
func Find(i interface{}, q bson.M) error {
	if !isPtr(i) {
		return NoPtr
	}

	s, err := GetSession()
	if err != nil {
		return err
	}
	defer s.Close()

	coll := getColl(s, typeName(i))

	query := coll.Find(q)

	if isSlice(reflect.TypeOf(i)) {
		err = query.All(i)
	} else {
		err = query.One(i)
	}
	return err
}

// Find a single record by id. Must pass a pointer to a struct.
func FindById(i interface{}, id string) error {
	return Find(i, bson.M{"_id": id})
}

// Updates a record. Uses the Id to identify the record to update. Must pass in a pointer
// to a struct.
func Update(i interface{}) error {
	if !isPtr(i) {
		return NoPtr
	}

	err := addCurrentDateTime(i, "UpdatedAt")
	if err != nil {
		return err
	}

	s, err := GetSession()
	if err != nil {
		return err
	}
	defer s.Close()

	id, err := getObjIdFromStruct(i)
	if err != nil {
		return err
	}

	return getColl(s, typeName(i)).Update(bson.M{"_id": id}, i)
}

// Deletes a record. Uses the Id to identify the record to delete. Must pass in a pointer
// to a struct.
func Delete(i interface{}) error {
	if !isPtr(i) {
		return NoPtr
	}

	s, err := GetSession()
	if err != nil {
		return err
	}
	defer s.Close()

	id, err := getObjIdFromStruct(i)
	if err != nil {
		return err
	}

	return getColl(s, typeName(i)).RemoveId(id)
}

// Returns a Mongo session. You must call Session.Close() when you're done.
func GetSession() (*mgo.Session, error) {
	var err error

	if mgoSession == nil {
		mgoSession, err = mgo.Dial(servers)
		if err != nil {
			return nil, err
		}
	}

	return mgoSession.Clone(), nil
}

// We pass in the session because that is a clone of the original and the
// caller will need to close it when finished.
func getColl(session *mgo.Session, coll string) *mgo.Collection {
	return session.DB(database).C(coll)
}

func getObjIdFromStruct(i interface{}) (bson.ObjectId, error) {
	v := reflect.ValueOf(i)

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return bson.ObjectId(""), errors.New("Can't delete record. Type must be a struct.")
	}

	f := v.FieldByName("Id")
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	return f.Interface().(bson.ObjectId), nil
}

func isPtr(i interface{}) bool {
	return reflect.ValueOf(i).Kind() == reflect.Ptr
}

func typeName(i interface{}) string {
	t := reflect.TypeOf(i)

	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	if isSlice(t) {
		t = t.Elem()

		if t.Kind() == reflect.Ptr {
			t = t.Elem()
		}
	}

	return t.Name()
}

// returns true if the interface is a slice
func isSlice(t reflect.Type) bool {
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}
	return t.Kind() == reflect.Slice
}

func addNewFields(i interface{}) error {
	err := addId(i)
	if err != nil {
		return err
	}

	if err := addCurrentDateTime(i, "CreatedAt"); err != nil {
		return err
	}

	return addCurrentDateTime(i, "UpdatedAt")
}

func addCurrentDateTime(i interface{}, name string) error {
	if !hasStructField(i, name) {
		return nil
	}

	now := time.Now()

	v := reflect.ValueOf(i)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	f := v.FieldByName(name)
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	if reflect.TypeOf(now) != f.Type() {
		return fmt.Errorf("%v must be time.Time type.", name)
	}

	if !f.CanSet() {
		return fmt.Errorf("Couldn't set time for field: %v", name)
	}

	f.Set(reflect.ValueOf(now))

	return nil
}

func hasStructField(i interface{}, field string) bool {
	t := reflect.TypeOf(i)
	if t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	if t.Kind() != reflect.Struct {
		return false
	}

	_, found := t.FieldByName(field)
	return found
}

func addId(i interface{}) error {
	v := reflect.ValueOf(i)

	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	if v.Kind() != reflect.Struct {
		return errors.New("Record must be a struct")
	}

	f := v.FieldByName("Id")
	if f.Kind() == reflect.Ptr {
		f = f.Elem()
	}

	if f.Kind() == reflect.String {
		if !f.Interface().(bson.ObjectId).Valid() {
			id := reflect.ValueOf(bson.NewObjectId())
			f.Set(id)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// package monit provides a metric reporting mechanism for webapps
//
// Basic usage:
//
//		m = monit.NewMonitor(monit.Config{
//			Host: "https://myhost.com/reporting/",
//			Base: map[string]interface{}{
//				"auth": "maybeINeedThis?"
//			},
//		})
//		m.Start()
package monit

import (
	"github.com/swhite24/envreader"
	"strconv"
	"time"
	"runtime"
	"encoding/json"
	"bytes"
	"net/http"
	"crypto/tls"
)

var (
	stats runtime.MemStats
	client http.Client
)

type (

	// Config contains configuration values for monit.  Values default to environment
	// variables where appropriate.
	Config struct {

		// Fully qualified url to send monitoring data to.
		// Can be provided in MONIT_HOST environment variable
		Host string

		// Interval, in seconds, at which to report metrics.
		// Can be provided in MONIT_INTERVAL environment variable
		Interval int

		// Base contains key value pairs to include as part of base object reported
		Base map[string]interface{}
	}

	// Monit exposes monitoring func
	Monit struct {
		config Config
		requests int
		cont bool
		start int64
	}
)

// NewMonitor provides an instance of Monit.
//
// Any zero-valued Config properties will use environment variables described above where appropriate.
func NewMonitor (c Config) (m Monit) {
	// Load environment
	vals := envreader.Read("MONIT_HOST", "MONIT_INTERVAL")

	// Check c
	if c.Host == "" { c.Host = vals["MONIT_HOST"] }
	if c.Interval == 0 {
		i, err := strconv.Atoi(vals["MONIT_INTERVAL"])
		if (err != nil) { panic(err) }
		c.Interval = i
	}
	if len(c.Base) == 0 { c.Base = make(map[string]interface{}) }

	m = Monit{ c , 0, true, time.Now().Unix() }

	return m
}

// Start starts a goroutine to report metrics to host based on Config value.
func (m *Monit) Start () {
	go (func (m *Monit) {
		for m.cont {
			// Sleep for specified interval
			time.Sleep(time.Duration(m.config.Interval) * time.Second)

			m.report()

			// Reset
			m.requests = 0
		}
	})(m)
}

func (m Monit) report () {
	// Get current stats
	m.getStat()

	// Get json buffer
	stat, _ := json.Marshal(m.config.Base)
	buf := bytes.NewBuffer(stat)

	// Issue request
	client.Post(m.config.Host, "application/json", buf)
}

func (m Monit) getStat () {
	runtime.ReadMemStats(&stats)

	// Mem_used in MB
	m.config.Base["app_used_memory"] = float64(stats.HeapAlloc) / 1000000
	m.config.Base["uptime"] = time.Now().Unix() - m.start
}

// Stop stops all reporting.  Call Start to begin again.
func (m *Monit) Stop () {
	m.cont = false
}

// Request increments count of requests to report for current interval.
func (m *Monit) Request () {
	m.requests = m.requests + 1
}

func init () {
	// Setup transport to ignore invalid certs
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{ InsecureSkipVerify: true },
	}
	client = http.Client{ Transport: transport }
}<commit_msg>Closing response body<commit_after>// package monit provides a metric reporting mechanism for webapps
//
// Basic usage:
//
//		m = monit.NewMonitor(monit.Config{
//			Host: "https://myhost.com/reporting/",
//			Base: map[string]interface{}{
//				"auth": "maybeINeedThis?"
//			},
//		})
//		m.Start()
package monit

import (
	"github.com/swhite24/envreader"
	"strconv"
	"time"
	"runtime"
	"encoding/json"
	"bytes"
	"net/http"
	"crypto/tls"
)

var (
	client http.Client
)

type (

	// Config contains configuration values for monit.  Values default to environment
	// variables where appropriate.
	Config struct {

		// Fully qualified url to send monitoring data to.
		// Can be provided in MONIT_HOST environment variable
		Host string

		// Interval, in seconds, at which to report metrics.
		// Can be provided in MONIT_INTERVAL environment variable
		Interval int

		// Base contains key value pairs to include as part of base object reported
		Base map[string]interface{}
	}

	// Monit exposes monitoring func
	Monit struct {
		config *Config
		requests int
		cont bool
		start int64
	}
)

// NewMonitor provides an instance of Monit.
//
// Any zero-valued Config properties will use environment variables described above where appropriate.
func NewMonitor (c Config) (m Monit) {
	// Load environment
	vals := envreader.Read("MONIT_HOST", "MONIT_INTERVAL")

	// Check c
	if c.Host == "" { c.Host = vals["MONIT_HOST"] }
	if c.Interval == 0 {
		i, err := strconv.Atoi(vals["MONIT_INTERVAL"])
		if (err != nil) { panic(err) }
		c.Interval = i
	}
	if len(c.Base) == 0 { c.Base = make(map[string]interface{}) }

	m = Monit{ &c , 0, true, time.Now().Unix() }

	return m
}

// Start starts a goroutine to report metrics to host based on Config value.
func (m *Monit) Start () {
	go (func (m *Monit) {
		for m.cont {
			// Sleep for specified interval
			time.Sleep(time.Duration(m.config.Interval) * time.Second)

			m.report()

			// Reset
			m.requests = 0
		}
	})(m)
}

func (m *Monit) report () {
	// Get current stats
	m.getStat()

	// Get json buffer
	stat, _ := json.Marshal(m.config.Base)
	buf := bytes.NewBuffer(stat)
	// Issue request
	r, _ := client.Post(m.config.Host, "application/json", buf)
	defer r.Body.Close()
}

func (m *Monit) getStat () {
	var stats runtime.MemStats
	runtime.ReadMemStats(&stats)

	// Mem_used in MB
	m.config.Base["app_used_memory"] = float64(stats.HeapAlloc) / 1000000
	m.config.Base["uptime"] = time.Now().Unix() - m.start
}

// Stop stops all reporting.  Call Start to begin again.
func (m *Monit) Stop () {
	m.cont = false
}

// Request increments count of requests to report for current interval.
func (m *Monit) Request () {
	m.requests = m.requests + 1
}

func init () {
	// Setup transport to ignore invalid certs
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{ InsecureSkipVerify: true },
	}
	client = http.Client{ Transport: transport }
}<|endoftext|>