func (ctrl *ShortIDController) Respond(shortID string) (res *ResPkg) {
res = NewResPkg()
// extract the ID
id, random, err := data.FullID(shortID)
if err != nil {
res.HttpStatus = http.StatusNotFound
res.Errors = append(res.Errors, err)
log.Error("shortid_controller Respond error")
log.Error(err)
return
}
shawty, err := ctrl.sh.GetByID(id, random)
if err != nil {
res.HttpStatus = http.StatusNotFound
res.Errors = append(res.Errors, err)
log.Error("shortid_controller Respond error")
log.Error(err)
return
}
ctrl.sh.IncHits(shawty.ID) // increase hit
res.Data["Domain"] = ctrl.config["SHAWTY_DOMAIN"]
res.Data["Shawty"] = shawty
res.HttpStatus = http.StatusMovedPermanently
return
}
func (controller *ShawtyJSController) ServeHTTP(w http.ResponseWriter, r *http.Request) {
defer closeReqBody(r)
var u = r.FormValue("url")
if u == "" {
u = r.Referer()
if u != "" {
// if it's in the referer then we don't want the browser to cache this
w.Header().Set("Cache-Control", "no-cache")
http.Redirect(w, r, "shawty.js?url="+url.QueryEscape(u), http.StatusTemporaryRedirect)
return
} else {
http.NotFound(w, r)
return
}
}
res := controller.GetJSResponse(u, utils.HttpRemoteIP(r), r.FormValue("bm") == "1")
tpl := getShawtyJs()
w.Header().Set("Content-Type", "application/javascript")
if err := tpl.Execute(w, res); err != nil {
log.Error("Cannot execute shawty javascript template")
log.Error(err)
http.Error(w, "Error", http.StatusInternalServerError)
}
}
func (page *HomeController) ServeHTTP(w http.ResponseWriter, r *http.Request) {
defer closeReqBody(r)
res := page.Index()
tpl := getIndexHtml()
if err := tpl.Execute(w, res); err != nil {
log.Error("Cannot execute index template")
log.Error(err)
http.Error(w, "Error", http.StatusInternalServerError)
}
}
func main() {
runtime.GOMAXPROCS(runtime.NumCPU())
// read configurations
confKeys := []string{
"SHAWTY_PORT", "SHAWTY_DB_TYPE", "SHAWTY_DB", "SHAWTY_DOMAIN",
"SHAWTY_MODE", "SHAWTY_LPM", "SHAWTY_LOG_DIR",
}
config := make(map[string]string)
for _, k := range confKeys {
config[k] = os.Getenv(k)
}
// setup logger
log.SetDir(config["SHAWTY_LOG_DIR"])
// setup data
random := utils.NewBestRand()
dbType := strings.ToLower(config["SHAWTY_DB_TYPE"])
var shawties data.Shawties
var err error
if dbType == "pg" || dbType == "postgres" || dbType == "postgresql" {
shawties, err = data.NewPgSh(random, config["SHAWTY_DB"])
} else if dbType == "mysql" {
shawties, err = data.NewMySh(random, config["SHAWTY_DB"])
} else {
err = errors.New("Unknown database type. Please check $SHAWTY_DB_TYPE")
}
if err != nil {
log.Error("Cannot create data source")
log.Error(err)
return
}
defer shawties.Close()
// register routes
home := web.NewHomeController(config)
shawtyjs := web.NewShawtyJSController(config, shawties)
shortID := web.NewShortIDController(config, shawties)
// setup HTTP server
router := mux.NewRouter()
router.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("static/"))))
router.Handle("/", home)
router.Handle("/shawty.js", shawtyjs)
router.Handle("/{shortID:[A-Za-z0-9]+}", shortID)
var port = config["SHAWTY_PORT"]
if port == "" {
port = "80"
}
l, err := net.Listen("tcp", "0.0.0.0:"+port)
if err != nil {
log.Errorf("Cannot listen at %s", port)
fmt.Println(err)
return
}
defer l.Close()
log.Infof("Listening at %s", port)
runMode := strings.ToLower(config["SHAWTY_MODE"])
switch runMode {
case "fcgi":
log.Info("Serving FCGI mode")
fcgi.Serve(l, router)
default:
log.Info("Serving HTTP mode")
http.Serve(l, router)
}
}