<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        int
    Listen      string
    File        string
    Directory   string
}

type Connection struct {
    File        string
    RemoteAddr  string
    ReadOnly    bool
}

var connections = make(map[string][]Connection)

/*
    Add a new connection to the list of connections for a file. Make sure there is only one writable connection per filename
    returns true if the connection was added correctly. false otherwise
 */
func addConnection(filename string, readOnly bool, remoteAddr string) bool {
    currentConnections, ok := connections[filename]
    if ok == false {
        currentConnections = make([]Connection, 4)
    }

    // If this a writable request, check to see if anybody else has a writable connection
    if !readOnly {
        for _, conn := range currentConnections {
            if !conn.ReadOnly {
                fmt.Printf("Error, too many writable connections. %s is already connected to %s\n", remoteAddr, filename)
                return false
            }
        }
    }

    newConnection := Connection{
        File: filename,
        RemoteAddr: remoteAddr,
        ReadOnly: readOnly,
    }

    connections[filename] = append(currentConnections, newConnection)
    return true
}

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 file specification

    var filename bytes.Buffer
    readOnly := false

    var current_directory = globalSettings.Directory
    var err error
    if current_directory == "" {
        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 != 0) {
        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_limit := 2048*1024    // set the buffer to 2mb

    buffer = make([]byte, buffer_limit)
    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])

        // Error out and drop the connection if there is an attempt to read too much
        if length > buffer_limit {
            fmt.Printf("E")

            file.Sync()
            return
        }

        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.IntFlag{
            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
    fmt.Printf("listen (%s) host (%s) port (%d)\n", globalSettings.Listen, globalSettings.Host, globalSettings.Port)
    hostingAddress := globalSettings.Listen
    if len(globalSettings.Listen) == 0 {
        if len(globalSettings.Host) == 0 || globalSettings.Port <= 0 {
            panic("You need to specify a host and port or specify a listen address (host:port)\n")
        }
        var port string
        fmt.Sprint(port, "%d", globalSettings.Port)
        hostingAddress = globalSettings.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, globalSettings)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options, globalSettings)
            break
        }
    }

}
<commit_msg>fixed #10 finally allowing users to specify their own value for the buffer limit<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        int
    Listen      string
    File        string
    Directory   string
    BufferLimit int
}

type Connection struct {
    File        string
    RemoteAddr  string
    ReadOnly    bool
}

var connections = make(map[string][]Connection)

/*
    Add a new connection to the list of connections for a file. Make sure there is only one writable connection per filename
    returns true if the connection was added correctly. false otherwise
 */
func addConnection(filename string, readOnly bool, remoteAddr string) bool {
    currentConnections, ok := connections[filename]
    if ok == false {
        currentConnections = make([]Connection, 4)
    }

    // If this a writable request, check to see if anybody else has a writable connection
    if !readOnly {
        for _, conn := range currentConnections {
            if !conn.ReadOnly {
                fmt.Printf("Error, too many writable connections. %s is already connected to %s\n", remoteAddr, filename)
                return false
            }
        }
    }

    newConnection := Connection{
        File: filename,
        RemoteAddr: remoteAddr,
        ReadOnly: readOnly,
    }

    connections[filename] = append(currentConnections, newConnection)
    return true
}

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

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 file specification

    var filename bytes.Buffer
    readOnly := false

    var current_directory = globalSettings.Directory
    var err error
    if current_directory == "" {
        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 != 0) {
        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_limit := globalSettings.BufferLimit*1024    // set the buffer to 2mb

    buffer = make([]byte, buffer_limit)
    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])

        // Error out and drop the connection if there is an attempt to read too much
        if length > buffer_limit {
            fmt.Printf("E")

            file.Sync()
            return
        }

        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.IntFlag{
            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",
        },
        cli.IntFlag{
            Name: "buffer",
            Value: globalSettings.BufferLimit,
            Usage: "The number of kilobytes in size of the maximum supported read request e.x. '2048'",
            Destination: &globalSettings.Port,
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    fmt.Printf("listen (%s) host (%s) port (%d)\n", globalSettings.Listen, globalSettings.Host, globalSettings.Port)
    hostingAddress := globalSettings.Listen
    if len(globalSettings.Listen) == 0 {
        if len(globalSettings.Host) == 0 || globalSettings.Port <= 0 {
            panic("You need to specify a host and port or specify a listen address (host:port)\n")
        }
        var port string
        fmt.Sprint(port, "%d", globalSettings.Port)
        hostingAddress = globalSettings.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, globalSettings)
            conn.Close()
            break
        case utils.NBD_COMMAND_EXPORT_NAME:
            go export_name(output, conn, payload_size, payload, options, globalSettings)
            break
        }
    }

}
<|endoftext|>


<commit_before><commit_msg>Adding ShortAlias<commit_after><|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"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"crypto/tls"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/boltdb"
	"github.com/pilosa/pilosa/encoding/proto"
	"github.com/pilosa/pilosa/gcnotify"
	"github.com/pilosa/pilosa/gopsutil"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/http"
	"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.Start 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

	Handler pilosa.Handler
	API     *pilosa.API
	ln      net.Listener

	serverOptions []pilosa.ServerOption
}

type CommandOption func(c *Command) error

func OptCommandServerOptions(opts ...pilosa.ServerOption) CommandOption {
	return func(c *Command) error {
		c.serverOptions = append(c.serverOptions, opts...)
		return nil
	}
}

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

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

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

	for _, opt := range opts {
		err := opt(c)
		if err != nil {
			panic(err)
			// TODO: Return error instead of panic?
		}
	}

	return c
}

// Start starts the pilosa server - it returns once the server is running.
func (m *Command) Start() (err error) {
	defer close(m.Started)

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return errors.Wrap(err, "setting up server")
	}

	// SetupNetworking
	err = m.setupNetworking()
	if err != nil {
		return errors.Wrap(err, "setting up networking")
	}
	go func() {
		err := m.Handler.Serve()
		if err != nil {
			m.logger.Printf("Handler serve error: %v", err)
		}
	}()

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return errors.Wrap(err, "opening server")
	}

	m.logger.Printf("Listening as %s\n", m.API.Node().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")
	}

	productName := "Pilosa"
	if pilosa.EnterpriseEnabled {
		productName += " Enterprise"
	}
	m.logger.Printf("%s %s, build time %s\n", productName, pilosa.Version, pilosa.BuildTime)

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

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

	// If port is 0, get auto-allocated port from listener
	if uri.Port == 0 {
		uri.SetPort(uint16(m.ln.Addr().(*net.TCPAddr).Port))
	}

	c := http.GetHTTPClient(TLSConfig)

	// Setup connection to primary store if this is a replica.
	var primaryTranslateStore pilosa.TranslateStore
	if m.Config.Translation.PrimaryURL != "" {
		primaryTranslateStore = http.NewTranslateStore(m.Config.Translation.PrimaryURL)
	}

	// Set Coordinator.
	coordinatorOpt := pilosa.OptServerIsCoordinator(false)
	if m.Config.Cluster.Coordinator || len(m.Config.Gossip.Seeds) == 0 {
		coordinatorOpt = pilosa.OptServerIsCoordinator(true)
	}

	serverOptions := []pilosa.ServerOption{
		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.OptServerSystemInfo(gopsutil.NewSystemInfo()),
		pilosa.OptServerGCNotifier(gcnotify.NewActiveGCNotifier()),
		pilosa.OptServerStatsClient(statsClient),
		pilosa.OptServerURI(uri),
		pilosa.OptServerInternalClient(http.NewInternalClientFromURI(uri, c)),
		pilosa.OptServerPrimaryTranslateStore(primaryTranslateStore),
		pilosa.OptServerClusterDisabled(m.Config.Cluster.Disabled, m.Config.Cluster.Hosts),
		pilosa.OptServerSerializer(proto.Serializer{}),
		coordinatorOpt,
	}

	serverOptions = append(serverOptions, m.serverOptions...)

	m.Server, err = pilosa.NewServer(serverOptions...)

	if err != nil {
		return errors.Wrap(err, "new server")
	}

	m.API, err = pilosa.NewAPI(pilosa.OptAPIServer(m.Server))
	if err != nil {
		return errors.Wrap(err, "new api")
	}

	m.Handler, err = http.NewHandler(
		http.OptHandlerAllowedOrigins(m.Config.Handler.AllowedOrigins),
		http.OptHandlerAPI(m.API),
		http.OptHandlerLogger(m.logger),
		http.OptHandlerListener(m.ln),
	)
	return errors.Wrap(err, "new handler")

}

// setupNetworking sets up internode communication based on the configuration.
func (m *Command) setupNetworking() error {
	if m.Config.Cluster.Disabled {
		return nil
	}

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

	// get the host portion of addr to use for binding
	gossipHost := m.API.Node().URI.Host
	m.gossipTransport, err = gossip.NewTransport(gossipHost, gossipPort, m.logger.Logger())
	if err != nil {
		return errors.Wrap(err, "getting transport")
	}

	gossipMemberSet, err := gossip.NewGossipMemberSet(
		m.Config.Gossip,
		m.API,
		gossip.WithLogger(m.logger.Logger()),
		gossip.WithTransport(m.gossipTransport),
	)
	if err != nil {
		return errors.Wrap(err, "getting memberset")
	}
	return errors.Wrap(gossipMemberSet.Open(), "opening gossip memberset")
}

// GossipTransport allows a caller to return the gossip transport created when
// setting up the GossipMemberSet. This is useful if one needs to determine the
// allocated ephemeral port programmatically. (usually used in tests)
func (m *Command) GossipTransport() *gossip.Transport {
	return m.gossipTransport
}

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

// 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].", name)
	}
}

// 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 server.NewStatsClient<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"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"crypto/tls"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/boltdb"
	"github.com/pilosa/pilosa/encoding/proto"
	"github.com/pilosa/pilosa/gcnotify"
	"github.com/pilosa/pilosa/gopsutil"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/http"
	"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.Start 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

	Handler pilosa.Handler
	API     *pilosa.API
	ln      net.Listener

	serverOptions []pilosa.ServerOption
}

type CommandOption func(c *Command) error

func OptCommandServerOptions(opts ...pilosa.ServerOption) CommandOption {
	return func(c *Command) error {
		c.serverOptions = append(c.serverOptions, opts...)
		return nil
	}
}

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

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

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

	for _, opt := range opts {
		err := opt(c)
		if err != nil {
			panic(err)
			// TODO: Return error instead of panic?
		}
	}

	return c
}

// Start starts the pilosa server - it returns once the server is running.
func (m *Command) Start() (err error) {
	defer close(m.Started)

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return errors.Wrap(err, "setting up server")
	}

	// SetupNetworking
	err = m.setupNetworking()
	if err != nil {
		return errors.Wrap(err, "setting up networking")
	}
	go func() {
		err := m.Handler.Serve()
		if err != nil {
			m.logger.Printf("Handler serve error: %v", err)
		}
	}()

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return errors.Wrap(err, "opening server")
	}

	m.logger.Printf("Listening as %s\n", m.API.Node().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")
	}

	productName := "Pilosa"
	if pilosa.EnterpriseEnabled {
		productName += " Enterprise"
	}
	m.logger.Printf("%s %s, build time %s\n", productName, pilosa.Version, pilosa.BuildTime)

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

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

	// If port is 0, get auto-allocated port from listener
	if uri.Port == 0 {
		uri.SetPort(uint16(m.ln.Addr().(*net.TCPAddr).Port))
	}

	c := http.GetHTTPClient(TLSConfig)

	// Setup connection to primary store if this is a replica.
	var primaryTranslateStore pilosa.TranslateStore
	if m.Config.Translation.PrimaryURL != "" {
		primaryTranslateStore = http.NewTranslateStore(m.Config.Translation.PrimaryURL)
	}

	// Set Coordinator.
	coordinatorOpt := pilosa.OptServerIsCoordinator(false)
	if m.Config.Cluster.Coordinator || len(m.Config.Gossip.Seeds) == 0 {
		coordinatorOpt = pilosa.OptServerIsCoordinator(true)
	}

	serverOptions := []pilosa.ServerOption{
		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.OptServerSystemInfo(gopsutil.NewSystemInfo()),
		pilosa.OptServerGCNotifier(gcnotify.NewActiveGCNotifier()),
		pilosa.OptServerStatsClient(statsClient),
		pilosa.OptServerURI(uri),
		pilosa.OptServerInternalClient(http.NewInternalClientFromURI(uri, c)),
		pilosa.OptServerPrimaryTranslateStore(primaryTranslateStore),
		pilosa.OptServerClusterDisabled(m.Config.Cluster.Disabled, m.Config.Cluster.Hosts),
		pilosa.OptServerSerializer(proto.Serializer{}),
		coordinatorOpt,
	}

	serverOptions = append(serverOptions, m.serverOptions...)

	m.Server, err = pilosa.NewServer(serverOptions...)

	if err != nil {
		return errors.Wrap(err, "new server")
	}

	m.API, err = pilosa.NewAPI(pilosa.OptAPIServer(m.Server))
	if err != nil {
		return errors.Wrap(err, "new api")
	}

	m.Handler, err = http.NewHandler(
		http.OptHandlerAllowedOrigins(m.Config.Handler.AllowedOrigins),
		http.OptHandlerAPI(m.API),
		http.OptHandlerLogger(m.logger),
		http.OptHandlerListener(m.ln),
	)
	return errors.Wrap(err, "new handler")

}

// setupNetworking sets up internode communication based on the configuration.
func (m *Command) setupNetworking() error {
	if m.Config.Cluster.Disabled {
		return nil
	}

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

	// get the host portion of addr to use for binding
	gossipHost := m.API.Node().URI.Host
	m.gossipTransport, err = gossip.NewTransport(gossipHost, gossipPort, m.logger.Logger())
	if err != nil {
		return errors.Wrap(err, "getting transport")
	}

	gossipMemberSet, err := gossip.NewGossipMemberSet(
		m.Config.Gossip,
		m.API,
		gossip.WithLogger(m.logger.Logger()),
		gossip.WithTransport(m.gossipTransport),
	)
	if err != nil {
		return errors.Wrap(err, "getting memberset")
	}
	return errors.Wrap(gossipMemberSet.Open(), "opening gossip memberset")
}

// GossipTransport allows a caller to return the gossip transport created when
// setting up the GossipMemberSet. This is useful if one needs to determine the
// allocated ephemeral port programmatically. (usually used in tests)
func (m *Command) GossipTransport() *gossip.Transport {
	return m.gossipTransport
}

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

// 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].", name)
	}
}

// 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 finalize

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"bytes"

	"github.com/cloudfoundry/libbuildpack"
)

type Staticfile struct {
	RootDir               string `yaml:"root"`
	HostDotFiles          bool   `yaml:"host_dot_files"`
	LocationInclude       string `yaml:"location_include"`
	DirectoryIndex        bool   `yaml:"directory"`
	SSI                   bool   `yaml:"ssi"`
	PushState             bool   `yaml:"pushstate"`
	HSTS                  bool   `yaml:"http_strict_transport_security"`
	HSTSIncludeSubDomains bool   `yaml:"http_strict_transport_security_include_subdomains"`
	HSTSPreload           bool   `yaml:"http_strict_transport_security_preload"`
	ForceHTTPS            bool   `yaml:"force_https"`
	BasicAuth             bool
	StatusCodes           map[string]string `yaml:"status_codes"`
}

type YAML interface {
	Load(string, interface{}) error
}

type Finalizer struct {
	BuildDir string
	DepDir   string
	Log      *libbuildpack.Logger
	Config   Staticfile
	YAML     YAML
}
type StaticfileTemp struct {
	RootDir               string            `yaml:"root,omitempty"`
	HostDotFiles          string            `yaml:"host_dot_files,omitempty"`
	LocationInclude       string            `yaml:"location_include"`
	DirectoryIndex        string            `yaml:"directory"`
	SSI                   string            `yaml:"ssi"`
	PushState             string            `yaml:"pushstate"`
	HSTS                  string            `yaml:"http_strict_transport_security"`
	HSTSIncludeSubDomains string            `yaml:"http_strict_transport_security_include_subdomains"`
	HSTSPreload           string            `yaml:"http_strict_transport_security_preload"`
	ForceHTTPS            string            `yaml:"force_https"`
	StatusCodes           map[string]string `yaml:"status_codes"`
}

var skipCopyFile = map[string]bool{
	"Staticfile":      true,
	"Staticfile.auth": true,
	"manifest.yml":    true,
	".profile":        true,
	".profile.d":      true,
	"stackato.yml":    true,
	".cloudfoundry":   true,
	"nginx":           true,
}

func Run(sf *Finalizer) error {
	var err error

	err = sf.LoadStaticfile()
	if err != nil {
		sf.Log.Error("Unable to load Staticfile: %s", err.Error())
		return err
	}

	appRootDir, err := sf.GetAppRootDir()
	if err != nil {
		sf.Log.Error("Invalid root directory: %s", err.Error())
		return err
	}

	sf.Warnings()

	err = sf.CopyFilesToPublic(appRootDir)
	if err != nil {
		sf.Log.Error("Unable to copy project files: %s", err.Error())
		return err
	}

	err = sf.ConfigureNginx()
	if err != nil {
		sf.Log.Error("Unable to configure nginx: %s", err.Error())
		return err
	}

	err = sf.WriteStartupFiles()
	if err != nil {
		sf.Log.Error("Unable to write startup file: %s", err.Error())
		return err
	}
	return nil
}

func (sf *Finalizer) WriteStartupFiles() error {
	profiledDir := filepath.Join(sf.DepDir, "profile.d")
	err := os.MkdirAll(profiledDir, 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(profiledDir, "staticfile.sh"), []byte(initScript), 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(sf.BuildDir, "start_logging.sh"), []byte(startLoggingScript), 0755)
	if err != nil {
		return err
	}

	bootScript := filepath.Join(sf.BuildDir, "boot.sh")
	return ioutil.WriteFile(bootScript, []byte(startCommand), 0755)
}

func (sf *Finalizer) LoadStaticfile() error {
	var hash StaticfileTemp
	conf := &sf.Config

	err := sf.YAML.Load(filepath.Join(sf.BuildDir, "Staticfile"), &hash)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	isEnabled := func(value string) bool {
		return (value == "enabled" || value == "true")
	}

	if hash.RootDir != "" {
		conf.RootDir = hash.RootDir
	}

	if isEnabled(hash.HostDotFiles) {
		sf.Log.BeginStep("Enabling hosting of dotfiles")
		conf.HostDotFiles = true
	}

	conf.LocationInclude = hash.LocationInclude
	if conf.LocationInclude != "" {
		sf.Log.BeginStep("Enabling location include file %s", conf.LocationInclude)
	}

	if hash.DirectoryIndex != "" {
		sf.Log.BeginStep("Enabling directory index for folders without index.html files")
		conf.DirectoryIndex = true
	}

	if isEnabled(hash.SSI) {
		sf.Log.BeginStep("Enabling SSI")
		conf.SSI = true
	}

	if isEnabled(hash.PushState) {
		sf.Log.BeginStep("Enabling pushstate")
		conf.PushState = true
	}

	if isEnabled(hash.HSTS) {
		sf.Log.BeginStep("Enabling HSTS")
		conf.HSTS = true
	}
	if isEnabled(hash.HSTSIncludeSubDomains) {
		sf.Log.BeginStep("Enabling HSTS includeSubDomains")
		conf.HSTSIncludeSubDomains = true
	}
	if isEnabled(hash.HSTSPreload) {
		sf.Log.BeginStep("Enabling HSTS Preload")
		conf.HSTSPreload = true
	}
	if isEnabled(hash.ForceHTTPS) {
		sf.Log.BeginStep("Enabling HTTPS redirect")
		conf.ForceHTTPS = true
	}
	if len(hash.StatusCodes) > 0 {
		sf.Log.BeginStep("Enabling custom pages for status_codes")
		conf.StatusCodes = sf.getStatusCodes(hash.StatusCodes)
	}

	if !conf.HSTS && (conf.HSTSIncludeSubDomains || conf.HSTSPreload) {
		sf.Log.Warning("http_strict_transport_security is not enabled while http_strict_transport_security_include_subdomains or http_strict_transport_security_preload have been enabled.")
		sf.Log.Protip("http_strict_transport_security_include_subdomains and http_strict_transport_security_preload do nothing without http_strict_transport_security enabled.", "https://docs.cloudfoundry.org/buildpacks/staticfile/index.html#strict-security")
	}

	authFile := filepath.Join(sf.BuildDir, "Staticfile.auth")
	_, err = os.Stat(authFile)
	if err == nil {
		conf.BasicAuth = true
		sf.Log.BeginStep("Enabling basic authentication using Staticfile.auth")
		sf.Log.Protip("Learn about basic authentication", "https://docs.cloudfoundry.org/buildpacks/staticfile/index.html#authentication")
	}

	return nil
}

func (sf *Finalizer) getStatusCodes(codes map[string]string) map[string]string {
	var versions map[string]string
	versions = make(map[string]string)
	for key, value := range codes {
		if strings.Contains(key, "4xx") {
			key = "400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 421 422 423 424 426 428 429 431 451"
		} else if strings.Contains(key, "5xx") {
			key = "500 501 502 503 504 505 506 507 508 510 511"
		}
		versions[key] = value
	}
	return versions
}

func (sf *Finalizer) GetAppRootDir() (string, error) {
	var rootDirRelative string

	if sf.Config.RootDir != "" {
		rootDirRelative = sf.Config.RootDir
	} else {
		rootDirRelative = "."
	}

	rootDirAbs, err := filepath.Abs(filepath.Join(sf.BuildDir, rootDirRelative))
	if err != nil {
		return "", err
	}

	sf.Log.BeginStep("Root folder %s", rootDirAbs)

	dirInfo, err := os.Stat(rootDirAbs)
	if err != nil {
		return "", fmt.Errorf("the application Staticfile specifies a root directory %s that does not exist", rootDirRelative)
	}

	if !dirInfo.IsDir() {
		return "", fmt.Errorf("the application Staticfile specifies a root directory %s that is a plain file, but was expected to be a directory", rootDirRelative)
	}

	return rootDirAbs, nil
}

func (sf *Finalizer) CopyFilesToPublic(appRootDir string) error {
	sf.Log.BeginStep("Copying project files into public")

	publicDir := filepath.Join(sf.BuildDir, "public")

	if publicDir == appRootDir {
		return nil
	}

	tmpDir, err := ioutil.TempDir("", "staticfile-buildpack.approot.")
	if err != nil {
		return err
	}

	files, err := ioutil.ReadDir(appRootDir)
	if err != nil {
		return err
	}

	for _, file := range files {
		if skipCopyFile[file.Name()] {
			continue
		}

		if strings.HasPrefix(file.Name(), ".") && !sf.Config.HostDotFiles {
			continue
		}

		err = os.Rename(filepath.Join(appRootDir, file.Name()), filepath.Join(tmpDir, file.Name()))
		if err != nil {
			return err
		}
	}

	if err := os.RemoveAll(publicDir); err != nil {
		return err
	}

	if err := os.Rename(tmpDir, publicDir); err != nil {
		return err
	}

	return nil
}

func (sf *Finalizer) Warnings() {
	if len(sf.Config.LocationInclude) > 0 && len(sf.Config.RootDir) == 0 {
		sf.Log.Warning("The location_include directive only works in conjunction with root.\nPlease specify root to use location_include")
	}

	if filepath.Clean(sf.Config.RootDir) == "." {
		found, _ := libbuildpack.FileExists(filepath.Join(sf.BuildDir, "nginx", "conf"))
		if found {
			sf.Log.Info("\n\n\n")
			sf.Log.Warning("You have an nginx/conf directory, but have not set *root*, or have set it to '.'.\nIf you are using the nginx/conf directory for nginx configuration, you probably need to also set the *root* directive.")
			sf.Log.Info("\n\n\n")
		}
	}
}

func (sf *Finalizer) ConfigureNginx() error {
	var err error

	sf.Log.BeginStep("Configuring nginx")

	nginxConf, err := sf.generateNginxConf()
	if err != nil {
		sf.Log.Error("Unable to generate nginx.conf: %s", err.Error())
		return err
	}

	confDir := filepath.Join(sf.BuildDir, "nginx", "conf")
	if err := os.MkdirAll(confDir, 0755); err != nil {
		return err
	}

	logsDir := filepath.Join(sf.BuildDir, "nginx", "logs")
	if err := os.MkdirAll(logsDir, 0755); err != nil {
		return err
	}

	confFiles := map[string]string{
		"nginx.conf": nginxConf,
		"mime.types": MimeTypes}

	for file, contents := range confFiles {
		confDest := filepath.Join(confDir, file)
		customConfFile := filepath.Join(sf.BuildDir, "public", file)

		_, err = os.Stat(customConfFile)
		if err == nil {
			err = os.Rename(customConfFile, confDest)
			if file == "nginx.conf" {
				sf.Log.Warning("overriding nginx.conf is deprecated and highly discouraged, as it breaks the functionality of the Staticfile and Staticfile.auth configuration directives. Please use the NGINX buildpack available at: https://github.com/cloudfoundry/nginx-buildpack")
			}
		} else {
			err = ioutil.WriteFile(confDest, []byte(contents), 0644)
		}

		if err != nil {
			return err
		}
	}

	if sf.Config.BasicAuth {
		authFile := filepath.Join(sf.BuildDir, "Staticfile.auth")
		err = libbuildpack.CopyFile(authFile, filepath.Join(confDir, ".htpasswd"))
		if err != nil {
			return err
		}
	}

	return nil
}

func (sf *Finalizer) generateNginxConf() (string, error) {
	buffer := new(bytes.Buffer)

	t := template.Must(template.New("nginx.conf").Parse(nginxConfTemplate))

	err := t.Execute(buffer, sf.Config)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<commit_msg>removes warning from finalize that is no longer applicable<commit_after>package finalize

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"bytes"

	"github.com/cloudfoundry/libbuildpack"
)

type Staticfile struct {
	RootDir               string `yaml:"root"`
	HostDotFiles          bool   `yaml:"host_dot_files"`
	LocationInclude       string `yaml:"location_include"`
	DirectoryIndex        bool   `yaml:"directory"`
	SSI                   bool   `yaml:"ssi"`
	PushState             bool   `yaml:"pushstate"`
	HSTS                  bool   `yaml:"http_strict_transport_security"`
	HSTSIncludeSubDomains bool   `yaml:"http_strict_transport_security_include_subdomains"`
	HSTSPreload           bool   `yaml:"http_strict_transport_security_preload"`
	ForceHTTPS            bool   `yaml:"force_https"`
	BasicAuth             bool
	StatusCodes           map[string]string `yaml:"status_codes"`
}

type YAML interface {
	Load(string, interface{}) error
}

type Finalizer struct {
	BuildDir string
	DepDir   string
	Log      *libbuildpack.Logger
	Config   Staticfile
	YAML     YAML
}
type StaticfileTemp struct {
	RootDir               string            `yaml:"root,omitempty"`
	HostDotFiles          string            `yaml:"host_dot_files,omitempty"`
	LocationInclude       string            `yaml:"location_include"`
	DirectoryIndex        string            `yaml:"directory"`
	SSI                   string            `yaml:"ssi"`
	PushState             string            `yaml:"pushstate"`
	HSTS                  string            `yaml:"http_strict_transport_security"`
	HSTSIncludeSubDomains string            `yaml:"http_strict_transport_security_include_subdomains"`
	HSTSPreload           string            `yaml:"http_strict_transport_security_preload"`
	ForceHTTPS            string            `yaml:"force_https"`
	StatusCodes           map[string]string `yaml:"status_codes"`
}

var skipCopyFile = map[string]bool{
	"Staticfile":      true,
	"Staticfile.auth": true,
	"manifest.yml":    true,
	".profile":        true,
	".profile.d":      true,
	"stackato.yml":    true,
	".cloudfoundry":   true,
	"nginx":           true,
}

func Run(sf *Finalizer) error {
	var err error

	err = sf.LoadStaticfile()
	if err != nil {
		sf.Log.Error("Unable to load Staticfile: %s", err.Error())
		return err
	}

	appRootDir, err := sf.GetAppRootDir()
	if err != nil {
		sf.Log.Error("Invalid root directory: %s", err.Error())
		return err
	}

	sf.Warnings()

	err = sf.CopyFilesToPublic(appRootDir)
	if err != nil {
		sf.Log.Error("Unable to copy project files: %s", err.Error())
		return err
	}

	err = sf.ConfigureNginx()
	if err != nil {
		sf.Log.Error("Unable to configure nginx: %s", err.Error())
		return err
	}

	err = sf.WriteStartupFiles()
	if err != nil {
		sf.Log.Error("Unable to write startup file: %s", err.Error())
		return err
	}
	return nil
}

func (sf *Finalizer) WriteStartupFiles() error {
	profiledDir := filepath.Join(sf.DepDir, "profile.d")
	err := os.MkdirAll(profiledDir, 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(profiledDir, "staticfile.sh"), []byte(initScript), 0755)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(filepath.Join(sf.BuildDir, "start_logging.sh"), []byte(startLoggingScript), 0755)
	if err != nil {
		return err
	}

	bootScript := filepath.Join(sf.BuildDir, "boot.sh")
	return ioutil.WriteFile(bootScript, []byte(startCommand), 0755)
}

func (sf *Finalizer) LoadStaticfile() error {
	var hash StaticfileTemp
	conf := &sf.Config

	err := sf.YAML.Load(filepath.Join(sf.BuildDir, "Staticfile"), &hash)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	isEnabled := func(value string) bool {
		return (value == "enabled" || value == "true")
	}

	if hash.RootDir != "" {
		conf.RootDir = hash.RootDir
	}

	if isEnabled(hash.HostDotFiles) {
		sf.Log.BeginStep("Enabling hosting of dotfiles")
		conf.HostDotFiles = true
	}

	conf.LocationInclude = hash.LocationInclude
	if conf.LocationInclude != "" {
		sf.Log.BeginStep("Enabling location include file %s", conf.LocationInclude)
	}

	if hash.DirectoryIndex != "" {
		sf.Log.BeginStep("Enabling directory index for folders without index.html files")
		conf.DirectoryIndex = true
	}

	if isEnabled(hash.SSI) {
		sf.Log.BeginStep("Enabling SSI")
		conf.SSI = true
	}

	if isEnabled(hash.PushState) {
		sf.Log.BeginStep("Enabling pushstate")
		conf.PushState = true
	}

	if isEnabled(hash.HSTS) {
		sf.Log.BeginStep("Enabling HSTS")
		conf.HSTS = true
	}
	if isEnabled(hash.HSTSIncludeSubDomains) {
		sf.Log.BeginStep("Enabling HSTS includeSubDomains")
		conf.HSTSIncludeSubDomains = true
	}
	if isEnabled(hash.HSTSPreload) {
		sf.Log.BeginStep("Enabling HSTS Preload")
		conf.HSTSPreload = true
	}
	if isEnabled(hash.ForceHTTPS) {
		sf.Log.BeginStep("Enabling HTTPS redirect")
		conf.ForceHTTPS = true
	}
	if len(hash.StatusCodes) > 0 {
		sf.Log.BeginStep("Enabling custom pages for status_codes")
		conf.StatusCodes = sf.getStatusCodes(hash.StatusCodes)
	}

	if !conf.HSTS && (conf.HSTSIncludeSubDomains || conf.HSTSPreload) {
		sf.Log.Warning("http_strict_transport_security is not enabled while http_strict_transport_security_include_subdomains or http_strict_transport_security_preload have been enabled.")
		sf.Log.Protip("http_strict_transport_security_include_subdomains and http_strict_transport_security_preload do nothing without http_strict_transport_security enabled.", "https://docs.cloudfoundry.org/buildpacks/staticfile/index.html#strict-security")
	}

	authFile := filepath.Join(sf.BuildDir, "Staticfile.auth")
	_, err = os.Stat(authFile)
	if err == nil {
		conf.BasicAuth = true
		sf.Log.BeginStep("Enabling basic authentication using Staticfile.auth")
		sf.Log.Protip("Learn about basic authentication", "https://docs.cloudfoundry.org/buildpacks/staticfile/index.html#authentication")
	}

	return nil
}

func (sf *Finalizer) getStatusCodes(codes map[string]string) map[string]string {
	var versions map[string]string
	versions = make(map[string]string)
	for key, value := range codes {
		if strings.Contains(key, "4xx") {
			key = "400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 421 422 423 424 426 428 429 431 451"
		} else if strings.Contains(key, "5xx") {
			key = "500 501 502 503 504 505 506 507 508 510 511"
		}
		versions[key] = value
	}
	return versions
}

func (sf *Finalizer) GetAppRootDir() (string, error) {
	var rootDirRelative string

	if sf.Config.RootDir != "" {
		rootDirRelative = sf.Config.RootDir
	} else {
		rootDirRelative = "."
	}

	rootDirAbs, err := filepath.Abs(filepath.Join(sf.BuildDir, rootDirRelative))
	if err != nil {
		return "", err
	}

	sf.Log.BeginStep("Root folder %s", rootDirAbs)

	dirInfo, err := os.Stat(rootDirAbs)
	if err != nil {
		return "", fmt.Errorf("the application Staticfile specifies a root directory %s that does not exist", rootDirRelative)
	}

	if !dirInfo.IsDir() {
		return "", fmt.Errorf("the application Staticfile specifies a root directory %s that is a plain file, but was expected to be a directory", rootDirRelative)
	}

	return rootDirAbs, nil
}

func (sf *Finalizer) CopyFilesToPublic(appRootDir string) error {
	sf.Log.BeginStep("Copying project files into public")

	publicDir := filepath.Join(sf.BuildDir, "public")

	if publicDir == appRootDir {
		return nil
	}

	tmpDir, err := ioutil.TempDir("", "staticfile-buildpack.approot.")
	if err != nil {
		return err
	}

	files, err := ioutil.ReadDir(appRootDir)
	if err != nil {
		return err
	}

	for _, file := range files {
		if skipCopyFile[file.Name()] {
			continue
		}

		if strings.HasPrefix(file.Name(), ".") && !sf.Config.HostDotFiles {
			continue
		}

		err = os.Rename(filepath.Join(appRootDir, file.Name()), filepath.Join(tmpDir, file.Name()))
		if err != nil {
			return err
		}
	}

	if err := os.RemoveAll(publicDir); err != nil {
		return err
	}

	if err := os.Rename(tmpDir, publicDir); err != nil {
		return err
	}

	return nil
}

func (sf *Finalizer) Warnings() {
	if filepath.Clean(sf.Config.RootDir) == "." {
		found, _ := libbuildpack.FileExists(filepath.Join(sf.BuildDir, "nginx", "conf"))
		if found {
			sf.Log.Info("\n\n\n")
			sf.Log.Warning("You have an nginx/conf directory, but have not set *root*, or have set it to '.'.\nIf you are using the nginx/conf directory for nginx configuration, you probably need to also set the *root* directive.")
			sf.Log.Info("\n\n\n")
		}
	}
}

func (sf *Finalizer) ConfigureNginx() error {
	var err error

	sf.Log.BeginStep("Configuring nginx")

	nginxConf, err := sf.generateNginxConf()
	if err != nil {
		sf.Log.Error("Unable to generate nginx.conf: %s", err.Error())
		return err
	}

	confDir := filepath.Join(sf.BuildDir, "nginx", "conf")
	if err := os.MkdirAll(confDir, 0755); err != nil {
		return err
	}

	logsDir := filepath.Join(sf.BuildDir, "nginx", "logs")
	if err := os.MkdirAll(logsDir, 0755); err != nil {
		return err
	}

	confFiles := map[string]string{
		"nginx.conf": nginxConf,
		"mime.types": MimeTypes}

	for file, contents := range confFiles {
		confDest := filepath.Join(confDir, file)
		customConfFile := filepath.Join(sf.BuildDir, "public", file)

		_, err = os.Stat(customConfFile)
		if err == nil {
			err = os.Rename(customConfFile, confDest)
			if file == "nginx.conf" {
				sf.Log.Warning("overriding nginx.conf is deprecated and highly discouraged, as it breaks the functionality of the Staticfile and Staticfile.auth configuration directives. Please use the NGINX buildpack available at: https://github.com/cloudfoundry/nginx-buildpack")
			}
		} else {
			err = ioutil.WriteFile(confDest, []byte(contents), 0644)
		}

		if err != nil {
			return err
		}
	}

	if sf.Config.BasicAuth {
		authFile := filepath.Join(sf.BuildDir, "Staticfile.auth")
		err = libbuildpack.CopyFile(authFile, filepath.Join(confDir, ".htpasswd"))
		if err != nil {
			return err
		}
	}

	return nil
}

func (sf *Finalizer) generateNginxConf() (string, error) {
	buffer := new(bytes.Buffer)

	t := template.Must(template.New("nginx.conf").Parse(nginxConfTemplate))

	err := t.Execute(buffer, sf.Config)
	if err != nil {
		return "", err
	}
	return buffer.String(), nil
}
<|endoftext|>


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

package server

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"github.com/andreaskoch/allmark/config"
	"github.com/andreaskoch/allmark/path"
	"github.com/andreaskoch/allmark/renderer"
	"github.com/andreaskoch/allmark/repository"
	"github.com/andreaskoch/allmark/util"
	"log"
	"math"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
)

var (
	routes map[string]string

	useTempDir = true
)

const (

	// Dynamic Routes
	ItemHandlerRoute      = "/"
	DebugHandlerRoute     = "/debug/index"
	WebSocketHandlerRoute = "/ws"

	// Static Routes
	ThemeFolderRoute = "/theme/"
)

type Server struct {
	repositoryPath string
	pathProvider   *path.Provider
	config         *config.Config
	renderer       *renderer.Renderer
}

func New(repositoryPath string, config *config.Config, useTempDir bool) *Server {

	return &Server{
		repositoryPath: repositoryPath,
		pathProvider:   path.NewProvider(repositoryPath, useTempDir),
		config:         config,
		renderer:       renderer.New(repositoryPath, config, useTempDir),
	}

}

func (server *Server) Serve() {

	server.renderer.Execute()

	// Initialize the routing table
	server.initializeRoutes()

	// start the websocket hub
	go h.run()

	// register handlers
	http.HandleFunc(ItemHandlerRoute, itemHandler)
	http.HandleFunc(DebugHandlerRoute, indexDebugger)
	http.Handle(WebSocketHandlerRoute, websocket.Handler(webSocketHandler))

	// serve theme files
	if themeFolder := server.config.ThemeFolder(); util.DirectoryExists(themeFolder) {
		http.Handle(ThemeFolderRoute, http.StripPrefix(ThemeFolderRoute, http.FileServer(http.Dir(themeFolder))))
	}

	// start http server: http
	httpBinding := server.getHttpBinding()
	fmt.Printf("Starting http server %q\n", httpBinding)

	if err := http.ListenAndServe(httpBinding, nil); err != nil {
		fmt.Fprintf(os.Stderr, "Server failed with error: %v", err)
	}
}

func (server *Server) getHttpBinding() string {

	// validate the port
	port := server.config.Server.Http.Port
	if port < 1 || port > math.MaxUint16 {
		panic(fmt.Sprintf("%q is an invalid value for a port. Ports can only be in the range of %v to %v,", port, 1, math.MaxUint16))
	}

	return fmt.Sprintf(":%v", port)
}

func (server *Server) initializeRoutes() {

	routes = make(map[string]string)

	go func() {

		for {
			select {
			case item := <-server.renderer.Rendered:
				if item != nil {
					fmt.Printf("Registering item %s\n", item)
					server.registerItem(item)
				}
			}
		}
	}()

}

func (server *Server) registerItem(item *repository.Item) {

	// recurse for child items
	for _, child := range item.Childs {
		server.registerItem(child)
	}

	// get the item route and
	// add it to the routing table
	server.registerRoute(item)

	// get the file routes and
	// add them to the routing table
	for _, file := range item.Files.Items() {
		server.registerRoute(file)
	}
}

func (server *Server) registerRoute(pather path.Pather) {

	if pather == nil {
		log.Printf("Cannot add a route for an uninitialized item %q.\n", pather.Path())
		return
	}

	route := server.pathProvider.GetWebRoute(pather)
	filePath := server.pathProvider.GetFilepath(pather)

	if strings.TrimSpace(route) == "" {
		log.Println("Cannot add an empty route to the routing table.")
		return
	}

	routes[route] = filePath
}

func openUrlInBrowser(url string) {
	var err error

	switch runtime.GOOS {
	case "linux":
		err = exec.Command("xdg-open", url).Start()
	case "windows":
		err = exec.Command(`C:\Windows\System32\rundll32.exe`, "url.dll,FileProtocolHandler", url).Start()
	case "darwin":
		err = exec.Command("open", url).Start()
	default:
		log.Print("unsupported platform")
	}

	if err != nil {
		panic(err)
	}
}
<commit_msg>Reconnected the websocket<commit_after>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package server

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"github.com/andreaskoch/allmark/config"
	"github.com/andreaskoch/allmark/path"
	"github.com/andreaskoch/allmark/renderer"
	"github.com/andreaskoch/allmark/repository"
	"github.com/andreaskoch/allmark/util"
	"log"
	"math"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
)

var (
	routes = make(map[string]string)

	useTempDir = true
)

const (

	// Dynamic Routes
	ItemHandlerRoute      = "/"
	DebugHandlerRoute     = "/debug/index"
	WebSocketHandlerRoute = "/ws"

	// Static Routes
	ThemeFolderRoute = "/theme/"
)

type Server struct {
	repositoryPath string
	pathProvider   *path.Provider
	config         *config.Config
	renderer       *renderer.Renderer
}

func New(repositoryPath string, config *config.Config, useTempDir bool) *Server {

	return &Server{
		repositoryPath: repositoryPath,
		pathProvider:   path.NewProvider(repositoryPath, useTempDir),
		config:         config,
		renderer:       renderer.New(repositoryPath, config, useTempDir),
	}

}

func (server *Server) Serve() {

	// start the renderer
	server.renderer.Execute()

	// start a change listener
	server.listenForChanges()

	// start the websocket hub
	go h.run()

	// register handlers
	http.HandleFunc(ItemHandlerRoute, itemHandler)
	http.HandleFunc(DebugHandlerRoute, indexDebugger)
	http.Handle(WebSocketHandlerRoute, websocket.Handler(webSocketHandler))

	// serve theme files
	if themeFolder := server.config.ThemeFolder(); util.DirectoryExists(themeFolder) {
		http.Handle(ThemeFolderRoute, http.StripPrefix(ThemeFolderRoute, http.FileServer(http.Dir(themeFolder))))
	}

	// start http server: http
	httpBinding := server.getHttpBinding()
	fmt.Printf("Starting http server %q\n", httpBinding)

	if err := http.ListenAndServe(httpBinding, nil); err != nil {
		fmt.Fprintf(os.Stderr, "Server failed with error: %v", err)
	}
}

func (server *Server) getHttpBinding() string {

	// validate the port
	port := server.config.Server.Http.Port
	if port < 1 || port > math.MaxUint16 {
		panic(fmt.Sprintf("%q is an invalid value for a port. Ports can only be in the range of %v to %v,", port, 1, math.MaxUint16))
	}

	return fmt.Sprintf(":%v", port)
}

func (server *Server) listenForChanges() {

	go func() {

		for {
			select {
			case item := <-server.renderer.Rendered:
				if item != nil {
					// register the item
					server.registerItem(item)

					// send update event to connected browsers
					h.broadcast <- UpdateMessage(item.Model)
				}
			}
		}
	}()

}

func (server *Server) registerItem(item *repository.Item) {

	// recurse for child items
	for _, child := range item.Childs {
		server.registerItem(child)
	}

	// get the item route and
	// add it to the routing table
	server.registerRoute(item)

	// get the file routes and
	// add them to the routing table
	for _, file := range item.Files.Items() {
		server.registerRoute(file)
	}
}

func (server *Server) registerRoute(pather path.Pather) {

	if pather == nil {
		log.Printf("Cannot add a route for an uninitialized item %q.\n", pather.Path())
		return
	}

	route := server.pathProvider.GetWebRoute(pather)
	filePath := server.pathProvider.GetFilepath(pather)

	if strings.TrimSpace(route) == "" {
		log.Println("Cannot add an empty route to the routing table.")
		return
	}

	routes[route] = filePath
}

func openUrlInBrowser(url string) {
	var err error

	switch runtime.GOOS {
	case "linux":
		err = exec.Command("xdg-open", url).Start()
	case "windows":
		err = exec.Command(`C:\Windows\System32\rundll32.exe`, "url.dll,FileProtocolHandler", url).Start()
	case "darwin":
		err = exec.Command("open", url).Start()
	default:
		log.Print("unsupported platform")
	}

	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

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

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

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

/*
TODO:
2) Add different input/output formats for the API
3) Add settings to the boards
4) Add settings to the threads
6) quote of the day
*/

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

type boards struct {
    Id int
    Name string
    Descr string
    ImageBoardClusterId string
    MaxThreadCount int //to be checked in insert thread
    MaxActiveThreadCount int //to be checked  in insert thread
    MaxPostsPerThread int // to be checked in insert thread
    AreAttachmentsAllowed bool // to be checked in insert post
    PostLimitsReachedActionId int // to be checked in insert post
}

type threads struct{
    Id int
    Name string
    Descr string
    BoardId 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)  ([]byte, error) {
    if req == nil || res == nil {
	return []byte{}, xerrors.NewSysErr()
    }

    values := req.URL.Query()
    api_key := values[`api_key`][0]
    rows, err := dbh.Query("select b.id, b.name, b.descr from boards b join image_board_clusters ibc on ibc.id = b.image_board_cluster_id where api_key = $1;", api_key)
    if err != nil {
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `002`, true)
    }
    defer rows.Close()

    var curr_boards []boards
    for rows.Next() {
	var board boards
	err = rows.Scan(&board.Id, &board.Name, &board.Descr)
	if err != nil {
	    return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `003`, true)
	}
	curr_boards = append(curr_boards, board)
    }
    bytes, err1 := json.Marshal(api_request{"ok", nil, &curr_boards})
    if err1 != nil {
	return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `004`, true)
    }
    return bytes, nil
}

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

    api_key := values[`api_key`][0]
    rows, err := dbh.Query(`select t.id, t.name from threads t
				join boards b on b.id = t.board_id
				join image_board_clusters ibc on ibc.id = b.image_board_cluster_id
			    where t.is_active = TRUE and t.board_id = $1 and ibc.api_key = $2;`, board_id[0], api_key)
    if err != nil {
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `006`, true)
    }
    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 []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `007`, true)
        }
        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 []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `008`, true)
    }

    return bytes, nil
}

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

    api_key := values[`api_key`][0]
    rows, err := dbh.Query(`select tp.id, tp.body
			    from thread_posts tp join threads t on t.id = tp.thread_id
						 join boards b on b.id = t.board_id
						 join image_board_clusters ibc on ibc.id = b.image_board_cluster_id
			    where tp.thread_id = $1 and ibc.api_key = $2;`, thread_id[0], api_key)
    if err != nil {
	glog.Error(err)
        return []byte{}, xerrors.NewSysErr()
    }
    defer rows.Close()

    var curr_posts []thread_posts
    for rows.Next() {
	glog.Info("new post for thread with id: ", thread_id[0])
        var curr_post thread_posts
        err = rows.Scan(&curr_post.Id, &curr_post.Body)
        if err != nil {
            return []byte{}, xerrors.NewSysErr()
        }
        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 []byte{}, xerrors.NewSysErr()
    }

    return bytes, nil
}

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

    thread_body_post, is_passed := values[`thread_post_body`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No thread_post_body given!`, `Invalid params: No thread_post_body given!`, `001`, true)
    }

    var is_limit_reached bool
    err := dbh.QueryRow("select (select count(*) from thread_posts  where thread_id = $1) > max_posts_per_thread  from threads where id = $1;", thread_id[0]).Scan(&is_limit_reached)
    if err != nil {
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `009`, true)
    }

    if is_limit_reached {
	dbh.QueryRow("UPDATE threads set is_active = false where id = $1",  thread_id[0]).Scan()
	return []byte{}, xerrors.NewUIErr(`Thread post limit reached!`, `Thread post limit reached!`, `010`, true)
    }

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

    _, 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 {
	glog.Error(err)
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `011`, true)
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `012`, true)
    }

    return bytes, nil
}

func addThread(res http.ResponseWriter, req *http.Request) ([]byte,error) {
    if req == nil || res == nil{
        return []byte{}, xerrors.NewSysErr()
    }
    values := req.URL.Query()

    thread_name, is_passed := values[`thread_name`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No thread_name given!`, `Invalid params: No thread_name given!`, `013`, true)
    }

    board_id, is_passed := values[`board_id`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No board_id given!`, `Invalid params: No board_id given!`, `014`, true)
    }

    var is_limit_reached bool
    err := dbh.QueryRow("select (select count(*) from threads  where board_id = $1) > thread_setting_max_thread_count  from boards where id = $1;", board_id[0]).Scan(&is_limit_reached)
    if err != nil {
	glog.Error("COULD NOT SELECT thread_count")
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `015`, true)
    }
    if is_limit_reached {
	return []byte{}, xerrors.NewUIErr(`Thread limit reached!`, `Thread limit reached!`, `016`, true)
    }

    _, err = dbh.Query("INSERT INTO threads(name, board_id, limits_reached_action_id) VALUES($1, $2, 1)", thread_name[0], board_id[0])

    if err != nil {
	glog.Error("INSERT FAILED")
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `017`, true)
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `018`, true)
    }

    return bytes, nil
}

var dbConnString = ``
var dbh *sql.DB

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

    var err error
    dbConnString = os.Getenv("ABC_DB_CONN_STRING") // DB will return error if empty string
    dbh, err = sql.Open("postgres", dbConnString)
    if err != nil {
	glog.Fatal(err)
    }

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

    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 = values[`api_key`]
					if !is_passed {
					    res.Write([]byte(`{"Status":"error","Msg":"Paremeter 'api_key' 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}`))
					    glog.Error("command: ", command[0])
					    return
					}

					res.Header().Set("Access-Control-Allow-Origin", "*")
					bytes, err := commands[command[0]](res, req)

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

					glog.Info(string(bytes))
					res.Write(bytes)
    })

    http.ListenAndServe(`:`+ os.Getenv("ABC_SERVER_ENDPOINT_URL"), nil)
}

<commit_msg>feat: thread posts are seen only if thread is active<commit_after>package main

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

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

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

/*
TODO:
2) Add different input/output formats for the API
6) quote of the day
*/

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

type boards struct {
    Id int
    Name string
    Descr string
    ImageBoardClusterId string
    MaxThreadCount int //to be checked in insert thread
    MaxActiveThreadCount int //to be checked  in insert thread
    MaxPostsPerThread int // to be checked in insert thread
    AreAttachmentsAllowed bool // to be checked in insert post
    PostLimitsReachedActionId int // to be checked in insert post
}

type threads struct{
    Id int
    Name string
    Descr string
    BoardId 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)  ([]byte, error) {
    if req == nil || res == nil {
	return []byte{}, xerrors.NewSysErr()
    }

    values := req.URL.Query()
    api_key := values[`api_key`][0]
    rows, err := dbh.Query("select b.id, b.name, b.descr from boards b join image_board_clusters ibc on ibc.id = b.image_board_cluster_id where api_key = $1;", api_key)
    if err != nil {
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `002`, true)
    }
    defer rows.Close()

    var curr_boards []boards
    for rows.Next() {
	var board boards
	err = rows.Scan(&board.Id, &board.Name, &board.Descr)
	if err != nil {
	    return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `003`, true)
	}
	curr_boards = append(curr_boards, board)
    }
    bytes, err1 := json.Marshal(api_request{"ok", nil, &curr_boards})
    if err1 != nil {
	return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `004`, true)
    }
    return bytes, nil
}

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

    api_key := values[`api_key`][0]
    rows, err := dbh.Query(`select t.id, t.name from threads t
				join boards b on b.id = t.board_id
				join image_board_clusters ibc on ibc.id = b.image_board_cluster_id
			    where t.is_active = TRUE and t.board_id = $1 and ibc.api_key = $2;`, board_id[0], api_key)
    if err != nil {
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `006`, true)
    }
    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 []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `007`, true)
        }
        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 []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `008`, true)
    }

    return bytes, nil
}

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

    api_key := values[`api_key`][0]
    rows, err := dbh.Query(`select tp.id, tp.body
			    from thread_posts tp join threads t on t.id = tp.thread_id
						 join boards b on b.id = t.board_id
						 join image_board_clusters ibc on ibc.id = b.image_board_cluster_id
			    where tp.thread_id = $1 and ibc.api_key = $2 and t.is_active = true;`, thread_id[0], api_key)
    if err != nil {
	glog.Error(err)
        return []byte{}, xerrors.NewSysErr()
    }
    defer rows.Close()

    var curr_posts []thread_posts
    for rows.Next() {
	glog.Info("new post for thread with id: ", thread_id[0])
        var curr_post thread_posts
        err = rows.Scan(&curr_post.Id, &curr_post.Body)
        if err != nil {
            return []byte{}, xerrors.NewSysErr()
        }
        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 []byte{}, xerrors.NewSysErr()
    }

    return bytes, nil
}

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

    thread_body_post, is_passed := values[`thread_post_body`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No thread_post_body given!`, `Invalid params: No thread_post_body given!`, `001`, true)
    }

    var is_limit_reached bool
    err := dbh.QueryRow("select (select count(*) from thread_posts  where thread_id = $1) > max_posts_per_thread  from threads where id = $1;", thread_id[0]).Scan(&is_limit_reached)
    if err != nil {
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `009`, true)
    }

    if is_limit_reached {
	dbh.QueryRow("UPDATE threads set is_active = false where id = $1",  thread_id[0]).Scan()
	return []byte{}, xerrors.NewUIErr(`Thread post limit reached!`, `Thread post limit reached!`, `010`, true)
    }

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

    _, 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 {
	glog.Error(err)
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `011`, true)
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `012`, true)
    }

    return bytes, nil
}

func addThread(res http.ResponseWriter, req *http.Request) ([]byte,error) {
    if req == nil || res == nil{
        return []byte{}, xerrors.NewSysErr()
    }
    values := req.URL.Query()

    thread_name, is_passed := values[`thread_name`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No thread_name given!`, `Invalid params: No thread_name given!`, `013`, true)
    }

    board_id, is_passed := values[`board_id`]
    if !is_passed {
        return []byte{}, xerrors.NewUIErr(`Invalid params: No board_id given!`, `Invalid params: No board_id given!`, `014`, true)
    }

    var is_limit_reached bool
    err := dbh.QueryRow("select (select count(*) from threads  where board_id = $1) > thread_setting_max_thread_count  from boards where id = $1;", board_id[0]).Scan(&is_limit_reached)
    if err != nil {
	glog.Error("COULD NOT SELECT thread_count")
	return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `015`, true)
    }
    if is_limit_reached {
	return []byte{}, xerrors.NewUIErr(`Thread limit reached!`, `Thread limit reached!`, `016`, true)
    }

    _, err = dbh.Query("INSERT INTO threads(name, board_id, limits_reached_action_id, max_posts_per_thread) VALUES($1, $2, 1, 10)", thread_name[0], board_id[0])

    if err != nil {
	glog.Error("INSERT FAILED")
        return []byte{}, xerrors.NewUIErr(err.Error(), err.Error(), `017`, true)
    }

    bytes, err1 := json.Marshal(api_request{"ok", nil, nil})
    if err1 != nil {
        return []byte{}, xerrors.NewUIErr(err1.Error(), err1.Error(), `018`, true)
    }

    return bytes, nil
}

var dbConnString = ``
var dbh *sql.DB

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

    var err error
    dbConnString = os.Getenv("ABC_DB_CONN_STRING") // DB will return error if empty string
    dbh, err = sql.Open("postgres", dbConnString)
    if err != nil {
	glog.Fatal(err)
    }

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

    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 = values[`api_key`]
					if !is_passed {
					    res.Write([]byte(`{"Status":"error","Msg":"Paremeter 'api_key' 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}`))
					    glog.Error("command: ", command[0])
					    return
					}

					res.Header().Set("Access-Control-Allow-Origin", "*")
					bytes, err := commands[command[0]](res, req)

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

					glog.Info(string(bytes))
					res.Write(bytes)
    })

    http.ListenAndServe(`:`+ os.Getenv("ABC_SERVER_ENDPOINT_URL"), nil)
}

<|endoftext|>


<commit_before><commit_msg>fix import<commit_after><|endoftext|>