113 lines
2.7 KiB
Go
113 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"html/template"
|
|
"net/http"
|
|
"waifu_gallery/s3"
|
|
|
|
_ "embed"
|
|
)
|
|
|
|
type LoadingPageContent struct {
|
|
Query string
|
|
GenerationEndpoint string
|
|
}
|
|
|
|
//go:embed loading.html
|
|
var loadingTmpl string
|
|
|
|
// You really wanna serve this endpoint behind WAF / Captcha or whatever and rate limit it to protect your precious tokens
|
|
func serveGenerationEndpoint(w http.ResponseWriter, r *http.Request) {
|
|
// HACK: Making this a /GET route to serve a CF challenge
|
|
|
|
// Validate user input
|
|
input, err := handleSubdomain(w, r)
|
|
if err != nil {
|
|
serveBadSubdomainError(w)
|
|
return
|
|
}
|
|
|
|
_, err = s3.GetMetadata(input)
|
|
// We want an error to happen here actually
|
|
if err == nil {
|
|
// Image already generated, bai bai
|
|
http.Redirect(w, r, "/", http.StatusMovedPermanently)
|
|
return
|
|
} else {
|
|
switch s3.ToErrorResponse(err).Code {
|
|
|
|
// Check that we error'd because the image doesn't exist
|
|
case "NoSuchKey":
|
|
// Generate picture or wait on existing generation job
|
|
// Probably not the best way but... if it works?
|
|
if ch, ok := generationQueue.Load(input); !ok {
|
|
// Check number of job requests
|
|
// TODO make this less horribly inefficient (lol)
|
|
var jobNbr int
|
|
generationQueue.Range(func(k, v interface{}) bool {
|
|
jobNbr++
|
|
return true
|
|
})
|
|
if config.MaxConcurrentJobs > 0 && jobNbr >= config.MaxConcurrentJobs {
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
w.Write([]byte("Too many images currently being generated, please try again in a few minutes."))
|
|
return
|
|
}
|
|
|
|
// Create channel
|
|
generationQueue.Store(input, make(chan string))
|
|
|
|
// Generate new picture
|
|
InfoLogger.Printf("generating with input '%s'\n", input)
|
|
err := createPicture(input) // where the magic happens
|
|
if err != nil {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
ErrorLogger.Println(err)
|
|
return
|
|
}
|
|
|
|
val, _ := generationQueue.Load(input)
|
|
generationQueue.Delete(input)
|
|
if c, ok := val.(chan string); ok {
|
|
// Non-blocking send to queue
|
|
select {
|
|
case c <- input:
|
|
default:
|
|
}
|
|
}
|
|
} else {
|
|
// Wait for job to finish if there's already a generation with the same input going on
|
|
if c, ok := ch.(chan string); ok {
|
|
<-c
|
|
}
|
|
}
|
|
default:
|
|
ErrorLogger.Println(err)
|
|
serveInternalError(w)
|
|
}
|
|
}
|
|
|
|
// Generation succeeded
|
|
w.Write([]byte("picture generated!"))
|
|
}
|
|
|
|
func writeWaitingPage(w http.ResponseWriter, input string) (err error) {
|
|
tmpl, err := template.New("loading").Parse(loadingTmpl)
|
|
if err != nil {
|
|
ErrorLogger.Println(err)
|
|
return
|
|
}
|
|
|
|
tmppData := LoadingPageContent{
|
|
Query: input,
|
|
GenerationEndpoint: "/generate",
|
|
}
|
|
err = tmpl.Execute(w, tmppData)
|
|
if err != nil {
|
|
ErrorLogger.Println(err)
|
|
return
|
|
}
|
|
|
|
return
|
|
}
|