<commit_before>package server

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/registry/auth"
	"github.com/endophage/gotuf/data"
	"github.com/endophage/gotuf/signed"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"

	"github.com/docker/notary/server/handlers"
	"github.com/docker/notary/utils"
)

func init() {
	data.SetDefaultExpiryTimes(
		map[string]int{
			"timestamp": 14,
		},
	)
}

// Run sets up and starts a TLS server that can be cancelled using the
// given configuration. The context it is passed is the context it should
// use directly for the TLS server, and generate children off for requests
func Run(ctx context.Context, addr, tlsCertFile, tlsKeyFile string, trust signed.CryptoService, authMethod string, authOpts interface{}) error {

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	var lsnr net.Listener
	lsnr, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return err
	}

	if tlsCertFile != "" && tlsKeyFile != "" {
		tlsConfig, err := utils.ConfigureServerTLS(&utils.ServerTLSOpts{
			ServerCertFile: tlsCertFile,
			ServerKeyFile:  tlsKeyFile,
		})
		if err != nil {
			return err
		}
		logrus.Info("Enabling TLS")
		lsnr = tls.NewListener(lsnr, tlsConfig)
	} else if tlsCertFile != "" || tlsKeyFile != "" {
		return fmt.Errorf("Partial TLS configuration found. Either include both a cert and key file in the configuration, or include neither to disable TLS.")
	}

	var ac auth.AccessController
	if authMethod == "token" {
		authOptions, ok := authOpts.(map[string]interface{})
		if !ok {
			return fmt.Errorf("auth.options must be a map[string]interface{}")
		}
		ac, err = auth.GetAccessController(authMethod, authOptions)
		if err != nil {
			return err
		}
	}
	hand := utils.RootHandlerFactory(ac, ctx, trust)

	r := mux.NewRouter()
	r.Methods("GET").Path("/v2/").Handler(hand(handlers.MainHandler))
	r.Methods("POST").Path("/v2/{imageName:.*}/_trust/tuf/").Handler(hand(handlers.AtomicUpdateHandler, "push", "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/{tufRole:(root|targets|snapshot)}.json").Handler(hand(handlers.GetHandler, "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/timestamp.json").Handler(hand(handlers.GetTimestampHandler, "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/timestamp.key").Handler(hand(handlers.GetTimestampKeyHandler, "push", "pull"))
	r.Methods("DELETE").Path("/v2/{imageName:.*}/_trust/tuf/").Handler(hand(handlers.DeleteHandler, "push", "pull"))
	r.Methods("GET", "POST", "PUT", "HEAD", "DELETE").Path("/{other:.*}").Handler(hand(utils.NotFoundHandler))
	svr := http.Server{
		Addr:    addr,
		Handler: r,
	}

	logrus.Info("Starting on ", addr)

	err = svr.Serve(lsnr)

	return err
}
<commit_msg>Add the health handler to the main server<commit_after>package server

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/health"
	"github.com/docker/distribution/registry/auth"
	"github.com/endophage/gotuf/data"
	"github.com/endophage/gotuf/signed"
	"github.com/gorilla/mux"
	"golang.org/x/net/context"

	"github.com/docker/notary/server/handlers"
	"github.com/docker/notary/utils"
)

func init() {
	data.SetDefaultExpiryTimes(
		map[string]int{
			"timestamp": 14,
		},
	)
}

// Run sets up and starts a TLS server that can be cancelled using the
// given configuration. The context it is passed is the context it should
// use directly for the TLS server, and generate children off for requests
func Run(ctx context.Context, addr, tlsCertFile, tlsKeyFile string, trust signed.CryptoService, authMethod string, authOpts interface{}) error {

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	var lsnr net.Listener
	lsnr, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return err
	}

	if tlsCertFile != "" && tlsKeyFile != "" {
		tlsConfig, err := utils.ConfigureServerTLS(&utils.ServerTLSOpts{
			ServerCertFile: tlsCertFile,
			ServerKeyFile:  tlsKeyFile,
		})
		if err != nil {
			return err
		}
		logrus.Info("Enabling TLS")
		lsnr = tls.NewListener(lsnr, tlsConfig)
	} else if tlsCertFile != "" || tlsKeyFile != "" {
		return fmt.Errorf("Partial TLS configuration found. Either include both a cert and key file in the configuration, or include neither to disable TLS.")
	}

	var ac auth.AccessController
	if authMethod == "token" {
		authOptions, ok := authOpts.(map[string]interface{})
		if !ok {
			return fmt.Errorf("auth.options must be a map[string]interface{}")
		}
		ac, err = auth.GetAccessController(authMethod, authOptions)
		if err != nil {
			return err
		}
	}
	hand := utils.RootHandlerFactory(ac, ctx, trust)

	r := mux.NewRouter()
	r.Methods("GET").Path("/v2/").Handler(hand(handlers.MainHandler))
	r.Methods("POST").Path("/v2/{imageName:.*}/_trust/tuf/").Handler(hand(handlers.AtomicUpdateHandler, "push", "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/{tufRole:(root|targets|snapshot)}.json").Handler(hand(handlers.GetHandler, "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/timestamp.json").Handler(hand(handlers.GetTimestampHandler, "pull"))
	r.Methods("GET").Path("/v2/{imageName:.*}/_trust/tuf/timestamp.key").Handler(hand(handlers.GetTimestampKeyHandler, "push", "pull"))
	r.Methods("DELETE").Path("/v2/{imageName:.*}/_trust/tuf/").Handler(hand(handlers.DeleteHandler, "push", "pull"))
	r.Methods("GET").Path("/_notary_server/health").Handler(hand(
		func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
			health.StatusHandler(w, r)
			return nil
		}))
	r.Methods("GET", "POST", "PUT", "HEAD", "DELETE").Path("/{other:.*}").Handler(hand(utils.NotFoundHandler))
	svr := http.Server{
		Addr:    addr,
		Handler: r,
	}

	logrus.Info("Starting on ", addr)

	err = svr.Serve(lsnr)

	return err
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io"
	"log"
	"net"

	"github.com/shadyoak/grpc-counter/service"
	"google.golang.org/grpc"
)

type CounterServer struct {
	clients clientList
	counter counter
	port    int
}

func (s *CounterServer) IncrementCounter(stream service.Counter_IncrementCounterServer) error {
	s.clients.addClient(stream)
	defer s.clients.removeClient(stream)

	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		} else if err != nil {
			log.Printf("receive error: %v", err)
			return err
		}

		count := s.counter.increment(in.Count)
		log.Println("current count:", count)
		val := service.CounterValue{count}

		if err := s.clients.notifyAllClients(stream, val); err != nil {
			log.Println("notify error:", err)
			return err
		}
	}
}

func New(port int) *CounterServer {
	return &CounterServer{
		clients: newClientList(),
		counter: 0,
		port:    port,
	}
}

func (s *CounterServer) Start() {
	port := fmt.Sprintf(":%v", s.port)
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	service.RegisterCounterServer(grpcServer, s)
	log.Println("counter server listening on port:", s.port)
	if err := grpcServer.Serve(lis); err != nil {
		log.Fatalf("unable to start counter server: %v", err)
	}
}
<commit_msg>Trimmed logging<commit_after>package server

import (
	"fmt"
	"io"
	"log"
	"net"

	"github.com/shadyoak/grpc-counter/service"
	"google.golang.org/grpc"
)

type CounterServer struct {
	clients clientList
	counter counter
	port    int
}

func (s *CounterServer) IncrementCounter(stream service.Counter_IncrementCounterServer) error {
	s.clients.addClient(stream)
	defer s.clients.removeClient(stream)

	for {
		in, err := stream.Recv()
		if err == io.EOF {
			return nil
		} else if err != nil {
			log.Printf("receive error: %v", err)
			return err
		}

		count := s.counter.increment(in.Count)
		//log.Println("current count:", count)
		val := service.CounterValue{count}

		if err := s.clients.notifyAllClients(stream, val); err != nil {
			log.Println("notify error:", err)
			return err
		}
	}
}

func New(port int) *CounterServer {
	return &CounterServer{
		clients: newClientList(),
		counter: 0,
		port:    port,
	}
}

func (s *CounterServer) Start() {
	port := fmt.Sprintf(":%v", s.port)
	lis, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	service.RegisterCounterServer(grpcServer, s)
	log.Println("counter server listening on port:", s.port)
	if err := grpcServer.Serve(lis); err != nil {
		log.Fatalf("unable to start counter server: %v", err)
	}
}
<|endoftext|>


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

import (
    "fmt"
    "net"
    "../utils"
    "bufio"
    "encoding/binary"
    "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, true)
    }
    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, false)
    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, false)
    if err != nil {
        return
    }

    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, true)

        //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 int(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, true)

            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, true)

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

            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, true)
    }

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

    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 {

        globalSettings.Host = c.GlobalString("host")
        globalSettings.Port = c.GlobalInt("port")
        globalSettings.Host = c.GlobalString("listen")
        globalSettings.Host = c.GlobalString("file")
        globalSettings.Host = c.GlobalString("directory")
        globalSettings.Port = c.GlobalInt("buffer")

        fmt.Println("%", globalSettings)

        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.BufferLimit,
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    fmt.Printf("Parameter Check: listen (%s) host (%s) port (%s)\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.Printf("the port is: %s\n", globalSettings.Port)

        port := string(globalSettings.Port)
        //fmt.Sprintf(port, "%d", globalSettings.Port)
        fmt.Printf("the port is now: %d\n", port)

        hostingAddress = globalSettings.Host + ":" + port
        fmt.Printf("The hosting address is %s, port is %s\n", hostingAddress, port)
    }

    fmt.Printf("aBlox server online at: %s\n", hostingAddress)
    listener, err := net.Listen("tcp", hostingAddress)

    utils.ErrorCheck(err, true)

    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, false)
        if err != nil {
            continue
        }

        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, false)
        if err != nil {
            continue
        }

        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, false)
            if err != nil {
                continue
            }
        }
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        offset = waiting_for
        waiting_for += int(payload_size)
        _, err = io.ReadFull(conn, data[offset:waiting_for])
        utils.ErrorCheck(err, false)
        if err != nil {
            continue
        }

        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 #37 fixed #38 fixed #39 cleaned up command line arguments, fixed error message, cleaned up cross linked command line arguments, dropped corruption on command line parsing<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"
    "os"
    "bytes"
    "io"
    "io/ioutil"
    "github.com/urfave/cli"
    "path/filepath"
    "strconv"
)

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        string
    Listen      string
    File        string
    Directory   string
    BufferLimit 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",
    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, true)
    }
    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, false)
    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, false)
    if err != nil {
        return
    }

    buffer_limit, _ := strconv.Atoi(globalSettings.BufferLimit)
    buffer_limit *= 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, true)

        //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 int(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, true)

            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, true)

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

            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, true)
    }

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

    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 {
        globalSettings.Host = c.GlobalString("host")
        globalSettings.Port = c.GlobalString("port")
        globalSettings.Listen = c.GlobalString("listen")
        globalSettings.File = c.GlobalString("file")
        globalSettings.Directory = c.GlobalString("directory")
        globalSettings.BufferLimit = c.GlobalString("buffer")

        if globalSettings.Listen == "" && (globalSettings.Host == "" || globalSettings.Port == "") {
            fmt.Println("Please specify either a full 'listen' parameter (e.g. 'localhost:8000', '192.168.1.2:8000) or a host and port\n")
        }
        return nil
    }

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

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    fmt.Printf("Parameter Check: listen (%s) host (%s) port (%s)\n", globalSettings.Listen, globalSettings.Host, globalSettings.Port)
    hostingAddress := globalSettings.Listen
    if len(globalSettings.Listen) == 0 {
        if len(globalSettings.Host) == 0 || len(globalSettings.Port) == 0 {
            panic("You need to specify a host and port or specify a listen address (host:port)\n")
        }
        fmt.Printf("the port is: %s\n", globalSettings.Port)

        port := string(globalSettings.Port)
        //fmt.Sprintf(port, "%d", globalSettings.Port)
        fmt.Printf("the port is now: %d\n", port)

        hostingAddress = globalSettings.Host + ":" + port
        fmt.Printf("The hosting address is %s, port is %s\n", hostingAddress, port)
    }

    fmt.Printf("aBlox server online at: %s\n", hostingAddress)
    listener, err := net.Listen("tcp", hostingAddress)

    utils.ErrorCheck(err, true)

    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, false)
        if err != nil {
            continue
        }

        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, false)
        if err != nil {
            continue
        }

        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, false)
            if err != nil {
                continue
            }
        }
        payload_size := int(binary.BigEndian.Uint32(data[16:]))

        offset = waiting_for
        waiting_for += int(payload_size)
        _, err = io.ReadFull(conn, data[offset:waiting_for])
        utils.ErrorCheck(err, false)
        if err != nil {
            continue
        }

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

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

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

}
<|endoftext|>


<commit_before>package server

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/packetzoom/logslammer/output"
)

type Server struct {
	Config *Config
	Buffer *buffer.Buffer

	mtx     sync.Mutex
	inputs  map[string]input.Input
	outputs map[string]output.Output
}

func signalCatcher() chan os.Signal {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	return c
}

func New(configFile string) (*Server, error) {
	config, err := LoadConfig(configFile)
	if err != nil {
		return nil, err
	}

	return &Server{
		Config:  config,
		Buffer:  buffer.New(),
		inputs:  make(map[string]input.Input),
		outputs: make(map[string]output.Output),
	}, nil
}

func (s *Server) Start() {
	log.Println("Starting server")

	// Start buffer
	log.Println("Starting buffer")
	go s.Buffer.Start()

	s.mtx.Lock()

	// Start inputs
	for name, config := range s.Config.Inputs {
		in, err := input.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := in.Init(config, s.Buffer); err != nil {
			log.Fatalf("Failed to init %s input: %v", name, err)
		}

		go func(name string, in input.Input) {
			log.Printf("Starting input %s", name)
			if err := in.Start(); err != nil {
				log.Fatalf("Error starting input %s: %v", name, err)
			}
		}(name, in)

		s.inputs[name] = in
	}

	// Start ouputs
	for name, config := range s.Config.Outputs {
		out, err := output.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := out.Init(config, s.Buffer); err != nil {
			log.Fatalf("Failed to init %s output: %v", name, err)
			continue
		}

		go func(name string, out output.Output) {
			log.Printf("Starting output %s", name)
			if err := out.Start(); err != nil {
				log.Fatalf("Error starting output %s: %v", name, err)
			}
		}(name, out)

		s.outputs[name] = out
	}

	s.mtx.Unlock()

	// Wait for kill signal
	<-signalCatcher()
	log.Printf("Received quit signal")

	// Stop Server
	s.Stop()
}

func (s *Server) Stop() {
	log.Println("Stopping server")

	s.mtx.Lock()

	// stop inputs
	for name, in := range s.inputs {
		log.Printf("Stopping input %s", name)
		if err := in.Stop(); err != nil {
			log.Printf("Error stopping %s input: %v", name, err)
		}
	}

	// stop ouputs
	for name, out := range s.outputs {
		log.Printf("Stopping output %s", name)
		if err := out.Stop(); err != nil {
			log.Printf("Error stopping %s output: %v", name, err)
		}
	}

	s.mtx.Unlock()

	log.Println("Stopping buffer")
	s.Buffer.Stop()
}
<commit_msg>Fix typo<commit_after>package server

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/packetzoom/logslammer/output"
)

type Server struct {
	Config *Config
	Buffer *buffer.Buffer

	mtx     sync.Mutex
	inputs  map[string]input.Input
	outputs map[string]output.Output
}

func signalCatcher() chan os.Signal {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	return c
}

func New(configFile string) (*Server, error) {
	config, err := LoadConfig(configFile)
	if err != nil {
		return nil, err
	}

	return &Server{
		Config:  config,
		Buffer:  buffer.New(),
		inputs:  make(map[string]input.Input),
		outputs: make(map[string]output.Output),
	}, nil
}

func (s *Server) Start() {
	log.Println("Starting server")

	// Start buffer
	log.Println("Starting buffer")
	go s.Buffer.Start()

	s.mtx.Lock()

	// Start inputs
	for name, config := range s.Config.Inputs {
		in, err := input.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := in.Init(config, s.Buffer); err != nil {
			log.Fatalf("Failed to init %s input: %v", name, err)
		}

		go func(name string, in input.Input) {
			log.Printf("Starting input %s", name)
			if err := in.Start(); err != nil {
				log.Fatalf("Error starting input %s: %v", name, err)
			}
		}(name, in)

		s.inputs[name] = in
	}

	// Start outputs
	for name, config := range s.Config.Outputs {
		out, err := output.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := out.Init(config, s.Buffer); err != nil {
			log.Fatalf("Failed to init %s output: %v", name, err)
			continue
		}

		go func(name string, out output.Output) {
			log.Printf("Starting output %s", name)
			if err := out.Start(); err != nil {
				log.Fatalf("Error starting output %s: %v", name, err)
			}
		}(name, out)

		s.outputs[name] = out
	}

	s.mtx.Unlock()

	// Wait for kill signal
	<-signalCatcher()
	log.Printf("Received quit signal")

	// Stop Server
	s.Stop()
}

func (s *Server) Stop() {
	log.Println("Stopping server")

	s.mtx.Lock()

	// stop inputs
	for name, in := range s.inputs {
		log.Printf("Stopping input %s", name)
		if err := in.Stop(); err != nil {
			log.Printf("Error stopping %s input: %v", name, err)
		}
	}

	// stop ouputs
	for name, out := range s.outputs {
		log.Printf("Stopping output %s", name)
		if err := out.Stop(); err != nil {
			log.Printf("Error stopping %s output: %v", name, err)
		}
	}

	s.mtx.Unlock()

	log.Println("Stopping buffer")
	s.Buffer.Stop()
}
<|endoftext|>


<commit_before>// Copyright 2014 CoreOS, Inc.
//
// 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

import (
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"os"
	"sync"
	"time"

	etcd "github.com/coreos/etcd/client"
	"github.com/coreos/go-systemd/activation"

	"github.com/coreos/fleet/agent"
	"github.com/coreos/fleet/api"
	"github.com/coreos/fleet/config"
	"github.com/coreos/fleet/engine"
	"github.com/coreos/fleet/heart"
	"github.com/coreos/fleet/log"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/pkg/lease"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
	"github.com/coreos/fleet/version"
)

const (
	// machineStateRefreshInterval is the amount of time the server will
	// wait before each attempt to refresh the local machine state
	machineStateRefreshInterval = time.Minute

	shutdownTimeout = time.Minute
)

type Server struct {
	agent          *agent.Agent
	aReconciler    *agent.AgentReconciler
	usPub          *agent.UnitStatePublisher
	usGen          *unit.UnitStateGenerator
	engine         *engine.Engine
	mach           *machine.CoreOSMachine
	hrt            heart.Heart
	mon            *Monitor
	api            *api.Server
	disableEngine  bool
	reconfigServer bool
	restartServer  bool

	engineReconcileInterval time.Duration

	killc chan struct{}  // used to signal monitor to shutdown server
	stopc chan struct{}  // used to terminate all other goroutines
	wg    sync.WaitGroup // used to co-ordinate shutdown
}

func New(cfg config.Config, listeners []net.Listener) (*Server, error) {
	agentTTL, err := time.ParseDuration(cfg.AgentTTL)
	if err != nil {
		return nil, err
	}

	mgr, err := systemd.NewSystemdUnitManager(systemd.DefaultUnitsDirectory)
	if err != nil {
		return nil, err
	}

	mach, err := newMachineFromConfig(cfg, mgr)
	if err != nil {
		return nil, err
	}

	tlsConfig, err := pkg.ReadTLSConfigFiles(cfg.EtcdCAFile, cfg.EtcdCertFile, cfg.EtcdKeyFile)
	if err != nil {
		return nil, err
	}

	eCfg := etcd.Config{
		Transport:               &http.Transport{TLSClientConfig: tlsConfig},
		Endpoints:               cfg.EtcdServers,
		HeaderTimeoutPerRequest: (time.Duration(cfg.EtcdRequestTimeout*1000) * time.Millisecond),
	}
	eClient, err := etcd.New(eCfg)
	if err != nil {
		return nil, err
	}

	kAPI := etcd.NewKeysAPI(eClient)
	reg := registry.NewEtcdRegistry(kAPI, cfg.EtcdKeyPrefix)

	pub := agent.NewUnitStatePublisher(reg, mach, agentTTL)
	gen := unit.NewUnitStateGenerator(mgr)

	a := agent.New(mgr, gen, reg, mach, agentTTL)

	var rStream pkg.EventStream
	if !cfg.DisableWatches {
		rStream = registry.NewEtcdEventStream(kAPI, cfg.EtcdKeyPrefix)
	}
	lManager := lease.NewEtcdLeaseManager(kAPI, cfg.EtcdKeyPrefix)

	ar := agent.NewReconciler(reg, rStream)

	e := engine.New(reg, lManager, rStream, mach)

	if len(listeners) == 0 {
		listeners, err = activation.Listeners(false)
		if err != nil {
			return nil, err
		}
	}

	hrt := heart.New(reg, mach)
	mon := NewMonitor(agentTTL)

	apiServer := api.NewServer(listeners, api.NewServeMux(reg, cfg.TokenLimit))
	apiServer.Serve()

	eIval := time.Duration(cfg.EngineReconcileInterval*1000) * time.Millisecond

	srv := Server{
		agent:       a,
		aReconciler: ar,
		usGen:       gen,
		usPub:       pub,
		engine:      e,
		mach:        mach,
		hrt:         hrt,
		mon:         mon,
		api:         apiServer,
		killc:       make(chan struct{}),
		stopc:       nil,
		engineReconcileInterval: eIval,
		disableEngine:           cfg.DisableEngine,
		reconfigServer:          false,
		restartServer:           false,
	}

	return &srv, nil
}

func newMachineFromConfig(cfg config.Config, mgr unit.UnitManager) (*machine.CoreOSMachine, error) {
	state := machine.MachineState{
		PublicIP: cfg.PublicIP,
		Metadata: cfg.Metadata(),
		Version:  version.Version,
	}

	mach := machine.NewCoreOSMachine(state, mgr)
	mach.Refresh()

	if mach.State().ID == "" {
		return nil, errors.New("unable to determine local machine ID")
	}

	return mach, nil
}

func (s *Server) Run() {
	log.Infof("Establishing etcd connectivity")

	var err error
	for sleep := time.Second; ; sleep = pkg.ExpBackoff(sleep, time.Minute) {
		if s.restartServer {
			_, err = s.hrt.Beat(s.mon.TTL)
			if err == nil {
				log.Infof("hrt.Beat() success")
				break
			}
		} else {
			_, err = s.hrt.Register(s.mon.TTL)
			if err == nil {
				log.Infof("hrt.Register() success")
				break
			}
		}
		log.Errorf("Server register machine failed: %v", err)
		time.Sleep(sleep)
	}

	go s.Supervise()

	log.Infof("Starting server components")
	s.stopc = make(chan struct{})
	s.wg = sync.WaitGroup{}
	beatc := make(chan *unit.UnitStateHeartbeat)

	components := []func(){
		func() { s.api.Available(s.stopc) },
		func() { s.mach.PeriodicRefresh(machineStateRefreshInterval, s.stopc) },
		func() { s.agent.Heartbeat(s.stopc) },
		func() { s.aReconciler.Run(s.agent, s.stopc) },
		func() { s.usGen.Run(beatc, s.stopc) },
		func() { s.usPub.Run(beatc, s.stopc) },
	}
	if s.disableEngine {
		log.Info("Not starting engine; disable-engine is set")
	} else {
		components = append(components, func() { s.engine.Run(s.engineReconcileInterval, s.stopc) })
	}
	for _, f := range components {
		f := f
		s.wg.Add(1)
		go func() {
			f()
			s.wg.Done()
		}()
	}
}

// Supervise monitors the life of the Server and coordinates its shutdown.
// A shutdown occurs when the monitor returns, either because a health check
// fails or a user triggers a shutdown. If the shutdown is due to a health
// check failure, the Server is restarted. Supervise will block shutdown until
// all components have finished shutting down or a timeout occurs; if this
// happens, the Server will not automatically be restarted.
func (s *Server) Supervise() {
	sd, err := s.mon.Monitor(s.hrt, s.killc)
	if sd {
		log.Infof("Server monitor triggered: told to shut down")
	} else {
		log.Errorf("Server monitor triggered: %v", err)
	}
	close(s.stopc)
	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(shutdownTimeout):
		log.Panicf("Timed out waiting for server to shut down. Panicking the server without cleanup.")
	}
	if !sd {
		log.Infof("Restarting server")
		s.SetRestartServer(true)
		s.Run()
		s.SetRestartServer(false)
	}
}

// Kill is used to gracefully terminate the server by triggering the Monitor to shut down
func (s *Server) Kill() {
	if !s.reconfigServer {
		close(s.killc)
	}
}

func (s *Server) Purge() {
	s.aReconciler.Purge(s.agent)
	s.usPub.Purge()
	s.engine.Purge()
	s.hrt.Clear()
}

func (s *Server) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Agent              *agent.Agent
		UnitStatePublisher *agent.UnitStatePublisher
		UnitStateGenerator *unit.UnitStateGenerator
	}{
		Agent:              s.agent,
		UnitStatePublisher: s.usPub,
		UnitStateGenerator: s.usGen,
	})
}

func (s *Server) GetApiServerListeners() []net.Listener {
	return s.api.GetListeners()
}

func (s *Server) SetReconfigServer(isReconfigServer bool) {
	s.reconfigServer = isReconfigServer
}

func (s *Server) SetRestartServer(isRestartServer bool) {
	s.restartServer = isRestartServer
}
<commit_msg>Change Panicf to panic<commit_after>// Copyright 2014 CoreOS, Inc.
//
// 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

import (
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"sync"
	"time"

	etcd "github.com/coreos/etcd/client"
	"github.com/coreos/go-systemd/activation"

	"github.com/coreos/fleet/agent"
	"github.com/coreos/fleet/api"
	"github.com/coreos/fleet/config"
	"github.com/coreos/fleet/engine"
	"github.com/coreos/fleet/heart"
	"github.com/coreos/fleet/log"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/pkg/lease"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
	"github.com/coreos/fleet/version"
)

const (
	// machineStateRefreshInterval is the amount of time the server will
	// wait before each attempt to refresh the local machine state
	machineStateRefreshInterval = time.Minute

	shutdownTimeout = time.Minute
)

type Server struct {
	agent          *agent.Agent
	aReconciler    *agent.AgentReconciler
	usPub          *agent.UnitStatePublisher
	usGen          *unit.UnitStateGenerator
	engine         *engine.Engine
	mach           *machine.CoreOSMachine
	hrt            heart.Heart
	mon            *Monitor
	api            *api.Server
	disableEngine  bool
	reconfigServer bool
	restartServer  bool

	engineReconcileInterval time.Duration

	killc chan struct{}  // used to signal monitor to shutdown server
	stopc chan struct{}  // used to terminate all other goroutines
	wg    sync.WaitGroup // used to co-ordinate shutdown
}

func New(cfg config.Config, listeners []net.Listener) (*Server, error) {
	agentTTL, err := time.ParseDuration(cfg.AgentTTL)
	if err != nil {
		return nil, err
	}

	mgr, err := systemd.NewSystemdUnitManager(systemd.DefaultUnitsDirectory)
	if err != nil {
		return nil, err
	}

	mach, err := newMachineFromConfig(cfg, mgr)
	if err != nil {
		return nil, err
	}

	tlsConfig, err := pkg.ReadTLSConfigFiles(cfg.EtcdCAFile, cfg.EtcdCertFile, cfg.EtcdKeyFile)
	if err != nil {
		return nil, err
	}

	eCfg := etcd.Config{
		Transport:               &http.Transport{TLSClientConfig: tlsConfig},
		Endpoints:               cfg.EtcdServers,
		HeaderTimeoutPerRequest: (time.Duration(cfg.EtcdRequestTimeout*1000) * time.Millisecond),
	}
	eClient, err := etcd.New(eCfg)
	if err != nil {
		return nil, err
	}

	kAPI := etcd.NewKeysAPI(eClient)
	reg := registry.NewEtcdRegistry(kAPI, cfg.EtcdKeyPrefix)

	pub := agent.NewUnitStatePublisher(reg, mach, agentTTL)
	gen := unit.NewUnitStateGenerator(mgr)

	a := agent.New(mgr, gen, reg, mach, agentTTL)

	var rStream pkg.EventStream
	if !cfg.DisableWatches {
		rStream = registry.NewEtcdEventStream(kAPI, cfg.EtcdKeyPrefix)
	}
	lManager := lease.NewEtcdLeaseManager(kAPI, cfg.EtcdKeyPrefix)

	ar := agent.NewReconciler(reg, rStream)

	e := engine.New(reg, lManager, rStream, mach)

	if len(listeners) == 0 {
		listeners, err = activation.Listeners(false)
		if err != nil {
			return nil, err
		}
	}

	hrt := heart.New(reg, mach)
	mon := NewMonitor(agentTTL)

	apiServer := api.NewServer(listeners, api.NewServeMux(reg, cfg.TokenLimit))
	apiServer.Serve()

	eIval := time.Duration(cfg.EngineReconcileInterval*1000) * time.Millisecond

	srv := Server{
		agent:       a,
		aReconciler: ar,
		usGen:       gen,
		usPub:       pub,
		engine:      e,
		mach:        mach,
		hrt:         hrt,
		mon:         mon,
		api:         apiServer,
		killc:       make(chan struct{}),
		stopc:       nil,
		engineReconcileInterval: eIval,
		disableEngine:           cfg.DisableEngine,
		reconfigServer:          false,
		restartServer:           false,
	}

	return &srv, nil
}

func newMachineFromConfig(cfg config.Config, mgr unit.UnitManager) (*machine.CoreOSMachine, error) {
	state := machine.MachineState{
		PublicIP: cfg.PublicIP,
		Metadata: cfg.Metadata(),
		Version:  version.Version,
	}

	mach := machine.NewCoreOSMachine(state, mgr)
	mach.Refresh()

	if mach.State().ID == "" {
		return nil, errors.New("unable to determine local machine ID")
	}

	return mach, nil
}

func (s *Server) Run() {
	log.Infof("Establishing etcd connectivity")

	var err error
	for sleep := time.Second; ; sleep = pkg.ExpBackoff(sleep, time.Minute) {
		if s.restartServer {
			_, err = s.hrt.Beat(s.mon.TTL)
			if err == nil {
				log.Infof("hrt.Beat() success")
				break
			}
		} else {
			_, err = s.hrt.Register(s.mon.TTL)
			if err == nil {
				log.Infof("hrt.Register() success")
				break
			}
		}
		log.Errorf("Server register machine failed: %v", err)
		time.Sleep(sleep)
	}

	go s.Supervise()

	log.Infof("Starting server components")
	s.stopc = make(chan struct{})
	s.wg = sync.WaitGroup{}
	beatc := make(chan *unit.UnitStateHeartbeat)

	components := []func(){
		func() { s.api.Available(s.stopc) },
		func() { s.mach.PeriodicRefresh(machineStateRefreshInterval, s.stopc) },
		func() { s.agent.Heartbeat(s.stopc) },
		func() { s.aReconciler.Run(s.agent, s.stopc) },
		func() { s.usGen.Run(beatc, s.stopc) },
		func() { s.usPub.Run(beatc, s.stopc) },
	}
	if s.disableEngine {
		log.Info("Not starting engine; disable-engine is set")
	} else {
		components = append(components, func() { s.engine.Run(s.engineReconcileInterval, s.stopc) })
	}
	for _, f := range components {
		f := f
		s.wg.Add(1)
		go func() {
			f()
			s.wg.Done()
		}()
	}
}

// Supervise monitors the life of the Server and coordinates its shutdown.
// A shutdown occurs when the monitor returns, either because a health check
// fails or a user triggers a shutdown. If the shutdown is due to a health
// check failure, the Server is restarted. Supervise will block shutdown until
// all components have finished shutting down or a timeout occurs; if this
// happens, the Server will not automatically be restarted.
func (s *Server) Supervise() {
	sd, err := s.mon.Monitor(s.hrt, s.killc)
	if sd {
		log.Infof("Server monitor triggered: told to shut down")
	} else {
		log.Errorf("Server monitor triggered: %v", err)
	}
	close(s.stopc)
	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(shutdownTimeout):
		log.Errorf("Timed out waiting for server to shut down. Panicking the server without cleanup.")
		panic("Failed server shutdown. Panic")
	}
	if !sd {
		log.Infof("Restarting server")
		s.SetRestartServer(true)
		s.Run()
		s.SetRestartServer(false)
	}
}

// Kill is used to gracefully terminate the server by triggering the Monitor to shut down
func (s *Server) Kill() {
	if !s.reconfigServer {
		close(s.killc)
	}
}

func (s *Server) Purge() {
	s.aReconciler.Purge(s.agent)
	s.usPub.Purge()
	s.engine.Purge()
	s.hrt.Clear()
}

func (s *Server) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Agent              *agent.Agent
		UnitStatePublisher *agent.UnitStatePublisher
		UnitStateGenerator *unit.UnitStateGenerator
	}{
		Agent:              s.agent,
		UnitStatePublisher: s.usPub,
		UnitStateGenerator: s.usGen,
	})
}

func (s *Server) GetApiServerListeners() []net.Listener {
	return s.api.GetListeners()
}

func (s *Server) SetReconfigServer(isReconfigServer bool) {
	s.reconfigServer = isReconfigServer
}

func (s *Server) SetRestartServer(isRestartServer bool) {
	s.restartServer = isRestartServer
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"
	// For pprof
	_ "net/http/pprof"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util"
	"github.com/pingcap/tidb/util/arena"
)

var (
	baseConnID uint32
)

var (
	errUnknownFieldType  = terror.ClassServer.New(codeUnknownFieldType, "unknown field type")
	errInvalidPayloadLen = terror.ClassServer.New(codeInvalidPayloadLen, "invalid payload length")
	errInvalidSequence   = terror.ClassServer.New(codeInvalidSequence, "invalid sequence")
	errInvalidType       = terror.ClassServer.New(codeInvalidType, "invalid type")
	errNotAllowedCommand = terror.ClassServer.New(codeNotAllowedCommand, "the used command is not allowed with this TiDB version")
	errAccessDenied      = terror.ClassServer.New(codeAccessDenied, mysql.MySQLErrName[mysql.ErrAccessDenied])
)

// Server is the MySQL protocol server
type Server struct {
	cfg               *config.Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn

	// When a critical error occurred, we don't want to exit the process, because there may be
	// a supervisor automatically restart it, then new client connection will be created, but we can't server it.
	// So we just stop the listener and store to force clients to chose other TiDB servers.
	stopListenerCh chan struct{}
}

// ConnectionCount gets current connection count.
func (s *Server) ConnectionCount() int {
	var cnt int
	s.rwlock.RLock()
	cnt = len(s.clients)
	s.rwlock.RUnlock()
	return cnt
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// newConn creates a new *clientConn from a net.Conn.
// It allocates a connection ID and random salt data for authentication.
func (s *Server) newConn(conn net.Conn) *clientConn {
	cc := &clientConn{
		conn:         conn,
		pkt:          newPacketIO(conn),
		server:       s,
		connectionID: atomic.AddUint32(&baseConnID, 1),
		collation:    mysql.DefaultCollationID,
		alloc:        arena.NewAllocator(32 * 1024),
	}
	log.Infof("[%d] new connection %s", cc.connectionID, conn.RemoteAddr().String())
	if s.cfg.TCPKeepAlive {
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			if err := tcpConn.SetKeepAlive(true); err != nil {
				log.Error("failed to set tcp keep alive option:", err)
			}
		}
	}
	cc.salt = util.RandomBuf(20)
	return cc
}

func (s *Server) skipAuth() bool {
	return s.cfg.SkipAuth
}

const tokenLimit = 1000

// NewServer creates a new Server.
func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(tokenLimit),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
		stopListenerCh:    make(chan struct{}, 1),
	}

	var err error
	if cfg.Socket != "" {
		cfg.SkipAuth = true
		s.listener, err = net.Listen("unix", cfg.Socket)
	} else {
		s.listener, err = net.Listen("tcp", s.cfg.Addr)
	}
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Init rand seed for randomBuf()
	rand.Seed(time.Now().UTC().UnixNano())
	log.Infof("Server run MySQL Protocol Listen at [%s]", s.cfg.Addr)
	return s, nil
}

// Run runs the server.
func (s *Server) Run() error {
	// Start HTTP API to report tidb info such as TPS.
	if s.cfg.ReportStatus {
		s.startStatusHTTP()
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Err.Error() == "use of closed network connection" {
					return nil
				}
			}
			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}
		if s.shouldStopListener() {
			conn.Close()
			break
		}
		go s.onConn(conn)
	}
	s.listener.Close()
	s.listener = nil
	for {
		log.Errorf("listener stopped, waiting for manual kill.")
		time.Sleep(time.Minute)
	}
}

func (s *Server) shouldStopListener() bool {
	select {
	case <-s.stopListenerCh:
		return true
	default:
		return false
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		s.listener.Close()
		s.listener = nil
	}
}

// onConn runs in its own goroutine, handles queries from this connection.
func (s *Server) onConn(c net.Conn) {
	conn := s.newConn(c)
	defer func() {
		log.Infof("[%d] close connection", conn.connectionID)
	}()

	if err := conn.handshake(); err != nil {
		// Some keep alive services will send request to TiDB and disconnect immediately.
		// So we use info log level.
		log.Infof("handshake error %s", errors.ErrorStack(err))
		c.Close()
		return
	}

	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	connections := len(s.clients)
	s.rwlock.Unlock()
	connGauge.Set(float64(connections))

	conn.Run()
}

// ShowProcessList implements the SessionManager interface.
func (s *Server) ShowProcessList() []util.ProcessInfo {
	var rs []util.ProcessInfo
	s.rwlock.RLock()
	for _, client := range s.clients {
		if client.killed {
			continue
		}
		rs = append(rs, client.ctx.ShowProcess())
	}
	s.rwlock.RUnlock()
	return rs
}

// Kill implements the SessionManager interface.
func (s *Server) Kill(connectionID uint64, query bool) {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	conn, ok := s.clients[uint32(connectionID)]
	if !ok {
		return
	}

	conn.ctx.Cancel()
	if !query {
		conn.killed = true
	}
}

// Server error codes.
const (
	codeUnknownFieldType  = 1
	codeInvalidPayloadLen = 2
	codeInvalidSequence   = 3
	codeInvalidType       = 4

	codeNotAllowedCommand = 1148
	codeAccessDenied      = mysql.ErrAccessDenied
)

func init() {
	serverMySQLErrCodes := map[terror.ErrCode]uint16{
		codeNotAllowedCommand: mysql.ErrNotAllowedCommand,
		codeAccessDenied:      mysql.ErrAccessDenied,
	}
	terror.ErrClassToMySQLCodes[terror.ClassServer] = serverMySQLErrCodes
}
<commit_msg>server: minor log info corrections (#4319)<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"
	// For pprof
	_ "net/http/pprof"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util"
	"github.com/pingcap/tidb/util/arena"
)

var (
	baseConnID uint32
)

var (
	errUnknownFieldType  = terror.ClassServer.New(codeUnknownFieldType, "unknown field type")
	errInvalidPayloadLen = terror.ClassServer.New(codeInvalidPayloadLen, "invalid payload length")
	errInvalidSequence   = terror.ClassServer.New(codeInvalidSequence, "invalid sequence")
	errInvalidType       = terror.ClassServer.New(codeInvalidType, "invalid type")
	errNotAllowedCommand = terror.ClassServer.New(codeNotAllowedCommand, "the used command is not allowed with this TiDB version")
	errAccessDenied      = terror.ClassServer.New(codeAccessDenied, mysql.MySQLErrName[mysql.ErrAccessDenied])
)

// Server is the MySQL protocol server
type Server struct {
	cfg               *config.Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn

	// When a critical error occurred, we don't want to exit the process, because there may be
	// a supervisor automatically restart it, then new client connection will be created, but we can't server it.
	// So we just stop the listener and store to force clients to chose other TiDB servers.
	stopListenerCh chan struct{}
}

// ConnectionCount gets current connection count.
func (s *Server) ConnectionCount() int {
	var cnt int
	s.rwlock.RLock()
	cnt = len(s.clients)
	s.rwlock.RUnlock()
	return cnt
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// newConn creates a new *clientConn from a net.Conn.
// It allocates a connection ID and random salt data for authentication.
func (s *Server) newConn(conn net.Conn) *clientConn {
	cc := &clientConn{
		conn:         conn,
		pkt:          newPacketIO(conn),
		server:       s,
		connectionID: atomic.AddUint32(&baseConnID, 1),
		collation:    mysql.DefaultCollationID,
		alloc:        arena.NewAllocator(32 * 1024),
	}
	log.Infof("[%d] new connection %s", cc.connectionID, conn.RemoteAddr().String())
	if s.cfg.TCPKeepAlive {
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			if err := tcpConn.SetKeepAlive(true); err != nil {
				log.Error("failed to set tcp keep alive option:", err)
			}
		}
	}
	cc.salt = util.RandomBuf(20)
	return cc
}

func (s *Server) skipAuth() bool {
	return s.cfg.SkipAuth
}

const tokenLimit = 1000

// NewServer creates a new Server.
func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(tokenLimit),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
		stopListenerCh:    make(chan struct{}, 1),
	}

	var err error
	if cfg.Socket != "" {
		cfg.SkipAuth = true
		if s.listener, err = net.Listen("unix", cfg.Socket); err == nil {
			log.Infof("Server is running MySQL Protocol through Socket [%s]", cfg.Socket)
		}
	} else {
		if s.listener, err = net.Listen("tcp", s.cfg.Addr); err == nil {
			log.Infof("Server is running MySQL Protocol at [%s]", s.cfg.Addr)
		}
	}
	if err != nil {
		return nil, errors.Trace(err)
	}

	// Init rand seed for randomBuf()
	rand.Seed(time.Now().UTC().UnixNano())
	return s, nil
}

// Run runs the server.
func (s *Server) Run() error {
	// Start HTTP API to report tidb info such as TPS.
	if s.cfg.ReportStatus {
		s.startStatusHTTP()
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Err.Error() == "use of closed network connection" {
					return nil
				}
			}
			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}
		if s.shouldStopListener() {
			conn.Close()
			break
		}
		go s.onConn(conn)
	}
	s.listener.Close()
	s.listener = nil
	for {
		log.Errorf("listener stopped, waiting for manual kill.")
		time.Sleep(time.Minute)
	}
}

func (s *Server) shouldStopListener() bool {
	select {
	case <-s.stopListenerCh:
		return true
	default:
		return false
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		s.listener.Close()
		s.listener = nil
	}
}

// onConn runs in its own goroutine, handles queries from this connection.
func (s *Server) onConn(c net.Conn) {
	conn := s.newConn(c)
	defer func() {
		log.Infof("[%d] close connection", conn.connectionID)
	}()

	if err := conn.handshake(); err != nil {
		// Some keep alive services will send request to TiDB and disconnect immediately.
		// So we use info log level.
		log.Infof("handshake error %s", errors.ErrorStack(err))
		c.Close()
		return
	}

	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	connections := len(s.clients)
	s.rwlock.Unlock()
	connGauge.Set(float64(connections))

	conn.Run()
}

// ShowProcessList implements the SessionManager interface.
func (s *Server) ShowProcessList() []util.ProcessInfo {
	var rs []util.ProcessInfo
	s.rwlock.RLock()
	for _, client := range s.clients {
		if client.killed {
			continue
		}
		rs = append(rs, client.ctx.ShowProcess())
	}
	s.rwlock.RUnlock()
	return rs
}

// Kill implements the SessionManager interface.
func (s *Server) Kill(connectionID uint64, query bool) {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	conn, ok := s.clients[uint32(connectionID)]
	if !ok {
		return
	}

	conn.ctx.Cancel()
	if !query {
		conn.killed = true
	}
}

// Server error codes.
const (
	codeUnknownFieldType  = 1
	codeInvalidPayloadLen = 2
	codeInvalidSequence   = 3
	codeInvalidType       = 4

	codeNotAllowedCommand = 1148
	codeAccessDenied      = mysql.ErrAccessDenied
)

func init() {
	serverMySQLErrCodes := map[terror.ErrCode]uint16{
		codeNotAllowedCommand: mysql.ErrNotAllowedCommand,
		codeAccessDenied:      mysql.ErrAccessDenied,
	}
	terror.ErrClassToMySQLCodes[terror.ClassServer] = serverMySQLErrCodes
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/opennetworkinglab/onos-warden/warden"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"io"
	"net"
	"sync"
)

type cluster struct {
	ad    *warden.ClusterAdvertisement
	agent warden.ClusterAgentService_AgentClustersServer
}

type request struct {
	req    *warden.ClusterRequest
	client warden.ClusterClientService_ServerClustersServer
}

type key struct {
	cId   string
	cType string
}

type wardenServer struct {
	lock sync.Mutex

	// mapping from key(ClusterId, ClusterType) to cluster resources
	clusters map[key]cluster

	// mapping from RequestId to key(ClusterId, ClusterType)
	requests map[string]key

	// registries of client and agent streams
	clients map[warden.ClusterClientService_ServerClustersServer]bool
	agents  map[warden.ClusterAgentService_AgentClustersServer]bool

	// setup a queue for incoming requests from the client
	//    - queue will be served by a worker that applies some "policy" / business logic and
	//      relays the requests to one of the selected agent
	incomingReq chan request
}

func (s *wardenServer) ServerClusters(stream warden.ClusterClientService_ServerClustersServer) error {
	s.lock.Lock()

	// register the stream so that we can send it new information to all active client
	s.clients[stream] = true

	// we can use the defer mechanism to prune the stream
	defer delete(s.clients, stream)

	// send what we have, i.e. send them existing clusters
	for _, cluster := range s.clusters {
		fmt.Println("Sending update", stream, cluster.ad)
		stream.Send(cluster.ad)
	}

	s.lock.Unlock()

	//FIXME revoke all duration == -1 requests if client disconnects

	// setup a go routing that will poll for requests from the client
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("client stream closed", stream)
			return nil
		}
		if err != nil {
			fmt.Println("client stream error", stream)
			return err
		}

		// enqueue the request
		s.incomingReq <- request{in, stream}
	}
	return nil
}

func (s *wardenServer) AgentClusters(stream warden.ClusterAgentService_AgentClustersServer) error {
	s.lock.Lock()

	// register the stream into the inventory of active agent
	s.agents[stream] = true

	s.lock.Unlock()

	// defer mechanism to prune the inventory
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()

		// remove cells from the warden map when agent disappears
		for id, cl := range s.clusters {
			//TODO maybe we should time these out instead? in case, the agent is coming right back
			if cl.agent == stream {
				delete(s.clusters, id)
			}
			if rId := cl.ad.RequestId; rId != "" {
				delete(s.requests, rId)
				//TODO need to send UNAVAILABLE
			}
		}
		delete(s.agents, stream)
		fmt.Println(s.clusters, s.agents)
	}()

	// setup polling loop for receiving new cluster advertisements
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("agent stream closed", stream)
			return nil
		}
		if err != nil {
			fmt.Println("agent stream error", stream, err)
			return err
		}
		fmt.Println(in)

		s.lock.Lock()

		// update the in-memory structures
		k := key{in.ClusterId, in.ClusterType}
		existing, ok := s.clusters[k]
		if ok && in.RequestId != existing.ad.RequestId {
			// reservation is no longer assocated with the old request
			delete(s.requests, existing.ad.RequestId)
		}
		s.clusters[k] = cluster{in, stream}
		if in.RequestId != "" {
			s.requests[in.RequestId] = k
		}

		fmt.Println(s.clusters)
		fmt.Println(s.requests)

		// relay the message about the updated resource
		for c := range s.clients {
			fmt.Println("Sending update", c, in)
			c.Send(in)
		}

		s.lock.Unlock()
	}

	return nil
}

func (s *wardenServer) processRequests() {
	for {
		request := <-s.incomingReq
		fmt.Println()
		req := request.req
		client := request.client

		func() {
			s.lock.Lock()
			defer s.lock.Unlock()

			// Check to see if we have already satisfied the request
			rId := req.RequestId
			k, ok := s.requests[rId]
			if ok {
				ad, ok := s.clusters[k]
				if ok {
					// unicast the cluster to client
					fmt.Println("send ad to client", ad)
					client.Send(ad.ad)
					ad.agent.Send(req)
					return
				} else {
					// this shouldn't happen, but we'll do some cleanup if it does
					delete(s.requests, rId)
				}
			}

			// If not, find an appropriate cluster for the request
			for _, c := range s.clusters {
				if req.ClusterType != "" && req.ClusterType != c.ad.ClusterType {
					continue
				}
				if req.ClusterId != "" && req.ClusterId != c.ad.ClusterId {
					continue
				}
				// find the first one that is available
				if c.ad.State == warden.ClusterAdvertisement_AVAILABLE {
					// relay the request to the agent that advertised it
					fmt.Println("sending request to cluster", c, req)
					c.agent.Send(req)
					s.requests[rId] = key{c.ad.ClusterId, c.ad.ClusterType}
					return
				}
			}
		}()

	}
}

func newServer() *wardenServer {
	s := new(wardenServer)
	s.clusters = make(map[key]cluster)
	s.requests = make(map[string]key)
	s.clients = make(map[warden.ClusterClientService_ServerClustersServer]bool)
	s.agents = make(map[warden.ClusterAgentService_AgentClustersServer]bool)
	s.incomingReq = make(chan request)
	return s
}

func main() {
	lis, err := net.Listen("tcp", ":1234")
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	s := newServer()
	go s.processRequests()
	warden.RegisterClusterClientServiceServer(grpcServer, s)
	warden.RegisterClusterAgentServiceServer(grpcServer, s)
	fmt.Println("starting to serve...")
	grpcServer.Serve(lis)
}
<commit_msg>Fixed bug that deletes all requests when any agent disconnects<commit_after>package main

import (
	"fmt"
	"github.com/opennetworkinglab/onos-warden/warden"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"io"
	"net"
	"sync"
)

type cluster struct {
	ad    *warden.ClusterAdvertisement
	agent warden.ClusterAgentService_AgentClustersServer
}

type request struct {
	req    *warden.ClusterRequest
	client warden.ClusterClientService_ServerClustersServer
}

type key struct {
	cId   string
	cType string
}

type wardenServer struct {
	lock sync.Mutex

	// mapping from key(ClusterId, ClusterType) to cluster resources
	clusters map[key]cluster

	// mapping from RequestId to key(ClusterId, ClusterType)
	requests map[string]key

	// registries of client and agent streams
	clients map[warden.ClusterClientService_ServerClustersServer]bool
	agents  map[warden.ClusterAgentService_AgentClustersServer]bool

	// setup a queue for incoming requests from the client
	//    - queue will be served by a worker that applies some "policy" / business logic and
	//      relays the requests to one of the selected agent
	incomingReq chan request
}

func (s *wardenServer) ServerClusters(stream warden.ClusterClientService_ServerClustersServer) error {
	s.lock.Lock()

	// register the stream so that we can send it new information to all active client
	s.clients[stream] = true

	// we can use the defer mechanism to prune the stream
	defer delete(s.clients, stream)

	// send what we have, i.e. send them existing clusters
	for _, cluster := range s.clusters {
		fmt.Println("Sending update", stream, cluster.ad)
		stream.Send(cluster.ad)
	}

	s.lock.Unlock()

	//FIXME revoke all duration == -1 requests if client disconnects

	// setup a go routing that will poll for requests from the client
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("client stream closed", stream)
			return nil
		}
		if err != nil {
			fmt.Println("client stream error", stream)
			return err
		}

		// enqueue the request
		s.incomingReq <- request{in, stream}
	}
	return nil
}

func (s *wardenServer) AgentClusters(stream warden.ClusterAgentService_AgentClustersServer) error {
	s.lock.Lock()

	// register the stream into the inventory of active agent
	s.agents[stream] = true

	s.lock.Unlock()

	// defer mechanism to prune the inventory
	defer func() {
		s.lock.Lock()
		defer s.lock.Unlock()

		// remove cells from the warden map when agent disappears
		for id, cl := range s.clusters {
			//TODO maybe we should time these out instead? in case, the agent is coming right back
			if cl.agent == stream {
				delete(s.clusters, id)
				if rId := cl.ad.RequestId; rId != "" {
					delete(s.requests, rId)
					//TODO need to send UNAVAILABLE
				}
			}
		}
		delete(s.agents, stream)
		fmt.Println(s.clusters, s.agents)
	}()

	// setup polling loop for receiving new cluster advertisements
	for {
		in, err := stream.Recv()
		if err == io.EOF {
			fmt.Println("agent stream closed", stream)
			return nil
		}
		if err != nil {
			fmt.Println("agent stream error", stream, err)
			return err
		}
		fmt.Println(in)

		s.lock.Lock()

		// update the in-memory structures
		k := key{in.ClusterId, in.ClusterType}
		existing, ok := s.clusters[k]
		if ok && in.RequestId != existing.ad.RequestId {
			// reservation is no longer assocated with the old request
			delete(s.requests, existing.ad.RequestId)
		}
		s.clusters[k] = cluster{in, stream}
		if in.RequestId != "" {
			s.requests[in.RequestId] = k
		}

		fmt.Println(s.clusters)
		fmt.Println(s.requests)

		// relay the message about the updated resource
		for c := range s.clients {
			fmt.Println("Sending update", c, in)
			c.Send(in)
		}

		s.lock.Unlock()
	}

	return nil
}

func (s *wardenServer) processRequests() {
	for {
		request := <-s.incomingReq
		req := request.req
		client := request.client

		func() {
			s.lock.Lock()
			defer s.lock.Unlock()

			// Check to see if we have already satisfied the request
			rId := req.RequestId
			k, ok := s.requests[rId]
			if ok {
				ad, ok := s.clusters[k]
				if ok {
					// unicast the cluster to client
					fmt.Println("send ad to client", ad)
					client.Send(ad.ad)
					ad.agent.Send(req)
					return
				} else {
					// this shouldn't happen, but we'll do some cleanup if it does
					delete(s.requests, rId)
				}
			}

			// If not, find an appropriate cluster for the request
			for _, c := range s.clusters {
				if req.ClusterType != "" && req.ClusterType != c.ad.ClusterType {
					continue
				}
				if req.ClusterId != "" && req.ClusterId != c.ad.ClusterId {
					continue
				}
				// find the first one that is available
				if c.ad.State == warden.ClusterAdvertisement_AVAILABLE {
					// relay the request to the agent that advertised it
					fmt.Println("sending request to cluster", c, req)
					c.agent.Send(req)
					s.requests[rId] = key{c.ad.ClusterId, c.ad.ClusterType}
					return
				}
			}
		}()

	}
}

func newServer() *wardenServer {
	s := new(wardenServer)
	s.clusters = make(map[key]cluster)
	s.requests = make(map[string]key)
	s.clients = make(map[warden.ClusterClientService_ServerClustersServer]bool)
	s.agents = make(map[warden.ClusterAgentService_AgentClustersServer]bool)
	s.incomingReq = make(chan request)
	return s
}

func main() {
	lis, err := net.Listen("tcp", ":1234")
	if err != nil {
		grpclog.Fatalf("failed to listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	s := newServer()
	go s.processRequests()
	warden.RegisterClusterClientServiceServer(grpcServer, s)
	warden.RegisterClusterAgentServiceServer(grpcServer, s)
	fmt.Println("starting to serve...")
	grpcServer.Serve(lis)
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/ops"
	"github.com/getlantern/proxy"
	"github.com/getlantern/proxy/filters"
	"github.com/getlantern/tlsdefaults"

	"github.com/getlantern/http-proxy/buffers"
	"github.com/getlantern/http-proxy/listeners"
)

var (
	testingLocal = false
	log          = golog.LoggerFor("server")
)

type listenerGenerator func(net.Listener) net.Listener

// Opts are used to configure a Server
type Opts struct {
	IdleTimeout time.Duration
	Filter      filters.Filter
	Dial        proxy.DialFunc
}

// Server is an HTTP proxy server.
type Server struct {
	// Allow is a function that determines whether or not to allow connections
	// from the given IP address. If unspecified, all connections are allowed.
	Allow              func(string) bool
	proxy              proxy.Proxy
	listenerGenerators []listenerGenerator
}

// New constructs a new HTTP proxy server using the given options
func New(opts *Opts) *Server {
	p, _ := proxy.New(&proxy.Opts{
		IdleTimeout:        opts.IdleTimeout,
		Dial:               opts.Dial,
		Filter:             opts.Filter,
		BufferSource:       buffers.Pool(),
		OKWaitsForUpstream: true,
		OnError: func(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
			status := http.StatusBadGateway
			if read {
				status = http.StatusBadRequest
			}
			return &http.Response{
				Request:    req,
				StatusCode: status,
				Body:       ioutil.NopCloser(strings.NewReader(err.Error())),
			}
		},
	})
	return &Server{
		proxy: p,
	}
}

func (s *Server) AddListenerWrappers(listenerGens ...listenerGenerator) {
	for _, g := range listenerGens {
		s.listenerGenerators = append(s.listenerGenerators, g)
	}
}

func (s *Server) ListenAndServeHTTP(addr string, readyCb func(addr string)) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Debugf("Listen http on %s", addr)
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) ListenAndServeHTTPS(addr, keyfile, certfile string, readyCb func(addr string)) error {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	listener, err := tlsdefaults.NewListener(s.wrapListenerIfNecessary(l), keyfile, certfile)
	if err != nil {
		return err
	}
	log.Debugf("Listen https on %s", addr)
	return s.serve(listener, readyCb)
}

func (s *Server) Serve(listener net.Listener, readyCb func(addr string)) error {
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) serve(listener net.Listener, readyCb func(addr string)) error {
	l := listeners.NewDefaultListener(listener)

	for _, wrap := range s.listenerGenerators {
		l = wrap(l)
	}

	if readyCb != nil {
		readyCb(l.Addr().String())
	}

	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				// delay code based on net/http.Server
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Errorf("http: Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return errors.New("Error accepting: %v", err)
		}
		tempDelay = 0
		s.handle(conn)
	}
}

func (s *Server) handle(conn net.Conn) {
	wrapConn, isWrapConn := conn.(listeners.WrapConn)
	if isWrapConn {
		wrapConn.OnState(http.StateNew)
	}
	go s.doHandle(conn, isWrapConn, wrapConn)
}

func (s *Server) doHandle(conn net.Conn, isWrapConn bool, wrapConn listeners.WrapConn) {
	clientIP := ""
	remoteAddr := conn.RemoteAddr()
	if remoteAddr != nil {
		clientIP, _, _ = net.SplitHostPort(remoteAddr.String())
	}
	op := ops.Begin("http_proxy_handle").Set("client_ip", clientIP)
	defer op.End()

	defer func() {
		p := recover()
		if p != nil {
			err := log.Errorf("Caught panic handling connection from %v: %v", conn.RemoteAddr(), p)
			if op != nil {
				op.FailIf(err)
			}
			safeClose(conn)
		}
	}()

	err := s.proxy.Handle(context.Background(), conn, conn)
	if err != nil {
		op.FailIf(log.Errorf("Error handling connection from %v: %v", conn.RemoteAddr(), err))
	}
	if isWrapConn {
		wrapConn.OnState(http.StateClosed)
	}
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection from %v: %v", conn.RemoteAddr(), p)
		}
	}()

	conn.Close()
}

func (s *Server) wrapListenerIfNecessary(l net.Listener) net.Listener {
	if s.Allow != nil {
		log.Debug("Wrapping listener with Allow")
		return &allowinglistener{l, s.Allow}
	}
	return l
}

type allowinglistener struct {
	wrapped net.Listener
	allow   func(string) bool
}

func (l *allowinglistener) Accept() (net.Conn, error) {
	conn, err := l.wrapped.Accept()
	if err != nil {
		return conn, err
	}

	ip := ""
	remoteAddr := conn.RemoteAddr()
	switch addr := remoteAddr.(type) {
	case *net.TCPAddr:
		ip = addr.IP.String()
	case *net.UDPAddr:
		ip = addr.IP.String()
	default:
		log.Errorf("Remote addr %v is of unknown type %v, unable to determine IP", remoteAddr, reflect.TypeOf(remoteAddr))
		return conn, err
	}
	if !l.allow(ip) {
		conn.Close()
		// Note - we don't return an error, because that causes http.Server to stop
		// serving.
	}

	return conn, err
}

func (l *allowinglistener) Close() error {
	return l.wrapped.Close()
}

func (l *allowinglistener) Addr() net.Addr {
	return l.wrapped.Addr()
}
<commit_msg>dont log DNS errors as errors<commit_after>package server

import (
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/ops"
	"github.com/getlantern/proxy"
	"github.com/getlantern/proxy/filters"
	"github.com/getlantern/tlsdefaults"

	"github.com/getlantern/http-proxy/buffers"
	"github.com/getlantern/http-proxy/listeners"
)

var (
	testingLocal = false
	log          = golog.LoggerFor("server")
)

type listenerGenerator func(net.Listener) net.Listener

// Opts are used to configure a Server
type Opts struct {
	IdleTimeout time.Duration
	Filter      filters.Filter
	Dial        proxy.DialFunc
}

// Server is an HTTP proxy server.
type Server struct {
	// Allow is a function that determines whether or not to allow connections
	// from the given IP address. If unspecified, all connections are allowed.
	Allow              func(string) bool
	proxy              proxy.Proxy
	listenerGenerators []listenerGenerator
}

// New constructs a new HTTP proxy server using the given options
func New(opts *Opts) *Server {
	p, _ := proxy.New(&proxy.Opts{
		IdleTimeout:        opts.IdleTimeout,
		Dial:               opts.Dial,
		Filter:             opts.Filter,
		BufferSource:       buffers.Pool(),
		OKWaitsForUpstream: true,
		OnError: func(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
			status := http.StatusBadGateway
			if read {
				status = http.StatusBadRequest
			}
			return &http.Response{
				Request:    req,
				StatusCode: status,
				Body:       ioutil.NopCloser(strings.NewReader(err.Error())),
			}
		},
	})
	return &Server{
		proxy: p,
	}
}

func (s *Server) AddListenerWrappers(listenerGens ...listenerGenerator) {
	for _, g := range listenerGens {
		s.listenerGenerators = append(s.listenerGenerators, g)
	}
}

func (s *Server) ListenAndServeHTTP(addr string, readyCb func(addr string)) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Debugf("Listen http on %s", addr)
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) ListenAndServeHTTPS(addr, keyfile, certfile string, readyCb func(addr string)) error {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	listener, err := tlsdefaults.NewListener(s.wrapListenerIfNecessary(l), keyfile, certfile)
	if err != nil {
		return err
	}
	log.Debugf("Listen https on %s", addr)
	return s.serve(listener, readyCb)
}

func (s *Server) Serve(listener net.Listener, readyCb func(addr string)) error {
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) serve(listener net.Listener, readyCb func(addr string)) error {
	l := listeners.NewDefaultListener(listener)

	for _, wrap := range s.listenerGenerators {
		l = wrap(l)
	}

	if readyCb != nil {
		readyCb(l.Addr().String())
	}

	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				// delay code based on net/http.Server
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Errorf("http: Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return errors.New("Error accepting: %v", err)
		}
		tempDelay = 0
		s.handle(conn)
	}
}

func (s *Server) handle(conn net.Conn) {
	wrapConn, isWrapConn := conn.(listeners.WrapConn)
	if isWrapConn {
		wrapConn.OnState(http.StateNew)
	}
	go s.doHandle(conn, isWrapConn, wrapConn)
}

func (s *Server) doHandle(conn net.Conn, isWrapConn bool, wrapConn listeners.WrapConn) {
	clientIP := ""
	remoteAddr := conn.RemoteAddr()
	if remoteAddr != nil {
		clientIP, _, _ = net.SplitHostPort(remoteAddr.String())
	}
	op := ops.Begin("http_proxy_handle").Set("client_ip", clientIP)
	defer op.End()

	defer func() {
		p := recover()
		if p != nil {
			err := log.Errorf("Caught panic handling connection from %v: %v", conn.RemoteAddr(), p)
			if op != nil {
				op.FailIf(err)
			}
			safeClose(conn)
		}
	}()

	err := s.proxy.Handle(context.Background(), conn, conn)
	if err != nil {
		var failErr error
		if strings.Contains(err.Error(), "no such host") {
			// We don't want to log no such host (DNS) errors as true errors.
			failErr = errors.New("DNS error handling connection from %v: %v", conn.RemoteAddr(), err)
		} else {
			failErr = log.Errorf("Error handling connection from %v: %v", conn.RemoteAddr(), err)
		}
		op.FailIf(failErr)
	}
	if isWrapConn {
		wrapConn.OnState(http.StateClosed)
	}
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection from %v: %v", conn.RemoteAddr(), p)
		}
	}()

	conn.Close()
}

func (s *Server) wrapListenerIfNecessary(l net.Listener) net.Listener {
	if s.Allow != nil {
		log.Debug("Wrapping listener with Allow")
		return &allowinglistener{l, s.Allow}
	}
	return l
}

type allowinglistener struct {
	wrapped net.Listener
	allow   func(string) bool
}

func (l *allowinglistener) Accept() (net.Conn, error) {
	conn, err := l.wrapped.Accept()
	if err != nil {
		return conn, err
	}

	ip := ""
	remoteAddr := conn.RemoteAddr()
	switch addr := remoteAddr.(type) {
	case *net.TCPAddr:
		ip = addr.IP.String()
	case *net.UDPAddr:
		ip = addr.IP.String()
	default:
		log.Errorf("Remote addr %v is of unknown type %v, unable to determine IP", remoteAddr, reflect.TypeOf(remoteAddr))
		return conn, err
	}
	if !l.allow(ip) {
		conn.Close()
		// Note - we don't return an error, because that causes http.Server to stop
		// serving.
	}

	return conn, err
}

func (l *allowinglistener) Close() error {
	return l.wrapped.Close()
}

func (l *allowinglistener) Addr() net.Addr {
	return l.wrapped.Addr()
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/spf13/viper"
	ctrlDio "github.com/yroffin/jarvis-go-ext/server/dio"
	"github.com/yroffin/jarvis-go-ext/server/utils/cron"
	"github.com/yroffin/jarvis-go-ext/server/utils/native"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
)

// Start : start the jarvis server
func Start() {
	e := echo.New()
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	api := e.Group("/api")
	{ // routes for /api
		dio := api.Group("/dio")
		{ // routes for /api/dio
			dio.Post("", ctrlDio.HandlePost)
		}
	}

	// init wiringPi library
	native.InitWiringPi()

	// init cron
	cron.Init("@every 60s")

	port := viper.GetString("jarvis.core.port")
	e.Run(standard.New(":" + port))
}
<commit_msg>Correction<commit_after>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/spf13/viper"
	ctrlDio "github.com/yroffin/jarvis-go-ext/server/dio"
	"github.com/yroffin/jarvis-go-ext/server/utils/cron"
	"github.com/yroffin/jarvis-go-ext/server/utils/logger"
	"github.com/yroffin/jarvis-go-ext/server/utils/native"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
)

// Start : start the jarvis server
func Start() {
	e := echo.New()
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	api := e.Group("/api")
	{ // routes for /api
		dio := api.Group("/dio")
		{ // routes for /api/dio
			dio.Post("", ctrlDio.HandlePost)
		}
	}

	// init wiringPi library
	native.InitWiringPi()

	// init cron
	cron.Init("@every 60s")

	// get prot from config
	intf := viper.GetString("jarvis.module.interface")
	port := viper.GetString("jarvis.module.port")

	logger.NewLogger().WithFields(log.Fields{
		"interface": intf,
		"port":      port,
	}).Info("DIO")

	e.Run(standard.New(intf + ":" + port))
}
<|endoftext|>


<commit_before>package crawl

import (
	"errors"
	"fmt"
	"log"
	"strconv"
	"time"

	"github.com/nzai/stockrecorder/config"
)

//	启动
func Start() {

	markets := config.GetArray("market", "markets")
	if len(markets) == 0 {
		//	未定义市场就直接返回
		return
	}

	for _, market := range markets {
		//	抓取市场数据
		err := marketJob(market)
		if err != nil {
			log.Fatalf("启动[%s]数据抓取任务失败:%v", market, err)
		}
	}
}

//	抓取市场数据任务
func marketJob(market string) error {
	//	任务启动时间(hour)
	endhour := config.GetString(market, "endhour", "")
	if endhour == "" {
		return errors.New(fmt.Sprintf("市场[%s]的starthour配置有误", market))
	}

	hour, err := strconv.Atoi(endhour)
	if err != nil {
		return err
	}

	now := time.Now().UTC()
	startTime := time.Date(now.Year(), now.Month(), now.Day(), hour%24, 0, 0, 0, now.Location())
	//	现在距离开始的时间间隔
	duration := startTime.Sub(now)
	if now.After(startTime) {
		//	今天的开始时间已经过了
		duration = duration + time.Hour*24
	}

	log.Printf("%s后开始抓取%s的数据", duration.String(), market)

	time.AfterFunc(duration, func() {
		//	到点后立即运行第一次
		go func(m string, h int) {
			//	抓取雅虎今日数据
			err := yahooToday(market, h)
			if err != nil {
				log.Fatal(err)
			}
		}(market, hour)

		//	之后每天运行一次
		ticker := time.NewTicker(time.Hour * 24)
		for _ = range ticker.C {
			//	抓取雅虎今日数据
			err := yahooToday(market, hour)
			if err != nil {
				log.Fatal(err)
			}
		}
	})

	return nil
}
<commit_msg>简化任务逻辑，任务启动函数只负责依次启动各市场数据的抓取<commit_after>package crawl

import (
	"log"

	"github.com/nzai/stockrecorder/config"
)

//	启动
func Start() {

	markets := config.GetArray("market", "markets")
	if len(markets) == 0 {
		//	未定义市场就直接返回
		return
	}

	for _, market := range markets {
		//	抓取市场数据
		err := marketAll(market)
		if err != nil {
			log.Fatalf("启动[%s]数据抓取任务失败:%v", market, err)
		}
	}
}
<|endoftext|>