diff --git a/heffalump/heffalump.go b/heffalump/heffalump.go index 298e5a9..9c0ea33 100644 --- a/heffalump/heffalump.go +++ b/heffalump/heffalump.go @@ -36,8 +36,16 @@ func NewHeffalump(mm MarkovMap, buffsize int) *Heffalump { } } +type ContentType int + +const ( + PlainText ContentType = iota + HTML + JSON +) + // WriteHell writes markov chain heffalump hell to the provided io.Writer -func (h *Heffalump) WriteHell(bw *bufio.Writer) (int64, error) { +func (h *Heffalump) WriteHell(bw *bufio.Writer, cType ContentType) (int64, error) { var n int64 var err error @@ -50,8 +58,21 @@ func (h *Heffalump) WriteHell(bw *bufio.Writer) (int64, error) { buf := h.pool.Get().([]byte) defer h.pool.Put(buf) - if _, err = bw.WriteString("\n\n"); err != nil { - return n, err + switch cType { + case PlainText: + break + case HTML: + if _, err = bw.WriteString("\n\n"); err != nil { + return n, err + } + break + case JSON: + if _, err = bw.WriteString("[\""); err != nil { + return n, err + } + break + default: + panic("unhandled default case") } if n, err = io.CopyBuffer(bw, h.mm, buf); err != nil { return n, nil diff --git a/internal/http/router.go b/internal/http/router.go index a89284c..a85930c 100644 --- a/internal/http/router.go +++ b/internal/http/router.go @@ -26,6 +26,20 @@ func getRealRemote(ctx *fasthttp.RequestCtx) string { return ctx.RemoteIP().String() } +func detectContentType(ctx *fasthttp.RequestCtx) (cType heffalump.ContentType) { + cType = heffalump.PlainText + + acceptHeader := string(ctx.Request.Header.Peek("Accept")) + + if strings.Contains(acceptHeader, "text/html") { + cType = heffalump.HTML + } else if strings.Contains(acceptHeader, "application/json") { + cType = heffalump.JSON + } + + return +} + func hellPot(ctx *fasthttp.RequestCtx) { path, pok := ctx.UserValue("path").(string) if len(path) < 1 || !pok { @@ -53,6 +67,7 @@ func hellPot(ctx *fasthttp.RequestCtx) { slog.Info().Msg("NEW") + var cType = detectContentType(ctx) s := time.Now() var n int64 @@ -61,7 +76,7 @@ func hellPot(ctx *fasthttp.RequestCtx) { var wn int64 for { - wn, err = heffalump.DefaultHeffalump.WriteHell(bw) + wn, err = heffalump.DefaultHeffalump.WriteHell(bw, cType) n += wn if err != nil { slog.Trace().Err(err).Msg("END_ON_ERR")