Merge pull request #137 from ShadowJonathan/with-json

yunginnanet · Apr 19, 2024 · 53afe56 · 53afe56
2 parents 6e183db + 54e492c
commit 53afe56
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/heffalump/heffalump.go b/heffalump/heffalump.go
@@ -41,8 +41,16 @@ func NewDefaultHeffalump() *Heffalump {
 	return NewHeffalump(NewDefaultMarkovMap(), DefaultBuffSize)
 }
 
+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
 
@@ -55,8 +63,24 @@ func (h *Heffalump) WriteHell(bw *bufio.Writer) (int64, error) {
 	buf := h.pool.Get().([]byte)
 	defer h.pool.Put(buf)
 
-	if _, err = bw.WriteString("<html>\n<body>\n"); err != nil {
-		return n, err
+	switch cType {
+	case PlainText:
+		if _, err = bw.WriteString("# Chapter 1\n"); err != nil {
+			return n, err
+		}
+		break
+	case HTML:
+		if _, err = bw.WriteString("<html>\n<body>\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
@@ -30,6 +30,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 {
@@ -57,6 +71,7 @@ func hellPot(ctx *fasthttp.RequestCtx) {
 
 	slog.Info().Msg("NEW")
 
+	var cType = detectContentType(ctx)
 	s := time.Now()
 	var n int64
 
@@ -65,7 +80,7 @@ func hellPot(ctx *fasthttp.RequestCtx) {
 		var wn int64
 
 		for {
-			wn, err = hellpotHeffalump.WriteHell(bw)
+			wn, err = hellpotHeffalump.WriteHell(bw, cType)
 			n += wn
 			if err != nil {
 				slog.Trace().Err(err).Msg("END_ON_ERR")