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

package main

import (
    "fmt"
    "time"
    "log"
    "crypto"
    "os"
    "bytes"
    "net/smtp"
    "net/textproto"
    "mime/multipart"
    "mime/quotedprintable"
    "html/template"
    "encoding/gob"
    "path/filepath"
    "net/http"
    "github.com/jmoiron/sqlx"
    _ "github.com/go-sql-driver/mysql"

    "golang.org/x/crypto/openpgp"
    "golang.org/x/crypto/openpgp/packet"
    "github.com/Sirupsen/logrus"
    "github.com/carbocation/interpose"
    "github.com/spf13/viper"
    "github.com/gorilla/mux"
    "github.com/gorilla/schema"
    "github.com/gorilla/sessions"
    "github.com/ubccr/goipa"
    "github.com/go-ini/ini"
)

const (
    MOKEY_COOKIE_SESSION = "mokey-session"
    MOKEY_COOKIE_SID     = "sid"
    MOKEY_COOKIE_USER    = "uid"
    MAX_PASS_LENGTH      = 8
)

type Application struct {
    dsn           string
    db            *sqlx.DB
    cookieStore   *sessions.CookieStore
    decoder       *schema.Decoder
    templates     map[string]*template.Template
    emails        map[string]*template.Template
    tmpldir       string
}

func NewDb() (*sqlx.DB, error) {
    db, err := sqlx.Open(viper.GetString("driver"), viper.GetString("dsn"))
    if err != nil {
        return nil, err
    }

    err = db.Ping()
    if err != nil {
        return nil, err
    }

    return db, nil
}

func NewApplication() (*Application, error) {
    db, err := NewDb()
    if err != nil {
        return nil, err
    }

    tmpldir := viper.GetString("templates")
    if len(tmpldir) == 0 {
        // default to directory of current executable
        path, err := filepath.EvalSymlinks(os.Args[0])
        if err != nil {
            log.Fatal(err)
        }
        dir, err := filepath.Abs(filepath.Dir(path))
        if err != nil {
            log.Fatal(err)
        }
        tmpldir = dir
    }

    logrus.Printf("Using template dir: %s", tmpldir)

    tmpls, err := filepath.Glob(tmpldir + "/templates/*.html")
    if err != nil {
        log.Fatal(err)
    }

    templates := make(map[string]*template.Template)
    for _, t := range tmpls {
        base := filepath.Base(t)
        if base != "layout.html" {
            templates[base] = template.Must(template.New("layout").ParseFiles(t,
                                                        tmpldir + "/templates/layout.html"))
        }
    }

    tmpls, err = filepath.Glob(tmpldir + "/templates/email/*.txt")
    if err != nil {
        log.Fatal(err)
    }

    emails := make(map[string]*template.Template)
    for _, t := range tmpls {
        base := filepath.Base(t)
        emails[base] = template.Must(template.New(base).ParseFiles(t))
    }

    app := &Application{}
    app.tmpldir = tmpldir
    app.db = db
    app.cookieStore = sessions.NewCookieStore([]byte(viper.GetString("secret")))
    app.decoder = schema.NewDecoder()
    //app.decoder.IgnoreUnknownKeys(true)
    app.templates = templates
    app.emails = emails

    return app, nil
}

func NewIpaClient(withKeytab bool) (*ipa.Client) {
    c := &ipa.Client{Host: viper.GetString("ipahost")}
    if withKeytab {
        c.KeyTab = viper.GetString("keytab")
    }

    return c
}

func quotedBody(body []byte) ([]byte, error) {
    var buf bytes.Buffer
    w := quotedprintable.NewWriter(&buf)
    _, err := w.Write(body)
    if err != nil {
        return nil, err
    }

    err = w.Close()
    if err != nil {
        return nil, err
    }

    return buf.Bytes(), nil
}

func sign(qtext []byte, header textproto.MIMEHeader) ([]byte, error) {
    var buf bytes.Buffer

    for k, vv := range header {
        for _, v := range vv {
            fmt.Fprintf(&buf, "%s: %s\r\n", k, v)
        }
    }
    fmt.Fprintf(&buf, "\r\n")
    _, err := buf.Write(qtext)
    if err != nil {
        return nil, err
    }

    file, err := os.Open(viper.GetString("pgp_key"))
    if err != nil {
        return nil, err
    }

    keyring, err := openpgp.ReadArmoredKeyRing(file)
    if err != nil {
        return nil, err
    }

    signingKey := keyring[0]

    if signingKey.PrivateKey.Encrypted {
        err = signingKey.PrivateKey.Decrypt([]byte(viper.GetString("pgp_passphrase")))
        if err != nil {
            return nil, err
        }
    }

    var sig bytes.Buffer
    err = openpgp.ArmoredDetachSign(&sig, signingKey, &buf, &packet.Config{DefaultHash: crypto.SHA256})
    if err != nil {
        return nil, err
    }

    return sig.Bytes(), nil
}

func (a *Application) SendEmail(email, subject, template string, data map[string]interface{}) (error) {
    logrus.WithFields(logrus.Fields{
        "email": email,
    }).Info("Sending email to user")

    if data == nil {
        data = make(map[string]interface{})
    }

    data["date"] = time.Now()
    data["contact"] = viper.GetString("email_from")
    data["sig"] = viper.GetString("email_sig")

    t := a.emails[template]
    var text bytes.Buffer
    err := t.ExecuteTemplate(&text, template, data)
    if err != nil {
        return err
    }

    qtext, err := quotedBody(text.Bytes())
    if err != nil {
        return err
    }

    header := make(textproto.MIMEHeader)
    header.Set("Mime-Version", "1.0")
    header.Set("Date", time.Now().Format(time.RFC1123Z))
    header.Set("To", email)
    header.Set("Subject", subject)
    header.Set("From", viper.GetString("email_from"))
    header.Set("Content-Type", "text/plain; charset=UTF-8")
    header.Set("Content-Transfer-Encoding", "quoted-printable")

    body := qtext

    if viper.GetBool("pgp_sign") {
        header.Del("Content-Transfer-Encoding")

        mhead := make(textproto.MIMEHeader)
        mhead.Add("Content-Type", "text/plain; charset=UTF-8")
        mhead.Add("Content-Transfer-Encoding", "quoted-printable")
        sig, err := sign(qtext, mhead)
        if err != nil {
            return err
        }

        var multipartBody bytes.Buffer
        mp := multipart.NewWriter(&multipartBody)
        boundary := mp.Boundary()
        mw, err := mp.CreatePart(mhead)
        if err != nil {
            return err
        }
        _, err = mw.Write(qtext)
        if err != nil {
            return err
        }

        mw, err = mp.CreatePart(textproto.MIMEHeader(
            map[string][]string {
                "Content-Type": []string{"application/pgp-signature; name=signature.asc;"},
            }))
        if err != nil {
            return err
        }

        _, err = mw.Write(sig)
        if err != nil {
            return err
        }

        err = mp.Close()
        if err != nil {
            return err
        }

        header.Set("Content-Type", fmt.Sprintf(`multipart/signed; boundary="%s"; protocol="application/pgp-signature"; micalg="pgp-sha256"`, boundary))
        body = multipartBody.Bytes()
    }

    c, err := smtp.Dial(fmt.Sprintf("%s:%d",viper.GetString("smtp_host"), viper.GetInt("smtp_port")))
    if err != nil {
        return err
    }
    defer c.Close()

    c.Mail(viper.GetString("email_from"))
    c.Rcpt(email)

    wc, err := c.Data()
    if err != nil {
        return err
    }
    defer wc.Close()

    var buf bytes.Buffer
    for k, vv := range header {
        for _, v := range vv {
            fmt.Fprintf(&buf, "%s: %s\r\n", k, v)
        }
    }
    fmt.Fprintf(&buf, "\r\n")

    if _, err = buf.WriteTo(wc); err != nil {
        return err
    }
    if _, err = wc.Write(body); err != nil {
        return err
    }

    return nil
}

func (a *Application) middlewareStruct() (*interpose.Middleware, error) {
    mw := interpose.New()
    mw.Use(Nosurf())

    mw.UseHandler(a.router())

    return mw, nil
}

func (a *Application) router() *mux.Router {
    router := mux.NewRouter()

    router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(http.StatusNotFound)
        renderTemplate(w, a.templates["404.html"], nil)
    })
    router.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir(fmt.Sprintf("%s/static", a.tmpldir)))))
    router.Path("/auth/login").Handler(RateLimit(a, LoginHandler(a))).Methods("GET", "POST")
    router.Path("/auth/logout").Handler(LogoutHandler(a)).Methods("GET")
    router.Path("/auth/forgotpw").Handler(RateLimit(a, ForgotPasswordHandler(a))).Methods("GET", "POST")
    router.Path("/auth/setup/{token:[0-9a-f]+}").Handler(SetupAccountHandler(a)).Methods("GET", "POST")
    router.Path("/auth/resetpw/{token:[0-9a-f]+}").Handler(ResetPasswordHandler(a)).Methods("GET", "POST")
    router.Path("/changepw").Handler(AuthRequired(a, ChangePasswordHandler(a))).Methods("GET", "POST")
    router.Path("/updatesec").Handler(AuthRequired(a, UpdateSecurityQuestionHandler(a))).Methods("GET", "POST")
    router.Path("/").Handler(AuthRequired(a, IndexHandler(a))).Methods("GET")

    return router
}

func init() {
    viper.SetDefault("port", 8080)
    viper.SetDefault("pgp_sign", false)
    viper.SetDefault("smtp_host", "localhost")
    viper.SetDefault("smtp_port", 25)
    viper.SetDefault("email_link_base", "http://localhost")
    viper.SetDefault("email_from", "helpdesk@example.com")
    viper.SetDefault("email_prefix", "mokey")
    viper.SetDefault("setup_max_age", 86400)
    viper.SetDefault("reset_max_age", 3600)
    viper.SetDefault("max_attempts", 10)
    viper.SetDefault("bind", "")
    viper.SetDefault("secret", "change-me")
    viper.SetDefault("driver", "mysql")
    viper.SetDefault("dsn", "/mokey?parseTime=true")
    viper.SetDefault("redis", ":6379")
    viper.SetDefault("rate_limit", "15")
    viper.SetDefault("rate_expire", "3600")

    gob.Register(&ipa.UserRecord{})
    gob.Register(&ipa.IpaDateTime{})

    if !viper.IsSet("ipahost") {
        cfg, err := ini.Load("/etc/ipa/default.conf")
        if err != nil {
            logrus.WithFields(logrus.Fields{
                "error": err,
            }).Warn("Failed parse /etc/ipa/default.conf. Defaulting to localhost")
            viper.SetDefault("ipahost", "localhost")
            return
        }

        ipaServer, err := cfg.Section("global").GetKey("server")
        if err != nil {
            logrus.WithFields(logrus.Fields{
                "error": err,
            }).Warn("Failed read ipahost from /etc/ipa/default.conf. Defaulting to localhost")
            viper.SetDefault("ipahost", "localhost")
            return
        }

        viper.SetDefault("ipahost", ipaServer)
    }
}

func Server() {
    app, err := NewApplication()
    if err != nil {
        logrus.Fatal(err.Error())
    }

    middle, err := app.middlewareStruct()
    if err != nil {
        logrus.Fatal(err.Error())
    }

    logrus.Printf("Running on http://%s:%d", viper.GetString("bind"), viper.GetInt("port"))

    http.Handle("/", middle)

    certFile := viper.GetString("cert")
    keyFile := viper.GetString("key")

    if certFile != "" && keyFile != "" {
        http.ListenAndServeTLS(fmt.Sprintf("%s:%d", viper.GetString("bind"), viper.GetInt("port")), certFile, keyFile, nil)
    } else {
        http.ListenAndServe(fmt.Sprintf("%s:%d", viper.GetString("bind"), viper.GetInt("port")), nil)
    }
}
<commit_msg>Fix parsing of template dir<commit_after>// Copyright 2015 mokey Authors. All rights reserved.
// Use of this source code is governed by a BSD style
// license that can be found in the LICENSE file.

package main

import (
    "fmt"
    "time"
    "log"
    "crypto"
    "os"
    "bytes"
    "net/smtp"
    "net/textproto"
    "mime/multipart"
    "mime/quotedprintable"
    "html/template"
    "encoding/gob"
    "path/filepath"
    "net/http"
    "github.com/jmoiron/sqlx"
    _ "github.com/go-sql-driver/mysql"

    "golang.org/x/crypto/openpgp"
    "golang.org/x/crypto/openpgp/packet"
    "github.com/Sirupsen/logrus"
    "github.com/carbocation/interpose"
    "github.com/spf13/viper"
    "github.com/gorilla/mux"
    "github.com/gorilla/schema"
    "github.com/gorilla/sessions"
    "github.com/ubccr/goipa"
    "github.com/go-ini/ini"
)

const (
    MOKEY_COOKIE_SESSION = "mokey-session"
    MOKEY_COOKIE_SID     = "sid"
    MOKEY_COOKIE_USER    = "uid"
    MAX_PASS_LENGTH      = 8
)

type Application struct {
    dsn           string
    db            *sqlx.DB
    cookieStore   *sessions.CookieStore
    decoder       *schema.Decoder
    templates     map[string]*template.Template
    emails        map[string]*template.Template
    tmpldir       string
}

func NewDb() (*sqlx.DB, error) {
    db, err := sqlx.Open(viper.GetString("driver"), viper.GetString("dsn"))
    if err != nil {
        return nil, err
    }

    err = db.Ping()
    if err != nil {
        return nil, err
    }

    return db, nil
}

func NewApplication() (*Application, error) {
    db, err := NewDb()
    if err != nil {
        return nil, err
    }

    tmpldir := viper.GetString("templates")
    if len(tmpldir) == 0 {
        // default to directory of current executable
        path, err := filepath.EvalSymlinks(os.Args[0])
        if err != nil {
            log.Fatal(err)
        }
        dir, err := filepath.Abs(filepath.Dir(path))
        if err != nil {
            log.Fatal(err)
        }
        tmpldir = dir + "/templates"
    }

    logrus.Printf("Using template dir: %s", tmpldir)

    tmpls, err := filepath.Glob(tmpldir + "/*.html")
    if err != nil {
        log.Fatal(err)
    }

    templates := make(map[string]*template.Template)
    for _, t := range tmpls {
        base := filepath.Base(t)
        if base != "layout.html" {
            templates[base] = template.Must(template.New("layout").ParseFiles(t,
                                                        tmpldir + "/layout.html"))
        }
    }

    tmpls, err = filepath.Glob(tmpldir + "/email/*.txt")
    if err != nil {
        log.Fatal(err)
    }

    emails := make(map[string]*template.Template)
    for _, t := range tmpls {
        base := filepath.Base(t)
        emails[base] = template.Must(template.New(base).ParseFiles(t))
    }

    app := &Application{}
    app.tmpldir = tmpldir
    app.db = db
    app.cookieStore = sessions.NewCookieStore([]byte(viper.GetString("secret")))
    app.decoder = schema.NewDecoder()
    //app.decoder.IgnoreUnknownKeys(true)
    app.templates = templates
    app.emails = emails

    return app, nil
}

func NewIpaClient(withKeytab bool) (*ipa.Client) {
    c := &ipa.Client{Host: viper.GetString("ipahost")}
    if withKeytab {
        c.KeyTab = viper.GetString("keytab")
    }

    return c
}

func quotedBody(body []byte) ([]byte, error) {
    var buf bytes.Buffer
    w := quotedprintable.NewWriter(&buf)
    _, err := w.Write(body)
    if err != nil {
        return nil, err
    }

    err = w.Close()
    if err != nil {
        return nil, err
    }

    return buf.Bytes(), nil
}

func sign(qtext []byte, header textproto.MIMEHeader) ([]byte, error) {
    var buf bytes.Buffer

    for k, vv := range header {
        for _, v := range vv {
            fmt.Fprintf(&buf, "%s: %s\r\n", k, v)
        }
    }
    fmt.Fprintf(&buf, "\r\n")
    _, err := buf.Write(qtext)
    if err != nil {
        return nil, err
    }

    file, err := os.Open(viper.GetString("pgp_key"))
    if err != nil {
        return nil, err
    }

    keyring, err := openpgp.ReadArmoredKeyRing(file)
    if err != nil {
        return nil, err
    }

    signingKey := keyring[0]

    if signingKey.PrivateKey.Encrypted {
        err = signingKey.PrivateKey.Decrypt([]byte(viper.GetString("pgp_passphrase")))
        if err != nil {
            return nil, err
        }
    }

    var sig bytes.Buffer
    err = openpgp.ArmoredDetachSign(&sig, signingKey, &buf, &packet.Config{DefaultHash: crypto.SHA256})
    if err != nil {
        return nil, err
    }

    return sig.Bytes(), nil
}

func (a *Application) SendEmail(email, subject, template string, data map[string]interface{}) (error) {
    logrus.WithFields(logrus.Fields{
        "email": email,
    }).Info("Sending email to user")

    if data == nil {
        data = make(map[string]interface{})
    }

    data["date"] = time.Now()
    data["contact"] = viper.GetString("email_from")
    data["sig"] = viper.GetString("email_sig")

    t := a.emails[template]
    var text bytes.Buffer
    err := t.ExecuteTemplate(&text, template, data)
    if err != nil {
        return err
    }

    qtext, err := quotedBody(text.Bytes())
    if err != nil {
        return err
    }

    header := make(textproto.MIMEHeader)
    header.Set("Mime-Version", "1.0")
    header.Set("Date", time.Now().Format(time.RFC1123Z))
    header.Set("To", email)
    header.Set("Subject", subject)
    header.Set("From", viper.GetString("email_from"))
    header.Set("Content-Type", "text/plain; charset=UTF-8")
    header.Set("Content-Transfer-Encoding", "quoted-printable")

    body := qtext

    if viper.GetBool("pgp_sign") {
        header.Del("Content-Transfer-Encoding")

        mhead := make(textproto.MIMEHeader)
        mhead.Add("Content-Type", "text/plain; charset=UTF-8")
        mhead.Add("Content-Transfer-Encoding", "quoted-printable")
        sig, err := sign(qtext, mhead)
        if err != nil {
            return err
        }

        var multipartBody bytes.Buffer
        mp := multipart.NewWriter(&multipartBody)
        boundary := mp.Boundary()
        mw, err := mp.CreatePart(mhead)
        if err != nil {
            return err
        }
        _, err = mw.Write(qtext)
        if err != nil {
            return err
        }

        mw, err = mp.CreatePart(textproto.MIMEHeader(
            map[string][]string {
                "Content-Type": []string{"application/pgp-signature; name=signature.asc;"},
            }))
        if err != nil {
            return err
        }

        _, err = mw.Write(sig)
        if err != nil {
            return err
        }

        err = mp.Close()
        if err != nil {
            return err
        }

        header.Set("Content-Type", fmt.Sprintf(`multipart/signed; boundary="%s"; protocol="application/pgp-signature"; micalg="pgp-sha256"`, boundary))
        body = multipartBody.Bytes()
    }

    c, err := smtp.Dial(fmt.Sprintf("%s:%d",viper.GetString("smtp_host"), viper.GetInt("smtp_port")))
    if err != nil {
        return err
    }
    defer c.Close()

    c.Mail(viper.GetString("email_from"))
    c.Rcpt(email)

    wc, err := c.Data()
    if err != nil {
        return err
    }
    defer wc.Close()

    var buf bytes.Buffer
    for k, vv := range header {
        for _, v := range vv {
            fmt.Fprintf(&buf, "%s: %s\r\n", k, v)
        }
    }
    fmt.Fprintf(&buf, "\r\n")

    if _, err = buf.WriteTo(wc); err != nil {
        return err
    }
    if _, err = wc.Write(body); err != nil {
        return err
    }

    return nil
}

func (a *Application) middlewareStruct() (*interpose.Middleware, error) {
    mw := interpose.New()
    mw.Use(Nosurf())

    mw.UseHandler(a.router())

    return mw, nil
}

func (a *Application) router() *mux.Router {
    router := mux.NewRouter()

    router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(http.StatusNotFound)
        renderTemplate(w, a.templates["404.html"], nil)
    })
    router.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir(fmt.Sprintf("%s/static", a.tmpldir)))))
    router.Path("/auth/login").Handler(RateLimit(a, LoginHandler(a))).Methods("GET", "POST")
    router.Path("/auth/logout").Handler(LogoutHandler(a)).Methods("GET")
    router.Path("/auth/forgotpw").Handler(RateLimit(a, ForgotPasswordHandler(a))).Methods("GET", "POST")
    router.Path("/auth/setup/{token:[0-9a-f]+}").Handler(SetupAccountHandler(a)).Methods("GET", "POST")
    router.Path("/auth/resetpw/{token:[0-9a-f]+}").Handler(ResetPasswordHandler(a)).Methods("GET", "POST")
    router.Path("/changepw").Handler(AuthRequired(a, ChangePasswordHandler(a))).Methods("GET", "POST")
    router.Path("/updatesec").Handler(AuthRequired(a, UpdateSecurityQuestionHandler(a))).Methods("GET", "POST")
    router.Path("/").Handler(AuthRequired(a, IndexHandler(a))).Methods("GET")

    return router
}

func init() {
    viper.SetDefault("port", 8080)
    viper.SetDefault("pgp_sign", false)
    viper.SetDefault("smtp_host", "localhost")
    viper.SetDefault("smtp_port", 25)
    viper.SetDefault("email_link_base", "http://localhost")
    viper.SetDefault("email_from", "helpdesk@example.com")
    viper.SetDefault("email_prefix", "mokey")
    viper.SetDefault("setup_max_age", 86400)
    viper.SetDefault("reset_max_age", 3600)
    viper.SetDefault("max_attempts", 10)
    viper.SetDefault("bind", "")
    viper.SetDefault("secret", "change-me")
    viper.SetDefault("driver", "mysql")
    viper.SetDefault("dsn", "/mokey?parseTime=true")
    viper.SetDefault("redis", ":6379")
    viper.SetDefault("rate_limit", "15")
    viper.SetDefault("rate_expire", "3600")

    gob.Register(&ipa.UserRecord{})
    gob.Register(&ipa.IpaDateTime{})

    if !viper.IsSet("ipahost") {
        cfg, err := ini.Load("/etc/ipa/default.conf")
        if err != nil {
            logrus.WithFields(logrus.Fields{
                "error": err,
            }).Warn("Failed parse /etc/ipa/default.conf. Defaulting to localhost")
            viper.SetDefault("ipahost", "localhost")
            return
        }

        ipaServer, err := cfg.Section("global").GetKey("server")
        if err != nil {
            logrus.WithFields(logrus.Fields{
                "error": err,
            }).Warn("Failed read ipahost from /etc/ipa/default.conf. Defaulting to localhost")
            viper.SetDefault("ipahost", "localhost")
            return
        }

        viper.SetDefault("ipahost", ipaServer)
    }
}

func Server() {
    app, err := NewApplication()
    if err != nil {
        logrus.Fatal(err.Error())
    }

    middle, err := app.middlewareStruct()
    if err != nil {
        logrus.Fatal(err.Error())
    }

    logrus.Printf("Running on http://%s:%d", viper.GetString("bind"), viper.GetInt("port"))

    http.Handle("/", middle)

    certFile := viper.GetString("cert")
    keyFile := viper.GetString("key")

    if certFile != "" && keyFile != "" {
        http.ListenAndServeTLS(fmt.Sprintf("%s:%d", viper.GetString("bind"), viper.GetInt("port")), certFile, keyFile, nil)
    } else {
        http.ListenAndServe(fmt.Sprintf("%s:%d", viper.GetString("bind"), viper.GetInt("port")), nil)
    }
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime/pprof"
)

func main() {
	var cpuprofile = flag.String("cpuprofile", "", "write CPU profile to file")
	var listen = flag.String("listen", ":8000", "listen address")
	var path = flag.String("path", "/tmp/restic", "data directory")
	var tls = flag.Bool("tls", false, "turn on TLS support")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		log.Println("CPU profiling enabled")
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	log.Println("Creating repository directories")
	dirs := []string{
		"data",
		"index",
		"keys",
		"locks",
		"snapshots",
		"tmp",
	}
	for _, d := range dirs {
		os.MkdirAll(filepath.Join(*path, d), 0700)
	}
	for i := 0; i < 256; i++ {
		os.MkdirAll(filepath.Join(*path, "data", fmt.Sprintf("%02x", i)), 0700)
	}

	context := &Context{*path}
	router := NewRouter()
	router.HeadFunc("/config", CheckConfig(context))
	router.GetFunc("/config", GetConfig(context))
	router.PostFunc("/config", SaveConfig(context))
	router.GetFunc("/:dir/", ListBlobs(context))
	router.HeadFunc("/:dir/:name", CheckBlob(context))
	router.GetFunc("/:type/:name", GetBlob(context))
	router.PostFunc("/:type/:name", SaveBlob(context))
	router.DeleteFunc("/:type/:name", DeleteBlob(context))

	var handler http.Handler
	htpasswdFile, err := NewHtpasswdFromFile(filepath.Join(*path, ".htpasswd"))
	if err != nil {
		log.Println("Authentication disabled")
		handler = router
	} else {
		log.Println("Authentication enabled")
		handler = AuthHandler(htpasswdFile, router)
	}

	if !*tls {
		log.Printf("Starting server on %s\n", *listen)
		err = http.ListenAndServe(*listen, handler)
	} else {
		privateKey := filepath.Join(*path, "private_key")
		publicKey := filepath.Join(*path, "public_key")
		log.Println("TLS enabled")
		log.Printf("Private key: %s", privateKey)
		log.Printf("Public key: %s", publicKey)
		log.Printf("Starting server on %s\n", *listen)
		err = http.ListenAndServeTLS(*listen, publicKey, privateKey, handler)
	}
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Stop logging date/time<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime/pprof"
)

func main() {
	log.SetFlags(0)

	var cpuprofile = flag.String("cpuprofile", "", "write CPU profile to file")
	var listen = flag.String("listen", ":8000", "listen address")
	var path = flag.String("path", "/tmp/restic", "data directory")
	var tls = flag.Bool("tls", false, "turn on TLS support")
	flag.Parse()

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		log.Println("CPU profiling enabled")
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	log.Println("Creating repository directories")
	dirs := []string{
		"data",
		"index",
		"keys",
		"locks",
		"snapshots",
		"tmp",
	}
	for _, d := range dirs {
		os.MkdirAll(filepath.Join(*path, d), 0700)
	}
	for i := 0; i < 256; i++ {
		os.MkdirAll(filepath.Join(*path, "data", fmt.Sprintf("%02x", i)), 0700)
	}

	context := &Context{*path}
	router := NewRouter()
	router.HeadFunc("/config", CheckConfig(context))
	router.GetFunc("/config", GetConfig(context))
	router.PostFunc("/config", SaveConfig(context))
	router.GetFunc("/:dir/", ListBlobs(context))
	router.HeadFunc("/:dir/:name", CheckBlob(context))
	router.GetFunc("/:type/:name", GetBlob(context))
	router.PostFunc("/:type/:name", SaveBlob(context))
	router.DeleteFunc("/:type/:name", DeleteBlob(context))

	var handler http.Handler
	htpasswdFile, err := NewHtpasswdFromFile(filepath.Join(*path, ".htpasswd"))
	if err != nil {
		log.Println("Authentication disabled")
		handler = router
	} else {
		log.Println("Authentication enabled")
		handler = AuthHandler(htpasswdFile, router)
	}

	if !*tls {
		log.Printf("Starting server on %s\n", *listen)
		err = http.ListenAndServe(*listen, handler)
	} else {
		privateKey := filepath.Join(*path, "private_key")
		publicKey := filepath.Join(*path, "public_key")
		log.Println("TLS enabled")
		log.Printf("Private key: %s", privateKey)
		log.Printf("Public key: %s", publicKey)
		log.Printf("Starting server on %s\n", *listen)
		err = http.ListenAndServeTLS(*listen, publicKey, privateKey, handler)
	}
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>/*
Server runs the http interface for this link shortener.

Author: Danver Braganza
*/

package main

import (
	"html/template"
	"log"
	"net/http"

	"github.com/gorilla/mux"

	"github.com/danverbraganza/shortlink/fetcher"
	"github.com/danverbraganza/shortlink/shortcut"
)

// ShortcutHandler is a struct that handles all requests.
type ShortcutHandler struct {
	index        shortcut.Index
	formTemplate *template.Template
}

// Get looks up the given shortcut requested in the index.
func (s ShortcutHandler) Get(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	query := vars["shortform"]

	results, sole, err := s.index.FindShortcut(query)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if sole {
		http.Redirect(w, r, results[0].URL, http.StatusSeeOther)
		return
	}
	s.ShowForm(w, r, results)
}

// ShowForm shows a nice form where the user can enter a new url.
func (s ShortcutHandler) ShowForm(w http.ResponseWriter, r *http.Request, partMatch []shortcut.Shortcut) {
	vars := mux.Vars(r)
	sf := vars["shortform"]
	s.formTemplate.Execute(w, struct {
		ShortForm string
		Shortcuts []shortcut.Shortcut
	}{sf, partMatch})
}

// ShowEmptyForm specifically shows an empty form.
func (s ShortcutHandler) ShowEmptyForm(w http.ResponseWriter, r *http.Request) {
	s.ShowForm(w, r, nil)
}

// Post saves the "first" url found as the "first" shortform found.
// You may pass shortforms as form params or in the url.
func (s ShortcutHandler) Post(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	r.ParseForm()

	shortform, ok := vars["shortform"]

	if !ok {
		// Shortform not provided in url: find it in the form.
		sc := r.Form["shortform"]
		if len(sc) == 0 {
			http.Error(w, "Shortcut was not supplied", http.StatusBadRequest)
			return
		}
		shortform = sc[0]
	}

	if _, sole, _ := s.index.FindShortcut(shortform); sole {
		http.Error(w, "Shortcut already exists.", http.StatusBadRequest)
		return
	}

	urls := r.Form["url"]
	if len(urls) == 0 {
		http.Error(w, "URL was not supplied", http.StatusBadRequest)
		return
	}
	url := urls[0]

	var description string
	if descriptions, ok := r.Form["description"]; ok {
		description = descriptions[0]
	}
	normalizedURL := shortcut.NormalizeURL(url)
	http.Redirect(w, r, normalizedURL, http.StatusSeeOther)
	go func() {
		if description == "" && r.Form["attempt"] != nil {
			description = fetcher.FindDescription(normalizedURL)
		}
		log.Print("Setting ", shortform, " to ", normalizedURL, ": ", description)
		s.index.AddShortcut(shortcut.Shortcut{
			ShortForm:   shortform,
			URL:         normalizedURL,
			Description: description,
		})
	}()
}

func main() {
	handler := ShortcutHandler{
		shortcut.NewIndex("links.bleve"),
		template.Must(template.ParseFiles("templates/form.tmpl")),
	}
	r := mux.NewRouter()

	r.HandleFunc("/index.html", handler.ShowEmptyForm)
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/index.html", http.StatusMovedPermanently)
	}).Methods("GET")

	r.HandleFunc("/{shortform:.*}", handler.Get).Methods("GET")
	r.HandleFunc("/", handler.Post).Methods("POST")
	r.HandleFunc("/{shortform:.*}", handler.Post).Methods("POST")
	http.Handle("/", r)

	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>Adding flags to server.go<commit_after>/*
Server runs the http interface for this link shortener.

Author: Danver Braganza
*/

package main

import (
	"flag"
	"html/template"
	"log"
	"net/http"
	"path"
	"strconv"

	"github.com/gorilla/mux"

	"github.com/danverbraganza/shortlink/fetcher"
	"github.com/danverbraganza/shortlink/shortcut"
)

// ShortcutHandler is a struct that handles all requests.
type ShortcutHandler struct {
	index        shortcut.Index
	formTemplate *template.Template
}

// Get looks up the given shortcut requested in the index.
func (s ShortcutHandler) Get(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	query := vars["shortform"]

	results, sole, err := s.index.FindShortcut(query)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if sole {
		http.Redirect(w, r, results[0].URL, http.StatusSeeOther)
		return
	}
	s.ShowForm(w, r, results)
}

// ShowForm shows a nice form where the user can enter a new url.
func (s ShortcutHandler) ShowForm(w http.ResponseWriter, r *http.Request, partMatch []shortcut.Shortcut) {
	vars := mux.Vars(r)
	sf := vars["shortform"]
	s.formTemplate.Execute(w, struct {
		ShortForm string
		Shortcuts []shortcut.Shortcut
	}{sf, partMatch})
}

// ShowEmptyForm specifically shows an empty form.
func (s ShortcutHandler) ShowEmptyForm(w http.ResponseWriter, r *http.Request) {
	s.ShowForm(w, r, nil)
}

// Post saves the "first" url found as the "first" shortform found.
// You may pass shortforms as form params or in the url.
func (s ShortcutHandler) Post(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	r.ParseForm()

	shortform, ok := vars["shortform"]

	if !ok {
		// Shortform not provided in url: find it in the form.
		sc := r.Form["shortform"]
		if len(sc) == 0 {
			http.Error(w, "Shortcut was not supplied", http.StatusBadRequest)
			return
		}
		shortform = sc[0]
	}

	if _, sole, _ := s.index.FindShortcut(shortform); sole {
		http.Error(w, "Shortcut already exists.", http.StatusBadRequest)
		return
	}

	urls := r.Form["url"]
	if len(urls) == 0 {
		http.Error(w, "URL was not supplied", http.StatusBadRequest)
		return
	}
	url := urls[0]

	var description string
	if descriptions, ok := r.Form["description"]; ok {
		description = descriptions[0]
	}
	normalizedURL := shortcut.NormalizeURL(url)
	http.Redirect(w, r, normalizedURL, http.StatusSeeOther)
	go func() {
		if description == "" && r.Form["attempt"] != nil {
			description = fetcher.FindDescription(normalizedURL)
		}
		log.Print("Setting ", shortform, " to ", normalizedURL, ": ", description)
		s.index.AddShortcut(shortcut.Shortcut{
			ShortForm:   shortform,
			URL:         normalizedURL,
			Description: description,
		})
	}()
}

func main() {
	indexFile := flag.String("indexfile", "links.bleve", "The location of the index file.")
	templateDir := flag.String("templates", "templates", "The location of the templates directory.")
	port := flag.Int("port", 8080, "The port to which to bind.")
	flag.Parse()

	handler := ShortcutHandler{
		shortcut.NewIndex(*indexFile),
		template.Must(template.ParseFiles(path.Join(*templateDir, "form.tmpl"))),
	}
	r := mux.NewRouter()

	r.HandleFunc("/index.html", handler.ShowEmptyForm)
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/index.html", http.StatusMovedPermanently)
	}).Methods("GET")

	r.HandleFunc("/{shortform:.*}", handler.Get).Methods("GET")
	r.HandleFunc("/", handler.Post).Methods("POST")
	r.HandleFunc("/{shortform:.*}", handler.Post).Methods("POST")
	http.Handle("/", r)

	log.Fatal(http.ListenAndServe(":" + strconv.Itoa(*port), nil))
}
<|endoftext|>


<commit_before>package falcore

import (
	"bufio"
	"crypto/rand"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"sync"
	"syscall"
	"time"
)

type Server struct {
	Addr               string
	Pipeline           *Pipeline
	CompletionCallback RequestCompletionCallback
	listener           net.Listener
	listenerFile       *os.File
	stopAccepting      chan int
	handlerWaitGroup   *sync.WaitGroup
	logPrefix          string
	AcceptReady        chan int
	sendfile           bool
	sockOpt            int
	bufferPool         *BufferPool
}

type RequestCompletionCallback func(req *Request, res *http.Response)

func NewServer(port int, pipeline *Pipeline) *Server {
	s := new(Server)
	s.Addr = fmt.Sprintf(":%v", port)
	s.Pipeline = pipeline
	s.stopAccepting = make(chan int)
	s.AcceptReady = make(chan int, 1)
	s.handlerWaitGroup = new(sync.WaitGroup)
	s.logPrefix = fmt.Sprintf("%d", syscall.Getpid())

	// openbsd/netbsd don't have TCP_NOPUSH so it's likely sendfile will be slower
	// without these socket options, just enable for linux, mac and freebsd.
	// TODO (Graham) windows has TransmitFile zero-copy mechanism, try to use it
	switch runtime.GOOS {
	case "linux":
		s.sendfile = true
		s.sockOpt = 0x3 // syscall.TCP_CORK
	case "freebsd", "darwin":
		s.sendfile = true
		s.sockOpt = 0x4 // syscall.TCP_NOPUSH
	default:
		s.sendfile = false
	}

	// buffer pool for reusing connection bufio.Readers
	s.bufferPool = NewBufferPool(100, 8192)

	return s
}

func (srv *Server) FdListen(fd int) error {
	var err error
	srv.listenerFile = os.NewFile(uintptr(fd), "")
	if srv.listener, err = net.FileListener(srv.listenerFile); err != nil {
		return err
	}
	if _, ok := srv.listener.(*net.TCPListener); !ok {
		return errors.New("Broken listener isn't TCP")
	}
	return nil
}

func (srv *Server) socketListen() error {
	var la *net.TCPAddr
	var err error
	if la, err = net.ResolveTCPAddr("tcp", srv.Addr); err != nil {
		return err
	}

	var l *net.TCPListener
	if l, err = net.ListenTCP("tcp", la); err != nil {
		return err
	}
	srv.listener = l
	// setup listener to be non-blocking if we're not on windows.
	// this is required for hot restart to work.
	return srv.setupNonBlockingListener(err, l)
}

func (srv *Server) ListenAndServe() error {
	if srv.Addr == "" {
		srv.Addr = ":http"
	}
	if srv.listener == nil {
		if err := srv.socketListen(); err != nil {
			return err
		}
	}
	return srv.serve()
}

func (srv *Server) SocketFd() int {
	return int(srv.listenerFile.Fd())
}

func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	if srv.Addr == "" {
		srv.Addr = ":https"
	}
	config := &tls.Config{
		Rand:       rand.Reader,
		Time:       time.Now,
		NextProtos: []string{"http/1.1"},
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	if srv.listener == nil {
		if err := srv.socketListen(); err != nil {
			return err
		}
	}

	srv.listener = tls.NewListener(srv.listener, config)

	return srv.serve()
}

func (srv *Server) StopAccepting() {
	close(srv.stopAccepting)
}

func (srv *Server) Port() int {
	if l := srv.listener; l != nil {
		a := l.Addr()
		if _, p, e := net.SplitHostPort(a.String()); e == nil && p != "" {
			server_port, _ := strconv.Atoi(p)
			return server_port
		}
	}
	return 0
}

func (srv *Server) serve() (e error) {
	var accept = true
	srv.AcceptReady <- 1
	for accept {
		var c net.Conn
		if l, ok := srv.listener.(*net.TCPListener); ok {
			l.SetDeadline(time.Now().Add(3 * time.Second))
		}
		c, e = srv.listener.Accept()
		if e != nil {
			if ope, ok := e.(*net.OpError); ok {
				if !(ope.Timeout() && ope.Temporary()) {
					Error("%s SERVER Accept Error: %v", srv.serverLogPrefix(), ope)
				}
			} else {
				Error("%s SERVER Accept Error: %v", srv.serverLogPrefix(), e)
			}
		} else {
			//Trace("Handling!")
			srv.handlerWaitGroup.Add(1)
			go srv.handler(c)
		}
		select {
		case <-srv.stopAccepting:
			accept = false
		default:
		}
	}
	Trace("Stopped accepting, waiting for handlers")
	// wait for handlers
	srv.handlerWaitGroup.Wait()
	return nil
}

func (srv *Server) sentinel(c net.Conn, connClosed chan int) {
	select {
	case <-srv.stopAccepting:
		c.SetReadDeadline(time.Now().Add(3 * time.Second))
	case <-connClosed:
	}
}

// For compatibility with net/http.Server or Google App Engine
// If you are using falcore.Server as a net/http.Handler, you should
// not call any of the Listen methods
func (srv *Server) ServeHTTP(wr http.ResponseWriter, req *http.Request) {
	// We can't get the connection in this case.
	// Need to be really careful about how we use this property elsewhere.
	request := newRequest(req, nil, time.Now())
	res := srv.handlerExecutePipeline(request)

	// Copy headers
	theHeader := wr.Header()
	for key, header := range res.Header {
		theHeader[key] = header
	}

	// Write headers
	wr.WriteHeader(res.StatusCode)

	// Write Body
	request.startPipelineStage("server.ResponseWrite")
	if res.Body != nil {
		defer res.Body.Close()
		io.Copy(wr, res.Body)
	}
	request.finishPipelineStage()
	request.finishRequest()

	srv.requestFinished(request, res)
}

func (srv *Server) handler(c net.Conn) {
	var startTime time.Time
	bpe := srv.bufferPool.Take(c)
	defer srv.bufferPool.Give(bpe)
	var closeSentinelChan = make(chan int)
	go srv.sentinel(c, closeSentinelChan)
	defer srv.connectionFinished(c, closeSentinelChan)
	var err error
	var req *http.Request
	// no keepalive (for now)
	reqCount := 0
	keepAlive := true
	for err == nil && keepAlive {
		if _, err := bpe.Br.Peek(1); err == nil {
			startTime = time.Now()
		}
		if req, err = http.ReadRequest(bpe.Br); err == nil {
			if req.Header.Get("Connection") != "Keep-Alive" {
				keepAlive = false
			}
			request := newRequest(req, c, startTime)
			reqCount++

			pssInit := new(PipelineStageStat)
			pssInit.Name = "server.Init"
			pssInit.StartTime = startTime
			pssInit.EndTime = time.Now()
			pssInit.Type = PipelineStageTypeOverhead
			request.appendPipelineStage(pssInit)

			// execute the pipeline
			var res = srv.handlerExecutePipeline(request, keepAlive)

			// shutting down?
			select {
			case <-srv.stopAccepting:
				keepAlive = false
				res.Close = true
			default:
			}

			// write response
			srv.handlerWriteResponse(request, res, c)

			if res.Close {
				keepAlive = false
			}
		} else {
			// EOF is socket closed
			if nerr, ok := err.(net.Error); err != io.EOF && !(ok && nerr.Timeout()) {
				Error("%s %v ERROR reading request: <%T %v>", srv.serverLogPrefix(), c.RemoteAddr(), err, err)
			}
		}
	}
	//Debug("%s Processed %v requests on connection %v", srv.serverLogPrefix(), reqCount, c.RemoteAddr())
}

func (srv *Server) handlerExecutePipeline(request *Request, boolean keepAlive)*http.Response {
	var res *http.Response
	// execute the pipeline
	if res = srv.Pipeline.execute(request); res == nil {
		res = StringResponse(request.HttpRequest, 404, nil, "Not Found")
	}

	// The res.Write omits Content-length on 0 length bodies, and by spec,
	// it SHOULD. While this is not MUST, it's kinda broken.  See sec 4.4
	// of rfc2616 and a 200 with a zero length does not satisfy any of the
	// 5 conditions if Connection: keep-alive is set :(
	// I'm forcing chunked which seems to work because I couldn't get the
	// content length to write if it was 0.
	// Specifically, the android http client waits forever if there's no
	// content-length instead of assuming zero at the end of headers. der.
	if res.ContentLength == 0 && len(res.TransferEncoding) == 0 && !((res.StatusCode-100 < 100) || res.StatusCode == 204 || res.StatusCode == 304) {
		request.HttpRequest.TransferEncoding = []string{"identity"}
	}
	if res.ContentLength < 0 {
		request.HttpRequest.TransferEncoding = []string{"chunked"}
	}

	// For HTTP/1.0 and Keep-Alive, sending the Connection: Keep-Alive response header is required
	// because close is default (opposite of 1.1)
	if keepAlive && !req.ProtoAtLeast(1, 1) {
		res.Header.Add("Connection", "Keep-Alive")
	}

	// cleanup
	request.HttpRequest.Body.Close()
	return res
}

func (srv *Server) handlerWriteResponse(request *Request, res *http.Response, c io.WriteCloser) {
	request.startPipelineStage("server.ResponseWrite")
	request.CurrentStage.Type = PipelineStageTypeOverhead
	if srv.sendfile {
		res.Write(c)
		if conn, ok := c.(net.Conn); ok {
			srv.cycleNonBlock(conn)
		}
	} else {
		wbuf := bufio.NewWriter(c)
		res.Write(wbuf)
		wbuf.Flush()
	}
	if res.Body != nil {
		res.Body.Close()
	}
	request.finishPipelineStage()
	request.finishRequest()
	srv.requestFinished(request, res)
}

func (srv *Server) serverLogPrefix() string {
	return srv.logPrefix
}

func (srv *Server) requestFinished(request *Request, res *http.Response) {
	if srv.CompletionCallback != nil {
		// Don't block the connecion for this
		go srv.CompletionCallback(request, res)
	}
}

func (srv *Server) connectionFinished(c net.Conn, closeChan chan int) {
	c.Close()
	close(closeChan)
	srv.handlerWaitGroup.Done()
}
<commit_msg>gofmt<commit_after>package falcore

import (
	"bufio"
	"crypto/rand"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"sync"
	"syscall"
	"time"
)

type Server struct {
	Addr               string
	Pipeline           *Pipeline
	CompletionCallback RequestCompletionCallback
	listener           net.Listener
	listenerFile       *os.File
	stopAccepting      chan int
	handlerWaitGroup   *sync.WaitGroup
	logPrefix          string
	AcceptReady        chan int
	sendfile           bool
	sockOpt            int
	bufferPool         *BufferPool
}

type RequestCompletionCallback func(req *Request, res *http.Response)

func NewServer(port int, pipeline *Pipeline) *Server {
	s := new(Server)
	s.Addr = fmt.Sprintf(":%v", port)
	s.Pipeline = pipeline
	s.stopAccepting = make(chan int)
	s.AcceptReady = make(chan int, 1)
	s.handlerWaitGroup = new(sync.WaitGroup)
	s.logPrefix = fmt.Sprintf("%d", syscall.Getpid())

	// openbsd/netbsd don't have TCP_NOPUSH so it's likely sendfile will be slower
	// without these socket options, just enable for linux, mac and freebsd.
	// TODO (Graham) windows has TransmitFile zero-copy mechanism, try to use it
	switch runtime.GOOS {
	case "linux":
		s.sendfile = true
		s.sockOpt = 0x3 // syscall.TCP_CORK
	case "freebsd", "darwin":
		s.sendfile = true
		s.sockOpt = 0x4 // syscall.TCP_NOPUSH
	default:
		s.sendfile = false
	}

	// buffer pool for reusing connection bufio.Readers
	s.bufferPool = NewBufferPool(100, 8192)

	return s
}

func (srv *Server) FdListen(fd int) error {
	var err error
	srv.listenerFile = os.NewFile(uintptr(fd), "")
	if srv.listener, err = net.FileListener(srv.listenerFile); err != nil {
		return err
	}
	if _, ok := srv.listener.(*net.TCPListener); !ok {
		return errors.New("Broken listener isn't TCP")
	}
	return nil
}

func (srv *Server) socketListen() error {
	var la *net.TCPAddr
	var err error
	if la, err = net.ResolveTCPAddr("tcp", srv.Addr); err != nil {
		return err
	}

	var l *net.TCPListener
	if l, err = net.ListenTCP("tcp", la); err != nil {
		return err
	}
	srv.listener = l
	// setup listener to be non-blocking if we're not on windows.
	// this is required for hot restart to work.
	return srv.setupNonBlockingListener(err, l)
}

func (srv *Server) ListenAndServe() error {
	if srv.Addr == "" {
		srv.Addr = ":http"
	}
	if srv.listener == nil {
		if err := srv.socketListen(); err != nil {
			return err
		}
	}
	return srv.serve()
}

func (srv *Server) SocketFd() int {
	return int(srv.listenerFile.Fd())
}

func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	if srv.Addr == "" {
		srv.Addr = ":https"
	}
	config := &tls.Config{
		Rand:       rand.Reader,
		Time:       time.Now,
		NextProtos: []string{"http/1.1"},
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	if srv.listener == nil {
		if err := srv.socketListen(); err != nil {
			return err
		}
	}

	srv.listener = tls.NewListener(srv.listener, config)

	return srv.serve()
}

func (srv *Server) StopAccepting() {
	close(srv.stopAccepting)
}

func (srv *Server) Port() int {
	if l := srv.listener; l != nil {
		a := l.Addr()
		if _, p, e := net.SplitHostPort(a.String()); e == nil && p != "" {
			server_port, _ := strconv.Atoi(p)
			return server_port
		}
	}
	return 0
}

func (srv *Server) serve() (e error) {
	var accept = true
	srv.AcceptReady <- 1
	for accept {
		var c net.Conn
		if l, ok := srv.listener.(*net.TCPListener); ok {
			l.SetDeadline(time.Now().Add(3 * time.Second))
		}
		c, e = srv.listener.Accept()
		if e != nil {
			if ope, ok := e.(*net.OpError); ok {
				if !(ope.Timeout() && ope.Temporary()) {
					Error("%s SERVER Accept Error: %v", srv.serverLogPrefix(), ope)
				}
			} else {
				Error("%s SERVER Accept Error: %v", srv.serverLogPrefix(), e)
			}
		} else {
			//Trace("Handling!")
			srv.handlerWaitGroup.Add(1)
			go srv.handler(c)
		}
		select {
		case <-srv.stopAccepting:
			accept = false
		default:
		}
	}
	Trace("Stopped accepting, waiting for handlers")
	// wait for handlers
	srv.handlerWaitGroup.Wait()
	return nil
}

func (srv *Server) sentinel(c net.Conn, connClosed chan int) {
	select {
	case <-srv.stopAccepting:
		c.SetReadDeadline(time.Now().Add(3 * time.Second))
	case <-connClosed:
	}
}

// For compatibility with net/http.Server or Google App Engine
// If you are using falcore.Server as a net/http.Handler, you should
// not call any of the Listen methods
func (srv *Server) ServeHTTP(wr http.ResponseWriter, req *http.Request) {
	// We can't get the connection in this case.
	// Need to be really careful about how we use this property elsewhere.
	request := newRequest(req, nil, time.Now())
	res := srv.handlerExecutePipeline(request)

	// Copy headers
	theHeader := wr.Header()
	for key, header := range res.Header {
		theHeader[key] = header
	}

	// Write headers
	wr.WriteHeader(res.StatusCode)

	// Write Body
	request.startPipelineStage("server.ResponseWrite")
	if res.Body != nil {
		defer res.Body.Close()
		io.Copy(wr, res.Body)
	}
	request.finishPipelineStage()
	request.finishRequest()

	srv.requestFinished(request, res)
}

func (srv *Server) handler(c net.Conn) {
	var startTime time.Time
	bpe := srv.bufferPool.Take(c)
	defer srv.bufferPool.Give(bpe)
	var closeSentinelChan = make(chan int)
	go srv.sentinel(c, closeSentinelChan)
	defer srv.connectionFinished(c, closeSentinelChan)
	var err error
	var req *http.Request
	// no keepalive (for now)
	reqCount := 0
	keepAlive := true
	for err == nil && keepAlive {
		if _, err := bpe.Br.Peek(1); err == nil {
			startTime = time.Now()
		}
		if req, err = http.ReadRequest(bpe.Br); err == nil {
			if req.Header.Get("Connection") != "Keep-Alive" {
				keepAlive = false
			}
			request := newRequest(req, c, startTime)
			reqCount++

			pssInit := new(PipelineStageStat)
			pssInit.Name = "server.Init"
			pssInit.StartTime = startTime
			pssInit.EndTime = time.Now()
			pssInit.Type = PipelineStageTypeOverhead
			request.appendPipelineStage(pssInit)

			// execute the pipeline
			var res = srv.handlerExecutePipeline(request, keepAlive)

			// shutting down?
			select {
			case <-srv.stopAccepting:
				keepAlive = false
				res.Close = true
			default:
			}

			// write response
			srv.handlerWriteResponse(request, res, c)

			if res.Close {
				keepAlive = false
			}
		} else {
			// EOF is socket closed
			if nerr, ok := err.(net.Error); err != io.EOF && !(ok && nerr.Timeout()) {
				Error("%s %v ERROR reading request: <%T %v>", srv.serverLogPrefix(), c.RemoteAddr(), err, err)
			}
		}
	}
	//Debug("%s Processed %v requests on connection %v", srv.serverLogPrefix(), reqCount, c.RemoteAddr())
}

func (srv *Server) handlerExecutePipeline(request *Request, boolean keepAlive) *http.Response {
	var res *http.Response
	// execute the pipeline
	if res = srv.Pipeline.execute(request); res == nil {
		res = StringResponse(request.HttpRequest, 404, nil, "Not Found")
	}

	// The res.Write omits Content-length on 0 length bodies, and by spec,
	// it SHOULD. While this is not MUST, it's kinda broken.  See sec 4.4
	// of rfc2616 and a 200 with a zero length does not satisfy any of the
	// 5 conditions if Connection: keep-alive is set :(
	// I'm forcing chunked which seems to work because I couldn't get the
	// content length to write if it was 0.
	// Specifically, the android http client waits forever if there's no
	// content-length instead of assuming zero at the end of headers. der.
	if res.ContentLength == 0 && len(res.TransferEncoding) == 0 && !((res.StatusCode-100 < 100) || res.StatusCode == 204 || res.StatusCode == 304) {
		request.HttpRequest.TransferEncoding = []string{"identity"}
	}
	if res.ContentLength < 0 {
		request.HttpRequest.TransferEncoding = []string{"chunked"}
	}

	// For HTTP/1.0 and Keep-Alive, sending the Connection: Keep-Alive response header is required
	// because close is default (opposite of 1.1)
	if keepAlive && !req.ProtoAtLeast(1, 1) {
		res.Header.Add("Connection", "Keep-Alive")
	}

	// cleanup
	request.HttpRequest.Body.Close()
	return res
}

func (srv *Server) handlerWriteResponse(request *Request, res *http.Response, c io.WriteCloser) {
	request.startPipelineStage("server.ResponseWrite")
	request.CurrentStage.Type = PipelineStageTypeOverhead
	if srv.sendfile {
		res.Write(c)
		if conn, ok := c.(net.Conn); ok {
			srv.cycleNonBlock(conn)
		}
	} else {
		wbuf := bufio.NewWriter(c)
		res.Write(wbuf)
		wbuf.Flush()
	}
	if res.Body != nil {
		res.Body.Close()
	}
	request.finishPipelineStage()
	request.finishRequest()
	srv.requestFinished(request, res)
}

func (srv *Server) serverLogPrefix() string {
	return srv.logPrefix
}

func (srv *Server) requestFinished(request *Request, res *http.Response) {
	if srv.CompletionCallback != nil {
		// Don't block the connecion for this
		go srv.CompletionCallback(request, res)
	}
}

func (srv *Server) connectionFinished(c net.Conn, closeChan chan int) {
	c.Close()
	close(closeChan)
	srv.handlerWaitGroup.Done()
}
<|endoftext|>


<commit_before>package wave

import "net/rpc"

// Server encapsulates the standard rpc.Server allowing for adding of aditional functions
// or hooking into existing rpc server functions.
type Server struct {
	*rpc.Server
}

// NewServer creates and returns a new instance of 'Server'
func NewServer() *Server {
	return new(Server)
}

// TODO: hook into `Register` and `RegisterName` functions to provide service discovery
// functionality
<commit_msg>fix server fun<commit_after>package wave

import "net/rpc"

// Server encapsulates the standard rpc.Server allowing for adding of aditional functions
// or hooking into existing rpc server functions.
type Server struct {
	*rpc.Server
}

// NewServer creates and returns a new instance of 'Server'
func NewServer() *Server {
	return &Server{
		Server: rpc.NewServer(),
	}
}

// TODO: hook into `Register` and `RegisterName` functions to provide service discovery
// functionality
<|endoftext|>


<commit_before>// Copyright 2016 orivil Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Package grace provides a graceful restarting http server.
package grace

import (
	"net/http"
	"os"
	"syscall"
	"os/signal"
	"net"
	"fmt"
	"os/exec"
	"crypto/tls"
	"sync"
	"flag"
	"runtime"
	"time"
	"github.com/fsnotify/fsnotify"
	"gopkg.in/orivil/log.v0"
)

const graceTag = "graceful"

var LogF = func(format string, args...interface{}) {

	as := make([]interface{}, len(args) + 1)
	as[0] = pid
	for idx, arg := range args {
		as[idx + 1] = arg
	}
	log.Printf("[process %d] " + format, as...)
}

var ErrF = func(format string, args...interface{}) {

	as := make([]interface{}, len(args) + 1)
	as[0] = pid
	for idx, arg := range args {
		as[idx + 1] = arg
	}
	log.ErrWarnF("[process %d] " + format, as...)
}

// IsSupportSocketFile can be overwritten.
var IsSupportSocketFile = func() bool {
	if runtime.GOOS == "windows" {
		return false
	} else {
		return true
	}
}

// IsSupportSignal can be overwritten.
var IsSupportSignal = func() bool {
	if runtime.GOOS == "windows" {
		return false
	} else {
		return true
	}
}

var pid = os.Getpid()

var isChild bool

func init() {
	flag.BoolVar(&isChild, graceTag, false, "")
	if !flag.Parsed() {
		flag.Parse()
	}
}

// ListenAndServe listens on the TCP network address addr
// and then calls Serve with handler to handle requests
// on incoming connections.
// Accepted connections are configured to enable TCP keep-alives.
// Handler is typically nil, in which case the DefaultServeMux is
// used.
//
// A trivial example server is:
//
//	package main
//
//	import (
//		"io"
//		"gopkg.in/orivil/grace.v0"
//		"gopkg.in/orivil/log.v0"
//	)
//
//	// hello world, the web server
//	func HelloServer(w http.ResponseWriter, req *http.Request) {
//		io.WriteString(w, "hello, world!\n")
//	}
//
//	func main() {
//		http.HandleFunc("/hello", HelloServer)
//		err := grace.ListenAndServe(":12345", nil)
//		if err != nil {
// 			log.ErrEmergency(err)
//		}
//	}
//
// If the server is graceful stopped, ListenAndServe will return a nil error.
func ListenAndServe(addr string, h http.Handler) error {

	httpServer := &http.Server{Addr: addr, Handler: h}
	return NewGraceServer(httpServer).ListenAndServe()
}

// ListenAndServeTLS acts identically to ListenAndServe, except that it
// expects HTTPS connections. Additionally, files containing a certificate and
// matching private key for the server must be provided. If the certificate
// is signed by a certificate authority, the certFile should be the concatenation
// of the server's certificate, any intermediates, and the CA's certificate.
//
// A trivial example server is:
//
//	import (
//		"log"
//		"gopkg.in/orivil/grace.v0"
//		"gopkg.in/orivil/log.v0"
//	)
//
//	func handler(w http.ResponseWriter, req *http.Request) {
//		w.Header().Set("Content-Type", "text/plain")
//		w.Write([]byte("This is an example server.\n"))
//	}
//
//	func main() {
//		http.HandleFunc("/", handler)
//		log.Printf("About to listen on 10443. Go to https://127.0.0.1:10443/")
//		err := http.ListenAndServeTLS(":10443", "cert.pem", "key.pem", nil)
//		if err != nil {
// 			log.ErrEmergency(err)
//		}
//	}
//
// One can use generate_cert.go in crypto/tls to generate cert.pem and key.pem.
//
// If the server is graceful stopped, ListenAndServeTLS will return a nil error.
func ListenAndServeTLS(addr, certFile, keyFile string, h http.Handler) error {

	httpServer := &http.Server{Addr: addr, Handler: h}
	return NewGraceServer(httpServer).ListenAndServeTLS(certFile, keyFile)
}

type GraceServer struct {
	*http.Server
	netListener net.Listener
	closeServer chan bool
	watcher *fsnotify.Watcher
}

func NewGraceServer(s *http.Server) *GraceServer {

	LogF("starting new server...\n")
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	return &GraceServer{
		Server: s,
		closeServer: make(chan bool),
		watcher: watcher,
	}
}

func (gs *GraceServer) newListener(addr string) (l net.Listener, err error) {
	if isChild {
		err := gs.killParentProcess()
		if err != nil {
			return nil, err
		} else {
			if IsSupportSocketFile() {
				f := os.NewFile(3, "")
				return net.FileListener(f)
			}
		}
	} else {
		gs.watch()
	}
	return net.Listen("tcp", addr)
}

// If the server is graceful stopped, ListenAndServeTLS will return a nil error.
func (gs *GraceServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := gs.Addr
	if addr == "" {
		addr = ":https"
	}

	l, err := gs.newListener(addr)
	if err != nil {
		return err
	}

	config := gs.TLSConfig
	if config == nil {
		config = &tls.Config{}
	}
	if !strSliceContains(config.NextProtos, "http/1.1") {
		config.NextProtos = append(config.NextProtos, "http/1.1")
	}

	c, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	config.Certificates = append(config.Certificates, c)
	gs.TLSConfig = config
	tlsListener := tls.NewListener(l, config)
	gs.netListener = tlsListener
	return gs.Serve()
}

// If the server is graceful stopped, ListenAndServe will return a nil error.
func (gs *GraceServer) ListenAndServe() error {
	addr := gs.Addr
	if addr == "" {
		addr = ":http"
	}

	l, err := gs.newListener(addr)
	if err != nil {
		return err
	}
	gs.netListener = l
	return gs.Serve()
}

func (gs *GraceServer) watch() {
	err := gs.watcher.Add(os.Args[0])
	if err != nil {
		ErrF(err.Error())
	}
}

func (gs *GraceServer) killParentProcess() error {

	ppid := os.Getppid()
	process, err := os.FindProcess(ppid)
	if err != nil {
		return err
	}

	LogF("killing parent process [ %d ]", ppid)
	if IsSupportSignal() {
		err = process.Signal(syscall.SIGINT)
	} else {
		err = process.Kill()
	}
	if err != nil {
		return err
	} else {
		wait := func(){
			// wait until parent process exited
			ticker := time.NewTicker(20 * time.Millisecond)
			for _ = range ticker.C {
				state, err := process.Wait()
				if err != nil || state.Exited() {
					gs.watch()
					return
				}
			}
		}

		if IsSupportSignal() {
			go wait()
		} else {
			wait()
		}
		return nil
	}
}

func (gs *GraceServer) Serve() error {
	errChan := make(chan error)
	go func() {
		gs.listenEvents()
		LogF("ready to serve http request...")
		errChan <- gs.Server.Serve(gs.netListener)
	}()

	select {
	case err := <-errChan:
		gs.watcher.Close()
		return err
	case <-gs.closeServer:
		gs.watcher.Close()
		gs.netListener.Close()
		return nil
	}
}

func (gs *GraceServer) listenEvents() {

	var sig os.Signal
	signalChan := make(chan os.Signal)

	signal.Notify(
		signalChan,
		syscall.SIGTERM,
		syscall.SIGHUP,
		syscall.SIGINT,
	)

	go func() {
		for sig = range signalChan {
			switch sig {

			case syscall.SIGTERM, syscall.SIGINT:
				gs.Stop()
				return

			case syscall.SIGHUP:
				gs.Restart()
				return

			default:
				LogF("unknown signal %v", sig)
			}
		}
	}()

	timer := time.NewTimer(0)
	<- timer.C
	go func() {
		for {
			select {
			case evt := <-gs.watcher.Events:
				// only trigger the last event
				if evt.Op & fsnotify.Write == fsnotify.Write {
					timer.Reset(200 * time.Millisecond)
				}
			case err := <-gs.watcher.Errors:
				ErrF("grace.GraceServer.listenEvents(): %v", err)
			}
		}
	}()

	go func() {
		for {
			<- timer.C
			gs.Restart()
		}
	}()
}

func (gs *GraceServer) Stop() {

	LogF("closing server...")
	gs.closeServer <- true
}

func (gs *GraceServer) Restart() {

	// un-watch the file, otherwise start new process will got error
	gs.watcher.Remove(os.Args[0])
	LogF("restarting http server...")
	err := gs.startNewProcess()
	if err != nil {

		// re-watch the file
		gs.watcher.Add(os.Args[0])
		ErrF("grace.GraceServer.restart(): %v", err)
		LogF("continue serving")
	}
}

func (gs *GraceServer) getGraceTagArgs() []string {
	res := make([]string, len(os.Args))
	for idx, arg := range os.Args {
		res[idx] = arg
	}
	if !isChild {
		// pass grace tag to the incoming process
		res = append(res, "-" + graceTag)
	}
	return res
}

func (gs *GraceServer) startNewProcess() error {
	if l, ok := gs.netListener.(CanGetFile); ok {
		args := gs.getGraceTagArgs()
		path := args[0]
		args = args[1:]
		cmd := exec.Command(path, args...)

		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if IsSupportSocketFile() {
			f, err := l.File()
			if err != nil {
				ErrF("grace.GraceServer.startNewProcess(): %v", err)
			} else {
				cmd.ExtraFiles = []*os.File{f}
			}
		}
		err := cmd.Run()
		if err != nil {
			return err
		}
		return nil
	} else {
		return fmt.Errorf("unknown net listener")
	}
}

type CanGetFile interface {
	File() (f *os.File, err error)
}

// wait group for net listener and opened net connects
var waitGroup = sync.WaitGroup{}

type Listener struct {
	net.Listener
}

func (l *Listener) Accept() (net.Conn, error) {

	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	} else {
		waitGroup.Add(1)
		return NewGracefulConn(c), nil
	}
}

func (l *Listener) Close() error {
	// stop accept new connect
	err := l.Listener.Close()

	// wait until all opened connects closed
	waitGroup.Wait()
	return err
}

type GracefulConn struct {
	net.Conn
}

func NewGracefulConn(c net.Conn) *GracefulConn {
	return &GracefulConn{
		Conn: c,
	}
}

func (w GracefulConn) Close() (e error) {
	e = w.Conn.Close()
	waitGroup.Done()
	return
}

func strSliceContains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}<commit_msg>fix error msg<commit_after>// Copyright 2016 orivil Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Package grace provides a graceful restarting http server.
package grace

import (
	"net/http"
	"os"
	"syscall"
	"os/signal"
	"net"
	"fmt"
	"os/exec"
	"crypto/tls"
	"sync"
	"flag"
	"runtime"
	"time"
	"github.com/fsnotify/fsnotify"
	"gopkg.in/orivil/log.v0"
)

const graceTag = "graceful"

var LogF = func(format string, args...interface{}) {

	as := make([]interface{}, len(args) + 1)
	as[0] = pid
	for idx, arg := range args {
		as[idx + 1] = arg
	}
	log.Printf("[process %d] " + format, as...)
}

var ErrF = func(format string, args...interface{}) {

	as := make([]interface{}, len(args) + 1)
	as[0] = pid
	for idx, arg := range args {
		as[idx + 1] = arg
	}
	log.ErrWarnF("[process %d] " + format, as...)
}

// TODO: from now on, only tested on windows and linux.
// IsSupportSocketFile can be overwritten.
var IsSupportSocketFile = func() bool {
	if runtime.GOOS == "windows" {
		return false
	} else {
		return true
	}
}

// IsSupportSignal can be overwritten.
var IsSupportSignal = func() bool {
	if runtime.GOOS == "windows" {
		return false
	} else {
		return true
	}
}

var pid = os.Getpid()

var isChild bool

func init() {
	flag.BoolVar(&isChild, graceTag, false, "")
	if !flag.Parsed() {
		flag.Parse()
	}
}

// ListenAndServe listens on the TCP network address addr
// and then calls Serve with handler to handle requests
// on incoming connections.
// Accepted connections are configured to enable TCP keep-alives.
// Handler is typically nil, in which case the DefaultServeMux is
// used.
//
// A trivial example server is:
//
//	package main
//
//	import (
//		"io"
//		"gopkg.in/orivil/grace.v0"
//		"gopkg.in/orivil/log.v0"
//	)
//
//	// hello world, the web server
//	func HelloServer(w http.ResponseWriter, req *http.Request) {
//		io.WriteString(w, "hello, world!\n")
//	}
//
//	func main() {
//		http.HandleFunc("/hello", HelloServer)
//		err := grace.ListenAndServe(":12345", nil)
//		if err != nil {
// 			log.ErrEmergency(err)
//		}
//	}
//
// If the server is graceful stopped, ListenAndServe will return a nil error.
func ListenAndServe(addr string, h http.Handler) error {

	httpServer := &http.Server{Addr: addr, Handler: h}
	return NewGraceServer(httpServer).ListenAndServe()
}

// ListenAndServeTLS acts identically to ListenAndServe, except that it
// expects HTTPS connections. Additionally, files containing a certificate and
// matching private key for the server must be provided. If the certificate
// is signed by a certificate authority, the certFile should be the concatenation
// of the server's certificate, any intermediates, and the CA's certificate.
//
// A trivial example server is:
//
//	import (
//		"log"
//		"gopkg.in/orivil/grace.v0"
//		"gopkg.in/orivil/log.v0"
//	)
//
//	func handler(w http.ResponseWriter, req *http.Request) {
//		w.Header().Set("Content-Type", "text/plain")
//		w.Write([]byte("This is an example server.\n"))
//	}
//
//	func main() {
//		http.HandleFunc("/", handler)
//		log.Printf("About to listen on 10443. Go to https://127.0.0.1:10443/")
//		err := http.ListenAndServeTLS(":10443", "cert.pem", "key.pem", nil)
//		if err != nil {
// 			log.ErrEmergency(err)
//		}
//	}
//
// One can use generate_cert.go in crypto/tls to generate cert.pem and key.pem.
//
// If the server is graceful stopped, ListenAndServeTLS will return a nil error.
func ListenAndServeTLS(addr, certFile, keyFile string, h http.Handler) error {

	httpServer := &http.Server{Addr: addr, Handler: h}
	return NewGraceServer(httpServer).ListenAndServeTLS(certFile, keyFile)
}

type GraceServer struct {
	*http.Server
	netListener net.Listener
	closeServer chan bool
	watcher *fsnotify.Watcher
}

func NewGraceServer(s *http.Server) *GraceServer {

	LogF("starting new server...\n")
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		panic(err)
	}
	return &GraceServer{
		Server: s,
		closeServer: make(chan bool),
		watcher: watcher,
	}
}

func (gs *GraceServer) newListener(addr string) (l net.Listener, err error) {
	if isChild {
		err := gs.killParentProcess()
		if err != nil {
			return nil, err
		} else {
			if IsSupportSocketFile() {
				f := os.NewFile(3, "")
				return net.FileListener(f)
			}
		}
	} else {
		gs.watch()
	}
	return net.Listen("tcp", addr)
}

// If the server is graceful stopped, ListenAndServeTLS will return a nil error.
func (gs *GraceServer) ListenAndServeTLS(certFile, keyFile string) error {
	addr := gs.Addr
	if addr == "" {
		addr = ":https"
	}

	l, err := gs.newListener(addr)
	if err != nil {
		return err
	}

	config := gs.TLSConfig
	if config == nil {
		config = &tls.Config{}
	}
	if !strSliceContains(config.NextProtos, "http/1.1") {
		config.NextProtos = append(config.NextProtos, "http/1.1")
	}

	c, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	config.Certificates = append(config.Certificates, c)
	gs.TLSConfig = config
	tlsListener := tls.NewListener(l, config)
	gs.netListener = tlsListener
	return gs.Serve()
}

// If the server is graceful stopped, ListenAndServe will return a nil error.
func (gs *GraceServer) ListenAndServe() error {
	addr := gs.Addr
	if addr == "" {
		addr = ":http"
	}

	l, err := gs.newListener(addr)
	if err != nil {
		return err
	}
	gs.netListener = l
	return gs.Serve()
}

func (gs *GraceServer) watch() {
	err := gs.watcher.Add(os.Args[0])
	if err != nil {
		ErrF(err.Error())
	}
}

func (gs *GraceServer) killParentProcess() error {

	ppid := os.Getppid()
	process, err := os.FindProcess(ppid)
	if err != nil {
		return err
	}

	LogF("killing parent process [ %d ]", ppid)
	if IsSupportSignal() {
		err = process.Signal(syscall.SIGINT)
	} else {
		err = process.Kill()
	}
	if err != nil {
		return err
	} else {
		wait := func(){
			// wait until parent process exited
			ticker := time.NewTicker(20 * time.Millisecond)
			for _ = range ticker.C {
				state, err := process.Wait()
				if err != nil || state.Exited() {
					gs.watch()
					return
				}
			}
		}

		if IsSupportSignal() {
			go wait()
		} else {
			wait()
		}
		return nil
	}
}

func (gs *GraceServer) Serve() error {
	errChan := make(chan error)
	go func() {
		gs.listenEvents()
		LogF("ready to serve http request...")
		errChan <- gs.Server.Serve(gs.netListener)
	}()

	select {
	case err := <-errChan:
		gs.watcher.Close()
		return err
	case <-gs.closeServer:
		gs.watcher.Close()
		gs.netListener.Close()
		return nil
	}
}

func (gs *GraceServer) listenEvents() {

	var sig os.Signal
	signalChan := make(chan os.Signal)

	signal.Notify(
		signalChan,
		syscall.SIGTERM,
		syscall.SIGHUP,
		syscall.SIGINT,
	)

	go func() {
		for sig = range signalChan {
			switch sig {

			case syscall.SIGTERM, syscall.SIGINT:
				gs.Stop()
				return

			case syscall.SIGHUP:
				gs.Restart()
				return

			default:
				LogF("unknown signal %v", sig)
			}
		}
	}()

	timer := time.NewTimer(0)
	<- timer.C
	go func() {
		for {
			select {
			case evt := <-gs.watcher.Events:
				// only trigger the last event
				if evt.Op & fsnotify.Write == fsnotify.Write {
					timer.Reset(200 * time.Millisecond)
				}
			case err := <-gs.watcher.Errors:
				if err != nil {
					ErrF("grace.GraceServer.listenEvents(): %v", err)
				}
			}
		}
	}()

	go func() {
		for {
			<- timer.C
			gs.Restart()
		}
	}()
}

func (gs *GraceServer) Stop() {

	LogF("closing server...")
	gs.closeServer <- true
}

func (gs *GraceServer) Restart() {

	// un-watch the file, otherwise start new process will got error
	gs.watcher.Remove(os.Args[0])
	LogF("restarting http server...")
	err := gs.startNewProcess()
	if err != nil {

		// re-watch the file
		gs.watcher.Add(os.Args[0])
		ErrF("grace.GraceServer.restart(): %v", err)
		LogF("continue serving")
	}
}

func (gs *GraceServer) getGraceTagArgs() []string {
	res := make([]string, len(os.Args))
	for idx, arg := range os.Args {
		res[idx] = arg
	}
	if !isChild {
		// pass grace tag to the incoming process
		res = append(res, "-" + graceTag)
	}
	return res
}

func (gs *GraceServer) startNewProcess() error {
	if l, ok := gs.netListener.(CanGetFile); ok {
		args := gs.getGraceTagArgs()
		path := args[0]
		args = args[1:]
		cmd := exec.Command(path, args...)

		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if IsSupportSocketFile() {
			f, err := l.File()
			if err != nil {
				ErrF("grace.GraceServer.startNewProcess(): %v", err)
			} else {
				cmd.ExtraFiles = []*os.File{f}
			}
		}
		err := cmd.Run()
		if err != nil {
			return err
		}
		return nil
	} else {
		return fmt.Errorf("unknown net listener")
	}
}

type CanGetFile interface {
	File() (f *os.File, err error)
}

// wait group for net listener and opened net connects
var waitGroup = sync.WaitGroup{}

type Listener struct {
	net.Listener
}

func (l *Listener) Accept() (net.Conn, error) {

	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	} else {
		waitGroup.Add(1)
		return NewGracefulConn(c), nil
	}
}

func (l *Listener) Close() error {
	// stop accept new connect
	err := l.Listener.Close()

	// wait until all opened connects closed
	waitGroup.Wait()
	return err
}

type GracefulConn struct {
	net.Conn
}

func NewGracefulConn(c net.Conn) *GracefulConn {
	return &GracefulConn{
		Conn: c,
	}
}

func (w GracefulConn) Close() (e error) {
	e = w.Conn.Close()
	waitGroup.Done()
	return
}

func strSliceContains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}<|endoftext|>


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

package main

import (
	"fmt"
	"net/http"

	"github.com/elazarl/goproxy"
	"github.com/jackwakefield/gopac"
	"github.com/spf13/viper"
)

type server struct {
	parser *gopac.Parser
	client *http.Client
	http   *goproxy.ProxyHttpServer
}

func newServer(parser *gopac.Parser, client *http.Client) (s *server) {
	s = &server{
		parser: parser,
		client: client,
		http:   goproxy.NewProxyHttpServer(),
	}

	// determine whether a proxy entry is available for the URL/host for each
	// HTTP request and if so, forward it to the proxy client
	s.http.OnRequest(s.isForwardable()).DoFunc(s.forward)

	return
}

func (s *server) isForwardable() goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		// find a proxy entry for the request URL/host
		entry, err := s.parser.FindProxy(req.URL.String(), req.URL.Host)

		// ensure no error occurred when finding an entry and that the entry
		// doesn't indicate the connection should be made directly
		if err == nil && entry != "DIRECT" {
			return true
		}

		return false
	}
}

func (s *server) forward(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	// retrieve a response for the request using the proxy client
	if resp, err := s.client.Get(r.URL.String()); err == nil {
		return r, resp
	}

	return r, nil
}

func (s *server) listen() error {
	serverPort := viper.GetInt(serverPortKey)
	address := fmt.Sprintf(":%d", serverPort)
	logger.Infof("Proxy server listening on address '%s'", address)

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

package main

import (
	"fmt"
	"net/http"
	"net/url"

	"github.com/jackwakefield/gopac"
	"github.com/spf13/viper"
	"gopkg.in/elazarl/goproxy.v1"
)

type server struct {
	parser    *gopac.Parser
	client    *http.Client
	http      *goproxy.ProxyHttpServer
	unproxied *http.Client
}

func newServer(parser *gopac.Parser, client *http.Client) (s *server) {
	s = &server{
		parser: parser,
		client: client,
		http:   goproxy.NewProxyHttpServer(),
	}

	s.http.NonproxyHandler = s.nonProxyHandler()

	// determine whether a proxy entry is available for the URL/host for each
	// HTTP request and if so, forward it to the proxy client
	s.http.OnRequest(s.isForwardable()).DoFunc(s.forward)

	return
}

func (s *server) isForwardable() goproxy.ReqConditionFunc {
	return func(req *http.Request, ctx *goproxy.ProxyCtx) bool {
		// find a proxy entry for the request URL/host
		entry, err := s.parser.FindProxy(req.URL.String(), req.URL.Host)

		// ensure no error occurred when finding an entry and that the entry
		// doesn't indicate the connection should be made directly
		if err == nil && entry != "DIRECT" {
			return true
		}

		return false
	}
}

func (s *server) forward(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
	// retrieve a response for the request using the proxy client
	if resp, err := s.client.Get(r.URL.String()); err == nil {
		return r, resp
	}

	return r, nil
}

func (s *server) nonProxyHandler() http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		// set the absolute URL from the host and relative URL and pass it back
		// to the proxy server
		var err error
		req.RequestURI = ""

		if req.URL, err = url.Parse("http://" + req.Host + req.URL.Path); err != nil {
			logger.Errorf("Failed to parse non-proxy URL (%s)", err)
			return
		}

		s.http.ServeHTTP(w, req)
	})
}

func (s *server) address() string {
	return fmt.Sprintf(":%d", viper.GetInt(serverPortKey))
}

func (s *server) listen() error {
	address := s.address()
	logger.Infof("Proxy server listening on address '%s'", address)

	return http.ListenAndServe(address, s.http)
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/jessemillar/rytsar-server/accessors"
	"github.com/jessemillar/rytsar-server/controllers"
	"github.com/zenazn/goji"
)

func main() {
	// Construct the dsn used for the database
	dsn := os.Getenv("RYTSAR_DB_USER") + ":" + os.Getenv("RYTSAR_DB_PASS") + "@tcp(" + os.Getenv("RYTSAR_DB_HOST") + ":" + os.Getenv("RYTSAR_DB_PORT") + ")/" + os.Getenv("RYTSAR_DB_NAME")

	// Construct a new AccessorGroup and connects it to the database
	ag := new(accessors.AccessorGroup)
	ag.ConnectToDB("mysql", dsn)

	// Constructs a new ControllerGroup and gives it the AccessorGroup
	cg := new(controllers.ControllerGroup)
	cg.Accessors = ag

	goji.Get("/health", cg.Health) // Service health
	goji.Get("/database/:latitude/:longitude/:radius", cg.DumpDatabase)
	goji.Get("/count/:latitude/:longitude/:radius", cg.CountNearbyEnemies)
	goji.Delete("/delete/:latitude/:longitude", cg.DeleteEnemy)
	goji.Serve()
}
<commit_msg>And this commit actually allows us to hit the /delete endpoint<commit_after>package main

import (
	"os"

	"github.com/jessemillar/rytsar-server/accessors"
	"github.com/jessemillar/rytsar-server/controllers"
	"github.com/zenazn/goji"
)

func main() {
	// Construct the dsn used for the database
	dsn := os.Getenv("RYTSAR_DB_USER") + ":" + os.Getenv("RYTSAR_DB_PASS") + "@tcp(" + os.Getenv("RYTSAR_DB_HOST") + ":" + os.Getenv("RYTSAR_DB_PORT") + ")/" + os.Getenv("RYTSAR_DB_NAME")

	// Construct a new AccessorGroup and connects it to the database
	ag := new(accessors.AccessorGroup)
	ag.ConnectToDB("mysql", dsn)

	// Constructs a new ControllerGroup and gives it the AccessorGroup
	cg := new(controllers.ControllerGroup)
	cg.Accessors = ag

	goji.Get("/health", cg.Health) // Service health
	goji.Get("/database/:latitude/:longitude/:radius", cg.DumpDatabase)
	goji.Get("/count/:latitude/:longitude/:radius", cg.CountNearbyEnemies)
	goji.Delete("/delete/:id", cg.DeleteEnemy)
	goji.Serve()
}
<|endoftext|>


<commit_before>// Copyright 2017 Axel Etcheverry. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package main

import (
	"crypto/md5"
	"fmt"
	"net/http"
	"os"
	"time"

	filetype "gopkg.in/h2non/filetype.v0"

	"io/ioutil"

	"github.com/euskadi31/image-service/httputil"
	"github.com/euskadi31/image-service/memfs"
	"github.com/justinas/alice"
	"github.com/rs/cors"
	"github.com/rs/xlog"
)

var c = cors.New(cors.Options{
	AllowedOrigins: []string{"*"},
	AllowedMethods: []string{"GET", "POST", "DELETE"},
	AllowedHeaders: []string{"*"},
})

// Server struct
type Server struct {
	mux    *http.ServeMux
	config *Configuration
	source SourceProvider
	cache  CacheProvider
}

// NewServer constructor
func NewServer(config *Configuration) *Server {
	var source SourceProvider
	var cache CacheProvider

	switch config.Image.Source.Provider {
	case "fs":
		source = NewFileSystemSourceProvider(config.Image.Source.FS)
	}

	switch config.Image.Cache.Provider {
	case "fs":
		cache = NewFileSystemCacheProvider(config.Image.Cache.FS)
	}

	return &Server{
		mux:    http.NewServeMux(),
		config: config,
		source: source,
		cache:  cache,
	}
}

// SetSourceProvider to server
func (s *Server) SetSourceProvider(source SourceProvider) {
	s.source = source
}

// SetCacheProvider to server
func (s *Server) SetCacheProvider(cache CacheProvider) {
	s.cache = cache
}

func (s Server) healthHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)

		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(`{"status": true}`))
}

func (s Server) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Not Found", http.StatusNotFound)
}

func (s Server) deleteHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		httputil.Failure(w, http.StatusBadRequest, httputil.ErrorMessage{
			Code:    0,
			Message: "Invalid URL path",
		})

		return
	}

	resource := &Resource{
		Path: r.URL.Path,
	}

	response := map[string]bool{
		"cache":  false,
		"source": false,
	}

	if from := r.URL.Query().Get("from"); from != "" {
		switch from {
		case "source":
			response["cache"] = s.cache.Del(resource)
			response["source"] = s.source.Del(resource)
		default:
			response["cache"] = s.cache.Del(resource)
		}
	}

	httputil.JSON(w, http.StatusOK, resource)
}

func (s Server) uploadHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		httputil.Failure(w, http.StatusBadRequest, httputil.ErrorMessage{
			Code:    0,
			Message: "Invalid URL path",
		})

		return
	}

	resource := &Resource{
		Path: r.URL.Path,
	}

	r.ParseMultipartForm(32 << 20)

	file, _, err := r.FormFile("image")
	if err != nil {
		httputil.FailureFromError(w, http.StatusBadRequest, err)

		return
	}

	defer file.Close()

	body, err := ioutil.ReadAll(file)
	if err != nil {
		httputil.FailureFromError(w, http.StatusBadRequest, err)

		return
	}

	resource.Body = body

	if err := s.source.Set(resource); err != nil {
		httputil.FailureFromError(w, http.StatusInternalServerError, err)

		return
	}

	// delete cache from source file
	go s.cache.Del(resource)

	mimeType := http.DetectContentType(body)

	// If cannot infer the type, infer it via magic numbers
	if mimeType == "application/octet-stream" {
		kind, err := filetype.Get(body)
		if err == nil && kind.MIME.Value != "" {
			mimeType = kind.MIME.Value
		}
	}

	h := md5.New()
	h.Write(body)

	httputil.JSON(w, http.StatusCreated, map[string]interface{}{
		"file": r.URL.Path,
		"size": len(body),
		"type": mimeType,
		"hash": fmt.Sprintf("%x", h.Sum(nil)),
	})
}

func (s Server) imageHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		http.Error(w, "Invalid URL path", http.StatusBadRequest)

		return
	}

	options, err := ParamsFromContext(r.Context())
	if err != nil {
		xlog.Errorf("Error while parsing options: %v", err.Error())
		http.Error(w, "Error while parsing options", http.StatusBadRequest)

		return
	}
	// xlog.Infof("options: %#v", options)

	resource := &Resource{
		Path:    r.URL.Path,
		Options: options,
	}

	w.Header().Set("Link", `</worker/client-hints.js>; rel="serviceworker"`)

	// fetch from cache
	if resource, ok := s.cache.Get(resource); ok {
		w.Header().Set("X-Image-From", "cache")

		ServeImage(w, r, resource)

		return
	}

	resource, err = s.source.Get(resource)
	if err != nil {
		if os.IsNotExist(err) {
			msg := fmt.Sprintf("File %s not found", r.URL.Path)

			xlog.Info(msg)
			http.Error(w, msg, http.StatusNotFound)

			return
		}

		xlog.Error(err)
		http.Error(w, err.Error(), http.StatusNotFound)

		return
	}

	if err := ProcessImage(resource); err != nil {
		xlog.Errorf("Error while processing the image: %v", err.Error())
		http.Error(w, "Error while processing the image", http.StatusInternalServerError)

		return
	}

	w.Header().Set("X-Image-From", "source")

	ServeImage(w, r, resource)

	// save resource in cache
	go func(r *Resource) {
		if err := s.cache.Set(r); err != nil {
			xlog.Error(err)
		}
	}(resource)
}

func (s Server) serviceWorker(w http.ResponseWriter, r *http.Request) {
	js := `
self.addEventListener('install', function(event) {
    console.log('Service Worker installing.');
});

self.addEventListener('activate', function(event) {
    console.log('Service Worker activating.');
});

// Listen to fetch events
self.addEventListener('fetch', function(event) {
    console.log(event.request.url);
    //throw JSON.stringify({data: event.request.url});

    if (/\.(jpg|jpeg|png)(\?+)?$/.test(event.request.url)) {
        console.log('Client-Hint');

        // Clone the request
		var req = event.request.clone();

        req.headers.set('DPR', window.devicePixelRatio);
        req.url = req.url + '&dpr=' + window.devicePixelRatio;

        event.respondWith(fetch(req/*, {
            mode: 'no-cors',
            headers: headers
        }*/));
    }
});
`
	w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	w.Write([]byte(js))
}

func (s Server) docHandler(w http.ResponseWriter, r *http.Request) {
	name := "doc/index.html"

	w.Header().Set("Content-Type", "text/html; charset=utf-8")

	body, err := Asset(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	info, err := AssetInfo(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	http.ServeContent(
		w,
		r,
		info.Name(),
		info.ModTime(),
		memfs.NewBuffer(&body),
	)
}

func (s Server) swaggerHandler(w http.ResponseWriter, r *http.Request) {
	name := "swagger.yaml"

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

	body, err := Asset(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	info, err := AssetInfo(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	http.ServeContent(
		w,
		r,
		info.Name,
		info.ModTime,
		memfs.NewBuffer(&body),
	)
}

// ListenAndServe service
func (s *Server) ListenAndServe() {
	middleware := alice.New(
		NewLoggerHandler(),
		NewImageExtensionFilterHandler(s.config),
		c.Handler,
	)

	readMiddleware := middleware.Append(
		NewParamsHandler(),
		NewClientHintsHandler(),
	)

	authMiddleware := middleware.Append(
		NewAuthHandler(s.config.Auth),
	)

	s.mux.HandleFunc("/favicon.ico", s.notFoundHandler)
	s.mux.HandleFunc("/health", s.healthHandler)
	s.mux.HandleFunc("/worker/client-hints.js", s.serviceWorker)

	if s.config.Doc.Enable {
		s.mux.HandleFunc("/swagger.yaml", s.swaggerHandler)
		s.mux.HandleFunc("/doc/", s.docHandler)

		/*fs := http.FileServer(http.Dir("doc"))
		s.mux.Handle("/doc/", http.StripPrefix("/doc/", fs))*/
	}

	s.mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			readMiddleware.ThenFunc(s.imageHandler).ServeHTTP(w, r)
		case http.MethodPost:
			authMiddleware.ThenFunc(s.uploadHandler).ServeHTTP(w, r)
		case http.MethodDelete:
			authMiddleware.ThenFunc(s.deleteHandler).ServeHTTP(w, r)
		default:
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		}
	})

	srv := &http.Server{
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 10 * time.Second,
		Handler:      s.mux,
	}

	srv.Addr = fmt.Sprintf("%s:%d", s.config.Server.Host, s.config.Server.Port)

	xlog.Infof("Server running on %s", srv.Addr)

	xlog.Fatal(srv.ListenAndServe())
}
<commit_msg>fix(server): bad code<commit_after>// Copyright 2017 Axel Etcheverry. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package main

import (
	"crypto/md5"
	"fmt"
	"net/http"
	"os"
	"time"

	filetype "gopkg.in/h2non/filetype.v0"

	"io/ioutil"

	"github.com/euskadi31/image-service/httputil"
	"github.com/euskadi31/image-service/memfs"
	"github.com/justinas/alice"
	"github.com/rs/cors"
	"github.com/rs/xlog"
)

var c = cors.New(cors.Options{
	AllowedOrigins: []string{"*"},
	AllowedMethods: []string{"GET", "POST", "DELETE"},
	AllowedHeaders: []string{"*"},
})

// Server struct
type Server struct {
	mux    *http.ServeMux
	config *Configuration
	source SourceProvider
	cache  CacheProvider
}

// NewServer constructor
func NewServer(config *Configuration) *Server {
	var source SourceProvider
	var cache CacheProvider

	switch config.Image.Source.Provider {
	case "fs":
		source = NewFileSystemSourceProvider(config.Image.Source.FS)
	}

	switch config.Image.Cache.Provider {
	case "fs":
		cache = NewFileSystemCacheProvider(config.Image.Cache.FS)
	}

	return &Server{
		mux:    http.NewServeMux(),
		config: config,
		source: source,
		cache:  cache,
	}
}

// SetSourceProvider to server
func (s *Server) SetSourceProvider(source SourceProvider) {
	s.source = source
}

// SetCacheProvider to server
func (s *Server) SetCacheProvider(cache CacheProvider) {
	s.cache = cache
}

func (s Server) healthHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)

		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(`{"status": true}`))
}

func (s Server) notFoundHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, "Not Found", http.StatusNotFound)
}

func (s Server) deleteHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		httputil.Failure(w, http.StatusBadRequest, httputil.ErrorMessage{
			Code:    0,
			Message: "Invalid URL path",
		})

		return
	}

	resource := &Resource{
		Path: r.URL.Path,
	}

	response := map[string]bool{
		"cache":  false,
		"source": false,
	}

	if from := r.URL.Query().Get("from"); from != "" {
		switch from {
		case "source":
			response["cache"] = s.cache.Del(resource)
			response["source"] = s.source.Del(resource)
		default:
			response["cache"] = s.cache.Del(resource)
		}
	}

	httputil.JSON(w, http.StatusOK, resource)
}

func (s Server) uploadHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		httputil.Failure(w, http.StatusBadRequest, httputil.ErrorMessage{
			Code:    0,
			Message: "Invalid URL path",
		})

		return
	}

	resource := &Resource{
		Path: r.URL.Path,
	}

	r.ParseMultipartForm(32 << 20)

	file, _, err := r.FormFile("image")
	if err != nil {
		httputil.FailureFromError(w, http.StatusBadRequest, err)

		return
	}

	defer file.Close()

	body, err := ioutil.ReadAll(file)
	if err != nil {
		httputil.FailureFromError(w, http.StatusBadRequest, err)

		return
	}

	resource.Body = body

	if err := s.source.Set(resource); err != nil {
		httputil.FailureFromError(w, http.StatusInternalServerError, err)

		return
	}

	// delete cache from source file
	go s.cache.Del(resource)

	mimeType := http.DetectContentType(body)

	// If cannot infer the type, infer it via magic numbers
	if mimeType == "application/octet-stream" {
		kind, err := filetype.Get(body)
		if err == nil && kind.MIME.Value != "" {
			mimeType = kind.MIME.Value
		}
	}

	h := md5.New()
	h.Write(body)

	httputil.JSON(w, http.StatusCreated, map[string]interface{}{
		"file": r.URL.Path,
		"size": len(body),
		"type": mimeType,
		"hash": fmt.Sprintf("%x", h.Sum(nil)),
	})
}

func (s Server) imageHandler(w http.ResponseWriter, r *http.Request) {
	if containsDotDot(r.URL.Path) {
		// Too many programs use r.URL.Path to construct the argument to
		// serveFile. Reject the request under the assumption that happened
		// here and ".." may not be wanted.
		// Note that name might not contain "..", for example if code (still
		// incorrectly) used filepath.Join(myDir, r.URL.Path).

		http.Error(w, "Invalid URL path", http.StatusBadRequest)

		return
	}

	options, err := ParamsFromContext(r.Context())
	if err != nil {
		xlog.Errorf("Error while parsing options: %v", err.Error())
		http.Error(w, "Error while parsing options", http.StatusBadRequest)

		return
	}
	// xlog.Infof("options: %#v", options)

	resource := &Resource{
		Path:    r.URL.Path,
		Options: options,
	}

	w.Header().Set("Link", `</worker/client-hints.js>; rel="serviceworker"`)

	// fetch from cache
	if resource, ok := s.cache.Get(resource); ok {
		w.Header().Set("X-Image-From", "cache")

		ServeImage(w, r, resource)

		return
	}

	resource, err = s.source.Get(resource)
	if err != nil {
		if os.IsNotExist(err) {
			msg := fmt.Sprintf("File %s not found", r.URL.Path)

			xlog.Info(msg)
			http.Error(w, msg, http.StatusNotFound)

			return
		}

		xlog.Error(err)
		http.Error(w, err.Error(), http.StatusNotFound)

		return
	}

	if err := ProcessImage(resource); err != nil {
		xlog.Errorf("Error while processing the image: %v", err.Error())
		http.Error(w, "Error while processing the image", http.StatusInternalServerError)

		return
	}

	w.Header().Set("X-Image-From", "source")

	ServeImage(w, r, resource)

	// save resource in cache
	go func(r *Resource) {
		if err := s.cache.Set(r); err != nil {
			xlog.Error(err)
		}
	}(resource)
}

func (s Server) serviceWorker(w http.ResponseWriter, r *http.Request) {
	js := `
self.addEventListener('install', function(event) {
    console.log('Service Worker installing.');
});

self.addEventListener('activate', function(event) {
    console.log('Service Worker activating.');
});

// Listen to fetch events
self.addEventListener('fetch', function(event) {
    console.log(event.request.url);
    //throw JSON.stringify({data: event.request.url});

    if (/\.(jpg|jpeg|png)(\?+)?$/.test(event.request.url)) {
        console.log('Client-Hint');

        // Clone the request
		var req = event.request.clone();

        req.headers.set('DPR', window.devicePixelRatio);
        req.url = req.url + '&dpr=' + window.devicePixelRatio;

        event.respondWith(fetch(req/*, {
            mode: 'no-cors',
            headers: headers
        }*/));
    }
});
`
	w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	w.Write([]byte(js))
}

func (s Server) docHandler(w http.ResponseWriter, r *http.Request) {
	name := "doc/index.html"

	w.Header().Set("Content-Type", "text/html; charset=utf-8")

	body, err := Asset(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	info, err := AssetInfo(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	http.ServeContent(
		w,
		r,
		info.Name(),
		info.ModTime(),
		memfs.NewBuffer(&body),
	)
}

func (s Server) swaggerHandler(w http.ResponseWriter, r *http.Request) {
	name := "swagger.yaml"

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

	body, err := Asset(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	info, err := AssetInfo(name)
	if err != nil {
		httputil.FailureFromError(w, 0, err)

		return
	}

	http.ServeContent(
		w,
		r,
		info.Name(),
		info.ModTime(),
		memfs.NewBuffer(&body),
	)
}

// ListenAndServe service
func (s *Server) ListenAndServe() {
	middleware := alice.New(
		NewLoggerHandler(),
		NewImageExtensionFilterHandler(s.config),
		c.Handler,
	)

	readMiddleware := middleware.Append(
		NewParamsHandler(),
		NewClientHintsHandler(),
	)

	authMiddleware := middleware.Append(
		NewAuthHandler(s.config.Auth),
	)

	s.mux.HandleFunc("/favicon.ico", s.notFoundHandler)
	s.mux.HandleFunc("/health", s.healthHandler)
	s.mux.HandleFunc("/worker/client-hints.js", s.serviceWorker)

	if s.config.Doc.Enable {
		s.mux.HandleFunc("/swagger.yaml", s.swaggerHandler)
		s.mux.HandleFunc("/doc/", s.docHandler)

		/*fs := http.FileServer(http.Dir("doc"))
		s.mux.Handle("/doc/", http.StripPrefix("/doc/", fs))*/
	}

	s.mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case http.MethodGet:
			readMiddleware.ThenFunc(s.imageHandler).ServeHTTP(w, r)
		case http.MethodPost:
			authMiddleware.ThenFunc(s.uploadHandler).ServeHTTP(w, r)
		case http.MethodDelete:
			authMiddleware.ThenFunc(s.deleteHandler).ServeHTTP(w, r)
		default:
			http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		}
	})

	srv := &http.Server{
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 10 * time.Second,
		Handler:      s.mux,
	}

	srv.Addr = fmt.Sprintf("%s:%d", s.config.Server.Host, s.config.Server.Port)

	xlog.Infof("Server running on %s", srv.Addr)

	xlog.Fatal(srv.ListenAndServe())
}
<|endoftext|>