<commit_before>package cli

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/pkg/cluster"
)

var flynnHostLogs = map[string]string{
	"flynn-host.log": "/var/log/flynn/flynn-host.log",

	// the following two entries are legacy paths from when flynn-host used
	// to log to stdout (which would be redirected to these files)
	"upstart-flynn-host.log": "/var/log/upstart/flynn-host.log",
	"tmp-flynn-host.log":     "/tmp/flynn-host.log",
}

var debugCmds = [][]string{
	{"ps", "faux"},
	{"ifconfig"},
	{"uname", "-a"},
	{"lsb_release", "-a"},
	{"date"},
	{"free", "-m"},
	{"df", "-h"},
	{"dpkg-query", "-W", "-f", "${Package}: ${Version}\n", "libvirt-bin"},
	{os.Args[0], "version"},
	{"virsh", "-c", "lxc:///", "list"},
	{"virsh", "-c", "lxc:///", "net-list"},
	{"route", "-n"},
	{"iptables", "-L", "-v", "-n", "--line-numbers"},
}

func init() {
	Register("collect-debug-info", runCollectDebugInfo, `
usage: flynn-host collect-debug-info [options]

Options:
  --tarball          Create a tarball instead of uploading to a gist
  --include-env      Include sensitive environment variables
  --job-lines=<num>  Maximum number of lines to fetch from each job [default: 1000]

Collect debug information into an anonymous gist or tarball`)
}

func runCollectDebugInfo(args *docopt.Args) error {
	log := log15.New()
	lines, err := strconv.Atoi(args.String["--job-lines"])
	if err != nil {
		return err
	}

	if args.Bool["--tarball"] {
		log.Info("creating a tarball containing logs and debug information")
	} else {
		log.Info("uploading logs and debug information to a private, anonymous gist")
	}
	log.Info("this may take a while depending on the size of your logs")

	gist := &Gist{
		Description: "Flynn debug information",
		Public:      false,
		Files:       make(map[string]File),
	}

	log.Info("getting flynn-host logs")
	for name, filepath := range flynnHostLogs {
		if err := gist.AddLocalFile(name, filepath); err != nil && !os.IsNotExist(err) {
			log.Error(fmt.Sprintf("error getting flynn-host log %q", filepath), "err", err)
		}
	}

	log.Info("getting job logs")
	if err := captureJobs(gist, args.Bool["--include-env"], lines); err != nil {
		log.Error("error getting job logs, falling back to on-disk logs", "err", err)
		debugCmds = append(debugCmds, []string{"bash", "-c", fmt.Sprintf("tail -n %d /var/log/flynn/*.log", lines)})
	}

	log.Info("getting system information")
	var debugOutput string
	for _, cmd := range debugCmds {
		output, err := captureCmd(cmd[0], cmd[1:]...)
		if err != nil {
			log.Error(fmt.Sprintf("error capturing output of %q", strings.Join(cmd, " ")), "err", err)
			continue
		}
		debugOutput += fmt.Sprintln("===>", strings.Replace(strings.Join(cmd, " "), "\n", `\n`, -1), "\n", output)
	}
	gist.AddFile("0-debug-output.log", debugOutput)

	if args.Bool["--tarball"] {
		path, err := gist.CreateTarball()
		if err != nil {
			log.Error("error creating tarball", "err", err)
			return err
		}
		log.Info(fmt.Sprintf("created tarball containing debug information at %s", path))
		return nil
	}

	if err := gist.Upload(log); err != nil {
		return err
	}

	log.Info(fmt.Sprintf("debug information uploaded to: %s", gist.URL))
	return nil
}

func captureCmd(name string, arg ...string) (string, error) {
	var buf bytes.Buffer
	c := exec.Command(name, arg...)
	c.Stdout = &buf
	if err := c.Run(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func captureJobs(gist *Gist, env bool, lines int) error {
	client := cluster.NewClient()

	jobs, err := jobList(client, true)
	if err != nil {
		return err
	}

	var buf bytes.Buffer
	printJobs(jobs, &buf)
	gist.AddFile("1-jobs.log", buf.String())

	for _, job := range jobs {
		var name string
		if system, ok := job.Job.Metadata["flynn-system-app"]; !ok || system != "true" {
			continue // Skip non-system applications
		}
		if app, ok := job.Job.Metadata["flynn-controller.app_name"]; ok {
			name += app + "-"
		}
		if typ, ok := job.Job.Metadata["flynn-controller.type"]; ok {
			name += typ + "-"
		}
		name += job.Job.ID + ".log"

		var content bytes.Buffer
		printJobDesc(&job, &content, env)
		fmt.Fprint(&content, "\n\n***** ***** ***** ***** ***** ***** ***** ***** ***** *****\n\n")
		stdoutR, stdoutW := io.Pipe()
		stderrR, stderrW := io.Pipe()
		go func() {
			getLog(job.HostID, job.Job.ID, client, false, true, stdoutW, stderrW)
			stdoutW.Close()
			stderrW.Close()
		}()
		tailLogs(stdoutR, stderrR, lines, &content, &content)

		gist.AddFile(name, content.String())
	}

	return nil
}
<commit_msg>host/cli: Collect entire log rather than only 1k lines<commit_after>package cli

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/pkg/cluster"
)

var flynnHostLogs = map[string]string{
	"flynn-host.log": "/var/log/flynn/flynn-host.log",

	// the following two entries are legacy paths from when flynn-host used
	// to log to stdout (which would be redirected to these files)
	"upstart-flynn-host.log": "/var/log/upstart/flynn-host.log",
	"tmp-flynn-host.log":     "/tmp/flynn-host.log",
}

var debugCmds = [][]string{
	{"ps", "faux"},
	{"ifconfig"},
	{"uname", "-a"},
	{"lsb_release", "-a"},
	{"date"},
	{"free", "-m"},
	{"df", "-h"},
	{"dpkg-query", "-W", "-f", "${Package}: ${Version}\n", "libvirt-bin"},
	{os.Args[0], "version"},
	{"virsh", "-c", "lxc:///", "list"},
	{"virsh", "-c", "lxc:///", "net-list"},
	{"route", "-n"},
	{"iptables", "-L", "-v", "-n", "--line-numbers"},
}

func init() {
	Register("collect-debug-info", runCollectDebugInfo, `
usage: flynn-host collect-debug-info [options]

Options:
  --tarball          Create a tarball instead of uploading to a gist
  --include-env      Include sensitive environment variables

Collect debug information into an anonymous gist or tarball`)
}

func runCollectDebugInfo(args *docopt.Args) error {
	log := log15.New()
	if args.Bool["--tarball"] {
		log.Info("creating a tarball containing logs and debug information")
	} else {
		log.Info("uploading logs and debug information to a private, anonymous gist")
	}
	log.Info("this may take a while depending on the size of your logs")

	gist := &Gist{
		Description: "Flynn debug information",
		Public:      false,
		Files:       make(map[string]File),
	}

	log.Info("getting flynn-host logs")
	for name, filepath := range flynnHostLogs {
		if err := gist.AddLocalFile(name, filepath); err != nil && !os.IsNotExist(err) {
			log.Error(fmt.Sprintf("error getting flynn-host log %q", filepath), "err", err)
		}
	}

	log.Info("getting job logs")
	if err := captureJobs(gist, args.Bool["--include-env"]); err != nil {
		log.Error("error getting job logs, falling back to on-disk logs", "err", err)
		debugCmds = append(debugCmds, []string{"bash", "-c", "tail -n+1 /var/log/flynn/*.log"})
	}

	log.Info("getting system information")
	var debugOutput string
	for _, cmd := range debugCmds {
		output, err := captureCmd(cmd[0], cmd[1:]...)
		if err != nil {
			log.Error(fmt.Sprintf("error capturing output of %q", strings.Join(cmd, " ")), "err", err)
			continue
		}
		debugOutput += fmt.Sprintln("===>", strings.Replace(strings.Join(cmd, " "), "\n", `\n`, -1), "\n", output)
	}
	gist.AddFile("0-debug-output.log", debugOutput)

	if args.Bool["--tarball"] {
		path, err := gist.CreateTarball()
		if err != nil {
			log.Error("error creating tarball", "err", err)
			return err
		}
		log.Info(fmt.Sprintf("created tarball containing debug information at %s", path))
		return nil
	}

	if err := gist.Upload(log); err != nil {
		return err
	}

	log.Info(fmt.Sprintf("debug information uploaded to: %s", gist.URL))
	return nil
}

func captureCmd(name string, arg ...string) (string, error) {
	var buf bytes.Buffer
	c := exec.Command(name, arg...)
	c.Stdout = &buf
	if err := c.Run(); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func captureJobs(gist *Gist, env bool) error {
	client := cluster.NewClient()

	jobs, err := jobList(client, true)
	if err != nil {
		return err
	}

	var buf bytes.Buffer
	printJobs(jobs, &buf)
	gist.AddFile("1-jobs.log", buf.String())

	for _, job := range jobs {
		var name string
		if system, ok := job.Job.Metadata["flynn-system-app"]; !ok || system != "true" {
			continue // Skip non-system applications
		}
		if app, ok := job.Job.Metadata["flynn-controller.app_name"]; ok {
			name += app + "-"
		}
		if typ, ok := job.Job.Metadata["flynn-controller.type"]; ok {
			name += typ + "-"
		}
		name += job.Job.ID + ".log"

		var content bytes.Buffer
		printJobDesc(&job, &content, env)
		fmt.Fprint(&content, "\n\n***** ***** ***** ***** ***** ***** ***** ***** ***** *****\n\n")
		getLog(job.HostID, job.Job.ID, client, false, true, &content, &content)
		gist.AddFile(name, content.String())
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Jacob Taylor jacob@ablox.io
// License: Apache2 - http://www.apache.org/licenses/LICENSE-2.0
package main

import (
    "fmt"
    "net"
    "../utils"
    "bufio"
    "encoding/binary"
    //"time"
    "os"
    "bytes"
    "io"
    "io/ioutil"
    "github.com/urfave/cli"
    "path/filepath"
)

const nbd_folder = "/sample_disks/"

var characters_per_line = 100
var newline = 0
var line_number = 0

// settings for the server
type Settings struct {
    ReadOnly    bool
    AutoFlush   bool
    Host        string
    Port        int32
    Listen      string
    File        string
    Directory   string
}

var globalSettings Settings = &Settings{
    ReadOnly: false,
    AutoFlush: true,
    Host: "localhost",
    Port: 8000,
    Listen: "",
    File: "",
    Directory: "sample_disks",
}

func send_export_list_item(output *bufio.Writer, options uint32, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, options, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer, options uint32) {
    send_message(output, options, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte, options uint32, globalSettings Settings) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    //todo add support for folder specifications
    //todo add support for file specificiation

    var filename bytes.Buffer
    current_directory, err := os.Getwd()
    utils.ErrorCheck(err)
    filename.WriteString(current_directory)
    filename.WriteString(nbd_folder)
    filename.Write(payload[:payload_size])

    fmt.Printf("Opening file: %s\n", filename.String())

    file, err := os.OpenFile(filename.String(), os.O_RDWR, 0644)

    utils.ErrorCheck(err)
    if err != nil {
        return
    }

    buffer := make([]byte, 256)
    offset := 0

    fs, err := file.Stat()
    file_size := uint64(fs.Size())

    binary.BigEndian.PutUint64(buffer[offset:], file_size)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 1)  // flags
    offset += 2

    // if requested, pad with 124 zeros
    if (options & utils.NBD_FLAG_NO_ZEROES) != utils.NBD_FLAG_NO_ZEROES {
        offset += 124
    }

    _, err = output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)

    buffer = make([]byte, 2048*1024)    // set the buffer to 2mb
    conn_reader := bufio.NewReader(conn)
    for {
        waiting_for := 28       // wait for at least the minimum payload size

        _, err := io.ReadFull(conn_reader, buffer[:waiting_for])
        if err == io.EOF {
            fmt.Printf("Abort detected, escaping processing loop\n")
            break
        }
        utils.ErrorCheck(err)

        //magic := binary.BigEndian.Uint32(buffer)
        command := binary.BigEndian.Uint32(buffer[4:8])
        //handle := binary.BigEndian.Uint64(buffer[8:16])
        from := binary.BigEndian.Uint64(buffer[16:24])
        length := binary.BigEndian.Uint32(buffer[24:28])

        newline += 1;
        if newline % characters_per_line == 0 {
            line_number++
            fmt.Printf("\n%5d: ", line_number * 100)
            newline -= characters_per_line
        }

        switch command {
        case utils.NBD_COMMAND_READ:
            fmt.Printf(".")

            _, err = file.ReadAt(buffer[16:16+length], int64(from))
            utils.ErrorCheck(err)

            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16+length])

            continue
        case utils.NBD_COMMAND_WRITE:
            fmt.Printf("W")

            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)

            _, err = file.WriteAt(buffer[28:28+length], int64(from))
            utils.ErrorCheck(err)

            if globalSettings.AutoFlush {
                file.Sync()
            }

            // let them know we are done
            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16])

            continue

        case utils.NBD_COMMAND_DISCONNECT:
            fmt.Printf("D")

            file.Sync()
            return
        }
    }
}

/*
First check for a specific file. If one is specified, use it. If not, check for a directory. If that is not
available, use the CWD.
 */
func send_export_list(output *bufio.Writer, options uint32, globalSettings Settings) {
    if globalSettings.File {
        _, file := filepath.Split(globalSettings.File)

        send_export_list_item(output, options, file)
        send_ack(output, options)
        return
    }

    var current_directory string
    var err error
    if globalSettings.Directory == "" {
        current_directory, err = os.Getwd()
        utils.ErrorCheck(err)
    }

    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    utils.ErrorCheck(err)

    for _, file := range files {
        send_export_list_item(output, options, file.Name())
    }

    send_ack(output, options)
}

func send_message(output *bufio.Writer, options uint32, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], options)  // put out the server options
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

var defaultOptions = []byte{0, 0}

func main() {

    app := cli.NewApp()
    app.Name = "AnyBlox"
    app.Usage = "block storage for the masses"
    app.Action = func(c *cli.Context) error {
        fmt.Println("Please specify either a full 'listen' parameter (e.g. 'localhost:8000', '192.168.1.2:8000) or a host and port\n")
        return nil
    }

    app.Flags = []cli.Flag {
        cli.StringFlag{
            Name: "host",
            Value: globalSettings.Host,
            Usage: "Hostname or IP address you want to serve traffic on. e.x. 'localhost', '192.168.1.2'",
            Destination: &globalSettings.Host,
        },
        cli.StringFlag{
            Name: "port",
            Value: globalSettings.Port,
            Usage: "Port you want to serve traffic on. e.x. '8000'",
            Destination: &globalSettings.Port,
        },
        cli.StringFlag{
            Name: "listen, l",
            Destination: &globalSettings.Listen,
            Usage: "Address and port the server should listen on. Listen will take priority over host and port parameters. hostname:port - e.x. 'localhost:8000', '192.168.1.2:8000'",
        },
        cli.StringFlag{
            Name: "file, f",
            Destination: &globalSettings.File,
            Value: "",
            Usage: "The file that should be shared by this server. 'file' overrides 'directory'. It is required to be a full absolute path that includes the filename",
        },
        cli.StringFlag{
            Name: "directory, d",
            Destination: &globalSettings.Directory,
            Value: globalSettings.Directory,
            Usage: "Specify a directory where the files to share are located. Default is 'sample_disks",
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    hostingAddress := globalSettings.Listen
    if len(globalSettings.Listen) == 0 {
        hostingAddress = globalSettings.Host + ":" + globalSettings.Port
    }

    fmt.Printf("About to listen on %s\n", hostingAddress)
    listener, err := net.Listen("tcp", hostingAddress)
    utils.ErrorCheck(err)

    fmt.Printf("aBlox server online\n")

    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    defer fmt.Printf("End of line\n")

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.WriteString("IHAVEOPT")      // Magic

        output.Write(defaultOptions)

        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        _, err = io.ReadFull(conn, data[:waiting_for])
        utils.ErrorCheck(err)

        options := binary.BigEndian.Uint32(data[:4])
        command := binary.BigEndian.Uint32(data[12:16])

        // If we are requesting an export, make sure we have the length of the data for the export name.
        if binary.BigEndian.Uint32(data[12:]) == utils.NBD_COMMAND_EXPORT_NAME {
            waiting_for += 4
            _, err = io.ReadFull(conn, data[16:20])
            utils.ErrorCheck(err)
        }
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        fmt.Printf("command is: %d\npayload_size is: %d\n", command, payload_size)
        offset = waiting_for
        waiting_for += int(payload_size)
        _, err = io.ReadFull(conn, data[offset:waiting_for])
        utils.ErrorCheck(err)

        payload := make([]byte, payload_size)
        if payload_size > 0 {
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output, options, globalSettings)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options, globalSettings)
            break
        }
    }

}
<commit_msg>typos, readonly filter on write, read only tracking and setting support. fixed #20<commit_after>// Copyright 2016 Jacob Taylor jacob@ablox.io
// License: Apache2 - http://www.apache.org/licenses/LICENSE-2.0
package main

import (
    "fmt"
    "net"
    "../utils"
    "bufio"
    "encoding/binary"
    //"time"
    "os"
    "bytes"
    "io"
    "io/ioutil"
    "github.com/urfave/cli"
    "path/filepath"
)

const nbd_folder = "/sample_disks/"

var characters_per_line = 100
var newline = 0
var line_number = 0

// settings for the server
type Settings struct {
    ReadOnly    bool
    AutoFlush   bool
    Host        string
    Port        int32
    Listen      string
    File        string
    Directory   string
}

var globalSettings Settings = &Settings{
    ReadOnly: false,
    AutoFlush: true,
    Host: "localhost",
    Port: 8000,
    Listen: "",
    File: "",
    Directory: "sample_disks",
}

func send_export_list_item(output *bufio.Writer, options uint32, export_name string) {
    data := make([]byte, 1024)
    length := len(export_name)
    offset := 0

    // length of export name
    binary.BigEndian.PutUint32(data[offset:], uint32(length))  // length of string
    offset += 4

    // export name
    copy(data[offset:], export_name)
    offset += length

    reply_type := uint32(2)     // reply_type: NBD_REP_SERVER
    send_message(output, options, reply_type, uint32(offset), data)
}

func send_ack(output *bufio.Writer, options uint32) {
    send_message(output, options, utils.NBD_COMMAND_ACK, 0, nil)
}

func export_name(output *bufio.Writer, conn net.Conn, payload_size int, payload []byte, options uint32, globalSettings Settings) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    //todo add support for folder specifications
    //todo add support for file specification

    var filename bytes.Buffer
    readOnly := false

    current_directory, err := os.Getwd()
    utils.ErrorCheck(err)
    filename.WriteString(current_directory)
    filename.WriteString(nbd_folder)
    filename.Write(payload[:payload_size])

    fmt.Printf("Opening file: %s\n", filename.String())

    fileMode := os.O_RDWR
    if globalSettings.ReadOnly || options && utils.NBD_OPT_READ_ONLY {
        fmt.Printf("Read Only is set\n")
        fileMode = os.O_RDONLY
        readOnly = true
    }

    file, err := os.OpenFile(filename.String(), fileMode, 0644)

    utils.ErrorCheck(err)
    if err != nil {
        return
    }

    buffer := make([]byte, 256)
    offset := 0

    fs, err := file.Stat()
    file_size := uint64(fs.Size())

    binary.BigEndian.PutUint64(buffer[offset:], file_size)  // size
    offset += 8

    binary.BigEndian.PutUint16(buffer[offset:], 1)  // flags
    offset += 2

    // if requested, pad with 124 zeros
    if (options & utils.NBD_FLAG_NO_ZEROES) != utils.NBD_FLAG_NO_ZEROES {
        offset += 124
    }

    _, err = output.Write(buffer[:offset])
    output.Flush()
    utils.ErrorCheck(err)

    buffer = make([]byte, 2048*1024)    // set the buffer to 2mb
    conn_reader := bufio.NewReader(conn)
    for {
        waiting_for := 28       // wait for at least the minimum payload size

        _, err := io.ReadFull(conn_reader, buffer[:waiting_for])
        if err == io.EOF {
            fmt.Printf("Abort detected, escaping processing loop\n")
            break
        }
        utils.ErrorCheck(err)

        //magic := binary.BigEndian.Uint32(buffer)
        command := binary.BigEndian.Uint32(buffer[4:8])
        //handle := binary.BigEndian.Uint64(buffer[8:16])
        from := binary.BigEndian.Uint64(buffer[16:24])
        length := binary.BigEndian.Uint32(buffer[24:28])

        newline += 1;
        if newline % characters_per_line == 0 {
            line_number++
            fmt.Printf("\n%5d: ", line_number * 100)
            newline -= characters_per_line
        }

        switch command {
        case utils.NBD_COMMAND_READ:
            fmt.Printf(".")

            _, err = file.ReadAt(buffer[16:16+length], int64(from))
            utils.ErrorCheck(err)

            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16+length])

            continue
        case utils.NBD_COMMAND_WRITE:
            if readOnly {
                fmt.Printf("E")
                fmt.Printf("\nAttempt to write to read only file blocked\n")

                continue
            }

            fmt.Printf("W")

            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)

            _, err = file.WriteAt(buffer[28:28+length], int64(from))
            utils.ErrorCheck(err)

            if globalSettings.AutoFlush {
                file.Sync()
            }

            // let them know we are done
            binary.BigEndian.PutUint32(buffer[:4], utils.NBD_REPLY_MAGIC)
            binary.BigEndian.PutUint32(buffer[4:8], 0)                      // error bits

            conn.Write(buffer[:16])

            continue

        case utils.NBD_COMMAND_DISCONNECT:
            fmt.Printf("D")

            file.Sync()
            return
        }
    }
}

/*
First check for a specific file. If one is specified, use it. If not, check for a directory. If that is not
available, use the CWD.
 */
func send_export_list(output *bufio.Writer, options uint32, globalSettings Settings) {
    if globalSettings.File {
        _, file := filepath.Split(globalSettings.File)

        send_export_list_item(output, options, file)
        send_ack(output, options)
        return
    }

    var current_directory string
    var err error
    if globalSettings.Directory == "" {
        current_directory, err = os.Getwd()
        utils.ErrorCheck(err)
    }

    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    utils.ErrorCheck(err)

    for _, file := range files {
        send_export_list_item(output, options, file.Name())
    }

    send_ack(output, options)
}

func send_message(output *bufio.Writer, options uint32, reply_type uint32, length uint32, data []byte ) {
    endian := binary.BigEndian
    buffer := make([]byte, 1024)
    offset := 0

    endian.PutUint64(buffer[offset:], utils.NBD_SERVER_SEND_REPLY_MAGIC)
    offset += 8

    endian.PutUint32(buffer[offset:], options)  // put out the server options
    offset += 4

    endian.PutUint32(buffer[offset:], reply_type)  // reply_type: NBD_REP_SERVER
    offset += 4

    endian.PutUint32(buffer[offset:], length)  // length of package
    offset += 4

    if data != nil {
        copy(buffer[offset:], data[0:length])
        offset += int(length)
    }

    data_to_send := buffer[:offset]
    output.Write(data_to_send)
    output.Flush()

    utils.LogData("Just sent:", offset, data_to_send)
}

var defaultOptions = []byte{0, 0}

func main() {

    app := cli.NewApp()
    app.Name = "AnyBlox"
    app.Usage = "block storage for the masses"
    app.Action = func(c *cli.Context) error {
        fmt.Println("Please specify either a full 'listen' parameter (e.g. 'localhost:8000', '192.168.1.2:8000) or a host and port\n")
        return nil
    }

    app.Flags = []cli.Flag {
        cli.StringFlag{
            Name: "host",
            Value: globalSettings.Host,
            Usage: "Hostname or IP address you want to serve traffic on. e.x. 'localhost', '192.168.1.2'",
            Destination: &globalSettings.Host,
        },
        cli.StringFlag{
            Name: "port",
            Value: globalSettings.Port,
            Usage: "Port you want to serve traffic on. e.x. '8000'",
            Destination: &globalSettings.Port,
        },
        cli.StringFlag{
            Name: "listen, l",
            Destination: &globalSettings.Listen,
            Usage: "Address and port the server should listen on. Listen will take priority over host and port parameters. hostname:port - e.x. 'localhost:8000', '192.168.1.2:8000'",
        },
        cli.StringFlag{
            Name: "file, f",
            Destination: &globalSettings.File,
            Value: "",
            Usage: "The file that should be shared by this server. 'file' overrides 'directory'. It is required to be a full absolute path that includes the filename",
        },
        cli.StringFlag{
            Name: "directory, d",
            Destination: &globalSettings.Directory,
            Value: globalSettings.Directory,
            Usage: "Specify a directory where the files to share are located. Default is 'sample_disks",
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    hostingAddress := globalSettings.Listen
    if len(globalSettings.Listen) == 0 {
        hostingAddress = globalSettings.Host + ":" + globalSettings.Port
    }

    fmt.Printf("About to listen on %s\n", hostingAddress)
    listener, err := net.Listen("tcp", hostingAddress)
    utils.ErrorCheck(err)

    fmt.Printf("aBlox server online\n")

    reply_magic := make([]byte, 4)
    binary.BigEndian.PutUint32(reply_magic, utils.NBD_REPLY_MAGIC)

    defer fmt.Printf("End of line\n")

    for {
        conn, err := listener.Accept()
        utils.ErrorCheck(err)

        fmt.Printf("We have a new connection from: %s\n", conn.RemoteAddr())
        output := bufio.NewWriter(conn)

        output.WriteString("NBDMAGIC")      // init password
        output.WriteString("IHAVEOPT")      // Magic

        output.Write(defaultOptions)

        output.Flush()

        // Fetch the data until we get the initial options
        data := make([]byte, 1024)
        offset := 0
        waiting_for := 16       // wait for at least the minimum payload size

        _, err = io.ReadFull(conn, data[:waiting_for])
        utils.ErrorCheck(err)

        options := binary.BigEndian.Uint32(data[:4])
        command := binary.BigEndian.Uint32(data[12:16])

        // If we are requesting an export, make sure we have the length of the data for the export name.
        if binary.BigEndian.Uint32(data[12:]) == utils.NBD_COMMAND_EXPORT_NAME {
            waiting_for += 4
            _, err = io.ReadFull(conn, data[16:20])
            utils.ErrorCheck(err)
        }
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        fmt.Printf("command is: %d\npayload_size is: %d\n", command, payload_size)
        offset = waiting_for
        waiting_for += int(payload_size)
        _, err = io.ReadFull(conn, data[offset:waiting_for])
        utils.ErrorCheck(err)

        payload := make([]byte, payload_size)
        if payload_size > 0 {
            copy(payload, data[20:])
        }

        utils.LogData("Payload is:", payload_size, payload)

        // At this point, we have the command, payload size, and payload.
        switch command {
        case utils.NBD_COMMAND_LIST:
            send_export_list(output, options, globalSettings)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options, globalSettings)
            break
        }
    }

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"github.com/golang/glog"
	"github.com/heidi-ann/hydra/cache"
	"github.com/heidi-ann/hydra/config"
	"github.com/heidi-ann/hydra/consensus"
	"github.com/heidi-ann/hydra/msgs"
	"github.com/heidi-ann/hydra/store"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

var keyval *store.Store
var disk *bufio.Writer
var c *cache.Cache
var cons_io *msgs.Io

type Peer struct {
	id      int
	address string
	handled bool // TOOD: replace with Mutex
}

var peers []Peer

var client_port = flag.Int("client-port", 8080, "port to listen on for clients")
var peer_port = flag.Int("peer-port", 8090, "port to listen on for peers")
var id = flag.Int("id", -1, "server ID")
var config_file = flag.String("config", "example.conf", "Server configuration file")

func handleRequest(req msgs.ClientRequest) msgs.ClientResponse {
	glog.Info("Handling ", req.Request)

	// check if already applied
	found, res := c.Check(req)
	if found {
		glog.Info("Request found in cache")
		return res // FAST PASS
	}

	// write to persistent storage
	n, err := disk.WriteString(req.Request)
	_ = disk.Flush()
	if err != nil {
		glog.Fatal(err)
	}
	glog.Infof("Written %b bytes to persistent storage", n)

	// CONSENESUS ALGORITHM HERE
	glog.Info("Passing request to consensus algorithm")
	(*cons_io).IncomingRequests <- req
	_ = <-(*cons_io).OutgoingRequests
	glog.Info("Request has been safely replicated by consensus algorithm")

	// check if request already applied
	found, res = c.Check(req)
	if found {
		glog.Info("Request found in cache and thus cannot be applied")
		return res
	}

	// apply request
	output := keyval.Process(req.Request)
	keyval.Print()

	// write response to request cache
	reply := msgs.ClientResponse{
		req.ClientID, req.RequestID, output}

	c.Add(reply)
	return reply
}

// iterative through peers and check there is a handler for each
// try to create one if not
func checkPeer() {
	for i := range peers {
		if !peers[i].handled {
			glog.Info("Peer ", i, " is not currently connected")
			cn, err := net.Dial("tcp", peers[i].address)

			if err != nil {
				glog.Warning(err)
				break
			}

			handlePeer(cn, true)
		}
	}
}

func handlePeer(cn net.Conn, _ bool) {
	addr := cn.RemoteAddr().String()
	glog.Info("Peer connection from ", addr)

	// handle requests
	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	// exchange peer ID's
	_, _ = writer.WriteString(strconv.Itoa(*id) + "\n")
	_ = writer.Flush()
	text, _ := reader.ReadString('\n')
	glog.Info("Received ", text)
	peer_id, err := strconv.Atoi(strings.Trim(text, "\n"))
	if err != nil {
		glog.Warning(err)
		return
	}

	glog.Infof("Ready to handle traffic from peer %d at %s ", peer_id, addr)

	peers[peer_id].handled = true

	close_err := make(chan error)
	go func() {
		for {
			// read request
			text, err := reader.ReadBytes(byte('\n'))
			glog.Info(string(text))
			glog.Info("Reading from peer ", peer_id)
			if err != nil {
				close_err <- err
				break
			}

			(*cons_io).Incoming.BytesToProtoMsg(text)

		}
	}()

	go func() {
		for {
			// send reply
			b, err := (*cons_io).OutgoingUnicast[peer_id].ProtoMsgToBytes()
			if err != nil {
				glog.Fatal("Could not marshal message")
			}
			glog.Info("Sending ", string(b))
			_, err = writer.Write(b)
			_, err = writer.Write([]byte("\n"))
			if err != nil {
				close_err <- err
				break
			}
			err = writer.Flush()
		}
	}()

	// block until connection fails
	err = <-close_err
	glog.Warning(err)

	// tidy up
	glog.Infof("No longer about to handle traffic from peer %d at %s ", id, addr)
	peers[peer_id].handled = false
	(*cons_io).Failure <- peer_id
	cn.Close()
}

func handleConnection(cn net.Conn) {
	glog.Info("Incoming client connection from ",
		cn.RemoteAddr().String())

	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	for {

		// read request
		glog.Info("Reading")
		text, err := reader.ReadBytes(byte('\n'))
		if err != nil {
			if err == io.EOF {
				break
			}
			glog.Fatal(err)
		}
		glog.Info(string(text))
		req := new(msgs.ClientRequest)
		msgs.Unmarshal(text, req)

		// construct reply
		reply := handleRequest(*req)
		b, err := msgs.Marshal(reply)
		if err != nil {
			glog.Fatal("error:", err)
		}
		glog.Info(string(b))

		// send reply
		// TODO: FIX currently all server send back replies
		glog.Info("Sending ", string(b))
		n, err := writer.Write(b)
		_, err = writer.Write([]byte("\n"))
		if err != nil {
			glog.Fatal(err)
		}

		// tidy up
		err = writer.Flush()
		glog.Info("Finished sending ", n, " bytes")

	}

	cn.Close()
}

func main() {
	// set up logging
	flag.Parse()
	defer glog.Flush()

	conf := config.ParseServerConfig(*config_file)
	if *id == -1 {
		glog.Fatal("ID is required")
	}

	glog.Info("Starting server ", *id)

	//set up state machine
	keyval = store.New()
	c = cache.Create()

	// setting up persistent log
	filename := "persistent_" + strconv.Itoa(*id) + ".log"
	glog.Info("Opening file: ", filename)
	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		glog.Fatal(err)
	}
	disk = bufio.NewWriter(file)
	defer disk.Flush()

	// check persistent storage for commands
	disk_reader := bufio.NewReader(file)
	for {
		str, err := disk_reader.ReadString('\n')
		if err != nil {
			glog.Info("No more commands in persistent storage")
			break
		}
		_ = keyval.Process(str)
	}

	// set up client server
	glog.Info("Starting up client server")
	listeningPort := ":" + strconv.Itoa(*client_port)
	ln, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle for incoming clients
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handleConnection(conn)
		}
	}()

	//set up peer state
	peers = make([]Peer, len(conf.Peers.Address))
	for i := range conf.Peers.Address {
		peers[i] = Peer{
			i, conf.Peers.Address[i], false}
	}
	cons_io = msgs.MakeIo(10, len(conf.Peers.Address))

	//set up peer server
	glog.Info("Starting up peer server")
	listeningPort = ":" + strconv.Itoa(*peer_port)
	lnPeers, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle local peer (without sending network traffic)
	peers[*id].handled = true
	from := &((*cons_io).Incoming)
	go from.Forward((*cons_io).OutgoingUnicast[*id])

	// regularly check if all peers are connected and reply if not
	go func() {
		for {
			checkPeer()
			time.Sleep(10 * time.Second)
		}
	}()

	// handle for incoming peers
	go func() {
		for {
			conn, err := lnPeers.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handlePeer(conn, false)
		}
	}()

	// setting up the consensus algorithm
	cons_config := consensus.Config{*id, len(conf.Peers.Address)}
	consensus.Init(cons_io, cons_config)
	cons_io.DumpPersistentStorage()

	// tidy up
	time.Sleep(30 * time.Second)
	glog.Info("Shutting down")
}
<commit_msg>adding persistent log code to server.go<commit_after>package main

import (
	"bufio"
	"flag"
	"github.com/golang/glog"
	"github.com/heidi-ann/hydra/cache"
	"github.com/heidi-ann/hydra/config"
	"github.com/heidi-ann/hydra/consensus"
	"github.com/heidi-ann/hydra/msgs"
	"github.com/heidi-ann/hydra/store"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

var keyval *store.Store
var disk *bufio.Writer
var c *cache.Cache
var cons_io *msgs.Io

type Peer struct {
	id      int
	address string
	handled bool // TOOD: replace with Mutex
}

var peers []Peer

var client_port = flag.Int("client-port", 8080, "port to listen on for clients")
var peer_port = flag.Int("peer-port", 8090, "port to listen on for peers")
var id = flag.Int("id", -1, "server ID")
var config_file = flag.String("config", "example.conf", "Server configuration file")

func handleRequest(req msgs.ClientRequest) msgs.ClientResponse {
	glog.Info("Handling ", req.Request)

	// check if already applied
	found, res := c.Check(req)
	if found {
		glog.Info("Request found in cache")
		return res // FAST PASS
	}

	// CONSENESUS ALGORITHM HERE
	glog.Info("Passing request to consensus algorithm")
	(*cons_io).IncomingRequests <- req
	_ = <-(*cons_io).OutgoingRequests
	glog.Info("Request has been safely replicated by consensus algorithm")

	// check if request already applied
	found, res = c.Check(req)
	if found {
		glog.Info("Request found in cache and thus cannot be applied")
		return res
	}

	// apply request
	output := keyval.Process(req.Request)
	keyval.Print()

	// write response to request cache
	reply := msgs.ClientResponse{
		req.ClientID, req.RequestID, output}

	c.Add(reply)
	return reply
}

// iterative through peers and check there is a handler for each
// try to create one if not
func checkPeer() {
	for i := range peers {
		if !peers[i].handled {
			glog.Info("Peer ", i, " is not currently connected")
			cn, err := net.Dial("tcp", peers[i].address)

			if err != nil {
				glog.Warning(err)
				break
			}

			handlePeer(cn, true)
		}
	}
}

func handlePeer(cn net.Conn, _ bool) {
	addr := cn.RemoteAddr().String()
	glog.Info("Peer connection from ", addr)

	// handle requests
	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	// exchange peer ID's
	_, _ = writer.WriteString(strconv.Itoa(*id) + "\n")
	_ = writer.Flush()
	text, _ := reader.ReadString('\n')
	glog.Info("Received ", text)
	peer_id, err := strconv.Atoi(strings.Trim(text, "\n"))
	if err != nil {
		glog.Warning(err)
		return
	}

	glog.Infof("Ready to handle traffic from peer %d at %s ", peer_id, addr)

	peers[peer_id].handled = true

	close_err := make(chan error)
	go func() {
		for {
			// read request
			text, err := reader.ReadBytes(byte('\n'))
			glog.Info(string(text))
			glog.Info("Reading from peer ", peer_id)
			if err != nil {
				close_err <- err
				break
			}

			(*cons_io).Incoming.BytesToProtoMsg(text)

		}
	}()

	go func() {
		for {
			// send reply
			b, err := (*cons_io).OutgoingUnicast[peer_id].ProtoMsgToBytes()
			if err != nil {
				glog.Fatal("Could not marshal message")
			}
			glog.Info("Sending ", string(b))
			_, err = writer.Write(b)
			_, err = writer.Write([]byte("\n"))
			if err != nil {
				close_err <- err
				break
			}
			err = writer.Flush()
		}
	}()

	// block until connection fails
	err = <-close_err
	glog.Warning(err)

	// tidy up
	glog.Infof("No longer about to handle traffic from peer %d at %s ", id, addr)
	peers[peer_id].handled = false
	(*cons_io).Failure <- peer_id
	cn.Close()
}

func handleConnection(cn net.Conn) {
	glog.Info("Incoming client connection from ",
		cn.RemoteAddr().String())

	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	for {

		// read request
		glog.Info("Reading")
		text, err := reader.ReadBytes(byte('\n'))
		if err != nil {
			if err == io.EOF {
				break
			}
			glog.Fatal(err)
		}
		glog.Info(string(text))
		req := new(msgs.ClientRequest)
		msgs.Unmarshal(text, req)

		// construct reply
		reply := handleRequest(*req)
		b, err := msgs.Marshal(reply)
		if err != nil {
			glog.Fatal("error:", err)
		}
		glog.Info(string(b))

		// send reply
		// TODO: FIX currently all server send back replies
		glog.Info("Sending ", string(b))
		n, err := writer.Write(b)
		_, err = writer.Write([]byte("\n"))
		if err != nil {
			glog.Fatal(err)
		}

		// tidy up
		err = writer.Flush()
		glog.Info("Finished sending ", n, " bytes")

	}

	cn.Close()
}

func main() {
	// set up logging
	flag.Parse()
	defer glog.Flush()

	conf := config.ParseServerConfig(*config_file)
	if *id == -1 {
		glog.Fatal("ID is required")
	}

	glog.Info("Starting server ", *id)

	//set up state machine
	keyval = store.New()
	c = cache.Create()

	// setup IO
	cons_io = msgs.MakeIo(10, len(conf.Peers.Address))

	// setting up persistent log
	filename := "persistent_" + strconv.Itoa(*id) + ".log"
	glog.Info("Opening file: ", filename)
	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		glog.Fatal(err)
	}
	disk = bufio.NewWriter(file)
	defer disk.Flush()

	// check persistent storage for commands
	found := false
	log := make([]msgs.Entry, 100) //TODO: Fix this
	view := 0
	disk_reader := bufio.NewReader(file)
	for {
		b, err := disk_reader.ReadBytes(byte('\n'))
		if err != nil {
			glog.Info("No more commands in persistent storage")
			break
		}
		found = true
		var update msgs.LogUpdate
		err = msgs.Unmarshal(b, &update)
		if err != nil {
			glog.Fatal("Cannot parse log update", err)
		}
		log[update.Index] = update.Entry
		//TODO: do this properly
		view = update.Entry.View
	}

	// write updates to persistent storage
	go func() {
		for {
			// get write requests
			select {
			//disgard view updates
			case view := <-(*cons_io).ViewPersist:
				glog.Info("Updating view to ", view)

			case log := <-(*cons_io).LogPersist:
				glog.Info("Updating log with ", log)
				b, err := msgs.Marshal(log)
				// write to persistent storage
				_, err = disk.Write(b)
				_, err = disk.Write([]byte("\n"))
				_ = disk.Flush()
				if err != nil {
					glog.Fatal(err)
				}
			}

		}
	}()

	// set up client server
	glog.Info("Starting up client server")
	listeningPort := ":" + strconv.Itoa(*client_port)
	ln, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle for incoming clients
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handleConnection(conn)
		}
	}()

	//set up peer state
	peers = make([]Peer, len(conf.Peers.Address))
	for i := range conf.Peers.Address {
		peers[i] = Peer{
			i, conf.Peers.Address[i], false}
	}

	//set up peer server
	glog.Info("Starting up peer server")
	listeningPort = ":" + strconv.Itoa(*peer_port)
	lnPeers, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle local peer (without sending network traffic)
	peers[*id].handled = true
	from := &((*cons_io).Incoming)
	go from.Forward((*cons_io).OutgoingUnicast[*id])

	// regularly check if all peers are connected and reply if not
	go func() {
		for {
			checkPeer()
			time.Sleep(10 * time.Second)
		}
	}()

	// handle for incoming peers
	go func() {
		for {
			conn, err := lnPeers.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handlePeer(conn, false)
		}
	}()

	// setting up the consensus algorithm
	cons_config := consensus.Config{*id, len(conf.Peers.Address)}
	if !found {
		consensus.Init(cons_io, cons_config)
	} else {
		consensus.Recover(cons_io, cons_config, view, log)
	}
	cons_io.DumpPersistentStorage()

	// tidy up
	time.Sleep(30 * time.Second)
	glog.Info("Shutting down")
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command server launches a stand-alone inverting proxy.
//
// Example usage:
//   go build -o ~/bin/inverting-proxy ./server/server.go
//   ~/bin/inverting-proxy --port 8081
package main

import (
	"bufio"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"time"

	"context"

	"github.com/google/inverting-proxy/agent/utils"
)

var (
	port = flag.Int("port", 0, "Port on which to listen")
)

// pendingRequest represents a frontend request
type pendingRequest struct {
	startTime time.Time
	req       *http.Request
	respChan  chan *http.Response
}

func newPendingRequest(r *http.Request) *pendingRequest {
	return &pendingRequest{
		startTime: time.Now(),
		req:       r,
		respChan:  make(chan *http.Response),
	}
}

type proxy struct {
	requestIDs    chan string
	randGenerator *rand.Rand

	// protects the map below
	sync.Mutex
	requests map[string]*pendingRequest
}

func newProxy() *proxy {
	return &proxy{
		requestIDs:    make(chan string),
		randGenerator: rand.New(rand.NewSource(time.Now().UnixNano())),
		requests:      make(map[string]*pendingRequest),
	}
}

func (p *proxy) handleAgentPostResponse(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	resp, err := http.ReadResponse(bufio.NewReader(r.Body), pending.req)
	if err != nil {
		log.Printf("Could not parse response to request %q: %v", requestID, err)
		http.Error(w, "Failure parsing request body", http.StatusBadRequest)
		return
	}
	// We want to track whether or not the body has finished being read so that we can
	// make sure that this method does not return until after that. However, we do not
	// want to block the sending of the response to the client while it is being read.
	//
	// To accommodate both goals, we replace the response body with a pipereader, start
	// forwarding the response immediately, and then copy the original body to the
	// corresponding pipewriter.
	respBody := resp.Body
	defer respBody.Close()

	pr, pw := io.Pipe()
	defer pw.Close()

	resp.Body = pr
	select {
	case <-r.Context().Done():
		return
	case pending.respChan <- resp:
	}
	if _, err := io.Copy(pw, respBody); err != nil {
		log.Printf("Could not read response to request %q: %v", requestID, err)
		http.Error(w, "Failure reading request body", http.StatusInternalServerError)
	}
}

func (p *proxy) handleAgentGetRequest(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	log.Printf("Returning pending request: %q", requestID)
	w.Header().Set(utils.HeaderRequestStartTime, pending.startTime.Format(time.RFC3339Nano))
	w.WriteHeader(http.StatusOK)
	pending.req.Write(w)
}

// waitForRequestIDs blocks until at least one request ID is available, and then returns
// a slice of all of the IDs available at that time.
//
// Note that any IDs returned by this method will never be returned again.
func (p *proxy) waitForRequestIDs(ctx context.Context) []string {
	var requestIDs []string
	select {
	case <-ctx.Done():
		return nil
	case <-time.After(30 * time.Second):
		return nil
	case id := <-p.requestIDs:
		requestIDs = append(requestIDs, id)
	}
	for {
		select {
		case id := <-p.requestIDs:
			requestIDs = append(requestIDs, id)
		default:
			return requestIDs
		}
	}
}

func (p *proxy) handleAgentListRequests(w http.ResponseWriter, r *http.Request) {
	requestIDs := p.waitForRequestIDs(r.Context())
	respJSON, err := json.Marshal(requestIDs)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failure serializing the request IDs: %v", err), http.StatusInternalServerError)
		return
	}
	log.Printf("Reporting pending requests: %s", respJSON)
	w.WriteHeader(http.StatusOK)
	w.Write(respJSON)
}

func (p *proxy) handleAgentRequest(w http.ResponseWriter, r *http.Request, backendID string) {
	requestID := r.Header.Get(utils.HeaderRequestID)
	if requestID == "" {
		log.Printf("Received new backend list request from %q", backendID)
		p.handleAgentListRequests(w, r)
		return
	}
	if r.Method == http.MethodPost {
		log.Printf("Received new backend post request from %q", backendID)
		p.handleAgentPostResponse(w, r, requestID)
		return
	}
	log.Printf("Received new backend get request from %q", backendID)
	p.handleAgentGetRequest(w, r, requestID)
}

func (p *proxy) newID() string {
	sum := sha256.Sum256([]byte(fmt.Sprintf("%d", p.randGenerator.Int63())))
	return fmt.Sprintf("%x", sum)
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if backendID := r.Header.Get(utils.HeaderBackendID); backendID != "" {
		p.handleAgentRequest(w, r, backendID)
		return
	}
	id := p.newID()
	log.Printf("Received new frontend request %q", id)
	pending := newPendingRequest(r)
	p.Lock()
	p.requests[id] = pending
	p.Unlock()

	// Enqueue the request
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting to enqueue the request ID for %q", id)
		return
	case p.requestIDs <- id:
	}
	log.Printf("Request %q enqueued after %s", id, time.Since(pending.startTime))

	// Pull out and copy the response
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting for the response to %q", id)
		return
	case resp := <-pending.respChan:
		defer resp.Body.Close()
		for key, vals := range resp.Header {
			w.Header()[key] = vals
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
		return
	}
	log.Printf("Response for %q received after %s", id, time.Since(pending.startTime))
}

func main() {
	flag.Parse()
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
	if err != nil {
		log.Fatalf("Failed to create the TCP listener for port %d: %v", *port, err)
	}
	log.Printf("Listening on %s", listener.Addr())
	log.Fatal(http.Serve(listener, newProxy()))
}
<commit_msg>Make response header copying compliant with RFC 2616<commit_after>/*
Copyright 2018 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command server launches a stand-alone inverting proxy.
//
// Example usage:
//   go build -o ~/bin/inverting-proxy ./server/server.go
//   ~/bin/inverting-proxy --port 8081
package main

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/google/inverting-proxy/agent/utils"
)

var (
	port = flag.Int("port", 0, "Port on which to listen")
)

// pendingRequest represents a frontend request
type pendingRequest struct {
	startTime time.Time
	req       *http.Request
	respChan  chan *http.Response
}

func newPendingRequest(r *http.Request) *pendingRequest {
	return &pendingRequest{
		startTime: time.Now(),
		req:       r,
		respChan:  make(chan *http.Response),
	}
}

type proxy struct {
	requestIDs    chan string
	randGenerator *rand.Rand

	// protects the map below
	sync.Mutex
	requests map[string]*pendingRequest
}

func newProxy() *proxy {
	return &proxy{
		requestIDs:    make(chan string),
		randGenerator: rand.New(rand.NewSource(time.Now().UnixNano())),
		requests:      make(map[string]*pendingRequest),
	}
}

func (p *proxy) handleAgentPostResponse(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	resp, err := http.ReadResponse(bufio.NewReader(r.Body), pending.req)
	if err != nil {
		log.Printf("Could not parse response to request %q: %v", requestID, err)
		http.Error(w, "Failure parsing request body", http.StatusBadRequest)
		return
	}
	// We want to track whether or not the body has finished being read so that we can
	// make sure that this method does not return until after that. However, we do not
	// want to block the sending of the response to the client while it is being read.
	//
	// To accommodate both goals, we replace the response body with a pipereader, start
	// forwarding the response immediately, and then copy the original body to the
	// corresponding pipewriter.
	respBody := resp.Body
	defer respBody.Close()

	pr, pw := io.Pipe()
	defer pw.Close()

	resp.Body = pr
	select {
	case <-r.Context().Done():
		return
	case pending.respChan <- resp:
	}
	if _, err := io.Copy(pw, respBody); err != nil {
		log.Printf("Could not read response to request %q: %v", requestID, err)
		http.Error(w, "Failure reading request body", http.StatusInternalServerError)
	}
}

func (p *proxy) handleAgentGetRequest(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	log.Printf("Returning pending request: %q", requestID)
	w.Header().Set(utils.HeaderRequestStartTime, pending.startTime.Format(time.RFC3339Nano))
	w.WriteHeader(http.StatusOK)
	pending.req.Write(w)
}

// waitForRequestIDs blocks until at least one request ID is available, and then returns
// a slice of all of the IDs available at that time.
//
// Note that any IDs returned by this method will never be returned again.
func (p *proxy) waitForRequestIDs(ctx context.Context) []string {
	var requestIDs []string
	select {
	case <-ctx.Done():
		return nil
	case <-time.After(30 * time.Second):
		return nil
	case id := <-p.requestIDs:
		requestIDs = append(requestIDs, id)
	}
	for {
		select {
		case id := <-p.requestIDs:
			requestIDs = append(requestIDs, id)
		default:
			return requestIDs
		}
	}
}

func (p *proxy) handleAgentListRequests(w http.ResponseWriter, r *http.Request) {
	requestIDs := p.waitForRequestIDs(r.Context())
	respJSON, err := json.Marshal(requestIDs)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failure serializing the request IDs: %v", err), http.StatusInternalServerError)
		return
	}
	log.Printf("Reporting pending requests: %s", respJSON)
	w.WriteHeader(http.StatusOK)
	w.Write(respJSON)
}

func (p *proxy) handleAgentRequest(w http.ResponseWriter, r *http.Request, backendID string) {
	requestID := r.Header.Get(utils.HeaderRequestID)
	if requestID == "" {
		log.Printf("Received new backend list request from %q", backendID)
		p.handleAgentListRequests(w, r)
		return
	}
	if r.Method == http.MethodPost {
		log.Printf("Received new backend post request from %q", backendID)
		p.handleAgentPostResponse(w, r, requestID)
		return
	}
	log.Printf("Received new backend get request from %q", backendID)
	p.handleAgentGetRequest(w, r, requestID)
}

func (p *proxy) newID() string {
	sum := sha256.Sum256([]byte(fmt.Sprintf("%d", p.randGenerator.Int63())))
	return fmt.Sprintf("%x", sum)
}

// isHopByHopHeader determines whether or not the given header name represents
// a header that is specific to a single network hop and thus should not be
// retransmitted by a proxy.
//
// See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers#hbh
func isHopByHopHeader(name string) bool {
	switch n := strings.ToLower(name); n {
	case "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade":
		return true
	default:
		return false
	}
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if backendID := r.Header.Get(utils.HeaderBackendID); backendID != "" {
		p.handleAgentRequest(w, r, backendID)
		return
	}
	id := p.newID()
	log.Printf("Received new frontend request %q", id)
	pending := newPendingRequest(r)
	p.Lock()
	p.requests[id] = pending
	p.Unlock()

	// Enqueue the request
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting to enqueue the request ID for %q", id)
		return
	case p.requestIDs <- id:
	}
	log.Printf("Request %q enqueued after %s", id, time.Since(pending.startTime))

	// Pull out and copy the response
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting for the response to %q", id)
		return
	case resp := <-pending.respChan:
		defer resp.Body.Close()
		// Copy all of the non-hop-by-hop headers to the proxied response
		for key, vals := range resp.Header {
			if !isHopByHopHeader(key) {
				w.Header()[key] = vals
			}
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
		return
	}
	log.Printf("Response for %q received after %s", id, time.Since(pending.startTime))
}

func main() {
	flag.Parse()
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
	if err != nil {
		log.Fatalf("Failed to create the TCP listener for port %d: %v", *port, err)
	}
	log.Printf("Listening on %s", listener.Addr())
	log.Fatal(http.Serve(listener, newProxy()))
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/zubairhamed/betwixt/api"
	"github.com/zubairhamed/canopus"
	. "github.com/zubairhamed/go-commons/network"
	"log"
	"net"
	"strings"
)

func NewDefaultCoapServer() *canopus.CoapServer {
	localAddr, err := net.ResolveUDPAddr("udp", ":5683")
	if err != nil {
		log.Println("Error starting CoAP Server: ", err)
	}
	return canopus.NewServer(localAddr, nil)
}

func NewDefaultServer() api.Server {
	return &DefaultServer{
		coapServer: NewDefaultCoapServer(),
		httpServer: NewDefaultHttpServer(),
		clients:    make(map[string]api.RegisteredClient),
		stats:      &ServerStatistics{},
	}
}

type DefaultServer struct {
	coapServer *canopus.CoapServer
	httpServer *HttpServer
	registry   api.Registry
	stats      *ServerStatistics
	clients    map[string]api.RegisteredClient
}

func (server *DefaultServer) Start() {
	coap := server.coapServer

	// Setup CoAP Routes
	SetupCoapRoutes(server)

	// Start CoAP Server
	go func() {
		coap.Start()
	}()

	// Setup HTTP Routes
	http := server.httpServer
	SetupHttpRoutes(server)

	// Start HTTP Server
	http.Start()
}

func (server *DefaultServer) UseRegistry(reg api.Registry) {
	server.registry = reg
}

func (server *DefaultServer) GetRegisteredClient(id string) api.RegisteredClient {
	return server.clients[id]
}

func (server *DefaultServer) update(id string) {
	for k, v := range server.clients {
		if v.GetId() == id {
			v.Update()
			server.clients[k] = v
		}
	}
}

func (server *DefaultServer) register(ep string, addr string, resources []*canopus.CoreResource) (string, error) {
	clientId := canopus.GenerateToken(8)
	newClient := NewRegisteredClient(ep, clientId, addr)

	objs := make(map[string][]string)
	for _, o := range resources {
		t := o.Target[1:len(o.Target)]
		sp := strings.Split(t, "/")

		if len(sp) > 1 {

		} else {

		}

		log.Println(t, sp[0], sp[1])

		objs[sp[0]] = append(objs[sp[0]], sp[1])
	}

	newClient.SetObjects(objs)

	server.clients[ep] = newClient

	return clientId, nil
}

func (server *DefaultServer) delete(id string) {
	for k, v := range server.clients {
		if v.GetId() == id {

			delete(server.clients, k)
			return
		}
	}
}
<commit_msg>http server accept port<commit_after>package server

import (
	"github.com/zubairhamed/betwixt/api"
	"github.com/zubairhamed/canopus"
	. "github.com/zubairhamed/go-commons/network"
	"log"
	"net"
	"strings"
)

func NewDefaultCoapServer() *canopus.CoapServer {
	localAddr, err := net.ResolveUDPAddr("udp", ":5683")
	if err != nil {
		log.Println("Error starting CoAP Server: ", err)
	}
	return canopus.NewServer(localAddr, nil)
}

func NewDefaultServer(port string) api.Server {
	return &DefaultServer{
		coapServer: NewDefaultCoapServer(),
		httpServer: NewDefaultHttpServer(port),
		clients:    make(map[string]api.RegisteredClient),
		stats:      &ServerStatistics{},
	}
}

type DefaultServer struct {
	coapServer *canopus.CoapServer
	httpServer *HttpServer
	registry   api.Registry
	stats      *ServerStatistics
	clients    map[string]api.RegisteredClient
}

func (server *DefaultServer) Start() {
	coap := server.coapServer

	// Setup CoAP Routes
	SetupCoapRoutes(server)

	// Start CoAP Server
	go func() {
		coap.Start()
	}()

	// Setup HTTP Routes
	http := server.httpServer
	SetupHttpRoutes(server)

	// Start HTTP Server
	http.Start()
}

func (server *DefaultServer) UseRegistry(reg api.Registry) {
	server.registry = reg
}

func (server *DefaultServer) GetRegisteredClient(id string) api.RegisteredClient {
	return server.clients[id]
}

func (server *DefaultServer) update(id string) {
	for k, v := range server.clients {
		if v.GetId() == id {
			v.Update()
			server.clients[k] = v
		}
	}
}

func (server *DefaultServer) register(ep string, addr string, resources []*canopus.CoreResource) (string, error) {
	clientId := canopus.GenerateToken(8)
	newClient := NewRegisteredClient(ep, clientId, addr)

	objs := make(map[string][]string)
	for _, o := range resources {
		t := o.Target[1:len(o.Target)]
		sp := strings.Split(t, "/")

		if len(sp) > 1 {

		} else {

		}

		log.Println(t, sp[0], sp[1])

		objs[sp[0]] = append(objs[sp[0]], sp[1])
	}

	newClient.SetObjects(objs)

	server.clients[ep] = newClient

	return clientId, nil
}

func (server *DefaultServer) delete(id string) {
	for k, v := range server.clients {
		if v.GetId() == id {

			delete(server.clients, k)
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

/*
This file initializes the
*/

import (
	// General
	"flag"
	"net/http"
	"os"
	"reflect"
	"strconv"

	"github.com/golang/glog"
	_ "github.com/gorilla/mux"
	"github.com/iambc/xerrors"

	//DB
	"database/sql"

	_ "github.com/lib/pq"
)

var api abcAPI

func main() {
	flag.Parse()
	var (
		dbConnString       = os.Getenv("ABC_DB_CONN_STRING")
		filesServerAddr    = os.Getenv("ABC_FILES_SERVER_URL")
		filesDir           = os.Getenv("ABC_FILES_DIR")
		serverEndpointAddr = os.Getenv("ABC_SERVER_ENDPOINT_URL")
	)

	db, err := sql.Open("postgres", dbConnString)
	if err != nil {
		glog.Fatal("Connection to the database has failed")
	}

	api.wr = &writerrdb{db}

	go http.ListenAndServe(":"+filesServerAddr, http.FileServer(http.Dir(filesDir)))
	http.HandleFunc("/api", QueryStringHandler)
	http.ListenAndServe(`:`+serverEndpointAddr, nil)
}

func QueryStringHandler(res http.ResponseWriter, req *http.Request) {
	res.Header().Set("Access-Control-Allow-Origin", "*")
	values := req.URL.Query()
	command, isPassed := values[`command`]
	if !isPassed {
		res.Write([]byte(`{"Status":"error","Msg":"Parameter 'command' is undefined.","Payload":null}`))
		return
	}

	var resp []byte
	var err error
	if command[0] == `getBoards` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.getBoards(apiKey[0])

	} else if command[0] == `getActiveThreadsForBoard` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		boardID, isPassed := values[`board_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'board_Id' is undefined.","Payload":null}`))
			return
		}
		boardIDInt, err := strconv.Atoi(boardID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		resp, err = api.getActiveThreadsForBoard(apiKey[0], boardIDInt)

	} else if command[0] == `getPostsForThread` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		threadID, isPassed := values[`thread_Id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_Id' is undefined.","Payload":null}`))
			return
		}
		threadIDInt, err := strconv.Atoi(threadID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		resp, err = api.getPostsForThread(apiKey[0], threadIDInt)

	} else if command[0] == `addPostToThread` {
		threadID, isPassed := values[`thread_Id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'board_Id' is undefined.","Payload":null}`))
			return
		}
		threadIDInt, err := strconv.Atoi(threadID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		threadBodyPost, isPassed := values[`thread_body_post`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_body_post' is undefined.","Payload":null}`))
			return
		}
		attachmentURL, isPassed := values[`attachment_Id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'attachment_Id' is undefined.","Payload":null}`))
			return
		}
		clientRemoteAddr, isPassed := values[`clientRemoteAddr`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'clientRemoteAddr' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.addPostToThread(threadIDInt, threadBodyPost[0], &attachmentURL[0], clientRemoteAddr[0])

	} else if command[0] == `addThread` {
		boardID, isPassed := values[`board_Id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'board_Id' is undefined.","Payload":null}`))
			return
		}
		boardIDInt, err := strconv.Atoi(boardID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		threadName, isPassed := values[`thread_name`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_name' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.addThread(boardIDInt, threadName[0])

	} else {
		res.Write([]byte(`{"Status":"error","Msg":"No such command exists.","Payload":null}`))
		glog.Error("command: ", command[0])
	}

	if err != nil {
		if string(reflect.TypeOf(err).Name()) == `XError` {
			res.Write([]byte(`{"Status":"` + err.(xerrors.XError).Code + `","Msg":"` + err.Error() + `","Payload":null}`))
		} else {
			res.Write([]byte(`{"Status":"000","Msg":"Application Error!","Payload":null}`))
		}
		glog.Error(err)
		return
	}
	res.Write(resp)

}
<commit_msg>fix: more bugs in parameter checking<commit_after>package main

/*
This file initializes the
*/

import (
	// General
	"flag"
	"net/http"
	"os"
	"reflect"
	"strconv"

	"github.com/golang/glog"
	_ "github.com/gorilla/mux"
	"github.com/iambc/xerrors"

	//DB
	"database/sql"

	_ "github.com/lib/pq"
)

var api abcAPI

func main() {
	flag.Parse()
	var (
		dbConnString       = os.Getenv("ABC_DB_CONN_STRING")
		filesServerAddr    = os.Getenv("ABC_FILES_SERVER_URL")
		filesDir           = os.Getenv("ABC_FILES_DIR")
		serverEndpointAddr = os.Getenv("ABC_SERVER_ENDPOINT_URL")
	)

	db, err := sql.Open("postgres", dbConnString)
	if err != nil {
		glog.Fatal("Connection to the database has failed")
	}

	api.wr = &writerrdb{db}

	go http.ListenAndServe(":"+filesServerAddr, http.FileServer(http.Dir(filesDir)))
	http.HandleFunc("/api", QueryStringHandler)
	http.ListenAndServe(`:`+serverEndpointAddr, nil)
}

func QueryStringHandler(res http.ResponseWriter, req *http.Request) {
	res.Header().Set("Access-Control-Allow-Origin", "*")
	values := req.URL.Query()
	command, isPassed := values[`command`]
	if !isPassed {
		res.Write([]byte(`{"Status":"error","Msg":"Parameter 'command' is undefined.","Payload":null}`))
		return
	}

	var resp []byte
	var err error
	if command[0] == `getBoards` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.getBoards(apiKey[0])

	} else if command[0] == `getActiveThreadsForBoard` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		boardID, isPassed := values[`board_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'board_Id' is undefined.","Payload":null}`))
			return
		}
		boardIDInt, err := strconv.Atoi(boardID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		resp, err = api.getActiveThreadsForBoard(apiKey[0], boardIDInt)

	} else if command[0] == `getPostsForThread` {
		apiKey, isPassed := values[`api_key`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'api_key' is undefined.","Payload":null}`))
			return
		}
		threadID, isPassed := values[`thread_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_id' is undefined.","Payload":null}`))
			return
		}
		threadIDInt, err := strconv.Atoi(threadID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		resp, err = api.getPostsForThread(apiKey[0], threadIDInt)

	} else if command[0] == `addPostToThread` {
		threadID, isPassed := values[`thread_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_id' is undefined.","Payload":null}`))
			return
		}
		threadIDInt, err := strconv.Atoi(threadID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		threadBodyPost, isPassed := values[`thread_body_post`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_body_post' is undefined.","Payload":null}`))
			return
		}
		attachmentURL, isPassed := values[`attachment_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'attachment_id' is undefined.","Payload":null}`))
			return
		}
		clientRemoteAddr, isPassed := values[`clientRemoteAddr`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'clientRemoteAddr' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.addPostToThread(threadIDInt, threadBodyPost[0], &attachmentURL[0], clientRemoteAddr[0])

	} else if command[0] == `addThread` {
		boardID, isPassed := values[`board_id`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'board_id' is undefined.","Payload":null}`))
			return
		}
		boardIDInt, err := strconv.Atoi(boardID[0])
		if err != nil {
			res.Write([]byte(`{"Status":"error","Msg":"Wrong value for parameter board_Id","Payload":null}`))
			return
		}
		threadName, isPassed := values[`thread_name`]
		if !isPassed {
			res.Write([]byte(`{"Status":"error","Msg":"Parameter 'thread_name' is undefined.","Payload":null}`))
			return
		}
		resp, err = api.addThread(boardIDInt, threadName[0])

	} else {
		res.Write([]byte(`{"Status":"error","Msg":"No such command exists.","Payload":null}`))
		glog.Error("command: ", command[0])
	}

	if err != nil {
		if string(reflect.TypeOf(err).Name()) == `XError` {
			res.Write([]byte(`{"Status":"` + err.(xerrors.XError).Code + `","Msg":"` + err.Error() + `","Payload":null}`))
		} else {
			res.Write([]byte(`{"Status":"000","Msg":"Application Error!","Payload":null}`))
		}
		glog.Error(err)
		return
	}
	res.Write(resp)

}
<|endoftext|>


<commit_before>// Copyright 2017 The qurl Authors. All rights reserved.

package server

import (
	"fmt"

	"bitbucket.org/ssgaas/api/routes"
	"github.com/labstack/echo"
)

// Start starts the HTTP server for the qurl API microservice.
func Start(address string, port string) {
	e := echo.New()
	e.DisableHTTP2 = true

	e.GET("/teapot", routes.TeaPot)

	// Start server
	serveraddress := fmt.Sprintf("%s:%s", address, port)
	e.Logger.Fatal(e.Start(serveraddress))
}
<commit_msg>Fix package import<commit_after>// Copyright 2017 The qurl Authors. All rights reserved.

package server

import (
	"fmt"

	"github.com/labstack/echo"
	"github.com/repejota/qurl/routes"
)

// Start starts the HTTP server for the qurl API microservice.
func Start(address string, port string) {
	e := echo.New()
	e.DisableHTTP2 = true

	e.GET("/teapot", routes.TeaPot)

	// Start server
	serveraddress := fmt.Sprintf("%s:%s", address, port)
	e.Logger.Fatal(e.Start(serveraddress))
}
<|endoftext|>


<commit_before>package main

import (
	// General
	"github.com/golang/glog"
	"flag"
	"github.com/iambc/xerrors"

	//API
	"net/http"
	"encoding/json"

	//DB
	"database/sql"
	_ "github.com/lib/pq"
)

type image_board_clusters struct {
    Id int
    Descr string
    LongDescr string
    BoardLimitCount int
}

type boards struct {
    Id int `json:"id"`
    Name string `json:"name"`
    Descr string
    ImageBoardCluster_id string
    BoardLimitCount int
}

type threads struct{
    Id int
    Name string
    Descr string
    Board_id int
    MaxPostsPerThread int
    AreAttachmentsAllowed bool
    LimitsReachedActionId int
}

type thread_posts struct{
    Id int
    Body string
    ThreadId int
    AttachmentUrl int
}

type thread_limits_reached_actions struct{
    Id	    int
    Name    string
    Descr   string
}

type api_request struct{
    Status  string
    Msg	    *string
    Payload interface{}
}

func getBoards(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
	return xerrors.NewSysErr()
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
	return xerrors.NewUiErr(err.Error(), err.Error())
    }
    rows, err := dbh.Query("select id, name from boards;")
    if err != nil {
	return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var curr_boards []boards
    for rows.Next() {
	var board boards
	err = rows.Scan(&board.Id, &board.Name)
	if err != nil {
	    return xerrors.NewUiErr(err.Error(), err.Error())
	}
	curr_boards = append(curr_boards, board)
    }
    bytes, err1 := json.Marshal(api_request{"ok", nil, &curr_boards})
    if err1 != nil {
	return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)
    return nil
}

func getActiveThreadsForBoard(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    board_id, is_passed := values[`board_id`]
    if !is_passed {
	res.Write([]byte(`Invalid params: No board_id given!`))
	return xerrors.NewUiErr(`Invalid params: No board_id given!`, `Invalid params: No board_id given!`)
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    rows, err := dbh.Query("select id, name from threads where board_id = $1;", board_id[0])
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var active_threads []threads
    for rows.Next() {
	glog.Info("Popped new thread")
        var thread threads
        err = rows.Scan(&thread.Id, &thread.Name)
        if err != nil {
            return xerrors.NewUiErr(err.Error(), err.Error())
        }
        active_threads = append(active_threads, thread)
    }
    var bytes []byte
    var err1 error
    if(len(active_threads) == 0){
        errMsg := "No objects returned."
        bytes, err1 = json.Marshal(api_request{"error", &errMsg, &active_threads})
    }else {
        bytes, err1 = json.Marshal(api_request{"ok", nil, &active_threads})
    }

    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

func getPostsForThread(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    thread_id, is_passed := values[`thread_id`]
    if !is_passed {
        res.Write([]byte(`Invalid params: No thread_id given!`))
        return xerrors.NewUiErr(`Invalid params: No thread_id given!`, `Invalid params: No thread_id given!`)
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    rows, err := dbh.Query("select id, body from thread_posts where thread_id = $1;", thread_id[0])
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var curr_posts []thread_posts
    for rows.Next() {
        glog.Info("Popped new thread")
        var curr_post thread_posts
        err = rows.Scan(&curr_post.Id, &curr_post.Body)
        if err != nil {
            return xerrors.NewUiErr(err.Error(), err.Error())
        }
        curr_posts = append(curr_posts, curr_post)
    }

    var bytes []byte
    var err1 error
    if(len(curr_posts) == 0){
	errMsg := "No objects returned."
	bytes, err1 = json.Marshal(api_request{"error", &errMsg, &curr_posts})
    }else {
	bytes, err1 = json.Marshal(api_request{"ok", nil, &curr_posts})
    }

    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

func addPostToThread(res http.ResponseWriter, req *http.Request) error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    thread_id, is_passed := values[`thread_id`]
    if !is_passed {
        res.Write([]byte(`{"Status":"error","Msg":"Invalid params: No thread_id given!","Payload":null}`))
        return xerrors.NewUiErr(`Invalid params: No thread_id given!`, `Invalid params: No thread_id given!`)
    }

    thread_body_post, is_passed := values[`thread_post_body`]
    if !is_passed {
        res.Write([]byte(`{"Status":"error","Msg":"Invalid params: No thread_post_body given!","Payload":null}`))
        return xerrors.NewUiErr(`Invalid params: No thread_post_body given!`, `Invalid params: No thread_post_body given!`)
    }

    attachment_urls, is_passed := values[`attachment_url`]
    var attachment_url *string
    if !is_passed{
	attachment_url = nil
    }else{
	attachment_url = &attachment_urls[0]
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    _, err = dbh.Query("INSERT INTO thread_posts(body, thread_id, attachment_url) VALUES($1, $2, $3)", thread_body_post[0], thread_id[0], attachment_url)

    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

// sample usage
func main() {
    flag.Parse()

    commands := map[string]func(http.ResponseWriter, *http.Request) error{
				"getBoards": getBoards,
				"getActiveThreadsForBoard": getActiveThreadsForBoard,
				"getPostsForThread": getPostsForThread,
				"addPostToThread": addPostToThread,
			       }

    http.HandleFunc("/api", func(res http.ResponseWriter, req *http.Request) {
					values := req.URL.Query()
					command, is_passed := values[`command`]
					if !is_passed {
					    res.Write([]byte(`{"Status":"error","Msg":"Paremeter 'command' is undefined.","Payload":null}`))
					    return
					}
					_, is_passed = commands[command[0]]
					if !is_passed{
					    res.Write([]byte(`{"Status":"error","Msg":"No such command exists.","Payload":null}`))
					    return
					}

					err := commands[command[0]](res, req)
					if err != nil{
					    glog.Error(err)
					}
    })

    http.ListenAndServe(`:8089`, nil)
}

<commit_msg>chore: added todos<commit_after>package main

import (
	// General
	"github.com/golang/glog"
	"flag"
	"github.com/iambc/xerrors"

	//API
	"net/http"
	"encoding/json"

	//DB
	"database/sql"
	_ "github.com/lib/pq"
)

/*
TODO:
1) Add api_key, api_authentication
2) Add different input/output formats for the API
3) Start adding settings to the boards, etc
*/

type image_board_clusters struct {
    Id int
    Descr string
    LongDescr string
    BoardLimitCount int
}

type boards struct {
    Id int `json:"id"`
    Name string `json:"name"`
    Descr string
    ImageBoardCluster_id string
    BoardLimitCount int
}

type threads struct{
    Id int
    Name string
    Descr string
    Board_id int
    MaxPostsPerThread int
    AreAttachmentsAllowed bool
    LimitsReachedActionId int
}

type thread_posts struct{
    Id int
    Body string
    ThreadId int
    AttachmentUrl int
}

type thread_limits_reached_actions struct{
    Id	    int
    Name    string
    Descr   string
}

type api_request struct{
    Status  string
    Msg	    *string
    Payload interface{}
}

func getBoards(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
	return xerrors.NewSysErr()
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
	return xerrors.NewUiErr(err.Error(), err.Error())
    }
    rows, err := dbh.Query("select id, name from boards;")
    if err != nil {
	return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var curr_boards []boards
    for rows.Next() {
	var board boards
	err = rows.Scan(&board.Id, &board.Name)
	if err != nil {
	    return xerrors.NewUiErr(err.Error(), err.Error())
	}
	curr_boards = append(curr_boards, board)
    }
    bytes, err1 := json.Marshal(api_request{"ok", nil, &curr_boards})
    if err1 != nil {
	return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)
    return nil
}

func getActiveThreadsForBoard(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    board_id, is_passed := values[`board_id`]
    if !is_passed {
	res.Write([]byte(`Invalid params: No board_id given!`))
	return xerrors.NewUiErr(`Invalid params: No board_id given!`, `Invalid params: No board_id given!`)
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    rows, err := dbh.Query("select id, name from threads where board_id = $1;", board_id[0])
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var active_threads []threads
    for rows.Next() {
	glog.Info("Popped new thread")
        var thread threads
        err = rows.Scan(&thread.Id, &thread.Name)
        if err != nil {
            return xerrors.NewUiErr(err.Error(), err.Error())
        }
        active_threads = append(active_threads, thread)
    }
    var bytes []byte
    var err1 error
    if(len(active_threads) == 0){
        errMsg := "No objects returned."
        bytes, err1 = json.Marshal(api_request{"error", &errMsg, &active_threads})
    }else {
        bytes, err1 = json.Marshal(api_request{"ok", nil, &active_threads})
    }

    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

func getPostsForThread(res http.ResponseWriter, req *http.Request)  error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    thread_id, is_passed := values[`thread_id`]
    if !is_passed {
        res.Write([]byte(`Invalid params: No thread_id given!`))
        return xerrors.NewUiErr(`Invalid params: No thread_id given!`, `Invalid params: No thread_id given!`)
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    rows, err := dbh.Query("select id, body from thread_posts where thread_id = $1;", thread_id[0])
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }
    defer rows.Close()

    var curr_posts []thread_posts
    for rows.Next() {
        glog.Info("Popped new thread")
        var curr_post thread_posts
        err = rows.Scan(&curr_post.Id, &curr_post.Body)
        if err != nil {
            return xerrors.NewUiErr(err.Error(), err.Error())
        }
        curr_posts = append(curr_posts, curr_post)
    }

    var bytes []byte
    var err1 error
    if(len(curr_posts) == 0){
	errMsg := "No objects returned."
	bytes, err1 = json.Marshal(api_request{"error", &errMsg, &curr_posts})
    }else {
	bytes, err1 = json.Marshal(api_request{"ok", nil, &curr_posts})
    }

    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

func addPostToThread(res http.ResponseWriter, req *http.Request) error {
    if req == nil {
        return xerrors.NewSysErr()
    }
    values := req.URL.Query()
    thread_id, is_passed := values[`thread_id`]
    if !is_passed {
        res.Write([]byte(`{"Status":"error","Msg":"Invalid params: No thread_id given!","Payload":null}`))
        return xerrors.NewUiErr(`Invalid params: No thread_id given!`, `Invalid params: No thread_id given!`)
    }

    thread_body_post, is_passed := values[`thread_post_body`]
    if !is_passed {
        res.Write([]byte(`{"Status":"error","Msg":"Invalid params: No thread_post_body given!","Payload":null}`))
        return xerrors.NewUiErr(`Invalid params: No thread_post_body given!`, `Invalid params: No thread_post_body given!`)
    }

    attachment_urls, is_passed := values[`attachment_url`]
    var attachment_url *string
    if !is_passed{
	attachment_url = nil
    }else{
	attachment_url = &attachment_urls[0]
    }

    dbh, err := sql.Open("postgres", "user=abc_api password=123 dbname=abc_dev_cluster sslmode=disable")
    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    _, err = dbh.Query("INSERT INTO thread_posts(body, thread_id, attachment_url) VALUES($1, $2, $3)", thread_body_post[0], thread_id[0], attachment_url)

    if err != nil {
        return xerrors.NewUiErr(err.Error(), err.Error())
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return xerrors.NewUiErr(err1.Error(), err1.Error())
    }
    res.Write(bytes)

    return nil
}

// sample usage
func main() {
    flag.Parse()

    commands := map[string]func(http.ResponseWriter, *http.Request) error{
				"getBoards": getBoards,
				"getActiveThreadsForBoard": getActiveThreadsForBoard,
				"getPostsForThread": getPostsForThread,
				"addPostToThread": addPostToThread,
			       }

    http.HandleFunc("/api", func(res http.ResponseWriter, req *http.Request) {
					values := req.URL.Query()
					command, is_passed := values[`command`]
					if !is_passed {
					    res.Write([]byte(`{"Status":"error","Msg":"Paremeter 'command' is undefined.","Payload":null}`))
					    return
					}
					_, is_passed = commands[command[0]]
					if !is_passed{
					    res.Write([]byte(`{"Status":"error","Msg":"No such command exists.","Payload":null}`))
					    return
					}

					err := commands[command[0]](res, req)
					if err != nil{
					    glog.Error(err)
					}
    })

    http.ListenAndServe(`:8089`, nil)
}

<|endoftext|>


<commit_before>// Copyright 2017 The Go 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 gke_test

import (
	"context"
	"strings"
	"testing"

	"cloud.google.com/go/compute/metadata"
	"golang.org/x/build/kubernetes"
	"golang.org/x/build/kubernetes/gke"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
	container "google.golang.org/api/container/v1"
)

// Tests NewClient and also Dialer.
func TestNewClient(t *testing.T) {
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		_, err := kc.GetPods(ctx)
		if err != nil {
			t.Fatal(err)
		}
	})
}

func TestDialPod(t *testing.T) {
	var passed bool
	var candidates int
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		pods, err := kc.GetPods(ctx)
		if err != nil {
			t.Fatal(err)
		}

		for _, pod := range pods {
			if pod.Status.Phase != "Running" {
				continue
			}
			for _, container := range pod.Spec.Containers {
				for _, port := range container.Ports {
					if strings.ToLower(string(port.Protocol)) == "udp" || port.ContainerPort == 0 {
						continue
					}
					candidates++
					c, err := kc.DialPod(ctx, pod.Name, port.ContainerPort)
					if err != nil {
						t.Logf("Dial %q/%q/%d: %v", cl.Name, pod.Name, port.ContainerPort, err)
						continue
					}
					c.Close()
					t.Logf("Dialed %q/%q/%d.", cl.Name, pod.Name, port.ContainerPort)
					passed = true
					return
				}
			}
		}
	})
	if candidates == 0 {
		t.Skip("no pods to dial")
	}
	if !passed {
		t.Errorf("dial failures")
	}
}

func TestDialService(t *testing.T) {
	var passed bool
	var candidates int
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		svcs, err := kc.GetServices(ctx)
		if err != nil {
			t.Fatal(err)
		}
		for _, svc := range svcs {
			eps, err := kc.GetServiceEndpoints(ctx, svc.Name, "")
			if err != nil {
				t.Fatal(err)
			}
			if len(eps) != 1 {
				continue
			}
			candidates++
			conn, err := kc.DialServicePort(ctx, svc.Name, "")
			if err != nil {
				t.Logf("%s: DialServicePort(%q) error: %v", cl.Name, svc.Name, err)
				continue
			}
			conn.Close()
			passed = true
			t.Logf("Dialed cluster %q service %q.", cl.Name, svc.Name)
			return
		}

	})
	if candidates == 0 {
		t.Skip("no services to dial")
	}
	if !passed {
		t.Errorf("dial failures")
	}
}

func foreachCluster(t *testing.T, fn func(*container.Cluster, *kubernetes.Client)) {
	if !metadata.OnGCE() {
		t.Skip("not on GCE; skipping")
	}
	ctx := context.Background()
	ts, err := google.DefaultTokenSource(ctx, compute.CloudPlatformScope)
	if err != nil {
		t.Fatal(err)
	}
	httpClient := oauth2.NewClient(ctx, ts)
	containerService, err := container.New(httpClient)
	if err != nil {
		t.Fatal(err)
	}
	proj, err := metadata.ProjectID()
	if err != nil {
		t.Fatal(err)
	}

	clusters, err := containerService.Projects.Zones.Clusters.List(proj, "-").Context(ctx).Do()
	if err != nil {
		t.Fatal(err)
	}

	if len(clusters.Clusters) == 0 {
		t.Skip("no GKE clusters")
	}
	for _, cl := range clusters.Clusters {
		kc, err := gke.NewClient(ctx, cl.Name, gke.OptZone(cl.Zone))
		if err != nil {
			t.Fatal(err)
		}
		fn(cl, kc)
		kc.Close()
	}
}

func TestGetNodes(t *testing.T) {
	var passed bool
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		nodes, err := kc.GetNodes(ctx)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("%d nodes in cluster %s", len(nodes), cl.Name)
		passed = true
	})
}
<commit_msg>kubernetes/gke: skip GKE tests if no metadata present<commit_after>// Copyright 2017 The Go 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 gke_test

import (
	"context"
	"strings"
	"testing"

	"cloud.google.com/go/compute/metadata"
	"golang.org/x/build/kubernetes"
	"golang.org/x/build/kubernetes/gke"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"
	container "google.golang.org/api/container/v1"
)

// Tests NewClient and also Dialer.
func TestNewClient(t *testing.T) {
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		_, err := kc.GetPods(ctx)
		if err != nil {
			t.Fatal(err)
		}
	})
}

func TestDialPod(t *testing.T) {
	var passed bool
	var candidates int
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		pods, err := kc.GetPods(ctx)
		if err != nil {
			t.Fatal(err)
		}

		for _, pod := range pods {
			if pod.Status.Phase != "Running" {
				continue
			}
			for _, container := range pod.Spec.Containers {
				for _, port := range container.Ports {
					if strings.ToLower(string(port.Protocol)) == "udp" || port.ContainerPort == 0 {
						continue
					}
					candidates++
					c, err := kc.DialPod(ctx, pod.Name, port.ContainerPort)
					if err != nil {
						t.Logf("Dial %q/%q/%d: %v", cl.Name, pod.Name, port.ContainerPort, err)
						continue
					}
					c.Close()
					t.Logf("Dialed %q/%q/%d.", cl.Name, pod.Name, port.ContainerPort)
					passed = true
					return
				}
			}
		}
	})
	if candidates == 0 {
		t.Skip("no pods to dial")
	}
	if !passed {
		t.Errorf("dial failures")
	}
}

func TestDialService(t *testing.T) {
	var passed bool
	var candidates int
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		svcs, err := kc.GetServices(ctx)
		if err != nil {
			t.Fatal(err)
		}
		for _, svc := range svcs {
			eps, err := kc.GetServiceEndpoints(ctx, svc.Name, "")
			if err != nil {
				t.Fatal(err)
			}
			if len(eps) != 1 {
				continue
			}
			candidates++
			conn, err := kc.DialServicePort(ctx, svc.Name, "")
			if err != nil {
				t.Logf("%s: DialServicePort(%q) error: %v", cl.Name, svc.Name, err)
				continue
			}
			conn.Close()
			passed = true
			t.Logf("Dialed cluster %q service %q.", cl.Name, svc.Name)
			return
		}

	})
	if candidates == 0 {
		t.Skip("no services to dial")
	}
	if !passed {
		t.Errorf("dial failures")
	}
}

func foreachCluster(t *testing.T, fn func(*container.Cluster, *kubernetes.Client)) {
	if !metadata.OnGCE() {
		t.Skip("not on GCE; skipping")
	}
	ctx := context.Background()
	ts, err := google.DefaultTokenSource(ctx, compute.CloudPlatformScope)
	if err != nil {
		t.Fatal(err)
	}
	httpClient := oauth2.NewClient(ctx, ts)
	containerService, err := container.New(httpClient)
	if err != nil {
		t.Fatal(err)
	}
	proj, err := metadata.ProjectID()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := ts.Token(); err != nil {
		val, err := metadata.InstanceAttributeValue("service-accounts/default/token")
		if val == "" {
			t.Skip("skipping on GCE instance without a service account")
		}
		t.Skipf("default token source doesn't work; skipping test: %v", err)
	}

	clusters, err := containerService.Projects.Zones.Clusters.List(proj, "-").Context(ctx).Do()
	if err != nil {
		t.Fatal(err)
	}

	if len(clusters.Clusters) == 0 {
		t.Skip("no GKE clusters")
	}
	for _, cl := range clusters.Clusters {
		kc, err := gke.NewClient(ctx, cl.Name, gke.OptZone(cl.Zone))
		if err != nil {
			t.Fatal(err)
		}
		fn(cl, kc)
		kc.Close()
	}
}

func TestGetNodes(t *testing.T) {
	var passed bool
	ctx := context.Background()
	foreachCluster(t, func(cl *container.Cluster, kc *kubernetes.Client) {
		if passed {
			return
		}
		nodes, err := kc.GetNodes(ctx)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("%d nodes in cluster %s", len(nodes), cl.Name)
		passed = true
	})
}
<|endoftext|>