<commit_before>package server

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"mime"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/golang/glog"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/core"
	"github.com/livepeer/go-livepeer/net"
)

const protoVerLPT = "Livepeer-Transcoder-1.0"
const transcodingErrorMimeType = "livepeer/transcoding-error"

var errSecret = errors.New("Invalid secret")

// Standalone Transcoder

// RunTranscoder is main routing of standalone transcoder
// Exiting it will terminate executable
func RunTranscoder(n *core.LivepeerNode, orchAddr string, capacity int) {
	expb := backoff.NewExponentialBackOff()
	expb.MaxInterval = time.Minute
	expb.MaxElapsedTime = 0
	backoff.Retry(func() error {
		glog.Info("Registering transcoder to ", orchAddr)
		err := runTranscoder(n, orchAddr, capacity)
		glog.Info("Unregistering transcoder: ", err)
		if _, fatal := err.(core.RemoteTranscoderFatalError); fatal {
			glog.Info("Terminating transcoder because of ", err)
			// Returning nil here will make `backoff` to stop trying to reconnect and exit
			return nil
		}
		// By returning error we tell `backoff` to try to connect again
		return err
	}, expb)
}

func checkTranscoderError(err error) error {
	if err != nil {
		s := status.Convert(err)
		if s.Message() == errSecret.Error() { // consider this unrecoverable
			return core.NewRemoteTranscoderFatalError(errSecret)
		}
		if status.Code(err) == codes.Canceled {
			return core.NewRemoteTranscoderFatalError(fmt.Errorf("Execution interrupted"))
		}
	}
	return err
}

func runTranscoder(n *core.LivepeerNode, orchAddr string, capacity int) error {
	tlsConfig := &tls.Config{InsecureSkipVerify: true}
	conn, err := grpc.Dial(orchAddr,
		grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
	if err != nil {
		glog.Error("Did not connect transcoder to orchesrator: ", err)
		return err
	}
	defer conn.Close()

	c := net.NewTranscoderClient(conn)
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	// Silence linter
	defer cancel()
	r, err := c.RegisterTranscoder(ctx, &net.RegisterRequest{Secret: n.OrchSecret, Capacity: int64(capacity)})
	if err := checkTranscoderError(err); err != nil {
		glog.Error("Could not register transcoder to orchestrator ", err)
		return err
	}

	// Catch interrupt signal to shut down transcoder
	exitc := make(chan os.Signal)
	signal.Notify(exitc, os.Interrupt)
	defer signal.Stop(exitc)
	go func() {
		select {
		case sig := <-exitc:
			glog.Infof("Exiting Livepeer Transcoder: %v", sig)
			// Cancelling context will close connection to orchestrator
			cancel()
			return
		}
	}()

	httpc := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}}
	var wg sync.WaitGroup
	for {
		notify, err := r.Recv()
		if err := checkTranscoderError(err); err != nil {
			glog.Infof(`End of stream recieve cylcle because of err="%v", waiting for running transcode jobs to complete`, err)
			wg.Wait()
			return err
		}
		wg.Add(1)
		go func() {
			runTranscode(n, orchAddr, httpc, notify)
			wg.Done()
		}()
	}
}

func runTranscode(n *core.LivepeerNode, orchAddr string, httpc *http.Client, notify *net.NotifySegment) {
	profiles, err := common.TxDataToVideoProfile(hex.EncodeToString(notify.Profiles))
	if err != nil {
		glog.Info("Unable to deserialize profiles ", err)
	}

	glog.Infof("Transcoding taskId=%d url=%s", notify.TaskId, notify.Url)
	var contentType string
	var body bytes.Buffer

	tData, err := n.Transcoder.Transcode(notify.Url, profiles)
	glog.V(common.VERBOSE).Infof("Transcoding done for taskId=%d url=%s err=%v", notify.TaskId, notify.Url, err)
	if err != nil {
		glog.Error("Unable to transcode ", err)
		body.Write([]byte(err.Error()))
		contentType = transcodingErrorMimeType
	} else {
		boundary := randName()
		w := multipart.NewWriter(&body)
		for _, v := range tData {
			w.SetBoundary(boundary)
			hdrs := textproto.MIMEHeader{
				"Content-Type":   {"video/MP2T"},
				"Content-Length": {strconv.Itoa(len(v))},
			}
			fw, err := w.CreatePart(hdrs)
			if err != nil {
				glog.Error("Could not create multipart part ", err)
			}
			io.Copy(fw, bytes.NewBuffer(v))
		}
		w.Close()
		contentType = "multipart/mixed; boundary=" + boundary
	}
	req, err := http.NewRequest("POST", "https://"+orchAddr+"/transcodeResults", &body)
	if err != nil {
		glog.Error("Error posting results ", err)
	}
	req.Header.Set("Authorization", protoVerLPT)
	req.Header.Set("Credentials", n.OrchSecret)
	req.Header.Set("Content-Type", contentType)
	req.Header.Set("TaskId", strconv.FormatInt(notify.TaskId, 10))
	resp, err := httpc.Do(req)
	if err != nil {
		glog.Error("Error submitting results ", err)
	}
	ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	glog.V(common.VERBOSE).Infof("Transcoding done results sent for taskId=%d url=%s err=%v", notify.TaskId, notify.Url, err)
}

// Orchestrator gRPC

func (h *lphttp) RegisterTranscoder(req *net.RegisterRequest, stream net.Transcoder_RegisterTranscoderServer) error {
	from := common.GetConnectionAddr(stream.Context())
	glog.Infof("Got a RegisterTranscoder request from transcoder=%s", from)

	if req.Secret != h.orchestrator.TranscoderSecret() {
		glog.Info(errSecret.Error())
		return errSecret
	}

	// blocks until stream is finished
	h.orchestrator.ServeTranscoder(stream, int(req.Capacity))
	return nil
}

// Orchestrator HTTP

func (h *lphttp) TranscodeResults(w http.ResponseWriter, r *http.Request) {
	orch := h.orchestrator

	authType := r.Header.Get("Authorization")
	creds := r.Header.Get("Credentials")
	if protoVerLPT != authType {
		glog.Error("Invalid auth type ", authType)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	if creds != orch.TranscoderSecret() {
		glog.Error("Invalid shared secret")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		glog.Error("Error getting mime type ", err)
		http.Error(w, err.Error(), http.StatusUnsupportedMediaType)
		return
	}

	tid, err := strconv.ParseInt(r.Header.Get("TaskId"), 10, 64)
	if err != nil {
		glog.Error("Could not parse task ID ", err)
		http.Error(w, "Invalid Task ID", http.StatusBadRequest)
		return
	}

	var res core.RemoteTranscoderResult
	if transcodingErrorMimeType == mediaType {
		w.Write([]byte("OK"))
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			glog.Error("Unable to read transcoding error body ", err)
			res.Err = err
		} else {
			res.Err = fmt.Errorf(string(body))
		}
		glog.Error("Trascoding error ", res.Err)
		orch.TranscoderResults(tid, &res)
		return
	}

	var segments [][]byte
	if "multipart/mixed" == mediaType {
		mr := multipart.NewReader(r.Body, params["boundary"])
		for {
			p, err := mr.NextPart()
			if err == io.EOF {
				break
			}
			if err != nil {
				glog.Error("Could not process multipart part ", err)
				res.Err = err
				break
			}
			body, err := ioutil.ReadAll(p)
			if err != nil {
				glog.Error("Error reading body ", err)
				res.Err = err
				break
			}
			segments = append(segments, body)
		}
		res.Segments = segments
		orch.TranscoderResults(tid, &res)
	}
	if res.Err != nil {
		http.Error(w, res.Err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write([]byte("OK"))
}

// utils

func randName() string {
	rand.Seed(time.Now().UnixNano())
	x := make([]byte, 10, 10)
	for i := 0; i < len(x); i++ {
		x[i] = byte(rand.Uint32())
	}
	return hex.EncodeToString(x)
}
<commit_msg>Watch for `SIGTERM` signal for graceful shutdown of transcoder<commit_after>package server

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"mime"
	"mime/multipart"
	"net/http"
	"net/textproto"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/golang/glog"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	"github.com/livepeer/go-livepeer/common"
	"github.com/livepeer/go-livepeer/core"
	"github.com/livepeer/go-livepeer/net"
)

const protoVerLPT = "Livepeer-Transcoder-1.0"
const transcodingErrorMimeType = "livepeer/transcoding-error"

var errSecret = errors.New("Invalid secret")

// Standalone Transcoder

// RunTranscoder is main routing of standalone transcoder
// Exiting it will terminate executable
func RunTranscoder(n *core.LivepeerNode, orchAddr string, capacity int) {
	expb := backoff.NewExponentialBackOff()
	expb.MaxInterval = time.Minute
	expb.MaxElapsedTime = 0
	backoff.Retry(func() error {
		glog.Info("Registering transcoder to ", orchAddr)
		err := runTranscoder(n, orchAddr, capacity)
		glog.Info("Unregistering transcoder: ", err)
		if _, fatal := err.(core.RemoteTranscoderFatalError); fatal {
			glog.Info("Terminating transcoder because of ", err)
			// Returning nil here will make `backoff` to stop trying to reconnect and exit
			return nil
		}
		// By returning error we tell `backoff` to try to connect again
		return err
	}, expb)
}

func checkTranscoderError(err error) error {
	if err != nil {
		s := status.Convert(err)
		if s.Message() == errSecret.Error() { // consider this unrecoverable
			return core.NewRemoteTranscoderFatalError(errSecret)
		}
		if status.Code(err) == codes.Canceled {
			return core.NewRemoteTranscoderFatalError(fmt.Errorf("Execution interrupted"))
		}
	}
	return err
}

func runTranscoder(n *core.LivepeerNode, orchAddr string, capacity int) error {
	tlsConfig := &tls.Config{InsecureSkipVerify: true}
	conn, err := grpc.Dial(orchAddr,
		grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)))
	if err != nil {
		glog.Error("Did not connect transcoder to orchesrator: ", err)
		return err
	}
	defer conn.Close()

	c := net.NewTranscoderClient(conn)
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	// Silence linter
	defer cancel()
	r, err := c.RegisterTranscoder(ctx, &net.RegisterRequest{Secret: n.OrchSecret, Capacity: int64(capacity)})
	if err := checkTranscoderError(err); err != nil {
		glog.Error("Could not register transcoder to orchestrator ", err)
		return err
	}

	// Catch interrupt signal to shut down transcoder
	exitc := make(chan os.Signal)
	signal.Notify(exitc, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(exitc)
	go func() {
		select {
		case sig := <-exitc:
			glog.Infof("Exiting Livepeer Transcoder: %v", sig)
			// Cancelling context will close connection to orchestrator
			cancel()
			return
		}
	}()

	httpc := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}}
	var wg sync.WaitGroup
	for {
		notify, err := r.Recv()
		if err := checkTranscoderError(err); err != nil {
			glog.Infof(`End of stream recieve cylcle because of err="%v", waiting for running transcode jobs to complete`, err)
			wg.Wait()
			return err
		}
		wg.Add(1)
		go func() {
			runTranscode(n, orchAddr, httpc, notify)
			wg.Done()
		}()
	}
}

func runTranscode(n *core.LivepeerNode, orchAddr string, httpc *http.Client, notify *net.NotifySegment) {
	profiles, err := common.TxDataToVideoProfile(hex.EncodeToString(notify.Profiles))
	if err != nil {
		glog.Info("Unable to deserialize profiles ", err)
	}

	glog.Infof("Transcoding taskId=%d url=%s", notify.TaskId, notify.Url)
	var contentType string
	var body bytes.Buffer

	tData, err := n.Transcoder.Transcode(notify.Url, profiles)
	glog.V(common.VERBOSE).Infof("Transcoding done for taskId=%d url=%s err=%v", notify.TaskId, notify.Url, err)
	if err != nil {
		glog.Error("Unable to transcode ", err)
		body.Write([]byte(err.Error()))
		contentType = transcodingErrorMimeType
	} else {
		boundary := randName()
		w := multipart.NewWriter(&body)
		for _, v := range tData {
			w.SetBoundary(boundary)
			hdrs := textproto.MIMEHeader{
				"Content-Type":   {"video/MP2T"},
				"Content-Length": {strconv.Itoa(len(v))},
			}
			fw, err := w.CreatePart(hdrs)
			if err != nil {
				glog.Error("Could not create multipart part ", err)
			}
			io.Copy(fw, bytes.NewBuffer(v))
		}
		w.Close()
		contentType = "multipart/mixed; boundary=" + boundary
	}
	req, err := http.NewRequest("POST", "https://"+orchAddr+"/transcodeResults", &body)
	if err != nil {
		glog.Error("Error posting results ", err)
	}
	req.Header.Set("Authorization", protoVerLPT)
	req.Header.Set("Credentials", n.OrchSecret)
	req.Header.Set("Content-Type", contentType)
	req.Header.Set("TaskId", strconv.FormatInt(notify.TaskId, 10))
	resp, err := httpc.Do(req)
	if err != nil {
		glog.Error("Error submitting results ", err)
	}
	ioutil.ReadAll(resp.Body)
	resp.Body.Close()
	glog.V(common.VERBOSE).Infof("Transcoding done results sent for taskId=%d url=%s err=%v", notify.TaskId, notify.Url, err)
}

// Orchestrator gRPC

func (h *lphttp) RegisterTranscoder(req *net.RegisterRequest, stream net.Transcoder_RegisterTranscoderServer) error {
	from := common.GetConnectionAddr(stream.Context())
	glog.Infof("Got a RegisterTranscoder request from transcoder=%s", from)

	if req.Secret != h.orchestrator.TranscoderSecret() {
		glog.Info(errSecret.Error())
		return errSecret
	}

	// blocks until stream is finished
	h.orchestrator.ServeTranscoder(stream, int(req.Capacity))
	return nil
}

// Orchestrator HTTP

func (h *lphttp) TranscodeResults(w http.ResponseWriter, r *http.Request) {
	orch := h.orchestrator

	authType := r.Header.Get("Authorization")
	creds := r.Header.Get("Credentials")
	if protoVerLPT != authType {
		glog.Error("Invalid auth type ", authType)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	if creds != orch.TranscoderSecret() {
		glog.Error("Invalid shared secret")
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	mediaType, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err != nil {
		glog.Error("Error getting mime type ", err)
		http.Error(w, err.Error(), http.StatusUnsupportedMediaType)
		return
	}

	tid, err := strconv.ParseInt(r.Header.Get("TaskId"), 10, 64)
	if err != nil {
		glog.Error("Could not parse task ID ", err)
		http.Error(w, "Invalid Task ID", http.StatusBadRequest)
		return
	}

	var res core.RemoteTranscoderResult
	if transcodingErrorMimeType == mediaType {
		w.Write([]byte("OK"))
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			glog.Error("Unable to read transcoding error body ", err)
			res.Err = err
		} else {
			res.Err = fmt.Errorf(string(body))
		}
		glog.Error("Trascoding error ", res.Err)
		orch.TranscoderResults(tid, &res)
		return
	}

	var segments [][]byte
	if "multipart/mixed" == mediaType {
		mr := multipart.NewReader(r.Body, params["boundary"])
		for {
			p, err := mr.NextPart()
			if err == io.EOF {
				break
			}
			if err != nil {
				glog.Error("Could not process multipart part ", err)
				res.Err = err
				break
			}
			body, err := ioutil.ReadAll(p)
			if err != nil {
				glog.Error("Error reading body ", err)
				res.Err = err
				break
			}
			segments = append(segments, body)
		}
		res.Segments = segments
		orch.TranscoderResults(tid, &res)
	}
	if res.Err != nil {
		http.Error(w, res.Err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write([]byte("OK"))
}

// utils

func randName() string {
	rand.Seed(time.Now().UnixNano())
	x := make([]byte, 10, 10)
	for i := 0; i < len(x); i++ {
		x[i] = byte(rand.Uint32())
	}
	return hex.EncodeToString(x)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"

	"github.com/garyburd/redigo/redis"
)

// Runner runs the code
type Runner struct {
	Ext     string `json:"ext"`
	Source  string `json:"source"`
	Version string `json:"version"`
}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	execArgs := []string{"run", "-i", "koderunr", r.Ext, r.Source}
	if r.Version != "" {
		execArgs = append(execArgs, r.Version)
	}

	cmd := exec.Command("docker", execArgs...)

	stdoutReader, stdoutWriter := io.Pipe()
	cmd.Stdout = stdoutWriter
	cmd.Stderr = stdoutWriter

	stdin, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v", err)
	}

	defer stdin.Close()
	defer stdoutWriter.Close()

	go func() {
		psc := redis.PubSubConn{Conn: conn}
		psc.Subscribe(uuid + "#stdin")
		defer psc.Close()

	StdinSubscriptionLoop:
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
				stdin.Write(n.Data)
			case error:
				break StdinSubscriptionLoop
			}
		}
		fmt.Println("Stdin subscription closed")
	}()

	// Doing the streaming
	go func() {
		buffer := make([]byte, 512)
		for {
			n, err := stdoutReader.Read(buffer)
			if err != nil {
				if err != io.EOF {
					stdoutReader.Close()
					fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				}
				break
			}

			data := buffer[0:n]
			output <- string(data)

			for i := 0; i < n; i++ {
				buffer[i] = 0
			}
		}
	}()

	cmd.Run()
}
<commit_msg>Remember to close the channel.<commit_after>package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"

	"github.com/garyburd/redigo/redis"
)

// Runner runs the code
type Runner struct {
	Ext     string `json:"ext"`
	Source  string `json:"source"`
	Version string `json:"version"`
}

// Run the code in the container
func (r *Runner) Run(output messages, conn redis.Conn, uuid string) {
	execArgs := []string{"run", "-i", "koderunr", r.Ext, r.Source}
	if r.Version != "" {
		execArgs = append(execArgs, r.Version)
	}

	cmd := exec.Command("docker", execArgs...)

	stdoutReader, stdoutWriter := io.Pipe()
	cmd.Stdout = stdoutWriter
	cmd.Stderr = stdoutWriter

	stdin, err := cmd.StdinPipe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v", err)
	}

	defer stdin.Close()
	defer stdoutWriter.Close()

	go func() {
		psc := redis.PubSubConn{Conn: conn}
		psc.Subscribe(uuid + "#stdin")
		defer psc.Close()

	StdinSubscriptionLoop:
		for {
			switch n := psc.Receive().(type) {
			case redis.Message:
				fmt.Printf("Message: %s %s\n", n.Channel, n.Data)
				stdin.Write(n.Data)
			case error:
				break StdinSubscriptionLoop
			}
		}
		fmt.Println("Stdin subscription closed")
	}()

	// Doing the streaming
	go func() {
		buffer := make([]byte, 512)
		for {
			n, err := stdoutReader.Read(buffer)
			if err != nil {
				if err != io.EOF {
					stdoutReader.Close()
					fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				}

				close(output)
				break
			}

			data := buffer[0:n]
			output <- string(data)

			for i := 0; i < n; i++ {
				buffer[i] = 0
			}
		}
	}()

	cmd.Run()
}
<|endoftext|>


<commit_before>package server

import (
	"io/ioutil"
	"net/http"
	"path/filepath"
	"screen-server/log"

	"github.com/ant0ine/go-json-rest/rest"
	"gopkg.in/yaml.v2"
)

type Config struct {
	DB         map[string][]string `yaml:"db,omitempty"`
	ServiceUrl []string            `yaml:"configserivce"`
}

var Settings Config

func check(e error) {
	if e != nil {
		logger.Log.Error("Server err: %v", e)
	}
}

func loadConfig() {
	filename, _ := filepath.Abs("./config.yaml")

	yamlFile, err := ioutil.ReadFile(filename)
	check(err)

	err = yaml.Unmarshal(yamlFile, &Settings)
	check(err)
}

//Open makes screen_server open
func Open() {
	logger.Init()
	logger.Log.Debug("Start Screen Server")
	logger.Log.Info("loading local config")
	loadConfig()
	logger.Log.Info("loaded")
	logger.Log.Info("--------------------")

	logger.Log.Info("connect mongo db")
	Mongo.InitDB()
	logger.Log.Info("connected")
	logger.Log.Info("--------------------")

	logger.Log.Info("loading resource&layouts")
	if Mongo.DB != nil {
		Resources.Load(Mongo.DB)
		Layouts.Load(Mongo.DB)
		logger.Log.Info("loaded")
	} else {
		logger.Log.Error("connect to mongodb failed")
	}
	logger.Log.Info("--------------------")

	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	router, err := rest.MakeRouter(
		&rest.Route{"GET", "/resource", getResourceHandler},
		&rest.Route{"POST", "/resource", setResourceHandler},
		&rest.Route{"GET", "/layouts", getLayoutsHandler},
		&rest.Route{"POST", "/layout", setLayoutHandler},
		&rest.Route{"POST", "/layout/current", setCurrentLayoutHandler},
		&rest.Route{"PATCH", "/layout/current", updateCurrentLayoutHandler},
		&rest.Route{"POST", "/layout/:id/resource", updateLayoutResourceHandler},
		&rest.Route{"POST", "/error", notifyErrorHandler},
		&rest.Route{"POST", "/control/heartbeat", setHeartbeatHandler},
		&rest.Route{"POST", "/display/heartbeat", setHeartbeatHandler},
	)

	check(err)

	api.SetApp(router)

	//Mock()

	server = newSocketServer()

	// ticker := time.NewTicker(time.Second * 30)
	//
	// go func() {
	// 	for t := range ticker.C {
	// 		logger.Log.Debug("time:", t)
	// 		logger.Log.Debug("clients:", clients.Map)
	// 	}
	// }()

	http.Handle("/", api.MakeHandler())
	http.Handle("/socket.io/", server)
	err = http.ListenAndServe(":8080", nil)
	check(err)

	logger.Log.Info("Start Screen Server")
}
<commit_msg>update<commit_after>package server

import (
	"io/ioutil"
	"net/http"
	"path/filepath"
	"screen-server/log"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"gopkg.in/yaml.v2"
)

type Config struct {
	DB         map[string][]string `yaml:"db,omitempty"`
	ServiceUrl []string            `yaml:"configserivce"`
}

var Settings Config

func check(e error) {
	if e != nil {
		logger.Log.Error("Server err: %v", e)
	}
}

func loadConfig() {
	filename, _ := filepath.Abs("./config.yaml")

	yamlFile, err := ioutil.ReadFile(filename)
	check(err)

	err = yaml.Unmarshal(yamlFile, &Settings)
	check(err)
}

//Open makes screen_server open
func Open() {
	logger.Init()
	logger.Log.Debug("Start Screen Server")
	logger.Log.Info("loading local config")
	loadConfig()
	logger.Log.Info("loaded")
	logger.Log.Info("--------------------")

	logger.Log.Info("connect mongo db")
	Mongo.InitDB()
	logger.Log.Info("connected")
	logger.Log.Info("--------------------")

	logger.Log.Info("loading resource&layouts")
	if Mongo.DB != nil {
		Resources.Load(Mongo.DB)
		Layouts.Load(Mongo.DB)
		logger.Log.Info("loaded")
	} else {
		logger.Log.Error("connect to mongodb failed")
	}
	logger.Log.Info("--------------------")

	api := rest.NewApi()
	api.Use(rest.DefaultDevStack...)
	router, err := rest.MakeRouter(
		&rest.Route{"GET", "/resource", getResourceHandler},
		&rest.Route{"POST", "/resource", setResourceHandler},
		&rest.Route{"GET", "/layouts", getLayoutsHandler},
		&rest.Route{"POST", "/layout", setLayoutHandler},
		&rest.Route{"POST", "/layout/current", setCurrentLayoutHandler},
		&rest.Route{"PATCH", "/layout/current", updateCurrentLayoutHandler},
		&rest.Route{"POST", "/layout/:id/resource", updateLayoutResourceHandler},
		&rest.Route{"POST", "/error", notifyErrorHandler},
		&rest.Route{"POST", "/control/heartbeat", setHeartbeatHandler},
		&rest.Route{"POST", "/display/heartbeat", setHeartbeatHandler},
	)

	check(err)

	api.SetApp(router)

	//Mock()

	server = newSocketServer()

	ticker := time.NewTicker(time.Second * 30)

	go func() {
		for t := range ticker.C {
			logger.Log.Debug("time:", t)
			logger.Log.Debug("clients:", clients.Map)
		}
	}()

	http.Handle("/", api.MakeHandler())
	http.Handle("/socket.io/", server)
	err = http.ListenAndServe(":8080", nil)
	check(err)

	logger.Log.Info("Start Screen Server")
}
<|endoftext|>


<commit_before>package server

import (
	"math/rand"
	"net"
	"net/http"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt"
	"github.com/influxdata/chronograf/canned"
	"github.com/influxdata/chronograf/influx"
	"github.com/influxdata/chronograf/layouts"
	clog "github.com/influxdata/chronograf/log"
	"github.com/influxdata/chronograf/uuid"
	client "github.com/influxdata/usage-client/v1"
	"github.com/tylerb/graceful"
)

var (
	startTime time.Time
	basepath  string
)

func init() {
	startTime = time.Now().UTC()
}

// Server for the chronograf API
type Server struct {
	Host string `long:"host" description:"the IP to listen on" default:"0.0.0.0" env:"HOST"`
	Port int    `long:"port" description:"the port to listen on for insecure connections, defaults to a random value" default:"8888" env:"PORT"`

	/* TODO: add in support for TLS
	TLSHost           string         `long:"tls-host" description:"the IP to listen on for tls, when not specified it's the same as --host" env:"TLS_HOST"`
	TLSPort           int            `long:"tls-port" description:"the port to listen on for secure connections, defaults to a random value" env:"TLS_PORT"`
	TLSCertificate    flags.Filename `long:"tls-certificate" description:"the certificate to use for secure connections" env:"TLS_CERTIFICATE"`
	TLSCertificateKey flags.Filename `long:"tls-key" description:"the private key to use for secure conections" env:"TLS_PRIVATE_KEY"`
	*/

	Develop            bool     `short:"d" long:"develop" description:"Run server in develop mode."`
	BoltPath           string   `short:"b" long:"bolt-path" description:"Full path to boltDB file (/var/lib/chronograf/chronograf-v1.db)" env:"BOLT_PATH" default:"chronograf-v1.db"`
	CannedPath         string   `short:"c" long:"canned-path" description:"Path to directory of pre-canned application layouts (/usr/share/chronograf/canned)" env:"CANNED_PATH" default:"canned"`
	TokenSecret        string   `short:"t" long:"token-secret" description:"Secret to sign tokens" env:"TOKEN_SECRET"`
	GithubClientID     string   `short:"i" long:"github-client-id" description:"Github Client ID for OAuth 2 support" env:"GH_CLIENT_ID"`
	GithubClientSecret string   `short:"s" long:"github-client-secret" description:"Github Client Secret for OAuth 2 support" env:"GH_CLIENT_SECRET"`
	GithubOrgs         []string `short:"o" long:"github-organization" description:"Github organization user is required to have active membership" env:"GH_ORGS" env-delim:","`
	ReportingDisabled  bool     `short:"r" long:"reporting-disabled" description:"Disable reporting of usage stats (os,arch,version,cluster_id,uptime) once every 24hr" env:"REPORTING_DISABLED"`
	LogLevel           string   `short:"l" long:"log-level" value-name:"choice" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic" default:"info" description:"Set the logging level" env:"LOG_LEVEL"`
	ShowVersion        bool     `short:"v" long:"version" description:"Show Chronograf version info"`
	BuildInfo          BuildInfo
	Listener           net.Listener
	Basepath           string       `long:"basepath" description:"A URL path prefix under which all chronograf routes will be mounted" json:"basePath"`
	handler            http.Handler
}

// BuildInfo is sent to the usage client to track versions and commits
type BuildInfo struct {
	Version string
	Commit  string
}

func (s *Server) useAuth() bool {
	return s.TokenSecret != "" && s.GithubClientID != "" && s.GithubClientSecret != ""
}

// Serve starts and runs the chronograf server
func (s *Server) Serve() error {
	logger := clog.New(clog.ParseLevel(s.LogLevel))
	service := openService(s.BoltPath, s.CannedPath, logger, s.useAuth())
	basepath = s.Basepath
	s.handler = NewMux(MuxOpts{
		Develop:            s.Develop,
		TokenSecret:        s.TokenSecret,
		GithubClientID:     s.GithubClientID,
		GithubClientSecret: s.GithubClientSecret,
		GithubOrgs:         s.GithubOrgs,
		Logger:             logger,
		UseAuth:            s.useAuth(),
	}, service)

	s.handler = Version(s.BuildInfo.Version, s.handler)

	var err error
	s.Listener, err = net.Listen("tcp", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	httpServer := &graceful.Server{Server: new(http.Server)}
	httpServer.SetKeepAlivesEnabled(true)
	httpServer.TCPKeepAlive = 5 * time.Second
	httpServer.Handler = s.handler

	if !s.ReportingDisabled {
		go reportUsageStats(s.BuildInfo, logger)
	}

	logger.
		WithField("component", "server").
		Info("Serving chronograf at http://", s.Listener.Addr())

	if err := httpServer.Serve(s.Listener); err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	logger.
		WithField("component", "server").
		Info("Stopped serving chronograf at http://", s.Listener.Addr())

	return nil
}

func openService(boltPath, cannedPath string, logger chronograf.Logger, useAuth bool) Service {
	db := bolt.NewClient()
	db.Path = boltPath
	if err := db.Open(); err != nil {
		logger.
			WithField("component", "boltstore").
			Fatal("Unable to open boltdb; is there a chronograf already running?  ", err)
	}

	// These apps are those handled from a directory
	apps := canned.NewApps(cannedPath, &uuid.V4{}, logger)
	// These apps are statically compiled into chronograf
	binApps := &canned.BinLayoutStore{
		Logger: logger,
	}

	// Acts as a front-end to both the bolt layouts, filesystem layouts and binary statically compiled layouts.
	// The idea here is that these stores form a hierarchy in which each is tried sequentially until
	// the operation has success.  So, the database is preferred over filesystem over binary data.
	layouts := &layouts.MultiLayoutStore{
		Stores: []chronograf.LayoutStore{
			db.LayoutStore,
			apps,
			binApps,
		},
	}

	return Service{
		ExplorationStore: db.ExplorationStore,
		SourcesStore:     db.SourcesStore,
		ServersStore:     db.ServersStore,
		UsersStore:       db.UsersStore,
		TimeSeries: &influx.Client{
			Logger: logger,
		},
		LayoutStore:     layouts,
		DashboardsStore: db.DashboardsStore,
		AlertRulesStore: db.AlertsStore,
		Logger:          logger,
		UseAuth:         useAuth,
	}
}

// reportUsageStats starts periodic server reporting.
func reportUsageStats(bi BuildInfo, logger chronograf.Logger) {
	rand.Seed(time.Now().UTC().UnixNano())
	serverID := strconv.FormatUint(uint64(rand.Int63()), 10)
	reporter := client.New("")
	u := &client.Usage{
		Product: "chronograf-ng",
		Data: []client.UsageData{
			{
				Values: client.Values{
					"os":         runtime.GOOS,
					"arch":       runtime.GOARCH,
					"version":    bi.Version,
					"cluster_id": serverID,
					"uptime":     time.Since(startTime).Seconds(),
				},
			},
		},
	}
	l := logger.WithField("component", "usage").
		WithField("reporting_addr", reporter.URL).
		WithField("freq", "24h").
		WithField("stats", "os,arch,version,cluster_id,uptime")
	l.Info("Reporting usage stats")
	_, _ = reporter.Save(u)

	ticker := time.NewTicker(24 * time.Hour)
	defer ticker.Stop()
	for {
		<-ticker.C
		l.Debug("Reporting usage stats")
		go reporter.Save(u)
	}
}
<commit_msg>Move Basepath prop to better match style<commit_after>package server

import (
	"math/rand"
	"net"
	"net/http"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt"
	"github.com/influxdata/chronograf/canned"
	"github.com/influxdata/chronograf/influx"
	"github.com/influxdata/chronograf/layouts"
	clog "github.com/influxdata/chronograf/log"
	"github.com/influxdata/chronograf/uuid"
	client "github.com/influxdata/usage-client/v1"
	"github.com/tylerb/graceful"
)

var (
	startTime time.Time
	basepath  string
)

func init() {
	startTime = time.Now().UTC()
}

// Server for the chronograf API
type Server struct {
	Host string `long:"host" description:"the IP to listen on" default:"0.0.0.0" env:"HOST"`
	Port int    `long:"port" description:"the port to listen on for insecure connections, defaults to a random value" default:"8888" env:"PORT"`

	/* TODO: add in support for TLS
	TLSHost           string         `long:"tls-host" description:"the IP to listen on for tls, when not specified it's the same as --host" env:"TLS_HOST"`
	TLSPort           int            `long:"tls-port" description:"the port to listen on for secure connections, defaults to a random value" env:"TLS_PORT"`
	TLSCertificate    flags.Filename `long:"tls-certificate" description:"the certificate to use for secure connections" env:"TLS_CERTIFICATE"`
	TLSCertificateKey flags.Filename `long:"tls-key" description:"the private key to use for secure conections" env:"TLS_PRIVATE_KEY"`
	*/

	Develop            bool     `short:"d" long:"develop" description:"Run server in develop mode."`
	BoltPath           string   `short:"b" long:"bolt-path" description:"Full path to boltDB file (/var/lib/chronograf/chronograf-v1.db)" env:"BOLT_PATH" default:"chronograf-v1.db"`
	CannedPath         string   `short:"c" long:"canned-path" description:"Path to directory of pre-canned application layouts (/usr/share/chronograf/canned)" env:"CANNED_PATH" default:"canned"`
	TokenSecret        string   `short:"t" long:"token-secret" description:"Secret to sign tokens" env:"TOKEN_SECRET"`
	GithubClientID     string   `short:"i" long:"github-client-id" description:"Github Client ID for OAuth 2 support" env:"GH_CLIENT_ID"`
	GithubClientSecret string   `short:"s" long:"github-client-secret" description:"Github Client Secret for OAuth 2 support" env:"GH_CLIENT_SECRET"`
	GithubOrgs         []string `short:"o" long:"github-organization" description:"Github organization user is required to have active membership" env:"GH_ORGS" env-delim:","`
	ReportingDisabled  bool     `short:"r" long:"reporting-disabled" description:"Disable reporting of usage stats (os,arch,version,cluster_id,uptime) once every 24hr" env:"REPORTING_DISABLED"`
	LogLevel           string   `short:"l" long:"log-level" value-name:"choice" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic" default:"info" description:"Set the logging level" env:"LOG_LEVEL"`
	ShowVersion        bool     `short:"v" long:"version" description:"Show Chronograf version info"`
	Basepath           string   `long:"basepath" description:"A URL path prefix under which all chronograf routes will be mounted"`
	BuildInfo          BuildInfo
	Listener           net.Listener
	handler            http.Handler
}

// BuildInfo is sent to the usage client to track versions and commits
type BuildInfo struct {
	Version string
	Commit  string
}

func (s *Server) useAuth() bool {
	return s.TokenSecret != "" && s.GithubClientID != "" && s.GithubClientSecret != ""
}

// Serve starts and runs the chronograf server
func (s *Server) Serve() error {
	logger := clog.New(clog.ParseLevel(s.LogLevel))
	service := openService(s.BoltPath, s.CannedPath, logger, s.useAuth())
	basepath = s.Basepath
	s.handler = NewMux(MuxOpts{
		Develop:            s.Develop,
		TokenSecret:        s.TokenSecret,
		GithubClientID:     s.GithubClientID,
		GithubClientSecret: s.GithubClientSecret,
		GithubOrgs:         s.GithubOrgs,
		Logger:             logger,
		UseAuth:            s.useAuth(),
	}, service)

	s.handler = Version(s.BuildInfo.Version, s.handler)

	var err error
	s.Listener, err = net.Listen("tcp", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	httpServer := &graceful.Server{Server: new(http.Server)}
	httpServer.SetKeepAlivesEnabled(true)
	httpServer.TCPKeepAlive = 5 * time.Second
	httpServer.Handler = s.handler

	if !s.ReportingDisabled {
		go reportUsageStats(s.BuildInfo, logger)
	}

	logger.
		WithField("component", "server").
		Info("Serving chronograf at http://", s.Listener.Addr())

	if err := httpServer.Serve(s.Listener); err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	logger.
		WithField("component", "server").
		Info("Stopped serving chronograf at http://", s.Listener.Addr())

	return nil
}

func openService(boltPath, cannedPath string, logger chronograf.Logger, useAuth bool) Service {
	db := bolt.NewClient()
	db.Path = boltPath
	if err := db.Open(); err != nil {
		logger.
			WithField("component", "boltstore").
			Fatal("Unable to open boltdb; is there a chronograf already running?  ", err)
	}

	// These apps are those handled from a directory
	apps := canned.NewApps(cannedPath, &uuid.V4{}, logger)
	// These apps are statically compiled into chronograf
	binApps := &canned.BinLayoutStore{
		Logger: logger,
	}

	// Acts as a front-end to both the bolt layouts, filesystem layouts and binary statically compiled layouts.
	// The idea here is that these stores form a hierarchy in which each is tried sequentially until
	// the operation has success.  So, the database is preferred over filesystem over binary data.
	layouts := &layouts.MultiLayoutStore{
		Stores: []chronograf.LayoutStore{
			db.LayoutStore,
			apps,
			binApps,
		},
	}

	return Service{
		ExplorationStore: db.ExplorationStore,
		SourcesStore:     db.SourcesStore,
		ServersStore:     db.ServersStore,
		UsersStore:       db.UsersStore,
		TimeSeries: &influx.Client{
			Logger: logger,
		},
		LayoutStore:     layouts,
		DashboardsStore: db.DashboardsStore,
		AlertRulesStore: db.AlertsStore,
		Logger:          logger,
		UseAuth:         useAuth,
	}
}

// reportUsageStats starts periodic server reporting.
func reportUsageStats(bi BuildInfo, logger chronograf.Logger) {
	rand.Seed(time.Now().UTC().UnixNano())
	serverID := strconv.FormatUint(uint64(rand.Int63()), 10)
	reporter := client.New("")
	u := &client.Usage{
		Product: "chronograf-ng",
		Data: []client.UsageData{
			{
				Values: client.Values{
					"os":         runtime.GOOS,
					"arch":       runtime.GOARCH,
					"version":    bi.Version,
					"cluster_id": serverID,
					"uptime":     time.Since(startTime).Seconds(),
				},
			},
		},
	}
	l := logger.WithField("component", "usage").
		WithField("reporting_addr", reporter.URL).
		WithField("freq", "24h").
		WithField("stats", "os,arch,version,cluster_id,uptime")
	l.Info("Reporting usage stats")
	_, _ = reporter.Save(u)

	ticker := time.NewTicker(24 * time.Hour)
	defer ticker.Stop()
	for {
		<-ticker.C
		l.Debug("Reporting usage stats")
		go reporter.Save(u)
	}
}
<|endoftext|>


<commit_before>//
// Copyright (C) 2014 Sebastian 'tokkee' Harl <sh@tokkee.org>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
// TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package server implements the core of the SysDB web server.
package server

import (
	"bytes"
	"errors"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"path/filepath"
	"strings"

	"github.com/sysdb/go/client"
	"github.com/sysdb/go/proto"
	"github.com/sysdb/go/sysdb"
)

// A Config specifies configuration values for a SysDB web server.
type Config struct {
	// Conns is a slice of connections to a SysDB server instance. The number of
	// elements specifies the maximum number of parallel queries to the backend.
	// Note that a client connection is not thread-safe but multiple idle
	// connections don't impose any load on the server.
	Conns []*client.Conn

	// TemplatePath specifies the relative or absolute location of template files.
	TemplatePath string

	// StaticPath specifies the relative or absolute location of static files.
	StaticPath string
}

// A Server implements an http.Handler that serves the SysDB user interface.
type Server struct {
	conns chan *client.Conn

	// Request multiplexer
	mux map[string]handler

	// Templates:
	main    *template.Template
	results map[string]*template.Template

	// Base directory of static files.
	basedir string
}

// New constructs a new SysDB web server using the specified configuration.
func New(cfg Config) (*Server, error) {
	if len(cfg.Conns) == 0 {
		return nil, errors.New("need at least one client connection")
	}

	s := &Server{
		conns:   make(chan *client.Conn, len(cfg.Conns)),
		results: make(map[string]*template.Template),
	}
	for _, c := range cfg.Conns {
		s.conns <- c
	}

	var err error
	s.main, err = cfg.parse("main.tmpl")
	if err != nil {
		return nil, err
	}

	types := []string{"host", "hosts", "service", "services", "metric", "metrics"}
	for _, t := range types {
		s.results[t], err = cfg.parse(t + ".tmpl")
		if err != nil {
			return nil, err
		}
	}

	s.basedir = cfg.StaticPath
	s.mux = map[string]handler{
		"images": s.static,
		"style":  s.static,
		"graph":  s.graph,
	}
	return s, nil
}

func (cfg Config) parse(name string) (*template.Template, error) {
	t := template.New(filepath.Base(name))
	return t.ParseFiles(filepath.Join(cfg.TemplatePath, name))
}

type request struct {
	r    *http.Request
	cmd  string
	args []string
}

type handler func(http.ResponseWriter, request)

// Content generators for HTML pages.
var content = map[string]func(request, *Server) (template.HTML, error){
	"": index,

	// Queries
	"host":     fetch,
	"service":  fetch,
	"metric":   fetch,
	"hosts":    listAll,
	"services": listAll,
	"metrics":  listAll,
	"lookup":   lookup,
}

// ServeHTTP implements the http.Handler interface and serves
// the SysDB user interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	if len(path) > 0 && path[0] == '/' {
		path = path[1:]
	}
	fields := strings.Split(path, "/")

	req := request{
		r:   r,
		cmd: fields[0],
	}
	if len(fields) > 1 {
		if fields[len(fields)-1] == "" {
			// Slash at the end of the URL
			fields = fields[:len(fields)-1]
		}
		if len(fields) > 1 {
			req.args = fields[1:]
		}
	}

	if h := s.mux[fields[0]]; h != nil {
		h(w, req)
		return
	}

	f, ok := content[req.cmd]
	if !ok {
		s.notfound(w, r)
		return
	}
	r.ParseForm()
	content, err := f(req, s)
	if err != nil {
		s.err(w, http.StatusBadRequest, fmt.Errorf("Error: %v", err))
		return
	}

	page := struct {
		Title   string
		Query   string
		Content template.HTML
	}{
		Title:   "SysDB - The System Database",
		Query:   r.FormValue("query"),
		Content: content,
	}

	var buf bytes.Buffer
	err = s.main.Execute(&buf, &page)
	if err != nil {
		s.internal(w, err)
		return
	}

	w.WriteHeader(http.StatusOK)
	io.Copy(w, &buf)
}

// static serves static content.
func (s *Server) static(w http.ResponseWriter, req request) {
	http.ServeFile(w, req.r, filepath.Clean(filepath.Join(s.basedir, req.r.URL.Path)))
}

// Content handlers.

func index(_ request, s *Server) (template.HTML, error) {
	return "

Welcome to the System Database.

", nil
}

func listAll(req request, s *Server) (template.HTML, error) {
	if len(req.args) != 0 {
		return "", fmt.Errorf("%s not found", strings.Title(req.cmd))
	}

	res, err := s.query(fmt.Sprintf("LIST %s", req.cmd))
	if err != nil {
		return "", err
	}
	// the template *must* exist
	return tmpl(s.results[req.cmd], res)
}

func lookup(req request, s *Server) (template.HTML, error) {
	if req.r.Method != "POST" {
		return "", errors.New("Method not allowed")
	}
	q := proto.EscapeString(req.r.FormValue("query"))
	if q == "''" {
		return "", errors.New("Empty query")
	}

	res, err := s.query(fmt.Sprintf("LOOKUP hosts MATCHING name =~ %s", q))
	if err != nil {
		return "", err
	}
	return tmpl(s.results["hosts"], res)
}

func fetch(req request, s *Server) (template.HTML, error) {
	if len(req.args) == 0 {
		return "", fmt.Errorf("%s not found", strings.Title(req.cmd))
	}

	var q string
	switch req.cmd {
	case "host":
		if len(req.args) != 1 {
			return "", fmt.Errorf("%s not found", strings.Title(req.cmd))
		}
		q = fmt.Sprintf("FETCH host %s", proto.EscapeString(req.args[0]))
	case "service", "metric":
		if len(req.args) < 2 {
			return "", fmt.Errorf("%s not found", strings.Title(req.cmd))
		}
		host := proto.EscapeString(req.args[0])
		name := proto.EscapeString(strings.Join(req.args[1:], "/"))
		q = fmt.Sprintf("FETCH %s %s.%s", req.cmd, host, name)
	default:
		panic("Unknown request: fetch(" + req.cmd + ")")
	}

	res, err := s.query(q)
	if err != nil {
		return "", err
	}
	return tmpl(s.results[req.cmd], res)
}

func tmpl(t *template.Template, data interface{}) (template.HTML, error) {
	var buf bytes.Buffer
	if err := t.Execute(&buf, data); err != nil {
		return "", fmt.Errorf("Template error: %v", err)
	}
	return template.HTML(buf.String()), nil
}

func html(s string) template.HTML {
	return template.HTML(template.HTMLEscapeString(s))
}

func (s *Server) query(cmd string) (interface{}, error) {
	c := <-s.conns
	defer func() { s.conns <- c }()

	m := &proto.Message{
		Type: proto.ConnectionQuery,
		Raw:  []byte(cmd),
	}
	if err := c.Send(m); err != nil {
		return nil, fmt.Errorf("Query %q: %v", cmd, err)
	}

	for {
		m, err := c.Receive()
		if err != nil {
			return nil, fmt.Errorf("Failed to receive server response: %v", err)
		}
		if m.Type == proto.ConnectionLog {
			log.Println(string(m.Raw[4:]))
			continue
		} else if m.Type == proto.ConnectionError {
			return nil, errors.New(string(m.Raw))
		}

		t, err := m.DataType()
		if err != nil {
			return nil, fmt.Errorf("Failed to unmarshal response: %v", err)
		}

		var res interface{}
		switch t {
		case proto.HostList:
			var hosts []sysdb.Host
			err = proto.Unmarshal(m, &hosts)
			res = hosts
		case proto.Host:
			var host sysdb.Host
			err = proto.Unmarshal(m, &host)
			res = host
		case proto.Timeseries:
			var ts sysdb.Timeseries
			err = proto.Unmarshal(m, &ts)
			res = ts
		default:
			return nil, fmt.Errorf("Unsupported data type %d", t)
		}
		if err != nil {
			return nil, fmt.Errorf("Failed to unmarshal response: %v", err)
		}
		return res, nil
	}
}

// vim: set tw=78 sw=4 sw=4 noexpandtab :
<commit_msg>Give content generators more control over the generated page.<commit_after>//
// Copyright (C) 2014 Sebastian 'tokkee' Harl <sh@tokkee.org>
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
// TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

// Package server implements the core of the SysDB web server.
package server

import (
	"bytes"
	"errors"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"path/filepath"
	"strings"

	"github.com/sysdb/go/client"
	"github.com/sysdb/go/proto"
	"github.com/sysdb/go/sysdb"
)

// A Config specifies configuration values for a SysDB web server.
type Config struct {
	// Conns is a slice of connections to a SysDB server instance. The number of
	// elements specifies the maximum number of parallel queries to the backend.
	// Note that a client connection is not thread-safe but multiple idle
	// connections don't impose any load on the server.
	Conns []*client.Conn

	// TemplatePath specifies the relative or absolute location of template files.
	TemplatePath string

	// StaticPath specifies the relative or absolute location of static files.
	StaticPath string
}

// A Server implements an http.Handler that serves the SysDB user interface.
type Server struct {
	conns chan *client.Conn

	// Request multiplexer
	mux map[string]handler

	// Templates:
	main    *template.Template
	results map[string]*template.Template

	// Base directory of static files.
	basedir string
}

// New constructs a new SysDB web server using the specified configuration.
func New(cfg Config) (*Server, error) {
	if len(cfg.Conns) == 0 {
		return nil, errors.New("need at least one client connection")
	}

	s := &Server{
		conns:   make(chan *client.Conn, len(cfg.Conns)),
		results: make(map[string]*template.Template),
	}
	for _, c := range cfg.Conns {
		s.conns <- c
	}

	var err error
	s.main, err = cfg.parse("main.tmpl")
	if err != nil {
		return nil, err
	}

	types := []string{"host", "hosts", "service", "services", "metric", "metrics"}
	for _, t := range types {
		s.results[t], err = cfg.parse(t + ".tmpl")
		if err != nil {
			return nil, err
		}
	}

	s.basedir = cfg.StaticPath
	s.mux = map[string]handler{
		"images": s.static,
		"style":  s.static,
		"graph":  s.graph,
	}
	return s, nil
}

func (cfg Config) parse(name string) (*template.Template, error) {
	t := template.New(filepath.Base(name))
	return t.ParseFiles(filepath.Join(cfg.TemplatePath, name))
}

type request struct {
	r    *http.Request
	cmd  string
	args []string
}

type handler func(http.ResponseWriter, request)

type page struct {
	Title   string
	Query   string
	Content template.HTML
}

// Content generators for HTML pages.
var content = map[string]func(request, *Server) (*page, error){
	"": index,

	// Queries
	"host":     fetch,
	"service":  fetch,
	"metric":   fetch,
	"hosts":    listAll,
	"services": listAll,
	"metrics":  listAll,
	"lookup":   lookup,
}

// ServeHTTP implements the http.Handler interface and serves
// the SysDB user interface.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path
	if len(path) > 0 && path[0] == '/' {
		path = path[1:]
	}
	fields := strings.Split(path, "/")

	req := request{
		r:   r,
		cmd: fields[0],
	}
	if len(fields) > 1 {
		if fields[len(fields)-1] == "" {
			// Slash at the end of the URL
			fields = fields[:len(fields)-1]
		}
		if len(fields) > 1 {
			req.args = fields[1:]
		}
	}

	if h := s.mux[fields[0]]; h != nil {
		h(w, req)
		return
	}

	f, ok := content[req.cmd]
	if !ok {
		s.notfound(w, r)
		return
	}
	r.ParseForm()
	page, err := f(req, s)
	if err != nil {
		s.err(w, http.StatusBadRequest, fmt.Errorf("Error: %v", err))
		return
	}

	page.Query = r.FormValue("query")
	if page.Title == "" {
		page.Title = "SysDB - The System Database"
	}

	var buf bytes.Buffer
	err = s.main.Execute(&buf, page)
	if err != nil {
		s.internal(w, err)
		return
	}

	w.WriteHeader(http.StatusOK)
	io.Copy(w, &buf)
}

// static serves static content.
func (s *Server) static(w http.ResponseWriter, req request) {
	http.ServeFile(w, req.r, filepath.Clean(filepath.Join(s.basedir, req.r.URL.Path)))
}

// Content handlers.

func index(_ request, s *Server) (*page, error) {
	return &page{Content: "

Welcome to the System Database.

"}, nil
}

func listAll(req request, s *Server) (*page, error) {
	if len(req.args) != 0 {
		return nil, fmt.Errorf("%s not found", strings.Title(req.cmd))
	}

	res, err := s.query(fmt.Sprintf("LIST %s", req.cmd))
	if err != nil {
		return nil, err
	}
	// the template *must* exist
	return tmpl(s.results[req.cmd], res)
}

func lookup(req request, s *Server) (*page, error) {
	if req.r.Method != "POST" {
		return nil, errors.New("Method not allowed")
	}
	q := proto.EscapeString(req.r.FormValue("query"))
	if q == "''" {
		return nil, errors.New("Empty query")
	}

	res, err := s.query(fmt.Sprintf("LOOKUP hosts MATCHING name =~ %s", q))
	if err != nil {
		return nil, err
	}
	return tmpl(s.results["hosts"], res)
}

func fetch(req request, s *Server) (*page, error) {
	if len(req.args) == 0 {
		return nil, fmt.Errorf("%s not found", strings.Title(req.cmd))
	}

	var q string
	switch req.cmd {
	case "host":
		if len(req.args) != 1 {
			return nil, fmt.Errorf("%s not found", strings.Title(req.cmd))
		}
		q = fmt.Sprintf("FETCH host %s", proto.EscapeString(req.args[0]))
	case "service", "metric":
		if len(req.args) < 2 {
			return nil, fmt.Errorf("%s not found", strings.Title(req.cmd))
		}
		host := proto.EscapeString(req.args[0])
		name := proto.EscapeString(strings.Join(req.args[1:], "/"))
		q = fmt.Sprintf("FETCH %s %s.%s", req.cmd, host, name)
	default:
		panic("Unknown request: fetch(" + req.cmd + ")")
	}

	res, err := s.query(q)
	if err != nil {
		return nil, err
	}
	return tmpl(s.results[req.cmd], res)
}

func tmpl(t *template.Template, data interface{}) (*page, error) {
	var buf bytes.Buffer
	if err := t.Execute(&buf, data); err != nil {
		return nil, fmt.Errorf("Template error: %v", err)
	}
	return &page{Content: template.HTML(buf.String())}, nil
}

func html(s string) template.HTML {
	return template.HTML(template.HTMLEscapeString(s))
}

func (s *Server) query(cmd string) (interface{}, error) {
	c := <-s.conns
	defer func() { s.conns <- c }()

	m := &proto.Message{
		Type: proto.ConnectionQuery,
		Raw:  []byte(cmd),
	}
	if err := c.Send(m); err != nil {
		return nil, fmt.Errorf("Query %q: %v", cmd, err)
	}

	for {
		m, err := c.Receive()
		if err != nil {
			return nil, fmt.Errorf("Failed to receive server response: %v", err)
		}
		if m.Type == proto.ConnectionLog {
			log.Println(string(m.Raw[4:]))
			continue
		} else if m.Type == proto.ConnectionError {
			return nil, errors.New(string(m.Raw))
		}

		t, err := m.DataType()
		if err != nil {
			return nil, fmt.Errorf("Failed to unmarshal response: %v", err)
		}

		var res interface{}
		switch t {
		case proto.HostList:
			var hosts []sysdb.Host
			err = proto.Unmarshal(m, &hosts)
			res = hosts
		case proto.Host:
			var host sysdb.Host
			err = proto.Unmarshal(m, &host)
			res = host
		case proto.Timeseries:
			var ts sysdb.Timeseries
			err = proto.Unmarshal(m, &ts)
			res = ts
		default:
			return nil, fmt.Errorf("Unsupported data type %d", t)
		}
		if err != nil {
			return nil, fmt.Errorf("Failed to unmarshal response: %v", err)
		}
		return res, nil
	}
}

// vim: set tw=78 sw=4 sw=4 noexpandtab :
<|endoftext|>


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

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

const nbd_folder = "/sample_disks/"

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

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

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

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

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

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

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

    defer conn.Close()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)
            //fmt.Printf("Done reading the data that should be written")

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

            file.Sync()

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

            conn.Write(buffer[:16])

            continue

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

            file.Sync()
            return
        }
    }
}

func send_export_list(output *bufio.Writer, options uint32) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(err)
    }
    for _, file := range files {
        send_export_list_item(output, options, file.Name())
    }

    send_ack(output, options)
}

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

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

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

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

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

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

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

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

var defaultOptions = []byte{0, 0}

func main() {

    app := cli.NewApp()
    app.Name = "AnyBlox"
    app.Usage = "block storage for the masses"
    app.Action = func(c *cli.Context) error {
        fmt.Println("boom!")
        return nil
    }

    app.Run(os.Args)

    if len(os.Args) <  3 {
        panic("missing arguments:  (ipaddress) (portnumber)")
        return
    }

    listener, err := net.Listen("tcp", os.Args[1] + ":" + os.Args[2])
    utils.ErrorCheck(err)

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

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

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

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

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

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

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

        output.Flush()

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

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

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

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

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

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

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

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

}
<commit_msg>adding basic flag support including host, port, and listen (host:port)<commit_after>// Copyright 2016 Jacob Taylor jacob@ablox.io
// License: Apache2 - http://www.apache.org/licenses/LICENSE-2.0
package main

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

const nbd_folder = "/sample_disks/"

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

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

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

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

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

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

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

    defer conn.Close()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            //waiting_for += int(length)                   // wait for the additional payload
            //fmt.Printf("About to read the data that we should be writing out.")
            _, err := io.ReadFull(conn_reader, buffer[28:28+length])
            if err == io.EOF {
                fmt.Printf("Abort detected, escaping processing loop\n")
                break
            }
            utils.ErrorCheck(err)
            //fmt.Printf("Done reading the data that should be written")

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

            file.Sync()

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

            conn.Write(buffer[:16])

            continue

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

            file.Sync()
            return
        }
    }
}

func send_export_list(output *bufio.Writer, options uint32) {
    current_directory, err := os.Getwd()
    files, err := ioutil.ReadDir(current_directory + nbd_folder)
    if err != nil {
        log.Fatal(err)
    }
    for _, file := range files {
        send_export_list_item(output, options, file.Name())
    }

    send_ack(output, options)
}

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

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

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

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

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

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

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

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

var defaultOptions = []byte{0, 0}

func main() {

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

    host := "localhost"
    port := "8000"
    listen := ""
    app.Flags = []cli.Flag {
        cli.StringFlag{
            Name: "host",
            Value: host,
            Usage: "Hostname or IP address you want to serve traffic on. e.x. 'localhost', '192.168.1.2'",
            Destination: &host,
        },
        cli.StringFlag{
            Name: "port",
            Value: port,
            Usage: "Port you want to serve traffic on. e.x. '8000'",
            Destination: &port,
        },
        cli.StringFlag{
            Name: "listen, l",
            Destination: &listen,
            Usage: "Address and port the server should listen on. Listen will take priority over host and port parameters. hostname:port - e.x. 'localhost:8000', '192.168.1.2:8000'",
        },
    }

    app.Run(os.Args)

    // Determine where the host should be listening to, depending on the arguments
    hostingAddress := listen
    if len(listen) == 0 {
        hostingAddress = host + ":" + port
    }

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

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

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

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

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

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

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

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

        output.Flush()

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

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

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

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

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

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

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

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

}
<|endoftext|>


<commit_before>package server

import (
    "github.com/qmsk/onewire/avrtemp"
    "fmt"
    "github.com/qmsk/onewire/hidraw"
    "log"
    "time"
)

type Device struct {
    hidraw          hidraw.DeviceInfo
    avrtemp         *avrtemp.Device
}

func (d Device) String() string {
    return d.hidraw.String()
}

type Stat struct {
    // set by Device.reader()
    Device          *Device
    ID              avrtemp.ID
    Time            time.Time
    Temperature     avrtemp.Temperature

    // set by Server.run()
    SensorConfig    *SensorConfig
}

func (stat Stat) String() string {
    return fmt.Sprintf("%v", stat.ID)
}

type Server struct {
    // XXX: these maps are all racy
    config          Config
    sensorConfig    map[string]*SensorConfig

    devices         map[string]*Device
    stats           map[string]Stat

    deviceChan          chan Device     // add/remove Device
    statChan            chan Stat       // read in from Devices
    influxChan          chan Stat       // write out to influx

    apiStatusChan       chan chan APIStatus
    apiStatChan         chan chan APIStat
}

func New() (*Server, error) {
    server := &Server{
        sensorConfig:   make(map[string]*SensorConfig),

        devices:        make(map[string]*Device),
        stats:          make(map[string]Stat),

        deviceChan:     make(chan Device),
        statChan:       make(chan Stat),

        apiStatusChan:  make(chan chan APIStatus),
        apiStatChan:    make(chan chan APIStat),
    }

    go server.run()

    return server, nil
}

func (s *Server) apiStatus(statusChan chan APIStatus) {
    defer close(statusChan)

    for name, device := range s.devices {
        status := APIStatus{
            Name:           name,
            HidrawDevice:   device.hidraw,
            Stats:          make(map[string]string),
        }

        if device.avrtemp != nil {
            status.AvrtempDevice = device.avrtemp.Status()
        }

        for statID, stat := range s.stats {
            if stat.Device != device {
                continue
            } else {
                // nil-safe
                status.Stats[statID] = s.sensorConfig[statID].String()
            }
        }

        statusChan <- status
    }
}

func (s *Server) apiStat(statChan chan APIStat) {
    defer close(statChan)

    for id, stat := range s.stats {
        apiStat := APIStat{
            ID:             id,
            Family:         stat.ID.Family(),
            Time:           stat.Time,
            Temperature:    stat.Temperature.Float64(),
        }

        if stat.SensorConfig != nil {
            apiStat.SensorName = stat.SensorConfig.String()
        }

        statChan <- apiStat
    }
}

func (s *Server) run() {
    var stat Stat

    for {
        select {
        case device := <-s.deviceChan:
            if device.avrtemp != nil {
                runningDevice := &device

                log.Printf("server.Server: Start avrtemp device %v\n", runningDevice)

                // run
                s.devices[device.String()] = runningDevice

                go runningDevice.reader(s.statChan)

            } else if runningDevice := s.devices[device.String()]; runningDevice != nil {
                log.Printf("server.Server: Stop device %v\n", runningDevice)

                // shutdown
                if runningDevice.avrtemp != nil {
                    runningDevice.avrtemp.Close()
                }

                delete(s.devices, device.String())
            }

        case stat = <-s.statChan:
            log.Printf("server.Server: Stat %v: %v\n", stat, stat.Temperature)

            stat.SensorConfig = s.sensorConfig[stat.ID.String()]

            s.stats[stat.String()] = stat

            s.influxChan <- stat

        case statusChan := <-s.apiStatusChan:
            s.apiStatus(statusChan)
        case statChan := <-s.apiStatChan:
            s.apiStat(statChan)
        }
    }
}

func (device *Device) reader(statChan chan Stat) {
    for {
        if report, err := device.avrtemp.Read(); err != nil {
            log.Printf("server.Device %v: avrtemp.Device %v: Read: %v\n", device, err)
            break
        } else {
            log.Printf("server.Device %v: avrtemp.Device %v: Read: %v\n", device, report)

            statChan <- Stat{
                Device:         device,
                ID:             report.ID,
                Time:           time.Now(),
                Temperature:    report.Temp,
            }
        }
    }
}

func (s *Server) AddHidrawDevice(deviceInfo hidraw.DeviceInfo) {
    device := Device{
        hidraw: deviceInfo,
    }

    if hidrawDevice, err := hidraw.Open(deviceInfo); err != nil {
        log.Printf("AddHidrawDevice %#v: hidraw.Open: %v\n", deviceInfo, err)
    } else if avrtempDevice, err := avrtemp.Open(hidrawDevice); err != nil {
        log.Printf("AddHidrawDevice %#v: avrtemp.Open: %v\n", deviceInfo, err)
    } else {
        log.Printf("AddHidrawDevice %#v: %#v\n", deviceInfo, avrtempDevice)

        device.avrtemp = avrtempDevice
    }

    s.deviceChan <- device
}

func (s *Server) RemoveHidrawDevice(deviceInfo hidraw.DeviceInfo) {
    log.Printf("RemoveHidrawDevice %v...\n", deviceInfo)

    s.deviceChan <- Device{hidraw: deviceInfo}
}

func (s *Server) MonitorHidraw(monitorChan chan hidraw.MonitorEvent) {
    for monitorEvent := range monitorChan {
        switch monitorEvent.Action {
        case "add":
            s.AddHidrawDevice(monitorEvent.DeviceInfo)
        case "remove":
            s.RemoveHidrawDevice(monitorEvent.DeviceInfo)
        default:
            log.Printf("MonitorHidraw: %v?! %v\n", monitorEvent.Action, monitorEvent.DeviceInfo)
        }
    }

    log.Printf("server.Server %v: MonitorHidraw: exit\n", s)
}
<commit_msg>server: docdoc Server state<commit_after>package server

import (
    "github.com/qmsk/onewire/avrtemp"
    "fmt"
    "github.com/qmsk/onewire/hidraw"
    "log"
    "time"
)

type Device struct {
    hidraw          hidraw.DeviceInfo
    avrtemp         *avrtemp.Device
}

func (d Device) String() string {
    return d.hidraw.String()
}

type Stat struct {
    // set by Device.reader()
    Device          *Device
    ID              avrtemp.ID
    Time            time.Time
    Temperature     avrtemp.Temperature

    // set by Server.run()
    SensorConfig    *SensorConfig
}

func (stat Stat) String() string {
    return fmt.Sprintf("%v", stat.ID)
}

type Server struct {
    // only modified at startup
    config          Config
    sensorConfig    map[string]*SensorConfig

    // state private to run()
    devices         map[string]*Device
    stats           map[string]Stat

    deviceChan          chan Device     // add/remove Device
    statChan            chan Stat       // read in from Devices
    influxChan          chan Stat       // write out to influx

    // HTTP API requests, handled by run()
    apiStatusChan       chan chan APIStatus
    apiStatChan         chan chan APIStat
}

func New() (*Server, error) {
    server := &Server{
        sensorConfig:   make(map[string]*SensorConfig),

        devices:        make(map[string]*Device),
        stats:          make(map[string]Stat),

        deviceChan:     make(chan Device),
        statChan:       make(chan Stat),

        apiStatusChan:  make(chan chan APIStatus),
        apiStatChan:    make(chan chan APIStat),
    }

    go server.run()

    return server, nil
}

func (s *Server) apiStatus(statusChan chan APIStatus) {
    defer close(statusChan)

    for name, device := range s.devices {
        status := APIStatus{
            Name:           name,
            HidrawDevice:   device.hidraw,
            Stats:          make(map[string]string),
        }

        if device.avrtemp != nil {
            status.AvrtempDevice = device.avrtemp.Status()
        }

        for statID, stat := range s.stats {
            if stat.Device != device {
                continue
            } else {
                // nil-safe
                status.Stats[statID] = s.sensorConfig[statID].String()
            }
        }

        statusChan <- status
    }
}

func (s *Server) apiStat(statChan chan APIStat) {
    defer close(statChan)

    for id, stat := range s.stats {
        apiStat := APIStat{
            ID:             id,
            Family:         stat.ID.Family(),
            Time:           stat.Time,
            Temperature:    stat.Temperature.Float64(),
        }

        if stat.SensorConfig != nil {
            apiStat.SensorName = stat.SensorConfig.String()
        }

        statChan <- apiStat
    }
}

func (s *Server) run() {
    var stat Stat

    for {
        select {
        case device := <-s.deviceChan:
            if device.avrtemp != nil {
                runningDevice := &device

                log.Printf("server.Server: Start avrtemp device %v\n", runningDevice)

                // run
                s.devices[device.String()] = runningDevice

                go runningDevice.reader(s.statChan)

            } else if runningDevice := s.devices[device.String()]; runningDevice != nil {
                log.Printf("server.Server: Stop device %v\n", runningDevice)

                // shutdown
                if runningDevice.avrtemp != nil {
                    runningDevice.avrtemp.Close()
                }

                delete(s.devices, device.String())
            }

        case stat = <-s.statChan:
            log.Printf("server.Server: Stat %v: %v\n", stat, stat.Temperature)

            stat.SensorConfig = s.sensorConfig[stat.ID.String()]

            s.stats[stat.String()] = stat

            s.influxChan <- stat

        case statusChan := <-s.apiStatusChan:
            s.apiStatus(statusChan)
        case statChan := <-s.apiStatChan:
            s.apiStat(statChan)
        }
    }
}

func (device *Device) reader(statChan chan Stat) {
    for {
        if report, err := device.avrtemp.Read(); err != nil {
            log.Printf("server.Device %v: avrtemp.Device %v: Read: %v\n", device, err)
            break
        } else {
            log.Printf("server.Device %v: avrtemp.Device %v: Read: %v\n", device, report)

            statChan <- Stat{
                Device:         device,
                ID:             report.ID,
                Time:           time.Now(),
                Temperature:    report.Temp,
            }
        }
    }
}

func (s *Server) AddHidrawDevice(deviceInfo hidraw.DeviceInfo) {
    device := Device{
        hidraw: deviceInfo,
    }

    if hidrawDevice, err := hidraw.Open(deviceInfo); err != nil {
        log.Printf("AddHidrawDevice %#v: hidraw.Open: %v\n", deviceInfo, err)
    } else if avrtempDevice, err := avrtemp.Open(hidrawDevice); err != nil {
        log.Printf("AddHidrawDevice %#v: avrtemp.Open: %v\n", deviceInfo, err)
    } else {
        log.Printf("AddHidrawDevice %#v: %#v\n", deviceInfo, avrtempDevice)

        device.avrtemp = avrtempDevice
    }

    s.deviceChan <- device
}

func (s *Server) RemoveHidrawDevice(deviceInfo hidraw.DeviceInfo) {
    log.Printf("RemoveHidrawDevice %v...\n", deviceInfo)

    s.deviceChan <- Device{hidraw: deviceInfo}
}

func (s *Server) MonitorHidraw(monitorChan chan hidraw.MonitorEvent) {
    for monitorEvent := range monitorChan {
        switch monitorEvent.Action {
        case "add":
            s.AddHidrawDevice(monitorEvent.DeviceInfo)
        case "remove":
            s.RemoveHidrawDevice(monitorEvent.DeviceInfo)
        default:
            log.Printf("MonitorHidraw: %v?! %v\n", monitorEvent.Action, monitorEvent.DeviceInfo)
        }
    }

    log.Printf("server.Server %v: MonitorHidraw: exit\n", s)
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"crypto/tls"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/gabstv/manners"
	"github.com/gabstv/sandpiper/pathtree"
	"github.com/gabstv/sandpiper/route"
	"github.com/gabstv/sandpiper/util"
	"golang.org/x/crypto/acme"
	"golang.org/x/crypto/acme/autocert"
)

// Server is the structure that controls, routes and certificates.
type Server interface {
	Add(r route.Route) error
	Run() error
	Close()
	Init()
	ServeHTTP(w http.ResponseWriter, r *http.Request)
	Routes() map[string]route.Route
	GetConfig() Config
	SetConfig(cfg Config)
}

type sServer struct {
	Cfg         Config
	trieDomains *pathtree.Trie
	domains     map[string]*route.Route
	Logger      *log.Logger
	closeChan   chan os.Signal
	htps        *http.Server
}

func (s *sServer) GetConfig() Config {
	return s.Cfg
}

func (s *sServer) SetConfig(cfg Config) {
	s.Cfg = cfg
}

// Default starts a server with the default configuration options
func Default(cfg *Config) Server {
	s := &sServer{}
	if cfg != nil {
		s.Cfg = *cfg
	}
	s.trieDomains = pathtree.NewTrie(".")
	s.domains = make(map[string]*route.Route, 0)
	s.Logger = log.New(os.Stderr, "[sp server] ", log.LstdFlags)
	return s
}

func (s *sServer) Routes() map[string]route.Route {
	mm := make(map[string]route.Route)
	for k, v := range s.domains {
		mm[k] = *v
	}
	return mm
}

func (s *sServer) startAPI(ctx context.Context) error {
	if s.Cfg.APIListen == "" {
		return nil
	}
	go runAPIV1(ctx, s, s.Cfg.APIListen, s.Cfg.APIKey, s.Cfg.Debug)
	if s.Cfg.APIDomain != "" {
		s.Add(route.Route{
			Autocert: s.Cfg.APIDomainAutocert,
			Domain:   s.Cfg.APIDomain,
			Server: route.RouteServer{
				OutAddress:  s.Cfg.APIListen,
				OutConnType: route.HTTP,
			},
			WsCFG: util.WsConfig{
				Enabled: false,
			},
		})
	}
	return nil
}

func (s *sServer) Add(r route.Route) error {
	rr := &route.Route{}
	*rr = r
	if rr.WsCFG.ReadBufferSize == 0 {
		rr.WsCFG.ReadBufferSize = 2048
	}
	if rr.WsCFG.WriteBufferSize == 0 {
		rr.WsCFG.WriteBufferSize = 2048
	}
	if rr.WsCFG.ReadDeadlineSeconds == 0 {
		rr.WsCFG.ReadDeadlineSeconds = time.Second * 60
	} else {
		if rr.WsCFG.ReadDeadlineSeconds < time.Millisecond {
			rr.WsCFG.ReadDeadlineSeconds = time.Duration(rr.WsCFG.ReadDeadlineSeconds) * time.Second
		}
	}
	err := s.trieDomains.Add(r.Domain, rr)
	if err != nil {
		return err
	}
	s.domains[r.Domain] = rr
	s.updateCertificates()
	return nil
}

func (s *sServer) updateCertificates() *autocert.Manager {
	autocdomains := make([]string, 0)
	for _, v := range s.domains {
		if v.Autocert {
			autocdomains = append(autocdomains, v.Domain)
		}
	}
	if len(autocdomains) < 1 {
		return nil
	}
	var m *autocert.Manager
	cpath := "/tmp/sandpiper"
	if s.Cfg.CachePath != "" {
		cpath = s.Cfg.CachePath
	}
	dcache := autocert.DirCache(cpath)
	m = &autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		HostPolicy: autocert.HostWhitelist(autocdomains...),
		Cache:      &dcache,
	}
	if s.Cfg.LetsEncryptURL != "" {
		m.Client = &acme.Client{DirectoryURL: s.Cfg.LetsEncryptURL}
	}

	certs := make(map[string]tls.Certificate)
	for k, v := range s.domains {
		if !v.Autocert && len(v.Certificate.KeyFile) > 0 && len(v.Certificate.CertFile) > 0 {
			ncert, err := tls.LoadX509KeyPair(v.Certificate.CertFile, v.Certificate.KeyFile)
			if err != nil {
				s.Logger.Println("Error loading certificate for", k, err.Error())
			} else {
				certs[k] = ncert
			}
		}
	}

	getcertfn := func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		if dom, ok := certs[clientHello.ServerName]; ok {
			return &dom, nil
		}
		return m.GetCertificate(clientHello)
	}
	if s.Cfg.LetsEncryptURL == "dev" {
		getcertfn = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if dom, ok := certs[clientHello.ServerName]; ok {
				return &dom, nil
			}
			cccert, err := createCert(clientHello)
			if err != nil {
				return nil, err
			}
			certs[clientHello.ServerName] = cccert
			return &cccert, nil
		}
	}
	if s.htps != nil {
		s.htps.TLSConfig = &tls.Config{GetCertificate: getcertfn}
	}
	return m
}

func (s *sServer) Run() error {
	s.Init()
	errc := make(chan error, 3)
	ctx, cancelf := context.WithCancel(context.Background())

	s.htps = &http.Server{
		Addr: s.Cfg.ListenAddrTLS,
	}
	autocertManager := s.updateCertificates()

	s.htps.Handler = s

	go func() {
		if autocertManager == nil || s.Cfg.DisableTLS {
			s.Logger.Println("Listening HTTP")
			errc <- http.ListenAndServe(s.Cfg.ListenAddr, s)
		} else {
			s.Logger.Println("Listening accepting HTTP requests to the SNI challenge")
			errc <- http.ListenAndServe(s.Cfg.ListenAddr, autocertManager.HTTPHandler(s))
		}
	}()

	certs := make([]util.Certificate, 0, len(s.domains))
	for _, v := range s.domains {
		if v.Certificate.CertFile != "" {
			certs = append(certs, v.Certificate)
		}
	}
	var wrapper *util.ServerWrapper
	if !s.Cfg.DisableTLS {
		go func() {
			if s.Cfg.Graceful {
				s.Logger.Println("Listening HTTPS (Graceful)")
				wrapper = util.NewGracefulServer(manners.NewWithServer(s.htps))
			} else {
				s.Logger.Println("Listening HTTPS (Vanilla)")
				wrapper = util.NewVanillaServer(s.htps)
			}
			errc <- util.ListenAndServeTLSSNI(wrapper, certs)
		}()
	}
	//
	// API
	if err := s.startAPI(ctx); err != nil {
		s.Logger.Println("START API ERROR:", err.Error())
	}
	//
	go func() {
		s.closeChan = make(chan os.Signal, 1)
		<-s.closeChan
		cancelf()
		if wrapper != nil {
			wrapper.Close()
		}
		errc <- nil
	}()
	//
	err := <-errc
	return err
}

func (s *sServer) Close() {
	s.closeChan <- os.Interrupt
}

func (s *sServer) Init() {
	// first start setting the number of cpu cores to use
	ncpu := runtime.NumCPU()
	if s.Cfg.NumCPU > 0 && s.Cfg.NumCPU < ncpu {
		ncpu = s.Cfg.NumCPU
	}
	runtime.GOMAXPROCS(ncpu)
}

func (s *sServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h := r.Host
	if s.Cfg.Debug {
		if ho := r.Header.Get("X-Sandpiper-Host"); ho != "" {
			h = ho
		}
		s.Logger.Println("Host: " + h)
	}
	res := s.trieDomains.Find(h)
	if res == nil {
		if len(s.Cfg.FallbackDomain) > 0 {
			if s.Cfg.Debug {
				s.Logger.Println("FALLBACK DOMAIN", s.Cfg.FallbackDomain)
			}
			dom := s.domains[s.Cfg.FallbackDomain]
			if dom != nil {
				if dom.AuthMode != "" {
					switch dom.AuthMode {
					case "apikey":
						if dom.AuthValue != r.Header.Get(dom.AuthKey) {
							w.WriteHeader(http.StatusUnauthorized)
							return
						}
					}
				}
				dom.ReverseProxy(w, r)
				return
			} else {
				if s.Cfg.Debug {
					s.Logger.Println("FALLBACK DOMAIN NOT FOUND")
				}
				http.Error(w, "fallback domain not found "+h, http.StatusInternalServerError)
				return
			}
		} else {
			if s.Cfg.Debug {
				s.Logger.Println("DOMAIN NOT FOUND")
			}
			http.Error(w, "domain not found "+h, http.StatusInternalServerError)
			return
		}
	}
	if res.EndRoute == nil {
		if s.Cfg.Debug {
			s.Logger.Println("ROUTE IS NULL")
		}
		http.Error(w, "route is null", http.StatusInternalServerError)
		return
	}
	if res.EndRoute.AuthMode != "" {
		switch res.EndRoute.AuthMode {
		case "apikey":
			if res.EndRoute.AuthValue != r.Header.Get(res.EndRoute.AuthKey) {
				w.WriteHeader(http.StatusUnauthorized)
				return
			}
		}
	}
	res.EndRoute.ReverseProxy(w, r)
}
<commit_msg>null check 2<commit_after>package server

import (
	"context"
	"crypto/tls"
	"log"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/gabstv/manners"
	"github.com/gabstv/sandpiper/pathtree"
	"github.com/gabstv/sandpiper/route"
	"github.com/gabstv/sandpiper/util"
	"golang.org/x/crypto/acme"
	"golang.org/x/crypto/acme/autocert"
)

// Server is the structure that controls, routes and certificates.
type Server interface {
	Add(r route.Route) error
	Run() error
	Close()
	Init()
	ServeHTTP(w http.ResponseWriter, r *http.Request)
	Routes() map[string]route.Route
	GetConfig() Config
	SetConfig(cfg Config)
}

type sServer struct {
	Cfg         Config
	trieDomains *pathtree.Trie
	domains     map[string]*route.Route
	Logger      *log.Logger
	closeChan   chan os.Signal
	htps        *http.Server
}

func (s *sServer) GetConfig() Config {
	return s.Cfg
}

func (s *sServer) SetConfig(cfg Config) {
	s.Cfg = cfg
}

// Default starts a server with the default configuration options
func Default(cfg *Config) Server {
	s := &sServer{}
	if cfg != nil {
		s.Cfg = *cfg
	}
	s.trieDomains = pathtree.NewTrie(".")
	s.domains = make(map[string]*route.Route, 0)
	s.Logger = log.New(os.Stderr, "[sp server] ", log.LstdFlags)
	return s
}

func (s *sServer) Routes() map[string]route.Route {
	mm := make(map[string]route.Route)
	for k, v := range s.domains {
		mm[k] = *v
	}
	return mm
}

func (s *sServer) startAPI(ctx context.Context) error {
	if s.Cfg.APIListen == "" {
		return nil
	}
	go runAPIV1(ctx, s, s.Cfg.APIListen, s.Cfg.APIKey, s.Cfg.Debug)
	if s.Cfg.APIDomain != "" {
		s.Add(route.Route{
			Autocert: s.Cfg.APIDomainAutocert,
			Domain:   s.Cfg.APIDomain,
			Server: route.RouteServer{
				OutAddress:  s.Cfg.APIListen,
				OutConnType: route.HTTP,
			},
			WsCFG: util.WsConfig{
				Enabled: false,
			},
		})
	}
	return nil
}

func (s *sServer) Add(r route.Route) error {
	rr := &route.Route{}
	*rr = r
	if rr.WsCFG.ReadBufferSize == 0 {
		rr.WsCFG.ReadBufferSize = 2048
	}
	if rr.WsCFG.WriteBufferSize == 0 {
		rr.WsCFG.WriteBufferSize = 2048
	}
	if rr.WsCFG.ReadDeadlineSeconds == 0 {
		rr.WsCFG.ReadDeadlineSeconds = time.Second * 60
	} else {
		if rr.WsCFG.ReadDeadlineSeconds < time.Millisecond {
			rr.WsCFG.ReadDeadlineSeconds = time.Duration(rr.WsCFG.ReadDeadlineSeconds) * time.Second
		}
	}
	err := s.trieDomains.Add(r.Domain, rr)
	if err != nil {
		return err
	}
	s.domains[r.Domain] = rr
	s.updateCertificates()
	return nil
}

func (s *sServer) updateCertificates() *autocert.Manager {
	autocdomains := make([]string, 0)
	for _, v := range s.domains {
		if v.Autocert {
			autocdomains = append(autocdomains, v.Domain)
		}
	}
	if len(autocdomains) < 1 {
		log.Println("no autocert domains")
		return nil
	}
	var m *autocert.Manager
	cpath := "/tmp/sandpiper"
	if s.Cfg.CachePath != "" {
		cpath = s.Cfg.CachePath
	}
	dcache := autocert.DirCache(cpath)
	m = &autocert.Manager{
		Prompt:     autocert.AcceptTOS,
		HostPolicy: autocert.HostWhitelist(autocdomains...),
		Cache:      &dcache,
	}
	if s.Cfg.LetsEncryptURL != "" {
		m.Client = &acme.Client{DirectoryURL: s.Cfg.LetsEncryptURL}
	}

	certs := make(map[string]tls.Certificate)
	for k, v := range s.domains {
		if !v.Autocert && len(v.Certificate.KeyFile) > 0 && len(v.Certificate.CertFile) > 0 {
			ncert, err := tls.LoadX509KeyPair(v.Certificate.CertFile, v.Certificate.KeyFile)
			if err != nil {
				s.Logger.Println("Error loading certificate for", k, err.Error())
			} else {
				certs[k] = ncert
			}
		}
	}

	getcertfn := func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		if dom, ok := certs[clientHello.ServerName]; ok {
			return &dom, nil
		}
		return m.GetCertificate(clientHello)
	}
	if s.Cfg.LetsEncryptURL == "dev" {
		getcertfn = func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			if dom, ok := certs[clientHello.ServerName]; ok {
				return &dom, nil
			}
			cccert, err := createCert(clientHello)
			if err != nil {
				return nil, err
			}
			certs[clientHello.ServerName] = cccert
			return &cccert, nil
		}
	}
	if s.htps != nil {
		s.htps.TLSConfig = &tls.Config{GetCertificate: getcertfn}
	} else {
		s.Logger.Println("s.htps WAS NIL")
		s.htps = &http.Server{
			Addr: s.Cfg.ListenAddrTLS,
		}
		s.htps.TLSConfig = &tls.Config{GetCertificate: getcertfn}
	}
	return m
}

func (s *sServer) Run() error {
	s.Init()
	errc := make(chan error, 3)
	ctx, cancelf := context.WithCancel(context.Background())

	s.htps = &http.Server{
		Addr: s.Cfg.ListenAddrTLS,
	}
	autocertManager := s.updateCertificates()

	s.htps.Handler = s

	go func() {
		if autocertManager == nil || s.Cfg.DisableTLS {
			s.Logger.Println("Listening HTTP")
			if autocertManager == nil {
				s.Logger.Println("autocertManager is nil")
			}
			errc <- http.ListenAndServe(s.Cfg.ListenAddr, s)
		} else {
			s.Logger.Println("Listening accepting HTTP requests to the SNI challenge")
			errc <- http.ListenAndServe(s.Cfg.ListenAddr, autocertManager.HTTPHandler(s))
		}
	}()

	certs := make([]util.Certificate, 0, len(s.domains))
	for _, v := range s.domains {
		if v.Certificate.CertFile != "" {
			certs = append(certs, v.Certificate)
		}
	}
	var wrapper *util.ServerWrapper
	if !s.Cfg.DisableTLS {
		go func() {
			if s.Cfg.Graceful {
				s.Logger.Println("Listening HTTPS (Graceful)")
				wrapper = util.NewGracefulServer(manners.NewWithServer(s.htps))
			} else {
				s.Logger.Println("Listening HTTPS (Vanilla)")
				wrapper = util.NewVanillaServer(s.htps)
			}
			errc <- util.ListenAndServeTLSSNI(wrapper, certs)
		}()
	}
	//
	// API
	if err := s.startAPI(ctx); err != nil {
		s.Logger.Println("START API ERROR:", err.Error())
	}
	//
	go func() {
		s.closeChan = make(chan os.Signal, 1)
		<-s.closeChan
		cancelf()
		if wrapper != nil {
			wrapper.Close()
		}
		errc <- nil
	}()
	//
	err := <-errc
	return err
}

func (s *sServer) Close() {
	s.closeChan <- os.Interrupt
}

func (s *sServer) Init() {
	// first start setting the number of cpu cores to use
	ncpu := runtime.NumCPU()
	if s.Cfg.NumCPU > 0 && s.Cfg.NumCPU < ncpu {
		ncpu = s.Cfg.NumCPU
	}
	runtime.GOMAXPROCS(ncpu)
}

func (s *sServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h := r.Host
	if s.Cfg.Debug {
		if ho := r.Header.Get("X-Sandpiper-Host"); ho != "" {
			h = ho
		}
		s.Logger.Println("Host: " + h)
	}
	res := s.trieDomains.Find(h)
	if res == nil {
		if len(s.Cfg.FallbackDomain) > 0 {
			if s.Cfg.Debug {
				s.Logger.Println("FALLBACK DOMAIN", s.Cfg.FallbackDomain)
			}
			dom := s.domains[s.Cfg.FallbackDomain]
			if dom != nil {
				if dom.AuthMode != "" {
					switch dom.AuthMode {
					case "apikey":
						if dom.AuthValue != r.Header.Get(dom.AuthKey) {
							w.WriteHeader(http.StatusUnauthorized)
							return
						}
					}
				}
				dom.ReverseProxy(w, r)
				return
			} else {
				if s.Cfg.Debug {
					s.Logger.Println("FALLBACK DOMAIN NOT FOUND")
				}
				http.Error(w, "fallback domain not found "+h, http.StatusInternalServerError)
				return
			}
		} else {
			if s.Cfg.Debug {
				s.Logger.Println("DOMAIN NOT FOUND")
			}
			http.Error(w, "domain not found "+h, http.StatusInternalServerError)
			return
		}
	}
	if res.EndRoute == nil {
		if s.Cfg.Debug {
			s.Logger.Println("ROUTE IS NULL")
		}
		http.Error(w, "route is null", http.StatusInternalServerError)
		return
	}
	if res.EndRoute.AuthMode != "" {
		switch res.EndRoute.AuthMode {
		case "apikey":
			if res.EndRoute.AuthValue != r.Header.Get(res.EndRoute.AuthKey) {
				w.WriteHeader(http.StatusUnauthorized)
				return
			}
		}
	}
	res.EndRoute.ReverseProxy(w, r)
}
<|endoftext|>