#gottem

yunginnanet · Sep 15, 2021 · 374858c · 374858c
1 parent a45440c
commit 374858c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 23 deletions.
diff --git a/heffalump/heffalump.go b/heffalump/heffalump.go
@@ -49,7 +49,10 @@ func (h *Heffalump) putBuffer(buf []byte) {
 
 // WriteHell writes markov chain heffalump hell to the provided io.Writer
 // https://github.com/carlmjohnson/heffalump
-func (h *Heffalump) WriteHell(bw *bufio.Writer) int64 {
+func (h *Heffalump) WriteHell(bw *bufio.Writer) (int64, error) {
+	var n int64
+	var err error
+
 	defer func() {
 		if r := recover(); r != nil {
 			log.Error().Interface("caller", r).Msg("panic recovered!")
@@ -59,15 +62,13 @@ func (h *Heffalump) WriteHell(bw *bufio.Writer) int64 {
 	buf := h.getBuffer()
 	defer h.putBuffer(buf)
 
-	if _, err := io.WriteString(bw, "<HTML>\n<BODY>\n"); err != nil {
-		log.Debug().Caller().Err(err).Msg("WriteString_fail")
+	if _, err = io.WriteString(bw, "<HTML>\n<BODY>\n"); err != nil {
+		return n, err
 	}
 
-	var n int64
-	var err error
 	if n, err = io.CopyBuffer(bw, h.mm, buf); err != nil {
-		log.Debug().Caller().Err(err).Msg("CopyBuffer_fail")
+		return n, nil
 	}
 
-	return n
+	return n, nil
 }
diff --git a/http/router.go b/http/router.go
@@ -35,21 +35,29 @@ func hellPot(ctx *fasthttp.RequestCtx) {
 	slog.Info().Msg("NEW")
 
 	s := time.Now()
-
 	var n int64
 
 	ctx.SetBodyStreamWriter(func(bw *bufio.Writer) {
-		n = heffalump.DefaultHeffalump.WriteHell(bw)
-	})
+		var err error
+		var wn int64
+
+		for {
+			wn, err = heffalump.DefaultHeffalump.WriteHell(bw)
+			n += wn
+			if err != nil {
+				slog.Debug().Err(err).Msg("END_ON_ERR")
+				break
+			}
+		}
 
+		slog.Info().
+			Int64("BYTES", n).
+			Dur("DURATION", time.Since(s)).
+			Msg("FINISH")
+	})
 
-	slog.Info().
-		Int64("BYTES", n).
-		Dur("DURATION", time.Since(s)).
-		Msg("FINISH")
 }
 
-
 func listenOnUnixSocket(addr string, r *router.Router) error {
 	var err error
 	var unixAddr *net.UnixAddr
@@ -82,13 +90,13 @@ func getSrv(r *router.Router) fasthttp.Server {
 			from fasthttp docs: "By default request read timeout is unlimited."
 			My thinking here is avoiding some sort of weird oversized GET query just in case.
 		*/
-		ReadTimeout: 5 * time.Second,
+		ReadTimeout:        5 * time.Second,
 		MaxRequestBodySize: 1 * 1024 * 1024,
 
 		// Help curb abuse of HellPot (we've always needed this badly)
-		MaxConnsPerIP: 10,
+		MaxConnsPerIP:      10,
 		MaxRequestsPerConn: 2,
-		Concurrency: config.MaxWorkers,
+		Concurrency:        config.MaxWorkers,
 
 		// only accept GET requests
 		GetOnly: true,
@@ -98,13 +106,9 @@ func getSrv(r *router.Router) fasthttp.Server {
 
 		CloseOnShutdown: true,
 
-
-
 		// No need to keepalive, our response is a sort of keep-alive ;)
 		DisableKeepalive: true,
 
-
-
 		Handler: r.Handler,
 	}
 }
@@ -121,7 +125,7 @@ func Serve() error {
 	}
 
 	srv := getSrv(r)
-	
+
 	if !config.UseUnixSocket {
 		log.Info().Str("caller", l).Msg("Listening and serving HTTP...")
 		return srv.ListenAndServe(l)