<commit_before>package main

import (
	"bytes"
	"database/sql"
	"errors"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/avabot/ava/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jbrukh/bayesian"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/labstack/echo"
	mw "github.com/avabot/ava/Godeps/_workspace/src/github.com/labstack/echo/middleware"
	_ "github.com/avabot/ava/Godeps/_workspace/src/github.com/lib/pq"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/subosito/twilio"
	"github.com/avabot/ava/Godeps/_workspace/src/golang.org/x/crypto/bcrypt"
	"github.com/avabot/ava/shared/datatypes"
	"github.com/avabot/ava/shared/language"
	"github.com/avabot/ava/shared/pkg"
)

var db *sqlx.DB
var bayes *bayesian.Classifier
var ErrInvalidCommand = errors.New("invalid command")
var ErrMissingPackage = errors.New("missing package")

func main() {
	rand.Seed(time.Now().UnixNano())
	app := cli.NewApp()
	app.Name = "ava"
	app.Usage = "general purpose ai platform"
	app.Version = "0.0.1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "server, s",
			Usage: "run server",
		},
		cli.StringFlag{
			Name:  "port, p",
			Value: "4000",
			Usage: "set port for server",
		},
		cli.BoolFlag{
			Name:  "install, i",
			Usage: "install packages in package.json",
		},
	}
	app.Action = func(c *cli.Context) {
		showHelp := true
		if c.Bool("install") {
			log.Println("TODO: install packages")
			showHelp = false
		}
		if c.Bool("server") {
			db = connectDB()
			startServer(c.String("port"))
			showHelp = false
		}
		if showHelp {
			cli.ShowAppHelp(c)
		}
	}
	app.Run(os.Args)
}

func startServer(port string) {
	var err error
	if err = checkRequiredEnvVars(); err != nil {
		log.Println("err:", err)
	}
	bayes, err = loadClassifier(bayes)
	if err != nil {
		log.Println("err: loading classifier:", err)
	}
	log.Println("booting local server")
	bootRPCServer(port)
	bootTwilio()
	bootDependencies()
	e := echo.New()
	initRoutes(e)
	log.Println("booted ava")
	e.Run(":" + port)
}

func bootRPCServer(port string) {
	ava := new(Ava)
	if err := rpc.Register(ava); err != nil {
		log.Println("register ava in rpc", err)
	}
	p, err := strconv.Atoi(port)
	if err != nil {
		log.Println("convert port to int", err)
	}
	pt := strconv.Itoa(p + 1)
	l, err := net.Listen("tcp", ":"+pt)
	log.Println("booting rpc server", pt)
	if err != nil {
		log.Println("err: rpc listen: ", err)
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				log.Println("err: rpc accept: ", err)
			}
			go rpc.ServeConn(conn)
		}
	}()
}

func connectDB() *sqlx.DB {
	log.Println("connecting to db")
	var d *sqlx.DB
	var err error
	if os.Getenv("AVA_ENV") == "production" {
		d, err = sqlx.Connect("postgres", os.Getenv("DATABASE_URL"))
	} else {
		d, err = sqlx.Connect("postgres",
			"user=egtann dbname=ava sslmode=disable")
	}
	if err != nil {
		log.Println("could not connect to db ", err.Error())
	}
	return d
}

func initRoutes(e *echo.Echo) {
	e.Use(mw.Logger())
	e.Use(mw.Gzip())
	e.Use(mw.Recover())
	e.SetDebug(true)
	e.Static("/public/css", "assets/css")
	e.Static("/public/images", "assets/images")

	e.Get("/", handlerIndex)
	e.Get("/signup", handlerSignup)
	e.Post("/signup", handlerSignupSubmit)
	e.Get("/login", handlerLogin)
	e.Post("/login", handlerLoginSubmit)
	e.Get("/success", handlerLoginSuccess)

	e.Post("/", handlerMain)
	e.Post("/twilio", handlerTwilio)
}

func handlerIndex(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplIndex, err := template.ParseFiles("assets/html/index.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	if err := tmplIndex.Execute(b, struct{}{}); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

// TODO
func handlerTwilio(c *echo.Context) error {
	log.Println(c.Form)
	return errors.New("not implemented")
}

func handlerMain(c *echo.Context) error {
	cmd := c.Form("cmd")
	if len(cmd) == 0 {
		return ErrInvalidCommand
	}
	var ret, pname, route, fid string
	var err error
	var uid, fidT int
	var ctxAdded bool
	var pw *pkg.PkgWrapper
	var m *datatypes.Message
	var u *datatypes.User
	in := &datatypes.Input{}
	si := &datatypes.StructuredInput{}
	if len(cmd) >= 5 && strings.ToLower(cmd)[0:5] == "train" {
		if err := train(bayes, cmd[6:]); err != nil {
			return err
		}
		goto Response
	}
	si, err = classify(bayes, cmd)
	if err != nil {
		log.Println("classifying sentence ", err)
	}
	uid, fid, fidT, err = validateParams(c)
	if err != nil {
		return err
	}
	in = &datatypes.Input{
		StructuredInput: si,
		UserId:          uid,
		FlexId:          fid,
		FlexIdType:      fidT,
	}
	m = &datatypes.Message{User: u, Input: in}
	u, err = getUser(in)
	if err != nil && err != ErrMissingUser {
		log.Println("getUser: ", err)
	}
	m, ctxAdded, err = addContext(m)
	if err != nil {
		log.Println("addContext: ", err)
	}
	ret, route, err = callPkg(m, ctxAdded)
	if err != nil && err != ErrMissingPackage {
		return err
	}
	if len(ret) == 0 {
		ret = language.Confused()
	}
	if pw != nil {
		pname = pw.P.Config.Name
	}
	in.StructuredInput = si
	if err := saveStructuredInput(in, ret, pname, route); err != nil {
		return err
	}
	if err := routeResponse(in, ret); err != nil {
		return err
	}
Response:
	err = c.HTML(http.StatusOK, ret)
	if err != nil {
		return err
	}
	return nil
}

func routeResponse(in *datatypes.Input, ret string) error {
	if in.FlexIdType != datatypes.FlexIdTypePhone {
		return errors.New("route type not implemented")
	}
	params := twilio.MessageParams{Body: ret}
	_, resp, err := tc.Messages.Send("+14242971568", in.FlexId, params)
	log.Println(resp)
	if err != nil {
		return err
	}
	return nil
}

func handlerLogin(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplLogin, err := template.ParseFiles("assets/html/login.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	var data struct{ Error string }
	if c.Get("err") != nil {
		data.Error = c.Get("err").(error).Error()
		c.Set("err", nil)
	}
	if err := tmplLogin.Execute(b, data); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerSignup(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplSignup, err := template.ParseFiles("assets/html/signup.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	data := struct{ Error string }{}
	if c.Get("err") != nil {
		data.Error = c.Get("err").(error).Error()
		c.Set("err", nil)
	}
	if err := tmplSignup.Execute(b, data); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerLoginSubmit(c *echo.Context) error {
	var u struct {
		Id       int
		Password []byte
	}
	var err error
	q := `SELECT id, password FROM users WHERE email=$1`
	err = db.Get(&u, q, c.Form("email"))
	if err == sql.ErrNoRows {
		err = errors.New("Invalid username/password combination")
		goto Response
	} else if err != nil {
		goto Response
	}
	err = bcrypt.CompareHashAndPassword(u.Password, []byte(c.Form("pw")))
	if err != nil {
		goto Response
	}
Response:
	if err != nil {
		c.Set("err", err)
		return handlerLogin(c)
	}
	return handlerLoginSuccess(c)
}

func handlerLoginSuccess(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplSignup, err := template.ParseFiles("assets/html/loginsuccess.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	if err := tmplSignup.Execute(b, struct{}{}); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerSignupSubmit(c *echo.Context) error {
	name := c.Form("name")
	email := c.Form("email")
	pw := c.Form("pw")
	var err error
	var hpw []byte
	var q string
	if len(name) == 0 {
		err = errors.New("You must enter a name.")
		goto Response
	}
	if len(email) == 0 || !strings.ContainsAny(email, "@") {
		err = errors.New("You must enter an email.")
		goto Response
	}
	if len(pw) < 8 {
		err = errors.New("Your password must be at least 8 characters.")
		goto Response
	}
	hpw, err = bcrypt.GenerateFromPassword([]byte(pw), 10)
	if err != nil {
		goto Response
	}
	q = `INSERT INTO users (name, email, password, locationid)
	      VALUES ($1, $2, $3, 0)`
	_, err = db.Exec(q, name, email, hpw)
	if err != nil && err.Error() ==
		`pq: duplicate key value violates unique constraint "users_email_key"` {
		err = errors.New("Sorry, that email is taken.")
	}
Response:
	if err != nil {
		c.Set("err", err)
		return handlerSignup(c)
	}
	return handlerLoginSuccess(c)
}

func validateParams(c *echo.Context) (int, string, int, error) {
	var uid, fidT int
	var fid string
	var err error
	fid = c.Form("flexid")
	uid, err = strconv.Atoi(c.Form("uid"))
	if err.Error() == `strconv.ParseInt: parsing "": invalid syntax` {
		uid = 0
	} else if err != nil {
		return uid, fid, fidT, err
	}
	fidT, err = strconv.Atoi(c.Form("flexidtype"))
	if err != nil && err.Error() == `strconv.ParseInt: parsing "": invalid syntax` {
		fidT = 0
	} else if err != nil {
		return uid, fid, fidT, err
	}
	return uid, fid, fidT, nil
}

func checkRequiredEnvVars() error {
	port := os.Getenv("PORT")
	_, err := strconv.Atoi(port)
	if err != nil {
		return errors.New("PORT is not set to an integer")
	}
	base := os.Getenv("BASE_URL")
	l := len(base)
	if l == 0 {
		return errors.New("BASE_URL not set")
	}
	if l < 4 || base[0:4] != "http" {
		return errors.New("BASE_URL invalid. Must include http/https")
	}
	return nil
}
<commit_msg>Add log statement<commit_after>package main

import (
	"bytes"
	"database/sql"
	"errors"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/avabot/ava/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jbrukh/bayesian"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/jmoiron/sqlx"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/labstack/echo"
	mw "github.com/avabot/ava/Godeps/_workspace/src/github.com/labstack/echo/middleware"
	_ "github.com/avabot/ava/Godeps/_workspace/src/github.com/lib/pq"
	"github.com/avabot/ava/Godeps/_workspace/src/github.com/subosito/twilio"
	"github.com/avabot/ava/Godeps/_workspace/src/golang.org/x/crypto/bcrypt"
	"github.com/avabot/ava/shared/datatypes"
	"github.com/avabot/ava/shared/language"
	"github.com/avabot/ava/shared/pkg"
)

var db *sqlx.DB
var bayes *bayesian.Classifier
var ErrInvalidCommand = errors.New("invalid command")
var ErrMissingPackage = errors.New("missing package")

func main() {
	rand.Seed(time.Now().UnixNano())
	app := cli.NewApp()
	app.Name = "ava"
	app.Usage = "general purpose ai platform"
	app.Version = "0.0.1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "server, s",
			Usage: "run server",
		},
		cli.StringFlag{
			Name:  "port, p",
			Value: "4000",
			Usage: "set port for server",
		},
		cli.BoolFlag{
			Name:  "install, i",
			Usage: "install packages in package.json",
		},
	}
	app.Action = func(c *cli.Context) {
		showHelp := true
		if c.Bool("install") {
			log.Println("TODO: install packages")
			showHelp = false
		}
		if c.Bool("server") {
			db = connectDB()
			startServer(c.String("port"))
			showHelp = false
		}
		if showHelp {
			cli.ShowAppHelp(c)
		}
	}
	app.Run(os.Args)
}

func startServer(port string) {
	var err error
	if err = checkRequiredEnvVars(); err != nil {
		log.Println("err:", err)
	}
	bayes, err = loadClassifier(bayes)
	if err != nil {
		log.Println("err: loading classifier:", err)
	}
	log.Println("booting local server")
	bootRPCServer(port)
	bootTwilio()
	bootDependencies()
	e := echo.New()
	initRoutes(e)
	log.Println("booted ava")
	e.Run(":" + port)
}

func bootRPCServer(port string) {
	ava := new(Ava)
	if err := rpc.Register(ava); err != nil {
		log.Println("register ava in rpc", err)
	}
	p, err := strconv.Atoi(port)
	if err != nil {
		log.Println("convert port to int", err)
	}
	pt := strconv.Itoa(p + 1)
	l, err := net.Listen("tcp", ":"+pt)
	log.Println("booting rpc server", pt)
	if err != nil {
		log.Println("err: rpc listen: ", err)
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				log.Println("err: rpc accept: ", err)
			}
			go rpc.ServeConn(conn)
		}
	}()
}

func connectDB() *sqlx.DB {
	log.Println("connecting to db")
	var d *sqlx.DB
	var err error
	log.Println("AVA_ENV =", os.Getenv("AVA_ENV"))
	if os.Getenv("AVA_ENV") == "production" {
		d, err = sqlx.Connect("postgres", os.Getenv("DATABASE_URL"))
	} else {
		d, err = sqlx.Connect("postgres",
			"user=egtann dbname=ava sslmode=disable")
	}
	if err != nil {
		log.Println("could not connect to db ", err.Error())
	}
	return d
}

func initRoutes(e *echo.Echo) {
	e.Use(mw.Logger())
	e.Use(mw.Gzip())
	e.Use(mw.Recover())
	e.SetDebug(true)
	e.Static("/public/css", "assets/css")
	e.Static("/public/images", "assets/images")

	e.Get("/", handlerIndex)
	e.Get("/signup", handlerSignup)
	e.Post("/signup", handlerSignupSubmit)
	e.Get("/login", handlerLogin)
	e.Post("/login", handlerLoginSubmit)
	e.Get("/success", handlerLoginSuccess)

	e.Post("/", handlerMain)
	e.Post("/twilio", handlerTwilio)
}

func handlerIndex(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplIndex, err := template.ParseFiles("assets/html/index.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	if err := tmplIndex.Execute(b, struct{}{}); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

// TODO
func handlerTwilio(c *echo.Context) error {
	log.Println(c.Form)
	return errors.New("not implemented")
}

func handlerMain(c *echo.Context) error {
	cmd := c.Form("cmd")
	if len(cmd) == 0 {
		return ErrInvalidCommand
	}
	var ret, pname, route, fid string
	var err error
	var uid, fidT int
	var ctxAdded bool
	var pw *pkg.PkgWrapper
	var m *datatypes.Message
	var u *datatypes.User
	in := &datatypes.Input{}
	si := &datatypes.StructuredInput{}
	if len(cmd) >= 5 && strings.ToLower(cmd)[0:5] == "train" {
		if err := train(bayes, cmd[6:]); err != nil {
			return err
		}
		goto Response
	}
	si, err = classify(bayes, cmd)
	if err != nil {
		log.Println("classifying sentence ", err)
	}
	uid, fid, fidT, err = validateParams(c)
	if err != nil {
		return err
	}
	in = &datatypes.Input{
		StructuredInput: si,
		UserId:          uid,
		FlexId:          fid,
		FlexIdType:      fidT,
	}
	m = &datatypes.Message{User: u, Input: in}
	u, err = getUser(in)
	if err != nil && err != ErrMissingUser {
		log.Println("getUser: ", err)
	}
	m, ctxAdded, err = addContext(m)
	if err != nil {
		log.Println("addContext: ", err)
	}
	ret, route, err = callPkg(m, ctxAdded)
	if err != nil && err != ErrMissingPackage {
		return err
	}
	if len(ret) == 0 {
		ret = language.Confused()
	}
	if pw != nil {
		pname = pw.P.Config.Name
	}
	in.StructuredInput = si
	if err := saveStructuredInput(in, ret, pname, route); err != nil {
		return err
	}
	if err := routeResponse(in, ret); err != nil {
		return err
	}
Response:
	err = c.HTML(http.StatusOK, ret)
	if err != nil {
		return err
	}
	return nil
}

func routeResponse(in *datatypes.Input, ret string) error {
	if in.FlexIdType != datatypes.FlexIdTypePhone {
		return errors.New("route type not implemented")
	}
	params := twilio.MessageParams{Body: ret}
	_, resp, err := tc.Messages.Send("+14242971568", in.FlexId, params)
	log.Println(resp)
	if err != nil {
		return err
	}
	return nil
}

func handlerLogin(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplLogin, err := template.ParseFiles("assets/html/login.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	var data struct{ Error string }
	if c.Get("err") != nil {
		data.Error = c.Get("err").(error).Error()
		c.Set("err", nil)
	}
	if err := tmplLogin.Execute(b, data); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerSignup(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplSignup, err := template.ParseFiles("assets/html/signup.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	data := struct{ Error string }{}
	if c.Get("err") != nil {
		data.Error = c.Get("err").(error).Error()
		c.Set("err", nil)
	}
	if err := tmplSignup.Execute(b, data); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerLoginSubmit(c *echo.Context) error {
	var u struct {
		Id       int
		Password []byte
	}
	var err error
	q := `SELECT id, password FROM users WHERE email=$1`
	err = db.Get(&u, q, c.Form("email"))
	if err == sql.ErrNoRows {
		err = errors.New("Invalid username/password combination")
		goto Response
	} else if err != nil {
		goto Response
	}
	err = bcrypt.CompareHashAndPassword(u.Password, []byte(c.Form("pw")))
	if err != nil {
		goto Response
	}
Response:
	if err != nil {
		c.Set("err", err)
		return handlerLogin(c)
	}
	return handlerLoginSuccess(c)
}

func handlerLoginSuccess(c *echo.Context) error {
	tmplLayout, err := template.ParseFiles("assets/html/layout.html")
	if err != nil {
		log.Fatalln(err)
	}
	tmplSignup, err := template.ParseFiles("assets/html/loginsuccess.html")
	if err != nil {
		log.Fatalln(err)
	}
	var s []byte
	b := bytes.NewBuffer(s)
	if err := tmplSignup.Execute(b, struct{}{}); err != nil {
		log.Fatalln(err)
	}
	b2 := bytes.NewBuffer(s)
	if err := tmplLayout.Execute(b2, b); err != nil {
		log.Fatalln(err)
	}
	err = c.HTML(http.StatusOK, "%s", b2)
	if err != nil {
		return err
	}
	return nil
}

func handlerSignupSubmit(c *echo.Context) error {
	name := c.Form("name")
	email := c.Form("email")
	pw := c.Form("pw")
	var err error
	var hpw []byte
	var q string
	if len(name) == 0 {
		err = errors.New("You must enter a name.")
		goto Response
	}
	if len(email) == 0 || !strings.ContainsAny(email, "@") {
		err = errors.New("You must enter an email.")
		goto Response
	}
	if len(pw) < 8 {
		err = errors.New("Your password must be at least 8 characters.")
		goto Response
	}
	hpw, err = bcrypt.GenerateFromPassword([]byte(pw), 10)
	if err != nil {
		goto Response
	}
	q = `INSERT INTO users (name, email, password, locationid)
	      VALUES ($1, $2, $3, 0)`
	_, err = db.Exec(q, name, email, hpw)
	if err != nil && err.Error() ==
		`pq: duplicate key value violates unique constraint "users_email_key"` {
		err = errors.New("Sorry, that email is taken.")
	}
Response:
	if err != nil {
		c.Set("err", err)
		return handlerSignup(c)
	}
	return handlerLoginSuccess(c)
}

func validateParams(c *echo.Context) (int, string, int, error) {
	var uid, fidT int
	var fid string
	var err error
	fid = c.Form("flexid")
	uid, err = strconv.Atoi(c.Form("uid"))
	if err.Error() == `strconv.ParseInt: parsing "": invalid syntax` {
		uid = 0
	} else if err != nil {
		return uid, fid, fidT, err
	}
	fidT, err = strconv.Atoi(c.Form("flexidtype"))
	if err != nil && err.Error() == `strconv.ParseInt: parsing "": invalid syntax` {
		fidT = 0
	} else if err != nil {
		return uid, fid, fidT, err
	}
	return uid, fid, fidT, nil
}

func checkRequiredEnvVars() error {
	port := os.Getenv("PORT")
	_, err := strconv.Atoi(port)
	if err != nil {
		return errors.New("PORT is not set to an integer")
	}
	base := os.Getenv("BASE_URL")
	l := len(base)
	if l == 0 {
		return errors.New("BASE_URL not set")
	}
	if l < 4 || base[0:4] != "http" {
		return errors.New("BASE_URL invalid. Must include http/https")
	}
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

var (
	S3Client   *s3.S3
	bucketName string
	region     string
)

func NewS3(bname, re string) error {
	bucketName = bname
	region = re

	creds := credentials.NewEnvCredentials()
	_, err := creds.Get()
	if err != nil {
		return err
	}

	config := &aws.Config{
		Endpoint:         aws.String(fmt.Sprintf("https://s3-%s.amazonaws.com", region)),
		S3ForcePathStyle: aws.Bool(true),
		Credentials:      creds,
	}

	sess, err := session.NewSession(config)
	if err != nil {
		fmt.Println("failed to create session")
		return err
	}

	S3Client = s3.New(sess)

	return nil
}

func UploadToS3(filePath string, key string) error {
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}

	defer file.Close()

	fileInfo, err := file.Stat()
	if err != nil {
		return err
	}
	var size int64 = fileInfo.Size()

	buffer := make([]byte, size)

	// read file content to buffer
	file.Read(buffer)

	fileBytes := bytes.NewReader(buffer) // convert to io.ReadSeeker type
	fileType := http.DetectContentType(buffer)

	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(key),        // required
		ACL:           aws.String("public-read"),
		Body:          fileBytes,
		ContentLength: aws.Int64(size),
		ContentType:   aws.String(fileType),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err = S3Client.PutObject(params)
	if err != nil {
		return err
	}

	return nil

}

func UploadToS3ByMultiPart(file multipart.File, fileName string) error {

	bs, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(fileName),   // required
		ACL:           aws.String("public-read"),
		Body:          file,
		ContentLength: aws.Int64(int64(len(bs))),
		ContentType:   aws.String("application/octet-stream"),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err = S3Client.PutObject(params)
	if err != nil {
		return err
	}

	return nil
}

func DeleteFromS3(key string) error {

	params := &s3.DeleteObjectInput{
		Bucket: aws.String(bucketName), // Required
		Key:    aws.String(key),        // Required
	}
	_, err := S3Client.DeleteObject(params)

	return err
}
<commit_msg>Modify aws<commit_after>package aws

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"mime/multipart"
	"net/http"
	"os"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

var (
	S3Client   *s3.S3
	bucketName string
	region     string
)

func init() {
	if os.Getenv("AWS_BUCKET_NAME") == "" {
		panic("AWS_BUCKET_NAME is missing?")
	}
	if os.Getenv("AWS_REGION") == "" {
		panic("AWS_REGION is missing?")
	}

	bucketName = os.Getenv("AWS_BUCKET_NAME")
	region = os.Getenv("AWS_REGION")

	creds := credentials.NewEnvCredentials()
	_, err := creds.Get()
	if err != nil {
		panic("failed to get env credentials")
	}

	config := &aws.Config{
		Endpoint:         aws.String(fmt.Sprintf("https://s3-%s.amazonaws.com", region)),
		S3ForcePathStyle: aws.Bool(true),
		Credentials:      creds,
	}

	sess, err := session.NewSession(config)
	if err != nil {
		panic("failed to create session")
	}

	S3Client = s3.New(sess)
}

func UploadToS3(filePath string, key string) error {

	file, err := os.Open(filePath)
	if err != nil {
		return err
	}

	defer file.Close()

	fileInfo, err := file.Stat()
	if err != nil {
		return err
	}
	var size int64 = fileInfo.Size()

	buffer := make([]byte, size)

	// read file content to buffer
	file.Read(buffer)

	fileBytes := bytes.NewReader(buffer) // convert to io.ReadSeeker type
	fileType := http.DetectContentType(buffer)

	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(key),        // required
		ACL:           aws.String("public-read"),
		Body:          fileBytes,
		ContentLength: aws.Int64(size),
		ContentType:   aws.String(fileType),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err = S3Client.PutObject(params)
	if err != nil {
		return err
	}

	return nil

}

func UploadToS3ByMultiPart(file multipart.File, fileName string) error {

	bs, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	return uploadWithBytes(bs, fileName)
}

func UploadToS3WithBytes(bytes []byte, fileName string) error {
	return uploadWithBytes(bytes, fileName)
}

func uploadWithBytes(bs []byte, fileName string) error {
	params := &s3.PutObjectInput{
		Bucket:        aws.String(bucketName), // required
		Key:           aws.String(fileName),   // required
		ACL:           aws.String("public-read"),
		Body:          bytes.NewReader(bs),
		ContentLength: aws.Int64(int64(len(bs))),
		ContentType:   aws.String("application/octet-stream"),
		Metadata: map[string]*string{
			"Key": aws.String("MetadataValue"), //required
		},
		// see more at http://godoc.org/github.com/aws/aws-sdk-go/service/s3#S3.PutObject
	}

	_, err := S3Client.PutObject(params)
	if err != nil {
		return err
	}

	return nil
}

func DeleteFromS3(key string) error {

	params := &s3.DeleteObjectInput{
		Bucket: aws.String(bucketName), // Required
		Key:    aws.String(key),        // Required
	}
	_, err := S3Client.DeleteObject(params)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"html/template"
	"io/ioutil"
	"net/http"

	"github.com/prometheus/common/log"
	"github.com/prometheus/prometheus/config"

	"gopkg.in/yaml.v2"
)

var (
	configFile    = flag.String("config.file", "prometheus.yml", "prometheus configuration file name.")
	listenAddress = flag.String("listen-address", ":12345", "The address to listen on for HTTP requests.")
)

func rewrite() {
	conf, err := config.LoadFile(*configFile)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
	}
	yamlData, yamlErr := yaml.Marshal(&conf)
	if yamlErr != nil {
		log.Fatalf("error: %v", yamlErr)
	}
	writeErr := ioutil.WriteFile(*configFile, yamlData, 0644)
	if writeErr != nil {
		log.Fatalf("error: %v", writeErr)
	}
}

func create(job_name string) {
	newcfg := &config.ScrapeConfig{}
	newcfg.JobName = job_name
	fsdc := &config.FileSDConfig{}
	fsdc.Names = []string{"aaa.yaml"}
	newcfg.FileSDConfigs = append(newcfg.FileSDConfigs, fsdc)
	rc := &config.RelabelConfig{}
	rc.Regex = config.MustNewRegexp("(.*):\\d+")
	newcfg.RelabelConfigs = append(newcfg.RelabelConfigs, rc)

	var newScrapeConfigs []*config.ScrapeConfig
	newScrapeConfigs = append(newScrapeConfigs, newcfg)
	conf, confErr := config.LoadFile(*configFile)
	if confErr != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, confErr)
	}
	for _, scfg := range conf.ScrapeConfigs {
		newScrapeConfigs = append(newScrapeConfigs, scfg)
	}
	conf.ScrapeConfigs = newScrapeConfigs
	newYamlData, err := yaml.Marshal(&conf)
	if err != nil {
		log.Fatalf("error: %v", err)
	}
	writeErr := ioutil.WriteFile(*configFile, newYamlData, 0644)
	if writeErr != nil {
		log.Fatalf("error: %v", writeErr)
	}
	rewrite()
}

func main() {
	flag.Parse()

	log.Infof("Loading configuration file %s", *configFile)
	log.Infof("Starting Server: %s", *listenAddress)

	t := template.Must(template.ParseFiles("templates/toya.html"))
	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		if r.Method == "GET" {
			conf, err := config.LoadFile(*configFile)
			if err != nil {
				log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
			}
			t.Execute(w, conf.ScrapeConfigs)
		}
	})
	http.HandleFunc("/create", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		if r.Method == "POST" {
			create(r.FormValue("job_name"))
			http.Redirect(w, r, "/", http.StatusFound)
		}
	})
	err := http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

}
<commit_msg>add exporter list<commit_after>package main

import (
	"flag"
	"html/template"
	"io/ioutil"
	"net/http"
	"path/filepath"

	"github.com/prometheus/common/log"
	"github.com/prometheus/prometheus/config"

	"gopkg.in/yaml.v2"
)

var (
	configFile      = flag.String("config.file", "prometheus.yml", "prometheus configuration file name.")
	exporterListDir = flag.String("exporter.list.dir", ".", "file_sd_configs names file dir")
	listenAddress   = flag.String("listen-address", ":12345", "The address to listen on for HTTP requests.")
)

func rewrite() {
	conf, err := config.LoadFile(*configFile)
	if err != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
	}
	yamlData, yamlErr := yaml.Marshal(&conf)
	if yamlErr != nil {
		log.Fatalf("error: %v", yamlErr)
	}
	writeErr := ioutil.WriteFile(*configFile, yamlData, 0644)
	if writeErr != nil {
		log.Fatalf("error: %v", writeErr)
	}
}

func create(job_name string, exporter_list string) {
	newcfg := &config.ScrapeConfig{}
	newcfg.JobName = job_name
	fsdc := &config.FileSDConfig{}
	exporterListFile := filepath.Join(*exporterListDir, job_name+".yml")
	exporterListFileWriteErr := ioutil.WriteFile(exporterListFile, []byte(exporter_list), 0644)
	if exporterListFileWriteErr != nil {
		log.Fatalf("error: %v", exporterListFileWriteErr)
	}
	fsdc.Names = []string{exporterListFile}
	newcfg.FileSDConfigs = append(newcfg.FileSDConfigs, fsdc)
	rc := &config.RelabelConfig{}
	rc.Regex = config.MustNewRegexp("(.*):\\d+")
	newcfg.RelabelConfigs = append(newcfg.RelabelConfigs, rc)

	var newScrapeConfigs []*config.ScrapeConfig
	newScrapeConfigs = append(newScrapeConfigs, newcfg)
	conf, confErr := config.LoadFile(*configFile)
	if confErr != nil {
		log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, confErr)
	}
	for _, scfg := range conf.ScrapeConfigs {
		newScrapeConfigs = append(newScrapeConfigs, scfg)
	}
	conf.ScrapeConfigs = newScrapeConfigs
	newYamlData, err := yaml.Marshal(&conf)
	if err != nil {
		log.Fatalf("error: %v", err)
	}
	writeErr := ioutil.WriteFile(*configFile, newYamlData, 0644)
	if writeErr != nil {
		log.Fatalf("error: %v", writeErr)
	}
	rewrite()
}

func main() {
	flag.Parse()

	log.Infof("Loading configuration file %s", *configFile)
	log.Infof("Starting Server: %s", *listenAddress)

	t := template.Must(template.ParseFiles("templates/toya.html"))
	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		if r.Method == "GET" {
			conf, err := config.LoadFile(*configFile)
			if err != nil {
				log.Errorf("Couldn't load configuration (-config.file=%s): %v", *configFile, err)
			}
			t.Execute(w, conf.ScrapeConfigs)
		}
	})
	http.HandleFunc("/create", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		if r.Method == "POST" {
			create(r.FormValue("job_name"), r.FormValue("exporter_list"))
			http.Redirect(w, r, "/", http.StatusFound)
		}
	})
	err := http.ListenAndServe(*listenAddress, nil)
	if err != nil {
		log.Fatalf("error: %v", err)
	}

}
<|endoftext|>


<commit_before>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// controlFlow tells the main loop what it should do next.
type controlFlow int

// controlFlow values.
const (
	controlStop  controlFlow = 0
	controlBegin controlFlow = 1
	controlNext  controlFlow = 2
)

// dataType specifies which type of data we are going to save for each node.
type dataType = Handler

// tree represents a radix tree.
type tree struct {
	prefix    string
	data      dataType
	children  [224]*tree
	parameter *tree
	wildcard  *tree
	kind      byte
}

// add adds a new element to the tree.
func (node *tree) add(path string, data dataType) {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			// This only occurs when the same parameter based route is added twice.
			// node: /post/:id|
			// path: /post/:id|
			if i == len(path) {
				node.data = data
				return
			}

			// When we hit a separator, we'll search for a fitting child.
			if path[i] == separator {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

		default:
			if i == len(path) {
				// The path already exists.
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					node.data = data
					return
				}

				// The path ended but the node prefix is longer.
				// node: /blog|feed
				// path: /blog|
				node.split(i-offset, "", data)
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				node.split(i-offset, path[i:], data)
				return
			}
		}

	next:
		i++
	}
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *tree) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	/// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.children[splitNode.prefix[0]-32] = splitNode
		return
	}

	node.children[splitNode.prefix[0]-32] = splitNode

	// Create new nodes with the remaining path
	node.append(path, data)
}

// clone clones the node with a new prefix.
func (node *tree) clone(prefix string) *tree {
	return &tree{
		prefix:    prefix,
		data:      node.data,
		children:  node.children,
		parameter: node.parameter,
		wildcard:  node.wildcard,
		kind:      node.kind,
	}
}

// reset resets the existing node data.
func (node *tree) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.children = [224]*tree{}
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *tree) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.children[separator-32] != nil || node.kind == wildcard {
		return
	}

	node.children[separator-32] = &tree{
		prefix: "/",
		data:   data,
	}
}

// append appends the given path to the tree.
func (node *tree) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			// If the node itself doesn't have a prefix (root node),
			// don't add a child and use the node itself.
			if node.prefix == "" {
				node.prefix = path
				node.data = data
				return
			}

			child := &tree{
				prefix: path,
				data:   data,
			}

			node.children[path[0]-32] = child
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &tree{
				prefix: path[1:paramEnd],
				kind:   path[paramStart],
			}

			switch child.kind {
			case parameter:
				child.addTrailingSlash(data)
				node.parameter = child
				node = child
				path = path[paramEnd:]
				continue

			case wildcard:
				child.data = data
				node.wildcard = child
				return
			}
		}

		// We know there's a parameter, but not directly at the start.

		// If the node itself doesn't have a prefix (root node),
		// don't add a child and use the node itself.
		if node.prefix == "" {
			node.prefix = path[:paramStart]
			path = path[paramStart:]
			continue
		}

		// Add a normal node with the path before the parameter start.
		child := &tree{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.children[path[0]-32] = child
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *tree) end(path string, data dataType, i int, offset int) (*tree, int, controlFlow) {
	child := node.children[path[i]-32]

	if child != nil {
		node = child
		offset = i
		return node, offset, controlNext
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.append(path[i:], data)
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// find finds the data for the given path and assigns it to ctx.handler, if available.
func (node *tree) find(path string, ctx *context) {
	var (
		i                  int
		offset             int
		lastWildcardOffset int
		lastWildcard       *tree
	)

	// Search tree for equal parts until we can no longer proceed
	for {
	begin:
		switch node.kind {
		case parameter:
			if i == len(path) {
				ctx.addParameter(node.prefix, path[offset:i])
				ctx.handler = node.data
				return
			}

			if path[i] == separator {
				ctx.addParameter(node.prefix, path[offset:i])
				node = node.children[separator-32]
				offset = i
				goto next
			}

		default:
			// We reached the end.
			if i == len(path) {
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					ctx.handler = node.data
					return
				}

				// node: /blog|feed
				// path: /blog|
				ctx.handler = nil
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				if node.wildcard != nil {
					lastWildcard = node.wildcard
					lastWildcardOffset = i
				}

				child := node.children[path[i]-32]

				if child != nil {
					node = child
					offset = i
					goto next
				}

				// node: /|:id
				// path: /|blog
				if node.parameter != nil {
					node = node.parameter
					offset = i
					goto begin
				}

				// node: /|*any
				// path: /|image.png
				if node.wildcard != nil {
					ctx.addParameter(node.wildcard.prefix, path[i:])
					ctx.handler = node.wildcard.data
					return
				}

				ctx.handler = nil
				return
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				if lastWildcard != nil {
					ctx.addParameter(lastWildcard.prefix, path[lastWildcardOffset:])
					ctx.handler = lastWildcard.data
					return
				}

				ctx.handler = nil
				return
			}
		}

	next:
		i++
	}
}

// each traverses the tree and calls the given function on every node.
func (node *tree) each(callback func(*tree)) {
	callback(node)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.each(callback)
	}

	if node.parameter != nil {
		node.parameter.each(callback)
	}

	if node.wildcard != nil {
		node.wildcard.each(callback)
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *tree) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *tree) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<commit_msg>Eliminate a jump in the hot path<commit_after>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// controlFlow tells the main loop what it should do next.
type controlFlow int

// controlFlow values.
const (
	controlStop  controlFlow = 0
	controlBegin controlFlow = 1
	controlNext  controlFlow = 2
)

// dataType specifies which type of data we are going to save for each node.
type dataType = Handler

// tree represents a radix tree.
type tree struct {
	prefix    string
	data      dataType
	children  [224]*tree
	parameter *tree
	wildcard  *tree
	kind      byte
}

// add adds a new element to the tree.
func (node *tree) add(path string, data dataType) {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			// This only occurs when the same parameter based route is added twice.
			// node: /post/:id|
			// path: /post/:id|
			if i == len(path) {
				node.data = data
				return
			}

			// When we hit a separator, we'll search for a fitting child.
			if path[i] == separator {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

		default:
			if i == len(path) {
				// The path already exists.
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					node.data = data
					return
				}

				// The path ended but the node prefix is longer.
				// node: /blog|feed
				// path: /blog|
				node.split(i-offset, "", data)
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				node.split(i-offset, path[i:], data)
				return
			}
		}

	next:
		i++
	}
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *tree) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	/// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.children[splitNode.prefix[0]-32] = splitNode
		return
	}

	node.children[splitNode.prefix[0]-32] = splitNode

	// Create new nodes with the remaining path
	node.append(path, data)
}

// clone clones the node with a new prefix.
func (node *tree) clone(prefix string) *tree {
	return &tree{
		prefix:    prefix,
		data:      node.data,
		children:  node.children,
		parameter: node.parameter,
		wildcard:  node.wildcard,
		kind:      node.kind,
	}
}

// reset resets the existing node data.
func (node *tree) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.children = [224]*tree{}
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *tree) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.children[separator-32] != nil || node.kind == wildcard {
		return
	}

	node.children[separator-32] = &tree{
		prefix: "/",
		data:   data,
	}
}

// append appends the given path to the tree.
func (node *tree) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			// If the node itself doesn't have a prefix (root node),
			// don't add a child and use the node itself.
			if node.prefix == "" {
				node.prefix = path
				node.data = data
				return
			}

			child := &tree{
				prefix: path,
				data:   data,
			}

			node.children[path[0]-32] = child
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &tree{
				prefix: path[1:paramEnd],
				kind:   path[paramStart],
			}

			switch child.kind {
			case parameter:
				child.addTrailingSlash(data)
				node.parameter = child
				node = child
				path = path[paramEnd:]
				continue

			case wildcard:
				child.data = data
				node.wildcard = child
				return
			}
		}

		// We know there's a parameter, but not directly at the start.

		// If the node itself doesn't have a prefix (root node),
		// don't add a child and use the node itself.
		if node.prefix == "" {
			node.prefix = path[:paramStart]
			path = path[paramStart:]
			continue
		}

		// Add a normal node with the path before the parameter start.
		child := &tree{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.children[path[0]-32] = child
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *tree) end(path string, data dataType, i int, offset int) (*tree, int, controlFlow) {
	child := node.children[path[i]-32]

	if child != nil {
		node = child
		offset = i
		return node, offset, controlNext
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.append(path[i:], data)
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// find finds the data for the given path and assigns it to ctx.handler, if available.
func (node *tree) find(path string, ctx *context) {
	var (
		i                  int
		offset             int
		lastWildcardOffset int
		lastWildcard       *tree
	)

	// Search tree for equal parts until we can no longer proceed
	for {
	begin:
		switch node.kind {
		case parameter:
			if i == len(path) {
				ctx.addParameter(node.prefix, path[offset:i])
				ctx.handler = node.data
				return
			}

			if path[i] == separator {
				ctx.addParameter(node.prefix, path[offset:i])
				node = node.children[separator-32]
				offset = i
				i++
				goto begin
			}

		default:
			// We reached the end.
			if i == len(path) {
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					ctx.handler = node.data
					return
				}

				// node: /blog|feed
				// path: /blog|
				ctx.handler = nil
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				if node.wildcard != nil {
					lastWildcard = node.wildcard
					lastWildcardOffset = i
				}

				child := node.children[path[i]-32]

				if child != nil {
					node = child
					offset = i
					i++
					goto begin
				}

				// node: /|:id
				// path: /|blog
				if node.parameter != nil {
					node = node.parameter
					offset = i
					goto begin
				}

				// node: /|*any
				// path: /|image.png
				if node.wildcard != nil {
					ctx.addParameter(node.wildcard.prefix, path[i:])
					ctx.handler = node.wildcard.data
					return
				}

				ctx.handler = nil
				return
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				if lastWildcard != nil {
					ctx.addParameter(lastWildcard.prefix, path[lastWildcardOffset:])
					ctx.handler = lastWildcard.data
					return
				}

				ctx.handler = nil
				return
			}
		}

		i++
	}
}

// each traverses the tree and calls the given function on every node.
func (node *tree) each(callback func(*tree)) {
	callback(node)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.each(callback)
	}

	if node.parameter != nil {
		node.parameter.each(callback)
	}

	if node.wildcard != nil {
		node.wildcard.each(callback)
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *tree) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *tree) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 ibelie, Chen Jie, Joungtao. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package tygo

import (
	"fmt"
	"log"
	"reflect"
	"sort"
	"strings"
)

var TYGO_PATH = reflect.TypeOf(Tygo{}).PkgPath()

type Tygo struct {
	cachedSize int
}

func (t *Tygo) SetCachedSize(cachedSize int) {
	t.cachedSize = cachedSize
}

func (t *Tygo) CachedSize() int {
	return t.cachedSize
}

type Type interface {
	WireType() WireType
	Identifier() string
	String() string
	IsPrimitive() bool
	IsIterative() bool
	Go() (string, map[string]string)
	Typyd() string
	Python() string
	Typescript() string
	Javascript() (string, map[string]string)
	ByteSizeGo(string, string, string, int, bool) (string, map[string]string)
	CachedSizeGo(string, string, string, int, bool) (string, map[string]string)
	SerializeGo(string, string, string, int, bool) (string, map[string]string)
	DeserializeGo(string, string, string, string, int, bool) (string, WireType, map[string]string)
}

type Enum struct {
	nameMax int
	sorted  []string
	Name    string
	Package string
	Values  map[string]int
}

type Field struct {
	Type
	Name string
	Hide bool
}

type Method struct {
	Name    string
	Params  []Type
	Results []Type
}

type Object struct {
	Name    string
	Package string
	Parent  *InstanceType
	Fields  []*Field
	Methods []*Method
}

type UnknownType string

type SimpleType uint

const (
	SimpleType_UNKNOWN SimpleType = iota
	SimpleType_NIL
	SimpleType_INT32
	SimpleType_INT64
	SimpleType_UINT32
	SimpleType_UINT64
	SimpleType_BYTES
	SimpleType_STRING
	SimpleType_SYMBOL
	SimpleType_BOOL
	SimpleType_FLOAT32
	SimpleType_FLOAT64
)

type FixedPointType struct {
	Precision uint
	Floor     int
}

type EnumType struct {
	*Enum
	Name string
}

type InstanceType struct {
	*Object
	IsPtr   bool
	Name    string
	PkgName string
	PkgPath string
}

type ListType struct {
	E Type
}

type DictType struct {
	K Type
	V Type
}

type VariantType struct {
	Ts []Type
}

func (t *Enum) IsPrimitive() bool {
	return false
}

func (t *Object) IsPrimitive() bool {
	return false
}

func (t UnknownType) IsPrimitive() bool {
	return false
}

func (t SimpleType) IsPrimitive() bool {
	switch t {
	case SimpleType_INT32:
		fallthrough
	case SimpleType_INT64:
		fallthrough
	case SimpleType_UINT32:
		fallthrough
	case SimpleType_UINT64:
		fallthrough
	case SimpleType_BOOL:
		fallthrough
	case SimpleType_FLOAT32:
		fallthrough
	case SimpleType_FLOAT64:
		return true
	default:
		return false
	}
}

func (t *FixedPointType) IsPrimitive() bool {
	return true
}

func (t *EnumType) IsPrimitive() bool {
	return true
}

func (t *InstanceType) IsPrimitive() bool {
	return false
}

func (t *ListType) IsPrimitive() bool {
	return false
}

func (t *DictType) IsPrimitive() bool {
	return false
}

func (t *VariantType) IsPrimitive() bool {
	return false
}

func (t *Enum) IsIterative() bool {
	return false
}

func (t *Object) IsIterative() bool {
	return false
}

func (t UnknownType) IsIterative() bool {
	return false
}

func (t SimpleType) IsIterative() bool {
	return false
}

func (t *FixedPointType) IsIterative() bool {
	return false
}

func (t *EnumType) IsIterative() bool {
	return false
}

func (t *InstanceType) IsIterative() bool {
	return false
}

func (t *ListType) IsIterative() bool {
	return !t.E.IsPrimitive()
}

func (t *DictType) IsIterative() bool {
	return true
}

func (t *VariantType) IsIterative() bool {
	return false
}

func (t *Enum) FullName() string {
	return t.Package + "/" + t.Name
}

func (t *Enum) Sorted() []string {
	if t.sorted == nil || len(t.sorted) != len(t.Values) {
		t.sorted = nil
		for name, _ := range t.Values {
			if t.nameMax < len(name) {
				t.nameMax = len(name)
			}
			t.sorted = append(t.sorted, name)
		}
	} else if t.nameMax <= 0 {
		for name, _ := range t.Values {
			if t.nameMax < len(name) {
				t.nameMax = len(name)
			}
		}
	}
	if !sort.IsSorted(t) {
		sort.Sort(t)
	}
	return t.sorted
}

func (t *Enum) Len() int {
	return len(t.sorted)
}

func (t *Enum) Swap(i, j int) {
	t.sorted[i], t.sorted[j] = t.sorted[j], t.sorted[i]
}

func (t *Enum) Less(i, j int) bool {
	return t.Values[t.sorted[i]] < t.Values[t.sorted[j]]
}

func (t *Enum) String() string {
	var values []string
	for _, name := range t.Sorted() {
		values = append(values, fmt.Sprintf("%s: %d", name, t.Values[name]))
	}
	return fmt.Sprintf("%s[%s]", t.Name, strings.Join(values, ", "))
}

func (t *Method) String() string {
	var params []string
	for _, param := range t.Params {
		params = append(params, param.String())
	}
	s := fmt.Sprintf("%s(%s)", t.Name, strings.Join(params, ", "))

	var results []string
	for _, result := range t.Results {
		results = append(results, result.String())
	}
	if len(results) == 1 {
		s = fmt.Sprintf("%s %s", s, results[0])
	} else if len(results) > 1 {
		s = fmt.Sprintf("%s (%s)", s, strings.Join(results, ", "))
	}

	return s
}

func (t *Object) FullName() string {
	return t.Package + "/" + t.Name
}

func (t *Object) HasParent() bool {
	if t.Parent.Name == "Tygo" && t.Parent.PkgName == "tygo" && t.Parent.PkgPath == TYGO_PATH {
		return false
	} else if t.Parent.PkgPath != "" {
		return true
	} else {
		name, num := t.Parent.Object.MaxFieldNum()
		return name != "" || num != 0
	}
}

func (t *Object) MaxFieldNum() (string, int) {
	if !t.HasParent() {
		return "", len(t.VisibleFields())
	} else if t.Parent.PkgPath == "" {
		name, num := t.Parent.Object.MaxFieldNum()
		return name, num + len(t.VisibleFields())
	} else {
		return t.Parent.Name, len(t.VisibleFields())
	}
}

func (t *Object) VisibleFields() (fields []*Field) {
	for _, field := range t.Fields {
		if !field.Hide {
			fields = append(fields, field)
		}
	}
	return
}

func (t *Object) AllFields(objects map[string]*Object, visible bool) []*Field {
	fields := t.Fields
	if visible {
		fields = t.VisibleFields()
	}
	if !t.HasParent() {
		return fields
	} else if t.Parent.PkgPath == "" {
		return append(t.Parent.Object.AllFields(objects, visible), fields...)
	} else if object, ok := objects[t.Parent.Name]; ok {
		return append(object.AllFields(objects, visible), fields...)
	} else {
		log.Fatalf("[Tygo][Object] Unsolved parent type: %s", t.Parent.Name)
		return fields
	}
}

func (t *Object) String() string {
	var fields []string

	fields = append(fields, fmt.Sprintf(`
	%s`, t.Parent))

	nameMax := 0
	for _, field := range t.Fields {
		length := len(field.Name)
		if field.Hide {
			length++
		}
		if nameMax < length {
			nameMax = length
		}
	}
	for _, field := range t.Fields {
		name := field.Name
		if field.Hide {
			name = "." + field.Name
		}
		fields = append(fields, fmt.Sprintf(`
	%s %s%s`, name, strings.Repeat(" ", nameMax-len(name)), field))
	}
	for _, method := range t.Methods {
		fields = append(fields, fmt.Sprintf(`
	%s`, method))
	}

	return fmt.Sprintf(`
%s{%s
}
`, t.Name, strings.Join(fields, ""))
}

func (t UnknownType) String() string {
	return string(t)
}

func SimpleType_FromString(s string) Type {
	switch s {
	case STR_NIL:
		return SimpleType_NIL
	case STR_INT32:
		return SimpleType_INT32
	case STR_INT64:
		return SimpleType_INT64
	case STR_UINT32:
		return SimpleType_UINT32
	case STR_UINT64:
		return SimpleType_UINT64
	case STR_BYTES:
		return SimpleType_BYTES
	case STR_STRING:
		return SimpleType_STRING
	case STR_SYMBOL:
		return SimpleType_SYMBOL
	case STR_BOOL:
		return SimpleType_BOOL
	case STR_FLOAT32:
		return SimpleType_FLOAT32
	case STR_FLOAT64:
		return SimpleType_FLOAT64
	default:
		return UnknownType(s)
	}
}

func (t SimpleType) String() string {
	switch t {
	case SimpleType_NIL:
		return STR_NIL
	case SimpleType_INT32:
		return STR_INT32
	case SimpleType_INT64:
		return STR_INT64
	case SimpleType_UINT32:
		return STR_UINT32
	case SimpleType_UINT64:
		return STR_UINT64
	case SimpleType_BYTES:
		return STR_BYTES
	case SimpleType_STRING:
		return STR_STRING
	case SimpleType_SYMBOL:
		return STR_SYMBOL
	case SimpleType_BOOL:
		return STR_BOOL
	case SimpleType_FLOAT32:
		return STR_FLOAT32
	case SimpleType_FLOAT64:
		return STR_FLOAT64
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value: %d", t)
		return STR_UNKNOWN
	}
}

func (t *FixedPointType) ToVarint(name string) string {
	value := name
	if t.Floor > 0 {
		value = fmt.Sprintf("(%s - %d)", value, t.Floor)
	} else if t.Floor < 0 {
		value = fmt.Sprintf("(%s + %d)", value, -t.Floor)
	}
	if precision := pow10(t.Precision); precision > 0 {
		value = fmt.Sprintf("%s * %d", value, precision)
	}
	return value
}

func (t *FixedPointType) FromVarint(name string) string {
	value := fmt.Sprintf("float64(%s)", name)
	if precision := pow10(t.Precision); precision > 0 {
		value = fmt.Sprintf("%s / %d", value, precision)
	}
	if t.Floor > 0 {
		value = fmt.Sprintf("%s + %d", value, t.Floor)
	} else if t.Floor < 0 {
		value = fmt.Sprintf("%s - %d", value, -t.Floor)
	}
	return value
}

func (t *FixedPointType) String() string {
	return fmt.Sprintf("fixedpoint<%d, %d>", t.Precision, t.Floor)
}

func (t *EnumType) String() string {
	return t.Name
}

func (t *InstanceType) String() string {
	s := ""
	if t.IsPtr {
		s += "*"
	}
	if t.PkgPath != "" {
		s += t.PkgPath + "."
	}
	s += t.Name
	return s
}

func (t *ListType) String() string {
	return fmt.Sprintf("[]%s", t.E)
}

func (t *DictType) String() string {
	return fmt.Sprintf("map[%s]%s", t.K, t.V)
}

func (t *VariantType) String() string {
	var ts []string
	for _, t := range t.Ts {
		ts = append(ts, t.String())
	}
	return fmt.Sprintf("variant<%s>", strings.Join(ts, ", "))
}

func (t *Enum) Identifier() string {
	return t.Name
}

func (t *Object) Identifier() string {
	return t.Name
}

func (t UnknownType) Identifier() string {
	return string(t)
}

func (t SimpleType) Identifier() string {
	switch t {
	case SimpleType_NIL:
		return "n"
	case SimpleType_INT32:
		return "i"
	case SimpleType_INT64:
		return "i2"
	case SimpleType_UINT32:
		return "ui"
	case SimpleType_UINT64:
		return "ui2"
	case SimpleType_BYTES:
		return "by"
	case SimpleType_STRING:
		return "s"
	case SimpleType_SYMBOL:
		return "sy"
	case SimpleType_BOOL:
		return "b"
	case SimpleType_FLOAT32:
		return "f"
	case SimpleType_FLOAT64:
		return "d"
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Identifier: %d", t)
		return ""
	}
}

func (t *EnumType) Identifier() string {
	return t.Name
}

func (t *InstanceType) Identifier() string {
	return t.Name
}

func (t *FixedPointType) Identifier() string {
	if t.Floor >= 0 {
		return fmt.Sprintf("FixedPoint_%d_%d", t.Precision, t.Floor)
	} else {
		return fmt.Sprintf("FixedPoint_%d_s%d", t.Precision, -t.Floor)
	}
}

func (t *ListType) Identifier() string {
	return "L" + t.E.Identifier()
}

func (t *DictType) Identifier() string {
	return "D" + t.K.Identifier() + t.V.Identifier()
}

func (t *VariantType) Identifier() string {
	var names []string
	for _, v := range t.Ts {
		names = append(names, v.Identifier())
	}
	sort.Strings(names)
	return "V" + shortName(strings.Join(names, ""))
}
<commit_msg>add type error information<commit_after>// Copyright 2017 ibelie, Chen Jie, Joungtao. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package tygo

import (
	"fmt"
	"log"
	"reflect"
	"sort"
	"strings"
)

var TYGO_PATH = reflect.TypeOf(Tygo{}).PkgPath()

type Tygo struct {
	cachedSize int
}

func (t *Tygo) SetCachedSize(cachedSize int) {
	t.cachedSize = cachedSize
}

func (t *Tygo) CachedSize() int {
	return t.cachedSize
}

type Type interface {
	WireType() WireType
	Identifier() string
	String() string
	IsPrimitive() bool
	IsIterative() bool
	Go() (string, map[string]string)
	Typyd() string
	Python() string
	Typescript() string
	Javascript() (string, map[string]string)
	ByteSizeGo(string, string, string, int, bool) (string, map[string]string)
	CachedSizeGo(string, string, string, int, bool) (string, map[string]string)
	SerializeGo(string, string, string, int, bool) (string, map[string]string)
	DeserializeGo(string, string, string, string, int, bool) (string, WireType, map[string]string)
}

type Enum struct {
	nameMax int
	sorted  []string
	Name    string
	Package string
	Values  map[string]int
}

type Field struct {
	Type
	Name string
	Hide bool
}

type Method struct {
	Name    string
	Params  []Type
	Results []Type
}

type Object struct {
	Name    string
	Package string
	Parent  *InstanceType
	Fields  []*Field
	Methods []*Method
}

type UnknownType string

type SimpleType uint

const (
	SimpleType_UNKNOWN SimpleType = iota
	SimpleType_NIL
	SimpleType_INT32
	SimpleType_INT64
	SimpleType_UINT32
	SimpleType_UINT64
	SimpleType_BYTES
	SimpleType_STRING
	SimpleType_SYMBOL
	SimpleType_BOOL
	SimpleType_FLOAT32
	SimpleType_FLOAT64
)

type FixedPointType struct {
	Precision uint
	Floor     int
}

type EnumType struct {
	*Enum
	Name string
}

type InstanceType struct {
	*Object
	IsPtr   bool
	Name    string
	PkgName string
	PkgPath string
}

type ListType struct {
	E Type
}

type DictType struct {
	K Type
	V Type
}

type VariantType struct {
	Ts []Type
}

func (t *Enum) IsPrimitive() bool {
	return false
}

func (t *Object) IsPrimitive() bool {
	return false
}

func (t UnknownType) IsPrimitive() bool {
	return false
}

func (t SimpleType) IsPrimitive() bool {
	switch t {
	case SimpleType_INT32:
		fallthrough
	case SimpleType_INT64:
		fallthrough
	case SimpleType_UINT32:
		fallthrough
	case SimpleType_UINT64:
		fallthrough
	case SimpleType_BOOL:
		fallthrough
	case SimpleType_FLOAT32:
		fallthrough
	case SimpleType_FLOAT64:
		return true
	default:
		return false
	}
}

func (t *FixedPointType) IsPrimitive() bool {
	return true
}

func (t *EnumType) IsPrimitive() bool {
	return true
}

func (t *InstanceType) IsPrimitive() bool {
	return false
}

func (t *ListType) IsPrimitive() bool {
	return false
}

func (t *DictType) IsPrimitive() bool {
	return false
}

func (t *VariantType) IsPrimitive() bool {
	return false
}

func (t *Enum) IsIterative() bool {
	return false
}

func (t *Object) IsIterative() bool {
	return false
}

func (t UnknownType) IsIterative() bool {
	return false
}

func (t SimpleType) IsIterative() bool {
	return false
}

func (t *FixedPointType) IsIterative() bool {
	return false
}

func (t *EnumType) IsIterative() bool {
	return false
}

func (t *InstanceType) IsIterative() bool {
	return false
}

func (t *ListType) IsIterative() bool {
	return !t.E.IsPrimitive()
}

func (t *DictType) IsIterative() bool {
	return true
}

func (t *VariantType) IsIterative() bool {
	return false
}

func (t *Enum) FullName() string {
	return t.Package + "/" + t.Name
}

func (t *Enum) Sorted() []string {
	if t.sorted == nil || len(t.sorted) != len(t.Values) {
		t.sorted = nil
		for name, _ := range t.Values {
			if t.nameMax < len(name) {
				t.nameMax = len(name)
			}
			t.sorted = append(t.sorted, name)
		}
	} else if t.nameMax <= 0 {
		for name, _ := range t.Values {
			if t.nameMax < len(name) {
				t.nameMax = len(name)
			}
		}
	}
	if !sort.IsSorted(t) {
		sort.Sort(t)
	}
	return t.sorted
}

func (t *Enum) Len() int {
	return len(t.sorted)
}

func (t *Enum) Swap(i, j int) {
	t.sorted[i], t.sorted[j] = t.sorted[j], t.sorted[i]
}

func (t *Enum) Less(i, j int) bool {
	return t.Values[t.sorted[i]] < t.Values[t.sorted[j]]
}

func (t *Enum) String() string {
	var values []string
	for _, name := range t.Sorted() {
		values = append(values, fmt.Sprintf("%s: %d", name, t.Values[name]))
	}
	return fmt.Sprintf("%s[%s]", t.Name, strings.Join(values, ", "))
}

func (t *Method) String() string {
	var params []string
	for _, param := range t.Params {
		params = append(params, param.String())
	}
	s := fmt.Sprintf("%s(%s)", t.Name, strings.Join(params, ", "))

	var results []string
	for _, result := range t.Results {
		results = append(results, result.String())
	}
	if len(results) == 1 {
		s = fmt.Sprintf("%s %s", s, results[0])
	} else if len(results) > 1 {
		s = fmt.Sprintf("%s (%s)", s, strings.Join(results, ", "))
	}

	return s
}

func (t *Object) FullName() string {
	return t.Package + "/" + t.Name
}

func (t *Object) HasParent() bool {
	if t.Parent.Name == "Tygo" && t.Parent.PkgName == "tygo" && t.Parent.PkgPath == TYGO_PATH {
		return false
	} else if t.Parent.PkgPath != "" {
		return true
	} else {
		name, num := t.Parent.Object.MaxFieldNum()
		return name != "" || num != 0
	}
}

func (t *Object) MaxFieldNum() (string, int) {
	if !t.HasParent() {
		return "", len(t.VisibleFields())
	} else if t.Parent.PkgPath == "" {
		name, num := t.Parent.Object.MaxFieldNum()
		return name, num + len(t.VisibleFields())
	} else {
		return t.Parent.Name, len(t.VisibleFields())
	}
}

func (t *Object) VisibleFields() (fields []*Field) {
	for _, field := range t.Fields {
		if !field.Hide {
			fields = append(fields, field)
		}
	}
	return
}

func (t *Object) AllFields(objects map[string]*Object, visible bool) []*Field {
	fields := t.Fields
	if visible {
		fields = t.VisibleFields()
	}
	if !t.HasParent() {
		return fields
	} else if t.Parent.PkgPath == "" {
		return append(t.Parent.Object.AllFields(objects, visible), fields...)
	} else if object, ok := objects[t.Parent.Name]; ok {
		return append(object.AllFields(objects, visible), fields...)
	} else {
		log.Fatalf("[Tygo][Object] Unsolved parent type: %s", t.Parent.Name)
		return fields
	}
}

func (t *Object) String() string {
	var fields []string

	fields = append(fields, fmt.Sprintf(`
	%s`, t.Parent))

	nameMax := 0
	for _, field := range t.Fields {
		length := len(field.Name)
		if field.Hide {
			length++
		}
		if nameMax < length {
			nameMax = length
		}
	}
	for _, field := range t.Fields {
		name := field.Name
		if field.Hide {
			name = "." + field.Name
		}
		fields = append(fields, fmt.Sprintf(`
	%s %s%s`, name, strings.Repeat(" ", nameMax-len(name)), field))
	}
	for _, method := range t.Methods {
		fields = append(fields, fmt.Sprintf(`
	%s`, method))
	}

	return fmt.Sprintf(`
%s{%s
}
`, t.Name, strings.Join(fields, ""))
}

func (t UnknownType) String() string {
	return string(t)
}

func SimpleType_FromString(s string) Type {
	switch s {
	case STR_NIL:
		return SimpleType_NIL
	case STR_INT32:
		return SimpleType_INT32
	case STR_INT64:
		return SimpleType_INT64
	case STR_UINT32:
		return SimpleType_UINT32
	case STR_UINT64:
		return SimpleType_UINT64
	case STR_BYTES:
		return SimpleType_BYTES
	case STR_STRING:
		return SimpleType_STRING
	case STR_SYMBOL:
		return SimpleType_SYMBOL
	case STR_BOOL:
		return SimpleType_BOOL
	case STR_FLOAT32:
		return SimpleType_FLOAT32
	case STR_FLOAT64:
		return SimpleType_FLOAT64
	default:
		log.Printf("[Tygo][SimpleType] Unknown simple type: %s", s)
		return UnknownType(s)
	}
}

func (t SimpleType) String() string {
	switch t {
	case SimpleType_NIL:
		return STR_NIL
	case SimpleType_INT32:
		return STR_INT32
	case SimpleType_INT64:
		return STR_INT64
	case SimpleType_UINT32:
		return STR_UINT32
	case SimpleType_UINT64:
		return STR_UINT64
	case SimpleType_BYTES:
		return STR_BYTES
	case SimpleType_STRING:
		return STR_STRING
	case SimpleType_SYMBOL:
		return STR_SYMBOL
	case SimpleType_BOOL:
		return STR_BOOL
	case SimpleType_FLOAT32:
		return STR_FLOAT32
	case SimpleType_FLOAT64:
		return STR_FLOAT64
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value: %d", t)
		return STR_UNKNOWN
	}
}

func (t *FixedPointType) ToVarint(name string) string {
	value := name
	if t.Floor > 0 {
		value = fmt.Sprintf("(%s - %d)", value, t.Floor)
	} else if t.Floor < 0 {
		value = fmt.Sprintf("(%s + %d)", value, -t.Floor)
	}
	if precision := pow10(t.Precision); precision > 0 {
		value = fmt.Sprintf("%s * %d", value, precision)
	}
	return value
}

func (t *FixedPointType) FromVarint(name string) string {
	value := fmt.Sprintf("float64(%s)", name)
	if precision := pow10(t.Precision); precision > 0 {
		value = fmt.Sprintf("%s / %d", value, precision)
	}
	if t.Floor > 0 {
		value = fmt.Sprintf("%s + %d", value, t.Floor)
	} else if t.Floor < 0 {
		value = fmt.Sprintf("%s - %d", value, -t.Floor)
	}
	return value
}

func (t *FixedPointType) String() string {
	return fmt.Sprintf("fixedpoint<%d, %d>", t.Precision, t.Floor)
}

func (t *EnumType) String() string {
	return t.Name
}

func (t *InstanceType) String() string {
	s := ""
	if t.IsPtr {
		s += "*"
	}
	if t.PkgPath != "" {
		s += t.PkgPath + "."
	}
	s += t.Name
	return s
}

func (t *ListType) String() string {
	return fmt.Sprintf("[]%s", t.E)
}

func (t *DictType) String() string {
	return fmt.Sprintf("map[%s]%s", t.K, t.V)
}

func (t *VariantType) String() string {
	var ts []string
	for _, t := range t.Ts {
		ts = append(ts, t.String())
	}
	return fmt.Sprintf("variant<%s>", strings.Join(ts, ", "))
}

func (t *Enum) Identifier() string {
	return t.Name
}

func (t *Object) Identifier() string {
	return t.Name
}

func (t UnknownType) Identifier() string {
	return string(t)
}

func (t SimpleType) Identifier() string {
	switch t {
	case SimpleType_NIL:
		return "n"
	case SimpleType_INT32:
		return "i"
	case SimpleType_INT64:
		return "i2"
	case SimpleType_UINT32:
		return "ui"
	case SimpleType_UINT64:
		return "ui2"
	case SimpleType_BYTES:
		return "by"
	case SimpleType_STRING:
		return "s"
	case SimpleType_SYMBOL:
		return "sy"
	case SimpleType_BOOL:
		return "b"
	case SimpleType_FLOAT32:
		return "f"
	case SimpleType_FLOAT64:
		return "d"
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Identifier: %d", t)
		return ""
	}
}

func (t *EnumType) Identifier() string {
	return t.Name
}

func (t *InstanceType) Identifier() string {
	return t.Name
}

func (t *FixedPointType) Identifier() string {
	if t.Floor >= 0 {
		return fmt.Sprintf("FixedPoint_%d_%d", t.Precision, t.Floor)
	} else {
		return fmt.Sprintf("FixedPoint_%d_s%d", t.Precision, -t.Floor)
	}
}

func (t *ListType) Identifier() string {
	return "L" + t.E.Identifier()
}

func (t *DictType) Identifier() string {
	return "D" + t.K.Identifier() + t.V.Identifier()
}

func (t *VariantType) Identifier() string {
	var names []string
	for _, v := range t.Ts {
		names = append(names, v.Identifier())
	}
	sort.Strings(names)
	return "V" + shortName(strings.Join(names, ""))
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/oschwald/geoip2-golang"
	"github.com/tcnksm/go-httpstat"
	"github.com/yanc0/beeping/sslcheck"
)

var VERSION = "0.5.0"
var MESSAGE = "BeePing instance - HTTP Ping as a Service (github.com/yanc0/beeping)"
var USERAGENT = "Beeping " + VERSION + " - https://github.com/yanc0/beeping"

var geodatfile *string
var instance *string
var listen *string
var port *string
var tlsmode *bool
var validatetarget *bool

type ErrorMessage struct {
	Message string `json:"message"`
}

type Beeping struct {
	Version string `json:"version"`
	Message string `json:"message"`
}

// Check defines the check to do
type Check struct {
	URL      string        `json:"url" binding:"required"`
	Pattern  string        `json:"pattern"`
	Header   string        `json:"header"`
	Insecure bool          `json:"insecure"`
	Timeout  time.Duration `json:"timeout"`
	Auth     string        `json:"auth"`
}

type Timeline struct {
	NameLookup    int64 `json:"name_lookup"`
	Connect       int64 `json:"connect"`
	Pretransfer   int64 `json:"pretransfer"`
	StartTransfer int64 `json:"starttransfer"`
}

type Geo struct {
	Country string `json:"country"`
	City    string `json:"city,omitempty"`
	IP      string `json:"ip"`
}

// Response defines the response to bring back
type Response struct {
	HTTPStatus      string `json:"http_status"`
	HTTPStatusCode  int    `json:"http_status_code"`
	HTTPBodyPattern bool   `json:"http_body_pattern"`
	HTTPHeader      bool   `json:"http_header"`
	HTTPRequestTime int64  `json:"http_request_time"`

	InstanceName string `json:"instance_name"`

	DNSLookup        int64 `json:"dns_lookup"`
	TCPConnection    int64 `json:"tcp_connection"`
	TLSHandshake     int64 `json:"tls_handshake,omitempty"`
	ServerProcessing int64 `json:"server_processing"`
	ContentTransfer  int64 `json:"content_transfer"`

	Timeline *Timeline          `json:"timeline"`
	Geo      *Geo               `json:"geo,omitempty"`
	SSL      *sslcheck.CheckSSL `json:"ssl,omitempty"`
}

func NewErrorMessage(message string) *ErrorMessage {
	var response = ErrorMessage{}
	response.Message = message
	return &response
}

func InvalidJSONResponse() *ErrorMessage {
	return NewErrorMessage("Invalid JSON sent")
}

func NewResponse() *Response {
	var response = Response{}
	response.Timeline = &Timeline{}
	return &response
}

func NewCheck() *Check {
	return &Check{Timeout: 10}
}

// Performs some validation checks on the target.
// Returns nil if valid, returns an error otherwise.
func (check *Check) validateTarget() error {
	targetURL, err := url.Parse(check.URL)
	if err != nil {
		return err
	}
	ip := net.ParseIP(targetURL.Hostname())
	if ip == nil {
		// Hostname provided is not an IP. Without whitelisting, it is not possible to tell
		// whether it is an internal hostname.
		return nil // For now, hostnames are not needed for this check.
	}

	// Check for local network IPs
	switch {
	// Loopback address
	case ip.IsLoopback():
		return fmt.Errorf("Disallowed target")
	// Link-local unicast
	case ip.IsLinkLocalUnicast():
		return fmt.Errorf("Disallowed target")
	// Link-local multicast
	case ip.IsLinkLocalMulticast():
		return fmt.Errorf("Disallowed target")
	// Private network (10.0.0.0/8)
	case len(ip) == 4 && ip[0] == 10:
		return fmt.Errorf("Disallowed target")
	// Private network (Carrier-grade NAT; 100.64.0.0/10)
	case len(ip) == 4 && ip[0] == 100 && ip[1] >= 64 && ip[1] <= 127:
		return fmt.Errorf("Disallowed target")
	// Private network (172.16.0.0/12)
	case len(ip) == 4 && ip[0] == 172 && ip[1] >= 16 && ip[1] <= 31:
		return fmt.Errorf("Disallowed target")
	// Private network (192.168.0.0/16)
	case len(ip) == 4 && ip[0] == 192 && ip[1] == 16:
		return fmt.Errorf("Disallowed target")
	// Private network (fc00::/7)
	case len(ip) == 16 && (ip[0] == 0xfc || ip[0] == 0xfd):
		return fmt.Errorf("Disallowed target")
	}

	return nil
}

func main() {
	geodatfile = flag.String("geodatfile", "/opt/GeoIP/GeoLite2-City.mmdb", "geoIP database path")
	instance = flag.String("instance", "", "beeping instance name (default hostname)")
	listen = flag.String("listen", "127.0.0.1", "The host to bind the server to")
	port = flag.String("port", "8080", "The port to bind the server to")
	tlsmode = flag.Bool("tlsmode", false, "Activate SSL/TLS versions and Cipher support checks (slow)")
	validatetarget = flag.Bool("validatetarget", true, "Perform some security checks on the target provided")
	flag.Parse()

	http.HandleFunc("/check", handlerCheck)
	http.HandleFunc("/", handlerDefault)

	if err := http.ListenAndServe(*listen+":"+*port, nil); err != nil {
		log.Fatalf("could not listen on %s:%s: %v", *listen, *port, err)
	}
	log.Println("[INFO] Listening on", *listen, *port)
}

func handlerDefault(w http.ResponseWriter, r *http.Request) {
	var beeping Beeping
	beeping.Version = VERSION
	beeping.Message = MESSAGE
	log.Println("[INFO] Beeping version", beeping.Version)
	jsonRes, _ := json.Marshal(beeping)
	w.Header().Set("Content-Type", "application/json")
	w.Write(jsonRes)
}

func handlerCheck(w http.ResponseWriter, r *http.Request) {
	var check = NewCheck()

	w.Header().Set("Content-Type", "application/json")

	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&check)
	if err != nil {
		log.Println("[WARN] Invalid JSON sent")
		jsonRes, _ := json.Marshal(InvalidJSONResponse())
		w.Write(jsonRes)
		return
	}

	if !*validatetarget {
		response, err := CheckHTTP(check)
		if err != nil {
			log.Println("[WARN] Check failed:", err.Error())
			jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
			w.Write(jsonRes)
			return
		}
		log.Println("[INFO] Successful check:", check.URL, "-", response.HTTPRequestTime, "ms")
		jsonRes, _ := json.Marshal(response)
		w.Write(jsonRes)
		return
	}

	if err := check.validateTarget(); err != nil {
		log.Println("[WARN] Invalid target:", err.Error())
		jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
		w.Write(jsonRes)
		return
	}

	response, err := CheckHTTP(check)
	if err != nil {
		log.Println("[WARN] Check failed:", err.Error())
		jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
		w.Write(jsonRes)
		return
	}

	log.Println("[INFO] Successful check:", check.URL, "-", response.HTTPRequestTime, "ms")
	jsonRes, _ := json.Marshal(response)
	w.Write(jsonRes)
}

// CheckHTTP do HTTP check and return a beeping response
func CheckHTTP(check *Check) (*Response, error) {
	var response = NewResponse()
	var conn net.Conn

	req, err := http.NewRequest("GET", check.URL, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("User-Agent", USERAGENT)

	if len(check.Auth) > 0 {
		req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(check.Auth)))
	}

	// Create go-httpstat powered context and pass it to http.Request
	var result httpstat.Result
	ctx := httpstat.WithHTTPStat(req.Context(), &result)

	// Add IP:PORT tracing to the context
	ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
		GotConn: func(i httptrace.GotConnInfo) {
			conn = i.Conn
		},
	})

	req = req.WithContext(ctx)

	// DefaultClient is not suitable cause it caches
	// tcp connection https://golang.org/pkg/net/http/#Client
	// Allow us to close Idle connections and reset network
	// metrics each time
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: check.Insecure,
		},
	}

	timeout := time.Duration(check.Timeout * time.Second)

	client := &http.Client{
		Transport: tr,
		Timeout:   timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	res.Body.Close()
	timeEndBody := time.Now()
	result.End(timeEndBody)
	var total = result.Total(timeEndBody)

	tr.CloseIdleConnections()

	pattern := strings.Contains(string(body), check.Pattern)

	header := true
	if check.Header != "" {
		key, value := splitCheckHeader(check.Header)
		if key != "" && value != "" && res.Header.Get(key) != value {
			header = false
		}
	}

	response.HTTPStatus = res.Status
	response.HTTPStatusCode = res.StatusCode
	response.HTTPBodyPattern = pattern
	response.HTTPHeader = header
	response.HTTPRequestTime = milliseconds(total)
	response.Timeline.NameLookup = milliseconds(result.NameLookup)
	response.Timeline.Connect = milliseconds(result.Connect)
	response.Timeline.Pretransfer = milliseconds(result.Pretransfer)
	response.Timeline.StartTransfer = milliseconds(result.StartTransfer)
	response.DNSLookup = milliseconds(result.DNSLookup)
	response.TCPConnection = milliseconds(result.TCPConnection)
	response.TLSHandshake = milliseconds(result.TLSHandshake)
	response.ServerProcessing = milliseconds(result.ServerProcessing)
	response.ContentTransfer = milliseconds(result.ContentTransfer(timeEndBody))

	if res.TLS != nil {
		cTLS := &sslcheck.CheckSSL{}
		if *tlsmode {
			cTLS.CheckCiphers(conn)
			cTLS.CheckVersions(conn)
		}
		cTLS.CertExpiryDate = res.TLS.PeerCertificates[0].NotAfter
		cTLS.CertExpiryDaysLeft = int64(cTLS.CertExpiryDate.Sub(time.Now()).Hours() / 24)
		cTLS.CertSignature = res.TLS.PeerCertificates[0].SignatureAlgorithm.String()
		response.SSL = cTLS
	}

	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
	if err != nil {
		log.Println("[WARN] Cannot parse IP address", err.Error())
	}

	_ = geoIPCountry(*geodatfile, ip, response)

	err = instanceName(*instance, response)
	if err != nil {
		log.Println("[WARN] Cannot set instance name", err.Error())
	}

	return response, nil
}

func milliseconds(d time.Duration) int64 {
	return d.Nanoseconds() / 1000 / 1000
}

func geoIPCountry(geodatabase string, ip string, response *Response) error {
	db, err := geoip2.Open(geodatabase)
	if err != nil {
		return err
	}
	defer db.Close()
	// If you are using strings that may be invalid, check that ip is not nil
	ipParse := net.ParseIP(ip)
	record, err := db.City(ipParse)
	if err != nil {
		return err
	}
	response.Geo = &Geo{}
	response.Geo.Country = record.Country.IsoCode
	response.Geo.IP = ip
	if record.Country.Names != nil {
		response.Geo.City = record.City.Names["en-EN"]
	}
	return nil
}

func instanceName(name string, response *Response) error {
	var err error
	response.InstanceName = name
	if name == "" {
		response.InstanceName, err = os.Hostname()
		if err != nil {
			return err
		}
	}

	return nil
}

func splitCheckHeader(header string) (string, string) {
	h := strings.SplitN(header, ":", 2)
	if len(h) == 2 {
		return strings.TrimSpace(h[0]), strings.TrimSpace(h[1])
	}
	return "", ""
}
<commit_msg>Fix: Show "Listening on" message again<commit_after>package main

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/oschwald/geoip2-golang"
	"github.com/tcnksm/go-httpstat"
	"github.com/yanc0/beeping/sslcheck"
)

var VERSION = "0.5.0"
var MESSAGE = "BeePing instance - HTTP Ping as a Service (github.com/yanc0/beeping)"
var USERAGENT = "Beeping " + VERSION + " - https://github.com/yanc0/beeping"

var geodatfile *string
var instance *string
var listen *string
var port *string
var tlsmode *bool
var validatetarget *bool

type ErrorMessage struct {
	Message string `json:"message"`
}

type Beeping struct {
	Version string `json:"version"`
	Message string `json:"message"`
}

// Check defines the check to do
type Check struct {
	URL      string        `json:"url" binding:"required"`
	Pattern  string        `json:"pattern"`
	Header   string        `json:"header"`
	Insecure bool          `json:"insecure"`
	Timeout  time.Duration `json:"timeout"`
	Auth     string        `json:"auth"`
}

type Timeline struct {
	NameLookup    int64 `json:"name_lookup"`
	Connect       int64 `json:"connect"`
	Pretransfer   int64 `json:"pretransfer"`
	StartTransfer int64 `json:"starttransfer"`
}

type Geo struct {
	Country string `json:"country"`
	City    string `json:"city,omitempty"`
	IP      string `json:"ip"`
}

// Response defines the response to bring back
type Response struct {
	HTTPStatus      string `json:"http_status"`
	HTTPStatusCode  int    `json:"http_status_code"`
	HTTPBodyPattern bool   `json:"http_body_pattern"`
	HTTPHeader      bool   `json:"http_header"`
	HTTPRequestTime int64  `json:"http_request_time"`

	InstanceName string `json:"instance_name"`

	DNSLookup        int64 `json:"dns_lookup"`
	TCPConnection    int64 `json:"tcp_connection"`
	TLSHandshake     int64 `json:"tls_handshake,omitempty"`
	ServerProcessing int64 `json:"server_processing"`
	ContentTransfer  int64 `json:"content_transfer"`

	Timeline *Timeline          `json:"timeline"`
	Geo      *Geo               `json:"geo,omitempty"`
	SSL      *sslcheck.CheckSSL `json:"ssl,omitempty"`
}

func NewErrorMessage(message string) *ErrorMessage {
	var response = ErrorMessage{}
	response.Message = message
	return &response
}

func InvalidJSONResponse() *ErrorMessage {
	return NewErrorMessage("Invalid JSON sent")
}

func NewResponse() *Response {
	var response = Response{}
	response.Timeline = &Timeline{}
	return &response
}

func NewCheck() *Check {
	return &Check{Timeout: 10}
}

// Performs some validation checks on the target.
// Returns nil if valid, returns an error otherwise.
func (check *Check) validateTarget() error {
	targetURL, err := url.Parse(check.URL)
	if err != nil {
		return err
	}
	ip := net.ParseIP(targetURL.Hostname())
	if ip == nil {
		// Hostname provided is not an IP. Without whitelisting, it is not possible to tell
		// whether it is an internal hostname.
		return nil // For now, hostnames are not needed for this check.
	}

	// Check for local network IPs
	switch {
	// Loopback address
	case ip.IsLoopback():
		return fmt.Errorf("Disallowed target")
	// Link-local unicast
	case ip.IsLinkLocalUnicast():
		return fmt.Errorf("Disallowed target")
	// Link-local multicast
	case ip.IsLinkLocalMulticast():
		return fmt.Errorf("Disallowed target")
	// Private network (10.0.0.0/8)
	case len(ip) == 4 && ip[0] == 10:
		return fmt.Errorf("Disallowed target")
	// Private network (Carrier-grade NAT; 100.64.0.0/10)
	case len(ip) == 4 && ip[0] == 100 && ip[1] >= 64 && ip[1] <= 127:
		return fmt.Errorf("Disallowed target")
	// Private network (172.16.0.0/12)
	case len(ip) == 4 && ip[0] == 172 && ip[1] >= 16 && ip[1] <= 31:
		return fmt.Errorf("Disallowed target")
	// Private network (192.168.0.0/16)
	case len(ip) == 4 && ip[0] == 192 && ip[1] == 16:
		return fmt.Errorf("Disallowed target")
	// Private network (fc00::/7)
	case len(ip) == 16 && (ip[0] == 0xfc || ip[0] == 0xfd):
		return fmt.Errorf("Disallowed target")
	}

	return nil
}

func main() {
	geodatfile = flag.String("geodatfile", "/opt/GeoIP/GeoLite2-City.mmdb", "geoIP database path")
	instance = flag.String("instance", "", "beeping instance name (default hostname)")
	listen = flag.String("listen", "127.0.0.1", "The host to bind the server to")
	port = flag.String("port", "8080", "The port to bind the server to")
	tlsmode = flag.Bool("tlsmode", false, "Activate SSL/TLS versions and Cipher support checks (slow)")
	validatetarget = flag.Bool("validatetarget", true, "Perform some security checks on the target provided")
	flag.Parse()

	http.HandleFunc("/check", handlerCheck)
	http.HandleFunc("/", handlerDefault)

	log.Println("[INFO] Listening on", *listen, *port)
	if err := http.ListenAndServe(*listen+":"+*port, nil); err != nil {
		log.Fatalf("could not listen on %s:%s: %v", *listen, *port, err)
	}
}

func handlerDefault(w http.ResponseWriter, r *http.Request) {
	var beeping Beeping
	beeping.Version = VERSION
	beeping.Message = MESSAGE
	log.Println("[INFO] Beeping version", beeping.Version)
	jsonRes, _ := json.Marshal(beeping)
	w.Header().Set("Content-Type", "application/json")
	w.Write(jsonRes)
}

func handlerCheck(w http.ResponseWriter, r *http.Request) {
	var check = NewCheck()

	w.Header().Set("Content-Type", "application/json")

	decoder := json.NewDecoder(r.Body)
	err := decoder.Decode(&check)
	if err != nil {
		log.Println("[WARN] Invalid JSON sent")
		jsonRes, _ := json.Marshal(InvalidJSONResponse())
		w.Write(jsonRes)
		return
	}

	if !*validatetarget {
		response, err := CheckHTTP(check)
		if err != nil {
			log.Println("[WARN] Check failed:", err.Error())
			jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
			w.Write(jsonRes)
			return
		}
		log.Println("[INFO] Successful check:", check.URL, "-", response.HTTPRequestTime, "ms")
		jsonRes, _ := json.Marshal(response)
		w.Write(jsonRes)
		return
	}

	if err := check.validateTarget(); err != nil {
		log.Println("[WARN] Invalid target:", err.Error())
		jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
		w.Write(jsonRes)
		return
	}

	response, err := CheckHTTP(check)
	if err != nil {
		log.Println("[WARN] Check failed:", err.Error())
		jsonRes, _ := json.Marshal(NewErrorMessage(err.Error()))
		w.Write(jsonRes)
		return
	}

	log.Println("[INFO] Successful check:", check.URL, "-", response.HTTPRequestTime, "ms")
	jsonRes, _ := json.Marshal(response)
	w.Write(jsonRes)
}

// CheckHTTP do HTTP check and return a beeping response
func CheckHTTP(check *Check) (*Response, error) {
	var response = NewResponse()
	var conn net.Conn

	req, err := http.NewRequest("GET", check.URL, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("User-Agent", USERAGENT)

	if len(check.Auth) > 0 {
		req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(check.Auth)))
	}

	// Create go-httpstat powered context and pass it to http.Request
	var result httpstat.Result
	ctx := httpstat.WithHTTPStat(req.Context(), &result)

	// Add IP:PORT tracing to the context
	ctx = httptrace.WithClientTrace(ctx, &httptrace.ClientTrace{
		GotConn: func(i httptrace.GotConnInfo) {
			conn = i.Conn
		},
	})

	req = req.WithContext(ctx)

	// DefaultClient is not suitable cause it caches
	// tcp connection https://golang.org/pkg/net/http/#Client
	// Allow us to close Idle connections and reset network
	// metrics each time
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: check.Insecure,
		},
	}

	timeout := time.Duration(check.Timeout * time.Second)

	client := &http.Client{
		Transport: tr,
		Timeout:   timeout,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	res.Body.Close()
	timeEndBody := time.Now()
	result.End(timeEndBody)
	var total = result.Total(timeEndBody)

	tr.CloseIdleConnections()

	pattern := strings.Contains(string(body), check.Pattern)

	header := true
	if check.Header != "" {
		key, value := splitCheckHeader(check.Header)
		if key != "" && value != "" && res.Header.Get(key) != value {
			header = false
		}
	}

	response.HTTPStatus = res.Status
	response.HTTPStatusCode = res.StatusCode
	response.HTTPBodyPattern = pattern
	response.HTTPHeader = header
	response.HTTPRequestTime = milliseconds(total)
	response.Timeline.NameLookup = milliseconds(result.NameLookup)
	response.Timeline.Connect = milliseconds(result.Connect)
	response.Timeline.Pretransfer = milliseconds(result.Pretransfer)
	response.Timeline.StartTransfer = milliseconds(result.StartTransfer)
	response.DNSLookup = milliseconds(result.DNSLookup)
	response.TCPConnection = milliseconds(result.TCPConnection)
	response.TLSHandshake = milliseconds(result.TLSHandshake)
	response.ServerProcessing = milliseconds(result.ServerProcessing)
	response.ContentTransfer = milliseconds(result.ContentTransfer(timeEndBody))

	if res.TLS != nil {
		cTLS := &sslcheck.CheckSSL{}
		if *tlsmode {
			cTLS.CheckCiphers(conn)
			cTLS.CheckVersions(conn)
		}
		cTLS.CertExpiryDate = res.TLS.PeerCertificates[0].NotAfter
		cTLS.CertExpiryDaysLeft = int64(cTLS.CertExpiryDate.Sub(time.Now()).Hours() / 24)
		cTLS.CertSignature = res.TLS.PeerCertificates[0].SignatureAlgorithm.String()
		response.SSL = cTLS
	}

	ip, _, err := net.SplitHostPort(conn.RemoteAddr().String())
	if err != nil {
		log.Println("[WARN] Cannot parse IP address", err.Error())
	}

	_ = geoIPCountry(*geodatfile, ip, response)

	err = instanceName(*instance, response)
	if err != nil {
		log.Println("[WARN] Cannot set instance name", err.Error())
	}

	return response, nil
}

func milliseconds(d time.Duration) int64 {
	return d.Nanoseconds() / 1000 / 1000
}

func geoIPCountry(geodatabase string, ip string, response *Response) error {
	db, err := geoip2.Open(geodatabase)
	if err != nil {
		return err
	}
	defer db.Close()
	// If you are using strings that may be invalid, check that ip is not nil
	ipParse := net.ParseIP(ip)
	record, err := db.City(ipParse)
	if err != nil {
		return err
	}
	response.Geo = &Geo{}
	response.Geo.Country = record.Country.IsoCode
	response.Geo.IP = ip
	if record.Country.Names != nil {
		response.Geo.City = record.City.Names["en-EN"]
	}
	return nil
}

func instanceName(name string, response *Response) error {
	var err error
	response.InstanceName = name
	if name == "" {
		response.InstanceName, err = os.Hostname()
		if err != nil {
			return err
		}
	}

	return nil
}

func splitCheckHeader(header string) (string, string) {
	h := strings.SplitN(header, ":", 2)
	if len(h) == 2 {
		return strings.TrimSpace(h[0]), strings.TrimSpace(h[1])
	}
	return "", ""
}
<|endoftext|>