<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"
    "log"
    "github.com/urfave/cli"
)

const nbd_folder = "/sample_disks/"

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

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) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    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 (options & utils.NBD_FLAG_NO_ZEROES) != utils.NBD_FLAG_NO_ZEROES {
        offset += 124               // pad with 124 zeroes
    }

    _, 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")

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, 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)
            //fmt.Printf("Done reading the data that should be written")

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

            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
        }
    }
}

func send_export_list(output *bufio.Writer, options uint32) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(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
    }

    host := "localhost"
    port := "8000"
    listen := ""
    app.Flags = []cli.Flag {
        cli.StringFlag{
            Name: "host",
            Value: host,
            Usage: "Hostname or IP address you want to serve traffic on. e.x. 'localhost', '192.168.1.2'",
            Destination: &host,
        },
        cli.StringFlag{
            Name: "port",
            Value: port,
            Usage: "Port you want to serve traffic on. e.x. '8000'",
            Destination: &port,
        },
        cli.StringFlag{
            Name: "listen, l",
            Destination: &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'",
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    hostingAddress := listen
    if len(listen) == 0 {
        hostingAddress = host + ":" + 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

        // S: handshake flags
        //output.Write([]byte{0, 3})          // Ubuntu
        output.Write(defaultOptions)        // Qemu
        // END: S: handshake flags

        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)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options)
            break
        }
    }

}
<commit_msg>cleaning up the spacing options code<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"
    "log"
    "github.com/urfave/cli"
)

const nbd_folder = "/sample_disks/"

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

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) {
    fmt.Printf("have request to bind to: %s\n", string(payload[:payload_size]))

    defer conn.Close()

    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 (options & utils.NBD_FLAG_NO_ZEROES) != utils.NBD_FLAG_NO_ZEROES {
        offset += 124               // pad with 124 zeroes
    }

    _, 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")

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, 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)
            //fmt.Printf("Done reading the data that should be written")

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

            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
        }
    }
}

func send_export_list(output *bufio.Writer, options uint32) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(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
    }

    host := "localhost"
    port := "8000"
    listen := ""
    app.Flags = []cli.Flag {
        cli.StringFlag{
            Name: "host",
            Value: host,
            Usage: "Hostname or IP address you want to serve traffic on. e.x. 'localhost', '192.168.1.2'",
            Destination: &host,
        },
        cli.StringFlag{
            Name: "port",
            Value: port,
            Usage: "Port you want to serve traffic on. e.x. '8000'",
            Destination: &port,
        },
        cli.StringFlag{
            Name: "listen, l",
            Destination: &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'",
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    hostingAddress := listen
    if len(listen) == 0 {
        hostingAddress = host + ":" + 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)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options)
            break
        }
    }

}
<|endoftext|>


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

// Package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.

package server

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"crypto/tls"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/boltdb"
	"github.com/pilosa/pilosa/gcnotify"
	"github.com/pilosa/pilosa/gopsutil"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/statik"
	"github.com/pilosa/pilosa/statsd"
	"github.com/pkg/errors"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

type loggerLogger interface {
	pilosa.Logger
	Logger() *log.Logger
}

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *Config

	// Gossip transport
	GossipTransport *gossip.Transport

	// Standard input/output
	*pilosa.CmdIO

	// Started will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}

	// Passed to the Gossip implementation.
	logOutput io.Writer
	logger    loggerLogger
}

// NewCommand returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Config: NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) { // TODO args WTF
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// SetupNetworking
	err = m.SetupNetworking()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}

	m.logger.Printf("Listening as %s\n", m.Server.URI)

	return nil
}

// Wait waits for the server to be closed or interrupted.
func (m *Command) Wait() error {
	// First SIGKILL causes server to shut down gracefully.
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	select {
	case sig := <-c:
		m.logger.Printf("Received %s; gracefully shutting down...\n", sig.String())

		// Second signal causes a hard shutdown.
		go func() { <-c; os.Exit(1) }()
		return errors.Wrap(m.Close(), "closing command")
	case <-m.Done:
		m.logger.Printf("Server closed externally")
		return nil
	}
}

// SetupLogger sets up the logger based on the configuration.
func (m *Command) SetupLogger() (pilosa.Logger, error) {
	if m.logger != nil {
		return m.logger, nil
	}
	var err error
	if m.Config.LogPath == "" {
		m.logOutput = m.Stderr
	} else {
		m.logOutput, err = os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return nil, errors.Wrap(err, "opening file")
		}
	}

	if m.Config.Verbose {
		m.logger = pilosa.NewVerboseLogger(m.logOutput)
	} else {
		m.logger = pilosa.NewStandardLogger(m.logOutput)
	}
	return m.logger, nil
}

// SetupServer uses the cluster configuration to set up this server.
func (m *Command) SetupServer() error {
	if m.logger == nil {
		_, err := m.SetupLogger()
		if err != nil {
			return errors.Wrap(err, "setting up logger")
		}

	}

	handler := pilosa.NewHandler()
	handler.Logger = m.logger
	handler.FileSystem = &statik.FileSystem{}
	handler.API = pilosa.NewAPI()
	handler.API.Logger = m.logger

	uri, err := pilosa.AddressWithDefaults(m.Config.Bind)
	if err != nil {
		return errors.Wrap(err, "processing bind address")
	}

	// Setup TLS
	var TLSConfig *tls.Config
	if uri.Scheme() == "https" {
		if m.Config.TLS.CertificatePath == "" {
			return errors.New("certificate path is required for TLS sockets")
		}
		if m.Config.TLS.CertificateKeyPath == "" {
			return errors.New("certificate key path is required for TLS sockets")
		}
		cert, err := tls.LoadX509KeyPair(m.Config.TLS.CertificatePath, m.Config.TLS.CertificateKeyPath)
		if err != nil {
			return errors.Wrap(err, "load x509 key pair")
		}
		TLSConfig = &tls.Config{
			Certificates:       []tls.Certificate{cert},
			InsecureSkipVerify: m.Config.TLS.SkipVerify,
		}
	}

	diagnosticsInterval := time.Duration(0)
	if m.Config.Metric.Diagnostics {
		diagnosticsInterval = time.Duration(DefaultDiagnosticsInterval)
	}

	statsClient, err := NewStatsClient(m.Config.Metric.Service, m.Config.Metric.Host)
	if err != nil {
		return errors.Wrap(err, "new stats client")
	}

	ln, err := getListener(*uri, TLSConfig)
	if err != nil {
		return errors.Wrap(err, "getting listener")
	}

	c := GetHTTPClient(TLSConfig)
	handler.API.RemoteClient = c

	m.Server, err = pilosa.NewServer(
		pilosa.OptServerAntiEntropyInterval(time.Duration(m.Config.AntiEntropy.Interval)),
		pilosa.OptServerLongQueryTime(time.Duration(m.Config.Cluster.LongQueryTime)),
		pilosa.OptServerDataDir(m.Config.DataDir),
		pilosa.OptServerReplicaN(m.Config.Cluster.ReplicaN),
		pilosa.OptServerMaxWritesPerRequest(m.Config.MaxWritesPerRequest),
		pilosa.OptServerMetricInterval(time.Duration(m.Config.Metric.PollInterval)),
		pilosa.OptServerDiagnosticsInterval(diagnosticsInterval),

		pilosa.OptServerLogger(m.logger),
		pilosa.OptServerAttrStoreFunc(boltdb.NewAttrStore),
		pilosa.OptServerHandler(handler),
		pilosa.OptServerSystemInfo(gopsutil.NewSystemInfo()),
		pilosa.OptServerGCNotifier(gcnotify.NewActiveGCNotifier()),
		pilosa.OptServerStatsClient(statsClient),
		pilosa.OptServerListener(ln),
		pilosa.OptServerURI(uri),
		pilosa.OptServerRemoteClient(c),
	)

	return errors.Wrap(err, "new server")
}

func GetHTTPClient(t *tls.Config) *http.Client {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext,
		MaxIdleConns:          1000,
		MaxIdleConnsPerHost:   200,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	if t != nil {
		transport.TLSClientConfig = t
	}
	return &http.Client{Transport: transport}
}

// SetupNetworking sets up internode communication based on the configuration.
func (m *Command) SetupNetworking() error {

	m.Server.NodeID = m.Server.LoadNodeID()

	if m.Config.Cluster.Disabled {
		m.Server.Cluster.Static = true
		m.Server.Cluster.Coordinator = m.Server.NodeID
		for _, address := range m.Config.Cluster.Hosts {
			uri, err := pilosa.NewURIFromAddress(address)
			if err != nil {
				return err
			}
			m.Server.Cluster.Nodes = append(m.Server.Cluster.Nodes, &pilosa.Node{
				URI: *uri,
			})
		}

		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.MemberSet = pilosa.NewStaticMemberSet(m.Server.Cluster.Nodes)
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		m.Server.Gossiper = pilosa.NopGossiper
		return nil
	}

	gossipPort, err := strconv.Atoi(m.Config.Gossip.Port)
	if err != nil {
		return err
	}

	// get the host portion of addr to use for binding
	gossipHost := m.Server.URI.Host()
	var transport *gossip.Transport
	if m.GossipTransport != nil {
		transport = m.GossipTransport
	} else {
		transport, err = gossip.NewTransport(gossipHost, gossipPort, m.logger.Logger())
		if err != nil {
			return err
		}
	}

	// Set Coordinator.
	if m.Config.Cluster.Coordinator || len(m.Config.Gossip.Seeds) == 0 {
		m.Server.Cluster.Coordinator = m.Server.NodeID
	}

	gossipEventReceiver := gossip.NewGossipEventReceiver(m.logger)
	m.Server.Cluster.EventReceiver = gossipEventReceiver
	gossipMemberSet, err := gossip.NewGossipMemberSet(m.Server.NodeID, m.Server.URI.Host(), m.Config.Gossip, gossipEventReceiver, m.Server, gossip.WithLogger(m.logger.Logger()), gossip.WithTransport(transport))
	if err != nil {
		return err
	}
	m.Server.Cluster.MemberSet = gossipMemberSet
	m.Server.Broadcaster = m.Server
	m.Server.BroadcastReceiver = gossipMemberSet
	m.Server.Gossiper = gossipMemberSet
	return nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	if closer, ok := m.logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}

// NewStatsClient creates a stats client from the config
func NewStatsClient(name string, host string) (pilosa.StatsClient, error) {
	switch name {
	case "expvar":
		return pilosa.NewExpvarStatsClient(), nil
	case "statsd":
		return statsd.NewStatsClient(host)
	case "nop", "none":
		return pilosa.NopStatsClient, nil
	default:
		return nil, errors.Errorf("'%v' not a valid stats client, choose from [expvar, statsd, none].")
	}
}

// getListener gets a net.Listener based on the config.
func getListener(uri pilosa.URI, tlsconf *tls.Config) (ln net.Listener, err error) {
	// If bind URI has the https scheme, enable TLS
	if uri.Scheme() == "https" && tlsconf != nil {
		ln, err = tls.Listen("tcp", uri.HostPort(), tlsconf)
		if err != nil {
			return nil, errors.Wrap(err, "tls.Listener")
		}
	} else if uri.Scheme() == "http" {
		// Open HTTP listener to determine port (if specified as :0).
		ln, err = net.Listen("tcp", uri.HostPort())
		if err != nil {
			return nil, errors.Wrap(err, "net.Listen")
		}
	} else {
		return nil, errors.Errorf("unsupported scheme: %s", uri.Scheme())
	}

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

// Package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.

package server

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"crypto/tls"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/boltdb"
	"github.com/pilosa/pilosa/gcnotify"
	"github.com/pilosa/pilosa/gopsutil"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/statik"
	"github.com/pilosa/pilosa/statsd"
	"github.com/pkg/errors"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

type loggerLogger interface {
	pilosa.Logger
	Logger() *log.Logger
}

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *Config

	// Gossip transport
	GossipTransport *gossip.Transport

	// Standard input/output
	*pilosa.CmdIO

	// Started will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}

	// Passed to the Gossip implementation.
	logOutput io.Writer
	logger    loggerLogger
}

// NewCommand returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Config: NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) { // TODO args WTF
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// SetupNetworking
	err = m.SetupNetworking()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}

	m.logger.Printf("Listening as %s\n", m.Server.URI)

	return nil
}

// Wait waits for the server to be closed or interrupted.
func (m *Command) Wait() error {
	// First SIGKILL causes server to shut down gracefully.
	c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	select {
	case sig := <-c:
		m.logger.Printf("Received %s; gracefully shutting down...\n", sig.String())

		// Second signal causes a hard shutdown.
		go func() { <-c; os.Exit(1) }()
		return errors.Wrap(m.Close(), "closing command")
	case <-m.Done:
		m.logger.Printf("Server closed externally")
		return nil
	}
}

// setupLogger sets up the logger based on the configuration.
func (m *Command) setupLogger() error {
	var err error
	if m.Config.LogPath == "" {
		m.logOutput = m.Stderr
	} else {
		m.logOutput, err = os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return errors.Wrap(err, "opening file")
		}
	}

	if m.Config.Verbose {
		m.logger = pilosa.NewVerboseLogger(m.logOutput)
	} else {
		m.logger = pilosa.NewStandardLogger(m.logOutput)
	}
	return nil
}

// SetupServer uses the cluster configuration to set up this server.
func (m *Command) SetupServer() error {
	err := m.setupLogger()
	if err != nil {
		return errors.Wrap(err, "setting up logger")
	}

	handler := pilosa.NewHandler()
	handler.Logger = m.logger
	handler.FileSystem = &statik.FileSystem{}
	handler.API = pilosa.NewAPI()
	handler.API.Logger = m.logger

	uri, err := pilosa.AddressWithDefaults(m.Config.Bind)
	if err != nil {
		return errors.Wrap(err, "processing bind address")
	}

	// Setup TLS
	var TLSConfig *tls.Config
	if uri.Scheme() == "https" {
		if m.Config.TLS.CertificatePath == "" {
			return errors.New("certificate path is required for TLS sockets")
		}
		if m.Config.TLS.CertificateKeyPath == "" {
			return errors.New("certificate key path is required for TLS sockets")
		}
		cert, err := tls.LoadX509KeyPair(m.Config.TLS.CertificatePath, m.Config.TLS.CertificateKeyPath)
		if err != nil {
			return errors.Wrap(err, "load x509 key pair")
		}
		TLSConfig = &tls.Config{
			Certificates:       []tls.Certificate{cert},
			InsecureSkipVerify: m.Config.TLS.SkipVerify,
		}
	}

	diagnosticsInterval := time.Duration(0)
	if m.Config.Metric.Diagnostics {
		diagnosticsInterval = time.Duration(DefaultDiagnosticsInterval)
	}

	statsClient, err := NewStatsClient(m.Config.Metric.Service, m.Config.Metric.Host)
	if err != nil {
		return errors.Wrap(err, "new stats client")
	}

	ln, err := getListener(*uri, TLSConfig)
	if err != nil {
		return errors.Wrap(err, "getting listener")
	}

	c := GetHTTPClient(TLSConfig)
	handler.API.RemoteClient = c

	m.Server, err = pilosa.NewServer(
		pilosa.OptServerAntiEntropyInterval(time.Duration(m.Config.AntiEntropy.Interval)),
		pilosa.OptServerLongQueryTime(time.Duration(m.Config.Cluster.LongQueryTime)),
		pilosa.OptServerDataDir(m.Config.DataDir),
		pilosa.OptServerReplicaN(m.Config.Cluster.ReplicaN),
		pilosa.OptServerMaxWritesPerRequest(m.Config.MaxWritesPerRequest),
		pilosa.OptServerMetricInterval(time.Duration(m.Config.Metric.PollInterval)),
		pilosa.OptServerDiagnosticsInterval(diagnosticsInterval),

		pilosa.OptServerLogger(m.logger),
		pilosa.OptServerAttrStoreFunc(boltdb.NewAttrStore),
		pilosa.OptServerHandler(handler),
		pilosa.OptServerSystemInfo(gopsutil.NewSystemInfo()),
		pilosa.OptServerGCNotifier(gcnotify.NewActiveGCNotifier()),
		pilosa.OptServerStatsClient(statsClient),
		pilosa.OptServerListener(ln),
		pilosa.OptServerURI(uri),
		pilosa.OptServerRemoteClient(c),
	)

	return errors.Wrap(err, "new server")
}

func GetHTTPClient(t *tls.Config) *http.Client {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext,
		MaxIdleConns:          1000,
		MaxIdleConnsPerHost:   200,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	if t != nil {
		transport.TLSClientConfig = t
	}
	return &http.Client{Transport: transport}
}

// SetupNetworking sets up internode communication based on the configuration.
func (m *Command) SetupNetworking() error {

	m.Server.NodeID = m.Server.LoadNodeID()

	if m.Config.Cluster.Disabled {
		m.Server.Cluster.Static = true
		m.Server.Cluster.Coordinator = m.Server.NodeID
		for _, address := range m.Config.Cluster.Hosts {
			uri, err := pilosa.NewURIFromAddress(address)
			if err != nil {
				return err
			}
			m.Server.Cluster.Nodes = append(m.Server.Cluster.Nodes, &pilosa.Node{
				URI: *uri,
			})
		}

		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.MemberSet = pilosa.NewStaticMemberSet(m.Server.Cluster.Nodes)
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		m.Server.Gossiper = pilosa.NopGossiper
		return nil
	}

	gossipPort, err := strconv.Atoi(m.Config.Gossip.Port)
	if err != nil {
		return err
	}

	// get the host portion of addr to use for binding
	gossipHost := m.Server.URI.Host()
	var transport *gossip.Transport
	if m.GossipTransport != nil {
		transport = m.GossipTransport
	} else {
		transport, err = gossip.NewTransport(gossipHost, gossipPort, m.logger.Logger())
		if err != nil {
			return err
		}
	}

	// Set Coordinator.
	if m.Config.Cluster.Coordinator || len(m.Config.Gossip.Seeds) == 0 {
		m.Server.Cluster.Coordinator = m.Server.NodeID
	}

	gossipEventReceiver := gossip.NewGossipEventReceiver(m.logger)
	m.Server.Cluster.EventReceiver = gossipEventReceiver
	gossipMemberSet, err := gossip.NewGossipMemberSet(m.Server.NodeID, m.Server.URI.Host(), m.Config.Gossip, gossipEventReceiver, m.Server, gossip.WithLogger(m.logger.Logger()), gossip.WithTransport(transport))
	if err != nil {
		return err
	}
	m.Server.Cluster.MemberSet = gossipMemberSet
	m.Server.Broadcaster = m.Server
	m.Server.BroadcastReceiver = gossipMemberSet
	m.Server.Gossiper = gossipMemberSet
	return nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	if closer, ok := m.logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}

// NewStatsClient creates a stats client from the config
func NewStatsClient(name string, host string) (pilosa.StatsClient, error) {
	switch name {
	case "expvar":
		return pilosa.NewExpvarStatsClient(), nil
	case "statsd":
		return statsd.NewStatsClient(host)
	case "nop", "none":
		return pilosa.NopStatsClient, nil
	default:
		return nil, errors.Errorf("'%v' not a valid stats client, choose from [expvar, statsd, none].")
	}
}

// getListener gets a net.Listener based on the config.
func getListener(uri pilosa.URI, tlsconf *tls.Config) (ln net.Listener, err error) {
	// If bind URI has the https scheme, enable TLS
	if uri.Scheme() == "https" && tlsconf != nil {
		ln, err = tls.Listen("tcp", uri.HostPort(), tlsconf)
		if err != nil {
			return nil, errors.Wrap(err, "tls.Listener")
		}
	} else if uri.Scheme() == "http" {
		// Open HTTP listener to determine port (if specified as :0).
		ln, err = net.Listen("tcp", uri.HostPort())
		if err != nil {
			return nil, errors.Wrap(err, "net.Listen")
		}
	} else {
		return nil, errors.Errorf("unsupported scheme: %s", uri.Scheme())
	}

	return ln, nil
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"
)

// FallbackServer is a grpc-fallback HTTP server.
type FallbackServer struct {
	backend string
	server  http.Server
	cc      *grpc.ClientConn
}

// NewServer creates a new grpc-fallback HTTP server on the
// given port that proxies to the given gRPC server backend.
func NewServer(port, backend string) *FallbackServer {
	if !strings.HasPrefix(port, ":") {
		port = ":" + port
	}

	return &FallbackServer{
		backend: backend,
		server: http.Server{
			Addr: port,
		},
	}
}

// Start starts the grpc-fallback HTTP server listening on its port,
// and opens a connection to the gRPC backend.
func (f *FallbackServer) Start() {
	var err error

	// setup connection to gRPC backend
	f.cc, err = f.dial()
	if err != nil {
		log.Fatal("Error dialing gRPC backend server:", err)
	}

	// setup grpc-fallback complient router
	r := mux.NewRouter()
	r.HandleFunc("/$rpc/{service:[.a-zA-Z0-9]+}/{method:[a-zA-Z]+}", f.handler).Headers("Content-Type", "application/x-protobuf")
	f.server.Handler = r

	log.Println("Fallback server listening on port:", f.server.Addr)
	err = f.server.ListenAndServe()
	if err != nil {
		log.Println("Error in fallback server while listening:", err)
	}
}

// StartBackground runs Start() in a goroutine.
func (f *FallbackServer) StartBackground() {
	go f.Start()
}

// Shutdown turns down the grpc-fallback HTTP server.
func (f *FallbackServer) Shutdown() {
	err := f.server.Shutdown(context.Background())
	if err != nil {
		log.Println("Error shutting down fallback server:", err)
	}
}

// handler is a generic HTTP handler that invokes the proper
// RPC given the grpc-fallback HTTP request.
func (f *FallbackServer) handler(w http.ResponseWriter, r *http.Request) {
	log.Println("Incoming grpc-fallback request:", r.RequestURI)
	v := mux.Vars(r)
	service := v["service"]
	method := v["method"]

	// invoke the desired RPC
	err := f.invoke(r.Body, w, service, method)
	if err != nil {
		code := 500

		// handle gRPC specific errors
		if st, ok := status.FromError(err); ok {
			code = httpStatusFromCode(st.Code())
		}

		w.WriteHeader(code)
		w.Write([]byte(err.Error()))
	}
}

// invoke buffers the request payload, invokes the RPC, and
// writes a succesful response to the given io.Writer.
func (f *FallbackServer) invoke(in io.Reader, out io.Writer, service, method string) error {
	res := &[]byte{}
	req, err := ioutil.ReadAll(in)
	if err != nil {
		return err
	}

	err = f.cc.Invoke(context.Background(), buildMethod(service, method), req, res)
	if err != nil {
		return err
	}

	_, err = out.Write(*res)

	return err
}

// dial creates a connection with the gRPC service backend.
func (f *FallbackServer) dial() (*grpc.ClientConn, error) {
	opts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.ForceCodec(rawCodec{})),
	}

	// default to basic CA, use insecure if on localhost
	auth := grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, ""))
	if strings.Contains(f.backend, "localhost") || strings.Contains(f.backend, "127.0.0.1") {
		auth = grpc.WithInsecure()
	}
	opts = append(opts, auth)

	cc, err := grpc.Dial(f.backend, opts...)
	if err != nil {
		log.Printf("Error dialing gRPC backend at %s: %v\n", f.backend, err)
		return nil, err
	}

	return cc, nil
}
<commit_msg>refactor some more<commit_after>package server

import (
	"context"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"

	"github.com/gorilla/mux"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"
)

// FallbackServer is a grpc-fallback HTTP server.
type FallbackServer struct {
	backend string
	server  http.Server
	cc      *grpc.ClientConn
}

// NewServer creates a new grpc-fallback HTTP server on the
// given port that proxies to the given gRPC server backend.
func NewServer(port, backend string) *FallbackServer {
	if !strings.HasPrefix(port, ":") {
		port = ":" + port
	}

	return &FallbackServer{
		backend: backend,
		server: http.Server{
			Addr: port,
		},
	}
}

// Start starts the grpc-fallback HTTP server listening on its port,
// and opens a connection to the gRPC backend.
func (f *FallbackServer) Start() {
	var err error

	// setup connection to gRPC backend
	f.cc, err = f.dial()
	if err != nil {
		log.Fatal("Error dialing gRPC backend server:", err)
	}

	// setup grpc-fallback complient router
	r := mux.NewRouter()
	r.HandleFunc("/$rpc/{service:[.a-zA-Z0-9]+}/{method:[a-zA-Z]+}", f.handler).Headers("Content-Type", "application/x-protobuf")
	f.server.Handler = r

	log.Println("Fallback server listening on port:", f.server.Addr)
	err = f.server.ListenAndServe()
	if err != nil {
		log.Println("Error in fallback server while listening:", err)
	}
}

// StartBackground runs Start() in a goroutine.
func (f *FallbackServer) StartBackground() {
	go f.Start()
}

// Shutdown turns down the grpc-fallback HTTP server.
func (f *FallbackServer) Shutdown() {
	err := f.server.Shutdown(context.Background())
	if err != nil {
		log.Println("Error shutting down fallback server:", err)
	}
}

// handler is a generic HTTP handler that invokes the proper
// RPC given the grpc-fallback HTTP request.
func (f *FallbackServer) handler(w http.ResponseWriter, r *http.Request) {
	log.Println("Incoming grpc-fallback request:", r.RequestURI)
	v := mux.Vars(r)
	service := v["service"]
	method := v["method"]

	// invoke the desired RPC
	err := f.invoke(r.Body, w, service, method)
	if err != nil {
		code := 500

		// handle gRPC specific errors
		if st, ok := status.FromError(err); ok {
			code = httpStatusFromCode(st.Code())
		}

		w.WriteHeader(code)
		w.Write([]byte(err.Error()))
	}
}

// invoke buffers the request payload, invokes the RPC, and
// writes a succesful response to the given io.Writer.
func (f *FallbackServer) invoke(in io.Reader, out io.Writer, service, method string) error {
	res := &[]byte{}
	req, err := ioutil.ReadAll(in)
	if err != nil {
		return err
	}

	err = f.cc.Invoke(context.Background(), buildMethod(service, method), req, res)
	if err != nil {
		return err
	}

	_, err = out.Write(*res)

	return err
}

// dial creates a connection with the gRPC service backend.
func (f *FallbackServer) dial() (*grpc.ClientConn, error) {
	opts := []grpc.DialOption{
		grpc.WithDefaultCallOptions(grpc.ForceCodec(rawCodec{})),
	}

	// default to basic CA, use insecure if on localhost
	auth := grpc.WithTransportCredentials(credentials.NewClientTLSFromCert(nil, ""))
	if strings.Contains(f.backend, "localhost") || strings.Contains(f.backend, "127.0.0.1") {
		auth = grpc.WithInsecure()
	}
	opts = append(opts, auth)

	return grpc.Dial(f.backend, opts...)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/blackjack/syslog"
	"github.com/boltdb/bolt"
	"github.com/satori/go.uuid"
	"golang.org/x/crypto/ssh"
)

var (
	AUTHKEYSFILE string = os.Getenv("HOME") + "/.ssh/authorized_keys"
	CLIENT       string = strings.Split(os.Getenv("SSH_CLIENT"), " ")[0]
	UUID         string
	USERNAME     string
	HOSTNAME     string
)

func Debug(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_DEBUG, format, a...)
}

func Log(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_NOTICE, format, a...)
}

func Warn(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_WARNING, format, a...)
}

func syntaxError() {
	fmt.Println("ERROR")
	os.Exit(1)
}

func dbFileLine(filename string) {
	file, err := os.Stat(filename)
	if err != nil {
		Warn("Error Stat %v: %v", filename, err)
		return
	}

	trimmed := strings.TrimPrefix(filename, "./")

	hash, _ := fingerprint(filename)

	mode := 0600
	if file.Mode()&0111 != 0 {
		mode = 0700
	}

	fmt.Printf("%s\t%o\t*\t%d\t%x\n", trimmed, mode, file.Size(), hash)
}

func dbDirLine(directory string) {
	trimmed := strings.TrimPrefix(directory, "./")
	fmt.Printf("%s/\t%d\t*\n", trimmed, 700)
}

func listDir(dirname string) {
	files, err := ioutil.ReadDir(dirname)
	if err != nil {
		Warn("Error reading directory %v", dirname)
		return
	}

	for _, file := range files {
		if file.Name() == ".git" {
			continue
		}

		fullname := dirname + "/" + file.Name()

		switch mode := file.Mode(); {
		case mode.IsDir():
			dbDirLine(fullname)
			listDir(fullname)
		case mode.IsRegular():
			dbFileLine(fullname)
		}
	}
}

func netskelDB() {
	servername, _ := os.Hostname()
	now := time.Now().Format("Mon, 2 Jan 2006 15:04:05 UTC")

	fmt.Printf("#\n# .netskeldb for %s at %v\n#\n# Generated %v by %v\n#\n", UUID, CLIENT, now, servername)

	// Force-inject the client itself
	dbDirLine("bin")
	dbFileLine("bin/netskel")

	os.Chdir("db")
	listDir(".")

	Log("Sent netskeldb to %s@%s at %s (%s)", USERNAME, HOSTNAME, CLIENT, UUID)
}

func fingerprint(filename string) ([]byte, error) {
	var result []byte

	file, err := os.Open(filename)
	if err != nil {
		return result, err
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return result, err
	}

	return hash.Sum(result), nil
}

func hexdump(filename string) {
	linelength := 30
	count := 0

	file, err := ioutil.ReadFile(filename)
	if err != nil {
		Warn("Error trying to hexdump %v: %v", filename, err)
		os.Exit(1)
	}

	for _, c := range file {
		fmt.Printf("%02x", c)
		count += 1
		if count >= linelength {
			count = 0
			fmt.Printf("\n")
		}
	}
	fmt.Printf("\n")
	Log("Sent %s (%d bytes) to %s@%s at %s (%s)", filename, len(file), USERNAME, HOSTNAME, CLIENT, UUID)
}

func addKey() {
	servername, err := os.Hostname()
	now := time.Now()
	nowFmt := now.Format("Mon Jan _2 15:04:05 2006")
	uuid := uuid.NewV4()
	cuuid := uuid.String()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		Warn("Error generating private key: %v", err)
		os.Exit(1)
	}
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	pemdata := pem.EncodeToMemory(privateKeyPEM)

	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		Warn("Error constructing public key: %v", err)
		os.Exit(1)
	}
	pubdata := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pub)))

	f, err := os.OpenFile(AUTHKEYSFILE, os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		Warn("Error opening %s for writing: %v", AUTHKEYSFILE, err)
		os.Exit(1)
	}

	defer f.Close()

	if _, err = f.WriteString("restrict " + pubdata + " " + HOSTNAME + " " + cuuid + " " + nowFmt + "\n"); err != nil {
		panic(err)
	}

	fmt.Printf("#\n# Netskel private key generated by %v for %v (%v)\n#\n", servername, HOSTNAME, CLIENT)
	fmt.Printf("# CLIENT_UUID %s\n#\n", uuid)
	fmt.Println(string(pemdata))

	secs := strconv.Itoa(int(now.Unix()))

	clientPut(cuuid, "hostname", HOSTNAME)
	clientPut(cuuid, "originalHostname", HOSTNAME)
	clientPut(cuuid, "created", secs)

	Log("Added %d byte public key to %s for %s@%s (%v) uuid %s", len(pubdata), AUTHKEYSFILE, USERNAME, HOSTNAME, CLIENT, uuid)
}

func clientHeartbeat() {
	now := time.Now()
	secs := strconv.Itoa(int(now.Unix()))

	clientPut(UUID, "inet", CLIENT)
	clientPut(UUID, "lastSeen", secs)
	clientPut(UUID, "hostname", HOSTNAME)
	clientPut(UUID, "username", HOSTNAME)

	Debug("Stored heartbeat for %v", UUID)
}

func clientPut(uuid, key, value string) {
	db, err := bolt.Open("clients.db", 0660, &bolt.Options{Timeout: 2 * time.Second})
	if err != nil {
		Warn("Unable to open client database: %v", err)
		return
	}
	defer db.Close()

	berr := db.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(uuid))
		if err != nil {
			return fmt.Errorf("Can't create bucket for %s: %v", uuid, err)
		}

		perr := b.Put([]byte(key), []byte(value))
		return perr
	})

	if berr != nil {
		Warn("clientPut error %v", berr)
	}
}

func parseUNH(nsCommand []string, uuidPosition, usernamePosition, hostnamePosition int) (string, string, string) {
	cuuid := "nouuid"
	username := "user"
	hostname := "unknown"

	if len(nsCommand) > uuidPosition {
		c, err := uuid.FromString(nsCommand[uuidPosition])
		if err != nil {
			Warn("Unable to parse client-supplied UUID %v: %v", nsCommand[uuidPosition], err)
		} else {
			cuuid = c.String()
		}
	}

	if len(nsCommand) > usernamePosition {
		username = nsCommand[usernamePosition]
	}

	if len(nsCommand) > hostnamePosition {
		hostname = nsCommand[hostnamePosition]
	}

	return cuuid, username, hostname
}

func main() {
	syslog.Openlog("netskel-server", syslog.LOG_PID, syslog.LOG_USER)

	if os.Args[0] != "server" {
		syntaxError()
	}

	nsCommand := strings.Split(os.Args[2], " ")
	command := nsCommand[0]

	Debug("Launched from %v with %v", CLIENT, nsCommand)

	switch command {
	case "netskeldb":
		UUID, USERNAME, HOSTNAME = parseUNH(nsCommand, 1, 2, 3)

		clientHeartbeat()
		netskelDB()

	case "md5":
		filename := nsCommand[1]
		hash, _ := fingerprint(filename)
		fmt.Println(hash)

	case "sendfile":
		filename := nsCommand[1]
		UUID, USERNAME, HOSTNAME = parseUNH(nsCommand, 2, 3, 4)

		if filename == "db/bin/netskel" {
			filename = "bin/netskel"
		}

		hexdump(filename)

	case "addkey":
		_, USERNAME, HOSTNAME = parseUNH(nsCommand, 1, 2, 3)
		addKey()

	default:
		syntaxError()
	}

	os.Exit(0)
}
<commit_msg>Store username not hostname here<commit_after>package main

import (
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/blackjack/syslog"
	"github.com/boltdb/bolt"
	"github.com/satori/go.uuid"
	"golang.org/x/crypto/ssh"
)

var (
	AUTHKEYSFILE string = os.Getenv("HOME") + "/.ssh/authorized_keys"
	CLIENT       string = strings.Split(os.Getenv("SSH_CLIENT"), " ")[0]
	UUID         string
	USERNAME     string
	HOSTNAME     string
)

func Debug(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_DEBUG, format, a...)
}

func Log(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_NOTICE, format, a...)
}

func Warn(format string, a ...interface{}) {
	syslog.Syslogf(syslog.LOG_WARNING, format, a...)
}

func syntaxError() {
	fmt.Println("ERROR")
	os.Exit(1)
}

func dbFileLine(filename string) {
	file, err := os.Stat(filename)
	if err != nil {
		Warn("Error Stat %v: %v", filename, err)
		return
	}

	trimmed := strings.TrimPrefix(filename, "./")

	hash, _ := fingerprint(filename)

	mode := 0600
	if file.Mode()&0111 != 0 {
		mode = 0700
	}

	fmt.Printf("%s\t%o\t*\t%d\t%x\n", trimmed, mode, file.Size(), hash)
}

func dbDirLine(directory string) {
	trimmed := strings.TrimPrefix(directory, "./")
	fmt.Printf("%s/\t%d\t*\n", trimmed, 700)
}

func listDir(dirname string) {
	files, err := ioutil.ReadDir(dirname)
	if err != nil {
		Warn("Error reading directory %v", dirname)
		return
	}

	for _, file := range files {
		if file.Name() == ".git" {
			continue
		}

		fullname := dirname + "/" + file.Name()

		switch mode := file.Mode(); {
		case mode.IsDir():
			dbDirLine(fullname)
			listDir(fullname)
		case mode.IsRegular():
			dbFileLine(fullname)
		}
	}
}

func netskelDB() {
	servername, _ := os.Hostname()
	now := time.Now().Format("Mon, 2 Jan 2006 15:04:05 UTC")

	fmt.Printf("#\n# .netskeldb for %s at %v\n#\n# Generated %v by %v\n#\n", UUID, CLIENT, now, servername)

	// Force-inject the client itself
	dbDirLine("bin")
	dbFileLine("bin/netskel")

	os.Chdir("db")
	listDir(".")

	Log("Sent netskeldb to %s@%s at %s (%s)", USERNAME, HOSTNAME, CLIENT, UUID)
}

func fingerprint(filename string) ([]byte, error) {
	var result []byte

	file, err := os.Open(filename)
	if err != nil {
		return result, err
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return result, err
	}

	return hash.Sum(result), nil
}

func hexdump(filename string) {
	linelength := 30
	count := 0

	file, err := ioutil.ReadFile(filename)
	if err != nil {
		Warn("Error trying to hexdump %v: %v", filename, err)
		os.Exit(1)
	}

	for _, c := range file {
		fmt.Printf("%02x", c)
		count += 1
		if count >= linelength {
			count = 0
			fmt.Printf("\n")
		}
	}
	fmt.Printf("\n")
	Log("Sent %s (%d bytes) to %s@%s at %s (%s)", filename, len(file), USERNAME, HOSTNAME, CLIENT, UUID)
}

func addKey() {
	servername, err := os.Hostname()
	now := time.Now()
	nowFmt := now.Format("Mon Jan _2 15:04:05 2006")
	uuid := uuid.NewV4()
	cuuid := uuid.String()

	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		Warn("Error generating private key: %v", err)
		os.Exit(1)
	}
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	pemdata := pem.EncodeToMemory(privateKeyPEM)

	pub, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		Warn("Error constructing public key: %v", err)
		os.Exit(1)
	}
	pubdata := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pub)))

	f, err := os.OpenFile(AUTHKEYSFILE, os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		Warn("Error opening %s for writing: %v", AUTHKEYSFILE, err)
		os.Exit(1)
	}

	defer f.Close()

	if _, err = f.WriteString("restrict " + pubdata + " " + HOSTNAME + " " + cuuid + " " + nowFmt + "\n"); err != nil {
		panic(err)
	}

	fmt.Printf("#\n# Netskel private key generated by %v for %v (%v)\n#\n", servername, HOSTNAME, CLIENT)
	fmt.Printf("# CLIENT_UUID %s\n#\n", uuid)
	fmt.Println(string(pemdata))

	secs := strconv.Itoa(int(now.Unix()))

	clientPut(cuuid, "hostname", HOSTNAME)
	clientPut(cuuid, "originalHostname", HOSTNAME)
	clientPut(cuuid, "created", secs)

	Log("Added %d byte public key to %s for %s@%s (%v) uuid %s", len(pubdata), AUTHKEYSFILE, USERNAME, HOSTNAME, CLIENT, uuid)
}

func clientHeartbeat() {
	now := time.Now()
	secs := strconv.Itoa(int(now.Unix()))

	clientPut(UUID, "inet", CLIENT)
	clientPut(UUID, "lastSeen", secs)
	clientPut(UUID, "hostname", HOSTNAME)
	clientPut(UUID, "username", USERNAME)

	Debug("Stored heartbeat for %v", UUID)
}

func clientPut(uuid, key, value string) {
	db, err := bolt.Open("clients.db", 0660, &bolt.Options{Timeout: 2 * time.Second})
	if err != nil {
		Warn("Unable to open client database: %v", err)
		return
	}
	defer db.Close()

	berr := db.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists([]byte(uuid))
		if err != nil {
			return fmt.Errorf("Can't create bucket for %s: %v", uuid, err)
		}

		perr := b.Put([]byte(key), []byte(value))
		return perr
	})

	if berr != nil {
		Warn("clientPut error %v", berr)
	}
}

func parseUNH(nsCommand []string, uuidPosition, usernamePosition, hostnamePosition int) (string, string, string) {
	cuuid := "nouuid"
	username := "user"
	hostname := "unknown"

	if len(nsCommand) > uuidPosition {
		c, err := uuid.FromString(nsCommand[uuidPosition])
		if err != nil {
			Warn("Unable to parse client-supplied UUID %v: %v", nsCommand[uuidPosition], err)
		} else {
			cuuid = c.String()
		}
	}

	if len(nsCommand) > usernamePosition {
		username = nsCommand[usernamePosition]
	}

	if len(nsCommand) > hostnamePosition {
		hostname = nsCommand[hostnamePosition]
	}

	return cuuid, username, hostname
}

func main() {
	syslog.Openlog("netskel-server", syslog.LOG_PID, syslog.LOG_USER)

	if os.Args[0] != "server" {
		syntaxError()
	}

	nsCommand := strings.Split(os.Args[2], " ")
	command := nsCommand[0]

	Debug("Launched from %v with %v", CLIENT, nsCommand)

	switch command {
	case "netskeldb":
		UUID, USERNAME, HOSTNAME = parseUNH(nsCommand, 1, 2, 3)

		clientHeartbeat()
		netskelDB()

	case "md5":
		filename := nsCommand[1]
		hash, _ := fingerprint(filename)
		fmt.Println(hash)

	case "sendfile":
		filename := nsCommand[1]
		UUID, USERNAME, HOSTNAME = parseUNH(nsCommand, 2, 3, 4)

		if filename == "db/bin/netskel" {
			filename = "bin/netskel"
		}

		hexdump(filename)

	case "addkey":
		_, USERNAME, HOSTNAME = parseUNH(nsCommand, 1, 2, 3)
		addKey()

	default:
		syntaxError()
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>// package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.
package server

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/httpbroadcast"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

const (
	// DefaultDataDir is the default data directory.
	DefaultDataDir = "~/.pilosa"
)

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *pilosa.Config

	// Profiling options.
	CPUProfile string
	CPUTime    time.Duration

	// Standard input/output
	*pilosa.CmdIO

	// running will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}
}

// NewMain returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Server: pilosa.NewServer(),
		Config: pilosa.NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) {
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}
	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", m.Server.Host)
	return nil
}

func (m *Command) SetupServer() error {
	cluster := pilosa.NewCluster()
	cluster.ReplicaN = m.Config.Cluster.ReplicaN

	for _, hostport := range m.Config.Cluster.Hosts {
		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: hostport})
	}
	// TODO: if InternalHosts is not provided then pilosa.Node.InternalHost is empty.
	// This will throw an error when trying to Broadcast messages over HTTP.
	// One option may be to fall back to using host from hostport + config.InternalPort.
	for i, internalhostport := range m.Config.Cluster.InternalHosts {
		cluster.Nodes[i].InternalHost = internalhostport
	}
	m.Server.Cluster = cluster

	// Setup logging output.
	if m.Config.LogPath == "" {
		m.Server.LogOutput = m.Stderr
	} else {
		logFile, err := os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return err
		}
		m.Server.LogOutput = logFile
	}

	// Configure holder.
	fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir)
	m.Server.Holder.Path = m.Config.DataDir
	m.Server.Holder.Stats = pilosa.NewExpvarStatsClient()

	var err error
	m.Server.Host, err = normalizeHost(m.Config.Host)
	if err != nil {
		return err
	}

	// Set internal port (string).
	internalPortStr := pilosa.DefaultInternalPort
	if m.Config.Cluster.InternalPort != "" {
		internalPortStr = m.Config.Cluster.InternalPort
	}

	switch m.Config.Cluster.Type {
	case "http":
		m.Server.Broadcaster = httpbroadcast.NewHTTPBroadcaster(m.Server, internalPortStr)
		m.Server.BroadcastReceiver = httpbroadcast.NewHTTPBroadcastReceiver(internalPortStr, m.Stderr)
		m.Server.Cluster.NodeSet = httpbroadcast.NewHTTPNodeSet()
		err := m.Server.Cluster.NodeSet.(*httpbroadcast.HTTPNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	case "gossip":
		gossipPort, err := strconv.Atoi(internalPortStr)
		if err != nil {
			return err
		}
		gossipSeed := pilosa.DefaultHost
		if m.Config.Cluster.GossipSeed != "" {
			gossipSeed = m.Config.Cluster.GossipSeed
		}
		// get the host portion of addr to use for binding
		gossipHost, _, err := net.SplitHostPort(m.Config.Host)
		if err != nil {
			gossipHost = m.Config.Host
		}
		gossipNodeSet := gossip.NewGossipNodeSet(m.Config.Host, gossipHost, gossipPort, gossipSeed, m.Server)
		m.Server.Cluster.NodeSet = gossipNodeSet
		m.Server.Broadcaster = gossipNodeSet
		m.Server.BroadcastReceiver = gossipNodeSet
	case "static", "":
		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.NodeSet = pilosa.NewStaticNodeSet()
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		err := m.Server.Cluster.NodeSet.(*pilosa.StaticNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("'%v' is not a supported value for broadcaster type.", m.Config.Cluster.Type)
	}

	// Set configuration options.
	m.Server.AntiEntropyInterval = time.Duration(m.Config.AntiEntropy.Interval)
	return nil
}

func normalizeHost(host string) (string, error) {
	if !strings.Contains(host, ":") {
		host = host + ":"
	} else if strings.Contains(host, "://") {
		if strings.HasPrefix(host, "http://") {
			host = host[7:]
		} else {
			return "", fmt.Errorf("invalid scheme or host: '%s'. use the format [http://]<host>:<port>", host)
		}
	}
	return host, nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	logOutput := m.Server.LogOutput
	if closer, ok := logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}
<commit_msg>fixed some comments<commit_after>// Package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.
package server

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/httpbroadcast"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

const (
	// DefaultDataDir is the default data directory.
	DefaultDataDir = "~/.pilosa"
)

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *pilosa.Config

	// Profiling options.
	CPUProfile string
	CPUTime    time.Duration

	// Standard input/output
	*pilosa.CmdIO

	// running will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}
}

// NewCommand returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Server: pilosa.NewServer(),
		Config: pilosa.NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) {
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}
	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", m.Server.Host)
	return nil
}

// SetupServer use the cluster configuration to setup this server
func (m *Command) SetupServer() error {
	cluster := pilosa.NewCluster()
	cluster.ReplicaN = m.Config.Cluster.ReplicaN

	for _, hostport := range m.Config.Cluster.Hosts {
		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: hostport})
	}
	// TODO: if InternalHosts is not provided then pilosa.Node.InternalHost is empty.
	// This will throw an error when trying to Broadcast messages over HTTP.
	// One option may be to fall back to using host from hostport + config.InternalPort.
	for i, internalhostport := range m.Config.Cluster.InternalHosts {
		cluster.Nodes[i].InternalHost = internalhostport
	}
	m.Server.Cluster = cluster

	// Setup logging output.
	if m.Config.LogPath == "" {
		m.Server.LogOutput = m.Stderr
	} else {
		logFile, err := os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return err
		}
		m.Server.LogOutput = logFile
	}

	// Configure holder.
	fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir)
	m.Server.Holder.Path = m.Config.DataDir
	m.Server.Holder.Stats = pilosa.NewExpvarStatsClient()

	var err error
	m.Server.Host, err = normalizeHost(m.Config.Host)
	if err != nil {
		return err
	}

	// Set internal port (string).
	internalPortStr := pilosa.DefaultInternalPort
	if m.Config.Cluster.InternalPort != "" {
		internalPortStr = m.Config.Cluster.InternalPort
	}

	switch m.Config.Cluster.Type {
	case "http":
		m.Server.Broadcaster = httpbroadcast.NewHTTPBroadcaster(m.Server, internalPortStr)
		m.Server.BroadcastReceiver = httpbroadcast.NewHTTPBroadcastReceiver(internalPortStr, m.Stderr)
		m.Server.Cluster.NodeSet = httpbroadcast.NewHTTPNodeSet()
		err := m.Server.Cluster.NodeSet.(*httpbroadcast.HTTPNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	case "gossip":
		gossipPort, err := strconv.Atoi(internalPortStr)
		if err != nil {
			return err
		}
		gossipSeed := pilosa.DefaultHost
		if m.Config.Cluster.GossipSeed != "" {
			gossipSeed = m.Config.Cluster.GossipSeed
		}
		// get the host portion of addr to use for binding
		gossipHost, _, err := net.SplitHostPort(m.Config.Host)
		if err != nil {
			gossipHost = m.Config.Host
		}
		gossipNodeSet := gossip.NewGossipNodeSet(m.Config.Host, gossipHost, gossipPort, gossipSeed, m.Server)
		m.Server.Cluster.NodeSet = gossipNodeSet
		m.Server.Broadcaster = gossipNodeSet
		m.Server.BroadcastReceiver = gossipNodeSet
	case "static", "":
		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.NodeSet = pilosa.NewStaticNodeSet()
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		err := m.Server.Cluster.NodeSet.(*pilosa.StaticNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("'%v' is not a supported value for broadcaster type", m.Config.Cluster.Type)
	}

	// Set configuration options.
	m.Server.AntiEntropyInterval = time.Duration(m.Config.AntiEntropy.Interval)
	return nil
}

func normalizeHost(host string) (string, error) {
	if !strings.Contains(host, ":") {
		host = host + ":"
	} else if strings.Contains(host, "://") {
		if strings.HasPrefix(host, "http://") {
			host = host[7:]
		} else {
			return "", fmt.Errorf("invalid scheme or host: '%s'. use the format [http://]<host>:<port>", host)
		}
	}
	return host, nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	logOutput := m.Server.LogOutput
	if closer, ok := logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"kmud/database"
	"kmud/game"
	"kmud/model"
	"kmud/utils"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"net"
	"strconv"
)

func login(conn net.Conn) *database.User {
	for {
		line := utils.GetUserInput(conn, "Username: ")

		if line == "" {
			return nil
		}

		user := model.M.GetUserByName(line)

		if user == nil {
			utils.WriteLine(conn, "User not found")
		} else if user.Online() {
			utils.WriteLine(conn, "That user is already online")
		} else {
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newUser(conn net.Conn) *database.User {
	for {
		name := utils.GetUserInput(conn, "Desired username: ")

		if name == "" {
			return nil
		}

		user := model.M.GetUserByName(name)

		if user != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			user = model.M.CreateUser(name)
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newCharacter(conn net.Conn, user *database.User) *database.Character {
	// TODO: character slot limit
	const SizeLimit = 12
	for {
		name := utils.GetUserInput(conn, "Desired character name: ")

		if name == "" {
			return nil
		}

		char := model.M.GetCharacterByName(name)

		if char != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			room := model.M.GetRooms()[0] // TODO: Better way to pick an initial character location
			return model.M.CreateCharacter(name, user.Id, room.Id)
		}
	}

	panic("Unexpected code path")
	return nil
}

func quit(conn net.Conn) {
	utils.WriteLine(conn, "Take luck!")
	conn.Close()
}

func mainMenu() utils.Menu {
	menu := utils.NewMenu("MUD")

	menu.AddAction("l", "[L]ogin")
	menu.AddAction("n", "[N]ew user")
	menu.AddAction("q", "[Q]uit")

	return menu
}

func userMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user.Id)

	menu := utils.NewMenu(user.PrettyName())
	menu.AddAction("l", "[L]ogout")
	menu.AddAction("a", "[A]dmin")
	menu.AddAction("n", "[N]ew character")
	if len(chars) > 0 {
		menu.AddAction("d", "[D]elete character")
	}

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.Id)
	}

	return menu
}

func deleteMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user.Id)

	menu := utils.NewMenu("Delete character")

	menu.AddAction("c", "[C]ancel")

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.Id)
	}

	return menu
}

func adminMenu() utils.Menu {
	menu := utils.NewMenu("Admin")
	menu.AddAction("u", "[U]sers")
	return menu
}

func userAdminMenu() utils.Menu {
	menu := utils.NewMenu("User Admin")

	for i, user := range model.M.GetUsers() {
		index := i + 1

		online := ""
		if user.Online() {
			online = "*"
		}

		actionText := fmt.Sprintf("[%v]%v", index, user.PrettyName()+online)
		menu.AddActionData(index, actionText, user.Id)
	}

	return menu
}

func userSpecificMenu(user *database.User) utils.Menu {
	menu := utils.NewMenu("User: " + user.PrettyName())
	menu.AddAction("d", "[D]elete")
	return menu
}

func handleConnection(session *mgo.Session, conn net.Conn) {
	defer conn.Close()
	defer session.Close()

	var user *database.User
	var character *database.Character

	defer func() {
		if r := recover(); r != nil {
			if user != nil {
				user.SetOnline(false)
			}

			fmt.Printf("Lost connection to client (%v/%v): %v, %v\n",
				user.PrettyName(),
				character.PrettyName(),
				conn.RemoteAddr(),
				r)
		}
	}()

	for {
		if user == nil {
			menu := mainMenu()
			choice, _ := menu.Exec(conn, utils.ColorModeNone)

			switch choice {
			case "l":
				user = login(conn)
			case "n":
				user = newUser(conn)
			case "":
				fallthrough
			case "q":
				quit(conn)
				return
			}
		} else if character == nil {
			menu := userMenu(user)
			choice, charId := menu.Exec(conn, user.GetColorMode())

			switch choice {
			case "":
				fallthrough
			case "l":
				user.SetOnline(false)
				user = nil
			case "a":
				adminMenu := adminMenu()
				for {
					choice, _ := adminMenu.Exec(conn, user.GetColorMode())
					if choice == "" {
						break
					} else if choice == "u" {
						for {
							userAdminMenu := userAdminMenu()
							choice, userId := userAdminMenu.Exec(conn, user.GetColorMode())
							if choice == "" {
								break
							} else {
								_, err := strconv.Atoi(choice)

								if err == nil {
									for {
										userMenu := userSpecificMenu(model.M.GetUser(userId))
										choice, _ = userMenu.Exec(conn, user.GetColorMode())
										if choice == "" {
											break
										} else if choice == "d" {
											model.M.DeleteUser(userId)
											break
										}
									}
								}
							}
						}
					}
				}
			case "n":
				character = newCharacter(conn, user)
			case "d":
				for {
					deleteMenu := deleteMenu(user)
					deleteChoice, deleteCharId := deleteMenu.Exec(conn, user.GetColorMode())

					if deleteChoice == "" || deleteChoice == "c" {
						break
					}

					_, err := strconv.Atoi(deleteChoice)

					if err == nil {
						model.M.DeleteCharacter(deleteCharId)
					}
				}

			default:
				_, err := strconv.Atoi(choice)

				if err == nil {
					character = model.M.GetCharacter(charId)
				}
			}
		} else {
			game.Exec(conn, user, character)
			character = nil
		}
	}
}

func Exec() {
	fmt.Printf("Connecting to database... ")
	session, err := mgo.Dial("localhost")

	utils.HandleError(err)

	fmt.Println("done.")

	listener, err := net.Listen("tcp", ":8945")
	utils.HandleError(err)

	err = model.Init(session.Copy())

	// If there are no rooms at all create one
	rooms := model.M.GetRooms()
	if len(rooms) == 0 {
		zones := model.M.GetZones()

		var zoneId bson.ObjectId

		if len(zones) == 0 {
			newZone := model.M.CreateZone("Default")
			zoneId = newZone.Id
		} else {
			zoneId = zones[0].Id
		}

		model.M.CreateRoom(zoneId)
	}

	fmt.Println("Server listening on port 8945")

	for {
		conn, err := listener.Accept()
		utils.HandleError(err)
		fmt.Println("Client connected:", conn.RemoteAddr())
		go handleConnection(session.Copy(), conn)
	}
}

// vim: nocindent
<commit_msg>Fixed some nil-pointer crashes<commit_after>package server

import (
	"fmt"
	"kmud/database"
	"kmud/game"
	"kmud/model"
	"kmud/utils"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"net"
	"strconv"
)

func login(conn net.Conn) *database.User {
	for {
		line := utils.GetUserInput(conn, "Username: ")

		if line == "" {
			return nil
		}

		user := model.M.GetUserByName(line)

		if user == nil {
			utils.WriteLine(conn, "User not found")
		} else if user.Online() {
			utils.WriteLine(conn, "That user is already online")
		} else {
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newUser(conn net.Conn) *database.User {
	for {
		name := utils.GetUserInput(conn, "Desired username: ")

		if name == "" {
			return nil
		}

		user := model.M.GetUserByName(name)

		if user != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			user = model.M.CreateUser(name)
			user.SetOnline(true)
			return user
		}
	}

	panic("Unexpected code path")
	return nil
}

func newCharacter(conn net.Conn, user *database.User) *database.Character {
	// TODO: character slot limit
	const SizeLimit = 12
	for {
		name := utils.GetUserInput(conn, "Desired character name: ")

		if name == "" {
			return nil
		}

		char := model.M.GetCharacterByName(name)

		if char != nil {
			utils.WriteLine(conn, "That name is unavailable")
		} else if err := utils.ValidateName(name); err != nil {
			utils.WriteLine(conn, err.Error())
		} else {
			room := model.M.GetRooms()[0] // TODO: Better way to pick an initial character location
			return model.M.CreateCharacter(name, user.Id, room.Id)
		}
	}

	panic("Unexpected code path")
	return nil
}

func quit(conn net.Conn) {
	utils.WriteLine(conn, "Take luck!")
	conn.Close()
}

func mainMenu() utils.Menu {
	menu := utils.NewMenu("MUD")

	menu.AddAction("l", "[L]ogin")
	menu.AddAction("n", "[N]ew user")
	menu.AddAction("q", "[Q]uit")

	return menu
}

func userMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user.Id)

	menu := utils.NewMenu(user.PrettyName())
	menu.AddAction("l", "[L]ogout")
	menu.AddAction("a", "[A]dmin")
	menu.AddAction("n", "[N]ew character")
	if len(chars) > 0 {
		menu.AddAction("d", "[D]elete character")
	}

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.Id)
	}

	return menu
}

func deleteMenu(user *database.User) utils.Menu {
	chars := model.M.GetUserCharacters(user.Id)

	menu := utils.NewMenu("Delete character")

	menu.AddAction("c", "[C]ancel")

	for i, char := range chars {
		index := i + 1
		actionText := fmt.Sprintf("[%v]%v", index, char.PrettyName())
		menu.AddActionData(index, actionText, char.Id)
	}

	return menu
}

func adminMenu() utils.Menu {
	menu := utils.NewMenu("Admin")
	menu.AddAction("u", "[U]sers")
	return menu
}

func userAdminMenu() utils.Menu {
	menu := utils.NewMenu("User Admin")

	for i, user := range model.M.GetUsers() {
		index := i + 1

		online := ""
		if user.Online() {
			online = "*"
		}

		actionText := fmt.Sprintf("[%v]%v", index, user.PrettyName()+online)
		menu.AddActionData(index, actionText, user.Id)
	}

	return menu
}

func userSpecificMenu(user *database.User) utils.Menu {
	menu := utils.NewMenu("User: " + user.PrettyName())
	menu.AddAction("d", "[D]elete")
	return menu
}

func handleConnection(session *mgo.Session, conn net.Conn) {
	defer conn.Close()
	defer session.Close()

	var user *database.User
	var character *database.Character

	defer func() {
		if r := recover(); r != nil {
			username := ""
			charname := ""

			if user != nil {
				user.SetOnline(false)
				username = user.PrettyName()
			}

			if character != nil {
				charname = character.PrettyName()
			}

			fmt.Printf("Lost connection to client (%v/%v): %v, %v\n",
				username,
				charname,
				conn.RemoteAddr(),
				r)
		}
	}()

	for {
		if user == nil {
			menu := mainMenu()
			choice, _ := menu.Exec(conn, utils.ColorModeNone)

			switch choice {
			case "l":
				user = login(conn)
			case "n":
				user = newUser(conn)
			case "":
				fallthrough
			case "q":
				quit(conn)
				return
			}
		} else if character == nil {
			menu := userMenu(user)
			choice, charId := menu.Exec(conn, user.GetColorMode())

			switch choice {
			case "":
				fallthrough
			case "l":
				user.SetOnline(false)
				user = nil
			case "a":
				adminMenu := adminMenu()
				for {
					choice, _ := adminMenu.Exec(conn, user.GetColorMode())
					if choice == "" {
						break
					} else if choice == "u" {
						for {
							userAdminMenu := userAdminMenu()
							choice, userId := userAdminMenu.Exec(conn, user.GetColorMode())
							if choice == "" {
								break
							} else {
								_, err := strconv.Atoi(choice)

								if err == nil {
									for {
										userMenu := userSpecificMenu(model.M.GetUser(userId))
										choice, _ = userMenu.Exec(conn, user.GetColorMode())
										if choice == "" {
											break
										} else if choice == "d" {
											model.M.DeleteUser(userId)
											break
										}
									}
								}
							}
						}
					}
				}
			case "n":
				character = newCharacter(conn, user)
			case "d":
				for {
					deleteMenu := deleteMenu(user)
					deleteChoice, deleteCharId := deleteMenu.Exec(conn, user.GetColorMode())

					if deleteChoice == "" || deleteChoice == "c" {
						break
					}

					_, err := strconv.Atoi(deleteChoice)

					if err == nil {
						model.M.DeleteCharacter(deleteCharId)
					}
				}

			default:
				_, err := strconv.Atoi(choice)

				if err == nil {
					character = model.M.GetCharacter(charId)
				}
			}
		} else {
			game.Exec(conn, user, character)
			character = nil
		}
	}
}

func Exec() {
	fmt.Printf("Connecting to database... ")
	session, err := mgo.Dial("localhost")

	utils.HandleError(err)

	fmt.Println("done.")

	listener, err := net.Listen("tcp", ":8945")
	utils.HandleError(err)

	err = model.Init(session.Copy())

	// If there are no rooms at all create one
	rooms := model.M.GetRooms()
	if len(rooms) == 0 {
		zones := model.M.GetZones()

		var zoneId bson.ObjectId

		if len(zones) == 0 {
			newZone := model.M.CreateZone("Default")
			zoneId = newZone.Id
		} else {
			zoneId = zones[0].Id
		}

		model.M.CreateRoom(zoneId)
	}

	fmt.Println("Server listening on port 8945")

	for {
		conn, err := listener.Accept()
		utils.HandleError(err)
		fmt.Println("Client connected:", conn.RemoteAddr())
		go handleConnection(session.Copy(), conn)
	}
}

// vim: nocindent
<|endoftext|>


<commit_before>package main

import "container/list"
import "encoding/json"
import "log"
import "net/http"
import "time"

type Message struct {
	Time time.Time
	Text string
}

type FetchResponse struct {
	Messages []Message
	Time     time.Time
}

type StoreRequest struct {
	StartTime time.Time
	Messages  chan<- []Message
}

type Store struct {
	Add chan Message
	Get chan StoreRequest
}

// TODO: Monotonic clock

func manage_store(store Store) {
	messages := list.New()
	message_count := 0
	max_messages := 1000
	waiting := list.New()
	for {
		select {
		case new_message := <-store.Add:
			messages.PushBack(new_message)
			for waiter := waiting.Front(); waiter != nil; waiter = waiter.Next() {
				waiter.Value.(StoreRequest).Messages <- []Message{new_message}
				close(waiter.Value.(StoreRequest).Messages)
			}
			waiting.Init()
			if message_count < max_messages {
				message_count++
			} else {
				messages.Remove(messages.Front())
			}
		case request := <-store.Get:
			if messages.Back() == nil || request.StartTime.After(messages.Back().Value.(Message).Time) {
				waiting.PushBack(request)
			} else {
				start := messages.Back()
				response_size := 1
				if messages.Front().Value.(Message).Time.After(request.StartTime) {
					start = messages.Front()
					response_size = message_count
				} else {
					for start.Prev().Value.(Message).Time.After(request.StartTime) {
						start = start.Prev()
						response_size++
					}
				}
				response_messages := make([]Message, 0, response_size)
				for m := start; m != nil; m = m.Next() {
					response_messages = append(response_messages, m.Value.(Message))
				}
				request.Messages <- response_messages
			}
		}
	}
}

func start_store() Store {
	store := Store{make(chan Message, 20), make(chan StoreRequest, 20)}
	go manage_store(store)
	return store
}

func start_server(store Store) {
	http.HandleFunc("/fetch", func(w http.ResponseWriter, r *http.Request) {
		var since time.Time
		url_since := r.FormValue("since")
		if url_since != "" {
			err := json.Unmarshal([]byte(url_since), &since)
			if err != nil {
				log.Print("fetch: parse since: ", err)
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte("Could not parse since as date"))
				return
			}
		}
		messages_from_store := make(chan []Message, 1)
		store.Get <- StoreRequest{since, messages_from_store}

		json_encoded, err := json.Marshal(FetchResponse{<-messages_from_store, time.Now()})
		if err != nil {
			log.Print("json encode: ", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Add("Content-Type", "application/json")
		w.Write(json_encoded)
	})

	http.HandleFunc("/speak", func(w http.ResponseWriter, r *http.Request) {
		store.Add <- Message{time.Now(), r.FormValue("text")}
	})

	log.Fatal(http.ListenAndServe(":8080", nil))
}

func main() {
	store := start_store()
	start_server(store)
}
<commit_msg>Wait when since == last message time<commit_after>package main

import "container/list"
import "encoding/json"
import "log"
import "net/http"
import "time"

type Message struct {
	Time time.Time
	Text string
}

type FetchResponse struct {
	Messages []Message
	Time     time.Time
}

type StoreRequest struct {
	StartTime time.Time
	Messages  chan<- []Message
}

type Store struct {
	Add chan Message
	Get chan StoreRequest
}

// TODO: Monotonic clock

func manage_store(store Store) {
	messages := list.New()
	message_count := 0
	max_messages := 1000
	waiting := list.New()
	for {
		select {
		case new_message := <-store.Add:
			messages.PushBack(new_message)
			for waiter := waiting.Front(); waiter != nil; waiter = waiter.Next() {
				waiter.Value.(StoreRequest).Messages <- []Message{new_message}
				close(waiter.Value.(StoreRequest).Messages)
			}
			waiting.Init()
			if message_count < max_messages {
				message_count++
			} else {
				messages.Remove(messages.Front())
			}
		case request := <-store.Get:
			if messages.Back() == nil || !request.StartTime.Before(messages.Back().Value.(Message).Time) {
				waiting.PushBack(request)
			} else {
				start := messages.Back()
				response_size := 1
				if messages.Front().Value.(Message).Time.After(request.StartTime) {
					start = messages.Front()
					response_size = message_count
				} else {
					for start.Prev().Value.(Message).Time.After(request.StartTime) {
						start = start.Prev()
						response_size++
					}
				}
				response_messages := make([]Message, 0, response_size)
				for m := start; m != nil; m = m.Next() {
					response_messages = append(response_messages, m.Value.(Message))
				}
				request.Messages <- response_messages
			}
		}
	}
}

func start_store() Store {
	store := Store{make(chan Message, 20), make(chan StoreRequest, 20)}
	go manage_store(store)
	return store
}

func start_server(store Store) {
	http.HandleFunc("/fetch", func(w http.ResponseWriter, r *http.Request) {
		var since time.Time
		url_since := r.FormValue("since")
		if url_since != "" {
			err := json.Unmarshal([]byte(url_since), &since)
			if err != nil {
				log.Print("fetch: parse since: ", err)
				w.WriteHeader(http.StatusBadRequest)
				w.Write([]byte("Could not parse since as date"))
				return
			}
		}
		messages_from_store := make(chan []Message, 1)
		store.Get <- StoreRequest{since, messages_from_store}

		json_encoded, err := json.Marshal(FetchResponse{<-messages_from_store, time.Now()})
		if err != nil {
			log.Print("json encode: ", err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.Header().Add("Content-Type", "application/json")
		w.Write(json_encoded)
	})

	http.HandleFunc("/speak", func(w http.ResponseWriter, r *http.Request) {
		store.Add <- Message{time.Now(), r.FormValue("text")}
	})

	log.Fatal(http.ListenAndServe(":8080", nil))
}

func main() {
	store := start_store()
	start_server(store)
}
<|endoftext|>